diff --git a/.githooks/prepare-commit-msg b/.githooks/prepare-commit-msg deleted file mode 100755 index fa51f80df..000000000 --- a/.githooks/prepare-commit-msg +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/sh -grep -qs "^Signed-off-by: " "$1" || { - name=$(git config user.name) - email=$(git config user.email) - - if [ -n "$name" ] && [ -n "$email" ]; then - echo "" >> "$1" - echo "Signed-off-by: $name <$email>" >> "$1" - fi -} \ No newline at end of file diff --git a/.github/actions/stamp-version/action.yml b/.github/actions/stamp-version/action.yml new file mode 100644 index 000000000..2a40301c2 --- /dev/null +++ b/.github/actions/stamp-version/action.yml @@ -0,0 +1,24 @@ +name: Stamp version +description: Resolve the release version and write python/cactus/_version.py +runs: + using: composite + steps: + - shell: bash + env: + TAG: ${{ github.event.release.tag_name }} + INPUT_VERSION: ${{ github.event.inputs.version }} + run: | + if [ -n "$TAG" ]; then + raw="${TAG#v}" + elif [ -n "$INPUT_VERSION" ]; then + raw="${INPUT_VERSION#v}" + else + raw=$(sed 's/^v//' CACTUS_VERSION) + fi + version=$(RAW="$raw" python3 -c 'import os; r=os.environ["RAW"]; print(r + ".0" if r.count(".") == 1 else r)') + if [ -z "$version" ]; then + echo "::error::No version source available (TAG, INPUT_VERSION, CACTUS_VERSION all empty)" + exit 1 + fi + printf '__version__ = "%s"\n' "$version" > python/cactus/_version.py + echo "Stamped version: $version" diff --git a/.github/assemble-docs.sh b/.github/assemble-docs.sh new file mode 100755 index 000000000..5c329452d --- /dev/null +++ b/.github/assemble-docs.sh @@ -0,0 +1,223 @@ +#!/usr/bin/env bash +# Assembles site_docs/ from scattered source files and rewrites internal links. +# Used by the release.yml workflow. +set -euo pipefail + +DOCS_VERSION="${1:-}" + +if sed --version >/dev/null 2>&1; then + sedi() { sed -i "$@"; } +else + sedi() { sed -i '' "$@"; } +fi + +rm -rf site_docs +mkdir -p site_docs/docs site_docs/python site_docs/apple site_docs/android \ + site_docs/flutter site_docs/rust site_docs/swift site_docs/kotlin \ + site_docs/react-native site_docs/blog site_docs/assets + +cp -r assets/* site_docs/assets/ + +echo "docs.cactuscompute.com" > site_docs/CNAME + +mkdir -p site_docs/stylesheets site_docs/javascripts +cp .github/docs-overrides/stylesheets/custom.css site_docs/stylesheets/custom.css +cp .github/docs-overrides/javascripts/mathjax.js site_docs/javascripts/mathjax.js + +[ -f CONTRIBUTING.md ] && cp CONTRIBUTING.md site_docs/CONTRIBUTING.md +[ -f DCO.md ] && cp DCO.md site_docs/DCO.md + +cp docs/*.md site_docs/docs/ + +cp python/README.md site_docs/python/README.md +cp apple/README.md site_docs/apple/README.md +cp android/README.md site_docs/android/README.md +cp bindings/flutter/README.md site_docs/flutter/README.md +cp bindings/swift/README.md site_docs/swift/README.md +cp bindings/kotlin/README.md site_docs/kotlin/README.md +cp bindings/react-native/README.md site_docs/react-native/README.md +cp bindings/rust/README.md site_docs/rust/README.md + +if [ -d blog ] && ls blog/*.md >/dev/null 2>&1; then + cp blog/*.md site_docs/blog/ +fi + +{ + echo '---' + echo 'title: "Cactus"' + echo 'description: "Energy-efficient AI inference engine for phones, wearables, Macs, and ARM devices."' + echo '---' + echo '' + cat README.md +} > site_docs/index.md + +sedi 's/^# Cactus$//' site_docs/index.md + +sedi 's|(/docs/cactus_engine\.md)|(docs/cactus_engine.md)|g' site_docs/index.md +sedi 's|(/docs/cactus_graph\.md)|(docs/cactus_graph.md)|g' site_docs/index.md +sedi 's|(/docs/cactus_index\.md)|(docs/cactus_index.md)|g' site_docs/index.md +sedi 's|(/docs/cactus_kernels\.md)|(docs/cactus_kernels.md)|g' site_docs/index.md +sedi 's|(/docs/cactus_quants\.md)|(docs/cactus_quants.md)|g' site_docs/index.md +sedi 's|(/docs/cactus_transpiler\.md)|(docs/cactus_transpiler.md)|g' site_docs/index.md +sedi 's|(/docs/cactus_hybrid\.md)|(docs/cactus_hybrid.md)|g' site_docs/index.md +sedi 's|(/docs/finetuning\.md)|(docs/finetuning.md)|g' site_docs/index.md +sedi 's|(/docs/compatibility\.md)|(docs/compatibility.md)|g' site_docs/index.md +sedi 's|(/docs/quickstart\.md)|(docs/quickstart.md)|g' site_docs/index.md +sedi 's|(/docs/choose-bindings\.md)|(docs/choose-bindings.md)|g' site_docs/index.md +sedi 's|(/CONTRIBUTING\.md)|(CONTRIBUTING.md)|g' site_docs/index.md +sedi 's|(/bindings/swift/)|(swift/README.md)|g' site_docs/index.md +sedi 's|(/bindings/kotlin/)|(kotlin/README.md)|g' site_docs/index.md +sedi 's|(/bindings/python/)|(python/README.md)|g' site_docs/index.md +sedi 's|(/bindings/react-native/)|(react-native/README.md)|g' site_docs/index.md +sedi 's|(/bindings/flutter/)|(flutter/README.md)|g' site_docs/index.md +sedi 's|(/bindings/rust/)|(rust/README.md)|g' site_docs/index.md +sedi 's|(/python/)|(python/README.md)|g' site_docs/index.md +sedi 's|(/apple/)|(apple/README.md)|g' site_docs/index.md +sedi 's|(/android/)|(android/README.md)|g' site_docs/index.md +sedi 's|(/blog/hybrid_transcription\.md)|(blog/hybrid_transcription.md)|g' site_docs/index.md +sedi 's|(/blog/lfm2_24b_a2b\.md)|(blog/lfm2_24b_a2b.md)|g' site_docs/index.md +sedi 's|(/blog/parakeet\.md)|(blog/parakeet.md)|g' site_docs/index.md +sedi 's|(/blog/lfm2\.5_350m\.md)|(blog/lfm2.5_350m.md)|g' site_docs/index.md +sedi 's|(/blog/gemma4\.md)|(blog/gemma4.md)|g' site_docs/index.md +sedi 's|(/blog/turboquant-h\.md)|(blog/turboquant-h.md)|g' site_docs/index.md + +for f in site_docs/docs/*.md; do + sedi 's|(/docs/cactus_engine\.md)|(cactus_engine.md)|g' "$f" + sedi 's|(/docs/cactus_graph\.md)|(cactus_graph.md)|g' "$f" + sedi 's|(/docs/cactus_index\.md)|(cactus_index.md)|g' "$f" + sedi 's|(/docs/cactus_kernels\.md)|(cactus_kernels.md)|g' "$f" + sedi 's|(/docs/cactus_quants\.md)|(cactus_quants.md)|g' "$f" + sedi 's|(/docs/cactus_transpiler\.md)|(cactus_transpiler.md)|g' "$f" + sedi 's|(/docs/cactus_hybrid\.md)|(cactus_hybrid.md)|g' "$f" + sedi 's|(/docs/finetuning\.md)|(finetuning.md)|g' "$f" + sedi 's|(/docs/compatibility\.md)|(compatibility.md)|g' "$f" + sedi 's|(/docs/quickstart\.md)|(quickstart.md)|g' "$f" + sedi 's|(/docs/choose-bindings\.md)|(choose-bindings.md)|g' "$f" + sedi 's|(/docs/index\.md)|(../index.md)|g' "$f" + sedi 's|(/blog/hybrid_transcription\.md)|(../blog/hybrid_transcription.md)|g' "$f" + sedi 's|(/blog/lfm2_24b_a2b\.md)|(../blog/lfm2_24b_a2b.md)|g' "$f" + sedi 's|(/blog/parakeet\.md)|(../blog/parakeet.md)|g' "$f" + sedi 's|(/blog/lfm2\.5_350m\.md)|(../blog/lfm2.5_350m.md)|g' "$f" + sedi 's|(/blog/gemma4\.md)|(../blog/gemma4.md)|g' "$f" + sedi 's|(/blog/turboquant-h\.md)|(../blog/turboquant-h.md)|g' "$f" + sedi 's|(/CONTRIBUTING\.md)|(../CONTRIBUTING.md)|g' "$f" + sedi 's|(/bindings/swift/)|(../swift/README.md)|g' "$f" + sedi 's|(/bindings/kotlin/)|(../kotlin/README.md)|g' "$f" + sedi 's|(/bindings/python/)|(../python/README.md)|g' "$f" + sedi 's|(/bindings/react-native/)|(../react-native/README.md)|g' "$f" + sedi 's|(/bindings/flutter/)|(../flutter/README.md)|g' "$f" + sedi 's|(/bindings/rust/)|(../rust/README.md)|g' "$f" + sedi 's|(/python/)|(../python/README.md)|g' "$f" + sedi 's|(/apple/)|(../apple/README.md)|g' "$f" + sedi 's|(/android/)|(../android/README.md)|g' "$f" +done + +for f in site_docs/python/README.md site_docs/apple/README.md site_docs/android/README.md site_docs/flutter/README.md site_docs/swift/README.md site_docs/kotlin/README.md site_docs/react-native/README.md site_docs/rust/README.md; do + sedi 's|(/docs/cactus_engine\.md)|(../docs/cactus_engine.md)|g' "$f" + sedi 's|(/docs/cactus_graph\.md)|(../docs/cactus_graph.md)|g' "$f" + sedi 's|(/docs/cactus_index\.md)|(../docs/cactus_index.md)|g' "$f" + sedi 's|(/docs/cactus_kernels\.md)|(../docs/cactus_kernels.md)|g' "$f" + sedi 's|(/docs/cactus_quants\.md)|(../docs/cactus_quants.md)|g' "$f" + sedi 's|(/docs/cactus_transpiler\.md)|(../docs/cactus_transpiler.md)|g' "$f" + sedi 's|(/docs/cactus_hybrid\.md)|(../docs/cactus_hybrid.md)|g' "$f" + sedi 's|(/docs/finetuning\.md)|(../docs/finetuning.md)|g' "$f" + sedi 's|(/docs/compatibility\.md)|(../docs/compatibility.md)|g' "$f" + sedi 's|(/docs/quickstart\.md)|(../docs/quickstart.md)|g' "$f" + sedi 's|(/docs/choose-bindings\.md)|(../docs/choose-bindings.md)|g' "$f" + sedi 's|(/docs/index\.md)|(../index.md)|g' "$f" + sedi 's|(/blog/hybrid_transcription\.md)|(../blog/hybrid_transcription.md)|g' "$f" + sedi 's|(/blog/lfm2_24b_a2b\.md)|(../blog/lfm2_24b_a2b.md)|g' "$f" + sedi 's|(/blog/parakeet\.md)|(../blog/parakeet.md)|g' "$f" + sedi 's|(/blog/lfm2\.5_350m\.md)|(../blog/lfm2.5_350m.md)|g' "$f" + sedi 's|(/blog/gemma4\.md)|(../blog/gemma4.md)|g' "$f" + sedi 's|(/blog/turboquant-h\.md)|(../blog/turboquant-h.md)|g' "$f" + sedi 's|(/CONTRIBUTING\.md)|(../CONTRIBUTING.md)|g' "$f" + sedi 's|(/bindings/swift/)|(../swift/README.md)|g' "$f" + sedi 's|(/bindings/kotlin/)|(../kotlin/README.md)|g' "$f" + sedi 's|(/bindings/python/)|(../python/README.md)|g' "$f" + sedi 's|(/bindings/react-native/)|(../react-native/README.md)|g' "$f" + sedi 's|(/bindings/flutter/)|(../flutter/README.md)|g' "$f" + sedi 's|(/bindings/rust/)|(../rust/README.md)|g' "$f" + sedi 's|(/python/)|(../python/README.md)|g' "$f" + sedi 's|(/apple/)|(../apple/README.md)|g' "$f" + sedi 's|(/android/)|(../android/README.md)|g' "$f" + sedi 's|(\.\.\/README\.md)|(../index.md)|g' "$f" +done + +if ls site_docs/blog/*.md >/dev/null 2>&1; then + for f in site_docs/blog/*.md; do + sedi 's|(/docs/cactus_engine\.md)|(../docs/cactus_engine.md)|g' "$f" + sedi 's|(/docs/cactus_graph\.md)|(../docs/cactus_graph.md)|g' "$f" + sedi 's|(/docs/cactus_index\.md)|(../docs/cactus_index.md)|g' "$f" + sedi 's|(/docs/cactus_kernels\.md)|(../docs/cactus_kernels.md)|g' "$f" + sedi 's|(/docs/cactus_quants\.md)|(../docs/cactus_quants.md)|g' "$f" + sedi 's|(/docs/cactus_transpiler\.md)|(../docs/cactus_transpiler.md)|g' "$f" + sedi 's|(/docs/cactus_hybrid\.md)|(../docs/cactus_hybrid.md)|g' "$f" + sedi 's|(/docs/finetuning\.md)|(../docs/finetuning.md)|g' "$f" + sedi 's|(/docs/compatibility\.md)|(../docs/compatibility.md)|g' "$f" + sedi 's|(/docs/quickstart\.md)|(../docs/quickstart.md)|g' "$f" + sedi 's|(/docs/choose-bindings\.md)|(../docs/choose-bindings.md)|g' "$f" + sedi 's|(/docs/index\.md)|(../index.md)|g' "$f" + sedi 's|(/CONTRIBUTING\.md)|(../CONTRIBUTING.md)|g' "$f" + sedi 's|(/bindings/swift/)|(../swift/README.md)|g' "$f" + sedi 's|(/bindings/kotlin/)|(../kotlin/README.md)|g' "$f" + sedi 's|(/bindings/python/)|(../python/README.md)|g' "$f" + sedi 's|(/bindings/react-native/)|(../react-native/README.md)|g' "$f" + sedi 's|(/bindings/flutter/)|(../flutter/README.md)|g' "$f" + sedi 's|(/bindings/rust/)|(../rust/README.md)|g' "$f" + sedi 's|(/python/)|(../python/README.md)|g' "$f" + sedi 's|(/apple/)|(../apple/README.md)|g' "$f" + sedi 's|(/android/)|(../android/README.md)|g' "$f" + sedi 's|(/blog/hybrid_transcription\.md)|(hybrid_transcription.md)|g' "$f" + sedi 's|(/blog/lfm2_24b_a2b\.md)|(lfm2_24b_a2b.md)|g' "$f" + sedi 's|(/blog/parakeet\.md)|(parakeet.md)|g' "$f" + sedi 's|(/blog/lfm2\.5_350m\.md)|(lfm2.5_350m.md)|g' "$f" + sedi 's|(/blog/gemma4\.md)|(gemma4.md)|g' "$f" + sedi 's|(/blog/turboquant-h\.md)|(turboquant-h.md)|g' "$f" + done +fi + +if [ -f site_docs/CONTRIBUTING.md ]; then + sedi 's|(/docs/cactus_engine\.md)|(docs/cactus_engine.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/cactus_graph\.md)|(docs/cactus_graph.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/cactus_index\.md)|(docs/cactus_index.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/cactus_kernels\.md)|(docs/cactus_kernels.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/cactus_quants\.md)|(docs/cactus_quants.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/cactus_transpiler\.md)|(docs/cactus_transpiler.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/cactus_hybrid\.md)|(docs/cactus_hybrid.md)|g' site_docs/CONTRIBUTING.md + sedi 's|(/docs/index\.md)|(index.md)|g' site_docs/CONTRIBUTING.md +fi + +if [ -n "$DOCS_VERSION" ]; then + { + echo "!!! note \"Version ${DOCS_VERSION}\"" + echo " You're viewing docs for **${DOCS_VERSION}**. If you are cloning the repository, make sure to check out this release: \`git checkout ${DOCS_VERSION}\`" + echo "" + cat site_docs/docs/quickstart.md + } > site_docs/docs/quickstart.tmp && mv site_docs/docs/quickstart.tmp site_docs/docs/quickstart.md +fi + +for nav_path in \ + "python/README.md" \ + "swift/README.md" \ + "kotlin/README.md" \ + "flutter/README.md" \ + "react-native/README.md" \ + "rust/README.md" \ + "blog/README.md" \ + "blog/hybrid_transcription.md" \ + "blog/lfm2_24b_a2b.md" \ + "blog/parakeet.md" \ + "blog/lfm2.5_350m.md" \ + "blog/gemma4.md" \ + "blog/turboquant-h.md" \ + "CONTRIBUTING.md" \ + "docs/compatibility.md"; do + if [ ! -f "site_docs/$nav_path" ]; then + grep -vF "$nav_path" mkdocs.yml > mkdocs.yml.tmp && mv mkdocs.yml.tmp mkdocs.yml + fi +done + +if ! ls site_docs/blog/*.md >/dev/null 2>&1; then + grep -v "^ - Blog:" mkdocs.yml > mkdocs.yml.tmp && mv mkdocs.yml.tmp mkdocs.yml +fi diff --git a/.github/docs-overrides/javascripts/mathjax.js b/.github/docs-overrides/javascripts/mathjax.js new file mode 100644 index 000000000..80e81ba59 --- /dev/null +++ b/.github/docs-overrides/javascripts/mathjax.js @@ -0,0 +1,12 @@ +window.MathJax = { + tex: { + inlineMath: [["\\(", "\\)"]], + displayMath: [["\\[", "\\]"]], + processEscapes: true, + processEnvironments: true + }, + options: { + ignoreHtmlClass: ".*|", + processHtmlClass: "arithmatex" + } +}; diff --git a/.github/docs-overrides/main.html b/.github/docs-overrides/main.html new file mode 100644 index 000000000..0af326afb --- /dev/null +++ b/.github/docs-overrides/main.html @@ -0,0 +1,8 @@ +{% extends "base.html" %} + +{% block outdated %} + You're not viewing the latest version. + + Click here to go to latest. + +{% endblock %} diff --git a/.github/docs-overrides/stylesheets/custom.css b/.github/docs-overrides/stylesheets/custom.css new file mode 100644 index 000000000..5ce45acc6 --- /dev/null +++ b/.github/docs-overrides/stylesheets/custom.css @@ -0,0 +1,50 @@ +@font-face { + font-family: 'Geist'; + src: url('https://cdn.jsdelivr.net/npm/geist@1.3.0/dist/fonts/geist-sans/Geist-Regular.woff2') format('woff2'); + font-weight: 400; + font-display: swap; +} +@font-face { + font-family: 'Geist'; + src: url('https://cdn.jsdelivr.net/npm/geist@1.3.0/dist/fonts/geist-sans/Geist-Medium.woff2') format('woff2'); + font-weight: 500; + font-display: swap; +} +@font-face { + font-family: 'Geist'; + src: url('https://cdn.jsdelivr.net/npm/geist@1.3.0/dist/fonts/geist-sans/Geist-Bold.woff2') format('woff2'); + font-weight: 700; + font-display: swap; +} +@font-face { + font-family: 'Geist Mono'; + src: url('https://cdn.jsdelivr.net/npm/geist@1.3.0/dist/fonts/geist-mono/GeistMono-Regular.woff2') format('woff2'); + font-weight: 400; + font-display: swap; +} + +:root { + --md-text-font: 'Geist', -apple-system, BlinkMacSystemFont, sans-serif; + --md-code-font: 'Geist Mono', 'SFMono-Regular', 'Menlo', monospace; +} + +[data-md-color-scheme="slate"] { + --md-default-bg-color: #000000; + --md-footer-bg-color: #000000; + --md-footer-bg-color--dark: #000000; +} + +[data-md-color-scheme="slate"] .md-header, +[data-md-color-scheme="slate"] .md-tabs, +[data-md-color-scheme="slate"] .md-sidebar { + background-color: #000000; +} + +.md-banner--warning { + background-color: #d4a017; + color: #000000; +} + +.md-banner--warning a { + color: #000000; +} diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml new file mode 100644 index 000000000..e05a224e4 --- /dev/null +++ b/.github/workflows/build.yml @@ -0,0 +1,46 @@ +name: Build + +on: + pull_request: + paths: + - cactus-engine/** + - cactus-graph/** + - cactus-kernels/** + - apple/** + - android/** + - .github/workflows/build.yml + push: + branches: [main] + paths: + - cactus-engine/** + - cactus-graph/** + - cactus-kernels/** + - apple/** + - android/** + - .github/workflows/build.yml + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + apple: + runs-on: macos-15 + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - run: bash cactus-engine/build.sh + - run: BUILD_STATIC=true BUILD_XCFRAMEWORK=false bash apple/build.sh + - run: bash android/build.sh + + linux: + runs-on: ubuntu-24.04-arm + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - run: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev + - run: bash cactus-engine/build.sh diff --git a/.github/workflows/cpp.yml b/.github/workflows/cpp.yml new file mode 100644 index 000000000..a75ad5312 --- /dev/null +++ b/.github/workflows/cpp.yml @@ -0,0 +1,42 @@ +name: Unit Tests + +on: + pull_request: + paths: + - cactus-kernels/** + - cactus-graph/** + - .github/workflows/cpp.yml + push: + branches: [main] + paths: + - cactus-kernels/** + - cactus-graph/** + - .github/workflows/cpp.yml + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit-tests: + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + include: + - os: macos-15 + deps: brew install ccache ninja + - os: ubuntu-24.04-arm + deps: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev ccache ninja-build + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - run: ${{ matrix.deps }} + - name: Run kernel tests + run: bash cactus-kernels/test.sh + - name: Run graph tests + run: bash cactus-graph/test.sh diff --git a/.github/workflows/dco.yml b/.github/workflows/dco.yml index 1d12852dd..fdbf6d45c 100644 --- a/.github/workflows/dco.yml +++ b/.github/workflows/dco.yml @@ -4,17 +4,35 @@ on: pull_request: types: [opened, synchronize, reopened] +permissions: + contents: read + +concurrency: + group: dco-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + jobs: dco-check: runs-on: ubuntu-latest + timeout-minutes: 5 name: DCO Check steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: fetch-depth: 0 - name: Check DCO - uses: dco-action/check@v1 - with: - token: ${{ secrets.GITHUB_TOKEN }} \ No newline at end of file + run: | + set -euo pipefail + + missing=0 + while read -r commit; do + author="$(git show -s --format='%an <%ae>' "$commit")" + if ! git show -s --format='%B' "$commit" | grep -Fqi "Signed-off-by: $author"; then + echo "::error::Commit $commit is missing Signed-off-by: $author" + missing=1 + fi + done < <(git rev-list --no-merges "${{ github.event.pull_request.base.sha }}..${{ github.event.pull_request.head.sha }}") + + exit "$missing" diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 000000000..5b76c753e --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,74 @@ +name: Deploy Docs + +on: + push: + branches: [main] + paths: + - docs/** + - blog/** + - assets/** + - mkdocs.yml + - CACTUS_VERSION + - CONTRIBUTING.md + - DCO.md + - .github/assemble-docs.sh + - .github/docs-overrides/** + - .github/workflows/docs.yml + - python/README.md + - apple/README.md + - android/README.md + - bindings/flutter/README.md + - bindings/kotlin/README.md + - bindings/react-native/README.md + - bindings/rust/README.md + - bindings/swift/README.md + - README.md + workflow_call: + inputs: + version: + type: string + required: false + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: docs-deploy + cancel-in-progress: false + +jobs: + deploy: + runs-on: ubuntu-latest + timeout-minutes: 15 + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + fetch-tags: true + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: pip install "mkdocs==1.6.1" "mkdocs-material==9.7.5" "mike==2.1.4" + + - name: Build and deploy docs + run: | + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + TAG="${{ inputs.version || '' }}" + [ -z "$TAG" ] && TAG=$(cat CACTUS_VERSION | tr -d '[:space:]') + bash .github/assemble-docs.sh "$TAG" + mike deploy "$TAG" latest --update-aliases --push + mike set-default latest --push + + git checkout gh-pages + echo "docs.cactuscompute.com" > CNAME + git add CNAME + git commit -m "Preserve CNAME for custom domain" --allow-empty || true + git push origin gh-pages + git checkout - diff --git a/.github/workflows/inference.yml b/.github/workflows/inference.yml new file mode 100644 index 000000000..d8cac3705 --- /dev/null +++ b/.github/workflows/inference.yml @@ -0,0 +1,112 @@ +name: Inference Tests + +on: + pull_request_review: + types: [submitted] + workflow_dispatch: + +permissions: + contents: read + actions: read + +concurrency: + group: inference-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + gate: + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + approved: ${{ steps.check.outputs.approved }} + steps: + - id: check + run: | + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + echo "approved=true" >> $GITHUB_OUTPUT + elif [ "${{ github.event.review.state }}" = "approved" ]; then + echo "approved=true" >> $GITHUB_OUTPUT + else + echo "approved=false" >> $GITHUB_OUTPUT + fi + + runner-check: + needs: gate + if: needs.gate.outputs.approved == 'true' + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + available: ${{ steps.check.outputs.available }} + steps: + - id: check + env: + GH_TOKEN: ${{ github.token }} + REPO: ${{ github.repository }} + run: | + set -euo pipefail + status=$(curl -sS -L -w "%{http_code}" -o runners.json \ + -H "Authorization: Bearer ${GH_TOKEN}" \ + -H "Accept: application/vnd.github+json" \ + -H "X-GitHub-Api-Version: 2022-11-28" \ + "https://api.github.com/repos/${REPO}/actions/runners?per_page=100" || true) + if [ "$status" != "200" ]; then + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::notice::Could not inspect self-hosted runners (HTTP $status); skipping inference tests instead of queueing." + exit 0 + fi + count=$(jq '[.runners[] | select(.status == "online" and .busy == false) | select(([.labels[].name] | contains(["self-hosted", "linux", "arm64"])))] | length' runners.json) + if [ "$count" -gt 0 ]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::notice::No idle self-hosted linux arm64 runner is available; skipping inference tests." + fi + + build: + needs: [gate, runner-check] + if: needs.gate.outputs.approved == 'true' && needs.runner-check.outputs.available == 'true' + runs-on: [self-hosted, linux, arm64] + timeout-minutes: 60 + steps: + - uses: actions/checkout@v4 + with: + clean: false + + - name: Setup Python environment + run: | + if [ ! -f venv/bin/activate ]; then + source ./setup + fi + + - name: Build libcactus + engine tests + run: | + bash cactus-engine/build.sh + cd cactus-engine/tests + rm -rf build && mkdir -p build && cd build + cmake .. -DCMAKE_BUILD_TYPE=Release + make -j$(nproc) + + model-tests: + needs: build + runs-on: [self-hosted, linux, arm64] + timeout-minutes: 60 + env: + CACTUS_NO_CLOUD_TELE: "1" + strategy: + fail-fast: false + matrix: + model: + - LiquidAI/LFM2-VL-450M + - openai/whisper-small + steps: + - uses: actions/checkout@v4 + with: + clean: false + + - name: Run tests for ${{ matrix.model }} + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + source venv/bin/activate + MODEL="${{ github.event.inputs.model || matrix.model }}" + python3 -m cactus test --component engine --model "$MODEL" diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml deleted file mode 100644 index c3b8fe8d2..000000000 --- a/.github/workflows/publish.yml +++ /dev/null @@ -1,34 +0,0 @@ -name: Publish Models - -on: - push: - tags: - - 'v*' - -jobs: - publish: - runs-on: macos-latest - steps: - - uses: actions/checkout@v4 - - - uses: actions/checkout@v4 - with: - repository: cactus-compute/cactus-pro - token: ${{ secrets.CACTUS_PRO_TOKEN }} - path: cactus-pro - - - uses: actions/setup-python@v5 - with: - python-version: '3.12' - - - run: pip install -r tools/requirements.txt - - - run: | - cd tools - python -m src.publish_to_hf \ - --version ${GITHUB_REF#refs/tags/} \ - --org Cactus-Compute \ - --precision MIXED \ - --bits 8 - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml new file mode 100644 index 000000000..9cf97fcd3 --- /dev/null +++ b/.github/workflows/pypi.yml @@ -0,0 +1,161 @@ +name: PyPI Publish + +on: + workflow_call: + workflow_dispatch: + inputs: + version: + description: "Version to publish (e.g. 2.0.0). Leave empty to read from CACTUS_VERSION." + required: false + type: string + +permissions: + contents: read + +jobs: + wheels-macos: + runs-on: macos-26 + timeout-minutes: 45 + env: + MACOSX_DEPLOYMENT_TARGET: "14.0" + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - uses: ./.github/actions/stamp-version + + - name: Install SDL2 (live microphone transcription) + run: brew install sdl2 + + - name: Build native library + run: | + pip install -e python/ + cactus build + mkdir -p python/cactus/bindings/lib + cp cactus-engine/build/libcactus_engine.dylib python/cactus/bindings/lib/ + + - uses: actions/setup-node@v4 + with: + node-version: "22" + + - name: Bundle coding agent into the wheel + run: | + (cd cactus-code && npm install --no-audit --no-fund) + bash cactus-code/scripts/bundle-for-pypi.sh + rm -rf python/cactus/code/packages/tui/native/win32 + rm -rf python/cactus/code/packages/tui/native/darwin/prebuilds/darwin-x64 + + - name: Build and repair wheel + working-directory: python + run: | + pip install build delocate wheel + python -m build --wheel -C="--build-option=--plat-name=macosx_14_0_arm64" + env -u MACOSX_DEPLOYMENT_TARGET delocate-wheel --require-archs arm64 -w wheelhouse dist/*.whl + whl="$(ls "$PWD"/wheelhouse/*.whl)" + work="$(mktemp -d)" + python -m wheel unpack "$whl" -d "$work" + pkg="$(ls -d "$work"/*/)" + cp "$(brew --prefix sdl3)/lib/libSDL3.0.dylib" "${pkg}cactus/.dylibs/libSDL3.dylib" + chmod +w "${pkg}cactus/.dylibs/libSDL3.dylib" + for lib in "${pkg}cactus/.dylibs/libSDL2-2.0.0.dylib" "${pkg}cactus/.dylibs/libSDL3.dylib"; do + vtool -set-build-version macos 14.0 14.0 -replace -output "$lib" "$lib" + codesign -f -s - "$lib" + done + sed -i '' -E 's/macosx_[0-9]+_[0-9]+_arm64/macosx_14_0_arm64/' "$(ls -d "${pkg}"*.dist-info)/WHEEL" + rm "$whl" + python -m wheel pack "$pkg" -d "$PWD/wheelhouse" + + - uses: actions/upload-artifact@v4 + with: + name: wheel-macos-arm64 + path: python/wheelhouse/*.whl + + wheels-linux: + runs-on: ubuntu-24.04-arm + timeout-minutes: 45 + container: + image: quay.io/pypa/manylinux_2_28_aarch64 + steps: + - uses: actions/checkout@v4 + + - name: Provision toolchain + run: | + dnf install -y epel-release + dnf install -y libcurl-devel openssl-devel SDL2-devel pkgconfig + echo "/opt/python/cp312-cp312/bin" >> "$GITHUB_PATH" + + - uses: ./.github/actions/stamp-version + + - name: Build native library + run: | + pip install -e python/ + cactus build + mkdir -p python/cactus/bindings/lib + cp cactus-engine/build/libcactus_engine.so python/cactus/bindings/lib/ + + - name: Install Node.js 22 + run: | + curl -fsSL https://nodejs.org/dist/v22.21.0/node-v22.21.0-linux-arm64.tar.xz -o /tmp/node.tar.xz + tar -xJf /tmp/node.tar.xz -C /opt + echo "/opt/node-v22.21.0-linux-arm64/bin" >> "$GITHUB_PATH" + + - name: Bundle coding agent into the wheel + run: | + (cd cactus-code && npm install --no-audit --no-fund) + bash cactus-code/scripts/bundle-for-pypi.sh + rm -rf python/cactus/code/packages/tui/native + + - name: Build and repair wheel + working-directory: python + run: | + pip install build auditwheel + python -m build --wheel -C="--build-option=--plat-name=linux_aarch64" + auditwheel repair dist/*.whl -w wheelhouse + + - uses: actions/upload-artifact@v4 + with: + name: wheel-linux-aarch64 + path: python/wheelhouse/*.whl + + sdist: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - uses: ./.github/actions/stamp-version + + - name: Build sdist + working-directory: python + run: | + pip install build + python -m build --sdist + + - uses: actions/upload-artifact@v4 + with: + name: sdist + path: python/dist/*.tar.gz + + publish: + needs: [wheels-macos, wheels-linux, sdist] + runs-on: ubuntu-24.04 + timeout-minutes: 15 + permissions: + id-token: write + steps: + - uses: actions/download-artifact@v4 + with: + path: dist + merge-multiple: true + + - uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: dist + attestations: false diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 000000000..4b15e6de8 --- /dev/null +++ b/.github/workflows/python.yml @@ -0,0 +1,47 @@ +name: Python Tests + +on: + pull_request: + paths: + - "python/**" + - ".github/workflows/python.yml" + push: + branches: [main] + paths: + - "python/**" + - ".github/workflows/python.yml" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-24.04-arm + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + + - name: Install build tools + run: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev + + - name: Build libcactus + run: bash cactus-engine/build.sh + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: pip install -e "python/[dev,convert]" + + - name: Run tests + run: | + cd python + pytest tests/ -v \ + --ignore=tests/test_model.py \ + --ignore=tests/test_server_live.py diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 000000000..c79543413 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,117 @@ +name: Release + +on: + release: + types: [published] + +permissions: + contents: read + +concurrency: + group: release-${{ github.ref }} + cancel-in-progress: false + +jobs: + pypi: + permissions: + id-token: write + contents: read + uses: ./.github/workflows/pypi.yml + + homebrew: + needs: pypi + runs-on: ubuntu-24.04 + timeout-minutes: 15 + env: + PACKAGE: cactus-compute + TAP: cactus-compute/homebrew-cactus + TAG: ${{ github.event.release.tag_name }} + steps: + - name: Resolve version + id: ver + run: | + raw="${TAG#v}" + version=$(python3 -c "raw='${raw}'; print(raw + '.0' if raw.count('.') == 1 else raw)") + echo "version=$version" >> "$GITHUB_OUTPUT" + + - name: Resolve sdist URL and SHA256 from PyPI + id: pypi + env: + VERSION: ${{ steps.ver.outputs.version }} + run: | + for i in $(seq 1 10); do + if json=$(curl -fsSL "https://pypi.org/pypi/${PACKAGE}/${VERSION}/json"); then + break + fi + echo "Waiting for PyPI metadata (attempt $i)..." + sleep 30 + done + [ -n "$json" ] || { echo "PyPI metadata not available for ${PACKAGE} ${VERSION}"; exit 1; } + pick() { echo "$json" | jq -er '.urls[] | select(.packagetype=="sdist") | .'"$1"; } + { + echo "sdist_url=$(pick url)" + echo "sdist_sha=$(pick digests.sha256)" + } >> "$GITHUB_OUTPUT" + + - name: Update tap formula + env: + HOMEBREW_TAP_TOKEN: ${{ secrets.HOMEBREW_TAP_TOKEN }} + VERSION: ${{ steps.ver.outputs.version }} + SDIST_URL: ${{ steps.pypi.outputs.sdist_url }} + SDIST_SHA: ${{ steps.pypi.outputs.sdist_sha }} + run: | + git clone "https://x-access-token:${HOMEBREW_TAP_TOKEN}@github.com/${TAP}.git" tap + mkdir -p tap/Formula + cat > tap/Formula/cactus.rb <= 22 on PATH. + cactus convert (torch) is opt-in, same as pip's [convert] extra: + #{Formula["python@3.12"].opt_bin}/python3.12 -m pip --python=#{libexec}/bin/python install "${PACKAGE}[convert]==#{version}" + EOS + end + + test do + assert_equal "cactus #{version}\n", shell_output("#{bin}/cactus --version") + assert_match "libcactus_engine", + Dir[libexec/"lib/python*/site-packages/cactus/bindings/lib/*"].join + end + end + EOF + cd tap + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + git add Formula/cactus.rb + git commit -m "Update cactus to ${TAG}" + git push origin main + + docs: + permissions: + contents: write + uses: ./.github/workflows/docs.yml + with: + version: ${{ github.event.release.tag_name }} diff --git a/.gitignore b/.gitignore index 64b09d249..b2a93af05 100644 --- a/.gitignore +++ b/.gitignore @@ -1,27 +1,32 @@ -build +**/build/ +*.a +*.so +*.dylib +*.bin +*.cg + +!cactus-engine/libs/curl/**/*.a +!android/mbedtls/**/*.a + +apple/build-*/ +apple/*.xcframework + .claude .vscode .DS_Store -*.gguf -weights/ -stage/ -.cache/ -android/build/ -android/libs/ -android/arm64-v8a/ -*.so - -apple/build/ -apple/build-static/ -apple/*.xcframework -apple/libcactus.a -*.a -!libs/libcactus_pro.a +__pycache__/ +*.pyc +*.egg-info +.pytest_cache/ venv -__pycache__ -*.egg-info -profile.txt -*.dylib -*.bin \ No newline at end of file +weights/ +transpiled/ +python/dist/ +python/cactus/bin/ +python/cactus/bindings/lib/ +site/ +site_docs/ + +python/cactus/code/ diff --git a/CACTUS_VERSION b/CACTUS_VERSION index 3e1ad720b..0ac852dde 100644 --- a/CACTUS_VERSION +++ b/CACTUS_VERSION @@ -1 +1 @@ -1.5.0 \ No newline at end of file +v2.0.1 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 78e1b8b74..5174d4485 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,198 +1,54 @@ -# Contributing to Cactus - -### Code Style - -- **C++ Standard**: Use C++20 features where appropriate -- **Formatting**: Follow the existing code style in the project -- **ARM NEON**: When writing SIMD code, ensure proper alignment and use appropriate intrinsics -- **Comments**: Add comments for complex algorithms, especially in kernel implementations - -### Performance Considerations - -Cactus is a high-performance inference library optimized for ARM processors. When contributing: - -1. **Benchmark Your Changes**: Test performance impact, especially for kernel functions -2. **Memory Efficiency**: Minimize memory allocations in hot paths -3. **SIMD Optimization**: Use ARM NEON intrinsics where beneficial -4. **Cache Awareness**: Consider cache line sizes and memory access patterns - -### Testing - -Before submitting a PR: +--- +title: "Contributing to Cactus" +description: "Guidelines for contributing to Cactus, the on-device AI inference engine. Covers code style, PR process, testing, and benchmarking." +keywords: ["contributing", "open source", "pull request", "code style", "C++20"] +--- -1. Ensure all existing tests pass -2. Add tests for new functionality -3. Test on ARM hardware if possible (Apple Silicon, Raspberry Pi, etc.) -4. Verify quantized operations maintain acceptable accuracy +# Contributing to Cactus -### Pull Request Process +Thank you for your interest in contributing to Cactus! This document covers the guidelines and process for making contributions. -1. **Fork** the repository and create your branch from `main` -2. **Make your changes** following the guidelines above -3. **Sign-off all commits** using DCO -4. **Update documentation** if you change APIs -5. **Open a Pull Request** with a clear title and description +## Code Guidelines -#### PR Description Template +- **C++ Standard**: Use C++20 features where appropriate. +- **Formatting**: Follow the existing code style in the project, one header per folder. +- **Comments**: Avoid comments, make your code read like plain english. +- **AI-Generated Code**: Do not blindly PR AI slop, this codebase is very complex, they miss details. +- **Update docs**: Please update docs when necessary, be intuitive and straightforward. +- **Keep It Simple**: Do not go beyond the scope of the GH issue, avoid bloated PRs, keep codes lean. +- **Benchmark Your Changes**: Test performance impact, Cactus is performance-critical. +- **Test everything**: A PR that fails to build is the biggest red flag, means it was not tested. -```markdown -## Description -Brief description of what this PR does +## Pull Request Process -## Type of Change -- [ ] Bug fix -- [ ] New feature -- [ ] Performance improvement -- [ ] Documentation update +1. Fork the repository and create a branch from `main`. +2. Make your changes, keeping the scope focused on the relevant GitHub issue. +3. Run `cactus test` to verify your changes build and pass all tests. +4. Update documentation if your changes affect the public API or user-facing behavior. +5. Submit a pull request with a clear description of what you changed and why. ## Testing -- [ ] Tests pass locally -- [ ] Tested on ARM hardware -- [ ] Benchmarked performance impact - -## Checklist -- [ ] All commits are signed-off (DCO) -- [ ] Code follows project style -- [ ] Comments added where necessary -- [ ] Documentation updated if needed -``` - -### Reporting Issues - -When reporting issues, please include: - -1. System information (OS, CPU architecture, ARM variant) -2. Cactus version or commit hash -3. Minimal code to reproduce the issue -4. Expected vs actual behavior -5. Any relevant logs or error messages - -### Areas of Contribution - -We especially welcome contributions in these areas: - -- **Kernel Optimizations**: SIMD implementations for ARM architectures -- **Quantization**: Improved quantization techniques (INT8, INT4) -- **Model Support**: Support for additional model architectures -- **NPU Integration**: Apple Neural Engine and other NPU backends -- **Documentation**: Tutorials, examples, and API documentation -- **Testing**: Test coverage and benchmarking infrastructure - -### Development Setup - -```bash -# Clone the repository -git clone https://github.com/yourusername/cactus.git -cd cactus - -# Setup the environment (installs dependencies and activates venv) -./setup -``` - -You can run these codes directly on M-series Macbooks since they are ARM-based. -Vanilla M3 CPU-only can run Qwen3-600m-INT8 at 60-70 toks/sec. - -### Running Tests - -```bash -cactus test # Run unit tests and benchmarks -cactus test --model # Use a specific model for tests -cactus test --ios # Run tests on connected iPhone -cactus test --android # Run tests on connected Android device -``` - -### Building - -```bash -cactus build # Build for ARM chips (libcactus.a) -cactus build --apple # Build for Apple platforms (iOS/macOS) -cactus build --android # Build for Android -``` - -### Downloading Models ```bash -cactus download # Download and convert model weights -cactus run # Download, build, and run playground -``` +# Run all tests +cactus test -### Python FFI for Researchers +# Run tests on a connected iOS device +cactus test --ios -Cactus provides Python bindings for quick scripting and research. After setup: +# Run tests on a connected Android device +cactus test --android -```bash -cactus build -cactus download LiquidAI/LFM2-VL-450M -cactus download openai/whisper-small -cd python && python example.py +# Test a specific model +cactus test --model LiquidAI/LFM2-VL-450M ``` -**Available functions:** +## Developer Certificate of Origin -```python -from cactus import ( - cactus_init, # Load a model - cactus_complete, # Text/VLM completion - cactus_transcribe, # Audio transcription (Whisper) - cactus_embed, # Text embeddings - cactus_image_embed, # Image embeddings - cactus_audio_embed, # Audio embeddings - cactus_reset, # Reset model state - cactus_destroy # Free model memory -) -``` - -**Quick example:** - -```python -import json -from cactus import cactus_init, cactus_complete, cactus_destroy - -# Load model -model = cactus_init("../weights/lfm2-vl-450m") - -# Text completion -messages = json.dumps([{"role": "user", "content": "What is 2+2?"}]) -response = cactus_complete(model, messages) -print(json.loads(response)["response"]) - -# VLM - describe image -messages = json.dumps([{"role": "user", "content": "Describe this image", "images": ["path/to/image.png"]}]) -response = cactus_complete(model, messages) - -cactus_destroy(model) -``` +All contributions must comply with the [Developer Certificate of Origin (DCO)](DCO.md). By submitting a contribution, you certify that you have the right to do so under the project's open source license. -**Whisper transcription:** - -```python -from cactus import cactus_init, cactus_transcribe, cactus_destroy - -whisper = cactus_init("../weights/whisper-small") -prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" -response = cactus_transcribe(whisper, "audio.wav", prompt=prompt) - -cactus_destroy(whisper) -``` - -See `python/example.py` for a complete working example. - -### Questions? - -If you have questions about contributing, feel free to: - -1. Open an issue for discussion -2. Check existing issues and PRs -3. Review the codebase documentation - -## License - -By contributing to Cactus, you agree that your contributions will be licensed under the same license as the project (check LICENSE file). - -## Enforcement - -Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project maintainers. All complaints will be reviewed and investigated promptly and fairly. - ---- +## See Also -Thank you for contributing to Cactus! \ No newline at end of file +- [Cactus Engine API](/docs/cactus_engine.md) — C API reference +- [Cactus Graph API](/docs/cactus_graph.md) — Computational graph API reference +- [Cactus Index API](/docs/cactus_index.md) — Vector database API reference diff --git a/README.md b/README.md index 95d00c444..00d2c928e 100644 --- a/README.md +++ b/README.md @@ -1,63 +1,52 @@ +# Cactus + Logo -Cross-platform & energy-efficient kernels, runtime and AI inference engine for mobile devices. +[![Docs][docs-shield]][docs-url] +[![Website][website-shield]][website-url] +[![GitHub][github-shield]][github-url] +[![HuggingFace][hf-shield]][hf-url] +[![Reddit][reddit-shield]][reddit-url] +[![Blog][blog-shield]][blog-url] + +A hybrid edge-cloud AI engine for mobile devices & wearables. ``` ┌─────────────────┐ -│ Cactus FFI │ ←── OpenAI compatible C API for integration (tools, RAG, cloud handoff) -└─────────────────┘ +│ Cactus Engine │ ←── OpenAI-compatible APIs for text, speech, and vision. +└─────────────────┘ │ ┌─────────────────┐ -│ Cactus Engine │ ←── High-level transformer engine (NPU support, INT4/INT8/FP16/MIXED) -└─────────────────┘ +│ Cactus Graph │ ←── Zero-copy computation graph +└─────────────────┘ │ ┌─────────────────┐ -│ Cactus Models │ ←── Implements SOTA models using Cactus Graphs -└─────────────────┘ +│ Cactus Kernels │ ←── CPU/GPU kernels for (Apple, Samsung, Pixel, etc.) +└─────────────────┘ │ -┌─────────────────┐ -│ Cactus Graph │ ←── Unified zero-copy computation graph (think NumPy for mobile) -└─────────────────┘ +┌─────────────────┐ +│ Cactus Quants │ ←── Custom rotation-based quantization technique +└─────────────────┘ │ ┌─────────────────┐ -│ Cactus Kernels │ ←── Low-level ARM-specific SIMD operations (think CUDA for mobile) +│Cactus Transpiler│ ←── Transpiles custom PyTorch model to Cactus. └─────────────────┘ ``` -# Cactus Graph & Kernel -```cpp -#include cactus.h - -CactusGraph graph; -auto a = graph.input({2, 3}, Precision::FP16); -auto b = graph.input({3, 4}, Precision::INT8); - -auto x1 = graph.matmul(a, b, false); -auto x2 = graph.transpose(x1); -auto result = graph.matmul(b, x2, true); - -float a_data[6] = {1.1f, 2.3f, 3.4f, 4.2f, 5.7f, 6.8f}; -float b_data[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; - -graph.set_input(a, a_data, Precision::FP16); -graph.set_input(b, b_data, Precision::INT8); - -graph.execute(); -void* output_data = graph.get_output(result); +## Quick Demo (Mac) -graph.hard_reset(); +- Step 1: `brew install cactus-compute/cactus/cactus` +- Step 2: `cactus run` -``` +## Cactus Engine -# Cactus Engine & FFI ```cpp -#include cactus.h - -cactus_set_pro_key(""); // email founders@cactuscompute.com for optional key +#include "cactus_engine.h" cactus_model_t model = cactus_init( - "path/to/weight/folder", // section to generate weigths below - "txt/or/md/file/or/dir/with/many", // nullptr if none, cactus does automatic fast RAG + "path/to/weight/folder", + "path to txt or dir of txts for auto-rag", + false ); const char* messages = R"([ @@ -72,116 +61,283 @@ const char* options = R"({ char response[4096]; int result = cactus_complete( - model, // model handle from cactus_init - messages, // JSON array of chat messages - response, // buffer to store response JSON - sizeof(response), // size of response buffer - options, // optional: generation options (nullptr for defaults) - nullptr, // optional: tools JSON for function calling - nullptr, // optional: streaming callback fn(token, id, user_data) - nullptr // optional: user data passed to callback + model, // model handle + messages, // JSON chat messages + response, // response buffer + sizeof(response), // buffer size + options, // generation options + nullptr, // tools JSON + nullptr, // streaming callback + nullptr, // user data + nullptr, // pcm audio buffer + 0 // pcm buffer size ); ``` -Example response from Gemma3-270m +Example response from Gemma4-E2B ```json { - "success": true, // when successfully generated locally - "error": null, // returns specific errors if success = false - "cloud_handoff": false, // true when model is unconfident, simply route to cloud - "response": "Hi there!", // null when error is not null or cloud_handoff = true - "function_calls": [], // parsed to [{"name":"set_alarm","arguments":{"hour":"10","minute":"0"}}] - "confidence": 0.8193, // how confident the model is with its response - "time_to_first_token_ms": 45.23, // latency (time to first token) - "total_time_ms": 163.67, // total execution time - "prefill_tps": 1621.89, // prefill tokens per second - "decode_tps": 168.42, // decode tokens per second - "ram_usage_mb": 245.67, // current process RAM usage in MB + "success": true, // generation succeeded + "error": null, // error details if failed + "cloud_handoff": false, // true if cloud model used + "response": "Hi there!", + "function_calls": [], // parsed tool calls + "segments": [], // transcription segments (empty for chat) + "confidence": 0.8193, // model confidence + "confidence_threshold": 0.7, // resolved handoff threshold (model-dependent) + "time_to_first_token_ms": 45.23, + "total_time_ms": 163.67, + "prefill_tps": 1621.89, + "decode_tps": 168.42, + "ram_usage_mb": 245.67, "prefill_tokens": 28, "decode_tokens": 50, "total_tokens": 78 } ``` -# Performance - -- **Models:** LFM2-VL-450m & Whisper-Small -- **Precision:** Cactus smartly blends INT4, INT8 and F16 for all weights. -- **Decode** = toks/sec, **P/D** = prefill/decode, **VLM** = 256×256 image, **STT** = 30s audio -- **Cactus Pro**: Uses NPU for realtime and large context (Apple for now), scores are marked with * - -| Device | Short Decode | 4k-P/D | VLM-TTFT | VLM-Dec | STT-TTFT | STT-Dec | -|--------|--------|--------|----------|---------|----------|---------| -| Mac M4 Pro | 170 | 989/150 | 0.2s/0.1s* | 168 | 1.0s/0.2s* | 92 | -| Mac M3 Pro | 140 | 890/123 | 0.3s/0.1s* | 149 | 1.5s/0.4s* | 81 | -| iPad/Mac M4 | 134 | 603/106 | 0.3s/0.1s* | 129 | 1.8s0.3s* | 70 | -| iPad/Mac M3 | 117 | 525/93 | 0.4s/0.1s* | 111 | 2.8s/0.7s* | 61 | -| iPhone 17 Pro | 126 | 428/84 | 0.5s/0.1s* | 120 | 3.0s/0.6s* | 80 | -| iPhone 16 Pro | 106 | 380/81 | 0.6s/0.2s* | 101 | 4.3s/0.7s* | 75 | -| iPhone 15 Pro | 90 | 330/75 | 0.7s/0.3s* | 92 | 4.5s/0.8s* | 70 | -| Galaxy S25 Ultra | 80 | 355/52 | 0.7s | 70 | 3.6s/- | 32 | -| Nothing 3 | 56 | 320/46 | 0.8s | 54 | 4.5s | 55 | -| Pixel 6a | 25 | 108/24 | 2.3s | 25 | 9.6 | 15 | -| Raspberry Pi 5 | 20 | 292/18 | 1.7s | 23 | 15s | 16 | - - -# Supported models - -- Cactus smartly and compactly blends INT4, INT8 and F16 for all weights. -- You could still quantize everything with one precision, but mixed is optimal - -| Model | Zipped Size | Completion | Tools | Vision | Embed | Speech | Pro | -|-------|------------------|------------|-------|--------|-------|--------|-----| -| google/gemma-3-270m-it | 252MB | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ | -| google/functiongemma-270m-it | 252MB | ✓ | ✓ | ✗ | ✗ | ✗ | ✗ | -| openai/whisper-small | 283MB | ✗ | ✗ | ✗ | ✓ | ✓ | Apple | -| LiquidAI/LFM2-350M | 244MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | -| LiquidAI/LFM2-VL-450M | 448MB | ✓ | ✗ | ✓ | ✓ | ✗ | Apple | -| nomic-ai/nomic-embed-text-v2-moe | 451MB | ✗ | ✗ | ✗ | ✓ | ✗ | ✗ | -| Qwen/Qwen3-0.6B | 514MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | -| Qwen/Qwen3-Embedding-0.6B | 514MB | ✗ | ✗ | ✗ | ✓ | ✗ | ✗ | -| LiquidAI/LFM2-700M | 498MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | -| google/gemma-3-1b-it | 642MB | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ | -| LiquidAI/LFM2.5-1.2B-Instruct | 474MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | -| LiquidAI/LFM2-1.2B-RAG | 474MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | -| LiquidAI/LFM2-1.2B-Tool | 474MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | -| openai/whisper-medium | 658MB | ✗ | ✗ | ✗ | ✓ | ✓ | Apple | -| LiquidAI/LFM2.5-VL-1.6B | 954MB | ✓ | ✗ | ✓ | ✓ | ✗ | Apple | -| Qwen/Qwen3-1.7B | 749MB | ✓ | ✓ | ✗ | ✓ | ✗ | ✗ | - -# Using this repo on a Mac - -```bash -git clone https://github.com/cactus-compute/cactus && cd cactus && source ./setup +## Cactus Graph + +```cpp +#include "cactus_graph.h" + +CactusGraph graph; +auto a = graph.input({2, 3}, Precision::FP16); +auto b = graph.input({3, 4}, Precision::INT8); + +auto x1 = graph.matmul(a, b, false); +auto x2 = graph.transpose(x1); +auto result = graph.matmul(b, x2, true); + +float a_data[6] = {1.1f, 2.3f, 3.4f, 4.2f, 5.7f, 6.8f}; +float b_data[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + +graph.set_input(a, a_data, Precision::FP16); +graph.set_input(b, b_data, Precision::INT8); + +graph.execute(); +void* output_data = graph.get_output(result); + +graph.hard_reset(); +``` + +## Inference Speed + +- LLM: Gemma-4-E2B-CQ4 (1k-context prefill / decode for 100 tokens) +- VLM: Gemma-4-E2B-CQ4 (256px image encode time / decode) +- Transcribe: Parakeet-TDT-0.6B-CQ4 (20s audio end-to-end transcribe time) +- 1k-Context RAM: peak MB during the LLM benchmark +- No speculative decode or MTP, pure decode + +Command: `cactus benchmark` [optional `--ios` or `--android`] + +| Device | LLM | VLM | Transcribe | RAM | +|--------|-----|-----|------------|---------------| +| Mac M5 Max | 2964tps / 154tps | 0.09s / 168tps | 0.15s | 1348MB | +| Mac M4 Pro | 1947tps / 97tps | 0.25s / 108tps | 0.28s | 1225MB | +| Mac M3 Pro | 1294tps / 64tps | 0.40s / 72tps | 0.37s | 735MB | +| iPad/Vision Pro M5 | 1336tps / 71tps | 0.25s / 80tps | 0.27s | 703MB | +| iPhone 17 Pro | 729tps / 37tps | 0.5s / 39tps | 0.51s | 644MB | +| iPhone 15 Pro | 511tps / 25tps | 1.16s / 27tps | 1.40s | 635MB | + +N/B: With 1k-context prefill and decode for 100 runs on M5 Max +- `LFM2.5-VL-1.6B` = 289toks/sec +- `Qwen3-1.7B` = 155toks/sec +- `LFM2.5-VL-450m` = 472toks/sec, image encodes in 43ms +- `LFM22.5-VL-230m` = 555toks/sec + +## Output Quality + +- Gemma-4-E2B-it accuracy across bit widths, averaged over 3 seeds. +- CQ3.26 and CQ2.54 are mixed-precision, CQ2/CQ3/CQ4 are uniformly quantized. +- Full results in [docs/cactus_quants.md](/docs/cactus_quants.md): + +| Task | F16 (Original) | CQ4 | CQ3.26 | CQ2.54 | CQ2 | +|-----------|-----|-----|--------|--------|-----| +| ARC-E | 73.80 | 73.73 | 74.20 | 68.20 | 50.80 | +| ARC-C | 56.47 | 52.47 | 51.53 | 37.20 | 24.73 | +| HellaSwag | 46.93 | 47.07 | 45.20 | 40.73 | 35.87 | +| WinoGrande | 61.00 | 61.13 | 59.60 | 60.13 | 51.27 | +| MMLU | 62.33 | 59.45 | 57.63 | 47.19 | 33.18 | +| GPQA | 34.34 | 34.34 | 31.82 | 30.81 | 23.23 | +| GSM8K | 73.67 | 71.20 | 66.20 | 22.00 | 0.40 | +| HumanEval | 54.88 | 57.11 | 53.66 | 15.24 | 1.02 | +| BFCL Simple | 92.00 | 92.42 | 91.50 | 82.25 | 18.75 | +| BFCL Multi | 89.00 | 88.33 | 89.00 | 52.50 | 13.67 | +| BFCL Parallel | 84.00 | 83.67 | 82.50 | 30.00 | 3.33 | +| BFCL Parallel-Multi | 78.00 | 83.33 | 82.00 | 37.00 | 1.33 | + +## Supported Models + +- Any HuggingFace model can be converted using `cactus convert [HF-Name]`, though experimental. +- Liquid, Gemma. whisper. parakeet and Qwen model families are especially tested. +- Some models have been pre-uploaded [here](https://huggingface.co/Cactus-Compute), just run `cactus download [HF-Name]`. +- `cactus run [HF-Name]` albeit first downloads or convert the model if not found. + +## Learn More + +| Reference | Language | Description | +|-----------|----------|-------------| +| [Cactus Engine](/docs/cactus_engine.md) | C | Chat completion, streaming, tool calling, transcription, embeddings, RAG, vision, vector index, cloud handoff | +| [Cactus Graph](/docs/cactus_graph.md) | C++ | Tensor operations, matrix multiplication, attention, normalization, activation functions | +| [Cactus Kernels](/docs/cactus_kernels.md) | C++ | ARM NEON SIMD kernels for matmul, attention, convolution, quantization, DSP, image processing | +| [Cactus Quants](/docs/cactus_quants.md) | C++ | Rotation-and-codebook quantization from 4-bit to 1-bit for all weight tensors | +| [Cactus Hybrid](/docs/cactus_hybrid.md) | C/Python | Route hard queries to the cloud automatically based on local model confidence | +| [Cactus Transpiler](/docs/cactus_transpiler.md) | Python | Convert any PyTorch model to a Cactus runtime graph for on-device inference | +| [Python Package](/python/) | Python | Python package and CLI | + +## Bindings + +- [Swift](/bindings/swift/) +- [Kotlin](/bindings/kotlin/) +- [Flutter](/bindings/flutter/) +- [React Native](/bindings/react-native/) +- [Python](/bindings/python/) +- [Rust](/bindings/rust/) + +## Using this repo + +``` +┌────────────────────────────────────────────────────────────────────────────────┐ +│ │ +│ Step 0: if on Linux (Ubuntu/Debian) │ +│ sudo apt-get install python3.12 python3.12-venv python3-pip cmake │ +│ build-essential libcurl4-openssl-dev │ +│ │ +│ Step 1: clone and setup │ +│ git clone https://github.com/cactus-compute/cactus && cd cactus │ +│ source ./setup │ +│ │ +│ Step 2: use the commands │ +│────────────────────────────────────────────────────────────────────────────────│ +│ │ +│ cactus auth manage cloud API key │ +│ --status show key status │ +│ --clear remove saved key │ +│ │ +│ cactus run [model|path] run a model (downloads if needed) │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --backend cpu|metal inference backend (default: auto) │ +│ --image image file for VLM inference │ +│ --audio audio file for audio chat │ +│ --system system prompt │ +│ --prompt send prompt immediately │ +│ --thinking enable thinking/reasoning mode │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild from source │ +│ │ +│ cactus transcribe [model] live microphone transcription with a model│ +│ --file audio file to transcribe (WAV) │ +│ --language language code (default: en) │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild from source │ +│ │ +│ cactus download [model] get a bundle (prebuilt, else build) │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild from source │ +│ │ +│ cactus convert [dir] HuggingFace -> runnable cactus bundle │ +│ (CQ weights + runtime graph) │ +│ --bits 1|2|3|4 CQ quantization (default: 4) │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild from source │ +│ --lora merge a LoRA adapter before converting │ +│ --weights-only stop after CQ weights (skip the graph) │ +│ --artifact-dir bundle output (default: weights/) │ +│ │ +│ cactus serve [model] OpenAI-compatible local HTTP server │ +│ --host bind address (default: 127.0.0.1) │ +│ --port port (default: 8080) │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --backend cpu|metal inference backend (default: auto) │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild from source │ +│ --no-cloud-handoff disable automatic cloud handoff │ +│ --confidence-threshold <0..1> handoff to cloud below this confidence │ +│ --cloud-timeout-ms max wait for cloud handoff │ +│ │ +│ cactus code run the AI coding agent (TUI / print) │ +│ --serve-model auto-start a server with this model │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --backend cpu|metal inference backend (default: auto) │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild from source │ +│ --host server address (default: 127.0.0.1) │ +│ --port server port (default: 8080) │ +│ --no-serve require a running server (no auto-start) │ +│ --no-cloud-handoff disable automatic cloud handoff │ +│ --confidence-threshold <0..1> handoff to cloud below this confidence │ +│ --cloud-timeout-ms max wait for cloud handoff │ +│ -- pass remaining args to the agent │ +│ │ +│ cactus list list downloaded models │ +│ │ +│ cactus build build cactus libraries │ +│ --apple Apple (iOS/macOS) │ +│ --android Android │ +│ --python shared lib for Python FFI │ +│ │ +│ cactus test run the test suite │ +│ --component kernels | graph | engine | all │ +│ (default: all) │ +│ --model default: google/gemma-4-E2B-it │ +│ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --backend cpu|metal inference backend (default: auto) │ +│ --token HuggingFace token (gated models) │ +│ --reconvert force local rebuild of test models │ +│ --suite run a single test suite by name │ +│ (resolved across components, │ +│ e.g. llm → engine) │ +│ --list list components and suites │ +│ --ios run on connected iPhone │ +│ --android run on connected Android │ +│ --enable-telemetry send cloud telemetry (off by default) │ +│ │ +│ cactus benchmark run the engine benchmark suite │ +│ --model default: google/gemma-4-E2B-it │ +│ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ +│ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ +│ --backend cpu|metal inference backend (default: auto) │ +│ --ios run on connected iPhone │ +│ --android run on connected Android │ +│ │ +│ cactus clean delete build artifacts, weights, venv │ +│ cactus --help show this help │ +│ │ +└────────────────────────────────────────────────────────────────────────────────┘ +``` + +## Citation + +If you use Cactus in your research, please cite it as follows: + +```bibtex +@software{cactus, + title = {Cactus: AI Inference Engine for Phones & Wearables}, + author = {Ndubuaku, Henry and Cactus Team}, + url = {https://github.com/cactus-compute/cactus}, + year = {2025} +} ``` -- `[model]` is a HuggingFace name from the table above (default: `google/gemma-3-270m-it`) -- Common flags: `--precision INT4|INT8|FP16` (default: INT4), `--token ` -- Always run `source ./setup` in any new terminal. - -| Command | Description | -|---------|-------------| -| `cactus run [model]` | Opens playground (auto downloads model) | -| `cactus download [model]` | Downloads model to `./weights` | -| `cactus convert [model] [dir]` | Converts model, supports LoRA merging (`--lora `) | -| `cactus build` | Builds for ARM (`--apple` or `--android`) | -| `cactus test` | Runs tests (`--ios` / `--android`, `--model [name/path]`) | -| `cactus clean` | Removes build artifacts | -| `cactus --help` | Shows all commands and flags | - -# Using in your apps - -- [Python for Mac](/python/) -- [React Native SDK](https://github.com/cactus-compute/cactus-react-native) -- [Swift Multiplatform SDK](https://github.com/mhayes853/swift-cactus) -- [Kotlin Multiplatform SDK](https://github.com/cactus-compute/cactus-kotlin) -- [Flutter SDK](https://github.com/cactus-compute/cactus-flutter) -- [Rust SDK](https://github.com/mrsarac/cactus-rs) - -# Try demo apps - -- [iOS Demo](https://apps.apple.com/gb/app/cactus-chat/id6744444212) -- [Android Demo](https://play.google.com/store/apps/details?id=com.rshemetsubuser.myapp) - -# Maintaining Organisations -1. [Cactus Compute, Inc](https://cactuscompute.com/) -2. [UCLA's BruinAI](https://bruinai.org/) \ No newline at end of file +**N/B:** Scroll all the way up and click the shields link for resources! + +[docs-shield]: https://img.shields.io/badge/Docs-555?style=for-the-badge&logo=readthedocs&logoColor=white +[docs-url]: https://cactus-compute.github.io/cactus/ + +[website-shield]: https://img.shields.io/badge/Website-555?style=for-the-badge&logo=safari&logoColor=white +[website-url]: https://cactuscompute.com/ + +[github-shield]: https://img.shields.io/badge/GitHub-555?style=for-the-badge&logo=github&logoColor=white +[github-url]: https://github.com/cactus-compute/cactus + +[hf-shield]: https://img.shields.io/badge/HuggingFace-555?style=for-the-badge&logo=huggingface&logoColor=white +[hf-url]: https://huggingface.co/Cactus-Compute + +[reddit-shield]: https://img.shields.io/badge/Reddit-555?style=for-the-badge&logo=reddit&logoColor=white +[reddit-url]: https://www.reddit.com/r/cactuscompute/ + +[blog-shield]: https://img.shields.io/badge/Blog-555?style=for-the-badge&logo=hashnode&logoColor=white +[blog-url]: https://cactuscompute.com/blog diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 000000000..bb3ce4b31 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,28 @@ +--- +title: "Security Policy" +description: "How to report security vulnerabilities in Cactus, the on-device AI inference engine." +keywords: ["security", "vulnerability", "responsible disclosure"] +--- + +# Security Policy + +## Reporting a Vulnerability + +If you discover a security vulnerability in Cactus, please report it responsibly. **Do not open a public GitHub issue for security vulnerabilities.** + +Instead, please email **security@cactuscompute.com** with: + +1. A description of the vulnerability. +2. Steps to reproduce the issue. +3. The potential impact. +4. Any suggested fixes (optional). + +We will acknowledge receipt within 48 hours and aim to provide a fix or mitigation plan within 7 days. + +## Scope + +This policy covers the Cactus inference engine, its bindings (Python, Swift, Kotlin, Flutter, Rust, React Native), and the official model weights hosted on [HuggingFace](https://huggingface.co/Cactus-Compute). + +## Supported Versions + +Security fixes are applied to the latest release. We recommend always running the latest version of Cactus. diff --git a/android/CMakeLists.txt b/android/CMakeLists.txt index d2260ddca..c62281b1c 100644 --- a/android/CMakeLists.txt +++ b/android/CMakeLists.txt @@ -1,141 +1,43 @@ cmake_minimum_required(VERSION 3.10) -project(Cactus LANGUAGES CXX) +project(CactusAndroid LANGUAGES CXX) -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED True) - -set(SOURCE_DIR ${CMAKE_SOURCE_DIR}/../cactus) +set(CMAKE_BUILD_TYPE Release CACHE STRING "") +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib) -if(NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE Release) - message(STATUS "CMAKE_BUILD_TYPE was not set, defaulting to Release.") +if(NOT ${ANDROID_ABI} STREQUAL "arm64-v8a") + message(FATAL_ERROR "Unsupported Android ABI: ${ANDROID_ABI}. Cactus only supports arm64-v8a.") endif() -file(GLOB ENGINE_SOURCES "${SOURCE_DIR}/engine/*.cpp") -file(GLOB GRAPH_SOURCES "${SOURCE_DIR}/graph/*.cpp") -file(GLOB KERNEL_SOURCES "${SOURCE_DIR}/kernel/*.cpp") -file(GLOB MODEL_SOURCES "${SOURCE_DIR}/models/*.cpp") -file(GLOB FFI_SOURCES "${SOURCE_DIR}/ffi/*.cpp") -file(GLOB NPU_SOURCES "${SOURCE_DIR}/npu/*.cpp") - -set(HEADER_FILES - ${SOURCE_DIR}/cactus.h - ${SOURCE_DIR}/engine/engine.h - ${SOURCE_DIR}/graph/graph.h - ${SOURCE_DIR}/kernel/kernel.h - ${SOURCE_DIR}/kernel/kernel_utils.h - ${SOURCE_DIR}/ffi/cactus_ffi.h - ${SOURCE_DIR}/ffi/cactus_utils.h - ${SOURCE_DIR}/models/model.h - ${SOURCE_DIR}/npu/npu.h -) - -set(COMMON_SOURCES - ${KERNEL_SOURCES} - ${GRAPH_SOURCES} - ${ENGINE_SOURCES} - ${FFI_SOURCES} - ${MODEL_SOURCES} - ${NPU_SOURCES} -) - -if(NOT COMMON_SOURCES) - message(FATAL_ERROR "No source files found! Check SOURCE_DIR: ${SOURCE_DIR}") +if(NOT TARGET cactus_engine) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../cactus-engine ${CMAKE_CURRENT_BINARY_DIR}/cactus-engine) endif() -find_library(LOG_LIB log) - -add_library( - cactus - SHARED - ${COMMON_SOURCES} - ${CMAKE_SOURCE_DIR}/cactus_jni.cpp -) - -add_library( - cactus_static - STATIC - ${COMMON_SOURCES} -) - -set(COMMON_INCLUDE_DIRS - ${SOURCE_DIR} - ${SOURCE_DIR}/engine - ${SOURCE_DIR}/graph - ${SOURCE_DIR}/kernel - ${SOURCE_DIR}/ffi - ${SOURCE_DIR}/models - ${SOURCE_DIR}/npu -) - -target_include_directories(cactus PRIVATE ${COMMON_INCLUDE_DIRS}) -target_include_directories(cactus_static PRIVATE ${COMMON_INCLUDE_DIRS}) +set(_DEPS cactus_engine cactus_graph cactus_kernels) -target_link_libraries(cactus ${LOG_LIB} android) +add_library(cactus_jni SHARED ${CMAKE_CURRENT_SOURCE_DIR}/cactus_jni.cpp) -set(COMMON_COMPILE_DEFINITIONS - PLATFORM_CPU_ONLY=1 +set_target_properties(cactus_jni PROPERTIES + OUTPUT_NAME cactus_engine + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON ) -target_compile_definitions(cactus PRIVATE ${COMMON_COMPILE_DEFINITIONS}) -target_compile_definitions(cactus_static PRIVATE ${COMMON_COMPILE_DEFINITIONS}) - -set(COMMON_COMPILE_OPTIONS - -pthread - -Wall - -Wextra - -pedantic - -Wno-c++20-designator - -Wno-missing-field-initializers - -fomit-frame-pointer - -funroll-loops - -ftree-vectorize +target_include_directories(cactus_jni PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../cactus-engine + ${CMAKE_CURRENT_SOURCE_DIR}/../cactus-engine/src ) -target_compile_options(cactus PRIVATE ${COMMON_COMPILE_OPTIONS}) -target_compile_options(cactus_static PRIVATE ${COMMON_COMPILE_OPTIONS}) - -if (${CMAKE_BUILD_TYPE} STREQUAL "Debug") - target_compile_options(cactus PRIVATE -DCACTUS_ANDROID_ENABLE_LOGGING) - target_compile_options(cactus_static PRIVATE -DCACTUS_ANDROID_ENABLE_LOGGING) -endif () +target_link_libraries(cactus_jni PRIVATE ${_DEPS}) -set(OPTIMIZATION_OPTIONS - -O3 - -DNDEBUG - -ffunction-sections - -fdata-sections -) - -target_compile_options(cactus PRIVATE ${OPTIMIZATION_OPTIONS}) -target_compile_options(cactus_static PRIVATE ${OPTIMIZATION_OPTIONS}) - -target_compile_options(cactus PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden) - -if (${ANDROID_ABI} STREQUAL "arm64-v8a") - set(ARM_OPTIONS - -march=armv8.2-a+fp16+simd+dotprod - ) - set(ARM_DEFINITIONS - __ARM_NEON=1 - __ARM_FEATURE_FP16_VECTOR_ARITHMETIC=1 - __ARM_FEATURE_DOTPROD=1 - ) - - target_compile_options(cactus PRIVATE ${ARM_OPTIONS}) - target_compile_options(cactus_static PRIVATE ${ARM_OPTIONS}) - - target_compile_definitions(cactus PRIVATE ${ARM_DEFINITIONS}) - target_compile_definitions(cactus_static PRIVATE ${ARM_DEFINITIONS}) -else() - message(FATAL_ERROR "Unsupported Android ABI: ${ANDROID_ABI}. Cactus only supports arm64-v8a due to ARM NEON requirements.") +set(_CACTUS_VERSION_FILE "${CMAKE_CURRENT_SOURCE_DIR}/../CACTUS_VERSION") +if(EXISTS "${_CACTUS_VERSION_FILE}") + file(READ "${_CACTUS_VERSION_FILE}" _CACTUS_VERSION_CONTENT) + string(STRIP "${_CACTUS_VERSION_CONTENT}" CACTUS_VERSION_STR) + target_compile_definitions(cactus_jni PRIVATE CACTUS_COMPILE_TIME_VERSION="${CACTUS_VERSION_STR}") endif() -target_link_options(cactus PRIVATE -Wl,--gc-sections) -target_link_options(cactus PRIVATE -Wl,--exclude-libs,ALL) -target_link_options(cactus PRIVATE -flto) -target_link_options(cactus PRIVATE -Wl,-z,max-page-size=16384) +if(CMAKE_BUILD_TYPE STREQUAL "Debug") + target_compile_definitions(cactus_jni PRIVATE CACTUS_ANDROID_ENABLE_LOGGING) +endif() -set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib) -set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib) -set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) +target_link_options(cactus_jni PRIVATE -Wl,-z,max-page-size=16384) diff --git a/android/Cactus.android.kt b/android/Cactus.android.kt deleted file mode 100644 index 3a6e302ca..000000000 --- a/android/Cactus.android.kt +++ /dev/null @@ -1,387 +0,0 @@ -package com.cactus - -import org.json.JSONArray -import org.json.JSONObject - -actual class Cactus private constructor(private var handle: Long) : AutoCloseable { - - actual companion object { - init { - System.loadLibrary("cactus") - } - - actual fun create(modelPath: String, corpusDir: String?): Cactus { - val handle = nativeInit(modelPath, corpusDir) - if (handle == 0L) { - throw CactusException(nativeGetLastError().ifEmpty { "Failed to initialize model" }) - } - return Cactus(handle) - } - - actual fun setTelemetryToken(token: String) = nativeSetTelemetryToken(token) - actual fun setProKey(key: String) = nativeSetProKey(key) - - @JvmStatic - private external fun nativeInit(modelPath: String, corpusDir: String?): Long - @JvmStatic - private external fun nativeGetLastError(): String - @JvmStatic - private external fun nativeSetTelemetryToken(token: String) - @JvmStatic - private external fun nativeSetProKey(key: String) - } - - actual fun complete(prompt: String, options: CompletionOptions): CompletionResult { - return complete(listOf(Message.user(prompt)), options, null, null) - } - - actual fun complete( - messages: List, - options: CompletionOptions, - tools: List>?, - callback: TokenCallback? - ): CompletionResult { - checkHandle() - val messagesJson = JSONArray(messages.map { it.toJson() }).toString() - val optionsJson = options.toJson() - val toolsJson = tools?.let { JSONArray(it.map { t -> JSONObject(t) }).toString() } - - val responseJson = nativeComplete(handle, messagesJson, optionsJson, toolsJson, callback) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return json.toCompletionResult() - } - - actual fun transcribe( - audioPath: String, - prompt: String?, - language: String?, - translate: Boolean - ): TranscriptionResult { - checkHandle() - val optionsJson = JSONObject().apply { - language?.let { put("language", it) } - put("translate", translate) - }.toString() - - val responseJson = nativeTranscribe(handle, audioPath, prompt, optionsJson, null) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return json.toTranscriptionResult() - } - - actual fun transcribe( - pcmData: ByteArray, - prompt: String?, - language: String?, - translate: Boolean - ): TranscriptionResult { - checkHandle() - val optionsJson = JSONObject().apply { - language?.let { put("language", it) } - put("translate", translate) - }.toString() - - val responseJson = nativeTranscribe(handle, null, prompt, optionsJson, pcmData) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return json.toTranscriptionResult() - } - - actual fun embed(text: String, normalize: Boolean): FloatArray { - checkHandle() - return nativeEmbed(handle, text, normalize) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to generate embedding" }) - } - - actual fun ragQuery(query: String, topK: Int): String { - checkHandle() - val responseJson = nativeRagQuery(handle, query, topK) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return responseJson - } - - actual fun tokenize(text: String): IntArray { - checkHandle() - return nativeTokenize(handle, text) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to tokenize" }) - } - - actual fun scoreWindow(tokens: IntArray, start: Int, end: Int, context: Int): String { - checkHandle() - val responseJson = nativeScoreWindow(handle, tokens, start, end, context) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - return responseJson - } - - actual fun imageEmbed(imagePath: String): FloatArray { - checkHandle() - return nativeImageEmbed(handle, imagePath) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to generate image embedding" }) - } - - actual fun audioEmbed(audioPath: String): FloatArray { - checkHandle() - return nativeAudioEmbed(handle, audioPath) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to generate audio embedding" }) - } - - actual fun createStreamTranscriber(): StreamTranscriber { - checkHandle() - val streamHandle = nativeStreamTranscribeInit(handle) - if (streamHandle == 0L) { - throw CactusException(nativeGetLastError().ifEmpty { "Failed to create stream transcriber" }) - } - return StreamTranscriber(streamHandle) - } - - actual fun reset() { - checkHandle() - nativeReset(handle) - } - - actual fun stop() { - checkHandle() - nativeStop(handle) - } - - actual override fun close() { - if (handle != 0L) { - nativeDestroy(handle) - handle = 0L - } - } - - private fun checkHandle() { - if (handle == 0L) { - throw CactusException("Model has been closed") - } - } - - private external fun nativeDestroy(handle: Long) - private external fun nativeReset(handle: Long) - private external fun nativeStop(handle: Long) - private external fun nativeComplete(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, callback: TokenCallback?): String - private external fun nativeTranscribe(handle: Long, audioPath: String?, prompt: String?, optionsJson: String?, pcmData: ByteArray?): String - private external fun nativeEmbed(handle: Long, text: String, normalize: Boolean): FloatArray? - private external fun nativeRagQuery(handle: Long, query: String, topK: Int): String - private external fun nativeTokenize(handle: Long, text: String): IntArray? - private external fun nativeScoreWindow(handle: Long, tokens: IntArray, start: Int, end: Int, context: Int): String - private external fun nativeImageEmbed(handle: Long, imagePath: String): FloatArray? - private external fun nativeAudioEmbed(handle: Long, audioPath: String): FloatArray? - private external fun nativeStreamTranscribeInit(handle: Long): Long -} - -actual class StreamTranscriber internal constructor(private var handle: Long) : AutoCloseable { - - actual fun insert(pcmData: ByteArray) { - checkHandle() - val result = nativeStreamTranscribeInsert(handle, pcmData) - if (result < 0) { - throw CactusException("Failed to insert audio data") - } - } - - actual fun process(language: String?): TranscriptionResult { - checkHandle() - val optionsJson = language?.let { JSONObject().put("language", it).toString() } - val responseJson = nativeStreamTranscribeProcess(handle, optionsJson) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - return json.toTranscriptionResult() - } - - actual fun finalize(): TranscriptionResult { - checkHandle() - val responseJson = nativeStreamTranscribeFinalize(handle) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - return json.toTranscriptionResult() - } - - actual override fun close() { - if (handle != 0L) { - nativeStreamTranscribeDestroy(handle) - handle = 0L - } - } - - private fun checkHandle() { - if (handle == 0L) { - throw CactusException("Stream transcriber has been closed") - } - } - - private external fun nativeStreamTranscribeInsert(handle: Long, pcmData: ByteArray): Int - private external fun nativeStreamTranscribeProcess(handle: Long, optionsJson: String?): String - private external fun nativeStreamTranscribeFinalize(handle: Long): String - private external fun nativeStreamTranscribeDestroy(handle: Long) -} - -actual class CactusIndex private constructor(private var handle: Long) : AutoCloseable { - - actual companion object { - init { - System.loadLibrary("cactus") - } - - actual fun create(indexDir: String, embeddingDim: Int): CactusIndex { - val handle = nativeIndexInit(indexDir, embeddingDim) - if (handle == 0L) { - throw CactusException("Failed to initialize index") - } - return CactusIndex(handle) - } - - @JvmStatic - private external fun nativeIndexInit(indexDir: String, embeddingDim: Int): Long - } - - actual fun add(ids: IntArray, documents: Array, embeddings: Array, metadatas: Array?) { - checkHandle() - val result = nativeIndexAdd(handle, ids, documents, metadatas, embeddings, embeddings[0].size) - if (result < 0) { - throw CactusException("Failed to add documents to index") - } - } - - actual fun delete(ids: IntArray) { - checkHandle() - val result = nativeIndexDelete(handle, ids) - if (result < 0) { - throw CactusException("Failed to delete documents from index") - } - } - - actual fun query(embedding: FloatArray, topK: Int): List { - checkHandle() - val responseJson = nativeIndexQuery(handle, embedding, topK, null) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - val results = json.getJSONArray("results") - return (0 until results.length()).map { i -> - val obj = results.getJSONObject(i) - IndexResult(obj.getInt("id"), obj.getDouble("score").toFloat()) - } - } - - actual fun compact() { - checkHandle() - val result = nativeIndexCompact(handle) - if (result < 0) { - throw CactusException("Failed to compact index") - } - } - - actual override fun close() { - if (handle != 0L) { - nativeIndexDestroy(handle) - handle = 0L - } - } - - private fun checkHandle() { - if (handle == 0L) { - throw CactusException("Index has been closed") - } - } - - private external fun nativeIndexAdd(handle: Long, ids: IntArray, documents: Array, metadatas: Array?, embeddings: Array, embeddingDim: Int): Int - private external fun nativeIndexDelete(handle: Long, ids: IntArray): Int - private external fun nativeIndexQuery(handle: Long, embedding: FloatArray, topK: Int, optionsJson: String?): String - private external fun nativeIndexCompact(handle: Long): Int - private external fun nativeIndexDestroy(handle: Long) -} - -private fun Message.toJson(): JSONObject = JSONObject().apply { - put("role", role) - put("content", content) -} - -private fun CompletionOptions.toJson(): String = JSONObject().apply { - put("temperature", temperature) - put("top_p", topP) - put("top_k", topK) - put("max_tokens", maxTokens) - put("stop", JSONArray(stopSequences)) - put("confidence_threshold", confidenceThreshold) -}.toString() - -private fun JSONObject.toCompletionResult(): CompletionResult { - val functionCalls = optJSONArray("function_calls")?.let { arr -> - (0 until arr.length()).map { arr.getJSONObject(it).toMap() } - } - return CompletionResult( - text = optString("text", ""), - functionCalls = functionCalls, - promptTokens = optInt("prompt_tokens", 0), - completionTokens = optInt("completion_tokens", 0), - timeToFirstToken = optDouble("time_to_first_token_ms", 0.0), - totalTime = optDouble("total_time_ms", 0.0), - prefillTokensPerSecond = optDouble("prefill_tokens_per_second", 0.0), - decodeTokensPerSecond = optDouble("decode_tokens_per_second", 0.0), - confidence = optDouble("confidence", 1.0), - needsCloudHandoff = optBoolean("cloud_handoff", false) - ) -} - -private fun JSONObject.toTranscriptionResult(): TranscriptionResult { - val segments = optJSONArray("segments")?.let { arr -> - (0 until arr.length()).map { arr.getJSONObject(it).toMap() } - } - return TranscriptionResult( - text = optString("text", ""), - segments = segments, - totalTime = optDouble("total_time_ms", 0.0) - ) -} - -private fun JSONObject.toMap(): Map { - val map = mutableMapOf() - keys().forEach { key -> - map[key] = when (val value = get(key)) { - is JSONObject -> value.toMap() - is JSONArray -> value.toList() - JSONObject.NULL -> Unit - else -> value - } - } - return map -} - -private fun JSONArray.toList(): List { - return (0 until length()).map { i -> - when (val value = get(i)) { - is JSONObject -> value.toMap() - is JSONArray -> value.toList() - JSONObject.NULL -> Unit - else -> value - } - } -} diff --git a/android/Cactus.common.kt b/android/Cactus.common.kt deleted file mode 100644 index 5eb0a6b7d..000000000 --- a/android/Cactus.common.kt +++ /dev/null @@ -1,105 +0,0 @@ -package com.cactus - -expect class Cactus : AutoCloseable { - companion object { - fun create(modelPath: String, corpusDir: String? = null): Cactus - fun setTelemetryToken(token: String) - fun setProKey(key: String) - } - - fun complete(prompt: String, options: CompletionOptions = CompletionOptions()): CompletionResult - fun complete( - messages: List, - options: CompletionOptions = CompletionOptions(), - tools: List>? = null, - callback: TokenCallback? = null - ): CompletionResult - - fun transcribe( - audioPath: String, - prompt: String? = null, - language: String? = null, - translate: Boolean = false - ): TranscriptionResult - - fun transcribe( - pcmData: ByteArray, - prompt: String? = null, - language: String? = null, - translate: Boolean = false - ): TranscriptionResult - - fun embed(text: String, normalize: Boolean = true): FloatArray - fun imageEmbed(imagePath: String): FloatArray - fun audioEmbed(audioPath: String): FloatArray - fun ragQuery(query: String, topK: Int = 5): String - fun tokenize(text: String): IntArray - fun scoreWindow(tokens: IntArray, start: Int, end: Int, context: Int): String - fun createStreamTranscriber(): StreamTranscriber - fun reset() - fun stop() - override fun close() -} - -expect class StreamTranscriber : AutoCloseable { - fun insert(pcmData: ByteArray) - fun process(language: String? = null): TranscriptionResult - fun finalize(): TranscriptionResult - override fun close() -} - -expect class CactusIndex : AutoCloseable { - companion object { - fun create(indexDir: String, embeddingDim: Int): CactusIndex - } - - fun add(ids: IntArray, documents: Array, embeddings: Array, metadatas: Array? = null) - fun delete(ids: IntArray) - fun query(embedding: FloatArray, topK: Int = 5): List - fun compact() - override fun close() -} - -data class Message(val role: String, val content: String) { - companion object { - fun system(content: String) = Message("system", content) - fun user(content: String) = Message("user", content) - fun assistant(content: String) = Message("assistant", content) - } -} - -data class CompletionOptions( - val temperature: Float = 0.7f, - val topP: Float = 0.9f, - val topK: Int = 40, - val maxTokens: Int = 512, - val stopSequences: List = emptyList(), - val confidenceThreshold: Float = 0f -) - -data class CompletionResult( - val text: String, - val functionCalls: List>?, - val promptTokens: Int, - val completionTokens: Int, - val timeToFirstToken: Double, - val totalTime: Double, - val prefillTokensPerSecond: Double, - val decodeTokensPerSecond: Double, - val confidence: Double, - val needsCloudHandoff: Boolean -) - -data class TranscriptionResult( - val text: String, - val segments: List>?, - val totalTime: Double -) - -fun interface TokenCallback { - fun onToken(token: String, tokenId: Int) -} - -data class IndexResult(val id: Int, val score: Float) - -class CactusException(message: String) : Exception(message) diff --git a/android/Cactus.ios.kt b/android/Cactus.ios.kt deleted file mode 100644 index f5444b0d9..000000000 --- a/android/Cactus.ios.kt +++ /dev/null @@ -1,592 +0,0 @@ -package com.cactus - -import cactus.* -import kotlinx.cinterop.* -import kotlinx.serialization.json.* -import platform.Foundation.NSLog - -@OptIn(ExperimentalForeignApi::class) -actual class Cactus private constructor(private var handle: COpaquePointer?) : AutoCloseable { - - actual companion object { - actual fun create(modelPath: String, corpusDir: String?): Cactus { - val handle = cactus_init(modelPath, corpusDir) - if (handle == null) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - return Cactus(handle) - } - - actual fun setTelemetryToken(token: String) { - cactus_set_telemetry_token(token) - } - - actual fun setProKey(key: String) { - cactus_set_pro_key(key) - } - } - - actual fun complete(prompt: String, options: CompletionOptions): CompletionResult { - return complete(listOf(Message.user(prompt)), options, null, null) - } - - actual fun complete( - messages: List, - options: CompletionOptions, - tools: List>?, - callback: TokenCallback? - ): CompletionResult { - checkHandle() - memScoped { - val buffer = allocArray(65536) - val messagesJson = serializeMessages(messages) - val optionsJson = serializeOptions(options) - val toolsJson = tools?.let { serializeTools(it) } - - val result = cactus_complete( - handle, - messagesJson, - buffer, - 65536u, - optionsJson, - toolsJson, - null, - null - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - val response = buffer.toKString() - return parseCompletionResult(response) - } - } - - actual fun transcribe( - audioPath: String, - prompt: String?, - language: String?, - translate: Boolean - ): TranscriptionResult { - checkHandle() - memScoped { - val buffer = allocArray(65536) - val optionsJson = buildJsonObject { - language?.let { put("language", it) } - put("translate", translate) - }.toString() - - val result = cactus_transcribe( - handle, - audioPath, - prompt, - buffer, - 65536u, - optionsJson, - null, - null, - null, - 0u - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return parseTranscriptionResult(buffer.toKString()) - } - } - - actual fun transcribe( - pcmData: ByteArray, - prompt: String?, - language: String?, - translate: Boolean - ): TranscriptionResult { - checkHandle() - memScoped { - val buffer = allocArray(65536) - val optionsJson = buildJsonObject { - language?.let { put("language", it) } - put("translate", translate) - }.toString() - - val pcmPtr = pcmData.refTo(0).getPointer(this) - - val result = cactus_transcribe( - handle, - null, - prompt, - buffer, - 65536u, - optionsJson, - null, - null, - pcmPtr.reinterpret(), - pcmData.size.toULong() - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return parseTranscriptionResult(buffer.toKString()) - } - } - - actual fun embed(text: String, normalize: Boolean): FloatArray { - checkHandle() - memScoped { - val buffer = allocArray(4096) - val dimPtr = alloc() - - val result = cactus_embed( - handle, - text, - buffer, - 4096u, - dimPtr.ptr, - normalize - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - val dim = dimPtr.value.toInt() - return FloatArray(dim) { buffer[it] } - } - } - - actual fun ragQuery(query: String, topK: Int): String { - checkHandle() - memScoped { - val buffer = allocArray(65536) - - val result = cactus_rag_query( - handle, - query, - buffer, - 65536u, - topK.toULong() - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return buffer.toKString() - } - } - - actual fun tokenize(text: String): IntArray { - checkHandle() - memScoped { - val buffer = allocArray(8192) - val tokenLen = alloc() - - val result = cactus_tokenize( - handle, - text, - buffer, - 8192u, - tokenLen.ptr - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return IntArray(tokenLen.value.toInt()) { buffer[it].toInt() } - } - } - - actual fun scoreWindow(tokens: IntArray, start: Int, end: Int, context: Int): String { - checkHandle() - memScoped { - val buffer = allocArray(65536) - val tokenBuffer = allocArray(tokens.size) - tokens.forEachIndexed { i, v -> tokenBuffer[i] = v.toUInt() } - - val result = cactus_score_window( - handle, - tokenBuffer, - tokens.size.toULong(), - start.toULong(), - end.toULong(), - context.toULong(), - buffer, - 65536u - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return buffer.toKString() - } - } - - actual fun imageEmbed(imagePath: String): FloatArray { - checkHandle() - memScoped { - val buffer = allocArray(4096) - val dimPtr = alloc() - - val result = cactus_image_embed( - handle, - imagePath, - buffer, - 4096u, - dimPtr.ptr - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return FloatArray(dimPtr.value.toInt()) { buffer[it] } - } - } - - actual fun audioEmbed(audioPath: String): FloatArray { - checkHandle() - memScoped { - val buffer = allocArray(4096) - val dimPtr = alloc() - - val result = cactus_audio_embed( - handle, - audioPath, - buffer, - 4096u, - dimPtr.ptr - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return FloatArray(dimPtr.value.toInt()) { buffer[it] } - } - } - - actual fun createStreamTranscriber(): StreamTranscriber { - checkHandle() - val streamHandle = cactus_stream_transcribe_init(handle) - if (streamHandle == null) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - return StreamTranscriber(streamHandle) - } - - actual fun reset() { - checkHandle() - cactus_reset(handle) - } - - actual fun stop() { - checkHandle() - cactus_stop(handle) - } - - actual override fun close() { - handle?.let { cactus_destroy(it) } - handle = null - } - - private fun checkHandle() { - if (handle == null) throw CactusException("Model has been closed") - } - - private fun serializeMessages(messages: List): String { - return buildJsonArray { - messages.forEach { msg -> - addJsonObject { - put("role", msg.role) - put("content", msg.content) - } - } - }.toString() - } - - private fun serializeOptions(options: CompletionOptions): String { - return buildJsonObject { - put("temperature", options.temperature) - put("top_p", options.topP) - put("top_k", options.topK) - put("max_tokens", options.maxTokens) - putJsonArray("stop") { options.stopSequences.forEach { add(it) } } - put("confidence_threshold", options.confidenceThreshold) - }.toString() - } - - private fun serializeTools(tools: List>): String { - return Json.encodeToString(tools) - } - - private fun parseCompletionResult(json: String): CompletionResult { - val obj = Json.parseToJsonElement(json).jsonObject - return CompletionResult( - text = obj["text"]?.jsonPrimitive?.contentOrNull ?: "", - functionCalls = obj["function_calls"]?.jsonArray?.map { it.jsonObject.toMap() }, - promptTokens = obj["prompt_tokens"]?.jsonPrimitive?.intOrNull ?: 0, - completionTokens = obj["completion_tokens"]?.jsonPrimitive?.intOrNull ?: 0, - timeToFirstToken = obj["time_to_first_token_ms"]?.jsonPrimitive?.doubleOrNull ?: 0.0, - totalTime = obj["total_time_ms"]?.jsonPrimitive?.doubleOrNull ?: 0.0, - prefillTokensPerSecond = obj["prefill_tokens_per_second"]?.jsonPrimitive?.doubleOrNull ?: 0.0, - decodeTokensPerSecond = obj["decode_tokens_per_second"]?.jsonPrimitive?.doubleOrNull ?: 0.0, - confidence = obj["confidence"]?.jsonPrimitive?.doubleOrNull ?: 1.0, - needsCloudHandoff = obj["cloud_handoff"]?.jsonPrimitive?.booleanOrNull ?: false - ) - } - - private fun parseTranscriptionResult(json: String): TranscriptionResult { - val obj = Json.parseToJsonElement(json).jsonObject - return TranscriptionResult( - text = obj["text"]?.jsonPrimitive?.contentOrNull ?: "", - segments = obj["segments"]?.jsonArray?.map { it.jsonObject.toMap() }, - totalTime = obj["total_time_ms"]?.jsonPrimitive?.doubleOrNull ?: 0.0 - ) - } - - private fun JsonObject.toMap(): Map { - return entries.associate { (k, v) -> - k to when (v) { - is JsonPrimitive -> v.contentOrNull ?: v.toString() - is JsonObject -> v.toMap() - is JsonArray -> v.map { it.toString() } - else -> v.toString() - } - } - } -} - -@OptIn(ExperimentalForeignApi::class) -actual class StreamTranscriber internal constructor(private var handle: COpaquePointer?) : AutoCloseable { - - actual fun insert(pcmData: ByteArray) { - checkHandle() - memScoped { - val pcmPtr = pcmData.refTo(0).getPointer(this) - val result = cactus_stream_transcribe_insert( - handle, - pcmPtr.reinterpret(), - pcmData.size.toULong() - ) - if (result < 0) { - throw CactusException("Failed to insert audio data") - } - } - } - - actual fun process(language: String?): TranscriptionResult { - checkHandle() - memScoped { - val buffer = allocArray(65536) - val optionsJson = language?.let { - buildJsonObject { put("language", it) }.toString() - } - - val result = cactus_stream_transcribe_process( - handle, - buffer, - 65536u, - optionsJson - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return parseTranscriptionResult(buffer.toKString()) - } - } - - actual fun finalize(): TranscriptionResult { - checkHandle() - memScoped { - val buffer = allocArray(65536) - - val result = cactus_stream_transcribe_finalize( - handle, - buffer, - 65536u - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return parseTranscriptionResult(buffer.toKString()) - } - } - - actual override fun close() { - handle?.let { cactus_stream_transcribe_destroy(it) } - handle = null - } - - private fun checkHandle() { - if (handle == null) throw CactusException("Stream transcriber has been closed") - } - - private fun parseTranscriptionResult(json: String): TranscriptionResult { - val obj = Json.parseToJsonElement(json).jsonObject - return TranscriptionResult( - text = obj["text"]?.jsonPrimitive?.contentOrNull ?: "", - segments = obj["segments"]?.jsonArray?.map { it.jsonObject.toMap() }, - totalTime = obj["total_time_ms"]?.jsonPrimitive?.doubleOrNull ?: 0.0 - ) - } - - private fun JsonObject.toMap(): Map { - return entries.associate { (k, v) -> - k to when (v) { - is JsonPrimitive -> v.contentOrNull ?: v.toString() - is JsonObject -> v.toMap() - is JsonArray -> v.map { it.toString() } - else -> v.toString() - } - } - } -} - -@OptIn(ExperimentalForeignApi::class) -actual class CactusIndex private constructor(private var handle: COpaquePointer?) : AutoCloseable { - - actual companion object { - actual fun create(indexDir: String, embeddingDim: Int): CactusIndex { - val handle = cactus_index_init(indexDir, embeddingDim.toULong()) - if (handle == null) { - throw CactusException("Failed to initialize index") - } - return CactusIndex(handle) - } - } - - actual fun add(ids: IntArray, documents: Array, embeddings: Array, metadatas: Array?) { - checkHandle() - memScoped { - val idPtr = allocArray(ids.size) - ids.forEachIndexed { i, v -> idPtr[i] = v } - - val docPtrs = allocArray>(documents.size) - documents.forEachIndexed { i, doc -> docPtrs[i] = doc.cstr.ptr } - - val metaPtrs = metadatas?.let { - val ptrs = allocArray>(it.size) - it.forEachIndexed { i, meta -> ptrs[i] = meta.cstr.ptr } - ptrs - } - - val embPtrs = allocArray>(embeddings.size) - embeddings.forEachIndexed { i, emb -> - val embArr = allocArray(emb.size) - emb.forEachIndexed { j, v -> embArr[j] = v } - embPtrs[i] = embArr - } - - val result = cactus_index_add( - handle, - idPtr, - docPtrs, - metaPtrs, - embPtrs, - ids.size.toULong(), - embeddings[0].size.toULong() - ) - - if (result < 0) { - throw CactusException("Failed to add documents to index") - } - } - } - - actual fun delete(ids: IntArray) { - checkHandle() - memScoped { - val idPtr = allocArray(ids.size) - ids.forEachIndexed { i, v -> idPtr[i] = v } - - val result = cactus_index_delete(handle, idPtr, ids.size.toULong()) - if (result < 0) { - throw CactusException("Failed to delete documents from index") - } - } - } - - actual fun query(embedding: FloatArray, topK: Int): List { - checkHandle() - memScoped { - val embArr = allocArray(embedding.size) - embedding.forEachIndexed { i, v -> embArr[i] = v } - val embPtr = alloc>() - embPtr.value = embArr - - val idBuffer = allocArray(topK) - val scoreBuffer = allocArray(topK) - val idBufferSize = alloc() - val scoreBufferSize = alloc() - idBufferSize.value = topK.toULong() - scoreBufferSize.value = topK.toULong() - - val idPtrPtr = alloc>() - idPtrPtr.value = idBuffer - val scorePtrPtr = alloc>() - scorePtrPtr.value = scoreBuffer - - val result = cactus_index_query( - handle, - embPtr.ptr, - 1u, - embedding.size.toULong(), - null, - idPtrPtr.ptr, - idBufferSize.ptr, - scorePtrPtr.ptr, - scoreBufferSize.ptr - ) - - if (result < 0) { - val error = cactus_get_last_error()?.toKString() ?: "Unknown error" - throw CactusException(error) - } - - return (0 until idBufferSize.value.toInt()).map { i -> - IndexResult(idBuffer[i], scoreBuffer[i]) - } - } - } - - actual fun compact() { - checkHandle() - val result = cactus_index_compact(handle) - if (result < 0) { - throw CactusException("Failed to compact index") - } - } - - actual override fun close() { - handle?.let { cactus_index_destroy(it) } - handle = null - } - - private fun checkHandle() { - if (handle == null) throw CactusException("Index has been closed") - } -} diff --git a/android/Cactus.kt b/android/Cactus.kt deleted file mode 100644 index 2316502ea..000000000 --- a/android/Cactus.kt +++ /dev/null @@ -1,441 +0,0 @@ -package com.cactus - -import org.json.JSONArray -import org.json.JSONObject -import java.io.Closeable - -class Cactus private constructor(private var handle: Long) : Closeable { - - companion object { - init { - System.loadLibrary("cactus") - } - - @JvmStatic - fun create(modelPath: String, corpusDir: String? = null): Cactus { - val handle = nativeInit(modelPath, corpusDir) - if (handle == 0L) { - throw CactusException(nativeGetLastError().ifEmpty { "Failed to initialize model" }) - } - return Cactus(handle) - } - - @JvmStatic - fun setTelemetryToken(token: String) = nativeSetTelemetryToken(token) - - @JvmStatic - fun setProKey(key: String) = nativeSetProKey(key) - - @JvmStatic - private external fun nativeInit(modelPath: String, corpusDir: String?): Long - @JvmStatic - private external fun nativeGetLastError(): String - @JvmStatic - private external fun nativeSetTelemetryToken(token: String) - @JvmStatic - private external fun nativeSetProKey(key: String) - } - - fun complete(prompt: String, options: CompletionOptions = CompletionOptions()): CompletionResult { - return complete(listOf(Message.user(prompt)), options) - } - - fun complete( - messages: List, - options: CompletionOptions = CompletionOptions(), - tools: List>? = null, - callback: TokenCallback? = null - ): CompletionResult { - checkHandle() - val messagesJson = JSONArray(messages.map { it.toJson() }).toString() - val toolsJson = tools?.let { JSONArray(it.map { t -> JSONObject(t) }).toString() } - - val responseJson = nativeComplete(handle, messagesJson, options.toJson(), toolsJson, callback) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return json.toCompletionResult() - } - - fun transcribe( - audioPath: String, - prompt: String? = null, - language: String? = null, - translate: Boolean = false - ): TranscriptionResult { - checkHandle() - val optionsJson = JSONObject().apply { - language?.let { put("language", it) } - put("translate", translate) - }.toString() - - val responseJson = nativeTranscribe(handle, audioPath, prompt, optionsJson, null) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return json.toTranscriptionResult() - } - - fun transcribe( - pcmData: ByteArray, - prompt: String? = null, - language: String? = null, - translate: Boolean = false - ): TranscriptionResult { - checkHandle() - val optionsJson = JSONObject().apply { - language?.let { put("language", it) } - put("translate", translate) - }.toString() - - val responseJson = nativeTranscribe(handle, null, prompt, optionsJson, pcmData) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return json.toTranscriptionResult() - } - - fun embed(text: String, normalize: Boolean = true): FloatArray { - checkHandle() - return nativeEmbed(handle, text, normalize) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to generate embedding" }) - } - - fun ragQuery(query: String, topK: Int = 5): String { - checkHandle() - val responseJson = nativeRagQuery(handle, query, topK) - val json = JSONObject(responseJson) - - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - - return responseJson - } - - fun reset() { - checkHandle() - nativeReset(handle) - } - - fun stop() { - checkHandle() - nativeStop(handle) - } - - fun tokenize(text: String): IntArray { - checkHandle() - return nativeTokenize(handle, text) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to tokenize" }) - } - - fun scoreWindow(tokens: IntArray, start: Int, end: Int, context: Int): String { - checkHandle() - val responseJson = nativeScoreWindow(handle, tokens, start, end, context) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - return responseJson - } - - fun imageEmbed(imagePath: String): FloatArray { - checkHandle() - return nativeImageEmbed(handle, imagePath) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to generate image embedding" }) - } - - fun audioEmbed(audioPath: String): FloatArray { - checkHandle() - return nativeAudioEmbed(handle, audioPath) - ?: throw CactusException(nativeGetLastError().ifEmpty { "Failed to generate audio embedding" }) - } - - fun createStreamTranscriber(): StreamTranscriber { - checkHandle() - val streamHandle = nativeStreamTranscribeInit(handle) - if (streamHandle == 0L) { - throw CactusException(nativeGetLastError().ifEmpty { "Failed to create stream transcriber" }) - } - return StreamTranscriber(streamHandle) - } - - override fun close() { - if (handle != 0L) { - nativeDestroy(handle) - handle = 0L - } - } - - private fun checkHandle() { - if (handle == 0L) { - throw CactusException("Model has been closed") - } - } - - private external fun nativeDestroy(handle: Long) - private external fun nativeReset(handle: Long) - private external fun nativeStop(handle: Long) - private external fun nativeComplete(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, callback: TokenCallback?): String - private external fun nativeTranscribe(handle: Long, audioPath: String?, prompt: String?, optionsJson: String?, pcmData: ByteArray?): String - private external fun nativeEmbed(handle: Long, text: String, normalize: Boolean): FloatArray? - private external fun nativeRagQuery(handle: Long, query: String, topK: Int): String - private external fun nativeTokenize(handle: Long, text: String): IntArray? - private external fun nativeScoreWindow(handle: Long, tokens: IntArray, start: Int, end: Int, context: Int): String - private external fun nativeImageEmbed(handle: Long, imagePath: String): FloatArray? - private external fun nativeAudioEmbed(handle: Long, audioPath: String): FloatArray? - private external fun nativeStreamTranscribeInit(handle: Long): Long -} - -class StreamTranscriber internal constructor(private var handle: Long) : Closeable { - - fun insert(pcmData: ByteArray) { - checkHandle() - val result = nativeStreamTranscribeInsert(handle, pcmData) - if (result < 0) { - throw CactusException("Failed to insert audio data") - } - } - - fun process(language: String? = null): TranscriptionResult { - checkHandle() - val optionsJson = language?.let { JSONObject().put("language", it).toString() } - val responseJson = nativeStreamTranscribeProcess(handle, optionsJson) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - return json.toTranscriptionResult() - } - - fun finalize(): TranscriptionResult { - checkHandle() - val responseJson = nativeStreamTranscribeFinalize(handle) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - return json.toTranscriptionResult() - } - - override fun close() { - if (handle != 0L) { - nativeStreamTranscribeDestroy(handle) - handle = 0L - } - } - - private fun checkHandle() { - if (handle == 0L) { - throw CactusException("Stream transcriber has been closed") - } - } - - private external fun nativeStreamTranscribeInsert(handle: Long, pcmData: ByteArray): Int - private external fun nativeStreamTranscribeProcess(handle: Long, optionsJson: String?): String - private external fun nativeStreamTranscribeFinalize(handle: Long): String - private external fun nativeStreamTranscribeDestroy(handle: Long) -} - -class CactusIndex private constructor(private var handle: Long) : Closeable { - - companion object { - init { - System.loadLibrary("cactus") - } - - @JvmStatic - fun create(indexDir: String, embeddingDim: Int): CactusIndex { - val handle = nativeIndexInit(indexDir, embeddingDim) - if (handle == 0L) { - throw CactusException("Failed to initialize index") - } - return CactusIndex(handle) - } - - @JvmStatic - private external fun nativeIndexInit(indexDir: String, embeddingDim: Int): Long - } - - fun add( - ids: IntArray, - documents: Array, - embeddings: Array, - metadatas: Array? = null - ) { - checkHandle() - val result = nativeIndexAdd(handle, ids, documents, metadatas, embeddings, embeddings[0].size) - if (result < 0) { - throw CactusException("Failed to add documents to index") - } - } - - fun delete(ids: IntArray) { - checkHandle() - val result = nativeIndexDelete(handle, ids) - if (result < 0) { - throw CactusException("Failed to delete documents from index") - } - } - - fun query(embedding: FloatArray, topK: Int = 5): List { - checkHandle() - val responseJson = nativeIndexQuery(handle, embedding, topK, null) - val json = JSONObject(responseJson) - if (json.has("error")) { - throw CactusException(json.getString("error")) - } - val results = json.getJSONArray("results") - return (0 until results.length()).map { i -> - val obj = results.getJSONObject(i) - IndexResult(obj.getInt("id"), obj.getDouble("score").toFloat()) - } - } - - fun compact() { - checkHandle() - val result = nativeIndexCompact(handle) - if (result < 0) { - throw CactusException("Failed to compact index") - } - } - - override fun close() { - if (handle != 0L) { - nativeIndexDestroy(handle) - handle = 0L - } - } - - private fun checkHandle() { - if (handle == 0L) { - throw CactusException("Index has been closed") - } - } - - private external fun nativeIndexAdd(handle: Long, ids: IntArray, documents: Array, metadatas: Array?, embeddings: Array, embeddingDim: Int): Int - private external fun nativeIndexDelete(handle: Long, ids: IntArray): Int - private external fun nativeIndexQuery(handle: Long, embedding: FloatArray, topK: Int, optionsJson: String?): String - private external fun nativeIndexCompact(handle: Long): Int - private external fun nativeIndexDestroy(handle: Long) -} - -data class IndexResult(val id: Int, val score: Float) - -data class Message(val role: String, val content: String) { - companion object { - fun system(content: String) = Message("system", content) - fun user(content: String) = Message("user", content) - fun assistant(content: String) = Message("assistant", content) - } - - fun toJson(): JSONObject = JSONObject().apply { - put("role", role) - put("content", content) - } -} - -data class CompletionOptions( - val temperature: Float = 0.7f, - val topP: Float = 0.9f, - val topK: Int = 40, - val maxTokens: Int = 512, - val stopSequences: List = emptyList(), - val confidenceThreshold: Float = 0f -) { - fun toJson(): String = JSONObject().apply { - put("temperature", temperature) - put("top_p", topP) - put("top_k", topK) - put("max_tokens", maxTokens) - put("stop", JSONArray(stopSequences)) - put("confidence_threshold", confidenceThreshold) - }.toString() -} - -data class CompletionResult( - val text: String, - val functionCalls: List>?, - val promptTokens: Int, - val completionTokens: Int, - val timeToFirstToken: Double, - val totalTime: Double, - val prefillTokensPerSecond: Double, - val decodeTokensPerSecond: Double, - val confidence: Double, - val needsCloudHandoff: Boolean -) - -data class TranscriptionResult( - val text: String, - val segments: List>?, - val totalTime: Double -) - -fun interface TokenCallback { - fun onToken(token: String, tokenId: Int) -} - -class CactusException(message: String) : Exception(message) - -private fun JSONObject.toCompletionResult(): CompletionResult { - val functionCalls = optJSONArray("function_calls")?.let { arr -> - (0 until arr.length()).map { arr.getJSONObject(it).toMap() } - } - return CompletionResult( - text = optString("text", ""), - functionCalls = functionCalls, - promptTokens = optInt("prompt_tokens", 0), - completionTokens = optInt("completion_tokens", 0), - timeToFirstToken = optDouble("time_to_first_token_ms", 0.0), - totalTime = optDouble("total_time_ms", 0.0), - prefillTokensPerSecond = optDouble("prefill_tokens_per_second", 0.0), - decodeTokensPerSecond = optDouble("decode_tokens_per_second", 0.0), - confidence = optDouble("confidence", 1.0), - needsCloudHandoff = optBoolean("cloud_handoff", false) - ) -} - -private fun JSONObject.toTranscriptionResult(): TranscriptionResult { - val segments = optJSONArray("segments")?.let { arr -> - (0 until arr.length()).map { arr.getJSONObject(it).toMap() } - } - return TranscriptionResult( - text = optString("text", ""), - segments = segments, - totalTime = optDouble("total_time_ms", 0.0) - ) -} - -private fun JSONObject.toMap(): Map { - val map = mutableMapOf() - keys().forEach { key -> - map[key] = when (val value = get(key)) { - is JSONObject -> value.toMap() - is JSONArray -> value.toList() - JSONObject.NULL -> Unit - else -> value - } - } - return map -} - -private fun JSONArray.toList(): List { - return (0 until length()).map { i -> - when (val value = get(i)) { - is JSONObject -> value.toMap() - is JSONArray -> value.toList() - JSONObject.NULL -> Unit - else -> value - } - } -} diff --git a/android/README.md b/android/README.md index dbff3d3d9..fa73fa740 100644 --- a/android/README.md +++ b/android/README.md @@ -1,300 +1,34 @@ -# Cactus for Android & Kotlin Multiplatform +# Android Build -Run AI models on-device with a simple Kotlin API. +Builds `libcactus_engine` for Android (arm64-v8a). -## Building +## Usage ```bash cactus build --android ``` -Build output: `android/build/lib/libcactus.so` - -see the main [README.md](../README.md) for how to use CLI & download weight - -## Integration - -### Android-only - -1. Copy `libcactus.so` to `app/src/main/jniLibs/arm64-v8a/` -2. Copy `Cactus.kt` to `app/src/main/java/com/cactus/` - -### Kotlin Multiplatform - -Source files: - -| File | Copy to | -|------|---------| -| `Cactus.common.kt` | `shared/src/commonMain/kotlin/com/cactus/` | -| `Cactus.android.kt` | `shared/src/androidMain/kotlin/com/cactus/` | -| `Cactus.ios.kt` | `shared/src/iosMain/kotlin/com/cactus/` | -| `cactus.def` | `shared/src/nativeInterop/cinterop/` | - -Binary files: - -| Platform | Location | -|----------|----------| -| Android | `libcactus.so` → `app/src/main/jniLibs/arm64-v8a/` | -| iOS | `libcactus-device.a` → link via cinterop | - -build.gradle.kts: - -```kotlin -kotlin { - androidTarget() - - listOf(iosArm64(), iosSimulatorArm64()).forEach { - it.compilations.getByName("main") { - cinterops { - create("cactus") { - defFile("src/nativeInterop/cinterop/cactus.def") - includeDirs("/path/to/cactus/ffi") - } - } - } - it.binaries.framework { - linkerOpts("-L/path/to/apple", "-lcactus-device") - } - } - - sourceSets { - commonMain.dependencies { - implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.6.0") - } - } -} -``` - -## Usage - -### Basic Completion - -```kotlin -import com.cactus.* +Or directly: -val model = Cactus.create("/path/to/model") -val result = model.complete("What is the capital of France?") -model.close() -``` - -### Chat Messages - -```kotlin -Cactus.create(modelPath).use { model -> - val result = model.complete( - messages = listOf( - Message.system("You are a helpful assistant."), - Message.user("What is 2 + 2?") - ) - ) - println(result.text) -} -``` - -### Completion Options - -```kotlin -val options = CompletionOptions( - temperature = 0.7f, - topP = 0.9f, - topK = 40, - maxTokens = 256, - stopSequences = listOf("\n\n") -) - -val result = model.complete("Write a haiku:", options) -``` - -### Streaming Tokens - -```kotlin -val result = model.complete( - messages = listOf(Message.user("Tell me a story")), - callback = TokenCallback { token, tokenId -> - print(token) - } -) -``` - -### Audio Transcription - -```kotlin -val result = model.transcribe("/path/to/audio.wav") - -val pcmData: ByteArray = ... // 16kHz mono PCM -val result = model.transcribe(pcmData) -``` - -### Embeddings - -```kotlin -val embedding = model.embed("Hello, world!") -val imageEmbedding = model.imageEmbed("/path/to/image.jpg") -val audioEmbedding = model.audioEmbed("/path/to/audio.wav") -``` - -### Tokenization - -```kotlin -val tokens = model.tokenize("Hello, world!") -val scores = model.scoreWindow(tokens, start = 0, end = tokens.size, context = 512) -``` - -### Streaming Transcription - -```kotlin -model.createStreamTranscriber().use { stream -> - stream.insert(audioChunk1) - stream.insert(audioChunk2) - val partial = stream.process() - println("Partial: ${partial.text}") - val final = stream.finalize() - println("Final: ${final.text}") -} -``` - -### RAG - -```kotlin -val model = Cactus.create( - modelPath = "/path/to/model", - corpusDir = "/path/to/documents" -) -val result = model.complete("What does the documentation say about X?") -``` - -### Vector Index - -```kotlin -CactusIndex.create("/path/to/index", embeddingDim = 384).use { index -> - val embeddings = arrayOf(model.embed("doc1"), model.embed("doc2")) - index.add( - ids = intArrayOf(1, 2), - documents = arrayOf("Document 1", "Document 2"), - embeddings = embeddings - ) - val results = index.query(model.embed("search query"), topK = 5) - results.forEach { println("ID: ${it.id}, Score: ${it.score}") } -} -``` - -## API Reference - -### Cactus - -```kotlin -object Cactus { - fun create(modelPath: String, corpusDir: String? = null): Cactus - fun setTelemetryToken(token: String) - fun setProKey(key: String) -} - -fun complete(prompt: String, options: CompletionOptions = CompletionOptions()): CompletionResult -fun complete(messages: List, options: CompletionOptions = CompletionOptions(), tools: List>? = null, callback: TokenCallback? = null): CompletionResult -fun transcribe(audioPath: String, prompt: String? = null, language: String? = null, translate: Boolean = false): TranscriptionResult -fun transcribe(pcmData: ByteArray, prompt: String? = null, language: String? = null, translate: Boolean = false): TranscriptionResult -fun embed(text: String, normalize: Boolean = true): FloatArray -fun imageEmbed(imagePath: String): FloatArray -fun audioEmbed(audioPath: String): FloatArray -fun ragQuery(query: String, topK: Int = 5): String -fun tokenize(text: String): IntArray -fun scoreWindow(tokens: IntArray, start: Int, end: Int, context: Int): String -fun createStreamTranscriber(): StreamTranscriber -fun reset() -fun stop() -fun close() -``` - -### Message - -```kotlin -data class Message(val role: String, val content: String) { - companion object { - fun system(content: String): Message - fun user(content: String): Message - fun assistant(content: String): Message - } -} -``` - -### CompletionOptions - -```kotlin -data class CompletionOptions( - val temperature: Float = 0.7f, - val topP: Float = 0.9f, - val topK: Int = 40, - val maxTokens: Int = 512, - val stopSequences: List = emptyList(), - val confidenceThreshold: Float = 0f -) -``` - -### CompletionResult - -```kotlin -data class CompletionResult( - val text: String, - val functionCalls: List>?, - val promptTokens: Int, - val completionTokens: Int, - val timeToFirstToken: Double, - val totalTime: Double, - val prefillTokensPerSecond: Double, - val decodeTokensPerSecond: Double, - val confidence: Double, - val needsCloudHandoff: Boolean -) -``` - -### TranscriptionResult - -```kotlin -data class TranscriptionResult( - val text: String, - val segments: List>?, - val totalTime: Double -) -``` - -### TokenCallback - -```kotlin -fun interface TokenCallback { - fun onToken(token: String, tokenId: Int) -} -``` - -### StreamTranscriber - -```kotlin -class StreamTranscriber : Closeable { - fun insert(pcmData: ByteArray) - fun process(language: String? = null): TranscriptionResult - fun finalize(): TranscriptionResult - fun close() -} +```bash +bash android/build.sh ``` -### CactusIndex +## Output -```kotlin -class CactusIndex : Closeable { - companion object { - fun create(indexDir: String, embeddingDim: Int): CactusIndex - } +- `libcactus_engine.so` — Shared library (JNI, for Android apps) +- `libcactus_engine.a` — Static library (for native test binaries) - fun add(ids: IntArray, documents: Array, embeddings: Array, metadatas: Array? = null) - fun delete(ids: IntArray) - fun query(embedding: FloatArray, topK: Int = 5): List - fun compact() - fun close() -} +## Options -data class IndexResult(val id: Int, val score: Float) -``` +| Variable | Default | Description | +|----------|---------|-------------| +| `ANDROID_NDK_HOME` | Auto-detected | Android NDK path | +| `ANDROID_PLATFORM` | `android-21` | Minimum API level | +| `CMAKE_BUILD_TYPE` | `Release` | CMake build type | +| `CACTUS_CURL_ROOT` | `cactus-engine/libs/curl` | Vendored libcurl path | ## Requirements -- Android API 24+ / arm64-v8a -- iOS 14+ / arm64 (KMP only) +- Android NDK (install via Android Studio > SDK Tools > NDK) +- CMake 3.10+ diff --git a/android/build.sh b/android/build.sh index e88b012bc..b139e6230 100755 --- a/android/build.sh +++ b/android/build.sh @@ -1,4 +1,5 @@ -#!/bin/bash -e +#!/bin/bash +set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" @@ -7,16 +8,24 @@ ANDROID_DIR="$PROJECT_ROOT/android" ANDROID_PLATFORM=${ANDROID_PLATFORM:-android-21} CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE:-Release} BUILD_DIR="$ANDROID_DIR/build" +CACTUS_CURL_ROOT=${CACTUS_CURL_ROOT:-"$PROJECT_ROOT/cactus-engine/libs/curl"} -if [ -z "$ANDROID_NDK_HOME" ]; then - if [ -n "$ANDROID_HOME" ]; then +if [ -z "${ANDROID_NDK_HOME:-}" ]; then + if [ -n "${ANDROID_HOME:-}" ]; then ANDROID_NDK_HOME=$(ls -d "$ANDROID_HOME/ndk/"* 2>/dev/null | sort -V | tail -1) - elif [ -d "$HOME/Library/Android/sdk" ]; then + fi + if [ -z "${ANDROID_NDK_HOME:-}" ] && [ -d "$HOME/Library/Android/sdk" ]; then ANDROID_NDK_HOME=$(ls -d "$HOME/Library/Android/sdk/ndk/"* 2>/dev/null | sort -V | tail -1) fi + if [ -z "${ANDROID_NDK_HOME:-}" ] && [ -d "/opt/homebrew/Caskroom/android-ndk" ]; then + ANDROID_NDK_HOME=$(ls -d /opt/homebrew/Caskroom/android-ndk/*/AndroidNDK*.app/Contents/NDK 2>/dev/null | sort -V | tail -1) + fi + if [ -z "${ANDROID_NDK_HOME:-}" ] && [ -d "/opt/homebrew/share/android-ndk" ]; then + ANDROID_NDK_HOME="/opt/homebrew/share/android-ndk" + fi fi -if [ -z "$ANDROID_NDK_HOME" ] || [ ! -d "$ANDROID_NDK_HOME" ]; then +if [ -z "${ANDROID_NDK_HOME:-}" ] || [ ! -d "$ANDROID_NDK_HOME" ]; then echo "Error: Android NDK not found." echo "Set ANDROID_NDK_HOME or install NDK via Android SDK Manager" exit 1 @@ -38,24 +47,26 @@ echo "Building Cactus for Android ($ABI)..." echo "Build type: $CMAKE_BUILD_TYPE" echo "Using $n_cpu CPU cores" echo "Android CMakeLists.txt: $ANDROID_DIR/CMakeLists.txt" +echo "Vendored libcurl root: $CACTUS_CURL_ROOT" cmake -DCMAKE_TOOLCHAIN_FILE="$CMAKE_TOOLCHAIN_FILE" \ -DANDROID_ABI="$ABI" \ -DANDROID_PLATFORM="$ANDROID_PLATFORM" \ -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ + -DCACTUS_CURL_ROOT="$CACTUS_CURL_ROOT" \ -S "$ANDROID_DIR" \ -B "$BUILD_DIR" >/dev/null cmake --build "$BUILD_DIR" --config "$CMAKE_BUILD_TYPE" -j "$n_cpu" >/dev/null -cp "$BUILD_DIR/lib/libcactus.so" "$ANDROID_DIR/" 2>/dev/null || \ - cp "$BUILD_DIR/libcactus.so" "$ANDROID_DIR" 2>/dev/null || \ - { echo "Error: Could not find libcactus.so"; exit 1; } +cp "$BUILD_DIR/lib/libcactus_engine.so" "$ANDROID_DIR/" 2>/dev/null || \ + cp "$BUILD_DIR/libcactus_engine.so" "$ANDROID_DIR" 2>/dev/null || \ + { echo "Error: Could not find libcactus_engine.so"; exit 1; } -cp "$BUILD_DIR/lib/libcactus_static.a" "$ANDROID_DIR/libcactus.a" 2>/dev/null || \ - cp "$BUILD_DIR/libcactus_static.a" "$ANDROID_DIR/libcactus.a" 2>/dev/null || \ - { echo "Warning: Could not find libcactus_static.a"; } +cp "$BUILD_DIR/lib/libcactus_engine.a" "$ANDROID_DIR/libcactus_engine.a" 2>/dev/null || \ + cp "$BUILD_DIR/libcactus_engine.a" "$ANDROID_DIR/libcactus_engine.a" 2>/dev/null || \ + { echo "Warning: Could not find libcactus_engine.a"; } echo "Build complete!" -echo "Shared library location: $ANDROID_DIR/libcactus.so" -echo "Static library location: $ANDROID_DIR/libcactus.a" \ No newline at end of file +echo "Shared library location: $ANDROID_DIR/libcactus_engine.so" +echo "Static library location: $ANDROID_DIR/libcactus_engine.a" diff --git a/android/cactus.def b/android/cactus.def deleted file mode 100644 index cf3e7ba7e..000000000 --- a/android/cactus.def +++ /dev/null @@ -1,3 +0,0 @@ -package = cactus -headers = cactus_ffi.h -headerFilter = cactus_ffi.h diff --git a/android/cactus_jni.cpp b/android/cactus_jni.cpp index 38b166ab4..44d06d188 100644 --- a/android/cactus_jni.cpp +++ b/android/cactus_jni.cpp @@ -1,296 +1,202 @@ #include -#include -#include -#include -#include "cactus_ffi.h" - -static constexpr size_t DEFAULT_BUFFER_SIZE = 65536; - -static const char* jstring_to_cstr(JNIEnv* env, jstring str) { - if (str == nullptr) return nullptr; - return env->GetStringUTFChars(str, nullptr); -} - -static void release_jstring(JNIEnv* env, jstring str, const char* cstr) { - if (str != nullptr && cstr != nullptr) { - env->ReleaseStringUTFChars(str, cstr); - } -} +#include "cactus_engine.h" struct TokenCallbackContext { - JNIEnv* env; + JavaVM* jvm; + jobject callback; + jmethodID method; +}; + +struct LogCallbackContext { + JavaVM* jvm; jobject callback; jmethodID method; }; +static LogCallbackContext* g_log_callback_ctx = nullptr; + +extern "C" { + static void token_callback_bridge(const char* token, uint32_t token_id, void* user_data) { - if (!user_data) return; auto* ctx = static_cast(user_data); - jstring jtoken = ctx->env->NewStringUTF(token); - ctx->env->CallVoidMethod(ctx->callback, ctx->method, jtoken, static_cast(token_id)); - ctx->env->DeleteLocalRef(jtoken); + JNIEnv* env = nullptr; + bool attached = false; + jint status = ctx->jvm->GetEnv(reinterpret_cast(&env), JNI_VERSION_1_6); + if (status == JNI_EDETACHED) { + if (ctx->jvm->AttachCurrentThread(&env, nullptr) != JNI_OK) return; + attached = true; + } + if (!env) return; + jstring jtoken = token ? env->NewStringUTF(token) : nullptr; + env->CallVoidMethod(ctx->callback, ctx->method, jtoken, static_cast(token_id)); + if (env->ExceptionCheck()) { env->ExceptionDescribe(); env->ExceptionClear(); } + if (jtoken) env->DeleteLocalRef(jtoken); + if (attached) ctx->jvm->DetachCurrentThread(); } -extern "C" { +static void log_callback_bridge(int level, const char* component, const char* message, void* user_data) { + auto* ctx = static_cast(user_data); + JNIEnv* env = nullptr; + bool attached = false; + jint status = ctx->jvm->GetEnv(reinterpret_cast(&env), JNI_VERSION_1_6); + if (status == JNI_EDETACHED) { + if (ctx->jvm->AttachCurrentThread(&env, nullptr) != JNI_OK) return; + attached = true; + } + if (!env) return; + jstring jcomponent = component ? env->NewStringUTF(component) : nullptr; + jstring jmessage = message ? env->NewStringUTF(message) : nullptr; + env->CallVoidMethod(ctx->callback, ctx->method, static_cast(level), jcomponent, jmessage); + if (env->ExceptionCheck()) { env->ExceptionDescribe(); env->ExceptionClear(); } + if (jcomponent) env->DeleteLocalRef(jcomponent); + if (jmessage) env->DeleteLocalRef(jmessage); + if (attached) ctx->jvm->DetachCurrentThread(); +} + +JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM*, void*) { + return JNI_VERSION_1_6; +} JNIEXPORT jlong JNICALL -Java_com_cactus_Cactus_nativeInit(JNIEnv* env, jobject, jstring modelPath, jstring corpusDir) { - const char* path = jstring_to_cstr(env, modelPath); - const char* corpus = jstring_to_cstr(env, corpusDir); - jlong handle = reinterpret_cast(cactus_init(path, corpus)); - release_jstring(env, modelPath, path); - release_jstring(env, corpusDir, corpus); +Java_com_cactus_CactusJNI_nativeInit(JNIEnv* env, jobject, jstring modelPath, jstring corpusDir, jboolean cacheIndex) { + const char* path = env->GetStringUTFChars(modelPath, nullptr); + const char* corpus = corpusDir ? env->GetStringUTFChars(corpusDir, nullptr) : nullptr; + jlong handle = reinterpret_cast(cactus_init(path, corpus, cacheIndex == JNI_TRUE)); + env->ReleaseStringUTFChars(modelPath, path); + if (corpus) env->ReleaseStringUTFChars(corpusDir, corpus); return handle; } JNIEXPORT void JNICALL -Java_com_cactus_Cactus_nativeDestroy(JNIEnv*, jobject, jlong handle) { - if (handle != 0) { - cactus_destroy(reinterpret_cast(handle)); - } +Java_com_cactus_CactusJNI_nativeDestroy(JNIEnv*, jobject, jlong handle) { + cactus_destroy(reinterpret_cast(handle)); } JNIEXPORT void JNICALL -Java_com_cactus_Cactus_nativeReset(JNIEnv*, jobject, jlong handle) { - if (handle != 0) { - cactus_reset(reinterpret_cast(handle)); - } +Java_com_cactus_CactusJNI_nativeReset(JNIEnv*, jobject, jlong handle) { + cactus_reset(reinterpret_cast(handle)); } JNIEXPORT void JNICALL -Java_com_cactus_Cactus_nativeStop(JNIEnv*, jobject, jlong handle) { - if (handle != 0) { - cactus_stop(reinterpret_cast(handle)); - } +Java_com_cactus_CactusJNI_nativeStop(JNIEnv*, jobject, jlong handle) { + cactus_stop(reinterpret_cast(handle)); } -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeComplete(JNIEnv* env, jobject, jlong handle, - jstring messagesJson, jstring optionsJson, - jstring toolsJson, jobject callback) { - if (handle == 0) { - return env->NewStringUTF("{\"error\":\"Model not initialized\"}"); - } - - const char* messages = jstring_to_cstr(env, messagesJson); - const char* options = jstring_to_cstr(env, optionsJson); - const char* tools = jstring_to_cstr(env, toolsJson); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeComplete(JNIEnv* env, jobject, jlong handle, + jstring messagesJson, jbyteArray responseBuffer, + jstring optionsJson, jstring toolsJson, + jobject callback, jbyteArray pcmData) { + const char* messages = env->GetStringUTFChars(messagesJson, nullptr); + const char* options = optionsJson ? env->GetStringUTFChars(optionsJson, nullptr) : nullptr; + const char* tools = toolsJson ? env->GetStringUTFChars(toolsJson, nullptr) : nullptr; - std::vector buffer(DEFAULT_BUFFER_SIZE); + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); TokenCallbackContext* ctx = nullptr; cactus_token_callback cb = nullptr; - - if (callback != nullptr) { - jclass callbackClass = env->GetObjectClass(callback); - jmethodID method = env->GetMethodID(callbackClass, "onToken", "(Ljava/lang/String;I)V"); - if (method != nullptr) { - ctx = new TokenCallbackContext{env, callback, method}; - cb = token_callback_bridge; - } + if (callback) { + JavaVM* jvm = nullptr; + env->GetJavaVM(&jvm); + jclass cls = env->GetObjectClass(callback); + jmethodID method = env->GetMethodID(cls, "onToken", "(Ljava/lang/String;I)V"); + ctx = new TokenCallbackContext{jvm, env->NewGlobalRef(callback), method}; + cb = token_callback_bridge; } - int result = cactus_complete( - reinterpret_cast(handle), - messages, - buffer.data(), - buffer.size(), - options, - tools, - cb, - ctx - ); - - delete ctx; - - release_jstring(env, messagesJson, messages); - release_jstring(env, optionsJson, options); - release_jstring(env, toolsJson, tools); - - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); - } - - return env->NewStringUTF(buffer.data()); -} - -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeTranscribe(JNIEnv* env, jobject, jlong handle, - jstring audioPath, jstring prompt, - jstring optionsJson, jbyteArray pcmData) { - if (handle == 0) { - return env->NewStringUTF("{\"error\":\"Model not initialized\"}"); - } - - const char* path = jstring_to_cstr(env, audioPath); - const char* promptStr = jstring_to_cstr(env, prompt); - const char* options = jstring_to_cstr(env, optionsJson); - - std::vector buffer(DEFAULT_BUFFER_SIZE); - - const uint8_t* pcmBuffer = nullptr; - size_t pcmSize = 0; jbyte* pcmBytes = nullptr; - - if (pcmData != nullptr) { + size_t pcmSize = 0; + if (pcmData) { pcmSize = env->GetArrayLength(pcmData); pcmBytes = env->GetByteArrayElements(pcmData, nullptr); - pcmBuffer = reinterpret_cast(pcmBytes); } - int result = cactus_transcribe( + int result = cactus_complete( reinterpret_cast(handle), - path, - promptStr, - buffer.data(), - buffer.size(), - options, - nullptr, - nullptr, - pcmBuffer, - pcmSize + messages, reinterpret_cast(buf), static_cast(bufSize), + options, tools, cb, ctx, + reinterpret_cast(pcmBytes), pcmSize ); - if (pcmBytes != nullptr) { - env->ReleaseByteArrayElements(pcmData, pcmBytes, JNI_ABORT); - } + if (ctx) { env->DeleteGlobalRef(ctx->callback); delete ctx; } + env->ReleaseByteArrayElements(responseBuffer, buf, 0); + if (pcmBytes) env->ReleaseByteArrayElements(pcmData, pcmBytes, JNI_ABORT); + env->ReleaseStringUTFChars(messagesJson, messages); + if (options) env->ReleaseStringUTFChars(optionsJson, options); + if (tools) env->ReleaseStringUTFChars(toolsJson, tools); - release_jstring(env, audioPath, path); - release_jstring(env, prompt, promptStr); - release_jstring(env, optionsJson, options); - - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); - } - - return env->NewStringUTF(buffer.data()); + return result; } -JNIEXPORT jfloatArray JNICALL -Java_com_cactus_Cactus_nativeEmbed(JNIEnv* env, jobject, jlong handle, - jstring text, jboolean normalize) { - if (handle == 0) { - return nullptr; - } - - const char* textStr = jstring_to_cstr(env, text); - - std::vector buffer(4096); - size_t embeddingDim = 0; - - int result = cactus_embed( - reinterpret_cast(handle), - textStr, - buffer.data(), - buffer.size(), - &embeddingDim, - normalize == JNI_TRUE - ); - - release_jstring(env, text, textStr); - - if (result < 0 || embeddingDim == 0) { - return nullptr; - } - - jfloatArray jarray = env->NewFloatArray(static_cast(embeddingDim)); - env->SetFloatArrayRegion(jarray, 0, static_cast(embeddingDim), buffer.data()); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativePrefill(JNIEnv* env, jobject, jlong handle, + jstring messagesJson, jbyteArray responseBuffer, + jstring optionsJson, jstring toolsJson, + jbyteArray pcmData) { + const char* messages = env->GetStringUTFChars(messagesJson, nullptr); + const char* options = optionsJson ? env->GetStringUTFChars(optionsJson, nullptr) : nullptr; + const char* tools = toolsJson ? env->GetStringUTFChars(toolsJson, nullptr) : nullptr; - return jarray; -} + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeRagQuery(JNIEnv* env, jobject, jlong handle, - jstring query, jint topK) { - if (handle == 0) { - return env->NewStringUTF("{\"error\":\"Model not initialized\"}"); + jbyte* pcmBytes = nullptr; + size_t pcmSize = 0; + if (pcmData) { + pcmSize = env->GetArrayLength(pcmData); + pcmBytes = env->GetByteArrayElements(pcmData, nullptr); } - const char* queryStr = jstring_to_cstr(env, query); - - std::vector buffer(DEFAULT_BUFFER_SIZE); - - int result = cactus_rag_query( + int result = cactus_prefill( reinterpret_cast(handle), - queryStr, - buffer.data(), - buffer.size(), - static_cast(topK) + messages, reinterpret_cast(buf), static_cast(bufSize), + options, tools, + reinterpret_cast(pcmBytes), pcmSize ); - release_jstring(env, query, queryStr); - - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); - } + env->ReleaseByteArrayElements(responseBuffer, buf, 0); + if (pcmBytes) env->ReleaseByteArrayElements(pcmData, pcmBytes, JNI_ABORT); + env->ReleaseStringUTFChars(messagesJson, messages); + if (options) env->ReleaseStringUTFChars(optionsJson, options); + if (tools) env->ReleaseStringUTFChars(toolsJson, tools); - return env->NewStringUTF(buffer.data()); -} - -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeGetLastError(JNIEnv* env, jobject) { - const char* error = cactus_get_last_error(); - return env->NewStringUTF(error ? error : ""); -} - -JNIEXPORT void JNICALL -Java_com_cactus_Cactus_nativeSetTelemetryToken(JNIEnv* env, jobject, jstring token) { - const char* tokenStr = jstring_to_cstr(env, token); - if (tokenStr != nullptr) { - cactus_set_telemetry_token(tokenStr); - release_jstring(env, token, tokenStr); - } -} - -JNIEXPORT void JNICALL -Java_com_cactus_Cactus_nativeSetProKey(JNIEnv* env, jobject, jstring key) { - const char* keyStr = jstring_to_cstr(env, key); - if (keyStr != nullptr) { - cactus_set_pro_key(keyStr); - release_jstring(env, key, keyStr); - } + return result; } -JNIEXPORT jintArray JNICALL -Java_com_cactus_Cactus_nativeTokenize(JNIEnv* env, jobject, jlong handle, jstring text) { - if (handle == 0) return nullptr; - - const char* textStr = jstring_to_cstr(env, text); - std::vector buffer(8192); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeTokenize(JNIEnv* env, jobject, jlong handle, + jstring text, jintArray tokenBuffer, jlongArray outTokenLen) { + const char* textStr = env->GetStringUTFChars(text, nullptr); + jsize bufLen = env->GetArrayLength(tokenBuffer); + jint* tokens = env->GetIntArrayElements(tokenBuffer, nullptr); size_t tokenLen = 0; int result = cactus_tokenize( reinterpret_cast(handle), textStr, - buffer.data(), - buffer.size(), + reinterpret_cast(tokens), + static_cast(bufLen), &tokenLen ); - release_jstring(env, text, textStr); + env->ReleaseIntArrayElements(tokenBuffer, tokens, 0); + env->ReleaseStringUTFChars(text, textStr); - if (result < 0 || tokenLen == 0) return nullptr; + jlong len = static_cast(tokenLen); + env->SetLongArrayRegion(outTokenLen, 0, 1, &len); - jintArray jarray = env->NewIntArray(static_cast(tokenLen)); - env->SetIntArrayRegion(jarray, 0, static_cast(tokenLen), reinterpret_cast(buffer.data())); - return jarray; + return result; } -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeScoreWindow(JNIEnv* env, jobject, jlong handle, - jintArray tokens, jint start, jint end, jint context) { - if (handle == 0) return env->NewStringUTF("{\"error\":\"Model not initialized\"}"); - +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeScoreWindow(JNIEnv* env, jobject, jlong handle, + jintArray tokens, jlong start, jlong end, + jlong context, jbyteArray responseBuffer) { jsize tokenLen = env->GetArrayLength(tokens); jint* tokenData = env->GetIntArrayElements(tokens, nullptr); - - std::vector buffer(DEFAULT_BUFFER_SIZE); + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); int result = cactus_score_window( reinterpret_cast(handle), @@ -299,213 +205,284 @@ Java_com_cactus_Cactus_nativeScoreWindow(JNIEnv* env, jobject, jlong handle, static_cast(start), static_cast(end), static_cast(context), - buffer.data(), - buffer.size() + reinterpret_cast(buf), + static_cast(bufSize) ); env->ReleaseIntArrayElements(tokens, tokenData, JNI_ABORT); + env->ReleaseByteArrayElements(responseBuffer, buf, 0); + + return result; +} - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeTranscribe(JNIEnv* env, jobject, jlong handle, + jstring audioPath, jstring prompt, + jbyteArray responseBuffer, + jstring optionsJson, jobject callback, + jbyteArray pcmData) { + const char* path = audioPath ? env->GetStringUTFChars(audioPath, nullptr) : nullptr; + const char* promptStr = env->GetStringUTFChars(prompt, nullptr); + const char* options = optionsJson ? env->GetStringUTFChars(optionsJson, nullptr) : nullptr; + + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); + + TokenCallbackContext* ctx = nullptr; + cactus_token_callback cb = nullptr; + if (callback) { + JavaVM* jvm = nullptr; + env->GetJavaVM(&jvm); + jclass cls = env->GetObjectClass(callback); + jmethodID method = env->GetMethodID(cls, "onToken", "(Ljava/lang/String;I)V"); + ctx = new TokenCallbackContext{jvm, env->NewGlobalRef(callback), method}; + cb = token_callback_bridge; + } + + jbyte* pcmBytes = nullptr; + size_t pcmSize = 0; + if (pcmData) { + pcmSize = env->GetArrayLength(pcmData); + pcmBytes = env->GetByteArrayElements(pcmData, nullptr); } - return env->NewStringUTF(buffer.data()); + int result = cactus_transcribe( + reinterpret_cast(handle), + path, promptStr, + reinterpret_cast(buf), static_cast(bufSize), + options, cb, ctx, + reinterpret_cast(pcmBytes), pcmSize + ); + + if (ctx) { env->DeleteGlobalRef(ctx->callback); delete ctx; } + env->ReleaseByteArrayElements(responseBuffer, buf, 0); + if (pcmBytes) env->ReleaseByteArrayElements(pcmData, pcmBytes, JNI_ABORT); + if (path) env->ReleaseStringUTFChars(audioPath, path); + env->ReleaseStringUTFChars(prompt, promptStr); + if (options) env->ReleaseStringUTFChars(optionsJson, options); + + return result; } JNIEXPORT jlong JNICALL -Java_com_cactus_Cactus_nativeStreamTranscribeInit(JNIEnv*, jobject, jlong handle) { - if (handle == 0) return 0; - return reinterpret_cast(cactus_stream_transcribe_init(reinterpret_cast(handle))); +Java_com_cactus_CactusJNI_nativeStreamTranscribeStart(JNIEnv* env, jobject, jlong handle, + jstring optionsJson) { + const char* options = optionsJson ? env->GetStringUTFChars(optionsJson, nullptr) : nullptr; + cactus_stream_transcribe_t stream = cactus_stream_transcribe_start( + reinterpret_cast(handle), options); + if (options) env->ReleaseStringUTFChars(optionsJson, options); + return reinterpret_cast(stream); } JNIEXPORT jint JNICALL -Java_com_cactus_Cactus_nativeStreamTranscribeInsert(JNIEnv* env, jobject, jlong streamHandle, jbyteArray pcmData) { - if (streamHandle == 0) return -1; +Java_com_cactus_CactusJNI_nativeStreamTranscribeProcess(JNIEnv* env, jobject, jlong stream, + jbyteArray pcmData, + jbyteArray responseBuffer) { + jbyte* pcmBytes = pcmData ? env->GetByteArrayElements(pcmData, nullptr) : nullptr; + size_t pcmSize = pcmData ? static_cast(env->GetArrayLength(pcmData)) : 0; - jsize pcmSize = env->GetArrayLength(pcmData); - jbyte* pcmBytes = env->GetByteArrayElements(pcmData, nullptr); + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); - int result = cactus_stream_transcribe_insert( - reinterpret_cast(streamHandle), - reinterpret_cast(pcmBytes), - static_cast(pcmSize) + int result = cactus_stream_transcribe_process( + reinterpret_cast(stream), + reinterpret_cast(pcmBytes), pcmSize, + reinterpret_cast(buf), static_cast(bufSize) ); - env->ReleaseByteArrayElements(pcmData, pcmBytes, JNI_ABORT); + env->ReleaseByteArrayElements(responseBuffer, buf, 0); + if (pcmBytes) env->ReleaseByteArrayElements(pcmData, pcmBytes, JNI_ABORT); + return result; } -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeStreamTranscribeProcess(JNIEnv* env, jobject, jlong streamHandle, jstring optionsJson) { - if (streamHandle == 0) return env->NewStringUTF("{\"error\":\"Stream not initialized\"}"); - - const char* options = jstring_to_cstr(env, optionsJson); - std::vector buffer(DEFAULT_BUFFER_SIZE); - - int result = cactus_stream_transcribe_process( - reinterpret_cast(streamHandle), - buffer.data(), - buffer.size(), - options +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeStreamTranscribeStop(JNIEnv* env, jobject, jlong stream, + jbyteArray responseBuffer) { + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); + + int result = cactus_stream_transcribe_stop( + reinterpret_cast(stream), + reinterpret_cast(buf), static_cast(bufSize) ); - release_jstring(env, optionsJson, options); + env->ReleaseByteArrayElements(responseBuffer, buf, 0); - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); - } - - return env->NewStringUTF(buffer.data()); + return result; } -JNIEXPORT jstring JNICALL -Java_com_cactus_Cactus_nativeStreamTranscribeFinalize(JNIEnv* env, jobject, jlong streamHandle) { - if (streamHandle == 0) return env->NewStringUTF("{\"error\":\"Stream not initialized\"}"); - - std::vector buffer(DEFAULT_BUFFER_SIZE); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeEmbed(JNIEnv* env, jobject, jlong handle, + jstring text, jfloatArray embeddingsBuffer, + jlongArray outEmbeddingDim, jboolean normalize) { + const char* textStr = env->GetStringUTFChars(text, nullptr); + jsize bufSize = env->GetArrayLength(embeddingsBuffer); + jfloat* buf = env->GetFloatArrayElements(embeddingsBuffer, nullptr); + size_t embeddingDim = 0; - int result = cactus_stream_transcribe_finalize( - reinterpret_cast(streamHandle), - buffer.data(), - buffer.size() + int result = cactus_embed( + reinterpret_cast(handle), + textStr, buf, static_cast(bufSize), + &embeddingDim, normalize == JNI_TRUE ); - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); - } + env->ReleaseFloatArrayElements(embeddingsBuffer, buf, 0); + env->ReleaseStringUTFChars(text, textStr); - return env->NewStringUTF(buffer.data()); -} + jlong dim = static_cast(embeddingDim); + env->SetLongArrayRegion(outEmbeddingDim, 0, 1, &dim); -JNIEXPORT void JNICALL -Java_com_cactus_Cactus_nativeStreamTranscribeDestroy(JNIEnv*, jobject, jlong streamHandle) { - if (streamHandle != 0) { - cactus_stream_transcribe_destroy(reinterpret_cast(streamHandle)); - } + return result; } -JNIEXPORT jfloatArray JNICALL -Java_com_cactus_Cactus_nativeImageEmbed(JNIEnv* env, jobject, jlong handle, jstring imagePath) { - if (handle == 0) return nullptr; - - const char* path = jstring_to_cstr(env, imagePath); - std::vector buffer(4096); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeImageEmbed(JNIEnv* env, jobject, jlong handle, + jstring imagePath, jfloatArray embeddingsBuffer, + jlongArray outEmbeddingDim) { + const char* path = env->GetStringUTFChars(imagePath, nullptr); + jsize bufSize = env->GetArrayLength(embeddingsBuffer); + jfloat* buf = env->GetFloatArrayElements(embeddingsBuffer, nullptr); size_t embeddingDim = 0; int result = cactus_image_embed( reinterpret_cast(handle), - path, - buffer.data(), - buffer.size(), - &embeddingDim + path, buf, static_cast(bufSize), &embeddingDim ); - release_jstring(env, imagePath, path); + env->ReleaseFloatArrayElements(embeddingsBuffer, buf, 0); + env->ReleaseStringUTFChars(imagePath, path); - if (result < 0 || embeddingDim == 0) return nullptr; + jlong dim = static_cast(embeddingDim); + env->SetLongArrayRegion(outEmbeddingDim, 0, 1, &dim); - jfloatArray jarray = env->NewFloatArray(static_cast(embeddingDim)); - env->SetFloatArrayRegion(jarray, 0, static_cast(embeddingDim), buffer.data()); - return jarray; + return result; } -JNIEXPORT jfloatArray JNICALL -Java_com_cactus_Cactus_nativeAudioEmbed(JNIEnv* env, jobject, jlong handle, jstring audioPath) { - if (handle == 0) return nullptr; - - const char* path = jstring_to_cstr(env, audioPath); - std::vector buffer(4096); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeAudioEmbed(JNIEnv* env, jobject, jlong handle, + jstring audioPath, jfloatArray embeddingsBuffer, + jlongArray outEmbeddingDim) { + const char* path = env->GetStringUTFChars(audioPath, nullptr); + jsize bufSize = env->GetArrayLength(embeddingsBuffer); + jfloat* buf = env->GetFloatArrayElements(embeddingsBuffer, nullptr); size_t embeddingDim = 0; int result = cactus_audio_embed( reinterpret_cast(handle), - path, - buffer.data(), - buffer.size(), - &embeddingDim + path, buf, static_cast(bufSize), &embeddingDim ); - release_jstring(env, audioPath, path); + env->ReleaseFloatArrayElements(embeddingsBuffer, buf, 0); + env->ReleaseStringUTFChars(audioPath, path); - if (result < 0 || embeddingDim == 0) return nullptr; + jlong dim = static_cast(embeddingDim); + env->SetLongArrayRegion(outEmbeddingDim, 0, 1, &dim); - jfloatArray jarray = env->NewFloatArray(static_cast(embeddingDim)); - env->SetFloatArrayRegion(jarray, 0, static_cast(embeddingDim), buffer.data()); - return jarray; + return result; +} + +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeRagQuery(JNIEnv* env, jobject, jlong handle, + jstring query, jbyteArray responseBuffer, jlong topK) { + const char* queryStr = env->GetStringUTFChars(query, nullptr); + jsize bufSize = env->GetArrayLength(responseBuffer); + jbyte* buf = env->GetByteArrayElements(responseBuffer, nullptr); + + int result = cactus_rag_query( + reinterpret_cast(handle), + queryStr, reinterpret_cast(buf), static_cast(bufSize), + static_cast(topK) + ); + + env->ReleaseByteArrayElements(responseBuffer, buf, 0); + env->ReleaseStringUTFChars(query, queryStr); + + return result; } JNIEXPORT jlong JNICALL -Java_com_cactus_CactusIndex_nativeIndexInit(JNIEnv* env, jobject, jstring indexDir, jint embeddingDim) { - const char* dir = jstring_to_cstr(env, indexDir); +Java_com_cactus_CactusJNI_nativeIndexInit(JNIEnv* env, jobject, jstring indexDir, jlong embeddingDim) { + const char* dir = env->GetStringUTFChars(indexDir, nullptr); jlong handle = reinterpret_cast(cactus_index_init(dir, static_cast(embeddingDim))); - release_jstring(env, indexDir, dir); + env->ReleaseStringUTFChars(indexDir, dir); return handle; } JNIEXPORT jint JNICALL -Java_com_cactus_CactusIndex_nativeIndexAdd(JNIEnv* env, jobject, jlong handle, - jintArray ids, jobjectArray documents, - jobjectArray metadatas, jobjectArray embeddings, - jint embeddingDim) { - if (handle == 0) return -1; - +Java_com_cactus_CactusJNI_nativeIndexAdd(JNIEnv* env, jobject, jlong handle, + jintArray ids, jobjectArray documents, + jobjectArray metadatas, jobjectArray embeddings, + jlong embeddingDim) { jsize count = env->GetArrayLength(ids); jint* idData = env->GetIntArrayElements(ids, nullptr); - std::vector docPtrs(count); - std::vector metaPtrs(count); - std::vector embPtrs(count); - std::vector docStrings(count); - std::vector metaStrings(count); - std::vector embArrays(count); + const char** docPtrs = new const char*[count]; + const char** metaPtrs = new const char*[count]; + const float** embPtrs = new const float*[count]; + jstring* docStrings = new jstring[count]; + jstring* metaStrings = new jstring[count]; + jfloatArray* embArrays = new jfloatArray[count]; + jsize acquired = 0; for (jsize i = 0; i < count; i++) { docStrings[i] = static_cast(env->GetObjectArrayElement(documents, i)); - docPtrs[i] = jstring_to_cstr(env, docStrings[i]); + if (!docStrings[i]) break; + docPtrs[i] = env->GetStringUTFChars(docStrings[i], nullptr); - if (metadatas != nullptr) { + if (metadatas) { metaStrings[i] = static_cast(env->GetObjectArrayElement(metadatas, i)); - metaPtrs[i] = jstring_to_cstr(env, metaStrings[i]); + if (!metaStrings[i]) { env->ReleaseStringUTFChars(docStrings[i], docPtrs[i]); env->DeleteLocalRef(docStrings[i]); break; } + metaPtrs[i] = env->GetStringUTFChars(metaStrings[i], nullptr); } else { + metaStrings[i] = nullptr; metaPtrs[i] = nullptr; } embArrays[i] = static_cast(env->GetObjectArrayElement(embeddings, i)); + if (!embArrays[i]) { + env->ReleaseStringUTFChars(docStrings[i], docPtrs[i]); env->DeleteLocalRef(docStrings[i]); + if (metaStrings[i]) { env->ReleaseStringUTFChars(metaStrings[i], metaPtrs[i]); env->DeleteLocalRef(metaStrings[i]); } + break; + } embPtrs[i] = env->GetFloatArrayElements(embArrays[i], nullptr); + acquired = i + 1; } - int result = cactus_index_add( - reinterpret_cast(handle), - reinterpret_cast(idData), - docPtrs.data(), - metadatas != nullptr ? metaPtrs.data() : nullptr, - embPtrs.data(), - static_cast(count), - static_cast(embeddingDim) - ); + int result; + if (acquired < count) { + if (env->ExceptionCheck()) env->ExceptionClear(); + result = -1; + } else { + result = cactus_index_add( + reinterpret_cast(handle), + reinterpret_cast(idData), + docPtrs, metadatas ? metaPtrs : nullptr, embPtrs, + static_cast(count), static_cast(embeddingDim) + ); + } - for (jsize i = 0; i < count; i++) { - release_jstring(env, docStrings[i], docPtrs[i]); - if (metadatas != nullptr) { - release_jstring(env, metaStrings[i], metaPtrs[i]); - } + for (jsize i = 0; i < acquired; i++) { + env->ReleaseStringUTFChars(docStrings[i], docPtrs[i]); + env->DeleteLocalRef(docStrings[i]); + if (metaStrings[i]) { env->ReleaseStringUTFChars(metaStrings[i], metaPtrs[i]); env->DeleteLocalRef(metaStrings[i]); } env->ReleaseFloatArrayElements(embArrays[i], const_cast(embPtrs[i]), JNI_ABORT); + env->DeleteLocalRef(embArrays[i]); } env->ReleaseIntArrayElements(ids, idData, JNI_ABORT); + delete[] docPtrs; + delete[] metaPtrs; + delete[] embPtrs; + delete[] docStrings; + delete[] metaStrings; + delete[] embArrays; + return result; } JNIEXPORT jint JNICALL -Java_com_cactus_CactusIndex_nativeIndexDelete(JNIEnv* env, jobject, jlong handle, jintArray ids) { - if (handle == 0) return -1; - +Java_com_cactus_CactusJNI_nativeIndexDelete(JNIEnv* env, jobject, jlong handle, jintArray ids) { jsize count = env->GetArrayLength(ids); jint* idData = env->GetIntArrayElements(ids, nullptr); @@ -519,66 +496,261 @@ Java_com_cactus_CactusIndex_nativeIndexDelete(JNIEnv* env, jobject, jlong handle return result; } -JNIEXPORT jstring JNICALL -Java_com_cactus_CactusIndex_nativeIndexQuery(JNIEnv* env, jobject, jlong handle, - jfloatArray embedding, jint topK, jstring optionsJson) { - if (handle == 0) return env->NewStringUTF("{\"error\":\"Index not initialized\"}"); +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeIndexGet(JNIEnv* env, jobject, jlong handle, + jintArray ids, + jobjectArray documentBuffers, jlongArray documentBufferSizes, + jobjectArray metadataBuffers, jlongArray metadataBufferSizes, + jobjectArray embeddingBuffers, jlongArray embeddingBufferSizes) { + jsize count = env->GetArrayLength(ids); + jint* idData = env->GetIntArrayElements(ids, nullptr); - jsize embDim = env->GetArrayLength(embedding); - jfloat* embData = env->GetFloatArrayElements(embedding, nullptr); - const char* options = jstring_to_cstr(env, optionsJson); + char** docBufs = new char*[count]; + size_t* docSizes = new size_t[count]; + char** metaBufs = new char*[count]; + size_t* metaSizes = new size_t[count]; + float** embBufs = new float*[count]; + size_t* embSizes = new size_t[count]; - std::vector idBuffer(topK); - std::vector scoreBuffer(topK); - size_t idBufferSize = topK; - size_t scoreBufferSize = topK; + jbyteArray* docArrays = new jbyteArray[count]; + jbyteArray* metaArrays = new jbyteArray[count]; + jfloatArray* embArrays = new jfloatArray[count]; - const float* embPtr = embData; - int* idPtr = idBuffer.data(); - float* scorePtr = scoreBuffer.data(); + jlong* docSizesIn = env->GetLongArrayElements(documentBufferSizes, nullptr); + jlong* metaSizesIn = env->GetLongArrayElements(metadataBufferSizes, nullptr); + jlong* embSizesIn = env->GetLongArrayElements(embeddingBufferSizes, nullptr); - int result = cactus_index_query( - reinterpret_cast(handle), - &embPtr, - 1, - static_cast(embDim), - options, - &idPtr, - &idBufferSize, - &scorePtr, - &scoreBufferSize - ); + jsize acquired = 0; + for (jsize i = 0; i < count; i++) { + docArrays[i] = static_cast(env->GetObjectArrayElement(documentBuffers, i)); + if (!docArrays[i]) break; + docBufs[i] = reinterpret_cast(env->GetByteArrayElements(docArrays[i], nullptr)); + docSizes[i] = static_cast(docSizesIn[i]); + + metaArrays[i] = static_cast(env->GetObjectArrayElement(metadataBuffers, i)); + if (!metaArrays[i]) { + env->ReleaseByteArrayElements(docArrays[i], reinterpret_cast(docBufs[i]), JNI_ABORT); env->DeleteLocalRef(docArrays[i]); + break; + } + metaBufs[i] = reinterpret_cast(env->GetByteArrayElements(metaArrays[i], nullptr)); + metaSizes[i] = static_cast(metaSizesIn[i]); + + embArrays[i] = static_cast(env->GetObjectArrayElement(embeddingBuffers, i)); + if (!embArrays[i]) { + env->ReleaseByteArrayElements(docArrays[i], reinterpret_cast(docBufs[i]), JNI_ABORT); env->DeleteLocalRef(docArrays[i]); + env->ReleaseByteArrayElements(metaArrays[i], reinterpret_cast(metaBufs[i]), JNI_ABORT); env->DeleteLocalRef(metaArrays[i]); + break; + } + embBufs[i] = env->GetFloatArrayElements(embArrays[i], nullptr); + embSizes[i] = static_cast(embSizesIn[i]); + acquired = i + 1; + } + + int result; + if (acquired < count) { + if (env->ExceptionCheck()) env->ExceptionClear(); + result = -1; + } else { + result = cactus_index_get( + reinterpret_cast(handle), + reinterpret_cast(idData), + static_cast(count), + docBufs, docSizes, metaBufs, metaSizes, embBufs, embSizes + ); + } + + for (jsize i = 0; i < acquired; i++) { + env->ReleaseByteArrayElements(docArrays[i], reinterpret_cast(docBufs[i]), 0); + env->DeleteLocalRef(docArrays[i]); + env->ReleaseByteArrayElements(metaArrays[i], reinterpret_cast(metaBufs[i]), 0); + env->DeleteLocalRef(metaArrays[i]); + env->ReleaseFloatArrayElements(embArrays[i], embBufs[i], 0); + env->DeleteLocalRef(embArrays[i]); + docSizesIn[i] = static_cast(docSizes[i]); + metaSizesIn[i] = static_cast(metaSizes[i]); + embSizesIn[i] = static_cast(embSizes[i]); + } - env->ReleaseFloatArrayElements(embedding, embData, JNI_ABORT); - release_jstring(env, optionsJson, options); + env->ReleaseLongArrayElements(documentBufferSizes, docSizesIn, 0); + env->ReleaseLongArrayElements(metadataBufferSizes, metaSizesIn, 0); + env->ReleaseLongArrayElements(embeddingBufferSizes, embSizesIn, 0); + env->ReleaseIntArrayElements(ids, idData, JNI_ABORT); + + delete[] docBufs; + delete[] docSizes; + delete[] metaBufs; + delete[] metaSizes; + delete[] embBufs; + delete[] embSizes; + delete[] docArrays; + delete[] metaArrays; + delete[] embArrays; + + return result; +} + +JNIEXPORT jint JNICALL +Java_com_cactus_CactusJNI_nativeIndexQuery(JNIEnv* env, jobject, jlong handle, + jobjectArray embeddings, jlong embeddingDim, + jstring optionsJson, + jobjectArray idBuffers, jlongArray idBufferSizes, + jobjectArray scoreBuffers, jlongArray scoreBufferSizes) { + jsize embCount = env->GetArrayLength(embeddings); + const char* options = optionsJson ? env->GetStringUTFChars(optionsJson, nullptr) : nullptr; + + const float** embPtrs = new const float*[embCount]; + jfloatArray* embArrays = new jfloatArray[embCount]; + + int** idPtrs = new int*[embCount]; + size_t* idSizes = new size_t[embCount]; + float** scorePtrs = new float*[embCount]; + size_t* scoreSizes = new size_t[embCount]; + + jintArray* idJArrays = new jintArray[embCount]; + jfloatArray* scoreJArrays = new jfloatArray[embCount]; + + jlong* idSizesIn = env->GetLongArrayElements(idBufferSizes, nullptr); + jlong* scoreSizesIn = env->GetLongArrayElements(scoreBufferSizes, nullptr); + + jsize acquired = 0; + for (jsize i = 0; i < embCount; i++) { + embArrays[i] = static_cast(env->GetObjectArrayElement(embeddings, i)); + if (!embArrays[i]) break; + embPtrs[i] = env->GetFloatArrayElements(embArrays[i], nullptr); - if (result < 0) { - const char* error = cactus_get_last_error(); - std::string errorJson = "{\"error\":\"" + std::string(error ? error : "Unknown error") + "\"}"; - return env->NewStringUTF(errorJson.c_str()); + idJArrays[i] = static_cast(env->GetObjectArrayElement(idBuffers, i)); + if (!idJArrays[i]) { + env->ReleaseFloatArrayElements(embArrays[i], const_cast(embPtrs[i]), JNI_ABORT); env->DeleteLocalRef(embArrays[i]); + break; + } + idPtrs[i] = env->GetIntArrayElements(idJArrays[i], nullptr); + idSizes[i] = static_cast(idSizesIn[i]); + + scoreJArrays[i] = static_cast(env->GetObjectArrayElement(scoreBuffers, i)); + if (!scoreJArrays[i]) { + env->ReleaseFloatArrayElements(embArrays[i], const_cast(embPtrs[i]), JNI_ABORT); env->DeleteLocalRef(embArrays[i]); + env->ReleaseIntArrayElements(idJArrays[i], idPtrs[i], JNI_ABORT); env->DeleteLocalRef(idJArrays[i]); + break; + } + scorePtrs[i] = env->GetFloatArrayElements(scoreJArrays[i], nullptr); + scoreSizes[i] = static_cast(scoreSizesIn[i]); + acquired = i + 1; } - std::string json = "{\"results\":["; - for (size_t i = 0; i < idBufferSize; i++) { - if (i > 0) json += ","; - json += "{\"id\":" + std::to_string(idBuffer[i]) + ",\"score\":" + std::to_string(scoreBuffer[i]) + "}"; + int result; + if (acquired < embCount) { + if (env->ExceptionCheck()) env->ExceptionClear(); + result = -1; + } else { + result = cactus_index_query( + reinterpret_cast(handle), + embPtrs, static_cast(embCount), static_cast(embeddingDim), + options, idPtrs, idSizes, scorePtrs, scoreSizes + ); + } + + for (jsize i = 0; i < acquired; i++) { + env->ReleaseFloatArrayElements(embArrays[i], const_cast(embPtrs[i]), JNI_ABORT); + env->DeleteLocalRef(embArrays[i]); + env->ReleaseIntArrayElements(idJArrays[i], idPtrs[i], 0); + env->DeleteLocalRef(idJArrays[i]); + env->ReleaseFloatArrayElements(scoreJArrays[i], scorePtrs[i], 0); + env->DeleteLocalRef(scoreJArrays[i]); + idSizesIn[i] = static_cast(idSizes[i]); + scoreSizesIn[i] = static_cast(scoreSizes[i]); } - json += "]}"; - return env->NewStringUTF(json.c_str()); + env->ReleaseLongArrayElements(idBufferSizes, idSizesIn, 0); + env->ReleaseLongArrayElements(scoreBufferSizes, scoreSizesIn, 0); + if (options) env->ReleaseStringUTFChars(optionsJson, options); + + delete[] embPtrs; + delete[] embArrays; + delete[] idPtrs; + delete[] idSizes; + delete[] scorePtrs; + delete[] scoreSizes; + delete[] idJArrays; + delete[] scoreJArrays; + + return result; } JNIEXPORT jint JNICALL -Java_com_cactus_CactusIndex_nativeIndexCompact(JNIEnv*, jobject, jlong handle) { - if (handle == 0) return -1; +Java_com_cactus_CactusJNI_nativeIndexCompact(JNIEnv*, jobject, jlong handle) { return cactus_index_compact(reinterpret_cast(handle)); } JNIEXPORT void JNICALL -Java_com_cactus_CactusIndex_nativeIndexDestroy(JNIEnv*, jobject, jlong handle) { - if (handle != 0) { - cactus_index_destroy(reinterpret_cast(handle)); +Java_com_cactus_CactusJNI_nativeIndexDestroy(JNIEnv*, jobject, jlong handle) { + cactus_index_destroy(reinterpret_cast(handle)); +} + +JNIEXPORT jstring JNICALL +Java_com_cactus_CactusJNI_nativeGetLastError(JNIEnv* env, jobject) { + return env->NewStringUTF(cactus_get_last_error()); +} + +JNIEXPORT void JNICALL +Java_com_cactus_CactusJNI_nativeLogSetLevel(JNIEnv*, jobject, jint level) { + cactus_log_set_level(static_cast(level)); +} + +JNIEXPORT void JNICALL +Java_com_cactus_CactusJNI_nativeLogSetCallback(JNIEnv* env, jobject, jobject callback) { + if (g_log_callback_ctx) { + env->DeleteGlobalRef(g_log_callback_ctx->callback); + delete g_log_callback_ctx; + g_log_callback_ctx = nullptr; + } + if (!callback) { + cactus_log_set_callback(nullptr, nullptr); + return; + } + JavaVM* jvm = nullptr; + env->GetJavaVM(&jvm); + jclass cls = env->GetObjectClass(callback); + jmethodID method = env->GetMethodID(cls, "onLog", "(ILjava/lang/String;Ljava/lang/String;)V"); + if (!method) { + env->ExceptionClear(); + cactus_log_set_callback(nullptr, nullptr); + return; } + g_log_callback_ctx = new LogCallbackContext{jvm, env->NewGlobalRef(callback), method}; + cactus_log_set_callback(log_callback_bridge, g_log_callback_ctx); +} + +JNIEXPORT void JNICALL +Java_com_cactus_CactusJNI_nativeSetTelemetryEnvironment(JNIEnv* env, jobject, + jstring framework, jstring cacheLocation, jstring version) { + const char* fw = framework ? env->GetStringUTFChars(framework, nullptr) : nullptr; + const char* cache = cacheLocation ? env->GetStringUTFChars(cacheLocation, nullptr) : nullptr; + const char* ver = version ? env->GetStringUTFChars(version, nullptr) : nullptr; + cactus_set_telemetry_environment(fw, cache, ver); + if (fw) env->ReleaseStringUTFChars(framework, fw); + if (cache) env->ReleaseStringUTFChars(cacheLocation, cache); + if (ver) env->ReleaseStringUTFChars(version, ver); +} + +JNIEXPORT void JNICALL +Java_com_cactus_CactusJNI_nativeSetAppId(JNIEnv* env, jobject, jstring appId) { + if (!appId) { + cactus_set_app_id(nullptr); + return; + } + const char* id = env->GetStringUTFChars(appId, nullptr); + cactus_set_app_id(id); + env->ReleaseStringUTFChars(appId, id); +} + +JNIEXPORT void JNICALL +Java_com_cactus_CactusJNI_nativeTelemetryFlush(JNIEnv*, jobject) { + cactus_telemetry_flush(); +} + +JNIEXPORT void JNICALL +Java_com_cactus_CactusJNI_nativeTelemetryShutdown(JNIEnv*, jobject) { + cactus_telemetry_shutdown(); } } diff --git a/android/mbedtls/arm64-v8a/lib/libmbedcrypto.a b/android/mbedtls/arm64-v8a/lib/libmbedcrypto.a new file mode 100644 index 000000000..b6d94095d Binary files /dev/null and b/android/mbedtls/arm64-v8a/lib/libmbedcrypto.a differ diff --git a/android/mbedtls/arm64-v8a/lib/libmbedtls.a b/android/mbedtls/arm64-v8a/lib/libmbedtls.a new file mode 100644 index 000000000..817a57a1f Binary files /dev/null and b/android/mbedtls/arm64-v8a/lib/libmbedtls.a differ diff --git a/android/mbedtls/arm64-v8a/lib/libmbedx509.a b/android/mbedtls/arm64-v8a/lib/libmbedx509.a new file mode 100644 index 000000000..1f55473a1 Binary files /dev/null and b/android/mbedtls/arm64-v8a/lib/libmbedx509.a differ diff --git a/apple/CMakeLists.txt b/apple/CMakeLists.txt new file mode 100644 index 000000000..d2795bdd4 --- /dev/null +++ b/apple/CMakeLists.txt @@ -0,0 +1,67 @@ +cmake_minimum_required(VERSION 3.10) +project(CactusApple LANGUAGES CXX) + +set(CMAKE_OSX_DEPLOYMENT_TARGET "13.3" CACHE STRING "") +set(CMAKE_XCODE_ATTRIBUTE_IPHONEOS_DEPLOYMENT_TARGET "16.4" CACHE STRING "") +set(CMAKE_BUILD_TYPE Release CACHE STRING "") +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib) + +if(NOT TARGET cactus_engine) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../cactus-engine ${CMAKE_CURRENT_BINARY_DIR}/cactus-engine) +endif() + +set(_DEPS cactus_engine cactus_graph cactus_kernels) + +if(BUILD_SHARED_LIBS) + file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/stub.cpp "") + add_library(cactus_framework SHARED ${CMAKE_CURRENT_BINARY_DIR}/stub.cpp) + set_target_properties(cactus_framework PROPERTIES OUTPUT_NAME cactus) + add_dependencies(cactus_framework ${_DEPS}) + + target_link_libraries(cactus_framework PRIVATE cactus_flags) + + target_link_options(cactus_framework PRIVATE + -force_load $ + -force_load $ + -force_load $ + ) + + target_link_libraries(cactus_framework PUBLIC + "-framework Accelerate" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" + ) + + if(CMAKE_SYSTEM_NAME STREQUAL "iOS") + string(TOLOWER "${CMAKE_OSX_SYSROOT}" _sysroot_lower) + if(_sysroot_lower MATCHES "simulator") + set(_CURL "${CACTUS_CURL_ROOT}/ios/simulator/libcurl.a") + else() + set(_CURL "${CACTUS_CURL_ROOT}/ios/device/libcurl.a") + endif() + else() + set(_CURL "${CACTUS_CURL_ROOT}/macos/libcurl.a") + endif() + if(EXISTS "${_CURL}") + target_link_libraries(cactus_framework PUBLIC "${_CURL}") + endif() + + set_target_properties(cactus_framework PROPERTIES + FRAMEWORK TRUE + FRAMEWORK_VERSION A + MACOSX_FRAMEWORK_IDENTIFIER com.cactuscompute.cactus + VERSION 1.0.0 + SOVERSION 1.0.0 + MACOSX_FRAMEWORK_INFO_PLIST "${CMAKE_CURRENT_SOURCE_DIR}/Info.plist" + ) + + add_custom_command(TARGET cactus_framework POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory "$/Headers" + COMMAND ${CMAKE_COMMAND} -E copy + "${CMAKE_CURRENT_SOURCE_DIR}/../cactus-engine/cactus_engine.h" + "$/Headers" + ) + + if(CMAKE_SYSTEM_NAME STREQUAL "iOS") + target_compile_definitions(cactus_framework PRIVATE TARGET_OS_IPHONE=1) + endif() +endif() diff --git a/apple/Cactus.swift b/apple/Cactus.swift deleted file mode 100644 index 852a7dd80..000000000 --- a/apple/Cactus.swift +++ /dev/null @@ -1,868 +0,0 @@ -import Foundation -import cactus - -public final class Cactus: @unchecked Sendable { - - - public struct CompletionResult { - public let text: String - public let functionCalls: [[String: Any]]? - public let promptTokens: Int - public let completionTokens: Int - public let timeToFirstToken: Double - public let totalTime: Double - public let prefillTokensPerSecond: Double - public let decodeTokensPerSecond: Double - public let confidence: Double - public let needsCloudHandoff: Bool - - init(json: [String: Any]) { - self.text = json["text"] as? String ?? "" - self.functionCalls = json["function_calls"] as? [[String: Any]] - self.promptTokens = json["prompt_tokens"] as? Int ?? 0 - self.completionTokens = json["completion_tokens"] as? Int ?? 0 - self.timeToFirstToken = json["time_to_first_token_ms"] as? Double ?? 0 - self.totalTime = json["total_time_ms"] as? Double ?? 0 - self.prefillTokensPerSecond = json["prefill_tokens_per_second"] as? Double ?? 0 - self.decodeTokensPerSecond = json["decode_tokens_per_second"] as? Double ?? 0 - self.confidence = json["confidence"] as? Double ?? 1.0 - self.needsCloudHandoff = json["cloud_handoff"] as? Bool ?? false - } - } - - public struct TranscriptionResult { - public let text: String - public let segments: [[String: Any]]? - public let totalTime: Double - - init(json: [String: Any]) { - self.text = json["text"] as? String ?? "" - self.segments = json["segments"] as? [[String: Any]] - self.totalTime = json["total_time_ms"] as? Double ?? 0 - } - } - - public struct Message { - public let role: String - public let content: String - - public init(role: String, content: String) { - self.role = role - self.content = content - } - - public static func system(_ content: String) -> Message { - Message(role: "system", content: content) - } - - public static func user(_ content: String) -> Message { - Message(role: "user", content: content) - } - - public static func assistant(_ content: String) -> Message { - Message(role: "assistant", content: content) - } - - func toDict() -> [String: String] { - ["role": role, "content": content] - } - } - - public struct CompletionOptions { - public var temperature: Float - public var topP: Float - public var topK: Int - public var maxTokens: Int - public var stopSequences: [String] - public var confidenceThreshold: Float - - public init( - temperature: Float = 0.7, - topP: Float = 0.9, - topK: Int = 40, - maxTokens: Int = 512, - stopSequences: [String] = [], - confidenceThreshold: Float = 0.0 - ) { - self.temperature = temperature - self.topP = topP - self.topK = topK - self.maxTokens = maxTokens - self.stopSequences = stopSequences - self.confidenceThreshold = confidenceThreshold - } - - public static let `default` = CompletionOptions() - - func toJSON() -> String? { - let dict: [String: Any] = [ - "temperature": temperature, - "top_p": topP, - "top_k": topK, - "max_tokens": maxTokens, - "stop": stopSequences, - "confidence_threshold": confidenceThreshold - ] - guard let data = try? JSONSerialization.data(withJSONObject: dict), - let json = String(data: data, encoding: .utf8) else { - return nil - } - return json - } - } - - public struct TranscriptionOptions { - public var language: String? - public var translateToEnglish: Bool - - public init(language: String? = nil, translateToEnglish: Bool = false) { - self.language = language - self.translateToEnglish = translateToEnglish - } - - public static let `default` = TranscriptionOptions() - - func toJSON() -> String? { - var dict: [String: Any] = [ - "translate": translateToEnglish - ] - if let lang = language { - dict["language"] = lang - } - guard let data = try? JSONSerialization.data(withJSONObject: dict), - let json = String(data: data, encoding: .utf8) else { - return nil - } - return json - } - } - - public enum CactusError: Error, LocalizedError { - case initializationFailed(String) - case completionFailed(String) - case transcriptionFailed(String) - case embeddingFailed(String) - case invalidResponse - - public var errorDescription: String? { - switch self { - case .initializationFailed(let msg): return "Initialization failed: \(msg)" - case .completionFailed(let msg): return "Completion failed: \(msg)" - case .transcriptionFailed(let msg): return "Transcription failed: \(msg)" - case .embeddingFailed(let msg): return "Embedding failed: \(msg)" - case .invalidResponse: return "Invalid response from model" - } - } - } - - - private let handle: OpaquePointer - private static let defaultBufferSize = 65536 - - public init(modelPath: String, corpusDir: String? = nil) throws { - guard let h = cactus_init(modelPath, corpusDir) else { - let error = String(cString: cactus_get_last_error()) - throw CactusError.initializationFailed(error.isEmpty ? "Unknown error" : error) - } - self.handle = h - } - - deinit { - cactus_destroy(handle) - } - - public func complete( - messages: [Message], - options: CompletionOptions = .default, - tools: [[String: Any]]? = nil, - onToken: ((String, UInt32) -> Void)? = nil - ) throws -> CompletionResult { - let messagesJSON = try serializeMessages(messages) - let optionsJSON = options.toJSON() - let toolsJSON = try serializeTools(tools) - - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - - let callbackContext = onToken.map { TokenCallbackContext(callback: $0) } - let contextPtr = callbackContext.map { Unmanaged.passUnretained($0).toOpaque() } - - let result = buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_complete( - handle, - messagesJSON, - bufferPtr.baseAddress, - bufferPtr.count, - optionsJSON, - toolsJSON, - onToken != nil ? tokenCallbackBridge : nil, - contextPtr - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.completionFailed(error.isEmpty ? "Unknown error" : error) - } - - let responseString = String(cString: buffer) - guard let json = parseJSON(responseString) else { - throw CactusError.invalidResponse - } - - if let errorMsg = json["error"] as? String { - throw CactusError.completionFailed(errorMsg) - } - - return CompletionResult(json: json) - } - - public func complete( - _ prompt: String, - options: CompletionOptions = .default - ) throws -> CompletionResult { - try complete(messages: [.user(prompt)], options: options) - } - - public func transcribe( - audioPath: String, - prompt: String? = nil, - options: TranscriptionOptions = .default - ) throws -> TranscriptionResult { - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - let optionsJSON = options.toJSON() - - let result = buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_transcribe( - handle, - audioPath, - prompt, - bufferPtr.baseAddress, - bufferPtr.count, - optionsJSON, - nil, - nil, - nil, - 0 - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.transcriptionFailed(error.isEmpty ? "Unknown error" : error) - } - - let responseString = String(cString: buffer) - guard let json = parseJSON(responseString) else { - throw CactusError.invalidResponse - } - - return TranscriptionResult(json: json) - } - - public func transcribe( - pcmData: Data, - prompt: String? = nil, - options: TranscriptionOptions = .default - ) throws -> TranscriptionResult { - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - let optionsJSON = options.toJSON() - - let result = pcmData.withUnsafeBytes { pcmPtr in - buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_transcribe( - handle, - nil, - prompt, - bufferPtr.baseAddress, - bufferPtr.count, - optionsJSON, - nil, - nil, - pcmPtr.baseAddress?.assumingMemoryBound(to: UInt8.self), - pcmData.count - ) - } - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.transcriptionFailed(error.isEmpty ? "Unknown error" : error) - } - - let responseString = String(cString: buffer) - guard let json = parseJSON(responseString) else { - throw CactusError.invalidResponse - } - - return TranscriptionResult(json: json) - } - - public func embed(text: String, normalize: Bool = true) throws -> [Float] { - var embeddingBuffer = [Float](repeating: 0, count: 4096) - var embeddingDim: Int = 0 - - let result = embeddingBuffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_embed( - handle, - text, - bufferPtr.baseAddress, - bufferPtr.count, - &embeddingDim, - normalize - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.embeddingFailed(error.isEmpty ? "Unknown error" : error) - } - - return Array(embeddingBuffer.prefix(embeddingDim)) - } - - public func ragQuery(_ query: String, topK: Int = 5) throws -> String { - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - - let result = buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_rag_query( - handle, - query, - bufferPtr.baseAddress, - bufferPtr.count, - topK - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.completionFailed(error.isEmpty ? "Unknown error" : error) - } - - return String(cString: buffer) - } - - public func tokenize(_ text: String) throws -> [UInt32] { - var tokenBuffer = [UInt32](repeating: 0, count: 8192) - var tokenLen: Int = 0 - - let result = tokenBuffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_tokenize( - handle, - text, - bufferPtr.baseAddress, - bufferPtr.count, - &tokenLen - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.completionFailed(error.isEmpty ? "Unknown error" : error) - } - - return Array(tokenBuffer.prefix(tokenLen)) - } - - public func scoreWindow(tokens: [UInt32], start: Int, end: Int, context: Int) throws -> String { - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - - let result = tokens.withUnsafeBufferPointer { tokenPtr in - buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_score_window( - handle, - tokenPtr.baseAddress, - tokenPtr.count, - start, - end, - context, - bufferPtr.baseAddress, - bufferPtr.count - ) - } - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.completionFailed(error.isEmpty ? "Unknown error" : error) - } - - return String(cString: buffer) - } - - public func imageEmbed(_ imagePath: String) throws -> [Float] { - var embeddingBuffer = [Float](repeating: 0, count: 4096) - var embeddingDim: Int = 0 - - let result = embeddingBuffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_image_embed( - handle, - imagePath, - bufferPtr.baseAddress, - bufferPtr.count, - &embeddingDim - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.embeddingFailed(error.isEmpty ? "Unknown error" : error) - } - - return Array(embeddingBuffer.prefix(embeddingDim)) - } - - public func audioEmbed(_ audioPath: String) throws -> [Float] { - var embeddingBuffer = [Float](repeating: 0, count: 4096) - var embeddingDim: Int = 0 - - let result = embeddingBuffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_audio_embed( - handle, - audioPath, - bufferPtr.baseAddress, - bufferPtr.count, - &embeddingDim - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw CactusError.embeddingFailed(error.isEmpty ? "Unknown error" : error) - } - - return Array(embeddingBuffer.prefix(embeddingDim)) - } - - public func createStreamTranscriber() throws -> StreamTranscriber { - guard let streamHandle = cactus_stream_transcribe_init(handle) else { - let error = String(cString: cactus_get_last_error()) - throw CactusError.transcriptionFailed(error.isEmpty ? "Unknown error" : error) - } - return StreamTranscriber(handle: streamHandle) - } - - public func reset() { - cactus_reset(handle) - } - - public func stop() { - cactus_stop(handle) - } - - - public static func setTelemetryToken(_ token: String) { - cactus_set_telemetry_token(token) - } - - public static func setProKey(_ key: String) { - cactus_set_pro_key(key) - } - - - private func serializeMessages(_ messages: [Message]) throws -> String { - let dicts = messages.map { $0.toDict() } - guard let data = try? JSONSerialization.data(withJSONObject: dicts), - let json = String(data: data, encoding: .utf8) else { - throw CactusError.completionFailed("Failed to serialize messages") - } - return json - } - - private func serializeTools(_ tools: [[String: Any]]?) throws -> String? { - guard let tools = tools else { return nil } - guard let data = try? JSONSerialization.data(withJSONObject: tools), - let json = String(data: data, encoding: .utf8) else { - throw CactusError.completionFailed("Failed to serialize tools") - } - return json - } - - private func parseJSON(_ string: String) -> [String: Any]? { - guard let data = string.data(using: .utf8), - let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { - return nil - } - return json - } -} - - -private class TokenCallbackContext { - let callback: (String, UInt32) -> Void - init(callback: @escaping (String, UInt32) -> Void) { - self.callback = callback - } -} - -private func tokenCallbackBridge(token: UnsafePointer?, tokenId: UInt32, userData: UnsafeMutableRawPointer?) { - guard let token = token, let userData = userData else { return } - let context = Unmanaged.fromOpaque(userData).takeUnretainedValue() - let tokenString = String(cString: token) - context.callback(tokenString, tokenId) -} - - -#if os(iOS) || os(macOS) || os(tvOS) || os(watchOS) || os(visionOS) -@available(iOS 13.0, macOS 10.15, tvOS 13.0, watchOS 6.0, *) -#endif -public extension Cactus { - - func complete( - messages: [Message], - options: CompletionOptions = .default, - tools: [[String: Any]]? = nil - ) async throws -> CompletionResult { - try await withCheckedThrowingContinuation { continuation in - DispatchQueue.global(qos: .userInitiated).async { - do { - let result = try self.complete(messages: messages, options: options, tools: tools) - continuation.resume(returning: result) - } catch { - continuation.resume(throwing: error) - } - } - } - } - - func completeStream( - messages: [Message], - options: CompletionOptions = .default - ) -> AsyncThrowingStream { - AsyncThrowingStream { continuation in - DispatchQueue.global(qos: .userInitiated).async { - do { - _ = try self.complete(messages: messages, options: options) { token, _ in - continuation.yield(token) - } - continuation.finish() - } catch { - continuation.finish(throwing: error) - } - } - } - } - - func transcribe( - audioPath: String, - prompt: String? = nil, - options: TranscriptionOptions = .default - ) async throws -> TranscriptionResult { - try await withCheckedThrowingContinuation { continuation in - DispatchQueue.global(qos: .userInitiated).async { - do { - let result = try self.transcribe(audioPath: audioPath, prompt: prompt, options: options) - continuation.resume(returning: result) - } catch { - continuation.resume(throwing: error) - } - } - } - } - - func embed(text: String, normalize: Bool = true) async throws -> [Float] { - try await withCheckedThrowingContinuation { continuation in - DispatchQueue.global(qos: .userInitiated).async { - do { - let result = try self.embed(text: text, normalize: normalize) - continuation.resume(returning: result) - } catch { - continuation.resume(throwing: error) - } - } - } - } -} - - -public final class StreamTranscriber: @unchecked Sendable { - - private var handle: OpaquePointer? - private static let defaultBufferSize = 65536 - - init(handle: OpaquePointer) { - self.handle = handle - } - - deinit { - close() - } - - public func insert(pcmData: Data) throws { - guard let handle = handle else { - throw Cactus.CactusError.transcriptionFailed("Stream transcriber has been closed") - } - - let result = pcmData.withUnsafeBytes { pcmPtr in - cactus_stream_transcribe_insert( - handle, - pcmPtr.baseAddress?.assumingMemoryBound(to: UInt8.self), - pcmData.count - ) - } - - if result < 0 { - throw Cactus.CactusError.transcriptionFailed("Failed to insert audio data") - } - } - - public func process(language: String? = nil) throws -> Cactus.TranscriptionResult { - guard let handle = handle else { - throw Cactus.CactusError.transcriptionFailed("Stream transcriber has been closed") - } - - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - var optionsJSON: String? = nil - if let lang = language { - let dict: [String: Any] = ["language": lang] - if let data = try? JSONSerialization.data(withJSONObject: dict), - let json = String(data: data, encoding: .utf8) { - optionsJSON = json - } - } - - let result = buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_stream_transcribe_process( - handle, - bufferPtr.baseAddress, - bufferPtr.count, - optionsJSON - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw Cactus.CactusError.transcriptionFailed(error.isEmpty ? "Unknown error" : error) - } - - let responseString = String(cString: buffer) - guard let data = responseString.data(using: .utf8), - let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { - throw Cactus.CactusError.invalidResponse - } - - return Cactus.TranscriptionResult(json: json) - } - - public func finalize() throws -> Cactus.TranscriptionResult { - guard let handle = handle else { - throw Cactus.CactusError.transcriptionFailed("Stream transcriber has been closed") - } - - var buffer = [CChar](repeating: 0, count: Self.defaultBufferSize) - - let result = buffer.withUnsafeMutableBufferPointer { bufferPtr in - cactus_stream_transcribe_finalize( - handle, - bufferPtr.baseAddress, - bufferPtr.count - ) - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw Cactus.CactusError.transcriptionFailed(error.isEmpty ? "Unknown error" : error) - } - - let responseString = String(cString: buffer) - guard let data = responseString.data(using: .utf8), - let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { - throw Cactus.CactusError.invalidResponse - } - - return Cactus.TranscriptionResult(json: json) - } - - public func close() { - if let handle = handle { - cactus_stream_transcribe_destroy(handle) - self.handle = nil - } - } -} - - -public final class CactusIndex: @unchecked Sendable { - - public struct IndexResult { - public let id: Int - public let score: Float - } - - public enum IndexError: Error, LocalizedError { - case initializationFailed(String) - case operationFailed(String) - - public var errorDescription: String? { - switch self { - case .initializationFailed(let msg): return "Index initialization failed: \(msg)" - case .operationFailed(let msg): return "Index operation failed: \(msg)" - } - } - } - - private var handle: OpaquePointer? - - public init(indexDir: String, embeddingDim: Int) throws { - guard let h = cactus_index_init(indexDir, embeddingDim) else { - throw IndexError.initializationFailed("Failed to initialize index") - } - self.handle = h - } - - deinit { - close() - } - - public func add( - ids: [Int], - documents: [String], - embeddings: [[Float]], - metadatas: [String]? = nil - ) throws { - guard let handle = handle else { - throw IndexError.operationFailed("Index has been closed") - } - - let count = ids.count - let embeddingDim = embeddings[0].count - - var idArray = ids.map { Int32($0) } - var docPtrs = documents.map { strdup($0) } - var metaPtrs: [UnsafeMutablePointer?]? = metadatas?.map { strdup($0) } - var embPtrs = embeddings.map { emb -> UnsafePointer? in - let ptr = UnsafeMutablePointer.allocate(capacity: emb.count) - ptr.initialize(from: emb, count: emb.count) - return UnsafePointer(ptr) - } - - let result = idArray.withUnsafeMutableBufferPointer { idPtr in - docPtrs.withUnsafeMutableBufferPointer { docPtr in - embPtrs.withUnsafeMutableBufferPointer { embPtr in - if let metaPtrs = metaPtrs { - var metaPtrsCopy = metaPtrs - return metaPtrsCopy.withUnsafeMutableBufferPointer { metaPtr in - cactus_index_add( - handle, - idPtr.baseAddress, - docPtr.baseAddress, - metaPtr.baseAddress, - embPtr.baseAddress, - count, - embeddingDim - ) - } - } else { - return cactus_index_add( - handle, - idPtr.baseAddress, - docPtr.baseAddress, - nil, - embPtr.baseAddress, - count, - embeddingDim - ) - } - } - } - } - - docPtrs.forEach { free($0) } - metaPtrs?.forEach { free($0) } - embPtrs.forEach { ptr in - if let ptr = ptr { - UnsafeMutablePointer(mutating: ptr).deallocate() - } - } - - if result < 0 { - throw IndexError.operationFailed("Failed to add documents to index") - } - } - - public func delete(ids: [Int]) throws { - guard let handle = handle else { - throw IndexError.operationFailed("Index has been closed") - } - - var idArray = ids.map { Int32($0) } - - let result = idArray.withUnsafeMutableBufferPointer { idPtr in - cactus_index_delete(handle, idPtr.baseAddress, ids.count) - } - - if result < 0 { - throw IndexError.operationFailed("Failed to delete documents from index") - } - } - - public func query(embedding: [Float], topK: Int = 5) throws -> [IndexResult] { - guard let handle = handle else { - throw IndexError.operationFailed("Index has been closed") - } - - var embeddingCopy = embedding - var idBuffer = [Int32](repeating: 0, count: topK) - var scoreBuffer = [Float](repeating: 0, count: topK) - var idBufferSize = topK - var scoreBufferSize = topK - - let result = embeddingCopy.withUnsafeMutableBufferPointer { embPtr in - idBuffer.withUnsafeMutableBufferPointer { idPtr in - scoreBuffer.withUnsafeMutableBufferPointer { scorePtr in - var embPtrPtr: UnsafePointer? = embPtr.baseAddress - var idPtrPtr: UnsafeMutablePointer? = idPtr.baseAddress - var scorePtrPtr: UnsafeMutablePointer? = scorePtr.baseAddress - - return withUnsafeMutablePointer(to: &embPtrPtr) { embPtrPtrPtr in - withUnsafeMutablePointer(to: &idPtrPtr) { idPtrPtrPtr in - withUnsafeMutablePointer(to: &scorePtrPtr) { scorePtrPtrPtr in - withUnsafeMutablePointer(to: &idBufferSize) { idSizePtr in - withUnsafeMutablePointer(to: &scoreBufferSize) { scoreSizePtr in - cactus_index_query( - handle, - embPtrPtrPtr, - 1, - embedding.count, - nil, - idPtrPtrPtr, - idSizePtr, - scorePtrPtrPtr, - scoreSizePtr - ) - } - } - } - } - } - } - } - } - - if result < 0 { - let error = String(cString: cactus_get_last_error()) - throw IndexError.operationFailed(error.isEmpty ? "Unknown error" : error) - } - - return (0..CFBundleVersion 1.0.0 MinimumOSVersion - 12.0 + 13.3 diff --git a/apple/README.md b/apple/README.md index ce445849e..98fe929e3 100644 --- a/apple/README.md +++ b/apple/README.md @@ -1,292 +1,36 @@ -# Cactus for Swift Multiplatform +# Apple Build -Run AI models on-device with a simple Swift API on iOS, macOS, and Android. +Builds `libcactus_engine` for iOS and macOS. -## Building +## Usage ```bash cactus build --apple ``` -Build outputs (in `apple/`): - -see the main [README.md](../README.md) for how to use CLI & download weight - -| File | Description | -|------|-------------| -| `cactus-ios.xcframework/` | iOS framework (device + simulator) | -| `cactus-macos.xcframework/` | macOS framework | -| `libcactus-device.a` | Static library for iOS device | -| `libcactus-simulator.a` | Static library for iOS simulator | - -For Android, build `libcactus.so` from the `android/` directory. - -## Integration - -### iOS/macOS: XCFramework (Recommended) - -1. Drag `cactus-ios.xcframework` (or `cactus-macos.xcframework`) into your Xcode project -2. Ensure "Embed & Sign" is selected in "Frameworks, Libraries, and Embedded Content" -3. Copy `Cactus.swift` into your project - -### iOS/macOS: Static Library - -1. Add `libcactus-device.a` (or `libcactus-simulator.a`) to "Link Binary With Libraries" -2. Create a folder with `cactus_ffi.h` and `module.modulemap`, add to Build Settings: - - "Header Search Paths" → path to folder - - "Import Paths" (Swift) → path to folder -3. Copy `Cactus.swift` into your project - -### Android (Swift SDK) +Or directly: -Requires [Swift SDK for Android](https://www.swift.org/documentation/articles/swift-sdk-for-android-getting-started.html). - -1. Copy files to your Swift project: - - `libcactus.so` → your library path - - `cactus_ffi.h` → your include path - - `module.android.modulemap` → rename to `module.modulemap` in include path - - `Cactus.swift` → your sources - -2. Build with Swift SDK for Android: ```bash -swift build --swift-sdk aarch64-unknown-linux-android28 \ - -Xswiftc -I/path/to/include \ - -Xlinker -L/path/to/lib \ - -Xlinker -lcactus -``` - -3. Bundle `libcactus.so` with your APK in `jniLibs/arm64-v8a/` - -## Usage - -### Basic Completion - -```swift -import Foundation - -let model = try Cactus(modelPath: "/path/to/model") -let result = try model.complete("What is the capital of France?") -``` - -### Chat Messages - -```swift -let result = try model.complete(messages: [ - .system("You are a helpful assistant."), - .user("What is 2 + 2?") -]) -``` - -### Completion Options - -```swift -let options = Cactus.CompletionOptions( - temperature: 0.7, - topP: 0.9, - topK: 40, - maxTokens: 256, - stopSequences: ["\n\n"] -) - -let result = try model.complete("Write a haiku:", options: options) -``` - -### Streaming Tokens - -```swift -let result = try model.complete( - messages: [.user("Tell me a story")], - onToken: { token, tokenId in - print(token, terminator: "") - fflush(stdout) - } -) -``` - -### Async/Await - -```swift -let result = try await model.complete(messages: [.user("Hello!")]) - -for try await token in model.completeStream(messages: [.user("Tell me a joke")]) { - print(token, terminator: "") -} -``` - -### Audio Transcription - -```swift -// From file -let result = try model.transcribe(audioPath: "/path/to/audio.wav") - -// From PCM data -let pcmData: Data = ... // 16kHz mono PCM -let result = try model.transcribe(pcmData: pcmData) -``` - -### Embeddings - -```swift -let embedding = try model.embed(text: "Hello, world!") -let imageEmbedding = try model.imageEmbed("/path/to/image.jpg") -let audioEmbedding = try model.audioEmbed("/path/to/audio.wav") -``` - -### Tokenization - -```swift -let tokens = try model.tokenize("Hello, world!") -let scores = try model.scoreWindow(tokens: tokens, start: 0, end: tokens.count, context: 512) -``` - -### Streaming Transcription - -```swift -let stream = try model.createStreamTranscriber() -try stream.insert(pcmData: audioChunk1) -try stream.insert(pcmData: audioChunk2) -let partial = try stream.process() -print("Partial: \(partial.text)") -let final = try stream.finalize() -print("Final: \(final.text)") -stream.close() -``` - -### RAG (Retrieval-Augmented Generation) - -```swift -let model = try Cactus( - modelPath: "/path/to/model", - corpusDir: "/path/to/documents" -) - -let result = try model.complete("What does the documentation say about X?") +bash apple/build.sh ``` -### Vector Index - -```swift -let index = try CactusIndex(indexDir: "/path/to/index", embeddingDim: 384) -let embeddings = [try model.embed(text: "doc1"), try model.embed(text: "doc2")] -try index.add( - ids: [1, 2], - documents: ["Document 1", "Document 2"], - embeddings: embeddings -) -let results = try index.query(embedding: try model.embed(text: "search query"), topK: 5) -results.forEach { print("ID: \($0.id), Score: \($0.score)") } -index.close() -``` - -## API Reference - -### Cactus +## Output -```swift -init(modelPath: String, corpusDir: String? = nil) throws +- `libcactus_engine-device.a` — iOS device (arm64) +- `libcactus_engine-simulator.a` — iOS simulator (arm64) +- `cactus-ios.xcframework/` — iOS XCFramework (device + simulator) +- `cactus-macos.xcframework/` — macOS XCFramework (arm64) -func complete(_ prompt: String, options: CompletionOptions = .default) throws -> CompletionResult -func complete(messages: [Message], options: CompletionOptions = .default, tools: [[String: Any]]? = nil, onToken: ((String, UInt32) -> Void)? = nil) throws -> CompletionResult +## Options -func transcribe(audioPath: String, prompt: String? = nil, options: TranscriptionOptions = .default) throws -> TranscriptionResult -func transcribe(pcmData: Data, prompt: String? = nil, options: TranscriptionOptions = .default) throws -> TranscriptionResult - -func embed(text: String, normalize: Bool = true) throws -> [Float] -func imageEmbed(_ imagePath: String) throws -> [Float] -func audioEmbed(_ audioPath: String) throws -> [Float] -func ragQuery(_ query: String, topK: Int = 5) throws -> String - -func tokenize(_ text: String) throws -> [UInt32] -func scoreWindow(tokens: [UInt32], start: Int, end: Int, context: Int) throws -> String -func createStreamTranscriber() throws -> StreamTranscriber - -func reset() // Clear KV cache -func stop() // Stop generation - -static func setTelemetryToken(_ token: String) -static func setProKey(_ key: String) -``` - -### CompletionResult - -```swift -struct CompletionResult { - let text: String - let functionCalls: [[String: Any]]? - let promptTokens: Int - let completionTokens: Int - let timeToFirstToken: Double - let totalTime: Double - let prefillTokensPerSecond: Double - let decodeTokensPerSecond: Double - let confidence: Double - let needsCloudHandoff: Bool -} -``` - -### Message - -```swift -struct Message { - static func system(_ content: String) -> Message - static func user(_ content: String) -> Message - static func assistant(_ content: String) -> Message -} -``` - -### CompletionOptions - -```swift -struct CompletionOptions { - var temperature: Float = 0.7 - var topP: Float = 0.9 - var topK: Int = 40 - var maxTokens: Int = 512 - var stopSequences: [String] = [] - var confidenceThreshold: Float = 0.0 - - static let `default` = CompletionOptions() -} -``` - -### StreamTranscriber - -```swift -class StreamTranscriber { - func insert(pcmData: Data) throws - func process(language: String? = nil) throws -> TranscriptionResult - func finalize() throws -> TranscriptionResult - func close() -} -``` - -### CactusIndex - -```swift -class CactusIndex { - init(indexDir: String, embeddingDim: Int) throws - - func add(ids: [Int], documents: [String], embeddings: [[Float]], metadatas: [String]? = nil) throws - func delete(ids: [Int]) throws - func query(embedding: [Float], topK: Int = 5) throws -> [IndexResult] - func compact() throws - func close() -} - -struct IndexResult { - let id: Int - let score: Float -} -``` +| Variable | Default | Description | +|----------|---------|-------------| +| `BUILD_STATIC` | `true` | Build static libraries | +| `BUILD_XCFRAMEWORK` | `true` | Build XCFrameworks | +| `CMAKE_BUILD_TYPE` | `Release` | CMake build type | +| `CACTUS_CURL_ROOT` | `cactus-engine/libs/curl` | Vendored libcurl path | ## Requirements -**Apple Platforms:** -- iOS 14.0+ / macOS 13.0+ / tvOS 14.0+ / watchOS 7.0+ -- Xcode 14.0+ -- Swift 5.7+ - -**Android:** -- Swift 6.0+ with [Swift SDK for Android](https://www.swift.org/documentation/articles/swift-sdk-for-android-getting-started.html) -- Android NDK 27d+ -- Android API 28+ / arm64-v8a +- Xcode with iOS SDK +- CMake 3.10+ diff --git a/apple/build.sh b/apple/build.sh index 4faae865f..f71ee5c5a 100755 --- a/apple/build.sh +++ b/apple/build.sh @@ -1,4 +1,7 @@ -#!/bin/bash -e +#!/bin/bash +IOS_MIN_VERSION=16.4 +MACOS_MIN_VERSION=13.3 +set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" @@ -7,6 +10,7 @@ APPLE_DIR="$ROOT_DIR/apple" CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE:-Release} BUILD_STATIC=${BUILD_STATIC:-true} BUILD_XCFRAMEWORK=${BUILD_XCFRAMEWORK:-true} +CACTUS_CURL_ROOT=${CACTUS_CURL_ROOT:-"$ROOT_DIR/cactus-engine/libs/curl"} if ! command -v cmake &> /dev/null; then echo "Error: cmake not found, please install it" @@ -26,88 +30,21 @@ echo "Build type: $CMAKE_BUILD_TYPE" echo "Using $n_cpu CPU cores" echo "Static library: $BUILD_STATIC" echo "XCFramework: $BUILD_XCFRAMEWORK" +echo "Vendored libcurl root: $CACTUS_CURL_ROOT" function cp_headers() { mkdir -p "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers" - cp "$ROOT_DIR/cactus/ffi/"*.h "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers/" 2>/dev/null || true - cp "$ROOT_DIR/cactus/engine/"*.h "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers/" 2>/dev/null || true - cp "$ROOT_DIR/cactus/graph/"*.h "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers/" 2>/dev/null || true - cp "$ROOT_DIR/cactus/kernel/"*.h "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers/" 2>/dev/null || true - cp "$ROOT_DIR/cactus/"*.h "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers/" 2>/dev/null || true + cp "$ROOT_DIR/cactus-engine/cactus_engine.h" "$ROOT_DIR/apple/$1/$2/cactus.framework/Headers/" + mkdir -p "$ROOT_DIR/apple/$1/$2/cactus.framework/Modules" + cp "$ROOT_DIR/apple/module.modulemap" "$ROOT_DIR/apple/$1/$2/cactus.framework/Modules/" } function create_ios_xcframework_info_plist() { - cat > "$ROOT_DIR/apple/cactus-ios.xcframework/Info.plist" << 'EOF' - - - - - AvailableLibraries - - - LibraryIdentifier - ios-arm64 - LibraryPath - cactus.framework - SupportedArchitectures - - arm64 - - SupportedPlatform - ios - - - LibraryIdentifier - ios-arm64-simulator - LibraryPath - cactus.framework - SupportedArchitectures - - arm64 - - SupportedPlatform - ios - SupportedPlatformVariant - simulator - - - CFBundlePackageType - XFWK - XCFrameworkFormatVersion - 1.0 - - -EOF + cp "$ROOT_DIR/apple/cactus-ios.Info.plist" "$ROOT_DIR/apple/cactus-ios.xcframework/Info.plist" } function create_macos_xcframework_info_plist() { - cat > "$ROOT_DIR/apple/cactus-macos.xcframework/Info.plist" << 'EOF' - - - - - AvailableLibraries - - - LibraryIdentifier - macos-arm64 - LibraryPath - cactus.framework - SupportedArchitectures - - arm64 - - SupportedPlatform - macos - - - CFBundlePackageType - XFWK - XCFrameworkFormatVersion - 1.0 - - -EOF + cp "$ROOT_DIR/apple/cactus-macos.Info.plist" "$ROOT_DIR/apple/cactus-macos.xcframework/Info.plist" } function build_static_library() { @@ -124,27 +61,19 @@ function build_static_library() { cmake -DCMAKE_SYSTEM_NAME=iOS \ -DCMAKE_OSX_ARCHITECTURES=arm64 \ - -DCMAKE_OSX_DEPLOYMENT_TARGET=13.0 \ + -DCMAKE_OSX_DEPLOYMENT_TARGET=$IOS_MIN_VERSION \ -DCMAKE_OSX_SYSROOT="$IOS_SDK_PATH" \ -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ -DBUILD_SHARED_LIBS=OFF \ + -DCACTUS_CURL_ROOT="$CACTUS_CURL_ROOT" \ -S "$APPLE_DIR" \ -B "$BUILD_DIR" >/dev/null cmake --build "$BUILD_DIR" --config "$CMAKE_BUILD_TYPE" -j "$n_cpu" >/dev/null mkdir -p "$APPLE_DIR" - - PRO_LIB="$ROOT_DIR/libs/libcactus_pro-ios.a" - if [ -f "$PRO_LIB" ]; then - echo "Merging libcactus_pro-ios.a into device static library..." - libtool -static -o "$APPLE_DIR/libcactus-device.a" "$BUILD_DIR/libcactus.a" "$PRO_LIB" - else - echo "Note: libcactus_pro-ios.a not found, building without NPU support" - cp "$BUILD_DIR/libcactus.a" "$APPLE_DIR/libcactus-device.a" - fi - - echo "Device static library built: $APPLE_DIR/libcactus-device.a" + cp "$BUILD_DIR/libcactus_engine.a" "$APPLE_DIR/libcactus_engine-device.a" + echo "Device static library built: $APPLE_DIR/libcactus_engine-device.a" echo "Building static library for iOS simulator..." BUILD_DIR_SIM="$APPLE_DIR/build-static-simulator" @@ -159,25 +88,18 @@ function build_static_library() { cmake -DCMAKE_SYSTEM_NAME=iOS \ -DCMAKE_OSX_ARCHITECTURES=arm64 \ - -DCMAKE_OSX_DEPLOYMENT_TARGET=13.0 \ + -DCMAKE_OSX_DEPLOYMENT_TARGET=$IOS_MIN_VERSION \ -DCMAKE_OSX_SYSROOT="$IOS_SIM_SDK_PATH" \ -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ -DBUILD_SHARED_LIBS=OFF \ + -DCACTUS_CURL_ROOT="$CACTUS_CURL_ROOT" \ -S "$APPLE_DIR" \ -B "$BUILD_DIR_SIM" >/dev/null cmake --build "$BUILD_DIR_SIM" --config "$CMAKE_BUILD_TYPE" -j "$n_cpu" >/dev/null - PRO_LIB_SIM="$ROOT_DIR/libs/libcactus_pro-ios-simulator.a" - if [ -f "$PRO_LIB_SIM" ]; then - echo "Merging libcactus_pro-ios-simulator.a into simulator static library..." - libtool -static -o "$APPLE_DIR/libcactus-simulator.a" "$BUILD_DIR_SIM/libcactus.a" "$PRO_LIB_SIM" - else - echo "Note: libcactus_pro-ios-simulator.a not found, building without NPU support" - cp "$BUILD_DIR_SIM/libcactus.a" "$APPLE_DIR/libcactus-simulator.a" - fi - - echo "Simulator static library built: $APPLE_DIR/libcactus-simulator.a" + cp "$BUILD_DIR_SIM/libcactus_engine.a" "$APPLE_DIR/libcactus_engine-simulator.a" + echo "Simulator static library built: $APPLE_DIR/libcactus_engine-simulator.a" } function build_framework() { @@ -190,9 +112,10 @@ function build_framework() { -DCMAKE_SYSTEM_NAME=$1 \ -DCMAKE_OSX_ARCHITECTURES="$2" \ -DCMAKE_OSX_SYSROOT=$3 \ - -DCMAKE_OSX_DEPLOYMENT_TARGET=13.0 \ + -DCMAKE_OSX_DEPLOYMENT_TARGET=$([ "$1" = "iOS" ] && echo "$IOS_MIN_VERSION" || echo "$MACOS_MIN_VERSION") \ -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ -DBUILD_SHARED_LIBS=ON \ + -DCACTUS_CURL_ROOT="$CACTUS_CURL_ROOT" \ -DCMAKE_XCODE_ATTRIBUTE_ONLY_ACTIVE_ARCH=NO \ -DCMAKE_IOS_INSTALL_COMBINED=YES >/dev/null @@ -285,12 +208,12 @@ echo "Total time: $((t1 - t0)) seconds" if [ "$BUILD_STATIC" = "true" ]; then rm -rf "$APPLE_DIR/build-static-device" "$APPLE_DIR/build-static-simulator" "$APPLE_DIR/build-static-macos" echo "Static libraries:" - echo " Device: $APPLE_DIR/libcactus-device.a" - echo " Simulator: $APPLE_DIR/libcactus-simulator.a" + echo " Device: $APPLE_DIR/libcactus_engine-device.a" + echo " Simulator: $APPLE_DIR/libcactus_engine-simulator.a" fi if [ "$BUILD_XCFRAMEWORK" = "true" ]; then echo "XCFrameworks:" echo " iOS: $APPLE_DIR/cactus-ios.xcframework" echo " macOS: $APPLE_DIR/cactus-macos.xcframework" -fi \ No newline at end of file +fi diff --git a/apple/cactus-ios.Info.plist b/apple/cactus-ios.Info.plist new file mode 100644 index 000000000..55424d634 --- /dev/null +++ b/apple/cactus-ios.Info.plist @@ -0,0 +1,39 @@ + + + + + AvailableLibraries + + + LibraryIdentifier + ios-arm64 + LibraryPath + cactus.framework + SupportedArchitectures + + arm64 + + SupportedPlatform + ios + + + LibraryIdentifier + ios-arm64-simulator + LibraryPath + cactus.framework + SupportedArchitectures + + arm64 + + SupportedPlatform + ios + SupportedPlatformVariant + simulator + + + CFBundlePackageType + XFWK + XCFrameworkFormatVersion + 1.0 + + diff --git a/apple/cactus-macos.Info.plist b/apple/cactus-macos.Info.plist new file mode 100644 index 000000000..fe06ebed8 --- /dev/null +++ b/apple/cactus-macos.Info.plist @@ -0,0 +1,25 @@ + + + + + AvailableLibraries + + + LibraryIdentifier + macos-arm64 + LibraryPath + cactus.framework + SupportedArchitectures + + arm64 + + SupportedPlatform + macos + + + CFBundlePackageType + XFWK + XCFrameworkFormatVersion + 1.0 + + diff --git a/apple/module.modulemap b/apple/module.modulemap index daa014138..706a1fd96 100644 --- a/apple/module.modulemap +++ b/apple/module.modulemap @@ -1,5 +1,4 @@ framework module cactus { - umbrella header "cactus_ffi.h" + header "cactus_engine.h" export * - module * { export * } } diff --git a/assets/logo_white.png b/assets/logo_white.png new file mode 100644 index 000000000..7f924127d Binary files /dev/null and b/assets/logo_white.png differ diff --git a/tests/assets/test.wav b/assets/test.wav similarity index 100% rename from tests/assets/test.wav rename to assets/test.wav diff --git a/bindings/README.md b/bindings/README.md new file mode 100644 index 000000000..e0926dcfa --- /dev/null +++ b/bindings/README.md @@ -0,0 +1,13 @@ +# Bindings + +- `python/` uses `ctypes` +- `swift/` uses a C module map +- `kotlin/` uses JNI on Android and Kotlin/Native cinterop on iOS (KMP) +- `flutter/` uses Dart FFI +- `rust/` uses raw `extern "C"` declarations +- `react-native/` contains a thin React Native bridge backed by the raw Kotlin and Swift bindings + +Platform-native build and packaging stay outside this folder: + +- [`android/`](/android/) builds Android native artifacts +- [`apple/`](/apple/) builds Apple native artifacts diff --git a/bindings/flutter/README.md b/bindings/flutter/README.md new file mode 100644 index 000000000..1e36bc559 --- /dev/null +++ b/bindings/flutter/README.md @@ -0,0 +1,62 @@ +# Flutter Bindings + +Dart FFI bindings to `cactus_engine.h`. + +## Integration + + +```bash +cactus build --apple +cactus build --android +``` + + + +1. Copy `android/libcactus_engine.so` to your app's `jniLibs/arm64-v8a/` +2. Add `apple/cactus-ios.xcframework` to your Xcode project (iOS) +3. Add `apple/cactus-macos.xcframework` to your Xcode project (macOS) +4. Copy `cactus.dart` into your Dart source tree +5. Add `ffi` to `pubspec.yaml` + + +## Usage + + +```dart +import 'dart:ffi'; +import 'cactus.dart'; +import 'package:ffi/ffi.dart'; + +final modelPath = '/path/to/model'.toNativeUtf8(); +final model = cactusInit(modelPath, nullptr, false); +calloc.free(modelPath); + +const messagesJson = '[{"role":"user","content":"Hello"}]'; +final msgs = messagesJson.toNativeUtf8(); +final buf = calloc(65536); +cactusComplete(model, msgs, buf.cast(), 65536, nullptr, nullptr, nullptr, nullptr, nullptr, 0); +final response = buf.cast().toDartString(); +calloc.free(msgs); +calloc.free(buf); +cactusDestroy(model); +``` + + +## Streaming transcription + +Push 16 kHz mono PCM16 as it arrives; each call returns `{"success":true,"confirmed":...,"pending":...}` (`confirmed` is final, `pending` is the volatile tail). + +```dart +final opts = '{"language":"en"}'.toNativeUtf8(); +final stream = cactusStreamTranscribeStart(model, opts); +final out = calloc(65536); +for (final chunk in pcmChunks) { // each: 16 kHz mono PCM16 bytes + final p = calloc(chunk.length)..asTypedList(chunk.length).setAll(0, chunk); + cactusStreamTranscribeProcess(stream, p, chunk.length, out.cast(), 65536); + calloc.free(p); + // parse {"confirmed":...,"pending":...} from out.cast().toDartString() +} +cactusStreamTranscribeStop(stream, out.cast(), 65536); +calloc.free(opts); +calloc.free(out); +``` diff --git a/bindings/flutter/cactus.dart b/bindings/flutter/cactus.dart new file mode 100644 index 000000000..f45fb6d25 --- /dev/null +++ b/bindings/flutter/cactus.dart @@ -0,0 +1,391 @@ +import 'dart:ffi'; +import 'dart:io'; + +import 'package:ffi/ffi.dart'; + +final DynamicLibrary _lib = (Platform.isAndroid || Platform.isLinux) + ? DynamicLibrary.open('libcactus_engine.so') + : DynamicLibrary.process(); + +typedef CactusModelT = Pointer; +typedef CactusIndexT = Pointer; +typedef CactusStreamTranscribeT = Pointer; + +typedef TokenCallbackNative = Void Function( + Pointer token, Uint32 tokenId, Pointer userData); +typedef LogCallbackNative = Void Function( + Int32 level, Pointer component, Pointer message, + Pointer userData); + +typedef _CactusInitN = Pointer Function( + Pointer modelPath, Pointer corpusDir, Bool cacheIndex); +typedef _CactusInitD = Pointer Function( + Pointer modelPath, Pointer corpusDir, bool cacheIndex); +final cactusInit = _lib.lookupFunction<_CactusInitN, _CactusInitD>('cactus_init'); + +typedef _CactusDestroyN = Void Function(Pointer model); +typedef _CactusDestroyD = void Function(Pointer model); +final cactusDestroy = + _lib.lookupFunction<_CactusDestroyN, _CactusDestroyD>('cactus_destroy'); + +typedef _CactusResetN = Void Function(Pointer model); +typedef _CactusResetD = void Function(Pointer model); +final cactusReset = + _lib.lookupFunction<_CactusResetN, _CactusResetD>('cactus_reset'); + +typedef _CactusStopN = Void Function(Pointer model); +typedef _CactusStopD = void Function(Pointer model); +final cactusStop = + _lib.lookupFunction<_CactusStopN, _CactusStopD>('cactus_stop'); + +typedef _CactusCompleteN = Int32 Function( + Pointer model, + Pointer messagesJson, + Pointer responseBuffer, + IntPtr bufferSize, + Pointer optionsJson, + Pointer toolsJson, + Pointer> callback, + Pointer userData, + Pointer pcmBuffer, + IntPtr pcmBufferSize); +typedef _CactusCompleteD = int Function( + Pointer model, + Pointer messagesJson, + Pointer responseBuffer, + int bufferSize, + Pointer optionsJson, + Pointer toolsJson, + Pointer> callback, + Pointer userData, + Pointer pcmBuffer, + int pcmBufferSize); +final cactusComplete = + _lib.lookupFunction<_CactusCompleteN, _CactusCompleteD>('cactus_complete'); + +typedef _CactusPrefillN = Int32 Function( + Pointer model, + Pointer messagesJson, + Pointer responseBuffer, + IntPtr bufferSize, + Pointer optionsJson, + Pointer toolsJson, + Pointer pcmBuffer, + IntPtr pcmBufferSize); +typedef _CactusPrefillD = int Function( + Pointer model, + Pointer messagesJson, + Pointer responseBuffer, + int bufferSize, + Pointer optionsJson, + Pointer toolsJson, + Pointer pcmBuffer, + int pcmBufferSize); +final cactusPrefill = + _lib.lookupFunction<_CactusPrefillN, _CactusPrefillD>('cactus_prefill'); + +typedef _CactusTokenizeN = Int32 Function( + Pointer model, + Pointer text, + Pointer tokenBuffer, + IntPtr tokenBufferLen, + Pointer outTokenLen); +typedef _CactusTokenizeD = int Function( + Pointer model, + Pointer text, + Pointer tokenBuffer, + int tokenBufferLen, + Pointer outTokenLen); +final cactusTokenize = + _lib.lookupFunction<_CactusTokenizeN, _CactusTokenizeD>('cactus_tokenize'); + +typedef _CactusScoreWindowN = Int32 Function( + Pointer model, + Pointer tokens, + IntPtr tokenLen, + IntPtr start, + IntPtr end, + IntPtr context, + Pointer responseBuffer, + IntPtr bufferSize); +typedef _CactusScoreWindowD = int Function( + Pointer model, + Pointer tokens, + int tokenLen, + int start, + int end, + int context, + Pointer responseBuffer, + int bufferSize); +final cactusScoreWindow = + _lib.lookupFunction<_CactusScoreWindowN, _CactusScoreWindowD>( + 'cactus_score_window'); + +typedef _CactusTranscribeN = Int32 Function( + Pointer model, + Pointer audioFilePath, + Pointer prompt, + Pointer responseBuffer, + IntPtr bufferSize, + Pointer optionsJson, + Pointer> callback, + Pointer userData, + Pointer pcmBuffer, + IntPtr pcmBufferSize); +typedef _CactusTranscribeD = int Function( + Pointer model, + Pointer audioFilePath, + Pointer prompt, + Pointer responseBuffer, + int bufferSize, + Pointer optionsJson, + Pointer> callback, + Pointer userData, + Pointer pcmBuffer, + int pcmBufferSize); +final cactusTranscribe = + _lib.lookupFunction<_CactusTranscribeN, _CactusTranscribeD>( + 'cactus_transcribe'); + +typedef _CactusStreamTranscribeStartN = Pointer Function( + Pointer model, Pointer optionsJson); +typedef _CactusStreamTranscribeStartD = Pointer Function( + Pointer model, Pointer optionsJson); +final cactusStreamTranscribeStart = _lib.lookupFunction< + _CactusStreamTranscribeStartN, + _CactusStreamTranscribeStartD>('cactus_stream_transcribe_start'); + +typedef _CactusStreamTranscribeProcessN = Int32 Function( + Pointer stream, + Pointer pcmBuffer, + IntPtr pcmBufferSize, + Pointer responseBuffer, + IntPtr bufferSize); +typedef _CactusStreamTranscribeProcessD = int Function( + Pointer stream, + Pointer pcmBuffer, + int pcmBufferSize, + Pointer responseBuffer, + int bufferSize); +final cactusStreamTranscribeProcess = _lib.lookupFunction< + _CactusStreamTranscribeProcessN, + _CactusStreamTranscribeProcessD>('cactus_stream_transcribe_process'); + +typedef _CactusStreamTranscribeStopN = Int32 Function( + Pointer stream, Pointer responseBuffer, IntPtr bufferSize); +typedef _CactusStreamTranscribeStopD = int Function( + Pointer stream, Pointer responseBuffer, int bufferSize); +final cactusStreamTranscribeStop = _lib.lookupFunction< + _CactusStreamTranscribeStopN, + _CactusStreamTranscribeStopD>('cactus_stream_transcribe_stop'); + +typedef _CactusEmbedN = Int32 Function( + Pointer model, + Pointer text, + Pointer embeddingsBuffer, + IntPtr bufferSize, + Pointer embeddingDim, + Bool normalize); +typedef _CactusEmbedD = int Function( + Pointer model, + Pointer text, + Pointer embeddingsBuffer, + int bufferSize, + Pointer embeddingDim, + bool normalize); +final cactusEmbed = + _lib.lookupFunction<_CactusEmbedN, _CactusEmbedD>('cactus_embed'); + +typedef _CactusImageEmbedN = Int32 Function( + Pointer model, + Pointer imagePath, + Pointer embeddingsBuffer, + IntPtr bufferSize, + Pointer embeddingDim); +typedef _CactusImageEmbedD = int Function( + Pointer model, + Pointer imagePath, + Pointer embeddingsBuffer, + int bufferSize, + Pointer embeddingDim); +final cactusImageEmbed = + _lib.lookupFunction<_CactusImageEmbedN, _CactusImageEmbedD>( + 'cactus_image_embed'); + +typedef _CactusAudioEmbedN = Int32 Function( + Pointer model, + Pointer audioPath, + Pointer embeddingsBuffer, + IntPtr bufferSize, + Pointer embeddingDim); +typedef _CactusAudioEmbedD = int Function( + Pointer model, + Pointer audioPath, + Pointer embeddingsBuffer, + int bufferSize, + Pointer embeddingDim); +final cactusAudioEmbed = + _lib.lookupFunction<_CactusAudioEmbedN, _CactusAudioEmbedD>( + 'cactus_audio_embed'); + +typedef _CactusRagQueryN = Int32 Function( + Pointer model, + Pointer query, + Pointer responseBuffer, + IntPtr bufferSize, + IntPtr topK); +typedef _CactusRagQueryD = int Function( + Pointer model, + Pointer query, + Pointer responseBuffer, + int bufferSize, + int topK); +final cactusRagQuery = + _lib.lookupFunction<_CactusRagQueryN, _CactusRagQueryD>( + 'cactus_rag_query'); + +typedef _CactusIndexInitN = Pointer Function( + Pointer indexDir, IntPtr embeddingDim); +typedef _CactusIndexInitD = Pointer Function( + Pointer indexDir, int embeddingDim); +final cactusIndexInit = + _lib.lookupFunction<_CactusIndexInitN, _CactusIndexInitD>( + 'cactus_index_init'); + +typedef _CactusIndexAddN = Int32 Function( + Pointer index, + Pointer ids, + Pointer> documents, + Pointer> metadatas, + Pointer> embeddings, + IntPtr count, + IntPtr embeddingDim); +typedef _CactusIndexAddD = int Function( + Pointer index, + Pointer ids, + Pointer> documents, + Pointer> metadatas, + Pointer> embeddings, + int count, + int embeddingDim); +final cactusIndexAdd = + _lib.lookupFunction<_CactusIndexAddN, _CactusIndexAddD>( + 'cactus_index_add'); + +typedef _CactusIndexDeleteN = Int32 Function( + Pointer index, Pointer ids, IntPtr idsCount); +typedef _CactusIndexDeleteD = int Function( + Pointer index, Pointer ids, int idsCount); +final cactusIndexDelete = + _lib.lookupFunction<_CactusIndexDeleteN, _CactusIndexDeleteD>( + 'cactus_index_delete'); + +typedef _CactusIndexGetN = Int32 Function( + Pointer index, + Pointer ids, + IntPtr idsCount, + Pointer> documentBuffers, + Pointer documentBufferSizes, + Pointer> metadataBuffers, + Pointer metadataBufferSizes, + Pointer> embeddingBuffers, + Pointer embeddingBufferSizes); +typedef _CactusIndexGetD = int Function( + Pointer index, + Pointer ids, + int idsCount, + Pointer> documentBuffers, + Pointer documentBufferSizes, + Pointer> metadataBuffers, + Pointer metadataBufferSizes, + Pointer> embeddingBuffers, + Pointer embeddingBufferSizes); +final cactusIndexGet = + _lib.lookupFunction<_CactusIndexGetN, _CactusIndexGetD>( + 'cactus_index_get'); + +typedef _CactusIndexQueryN = Int32 Function( + Pointer index, + Pointer> embeddings, + IntPtr embeddingsCount, + IntPtr embeddingDim, + Pointer optionsJson, + Pointer> idBuffers, + Pointer idBufferSizes, + Pointer> scoreBuffers, + Pointer scoreBufferSizes); +typedef _CactusIndexQueryD = int Function( + Pointer index, + Pointer> embeddings, + int embeddingsCount, + int embeddingDim, + Pointer optionsJson, + Pointer> idBuffers, + Pointer idBufferSizes, + Pointer> scoreBuffers, + Pointer scoreBufferSizes); +final cactusIndexQuery = + _lib.lookupFunction<_CactusIndexQueryN, _CactusIndexQueryD>( + 'cactus_index_query'); + +typedef _CactusIndexCompactN = Int32 Function(Pointer index); +typedef _CactusIndexCompactD = int Function(Pointer index); +final cactusIndexCompact = + _lib.lookupFunction<_CactusIndexCompactN, _CactusIndexCompactD>( + 'cactus_index_compact'); + +typedef _CactusIndexDestroyN = Void Function(Pointer index); +typedef _CactusIndexDestroyD = void Function(Pointer index); +final cactusIndexDestroy = + _lib.lookupFunction<_CactusIndexDestroyN, _CactusIndexDestroyD>( + 'cactus_index_destroy'); + +typedef _CactusGetLastErrorN = Pointer Function(); +typedef _CactusGetLastErrorD = Pointer Function(); +final cactusGetLastError = + _lib.lookupFunction<_CactusGetLastErrorN, _CactusGetLastErrorD>( + 'cactus_get_last_error'); + +typedef _CactusLogSetLevelN = Void Function(Int32 level); +typedef _CactusLogSetLevelD = void Function(int level); +final cactusLogSetLevel = + _lib.lookupFunction<_CactusLogSetLevelN, _CactusLogSetLevelD>( + 'cactus_log_set_level'); + +typedef _CactusLogSetCallbackN = Void Function( + Pointer> callback, + Pointer userData); +typedef _CactusLogSetCallbackD = void Function( + Pointer> callback, + Pointer userData); +final cactusLogSetCallback = + _lib.lookupFunction<_CactusLogSetCallbackN, _CactusLogSetCallbackD>( + 'cactus_log_set_callback'); + +typedef _CactusSetTelemetryEnvironmentN = Void Function( + Pointer framework, Pointer cacheLocation, + Pointer version); +typedef _CactusSetTelemetryEnvironmentD = void Function( + Pointer framework, Pointer cacheLocation, + Pointer version); +final cactusSetTelemetryEnvironment = _lib.lookupFunction< + _CactusSetTelemetryEnvironmentN, + _CactusSetTelemetryEnvironmentD>('cactus_set_telemetry_environment'); + +typedef _CactusSetAppIdN = Void Function(Pointer appId); +typedef _CactusSetAppIdD = void Function(Pointer appId); +final cactusSetAppId = + _lib.lookupFunction<_CactusSetAppIdN, _CactusSetAppIdD>( + 'cactus_set_app_id'); + +typedef _CactusTelemetryFlushN = Void Function(); +typedef _CactusTelemetryFlushD = void Function(); +final cactusTelemetryFlush = + _lib.lookupFunction<_CactusTelemetryFlushN, _CactusTelemetryFlushD>( + 'cactus_telemetry_flush'); + +typedef _CactusTelemetryShutdownN = Void Function(); +typedef _CactusTelemetryShutdownD = void Function(); +final cactusTelemetryShutdown = + _lib.lookupFunction<_CactusTelemetryShutdownN, _CactusTelemetryShutdownD>( + 'cactus_telemetry_shutdown'); diff --git a/bindings/kotlin/Cactus.android.kt b/bindings/kotlin/Cactus.android.kt new file mode 100644 index 000000000..b51eea024 --- /dev/null +++ b/bindings/kotlin/Cactus.android.kt @@ -0,0 +1,156 @@ +package com.cactus + +private fun check(rc: Int) { + if (rc < 0) throw RuntimeException(CactusJNI.nativeGetLastError().ifEmpty { "Unknown error" }) +} + +actual fun cactusInit(modelPath: String, corpusDir: String?, cacheIndex: Boolean): Long { + val handle = CactusJNI.nativeInit(modelPath, corpusDir, cacheIndex) + if (handle == 0L) throw RuntimeException(CactusJNI.nativeGetLastError().ifEmpty { "Failed to initialize model" }) + return handle +} + +actual fun cactusDestroy(handle: Long) = CactusJNI.nativeDestroy(handle) +actual fun cactusReset(handle: Long) = CactusJNI.nativeReset(handle) +actual fun cactusStop(handle: Long) = CactusJNI.nativeStop(handle) + +actual fun cactusComplete(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): String { + val buffer = ByteArray(1024 * 1024) + check(CactusJNI.nativeComplete(handle, messagesJson, buffer, optionsJson, toolsJson, callback, pcmData)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusPrefill(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, pcmData: ByteArray?): String { + val buffer = ByteArray(1024 * 1024) + check(CactusJNI.nativePrefill(handle, messagesJson, buffer, optionsJson, toolsJson, pcmData)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusTokenize(handle: Long, text: String): IntArray { + val tokenBuffer = IntArray(8192) + val outLen = LongArray(1) + check(CactusJNI.nativeTokenize(handle, text, tokenBuffer, outLen)) + return tokenBuffer.copyOf(outLen[0].toInt()) +} + +actual fun cactusScoreWindow(handle: Long, tokens: IntArray, start: Long, end: Long, context: Long): String { + val buffer = ByteArray(1024 * 1024) + check(CactusJNI.nativeScoreWindow(handle, tokens, start, end, context, buffer)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusTranscribe(handle: Long, audioPath: String?, prompt: String, optionsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): String { + val buffer = ByteArray(1024 * 1024) + check(CactusJNI.nativeTranscribe(handle, audioPath, prompt, buffer, optionsJson, callback, pcmData)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusStreamTranscribeStart(handle: Long, optionsJson: String?): Long { + val stream = CactusJNI.nativeStreamTranscribeStart(handle, optionsJson) + if (stream == 0L) throw RuntimeException(CactusJNI.nativeGetLastError().ifEmpty { "Failed to start streaming transcription" }) + return stream +} + +actual fun cactusStreamTranscribeProcess(stream: Long, pcmData: ByteArray?): String { + val buffer = ByteArray(65536) + check(CactusJNI.nativeStreamTranscribeProcess(stream, pcmData, buffer)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusStreamTranscribeStop(stream: Long): String { + val buffer = ByteArray(65536) + check(CactusJNI.nativeStreamTranscribeStop(stream, buffer)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusEmbed(handle: Long, text: String, normalize: Boolean): FloatArray { + val buffer = FloatArray(4096) + val outDim = LongArray(1) + check(CactusJNI.nativeEmbed(handle, text, buffer, outDim, normalize)) + return buffer.copyOf(outDim[0].toInt()) +} + +actual fun cactusImageEmbed(handle: Long, imagePath: String): FloatArray { + val buffer = FloatArray(4096) + val outDim = LongArray(1) + check(CactusJNI.nativeImageEmbed(handle, imagePath, buffer, outDim)) + return buffer.copyOf(outDim[0].toInt()) +} + +actual fun cactusAudioEmbed(handle: Long, audioPath: String): FloatArray { + val buffer = FloatArray(4096) + val outDim = LongArray(1) + check(CactusJNI.nativeAudioEmbed(handle, audioPath, buffer, outDim)) + return buffer.copyOf(outDim[0].toInt()) +} + +actual fun cactusRagQuery(handle: Long, query: String, topK: Long): String { + val buffer = ByteArray(1024 * 1024) + check(CactusJNI.nativeRagQuery(handle, query, buffer, topK)) + return buffer.decodeToString().trimEnd('\u0000') +} + +actual fun cactusIndexInit(indexDir: String, embeddingDim: Long): Long { + val handle = CactusJNI.nativeIndexInit(indexDir, embeddingDim) + if (handle == 0L) throw RuntimeException(CactusJNI.nativeGetLastError().ifEmpty { "Failed to initialize index" }) + return handle +} + +actual fun cactusIndexAdd(handle: Long, ids: IntArray, documents: Array, metadatas: Array?, embeddings: Array, embeddingDim: Long): Int { + val rc = CactusJNI.nativeIndexAdd(handle, ids, documents, metadatas, embeddings, embeddingDim) + check(rc) + return rc +} + +actual fun cactusIndexDelete(handle: Long, ids: IntArray): Int { + val rc = CactusJNI.nativeIndexDelete(handle, ids) + check(rc) + return rc +} + +actual fun cactusIndexGet(handle: Long, ids: IntArray): String { + val count = ids.size + val docBuffers = Array(count) { ByteArray(4096) } + val docSizes = LongArray(count) { 4096L } + val metaBuffers = Array(count) { ByteArray(4096) } + val metaSizes = LongArray(count) { 4096L } + val embBuffers = Array(count) { FloatArray(4096) } + val embSizes = LongArray(count) { 4096L } + check(CactusJNI.nativeIndexGet(handle, ids, docBuffers, docSizes, metaBuffers, metaSizes, embBuffers, embSizes)) + val results = ids.indices.map { i -> + val doc = docBuffers[i].decodeToString().trimEnd('\u0000') + val meta = metaBuffers[i].decodeToString().trimEnd('\u0000') + """{"document":"$doc","metadata":"$meta"}""" + } + return """{"results":[${results.joinToString(",")}]}""" +} + +actual fun cactusIndexQuery(handle: Long, embedding: FloatArray, optionsJson: String?): String { + val idBuffers = arrayOf(IntArray(1000)) + val idSizes = longArrayOf(1000L) + val scoreBuffers = arrayOf(FloatArray(1000)) + val scoreSizes = longArrayOf(1000L) + check(CactusJNI.nativeIndexQuery(handle, arrayOf(embedding), embedding.size.toLong(), optionsJson, idBuffers, idSizes, scoreBuffers, scoreSizes)) + val results = (0 until idSizes[0].toInt()).map { i -> + """{"id":${idBuffers[0][i]},"score":${scoreBuffers[0][i]}}""" + } + return """{"results":[${results.joinToString(",")}]}""" +} + +actual fun cactusIndexCompact(handle: Long): Int { + val rc = CactusJNI.nativeIndexCompact(handle) + check(rc) + return rc +} + +actual fun cactusIndexDestroy(handle: Long) = CactusJNI.nativeIndexDestroy(handle) +actual fun cactusGetLastError(): String = CactusJNI.nativeGetLastError() +actual fun cactusLogSetLevel(level: Int) = CactusJNI.nativeLogSetLevel(level) +actual fun cactusLogSetCallback(callback: CactusLogCallback?) = CactusJNI.nativeLogSetCallback(callback) + +actual fun cactusSetTelemetryEnvironment(framework: String?, cacheLocation: String?, version: String?) = + CactusJNI.nativeSetTelemetryEnvironment(framework, cacheLocation, version) + +actual fun cactusSetAppId(appId: String) = CactusJNI.nativeSetAppId(appId) +actual fun cactusTelemetryFlush() = CactusJNI.nativeTelemetryFlush() +actual fun cactusTelemetryShutdown() = CactusJNI.nativeTelemetryShutdown() diff --git a/bindings/kotlin/Cactus.common.kt b/bindings/kotlin/Cactus.common.kt new file mode 100644 index 000000000..dcc9be275 --- /dev/null +++ b/bindings/kotlin/Cactus.common.kt @@ -0,0 +1,32 @@ +package com.cactus + +expect fun cactusInit(modelPath: String, corpusDir: String?, cacheIndex: Boolean): Long +expect fun cactusDestroy(handle: Long) +expect fun cactusReset(handle: Long) +expect fun cactusStop(handle: Long) +expect fun cactusComplete(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): String +expect fun cactusPrefill(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, pcmData: ByteArray?): String +expect fun cactusTokenize(handle: Long, text: String): IntArray +expect fun cactusScoreWindow(handle: Long, tokens: IntArray, start: Long, end: Long, context: Long): String +expect fun cactusTranscribe(handle: Long, audioPath: String?, prompt: String, optionsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): String +expect fun cactusStreamTranscribeStart(handle: Long, optionsJson: String?): Long +expect fun cactusStreamTranscribeProcess(stream: Long, pcmData: ByteArray?): String +expect fun cactusStreamTranscribeStop(stream: Long): String +expect fun cactusEmbed(handle: Long, text: String, normalize: Boolean): FloatArray +expect fun cactusImageEmbed(handle: Long, imagePath: String): FloatArray +expect fun cactusAudioEmbed(handle: Long, audioPath: String): FloatArray +expect fun cactusRagQuery(handle: Long, query: String, topK: Long): String +expect fun cactusIndexInit(indexDir: String, embeddingDim: Long): Long +expect fun cactusIndexAdd(handle: Long, ids: IntArray, documents: Array, metadatas: Array?, embeddings: Array, embeddingDim: Long): Int +expect fun cactusIndexDelete(handle: Long, ids: IntArray): Int +expect fun cactusIndexGet(handle: Long, ids: IntArray): String +expect fun cactusIndexQuery(handle: Long, embedding: FloatArray, optionsJson: String?): String +expect fun cactusIndexCompact(handle: Long): Int +expect fun cactusIndexDestroy(handle: Long) +expect fun cactusGetLastError(): String +expect fun cactusLogSetLevel(level: Int) +expect fun cactusLogSetCallback(callback: CactusLogCallback?) +expect fun cactusSetTelemetryEnvironment(framework: String?, cacheLocation: String?, version: String?) +expect fun cactusSetAppId(appId: String) +expect fun cactusTelemetryFlush() +expect fun cactusTelemetryShutdown() diff --git a/bindings/kotlin/Cactus.ios.kt b/bindings/kotlin/Cactus.ios.kt new file mode 100644 index 000000000..98ad5ec07 --- /dev/null +++ b/bindings/kotlin/Cactus.ios.kt @@ -0,0 +1,357 @@ +package com.cactus + +import cactus.* +import kotlinx.cinterop.* + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusInit(modelPath: String, corpusDir: String?, cacheIndex: Boolean): Long { + val ptr = cactus_init(modelPath, corpusDir, cacheIndex) + ?: throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Failed to initialize model") + return ptr.rawValue.toLong() +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusDestroy(handle: Long) { + cactus_destroy(handle.toCPointer()) +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusReset(handle: Long) { cactus_reset(handle.toCPointer()) } + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusStop(handle: Long) { cactus_stop(handle.toCPointer()) } + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusComplete(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): String { + memScoped { + val buffer = allocArray(1048576) + val callbackRef = callback?.let { StableRef.create(it) } + val pcmPtr = pcmData?.refTo(0)?.getPointer(this) + try { + val result = cactus_complete( + handle.toCPointer(), + messagesJson, + buffer, + 1048576u, + optionsJson, + toolsJson, + callbackRef?.let { + staticCFunction?, UInt, COpaquePointer?, Unit> { token, tokenId, userData -> + if (token != null && userData != null) { + userData.asStableRef().get().onToken(token.toKString(), tokenId.toInt()) + } + } + }, + callbackRef?.asCPointer(), + pcmPtr?.reinterpret(), + pcmData?.size?.toULong() ?: 0u + ) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } finally { + callbackRef?.dispose() + } + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusPrefill(handle: Long, messagesJson: String, optionsJson: String?, toolsJson: String?, pcmData: ByteArray?): String { + memScoped { + val buffer = allocArray(1048576) + val pcmPtr = pcmData?.refTo(0)?.getPointer(this) + val result = cactus_prefill( + handle.toCPointer(), + messagesJson, + buffer, + 1048576u, + optionsJson, + toolsJson, + pcmPtr?.reinterpret(), + pcmData?.size?.toULong() ?: 0u + ) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusTokenize(handle: Long, text: String): IntArray { + memScoped { + val buffer = allocArray(8192) + val tokenLen = alloc() + val result = cactus_tokenize(handle.toCPointer(), text, buffer, 8192u, tokenLen.ptr) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return IntArray(tokenLen.value.toInt()) { buffer[it].toInt() } + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusScoreWindow(handle: Long, tokens: IntArray, start: Long, end: Long, context: Long): String { + memScoped { + val buffer = allocArray(1048576) + val tokenBuffer = allocArray(tokens.size) + tokens.forEachIndexed { i, v -> tokenBuffer[i] = v.toUInt() } + val result = cactus_score_window( + handle.toCPointer(), tokenBuffer, tokens.size.toULong(), + start.toULong(), end.toULong(), context.toULong(), + buffer, 1048576u + ) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusTranscribe(handle: Long, audioPath: String?, prompt: String, optionsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): String { + memScoped { + val buffer = allocArray(1048576) + val callbackRef = callback?.let { StableRef.create(it) } + val pcmPtr = pcmData?.refTo(0)?.getPointer(this) + try { + val result = cactus_transcribe( + handle.toCPointer(), + audioPath, + prompt, + buffer, + 1048576u, + optionsJson, + callbackRef?.let { + staticCFunction?, UInt, COpaquePointer?, Unit> { token, tokenId, userData -> + if (token != null && userData != null) { + userData.asStableRef().get().onToken(token.toKString(), tokenId.toInt()) + } + } + }, + callbackRef?.asCPointer(), + pcmPtr?.reinterpret(), + pcmData?.size?.toULong() ?: 0u + ) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } finally { + callbackRef?.dispose() + } + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusStreamTranscribeStart(handle: Long, optionsJson: String?): Long { + val stream = cactus_stream_transcribe_start(handle.toCPointer(), optionsJson) + ?: throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Failed to start streaming transcription") + return stream.rawValue.toLong() +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusStreamTranscribeProcess(stream: Long, pcmData: ByteArray?): String { + memScoped { + val buffer = allocArray(65536) + val pcmPtr = pcmData?.refTo(0)?.getPointer(this) + val result = cactus_stream_transcribe_process( + stream.toCPointer(), + pcmPtr?.reinterpret(), + pcmData?.size?.toULong() ?: 0u, + buffer, + 65536u + ) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusStreamTranscribeStop(stream: Long): String { + memScoped { + val buffer = allocArray(65536) + val result = cactus_stream_transcribe_stop(stream.toCPointer(), buffer, 65536u) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusEmbed(handle: Long, text: String, normalize: Boolean): FloatArray { + memScoped { + val buffer = allocArray(4096) + val dimPtr = alloc() + val result = cactus_embed(handle.toCPointer(), text, buffer, 4096u, dimPtr.ptr, normalize) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return FloatArray(dimPtr.value.toInt()) { buffer[it] } + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusImageEmbed(handle: Long, imagePath: String): FloatArray { + memScoped { + val buffer = allocArray(4096) + val dimPtr = alloc() + val result = cactus_image_embed(handle.toCPointer(), imagePath, buffer, 4096u, dimPtr.ptr) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return FloatArray(dimPtr.value.toInt()) { buffer[it] } + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusAudioEmbed(handle: Long, audioPath: String): FloatArray { + memScoped { + val buffer = allocArray(4096) + val dimPtr = alloc() + val result = cactus_audio_embed(handle.toCPointer(), audioPath, buffer, 4096u, dimPtr.ptr) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return FloatArray(dimPtr.value.toInt()) { buffer[it] } + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusRagQuery(handle: Long, query: String, topK: Long): String { + memScoped { + val buffer = allocArray(1048576) + val result = cactus_rag_query(handle.toCPointer(), query, buffer, 1048576u, topK.toULong()) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return buffer.toKString() + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexInit(indexDir: String, embeddingDim: Long): Long { + val ptr = cactus_index_init(indexDir, embeddingDim.toULong()) + ?: throw RuntimeException("Failed to initialize index") + return ptr.rawValue.toLong() +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexAdd(handle: Long, ids: IntArray, documents: Array, metadatas: Array?, embeddings: Array, embeddingDim: Long): Int { + memScoped { + val idPtr = allocArray(ids.size) + ids.forEachIndexed { i, v -> idPtr[i] = v } + val docPtrs = allocArray>(documents.size) + documents.forEachIndexed { i, doc -> docPtrs[i] = doc.cstr.ptr } + val metaPtrs = metadatas?.let { + val ptrs = allocArray>(it.size) + it.forEachIndexed { i, meta -> ptrs[i] = meta.cstr.ptr } + ptrs + } + val embPtrs = allocArray>(embeddings.size) + embeddings.forEachIndexed { i, emb -> + val embArr = allocArray(emb.size) + emb.forEachIndexed { j, v -> embArr[j] = v } + embPtrs[i] = embArr + } + val result = cactus_index_add( + handle.toCPointer(), idPtr, docPtrs, metaPtrs, embPtrs, + ids.size.toULong(), embeddingDim.toULong() + ) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return result + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexDelete(handle: Long, ids: IntArray): Int { + memScoped { + val idPtr = allocArray(ids.size) + ids.forEachIndexed { i, v -> idPtr[i] = v } + val result = cactus_index_delete(handle.toCPointer(), idPtr, ids.size.toULong()) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return result + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexGet(handle: Long, ids: IntArray): String { + memScoped { + val count = ids.size + val idPtr = allocArray(count) + ids.forEachIndexed { i, v -> idPtr[i] = v } + val docBufs = allocArray>(count) + val docSizes = allocArray(count) + val metaBufs = allocArray>(count) + val metaSizes = allocArray(count) + val embBufs = allocArray>(count) + val embSizes = allocArray(count) + val docAllocs = (0 until count).map { allocArray(4096) } + val metaAllocs = (0 until count).map { allocArray(4096) } + val embAllocs = (0 until count).map { allocArray(4096) } + for (i in 0 until count) { + docBufs[i] = docAllocs[i]; docSizes[i] = 4096u + metaBufs[i] = metaAllocs[i]; metaSizes[i] = 4096u + embBufs[i] = embAllocs[i]; embSizes[i] = 4096u + } + cactus_index_get(handle.toCPointer(), idPtr, count.toULong(), docBufs, docSizes, metaBufs, metaSizes, embBufs, embSizes) + val results = (0 until count).map { i -> + val doc = docAllocs[i].toKString() + val meta = metaAllocs[i].toKString() + """{"document":"$doc","metadata":"$meta"}""" + } + return """{"results":[${results.joinToString(",")}]}""" + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexQuery(handle: Long, embedding: FloatArray, optionsJson: String?): String { + memScoped { + val embArr = allocArray(embedding.size) + embedding.forEachIndexed { i, v -> embArr[i] = v } + val embPtrs = allocArray>(1) + embPtrs[0] = embArr + val idBuf = allocArray(1000) + val idPtrs = allocArray>(1) + idPtrs[0] = idBuf + val idSizes = allocArray(1); idSizes[0] = 1000u + val scoreBuf = allocArray(1000) + val scorePtrs = allocArray>(1) + scorePtrs[0] = scoreBuf + val scoreSizes = allocArray(1); scoreSizes[0] = 1000u + cactus_index_query(handle.toCPointer(), embPtrs, 1u, embedding.size.toULong(), optionsJson, idPtrs, idSizes, scorePtrs, scoreSizes) + val results = (0 until idSizes[0].toInt()).map { i -> + """{"id":${idBuf[i]},"score":${scoreBuf[i]}}""" + } + return """{"results":[${results.joinToString(",")}]}""" + } +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexCompact(handle: Long): Int { + val result = cactus_index_compact(handle.toCPointer()) + if (result < 0) throw RuntimeException(cactus_get_last_error()?.toKString() ?: "Unknown error") + return result +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusIndexDestroy(handle: Long) { + cactus_index_destroy(handle.toCPointer()) +} + +actual fun cactusGetLastError(): String = cactus_get_last_error()?.toKString() ?: "" + +actual fun cactusLogSetLevel(level: Int) { + cactus_log_set_level(level) +} + +private var _logCallbackRef: StableRef? = null + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusLogSetCallback(callback: CactusLogCallback?) { + _logCallbackRef?.dispose() + _logCallbackRef = null + if (callback == null) { + cactus_log_set_callback(null, null) + return + } + val ref = StableRef.create(callback) + _logCallbackRef = ref + cactus_log_set_callback(staticCFunction { level, component, message, userData -> + val cb = userData!!.asStableRef().get() + cb.onLog(level, component?.toKString() ?: "", message?.toKString() ?: "") + }, ref.asCPointer()) +} + +@OptIn(ExperimentalForeignApi::class) +actual fun cactusSetTelemetryEnvironment(framework: String?, cacheLocation: String?, version: String?) { + cactus_set_telemetry_environment(framework, cacheLocation, version) +} + +actual fun cactusSetAppId(appId: String) { cactus_set_app_id(appId) } + +actual fun cactusTelemetryFlush() { cactus_telemetry_flush() } + +actual fun cactusTelemetryShutdown() { cactus_telemetry_shutdown() } diff --git a/bindings/kotlin/Cactus.kt b/bindings/kotlin/Cactus.kt new file mode 100644 index 000000000..875076ade --- /dev/null +++ b/bindings/kotlin/Cactus.kt @@ -0,0 +1,38 @@ +package com.cactus + +object CactusJNI { + init { + System.loadLibrary("cactus_engine") + } + + @JvmStatic external fun nativeInit(modelPath: String, corpusDir: String?, cacheIndex: Boolean): Long + @JvmStatic external fun nativeDestroy(handle: Long) + @JvmStatic external fun nativeReset(handle: Long) + @JvmStatic external fun nativeStop(handle: Long) + @JvmStatic external fun nativeComplete(handle: Long, messagesJson: String, responseBuffer: ByteArray, optionsJson: String?, toolsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): Int + @JvmStatic external fun nativePrefill(handle: Long, messagesJson: String, responseBuffer: ByteArray, optionsJson: String?, toolsJson: String?, pcmData: ByteArray?): Int + @JvmStatic external fun nativeTokenize(handle: Long, text: String, tokenBuffer: IntArray, outTokenLen: LongArray): Int + @JvmStatic external fun nativeScoreWindow(handle: Long, tokens: IntArray, start: Long, end: Long, context: Long, responseBuffer: ByteArray): Int + @JvmStatic external fun nativeTranscribe(handle: Long, audioPath: String?, prompt: String, responseBuffer: ByteArray, optionsJson: String?, callback: CactusTokenCallback?, pcmData: ByteArray?): Int + @JvmStatic external fun nativeStreamTranscribeStart(handle: Long, optionsJson: String?): Long + @JvmStatic external fun nativeStreamTranscribeProcess(stream: Long, pcmData: ByteArray?, responseBuffer: ByteArray): Int + @JvmStatic external fun nativeStreamTranscribeStop(stream: Long, responseBuffer: ByteArray): Int + @JvmStatic external fun nativeEmbed(handle: Long, text: String, embeddingsBuffer: FloatArray, outEmbeddingDim: LongArray, normalize: Boolean): Int + @JvmStatic external fun nativeImageEmbed(handle: Long, imagePath: String, embeddingsBuffer: FloatArray, outEmbeddingDim: LongArray): Int + @JvmStatic external fun nativeAudioEmbed(handle: Long, audioPath: String, embeddingsBuffer: FloatArray, outEmbeddingDim: LongArray): Int + @JvmStatic external fun nativeRagQuery(handle: Long, query: String, responseBuffer: ByteArray, topK: Long): Int + @JvmStatic external fun nativeIndexInit(indexDir: String, embeddingDim: Long): Long + @JvmStatic external fun nativeIndexAdd(handle: Long, ids: IntArray, documents: Array, metadatas: Array?, embeddings: Array, embeddingDim: Long): Int + @JvmStatic external fun nativeIndexDelete(handle: Long, ids: IntArray): Int + @JvmStatic external fun nativeIndexGet(handle: Long, ids: IntArray, documentBuffers: Array, documentBufferSizes: LongArray, metadataBuffers: Array, metadataBufferSizes: LongArray, embeddingBuffers: Array, embeddingBufferSizes: LongArray): Int + @JvmStatic external fun nativeIndexQuery(handle: Long, embeddings: Array, embeddingDim: Long, optionsJson: String?, idBuffers: Array, idBufferSizes: LongArray, scoreBuffers: Array, scoreBufferSizes: LongArray): Int + @JvmStatic external fun nativeIndexCompact(handle: Long): Int + @JvmStatic external fun nativeIndexDestroy(handle: Long) + @JvmStatic external fun nativeGetLastError(): String + @JvmStatic external fun nativeLogSetLevel(level: Int) + @JvmStatic external fun nativeLogSetCallback(callback: CactusLogCallback?) + @JvmStatic external fun nativeSetTelemetryEnvironment(framework: String?, cacheLocation: String?, version: String?) + @JvmStatic external fun nativeSetAppId(appId: String) + @JvmStatic external fun nativeTelemetryFlush() + @JvmStatic external fun nativeTelemetryShutdown() +} diff --git a/bindings/kotlin/CactusCallbacks.kt b/bindings/kotlin/CactusCallbacks.kt new file mode 100644 index 000000000..c7e4f1811 --- /dev/null +++ b/bindings/kotlin/CactusCallbacks.kt @@ -0,0 +1,9 @@ +package com.cactus + +fun interface CactusTokenCallback { + fun onToken(token: String, tokenId: Int) +} + +fun interface CactusLogCallback { + fun onLog(level: Int, component: String, message: String) +} diff --git a/bindings/kotlin/README.md b/bindings/kotlin/README.md new file mode 100644 index 000000000..b0301543b --- /dev/null +++ b/bindings/kotlin/README.md @@ -0,0 +1,86 @@ +# Kotlin Bindings + +JNI bridge to `cactus_engine.h` for Android, with KMP support for iOS. + +The JNI bridge itself (`android/cactus_jni.cpp`, compiled into `libcactus_engine.so`) is JVM-language-agnostic. Java consumers can use it by writing an equivalent `CactusJNI.java` with `native` declarations matching the signatures in `Cactus.kt` — same library, same `Java_com_cactus_CactusJNI_*` symbols. + +## Android Integration + + +```bash +cactus build --android +``` + + + +1. Copy `android/libcactus_engine.so` to `app/src/main/jniLibs/arm64-v8a/` +2. Copy `Cactus.kt` and `CactusCallbacks.kt` to your Kotlin source tree + + + +```kotlin +val handle = CactusJNI.nativeInit("/path/to/model", null, false) +val buf = ByteArray(65536) +CactusJNI.nativeComplete(handle, messagesJson, buf, null, null, null, null) +val response = String(buf, 0, buf.indexOf(0)) +CactusJNI.nativeDestroy(handle) +``` + + +Streaming transcription — push 16 kHz mono PCM16, read `{"success":true,"confirmed":...,"pending":...}` back each call: + +```kotlin +val stream = CactusJNI.nativeStreamTranscribeStart(handle, "{\"language\":\"en\"}") +val out = ByteArray(65536) +for (chunk in pcmChunks) { // each: 16 kHz mono PCM16 bytes + CactusJNI.nativeStreamTranscribeProcess(stream, chunk, out) + // parse confirmed/pending from String(out, 0, out.indexOf(0)) +} +CactusJNI.nativeStreamTranscribeStop(stream, out) +``` + +## Kotlin Multiplatform + +``` +commonMain/ Cactus.common.kt expect declarations +commonMain/ CactusCallbacks.kt shared callback interfaces +androidMain/ Cactus.android.kt actual via JNI +iosMain/ Cactus.ios.kt actual via cinterop +``` + +### build.gradle.kts + +```kotlin +kotlin { + androidTarget() + + listOf(iosArm64(), iosSimulatorArm64()).forEach { target -> + target.compilations.getByName("main") { + cinterops { + val cactus by creating { + defFile("src/nativeInterop/cinterop/cactus.def") + includeDirs("/path/to/cactus/cactus-engine") + } + } + } + val libSuffix = if (target.name == "iosSimulatorArm64") "simulator" else "device" + target.binaries.framework { + linkerOpts("-L/path/to/cactus/apple", "-lcactus_engine-$libSuffix") + } + } +} +``` + +### Usage (shared code) + +```kotlin +val model = cactusInit("/path/to/model", null, false) +val result = cactusComplete(model, messagesJson, null, null, null, null) + +// streaming transcription: start, feed PCM16 chunks, stop +val stream = cactusStreamTranscribeStart(model, "{\"language\":\"en\"}") +val out = cactusStreamTranscribeProcess(stream, pcmChunk) // {"success":true,"confirmed":...,"pending":...} +val tail = cactusStreamTranscribeStop(stream) + +cactusDestroy(model) +``` diff --git a/bindings/kotlin/cactus.def b/bindings/kotlin/cactus.def new file mode 100644 index 000000000..32086abd6 --- /dev/null +++ b/bindings/kotlin/cactus.def @@ -0,0 +1,3 @@ +package = cactus +headers = cactus_engine.h +headerFilter = cactus_engine.h diff --git a/bindings/python/README.md b/bindings/python/README.md new file mode 100644 index 000000000..1963c98ed --- /dev/null +++ b/bindings/python/README.md @@ -0,0 +1,10 @@ +# Python Bindings + +ctypes FFI to `cactus_engine.h`. The binding module lives at +`/python/cactus/bindings/cactus.py`. + +It loads `libcactus_engine.{so,dylib}` from `/python/cactus/bindings/lib/` (bundled) +or `cactus-engine/build/` (dev build), populated by `cactus build --python`. + +For installation, the high-level API, the CLI, and the server, +see [`/python/README.md`](/python/). diff --git a/bindings/react-native/README.md b/bindings/react-native/README.md new file mode 100644 index 000000000..cd084b305 --- /dev/null +++ b/bindings/react-native/README.md @@ -0,0 +1,48 @@ +# React Native Bindings + +Native bridge modules over the cactus C API for iOS and Android. + +## Integration + + +```bash +cactus build --apple +cactus build --android +``` + + + +1. Add the bridge files from this folder to your React Native app +2. Add the raw bindings they depend on: [`bindings/kotlin/`](/bindings/kotlin/) (Android), [`bindings/swift/`](/bindings/swift/) (Apple) +3. Register the Android package in `MainApplication.kt`: + ```kotlin + override fun getPackages() = PackageList(this).packages.apply { + add(com.cactus.reactnative.CactusPackage()) + } + ``` + + +## Usage + + +```ts +import Cactus from './index'; + +const handle = await Cactus.init('/path/to/model', null, false); +const result = await Cactus.complete(handle, messagesJson, null, null, null, false); +await Cactus.destroy(handle); +``` + + +## Streaming transcription + +Push 16 kHz mono PCM16 (base64) as it arrives; each call returns `{"success":true,"confirmed":...,"pending":...}` (`confirmed` is final, `pending` is the volatile tail). + +```ts +const stream = await Cactus.streamTranscribeStart(handle, '{"language":"en"}'); +for (const chunkBase64 of pcmChunks) { + const json = await Cactus.streamTranscribeProcess(stream, chunkBase64); + // parse json -> append "confirmed", show "pending" as a live preview +} +const tail = await Cactus.streamTranscribeStop(stream); +``` diff --git a/bindings/react-native/android/CactusModule.kt b/bindings/react-native/android/CactusModule.kt new file mode 100644 index 000000000..534022141 --- /dev/null +++ b/bindings/react-native/android/CactusModule.kt @@ -0,0 +1,503 @@ +package com.cactus.reactnative + +import android.util.Base64 +import com.cactus.CactusJNI +import com.facebook.react.bridge.Arguments +import com.facebook.react.bridge.Promise +import com.facebook.react.bridge.ReactApplicationContext +import com.facebook.react.bridge.ReactContextBaseJavaModule +import com.facebook.react.bridge.ReactMethod +import com.facebook.react.bridge.ReadableArray +import com.facebook.react.modules.core.DeviceEventManagerModule +import com.cactus.CactusTokenCallback +import org.json.JSONArray +import org.json.JSONObject + +class CactusModule(reactContext: ReactApplicationContext) : ReactContextBaseJavaModule(reactContext) { + companion object { + private const val DEFAULT_BUFFER_SIZE = 65536 + private const val LARGE_BUFFER_SIZE = 1 shl 20 + private const val DEFAULT_EMBEDDING_BUFFER_SIZE = 4096 + private const val DEFAULT_TOKEN_BUFFER_SIZE = 8192 + private const val DEFAULT_INDEX_RESULT_CAPACITY = 1000 + private const val DEFAULT_INDEX_DOC_BUFFER_SIZE = 4096 + private const val DEFAULT_INDEX_EMBED_BUFFER_SIZE = 4096 + } + + override fun getName(): String = "Cactus" + + private fun parseHandle(handle: String, promise: Promise): Long? { + val parsed = handle.toLongOrNull() + if (parsed == null) { + promise.reject("CACTUS_ERROR", "Invalid native handle") + } + return parsed + } + + private fun fail(promise: Promise, defaultMessage: String) { + val message = CactusJNI.nativeGetLastError().ifEmpty { defaultMessage } + promise.reject("CACTUS_ERROR", message) + } + + private fun decodeNullTerminatedUtf8(buffer: ByteArray): String { + val end = buffer.indexOf(0).let { if (it >= 0) it else buffer.size } + return buffer.copyOf(end).toString(Charsets.UTF_8) + } + + private fun decodeBase64OrNull(value: String?): ByteArray? = + if (value == null) null else Base64.decode(value, Base64.DEFAULT) + + private fun readableArrayToIntArray(values: ReadableArray): IntArray = + IntArray(values.size()) { index -> values.getDouble(index).toInt() } + + private fun readableArrayToFloatArray(values: ReadableArray): FloatArray = + FloatArray(values.size()) { index -> values.getDouble(index).toFloat() } + + private fun readableArrayToStringArray(values: ReadableArray): Array = + Array(values.size()) { index -> values.getString(index) ?: "" } + + private fun readableNestedFloatArrays(values: ReadableArray): Array = + Array(values.size()) { index -> + readableArrayToFloatArray(values.getArray(index) ?: throw IllegalArgumentException("embeddings[$index] is null")) + } + + @ReactMethod + fun init(modelPath: String, corpusDir: String?, cacheIndex: Boolean, promise: Promise) { + val handle = CactusJNI.nativeInit(modelPath, corpusDir, cacheIndex) + if (handle == 0L) { + fail(promise, "Failed to initialize model") + return + } + promise.resolve(handle.toString()) + } + + @ReactMethod + fun destroy(handle: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + CactusJNI.nativeDestroy(nativeHandle) + promise.resolve(null) + } + + @ReactMethod + fun reset(handle: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + CactusJNI.nativeReset(nativeHandle) + promise.resolve(null) + } + + @ReactMethod + fun stop(handle: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + CactusJNI.nativeStop(nativeHandle) + promise.resolve(null) + } + + @ReactMethod + fun prefill( + handle: String, + messagesJson: String, + optionsJson: String?, + toolsJson: String?, + pcmDataBase64: String?, + promise: Promise, + ) { + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val nativeHandle = parseHandle(handle, promise) ?: return + val rc = CactusJNI.nativePrefill( + nativeHandle, + messagesJson, + buffer, + optionsJson, + toolsJson, + decodeBase64OrNull(pcmDataBase64), + ) + if (rc < 0) { + fail(promise, "Prefill failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + private fun emitToken(token: String, tokenId: Int) { + reactApplicationContext + .getJSModule(DeviceEventManagerModule.RCTDeviceEventEmitter::class.java) + .emit("onToken", Arguments.createMap().apply { + putString("token", token) + putInt("tokenId", tokenId) + }) + } + + @ReactMethod + fun addListener(@Suppress("UNUSED_PARAMETER") eventName: String) {} + + @ReactMethod + fun removeListeners(@Suppress("UNUSED_PARAMETER") count: Int) {} + + @ReactMethod + fun complete( + handle: String, + messagesJson: String, + optionsJson: String?, + toolsJson: String?, + pcmDataBase64: String?, + streamTokens: Boolean, + promise: Promise, + ) { + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val nativeHandle = parseHandle(handle, promise) ?: return + val callback = if (streamTokens) CactusTokenCallback { token, tokenId -> emitToken(token, tokenId) } else null + val rc = CactusJNI.nativeComplete( + nativeHandle, + messagesJson, + buffer, + optionsJson, + toolsJson, + callback, + decodeBase64OrNull(pcmDataBase64), + ) + if (rc < 0) { + fail(promise, "Completion failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + fun tokenize(handle: String, text: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val buffer = IntArray(DEFAULT_TOKEN_BUFFER_SIZE) + val outLen = LongArray(1) + val rc = CactusJNI.nativeTokenize(nativeHandle, text, buffer, outLen) + if (rc < 0) { + fail(promise, "Tokenization failed") + return + } + val result = Arguments.createArray() + for (i in 0 until outLen[0].toInt()) { + result.pushInt(buffer[i]) + } + promise.resolve(result) + } + + @ReactMethod + fun scoreWindow(handle: String, tokens: ReadableArray, start: Double, end: Double, context: Double, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val rc = CactusJNI.nativeScoreWindow( + nativeHandle, + readableArrayToIntArray(tokens), + start.toLong(), + end.toLong(), + context.toLong(), + buffer, + ) + if (rc < 0) { + fail(promise, "Score window failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + fun transcribe( + handle: String, + audioPath: String?, + prompt: String?, + optionsJson: String?, + pcmDataBase64: String?, + promise: Promise, + ) { + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val nativeHandle = parseHandle(handle, promise) ?: return + val rc = CactusJNI.nativeTranscribe( + nativeHandle, + audioPath, + prompt ?: "", + buffer, + optionsJson, + null, + decodeBase64OrNull(pcmDataBase64), + ) + if (rc < 0) { + fail(promise, "Transcription failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + fun streamTranscribeStart(handle: String, optionsJson: String?, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val stream = CactusJNI.nativeStreamTranscribeStart(nativeHandle, optionsJson) + if (stream == 0L) { + fail(promise, "Failed to start streaming transcription") + return + } + promise.resolve(stream.toString()) + } + + @ReactMethod + fun streamTranscribeProcess(streamHandle: String, pcmDataBase64: String?, promise: Promise) { + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val stream = parseHandle(streamHandle, promise) ?: return + val rc = CactusJNI.nativeStreamTranscribeProcess(stream, decodeBase64OrNull(pcmDataBase64), buffer) + if (rc < 0) { + fail(promise, "Stream transcription failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + fun streamTranscribeStop(streamHandle: String, promise: Promise) { + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val stream = parseHandle(streamHandle, promise) ?: return + val rc = CactusJNI.nativeStreamTranscribeStop(stream, buffer) + if (rc < 0) { + fail(promise, "Stream transcription stop failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + fun embed(handle: String, text: String, normalize: Boolean, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val buffer = FloatArray(DEFAULT_EMBEDDING_BUFFER_SIZE) + val outDim = LongArray(1) + val rc = CactusJNI.nativeEmbed(nativeHandle, text, buffer, outDim, normalize) + if (rc < 0) { + fail(promise, "Embedding failed") + return + } + val result = Arguments.createArray() + for (i in 0 until outDim[0].toInt()) { + result.pushDouble(buffer[i].toDouble()) + } + promise.resolve(result) + } + + @ReactMethod + fun imageEmbed(handle: String, imagePath: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val buffer = FloatArray(DEFAULT_EMBEDDING_BUFFER_SIZE) + val outDim = LongArray(1) + val rc = CactusJNI.nativeImageEmbed(nativeHandle, imagePath, buffer, outDim) + if (rc < 0) { + fail(promise, "Image embedding failed") + return + } + val result = Arguments.createArray() + for (i in 0 until outDim[0].toInt()) { + result.pushDouble(buffer[i].toDouble()) + } + promise.resolve(result) + } + + @ReactMethod + fun audioEmbed(handle: String, audioPath: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val buffer = FloatArray(DEFAULT_EMBEDDING_BUFFER_SIZE) + val outDim = LongArray(1) + val rc = CactusJNI.nativeAudioEmbed(nativeHandle, audioPath, buffer, outDim) + if (rc < 0) { + fail(promise, "Audio embedding failed") + return + } + val result = Arguments.createArray() + for (i in 0 until outDim[0].toInt()) { + result.pushDouble(buffer[i].toDouble()) + } + promise.resolve(result) + } + + @ReactMethod + fun ragQuery(handle: String, query: String, topK: Double, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val buffer = ByteArray(DEFAULT_BUFFER_SIZE) + val rc = CactusJNI.nativeRagQuery(nativeHandle, query, buffer, topK.toLong()) + if (rc < 0) { + fail(promise, "RAG query failed") + return + } + promise.resolve(decodeNullTerminatedUtf8(buffer)) + } + + @ReactMethod + fun indexInit(indexDir: String, embeddingDim: Double, promise: Promise) { + val handle = CactusJNI.nativeIndexInit(indexDir, embeddingDim.toLong()) + if (handle == 0L) { + fail(promise, "Failed to initialize index") + return + } + promise.resolve(handle.toString()) + } + + @ReactMethod + fun indexAdd( + handle: String, + ids: ReadableArray, + documents: ReadableArray, + embeddings: ReadableArray, + metadatas: ReadableArray?, + promise: Promise, + ) { + val nativeHandle = parseHandle(handle, promise) ?: return + val idsArray = readableArrayToIntArray(ids) + val docsArray = readableArrayToStringArray(documents) + val embArray = readableNestedFloatArrays(embeddings) + val metaArray = metadatas?.let { readableArrayToStringArray(it) } + val embeddingDim = if (embArray.isNotEmpty()) embArray[0].size.toLong() else 0L + val rc = CactusJNI.nativeIndexAdd(nativeHandle, idsArray, docsArray, metaArray, embArray, embeddingDim) + if (rc < 0) { + fail(promise, "Failed to add to index") + return + } + promise.resolve(null) + } + + @ReactMethod + fun indexDelete(handle: String, ids: ReadableArray, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val rc = CactusJNI.nativeIndexDelete(nativeHandle, readableArrayToIntArray(ids)) + if (rc < 0) { + fail(promise, "Failed to delete from index") + return + } + promise.resolve(null) + } + + @ReactMethod + fun indexGet(handle: String, ids: ReadableArray, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val idArray = readableArrayToIntArray(ids) + val count = idArray.size + val docBuffers = Array(count) { ByteArray(DEFAULT_INDEX_DOC_BUFFER_SIZE) } + val docSizes = LongArray(count) { DEFAULT_INDEX_DOC_BUFFER_SIZE.toLong() } + val metaBuffers = Array(count) { ByteArray(DEFAULT_INDEX_DOC_BUFFER_SIZE) } + val metaSizes = LongArray(count) { DEFAULT_INDEX_DOC_BUFFER_SIZE.toLong() } + val embBuffers = Array(count) { FloatArray(DEFAULT_INDEX_EMBED_BUFFER_SIZE) } + val embSizes = LongArray(count) { DEFAULT_INDEX_EMBED_BUFFER_SIZE.toLong() } + val rc = CactusJNI.nativeIndexGet( + nativeHandle, + idArray, + docBuffers, + docSizes, + metaBuffers, + metaSizes, + embBuffers, + embSizes, + ) + if (rc < 0) { + fail(promise, "Failed to get from index") + return + } + val results = JSONArray() + for (i in 0 until count) { + val item = JSONObject() + val document = decodeNullTerminatedUtf8(docBuffers[i]) + val metadataRaw = decodeNullTerminatedUtf8(metaBuffers[i]) + val embDim = embSizes[i].toInt() + val embeddingJson = JSONArray() + for (value in embBuffers[i].copyOf(embDim)) { + embeddingJson.put(value.toDouble()) + } + item.put("document", document) + if (metadataRaw.isEmpty()) { + item.put("metadata", JSONObject.NULL) + } else { + item.put("metadata", metadataRaw) + } + item.put("embedding", embeddingJson) + results.put(item) + } + promise.resolve(JSONObject().put("results", results).toString()) + } + + @ReactMethod + fun indexQuery(handle: String, embedding: ReadableArray, optionsJson: String?, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val emb = readableArrayToFloatArray(embedding) + val idBuffers = arrayOf(IntArray(DEFAULT_INDEX_RESULT_CAPACITY)) + val idSizes = LongArray(1) { DEFAULT_INDEX_RESULT_CAPACITY.toLong() } + val scoreBuffers = arrayOf(FloatArray(DEFAULT_INDEX_RESULT_CAPACITY)) + val scoreSizes = LongArray(1) { DEFAULT_INDEX_RESULT_CAPACITY.toLong() } + val rc = CactusJNI.nativeIndexQuery( + nativeHandle, + arrayOf(emb), + emb.size.toLong(), + optionsJson, + idBuffers, + idSizes, + scoreBuffers, + scoreSizes, + ) + if (rc < 0) { + fail(promise, "Index query failed") + return + } + val results = JSONArray() + for (i in 0 until idSizes[0].toInt()) { + results.put( + JSONObject() + .put("id", idBuffers[0][i]) + .put("score", scoreBuffers[0][i].toDouble()), + ) + } + promise.resolve(JSONObject().put("results", results).toString()) + } + + @ReactMethod + fun indexCompact(handle: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + val rc = CactusJNI.nativeIndexCompact(nativeHandle) + if (rc < 0) { + fail(promise, "Failed to compact index") + return + } + promise.resolve(null) + } + + @ReactMethod + fun indexDestroy(handle: String, promise: Promise) { + val nativeHandle = parseHandle(handle, promise) ?: return + CactusJNI.nativeIndexDestroy(nativeHandle) + promise.resolve(null) + } + + @ReactMethod + fun logSetLevel(level: Double, promise: Promise) { + CactusJNI.nativeLogSetLevel(level.toInt()) + promise.resolve(null) + } + + @ReactMethod + fun setTelemetryEnvironment(framework: String?, cacheLocation: String?, version: String?, promise: Promise) { + CactusJNI.nativeSetTelemetryEnvironment(framework, cacheLocation, version) + promise.resolve(null) + } + + @ReactMethod + fun setAppId(appId: String, promise: Promise) { + CactusJNI.nativeSetAppId(appId) + promise.resolve(null) + } + + @ReactMethod + fun telemetryFlush(promise: Promise) { + CactusJNI.nativeTelemetryFlush() + promise.resolve(null) + } + + @ReactMethod + fun telemetryShutdown(promise: Promise) { + CactusJNI.nativeTelemetryShutdown() + promise.resolve(null) + } + + @ReactMethod + fun getLastError(promise: Promise) { + promise.resolve(CactusJNI.nativeGetLastError()) + } +} diff --git a/bindings/react-native/android/CactusPackage.kt b/bindings/react-native/android/CactusPackage.kt new file mode 100644 index 000000000..066db81ab --- /dev/null +++ b/bindings/react-native/android/CactusPackage.kt @@ -0,0 +1,14 @@ +package com.cactus.reactnative + +import com.facebook.react.ReactPackage +import com.facebook.react.bridge.NativeModule +import com.facebook.react.bridge.ReactApplicationContext +import com.facebook.react.uimanager.ViewManager + +class CactusPackage : ReactPackage { + override fun createNativeModules(reactContext: ReactApplicationContext): List = + listOf(CactusModule(reactContext)) + + override fun createViewManagers(reactContext: ReactApplicationContext): List> = + emptyList() +} diff --git a/bindings/react-native/apple/Cactus.m b/bindings/react-native/apple/Cactus.m new file mode 100644 index 000000000..cbc31f1a2 --- /dev/null +++ b/bindings/react-native/apple/Cactus.m @@ -0,0 +1,41 @@ +#import +#import + +@interface RCT_EXTERN_MODULE(Cactus, RCTEventEmitter) + +RCT_EXTERN_METHOD(init:(NSString *)modelPath corpusDir:(NSString * _Nullable)corpusDir cacheIndex:(BOOL)cacheIndex resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(destroy:(NSString *)handle resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(reset:(NSString *)handle resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(stop:(NSString *)handle resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(prefill:(NSString *)handle messagesJson:(NSString *)messagesJson optionsJson:(NSString * _Nullable)optionsJson toolsJson:(NSString * _Nullable)toolsJson pcmDataBase64:(NSString * _Nullable)pcmDataBase64 resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(complete:(NSString *)handle messagesJson:(NSString *)messagesJson optionsJson:(NSString * _Nullable)optionsJson toolsJson:(NSString * _Nullable)toolsJson pcmDataBase64:(NSString * _Nullable)pcmDataBase64 streamTokens:(BOOL)streamTokens resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(tokenize:(NSString *)handle text:(NSString *)text resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(scoreWindow:(NSString *)handle tokens:(NSArray *)tokens start:(nonnull NSNumber *)start end:(nonnull NSNumber *)end context:(nonnull NSNumber *)context resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(transcribe:(NSString *)handle audioPath:(NSString * _Nullable)audioPath prompt:(NSString * _Nullable)prompt optionsJson:(NSString * _Nullable)optionsJson pcmDataBase64:(NSString * _Nullable)pcmDataBase64 resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(streamTranscribeStart:(NSString *)handle optionsJson:(NSString * _Nullable)optionsJson resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(streamTranscribeProcess:(NSString *)streamHandle pcmDataBase64:(NSString * _Nullable)pcmDataBase64 resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(streamTranscribeStop:(NSString *)streamHandle resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(embed:(NSString *)handle text:(NSString *)text normalize:(BOOL)normalize resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(imageEmbed:(NSString *)handle imagePath:(NSString *)imagePath resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(audioEmbed:(NSString *)handle audioPath:(NSString *)audioPath resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(ragQuery:(NSString *)handle query:(NSString *)query topK:(nonnull NSNumber *)topK resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexInit:(NSString *)indexDir embeddingDim:(nonnull NSNumber *)embeddingDim resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexAdd:(NSString *)handle ids:(NSArray *)ids documents:(NSArray *)documents embeddings:(NSArray *> *)embeddings metadatas:(NSArray * _Nullable)metadatas resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexDelete:(NSString *)handle ids:(NSArray *)ids resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexGet:(NSString *)handle ids:(NSArray *)ids resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexQuery:(NSString *)handle embedding:(NSArray *)embedding optionsJson:(NSString * _Nullable)optionsJson resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexCompact:(NSString *)handle resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(indexDestroy:(NSString *)handle resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(logSetLevel:(nonnull NSNumber *)level resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(setTelemetryEnvironment:(NSString * _Nullable)framework cacheLocation:(NSString * _Nullable)cacheLocation version:(NSString * _Nullable)version resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(setAppId:(NSString *)appId resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(telemetryFlush:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(telemetryShutdown:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) +RCT_EXTERN_METHOD(getLastError:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) + ++ (BOOL)requiresMainQueueSetup +{ + return NO; +} + +@end diff --git a/bindings/react-native/apple/Cactus.swift b/bindings/react-native/apple/Cactus.swift new file mode 100644 index 000000000..7a49191c3 --- /dev/null +++ b/bindings/react-native/apple/Cactus.swift @@ -0,0 +1,635 @@ +import Foundation +import React +import cactus + +@objc(Cactus) +final class Cactus: RCTEventEmitter { + private let defaultBufferSize = 65_536 + private let largeBufferSize = 1 << 20 + private let defaultEmbeddingBufferSize = 4096 + private let defaultTokenBufferSize = 8192 + private let defaultIndexResultCapacity = 1000 + private let defaultIndexDocBufferSize = 4096 + private let defaultIndexEmbeddingBufferSize = 4096 + + @objc + override static func requiresMainQueueSetup() -> Bool { + false + } + + override func supportedEvents() -> [String] { + ["onToken"] + } + + private func lastError(_ fallback: String) -> String { + guard let ptr = cactus_get_last_error() else { return fallback } + return String(cString: ptr) + } + + private func reject(_ rejecter: RCTPromiseRejectBlock, _ fallback: String) { + rejecter("CACTUS_ERROR", lastError(fallback), nil) + } + + private func decodeBase64(_ value: String?) -> Data? { + guard let value else { return nil } + return Data(base64Encoded: value) + } + + private func stringFromBuffer(_ buffer: [CChar]) -> String { + String(cString: buffer) + } + + private func encodeHandle(_ handle: UnsafeMutableRawPointer) -> String { + String(UInt(bitPattern: handle)) + } + + private func decodeHandle(_ handle: String) -> UnsafeMutableRawPointer? { + guard let bits = UInt(handle) else { return nil } + return UnsafeMutableRawPointer(bitPattern: bits) + } + + @objc(init:corpusDir:cacheIndex:resolver:rejecter:) + func initialize(_ modelPath: String, corpusDir: String?, cacheIndex: Bool, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + let handle = cactus_init(modelPath, corpusDir, cacheIndex) + guard let handle else { + self.reject(reject, "Failed to initialize model") + return + } + resolve(encodeHandle(handle)) + } + + @objc(destroy:resolver:rejecter:) + func destroy(_ handle: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + cactus_destroy(nativeHandle) + resolve(nil) + } + + @objc(reset:resolver:rejecter:) + func reset(_ handle: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + cactus_reset(nativeHandle) + resolve(nil) + } + + @objc(stop:resolver:rejecter:) + func stop(_ handle: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + cactus_stop(nativeHandle) + resolve(nil) + } + + @objc(prefill:messagesJson:optionsJson:toolsJson:pcmDataBase64:resolver:rejecter:) + func prefill( + _ handle: String, + messagesJson: String, + optionsJson: String?, + toolsJson: String?, + pcmDataBase64: String?, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + let rc = decodeBase64(pcmDataBase64)?.withUnsafeBytes { rawBuffer in + cactus_prefill(nativeHandle, messagesJson, &buffer, buffer.count, optionsJson, toolsJson, rawBuffer.bindMemory(to: UInt8.self).baseAddress, rawBuffer.count) + } ?? cactus_prefill(nativeHandle, messagesJson, &buffer, buffer.count, optionsJson, toolsJson, nil, 0) + guard rc >= 0 else { + self.reject(reject, "Prefill failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(complete:messagesJson:optionsJson:toolsJson:pcmDataBase64:streamTokens:resolver:rejecter:) + func complete( + _ handle: String, + messagesJson: String, + optionsJson: String?, + toolsJson: String?, + pcmDataBase64: String?, + streamTokens: Bool, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + + var cb: cactus_token_callback? = nil + var userData: UnsafeMutableRawPointer? = nil + if streamTokens { + let emitter = Unmanaged.passRetained(self).toOpaque() + userData = emitter + cb = { tokenPtr, tokenId, ctx in + guard let ctx = ctx, let tokenPtr = tokenPtr else { return } + let emitter = Unmanaged.fromOpaque(ctx).takeUnretainedValue() + let token = String(cString: tokenPtr) + emitter.sendEvent(withName: "onToken", body: ["token": token, "tokenId": tokenId]) + } + } + + let rc = decodeBase64(pcmDataBase64)?.withUnsafeBytes { rawBuffer in + cactus_complete(nativeHandle, messagesJson, &buffer, buffer.count, optionsJson, toolsJson, cb, userData, rawBuffer.bindMemory(to: UInt8.self).baseAddress, rawBuffer.count) + } ?? cactus_complete(nativeHandle, messagesJson, &buffer, buffer.count, optionsJson, toolsJson, cb, userData, nil, 0) + + if streamTokens, let userData = userData { + Unmanaged.fromOpaque(userData).release() + } + + guard rc >= 0 else { + self.reject(reject, "Completion failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(tokenize:text:resolver:rejecter:) + func tokenize(_ handle: String, text: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var buffer = [UInt32](repeating: 0, count: defaultTokenBufferSize) + var outLen: Int = 0 + let rc = cactus_tokenize(nativeHandle, text, &buffer, buffer.count, &outLen) + guard rc >= 0 else { + self.reject(reject, "Tokenization failed") + return + } + resolve(buffer.prefix(outLen).map(Int.init)) + } + + @objc(scoreWindow:tokens:start:end:context:resolver:rejecter:) + func scoreWindow( + _ handle: String, + tokens: [NSNumber], + start: NSNumber, + end: NSNumber, + context: NSNumber, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var tokenValues = tokens.map { UInt32(truncating: $0) } + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + let rc = cactus_score_window(nativeHandle, &tokenValues, tokenValues.count, start.intValue, end.intValue, context.intValue, &buffer, buffer.count) + guard rc >= 0 else { + self.reject(reject, "Score window failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(transcribe:audioPath:prompt:optionsJson:pcmDataBase64:resolver:rejecter:) + func transcribe( + _ handle: String, + audioPath: String?, + prompt: String?, + optionsJson: String?, + pcmDataBase64: String?, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + let promptValue = prompt ?? "" + let rc = decodeBase64(pcmDataBase64)?.withUnsafeBytes { rawBuffer in + cactus_transcribe(nativeHandle, audioPath, promptValue, &buffer, buffer.count, optionsJson, nil, nil, rawBuffer.bindMemory(to: UInt8.self).baseAddress, rawBuffer.count) + } ?? cactus_transcribe(nativeHandle, audioPath, promptValue, &buffer, buffer.count, optionsJson, nil, nil, nil, 0) + guard rc >= 0 else { + self.reject(reject, "Transcription failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(streamTranscribeStart:optionsJson:resolver:rejecter:) + func streamTranscribeStart( + _ handle: String, + optionsJson: String?, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + guard let stream = cactus_stream_transcribe_start(nativeHandle, optionsJson) else { + self.reject(reject, "Failed to start streaming transcription") + return + } + resolve(encodeHandle(stream)) + } + + @objc(streamTranscribeProcess:pcmDataBase64:resolver:rejecter:) + func streamTranscribeProcess( + _ streamHandle: String, + pcmDataBase64: String?, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + guard let stream = decodeHandle(streamHandle) else { + self.reject(reject, "Invalid stream handle") + return + } + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + let rc = decodeBase64(pcmDataBase64)?.withUnsafeBytes { rawBuffer in + cactus_stream_transcribe_process(stream, rawBuffer.bindMemory(to: UInt8.self).baseAddress, rawBuffer.count, &buffer, buffer.count) + } ?? cactus_stream_transcribe_process(stream, nil, 0, &buffer, buffer.count) + guard rc >= 0 else { + self.reject(reject, "Stream transcription failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(streamTranscribeStop:resolver:rejecter:) + func streamTranscribeStop( + _ streamHandle: String, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + guard let stream = decodeHandle(streamHandle) else { + self.reject(reject, "Invalid stream handle") + return + } + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + let rc = cactus_stream_transcribe_stop(stream, &buffer, buffer.count) + guard rc >= 0 else { + self.reject(reject, "Stream transcription stop failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(embed:text:normalize:resolver:rejecter:) + func embed(_ handle: String, text: String, normalize: Bool, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var buffer = [Float](repeating: 0, count: defaultEmbeddingBufferSize) + var dim: Int = 0 + let rc = cactus_embed(nativeHandle, text, &buffer, buffer.count, &dim, normalize) + guard rc >= 0 else { + self.reject(reject, "Embedding failed") + return + } + resolve(buffer.prefix(dim).map(Double.init)) + } + + @objc(imageEmbed:imagePath:resolver:rejecter:) + func imageEmbed(_ handle: String, imagePath: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var buffer = [Float](repeating: 0, count: defaultEmbeddingBufferSize) + var dim: Int = 0 + let rc = cactus_image_embed(nativeHandle, imagePath, &buffer, buffer.count, &dim) + guard rc >= 0 else { + self.reject(reject, "Image embedding failed") + return + } + resolve(buffer.prefix(dim).map(Double.init)) + } + + @objc(audioEmbed:audioPath:resolver:rejecter:) + func audioEmbed(_ handle: String, audioPath: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var buffer = [Float](repeating: 0, count: defaultEmbeddingBufferSize) + var dim: Int = 0 + let rc = cactus_audio_embed(nativeHandle, audioPath, &buffer, buffer.count, &dim) + guard rc >= 0 else { + self.reject(reject, "Audio embedding failed") + return + } + resolve(buffer.prefix(dim).map(Double.init)) + } + + @objc(ragQuery:query:topK:resolver:rejecter:) + func ragQuery(_ handle: String, query: String, topK: NSNumber, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var buffer = [CChar](repeating: 0, count: defaultBufferSize) + let rc = cactus_rag_query(nativeHandle, query, &buffer, buffer.count, topK.intValue) + guard rc >= 0 else { + self.reject(reject, "RAG query failed") + return + } + resolve(stringFromBuffer(buffer)) + } + + @objc(indexInit:embeddingDim:resolver:rejecter:) + func indexInit(_ indexDir: String, embeddingDim: NSNumber, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + let handle = cactus_index_init(indexDir, embeddingDim.intValue) + guard let handle else { + self.reject(reject, "Failed to initialize index") + return + } + resolve(encodeHandle(handle)) + } + + @objc(indexAdd:ids:documents:embeddings:metadatas:resolver:rejecter:) + func indexAdd( + _ handle: String, + ids: [NSNumber], + documents: [String], + embeddings: [[NSNumber]], + metadatas: [String]?, + resolver resolve: RCTPromiseResolveBlock, + rejecter reject: RCTPromiseRejectBlock + ) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var idValues = ids.map { Int32(truncating: $0) } + guard documents.count == idValues.count, embeddings.count == idValues.count, metadatas == nil || metadatas!.count == idValues.count else { + self.reject(reject, "ids, documents, embeddings and metadatas must have equal length") + return + } + let embeddingDim = embeddings.first?.count ?? 0 + guard embeddings.allSatisfy({ $0.count == embeddingDim }) else { + self.reject(reject, "Embedding rows must all have the same length") + return + } + var docStrings = documents.map { strdup($0) } + var docPointers = docStrings.map { UnsafePointer($0) } + var metaStrings = metadatas?.map { strdup($0) } + var metaPointers = metaStrings?.map { UnsafePointer($0) } + var embeddingPointers = embeddings.map { row -> UnsafeMutablePointer in + let ptr = UnsafeMutablePointer.allocate(capacity: row.count) + for (index, value) in row.enumerated() { + ptr[index] = Float(truncating: value) + } + return ptr + } + var embeddingConstPointers: [UnsafePointer?] = embeddingPointers.map { UnsafePointer($0) } + defer { + docStrings.forEach { free($0) } + metaStrings?.forEach { free($0) } + embeddingPointers.forEach { $0.deallocate() } + } + var idsCopy = idValues + let rc = idsCopy.withUnsafeMutableBufferPointer { idBuffer in + docPointers.withUnsafeMutableBufferPointer { docBuffer in + embeddingConstPointers.withUnsafeMutableBufferPointer { embBuffer in + if var metaPointers { + return metaPointers.withUnsafeMutableBufferPointer { metaBuffer in + cactus_index_add( + nativeHandle, + idBuffer.baseAddress, + docBuffer.baseAddress, + metaBuffer.baseAddress, + embBuffer.baseAddress, + idValues.count, + embeddingDim + ) + } + } + return cactus_index_add( + nativeHandle, + idBuffer.baseAddress, + docBuffer.baseAddress, + nil, + embBuffer.baseAddress, + idValues.count, + embeddingDim + ) + } + } + } + guard rc >= 0 else { + self.reject(reject, "Failed to add to index") + return + } + resolve(nil) + } + + @objc(indexDelete:ids:resolver:rejecter:) + func indexDelete(_ handle: String, ids: [NSNumber], resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var idValues = ids.map { Int32(truncating: $0) } + let rc = cactus_index_delete(nativeHandle, &idValues, idValues.count) + guard rc >= 0 else { + self.reject(reject, "Failed to delete from index") + return + } + resolve(nil) + } + + @objc(indexGet:ids:resolver:rejecter:) + func indexGet(_ handle: String, ids: [NSNumber], resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + var idValues = ids.map { Int32(truncating: $0) } + let count = idValues.count + var docPointers: [UnsafeMutablePointer?] = (0...allocate(capacity: defaultIndexDocBufferSize) } + var docSizes = Array(repeating: defaultIndexDocBufferSize, count: count) + var metaPointers: [UnsafeMutablePointer?] = (0...allocate(capacity: defaultIndexDocBufferSize) } + var metaSizes = Array(repeating: defaultIndexDocBufferSize, count: count) + var embPointers: [UnsafeMutablePointer?] = (0...allocate(capacity: defaultIndexEmbeddingBufferSize) } + var embSizes = Array(repeating: defaultIndexEmbeddingBufferSize, count: count) + defer { + docPointers.forEach { $0?.deallocate() } + metaPointers.forEach { $0?.deallocate() } + embPointers.forEach { $0?.deallocate() } + } + for i in 0..= 0 else { + self.reject(reject, "Failed to get from index") + return + } + var results = [[String: Any]]() + for i in 0..? = UnsafePointer(embBuffer.baseAddress) + return idStorage.withUnsafeMutableBufferPointer { idBuffer in + var idPointer: UnsafeMutablePointer? = idBuffer.baseAddress + return scoreStorage.withUnsafeMutableBufferPointer { scoreBuffer in + var scorePointer: UnsafeMutablePointer? = scoreBuffer.baseAddress + return cactus_index_query( + nativeHandle, + &embPointer, + 1, + embValues.count, + optionsJson, + &idPointer, + &idSizes, + &scorePointer, + &scoreSizes + ) + } + } + } + guard rc >= 0 else { + self.reject(reject, "Index query failed") + return + } + let results = (0..= 0 else { + self.reject(reject, "Failed to compact index") + return + } + resolve(nil) + } + + @objc(indexDestroy:resolver:rejecter:) + func indexDestroy(_ handle: String, resolver resolve: RCTPromiseResolveBlock, rejecter reject: RCTPromiseRejectBlock) { + guard let nativeHandle = decodeHandle(handle) else { + self.reject(reject, "Invalid native handle") + return + } + cactus_index_destroy(nativeHandle) + resolve(nil) + } + + @objc(logSetLevel:resolver:rejecter:) + func logSetLevel(_ level: NSNumber, resolver resolve: RCTPromiseResolveBlock, rejecter _: RCTPromiseRejectBlock) { + cactus_log_set_level(level.int32Value) + resolve(nil) + } + + @objc(setTelemetryEnvironment:cacheLocation:version:resolver:rejecter:) + func setTelemetryEnvironment( + _ framework: String?, + cacheLocation: String?, + version: String?, + resolver resolve: RCTPromiseResolveBlock, + rejecter _: RCTPromiseRejectBlock + ) { + cactus_set_telemetry_environment(framework, cacheLocation, version) + resolve(nil) + } + + @objc(setAppId:resolver:rejecter:) + func setAppId(_ appId: String, resolver resolve: RCTPromiseResolveBlock, rejecter _: RCTPromiseRejectBlock) { + cactus_set_app_id(appId) + resolve(nil) + } + + @objc(telemetryFlush:rejecter:) + func telemetryFlush(_ resolve: RCTPromiseResolveBlock, rejecter _: RCTPromiseRejectBlock) { + cactus_telemetry_flush() + resolve(nil) + } + + @objc(telemetryShutdown:rejecter:) + func telemetryShutdown(_ resolve: RCTPromiseResolveBlock, rejecter _: RCTPromiseRejectBlock) { + cactus_telemetry_shutdown() + resolve(nil) + } + + @objc(getLastError:rejecter:) + func getLastError(_ resolve: RCTPromiseResolveBlock, rejecter _: RCTPromiseRejectBlock) { + guard let ptr = cactus_get_last_error() else { + resolve(nil) + return + } + resolve(String(cString: ptr)) + } +} diff --git a/bindings/react-native/index.ts b/bindings/react-native/index.ts new file mode 100644 index 000000000..5c5c64d20 --- /dev/null +++ b/bindings/react-native/index.ts @@ -0,0 +1,89 @@ +import {NativeModules, NativeEventEmitter} from 'react-native'; + +export interface CactusNativeModule { + init(modelPath: string, corpusDir: string | null, cacheIndex: boolean): Promise; + destroy(handle: string): Promise; + reset(handle: string): Promise; + stop(handle: string): Promise; + prefill( + handle: string, + messagesJson: string, + optionsJson: string | null, + toolsJson: string | null, + pcmDataBase64: string | null, + ): Promise; + complete( + handle: string, + messagesJson: string, + optionsJson: string | null, + toolsJson: string | null, + pcmDataBase64: string | null, + streamTokens: boolean, + ): Promise; + tokenize(handle: string, text: string): Promise; + scoreWindow( + handle: string, + tokens: number[], + start: number, + end: number, + context: number, + ): Promise; + transcribe( + handle: string, + audioPath: string | null, + prompt: string | null, + optionsJson: string | null, + pcmDataBase64: string | null, + ): Promise; + streamTranscribeStart(handle: string, optionsJson: string | null): Promise; + streamTranscribeProcess( + streamHandle: string, + pcmDataBase64: string | null, + ): Promise; + streamTranscribeStop(streamHandle: string): Promise; + embed(handle: string, text: string, normalize: boolean): Promise; + imageEmbed(handle: string, imagePath: string): Promise; + audioEmbed(handle: string, audioPath: string): Promise; + ragQuery(handle: string, query: string, topK: number): Promise; + indexInit(indexDir: string, embeddingDim: number): Promise; + indexAdd( + handle: string, + ids: number[], + documents: string[], + embeddings: number[][], + metadatas: string[] | null, + ): Promise; + indexDelete(handle: string, ids: number[]): Promise; + indexGet(handle: string, ids: number[]): Promise; + indexQuery(handle: string, embedding: number[], optionsJson: string | null): Promise; + indexCompact(handle: string): Promise; + indexDestroy(handle: string): Promise; + logSetLevel(level: number): Promise; + setTelemetryEnvironment( + framework: string | null, + cacheLocation: string | null, + version: string | null, + ): Promise; + setAppId(appId: string): Promise; + telemetryFlush(): Promise; + telemetryShutdown(): Promise; + getLastError(): Promise; +} + +const LINKING_ERROR = + 'Cactus native module not linked. ' + + 'Add the bindings/react-native native bridge files and register NativeModules.Cactus.'; + +const nativeModule = NativeModules.Cactus as CactusNativeModule | undefined; + +if (!nativeModule) { + throw new Error(LINKING_ERROR); +} + +export const Cactus = nativeModule; + +export const CactusEvents = new NativeEventEmitter(NativeModules.Cactus); + +export type TokenEvent = {token: string; tokenId: number}; + +export default Cactus; diff --git a/bindings/rust/README.md b/bindings/rust/README.md new file mode 100644 index 000000000..f264dd51f --- /dev/null +++ b/bindings/rust/README.md @@ -0,0 +1,67 @@ +# Rust Bindings + +Raw `extern "C"` declarations for `cactus_engine.h`. + +## Integration + + +```bash +cactus build +``` + + +Copy `cactus.rs` into your project (it carries its own `#[link]` +attributes) and point Cargo at the build directory: + +```rust +// build.rs +println!("cargo:rustc-link-search=native=/path/to/cactus/cactus-engine/build"); +``` + +## Usage + +```rust +use std::ffi::CString; +use std::os::raw::c_char; + +mod cactus; + +fn main() { + unsafe { + let path = CString::new("/path/to/model").unwrap(); + let model = cactus::cactus_init(path.as_ptr(), std::ptr::null(), false); + + let messages = CString::new(r#"[{"role":"user","content":"Hello"}]"#).unwrap(); + let mut response = vec![0u8; 65536]; + cactus::cactus_complete( + model, + messages.as_ptr(), + response.as_mut_ptr() as *mut c_char, + response.len(), + std::ptr::null(), std::ptr::null(), + None, std::ptr::null_mut(), + std::ptr::null(), 0, + ); + + cactus::cactus_destroy(model); + } +} +``` + +## Streaming transcription + +Push 16 kHz mono PCM16 as it arrives; each call returns `{"success":true,"confirmed":...,"pending":...}` (`confirmed` is final, `pending` is the volatile tail). All `unsafe`, like the rest: + +```rust +let opts = CString::new(r#"{"language":"en"}"#).unwrap(); +let stream = cactus::cactus_stream_transcribe_start(model, opts.as_ptr()); +let mut buf = vec![0u8; 65536]; +for chunk in pcm_chunks { // each: 16 kHz mono i16 bytes + cactus::cactus_stream_transcribe_process( + stream, chunk.as_ptr(), chunk.len(), + buf.as_mut_ptr() as *mut c_char, buf.len(), + ); + // append "confirmed"; show "pending" as a live preview +} +cactus::cactus_stream_transcribe_stop(stream, buf.as_mut_ptr() as *mut c_char, buf.len()); +``` diff --git a/bindings/rust/cactus.rs b/bindings/rust/cactus.rs new file mode 100644 index 000000000..3a956a618 --- /dev/null +++ b/bindings/rust/cactus.rs @@ -0,0 +1,54 @@ +use std::os::raw::{c_char, c_float, c_int, c_void}; + +pub type CactusModelT = *mut c_void; +pub type CactusIndexT = *mut c_void; +pub type CactusStreamTranscribeT = *mut c_void; +pub type CactusTokenCallback = Option; +pub type CactusLogCallback = Option; + +#[cfg_attr(target_os = "macos", link(name = "Accelerate", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "Metal", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "MetalPerformanceShaders", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "Foundation", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "Security", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "SystemConfiguration", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "CFNetwork", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "curl"))] +#[link(name = "cactus_engine", kind = "static")] +unsafe extern "C" { + pub fn cactus_init(model_path: *const c_char, corpus_dir: *const c_char, cache_index: bool) -> CactusModelT; + pub fn cactus_destroy(model: CactusModelT); + pub fn cactus_reset(model: CactusModelT); + pub fn cactus_stop(model: CactusModelT); + + pub fn cactus_complete(model: CactusModelT, messages_json: *const c_char, response_buffer: *mut c_char, buffer_size: usize, options_json: *const c_char, tools_json: *const c_char, callback: CactusTokenCallback, user_data: *mut c_void, pcm_buffer: *const u8, pcm_buffer_size: usize) -> c_int; + pub fn cactus_prefill(model: CactusModelT, messages_json: *const c_char, response_buffer: *mut c_char, buffer_size: usize, options_json: *const c_char, tools_json: *const c_char, pcm_buffer: *const u8, pcm_buffer_size: usize) -> c_int; + pub fn cactus_tokenize(model: CactusModelT, text: *const c_char, token_buffer: *mut u32, token_buffer_len: usize, out_token_len: *mut usize) -> c_int; + pub fn cactus_score_window(model: CactusModelT, tokens: *const u32, token_len: usize, start: usize, end: usize, context: usize, response_buffer: *mut c_char, buffer_size: usize) -> c_int; + pub fn cactus_transcribe(model: CactusModelT, audio_file_path: *const c_char, prompt: *const c_char, response_buffer: *mut c_char, buffer_size: usize, options_json: *const c_char, callback: CactusTokenCallback, user_data: *mut c_void, pcm_buffer: *const u8, pcm_buffer_size: usize) -> c_int; + pub fn cactus_stream_transcribe_start(model: CactusModelT, options_json: *const c_char) -> CactusStreamTranscribeT; + pub fn cactus_stream_transcribe_process(stream: CactusStreamTranscribeT, pcm_buffer: *const u8, pcm_buffer_size: usize, response_buffer: *mut c_char, buffer_size: usize) -> c_int; + pub fn cactus_stream_transcribe_stop(stream: CactusStreamTranscribeT, response_buffer: *mut c_char, buffer_size: usize) -> c_int; + + pub fn cactus_embed(model: CactusModelT, text: *const c_char, embeddings_buffer: *mut c_float, buffer_size: usize, embedding_dim: *mut usize, normalize: bool) -> c_int; + pub fn cactus_image_embed(model: CactusModelT, image_path: *const c_char, embeddings_buffer: *mut c_float, buffer_size: usize, embedding_dim: *mut usize) -> c_int; + pub fn cactus_audio_embed(model: CactusModelT, audio_path: *const c_char, embeddings_buffer: *mut c_float, buffer_size: usize, embedding_dim: *mut usize) -> c_int; + + pub fn cactus_rag_query(model: CactusModelT, query: *const c_char, response_buffer: *mut c_char, buffer_size: usize, top_k: usize) -> c_int; + + pub fn cactus_index_init(index_dir: *const c_char, embedding_dim: usize) -> CactusIndexT; + pub fn cactus_index_add(index: CactusIndexT, ids: *const c_int, documents: *const *const c_char, metadatas: *const *const c_char, embeddings: *const *const c_float, count: usize, embedding_dim: usize) -> c_int; + pub fn cactus_index_delete(index: CactusIndexT, ids: *const c_int, ids_count: usize) -> c_int; + pub fn cactus_index_get(index: CactusIndexT, ids: *const c_int, ids_count: usize, document_buffers: *mut *mut c_char, document_buffer_sizes: *mut usize, metadata_buffers: *mut *mut c_char, metadata_buffer_sizes: *mut usize, embedding_buffers: *mut *mut c_float, embedding_buffer_sizes: *mut usize) -> c_int; + pub fn cactus_index_query(index: CactusIndexT, embeddings: *const *const c_float, embeddings_count: usize, embedding_dim: usize, options_json: *const c_char, id_buffers: *mut *mut c_int, id_buffer_sizes: *mut usize, score_buffers: *mut *mut c_float, score_buffer_sizes: *mut usize) -> c_int; + pub fn cactus_index_compact(index: CactusIndexT) -> c_int; + pub fn cactus_index_destroy(index: CactusIndexT); + + pub fn cactus_get_last_error() -> *const c_char; + pub fn cactus_log_set_level(level: c_int); + pub fn cactus_log_set_callback(callback: CactusLogCallback, user_data: *mut c_void); + pub fn cactus_set_telemetry_environment(framework: *const c_char, cache_location: *const c_char, version: *const c_char); + pub fn cactus_set_app_id(app_id: *const c_char); + pub fn cactus_telemetry_flush(); + pub fn cactus_telemetry_shutdown(); +} diff --git a/bindings/swift/Cactus.swift b/bindings/swift/Cactus.swift new file mode 100644 index 000000000..44a4cbfe7 --- /dev/null +++ b/bindings/swift/Cactus.swift @@ -0,0 +1 @@ +@_exported import cactus diff --git a/bindings/swift/README.md b/bindings/swift/README.md new file mode 100644 index 000000000..0cd20bbf9 --- /dev/null +++ b/bindings/swift/README.md @@ -0,0 +1,47 @@ +# Swift Bindings + +C module map import of `cactus_engine.h`. Works on iOS and macOS. + +## Integration + + +```bash +cactus build --apple +``` + + + +**XCFramework**: drag `apple/cactus-ios.xcframework` (or `apple/cactus-macos.xcframework`) into Xcode (Embed & Sign). The framework bundles `cactus_engine.h`, so `import cactus` works directly. + +**Static library**: link `apple/libcactus_engine-device.a` (or `apple/libcactus_engine-simulator.a`), copy `bindings/swift/module.modulemap` into your project, and add `cactus-engine/` to Header Search Paths so the module map can find `cactus_engine.h`. + + +## Usage + + +```swift +import cactus + +let model = cactus_init("/path/to/model", nil, false) +var buf = [CChar](repeating: 0, count: 65536) +cactus_complete(model, messagesJson, &buf, buf.count, nil, nil, nil, nil, nil, 0) +let response = String(cString: buf) +cactus_destroy(model) +``` + + +### Streaming transcription + +The streaming functions are exposed the same way. Push 16 kHz mono PCM16 as it arrives and read back `{"success":true,"confirmed":...,"pending":...}` — `confirmed` words are final, `pending` is the volatile tail. + +```swift +let stream = cactus_stream_transcribe_start(model, "{\"language\":\"en\"}") +var buf = [CChar](repeating: 0, count: 65536) +for chunk in pcmChunks { // each chunk: 16 kHz mono 16-bit PCM + chunk.withUnsafeBytes { raw in + _ = cactus_stream_transcribe_process(stream, raw.bindMemory(to: UInt8.self).baseAddress, raw.count, &buf, buf.count) + } + // parse String(cString: buf) -> append "confirmed", show "pending" live +} +cactus_stream_transcribe_stop(stream, &buf, buf.count) +``` diff --git a/bindings/swift/ffi/cactus_ffi_shim.h b/bindings/swift/ffi/cactus_ffi_shim.h new file mode 100644 index 000000000..ae623eeb8 --- /dev/null +++ b/bindings/swift/ffi/cactus_ffi_shim.h @@ -0,0 +1 @@ +#include "../../../cactus-engine/cactus_engine.h" diff --git a/bindings/swift/ffi/module.modulemap b/bindings/swift/ffi/module.modulemap new file mode 100644 index 000000000..0fc94f2dd --- /dev/null +++ b/bindings/swift/ffi/module.modulemap @@ -0,0 +1,4 @@ +module cactus [system] { + header "cactus_ffi_shim.h" + export * +} diff --git a/apple/module.android.modulemap b/bindings/swift/module.modulemap similarity index 51% rename from apple/module.android.modulemap rename to bindings/swift/module.modulemap index b68f641c2..af75a51b2 100644 --- a/apple/module.android.modulemap +++ b/bindings/swift/module.modulemap @@ -1,4 +1,4 @@ module cactus { - header "cactus_ffi.h" + header "cactus_engine.h" export * } diff --git a/blog/README.md b/blog/README.md new file mode 100644 index 000000000..420c7d498 --- /dev/null +++ b/blog/README.md @@ -0,0 +1,10 @@ +# Cactus Blog + +| Post | Author | Description | +|------|--------|-------------| +| [TurboQuant-H](turboquant-h.md) | Karen Mosoyan & Henry Ndubuaku | Hadamard rotation for 2-bit per-layer embedding quantization: 4x PLI compression, 40% total model size reduction | +| [Gemma 4 on Cactus](gemma4.md) | Henry Ndubuaku | Native multimodal voice, vision, and audio on-device with hybrid cloud handoff | +| [Hybrid Transcription](hybrid_transcription.md) | Roman Shemet | Sub-150ms transcription with cloud-level accuracy using on-device/cloud hybrid inference | +| [On-Device Coding Agents](lfm2_24b_a2b.md) | Noah Cylich & Henry Ndubuaku | Running LFM2-24B MoE locally on Mac for coding use cases | +| [Ridiculously Fast Transcription](parakeet.md) | Satyajit Kumar & Henry Ndubuaku | 6M tok/sec decode speed with Parakeet CTC 1.1B | +| [LFM-2.5-350m on Cactus](lfm2.5_350m.md) | Henry Ndubuaku | 140 tok/sec single-core INT8 inference across seven devices, from Vision Pro to Raspberry Pi 5 | \ No newline at end of file diff --git a/blog/gemma4.md b/blog/gemma4.md new file mode 100644 index 000000000..cd07ff058 --- /dev/null +++ b/blog/gemma4.md @@ -0,0 +1,141 @@ +--- +title: "Gemma 4 on Cactus: The first model you can talk to, show things, and trust to know when it needs help" +description: "Gemma 4 runs natively on your device with real-time voice, vision, and audio, and routes hard problems to the cloud when it should." +keywords: ["gemma4", "multimodal", "hybrid", "realtime inference", "on-device inference", "Google", "DeepMind"] +author: "Henry Ndubuaku & The Cactus Jacks" +date: 2026-04-01 +tags: ["gemma4", "cactus", "on-device", "hybrid inference", "multimodal", "voice AI"] +--- + +Gemma 4 is the first on-device model that genuinely works across text, vision, and audio in a single architecture, and does it well enough that you stop thinking about the model and start thinking about what to build with it. + +We shipped day-one support in Cactus. Here's what that means in practice and why we think this changes the trajectory of on-device AI. + +## What's actually new here + +Most "multimodal" models bolt on vision or audio as an afterthought. Gemma 4 is different. The vision encoder, audio encoder, and language model were trained together from the start. The audio conformer isn't a whisper-style bolt-on. It's a 300M-parameter encoder that feeds directly into the transformer's residual stream. Same for vision. The model doesn't "transcribe then think." It reasons over the raw modality. + +This matters because latency compounds. A pipeline that transcribes audio, then feeds text to an LLM, then generates a response has three serialized steps. Gemma 4 collapses that into one forward pass. On a modern ARM device, that means 0.3 seconds from the end of a 30-second voice clip to the first token of a response. Not 0.3 seconds for the transcription step. 0.3 seconds total. + +## Performance + +Cactus targets ARM across platforms: Apple Silicon Macs, iPhones, iPads, Vision Pro, and Android devices with ARM64 chipsets (Snapdragon, Dimensity, Tensor). On Apple hardware we additionally leverage Metal GPU acceleration, while on Android and Linux the runtime uses NEON, i8mm, and dot-product intrinsics. + +| Metric (M5 Mac/iPad/Vision Pro) | E2B | +|---|---| +| 4096-token prefill | 660 tok/s | +| 1024-token decode | 40 tok/s | +| 30s audio end-to-end | 0.3s | +| Image encode | 0.7s | + +40 tok/s decode is faster than most people read. You're not waiting for the model. The model is waiting for you. + +## The benchmarks are hard to believe + +We normally don't lead with benchmarks because they rarely reflect real-world use. But these numbers demand attention. + +### LLM + +| Benchmark | E4B | E2B | Gemma 3 27B (no think) | +|---|---|---|---| +| MMLU Pro | 69.4% | 60.0% | 67.6% | +| AIME 2026 (no tools) | 42.5% | 37.5% | 20.8% | +| LiveCodeBench v6 | 52.0% | 44.0% | 29.1% | +| Codeforces ELO | 940 | 633 | 110 | +| GPQA Diamond | 58.6% | 43.4% | 42.4% | +| Tau2 (avg over 3) | 42.2% | 24.5% | 16.2% | +| BigBench Extra Hard | 33.1% | 21.9% | 19.3% | +| MMMLU | 76.6% | 67.4% | 70.7% | +| MRCR v2 8-needle 128k (avg) | 25.4% | 19.1% | 13.5% | + +E4B, a 4.5B-effective-parameter model running on your phone, outperforms Gemma 3 27B on nearly every benchmark. E2B, the smaller variant at 2.3B effective parameters, matches or beats it on half of them. The AIME and LiveCodeBench numbers are particularly striking: these are hard reasoning tasks where you'd expect scale to dominate. + +### Vision + +| Benchmark | E4B | E2B | Gemma 3 27B (no think) | +|---|---|---|---| +| MMMU Pro | 52.6% | 44.2% | 49.7% | +| OmniDocBench 1.5 (edit dist, lower=better) | 0.181 | 0.290 | 0.365 | +| MATH-Vision | 59.5% | 52.4% | 46.0% | +| MedXPertQA MM | 28.7% | 23.5% | - | + +E4B beats Gemma 3 27B on vision tasks across the board. Document understanding (OmniDocBench) is 2x better. This isn't a toy. It's genuinely useful for reading receipts, parsing handwritten notes, understanding diagrams. + +### Audio + +| Benchmark | E4B | E2B | +|---|---|---| +| CoVoST | 35.54 | 33.47 | +| FLEURS (lower=better) | 0.08 | 0.09 | + +## Model Architecture + +| Property | E2B | E4B | +|---|---|---| +| Total Parameters | 2.3B effective (5.1B w/ embeddings) | 4.5B effective (8B w/ embeddings) | +| Layers | 35 | 42 | +| Sliding Window | 512 tokens | 512 tokens | +| Context Length | 128K tokens | 128K tokens | +| Vocabulary Size | 262K | 262K | +| Supported Modalities | Text, Image, Audio | Text, Image, Audio | +| Vision Encoder | ~150M params | ~150M params | +| Audio Encoder | ~300M params | ~300M params | + +The architecture uses per-layer embeddings with AltUp, a technique that keeps most of the vocabulary knowledge in a shared embedding table while giving each layer a small specialized projection. This is how they fit 262K vocabulary tokens into a model that runs on a phone. The sliding window attention at 512 tokens with global attention every 5 layers gives you 128K context without the quadratic memory cost. + +## What this unlocks + +### Voice control that actually works + +We've all used voice assistants that feel like speaking into a form field. The voice gets transcribed, the text gets processed, and you get a response that could just as easily have been typed. Gemma 4 doesn't work that way. Because it reasons directly over the audio signal, it picks up on tone, hesitation, emphasis. It knows the difference between "delete that" spoken confidently and "delete that?" spoken as a question. + +On headsets running Cactus (Vision Pro, Quest, or any AR/VR platform with ARM compute), this means spatial voice control that responds to how you speak, not just what you say. Ask it to "move that over there" while looking at a 3D object and gesturing, and the model has the audio understanding to disambiguate "that" from context. The 0.3s latency means the interaction feels physical. You speak, things happen. + +### Voice agents that run locally + +If you're building a voice agent (a customer service bot, an in-car assistant, a medical triage system) you currently have two bad options. You can run everything in the cloud and deal with latency, cost, and privacy issues. Or you can cobble together a local pipeline of ASR + LLM + TTS and deal with error propagation and integration pain. + +Gemma 4 on Cactus gives you a third option. The model handles voice understanding natively. You pipe audio in, you get structured responses out. Tool calling works. The model can decide to call functions, hit APIs, or route to a cloud model when the task exceeds its capability. And because it runs locally, there's no per-request cost and no audio leaving the device. + +For healthcare, legal, and finance applications where audio data is sensitive, this isn't a nice-to-have. It's a compliance requirement. + +### Hybrid inference: knowing when to ask for help + +This is the part we're most excited about. Gemma 4 is good, but it's not omniscient. A 2B-parameter model isn't going to write a production database migration or debug a complex distributed system. What it can do is recognize when a task is beyond its capability and route it to a frontier cloud model. + +Cactus implements this as cloud handoff. The on-device model evaluates the complexity of the request, and if it determines it can't handle it confidently, it signals for handoff. The request goes to a cloud model (Claude, GPT-4, Gemini, your choice), the response comes back, and the user sees a seamless interaction. The on-device model handles the 80% of requests that are straightforward. The cloud handles the 20% that need heavy lifting. + +This means you can build apps that feel like they have frontier-model intelligence while keeping costs at a fraction of full cloud inference. And the user's casual conversations, voice notes, and image queries never leave their device. + +### Building on any platform + +Cactus runs Gemma 4 on macOS, iOS, Android, and Linux. + +## Try it + +```bash +brew install cactus-compute/cactus/cactus +cactus run google/gemma-4-E2B-it +``` + +The weights download automatically. INT4 quantized, optimized for ARM. Both E2B and E4B are available. + +## Build with it + +Cactus supports React Native, Flutter, Swift, Kotlin, Python, Rust, and C++. Pick the binding that fits your stack: + +| Platform | Binding | Install | +|---|---|---| +| React Native | [React Native bindings](/bindings/react-native/) | Native bridge over the C API | +| Flutter | [cactus-flutter](/bindings/flutter/) | Dart FFI bindings | +| Swift (iOS/macOS) | [cactus-swift](/bindings/swift/) | XCFramework with Metal support | +| Kotlin (Android) | [cactus-kotlin](/bindings/kotlin/) | JNI + Kotlin Multiplatform | +| Python | [cactus-compute](/python/) | `pip install cactus-compute` | +| Rust | [Rust bindings](/bindings/rust/) | Copy `cactus.rs`, link `libcactus_engine.a` | +| C++ | [cactus.h](/docs/cactus_engine.md) | Single header, link `libcactus_engine.a` | + +Full quickstart with code examples for every binding: [Quickstart](/docs/quickstart.md) + +Pre-quantized weights for Gemma 4 and 30+ other models: [huggingface.co/Cactus-Compute](https://huggingface.co/Cactus-Compute) + +If you're building something with on-device multimodal AI (voice agents, VR interfaces, local-first apps) we want to hear about it. Open an issue on [GitHub](https://github.com/cactus-compute/cactus), or just start building. diff --git a/blog/hybrid_transcription.md b/blog/hybrid_transcription.md new file mode 100644 index 000000000..e92ce139c --- /dev/null +++ b/blog/hybrid_transcription.md @@ -0,0 +1,81 @@ +--- +title: "Sub-150ms Transcription with Cloud-Level Accuracy: Why We Built a Hybrid Engine" +description: "How Cactus combines on-device and cloud inference for real-time speech transcription with sub-150ms latency and automatic cloud handoff for noisy audio." +keywords: ["hybrid AI", "speech transcription", "on-device inference", "cloud handoff", "Whisper", "voice interface", "edge AI"] +author: "Roman Shemet" +date: 2026-02-01 +tags: ["transcription", "hybrid AI", "edge AI", "whisper", "cloud handoff"] +--- + +# Sub-150ms transcription with cloud-level accuracy: Why we built a hybrid engine + +*By Roman Shemet* + +## The Voice Interface Dilemma +Voice is no longer a novelty; it is quickly becoming the default interface. From meeting transcriptions and dictation to voice notes and ambient computing, users expect to speak to their devices. + +But building a seamless voice experience is hard. If you've ever built a voice product, you are intimately familiar with the iron triangle of AI: Cheap, Accurate, Fast. + +Historically, you could only pick two: +- **Fast and Accurate?** It won't be cheap. You'll need heavy cloud compute. +- **Fast and Cheap?** It won't be accurate. You're running tiny, compromised models. +- **Cheap and Accurate?** It won't be fast. You're waiting in API queues. + +The logic is simple. High-accuracy transcription requires large models. Large models require massive cloud resources. This introduces network latency and high infrastructure costs. + +Or, you can run small models on-device. The latency is practically zero, and compute is free. But smaller models struggle with noisy environments, accents, and complex vocabulary. + +You either pay the "cloud tax" in latency and cost, or you pay the "edge tax" in accuracy. + +## Enter Hybrid AI: The Best of Both Worlds +We didn't want to choose, so we spent the last few months building a third option at Cactus. We call it Hybrid AI. + +> What if we had small, fast on-device models that were self-aware enough to **know** when they are struggling? + +Instead of routing 100% of user audio to an expensive cloud API, the Cactus engine processes speech locally by default. Our on-device inference provides real-time transcription with sub-150ms latency. + +However, when the engine detects messy audio (think background noise, static, or multiple speakers), it automatically hands that specific segment off to a larger model in the cloud to clean it up. + +It's like having a brilliant intern who processes 80% of the work instantly, but knows exactly when to call the senior engineer to review the tricky edge cases. + +## See it in Action +Here is a look at the hybrid engine running locally, seamlessly managing the local-to-cloud transcription flow: + +[![Realtime Hybrid Transcription With Cactus](https://img.youtube.com/vi/0SldBm8rGkE/maxresdefault.jpg)](https://youtu.be/0SldBm8rGkE?si=TLLYnDuGEG8Rqexw) + +As a Mac user, you can pull down the CLI and test the latency and accuracy locally right now: + +```bash +brew install cactus-compute/cactus/cactus +cactus transcribe --file path/to/audio.wav +``` + +## Built for the Edge, from the Ground Up +Achieving this required relying on our industry-leading on-device engine. We built Cactus from the ground up exactly for this. + +Because the engine sits directly on the metal, it is optimized for edge constraints. It targets ARM-based CPUs natively, resulting in super low RAM usage (e.g. Moonshine transcription for example, runs with sub-10MB RAM utilization). This is a critical factor for mobile and wearable devices where memory is strictly rationed by the OS. + +By keeping the footprint minimal, the engine operates silently in the background without draining battery or causing OS-level memory warnings, making it viable for continuous, ambient listening. + +## A Single API for Every Platform +While the underlying engine is low-level C++, we know that product teams need to move fast in their native environments. We designed the integration to be a drop-in experience regardless of your stack. + +The interface remains the same whether you are a low-level C++ engineer, or a mobile developer. You can integrate the Cactus Hybrid AI engine directly into your stack: +- **Mobile:** React Native, Flutter, Kotlin, Swift +- **Systems:** C++, Rust +- **Scripting:** Python + +At the top level, the implementation is beautifully simple: initialize the model, pass the audio stream, and receive real-time text. + +A few levels deeper, developers have granular control. You can configure the engine's behavior to define a specific Word Error Rate (WER) threshold, target a strict cloud-handoff ratio, or establish hard cost limits per user. + +Check out the open-source engine and documentation on GitHub. + +## See Also + +- [Cactus Engine API Reference](/docs/cactus_engine.md) — Transcription API: one-shot `cactus_transcribe` and streaming `cactus_stream_transcribe_*` +- [Python Binding](/python/) — Python bindings +- [Swift Binding](/bindings/swift/) — Swift API +- [Kotlin Binding](/bindings/kotlin/) — Kotlin API +- [Flutter Binding](/bindings/flutter/) — Flutter bindings +- [LFM2 24B Review](/blog/lfm2_24b_a2b.md) — Running large MoE models locally with Cactus \ No newline at end of file diff --git a/blog/lfm2.5_350m.md b/blog/lfm2.5_350m.md new file mode 100644 index 000000000..3c314fa2f --- /dev/null +++ b/blog/lfm2.5_350m.md @@ -0,0 +1,204 @@ +--- +title: "LFM-2.5-350m on Cactus: 140 tok/sec, Single Core, 355 MB" +description: "Benchmarking Liquid's LFM-2.5-350m across seven devices with Cactus. INT8 quantization, single-core CPU decode, zero-copy loading, and why this configuration makes on-device inference practical." +keywords: ["LFM-2.5", "on-device inference", "INT8 quantization", "ARM CPU", "mobile AI", "edge inference", "Cactus"] +author: "Henry Ndubuaku" +date: 2026-03-31 +tags: ["LFM", "INT8", "edge AI", "mobile inference", "ARM", "benchmarks"] +--- + +# LFM-2.5-350m on Cactus: 140 tok/sec, Single Core, 355 MB + +*By Henry Ndubuaku* + +We ran Liquid's LFM-2.5-350m through the Cactus inference engine across seven devices, from a Vision Pro to a Raspberry Pi 5. +The configuration: 1024-token prefill, 100-token decode, INT8 quantization, CPU-only, single-core decode. +The model file is 355 MB on Cactus Compressed format. + +## Device Benchmarks + +### Apple +| Device | Prefill (tok/s) | Decode (tok/s) | RAM | +|---|---|---|---| +| Vision Pro | 2,067 | 140 | 85 MB | +| iPhone 17 Pro | 797 | 140 | 75 MB | +| iPhone 13 Mini | 496 | 88 | 56 MB | + +### Android +| Device | Prefill (tok/s) | Decode (tok/s) | RAM | +|---|---|---|---| +| Galaxy S25 Ultra | 660 | 100 | 322 MB | +| Google Pixel 6a | 208 | 42 | 328 MB | +| Galaxy A56 | 110 | 24 | 330 MB | + +### Other +| Device | Prefill (tok/s) | Decode (tok/s) | RAM | +|---|---|---|---| +| Raspberry Pi 5 | 200 | 30 | 300 MB | + +Two things stand out in this data. First, the decode speeds: 140 tok/sec on a single core is fast, +most people read at roughly 4-5 words per second, so even the Raspberry Pi's 30 tok/sec is well ahead of human consumption. +Second, the RAM gap between Apple and Android: 56-85 MB versus 300-330 MB for the same 355 MB model. +That gap is not a bug. It's a direct consequence of how the two platforms handle memory-mapped files, which we'll get into in the Cactus format section below. + +## Why CPU, Not GPU or Dedicated Accelerators + +It is tempting to reach for the GPU. On paper, mobile GPUs have 10-50x the FLOPS of a single CPU core. +But inference on a 350M-parameter model is not a FLOPS problem, it's a memory bandwidth problem. +During decode, you generate one token at a time, which means you read the entire weight matrix for each token but only perform a single matrix-vector multiply. +The arithmetic intensity (FLOPS per byte loaded) is extremely low. Without batched decode, GPU compute units spend most of their time stalled on memory fetches. +GPU excels when arithmetic intensity is high, large batch sizes, long matrix-matrix multiplies. +For single-token decode, the GPU's compute units sit mostly idle while waiting for memory. + +There's also the practical problem: mobile GPU inference requires Metal on iOS and Vulkan or OpenCL on Android. +That's two separate shader codepaths to maintain, with different memory models, different synchronization primitives, and different performance characteristics per vendor. +And GPU inference holds a persistent GPU context, which competes with the UI rendering pipeline and drains battery even when the inference workload is light. +GPUs are also energy-inefficient. On Macs and PCs this barely matters since they're plugged into a power source most of the time, and GPU inference works well there. But mobile devices run on battery, and every watt spent on GPU inference cuts directly into the user's runtime. The GPU path that works fine on a MacBook becomes a battery killer on a phone. + +Dedicated AI accelerators are a more interesting target. Apple's Neural Engine and Qualcomm's Hexagon DSP both offer high-throughput inference at good power efficiency. +The problem is access. Apple's Neural Engine requires conversion into a proprietary model format, which imposes its own quantization and graph constraints. +You don't control the execution schedule, the memory layout, or the precision semantics. +The quantisation techniques available through that toolchain severely lag the state-of-the-art methods used by GGML and Cactus. +Qualcomm's QNN SDK has similar restrictions (limited quantisation scheme support) and not all Android devices ship Qualcomm chips. +Both lock you into a vendor-specific format that may not support your model architecture at all, and neither works on budget devices that lack the accelerator entirely. + +CPU is the universal target. Every ARM device, flagship, budget, wearable, Raspberry Pi, and most recent ones support DOTPROD or I8MM extensions. +A single well-optimized CPU codepath works everywhere, the main challenge is energy-inefficiency compared to dedicated accelerators. + +ARM is also closing this gap. The progression from NEON to DOTPROD to I8MM to SME2 reflects a deliberate push to bring matrix-processing efficiency onto the CPU itself. Each generation narrows the energy and throughput gap with dedicated accelerators. SME2, landing on upcoming ARMv9.2-A cores, adds streaming matrix operations that approach accelerator-class efficiency for INT8 workloads while remaining fully programmable. The bet is that general-purpose CPU silicon will continue absorbing the capabilities that once justified a separate accelerator. + +## Why Single Core for Decode + +Decode is memory-bandwidth-bound. For a 350M INT8 model, each decode step reads roughly 355 MB of weights to produce a single token. +On an iPhone 17 Pro with ~60 GB/s memory bandwidth, the theoretical maximum decode speed is around 169 tok/sec, +our measured 140 tok/sec is 83% of that ceiling, the haircut comes from our choice to stream weights from storage for memory efficiency, reducing OS chances of throttling background inference. +Adding more cores doesn't give you more memory bandwidth; it just adds synchronization overhead. + +Prefill is different. With 1024 tokens in the input, each weight load amortizes across 1024 multiply-accumulate operations. +The arithmetic intensity is 1024x higher than decode, making prefill genuinely compute-bound. +That's why we use multi-threaded prefill, the Vision Pro's 2,067 tok/sec prefill speed reflects all available cores working in parallel. + +But there's a deeper reason single-core decode matters for mobile: background execution. +On both iOS and Android, the OS aggressively throttles multi-core workloads that run in the background. +iOS will suspend or terminate background tasks that consume excessive CPU across multiple cores. +Android's thermal management will cap frequency on the performance cores. +We found that single-core decode workload, by contrast, looks like a lightweight background task to the OS scheduler. +It doesn't trigger thermal throttling, doesn't compete with foreground apps, and doesn't get killed. +This is the difference between inference as a demo and inference as a product feature, the model needs to run while the user is doing something else. + +## Why INT8 + +At 350M parameters, quantization precision matters more than at larger scales. +Quantization works by approximating continuous weight values with discrete integers, and the approximation error is relative to the dynamic range of each weight group. +Smaller models have fewer parameters absorbing this error, so each parameter carries more of the model's capacity. +INT4 quantization at 350M parameters produces measurable accuracy degradation, our internal testing shows 2-4 point drops on MMLU compared to INT8, which is near-lossless at this scale. + +Cactus uses grouped affine quantization: weights are divided into groups of 32 elements, and each group gets its own scale factor. +For each group, the scale is computed as `max(|w|) / 127`, and each weight is rounded to the nearest integer in [-128, 127]. +During inference, the INT8 values are multiplied by their group's scale to reconstruct the approximate FP32 value. +Per-group scaling is critical, a single outlier weight in a per-tensor scheme would compress the effective precision of every other weight in that tensor. + +The performance benefit is concrete. On ARM CPUs with I8MM support (ARMv8.2-a and later), the `SMMLA` instruction computes a 2x8 by 8x2 INT8 matrix multiply in a single instruction, accumulating into INT32. +This is a native hardware operation, not a software emulation of low-precision arithmetic. +On devices without I8MM (like the Pixel 6a's Tensor G1), Cactus falls back to DOTPROD instructions, which compute 4-element INT8 dot products. +Both paths are substantially faster than FP16 arithmetic because they process more elements per instruction and use less memory bandwidth per parameter. + +## The Cactus Format and Zero-Copy Loading + +The RAM numbers in our benchmarks tell a story about memory architecture. +The LFM-2.5-350m model file is 355 MB. On Apple devices, runtime RAM usage is 56-85 MB. On Android, it's 300-330 MB. +The model is the same. The quantization is the same. The difference is in how the operating system handles memory-mapped files. + +The Cactus format (identified by the magic number `0x54434143`, "CACT") is a binary container designed around `mmap()`. +Rather than loading the entire model into a heap allocation, Cactus maps the file directly into virtual address space with `mmap(PROT_READ, MAP_SHARED)`. +The weights stay on disk. When the inference engine accesses a weight tensor, the OS kernel page-faults it into physical memory on demand. +When memory pressure rises, the OS can evict those pages, they're backed by the file, so no writeback is needed. + +On Apple Silicon, this works exceptionally well. +The unified memory architecture means there's no distinction between CPU memory and file cache, +a memory-mapped page that's been paged in *is* the physical memory the CPU reads from. +The OS can manage a working set that's a fraction of the total model size because not all layers are active simultaneously during a single decode step. + +On Android, the behavior varies by kernel version and vendor. +Many Android devices aggressively copy memory-mapped pages into anonymous memory due to SELinux policies and vendor-specific memory management. +The result is that `mmap()` often degrades to something closer to `malloc() + read()`, explaining the 300+ MB footprint. + +The format itself is designed for SIMD-friendly access. +Weight data is aligned to 32-byte boundaries (matching ARM NEON register width), +and an optional block-interleaved layout groups 4 columns of weights together so that a single vector load fetches data for 4 output neurons simultaneously. +This layout eliminates the need for runtime transposition and ensures every memory access is a sequential, aligned read, optimal for both cache prefetching and NEON/I8MM instruction operands. + +## Model Quality + +A fast small model is only useful if it's accurate enough. Here's how the LFM-2.5 family compares to similarly-sized open models: + +| Benchmark | LFM2-350M | LFM2-700M | LFM2-1.2B | Qwen3-0.6B | Qwen3-1.7B | Llama-3.2-1B | gemma-3-1b-it | +|---|---|---|---|---|---|---|---| +| MMLU | 43.43 | 49.90 | **55.23** | 44.93 | **59.11** | 46.60 | 40.08 | +| GPQA | **27.46** | **28.48** | **31.47** | 22.14 | 27.72 | 28.84 | 21.07 | +| IFEval | **65.12** | **72.23** | **74.89** | 64.24 | 73.98 | 52.39 | 62.90 | +| IFBench | 16.41 | **20.56** | **20.70** | 19.75 | **21.27** | 16.86 | 17.72 | +| GSM8K | 30.10 | 46.40 | 58.30 | 36.47 | 51.40 | 35.71 | **59.59** | +| MGSM | 29.52 | 45.36 | 55.04 | 41.28 | **66.56** | 29.12 | 43.60 | +| MMMLU | **37.99** | **43.28** | **46.73** | 30.84 | 46.51 | 38.15 | 34.43 | + +The numbers worth examining closely: + +**IFEval (instruction following):** LFM2-350M scores 65.12, versus 52.39 for Llama-3.2-1B (a model 3x its size) and 64.24 for Qwen3-0.6B (nearly 2x its size). For on-device applications, +where the model is typically following structured instructions from an app rather than engaging in open-ended conversation, +instruction following is arguably the most important benchmark. The LFM2-1.2B's 74.89 is within a point of Qwen3-1.7B's 73.98, at 70% of the parameters. + +**GPQA (graduate-level reasoning):** LFM2-350M at 27.46 beats both Qwen3-0.6B (22.14) and gemma-3-1b-it (21.07). +This is notable because GPQA tests reasoning depth, not just pattern matching, it suggests the LFM architecture is efficient at encoding reasoning capability per parameter. + +**MMMLU (multilingual knowledge):** LFM2-350M scores 37.99 versus Qwen3-0.6B's 30.84. +A 7-point lead over a model nearly twice the size suggests strong multilingual data efficiency. The full LFM2-1.2B hits 46.73, matching Qwen3-1.7B's 46.51. + +**Where LFM2-350M trails:** GSM8K (30.10) and MGSM (29.52) math benchmarks. gemma-3-1b-it leads GSM8K at 59.59, and Qwen3-1.7B leads MGSM at 66.56. +Math is where the parameter count ceiling hits hardest, since multi-step arithmetic requires retaining intermediate state that larger models can distribute across more capacity. +For pure math tasks at this scale, a larger model is the right answer. + +## What This Adds Up To + +The LFM-2.5-350m at INT8 on Cactus is a 355 MB model that decodes at 88 tok/sec on a three-year-old iPhone 13 Mini using 56 MB of RAM on a single CPU core. +It follows instructions better than models 3x its size. It runs as a background process without triggering OS throttling. +These are not theoretical numbers on a development board with a fan. They're measured on production consumer devices, under the constraints that real mobile apps face; +memory pressure from other apps, thermal limits, background execution policies, 4G download speeds for model delivery. +The question for on-device inference has always been whether the quality-speed-size tradeoff lands in a useful region. +At 350M parameters, LFM-2.5 on Cactus suggests it does: accurate enough for structured on-device tasks, fast enough to be invisible to the user, and small enough to ship inside an app. + +## Where is Cactus headed? + +On-device inference on a Mac or PC is a solved problem. You have gigabytes of RAM, a fast SSD, a multi-core CPU with wide SIMD, and a stable OS that doesn't kill your process for using too much battery. +Any reasonably written inference engine will run well on a MacBook and GPUs are a great idea, given these devices are plugged in anyway. + +Mobile devices, wearables, and custom hardware are a different story entirely. The ecosystem is incredibly fragmented. +You're dealing with hundreds of ARM chip variants across Apple, Qualcomm, MediaTek, Samsung, and Broadcom, each with different ISA extensions, different memory bandwidth, different thermal envelopes. +Android alone ships on devices ranging from 3 GB of RAM to 16 GB, with kernel versions spanning half a decade. +Wearables have even tighter constraints. Custom embedded hardware has its own. +A model that runs fine on a Galaxy S25 Ultra may crash on a Galaxy A56, not because the code is wrong, but because the memory management behavior differs at the kernel level. + +This fragmentation is what makes on-device AI difficult in production. It's a research and engineering problem that compounds at scale: every new device, every OS update, +every vendor-specific memory policy is another edge case that can break inference for some slice of your user base. +For enterprises shipping AI features to millions of devices, this is a hard pill. +The choice often comes down to small models that fit everywhere but are too weak to be useful, or capable models with large files that fail on half the device population. + +Cactus exists because we took on this specific slice of the problem; I spent the last few years building foundation models and inference for tiny devices. +The team's approach is systematic: profile every major ARM variant, write dedicated kernel paths for each ISA extension tier (NEON, DOTPROD, I8MM, SME2), +design the model format around the worst-case memory behavior (Android's mmap limitations), and test on real production devices, not emulators. +This is not trivial; battling with chip manufacturers with different brand visions, model developers chasing benchmarks over usability and more. + +The result is that businesses shipping with Cactus don't need to worry about inference failing silently on a budget Samsung, or a model that's accurate on benchmarks but too large for half their install base. +The kernel, the format, and the quantization strategy are all designed so that the same model file works correctly across the full device spectrum, with predictable RAM, predictable speed, and predictable quality. +On-device market is quite noisy, with hobbyists driving decisions, but production use is largely unsolved, Cactus is a bit unconventional but rightfully so. + +LFM-2.5-350m is a step in the right direction, hence this dedicated post from me personally. +Not just that LFM-2.5-350m is fast on one device, but on Cactus, it brings frontier intelligence to very different devices, +from a Vision Pro to a Raspberry Pi 5, with no special casing and no device-specific builds. + +## See Also + +- [Cactus Engine API Reference](/docs/cactus_engine.md) - Full C API docs for completion, tool calling, and cloud handoff +- [LFM2-24B-A2B](/blog/lfm2_24b_a2b.md) - Running LFM2 24B MoE locally on Mac for coding use cases +- [Hybrid Transcription](/blog/hybrid_transcription.md) - On-device/cloud hybrid speech transcription with Cactus +- [Ridiculously Fast Transcription](/blog/parakeet.md) - 6M tok/sec decode speed with Parakeet CTC 1.1B \ No newline at end of file diff --git a/blog/lfm2_24b_a2b.md b/blog/lfm2_24b_a2b.md new file mode 100644 index 000000000..098889abf --- /dev/null +++ b/blog/lfm2_24b_a2b.md @@ -0,0 +1,220 @@ +--- +title: "The Sweet Spot for Mac Code Use: Reviewing LFM2 24B MoE A2B with Cactus" +description: "Review of LiquidAI's LFM2-24B-A2B mixture-of-experts model running locally on Mac with Cactus. Architecture breakdown, benchmarks, and coding agent use cases." +keywords: ["LFM2", "mixture of experts", "MoE", "on-device coding", "Mac inference", "LiquidAI", "Apple Silicon"] +author: "Noah Cylich and Henry Ndubuaku" +date: 2026-02-15 +tags: ["LFM2", "MoE", "coding agents", "Apple Silicon", "Python", "function calling"] +--- + +# The Sweet Spot for Mac Code Use: Reviewing LFM2 24B MoE A2B with Cactus + +*By Noah Cylich and Henry Ndubuaku* + +[![Video Title](https://img.youtube.com/vi/-duR4gh_O10/maxresdefault.jpg)](https://youtu.be/-duR4gh_O10) + +LFM2-24B-A2B is a really great next step to see over the LFM2-8B-A1B model. The model features 24B total parameters, but only activates a sparse subset of 2B during inference. This allows it to be competitive in inference speed to 2B dense models, while delivering far greater performance. + +> "LFM2-24B-A2B excels at coding, keen to see on-device coding agents built with these." +> — Henry Ndubuaku, Cactus Co-founder & CTO + +## Architecture Breakdown + +Going into more depth about the model, there's really a lot to appreciate with all the architectural work LFM has accomplished, here's the breakdown: + +1. **GQA**: Grouped-query attention is the industry standard choice for efficient LLMs, their choice of a group size of 4 means that the KV cache is 4x smaller than standard, baseline attention. +2. **Gated Convolution**: This is the signature design choice of the Liquid series of models and efficiently adds parameters and expressiveness without much compute cost. +3. **Efficient Vocab**: The small vocab size of 65k is actually a strength for these models, as the final matmul vocab projection is the slowest static part of every model and is extremely parameter efficient. Gemma3 270m for instance dedicated 170m params just to its vocab projection since it has a vocab of 250k tokens. +4. **MoE**: Mixture of experts is the most important choice for this model that really separates it from Liquid's prior work, it scales up parameters without sacrificing speed. + +Together with Cactus, these choices enable lightning fast inference at low energy. Ultimately, despite being 24B params, only 200mb of running memory is used, while generating 25 TPS with our m4 pro chips with 48gb of ram. + + +## Model Architecture Diagram + +``` + ┌───────────────┐ + │ Linear │ + │Tied w/ Embed. │ + └───────┬───────┘ + │ + ┌─────┴─────┐ ┌───────────────────────┐ + │ Norm │ │ Gated Short │ + └─────┬─────┘ │ Convolution Block │ + │ │ │ + ┌──────────────────┐ ┌─────────⊕─────────────┐ │ ↑ │ + │ SwiGLU Expert │ │ │ │ │ ┌───┴───┐ │ + │ │ │ ┌───────┴───────────┐ │ │ │Linear │ │ + │ ↑ │ │ │ MoE Block │ │ │ └───┬───┘ │ + │ ┌───┴───┐ │ │ │ ↑ │ │ │ │ │ + │ │Linear │ │ │ │ ⊕ │ │ │ ┌──────►⊗ │ + │ └───┬───┘ │ │ │ ↗ ↗ ↖ ↖ │ │ │ │ ↑ │ + │ │ │ │ │ ⊗ ⊗ ⊗ ⊗ │ │ │ ┌───┴────┐ │ │ + │ ⊗ ◄───┐ │◄------------┤ │ ╎ ╎ ╎ ╎ │ │ ┌------►│ │ Conv1D │ │ │ + │ ↑ │ │ │ │ ↑ ↑ ↑ ↑ │ │ ╎ │ └───┬────┘ │ │ + │ │ ┌──┴──┐│ │ │ E1...E4...E9...E64│ │ ╎ │ │ │ │ + │ │ │SiLU ││ │ │ ↑ ↑ ↑ ↑ │ │ ╎ │ ⊗ ◄───┐ │ │ + │ │ └──┬──┘│ │ ├─┴────┴─────┴────┴─┤ │ ╎ │ ↑ │ │ │ + │ ┌───┴───┐ │ │ │ │ ┌───────────┐ │ │ ╎ │ B X C │ + │ │Linear ├─┘ │ │ │ │ Router │ │ │ ╎ │ ↑ ↑ ↑ │ + │ └───┬───┘ │ │ │ └─────┬─────┘ │ │ ╎ │ ┌──┴─────┴─┴──────┐ │ + │ ↑ │ │ │ │ │ │ ╎ │ │ Linear │ │ + └────────│─────────┘ │ └─────────┴─────────┘ │ ╎ │ └────────┬────────┘ │ + │ │ │ ╎ │ ↑ | + └───────────│───────────┘ × Num of Layers ╎ └───────────│───────────┘ + │ ╎ + ┌─────┴─────┐ ╎ + │ Norm │ ╎ ┌───────────────────────┐ + └─────┬─────┘ ╎ │ GQA Block │ + │ ╎ │ │ + ┌─────────⊕─────────┐ ╎ │ ↑ │ + │ │ │ ╎ │ ┌───┴───┐ │ + │ ┌───────┴───────┐ │ │ │ │Linear │ │ + │ │Sequence Block │ │ │ │ └───┬───┘ │ + │ └───────┬───────┘ ├-------------------┤ OR │ │ │ + │ │ │ │ │ ┌─────────┴─────────┐ │ + └─────────│─────────┘ │ │ │ Grouped Query │ │ + │ └------►│ │ Attention │ │ + ┌─────┴─────┐ │ └─┬───────┬───────┬─┘ │ + │ Norm │ │ Q K V │ + └─────┬─────┘ │ ↑ ↑ ↑ │ + │ │ ┌─┴──┐ ┌─┴──┐ │ │ + ┌─────┴─────┐ │ │Norm│ │Norm│ │ │ + │ Embedding │ │ └─┬──┘ └─┬──┘ │ │ + └─────┬─────┘ │ ↑ ↑ │ │ + │ │ ┌─┴───────┴───────┴─┐ │ + Input │ │ Linear │ │ + │ └─────────┬─────────┘ │ + │ ↑ │ + └───────────│───────────┘ +``` + +## Getting Started with LFM2-24B on Cactus + +Ready to run LFM2-24B locally on your Mac? Here's how to get up and running. + +### Prerequisites + +- macOS with Apple Silicon and 16GB+ RAM (M1 or later recommended; M4 Pro with 48GB RAM for best results) +- Python 3.10+ +- CMake (`brew install cmake`) +- Git + +### 1. Clone and Build + +```bash +git clone https://github.com/cactus-compute/cactus.git +cd cactus +source ./setup + +# Build the Cactus engine (shared library for Python FFI) +cactus build --python +``` + +### 2. Download the Model + +Cactus handles downloading and converting HuggingFace models to its optimized binary format with INT4/INT8 quantization, all in one command: + +```bash +cactus download LiquidAI/LFM2-24B-A2B +``` + +### 3. Chat Interactively + +The fastest way to start chatting with the model: + +```bash +cactus run LiquidAI/LFM2-24B-A2B +``` + +This builds, downloads (if needed), and launches an interactive chat session. + +### 4. Use the Python API + +For building your own applications and agents, use the Python FFI bindings directly: + +```python +import json +from cactus import get_bundle_dir, cactus_init, cactus_complete, cactus_reset, cactus_destroy + +model = cactus_init(str(get_bundle_dir("LiquidAI/LFM2-24B-A2B")), None, False) + +# Simple chat completion +messages = json.dumps([{"role": "user", "content": "Write a Python function to sort a list"}]) +response = cactus_complete(model, messages, None, None, None) + +print(response["response"]) # Generated text +print(f"{response['decode_tps']:.1f} tokens/sec") + +# Streaming with a callback +def on_token(token, token_id): + print(token, end="", flush=True) + +cactus_complete(model, messages, None, None, on_token) + +# Clean up +cactus_reset(model) +cactus_destroy(model) +``` + +### 5. Function Calling for Agents + +Cactus supports tool use out of the box, a key building block for on-device coding agents: + +```python +tools = [ + { + "type": "function", + "function": { + "name": "run_code", + "description": "Execute Python code and return the output", + "parameters": { + "type": "object", + "properties": { + "code": {"type": "string", "description": "Python code to execute"} + }, + "required": ["code"] + } + } + } +] + +messages = json.dumps([{"role": "user", "content": "Calculate the factorial of 10"}]) +response = cactus_complete(model, messages, None, json.dumps(tools), None) + +if response["function_calls"]: + print(response["function_calls"]) # Model's tool invocation +``` + +### Cloud Handoff + +Cactus measures model confidence during generation. When the model isn't confident enough for a query, the response signals `cloud_handoff: true`, letting your agent route complex requests to a cloud API while keeping simple ones fast and local: + +```python +options = json.dumps({"confidence_threshold": 0.7}) +response = cactus_complete(model, messages, options, None, None) + +if response["cloud_handoff"]: + # Route to cloud API for this query + pass +else: + print(response["response"]) +``` + +This hybrid local-cloud pattern is what makes on-device coding agents practical: fast local inference for the majority of tasks, with automatic escalation when needed. + +## Conclusion + +LFM2-24B-A2B represents a compelling sweet spot for on-device coding. The MoE architecture activates just 2B of its 24B parameters per token, delivering quality that punches well above its compute class while keeping inference fast and memory-lean at ~200MB of running RAM. Paired with Cactus's INT4 quantization, SIMD-optimized kernels, and built-in function calling, this is a model you can actually build local coding agents on top of today. + +The pieces are coming together: models that are smart enough to be useful, efficient enough to run on a laptop, and runtimes that make it all accessible through a few lines of Python. Whether you're building a code assistant that works offline, a privacy-first dev tool, or just experimenting with what's possible without a cloud API, LFM2-24B with Cactus is a great place to start. + +Give it a try, build something, and let us know what you think. + +## See Also + +- [Cactus Engine API Reference](/docs/cactus_engine.md) — Full C API docs for completion, tool calling, and cloud handoff +- [Python Binding](/python/) — Python bindings used in the examples above +- [Fine-tuning Guide](/docs/finetuning.md) — Deploy your own LoRA fine-tunes to mobile +- [Hybrid Transcription](/blog/hybrid_transcription.md) — On-device/cloud hybrid speech transcription with Cactus +- [Runtime Compatibility](/docs/compatibility.md) — Weight versioning across Cactus releases \ No newline at end of file diff --git a/blog/parakeet.md b/blog/parakeet.md new file mode 100644 index 000000000..508c453fa --- /dev/null +++ b/blog/parakeet.md @@ -0,0 +1,291 @@ +--- +title: "Ridiculously Fast On-Device Transcription: Reviewing Parakeet CTC 1.1B with Cactus" +description: "Review of NVIDIA's Parakeet-CTC-1.1B model running locally on Mac with Cactus. Architecture breakdown, benchmarks, and transcription use cases." +keywords: ["Parakeet CTC 1.1B", "FastConformer", "speech-to-text", "on-device transcription", "Apple Silicon"] +author: "Satyajit Kumar and Henry Ndubuaku" +date: 2026-02-26 +tags: ["Parakeet", "ASR", "speech-to-text", "Apple Silicon", "Transcription"] +--- + +# Ridiculously Fast On-Device Transcription: Reviewing Parakeet CTC 1.1B with Cactus + +*By Satyajit Kumar and Henry Ndubuaku* + +[![Video Title](https://img.youtube.com/vi/x0t83tDr5S0/maxresdefault.jpg)](https://youtu.be/x0t83tDr5S0) + +Parakeet CTC 1.1B is NVIDIA’s non-autoregressive English speech-to-text model built on FastConformer. At only **1.1 billion parameters**, it is small enough to run entirely on-device while still delivering state-of-the-art transcription quality. It uses Limited Context Attention in the encoder and a lightweight CTC projection head instead of an autoregressive decoder, which makes the decoding stage extremely efficient. Using Cactus we achieve up to **6 million tokens/second** decode speed with **sub-200 ms end-to-end latency** on Apple Silicon, fast enough for real-time, always-on transcription without a cloud round-trip. + +## Architecture Details + +Parakeet CTC 1.1B is built on NVIDIA's FastConformer encoder and optimized for non-autoregressive ASR. At a high level: + +1. **Audio front-end (mel + subsampling):** Input audio is converted to log-mel features, then an 8x depthwise-separable convolutional subsampler reduces sequence length before the encoder stack. +2. **FastConformer encoder blocks:** The encoder combines Conformer layers with **Limited Context Attention (LCA)** for local efficiency and periodic **Global Tokens (GT)** so long-range context is still preserved. +3. **CTC projection head:** Instead of an autoregressive decoder, Parakeet projects encoder states directly to token logits and uses **CTC** decoding (blank/repeat collapse), making inference highly parallel and low latency. + +This architecture is why Parakeet works well for both real-time and batch transcription: most compute is in the encoder pass, and decoding stays lightweight. + +## Model Architecture Diagram + +```text + ┌───────────────────────┐ + │ CTC Collapse │ + │ remove blanks / merge │ + │ repeated labels │ + └───────────┬───────────┘ + ▲ + ┌───────────┴───────────┐ + │ CTC Projection Head │ + │ Conv1D / Linear → V │ + └───────────┬───────────┘ + ▲ + ┌───────────┴───────────┐ + │ Norm │ + └───────────┬───────────┘ + ▲ + ┌────────────────────────⊕───────────────────────┐ + │ │ │ + │ FastConformer Encoder Stack │ + │ × Num Layers │ + │ │ + │ ┌────────────────────────────────────────┐ │ + │ │ FastConformer Block │ │ + │ │ │ │ + │ │ ┌──────────────┐ │ │ + │ │ │ FFN │ │ │ + │ │ │ Linear │ │ │ + │ │ │ SwiGLU/Act │ │ │ + │ │ │ Linear │ │ │ + │ │ └──────┬───────┘ │ │ + │ │ │ │ │ + │ │ ⊕ │ │ + │ │ │ │ │ + │ │ ┌──────┴───────┐ │ │ + │ │ │ Conv Module │ │ │ + │ │ │ Pointwise │ │ │ + │ │ │ Depthwise │ │ │ + │ │ │ Pointwise │ │ │ + │ │ └──────┬───────┘ │ │ + │ │ │ │ │ + │ │ ⊕ │ │ + │ │ │ │ │ + │ │ ┌───────────────┴──────────────┐ │ │ + │ │ │ Limited Context Attention │ │ │ + │ │ │ local / sliding window │ │ │ + │ │ │ │ │ │ + │ │ │ Q K V │ │ │ + │ │ │ ↑ ↑ ↑ │ │ │ + │ │ │ ┌────┴────────┴────────┴───┐ │ │ │ + │ │ │ │ Linear │ │ │ │ + │ │ │ └─────────────┬────────────┘ │ │ │ + │ │ └───────────────┼──────────────┘ │ │ + │ │ │ │ │ + │ │ ⊕ │ │ + │ │ │ │ │ + │ │ ┌───────────────┴──────────────┐ │ │ + │ │ │ FFN │ │ │ + │ │ │ Linear → Act → Linear │ │ │ + │ │ └──────────────────────────────┘ │ │ + │ │ │ │ + │ └────────────────────────────────────────┘ │ + └────────────────────────┬───────────────────────┘ + ▲ + ┌───────────┴───────────┐ + │ Conv Subsampling / │ + │ Sequence Reduction │ + │ (time downsample) │ + └───────────┬───────────┘ + ▲ + ┌───────────┴───────────┐ + │ Mel-Spectrogram / │ + │ Acoustic Features │ + └───────────┬───────────┘ + ▲ + ┌───────────┴───────────┐ + │ 16 kHz Audio │ + │ Waveform In │ + └───────────────────────┘ +``` + +## Getting Started with Parakeet-CTC-1.1B on Cactus + +### Quick Start (Homebrew) + +The fastest way to try Parakeet: two commands, sub-200 ms latency: + +```bash +brew install cactus-compute/cactus/cactus +cactus transcribe nvidia/parakeet-ctc-1.1b +``` + +That's it. Cactus downloads the 1.1B model, quantizes it, and starts a live transcription session from your microphone. To transcribe a file instead: + +```bash +cactus transcribe nvidia/parakeet-ctc-1.1b --file /path/to/your/file.wav +``` + +### Building from Source + +If you need the Python, Rust, or C libraries for integration, build from source: + +#### Prerequisites + +- macOS with Apple Silicon and 16GB+ RAM (M1 or later recommended) +- Python 3.10+ +- CMake (`brew install cmake`) +- Git + +#### Clone and Build + +```bash +git clone https://github.com/cactus-compute/cactus.git +cd cactus + +# Build the Cactus engine (shared library for Python FFI) +cactus build --python +``` + +#### Download the Model + +Cactus handles downloading and converting HuggingFace models to its optimized binary format with INT4/INT8 quantization, all in one command: + +```bash +cactus download nvidia/parakeet-ctc-1.1b +``` + +### 4. Use the [Python Binding](/python/) + +For integrating Parakeet into your own applications, use the Python FFI bindings directly: + +```python +from cactus import get_bundle_dir, cactus_init, cactus_transcribe, cactus_destroy + +model = cactus_init(str(get_bundle_dir("nvidia/parakeet-ctc-1.1b")), None, False) + +result = cactus_transcribe(model, "/path/to/audio.wav") + +print("\n\nFinal transcript:") +print(result["response"]) +print(f"Decode speed: {result['decode_tps']:.1f} tokens/sec") + +cactus_destroy(model) +``` + +### 5. Use the [C API](/docs/cactus_engine.md) + +The C API is the base layer all other bindings build on. Link against `libcactus_engine` and include the FFI header: + +```c +#include "cactus_engine.h" +#include +#include + +int main() { + cactus_model_t model = cactus_init("weights/parakeet-ctc-1.1b", NULL, false); + + char response[16384]; + int rc = cactus_transcribe( + model, "audio.wav", NULL, + response, sizeof(response), + NULL, NULL, NULL, NULL, 0 + ); + + if (rc >= 0) printf("Transcript: %s\n", response); + + cactus_destroy(model); + return 0; +} +``` + +### 6. Use the [Rust Binding](/bindings/rust/) + +Copy `cactus.rs` into your project (see [the README](/bindings/rust/)), link `libcactus_engine.a` from `cactus build`, and call the FFI bindings directly: + +```rust +use std::ffi::CString; +use std::os::raw::c_char; +use std::ptr; + +mod cactus; + +fn main() { + let model_path = CString::new("weights/parakeet-ctc-1.1b").unwrap(); + let audio_path = CString::new("audio.wav").unwrap(); + + let model = unsafe { + cactus::cactus_init(model_path.as_ptr(), ptr::null(), false) + }; + + let mut buf = vec![0u8; 16384]; + let rc = unsafe { + cactus::cactus_transcribe( + model, + audio_path.as_ptr(), + ptr::null(), + buf.as_mut_ptr() as *mut c_char, + buf.len(), + ptr::null(), None, ptr::null_mut(), + ptr::null(), 0, + ) + }; + + if rc >= 0 { + let response = unsafe { std::ffi::CStr::from_ptr(buf.as_ptr() as *const c_char).to_string_lossy() }; + println!("Transcript: {}", response); + } + + unsafe { cactus::cactus_destroy(model) }; +} +``` + +### 7. Use the [Swift Binding](/bindings/swift/) + +The Swift binding exposes top-level functions that map directly to the C FFI: + +```swift +import cactus + +let model = cactus_init("weights/parakeet-ctc-1.1b", nil, false) +var buf = [CChar](repeating: 0, count: 65536) +cactus_transcribe(model, "/path/to/audio.wav", nil, &buf, buf.count, nil, nil, nil, nil, 0) +let resultJson = String(cString: buf) +print(resultJson) + +cactus_destroy(model) +``` + +### 8. Use the [Kotlin Binding](/bindings/kotlin/) + +The Kotlin binding exposes top-level functions that map directly to the C FFI: + +```kotlin +import com.cactus.* + +val model = cactusInit("weights/parakeet-ctc-1.1b", null, false) +val resultJson = cactusTranscribe(model, "/path/to/audio.wav", "", null, null, null) +println(resultJson) + +cactusDestroy(model) +``` + +### 9. Use the [Flutter Binding](/bindings/flutter/) + +The Flutter binding brings Cactus transcription to iOS, macOS, and Android: + +```dart +import 'cactus.dart'; + +final model = cactusInit('weights/parakeet-ctc-1.1b', null, false); +final resultJson = cactusTranscribe(model, '/path/to/audio.wav', null, null, null, null); +print(resultJson); + +cactusDestroy(model); +``` + +## See Also + +- [Cactus Engine API Reference](/docs/cactus_engine.md) — Full C API docs for completion, tool calling, and cloud handoff +- [Python Binding](/python/) — Python bindings used in the examples above +- [Hybrid Transcription](/blog/hybrid_transcription.md) — On-device/cloud hybrid speech transcription with Cactus +- [LFM2-24B-A2B](/blog/lfm2_24b_a2b.md) - Reviewing LFM2 24B MoE A2B with Cactus +- [Runtime Compatibility](/docs/compatibility.md) — Weight versioning across Cactus releases diff --git a/blog/turboquant-h.md b/blog/turboquant-h.md new file mode 100644 index 000000000..fccb48c89 --- /dev/null +++ b/blog/turboquant-h.md @@ -0,0 +1,296 @@ +--- +title: "TurboQuant-H: Hadamard Rotation for 2-Bit Embedding Quantization" +description: "TurboQuant compresses KV cache to 1-3 bits. We introduce TurboQuant-H, a simplified offline variant using Hadamard rotation and per-group Lloyd-Max codebooks, applied to per-layer embedding tables where it matters most: models where embeddings are 60%+ of total weight storage." +keywords: ["TurboQuant-H", "TurboQuant", "vector quantization", "embeddings", "2-bit", "on-device AI", "Gemma 4", "Gemma 3n", "per-layer embeddings", "AltUp", "Hadamard", "memory reduction"] +author: "Karen Mosoyan & Henry Ndubuaku" +date: 2026-04-21 +tags: ["quantization", "embeddings", "TurboQuant-H", "Gemma", "on-device", "memory optimization"] +--- + +# TurboQuant-H: Hadamard Rotation for 2-Bit Embedding Quantization + +*By Karen Mosoyan & Henry Ndubuaku* + +*(karen@cactuscompute.com, henry@cactuscompute.com)* + +## Abstract + +TurboQuant (Zandieh et al., ICLR 2026) compresses KV cache vectors to 1-3 bits via random orthogonal rotation, optimal scalar quantization, and QJL bias correction. We introduce **TurboQuant-H**, a simplified offline variant that replaces random rotation with Hadamard rotation, uses per-group Lloyd-Max codebooks, and drops the QJL correction stage. We apply TurboQuant-H to per-layer input (PLI) embedding tables in Gemma 4 E2B, where embeddings constitute 60.6% of total model weight. On Gemma 4 E2B, TurboQuant-H compresses PLI weights from 2,496 MB to 624 MB (4x) at 2.125 effective bits per dimension, reducing total LLM storage by 40% (4,790 MB → 2,918 MB) with a perplexity increase of 0.06 (1.85 → 1.91) and no measured speed regression. + +## 1. Introduction + +TurboQuant (ICLR 2026) compresses KV cache vectors to 1-3 bits with near-zero quality loss. The technique is elegant: rotate vectors with a random orthogonal matrix, exploit the resulting Beta distribution to apply optimal scalar quantizers per coordinate, then correct inner product bias with a 1-bit QJL residual. The paper demonstrates quality neutrality at 3.5 bits and marginal degradation at 2.5 bits on Llama-3.1-8B and Ministral-7B. + +But TurboQuant was designed for KV cache, vectors generated at runtime during inference. There's a catch: mobile devices and wearables need small models, which we found to significantly degrade when KV cache goes below INT4. We in fact keep KV cache at INT8 on Cactus to ensure correctness. This makes applying TurboQuant to Cactus KV workloads tricky. + +However, with the emergence of per-layer embedding architectures (each layer has its own embedding lookup), these embeddings dominate the parameter count of models like the Gemma E-series. For instance, Gemma E2B has 2.3B effective parameters but 5.1B total, because the per-layer embeddings alone account for the difference. That bloats memory and storage footprint by more than 2x. There is a need to re-visit embedding quantisation. + +## 2. Background: Per-Layer Embeddings Dominate Model Storage + +Most quantization research focuses on linear layer weights and activations. Embeddings are treated as untouchable lookup tables, typically kept at FP16 or at best INT8 while everything else goes to INT4. This made sense when embeddings were a small fraction of total parameters. That assumption broke with per-layer embedding architectures. + +Gemma 4 E2B uses AltUp, a technique where each of the 35 transformer layers gets its own embedding projection from the 262K-token vocabulary. Instead of one shared embedding table, you have a shared table plus a per-layer table. The numbers on Cactus's current INT4 weights: + +| Component | Size | % of Model | +|---|---|---| +| `token_embeddings` (shared) | 408 MB | 8.7% | +| `embed_tokens_per_layer` (35 layers) | 2,496 MB | 52.1% | +| **Total embedding storage** | **2,904 MB** | **60.6%** | +| All other weights (attention, FFN, norms, encoders) | 1,886 MB | 39.4% | +| **Total model** | **4,790 MB** | 100% | + +The per-layer embedding table is 2.5 GB. More than half the model. This is not unique to Gemma 4. The AltUp design pattern, where per-layer vocabulary projections replace a single shared embedding, is becoming standard for models that need large vocabularies (262K tokens for multilingual coverage) without proportionally large hidden dimensions. Gemma 3n uses the same architecture. Any model that follows this pattern will be embedding-dominated. + +## 3. TurboQuant-H + +TurboQuant-H shares the core insight from TurboQuant; rotation concentrates coordinates into a well-behaved distribution, enabling aggressive scalar quantization, but simplifies the pipeline for offline weight quantization. + +### 3.1 Comparison with TurboQuant + +| | TurboQuant (Zandieh et al.) | TurboQuant-H (this work) | +|---|---|---| +| **Target** | KV cache (runtime activations) | Embedding weight tables (offline) | +| **Rotation** | Random orthogonal matrix via QR of Gaussian, $O(d^2)$ | Normalized Hadamard matrix, $O(N \log N)$, symmetric = self-inverse | +| **Quantizer** | Per-coordinate scalar quantizer (precomputed for Beta distribution) | Per-position Lloyd-Max codebook (trained on actual weight distribution) | +| **Codebook** | Implicit (quantization levels derived from Beta CDF) | Explicit FP16 centroids per position group (0.125 bits overhead at group-128) | +| **Bias correction** | Two-stage: MSE quantizer at $b-1$ bits + 1-bit QJL residual | Single-stage: no QJL correction | +| **When it runs** | Every forward pass during inference | Once during weight conversion | +| **Bit width** | 2.5-bit and 3.5-bit | 2-bit (+0.125 codebook overhead ≈ 2.125 effective) | + +### 3.2 Formal Description + +Let $\mathbf{E} \in \mathbb{R}^{V \times D}$ be the PLI embedding matrix with vocabulary size $V = 262{,}144$ and embedding dimension $D = 8{,}190$. We partition each row into $P = \lceil D/G \rceil = 64$ positional groups of $G = 128$ contiguous elements. The $p$-th positional group of vocabulary row $v$ is denoted $\mathbf{x}_{v,p} \in \mathbb{R}^G$. + +**Quantization** (offline, during weight conversion): + +**Step 1: Hadamard rotation.** For each row $v$ and position $p$, rotate: + +$$\hat{\mathbf{x}}_{v,p} = \bar{\mathbf{H}}_G \cdot \mathbf{x}_{v,p}$$ + +where $\bar{\mathbf{H}}_G = \frac{1}{\sqrt{G}} \mathbf{H}_G$ is the $G \times G$ normalized Hadamard matrix, satisfying $\bar{\mathbf{H}}_G^T \bar{\mathbf{H}}_G = \mathbf{I}$ and $\bar{\mathbf{H}}_G = \bar{\mathbf{H}}_G^T$. The $\frac{1}{\sqrt{G}}$ normalization ensures the transform is its own inverse. + +**Step 2: Codebook training.** For each positional group $p \in \{1, \ldots, P\}$, collect the rotated values from all $V$ vocabulary rows and train a Lloyd-Max codebook $\mathcal{C}_p = \{c_1, c_2, c_3, c_4\}$ (at $b = 2$ bits, 4 centroids) by minimizing: + +$$\mathcal{C}_p^* = \arg\min_{\mathcal{C}} \sum_{v=1}^{V} \sum_{i=1}^{G} \min_{c \in \mathcal{C}} \left( \hat{x}_{v,p,i} - c \right)^2$$ + +This trains each codebook on $V \times G = 262{,}144 \times 128 \approx 33.6\text{M}$ data points, giving the Lloyd-Max algorithm sufficient statistics for accurate centroid placement. Each positional group gets its own codebook because the weight distributions vary across positions in the embedding dimension. We tried a single joint codebook shared across all positions; per-position was consistently better. + +The codebook values are stored at FP16 (16 bits per centroid), contributing: + +$$\frac{2^b \cdot 16}{G} = \frac{4 \cdot 16}{128} = 0.125 \text{ bits/element overhead}$$ + +**Step 3: Quantize by proximity.** Each rotated element maps to its nearest centroid: + +$$q_{v,p,i} = \arg\min_{j \in \{1,\ldots,2^b\}} \left| \hat{x}_{v,p,i} - c_j \right|$$ + +Store the 2-bit indices $q_{v,p,i}$ and the $P = 64$ FP16 codebooks $\{\mathcal{C}_1, \ldots, \mathcal{C}_P\}$. + +**Dequantization** (at inference): + +$$\tilde{\mathbf{x}}_{v,p} = \bar{\mathbf{H}}_G \cdot \text{scatter}(\mathcal{C}_p, \mathbf{q}_{v,p})$$ + +where $\text{scatter}(\mathcal{C}_p, \mathbf{q}_{v,p})_i = c_{q_{v,p,i}}$ maps indices back to centroids. Since $\bar{\mathbf{H}}_G$ is symmetric and orthogonal, the inverse rotation is the same forward transform. No transpose is needed. + +**Effective bit rate:** + +$$b_{\text{eff}} = b + \frac{2^b \cdot 16}{G} = 2 + 0.125 = 2.125 \text{ bits/element}$$ + +### 3.3 The Quantization Pipeline + +``` +QUANTIZATION (offline, during cactus convert) +============================================== + +PLI Matrix E (262K x 8190) + | + v ++-----------------------+ +| Partition into | Each row -> 64 positional groups +| groups of G=128 | of 128 contiguous elements ++-----------+-----------+ + | + v ++-----------------------+ +| Hadamard rotation | x_hat = (1/sqrt(G)) * H_128 * x +| per group | O(G log G) butterfly ++-----------+-----------+ + | + v ++-----------------------+ +| Lloyd-Max codebook | Train 4 centroids (2-bit) per position +| per position | across all 262K vocab rows +| | C_p = {c1, c2, c3, c4} in FP16 ++-----------+-----------+ + | + v ++-----------------------+ +| Quantize by | q = argmin_j |x_hat_i - c_j| +| proximity | Store 2-bit indices per element ++-----------+-----------+ + | + v +Output: 2-bit index tensor + 64 FP16 codebooks +Effective: 2.125 bits/element + + +DEQUANTIZATION (at inference, per token) +========================================= + +Token IDs + | + v ++-----------------------+ +| Gather 2-bit indices | Look up row from compressed table +| + codebook per pos. | ~3.8x less bandwidth than INT8 ++-----------+-----------+ + | + v ++-----------------------+ +| Scatter codebook | Replace 2-bit indices with FP16 +| values | centroid values from C_p ++-----------+-----------+ + | + v ++-----------------------+ +| Hadamard rotation | x_tilde = (1/sqrt(G)) * H_128 * scatter(...) +| (same as forward, | H_bar is symmetric: H_bar = H_bar^T = H_bar^-1 +| no transpose needed) | O(G log G) butterfly per group ++-----------+-----------+ + | + v +FP16 embedding -> feed to transformer layer +``` + +### 3.4 Design Decisions + +**Why Hadamard instead of random orthogonal?** The normalized Hadamard matrix $\bar{\mathbf{H}}_G = \frac{1}{\sqrt{G}} \mathbf{H}_G$ is deterministic, $O(N \log N)$ to apply via the butterfly factorization (same structure as the FFT), and its own inverse ($\bar{\mathbf{H}} = \bar{\mathbf{H}}^T = \bar{\mathbf{H}}^{-1}$). For offline weight quantization we don't need the data-oblivious guarantees of a random rotation, we have full access to the weight data at conversion time. The Hadamard rotation still concentrates coordinates, which is all we need to make low-bit scalar quantization work. Note: the unnormalized Hadamard satisfies $\mathbf{H}\mathbf{H}^T = G \cdot \mathbf{I}$, so the $\frac{1}{\sqrt{G}}$ factor is essential for the self-inverse property. + +**Why no QJL correction?** TurboQuant's second stage exists because MSE-optimal quantizers introduce multiplicative bias in inner product estimation. At 1-bit, $\mathbb{E}[\langle \mathbf{y}, Q(\mathbf{x}) \rangle] = \frac{2}{\pi} \langle \mathbf{y}, \mathbf{x} \rangle$, a 36% shrinkage. The QJL residual corrects this at the cost of 1 additional bit per dimension. But we're quantizing at 2 bits with a trained codebook, not a precomputed scalar quantizer. The per-position Lloyd-Max codebook already minimizes distortion over the actual weight distribution, and the Hadamard rotation ensures the codebook sees well-spread inputs. At 2 bits with group-128, the inner product bias is small enough that the downstream perplexity impact is negligible (PPL 1.91 vs 1.85). Adding QJL would cost an extra bit per dimension for a correction that isn't needed at this operating point. + +**Why per-position codebooks instead of per-coordinate?** TurboQuant can use a single precomputed quantizer because random rotation makes all coordinates identically distributed (each coordinate of a uniform unit-sphere vector follows a known distribution that converges to $\mathcal{N}(0, 1/d)$ in high dimensions). Hadamard rotation concentrates coordinates but doesn't make them identically distributed — there are structured patterns from the butterfly network. Per-position codebooks (one codebook per group of 128 dimensions, trained across all 262K vocabulary rows at that position) adapt to these patterns. We tried a single joint codebook shared across all positions; per-position was consistently better. + +**Why group size 128?** We swept group sizes from 32 to 512. The qualitative tradeoffs: + +| Group size | Codebook overhead (bits/elem) | Hadamard cost | Quality | +|---|---|---|---| +| 32 | 0.500 | fastest | degraded (high overhead eats bit budget) | +| 64 | 0.250 | fast | good | +| **128** | **0.125** | **fast** | **best (sweet spot)** | +| 256 | 0.063 | moderate | slightly worse (less uniform within group) | +| 512 | 0.031 | slow | worse (distribution spreads too thin) | + +Group-128 gives 0.125 bits overhead, a fast butterfly, and the best quality. Detailed per-group-size PPL numbers are a target for future work. + +## 4. Results + +### 4.1 Perplexity + +Evaluated on Gemma 4 E2B. Evaluation set: 128 self-generated WildChat-1M completions from our `trajectories.jsonl` calibration set, completion-only NLL, 24,438 scored tokens. + +| Variant | Avg bits | PPL | +|---|---|---| +| HuggingFace BF16 | 16 | 1.2892 | +| Cactus default (INT4 linears + INT8 PLI + INT8 token-emb) | ~6.3 | 1.8547 | +| **Cactus + TurboQuant-H PLI** | **~3.8** | **1.9111** | + +Perplexity moves from 1.85 to 1.91. A delta of 0.06 PPL on a 24K-token eval set, within noise for practical use. No measured speed regression. + +### 4.2 Disk Footprint + +| Variant | Size (MB) | Factor | +|---|---|---| +| HuggingFace FP16 snapshot | ~10,240 | 1.00× | +| Cactus default (INT4 linears + INT8 PLI + INT8 emb) | ~4,790 | 0.47× | +| **Cactus + TurboQuant-H PLI** | **2,918** | **0.29×** | + +The PLI table specifically: **2,496 MB → 624 MB, a 4× reduction.** + +Total LLM weight reduction: **40%** from the Cactus baseline. Including the vision and audio encoders (untouched by this change), the overall model reduction is **30%**. + +For Gemma 4 E2B, that's the difference between a 4.8 GB model and a 2.9 GB model. On a 4 GB RAM Android device, that's the difference between fitting and not fitting. + +For Gemma-270m, where the embedding table (295M params) is larger than all other weights combined, we expect the same technique to cut total model size roughly in half. Validation on Gemma-270m is planned future work. + +## 5. Inference Path + +### 5.1 Before (Cactus default) + +``` +token_ids → gather from INT8 table → dequantize to FP16 → feed to transformer +``` + +### 5.2 After (Cactus TurboQuant-H) + +``` +token_ids → gather 2-bit indices → scatter codebook → Hadamard rotation → FP16 → transformer +``` + +### 5.3 Overhead Analysis + +The Hadamard butterfly on a group of $G = 128$ elements has $\log_2(G) = 7$ stages. Each stage performs $G/2 = 64$ add/subtract pairs (the butterfly operations). Over 7 stages that is $7 \times 64 = 448$ add/sub operations per group. On ARM NEON (128-bit SIMD, 8 FP16 lanes), each stage processes the group in $G / (2 \times 8) = 8$ vector instructions, giving $7 \times 8 = 56$ vector operations per group. + +For a single PLI embedding row of $D = 8{,}190$ elements: + +$$\text{Groups per row} = \lceil 8190 / 128 \rceil = 64$$ + +$$\text{Total vector ops} = 64 \times 56 = 3{,}584$$ + +At one vector add/sub per cycle on a 2 GHz A15 NEON unit, this completes in under 2 microseconds per embedding lookup. The gather from a table that is $8 / 2.125 \approx 3.8\times$ smaller in memory more than compensates for the rotation cost in bandwidth savings. + +## 6. Related Work + +**TurboQuant** (Zandieh et al., ICLR 2026) introduced data-oblivious vector quantization with random rotation and QJL correction for KV cache compression, achieving near-optimal distortion rates within a constant factor of $\approx 2.7$ from information-theoretic lower bounds. + +**QuIP#** (Tseng et al., 2024) uses the randomized Hadamard transform for weight quantization of linear layers, but does not address embedding tables or per-layer embeddings. + +**GPTQ, AWQ** focus on linear layer weight quantization with calibration data. These methods do not handle embedding tables, which are pure lookup operations with no gradient flow during inference. + +To our knowledge, TurboQuant-H is the first application of rotation-based vector quantization specifically to per-layer embedding tables, which is where the technique yields the largest storage benefit due to the embedding-dominated weight distribution in AltUp architectures. + +## 7. Next Steps + +This is a research preview of ongoing work at Cactus. Here's what's coming: + +1. **Downstream task validation.** The 0.06 PPL increase is promising, but we're running full evals (MMLU, IFEval, GPQA) across Gemma-270m, Gemma 3n, and Gemma 4 E4B to confirm TurboQuant-H holds up on real tasks, not just perplexity. +2. **Extend to shared `token_embeddings`.** TurboQuant-H currently targets PLI tables only. Applying it to the shared embedding table (408 MB) is straightforward and would push total compression further. +3. **Quantitative group size sweep.** Section 3.4 reports qualitative tradeoffs. We're collecting per-group-size PPL numbers to give a rigorous recommendation. +4. **1-bit with QJL correction.** At 1-bit, the theoretical compression reaches 8x on PLI tables. We're evaluating whether reintroducing the QJL residual stage at this extreme bit width recovers enough quality to be practical. +5. **Ship TurboQuant-H weights.** Integrate the quantization path into `cactus convert` and publish pre-quantized weights on HuggingFace for all supported Gemma E-series models. + +## Try It + +Run Gemma 4 today on Cactus: + +```bash +brew install cactus-compute/cactus/cactus +cactus run google/gemma-4-E2B-it +``` + +TurboQuant-H PLI weights will ship in an upcoming release. If you're working on embedding quantization or have thoughts on extending this to the shared token embedding table, open an issue on [GitHub](https://github.com/cactus-compute/cactus). + +## Citation + +If you use TurboQuant-H in your research, please cite: + +```bibtex +@article{turboquant-h, + title = {TurboQuant-H: Hadamard Rotation for 2-Bit Embedding + Quantization in Embedding-Dominated Models}, + author = {Mosoyan, Karen and Ndubuaku, Henry}, + year = {2026}, + url = {https://docs.cactuscompute.com/latest/blog/turboquant-h/}, + note = {Cactus Compute Research Preview} +} +``` + +## References + +- Zandieh, Daliri, Hadian, Mirrokni. [TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate](https://arxiv.org/abs/2504.19874). ICLR 2026. +- Tseng et al. QuIP#: Even Better LLM Quantization with Hadamard Incoherence and Lattice Codebooks. 2024. +- Google Gemma Team. Gemma 3n / Gemma 4 Technical Reports. 2026. + +## See Also + +- [Gemma 4 on Cactus](/blog/gemma4.md) — Day-one multimodal support with vision, audio, and hybrid inference +- [LFM-2.5-350m on Cactus](/blog/lfm2.5_350m.md) — INT8 quantization deep dive and zero-copy loading +- [Cactus Engine API](/docs/cactus_engine.md) — Full C API reference diff --git a/cactus-code/.gitignore b/cactus-code/.gitignore new file mode 100644 index 000000000..76add878f --- /dev/null +++ b/cactus-code/.gitignore @@ -0,0 +1,2 @@ +node_modules +dist \ No newline at end of file diff --git a/cactus-code/LICENSE b/cactus-code/LICENSE new file mode 100644 index 000000000..b0a8e9b81 --- /dev/null +++ b/cactus-code/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Mario Zechner + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/cactus-code/README.md b/cactus-code/README.md new file mode 100644 index 000000000..7aa51eca0 --- /dev/null +++ b/cactus-code/README.md @@ -0,0 +1,17 @@ +# cactus-code + +The terminal coding agent behind `cactus code` — an interactive AI assistant +for a local [Cactus](https://github.com/cactus-compute/cactus) server. + +## Credits + +Based on **[pi](https://github.com/earendil-works/pi)** by **Mario Zechner** +and contributors. The agent — its TUI, sessions, tools, and compaction — is +their work. Thank you. + +Want the original, full-featured agent? Use upstream: +https://github.com/earendil-works/pi + +## License + +MIT — see [LICENSE](./LICENSE), © 2025 Mario Zechner. Same license as upstream. diff --git a/cactus-code/package-lock.json b/cactus-code/package-lock.json new file mode 100644 index 000000000..681a5975c --- /dev/null +++ b/cactus-code/package-lock.json @@ -0,0 +1,1555 @@ +{ + "name": "pi-monorepo", + "version": "0.0.3", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "pi-monorepo", + "version": "0.0.3", + "workspaces": [ + "packages/*" + ], + "devDependencies": { + "@types/node": "22.19.19", + "@typescript/native-preview": "7.0.0-dev.20260120.1", + "shx": "0.4.0" + }, + "engines": { + "node": ">=22.19.0" + } + }, + "node_modules/@earendil-works/pi-agent-core": { + "resolved": "packages/agent", + "link": true + }, + "node_modules/@earendil-works/pi-ai": { + "resolved": "packages/ai", + "link": true + }, + "node_modules/@earendil-works/pi-coding-agent": { + "resolved": "packages/coding-agent", + "link": true + }, + "node_modules/@earendil-works/pi-tui": { + "resolved": "packages/tui", + "link": true + }, + "node_modules/@mariozechner/clipboard": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard/-/clipboard-0.3.9.tgz", + "integrity": "sha512-ABnA53mdfkGZwOFUdZNv2S0CWGO/EIuPj8Vv9xmBFmSYg/qFc7ihO6q5FcQjvoE67kZpWkEc4AhD6B/os04yuA==", + "license": "MIT", + "optional": true, + "engines": { + "node": ">= 10" + }, + "optionalDependencies": { + "@mariozechner/clipboard-darwin-arm64": "0.3.9", + "@mariozechner/clipboard-darwin-universal": "0.3.9", + "@mariozechner/clipboard-darwin-x64": "0.3.9", + "@mariozechner/clipboard-linux-arm64-gnu": "0.3.9", + "@mariozechner/clipboard-linux-arm64-musl": "0.3.9", + "@mariozechner/clipboard-linux-riscv64-gnu": "0.3.9", + "@mariozechner/clipboard-linux-x64-gnu": "0.3.9", + "@mariozechner/clipboard-linux-x64-musl": "0.3.9", + "@mariozechner/clipboard-win32-arm64-msvc": "0.3.9", + "@mariozechner/clipboard-win32-x64-msvc": "0.3.9" + } + }, + "node_modules/@mariozechner/clipboard-darwin-arm64": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-darwin-arm64/-/clipboard-darwin-arm64-0.3.9.tgz", + "integrity": "sha512-BfgV7vCEWZwJwZJw03r6bP5+tf0iI/ANuQYCxi9RNn7FrWB3yzGuMKCrNLRl6V761vXRdL8+OqZ0wd4TqlsNOQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-darwin-universal": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-darwin-universal/-/clipboard-darwin-universal-0.3.9.tgz", + "integrity": "sha512-BGGR4iA9Z2shAjI65eI5xtyb3LYNlDW9X3gxKxDbqtbnREohsrqznov6zpKoIrsRWpzlYVEdKphS7ksJ0/ndSQ==", + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-darwin-x64": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-darwin-x64/-/clipboard-darwin-x64-0.3.9.tgz", + "integrity": "sha512-4kURmCbS6nt8uYhtmWpUcJWyPHfmAr5dTpXD1nO3pIfa+TSQ9DbrGOYCKH+aEFW47XhQ4Vp8ZTszie+wfFvDKg==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-linux-arm64-gnu": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-linux-arm64-gnu/-/clipboard-linux-arm64-gnu-0.3.9.tgz", + "integrity": "sha512-g59OkUGP2DDfCOIKypHeYgv2M55u/cKvXa5dSxFbEJ34XvIQMdcVmpKCkGUro3ZgefXiGVdwguvTMQGpHWzIXw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-linux-arm64-musl": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-linux-arm64-musl/-/clipboard-linux-arm64-musl-0.3.9.tgz", + "integrity": "sha512-AGuJdgKsmJdm4Pych7kv3sqe591ERRaAHW3xjLooiFzn8J+PxUyof++7YZrB5Y5tpnTO+K18Og3taj2NpluCRQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-linux-riscv64-gnu": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-linux-riscv64-gnu/-/clipboard-linux-riscv64-gnu-0.3.9.tgz", + "integrity": "sha512-DXBEAiuMpk7dhS1a9NzNxVAFi1vaKoPu7rQNgY8LIDLGrK3lnIp3nT10DUum+PKVJoJppIP+NAA8IZe4DMNDPw==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-linux-x64-gnu": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-linux-x64-gnu/-/clipboard-linux-x64-gnu-0.3.9.tgz", + "integrity": "sha512-WORrMLd6EpElEME7JRKfSaY34nW1P5LbdgK5YNCS1ncG2LqmITsSMEJ8nh2mpvxb3TxqbOOKgY7k9eMJYlW9Mw==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-linux-x64-musl": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-linux-x64-musl/-/clipboard-linux-x64-musl-0.3.9.tgz", + "integrity": "sha512-/DHn+1DrfL6oRaPPWXaOKvonFFrni666fxd+zFqiQEfvBH0tsHVWjq9iqBk0oDp0qaPA72lIMy5BptxISBEhZQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-win32-arm64-msvc": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-win32-arm64-msvc/-/clipboard-win32-arm64-msvc-0.3.9.tgz", + "integrity": "sha512-O5FHD3ErkMwMhNzAfu3ggy0ug4z7btZuoQgwwxlzPrwV2bxlD6WDpqBY4NCgICAgZdDKdp+loUEKVAVt8aYnhQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@mariozechner/clipboard-win32-x64-msvc": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@mariozechner/clipboard-win32-x64-msvc/-/clipboard-win32-x64-msvc-0.3.9.tgz", + "integrity": "sha512-ihQC3EufqEY81vhXBgVBtK4prL+wc62zJsSvxrgz7K1hsdt6OObz6v9p3Rn1OG3GJksTTKMJF0u/guMISHPhSA==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@nodelib/fs.scandir": { + "version": "2.1.5", + "resolved": "https://registry.npmjs.org/@nodelib/fs.scandir/-/fs.scandir-2.1.5.tgz", + "integrity": "sha512-vq24Bq3ym5HEQm2NKCr3yXDwjc7vTsEThRDnkp2DK9p1uqLR+DHurm/NOTo0KG7HYHU7eppKZj3MyqYuMBf62g==", + "dev": true, + "license": "MIT", + "dependencies": { + "@nodelib/fs.stat": "2.0.5", + "run-parallel": "^1.1.9" + }, + "engines": { + "node": ">= 8" + } + }, + "node_modules/@nodelib/fs.stat": { + "version": "2.0.5", + "resolved": "https://registry.npmjs.org/@nodelib/fs.stat/-/fs.stat-2.0.5.tgz", + "integrity": "sha512-RkhPPp2zrqDAQA/2jNhnztcPAlv64XdhIp7a7454A5ovI7Bukxgt7MX7udwAu3zg1DcpPU0rz3VV1SeaqvY4+A==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 8" + } + }, + "node_modules/@nodelib/fs.walk": { + "version": "1.2.8", + "resolved": "https://registry.npmjs.org/@nodelib/fs.walk/-/fs.walk-1.2.8.tgz", + "integrity": "sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==", + "dev": true, + "license": "MIT", + "dependencies": { + "@nodelib/fs.scandir": "2.1.5", + "fastq": "^1.6.0" + }, + "engines": { + "node": ">= 8" + } + }, + "node_modules/@silvia-odwyer/photon-node": { + "version": "0.3.4", + "resolved": "https://registry.npmjs.org/@silvia-odwyer/photon-node/-/photon-node-0.3.4.tgz", + "integrity": "sha512-bnly4BKB3KDTFxrUIcgCLbaeVVS8lrAkri1pEzskpmxu9MdfGQTy8b8EgcD83ywD3RPMsIulY8xJH5Awa+t9fA==", + "license": "Apache-2.0" + }, + "node_modules/@types/cross-spawn": { + "version": "6.0.6", + "resolved": "https://registry.npmjs.org/@types/cross-spawn/-/cross-spawn-6.0.6.tgz", + "integrity": "sha512-fXRhhUkG4H3TQk5dBhQ7m/JDdSNHKwR2BBia62lhwEIq9xGiQKLxd6LymNhn47SjXhsUEPmxi+PKw2OkW4LLjA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*" + } + }, + "node_modules/@types/diff": { + "version": "7.0.2", + "resolved": "https://registry.npmjs.org/@types/diff/-/diff-7.0.2.tgz", + "integrity": "sha512-JSWRMozjFKsGlEjiiKajUjIJVKuKdE3oVy2DNtK+fUo8q82nhFZ2CPQwicAIkXrofahDXrWJ7mjelvZphMS98Q==", + "dev": true, + "license": "MIT" + }, + "node_modules/@types/hosted-git-info": { + "version": "3.0.5", + "resolved": "https://registry.npmjs.org/@types/hosted-git-info/-/hosted-git-info-3.0.5.tgz", + "integrity": "sha512-Dmngh7U003cOHPhKGyA7LWqrnvcTyILNgNPmNCxlx7j8MIi54iBliiT8XqVLIQ3GchoOjVAyBzNJVyuaJjqokg==", + "dev": true, + "license": "MIT" + }, + "node_modules/@types/ms": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/@types/ms/-/ms-2.1.0.tgz", + "integrity": "sha512-GsCCIZDE/p3i96vtEqx+7dBUGXrc7zeSK3wwPHIaRThS+9OhWIXRqzs4d6k1SVU8g91DrNRWxWUGhp5KXQb2VA==", + "dev": true, + "license": "MIT" + }, + "node_modules/@types/node": { + "version": "22.19.19", + "resolved": "https://registry.npmjs.org/@types/node/-/node-22.19.19.tgz", + "integrity": "sha512-dyh/xO2Fh5bYrfWaaqGrRQQGkNdmYw6AmaAUvYeUMNTWQtvb796ikLdmTchRmOlOiIJ1TDXfWgVx1QkUlQ6Hew==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/@types/proper-lockfile": { + "version": "4.1.4", + "resolved": "https://registry.npmjs.org/@types/proper-lockfile/-/proper-lockfile-4.1.4.tgz", + "integrity": "sha512-uo2ABllncSqg9F1D4nugVl9v93RmjxF6LJzQLMLDdPaXCUIDPeOJ21Gbqi43xNKzBi/WQ0Q0dICqufzQbMjipQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/retry": "*" + } + }, + "node_modules/@types/retry": { + "version": "0.12.5", + "resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.5.tgz", + "integrity": "sha512-3xSjTp3v03X/lSQLkczaN9UIEwJMoMCA1+Nb5HfbJEQWogdeQIyVtTvxPXDQjZ5zws8rFQfVfRdz03ARihPJgw==", + "dev": true, + "license": "MIT" + }, + "node_modules/@types/semver": { + "version": "7.7.1", + "resolved": "https://registry.npmjs.org/@types/semver/-/semver-7.7.1.tgz", + "integrity": "sha512-FmgJfu+MOcQ370SD0ev7EI8TlCAfKYU+B4m5T3yXc1CiRN94g/SZPtsCkk506aUDtlMnFZvasDwHHUcZUEaYuA==", + "dev": true, + "license": "MIT" + }, + "node_modules/@typescript/native-preview": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview/-/native-preview-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-nnEf37C9ue7OBRnF2zmV/OCBmV5Y7T/K4mCHa+nxgiXcF/1w8sA0cgdFl+gHQ0mysqUJ+Bu5btAMeWgpLyjrgg==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsgo": "bin/tsgo.js" + }, + "optionalDependencies": { + "@typescript/native-preview-darwin-arm64": "7.0.0-dev.20260120.1", + "@typescript/native-preview-darwin-x64": "7.0.0-dev.20260120.1", + "@typescript/native-preview-linux-arm": "7.0.0-dev.20260120.1", + "@typescript/native-preview-linux-arm64": "7.0.0-dev.20260120.1", + "@typescript/native-preview-linux-x64": "7.0.0-dev.20260120.1", + "@typescript/native-preview-win32-arm64": "7.0.0-dev.20260120.1", + "@typescript/native-preview-win32-x64": "7.0.0-dev.20260120.1" + } + }, + "node_modules/@typescript/native-preview-darwin-arm64": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-darwin-arm64/-/native-preview-darwin-arm64-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-r3pWFuR2H7mn6ScwpH5jJljKQqKto0npVuJSk6pRwFwexpTyxOGmJTZJ1V0AWiisaNxU2+CNAqWFJSJYIE/QTg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ] + }, + "node_modules/@typescript/native-preview-darwin-x64": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-darwin-x64/-/native-preview-darwin-x64-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-cuC1+wLbUP+Ip2UT94G134fqRdp5w3b3dhcCO6/FQ4yXxvRNyv/WK+upHBUFDaeSOeHgDTyO9/QFYUWwC4If1A==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ] + }, + "node_modules/@typescript/native-preview-linux-arm": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-linux-arm/-/native-preview-linux-arm-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-vN6OYVySol/kQZjJGmAzd6L30SyVlCgmCXS8WjUYtE5clN0YrzQHop16RK29fYZHMxpkOniVBtRPxUYQANZBlQ==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ] + }, + "node_modules/@typescript/native-preview-linux-arm64": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-linux-arm64/-/native-preview-linux-arm64-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-zZGvEGY7wcHYefMZ87KNmvjN3NLIhsCMHEpHZiGCS3khKf+8z6ZsanrzCjOTodvL01VPyBzHxV1EtkSxAcLiQg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ] + }, + "node_modules/@typescript/native-preview-linux-x64": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-linux-x64/-/native-preview-linux-x64-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-JBfNhWd/asd5MDeS3VgRvE24pGKBkmvLub6tsux6ypr+Yhy+o0WaAEzVpmlRYZUqss2ai5tvOu4dzPBXzZAtFw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ] + }, + "node_modules/@typescript/native-preview-win32-arm64": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-win32-arm64/-/native-preview-win32-arm64-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-tTndRtYCq2xwgE0VkTi9ACNiJaV43+PqvBqCxk8ceYi3X36Ve+CCnwlZfZJ4k9NxZthtrAwF/kUmpC9iIYbq1w==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "win32" + ] + }, + "node_modules/@typescript/native-preview-win32-x64": { + "version": "7.0.0-dev.20260120.1", + "resolved": "https://registry.npmjs.org/@typescript/native-preview-win32-x64/-/native-preview-win32-x64-7.0.0-dev.20260120.1.tgz", + "integrity": "sha512-oZia7hFL6k9pVepfonuPI86Jmyz6WlJKR57tWCDwRNmpA7odxuTq1PbvcYgy1z4+wHF1nnKKJY0PMAiq6ac18w==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "win32" + ] + }, + "node_modules/balanced-match": { + "version": "4.0.4", + "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz", + "integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==", + "license": "MIT", + "engines": { + "node": "18 || 20 || >=22" + } + }, + "node_modules/brace-expansion": { + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", + "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", + "license": "MIT", + "dependencies": { + "balanced-match": "^4.0.2" + }, + "engines": { + "node": "18 || 20 || >=22" + } + }, + "node_modules/braces": { + "version": "3.0.3", + "resolved": "https://registry.npmjs.org/braces/-/braces-3.0.3.tgz", + "integrity": "sha512-yQbXgO/OSZVD2IsiLlro+7Hf6Q18EJrKSEsdoMzKePKXct3gvD8oLcOQdIzGupr5Fj+EDe8gO/lxc1BzfMpxvA==", + "dev": true, + "license": "MIT", + "dependencies": { + "fill-range": "^7.1.1" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/chalk": { + "version": "5.6.2", + "resolved": "https://registry.npmjs.org/chalk/-/chalk-5.6.2.tgz", + "integrity": "sha512-7NzBL0rN6fMUW+f7A6Io4h40qQlG+xGmtMxfbnH/K7TAtt8JQWVQK+6g0UXKMeVJoyV5EkkNsErQ8pVD3bLHbA==", + "license": "MIT", + "engines": { + "node": "^12.17.0 || ^14.13 || >=16.0.0" + }, + "funding": { + "url": "https://github.com/chalk/chalk?sponsor=1" + } + }, + "node_modules/cross-spawn": { + "version": "7.0.6", + "resolved": "https://registry.npmjs.org/cross-spawn/-/cross-spawn-7.0.6.tgz", + "integrity": "sha512-uV2QOWP2nWzsy2aMp8aRibhi9dlzF5Hgh5SHaB9OiTGEyDTiJJyx0uy51QXdyWbtAHNua4XJzUKca3OzKUd3vA==", + "license": "MIT", + "dependencies": { + "path-key": "^3.1.0", + "shebang-command": "^2.0.0", + "which": "^2.0.1" + }, + "engines": { + "node": ">= 8" + } + }, + "node_modules/diff": { + "version": "8.0.4", + "resolved": "https://registry.npmjs.org/diff/-/diff-8.0.4.tgz", + "integrity": "sha512-DPi0FmjiSU5EvQV0++GFDOJ9ASQUVFh5kD+OzOnYdi7n3Wpm9hWWGfB/O2blfHcMVTL5WkQXSnRiK9makhrcnw==", + "license": "BSD-3-Clause", + "engines": { + "node": ">=0.3.1" + } + }, + "node_modules/end-of-stream": { + "version": "1.4.5", + "resolved": "https://registry.npmjs.org/end-of-stream/-/end-of-stream-1.4.5.tgz", + "integrity": "sha512-ooEGc6HP26xXq/N+GCGOT0JKCLDGrq2bQUZrQ7gyrJiZANJ/8YDTxTpQBXGMn+WbIQXNVpyWymm7KYVICQnyOg==", + "dev": true, + "license": "MIT", + "dependencies": { + "once": "^1.4.0" + } + }, + "node_modules/es-errors": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/es-errors/-/es-errors-1.3.0.tgz", + "integrity": "sha512-Zf5H2Kxt2xjTvbJvP2ZWLEICxA6j+hAmMzIlypy4xcBg1vKVnx89Wy0GbS+kf5cwCVFFzdCFh2XSCFNULS6csw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/execa": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/execa/-/execa-1.0.0.tgz", + "integrity": "sha512-adbxcyWV46qiHyvSp50TKt05tB4tK3HcmF7/nxfAdhnox83seTDbwnaqKO4sXRy7roHAIFqJP/Rw/AuEbX61LA==", + "dev": true, + "license": "MIT", + "dependencies": { + "cross-spawn": "^6.0.0", + "get-stream": "^4.0.0", + "is-stream": "^1.1.0", + "npm-run-path": "^2.0.0", + "p-finally": "^1.0.0", + "signal-exit": "^3.0.0", + "strip-eof": "^1.0.0" + }, + "engines": { + "node": ">=6" + } + }, + "node_modules/execa/node_modules/cross-spawn": { + "version": "6.0.6", + "resolved": "https://registry.npmjs.org/cross-spawn/-/cross-spawn-6.0.6.tgz", + "integrity": "sha512-VqCUuhcd1iB+dsv8gxPttb5iZh/D0iubSP21g36KXdEuf6I5JiioesUVjpCdHV9MZRUfVFlvwtIUyPfxo5trtw==", + "dev": true, + "license": "MIT", + "dependencies": { + "nice-try": "^1.0.4", + "path-key": "^2.0.1", + "semver": "^5.5.0", + "shebang-command": "^1.2.0", + "which": "^1.2.9" + }, + "engines": { + "node": ">=4.8" + } + }, + "node_modules/execa/node_modules/path-key": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/path-key/-/path-key-2.0.1.tgz", + "integrity": "sha512-fEHGKCSmUSDPv4uoj8AlD+joPlq3peND+HRYyxFz4KPw4z926S/b8rIuFs2FYJg3BwsxJf6A9/3eIdLaYC+9Dw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=4" + } + }, + "node_modules/execa/node_modules/semver": { + "version": "5.7.2", + "resolved": "https://registry.npmjs.org/semver/-/semver-5.7.2.tgz", + "integrity": "sha512-cBznnQ9KjJqU67B52RMC65CMarK2600WFnbkcaiwWq3xy/5haFJlshgnpjovMVJ+Hff49d8GEn0b87C5pDQ10g==", + "dev": true, + "license": "ISC", + "bin": { + "semver": "bin/semver" + } + }, + "node_modules/execa/node_modules/shebang-command": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/shebang-command/-/shebang-command-1.2.0.tgz", + "integrity": "sha512-EV3L1+UQWGor21OmnvojK36mhg+TyIKDh3iFBKBohr5xeXIhNBcx8oWdgkTEEQ+BEFFYdLRuqMfd5L84N1V5Vg==", + "dev": true, + "license": "MIT", + "dependencies": { + "shebang-regex": "^1.0.0" + }, + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/execa/node_modules/shebang-regex": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/shebang-regex/-/shebang-regex-1.0.0.tgz", + "integrity": "sha512-wpoSFAxys6b2a2wHZ1XpDSgD7N9iVjg29Ph9uV/uaP9Ex/KXlkTZTeddxDPSYQpgvzKLGJke2UU0AzoGCjNIvQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/execa/node_modules/which": { + "version": "1.3.1", + "resolved": "https://registry.npmjs.org/which/-/which-1.3.1.tgz", + "integrity": "sha512-HxJdYWq1MTIQbJ3nw0cqssHoTNU267KlrDuGZ1WYlxDStUtKUhOaJmh112/TZmHxxUfuJqPXSOm7tDyas0OSIQ==", + "dev": true, + "license": "ISC", + "dependencies": { + "isexe": "^2.0.0" + }, + "bin": { + "which": "bin/which" + } + }, + "node_modules/fast-glob": { + "version": "3.3.3", + "resolved": "https://registry.npmjs.org/fast-glob/-/fast-glob-3.3.3.tgz", + "integrity": "sha512-7MptL8U0cqcFdzIzwOTHoilX9x5BrNqye7Z/LuC7kCMRio1EMSyqRK3BEAUD7sXRq4iT4AzTVuZdhgQ2TCvYLg==", + "dev": true, + "license": "MIT", + "dependencies": { + "@nodelib/fs.stat": "^2.0.2", + "@nodelib/fs.walk": "^1.2.3", + "glob-parent": "^5.1.2", + "merge2": "^1.3.0", + "micromatch": "^4.0.8" + }, + "engines": { + "node": ">=8.6.0" + } + }, + "node_modules/fastq": { + "version": "1.20.1", + "resolved": "https://registry.npmjs.org/fastq/-/fastq-1.20.1.tgz", + "integrity": "sha512-GGToxJ/w1x32s/D2EKND7kTil4n8OVk/9mycTc4VDza13lOvpUZTGX3mFSCtV9ksdGBVzvsyAVLM6mHFThxXxw==", + "dev": true, + "license": "ISC", + "dependencies": { + "reusify": "^1.0.4" + } + }, + "node_modules/fill-range": { + "version": "7.1.1", + "resolved": "https://registry.npmjs.org/fill-range/-/fill-range-7.1.1.tgz", + "integrity": "sha512-YsGpe3WHLK8ZYi4tWDg2Jy3ebRz2rXowDxnld4bkQB00cc/1Zw9AWnC0i9ztDJitivtQvaI9KaLyKrc+hBW0yg==", + "dev": true, + "license": "MIT", + "dependencies": { + "to-regex-range": "^5.0.1" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/function-bind": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz", + "integrity": "sha512-7XHNxH7qX9xG5mIwxkhumTox/MIRNcOgDrxWsMt2pAr23WHp6MrRlN7FBSFpCpr+oVO0F744iUgR82nJMfG2SA==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/get-east-asian-width": { + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/get-east-asian-width/-/get-east-asian-width-1.6.0.tgz", + "integrity": "sha512-QRbvDIbx6YklUe6RxeTeleMR0yv3cYH6PsPZHcnVn7xv7zO1BHN8r0XETu8n6Ye3Q+ahtSarc3WgtNWmehIBfA==", + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, + "node_modules/get-stream": { + "version": "4.1.0", + "resolved": "https://registry.npmjs.org/get-stream/-/get-stream-4.1.0.tgz", + "integrity": "sha512-GMat4EJ5161kIy2HevLlr4luNjBgvmj413KaQA7jt4V8B4RDsfpHk7WQ9GVqfYyyx8OS/L66Kox+rJRNklLK7w==", + "dev": true, + "license": "MIT", + "dependencies": { + "pump": "^3.0.0" + }, + "engines": { + "node": ">=6" + } + }, + "node_modules/glob": { + "version": "13.0.6", + "resolved": "https://registry.npmjs.org/glob/-/glob-13.0.6.tgz", + "integrity": "sha512-Wjlyrolmm8uDpm/ogGyXZXb1Z+Ca2B8NbJwqBVg0axK9GbBeoS7yGV6vjXnYdGm6X53iehEuxxbyiKp8QmN4Vw==", + "license": "BlueOak-1.0.0", + "dependencies": { + "minimatch": "^10.2.2", + "minipass": "^7.1.3", + "path-scurry": "^2.0.2" + }, + "engines": { + "node": "18 || 20 || >=22" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" + } + }, + "node_modules/glob-parent": { + "version": "5.1.2", + "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", + "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", + "dev": true, + "license": "ISC", + "dependencies": { + "is-glob": "^4.0.1" + }, + "engines": { + "node": ">= 6" + } + }, + "node_modules/graceful-fs": { + "version": "4.2.11", + "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz", + "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==", + "license": "ISC" + }, + "node_modules/hasown": { + "version": "2.0.3", + "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.3.tgz", + "integrity": "sha512-ej4AhfhfL2Q2zpMmLo7U1Uv9+PyhIZpgQLGT1F9miIGmiCJIoCgSmczFdrc97mWT4kVY72KA+WnnhJ5pghSvSg==", + "dev": true, + "license": "MIT", + "dependencies": { + "function-bind": "^1.1.2" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/highlight.js": { + "version": "10.7.3", + "resolved": "https://registry.npmjs.org/highlight.js/-/highlight.js-10.7.3.tgz", + "integrity": "sha512-tzcUFauisWKNHaRkN4Wjl/ZA07gENAjFl3J/c480dprkGTg5EQstgaNFqBfUqCq54kZRIEcreTsAgF/m2quD7A==", + "license": "BSD-3-Clause", + "engines": { + "node": "*" + } + }, + "node_modules/hosted-git-info": { + "version": "9.0.3", + "resolved": "https://registry.npmjs.org/hosted-git-info/-/hosted-git-info-9.0.3.tgz", + "integrity": "sha512-Hc+ghLoSt6QaYZUv0WBiIvmMDZuZZ7oaDvdH8MbfOO4lOsxdXLEvuC6ePoGs9H1X9oCLyq6+NVN0MKqD+ydxyg==", + "license": "ISC", + "dependencies": { + "lru-cache": "^11.1.0" + }, + "engines": { + "node": "^20.17.0 || >=22.9.0" + } + }, + "node_modules/ignore": { + "version": "7.0.5", + "resolved": "https://registry.npmjs.org/ignore/-/ignore-7.0.5.tgz", + "integrity": "sha512-Hs59xBNfUIunMFgWAbGX5cq6893IbWg4KnrjbYwX3tx0ztorVgTDA6B2sxf8ejHJ4wz8BqGUMYlnzNBer5NvGg==", + "license": "MIT", + "engines": { + "node": ">= 4" + } + }, + "node_modules/interpret": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/interpret/-/interpret-1.4.0.tgz", + "integrity": "sha512-agE4QfB2Lkp9uICn7BAqoscw4SZP9kTE2hxiFI3jBPmXJfdqiahTbUuKGsMoN2GtqL9AxhYioAcVvgsb1HvRbA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.10" + } + }, + "node_modules/is-core-module": { + "version": "2.16.2", + "resolved": "https://registry.npmjs.org/is-core-module/-/is-core-module-2.16.2.tgz", + "integrity": "sha512-evOr8xfXKxE6qSR0hSXL2r3sd7ALj8+7jQEUvPYcm5sgZFdJ+AYzT6yNmJenvIYQBgIGwfwz08sL8zoL7yq2BA==", + "dev": true, + "license": "MIT", + "dependencies": { + "hasown": "^2.0.3" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/is-extglob": { + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/is-extglob/-/is-extglob-2.1.1.tgz", + "integrity": "sha512-SbKbANkN603Vi4jEZv49LeVJMn4yGwsbzZworEoyEiutsN3nJYdbO36zfhGJ6QEDpOZIFkDtnq5JRxmvl3jsoQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/is-glob": { + "version": "4.0.3", + "resolved": "https://registry.npmjs.org/is-glob/-/is-glob-4.0.3.tgz", + "integrity": "sha512-xelSayHH36ZgE7ZWhli7pW34hNbNl8Ojv5KVmkJD4hBdD3th8Tfk9vYasLM+mXWOZhFkgZfxhLSnrwRr4elSSg==", + "dev": true, + "license": "MIT", + "dependencies": { + "is-extglob": "^2.1.1" + }, + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/is-number": { + "version": "7.0.0", + "resolved": "https://registry.npmjs.org/is-number/-/is-number-7.0.0.tgz", + "integrity": "sha512-41Cifkg6e8TylSpdtTpeLVMqvSBEVzTttHvERD741+pnZ8ANv0004MRL43QKPDlK9cGvNp6NZWZUBlbGXYxxng==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.12.0" + } + }, + "node_modules/is-stream": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/is-stream/-/is-stream-1.1.0.tgz", + "integrity": "sha512-uQPm8kcs47jx38atAcWTVxyltQYoPT68y9aWYdV6yWXSyW8mzSat0TL6CiWdZeCdF3KrAvpVtnHbTv4RN+rqdQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/isexe": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/isexe/-/isexe-2.0.0.tgz", + "integrity": "sha512-RHxMLp9lnKHGHRng9QFhRCMbYAcVpn69smSGcq3f36xjgVVWThj4qqLbTLlq7Ssj8B+fIQ1EuCEGI2lKsyQeIw==", + "license": "ISC" + }, + "node_modules/jiti": { + "version": "2.7.0", + "resolved": "https://registry.npmjs.org/jiti/-/jiti-2.7.0.tgz", + "integrity": "sha512-AC/7JofJvZGrrneWNaEnJeOLUx+JlGt7tNa0wZiRPT4MY1wmfKjt2+6O2p2uz2+skll8OZZmJMNqeke7kKbNgQ==", + "license": "MIT", + "bin": { + "jiti": "lib/jiti-cli.mjs" + } + }, + "node_modules/lru-cache": { + "version": "11.4.0", + "resolved": "https://registry.npmjs.org/lru-cache/-/lru-cache-11.4.0.tgz", + "integrity": "sha512-W+R+kFL4HgVxONq2bhXPi3bGpzGe/yEhVOp233qw9wCRtgncJ15P3bC+e4zZMu4Cq7d+WAJjXGW0uUkifhcatA==", + "license": "BlueOak-1.0.0", + "engines": { + "node": "20 || >=22" + } + }, + "node_modules/marked": { + "version": "18.0.5", + "resolved": "https://registry.npmjs.org/marked/-/marked-18.0.5.tgz", + "integrity": "sha512-S6GcvALHg6K4ohtu4E7x0a1AqhAjp6cV8KhLSyN9qVapnzJkusVBxZRcIU9AeYsbe6P1hKDusSbEOzGyyuce6w==", + "license": "MIT", + "bin": { + "marked": "bin/marked.js" + }, + "engines": { + "node": ">= 20" + } + }, + "node_modules/merge2": { + "version": "1.4.1", + "resolved": "https://registry.npmjs.org/merge2/-/merge2-1.4.1.tgz", + "integrity": "sha512-8q7VEgMJW4J8tcfVPy8g09NcQwZdbwFEqhe/WZkoIzjn/3TGDwtOCYtXGxA3O8tPzpczCCDgv+P2P5y00ZJOOg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 8" + } + }, + "node_modules/micromatch": { + "version": "4.0.8", + "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.8.tgz", + "integrity": "sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==", + "dev": true, + "license": "MIT", + "dependencies": { + "braces": "^3.0.3", + "picomatch": "^2.3.1" + }, + "engines": { + "node": ">=8.6" + } + }, + "node_modules/minimatch": { + "version": "10.2.5", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.5.tgz", + "integrity": "sha512-MULkVLfKGYDFYejP07QOurDLLQpcjk7Fw+7jXS2R2czRQzR56yHRveU5NDJEOviH+hETZKSkIk5c+T23GjFUMg==", + "license": "BlueOak-1.0.0", + "dependencies": { + "brace-expansion": "^5.0.5" + }, + "engines": { + "node": "18 || 20 || >=22" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" + } + }, + "node_modules/minimist": { + "version": "1.2.8", + "resolved": "https://registry.npmjs.org/minimist/-/minimist-1.2.8.tgz", + "integrity": "sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/minipass": { + "version": "7.1.3", + "resolved": "https://registry.npmjs.org/minipass/-/minipass-7.1.3.tgz", + "integrity": "sha512-tEBHqDnIoM/1rXME1zgka9g6Q2lcoCkxHLuc7ODJ5BxbP5d4c2Z5cGgtXAku59200Cx7diuHTOYfSBD8n6mm8A==", + "license": "BlueOak-1.0.0", + "engines": { + "node": ">=16 || 14 >=14.17" + } + }, + "node_modules/nice-try": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/nice-try/-/nice-try-1.0.5.tgz", + "integrity": "sha512-1nh45deeb5olNY7eX82BkPO7SSxR5SSYJiPTrTdFUVYwAl8CKMA5N9PjTYkHiRjisVcxcQ1HXdLhx2qxxJzLNQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/npm-run-path": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/npm-run-path/-/npm-run-path-2.0.2.tgz", + "integrity": "sha512-lJxZYlT4DW/bRUtFh1MQIWqmLwQfAxnqWG4HhEdjMlkrJYnJn0Jrr2u3mgxqaWsdiBc76TYkTG/mhrnYTuzfHw==", + "dev": true, + "license": "MIT", + "dependencies": { + "path-key": "^2.0.0" + }, + "engines": { + "node": ">=4" + } + }, + "node_modules/npm-run-path/node_modules/path-key": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/path-key/-/path-key-2.0.1.tgz", + "integrity": "sha512-fEHGKCSmUSDPv4uoj8AlD+joPlq3peND+HRYyxFz4KPw4z926S/b8rIuFs2FYJg3BwsxJf6A9/3eIdLaYC+9Dw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=4" + } + }, + "node_modules/once": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", + "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==", + "dev": true, + "license": "ISC", + "dependencies": { + "wrappy": "1" + } + }, + "node_modules/openai": { + "version": "6.26.0", + "resolved": "https://registry.npmjs.org/openai/-/openai-6.26.0.tgz", + "integrity": "sha512-zd23dbWTjiJ6sSAX6s0HrCZi41JwTA1bQVs0wLQPZ2/5o2gxOJA5wh7yOAUgwYybfhDXyhwlpeQf7Mlgx8EOCA==", + "license": "Apache-2.0", + "bin": { + "openai": "bin/cli" + }, + "peerDependencies": { + "ws": "^8.18.0", + "zod": "^3.25 || ^4.0" + }, + "peerDependenciesMeta": { + "ws": { + "optional": true + }, + "zod": { + "optional": true + } + } + }, + "node_modules/p-finally": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/p-finally/-/p-finally-1.0.0.tgz", + "integrity": "sha512-LICb2p9CB7FS+0eR1oqWnHhp0FljGLZCWBE9aix0Uye9W8LTQPwMTYVGWQWIw9RdQiDg4+epXQODwIYJtSJaow==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=4" + } + }, + "node_modules/partial-json": { + "version": "0.1.7", + "resolved": "https://registry.npmjs.org/partial-json/-/partial-json-0.1.7.tgz", + "integrity": "sha512-Njv/59hHaokb/hRUjce3Hdv12wd60MtM9Z5Olmn+nehe0QDAsRtRbJPvJ0Z91TusF0SuZRIvnM+S4l6EIP8leA==", + "license": "MIT" + }, + "node_modules/path-key": { + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/path-key/-/path-key-3.1.1.tgz", + "integrity": "sha512-ojmeN0qd+y0jszEtoY48r0Peq5dwMEkIlCOu6Q5f41lfkswXuKtYrhgoTpLnyIcHm24Uhqx+5Tqm2InSwLhE6Q==", + "license": "MIT", + "engines": { + "node": ">=8" + } + }, + "node_modules/path-parse": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/path-parse/-/path-parse-1.0.7.tgz", + "integrity": "sha512-LDJzPVEEEPR+y48z93A0Ed0yXb8pAByGWo/k5YYdYgpY2/2EsOsksJrq7lOHxryrVOn1ejG6oAp8ahvOIQD8sw==", + "dev": true, + "license": "MIT" + }, + "node_modules/path-scurry": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/path-scurry/-/path-scurry-2.0.2.tgz", + "integrity": "sha512-3O/iVVsJAPsOnpwWIeD+d6z/7PmqApyQePUtCndjatj/9I5LylHvt5qluFaBT3I5h3r1ejfR056c+FCv+NnNXg==", + "license": "BlueOak-1.0.0", + "dependencies": { + "lru-cache": "^11.0.0", + "minipass": "^7.1.2" + }, + "engines": { + "node": "18 || 20 || >=22" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" + } + }, + "node_modules/picomatch": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.2.tgz", + "integrity": "sha512-V7+vQEJ06Z+c5tSye8S+nHUfI51xoXIXjHQ99cQtKUkQqqO1kO/KCJUfZXuB47h/YBlDhah2H3hdUGXn8ie0oA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=8.6" + }, + "funding": { + "url": "https://github.com/sponsors/jonschlinkert" + } + }, + "node_modules/proper-lockfile": { + "version": "4.1.2", + "resolved": "https://registry.npmjs.org/proper-lockfile/-/proper-lockfile-4.1.2.tgz", + "integrity": "sha512-TjNPblN4BwAWMXU8s9AEz4JmQxnD1NNL7bNOY/AKUzyamc379FWASUhc/K1pL2noVb+XmZKLL68cjzLsiOAMaA==", + "license": "MIT", + "dependencies": { + "graceful-fs": "^4.2.4", + "retry": "^0.12.0", + "signal-exit": "^3.0.2" + } + }, + "node_modules/proper-lockfile/node_modules/retry": { + "version": "0.12.0", + "resolved": "https://registry.npmjs.org/retry/-/retry-0.12.0.tgz", + "integrity": "sha512-9LkiTwjUh6rT555DtE9rTX+BKByPfrMzEAtnlEtdEwr3Nkffwiihqe2bWADg+OQRjt9gl6ICdmB/ZFDCGAtSow==", + "license": "MIT", + "engines": { + "node": ">= 4" + } + }, + "node_modules/pump": { + "version": "3.0.4", + "resolved": "https://registry.npmjs.org/pump/-/pump-3.0.4.tgz", + "integrity": "sha512-VS7sjc6KR7e1ukRFhQSY5LM2uBWAUPiOPa/A3mkKmiMwSmRFUITt0xuj+/lesgnCv+dPIEYlkzrcyXgquIHMcA==", + "dev": true, + "license": "MIT", + "dependencies": { + "end-of-stream": "^1.1.0", + "once": "^1.3.1" + } + }, + "node_modules/queue-microtask": { + "version": "1.2.3", + "resolved": "https://registry.npmjs.org/queue-microtask/-/queue-microtask-1.2.3.tgz", + "integrity": "sha512-NuaNSa6flKT5JaSYQzJok04JzTL1CA6aGhv5rfLW3PgqA+M2ChpZQnAC8h8i4ZFkBS8X5RqkDBHA7r4hej3K9A==", + "dev": true, + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "MIT" + }, + "node_modules/rechoir": { + "version": "0.6.2", + "resolved": "https://registry.npmjs.org/rechoir/-/rechoir-0.6.2.tgz", + "integrity": "sha512-HFM8rkZ+i3zrV+4LQjwQ0W+ez98pApMGM3HUrN04j3CqzPOzl9nmP15Y8YXNm8QHGv/eacOVEjqhmWpkRV0NAw==", + "dev": true, + "dependencies": { + "resolve": "^1.1.6" + }, + "engines": { + "node": ">= 0.10" + } + }, + "node_modules/resolve": { + "version": "1.22.12", + "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.22.12.tgz", + "integrity": "sha512-TyeJ1zif53BPfHootBGwPRYT1RUt6oGWsaQr8UyZW/eAm9bKoijtvruSDEmZHm92CwS9nj7/fWttqPCgzep8CA==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-errors": "^1.3.0", + "is-core-module": "^2.16.1", + "path-parse": "^1.0.7", + "supports-preserve-symlinks-flag": "^1.0.0" + }, + "bin": { + "resolve": "bin/resolve" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/reusify": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/reusify/-/reusify-1.1.0.tgz", + "integrity": "sha512-g6QUff04oZpHs0eG5p83rFLhHeV00ug/Yf9nZM6fLeUrPguBTkTQOdpAWWspMh55TZfVQDPaN3NQJfbVRAxdIw==", + "dev": true, + "license": "MIT", + "engines": { + "iojs": ">=1.0.0", + "node": ">=0.10.0" + } + }, + "node_modules/run-parallel": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/run-parallel/-/run-parallel-1.2.0.tgz", + "integrity": "sha512-5l4VyZR86LZ/lDxZTR6jqL8AFE2S0IFLMP26AbjsLVADxHdhB/c0GUsH+y39UfCi3dzz8OlQuPmnaJOMoDHQBA==", + "dev": true, + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/feross" + }, + { + "type": "patreon", + "url": "https://www.patreon.com/feross" + }, + { + "type": "consulting", + "url": "https://feross.org/support" + } + ], + "license": "MIT", + "dependencies": { + "queue-microtask": "^1.2.2" + } + }, + "node_modules/semver": { + "version": "7.8.0", + "resolved": "https://registry.npmjs.org/semver/-/semver-7.8.0.tgz", + "integrity": "sha512-AcM7dV/5ul4EekoQ29Agm5vri8JNqRyj39o0qpX6vDF2GZrtutZl5RwgD1XnZjiTAfncsJhMI48QQH3sN87YNA==", + "license": "ISC", + "bin": { + "semver": "bin/semver.js" + }, + "engines": { + "node": ">=10" + } + }, + "node_modules/shebang-command": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/shebang-command/-/shebang-command-2.0.0.tgz", + "integrity": "sha512-kHxr2zZpYtdmrN1qDjrrX/Z1rR1kG8Dx+gkpK1G4eXmvXswmcE1hTWBWYUzlraYw1/yZp6YuDY77YtvbN0dmDA==", + "license": "MIT", + "dependencies": { + "shebang-regex": "^3.0.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/shebang-regex": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/shebang-regex/-/shebang-regex-3.0.0.tgz", + "integrity": "sha512-7++dFhtcx3353uBaq8DDR4NuxBetBzC7ZQOhmTQInHEd6bSrXdiEyzCvG07Z44UYdLShWUyXt5M/yhz8ekcb1A==", + "license": "MIT", + "engines": { + "node": ">=8" + } + }, + "node_modules/shelljs": { + "version": "0.9.2", + "resolved": "https://registry.npmjs.org/shelljs/-/shelljs-0.9.2.tgz", + "integrity": "sha512-S3I64fEiKgTZzKCC46zT/Ib9meqofLrQVbpSswtjFfAVDW+AZ54WTnAM/3/yENoxz/V1Cy6u3kiiEbQ4DNphvw==", + "dev": true, + "license": "BSD-3-Clause", + "dependencies": { + "execa": "^1.0.0", + "fast-glob": "^3.3.2", + "interpret": "^1.0.0", + "rechoir": "^0.6.2" + }, + "bin": { + "shjs": "bin/shjs" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/shx": { + "version": "0.4.0", + "resolved": "https://registry.npmjs.org/shx/-/shx-0.4.0.tgz", + "integrity": "sha512-Z0KixSIlGPpijKgcH6oCMCbltPImvaKy0sGH8AkLRXw1KyzpKtaCTizP2xen+hNDqVF4xxgvA0KXSb9o4Q6hnA==", + "dev": true, + "license": "MIT", + "dependencies": { + "minimist": "^1.2.8", + "shelljs": "^0.9.2" + }, + "bin": { + "shx": "lib/cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/signal-exit": { + "version": "3.0.7", + "resolved": "https://registry.npmjs.org/signal-exit/-/signal-exit-3.0.7.tgz", + "integrity": "sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ==", + "license": "ISC" + }, + "node_modules/strip-eof": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/strip-eof/-/strip-eof-1.0.0.tgz", + "integrity": "sha512-7FCwGGmx8mD5xQd3RPUvnSpUXHM3BWuzjtpD4TXsfcZ9EL4azvVVUscFYwD9nx8Kh+uCBC00XBtAykoMHwTh8Q==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/supports-preserve-symlinks-flag": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/supports-preserve-symlinks-flag/-/supports-preserve-symlinks-flag-1.0.0.tgz", + "integrity": "sha512-ot0WnXS9fgdkgIcePe6RHNk1WA8+muPa6cSjeR3V8K27q9BB1rTE3R1p7Hv0z1ZyAc8s6Vvv8DIyWf681MAt0w==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/to-regex-range": { + "version": "5.0.1", + "resolved": "https://registry.npmjs.org/to-regex-range/-/to-regex-range-5.0.1.tgz", + "integrity": "sha512-65P7iz6X5yEr1cwcgvQxbbIw7Uk3gOy5dIdtZ4rDveLqhrdJP+Li/Hx6tyK0NEb+2GCyneCMJiGqrADCSNk8sQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "is-number": "^7.0.0" + }, + "engines": { + "node": ">=8.0" + } + }, + "node_modules/typebox": { + "version": "1.1.38", + "resolved": "https://registry.npmjs.org/typebox/-/typebox-1.1.38.tgz", + "integrity": "sha512-pZ0aQPmMmXoUvSbeuWf/Hzsc+avNw/Zd6VeE8CFgkVGWyuHPJvqeJJDeJqLve+K70LvjYIoleGcoJHPT17cWoA==", + "license": "MIT" + }, + "node_modules/undici": { + "version": "8.5.0", + "resolved": "https://registry.npmjs.org/undici/-/undici-8.5.0.tgz", + "integrity": "sha512-xamtWoB1EshgjpmlXd7GGm2VfdDtw1+rD8uhry8pSNW3If6S8E0m2T2+orSKeZXEn/aPJMviCpDBA65WJt8zhg==", + "license": "MIT", + "engines": { + "node": ">=22.19.0" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/which": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz", + "integrity": "sha512-BLI3Tl1TW3Pvl70l3yq3Y64i+awpwXqsGBYWkkqMtnbXgrMD+yj7rhW0kuEDxzJaYXGjEW5ogapKNMEKNMjibA==", + "license": "ISC", + "dependencies": { + "isexe": "^2.0.0" + }, + "bin": { + "node-which": "bin/node-which" + }, + "engines": { + "node": ">= 8" + } + }, + "node_modules/wrappy": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz", + "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", + "dev": true, + "license": "ISC" + }, + "node_modules/yaml": { + "version": "2.9.0", + "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.9.0.tgz", + "integrity": "sha512-2AvhNX3mb8zd6Zy7INTtSpl1F15HW6Wnqj0srWlkKLcpYl/gMIMJiyuGq2KeI2YFxUPjdlB+3Lc10seMLtL4cA==", + "license": "ISC", + "bin": { + "yaml": "bin.mjs" + }, + "engines": { + "node": ">= 14.6" + }, + "funding": { + "url": "https://github.com/sponsors/eemeli" + } + }, + "packages/agent": { + "name": "@earendil-works/pi-agent-core", + "version": "0.80.2", + "license": "MIT", + "dependencies": { + "@earendil-works/pi-ai": "^0.80.2", + "ignore": "7.0.5", + "typebox": "1.1.38", + "yaml": "2.9.0" + }, + "devDependencies": { + "@types/node": "24.12.4" + }, + "engines": { + "node": ">=22.19.0" + } + }, + "packages/agent/node_modules/@types/node": { + "version": "24.12.4", + "resolved": "https://registry.npmjs.org/@types/node/-/node-24.12.4.tgz", + "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~7.16.0" + } + }, + "packages/agent/node_modules/undici-types": { + "version": "7.16.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-7.16.0.tgz", + "integrity": "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw==", + "dev": true, + "license": "MIT" + }, + "packages/ai": { + "name": "@earendil-works/pi-ai", + "version": "0.80.2", + "license": "MIT", + "dependencies": { + "openai": "6.26.0", + "partial-json": "0.1.7", + "typebox": "1.1.38" + }, + "devDependencies": { + "@types/node": "24.12.4" + }, + "engines": { + "node": ">=22.19.0" + } + }, + "packages/ai/node_modules/@types/node": { + "version": "24.12.4", + "resolved": "https://registry.npmjs.org/@types/node/-/node-24.12.4.tgz", + "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~7.16.0" + } + }, + "packages/ai/node_modules/undici-types": { + "version": "7.16.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-7.16.0.tgz", + "integrity": "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw==", + "dev": true, + "license": "MIT" + }, + "packages/coding-agent": { + "name": "@earendil-works/pi-coding-agent", + "version": "0.80.2", + "license": "MIT", + "dependencies": { + "@earendil-works/pi-agent-core": "^0.80.2", + "@earendil-works/pi-ai": "^0.80.2", + "@earendil-works/pi-tui": "^0.80.2", + "@silvia-odwyer/photon-node": "0.3.4", + "chalk": "5.6.2", + "cross-spawn": "7.0.6", + "diff": "8.0.4", + "glob": "13.0.6", + "highlight.js": "10.7.3", + "hosted-git-info": "9.0.3", + "ignore": "7.0.5", + "jiti": "2.7.0", + "minimatch": "10.2.5", + "proper-lockfile": "4.1.2", + "semver": "7.8.0", + "typebox": "1.1.38", + "undici": "8.5.0", + "yaml": "2.9.0" + }, + "bin": { + "cactus": "dist/cli.js" + }, + "devDependencies": { + "@types/cross-spawn": "6.0.6", + "@types/diff": "7.0.2", + "@types/hosted-git-info": "3.0.5", + "@types/ms": "2.1.0", + "@types/node": "24.12.4", + "@types/proper-lockfile": "4.1.4", + "@types/semver": "7.7.1", + "shx": "0.4.0" + }, + "engines": { + "node": ">=22.19.0" + }, + "optionalDependencies": { + "@mariozechner/clipboard": "0.3.9" + } + }, + "packages/coding-agent/examples/extensions/custom-provider-anthropic": { + "name": "pi-extension-custom-provider-anthropic", + "version": "0.80.2", + "extraneous": true, + "dependencies": { + "@anthropic-ai/sdk": "0.52.0" + } + }, + "packages/coding-agent/examples/extensions/custom-provider-gitlab-duo": { + "name": "pi-extension-custom-provider-gitlab-duo", + "version": "0.80.2", + "extraneous": true + }, + "packages/coding-agent/examples/extensions/gondolin": { + "name": "pi-extension-gondolin", + "version": "0.80.2", + "extraneous": true, + "dependencies": { + "@earendil-works/gondolin": "0.12.0" + } + }, + "packages/coding-agent/examples/extensions/sandbox": { + "name": "pi-extension-sandbox", + "version": "1.10.2", + "extraneous": true, + "dependencies": { + "@anthropic-ai/sandbox-runtime": "0.0.26" + } + }, + "packages/coding-agent/examples/extensions/with-deps": { + "name": "pi-extension-with-deps", + "version": "0.80.2", + "extraneous": true, + "dependencies": { + "ms": "2.1.3" + }, + "devDependencies": { + "@types/ms": "2.1.0" + } + }, + "packages/coding-agent/node_modules/@types/node": { + "version": "24.12.4", + "resolved": "https://registry.npmjs.org/@types/node/-/node-24.12.4.tgz", + "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~7.16.0" + } + }, + "packages/coding-agent/node_modules/undici-types": { + "version": "7.16.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-7.16.0.tgz", + "integrity": "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw==", + "dev": true, + "license": "MIT" + }, + "packages/tui": { + "name": "@earendil-works/pi-tui", + "version": "0.80.2", + "license": "MIT", + "dependencies": { + "get-east-asian-width": "1.6.0", + "marked": "18.0.5" + }, + "devDependencies": { + "@types/node": "24.12.4" + }, + "engines": { + "node": ">=22.19.0" + } + }, + "packages/tui/node_modules/@types/node": { + "version": "24.12.4", + "resolved": "https://registry.npmjs.org/@types/node/-/node-24.12.4.tgz", + "integrity": "sha512-GUUEShf+PBCGW2KaXwcIt3Yk+e3pkKwWKb9GSyM9WQVE+ep2jzmHdGsHzu4wgcZy5fN9FBdVzjpBQsYlpfpgLA==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~7.16.0" + } + }, + "packages/tui/node_modules/undici-types": { + "version": "7.16.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-7.16.0.tgz", + "integrity": "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/cactus-code/package.json b/cactus-code/package.json new file mode 100644 index 000000000..0b70eeaab --- /dev/null +++ b/cactus-code/package.json @@ -0,0 +1,24 @@ +{ + "name": "pi-monorepo", + "private": true, + "type": "module", + "workspaces": [ + "packages/*" + ], + "scripts": { + "clean": "npm run clean --workspaces", + "build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../coding-agent && npm run build" + }, + "devDependencies": { + "@types/node": "22.19.19", + "@typescript/native-preview": "7.0.0-dev.20260120.1", + "shx": "0.4.0" + }, + "engines": { + "node": ">=22.19.0" + }, + "version": "0.0.3", + "overrides": { + "rimraf": "6.1.2" + } +} diff --git a/cactus-code/packages/agent/package.json b/cactus-code/packages/agent/package.json new file mode 100644 index 000000000..82003a79d --- /dev/null +++ b/cactus-code/packages/agent/package.json @@ -0,0 +1,54 @@ +{ + "name": "@earendil-works/pi-agent-core", + "version": "0.80.2", + "description": "General-purpose agent with transport abstraction, state management, and attachment support", + "type": "module", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + }, + "./node": { + "types": "./dist/node.d.ts", + "import": "./dist/node.js" + }, + "./package.json": "./package.json" + }, + "files": [ + "dist", + "README.md" + ], + "scripts": { + "clean": "shx rm -rf dist", + "build": "tsgo -p tsconfig.build.json", + "prepublishOnly": "npm run clean && npm run build" + }, + "dependencies": { + "@earendil-works/pi-ai": "^0.80.2", + "ignore": "7.0.5", + "typebox": "1.1.38", + "yaml": "2.9.0" + }, + "keywords": [ + "ai", + "agent", + "llm", + "transport", + "state-management" + ], + "author": "Mario Zechner", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/agent" + }, + "engines": { + "node": ">=22.19.0" + }, + "devDependencies": { + "@types/node": "24.12.4" + } +} diff --git a/cactus-code/packages/agent/src/agent-loop.ts b/cactus-code/packages/agent/src/agent-loop.ts new file mode 100644 index 000000000..d93458d07 --- /dev/null +++ b/cactus-code/packages/agent/src/agent-loop.ts @@ -0,0 +1,748 @@ +/** + * Agent loop that works with AgentMessage throughout. + * Transforms to Message[] only at the LLM call boundary. + */ + +import { + type AssistantMessage, + type Context, + EventStream, + streamSimple, + type ToolResultMessage, + validateToolArguments, +} from "@earendil-works/pi-ai/compat"; +import type { + AgentContext, + AgentEvent, + AgentLoopConfig, + AgentMessage, + AgentTool, + AgentToolCall, + AgentToolResult, + StreamFn, +} from "./types.ts"; + +export type AgentEventSink = (event: AgentEvent) => Promise | void; + +/** + * Start an agent loop with a new prompt message. + * The prompt is added to the context and events are emitted for it. + */ +export function agentLoop( + prompts: AgentMessage[], + context: AgentContext, + config: AgentLoopConfig, + signal?: AbortSignal, + streamFn?: StreamFn, +): EventStream { + const stream = createAgentStream(); + + void runAgentLoop( + prompts, + context, + config, + async (event) => { + stream.push(event); + }, + signal, + streamFn, + ).then((messages) => { + stream.end(messages); + }); + + return stream; +} + +/** + * Continue an agent loop from the current context without adding a new message. + * Used for retries - context already has user message or tool results. + * + * **Important:** The last message in context must convert to a `user` or `toolResult` message + * via `convertToLlm`. If it doesn't, the LLM provider will reject the request. + * This cannot be validated here since `convertToLlm` is only called once per turn. + */ +export function agentLoopContinue( + context: AgentContext, + config: AgentLoopConfig, + signal?: AbortSignal, + streamFn?: StreamFn, +): EventStream { + if (context.messages.length === 0) { + throw new Error("Cannot continue: no messages in context"); + } + + if (context.messages[context.messages.length - 1].role === "assistant") { + throw new Error("Cannot continue from message role: assistant"); + } + + const stream = createAgentStream(); + + void runAgentLoopContinue( + context, + config, + async (event) => { + stream.push(event); + }, + signal, + streamFn, + ).then((messages) => { + stream.end(messages); + }); + + return stream; +} + +export async function runAgentLoop( + prompts: AgentMessage[], + context: AgentContext, + config: AgentLoopConfig, + emit: AgentEventSink, + signal?: AbortSignal, + streamFn?: StreamFn, +): Promise { + const newMessages: AgentMessage[] = [...prompts]; + const currentContext: AgentContext = { + ...context, + messages: [...context.messages, ...prompts], + }; + + await emit({ type: "agent_start" }); + await emit({ type: "turn_start" }); + for (const prompt of prompts) { + await emit({ type: "message_start", message: prompt }); + await emit({ type: "message_end", message: prompt }); + } + + await runLoop(currentContext, newMessages, config, signal, emit, streamFn); + return newMessages; +} + +export async function runAgentLoopContinue( + context: AgentContext, + config: AgentLoopConfig, + emit: AgentEventSink, + signal?: AbortSignal, + streamFn?: StreamFn, +): Promise { + if (context.messages.length === 0) { + throw new Error("Cannot continue: no messages in context"); + } + + if (context.messages[context.messages.length - 1].role === "assistant") { + throw new Error("Cannot continue from message role: assistant"); + } + + const newMessages: AgentMessage[] = []; + const currentContext: AgentContext = { ...context }; + + await emit({ type: "agent_start" }); + await emit({ type: "turn_start" }); + + await runLoop(currentContext, newMessages, config, signal, emit, streamFn); + return newMessages; +} + +function createAgentStream(): EventStream { + return new EventStream( + (event: AgentEvent) => event.type === "agent_end", + (event: AgentEvent) => (event.type === "agent_end" ? event.messages : []), + ); +} + +/** + * Main loop logic shared by agentLoop and agentLoopContinue. + */ +async function runLoop( + initialContext: AgentContext, + newMessages: AgentMessage[], + initialConfig: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, + streamFn?: StreamFn, +): Promise { + let currentContext = initialContext; + let config = initialConfig; + let firstTurn = true; + // Check for steering messages at start (user may have typed while waiting) + let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || []; + + // Outer loop: continues when queued follow-up messages arrive after agent would stop + while (true) { + let hasMoreToolCalls = true; + + // Inner loop: process tool calls and steering messages + while (hasMoreToolCalls || pendingMessages.length > 0) { + if (!firstTurn) { + await emit({ type: "turn_start" }); + } else { + firstTurn = false; + } + + // Process pending messages (inject before next assistant response) + if (pendingMessages.length > 0) { + for (const message of pendingMessages) { + await emit({ type: "message_start", message }); + await emit({ type: "message_end", message }); + currentContext.messages.push(message); + newMessages.push(message); + } + pendingMessages = []; + } + + // Stream assistant response + const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFn); + newMessages.push(message); + + if (message.stopReason === "error" || message.stopReason === "aborted") { + await emit({ type: "turn_end", message, toolResults: [] }); + await emit({ type: "agent_end", messages: newMessages }); + return; + } + + // Check for tool calls + const toolCalls = message.content.filter((c) => c.type === "toolCall"); + + const toolResults: ToolResultMessage[] = []; + hasMoreToolCalls = false; + if (toolCalls.length > 0) { + const executedToolBatch = await executeToolCalls(currentContext, message, config, signal, emit); + toolResults.push(...executedToolBatch.messages); + hasMoreToolCalls = !executedToolBatch.terminate; + + for (const result of toolResults) { + currentContext.messages.push(result); + newMessages.push(result); + } + } + + await emit({ type: "turn_end", message, toolResults }); + + const nextTurnContext = { + message, + toolResults, + context: currentContext, + newMessages, + }; + const nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext); + if (nextTurnSnapshot) { + currentContext = nextTurnSnapshot.context ?? currentContext; + config = { + ...config, + model: nextTurnSnapshot.model ?? config.model, + reasoning: + nextTurnSnapshot.thinkingLevel === undefined + ? config.reasoning + : nextTurnSnapshot.thinkingLevel === "off" + ? undefined + : nextTurnSnapshot.thinkingLevel, + }; + } + + if ( + await config.shouldStopAfterTurn?.({ + message, + toolResults, + context: currentContext, + newMessages, + }) + ) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + + pendingMessages = (await config.getSteeringMessages?.()) || []; + } + + // Agent would stop here. Check for follow-up messages. + const followUpMessages = (await config.getFollowUpMessages?.()) || []; + if (followUpMessages.length > 0) { + // Set as pending so inner loop processes them + pendingMessages = followUpMessages; + continue; + } + + // No more messages, exit + break; + } + + await emit({ type: "agent_end", messages: newMessages }); +} + +/** + * Stream an assistant response from the LLM. + * This is where AgentMessage[] gets transformed to Message[] for the LLM. + */ +async function streamAssistantResponse( + context: AgentContext, + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, + streamFn?: StreamFn, +): Promise { + // Apply context transform if configured (AgentMessage[] → AgentMessage[]) + let messages = context.messages; + if (config.transformContext) { + messages = await config.transformContext(messages, signal); + } + + // Convert to LLM-compatible messages (AgentMessage[] → Message[]) + const llmMessages = await config.convertToLlm(messages); + + // Build LLM context + const llmContext: Context = { + systemPrompt: context.systemPrompt, + messages: llmMessages, + tools: context.tools, + }; + + const streamFunction = streamFn || streamSimple; + + // Resolve API key (important for expiring tokens) + const resolvedApiKey = + (config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey; + + const response = await streamFunction(config.model, llmContext, { + ...config, + apiKey: resolvedApiKey, + signal, + }); + + let partialMessage: AssistantMessage | null = null; + let addedPartial = false; + + for await (const event of response) { + switch (event.type) { + case "start": + partialMessage = event.partial; + context.messages.push(partialMessage); + addedPartial = true; + await emit({ type: "message_start", message: { ...partialMessage } }); + break; + + case "text_start": + case "text_delta": + case "text_end": + case "thinking_start": + case "thinking_delta": + case "thinking_end": + case "toolcall_start": + case "toolcall_delta": + case "toolcall_end": + if (partialMessage) { + partialMessage = event.partial; + context.messages[context.messages.length - 1] = partialMessage; + await emit({ + type: "message_update", + assistantMessageEvent: event, + message: { ...partialMessage }, + }); + } + break; + + case "done": + case "error": { + const finalMessage = await response.result(); + if (addedPartial) { + context.messages[context.messages.length - 1] = finalMessage; + } else { + context.messages.push(finalMessage); + } + if (!addedPartial) { + await emit({ type: "message_start", message: { ...finalMessage } }); + } + await emit({ type: "message_end", message: finalMessage }); + return finalMessage; + } + } + } + + const finalMessage = await response.result(); + if (addedPartial) { + context.messages[context.messages.length - 1] = finalMessage; + } else { + context.messages.push(finalMessage); + await emit({ type: "message_start", message: { ...finalMessage } }); + } + await emit({ type: "message_end", message: finalMessage }); + return finalMessage; +} + +/** + * Execute tool calls from an assistant message. + */ +async function executeToolCalls( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall"); + const hasSequentialToolCall = toolCalls.some( + (tc) => currentContext.tools?.find((t) => t.name === tc.name)?.executionMode === "sequential", + ); + if (config.toolExecution === "sequential" || hasSequentialToolCall) { + return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit); + } + return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit); +} + +type ExecutedToolCallBatch = { + messages: ToolResultMessage[]; + terminate: boolean; +}; + +async function executeToolCallsSequential( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + toolCalls: AgentToolCall[], + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const finalizedCalls: FinalizedToolCallOutcome[] = []; + const messages: ToolResultMessage[] = []; + + for (const toolCall of toolCalls) { + await emit({ + type: "tool_execution_start", + toolCallId: toolCall.id, + toolName: toolCall.name, + args: toolCall.arguments, + }); + + const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal); + let finalized: FinalizedToolCallOutcome; + if (preparation.kind === "immediate") { + finalized = { + toolCall, + result: preparation.result, + isError: preparation.isError, + }; + } else { + const executed = await executePreparedToolCall(preparation, signal, emit); + finalized = await finalizeExecutedToolCall( + currentContext, + assistantMessage, + preparation, + executed, + config, + signal, + ); + } + + await emitToolExecutionEnd(finalized, emit); + const toolResultMessage = createToolResultMessage(finalized); + await emitToolResultMessage(toolResultMessage, emit); + finalizedCalls.push(finalized); + messages.push(toolResultMessage); + + if (signal?.aborted) { + break; + } + } + + return { + messages, + terminate: shouldTerminateToolBatch(finalizedCalls), + }; +} + +async function executeToolCallsParallel( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + toolCalls: AgentToolCall[], + config: AgentLoopConfig, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const finalizedCalls: FinalizedToolCallEntry[] = []; + + for (const toolCall of toolCalls) { + await emit({ + type: "tool_execution_start", + toolCallId: toolCall.id, + toolName: toolCall.name, + args: toolCall.arguments, + }); + + const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal); + if (preparation.kind === "immediate") { + const finalized = { + toolCall, + result: preparation.result, + isError: preparation.isError, + } satisfies FinalizedToolCallOutcome; + await emitToolExecutionEnd(finalized, emit); + finalizedCalls.push(finalized); + if (signal?.aborted) { + break; + } + continue; + } + + finalizedCalls.push(async () => { + const executed = await executePreparedToolCall(preparation, signal, emit); + const finalized = await finalizeExecutedToolCall( + currentContext, + assistantMessage, + preparation, + executed, + config, + signal, + ); + await emitToolExecutionEnd(finalized, emit); + return finalized; + }); + if (signal?.aborted) { + break; + } + } + + const orderedFinalizedCalls = await Promise.all( + finalizedCalls.map((entry) => (typeof entry === "function" ? entry() : Promise.resolve(entry))), + ); + const messages: ToolResultMessage[] = []; + for (const finalized of orderedFinalizedCalls) { + const toolResultMessage = createToolResultMessage(finalized); + await emitToolResultMessage(toolResultMessage, emit); + messages.push(toolResultMessage); + } + + return { + messages, + terminate: shouldTerminateToolBatch(orderedFinalizedCalls), + }; +} + +type PreparedToolCall = { + kind: "prepared"; + toolCall: AgentToolCall; + tool: AgentTool; + args: unknown; +}; + +type ImmediateToolCallOutcome = { + kind: "immediate"; + result: AgentToolResult; + isError: boolean; +}; + +type ExecutedToolCallOutcome = { + result: AgentToolResult; + isError: boolean; +}; + +type FinalizedToolCallOutcome = { + toolCall: AgentToolCall; + result: AgentToolResult; + isError: boolean; +}; + +type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise); + +function shouldTerminateToolBatch(finalizedCalls: FinalizedToolCallOutcome[]): boolean { + return finalizedCalls.length > 0 && finalizedCalls.every((finalized) => finalized.result.terminate === true); +} + +function prepareToolCallArguments(tool: AgentTool, toolCall: AgentToolCall): AgentToolCall { + if (!tool.prepareArguments) { + return toolCall; + } + const preparedArguments = tool.prepareArguments(toolCall.arguments); + if (preparedArguments === toolCall.arguments) { + return toolCall; + } + return { + ...toolCall, + arguments: preparedArguments as Record, + }; +} + +async function prepareToolCall( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + toolCall: AgentToolCall, + config: AgentLoopConfig, + signal: AbortSignal | undefined, +): Promise { + const tool = currentContext.tools?.find((t) => t.name === toolCall.name); + if (!tool) { + return { + kind: "immediate", + result: createErrorToolResult(`Tool ${toolCall.name} not found`), + isError: true, + }; + } + + try { + const preparedToolCall = prepareToolCallArguments(tool, toolCall); + const validatedArgs = validateToolArguments(tool, preparedToolCall); + if (config.beforeToolCall) { + const beforeResult = await config.beforeToolCall( + { + assistantMessage, + toolCall, + args: validatedArgs, + context: currentContext, + }, + signal, + ); + if (signal?.aborted) { + return { + kind: "immediate", + result: createErrorToolResult("Operation aborted"), + isError: true, + }; + } + if (beforeResult?.block) { + return { + kind: "immediate", + result: createErrorToolResult(beforeResult.reason || "Tool execution was blocked"), + isError: true, + }; + } + } + if (signal?.aborted) { + return { + kind: "immediate", + result: createErrorToolResult("Operation aborted"), + isError: true, + }; + } + return { + kind: "prepared", + toolCall, + tool, + args: validatedArgs, + }; + } catch (error) { + return { + kind: "immediate", + result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + isError: true, + }; + } +} + +async function executePreparedToolCall( + prepared: PreparedToolCall, + signal: AbortSignal | undefined, + emit: AgentEventSink, +): Promise { + const updateEvents: Promise[] = []; + let acceptingUpdates = true; + + try { + const result = await prepared.tool.execute( + prepared.toolCall.id, + prepared.args as never, + signal, + (partialResult) => { + if (!acceptingUpdates) return; + updateEvents.push( + Promise.resolve( + emit({ + type: "tool_execution_update", + toolCallId: prepared.toolCall.id, + toolName: prepared.toolCall.name, + args: prepared.toolCall.arguments, + partialResult, + }), + ), + ); + }, + ); + acceptingUpdates = false; + await Promise.all(updateEvents); + return { result, isError: false }; + } catch (error) { + acceptingUpdates = false; + await Promise.all(updateEvents); + return { + result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + isError: true, + }; + } finally { + acceptingUpdates = false; + } +} + +async function finalizeExecutedToolCall( + currentContext: AgentContext, + assistantMessage: AssistantMessage, + prepared: PreparedToolCall, + executed: ExecutedToolCallOutcome, + config: AgentLoopConfig, + signal: AbortSignal | undefined, +): Promise { + let result = executed.result; + let isError = executed.isError; + + if (config.afterToolCall) { + try { + const afterResult = await config.afterToolCall( + { + assistantMessage, + toolCall: prepared.toolCall, + args: prepared.args, + result, + isError, + context: currentContext, + }, + signal, + ); + if (afterResult) { + result = { + content: afterResult.content ?? result.content, + details: afterResult.details ?? result.details, + terminate: afterResult.terminate ?? result.terminate, + }; + isError = afterResult.isError ?? isError; + } + } catch (error) { + result = createErrorToolResult(error instanceof Error ? error.message : String(error)); + isError = true; + } + } + + return { + toolCall: prepared.toolCall, + result, + isError, + }; +} + +function createErrorToolResult(message: string): AgentToolResult { + return { + content: [{ type: "text", text: message }], + details: {}, + }; +} + +async function emitToolExecutionEnd(finalized: FinalizedToolCallOutcome, emit: AgentEventSink): Promise { + await emit({ + type: "tool_execution_end", + toolCallId: finalized.toolCall.id, + toolName: finalized.toolCall.name, + result: finalized.result, + isError: finalized.isError, + }); +} + +function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResultMessage { + return { + role: "toolResult", + toolCallId: finalized.toolCall.id, + toolName: finalized.toolCall.name, + content: finalized.result.content, + details: finalized.result.details, + isError: finalized.isError, + timestamp: Date.now(), + }; +} + +async function emitToolResultMessage(toolResultMessage: ToolResultMessage, emit: AgentEventSink): Promise { + await emit({ type: "message_start", message: toolResultMessage }); + await emit({ type: "message_end", message: toolResultMessage }); +} diff --git a/cactus-code/packages/agent/src/agent.ts b/cactus-code/packages/agent/src/agent.ts new file mode 100644 index 000000000..54020435e --- /dev/null +++ b/cactus-code/packages/agent/src/agent.ts @@ -0,0 +1,557 @@ +import { + type ImageContent, + type Message, + type Model, + type SimpleStreamOptions, + streamSimple, + type TextContent, + type ThinkingBudgets, + type Transport, +} from "@earendil-works/pi-ai/compat"; +import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts"; +import type { + AfterToolCallContext, + AfterToolCallResult, + AgentContext, + AgentEvent, + AgentLoopConfig, + AgentLoopTurnUpdate, + AgentMessage, + AgentState, + AgentTool, + BeforeToolCallContext, + BeforeToolCallResult, + QueueMode, + StreamFn, + ToolExecutionMode, +} from "./types.ts"; + +export type { QueueMode } from "./types.ts"; + +function defaultConvertToLlm(messages: AgentMessage[]): Message[] { + return messages.filter( + (message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult", + ); +} + +const EMPTY_USAGE = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +const DEFAULT_MODEL = { + id: "unknown", + name: "unknown", + api: "unknown", + provider: "unknown", + baseUrl: "", + reasoning: false, + input: [], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 0, + maxTokens: 0, +} satisfies Model; + +type MutableAgentState = Omit & { + isStreaming: boolean; + streamingMessage?: AgentMessage; + pendingToolCalls: Set; + errorMessage?: string; +}; + +function createMutableAgentState( + initialState?: Partial>, +): MutableAgentState { + let tools = initialState?.tools?.slice() ?? []; + let messages = initialState?.messages?.slice() ?? []; + + return { + systemPrompt: initialState?.systemPrompt ?? "", + model: initialState?.model ?? DEFAULT_MODEL, + thinkingLevel: initialState?.thinkingLevel ?? "off", + get tools() { + return tools; + }, + set tools(nextTools: AgentTool[]) { + tools = nextTools.slice(); + }, + get messages() { + return messages; + }, + set messages(nextMessages: AgentMessage[]) { + messages = nextMessages.slice(); + }, + isStreaming: false, + streamingMessage: undefined, + pendingToolCalls: new Set(), + errorMessage: undefined, + }; +} + +/** Options for constructing an {@link Agent}. */ +export interface AgentOptions { + initialState?: Partial>; + convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise; + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; + streamFn?: StreamFn; + getApiKey?: (provider: string) => Promise | string | undefined; + onPayload?: SimpleStreamOptions["onPayload"]; + onResponse?: SimpleStreamOptions["onResponse"]; + beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; + afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; + prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; + steeringMode?: QueueMode; + followUpMode?: QueueMode; + sessionId?: string; + thinkingBudgets?: ThinkingBudgets; + transport?: Transport; + maxRetryDelayMs?: number; + toolExecution?: ToolExecutionMode; +} + +class PendingMessageQueue { + private messages: AgentMessage[] = []; + public mode: QueueMode; + + constructor(mode: QueueMode) { + this.mode = mode; + } + + enqueue(message: AgentMessage): void { + this.messages.push(message); + } + + hasItems(): boolean { + return this.messages.length > 0; + } + + drain(): AgentMessage[] { + if (this.mode === "all") { + const drained = this.messages.slice(); + this.messages = []; + return drained; + } + + const first = this.messages[0]; + if (!first) { + return []; + } + this.messages = this.messages.slice(1); + return [first]; + } + + clear(): void { + this.messages = []; + } +} + +type ActiveRun = { + promise: Promise; + resolve: () => void; + abortController: AbortController; +}; + +/** + * Stateful wrapper around the low-level agent loop. + * + * `Agent` owns the current transcript, emits lifecycle events, executes tools, + * and exposes queueing APIs for steering and follow-up messages. + */ +export class Agent { + private _state: MutableAgentState; + private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise | void>(); + private readonly steeringQueue: PendingMessageQueue; + private readonly followUpQueue: PendingMessageQueue; + + public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; + public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; + public streamFn: StreamFn; + public getApiKey?: (provider: string) => Promise | string | undefined; + public onPayload?: SimpleStreamOptions["onPayload"]; + public onResponse?: SimpleStreamOptions["onResponse"]; + public beforeToolCall?: ( + context: BeforeToolCallContext, + signal?: AbortSignal, + ) => Promise; + public afterToolCall?: ( + context: AfterToolCallContext, + signal?: AbortSignal, + ) => Promise; + public prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; + private activeRun?: ActiveRun; + /** Session identifier forwarded to providers for cache-aware backends. */ + public sessionId?: string; + /** Optional per-level thinking token budgets forwarded to the stream function. */ + public thinkingBudgets?: ThinkingBudgets; + /** Preferred transport forwarded to the stream function. */ + public transport: Transport; + /** Optional cap for provider-requested retry delays. */ + public maxRetryDelayMs?: number; + /** Tool execution strategy for assistant messages that contain multiple tool calls. */ + public toolExecution: ToolExecutionMode; + + constructor(options: AgentOptions = {}) { + this._state = createMutableAgentState(options.initialState); + this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm; + this.transformContext = options.transformContext; + this.streamFn = options.streamFn ?? streamSimple; + this.getApiKey = options.getApiKey; + this.onPayload = options.onPayload; + this.onResponse = options.onResponse; + this.beforeToolCall = options.beforeToolCall; + this.afterToolCall = options.afterToolCall; + this.prepareNextTurn = options.prepareNextTurn; + this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); + this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time"); + this.sessionId = options.sessionId; + this.thinkingBudgets = options.thinkingBudgets; + this.transport = options.transport ?? "auto"; + this.maxRetryDelayMs = options.maxRetryDelayMs; + this.toolExecution = options.toolExecution ?? "parallel"; + } + + /** + * Subscribe to agent lifecycle events. + * + * Listener promises are awaited in subscription order and are included in + * the current run's settlement. Listeners also receive the active abort + * signal for the current run. + * + * `agent_end` is the final emitted event for a run, but the agent does not + * become idle until all awaited listeners for that event have settled. + */ + subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise | void): () => void { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + } + + /** + * Current agent state. + * + * Assigning `state.tools` or `state.messages` copies the provided top-level array. + */ + get state(): AgentState { + return this._state; + } + + /** Controls how queued steering messages are drained. */ + set steeringMode(mode: QueueMode) { + this.steeringQueue.mode = mode; + } + + get steeringMode(): QueueMode { + return this.steeringQueue.mode; + } + + /** Controls how queued follow-up messages are drained. */ + set followUpMode(mode: QueueMode) { + this.followUpQueue.mode = mode; + } + + get followUpMode(): QueueMode { + return this.followUpQueue.mode; + } + + /** Queue a message to be injected after the current assistant turn finishes. */ + steer(message: AgentMessage): void { + this.steeringQueue.enqueue(message); + } + + /** Queue a message to run only after the agent would otherwise stop. */ + followUp(message: AgentMessage): void { + this.followUpQueue.enqueue(message); + } + + /** Remove all queued steering messages. */ + clearSteeringQueue(): void { + this.steeringQueue.clear(); + } + + /** Remove all queued follow-up messages. */ + clearFollowUpQueue(): void { + this.followUpQueue.clear(); + } + + /** Remove all queued steering and follow-up messages. */ + clearAllQueues(): void { + this.clearSteeringQueue(); + this.clearFollowUpQueue(); + } + + /** Returns true when either queue still contains pending messages. */ + hasQueuedMessages(): boolean { + return this.steeringQueue.hasItems() || this.followUpQueue.hasItems(); + } + + /** Active abort signal for the current run, if any. */ + get signal(): AbortSignal | undefined { + return this.activeRun?.abortController.signal; + } + + /** Abort the current run, if one is active. */ + abort(): void { + this.activeRun?.abortController.abort(); + } + + /** + * Resolve when the current run and all awaited event listeners have finished. + * + * This resolves after `agent_end` listeners settle. + */ + waitForIdle(): Promise { + return this.activeRun?.promise ?? Promise.resolve(); + } + + /** Clear transcript state, runtime state, and queued messages. */ + reset(): void { + this._state.messages = []; + this._state.isStreaming = false; + this._state.streamingMessage = undefined; + this._state.pendingToolCalls = new Set(); + this._state.errorMessage = undefined; + this.clearFollowUpQueue(); + this.clearSteeringQueue(); + } + + /** Start a new prompt from text, a single message, or a batch of messages. */ + async prompt(message: AgentMessage | AgentMessage[]): Promise; + async prompt(input: string, images?: ImageContent[]): Promise; + async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise { + if (this.activeRun) { + throw new Error( + "Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.", + ); + } + const messages = this.normalizePromptInput(input, images); + await this.runPromptMessages(messages); + } + + /** Continue from the current transcript. The last message must be a user or tool-result message. */ + async continue(): Promise { + if (this.activeRun) { + throw new Error("Agent is already processing. Wait for completion before continuing."); + } + + const lastMessage = this._state.messages[this._state.messages.length - 1]; + if (!lastMessage) { + throw new Error("No messages to continue from"); + } + + if (lastMessage.role === "assistant") { + const queuedSteering = this.steeringQueue.drain(); + if (queuedSteering.length > 0) { + await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true }); + return; + } + + const queuedFollowUps = this.followUpQueue.drain(); + if (queuedFollowUps.length > 0) { + await this.runPromptMessages(queuedFollowUps); + return; + } + + throw new Error("Cannot continue from message role: assistant"); + } + + await this.runContinuation(); + } + + private normalizePromptInput( + input: string | AgentMessage | AgentMessage[], + images?: ImageContent[], + ): AgentMessage[] { + if (Array.isArray(input)) { + return input; + } + + if (typeof input !== "string") { + return [input]; + } + + const content: Array = [{ type: "text", text: input }]; + if (images && images.length > 0) { + content.push(...images); + } + return [{ role: "user", content, timestamp: Date.now() }]; + } + + private async runPromptMessages( + messages: AgentMessage[], + options: { skipInitialSteeringPoll?: boolean } = {}, + ): Promise { + await this.runWithLifecycle(async (signal) => { + await runAgentLoop( + messages, + this.createContextSnapshot(), + this.createLoopConfig(options), + (event) => this.processEvents(event), + signal, + this.streamFn, + ); + }); + } + + private async runContinuation(): Promise { + await this.runWithLifecycle(async (signal) => { + await runAgentLoopContinue( + this.createContextSnapshot(), + this.createLoopConfig(), + (event) => this.processEvents(event), + signal, + this.streamFn, + ); + }); + } + + private createContextSnapshot(): AgentContext { + return { + systemPrompt: this._state.systemPrompt, + messages: this._state.messages.slice(), + tools: this._state.tools.slice(), + }; + } + + private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig { + let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; + return { + model: this._state.model, + reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, + sessionId: this.sessionId, + onPayload: this.onPayload, + onResponse: this.onResponse, + transport: this.transport, + thinkingBudgets: this.thinkingBudgets, + maxRetryDelayMs: this.maxRetryDelayMs, + toolExecution: this.toolExecution, + beforeToolCall: this.beforeToolCall, + afterToolCall: this.afterToolCall, + prepareNextTurn: this.prepareNextTurn ? async () => await this.prepareNextTurn?.(this.signal) : undefined, + convertToLlm: this.convertToLlm, + transformContext: this.transformContext, + getApiKey: this.getApiKey, + getSteeringMessages: async () => { + if (skipInitialSteeringPoll) { + skipInitialSteeringPoll = false; + return []; + } + return this.steeringQueue.drain(); + }, + getFollowUpMessages: async () => this.followUpQueue.drain(), + }; + } + + private async runWithLifecycle(executor: (signal: AbortSignal) => Promise): Promise { + if (this.activeRun) { + throw new Error("Agent is already processing."); + } + + const abortController = new AbortController(); + let resolvePromise = () => {}; + const promise = new Promise((resolve) => { + resolvePromise = resolve; + }); + this.activeRun = { promise, resolve: resolvePromise, abortController }; + + this._state.isStreaming = true; + this._state.streamingMessage = undefined; + this._state.errorMessage = undefined; + + try { + await executor(abortController.signal); + } catch (error) { + await this.handleRunFailure(error, abortController.signal.aborted); + } finally { + this.finishRun(); + } + } + + private async handleRunFailure(error: unknown, aborted: boolean): Promise { + const failureMessage = { + role: "assistant", + content: [{ type: "text", text: "" }], + api: this._state.model.api, + provider: this._state.model.provider, + model: this._state.model.id, + usage: EMPTY_USAGE, + stopReason: aborted ? "aborted" : "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + } satisfies AgentMessage; + await this.processEvents({ type: "message_start", message: failureMessage }); + await this.processEvents({ type: "message_end", message: failureMessage }); + await this.processEvents({ type: "turn_end", message: failureMessage, toolResults: [] }); + await this.processEvents({ type: "agent_end", messages: [failureMessage] }); + } + + private finishRun(): void { + this._state.isStreaming = false; + this._state.streamingMessage = undefined; + this._state.pendingToolCalls = new Set(); + this.activeRun?.resolve(); + this.activeRun = undefined; + } + + /** + * Reduce internal state for a loop event, then await listeners. + * + * `agent_end` only means no further loop events will be emitted. The run is + * considered idle later, after all awaited listeners for `agent_end` finish + * and `finishRun()` clears runtime-owned state. + */ + private async processEvents(event: AgentEvent): Promise { + switch (event.type) { + case "message_start": + this._state.streamingMessage = event.message; + break; + + case "message_update": + this._state.streamingMessage = event.message; + break; + + case "message_end": + this._state.streamingMessage = undefined; + this._state.messages.push(event.message); + break; + + case "tool_execution_start": { + const pendingToolCalls = new Set(this._state.pendingToolCalls); + pendingToolCalls.add(event.toolCallId); + this._state.pendingToolCalls = pendingToolCalls; + break; + } + + case "tool_execution_end": { + const pendingToolCalls = new Set(this._state.pendingToolCalls); + pendingToolCalls.delete(event.toolCallId); + this._state.pendingToolCalls = pendingToolCalls; + break; + } + + case "turn_end": + if (event.message.role === "assistant" && event.message.errorMessage) { + this._state.errorMessage = event.message.errorMessage; + } + break; + + case "agent_end": + this._state.streamingMessage = undefined; + break; + } + + const signal = this.activeRun?.abortController.signal; + if (!signal) { + throw new Error("Agent listener invoked outside active run"); + } + for (const listener of this.listeners) { + await listener(event, signal); + } + } +} diff --git a/cactus-code/packages/agent/src/harness/agent-harness.ts b/cactus-code/packages/agent/src/harness/agent-harness.ts new file mode 100644 index 000000000..1d09b054e --- /dev/null +++ b/cactus-code/packages/agent/src/harness/agent-harness.ts @@ -0,0 +1,1029 @@ +import type { AssistantMessage, ImageContent, Model, Models, UserMessage } from "@earendil-works/pi-ai"; +import { runAgentLoop } from "../agent-loop.ts"; +import type { + AgentContext, + AgentEvent, + AgentLoopConfig, + AgentMessage, + AgentTool, + QueueMode, + StreamFn, + ThinkingLevel, +} from "../types.ts"; +import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.ts"; +import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.ts"; +import { convertToLlm } from "./messages.ts"; +import { formatPromptTemplateInvocation } from "./prompt-templates.ts"; +import { formatSkillInvocation } from "./skills.ts"; +import type { + AbortResult, + AgentHarnessEvent, + AgentHarnessEventResultMap, + AgentHarnessOptions, + AgentHarnessOwnEvent, + AgentHarnessPhase, + AgentHarnessResources, + AgentHarnessStreamOptions, + AgentHarnessStreamOptionsPatch, + ExecutionEnv, + NavigateTreeResult, + PendingSessionWrite, + PromptTemplate, + Session, + Skill, +} from "./types.ts"; +import { AgentHarnessError, BranchSummaryError, CompactionError, SessionError, toError } from "./types.ts"; + +function createUserMessage(text: string, images?: ImageContent[]): UserMessage { + const content: Array<{ type: "text"; text: string } | ImageContent> = [{ type: "text", text }]; + if (images) content.push(...images); + return { role: "user", content, timestamp: Date.now() }; +} + +function createFailureMessage(model: Model, error: unknown, aborted: boolean): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text: "" }], + api: model.api, + provider: model.provider, + model: model.id, + stopReason: aborted ? "aborted" : "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; +} + +function cloneStreamOptions(streamOptions?: AgentHarnessStreamOptions): AgentHarnessStreamOptions { + return { + ...streamOptions, + headers: streamOptions?.headers ? { ...streamOptions.headers } : undefined, + metadata: streamOptions?.metadata ? { ...streamOptions.metadata } : undefined, + }; +} + +function findDuplicateNames(names: string[]): string[] { + const seen = new Set(); + const duplicates = new Set(); + for (const name of names) { + if (seen.has(name)) duplicates.add(name); + seen.add(name); + } + return [...duplicates]; +} + +function applyStreamOptionsPatch( + base: AgentHarnessStreamOptions, + patch?: AgentHarnessStreamOptionsPatch, +): AgentHarnessStreamOptions { + const result = cloneStreamOptions(base); + if (!patch) return result; + + if (Object.hasOwn(patch, "transport")) result.transport = patch.transport; + if (Object.hasOwn(patch, "timeoutMs")) result.timeoutMs = patch.timeoutMs; + if (Object.hasOwn(patch, "maxRetries")) result.maxRetries = patch.maxRetries; + if (Object.hasOwn(patch, "maxRetryDelayMs")) result.maxRetryDelayMs = patch.maxRetryDelayMs; + if (Object.hasOwn(patch, "cacheRetention")) result.cacheRetention = patch.cacheRetention; + + if (Object.hasOwn(patch, "headers")) { + if (patch.headers === undefined) { + result.headers = undefined; + } else { + const headers = { ...(result.headers ?? {}) }; + for (const [key, value] of Object.entries(patch.headers)) { + if (value === undefined) delete headers[key]; + else headers[key] = value; + } + result.headers = Object.keys(headers).length > 0 ? headers : undefined; + } + } + + if (Object.hasOwn(patch, "metadata")) { + if (patch.metadata === undefined) { + result.metadata = undefined; + } else { + const metadata = { ...(result.metadata ?? {}) }; + for (const [key, value] of Object.entries(patch.metadata)) { + if (value === undefined) delete metadata[key]; + else metadata[key] = value; + } + result.metadata = Object.keys(metadata).length > 0 ? metadata : undefined; + } + } + + return result; +} + +const SUBSCRIBER_EVENT_TYPE = "*"; + +type AgentHarnessHandler = (event: any, signal?: AbortSignal) => Promise | any; + +function normalizeHarnessError(error: unknown, fallbackCode: AgentHarnessError["code"]): AgentHarnessError { + if (error instanceof AgentHarnessError) return error; + const cause = toError(error); + if (cause instanceof SessionError) return new AgentHarnessError("session", cause.message, cause); + if (cause instanceof CompactionError) return new AgentHarnessError("compaction", cause.message, cause); + if (cause instanceof BranchSummaryError) return new AgentHarnessError("branch_summary", cause.message, cause); + return new AgentHarnessError(fallbackCode, cause.message, cause); +} + +function normalizeHookError(error: unknown): AgentHarnessError { + return normalizeHarnessError(error, "hook"); +} + +interface AgentHarnessTurnState< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, + TTool extends AgentTool = AgentTool, +> { + messages: AgentMessage[]; + resources: AgentHarnessResources; + streamOptions: AgentHarnessStreamOptions; + sessionId: string; + systemPrompt: string; + model: Model; + thinkingLevel: ThinkingLevel; + tools: TTool[]; + activeTools: TTool[]; +} + +export class AgentHarness< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, + TTool extends AgentTool = AgentTool, +> { + readonly env: ExecutionEnv; + private session: Session; + readonly models: Models; + private phase: AgentHarnessPhase = "idle"; + private runAbortController?: AbortController; + private runPromise?: Promise; + private pendingSessionWrites: PendingSessionWrite[] = []; + private model: Model; + private thinkingLevel: ThinkingLevel; + private systemPrompt: AgentHarnessOptions["systemPrompt"]; + private streamOptions: AgentHarnessStreamOptions; + private resources: AgentHarnessResources; + private tools = new Map(); + private activeToolNames: string[]; + private steerQueue: UserMessage[] = []; + private steeringQueueMode: QueueMode; + private followUpQueue: UserMessage[] = []; + private followUpQueueMode: QueueMode; + private nextTurnQueue: AgentMessage[] = []; + private handlers = new Map>(); + + constructor(options: AgentHarnessOptions) { + this.env = options.env; + this.session = options.session; + this.models = options.models; + this.resources = options.resources ?? {}; + this.streamOptions = cloneStreamOptions(options.streamOptions); + this.systemPrompt = options.systemPrompt; + this.validateUniqueNames( + (options.tools ?? []).map((tool) => tool.name), + "Duplicate tool name(s)", + ); + for (const tool of options.tools ?? []) { + this.tools.set(tool.name, tool); + } + this.model = options.model; + this.thinkingLevel = options.thinkingLevel ?? "off"; + this.activeToolNames = options.activeToolNames + ? [...options.activeToolNames] + : (options.tools ?? []).map((tool) => tool.name); + this.validateUniqueNames(this.activeToolNames, "Duplicate active tool name(s)"); + this.validateToolNames(this.activeToolNames); + this.steeringQueueMode = options.steeringMode ?? "one-at-a-time"; + this.followUpQueueMode = options.followUpMode ?? "one-at-a-time"; + } + + private getHandlers(type: string): Set | undefined { + return this.handlers.get(type); + } + + private async emitOwn(event: AgentHarnessOwnEvent, signal?: AbortSignal): Promise { + for (const listener of this.getHandlers(SUBSCRIBER_EVENT_TYPE) ?? []) { + try { + await listener(event, signal); + } catch (error) { + throw normalizeHookError(error); + } + } + } + + private async emitAny(event: AgentHarnessEvent, signal?: AbortSignal): Promise { + for (const listener of this.getHandlers(SUBSCRIBER_EVENT_TYPE) ?? []) { + try { + await listener(event, signal); + } catch (error) { + throw normalizeHookError(error); + } + } + } + + private async emitHook( + event: Extract, + ): Promise { + const handlers = this.getHandlers(event.type as TType); + if (!handlers || handlers.size === 0) return undefined; + let lastResult: AgentHarnessEventResultMap[TType] | undefined; + for (const handler of handlers) { + try { + const result = await handler(event); + if (result !== undefined) { + lastResult = result; + } + } catch (error) { + throw normalizeHookError(error); + } + } + return lastResult; + } + + private async emitBeforeProviderRequest( + model: Model, + sessionId: string, + streamOptions: AgentHarnessStreamOptions, + ): Promise { + const handlers = this.getHandlers("before_provider_request"); + let current = cloneStreamOptions(streamOptions); + if (!handlers || handlers.size === 0) return current; + for (const handler of handlers) { + try { + const result = await handler({ + type: "before_provider_request", + model, + sessionId, + streamOptions: cloneStreamOptions(current), + }); + if (result?.streamOptions) { + current = applyStreamOptionsPatch(current, result.streamOptions); + } + } catch (error) { + throw normalizeHookError(error); + } + } + return current; + } + + private async emitBeforeProviderPayload(model: Model, payload: unknown): Promise { + const handlers = this.getHandlers("before_provider_payload"); + let current = payload; + if (!handlers || handlers.size === 0) return current; + for (const handler of handlers) { + try { + const result = await handler({ type: "before_provider_payload", model, payload: current }); + if (result !== undefined) { + current = result.payload; + } + } catch (error) { + throw normalizeHookError(error); + } + } + return current; + } + + private async emitQueueUpdate(): Promise { + await this.emitOwn({ + type: "queue_update", + steer: [...this.steerQueue], + followUp: [...this.followUpQueue], + nextTurn: [...this.nextTurnQueue], + }); + } + + private startRunPromise(): () => void { + let finish = () => {}; + this.runPromise = new Promise((resolve) => { + finish = resolve; + }); + return () => { + this.runPromise = undefined; + finish(); + }; + } + + private async createTurnState(): Promise> { + const context = await this.session.buildContext(); + const resources = this.getResources(); + const sessionMetadata = await this.session.getMetadata(); + const tools = [...this.tools.values()]; + const activeTools = this.activeToolNames + .map((name) => this.tools.get(name)) + .filter((tool): tool is TTool => tool !== undefined); + let systemPrompt = "You are a helpful assistant."; + if (typeof this.systemPrompt === "string") { + systemPrompt = this.systemPrompt; + } else if (this.systemPrompt) { + systemPrompt = await this.systemPrompt({ + env: this.env, + session: this.session, + model: this.model, + thinkingLevel: this.thinkingLevel, + activeTools, + resources, + }); + } + return { + messages: context.messages, + resources, + streamOptions: cloneStreamOptions(this.streamOptions), + sessionId: sessionMetadata.id, + systemPrompt, + model: this.model, + thinkingLevel: this.thinkingLevel, + tools, + activeTools, + }; + } + + private createContext( + turnState: AgentHarnessTurnState, + systemPrompt?: string, + ): AgentContext { + return { + systemPrompt: systemPrompt ?? turnState.systemPrompt, + messages: turnState.messages.slice(), + tools: turnState.activeTools.slice(), + }; + } + + private createStreamFn(getTurnState: () => AgentHarnessTurnState): StreamFn { + return async (model, context, streamOptions) => { + const turnState = getTurnState(); + const snapshotOptions: AgentHarnessStreamOptions = { ...turnState.streamOptions }; + const requestOptions = await this.emitBeforeProviderRequest(model, turnState.sessionId, snapshotOptions); + return this.models.streamSimple(model, context, { + cacheRetention: requestOptions.cacheRetention, + headers: requestOptions.headers, + maxRetries: requestOptions.maxRetries, + maxRetryDelayMs: requestOptions.maxRetryDelayMs, + metadata: requestOptions.metadata, + onPayload: async (payload) => await this.emitBeforeProviderPayload(model, payload), + onResponse: async (response) => { + const headers = { ...(response.headers as Record) }; + await this.emitOwn( + { type: "after_provider_response", status: response.status, headers }, + streamOptions?.signal, + ); + }, + reasoning: streamOptions?.reasoning, + signal: streamOptions?.signal, + sessionId: turnState.sessionId, + timeoutMs: requestOptions.timeoutMs, + transport: requestOptions.transport, + }); + }; + } + + private async drainQueuedMessages(queue: AgentMessage[], mode: QueueMode): Promise { + const messages = mode === "all" ? queue.splice(0) : queue.splice(0, 1); + if (messages.length === 0) return messages; + try { + await this.emitQueueUpdate(); + return messages; + } catch (error) { + queue.unshift(...messages); + throw normalizeHookError(error); + } + } + + private createLoopConfig( + getTurnState: () => AgentHarnessTurnState, + setTurnState: (turnState: AgentHarnessTurnState) => void, + ): AgentLoopConfig { + const turnState = getTurnState(); + return { + model: turnState.model, + reasoning: turnState.thinkingLevel === "off" ? undefined : turnState.thinkingLevel, + convertToLlm, + transformContext: async (messages) => { + const result = await this.emitHook({ type: "context", messages: [...messages] }); + return result?.messages ?? messages; + }, + beforeToolCall: async ({ toolCall, args }) => { + const result = await this.emitHook({ + type: "tool_call", + toolCallId: toolCall.id, + toolName: toolCall.name, + input: args as Record, + }); + return result ? { block: result.block, reason: result.reason } : undefined; + }, + afterToolCall: async ({ toolCall, args, result, isError }) => { + const patch = await this.emitHook({ + type: "tool_result", + toolCallId: toolCall.id, + toolName: toolCall.name, + input: args as Record, + content: result.content, + details: result.details, + isError, + }); + return patch + ? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate } + : undefined; + }, + prepareNextTurn: async () => { + await this.flushPendingSessionWrites(); + const nextTurnState = await this.createTurnState(); + setTurnState(nextTurnState); + return { + context: this.createContext(nextTurnState), + model: nextTurnState.model, + thinkingLevel: nextTurnState.thinkingLevel, + }; + }, + getSteeringMessages: async () => this.drainQueuedMessages(this.steerQueue, this.steeringQueueMode), + getFollowUpMessages: async () => this.drainQueuedMessages(this.followUpQueue, this.followUpQueueMode), + }; + } + + private validateUniqueNames(names: string[], message: string): void { + const duplicates = findDuplicateNames(names); + if (duplicates.length > 0) + throw new AgentHarnessError("invalid_argument", `${message}: ${duplicates.join(", ")}`); + } + + private validateToolNames(toolNames: string[], tools: Map = this.tools): void { + this.validateUniqueNames(toolNames, "Duplicate active tool name(s)"); + const missing = toolNames.filter((name) => !tools.has(name)); + if (missing.length > 0) throw new AgentHarnessError("invalid_argument", `Unknown tool(s): ${missing.join(", ")}`); + } + + private async flushPendingSessionWrites(): Promise { + while (this.pendingSessionWrites.length > 0) { + const write = this.pendingSessionWrites[0]!; + if (write.type === "message") { + await this.session.appendMessage(write.message); + } else if (write.type === "model_change") { + await this.session.appendModelChange(write.provider, write.modelId); + } else if (write.type === "thinking_level_change") { + await this.session.appendThinkingLevelChange(write.thinkingLevel); + } else if (write.type === "active_tools_change") { + await this.session.appendActiveToolsChange(write.activeToolNames); + } else if (write.type === "custom") { + await this.session.appendCustomEntry(write.customType, write.data); + } else if (write.type === "custom_message") { + await this.session.appendCustomMessageEntry(write.customType, write.content, write.display, write.details); + } else if (write.type === "label") { + await this.session.appendLabel(write.targetId, write.label); + } else if (write.type === "session_info") { + await this.session.appendSessionName(write.name ?? ""); + } else if (write.type === "leaf") { + await this.session.getStorage().setLeafId(write.targetId); + } + this.pendingSessionWrites.shift(); + } + } + + private async handleAgentEvent(event: AgentEvent, signal?: AbortSignal): Promise { + if (event.type === "message_end") { + await this.session.appendMessage(event.message); + await this.emitAny(event, signal); + return; + } + if (event.type === "turn_end") { + let eventError: unknown; + try { + await this.emitAny(event, signal); + } catch (error) { + eventError = error; + } + const hadPendingMutations = this.pendingSessionWrites.length > 0; + await this.flushPendingSessionWrites(); + if (eventError) throw eventError; + await this.emitOwn({ type: "save_point", hadPendingMutations }); + return; + } + if (event.type === "agent_end") { + await this.flushPendingSessionWrites(); + this.phase = "idle"; + await this.emitAny(event, signal); + await this.emitOwn({ type: "settled", nextTurnCount: this.nextTurnQueue.length }, signal); + return; + } + await this.emitAny(event, signal); + } + + private async emitRunFailure( + model: Model, + error: unknown, + aborted: boolean, + signal: AbortSignal, + ): Promise { + const failureMessage = createFailureMessage(model, error, aborted); + await this.handleAgentEvent({ type: "message_start", message: failureMessage }, signal); + await this.handleAgentEvent({ type: "message_end", message: failureMessage }, signal); + await this.handleAgentEvent({ type: "turn_end", message: failureMessage, toolResults: [] }, signal); + await this.handleAgentEvent({ type: "agent_end", messages: [failureMessage] }, signal); + return [failureMessage]; + } + + private async executeTurn( + turnState: AgentHarnessTurnState, + text: string, + options?: { images?: ImageContent[] }, + ): Promise { + let activeTurnState = turnState; + let messages: AgentMessage[] = [createUserMessage(text, options?.images)]; + if (this.nextTurnQueue.length > 0) { + const queuedMessages = this.nextTurnQueue.splice(0); + try { + await this.emitQueueUpdate(); + } catch (error) { + this.nextTurnQueue.unshift(...queuedMessages); + throw normalizeHookError(error); + } + messages = [...queuedMessages, messages[0]!]; + } + const beforeResult = await this.emitHook({ + type: "before_agent_start", + prompt: text, + images: options?.images, + systemPrompt: turnState.systemPrompt, + resources: turnState.resources, + }); + if (beforeResult?.messages) messages = [...messages, ...beforeResult.messages]; + + const abortController = new AbortController(); + const getTurnState = () => activeTurnState; + const setTurnState = (nextTurnState: AgentHarnessTurnState) => { + activeTurnState = nextTurnState; + }; + this.runAbortController = abortController; + const runResultPromise = (async () => { + try { + return await runAgentLoop( + messages, + this.createContext(turnState, beforeResult?.systemPrompt), + this.createLoopConfig(getTurnState, setTurnState), + (event) => this.handleAgentEvent(event, abortController.signal), + abortController.signal, + this.createStreamFn(getTurnState), + ); + } catch (error) { + try { + return await this.emitRunFailure( + activeTurnState.model, + error, + abortController.signal.aborted, + abortController.signal, + ); + } catch (failureError) { + const cause = new AggregateError( + [toError(error), toError(failureError)], + "Agent run failed and failure reporting failed", + ); + throw new AgentHarnessError("unknown", cause.message, cause); + } + } + })(); + try { + const newMessages = await runResultPromise; + for (let i = newMessages.length - 1; i >= 0; i--) { + const message = newMessages[i]!; + if (message.role === "assistant") { + return message; + } + } + throw new AgentHarnessError("invalid_state", "AgentHarness prompt completed without an assistant message"); + } finally { + try { + await this.flushPendingSessionWrites(); + } finally { + this.runAbortController = undefined; + } + } + } + + async prompt(text: string, options?: { images?: ImageContent[] }): Promise { + if (this.phase !== "idle") throw new AgentHarnessError("busy", "AgentHarness is busy"); + this.phase = "turn"; + const finishRunPromise = this.startRunPromise(); + try { + const turnState = await this.createTurnState(); + return await this.executeTurn(turnState, text, options); + } catch (error) { + this.phase = "idle"; + throw normalizeHarnessError(error, "unknown"); + } finally { + finishRunPromise(); + } + } + + async skill(name: string, additionalInstructions?: string): Promise { + if (this.phase !== "idle") throw new AgentHarnessError("busy", "AgentHarness is busy"); + this.phase = "turn"; + const finishRunPromise = this.startRunPromise(); + try { + const turnState = await this.createTurnState(); + const skill = (turnState.resources.skills ?? []).find((candidate) => candidate.name === name); + if (!skill) throw new AgentHarnessError("invalid_argument", `Unknown skill: ${name}`); + return await this.executeTurn(turnState, formatSkillInvocation(skill, additionalInstructions)); + } catch (error) { + this.phase = "idle"; + throw normalizeHarnessError(error, "unknown"); + } finally { + finishRunPromise(); + } + } + + async promptFromTemplate(name: string, args: string[] = []): Promise { + if (this.phase !== "idle") throw new AgentHarnessError("busy", "AgentHarness is busy"); + this.phase = "turn"; + const finishRunPromise = this.startRunPromise(); + try { + const turnState = await this.createTurnState(); + const template = (turnState.resources.promptTemplates ?? []).find((candidate) => candidate.name === name); + if (!template) throw new AgentHarnessError("invalid_argument", `Unknown prompt template: ${name}`); + return await this.executeTurn(turnState, formatPromptTemplateInvocation(template, args)); + } catch (error) { + this.phase = "idle"; + throw normalizeHarnessError(error, "unknown"); + } finally { + finishRunPromise(); + } + } + + async steer(text: string, options?: { images?: ImageContent[] }): Promise { + if (this.phase === "idle") throw new AgentHarnessError("invalid_state", "Cannot steer while idle"); + this.steerQueue.push(createUserMessage(text, options?.images)); + await this.emitQueueUpdate(); + } + + async followUp(text: string, options?: { images?: ImageContent[] }): Promise { + if (this.phase === "idle") throw new AgentHarnessError("invalid_state", "Cannot follow up while idle"); + this.followUpQueue.push(createUserMessage(text, options?.images)); + await this.emitQueueUpdate(); + } + + async nextTurn(text: string, options?: { images?: ImageContent[] }): Promise { + this.nextTurnQueue.push(createUserMessage(text, options?.images)); + await this.emitQueueUpdate(); + } + + async appendMessage(message: AgentMessage): Promise { + try { + if (this.phase === "idle") { + await this.session.appendMessage(message); + } else { + this.pendingSessionWrites.push({ type: "message", message }); + } + } catch (error) { + throw normalizeHarnessError(error, "session"); + } + } + + async compact( + customInstructions?: string, + ): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> { + if (this.phase !== "idle") throw new AgentHarnessError("busy", "compact() requires idle harness"); + this.phase = "compaction"; + try { + const model = this.model; + if (!model) throw new AgentHarnessError("invalid_state", "No model set for compaction"); + const branchEntries = await this.session.getBranch(); + const preparationResult = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS); + if (!preparationResult.ok) throw preparationResult.error; + const preparation = preparationResult.value; + if (!preparation) throw new AgentHarnessError("compaction", "Nothing to compact"); + const hookResult = await this.emitHook({ + type: "session_before_compact", + preparation, + branchEntries, + customInstructions, + signal: new AbortController().signal, + }); + if (hookResult?.cancel) throw new AgentHarnessError("compaction", "Compaction cancelled"); + const provided = hookResult?.compaction; + const compactResult = provided + ? { ok: true as const, value: provided } + : await compact(preparation, this.models, model, customInstructions, undefined, this.thinkingLevel); + if (!compactResult.ok) throw compactResult.error; + const result = compactResult.value; + const entryId = await this.session.appendCompaction( + result.summary, + result.firstKeptEntryId, + result.tokensBefore, + result.details, + provided !== undefined, + ); + const entry = await this.session.getEntry(entryId); + if (entry?.type === "compaction") { + await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined }); + } + return result; + } catch (error) { + throw normalizeHarnessError(error, "compaction"); + } finally { + this.phase = "idle"; + } + } + + async navigateTree( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ): Promise { + if (this.phase !== "idle") throw new AgentHarnessError("busy", "navigateTree() requires idle harness"); + this.phase = "branch_summary"; + try { + const oldLeafId = await this.session.getLeafId(); + if (oldLeafId === targetId) return { cancelled: false }; + const targetEntry = await this.session.getEntry(targetId); + if (!targetEntry) throw new AgentHarnessError("invalid_argument", `Entry ${targetId} not found`); + const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.session, oldLeafId, targetId); + const preparation = { + targetId, + oldLeafId, + commonAncestorId, + entriesToSummarize: entries, + userWantsSummary: options?.summarize ?? false, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }; + const signal = new AbortController().signal; + const hookResult = await this.emitHook({ type: "session_before_tree", preparation, signal }); + if (hookResult?.cancel) return { cancelled: true }; + let summaryEntry: NavigateTreeResult["summaryEntry"]; + let summaryText: string | undefined = hookResult?.summary?.summary; + let summaryDetails: unknown = hookResult?.summary?.details; + if (!summaryText && options?.summarize && entries.length > 0) { + const model = this.model; + if (!model) throw new AgentHarnessError("invalid_state", "No model set for branch summary"); + const branchSummary = await generateBranchSummary(entries, { + models: this.models, + model, + signal: new AbortController().signal, + customInstructions: hookResult?.customInstructions ?? options?.customInstructions, + replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, + }); + if (!branchSummary.ok) { + if (branchSummary.error.code === "aborted") return { cancelled: true }; + throw new AgentHarnessError("branch_summary", branchSummary.error.message, branchSummary.error); + } + summaryText = branchSummary.value.summary; + summaryDetails = { + readFiles: branchSummary.value.readFiles, + modifiedFiles: branchSummary.value.modifiedFiles, + }; + } + let editorText: string | undefined; + let newLeafId: string | null; + if (targetEntry.type === "message" && targetEntry.message.role === "user") { + newLeafId = targetEntry.parentId; + const content = targetEntry.message.content; + editorText = + typeof content === "string" + ? content + : content + .filter((c): c is { readonly type: "text"; readonly text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } else if (targetEntry.type === "custom_message") { + newLeafId = targetEntry.parentId; + editorText = + typeof targetEntry.content === "string" + ? targetEntry.content + : targetEntry.content + .filter((c): c is { readonly type: "text"; readonly text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } else { + newLeafId = targetId; + } + const summaryId = await this.session.moveTo( + newLeafId, + summaryText + ? { summary: summaryText, details: summaryDetails, fromHook: hookResult?.summary !== undefined } + : undefined, + ); + if (summaryId) { + const entry = await this.session.getEntry(summaryId); + if (entry?.type === "branch_summary") summaryEntry = entry; + } + await this.emitOwn({ + type: "session_tree", + newLeafId: await this.session.getLeafId(), + oldLeafId, + summaryEntry, + fromHook: hookResult?.summary !== undefined, + }); + return { cancelled: false, editorText, summaryEntry }; + } catch (error) { + throw normalizeHarnessError(error, "branch_summary"); + } finally { + this.phase = "idle"; + } + } + + getModel(): Model { + return this.model; + } + + async setModel(model: Model): Promise { + try { + const previousModel = this.model; + if (this.phase === "idle") { + await this.session.appendModelChange(model.provider, model.id); + } else { + this.pendingSessionWrites.push({ type: "model_change", provider: model.provider, modelId: model.id }); + } + this.model = model; + await this.emitOwn({ type: "model_update", model, previousModel, source: "set" }); + } catch (error) { + throw normalizeHarnessError(error, "session"); + } + } + + getThinkingLevel(): ThinkingLevel { + return this.thinkingLevel; + } + + async setThinkingLevel(level: ThinkingLevel): Promise { + try { + const previousLevel = this.thinkingLevel; + if (this.phase === "idle") { + await this.session.appendThinkingLevelChange(level); + } else { + this.pendingSessionWrites.push({ type: "thinking_level_change", thinkingLevel: level }); + } + this.thinkingLevel = level; + await this.emitOwn({ type: "thinking_level_update", level, previousLevel }); + } catch (error) { + throw normalizeHarnessError(error, "session"); + } + } + + getTools(): TTool[] { + return [...this.tools.values()]; + } + + async setTools(tools: TTool[], activeToolNames?: string[]): Promise { + try { + this.validateUniqueNames( + tools.map((tool) => tool.name), + "Duplicate tool name(s)", + ); + const nextTools = new Map(tools.map((tool) => [tool.name, tool])); + const nextActiveToolNames = activeToolNames ? [...activeToolNames] : this.activeToolNames; + this.validateToolNames(nextActiveToolNames, nextTools); + const previousToolNames = [...this.tools.keys()]; + const previousActiveToolNames = [...this.activeToolNames]; + if (this.phase === "idle") { + await this.session.appendActiveToolsChange(nextActiveToolNames); + } else { + this.pendingSessionWrites.push({ type: "active_tools_change", activeToolNames: [...nextActiveToolNames] }); + } + this.tools = nextTools; + this.activeToolNames = [...nextActiveToolNames]; + await this.emitOwn({ + type: "tools_update", + toolNames: [...this.tools.keys()], + previousToolNames, + activeToolNames: [...this.activeToolNames], + previousActiveToolNames, + source: "set", + }); + } catch (error) { + throw normalizeHarnessError(error, "invalid_argument"); + } + } + + getActiveTools(): TTool[] { + return this.activeToolNames.map((name) => this.tools.get(name)!); + } + + async setActiveTools(toolNames: string[]): Promise { + try { + this.validateToolNames(toolNames); + const previousToolNames = [...this.tools.keys()]; + const previousActiveToolNames = [...this.activeToolNames]; + if (this.phase === "idle") { + await this.session.appendActiveToolsChange(toolNames); + } else { + this.pendingSessionWrites.push({ type: "active_tools_change", activeToolNames: [...toolNames] }); + } + this.activeToolNames = [...toolNames]; + await this.emitOwn({ + type: "tools_update", + toolNames: [...this.tools.keys()], + previousToolNames, + activeToolNames: [...this.activeToolNames], + previousActiveToolNames, + source: "set", + }); + } catch (error) { + throw normalizeHarnessError(error, "invalid_argument"); + } + } + + getSteeringMode(): QueueMode { + return this.steeringQueueMode; + } + + async setSteeringMode(mode: QueueMode): Promise { + this.steeringQueueMode = mode; + } + + getFollowUpMode(): QueueMode { + return this.followUpQueueMode; + } + + async setFollowUpMode(mode: QueueMode): Promise { + this.followUpQueueMode = mode; + } + + getResources(): AgentHarnessResources { + return { + skills: this.resources.skills?.slice(), + promptTemplates: this.resources.promptTemplates?.slice(), + }; + } + + async setResources(resources: AgentHarnessResources): Promise { + const previousResources = this.getResources(); + this.resources = { + skills: resources.skills?.slice(), + promptTemplates: resources.promptTemplates?.slice(), + }; + await this.emitOwn({ type: "resources_update", resources: this.getResources(), previousResources }); + } + + getStreamOptions(): AgentHarnessStreamOptions { + return cloneStreamOptions(this.streamOptions); + } + + async setStreamOptions(streamOptions: AgentHarnessStreamOptions): Promise { + this.streamOptions = cloneStreamOptions(streamOptions); + } + + async abort(): Promise { + const clearedSteer = [...this.steerQueue]; + const clearedFollowUp = [...this.followUpQueue]; + this.steerQueue = []; + this.followUpQueue = []; + this.runAbortController?.abort(); + const errors: Error[] = []; + try { + await this.emitQueueUpdate(); + } catch (error) { + errors.push(toError(error)); + } + try { + await this.waitForIdle(); + } catch (error) { + errors.push(toError(error)); + } + try { + await this.emitOwn({ type: "abort", clearedSteer, clearedFollowUp }); + } catch (error) { + errors.push(toError(error)); + } + if (errors.length > 0) { + const cause = errors.length === 1 ? errors[0]! : new AggregateError(errors, "Abort completed with errors"); + throw normalizeHarnessError(cause, "hook"); + } + return { clearedSteer, clearedFollowUp }; + } + + async waitForIdle(): Promise { + await this.runPromise; + } + + subscribe( + listener: (event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void, + ): () => void { + let handlers = this.handlers.get(SUBSCRIBER_EVENT_TYPE); + if (!handlers) { + handlers = new Set(); + this.handlers.set(SUBSCRIBER_EVENT_TYPE, handlers); + } + handlers.add(listener as AgentHarnessHandler); + return () => handlers!.delete(listener as AgentHarnessHandler); + } + + on( + type: TType, + handler: ( + event: Extract, + ) => Promise | AgentHarnessEventResultMap[TType], + ): () => void { + let handlers = this.handlers.get(type); + if (!handlers) { + handlers = new Set(); + this.handlers.set(type, handlers); + } + handlers.add(handler as AgentHarnessHandler); + return () => handlers!.delete(handler as AgentHarnessHandler); + } +} diff --git a/cactus-code/packages/agent/src/harness/compaction/branch-summarization.ts b/cactus-code/packages/agent/src/harness/compaction/branch-summarization.ts new file mode 100644 index 000000000..fdf1df498 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/compaction/branch-summarization.ts @@ -0,0 +1,261 @@ +import type { Model, Models } from "@earendil-works/pi-ai"; + +import type { AgentMessage } from "../../types.ts"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.ts"; +import type { BranchSummaryResult, Session, SessionTreeEntry } from "../types.ts"; +import { BranchSummaryError, err, ok, type Result, SessionError } from "../types.ts"; +import { estimateTokens, SUMMARIZATION_SYSTEM_PROMPT } from "./compaction.ts"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + serializeConversation, +} from "./utils.ts"; + +/** File-operation details stored on generated branch summary entries. */ +export interface BranchSummaryDetails { + /** Files read while exploring the summarized branch. */ + readFiles: string[]; + /** Files modified while exploring the summarized branch. */ + modifiedFiles: string[]; +} + +export type { FileOperations } from "./utils.ts"; + +/** Prepared branch content for summarization. */ +export interface BranchPreparation { + /** Messages selected for the branch summary. */ + messages: AgentMessage[]; + /** File operations extracted from the branch. */ + fileOps: FileOperations; + /** Estimated token count for selected messages. */ + totalTokens: number; +} + +/** Entries selected for branch summarization. */ +export interface CollectEntriesResult { + /** Entries to summarize in chronological order. */ + entries: SessionTreeEntry[]; + /** Deepest common ancestor between the previous leaf and target entry. */ + commonAncestorId: string | null; +} + +/** Options for generating a branch summary. */ +export interface GenerateBranchSummaryOptions { + /** Provider collection the summarization request goes through; owns auth resolution. */ + models: Models; + /** Model used for summarization. */ + model: Model; + /** Abort signal for the summarization request. */ + signal: AbortSignal; + /** Optional instructions appended to or replacing the default prompt. */ + customInstructions?: string; + /** Replace the default prompt with custom instructions instead of appending them. */ + replaceInstructions?: boolean; + /** Tokens reserved for prompt and model output. Defaults to 16384. */ + reserveTokens?: number; +} + +/** Collect entries that should be summarized before navigating to a different session tree entry. */ +export async function collectEntriesForBranchSummary( + session: Session, + oldLeafId: string | null, + targetId: string, +): Promise { + if (!oldLeafId) { + return { entries: [], commonAncestorId: null }; + } + const oldPath = new Set((await session.getBranch(oldLeafId)).map((e) => e.id)); + const targetPath = await session.getBranch(targetId); + let commonAncestorId: string | null = null; + for (let i = targetPath.length - 1; i >= 0; i--) { + if (oldPath.has(targetPath[i].id)) { + commonAncestorId = targetPath[i].id; + break; + } + } + const entries: SessionTreeEntry[] = []; + let current: string | null = oldLeafId; + + while (current && current !== commonAncestorId) { + const entry = await session.getEntry(current); + if (!entry) throw new SessionError("invalid_session", `Entry ${current} not found`); + entries.push(entry as SessionTreeEntry); + current = entry.parentId; + } + entries.reverse(); + + return { entries, commonAncestorId }; +} +function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined { + switch (entry.type) { + case "message": + if (entry.message.role === "toolResult") return undefined; + return entry.message; + + case "custom_message": + return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + + case "branch_summary": + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); + + case "compaction": + return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); + case "thinking_level_change": + case "model_change": + case "active_tools_change": + case "custom": + case "label": + case "session_info": + case "leaf": + return undefined; + } +} + +/** Prepare branch entries for summarization within an optional token budget. */ +export function prepareBranchEntries(entries: SessionTreeEntry[], tokenBudget: number = 0): BranchPreparation { + const messages: AgentMessage[] = []; + const fileOps = createFileOps(); + let totalTokens = 0; + for (const entry of entries) { + if (entry.type === "branch_summary" && !entry.fromHook && entry.details) { + const details = entry.details as BranchSummaryDetails; + if (Array.isArray(details.readFiles)) { + for (const f of details.readFiles) fileOps.read.add(f); + } + if (Array.isArray(details.modifiedFiles)) { + for (const f of details.modifiedFiles) { + fileOps.edited.add(f); + } + } + } + } + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + const message = getMessageFromEntry(entry); + if (!message) continue; + extractFileOpsFromMessage(message, fileOps); + + const tokens = estimateTokens(message); + if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { + if (entry.type === "compaction" || entry.type === "branch_summary") { + if (totalTokens < tokenBudget * 0.9) { + messages.unshift(message); + totalTokens += tokens; + } + } + break; + } + + messages.unshift(message); + totalTokens += tokens; + } + + return { messages, fileOps, totalTokens }; +} + +const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here. +Summary of that exploration: + +`; + +const BRANCH_SUMMARY_PROMPT = `Create a structured summary of this conversation branch for context when returning later. + +Use this EXACT format: + +## Goal +[What was the user trying to accomplish in this branch?] + +## Constraints & Preferences +- [Any constraints, preferences, or requirements mentioned] +- [Or "(none)" if none were mentioned] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Work that was started but not finished] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [What should happen next to continue this work] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +/** Generate a summary for abandoned branch entries. */ +export async function generateBranchSummary( + entries: SessionTreeEntry[], + options: GenerateBranchSummaryOptions, +): Promise> { + const { models, model, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; + const contextWindow = model.contextWindow || 128000; + const tokenBudget = contextWindow - reserveTokens; + + const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget); + + if (messages.length === 0) { + return ok({ summary: "No content to summarize", readFiles: [], modifiedFiles: [] }); + } + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + let instructions: string; + if (replaceInstructions && customInstructions) { + instructions = customInstructions; + } else if (customInstructions) { + instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`; + } else { + instructions = BRANCH_SUMMARY_PROMPT; + } + const promptText = `\n${conversationText}\n\n\n${instructions}`; + + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + const response = await models.completeSimple( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + { signal, maxTokens: 2048 }, + ); + if (response.stopReason === "aborted") { + return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted")); + } + if (response.stopReason === "error") { + return err( + new BranchSummaryError( + "summarization_failed", + `Branch summary failed: ${response.errorMessage || "Unknown error"}`, + ), + ); + } + + let summary = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + summary = BRANCH_SUMMARY_PREAMBLE + summary; + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + return ok({ + summary: summary || "No summary generated", + readFiles, + modifiedFiles, + }); +} diff --git a/cactus-code/packages/agent/src/harness/compaction/compaction.ts b/cactus-code/packages/agent/src/harness/compaction/compaction.ts new file mode 100644 index 000000000..2d6d55837 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/compaction/compaction.ts @@ -0,0 +1,747 @@ +import type { AssistantMessage, ImageContent, Model, Models, TextContent, Usage } from "@earendil-works/pi-ai"; +import type { AgentMessage, ThinkingLevel } from "../../types.ts"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.ts"; +import { buildSessionContext } from "../session/session.ts"; +import { type CompactionEntry, CompactionError, err, ok, type Result, type SessionTreeEntry } from "../types.ts"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + serializeConversation, +} from "./utils.ts"; + +/** File-operation details stored on generated compaction entries. */ +export interface CompactionDetails { + /** Files read in the compacted history. */ + readFiles: string[]; + /** Files modified in the compacted history. */ + modifiedFiles: string[]; +} +function safeJsonStringify(value: unknown): string { + try { + return JSON.stringify(value) ?? "undefined"; + } catch { + return "[unserializable]"; + } +} + +function extractFileOperations( + messages: AgentMessage[], + entries: SessionTreeEntry[], + prevCompactionIndex: number, +): FileOperations { + const fileOps = createFileOps(); + if (prevCompactionIndex >= 0) { + const prevCompaction = entries[prevCompactionIndex] as CompactionEntry; + if (!prevCompaction.fromHook && prevCompaction.details) { + const details = prevCompaction.details as CompactionDetails; + if (Array.isArray(details.readFiles)) { + for (const f of details.readFiles) fileOps.read.add(f); + } + if (Array.isArray(details.modifiedFiles)) { + for (const f of details.modifiedFiles) fileOps.edited.add(f); + } + } + } + for (const msg of messages) { + extractFileOpsFromMessage(msg, fileOps); + } + + return fileOps; +} +function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined { + if (entry.type === "message") { + return entry.message as AgentMessage; + } + if (entry.type === "custom_message") { + return createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + entry.display, + entry.details, + entry.timestamp, + ); + } + if (entry.type === "branch_summary") { + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); + } + if (entry.type === "compaction") { + return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); + } + return undefined; +} + +function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage | undefined { + if (entry.type === "compaction") { + return undefined; + } + return getMessageFromEntry(entry); +} + +/** Generated compaction data ready to be persisted as a compaction entry. */ +export interface CompactionResult { + /** Summary text that replaces compacted history in future context. */ + summary: string; + /** Entry id where retained history starts. */ + firstKeptEntryId: string; + /** Estimated context tokens before compaction. */ + tokensBefore: number; + /** Optional implementation-specific details stored with the compaction entry. */ + details?: T; +} + +/** Compaction thresholds and retention settings. */ +export interface CompactionSettings { + /** Enable automatic compaction decisions. */ + enabled: boolean; + /** Tokens reserved for summary prompt and output. */ + reserveTokens: number; + /** Approximate recent-context tokens to keep after compaction. */ + keepRecentTokens: number; +} + +/** Default compaction settings used by the harness. */ +export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { + enabled: true, + reserveTokens: 16384, + keepRecentTokens: 20000, +}; + +/** Calculate total context tokens from provider usage. */ +export function calculateContextTokens(usage: Usage): number { + return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; +} +function getAssistantUsage(msg: AgentMessage): Usage | undefined { + if (msg.role === "assistant" && "usage" in msg) { + const assistantMsg = msg as AssistantMessage; + if ( + assistantMsg.stopReason !== "aborted" && + assistantMsg.stopReason !== "error" && + assistantMsg.usage && + calculateContextTokens(assistantMsg.usage) > 0 + ) { + return assistantMsg.usage; + } + } + return undefined; +} + +/** Return usage from the last valid assistant message in session entries. */ +export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | undefined { + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry.type === "message") { + const usage = getAssistantUsage(entry.message as AgentMessage); + if (usage) return usage; + } + } + return undefined; +} + +/** Estimated context-token usage for a message list. */ +export interface ContextUsageEstimate { + /** Estimated total context tokens. */ + tokens: number; + /** Tokens reported by the most recent assistant usage block. */ + usageTokens: number; + /** Estimated tokens after the most recent assistant usage block. */ + trailingTokens: number; + /** Index of the message that provided usage, or null when none exists. */ + lastUsageIndex: number | null; +} + +function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const usage = getAssistantUsage(messages[i]); + if (usage) return { usage, index: i }; + } + return undefined; +} + +/** Estimate context tokens for messages using provider usage when available. */ +export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate { + const usageInfo = getLastAssistantUsageInfo(messages); + + if (!usageInfo) { + let estimated = 0; + for (const message of messages) { + estimated += estimateTokens(message); + } + return { + tokens: estimated, + usageTokens: 0, + trailingTokens: estimated, + lastUsageIndex: null, + }; + } + + const usageTokens = calculateContextTokens(usageInfo.usage); + let trailingTokens = 0; + for (let i = usageInfo.index + 1; i < messages.length; i++) { + trailingTokens += estimateTokens(messages[i]); + } + + return { + tokens: usageTokens + trailingTokens, + usageTokens, + trailingTokens, + lastUsageIndex: usageInfo.index, + }; +} + +/** Return whether context usage exceeds the configured compaction threshold. */ +export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean { + if (!settings.enabled) return false; + return contextTokens > contextWindow - settings.reserveTokens; +} + +const ESTIMATED_IMAGE_CHARS = 4800; + +function estimateTextAndImageContentChars(content: string | Array<{ type: string; text?: string }>): number { + if (typeof content === "string") { + return content.length; + } + + let chars = 0; + for (const block of content) { + if (block.type === "text" && block.text) { + chars += block.text.length; + } else if (block.type === "image") { + chars += ESTIMATED_IMAGE_CHARS; + } + } + return chars; +} + +/** Estimate token count for one message using a conservative character heuristic. */ +export function estimateTokens(message: AgentMessage): number { + let chars = 0; + + switch (message.role) { + case "user": { + chars = estimateTextAndImageContentChars( + (message as { content: string | Array<{ type: string; text?: string }> }).content, + ); + return Math.ceil(chars / 4); + } + case "assistant": { + const assistant = message as AssistantMessage; + for (const block of assistant.content) { + if (block.type === "text") { + chars += block.text.length; + } else if (block.type === "thinking") { + chars += block.thinking.length; + } else if (block.type === "toolCall") { + chars += block.name.length + safeJsonStringify(block.arguments).length; + } + } + return Math.ceil(chars / 4); + } + case "custom": + case "toolResult": { + chars = estimateTextAndImageContentChars(message.content); + return Math.ceil(chars / 4); + } + case "bashExecution": { + chars = message.command.length + message.output.length; + return Math.ceil(chars / 4); + } + case "branchSummary": + case "compactionSummary": { + chars = message.summary.length; + return Math.ceil(chars / 4); + } + } + + return 0; +} +function findValidCutPoints(entries: SessionTreeEntry[], startIndex: number, endIndex: number): number[] { + const cutPoints: number[] = []; + for (let i = startIndex; i < endIndex; i++) { + const entry = entries[i]; + switch (entry.type) { + case "message": { + const role = entry.message.role; + switch (role) { + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + case "user": + case "assistant": + cutPoints.push(i); + break; + case "toolResult": + break; + } + break; + } + case "thinking_level_change": + case "model_change": + case "active_tools_change": + case "compaction": + case "branch_summary": + case "custom": + case "custom_message": + case "label": + case "session_info": + case "leaf": + break; + } + if (entry.type === "branch_summary" || entry.type === "custom_message") { + cutPoints.push(i); + } + } + return cutPoints; +} + +/** Find the user-visible message that starts the turn containing an entry. */ +export function findTurnStartIndex(entries: SessionTreeEntry[], entryIndex: number, startIndex: number): number { + for (let i = entryIndex; i >= startIndex; i--) { + const entry = entries[i]; + if (entry.type === "branch_summary" || entry.type === "custom_message") { + return i; + } + if (entry.type === "message") { + const role = entry.message.role; + if (role === "user" || role === "bashExecution") { + return i; + } + } + } + return -1; +} + +/** Cut point selected for compaction. */ +export interface CutPointResult { + /** Index of the first entry retained after compaction. */ + firstKeptEntryIndex: number; + /** Index of the turn-start entry when the cut splits a turn, otherwise -1. */ + turnStartIndex: number; + /** Whether the selected cut point splits an in-progress turn. */ + isSplitTurn: boolean; +} + +/** Find the compaction cut point that keeps approximately the requested recent-token budget. */ +export function findCutPoint( + entries: SessionTreeEntry[], + startIndex: number, + endIndex: number, + keepRecentTokens: number, +): CutPointResult { + const cutPoints = findValidCutPoints(entries, startIndex, endIndex); + + if (cutPoints.length === 0) { + return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false }; + } + let accumulatedTokens = 0; + let cutIndex = cutPoints[0]; + + for (let i = endIndex - 1; i >= startIndex; i--) { + const entry = entries[i]; + if (entry.type !== "message") continue; + const messageTokens = estimateTokens(entry.message as AgentMessage); + accumulatedTokens += messageTokens; + if (accumulatedTokens >= keepRecentTokens) { + for (let c = 0; c < cutPoints.length; c++) { + if (cutPoints[c] >= i) { + cutIndex = cutPoints[c]; + break; + } + } + break; + } + } + while (cutIndex > startIndex) { + const prevEntry = entries[cutIndex - 1]; + if (prevEntry.type === "compaction") { + break; + } + if (prevEntry.type === "message") { + break; + } + cutIndex--; + } + const cutEntry = entries[cutIndex]; + const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user"; + const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); + + return { + firstKeptEntryIndex: cutIndex, + turnStartIndex, + isSplitTurn: !isUserMessage && turnStartIndex !== -1, + }; +} + +export const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified. + +Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`; + +const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work. + +Use this EXACT format: + +## Goal +[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.] + +## Constraints & Preferences +- [Any constraints, preferences, or requirements mentioned by user] +- [Or "(none)" if none were mentioned] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Current work] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [Ordered list of what should happen next] + +## Critical Context +- [Any data, examples, or references needed to continue] +- [Or "(none)" if not applicable] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in tags. + +Update the existing structured summary with new information. RULES: +- PRESERVE all existing information from the previous summary +- ADD new progress, decisions, and context from the new messages +- UPDATE the Progress section: move items from "In Progress" to "Done" when completed +- UPDATE "Next Steps" based on what was accomplished +- PRESERVE exact file paths, function names, and error messages +- If something is no longer relevant, you may remove it + +Use this EXACT format: + +## Goal +[Preserve existing goals, add new ones if the task expanded] + +## Constraints & Preferences +- [Preserve existing, add new ones discovered] + +## Progress +### Done +- [x] [Include previously done items AND newly completed items] + +### In Progress +- [ ] [Current work - update based on progress] + +### Blocked +- [Current blockers - remove if resolved] + +## Key Decisions +- **[Decision]**: [Brief rationale] (preserve all previous, add new) + +## Next Steps +1. [Update based on current state] + +## Critical Context +- [Preserve important context, add new if needed] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +/** Generate or update a conversation summary for compaction. */ +export async function generateSummary( + currentMessages: AgentMessage[], + models: Models, + model: Model, + reserveTokens: number, + signal?: AbortSignal, + customInstructions?: string, + previousSummary?: string, + thinkingLevel?: ThinkingLevel, +): Promise> { + const maxTokens = Math.min( + Math.floor(0.8 * reserveTokens), + model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, + ); + let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; + if (customInstructions) { + basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; + } + const llmMessages = convertToLlm(currentMessages); + const conversationText = serializeConversation(llmMessages); + let promptText = `\n${conversationText}\n\n\n`; + if (previousSummary) { + promptText += `\n${previousSummary}\n\n\n`; + } + promptText += basePrompt; + + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + const completionOptions = + model.reasoning && thinkingLevel && thinkingLevel !== "off" + ? { maxTokens, signal, reasoning: thinkingLevel } + : { maxTokens, signal }; + + const response = await models.completeSimple( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + completionOptions, + ); + if (response.stopReason === "aborted") { + return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted")); + } + if (response.stopReason === "error") { + return err( + new CompactionError( + "summarization_failed", + `Summarization failed: ${response.errorMessage || "Unknown error"}`, + ), + ); + } + + const textContent = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + + return ok(textContent); +} + +/** Prepared inputs for a compaction run. */ +export interface CompactionPreparation { + /** Entry id where retained history starts. */ + firstKeptEntryId: string; + /** Messages summarized into the history summary. */ + messagesToSummarize: AgentMessage[]; + /** Prefix messages summarized separately when compaction splits a turn. */ + turnPrefixMessages: AgentMessage[]; + /** Whether compaction splits a turn. */ + isSplitTurn: boolean; + /** Estimated context tokens before compaction. */ + tokensBefore: number; + /** Previous compaction summary used for iterative updates. */ + previousSummary?: string; + /** File operations extracted from summarized history. */ + fileOps: FileOperations; + /** Settings used to prepare compaction. */ + settings: CompactionSettings; +} + +/** Prepare session entries for compaction, or return undefined when compaction is not applicable. */ +export function prepareCompaction( + pathEntries: SessionTreeEntry[], + settings: CompactionSettings, +): Result { + if (pathEntries.length === 0 || pathEntries[pathEntries.length - 1].type === "compaction") { + return ok(undefined); + } + + let prevCompactionIndex = -1; + for (let i = pathEntries.length - 1; i >= 0; i--) { + if (pathEntries[i].type === "compaction") { + prevCompactionIndex = i; + break; + } + } + + let previousSummary: string | undefined; + let boundaryStart = 0; + if (prevCompactionIndex >= 0) { + const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; + previousSummary = prevCompaction.summary; + const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); + boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; + } + const boundaryEnd = pathEntries.length; + + const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens; + + const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens); + const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex]; + if (!firstKeptEntry?.id) { + return err(new CompactionError("invalid_session", "First kept entry has no UUID - session may need migration")); + } + const firstKeptEntryId = firstKeptEntry.id; + + const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; + const messagesToSummarize: AgentMessage[] = []; + for (let i = boundaryStart; i < historyEnd; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) messagesToSummarize.push(msg); + } + const turnPrefixMessages: AgentMessage[] = []; + if (cutPoint.isSplitTurn) { + for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) turnPrefixMessages.push(msg); + } + } + const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); + if (cutPoint.isSplitTurn) { + for (const msg of turnPrefixMessages) { + extractFileOpsFromMessage(msg, fileOps); + } + } + + return ok({ + firstKeptEntryId, + messagesToSummarize, + turnPrefixMessages, + isSplitTurn: cutPoint.isSplitTurn, + tokensBefore, + previousSummary, + fileOps, + settings, + }); +} + +const TURN_PREFIX_SUMMARIZATION_PROMPT = `This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained. + +Summarize the prefix to provide context for the retained suffix: + +## Original Request +[What did the user ask for in this turn?] + +## Early Progress +- [Key decisions and work done in the prefix] + +## Context for Suffix +- [Information needed to understand the retained recent work] + +Be concise. Focus on what's needed to understand the kept suffix.`; + +export { serializeConversation } from "./utils.ts"; + +/** Generate compaction summary data from prepared session history. */ +export async function compact( + preparation: CompactionPreparation, + models: Models, + model: Model, + customInstructions?: string, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise> { + const { + firstKeptEntryId, + messagesToSummarize, + turnPrefixMessages, + isSplitTurn, + tokensBefore, + previousSummary, + fileOps, + settings, + } = preparation; + + if (!firstKeptEntryId) { + return err(new CompactionError("invalid_session", "First kept entry has no UUID - session may need migration")); + } + + let summary: string; + + if (isSplitTurn && turnPrefixMessages.length > 0) { + const [historyResult, turnPrefixResult] = await Promise.all([ + messagesToSummarize.length > 0 + ? generateSummary( + messagesToSummarize, + models, + model, + settings.reserveTokens, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ) + : Promise.resolve(ok("No prior history.")), + generateTurnPrefixSummary(turnPrefixMessages, models, model, settings.reserveTokens, signal, thinkingLevel), + ]); + if (!historyResult.ok) return err(historyResult.error); + if (!turnPrefixResult.ok) return err(turnPrefixResult.error); + summary = `${historyResult.value}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value}`; + } else { + const summaryResult = await generateSummary( + messagesToSummarize, + models, + model, + settings.reserveTokens, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ); + if (!summaryResult.ok) return err(summaryResult.error); + summary = summaryResult.value; + } + + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + return ok({ + summary, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles } as CompactionDetails, + }); +} +async function generateTurnPrefixSummary( + messages: AgentMessage[], + models: Models, + model: Model, + reserveTokens: number, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise> { + const maxTokens = Math.min( + Math.floor(0.5 * reserveTokens), + model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, + ); + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + const response = await models.completeSimple( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + model.reasoning && thinkingLevel && thinkingLevel !== "off" + ? { maxTokens, signal, reasoning: thinkingLevel } + : { maxTokens, signal }, + ); + if (response.stopReason === "aborted") { + return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted")); + } + if (response.stopReason === "error") { + return err( + new CompactionError( + "summarization_failed", + `Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`, + ), + ); + } + + return ok( + response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"), + ); +} diff --git a/cactus-code/packages/agent/src/harness/compaction/utils.ts b/cactus-code/packages/agent/src/harness/compaction/utils.ts new file mode 100644 index 000000000..07535b79b --- /dev/null +++ b/cactus-code/packages/agent/src/harness/compaction/utils.ts @@ -0,0 +1,144 @@ +import type { Message } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "../../types.ts"; + +/** File paths touched by a session branch or compaction range. */ +export interface FileOperations { + /** Files read but not necessarily modified. */ + read: Set; + /** Files written by full-file write operations. */ + written: Set; + /** Files modified by edit operations. */ + edited: Set; +} + +/** Create an empty file-operation accumulator. */ +export function createFileOps(): FileOperations { + return { + read: new Set(), + written: new Set(), + edited: new Set(), + }; +} + +/** Add file operations from assistant tool calls to an accumulator. */ +export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void { + if (message.role !== "assistant") return; + if (!("content" in message) || !Array.isArray(message.content)) return; + + for (const block of message.content) { + if (typeof block !== "object" || block === null) continue; + if (!("type" in block) || block.type !== "toolCall") continue; + if (!("arguments" in block) || !("name" in block)) continue; + + const args = block.arguments as Record | undefined; + if (!args) continue; + + const path = typeof args.path === "string" ? args.path : undefined; + if (!path) continue; + + switch (block.name) { + case "read": + fileOps.read.add(path); + break; + case "write": + fileOps.written.add(path); + break; + case "edit": + fileOps.edited.add(path); + break; + } + } +} + +/** Compute sorted read-only and modified file lists from accumulated operations. */ +export function computeFileLists(fileOps: FileOperations): { readFiles: string[]; modifiedFiles: string[] } { + const modified = new Set([...fileOps.edited, ...fileOps.written]); + const readOnly = [...fileOps.read].filter((f) => !modified.has(f)).sort(); + const modifiedFiles = [...modified].sort(); + return { readFiles: readOnly, modifiedFiles }; +} + +/** Format file lists as summary metadata tags. */ +export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string { + const sections: string[] = []; + if (readFiles.length > 0) { + sections.push(`\n${readFiles.join("\n")}\n`); + } + if (modifiedFiles.length > 0) { + sections.push(`\n${modifiedFiles.join("\n")}\n`); + } + if (sections.length === 0) return ""; + return `\n\n${sections.join("\n\n")}`; +} + +const TOOL_RESULT_MAX_CHARS = 2000; + +function safeJsonStringify(value: unknown): string { + try { + return JSON.stringify(value) ?? "undefined"; + } catch { + return "[unserializable]"; + } +} + +function truncateForSummary(text: string, maxChars: number): string { + if (text.length <= maxChars) return text; + const truncatedChars = text.length - maxChars; + return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`; +} + +/** Serialize LLM messages to plain text for summarization prompts. */ +export function serializeConversation(messages: Message[]): string { + const parts: string[] = []; + + for (const msg of messages) { + if (msg.role === "user") { + const content = + typeof msg.content === "string" + ? msg.content + : msg.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (content) parts.push(`[User]: ${content}`); + } else if (msg.role === "assistant") { + const textParts: string[] = []; + const thinkingParts: string[] = []; + const toolCalls: string[] = []; + + for (const block of msg.content) { + if (block.type === "text") { + textParts.push(block.text); + } else if (block.type === "thinking") { + thinkingParts.push(block.thinking); + } else if (block.type === "toolCall") { + const args = block.arguments as Record; + const argsStr = Object.entries(args) + .map(([k, v]) => `${k}=${safeJsonStringify(v)}`) + .join(", "); + toolCalls.push(`${block.name}(${argsStr})`); + } + } + + if (thinkingParts.length > 0) { + parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); + } + if (textParts.length > 0) { + parts.push(`[Assistant]: ${textParts.join("\n")}`); + } + if (toolCalls.length > 0) { + parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`); + } + } else if (msg.role === "toolResult") { + const content = msg.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (content) { + parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); + } + } + } + + return parts.join("\n\n"); +} diff --git a/cactus-code/packages/agent/src/harness/env/nodejs.ts b/cactus-code/packages/agent/src/harness/env/nodejs.ts new file mode 100644 index 000000000..3d929c824 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/env/nodejs.ts @@ -0,0 +1,550 @@ +import { spawn } from "node:child_process"; +import { randomUUID } from "node:crypto"; +import { constants, createReadStream } from "node:fs"; +import { + access, + appendFile, + lstat, + mkdir, + mkdtemp, + readdir, + readFile, + realpath, + rm, + writeFile, +} from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { isAbsolute, join, resolve } from "node:path"; +import { createInterface } from "node:readline"; +import { + type ExecutionEnv, + ExecutionError, + err, + FileError, + type FileInfo, + type FileKind, + ok, + type Result, + toError, +} from "../types.ts"; + +function resolvePath(cwd: string, path: string): string { + return isAbsolute(path) ? path : resolve(cwd, path); +} + +function fileKindFromStats(stats: { + isFile(): boolean; + isDirectory(): boolean; + isSymbolicLink(): boolean; +}): FileKind | undefined { + if (stats.isFile()) return "file"; + if (stats.isDirectory()) return "directory"; + if (stats.isSymbolicLink()) return "symlink"; + return undefined; +} + +function fileInfoFromStats( + path: string, + stats: { isFile(): boolean; isDirectory(): boolean; isSymbolicLink(): boolean; size: number; mtimeMs: number }, +): Result { + const kind = fileKindFromStats(stats); + if (!kind) return err(new FileError("invalid", "Unsupported file type", path)); + return ok({ + name: path.replace(/\/+$/, "").split("/").pop() ?? path, + path, + kind, + size: stats.size, + mtimeMs: stats.mtimeMs, + }); +} + +function isNodeError(error: unknown): error is NodeJS.ErrnoException { + return error instanceof Error && "code" in error; +} + +function toFileError(error: unknown, path?: string): FileError { + if (error instanceof FileError) return error; + const cause = toError(error); + if (isNodeError(error)) { + const message = error.message; + switch (error.code) { + case "ABORT_ERR": + return new FileError("aborted", message, path, cause); + case "ENOENT": + return new FileError("not_found", message, path, cause); + case "EACCES": + case "EPERM": + return new FileError("permission_denied", message, path, cause); + case "ENOTDIR": + return new FileError("not_directory", message, path, cause); + case "EISDIR": + return new FileError("is_directory", message, path, cause); + case "EINVAL": + return new FileError("invalid", message, path, cause); + } + } + return new FileError("unknown", cause.message, path, cause); +} + +function abortResult(signal: AbortSignal | undefined, path?: string): Result | undefined { + return signal?.aborted ? err(new FileError("aborted", "aborted", path)) : undefined; +} + +async function pathExists(path: string): Promise { + try { + await access(path, constants.F_OK); + return true; + } catch { + return false; + } +} + +async function runCommand( + command: string, + args: string[], + timeoutMs: number, +): Promise<{ stdout: string; status: number | null }> { + return await new Promise((resolve) => { + let stdout = ""; + let child: ReturnType; + try { + child = spawn(command, args, { + stdio: ["ignore", "pipe", "ignore"], + windowsHide: true, + }); + } catch { + resolve({ stdout: "", status: null }); + return; + } + const timeout = setTimeout(() => { + if (child.pid) killProcessTree(child.pid); + }, timeoutMs); + child.stdout?.setEncoding("utf8"); + child.stdout?.on("data", (chunk: string) => { + stdout += chunk; + }); + child.on("error", () => { + clearTimeout(timeout); + resolve({ stdout: "", status: null }); + }); + child.on("close", (status) => { + clearTimeout(timeout); + resolve({ stdout, status }); + }); + }); +} + +async function findBashOnPath(): Promise { + const result = + process.platform === "win32" + ? await runCommand("where", ["bash.exe"], 5000) + : await runCommand("which", ["bash"], 5000); + if (result.status !== 0 || !result.stdout) return null; + const firstMatch = result.stdout.trim().split(/\r?\n/)[0]; + return firstMatch && (await pathExists(firstMatch)) ? firstMatch : null; +} + +interface ShellConfig { + shell: string; + args: string[]; + commandTransport?: "argv" | "stdin"; +} + +function isLegacyWslBashPath(path: string): boolean { + const normalized = path.replace(/\//g, "\\").toLowerCase(); + return /^[a-z]:\\windows\\(?:system32|sysnative)\\bash\.exe$/.test(normalized); +} + +function getBashShellConfig(shell: string): ShellConfig { + return isLegacyWslBashPath(shell) ? { shell, args: ["-s"], commandTransport: "stdin" } : { shell, args: ["-c"] }; +} + +async function getShellConfig(customShellPath?: string): Promise> { + if (customShellPath) { + if (await pathExists(customShellPath)) { + return ok(getBashShellConfig(customShellPath)); + } + return err(new ExecutionError("shell_unavailable", `Custom shell path not found: ${customShellPath}`)); + } + if (process.platform === "win32") { + const candidates: string[] = []; + const programFiles = process.env.ProgramFiles; + if (programFiles) candidates.push(`${programFiles}\\Git\\bin\\bash.exe`); + const programFilesX86 = process.env["ProgramFiles(x86)"]; + if (programFilesX86) candidates.push(`${programFilesX86}\\Git\\bin\\bash.exe`); + for (const candidate of candidates) { + if (await pathExists(candidate)) { + return ok(getBashShellConfig(candidate)); + } + } + const bashOnPath = await findBashOnPath(); + if (bashOnPath) { + return ok(getBashShellConfig(bashOnPath)); + } + return err(new ExecutionError("shell_unavailable", "No bash shell found")); + } + + if (await pathExists("/bin/bash")) { + return ok(getBashShellConfig("/bin/bash")); + } + const bashOnPath = await findBashOnPath(); + if (bashOnPath) { + return ok(getBashShellConfig(bashOnPath)); + } + return ok({ shell: "sh", args: ["-c"] }); +} + +function getShellEnv(baseEnv?: NodeJS.ProcessEnv, extraEnv?: Record): NodeJS.ProcessEnv { + return { + ...process.env, + ...baseEnv, + ...extraEnv, + }; +} + +function killProcessTree(pid: number): void { + if (process.platform === "win32") { + try { + spawn("taskkill", ["/F", "/T", "/PID", String(pid)], { + stdio: "ignore", + detached: true, + windowsHide: true, + }); + } catch { + // Ignore errors. + } + return; + } + + try { + process.kill(-pid, "SIGKILL"); + } catch { + try { + process.kill(pid, "SIGKILL"); + } catch { + // Process already dead. + } + } +} + +export class NodeExecutionEnv implements ExecutionEnv { + cwd: string; + private shellPath?: string; + private shellEnv?: NodeJS.ProcessEnv; + + constructor(options: { cwd: string; shellPath?: string; shellEnv?: NodeJS.ProcessEnv }) { + this.cwd = options.cwd; + this.shellPath = options.shellPath; + this.shellEnv = options.shellEnv; + } + + async absolutePath(path: string): Promise> { + return ok(resolvePath(this.cwd, path)); + } + + async joinPath(parts: string[]): Promise> { + return ok(join(...parts)); + } + + async exec( + command: string, + options?: { + cwd?: string; + env?: Record; + timeout?: number; + abortSignal?: AbortSignal; + onStdout?: (chunk: string) => void; + onStderr?: (chunk: string) => void; + }, + ): Promise> { + if (options?.abortSignal?.aborted) return err(new ExecutionError("aborted", "aborted")); + + const cwd = options?.cwd ? resolvePath(this.cwd, options.cwd) : this.cwd; + const shellConfig = await getShellConfig(this.shellPath); + if (!shellConfig.ok) return shellConfig; + + return await new Promise((resolvePromise) => { + let stdout = ""; + let stderr = ""; + let settled = false; + let timedOut = false; + let callbackError: ExecutionError | undefined; + let child: ReturnType | undefined; + let timeoutId: ReturnType | undefined; + + const onAbort = () => { + if (child?.pid) { + killProcessTree(child.pid); + } + }; + + const settle = (result: Result<{ stdout: string; stderr: string; exitCode: number }, ExecutionError>) => { + if (timeoutId) clearTimeout(timeoutId); + if (options?.abortSignal) options.abortSignal.removeEventListener("abort", onAbort); + if (settled) return; + settled = true; + resolvePromise(result); + }; + + try { + const commandFromStdin = shellConfig.value.commandTransport === "stdin"; + child = spawn( + shellConfig.value.shell, + commandFromStdin ? shellConfig.value.args : [...shellConfig.value.args, command], + { + cwd, + detached: process.platform !== "win32", + env: getShellEnv(this.shellEnv, options?.env), + stdio: [commandFromStdin ? "pipe" : "ignore", "pipe", "pipe"], + windowsHide: true, + }, + ); + if (commandFromStdin) { + child.stdin?.on("error", () => {}); + child.stdin?.end(command); + } + } catch (error) { + const cause = toError(error); + settle(err(new ExecutionError("spawn_error", cause.message, cause))); + return; + } + + timeoutId = + typeof options?.timeout === "number" + ? setTimeout(() => { + timedOut = true; + if (child?.pid) { + killProcessTree(child.pid); + } + }, options.timeout * 1000) + : undefined; + + if (options?.abortSignal) { + if (options.abortSignal.aborted) { + onAbort(); + } else { + options.abortSignal.addEventListener("abort", onAbort, { once: true }); + } + } + + child.stdout?.setEncoding("utf8"); + child.stderr?.setEncoding("utf8"); + child.stdout?.on("data", (chunk: string) => { + stdout += chunk; + try { + options?.onStdout?.(chunk); + } catch (error) { + const cause = toError(error); + callbackError = new ExecutionError("callback_error", cause.message, cause); + onAbort(); + } + }); + child.stderr?.on("data", (chunk: string) => { + stderr += chunk; + try { + options?.onStderr?.(chunk); + } catch (error) { + const cause = toError(error); + callbackError = new ExecutionError("callback_error", cause.message, cause); + onAbort(); + } + }); + + child.on("error", (error) => { + settle(err(new ExecutionError("spawn_error", error.message, error))); + }); + + child.on("close", (code) => { + if (callbackError) { + settle(err(callbackError)); + return; + } + if (timedOut) { + settle(err(new ExecutionError("timeout", `timeout:${options?.timeout}`))); + return; + } + if (options?.abortSignal?.aborted) { + settle(err(new ExecutionError("aborted", "aborted"))); + return; + } + settle(ok({ stdout, stderr, exitCode: code ?? 0 })); + }); + }); + } + + async readTextFile(path: string, abortSignal?: AbortSignal): Promise> { + const resolved = resolvePath(this.cwd, path); + const aborted = abortResult(abortSignal, resolved); + if (aborted) return aborted; + try { + return ok(await readFile(resolved, { encoding: "utf8", signal: abortSignal })); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async readTextLines( + path: string, + options?: { maxLines?: number; abortSignal?: AbortSignal }, + ): Promise> { + const resolved = resolvePath(this.cwd, path); + const aborted = abortResult(options?.abortSignal, resolved); + if (aborted) return aborted; + if (options?.maxLines !== undefined && options.maxLines <= 0) return ok([]); + let stream: ReturnType | undefined; + let lineReader: ReturnType | undefined; + try { + stream = createReadStream(resolved, { encoding: "utf8", signal: options?.abortSignal }); + lineReader = createInterface({ input: stream, crlfDelay: Infinity }); + const lines: string[] = []; + for await (const line of lineReader) { + const loopAbort = abortResult(options?.abortSignal, resolved); + if (loopAbort) return loopAbort; + lines.push(line); + if (options?.maxLines !== undefined && lines.length >= options.maxLines) break; + } + const afterReadAbort = abortResult(options?.abortSignal, resolved); + if (afterReadAbort) return afterReadAbort; + return ok(lines); + } catch (error) { + return err(toFileError(error, resolved)); + } finally { + lineReader?.close(); + stream?.destroy(); + } + } + + async readBinaryFile(path: string, abortSignal?: AbortSignal): Promise> { + const resolved = resolvePath(this.cwd, path); + const aborted = abortResult(abortSignal, resolved); + if (aborted) return aborted; + try { + return ok(await readFile(resolved, { signal: abortSignal })); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async writeFile( + path: string, + content: string | Uint8Array, + abortSignal?: AbortSignal, + ): Promise> { + const resolved = resolvePath(this.cwd, path); + const aborted = abortResult(abortSignal, resolved); + if (aborted) return aborted; + try { + await mkdir(resolve(resolved, ".."), { recursive: true }); + const afterMkdirAbort = abortResult(abortSignal, resolved); + if (afterMkdirAbort) return afterMkdirAbort; + await writeFile(resolved, content, { signal: abortSignal }); + return ok(undefined); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async appendFile(path: string, content: string | Uint8Array): Promise> { + const resolved = resolvePath(this.cwd, path); + try { + await mkdir(resolve(resolved, ".."), { recursive: true }); + await appendFile(resolved, content); + return ok(undefined); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async fileInfo(path: string): Promise> { + const resolved = resolvePath(this.cwd, path); + try { + return fileInfoFromStats(resolved, await lstat(resolved)); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async listDir(path: string, abortSignal?: AbortSignal): Promise> { + const resolved = resolvePath(this.cwd, path); + const aborted = abortResult(abortSignal, resolved); + if (aborted) return aborted; + try { + const entries = await readdir(resolved, { withFileTypes: true }); + const infos: FileInfo[] = []; + for (const entry of entries) { + const loopAbort = abortResult(abortSignal, resolved); + if (loopAbort) return loopAbort; + const entryPath = resolve(resolved, entry.name); + try { + const info = fileInfoFromStats(entryPath, await lstat(entryPath)); + if (info.ok) infos.push(info.value); + } catch (error) { + return err(toFileError(error, entryPath)); + } + } + return ok(infos); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async canonicalPath(path: string): Promise> { + const resolved = resolvePath(this.cwd, path); + try { + return ok(await realpath(resolved)); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async exists(path: string): Promise> { + const result = await this.fileInfo(path); + if (result.ok) return ok(true); + if (result.error.code === "not_found") return ok(false); + return err(result.error); + } + + async createDir(path: string, options?: { recursive?: boolean }): Promise> { + const resolved = resolvePath(this.cwd, path); + try { + await mkdir(resolved, { recursive: options?.recursive ?? true }); + return ok(undefined); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async remove(path: string, options?: { recursive?: boolean; force?: boolean }): Promise> { + const resolved = resolvePath(this.cwd, path); + try { + await rm(resolved, { recursive: options?.recursive ?? false, force: options?.force ?? false }); + return ok(undefined); + } catch (error) { + return err(toFileError(error, resolved)); + } + } + + async createTempDir(prefix: string = "tmp-"): Promise> { + try { + return ok(await mkdtemp(join(tmpdir(), prefix))); + } catch (error) { + return err(toFileError(error)); + } + } + + async createTempFile(options?: { prefix?: string; suffix?: string }): Promise> { + const dir = await this.createTempDir("tmp-"); + if (!dir.ok) return dir; + const filePath = join(dir.value, `${options?.prefix ?? ""}${randomUUID()}${options?.suffix ?? ""}`); + try { + await writeFile(filePath, ""); + return ok(filePath); + } catch (error) { + return err(toFileError(error, filePath)); + } + } + + async cleanup(): Promise { + // nothing to clean up for the local node implementation + } +} diff --git a/cactus-code/packages/agent/src/harness/messages.ts b/cactus-code/packages/agent/src/harness/messages.ts new file mode 100644 index 000000000..36ce96a1e --- /dev/null +++ b/cactus-code/packages/agent/src/harness/messages.ts @@ -0,0 +1,164 @@ +import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "../types.ts"; + +export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary: + + +`; + +export const COMPACTION_SUMMARY_SUFFIX = ` +`; + +export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from: + + +`; + +export const BRANCH_SUMMARY_SUFFIX = ``; + +export interface BashExecutionMessage { + role: "bashExecution"; + command: string; + output: string; + exitCode: number | undefined; + cancelled: boolean; + truncated: boolean; + fullOutputPath?: string; + timestamp: number; + excludeFromContext?: boolean; +} + +export interface CustomMessage { + role: "custom"; + customType: string; + content: string | (TextContent | ImageContent)[]; + display: boolean; + details?: T; + timestamp: number; +} + +export interface BranchSummaryMessage { + role: "branchSummary"; + summary: string; + fromId: string; + timestamp: number; +} + +export interface CompactionSummaryMessage { + role: "compactionSummary"; + summary: string; + tokensBefore: number; + timestamp: number; +} + +declare module "../types.ts" { + interface CustomAgentMessages { + bashExecution: BashExecutionMessage; + custom: CustomMessage; + branchSummary: BranchSummaryMessage; + compactionSummary: CompactionSummaryMessage; + } +} + +export function bashExecutionToText(msg: BashExecutionMessage): string { + let text = `Ran \`${msg.command}\`\n`; + if (msg.output) { + text += `\`\`\`\n${msg.output}\n\`\`\``; + } else { + text += "(no output)"; + } + if (msg.cancelled) { + text += "\n\n(command cancelled)"; + } else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) { + text += `\n\nCommand exited with code ${msg.exitCode}`; + } + if (msg.truncated && msg.fullOutputPath) { + text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`; + } + return text; +} + +export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage { + return { + role: "branchSummary", + summary, + fromId, + timestamp: new Date(timestamp).getTime(), + }; +} + +export function createCompactionSummaryMessage( + summary: string, + tokensBefore: number, + timestamp: string, +): CompactionSummaryMessage { + return { + role: "compactionSummary", + summary, + tokensBefore, + timestamp: new Date(timestamp).getTime(), + }; +} + +export function createCustomMessage( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details: unknown | undefined, + timestamp: string, +): CustomMessage { + return { + role: "custom", + customType, + content, + display, + details, + timestamp: new Date(timestamp).getTime(), + }; +} + +export function convertToLlm(messages: AgentMessage[]): Message[] { + return messages + .map((m): Message | undefined => { + switch (m.role) { + case "bashExecution": + if (m.excludeFromContext) { + return undefined; + } + return { + role: "user", + content: [{ type: "text", text: bashExecutionToText(m) }], + timestamp: m.timestamp, + }; + case "custom": { + const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; + return { + role: "user", + content, + timestamp: m.timestamp, + }; + } + case "branchSummary": + return { + role: "user", + content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }], + timestamp: m.timestamp, + }; + case "compactionSummary": + return { + role: "user", + content: [ + { type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX }, + ], + timestamp: m.timestamp, + }; + case "user": + case "assistant": + case "toolResult": + return m; + default: + return undefined; + } + }) + .filter((m): m is Message => m !== undefined); +} diff --git a/cactus-code/packages/agent/src/harness/prompt-templates.ts b/cactus-code/packages/agent/src/harness/prompt-templates.ts new file mode 100644 index 000000000..6c1b6be0b --- /dev/null +++ b/cactus-code/packages/agent/src/harness/prompt-templates.ts @@ -0,0 +1,267 @@ +import { parse } from "yaml"; +import { type ExecutionEnv, type FileInfo, type PromptTemplate, type Result, toError } from "./types.ts"; + +export type PromptTemplateDiagnosticCode = "file_info_failed" | "list_failed" | "read_failed" | "parse_failed"; + +/** Warning produced while loading prompt templates. */ +export interface PromptTemplateDiagnostic { + /** Diagnostic severity. Currently only warnings are emitted. */ + type: "warning"; + /** Stable diagnostic code. */ + code: PromptTemplateDiagnosticCode; + /** Human-readable diagnostic message. */ + message: string; + /** Path associated with the diagnostic. */ + path: string; +} + +interface PromptTemplateFrontmatter { + description?: string; + "argument-hint"?: string; + [key: string]: unknown; +} + +/** + * Load prompt templates from one or more paths. + * + * Directory inputs load direct `.md` children non-recursively. File inputs load explicit `.md` files. Missing paths and + * non-markdown files are skipped. Read and parse failures are returned as diagnostics. + */ +export async function loadPromptTemplates( + env: ExecutionEnv, + paths: string | string[], +): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> { + const promptTemplates: PromptTemplate[] = []; + const diagnostics: PromptTemplateDiagnostic[] = []; + for (const path of Array.isArray(paths) ? paths : [paths]) { + const infoResult = await env.fileInfo(path); + if (!infoResult.ok) { + if (infoResult.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: infoResult.error.message, + path, + }); + } + continue; + } + const info = infoResult.value; + const kind = await resolveKind(env, info, diagnostics); + if (kind === "directory") { + const result = await loadTemplatesFromDir(env, info.path); + promptTemplates.push(...result.promptTemplates); + diagnostics.push(...result.diagnostics); + } else if (kind === "file" && info.name.endsWith(".md")) { + const result = await loadTemplateFromFile(env, info.path); + if (result.promptTemplate) promptTemplates.push(result.promptTemplate); + diagnostics.push(...result.diagnostics); + } + } + return { promptTemplates, diagnostics }; +} + +/** + * Load prompt templates from source-tagged paths. + * + * Source values are preserved exactly and attached to every loaded prompt template and diagnostic. The agent package does + * not interpret source values; applications define their own provenance shape. + */ +export async function loadSourcedPromptTemplates( + env: ExecutionEnv, + inputs: Array<{ path: string; source: TSource }>, + mapPromptTemplate?: (promptTemplate: PromptTemplate, source: TSource) => TPromptTemplate, +): Promise<{ + promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }>; + diagnostics: Array; +}> { + const promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }> = []; + const diagnostics: Array = []; + for (const input of inputs) { + const result = await loadPromptTemplates(env, input.path); + for (const promptTemplate of result.promptTemplates) { + promptTemplates.push({ + promptTemplate: mapPromptTemplate + ? mapPromptTemplate(promptTemplate, input.source) + : (promptTemplate as TPromptTemplate), + source: input.source, + }); + } + for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source }); + } + return { promptTemplates, diagnostics }; +} + +async function loadTemplatesFromDir( + env: ExecutionEnv, + dir: string, +): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> { + const promptTemplates: PromptTemplate[] = []; + const diagnostics: PromptTemplateDiagnostic[] = []; + const entriesResult = await env.listDir(dir); + if (!entriesResult.ok) { + diagnostics.push({ + type: "warning", + code: "list_failed", + message: entriesResult.error.message, + path: dir, + }); + return { promptTemplates, diagnostics }; + } + const entries = entriesResult.value; + + for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) { + const kind = await resolveKind(env, entry, diagnostics); + if (kind !== "file" || !entry.name.endsWith(".md")) continue; + const result = await loadTemplateFromFile(env, entry.path); + if (result.promptTemplate) promptTemplates.push(result.promptTemplate); + diagnostics.push(...result.diagnostics); + } + return { promptTemplates, diagnostics }; +} + +async function loadTemplateFromFile( + env: ExecutionEnv, + filePath: string, +): Promise<{ promptTemplate: PromptTemplate | null; diagnostics: PromptTemplateDiagnostic[] }> { + const diagnostics: PromptTemplateDiagnostic[] = []; + const rawContent = await env.readTextFile(filePath); + if (!rawContent.ok) { + diagnostics.push({ + type: "warning", + code: "read_failed", + message: rawContent.error.message, + path: filePath, + }); + return { promptTemplate: null, diagnostics }; + } + + const parsed = parseFrontmatter(rawContent.value); + if (!parsed.ok) { + diagnostics.push({ + type: "warning", + code: "parse_failed", + message: parsed.error.message, + path: filePath, + }); + return { promptTemplate: null, diagnostics }; + } + + const { frontmatter, body } = parsed.value; + const firstLine = body.split("\n").find((line) => line.trim()); + let description = typeof frontmatter.description === "string" ? frontmatter.description : ""; + if (!description && firstLine) { + description = firstLine.slice(0, 60); + if (firstLine.length > 60) description += "..."; + } + return { + promptTemplate: { + name: basenameEnvPath(filePath).replace(/\.md$/i, ""), + description, + content: body, + }, + diagnostics, + }; +} + +async function resolveKind( + env: ExecutionEnv, + info: FileInfo, + diagnostics: PromptTemplateDiagnostic[], +): Promise<"file" | "directory" | undefined> { + if (info.kind === "file" || info.kind === "directory") return info.kind; + const canonicalPath = await env.canonicalPath(info.path); + if (!canonicalPath.ok) { + if (canonicalPath.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: canonicalPath.error.message, + path: info.path, + }); + } + return undefined; + } + const target = await env.fileInfo(canonicalPath.value); + if (!target.ok) { + if (target.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: target.error.message, + path: info.path, + }); + } + return undefined; + } + return target.value.kind === "file" || target.value.kind === "directory" ? target.value.kind : undefined; +} + +function parseFrontmatter>( + content: string, +): Result<{ frontmatter: T; body: string }, Error> { + try { + const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + if (!normalized.startsWith("---")) return { ok: true, value: { frontmatter: {} as T, body: normalized } }; + const endIndex = normalized.indexOf("\n---", 3); + if (endIndex === -1) return { ok: true, value: { frontmatter: {} as T, body: normalized } }; + const yamlString = normalized.slice(4, endIndex); + const body = normalized.slice(endIndex + 4).trim(); + return { ok: true, value: { frontmatter: (parse(yamlString) ?? {}) as T, body } }; + } catch (error) { + return { ok: false, error: toError(error) }; + } +} + +function basenameEnvPath(path: string): string { + const normalized = path.replace(/\/+$/, ""); + const slashIndex = normalized.lastIndexOf("/"); + return slashIndex === -1 ? normalized : normalized.slice(slashIndex + 1); +} + +/** Parse an argument string using simple shell-style single and double quotes. */ +export function parseCommandArgs(argsString: string): string[] { + const args: string[] = []; + let current = ""; + let inQuote: string | null = null; + + for (let i = 0; i < argsString.length; i++) { + const char = argsString[i]!; + if (inQuote) { + if (char === inQuote) inQuote = null; + else current += char; + } else if (char === '"' || char === "'") { + inQuote = char; + } else if (char === " " || char === "\t") { + if (current) { + args.push(current); + current = ""; + } + } else { + current += char; + } + } + if (current) args.push(current); + return args; +} + +/** Substitute prompt template placeholders (`$1`, `$@`, `$ARGUMENTS`, `${@:N}`, `${@:N:L}`) with command arguments. */ +export function substituteArgs(content: string, args: string[]): string { + let result = content; + result = result.replace(/\$(\d+)/g, (_, num: string) => args[parseInt(num, 10) - 1] ?? ""); + result = result.replace(/\$\{@:(\d+)(?::(\d+))?\}/g, (_, startStr: string, lengthStr?: string) => { + let start = parseInt(startStr, 10) - 1; + if (start < 0) start = 0; + if (lengthStr) return args.slice(start, start + parseInt(lengthStr, 10)).join(" "); + return args.slice(start).join(" "); + }); + const allArgs = args.join(" "); + result = result.replace(/\$ARGUMENTS/g, allArgs); + result = result.replace(/\$@/g, allArgs); + return result; +} + +/** Format a prompt template invocation with positional arguments. */ +export function formatPromptTemplateInvocation(template: PromptTemplate, args: string[] = []): string { + return substituteArgs(template.content, args); +} diff --git a/cactus-code/packages/agent/src/harness/session/jsonl-repo.ts b/cactus-code/packages/agent/src/harness/session/jsonl-repo.ts new file mode 100644 index 000000000..1a08d4d35 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/jsonl-repo.ts @@ -0,0 +1,177 @@ +import type { + FileSystem, + JsonlSessionCreateOptions, + JsonlSessionListOptions, + JsonlSessionMetadata, + JsonlSessionRepoApi, + Session, +} from "../types.ts"; +import { SessionError, toError } from "../types.ts"; +import { JsonlSessionStorage, loadJsonlSessionMetadata } from "./jsonl-storage.ts"; +import { + createSessionId, + createTimestamp, + getEntriesToFork, + getFileSystemResultOrThrow, + toSession, +} from "./repo-utils.ts"; + +type JsonlSessionRepoFileSystem = Pick< + FileSystem, + | "cwd" + | "absolutePath" + | "joinPath" + | "readTextFile" + | "readTextLines" + | "writeFile" + | "appendFile" + | "listDir" + | "exists" + | "createDir" + | "remove" +>; + +function encodeCwd(cwd: string): string { + return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`; +} + +export class JsonlSessionRepo implements JsonlSessionRepoApi { + private readonly fs: JsonlSessionRepoFileSystem; + private readonly sessionsRootInput: string; + private sessionsRoot: string | undefined; + + constructor(options: { fs: JsonlSessionRepoFileSystem; sessionsRoot: string }) { + this.fs = options.fs; + this.sessionsRootInput = options.sessionsRoot; + } + + private async getSessionsRoot(): Promise { + if (!this.sessionsRoot) { + this.sessionsRoot = getFileSystemResultOrThrow( + await this.fs.absolutePath(this.sessionsRootInput), + `Failed to resolve sessions root ${this.sessionsRootInput}`, + ); + } + return this.sessionsRoot; + } + + private async getSessionDir(cwd: string): Promise { + return getFileSystemResultOrThrow( + await this.fs.joinPath([await this.getSessionsRoot(), encodeCwd(cwd)]), + `Failed to resolve session directory for ${cwd}`, + ); + } + + private async createSessionFilePath(cwd: string, sessionId: string, timestamp: string): Promise { + return getFileSystemResultOrThrow( + await this.fs.joinPath([ + await this.getSessionDir(cwd), + `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`, + ]), + `Failed to resolve session file path for ${sessionId}`, + ); + } + + async create(options: JsonlSessionCreateOptions): Promise> { + const id = options.id ?? createSessionId(); + const createdAt = createTimestamp(); + const sessionDir = await this.getSessionDir(options.cwd); + getFileSystemResultOrThrow( + await this.fs.createDir(sessionDir, { recursive: true }), + `Failed to create session directory ${sessionDir}`, + ); + const filePath = await this.createSessionFilePath(options.cwd, id, createdAt); + const storage = await JsonlSessionStorage.create(this.fs, filePath, { + cwd: options.cwd, + sessionId: id, + parentSessionPath: options.parentSessionPath, + }); + return toSession(storage); + } + + async open(metadata: JsonlSessionMetadata): Promise> { + if ( + !getFileSystemResultOrThrow(await this.fs.exists(metadata.path), `Failed to check session ${metadata.path}`) + ) { + throw new SessionError("not_found", `Session not found: ${metadata.path}`); + } + const storage = await JsonlSessionStorage.open(this.fs, metadata.path); + return toSession(storage); + } + + async list(options: JsonlSessionListOptions = {}): Promise { + const dirs = options.cwd ? [await this.getSessionDir(options.cwd)] : await this.listSessionDirs(); + const sessions: JsonlSessionMetadata[] = []; + for (const dir of dirs) { + if (!getFileSystemResultOrThrow(await this.fs.exists(dir), `Failed to check session directory ${dir}`)) { + continue; + } + const files = getFileSystemResultOrThrow( + await this.fs.listDir(dir), + `Failed to list sessions in ${dir}`, + ).filter((file) => file.kind !== "directory" && file.name.endsWith(".jsonl")); + for (const file of files) { + try { + sessions.push(await loadJsonlSessionMetadata(this.fs, file.path)); + } catch (error) { + const cause = toError(error); + if (!(cause instanceof SessionError) || cause.code !== "invalid_session") throw cause; + } + } + } + sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()); + return sessions; + } + + async delete(metadata: JsonlSessionMetadata): Promise { + getFileSystemResultOrThrow( + await this.fs.remove(metadata.path, { force: true }), + `Failed to delete session ${metadata.path}`, + ); + } + + async fork( + sourceMetadata: JsonlSessionMetadata, + options: JsonlSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + const forkedEntries = await getEntriesToFork(source.getStorage(), options); + const id = options.id ?? createSessionId(); + const createdAt = createTimestamp(); + const sessionDir = await this.getSessionDir(options.cwd); + getFileSystemResultOrThrow( + await this.fs.createDir(sessionDir, { recursive: true }), + `Failed to create session directory ${sessionDir}`, + ); + const storage = await JsonlSessionStorage.create( + this.fs, + await this.createSessionFilePath(options.cwd, id, createdAt), + { + cwd: options.cwd, + sessionId: id, + parentSessionPath: options.parentSessionPath ?? sourceMetadata.path, + }, + ); + for (const entry of forkedEntries) { + await storage.appendEntry(entry); + } + return toSession(storage); + } + + private async listSessionDirs(): Promise { + const sessionsRoot = await this.getSessionsRoot(); + if ( + !getFileSystemResultOrThrow( + await this.fs.exists(sessionsRoot), + `Failed to check sessions root ${sessionsRoot}`, + ) + ) { + return []; + } + const entries = getFileSystemResultOrThrow( + await this.fs.listDir(sessionsRoot), + `Failed to list sessions root ${sessionsRoot}`, + ); + return entries.filter((entry) => entry.kind === "directory").map((entry) => entry.path); + } +} diff --git a/cactus-code/packages/agent/src/harness/session/jsonl-storage.ts b/cactus-code/packages/agent/src/harness/session/jsonl-storage.ts new file mode 100644 index 000000000..54ad9eb03 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/jsonl-storage.ts @@ -0,0 +1,293 @@ +import type { FileSystem, JsonlSessionMetadata, LeafEntry, SessionStorage, SessionTreeEntry } from "../types.ts"; +import { SessionError, toError } from "../types.ts"; +import { getFileSystemResultOrThrow } from "./repo-utils.ts"; +import { uuidv7 } from "./uuid.ts"; + +type JsonlSessionStorageFileSystem = Pick; + +interface SessionHeader { + type: "session"; + version: 3; + id: string; + timestamp: string; + cwd: string; + parentSession?: string; +} + +function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) { + labelsById.set(entry.targetId, label); + } else { + labelsById.delete(entry.targetId); + } +} + +function buildLabelsById(entries: SessionTreeEntry[]): Map { + const labelsById = new Map(); + for (const entry of entries) { + updateLabelCache(labelsById, entry); + } + return labelsById; +} + +function generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + const id = uuidv7().slice(0, 8); + if (!byId.has(id)) return id; + } + return uuidv7(); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function invalidSession(filePath: string, message: string, cause?: Error): SessionError { + return new SessionError("invalid_session", `Invalid JSONL session file ${filePath}: ${message}`, cause); +} + +function invalidEntry(filePath: string, lineNumber: number, message: string, cause?: Error): SessionError { + return new SessionError( + "invalid_entry", + `Invalid JSONL session file ${filePath}: line ${lineNumber} ${message}`, + cause, + ); +} + +function parseHeaderLine(line: string, filePath: string): SessionHeader { + let parsed: unknown; + try { + parsed = JSON.parse(line); + } catch (error) { + throw invalidSession(filePath, "first line is not a valid session header", toError(error)); + } + if (!isRecord(parsed)) throw invalidSession(filePath, "first line is not a valid session header"); + if (parsed.type !== "session") throw invalidSession(filePath, "first line is not a valid session header"); + if (parsed.version !== 3) throw invalidSession(filePath, "unsupported session version"); + if (typeof parsed.id !== "string" || !parsed.id) throw invalidSession(filePath, "session header is missing id"); + if (typeof parsed.timestamp !== "string" || !parsed.timestamp) { + throw invalidSession(filePath, "session header is missing timestamp"); + } + if (typeof parsed.cwd !== "string" || !parsed.cwd) throw invalidSession(filePath, "session header is missing cwd"); + if (parsed.parentSession !== undefined && typeof parsed.parentSession !== "string") { + throw invalidSession(filePath, "session header parentSession must be a string"); + } + return { + type: "session", + version: 3, + id: parsed.id, + timestamp: parsed.timestamp, + cwd: parsed.cwd, + parentSession: parsed.parentSession, + }; +} + +function parseEntryLine(line: string, filePath: string, lineNumber: number): SessionTreeEntry { + let parsed: unknown; + try { + parsed = JSON.parse(line); + } catch (error) { + throw invalidEntry(filePath, lineNumber, "is not valid JSON", toError(error)); + } + if (!isRecord(parsed)) throw invalidEntry(filePath, lineNumber, "is not a valid session entry"); + if (typeof parsed.type !== "string") throw invalidEntry(filePath, lineNumber, "is missing entry type"); + if (typeof parsed.id !== "string" || !parsed.id) throw invalidEntry(filePath, lineNumber, "is missing entry id"); + if (parsed.parentId !== null && typeof parsed.parentId !== "string") { + throw invalidEntry(filePath, lineNumber, "has invalid parentId"); + } + if (typeof parsed.timestamp !== "string" || !parsed.timestamp) { + throw invalidEntry(filePath, lineNumber, "is missing timestamp"); + } + if (parsed.type === "leaf" && parsed.targetId !== null && typeof parsed.targetId !== "string") { + throw invalidEntry(filePath, lineNumber, "has invalid targetId"); + } + return parsed as unknown as SessionTreeEntry; +} + +function leafIdAfterEntry(entry: SessionTreeEntry): string | null { + return entry.type === "leaf" ? entry.targetId : entry.id; +} + +function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata { + return { + id: header.id, + createdAt: header.timestamp, + cwd: header.cwd, + path, + parentSessionPath: header.parentSession, + }; +} + +export async function loadJsonlSessionMetadata( + fs: JsonlSessionStorageFileSystem, + filePath: string, +): Promise { + const lines = getFileSystemResultOrThrow( + await fs.readTextLines(filePath, { maxLines: 1 }), + `Failed to read session header ${filePath}`, + ); + const line = lines[0]; + if (line?.trim()) return headerToSessionMetadata(parseHeaderLine(line, filePath), filePath); + throw invalidSession(filePath, "missing session header"); +} + +async function loadJsonlStorage( + fs: JsonlSessionStorageFileSystem, + filePath: string, +): Promise<{ + header: SessionHeader; + entries: SessionTreeEntry[]; + leafId: string | null; +}> { + const content = getFileSystemResultOrThrow(await fs.readTextFile(filePath), `Failed to read session ${filePath}`); + const lines = content.split("\n").filter((line) => line.trim()); + if (lines.length === 0) { + throw invalidSession(filePath, "missing session header"); + } + + const header = parseHeaderLine(lines[0]!, filePath); + const entries: SessionTreeEntry[] = []; + let leafId: string | null = null; + for (let i = 1; i < lines.length; i++) { + const entry = parseEntryLine(lines[i]!, filePath, i + 1); + entries.push(entry); + leafId = leafIdAfterEntry(entry); + } + return { header, entries, leafId }; +} + +export class JsonlSessionStorage implements SessionStorage { + private readonly fs: JsonlSessionStorageFileSystem; + private readonly filePath: string; + private readonly metadata: JsonlSessionMetadata; + private entries: SessionTreeEntry[]; + private byId: Map; + private labelsById: Map; + private currentLeafId: string | null; + + private constructor( + fs: JsonlSessionStorageFileSystem, + filePath: string, + header: SessionHeader, + entries: SessionTreeEntry[], + leafId: string | null, + ) { + this.fs = fs; + this.filePath = filePath; + this.metadata = headerToSessionMetadata(header, this.filePath); + this.entries = entries; + this.byId = new Map(entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(entries); + this.currentLeafId = leafId; + } + + static async open(fs: JsonlSessionStorageFileSystem, filePath: string): Promise { + const loaded = await loadJsonlStorage(fs, filePath); + return new JsonlSessionStorage(fs, filePath, loaded.header, loaded.entries, loaded.leafId); + } + + static async create( + fs: JsonlSessionStorageFileSystem, + filePath: string, + options: { + cwd: string; + sessionId: string; + parentSessionPath?: string; + }, + ): Promise { + const header: SessionHeader = { + type: "session", + version: 3, + id: options.sessionId, + timestamp: new Date().toISOString(), + cwd: options.cwd, + parentSession: options.parentSessionPath, + }; + getFileSystemResultOrThrow( + await fs.writeFile(filePath, `${JSON.stringify(header)}\n`), + `Failed to create session ${filePath}`, + ); + return new JsonlSessionStorage(fs, filePath, header, [], null); + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + if (this.currentLeafId !== null && !this.byId.has(this.currentLeafId)) { + throw new SessionError("invalid_session", `Entry ${this.currentLeafId} not found`); + } + return this.currentLeafId; + } + + async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !this.byId.has(leafId)) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + const entry: LeafEntry = { + type: "leaf", + id: generateEntryId(this.byId), + parentId: this.currentLeafId, + timestamp: new Date().toISOString(), + targetId: leafId, + }; + getFileSystemResultOrThrow( + await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`), + `Failed to append session leaf ${entry.id}`, + ); + this.entries.push(entry); + this.byId.set(entry.id, entry); + this.currentLeafId = leafId; + } + + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + + async appendEntry(entry: SessionTreeEntry): Promise { + getFileSystemResultOrThrow( + await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`), + `Failed to append session entry ${entry.id}`, + ); + this.entries.push(entry); + this.byId.set(entry.id, entry); + updateLabelCache(this.labelsById, entry); + this.currentLeafId = leafIdAfterEntry(entry); + } + + async getEntry(id: string): Promise { + return this.byId.get(id); + } + + async findEntries( + type: TType, + ): Promise>> { + return this.entries.filter((entry): entry is Extract => entry.type === type); + } + + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + + async getPathToRoot(leafId: string | null): Promise { + if (leafId === null) return []; + const path: SessionTreeEntry[] = []; + let current = this.byId.get(leafId); + if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); + while (current) { + path.unshift(current); + if (!current.parentId) break; + const parent = this.byId.get(current.parentId); + if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); + current = parent; + } + return path; + } + + async getEntries(): Promise { + return [...this.entries]; + } +} diff --git a/cactus-code/packages/agent/src/harness/session/memory-repo.ts b/cactus-code/packages/agent/src/harness/session/memory-repo.ts new file mode 100644 index 000000000..145bc2b7c --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/memory-repo.ts @@ -0,0 +1,50 @@ +import { type Session, SessionError, type SessionMetadata, type SessionRepo } from "../types.ts"; +import { InMemorySessionStorage } from "./memory-storage.ts"; +import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./repo-utils.ts"; + +export class InMemorySessionRepo implements SessionRepo { + private sessions = new Map>(); + + async create(options: { id?: string } = {}): Promise> { + const metadata: SessionMetadata = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + }; + const storage = new InMemorySessionStorage({ metadata }); + const session = toSession(storage); + this.sessions.set(metadata.id, session); + return session; + } + + async open(metadata: SessionMetadata): Promise> { + const session = this.sessions.get(metadata.id); + if (!session) { + throw new SessionError("not_found", `Session not found: ${metadata.id}`); + } + return session; + } + + async list(): Promise { + return Promise.all([...this.sessions.values()].map((session) => session.getMetadata())); + } + + async delete(metadata: SessionMetadata): Promise { + this.sessions.delete(metadata.id); + } + + async fork( + sourceMetadata: SessionMetadata, + options: { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + const forkedEntries = await getEntriesToFork(source.getStorage(), options); + const metadata: SessionMetadata = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + }; + const storage = new InMemorySessionStorage({ metadata, entries: forkedEntries }); + const session = toSession(storage); + this.sessions.set(metadata.id, session); + return session; + } +} diff --git a/cactus-code/packages/agent/src/harness/session/memory-storage.ts b/cactus-code/packages/agent/src/harness/session/memory-storage.ts new file mode 100644 index 000000000..c6428cac6 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/memory-storage.ts @@ -0,0 +1,131 @@ +import { + type LeafEntry, + SessionError, + type SessionMetadata, + type SessionStorage, + type SessionTreeEntry, +} from "../types.ts"; +import { uuidv7 } from "./uuid.ts"; + +function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) { + labelsById.set(entry.targetId, label); + } else { + labelsById.delete(entry.targetId); + } +} + +function buildLabelsById(entries: SessionTreeEntry[]): Map { + const labelsById = new Map(); + for (const entry of entries) { + updateLabelCache(labelsById, entry); + } + return labelsById; +} + +function generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + const id = uuidv7().slice(0, 8); + if (!byId.has(id)) return id; + } + return uuidv7(); +} + +function leafIdAfterEntry(entry: SessionTreeEntry): string | null { + return entry.type === "leaf" ? entry.targetId : entry.id; +} + +export class InMemorySessionStorage + implements SessionStorage +{ + private readonly metadata: TMetadata; + private entries: SessionTreeEntry[]; + private byId: Map; + private labelsById: Map; + private leafId: string | null; + + constructor(options?: { entries?: SessionTreeEntry[]; metadata?: TMetadata }) { + this.entries = options?.entries ? [...options.entries] : []; + this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(this.entries); + this.leafId = null; + for (const entry of this.entries) this.leafId = leafIdAfterEntry(entry); + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); + } + this.metadata = options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata); + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); + } + return this.leafId; + } + + async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !this.byId.has(leafId)) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + const entry: LeafEntry = { + type: "leaf", + id: generateEntryId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + targetId: leafId, + }; + this.entries.push(entry); + this.byId.set(entry.id, entry); + this.leafId = leafId; + } + + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + + async appendEntry(entry: SessionTreeEntry): Promise { + this.entries.push(entry); + this.byId.set(entry.id, entry); + updateLabelCache(this.labelsById, entry); + this.leafId = leafIdAfterEntry(entry); + } + + async getEntry(id: string): Promise { + return this.byId.get(id); + } + + async findEntries( + type: TType, + ): Promise>> { + return this.entries.filter((entry): entry is Extract => entry.type === type); + } + + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + + async getPathToRoot(leafId: string | null): Promise { + if (leafId === null) return []; + const path: SessionTreeEntry[] = []; + let current = this.byId.get(leafId); + if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`); + while (current) { + path.unshift(current); + if (!current.parentId) break; + const parent = this.byId.get(current.parentId); + if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); + current = parent; + } + return path; + } + + async getEntries(): Promise { + return [...this.entries]; + } +} diff --git a/cactus-code/packages/agent/src/harness/session/repo-utils.ts b/cactus-code/packages/agent/src/harness/session/repo-utils.ts new file mode 100644 index 000000000..a25b62def --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/repo-utils.ts @@ -0,0 +1,51 @@ +import { + type FileError, + type Result, + SessionError, + type SessionMetadata, + type SessionStorage, + type SessionTreeEntry, +} from "../types.ts"; +import { Session } from "./session.ts"; +import { uuidv7 } from "./uuid.ts"; + +export function createSessionId(): string { + return uuidv7(); +} + +export function createTimestamp(): string { + return new Date().toISOString(); +} + +export function toSession(storage: SessionStorage): Session { + return new Session(storage); +} + +export function getFileSystemResultOrThrow(result: Result, message: string): TValue { + if (!result.ok) { + const code = result.error.code === "not_found" ? "not_found" : "storage"; + throw new SessionError(code, `${message}: ${result.error.message}`, result.error); + } + return result.value; +} + +export async function getEntriesToFork( + storage: SessionStorage, + options: { entryId?: string; position?: "before" | "at" }, +): Promise { + if (!options.entryId) return storage.getEntries(); + const target = await storage.getEntry(options.entryId); + if (!target) { + throw new SessionError("invalid_fork_target", `Entry ${options.entryId} not found`); + } + let effectiveLeafId: string | null; + if ((options.position ?? "before") === "at") { + effectiveLeafId = target.id; + } else { + if (target.type !== "message" || target.message.role !== "user") { + throw new SessionError("invalid_fork_target", `Entry ${options.entryId} is not a user message`); + } + effectiveLeafId = target.parentId; + } + return storage.getPathToRoot(effectiveLeafId); +} diff --git a/cactus-code/packages/agent/src/harness/session/session.ts b/cactus-code/packages/agent/src/harness/session/session.ts new file mode 100644 index 000000000..ce369dfff --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/session.ts @@ -0,0 +1,267 @@ +import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "../../types.ts"; +import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.ts"; +import type { + ActiveToolsChangeEntry, + BranchSummaryEntry, + CompactionEntry, + CustomEntry, + CustomMessageEntry, + LabelEntry, + MessageEntry, + ModelChangeEntry, + SessionContext, + SessionInfoEntry, + SessionMetadata, + SessionStorage, + SessionTreeEntry, + ThinkingLevelChangeEntry, +} from "../types.ts"; +import { SessionError } from "../types.ts"; + +export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext { + let thinkingLevel = "off"; + let model: { provider: string; modelId: string } | null = null; + let activeToolNames: string[] | null = null; + let compaction: CompactionEntry | null = null; + + for (const entry of pathEntries) { + if (entry.type === "thinking_level_change") { + thinkingLevel = entry.thinkingLevel; + } else if (entry.type === "model_change") { + model = { provider: entry.provider, modelId: entry.modelId }; + } else if (entry.type === "message" && entry.message.role === "assistant") { + model = { provider: entry.message.provider, modelId: entry.message.model }; + } else if (entry.type === "active_tools_change") { + activeToolNames = [...entry.activeToolNames]; + } else if (entry.type === "compaction") { + compaction = entry; + } + } + + const messages: AgentMessage[] = []; + const appendMessage = (entry: SessionTreeEntry) => { + if (entry.type === "message") { + messages.push(entry.message as AgentMessage); + } else if (entry.type === "custom_message") { + messages.push( + createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + entry.display, + entry.details, + entry.timestamp, + ), + ); + } else if (entry.type === "branch_summary" && entry.summary) { + messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)); + } + }; + + if (compaction) { + messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp)); + const compactionIdx = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id); + let foundFirstKept = false; + for (let i = 0; i < compactionIdx; i++) { + const entry = pathEntries[i]!; + if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; + if (foundFirstKept) appendMessage(entry); + } + for (let i = compactionIdx + 1; i < pathEntries.length; i++) { + appendMessage(pathEntries[i]!); + } + } else { + for (const entry of pathEntries) { + appendMessage(entry); + } + } + + return { messages, thinkingLevel, model, activeToolNames }; +} + +export class Session { + private storage: SessionStorage; + + constructor(storage: SessionStorage) { + this.storage = storage; + } + + getMetadata(): Promise { + return this.storage.getMetadata(); + } + + getStorage(): SessionStorage { + return this.storage; + } + + getLeafId(): Promise { + return this.storage.getLeafId(); + } + + getEntry(id: string): Promise { + return this.storage.getEntry(id); + } + + getEntries(): Promise { + return this.storage.getEntries(); + } + + async getBranch(fromId?: string): Promise { + const leafId = fromId ?? (await this.storage.getLeafId()); + return this.storage.getPathToRoot(leafId); + } + + async buildContext(): Promise { + return buildSessionContext(await this.getBranch()); + } + + getLabel(id: string): Promise { + return this.storage.getLabel(id); + } + + async getSessionName(): Promise { + const entries = await this.storage.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; + } + + private async appendTypedEntry(entry: TEntry): Promise { + await this.storage.appendEntry(entry); + return entry.id; + } + + async appendMessage(message: AgentMessage): Promise { + return this.appendTypedEntry({ + type: "message", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + message, + } satisfies MessageEntry); + } + + async appendThinkingLevelChange(thinkingLevel: string): Promise { + return this.appendTypedEntry({ + type: "thinking_level_change", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + thinkingLevel, + } satisfies ThinkingLevelChangeEntry); + } + + async appendModelChange(provider: string, modelId: string): Promise { + return this.appendTypedEntry({ + type: "model_change", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + provider, + modelId, + } satisfies ModelChangeEntry); + } + + async appendActiveToolsChange(activeToolNames: string[]): Promise { + return this.appendTypedEntry({ + type: "active_tools_change", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + activeToolNames: [...activeToolNames], + } satisfies ActiveToolsChangeEntry); + } + + async appendCompaction( + summary: string, + firstKeptEntryId: string, + tokensBefore: number, + details?: T, + fromHook?: boolean, + ): Promise { + return this.appendTypedEntry({ + type: "compaction", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + summary, + firstKeptEntryId, + tokensBefore, + details, + fromHook, + } satisfies CompactionEntry); + } + + async appendCustomEntry(customType: string, data?: unknown): Promise { + return this.appendTypedEntry({ + type: "custom", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + customType, + data, + } satisfies CustomEntry); + } + + async appendCustomMessageEntry( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details?: T, + ): Promise { + return this.appendTypedEntry({ + type: "custom_message", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + customType, + content, + display, + details, + } satisfies CustomMessageEntry); + } + + async appendLabel(targetId: string, label: string | undefined): Promise { + if (!(await this.storage.getEntry(targetId))) { + throw new SessionError("not_found", `Entry ${targetId} not found`); + } + return this.appendTypedEntry({ + type: "label", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + targetId, + label, + } satisfies LabelEntry); + } + + async appendSessionName(name: string): Promise { + const sanitizedName = name.replace(/[\r\n]+/g, " ").trim(); + return this.appendTypedEntry({ + type: "session_info", + id: await this.storage.createEntryId(), + parentId: await this.storage.getLeafId(), + timestamp: new Date().toISOString(), + name: sanitizedName, + } satisfies SessionInfoEntry); + } + + async moveTo( + entryId: string | null, + summary?: { summary: string; details?: unknown; fromHook?: boolean }, + ): Promise { + if (entryId !== null && !(await this.storage.getEntry(entryId))) { + throw new SessionError("not_found", `Entry ${entryId} not found`); + } + await this.storage.setLeafId(entryId); + if (!summary) return undefined; + return this.appendTypedEntry({ + type: "branch_summary", + id: await this.storage.createEntryId(), + parentId: entryId, + timestamp: new Date().toISOString(), + fromId: entryId ?? "root", + summary: summary.summary, + details: summary.details, + fromHook: summary.fromHook, + } satisfies BranchSummaryEntry); + } +} diff --git a/cactus-code/packages/agent/src/harness/session/uuid.ts b/cactus-code/packages/agent/src/harness/session/uuid.ts new file mode 100644 index 000000000..0c4cef8e0 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/session/uuid.ts @@ -0,0 +1,54 @@ +let lastTimestamp = -Infinity; +let sequence = 0; + +function fillRandomBytes(bytes: Uint8Array): void { + const crypto = globalThis.crypto; + if (crypto?.getRandomValues) { + crypto.getRandomValues(bytes); + return; + } + for (let i = 0; i < bytes.length; i++) { + bytes[i] = Math.floor(Math.random() * 256); + } +} + +export function uuidv7(): string { + const random = new Uint8Array(16); + fillRandomBytes(random); + const timestamp = Date.now(); + + if (timestamp > lastTimestamp) { + sequence = random[6] * 0x1000000 + random[7] * 0x10000 + random[8] * 0x100 + random[9]; + lastTimestamp = timestamp; + } else { + sequence = (sequence + 1) >>> 0; + if (sequence === 0) { + lastTimestamp++; + } + } + + const bytes = new Uint8Array(16); + bytes[0] = (lastTimestamp / 0x10000000000) & 0xff; + bytes[1] = (lastTimestamp / 0x100000000) & 0xff; + bytes[2] = (lastTimestamp / 0x1000000) & 0xff; + bytes[3] = (lastTimestamp / 0x10000) & 0xff; + bytes[4] = (lastTimestamp / 0x100) & 0xff; + bytes[5] = lastTimestamp & 0xff; + bytes[6] = 0x70 | ((sequence >>> 28) & 0x0f); + bytes[7] = (sequence >>> 20) & 0xff; + bytes[8] = 0x80 | ((sequence >>> 14) & 0x3f); + bytes[9] = (sequence >>> 6) & 0xff; + bytes[10] = ((sequence & 0x3f) << 2) | (random[10] & 0x03); + bytes[11] = random[11]; + bytes[12] = random[12]; + bytes[13] = random[13]; + bytes[14] = random[14]; + bytes[15] = random[15]; + + return formatUuid(bytes); +} + +function formatUuid(bytes: Uint8Array): string { + const hex = Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")); + return `${hex.slice(0, 4).join("")}-${hex.slice(4, 6).join("")}-${hex.slice(6, 8).join("")}-${hex.slice(8, 10).join("")}-${hex.slice(10, 16).join("")}`; +} diff --git a/cactus-code/packages/agent/src/harness/skills.ts b/cactus-code/packages/agent/src/harness/skills.ts new file mode 100644 index 000000000..a1a4c0c4a --- /dev/null +++ b/cactus-code/packages/agent/src/harness/skills.ts @@ -0,0 +1,375 @@ +import ignore from "ignore"; +import { parse } from "yaml"; +import { type ExecutionEnv, type FileInfo, type Result, type Skill, toError } from "./types.ts"; + +const MAX_NAME_LENGTH = 64; +const MAX_DESCRIPTION_LENGTH = 1024; +const IGNORE_FILE_NAMES = [".gitignore", ".ignore", ".fdignore"]; + +type IgnoreMatcher = ReturnType; + +export type SkillDiagnosticCode = + | "file_info_failed" + | "list_failed" + | "read_failed" + | "parse_failed" + | "invalid_metadata"; + +/** Warning produced while loading skills. */ +export interface SkillDiagnostic { + /** Diagnostic severity. Currently only warnings are emitted. */ + type: "warning"; + /** Stable diagnostic code. */ + code: SkillDiagnosticCode; + /** Human-readable diagnostic message. */ + message: string; + /** Path associated with the diagnostic. */ + path: string; +} + +interface SkillFrontmatter { + name?: string; + description?: string; + "disable-model-invocation"?: boolean; + [key: string]: unknown; +} + +/** Format a skill invocation prompt, optionally appending additional user instructions. */ +export function formatSkillInvocation(skill: Skill, additionalInstructions?: string): string { + const skillBlock = `\nReferences are relative to ${dirnameEnvPath(skill.filePath)}.\n\n${skill.content}\n`; + return additionalInstructions ? `${skillBlock}\n\n${additionalInstructions}` : skillBlock; +} + +/** + * Load skills from one or more directories. + * + * Traverses directories recursively, loads `SKILL.md` files, loads direct root `.md` files as skills, honors ignore files, + * and returns diagnostics for invalid skill files. Missing input directories are skipped. + */ +export async function loadSkills( + env: ExecutionEnv, + dirs: string | string[], +): Promise<{ skills: Skill[]; diagnostics: SkillDiagnostic[] }> { + const skills: Skill[] = []; + const diagnostics: SkillDiagnostic[] = []; + for (const dir of Array.isArray(dirs) ? dirs : [dirs]) { + const rootInfoResult = await env.fileInfo(dir); + if (!rootInfoResult.ok) { + if (rootInfoResult.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: rootInfoResult.error.message, + path: dir, + }); + } + continue; + } + const rootInfo = rootInfoResult.value; + if ((await resolveKind(env, rootInfo, diagnostics)) !== "directory") continue; + const result = await loadSkillsFromDirInternal(env, rootInfo.path, true, ignore(), rootInfo.path); + skills.push(...result.skills); + diagnostics.push(...result.diagnostics); + } + return { skills, diagnostics }; +} + +/** + * Load skills from source-tagged directories. + * + * Source values are preserved exactly and attached to every loaded skill and diagnostic. The agent package does not + * interpret source values; applications define their own provenance shape. + */ +export async function loadSourcedSkills( + env: ExecutionEnv, + inputs: Array<{ path: string; source: TSource }>, + mapSkill?: (skill: Skill, source: TSource) => TSkill, +): Promise<{ + skills: Array<{ skill: TSkill; source: TSource }>; + diagnostics: Array; +}> { + const skills: Array<{ skill: TSkill; source: TSource }> = []; + const diagnostics: Array = []; + for (const input of inputs) { + const result = await loadSkills(env, input.path); + for (const skill of result.skills) { + skills.push({ skill: mapSkill ? mapSkill(skill, input.source) : (skill as TSkill), source: input.source }); + } + for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source }); + } + return { skills, diagnostics }; +} + +async function loadSkillsFromDirInternal( + env: ExecutionEnv, + dir: string, + includeRootFiles: boolean, + ignoreMatcher: IgnoreMatcher, + rootDir: string, +): Promise<{ skills: Skill[]; diagnostics: SkillDiagnostic[] }> { + const skills: Skill[] = []; + const diagnostics: SkillDiagnostic[] = []; + + const dirInfoResult = await env.fileInfo(dir); + if (!dirInfoResult.ok) { + if (dirInfoResult.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: dirInfoResult.error.message, + path: dir, + }); + } + return { skills, diagnostics }; + } + const dirInfo = dirInfoResult.value; + if ((await resolveKind(env, dirInfo, diagnostics)) !== "directory") return { skills, diagnostics }; + + await addIgnoreRules(env, ignoreMatcher, dir, rootDir, diagnostics); + + const entriesResult = await env.listDir(dir); + if (!entriesResult.ok) { + diagnostics.push({ type: "warning", code: "list_failed", message: entriesResult.error.message, path: dir }); + return { skills, diagnostics }; + } + const entries = entriesResult.value; + + for (const entry of entries) { + if (entry.name !== "SKILL.md") continue; + const fullPath = entry.path; + const kind = await resolveKind(env, entry, diagnostics); + if (kind !== "file") continue; + const relPath = relativeEnvPath(rootDir, fullPath); + if (ignoreMatcher.ignores(relPath)) continue; + + const result = await loadSkillFromFile(env, fullPath); + if (result.skill) skills.push(result.skill); + diagnostics.push(...result.diagnostics); + return { skills, diagnostics }; + } + + for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) { + if (entry.name.startsWith(".") || entry.name === "node_modules") continue; + const fullPath = entry.path; + const kind = await resolveKind(env, entry, diagnostics); + if (!kind) continue; + + const relPath = relativeEnvPath(rootDir, fullPath); + const ignorePath = kind === "directory" ? `${relPath}/` : relPath; + if (ignoreMatcher.ignores(ignorePath)) continue; + + if (kind === "directory") { + const result = await loadSkillsFromDirInternal(env, fullPath, false, ignoreMatcher, rootDir); + skills.push(...result.skills); + diagnostics.push(...result.diagnostics); + continue; + } + + if (kind !== "file" || !includeRootFiles || !entry.name.endsWith(".md")) continue; + const result = await loadSkillFromFile(env, fullPath); + if (result.skill) skills.push(result.skill); + diagnostics.push(...result.diagnostics); + } + + return { skills, diagnostics }; +} + +async function addIgnoreRules( + env: ExecutionEnv, + ig: IgnoreMatcher, + dir: string, + rootDir: string, + diagnostics: SkillDiagnostic[], +): Promise { + const relativeDir = relativeEnvPath(rootDir, dir); + const prefix = relativeDir ? `${relativeDir}/` : ""; + + for (const filename of IGNORE_FILE_NAMES) { + const ignorePath = joinEnvPath(dir, filename); + const info = await env.fileInfo(ignorePath); + if (!info.ok) { + if (info.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: info.error.message, + path: ignorePath, + }); + } + continue; + } + if (info.value.kind !== "file") continue; + const content = await env.readTextFile(ignorePath); + if (!content.ok) { + diagnostics.push({ type: "warning", code: "read_failed", message: content.error.message, path: ignorePath }); + continue; + } + const patterns = content.value + .split(/\r?\n/) + .map((line) => prefixIgnorePattern(line, prefix)) + .filter((line): line is string => Boolean(line)); + if (patterns.length > 0) ig.add(patterns); + } +} + +function prefixIgnorePattern(line: string, prefix: string): string | null { + const trimmed = line.trim(); + if (!trimmed) return null; + if (trimmed.startsWith("#") && !trimmed.startsWith("\\#")) return null; + + let pattern = line; + let negated = false; + if (pattern.startsWith("!")) { + negated = true; + pattern = pattern.slice(1); + } else if (pattern.startsWith("\\!")) { + pattern = pattern.slice(1); + } + if (pattern.startsWith("/")) pattern = pattern.slice(1); + const prefixed = prefix ? `${prefix}${pattern}` : pattern; + return negated ? `!${prefixed}` : prefixed; +} + +async function loadSkillFromFile( + env: ExecutionEnv, + filePath: string, +): Promise<{ skill: Skill | null; diagnostics: SkillDiagnostic[] }> { + const diagnostics: SkillDiagnostic[] = []; + const rawContent = await env.readTextFile(filePath); + if (!rawContent.ok) { + diagnostics.push({ type: "warning", code: "read_failed", message: rawContent.error.message, path: filePath }); + return { skill: null, diagnostics }; + } + + const parsed = parseFrontmatter(rawContent.value); + if (!parsed.ok) { + diagnostics.push({ type: "warning", code: "parse_failed", message: parsed.error.message, path: filePath }); + return { skill: null, diagnostics }; + } + + const { frontmatter, body } = parsed.value; + const skillDir = dirnameEnvPath(filePath); + const parentDirName = basenameEnvPath(skillDir); + const description = typeof frontmatter.description === "string" ? frontmatter.description : undefined; + + for (const error of validateDescription(description)) { + diagnostics.push({ type: "warning", code: "invalid_metadata", message: error, path: filePath }); + } + + const frontmatterName = typeof frontmatter.name === "string" ? frontmatter.name : undefined; + const name = frontmatterName || parentDirName; + for (const error of validateName(name, parentDirName)) { + diagnostics.push({ type: "warning", code: "invalid_metadata", message: error, path: filePath }); + } + + if (!description || description.trim() === "") { + return { skill: null, diagnostics }; + } + + return { + skill: { + name, + description, + content: body, + filePath, + disableModelInvocation: frontmatter["disable-model-invocation"] === true, + }, + diagnostics, + }; +} + +function validateName(name: string, parentDirName: string): string[] { + const errors: string[] = []; + if (name !== parentDirName) errors.push(`name "${name}" does not match parent directory "${parentDirName}"`); + if (name.length > MAX_NAME_LENGTH) errors.push(`name exceeds ${MAX_NAME_LENGTH} characters (${name.length})`); + if (!/^[a-z0-9-]+$/.test(name)) { + errors.push("name contains invalid characters (must be lowercase a-z, 0-9, hyphens only)"); + } + if (name.startsWith("-") || name.endsWith("-")) errors.push("name must not start or end with a hyphen"); + if (name.includes("--")) errors.push("name must not contain consecutive hyphens"); + return errors; +} + +function validateDescription(description: string | undefined): string[] { + const errors: string[] = []; + if (!description || description.trim() === "") { + errors.push("description is required"); + } else if (description.length > MAX_DESCRIPTION_LENGTH) { + errors.push(`description exceeds ${MAX_DESCRIPTION_LENGTH} characters (${description.length})`); + } + return errors; +} + +function parseFrontmatter>( + content: string, +): Result<{ frontmatter: T; body: string }, Error> { + try { + const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + if (!normalized.startsWith("---")) return { ok: true, value: { frontmatter: {} as T, body: normalized } }; + const endIndex = normalized.indexOf("\n---", 3); + if (endIndex === -1) return { ok: true, value: { frontmatter: {} as T, body: normalized } }; + const yamlString = normalized.slice(4, endIndex); + const body = normalized.slice(endIndex + 4).trim(); + return { ok: true, value: { frontmatter: (parse(yamlString) ?? {}) as T, body } }; + } catch (error) { + return { ok: false, error: toError(error) }; + } +} + +async function resolveKind( + env: ExecutionEnv, + info: FileInfo, + diagnostics: SkillDiagnostic[], +): Promise<"file" | "directory" | undefined> { + if (info.kind === "file" || info.kind === "directory") return info.kind; + const canonicalPath = await env.canonicalPath(info.path); + if (!canonicalPath.ok) { + if (canonicalPath.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: canonicalPath.error.message, + path: info.path, + }); + } + return undefined; + } + const target = await env.fileInfo(canonicalPath.value); + if (!target.ok) { + if (target.error.code !== "not_found") { + diagnostics.push({ + type: "warning", + code: "file_info_failed", + message: target.error.message, + path: info.path, + }); + } + return undefined; + } + return target.value.kind === "file" || target.value.kind === "directory" ? target.value.kind : undefined; +} + +function joinEnvPath(base: string, child: string): string { + return `${base.replace(/\/+$/, "")}/${child.replace(/^\/+/, "")}`; +} + +function dirnameEnvPath(path: string): string { + const normalized = path.replace(/\/+$/, ""); + const slashIndex = normalized.lastIndexOf("/"); + return slashIndex <= 0 ? "/" : normalized.slice(0, slashIndex); +} + +function basenameEnvPath(path: string): string { + const normalized = path.replace(/\/+$/, ""); + const slashIndex = normalized.lastIndexOf("/"); + return slashIndex === -1 ? normalized : normalized.slice(slashIndex + 1); +} + +function relativeEnvPath(root: string, path: string): string { + const normalizedRoot = root.replace(/\/+$/, ""); + const normalizedPath = path.replace(/\/+$/, ""); + if (normalizedPath === normalizedRoot) return ""; + return normalizedPath.startsWith(`${normalizedRoot}/`) + ? normalizedPath.slice(normalizedRoot.length + 1) + : normalizedPath.replace(/^\/+/, ""); +} diff --git a/cactus-code/packages/agent/src/harness/system-prompt.ts b/cactus-code/packages/agent/src/harness/system-prompt.ts new file mode 100644 index 000000000..ec121da2c --- /dev/null +++ b/cactus-code/packages/agent/src/harness/system-prompt.ts @@ -0,0 +1,34 @@ +import type { Skill } from "./types.ts"; + +export function formatSkillsForSystemPrompt(skills: Skill[]): string { + const visibleSkills = skills.filter((skill) => !skill.disableModelInvocation); + if (visibleSkills.length === 0) return ""; + + const lines = [ + "The following skills provide specialized instructions for specific tasks.", + "Read the full skill file when the task matches its description.", + "When a skill file references a relative path, resolve it against the skill directory (parent of SKILL.md / dirname of the path) and use that absolute path in tool commands.", + "", + "", + ]; + + for (const skill of visibleSkills) { + lines.push(" "); + lines.push(` ${escapeXml(skill.name)}`); + lines.push(` ${escapeXml(skill.description)}`); + lines.push(` ${escapeXml(skill.filePath)}`); + lines.push(" "); + } + + lines.push(""); + return lines.join("\n"); +} + +function escapeXml(value: string): string { + return value + .replace(/&/g, "&") + .replace(//g, ">") + .replace(/"/g, """) + .replace(/'/g, "'"); +} diff --git a/cactus-code/packages/agent/src/harness/types.ts b/cactus-code/packages/agent/src/harness/types.ts new file mode 100644 index 000000000..29048e617 --- /dev/null +++ b/cactus-code/packages/agent/src/harness/types.ts @@ -0,0 +1,836 @@ +import type { ImageContent, Model, Models, SimpleStreamOptions, TextContent, Transport } from "@earendil-works/pi-ai"; +import type { AgentEvent, AgentMessage, AgentTool, QueueMode, ThinkingLevel } from "../index.ts"; +import type { Session } from "./session/session.ts"; + +/** Result of a fallible operation. Expected failures are returned as `ok: false` instead of thrown. */ +export type Result = { ok: true; value: TValue } | { ok: false; error: TError }; + +/** Create a successful {@link Result}. */ +export function ok(value: TValue): Result { + return { ok: true, value }; +} + +/** Create a failed {@link Result}. */ +export function err(error: TError): Result { + return { ok: false, error }; +} + +/** Return the success value or throw the failure error. Intended for tests and explicit adapter boundaries. */ +export function getOrThrow(result: Result): TValue { + if (!result.ok) throw result.error; + return result.value; +} + +/** Return the success value or `undefined`. Only object values are allowed to avoid truthiness bugs with primitives. */ +export function getOrUndefined(result: Result): TValue | undefined { + return result.ok ? result.value : undefined; +} + +/** Normalize unknown thrown values into Error instances before using them as typed error causes. */ +export function toError(error: unknown): Error { + if (error instanceof Error) return error; + if (typeof error === "string") return new Error(error); + try { + return new Error(JSON.stringify(error)); + } catch { + return new Error(String(error)); + } +} + +/** + * Skill loaded from a `SKILL.md` file or provided by an application. + * + * `name`, `description`, and `filePath` are inserted into the system prompt in an XML-formatted block as suggested by agentskills.io. + * Use {@link formatSkillsForSystemPrompt} to generate the spec-compatible system prompt block. + */ +export interface Skill { + /** Stable skill name used for lookup and model-visible listings. */ + name: string; + /** Short model-visible description of when to use the skill. */ + description: string; + /** Full skill instructions. */ + content: string; + /** Absolute path to the skill file. Used for model-visible location and resolving relative references. */ + filePath: string; + /** Exclude this skill from model-visible skill lists while still allowing explicit application invocation. */ + disableModelInvocation?: boolean; +} + +/** Prompt template that can be formatted into a prompt for explicit invocation. */ +export interface PromptTemplate { + /** Stable template name used for lookup or application command routing. */ + name: string; + /** Optional description for command lists or autocomplete. */ + description?: string; + /** Template content. Argument placeholders are formatted by `formatPromptTemplateInvocation`. */ + content: string; +} + +/** Resources made available to explicit invocation methods and system-prompt callbacks. */ +export interface AgentHarnessResources< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> { + /** Prompt templates available for explicit invocation. */ + promptTemplates?: TPromptTemplate[]; + /** Skills available to the model and explicit skill invocation. */ + skills?: TSkill[]; +} + +/** Curated provider request options owned by the harness and snapshotted per turn. */ +export interface AgentHarnessStreamOptions { + /** Preferred transport forwarded to the stream function. */ + transport?: Transport; + /** Provider request timeout in milliseconds. */ + timeoutMs?: number; + /** Maximum provider retry attempts. */ + maxRetries?: number; + /** Optional cap for provider-requested retry delays. */ + maxRetryDelayMs?: number; + /** Additional request headers merged with auth and lifecycle headers. */ + headers?: Record; + /** Provider metadata forwarded with requests. */ + metadata?: SimpleStreamOptions["metadata"]; + /** Provider cache retention hint. */ + cacheRetention?: SimpleStreamOptions["cacheRetention"]; +} + +/** Per-request stream option patch returned by provider hooks. */ +export interface AgentHarnessStreamOptionsPatch + extends Omit, "headers" | "metadata"> { + /** Header patch. `undefined` values delete keys; explicit `headers: undefined` clears all headers. */ + headers?: Record; + /** Metadata patch. `undefined` values delete keys; explicit `metadata: undefined` clears all metadata. */ + metadata?: Record; +} + +/** Kind of filesystem object as addressed by a {@link FileSystem}. Symlinks are not followed automatically. */ +export type FileKind = "file" | "directory" | "symlink"; + +/** Stable, backend-independent file error codes returned by {@link FileSystem} file operations. */ +export type FileErrorCode = + | "aborted" + | "not_found" + | "permission_denied" + | "not_directory" + | "is_directory" + | "invalid" + | "not_supported" + | "unknown"; + +/** Error returned by {@link FileSystem} file operations. */ +export class FileError extends Error { + /** Backend-independent error code. */ + public code: FileErrorCode; + /** Absolute addressed path associated with the failure, when available. */ + public path?: string; + + constructor(code: FileErrorCode, message: string, path?: string, cause?: Error) { + super(message, cause === undefined ? undefined : { cause }); + this.name = "FileError"; + this.code = code; + this.path = path; + } +} + +/** Stable, backend-independent execution error codes returned by {@link ExecutionEnv.exec}. */ +export type ExecutionErrorCode = + | "aborted" + | "timeout" + | "shell_unavailable" + | "spawn_error" + | "callback_error" + | "unknown"; + +/** Error returned by {@link ExecutionEnv.exec}. */ +export class ExecutionError extends Error { + /** Backend-independent error code. */ + public code: ExecutionErrorCode; + + constructor(code: ExecutionErrorCode, message: string, cause?: Error) { + super(message, cause === undefined ? undefined : { cause }); + this.name = "ExecutionError"; + this.code = code; + } +} + +/** Stable compaction error codes returned by compaction helpers. */ +export type CompactionErrorCode = "aborted" | "summarization_failed" | "invalid_session" | "unknown"; + +/** Error returned by compaction helpers. */ +export class CompactionError extends Error { + /** Backend-independent error code. */ + public code: CompactionErrorCode; + + constructor(code: CompactionErrorCode, message: string, cause?: Error) { + super(message, cause === undefined ? undefined : { cause }); + this.name = "CompactionError"; + this.code = code; + } +} + +/** Stable branch-summary error codes returned by branch summarization helpers. */ +export type BranchSummaryErrorCode = "aborted" | "summarization_failed" | "invalid_session"; + +/** Error returned by branch summarization helpers. */ +export class BranchSummaryError extends Error { + /** Backend-independent error code. */ + public code: BranchSummaryErrorCode; + + constructor(code: BranchSummaryErrorCode, message: string, cause?: Error) { + super(message, cause === undefined ? undefined : { cause }); + this.name = "BranchSummaryError"; + this.code = code; + } +} + +export type SessionErrorCode = + | "not_found" + | "invalid_session" + | "invalid_entry" + | "invalid_fork_target" + | "storage" + | "unknown"; + +/** Error thrown by session storage, repositories, and session tree operations. */ +export class SessionError extends Error { + /** Session subsystem error code. */ + public code: SessionErrorCode; + + constructor(code: SessionErrorCode, message: string, cause?: Error) { + super(message, cause === undefined ? undefined : { cause }); + this.name = "SessionError"; + this.code = code; + } +} + +export type AgentHarnessErrorCode = + | "busy" + | "invalid_state" + | "invalid_argument" + | "session" + | "hook" + | "auth" + | "compaction" + | "branch_summary" + | "unknown"; + +/** Public AgentHarness failure with a stable top-level classification. */ +export class AgentHarnessError extends Error { + public code: AgentHarnessErrorCode; + + constructor(code: AgentHarnessErrorCode, message: string, cause?: Error) { + super(message, cause === undefined ? undefined : { cause }); + this.name = "AgentHarnessError"; + this.code = code; + } +} + +/** Metadata for one filesystem object in a {@link FileSystem}. */ +export interface FileInfo { + /** Basename of {@link path}. */ + name: string; + /** Absolute, syntactically normalized addressed path in the execution environment. Symlinks are not followed. */ + path: string; + /** Object kind. Symlink targets are not followed; use {@link FileSystem.canonicalPath} explicitly. */ + kind: FileKind; + /** Size in bytes for the addressed filesystem object. */ + size: number; + /** Modification time as milliseconds since Unix epoch. */ + mtimeMs: number; +} + +/** + * Filesystem capability used by the harness. + * + * Paths passed to methods may be absolute or relative to {@link cwd}. Paths returned by file operations are addressed paths + * in the filesystem namespace, but are not canonicalized through symlinks unless returned by {@link canonicalPath}. + * + * Operation methods must never throw or reject. All filesystem failures, including unexpected backend failures, must be + * encoded in the returned {@link Result}. Implementations must preserve this invariant. + */ +export interface FileSystem { + /** Current working directory for relative paths. */ + cwd: string; + + /** Return an absolute addressed path without requiring it to exist and without resolving symlinks. */ + absolutePath(path: string, abortSignal?: AbortSignal): Promise>; + /** Join path segments in the filesystem namespace without requiring the result to exist. */ + joinPath(parts: string[], abortSignal?: AbortSignal): Promise>; + /** Read a UTF-8 text file. */ + readTextFile(path: string, abortSignal?: AbortSignal): Promise>; + /** Read UTF-8 text lines. Implementations should stop once `maxLines` lines have been read. */ + readTextLines( + path: string, + options?: { maxLines?: number; abortSignal?: AbortSignal }, + ): Promise>; + /** Read a binary file. */ + readBinaryFile(path: string, abortSignal?: AbortSignal): Promise>; + /** Create or overwrite a file, creating parent directories when supported. */ + writeFile(path: string, content: string | Uint8Array, abortSignal?: AbortSignal): Promise>; + /** Create or append to a file, creating parent directories when supported. */ + appendFile(path: string, content: string | Uint8Array, abortSignal?: AbortSignal): Promise>; + /** Return metadata for the addressed path without following symlinks. */ + fileInfo(path: string, abortSignal?: AbortSignal): Promise>; + /** List direct children of a directory without following symlinks. */ + listDir(path: string, abortSignal?: AbortSignal): Promise>; + /** Return the canonical path for an existing path, resolving symlinks where supported. */ + canonicalPath(path: string, abortSignal?: AbortSignal): Promise>; + /** Return false for missing paths. Other errors, such as permission failures, return a {@link FileError}. */ + exists(path: string, abortSignal?: AbortSignal): Promise>; + /** Create a directory. Defaults: `recursive: true`, no abort signal. */ + createDir( + path: string, + options?: { recursive?: boolean; abortSignal?: AbortSignal }, + ): Promise>; + /** Remove a file or directory. Defaults: `recursive: false`, `force: false`, no abort signal. */ + remove( + path: string, + options?: { recursive?: boolean; force?: boolean; abortSignal?: AbortSignal }, + ): Promise>; + /** Create a temporary directory and return its absolute path. Defaults: `prefix: "tmp-"`, no abort signal. */ + createTempDir(prefix?: string, abortSignal?: AbortSignal): Promise>; + /** Create a temporary file and return its absolute path. Defaults: `prefix: ""`, `suffix: ""`, no abort signal. */ + createTempFile(options?: { + prefix?: string; + suffix?: string; + abortSignal?: AbortSignal; + }): Promise>; + + /** Release filesystem resources. Must be best-effort and must not throw or reject. */ + cleanup(): Promise; +} + +/** Options for {@link Shell.exec}. */ +export interface ShellExecOptions { + /** Working directory for the command. Relative paths are resolved against {@link ExecutionEnv.cwd}. Defaults to {@link ExecutionEnv.cwd}. */ + cwd?: string; + /** Additional environment variables for the command. Values override the environment defaults. Defaults to no overrides. */ + env?: Record; + /** Timeout in seconds. Implementations should return a timeout error when the command exceeds this duration. Defaults to no timeout. */ + timeout?: number; + /** Abort signal used to terminate the command. Defaults to no abort signal. */ + abortSignal?: AbortSignal; + /** Called with stdout chunks as they are produced. */ + onStdout?: (chunk: string) => void; + /** Called with stderr chunks as they are produced. */ + onStderr?: (chunk: string) => void; +} + +/** Shell execution capability used by the harness. */ +export interface Shell { + /** Execute a shell command in {@link FileSystem.cwd} unless `options.cwd` is provided. */ + exec( + command: string, + options?: ShellExecOptions, + ): Promise>; + /** Release shell resources. Must be best-effort and must not throw or reject. */ + cleanup(): Promise; +} + +/** Filesystem and process execution environment used by the harness. */ +export interface ExecutionEnv extends FileSystem, Shell {} + +export interface SessionTreeEntryBase { + type: string; + id: string; + parentId: string | null; + timestamp: string; +} + +export interface MessageEntry extends SessionTreeEntryBase { + type: "message"; + message: AgentMessage; +} + +export interface ThinkingLevelChangeEntry extends SessionTreeEntryBase { + type: "thinking_level_change"; + thinkingLevel: string; +} + +export interface ModelChangeEntry extends SessionTreeEntryBase { + type: "model_change"; + provider: string; + modelId: string; +} + +export interface ActiveToolsChangeEntry extends SessionTreeEntryBase { + type: "active_tools_change"; + activeToolNames: string[]; +} + +export interface CompactionEntry extends SessionTreeEntryBase { + type: "compaction"; + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + details?: T; + fromHook?: boolean; +} + +export interface BranchSummaryEntry extends SessionTreeEntryBase { + type: "branch_summary"; + fromId: string; + summary: string; + details?: T; + fromHook?: boolean; +} + +export interface CustomEntry extends SessionTreeEntryBase { + type: "custom"; + customType: string; + data?: T; +} + +export interface CustomMessageEntry extends SessionTreeEntryBase { + type: "custom_message"; + customType: string; + content: string | (TextContent | ImageContent)[]; + details?: T; + display: boolean; +} + +export interface LabelEntry extends SessionTreeEntryBase { + type: "label"; + targetId: string; + label: string | undefined; +} + +export interface SessionInfoEntry extends SessionTreeEntryBase { + type: "session_info"; // legacy name, kept for backwards compatibility + name?: string; +} + +export interface LeafEntry extends SessionTreeEntryBase { + type: "leaf"; + targetId: string | null; +} + +export type SessionTreeEntry = + | MessageEntry + | ThinkingLevelChangeEntry + | ModelChangeEntry + | ActiveToolsChangeEntry + | CompactionEntry + | BranchSummaryEntry + | CustomEntry + | CustomMessageEntry + | LabelEntry + | SessionInfoEntry + | LeafEntry; + +export interface SessionContext { + messages: AgentMessage[]; + thinkingLevel: string; + model: { provider: string; modelId: string } | null; + activeToolNames: string[] | null; +} + +export interface SessionMetadata { + id: string; + createdAt: string; +} + +export interface JsonlSessionMetadata extends SessionMetadata { + cwd: string; + path: string; + parentSessionPath?: string; +} + +export interface SessionStorage { + getMetadata(): Promise; + getLeafId(): Promise; + /** Persist a leaf entry that records the active session-tree leaf. */ + setLeafId(leafId: string | null): Promise; + createEntryId(): Promise; + appendEntry(entry: SessionTreeEntry): Promise; + getEntry(id: string): Promise; + findEntries( + type: TType, + ): Promise>>; + getLabel(id: string): Promise; + getPathToRoot(leafId: string | null): Promise; + getEntries(): Promise; +} + +export type { Session } from "./session/session.ts"; + +export interface SessionCreateOptions { + id?: string; +} + +export interface SessionForkOptions { + entryId?: string; + position?: "before" | "at"; + id?: string; +} + +export interface SessionRepo< + TMetadata extends SessionMetadata = SessionMetadata, + TCreateOptions extends SessionCreateOptions = SessionCreateOptions, + TListOptions = void, +> { + create(options: TCreateOptions): Promise>; + open(metadata: TMetadata): Promise>; + list(options?: TListOptions): Promise; + delete(metadata: TMetadata): Promise; + fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise>; +} + +export interface JsonlSessionCreateOptions extends SessionCreateOptions { + cwd: string; + parentSessionPath?: string; +} + +export interface JsonlSessionListOptions { + cwd?: string; +} + +export interface JsonlSessionRepoApi + extends SessionRepo {} + +export type AgentHarnessPhase = "idle" | "turn" | "compaction" | "branch_summary" | "retry"; + +export type PendingSessionWrite = SessionTreeEntry extends infer TEntry + ? TEntry extends SessionTreeEntry + ? Omit + : never + : never; + +export interface QueueUpdateEvent { + type: "queue_update"; + steer: AgentMessage[]; + followUp: AgentMessage[]; + nextTurn: AgentMessage[]; +} + +export interface SavePointEvent { + type: "save_point"; + hadPendingMutations: boolean; +} + +export interface AbortEvent { + type: "abort"; + clearedSteer: AgentMessage[]; + clearedFollowUp: AgentMessage[]; +} + +export interface SettledEvent { + type: "settled"; + nextTurnCount: number; +} + +export interface BeforeAgentStartEvent< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> { + type: "before_agent_start"; + prompt: string; + images?: ImageContent[]; + systemPrompt: string; + resources: AgentHarnessResources; +} + +export interface ContextEvent { + type: "context"; + messages: AgentMessage[]; +} + +export interface BeforeProviderRequestEvent { + type: "before_provider_request"; + model: Model; + sessionId: string; + streamOptions: AgentHarnessStreamOptions; +} + +export interface BeforeProviderPayloadEvent { + type: "before_provider_payload"; + model: Model; + payload: unknown; +} + +export interface AfterProviderResponseEvent { + type: "after_provider_response"; + status: number; + headers: Record; +} + +export interface ToolCallEvent { + type: "tool_call"; + toolCallId: string; + toolName: string; + input: Record; +} + +export interface ToolResultEvent { + type: "tool_result"; + toolCallId: string; + toolName: string; + input: Record; + content: Array; + details: unknown; + isError: boolean; +} + +export interface SessionBeforeCompactEvent { + type: "session_before_compact"; + preparation: CompactionPreparation; + branchEntries: SessionTreeEntry[]; + customInstructions?: string; + signal: AbortSignal; +} + +export interface SessionCompactEvent { + type: "session_compact"; + compactionEntry: CompactionEntry; + fromHook: boolean; +} + +export interface SessionBeforeTreeEvent { + type: "session_before_tree"; + preparation: TreePreparation; + signal: AbortSignal; +} + +export interface SessionTreeEvent { + type: "session_tree"; + newLeafId: string | null; + oldLeafId: string | null; + summaryEntry?: BranchSummaryEntry; + fromHook?: boolean; +} + +export interface ModelUpdateEvent { + type: "model_update"; + model: Model; + previousModel: Model | undefined; + source: "set" | "restore"; +} + +export interface ThinkingLevelUpdateEvent { + type: "thinking_level_update"; + level: ThinkingLevel; + previousLevel: ThinkingLevel; +} + +export interface ToolsUpdateEvent { + type: "tools_update"; + toolNames: string[]; + previousToolNames: string[]; + activeToolNames: string[]; + previousActiveToolNames: string[]; + source: "set" | "restore"; +} + +export interface ResourcesUpdateEvent< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> { + type: "resources_update"; + resources: AgentHarnessResources; + previousResources: AgentHarnessResources; +} + +export type AgentHarnessOwnEvent< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, +> = + | QueueUpdateEvent + | SavePointEvent + | AbortEvent + | SettledEvent + | BeforeAgentStartEvent + | ContextEvent + | BeforeProviderRequestEvent + | BeforeProviderPayloadEvent + | AfterProviderResponseEvent + | ToolCallEvent + | ToolResultEvent + | SessionBeforeCompactEvent + | SessionCompactEvent + | SessionBeforeTreeEvent + | SessionTreeEvent + | ModelUpdateEvent + | ThinkingLevelUpdateEvent + | ResourcesUpdateEvent + | ToolsUpdateEvent; + +export type AgentHarnessEvent = + | AgentEvent + | AgentHarnessOwnEvent; + +export interface BeforeAgentStartResult { + messages?: AgentMessage[]; + systemPrompt?: string; +} + +export interface ContextResult { + messages: AgentMessage[]; +} + +export interface BeforeProviderRequestResult { + streamOptions?: AgentHarnessStreamOptionsPatch; +} + +export interface BeforeProviderPayloadResult { + payload: unknown; +} + +export interface ToolCallResult { + block?: boolean; + reason?: string; +} + +export interface ToolResultPatch { + content?: Array; + details?: unknown; + isError?: boolean; + terminate?: boolean; +} + +export interface SessionBeforeCompactResult { + cancel?: boolean; + compaction?: CompactResult; +} + +export interface SessionBeforeTreeResult { + cancel?: boolean; + summary?: { summary: string; details?: unknown }; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; +} + +export type AgentHarnessEventResultMap = { + before_agent_start: BeforeAgentStartResult | undefined; + context: ContextResult | undefined; + before_provider_request: BeforeProviderRequestResult | undefined; + before_provider_payload: BeforeProviderPayloadResult | undefined; + after_provider_response: undefined; + tool_call: ToolCallResult | undefined; + tool_result: ToolResultPatch | undefined; + session_before_compact: SessionBeforeCompactResult | undefined; + session_compact: undefined; + session_before_tree: SessionBeforeTreeResult | undefined; + session_tree: undefined; + model_update: undefined; + thinking_level_update: undefined; + resources_update: undefined; + tools_update: undefined; + queue_update: undefined; + save_point: undefined; + abort: undefined; + settled: undefined; +}; + +export interface AgentHarnessPromptOptions { + images?: ImageContent[]; +} + +export interface AbortResult { + clearedSteer: AgentMessage[]; + clearedFollowUp: AgentMessage[]; +} + +export interface CompactResult { + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + details?: unknown; +} + +export interface NavigateTreeResult { + cancelled: boolean; + editorText?: string; + summaryEntry?: BranchSummaryEntry; +} + +export interface CompactionSettings { + enabled: boolean; + reserveTokens: number; + keepRecentTokens: number; +} + +export interface CompactionPreparation { + firstKeptEntryId: string; + messagesToSummarize: AgentMessage[]; + turnPrefixMessages: AgentMessage[]; + isSplitTurn: boolean; + tokensBefore: number; + previousSummary?: string; + fileOps: FileOperations; + settings: CompactionSettings; +} + +export interface FileOperations { + read: Set; + written: Set; + edited: Set; +} + +export interface TreePreparation { + targetId: string; + oldLeafId: string | null; + commonAncestorId: string | null; + entriesToSummarize: SessionTreeEntry[]; + userWantsSummary: boolean; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; +} + +export interface GenerateBranchSummaryOptions { + model: Model; + apiKey: string; + headers?: Record; + signal: AbortSignal; + customInstructions?: string; + replaceInstructions?: boolean; + reserveTokens?: number; +} + +export interface BranchSummaryResult { + summary: string; + readFiles: string[]; + modifiedFiles: string[]; +} + +export interface AgentHarnessOptions< + TSkill extends Skill = Skill, + TPromptTemplate extends PromptTemplate = PromptTemplate, + TTool extends AgentTool = AgentTool, +> { + env: ExecutionEnv; + session: Session; + /** + * Provider collection used for all model requests (turn streaming, + * compaction, branch summarization). Auth resolves through the providers' + * auth. + */ + models: Models; + tools?: TTool[]; + /** + * Concrete resources available to explicit invocation methods and system-prompt callbacks. + * Applications own loading/reloading resources and should call `setResources()` with new values. + */ + resources?: AgentHarnessResources; + systemPrompt?: + | string + | ((context: { + env: ExecutionEnv; + session: Session; + model: Model; + thinkingLevel: ThinkingLevel; + activeTools: TTool[]; + resources: AgentHarnessResources; + }) => string | Promise); + /** Curated stream/provider request options. Snapshotted at turn start. */ + streamOptions?: AgentHarnessStreamOptions; + model: Model; + thinkingLevel?: ThinkingLevel; + activeToolNames?: string[]; + steeringMode?: QueueMode; + followUpMode?: QueueMode; +} + +export type { AgentHarness } from "./agent-harness.ts"; diff --git a/cactus-code/packages/agent/src/harness/utils/shell-output.ts b/cactus-code/packages/agent/src/harness/utils/shell-output.ts new file mode 100644 index 000000000..a98b51a5f --- /dev/null +++ b/cactus-code/packages/agent/src/harness/utils/shell-output.ts @@ -0,0 +1,135 @@ +import { type ExecutionEnv, ExecutionError, err, ok, type Result, type ShellExecOptions, toError } from "../types.ts"; +import { DEFAULT_MAX_BYTES, truncateTail } from "./truncate.ts"; + +export interface ShellCaptureOptions extends Omit { + onChunk?: (chunk: string) => void; +} + +export interface ShellCaptureResult { + output: string; + exitCode: number | undefined; + cancelled: boolean; + truncated: boolean; + fullOutputPath?: string; +} + +function toExecutionError(error: unknown): ExecutionError { + if (error instanceof ExecutionError) return error; + const cause = toError(error); + return new ExecutionError("unknown", cause.message, cause); +} + +export function sanitizeBinaryOutput(str: string): string { + return Array.from(str) + .filter((char) => { + const code = char.codePointAt(0); + if (code === undefined) return false; + if (code === 0x09 || code === 0x0a || code === 0x0d) return true; + if (code <= 0x1f) return false; + if (code >= 0xfff9 && code <= 0xfffb) return false; + return true; + }) + .join(""); +} + +export async function executeShellWithCapture( + env: ExecutionEnv, + command: string, + options?: ShellCaptureOptions, +): Promise> { + const outputChunks: string[] = []; + let outputBytes = 0; + const maxOutputBytes = DEFAULT_MAX_BYTES * 2; + const encoder = new TextEncoder(); + + let totalBytes = 0; + let fullOutputPath: string | undefined; + let writeChain: Promise> = Promise.resolve(ok(undefined)); + let captureError: ExecutionError | undefined; + + const appendFullOutput = (text: string): void => { + if (!fullOutputPath || captureError) return; + const path = fullOutputPath; + writeChain = writeChain.then(async (previous) => { + if (!previous.ok) return previous; + const appendResult = await env.appendFile(path, text, options?.abortSignal); + return appendResult.ok ? ok(undefined) : err(toExecutionError(appendResult.error)); + }); + }; + + const ensureFullOutputFile = (initialContent: string): void => { + if (fullOutputPath || captureError) return; + writeChain = writeChain.then(async (previous) => { + if (!previous.ok) return previous; + const tempFile = await env.createTempFile({ + prefix: "bash-", + suffix: ".log", + abortSignal: options?.abortSignal, + }); + if (!tempFile.ok) return err(toExecutionError(tempFile.error)); + fullOutputPath = tempFile.value; + const appendResult = await env.appendFile(tempFile.value, initialContent, options?.abortSignal); + return appendResult.ok ? ok(undefined) : err(toExecutionError(appendResult.error)); + }); + }; + + const onChunk = (chunk: string) => { + try { + totalBytes += encoder.encode(chunk).byteLength; + const text = sanitizeBinaryOutput(chunk).replace(/\r/g, ""); + if (totalBytes > DEFAULT_MAX_BYTES && !fullOutputPath) { + ensureFullOutputFile(outputChunks.join("") + text); + } else { + appendFullOutput(text); + } + outputChunks.push(text); + outputBytes += text.length; + while (outputBytes > maxOutputBytes && outputChunks.length > 1) { + const removed = outputChunks.shift()!; + outputBytes -= removed.length; + } + options?.onChunk?.(text); + } catch (error) { + captureError = toExecutionError(error); + } + }; + + try { + const result = await env.exec(command, { + ...(options ?? {}), + onStdout: onChunk, + onStderr: onChunk, + }); + const tailOutput = outputChunks.join(""); + const truncationResult = truncateTail(tailOutput); + if (truncationResult.truncated && !fullOutputPath) { + ensureFullOutputFile(tailOutput); + } + const writeResult = await writeChain; + if (!writeResult.ok) return err(writeResult.error); + if (captureError) return err(captureError); + + if (!result.ok) { + if (result.error.code === "aborted" || options?.abortSignal?.aborted) { + return ok({ + output: truncationResult.truncated ? truncationResult.content : tailOutput, + exitCode: undefined, + cancelled: true, + truncated: truncationResult.truncated, + fullOutputPath, + }); + } + return err(result.error); + } + const cancelled = options?.abortSignal?.aborted ?? false; + return ok({ + output: truncationResult.truncated ? truncationResult.content : tailOutput, + exitCode: cancelled ? undefined : result.value.exitCode, + cancelled, + truncated: truncationResult.truncated, + fullOutputPath, + }); + } catch (error) { + return err(toExecutionError(error)); + } +} diff --git a/cactus-code/packages/agent/src/harness/utils/truncate.ts b/cactus-code/packages/agent/src/harness/utils/truncate.ts new file mode 100644 index 000000000..2372a941b --- /dev/null +++ b/cactus-code/packages/agent/src/harness/utils/truncate.ts @@ -0,0 +1,344 @@ +/** + * Shared truncation utilities for tool outputs. + * + * Truncation is based on two independent limits - whichever is hit first wins: + * - Line limit (default: 2000 lines) + * - Byte limit (default: 50KB) + * + * Never returns partial lines (except bash tail truncation edge case). + */ + +export const DEFAULT_MAX_LINES = 2000; +export const DEFAULT_MAX_BYTES = 50 * 1024; // 50KB +export const GREP_MAX_LINE_LENGTH = 500; // Max chars per grep match line + +export interface TruncationResult { + /** The truncated content */ + content: string; + /** Whether truncation occurred */ + truncated: boolean; + /** Which limit was hit: "lines", "bytes", or null if not truncated */ + truncatedBy: "lines" | "bytes" | null; + /** Total number of lines in the original content */ + totalLines: number; + /** Total number of bytes in the original content */ + totalBytes: number; + /** Number of complete lines in the truncated output */ + outputLines: number; + /** Number of bytes in the truncated output */ + outputBytes: number; + /** Whether the last line was partially truncated (only for tail truncation edge case) */ + lastLinePartial: boolean; + /** Whether the first line exceeded the byte limit (for head truncation) */ + firstLineExceedsLimit: boolean; + /** The max lines limit that was applied */ + maxLines: number; + /** The max bytes limit that was applied */ + maxBytes: number; +} + +export interface TruncationOptions { + /** Maximum number of lines (default: 2000) */ + maxLines?: number; + /** Maximum number of bytes (default: 50KB) */ + maxBytes?: number; +} + +interface RuntimeBuffer { + byteLength(content: string, encoding: "utf8"): number; +} + +const runtimeBuffer = (globalThis as { Buffer?: RuntimeBuffer }).Buffer; +const nonAsciiPattern = /[^\x00-\x7f]/; + +function utf8ByteLength(content: string): number { + if (runtimeBuffer) return runtimeBuffer.byteLength(content, "utf8"); + + const firstNonAscii = content.search(nonAsciiPattern); + if (firstNonAscii === -1) return content.length; + + let bytes = firstNonAscii; + for (let i = firstNonAscii; i < content.length; i++) { + const code = content.charCodeAt(i); + if (code <= 0x7f) { + bytes += 1; + } else if (code <= 0x7ff) { + bytes += 2; + } else if (code >= 0xd800 && code <= 0xdbff && i + 1 < content.length) { + const next = content.charCodeAt(i + 1); + if (next >= 0xdc00 && next <= 0xdfff) { + bytes += 4; + i++; + } else { + bytes += 3; + } + } else { + bytes += 3; + } + } + return bytes; +} + +function replaceUnpairedSurrogates(content: string): string { + let output = ""; + for (let i = 0; i < content.length; i++) { + const code = content.charCodeAt(i); + if (code >= 0xd800 && code <= 0xdbff) { + if (i + 1 < content.length) { + const next = content.charCodeAt(i + 1); + if (next >= 0xdc00 && next <= 0xdfff) { + output += content[i] + content[i + 1]; + i++; + continue; + } + } + output += "�"; + } else if (code >= 0xdc00 && code <= 0xdfff) { + output += "�"; + } else { + output += content[i]; + } + } + return output; +} + +/** + * Format bytes as human-readable size. + */ +export function formatSize(bytes: number): string { + if (bytes < 1024) { + return `${bytes}B`; + } else if (bytes < 1024 * 1024) { + return `${(bytes / 1024).toFixed(1)}KB`; + } else { + return `${(bytes / (1024 * 1024)).toFixed(1)}MB`; + } +} + +/** + * Truncate content from the head (keep first N lines/bytes). + * Suitable for file reads where you want to see the beginning. + * + * Never returns partial lines. If first line exceeds byte limit, + * returns empty content with firstLineExceedsLimit=true. + */ +export function truncateHead(content: string, options: TruncationOptions = {}): TruncationResult { + const maxLines = options.maxLines ?? DEFAULT_MAX_LINES; + const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + + const totalBytes = utf8ByteLength(content); + const lines = content.split("\n"); + const totalLines = lines.length; + + // Check if no truncation needed + if (totalLines <= maxLines && totalBytes <= maxBytes) { + return { + content, + truncated: false, + truncatedBy: null, + totalLines, + totalBytes, + outputLines: totalLines, + outputBytes: totalBytes, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; + } + + // Check if first line alone exceeds byte limit + const firstLineBytes = utf8ByteLength(lines[0]); + if (firstLineBytes > maxBytes) { + return { + content: "", + truncated: true, + truncatedBy: "bytes", + totalLines, + totalBytes, + outputLines: 0, + outputBytes: 0, + lastLinePartial: false, + firstLineExceedsLimit: true, + maxLines, + maxBytes, + }; + } + + // Collect complete lines that fit + const outputLinesArr: string[] = []; + let outputBytesCount = 0; + let truncatedBy: "lines" | "bytes" = "lines"; + + for (let i = 0; i < lines.length && i < maxLines; i++) { + const line = lines[i]; + const lineBytes = utf8ByteLength(line) + (i > 0 ? 1 : 0); // +1 for newline + + if (outputBytesCount + lineBytes > maxBytes) { + truncatedBy = "bytes"; + break; + } + + outputLinesArr.push(line); + outputBytesCount += lineBytes; + } + + // If we exited due to line limit + if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) { + truncatedBy = "lines"; + } + + const outputContent = outputLinesArr.join("\n"); + const finalOutputBytes = utf8ByteLength(outputContent); + + return { + content: outputContent, + truncated: true, + truncatedBy, + totalLines, + totalBytes, + outputLines: outputLinesArr.length, + outputBytes: finalOutputBytes, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; +} + +/** + * Truncate content from the tail (keep last N lines/bytes). + * Suitable for bash output where you want to see the end (errors, final results). + * + * May return partial first line if the last line of original content exceeds byte limit. + */ +export function truncateTail(content: string, options: TruncationOptions = {}): TruncationResult { + const maxLines = options.maxLines ?? DEFAULT_MAX_LINES; + const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + + const totalBytes = utf8ByteLength(content); + const lines = content.split("\n"); + if (lines.length > 1 && lines[lines.length - 1] === "") lines.pop(); + const totalLines = lines.length; + + // Check if no truncation needed + if (totalLines <= maxLines && totalBytes <= maxBytes) { + return { + content, + truncated: false, + truncatedBy: null, + totalLines, + totalBytes, + outputLines: totalLines, + outputBytes: totalBytes, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; + } + + // Work backwards from the end + const outputLinesArr: string[] = []; + let outputBytesCount = 0; + let truncatedBy: "lines" | "bytes" = "lines"; + let lastLinePartial = false; + + for (let i = lines.length - 1; i >= 0 && outputLinesArr.length < maxLines; i--) { + const line = lines[i]; + const lineBytes = utf8ByteLength(line) + (outputLinesArr.length > 0 ? 1 : 0); // +1 for newline + + if (outputBytesCount + lineBytes > maxBytes) { + truncatedBy = "bytes"; + // Edge case: if we haven't added ANY lines yet and this line exceeds maxBytes, + // take the end of the line (partial) + if (outputLinesArr.length === 0) { + const truncatedLine = truncateStringToBytesFromEnd(line, maxBytes); + outputLinesArr.unshift(truncatedLine); + outputBytesCount = utf8ByteLength(truncatedLine); + lastLinePartial = true; + } + break; + } + + outputLinesArr.unshift(line); + outputBytesCount += lineBytes; + } + + // If we exited due to line limit + if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) { + truncatedBy = "lines"; + } + + const outputContent = outputLinesArr.join("\n"); + const finalOutputBytes = utf8ByteLength(outputContent); + + return { + content: outputContent, + truncated: true, + truncatedBy, + totalLines, + totalBytes, + outputLines: outputLinesArr.length, + outputBytes: finalOutputBytes, + lastLinePartial, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; +} + +/** + * Truncate a string to fit within a byte limit (from the end). + * Handles multi-byte UTF-8 characters correctly. + */ +function truncateStringToBytesFromEnd(str: string, maxBytes: number): string { + if (maxBytes <= 0) return ""; + + let outputBytes = 0; + let start = str.length; + let needsReplacement = false; + for (let i = str.length; i > 0; ) { + let characterStart = i - 1; + const code = str.charCodeAt(characterStart); + let characterBytes: number; + let unpairedSurrogate = false; + if (code >= 0xdc00 && code <= 0xdfff && characterStart > 0) { + const previous = str.charCodeAt(characterStart - 1); + if (previous >= 0xd800 && previous <= 0xdbff) { + characterStart--; + characterBytes = 4; + } else { + characterBytes = 3; + unpairedSurrogate = true; + } + } else if (code >= 0xd800 && code <= 0xdfff) { + characterBytes = 3; + unpairedSurrogate = true; + } else { + characterBytes = code <= 0x7f ? 1 : code <= 0x7ff ? 2 : 3; + } + if (outputBytes + characterBytes > maxBytes) break; + outputBytes += characterBytes; + start = characterStart; + needsReplacement ||= unpairedSurrogate; + i = characterStart; + } + + const output = str.slice(start); + return needsReplacement ? replaceUnpairedSurrogates(output) : output; +} + +/** + * Truncate a single line to max characters, adding [truncated] suffix. + * Used for grep match lines. + */ +export function truncateLine( + line: string, + maxChars: number = GREP_MAX_LINE_LENGTH, +): { text: string; wasTruncated: boolean } { + if (line.length <= maxChars) { + return { text: line, wasTruncated: false }; + } + return { text: `${line.slice(0, maxChars)}... [truncated]`, wasTruncated: true }; +} diff --git a/cactus-code/packages/agent/src/index.ts b/cactus-code/packages/agent/src/index.ts new file mode 100644 index 000000000..fd3c2b9fc --- /dev/null +++ b/cactus-code/packages/agent/src/index.ts @@ -0,0 +1,44 @@ +// Core Agent +export * from "./agent.ts"; +// Loop functions +export * from "./agent-loop.ts"; +export * from "./harness/agent-harness.ts"; +export { + type BranchPreparation, + type BranchSummaryDetails, + type CollectEntriesResult, + collectEntriesForBranchSummary, + generateBranchSummary, + prepareBranchEntries, +} from "./harness/compaction/branch-summarization.ts"; +export { + calculateContextTokens, + compact, + DEFAULT_COMPACTION_SETTINGS, + estimateContextTokens, + estimateTokens, + findCutPoint, + findTurnStartIndex, + generateSummary, + getLastAssistantUsage, + prepareCompaction, + serializeConversation, + shouldCompact, +} from "./harness/compaction/compaction.ts"; +export * from "./harness/messages.ts"; +export * from "./harness/prompt-templates.ts"; +export * from "./harness/session/jsonl-repo.ts"; +export * from "./harness/session/memory-repo.ts"; +export * from "./harness/session/repo-utils.ts"; +export * from "./harness/session/session.ts"; +export { uuidv7 } from "./harness/session/uuid.ts"; +export * from "./harness/skills.ts"; +export * from "./harness/system-prompt.ts"; +// Harness +export * from "./harness/types.ts"; +export * from "./harness/utils/shell-output.ts"; +export * from "./harness/utils/truncate.ts"; +// Proxy utilities +export * from "./proxy.ts"; +// Types +export * from "./types.ts"; diff --git a/cactus-code/packages/agent/src/node.ts b/cactus-code/packages/agent/src/node.ts new file mode 100644 index 000000000..b6f53a5e0 --- /dev/null +++ b/cactus-code/packages/agent/src/node.ts @@ -0,0 +1,2 @@ +export { NodeExecutionEnv } from "./harness/env/nodejs.ts"; +export * from "./index.ts"; diff --git a/cactus-code/packages/agent/src/proxy.ts b/cactus-code/packages/agent/src/proxy.ts new file mode 100644 index 000000000..5f0925c91 --- /dev/null +++ b/cactus-code/packages/agent/src/proxy.ts @@ -0,0 +1,367 @@ +/** + * Proxy stream function for apps that route LLM calls through a server. + * The server manages auth and proxies requests to LLM providers. + */ + +// Internal import for JSON parsing utility +import { + type AssistantMessage, + type AssistantMessageEvent, + type Context, + EventStream, + type Model, + parseStreamingJson, + type SimpleStreamOptions, + type StopReason, + type ToolCall, +} from "@earendil-works/pi-ai"; + +// Create stream class matching ProxyMessageEventStream +class ProxyMessageEventStream extends EventStream { + constructor() { + super( + (event) => event.type === "done" || event.type === "error", + (event) => { + if (event.type === "done") return event.message; + if (event.type === "error") return event.error; + throw new Error("Unexpected event type"); + }, + ); + } +} + +/** + * Proxy event types - server sends these with partial field stripped to reduce bandwidth. + */ +export type ProxyAssistantMessageEvent = + | { type: "start" } + | { type: "text_start"; contentIndex: number } + | { type: "text_delta"; contentIndex: number; delta: string } + | { type: "text_end"; contentIndex: number; contentSignature?: string } + | { type: "thinking_start"; contentIndex: number } + | { type: "thinking_delta"; contentIndex: number; delta: string } + | { type: "thinking_end"; contentIndex: number; contentSignature?: string } + | { type: "toolcall_start"; contentIndex: number; id: string; toolName: string } + | { type: "toolcall_delta"; contentIndex: number; delta: string } + | { type: "toolcall_end"; contentIndex: number } + | { + type: "done"; + reason: Extract; + usage: AssistantMessage["usage"]; + } + | { + type: "error"; + reason: Extract; + errorMessage?: string; + usage: AssistantMessage["usage"]; + }; + +type ProxySerializableStreamOptions = Pick< + SimpleStreamOptions, + | "temperature" + | "maxTokens" + | "reasoning" + | "cacheRetention" + | "sessionId" + | "headers" + | "metadata" + | "transport" + | "thinkingBudgets" + | "maxRetryDelayMs" +>; + +export interface ProxyStreamOptions extends ProxySerializableStreamOptions { + /** Local abort signal for the proxy request */ + signal?: AbortSignal; + /** Auth token for the proxy server */ + authToken: string; + /** Proxy server URL (e.g., "https://genai.example.com") */ + proxyUrl: string; +} + +/** + * Stream function that proxies through a server instead of calling LLM providers directly. + * The server strips the partial field from delta events to reduce bandwidth. + * We reconstruct the partial message client-side. + * + * Use this as the `streamFn` option when creating an Agent that needs to go through a proxy. + * + * @example + * ```typescript + * const agent = new Agent({ + * streamFn: (model, context, options) => + * streamProxy(model, context, { + * ...options, + * authToken: await getAuthToken(), + * proxyUrl: "https://genai.example.com", + * }), + * }); + * ``` + */ +function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializableStreamOptions { + return { + temperature: options.temperature, + maxTokens: options.maxTokens, + reasoning: options.reasoning, + cacheRetention: options.cacheRetention, + sessionId: options.sessionId, + headers: options.headers, + metadata: options.metadata, + transport: options.transport, + thinkingBudgets: options.thinkingBudgets, + maxRetryDelayMs: options.maxRetryDelayMs, + }; +} + +export function streamProxy(model: Model, context: Context, options: ProxyStreamOptions): ProxyMessageEventStream { + const stream = new ProxyMessageEventStream(); + + (async () => { + // Initialize the partial message that we'll build up from events + const partial: AssistantMessage = { + role: "assistant", + stopReason: "stop", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: Date.now(), + }; + + let reader: ReadableStreamDefaultReader | undefined; + + const abortHandler = () => { + if (reader) { + reader.cancel("Request aborted by user").catch(() => {}); + } + }; + + if (options.signal) { + options.signal.addEventListener("abort", abortHandler); + } + + try { + const response = await fetch(`${options.proxyUrl}/api/stream`, { + method: "POST", + headers: { + Authorization: `Bearer ${options.authToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model, + context, + options: buildProxyRequestOptions(options), + }), + signal: options.signal, + }); + + if (!response.ok) { + let errorMessage = `Proxy error: ${response.status} ${response.statusText}`; + try { + const errorData = (await response.json()) as { error?: string }; + if (errorData.error) { + errorMessage = `Proxy error: ${errorData.error}`; + } + } catch { + // Couldn't parse error response + } + throw new Error(errorMessage); + } + + reader = response.body!.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + if (options.signal?.aborted) { + throw new Error("Request aborted by user"); + } + + buffer += decoder.decode(value, { stream: true }); + const lines = buffer.split("\n"); + buffer = lines.pop() || ""; + + for (const line of lines) { + if (line.startsWith("data: ")) { + const data = line.slice(6).trim(); + if (data) { + const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent; + const event = processProxyEvent(proxyEvent, partial); + if (event) { + stream.push(event); + } + } + } + } + } + + if (options.signal?.aborted) { + throw new Error("Request aborted by user"); + } + + stream.end(); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + const reason = options.signal?.aborted ? "aborted" : "error"; + partial.stopReason = reason; + partial.errorMessage = errorMessage; + stream.push({ + type: "error", + reason, + error: partial, + }); + stream.end(); + } finally { + if (options.signal) { + options.signal.removeEventListener("abort", abortHandler); + } + } + })(); + + return stream; +} + +/** + * Process a proxy event and update the partial message. + */ +function processProxyEvent( + proxyEvent: ProxyAssistantMessageEvent, + partial: AssistantMessage, +): AssistantMessageEvent | undefined { + switch (proxyEvent.type) { + case "start": + return { type: "start", partial }; + + case "text_start": + partial.content[proxyEvent.contentIndex] = { type: "text", text: "" }; + return { type: "text_start", contentIndex: proxyEvent.contentIndex, partial }; + + case "text_delta": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "text") { + content.text += proxyEvent.delta; + return { + type: "text_delta", + contentIndex: proxyEvent.contentIndex, + delta: proxyEvent.delta, + partial, + }; + } + throw new Error("Received text_delta for non-text content"); + } + + case "text_end": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "text") { + content.textSignature = proxyEvent.contentSignature; + return { + type: "text_end", + contentIndex: proxyEvent.contentIndex, + content: content.text, + partial, + }; + } + throw new Error("Received text_end for non-text content"); + } + + case "thinking_start": + partial.content[proxyEvent.contentIndex] = { type: "thinking", thinking: "" }; + return { type: "thinking_start", contentIndex: proxyEvent.contentIndex, partial }; + + case "thinking_delta": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "thinking") { + content.thinking += proxyEvent.delta; + return { + type: "thinking_delta", + contentIndex: proxyEvent.contentIndex, + delta: proxyEvent.delta, + partial, + }; + } + throw new Error("Received thinking_delta for non-thinking content"); + } + + case "thinking_end": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "thinking") { + content.thinkingSignature = proxyEvent.contentSignature; + return { + type: "thinking_end", + contentIndex: proxyEvent.contentIndex, + content: content.thinking, + partial, + }; + } + throw new Error("Received thinking_end for non-thinking content"); + } + + case "toolcall_start": + partial.content[proxyEvent.contentIndex] = { + type: "toolCall", + id: proxyEvent.id, + name: proxyEvent.toolName, + arguments: {}, + partialJson: "", + } satisfies ToolCall & { partialJson: string } as ToolCall; + return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial }; + + case "toolcall_delta": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "toolCall") { + (content as any).partialJson += proxyEvent.delta; + content.arguments = parseStreamingJson((content as any).partialJson) || {}; + partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity + return { + type: "toolcall_delta", + contentIndex: proxyEvent.contentIndex, + delta: proxyEvent.delta, + partial, + }; + } + throw new Error("Received toolcall_delta for non-toolCall content"); + } + + case "toolcall_end": { + const content = partial.content[proxyEvent.contentIndex]; + if (content?.type === "toolCall") { + delete (content as any).partialJson; + return { + type: "toolcall_end", + contentIndex: proxyEvent.contentIndex, + toolCall: content, + partial, + }; + } + return undefined; + } + + case "done": + partial.stopReason = proxyEvent.reason; + partial.usage = proxyEvent.usage; + return { type: "done", reason: proxyEvent.reason, message: partial }; + + case "error": + partial.stopReason = proxyEvent.reason; + partial.errorMessage = proxyEvent.errorMessage; + partial.usage = proxyEvent.usage; + return { type: "error", reason: proxyEvent.reason, error: partial }; + + default: { + const _exhaustiveCheck: never = proxyEvent; + console.warn(`Unhandled proxy event type: ${(proxyEvent as any).type}`); + return undefined; + } + } +} diff --git a/cactus-code/packages/agent/src/types.ts b/cactus-code/packages/agent/src/types.ts new file mode 100644 index 000000000..abfa3de6f --- /dev/null +++ b/cactus-code/packages/agent/src/types.ts @@ -0,0 +1,428 @@ +import type { + Api, + AssistantMessage, + AssistantMessageEvent, + AssistantMessageEventStream, + Context, + ImageContent, + Message, + Model, + SimpleStreamOptions, + TextContent, + Tool, + ToolResultMessage, +} from "@earendil-works/pi-ai"; +import type { Static, TSchema } from "typebox"; + +/** + * Stream function used by the agent loop. `Models.streamSimple` satisfies + * this shape. + * + * Contract: + * - Must not throw or return a rejected promise for request/model/runtime failures. + * - Must return an AssistantMessageEventStream. + * - Failures must be encoded in the returned stream via protocol events and a + * final AssistantMessage with stopReason "error" or "aborted" and errorMessage. + */ +export type StreamFn = ( + model: Model, + context: Context, + options?: SimpleStreamOptions, +) => AssistantMessageEventStream | Promise; + +/** + * Configuration for how tool calls from a single assistant message are executed. + * + * - "sequential": each tool call is prepared, executed, and finalized before the next one starts. + * - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently. + * `tool_execution_end` is emitted in tool completion order after each tool is finalized, + * while tool-result message artifacts are emitted later in assistant source order. + */ +export type ToolExecutionMode = "sequential" | "parallel"; + +/** + * Controls how many queued user messages are injected when the agent loop reaches a queue drain point. + * + * - "all": drain and inject every queued message at that point. + * - "one-at-a-time": drain and inject only the oldest queued message, leaving the rest queued for later drain points. + */ +export type QueueMode = "all" | "one-at-a-time"; + +/** A single tool call content block emitted by an assistant message. */ +export type AgentToolCall = Extract; + +/** + * Result returned from `beforeToolCall`. + * + * Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead. + * `reason` becomes the text shown in that error result. If omitted, a default blocked message is used. + */ +export interface BeforeToolCallResult { + block?: boolean; + reason?: string; +} + +/** + * Partial override returned from `afterToolCall`. + * + * Merge semantics are field-by-field: + * - `content`: if provided, replaces the tool result content array in full + * - `details`: if provided, replaces the tool result details value in full + * - `isError`: if provided, replaces the tool result error flag + * - `terminate`: if provided, replaces the early-termination hint + * + * Omitted fields keep the original executed tool result values. + * There is no deep merge for `content` or `details`. + */ +export interface AfterToolCallResult { + content?: (TextContent | ImageContent)[]; + details?: unknown; + isError?: boolean; + /** + * Hint that the agent should stop after the current tool batch. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; +} + +/** Context passed to `beforeToolCall`. */ +export interface BeforeToolCallContext { + /** The assistant message that requested the tool call. */ + assistantMessage: AssistantMessage; + /** The raw tool call block from `assistantMessage.content`. */ + toolCall: AgentToolCall; + /** Validated tool arguments for the target tool schema. */ + args: unknown; + /** Current agent context at the time the tool call is prepared. */ + context: AgentContext; +} + +/** Context passed to `afterToolCall`. */ +export interface AfterToolCallContext { + /** The assistant message that requested the tool call. */ + assistantMessage: AssistantMessage; + /** The raw tool call block from `assistantMessage.content`. */ + toolCall: AgentToolCall; + /** Validated tool arguments for the target tool schema. */ + args: unknown; + /** The executed tool result before any `afterToolCall` overrides are applied. */ + result: AgentToolResult; + /** Whether the executed tool result is currently treated as an error. */ + isError: boolean; + /** Current agent context at the time the tool call is finalized. */ + context: AgentContext; +} + +/** Context passed to `shouldStopAfterTurn`. */ +export interface ShouldStopAfterTurnContext { + /** The assistant message that completed the turn. */ + message: AssistantMessage; + /** Tool result messages passed to the preceding `turn_end` event. */ + toolResults: ToolResultMessage[]; + /** Current agent context after the turn's assistant message and tool results have been appended. */ + context: AgentContext; + /** Messages that this loop invocation will return if it exits at this point. Prompt runs include the initial prompt messages; continuation runs do not include pre-existing context messages. */ + newMessages: AgentMessage[]; +} + +/** Replacement runtime state used by the agent loop before starting another provider request. */ +export interface AgentLoopTurnUpdate { + /** Context for the next provider request. */ + context?: AgentContext; + /** Model for the next provider request. */ + model?: Model; + /** Thinking level for the next provider request. */ + thinkingLevel?: ThinkingLevel; +} + +export interface PrepareNextTurnContext extends ShouldStopAfterTurnContext {} + +export interface AgentLoopConfig extends SimpleStreamOptions { + model: Model; + + /** + * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call. + * + * Each AgentMessage must be converted to a UserMessage, AssistantMessage, or ToolResultMessage + * that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications, + * status messages) should be filtered out. + * + * Contract: must not throw or reject. Return a safe fallback value instead. + * Throwing interrupts the low-level agent loop without producing a normal event sequence. + * + * @example + * ```typescript + * convertToLlm: (messages) => messages.flatMap(m => { + * if (m.role === "custom") { + * // Convert custom message to user message + * return [{ role: "user", content: m.content, timestamp: m.timestamp }]; + * } + * if (m.role === "notification") { + * // Filter out UI-only messages + * return []; + * } + * // Pass through standard LLM messages + * return [m]; + * }) + * ``` + */ + convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; + + /** + * Optional transform applied to the context before `convertToLlm`. + * + * Use this for operations that work at the AgentMessage level: + * - Context window management (pruning old messages) + * - Injecting context from external sources + * + * Contract: must not throw or reject. Return the original messages or another + * safe fallback value instead. + * + * @example + * ```typescript + * transformContext: async (messages) => { + * if (estimateTokens(messages) > MAX_TOKENS) { + * return pruneOldMessages(messages); + * } + * return messages; + * } + * ``` + */ + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; + + /** + * Resolves an API key dynamically for each LLM call. + * + * Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire + * during long-running tool execution phases. + * + * Contract: must not throw or reject. Return undefined when no key is available. + */ + getApiKey?: (provider: string) => Promise | string | undefined; + + /** + * Called after each turn fully completes and `turn_end` has been emitted. + * + * If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues, + * without starting another LLM call. The current assistant response and any tool executions finish normally. + * + * Use this to request a graceful stop after the current turn, e.g. before context gets too full. + * + * Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence. + */ + shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise; + + /** + * Called after `turn_end` and before the loop decides whether another provider request should start. + * Return replacement context/model/thinking state to affect the next turn in this run. + * Return undefined to keep using the current context/config. + */ + prepareNextTurn?: ( + context: PrepareNextTurnContext, + ) => AgentLoopTurnUpdate | undefined | Promise; + + /** + * Returns steering messages to inject into the conversation mid-run. + * + * Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first. + * If messages are returned, they are added to the context before the next LLM call. + * Tool calls from the current assistant message are not skipped. + * + * Use this for "steering" the agent while it's working. + * + * Contract: must not throw or reject. Return [] when no steering messages are available. + */ + getSteeringMessages?: () => Promise; + + /** + * Returns follow-up messages to process after the agent would otherwise stop. + * + * Called when the agent has no more tool calls and no steering messages. + * If messages are returned, they're added to the context and the agent + * continues with another turn. + * + * Use this for follow-up messages that should wait until the agent finishes. + * + * Contract: must not throw or reject. Return [] when no follow-up messages are available. + */ + getFollowUpMessages?: () => Promise; + + /** + * Tool execution mode. + * - "sequential": execute tool calls one by one + * - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently; + * emit `tool_execution_end` in tool completion order after each tool is finalized, + * then emit tool-result message artifacts later in assistant source order + * + * Default: "parallel" + */ + toolExecution?: ToolExecutionMode; + + /** + * Called before a tool is executed, after arguments have been validated. + * + * Return `{ block: true }` to prevent execution. The loop emits an error tool result instead. + * The hook receives the agent abort signal and is responsible for honoring it. + */ + beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; + + /** + * Called after a tool finishes executing, before `tool_execution_end` and tool-result message events are emitted. + * + * Return an `AfterToolCallResult` to override parts of the executed tool result: + * - `content` replaces the full content array + * - `details` replaces the full details payload + * - `isError` replaces the error flag + * - `terminate` replaces the early-termination hint + * + * Any omitted fields keep their original values. No deep merge is performed. + * The hook receives the agent abort signal and is responsible for honoring it. + */ + afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; +} + +/** + * Thinking/reasoning level for models that support it. + * Note: "xhigh" is only supported by selected model families. Use model thinking-level metadata + * from @earendil-works/pi-ai to detect support for a concrete model. + */ +export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh"; + +/** + * Extensible interface for custom app messages. + * Apps can extend via declaration merging: + * + * @example + * ```typescript + * declare module "@mariozechner/agent" { + * interface CustomAgentMessages { + * artifact: ArtifactMessage; + * notification: NotificationMessage; + * } + * } + * ``` + */ +export interface CustomAgentMessages { + // Empty by default - apps extend via declaration merging +} + +/** + * AgentMessage: Union of LLM messages + custom messages. + * This abstraction allows apps to add custom message types while maintaining + * type safety and compatibility with the base LLM messages. + */ +export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages]; + +/** + * Public agent state. + * + * `tools` and `messages` use accessor properties so implementations can copy + * assigned arrays before storing them. + */ +export interface AgentState { + /** System prompt sent with each model request. */ + systemPrompt: string; + /** Active model used for future turns. */ + model: Model; + /** Requested reasoning level for future turns. */ + thinkingLevel: ThinkingLevel; + /** Available tools. Assigning a new array copies the top-level array. */ + set tools(tools: AgentTool[]); + get tools(): AgentTool[]; + /** Conversation transcript. Assigning a new array copies the top-level array. */ + set messages(messages: AgentMessage[]); + get messages(): AgentMessage[]; + /** + * True while the agent is processing a prompt or continuation. + * + * This remains true until awaited `agent_end` listeners settle. + */ + readonly isStreaming: boolean; + /** Partial assistant message for the current streamed response, if any. */ + readonly streamingMessage?: AgentMessage; + /** Tool call ids currently executing. */ + readonly pendingToolCalls: ReadonlySet; + /** Error message from the most recent failed or aborted assistant turn, if any. */ + readonly errorMessage?: string; +} + +/** Final or partial result produced by a tool. */ +export interface AgentToolResult { + /** Text or image content returned to the model. */ + content: (TextContent | ImageContent)[]; + /** Arbitrary structured details for logs or UI rendering. */ + details: T; + /** + * Hint that the agent should stop after the current tool batch. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; +} + +/** + * Callback used by tools to stream partial execution updates. + * + * The callback is scoped to the current `execute()` invocation. Calls made after + * the tool promise settles are ignored. + */ +export type AgentToolUpdateCallback = (partialResult: AgentToolResult) => void; + +/** Tool definition used by the agent runtime. */ +export interface AgentTool extends Tool { + /** Human-readable label for UI display. */ + label: string; + /** + * Optional compatibility shim for raw tool-call arguments before schema validation. + * Must return an object that matches `TParameters`. + */ + prepareArguments?: (args: unknown) => Static; + /** Execute the tool call. Throw on failure instead of encoding errors in `content`. */ + execute: ( + toolCallId: string, + params: Static, + signal?: AbortSignal, + onUpdate?: AgentToolUpdateCallback, + ) => Promise>; + /** + * Per-tool execution mode override. + * - "sequential": this tool must execute one at a time with other tool calls. + * - "parallel": this tool can execute concurrently with other tool calls. + * + * If omitted, the default execution mode applies. + */ + executionMode?: ToolExecutionMode; +} + +/** Context snapshot passed into the low-level agent loop. */ +export interface AgentContext { + /** System prompt included with the request. */ + systemPrompt: string; + /** Transcript visible to the model. */ + messages: AgentMessage[]; + /** Tools available for this run. */ + tools?: AgentTool[]; +} + +/** + * Events emitted by the Agent for UI updates. + * + * `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()` + * listeners for that event are still part of run settlement. The agent becomes + * idle only after those listeners finish. + */ +export type AgentEvent = + // Agent lifecycle + | { type: "agent_start" } + | { type: "agent_end"; messages: AgentMessage[] } + // Turn lifecycle - a turn is one assistant response + any tool calls/results + | { type: "turn_start" } + | { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] } + // Message lifecycle - emitted for user, assistant, and toolResult messages + | { type: "message_start"; message: AgentMessage } + // Only emitted for assistant messages during streaming + | { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent } + | { type: "message_end"; message: AgentMessage } + // Tool execution lifecycle + | { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any } + | { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any } + | { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean }; diff --git a/cactus-code/packages/agent/tsconfig.build.json b/cactus-code/packages/agent/tsconfig.build.json new file mode 100644 index 000000000..979dbc899 --- /dev/null +++ b/cactus-code/packages/agent/tsconfig.build.json @@ -0,0 +1,13 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "paths": { + "@earendil-works/pi-ai": ["../ai/dist/index.d.ts"], + "@earendil-works/pi-ai/*": ["../ai/dist/*.d.ts", "../ai/dist/providers/*.d.ts"] + }, + "rootDir": "./src" + }, + "include": ["src/**/*.ts"], + "exclude": ["node_modules", "dist", "**/*.d.ts", "src/**/*.d.ts"] +} diff --git a/cactus-code/packages/ai/package.json b/cactus-code/packages/ai/package.json new file mode 100644 index 000000000..6c63411de --- /dev/null +++ b/cactus-code/packages/ai/package.json @@ -0,0 +1,70 @@ +{ + "name": "@earendil-works/pi-ai", + "version": "0.80.2", + "description": "Unified LLM API with automatic model discovery and provider configuration", + "type": "module", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "sideEffects": [ + "./dist/compat.js" + ], + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + }, + "./compat": { + "types": "./dist/compat.d.ts", + "import": "./dist/compat.js" + }, + "./providers/*": { + "types": "./dist/providers/*.d.ts", + "import": "./dist/providers/*.js" + }, + "./api/*": { + "types": "./dist/api/*.d.ts", + "import": "./dist/api/*.js" + }, + "./oauth": { + "types": "./dist/oauth.d.ts", + "import": "./dist/oauth.js" + } + }, + "files": [ + "dist", + "README.md" + ], + "scripts": { + "clean": "shx rm -rf dist", + "build": "tsgo -p tsconfig.build.json", + "prepublishOnly": "npm run clean && npm run build" + }, + "dependencies": { + "openai": "6.26.0", + "partial-json": "0.1.7", + "typebox": "1.1.38" + }, + "keywords": [ + "ai", + "llm", + "openai", + "anthropic", + "gemini", + "bedrock", + "unified", + "api" + ], + "author": "Mario Zechner", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/ai" + }, + "engines": { + "node": ">=22.19.0" + }, + "devDependencies": { + "@types/node": "24.12.4" + } +} diff --git a/cactus-code/packages/ai/src/api/github-copilot-headers.ts b/cactus-code/packages/ai/src/api/github-copilot-headers.ts new file mode 100644 index 000000000..602750c49 --- /dev/null +++ b/cactus-code/packages/ai/src/api/github-copilot-headers.ts @@ -0,0 +1,37 @@ +import type { Message } from "../types.ts"; + +// Copilot expects X-Initiator to indicate whether the request is user-initiated +// or agent-initiated (e.g. follow-up after assistant/tool messages). +export function inferCopilotInitiator(messages: Message[]): "user" | "agent" { + const last = messages[messages.length - 1]; + return last && last.role !== "user" ? "agent" : "user"; +} + +// Copilot requires Copilot-Vision-Request header when sending images +export function hasCopilotVisionInput(messages: Message[]): boolean { + return messages.some((msg) => { + if (msg.role === "user" && Array.isArray(msg.content)) { + return msg.content.some((c) => c.type === "image"); + } + if (msg.role === "toolResult" && Array.isArray(msg.content)) { + return msg.content.some((c) => c.type === "image"); + } + return false; + }); +} + +export function buildCopilotDynamicHeaders(params: { + messages: Message[]; + hasImages: boolean; +}): Record { + const headers: Record = { + "X-Initiator": inferCopilotInitiator(params.messages), + "Openai-Intent": "conversation-edits", + }; + + if (params.hasImages) { + headers["Copilot-Vision-Request"] = "true"; + } + + return headers; +} diff --git a/cactus-code/packages/ai/src/api/lazy.ts b/cactus-code/packages/ai/src/api/lazy.ts new file mode 100644 index 000000000..fe1836aea --- /dev/null +++ b/cactus-code/packages/ai/src/api/lazy.ts @@ -0,0 +1,70 @@ +import type { Api, AssistantMessage, AssistantMessageEvent, Model, ProviderStreams } from "../types.ts"; +import { AssistantMessageEventStream } from "../utils/event-stream.ts"; + +function createSetupErrorMessage(model: Model, error: unknown): AssistantMessage { + return { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + }; +} + +function forwardStream(target: AssistantMessageEventStream, source: AsyncIterable): void { + (async () => { + for await (const event of source) { + target.push(event); + } + target.end(); + })(); +} + +/** + * Returns a stream synchronously while running async setup (auth resolution, + * lazy module loading) behind it. Setup failures terminate the stream with an + * error event. + */ +export function lazyStream( + model: Model, + setup: () => Promise>, +): AssistantMessageEventStream { + const outer = new AssistantMessageEventStream(); + + setup() + .then((inner) => { + forwardStream(outer, inner); + }) + .catch((error) => { + const message = createSetupErrorMessage(model, error); + outer.push({ type: "error", reason: "error", error: message }); + outer.end(message); + }); + + return outer; +} + +/** + * Wraps a dynamically imported API implementation module as `ProviderStreams`. + * The module loads on first stream call; the host's import cache deduplicates + * loads. Load failures terminate the returned stream with an error event. + */ +export function lazyApi(load: () => Promise): ProviderStreams { + return { + stream: (model, context, options) => + lazyStream(model, async () => (await load()).stream(model, context, options)), + streamSimple: (model, context, options) => + lazyStream(model, async () => (await load()).streamSimple(model, context, options)), + }; +} diff --git a/cactus-code/packages/ai/src/api/openai-completions.lazy.ts b/cactus-code/packages/ai/src/api/openai-completions.lazy.ts new file mode 100644 index 000000000..6f6c6f61a --- /dev/null +++ b/cactus-code/packages/ai/src/api/openai-completions.lazy.ts @@ -0,0 +1,4 @@ +import type { ProviderStreams } from "../types.ts"; +import { lazyApi } from "./lazy.ts"; + +export const openAICompletionsApi = (): ProviderStreams => lazyApi(() => import("./openai-completions.ts")); diff --git a/cactus-code/packages/ai/src/api/openai-completions.ts b/cactus-code/packages/ai/src/api/openai-completions.ts new file mode 100644 index 000000000..2dae2c6e3 --- /dev/null +++ b/cactus-code/packages/ai/src/api/openai-completions.ts @@ -0,0 +1,1288 @@ +import OpenAI from "openai"; +import type { + ChatCompletionAssistantMessageParam, + ChatCompletionChunk, + ChatCompletionContentPart, + ChatCompletionContentPartImage, + ChatCompletionContentPartText, + ChatCompletionDeveloperMessageParam, + ChatCompletionMessageParam, + ChatCompletionSystemMessageParam, + ChatCompletionToolMessageParam, +} from "openai/resources/chat/completions.js"; +import { calculateCost, clampThinkingLevel } from "../models.ts"; +import type { + AssistantMessage, + CacheRetention, + ChatTemplateKwargValue, + Context, + ImageContent, + Message, + Model, + OpenAICompletionsCompat, + ProviderEnv, + ProviderHeaders, + SimpleStreamOptions, + StopReason, + StreamFunction, + StreamOptions, + TextContent, + ThinkingContent, + Tool, + ToolCall, + ToolResultMessage, +} from "../types.ts"; +import { AssistantMessageEventStream } from "../utils/event-stream.ts"; +import { headersToRecord } from "../utils/headers.ts"; +import { parseStreamingJson } from "../utils/json-parse.ts"; +import { getProviderEnvValue } from "../utils/provider-env.ts"; +import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; +import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.ts"; +import { clampOpenAIPromptCacheKey } from "./openai-prompt-cache.ts"; +import { buildBaseOptions } from "./simple-options.ts"; +import { transformMessages } from "./transform-messages.ts"; + +/** + * Check if conversation messages contain tool calls or tool results. + * This is needed because Anthropic (via proxy) requires the tools param + * to be present when messages include tool_calls or tool role messages. + */ +function hasHeader(headers: ProviderHeaders | undefined, name: string): boolean { + if (!headers) return false; + const expected = name.toLowerCase(); + for (const [key, value] of Object.entries(headers)) { + if (key.toLowerCase() === expected && value !== null && value.trim().length > 0) return true; + } + return false; +} + +function getClientApiKey(provider: string, apiKey: string | undefined, headers: ProviderHeaders | undefined): string { + if (apiKey) return apiKey; + if (hasHeader(headers, "authorization") || hasHeader(headers, "cf-aig-authorization")) return "unused"; + throw new Error(`No API key for provider: ${provider}`); +} + +function hasToolHistory(messages: Message[]): boolean { + for (const msg of messages) { + if (msg.role === "toolResult") { + return true; + } + if (msg.role === "assistant") { + if (msg.content.some((block) => block.type === "toolCall")) { + return true; + } + } + } + return false; +} + +function isTextContentBlock(block: { type: string }): block is TextContent { + return block.type === "text"; +} + +function isThinkingContentBlock(block: { type: string }): block is ThinkingContent { + return block.type === "thinking"; +} + +function isToolCallBlock(block: { type: string }): block is ToolCall { + return block.type === "toolCall"; +} + +function isImageContentBlock(block: { type: string }): block is ImageContent { + return block.type === "image"; +} + +function isEncryptedReasoningDetail(detail: unknown): detail is OpenAIEncryptedReasoningDetail { + if (typeof detail !== "object" || detail === null) { + return false; + } + const candidate = detail as Record; + return ( + candidate.type === "reasoning.encrypted" && + typeof candidate.id === "string" && + candidate.id.length > 0 && + typeof candidate.data === "string" && + candidate.data.length > 0 + ); +} + +export interface OpenAICompletionsOptions extends StreamOptions { + toolChoice?: "auto" | "none" | "required" | { type: "function"; function: { name: string } }; + reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh"; +} + +interface OpenAICompatCacheControl { + type: "ephemeral"; + ttl?: string; +} + +type ResolvedOpenAICompletionsCompat = Omit, "cacheControlFormat"> & { + cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; +}; + +type ResolvedChatTemplateKwargValue = string | number | boolean | null; + +type ChatCompletionInstructionMessageParam = ChatCompletionDeveloperMessageParam | ChatCompletionSystemMessageParam; + +type OpenAIEncryptedReasoningDetail = { + type: "reasoning.encrypted"; + id: string; + data: string; +}; + +type ChatCompletionTextPartWithCacheControl = ChatCompletionContentPartText & { + cache_control?: OpenAICompatCacheControl; +}; + +type ChatCompletionToolWithCacheControl = OpenAI.Chat.Completions.ChatCompletionTool & { + cache_control?: OpenAICompatCacheControl; +}; + +function resolveCacheRetention(cacheRetention?: CacheRetention, env?: ProviderEnv): CacheRetention { + if (cacheRetention) { + return cacheRetention; + } + if (getProviderEnvValue("PI_CACHE_RETENTION", env) === "long") { + return "long"; + } + return "short"; +} + +export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptions> = ( + model: Model<"openai-completions">, + context: Context, + options?: OpenAICompletionsOptions, +): AssistantMessageEventStream => { + const stream = new AssistantMessageEventStream(); + + (async () => { + const output: AssistantMessage = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; + + try { + const apiKey = getClientApiKey(model.provider, options?.apiKey, options?.headers); + const compat = getCompat(model); + const cacheRetention = resolveCacheRetention(options?.cacheRetention, options?.env); + const cacheSessionId = cacheRetention === "none" ? undefined : options?.sessionId; + const client = createClient(model, context, apiKey, options?.headers, cacheSessionId, compat); + let params = buildParams(model, context, options, compat, cacheRetention); + const nextParams = await options?.onPayload?.(params, model); + if (nextParams !== undefined) { + params = nextParams as OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming; + } + const requestOptions = { + ...(options?.signal ? { signal: options.signal } : {}), + ...(options?.timeoutMs !== undefined ? { timeout: options.timeoutMs } : {}), + maxRetries: options?.maxRetries ?? 0, + }; + const { data: openaiStream, response } = await client.chat.completions + .create(params, requestOptions) + .withResponse(); + await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); + stream.push({ type: "start", partial: output }); + + interface StreamingToolCallBlock extends ToolCall { + partialArgs?: string; + streamIndex?: number; + } + type StreamingBlock = TextContent | ThinkingContent | StreamingToolCallBlock; + type StreamingToolCallDelta = NonNullable[number]; + + let textBlock: TextContent | null = null; + let thinkingBlock: ThinkingContent | null = null; + let hasFinishReason = false; + const toolCallBlocksByIndex = new Map(); + const toolCallBlocksById = new Map(); + const pendingReasoningDetailsByToolCallId = new Map(); + const blocks = output.content as StreamingBlock[]; + const getContentIndex = (block: StreamingBlock) => blocks.indexOf(block); + const finishBlock = (block: StreamingBlock) => { + const contentIndex = getContentIndex(block); + if (contentIndex === -1) { + return; + } + if (block.type === "text") { + stream.push({ + type: "text_end", + contentIndex, + content: block.text, + partial: output, + }); + } else if (block.type === "thinking") { + stream.push({ + type: "thinking_end", + contentIndex, + content: block.thinking, + partial: output, + }); + } else if (block.type === "toolCall") { + block.arguments = parseStreamingJson(block.partialArgs); + // Finalize in-place and strip the scratch buffers so replay only + // carries parsed arguments. + delete block.partialArgs; + delete block.streamIndex; + stream.push({ + type: "toolcall_end", + contentIndex, + toolCall: block, + partial: output, + }); + } + }; + const ensureTextBlock = () => { + if (!textBlock) { + textBlock = { type: "text", text: "" }; + blocks.push(textBlock); + stream.push({ type: "text_start", contentIndex: getContentIndex(textBlock), partial: output }); + } + return textBlock; + }; + const ensureThinkingBlock = (thinkingSignature: string) => { + if (!thinkingBlock) { + thinkingBlock = { + type: "thinking", + thinking: "", + thinkingSignature, + }; + blocks.push(thinkingBlock); + stream.push({ type: "thinking_start", contentIndex: getContentIndex(thinkingBlock), partial: output }); + } + return thinkingBlock; + }; + const applyPendingReasoningDetail = (block: StreamingToolCallBlock) => { + if (!block.id) { + return; + } + const pendingReasoningDetail = pendingReasoningDetailsByToolCallId.get(block.id); + if (pendingReasoningDetail) { + block.thoughtSignature = pendingReasoningDetail; + pendingReasoningDetailsByToolCallId.delete(block.id); + } + }; + const ensureToolCallBlock = (toolCall: StreamingToolCallDelta) => { + const streamIndex = typeof toolCall.index === "number" ? toolCall.index : undefined; + let block = streamIndex !== undefined ? toolCallBlocksByIndex.get(streamIndex) : undefined; + if (!block && toolCall.id) { + block = toolCallBlocksById.get(toolCall.id); + } + if (!block) { + block = { + type: "toolCall", + id: toolCall.id || "", + name: toolCall.function?.name || "", + arguments: {}, + partialArgs: "", + streamIndex, + }; + if (streamIndex !== undefined) { + toolCallBlocksByIndex.set(streamIndex, block); + } + if (toolCall.id) { + toolCallBlocksById.set(toolCall.id, block); + } + blocks.push(block); + stream.push({ + type: "toolcall_start", + contentIndex: getContentIndex(block), + partial: output, + }); + } + if (streamIndex !== undefined && block.streamIndex === undefined) { + block.streamIndex = streamIndex; + toolCallBlocksByIndex.set(streamIndex, block); + } + if (toolCall.id) { + toolCallBlocksById.set(toolCall.id, block); + } + applyPendingReasoningDetail(block); + return block; + }; + + for await (const chunk of openaiStream) { + if (!chunk || typeof chunk !== "object") continue; + + // OpenAI documents ChatCompletionChunk.id as the unique chat completion identifier, + // and each chunk in a streamed completion carries the same id. + output.responseId ||= chunk.id; + if (typeof chunk.model === "string" && chunk.model.length > 0 && chunk.model !== model.id) { + output.responseModel ||= chunk.model; + } + if (chunk.usage) { + output.usage = parseChunkUsage(chunk.usage, model); + } + if (typeof (chunk as any).cloud_handoff === "boolean") { + output.cloudHandoff = (chunk as any).cloud_handoff; + } + + const choice = Array.isArray(chunk.choices) ? chunk.choices[0] : undefined; + if (!choice) continue; + + // Fallback: some providers (e.g., Moonshot) return usage + // in choice.usage instead of the standard chunk.usage + if (!chunk.usage && (choice as any).usage) { + output.usage = parseChunkUsage((choice as any).usage, model); + } + + if (choice.finish_reason) { + const finishReasonResult = mapStopReason(choice.finish_reason); + output.stopReason = finishReasonResult.stopReason; + if (finishReasonResult.errorMessage) { + output.errorMessage = finishReasonResult.errorMessage; + } + hasFinishReason = true; + } + + if (choice.delta) { + if ( + choice.delta.content !== null && + choice.delta.content !== undefined && + choice.delta.content.length > 0 + ) { + const block = ensureTextBlock(); + block.text += choice.delta.content; + stream.push({ + type: "text_delta", + contentIndex: getContentIndex(block), + delta: choice.delta.content, + partial: output, + }); + } + + // Some endpoints return reasoning in reasoning_content (llama.cpp), + // or reasoning (other openai compatible endpoints) + // Use the first non-empty reasoning field to avoid duplication + // (e.g., chutes.ai returns both reasoning_content and reasoning with same content) + const reasoningFields = ["reasoning_content", "reasoning", "reasoning_text"]; + const deltaFields = choice.delta as Record; + let foundReasoningField: string | null = null; + for (const field of reasoningFields) { + const value = deltaFields[field]; + if (typeof value === "string" && value.length > 0) { + foundReasoningField = field; + break; + } + } + + if (foundReasoningField) { + const delta = deltaFields[foundReasoningField]; + if (typeof delta === "string" && delta.length > 0) { + const thinkingSignature = + model.provider === "opencode-go" && foundReasoningField === "reasoning" + ? "reasoning_content" + : foundReasoningField; + const block = ensureThinkingBlock(thinkingSignature); + block.thinking += delta; + stream.push({ + type: "thinking_delta", + contentIndex: getContentIndex(block), + delta, + partial: output, + }); + } + } + + if (choice?.delta?.tool_calls) { + for (const toolCall of choice.delta.tool_calls) { + const block = ensureToolCallBlock(toolCall); + if (!block.id && toolCall.id) { + block.id = toolCall.id; + toolCallBlocksById.set(toolCall.id, block); + } + if (!block.name && toolCall.function?.name) { + block.name = toolCall.function.name; + } + + let delta = ""; + if (toolCall.function?.arguments) { + delta = toolCall.function.arguments; + block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments; + block.arguments = parseStreamingJson(block.partialArgs); + } + stream.push({ + type: "toolcall_delta", + contentIndex: getContentIndex(block), + delta, + partial: output, + }); + } + } + + const reasoningDetails = (choice.delta as { reasoning_details?: unknown }).reasoning_details; + if (Array.isArray(reasoningDetails)) { + for (const detail of reasoningDetails) { + if (isEncryptedReasoningDetail(detail)) { + const serializedDetail = JSON.stringify(detail); + const matchingToolCall = toolCallBlocksById.get(detail.id); + if (matchingToolCall) { + matchingToolCall.thoughtSignature = serializedDetail; + } else { + pendingReasoningDetailsByToolCallId.set(detail.id, serializedDetail); + } + } + } + } + } + } + + for (const block of blocks) { + finishBlock(block); + } + if (options?.signal?.aborted) { + throw new Error("Request was aborted"); + } + + if (output.stopReason === "aborted") { + throw new Error("Request was aborted"); + } + if (output.stopReason === "error") { + throw new Error(output.errorMessage || "Provider returned an error stop reason"); + } + if (!hasFinishReason) { + throw new Error("Stream ended without finish_reason"); + } + + stream.push({ type: "done", reason: output.stopReason, message: output }); + stream.end(); + } catch (error) { + for (const block of output.content) { + delete (block as { index?: number }).index; + // Streaming scratch buffers are only used during parsing; never persist them. + delete (block as { partialArgs?: string }).partialArgs; + delete (block as { streamIndex?: number }).streamIndex; + } + output.stopReason = options?.signal?.aborted ? "aborted" : "error"; + output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); + // Some providers via OpenRouter give additional information in this field. + const rawMetadata = (error as any)?.error?.metadata?.raw; + if (rawMetadata) output.errorMessage += `\n${rawMetadata}`; + stream.push({ type: "error", reason: output.stopReason, error: output }); + stream.end(); + } + })(); + + return stream; +}; + +export const streamSimple: StreamFunction<"openai-completions", SimpleStreamOptions> = ( + model: Model<"openai-completions">, + context: Context, + options?: SimpleStreamOptions, +): AssistantMessageEventStream => { + getClientApiKey(model.provider, options?.apiKey, options?.headers); + + const base = buildBaseOptions(model, options, options?.apiKey); + const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; + const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning; + const toolChoice = (options as OpenAICompletionsOptions | undefined)?.toolChoice; + + return stream(model, context, { + ...base, + reasoningEffort, + toolChoice, + } satisfies OpenAICompletionsOptions); +}; + +function createClient( + model: Model<"openai-completions">, + context: Context, + apiKey: string, + optionsHeaders?: ProviderHeaders, + sessionId?: string, + compat: ResolvedOpenAICompletionsCompat = getCompat(model), +) { + const headers: ProviderHeaders = { ...model.headers }; + if (model.provider === "github-copilot") { + const hasImages = hasCopilotVisionInput(context.messages); + const copilotHeaders = buildCopilotDynamicHeaders({ + messages: context.messages, + hasImages, + }); + Object.assign(headers, copilotHeaders); + } + + if (sessionId && compat.sendSessionAffinityHeaders) { + headers.session_id = sessionId; + headers["x-client-request-id"] = sessionId; + headers["x-session-affinity"] = sessionId; + } + + // Merge options headers last so they can override defaults + if (optionsHeaders) { + Object.assign(headers, optionsHeaders); + } + + return new OpenAI({ + apiKey, + baseURL: model.baseUrl, + dangerouslyAllowBrowser: true, + defaultHeaders: headers, + }); +} + +function buildParams( + model: Model<"openai-completions">, + context: Context, + options?: OpenAICompletionsOptions, + compat: ResolvedOpenAICompletionsCompat = getCompat(model), + cacheRetention: CacheRetention = resolveCacheRetention(options?.cacheRetention, options?.env), +) { + const messages = convertMessages(model, context, compat); + const cacheControl = getCompatCacheControl(compat, cacheRetention); + + const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: model.id, + messages, + stream: true, + prompt_cache_key: + (model.baseUrl.includes("api.openai.com") && cacheRetention !== "none") || + (cacheRetention === "long" && compat.supportsLongCacheRetention) + ? clampOpenAIPromptCacheKey(options?.sessionId) + : undefined, + prompt_cache_retention: cacheRetention === "long" && compat.supportsLongCacheRetention ? "24h" : undefined, + }; + + if (compat.supportsUsageInStreaming !== false) { + (params as any).stream_options = { include_usage: true }; + } + + if (compat.supportsStore) { + params.store = false; + } + + if (options?.maxTokens) { + if (compat.maxTokensField === "max_tokens") { + (params as any).max_tokens = options.maxTokens; + } else { + params.max_completion_tokens = options.maxTokens; + } + } + + if (options?.temperature !== undefined) { + params.temperature = options.temperature; + } + + if (context.tools && context.tools.length > 0) { + params.tools = convertTools(context.tools, compat); + if (compat.zaiToolStream) { + (params as any).tool_stream = true; + } + } else if (hasToolHistory(context.messages)) { + // Anthropic (via LiteLLM/proxy) requires tools param when conversation has tool_calls/tool_results + params.tools = []; + } + + if (cacheControl) { + applyAnthropicCacheControl(messages, params.tools, cacheControl); + } + + if (options?.toolChoice) { + params.tool_choice = options.toolChoice; + } + + if (compat.thinkingFormat === "zai" && model.reasoning) { + const zaiParams = params as Omit & { + thinking?: { type: "enabled" | "disabled" }; + reasoning_effort?: string; + }; + zaiParams.thinking = { type: options?.reasoningEffort ? "enabled" : "disabled" }; + if (options?.reasoningEffort && compat.supportsReasoningEffort) { + const mappedEffort = model.thinkingLevelMap?.[options.reasoningEffort]; + const effort = mappedEffort === undefined ? options.reasoningEffort : mappedEffort; + if (typeof effort === "string") { + zaiParams.reasoning_effort = effort; + } + } + } else if (compat.thinkingFormat === "qwen" && model.reasoning) { + (params as any).enable_thinking = !!options?.reasoningEffort; + } else if (compat.thinkingFormat === "qwen-chat-template" && model.reasoning) { + (params as any).chat_template_kwargs = { + enable_thinking: !!options?.reasoningEffort, + preserve_thinking: true, + }; + } else if (compat.thinkingFormat === "chat-template" && model.reasoning) { + const chatTemplateKwargs = buildChatTemplateKwargs(model, options, compat); + if (chatTemplateKwargs) { + (params as any).chat_template_kwargs = chatTemplateKwargs; + } + } else if (compat.thinkingFormat === "deepseek" && model.reasoning) { + if (options?.reasoningEffort) { + (params as any).thinking = { type: "enabled" }; + } else if (model.thinkingLevelMap?.off !== null) { + (params as any).thinking = { type: "disabled" }; + } + if (options?.reasoningEffort && compat.supportsReasoningEffort) { + (params as any).reasoning_effort = + model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort; + } + } else if (compat.thinkingFormat === "openrouter" && model.reasoning) { + // OpenRouter normalizes reasoning across providers via a nested reasoning object. + const openRouterParams = params as typeof params & { reasoning?: { effort?: string } }; + if (options?.reasoningEffort) { + openRouterParams.reasoning = { + effort: model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort, + }; + } else if (model.thinkingLevelMap?.off !== null) { + openRouterParams.reasoning = { effort: model.thinkingLevelMap?.off ?? "none" }; + } + } else if (compat.thinkingFormat === "ant-ling" && model.reasoning && options?.reasoningEffort) { + const effort = model.thinkingLevelMap?.[options.reasoningEffort]; + if (typeof effort === "string") { + (params as typeof params & { reasoning?: { effort: string } }).reasoning = { effort }; + } + } else if (compat.thinkingFormat === "together" && model.reasoning) { + const togetherParams = params as Omit & { + reasoning?: { enabled: boolean }; + reasoning_effort?: string; + }; + togetherParams.reasoning = { enabled: !!options?.reasoningEffort }; + if (options?.reasoningEffort && compat.supportsReasoningEffort) { + togetherParams.reasoning_effort = model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort; + } + } else if (compat.thinkingFormat === "string-thinking" && model.reasoning) { + const stringThinkingParams = params as typeof params & { thinking?: string }; + if (options?.reasoningEffort) { + stringThinkingParams.thinking = model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort; + } else if (model.thinkingLevelMap?.off !== null) { + stringThinkingParams.thinking = model.thinkingLevelMap?.off ?? "none"; + } + } else if (options?.reasoningEffort && model.reasoning && compat.supportsReasoningEffort) { + // OpenAI-style reasoning_effort + (params as any).reasoning_effort = model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort; + } else if (!options?.reasoningEffort && model.reasoning && compat.supportsReasoningEffort) { + const offValue = model.thinkingLevelMap?.off; + if (typeof offValue === "string") { + (params as any).reasoning_effort = offValue; + } + } + + // OpenRouter provider routing preferences + if (model.compat?.openRouterRouting) { + (params as any).provider = model.compat.openRouterRouting; + } + + // Vercel AI Gateway provider routing preferences + if (model.compat?.vercelGatewayRouting) { + const routing = model.compat.vercelGatewayRouting; + if (routing.only || routing.order) { + const gatewayOptions: Record = {}; + if (routing.only) gatewayOptions.only = routing.only; + if (routing.order) gatewayOptions.order = routing.order; + (params as any).providerOptions = { gateway: gatewayOptions }; + } + } + + return params; +} + +function buildChatTemplateKwargs( + model: Model<"openai-completions">, + options: OpenAICompletionsOptions | undefined, + compat: ResolvedOpenAICompletionsCompat, +): Record | undefined { + const kwargs: Record = {}; + + for (const [key, value] of Object.entries(compat.chatTemplateKwargs)) { + const resolved = resolveChatTemplateKwargValue(model, options, value); + if (resolved !== undefined) { + kwargs[key] = resolved; + } + } + + return Object.keys(kwargs).length > 0 ? kwargs : undefined; +} + +function resolveChatTemplateKwargValue( + model: Model<"openai-completions">, + options: OpenAICompletionsOptions | undefined, + value: ChatTemplateKwargValue, +): ResolvedChatTemplateKwargValue | undefined { + if (typeof value !== "object" || value === null) { + return value; + } + + const reasoningEffort = options?.reasoningEffort; + if (!reasoningEffort && value.omitWhenOff) { + return undefined; + } + if (value.$var === "thinking.enabled") { + return !!reasoningEffort; + } + + const mappedValue = reasoningEffort ? model.thinkingLevelMap?.[reasoningEffort] : model.thinkingLevelMap?.off; + return mappedValue === undefined ? reasoningEffort : typeof mappedValue === "string" ? mappedValue : undefined; +} + +function getCompatCacheControl( + compat: ResolvedOpenAICompletionsCompat, + cacheRetention: CacheRetention, +): OpenAICompatCacheControl | undefined { + if (compat.cacheControlFormat !== "anthropic" || cacheRetention === "none") { + return undefined; + } + + const ttl = cacheRetention === "long" && compat.supportsLongCacheRetention ? "1h" : undefined; + return { type: "ephemeral", ...(ttl ? { ttl } : {}) }; +} + +function applyAnthropicCacheControl( + messages: ChatCompletionMessageParam[], + tools: OpenAI.Chat.Completions.ChatCompletionTool[] | undefined, + cacheControl: OpenAICompatCacheControl, +): void { + addCacheControlToSystemPrompt(messages, cacheControl); + addCacheControlToLastTool(tools, cacheControl); + addCacheControlToLastConversationMessage(messages, cacheControl); +} + +function addCacheControlToSystemPrompt( + messages: ChatCompletionMessageParam[], + cacheControl: OpenAICompatCacheControl, +): void { + for (const message of messages) { + if (message.role === "system" || message.role === "developer") { + addCacheControlToInstructionMessage(message, cacheControl); + return; + } + } +} + +function addCacheControlToLastConversationMessage( + messages: ChatCompletionMessageParam[], + cacheControl: OpenAICompatCacheControl, +): void { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]; + if (message.role === "user" || message.role === "assistant") { + if (addCacheControlToMessage(message, cacheControl)) { + return; + } + } + } +} + +function addCacheControlToLastTool( + tools: OpenAI.Chat.Completions.ChatCompletionTool[] | undefined, + cacheControl: OpenAICompatCacheControl, +): void { + if (!tools || tools.length === 0) { + return; + } + + const lastTool = tools[tools.length - 1] as ChatCompletionToolWithCacheControl; + lastTool.cache_control = cacheControl; +} + +function addCacheControlToInstructionMessage( + message: ChatCompletionInstructionMessageParam, + cacheControl: OpenAICompatCacheControl, +): boolean { + return addCacheControlToTextContent(message, cacheControl); +} + +function addCacheControlToMessage( + message: ChatCompletionMessageParam, + cacheControl: OpenAICompatCacheControl, +): boolean { + if (message.role === "user" || message.role === "assistant") { + return addCacheControlToTextContent(message, cacheControl); + } + return false; +} + +function addCacheControlToTextContent( + message: + | ChatCompletionInstructionMessageParam + | ChatCompletionAssistantMessageParam + | Extract, + cacheControl: OpenAICompatCacheControl, +): boolean { + const content = message.content; + if (typeof content === "string") { + if (content.length === 0) { + return false; + } + message.content = [ + { + type: "text", + text: content, + cache_control: cacheControl, + }, + ] as ChatCompletionTextPartWithCacheControl[]; + return true; + } + + if (!Array.isArray(content)) { + return false; + } + + for (let i = content.length - 1; i >= 0; i--) { + const part = content[i]; + if (part?.type === "text") { + const textPart = part as ChatCompletionTextPartWithCacheControl; + textPart.cache_control = cacheControl; + return true; + } + } + + return false; +} + +export function convertMessages( + model: Model<"openai-completions">, + context: Context, + compat: ResolvedOpenAICompletionsCompat, +): ChatCompletionMessageParam[] { + const params: ChatCompletionMessageParam[] = []; + + const normalizeToolCallId = (id: string): string => { + // Handle pipe-separated IDs from OpenAI Responses API + // Format: {call_id}|{id} where {id} can be 400+ chars with special chars (+, /, =) + // These come from providers like github-copilot, openai-codex, opencode + // Extract just the call_id part and normalize it + if (id.includes("|")) { + const [callId] = id.split("|"); + // Sanitize to allowed chars and truncate to 40 chars (OpenAI limit) + return callId.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 40); + } + + if (model.provider === "openai") return id.length > 40 ? id.slice(0, 40) : id; + return id; + }; + + const transformedMessages = transformMessages(context.messages, model, (id) => normalizeToolCallId(id)); + + if (context.systemPrompt) { + const useDeveloperRole = model.reasoning && compat.supportsDeveloperRole; + const role = useDeveloperRole ? "developer" : "system"; + params.push({ role: role, content: sanitizeSurrogates(context.systemPrompt) }); + } + + let lastRole: string | null = null; + + for (let i = 0; i < transformedMessages.length; i++) { + const msg = transformedMessages[i]; + // Some providers don't allow user messages directly after tool results + // Insert a synthetic assistant message to bridge the gap + if (compat.requiresAssistantAfterToolResult && lastRole === "toolResult" && msg.role === "user") { + params.push({ + role: "assistant", + content: "I have processed the tool results.", + }); + } + + if (msg.role === "user") { + if (typeof msg.content === "string") { + params.push({ + role: "user", + content: sanitizeSurrogates(msg.content), + }); + } else { + const content: ChatCompletionContentPart[] = msg.content.map((item): ChatCompletionContentPart => { + if (item.type === "text") { + return { + type: "text", + text: sanitizeSurrogates(item.text), + } satisfies ChatCompletionContentPartText; + } else { + return { + type: "image_url", + image_url: { + url: `data:${item.mimeType};base64,${item.data}`, + }, + } satisfies ChatCompletionContentPartImage; + } + }); + if (content.length === 0) continue; + params.push({ + role: "user", + content, + }); + } + } else if (msg.role === "assistant") { + // Some providers don't accept null content, use empty string instead + const assistantMsg: ChatCompletionAssistantMessageParam = { + role: "assistant", + content: compat.requiresAssistantAfterToolResult ? "" : null, + }; + + const assistantTextParts = msg.content + .filter(isTextContentBlock) + .filter((block) => block.text.trim().length > 0) + .map( + (block) => + ({ + type: "text", + text: sanitizeSurrogates(block.text), + }) satisfies ChatCompletionContentPartText, + ); + const assistantText = assistantTextParts.map((part) => part.text).join(""); + + const nonEmptyThinkingBlocks = msg.content + .filter(isThinkingContentBlock) + .filter((block) => block.thinking.trim().length > 0); + if (nonEmptyThinkingBlocks.length > 0) { + if (compat.requiresThinkingAsText) { + // Convert thinking blocks to plain text (no tags to avoid model mimicking them) + const thinkingText = nonEmptyThinkingBlocks + .map((block) => sanitizeSurrogates(block.thinking)) + .join("\n\n"); + assistantMsg.content = [{ type: "text", text: thinkingText }, ...assistantTextParts]; + } else { + // Always send assistant content as a plain string (OpenAI Chat Completions + // API standard format). Sending as an array of {type:"text", text:"..."} + // objects is non-standard and causes some models (e.g. DeepSeek V3.2 via + // NVIDIA NIM) to mirror the content-block structure literally in their + // output, producing recursive nesting like [{'type':'text','text':'[{...}]'}]. + if (assistantText.length > 0) { + assistantMsg.content = assistantText; + } + + // Use the signature from the first thinking block if available (for llama.cpp server + gpt-oss) + let signature = nonEmptyThinkingBlocks[0].thinkingSignature; + if (model.provider === "opencode-go" && signature === "reasoning") { + signature = "reasoning_content"; + } + if (signature && signature.length > 0) { + (assistantMsg as any)[signature] = nonEmptyThinkingBlocks.map((block) => block.thinking).join("\n"); + } + } + } else if (assistantText.length > 0) { + // Always send assistant content as a plain string (OpenAI Chat Completions + // API standard format). Sending as an array of {type:"text", text:"..."} + // objects is non-standard and causes some models (e.g. DeepSeek V3.2 via + // NVIDIA NIM) to mirror the content-block structure literally in their + // output, producing recursive nesting like [{'type':'text','text':'[{...}]'}]. + assistantMsg.content = assistantText; + } + + const toolCalls = msg.content.filter(isToolCallBlock); + if (toolCalls.length > 0) { + assistantMsg.tool_calls = toolCalls.map((tc) => ({ + id: tc.id, + type: "function" as const, + function: { + name: tc.name, + arguments: JSON.stringify(tc.arguments), + }, + })); + const reasoningDetails = toolCalls + .filter((tc) => tc.thoughtSignature) + .map((tc) => { + try { + return JSON.parse(tc.thoughtSignature!); + } catch { + return null; + } + }) + .filter(Boolean); + if (reasoningDetails.length > 0) { + (assistantMsg as any).reasoning_details = reasoningDetails; + } + } + if ( + compat.requiresReasoningContentOnAssistantMessages && + model.reasoning && + (assistantMsg as { reasoning_content?: string }).reasoning_content === undefined + ) { + (assistantMsg as { reasoning_content?: string }).reasoning_content = ""; + } + // Skip assistant messages that have no content and no tool calls. + // Some providers require "either content or tool_calls, but not none". + // Other providers also don't accept empty assistant messages. + // This handles aborted assistant responses that got no content. + const content = assistantMsg.content; + const hasContent = + content !== null && + content !== undefined && + (typeof content === "string" ? content.length > 0 : content.length > 0); + if (!hasContent && !assistantMsg.tool_calls) { + continue; + } + params.push(assistantMsg); + } else if (msg.role === "toolResult") { + const imageBlocks: Array<{ type: "image_url"; image_url: { url: string } }> = []; + let j = i; + + for (; j < transformedMessages.length && transformedMessages[j].role === "toolResult"; j++) { + const toolMsg = transformedMessages[j] as ToolResultMessage; + + // Extract text and image content + const textResult = toolMsg.content + .filter(isTextContentBlock) + .map((block) => block.text) + .join("\n"); + const hasImages = toolMsg.content.some((c) => c.type === "image"); + + // Always send tool result with text (or placeholder if only images) + const hasText = textResult.length > 0; + // Some providers require the 'name' field in tool results + const toolResultMsg: ChatCompletionToolMessageParam = { + role: "tool", + content: sanitizeSurrogates(hasText ? textResult : "(see attached image)"), + tool_call_id: toolMsg.toolCallId, + }; + if (compat.requiresToolResultName && toolMsg.toolName) { + (toolResultMsg as any).name = toolMsg.toolName; + } + params.push(toolResultMsg); + + if (hasImages && model.input.includes("image")) { + for (const block of toolMsg.content) { + if (isImageContentBlock(block)) { + imageBlocks.push({ + type: "image_url", + image_url: { + url: `data:${block.mimeType};base64,${block.data}`, + }, + }); + } + } + } + } + + i = j - 1; + + if (imageBlocks.length > 0) { + if (compat.requiresAssistantAfterToolResult) { + params.push({ + role: "assistant", + content: "I have processed the tool results.", + }); + } + + params.push({ + role: "user", + content: [ + { + type: "text", + text: "Attached image(s) from tool result:", + }, + ...imageBlocks, + ], + }); + lastRole = "user"; + } else { + lastRole = "toolResult"; + } + continue; + } + + lastRole = msg.role; + } + + return params; +} + +function convertTools( + tools: Tool[], + compat: ResolvedOpenAICompletionsCompat, +): OpenAI.Chat.Completions.ChatCompletionTool[] { + return tools.map((tool) => ({ + type: "function", + function: { + name: tool.name, + description: tool.description, + parameters: tool.parameters as any, // TypeBox already generates JSON Schema + // Only include strict if provider supports it. Some reject unknown fields. + ...(compat.supportsStrictMode !== false && { strict: false }), + }, + })); +} + +function parseChunkUsage( + rawUsage: { + prompt_tokens?: number; + completion_tokens?: number; + prompt_cache_hit_tokens?: number; + prompt_tokens_details?: { cached_tokens?: number; cache_write_tokens?: number }; + }, + model: Model<"openai-completions">, +): AssistantMessage["usage"] { + const promptTokens = rawUsage.prompt_tokens || 0; + const cacheReadTokens = rawUsage.prompt_tokens_details?.cached_tokens ?? rawUsage.prompt_cache_hit_tokens ?? 0; + const cacheWriteTokens = rawUsage.prompt_tokens_details?.cache_write_tokens || 0; + + // Follow documented OpenAI/OpenRouter semantics: cached_tokens is cache-read + // tokens (hits). OpenAI does not document or emit cache_write_tokens, but + // OpenRouter-compatible providers can include it as a separate write count. + // OpenRouter's own provider/tests affirm the separate mapping: + // https://github.com/OpenRouterTeam/ai-sdk-provider/pull/409 + // Do not subtract writes from cached_tokens, otherwise spec-compliant + // providers are under-reported. DS4 mirrors this contract too: + // https://github.com/antirez/ds4/pull/29 + const input = Math.max(0, promptTokens - cacheReadTokens - cacheWriteTokens); + // OpenAI completion_tokens already includes reasoning_tokens. + const outputTokens = rawUsage.completion_tokens || 0; + const usage: AssistantMessage["usage"] = { + input, + output: outputTokens, + cacheRead: cacheReadTokens, + cacheWrite: cacheWriteTokens, + totalTokens: input + outputTokens + cacheReadTokens + cacheWriteTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + calculateCost(model, usage); + return usage; +} + +function mapStopReason(reason: ChatCompletionChunk.Choice["finish_reason"] | string): { + stopReason: StopReason; + errorMessage?: string; +} { + if (reason === null) return { stopReason: "stop" }; + switch (reason) { + case "stop": + case "end": + return { stopReason: "stop" }; + case "length": + return { stopReason: "length" }; + case "function_call": + case "tool_calls": + return { stopReason: "toolUse" }; + case "content_filter": + return { stopReason: "error", errorMessage: "Provider finish_reason: content_filter" }; + case "network_error": + return { stopReason: "error", errorMessage: "Provider finish_reason: network_error" }; + default: + return { + stopReason: "error", + errorMessage: `Provider finish_reason: ${reason}`, + }; + } +} + +/** + * Auto-detect compatibility settings from provider name and baseUrl. + * Used as the base when model.compat is not set; explicit model.compat + * entries override these detected values. + */ +function detectCompat(model: Model<"openai-completions">): ResolvedOpenAICompletionsCompat { + const provider = model.provider; + const baseUrl = model.baseUrl; + + const isZai = + provider === "zai" || + provider === "zai-coding-cn" || + baseUrl.includes("api.z.ai") || + baseUrl.includes("open.bigmodel.cn"); + const isTogether = + provider === "together" || baseUrl.includes("api.together.ai") || baseUrl.includes("api.together.xyz"); + const isMoonshot = provider === "moonshotai" || provider === "moonshotai-cn" || baseUrl.includes("api.moonshot."); + const isOpenRouter = provider === "openrouter" || baseUrl.includes("openrouter.ai"); + const isCloudflareWorkersAI = provider === "cloudflare-workers-ai" || baseUrl.includes("api.cloudflare.com"); + const isCloudflareAiGateway = provider === "cloudflare-ai-gateway" || baseUrl.includes("gateway.ai.cloudflare.com"); + const isNvidia = provider === "nvidia" || baseUrl.includes("integrate.api.nvidia.com"); + const isAntLing = provider === "ant-ling" || baseUrl.includes("api.ant-ling.com"); + + const isNonStandard = + isNvidia || + provider === "cerebras" || + baseUrl.includes("cerebras.ai") || + provider === "xai" || + baseUrl.includes("api.x.ai") || + isTogether || + baseUrl.includes("chutes.ai") || + baseUrl.includes("deepseek.com") || + isZai || + isMoonshot || + provider === "opencode" || + baseUrl.includes("opencode.ai") || + isCloudflareWorkersAI || + isCloudflareAiGateway || + isAntLing; + + const useMaxTokens = + baseUrl.includes("chutes.ai") || isMoonshot || isCloudflareAiGateway || isTogether || isNvidia || isAntLing; + + const isGrok = provider === "xai" || baseUrl.includes("api.x.ai"); + const isDeepSeek = provider === "deepseek" || baseUrl.includes("deepseek.com"); + const isOpenRouterDeveloperRoleModel = + isOpenRouter && (model.id.startsWith("anthropic/") || model.id.startsWith("openai/")); + const cacheControlFormat = provider === "openrouter" && model.id.startsWith("anthropic/") ? "anthropic" : undefined; + + return { + supportsStore: !isNonStandard, + supportsDeveloperRole: isOpenRouterDeveloperRoleModel || (!isNonStandard && !isOpenRouter), + supportsReasoningEffort: + !isGrok && !isZai && !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia && !isAntLing, + supportsUsageInStreaming: true, + maxTokensField: useMaxTokens ? "max_tokens" : "max_completion_tokens", + requiresToolResultName: false, + requiresAssistantAfterToolResult: false, + requiresThinkingAsText: false, + requiresReasoningContentOnAssistantMessages: isDeepSeek, + thinkingFormat: isDeepSeek + ? "deepseek" + : isZai + ? "zai" + : isTogether + ? "together" + : isAntLing + ? "ant-ling" + : isOpenRouter + ? "openrouter" + : "openai", + openRouterRouting: {}, + vercelGatewayRouting: {}, + chatTemplateKwargs: {}, + zaiToolStream: false, + supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia, + cacheControlFormat, + sendSessionAffinityHeaders: false, + supportsLongCacheRetention: !( + isTogether || + isCloudflareWorkersAI || + isCloudflareAiGateway || + isNvidia || + isAntLing + ), + }; +} + +/** + * Get resolved compatibility settings for a model. + * Auto-detects from provider/URL then overrides with explicit model.compat. + */ +function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletionsCompat { + const detected = detectCompat(model); + if (!model.compat) return detected; + + return { + supportsStore: model.compat.supportsStore ?? detected.supportsStore, + supportsDeveloperRole: model.compat.supportsDeveloperRole ?? detected.supportsDeveloperRole, + supportsReasoningEffort: model.compat.supportsReasoningEffort ?? detected.supportsReasoningEffort, + supportsUsageInStreaming: model.compat.supportsUsageInStreaming ?? detected.supportsUsageInStreaming, + maxTokensField: model.compat.maxTokensField ?? detected.maxTokensField, + requiresToolResultName: model.compat.requiresToolResultName ?? detected.requiresToolResultName, + requiresAssistantAfterToolResult: + model.compat.requiresAssistantAfterToolResult ?? detected.requiresAssistantAfterToolResult, + requiresThinkingAsText: model.compat.requiresThinkingAsText ?? detected.requiresThinkingAsText, + requiresReasoningContentOnAssistantMessages: + model.compat.requiresReasoningContentOnAssistantMessages ?? + detected.requiresReasoningContentOnAssistantMessages, + thinkingFormat: model.compat.thinkingFormat ?? detected.thinkingFormat, + openRouterRouting: model.compat.openRouterRouting ?? {}, + vercelGatewayRouting: model.compat.vercelGatewayRouting ?? detected.vercelGatewayRouting, + chatTemplateKwargs: model.compat.chatTemplateKwargs ?? detected.chatTemplateKwargs, + zaiToolStream: model.compat.zaiToolStream ?? detected.zaiToolStream, + supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode, + cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat, + sendSessionAffinityHeaders: model.compat.sendSessionAffinityHeaders ?? detected.sendSessionAffinityHeaders, + supportsLongCacheRetention: model.compat.supportsLongCacheRetention ?? detected.supportsLongCacheRetention, + }; +} diff --git a/cactus-code/packages/ai/src/api/openai-prompt-cache.ts b/cactus-code/packages/ai/src/api/openai-prompt-cache.ts new file mode 100644 index 000000000..1c3355d52 --- /dev/null +++ b/cactus-code/packages/ai/src/api/openai-prompt-cache.ts @@ -0,0 +1,8 @@ +export const OPENAI_PROMPT_CACHE_KEY_MAX_LENGTH = 64; + +export function clampOpenAIPromptCacheKey(key: string | undefined): string | undefined { + if (key === undefined) return undefined; + const chars = Array.from(key); + if (chars.length <= OPENAI_PROMPT_CACHE_KEY_MAX_LENGTH) return key; + return chars.slice(0, OPENAI_PROMPT_CACHE_KEY_MAX_LENGTH).join(""); +} diff --git a/cactus-code/packages/ai/src/api/simple-options.ts b/cactus-code/packages/ai/src/api/simple-options.ts new file mode 100644 index 000000000..773e8d95a --- /dev/null +++ b/cactus-code/packages/ai/src/api/simple-options.ts @@ -0,0 +1,54 @@ +import type { Api, Model, SimpleStreamOptions, StreamOptions, ThinkingBudgets, ThinkingLevel } from "../types.ts"; + +export function buildBaseOptions(_model: Model, options?: SimpleStreamOptions, apiKey?: string): StreamOptions { + return { + temperature: options?.temperature, + maxTokens: options?.maxTokens, + signal: options?.signal, + apiKey: apiKey || options?.apiKey, + transport: options?.transport, + cacheRetention: options?.cacheRetention, + sessionId: options?.sessionId, + headers: options?.headers, + onPayload: options?.onPayload, + onResponse: options?.onResponse, + timeoutMs: options?.timeoutMs, + websocketConnectTimeoutMs: options?.websocketConnectTimeoutMs, + maxRetries: options?.maxRetries, + maxRetryDelayMs: options?.maxRetryDelayMs, + metadata: options?.metadata, + env: options?.env, + }; +} + +export function clampReasoning(effort: ThinkingLevel | undefined): Exclude | undefined { + return effort === "xhigh" ? "high" : effort; +} + +export function adjustMaxTokensForThinking( + // Undefined means no explicit caller cap. Use the model cap and fit thinking inside it. + baseMaxTokens: number | undefined, + modelMaxTokens: number, + reasoningLevel: ThinkingLevel, + customBudgets?: ThinkingBudgets, +): { maxTokens: number; thinkingBudget: number } { + const defaultBudgets: ThinkingBudgets = { + minimal: 1024, + low: 2048, + medium: 8192, + high: 16384, + }; + const budgets = { ...defaultBudgets, ...customBudgets }; + + const minOutputTokens = 1024; + const level = clampReasoning(reasoningLevel)!; + let thinkingBudget = budgets[level]!; + const maxTokens = + baseMaxTokens === undefined ? modelMaxTokens : Math.min(baseMaxTokens + thinkingBudget, modelMaxTokens); + + if (maxTokens <= thinkingBudget) { + thinkingBudget = Math.max(0, maxTokens - minOutputTokens); + } + + return { maxTokens, thinkingBudget }; +} diff --git a/cactus-code/packages/ai/src/api/transform-messages.ts b/cactus-code/packages/ai/src/api/transform-messages.ts new file mode 100644 index 000000000..4029a0643 --- /dev/null +++ b/cactus-code/packages/ai/src/api/transform-messages.ts @@ -0,0 +1,220 @@ +import type { + Api, + AssistantMessage, + ImageContent, + Message, + Model, + TextContent, + ToolCall, + ToolResultMessage, +} from "../types.ts"; + +const NON_VISION_USER_IMAGE_PLACEHOLDER = "(image omitted: model does not support images)"; +const NON_VISION_TOOL_IMAGE_PLACEHOLDER = "(tool image omitted: model does not support images)"; + +function replaceImagesWithPlaceholder(content: (TextContent | ImageContent)[], placeholder: string): TextContent[] { + const result: TextContent[] = []; + let previousWasPlaceholder = false; + + for (const block of content) { + if (block.type === "image") { + if (!previousWasPlaceholder) { + result.push({ type: "text", text: placeholder }); + } + previousWasPlaceholder = true; + continue; + } + + result.push(block); + previousWasPlaceholder = block.text === placeholder; + } + + return result; +} + +function downgradeUnsupportedImages(messages: Message[], model: Model): Message[] { + if (model.input.includes("image")) { + return messages; + } + + return messages.map((msg) => { + if (msg.role === "user" && Array.isArray(msg.content)) { + return { + ...msg, + content: replaceImagesWithPlaceholder(msg.content, NON_VISION_USER_IMAGE_PLACEHOLDER), + }; + } + + if (msg.role === "toolResult") { + return { + ...msg, + content: replaceImagesWithPlaceholder(msg.content, NON_VISION_TOOL_IMAGE_PLACEHOLDER), + }; + } + + return msg; + }); +} + +/** + * Normalize tool call ID for cross-provider compatibility. + * OpenAI Responses API generates IDs that are 450+ chars with special characters like `|`. + * Anthropic APIs require IDs matching ^[a-zA-Z0-9_-]+$ (max 64 chars). + */ +export function transformMessages( + messages: Message[], + model: Model, + normalizeToolCallId?: (id: string, model: Model, source: AssistantMessage) => string, +): Message[] { + // Build a map of original tool call IDs to normalized IDs + const toolCallIdMap = new Map(); + const imageAwareMessages = downgradeUnsupportedImages(messages, model); + + // First pass: transform messages (unsupported image downgrade, thinking blocks, tool call ID normalization) + const transformed = imageAwareMessages.map((msg) => { + // User messages pass through unchanged + if (msg.role === "user") { + return msg; + } + + // Handle toolResult messages - normalize toolCallId if we have a mapping + if (msg.role === "toolResult") { + const normalizedId = toolCallIdMap.get(msg.toolCallId); + if (normalizedId && normalizedId !== msg.toolCallId) { + return { ...msg, toolCallId: normalizedId }; + } + return msg; + } + + // Assistant messages need transformation check + if (msg.role === "assistant") { + const assistantMsg = msg as AssistantMessage; + const isSameModel = + assistantMsg.provider === model.provider && + assistantMsg.api === model.api && + assistantMsg.model === model.id; + + const transformedContent = assistantMsg.content.flatMap((block) => { + if (block.type === "thinking") { + // Redacted thinking is opaque encrypted content, only valid for the same model. + // Drop it for cross-model to avoid API errors. + if (block.redacted) { + return isSameModel ? block : []; + } + // For same model: keep thinking blocks with signatures (needed for replay) + // even if the thinking text is empty (OpenAI encrypted reasoning) + if (isSameModel && block.thinkingSignature) return block; + // Skip empty thinking blocks, convert others to plain text + if (!block.thinking || block.thinking.trim() === "") return []; + if (isSameModel) return block; + return { + type: "text" as const, + text: block.thinking, + }; + } + + if (block.type === "text") { + if (isSameModel) return block; + return { + type: "text" as const, + text: block.text, + }; + } + + if (block.type === "toolCall") { + const toolCall = block as ToolCall; + let normalizedToolCall: ToolCall = toolCall; + + if (!isSameModel && toolCall.thoughtSignature) { + normalizedToolCall = { ...toolCall }; + delete (normalizedToolCall as { thoughtSignature?: string }).thoughtSignature; + } + + if (!isSameModel && normalizeToolCallId) { + const normalizedId = normalizeToolCallId(toolCall.id, model, assistantMsg); + if (normalizedId !== toolCall.id) { + toolCallIdMap.set(toolCall.id, normalizedId); + normalizedToolCall = { ...normalizedToolCall, id: normalizedId }; + } + } + + return normalizedToolCall; + } + + return block; + }); + + return { + ...assistantMsg, + content: transformedContent, + }; + } + return msg; + }); + + // Second pass: insert synthetic empty tool results for orphaned tool calls + // This preserves thinking signatures and satisfies API requirements + const result: Message[] = []; + let pendingToolCalls: ToolCall[] = []; + let existingToolResultIds = new Set(); + const insertSyntheticToolResults = () => { + if (pendingToolCalls.length > 0) { + for (const tc of pendingToolCalls) { + if (!existingToolResultIds.has(tc.id)) { + result.push({ + role: "toolResult", + toolCallId: tc.id, + toolName: tc.name, + content: [{ type: "text", text: "No result provided" }], + isError: true, + timestamp: Date.now(), + } as ToolResultMessage); + } + } + pendingToolCalls = []; + existingToolResultIds = new Set(); + } + }; + + for (let i = 0; i < transformed.length; i++) { + const msg = transformed[i]; + + if (msg.role === "assistant") { + // If we have pending orphaned tool calls from a previous assistant, insert synthetic results now + insertSyntheticToolResults(); + + // Skip errored/aborted assistant messages entirely. + // These are incomplete turns that shouldn't be replayed: + // - May have partial content (reasoning without message, incomplete tool calls) + // - Replaying them can cause API errors (e.g., OpenAI "reasoning without following item") + // - The model should retry from the last valid state + const assistantMsg = msg as AssistantMessage; + if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") { + continue; + } + + // Track tool calls from this assistant message + const toolCalls = assistantMsg.content.filter((b) => b.type === "toolCall") as ToolCall[]; + if (toolCalls.length > 0) { + pendingToolCalls = toolCalls; + existingToolResultIds = new Set(); + } + + result.push(msg); + } else if (msg.role === "toolResult") { + existingToolResultIds.add(msg.toolCallId); + result.push(msg); + } else if (msg.role === "user") { + // User message interrupts tool flow - insert synthetic results for orphaned calls + insertSyntheticToolResults(); + result.push(msg); + } else { + result.push(msg); + } + } + + // If the conversation ends with unresolved tool calls, synthesize results now. + insertSyntheticToolResults(); + + return result; +} diff --git a/cactus-code/packages/ai/src/auth/context.ts b/cactus-code/packages/ai/src/auth/context.ts new file mode 100644 index 000000000..30e088bfa --- /dev/null +++ b/cactus-code/packages/ai/src/auth/context.ts @@ -0,0 +1,45 @@ +import type { AuthContext } from "./types.ts"; + +interface NodeFsModule { + access(path: string): Promise; +} + +interface NodeOsModule { + homedir(): string; +} + +// Variable specifier so browser bundlers do not try to resolve node builtins. +const importNodeModule = (specifier: string): Promise => import(specifier); + +function getProcessEnv(): Record | undefined { + const proc = (globalThis as { process?: { env?: Record } }).process; + return proc?.env; +} + +/** + * Default auth context: env vars from `process.env` (undefined in browsers), + * file existence via node:fs (always false in browsers). + */ +export function defaultProviderAuthContext(): AuthContext { + return { + async env(name: string): Promise { + const value = getProcessEnv()?.[name]; + return typeof value === "string" && value.trim().length > 0 ? value : undefined; + }, + + async fileExists(path: string): Promise { + try { + const fs = (await importNodeModule("node:fs/promises")) as NodeFsModule; + let resolved = path; + if (resolved.startsWith("~")) { + const os = (await importNodeModule("node:os")) as NodeOsModule; + resolved = os.homedir() + resolved.slice(1); + } + await fs.access(resolved); + return true; + } catch { + return false; + } + }, + }; +} diff --git a/cactus-code/packages/ai/src/auth/credential-store.ts b/cactus-code/packages/ai/src/auth/credential-store.ts new file mode 100644 index 000000000..beeb9d858 --- /dev/null +++ b/cactus-code/packages/ai/src/auth/credential-store.ts @@ -0,0 +1,47 @@ +import type { Credential, CredentialStore } from "./types.ts"; + +/** + * Default in-memory credential store. Apps inject persistent stores. + * Keyed by `Provider.id`, one credential per provider; see `CredentialStore`. + * Writes are serialized per provider through a promise chain. + */ +export class InMemoryCredentialStore implements CredentialStore { + private credentials = new Map(); + private chains = new Map>(); + + /** Serialize tasks per provider id. */ + private enqueue(providerId: string, task: () => Promise): Promise { + const previous = this.chains.get(providerId) ?? Promise.resolve(); + const next = (async () => { + await previous.catch(() => {}); + return task(); + })(); + this.chains.set( + providerId, + next.catch(() => {}), + ); + return next; + } + + async read(providerId: string): Promise { + return this.credentials.get(providerId); + } + + modify( + providerId: string, + fn: (current: Credential | undefined) => Promise, + ): Promise { + return this.enqueue(providerId, async () => { + const current = this.credentials.get(providerId); + const next = await fn(current); + if (next !== undefined) this.credentials.set(providerId, next); + return next ?? current; + }); + } + + delete(providerId: string): Promise { + return this.enqueue(providerId, async () => { + this.credentials.delete(providerId); + }); + } +} diff --git a/cactus-code/packages/ai/src/auth/helpers.ts b/cactus-code/packages/ai/src/auth/helpers.ts new file mode 100644 index 000000000..d9a34ad27 --- /dev/null +++ b/cactus-code/packages/ai/src/auth/helpers.ts @@ -0,0 +1,46 @@ +import type { ApiKeyAuth, OAuthAuth } from "./types.ts"; + +/** + * Standard api-key auth: a stored credential key wins, otherwise the first + * set env var resolves. Includes a `login` that prompts for the key. + * Providers with non-standard resolution (provider env, ambient files, IAM) + * write their own `ApiKeyAuth`. + */ +export function envApiKeyAuth(name: string, envVars: readonly string[]): ApiKeyAuth { + return { + name, + login: async (callbacks) => { + const key = await callbacks.prompt({ type: "secret", message: `Enter ${name}` }); + return { type: "api_key", key }; + }, + resolve: async ({ ctx, credential }) => { + if (credential?.key) return { auth: { apiKey: credential.key }, source: "stored credential" }; + for (const envVar of envVars) { + const value = await ctx.env(envVar); + if (value) return { auth: { apiKey: value }, source: envVar }; + } + return undefined; + }, + }; +} + +/** + * Wraps a dynamically imported `OAuthAuth` so provider definitions can + * advertise OAuth without importing the implementation. The flow loads on + * first `login`/`refresh`/`toAuth` call; callers keep Node-only flow code out + * of bundles by loading through a bundler-opaque dynamic import (variable + * specifier, see the bedrock lazy wrapper). + */ +export function lazyOAuth(input: { name: string; load: () => Promise }): OAuthAuth { + let promise: Promise | undefined; + const loaded = () => { + promise ??= input.load(); + return promise; + }; + return { + name: input.name, + login: async (callbacks) => (await loaded()).login(callbacks), + refresh: async (credential) => (await loaded()).refresh(credential), + toAuth: async (credential) => (await loaded()).toAuth(credential), + }; +} diff --git a/cactus-code/packages/ai/src/auth/resolve.ts b/cactus-code/packages/ai/src/auth/resolve.ts new file mode 100644 index 000000000..81d7a2707 --- /dev/null +++ b/cactus-code/packages/ai/src/auth/resolve.ts @@ -0,0 +1,141 @@ +import type { Api, ImagesApi, ImagesModel, Model, ProviderEnv } from "../types.ts"; +import type { + ApiKeyAuth, + ApiKeyCredential, + AuthContext, + AuthResult, + Credential, + CredentialStore, + OAuthAuth, + OAuthCredential, + ProviderAuth, +} from "./types.ts"; + +export type ModelsErrorCode = "model_source" | "model_validation" | "provider" | "stream" | "auth" | "oauth"; + +export interface AuthResolutionOverrides { + apiKey?: string; + env?: ProviderEnv; +} + +export class ModelsError extends Error { + readonly code: ModelsErrorCode; + + constructor(code: ModelsErrorCode, message: string, options?: { cause?: unknown }) { + super(message, options); + this.name = "ModelsError"; + this.code = code; + } +} + +/** Model shape auth resolution receives: chat or image-generation models. */ +export type AuthModel = Model | ImagesModel; + +/** + * Auth resolution shared by the `Models` and `ImagesModels` collections. + * A stored credential owns the provider: ambient/env is consulted only when + * nothing is stored. No silent env fallback after a failed refresh or for a + * credential type without a matching handler. + */ +export async function resolveProviderAuth( + provider: { id: string; auth: ProviderAuth }, + model: AuthModel, + credentials: CredentialStore, + authContext: AuthContext, + overrides?: AuthResolutionOverrides, +): Promise { + const requestAuthContext = overrides?.env ? overlayEnvAuthContext(authContext, overrides.env) : authContext; + + if (overrides?.apiKey !== undefined && provider.auth.apiKey) { + return resolveApiKey(requestAuthContext, provider.auth.apiKey, model, { + type: "api_key", + key: overrides.apiKey, + env: overrides.env, + }); + } + + const stored = await readCredential(credentials, provider.id); + if (stored) { + if (stored.type === "oauth" && provider.auth.oauth) { + return resolveStoredOAuth(credentials, provider.id, provider.auth.oauth, stored); + } + if (stored.type === "api_key" && provider.auth.apiKey) { + const credential = overrides?.env ? { ...stored, env: { ...stored.env, ...overrides.env } } : stored; + return resolveApiKey(requestAuthContext, provider.auth.apiKey, model, credential); + } + return undefined; + } + + // Ambient (env vars, AWS profiles, ADC files). + return provider.auth.apiKey ? resolveApiKey(requestAuthContext, provider.auth.apiKey, model, undefined) : undefined; +} + +function overlayEnvAuthContext(base: AuthContext, env: ProviderEnv): AuthContext { + return { + env: async (name) => env[name] || (await base.env(name)), + fileExists: (path) => base.fileExists(path), + }; +} + +/** + * OAuth resolution with double-checked locking (same pattern as today's + * AuthStorage): valid tokens cost zero locks; expired tokens lock, re-check + * expiry under the lock, refresh once globally, and persist the rotated + * credential before release. + */ +async function resolveStoredOAuth( + credentials: CredentialStore, + providerId: string, + oauth: OAuthAuth, + stored: OAuthCredential, +): Promise { + let credential = stored; + + if (Date.now() >= credential.expires) { + // Optimistic check said expired; the authoritative check runs under the lock. + let post: Credential | undefined; + try { + post = await credentials.modify(providerId, async (current) => { + if (current?.type !== "oauth") return undefined; // logged out meanwhile + if (Date.now() < current.expires) return undefined; // another process/request refreshed + try { + return await oauth.refresh(current); + } catch (error) { + throw new ModelsError("oauth", `OAuth refresh failed for ${providerId}`, { cause: error }); + } + }); + } catch (error) { + if (error instanceof ModelsError) throw error; + throw new ModelsError("auth", `Credential store modify failed for ${providerId}`, { cause: error }); + } + if (post?.type !== "oauth") return undefined; // logged out meanwhile + credential = post; + } + + try { + return { auth: await oauth.toAuth(credential), source: "OAuth" }; + } catch (error) { + throw new ModelsError("oauth", `OAuth auth derivation failed for ${providerId}`, { cause: error }); + } +} + +async function resolveApiKey( + authContext: AuthContext, + apiKey: ApiKeyAuth, + model: AuthModel, + credential: ApiKeyCredential | undefined, +): Promise { + try { + return await apiKey.resolve({ model, ctx: authContext, credential }); + } catch (error) { + throw new ModelsError("auth", `API key auth failed for provider ${model.provider}`, { cause: error }); + } +} + +async function readCredential(credentials: CredentialStore, providerId: string): Promise { + try { + return await credentials.read(providerId); + } catch (error) { + throw new ModelsError("auth", `Credential store read failed for ${providerId}`, { cause: error }); + } +} diff --git a/cactus-code/packages/ai/src/auth/types.ts b/cactus-code/packages/ai/src/auth/types.ts new file mode 100644 index 000000000..9710dbcbc --- /dev/null +++ b/cactus-code/packages/ai/src/auth/types.ts @@ -0,0 +1,182 @@ +import type { Api, ImagesApi, ImagesModel, Model, ProviderEnv, ProviderHeaders } from "../types.ts"; +import type { OAuthCredentials } from "../utils/oauth/types.ts"; + +/** + * Request auth for a single model request. If a value cannot be expressed as + * `apiKey`, `headers`, or `baseUrl`, it is provider config, not auth. + */ +export interface ModelAuth { + apiKey?: string; + headers?: ProviderHeaders; + baseUrl?: string; +} + +/** + * Stored api-key credential. `env` holds provider-scoped environment/config + * values such as Cloudflare account/gateway ids. + */ +export interface ApiKeyCredential { + type: "api_key"; + key?: string; + env?: ProviderEnv; +} + +/** Stored OAuth credential (`access`, `refresh`, `expires` from OAuthCredentials). */ +export interface OAuthCredential extends OAuthCredentials { + type: "oauth"; +} + +/** One type-tagged credential per provider — the shape of today's auth.json. */ +export type Credential = ApiKeyCredential | OAuthCredential; + +/** + * App-owned credential storage, keyed by `Provider.id`, one credential per + * provider. `modify` is the only write path, so every mutation is a + * serialized read-modify-write; `Models.getAuth()` runs OAuth refresh inside + * `modify` so concurrent requests cannot double-refresh a rotated token. The + * app persists a credential after login via + * `modify(provider.id, async () => credential)`. Login/logout orchestration + * is app-owned. + * + * Error semantics: `read` resolves `undefined` for missing entries. Methods + * reject only on storage failure; `Models` wraps such rejections in + * `ModelsError` with code "auth". Best-effort stores that serve an in-memory + * view and record persistence errors internally (like coding-agent's + * AuthStorage) are valid implementations. + */ +export interface CredentialStore { + /** + * Read the stored credential, possibly expired. Display/status use; + * resolved request auth comes from `Models.getAuth()`. + */ + read(providerId: string): Promise; + + /** + * Serialized write — the only write path. `fn` sees the current credential + * because correct writes (refresh, login-during-refresh) depend on it; + * return the new credential, or undefined to leave the entry unchanged. + * Mutual exclusion per provider id, cross-process too where the backing + * store supports it (e.g. a file lock). Resolves with the post-write + * credential. Rejections from `fn` propagate. + */ + modify( + providerId: string, + fn: (current: Credential | undefined) => Promise, + ): Promise; + + /** Remove a credential (logout). Implementations serialize this against `modify`. */ + delete(providerId: string): Promise; +} + +/** Environment access for auth resolution. Injectable for tests and browsers. */ +export interface AuthContext { + env(name: string): Promise; + /** Check whether a file exists. Supports a leading `~`. Always false in browsers. */ + fileExists(path: string): Promise; +} + +/** Result of resolving auth for a model. */ +export interface AuthResult { + auth: ModelAuth; + /** Provider-scoped environment/config values resolved from credentials and ambient context. */ + env?: ProviderEnv; + /** Human-readable label for status UI: "ANTHROPIC_API_KEY", "OAuth", "~/.aws/credentials". */ + source?: string; +} + +/** + * Prompt shown to the user during login. `signal` lets the flow cancel a + * pending prompt when an out-of-band event resolves the step, e.g. a + * `manual_code` prompt raced against a callback server, aborted when the + * callback wins. + */ +export type AuthPrompt = { signal?: AbortSignal } & ( + | { type: "text"; message: string; placeholder?: string } + | { type: "secret"; message: string; placeholder?: string } + | { type: "select"; message: string; options: readonly { id: string; label: string; description?: string }[] } + | { type: "manual_code"; message: string; placeholder?: string } +); + +export type AuthEvent = + | { type: "auth_url"; url: string; instructions?: string } + | { + type: "device_code"; + userCode: string; + verificationUri: string; + intervalSeconds?: number; + expiresInSeconds?: number; + } + | { type: "progress"; message: string }; + +/** + * Login interaction callbacks serving both api-key and OAuth flows. + * + * `prompt()` returns the entered/selected string (`select` returns the option + * id). Rejects on cancel/abort. `signal` aborts the whole login flow; + * per-prompt cancellation uses `AuthPrompt.signal`. + */ +export interface AuthLoginCallbacks { + signal?: AbortSignal; + + prompt(prompt: AuthPrompt): Promise; + notify(event: AuthEvent): void; +} + +/** + * Api-key auth: stored key/provider env plus ambient sources (env vars, AWS + * profiles, ADC files). Ambient-only providers omit `login`. + */ +export interface ApiKeyAuth { + /** Display name, e.g. "Anthropic API key". */ + name: string; + + /** Interactive setup (prompt for key/provider env). Absent = ambient-only. */ + login?(callbacks: AuthLoginCallbacks): Promise; + + /** + * Resolve auth from the stored credential and/or ambient sources, merging + * per field (`credential.key ?? env("...")`, `credential.env?.NAME ?? env("...")`). + * undefined = not configured. Receives the chat or image-generation model + * the request is for (both carry `provider` and `baseUrl`). + */ + resolve(input: { + model: Model | ImagesModel; + ctx: AuthContext; + credential?: ApiKeyCredential; + }): Promise; +} + +/** + * OAuth auth. The `refresh`/`toAuth` split lets `Models` own the locked + * refresh pattern: `refresh` produces a credential, `toAuth` derives request + * auth from whatever credential ends up stored. + */ +export interface OAuthAuth { + /** Display name, e.g. "Anthropic (Claude Pro/Max)". */ + name: string; + + login(callbacks: AuthLoginCallbacks): Promise; + + /** + * Exchange the refresh token. Network call; throws on failure + * (invalid_grant etc.). `Models` runs this under the store lock. + */ + refresh(credential: OAuthCredential): Promise; + + /** + * Side-effect-free derivation of request auth from a valid credential. + * Covers per-credential baseUrl (GitHub Copilot). Async so lazy wrappers + * can load the implementation on first use. + */ + toAuth(credential: OAuthCredential): Promise; +} + +/** + * Provider auth. At least one of `apiKey`/`oauth` must be present: even + * ambient-credential providers and keyless local servers provide `apiKey` + * auth whose `resolve()` reports whether the provider is configured. + */ +export interface ProviderAuth { + apiKey?: ApiKeyAuth; + oauth?: OAuthAuth; +} diff --git a/cactus-code/packages/ai/src/compat.ts b/cactus-code/packages/ai/src/compat.ts new file mode 100644 index 000000000..e5d3f7761 --- /dev/null +++ b/cactus-code/packages/ai/src/compat.ts @@ -0,0 +1,246 @@ +/** + * Temporary compatibility entrypoint preserving the old global pi-ai API + * surface: api-dispatch `stream()`/`complete()` with env API key injection, + * the api-registry, generated catalog reads (`getModel`/`getModels`/ + * `getProviders`), per-API lazy stream wrappers, and image generation. + * + * Existing apps switch imports from "@earendil-works/pi-ai" to + * "@earendil-works/pi-ai/compat" unchanged; new code uses `createModels()` + * and the provider factories. This module is deleted with the coding-agent + * ModelManager migration. + */ + +export * from "./api/openai-completions.lazy.ts"; +export * from "./env-api-keys.ts"; +export * from "./index.ts"; + +import { openAICompletionsApi } from "./api/openai-completions.lazy.ts"; +import { getEnvApiKey } from "./env-api-keys.ts"; +import { builtinModels, getBuiltinModel, getBuiltinModels, getBuiltinProviders } from "./providers/all.ts"; +import { createFauxCore, type FauxProviderRegistration, type RegisterFauxProviderOptions } from "./providers/faux.ts"; +import type { + Api, + ApiStreamOptions, + AssistantMessage, + AssistantMessageEventStream, + Context, + Model, + ProviderStreamOptions, + ProviderStreams, + SimpleStreamOptions, + StreamFunction, + StreamOptions, +} from "./types.ts"; + +/** @deprecated Static catalog read. Use `getBuiltinModel` from "@earendil-works/pi-ai/providers/all" or `Models.getModel()`. */ +export const getModel = getBuiltinModel; + +/** @deprecated Static catalog read. Use `getBuiltinModels` from "@earendil-works/pi-ai/providers/all" or `Models.getModels()`. */ +export const getModels = getBuiltinModels; + +/** @deprecated Static catalog read. Use `getBuiltinProviders` from "@earendil-works/pi-ai/providers/all" or `Models.getProviders()`. */ +export const getProviders = getBuiltinProviders; + +export type ApiStreamFunction = ( + model: Model, + context: Context, + options?: StreamOptions, +) => AssistantMessageEventStream; + +export type ApiStreamSimpleFunction = ( + model: Model, + context: Context, + options?: SimpleStreamOptions, +) => AssistantMessageEventStream; + +export interface ApiProvider { + api: TApi; + stream: StreamFunction; + streamSimple: StreamFunction; +} + +interface ApiProviderInternal { + api: Api; + stream: ApiStreamFunction; + streamSimple: ApiStreamSimpleFunction; +} + +type RegisteredApiProvider = { + provider: ApiProviderInternal; + sourceId?: string; +}; + +const apiProviderRegistry = new Map(); + +function wrapStream( + api: TApi, + stream: StreamFunction, +): ApiStreamFunction { + return (model, context, options) => { + if (model.api !== api) { + throw new Error(`Mismatched api: ${model.api} expected ${api}`); + } + return stream(model as Model, context, options as TOptions); + }; +} + +function wrapStreamSimple( + api: TApi, + streamSimple: StreamFunction, +): ApiStreamSimpleFunction { + return (model, context, options) => { + if (model.api !== api) { + throw new Error(`Mismatched api: ${model.api} expected ${api}`); + } + return streamSimple(model as Model, context, options); + }; +} + +export function registerApiProvider( + provider: ApiProvider, + sourceId?: string, +): void { + apiProviderRegistry.set(provider.api, { + provider: { + api: provider.api, + stream: wrapStream(provider.api, provider.stream), + streamSimple: wrapStreamSimple(provider.api, provider.streamSimple), + }, + sourceId, + }); +} + +export function getApiProvider(api: Api): ApiProviderInternal | undefined { + return apiProviderRegistry.get(api)?.provider; +} + +export function getApiProviders(): ApiProviderInternal[] { + return Array.from(apiProviderRegistry.values(), (entry) => entry.provider); +} + +export function unregisterApiProviders(sourceId: string): void { + for (const [api, entry] of apiProviderRegistry.entries()) { + if (entry.sourceId === sourceId) { + apiProviderRegistry.delete(api); + } + } +} + +function clearApiProviders(): void { + apiProviderRegistry.clear(); +} + +export function registerFauxProvider(options: RegisterFauxProviderOptions = {}): FauxProviderRegistration { + const core = createFauxCore(options); + const sourceId = `faux-provider-${Math.random().toString(36).slice(2, 10)}`; + registerApiProvider({ api: core.api, stream: core.stream, streamSimple: core.streamSimple }, sourceId); + return { + api: core.api, + models: core.models, + getModel: core.getModel, + state: core.state, + setResponses: core.setResponses, + appendResponses: core.appendResponses, + getPendingResponseCount: core.getPendingResponseCount, + unregister() { + unregisterApiProviders(sourceId); + }, + }; +} + +const BUILTIN_APIS: [Api, ProviderStreams][] = [["openai-completions", openAICompletionsApi()]]; + +const builtinApiProviderInstances = new Map>(); + +/** + * Registers the builtin API implementations into the api-registry without + * clobbering existing entries: compat may load after a test or extension has + * already registered an override for a builtin api id. + */ +export function registerBuiltInApiProviders(): void { + for (const [api, streams] of BUILTIN_APIS) { + if (!getApiProvider(api)) { + registerApiProvider({ api, stream: streams.stream, streamSimple: streams.streamSimple }); + } + builtinApiProviderInstances.set(api, getApiProvider(api)); + } +} + +export function resetApiProviders(): void { + clearApiProviders(); + builtinApiProviderInstances.clear(); + registerBuiltInApiProviders(); +} + +registerBuiltInApiProviders(); + +const compatModels = builtinModels(); + +function hasExplicitApiKey(apiKey: string | undefined): apiKey is string { + return typeof apiKey === "string" && apiKey.trim().length > 0; +} + +function withEnvApiKey( + model: Model, + options: TOptions | undefined, +): TOptions | undefined { + if (hasExplicitApiKey(options?.apiKey)) return options; + const apiKey = getEnvApiKey(model.provider, options?.env); + if (!apiKey) return options; + return { ...options, apiKey } as TOptions; +} + +function shouldUseBuiltinModels(model: Model): boolean { + const builtin = compatModels.getModel(model.provider, model.id); + return builtin?.api === model.api && getApiProvider(model.api) === builtinApiProviderInstances.get(model.api); +} + +function resolveApiProvider(api: Api) { + const provider = getApiProvider(api); + if (!provider) { + throw new Error(`No API provider registered for api: ${api}`); + } + return provider; +} + +export function stream( + model: Model, + context: Context, + options?: ProviderStreamOptions, +): AssistantMessageEventStream { + if (shouldUseBuiltinModels(model)) { + return compatModels.stream(model, context, options as ApiStreamOptions | undefined); + } + const provider = resolveApiProvider(model.api); + return provider.stream(model, context, withEnvApiKey(model, options) as StreamOptions); +} + +export async function complete( + model: Model, + context: Context, + options?: ProviderStreamOptions, +): Promise { + const s = stream(model, context, options); + return s.result(); +} + +export function streamSimple( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): AssistantMessageEventStream { + if (shouldUseBuiltinModels(model)) { + return compatModels.streamSimple(model, context, options); + } + const provider = resolveApiProvider(model.api); + return provider.streamSimple(model, context, withEnvApiKey(model, options)); +} + +export async function completeSimple( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): Promise { + const s = streamSimple(model, context, options); + return s.result(); +} diff --git a/cactus-code/packages/ai/src/env-api-keys.ts b/cactus-code/packages/ai/src/env-api-keys.ts new file mode 100644 index 000000000..7bd955e1f --- /dev/null +++ b/cactus-code/packages/ai/src/env-api-keys.ts @@ -0,0 +1,177 @@ +// NEVER convert to top-level imports - breaks browser/Vite builds +let _existsSync: typeof import("node:fs").existsSync | null = null; +let _homedir: typeof import("node:os").homedir | null = null; +let _join: typeof import("node:path").join | null = null; + +type DynamicImport = (specifier: string) => Promise; + +const dynamicImport: DynamicImport = (specifier) => import(specifier); +const NODE_FS_SPECIFIER = "node:" + "fs"; +const NODE_OS_SPECIFIER = "node:" + "os"; +const NODE_PATH_SPECIFIER = "node:" + "path"; + +// Eagerly load in Node.js/Bun environment only +if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) { + dynamicImport(NODE_FS_SPECIFIER).then((m) => { + _existsSync = (m as typeof import("node:fs")).existsSync; + }); + dynamicImport(NODE_OS_SPECIFIER).then((m) => { + _homedir = (m as typeof import("node:os")).homedir; + }); + dynamicImport(NODE_PATH_SPECIFIER).then((m) => { + _join = (m as typeof import("node:path")).join; + }); +} + +import type { KnownProvider, ProviderEnv } from "./types.ts"; +import { getProviderEnvValue } from "./utils/provider-env.ts"; + +let cachedVertexAdcCredentialsExists: boolean | null = null; + +function hasVertexAdcCredentials(env?: ProviderEnv): boolean { + const explicitCredentialsPath = env?.GOOGLE_APPLICATION_CREDENTIALS; + if (explicitCredentialsPath) { + return _existsSync ? _existsSync(explicitCredentialsPath) : false; + } + + if (cachedVertexAdcCredentialsExists === null) { + // If node modules haven't loaded yet (async import race at startup), + // return false WITHOUT caching so the next call retries once they're ready. + // Only cache false permanently in a browser environment where fs is never available. + if (!_existsSync || !_homedir || !_join) { + const isNode = typeof process !== "undefined" && (process.versions?.node || process.versions?.bun); + if (!isNode) { + // Definitively in a browser — safe to cache false permanently + cachedVertexAdcCredentialsExists = false; + } + return false; + } + + // Check GOOGLE_APPLICATION_CREDENTIALS env var first (standard way) + const gacPath = getProviderEnvValue("GOOGLE_APPLICATION_CREDENTIALS", env); + if (gacPath) { + cachedVertexAdcCredentialsExists = _existsSync(gacPath); + } else { + // Fall back to default ADC path (lazy evaluation) + cachedVertexAdcCredentialsExists = _existsSync( + _join(_homedir(), ".config", "gcloud", "application_default_credentials.json"), + ); + } + } + return cachedVertexAdcCredentialsExists; +} + +function getApiKeyEnvVars(provider: string): readonly string[] | undefined { + if (provider === "github-copilot") { + return ["COPILOT_GITHUB_TOKEN"]; + } + + // ANTHROPIC_OAUTH_TOKEN takes precedence over ANTHROPIC_API_KEY + if (provider === "anthropic") { + return ["ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY"]; + } + + const envMap: Record = { + "ant-ling": "ANT_LING_API_KEY", + openai: "OPENAI_API_KEY", + "azure-openai-responses": "AZURE_OPENAI_API_KEY", + nvidia: "NVIDIA_API_KEY", + deepseek: "DEEPSEEK_API_KEY", + google: "GEMINI_API_KEY", + "google-vertex": "GOOGLE_CLOUD_API_KEY", + groq: "GROQ_API_KEY", + cerebras: "CEREBRAS_API_KEY", + xai: "XAI_API_KEY", + openrouter: "OPENROUTER_API_KEY", + "vercel-ai-gateway": "AI_GATEWAY_API_KEY", + zai: "ZAI_API_KEY", + "zai-coding-cn": "ZAI_CODING_CN_API_KEY", + mistral: "MISTRAL_API_KEY", + minimax: "MINIMAX_API_KEY", + "minimax-cn": "MINIMAX_CN_API_KEY", + moonshotai: "MOONSHOT_API_KEY", + "moonshotai-cn": "MOONSHOT_API_KEY", + huggingface: "HF_TOKEN", + fireworks: "FIREWORKS_API_KEY", + together: "TOGETHER_API_KEY", + opencode: "OPENCODE_API_KEY", + "opencode-go": "OPENCODE_API_KEY", + "kimi-coding": "KIMI_API_KEY", + "cloudflare-workers-ai": "CLOUDFLARE_API_KEY", + "cloudflare-ai-gateway": "CLOUDFLARE_API_KEY", + xiaomi: "XIAOMI_API_KEY", + "xiaomi-token-plan-cn": "XIAOMI_TOKEN_PLAN_CN_API_KEY", + "xiaomi-token-plan-ams": "XIAOMI_TOKEN_PLAN_AMS_API_KEY", + "xiaomi-token-plan-sgp": "XIAOMI_TOKEN_PLAN_SGP_API_KEY", + }; + + const envVar = envMap[provider]; + return envVar ? [envVar] : undefined; +} + +/** + * Find configured environment variables that can provide an API key for a provider. + * + * This only reports actual API key variables. It intentionally excludes ambient + * credential sources such as AWS profiles, AWS IAM credentials, and Google + * Application Default Credentials. + */ +export function findEnvKeys(provider: KnownProvider, env?: ProviderEnv): string[] | undefined; +export function findEnvKeys(provider: string, env?: ProviderEnv): string[] | undefined; +export function findEnvKeys(provider: string, env?: ProviderEnv): string[] | undefined { + const envVars = getApiKeyEnvVars(provider); + if (!envVars) return undefined; + + const found = envVars.filter((envVar) => !!getProviderEnvValue(envVar, env)); + return found.length > 0 ? found : undefined; +} + +/** + * Get API key for provider from known environment variables, e.g. OPENAI_API_KEY. + * + * Will not return API keys for providers that require OAuth tokens. + */ +export function getEnvApiKey(provider: KnownProvider, env?: ProviderEnv): string | undefined; +export function getEnvApiKey(provider: string, env?: ProviderEnv): string | undefined; +export function getEnvApiKey(provider: string, env?: ProviderEnv): string | undefined { + const envKeys = findEnvKeys(provider, env); + if (envKeys?.[0]) { + return getProviderEnvValue(envKeys[0], env); + } + + // Vertex AI supports either an explicit API key or Application Default Credentials. + // Auth is configured via `gcloud auth application-default login`. + if (provider === "google-vertex") { + const hasCredentials = hasVertexAdcCredentials(env); + const hasProject = !!( + getProviderEnvValue("GOOGLE_CLOUD_PROJECT", env) || getProviderEnvValue("GCLOUD_PROJECT", env) + ); + const hasLocation = !!getProviderEnvValue("GOOGLE_CLOUD_LOCATION", env); + + if (hasCredentials && hasProject && hasLocation) { + return ""; + } + } + + if (provider === "amazon-bedrock") { + // Amazon Bedrock supports multiple credential sources: + // 1. AWS_PROFILE - named profile from ~/.aws/credentials + // 2. AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY - standard IAM keys + // 3. AWS_BEARER_TOKEN_BEDROCK - Bedrock bearer token + // 4. AWS_CONTAINER_CREDENTIALS_RELATIVE_URI - ECS task roles + // 5. AWS_CONTAINER_CREDENTIALS_FULL_URI - ECS task roles (full URI) + // 6. AWS_WEB_IDENTITY_TOKEN_FILE - IRSA (IAM Roles for Service Accounts) + if ( + getProviderEnvValue("AWS_PROFILE", env) || + (getProviderEnvValue("AWS_ACCESS_KEY_ID", env) && getProviderEnvValue("AWS_SECRET_ACCESS_KEY", env)) || + getProviderEnvValue("AWS_BEARER_TOKEN_BEDROCK", env) || + getProviderEnvValue("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", env) || + getProviderEnvValue("AWS_CONTAINER_CREDENTIALS_FULL_URI", env) || + getProviderEnvValue("AWS_WEB_IDENTITY_TOKEN_FILE", env) + ) { + return ""; + } + } + + return undefined; +} diff --git a/cactus-code/packages/ai/src/index.ts b/cactus-code/packages/ai/src/index.ts new file mode 100644 index 000000000..ddfbb6ef6 --- /dev/null +++ b/cactus-code/packages/ai/src/index.ts @@ -0,0 +1,37 @@ +export type { Static, TSchema } from "typebox"; +export { Type } from "typebox"; + +// Core only, side-effect free: no generated catalogs, no provider factories, +// no api-registry, no OAuth implementations, no compat. Provider factories +// live under "@earendil-works/pi-ai/providers/*", API implementations under +// "@earendil-works/pi-ai/api/*", the old global API under +// "@earendil-works/pi-ai/compat". +export * from "./api/lazy.ts"; +export type { OpenAICompletionsOptions } from "./api/openai-completions.ts"; +export * from "./auth/context.ts"; +export * from "./auth/credential-store.ts"; +export * from "./auth/helpers.ts"; +export * from "./auth/types.ts"; +export * from "./models.ts"; +export * from "./providers/faux.ts"; +export * from "./session-resources.ts"; +export * from "./types.ts"; +export * from "./utils/diagnostics.ts"; +export * from "./utils/event-stream.ts"; +export * from "./utils/json-parse.ts"; +export type { + OAuthAuthInfo, + OAuthCredentials, + OAuthDeviceCodeInfo, + OAuthLoginCallbacks, + OAuthPrompt, + OAuthProvider, + OAuthProviderId, + OAuthProviderInfo, + OAuthProviderInterface, + OAuthSelectOption, + OAuthSelectPrompt, +} from "./utils/oauth/types.ts"; +export * from "./utils/overflow.ts"; +export * from "./utils/typebox-helpers.ts"; +export * from "./utils/validation.ts"; diff --git a/cactus-code/packages/ai/src/models.generated.ts b/cactus-code/packages/ai/src/models.generated.ts new file mode 100644 index 000000000..c06f3454a --- /dev/null +++ b/cactus-code/packages/ai/src/models.generated.ts @@ -0,0 +1,3 @@ +export const MODELS = { + cactus: {} as Record, +} as const; diff --git a/cactus-code/packages/ai/src/models.ts b/cactus-code/packages/ai/src/models.ts new file mode 100644 index 000000000..f9cf27d36 --- /dev/null +++ b/cactus-code/packages/ai/src/models.ts @@ -0,0 +1,441 @@ +import { lazyStream } from "./api/lazy.ts"; +import { defaultProviderAuthContext as defaultAuthContext } from "./auth/context.ts"; +import { InMemoryCredentialStore } from "./auth/credential-store.ts"; +import { ModelsError, resolveProviderAuth } from "./auth/resolve.ts"; +import type { AuthContext, AuthResult, CredentialStore, ProviderAuth } from "./auth/types.ts"; +import type { + Api, + ApiStreamOptions, + AssistantMessage, + AssistantMessageEventStream, + Context, + Model, + ModelThinkingLevel, + ProviderHeaders, + ProviderStreams, + SimpleStreamOptions, + StreamOptions, + Usage, +} from "./types.ts"; + +export { type AuthModel, ModelsError, type ModelsErrorCode } from "./auth/resolve.ts"; + +/** + * A provider is the concrete runtime unit. It owns id/name/base metadata, + * auth methods, model listing, and stream behavior. + * + * `TApi` lets concrete provider factories declare which APIs their models + * use (e.g. `openaiProvider(): Provider<"openai-responses" | "openai-completions">`), + * giving typed model lists to direct factory users. Inside a `Models` + * collection providers are held as `Provider`. + */ +export interface Provider { + readonly id: string; + readonly name: string; + + readonly baseUrl?: string; + readonly headers?: ProviderHeaders; + + /** + * Required: at least one of `apiKey`/`oauth`. Every provider has auth + * semantics — even providers with only ambient credentials (env vars, AWS + * profiles, ADC files) and keyless local servers provide `apiKey` auth + * whose `resolve()` reports whether the provider is configured. + * `Models.getAuth()` returns undefined when the provider is unconfigured. + */ + readonly auth: ProviderAuth; + + /** + * Current known models, sync. Static providers return their catalog; + * dynamic providers return the list as of the last `refreshModels()` + * (empty before the first). Must not throw; `Models` treats a throwing + * implementation as having no models. + */ + getModels(): readonly Model[]; + + /** + * Dynamic providers only: fetch and update the model list. Side-effect-free + * discovery (no loading/downloading); provider-specific model lifecycle + * belongs in app commands. Concurrent calls share one in-flight fetch. + * May reject (network); on rejection the model list stays at its last-known + * state and a later call retries. + */ + refreshModels?(): Promise; + + stream( + model: Model, + context: Context, + options?: ApiStreamOptions, + ): AssistantMessageEventStream; + + streamSimple(model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream; +} + +/** + * Runtime collection of providers plus auth application and stream + * convenience. Providers own stream behavior; `Models` resolves auth and + * delegates each request to the provider that owns the model. + */ +export interface Models { + getProviders(): readonly Provider[]; + getProvider(id: string): Provider | undefined; + + /** + * Sync read of last-known models from one provider or all providers. + * Best-effort: a provider whose `getModels()` throws yields no models. + */ + getModels(provider?: string): readonly Model[]; + + /** + * Sync runtime model lookup against last-known lists. Dynamic model lists + * are typed as `Model`; narrow with the `hasApi()` type guard. + */ + getModel(provider: string, id: string): Model | undefined; + + /** + * Ask dynamic providers to re-fetch their model lists. With a provider id, + * rejects with `ModelsError` ("model_source") on that provider's fetch + * failure; without one, refreshes all providers concurrently best-effort. + * Static providers (no `refreshModels`) are no-ops. + */ + refresh(provider?: string): Promise; + + /** + * Resolve request auth for a model. Includes a source label for status UI. + * Resolves `undefined` when the provider is unknown or unconfigured. + * Rejects with `ModelsError`: code "oauth" when a token refresh fails (the + * stored credential is preserved for retry; re-login fixes it), code "auth" + * when api-key resolution or the credential store fails. Request paths + * surface rejections as stream errors; status/availability UIs catch them + * and render "needs re-login" instead of treating them as unconfigured. + */ + getAuth(model: Model): Promise; + + stream( + model: Model, + context: Context, + options?: ApiStreamOptions, + ): AssistantMessageEventStream; + + complete( + model: Model, + context: Context, + options?: ApiStreamOptions, + ): Promise; + + streamSimple(model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream; + completeSimple(model: Model, context: Context, options?: SimpleStreamOptions): Promise; +} + +export interface MutableModels extends Models { + /** Upsert/replace by provider.id. Provider ids are unique. */ + setProvider(provider: Provider): void; + deleteProvider(id: string): void; + clearProviders(): void; +} + +export interface CreateModelsOptions { + credentials?: CredentialStore; + authContext?: AuthContext; +} + +class ModelsImpl implements MutableModels { + private providers = new Map(); + private credentials: CredentialStore; + private authContext: AuthContext; + + constructor(options?: CreateModelsOptions) { + this.credentials = options?.credentials ?? new InMemoryCredentialStore(); + this.authContext = options?.authContext ?? defaultAuthContext(); + } + + setProvider(provider: Provider): void { + this.providers.set(provider.id, provider); + } + + deleteProvider(id: string): void { + this.providers.delete(id); + } + + clearProviders(): void { + this.providers.clear(); + } + + getProviders(): readonly Provider[] { + return Array.from(this.providers.values()); + } + + getProvider(id: string): Provider | undefined { + return this.providers.get(id); + } + + getModels(provider?: string): readonly Model[] { + if (provider !== undefined) { + const entry = this.providers.get(provider); + if (!entry) return []; + try { + return entry.getModels(); + } catch { + return []; + } + } + + const models: Model[] = []; + for (const entry of this.providers.values()) { + try { + models.push(...entry.getModels()); + } catch { + // Best-effort: ill-behaved providers yield no models. + } + } + return models; + } + + getModel(provider: string, id: string): Model | undefined { + return this.getModels(provider).find((model) => model.id === id); + } + + async refresh(provider?: string): Promise { + if (provider !== undefined) { + const entry = this.providers.get(provider); + if (!entry?.refreshModels) return; + try { + await entry.refreshModels(); + } catch (error) { + if (error instanceof ModelsError) throw error; + throw new ModelsError("model_source", `Model refresh failed for ${provider}`, { cause: error }); + } + return; + } + + // Cannot reject: the async mapper turns even sync throws from ill-behaved + // providers into rejections, and allSettled captures all of them. + await Promise.allSettled(Array.from(this.providers.values(), async (entry) => entry.refreshModels?.())); + } + + async getAuth(model: Model): Promise { + const provider = this.providers.get(model.provider); + if (!provider) return undefined; + return resolveProviderAuth(provider, model, this.credentials, this.authContext); + } + + private requireProvider(model: Model): Provider { + const provider = this.providers.get(model.provider); + if (!provider) { + throw new ModelsError("provider", `Unknown provider: ${model.provider}`); + } + return provider; + } + + private async applyAuth( + model: Model, + options: TOptions | undefined, + ): Promise<{ requestModel: Model; requestOptions: TOptions | undefined }> { + const resolution = await resolveProviderAuth( + this.requireProvider(model), + model, + this.credentials, + this.authContext, + { + apiKey: options?.apiKey, + env: options?.env, + }, + ); + const auth = resolution?.auth; + if (!auth) return { requestModel: model, requestOptions: options }; + + const requestModel = auth.baseUrl ? { ...model, baseUrl: auth.baseUrl } : model; + + // Explicit request options win per-field; headers/env merge per key. + const apiKey = options?.apiKey ?? auth.apiKey; + const headers = auth.headers || options?.headers ? { ...auth.headers, ...options?.headers } : undefined; + const env = resolution.env || options?.env ? { ...(resolution.env ?? {}), ...(options?.env ?? {}) } : undefined; + const requestOptions = { ...options, apiKey, headers, env } as TOptions; + + return { requestModel, requestOptions }; + } + + stream( + model: Model, + context: Context, + options?: ApiStreamOptions, + ): AssistantMessageEventStream { + return lazyStream(model, async () => { + const provider = this.requireProvider(model); + const { requestModel, requestOptions } = await this.applyAuth(model, options as StreamOptions | undefined); + return provider.stream(requestModel as Model, context, requestOptions as ApiStreamOptions); + }); + } + + async complete( + model: Model, + context: Context, + options?: ApiStreamOptions, + ): Promise { + return this.stream(model, context, options).result(); + } + + streamSimple(model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream { + return lazyStream(model, async () => { + const provider = this.requireProvider(model); + const { requestModel, requestOptions } = await this.applyAuth(model, options); + return provider.streamSimple(requestModel, context, requestOptions); + }); + } + + async completeSimple(model: Model, context: Context, options?: SimpleStreamOptions): Promise { + return this.streamSimple(model, context, options).result(); + } +} + +export function createModels(options?: CreateModelsOptions): MutableModels { + return new ModelsImpl(options); +} + +export interface CreateProviderOptions { + id: string; + /** Display name. Default: `id`. */ + name?: string; + baseUrl?: string; + headers?: ProviderHeaders; + /** Required — every provider has auth semantics, even ambient/keyless ones. */ + auth: ProviderAuth; + /** Initial model list (empty for purely dynamic providers). */ + models: readonly Model[]; + /** + * Dynamic providers: fetch the current list. Stored on success; concurrent + * calls share one in-flight fetch. May reject: the stored list then stays + * at its last-known state, the rejection propagates to the caller of + * `refreshModels()` (wrapped as ModelsError "model_source" by + * `Models.refresh(provider)`), and a later call retries. + */ + refreshModels?: () => Promise[]>; + /** Single implementation, or map keyed by `model.api` for mixed-API providers. */ + api: ProviderStreams | Partial>; +} + +/** + * Builds a provider from parts. Built-in provider factories and models.json + * custom providers both go through this. A single `api` streams all models; + * an `api` map dispatches on `model.api`, and a model whose api has no entry + * produces a stream error. + */ +export function createProvider(input: CreateProviderOptions): Provider { + let models = input.models; + let inflightRefresh: Promise | undefined; + const refreshModels = input.refreshModels; + const single = + typeof (input.api as ProviderStreams).stream === "function" ? (input.api as ProviderStreams) : undefined; + const byApi = single ? undefined : (input.api as Partial>); + + const apiFor = (model: Model): ProviderStreams | undefined => single ?? byApi?.[model.api]; + + const dispatch = ( + model: Model, + run: (streams: ProviderStreams) => AssistantMessageEventStream, + ): AssistantMessageEventStream => { + const streams = apiFor(model); + if (!streams) { + return lazyStream(model, async () => { + throw new ModelsError("stream", `Provider ${input.id} has no API implementation for "${model.api}"`); + }); + } + return run(streams); + }; + + return { + id: input.id, + name: input.name ?? input.id, + baseUrl: input.baseUrl, + headers: input.headers, + auth: input.auth, + getModels: () => models, + refreshModels: refreshModels + ? () => { + inflightRefresh ??= (async () => { + try { + models = await refreshModels(); + } finally { + inflightRefresh = undefined; + } + })(); + return inflightRefresh; + } + : undefined, + stream: (model, context, options) => dispatch(model, (streams) => streams.stream(model, context, options)), + streamSimple: (model, context, options) => + dispatch(model, (streams) => streams.streamSimple(model, context, options)), + }; +} + +/** + * Runtime-checked narrowing for dynamically looked-up models: + * + * ```ts + * const model = models.getModel("anthropic", "claude-opus-4-7"); + * if (model && hasApi(model, "anthropic-messages")) { + * // model: Model<"anthropic-messages">, stream options fully typed + * } + * ``` + */ +export function hasApi(model: Model, api: TApi): model is Model { + return model.api === api; +} + +export function calculateCost(model: Model, usage: Usage): Usage["cost"] { + // Anthropic charges 2x base input for 1h cache writes. + const longWrite = usage.cacheWrite1h ?? 0; + const shortWrite = usage.cacheWrite - longWrite; + usage.cost.input = (model.cost.input / 1000000) * usage.input; + usage.cost.output = (model.cost.output / 1000000) * usage.output; + usage.cost.cacheRead = (model.cost.cacheRead / 1000000) * usage.cacheRead; + usage.cost.cacheWrite = (model.cost.cacheWrite * shortWrite + model.cost.input * 2 * longWrite) / 1000000; + usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite; + return usage.cost; +} + +const EXTENDED_THINKING_LEVELS: ModelThinkingLevel[] = ["off", "minimal", "low", "medium", "high", "xhigh"]; + +export function getSupportedThinkingLevels(model: Model): ModelThinkingLevel[] { + if (!model.reasoning) return ["off"]; + + return EXTENDED_THINKING_LEVELS.filter((level) => { + const mapped = model.thinkingLevelMap?.[level]; + if (mapped === null) return false; + if (level === "xhigh") return mapped !== undefined; + return true; + }); +} + +export function clampThinkingLevel( + model: Model, + level: ModelThinkingLevel, +): ModelThinkingLevel { + const availableLevels = getSupportedThinkingLevels(model); + if (availableLevels.includes(level)) return level; + + const requestedIndex = EXTENDED_THINKING_LEVELS.indexOf(level); + if (requestedIndex === -1) return availableLevels[0] ?? "off"; + + for (let i = requestedIndex; i < EXTENDED_THINKING_LEVELS.length; i++) { + const candidate = EXTENDED_THINKING_LEVELS[i]; + if (availableLevels.includes(candidate)) return candidate; + } + for (let i = requestedIndex - 1; i >= 0; i--) { + const candidate = EXTENDED_THINKING_LEVELS[i]; + if (availableLevels.includes(candidate)) return candidate; + } + return availableLevels[0] ?? "off"; +} + +/** + * Check if two models are equal by comparing both their id and provider. + * Returns false if either model is null or undefined. + */ +export function modelsAreEqual( + a: Model | null | undefined, + b: Model | null | undefined, +): boolean { + if (!a || !b) return false; + return a.id === b.id && a.provider === b.provider; +} diff --git a/cactus-code/packages/ai/src/oauth.ts b/cactus-code/packages/ai/src/oauth.ts new file mode 100644 index 000000000..489c8b89e --- /dev/null +++ b/cactus-code/packages/ai/src/oauth.ts @@ -0,0 +1,61 @@ +import type { + OAuthCredentials, + OAuthProviderId, + OAuthProviderInfo, + OAuthProviderInterface, +} from "./utils/oauth/types.ts"; + +export * from "./utils/oauth/types.ts"; + +const oauthProviderRegistry = new Map(); + +export function getOAuthProvider(id: OAuthProviderId): OAuthProviderInterface | undefined { + return oauthProviderRegistry.get(id); +} + +export function registerOAuthProvider(provider: OAuthProviderInterface): void { + oauthProviderRegistry.set(provider.id, provider); +} + +export function unregisterOAuthProvider(id: string): void { + oauthProviderRegistry.delete(id); +} + +export function resetOAuthProviders(): void { + oauthProviderRegistry.clear(); +} + +export function getOAuthProviders(): OAuthProviderInterface[] { + return Array.from(oauthProviderRegistry.values()); +} + +export function getOAuthProviderInfoList(): OAuthProviderInfo[] { + return getOAuthProviders().map((p) => ({ id: p.id, name: p.name, available: true })); +} + +export async function refreshOAuthToken( + providerId: OAuthProviderId, + credentials: OAuthCredentials, +): Promise { + const provider = getOAuthProvider(providerId); + if (!provider) throw new Error(`Unknown OAuth provider: ${providerId}`); + return provider.refreshToken(credentials); +} + +export async function getOAuthApiKey( + providerId: OAuthProviderId, + credentials: Record, +): Promise<{ newCredentials: OAuthCredentials; apiKey: string } | null> { + const provider = getOAuthProvider(providerId); + if (!provider) throw new Error(`Unknown OAuth provider: ${providerId}`); + let creds = credentials[providerId]; + if (!creds) return null; + if (Date.now() >= creds.expires) { + try { + creds = await provider.refreshToken(creds); + } catch (_error) { + throw new Error(`Failed to refresh OAuth token for ${providerId}`); + } + } + return { newCredentials: creds, apiKey: provider.getApiKey(creds) }; +} diff --git a/cactus-code/packages/ai/src/providers/all.ts b/cactus-code/packages/ai/src/providers/all.ts new file mode 100644 index 000000000..3eebad1cd --- /dev/null +++ b/cactus-code/packages/ai/src/providers/all.ts @@ -0,0 +1,42 @@ +import { MODELS } from "../models.generated.ts"; +import { type CreateModelsOptions, createModels, type MutableModels, type Provider } from "../models.ts"; +import type { Api, KnownProvider, Model } from "../types.ts"; +import { cactusProvider } from "./cactus.ts"; + +type BuiltinModelApi< + TProvider extends KnownProvider, + TModelId extends keyof (typeof MODELS)[TProvider], +> = (typeof MODELS)[TProvider][TModelId] extends { api: infer TApi } ? (TApi extends Api ? TApi : never) : never; + +export function getBuiltinModel( + provider: TProvider, + modelId: TModelId, +): Model> { + const models = MODELS[provider] as Record> | undefined; + return models?.[modelId as string] as Model>; +} + +export function getBuiltinProviders(): KnownProvider[] { + return Object.keys(MODELS) as KnownProvider[]; +} + +export function getBuiltinModels( + provider: TProvider, +): Model>[] { + const models = MODELS[provider] as Record> | undefined; + return models + ? (Object.values(models) as Model>[]) + : []; +} + +export function builtinProviders(): Provider[] { + return [cactusProvider()]; +} + +export function builtinModels(options?: CreateModelsOptions): MutableModels { + const models = createModels(options); + for (const provider of builtinProviders()) { + models.setProvider(provider); + } + return models; +} diff --git a/cactus-code/packages/ai/src/providers/cactus.ts b/cactus-code/packages/ai/src/providers/cactus.ts new file mode 100644 index 000000000..80fa3dd4f --- /dev/null +++ b/cactus-code/packages/ai/src/providers/cactus.ts @@ -0,0 +1,95 @@ +import { openAICompletionsApi } from "../api/openai-completions.lazy.ts"; +import type { ApiKeyAuth } from "../auth/types.ts"; +import { createProvider, type Provider } from "../models.ts"; +import type { Model } from "../types.ts"; + +const DEFAULT_BASE_URL = "http://127.0.0.1:8080/v1"; +const DEFAULT_CONTEXT_WINDOW = 32768; +const DEFAULT_MAX_TOKENS = 8192; +const MAX_CONTEXT_WINDOW = 8_000_000; + +function sanitizeContextWindow(cw: unknown): number { + return typeof cw === "number" && Number.isFinite(cw) && cw > 0 && cw <= MAX_CONTEXT_WINDOW + ? Math.floor(cw) + : DEFAULT_CONTEXT_WINDOW; +} + +export function cactusBaseUrl(): string { + return process.env.CACTUS_BASE_URL?.trim().replace(/\/+$/, "") || DEFAULT_BASE_URL; +} + +function cactusAuth(): ApiKeyAuth { + return { + name: "Cactus (local, no auth)", + resolve: async () => ({ auth: { apiKey: "cactus" }, source: "local server" }), + }; +} + +function isVisionModel(id: string, modelType: string): boolean { + const s = `${id} ${modelType}`.toLowerCase(); + return /(?:^|[^a-z])vl(?:[^a-z]|$)|gemma-?4|gemma-?3n|qwen3\.5/.test(s); +} + +function isThinkingModel(id: string, modelType: string): boolean { + const s = `${id} ${modelType}`.toLowerCase(); + return /gemma-?4|qwen3/.test(s); +} + +interface OpenAIModelObject { + id?: string; + context_window?: number; + model_type?: string; +} +interface OpenAIModelsResponse { + data?: OpenAIModelObject[]; +} + +export function cactusModelFromId( + baseUrl: string, + id: string, + contextWindow: number = DEFAULT_CONTEXT_WINDOW, + modelType = "", +): Model<"openai-completions"> { + return { + id, + name: id, + api: "openai-completions", + provider: "cactus", + baseUrl, + reasoning: isThinkingModel(id, modelType), + input: isVisionModel(id, modelType) ? ["text", "image"] : ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: sanitizeContextWindow(contextWindow), + maxTokens: DEFAULT_MAX_TOKENS, + }; +} + +export async function fetchCactusModels(baseUrl: string = cactusBaseUrl()): Promise[]> { + const res = await fetch(`${baseUrl}/models`); + if (!res.ok) throw new Error(`Cactus model discovery failed: HTTP ${res.status}`); + const body = (await res.json()) as OpenAIModelsResponse; + const data = Array.isArray(body?.data) ? body.data : []; + return data + .filter((m): m is OpenAIModelObject & { id: string } => typeof m?.id === "string" && m.id.length > 0) + .map((m) => + cactusModelFromId( + baseUrl, + m.id, + typeof m.context_window === "number" ? m.context_window : DEFAULT_CONTEXT_WINDOW, + typeof m.model_type === "string" ? m.model_type : "", + ), + ); +} + +export function cactusProvider(): Provider<"openai-completions"> { + const baseUrl = cactusBaseUrl(); + return createProvider({ + id: "cactus", + name: "Cactus", + baseUrl, + auth: { apiKey: cactusAuth() }, + models: [], + refreshModels: () => fetchCactusModels(baseUrl), + api: openAICompletionsApi(), + }); +} diff --git a/cactus-code/packages/ai/src/providers/faux.ts b/cactus-code/packages/ai/src/providers/faux.ts new file mode 100644 index 000000000..4a26f1ad6 --- /dev/null +++ b/cactus-code/packages/ai/src/providers/faux.ts @@ -0,0 +1,538 @@ +import { createProvider, type Provider } from "../models.ts"; +import type { + AssistantMessage, + AssistantMessageEventStream, + Context, + ImageContent, + Message, + Model, + SimpleStreamOptions, + StreamFunction, + StreamOptions, + TextContent, + ThinkingContent, + ToolCall, + ToolResultMessage, + Usage, +} from "../types.ts"; +import { createAssistantMessageEventStream } from "../utils/event-stream.ts"; + +const DEFAULT_API = "faux"; +const DEFAULT_PROVIDER = "faux"; +const DEFAULT_MODEL_ID = "faux-1"; +const DEFAULT_MODEL_NAME = "Faux Model"; +const DEFAULT_BASE_URL = "http://localhost:0"; +const DEFAULT_MIN_TOKEN_SIZE = 3; +const DEFAULT_MAX_TOKEN_SIZE = 5; + +const DEFAULT_USAGE: Usage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +export interface FauxModelDefinition { + id: string; + name?: string; + reasoning?: boolean; + input?: ("text" | "image")[]; + cost?: { input: number; output: number; cacheRead: number; cacheWrite: number }; + contextWindow?: number; + maxTokens?: number; +} + +export type FauxContentBlock = TextContent | ThinkingContent | ToolCall; + +export function fauxText(text: string): TextContent { + return { type: "text", text }; +} + +export function fauxThinking(thinking: string): ThinkingContent { + return { type: "thinking", thinking }; +} + +export function fauxToolCall(name: string, arguments_: ToolCall["arguments"], options: { id?: string } = {}): ToolCall { + return { + type: "toolCall", + id: options.id ?? randomId("tool"), + name, + arguments: arguments_, + }; +} + +function normalizeFauxAssistantContent(content: string | FauxContentBlock | FauxContentBlock[]): FauxContentBlock[] { + if (typeof content === "string") { + return [fauxText(content)]; + } + return Array.isArray(content) ? content : [content]; +} + +export function fauxAssistantMessage( + content: string | FauxContentBlock | FauxContentBlock[], + options: { + stopReason?: AssistantMessage["stopReason"]; + errorMessage?: string; + responseId?: string; + timestamp?: number; + } = {}, +): AssistantMessage { + return { + role: "assistant", + content: normalizeFauxAssistantContent(content), + api: DEFAULT_API, + provider: DEFAULT_PROVIDER, + model: DEFAULT_MODEL_ID, + usage: DEFAULT_USAGE, + stopReason: options.stopReason ?? "stop", + errorMessage: options.errorMessage, + responseId: options.responseId, + timestamp: options.timestamp ?? Date.now(), + }; +} + +export type FauxResponseFactory = ( + context: Context, + options: StreamOptions | undefined, + state: { callCount: number }, + model: Model, +) => AssistantMessage | Promise; + +export type FauxResponseStep = AssistantMessage | FauxResponseFactory; + +export interface RegisterFauxProviderOptions { + api?: string; + provider?: string; + models?: FauxModelDefinition[]; + tokensPerSecond?: number; + tokenSize?: { + min?: number; + max?: number; + }; +} + +export interface FauxProviderRegistration { + api: string; + models: [Model, ...Model[]]; + getModel(): Model; + getModel(modelId: string): Model | undefined; + state: { callCount: number }; + setResponses: (responses: FauxResponseStep[]) => void; + appendResponses: (responses: FauxResponseStep[]) => void; + getPendingResponseCount: () => number; + unregister: () => void; +} + +export interface FauxProviderHandle { + provider: Provider; + api: string; + models: [Model, ...Model[]]; + getModel(): Model; + getModel(modelId: string): Model | undefined; + state: { callCount: number }; + setResponses: (responses: FauxResponseStep[]) => void; + appendResponses: (responses: FauxResponseStep[]) => void; + getPendingResponseCount: () => number; +} + +function estimateTokens(text: string): number { + return Math.ceil(text.length / 4); +} + +function randomId(prefix: string): string { + return `${prefix}:${Date.now()}:${Math.random().toString(36).slice(2)}`; +} + +function contentToText(content: string | Array): string { + if (typeof content === "string") { + return content; + } + return content + .map((block) => { + if (block.type === "text") { + return block.text; + } + return `[image:${block.mimeType}:${block.data.length}]`; + }) + .join("\n"); +} + +function assistantContentToText(content: Array): string { + return content + .map((block) => { + if (block.type === "text") { + return block.text; + } + if (block.type === "thinking") { + return block.thinking; + } + return `${block.name}:${JSON.stringify(block.arguments)}`; + }) + .join("\n"); +} + +function toolResultToText(message: ToolResultMessage): string { + return [message.toolName, ...message.content.map((block) => contentToText([block]))].join("\n"); +} + +function messageToText(message: Message): string { + if (message.role === "user") { + return contentToText(message.content); + } + if (message.role === "assistant") { + return assistantContentToText(message.content); + } + return toolResultToText(message); +} + +function serializeContext(context: Context): string { + const parts: string[] = []; + if (context.systemPrompt) { + parts.push(`system:${context.systemPrompt}`); + } + for (const message of context.messages) { + parts.push(`${message.role}:${messageToText(message)}`); + } + if (context.tools?.length) { + parts.push(`tools:${JSON.stringify(context.tools)}`); + } + return parts.join("\n\n"); +} + +function commonPrefixLength(a: string, b: string): number { + const length = Math.min(a.length, b.length); + let index = 0; + while (index < length && a[index] === b[index]) { + index++; + } + return index; +} + +function withUsageEstimate( + message: AssistantMessage, + context: Context, + options: StreamOptions | undefined, + promptCache: Map, +): AssistantMessage { + const promptText = serializeContext(context); + const promptTokens = estimateTokens(promptText); + const outputTokens = estimateTokens(assistantContentToText(message.content)); + let input = promptTokens; + let cacheRead = 0; + let cacheWrite = 0; + const sessionId = options?.sessionId; + + if (sessionId && options?.cacheRetention !== "none") { + const previousPrompt = promptCache.get(sessionId); + if (previousPrompt) { + const cachedChars = commonPrefixLength(previousPrompt, promptText); + cacheRead = estimateTokens(previousPrompt.slice(0, cachedChars)); + cacheWrite = estimateTokens(promptText.slice(cachedChars)); + input = Math.max(0, promptTokens - cacheRead); + } else { + cacheWrite = promptTokens; + } + promptCache.set(sessionId, promptText); + } + + return { + ...message, + usage: { + input, + output: outputTokens, + cacheRead, + cacheWrite, + totalTokens: input + outputTokens + cacheRead + cacheWrite, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; +} + +function splitStringByTokenSize(text: string, minTokenSize: number, maxTokenSize: number): string[] { + const chunks: string[] = []; + let index = 0; + while (index < text.length) { + const tokenSize = minTokenSize + Math.floor(Math.random() * (maxTokenSize - minTokenSize + 1)); + const charSize = Math.max(1, tokenSize * 4); + chunks.push(text.slice(index, index + charSize)); + index += charSize; + } + return chunks.length > 0 ? chunks : [""]; +} + +function cloneMessage(message: AssistantMessage, api: string, provider: string, modelId: string): AssistantMessage { + const cloned = structuredClone(message); + return { + ...cloned, + api, + provider, + model: modelId, + timestamp: cloned.timestamp ?? Date.now(), + usage: cloned.usage ?? DEFAULT_USAGE, + }; +} + +function createErrorMessage(error: unknown, api: string, provider: string, modelId: string): AssistantMessage { + return { + role: "assistant", + content: [], + api, + provider, + model: modelId, + usage: DEFAULT_USAGE, + stopReason: "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + }; +} + +function createAbortedMessage(partial: AssistantMessage): AssistantMessage { + return { + ...partial, + stopReason: "aborted", + errorMessage: "Request was aborted", + timestamp: Date.now(), + }; +} + +function scheduleChunk(chunk: string, tokensPerSecond: number | undefined): Promise { + if (!tokensPerSecond || tokensPerSecond <= 0) { + return new Promise((resolve) => queueMicrotask(resolve)); + } + const delayMs = (estimateTokens(chunk) / tokensPerSecond) * 1000; + return new Promise((resolve) => setTimeout(resolve, delayMs)); +} + +async function streamWithDeltas( + stream: AssistantMessageEventStream, + message: AssistantMessage, + minTokenSize: number, + maxTokenSize: number, + tokensPerSecond: number | undefined, + signal: AbortSignal | undefined, +): Promise { + const partial: AssistantMessage = { ...message, content: [] }; + if (signal?.aborted) { + const aborted = createAbortedMessage(partial); + stream.push({ type: "error", reason: "aborted", error: aborted }); + stream.end(aborted); + return; + } + + stream.push({ type: "start", partial: { ...partial } }); + + for (let index = 0; index < message.content.length; index++) { + if (signal?.aborted) { + const aborted = createAbortedMessage(partial); + stream.push({ type: "error", reason: "aborted", error: aborted }); + stream.end(aborted); + return; + } + + const block = message.content[index]; + + if (block.type === "thinking") { + partial.content = [...partial.content, { type: "thinking", thinking: "" }]; + stream.push({ type: "thinking_start", contentIndex: index, partial: { ...partial } }); + for (const chunk of splitStringByTokenSize(block.thinking, minTokenSize, maxTokenSize)) { + await scheduleChunk(chunk, tokensPerSecond); + if (signal?.aborted) { + const aborted = createAbortedMessage(partial); + stream.push({ type: "error", reason: "aborted", error: aborted }); + stream.end(aborted); + return; + } + (partial.content[index] as ThinkingContent).thinking += chunk; + stream.push({ type: "thinking_delta", contentIndex: index, delta: chunk, partial: { ...partial } }); + } + stream.push({ + type: "thinking_end", + contentIndex: index, + content: block.thinking, + partial: { ...partial }, + }); + continue; + } + + if (block.type === "text") { + partial.content = [...partial.content, { type: "text", text: "" }]; + stream.push({ type: "text_start", contentIndex: index, partial: { ...partial } }); + for (const chunk of splitStringByTokenSize(block.text, minTokenSize, maxTokenSize)) { + await scheduleChunk(chunk, tokensPerSecond); + if (signal?.aborted) { + const aborted = createAbortedMessage(partial); + stream.push({ type: "error", reason: "aborted", error: aborted }); + stream.end(aborted); + return; + } + (partial.content[index] as TextContent).text += chunk; + stream.push({ type: "text_delta", contentIndex: index, delta: chunk, partial: { ...partial } }); + } + stream.push({ type: "text_end", contentIndex: index, content: block.text, partial: { ...partial } }); + continue; + } + + partial.content = [...partial.content, { type: "toolCall", id: block.id, name: block.name, arguments: {} }]; + stream.push({ type: "toolcall_start", contentIndex: index, partial: { ...partial } }); + for (const chunk of splitStringByTokenSize(JSON.stringify(block.arguments), minTokenSize, maxTokenSize)) { + await scheduleChunk(chunk, tokensPerSecond); + if (signal?.aborted) { + const aborted = createAbortedMessage(partial); + stream.push({ type: "error", reason: "aborted", error: aborted }); + stream.end(aborted); + return; + } + stream.push({ type: "toolcall_delta", contentIndex: index, delta: chunk, partial: { ...partial } }); + } + (partial.content[index] as ToolCall).arguments = block.arguments; + stream.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: { ...partial } }); + } + + if (message.stopReason === "error" || message.stopReason === "aborted") { + stream.push({ type: "error", reason: message.stopReason, error: message }); + stream.end(message); + return; + } + + stream.push({ type: "done", reason: message.stopReason, message }); + stream.end(message); +} + +export function createFauxCore(options: RegisterFauxProviderOptions) { + const api = options.api ?? randomId(DEFAULT_API); + const provider = options.provider ?? DEFAULT_PROVIDER; + const minTokenSize = Math.max( + 1, + Math.min(options.tokenSize?.min ?? DEFAULT_MIN_TOKEN_SIZE, options.tokenSize?.max ?? DEFAULT_MAX_TOKEN_SIZE), + ); + const maxTokenSize = Math.max(minTokenSize, options.tokenSize?.max ?? DEFAULT_MAX_TOKEN_SIZE); + let pendingResponses: FauxResponseStep[] = []; + const tokensPerSecond = options.tokensPerSecond; + const state = { callCount: 0 }; + const promptCache = new Map(); + + const modelDefinitions = options.models?.length + ? options.models + : [ + { + id: DEFAULT_MODEL_ID, + name: DEFAULT_MODEL_NAME, + reasoning: false, + input: ["text", "image"] as ("text" | "image")[], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 16384, + }, + ]; + const models = modelDefinitions.map((definition) => ({ + id: definition.id, + name: definition.name ?? definition.id, + api, + provider, + baseUrl: DEFAULT_BASE_URL, + reasoning: definition.reasoning ?? false, + input: definition.input ?? ["text", "image"], + cost: definition.cost ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: definition.contextWindow ?? 128000, + maxTokens: definition.maxTokens ?? 16384, + })) as [Model, ...Model[]]; + + const stream: StreamFunction = (requestModel, context, streamOptions) => { + const outer = createAssistantMessageEventStream(); + const step = pendingResponses.shift(); + state.callCount++; + + queueMicrotask(async () => { + try { + await streamOptions?.onResponse?.({ status: 200, headers: {} }, requestModel); + if (!step) { + let message = createErrorMessage( + new Error("No more faux responses queued"), + api, + provider, + requestModel.id, + ); + message = withUsageEstimate(message, context, streamOptions, promptCache); + outer.push({ type: "error", reason: "error", error: message }); + outer.end(message); + return; + } + + const resolved = + typeof step === "function" ? await step(context, streamOptions, state, requestModel) : step; + let message = cloneMessage(resolved, api, provider, requestModel.id); + message = withUsageEstimate(message, context, streamOptions, promptCache); + await streamWithDeltas(outer, message, minTokenSize, maxTokenSize, tokensPerSecond, streamOptions?.signal); + } catch (error) { + const message = createErrorMessage(error, api, provider, requestModel.id); + outer.push({ type: "error", reason: "error", error: message }); + outer.end(message); + } + }); + + return outer; + }; + + const streamSimple: StreamFunction = (streamModel, context, streamOptions) => + stream(streamModel, context, streamOptions); + + function getModel(): Model; + function getModel(requestedModelId: string): Model | undefined; + function getModel(requestedModelId?: string): Model | undefined { + if (!requestedModelId) { + return models[0]; + } + return models.find((candidate) => candidate.id === requestedModelId); + } + + return { + api, + provider, + models, + stream, + streamSimple, + getModel, + state, + setResponses(responses: FauxResponseStep[]) { + pendingResponses = [...responses]; + }, + appendResponses(responses: FauxResponseStep[]) { + pendingResponses.push(...responses); + }, + getPendingResponseCount() { + return pendingResponses.length; + }, + }; +} + +/** + * Faux provider for tests built on explicit `Models` collections: + * + * ```ts + * const faux = fauxProvider(); + * const models = createModels(); + * models.setProvider(faux.provider); + * faux.setResponses([fauxAssistantMessage("hi")]); + * ``` + */ +export function fauxProvider(options: RegisterFauxProviderOptions = {}): FauxProviderHandle { + const core = createFauxCore(options); + const provider = createProvider({ + id: core.provider, + auth: { apiKey: { name: "Faux", resolve: async () => ({ auth: {} }) } }, + models: core.models, + api: { stream: core.stream, streamSimple: core.streamSimple }, + }); + return { + provider, + api: core.api, + models: core.models, + getModel: core.getModel, + state: core.state, + setResponses: core.setResponses, + appendResponses: core.appendResponses, + getPendingResponseCount: core.getPendingResponseCount, + }; +} diff --git a/cactus-code/packages/ai/src/session-resources.ts b/cactus-code/packages/ai/src/session-resources.ts new file mode 100644 index 000000000..d63009268 --- /dev/null +++ b/cactus-code/packages/ai/src/session-resources.ts @@ -0,0 +1,24 @@ +export type SessionResourceCleanup = (sessionId?: string) => void; + +const sessionResourceCleanups = new Set(); + +export function registerSessionResourceCleanup(cleanup: SessionResourceCleanup): () => void { + sessionResourceCleanups.add(cleanup); + return () => { + sessionResourceCleanups.delete(cleanup); + }; +} + +export function cleanupSessionResources(sessionId?: string): void { + const errors: unknown[] = []; + for (const cleanup of sessionResourceCleanups) { + try { + cleanup(sessionId); + } catch (error) { + errors.push(error); + } + } + if (errors.length > 0) { + throw new AggregateError(errors, "Failed to cleanup session resources"); + } +} diff --git a/cactus-code/packages/ai/src/types.ts b/cactus-code/packages/ai/src/types.ts new file mode 100644 index 000000000..0d388db7b --- /dev/null +++ b/cactus-code/packages/ai/src/types.ts @@ -0,0 +1,638 @@ +import type { OpenAICompletionsOptions } from "./api/openai-completions.ts"; +import type { AssistantMessageDiagnostic } from "./utils/diagnostics.ts"; +import type { AssistantMessageEventStream } from "./utils/event-stream.ts"; + +export type { AssistantMessageEventStream } from "./utils/event-stream.ts"; + +export type KnownApi = "openai-completions"; + +export type Api = KnownApi | (string & {}); + +export type KnownImagesApi = "openrouter-images"; + +export type ImagesApi = KnownImagesApi | (string & {}); + +export type KnownProvider = "cactus"; +export type ProviderId = KnownProvider | string; + +export type KnownImagesProvider = "openrouter"; + +export type ImagesProviderId = KnownImagesProvider | string; + +export type ThinkingLevel = "minimal" | "low" | "medium" | "high" | "xhigh"; +export type ModelThinkingLevel = "off" | ThinkingLevel; +export type ThinkingLevelMap = Partial>; +export type ChatTemplateKwargValue = + | string + | number + | boolean + | null + | { + $var: "thinking.enabled" | "thinking.effort"; + omitWhenOff?: boolean; + }; + +/** Token budgets for each thinking level (token-based providers only) */ +export interface ThinkingBudgets { + minimal?: number; + low?: number; + medium?: number; + high?: number; +} + +// Base options all providers share +export type CacheRetention = "none" | "short" | "long"; + +export type Transport = "sse" | "websocket" | "websocket-cached" | "auto"; + +/** Provider-scoped environment overrides. Values take precedence over process.env. */ +export type ProviderEnv = Record; +export type ProviderHeaders = Record; + +export interface ProviderResponse { + status: number; + headers: Record; +} + +export interface StreamOptions { + temperature?: number; + maxTokens?: number; + signal?: AbortSignal; + apiKey?: string; + /** + * Preferred transport for providers that support multiple transports. + * Providers that do not support this option ignore it. + */ + transport?: Transport; + /** + * Prompt cache retention preference. Providers map this to their supported values. + * Default: "short". + */ + cacheRetention?: CacheRetention; + /** + * Optional session identifier for providers that support session-based caching. + * Providers can use this to enable prompt caching, request routing, or other + * session-aware features. Ignored by providers that don't support it. + */ + sessionId?: string; + /** + * Optional callback for inspecting or replacing provider payloads before sending. + * Return undefined to keep the payload unchanged. + */ + onPayload?: (payload: unknown, model: Model) => unknown | undefined | Promise; + /** + * Optional callback invoked after an HTTP response is received and before + * its body stream is consumed. + */ + onResponse?: (response: ProviderResponse, model: Model) => void | Promise; + /** + * Optional custom HTTP headers to include in API requests. + * Merged with provider defaults; caller values override default headers. + * On AWS Bedrock these are injected via a Smithy `build`-step middleware so + * they are covered by SigV4 signing; reserved headers (`x-amz-*`, + * `authorization`, `host`) are silently ignored to preserve SigV4 / bearer auth. + * A null value suppresses a provider/API default header with the same name. + */ + headers?: ProviderHeaders; + /** + * HTTP request timeout in milliseconds for providers/SDKs that support it. + * For example, OpenAI and Anthropic SDK clients default to 10 minutes. + */ + timeoutMs?: number; + /** + * WebSocket connect timeout in milliseconds for providers that support + * WebSocket transports. This covers the connection/open handshake only; + * stream idleness after connection uses timeoutMs. + */ + websocketConnectTimeoutMs?: number; + /** + * Maximum retry attempts for providers/SDKs that support client-side retries. + * For example, OpenAI and Anthropic SDK clients default to 2. + */ + maxRetries?: number; + /** + * Maximum delay in milliseconds to wait for a retry when the server requests a long wait. + * If the server's requested delay exceeds this value, the request fails immediately + * with an error containing the requested delay, allowing higher-level retry logic + * to handle it with user visibility. + * Default: 60000 (60 seconds). Set to 0 to disable the cap. + */ + maxRetryDelayMs?: number; + /** + * Optional metadata to include in API requests. + * Providers extract the fields they understand and ignore the rest. + * For example, Anthropic uses `user_id` for abuse tracking and rate limiting. + */ + metadata?: Record; + /** + * Provider-scoped environment values. These take precedence over process.env for + * provider configuration such as regional settings, endpoint placeholders, and + * proxy variables. + */ + env?: ProviderEnv; +} + +export type ProviderStreamOptions = StreamOptions & Record; + +/** + * Maps known APIs to their full provider-specific stream option types. + * Type-only imports from API implementation modules are erased at emit, so + * this is tree-shake safe. + */ +export interface ApiOptionsMap { + "openai-completions": OpenAICompletionsOptions; +} + +/** + * Full stream options for an API. Known APIs resolve to their concrete option + * type; custom API strings fall back to the generic shape. + */ +export type ApiStreamOptions = TApi extends keyof ApiOptionsMap + ? ApiOptionsMap[TApi] + : StreamOptions & Record; + +/** + * The uniform stream contract of an API implementation module: every module + * under `src/api/` exports exactly `stream` and `streamSimple`, so the module + * itself satisfies this interface. Lazy wrappers (`lazyApi()`) and provider + * factories pass these around as values. This is the untyped dispatch shape; + * per-API option typing lives on the implementation modules themselves and on + * `Provider.stream()` via `ApiStreamOptions`. + */ +export interface ProviderStreams { + stream(model: Model, context: Context, options?: StreamOptions): AssistantMessageEventStream; + streamSimple(model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream; +} + +/** + * The uniform contract of an image-generation API implementation module: + * every image API module under `src/api/` exports exactly `generateImages`, + * so the module itself satisfies this interface. Lazy wrappers and image + * provider factories pass these around as values. + */ +export interface ProviderImages { + generateImages( + model: ImagesModel, + context: ImagesContext, + options?: ImagesOptions, + ): Promise; +} + +export interface ImagesOptions { + signal?: AbortSignal; + apiKey?: string; + /** + * Provider-scoped environment values. These take precedence over process.env for + * provider configuration such as endpoint placeholders and proxy variables. + */ + env?: ProviderEnv; + /** + * Optional callback for inspecting or replacing provider payloads before sending. + * Return undefined to keep the payload unchanged. + */ + onPayload?: (payload: unknown, model: ImagesModel) => unknown | undefined | Promise; + /** + * Optional callback invoked after an HTTP response is received. + */ + onResponse?: (response: ProviderResponse, model: ImagesModel) => void | Promise; + /** + * Optional custom HTTP headers to include in API requests. + * Merged with provider defaults; can override default headers. + * A null value suppresses a provider/API default header with the same name. + */ + headers?: ProviderHeaders; + /** + * HTTP request timeout in milliseconds for providers/SDKs that support it. + */ + timeoutMs?: number; + /** + * Maximum retry attempts for providers/SDKs that support client-side retries. + */ + maxRetries?: number; + /** + * Maximum delay in milliseconds to wait for a retry when the server requests a long wait. + * If the server's requested delay exceeds this value, the request fails immediately + * with an error containing the requested delay, allowing higher-level retry logic + * to handle it with user visibility. + * Default: 60000 (60 seconds). Set to 0 to disable the cap. + */ + maxRetryDelayMs?: number; + /** + * Optional metadata to include in API requests. + * Providers extract the fields they understand and ignore the rest. + */ + metadata?: Record; +} + +export type ProviderImagesOptions = ImagesOptions & Record; + +// Unified options with reasoning passed to streamSimple() and completeSimple() +export interface SimpleStreamOptions extends StreamOptions { + reasoning?: ThinkingLevel; + /** Custom token budgets for thinking levels (token-based providers only) */ + thinkingBudgets?: ThinkingBudgets; +} + +// Generic StreamFunction with typed options. +// +// Contract: +// - Must return an AssistantMessageEventStream. +// - Once invoked, request/model/runtime failures should be encoded in the +// returned stream, not thrown. +// - Error termination must produce an AssistantMessage with stopReason +// "error" or "aborted" and errorMessage, emitted via the stream protocol. +export type StreamFunction = ( + model: Model, + context: Context, + options?: TOptions, +) => AssistantMessageEventStream; + +export type ImagesFunction = ( + model: ImagesModel, + context: ImagesContext, + options?: TOptions, +) => Promise; + +export interface TextSignatureV1 { + v: 1; + id: string; + phase?: "commentary" | "final_answer"; +} + +export interface TextContent { + type: "text"; + text: string; + textSignature?: string; // e.g., for OpenAI responses, message metadata (legacy id string or TextSignatureV1 JSON) +} + +export interface ThinkingContent { + type: "thinking"; + thinking: string; + thinkingSignature?: string; // e.g., for OpenAI responses, the reasoning item ID + /** When true, the thinking content was redacted by safety filters. The opaque + * encrypted payload is stored in `thinkingSignature` so it can be passed back + * to the API for multi-turn continuity. */ + redacted?: boolean; +} + +export interface ImageContent { + type: "image"; + data: string; // base64 encoded image data + mimeType: string; // e.g., "image/jpeg", "image/png" +} + +export interface ToolCall { + type: "toolCall"; + id: string; + name: string; + arguments: Record; + thoughtSignature?: string; // Google-specific: opaque signature for reusing thought context +} + +export interface Usage { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + /** Subset of `cacheWrite` written with 1h retention. Only Anthropic reports this split. */ + cacheWrite1h?: number; + totalTokens: number; + cost: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; +} + +export type StopReason = "stop" | "length" | "toolUse" | "error" | "aborted"; + +export interface UserMessage { + role: "user"; + content: string | (TextContent | ImageContent)[]; + timestamp: number; // Unix timestamp in milliseconds +} + +export interface AssistantMessage { + role: "assistant"; + content: (TextContent | ThinkingContent | ToolCall)[]; + api: Api; + provider: ProviderId; + model: string; + responseModel?: string; // Concrete `chunk.model` when different from the requested `model` (e.g. OpenRouter `auto` -> `anthropic/...`) + responseId?: string; // Provider-specific response/message identifier when the upstream API exposes one + diagnostics?: AssistantMessageDiagnostic[]; // Redacted provider/runtime diagnostics for failures and recoveries. + cloudHandoff?: boolean; // Cactus hybrid: true if this response was produced by the cloud model, false if local, undefined if unknown + usage: Usage; + stopReason: StopReason; + errorMessage?: string; + timestamp: number; // Unix timestamp in milliseconds +} + +export interface ToolResultMessage { + role: "toolResult"; + toolCallId: string; + toolName: string; + content: (TextContent | ImageContent)[]; // Supports text and images + details?: TDetails; + isError: boolean; + timestamp: number; // Unix timestamp in milliseconds +} + +export type Message = UserMessage | AssistantMessage | ToolResultMessage; + +export type ImagesInputContent = TextContent | ImageContent; +export type ImagesOutputContent = TextContent | ImageContent; + +export interface ImagesContext { + input: ImagesInputContent[]; +} + +export type ImagesStopReason = "stop" | "error" | "aborted"; + +export interface AssistantImages { + api: ImagesApi; + provider: ImagesProviderId; + model: string; + output: ImagesOutputContent[]; + responseId?: string; + usage?: Usage; + stopReason: ImagesStopReason; + errorMessage?: string; + timestamp: number; // Unix timestamp in milliseconds +} + +import type { TSchema } from "typebox"; + +export interface Tool { + name: string; + description: string; + parameters: TParameters; +} + +export interface Context { + systemPrompt?: string; + messages: Message[]; + tools?: Tool[]; +} + +/** + * Event protocol for AssistantMessageEventStream. + * + * Streams should emit `start` before partial updates, then terminate with either: + * - `done` carrying the final successful AssistantMessage, or + * - `error` carrying the final AssistantMessage with stopReason "error" or "aborted" + * and errorMessage. + */ +export type AssistantMessageEvent = + | { type: "start"; partial: AssistantMessage } + | { type: "text_start"; contentIndex: number; partial: AssistantMessage } + | { type: "text_delta"; contentIndex: number; delta: string; partial: AssistantMessage } + | { type: "text_end"; contentIndex: number; content: string; partial: AssistantMessage } + | { type: "thinking_start"; contentIndex: number; partial: AssistantMessage } + | { type: "thinking_delta"; contentIndex: number; delta: string; partial: AssistantMessage } + | { type: "thinking_end"; contentIndex: number; content: string; partial: AssistantMessage } + | { type: "toolcall_start"; contentIndex: number; partial: AssistantMessage } + | { type: "toolcall_delta"; contentIndex: number; delta: string; partial: AssistantMessage } + | { type: "toolcall_end"; contentIndex: number; toolCall: ToolCall; partial: AssistantMessage } + | { type: "done"; reason: Extract; message: AssistantMessage } + | { type: "error"; reason: Extract; error: AssistantMessage }; + +/** + * Compatibility settings for OpenAI-compatible completions APIs. + * Use this to override URL-based auto-detection for custom providers. + */ +export interface OpenAICompletionsCompat { + /** Whether the provider supports the `store` field. Default: auto-detected from URL. */ + supportsStore?: boolean; + /** Whether the provider supports the `developer` role (vs `system`). Default: auto-detected from URL. */ + supportsDeveloperRole?: boolean; + /** Whether the provider supports `reasoning_effort`. Default: auto-detected from URL. */ + supportsReasoningEffort?: boolean; + /** Whether the provider supports `stream_options: { include_usage: true }` for token usage in streaming responses. Default: true. */ + supportsUsageInStreaming?: boolean; + /** Which field to use for max tokens. Default: auto-detected from URL. */ + maxTokensField?: "max_completion_tokens" | "max_tokens"; + /** Whether tool results require the `name` field. Default: auto-detected from URL. */ + requiresToolResultName?: boolean; + /** Whether a user message after tool results requires an assistant message in between. Default: auto-detected from URL. */ + requiresAssistantAfterToolResult?: boolean; + /** Whether thinking blocks must be converted to text blocks with delimiters. Default: auto-detected from URL. */ + requiresThinkingAsText?: boolean; + /** Whether all replayed assistant messages must include an empty reasoning_content field when reasoning is enabled. Default: auto-detected from URL. */ + requiresReasoningContentOnAssistantMessages?: boolean; + /** Format for reasoning/thinking parameter. "openai" uses reasoning_effort, "openrouter" uses reasoning: { effort }, "deepseek" uses thinking: { type } plus reasoning_effort when supported, "together" uses reasoning: { enabled } plus reasoning_effort when supported, "zai" uses thinking: { type }, "qwen" uses top-level enable_thinking: boolean, "qwen-chat-template" uses chat_template_kwargs.enable_thinking and preserve_thinking, "chat-template" uses configurable chat_template_kwargs, "string-thinking" uses top-level thinking: string, and "ant-ling" uses reasoning: { effort } only when the mapped effort is non-null. Default: "openai". */ + thinkingFormat?: + | "openai" + | "openrouter" + | "deepseek" + | "together" + | "zai" + | "qwen" + | "chat-template" + | "qwen-chat-template" + | "string-thinking" + | "ant-ling"; + /** Kwargs to send as `chat_template_kwargs` when `thinkingFormat` is `chat-template`. Use `{ "$var": "thinking.enabled" }` or `{ "$var": "thinking.effort" }` for pi-controlled thinking values. */ + chatTemplateKwargs?: Record; + /** OpenRouter-compatible routing preferences sent as the `provider` request field. */ + openRouterRouting?: OpenRouterRouting; + /** Vercel AI Gateway routing preferences. Only used when baseUrl points to Vercel AI Gateway. */ + vercelGatewayRouting?: VercelGatewayRouting; + /** Whether z.ai supports top-level `tool_stream: true` for streaming tool call deltas. Default: false. */ + zaiToolStream?: boolean; + /** Whether the provider supports the `strict` field in tool definitions. Default: true. */ + supportsStrictMode?: boolean; + /** Cache control convention for prompt caching. "anthropic" applies Anthropic-style `cache_control` markers to the system prompt, last tool definition, and last user/assistant text content. */ + cacheControlFormat?: "anthropic"; + /** Whether to send known session-affinity headers (`session_id`, `x-client-request-id`, `x-session-affinity`) from `options.sessionId` when caching is enabled. Default: false. */ + sendSessionAffinityHeaders?: boolean; + /** Whether the provider supports long prompt cache retention (`prompt_cache_retention: "24h"` or Anthropic-style `cache_control.ttl: "1h"`, depending on format). Default: true. */ + supportsLongCacheRetention?: boolean; +} + +/** Compatibility settings for OpenAI Responses APIs. */ +export interface OpenAIResponsesCompat { + /** Whether the provider supports the `developer` role (vs `system`). Default: true. */ + supportsDeveloperRole?: boolean; + /** Whether to send the OpenAI `session_id` cache-affinity header from `options.sessionId` when caching is enabled. Default: true. */ + sendSessionIdHeader?: boolean; + /** Whether the provider supports `prompt_cache_retention: "24h"`. Default: true. */ + supportsLongCacheRetention?: boolean; +} + +/** Compatibility settings for Anthropic Messages-compatible APIs. */ +export interface AnthropicMessagesCompat { + /** + * Whether the provider accepts per-tool `eager_input_streaming`. + * When false, the Anthropic provider omits `tools[].eager_input_streaming` + * and sends the legacy `fine-grained-tool-streaming-2025-05-14` beta header + * for tool-enabled requests. + * Default: true. + */ + supportsEagerToolInputStreaming?: boolean; + /** Whether the provider supports Anthropic long cache retention (`cache_control.ttl: "1h"`). Default: true. */ + supportsLongCacheRetention?: boolean; + /** + * Whether to send the `x-session-affinity` header from `options.sessionId` + * when caching is enabled. Required for providers like Fireworks that use + * session affinity for prompt cache routing (requests to the same replica + * maximize cache hits). + * Default: false. + */ + sendSessionAffinityHeaders?: boolean; + /** + * Whether the provider supports Anthropic-style `cache_control` markers on + * tool definitions. When false, `cache_control` is omitted from tool params. + * Some Anthropic-compatible providers (e.g., Fireworks) do not support this + * field on tools and may reject or ignore it. + * Default: true. + */ + supportsCacheControlOnTools?: boolean; + /** + * Whether the model accepts the Anthropic `temperature` request field. + * Claude Opus 4.7+ rejects non-default temperature values. + * Default: true. + */ + supportsTemperature?: boolean; + /** + * Whether to force adaptive thinking (`thinking.type: "adaptive"` plus + * `output_config.effort`) regardless of the model id. Built-in models that + * require adaptive thinking set this in generated metadata. Custom + * Anthropic-compatible providers can set this to `true` for any model whose + * upstream requires the adaptive format. Set to `false` to + * opt out on overridden built-in models. + * Default: false. + */ + forceAdaptiveThinking?: boolean; + /** Whether to replay empty thinking signatures as `signature: ""` instead of converting thinking to text. Default: false. */ + allowEmptySignature?: boolean; +} + +/** + * OpenRouter provider routing preferences. + * Controls which upstream providers OpenRouter routes requests to. + * Sent as the `provider` field in the OpenRouter API request body. + * @see https://openrouter.ai/docs/guides/routing/provider-selection + */ +export interface OpenRouterRouting { + /** Whether to allow backup providers to serve requests. Default: true. */ + allow_fallbacks?: boolean; + /** Whether to filter providers to only those that support all parameters in the request. Default: false. */ + require_parameters?: boolean; + /** Data collection setting. "allow" (default): allow providers that may store/train on data. "deny": only use providers that don't collect user data. */ + data_collection?: "deny" | "allow"; + /** Whether to restrict routing to only ZDR (Zero Data Retention) endpoints. */ + zdr?: boolean; + /** Whether to restrict routing to only models that allow text distillation. */ + enforce_distillable_text?: boolean; + /** An ordered list of provider names/slugs to try in sequence, falling back to the next if unavailable. */ + order?: string[]; + /** List of provider names/slugs to exclusively allow for this request. */ + only?: string[]; + /** List of provider names/slugs to skip for this request. */ + ignore?: string[]; + /** A list of quantization levels to filter providers by (e.g., ["fp16", "bf16", "fp8", "fp6", "int8", "int4", "fp4", "fp32"]). */ + quantizations?: string[]; + /** Sorting strategy. Can be a string (e.g., "price", "throughput", "latency") or an object with `by` and `partition`. */ + sort?: + | string + | { + /** The sorting metric: "price", "throughput", "latency". */ + by?: string; + /** Partitioning strategy: "model" (default) or "none". */ + partition?: string | null; + }; + /** Maximum price per million tokens (USD). */ + max_price?: { + /** Price per million prompt tokens. */ + prompt?: number | string; + /** Price per million completion tokens. */ + completion?: number | string; + /** Price per image. */ + image?: number | string; + /** Price per audio unit. */ + audio?: number | string; + /** Price per request. */ + request?: number | string; + }; + /** Preferred minimum throughput (tokens/second). Can be a number (applies to p50) or an object with percentile-specific cutoffs. */ + preferred_min_throughput?: + | number + | { + /** Minimum tokens/second at the 50th percentile. */ + p50?: number; + /** Minimum tokens/second at the 75th percentile. */ + p75?: number; + /** Minimum tokens/second at the 90th percentile. */ + p90?: number; + /** Minimum tokens/second at the 99th percentile. */ + p99?: number; + }; + /** Preferred maximum latency (seconds). Can be a number (applies to p50) or an object with percentile-specific cutoffs. */ + preferred_max_latency?: + | number + | { + /** Maximum latency in seconds at the 50th percentile. */ + p50?: number; + /** Maximum latency in seconds at the 75th percentile. */ + p75?: number; + /** Maximum latency in seconds at the 90th percentile. */ + p90?: number; + /** Maximum latency in seconds at the 99th percentile. */ + p99?: number; + }; +} + +/** + * Vercel AI Gateway routing preferences. + * Controls which upstream providers the gateway routes requests to. + * @see https://vercel.com/docs/ai-gateway/models-and-providers/provider-options + */ +export interface VercelGatewayRouting { + /** List of provider slugs to exclusively use for this request (e.g., ["bedrock", "anthropic"]). */ + only?: string[]; + /** List of provider slugs to try in order (e.g., ["anthropic", "openai"]). */ + order?: string[]; +} + +// Model interface for the unified model system +export interface Model { + id: string; + name: string; + api: TApi; + provider: ProviderId; + baseUrl: string; + reasoning: boolean; + /** + * Maps pi thinking levels to provider/model-specific values. + * Missing keys use provider defaults. null marks a level as unsupported. + */ + thinkingLevelMap?: ThinkingLevelMap; + input: ("text" | "image")[]; + cost: { + input: number; // $/million tokens + output: number; // $/million tokens + cacheRead: number; // $/million tokens + cacheWrite: number; // $/million tokens + }; + contextWindow: number; + maxTokens: number; + headers?: Record; + /** Compatibility overrides for OpenAI-compatible APIs. If not set, auto-detected from baseUrl. */ + compat?: TApi extends "openai-completions" + ? OpenAICompletionsCompat + : TApi extends "openai-responses" + ? OpenAIResponsesCompat + : TApi extends "anthropic-messages" + ? AnthropicMessagesCompat + : never; +} + +export interface ImagesModel + extends Omit, "api" | "provider" | "reasoning" | "contextWindow" | "maxTokens" | "compat"> { + api: TApi; + provider: ImagesProviderId; + output: ("text" | "image")[]; +} diff --git a/cactus-code/packages/ai/src/utils/abort-signals.ts b/cactus-code/packages/ai/src/utils/abort-signals.ts new file mode 100644 index 000000000..8e1a9ade9 --- /dev/null +++ b/cactus-code/packages/ai/src/utils/abort-signals.ts @@ -0,0 +1,41 @@ +export interface CombinedAbortSignal { + signal?: AbortSignal; + cleanup: () => void; +} + +export function combineAbortSignals(signals: readonly (AbortSignal | undefined)[]): CombinedAbortSignal { + const activeSignals = signals.filter((signal): signal is AbortSignal => signal !== undefined); + if (activeSignals.length === 0) { + return { cleanup: () => {} }; + } + if (activeSignals.length === 1) { + return { signal: activeSignals[0], cleanup: () => {} }; + } + + const controller = new AbortController(); + const listeners: Array<{ signal: AbortSignal; listener: () => void }> = []; + const abort = (signal: AbortSignal) => { + if (!controller.signal.aborted) { + controller.abort(signal.reason); + } + }; + + for (const signal of activeSignals) { + if (signal.aborted) { + abort(signal); + break; + } + const listener = () => abort(signal); + signal.addEventListener("abort", listener, { once: true }); + listeners.push({ signal, listener }); + } + + return { + signal: controller.signal, + cleanup: () => { + for (const { signal, listener } of listeners) { + signal.removeEventListener("abort", listener); + } + }, + }; +} diff --git a/cactus-code/packages/ai/src/utils/diagnostics.ts b/cactus-code/packages/ai/src/utils/diagnostics.ts new file mode 100644 index 000000000..89d135619 --- /dev/null +++ b/cactus-code/packages/ai/src/utils/diagnostics.ts @@ -0,0 +1,45 @@ +export interface DiagnosticErrorInfo { + name?: string; + message: string; + stack?: string; + code?: string | number; +} + +export interface AssistantMessageDiagnostic { + type: string; + timestamp: number; + error?: DiagnosticErrorInfo; + details?: Record; +} + +export function formatThrownValue(value: unknown): string { + if (value instanceof Error) return value.message || value.name; + if (typeof value === "string") return value; + return String(value); +} + +export function extractDiagnosticError(error: unknown): DiagnosticErrorInfo { + if (!(error instanceof Error)) return { name: "ThrownValue", message: formatThrownValue(error) }; + const code = (error as Error & { code?: unknown }).code; + return { + name: error.name || undefined, + message: error.message || error.name, + stack: error.stack, + code: typeof code === "string" || typeof code === "number" ? code : undefined, + }; +} + +export function createAssistantMessageDiagnostic( + type: string, + error: unknown, + details?: Record, +): AssistantMessageDiagnostic { + return { type, timestamp: Date.now(), error: extractDiagnosticError(error), details }; +} + +export function appendAssistantMessageDiagnostic( + message: T, + diagnostic: AssistantMessageDiagnostic, +): void { + message.diagnostics = [...(message.diagnostics ?? []), diagnostic]; +} diff --git a/cactus-code/packages/ai/src/utils/event-stream.ts b/cactus-code/packages/ai/src/utils/event-stream.ts new file mode 100644 index 000000000..e5eef921d --- /dev/null +++ b/cactus-code/packages/ai/src/utils/event-stream.ts @@ -0,0 +1,88 @@ +import type { AssistantMessage, AssistantMessageEvent } from "../types.ts"; + +// Generic event stream class for async iteration +export class EventStream implements AsyncIterable { + private queue: T[] = []; + private waiting: ((value: IteratorResult) => void)[] = []; + private done = false; + private finalResultPromise: Promise; + private resolveFinalResult!: (result: R) => void; + private isComplete: (event: T) => boolean; + private extractResult: (event: T) => R; + + constructor(isComplete: (event: T) => boolean, extractResult: (event: T) => R) { + this.isComplete = isComplete; + this.extractResult = extractResult; + this.finalResultPromise = new Promise((resolve) => { + this.resolveFinalResult = resolve; + }); + } + + push(event: T): void { + if (this.done) return; + + if (this.isComplete(event)) { + this.done = true; + this.resolveFinalResult(this.extractResult(event)); + } + + // Deliver to waiting consumer or queue it + const waiter = this.waiting.shift(); + if (waiter) { + waiter({ value: event, done: false }); + } else { + this.queue.push(event); + } + } + + end(result?: R): void { + this.done = true; + if (result !== undefined) { + this.resolveFinalResult(result); + } + // Notify all waiting consumers that we're done + while (this.waiting.length > 0) { + const waiter = this.waiting.shift()!; + waiter({ value: undefined as any, done: true }); + } + } + + async *[Symbol.asyncIterator](): AsyncIterator { + while (true) { + if (this.queue.length > 0) { + yield this.queue.shift()!; + } else if (this.done) { + return; + } else { + const result = await new Promise>((resolve) => this.waiting.push(resolve)); + if (result.done) return; + yield result.value; + } + } + } + + result(): Promise { + return this.finalResultPromise; + } +} + +export class AssistantMessageEventStream extends EventStream { + constructor() { + super( + (event) => event.type === "done" || event.type === "error", + (event) => { + if (event.type === "done") { + return event.message; + } else if (event.type === "error") { + return event.error; + } + throw new Error("Unexpected event type for final result"); + }, + ); + } +} + +/** Factory function for AssistantMessageEventStream (for use in extensions) */ +export function createAssistantMessageEventStream(): AssistantMessageEventStream { + return new AssistantMessageEventStream(); +} diff --git a/cactus-code/packages/ai/src/utils/hash.ts b/cactus-code/packages/ai/src/utils/hash.ts new file mode 100644 index 000000000..1ff55e8b4 --- /dev/null +++ b/cactus-code/packages/ai/src/utils/hash.ts @@ -0,0 +1,13 @@ +/** Fast deterministic hash to shorten long strings */ +export function shortHash(str: string): string { + let h1 = 0xdeadbeef; + let h2 = 0x41c6ce57; + for (let i = 0; i < str.length; i++) { + const ch = str.charCodeAt(i); + h1 = Math.imul(h1 ^ ch, 2654435761); + h2 = Math.imul(h2 ^ ch, 1597334677); + } + h1 = Math.imul(h1 ^ (h1 >>> 16), 2246822507) ^ Math.imul(h2 ^ (h2 >>> 13), 3266489909); + h2 = Math.imul(h2 ^ (h2 >>> 16), 2246822507) ^ Math.imul(h1 ^ (h1 >>> 13), 3266489909); + return (h2 >>> 0).toString(36) + (h1 >>> 0).toString(36); +} diff --git a/cactus-code/packages/ai/src/utils/headers.ts b/cactus-code/packages/ai/src/utils/headers.ts new file mode 100644 index 000000000..2d923d287 --- /dev/null +++ b/cactus-code/packages/ai/src/utils/headers.ts @@ -0,0 +1,18 @@ +import type { ProviderHeaders } from "../types.ts"; + +export function headersToRecord(headers: Headers): Record { + const result: Record = {}; + for (const [key, value] of headers.entries()) { + result[key] = value; + } + return result; +} + +export function providerHeadersToRecord(headers: ProviderHeaders | undefined): Record | undefined { + if (!headers) return undefined; + const result: Record = {}; + for (const [key, value] of Object.entries(headers)) { + if (value !== null) result[key] = value; + } + return Object.keys(result).length > 0 ? result : undefined; +} diff --git a/cactus-code/packages/ai/src/utils/json-parse.ts b/cactus-code/packages/ai/src/utils/json-parse.ts new file mode 100644 index 000000000..f39f9f09d --- /dev/null +++ b/cactus-code/packages/ai/src/utils/json-parse.ts @@ -0,0 +1,124 @@ +import { parse as partialParse } from "partial-json"; + +const VALID_JSON_ESCAPES = new Set(['"', "\\", "/", "b", "f", "n", "r", "t", "u"]); + +function isControlCharacter(char: string): boolean { + const codePoint = char.codePointAt(0); + return codePoint !== undefined && codePoint >= 0x00 && codePoint <= 0x1f; +} + +function escapeControlCharacter(char: string): string { + switch (char) { + case "\b": + return "\\b"; + case "\f": + return "\\f"; + case "\n": + return "\\n"; + case "\r": + return "\\r"; + case "\t": + return "\\t"; + default: + return `\\u${char.codePointAt(0)?.toString(16).padStart(4, "0") ?? "0000"}`; + } +} + +/** + * Repairs malformed JSON string literals by: + * - escaping raw control characters inside strings + * - doubling backslashes before invalid escape characters + */ +export function repairJson(json: string): string { + let repaired = ""; + let inString = false; + + for (let index = 0; index < json.length; index++) { + const char = json[index]; + + if (!inString) { + repaired += char; + if (char === '"') { + inString = true; + } + continue; + } + + if (char === '"') { + repaired += char; + inString = false; + continue; + } + + if (char === "\\") { + const nextChar = json[index + 1]; + if (nextChar === undefined) { + repaired += "\\\\"; + continue; + } + + if (nextChar === "u") { + const unicodeDigits = json.slice(index + 2, index + 6); + if (/^[0-9a-fA-F]{4}$/.test(unicodeDigits)) { + repaired += `\\u${unicodeDigits}`; + index += 5; + continue; + } + } + + if (VALID_JSON_ESCAPES.has(nextChar)) { + repaired += `\\${nextChar}`; + index += 1; + continue; + } + + repaired += "\\\\"; + continue; + } + + repaired += isControlCharacter(char) ? escapeControlCharacter(char) : char; + } + + return repaired; +} + +export function parseJsonWithRepair(json: string): T { + try { + return JSON.parse(json) as T; + } catch (error) { + const repairedJson = repairJson(json); + if (repairedJson !== json) { + return JSON.parse(repairedJson) as T; + } + throw error; + } +} + +/** + * Attempts to parse potentially incomplete JSON during streaming. + * Always returns a valid object, even if the JSON is incomplete. + * + * @param partialJson The partial JSON string from streaming + * @returns Parsed object or empty object if parsing fails + */ +export function parseStreamingJson>(partialJson: string | undefined): T { + if (!partialJson || partialJson.trim() === "") { + return {} as T; + } + + try { + return parseJsonWithRepair(partialJson); + } catch { + try { + const result = partialParse(partialJson); + return (result ?? {}) as T; + } catch { + try { + const result = partialParse(repairJson(partialJson)); + return (result ?? {}) as T; + } catch { + return {} as T; + } + } + } +} diff --git a/cactus-code/packages/ai/src/utils/oauth/types.ts b/cactus-code/packages/ai/src/utils/oauth/types.ts new file mode 100644 index 000000000..008be405e --- /dev/null +++ b/cactus-code/packages/ai/src/utils/oauth/types.ts @@ -0,0 +1,79 @@ +import type { Api, Model } from "../../types.ts"; + +export type OAuthCredentials = { + refresh: string; + access: string; + expires: number; + [key: string]: unknown; +}; + +export type OAuthProviderId = string; + +/** @deprecated Use OAuthProviderId instead */ +export type OAuthProvider = OAuthProviderId; + +export type OAuthPrompt = { + message: string; + placeholder?: string; + allowEmpty?: boolean; +}; + +export type OAuthAuthInfo = { + url: string; + instructions?: string; +}; + +export type OAuthDeviceCodeInfo = { + userCode: string; + verificationUri: string; + intervalSeconds?: number; + expiresInSeconds?: number; +}; + +export type OAuthSelectOption = { + id: string; + label: string; +}; + +export type OAuthSelectPrompt = { + message: string; + options: OAuthSelectOption[]; +}; + +export interface OAuthLoginCallbacks { + onAuth: (info: OAuthAuthInfo) => void; + onDeviceCode: (info: OAuthDeviceCodeInfo) => void; + onPrompt: (prompt: OAuthPrompt) => Promise; + onProgress?: (message: string) => void; + onManualCodeInput?: () => Promise; + /** Show an interactive selector and return the selected option id, or undefined on cancel. */ + onSelect: (prompt: OAuthSelectPrompt) => Promise; + signal?: AbortSignal; +} + +export interface OAuthProviderInterface { + readonly id: OAuthProviderId; + readonly name: string; + + /** Run the login flow, return credentials to persist */ + login(callbacks: OAuthLoginCallbacks): Promise; + + /** Whether login uses a local callback server and supports manual code input. */ + usesCallbackServer?: boolean; + + /** Refresh expired credentials, return updated credentials to persist */ + refreshToken(credentials: OAuthCredentials): Promise; + + /** Convert credentials to API key string for the provider */ + getApiKey(credentials: OAuthCredentials): string; + + /** Optional: modify models for this provider (e.g., update baseUrl) */ + modifyModels?(models: Model[], credentials: OAuthCredentials): Model[]; +} + +/** @deprecated Use OAuthProviderInterface instead */ +export interface OAuthProviderInfo { + id: OAuthProviderId; + name: string; + available: boolean; +} diff --git a/cactus-code/packages/ai/src/utils/overflow.ts b/cactus-code/packages/ai/src/utils/overflow.ts new file mode 100644 index 000000000..623b873ab --- /dev/null +++ b/cactus-code/packages/ai/src/utils/overflow.ts @@ -0,0 +1,162 @@ +import type { AssistantMessage } from "../types.ts"; + +/** + * Regex patterns to detect context overflow errors from different providers. + * + * These patterns match error messages returned when the input exceeds + * the model's context window. + * + * Provider-specific patterns (with example error messages): + * + * - Anthropic: "prompt is too long: 213462 tokens > 200000 maximum" + * - Anthropic: "413 {\"error\":{\"type\":\"request_too_large\",\"message\":\"Request exceeds the maximum size\"}}" + * - OpenAI: "Your input exceeds the context window of this model" + * - OpenAI/LiteLLM: "Requested token count exceeds the model's maximum context length of 131072 tokens" + * - OpenAI-compatible: "Input length (265330) exceeds model's maximum context length (262144)." + * - Google: "The input token count (1196265) exceeds the maximum number of tokens allowed (1048575)" + * - xAI: "This model's maximum prompt length is 131072 but the request contains 537812 tokens" + * - Groq: "Please reduce the length of the messages or completion" + * - OpenRouter: "This endpoint's maximum context length is X tokens. However, you requested about Y tokens" + * - OpenRouter/Poolside: "Input length X exceeds the maximum allowed input length of Y tokens." + * - Together AI: "The input (X tokens) is longer than the model's context length (Y tokens)." + * - llama.cpp: "the request exceeds the available context size, try increasing it" + * - LM Studio: "tokens to keep from the initial prompt is greater than the context length" + * - GitHub Copilot: "prompt token count of X exceeds the limit of Y" + * - MiniMax: "invalid params, context window exceeds limit" + * - Kimi For Coding: "Your request exceeded model token limit: X (requested: Y)" + * - Cerebras: "400/413 status code (no body)" + * - Mistral: "Prompt contains X tokens ... too large for model with Y maximum context length" + * - z.ai: Does NOT error, accepts overflow silently - handled via usage.input > contextWindow + * - Xiaomi MiMo: Truncates input to fill contextWindow exactly, then returns finish_reason "length" + * with output=0 (no room left to generate). Detected via stopReason "length" + zero output + + * input filling the context window. + * - Ollama: Some deployments truncate silently, others return errors like "prompt too long; exceeded max context length by X tokens" + */ +const OVERFLOW_PATTERNS = [ + /prompt is too long/i, // Anthropic token overflow + /request_too_large/i, // Anthropic request byte-size overflow (HTTP 413) + /input is too long for requested model/i, // Amazon Bedrock + /exceeds the context window/i, // OpenAI (Completions & Responses API) + /exceeds (?:the )?(?:model'?s )?maximum context length(?: of [\d,]+ tokens?|\s*\([\d,]+\))/i, // OpenAI-compatible proxies (LiteLLM) + /input token count.*exceeds the maximum/i, // Google (Gemini) + /maximum prompt length is \d+/i, // xAI (Grok) + /reduce the length of the messages/i, // Groq + /maximum context length is \d+ tokens/i, // OpenRouter (most backends) + /exceeds (?:the )?maximum allowed input length of [\d,]+ tokens?/i, // OpenRouter/Poolside + /input \(\d+ tokens\) is longer than the model'?s context length \(\d+ tokens\)/i, // Together AI + /exceeds the limit of \d+/i, // GitHub Copilot + /exceeds the available context size/i, // llama.cpp server + /greater than the context length/i, // LM Studio + /context window exceeds limit/i, // MiniMax + /exceeded model token limit/i, // Kimi For Coding + /too large for model with \d+ maximum context length/i, // Mistral + /model_context_window_exceeded/i, // z.ai non-standard finish_reason surfaced as error text + /prompt too long; exceeded (?:max )?context length/i, // Ollama explicit overflow error + /context[_ ]length[_ ]exceeded/i, // Generic fallback + /too many tokens/i, // Generic fallback + /token limit exceeded/i, // Generic fallback + /^4(?:00|13)\s*(?:status code)?\s*\(no body\)/i, // Cerebras: 400/413 with no body +]; + +/** + * Patterns that indicate non-overflow errors (e.g. rate limiting, server errors). + * Error messages matching any of these are excluded from overflow detection + * even if they also match an OVERFLOW_PATTERN. + * + * Example: Bedrock formats throttling errors as "ThrottlingException: Too many tokens, + * please wait before trying again." which would match the /too many tokens/i overflow + * pattern without this exclusion. + */ +const NON_OVERFLOW_PATTERNS = [ + /^(Throttling error|Service unavailable):/i, // AWS Bedrock non-overflow errors (human-readable prefixes from formatBedrockError) + /rate limit/i, // Generic rate limiting + /too many requests/i, // Generic HTTP 429 style +]; + +/** + * Check if an assistant message represents a context overflow error. + * + * This handles two cases: + * 1. Error-based overflow: Most providers return stopReason "error" with a + * specific error message pattern. + * 2. Silent overflow: Some providers accept overflow requests and return + * successfully. For these, we check if usage.input exceeds the context window. + * + * ## Reliability by Provider + * + * **Reliable detection (returns error with detectable message):** + * - Anthropic: "prompt is too long: X tokens > Y maximum" or "request_too_large" + * - OpenAI (Completions & Responses): "exceeds the context window", "exceeds the model's maximum context length of X tokens", or "exceeds model's maximum context length (X)" + * - Google Gemini: "input token count exceeds the maximum" + * - xAI (Grok): "maximum prompt length is X but request contains Y" + * - Groq: "reduce the length of the messages" + * - Cerebras: 400/413 status code (no body) + * - Mistral: "Prompt contains X tokens ... too large for model with Y maximum context length" + * - OpenRouter (most backends): "maximum context length is X tokens" + * - OpenRouter/Poolside: "Input length X exceeds the maximum allowed input length of Y tokens." + * - Together AI: "The input (X tokens) is longer than the model's context length (Y tokens)." + * - llama.cpp: "exceeds the available context size" + * - LM Studio: "greater than the context length" + * - Kimi For Coding: "exceeded model token limit: X (requested: Y)" + * + * **Unreliable detection:** + * - z.ai: Sometimes accepts overflow silently (detectable via usage.input > contextWindow), + * sometimes returns rate limit errors. Pass contextWindow param to detect silent overflow. + * - Xiaomi MiMo: Truncates input to fit contextWindow then returns stopReason "length" with + * output=0. Pass contextWindow param to detect via the "filled context + zero output" signal. + * - Ollama: May truncate input silently for some setups, but may also return explicit + * overflow errors that match the patterns above. Silent truncation still cannot be + * detected here because we do not know the expected token count. + * + * ## Custom Providers + * + * If you've added custom models via settings.json, this function may not detect + * overflow errors from those providers. To add support: + * + * 1. Send a request that exceeds the model's context window + * 2. Check the errorMessage in the response + * 3. Create a regex pattern that matches the error + * 4. The pattern should be added to OVERFLOW_PATTERNS in this file, or + * check the errorMessage yourself before calling this function + * + * @param message - The assistant message to check + * @param contextWindow - Optional context window size for detecting silent overflow (z.ai) + * @returns true if the message indicates a context overflow + */ +export function isContextOverflow(message: AssistantMessage, contextWindow?: number): boolean { + // Case 1: Check error message patterns + if (message.stopReason === "error" && message.errorMessage) { + // Skip messages matching known non-overflow patterns (e.g. throttling / rate-limit) + const isNonOverflow = NON_OVERFLOW_PATTERNS.some((p) => p.test(message.errorMessage!)); + if (!isNonOverflow && OVERFLOW_PATTERNS.some((p) => p.test(message.errorMessage!))) { + return true; + } + } + + // Case 2: Silent overflow (z.ai style) - successful but usage exceeds context + if (contextWindow && message.stopReason === "stop") { + const inputTokens = message.usage.input + message.usage.cacheRead; + if (inputTokens > contextWindow) { + return true; + } + } + + // Case 3: Length-stop overflow (Xiaomi MiMo style) - server truncates oversized input + // to fit the context window, leaving no room for output. Returns stopReason "length" + // with output=0 and input+cacheRead filling the context window. + if (contextWindow && message.stopReason === "length" && message.usage.output === 0) { + const inputTokens = message.usage.input + message.usage.cacheRead; + if (inputTokens >= contextWindow * 0.99) { + return true; + } + } + + return false; +} + +/** + * Get the overflow patterns for testing purposes. + */ +export function getOverflowPatterns(): RegExp[] { + return [...OVERFLOW_PATTERNS]; +} diff --git a/cactus-code/packages/ai/src/utils/provider-env.ts b/cactus-code/packages/ai/src/utils/provider-env.ts new file mode 100644 index 000000000..db4960672 --- /dev/null +++ b/cactus-code/packages/ai/src/utils/provider-env.ts @@ -0,0 +1,52 @@ +import type { ProviderEnv } from "../types.ts"; + +let procEnvCache: Map | null = null; + +/** + * Fallback for https://github.com/oven-sh/bun/issues/27802. + * Bun compiled binaries can expose an empty process.env inside Linux sandboxes + * even though /proc/self/environ contains the environment. + * + * This intentionally duplicates restoreSandboxEnv() in + * packages/coding-agent/src/bun/restore-sandbox-env.ts. The ai package can be + * used directly, without going through that entrypoint, so provider env lookup + * must not depend on process.env having been patched. + */ +function getBunSandboxEnvValue(name: string): string | undefined { + if (typeof process === "undefined" || !process.versions?.bun || Object.keys(process.env).length > 0) { + return undefined; + } + + if (procEnvCache === null) { + procEnvCache = new Map(); + try { + const { readFileSync } = require("node:fs") as { + readFileSync(path: string, encoding: BufferEncoding): string; + }; + const data = readFileSync("/proc/self/environ", "utf-8"); + for (const entry of data.split("\0")) { + const idx = entry.indexOf("="); + if (idx > 0) { + procEnvCache.set(entry.slice(0, idx), entry.slice(idx + 1)); + } + } + } catch { + // /proc/self/environ may not exist or may not be readable. + } + } + + return procEnvCache.get(name); +} + +/** + * Resolve a provider env value from scoped overrides, normal process.env, then + * the duplicated Bun sandbox fallback for direct pi-ai consumers. + */ +export function getProviderEnvValue(name: string, env?: ProviderEnv): string | undefined { + return ( + env?.[name] || + (typeof process !== "undefined" ? process.env[name] : undefined) || + getBunSandboxEnvValue(name) || + undefined + ); +} diff --git a/cactus-code/packages/ai/src/utils/sanitize-unicode.ts b/cactus-code/packages/ai/src/utils/sanitize-unicode.ts new file mode 100644 index 000000000..d869ee9dc --- /dev/null +++ b/cactus-code/packages/ai/src/utils/sanitize-unicode.ts @@ -0,0 +1,25 @@ +/** + * Removes unpaired Unicode surrogate characters from a string. + * + * Unpaired surrogates (high surrogates 0xD800-0xDBFF without matching low surrogates 0xDC00-0xDFFF, + * or vice versa) cause JSON serialization errors in many API providers. + * + * Valid emoji and other characters outside the Basic Multilingual Plane use properly paired + * surrogates and will NOT be affected by this function. + * + * @param text - The text to sanitize + * @returns The sanitized text with unpaired surrogates removed + * + * @example + * // Valid emoji (properly paired surrogates) are preserved + * sanitizeSurrogates("Hello 🙈 World") // => "Hello 🙈 World" + * + * // Unpaired high surrogate is removed + * const unpaired = String.fromCharCode(0xD83D); // high surrogate without low + * sanitizeSurrogates(`Text ${unpaired} here`) // => "Text here" + */ +export function sanitizeSurrogates(text: string): string { + // Replace unpaired high surrogates (0xD800-0xDBFF not followed by low surrogate) + // Replace unpaired low surrogates (0xDC00-0xDFFF not preceded by high surrogate) + return text.replace(/[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?; // "add" | "subtract" | "multiply" | "divide" + */ +export function StringEnum( + values: T, + options?: { description?: string; default?: T[number] }, +): TUnsafe { + return Type.Unsafe({ + type: "string", + enum: values as any, + ...(options?.description && { description: options.description }), + ...(options?.default && { default: options.default }), + }); +} diff --git a/cactus-code/packages/ai/src/utils/validation.ts b/cactus-code/packages/ai/src/utils/validation.ts new file mode 100644 index 000000000..060898c74 --- /dev/null +++ b/cactus-code/packages/ai/src/utils/validation.ts @@ -0,0 +1,324 @@ +import { Compile } from "typebox/compile"; +import type { TLocalizedValidationError } from "typebox/error"; +import { Value } from "typebox/value"; +import type { Tool, ToolCall } from "../types.ts"; + +const validatorCache = new WeakMap>(); +const TYPEBOX_KIND = Symbol.for("TypeBox.Kind"); + +interface JsonSchemaObject { + type?: string | string[]; + properties?: Record; + items?: JsonSchemaObject | JsonSchemaObject[]; + additionalProperties?: boolean | JsonSchemaObject; + allOf?: JsonSchemaObject[]; + anyOf?: JsonSchemaObject[]; + oneOf?: JsonSchemaObject[]; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function isJsonSchemaObject(value: unknown): value is JsonSchemaObject { + return isRecord(value); +} + +function hasTypeBoxMetadata(schema: unknown): boolean { + return isRecord(schema) && Object.getOwnPropertySymbols(schema).includes(TYPEBOX_KIND); +} + +function getSchemaTypes(schema: JsonSchemaObject): string[] { + if (typeof schema.type === "string") { + return [schema.type]; + } + if (Array.isArray(schema.type)) { + return schema.type.filter((type): type is string => typeof type === "string"); + } + return []; +} + +function matchesJsonType(value: unknown, type: string): boolean { + switch (type) { + case "number": + return typeof value === "number"; + case "integer": + return typeof value === "number" && Number.isInteger(value); + case "boolean": + return typeof value === "boolean"; + case "string": + return typeof value === "string"; + case "null": + return value === null; + case "array": + return Array.isArray(value); + case "object": + return isRecord(value) && !Array.isArray(value); + default: + return false; + } +} + +function isValidatorSchema(value: unknown): value is Tool["parameters"] { + return isRecord(value); +} + +function getSubSchemaValidator(schema: JsonSchemaObject): ReturnType | undefined { + if (!isValidatorSchema(schema)) { + return undefined; + } + try { + return getValidator(schema); + } catch { + return undefined; + } +} + +function coercePrimitiveByType(value: unknown, type: string): unknown { + switch (type) { + case "number": { + if (value === null) { + return 0; + } + if (typeof value === "string" && value.trim() !== "") { + const parsed = Number(value); + if (Number.isFinite(parsed)) { + return parsed; + } + } + if (typeof value === "boolean") { + return value ? 1 : 0; + } + return value; + } + case "integer": { + if (value === null) { + return 0; + } + if (typeof value === "string" && value.trim() !== "") { + const parsed = Number(value); + if (Number.isInteger(parsed)) { + return parsed; + } + } + if (typeof value === "boolean") { + return value ? 1 : 0; + } + return value; + } + case "boolean": { + if (value === null) { + return false; + } + if (typeof value === "string") { + if (value === "true") { + return true; + } + if (value === "false") { + return false; + } + } + if (typeof value === "number") { + if (value === 1) { + return true; + } + if (value === 0) { + return false; + } + } + return value; + } + case "string": { + if (value === null) { + return ""; + } + if (typeof value === "number" || typeof value === "boolean") { + return String(value); + } + return value; + } + case "null": { + if (value === "" || value === 0 || value === false) { + return null; + } + return value; + } + default: + return value; + } +} + +function applySchemaObjectCoercion(value: Record, schema: JsonSchemaObject): void { + const properties = schema.properties; + const definedKeys = new Set(properties ? Object.keys(properties) : []); + + if (properties) { + for (const [key, propertySchema] of Object.entries(properties)) { + if (!(key in value)) { + continue; + } + value[key] = coerceWithJsonSchema(value[key], propertySchema); + } + } + + if (schema.additionalProperties && isJsonSchemaObject(schema.additionalProperties)) { + for (const [key, propertyValue] of Object.entries(value)) { + if (definedKeys.has(key)) { + continue; + } + value[key] = coerceWithJsonSchema(propertyValue, schema.additionalProperties); + } + } +} + +function applySchemaArrayCoercion(value: unknown[], schema: JsonSchemaObject): void { + if (Array.isArray(schema.items)) { + for (let index = 0; index < value.length; index++) { + const itemSchema = schema.items[index]; + if (!itemSchema) { + continue; + } + value[index] = coerceWithJsonSchema(value[index], itemSchema); + } + return; + } + + if (isJsonSchemaObject(schema.items)) { + for (let index = 0; index < value.length; index++) { + value[index] = coerceWithJsonSchema(value[index], schema.items); + } + } +} + +function coerceWithUnionSchema(value: unknown, schemas: JsonSchemaObject[]): unknown { + for (const schema of schemas) { + const candidate = structuredClone(value); + const coerced = coerceWithJsonSchema(candidate, schema); + const validator = getSubSchemaValidator(schema); + if (validator?.Check(coerced)) { + return coerced; + } + } + return value; +} + +function coerceWithJsonSchema(value: unknown, schema: JsonSchemaObject): unknown { + let nextValue = value; + + if (Array.isArray(schema.allOf)) { + for (const nested of schema.allOf) { + nextValue = coerceWithJsonSchema(nextValue, nested); + } + } + + if (Array.isArray(schema.anyOf)) { + nextValue = coerceWithUnionSchema(nextValue, schema.anyOf); + } + + if (Array.isArray(schema.oneOf)) { + nextValue = coerceWithUnionSchema(nextValue, schema.oneOf); + } + + const schemaTypes = getSchemaTypes(schema); + const matchesUnionMember = + schemaTypes.length > 1 && schemaTypes.some((schemaType) => matchesJsonType(nextValue, schemaType)); + if (schemaTypes.length > 0 && !matchesUnionMember) { + for (const schemaType of schemaTypes) { + const candidate = coercePrimitiveByType(nextValue, schemaType); + if (candidate !== nextValue) { + nextValue = candidate; + break; + } + } + } + + if (schemaTypes.includes("object") && isRecord(nextValue) && !Array.isArray(nextValue)) { + applySchemaObjectCoercion(nextValue, schema); + } + + if (schemaTypes.includes("array") && Array.isArray(nextValue)) { + applySchemaArrayCoercion(nextValue, schema); + } + + return nextValue; +} + +function getValidator(schema: Tool["parameters"]): ReturnType { + const key = schema as object; + const cached = validatorCache.get(key); + if (cached) { + return cached; + } + const validator = Compile(schema); + validatorCache.set(key, validator); + return validator; +} + +function formatValidationPath(error: TLocalizedValidationError): string { + if (error.keyword === "required") { + const requiredProperties = (error.params as { requiredProperties?: string[] }).requiredProperties; + const requiredProperty = requiredProperties?.[0]; + if (requiredProperty) { + const basePath = error.instancePath.replace(/^\//, "").replace(/\//g, "."); + return basePath ? `${basePath}.${requiredProperty}` : requiredProperty; + } + } + const path = error.instancePath.replace(/^\//, "").replace(/\//g, "."); + return path || "root"; +} + +/** + * Finds a tool by name and validates the tool call arguments against its TypeBox schema + * @param tools Array of tool definitions + * @param toolCall The tool call from the LLM + * @returns The validated arguments + * @throws Error if tool is not found or validation fails + */ +export function validateToolCall(tools: Tool[], toolCall: ToolCall): any { + const tool = tools.find((t) => t.name === toolCall.name); + if (!tool) { + throw new Error(`Tool "${toolCall.name}" not found`); + } + return validateToolArguments(tool, toolCall); +} + +/** + * Validates tool call arguments against the tool's TypeBox schema + * @param tool The tool definition with TypeBox schema + * @param toolCall The tool call from the LLM + * @returns The validated (and potentially coerced) arguments + * @throws Error with formatted message if validation fails + */ +export function validateToolArguments(tool: Tool, toolCall: ToolCall): any { + const args = structuredClone(toolCall.arguments); + Value.Convert(tool.parameters, args); + + const validator = getValidator(tool.parameters); + if (!hasTypeBoxMetadata(tool.parameters) && isJsonSchemaObject(tool.parameters)) { + const coerced = coerceWithJsonSchema(args, tool.parameters); + if (coerced !== args) { + if (isRecord(args) && isRecord(coerced)) { + for (const key of Object.keys(args)) { + delete args[key]; + } + Object.assign(args, coerced); + } else { + return validator.Check(coerced) ? coerced : args; + } + } + } + + if (validator.Check(args)) { + return args; + } + + const errors = + validator + .Errors(args) + .map((error) => ` - ${formatValidationPath(error)}: ${error.message}`) + .join("\n") || "Unknown validation error"; + + const errorMessage = `Validation failed for tool "${toolCall.name}":\n${errors}\n\nReceived arguments:\n${JSON.stringify(toolCall.arguments, null, 2)}`; + + throw new Error(errorMessage); +} diff --git a/cactus-code/packages/ai/tsconfig.build.json b/cactus-code/packages/ai/tsconfig.build.json new file mode 100644 index 000000000..6089faa7f --- /dev/null +++ b/cactus-code/packages/ai/tsconfig.build.json @@ -0,0 +1,9 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "rootDir": "./src" + }, + "include": ["src/**/*.ts"], + "exclude": ["node_modules", "dist", "**/*.d.ts", "src/**/*.d.ts"] +} \ No newline at end of file diff --git a/cactus-code/packages/coding-agent/package.json b/cactus-code/packages/coding-agent/package.json new file mode 100644 index 000000000..34382a6d6 --- /dev/null +++ b/cactus-code/packages/coding-agent/package.json @@ -0,0 +1,90 @@ +{ + "name": "@earendil-works/pi-coding-agent", + "version": "0.80.2", + "description": "Coding agent CLI with read, bash, edit, write tools and session management", + "type": "module", + "piConfig": { + "name": "cactus", + "configDir": ".cactus" + }, + "bin": { + "cactus": "dist/cli.js" + }, + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + } + }, + "files": [ + "dist" + ], + "scripts": { + "clean": "shx rm -rf dist", + "build": "tsgo -p tsconfig.build.json && shx chmod +x dist/cli.js && npm run copy-assets", + "build:binary": "npm --prefix ../tui run build && npm --prefix ../ai run build && npm --prefix ../agent run build && npm run build && bun build --compile ./dist/bun/cli.js ./src/utils/image-resize-worker.ts --outfile dist/pi && npm run copy-binary-assets", + "copy-assets": "shx mkdir -p dist/modes/interactive/theme && shx cp src/modes/interactive/theme/*.json dist/modes/interactive/theme/ && shx mkdir -p dist/modes/interactive/assets && shx cp src/modes/interactive/assets/*.png dist/modes/interactive/assets/", + "copy-binary-assets": "shx cp package.json dist/ && shx cp README.md dist/ && shx cp CHANGELOG.md dist/ && shx mkdir -p dist/theme && shx cp src/modes/interactive/theme/*.json dist/theme/ && shx mkdir -p dist/assets && shx cp src/modes/interactive/assets/*.png dist/assets/ && shx mkdir -p dist/export-html/vendor && shx cp src/core/export-html/template.html dist/export-html/ && shx cp src/core/export-html/vendor/*.js dist/export-html/vendor/ && shx cp -r docs dist/ && shx cp -r examples dist/ && shx cp ../../node_modules/@silvia-odwyer/photon-node/photon_rs_bg.wasm dist/", + "shrinkwrap": "node ../../scripts/generate-coding-agent-shrinkwrap.mjs", + "prepublishOnly": "npm run clean && npm run build && npm run shrinkwrap" + }, + "dependencies": { + "@earendil-works/pi-agent-core": "^0.80.2", + "@earendil-works/pi-ai": "^0.80.2", + "@earendil-works/pi-tui": "^0.80.2", + "@silvia-odwyer/photon-node": "0.3.4", + "chalk": "5.6.2", + "cross-spawn": "7.0.6", + "diff": "8.0.4", + "glob": "13.0.6", + "highlight.js": "10.7.3", + "hosted-git-info": "9.0.3", + "ignore": "7.0.5", + "jiti": "2.7.0", + "minimatch": "10.2.5", + "proper-lockfile": "4.1.2", + "semver": "7.8.0", + "typebox": "1.1.38", + "undici": "8.5.0", + "yaml": "2.9.0" + }, + "overrides": { + "rimraf": "6.1.2", + "gaxios": { + "rimraf": "6.1.2" + } + }, + "optionalDependencies": { + "@mariozechner/clipboard": "0.3.9" + }, + "devDependencies": { + "@types/cross-spawn": "6.0.6", + "@types/diff": "7.0.2", + "@types/hosted-git-info": "3.0.5", + "@types/ms": "2.1.0", + "@types/node": "24.12.4", + "@types/proper-lockfile": "4.1.4", + "@types/semver": "7.7.1", + "shx": "0.4.0" + }, + "keywords": [ + "coding-agent", + "ai", + "llm", + "cli", + "tui", + "agent" + ], + "author": "Mario Zechner", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/coding-agent" + }, + "engines": { + "node": ">=22.19.0" + } +} diff --git a/cactus-code/packages/coding-agent/src/cli.ts b/cactus-code/packages/coding-agent/src/cli.ts new file mode 100644 index 000000000..b0b11ca5c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli.ts @@ -0,0 +1,20 @@ +#!/usr/bin/env node +/** + * CLI entry point for the refactored coding agent. + * Uses main.ts with AgentSession and new mode modules. + * + * Test with: npx tsx src/cli-new.ts [args...] + */ +import { APP_NAME } from "./config.ts"; +import { configureHttpDispatcher } from "./core/http-dispatcher.ts"; +import { main } from "./main.ts"; + +process.title = APP_NAME; +process.env.PI_CODING_AGENT = "true"; +process.emitWarning = (() => {}) as typeof process.emitWarning; + +// Configure undici's global dispatcher before provider SDKs issue requests. +// Runtime settings are applied once SettingsManager has loaded global/project settings. +configureHttpDispatcher(); + +main(process.argv.slice(2)); diff --git a/cactus-code/packages/coding-agent/src/cli/args.ts b/cactus-code/packages/coding-agent/src/cli/args.ts new file mode 100644 index 000000000..0d0a6c60d --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/args.ts @@ -0,0 +1,292 @@ +/** + * CLI argument parsing and help display + */ + +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import chalk from "chalk"; +import { APP_NAME, CONFIG_DIR_NAME, ENV_AGENT_DIR, ENV_SESSION_DIR } from "../config.ts"; +import type { ExtensionFlag } from "../core/extensions/types.ts"; + +export type Mode = "text" | "json"; + +export interface Args { + provider?: string; + model?: string; + apiKey?: string; + systemPrompt?: string; + appendSystemPrompt?: string[]; + thinking?: ThinkingLevel; + continue?: boolean; + resume?: boolean; + help?: boolean; + version?: boolean; + mode?: Mode; + name?: string; + noSession?: boolean; + session?: string; + sessionId?: string; + fork?: string; + sessionDir?: string; + models?: string[]; + tools?: string[]; + excludeTools?: string[]; + noTools?: boolean; + noBuiltinTools?: boolean; + print?: boolean; + noSkills?: boolean; + skills?: string[]; + promptTemplates?: string[]; + noPromptTemplates?: boolean; + themes?: string[]; + noThemes?: boolean; + noContextFiles?: boolean; + listModels?: string | true; + offline?: boolean; + verbose?: boolean; + projectTrustOverride?: boolean; + messages: string[]; + fileArgs: string[]; + /** Unknown flags (potentially extension flags) - map of flag name to value */ + unknownFlags: Map; + diagnostics: Array<{ type: "warning" | "error"; message: string }>; +} + +const VALID_THINKING_LEVELS = ["off", "minimal", "low", "medium", "high", "xhigh"] as const; + +export function isValidThinkingLevel(level: string): level is ThinkingLevel { + return VALID_THINKING_LEVELS.includes(level as ThinkingLevel); +} + +export function parseArgs(args: string[]): Args { + const result: Args = { + messages: [], + fileArgs: [], + unknownFlags: new Map(), + diagnostics: [], + }; + + for (let i = 0; i < args.length; i++) { + const arg = args[i]; + + if (arg === "--help" || arg === "-h") { + result.help = true; + } else if (arg === "--version" || arg === "-v") { + result.version = true; + } else if (arg === "--mode" && i + 1 < args.length) { + const mode = args[++i]; + if (mode === "text" || mode === "json") { + result.mode = mode; + } + } else if (arg === "--continue" || arg === "-c") { + result.continue = true; + } else if (arg === "--resume" || arg === "-r") { + result.resume = true; + } else if (arg === "--provider" && i + 1 < args.length) { + result.provider = args[++i]; + } else if (arg === "--model" && i + 1 < args.length) { + result.model = args[++i]; + } else if (arg === "--api-key" && i + 1 < args.length) { + result.apiKey = args[++i]; + } else if (arg === "--system-prompt" && i + 1 < args.length) { + result.systemPrompt = args[++i]; + } else if (arg === "--append-system-prompt" && i + 1 < args.length) { + result.appendSystemPrompt = result.appendSystemPrompt ?? []; + result.appendSystemPrompt.push(args[++i]); + } else if (arg === "--name" || arg === "-n") { + if (i + 1 < args.length) { + result.name = args[++i]; + } else { + result.diagnostics.push({ type: "error", message: "--name requires a value" }); + } + } else if (arg === "--no-session") { + result.noSession = true; + } else if (arg === "--session" && i + 1 < args.length) { + result.session = args[++i]; + } else if (arg === "--session-id" && i + 1 < args.length) { + result.sessionId = args[++i]; + } else if (arg === "--fork" && i + 1 < args.length) { + result.fork = args[++i]; + } else if (arg === "--session-dir" && i + 1 < args.length) { + result.sessionDir = args[++i]; + } else if (arg === "--models" && i + 1 < args.length) { + result.models = args[++i].split(",").map((s) => s.trim()); + } else if (arg === "--no-tools" || arg === "-nt") { + result.noTools = true; + } else if (arg === "--no-builtin-tools" || arg === "-nbt") { + result.noBuiltinTools = true; + } else if ((arg === "--tools" || arg === "-t") && i + 1 < args.length) { + result.tools = args[++i] + .split(",") + .map((s) => s.trim()) + .filter((name) => name.length > 0); + } else if ((arg === "--exclude-tools" || arg === "-xt") && i + 1 < args.length) { + result.excludeTools = args[++i] + .split(",") + .map((s) => s.trim()) + .filter((name) => name.length > 0); + } else if (arg === "--thinking" && i + 1 < args.length) { + const level = args[++i]; + if (isValidThinkingLevel(level)) { + result.thinking = level; + } else { + result.diagnostics.push({ + type: "warning", + message: `Invalid thinking level "${level}". Valid values: ${VALID_THINKING_LEVELS.join(", ")}`, + }); + } + } else if (arg === "--print" || arg === "-p") { + result.print = true; + const next = args[i + 1]; + if (next !== undefined && !next.startsWith("@") && (!next.startsWith("-") || next.startsWith("---"))) { + result.messages.push(next); + i++; + } + } else if (arg === "--skill" && i + 1 < args.length) { + result.skills = result.skills ?? []; + result.skills.push(args[++i]); + } else if (arg === "--prompt-template" && i + 1 < args.length) { + result.promptTemplates = result.promptTemplates ?? []; + result.promptTemplates.push(args[++i]); + } else if (arg === "--theme" && i + 1 < args.length) { + result.themes = result.themes ?? []; + result.themes.push(args[++i]); + } else if (arg === "--no-skills" || arg === "-ns") { + result.noSkills = true; + } else if (arg === "--no-prompt-templates" || arg === "-np") { + result.noPromptTemplates = true; + } else if (arg === "--no-themes") { + result.noThemes = true; + } else if (arg === "--no-context-files" || arg === "-nc") { + result.noContextFiles = true; + } else if (arg === "--list-models") { + // Check if next arg is a search pattern (not a flag or file arg) + if (i + 1 < args.length && !args[i + 1].startsWith("-") && !args[i + 1].startsWith("@")) { + result.listModels = args[++i]; + } else { + result.listModels = true; + } + } else if (arg === "--verbose") { + result.verbose = true; + } else if (arg === "--approve" || arg === "-a") { + result.projectTrustOverride = true; + } else if (arg === "--no-approve" || arg === "-na") { + result.projectTrustOverride = false; + } else if (arg === "--offline") { + result.offline = true; + } else if (arg.startsWith("@")) { + result.fileArgs.push(arg.slice(1)); // Remove @ prefix + } else if (arg.startsWith("--")) { + const eqIndex = arg.indexOf("="); + if (eqIndex !== -1) { + result.unknownFlags.set(arg.slice(2, eqIndex), arg.slice(eqIndex + 1)); + } else { + const flagName = arg.slice(2); + const next = args[i + 1]; + if (next !== undefined && !next.startsWith("-") && !next.startsWith("@")) { + result.unknownFlags.set(flagName, next); + i++; + } else { + result.unknownFlags.set(flagName, true); + } + } + } else if (arg.startsWith("-") && !arg.startsWith("--")) { + result.diagnostics.push({ type: "error", message: `Unknown option: ${arg}` }); + } else if (!arg.startsWith("-")) { + result.messages.push(arg); + } + } + + return result; +} + +export function printHelp(extensionFlags?: ExtensionFlag[]): void { + const extensionFlagsText = + extensionFlags && extensionFlags.length > 0 + ? `\n${chalk.bold("Extension CLI Flags:")}\n${extensionFlags + .map((flag) => { + const value = flag.type === "string" ? " " : ""; + const description = flag.description ?? `Registered by ${flag.extensionPath}`; + return ` --${flag.name}${value}`.padEnd(30) + description; + }) + .join("\n")}\n` + : ""; + console.log(`${chalk.bold(`${APP_NAME} code`)} - AI coding assistant for the local Cactus model + +${chalk.bold("Usage:")} + ${APP_NAME} code [options] [@files...] [messages...] + +${chalk.bold("Options:")} + --model Model pattern or ID (defaults to the model the Cactus server serves) + --system-prompt System prompt (default: coding assistant prompt) + --append-system-prompt Append text or file contents to the system prompt (can be used multiple times) + --mode Output mode: text (default) or json + --print, -p Non-interactive mode: process prompt and exit + --continue, -c Continue previous session + --resume, -r Select a session to resume + --session Use specific session file or partial UUID + --session-id Use exact project session ID, creating it if missing + --fork Fork specific session file or partial UUID into a new session + --session-dir Directory for session storage and lookup + --no-session Don't save session (ephemeral) + --name, -n Set session display name + --models Comma-separated model patterns for Ctrl+P cycling + Supports glob patterns and fuzzy matching + --no-tools, -nt Disable all tools by default (built-in and extension) + --no-builtin-tools, -nbt Disable built-in tools by default but keep extension/custom tools enabled + --tools, -t Comma-separated allowlist of tool names to enable + Applies to built-in, extension, and custom tools + --exclude-tools, -xt Comma-separated denylist of tool names to disable + Applies to built-in, extension, and custom tools + --thinking Set thinking level: off, minimal, low, medium, high, xhigh + --skill Load a skill file or directory (can be used multiple times) + --no-skills, -ns Disable skills discovery and loading + --prompt-template Load a prompt template file or directory (can be used multiple times) + --no-prompt-templates, -np Disable prompt template discovery and loading + --theme Load a theme file or directory (can be used multiple times) + --no-themes Disable theme discovery and loading + --no-context-files, -nc Disable AGENTS.md and CLAUDE.md discovery and loading + --list-models [search] List available models (with optional fuzzy search) + --verbose Force verbose startup (overrides quietStartup setting) + --approve, -a Trust project-local files for this run + --no-approve, -na Ignore project-local files for this run + --offline Disable startup network operations + --help, -h Show this help + --version, -v Show version number + +${chalk.bold("Examples:")} + # Interactive mode + ${APP_NAME} code + + # Interactive mode with an initial prompt + ${APP_NAME} code "List all .ts files in src/" + + # Include files in the initial message + ${APP_NAME} code @prompt.md @image.png "What color is the sky?" + + # Non-interactive mode (process and exit) + ${APP_NAME} code -p "List all .ts files in src/" + + # Continue the previous session + ${APP_NAME} code --continue "What did we discuss?" + + # Start with a specific thinking level + ${APP_NAME} code --thinking high "Solve this complex problem" + + # Read-only mode (no file modifications) + ${APP_NAME} code --tools read,grep,find,ls -p "Review the code in src/" + +${chalk.bold("Environment Variables:")} + CACTUS_BASE_URL - Cactus server URL (default: http://127.0.0.1:8080/v1) + ${ENV_AGENT_DIR.padEnd(32)} - Config directory (default: ~/${CONFIG_DIR_NAME}/agent) + ${ENV_SESSION_DIR.padEnd(32)} - Session storage directory (overridden by --session-dir) + +${chalk.bold("Built-in Tools:")} + read - Read file contents + bash - Execute bash commands + edit - Edit files with find/replace + write - Write files (creates/overwrites) + grep - Search file contents + find - Find files by glob pattern + ls - List directory contents +`); +} diff --git a/cactus-code/packages/coding-agent/src/cli/file-processor.ts b/cactus-code/packages/coding-agent/src/cli/file-processor.ts new file mode 100644 index 000000000..e1da052a8 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/file-processor.ts @@ -0,0 +1,99 @@ +/** + * Process @file CLI arguments into text content and image attachments + */ + +import { access, readFile, stat } from "node:fs/promises"; +import type { ImageContent } from "@earendil-works/pi-ai"; +import chalk from "chalk"; +import { resolve } from "path"; +import { resolveReadPath } from "../core/tools/path-utils.ts"; +import { formatDimensionNote, resizeImage } from "../utils/image-resize.ts"; +import { detectSupportedImageMimeTypeFromFile } from "../utils/mime.ts"; + +export interface ProcessedFiles { + text: string; + images: ImageContent[]; +} + +export interface ProcessFileOptions { + /** Whether to auto-resize images to 2000x2000 max. Default: true */ + autoResizeImages?: boolean; +} + +/** Process @file arguments into text content and image attachments */ +export async function processFileArguments(fileArgs: string[], options?: ProcessFileOptions): Promise { + const autoResizeImages = options?.autoResizeImages ?? true; + let text = ""; + const images: ImageContent[] = []; + + for (const fileArg of fileArgs) { + // Expand and resolve path (handles ~ expansion and macOS screenshot Unicode spaces) + const absolutePath = resolve(resolveReadPath(fileArg, process.cwd())); + + // Check if file exists + try { + await access(absolutePath); + } catch { + console.error(chalk.red(`Error: File not found: ${absolutePath}`)); + process.exit(1); + } + + // Check if file is empty + const stats = await stat(absolutePath); + if (stats.size === 0) { + // Skip empty files + continue; + } + + const mimeType = await detectSupportedImageMimeTypeFromFile(absolutePath); + + if (mimeType) { + // Handle image file + const content = await readFile(absolutePath); + + let attachment: ImageContent; + let dimensionNote: string | undefined; + + if (autoResizeImages) { + const resized = await resizeImage(content, mimeType); + if (!resized) { + text += `[Image omitted: could not be resized below the inline image size limit.]\n`; + continue; + } + dimensionNote = formatDimensionNote(resized); + attachment = { + type: "image", + mimeType: resized.mimeType, + data: resized.data, + }; + } else { + attachment = { + type: "image", + mimeType, + data: content.toString("base64"), + }; + } + + images.push(attachment); + + // Add text reference to image with optional dimension note + if (dimensionNote) { + text += `${dimensionNote}\n`; + } else { + text += `\n`; + } + } else { + // Handle text file + try { + const content = await readFile(absolutePath, "utf-8"); + text += `\n${content}\n\n`; + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + console.error(chalk.red(`Error: Could not read file ${absolutePath}: ${message}`)); + process.exit(1); + } + } + } + + return { text, images }; +} diff --git a/cactus-code/packages/coding-agent/src/cli/initial-message.ts b/cactus-code/packages/coding-agent/src/cli/initial-message.ts new file mode 100644 index 000000000..c73f18a2f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/initial-message.ts @@ -0,0 +1,43 @@ +import type { ImageContent } from "@earendil-works/pi-ai"; +import type { Args } from "./args.ts"; + +export interface InitialMessageInput { + parsed: Args; + fileText?: string; + fileImages?: ImageContent[]; + stdinContent?: string; +} + +export interface InitialMessageResult { + initialMessage?: string; + initialImages?: ImageContent[]; +} + +/** + * Combine stdin content, @file text, and the first CLI message into a single + * initial prompt for non-interactive mode. + */ +export function buildInitialMessage({ + parsed, + fileText, + fileImages, + stdinContent, +}: InitialMessageInput): InitialMessageResult { + const parts: string[] = []; + if (stdinContent !== undefined) { + parts.push(stdinContent); + } + if (fileText) { + parts.push(fileText); + } + + if (parsed.messages.length > 0) { + parts.push(parsed.messages[0]); + parsed.messages.shift(); + } + + return { + initialMessage: parts.length > 0 ? parts.join("") : undefined, + initialImages: fileImages && fileImages.length > 0 ? fileImages : undefined, + }; +} diff --git a/cactus-code/packages/coding-agent/src/cli/list-models.ts b/cactus-code/packages/coding-agent/src/cli/list-models.ts new file mode 100644 index 000000000..04840daab --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/list-models.ts @@ -0,0 +1,112 @@ +/** + * List available models with optional fuzzy search + */ + +import type { Api, Model } from "@earendil-works/pi-ai"; +import { fuzzyFilter } from "@earendil-works/pi-tui"; +import chalk from "chalk"; +import { formatNoModelsAvailableMessage } from "../core/auth-guidance.ts"; +import type { ModelRegistry } from "../core/model-registry.ts"; + +/** + * Format a number as human-readable (e.g., 200000 -> "200K", 1000000 -> "1M") + */ +function formatTokenCount(count: number): string { + if (!Number.isFinite(count)) return "?"; + if (count >= 1_000_000) { + const millions = count / 1_000_000; + return millions % 1 === 0 ? `${millions}M` : `${millions.toFixed(1)}M`; + } + if (count >= 1_000) { + const thousands = count / 1_000; + return thousands % 1 === 0 ? `${thousands}K` : `${thousands.toFixed(1)}K`; + } + return count.toString(); +} + +/** + * List available models, optionally filtered by search pattern + */ +export async function listModels(modelRegistry: ModelRegistry, searchPattern?: string): Promise { + const loadError = modelRegistry.getError(); + if (loadError) { + console.error(chalk.yellow(`Warning: errors loading models.json:\n${loadError}`)); + } + + const models = modelRegistry.getAvailable(); + + if (models.length === 0) { + console.log(formatNoModelsAvailableMessage()); + return; + } + + // Apply fuzzy filter if search pattern provided + let filteredModels: Model[] = models; + if (searchPattern) { + filteredModels = fuzzyFilter(models, searchPattern, (m) => `${m.provider} ${m.id}`); + } + + if (filteredModels.length === 0) { + console.log(`No models matching "${searchPattern}"`); + return; + } + + // Sort by provider, then by model id + filteredModels.sort((a, b) => { + const providerCmp = a.provider.localeCompare(b.provider); + if (providerCmp !== 0) return providerCmp; + return a.id.localeCompare(b.id); + }); + + // Calculate column widths + const rows = filteredModels.map((m) => ({ + provider: m.provider, + model: m.id, + context: formatTokenCount(m.contextWindow), + maxOut: formatTokenCount(m.maxTokens), + thinking: m.reasoning ? "yes" : "no", + images: m.input.includes("image") ? "yes" : "no", + })); + + const headers = { + provider: "provider", + model: "model", + context: "context", + maxOut: "max-out", + thinking: "thinking", + images: "images", + }; + + const widths = { + provider: Math.max(headers.provider.length, ...rows.map((r) => r.provider.length)), + model: Math.max(headers.model.length, ...rows.map((r) => r.model.length)), + context: Math.max(headers.context.length, ...rows.map((r) => r.context.length)), + maxOut: Math.max(headers.maxOut.length, ...rows.map((r) => r.maxOut.length)), + thinking: Math.max(headers.thinking.length, ...rows.map((r) => r.thinking.length)), + images: Math.max(headers.images.length, ...rows.map((r) => r.images.length)), + }; + + // Print header + const headerLine = [ + headers.provider.padEnd(widths.provider), + headers.model.padEnd(widths.model), + headers.context.padEnd(widths.context), + headers.maxOut.padEnd(widths.maxOut), + headers.thinking.padEnd(widths.thinking), + headers.images.padEnd(widths.images), + ].join(" "); + console.log(headerLine); + + // Print rows + for (const row of rows) { + const line = [ + row.provider.padEnd(widths.provider), + row.model.padEnd(widths.model), + row.context.padEnd(widths.context), + row.maxOut.padEnd(widths.maxOut), + row.thinking.padEnd(widths.thinking), + row.images.padEnd(widths.images), + ].join(" "); + console.log(line); + } +} diff --git a/cactus-code/packages/coding-agent/src/cli/project-trust.ts b/cactus-code/packages/coding-agent/src/cli/project-trust.ts new file mode 100644 index 000000000..b27871a95 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/project-trust.ts @@ -0,0 +1,62 @@ +import chalk from "chalk"; +import type { ProjectTrustContext } from "../core/extensions/types.ts"; +import type { AppMode } from "../core/project-trust.ts"; +import type { SettingsManager } from "../core/settings-manager.ts"; +import { showStartupInput, showStartupSelector } from "./startup-ui.ts"; + +export function createProjectTrustContext(options: { + cwd: string; + mode: AppMode; + settingsManager: SettingsManager; + hasUI: boolean; +}): ProjectTrustContext { + return { + cwd: options.cwd, + mode: options.mode === "interactive" ? "tui" : options.mode, + hasUI: options.hasUI, + ui: { + select: async (title, selectOptions) => { + if (!options.hasUI) { + return undefined; + } + if (options.mode !== "interactive") { + return undefined; + } + return showStartupSelector( + options.settingsManager, + title, + selectOptions.map((option) => ({ label: option, value: option })), + ); + }, + confirm: async (title, message) => { + if (!options.hasUI) { + return false; + } + if (options.mode !== "interactive") { + return false; + } + return ( + (await showStartupSelector(options.settingsManager, `${title}\n${message}`, [ + { label: "Yes", value: true }, + { label: "No", value: false }, + ])) ?? false + ); + }, + input: async (title, placeholder) => { + if (!options.hasUI) { + return undefined; + } + if (options.mode !== "interactive") { + return undefined; + } + return showStartupInput(options.settingsManager, title, placeholder); + }, + notify: (message, type = "info") => { + if (options.mode !== "interactive") { + const color = type === "error" ? chalk.red : type === "warning" ? chalk.yellow : chalk.cyan; + console.error(color(message)); + } + }, + }, + }; +} diff --git a/cactus-code/packages/coding-agent/src/cli/session-picker.ts b/cactus-code/packages/coding-agent/src/cli/session-picker.ts new file mode 100644 index 000000000..793f1d0a3 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/session-picker.ts @@ -0,0 +1,55 @@ +/** + * TUI session selector for --resume flag + */ + +import { setKeybindings } from "@earendil-works/pi-tui"; +import { KeybindingsManager } from "../core/keybindings.ts"; +import type { SessionInfo, SessionListProgress } from "../core/session-manager.ts"; +import type { SettingsManager } from "../core/settings-manager.ts"; +import { SessionSelectorComponent } from "../modes/interactive/components/session-selector.ts"; +import { createStartupTui, startStartupTui } from "./startup-ui.ts"; + +type SessionsLoader = (onProgress?: SessionListProgress) => Promise; + +/** Show TUI session selector and return selected session path or null if cancelled */ +export async function selectSession( + currentSessionsLoader: SessionsLoader, + allSessionsLoader: SessionsLoader, + settingsManager: SettingsManager, +): Promise { + const ui = await createStartupTui(settingsManager); + return new Promise((resolve) => { + const keybindings = KeybindingsManager.create(); + setKeybindings(keybindings); + let resolved = false; + + const selector = new SessionSelectorComponent( + currentSessionsLoader, + allSessionsLoader, + (path: string) => { + if (!resolved) { + resolved = true; + ui.stop(); + resolve(path); + } + }, + () => { + if (!resolved) { + resolved = true; + ui.stop(); + resolve(null); + } + }, + () => { + ui.stop(); + process.exit(0); + }, + () => ui.requestRender(), + { showRenameHint: false, keybindings }, + ); + + ui.addChild(selector); + ui.setFocus(selector.getSessionList()); + startStartupTui(ui, settingsManager); + }); +} diff --git a/cactus-code/packages/coding-agent/src/cli/startup-ui.ts b/cactus-code/packages/coding-agent/src/cli/startup-ui.ts new file mode 100644 index 000000000..73c0271c8 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/cli/startup-ui.ts @@ -0,0 +1,239 @@ +import { ProcessTerminal, setKeybindings, TUI } from "@earendil-works/pi-tui"; +import { existsSync } from "fs"; +import { APP_NAME, CONFIG_DIR_NAME, ENV_AGENT_DIR, getAgentDir, getSettingsPath, PACKAGE_NAME } from "../config.ts"; +import { areExperimentalFeaturesEnabled } from "../core/experimental.ts"; +import { KeybindingsManager } from "../core/keybindings.ts"; +import { DefaultPackageManager, type ResolvedResource } from "../core/package-manager.ts"; +import { SettingsManager } from "../core/settings-manager.ts"; +import { ExtensionInputComponent } from "../modes/interactive/components/extension-input.ts"; +import { ExtensionSelectorComponent } from "../modes/interactive/components/extension-selector.ts"; +import { + FirstTimeSetupComponent, + type FirstTimeSetupResult, +} from "../modes/interactive/components/first-time-setup.ts"; +import { + detectTerminalBackgroundFromEnv, + detectTerminalThemeForAuto, + initTheme, + loadThemeFromPath, + parseAutoThemeSetting, + resolveThemeSetting, + setRegisteredThemes, + setTheme, + type Theme, +} from "../modes/interactive/theme/theme.ts"; + +const OFFICIAL_PACKAGE_NAME = "@earendil-works/pi-coding-agent"; +const OFFICIAL_APP_NAME = "pi"; +const OFFICIAL_CONFIG_DIR_NAME = ".pi"; + +interface DistributionMetadata { + packageName: string; + appName: string; + configDirName: string; +} + +function isOfficialDistribution({ packageName, appName, configDirName }: DistributionMetadata): boolean { + return ( + packageName === OFFICIAL_PACKAGE_NAME && + appName === OFFICIAL_APP_NAME && + configDirName === OFFICIAL_CONFIG_DIR_NAME + ); +} + +function loadThemes(resources: ResolvedResource[]): Theme[] { + const themes: Theme[] = []; + const seen = new Set(); + for (const resource of resources) { + if (!resource.enabled) continue; + try { + const loadedTheme = loadThemeFromPath(resource.path); + if (loadedTheme.name) { + if (seen.has(loadedTheme.name)) continue; + seen.add(loadedTheme.name); + } + themes.push(loadedTheme); + } catch { + // Startup prompts should not fail because a theme is broken. The normal + // resource loader reports theme diagnostics later in startup. + } + } + return themes; +} + +async function loadStartupThemes(settingsManager: SettingsManager): Promise { + const globalSettingsManager = SettingsManager.inMemory(settingsManager.getGlobalSettings(), { + projectTrusted: false, + }); + const packageManager = new DefaultPackageManager({ + cwd: process.cwd(), + agentDir: getAgentDir(), + settingsManager: globalSettingsManager, + }); + const resolvedPaths = await packageManager.resolve(async () => "skip"); + return loadThemes(resolvedPaths.themes); +} + +export async function createStartupTui(settingsManager: SettingsManager): Promise { + setRegisteredThemes(await loadStartupThemes(settingsManager)); + const terminalTheme = detectTerminalBackgroundFromEnv().theme; + initTheme(resolveThemeSetting(settingsManager.getThemeSetting(), terminalTheme) ?? terminalTheme); + setKeybindings(KeybindingsManager.create()); + const ui = new TUI(new ProcessTerminal(), settingsManager.getShowHardwareCursor()); + ui.setClearOnShrink(settingsManager.getClearOnShrink()); + return ui; +} + +export function startStartupTui(ui: TUI, settingsManager: SettingsManager): void { + ui.start(); + void applyDetectedStartupTheme(ui, settingsManager); +} + +async function applyDetectedStartupTheme(ui: TUI, settingsManager: SettingsManager): Promise { + const themeSetting = settingsManager.getThemeSetting(); + if (themeSetting && !parseAutoThemeSetting(themeSetting)) return; + + const terminalTheme = await detectTerminalThemeForAuto({ ui, timeoutMs: 100 }); + setTheme(resolveThemeSetting(themeSetting, terminalTheme) ?? terminalTheme); + ui.invalidate(); + ui.requestRender(); +} + +async function clearStartupTui(ui: TUI): Promise { + ui.clear(); + ui.requestRender(); + await new Promise((resolve) => setTimeout(resolve, 25)); +} + +/** + * First-time setup runs when all of these hold: + * - this is the official Pi distribution (not a fork/rebrand) + * - experimental features are enabled (PI_EXPERIMENTAL=1) + * - the default agent directory is used (no custom agent dir override) + * - setup was not completed before (settings.json does not exist) + */ +export function shouldRunFirstTimeSetup(settingsPath: string = getSettingsPath()): boolean { + if ( + !isOfficialDistribution({ + packageName: PACKAGE_NAME, + appName: APP_NAME, + configDirName: CONFIG_DIR_NAME, + }) + ) { + return false; + } + if (!areExperimentalFeaturesEnabled()) { + return false; + } + if (process.env[ENV_AGENT_DIR]) { + return false; + } + return !existsSync(settingsPath); +} + +export async function showStartupSelector( + settingsManager: SettingsManager, + title: string, + options: Array<{ label: string; value: T }>, +): Promise { + const ui = await createStartupTui(settingsManager); + return new Promise((resolve) => { + let settled = false; + const finish = async (result: T | undefined) => { + if (settled) { + return; + } + settled = true; + await clearStartupTui(ui); + ui.stop(); + resolve(result); + }; + + const selector = new ExtensionSelectorComponent( + title, + options.map((option) => option.label), + (option) => void finish(options.find((entry) => entry.label === option)?.value), + () => void finish(undefined), + { tui: ui }, + ); + ui.addChild(selector); + ui.setFocus(selector); + startStartupTui(ui, settingsManager); + }); +} + +/** Show the first-time setup dialog and persist the result */ +export async function showFirstTimeSetup(settingsManager: SettingsManager): Promise { + const ui = await createStartupTui(settingsManager); + return new Promise((resolve) => { + let settled = false; + const finish = async (result: FirstTimeSetupResult | undefined) => { + if (settled) { + return; + } + settled = true; + if (result) { + settingsManager.setTheme(result.theme); + settingsManager.setEnableAnalytics(result.shareAnalytics); + await settingsManager.flush(); + } + await clearStartupTui(ui); + ui.stop(); + resolve(); + }; + + const showSetup = async () => { + ui.start(); + const detectedTheme = await detectTerminalThemeForAuto({ ui, timeoutMs: 100 }); + setTheme(detectedTheme); + const component = new FirstTimeSetupComponent({ + detectedTheme, + onThemePreview: (themeName) => { + setTheme(themeName); + ui.requestRender(); + }, + onSubmit: (result) => void finish(result), + onCancel: () => void finish(undefined), + }); + ui.addChild(component); + ui.setFocus(component); + ui.requestRender(); + }; + + void showSetup(); + }); +} + +export async function showStartupInput( + settingsManager: SettingsManager, + title: string, + placeholder?: string, +): Promise { + const ui = await createStartupTui(settingsManager); + return new Promise((resolve) => { + let settled = false; + const finish = async (result: string | undefined) => { + if (settled) { + return; + } + settled = true; + input.dispose(); + await clearStartupTui(ui); + ui.stop(); + resolve(result); + }; + + const input = new ExtensionInputComponent( + title, + placeholder, + (value) => void finish(value), + () => void finish(undefined), + { + tui: ui, + }, + ); + ui.addChild(input); + ui.setFocus(input); + startStartupTui(ui, settingsManager); + }); +} diff --git a/cactus-code/packages/coding-agent/src/config.ts b/cactus-code/packages/coding-agent/src/config.ts new file mode 100644 index 000000000..94a23da5c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/config.ts @@ -0,0 +1,554 @@ +import { accessSync, constants, existsSync, readFileSync, realpathSync } from "fs"; +import { homedir } from "os"; +import { basename, dirname, join, resolve, sep, win32 } from "path"; +import { fileURLToPath } from "url"; +import { spawnProcessSync } from "./utils/child-process.ts"; +import { normalizePath } from "./utils/paths.ts"; + +// ============================================================================= +// Package Detection +// ============================================================================= + +const __filename = fileURLToPath(import.meta.url); +const __dirname = dirname(__filename); + +/** + * Detect if we're running as a Bun compiled binary. + * Bun binaries have import.meta.url containing "$bunfs", "~BUN", or "%7EBUN" (Bun's virtual filesystem path) + */ +export const isBunBinary = + import.meta.url.includes("$bunfs") || import.meta.url.includes("~BUN") || import.meta.url.includes("%7EBUN"); + +/** Detect if Bun is the runtime (compiled binary or bun run) */ +export const isBunRuntime = !!process.versions.bun; + +// ============================================================================= +// Install Method Detection +// ============================================================================= + +export type InstallMethod = "bun-binary" | "npm" | "pnpm" | "yarn" | "bun" | "unknown"; + +interface SelfUpdateCommandStep { + command: string; + args: string[]; + display: string; +} + +export interface SelfUpdateCommand extends SelfUpdateCommandStep { + steps?: SelfUpdateCommandStep[]; +} + +export type SelfUpdatePackageTarget = string | { packageName: string; installSpec?: string }; + +function normalizeSelfUpdatePackageTarget(target: SelfUpdatePackageTarget): { + packageName: string; + installSpec: string; +} { + if (typeof target === "string") { + return { packageName: target, installSpec: target }; + } + return { packageName: target.packageName, installSpec: target.installSpec ?? target.packageName }; +} + +function makeSelfUpdateCommand( + installStep: SelfUpdateCommandStep, + uninstallStep?: SelfUpdateCommandStep, +): SelfUpdateCommand { + if (!uninstallStep) return installStep; + return { + ...installStep, + display: `${uninstallStep.display} && ${installStep.display}`, + steps: [uninstallStep, installStep], + }; +} + +function makeSelfUpdateCommandStep(command: string, args: string[]): SelfUpdateCommandStep { + return { + command, + args, + display: [command, ...args].map((arg) => (/\s/.test(arg) ? `"${arg}"` : arg)).join(" "), + }; +} + +export function detectInstallMethod(): InstallMethod { + if (isBunBinary) { + return "bun-binary"; + } + + const resolvedPath = `${__dirname}\0${process.execPath || ""}`.toLowerCase().replace(/\\/g, "/"); + + if (resolvedPath.includes("/pnpm/") || resolvedPath.includes("/.pnpm/")) { + return "pnpm"; + } + if (resolvedPath.includes("/yarn/") || resolvedPath.includes("/.yarn/")) { + return "yarn"; + } + if (isBunRuntime || resolvedPath.includes("/install/global/node_modules/")) { + return "bun"; + } + if (resolvedPath.includes("/npm/") || resolvedPath.includes("/node_modules/")) { + return "npm"; + } + + return "unknown"; +} + +function getInferredNpmInstall(): { root: string; prefix: string } | undefined { + const packageDir = getPackageDir(); + const path = process.platform === "win32" || packageDir.includes("\\") ? win32 : { basename, dirname }; + const parent = path.dirname(packageDir); + let root: string | undefined; + if (path.basename(parent).startsWith("@") && path.basename(path.dirname(parent)) === "node_modules") { + root = path.dirname(parent); + } else if (path.basename(parent) === "node_modules") { + root = parent; + } + if (!root) return undefined; + const rootParent = path.dirname(root); + if (path.basename(rootParent) === "lib") return { root, prefix: path.dirname(rootParent) }; + // Windows global npm prefixes use `\\node_modules`, which is + // indistinguishable from local project installs by path shape alone. Do not + // infer unsupported Windows custom prefixes without `npm root -g` evidence. + return undefined; +} + +function getSelfUpdateCommandForMethod( + method: InstallMethod, + installedPackageName: string, + updatePackageTarget: SelfUpdatePackageTarget = installedPackageName, + npmCommand?: string[], +): SelfUpdateCommand | undefined { + const target = normalizeSelfUpdatePackageTarget(updatePackageTarget); + switch (method) { + case "bun-binary": + return undefined; + case "pnpm": { + const match = readCommandOutput("pnpm", ["root", "-g"]) + ? undefined + : /^(.*[\\/]global[\\/][^\\/]+)[\\/]\.pnpm[\\/]/.exec(getPackageDir()); + const binDirArgs = match + ? [`--config.global-bin-dir=${process.env.PNPM_HOME || dirname(dirname(match[1]))}`] + : []; + return makeSelfUpdateCommand( + makeSelfUpdateCommandStep("pnpm", [ + "install", + "-g", + "--ignore-scripts", + "--config.minimumReleaseAge=0", + ...binDirArgs, + target.installSpec, + ]), + target.packageName === installedPackageName + ? undefined + : makeSelfUpdateCommandStep("pnpm", ["remove", "-g", ...binDirArgs, installedPackageName]), + ); + } + case "yarn": + return makeSelfUpdateCommand( + makeSelfUpdateCommandStep("yarn", ["global", "add", "--ignore-scripts", target.installSpec]), + target.packageName === installedPackageName + ? undefined + : makeSelfUpdateCommandStep("yarn", ["global", "remove", installedPackageName]), + ); + case "bun": + return makeSelfUpdateCommand( + makeSelfUpdateCommandStep("bun", [ + "install", + "-g", + "--ignore-scripts", + "--minimum-release-age=0", + target.installSpec, + ]), + target.packageName === installedPackageName + ? undefined + : makeSelfUpdateCommandStep("bun", ["uninstall", "-g", installedPackageName]), + ); + case "npm": { + const [command = "npm", ...npmArgs] = npmCommand ?? []; + const inferred = npmCommand?.length ? undefined : getInferredNpmInstall(); + const prefixArgs = [...npmArgs, ...(inferred ? ["--prefix", inferred.prefix] : [])]; + const installStep = makeSelfUpdateCommandStep(command, [ + ...prefixArgs, + "install", + "-g", + "--ignore-scripts", + "--min-release-age=0", + target.installSpec, + ]); + const uninstallStep = + target.packageName === installedPackageName + ? undefined + : makeSelfUpdateCommandStep(command, [...prefixArgs, "uninstall", "-g", installedPackageName]); + return makeSelfUpdateCommand(installStep, uninstallStep); + } + case "unknown": + return undefined; + } +} + +function readCommandOutput( + command: string, + args: string[], + options: { requireSuccess?: boolean } = {}, +): string | undefined { + const result = spawnProcessSync(command, args, { + encoding: "utf-8", + stdio: ["ignore", "pipe", "pipe"], + }); + if (result.status === 0) return result.stdout.trim() || undefined; + if (options.requireSuccess) { + const reason = result.error?.message || result.stderr.trim() || `exit code ${result.status ?? "unknown"}`; + throw new Error(`Failed to run ${[command, ...args].join(" ")}: ${reason}`); + } + return undefined; +} + +function getGlobalPackageRoots(method: InstallMethod, _packageName: string, npmCommand?: string[]): string[] { + switch (method) { + case "npm": { + const configured = !!npmCommand?.length; + const [command = "npm", ...npmArgs] = npmCommand ?? []; + if (configured && command === "bun") { + const bunBin = readCommandOutput(command, [...npmArgs, "pm", "bin", "-g"], { + requireSuccess: true, + }); + const roots = [join(homedir(), ".bun", "install", "global", "node_modules")]; + if (bunBin) { + roots.push(join(dirname(bunBin), "install", "global", "node_modules")); + } + return roots; + } + const root = readCommandOutput(command, [...npmArgs, "root", "-g"], { + requireSuccess: configured, + }); + const inferred = configured ? undefined : getInferredNpmInstall(); + return [root, inferred?.root].filter((x): x is string => !!x); + } + case "pnpm": { + const root = readCommandOutput("pnpm", ["root", "-g"]); + if (root) return [root, dirname(root)]; + const match = /^(.*[\\/]global[\\/][^\\/]+)[\\/]\.pnpm[\\/]/.exec(getPackageDir()); + return match ? [match[1]] : []; + } + case "yarn": { + const dir = readCommandOutput("yarn", ["global", "dir"]); + return dir ? [dir, join(dir, "node_modules")] : []; + } + case "bun": { + const bunBin = readCommandOutput("bun", ["pm", "bin", "-g"]); + const roots = [join(homedir(), ".bun", "install", "global", "node_modules")]; + if (bunBin) { + roots.push(join(dirname(bunBin), "install", "global", "node_modules")); + } + return roots; + } + case "bun-binary": + case "unknown": + return []; + } +} + +function normalizeExistingPathForComparison(path: string, resolveSymlinks: boolean): string | undefined { + const resolvedPath = resolve(path); + if (!existsSync(resolvedPath)) { + return undefined; + } + let normalizedPath = resolvedPath; + if (resolveSymlinks) { + try { + normalizedPath = realpathSync(resolvedPath); + } catch { + return undefined; + } + } + if (process.platform === "win32") { + normalizedPath = normalizedPath.toLowerCase(); + } + return normalizedPath; +} + +function getPathComparisonCandidates(path: string): string[] { + return Array.from( + new Set( + [normalizeExistingPathForComparison(path, false), normalizeExistingPathForComparison(path, true)].filter( + (candidate): candidate is string => !!candidate, + ), + ), + ); +} + +function getEntrypointPackageDir(): string | undefined { + const entrypoint = process.argv[1]; + if (!entrypoint) return undefined; + let dir = dirname(entrypoint); + while (dir !== dirname(dir)) { + if (existsSync(join(dir, "package.json"))) { + return dir; + } + dir = dirname(dir); + } + return undefined; +} + +function isSelfUpdatePathWritable(): boolean { + const packageDir = getPackageDir(); + try { + accessSync(packageDir, constants.W_OK); + accessSync(dirname(packageDir), constants.W_OK); + return true; + } catch { + return false; + } +} + +function isManagedByGlobalPackageManager(method: InstallMethod, packageName: string, npmCommand?: string[]): boolean { + const packageDirs = [getPackageDir(), getEntrypointPackageDir()].filter((dir): dir is string => !!dir); + const packageDirCandidates = packageDirs.flatMap((dir) => getPathComparisonCandidates(dir)); + return getGlobalPackageRoots(method, packageName, npmCommand).some((root) => { + return getPathComparisonCandidates(root).some((normalizedRoot) => { + const rootPrefix = normalizedRoot.endsWith(sep) ? normalizedRoot : `${normalizedRoot}${sep}`; + return packageDirCandidates.some((packageDir) => packageDir.startsWith(rootPrefix)); + }); + }); +} + +export function getSelfUpdateCommand( + packageName: string, + npmCommand?: string[], + updatePackageTarget: SelfUpdatePackageTarget = packageName, +): SelfUpdateCommand | undefined { + const method = detectInstallMethod(); + const command = getSelfUpdateCommandForMethod(method, packageName, updatePackageTarget, npmCommand); + if (!command || !isManagedByGlobalPackageManager(method, packageName, npmCommand) || !isSelfUpdatePathWritable()) { + return undefined; + } + return command; +} + +export function getSelfUpdateUnavailableInstruction( + packageName: string, + npmCommand?: string[], + updatePackageTarget: SelfUpdatePackageTarget = packageName, +): string { + const method = detectInstallMethod(); + const target = normalizeSelfUpdatePackageTarget(updatePackageTarget); + if (method === "bun-binary") { + return `Download from: https://github.com/earendil-works/pi-mono/releases/latest`; + } + const command = getSelfUpdateCommandForMethod(method, packageName, target, npmCommand); + if (command) { + if (isManagedByGlobalPackageManager(method, packageName, npmCommand) && !isSelfUpdatePathWritable()) { + return `This installation is managed by a global ${method} install, but the install path is not writable. Update it yourself with: ${command.display}`; + } + return `This installation is not managed by a global ${method} install. Update it with the package manager, wrapper, or source checkout that provides it.`; + } + return `Update ${target.installSpec} using the package manager, wrapper, or source checkout that provides this installation.`; +} + +export function getUpdateInstruction(packageName: string): string { + const method = detectInstallMethod(); + const command = getSelfUpdateCommandForMethod(method, packageName); + if (command) { + return `Run: ${command.display}`; + } + return getSelfUpdateUnavailableInstruction(packageName); +} + +// ============================================================================= +// Package Asset Paths (shipped with executable) +// ============================================================================= + +/** + * Get the base directory for resolving package assets (themes, package.json, README.md, CHANGELOG.md). + * - For Bun binary: returns the directory containing the executable + * - For Node.js (dist/): returns __dirname (the dist/ directory) + * - For tsx (src/): returns parent directory (the package root) + */ +export function getPackageDir(): string { + // Allow override via environment variable (useful for Nix/Guix where store paths tokenize poorly) + const envDir = process.env.PI_PACKAGE_DIR; + if (envDir) { + return normalizePath(envDir); + } + + if (isBunBinary) { + // Bun binary: process.execPath points to the compiled executable + return dirname(process.execPath); + } + // Node.js: walk up from __dirname until we find package.json + let dir = __dirname; + while (dir !== dirname(dir)) { + if (existsSync(join(dir, "package.json"))) { + return dir; + } + dir = dirname(dir); + } + // Fallback (shouldn't happen) + return __dirname; +} + +/** + * Get path to built-in themes directory (shipped with package) + * - For Bun binary: theme/ next to executable + * - For Node.js (dist/): dist/modes/interactive/theme/ + * - For tsx (src/): src/modes/interactive/theme/ + */ +export function getThemesDir(): string { + if (isBunBinary) { + return join(getPackageDir(), "theme"); + } + // Theme is in modes/interactive/theme/ relative to src/ or dist/ + const packageDir = getPackageDir(); + const srcOrDist = existsSync(join(packageDir, "src")) ? "src" : "dist"; + return join(packageDir, srcOrDist, "modes", "interactive", "theme"); +} + + +/** Get path to package.json */ +export function getPackageJsonPath(): string { + return join(getPackageDir(), "package.json"); +} + +/** Get path to README.md */ +export function getReadmePath(): string { + return resolve(join(getPackageDir(), "README.md")); +} + +/** Get path to docs directory */ +export function getDocsPath(): string { + return resolve(join(getPackageDir(), "docs")); +} + +/** Get path to examples directory */ +export function getExamplesPath(): string { + return resolve(join(getPackageDir(), "examples")); +} + +/** + * Get path to built-in interactive assets directory. + * - For Bun binary: assets/ next to executable + * - For Node.js (dist/): dist/modes/interactive/assets/ + * - For tsx (src/): src/modes/interactive/assets/ + */ +export function getInteractiveAssetsDir(): string { + if (isBunBinary) { + return join(getPackageDir(), "assets"); + } + const packageDir = getPackageDir(); + const srcOrDist = existsSync(join(packageDir, "src")) ? "src" : "dist"; + return join(packageDir, srcOrDist, "modes", "interactive", "assets"); +} + +/** Get path to a bundled interactive asset */ +export function getBundledInteractiveAssetPath(name: string): string { + return join(getInteractiveAssetsDir(), name); +} + +// ============================================================================= +// App Config (from package.json piConfig) +// ============================================================================= + +interface PackageJson { + name?: string; + version?: string; + piConfig?: { + name?: string; + configDir?: string; + }; +} + +let pkg: PackageJson = {}; +try { + pkg = JSON.parse(readFileSync(getPackageJsonPath(), "utf-8")) as PackageJson; +} catch (e: unknown) { + const err = e as NodeJS.ErrnoException; + if (err.code !== "ENOENT") throw e; +} + +const piConfigName: string | undefined = pkg.piConfig?.name; +export const PACKAGE_NAME: string = pkg.name || "@earendil-works/pi-coding-agent"; +export const APP_NAME: string = piConfigName || "cactus"; +export const APP_TITLE: string = piConfigName || "cactus"; +export const CONFIG_DIR_NAME: string = pkg.piConfig?.configDir || ".cactus"; + +function readCactusVersion(): string | undefined { + const pkgDir = getPackageDir(); + for (const rel of [ + ["..", "..", "CACTUS_VERSION"], // bundled wheel: /CACTUS_VERSION + ["..", "..", "..", "CACTUS_VERSION"], // dev checkout: /CACTUS_VERSION + ]) { + try { + const value = readFileSync(join(pkgDir, ...rel), "utf-8").trim(); + if (value) return value.replace(/^v/i, ""); + } catch { + } + } + return undefined; +} + +export const VERSION: string = readCactusVersion() || pkg.version || "0.0.0"; + +export const ENV_AGENT_DIR = `${APP_NAME.toUpperCase()}_CODING_AGENT_DIR`; +export const ENV_SESSION_DIR = `${APP_NAME.toUpperCase()}_CODING_AGENT_SESSION_DIR`; + +export function expandTildePath(path: string): string { + return normalizePath(path); +} + +// ============================================================================= +// User Config Paths (~/.cactus/agent/*) +// ============================================================================= + +/** Get the agent config directory (e.g., ~/.cactus/agent/) */ +export function getAgentDir(): string { + const envDir = process.env[ENV_AGENT_DIR]; + if (envDir) { + return expandTildePath(envDir); + } + return join(homedir(), CONFIG_DIR_NAME, "agent"); +} + +/** Get path to user's custom themes directory */ +export function getCustomThemesDir(): string { + return join(getAgentDir(), "themes"); +} + +/** Get path to models.json */ +export function getModelsPath(): string { + return join(getAgentDir(), "models.json"); +} + +/** Get path to auth.json */ +export function getAuthPath(): string { + return join(getAgentDir(), "auth.json"); +} + +/** Get path to settings.json */ +export function getSettingsPath(): string { + return join(getAgentDir(), "settings.json"); +} + +/** Get path to tools directory */ +export function getToolsDir(): string { + return join(getAgentDir(), "tools"); +} + +/** Get path to managed binaries directory (fd, rg) */ +export function getBinDir(): string { + return join(getAgentDir(), "bin"); +} + +/** Get path to prompt templates directory */ +export function getPromptsDir(): string { + return join(getAgentDir(), "prompts"); +} + +/** Get path to sessions directory */ +export function getSessionsDir(): string { + return join(getAgentDir(), "sessions"); +} + +/** Get path to debug log file */ +export function getDebugLogPath(): string { + return join(getAgentDir(), `${APP_NAME}-debug.log`); +} diff --git a/cactus-code/packages/coding-agent/src/core/agent-session-runtime.ts b/cactus-code/packages/coding-agent/src/core/agent-session-runtime.ts new file mode 100644 index 000000000..7f29275aa --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/agent-session-runtime.ts @@ -0,0 +1,433 @@ +import { copyFileSync, existsSync, mkdirSync } from "node:fs"; +import { basename, join, resolve } from "node:path"; +import { resolvePath } from "../utils/paths.ts"; +import type { AgentSession } from "./agent-session.ts"; +import type { AgentSessionRuntimeDiagnostic, AgentSessionServices } from "./agent-session-services.ts"; +import type { + ProjectTrustContext, + ReplacedSessionContext, + SessionShutdownEvent, + SessionStartEvent, +} from "./extensions/index.ts"; +import { emitSessionShutdownEvent } from "./extensions/runner.ts"; +import type { CreateAgentSessionResult } from "./sdk.ts"; +import { assertSessionCwdExists } from "./session-cwd.ts"; +import { SessionManager } from "./session-manager.ts"; + +/** + * Result returned by runtime creation. + * + * The caller gets the created session, its cwd-bound services, and all + * diagnostics collected during setup. + */ +export interface CreateAgentSessionRuntimeResult extends CreateAgentSessionResult { + services: AgentSessionServices; + diagnostics: AgentSessionRuntimeDiagnostic[]; +} + +/** + * Creates a full runtime for a target cwd and session manager. + * + * The factory closes over process-global fixed inputs, recreates cwd-bound + * services for the effective cwd, resolves session options against those + * services, and finally creates the AgentSession. + */ +export type CreateAgentSessionRuntimeFactory = (options: { + cwd: string; + agentDir: string; + sessionManager: SessionManager; + sessionStartEvent?: SessionStartEvent; + projectTrustContext?: ProjectTrustContext; +}) => Promise; + +/** + * Thrown when /import references a JSONL file path that does not exist. + */ +export class SessionImportFileNotFoundError extends Error { + readonly filePath: string; + + constructor(filePath: string) { + super(`File not found: ${filePath}`); + this.name = "SessionImportFileNotFoundError"; + this.filePath = filePath; + } +} + +function extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { + if (typeof content === "string") { + return content; + } + + return content + .filter((part): part is { type: "text"; text: string } => part.type === "text" && typeof part.text === "string") + .map((part) => part.text) + .join(""); +} + +/** + * Owns the current AgentSession plus its cwd-bound services. + * + * Session replacement methods tear down the current runtime first, then create + * and apply the next runtime. If creation fails, the error is propagated to the + * caller. The caller is responsible for user-facing error handling. + */ +export class AgentSessionRuntime { + private rebindSession?: (session: AgentSession) => Promise; + private beforeSessionInvalidate?: () => void; + private _session: AgentSession; + private _services: AgentSessionServices; + private readonly createRuntime: CreateAgentSessionRuntimeFactory; + private _diagnostics: AgentSessionRuntimeDiagnostic[]; + private _modelFallbackMessage?: string; + + constructor( + _session: AgentSession, + _services: AgentSessionServices, + createRuntime: CreateAgentSessionRuntimeFactory, + _diagnostics: AgentSessionRuntimeDiagnostic[] = [], + _modelFallbackMessage?: string, + ) { + this._session = _session; + this._services = _services; + this.createRuntime = createRuntime; + this._diagnostics = _diagnostics; + this._modelFallbackMessage = _modelFallbackMessage; + } + + get services(): AgentSessionServices { + return this._services; + } + + get session(): AgentSession { + return this._session; + } + + get cwd(): string { + return this._services.cwd; + } + + get diagnostics(): readonly AgentSessionRuntimeDiagnostic[] { + return this._diagnostics; + } + + get modelFallbackMessage(): string | undefined { + return this._modelFallbackMessage; + } + + setRebindSession(rebindSession?: (session: AgentSession) => Promise): void { + this.rebindSession = rebindSession; + } + + /** + * Set a synchronous callback that runs after `session_shutdown` handlers finish + * but before the current session is invalidated. + * + * This is for host-owned UI teardown that must not yield to the event loop, + * such as detaching extension-provided TUI components before the old extension + * context becomes stale. + */ + setBeforeSessionInvalidate(beforeSessionInvalidate?: () => void): void { + this.beforeSessionInvalidate = beforeSessionInvalidate; + } + + private async emitBeforeSwitch( + reason: "new" | "resume", + targetSessionFile?: string, + ): Promise<{ cancelled: boolean }> { + const runner = this.session.extensionRunner; + if (!runner.hasHandlers("session_before_switch")) { + return { cancelled: false }; + } + + const result = await runner.emit({ + type: "session_before_switch", + reason, + targetSessionFile, + }); + return { cancelled: result?.cancel === true }; + } + + private async emitBeforeFork( + entryId: string, + options: { position: "before" | "at" }, + ): Promise<{ cancelled: boolean }> { + const runner = this.session.extensionRunner; + if (!runner.hasHandlers("session_before_fork")) { + return { cancelled: false }; + } + + const result = await runner.emit({ + type: "session_before_fork", + entryId, + ...options, + }); + return { cancelled: result?.cancel === true }; + } + + private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise { + await emitSessionShutdownEvent(this.session.extensionRunner, { + type: "session_shutdown", + reason, + targetSessionFile, + }); + this.beforeSessionInvalidate?.(); + this.session.dispose(); + } + + private apply(result: CreateAgentSessionRuntimeResult): void { + this._session = result.session; + this._services = result.services; + this._diagnostics = result.diagnostics; + this._modelFallbackMessage = result.modelFallbackMessage; + } + + private async finishSessionReplacement(withSession?: (ctx: ReplacedSessionContext) => Promise): Promise { + if (this.rebindSession) { + await this.rebindSession(this.session); + } + if (withSession) { + await withSession(this.session.createReplacedSessionContext()); + } + } + + async switchSession( + sessionPath: string, + options?: { + cwdOverride?: string; + withSession?: (ctx: ReplacedSessionContext) => Promise; + projectTrustContextFactory?: (cwd: string) => ProjectTrustContext; + }, + ): Promise<{ cancelled: boolean }> { + const beforeResult = await this.emitBeforeSwitch("resume", sessionPath); + if (beforeResult.cancelled) { + return beforeResult; + } + + const previousSessionFile = this.session.sessionFile; + const sessionManager = SessionManager.open(sessionPath, undefined, options?.cwdOverride); + assertSessionCwdExists(sessionManager, this.cwd); + await this.teardownCurrent("resume", sessionManager.getSessionFile()); + this.apply( + await this.createRuntime({ + cwd: sessionManager.getCwd(), + agentDir: this.services.agentDir, + sessionManager, + sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile }, + projectTrustContext: options?.projectTrustContextFactory?.(sessionManager.getCwd()), + }), + ); + await this.finishSessionReplacement(options?.withSession); + return { cancelled: false }; + } + + async newSession(options?: { + parentSession?: string; + setup?: (sessionManager: SessionManager) => Promise; + withSession?: (ctx: ReplacedSessionContext) => Promise; + }): Promise<{ cancelled: boolean }> { + const beforeResult = await this.emitBeforeSwitch("new"); + if (beforeResult.cancelled) { + return beforeResult; + } + + const previousSessionFile = this.session.sessionFile; + const sessionDir = this.session.sessionManager.getSessionDir(); + const sessionManager = this.session.sessionManager.isPersisted() + ? SessionManager.create(this.cwd, sessionDir) + : SessionManager.inMemory(this.cwd); + if (options?.parentSession) { + sessionManager.newSession({ parentSession: options.parentSession }); + } + + await this.teardownCurrent("new", sessionManager.getSessionFile()); + this.apply( + await this.createRuntime({ + cwd: this.cwd, + agentDir: this.services.agentDir, + sessionManager, + sessionStartEvent: { type: "session_start", reason: "new", previousSessionFile }, + }), + ); + if (options?.setup) { + await options.setup(this.session.sessionManager); + this.session.agent.state.messages = this.session.sessionManager.buildSessionContext().messages; + } + await this.finishSessionReplacement(options?.withSession); + return { cancelled: false }; + } + + async fork( + entryId: string, + options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise }, + ): Promise<{ cancelled: boolean; selectedText?: string }> { + const position = options?.position ?? "before"; + const beforeResult = await this.emitBeforeFork(entryId, { position }); + if (beforeResult.cancelled) { + return { cancelled: true }; + } + let targetLeafId: string | null; + let selectedText: string | undefined; + + const selectedEntry = this.session.sessionManager.getEntry(entryId); + if (!selectedEntry) { + throw new Error("Invalid entry ID for forking"); + } + + if (position === "at") { + targetLeafId = selectedEntry.id; + } else { + if (selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { + throw new Error("Invalid entry ID for forking"); + } + targetLeafId = selectedEntry.parentId; + selectedText = extractUserMessageText(selectedEntry.message.content); + } + + const previousSessionFile = this.session.sessionFile; + if (this.session.sessionManager.isPersisted()) { + const currentSessionFile = this.session.sessionFile; + if (!currentSessionFile) { + throw new Error("Persisted session is missing a session file"); + } + const sessionDir = this.session.sessionManager.getSessionDir(); + if (!targetLeafId) { + const sessionManager = SessionManager.create(this.cwd, sessionDir); + sessionManager.newSession({ parentSession: currentSessionFile }); + await this.teardownCurrent("fork", sessionManager.getSessionFile()); + this.apply( + await this.createRuntime({ + cwd: this.cwd, + agentDir: this.services.agentDir, + sessionManager, + sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile }, + }), + ); + await this.finishSessionReplacement(options?.withSession); + return { cancelled: false, selectedText }; + } + + const sessionManager = SessionManager.open(currentSessionFile, sessionDir); + const forkedSessionPath = sessionManager.createBranchedSession(targetLeafId); + if (!forkedSessionPath) { + throw new Error("Failed to create forked session"); + } + await this.teardownCurrent("fork", sessionManager.getSessionFile()); + this.apply( + await this.createRuntime({ + cwd: sessionManager.getCwd(), + agentDir: this.services.agentDir, + sessionManager, + sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile }, + }), + ); + await this.finishSessionReplacement(options?.withSession); + return { cancelled: false, selectedText }; + } + + const sessionManager = this.session.sessionManager; + if (!targetLeafId) { + sessionManager.newSession({ parentSession: this.session.sessionFile }); + } else { + sessionManager.createBranchedSession(targetLeafId); + } + await this.teardownCurrent("fork", sessionManager.getSessionFile()); + this.apply( + await this.createRuntime({ + cwd: this.cwd, + agentDir: this.services.agentDir, + sessionManager, + sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile }, + }), + ); + await this.finishSessionReplacement(options?.withSession); + return { cancelled: false, selectedText }; + } + + /** + * Import a session JSONL file and switch runtime state to the imported session. + * + * @returns `{ cancelled: true }` when cancelled by `session_before_switch`, otherwise `{ cancelled: false }`. + * @throws {SessionImportFileNotFoundError} When the input path does not exist. + * @throws {MissingSessionCwdError} When the imported session cwd cannot be resolved and no override is provided. + */ + async importFromJsonl(inputPath: string, cwdOverride?: string): Promise<{ cancelled: boolean }> { + const resolvedPath = resolvePath(inputPath); + if (!existsSync(resolvedPath)) { + throw new SessionImportFileNotFoundError(resolvedPath); + } + + const sessionDir = this.session.sessionManager.getSessionDir(); + if (!existsSync(sessionDir)) { + mkdirSync(sessionDir, { recursive: true }); + } + + const destinationPath = join(sessionDir, basename(resolvedPath)); + const beforeResult = await this.emitBeforeSwitch("resume", destinationPath); + if (beforeResult.cancelled) { + return beforeResult; + } + + const previousSessionFile = this.session.sessionFile; + if (resolve(destinationPath) !== resolvedPath) { + copyFileSync(resolvedPath, destinationPath); + } + + const sessionManager = SessionManager.open(destinationPath, sessionDir, cwdOverride); + assertSessionCwdExists(sessionManager, this.cwd); + await this.teardownCurrent("resume", sessionManager.getSessionFile()); + this.apply( + await this.createRuntime({ + cwd: sessionManager.getCwd(), + agentDir: this.services.agentDir, + sessionManager, + sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile }, + }), + ); + await this.finishSessionReplacement(); + return { cancelled: false }; + } + + async dispose(): Promise { + await emitSessionShutdownEvent(this.session.extensionRunner, { + type: "session_shutdown", + reason: "quit", + }); + this.beforeSessionInvalidate?.(); + this.session.dispose(); + } +} + +/** + * Create the initial runtime from a runtime factory and initial session target. + * + * The same factory is stored on the returned AgentSessionRuntime and reused for + * later /new, /resume, /fork, and import flows. + */ +export async function createAgentSessionRuntime( + createRuntime: CreateAgentSessionRuntimeFactory, + options: { + cwd: string; + agentDir: string; + sessionManager: SessionManager; + sessionStartEvent?: SessionStartEvent; + }, +): Promise { + assertSessionCwdExists(options.sessionManager, options.cwd); + const result = await createRuntime(options); + return new AgentSessionRuntime( + result.session, + result.services, + createRuntime, + result.diagnostics, + result.modelFallbackMessage, + ); +} + +export { + type AgentSessionRuntimeDiagnostic, + type AgentSessionServices, + type CreateAgentSessionFromServicesOptions, + type CreateAgentSessionServicesOptions, + createAgentSessionFromServices, + createAgentSessionServices, +} from "./agent-session-services.ts"; diff --git a/cactus-code/packages/coding-agent/src/core/agent-session-services.ts b/cactus-code/packages/coding-agent/src/core/agent-session-services.ts new file mode 100644 index 000000000..62ca3aec5 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/agent-session-services.ts @@ -0,0 +1,242 @@ +import { join } from "node:path"; +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { Model } from "@earendil-works/pi-ai"; +import { cactusBaseUrl, cactusModelFromId, fetchCactusModels } from "@earendil-works/pi-ai/providers/cactus"; +import { getAgentDir } from "../config.ts"; +import { resolvePath } from "../utils/paths.ts"; +import { AuthStorage } from "./auth-storage.ts"; +import type { SessionStartEvent, ToolDefinition } from "./extensions/index.ts"; +import { ModelRegistry } from "./model-registry.ts"; +import { + DefaultResourceLoader, + type DefaultResourceLoaderOptions, + type ResourceLoader, + type ResourceLoaderReloadOptions, +} from "./resource-loader.ts"; +import { type CreateAgentSessionOptions, type CreateAgentSessionResult, createAgentSession } from "./sdk.ts"; +import type { SessionManager } from "./session-manager.ts"; +import { SettingsManager } from "./settings-manager.ts"; + +/** + * Non-fatal issues collected while creating services or sessions. + * + * Runtime creation returns diagnostics to the caller instead of printing or + * exiting. The app layer decides whether warnings should be shown and whether + * errors should abort startup. + */ +export interface AgentSessionRuntimeDiagnostic { + type: "info" | "warning" | "error"; + message: string; +} + +/** + * Inputs for creating cwd-bound runtime services. + * + * These services are recreated whenever the effective session cwd changes. + * CLI-provided resource paths should be resolved to absolute paths before they + * reach this function, so later cwd switches do not reinterpret them. + */ +export interface CreateAgentSessionServicesOptions { + cwd: string; + agentDir?: string; + authStorage?: AuthStorage; + settingsManager?: SettingsManager; + modelRegistry?: ModelRegistry; + extensionFlagValues?: Map; + resourceLoaderOptions?: Omit; + resourceLoaderReloadOptions?: ResourceLoaderReloadOptions; +} + +/** + * Inputs for creating an AgentSession from already-created services. + * + * Use this after services exist and any cwd-bound model/tool/session options + * have been resolved against those services. + */ +export interface CreateAgentSessionFromServicesOptions { + services: AgentSessionServices; + sessionManager: SessionManager; + sessionStartEvent?: SessionStartEvent; + model?: Model; + thinkingLevel?: ThinkingLevel; + scopedModels?: Array<{ model: Model; thinkingLevel?: ThinkingLevel }>; + tools?: string[]; + excludeTools?: CreateAgentSessionOptions["excludeTools"]; + noTools?: CreateAgentSessionOptions["noTools"]; + customTools?: ToolDefinition[]; +} + +/** + * Coherent cwd-bound runtime services for one effective session cwd. + * + * This is infrastructure only. The AgentSession itself is created separately so + * session options can be resolved against these services first. + */ +export interface AgentSessionServices { + cwd: string; + agentDir: string; + authStorage: AuthStorage; + settingsManager: SettingsManager; + modelRegistry: ModelRegistry; + resourceLoader: ResourceLoader; + diagnostics: AgentSessionRuntimeDiagnostic[]; +} + +function applyExtensionFlagValues( + resourceLoader: ResourceLoader, + extensionFlagValues: Map | undefined, +): AgentSessionRuntimeDiagnostic[] { + if (!extensionFlagValues) { + return []; + } + + const diagnostics: AgentSessionRuntimeDiagnostic[] = []; + const extensionsResult = resourceLoader.getExtensions(); + const registeredFlags = new Map(); + for (const extension of extensionsResult.extensions) { + for (const [name, flag] of extension.flags) { + registeredFlags.set(name, { type: flag.type }); + } + } + + const unknownFlags: string[] = []; + for (const [name, value] of extensionFlagValues) { + const flag = registeredFlags.get(name); + if (!flag) { + unknownFlags.push(name); + continue; + } + if (flag.type === "boolean") { + extensionsResult.runtime.flagValues.set(name, true); + continue; + } + if (typeof value === "string") { + extensionsResult.runtime.flagValues.set(name, value); + continue; + } + diagnostics.push({ + type: "error", + message: `Extension flag "--${name}" requires a value`, + }); + } + + if (unknownFlags.length > 0) { + diagnostics.push({ + type: "error", + message: `Unknown option${unknownFlags.length === 1 ? "" : "s"}: ${unknownFlags.map((name) => `--${name}`).join(", ")}`, + }); + } + + return diagnostics; +} + +/** + * Create cwd-bound runtime services. + * + * Returns services plus diagnostics. It does not create an AgentSession. + */ +export async function createAgentSessionServices( + options: CreateAgentSessionServicesOptions, +): Promise { + const cwd = resolvePath(options.cwd); + const agentDir = options.agentDir ? resolvePath(options.agentDir) : getAgentDir(); + const authStorage = options.authStorage ?? AuthStorage.create(join(agentDir, "auth.json")); + const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir); + const modelRegistry = options.modelRegistry ?? ModelRegistry.create(authStorage, join(agentDir, "models.json")); + const resourceLoader = new DefaultResourceLoader({ + ...(options.resourceLoaderOptions ?? {}), + cwd, + agentDir, + settingsManager, + }); + await resourceLoader.reload(options.resourceLoaderReloadOptions); + + const diagnostics: AgentSessionRuntimeDiagnostic[] = []; + + const cactusUrl = cactusBaseUrl(); + let cactusModels: Model<"openai-completions">[] = []; + try { + cactusModels = await fetchCactusModels(cactusUrl); + } catch (error) { + diagnostics.push({ + type: "warning", + message: `Cactus model discovery failed at ${cactusUrl}: ${ + error instanceof Error ? error.message : String(error) + }. Start the local Cactus server and retry.`, + }); + } + if (cactusModels.length === 0) { + cactusModels = [cactusModelFromId(cactusUrl, "cactus")]; + } + modelRegistry.registerProvider("cactus", { + name: "Cactus", + baseUrl: cactusUrl, + api: "openai-completions", + apiKey: "cactus", + models: cactusModels.map((m) => ({ + id: m.id, + name: m.name, + api: m.api, + baseUrl: m.baseUrl, + reasoning: m.reasoning, + input: m.input, + cost: m.cost, + contextWindow: m.contextWindow, + maxTokens: m.maxTokens, + })), + }); + + const extensionsResult = resourceLoader.getExtensions(); + for (const { name, config, extensionPath } of extensionsResult.runtime.pendingProviderRegistrations) { + try { + modelRegistry.registerProvider(name, config); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + diagnostics.push({ + type: "error", + message: `Extension "${extensionPath}" error: ${message}`, + }); + } + } + extensionsResult.runtime.pendingProviderRegistrations = []; + diagnostics.push(...applyExtensionFlagValues(resourceLoader, options.extensionFlagValues)); + + return { + cwd, + agentDir, + authStorage, + settingsManager, + modelRegistry, + resourceLoader, + diagnostics, + }; +} + +/** + * Create an AgentSession from previously created services. + * + * This keeps session creation separate from service creation so callers can + * resolve model, thinking, tools, and other session inputs against the target + * cwd before constructing the session. + */ +export async function createAgentSessionFromServices( + options: CreateAgentSessionFromServicesOptions, +): Promise { + return createAgentSession({ + cwd: options.services.cwd, + agentDir: options.services.agentDir, + authStorage: options.services.authStorage, + settingsManager: options.services.settingsManager, + modelRegistry: options.services.modelRegistry, + resourceLoader: options.services.resourceLoader, + sessionManager: options.sessionManager, + model: options.model, + thinkingLevel: options.thinkingLevel, + scopedModels: options.scopedModels, + tools: options.tools, + excludeTools: options.excludeTools, + noTools: options.noTools, + customTools: options.customTools, + sessionStartEvent: options.sessionStartEvent, + }); +} diff --git a/cactus-code/packages/coding-agent/src/core/agent-session.ts b/cactus-code/packages/coding-agent/src/core/agent-session.ts new file mode 100644 index 000000000..83a541da6 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/agent-session.ts @@ -0,0 +1,3149 @@ +/** + * AgentSession - Core abstraction for agent lifecycle and session management. + * + * This class is shared between all run modes (interactive, print, rpc). + * It encapsulates: + * - Agent state access + * - Event subscription with automatic session persistence + * - Model and thinking level management + * - Compaction (manual and auto) + * - Bash execution + * - Session switching and branching + * + * Modes use this class and add their own I/O layer on top. + */ + +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { basename, dirname } from "node:path"; +import type { + Agent, + AgentEvent, + AgentMessage, + AgentState, + AgentTool, + ThinkingLevel, +} from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@earendil-works/pi-ai/compat"; +import { + clampThinkingLevel, + cleanupSessionResources, + getSupportedThinkingLevels, + isContextOverflow, + modelsAreEqual, + resetApiProviders, + streamSimple, +} from "@earendil-works/pi-ai/compat"; +import { getThemeByName, theme } from "../modes/interactive/theme/theme.ts"; +import { stripFrontmatter } from "../utils/frontmatter.ts"; +import { resolvePath } from "../utils/paths.ts"; +import { sleep } from "../utils/sleep.ts"; +import { formatNoApiKeyFoundMessage, formatNoModelSelectedMessage } from "./auth-guidance.ts"; +import { type BashResult, executeBashWithOperations } from "./bash-executor.ts"; +import { + type CompactionResult, + calculateContextTokens, + collectEntriesForBranchSummary, + compact, + estimateContextTokens, + estimateTokens, + generateBranchSummary, + prepareCompaction, + resolveCompactionSettings, + shouldCompact, +} from "./compaction/index.ts"; +import { DEFAULT_THINKING_LEVEL } from "./defaults.ts"; +import { + type ContextUsage, + type ExtensionCommandContextActions, + type ExtensionErrorListener, + type ExtensionMode, + ExtensionRunner, + type ExtensionUIContext, + type InputSource, + type MessageEndEvent, + type MessageStartEvent, + type MessageUpdateEvent, + type ReplacedSessionContext, + type SessionBeforeCompactResult, + type SessionBeforeTreeResult, + type SessionStartEvent, + type ShutdownHandler, + type ToolDefinition, + type ToolExecutionEndEvent, + type ToolExecutionStartEvent, + type ToolExecutionUpdateEvent, + type ToolInfo, + type TreePreparation, + type TurnEndEvent, + type TurnStartEvent, + wrapRegisteredTools, +} from "./extensions/index.ts"; +import { emitSessionShutdownEvent } from "./extensions/runner.ts"; +import type { BashExecutionMessage, CustomMessage } from "./messages.ts"; +import type { ModelRegistry } from "./model-registry.ts"; +import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts"; +import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts"; +import type { BranchSummaryEntry, CompactionEntry, SessionManager } from "./session-manager.ts"; +import { CURRENT_SESSION_VERSION, getLatestCompactionEntry, type SessionHeader } from "./session-manager.ts"; +import type { SettingsManager } from "./settings-manager.ts"; +import type { SlashCommandInfo } from "./slash-commands.ts"; +import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.ts"; +import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-prompt.ts"; +import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts"; +import { createAllToolDefinitions } from "./tools/index.ts"; +import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts"; + +// ============================================================================ +// Skill Block Parsing +// ============================================================================ + +/** Parsed skill block from a user message */ +export interface ParsedSkillBlock { + name: string; + location: string; + content: string; + userMessage: string | undefined; +} + +/** + * Parse a skill block from message text. + * Returns null if the text doesn't contain a skill block. + */ +export function parseSkillBlock(text: string): ParsedSkillBlock | null { + const match = text.match(/^\n([\s\S]*?)\n<\/skill>(?:\n\n([\s\S]+))?$/); + if (!match) return null; + return { + name: match[1], + location: match[2], + content: match[3], + userMessage: match[4]?.trim() || undefined, + }; +} + +/** Session-specific events that extend the core AgentEvent */ +export type AgentSessionEvent = + | Exclude + | { + type: "agent_end"; + messages: AgentMessage[]; + willRetry: boolean; + } + | { + type: "queue_update"; + steering: readonly string[]; + followUp: readonly string[]; + } + | { type: "compaction_start"; reason: "manual" | "threshold" | "overflow" } + | { type: "session_info_changed"; name: string | undefined } + | { type: "thinking_level_changed"; level: ThinkingLevel } + | { + type: "compaction_end"; + reason: "manual" | "threshold" | "overflow"; + result: CompactionResult | undefined; + aborted: boolean; + willRetry: boolean; + errorMessage?: string; + } + | { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string } + | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string }; + +/** Listener function for agent session events */ +export type AgentSessionEventListener = (event: AgentSessionEvent) => void; + +// ============================================================================ +// Types +// ============================================================================ + +export interface AgentSessionConfig { + agent: Agent; + sessionManager: SessionManager; + settingsManager: SettingsManager; + cwd: string; + /** Models to cycle through with Ctrl+P (from --models flag) */ + scopedModels?: Array<{ model: Model; thinkingLevel?: ThinkingLevel }>; + /** Resource loader for skills, prompts, themes, context files, system prompt */ + resourceLoader: ResourceLoader; + /** SDK custom tools registered outside extensions */ + customTools?: ToolDefinition[]; + /** Model registry for API key resolution and model discovery */ + modelRegistry: ModelRegistry; + /** Initial active built-in tool names. Default: [read, bash, edit, write] */ + initialActiveToolNames?: string[]; + /** Optional allowlist of tool names. When provided, only these tool names are exposed. */ + allowedToolNames?: string[]; + /** Optional denylist of tool names. When provided, these tool names are not exposed. */ + excludedToolNames?: string[]; + /** + * Override base tools (useful for custom runtimes). + * + * These are synthesized into minimal ToolDefinitions internally so AgentSession can keep + * a definition-first registry even when callers provide plain AgentTool instances. + */ + baseToolsOverride?: Record; + /** Mutable ref used by Agent to access the current ExtensionRunner */ + extensionRunnerRef?: { current?: ExtensionRunner }; + /** Session start event metadata emitted when extensions bind to this runtime. */ + sessionStartEvent?: SessionStartEvent; +} + +export interface ExtensionBindings { + uiContext?: ExtensionUIContext; + mode?: ExtensionMode; + commandContextActions?: ExtensionCommandContextActions; + abortHandler?: () => void; + shutdownHandler?: ShutdownHandler; + onError?: ExtensionErrorListener; +} + +/** Options for AgentSession.prompt() */ +export interface PromptOptions { + /** Whether to expand file-based prompt templates (default: true) */ + expandPromptTemplates?: boolean; + /** Image attachments */ + images?: ImageContent[]; + /** When streaming, how to queue the message: "steer" (interrupt) or "followUp" (wait). Required if streaming. */ + streamingBehavior?: "steer" | "followUp"; + /** Source of input for extension input event handlers. Defaults to "interactive". */ + source?: InputSource; + /** Internal hook used by RPC mode to observe prompt preflight acceptance or rejection. */ + preflightResult?: (success: boolean) => void; +} + +/** Result from cycleModel() */ +export interface ModelCycleResult { + model: Model; + thinkingLevel: ThinkingLevel; + /** Whether cycling through scoped models (--models flag) or all available */ + isScoped: boolean; +} + +/** Session statistics for /session command */ +export interface SessionStats { + sessionFile: string | undefined; + sessionId: string; + userMessages: number; + assistantMessages: number; + toolCalls: number; + toolResults: number; + totalMessages: number; + tokens: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; + cost: number; + contextUsage?: ContextUsage; +} + +interface ToolDefinitionEntry { + definition: ToolDefinition; + sourceInfo: SourceInfo; +} + +function estimateMessagesTokens(messages: AgentMessage[]): number { + let tokens = 0; + for (const message of messages) { + tokens += estimateTokens(message); + } + return tokens; +} + +// ============================================================================ +// Constants +// ============================================================================ + +/** Standard thinking levels */ +const THINKING_LEVELS: ThinkingLevel[] = ["off", "minimal", "low", "medium", "high"]; + +// ============================================================================ +// AgentSession Class +// ============================================================================ + +export class AgentSession { + readonly agent: Agent; + readonly sessionManager: SessionManager; + readonly settingsManager: SettingsManager; + + private _scopedModels: Array<{ model: Model; thinkingLevel?: ThinkingLevel }>; + + // Event subscription state + private _unsubscribeAgent?: () => void; + private _eventListeners: AgentSessionEventListener[] = []; + + /** Tracks pending steering messages for UI display. Removed when delivered. */ + private _steeringMessages: string[] = []; + /** Tracks pending follow-up messages for UI display. Removed when delivered. */ + private _followUpMessages: string[] = []; + /** Messages queued to be included with the next user prompt as context ("asides"). */ + private _pendingNextTurnMessages: CustomMessage[] = []; + + // Compaction state + private _compactionAbortController: AbortController | undefined = undefined; + private _autoCompactionAbortController: AbortController | undefined = undefined; + private _overflowRecoveryAttempted = false; + + // Branch summarization state + private _branchSummaryAbortController: AbortController | undefined = undefined; + + // Retry state + private _retryAbortController: AbortController | undefined = undefined; + private _retryAttempt = 0; + + // Bash execution state + private _bashAbortController: AbortController | undefined = undefined; + private _pendingBashMessages: BashExecutionMessage[] = []; + + // Extension system + private _extensionRunner!: ExtensionRunner; + private _turnIndex = 0; + + private _resourceLoader: ResourceLoader; + private _customTools: ToolDefinition[]; + private _baseToolDefinitions: Map = new Map(); + private _cwd: string; + private _extensionRunnerRef?: { current?: ExtensionRunner }; + private _initialActiveToolNames?: string[]; + private _allowedToolNames?: Set; + private _excludedToolNames?: Set; + private _baseToolsOverride?: Record; + private _sessionStartEvent: SessionStartEvent; + private _extensionUIContext?: ExtensionUIContext; + private _extensionMode: ExtensionMode = "print"; + private _extensionCommandContextActions?: ExtensionCommandContextActions; + private _extensionAbortHandler?: () => void; + private _extensionShutdownHandler?: ShutdownHandler; + private _extensionErrorListener?: ExtensionErrorListener; + private _extensionErrorUnsubscriber?: () => void; + + // Model registry for API key resolution + private _modelRegistry: ModelRegistry; + + // Tool registry for extension getTools/setTools + private _toolRegistry: Map = new Map(); + private _toolDefinitions: Map = new Map(); + private _toolPromptSnippets: Map = new Map(); + private _toolPromptGuidelines: Map = new Map(); + + // Base system prompt (without extension appends) - used to apply fresh appends each turn + private _baseSystemPrompt = ""; + private _baseSystemPromptOptions!: BuildSystemPromptOptions; + + constructor(config: AgentSessionConfig) { + this.agent = config.agent; + this.sessionManager = config.sessionManager; + this.settingsManager = config.settingsManager; + this._scopedModels = config.scopedModels ?? []; + this._resourceLoader = config.resourceLoader; + this._customTools = config.customTools ?? []; + this._cwd = config.cwd; + this._modelRegistry = config.modelRegistry; + this._extensionRunnerRef = config.extensionRunnerRef; + this._initialActiveToolNames = config.initialActiveToolNames; + this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined; + this._excludedToolNames = config.excludedToolNames ? new Set(config.excludedToolNames) : undefined; + this._baseToolsOverride = config.baseToolsOverride; + this._sessionStartEvent = config.sessionStartEvent ?? { type: "session_start", reason: "startup" }; + + // Always subscribe to agent events for internal handling + // (session persistence, extensions, auto-compaction, retry logic) + this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent); + this._installAgentToolHooks(); + + this._buildRuntime({ + activeToolNames: this._initialActiveToolNames, + includeAllExtensionTools: true, + }); + } + + /** Model registry for API key resolution and model discovery */ + get modelRegistry(): ModelRegistry { + return this._modelRegistry; + } + + private async _getRequiredRequestAuth(model: Model): Promise<{ + apiKey: string; + headers?: Record; + env?: Record; + }> { + const result = await this._modelRegistry.getApiKeyAndHeaders(model); + if (!result.ok) { + if (result.error.startsWith("No API key found")) { + throw new Error(formatNoApiKeyFoundMessage(model.provider)); + } + throw new Error(result.error); + } + if (result.apiKey) { + return { apiKey: result.apiKey, headers: result.headers, env: result.env }; + } + + const isOAuth = this._modelRegistry.isUsingOAuth(model); + if (isOAuth) { + throw new Error( + `Authentication failed for "${model.provider}". ` + + `Credentials may have expired or network is unavailable. ` + + `Run '/login ${model.provider}' to re-authenticate.`, + ); + } + throw new Error(formatNoApiKeyFoundMessage(model.provider)); + } + + private async _getCompactionRequestAuth(model: Model): Promise<{ + apiKey?: string; + headers?: Record; + env?: Record; + }> { + if (this.agent.streamFn === streamSimple) { + return this._getRequiredRequestAuth(model); + } + + const result = await this._modelRegistry.getApiKeyAndHeaders(model); + return result.ok ? { apiKey: result.apiKey, headers: result.headers, env: result.env } : {}; + } + + /** + * Install tool hooks once on the Agent instance. + * + * The callbacks read `this._extensionRunner` at execution time, so extension reload swaps in the + * new runner without reinstalling hooks. Extension-specific tool wrappers are still used to adapt + * registered tool execution to the extension context. Tool call and tool result interception now + * happens here instead of in wrappers. + */ + private _installAgentToolHooks(): void { + this.agent.beforeToolCall = async ({ toolCall, args }) => { + const runner = this._extensionRunner; + if (!runner.hasHandlers("tool_call")) { + return undefined; + } + + try { + return await runner.emitToolCall({ + type: "tool_call", + toolName: toolCall.name, + toolCallId: toolCall.id, + input: args as Record, + }); + } catch (err) { + if (err instanceof Error) { + throw err; + } + throw new Error(`Extension failed, blocking execution: ${String(err)}`); + } + }; + + this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => { + const runner = this._extensionRunner; + if (!runner.hasHandlers("tool_result")) { + return undefined; + } + + const hookResult = await runner.emitToolResult({ + type: "tool_result", + toolName: toolCall.name, + toolCallId: toolCall.id, + input: args as Record, + content: result.content, + details: result.details, + isError, + }); + + if (!hookResult) { + return undefined; + } + + return { + content: hookResult.content, + details: hookResult.details, + isError: hookResult.isError ?? isError, + }; + }; + } + + // ========================================================================= + // Event Subscription + // ========================================================================= + + /** Emit an event to all listeners */ + private _emit(event: AgentSessionEvent): void { + for (const l of this._eventListeners) { + l(event); + } + } + + private _emitQueueUpdate(): void { + this._emit({ + type: "queue_update", + steering: [...this._steeringMessages], + followUp: [...this._followUpMessages], + }); + } + + // Track last assistant message for auto-compaction check + private _lastAssistantMessage: AssistantMessage | undefined = undefined; + + /** Internal handler for agent events - shared by subscribe and reconnect */ + private _handleAgentEvent = async (event: AgentEvent): Promise => { + // When a user message starts, check if it's from either queue and remove it BEFORE emitting + // This ensures the UI sees the updated queue state + if (event.type === "message_start" && event.message.role === "user") { + this._overflowRecoveryAttempted = false; + const messageText = this._getUserMessageText(event.message); + if (messageText) { + // Check steering queue first + const steeringIndex = this._steeringMessages.indexOf(messageText); + if (steeringIndex !== -1) { + this._steeringMessages.splice(steeringIndex, 1); + this._emitQueueUpdate(); + } else { + // Check follow-up queue + const followUpIndex = this._followUpMessages.indexOf(messageText); + if (followUpIndex !== -1) { + this._followUpMessages.splice(followUpIndex, 1); + this._emitQueueUpdate(); + } + } + } + } + + // Emit to extensions first + await this._emitExtensionEvent(event); + + // Notify all listeners + this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event); + + // Handle session persistence + if (event.type === "message_end") { + // Check if this is a custom message from extensions + if (event.message.role === "custom") { + // Persist as CustomMessageEntry + this.sessionManager.appendCustomMessageEntry( + event.message.customType, + event.message.content, + event.message.display, + event.message.details, + ); + } else if ( + event.message.role === "user" || + event.message.role === "assistant" || + event.message.role === "toolResult" + ) { + // Regular LLM message - persist as SessionMessageEntry + this.sessionManager.appendMessage(event.message); + } + // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere + + // Track assistant message for auto-compaction (checked on agent_end) + if (event.message.role === "assistant") { + this._lastAssistantMessage = event.message; + + const assistantMsg = event.message as AssistantMessage; + if (assistantMsg.stopReason !== "error") { + this._overflowRecoveryAttempted = false; + } + + // Reset retry counter immediately on successful assistant response + // This prevents accumulation across multiple LLM calls within a turn + if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) { + this._emit({ + type: "auto_retry_end", + success: true, + attempt: this._retryAttempt, + }); + this._retryAttempt = 0; + } + } + } + }; + + private _willRetryAfterAgentEnd(event: Extract): boolean { + const settings = this.settingsManager.getRetrySettings(); + if (!settings.enabled || this._retryAttempt >= settings.maxRetries) { + return false; + } + + for (let i = event.messages.length - 1; i >= 0; i--) { + const message = event.messages[i]; + if (message.role === "assistant") { + return this._isRetryableError(message as AssistantMessage); + } + } + return false; + } + + /** Extract text content from a message */ + private _getUserMessageText(message: Message): string { + if (message.role !== "user") return ""; + const content = message.content; + if (typeof content === "string") return content; + const textBlocks = content.filter((c) => c.type === "text"); + return textBlocks.map((c) => (c as TextContent).text).join(""); + } + + /** Find the last assistant message in agent state (including aborted ones) */ + private _findLastAssistantMessage(): AssistantMessage | undefined { + const messages = this.agent.state.messages; + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]; + if (msg.role === "assistant") { + return msg as AssistantMessage; + } + } + return undefined; + } + + private _replaceMessageInPlace(target: AgentMessage, replacement: AgentMessage): void { + // Agent-core stores the finalized message object in its state before emitting message_end. + // SessionManager persistence happens later in _handleAgentEvent() with event.message. + // Mutating this object in place keeps agent state, later turn/agent events, listeners, + // and the eventual SessionManager.appendMessage(event.message) persistence in sync. + if (target === replacement) { + return; + } + + const targetRecord = target as unknown as Record; + for (const key of Object.keys(targetRecord)) { + delete targetRecord[key]; + } + Object.assign(targetRecord, replacement); + } + + /** Emit extension events based on agent events */ + private async _emitExtensionEvent(event: AgentEvent): Promise { + if (event.type === "agent_start") { + this._turnIndex = 0; + await this._extensionRunner.emit({ type: "agent_start" }); + } else if (event.type === "agent_end") { + await this._extensionRunner.emit({ type: "agent_end", messages: event.messages }); + } else if (event.type === "turn_start") { + const extensionEvent: TurnStartEvent = { + type: "turn_start", + turnIndex: this._turnIndex, + timestamp: Date.now(), + }; + await this._extensionRunner.emit(extensionEvent); + } else if (event.type === "turn_end") { + const extensionEvent: TurnEndEvent = { + type: "turn_end", + turnIndex: this._turnIndex, + message: event.message, + toolResults: event.toolResults, + }; + await this._extensionRunner.emit(extensionEvent); + this._turnIndex++; + } else if (event.type === "message_start") { + const extensionEvent: MessageStartEvent = { + type: "message_start", + message: event.message, + }; + await this._extensionRunner.emit(extensionEvent); + } else if (event.type === "message_update") { + const extensionEvent: MessageUpdateEvent = { + type: "message_update", + message: event.message, + assistantMessageEvent: event.assistantMessageEvent, + }; + await this._extensionRunner.emit(extensionEvent); + } else if (event.type === "message_end") { + const extensionEvent: MessageEndEvent = { + type: "message_end", + message: event.message, + }; + const replacement = await this._extensionRunner.emitMessageEnd(extensionEvent); + if (replacement) { + this._replaceMessageInPlace(event.message, replacement); + } + } else if (event.type === "tool_execution_start") { + const extensionEvent: ToolExecutionStartEvent = { + type: "tool_execution_start", + toolCallId: event.toolCallId, + toolName: event.toolName, + args: event.args, + }; + await this._extensionRunner.emit(extensionEvent); + } else if (event.type === "tool_execution_update") { + const extensionEvent: ToolExecutionUpdateEvent = { + type: "tool_execution_update", + toolCallId: event.toolCallId, + toolName: event.toolName, + args: event.args, + partialResult: event.partialResult, + }; + await this._extensionRunner.emit(extensionEvent); + } else if (event.type === "tool_execution_end") { + const extensionEvent: ToolExecutionEndEvent = { + type: "tool_execution_end", + toolCallId: event.toolCallId, + toolName: event.toolName, + result: event.result, + isError: event.isError, + }; + await this._extensionRunner.emit(extensionEvent); + } + } + + /** + * Subscribe to agent events. + * Session persistence is handled internally (saves messages on message_end). + * Multiple listeners can be added. Returns unsubscribe function for this listener. + */ + subscribe(listener: AgentSessionEventListener): () => void { + this._eventListeners.push(listener); + + // Return unsubscribe function for this specific listener + return () => { + const index = this._eventListeners.indexOf(listener); + if (index !== -1) { + this._eventListeners.splice(index, 1); + } + }; + } + + /** + * Temporarily disconnect from agent events. + * User listeners are preserved and will receive events again after resubscribe(). + * Used internally during operations that need to pause event processing. + */ + private _disconnectFromAgent(): void { + if (this._unsubscribeAgent) { + this._unsubscribeAgent(); + this._unsubscribeAgent = undefined; + } + } + + /** + * Reconnect to agent events after _disconnectFromAgent(). + * Preserves all existing listeners. + */ + private _reconnectToAgent(): void { + if (this._unsubscribeAgent) return; // Already connected + this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent); + } + + /** + * Remove all listeners and disconnect from agent. + * Call this when completely done with the session. + */ + dispose(): void { + try { + this.abortRetry(); + this.abortCompaction(); + this.abortBranchSummary(); + this.abortBash(); + this.agent.abort(); + } catch { + // Dispose must succeed even if an abort hook throws. + } + + this._extensionRunner.invalidate( + "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", + ); + this._disconnectFromAgent(); + this._eventListeners = []; + cleanupSessionResources(this.sessionId); + } + + // ========================================================================= + // Read-only State Access + // ========================================================================= + + /** Full agent state */ + get state(): AgentState { + return this.agent.state; + } + + /** Current model (may be undefined if not yet selected) */ + get model(): Model | undefined { + return this.agent.state.model; + } + + /** Current thinking level */ + get thinkingLevel(): ThinkingLevel { + return this.agent.state.thinkingLevel; + } + + /** Whether agent is currently streaming a response */ + get isStreaming(): boolean { + return this.agent.state.isStreaming; + } + + /** Current effective system prompt (includes any per-turn extension modifications) */ + get systemPrompt(): string { + return this.agent.state.systemPrompt; + } + + /** Current retry attempt (0 if not retrying) */ + get retryAttempt(): number { + return this._retryAttempt; + } + + /** + * Get the names of currently active tools. + * Returns the names of tools currently set on the agent. + */ + getActiveToolNames(): string[] { + return this.agent.state.tools.map((t) => t.name); + } + + /** + * Get all configured tools with name, description, parameter schema, prompt guidelines, and source metadata. + */ + getAllTools(): ToolInfo[] { + return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({ + name: definition.name, + description: definition.description, + parameters: definition.parameters, + promptGuidelines: definition.promptGuidelines, + sourceInfo, + })); + } + + getToolDefinition(name: string): ToolDefinition | undefined { + return this._toolDefinitions.get(name)?.definition; + } + + /** + * Set active tools by name. + * Only tools in the registry can be enabled. Unknown tool names are ignored. + * Also rebuilds the system prompt to reflect the new tool set. + * Changes take effect on the next agent turn. + */ + setActiveToolsByName(toolNames: string[]): void { + const tools: AgentTool[] = []; + const validToolNames: string[] = []; + for (const name of toolNames) { + const tool = this._toolRegistry.get(name); + if (tool) { + tools.push(tool); + validToolNames.push(name); + } + } + this.agent.state.tools = tools; + + // Rebuild base system prompt with new tool set + this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames); + this.agent.state.systemPrompt = this._baseSystemPrompt; + } + + /** Whether compaction or branch summarization is currently running */ + get isCompacting(): boolean { + return ( + this._autoCompactionAbortController !== undefined || + this._compactionAbortController !== undefined || + this._branchSummaryAbortController !== undefined + ); + } + + /** All messages including custom types like BashExecutionMessage */ + get messages(): AgentMessage[] { + return this.agent.state.messages; + } + + /** Current steering mode */ + get steeringMode(): "all" | "one-at-a-time" { + return this.agent.steeringMode; + } + + /** Current follow-up mode */ + get followUpMode(): "all" | "one-at-a-time" { + return this.agent.followUpMode; + } + + /** Current session file path, or undefined if sessions are disabled */ + get sessionFile(): string | undefined { + return this.sessionManager.getSessionFile(); + } + + /** Current session ID */ + get sessionId(): string { + return this.sessionManager.getSessionId(); + } + + /** Current session display name, if set */ + get sessionName(): string | undefined { + return this.sessionManager.getSessionName(); + } + + /** Scoped models for cycling (from --models flag) */ + get scopedModels(): ReadonlyArray<{ model: Model; thinkingLevel?: ThinkingLevel }> { + return this._scopedModels; + } + + /** Update scoped models for cycling */ + setScopedModels(scopedModels: Array<{ model: Model; thinkingLevel?: ThinkingLevel }>): void { + this._scopedModels = scopedModels; + } + + /** File-based prompt templates */ + get promptTemplates(): ReadonlyArray { + return this._resourceLoader.getPrompts().prompts; + } + + private _normalizePromptSnippet(text: string | undefined): string | undefined { + if (!text) return undefined; + const oneLine = text + .replace(/[\r\n]+/g, " ") + .replace(/\s+/g, " ") + .trim(); + return oneLine.length > 0 ? oneLine : undefined; + } + + private _normalizePromptGuidelines(guidelines: string[] | undefined): string[] { + if (!guidelines || guidelines.length === 0) { + return []; + } + + const unique = new Set(); + for (const guideline of guidelines) { + const normalized = guideline.trim(); + if (normalized.length > 0) { + unique.add(normalized); + } + } + return Array.from(unique); + } + + private _rebuildSystemPrompt(toolNames: string[]): string { + const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name)); + const toolSnippets: Record = {}; + const promptGuidelines: string[] = []; + for (const name of validToolNames) { + const snippet = this._toolPromptSnippets.get(name); + if (snippet) { + toolSnippets[name] = snippet; + } + + const toolGuidelines = this._toolPromptGuidelines.get(name); + if (toolGuidelines) { + promptGuidelines.push(...toolGuidelines); + } + } + + const loaderSystemPrompt = this._resourceLoader.getSystemPrompt(); + const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt(); + const appendSystemPrompt = + loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined; + const loadedSkills = this._resourceLoader.getSkills().skills; + const loadedContextFiles = this._resourceLoader.getAgentsFiles().agentsFiles; + + this._baseSystemPromptOptions = { + cwd: this._cwd, + skills: loadedSkills, + contextFiles: loadedContextFiles, + customPrompt: loaderSystemPrompt, + appendSystemPrompt, + selectedTools: validToolNames, + toolSnippets, + promptGuidelines, + }; + return buildSystemPrompt(this._baseSystemPromptOptions); + } + + // ========================================================================= + // Prompting + // ========================================================================= + + private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise { + try { + await this.agent.prompt(messages); + while (await this._handlePostAgentRun()) { + await this.agent.continue(); + } + } finally { + this._flushPendingBashMessages(); + } + } + + private async _handlePostAgentRun(): Promise { + const msg = this._lastAssistantMessage; + this._lastAssistantMessage = undefined; + if (!msg) { + return false; + } + + if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) { + return true; + } + + if (msg.stopReason === "error" && this._retryAttempt > 0) { + this._emit({ + type: "auto_retry_end", + success: false, + attempt: this._retryAttempt, + finalError: msg.errorMessage, + }); + this._retryAttempt = 0; + } + + if (await this._checkCompaction(msg)) { + return true; + } + + // The agent loop drains both queues before emitting agent_end. Any messages + // here were queued by agent_end extension handlers and need a continuation. + return this.agent.hasQueuedMessages(); + } + + /** + * Send a prompt to the agent. + * - Handles extension commands (registered via pi.registerCommand) immediately, even during streaming + * - Expands file-based prompt templates by default + * - During streaming, queues via steer() or followUp() based on streamingBehavior option + * - Validates model and API key before sending (when not streaming) + * @throws Error if streaming and no streamingBehavior specified + * @throws Error if no model selected or no API key available (when not streaming) + */ + async prompt(text: string, options?: PromptOptions): Promise { + const expandPromptTemplates = options?.expandPromptTemplates ?? true; + const preflightResult = options?.preflightResult; + let messages: AgentMessage[] | undefined; + + try { + // Handle extension commands first (execute immediately, even during streaming) + // Extension commands manage their own LLM interaction via pi.sendMessage() + if (expandPromptTemplates && text.startsWith("/")) { + const handled = await this._tryExecuteExtensionCommand(text); + if (handled) { + // Extension command executed, no prompt to send + preflightResult?.(true); + return; + } + } + + // Emit input event for extension interception (before skill/template expansion) + let currentText = text; + let currentImages = options?.images; + if (this._extensionRunner.hasHandlers("input")) { + const inputResult = await this._extensionRunner.emitInput( + currentText, + currentImages, + options?.source ?? "interactive", + this.isStreaming ? options?.streamingBehavior : undefined, + ); + if (inputResult.action === "handled") { + preflightResult?.(true); + return; + } + if (inputResult.action === "transform") { + currentText = inputResult.text; + currentImages = inputResult.images ?? currentImages; + } + } + + // Expand skill commands (/skill:name args) and prompt templates (/template args) + let expandedText = currentText; + if (expandPromptTemplates) { + expandedText = this._expandSkillCommand(expandedText); + expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]); + } + + // If streaming, queue via steer() or followUp() based on option + if (this.isStreaming) { + if (!options?.streamingBehavior) { + throw new Error( + "Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", + ); + } + if (options.streamingBehavior === "followUp") { + await this._queueFollowUp(expandedText, currentImages); + } else { + await this._queueSteer(expandedText, currentImages); + } + preflightResult?.(true); + return; + } + + // Flush any pending bash messages before the new prompt + this._flushPendingBashMessages(); + + // Validate model + if (!this.model) { + throw new Error(formatNoModelSelectedMessage()); + } + + if (!this._modelRegistry.hasConfiguredAuth(this.model)) { + const isOAuth = this._modelRegistry.isUsingOAuth(this.model); + if (isOAuth) { + throw new Error( + `Authentication failed for "${this.model.provider}". ` + + `Credentials may have expired or network is unavailable. ` + + `Run '/login ${this.model.provider}' to re-authenticate.`, + ); + } + throw new Error(formatNoApiKeyFoundMessage(this.model.provider)); + } + + // Check if we need to compact before sending (catches aborted responses) + const lastAssistant = this._findLastAssistantMessage(); + if (lastAssistant && (await this._checkCompaction(lastAssistant, false))) { + try { + await this.agent.continue(); + while (await this._handlePostAgentRun()) { + await this.agent.continue(); + } + } finally { + this._flushPendingBashMessages(); + } + } + + // Build messages array (custom message if any, then user message) + messages = []; + + // Add user message + const userContent: (TextContent | ImageContent)[] = [{ type: "text", text: expandedText }]; + if (currentImages) { + userContent.push(...currentImages); + } + messages.push({ + role: "user", + content: userContent, + timestamp: Date.now(), + }); + + // Inject any pending "nextTurn" messages as context alongside the user message + for (const msg of this._pendingNextTurnMessages) { + messages.push(msg); + } + this._pendingNextTurnMessages = []; + + // Emit before_agent_start extension event + const result = await this._extensionRunner.emitBeforeAgentStart( + expandedText, + currentImages, + this._baseSystemPrompt, + this._baseSystemPromptOptions, + ); + // Add all custom messages from extensions + if (result?.messages) { + for (const msg of result.messages) { + messages.push({ + role: "custom", + customType: msg.customType, + content: msg.content, + display: msg.display, + details: msg.details, + timestamp: Date.now(), + }); + } + } + // Apply extension-modified system prompt, or reset to base + if (result?.systemPrompt) { + this.agent.state.systemPrompt = result.systemPrompt; + } else { + // Ensure we're using the base prompt (in case previous turn had modifications) + this.agent.state.systemPrompt = this._baseSystemPrompt; + } + } catch (error) { + preflightResult?.(false); + throw error; + } + + if (!messages) { + return; + } + + preflightResult?.(true); + await this._runAgentPrompt(messages); + } + + /** + * Try to execute an extension command. Returns true if command was found and executed. + */ + private async _tryExecuteExtensionCommand(text: string): Promise { + // Parse command name and args + const spaceIndex = text.indexOf(" "); + const commandName = spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex); + const args = spaceIndex === -1 ? "" : text.slice(spaceIndex + 1); + + const command = this._extensionRunner.getCommand(commandName); + if (!command) return false; + + // Get command context from extension runner (includes session control methods) + const ctx = this._extensionRunner.createCommandContext(); + + try { + await command.handler(args, ctx); + return true; + } catch (err) { + // Emit error via extension runner + this._extensionRunner.emitError({ + extensionPath: `command:${commandName}`, + event: "command", + error: err instanceof Error ? err.message : String(err), + }); + return true; + } + } + + /** + * Expand skill commands (/skill:name args) to their full content. + * Returns the expanded text, or the original text if not a skill command or skill not found. + * Emits errors via extension runner if file read fails. + */ + private _expandSkillCommand(text: string): string { + if (!text.startsWith("/skill:")) return text; + + const spaceIndex = text.indexOf(" "); + const skillName = spaceIndex === -1 ? text.slice(7) : text.slice(7, spaceIndex); + const args = spaceIndex === -1 ? "" : text.slice(spaceIndex + 1).trim(); + + const skill = this.resourceLoader.getSkills().skills.find((s) => s.name === skillName); + if (!skill) return text; // Unknown skill, pass through + + try { + const content = readFileSync(skill.filePath, "utf-8"); + const body = stripFrontmatter(content).trim(); + const skillBlock = `\nReferences are relative to ${skill.baseDir}.\n\n${body}\n`; + return args ? `${skillBlock}\n\n${args}` : skillBlock; + } catch (err) { + // Emit error like extension commands do + this._extensionRunner.emitError({ + extensionPath: skill.filePath, + event: "skill_expansion", + error: err instanceof Error ? err.message : String(err), + }); + return text; // Return original on error + } + } + + /** + * Queue a steering message while the agent is running. + * Delivered after the current assistant turn finishes executing its tool calls, + * before the next LLM call. + * Expands skill commands and prompt templates. Errors on extension commands. + * @param images Optional image attachments to include with the message + * @throws Error if text is an extension command + */ + async steer(text: string, images?: ImageContent[]): Promise { + // Check for extension commands (cannot be queued) + if (text.startsWith("/")) { + this._throwIfExtensionCommand(text); + } + + // Expand skill commands and prompt templates + let expandedText = this._expandSkillCommand(text); + expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]); + + await this._queueSteer(expandedText, images); + } + + /** + * Queue a follow-up message to be processed after the agent finishes. + * Delivered only when agent has no more tool calls or steering messages. + * Expands skill commands and prompt templates. Errors on extension commands. + * @param images Optional image attachments to include with the message + * @throws Error if text is an extension command + */ + async followUp(text: string, images?: ImageContent[]): Promise { + // Check for extension commands (cannot be queued) + if (text.startsWith("/")) { + this._throwIfExtensionCommand(text); + } + + // Expand skill commands and prompt templates + let expandedText = this._expandSkillCommand(text); + expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]); + + await this._queueFollowUp(expandedText, images); + } + + /** + * Internal: Queue a steering message (already expanded, no extension command check). + */ + private async _queueSteer(text: string, images?: ImageContent[]): Promise { + this._steeringMessages.push(text); + this._emitQueueUpdate(); + const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; + if (images) { + content.push(...images); + } + this.agent.steer({ + role: "user", + content, + timestamp: Date.now(), + }); + } + + /** + * Internal: Queue a follow-up message (already expanded, no extension command check). + */ + private async _queueFollowUp(text: string, images?: ImageContent[]): Promise { + this._followUpMessages.push(text); + this._emitQueueUpdate(); + const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; + if (images) { + content.push(...images); + } + this.agent.followUp({ + role: "user", + content, + timestamp: Date.now(), + }); + } + + /** + * Throw an error if the text is an extension command. + */ + private _throwIfExtensionCommand(text: string): void { + const spaceIndex = text.indexOf(" "); + const commandName = spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex); + const command = this._extensionRunner.getCommand(commandName); + + if (command) { + throw new Error( + `Extension command "/${commandName}" cannot be queued. Use prompt() or execute the command when not streaming.`, + ); + } + } + + /** + * Send a custom message to the session. Creates a CustomMessageEntry. + * + * Handles three cases: + * - Streaming: queues message, processed when loop pulls from queue + * - Not streaming + triggerTurn: appends to state/session, starts new turn + * - Not streaming + no trigger: appends to state/session, no turn + * + * @param message Custom message with customType, content, display, details + * @param options.triggerTurn If true and not streaming, triggers a new LLM turn + * @param options.deliverAs Delivery mode: "steer", "followUp", or "nextTurn" + */ + async sendCustomMessage( + message: Pick, "customType" | "content" | "display" | "details">, + options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, + ): Promise { + const appMessage = { + role: "custom" as const, + customType: message.customType, + content: message.content, + display: message.display, + details: message.details, + timestamp: Date.now(), + } satisfies CustomMessage; + if (options?.deliverAs === "nextTurn") { + this._pendingNextTurnMessages.push(appMessage); + } else if (this.isStreaming) { + if (options?.deliverAs === "followUp") { + this.agent.followUp(appMessage); + } else { + this.agent.steer(appMessage); + } + } else if (options?.triggerTurn) { + await this._runAgentPrompt(appMessage); + } else { + this.agent.state.messages.push(appMessage); + this.sessionManager.appendCustomMessageEntry( + message.customType, + message.content, + message.display, + message.details, + ); + this._emit({ type: "message_start", message: appMessage }); + this._emit({ type: "message_end", message: appMessage }); + } + } + + /** + * Send a user message to the agent. Always triggers a turn. + * When the agent is streaming, use deliverAs to specify how to queue the message. + * + * @param content User message content (string or content array) + * @param options.deliverAs Delivery mode when streaming: "steer" or "followUp" + */ + async sendUserMessage( + content: string | (TextContent | ImageContent)[], + options?: { deliverAs?: "steer" | "followUp" }, + ): Promise { + // Normalize content to text string + optional images + let text: string; + let images: ImageContent[] | undefined; + + if (typeof content === "string") { + text = content; + } else { + const textParts: string[] = []; + images = []; + for (const part of content) { + if (part.type === "text") { + textParts.push(part.text); + } else { + images.push(part); + } + } + text = textParts.join("\n"); + if (images.length === 0) images = undefined; + } + + // Use prompt() with expandPromptTemplates: false to skip command handling and template expansion + await this.prompt(text, { + expandPromptTemplates: false, + streamingBehavior: options?.deliverAs, + images, + source: "extension", + }); + } + + /** + * Clear all queued messages and return them. + * Useful for restoring to editor when user aborts. + * @returns Object with steering and followUp arrays + */ + clearQueue(): { steering: string[]; followUp: string[] } { + const steering = [...this._steeringMessages]; + const followUp = [...this._followUpMessages]; + this._steeringMessages = []; + this._followUpMessages = []; + this.agent.clearAllQueues(); + this._emitQueueUpdate(); + return { steering, followUp }; + } + + /** Number of pending messages (includes both steering and follow-up) */ + get pendingMessageCount(): number { + return this._steeringMessages.length + this._followUpMessages.length; + } + + /** Get pending steering messages (read-only) */ + getSteeringMessages(): readonly string[] { + return this._steeringMessages; + } + + /** Get pending follow-up messages (read-only) */ + getFollowUpMessages(): readonly string[] { + return this._followUpMessages; + } + + get resourceLoader(): ResourceLoader { + return this._resourceLoader; + } + + /** + * Abort current operation and wait for agent to become idle. + */ + async abort(): Promise { + this.abortRetry(); + this.agent.abort(); + await this.agent.waitForIdle(); + } + + // ========================================================================= + // Model Management + // ========================================================================= + + private async _emitModelSelect( + nextModel: Model, + previousModel: Model | undefined, + source: "set" | "cycle" | "restore", + ): Promise { + if (modelsAreEqual(previousModel, nextModel)) return; + await this._extensionRunner.emit({ + type: "model_select", + model: nextModel, + previousModel, + source, + }); + } + + /** + * Set model directly. + * Validates that auth is configured, saves to session and settings. + * @throws Error if no auth is configured for the model + */ + async setModel(model: Model): Promise { + if (!this._modelRegistry.hasConfiguredAuth(model)) { + throw new Error(`No API key for ${model.provider}/${model.id}`); + } + + const previousModel = this.model; + const thinkingLevel = this._getThinkingLevelForModelSwitch(); + this.agent.state.model = model; + this.sessionManager.appendModelChange(model.provider, model.id); + this.settingsManager.setDefaultModelAndProvider(model.provider, model.id); + + // Re-clamp thinking level for new model's capabilities + this.setThinkingLevel(thinkingLevel); + + await this._emitModelSelect(model, previousModel, "set"); + } + + /** + * Cycle to next/previous model. + * Uses scoped models (from --models flag) if available, otherwise all available models. + * @param direction - "forward" (default) or "backward" + * @returns The new model info, or undefined if only one model available + */ + async cycleModel(direction: "forward" | "backward" = "forward"): Promise { + if (this._scopedModels.length > 0) { + return this._cycleScopedModel(direction); + } + return this._cycleAvailableModel(direction); + } + + private async _cycleScopedModel(direction: "forward" | "backward"): Promise { + const scopedModels = this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model)); + if (scopedModels.length <= 1) return undefined; + + const currentModel = this.model; + let currentIndex = scopedModels.findIndex((sm) => modelsAreEqual(sm.model, currentModel)); + + if (currentIndex === -1) currentIndex = 0; + const len = scopedModels.length; + const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; + const next = scopedModels[nextIndex]; + const thinkingLevel = this._getThinkingLevelForModelSwitch(next.thinkingLevel); + + // Apply model + this.agent.state.model = next.model; + this.sessionManager.appendModelChange(next.model.provider, next.model.id); + this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id); + + // Apply thinking level. + // - Explicit scoped model thinking level overrides current session level + // - Undefined scoped model thinking level inherits the current session preference + // setThinkingLevel clamps to model capabilities. + this.setThinkingLevel(thinkingLevel); + + await this._emitModelSelect(next.model, currentModel, "cycle"); + + return { model: next.model, thinkingLevel: this.thinkingLevel, isScoped: true }; + } + + private async _cycleAvailableModel(direction: "forward" | "backward"): Promise { + const availableModels = await this._modelRegistry.getAvailable(); + if (availableModels.length <= 1) return undefined; + + const currentModel = this.model; + let currentIndex = availableModels.findIndex((m) => modelsAreEqual(m, currentModel)); + + if (currentIndex === -1) currentIndex = 0; + const len = availableModels.length; + const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; + const nextModel = availableModels[nextIndex]; + + const thinkingLevel = this._getThinkingLevelForModelSwitch(); + this.agent.state.model = nextModel; + this.sessionManager.appendModelChange(nextModel.provider, nextModel.id); + this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id); + + // Re-clamp thinking level for new model's capabilities + this.setThinkingLevel(thinkingLevel); + + await this._emitModelSelect(nextModel, currentModel, "cycle"); + + return { model: nextModel, thinkingLevel: this.thinkingLevel, isScoped: false }; + } + + // ========================================================================= + // Thinking Level Management + // ========================================================================= + + /** + * Set thinking level. + * Clamps to model capabilities based on available thinking levels. + * Saves to session and settings only if the level actually changes. + */ + setThinkingLevel(level: ThinkingLevel): void { + const availableLevels = this.getAvailableThinkingLevels(); + const effectiveLevel = availableLevels.includes(level) ? level : this._clampThinkingLevel(level, availableLevels); + + // Only persist if actually changing + const previousLevel = this.agent.state.thinkingLevel; + const isChanging = effectiveLevel !== previousLevel; + + this.agent.state.thinkingLevel = effectiveLevel; + + if (isChanging) { + this.sessionManager.appendThinkingLevelChange(effectiveLevel); + if (this.supportsThinking() || effectiveLevel !== "off") { + this.settingsManager.setDefaultThinkingLevel(effectiveLevel); + } + this._emit({ type: "thinking_level_changed", level: effectiveLevel }); + void this._extensionRunner.emit({ + type: "thinking_level_select", + level: effectiveLevel, + previousLevel, + }); + } + } + + /** + * Cycle to next thinking level. + * @returns New level, or undefined if model doesn't support thinking + */ + cycleThinkingLevel(): ThinkingLevel | undefined { + if (!this.supportsThinking()) return undefined; + + const levels = this.getAvailableThinkingLevels(); + const currentIndex = levels.indexOf(this.thinkingLevel); + const nextIndex = (currentIndex + 1) % levels.length; + const nextLevel = levels[nextIndex]; + + this.setThinkingLevel(nextLevel); + return nextLevel; + } + + /** + * Get available thinking levels for current model. + * The provider will clamp to what the specific model supports internally. + */ + getAvailableThinkingLevels(): ThinkingLevel[] { + if (!this.model) return THINKING_LEVELS; + return getSupportedThinkingLevels(this.model) as ThinkingLevel[]; + } + + /** + * Check if current model supports thinking/reasoning. + */ + supportsThinking(): boolean { + return !!this.model?.reasoning; + } + + private _getThinkingLevelForModelSwitch(explicitLevel?: ThinkingLevel): ThinkingLevel { + if (explicitLevel !== undefined) { + return explicitLevel; + } + if (!this.supportsThinking()) { + return this.settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL; + } + return this.thinkingLevel; + } + + private _clampThinkingLevel(level: ThinkingLevel, _availableLevels: ThinkingLevel[]): ThinkingLevel { + return this.model ? (clampThinkingLevel(this.model, level) as ThinkingLevel) : "off"; + } + + // ========================================================================= + // Queue Mode Management + // ========================================================================= + + private syncQueueModesFromSettings(): void { + this.agent.steeringMode = this.settingsManager.getSteeringMode(); + this.agent.followUpMode = this.settingsManager.getFollowUpMode(); + } + + /** + * Set steering message mode. + * Saves to settings. + */ + setSteeringMode(mode: "all" | "one-at-a-time"): void { + this.agent.steeringMode = mode; + this.settingsManager.setSteeringMode(mode); + } + + /** + * Set follow-up message mode. + * Saves to settings. + */ + setFollowUpMode(mode: "all" | "one-at-a-time"): void { + this.agent.followUpMode = mode; + this.settingsManager.setFollowUpMode(mode); + } + + // ========================================================================= + // Compaction + // ========================================================================= + + /** + * Manually compact the session context. + * Aborts current agent operation first. + * @param customInstructions Optional instructions for the compaction summary + */ + async compact(customInstructions?: string): Promise { + this._disconnectFromAgent(); + await this.abort(); + this._compactionAbortController = new AbortController(); + this._emit({ type: "compaction_start", reason: "manual" }); + + try { + if (!this.model) { + throw new Error(formatNoModelSelectedMessage()); + } + + const { apiKey, headers, env } = await this._getCompactionRequestAuth(this.model); + + const pathEntries = this.sessionManager.getBranch(); + const settings = resolveCompactionSettings( + this.settingsManager.getCompactionSettings(), + this.model?.contextWindow ?? 0, + ); + + const preparation = prepareCompaction(pathEntries, settings); + if (!preparation) { + // Check why we can't compact + const lastEntry = pathEntries[pathEntries.length - 1]; + if (lastEntry?.type === "compaction") { + throw new Error("Already compacted"); + } + throw new Error("Nothing to compact (session too small)"); + } + + let extensionCompaction: CompactionResult | undefined; + let fromExtension = false; + + if (this._extensionRunner.hasHandlers("session_before_compact")) { + const result = (await this._extensionRunner.emit({ + type: "session_before_compact", + preparation, + branchEntries: pathEntries, + customInstructions, + reason: "manual", + willRetry: false, + signal: this._compactionAbortController.signal, + })) as SessionBeforeCompactResult | undefined; + + if (result?.cancel) { + throw new Error("Compaction cancelled"); + } + + if (result?.compaction) { + extensionCompaction = result.compaction; + fromExtension = true; + } + } + + let summary: string; + let firstKeptEntryId: string; + let tokensBefore: number; + let details: unknown; + + if (extensionCompaction) { + // Extension provided compaction content + summary = extensionCompaction.summary; + firstKeptEntryId = extensionCompaction.firstKeptEntryId; + tokensBefore = extensionCompaction.tokensBefore; + details = extensionCompaction.details; + } else { + // Generate compaction result + const result = await compact( + preparation, + this.model, + apiKey, + headers, + customInstructions, + this._compactionAbortController.signal, + this.thinkingLevel, + this.agent.streamFn, + env, + ); + summary = result.summary; + firstKeptEntryId = result.firstKeptEntryId; + tokensBefore = result.tokensBefore; + details = result.details; + } + + if (this._compactionAbortController.signal.aborted) { + throw new Error("Compaction cancelled"); + } + + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); + const newEntries = this.sessionManager.getEntries(); + const sessionContext = this.sessionManager.buildSessionContext(); + this.agent.state.messages = sessionContext.messages; + const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); + + // Get the saved compaction entry for the extension event + const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as + | CompactionEntry + | undefined; + + if (this._extensionRunner && savedCompactionEntry) { + await this._extensionRunner.emit({ + type: "session_compact", + compactionEntry: savedCompactionEntry, + fromExtension, + reason: "manual", + willRetry: false, + }); + } + + const compactionResult: CompactionResult = { + summary, + firstKeptEntryId, + tokensBefore, + estimatedTokensAfter, + details, + }; + this._emit({ + type: "compaction_end", + reason: "manual", + result: compactionResult, + aborted: false, + willRetry: false, + }); + return compactionResult; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const aborted = message === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError"); + this._emit({ + type: "compaction_end", + reason: "manual", + result: undefined, + aborted, + willRetry: false, + errorMessage: aborted ? undefined : `Compaction failed: ${message}`, + }); + throw error; + } finally { + this._compactionAbortController = undefined; + this._reconnectToAgent(); + } + } + + /** + * Cancel in-progress compaction (manual or auto). + */ + abortCompaction(): void { + this._compactionAbortController?.abort(); + this._autoCompactionAbortController?.abort(); + } + + /** + * Cancel in-progress branch summarization. + */ + abortBranchSummary(): void { + this._branchSummaryAbortController?.abort(); + } + + /** + * Check if compaction is needed and run it. + * Called after agent_end and before prompt submission. + * + * Two cases: + * 1. Overflow: LLM returned context overflow error, remove error message from agent state, compact, auto-retry + * 2. Threshold: Context over threshold, compact, NO auto-retry (user continues manually) + * + * @param assistantMessage The assistant message to check + * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true + */ + private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise { + const settings = resolveCompactionSettings( + this.settingsManager.getCompactionSettings(), + this.model?.contextWindow ?? 0, + ); + if (!settings.enabled) return false; + + // Skip if message was aborted (user cancelled) - unless skipAbortedCheck is false + if (skipAbortedCheck && assistantMessage.stopReason === "aborted") return false; + + const contextWindow = this.model?.contextWindow ?? 0; + + // Skip overflow check if the message came from a different model. + // This handles the case where user switched from a smaller-context model (e.g. opus) + // to a larger-context model (e.g. codex) - the overflow error from the old model + // shouldn't trigger compaction for the new model. + const sameModel = + this.model && assistantMessage.provider === this.model.provider && assistantMessage.model === this.model.id; + + // Skip compaction checks if this assistant message is older than the latest + // compaction boundary. This prevents a stale pre-compaction usage/error + // from retriggering compaction on the first prompt after compaction. + const compactionEntry = getLatestCompactionEntry(this.sessionManager.getBranch()); + const assistantIsFromBeforeCompaction = + compactionEntry !== null && assistantMessage.timestamp <= new Date(compactionEntry.timestamp).getTime(); + if (assistantIsFromBeforeCompaction) { + return false; + } + + // Case 1: Overflow - LLM returned context overflow error, or reported usage exceeded + // the configured window. A successful response over the configured window should compact + // but must not retry: the assistant answer already completed and agent.continue() cannot + // continue from an assistant message. + if (sameModel && isContextOverflow(assistantMessage, contextWindow)) { + const willRetry = assistantMessage.stopReason !== "stop"; + + if (!willRetry) { + return await this._runAutoCompaction("overflow", false); + } + + if (this._overflowRecoveryAttempted) { + this._emit({ + type: "compaction_end", + reason: "overflow", + result: undefined, + aborted: false, + willRetry: false, + errorMessage: + "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", + }); + return false; + } + + this._overflowRecoveryAttempted = true; + // Remove the error message from agent state (it IS saved to session for history, + // but we don't want it in context for the retry) + const messages = this.agent.state.messages; + if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { + this.agent.state.messages = messages.slice(0, -1); + } + return await this._runAutoCompaction("overflow", willRetry); + } + + // Case 2: Threshold - context is getting large + // For error messages or all-zero usage messages, estimate from the last valid response. + // This ensures sessions that hit persistent API errors (e.g. 529) or malformed zero-usage + // responses can still compact and do not reset context accounting. + let contextTokens: number; + const directContextTokens = assistantMessage.usage ? calculateContextTokens(assistantMessage.usage) : 0; + if (assistantMessage.stopReason === "error" || directContextTokens === 0) { + const messages = this.agent.state.messages; + const estimate = estimateContextTokens(messages); + if (estimate.lastUsageIndex === null) return false; // No usage data at all + // Verify the usage source is post-compaction. Kept pre-compaction messages + // have stale usage reflecting the old (larger) context and would falsely + // trigger compaction right after one just finished. + const usageMsg = messages[estimate.lastUsageIndex]; + if ( + compactionEntry && + usageMsg.role === "assistant" && + (usageMsg as AssistantMessage).timestamp <= new Date(compactionEntry.timestamp).getTime() + ) { + return false; + } + contextTokens = estimate.tokens; + } else { + contextTokens = directContextTokens; + } + if (shouldCompact(contextTokens, contextWindow, settings)) { + return await this._runAutoCompaction("threshold", false); + } + return false; + } + + /** + * Internal: Run auto-compaction with events. + */ + private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise { + const settings = resolveCompactionSettings( + this.settingsManager.getCompactionSettings(), + this.model?.contextWindow ?? 0, + ); + let started = false; + + try { + if (!this.model) { + return false; + } + + let apiKey: string | undefined; + let headers: Record | undefined; + let env: Record | undefined; + if (this.agent.streamFn === streamSimple) { + const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model); + if (!authResult.ok || !authResult.apiKey) { + return false; + } + apiKey = authResult.apiKey; + headers = authResult.headers; + env = authResult.env; + } else { + ({ apiKey, headers, env } = await this._getCompactionRequestAuth(this.model)); + } + + const pathEntries = this.sessionManager.getBranch(); + + const preparation = prepareCompaction(pathEntries, settings); + if (!preparation) { + return false; + } + + this._emit({ type: "compaction_start", reason }); + this._autoCompactionAbortController = new AbortController(); + started = true; + + let extensionCompaction: CompactionResult | undefined; + let fromExtension = false; + + if (this._extensionRunner.hasHandlers("session_before_compact")) { + const extensionResult = (await this._extensionRunner.emit({ + type: "session_before_compact", + preparation, + branchEntries: pathEntries, + customInstructions: undefined, + reason, + willRetry, + signal: this._autoCompactionAbortController.signal, + })) as SessionBeforeCompactResult | undefined; + + if (extensionResult?.cancel) { + this._emit({ + type: "compaction_end", + reason, + result: undefined, + aborted: true, + willRetry: false, + }); + return false; + } + + if (extensionResult?.compaction) { + extensionCompaction = extensionResult.compaction; + fromExtension = true; + } + } + + let summary: string; + let firstKeptEntryId: string; + let tokensBefore: number; + let details: unknown; + + if (extensionCompaction) { + // Extension provided compaction content + summary = extensionCompaction.summary; + firstKeptEntryId = extensionCompaction.firstKeptEntryId; + tokensBefore = extensionCompaction.tokensBefore; + details = extensionCompaction.details; + } else { + // Generate compaction result + const compactResult = await compact( + preparation, + this.model, + apiKey, + headers, + undefined, + this._autoCompactionAbortController.signal, + this.thinkingLevel, + this.agent.streamFn, + env, + ); + summary = compactResult.summary; + firstKeptEntryId = compactResult.firstKeptEntryId; + tokensBefore = compactResult.tokensBefore; + details = compactResult.details; + } + + if (this._autoCompactionAbortController.signal.aborted) { + this._emit({ + type: "compaction_end", + reason, + result: undefined, + aborted: true, + willRetry: false, + }); + return false; + } + + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); + const newEntries = this.sessionManager.getEntries(); + const sessionContext = this.sessionManager.buildSessionContext(); + this.agent.state.messages = sessionContext.messages; + const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); + + // Get the saved compaction entry for the extension event + const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as + | CompactionEntry + | undefined; + + if (this._extensionRunner && savedCompactionEntry) { + await this._extensionRunner.emit({ + type: "session_compact", + compactionEntry: savedCompactionEntry, + fromExtension, + reason, + willRetry, + }); + } + + const result: CompactionResult = { + summary, + firstKeptEntryId, + tokensBefore, + estimatedTokensAfter, + details, + }; + this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); + + if (willRetry) { + const messages = this.agent.state.messages; + const lastMsg = messages[messages.length - 1]; + if (lastMsg?.role === "assistant" && (lastMsg as AssistantMessage).stopReason === "error") { + this.agent.state.messages = messages.slice(0, -1); + } + return true; + } + + // Auto-compaction can complete while follow-up/steering/custom messages are waiting. + // Continue once so queued messages are delivered. + return this.agent.hasQueuedMessages(); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : "compaction failed"; + if (started) { + this._emit({ + type: "compaction_end", + reason, + result: undefined, + aborted: false, + willRetry: false, + errorMessage: + reason === "overflow" + ? `Context overflow recovery failed: ${errorMessage}` + : `Auto-compaction failed: ${errorMessage}`, + }); + } + return false; + } finally { + this._autoCompactionAbortController = undefined; + } + } + + /** + * Toggle auto-compaction setting. + */ + setAutoCompactionEnabled(enabled: boolean): void { + this.settingsManager.setCompactionEnabled(enabled); + } + + /** Whether auto-compaction is enabled */ + get autoCompactionEnabled(): boolean { + return this.settingsManager.getCompactionEnabled(); + } + + async bindExtensions(bindings: ExtensionBindings): Promise { + if (bindings.uiContext !== undefined) { + this._extensionUIContext = bindings.uiContext; + } + if (bindings.mode !== undefined) { + this._extensionMode = bindings.mode; + } + if (bindings.commandContextActions !== undefined) { + this._extensionCommandContextActions = bindings.commandContextActions; + } + if (bindings.abortHandler !== undefined) { + this._extensionAbortHandler = bindings.abortHandler; + } + if (bindings.shutdownHandler !== undefined) { + this._extensionShutdownHandler = bindings.shutdownHandler; + } + if (bindings.onError !== undefined) { + this._extensionErrorListener = bindings.onError; + } + + this._applyExtensionBindings(this._extensionRunner); + await this._extensionRunner.emit(this._sessionStartEvent); + await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup"); + } + + private async extendResourcesFromExtensions(reason: "startup" | "reload"): Promise { + if (!this._extensionRunner.hasHandlers("resources_discover")) { + return; + } + + const { skillPaths, promptPaths, themePaths } = await this._extensionRunner.emitResourcesDiscover( + this._cwd, + reason, + ); + + if (skillPaths.length === 0 && promptPaths.length === 0 && themePaths.length === 0) { + return; + } + + const extensionPaths: ResourceExtensionPaths = { + skillPaths: this.buildExtensionResourcePaths(skillPaths), + promptPaths: this.buildExtensionResourcePaths(promptPaths), + themePaths: this.buildExtensionResourcePaths(themePaths), + }; + + this._resourceLoader.extendResources(extensionPaths); + this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); + this.agent.state.systemPrompt = this._baseSystemPrompt; + } + + private buildExtensionResourcePaths(entries: Array<{ path: string; extensionPath: string }>): Array<{ + path: string; + metadata: { source: string; scope: "temporary"; origin: "top-level"; baseDir?: string }; + }> { + return entries.map((entry) => { + const source = this.getExtensionSourceLabel(entry.extensionPath); + const baseDir = entry.extensionPath.startsWith("<") ? undefined : dirname(entry.extensionPath); + return { + path: entry.path, + metadata: { + source, + scope: "temporary", + origin: "top-level", + baseDir, + }, + }; + }); + } + + private getExtensionSourceLabel(extensionPath: string): string { + if (extensionPath.startsWith("<")) { + return `extension:${extensionPath.replace(/[<>]/g, "")}`; + } + const base = basename(extensionPath); + const name = base.replace(/\.(ts|js)$/, ""); + return `extension:${name}`; + } + + private _applyExtensionBindings(runner: ExtensionRunner): void { + runner.setUIContext(this._extensionUIContext, this._extensionMode); + runner.bindCommandContext(this._extensionCommandContextActions); + + this._extensionErrorUnsubscriber?.(); + this._extensionErrorUnsubscriber = this._extensionErrorListener + ? runner.onError(this._extensionErrorListener) + : undefined; + } + + private _refreshCurrentModelFromRegistry(): void { + const currentModel = this.model; + if (!currentModel) { + return; + } + + const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id); + if (!refreshedModel || refreshedModel === currentModel) { + return; + } + + this.agent.state.model = refreshedModel; + } + + private _bindExtensionCore(runner: ExtensionRunner): void { + const getCommands = (): SlashCommandInfo[] => { + const extensionCommands: SlashCommandInfo[] = runner.getRegisteredCommands().map((command) => ({ + name: command.invocationName, + description: command.description, + source: "extension", + sourceInfo: command.sourceInfo, + })); + + const templates: SlashCommandInfo[] = this.promptTemplates.map((template) => ({ + name: template.name, + description: template.description, + source: "prompt", + sourceInfo: template.sourceInfo, + })); + + const skills: SlashCommandInfo[] = this._resourceLoader.getSkills().skills.map((skill) => ({ + name: `skill:${skill.name}`, + description: skill.description, + source: "skill", + sourceInfo: skill.sourceInfo, + })); + + return [...extensionCommands, ...templates, ...skills]; + }; + + runner.bindCore( + { + sendMessage: (message, options) => { + this.sendCustomMessage(message, options).catch((err) => { + runner.emitError({ + extensionPath: "", + event: "send_message", + error: err instanceof Error ? err.message : String(err), + }); + }); + }, + sendUserMessage: (content, options) => { + this.sendUserMessage(content, options).catch((err) => { + runner.emitError({ + extensionPath: "", + event: "send_user_message", + error: err instanceof Error ? err.message : String(err), + }); + }); + }, + appendEntry: (customType, data) => { + this.sessionManager.appendCustomEntry(customType, data); + }, + setSessionName: (name) => { + this.setSessionName(name); + }, + getSessionName: () => { + return this.sessionManager.getSessionName(); + }, + setLabel: (entryId, label) => { + this.sessionManager.appendLabelChange(entryId, label); + }, + getActiveTools: () => this.getActiveToolNames(), + getAllTools: () => this.getAllTools(), + setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames), + refreshTools: () => this._refreshToolRegistry(), + getCommands, + setModel: async (model) => { + if (!this.modelRegistry.hasConfiguredAuth(model)) return false; + await this.setModel(model); + return true; + }, + getThinkingLevel: () => this.thinkingLevel, + setThinkingLevel: (level) => this.setThinkingLevel(level), + }, + { + getModel: () => this.model, + isIdle: () => !this.isStreaming, + isProjectTrusted: () => this.settingsManager.isProjectTrusted(), + getSignal: () => this.agent.signal, + abort: () => { + if (this._extensionAbortHandler) { + this._extensionAbortHandler(); + return; + } + void this.abort(); + }, + hasPendingMessages: () => this.pendingMessageCount > 0, + shutdown: () => { + this._extensionShutdownHandler?.(); + }, + getContextUsage: () => this.getContextUsage(), + compact: (options) => { + void (async () => { + try { + const result = await this.compact(options?.customInstructions); + options?.onComplete?.(result); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + options?.onError?.(err); + } + })(); + }, + getSystemPrompt: () => this.systemPrompt, + getSystemPromptOptions: () => this._baseSystemPromptOptions, + }, + { + registerProvider: (name, config) => { + this._modelRegistry.registerProvider(name, config); + this._refreshCurrentModelFromRegistry(); + }, + unregisterProvider: (name) => { + this._modelRegistry.unregisterProvider(name); + this._refreshCurrentModelFromRegistry(); + }, + }, + ); + } + + private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void { + const previousRegistryNames = new Set(this._toolRegistry.keys()); + const previousActiveToolNames = this.getActiveToolNames(); + const allowedToolNames = this._allowedToolNames; + const excludedToolNames = this._excludedToolNames; + const isAllowedTool = (name: string): boolean => + (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name); + + const registeredTools = this._extensionRunner.getAllRegisteredTools(); + const allCustomTools = [ + ...registeredTools, + ...this._customTools.map((definition) => ({ + definition, + sourceInfo: createSyntheticSourceInfo(``, { source: "sdk" }), + })), + ].filter((tool) => isAllowedTool(tool.definition.name)); + const definitionRegistry = new Map( + Array.from(this._baseToolDefinitions.entries()) + .filter(([name]) => isAllowedTool(name)) + .map(([name, definition]) => [ + name, + { + definition, + sourceInfo: createSyntheticSourceInfo(``, { source: "builtin" }), + }, + ]), + ); + for (const tool of allCustomTools) { + definitionRegistry.set(tool.definition.name, { + definition: tool.definition, + sourceInfo: tool.sourceInfo, + }); + } + this._toolDefinitions = definitionRegistry; + this._toolPromptSnippets = new Map( + Array.from(definitionRegistry.values()) + .map(({ definition }) => { + const snippet = this._normalizePromptSnippet(definition.promptSnippet); + return snippet ? ([definition.name, snippet] as const) : undefined; + }) + .filter((entry): entry is readonly [string, string] => entry !== undefined), + ); + this._toolPromptGuidelines = new Map( + Array.from(definitionRegistry.values()) + .map(({ definition }) => { + const guidelines = this._normalizePromptGuidelines(definition.promptGuidelines); + return guidelines.length > 0 ? ([definition.name, guidelines] as const) : undefined; + }) + .filter((entry): entry is readonly [string, string[]] => entry !== undefined), + ); + const runner = this._extensionRunner; + const wrappedExtensionTools = wrapRegisteredTools(allCustomTools, runner); + const wrappedBuiltInTools = wrapRegisteredTools( + Array.from(this._baseToolDefinitions.values()) + .filter((definition) => isAllowedTool(definition.name)) + .map((definition) => ({ + definition, + sourceInfo: createSyntheticSourceInfo(``, { source: "builtin" }), + })), + runner, + ); + + const toolRegistry = new Map(wrappedBuiltInTools.map((tool) => [tool.name, tool])); + for (const tool of wrappedExtensionTools as AgentTool[]) { + toolRegistry.set(tool.name, tool); + } + this._toolRegistry = toolRegistry; + + const nextActiveToolNames = ( + options?.activeToolNames ? [...options.activeToolNames] : [...previousActiveToolNames] + ).filter((name) => isAllowedTool(name)); + + if (allowedToolNames) { + for (const toolName of this._toolRegistry.keys()) { + if (allowedToolNames.has(toolName)) { + nextActiveToolNames.push(toolName); + } + } + } else if (options?.includeAllExtensionTools) { + for (const tool of wrappedExtensionTools) { + nextActiveToolNames.push(tool.name); + } + } else if (!options?.activeToolNames) { + for (const toolName of this._toolRegistry.keys()) { + if (!previousRegistryNames.has(toolName)) { + nextActiveToolNames.push(toolName); + } + } + } + + this.setActiveToolsByName([...new Set(nextActiveToolNames)]); + } + + private _buildRuntime(options: { + activeToolNames?: string[]; + flagValues?: Map; + includeAllExtensionTools?: boolean; + }): void { + const autoResizeImages = this.settingsManager.getImageAutoResize(); + const shellCommandPrefix = this.settingsManager.getShellCommandPrefix(); + const shellPath = this.settingsManager.getShellPath(); + const baseToolDefinitions = this._baseToolsOverride + ? Object.fromEntries( + Object.entries(this._baseToolsOverride).map(([name, tool]) => [ + name, + createToolDefinitionFromAgentTool(tool), + ]), + ) + : createAllToolDefinitions(this._cwd, { + read: { autoResizeImages }, + bash: { commandPrefix: shellCommandPrefix, shellPath }, + }); + + this._baseToolDefinitions = new Map( + Object.entries(baseToolDefinitions).map(([name, tool]) => [name, tool as ToolDefinition]), + ); + + const extensionsResult = this._resourceLoader.getExtensions(); + if (options.flagValues) { + for (const [name, value] of options.flagValues) { + extensionsResult.runtime.flagValues.set(name, value); + } + } + + this._extensionRunner = new ExtensionRunner( + extensionsResult.extensions, + extensionsResult.runtime, + this._cwd, + this.sessionManager, + this._modelRegistry, + ); + if (this._extensionRunnerRef) { + this._extensionRunnerRef.current = this._extensionRunner; + } + this._bindExtensionCore(this._extensionRunner); + this._applyExtensionBindings(this._extensionRunner); + + const defaultActiveToolNames = this._baseToolsOverride + ? Object.keys(this._baseToolsOverride) + : ["read", "bash", "edit", "write"]; + const baseActiveToolNames = options.activeToolNames ?? defaultActiveToolNames; + this._refreshToolRegistry({ + activeToolNames: baseActiveToolNames, + includeAllExtensionTools: options.includeAllExtensionTools, + }); + } + + async reload(options?: { beforeSessionStart?: () => void | Promise }): Promise { + const previousFlagValues = this._extensionRunner.getFlagValues(); + await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" }); + await this.settingsManager.reload(); + this.syncQueueModesFromSettings(); + resetApiProviders(); + await this._resourceLoader.reload(); + this._buildRuntime({ + activeToolNames: this.getActiveToolNames(), + flagValues: previousFlagValues, + includeAllExtensionTools: true, + }); + + const hasBindings = + this._extensionUIContext || + this._extensionCommandContextActions || + this._extensionShutdownHandler || + this._extensionErrorListener; + if (hasBindings) { + await options?.beforeSessionStart?.(); + await this._extensionRunner.emit({ type: "session_start", reason: "reload" }); + await this.extendResourcesFromExtensions("reload"); + } + } + + // ========================================================================= + // Auto-Retry + // ========================================================================= + + private _isNonRetryableProviderLimitError(errorMessage: string): boolean { + return /GoUsageLimitError|FreeUsageLimitError|Monthly usage limit reached|available balance|insufficient_quota|out of budget|quota exceeded|billing/i.test( + errorMessage, + ); + } + + /** + * Check if an error is retryable (overloaded, rate limit, server errors). + * Context overflow errors are NOT retryable (handled by compaction instead). + */ + private _isRetryableError(message: AssistantMessage): boolean { + if (message.stopReason !== "error" || !message.errorMessage) return false; + + // Context overflow is handled by compaction, not retry + const contextWindow = this.model?.contextWindow ?? 0; + if (isContextOverflow(message, contextWindow)) return false; + + const err = message.errorMessage; + if (this._isNonRetryableProviderLimitError(err)) return false; + // Match: overloaded_error, provider returned error, rate limit, 429, 500, 502, 503, 504, service unavailable, network/connection errors (including connection lost), WebSocket transport closes/errors, fetch failed, premature stream endings, HTTP/2 closed before response, terminated, retry delay exceeded + return /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|connection.?lost|websocket.?closed|websocket.?error|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|ended without|stream ended before message_stop|http2 request did not get a response|timed? out|timeout|terminated|retry delay/i.test( + err, + ); + } + + /** + * Prepare a retryable error for continuation with exponential backoff. + * @returns true if the caller should continue the agent, false otherwise + */ + private async _prepareRetry(message: AssistantMessage): Promise { + const settings = this.settingsManager.getRetrySettings(); + if (!settings.enabled) { + return false; + } + + this._retryAttempt++; + + if (this._retryAttempt > settings.maxRetries) { + // Preserve the completed attempt count so post-run handling can emit the final failure. + this._retryAttempt--; + return false; + } + + const delayMs = settings.baseDelayMs * 2 ** (this._retryAttempt - 1); + + this._emit({ + type: "auto_retry_start", + attempt: this._retryAttempt, + maxAttempts: settings.maxRetries, + delayMs, + errorMessage: message.errorMessage || "Unknown error", + }); + + // Remove error message from agent state (keep in session for history) + const messages = this.agent.state.messages; + if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { + this.agent.state.messages = messages.slice(0, -1); + } + + // Wait with exponential backoff (abortable) + this._retryAbortController = new AbortController(); + try { + await sleep(delayMs, this._retryAbortController.signal); + } catch { + // Aborted during sleep - emit end event so UI can clean up + const attempt = this._retryAttempt; + this._retryAttempt = 0; + this._emit({ + type: "auto_retry_end", + success: false, + attempt, + finalError: "Retry cancelled", + }); + return false; + } finally { + this._retryAbortController = undefined; + } + + return true; + } + + /** + * Cancel in-progress retry. + */ + abortRetry(): void { + this._retryAbortController?.abort(); + } + + /** Whether auto-retry is currently in progress */ + get isRetrying(): boolean { + return this._retryAbortController !== undefined; + } + + /** Whether auto-retry is enabled */ + get autoRetryEnabled(): boolean { + return this.settingsManager.getRetryEnabled(); + } + + /** + * Toggle auto-retry setting. + */ + setAutoRetryEnabled(enabled: boolean): void { + this.settingsManager.setRetryEnabled(enabled); + } + + // ========================================================================= + // Bash Execution + // ========================================================================= + + /** + * Execute a bash command. + * Adds result to agent context and session. + * @param command The bash command to execute + * @param onChunk Optional streaming callback for output + * @param options.excludeFromContext If true, command output won't be sent to LLM (!! prefix) + * @param options.operations Custom BashOperations for remote execution + */ + async executeBash( + command: string, + onChunk?: (chunk: string) => void, + options?: { excludeFromContext?: boolean; operations?: BashOperations }, + ): Promise { + this._bashAbortController = new AbortController(); + + // Apply command prefix if configured (e.g., "shopt -s expand_aliases" for alias support) + const prefix = this.settingsManager.getShellCommandPrefix(); + const shellPath = this.settingsManager.getShellPath(); + const resolvedCommand = prefix ? `${prefix}\n${command}` : command; + + try { + const result = await executeBashWithOperations( + resolvedCommand, + this.sessionManager.getCwd(), + options?.operations ?? createLocalBashOperations({ shellPath }), + { + onChunk, + signal: this._bashAbortController.signal, + }, + ); + + this.recordBashResult(command, result, options); + return result; + } finally { + this._bashAbortController = undefined; + } + } + + /** + * Record a bash execution result in session history. + * Used by executeBash and by extensions that handle bash execution themselves. + */ + recordBashResult(command: string, result: BashResult, options?: { excludeFromContext?: boolean }): void { + const bashMessage: BashExecutionMessage = { + role: "bashExecution", + command, + output: result.output, + exitCode: result.exitCode, + cancelled: result.cancelled, + truncated: result.truncated, + fullOutputPath: result.fullOutputPath, + timestamp: Date.now(), + excludeFromContext: options?.excludeFromContext, + }; + + // If agent is streaming, defer adding to avoid breaking tool_use/tool_result ordering + if (this.isStreaming) { + // Queue for later - will be flushed on agent_end + this._pendingBashMessages.push(bashMessage); + } else { + // Add to agent state immediately + this.agent.state.messages.push(bashMessage); + + // Save to session + this.sessionManager.appendMessage(bashMessage); + } + } + + /** + * Cancel running bash command. + */ + abortBash(): void { + this._bashAbortController?.abort(); + } + + /** Whether a bash command is currently running */ + get isBashRunning(): boolean { + return this._bashAbortController !== undefined; + } + + /** Whether there are pending bash messages waiting to be flushed */ + get hasPendingBashMessages(): boolean { + return this._pendingBashMessages.length > 0; + } + + /** + * Flush pending bash messages to agent state and session. + * Called after agent turn completes to maintain proper message ordering. + */ + private _flushPendingBashMessages(): void { + if (this._pendingBashMessages.length === 0) return; + + for (const bashMessage of this._pendingBashMessages) { + // Add to agent state + this.agent.state.messages.push(bashMessage); + + // Save to session + this.sessionManager.appendMessage(bashMessage); + } + + this._pendingBashMessages = []; + } + + // ========================================================================= + // Session Management + // ========================================================================= + + /** + * Set a display name for the current session. + */ + setSessionName(name: string): void { + this.sessionManager.appendSessionInfo(name); + this._emit({ type: "session_info_changed", name: this.sessionManager.getSessionName() }); + } + + // ========================================================================= + // Tree Navigation + // ========================================================================= + + /** + * Navigate to a different node in the session tree. + * Unlike fork() which creates a new session file, this stays in the same file. + * + * @param targetId The entry ID to navigate to + * @param options.summarize Whether user wants to summarize abandoned branch + * @param options.customInstructions Custom instructions for summarizer + * @param options.replaceInstructions If true, customInstructions replaces the default prompt + * @param options.label Label to attach to the branch summary entry + * @returns Result with editorText (if user message) and cancelled status + */ + async navigateTree( + targetId: string, + options: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string } = {}, + ): Promise<{ editorText?: string; cancelled: boolean; aborted?: boolean; summaryEntry?: BranchSummaryEntry }> { + const oldLeafId = this.sessionManager.getLeafId(); + + // No-op if already at target + if (targetId === oldLeafId) { + return { cancelled: false }; + } + + // Model required for summarization + if (options.summarize && !this.model) { + throw new Error("No model available for summarization"); + } + + const targetEntry = this.sessionManager.getEntry(targetId); + if (!targetEntry) { + throw new Error(`Entry ${targetId} not found`); + } + + // Collect entries to summarize (from old leaf to common ancestor) + const { entries: entriesToSummarize, commonAncestorId } = collectEntriesForBranchSummary( + this.sessionManager, + oldLeafId, + targetId, + ); + + // Prepare event data - mutable so extensions can override + let customInstructions = options.customInstructions; + let replaceInstructions = options.replaceInstructions; + let label = options.label; + + const preparation: TreePreparation = { + targetId, + oldLeafId, + commonAncestorId, + entriesToSummarize, + userWantsSummary: options.summarize ?? false, + customInstructions, + replaceInstructions, + label, + }; + + // Set up abort controller for summarization + this._branchSummaryAbortController = new AbortController(); + + try { + let extensionSummary: { summary: string; details?: unknown } | undefined; + let fromExtension = false; + + // Emit session_before_tree event + if (this._extensionRunner.hasHandlers("session_before_tree")) { + const result = (await this._extensionRunner.emit({ + type: "session_before_tree", + preparation, + signal: this._branchSummaryAbortController.signal, + })) as SessionBeforeTreeResult | undefined; + + if (result?.cancel) { + return { cancelled: true }; + } + + if (result?.summary && options.summarize) { + extensionSummary = result.summary; + fromExtension = true; + } + + // Allow extensions to override instructions and label + if (result?.customInstructions !== undefined) { + customInstructions = result.customInstructions; + } + if (result?.replaceInstructions !== undefined) { + replaceInstructions = result.replaceInstructions; + } + if (result?.label !== undefined) { + label = result.label; + } + } + + // Run default summarizer if needed + let summaryText: string | undefined; + let summaryDetails: unknown; + if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { + const model = this.model!; + const { apiKey, headers, env } = await this._getRequiredRequestAuth(model); + const branchSummarySettings = this.settingsManager.getBranchSummarySettings(); + const result = await generateBranchSummary(entriesToSummarize, { + model, + apiKey, + headers, + env, + signal: this._branchSummaryAbortController.signal, + customInstructions, + replaceInstructions, + reserveTokens: branchSummarySettings.reserveTokens, + streamFn: this.agent.streamFn, + }); + if (result.aborted) { + return { cancelled: true, aborted: true }; + } + if (result.error) { + throw new Error(result.error); + } + summaryText = result.summary; + summaryDetails = { + readFiles: result.readFiles || [], + modifiedFiles: result.modifiedFiles || [], + }; + } else if (extensionSummary) { + summaryText = extensionSummary.summary; + summaryDetails = extensionSummary.details; + } + + // Determine the new leaf position based on target type + let newLeafId: string | null; + let editorText: string | undefined; + + if (targetEntry.type === "message" && targetEntry.message.role === "user") { + // User message: leaf = parent (null if root), text goes to editor + newLeafId = targetEntry.parentId; + editorText = this._extractUserMessageText(targetEntry.message.content); + } else if (targetEntry.type === "custom_message") { + // Custom message: leaf = parent (null if root), text goes to editor + newLeafId = targetEntry.parentId; + editorText = + typeof targetEntry.content === "string" + ? targetEntry.content + : targetEntry.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } else { + // Non-user message: leaf = selected node + newLeafId = targetId; + } + + // Switch leaf (with or without summary) + // Summary is attached at the navigation target position (newLeafId), not the old branch + let summaryEntry: BranchSummaryEntry | undefined; + if (summaryText) { + // Create summary at target position (can be null for root) + const summaryId = this.sessionManager.branchWithSummary( + newLeafId, + summaryText, + summaryDetails, + fromExtension, + ); + summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; + + // Attach label to the summary entry + if (label) { + this.sessionManager.appendLabelChange(summaryId, label); + } + } else if (newLeafId === null) { + // No summary, navigating to root - reset leaf + this.sessionManager.resetLeaf(); + } else { + // No summary, navigating to non-root + this.sessionManager.branch(newLeafId); + } + + // Attach label to target entry when not summarizing (no summary entry to label) + if (label && !summaryText) { + this.sessionManager.appendLabelChange(targetId, label); + } + + // Update agent state + const sessionContext = this.sessionManager.buildSessionContext(); + this.agent.state.messages = sessionContext.messages; + + // Emit session_tree event + await this._extensionRunner.emit({ + type: "session_tree", + newLeafId: this.sessionManager.getLeafId(), + oldLeafId, + summaryEntry, + fromExtension: summaryText ? fromExtension : undefined, + }); + + // Emit to custom tools + + return { editorText, cancelled: false, summaryEntry }; + } finally { + this._branchSummaryAbortController = undefined; + } + } + + /** + * Get all user messages from session for fork selector. + */ + getUserMessagesForForking(): Array<{ entryId: string; text: string }> { + const entries = this.sessionManager.getEntries(); + const result: Array<{ entryId: string; text: string }> = []; + + for (const entry of entries) { + if (entry.type !== "message") continue; + if (entry.message.role !== "user") continue; + + const text = this._extractUserMessageText(entry.message.content); + if (text) { + result.push({ entryId: entry.id, text }); + } + } + + return result; + } + + private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { + if (typeof content === "string") return content; + if (Array.isArray(content)) { + return content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } + return ""; + } + + /** + * Get session statistics. + */ + getSessionStats(): SessionStats { + const state = this.state; + const userMessages = state.messages.filter((m) => m.role === "user").length; + const assistantMessages = state.messages.filter((m) => m.role === "assistant").length; + const toolResults = state.messages.filter((m) => m.role === "toolResult").length; + + let toolCalls = 0; + let totalInput = 0; + let totalOutput = 0; + let totalCacheRead = 0; + let totalCacheWrite = 0; + let totalCost = 0; + + for (const message of state.messages) { + if (message.role === "assistant") { + const assistantMsg = message as AssistantMessage; + toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; + totalInput += assistantMsg.usage.input; + totalOutput += assistantMsg.usage.output; + totalCacheRead += assistantMsg.usage.cacheRead; + totalCacheWrite += assistantMsg.usage.cacheWrite; + totalCost += assistantMsg.usage.cost.total; + } + } + + return { + sessionFile: this.sessionFile, + sessionId: this.sessionId, + userMessages, + assistantMessages, + toolCalls, + toolResults, + totalMessages: state.messages.length, + tokens: { + input: totalInput, + output: totalOutput, + cacheRead: totalCacheRead, + cacheWrite: totalCacheWrite, + total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, + }, + cost: totalCost, + contextUsage: this.getContextUsage(), + }; + } + + getContextUsage(): ContextUsage | undefined { + const model = this.model; + if (!model) return undefined; + + const contextWindow = model.contextWindow ?? 0; + if (contextWindow <= 0) return undefined; + + // After compaction, the last assistant usage reflects pre-compaction context size. + // We can only trust usage from an assistant that responded after the latest compaction. + // If no such assistant exists, context token count is unknown until the next LLM response. + const branchEntries = this.sessionManager.getBranch(); + const latestCompaction = getLatestCompactionEntry(branchEntries); + + if (latestCompaction) { + // Check if there's a valid assistant usage after the compaction boundary + const compactionIndex = branchEntries.lastIndexOf(latestCompaction); + let hasPostCompactionUsage = false; + for (let i = branchEntries.length - 1; i > compactionIndex; i--) { + const entry = branchEntries[i]; + if (entry.type === "message" && entry.message.role === "assistant") { + const assistant = entry.message; + if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") { + const contextTokens = calculateContextTokens(assistant.usage); + if (contextTokens > 0) { + hasPostCompactionUsage = true; + break; + } + } + } + } + + if (!hasPostCompactionUsage) { + return { tokens: null, contextWindow, percent: null }; + } + } + + const estimate = estimateContextTokens(this.messages); + const percent = (estimate.tokens / contextWindow) * 100; + + return { + tokens: estimate.tokens, + contextWindow, + percent, + }; + } + + /** + * Export session to HTML. + * @param outputPath Optional output path (defaults to session directory) + * @returns Path to exported file + */ + /** + * Export the current session branch to a JSONL file. + * Writes the session header followed by all entries on the current branch path. + * @param outputPath Target file path. If omitted, generates a timestamped file in cwd. + * @returns The resolved output file path. + */ + exportToJsonl(outputPath?: string): string { + const filePath = resolvePath( + outputPath ?? `session-${new Date().toISOString().replace(/[:.]/g, "-")}.jsonl`, + process.cwd(), + ); + const dir = dirname(filePath); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + + const header: SessionHeader = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: this.sessionManager.getSessionId(), + timestamp: new Date().toISOString(), + cwd: this.sessionManager.getCwd(), + }; + + const branchEntries = this.sessionManager.getBranch(); + const lines = [JSON.stringify(header)]; + + // Re-chain parentIds to form a linear sequence + let prevId: string | null = null; + for (const entry of branchEntries) { + const linear = { ...entry, parentId: prevId }; + lines.push(JSON.stringify(linear)); + prevId = entry.id; + } + + writeFileSync(filePath, `${lines.join("\n")}\n`); + return filePath; + } + + // ========================================================================= + // Utilities + // ========================================================================= + + /** + * Get text content of last assistant message. + * Useful for /copy command. + * @returns Text content, or undefined if no assistant message exists + */ + getLastAssistantText(): string | undefined { + const lastAssistant = this.messages + .slice() + .reverse() + .find((m) => { + if (m.role !== "assistant") return false; + const msg = m as AssistantMessage; + // Skip aborted messages with no content + if (msg.stopReason === "aborted" && msg.content.length === 0) return false; + return true; + }); + + if (!lastAssistant) return undefined; + + let text = ""; + for (const content of (lastAssistant as AssistantMessage).content) { + if (content.type === "text") { + text += content.text; + } + } + + return text.trim() || undefined; + } + + // ========================================================================= + // Extension System + // ========================================================================= + + createReplacedSessionContext(): ReplacedSessionContext { + const context = Object.defineProperties( + {}, + Object.getOwnPropertyDescriptors(this._extensionRunner.createCommandContext()), + ) as ReplacedSessionContext; + context.sendMessage = (message, options) => this.sendCustomMessage(message, options); + context.sendUserMessage = (content, options) => this.sendUserMessage(content, options); + return context; + } + + /** + * Check if extensions have handlers for a specific event type. + */ + hasExtensionHandlers(eventType: string): boolean { + return this._extensionRunner.hasHandlers(eventType); + } + + /** + * Get the extension runner (for setting UI context and error handlers). + */ + get extensionRunner(): ExtensionRunner { + return this._extensionRunner; + } +} diff --git a/cactus-code/packages/coding-agent/src/core/auth-guidance.ts b/cactus-code/packages/coding-agent/src/core/auth-guidance.ts new file mode 100644 index 000000000..9782cdaff --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/auth-guidance.ts @@ -0,0 +1,25 @@ +import { join } from "node:path"; +import { getDocsPath } from "../config.ts"; + +const UNKNOWN_PROVIDER = "unknown"; + +export function getProviderLoginHelp(): string { + return [ + "Use /login to log into a provider via OAuth or API key. See:", + ` ${join(getDocsPath(), "providers.md")}`, + ` ${join(getDocsPath(), "models.md")}`, + ].join("\n"); +} + +export function formatNoModelsAvailableMessage(): string { + return `No models available. ${getProviderLoginHelp()}`; +} + +export function formatNoModelSelectedMessage(): string { + return `No model selected.\n\n${getProviderLoginHelp()}\n\nThen use /model to select a model.`; +} + +export function formatNoApiKeyFoundMessage(provider: string): string { + const providerDisplay = provider === UNKNOWN_PROVIDER ? "the selected model" : provider; + return `No API key found for ${providerDisplay}.\n\n${getProviderLoginHelp()}`; +} diff --git a/cactus-code/packages/coding-agent/src/core/auth-storage.ts b/cactus-code/packages/coding-agent/src/core/auth-storage.ts new file mode 100644 index 000000000..31a0ec9c3 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/auth-storage.ts @@ -0,0 +1,528 @@ +/** + * Credential storage for API keys and OAuth tokens. + * Handles loading, saving, and refreshing credentials from auth.json. + * + * Uses file locking to prevent race conditions when multiple pi instances + * try to refresh tokens simultaneously. + */ + +import { + findEnvKeys, + getEnvApiKey, + type OAuthCredentials, + type OAuthLoginCallbacks, + type OAuthProviderId, +} from "@earendil-works/pi-ai/compat"; +import { getOAuthApiKey, getOAuthProvider, getOAuthProviders } from "@earendil-works/pi-ai/oauth"; +import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; +import { dirname, join } from "path"; +import lockfile from "proper-lockfile"; +import { getAgentDir } from "../config.ts"; +import { normalizePath } from "../utils/paths.ts"; +import { resolveConfigValue } from "./resolve-config-value.ts"; + +export type ApiKeyCredential = { + type: "api_key"; + key: string; + env?: Record; +}; + +export type OAuthCredential = { + type: "oauth"; +} & OAuthCredentials; + +export type AuthCredential = ApiKeyCredential | OAuthCredential; + +export type AuthStorageData = Record; + +export type AuthStatus = { + configured: boolean; + source?: "stored" | "runtime" | "environment" | "fallback" | "models_json_key" | "models_json_command"; + label?: string; +}; + +export interface GetApiKeyOptions { + includeFallback?: boolean; +} + +type LockResult = { + result: T; + next?: string; +}; + +const AUTH_FILE_WRITE_OPTIONS = { encoding: "utf-8", mode: 0o600 } as const; + +export interface AuthStorageBackend { + withLock(fn: (current: string | undefined) => LockResult): T; + withLockAsync(fn: (current: string | undefined) => Promise>): Promise; +} + +export class FileAuthStorageBackend implements AuthStorageBackend { + private authPath: string; + + constructor(authPath: string = join(getAgentDir(), "auth.json")) { + this.authPath = normalizePath(authPath); + } + + private ensureParentDir(): void { + const dir = dirname(this.authPath); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true, mode: 0o700 }); + } + } + + private ensureFileExists(): void { + if (!existsSync(this.authPath)) { + writeFileSync(this.authPath, "{}", AUTH_FILE_WRITE_OPTIONS); + chmodSync(this.authPath, 0o600); + } + } + + private acquireLockSyncWithRetry(path: string): () => void { + const maxAttempts = 10; + const delayMs = 20; + let lastError: unknown; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { + try { + return lockfile.lockSync(path, { realpath: false }); + } catch (error) { + const code = + typeof error === "object" && error !== null && "code" in error + ? String((error as { code?: unknown }).code) + : undefined; + if (code !== "ELOCKED" || attempt === maxAttempts) { + throw error; + } + lastError = error; + const start = Date.now(); + while (Date.now() - start < delayMs) { + // Sleep synchronously to avoid changing callers to async. + } + } + } + + throw (lastError as Error) ?? new Error("Failed to acquire auth storage lock"); + } + + withLock(fn: (current: string | undefined) => LockResult): T { + this.ensureParentDir(); + this.ensureFileExists(); + + let release: (() => void) | undefined; + try { + release = this.acquireLockSyncWithRetry(this.authPath); + const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined; + const { result, next } = fn(current); + if (next !== undefined) { + writeFileSync(this.authPath, next, AUTH_FILE_WRITE_OPTIONS); + chmodSync(this.authPath, 0o600); + } + return result; + } finally { + if (release) { + release(); + } + } + } + + async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { + this.ensureParentDir(); + this.ensureFileExists(); + + let release: (() => Promise) | undefined; + let lockCompromised = false; + let lockCompromisedError: Error | undefined; + const throwIfCompromised = () => { + if (lockCompromised) { + throw lockCompromisedError ?? new Error("Auth storage lock was compromised"); + } + }; + + try { + release = await lockfile.lock(this.authPath, { + retries: { + retries: 10, + factor: 2, + minTimeout: 100, + maxTimeout: 10000, + randomize: true, + }, + stale: 30000, + onCompromised: (err) => { + lockCompromised = true; + lockCompromisedError = err; + }, + }); + + throwIfCompromised(); + const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined; + const { result, next } = await fn(current); + throwIfCompromised(); + if (next !== undefined) { + writeFileSync(this.authPath, next, AUTH_FILE_WRITE_OPTIONS); + chmodSync(this.authPath, 0o600); + } + throwIfCompromised(); + return result; + } finally { + if (release) { + try { + await release(); + } catch { + // Ignore unlock errors when lock is compromised. + } + } + } + } +} + +export class InMemoryAuthStorageBackend implements AuthStorageBackend { + private value: string | undefined; + + withLock(fn: (current: string | undefined) => LockResult): T { + const { result, next } = fn(this.value); + if (next !== undefined) { + this.value = next; + } + return result; + } + + async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { + const { result, next } = await fn(this.value); + if (next !== undefined) { + this.value = next; + } + return result; + } +} + +/** + * Credential storage backed by a JSON file. + */ +export class AuthStorage { + private data: AuthStorageData = {}; + private runtimeOverrides: Map = new Map(); + private loadError: Error | null = null; + private errors: Error[] = []; + private storage: AuthStorageBackend; + + private constructor(storage: AuthStorageBackend) { + this.storage = storage; + this.reload(); + } + + static create(authPath?: string): AuthStorage { + return new AuthStorage(new FileAuthStorageBackend(authPath ?? join(getAgentDir(), "auth.json"))); + } + + static fromStorage(storage: AuthStorageBackend): AuthStorage { + return new AuthStorage(storage); + } + + static inMemory(data: AuthStorageData = {}): AuthStorage { + const storage = new InMemoryAuthStorageBackend(); + storage.withLock(() => ({ result: undefined, next: JSON.stringify(data, null, 2) })); + return AuthStorage.fromStorage(storage); + } + + /** + * Set a runtime API key override (not persisted to disk). + * Used for CLI --api-key flag. + */ + setRuntimeApiKey(provider: string, apiKey: string): void { + this.runtimeOverrides.set(provider, apiKey); + } + + /** + * Remove a runtime API key override. + */ + removeRuntimeApiKey(provider: string): void { + this.runtimeOverrides.delete(provider); + } + + private recordError(error: unknown): void { + const normalizedError = error instanceof Error ? error : new Error(String(error)); + this.errors.push(normalizedError); + } + + private parseStorageData(content: string | undefined): AuthStorageData { + if (!content) { + return {}; + } + return JSON.parse(content) as AuthStorageData; + } + + /** + * Reload credentials from storage. + */ + reload(): void { + let content: string | undefined; + try { + this.storage.withLock((current) => { + content = current; + return { result: undefined }; + }); + this.data = this.parseStorageData(content); + this.loadError = null; + } catch (error) { + this.loadError = error as Error; + this.recordError(error); + } + } + + private persistProviderChange(provider: string, credential: AuthCredential | undefined): void { + if (this.loadError) { + return; + } + + try { + this.storage.withLock((current) => { + const currentData = this.parseStorageData(current); + const merged: AuthStorageData = { ...currentData }; + if (credential) { + merged[provider] = credential; + } else { + delete merged[provider]; + } + return { result: undefined, next: JSON.stringify(merged, null, 2) }; + }); + } catch (error) { + this.recordError(error); + } + } + + /** + * Get credential for a provider. + */ + get(provider: string): AuthCredential | undefined { + return this.data[provider] ?? undefined; + } + + /** + * Get provider-scoped environment values for an API key credential. + */ + getProviderEnv(provider: string): Record | undefined { + const cred = this.data[provider]; + return cred?.type === "api_key" && cred.env ? { ...cred.env } : undefined; + } + + /** + * Set credential for a provider. + */ + set(provider: string, credential: AuthCredential): void { + this.data[provider] = credential; + this.persistProviderChange(provider, credential); + } + + /** + * Remove credential for a provider. + */ + remove(provider: string): void { + delete this.data[provider]; + this.persistProviderChange(provider, undefined); + } + + /** + * List all providers with credentials. + */ + list(): string[] { + return Object.keys(this.data); + } + + /** + * Check if credentials exist for a provider in auth.json. + */ + has(provider: string): boolean { + return provider in this.data; + } + + /** + * Check if any form of auth is configured for a provider. + * Unlike getApiKey(), this doesn't refresh OAuth tokens. + */ + hasAuth(provider: string): boolean { + if (this.runtimeOverrides.has(provider)) return true; + if (this.data[provider]) return true; + if (getEnvApiKey(provider)) return true; + return false; + } + + /** + * Return auth status without exposing credential values or refreshing tokens. + */ + getAuthStatus(provider: string): AuthStatus { + if (this.data[provider]) { + return { configured: true, source: "stored" }; + } + + if (this.runtimeOverrides.has(provider)) { + return { configured: false, source: "runtime", label: "--api-key" }; + } + + const envKeys = findEnvKeys(provider); + if (envKeys?.[0]) { + return { configured: false, source: "environment", label: envKeys[0] }; + } + + return { configured: false }; + } + + /** + * Get all credentials (for passing to getOAuthApiKey). + */ + getAll(): AuthStorageData { + return { ...this.data }; + } + + drainErrors(): Error[] { + const drained = [...this.errors]; + this.errors = []; + return drained; + } + + /** + * Login to an OAuth provider. + */ + async login(providerId: OAuthProviderId, callbacks: OAuthLoginCallbacks): Promise { + const provider = getOAuthProvider(providerId); + if (!provider) { + throw new Error(`Unknown OAuth provider: ${providerId}`); + } + + const credentials = await provider.login(callbacks); + this.set(providerId, { type: "oauth", ...credentials }); + } + + /** + * Logout from a provider. + */ + logout(provider: string): void { + this.remove(provider); + } + + /** + * Refresh OAuth token with backend locking to prevent race conditions. + * Multiple pi instances may try to refresh simultaneously when tokens expire. + */ + private async refreshOAuthTokenWithLock( + providerId: OAuthProviderId, + ): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> { + const provider = getOAuthProvider(providerId); + if (!provider) { + return null; + } + + const result = await this.storage.withLockAsync(async (current) => { + const currentData = this.parseStorageData(current); + this.data = currentData; + this.loadError = null; + + const cred = currentData[providerId]; + if (cred?.type !== "oauth") { + return { result: null }; + } + + if (Date.now() < cred.expires) { + return { result: { apiKey: provider.getApiKey(cred), newCredentials: cred } }; + } + + const oauthCreds: Record = {}; + for (const [key, value] of Object.entries(currentData)) { + if (value.type === "oauth") { + oauthCreds[key] = value; + } + } + + const refreshed = await getOAuthApiKey(providerId, oauthCreds); + if (!refreshed) { + return { result: null }; + } + + const merged: AuthStorageData = { + ...currentData, + [providerId]: { type: "oauth", ...refreshed.newCredentials }, + }; + this.data = merged; + this.loadError = null; + return { result: refreshed, next: JSON.stringify(merged, null, 2) }; + }); + + return result; + } + + /** + * Get API key for a provider. + * Priority: + * 1. Runtime override (CLI --api-key) + * 2. API key from auth.json + * 3. OAuth token from auth.json (auto-refreshed with locking) + * 4. Environment variable + */ + async getApiKey(providerId: string, options: GetApiKeyOptions = {}): Promise { + // Runtime override takes highest priority + const runtimeKey = this.runtimeOverrides.get(providerId); + if (runtimeKey) { + return runtimeKey; + } + + const cred = this.data[providerId]; + + if (cred?.type === "api_key") { + return resolveConfigValue(cred.key, cred.env); + } + + if (cred?.type === "oauth") { + const provider = getOAuthProvider(providerId); + if (!provider) { + // Unknown OAuth provider, can't get API key + return undefined; + } + + // Check if token needs refresh + const needsRefresh = Date.now() >= cred.expires; + + if (needsRefresh) { + // Use locked refresh to prevent race conditions + try { + const result = await this.refreshOAuthTokenWithLock(providerId); + if (result) { + return result.apiKey; + } + } catch (error) { + this.recordError(error); + // Refresh failed - re-read file to check if another instance succeeded + this.reload(); + const updatedCred = this.data[providerId]; + + if (updatedCred?.type === "oauth" && Date.now() < updatedCred.expires) { + // Another instance refreshed successfully, use those credentials + return provider.getApiKey(updatedCred); + } + + // Refresh truly failed - return undefined so model discovery skips this provider + // User can /login to re-authenticate (credentials preserved for retry) + return undefined; + } + } else { + // Token not expired, use current access token + return provider.getApiKey(cred); + } + } + + if (options.includeFallback === false) return undefined; + + // Fall back to environment variable + const envKey = getEnvApiKey(providerId); + if (envKey) return envKey; + + return undefined; + } + + /** + * Get all registered OAuth providers + */ + getOAuthProviders() { + return getOAuthProviders(); + } +} diff --git a/cactus-code/packages/coding-agent/src/core/bash-executor.ts b/cactus-code/packages/coding-agent/src/core/bash-executor.ts new file mode 100644 index 000000000..875a1556d --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/bash-executor.ts @@ -0,0 +1,156 @@ +/** + * Bash command execution with streaming support and cancellation. + * + * This module provides a unified bash execution implementation used by: + * - AgentSession.executeBash() for interactive and RPC modes + * - Direct calls from modes that need bash execution + */ + +import { randomBytes } from "node:crypto"; +import { createWriteStream, type WriteStream } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { stripAnsi } from "../utils/ansi.ts"; +import { sanitizeBinaryOutput } from "../utils/shell.ts"; +import type { BashOperations } from "./tools/bash.ts"; +import { DEFAULT_MAX_BYTES, truncateTail } from "./tools/truncate.ts"; + +// ============================================================================ +// Types +// ============================================================================ + +export interface BashExecutorOptions { + /** Callback for streaming output chunks (already sanitized) */ + onChunk?: (chunk: string) => void; + /** AbortSignal for cancellation */ + signal?: AbortSignal; +} + +export interface BashResult { + /** Combined stdout + stderr output (sanitized, possibly truncated) */ + output: string; + /** Process exit code (undefined if killed/cancelled) */ + exitCode: number | undefined; + /** Whether the command was cancelled via signal */ + cancelled: boolean; + /** Whether the output was truncated */ + truncated: boolean; + /** Path to temp file containing full output (if output exceeded truncation threshold) */ + fullOutputPath?: string; +} + +// ============================================================================ +// Implementation +// ============================================================================ + +/** + * Execute a bash command using custom BashOperations. + * Used for remote execution (SSH, containers, etc.). + */ +export async function executeBashWithOperations( + command: string, + cwd: string, + operations: BashOperations, + options?: BashExecutorOptions, +): Promise { + const outputChunks: string[] = []; + let outputBytes = 0; + const maxOutputBytes = DEFAULT_MAX_BYTES * 2; + + let tempFilePath: string | undefined; + let tempFileStream: WriteStream | undefined; + let totalBytes = 0; + + const ensureTempFile = () => { + if (tempFilePath) { + return; + } + const id = randomBytes(8).toString("hex"); + tempFilePath = join(tmpdir(), `cactus-bash-${id}.log`); + tempFileStream = createWriteStream(tempFilePath); + for (const chunk of outputChunks) { + tempFileStream.write(chunk); + } + }; + + const decoder = new TextDecoder(); + + const onData = (data: Buffer) => { + totalBytes += data.length; + + // Sanitize: strip ANSI, replace binary garbage, normalize newlines + const text = sanitizeBinaryOutput(stripAnsi(decoder.decode(data, { stream: true }))).replace(/\r/g, ""); + + // Start writing to temp file if exceeds threshold + if (totalBytes > DEFAULT_MAX_BYTES) { + ensureTempFile(); + } + + if (tempFileStream) { + tempFileStream.write(text); + } + + // Keep rolling buffer + outputChunks.push(text); + outputBytes += text.length; + while (outputBytes > maxOutputBytes && outputChunks.length > 1) { + const removed = outputChunks.shift()!; + outputBytes -= removed.length; + } + + // Stream to callback + if (options?.onChunk) { + options.onChunk(text); + } + }; + + try { + const result = await operations.exec(command, cwd, { + onData, + signal: options?.signal, + }); + + const fullOutput = outputChunks.join(""); + const truncationResult = truncateTail(fullOutput); + if (truncationResult.truncated) { + ensureTempFile(); + } + if (tempFileStream) { + tempFileStream.end(); + } + const cancelled = options?.signal?.aborted ?? false; + + return { + output: truncationResult.truncated ? truncationResult.content : fullOutput, + exitCode: cancelled ? undefined : (result.exitCode ?? undefined), + cancelled, + truncated: truncationResult.truncated, + fullOutputPath: tempFilePath, + }; + } catch (err) { + // Check if it was an abort + if (options?.signal?.aborted) { + const fullOutput = outputChunks.join(""); + const truncationResult = truncateTail(fullOutput); + if (truncationResult.truncated) { + ensureTempFile(); + } + if (tempFileStream) { + tempFileStream.end(); + } + return { + output: truncationResult.truncated ? truncationResult.content : fullOutput, + exitCode: undefined, + cancelled: true, + truncated: truncationResult.truncated, + fullOutputPath: tempFilePath, + }; + } + + if (tempFileStream) { + tempFileStream.end(); + } + + throw err; + } +} diff --git a/cactus-code/packages/coding-agent/src/core/compaction/branch-summarization.ts b/cactus-code/packages/coding-agent/src/core/compaction/branch-summarization.ts new file mode 100644 index 000000000..3f557c01b --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -0,0 +1,371 @@ +/** + * Branch summarization for tree navigation. + * + * When navigating to a different point in the session tree, this generates + * a summary of the branch being left so context isn't lost. + */ + +import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; +import type { Model, SimpleStreamOptions } from "@earendil-works/pi-ai/compat"; +import { completeSimple } from "@earendil-works/pi-ai/compat"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.ts"; +import type { ReadonlySessionManager, SessionEntry } from "../session-manager.ts"; +import { estimateTokens } from "./compaction.ts"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + SUMMARIZATION_SYSTEM_PROMPT, + serializeConversation, +} from "./utils.ts"; + +// ============================================================================ +// Types +// ============================================================================ + +export interface BranchSummaryResult { + summary?: string; + readFiles?: string[]; + modifiedFiles?: string[]; + aborted?: boolean; + error?: string; +} + +/** Details stored in BranchSummaryEntry.details for file tracking */ +export interface BranchSummaryDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +export type { FileOperations } from "./utils.ts"; + +export interface BranchPreparation { + /** Messages extracted for summarization, in chronological order */ + messages: AgentMessage[]; + /** File operations extracted from tool calls */ + fileOps: FileOperations; + /** Total estimated tokens in messages */ + totalTokens: number; +} + +export interface CollectEntriesResult { + /** Entries to summarize, in chronological order */ + entries: SessionEntry[]; + /** Common ancestor between old and new position, if any */ + commonAncestorId: string | null; +} + +export interface GenerateBranchSummaryOptions { + /** Model to use for summarization */ + model: Model; + /** API key for the model */ + apiKey: string; + /** Request headers for the model */ + headers?: Record; + /** Provider-scoped environment values for the model */ + env?: Record; + /** Abort signal for cancellation */ + signal: AbortSignal; + /** Optional custom instructions for summarization */ + customInstructions?: string; + /** If true, customInstructions replaces the default prompt instead of being appended */ + replaceInstructions?: boolean; + /** Tokens reserved for prompt + LLM response (default 16384) */ + reserveTokens?: number; + /** Optional session stream function. Used to preserve SDK request behavior without mutating agent state. */ + streamFn?: StreamFn; +} + +// ============================================================================ +// Entry Collection +// ============================================================================ + +/** + * Collect entries that should be summarized when navigating from one position to another. + * + * Walks from oldLeafId back to the common ancestor with targetId, collecting entries + * along the way. Does NOT stop at compaction boundaries - those are included and their + * summaries become context. + * + * @param session - Session manager (read-only access) + * @param oldLeafId - Current position (where we're navigating from) + * @param targetId - Target position (where we're navigating to) + * @returns Entries to summarize and the common ancestor + */ +export function collectEntriesForBranchSummary( + session: ReadonlySessionManager, + oldLeafId: string | null, + targetId: string, +): CollectEntriesResult { + // If no old position, nothing to summarize + if (!oldLeafId) { + return { entries: [], commonAncestorId: null }; + } + + // Find common ancestor (deepest node that's on both paths) + const oldPath = new Set(session.getBranch(oldLeafId).map((e) => e.id)); + const targetPath = session.getBranch(targetId); + + // targetPath is root-first, so iterate backwards to find deepest common ancestor + let commonAncestorId: string | null = null; + for (let i = targetPath.length - 1; i >= 0; i--) { + if (oldPath.has(targetPath[i].id)) { + commonAncestorId = targetPath[i].id; + break; + } + } + + // Collect entries from old leaf back to common ancestor + const entries: SessionEntry[] = []; + let current: string | null = oldLeafId; + + while (current && current !== commonAncestorId) { + const entry = session.getEntry(current); + if (!entry) break; + entries.push(entry); + current = entry.parentId; + } + + // Reverse to get chronological order + entries.reverse(); + + return { entries, commonAncestorId }; +} + +// ============================================================================ +// Entry to Message Conversion +// ============================================================================ + +/** + * Extract AgentMessage from a session entry. + * Similar to getMessageFromEntry in compaction.ts but also handles compaction entries. + */ +function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { + switch (entry.type) { + case "message": + // Skip tool results - context is in assistant's tool call + if (entry.message.role === "toolResult") return undefined; + return entry.message; + + case "custom_message": + return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + + case "branch_summary": + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); + + case "compaction": + return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); + + // These don't contribute to conversation content + case "thinking_level_change": + case "model_change": + case "custom": + case "label": + case "session_info": + return undefined; + } +} + +/** + * Prepare entries for summarization with token budget. + * + * Walks entries from NEWEST to OLDEST, adding messages until we hit the token budget. + * This ensures we keep the most recent context when the branch is too long. + * + * Also collects file operations from: + * - Tool calls in assistant messages + * - Existing branch_summary entries' details (for cumulative tracking) + * + * @param entries - Entries in chronological order + * @param tokenBudget - Maximum tokens to include (0 = no limit) + */ +export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: number = 0): BranchPreparation { + const messages: AgentMessage[] = []; + const fileOps = createFileOps(); + let totalTokens = 0; + + // First pass: collect file ops from ALL entries (even if they don't fit in token budget) + // This ensures we capture cumulative file tracking from nested branch summaries + // Only extract from pi-generated summaries (fromHook !== true), not extension-generated ones + for (const entry of entries) { + if (entry.type === "branch_summary" && !entry.fromHook && entry.details) { + const details = entry.details as BranchSummaryDetails; + if (Array.isArray(details.readFiles)) { + for (const f of details.readFiles) fileOps.read.add(f); + } + if (Array.isArray(details.modifiedFiles)) { + // Modified files go into both edited and written for proper deduplication + for (const f of details.modifiedFiles) { + fileOps.edited.add(f); + } + } + } + } + + // Second pass: walk from newest to oldest, adding messages until token budget + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + const message = getMessageFromEntry(entry); + if (!message) continue; + + // Extract file ops from assistant messages (tool calls) + extractFileOpsFromMessage(message, fileOps); + + const tokens = estimateTokens(message); + + // Check budget before adding + if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { + // If this is a summary entry, try to fit it anyway as it's important context + if (entry.type === "compaction" || entry.type === "branch_summary") { + if (totalTokens < tokenBudget * 0.9) { + messages.unshift(message); + totalTokens += tokens; + } + } + // Stop - we've hit the budget + break; + } + + messages.unshift(message); + totalTokens += tokens; + } + + return { messages, fileOps, totalTokens }; +} + +// ============================================================================ +// Summary Generation +// ============================================================================ + +const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here. +Summary of that exploration: + +`; + +const BRANCH_SUMMARY_PROMPT = `Create a structured summary of this conversation branch for context when returning later. + +Use this EXACT format: + +## Goal +[What was the user trying to accomplish in this branch?] + +## Constraints & Preferences +- [Any constraints, preferences, or requirements mentioned] +- [Or "(none)" if none were mentioned] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Work that was started but not finished] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [What should happen next to continue this work] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +/** + * Generate a summary of abandoned branch entries. + * + * @param entries - Session entries to summarize (chronological order) + * @param options - Generation options + */ +export async function generateBranchSummary( + entries: SessionEntry[], + options: GenerateBranchSummaryOptions, +): Promise { + const { + model, + apiKey, + headers, + env, + signal, + customInstructions, + replaceInstructions, + reserveTokens = 16384, + streamFn, + } = options; + + // Token budget = context window minus reserved space for prompt + response + const contextWindow = model.contextWindow || 128000; + const tokenBudget = contextWindow - reserveTokens; + + const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget); + + if (messages.length === 0) { + return { summary: "No content to summarize" }; + } + + // Transform to LLM-compatible messages, then serialize to text + // Serialization prevents the model from treating it as a conversation to continue + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + + // Build prompt + let instructions: string; + if (replaceInstructions && customInstructions) { + instructions = customInstructions; + } else if (customInstructions) { + instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`; + } else { + instructions = BRANCH_SUMMARY_PROMPT; + } + const promptText = `\n${conversationText}\n\n\n${instructions}`; + + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + // Call LLM for summarization. Prefer the session stream function so SDK + // request behavior (timeouts, retries, attribution headers) stays consistent + // without running through agent state/events. + const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }; + const requestOptions: SimpleStreamOptions = { apiKey, headers, env, signal, maxTokens: 2048 }; + const response = streamFn + ? await (await streamFn(model, context, requestOptions)).result() + : await completeSimple(model, context, requestOptions); + + // Check if aborted or errored + if (response.stopReason === "aborted") { + return { aborted: true }; + } + if (response.stopReason === "error") { + return { error: response.errorMessage || "Summarization failed" }; + } + + let summary = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + + // Prepend preamble to provide context about the branch summary + summary = BRANCH_SUMMARY_PREAMBLE + summary; + + // Compute file lists and append to summary + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + return { + summary: summary || "No summary generated", + readFiles, + modifiedFiles, + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/compaction/compaction.ts b/cactus-code/packages/coding-agent/src/core/compaction/compaction.ts new file mode 100644 index 000000000..8d1877183 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/compaction/compaction.ts @@ -0,0 +1,905 @@ +/** + * Context compaction for long sessions. + * + * Pure functions for compaction logic. The session manager handles I/O, + * and after compaction the session is reloaded. + */ + +import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; +import { completeSimple } from "@earendil-works/pi-ai/compat"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.ts"; +import { buildSessionContext, type CompactionEntry, type SessionEntry } from "../session-manager.ts"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + SUMMARIZATION_SYSTEM_PROMPT, + serializeConversation, +} from "./utils.ts"; + +// ============================================================================ +// File Operation Tracking +// ============================================================================ + +/** Details stored in CompactionEntry.details for file tracking */ +export interface CompactionDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +/** + * Extract file operations from messages and previous compaction entries. + */ +function extractFileOperations( + messages: AgentMessage[], + entries: SessionEntry[], + prevCompactionIndex: number, +): FileOperations { + const fileOps = createFileOps(); + + // Collect from previous compaction's details (if pi-generated) + if (prevCompactionIndex >= 0) { + const prevCompaction = entries[prevCompactionIndex] as CompactionEntry; + if (!prevCompaction.fromHook && prevCompaction.details) { + // fromHook field kept for session file compatibility + const details = prevCompaction.details as CompactionDetails; + if (Array.isArray(details.readFiles)) { + for (const f of details.readFiles) fileOps.read.add(f); + } + if (Array.isArray(details.modifiedFiles)) { + for (const f of details.modifiedFiles) fileOps.edited.add(f); + } + } + } + + // Extract from tool calls in messages + for (const msg of messages) { + extractFileOpsFromMessage(msg, fileOps); + } + + return fileOps; +} + +// ============================================================================ +// Message Extraction +// ============================================================================ + +/** + * Extract AgentMessage from an entry if it produces one. + * Returns undefined for entries that don't contribute to LLM context. + */ +function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { + if (entry.type === "message") { + return entry.message; + } + if (entry.type === "custom_message") { + return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + } + if (entry.type === "branch_summary") { + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); + } + if (entry.type === "compaction") { + return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); + } + return undefined; +} + +function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | undefined { + if (entry.type === "compaction") { + return undefined; + } + return getMessageFromEntry(entry); +} + +/** Result from compact() - SessionManager adds uuid/parentUuid when saving */ +export interface CompactionResult { + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + estimatedTokensAfter?: number; + /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ + details?: T; +} + +// ============================================================================ +// Types +// ============================================================================ + +export interface CompactionSettings { + enabled: boolean; + reserveTokens: number; + keepRecentTokens: number; +} + +export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { + enabled: true, + reserveTokens: 16384, + keepRecentTokens: 20000, +}; + +// ============================================================================ +// Token calculation +// ============================================================================ + +/** + * Calculate total context tokens from usage. + * Uses the native totalTokens field when available, falls back to computing from components. + */ +export function calculateContextTokens(usage: Usage): number { + return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; +} + +/** + * Get usage from an assistant message if available. + * Skips aborted, error, and all-zero usage messages as they don't have valid usage data. + */ +function getAssistantUsage(msg: AgentMessage): Usage | undefined { + if (msg.role === "assistant" && "usage" in msg) { + const assistantMsg = msg as AssistantMessage; + if ( + assistantMsg.stopReason !== "aborted" && + assistantMsg.stopReason !== "error" && + assistantMsg.usage && + calculateContextTokens(assistantMsg.usage) > 0 + ) { + return assistantMsg.usage; + } + } + return undefined; +} + +/** + * Find the last valid assistant message usage from session entries. + */ +export function getLastAssistantUsage(entries: SessionEntry[]): Usage | undefined { + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry.type === "message") { + const usage = getAssistantUsage(entry.message); + if (usage) return usage; + } + } + return undefined; +} + +export interface ContextUsageEstimate { + tokens: number; + usageTokens: number; + trailingTokens: number; + lastUsageIndex: number | null; +} + +function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const usage = getAssistantUsage(messages[i]); + if (usage) return { usage, index: i }; + } + return undefined; +} + +/** + * Estimate context tokens from messages, using the last assistant usage when available. + * If there are messages after the last usage, estimate their tokens with estimateTokens. + */ +export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate { + const usageInfo = getLastAssistantUsageInfo(messages); + + if (!usageInfo) { + let estimated = 0; + for (const message of messages) { + estimated += estimateTokens(message); + } + return { + tokens: estimated, + usageTokens: 0, + trailingTokens: estimated, + lastUsageIndex: null, + }; + } + + const usageTokens = calculateContextTokens(usageInfo.usage); + let trailingTokens = 0; + for (let i = usageInfo.index + 1; i < messages.length; i++) { + trailingTokens += estimateTokens(messages[i]); + } + + return { + tokens: usageTokens + trailingTokens, + usageTokens, + trailingTokens, + lastUsageIndex: usageInfo.index, + }; +} + +export function resolveCompactionSettings(settings: CompactionSettings, contextWindow: number): CompactionSettings { + if (!Number.isFinite(contextWindow) || contextWindow <= 0) { + return { ...settings, enabled: false }; + } + const reserveTokens = Math.max(256, Math.min(settings.reserveTokens, Math.floor(contextWindow * 0.3))); + const keepRecentTokens = Math.max(256, Math.min(settings.keepRecentTokens, Math.floor(contextWindow * 0.4))); + if (reserveTokens >= contextWindow) { + return { ...settings, enabled: false }; + } + return { ...settings, reserveTokens, keepRecentTokens }; +} + +/** + * Check if compaction should trigger based on context usage. + */ +export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean { + if (!settings.enabled) return false; + return contextTokens > contextWindow - settings.reserveTokens; +} + +// ============================================================================ +// Cut point detection +// ============================================================================ + +const ESTIMATED_IMAGE_CHARS = 4800; + +function estimateTextAndImageContentChars(content: string | Array<{ type: string; text?: string }>): number { + if (typeof content === "string") { + return content.length; + } + + let chars = 0; + for (const block of content) { + if (block.type === "text" && block.text) { + chars += block.text.length; + } else if (block.type === "image") { + chars += ESTIMATED_IMAGE_CHARS; + } + } + return chars; +} + +/** + * Estimate token count for a message using chars/4 heuristic. + * This is conservative (overestimates tokens). + */ +export function estimateTokens(message: AgentMessage): number { + let chars = 0; + + switch (message.role) { + case "user": { + chars = estimateTextAndImageContentChars( + (message as { content: string | Array<{ type: string; text?: string }> }).content, + ); + return Math.ceil(chars / 4); + } + case "assistant": { + const assistant = message as AssistantMessage; + for (const block of assistant.content) { + if (block.type === "text") { + chars += block.text.length; + } else if (block.type === "thinking") { + chars += block.thinking.length; + } else if (block.type === "toolCall") { + chars += block.name.length + JSON.stringify(block.arguments).length; + } + } + return Math.ceil(chars / 4); + } + case "custom": + case "toolResult": { + chars = estimateTextAndImageContentChars(message.content); + return Math.ceil(chars / 4); + } + case "bashExecution": { + chars = message.command.length + message.output.length; + return Math.ceil(chars / 4); + } + case "branchSummary": + case "compactionSummary": { + chars = message.summary.length; + return Math.ceil(chars / 4); + } + } + + return 0; +} + +/** + * Find valid cut points: indices of user, assistant, custom, or bashExecution messages. + * Never cut at tool results (they must follow their tool call). + * When we cut at an assistant message with tool calls, its tool results follow it + * and will be kept. + * BashExecutionMessage is treated like a user message (user-initiated context). + */ +function findValidCutPoints(entries: SessionEntry[], startIndex: number, endIndex: number): number[] { + const cutPoints: number[] = []; + for (let i = startIndex; i < endIndex; i++) { + const entry = entries[i]; + switch (entry.type) { + case "message": { + const role = entry.message.role; + switch (role) { + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + case "user": + case "assistant": + cutPoints.push(i); + break; + case "toolResult": + break; + } + break; + } + case "thinking_level_change": + case "model_change": + case "compaction": + case "branch_summary": + case "custom": + case "custom_message": + case "label": + case "session_info": + break; + } + + // branch_summary and custom_message are user-role messages, valid cut points + if (entry.type === "branch_summary" || entry.type === "custom_message") { + cutPoints.push(i); + } + } + return cutPoints; +} + +/** + * Find the user message (or bashExecution) that starts the turn containing the given entry index. + * Returns -1 if no turn start found before the index. + * BashExecutionMessage is treated like a user message for turn boundaries. + */ +export function findTurnStartIndex(entries: SessionEntry[], entryIndex: number, startIndex: number): number { + for (let i = entryIndex; i >= startIndex; i--) { + const entry = entries[i]; + // branch_summary and custom_message are user-role messages, can start a turn + if (entry.type === "branch_summary" || entry.type === "custom_message") { + return i; + } + if (entry.type === "message") { + const role = entry.message.role; + if (role === "user" || role === "bashExecution") { + return i; + } + } + } + return -1; +} + +export interface CutPointResult { + /** Index of first entry to keep */ + firstKeptEntryIndex: number; + /** Index of user message that starts the turn being split, or -1 if not splitting */ + turnStartIndex: number; + /** Whether this cut splits a turn (cut point is not a user message) */ + isSplitTurn: boolean; +} + +/** + * Find the cut point in session entries that keeps approximately `keepRecentTokens`. + * + * Algorithm: Walk backwards from newest, accumulating estimated message sizes. + * Stop when we've accumulated >= keepRecentTokens. Cut at that point. + * + * Can cut at user OR assistant messages (never tool results). When cutting at an + * assistant message with tool calls, its tool results come after and will be kept. + * + * Returns CutPointResult with: + * - firstKeptEntryIndex: the entry index to start keeping from + * - turnStartIndex: if cutting mid-turn, the user message that started that turn + * - isSplitTurn: whether we're cutting in the middle of a turn + * + * Only considers entries between `startIndex` and `endIndex` (exclusive). + */ +export function findCutPoint( + entries: SessionEntry[], + startIndex: number, + endIndex: number, + keepRecentTokens: number, +): CutPointResult { + const cutPoints = findValidCutPoints(entries, startIndex, endIndex); + + if (cutPoints.length === 0) { + return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false }; + } + + // Walk backwards from newest, accumulating estimated message sizes + let accumulatedTokens = 0; + let cutIndex = cutPoints[0]; // Default: keep from first message (not header) + + for (let i = endIndex - 1; i >= startIndex; i--) { + const entry = entries[i]; + if (entry.type !== "message") continue; + + // Estimate this message's size + const messageTokens = estimateTokens(entry.message); + accumulatedTokens += messageTokens; + + // Check if we've exceeded the budget + if (accumulatedTokens >= keepRecentTokens) { + // Find the closest valid cut point at or after this entry + for (let c = 0; c < cutPoints.length; c++) { + if (cutPoints[c] >= i) { + cutIndex = cutPoints[c]; + break; + } + } + break; + } + } + + // Scan backwards from cutIndex to include any non-message entries (bash, settings, etc.) + while (cutIndex > startIndex) { + const prevEntry = entries[cutIndex - 1]; + // Stop at session header or compaction boundaries + if (prevEntry.type === "compaction") { + break; + } + if (prevEntry.type === "message") { + // Stop if we hit any message + break; + } + // Include this non-message entry (bash, settings change, etc.) + cutIndex--; + } + + // Determine if this is a split turn + const cutEntry = entries[cutIndex]; + const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user"; + const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); + + return { + firstKeptEntryIndex: cutIndex, + turnStartIndex, + isSplitTurn: !isUserMessage && turnStartIndex !== -1, + }; +} + +// ============================================================================ +// Summarization +// ============================================================================ + +const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work. + +Use this EXACT format: + +## Goal +[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.] + +## Constraints & Preferences +- [Any constraints, preferences, or requirements mentioned by user] +- [Or "(none)" if none were mentioned] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Current work] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [Ordered list of what should happen next] + +## Critical Context +- [Any data, examples, or references needed to continue] +- [Or "(none)" if not applicable] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in tags. + +Update the existing structured summary with new information. RULES: +- PRESERVE all existing information from the previous summary +- ADD new progress, decisions, and context from the new messages +- UPDATE the Progress section: move items from "In Progress" to "Done" when completed +- UPDATE "Next Steps" based on what was accomplished +- PRESERVE exact file paths, function names, and error messages +- If something is no longer relevant, you may remove it + +Use this EXACT format: + +## Goal +[Preserve existing goals, add new ones if the task expanded] + +## Constraints & Preferences +- [Preserve existing, add new ones discovered] + +## Progress +### Done +- [x] [Include previously done items AND newly completed items] + +### In Progress +- [ ] [Current work - update based on progress] + +### Blocked +- [Current blockers - remove if resolved] + +## Key Decisions +- **[Decision]**: [Brief rationale] (preserve all previous, add new) + +## Next Steps +1. [Update based on current state] + +## Critical Context +- [Preserve important context, add new if needed] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +function createSummarizationOptions( + model: Model, + maxTokens: number, + apiKey: string | undefined, + headers: Record | undefined, + env: Record | undefined, + signal: AbortSignal | undefined, + thinkingLevel: ThinkingLevel | undefined, +): SimpleStreamOptions { + const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers, env }; + if (model.reasoning && thinkingLevel && thinkingLevel !== "off") { + options.reasoning = thinkingLevel; + } + return options; +} + +async function completeSummarization( + model: Model, + context: Context, + options: SimpleStreamOptions, + streamFn?: StreamFn, +): Promise { + if (!streamFn) { + return completeSimple(model, context, options); + } + const stream = await streamFn(model, context, options); + return stream.result(); +} + +/** + * Generate a summary of the conversation using the LLM. + * If previousSummary is provided, uses the update prompt to merge. + */ +export async function generateSummary( + currentMessages: AgentMessage[], + model: Model, + reserveTokens: number, + apiKey: string | undefined, + headers?: Record, + signal?: AbortSignal, + customInstructions?: string, + previousSummary?: string, + thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, + env?: Record, +): Promise { + const maxTokens = Math.min( + Math.floor(0.8 * reserveTokens), + model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, + ); + + // Use update prompt if we have a previous summary, otherwise initial prompt + let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; + if (customInstructions) { + basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; + } + + // Serialize conversation to text so model doesn't try to continue it + // Convert to LLM messages first (handles custom types like bashExecution, custom, etc.) + const llmMessages = convertToLlm(currentMessages); + const conversationText = serializeConversation(llmMessages); + + // Build the prompt with conversation wrapped in tags + let promptText = `\n${conversationText}\n\n\n`; + if (previousSummary) { + promptText += `\n${previousSummary}\n\n\n`; + } + promptText += basePrompt; + + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + const completionOptions = createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel); + + const response = await completeSummarization( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + completionOptions, + streamFn, + ); + + if (response.stopReason === "error") { + throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`); + } + + const textContent = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + + return textContent; +} + +// ============================================================================ +// Compaction Preparation (for extensions) +// ============================================================================ + +export interface CompactionPreparation { + /** UUID of first entry to keep */ + firstKeptEntryId: string; + /** Messages that will be summarized and discarded */ + messagesToSummarize: AgentMessage[]; + /** Messages that will be turned into turn prefix summary (if splitting) */ + turnPrefixMessages: AgentMessage[]; + /** Whether this is a split turn (cut point in middle of turn) */ + isSplitTurn: boolean; + tokensBefore: number; + /** Summary from previous compaction, for iterative update */ + previousSummary?: string; + /** File operations extracted from messagesToSummarize */ + fileOps: FileOperations; + /** Compaction settions from settings.jsonl */ + settings: CompactionSettings; +} + +export function prepareCompaction( + pathEntries: SessionEntry[], + settings: CompactionSettings, +): CompactionPreparation | undefined { + if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") { + return undefined; + } + + let prevCompactionIndex = -1; + for (let i = pathEntries.length - 1; i >= 0; i--) { + if (pathEntries[i].type === "compaction") { + prevCompactionIndex = i; + break; + } + } + + let previousSummary: string | undefined; + let boundaryStart = 0; + if (prevCompactionIndex >= 0) { + const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; + previousSummary = prevCompaction.summary; + const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); + boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; + } + const boundaryEnd = pathEntries.length; + + const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens; + + const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens); + + // Get UUID of first kept entry + const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex]; + if (!firstKeptEntry?.id) { + return undefined; // Session needs migration + } + const firstKeptEntryId = firstKeptEntry.id; + + const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; + + // Messages to summarize (will be discarded after summary) + const messagesToSummarize: AgentMessage[] = []; + for (let i = boundaryStart; i < historyEnd; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) messagesToSummarize.push(msg); + } + + // Messages for turn prefix summary (if splitting a turn) + const turnPrefixMessages: AgentMessage[] = []; + if (cutPoint.isSplitTurn) { + for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + if (msg) turnPrefixMessages.push(msg); + } + } + + if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0) { + return undefined; + } + + // Extract file operations from messages and previous compaction + const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); + + // Also extract file ops from turn prefix if splitting + if (cutPoint.isSplitTurn) { + for (const msg of turnPrefixMessages) { + extractFileOpsFromMessage(msg, fileOps); + } + } + + return { + firstKeptEntryId, + messagesToSummarize, + turnPrefixMessages, + isSplitTurn: cutPoint.isSplitTurn, + tokensBefore, + previousSummary, + fileOps, + settings, + }; +} + +// ============================================================================ +// Main compaction function +// ============================================================================ + +const TURN_PREFIX_SUMMARIZATION_PROMPT = `This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained. + +Summarize the prefix to provide context for the retained suffix: + +## Original Request +[What did the user ask for in this turn?] + +## Early Progress +- [Key decisions and work done in the prefix] + +## Context for Suffix +- [Information needed to understand the retained recent work] + +Be concise. Focus on what's needed to understand the kept suffix.`; + +/** + * Generate summaries for compaction using prepared data. + * Returns CompactionResult - SessionManager adds uuid/parentUuid when saving. + * + * @param preparation - Pre-calculated preparation from prepareCompaction() + * @param customInstructions - Optional custom focus for the summary + */ +export async function compact( + preparation: CompactionPreparation, + model: Model, + apiKey: string | undefined, + headers?: Record, + customInstructions?: string, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, + env?: Record, +): Promise { + const { + firstKeptEntryId, + messagesToSummarize, + turnPrefixMessages, + isSplitTurn, + tokensBefore, + previousSummary, + fileOps, + settings, + } = preparation; + + // Generate summaries (can be parallel if both needed) and merge into one + let summary: string; + + if (isSplitTurn && turnPrefixMessages.length > 0) { + // Generate both summaries in parallel + const [historyResult, turnPrefixResult] = await Promise.all([ + messagesToSummarize.length > 0 + ? generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + streamFn, + env, + ) + : Promise.resolve("No prior history."), + generateTurnPrefixSummary( + turnPrefixMessages, + model, + settings.reserveTokens, + apiKey, + headers, + env, + signal, + thinkingLevel, + streamFn, + ), + ]); + // Merge into single summary + summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; + } else { + // Just generate history summary + summary = await generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + streamFn, + env, + ); + } + + // Compute file lists and append to summary + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + if (!firstKeptEntryId) { + throw new Error("First kept entry has no UUID - session may need migration"); + } + + return { + summary, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles } as CompactionDetails, + }; +} + +/** + * Generate a summary for a turn prefix (when splitting a turn). + */ +async function generateTurnPrefixSummary( + messages: AgentMessage[], + model: Model, + reserveTokens: number, + apiKey: string | undefined, + headers?: Record, + env?: Record, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, +): Promise { + const maxTokens = Math.min( + Math.floor(0.5 * reserveTokens), + model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, + ); // Smaller budget for turn prefix + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + const response = await completeSummarization( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel), + streamFn, + ); + + if (response.stopReason === "error") { + throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); + } + + return response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); +} diff --git a/cactus-code/packages/coding-agent/src/core/compaction/index.ts b/cactus-code/packages/coding-agent/src/core/compaction/index.ts new file mode 100644 index 000000000..7fae5f2c3 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/compaction/index.ts @@ -0,0 +1,7 @@ +/** + * Compaction and summarization utilities. + */ + +export * from "./branch-summarization.ts"; +export * from "./compaction.ts"; +export * from "./utils.ts"; diff --git a/cactus-code/packages/coding-agent/src/core/compaction/utils.ts b/cactus-code/packages/coding-agent/src/core/compaction/utils.ts new file mode 100644 index 000000000..6cfc16227 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/compaction/utils.ts @@ -0,0 +1,170 @@ +/** + * Shared utilities for compaction and branch summarization. + */ + +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { Message } from "@earendil-works/pi-ai"; + +// ============================================================================ +// File Operation Tracking +// ============================================================================ + +export interface FileOperations { + read: Set; + written: Set; + edited: Set; +} + +export function createFileOps(): FileOperations { + return { + read: new Set(), + written: new Set(), + edited: new Set(), + }; +} + +/** + * Extract file operations from tool calls in an assistant message. + */ +export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void { + if (message.role !== "assistant") return; + if (!("content" in message) || !Array.isArray(message.content)) return; + + for (const block of message.content) { + if (typeof block !== "object" || block === null) continue; + if (!("type" in block) || block.type !== "toolCall") continue; + if (!("arguments" in block) || !("name" in block)) continue; + + const args = block.arguments as Record | undefined; + if (!args) continue; + + const path = typeof args.path === "string" ? args.path : undefined; + if (!path) continue; + + switch (block.name) { + case "read": + fileOps.read.add(path); + break; + case "write": + fileOps.written.add(path); + break; + case "edit": + fileOps.edited.add(path); + break; + } + } +} + +/** + * Compute final file lists from file operations. + * Returns readFiles (files only read, not modified) and modifiedFiles. + */ +export function computeFileLists(fileOps: FileOperations): { readFiles: string[]; modifiedFiles: string[] } { + const modified = new Set([...fileOps.edited, ...fileOps.written]); + const readOnly = [...fileOps.read].filter((f) => !modified.has(f)).sort(); + const modifiedFiles = [...modified].sort(); + return { readFiles: readOnly, modifiedFiles }; +} + +/** + * Format file operations as XML tags for summary. + */ +export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string { + const sections: string[] = []; + if (readFiles.length > 0) { + sections.push(`\n${readFiles.join("\n")}\n`); + } + if (modifiedFiles.length > 0) { + sections.push(`\n${modifiedFiles.join("\n")}\n`); + } + if (sections.length === 0) return ""; + return `\n\n${sections.join("\n\n")}`; +} + +// ============================================================================ +// Message Serialization +// ============================================================================ + +/** Maximum characters for a tool result in serialized summaries. */ +const TOOL_RESULT_MAX_CHARS = 2000; + +/** + * Truncate text to a maximum character length for summarization. + * Keeps the beginning and appends a truncation marker. + */ +function truncateForSummary(text: string, maxChars: number): string { + if (text.length <= maxChars) return text; + const truncatedChars = text.length - maxChars; + return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`; +} + +/** + * Serialize LLM messages to text for summarization. + * This prevents the model from treating it as a conversation to continue. + * Call convertToLlm() first to handle custom message types. + * + * Tool results are truncated to keep the summarization request within + * reasonable token budgets. Full content is not needed for summarization. + */ +export function serializeConversation(messages: Message[]): string { + const parts: string[] = []; + + for (const msg of messages) { + if (msg.role === "user") { + const content = + typeof msg.content === "string" + ? msg.content + : msg.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (content) parts.push(`[User]: ${content}`); + } else if (msg.role === "assistant") { + const textParts: string[] = []; + const thinkingParts: string[] = []; + const toolCalls: string[] = []; + + for (const block of msg.content) { + if (block.type === "text") { + textParts.push(block.text); + } else if (block.type === "thinking") { + thinkingParts.push(block.thinking); + } else if (block.type === "toolCall") { + const args = block.arguments as Record; + const argsStr = Object.entries(args) + .map(([k, v]) => `${k}=${JSON.stringify(v)}`) + .join(", "); + toolCalls.push(`${block.name}(${argsStr})`); + } + } + + if (thinkingParts.length > 0) { + parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); + } + if (textParts.length > 0) { + parts.push(`[Assistant]: ${textParts.join("\n")}`); + } + if (toolCalls.length > 0) { + parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`); + } + } else if (msg.role === "toolResult") { + const content = msg.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (content) { + parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); + } + } + } + + return parts.join("\n\n"); +} + +// ============================================================================ +// Summarization System Prompt +// ============================================================================ + +export const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified. + +Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`; diff --git a/cactus-code/packages/coding-agent/src/core/defaults.ts b/cactus-code/packages/coding-agent/src/core/defaults.ts new file mode 100644 index 000000000..fddc7d14f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/defaults.ts @@ -0,0 +1,3 @@ +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; + +export const DEFAULT_THINKING_LEVEL: ThinkingLevel = "medium"; diff --git a/cactus-code/packages/coding-agent/src/core/diagnostics.ts b/cactus-code/packages/coding-agent/src/core/diagnostics.ts new file mode 100644 index 000000000..20fb80243 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/diagnostics.ts @@ -0,0 +1,15 @@ +export interface ResourceCollision { + resourceType: "extension" | "skill" | "prompt" | "theme"; + name: string; // skill name, command/tool/flag name, prompt name, theme name + winnerPath: string; + loserPath: string; + winnerSource?: string; // e.g., "npm:foo", "git:...", "local" + loserSource?: string; +} + +export interface ResourceDiagnostic { + type: "warning" | "error" | "collision"; + message: string; + path?: string; + collision?: ResourceCollision; +} diff --git a/cactus-code/packages/coding-agent/src/core/event-bus.ts b/cactus-code/packages/coding-agent/src/core/event-bus.ts new file mode 100644 index 000000000..a4c87b9f0 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/event-bus.ts @@ -0,0 +1,33 @@ +import { EventEmitter } from "node:events"; + +export interface EventBus { + emit(channel: string, data: unknown): void; + on(channel: string, handler: (data: unknown) => void): () => void; +} + +export interface EventBusController extends EventBus { + clear(): void; +} + +export function createEventBus(): EventBusController { + const emitter = new EventEmitter(); + return { + emit: (channel, data) => { + emitter.emit(channel, data); + }, + on: (channel, handler) => { + const safeHandler = async (data: unknown) => { + try { + await handler(data); + } catch (err) { + console.error(`Event handler error (${channel}):`, err); + } + }; + emitter.on(channel, safeHandler); + return () => emitter.off(channel, safeHandler); + }, + clear: () => { + emitter.removeAllListeners(); + }, + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/exec.ts b/cactus-code/packages/coding-agent/src/core/exec.ts new file mode 100644 index 000000000..5afe9ec36 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/exec.ts @@ -0,0 +1,107 @@ +/** + * Shared command execution utilities for extensions and custom tools. + */ + +import { spawn } from "node:child_process"; +import { waitForChildProcess } from "../utils/child-process.ts"; + +/** + * Options for executing shell commands. + */ +export interface ExecOptions { + /** AbortSignal to cancel the command */ + signal?: AbortSignal; + /** Timeout in milliseconds */ + timeout?: number; + /** Working directory */ + cwd?: string; +} + +/** + * Result of executing a shell command. + */ +export interface ExecResult { + stdout: string; + stderr: string; + code: number; + killed: boolean; +} + +/** + * Execute a shell command and return stdout/stderr/code. + * Supports timeout and abort signal. + */ +export async function execCommand( + command: string, + args: string[], + cwd: string, + options?: ExecOptions, +): Promise { + return new Promise((resolve) => { + const proc = spawn(command, args, { + cwd, + shell: false, + stdio: ["ignore", "pipe", "pipe"], + }); + + let stdout = ""; + let stderr = ""; + let killed = false; + let timeoutId: NodeJS.Timeout | undefined; + + const killProcess = () => { + if (!killed) { + killed = true; + proc.kill("SIGTERM"); + // Force kill after 5 seconds if SIGTERM doesn't work + setTimeout(() => { + if (!proc.killed) { + proc.kill("SIGKILL"); + } + }, 5000); + } + }; + + // Handle abort signal + if (options?.signal) { + if (options.signal.aborted) { + killProcess(); + } else { + options.signal.addEventListener("abort", killProcess, { once: true }); + } + } + + // Handle timeout + if (options?.timeout && options.timeout > 0) { + timeoutId = setTimeout(() => { + killProcess(); + }, options.timeout); + } + + proc.stdout?.on("data", (data) => { + stdout += data.toString(); + }); + + proc.stderr?.on("data", (data) => { + stderr += data.toString(); + }); + + // Wait for process termination without hanging on inherited stdio handles + // held open by detached descendants. + waitForChildProcess(proc) + .then((code) => { + if (timeoutId) clearTimeout(timeoutId); + if (options?.signal) { + options.signal.removeEventListener("abort", killProcess); + } + resolve({ stdout, stderr, code: code ?? 0, killed }); + }) + .catch((_err) => { + if (timeoutId) clearTimeout(timeoutId); + if (options?.signal) { + options.signal.removeEventListener("abort", killProcess); + } + resolve({ stdout, stderr, code: 1, killed }); + }); + }); +} diff --git a/cactus-code/packages/coding-agent/src/core/experimental.ts b/cactus-code/packages/coding-agent/src/core/experimental.ts new file mode 100644 index 000000000..12d33c74c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/experimental.ts @@ -0,0 +1,3 @@ +export function areExperimentalFeaturesEnabled(): boolean { + return process.env.PI_EXPERIMENTAL === "1"; +} diff --git a/cactus-code/packages/coding-agent/src/core/extensions/index.ts b/cactus-code/packages/coding-agent/src/core/extensions/index.ts new file mode 100644 index 000000000..3ed7a6625 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/extensions/index.ts @@ -0,0 +1,178 @@ +/** + * Extension system for lifecycle events and custom tools. + */ + +export type { SlashCommandInfo, SlashCommandSource } from "../slash-commands.ts"; +export type { SourceInfo } from "../source-info.ts"; +export { + createExtensionRuntime, + discoverAndLoadExtensions, + loadExtensionFromFactory, + loadExtensions, +} from "./loader.ts"; +export type { + ExtensionErrorListener, + ForkHandler, + NavigateTreeHandler, + NewSessionHandler, + ShutdownHandler, + SwitchSessionHandler, +} from "./runner.ts"; +export { ExtensionRunner } from "./runner.ts"; +export type { + AfterProviderResponseEvent, + AgentEndEvent, + AgentStartEvent, + // Re-exports + AgentToolResult, + AgentToolUpdateCallback, + AppendEntryHandler, + // App keybindings (for custom editors) + AppKeybinding, + AutocompleteProviderFactory, + // Events - Tool (ToolCallEvent types) + BashToolCallEvent, + BashToolResultEvent, + BeforeAgentStartEvent, + BeforeAgentStartEventResult, + BeforeProviderRequestEvent, + BeforeProviderRequestEventResult, + BuildSystemPromptOptions, + // Context + CompactOptions, + // Events - Agent + ContextEvent, + // Event Results + ContextEventResult, + ContextUsage, + CustomToolCallEvent, + CustomToolResultEvent, + EditorFactory, + EditToolCallEvent, + EditToolResultEvent, + ExecOptions, + ExecResult, + Extension, + ExtensionActions, + // API + ExtensionAPI, + ExtensionCommandContext, + ExtensionCommandContextActions, + ExtensionContext, + ExtensionContextActions, + // Errors + ExtensionError, + ExtensionEvent, + ExtensionFactory, + ExtensionFlag, + ExtensionHandler, + ExtensionMode, + // Runtime + ExtensionRuntime, + ExtensionShortcut, + ExtensionUIContext, + ExtensionUIDialogOptions, + ExtensionWidgetOptions, + FindToolCallEvent, + FindToolResultEvent, + GetActiveToolsHandler, + GetAllToolsHandler, + GetCommandsHandler, + GetThinkingLevelHandler, + GrepToolCallEvent, + GrepToolResultEvent, + // Events - Input + InputEvent, + InputEventResult, + InputSource, + KeybindingsManager, + LoadExtensionsResult, + LsToolCallEvent, + LsToolResultEvent, + // Events - Message + MessageEndEvent, + // Message Rendering + MessageRenderer, + MessageRenderOptions, + MessageStartEvent, + MessageUpdateEvent, + ModelSelectEvent, + ModelSelectSource, + ProjectTrustContext, + ProjectTrustEvent, + ProjectTrustEventDecision, + ProjectTrustEventResult, + ProjectTrustHandler, + // Provider Registration + ProviderConfig, + ProviderModelConfig, + ReadToolCallEvent, + ReadToolResultEvent, + // Commands + RegisteredCommand, + RegisteredTool, + ReplacedSessionContext, + ResolvedCommand, + // Events - Resources + ResourcesDiscoverEvent, + ResourcesDiscoverResult, + SendMessageHandler, + SendUserMessageHandler, + SessionBeforeCompactEvent, + SessionBeforeCompactResult, + SessionBeforeForkEvent, + SessionBeforeForkResult, + SessionBeforeSwitchEvent, + SessionBeforeSwitchResult, + SessionBeforeTreeEvent, + SessionBeforeTreeResult, + SessionCompactEvent, + SessionEvent, + SessionShutdownEvent, + // Events - Session + SessionStartEvent, + SessionTreeEvent, + SetActiveToolsHandler, + SetLabelHandler, + SetModelHandler, + SetThinkingLevelHandler, + TerminalInputHandler, + // Events - Tool + ToolCallEvent, + ToolCallEventResult, + // Tools + ToolDefinition, + // Events - Tool Execution + ToolExecutionEndEvent, + // Tool execution mode + ToolExecutionMode, + ToolExecutionStartEvent, + ToolExecutionUpdateEvent, + ToolInfo, + ToolRenderResultOptions, + ToolResultEvent, + ToolResultEventResult, + TreePreparation, + TurnEndEvent, + TurnStartEvent, + // Events - User Bash + UserBashEvent, + UserBashEventResult, + WidgetPlacement, + WorkingIndicatorOptions, + WriteToolCallEvent, + WriteToolResultEvent, +} from "./types.ts"; +// Type guards +export { + defineTool, + isBashToolResult, + isEditToolResult, + isFindToolResult, + isGrepToolResult, + isLsToolResult, + isReadToolResult, + isToolCallEventType, + isWriteToolResult, +} from "./types.ts"; +export { wrapRegisteredTool, wrapRegisteredTools } from "./wrapper.ts"; diff --git a/cactus-code/packages/coding-agent/src/core/extensions/loader.ts b/cactus-code/packages/coding-agent/src/core/extensions/loader.ts new file mode 100644 index 000000000..a93f7c852 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/extensions/loader.ts @@ -0,0 +1,687 @@ +/** + * Extension loader - loads TypeScript extension modules using jiti. + * + */ + +import * as fs from "node:fs"; +import { createRequire } from "node:module"; +import * as path from "node:path"; +import { fileURLToPath } from "node:url"; +import * as _bundledPiAgentCore from "@earendil-works/pi-agent-core"; +import * as _bundledPiAiCompat from "@earendil-works/pi-ai/compat"; +import * as _bundledPiAiOauth from "@earendil-works/pi-ai/oauth"; +import type { KeyId } from "@earendil-works/pi-tui"; +import * as _bundledPiTui from "@earendil-works/pi-tui"; +import { createJiti } from "jiti/static"; +// Static imports of packages that extensions may use. +// These MUST be static so Bun bundles them into the compiled binary. +// The virtualModules option then makes them available to extensions. +import * as _bundledTypebox from "typebox"; +import * as _bundledTypeboxCompile from "typebox/compile"; +import * as _bundledTypeboxValue from "typebox/value"; +import { CONFIG_DIR_NAME, getAgentDir, isBunBinary } from "../../config.ts"; +// NOTE: This import works because loader.ts exports are NOT re-exported from index.ts, +// avoiding a circular dependency. Extensions can import from @earendil-works/pi-coding-agent. +import * as _bundledPiCodingAgent from "../../index.ts"; +import { resolvePath } from "../../utils/paths.ts"; +import { createEventBus, type EventBus } from "../event-bus.ts"; +import type { ExecOptions } from "../exec.ts"; +import { execCommand } from "../exec.ts"; +import { createSyntheticSourceInfo } from "../source-info.ts"; +import type { + Extension, + ExtensionAPI, + ExtensionFactory, + ExtensionRuntime, + LoadExtensionsResult, + MessageRenderer, + ProviderConfig, + RegisteredCommand, + ToolDefinition, +} from "./types.ts"; + +/** Modules available to extensions via virtualModules (for compiled Bun binary) */ +const VIRTUAL_MODULES: Record = { + typebox: _bundledTypebox, + "typebox/compile": _bundledTypeboxCompile, + "typebox/value": _bundledTypeboxValue, + "@sinclair/typebox": _bundledTypebox, + "@sinclair/typebox/compile": _bundledTypeboxCompile, + "@sinclair/typebox/value": _bundledTypeboxValue, + "@earendil-works/pi-agent-core": _bundledPiAgentCore, + "@earendil-works/pi-tui": _bundledPiTui, + // Extensions resolve the pi-ai root to the compat entrypoint (a strict + // superset of the core entrypoint): existing extensions using the old + // global API keep working at runtime until compat is removed. + "@earendil-works/pi-ai": _bundledPiAiCompat, + "@earendil-works/pi-ai/compat": _bundledPiAiCompat, + "@earendil-works/pi-ai/oauth": _bundledPiAiOauth, + "@earendil-works/pi-coding-agent": _bundledPiCodingAgent, + "@mariozechner/pi-agent-core": _bundledPiAgentCore, + "@mariozechner/pi-tui": _bundledPiTui, + "@mariozechner/pi-ai": _bundledPiAiCompat, + "@mariozechner/pi-ai/compat": _bundledPiAiCompat, + "@mariozechner/pi-ai/oauth": _bundledPiAiOauth, + "@mariozechner/pi-coding-agent": _bundledPiCodingAgent, +}; + +const require = createRequire(import.meta.url); + +/** + * Get aliases for jiti (used in Node.js/development mode). + * In Bun binary mode, virtualModules is used instead. + */ +let _aliases: Record | null = null; + +function getAliases(): Record { + if (_aliases) return _aliases; + + const __dirname = path.dirname(fileURLToPath(import.meta.url)); + const packageIndex = path.resolve(__dirname, "../..", "index.js"); + + const typeboxEntry = require.resolve("typebox"); + const typeboxCompileEntry = require.resolve("typebox/compile"); + const typeboxValueEntry = require.resolve("typebox/value"); + + const packagesRoot = path.resolve(__dirname, "../../../../"); + const resolveWorkspaceOrImport = (workspaceRelativePath: string, specifier: string): string => { + const workspacePath = path.join(packagesRoot, workspaceRelativePath); + if (fs.existsSync(workspacePath)) { + return workspacePath; + } + return fileURLToPath(import.meta.resolve(specifier)); + }; + + const piCodingAgentEntry = packageIndex; + const piAgentCoreEntry = resolveWorkspaceOrImport("agent/dist/index.js", "@earendil-works/pi-agent-core"); + const piTuiEntry = resolveWorkspaceOrImport("tui/dist/index.js", "@earendil-works/pi-tui"); + // Extensions resolve the pi-ai root to the compat entrypoint (a strict + // superset of the core entrypoint): existing extensions using the old + // global API keep working at runtime until compat is removed. + const piAiCompatEntry = resolveWorkspaceOrImport("ai/dist/compat.js", "@earendil-works/pi-ai/compat"); + const piAiOauthEntry = resolveWorkspaceOrImport("ai/dist/oauth.js", "@earendil-works/pi-ai/oauth"); + + _aliases = { + "@earendil-works/pi-coding-agent": piCodingAgentEntry, + "@earendil-works/pi-agent-core": piAgentCoreEntry, + "@earendil-works/pi-tui": piTuiEntry, + "@earendil-works/pi-ai": piAiCompatEntry, + "@earendil-works/pi-ai/compat": piAiCompatEntry, + "@earendil-works/pi-ai/oauth": piAiOauthEntry, + "@mariozechner/pi-coding-agent": piCodingAgentEntry, + "@mariozechner/pi-agent-core": piAgentCoreEntry, + "@mariozechner/pi-tui": piTuiEntry, + "@mariozechner/pi-ai": piAiCompatEntry, + "@mariozechner/pi-ai/compat": piAiCompatEntry, + "@mariozechner/pi-ai/oauth": piAiOauthEntry, + typebox: typeboxEntry, + "typebox/compile": typeboxCompileEntry, + "typebox/value": typeboxValueEntry, + "@sinclair/typebox": typeboxEntry, + "@sinclair/typebox/compile": typeboxCompileEntry, + "@sinclair/typebox/value": typeboxValueEntry, + }; + + return _aliases; +} + +type HandlerFn = (...args: unknown[]) => Promise; + +let extensionCacheCwd: string | undefined; +let extensionCacheGeneration = 0; +const extensionCache = new Map(); + +interface ExtensionCacheToken { + cwd: string; + generation: number; +} + +export function clearExtensionCache(): void { + extensionCache.clear(); + extensionCacheCwd = undefined; + extensionCacheGeneration++; +} + +function useExtensionCacheCwd(cwd: string): ExtensionCacheToken { + const resolvedCwd = resolvePath(cwd); + if (extensionCacheCwd !== undefined && extensionCacheCwd !== resolvedCwd) { + clearExtensionCache(); + } + extensionCacheCwd = resolvedCwd; + return { cwd: resolvedCwd, generation: extensionCacheGeneration }; +} + +/** + * Create a runtime with throwing stubs for action methods. + * Runner.bindCore() replaces these with real implementations. + */ +export function createExtensionRuntime(): ExtensionRuntime { + const notInitialized = () => { + throw new Error("Extension runtime not initialized. Action methods cannot be called during extension loading."); + }; + const state: { staleMessage?: string } = {}; + const assertActive = () => { + if (state.staleMessage) { + throw new Error(state.staleMessage); + } + }; + + const runtime: ExtensionRuntime = { + sendMessage: notInitialized, + sendUserMessage: notInitialized, + appendEntry: notInitialized, + setSessionName: notInitialized, + getSessionName: notInitialized, + setLabel: notInitialized, + getActiveTools: notInitialized, + getAllTools: notInitialized, + setActiveTools: notInitialized, + // registerTool() is valid during extension load; refresh is only needed post-bind. + refreshTools: () => {}, + getCommands: notInitialized, + setModel: () => Promise.reject(new Error("Extension runtime not initialized")), + getThinkingLevel: notInitialized, + setThinkingLevel: notInitialized, + flagValues: new Map(), + pendingProviderRegistrations: [], + assertActive, + invalidate: (message) => { + state.staleMessage ??= + message ?? + "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload()."; + }, + // Pre-bind: queue registrations so bindCore() can flush them once the + // model registry is available. bindCore() replaces both with direct calls. + registerProvider: (name, config, extensionPath = "") => { + runtime.pendingProviderRegistrations.push({ name, config, extensionPath }); + }, + unregisterProvider: (name) => { + runtime.pendingProviderRegistrations = runtime.pendingProviderRegistrations.filter((r) => r.name !== name); + }, + }; + + return runtime; +} + +/** + * Create the ExtensionAPI for an extension. + * Registration methods write to the extension object. + * Action methods delegate to the shared runtime. + */ +function createExtensionAPI( + extension: Extension, + runtime: ExtensionRuntime, + cwd: string, + eventBus: EventBus, +): ExtensionAPI { + const api = { + // Registration methods - write to extension + on(event: string, handler: HandlerFn): void { + runtime.assertActive(); + const list = extension.handlers.get(event) ?? []; + list.push(handler); + extension.handlers.set(event, list); + }, + + registerTool(tool: ToolDefinition): void { + runtime.assertActive(); + extension.tools.set(tool.name, { + definition: tool, + sourceInfo: extension.sourceInfo, + }); + runtime.refreshTools(); + }, + + registerCommand(name: string, options: Omit): void { + runtime.assertActive(); + extension.commands.set(name, { + name, + sourceInfo: extension.sourceInfo, + ...options, + }); + }, + + registerShortcut( + shortcut: KeyId, + options: { + description?: string; + handler: (ctx: import("./types.ts").ExtensionContext) => Promise | void; + }, + ): void { + runtime.assertActive(); + extension.shortcuts.set(shortcut, { shortcut, extensionPath: extension.path, ...options }); + }, + + registerFlag( + name: string, + options: { description?: string; type: "boolean" | "string"; default?: boolean | string }, + ): void { + runtime.assertActive(); + extension.flags.set(name, { name, extensionPath: extension.path, ...options }); + if (options.default !== undefined && !runtime.flagValues.has(name)) { + runtime.flagValues.set(name, options.default); + } + }, + + registerMessageRenderer(customType: string, renderer: MessageRenderer): void { + runtime.assertActive(); + extension.messageRenderers.set(customType, renderer as MessageRenderer); + }, + + // Flag access - checks extension registered it, reads from runtime + getFlag(name: string): boolean | string | undefined { + runtime.assertActive(); + if (!extension.flags.has(name)) return undefined; + return runtime.flagValues.get(name); + }, + + // Action methods - delegate to shared runtime + sendMessage(message, options): void { + runtime.assertActive(); + runtime.sendMessage(message, options); + }, + + sendUserMessage(content, options): void { + runtime.assertActive(); + runtime.sendUserMessage(content, options); + }, + + appendEntry(customType: string, data?: unknown): void { + runtime.assertActive(); + runtime.appendEntry(customType, data); + }, + + setSessionName(name: string): void { + runtime.assertActive(); + runtime.setSessionName(name); + }, + + getSessionName(): string | undefined { + runtime.assertActive(); + return runtime.getSessionName(); + }, + + setLabel(entryId: string, label: string | undefined): void { + runtime.assertActive(); + runtime.setLabel(entryId, label); + }, + + exec(command: string, args: string[], options?: ExecOptions) { + runtime.assertActive(); + return execCommand(command, args, options?.cwd ?? cwd, options); + }, + + getActiveTools(): string[] { + runtime.assertActive(); + return runtime.getActiveTools(); + }, + + getAllTools() { + runtime.assertActive(); + return runtime.getAllTools(); + }, + + setActiveTools(toolNames: string[]): void { + runtime.assertActive(); + runtime.setActiveTools(toolNames); + }, + + getCommands() { + runtime.assertActive(); + return runtime.getCommands(); + }, + + setModel(model) { + runtime.assertActive(); + return runtime.setModel(model); + }, + + getThinkingLevel() { + runtime.assertActive(); + return runtime.getThinkingLevel(); + }, + + setThinkingLevel(level) { + runtime.assertActive(); + runtime.setThinkingLevel(level); + }, + + registerProvider(name: string, config: ProviderConfig) { + runtime.assertActive(); + runtime.registerProvider(name, config, extension.path); + }, + + unregisterProvider(name: string) { + runtime.assertActive(); + runtime.unregisterProvider(name, extension.path); + }, + + events: eventBus, + } as ExtensionAPI; + + return api; +} + +function isCurrentCacheToken(cacheToken: ExtensionCacheToken | undefined): cacheToken is ExtensionCacheToken { + return ( + cacheToken !== undefined && + extensionCacheCwd === cacheToken.cwd && + extensionCacheGeneration === cacheToken.generation + ); +} + +async function loadExtensionModule(extensionPath: string, cacheToken?: ExtensionCacheToken) { + if (isCurrentCacheToken(cacheToken)) { + const cachedFactory = extensionCache.get(extensionPath); + if (cachedFactory) { + return cachedFactory; + } + } + + const jiti = createJiti(import.meta.url, { + moduleCache: false, + // In Bun binary: use virtualModules for bundled packages (no filesystem resolution) + // Also disable tryNative so jiti handles ALL imports (not just the entry point) + // In Node.js/dev: use aliases to resolve to node_modules paths + ...(isBunBinary ? { virtualModules: VIRTUAL_MODULES, tryNative: false } : { alias: getAliases() }), + }); + + const module = await jiti.import(extensionPath, { default: true }); + const factory = module as ExtensionFactory; + if (typeof factory !== "function") { + return undefined; + } + if (isCurrentCacheToken(cacheToken)) { + extensionCache.set(extensionPath, factory); + } + return factory; +} + +/** + * Create an Extension object with empty collections. + */ +function createExtension(extensionPath: string, resolvedPath: string): Extension { + const source = + extensionPath.startsWith("<") && extensionPath.endsWith(">") + ? extensionPath.slice(1, -1).split(":")[0] || "temporary" + : "local"; + const baseDir = extensionPath.startsWith("<") ? undefined : path.dirname(resolvedPath); + + return { + path: extensionPath, + resolvedPath, + sourceInfo: createSyntheticSourceInfo(extensionPath, { source, baseDir }), + handlers: new Map(), + tools: new Map(), + messageRenderers: new Map(), + commands: new Map(), + flags: new Map(), + shortcuts: new Map(), + }; +} + +async function loadExtension( + extensionPath: string, + cwd: string, + eventBus: EventBus, + runtime: ExtensionRuntime, + cacheToken?: ExtensionCacheToken, +): Promise<{ extension: Extension | null; error: string | null }> { + const resolvedPath = resolvePath(extensionPath, cwd, { normalizeUnicodeSpaces: true }); + + try { + const factory = await loadExtensionModule(resolvedPath, cacheToken); + if (!factory) { + return { extension: null, error: `Extension does not export a valid factory function: ${extensionPath}` }; + } + + const extension = createExtension(extensionPath, resolvedPath); + const api = createExtensionAPI(extension, runtime, cwd, eventBus); + await factory(api); + + return { extension, error: null }; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + return { extension: null, error: `Failed to load extension: ${message}` }; + } +} + +/** + * Create an Extension from an inline factory function. + */ +export async function loadExtensionFromFactory( + factory: ExtensionFactory, + cwd: string, + eventBus: EventBus, + runtime: ExtensionRuntime, + extensionPath = "", +): Promise { + const extension = createExtension(extensionPath, extensionPath); + const resolvedCwd = resolvePath(cwd); + const api = createExtensionAPI(extension, runtime, resolvedCwd, eventBus); + await factory(api); + return extension; +} + +/** + * Load extensions from paths. + */ +async function loadExtensionsInternal( + paths: string[], + cwd: string, + eventBus?: EventBus, + runtime?: ExtensionRuntime, + useCache = false, +): Promise { + const extensions: Extension[] = []; + const errors: Array<{ path: string; error: string }> = []; + const cacheToken = useCache ? useExtensionCacheCwd(cwd) : undefined; + const resolvedCwd = cacheToken?.cwd ?? resolvePath(cwd); + const resolvedEventBus = eventBus ?? createEventBus(); + const resolvedRuntime = runtime ?? createExtensionRuntime(); + + for (const extPath of paths) { + const { extension, error } = await loadExtension( + extPath, + resolvedCwd, + resolvedEventBus, + resolvedRuntime, + cacheToken, + ); + + if (error) { + errors.push({ path: extPath, error }); + continue; + } + + if (extension) { + extensions.push(extension); + } + } + + return { + extensions, + errors, + runtime: resolvedRuntime, + }; +} + +export async function loadExtensions( + paths: string[], + cwd: string, + eventBus?: EventBus, + runtime?: ExtensionRuntime, +): Promise { + return loadExtensionsInternal(paths, cwd, eventBus, runtime); +} + +export async function loadExtensionsCached( + paths: string[], + cwd: string, + eventBus?: EventBus, + runtime?: ExtensionRuntime, +): Promise { + return loadExtensionsInternal(paths, cwd, eventBus, runtime, true); +} + +interface PiManifest { + extensions?: string[]; + themes?: string[]; + skills?: string[]; + prompts?: string[]; +} + +function readPiManifest(packageJsonPath: string): PiManifest | null { + try { + const content = fs.readFileSync(packageJsonPath, "utf-8"); + const pkg = JSON.parse(content); + if (pkg.pi && typeof pkg.pi === "object") { + return pkg.pi as PiManifest; + } + return null; + } catch { + return null; + } +} + +function isExtensionFile(name: string): boolean { + return name.endsWith(".ts") || name.endsWith(".js"); +} + +/** + * Resolve extension entry points from a directory. + * + * Checks for: + * 1. package.json with "pi.extensions" field -> returns declared paths + * 2. index.ts or index.js -> returns the index file + * + * Returns resolved paths or null if no entry points found. + */ +function resolveExtensionEntries(dir: string): string[] | null { + // Check for package.json with "pi" field first + const packageJsonPath = path.join(dir, "package.json"); + if (fs.existsSync(packageJsonPath)) { + const manifest = readPiManifest(packageJsonPath); + if (manifest?.extensions?.length) { + const entries: string[] = []; + for (const extPath of manifest.extensions) { + const resolvedExtPath = path.resolve(dir, extPath); + if (fs.existsSync(resolvedExtPath)) { + entries.push(resolvedExtPath); + } + } + if (entries.length > 0) { + return entries; + } + } + } + + // Check for index.ts or index.js + const indexTs = path.join(dir, "index.ts"); + const indexJs = path.join(dir, "index.js"); + if (fs.existsSync(indexTs)) { + return [indexTs]; + } + if (fs.existsSync(indexJs)) { + return [indexJs]; + } + + return null; +} + +/** + * Discover extensions in a directory. + * + * Discovery rules: + * 1. Direct files: `extensions/*.ts` or `*.js` → load + * 2. Subdirectory with index: `extensions/* /index.ts` or `index.js` → load + * 3. Subdirectory with package.json: `extensions/* /package.json` with "pi" field → load what it declares + * + * No recursion beyond one level. Complex packages must use package.json manifest. + */ +function discoverExtensionsInDir(dir: string): string[] { + if (!fs.existsSync(dir)) { + return []; + } + + const discovered: string[] = []; + + try { + const entries = fs.readdirSync(dir, { withFileTypes: true }); + + for (const entry of entries) { + const entryPath = path.join(dir, entry.name); + + // 1. Direct files: *.ts or *.js + if ((entry.isFile() || entry.isSymbolicLink()) && isExtensionFile(entry.name)) { + discovered.push(entryPath); + continue; + } + + // 2 & 3. Subdirectories + if (entry.isDirectory() || entry.isSymbolicLink()) { + const entries = resolveExtensionEntries(entryPath); + if (entries) { + discovered.push(...entries); + } + } + } + } catch { + return []; + } + + return discovered; +} + +/** + * Discover and load extensions from standard locations. + */ +export async function discoverAndLoadExtensions( + configuredPaths: string[], + cwd: string, + agentDir: string = getAgentDir(), + eventBus?: EventBus, +): Promise { + const resolvedCwd = resolvePath(cwd); + const resolvedAgentDir = resolvePath(agentDir); + const allPaths: string[] = []; + const seen = new Set(); + + const addPaths = (paths: string[]) => { + for (const p of paths) { + const resolved = path.resolve(p); + if (!seen.has(resolved)) { + seen.add(resolved); + allPaths.push(p); + } + } + }; + + // 1. Project-local extensions: cwd/${CONFIG_DIR_NAME}/extensions/ + const localExtDir = path.join(resolvedCwd, CONFIG_DIR_NAME, "extensions"); + addPaths(discoverExtensionsInDir(localExtDir)); + + // 2. Global extensions: agentDir/extensions/ + const globalExtDir = path.join(resolvedAgentDir, "extensions"); + addPaths(discoverExtensionsInDir(globalExtDir)); + + // 3. Explicitly configured paths + for (const p of configuredPaths) { + const resolved = resolvePath(p, resolvedCwd, { normalizeUnicodeSpaces: true }); + if (fs.existsSync(resolved) && fs.statSync(resolved).isDirectory()) { + // Check for package.json with pi manifest or index.ts + const entries = resolveExtensionEntries(resolved); + if (entries) { + addPaths(entries); + continue; + } + // No explicit entries - discover individual files in directory + addPaths(discoverExtensionsInDir(resolved)); + continue; + } + + addPaths([resolved]); + } + + return loadExtensions(allPaths, resolvedCwd, eventBus); +} diff --git a/cactus-code/packages/coding-agent/src/core/extensions/runner.ts b/cactus-code/packages/coding-agent/src/core/extensions/runner.ts new file mode 100644 index 000000000..9cb0e6570 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/extensions/runner.ts @@ -0,0 +1,1135 @@ +/** + * Extension runner - executes extensions and manages their lifecycle. + */ + +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { ImageContent, Model } from "@earendil-works/pi-ai"; +import type { KeyId } from "@earendil-works/pi-tui"; +import { type Theme, theme } from "../../modes/interactive/theme/theme.ts"; +import type { ResourceDiagnostic } from "../diagnostics.ts"; +import type { KeybindingsConfig } from "../keybindings.ts"; +import type { ModelRegistry } from "../model-registry.ts"; +import type { SessionManager } from "../session-manager.ts"; +import type { BuildSystemPromptOptions } from "../system-prompt.ts"; +import type { + BeforeAgentStartEvent, + BeforeAgentStartEventResult, + BeforeProviderRequestEvent, + CompactOptions, + ContextEvent, + ContextEventResult, + ContextUsage, + Extension, + ExtensionActions, + ExtensionCommandContext, + ExtensionCommandContextActions, + ExtensionContext, + ExtensionContextActions, + ExtensionError, + ExtensionEvent, + ExtensionFlag, + ExtensionMode, + ExtensionRuntime, + ExtensionShortcut, + ExtensionUIContext, + InputEvent, + InputEventResult, + InputSource, + LoadExtensionsResult, + MessageEndEvent, + MessageEndEventResult, + MessageRenderer, + ProjectTrustContext, + ProjectTrustEvent, + ProjectTrustEventResult, + ProviderConfig, + RegisteredCommand, + RegisteredTool, + ReplacedSessionContext, + ResolvedCommand, + ResourcesDiscoverEvent, + ResourcesDiscoverResult, + SessionBeforeCompactResult, + SessionBeforeForkResult, + SessionBeforeSwitchResult, + SessionBeforeTreeResult, + SessionShutdownEvent, + ToolCallEvent, + ToolCallEventResult, + ToolResultEvent, + ToolResultEventResult, + UserBashEvent, + UserBashEventResult, +} from "./types.ts"; + +// Extension shortcuts compete with canonical keybinding ids from keybindings.json. +// Only editor-global shortcuts are reserved here. Picker-specific bindings are not. +const RESERVED_KEYBINDINGS_FOR_EXTENSION_CONFLICTS = [ + "app.interrupt", + "app.clear", + "app.exit", + "app.suspend", + "app.thinking.cycle", + "app.model.cycleForward", + "app.model.cycleBackward", + "app.model.select", + "app.tools.expand", + "app.thinking.toggle", + "app.editor.external", + "app.message.followUp", + "tui.input.submit", + "tui.select.confirm", + "tui.select.cancel", + "tui.input.copy", + "tui.editor.deleteToLineEnd", +] as const; + +type BuiltInKeyBindings = Partial>; + +const buildBuiltinKeybindings = (resolvedKeybindings: KeybindingsConfig): BuiltInKeyBindings => { + const builtinKeybindings = {} as BuiltInKeyBindings; + for (const [keybinding, keys] of Object.entries(resolvedKeybindings)) { + if (keys === undefined) continue; + const keyList = Array.isArray(keys) ? keys : [keys]; + const restrictOverride = (RESERVED_KEYBINDINGS_FOR_EXTENSION_CONFLICTS as readonly string[]).includes(keybinding); + for (const key of keyList) { + const normalizedKey = key.toLowerCase() as KeyId; + // If multiple actions bind the same key, the reserved action wins so extensions + // remain blocked by reserved shortcuts regardless of iteration order. + const existing = builtinKeybindings[normalizedKey]; + if (existing?.restrictOverride && !restrictOverride) continue; + builtinKeybindings[normalizedKey] = { + keybinding, + restrictOverride, + }; + } + } + return builtinKeybindings; +}; + +/** Combined result from all before_agent_start handlers */ +interface BeforeAgentStartCombinedResult { + messages?: NonNullable[]; + systemPrompt?: string; +} + +/** + * Events handled by the generic emit() method. + * Events with dedicated emitXxx() methods are excluded for stronger type safety. + */ +type RunnerEmitEvent = Exclude< + ExtensionEvent, + | ToolCallEvent + | ProjectTrustEvent + | ToolResultEvent + | UserBashEvent + | ContextEvent + | BeforeProviderRequestEvent + | BeforeAgentStartEvent + | MessageEndEvent + | ResourcesDiscoverEvent + | InputEvent +>; + +type SessionBeforeEvent = Extract< + RunnerEmitEvent, + { type: "session_before_switch" | "session_before_fork" | "session_before_compact" | "session_before_tree" } +>; + +type SessionBeforeEventResult = + | SessionBeforeSwitchResult + | SessionBeforeForkResult + | SessionBeforeCompactResult + | SessionBeforeTreeResult; + +type RunnerEmitResult = TEvent extends { type: "session_before_switch" } + ? SessionBeforeSwitchResult | undefined + : TEvent extends { type: "session_before_fork" } + ? SessionBeforeForkResult | undefined + : TEvent extends { type: "session_before_compact" } + ? SessionBeforeCompactResult | undefined + : TEvent extends { type: "session_before_tree" } + ? SessionBeforeTreeResult | undefined + : undefined; + +export type ExtensionErrorListener = (error: ExtensionError) => void; + +export type NewSessionHandler = (options?: { + parentSession?: string; + setup?: (sessionManager: SessionManager) => Promise; + withSession?: (ctx: ReplacedSessionContext) => Promise; +}) => Promise<{ cancelled: boolean }>; + +export type ForkHandler = ( + entryId: string, + options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise }, +) => Promise<{ cancelled: boolean }>; + +export type NavigateTreeHandler = ( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, +) => Promise<{ cancelled: boolean }>; + +export type SwitchSessionHandler = ( + sessionPath: string, + options?: { withSession?: (ctx: ReplacedSessionContext) => Promise }, +) => Promise<{ cancelled: boolean }>; + +export type ReloadHandler = () => Promise; + +export type ShutdownHandler = () => void; + +/** + * Helper function to emit session_shutdown event to extensions. + * Returns true if the event was emitted, false if there were no handlers. + */ +export async function emitSessionShutdownEvent( + extensionRunner: ExtensionRunner, + event: SessionShutdownEvent, +): Promise { + if (extensionRunner.hasHandlers("session_shutdown")) { + await extensionRunner.emit(event); + return true; + } + return false; +} + +export async function emitProjectTrustEvent( + extensionsResult: LoadExtensionsResult, + event: ProjectTrustEvent, + ctx: ProjectTrustContext, +): Promise<{ result?: ProjectTrustEventResult; errors: ExtensionError[] }> { + const errors: ExtensionError[] = []; + for (const ext of extensionsResult.extensions) { + // A single extension may register multiple handlers for the same event. + // The first project_trust handler that returns yes/no wins; undecided falls through. + const handlers = ext.handlers.get("project_trust"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const handlerResult = (await handler(event, ctx)) as ProjectTrustEventResult; + if (handlerResult.trusted === "undecided") { + continue; + } + return { result: handlerResult, errors }; + } catch (error) { + errors.push({ + extensionPath: ext.path, + event: event.type, + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + }); + } + } + } + return { errors }; +} + +const noOpUIContext: ExtensionUIContext = { + select: async () => undefined, + confirm: async () => false, + input: async () => undefined, + notify: () => {}, + onTerminalInput: () => () => {}, + setStatus: () => {}, + setWorkingMessage: () => {}, + setWorkingVisible: () => {}, + setWorkingIndicator: () => {}, + setHiddenThinkingLabel: () => {}, + setWidget: () => {}, + setFooter: () => {}, + setHeader: () => {}, + setTitle: () => {}, + custom: async () => undefined as never, + pasteToEditor: () => {}, + setEditorText: () => {}, + getEditorText: () => "", + editor: async () => undefined, + addAutocompleteProvider: () => {}, + setEditorComponent: () => {}, + getEditorComponent: () => undefined, + get theme() { + return theme; + }, + getAllThemes: () => [], + getTheme: () => undefined, + setTheme: (_theme: string | Theme) => ({ success: false, error: "UI not available" }), + getToolsExpanded: () => false, + setToolsExpanded: () => {}, +}; + +export class ExtensionRunner { + private extensions: Extension[]; + private runtime: ExtensionRuntime; + private uiContext: ExtensionUIContext; + private mode: ExtensionMode = "print"; + private cwd: string; + private sessionManager: SessionManager; + private modelRegistry: ModelRegistry; + private errorListeners: Set = new Set(); + private getModel: () => Model | undefined = () => undefined; + private isIdleFn: () => boolean = () => true; + private isProjectTrustedFn: () => boolean = () => true; + private getSignalFn: () => AbortSignal | undefined = () => undefined; + private waitForIdleFn: () => Promise = async () => {}; + private abortFn: () => void = () => {}; + private hasPendingMessagesFn: () => boolean = () => false; + private getContextUsageFn: () => ContextUsage | undefined = () => undefined; + private compactFn: (options?: CompactOptions) => void = () => {}; + private getSystemPromptFn: () => string = () => ""; + private getSystemPromptOptionsFn: () => BuildSystemPromptOptions = () => ({ cwd: this.cwd }); + private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false }); + private forkHandler: ForkHandler = async () => ({ cancelled: false }); + private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false }); + private switchSessionHandler: SwitchSessionHandler = async () => ({ cancelled: false }); + private reloadHandler: ReloadHandler = async () => {}; + private shutdownHandler: ShutdownHandler = () => {}; + private shortcutDiagnostics: ResourceDiagnostic[] = []; + private commandDiagnostics: ResourceDiagnostic[] = []; + private staleMessage: string | undefined; + + constructor( + extensions: Extension[], + runtime: ExtensionRuntime, + cwd: string, + sessionManager: SessionManager, + modelRegistry: ModelRegistry, + ) { + this.extensions = extensions; + this.runtime = runtime; + this.uiContext = noOpUIContext; + this.cwd = cwd; + this.sessionManager = sessionManager; + this.modelRegistry = modelRegistry; + } + + bindCore( + actions: ExtensionActions, + contextActions: ExtensionContextActions, + providerActions?: { + registerProvider?: (name: string, config: ProviderConfig) => void; + unregisterProvider?: (name: string) => void; + }, + ): void { + // Copy actions into the shared runtime (all extension APIs reference this) + this.runtime.sendMessage = actions.sendMessage; + this.runtime.sendUserMessage = actions.sendUserMessage; + this.runtime.appendEntry = actions.appendEntry; + this.runtime.setSessionName = actions.setSessionName; + this.runtime.getSessionName = actions.getSessionName; + this.runtime.setLabel = actions.setLabel; + this.runtime.getActiveTools = actions.getActiveTools; + this.runtime.getAllTools = actions.getAllTools; + this.runtime.setActiveTools = actions.setActiveTools; + this.runtime.refreshTools = actions.refreshTools; + this.runtime.getCommands = actions.getCommands; + this.runtime.setModel = actions.setModel; + this.runtime.getThinkingLevel = actions.getThinkingLevel; + this.runtime.setThinkingLevel = actions.setThinkingLevel; + + // Context actions (required) + this.getModel = contextActions.getModel; + this.isIdleFn = contextActions.isIdle; + this.isProjectTrustedFn = contextActions.isProjectTrusted; + this.getSignalFn = contextActions.getSignal; + this.abortFn = contextActions.abort; + this.hasPendingMessagesFn = contextActions.hasPendingMessages; + this.shutdownHandler = contextActions.shutdown; + this.getContextUsageFn = contextActions.getContextUsage; + this.compactFn = contextActions.compact; + this.getSystemPromptFn = contextActions.getSystemPrompt; + this.getSystemPromptOptionsFn = contextActions.getSystemPromptOptions ?? (() => ({ cwd: this.cwd })); + + // Flush provider registrations queued during extension loading + for (const { name, config, extensionPath } of this.runtime.pendingProviderRegistrations) { + try { + if (providerActions?.registerProvider) { + providerActions.registerProvider(name, config); + } else { + this.modelRegistry.registerProvider(name, config); + } + } catch (err) { + this.emitError({ + extensionPath, + event: "register_provider", + error: err instanceof Error ? err.message : String(err), + stack: err instanceof Error ? err.stack : undefined, + }); + } + } + this.runtime.pendingProviderRegistrations = []; + + // From this point on, provider registration/unregistration takes effect immediately + // without requiring a /reload. + this.runtime.registerProvider = (name, config) => { + if (providerActions?.registerProvider) { + providerActions.registerProvider(name, config); + return; + } + this.modelRegistry.registerProvider(name, config); + }; + this.runtime.unregisterProvider = (name) => { + if (providerActions?.unregisterProvider) { + providerActions.unregisterProvider(name); + return; + } + this.modelRegistry.unregisterProvider(name); + }; + } + + bindCommandContext(actions?: ExtensionCommandContextActions): void { + if (actions) { + this.waitForIdleFn = actions.waitForIdle; + this.newSessionHandler = actions.newSession; + this.forkHandler = actions.fork; + this.navigateTreeHandler = actions.navigateTree; + this.switchSessionHandler = actions.switchSession; + this.reloadHandler = actions.reload; + return; + } + + this.waitForIdleFn = async () => {}; + this.newSessionHandler = async () => ({ cancelled: false }); + this.forkHandler = async () => ({ cancelled: false }); + this.navigateTreeHandler = async () => ({ cancelled: false }); + this.switchSessionHandler = async () => ({ cancelled: false }); + this.reloadHandler = async () => {}; + } + + setUIContext(uiContext?: ExtensionUIContext, mode: ExtensionMode = "print"): void { + this.uiContext = uiContext ?? noOpUIContext; + this.mode = mode; + } + + getUIContext(): ExtensionUIContext { + return this.uiContext; + } + + hasUI(): boolean { + return this.uiContext !== noOpUIContext; + } + + getExtensionPaths(): string[] { + return this.extensions.map((e) => e.path); + } + + /** Get all registered tools from all extensions (first registration per name wins). */ + getAllRegisteredTools(): RegisteredTool[] { + const toolsByName = new Map(); + for (const ext of this.extensions) { + for (const tool of ext.tools.values()) { + if (!toolsByName.has(tool.definition.name)) { + toolsByName.set(tool.definition.name, tool); + } + } + } + return Array.from(toolsByName.values()); + } + + /** Get a tool definition by name. Returns undefined if not found. */ + getToolDefinition(toolName: string): RegisteredTool["definition"] | undefined { + for (const ext of this.extensions) { + const tool = ext.tools.get(toolName); + if (tool) { + return tool.definition; + } + } + return undefined; + } + + getFlags(): Map { + const allFlags = new Map(); + for (const ext of this.extensions) { + for (const [name, flag] of ext.flags) { + if (!allFlags.has(name)) { + allFlags.set(name, flag); + } + } + } + return allFlags; + } + + setFlagValue(name: string, value: boolean | string): void { + this.runtime.flagValues.set(name, value); + } + + getFlagValues(): Map { + return new Map(this.runtime.flagValues); + } + + getShortcuts(resolvedKeybindings: KeybindingsConfig): Map { + this.shortcutDiagnostics = []; + const builtinKeybindings = buildBuiltinKeybindings(resolvedKeybindings); + const extensionShortcuts = new Map(); + + const addDiagnostic = (message: string, extensionPath: string) => { + this.shortcutDiagnostics.push({ type: "warning", message, path: extensionPath }); + if (!this.hasUI()) { + console.warn(message); + } + }; + + for (const ext of this.extensions) { + for (const [key, shortcut] of ext.shortcuts) { + const normalizedKey = key.toLowerCase() as KeyId; + + const builtInKeybinding = builtinKeybindings[normalizedKey]; + if (builtInKeybinding?.restrictOverride === true) { + addDiagnostic( + `Extension shortcut '${key}' from ${shortcut.extensionPath} conflicts with built-in shortcut. Skipping.`, + shortcut.extensionPath, + ); + continue; + } + + if (builtInKeybinding?.restrictOverride === false) { + addDiagnostic( + `Extension shortcut conflict: '${key}' is built-in shortcut for ${builtInKeybinding.keybinding} and ${shortcut.extensionPath}. Using ${shortcut.extensionPath}.`, + shortcut.extensionPath, + ); + } + + const existingExtensionShortcut = extensionShortcuts.get(normalizedKey); + if (existingExtensionShortcut) { + addDiagnostic( + `Extension shortcut conflict: '${key}' registered by both ${existingExtensionShortcut.extensionPath} and ${shortcut.extensionPath}. Using ${shortcut.extensionPath}.`, + shortcut.extensionPath, + ); + } + extensionShortcuts.set(normalizedKey, shortcut); + } + } + return extensionShortcuts; + } + + getShortcutDiagnostics(): ResourceDiagnostic[] { + return this.shortcutDiagnostics; + } + + invalidate( + message = "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", + ): void { + if (!this.staleMessage) { + this.staleMessage = message; + this.runtime.invalidate(message); + } + } + + private assertActive(): void { + if (this.staleMessage) { + throw new Error(this.staleMessage); + } + } + + onError(listener: ExtensionErrorListener): () => void { + this.errorListeners.add(listener); + return () => this.errorListeners.delete(listener); + } + + emitError(error: ExtensionError): void { + for (const listener of this.errorListeners) { + listener(error); + } + } + + hasHandlers(eventType: string): boolean { + for (const ext of this.extensions) { + const handlers = ext.handlers.get(eventType); + if (handlers && handlers.length > 0) { + return true; + } + } + return false; + } + + getMessageRenderer(customType: string): MessageRenderer | undefined { + for (const ext of this.extensions) { + const renderer = ext.messageRenderers.get(customType); + if (renderer) { + return renderer; + } + } + return undefined; + } + + private resolveRegisteredCommands(): ResolvedCommand[] { + const commands: RegisteredCommand[] = []; + const counts = new Map(); + + for (const ext of this.extensions) { + for (const command of ext.commands.values()) { + commands.push(command); + counts.set(command.name, (counts.get(command.name) ?? 0) + 1); + } + } + + const seen = new Map(); + const takenInvocationNames = new Set(); + + return commands.map((command) => { + const occurrence = (seen.get(command.name) ?? 0) + 1; + seen.set(command.name, occurrence); + + let invocationName = (counts.get(command.name) ?? 0) > 1 ? `${command.name}:${occurrence}` : command.name; + + if (takenInvocationNames.has(invocationName)) { + let suffix = occurrence; + do { + suffix++; + invocationName = `${command.name}:${suffix}`; + } while (takenInvocationNames.has(invocationName)); + } + + takenInvocationNames.add(invocationName); + return { + ...command, + invocationName, + }; + }); + } + + getRegisteredCommands(): ResolvedCommand[] { + this.commandDiagnostics = []; + return this.resolveRegisteredCommands(); + } + + getCommandDiagnostics(): ResourceDiagnostic[] { + return this.commandDiagnostics; + } + + getCommand(name: string): ResolvedCommand | undefined { + return this.resolveRegisteredCommands().find((command) => command.invocationName === name); + } + + /** + * Request a graceful shutdown. Called by extension tools and event handlers. + * The actual shutdown behavior is provided by the mode via bindExtensions(). + */ + shutdown(): void { + this.shutdownHandler(); + } + + /** + * Create an ExtensionContext for use in event handlers and tool execution. + * Context values are resolved at call time, so changes via bindCore/bindUI are reflected. + */ + createContext(): ExtensionContext { + const runner = this; + const getModel = this.getModel; + return { + get ui() { + runner.assertActive(); + return runner.uiContext; + }, + get mode() { + runner.assertActive(); + return runner.mode; + }, + get hasUI() { + runner.assertActive(); + return runner.hasUI(); + }, + get cwd() { + runner.assertActive(); + return runner.cwd; + }, + get sessionManager() { + runner.assertActive(); + return runner.sessionManager; + }, + get modelRegistry() { + runner.assertActive(); + return runner.modelRegistry; + }, + get model() { + runner.assertActive(); + return getModel(); + }, + isIdle: () => { + runner.assertActive(); + return runner.isIdleFn(); + }, + isProjectTrusted: () => { + runner.assertActive(); + return runner.isProjectTrustedFn(); + }, + get signal() { + runner.assertActive(); + return runner.getSignalFn(); + }, + abort: () => { + runner.assertActive(); + runner.abortFn(); + }, + hasPendingMessages: () => { + runner.assertActive(); + return runner.hasPendingMessagesFn(); + }, + shutdown: () => { + runner.assertActive(); + runner.shutdownHandler(); + }, + getContextUsage: () => { + runner.assertActive(); + return runner.getContextUsageFn(); + }, + compact: (options) => { + runner.assertActive(); + runner.compactFn(options); + }, + getSystemPrompt: () => { + runner.assertActive(); + return runner.getSystemPromptFn(); + }, + }; + } + + createCommandContext(): ExtensionCommandContext { + // Use property descriptors instead of object spread so the guarded getters from + // createContext() stay lazy. A spread would eagerly read them once and freeze the + // old values into the returned object, bypassing stale-instance checks. + const context = Object.defineProperties( + {}, + Object.getOwnPropertyDescriptors(this.createContext()), + ) as ExtensionCommandContext; + context.getSystemPromptOptions = () => { + this.assertActive(); + return this.getSystemPromptOptionsFn(); + }; + context.waitForIdle = () => { + this.assertActive(); + return this.waitForIdleFn(); + }; + context.newSession = (options) => { + this.assertActive(); + return this.newSessionHandler(options); + }; + context.fork = (entryId, options) => { + this.assertActive(); + return this.forkHandler(entryId, options); + }; + context.navigateTree = (targetId, options) => { + this.assertActive(); + return this.navigateTreeHandler(targetId, options); + }; + context.switchSession = (sessionPath, options) => { + this.assertActive(); + return this.switchSessionHandler(sessionPath, options); + }; + context.reload = () => { + this.assertActive(); + return this.reloadHandler(); + }; + return context; + } + + private isSessionBeforeEvent(event: RunnerEmitEvent): event is SessionBeforeEvent { + return ( + event.type === "session_before_switch" || + event.type === "session_before_fork" || + event.type === "session_before_compact" || + event.type === "session_before_tree" + ); + } + + async emit(event: TEvent): Promise> { + const ctx = this.createContext(); + let result: SessionBeforeEventResult | undefined; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get(event.type); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const handlerResult = await handler(event, ctx); + + if (this.isSessionBeforeEvent(event) && handlerResult) { + result = handlerResult as SessionBeforeEventResult; + if (result.cancel) { + return result as RunnerEmitResult; + } + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: event.type, + error: message, + stack, + }); + } + } + } + + return result as RunnerEmitResult; + } + + async emitMessageEnd(event: MessageEndEvent): Promise { + const ctx = this.createContext(); + let currentMessage = event.message; + let modified = false; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("message_end"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const currentEvent: MessageEndEvent = { ...event, message: currentMessage }; + const handlerResult = (await handler(currentEvent, ctx)) as MessageEndEventResult | undefined; + if (!handlerResult?.message) continue; + + if (handlerResult.message.role !== currentMessage.role) { + this.emitError({ + extensionPath: ext.path, + event: "message_end", + error: "message_end handlers must return a message with the same role", + }); + continue; + } + + currentMessage = handlerResult.message; + modified = true; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "message_end", + error: message, + stack, + }); + } + } + } + + return modified ? currentMessage : undefined; + } + + async emitToolResult(event: ToolResultEvent): Promise { + const ctx = this.createContext(); + const currentEvent: ToolResultEvent = { ...event }; + let modified = false; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("tool_result"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const handlerResult = (await handler(currentEvent, ctx)) as ToolResultEventResult | undefined; + if (!handlerResult) continue; + + if (handlerResult.content !== undefined) { + currentEvent.content = handlerResult.content; + modified = true; + } + if (handlerResult.details !== undefined) { + currentEvent.details = handlerResult.details; + modified = true; + } + if (handlerResult.isError !== undefined) { + currentEvent.isError = handlerResult.isError; + modified = true; + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "tool_result", + error: message, + stack, + }); + } + } + } + + if (!modified) { + return undefined; + } + + return { + content: currentEvent.content, + details: currentEvent.details, + isError: currentEvent.isError, + }; + } + + async emitToolCall(event: ToolCallEvent): Promise { + const ctx = this.createContext(); + let result: ToolCallEventResult | undefined; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("tool_call"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + const handlerResult = await handler(event, ctx); + + if (handlerResult) { + result = handlerResult as ToolCallEventResult; + if (result.block) { + return result; + } + } + } + } + + return result; + } + + async emitUserBash(event: UserBashEvent): Promise { + const ctx = this.createContext(); + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("user_bash"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const handlerResult = await handler(event, ctx); + if (handlerResult) { + return handlerResult as UserBashEventResult; + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "user_bash", + error: message, + stack, + }); + } + } + } + + return undefined; + } + + async emitContext(messages: AgentMessage[]): Promise { + const ctx = this.createContext(); + let currentMessages = structuredClone(messages); + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("context"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const event: ContextEvent = { type: "context", messages: currentMessages }; + const handlerResult = await handler(event, ctx); + + if (handlerResult && (handlerResult as ContextEventResult).messages) { + currentMessages = (handlerResult as ContextEventResult).messages!; + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "context", + error: message, + stack, + }); + } + } + } + + return currentMessages; + } + + async emitBeforeProviderRequest(payload: unknown): Promise { + const ctx = this.createContext(); + let currentPayload = payload; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("before_provider_request"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const event: BeforeProviderRequestEvent = { + type: "before_provider_request", + payload: currentPayload, + }; + const handlerResult = await handler(event, ctx); + if (handlerResult !== undefined) { + currentPayload = handlerResult; + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "before_provider_request", + error: message, + stack, + }); + } + } + } + + return currentPayload; + } + + async emitBeforeAgentStart( + prompt: string, + images: ImageContent[] | undefined, + systemPrompt: string, + systemPromptOptions: BuildSystemPromptOptions, + ): Promise { + let currentSystemPrompt = systemPrompt; + const ctx = Object.defineProperties( + {}, + Object.getOwnPropertyDescriptors(this.createContext()), + ) as ExtensionContext; + ctx.getSystemPrompt = () => { + this.assertActive(); + return currentSystemPrompt; + }; + const messages: NonNullable[] = []; + let systemPromptModified = false; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("before_agent_start"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const event: BeforeAgentStartEvent = { + type: "before_agent_start", + prompt, + images, + systemPrompt: currentSystemPrompt, + systemPromptOptions, + }; + const handlerResult = await handler(event, ctx); + + if (handlerResult) { + const result = handlerResult as BeforeAgentStartEventResult; + if (result.message) { + messages.push(result.message); + } + if (result.systemPrompt !== undefined) { + currentSystemPrompt = result.systemPrompt; + systemPromptModified = true; + } + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "before_agent_start", + error: message, + stack, + }); + } + } + } + + if (messages.length > 0 || systemPromptModified) { + return { + messages: messages.length > 0 ? messages : undefined, + systemPrompt: systemPromptModified ? currentSystemPrompt : undefined, + }; + } + + return undefined; + } + + async emitResourcesDiscover( + cwd: string, + reason: ResourcesDiscoverEvent["reason"], + ): Promise<{ + skillPaths: Array<{ path: string; extensionPath: string }>; + promptPaths: Array<{ path: string; extensionPath: string }>; + themePaths: Array<{ path: string; extensionPath: string }>; + }> { + const ctx = this.createContext(); + const skillPaths: Array<{ path: string; extensionPath: string }> = []; + const promptPaths: Array<{ path: string; extensionPath: string }> = []; + const themePaths: Array<{ path: string; extensionPath: string }> = []; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("resources_discover"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const event: ResourcesDiscoverEvent = { type: "resources_discover", cwd, reason }; + const handlerResult = await handler(event, ctx); + const result = handlerResult as ResourcesDiscoverResult | undefined; + + if (result?.skillPaths?.length) { + skillPaths.push(...result.skillPaths.map((path) => ({ path, extensionPath: ext.path }))); + } + if (result?.promptPaths?.length) { + promptPaths.push(...result.promptPaths.map((path) => ({ path, extensionPath: ext.path }))); + } + if (result?.themePaths?.length) { + themePaths.push(...result.themePaths.map((path) => ({ path, extensionPath: ext.path }))); + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "resources_discover", + error: message, + stack, + }); + } + } + } + + return { skillPaths, promptPaths, themePaths }; + } + + /** Emit input event. Transforms chain, "handled" short-circuits. */ + async emitInput( + text: string, + images: ImageContent[] | undefined, + source: InputSource, + streamingBehavior?: "steer" | "followUp", + ): Promise { + const ctx = this.createContext(); + let currentText = text; + let currentImages = images; + + for (const ext of this.extensions) { + for (const handler of ext.handlers.get("input") ?? []) { + try { + const event: InputEvent = { + type: "input", + text: currentText, + images: currentImages, + source, + streamingBehavior, + }; + const result = (await handler(event, ctx)) as InputEventResult | undefined; + if (result?.action === "handled") return result; + if (result?.action === "transform") { + currentText = result.text; + currentImages = result.images ?? currentImages; + } + } catch (err) { + this.emitError({ + extensionPath: ext.path, + event: "input", + error: err instanceof Error ? err.message : String(err), + stack: err instanceof Error ? err.stack : undefined, + }); + } + } + } + return currentText !== text || currentImages !== images + ? { action: "transform", text: currentText, images: currentImages } + : { action: "continue" }; + } +} diff --git a/cactus-code/packages/coding-agent/src/core/extensions/types.ts b/cactus-code/packages/coding-agent/src/core/extensions/types.ts new file mode 100644 index 000000000..7234d4e4f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/extensions/types.ts @@ -0,0 +1,1614 @@ +/** + * Extension system types. + * + * Extensions are TypeScript modules that can: + * - Subscribe to agent lifecycle events + * - Register LLM-callable tools + * - Register commands, keyboard shortcuts, and CLI flags + * - Interact with the user via UI primitives + */ + +import type { + AgentMessage, + AgentToolResult, + AgentToolUpdateCallback, + ThinkingLevel, + ToolExecutionMode, +} from "@earendil-works/pi-agent-core"; +import type { + Api, + AssistantMessageEvent, + AssistantMessageEventStream, + Context, + ImageContent, + Model, + OAuthCredentials, + OAuthLoginCallbacks, + SimpleStreamOptions, + TextContent, + ToolResultMessage, +} from "@earendil-works/pi-ai"; +import type { + AutocompleteItem, + AutocompleteProvider, + Component, + EditorComponent, + EditorTheme, + KeyId, + OverlayHandle, + OverlayOptions, + TUI, +} from "@earendil-works/pi-tui"; +import type { Static, TSchema } from "typebox"; +import type { Theme } from "../../modes/interactive/theme/theme.ts"; +import type { BashResult } from "../bash-executor.ts"; +import type { CompactionPreparation, CompactionResult } from "../compaction/index.ts"; +import type { EventBus } from "../event-bus.ts"; +import type { ExecOptions, ExecResult } from "../exec.ts"; +import type { ReadonlyFooterDataProvider } from "../footer-data-provider.ts"; +import type { KeybindingsManager } from "../keybindings.ts"; +import type { CustomMessage } from "../messages.ts"; +import type { ModelRegistry } from "../model-registry.ts"; +import type { + BranchSummaryEntry, + CompactionEntry, + ReadonlySessionManager, + SessionEntry, + SessionManager, +} from "../session-manager.ts"; +import type { SlashCommandInfo } from "../slash-commands.ts"; +import type { SourceInfo } from "../source-info.ts"; +import type { BuildSystemPromptOptions } from "../system-prompt.ts"; +import type { BashOperations } from "../tools/bash.ts"; +import type { EditToolDetails } from "../tools/edit.ts"; +import type { + BashToolDetails, + BashToolInput, + EditToolInput, + FindToolDetails, + FindToolInput, + GrepToolDetails, + GrepToolInput, + LsToolDetails, + LsToolInput, + ReadToolDetails, + ReadToolInput, + WriteToolInput, +} from "../tools/index.ts"; + +export type { ExecOptions, ExecResult } from "../exec.ts"; +export type { BuildSystemPromptOptions } from "../system-prompt.ts"; +export type { AgentToolResult, AgentToolUpdateCallback, ToolExecutionMode }; +export type { AppKeybinding, KeybindingsManager } from "../keybindings.ts"; + +// ============================================================================ +// UI Context +// ============================================================================ + +/** Options for extension UI dialogs. */ +export interface ExtensionUIDialogOptions { + /** AbortSignal to programmatically dismiss the dialog. */ + signal?: AbortSignal; + /** Timeout in milliseconds. Dialog auto-dismisses with live countdown display. */ + timeout?: number; +} + +/** Placement for extension widgets. */ +export type WidgetPlacement = "aboveEditor" | "belowEditor"; + +/** Options for extension widgets. */ +export interface ExtensionWidgetOptions { + /** Where the widget is rendered. Defaults to "aboveEditor". */ + placement?: WidgetPlacement; +} + +/** Raw terminal input listener for extensions. */ +export type TerminalInputHandler = (data: string) => { consume?: boolean; data?: string } | undefined; + +/** Working indicator configuration for the interactive streaming loader. */ +export interface WorkingIndicatorOptions { + /** Animation frames. Use an empty array to hide the indicator entirely. Custom frames are rendered verbatim. */ + frames?: string[]; + /** Frame interval in milliseconds for animated indicators. */ + intervalMs?: number; +} + +/** Wrap the current autocomplete provider with additional behavior. */ +export type AutocompleteProviderFactory = (current: AutocompleteProvider) => AutocompleteProvider; +export type EditorFactory = (tui: TUI, theme: EditorTheme, keybindings: KeybindingsManager) => EditorComponent; + +/** + * UI context for extensions to request interactive UI. + * Each mode (interactive, RPC, print) provides its own implementation. + */ +export interface ExtensionUIContext { + /** Show a selector and return the user's choice. */ + select(title: string, options: string[], opts?: ExtensionUIDialogOptions): Promise; + + /** Show a confirmation dialog. */ + confirm(title: string, message: string, opts?: ExtensionUIDialogOptions): Promise; + + /** Show a text input dialog. */ + input(title: string, placeholder?: string, opts?: ExtensionUIDialogOptions): Promise; + + /** Show a notification to the user. */ + notify(message: string, type?: "info" | "warning" | "error"): void; + + /** Listen to raw terminal input (interactive mode only). Returns an unsubscribe function. */ + onTerminalInput(handler: TerminalInputHandler): () => void; + + /** Set status text in the footer/status bar. Pass undefined to clear. */ + setStatus(key: string, text: string | undefined): void; + + /** Set the working/loading message shown during streaming. Call with no argument to restore default. */ + setWorkingMessage(message?: string): void; + + /** Show or hide the built-in interactive working loader row during streaming. */ + setWorkingVisible(visible: boolean): void; + + /** + * Configure the interactive working indicator shown during streaming. + * + * - Omit the argument to restore the default animated spinner. + * - Use `frames: ["●"]` for a static indicator. + * - Use `frames: []` to hide the indicator entirely. + * - Custom frames are rendered as provided, so extensions must add their own colors. + */ + setWorkingIndicator(options?: WorkingIndicatorOptions): void; + + /** Set the label shown for hidden thinking blocks. Call with no argument to restore default. */ + setHiddenThinkingLabel(label?: string): void; + + /** Set a widget to display above or below the editor. Accepts string array or component factory. */ + setWidget(key: string, content: string[] | undefined, options?: ExtensionWidgetOptions): void; + setWidget( + key: string, + content: ((tui: TUI, theme: Theme) => Component & { dispose?(): void }) | undefined, + options?: ExtensionWidgetOptions, + ): void; + + /** Set a custom footer component, or undefined to restore the built-in footer. + * + * The factory receives a FooterDataProvider for data not otherwise accessible: + * git branch and extension statuses from setStatus(). Token stats, model info, + * etc. are available via ctx.sessionManager and ctx.model. + */ + setFooter( + factory: + | ((tui: TUI, theme: Theme, footerData: ReadonlyFooterDataProvider) => Component & { dispose?(): void }) + | undefined, + ): void; + + /** Set a custom header component (shown at startup, above chat), or undefined to restore the built-in header. */ + setHeader(factory: ((tui: TUI, theme: Theme) => Component & { dispose?(): void }) | undefined): void; + + /** Set the terminal window/tab title. */ + setTitle(title: string): void; + + /** Show a custom component with keyboard focus. */ + custom( + factory: ( + tui: TUI, + theme: Theme, + keybindings: KeybindingsManager, + done: (result: T) => void, + ) => (Component & { dispose?(): void }) | Promise, + options?: { + overlay?: boolean; + /** Overlay positioning/sizing options. Can be static or a function for dynamic updates. */ + overlayOptions?: OverlayOptions | (() => OverlayOptions); + /** Called with the overlay handle after the overlay is shown. Use to control visibility. */ + onHandle?: (handle: OverlayHandle) => void; + }, + ): Promise; + + /** Paste text into the editor, triggering paste handling (collapse for large content). */ + pasteToEditor(text: string): void; + + /** Set the text in the core input editor. */ + setEditorText(text: string): void; + + /** Get the current text from the core input editor. */ + getEditorText(): string; + + /** Show a multi-line editor for text editing. */ + editor(title: string, prefill?: string): Promise; + + /** Stack additional autocomplete behavior on top of the built-in provider. */ + addAutocompleteProvider(factory: AutocompleteProviderFactory): void; + + /** + * Set a custom editor component via factory function. + * Pass undefined to restore the default editor. + * + * The factory receives: + * - `theme`: EditorTheme for styling borders and autocomplete + * - `keybindings`: KeybindingsManager for app-level keybindings + * + * For full app keybinding support (escape, ctrl+d, model switching, etc.), + * extend `CustomEditor` from `@earendil-works/pi-coding-agent` and call + * `super.handleInput(data)` for keys you don't handle. + * + * @example + * ```ts + * import { CustomEditor } from "@earendil-works/pi-coding-agent"; + * + * class VimEditor extends CustomEditor { + * private mode: "normal" | "insert" = "insert"; + * + * handleInput(data: string): void { + * if (this.mode === "normal") { + * // Handle vim normal mode keys... + * if (data === "i") { this.mode = "insert"; return; } + * } + * super.handleInput(data); // App keybindings + text editing + * } + * } + * + * ctx.ui.setEditorComponent((tui, theme, keybindings) => + * new VimEditor(tui, theme, keybindings) + * ); + * ``` + */ + setEditorComponent(factory: EditorFactory | undefined): void; + + /** Get the currently configured custom editor factory, or undefined when using the default editor. */ + getEditorComponent(): EditorFactory | undefined; + + /** Get the current theme for styling. */ + readonly theme: Theme; + + /** Get all available themes with their names and file paths. */ + getAllThemes(): { name: string; path: string | undefined }[]; + + /** Load a theme by name without switching to it. Returns undefined if not found. */ + getTheme(name: string): Theme | undefined; + + /** Set the current theme by name or Theme object. */ + setTheme(theme: string | Theme): { success: boolean; error?: string }; + + /** Get current tool output expansion state. */ + getToolsExpanded(): boolean; + + /** Set tool output expansion state. */ + setToolsExpanded(expanded: boolean): void; +} + +// ============================================================================ +// Extension Context +// ============================================================================ + +export interface ContextUsage { + /** Estimated context tokens, or null if unknown (e.g. right after compaction, before next LLM response). */ + tokens: number | null; + contextWindow: number; + /** Context usage as percentage of context window, or null if tokens is unknown. */ + percent: number | null; +} + +export interface CompactOptions { + customInstructions?: string; + onComplete?: (result: CompactionResult) => void; + onError?: (error: Error) => void; +} + +/** + * Context passed to extension event handlers. + */ +export type ExtensionMode = "tui" | "rpc" | "json" | "print"; + +export interface ExtensionContext { + /** UI methods for user interaction */ + ui: ExtensionUIContext; + /** Current run mode. Use "tui" to guard terminal-only UI such as custom components. */ + mode: ExtensionMode; + /** Whether dialog-capable UI is available (true in TUI and RPC modes) */ + hasUI: boolean; + /** Current working directory */ + cwd: string; + /** Session manager (read-only) */ + sessionManager: ReadonlySessionManager; + /** Model registry for API key resolution */ + modelRegistry: ModelRegistry; + /** Current model (may be undefined) */ + model: Model | undefined; + /** Whether the agent is idle (not streaming) */ + isIdle(): boolean; + /** Whether project-local trust is active for this context. */ + isProjectTrusted(): boolean; + /** The current abort signal, or undefined when the agent is not streaming. */ + signal: AbortSignal | undefined; + /** Abort the current agent operation */ + abort(): void; + /** Whether there are queued messages waiting */ + hasPendingMessages(): boolean; + /** Gracefully shutdown pi and exit. Available in all contexts. */ + shutdown(): void; + /** Get current context usage for the active model. */ + getContextUsage(): ContextUsage | undefined; + /** Trigger compaction without awaiting completion. */ + compact(options?: CompactOptions): void; + /** Get the current effective system prompt. */ + getSystemPrompt(): string; +} + +/** + * Extended context for command handlers. + * Includes session control methods only safe in user-initiated commands. + */ +export interface ExtensionCommandContext extends ExtensionContext { + /** Get the current base system-prompt construction options. */ + getSystemPromptOptions(): BuildSystemPromptOptions; + + /** Wait for the agent to finish streaming */ + waitForIdle(): Promise; + + /** Start a new session, optionally with initialization. */ + newSession(options?: { + parentSession?: string; + setup?: (sessionManager: SessionManager) => Promise; + withSession?: (ctx: ReplacedSessionContext) => Promise; + }): Promise<{ cancelled: boolean }>; + + /** Fork from a specific entry, creating a new session file. */ + fork( + entryId: string, + options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise }, + ): Promise<{ cancelled: boolean }>; + + /** Navigate to a different point in the session tree. */ + navigateTree( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ): Promise<{ cancelled: boolean }>; + + /** Switch to a different session file. */ + switchSession( + sessionPath: string, + options?: { withSession?: (ctx: ReplacedSessionContext) => Promise }, + ): Promise<{ cancelled: boolean }>; + + /** Reload extensions, skills, prompts, and themes. */ + reload(): Promise; +} + +/** + * Fresh command-capable context bound to the replacement session after a session switch. + * + * This is passed to `withSession()` callbacks on `newSession()`, `fork()`, and `switchSession()`. + */ +export interface ReplacedSessionContext extends ExtensionCommandContext { + sendMessage( + message: Pick, "customType" | "content" | "display" | "details">, + options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, + ): Promise; + + sendUserMessage( + content: string | (TextContent | ImageContent)[], + options?: { deliverAs?: "steer" | "followUp" }, + ): Promise; +} + +// ============================================================================ +// Tool Types +// ============================================================================ + +/** Rendering options for tool results */ +export interface ToolRenderResultOptions { + /** Whether the result view is expanded */ + expanded: boolean; + /** Whether this is a partial/streaming result */ + isPartial: boolean; +} + +/** Context passed to tool renderers. */ +export interface ToolRenderContext { + /** Current tool call arguments. Shared across call/result renders for the same tool call. */ + args: TArgs; + /** Unique id for this tool execution. Stable across call/result renders for the same tool call. */ + toolCallId: string; + /** Invalidate just this tool execution component for redraw. */ + invalidate: () => void; + /** Previously returned component for this render slot, if any. */ + lastComponent: Component | undefined; + /** Shared renderer state for this tool row. Initialized by tool-execution.ts. */ + state: TState; + /** Working directory for this tool execution. */ + cwd: string; + /** Whether the tool execution has started. */ + executionStarted: boolean; + /** Whether the tool call arguments are complete. */ + argsComplete: boolean; + /** Whether the tool result is partial/streaming. */ + isPartial: boolean; + /** Whether the result view is expanded. */ + expanded: boolean; + /** Whether inline images are currently shown in the TUI. */ + showImages: boolean; + /** Whether the current result is an error. */ + isError: boolean; +} + +/** + * Tool definition for registerTool(). + */ +export interface ToolDefinition { + /** Tool name (used in LLM tool calls) */ + name: string; + /** Human-readable label for UI */ + label: string; + /** Description for LLM */ + description: string; + /** Optional one-line snippet for the Available tools section in the default system prompt. Custom tools are omitted from that section when this is not provided. */ + promptSnippet?: string; + /** Optional guideline bullets appended to the default system prompt Guidelines section when this tool is active. */ + promptGuidelines?: string[]; + /** Parameter schema (TypeBox) */ + parameters: TParams; + /** Controls whether ToolExecutionComponent renders the standard colored shell or the tool renders its own framing. */ + renderShell?: "default" | "self"; + + /** Optional compatibility shim to prepare raw tool call arguments before schema validation. Must return an object conforming to TParams. */ + prepareArguments?: (args: unknown) => Static; + + /** + * Per-tool execution mode override. + * - "sequential": this tool must execute one at a time with other tool calls. + * - "parallel": this tool can execute concurrently with other tool calls. + * + * If omitted, the default execution mode applies. + */ + executionMode?: ToolExecutionMode; + + /** Execute the tool. */ + execute( + toolCallId: string, + params: Static, + signal: AbortSignal | undefined, + onUpdate: AgentToolUpdateCallback | undefined, + ctx: ExtensionContext, + ): Promise>; + + /** Custom rendering for tool call display */ + renderCall?: (args: Static, theme: Theme, context: ToolRenderContext>) => Component; + + /** Custom rendering for tool result display */ + renderResult?: ( + result: AgentToolResult, + options: ToolRenderResultOptions, + theme: Theme, + context: ToolRenderContext>, + ) => Component; +} + +type AnyToolDefinition = ToolDefinition; + +/** + * Preserve parameter inference for standalone tool definitions. + * + * Use this when assigning a tool to a variable or passing it through arrays such + * as `customTools`, where contextual typing would otherwise widen params to + * `unknown`. + */ +export function defineTool( + tool: ToolDefinition, +): ToolDefinition & AnyToolDefinition { + return tool as ToolDefinition & AnyToolDefinition; +} + +// ============================================================================ +// Startup/Resource Events +// ============================================================================ + +export interface ProjectTrustEvent { + type: "project_trust"; + cwd: string; +} + +export type ProjectTrustEventDecision = "yes" | "no" | "undecided"; + +export interface ProjectTrustEventResult { + trusted: ProjectTrustEventDecision; + remember?: boolean; +} + +export interface ProjectTrustContext { + cwd: string; + mode: ExtensionMode; + hasUI: boolean; + ui: Pick; +} + +export type ProjectTrustHandler = ( + event: ProjectTrustEvent, + ctx: ProjectTrustContext, +) => Promise | ProjectTrustEventResult; + +/** Fired after session_start to allow extensions to provide additional resource paths. */ +export interface ResourcesDiscoverEvent { + type: "resources_discover"; + cwd: string; + reason: "startup" | "reload"; +} + +/** Result from resources_discover event handler */ +export interface ResourcesDiscoverResult { + skillPaths?: string[]; + promptPaths?: string[]; + themePaths?: string[]; +} + +// ============================================================================ +// Session Events +// ============================================================================ + +/** Fired when a session is started, loaded, or reloaded */ +export interface SessionStartEvent { + type: "session_start"; + /** Why this session start happened. */ + reason: "startup" | "reload" | "new" | "resume" | "fork"; + /** Previously active session file. Present for "new", "resume", and "fork". */ + previousSessionFile?: string; +} + +/** Fired before switching to another session (can be cancelled) */ +export interface SessionBeforeSwitchEvent { + type: "session_before_switch"; + reason: "new" | "resume"; + targetSessionFile?: string; +} + +/** Fired before forking a session (can be cancelled) */ +export interface SessionBeforeForkEvent { + type: "session_before_fork"; + entryId: string; + position: "before" | "at"; +} + +/** Fired before context compaction (can be cancelled or customized) */ +export interface SessionBeforeCompactEvent { + type: "session_before_compact"; + preparation: CompactionPreparation; + branchEntries: SessionEntry[]; + customInstructions?: string; + /** What triggered the compaction: manual /compact, the context threshold, or context overflow recovery */ + reason: "manual" | "threshold" | "overflow"; + /** True when the aborted turn is retried after this compaction (overflow recovery) */ + willRetry: boolean; + signal: AbortSignal; +} + +/** Fired after context compaction */ +export interface SessionCompactEvent { + type: "session_compact"; + compactionEntry: CompactionEntry; + fromExtension: boolean; + /** What triggered the compaction: manual /compact, the context threshold, or context overflow recovery */ + reason: "manual" | "threshold" | "overflow"; + /** True when the aborted turn is retried after this compaction (overflow recovery) */ + willRetry: boolean; +} + +/** Fired before an extension runtime is torn down due to quit, reload, or session replacement. */ +export interface SessionShutdownEvent { + type: "session_shutdown"; + reason: "quit" | "reload" | "new" | "resume" | "fork"; + /** Destination session file when shutting down due to session replacement. */ + targetSessionFile?: string; +} + +/** Preparation data for tree navigation */ +export interface TreePreparation { + targetId: string; + oldLeafId: string | null; + commonAncestorId: string | null; + entriesToSummarize: SessionEntry[]; + userWantsSummary: boolean; + /** Custom instructions for summarization */ + customInstructions?: string; + /** If true, customInstructions replaces the default prompt instead of being appended */ + replaceInstructions?: boolean; + /** Label to attach to the branch summary entry */ + label?: string; +} + +/** Fired before navigating in the session tree (can be cancelled) */ +export interface SessionBeforeTreeEvent { + type: "session_before_tree"; + preparation: TreePreparation; + signal: AbortSignal; +} + +/** Fired after navigating in the session tree */ +export interface SessionTreeEvent { + type: "session_tree"; + newLeafId: string | null; + oldLeafId: string | null; + summaryEntry?: BranchSummaryEntry; + fromExtension?: boolean; +} + +export type SessionEvent = + | SessionStartEvent + | SessionBeforeSwitchEvent + | SessionBeforeForkEvent + | SessionBeforeCompactEvent + | SessionCompactEvent + | SessionShutdownEvent + | SessionBeforeTreeEvent + | SessionTreeEvent; + +// ============================================================================ +// Agent Events +// ============================================================================ + +/** Fired before each LLM call. Can modify messages. */ +export interface ContextEvent { + type: "context"; + messages: AgentMessage[]; +} + +/** Fired before a provider request is sent. Can replace the payload. */ +export interface BeforeProviderRequestEvent { + type: "before_provider_request"; + payload: unknown; +} + +/** Fired after a provider response is received and before the response stream is consumed. */ +export interface AfterProviderResponseEvent { + type: "after_provider_response"; + status: number; + headers: Record; +} + +/** Fired after user submits prompt but before agent loop. */ +export interface BeforeAgentStartEvent { + type: "before_agent_start"; + /** The raw user prompt text (after expansion). */ + prompt: string; + /** Images attached to the user prompt, if any. */ + images?: ImageContent[]; + /** The fully assembled system prompt string. */ + systemPrompt: string; + /** Structured options used to build the system prompt. Extensions can inspect this to understand what Pi loaded without re-discovering resources. */ + systemPromptOptions: BuildSystemPromptOptions; +} + +/** Fired when an agent loop starts */ +export interface AgentStartEvent { + type: "agent_start"; +} + +/** Fired when an agent loop ends */ +export interface AgentEndEvent { + type: "agent_end"; + messages: AgentMessage[]; +} + +/** Fired at the start of each turn */ +export interface TurnStartEvent { + type: "turn_start"; + turnIndex: number; + timestamp: number; +} + +/** Fired at the end of each turn */ +export interface TurnEndEvent { + type: "turn_end"; + turnIndex: number; + message: AgentMessage; + toolResults: ToolResultMessage[]; +} + +/** Fired when a message starts (user, assistant, or toolResult) */ +export interface MessageStartEvent { + type: "message_start"; + message: AgentMessage; +} + +/** Fired during assistant message streaming with token-by-token updates */ +export interface MessageUpdateEvent { + type: "message_update"; + message: AgentMessage; + assistantMessageEvent: AssistantMessageEvent; +} + +/** Fired when a message ends */ +export interface MessageEndEvent { + type: "message_end"; + message: AgentMessage; +} + +/** Fired when a tool starts executing */ +export interface ToolExecutionStartEvent { + type: "tool_execution_start"; + toolCallId: string; + toolName: string; + args: any; +} + +/** Fired during tool execution with partial/streaming output */ +export interface ToolExecutionUpdateEvent { + type: "tool_execution_update"; + toolCallId: string; + toolName: string; + args: any; + partialResult: any; +} + +/** Fired when a tool finishes executing */ +export interface ToolExecutionEndEvent { + type: "tool_execution_end"; + toolCallId: string; + toolName: string; + result: any; + isError: boolean; +} + +// ============================================================================ +// Model Events +// ============================================================================ + +export type ModelSelectSource = "set" | "cycle" | "restore"; + +/** Fired when a new model is selected */ +export interface ModelSelectEvent { + type: "model_select"; + model: Model; + previousModel: Model | undefined; + source: ModelSelectSource; +} + +/** Fired when a new thinking level is selected */ +export interface ThinkingLevelSelectEvent { + type: "thinking_level_select"; + level: ThinkingLevel; + previousLevel: ThinkingLevel; +} + +// ============================================================================ +// User Bash Events +// ============================================================================ + +/** Fired when user executes a bash command via ! or !! prefix */ +export interface UserBashEvent { + type: "user_bash"; + /** The command to execute */ + command: string; + /** True if !! prefix was used (excluded from LLM context) */ + excludeFromContext: boolean; + /** Current working directory */ + cwd: string; +} + +// ============================================================================ +// Input Events +// ============================================================================ + +/** Source of user input */ +export type InputSource = "interactive" | "rpc" | "extension"; + +/** Fired when user input is received, before agent processing */ +export interface InputEvent { + type: "input"; + /** The input text */ + text: string; + /** Attached images, if any */ + images?: ImageContent[]; + /** Where the input came from */ + source: InputSource; + /** How the input will be delivered during streaming, or undefined when idle */ + streamingBehavior?: "steer" | "followUp"; +} + +/** Result from input event handler */ +export type InputEventResult = + | { action: "continue" } + | { action: "transform"; text: string; images?: ImageContent[] } + | { action: "handled" }; + +// ============================================================================ +// Tool Events +// ============================================================================ + +interface ToolCallEventBase { + type: "tool_call"; + toolCallId: string; +} + +export interface BashToolCallEvent extends ToolCallEventBase { + toolName: "bash"; + input: BashToolInput; +} + +export interface ReadToolCallEvent extends ToolCallEventBase { + toolName: "read"; + input: ReadToolInput; +} + +export interface EditToolCallEvent extends ToolCallEventBase { + toolName: "edit"; + input: EditToolInput; +} + +export interface WriteToolCallEvent extends ToolCallEventBase { + toolName: "write"; + input: WriteToolInput; +} + +export interface GrepToolCallEvent extends ToolCallEventBase { + toolName: "grep"; + input: GrepToolInput; +} + +export interface FindToolCallEvent extends ToolCallEventBase { + toolName: "find"; + input: FindToolInput; +} + +export interface LsToolCallEvent extends ToolCallEventBase { + toolName: "ls"; + input: LsToolInput; +} + +export interface CustomToolCallEvent extends ToolCallEventBase { + toolName: string; + input: Record; +} + +/** + * Fired before a tool executes. Can block. + * + * `event.input` is mutable. Mutate it in place to patch tool arguments before execution. + * Later `tool_call` handlers see earlier mutations. No re-validation is performed after mutation. + */ +export type ToolCallEvent = + | BashToolCallEvent + | ReadToolCallEvent + | EditToolCallEvent + | WriteToolCallEvent + | GrepToolCallEvent + | FindToolCallEvent + | LsToolCallEvent + | CustomToolCallEvent; + +interface ToolResultEventBase { + type: "tool_result"; + toolCallId: string; + input: Record; + content: (TextContent | ImageContent)[]; + isError: boolean; +} + +export interface BashToolResultEvent extends ToolResultEventBase { + toolName: "bash"; + details: BashToolDetails | undefined; +} + +export interface ReadToolResultEvent extends ToolResultEventBase { + toolName: "read"; + details: ReadToolDetails | undefined; +} + +export interface EditToolResultEvent extends ToolResultEventBase { + toolName: "edit"; + details: EditToolDetails | undefined; +} + +export interface WriteToolResultEvent extends ToolResultEventBase { + toolName: "write"; + details: undefined; +} + +export interface GrepToolResultEvent extends ToolResultEventBase { + toolName: "grep"; + details: GrepToolDetails | undefined; +} + +export interface FindToolResultEvent extends ToolResultEventBase { + toolName: "find"; + details: FindToolDetails | undefined; +} + +export interface LsToolResultEvent extends ToolResultEventBase { + toolName: "ls"; + details: LsToolDetails | undefined; +} + +export interface CustomToolResultEvent extends ToolResultEventBase { + toolName: string; + details: unknown; +} + +/** Fired after a tool executes. Can modify result. */ +export type ToolResultEvent = + | BashToolResultEvent + | ReadToolResultEvent + | EditToolResultEvent + | WriteToolResultEvent + | GrepToolResultEvent + | FindToolResultEvent + | LsToolResultEvent + | CustomToolResultEvent; + +// Type guards for ToolResultEvent +export function isBashToolResult(e: ToolResultEvent): e is BashToolResultEvent { + return e.toolName === "bash"; +} +export function isReadToolResult(e: ToolResultEvent): e is ReadToolResultEvent { + return e.toolName === "read"; +} +export function isEditToolResult(e: ToolResultEvent): e is EditToolResultEvent { + return e.toolName === "edit"; +} +export function isWriteToolResult(e: ToolResultEvent): e is WriteToolResultEvent { + return e.toolName === "write"; +} +export function isGrepToolResult(e: ToolResultEvent): e is GrepToolResultEvent { + return e.toolName === "grep"; +} +export function isFindToolResult(e: ToolResultEvent): e is FindToolResultEvent { + return e.toolName === "find"; +} +export function isLsToolResult(e: ToolResultEvent): e is LsToolResultEvent { + return e.toolName === "ls"; +} + +/** + * Type guard for narrowing ToolCallEvent by tool name. + * + * Built-in tools narrow automatically (no type params needed): + * ```ts + * if (isToolCallEventType("bash", event)) { + * event.input.command; // string + * } + * ``` + * + * Custom tools require explicit type parameters: + * ```ts + * if (isToolCallEventType<"my_tool", MyToolInput>("my_tool", event)) { + * event.input.action; // typed + * } + * ``` + * + * Note: Direct narrowing via `event.toolName === "bash"` doesn't work because + * CustomToolCallEvent.toolName is `string` which overlaps with all literals. + */ +export function isToolCallEventType(toolName: "bash", event: ToolCallEvent): event is BashToolCallEvent; +export function isToolCallEventType(toolName: "read", event: ToolCallEvent): event is ReadToolCallEvent; +export function isToolCallEventType(toolName: "edit", event: ToolCallEvent): event is EditToolCallEvent; +export function isToolCallEventType(toolName: "write", event: ToolCallEvent): event is WriteToolCallEvent; +export function isToolCallEventType(toolName: "grep", event: ToolCallEvent): event is GrepToolCallEvent; +export function isToolCallEventType(toolName: "find", event: ToolCallEvent): event is FindToolCallEvent; +export function isToolCallEventType(toolName: "ls", event: ToolCallEvent): event is LsToolCallEvent; +export function isToolCallEventType>( + toolName: TName, + event: ToolCallEvent, +): event is ToolCallEvent & { toolName: TName; input: TInput }; +export function isToolCallEventType(toolName: string, event: ToolCallEvent): boolean { + return event.toolName === toolName; +} + +/** Union of all event types */ +export type ExtensionEvent = + | ProjectTrustEvent + | ResourcesDiscoverEvent + | SessionEvent + | ContextEvent + | BeforeProviderRequestEvent + | AfterProviderResponseEvent + | BeforeAgentStartEvent + | AgentStartEvent + | AgentEndEvent + | TurnStartEvent + | TurnEndEvent + | MessageStartEvent + | MessageUpdateEvent + | MessageEndEvent + | ToolExecutionStartEvent + | ToolExecutionUpdateEvent + | ToolExecutionEndEvent + | ModelSelectEvent + | ThinkingLevelSelectEvent + | UserBashEvent + | InputEvent + | ToolCallEvent + | ToolResultEvent; + +// ============================================================================ +// Event Results +// ============================================================================ + +export interface ContextEventResult { + messages?: AgentMessage[]; +} + +export type BeforeProviderRequestEventResult = unknown; + +export interface ToolCallEventResult { + /** Block tool execution. To modify arguments, mutate `event.input` in place instead. */ + block?: boolean; + reason?: string; +} + +/** Result from user_bash event handler */ +export interface UserBashEventResult { + /** Custom operations to use for execution */ + operations?: BashOperations; + /** Full replacement: extension handled execution, use this result */ + result?: BashResult; +} + +export interface ToolResultEventResult { + content?: (TextContent | ImageContent)[]; + details?: unknown; + isError?: boolean; +} + +export interface MessageEndEventResult { + /** Replace the finalized message. The replacement must keep the original message role. */ + message?: AgentMessage; +} + +export interface BeforeAgentStartEventResult { + message?: Pick; + /** Replace the system prompt for this turn. If multiple extensions return this, they are chained. */ + systemPrompt?: string; +} + +export interface SessionBeforeSwitchResult { + cancel?: boolean; +} + +export interface SessionBeforeForkResult { + cancel?: boolean; + skipConversationRestore?: boolean; +} + +export interface SessionBeforeCompactResult { + cancel?: boolean; + compaction?: CompactionResult; +} + +export interface SessionBeforeTreeResult { + cancel?: boolean; + summary?: { + summary: string; + details?: unknown; + }; + /** Override custom instructions for summarization */ + customInstructions?: string; + /** Override whether customInstructions replaces the default prompt */ + replaceInstructions?: boolean; + /** Override label to attach to the branch summary entry */ + label?: string; +} + +// ============================================================================ +// Message Rendering +// ============================================================================ + +export interface MessageRenderOptions { + expanded: boolean; +} + +export type MessageRenderer = ( + message: CustomMessage, + options: MessageRenderOptions, + theme: Theme, +) => Component | undefined; + +// ============================================================================ +// Command Registration +// ============================================================================ + +export interface RegisteredCommand { + name: string; + sourceInfo: SourceInfo; + description?: string; + getArgumentCompletions?: (argumentPrefix: string) => AutocompleteItem[] | null | Promise; + handler: (args: string, ctx: ExtensionCommandContext) => Promise; +} + +export interface ResolvedCommand extends RegisteredCommand { + invocationName: string; +} + +// ============================================================================ +// Extension API +// ============================================================================ + +/** Handler function type for events */ +// biome-ignore lint/suspicious/noConfusingVoidType: void allows bare return statements +export type ExtensionHandler = (event: E, ctx: ExtensionContext) => Promise | R | void; + +/** + * ExtensionAPI passed to extension factory functions. + */ +export interface ExtensionAPI { + // ========================================================================= + // Event Subscription + // ========================================================================= + + on(event: "project_trust", handler: ProjectTrustHandler): void; + on(event: "resources_discover", handler: ExtensionHandler): void; + on(event: "session_start", handler: ExtensionHandler): void; + on( + event: "session_before_switch", + handler: ExtensionHandler, + ): void; + on(event: "session_before_fork", handler: ExtensionHandler): void; + on( + event: "session_before_compact", + handler: ExtensionHandler, + ): void; + on(event: "session_compact", handler: ExtensionHandler): void; + on(event: "session_shutdown", handler: ExtensionHandler): void; + on(event: "session_before_tree", handler: ExtensionHandler): void; + on(event: "session_tree", handler: ExtensionHandler): void; + on(event: "context", handler: ExtensionHandler): void; + on( + event: "before_provider_request", + handler: ExtensionHandler, + ): void; + on(event: "after_provider_response", handler: ExtensionHandler): void; + on(event: "before_agent_start", handler: ExtensionHandler): void; + on(event: "agent_start", handler: ExtensionHandler): void; + on(event: "agent_end", handler: ExtensionHandler): void; + on(event: "turn_start", handler: ExtensionHandler): void; + on(event: "turn_end", handler: ExtensionHandler): void; + on(event: "message_start", handler: ExtensionHandler): void; + on(event: "message_update", handler: ExtensionHandler): void; + on(event: "message_end", handler: ExtensionHandler): void; + on(event: "tool_execution_start", handler: ExtensionHandler): void; + on(event: "tool_execution_update", handler: ExtensionHandler): void; + on(event: "tool_execution_end", handler: ExtensionHandler): void; + on(event: "model_select", handler: ExtensionHandler): void; + on(event: "thinking_level_select", handler: ExtensionHandler): void; + on(event: "tool_call", handler: ExtensionHandler): void; + on(event: "tool_result", handler: ExtensionHandler): void; + on(event: "user_bash", handler: ExtensionHandler): void; + on(event: "input", handler: ExtensionHandler): void; + + // ========================================================================= + // Tool Registration + // ========================================================================= + + /** Register a tool that the LLM can call. */ + registerTool( + tool: ToolDefinition, + ): void; + + // ========================================================================= + // Command, Shortcut, Flag Registration + // ========================================================================= + + /** Register a custom command. */ + registerCommand(name: string, options: Omit): void; + + /** Register a keyboard shortcut. */ + registerShortcut( + shortcut: KeyId, + options: { + description?: string; + handler: (ctx: ExtensionContext) => Promise | void; + }, + ): void; + + /** Register a CLI flag. */ + registerFlag( + name: string, + options: { + description?: string; + type: "boolean" | "string"; + default?: boolean | string; + }, + ): void; + + /** Get the value of a registered CLI flag. */ + getFlag(name: string): boolean | string | undefined; + + // ========================================================================= + // Message Rendering + // ========================================================================= + + /** Register a custom renderer for CustomMessageEntry. */ + registerMessageRenderer(customType: string, renderer: MessageRenderer): void; + + // ========================================================================= + // Actions + // ========================================================================= + + /** Send a custom message to the session. */ + sendMessage( + message: Pick, "customType" | "content" | "display" | "details">, + options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, + ): void; + + /** + * Send a user message to the agent. Always triggers a turn. + * When the agent is streaming, use deliverAs to specify how to queue the message. + */ + sendUserMessage( + content: string | (TextContent | ImageContent)[], + options?: { deliverAs?: "steer" | "followUp" }, + ): void; + + /** Append a custom entry to the session for state persistence (not sent to LLM). */ + appendEntry(customType: string, data?: T): void; + + // ========================================================================= + // Session Metadata + // ========================================================================= + + /** Set the session display name (shown in session selector). */ + setSessionName(name: string): void; + + /** Get the current session name, if set. */ + getSessionName(): string | undefined; + + /** Set or clear a label on an entry. Labels are user-defined markers for bookmarking/navigation. */ + setLabel(entryId: string, label: string | undefined): void; + + /** Execute a shell command. */ + exec(command: string, args: string[], options?: ExecOptions): Promise; + + /** Get the list of currently active tool names. */ + getActiveTools(): string[]; + + /** Get all configured tools with parameter schema, prompt guidelines, and source metadata. */ + getAllTools(): ToolInfo[]; + + /** Set the active tools by name. */ + setActiveTools(toolNames: string[]): void; + + /** Get available slash commands in the current session. */ + getCommands(): SlashCommandInfo[]; + + // ========================================================================= + // Model and Thinking Level + // ========================================================================= + + /** Set the current model. Returns false if no API key available. */ + setModel(model: Model): Promise; + + /** Get current thinking level. */ + getThinkingLevel(): ThinkingLevel; + + /** Set thinking level (clamped to model capabilities). */ + setThinkingLevel(level: ThinkingLevel): void; + + // ========================================================================= + // Provider Registration + // ========================================================================= + + /** + * Register or override a model provider. + * + * If `models` is provided: replaces all existing models for this provider. + * If only `baseUrl` is provided: overrides the URL for existing models. + * If `oauth` is provided: registers OAuth provider for /login support. + * If `streamSimple` is provided: registers a custom API stream handler. + * + * During initial extension load this call is queued and applied once the + * runner has bound its context. After that it takes effect immediately, so + * it is safe to call from command handlers or event callbacks without + * requiring a `/reload`. + * + * @example + * // Register a new provider with custom models + * pi.registerProvider("my-proxy", { + * baseUrl: "https://proxy.example.com", + * apiKey: "$PROXY_API_KEY", + * api: "anthropic-messages", + * models: [ + * { + * id: "claude-sonnet-4-20250514", + * name: "Claude 4 Sonnet (proxy)", + * reasoning: false, + * input: ["text", "image"], + * cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + * contextWindow: 200000, + * maxTokens: 16384 + * } + * ] + * }); + * + * @example + * // Override baseUrl for an existing provider + * pi.registerProvider("anthropic", { + * baseUrl: "https://proxy.example.com" + * }); + * + * @example + * // Register provider with OAuth support + * pi.registerProvider("corporate-ai", { + * baseUrl: "https://ai.corp.com", + * api: "openai-responses", + * models: [...], + * oauth: { + * name: "Corporate AI (SSO)", + * async login(callbacks) { ... }, + * async refreshToken(credentials) { ... }, + * getApiKey(credentials) { return credentials.access; } + * } + * }); + */ + registerProvider(name: string, config: ProviderConfig): void; + + /** + * Unregister a previously registered provider. + * + * Removes all models belonging to the named provider and restores any + * built-in models that were overridden by it. Has no effect if the provider + * is not currently registered. + * + * Like `registerProvider`, this takes effect immediately when called after + * the initial load phase. + * + * @example + * pi.unregisterProvider("my-proxy"); + */ + unregisterProvider(name: string): void; + + /** Shared event bus for extension communication. */ + events: EventBus; +} + +// ============================================================================ +// Provider Registration Types +// ============================================================================ + +/** Configuration for registering a provider via pi.registerProvider(). */ +export interface ProviderConfig { + /** Display name for the provider in UI. */ + name?: string; + /** Base URL for the API endpoint. Required when defining models. */ + baseUrl?: string; + /** API key literal, env interpolation ($ENV_VAR or ${ENV_VAR}), or leading !command. Required when defining models (unless oauth provided). */ + apiKey?: string; + /** API type. Required at provider or model level when defining models. */ + api?: Api; + /** Optional streamSimple handler for custom APIs. */ + streamSimple?: (model: Model, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream; + /** Custom headers to include in requests. */ + headers?: Record; + /** If true, adds Authorization: Bearer header with the resolved API key. */ + authHeader?: boolean; + /** Models to register. If provided, replaces all existing models for this provider. */ + models?: ProviderModelConfig[]; + /** OAuth provider for /login support. The `id` is set automatically from the provider name. */ + oauth?: { + /** Display name for the provider in login UI. */ + name: string; + /** Run the login flow, return credentials to persist. */ + login(callbacks: OAuthLoginCallbacks): Promise; + /** Refresh expired credentials, return updated credentials to persist. */ + refreshToken(credentials: OAuthCredentials): Promise; + /** Convert credentials to API key string for the provider. */ + getApiKey(credentials: OAuthCredentials): string; + /** Optional: modify models for this provider (e.g., update baseUrl based on credentials). */ + modifyModels?(models: Model[], credentials: OAuthCredentials): Model[]; + }; +} + +/** Configuration for a model within a provider. */ +export interface ProviderModelConfig { + /** Model ID (e.g., "claude-sonnet-4-20250514"). */ + id: string; + /** Display name (e.g., "Claude 4 Sonnet"). */ + name: string; + /** API type override for this model. */ + api?: Api; + /** API endpoint URL override for this model. */ + baseUrl?: string; + /** Whether the model supports extended thinking. */ + reasoning: boolean; + /** Maps pi thinking levels to provider/model-specific values; null marks a level unsupported. */ + thinkingLevelMap?: Model["thinkingLevelMap"]; + /** Supported input types. */ + input: ("text" | "image")[]; + /** Cost per token (for tracking, can be 0). */ + cost: { input: number; output: number; cacheRead: number; cacheWrite: number }; + /** Maximum context window size in tokens. */ + contextWindow: number; + /** Maximum output tokens. */ + maxTokens: number; + /** Custom headers for this model. */ + headers?: Record; + /** OpenAI compatibility settings. */ + compat?: Model["compat"]; +} + +/** Extension factory function type. Supports both sync and async initialization. */ +export type ExtensionFactory = (pi: ExtensionAPI) => void | Promise; + +// ============================================================================ +// Loaded Extension Types +// ============================================================================ + +export interface RegisteredTool { + definition: ToolDefinition; + sourceInfo: SourceInfo; +} + +export interface ExtensionFlag { + name: string; + description?: string; + type: "boolean" | "string"; + default?: boolean | string; + extensionPath: string; +} + +export interface ExtensionShortcut { + shortcut: KeyId; + description?: string; + handler: (ctx: ExtensionContext) => Promise | void; + extensionPath: string; +} + +type HandlerFn = (...args: unknown[]) => Promise; + +export type SendMessageHandler = ( + message: Pick, "customType" | "content" | "display" | "details">, + options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, +) => void; + +export type SendUserMessageHandler = ( + content: string | (TextContent | ImageContent)[], + options?: { deliverAs?: "steer" | "followUp" }, +) => void; + +export type AppendEntryHandler = (customType: string, data?: T) => void; + +export type SetSessionNameHandler = (name: string) => void; + +export type GetSessionNameHandler = () => string | undefined; + +export type GetActiveToolsHandler = () => string[]; + +/** Tool info with name, description, parameter schema, prompt guidelines, and source metadata. */ +export type ToolInfo = Pick & { + sourceInfo: SourceInfo; +}; + +export type GetAllToolsHandler = () => ToolInfo[]; + +export type GetCommandsHandler = () => SlashCommandInfo[]; + +export type SetActiveToolsHandler = (toolNames: string[]) => void; + +export type RefreshToolsHandler = () => void; + +export type SetModelHandler = (model: Model) => Promise; + +export type GetThinkingLevelHandler = () => ThinkingLevel; + +export type SetThinkingLevelHandler = (level: ThinkingLevel) => void; + +export type SetLabelHandler = (entryId: string, label: string | undefined) => void; + +/** + * Shared state created by loader, used during registration and runtime. + * Contains flag values (defaults set during registration, CLI values set after). + */ +export interface ExtensionRuntimeState { + flagValues: Map; + /** Provider registrations queued during extension loading, processed when runner binds */ + pendingProviderRegistrations: Array<{ name: string; config: ProviderConfig; extensionPath: string }>; + /** Throws when this extension instance is stale after runtime replacement. */ + assertActive: () => void; + /** Marks this extension instance as stale after runtime replacement or reload. */ + invalidate: (message?: string) => void; + /** + * Register or unregister a provider. + * + * Before bindCore(): queues registrations / removes from queue. + * After bindCore(): calls ModelRegistry directly for immediate effect. + */ + registerProvider: (name: string, config: ProviderConfig, extensionPath?: string) => void; + unregisterProvider: (name: string, extensionPath?: string) => void; +} + +/** + * Action implementations for pi.* API methods. + * Provided to runner.initialize(), copied into the shared runtime. + */ +export interface ExtensionActions { + sendMessage: SendMessageHandler; + sendUserMessage: SendUserMessageHandler; + appendEntry: AppendEntryHandler; + setSessionName: SetSessionNameHandler; + getSessionName: GetSessionNameHandler; + setLabel: SetLabelHandler; + getActiveTools: GetActiveToolsHandler; + getAllTools: GetAllToolsHandler; + setActiveTools: SetActiveToolsHandler; + refreshTools: RefreshToolsHandler; + getCommands: GetCommandsHandler; + setModel: SetModelHandler; + getThinkingLevel: GetThinkingLevelHandler; + setThinkingLevel: SetThinkingLevelHandler; +} + +/** + * Actions for ExtensionContext (ctx.* in event handlers). + * Required by all modes. + */ +export interface ExtensionContextActions { + getModel: () => Model | undefined; + isIdle: () => boolean; + isProjectTrusted: () => boolean; + getSignal: () => AbortSignal | undefined; + abort: () => void; + hasPendingMessages: () => boolean; + shutdown: () => void; + getContextUsage: () => ContextUsage | undefined; + compact: (options?: CompactOptions) => void; + getSystemPrompt: () => string; + getSystemPromptOptions?: () => BuildSystemPromptOptions; +} + +/** + * Actions for ExtensionCommandContext (ctx.* in command handlers). + * Only needed for interactive mode where extension commands are invokable. + */ +export interface ExtensionCommandContextActions { + waitForIdle: () => Promise; + newSession: (options?: { + parentSession?: string; + setup?: (sessionManager: SessionManager) => Promise; + withSession?: (ctx: ReplacedSessionContext) => Promise; + }) => Promise<{ cancelled: boolean }>; + fork: ( + entryId: string, + options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise }, + ) => Promise<{ cancelled: boolean }>; + navigateTree: ( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ) => Promise<{ cancelled: boolean }>; + switchSession: ( + sessionPath: string, + options?: { withSession?: (ctx: ReplacedSessionContext) => Promise }, + ) => Promise<{ cancelled: boolean }>; + reload: () => Promise; +} + +/** + * Full runtime = state + actions. + * Created by loader with throwing action stubs, completed by runner.initialize(). + */ +export interface ExtensionRuntime extends ExtensionRuntimeState, ExtensionActions {} + +/** Loaded extension with all registered items. */ +export interface Extension { + path: string; + resolvedPath: string; + sourceInfo: SourceInfo; + handlers: Map; + tools: Map; + messageRenderers: Map; + commands: Map; + flags: Map; + shortcuts: Map; +} + +/** Result of loading extensions. */ +export interface LoadExtensionsResult { + extensions: Extension[]; + errors: Array<{ path: string; error: string }>; + /** Shared runtime - actions are throwing stubs until runner.initialize() */ + runtime: ExtensionRuntime; +} + +// ============================================================================ +// Extension Error +// ============================================================================ + +export interface ExtensionError { + extensionPath: string; + event: string; + error: string; + stack?: string; +} diff --git a/cactus-code/packages/coding-agent/src/core/extensions/wrapper.ts b/cactus-code/packages/coding-agent/src/core/extensions/wrapper.ts new file mode 100644 index 000000000..6a84038c2 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/extensions/wrapper.ts @@ -0,0 +1,30 @@ +/** + * Tool wrappers for extension-registered tools. + * + * These wrappers only adapt tool execution so extension tools receive the runner context. + * Tool call and tool result interception is handled by AgentSession via agent-core hooks. + */ + +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { wrapToolDefinition, wrapToolDefinitions } from "../tools/tool-definition-wrapper.ts"; +import type { ExtensionRunner } from "./runner.ts"; +import type { RegisteredTool } from "./types.ts"; + +/** + * Wrap a RegisteredTool into an AgentTool. + * Uses the runner's createContext() for consistent context across tools and event handlers. + */ +export function wrapRegisteredTool(registeredTool: RegisteredTool, runner: ExtensionRunner): AgentTool { + return wrapToolDefinition(registeredTool.definition, () => runner.createContext()); +} + +/** + * Wrap all registered tools into AgentTools. + * Uses the runner's createContext() for consistent context across tools and event handlers. + */ +export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] { + return wrapToolDefinitions( + registeredTools.map((registeredTool) => registeredTool.definition), + () => runner.createContext(), + ); +} diff --git a/cactus-code/packages/coding-agent/src/core/footer-data-provider.ts b/cactus-code/packages/coding-agent/src/core/footer-data-provider.ts new file mode 100644 index 000000000..f15001c2e --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/footer-data-provider.ts @@ -0,0 +1,388 @@ +import { type ExecFileException, execFile, spawnSync } from "child_process"; +import { existsSync, type FSWatcher, readFileSync, type Stats, statSync, unwatchFile, watchFile } from "fs"; +import { dirname, join, resolve } from "path"; +import { closeWatcher, FS_WATCH_RETRY_DELAY_MS, watchWithErrorHandler } from "../utils/fs-watch.ts"; + +type GitPaths = { + repoDir: string; + commonGitDir: string; + headPath: string; +}; + +/** + * Find git metadata paths by walking up from cwd. + * Handles both regular git repos (.git is a directory) and worktrees (.git is a file). + */ +function findGitPaths(cwd: string): GitPaths | null { + let dir = cwd; + while (true) { + const gitPath = join(dir, ".git"); + if (existsSync(gitPath)) { + try { + const stat = statSync(gitPath); + if (stat.isFile()) { + const content = readFileSync(gitPath, "utf8").trim(); + if (content.startsWith("gitdir: ")) { + const gitDir = resolve(dir, content.slice(8).trim()); + const headPath = join(gitDir, "HEAD"); + if (!existsSync(headPath)) return null; + const commonDirPath = join(gitDir, "commondir"); + const commonGitDir = existsSync(commonDirPath) + ? resolve(gitDir, readFileSync(commonDirPath, "utf8").trim()) + : gitDir; + return { repoDir: dir, commonGitDir, headPath }; + } + } else if (stat.isDirectory()) { + const headPath = join(gitPath, "HEAD"); + if (!existsSync(headPath)) return null; + return { repoDir: dir, commonGitDir: gitPath, headPath }; + } + } catch { + return null; + } + } + const parent = dirname(dir); + if (parent === dir) return null; + dir = parent; + } +} + +/** Ask git for the current branch. Returns null on detached HEAD or if git is unavailable. */ +function resolveBranchWithGitSync(repoDir: string): string | null { + const result = spawnSync("git", ["--no-optional-locks", "symbolic-ref", "--quiet", "--short", "HEAD"], { + cwd: repoDir, + encoding: "utf8", + stdio: ["ignore", "pipe", "ignore"], + }); + const branch = result.status === 0 ? result.stdout.trim() : ""; + return branch || null; +} + +/** Ask git for the current branch asynchronously. Returns null on detached HEAD or if git is unavailable. */ +function resolveBranchWithGitAsync(repoDir: string): Promise { + return new Promise((resolvePromise) => { + execFile( + "git", + ["--no-optional-locks", "symbolic-ref", "--quiet", "--short", "HEAD"], + { + cwd: repoDir, + encoding: "utf8", + }, + (error: ExecFileException | null, stdout: string) => { + if (error) { + resolvePromise(null); + return; + } + const branch = stdout.trim(); + resolvePromise(branch || null); + }, + ); + }); +} + +function isWslEnvironment(): boolean { + return process.platform === "linux" && !!(process.env.WSL_DISTRO_NAME || process.env.WSL_INTEROP); +} + +function isWindowsMountedRepoPath(repoDir: string): boolean { + return /^\/mnt\/[a-z](?:\/|$)/i.test(repoDir); +} + +function shouldPollGitHead(repoDir: string): boolean { + return isWslEnvironment() && isWindowsMountedRepoPath(repoDir); +} + +/** + * Provides git branch and extension statuses - data not otherwise accessible to extensions. + * Token stats, model info available via ctx.sessionManager and ctx.model. + */ +export class FooterDataProvider { + private cwd: string; + private static readonly WATCH_DEBOUNCE_MS = 500; + + private extensionStatuses = new Map(); + private cachedBranch: string | null | undefined = undefined; + private gitPaths: GitPaths | null | undefined = undefined; + private headWatcher: FSWatcher | null = null; + private headWatchFilePath: string | null = null; + private headWatchFileListener: ((current: Stats, previous: Stats) => void) | null = null; + private reftableWatcher: FSWatcher | null = null; + private reftableTablesListWatcher: FSWatcher | null = null; + private reftableTablesListPath: string | null = null; + private branchChangeCallbacks = new Set<() => void>(); + private availableProviderCount = 0; + private refreshTimer: ReturnType | null = null; + private gitWatcherRetryTimer: ReturnType | null = null; + private refreshInFlight = false; + private refreshPending = false; + private disposed = false; + + constructor(cwd: string) { + this.cwd = cwd; + this.gitPaths = findGitPaths(cwd); + this.setupGitWatcher(); + } + + /** Current git branch, null if not in repo, "detached" if detached HEAD */ + getGitBranch(): string | null { + if (this.cachedBranch === undefined) { + this.cachedBranch = this.resolveGitBranchSync(); + } + return this.cachedBranch; + } + + /** Extension status texts set via ctx.ui.setStatus() */ + getExtensionStatuses(): ReadonlyMap { + return this.extensionStatuses; + } + + /** Subscribe to git branch changes. Returns unsubscribe function. */ + onBranchChange(callback: () => void): () => void { + this.branchChangeCallbacks.add(callback); + return () => this.branchChangeCallbacks.delete(callback); + } + + /** Internal: set extension status */ + setExtensionStatus(key: string, text: string | undefined): void { + if (text === undefined) { + this.extensionStatuses.delete(key); + } else { + this.extensionStatuses.set(key, text); + } + } + + /** Internal: clear extension statuses */ + clearExtensionStatuses(): void { + this.extensionStatuses.clear(); + } + + /** Number of unique providers with available models (for footer display) */ + getAvailableProviderCount(): number { + return this.availableProviderCount; + } + + /** Internal: update available provider count */ + setAvailableProviderCount(count: number): void { + this.availableProviderCount = count; + } + + setCwd(cwd: string): void { + if (this.cwd === cwd) { + return; + } + + this.cwd = cwd; + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + this.refreshTimer = null; + } + this.clearGitWatchers(); + this.cachedBranch = undefined; + this.gitPaths = findGitPaths(cwd); + this.setupGitWatcher(); + this.notifyBranchChange(); + } + + /** Internal: cleanup */ + dispose(): void { + this.disposed = true; + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + this.refreshTimer = null; + } + this.clearGitWatchers(); + this.branchChangeCallbacks.clear(); + } + + private notifyBranchChange(): void { + for (const cb of this.branchChangeCallbacks) cb(); + } + + private scheduleRefresh(): void { + if (this.disposed || this.refreshTimer) return; + if (this.refreshInFlight) { + this.refreshPending = true; + return; + } + this.refreshTimer = setTimeout(() => { + this.refreshTimer = null; + void this.refreshGitBranchAsync(); + }, FooterDataProvider.WATCH_DEBOUNCE_MS); + } + + private async refreshGitBranchAsync(): Promise { + if (this.disposed) return; + if (this.refreshInFlight) { + this.refreshPending = true; + return; + } + + this.refreshInFlight = true; + try { + const nextBranch = await this.resolveGitBranchAsync(); + if (this.disposed) return; + if (this.cachedBranch !== undefined && this.cachedBranch !== nextBranch) { + this.cachedBranch = nextBranch; + this.notifyBranchChange(); + return; + } + this.cachedBranch = nextBranch; + } finally { + this.refreshInFlight = false; + if (this.refreshPending && !this.disposed) { + this.refreshPending = false; + this.scheduleRefresh(); + } + } + } + + private resolveGitBranchSync(): string | null { + try { + if (!this.gitPaths) return null; + const content = readFileSync(this.gitPaths.headPath, "utf8").trim(); + if (content.startsWith("ref: refs/heads/")) { + const branch = content.slice(16); + return branch === ".invalid" ? (resolveBranchWithGitSync(this.gitPaths.repoDir) ?? "detached") : branch; + } + return "detached"; + } catch { + return null; + } + } + + private async resolveGitBranchAsync(): Promise { + try { + if (!this.gitPaths) return null; + const content = readFileSync(this.gitPaths.headPath, "utf8").trim(); + if (content.startsWith("ref: refs/heads/")) { + const branch = content.slice(16); + return branch === ".invalid" + ? ((await resolveBranchWithGitAsync(this.gitPaths.repoDir)) ?? "detached") + : branch; + } + return "detached"; + } catch { + return null; + } + } + + private clearGitWatchers(): void { + closeWatcher(this.headWatcher); + this.headWatcher = null; + if (this.headWatchFilePath && this.headWatchFileListener) { + unwatchFile(this.headWatchFilePath, this.headWatchFileListener); + this.headWatchFilePath = null; + this.headWatchFileListener = null; + } + closeWatcher(this.reftableWatcher); + this.reftableWatcher = null; + closeWatcher(this.reftableTablesListWatcher); + this.reftableTablesListWatcher = null; + if (this.reftableTablesListPath) { + unwatchFile(this.reftableTablesListPath); + this.reftableTablesListPath = null; + } + if (this.gitWatcherRetryTimer) { + clearTimeout(this.gitWatcherRetryTimer); + this.gitWatcherRetryTimer = null; + } + } + + private scheduleGitWatcherRetry(): void { + if (this.disposed || this.gitWatcherRetryTimer) { + return; + } + + this.gitWatcherRetryTimer = setTimeout(() => { + this.gitWatcherRetryTimer = null; + this.setupGitWatcher(); + }, FS_WATCH_RETRY_DELAY_MS); + } + + private handleGitWatcherError(): void { + this.clearGitWatchers(); + this.scheduleGitWatcherRetry(); + } + + private setupGitWatcher(): void { + this.clearGitWatchers(); + if (!this.gitPaths) return; + + const pollGitHead = shouldPollGitHead(this.gitPaths.repoDir); + + // Watch the directory containing HEAD, not HEAD itself. + // Git uses atomic writes (write temp, rename over HEAD), which changes the inode. + // fs.watch on a file stops working after the inode changes. + this.headWatcher = watchWithErrorHandler( + dirname(this.gitPaths.headPath), + (_eventType, filename) => { + if (!filename || filename === "HEAD") { + this.scheduleRefresh(); + } + }, + () => this.handleGitWatcherError(), + ); + if (pollGitHead) { + this.headWatchFilePath = this.gitPaths.headPath; + this.headWatchFileListener = (current, previous) => { + if ( + current.mtimeMs !== previous.mtimeMs || + current.ctimeMs !== previous.ctimeMs || + current.size !== previous.size + ) { + this.scheduleRefresh(); + } + }; + watchFile(this.headWatchFilePath, { interval: 1000 }, this.headWatchFileListener); + } + if (!this.headWatcher && !pollGitHead) { + return; + } + + // In reftable repos, branch switches update files in the reftable directory + // instead of HEAD. Watch it separately so the footer picks up those changes. + const reftableDir = join(this.gitPaths.commonGitDir, "reftable"); + if (existsSync(reftableDir)) { + this.reftableWatcher = watchWithErrorHandler( + reftableDir, + () => { + this.scheduleRefresh(); + }, + () => this.handleGitWatcherError(), + ); + if (!this.reftableWatcher) { + return; + } + + const tablesListPath = join(reftableDir, "tables.list"); + if (existsSync(tablesListPath)) { + this.reftableTablesListPath = tablesListPath; + this.reftableTablesListWatcher = watchWithErrorHandler( + tablesListPath, + () => { + this.scheduleRefresh(); + }, + () => this.handleGitWatcherError(), + ); + if (!this.reftableTablesListWatcher) { + return; + } + watchFile(tablesListPath, { interval: 250 }, (current, previous) => { + if ( + current.mtimeMs !== previous.mtimeMs || + current.ctimeMs !== previous.ctimeMs || + current.size !== previous.size + ) { + this.scheduleRefresh(); + } + }); + } + } + } +} + +/** Read-only view for extensions - excludes setExtensionStatus, setAvailableProviderCount and dispose */ +export type ReadonlyFooterDataProvider = Pick< + FooterDataProvider, + "getGitBranch" | "getExtensionStatuses" | "getAvailableProviderCount" | "onBranchChange" +>; diff --git a/cactus-code/packages/coding-agent/src/core/http-dispatcher.ts b/cactus-code/packages/coding-agent/src/core/http-dispatcher.ts new file mode 100644 index 000000000..0910f4d6c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/http-dispatcher.ts @@ -0,0 +1,73 @@ +import * as undici from "undici"; + +export const DEFAULT_HTTP_IDLE_TIMEOUT_MS = 300_000; + +export const HTTP_IDLE_TIMEOUT_CHOICES = [ + { label: "30 sec", timeoutMs: 30_000 }, + { label: "1 min", timeoutMs: 60_000 }, + { label: "2 min", timeoutMs: 120_000 }, + { label: "5 min", timeoutMs: 300_000 }, + { label: "disabled", timeoutMs: 0 }, +] as const; + +const originalGlobalFetch = globalThis.fetch; +let installedGlobalFetch: typeof globalThis.fetch | undefined; + +export function parseHttpIdleTimeoutMs(value: unknown): number | undefined { + if (typeof value === "string") { + const trimmed = value.trim(); + if (trimmed.toLowerCase() === "disabled") { + return 0; + } + if (trimmed.length === 0) { + return undefined; + } + return parseHttpIdleTimeoutMs(Number(trimmed)); + } + + if (typeof value !== "number" || !Number.isFinite(value) || value < 0) { + return undefined; + } + return Math.floor(value); +} + +export function formatHttpIdleTimeoutMs(timeoutMs: number): string { + const choice = HTTP_IDLE_TIMEOUT_CHOICES.find((item) => item.timeoutMs === timeoutMs); + if (choice) { + return choice.label; + } + return `${timeoutMs / 1000} sec`; +} + +export function applyHttpProxySettings(httpProxy: string | undefined): void { + const proxy = httpProxy?.trim(); + if (!proxy) return; + process.env.HTTP_PROXY ??= proxy; + process.env.HTTPS_PROXY ??= proxy; +} + +export function configureHttpDispatcher(timeoutMs: number = DEFAULT_HTTP_IDLE_TIMEOUT_MS): void { + const normalizedTimeoutMs = parseHttpIdleTimeoutMs(timeoutMs); + if (normalizedTimeoutMs === undefined) { + throw new Error(`Invalid HTTP idle timeout: ${String(timeoutMs)}`); + } + undici.setGlobalDispatcher( + new undici.EnvHttpProxyAgent({ + allowH2: false, + bodyTimeout: normalizedTimeoutMs, + headersTimeout: normalizedTimeoutMs, + }), + ); + // Keep fetch and the dispatcher on the same undici implementation. Node 26.0's + // bundled fetch can otherwise consume compressed responses through npm undici's + // dispatcher without decompressing them, causing response.json() failures. + // If a caller replaced fetch after module load, preserve that deliberate override. + const shouldInstallGlobals = + installedGlobalFetch === undefined + ? globalThis.fetch === originalGlobalFetch + : globalThis.fetch === installedGlobalFetch; + if (shouldInstallGlobals) { + undici.install?.(); + installedGlobalFetch = globalThis.fetch; + } +} diff --git a/cactus-code/packages/coding-agent/src/core/index.ts b/cactus-code/packages/coding-agent/src/core/index.ts new file mode 100644 index 000000000..b7654f422 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/index.ts @@ -0,0 +1,78 @@ +/** + * Core modules shared between all run modes. + */ + +export { + AgentSession, + type AgentSessionConfig, + type AgentSessionEvent, + type AgentSessionEventListener, + type ModelCycleResult, + type PromptOptions, + type SessionStats, +} from "./agent-session.ts"; +export { + AgentSessionRuntime, + type CreateAgentSessionRuntimeFactory, + type CreateAgentSessionRuntimeResult, + createAgentSessionRuntime, +} from "./agent-session-runtime.ts"; +export { + type AgentSessionRuntimeDiagnostic, + type AgentSessionServices, + type CreateAgentSessionFromServicesOptions, + type CreateAgentSessionServicesOptions, + createAgentSessionFromServices, + createAgentSessionServices, +} from "./agent-session-services.ts"; +export { type BashExecutorOptions, type BashResult, executeBashWithOperations } from "./bash-executor.ts"; +export type { CompactionResult } from "./compaction/index.ts"; +export { createEventBus, type EventBus, type EventBusController } from "./event-bus.ts"; +export { areExperimentalFeaturesEnabled } from "./experimental.ts"; +// Extensions system +export { + type AgentEndEvent, + type AgentStartEvent, + type AgentToolResult, + type AgentToolUpdateCallback, + type BeforeAgentStartEvent, + type BeforeAgentStartEventResult, + type BuildSystemPromptOptions, + type ContextEvent, + defineTool, + discoverAndLoadExtensions, + type ExecOptions, + type ExecResult, + type Extension, + type ExtensionAPI, + type ExtensionCommandContext, + type ExtensionContext, + type ExtensionError, + type ExtensionEvent, + type ExtensionFactory, + type ExtensionFlag, + type ExtensionHandler, + ExtensionRunner, + type ExtensionShortcut, + type ExtensionUIContext, + type LoadExtensionsResult, + type MessageRenderer, + type RegisteredCommand, + type SessionBeforeCompactEvent, + type SessionBeforeForkEvent, + type SessionBeforeSwitchEvent, + type SessionBeforeTreeEvent, + type SessionCompactEvent, + type SessionShutdownEvent, + type SessionStartEvent, + type SessionTreeEvent, + type ToolCallEvent, + type ToolCallEventResult, + type ToolDefinition, + type ToolRenderResultOptions, + type ToolResultEvent, + type TurnEndEvent, + type TurnStartEvent, + type WorkingIndicatorOptions, +} from "./extensions/index.ts"; +export { createSyntheticSourceInfo } from "./source-info.ts"; diff --git a/cactus-code/packages/coding-agent/src/core/keybindings.ts b/cactus-code/packages/coding-agent/src/core/keybindings.ts new file mode 100644 index 000000000..7c0de6431 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/keybindings.ts @@ -0,0 +1,370 @@ +import { + type Keybinding, + type KeybindingDefinitions, + type KeybindingsConfig, + type KeyId, + TUI_KEYBINDINGS, + KeybindingsManager as TuiKeybindingsManager, +} from "@earendil-works/pi-tui"; +import { existsSync, readFileSync } from "fs"; +import { join } from "path"; +import { getAgentDir } from "../config.ts"; + +export interface AppKeybindings { + "app.interrupt": true; + "app.clear": true; + "app.exit": true; + "app.suspend": true; + "app.thinking.cycle": true; + "app.model.cycleForward": true; + "app.model.cycleBackward": true; + "app.model.select": true; + "app.tools.expand": true; + "app.thinking.toggle": true; + "app.session.toggleNamedFilter": true; + "app.editor.external": true; + "app.message.followUp": true; + "app.message.dequeue": true; + "app.clipboard.pasteImage": true; + "app.session.new": true; + "app.session.tree": true; + "app.session.fork": true; + "app.session.resume": true; + "app.tree.foldOrUp": true; + "app.tree.unfoldOrDown": true; + "app.tree.editLabel": true; + "app.tree.toggleLabelTimestamp": true; + "app.session.togglePath": true; + "app.session.toggleSort": true; + "app.session.rename": true; + "app.session.delete": true; + "app.session.deleteNoninvasive": true; + "app.models.save": true; + "app.models.enableAll": true; + "app.models.clearAll": true; + "app.models.toggleProvider": true; + "app.models.reorderUp": true; + "app.models.reorderDown": true; + "app.tree.filter.default": true; + "app.tree.filter.noTools": true; + "app.tree.filter.userOnly": true; + "app.tree.filter.labeledOnly": true; + "app.tree.filter.all": true; + "app.tree.filter.cycleForward": true; + "app.tree.filter.cycleBackward": true; +} + +export type AppKeybinding = keyof AppKeybindings; + +declare module "@earendil-works/pi-tui" { + interface Keybindings extends AppKeybindings {} +} + +export const KEYBINDINGS = { + ...TUI_KEYBINDINGS, + "app.interrupt": { defaultKeys: "escape", description: "Cancel or abort" }, + "app.clear": { defaultKeys: "ctrl+c", description: "Clear editor" }, + "app.exit": { defaultKeys: "ctrl+d", description: "Exit when editor is empty" }, + "app.suspend": { + defaultKeys: process.platform === "win32" ? [] : "ctrl+z", + description: "Suspend to background", + }, + "app.thinking.cycle": { + defaultKeys: "shift+tab", + description: "Cycle thinking level", + }, + "app.model.cycleForward": { + defaultKeys: "ctrl+p", + description: "Cycle to next model", + }, + "app.model.cycleBackward": { + defaultKeys: "shift+ctrl+p", + description: "Cycle to previous model", + }, + "app.model.select": { defaultKeys: "ctrl+l", description: "Open model selector" }, + "app.tools.expand": { defaultKeys: "ctrl+o", description: "Toggle tool output" }, + "app.thinking.toggle": { + defaultKeys: "ctrl+t", + description: "Toggle thinking blocks", + }, + "app.session.toggleNamedFilter": { + defaultKeys: "ctrl+n", + description: "Toggle named session filter", + }, + "app.editor.external": { + defaultKeys: "ctrl+g", + description: "Open external editor", + }, + "app.message.followUp": { + defaultKeys: "alt+enter", + description: "Queue follow-up message", + }, + "app.message.dequeue": { + defaultKeys: "alt+up", + description: "Restore queued messages", + }, + "app.clipboard.pasteImage": { + defaultKeys: process.platform === "win32" ? "alt+v" : "ctrl+v", + description: "Paste image from clipboard", + }, + "app.session.new": { defaultKeys: [], description: "Start a new session" }, + "app.session.tree": { defaultKeys: [], description: "Open session tree" }, + "app.session.fork": { defaultKeys: [], description: "Fork current session" }, + "app.session.resume": { defaultKeys: [], description: "Resume a session" }, + "app.tree.foldOrUp": { + defaultKeys: ["ctrl+left", "alt+left"], + description: "Fold tree branch or move up", + }, + "app.tree.unfoldOrDown": { + defaultKeys: ["ctrl+right", "alt+right"], + description: "Unfold tree branch or move down", + }, + "app.tree.editLabel": { + defaultKeys: "shift+l", + description: "Edit tree label", + }, + "app.tree.toggleLabelTimestamp": { + defaultKeys: "shift+t", + description: "Toggle tree label timestamps", + }, + "app.session.togglePath": { + defaultKeys: "ctrl+p", + description: "Toggle session path display", + }, + "app.session.toggleSort": { + defaultKeys: "ctrl+s", + description: "Toggle session sort mode", + }, + "app.session.rename": { + defaultKeys: "ctrl+r", + description: "Rename session", + }, + "app.session.delete": { + defaultKeys: "ctrl+d", + description: "Delete session", + }, + "app.session.deleteNoninvasive": { + defaultKeys: "ctrl+backspace", + description: "Delete session when query is empty", + }, + "app.models.save": { + defaultKeys: "ctrl+s", + description: "Save model selection", + }, + "app.models.enableAll": { + defaultKeys: "ctrl+a", + description: "Enable all models", + }, + "app.models.clearAll": { + defaultKeys: "ctrl+x", + description: "Clear all models", + }, + "app.models.toggleProvider": { + defaultKeys: "ctrl+p", + description: "Toggle all models for provider", + }, + "app.models.reorderUp": { + defaultKeys: "alt+up", + description: "Move model up in order", + }, + "app.models.reorderDown": { + defaultKeys: "alt+down", + description: "Move model down in order", + }, + "app.tree.filter.default": { + defaultKeys: "ctrl+d", + description: "Tree filter: default view", + }, + "app.tree.filter.noTools": { + defaultKeys: "ctrl+t", + description: "Tree filter: hide tool results", + }, + "app.tree.filter.userOnly": { + defaultKeys: "ctrl+u", + description: "Tree filter: user messages only", + }, + "app.tree.filter.labeledOnly": { + defaultKeys: "ctrl+l", + description: "Tree filter: labeled entries only", + }, + "app.tree.filter.all": { + defaultKeys: "ctrl+a", + description: "Tree filter: show all entries", + }, + "app.tree.filter.cycleForward": { + defaultKeys: "ctrl+o", + description: "Tree filter: cycle forward", + }, + "app.tree.filter.cycleBackward": { + defaultKeys: "shift+ctrl+o", + description: "Tree filter: cycle backward", + }, +} as const satisfies KeybindingDefinitions; + +const KEYBINDING_NAME_MIGRATIONS = { + cursorUp: "tui.editor.cursorUp", + cursorDown: "tui.editor.cursorDown", + cursorLeft: "tui.editor.cursorLeft", + cursorRight: "tui.editor.cursorRight", + cursorWordLeft: "tui.editor.cursorWordLeft", + cursorWordRight: "tui.editor.cursorWordRight", + cursorLineStart: "tui.editor.cursorLineStart", + cursorLineEnd: "tui.editor.cursorLineEnd", + jumpForward: "tui.editor.jumpForward", + jumpBackward: "tui.editor.jumpBackward", + pageUp: "tui.editor.pageUp", + pageDown: "tui.editor.pageDown", + deleteCharBackward: "tui.editor.deleteCharBackward", + deleteCharForward: "tui.editor.deleteCharForward", + deleteWordBackward: "tui.editor.deleteWordBackward", + deleteWordForward: "tui.editor.deleteWordForward", + deleteToLineStart: "tui.editor.deleteToLineStart", + deleteToLineEnd: "tui.editor.deleteToLineEnd", + yank: "tui.editor.yank", + yankPop: "tui.editor.yankPop", + undo: "tui.editor.undo", + newLine: "tui.input.newLine", + submit: "tui.input.submit", + tab: "tui.input.tab", + copy: "tui.input.copy", + selectUp: "tui.select.up", + selectDown: "tui.select.down", + selectPageUp: "tui.select.pageUp", + selectPageDown: "tui.select.pageDown", + selectConfirm: "tui.select.confirm", + selectCancel: "tui.select.cancel", + interrupt: "app.interrupt", + clear: "app.clear", + exit: "app.exit", + suspend: "app.suspend", + cycleThinkingLevel: "app.thinking.cycle", + cycleModelForward: "app.model.cycleForward", + cycleModelBackward: "app.model.cycleBackward", + selectModel: "app.model.select", + expandTools: "app.tools.expand", + toggleThinking: "app.thinking.toggle", + toggleSessionNamedFilter: "app.session.toggleNamedFilter", + externalEditor: "app.editor.external", + followUp: "app.message.followUp", + dequeue: "app.message.dequeue", + pasteImage: "app.clipboard.pasteImage", + newSession: "app.session.new", + tree: "app.session.tree", + fork: "app.session.fork", + resume: "app.session.resume", + treeFoldOrUp: "app.tree.foldOrUp", + treeUnfoldOrDown: "app.tree.unfoldOrDown", + treeEditLabel: "app.tree.editLabel", + treeToggleLabelTimestamp: "app.tree.toggleLabelTimestamp", + toggleSessionPath: "app.session.togglePath", + toggleSessionSort: "app.session.toggleSort", + renameSession: "app.session.rename", + deleteSession: "app.session.delete", + deleteSessionNoninvasive: "app.session.deleteNoninvasive", +} as const satisfies Record; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function isLegacyKeybindingName(key: string): key is keyof typeof KEYBINDING_NAME_MIGRATIONS { + return key in KEYBINDING_NAME_MIGRATIONS; +} + +function toKeybindingsConfig(value: unknown): KeybindingsConfig { + if (!isRecord(value)) return {}; + + const config: KeybindingsConfig = {}; + for (const [key, binding] of Object.entries(value)) { + if (typeof binding === "string") { + config[key] = binding as KeyId; + continue; + } + if (Array.isArray(binding) && binding.every((entry) => typeof entry === "string")) { + config[key] = binding as KeyId[]; + } + } + return config; +} + +export function migrateKeybindingsConfig(rawConfig: Record): { + config: Record; + migrated: boolean; +} { + const config: Record = {}; + let migrated = false; + + for (const [key, value] of Object.entries(rawConfig)) { + const nextKey = isLegacyKeybindingName(key) ? KEYBINDING_NAME_MIGRATIONS[key] : key; + if (nextKey !== key) { + migrated = true; + } + if (key !== nextKey && Object.hasOwn(rawConfig, nextKey)) { + migrated = true; + continue; + } + config[nextKey] = value; + } + + return { config: orderKeybindingsConfig(config), migrated }; +} + +function orderKeybindingsConfig(config: Record): Record { + const ordered: Record = {}; + for (const keybinding of Object.keys(KEYBINDINGS)) { + if (Object.hasOwn(config, keybinding)) { + ordered[keybinding] = config[keybinding]; + } + } + + const extras = Object.keys(config) + .filter((key) => !Object.hasOwn(ordered, key)) + .sort(); + for (const key of extras) { + ordered[key] = config[key]; + } + + return ordered; +} + +function loadRawConfig(path: string): Record | undefined { + if (!existsSync(path)) return undefined; + try { + const parsed = JSON.parse(readFileSync(path, "utf-8")) as unknown; + return isRecord(parsed) ? parsed : undefined; + } catch { + return undefined; + } +} + +export class KeybindingsManager extends TuiKeybindingsManager { + private configPath: string | undefined; + + constructor(userBindings: KeybindingsConfig = {}, configPath?: string) { + super(KEYBINDINGS, userBindings); + this.configPath = configPath; + } + + static create(agentDir: string = getAgentDir()): KeybindingsManager { + const configPath = join(agentDir, "keybindings.json"); + const userBindings = KeybindingsManager.loadFromFile(configPath); + return new KeybindingsManager(userBindings, configPath); + } + + reload(): void { + if (!this.configPath) return; + this.setUserBindings(KeybindingsManager.loadFromFile(this.configPath)); + } + + getEffectiveConfig(): KeybindingsConfig { + return this.getResolvedBindings(); + } + + private static loadFromFile(path: string): KeybindingsConfig { + const rawConfig = loadRawConfig(path); + if (!rawConfig) return {}; + return toKeybindingsConfig(migrateKeybindingsConfig(rawConfig).config); + } +} + +export type { Keybinding, KeyId, KeybindingsConfig }; diff --git a/cactus-code/packages/coding-agent/src/core/messages.ts b/cactus-code/packages/coding-agent/src/core/messages.ts new file mode 100644 index 000000000..81ffabb58 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/messages.ts @@ -0,0 +1,195 @@ +/** + * Custom message types and transformers for the coding agent. + * + * Extends the base AgentMessage type with coding-agent specific message types, + * and provides a transformer to convert them to LLM-compatible messages. + */ + +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai"; + +export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary: + + +`; + +export const COMPACTION_SUMMARY_SUFFIX = ` +`; + +export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from: + + +`; + +export const BRANCH_SUMMARY_SUFFIX = ``; + +/** + * Message type for bash executions via the ! command. + */ +export interface BashExecutionMessage { + role: "bashExecution"; + command: string; + output: string; + exitCode: number | undefined; + cancelled: boolean; + truncated: boolean; + fullOutputPath?: string; + timestamp: number; + /** If true, this message is excluded from LLM context (!! prefix) */ + excludeFromContext?: boolean; +} + +/** + * Message type for extension-injected messages via sendMessage(). + * These are custom messages that extensions can inject into the conversation. + */ +export interface CustomMessage { + role: "custom"; + customType: string; + content: string | (TextContent | ImageContent)[]; + display: boolean; + details?: T; + timestamp: number; +} + +export interface BranchSummaryMessage { + role: "branchSummary"; + summary: string; + fromId: string; + timestamp: number; +} + +export interface CompactionSummaryMessage { + role: "compactionSummary"; + summary: string; + tokensBefore: number; + timestamp: number; +} + +// Extend CustomAgentMessages via declaration merging +declare module "@earendil-works/pi-agent-core" { + interface CustomAgentMessages { + bashExecution: BashExecutionMessage; + custom: CustomMessage; + branchSummary: BranchSummaryMessage; + compactionSummary: CompactionSummaryMessage; + } +} + +/** + * Convert a BashExecutionMessage to user message text for LLM context. + */ +export function bashExecutionToText(msg: BashExecutionMessage): string { + let text = `Ran \`${msg.command}\`\n`; + if (msg.output) { + text += `\`\`\`\n${msg.output}\n\`\`\``; + } else { + text += "(no output)"; + } + if (msg.cancelled) { + text += "\n\n(command cancelled)"; + } else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) { + text += `\n\nCommand exited with code ${msg.exitCode}`; + } + if (msg.truncated && msg.fullOutputPath) { + text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`; + } + return text; +} + +export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage { + return { + role: "branchSummary", + summary, + fromId, + timestamp: new Date(timestamp).getTime(), + }; +} + +export function createCompactionSummaryMessage( + summary: string, + tokensBefore: number, + timestamp: string, +): CompactionSummaryMessage { + return { + role: "compactionSummary", + summary: summary, + tokensBefore, + timestamp: new Date(timestamp).getTime(), + }; +} + +/** Convert CustomMessageEntry to AgentMessage format */ +export function createCustomMessage( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details: unknown | undefined, + timestamp: string, +): CustomMessage { + return { + role: "custom", + customType, + content, + display, + details, + timestamp: new Date(timestamp).getTime(), + }; +} + +/** + * Transform AgentMessages (including custom types) to LLM-compatible Messages. + * + * This is used by: + * - Agent's transormToLlm option (for prompt calls and queued messages) + * - Compaction's generateSummary (for summarization) + * - Custom extensions and tools + */ +export function convertToLlm(messages: AgentMessage[]): Message[] { + return messages + .map((m): Message | undefined => { + switch (m.role) { + case "bashExecution": + // Skip messages excluded from context (!! prefix) + if (m.excludeFromContext) { + return undefined; + } + return { + role: "user", + content: [{ type: "text", text: bashExecutionToText(m) }], + timestamp: m.timestamp, + }; + case "custom": { + const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; + return { + role: "user", + content, + timestamp: m.timestamp, + }; + } + case "branchSummary": + return { + role: "user", + content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }], + timestamp: m.timestamp, + }; + case "compactionSummary": + return { + role: "user", + content: [ + { type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX }, + ], + timestamp: m.timestamp, + }; + case "user": + case "assistant": + case "toolResult": + return m; + default: + // biome-ignore lint/correctness/noSwitchDeclarations: fine + const _exhaustiveCheck: never = m; + return undefined; + } + }) + .filter((m) => m !== undefined); +} diff --git a/cactus-code/packages/coding-agent/src/core/model-registry.ts b/cactus-code/packages/coding-agent/src/core/model-registry.ts new file mode 100644 index 000000000..70f393948 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/model-registry.ts @@ -0,0 +1,990 @@ +/** + * Model registry - manages built-in and custom models, provides API key resolution. + */ + +import { + type AnthropicMessagesCompat, + type Api, + type AssistantMessageEventStream, + type Context, + getModels, + getProviders, + type KnownProvider, + type Model, + type OAuthProviderInterface, + type OpenAICompletionsCompat, + type OpenAIResponsesCompat, + registerApiProvider, + resetApiProviders, + type SimpleStreamOptions, +} from "@earendil-works/pi-ai/compat"; +import { registerOAuthProvider, resetOAuthProviders } from "@earendil-works/pi-ai/oauth"; +import { existsSync, readFileSync } from "fs"; +import { join } from "path"; +import { type Static, Type } from "typebox"; +import { Compile } from "typebox/compile"; +import type { TLocalizedValidationError } from "typebox/error"; +import { getAgentDir } from "../config.ts"; +import { stripJsonComments } from "../utils/json.ts"; +import { normalizePath } from "../utils/paths.ts"; +import type { AuthStatus, AuthStorage } from "./auth-storage.ts"; +import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "./provider-display-names.ts"; +import { + clearConfigValueCache, + getConfigValueEnvVarNames, + isCommandConfigValue, + isConfigValueConfigured, + resolveConfigValueOrThrow, + resolveConfigValueUncached, + resolveHeadersOrThrow, +} from "./resolve-config-value.ts"; + +// Schema for OpenRouter routing preferences +const PercentileCutoffsSchema = Type.Object({ + p50: Type.Optional(Type.Number()), + p75: Type.Optional(Type.Number()), + p90: Type.Optional(Type.Number()), + p99: Type.Optional(Type.Number()), +}); + +const OpenRouterRoutingSchema = Type.Object({ + allow_fallbacks: Type.Optional(Type.Boolean()), + require_parameters: Type.Optional(Type.Boolean()), + data_collection: Type.Optional(Type.Union([Type.Literal("deny"), Type.Literal("allow")])), + zdr: Type.Optional(Type.Boolean()), + enforce_distillable_text: Type.Optional(Type.Boolean()), + order: Type.Optional(Type.Array(Type.String())), + only: Type.Optional(Type.Array(Type.String())), + ignore: Type.Optional(Type.Array(Type.String())), + quantizations: Type.Optional(Type.Array(Type.String())), + sort: Type.Optional( + Type.Union([ + Type.String(), + Type.Object({ + by: Type.Optional(Type.String()), + partition: Type.Optional(Type.Union([Type.String(), Type.Null()])), + }), + ]), + ), + max_price: Type.Optional( + Type.Object({ + prompt: Type.Optional(Type.Union([Type.Number(), Type.String()])), + completion: Type.Optional(Type.Union([Type.Number(), Type.String()])), + image: Type.Optional(Type.Union([Type.Number(), Type.String()])), + audio: Type.Optional(Type.Union([Type.Number(), Type.String()])), + request: Type.Optional(Type.Union([Type.Number(), Type.String()])), + }), + ), + preferred_min_throughput: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), + preferred_max_latency: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])), +}); + +// Schema for Vercel AI Gateway routing preferences +const VercelGatewayRoutingSchema = Type.Object({ + only: Type.Optional(Type.Array(Type.String())), + order: Type.Optional(Type.Array(Type.String())), +}); + +// Schema for thinking level support and provider-specific values +const ThinkingLevelMapValueSchema = Type.Union([Type.String(), Type.Null()]); +const ThinkingLevelMapSchema = Type.Object({ + off: Type.Optional(ThinkingLevelMapValueSchema), + minimal: Type.Optional(ThinkingLevelMapValueSchema), + low: Type.Optional(ThinkingLevelMapValueSchema), + medium: Type.Optional(ThinkingLevelMapValueSchema), + high: Type.Optional(ThinkingLevelMapValueSchema), + xhigh: Type.Optional(ThinkingLevelMapValueSchema), +}); + +const ChatTemplateKwargScalarSchema = Type.Union([Type.String(), Type.Number(), Type.Boolean(), Type.Null()]); +const ChatTemplateKwargVariableSchema = Type.Object({ + $var: Type.Union([Type.Literal("thinking.enabled"), Type.Literal("thinking.effort")]), + omitWhenOff: Type.Optional(Type.Boolean()), +}); +const ChatTemplateKwargSchema = Type.Union([ChatTemplateKwargScalarSchema, ChatTemplateKwargVariableSchema]); + +const OpenAICompletionsCompatSchema = Type.Object({ + supportsStore: Type.Optional(Type.Boolean()), + supportsDeveloperRole: Type.Optional(Type.Boolean()), + supportsReasoningEffort: Type.Optional(Type.Boolean()), + supportsUsageInStreaming: Type.Optional(Type.Boolean()), + maxTokensField: Type.Optional(Type.Union([Type.Literal("max_completion_tokens"), Type.Literal("max_tokens")])), + requiresToolResultName: Type.Optional(Type.Boolean()), + requiresAssistantAfterToolResult: Type.Optional(Type.Boolean()), + requiresThinkingAsText: Type.Optional(Type.Boolean()), + requiresReasoningContentOnAssistantMessages: Type.Optional(Type.Boolean()), + thinkingFormat: Type.Optional( + Type.Union([ + Type.Literal("openai"), + Type.Literal("openrouter"), + Type.Literal("together"), + Type.Literal("deepseek"), + Type.Literal("zai"), + Type.Literal("qwen"), + Type.Literal("chat-template"), + Type.Literal("qwen-chat-template"), + Type.Literal("string-thinking"), + Type.Literal("ant-ling"), + ]), + ), + chatTemplateKwargs: Type.Optional(Type.Record(Type.String(), ChatTemplateKwargSchema)), + cacheControlFormat: Type.Optional(Type.Literal("anthropic")), + openRouterRouting: Type.Optional(OpenRouterRoutingSchema), + vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema), + supportsStrictMode: Type.Optional(Type.Boolean()), + supportsLongCacheRetention: Type.Optional(Type.Boolean()), +}); + +const OpenAIResponsesCompatSchema = Type.Object({ + supportsDeveloperRole: Type.Optional(Type.Boolean()), + sendSessionIdHeader: Type.Optional(Type.Boolean()), + supportsLongCacheRetention: Type.Optional(Type.Boolean()), +}); + +const AnthropicMessagesCompatSchema = Type.Object({ + supportsEagerToolInputStreaming: Type.Optional(Type.Boolean()), + supportsLongCacheRetention: Type.Optional(Type.Boolean()), + sendSessionAffinityHeaders: Type.Optional(Type.Boolean()), + supportsCacheControlOnTools: Type.Optional(Type.Boolean()), + forceAdaptiveThinking: Type.Optional(Type.Boolean()), +}); + +const ProviderCompatSchema = Type.Union([ + OpenAICompletionsCompatSchema, + OpenAIResponsesCompatSchema, + AnthropicMessagesCompatSchema, +]); + +// Schema for custom model definition +// Most fields are optional with sensible defaults for local models (Ollama, LM Studio, etc.) +const ModelDefinitionSchema = Type.Object({ + id: Type.String({ minLength: 1 }), + name: Type.Optional(Type.String({ minLength: 1 })), + api: Type.Optional(Type.String({ minLength: 1 })), + baseUrl: Type.Optional(Type.String({ minLength: 1 })), + reasoning: Type.Optional(Type.Boolean()), + thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), + input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), + cost: Type.Optional( + Type.Object({ + input: Type.Number(), + output: Type.Number(), + cacheRead: Type.Number(), + cacheWrite: Type.Number(), + }), + ), + contextWindow: Type.Optional(Type.Number()), + maxTokens: Type.Optional(Type.Number()), + headers: Type.Optional(Type.Record(Type.String(), Type.String())), + compat: Type.Optional(ProviderCompatSchema), +}); + +// Schema for per-model overrides (all fields optional, merged with built-in model) +const ModelOverrideSchema = Type.Object({ + name: Type.Optional(Type.String({ minLength: 1 })), + reasoning: Type.Optional(Type.Boolean()), + thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema), + input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))), + cost: Type.Optional( + Type.Object({ + input: Type.Optional(Type.Number()), + output: Type.Optional(Type.Number()), + cacheRead: Type.Optional(Type.Number()), + cacheWrite: Type.Optional(Type.Number()), + }), + ), + contextWindow: Type.Optional(Type.Number()), + maxTokens: Type.Optional(Type.Number()), + headers: Type.Optional(Type.Record(Type.String(), Type.String())), + compat: Type.Optional(ProviderCompatSchema), +}); + +type ModelOverride = Static; + +const ProviderConfigSchema = Type.Object({ + name: Type.Optional(Type.String({ minLength: 1 })), + baseUrl: Type.Optional(Type.String({ minLength: 1 })), + apiKey: Type.Optional(Type.String({ minLength: 1 })), + api: Type.Optional(Type.String({ minLength: 1 })), + headers: Type.Optional(Type.Record(Type.String(), Type.String())), + compat: Type.Optional(ProviderCompatSchema), + authHeader: Type.Optional(Type.Boolean()), + models: Type.Optional(Type.Array(ModelDefinitionSchema)), + modelOverrides: Type.Optional(Type.Record(Type.String(), ModelOverrideSchema)), +}); + +const ModelsConfigSchema = Type.Object({ + providers: Type.Record(Type.String(), ProviderConfigSchema), +}); + +const validateModelsConfig = Compile(ModelsConfigSchema); + +type ModelsConfig = Static; + +function formatValidationPath(error: TLocalizedValidationError): string { + if (error.keyword === "required") { + const requiredProperties = (error.params as { requiredProperties?: string[] }).requiredProperties; + const requiredProperty = requiredProperties?.[0]; + if (requiredProperty) { + const basePath = error.instancePath.replace(/^\//, "").replace(/\//g, "."); + return basePath ? `${basePath}.${requiredProperty}` : requiredProperty; + } + } + const path = error.instancePath.replace(/^\//, "").replace(/\//g, "."); + return path || "root"; +} + +/** Provider override config (baseUrl, compat) without request auth/headers */ +interface ProviderOverride { + baseUrl?: string; + compat?: Model["compat"]; +} + +interface ProviderRequestConfig { + apiKey?: string; + headers?: Record; + authHeader?: boolean; +} + +export type ResolvedRequestAuth = + | { + ok: true; + apiKey?: string; + headers?: Record; + env?: Record; + } + | { + ok: false; + error: string; + }; + +/** Result of loading custom models from models.json */ +interface CustomModelsResult { + models: Model[]; + /** Providers with baseUrl/headers/apiKey overrides for built-in models */ + overrides: Map; + /** Per-model overrides: provider -> modelId -> override */ + modelOverrides: Map>; + error: string | undefined; +} + +function emptyCustomModelsResult(error?: string): CustomModelsResult { + return { models: [], overrides: new Map(), modelOverrides: new Map(), error }; +} + +function mergeCompat( + baseCompat: Model["compat"], + overrideCompat: ModelOverride["compat"], +): Model["compat"] | undefined { + if (!overrideCompat) return baseCompat; + + const base = baseCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat | undefined; + const override = overrideCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat; + const merged = { ...base, ...override } as OpenAICompletionsCompat | OpenAIResponsesCompat | AnthropicMessagesCompat; + + const baseCompletions = base as OpenAICompletionsCompat | undefined; + const overrideCompletions = override as OpenAICompletionsCompat; + const mergedCompletions = merged as OpenAICompletionsCompat; + + if (baseCompletions?.openRouterRouting || overrideCompletions.openRouterRouting) { + mergedCompletions.openRouterRouting = { + ...baseCompletions?.openRouterRouting, + ...overrideCompletions.openRouterRouting, + }; + } + + if (baseCompletions?.vercelGatewayRouting || overrideCompletions.vercelGatewayRouting) { + mergedCompletions.vercelGatewayRouting = { + ...baseCompletions?.vercelGatewayRouting, + ...overrideCompletions.vercelGatewayRouting, + }; + } + + if (baseCompletions?.chatTemplateKwargs || overrideCompletions.chatTemplateKwargs) { + mergedCompletions.chatTemplateKwargs = { + ...baseCompletions?.chatTemplateKwargs, + ...overrideCompletions.chatTemplateKwargs, + }; + } + + return merged as Model["compat"]; +} + +/** + * Deep merge a model override into a model. + * Handles nested objects (cost, compat) by merging rather than replacing. + */ +function applyModelOverride(model: Model, override: ModelOverride): Model { + const result = { ...model }; + + // Simple field overrides + if (override.name !== undefined) result.name = override.name; + if (override.reasoning !== undefined) result.reasoning = override.reasoning; + if (override.thinkingLevelMap !== undefined) { + result.thinkingLevelMap = { ...model.thinkingLevelMap, ...override.thinkingLevelMap }; + } + if (override.input !== undefined) result.input = override.input as ("text" | "image")[]; + if (override.contextWindow !== undefined) result.contextWindow = override.contextWindow; + if (override.maxTokens !== undefined) result.maxTokens = override.maxTokens; + + // Merge cost (partial override) + if (override.cost) { + result.cost = { + input: override.cost.input ?? model.cost.input, + output: override.cost.output ?? model.cost.output, + cacheRead: override.cost.cacheRead ?? model.cost.cacheRead, + cacheWrite: override.cost.cacheWrite ?? model.cost.cacheWrite, + }; + } + + // Deep merge compat + result.compat = mergeCompat(model.compat, override.compat); + + return result; +} + +/** Clear the config value command cache. Exported for testing. */ +export const clearApiKeyCache = clearConfigValueCache; + +/** + * Model registry - loads and manages models, resolves API keys via AuthStorage. + */ +export class ModelRegistry { + private models: Model[] = []; + private providerRequestConfigs: Map = new Map(); + private modelRequestHeaders: Map> = new Map(); + private registeredProviders: Map = new Map(); + private loadError: string | undefined = undefined; + readonly authStorage: AuthStorage; + private modelsJsonPath: string | undefined; + + private constructor(authStorage: AuthStorage, modelsJsonPath: string | undefined) { + this.authStorage = authStorage; + this.modelsJsonPath = modelsJsonPath ? normalizePath(modelsJsonPath) : undefined; + this.loadModels(); + } + + static create(authStorage: AuthStorage, modelsJsonPath: string = join(getAgentDir(), "models.json")): ModelRegistry { + return new ModelRegistry(authStorage, modelsJsonPath); + } + + static inMemory(authStorage: AuthStorage): ModelRegistry { + return new ModelRegistry(authStorage, undefined); + } + + /** + * Reload models from disk (built-in + custom from models.json). + */ + refresh(): void { + this.providerRequestConfigs.clear(); + this.modelRequestHeaders.clear(); + this.loadError = undefined; + + // Ensure dynamic API/OAuth registrations are rebuilt from current provider state. + resetApiProviders(); + resetOAuthProviders(); + + this.loadModels(); + + for (const [providerName, config] of this.registeredProviders.entries()) { + this.applyProviderConfig(providerName, config); + } + } + + /** + * Get any error from loading models.json (undefined if no error). + */ + getError(): string | undefined { + return this.loadError; + } + + private loadModels(): void { + // Load custom models and overrides from models.json + const { + models: customModels, + overrides, + modelOverrides, + error, + } = this.modelsJsonPath ? this.loadCustomModels(this.modelsJsonPath) : emptyCustomModelsResult(); + + if (error) { + this.loadError = error; + // Keep built-in models even if custom models failed to load + } + + const builtInModels = this.loadBuiltInModels(overrides, modelOverrides); + let combined = this.mergeCustomModels(builtInModels, customModels); + + // Let OAuth providers modify their models (e.g., update baseUrl) + for (const oauthProvider of this.authStorage.getOAuthProviders()) { + const cred = this.authStorage.get(oauthProvider.id); + if (cred?.type === "oauth" && oauthProvider.modifyModels) { + combined = oauthProvider.modifyModels(combined, cred); + } + } + + this.models = combined; + } + + /** Load built-in models and apply provider/model overrides */ + private loadBuiltInModels( + overrides: Map, + modelOverrides: Map>, + ): Model[] { + return getProviders().flatMap((provider) => { + const models = getModels(provider as KnownProvider) as Model[]; + const providerOverride = overrides.get(provider); + const perModelOverrides = modelOverrides.get(provider); + + return models.map((m) => { + let model = m; + + // Apply provider-level baseUrl/headers/compat override + if (providerOverride) { + model = { + ...model, + baseUrl: providerOverride.baseUrl ?? model.baseUrl, + compat: mergeCompat(model.compat, providerOverride.compat), + }; + } + + // Apply per-model override + const modelOverride = perModelOverrides?.get(m.id); + if (modelOverride) { + model = applyModelOverride(model, modelOverride); + } + + return model; + }); + }); + } + + /** Merge custom models into built-in list by provider+id (custom wins on conflicts). */ + private mergeCustomModels(builtInModels: Model[], customModels: Model[]): Model[] { + const merged = [...builtInModels]; + for (const customModel of customModels) { + const existingIndex = merged.findIndex((m) => m.provider === customModel.provider && m.id === customModel.id); + if (existingIndex >= 0) { + merged[existingIndex] = customModel; + } else { + merged.push(customModel); + } + } + return merged; + } + + private loadCustomModels(modelsJsonPath: string): CustomModelsResult { + if (!existsSync(modelsJsonPath)) { + return emptyCustomModelsResult(); + } + + try { + const content = readFileSync(modelsJsonPath, "utf-8"); + const parsed = JSON.parse(stripJsonComments(content)) as unknown; + + if (!validateModelsConfig.Check(parsed)) { + const errors = + validateModelsConfig + .Errors(parsed) + .map((error) => ` - ${formatValidationPath(error)}: ${error.message}`) + .join("\n") || "Unknown schema error"; + return emptyCustomModelsResult(`Invalid models.json schema:\n${errors}\n\nFile: ${modelsJsonPath}`); + } + + const config = parsed as ModelsConfig; + + // Additional validation + this.validateConfig(config); + + const overrides = new Map(); + const modelOverrides = new Map>(); + + for (const [providerName, providerConfig] of Object.entries(config.providers)) { + if (providerConfig.baseUrl || providerConfig.compat) { + overrides.set(providerName, { + baseUrl: providerConfig.baseUrl, + compat: providerConfig.compat, + }); + } + + this.storeProviderRequestConfig(providerName, providerConfig); + + if (providerConfig.modelOverrides) { + modelOverrides.set(providerName, new Map(Object.entries(providerConfig.modelOverrides))); + for (const [modelId, modelOverride] of Object.entries(providerConfig.modelOverrides)) { + this.storeModelHeaders(providerName, modelId, modelOverride.headers); + } + } + } + + return { models: this.parseModels(config), overrides, modelOverrides, error: undefined }; + } catch (error) { + if (error instanceof SyntaxError) { + return emptyCustomModelsResult(`Failed to parse models.json: ${error.message}\n\nFile: ${modelsJsonPath}`); + } + return emptyCustomModelsResult( + `Failed to load models.json: ${error instanceof Error ? error.message : error}\n\nFile: ${modelsJsonPath}`, + ); + } + } + + private validateConfig(config: ModelsConfig): void { + const builtInProviders = new Set(getProviders()); + + for (const [providerName, providerConfig] of Object.entries(config.providers)) { + const isBuiltIn = builtInProviders.has(providerName); + const hasProviderApi = !!providerConfig.api; + const models = providerConfig.models ?? []; + const hasModelOverrides = + providerConfig.modelOverrides && Object.keys(providerConfig.modelOverrides).length > 0; + + if (models.length === 0) { + // Override-only config: needs baseUrl, headers, compat, modelOverrides, or some combination. + if (!providerConfig.baseUrl && !providerConfig.headers && !providerConfig.compat && !hasModelOverrides) { + throw new Error( + `Provider ${providerName}: must specify "baseUrl", "headers", "compat", "modelOverrides", or "models".`, + ); + } + } else if (!isBuiltIn) { + // Non-built-in providers with custom models require an endpoint. + // Auth can come from auth.json, --api-key, or provider request config. + if (!providerConfig.baseUrl) { + throw new Error(`Provider ${providerName}: "baseUrl" is required when defining custom models.`); + } + } + // Built-in providers with custom models: baseUrl/apiKey/api are optional, + // inherited from built-in models. Auth comes from env vars / auth storage. + + for (const modelDef of models) { + const hasModelApi = !!modelDef.api; + + if (!hasProviderApi && !hasModelApi && !isBuiltIn) { + throw new Error( + `Provider ${providerName}, model ${modelDef.id}: no "api" specified. Set at provider or model level.`, + ); + } + // For built-in providers, api is optional — inherited from built-in models. + + if (!modelDef.id) throw new Error(`Provider ${providerName}: model missing "id"`); + // Validate contextWindow/maxTokens only if provided (they have defaults) + if (modelDef.contextWindow !== undefined && modelDef.contextWindow <= 0) + throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid contextWindow`); + if (modelDef.maxTokens !== undefined && modelDef.maxTokens <= 0) + throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid maxTokens`); + } + } + } + + private parseModels(config: ModelsConfig): Model[] { + const models: Model[] = []; + const builtInProviders = new Set(getProviders()); + + // Cache built-in defaults (api, baseUrl) per provider, extracted from first model. + const builtInDefaultsCache = new Map(); + const getBuiltInDefaults = (providerName: string): { api: string; baseUrl: string } | undefined => { + if (!builtInProviders.has(providerName)) return undefined; + if (builtInDefaultsCache.has(providerName)) return builtInDefaultsCache.get(providerName); + const builtIn = getModels(providerName as KnownProvider) as Model[]; + if (builtIn.length === 0) return undefined; + const defaults = { api: builtIn[0].api, baseUrl: builtIn[0].baseUrl }; + builtInDefaultsCache.set(providerName, defaults); + return defaults; + }; + + for (const [providerName, providerConfig] of Object.entries(config.providers)) { + const modelDefs = providerConfig.models ?? []; + if (modelDefs.length === 0) continue; // Override-only, no custom models + + const builtInDefaults = getBuiltInDefaults(providerName); + + for (const modelDef of modelDefs) { + const api = modelDef.api ?? providerConfig.api ?? builtInDefaults?.api; + if (!api) continue; + + const baseUrl = modelDef.baseUrl ?? providerConfig.baseUrl ?? builtInDefaults?.baseUrl; + if (!baseUrl) continue; + + const compat = mergeCompat(providerConfig.compat, modelDef.compat); + this.storeModelHeaders(providerName, modelDef.id, modelDef.headers); + + const defaultCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; + models.push({ + id: modelDef.id, + name: modelDef.name ?? modelDef.id, + api: api as Api, + provider: providerName, + baseUrl, + reasoning: modelDef.reasoning ?? false, + thinkingLevelMap: modelDef.thinkingLevelMap, + input: (modelDef.input ?? ["text"]) as ("text" | "image")[], + cost: modelDef.cost ?? defaultCost, + contextWindow: modelDef.contextWindow ?? 128000, + maxTokens: modelDef.maxTokens ?? 16384, + headers: undefined, + compat, + } as Model); + } + } + + return models; + } + + /** + * Get all models (built-in + custom). + * If models.json had errors, returns only built-in models. + */ + getAll(): Model[] { + return this.models; + } + + /** + * Get only models that have auth configured. + * This is a fast check that doesn't refresh OAuth tokens. + */ + getAvailable(): Model[] { + return this.models.filter((m) => this.hasConfiguredAuth(m)); + } + + /** + * Find a model by provider and ID. + */ + find(provider: string, modelId: string): Model | undefined { + return this.models.find((m) => m.provider === provider && m.id === modelId); + } + + /** + * Get API key for a model. + */ + hasConfiguredAuth(model: Model): boolean { + const providerApiKey = this.providerRequestConfigs.get(model.provider)?.apiKey; + return ( + this.authStorage.hasAuth(model.provider) || + (providerApiKey !== undefined && isConfigValueConfigured(providerApiKey)) + ); + } + + private getModelRequestKey(provider: string, modelId: string): string { + return `${provider}:${modelId}`; + } + + private storeProviderRequestConfig( + providerName: string, + config: { + apiKey?: string; + headers?: Record; + authHeader?: boolean; + }, + ): void { + if (!config.apiKey && !config.headers && !config.authHeader) { + return; + } + + this.providerRequestConfigs.set(providerName, { + apiKey: config.apiKey, + headers: config.headers, + authHeader: config.authHeader, + }); + } + + private storeModelHeaders(providerName: string, modelId: string, headers?: Record): void { + const key = this.getModelRequestKey(providerName, modelId); + if (!headers || Object.keys(headers).length === 0) { + this.modelRequestHeaders.delete(key); + return; + } + this.modelRequestHeaders.set(key, headers); + } + + /** + * Get API key and request headers for a model. + */ + async getApiKeyAndHeaders(model: Model): Promise { + try { + const providerConfig = this.providerRequestConfigs.get(model.provider); + const providerEnv = this.authStorage.getProviderEnv(model.provider); + const apiKeyFromAuthStorage = await this.authStorage.getApiKey(model.provider, { includeFallback: false }); + const apiKey = + apiKeyFromAuthStorage ?? + (providerConfig?.apiKey + ? resolveConfigValueOrThrow( + providerConfig.apiKey, + `API key for provider "${model.provider}"`, + providerEnv, + ) + : undefined); + + const providerHeaders = resolveHeadersOrThrow( + providerConfig?.headers, + `provider "${model.provider}"`, + providerEnv, + ); + const modelHeaders = resolveHeadersOrThrow( + this.modelRequestHeaders.get(this.getModelRequestKey(model.provider, model.id)), + `model "${model.provider}/${model.id}"`, + providerEnv, + ); + + let headers = + model.headers || providerHeaders || modelHeaders + ? { ...model.headers, ...providerHeaders, ...modelHeaders } + : undefined; + + if (providerConfig?.authHeader) { + if (!apiKey) { + return { ok: false, error: `No API key found for "${model.provider}"` }; + } + headers = { ...headers, Authorization: `Bearer ${apiKey}` }; + } + + return { + ok: true, + apiKey, + headers: headers && Object.keys(headers).length > 0 ? headers : undefined, + env: providerEnv && Object.keys(providerEnv).length > 0 ? providerEnv : undefined, + }; + } catch (error) { + return { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + } + + /** + * Return auth status for a provider, including request auth configured in models.json. + * This intentionally does not execute command-backed config values. + */ + getProviderAuthStatus(provider: string): AuthStatus { + const authStatus = this.authStorage.getAuthStatus(provider); + if (authStatus.source) { + return authStatus; + } + + const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; + if (!providerApiKey) { + return authStatus; + } + + if (isCommandConfigValue(providerApiKey)) { + return { configured: true, source: "models_json_command" }; + } + + const envVarNames = getConfigValueEnvVarNames(providerApiKey); + if (envVarNames.length > 0) { + return isConfigValueConfigured(providerApiKey) + ? { configured: true, source: "environment", label: envVarNames.join(", ") } + : { configured: false }; + } + + return { configured: true, source: "models_json_key" }; + } + + /** + * Get display name for a provider. + */ + getProviderDisplayName(provider: string): string { + const registeredProvider = this.registeredProviders.get(provider); + const oauthProvider = this.authStorage.getOAuthProviders().find((p) => p.id === provider); + + return ( + registeredProvider?.name ?? + registeredProvider?.oauth?.name ?? + oauthProvider?.name ?? + BUILT_IN_PROVIDER_DISPLAY_NAMES[provider] ?? + provider + ); + } + + /** + * Get API key for a provider. + */ + async getApiKeyForProvider(provider: string): Promise { + const apiKey = await this.authStorage.getApiKey(provider); + if (apiKey !== undefined) { + return apiKey; + } + + const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey; + return providerApiKey + ? resolveConfigValueUncached(providerApiKey, this.authStorage.getProviderEnv(provider)) + : undefined; + } + + /** + * Check if a model is using OAuth credentials (subscription). + */ + isUsingOAuth(model: Model): boolean { + const cred = this.authStorage.get(model.provider); + return cred?.type === "oauth"; + } + + /** + * Register a provider dynamically (from extensions). + * + * If provider has models: replaces all existing models for this provider. + * If provider has only baseUrl/headers: overrides existing models' URLs. + * If provider has oauth: registers OAuth provider for /login support. + */ + registerProvider(providerName: string, config: ProviderConfigInput): void { + this.validateProviderConfig(providerName, config); + this.applyProviderConfig(providerName, config); + this.upsertRegisteredProvider(providerName, config); + } + + /** + * Unregister a previously registered provider. + * + * Removes the provider from the registry and reloads models from disk so that + * built-in models overridden by this provider are restored to their original state. + * Also resets dynamic OAuth and API stream registrations before reapplying + * remaining dynamic providers. + * Has no effect if the provider was never registered. + */ + unregisterProvider(providerName: string): void { + if (!this.registeredProviders.has(providerName)) return; + this.registeredProviders.delete(providerName); + this.refresh(); + } + + /** + * Upsert a provider config into registeredProviders. + * If the provider is already registered, defined values in the incoming config + * override existing ones; undefined values are preserved from the stored config. + * If the provider is not registered, the incoming config is stored as-is. + */ + private upsertRegisteredProvider(providerName: string, config: ProviderConfigInput): void { + const existing = this.registeredProviders.get(providerName); + if (!existing) { + this.registeredProviders.set(providerName, config); + return; + } + for (const k of Object.keys(config) as (keyof ProviderConfigInput)[]) { + if (config[k] !== undefined) { + (existing as Record)[k] = config[k]; + } + } + } + + private validateProviderConfig(providerName: string, config: ProviderConfigInput): void { + if (config.streamSimple && !config.api) { + throw new Error(`Provider ${providerName}: "api" is required when registering streamSimple.`); + } + + if (!config.models || config.models.length === 0) { + return; + } + + if (!config.baseUrl) { + throw new Error(`Provider ${providerName}: "baseUrl" is required when defining models.`); + } + if (!config.apiKey && !config.oauth) { + throw new Error(`Provider ${providerName}: "apiKey" or "oauth" is required when defining models.`); + } + + for (const modelDef of config.models) { + const api = modelDef.api || config.api; + if (!api) { + throw new Error(`Provider ${providerName}, model ${modelDef.id}: no "api" specified.`); + } + } + } + + private applyProviderConfig(providerName: string, config: ProviderConfigInput): void { + // Register OAuth provider if provided + if (config.oauth) { + // Ensure the OAuth provider ID matches the provider name + const oauthProvider: OAuthProviderInterface = { + ...config.oauth, + id: providerName, + }; + registerOAuthProvider(oauthProvider); + } + + if (config.streamSimple) { + const streamSimple = config.streamSimple; + registerApiProvider( + { + api: config.api!, + stream: (model, context, options) => streamSimple(model, context, options as SimpleStreamOptions), + streamSimple, + }, + `provider:${providerName}`, + ); + } + + this.storeProviderRequestConfig(providerName, config); + + if (config.models && config.models.length > 0) { + // Full replacement: remove existing models for this provider + this.models = this.models.filter((m) => m.provider !== providerName); + + // Parse and add new models + for (const modelDef of config.models) { + const api = modelDef.api || config.api; + this.storeModelHeaders(providerName, modelDef.id, modelDef.headers); + + this.models.push({ + id: modelDef.id, + name: modelDef.name, + api: api as Api, + provider: providerName, + baseUrl: modelDef.baseUrl ?? config.baseUrl!, + reasoning: modelDef.reasoning, + thinkingLevelMap: modelDef.thinkingLevelMap, + input: modelDef.input as ("text" | "image")[], + cost: modelDef.cost, + contextWindow: modelDef.contextWindow, + maxTokens: modelDef.maxTokens, + headers: undefined, + compat: modelDef.compat, + } as Model); + } + + // Apply OAuth modifyModels if credentials exist (e.g., to update baseUrl) + if (config.oauth?.modifyModels) { + const cred = this.authStorage.get(providerName); + if (cred?.type === "oauth") { + this.models = config.oauth.modifyModels(this.models, cred); + } + } + } else if (config.baseUrl || config.headers) { + // Override-only: update baseUrl for existing models. Request headers are resolved per request. + this.models = this.models.map((m) => { + if (m.provider !== providerName) return m; + return { + ...m, + baseUrl: config.baseUrl ?? m.baseUrl, + }; + }); + } + } +} + +/** + * Input type for registerProvider API. + */ +export interface ProviderConfigInput { + name?: string; + baseUrl?: string; + apiKey?: string; + api?: Api; + streamSimple?: (model: Model, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream; + headers?: Record; + authHeader?: boolean; + /** OAuth provider for /login support */ + oauth?: Omit; + models?: Array<{ + id: string; + name: string; + api?: Api; + baseUrl?: string; + reasoning: boolean; + thinkingLevelMap?: Model["thinkingLevelMap"]; + input: ("text" | "image")[]; + cost: { input: number; output: number; cacheRead: number; cacheWrite: number }; + contextWindow: number; + maxTokens: number; + headers?: Record; + compat?: Model["compat"]; + }>; +} diff --git a/cactus-code/packages/coding-agent/src/core/model-resolver.ts b/cactus-code/packages/coding-agent/src/core/model-resolver.ts new file mode 100644 index 000000000..5616bc12f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/model-resolver.ts @@ -0,0 +1,647 @@ +/** + * Model resolution, scoping, and initial selection + */ + +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { type Api, type KnownProvider, type Model, modelsAreEqual } from "@earendil-works/pi-ai"; +import chalk from "chalk"; +import { minimatch } from "minimatch"; +import { isValidThinkingLevel } from "../cli/args.ts"; +import { DEFAULT_THINKING_LEVEL } from "./defaults.ts"; +import type { ModelRegistry } from "./model-registry.ts"; + +/** Default model IDs for each known provider */ +export const defaultModelPerProvider: Record = { + cactus: "cactus", +}; + +export interface ScopedModel { + model: Model; + /** Thinking level if explicitly specified in pattern (e.g., "model:high"), undefined otherwise */ + thinkingLevel?: ThinkingLevel; +} + +/** + * Helper to check if a model ID looks like an alias (no date suffix) + * Dates are typically in format: -20241022 or -20250929 + */ +function isAlias(id: string): boolean { + // Check if ID ends with -latest + if (id.endsWith("-latest")) return true; + + // Check if ID ends with a date pattern (-YYYYMMDD) + const datePattern = /-\d{8}$/; + return !datePattern.test(id); +} + +/** + * Find an exact model reference match. + * Supports either a bare model id or a canonical provider/modelId reference. + * When matching by bare id, ambiguous matches across providers are rejected. + */ +export function findExactModelReferenceMatch( + modelReference: string, + availableModels: Model[], +): Model | undefined { + const trimmedReference = modelReference.trim(); + if (!trimmedReference) { + return undefined; + } + + const normalizedReference = trimmedReference.toLowerCase(); + + const canonicalMatches = availableModels.filter( + (model) => `${model.provider}/${model.id}`.toLowerCase() === normalizedReference, + ); + if (canonicalMatches.length === 1) { + return canonicalMatches[0]; + } + if (canonicalMatches.length > 1) { + return undefined; + } + + const slashIndex = trimmedReference.indexOf("/"); + if (slashIndex !== -1) { + const provider = trimmedReference.substring(0, slashIndex).trim(); + const modelId = trimmedReference.substring(slashIndex + 1).trim(); + if (provider && modelId) { + const providerMatches = availableModels.filter( + (model) => + model.provider.toLowerCase() === provider.toLowerCase() && + model.id.toLowerCase() === modelId.toLowerCase(), + ); + if (providerMatches.length === 1) { + return providerMatches[0]; + } + if (providerMatches.length > 1) { + return undefined; + } + } + } + + const idMatches = availableModels.filter((model) => model.id.toLowerCase() === normalizedReference); + return idMatches.length === 1 ? idMatches[0] : undefined; +} + +/** + * Try to match a pattern to a model from the available models list. + * Returns the matched model or undefined if no match found. + */ +function tryMatchModel(modelPattern: string, availableModels: Model[]): Model | undefined { + const exactMatch = findExactModelReferenceMatch(modelPattern, availableModels); + if (exactMatch) { + return exactMatch; + } + + // No exact match - fall back to partial matching + const matches = availableModels.filter( + (m) => + m.id.toLowerCase().includes(modelPattern.toLowerCase()) || + m.name?.toLowerCase().includes(modelPattern.toLowerCase()), + ); + + if (matches.length === 0) { + return undefined; + } + + // Separate into aliases and dated versions + const aliases = matches.filter((m) => isAlias(m.id)); + const datedVersions = matches.filter((m) => !isAlias(m.id)); + + if (aliases.length > 0) { + // Prefer alias - if multiple aliases, pick the one that sorts highest + aliases.sort((a, b) => b.id.localeCompare(a.id)); + return aliases[0]; + } else { + // No alias found, pick latest dated version + datedVersions.sort((a, b) => b.id.localeCompare(a.id)); + return datedVersions[0]; + } +} + +export interface ParsedModelResult { + model: Model | undefined; + /** Thinking level if explicitly specified in pattern, undefined otherwise */ + thinkingLevel?: ThinkingLevel; + warning: string | undefined; +} + +function buildFallbackModel(provider: string, modelId: string, availableModels: Model[]): Model | undefined { + const providerModels = availableModels.filter((m) => m.provider === provider); + if (providerModels.length === 0) return undefined; + + const defaultId = defaultModelPerProvider[provider as KnownProvider]; + const baseModel = defaultId + ? (providerModels.find((m) => m.id === defaultId) ?? providerModels[0]) + : providerModels[0]; + + return { + ...baseModel, + id: modelId, + name: modelId, + }; +} + +/** + * Parse a pattern to extract model and thinking level. + * Handles models with colons in their IDs (e.g., OpenRouter's :exacto suffix). + * + * Algorithm: + * 1. Try to match full pattern as a model + * 2. If found, return it with "off" thinking level + * 3. If not found and has colons, split on last colon: + * - If suffix is valid thinking level, use it and recurse on prefix + * - If suffix is invalid, warn and recurse on prefix with "off" + * + * @internal Exported for testing + */ +export function parseModelPattern( + pattern: string, + availableModels: Model[], + options?: { allowInvalidThinkingLevelFallback?: boolean }, +): ParsedModelResult { + // Try exact match first + const exactMatch = tryMatchModel(pattern, availableModels); + if (exactMatch) { + return { model: exactMatch, thinkingLevel: undefined, warning: undefined }; + } + + // No match - try splitting on last colon if present + const lastColonIndex = pattern.lastIndexOf(":"); + if (lastColonIndex === -1) { + // No colons, pattern simply doesn't match any model + return { model: undefined, thinkingLevel: undefined, warning: undefined }; + } + + const prefix = pattern.substring(0, lastColonIndex); + const suffix = pattern.substring(lastColonIndex + 1); + + if (isValidThinkingLevel(suffix)) { + // Valid thinking level - recurse on prefix and use this level + const result = parseModelPattern(prefix, availableModels, options); + if (result.model) { + // Only use this thinking level if no warning from inner recursion + return { + model: result.model, + thinkingLevel: result.warning ? undefined : suffix, + warning: result.warning, + }; + } + return result; + } else { + // Invalid suffix + const allowFallback = options?.allowInvalidThinkingLevelFallback ?? true; + if (!allowFallback) { + // In strict mode (CLI --model parsing), treat it as part of the model id and fail. + // This avoids accidentally resolving to a different model. + return { model: undefined, thinkingLevel: undefined, warning: undefined }; + } + + // Scope mode: recurse on prefix and warn + const result = parseModelPattern(prefix, availableModels, options); + if (result.model) { + return { + model: result.model, + thinkingLevel: undefined, + warning: `Invalid thinking level "${suffix}" in pattern "${pattern}". Using default instead.`, + }; + } + return result; + } +} + +/** + * Resolve model patterns to actual Model objects with optional thinking levels + * Format: "pattern:level" where :level is optional + * For each pattern, finds all matching models and picks the best version: + * 1. Prefer alias (e.g., claude-sonnet-4-5) over dated versions (claude-sonnet-4-5-20250929) + * 2. If no alias, pick the latest dated version + * + * Supports models with colons in their IDs (e.g., OpenRouter's model:exacto). + * The algorithm tries to match the full pattern first, then progressively + * strips colon-suffixes to find a match. + */ +export async function resolveModelScope(patterns: string[], modelRegistry: ModelRegistry): Promise { + const availableModels = await modelRegistry.getAvailable(); + const scopedModels: ScopedModel[] = []; + + for (const pattern of patterns) { + // Check if pattern contains glob characters + if (pattern.includes("*") || pattern.includes("?") || pattern.includes("[")) { + // Extract optional thinking level suffix (e.g., "provider/*:high") + const colonIdx = pattern.lastIndexOf(":"); + let globPattern = pattern; + let thinkingLevel: ThinkingLevel | undefined; + + if (colonIdx !== -1) { + const suffix = pattern.substring(colonIdx + 1); + if (isValidThinkingLevel(suffix)) { + thinkingLevel = suffix; + globPattern = pattern.substring(0, colonIdx); + } + } + + // Match against "provider/modelId" format OR just model ID + // This allows "*sonnet*" to match without requiring "anthropic/*sonnet*" + const matchingModels = availableModels.filter((m) => { + const fullId = `${m.provider}/${m.id}`; + return minimatch(fullId, globPattern, { nocase: true }) || minimatch(m.id, globPattern, { nocase: true }); + }); + + if (matchingModels.length === 0) { + console.warn(chalk.yellow(`Warning: No models match pattern "${pattern}"`)); + continue; + } + + for (const model of matchingModels) { + if (!scopedModels.find((sm) => modelsAreEqual(sm.model, model))) { + scopedModels.push({ model, thinkingLevel }); + } + } + continue; + } + + const { model, thinkingLevel, warning } = parseModelPattern(pattern, availableModels); + + if (warning) { + console.warn(chalk.yellow(`Warning: ${warning}`)); + } + + if (!model) { + console.warn(chalk.yellow(`Warning: No models match pattern "${pattern}"`)); + continue; + } + + // Avoid duplicates + if (!scopedModels.find((sm) => modelsAreEqual(sm.model, model))) { + scopedModels.push({ model, thinkingLevel }); + } + } + + return scopedModels; +} + +export interface ResolveCliModelResult { + model: Model | undefined; + thinkingLevel?: ThinkingLevel; + warning: string | undefined; + /** + * Error message suitable for CLI display. + * When set, model will be undefined. + */ + error: string | undefined; +} + +/** + * Resolve a single model from CLI flags. + * + * Supports: + * - --provider --model + * - --model / + * - Fuzzy matching (same rules as model scoping: exact id, then partial id/name) + * + * Note: This does not apply the thinking level by itself, but it may *parse* and + * return a thinking level from ":" so the caller can apply it. + */ +export function resolveCliModel(options: { + cliProvider?: string; + cliModel?: string; + cliThinking?: ThinkingLevel; + modelRegistry: ModelRegistry; +}): ResolveCliModelResult { + const { cliProvider, cliModel, cliThinking, modelRegistry } = options; + + if (!cliModel) { + return { model: undefined, warning: undefined, error: undefined }; + } + + // Important: use *all* models here, not just models with pre-configured auth. + // This allows "--api-key" to be used for first-time setup. + const availableModels = modelRegistry.getAll(); + if (availableModels.length === 0) { + return { + model: undefined, + warning: undefined, + error: "No models available. Check your installation or add models to models.json.", + }; + } + + // Build canonical provider lookup (case-insensitive) + const providerMap = new Map(); + for (const m of availableModels) { + providerMap.set(m.provider.toLowerCase(), m.provider); + } + + let provider = cliProvider ? providerMap.get(cliProvider.toLowerCase()) : undefined; + if (cliProvider && !provider) { + return { + model: undefined, + warning: undefined, + error: `Unknown provider "${cliProvider}". Use --list-models to see available providers/models.`, + }; + } + + // If no explicit --provider, try to interpret "provider/model" format first. + // When the prefix before the first slash matches a known provider, prefer that + // interpretation over matching models whose IDs literally contain slashes + // (e.g. "zai/glm-5" should resolve to provider=zai, model=glm-5, not to a + // vercel-ai-gateway model with id "zai/glm-5"). + let pattern = cliModel; + let inferredProvider = false; + + if (!provider) { + const slashIndex = cliModel.indexOf("/"); + if (slashIndex !== -1) { + const maybeProvider = cliModel.substring(0, slashIndex); + const canonical = providerMap.get(maybeProvider.toLowerCase()); + if (canonical) { + provider = canonical; + pattern = cliModel.substring(slashIndex + 1); + inferredProvider = true; + } + } + } + + // If no provider was inferred from the slash, try exact matches without provider inference. + // This handles models whose IDs naturally contain slashes (e.g. OpenRouter-style IDs). + if (!provider) { + const lower = cliModel.toLowerCase(); + const exact = availableModels.find( + (m) => m.id.toLowerCase() === lower || `${m.provider}/${m.id}`.toLowerCase() === lower, + ); + if (exact) { + return { model: exact, warning: undefined, thinkingLevel: undefined, error: undefined }; + } + } + + if (cliProvider && provider) { + // If both were provided, tolerate --model / by stripping the provider prefix + const prefix = `${provider}/`; + if (cliModel.toLowerCase().startsWith(prefix.toLowerCase())) { + pattern = cliModel.substring(prefix.length); + } + } + + const candidates = provider ? availableModels.filter((m) => m.provider === provider) : availableModels; + const { model, thinkingLevel, warning } = parseModelPattern(pattern, candidates, { + allowInvalidThinkingLevelFallback: false, + }); + + if (model) { + // If provider inference matched an unauthenticated provider/model pair, prefer + // one exact raw model-id match that is authenticated. This keeps + // "provider/model" syntax preferred when usable, but handles models whose + // literal id starts with a known provider name (for example + // commandcode model id "xiaomi/mimo-v2.5-pro"). + if (inferredProvider) { + const rawExactMatches = availableModels.filter( + (m) => m.id.toLowerCase() === cliModel.toLowerCase() && !modelsAreEqual(m, model), + ); + if (rawExactMatches.length > 0 && !modelRegistry.hasConfiguredAuth(model)) { + const authenticatedRawMatches = rawExactMatches.filter((m) => modelRegistry.hasConfiguredAuth(m)); + if (authenticatedRawMatches.length === 1) { + return { + model: authenticatedRawMatches[0], + thinkingLevel: undefined, + warning: undefined, + error: undefined, + }; + } + } + } + return { model, thinkingLevel, warning, error: undefined }; + } + + // If we inferred a provider from the slash but found no match within that provider, + // fall back to matching the full input as a raw model id across all models. + // This handles OpenRouter-style IDs like "openai/gpt-4o:extended" where "openai" + // looks like a provider but the full string is actually a model id on openrouter. + if (inferredProvider) { + const lower = cliModel.toLowerCase(); + const exact = availableModels.find( + (m) => m.id.toLowerCase() === lower || `${m.provider}/${m.id}`.toLowerCase() === lower, + ); + if (exact) { + return { model: exact, warning: undefined, thinkingLevel: undefined, error: undefined }; + } + // Also try parseModelPattern on the full input against all models + const fallback = parseModelPattern(cliModel, availableModels, { + allowInvalidThinkingLevelFallback: false, + }); + if (fallback.model) { + return { + model: fallback.model, + thinkingLevel: fallback.thinkingLevel, + warning: fallback.warning, + error: undefined, + }; + } + } + + if (provider) { + // Parse thinking level suffix from the pattern before building the fallback model, + // but only when --thinking is not explicitly provided. + // e.g. "zai-org/GLM-5.1-FP8:high" → modelId="zai-org/GLM-5.1-FP8", fallbackThinking="high" + let fallbackPattern = pattern; + let fallbackThinking: ThinkingLevel | undefined; + if (!cliThinking) { + const lastColon = pattern.lastIndexOf(":"); + if (lastColon !== -1) { + const suffix = pattern.substring(lastColon + 1); + if (isValidThinkingLevel(suffix)) { + fallbackPattern = pattern.substring(0, lastColon); + fallbackThinking = suffix; + } + } + } + + const fallbackModel = buildFallbackModel(provider, fallbackPattern, availableModels); + if (fallbackModel) { + const requestedThinking = cliThinking ?? fallbackThinking; + const model = + requestedThinking && requestedThinking !== "off" ? { ...fallbackModel, reasoning: true } : fallbackModel; + const fallbackWarning = warning + ? `${warning} Model "${fallbackPattern}" not found for provider "${provider}". Using custom model id.` + : `Model "${fallbackPattern}" not found for provider "${provider}". Using custom model id.`; + return { model, thinkingLevel: fallbackThinking, warning: fallbackWarning, error: undefined }; + } + } + + const display = provider ? `${provider}/${pattern}` : cliModel; + return { + model: undefined, + thinkingLevel: undefined, + warning, + error: `Model "${display}" not found. Use --list-models to see available models.`, + }; +} + +export interface InitialModelResult { + model: Model | undefined; + thinkingLevel: ThinkingLevel; + fallbackMessage: string | undefined; +} + +/** + * Find the initial model to use based on priority: + * 1. CLI args (provider + model) + * 2. First model from scoped models (if not continuing/resuming) + * 3. Restored from session (if continuing/resuming) + * 4. Saved default from settings + * 5. First available model with valid API key + */ +export async function findInitialModel(options: { + cliProvider?: string; + cliModel?: string; + scopedModels: ScopedModel[]; + isContinuing: boolean; + defaultProvider?: string; + defaultModelId?: string; + defaultThinkingLevel?: ThinkingLevel; + modelRegistry: ModelRegistry; +}): Promise { + const { + cliProvider, + cliModel, + scopedModels, + isContinuing, + defaultProvider, + defaultModelId, + defaultThinkingLevel, + modelRegistry, + } = options; + + let model: Model | undefined; + let thinkingLevel: ThinkingLevel = DEFAULT_THINKING_LEVEL; + + // 1. CLI args take priority + if (cliProvider && cliModel) { + const resolved = resolveCliModel({ + cliProvider, + cliModel, + modelRegistry, + }); + if (resolved.error) { + console.error(chalk.red(resolved.error)); + process.exit(1); + } + if (resolved.model) { + return { model: resolved.model, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; + } + } + + // 2. Use first model from scoped models (skip if continuing/resuming) + if (scopedModels.length > 0 && !isContinuing) { + return { + model: scopedModels[0].model, + thinkingLevel: scopedModels[0].thinkingLevel ?? defaultThinkingLevel ?? DEFAULT_THINKING_LEVEL, + fallbackMessage: undefined, + }; + } + + // 3. Try saved default from settings + if (defaultProvider && defaultModelId) { + const found = modelRegistry.find(defaultProvider, defaultModelId); + if (found) { + model = found; + if (defaultThinkingLevel) { + thinkingLevel = defaultThinkingLevel; + } + return { model, thinkingLevel, fallbackMessage: undefined }; + } + } + + // 4. Try first available model with valid API key + const availableModels = await modelRegistry.getAvailable(); + + if (availableModels.length > 0) { + // Try to find a default model from known providers + for (const provider of Object.keys(defaultModelPerProvider) as KnownProvider[]) { + const defaultId = defaultModelPerProvider[provider]; + const match = availableModels.find((m) => m.provider === provider && m.id === defaultId); + if (match) { + return { model: match, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; + } + } + + // If no default found, use first available + return { model: availableModels[0], thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; + } + + // 5. No model found + return { model: undefined, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; +} + +/** + * Restore model from session, with fallback to available models + */ +export async function restoreModelFromSession( + savedProvider: string, + savedModelId: string, + currentModel: Model | undefined, + shouldPrintMessages: boolean, + modelRegistry: ModelRegistry, +): Promise<{ model: Model | undefined; fallbackMessage: string | undefined }> { + const restoredModel = modelRegistry.find(savedProvider, savedModelId); + + // Check if restored model exists and still has auth configured + const hasConfiguredAuth = restoredModel ? modelRegistry.hasConfiguredAuth(restoredModel) : false; + + if (restoredModel && hasConfiguredAuth) { + if (shouldPrintMessages) { + console.log(chalk.dim(`Restored model: ${savedProvider}/${savedModelId}`)); + } + return { model: restoredModel, fallbackMessage: undefined }; + } + + // Model not found or no API key - fall back + const reason = !restoredModel ? "model no longer exists" : "no auth configured"; + + if (shouldPrintMessages) { + console.error(chalk.yellow(`Warning: Could not restore model ${savedProvider}/${savedModelId} (${reason}).`)); + } + + // If we already have a model, use it as fallback + if (currentModel) { + if (shouldPrintMessages) { + console.log(chalk.dim(`Falling back to: ${currentModel.provider}/${currentModel.id}`)); + } + return { + model: currentModel, + fallbackMessage: `Could not restore model ${savedProvider}/${savedModelId} (${reason}). Using ${currentModel.provider}/${currentModel.id}.`, + }; + } + + // Try to find any available model + const availableModels = await modelRegistry.getAvailable(); + + if (availableModels.length > 0) { + // Try to find a default model from known providers + let fallbackModel: Model | undefined; + for (const provider of Object.keys(defaultModelPerProvider) as KnownProvider[]) { + const defaultId = defaultModelPerProvider[provider]; + const match = availableModels.find((m) => m.provider === provider && m.id === defaultId); + if (match) { + fallbackModel = match; + break; + } + } + + // If no default found, use first available + if (!fallbackModel) { + fallbackModel = availableModels[0]; + } + + if (shouldPrintMessages) { + console.log(chalk.dim(`Falling back to: ${fallbackModel.provider}/${fallbackModel.id}`)); + } + + return { + model: fallbackModel, + fallbackMessage: `Could not restore model ${savedProvider}/${savedModelId} (${reason}). Using ${fallbackModel.provider}/${fallbackModel.id}.`, + }; + } + + // No models available + return { model: undefined, fallbackMessage: undefined }; +} diff --git a/cactus-code/packages/coding-agent/src/core/output-guard.ts b/cactus-code/packages/coding-agent/src/core/output-guard.ts new file mode 100644 index 000000000..8781a51dc --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/output-guard.ts @@ -0,0 +1,108 @@ +interface StdoutTakeoverState { + rawStdoutWrite: (chunk: string, callback?: (error?: Error | null) => void) => boolean; + rawStderrWrite: (chunk: string, callback?: (error?: Error | null) => void) => boolean; + originalStdoutWrite: typeof process.stdout.write; +} + +let stdoutTakeoverState: StdoutTakeoverState | undefined; + +const RAW_STDOUT_RETRY_DELAY_MS = 10; + +let rawStdoutWriteTail: Promise = Promise.resolve(); + +function getRawStdoutWrite(): StdoutTakeoverState["rawStdoutWrite"] { + if (stdoutTakeoverState) { + return stdoutTakeoverState.rawStdoutWrite; + } + return process.stdout.write.bind(process.stdout) as StdoutTakeoverState["rawStdoutWrite"]; +} + +async function writeRawStdoutChunk(text: string): Promise { + while (true) { + try { + await new Promise((resolve, reject) => { + try { + getRawStdoutWrite()(text, (error) => { + if (error) reject(error); + else resolve(); + }); + } catch (error) { + reject(error instanceof Error ? error : new Error(String(error))); + } + }); + return; + } catch (error) { + const writeError = error instanceof Error ? error : new Error(String(error)); + const code = (writeError as Error & { code?: unknown }).code; + if (code !== "ENOBUFS" && code !== "EAGAIN" && code !== "EWOULDBLOCK") { + throw writeError; + } + await new Promise((resolve) => setTimeout(resolve, RAW_STDOUT_RETRY_DELAY_MS)); + } + } +} + +export function takeOverStdout(): void { + if (stdoutTakeoverState) { + return; + } + + const rawStdoutWrite = process.stdout.write.bind(process.stdout) as StdoutTakeoverState["rawStdoutWrite"]; + const rawStderrWrite = process.stderr.write.bind(process.stderr) as StdoutTakeoverState["rawStderrWrite"]; + const originalStdoutWrite = process.stdout.write; + + process.stdout.write = (( + chunk: string | Uint8Array, + encodingOrCallback?: BufferEncoding | ((error?: Error | null) => void), + callback?: (error?: Error | null) => void, + ): boolean => { + if (typeof encodingOrCallback === "function") { + return rawStderrWrite(String(chunk), encodingOrCallback); + } + return rawStderrWrite(String(chunk), callback); + }) as typeof process.stdout.write; + + stdoutTakeoverState = { + rawStdoutWrite, + rawStderrWrite, + originalStdoutWrite, + }; +} + +export function restoreStdout(): void { + if (!stdoutTakeoverState) { + return; + } + + process.stdout.write = stdoutTakeoverState.originalStdoutWrite; + stdoutTakeoverState = undefined; +} + +export function isStdoutTakenOver(): boolean { + return stdoutTakeoverState !== undefined; +} + +export function writeRawStdout(text: string): void { + if (text.length === 0) { + return; + } + rawStdoutWriteTail = rawStdoutWriteTail.then(() => writeRawStdoutChunk(text)); + void rawStdoutWriteTail.catch(() => { + process.exit(1); + }); +} + +export async function waitForRawStdoutBackpressure(): Promise { + while (true) { + const tail = rawStdoutWriteTail; + await tail; + if (tail === rawStdoutWriteTail) { + return; + } + } +} + +export async function flushRawStdout(): Promise { + await waitForRawStdoutBackpressure(); + await writeRawStdoutChunk(""); +} diff --git a/cactus-code/packages/coding-agent/src/core/package-manager.ts b/cactus-code/packages/coding-agent/src/core/package-manager.ts new file mode 100644 index 000000000..ae81663d0 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/package-manager.ts @@ -0,0 +1,2561 @@ +import type { ChildProcess, ChildProcessByStdio } from "node:child_process"; +import { createHash } from "node:crypto"; +import { chmodSync, existsSync, mkdirSync, readdirSync, readFileSync, rmSync, statSync, writeFileSync } from "node:fs"; +import { homedir } from "node:os"; + +function getEnv(): NodeJS.ProcessEnv { + if (process.platform !== "linux" || Object.keys(process.env).length > 0) { + return process.env; + } + try { + const data = readFileSync("/proc/self/environ", "utf-8"); + const env: NodeJS.ProcessEnv = {}; + for (const entry of data.split("\0")) { + const idx = entry.indexOf("="); + if (idx > 0) { + env[entry.slice(0, idx)] = entry.slice(idx + 1); + } + } + return env; + } catch { + return process.env; + } +} + +import { basename, dirname, join, relative, resolve, sep } from "node:path"; +import type { Readable } from "node:stream"; +import { globSync } from "glob"; +import ignore from "ignore"; +import { minimatch } from "minimatch"; +import { maxSatisfying, rcompare, satisfies, valid, validRange } from "semver"; +import { CONFIG_DIR_NAME } from "../config.ts"; +import { spawnProcess, spawnProcessSync } from "../utils/child-process.ts"; +import { type GitSource, parseGitUrl } from "../utils/git.ts"; +import { canonicalizePath, isLocalPath, markPathIgnoredByCloudSync, resolvePath } from "../utils/paths.ts"; +import { isStdoutTakenOver } from "./output-guard.ts"; +import type { PackageSource, SettingsManager } from "./settings-manager.ts"; + +const NETWORK_TIMEOUT_MS = 10000; +const UPDATE_CHECK_CONCURRENCY = 4; +const GIT_UPDATE_CONCURRENCY = 4; + +function isOfflineModeEnabled(): boolean { + const value = process.env.PI_OFFLINE; + if (!value) return false; + return value === "1" || value.toLowerCase() === "true" || value.toLowerCase() === "yes"; +} + +function isExactNpmVersion(version: string | undefined): boolean { + return valid(version ?? "") !== null; +} + +function getNpmVersionRange(version: string | undefined): string | undefined { + return version ? (validRange(version) ?? undefined) : undefined; +} + +export interface PathMetadata { + source: string; + scope: SourceScope; + origin: "package" | "top-level"; + baseDir?: string; +} + +export interface ResolvedResource { + path: string; + enabled: boolean; + metadata: PathMetadata; +} + +export interface ResolvedPaths { + extensions: ResolvedResource[]; + skills: ResolvedResource[]; + prompts: ResolvedResource[]; + themes: ResolvedResource[]; +} + +export type MissingSourceAction = "install" | "skip" | "error"; + +export interface ProgressEvent { + type: "start" | "progress" | "complete" | "error"; + action: "install" | "remove" | "update" | "clone" | "pull"; + source: string; + message?: string; +} + +export type ProgressCallback = (event: ProgressEvent) => void; + +export interface PackageUpdate { + source: string; + displayName: string; + type: "npm" | "git"; + scope: Exclude; +} + +export interface ConfiguredPackage { + source: string; + scope: "user" | "project"; + filtered: boolean; + installedPath?: string; +} + +export interface PackageManager { + resolve(onMissing?: (source: string) => Promise): Promise; + install(source: string, options?: { local?: boolean }): Promise; + installAndPersist(source: string, options?: { local?: boolean }): Promise; + remove(source: string, options?: { local?: boolean }): Promise; + removeAndPersist(source: string, options?: { local?: boolean }): Promise; + update(source?: string): Promise; + listConfiguredPackages(): ConfiguredPackage[]; + resolveExtensionSources( + sources: string[], + options?: { local?: boolean; temporary?: boolean }, + ): Promise; + addSourceToSettings(source: string, options?: { local?: boolean }): boolean; + removeSourceFromSettings(source: string, options?: { local?: boolean }): boolean; + setProgressCallback(callback: ProgressCallback | undefined): void; + getInstalledPath(source: string, scope: "user" | "project"): string | undefined; +} + +interface PackageManagerOptions { + cwd: string; + agentDir: string; + settingsManager: SettingsManager; +} + +type SourceScope = "user" | "project" | "temporary"; + +type NpmSource = { + type: "npm"; + spec: string; + name: string; + version?: string; + range?: string; + pinned: boolean; +}; + +type LocalSource = { + type: "local"; + path: string; +}; + +type ParsedSource = NpmSource | GitSource | LocalSource; + +type InstalledSourceScope = Exclude; + +interface ConfiguredUpdateSource { + source: string; + scope: InstalledSourceScope; +} + +interface NpmUpdateTarget extends ConfiguredUpdateSource { + parsed: NpmSource; +} + +interface GitUpdateTarget extends ConfiguredUpdateSource { + parsed: GitSource; +} + +interface PiManifest { + extensions?: string[]; + skills?: string[]; + prompts?: string[]; + themes?: string[]; +} + +interface ResourceAccumulator { + extensions: Map; + skills: Map; + prompts: Map; + themes: Map; +} + +/** + * Compute a numeric precedence rank for a resource based on its metadata. + * Lower rank = higher precedence. Used to sort resolved resources so that + * name-collision resolution ("first wins") produces the correct outcome. + * + * Precedence (highest to lowest): + * 0 project + settings entry (source: "local", scope: "project") + * 1 project + auto-discovered (source: "auto", scope: "project") + * 2 user + settings entry (source: "local", scope: "user") + * 3 user + auto-discovered (source: "auto", scope: "user") + * 4 package resource (origin: "package") + */ +function resourcePrecedenceRank(m: PathMetadata): number { + if (m.origin === "package") return 4; + const scopeBase = m.scope === "project" ? 0 : 2; + return scopeBase + (m.source === "local" ? 0 : 1); +} + +interface PackageFilter { + extensions?: string[]; + skills?: string[]; + prompts?: string[]; + themes?: string[]; +} + +type ResourceType = "extensions" | "skills" | "prompts" | "themes"; + +const RESOURCE_TYPES: ResourceType[] = ["extensions", "skills", "prompts", "themes"]; + +const FILE_PATTERNS: Record = { + extensions: /\.(ts|js)$/, + skills: /\.md$/, + prompts: /\.md$/, + themes: /\.json$/, +}; + +const IGNORE_FILE_NAMES = [".gitignore", ".ignore", ".fdignore"]; + +type IgnoreMatcher = ReturnType; + +function toPosixPath(p: string): string { + return p.split(sep).join("/"); +} + +function getHomeDir(): string { + return process.env.HOME || homedir(); +} + +export function getExtensionTempFolder(agentDir: string): string { + const tempFolder = join(agentDir, "tmp", "extensions"); + mkdirSync(tempFolder, { recursive: true, mode: 0o700 }); + chmodSync(tempFolder, 0o700); + return tempFolder; +} + +function prefixIgnorePattern(line: string, prefix: string): string | null { + const trimmed = line.trim(); + if (!trimmed) return null; + if (trimmed.startsWith("#") && !trimmed.startsWith("\\#")) return null; + + let pattern = line; + let negated = false; + + if (pattern.startsWith("!")) { + negated = true; + pattern = pattern.slice(1); + } else if (pattern.startsWith("\\!")) { + pattern = pattern.slice(1); + } + + if (pattern.startsWith("/")) { + pattern = pattern.slice(1); + } + + const prefixed = prefix ? `${prefix}${pattern}` : pattern; + return negated ? `!${prefixed}` : prefixed; +} + +function addIgnoreRules(ig: IgnoreMatcher, dir: string, rootDir: string): void { + const relativeDir = relative(rootDir, dir); + const prefix = relativeDir ? `${toPosixPath(relativeDir)}/` : ""; + + for (const filename of IGNORE_FILE_NAMES) { + const ignorePath = join(dir, filename); + if (!existsSync(ignorePath)) continue; + try { + const content = readFileSync(ignorePath, "utf-8"); + const patterns = content + .split(/\r?\n/) + .map((line) => prefixIgnorePattern(line, prefix)) + .filter((line): line is string => Boolean(line)); + if (patterns.length > 0) { + ig.add(patterns); + } + } catch {} + } +} + +function isPattern(s: string): boolean { + return s.startsWith("!") || s.startsWith("+") || s.startsWith("-") || s.includes("*") || s.includes("?"); +} + +function isOverridePattern(s: string): boolean { + return s.startsWith("!") || s.startsWith("+") || s.startsWith("-"); +} + +function hasGlobPattern(s: string): boolean { + return s.includes("*") || s.includes("?"); +} + +function splitPatterns(entries: string[]): { plain: string[]; patterns: string[] } { + const plain: string[] = []; + const patterns: string[] = []; + for (const entry of entries) { + if (isPattern(entry)) { + patterns.push(entry); + } else { + plain.push(entry); + } + } + return { plain, patterns }; +} + +function collectFiles( + dir: string, + filePattern: RegExp, + skipNodeModules = true, + ignoreMatcher?: IgnoreMatcher, + rootDir?: string, +): string[] { + const files: string[] = []; + if (!existsSync(dir)) return files; + + const root = rootDir ?? dir; + const ig = ignoreMatcher ?? ignore(); + addIgnoreRules(ig, dir, root); + + try { + const entries = readdirSync(dir, { withFileTypes: true }); + for (const entry of entries) { + if (entry.name.startsWith(".")) continue; + if (skipNodeModules && entry.name === "node_modules") continue; + + const fullPath = join(dir, entry.name); + let isDir = entry.isDirectory(); + let isFile = entry.isFile(); + + if (entry.isSymbolicLink()) { + try { + const stats = statSync(fullPath); + isDir = stats.isDirectory(); + isFile = stats.isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(root, fullPath)); + const ignorePath = isDir ? `${relPath}/` : relPath; + if (ig.ignores(ignorePath)) continue; + + if (isDir) { + files.push(...collectFiles(fullPath, filePattern, skipNodeModules, ig, root)); + } else if (isFile && filePattern.test(entry.name)) { + files.push(fullPath); + } + } + } catch { + // Ignore errors + } + + return files; +} + +type SkillDiscoveryMode = "pi" | "agents"; + +function collectSkillEntries( + dir: string, + mode: SkillDiscoveryMode, + ignoreMatcher?: IgnoreMatcher, + rootDir?: string, +): string[] { + const entries: string[] = []; + if (!existsSync(dir)) return entries; + + const root = rootDir ?? dir; + const ig = ignoreMatcher ?? ignore(); + addIgnoreRules(ig, dir, root); + + try { + const dirEntries = readdirSync(dir, { withFileTypes: true }); + + for (const entry of dirEntries) { + if (entry.name !== "SKILL.md") { + continue; + } + + const fullPath = join(dir, entry.name); + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + isFile = statSync(fullPath).isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(root, fullPath)); + if (isFile && !ig.ignores(relPath)) { + entries.push(fullPath); + return entries; + } + } + + for (const entry of dirEntries) { + if (entry.name.startsWith(".")) continue; + if (entry.name === "node_modules") continue; + + const fullPath = join(dir, entry.name); + let isDir = entry.isDirectory(); + let isFile = entry.isFile(); + + if (entry.isSymbolicLink()) { + try { + const stats = statSync(fullPath); + isDir = stats.isDirectory(); + isFile = stats.isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(root, fullPath)); + if (mode === "pi" && dir === root && isFile && entry.name.endsWith(".md") && !ig.ignores(relPath)) { + entries.push(fullPath); + continue; + } + + if (!isDir) continue; + if (ig.ignores(`${relPath}/`)) continue; + + entries.push(...collectSkillEntries(fullPath, mode, ig, root)); + } + } catch { + // Ignore errors + } + + return entries; +} + +function collectAutoSkillEntries(dir: string, mode: SkillDiscoveryMode): string[] { + return collectSkillEntries(dir, mode); +} + +function findGitRepoRoot(startDir: string): string | null { + let dir = resolve(startDir); + while (true) { + if (existsSync(join(dir, ".git"))) { + return dir; + } + const parent = dirname(dir); + if (parent === dir) { + return null; + } + dir = parent; + } +} + +function collectAncestorAgentsSkillDirs(startDir: string): string[] { + const skillDirs: string[] = []; + const resolvedStartDir = resolve(startDir); + const gitRepoRoot = findGitRepoRoot(resolvedStartDir); + + let dir = resolvedStartDir; + while (true) { + skillDirs.push(join(dir, ".agents", "skills")); + if (gitRepoRoot && dir === gitRepoRoot) { + break; + } + const parent = dirname(dir); + if (parent === dir) { + break; + } + dir = parent; + } + + return skillDirs; +} + +function collectAutoPromptEntries(dir: string): string[] { + const entries: string[] = []; + if (!existsSync(dir)) return entries; + + const ig = ignore(); + addIgnoreRules(ig, dir, dir); + + try { + const dirEntries = readdirSync(dir, { withFileTypes: true }); + for (const entry of dirEntries) { + if (entry.name.startsWith(".")) continue; + if (entry.name === "node_modules") continue; + + const fullPath = join(dir, entry.name); + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + isFile = statSync(fullPath).isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(dir, fullPath)); + if (ig.ignores(relPath)) continue; + + if (isFile && entry.name.endsWith(".md")) { + entries.push(fullPath); + } + } + } catch { + // Ignore errors + } + + return entries; +} + +function collectAutoThemeEntries(dir: string): string[] { + const entries: string[] = []; + if (!existsSync(dir)) return entries; + + const ig = ignore(); + addIgnoreRules(ig, dir, dir); + + try { + const dirEntries = readdirSync(dir, { withFileTypes: true }); + for (const entry of dirEntries) { + if (entry.name.startsWith(".")) continue; + if (entry.name === "node_modules") continue; + + const fullPath = join(dir, entry.name); + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + isFile = statSync(fullPath).isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(dir, fullPath)); + if (ig.ignores(relPath)) continue; + + if (isFile && entry.name.endsWith(".json")) { + entries.push(fullPath); + } + } + } catch { + // Ignore errors + } + + return entries; +} + +function readPiManifestFile(packageJsonPath: string): PiManifest | null { + try { + const content = readFileSync(packageJsonPath, "utf-8"); + const pkg = JSON.parse(content) as { pi?: PiManifest }; + return pkg.pi ?? null; + } catch { + return null; + } +} + +function resolveExtensionEntries(dir: string): string[] | null { + const packageJsonPath = join(dir, "package.json"); + if (existsSync(packageJsonPath)) { + const manifest = readPiManifestFile(packageJsonPath); + if (manifest?.extensions?.length) { + const entries: string[] = []; + for (const extPath of manifest.extensions) { + const resolvedExtPath = resolve(dir, extPath); + if (existsSync(resolvedExtPath)) { + entries.push(resolvedExtPath); + } + } + if (entries.length > 0) { + return entries; + } + } + } + + const indexTs = join(dir, "index.ts"); + const indexJs = join(dir, "index.js"); + if (existsSync(indexTs)) { + return [indexTs]; + } + if (existsSync(indexJs)) { + return [indexJs]; + } + + return null; +} + +function collectAutoExtensionEntries(dir: string): string[] { + const entries: string[] = []; + if (!existsSync(dir)) return entries; + + // First check if this directory itself has explicit extension entries (package.json or index) + const rootEntries = resolveExtensionEntries(dir); + if (rootEntries) { + return rootEntries; + } + + // Otherwise, discover extensions from directory contents + const ig = ignore(); + addIgnoreRules(ig, dir, dir); + + try { + const dirEntries = readdirSync(dir, { withFileTypes: true }); + for (const entry of dirEntries) { + if (entry.name.startsWith(".")) continue; + if (entry.name === "node_modules") continue; + + const fullPath = join(dir, entry.name); + let isDir = entry.isDirectory(); + let isFile = entry.isFile(); + + if (entry.isSymbolicLink()) { + try { + const stats = statSync(fullPath); + isDir = stats.isDirectory(); + isFile = stats.isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(dir, fullPath)); + const ignorePath = isDir ? `${relPath}/` : relPath; + if (ig.ignores(ignorePath)) continue; + + if (isFile && (entry.name.endsWith(".ts") || entry.name.endsWith(".js"))) { + entries.push(fullPath); + } else if (isDir) { + const resolvedEntries = resolveExtensionEntries(fullPath); + if (resolvedEntries) { + entries.push(...resolvedEntries); + } + } + } + } catch { + // Ignore errors + } + + return entries; +} + +/** + * Collect resource files from a directory based on resource type. + * Extensions use smart discovery (index.ts in subdirs), others use recursive collection. + */ +function collectResourceFiles(dir: string, resourceType: ResourceType): string[] { + if (resourceType === "skills") { + return collectSkillEntries(dir, "pi"); + } + if (resourceType === "extensions") { + return collectAutoExtensionEntries(dir); + } + return collectFiles(dir, FILE_PATTERNS[resourceType]); +} + +function matchesAnyPattern(filePath: string, patterns: string[], baseDir: string): boolean { + const rel = toPosixPath(relative(baseDir, filePath)); + const name = basename(filePath); + const filePathPosix = toPosixPath(filePath); + const isSkillFile = name === "SKILL.md"; + const parentDir = isSkillFile ? dirname(filePath) : undefined; + const parentRel = isSkillFile ? toPosixPath(relative(baseDir, parentDir!)) : undefined; + const parentName = isSkillFile ? basename(parentDir!) : undefined; + const parentDirPosix = isSkillFile ? toPosixPath(parentDir!) : undefined; + + return patterns.some((pattern) => { + const normalizedPattern = toPosixPath(pattern); + if ( + minimatch(rel, normalizedPattern) || + minimatch(name, normalizedPattern) || + minimatch(filePathPosix, normalizedPattern) + ) { + return true; + } + if (!isSkillFile) return false; + return ( + minimatch(parentRel!, normalizedPattern) || + minimatch(parentName!, normalizedPattern) || + minimatch(parentDirPosix!, normalizedPattern) + ); + }); +} + +function normalizeExactPattern(pattern: string): string { + const normalized = pattern.startsWith("./") || pattern.startsWith(".\\") ? pattern.slice(2) : pattern; + return toPosixPath(normalized); +} + +function matchesAnyExactPattern(filePath: string, patterns: string[], baseDir: string): boolean { + if (patterns.length === 0) return false; + const rel = toPosixPath(relative(baseDir, filePath)); + const name = basename(filePath); + const filePathPosix = toPosixPath(filePath); + const isSkillFile = name === "SKILL.md"; + const parentDir = isSkillFile ? dirname(filePath) : undefined; + const parentRel = isSkillFile ? toPosixPath(relative(baseDir, parentDir!)) : undefined; + const parentDirPosix = isSkillFile ? toPosixPath(parentDir!) : undefined; + + return patterns.some((pattern) => { + const normalized = normalizeExactPattern(pattern); + if (normalized === rel || normalized === filePathPosix) { + return true; + } + if (!isSkillFile) return false; + return normalized === parentRel || normalized === parentDirPosix; + }); +} + +function getOverridePatterns(entries: string[]): string[] { + return entries.filter((pattern) => pattern.startsWith("!") || pattern.startsWith("+") || pattern.startsWith("-")); +} + +function isEnabledByOverrides(filePath: string, patterns: string[], baseDir: string): boolean { + const overrides = getOverridePatterns(patterns); + const excludes = overrides.filter((pattern) => pattern.startsWith("!")).map((pattern) => pattern.slice(1)); + const forceIncludes = overrides.filter((pattern) => pattern.startsWith("+")).map((pattern) => pattern.slice(1)); + const forceExcludes = overrides.filter((pattern) => pattern.startsWith("-")).map((pattern) => pattern.slice(1)); + + let enabled = true; + if (excludes.length > 0 && matchesAnyPattern(filePath, excludes, baseDir)) { + enabled = false; + } + if (forceIncludes.length > 0 && matchesAnyExactPattern(filePath, forceIncludes, baseDir)) { + enabled = true; + } + if (forceExcludes.length > 0 && matchesAnyExactPattern(filePath, forceExcludes, baseDir)) { + enabled = false; + } + return enabled; +} + +/** + * Apply patterns to paths and return a Set of enabled paths. + * Pattern types: + * - Plain patterns: include matching paths + * - `!pattern`: exclude matching paths + * - `+path`: force-include exact path (overrides exclusions) + * - `-path`: force-exclude exact path (overrides force-includes) + */ +function applyPatterns(allPaths: string[], patterns: string[], baseDir: string): Set { + const includes: string[] = []; + const excludes: string[] = []; + const forceIncludes: string[] = []; + const forceExcludes: string[] = []; + + for (const p of patterns) { + if (p.startsWith("+")) { + forceIncludes.push(p.slice(1)); + } else if (p.startsWith("-")) { + forceExcludes.push(p.slice(1)); + } else if (p.startsWith("!")) { + excludes.push(p.slice(1)); + } else { + includes.push(p); + } + } + + // Step 1: Apply includes (or all if no includes) + let result: string[]; + if (includes.length === 0) { + result = [...allPaths]; + } else { + result = allPaths.filter((filePath) => matchesAnyPattern(filePath, includes, baseDir)); + } + + // Step 2: Apply excludes + if (excludes.length > 0) { + result = result.filter((filePath) => !matchesAnyPattern(filePath, excludes, baseDir)); + } + + // Step 3: Force-include (add back from allPaths, overriding exclusions) + if (forceIncludes.length > 0) { + for (const filePath of allPaths) { + if (!result.includes(filePath) && matchesAnyExactPattern(filePath, forceIncludes, baseDir)) { + result.push(filePath); + } + } + } + + // Step 4: Force-exclude (remove even if included or force-included) + if (forceExcludes.length > 0) { + result = result.filter((filePath) => !matchesAnyExactPattern(filePath, forceExcludes, baseDir)); + } + + return new Set(result); +} + +export class DefaultPackageManager implements PackageManager { + private cwd: string; + private agentDir: string; + private settingsManager: SettingsManager; + private globalNpmRoot: string | undefined; + private globalNpmRootCommandKey: string | undefined; + private progressCallback: ProgressCallback | undefined; + + constructor(options: PackageManagerOptions) { + this.cwd = resolvePath(options.cwd); + this.agentDir = resolvePath(options.agentDir); + this.settingsManager = options.settingsManager; + } + + setProgressCallback(callback: ProgressCallback | undefined): void { + this.progressCallback = callback; + } + + addSourceToSettings(source: string, options?: { local?: boolean }): boolean { + const scope: SourceScope = options?.local ? "project" : "user"; + const currentSettings = + scope === "project" ? this.settingsManager.getProjectSettings() : this.settingsManager.getGlobalSettings(); + const currentPackages = currentSettings.packages ?? []; + const normalizedSource = this.normalizePackageSourceForSettings(source, scope); + const matchIndex = currentPackages.findIndex((existing) => this.packageSourcesMatch(existing, source, scope)); + if (matchIndex !== -1) { + const existing = currentPackages[matchIndex]; + if (this.getPackageSourceString(existing) === normalizedSource) { + return false; + } + const nextPackages = [...currentPackages]; + nextPackages[matchIndex] = + typeof existing === "string" ? normalizedSource : { ...existing, source: normalizedSource }; + if (scope === "project") { + this.settingsManager.setProjectPackages(nextPackages); + } else { + this.settingsManager.setPackages(nextPackages); + } + return true; + } + const nextPackages = [...currentPackages, normalizedSource]; + if (scope === "project") { + this.settingsManager.setProjectPackages(nextPackages); + } else { + this.settingsManager.setPackages(nextPackages); + } + return true; + } + + removeSourceFromSettings(source: string, options?: { local?: boolean }): boolean { + const scope: SourceScope = options?.local ? "project" : "user"; + const currentSettings = + scope === "project" ? this.settingsManager.getProjectSettings() : this.settingsManager.getGlobalSettings(); + const currentPackages = currentSettings.packages ?? []; + const nextPackages = currentPackages.filter((existing) => !this.packageSourcesMatch(existing, source, scope)); + const changed = nextPackages.length !== currentPackages.length; + if (!changed) { + return false; + } + if (scope === "project") { + this.settingsManager.setProjectPackages(nextPackages); + } else { + this.settingsManager.setPackages(nextPackages); + } + return true; + } + + getInstalledPath(source: string, scope: "user" | "project"): string | undefined { + const parsed = this.parseSource(source); + if (parsed.type === "npm") { + const path = this.getNpmInstallPath(parsed, scope); + return existsSync(path) ? path : undefined; + } + if (parsed.type === "git") { + const path = this.getGitInstallPath(parsed, scope); + return existsSync(path) ? path : undefined; + } + if (parsed.type === "local") { + const baseDir = this.getBaseDirForScope(scope); + const path = this.resolvePathFromBase(parsed.path, baseDir); + return existsSync(path) ? path : undefined; + } + return undefined; + } + + private emitProgress(event: ProgressEvent): void { + this.progressCallback?.(event); + } + + private async withProgress( + action: ProgressEvent["action"], + source: string, + message: string, + operation: () => Promise, + ): Promise { + this.emitProgress({ type: "start", action, source, message }); + try { + await operation(); + this.emitProgress({ type: "complete", action, source }); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + this.emitProgress({ type: "error", action, source, message: errorMessage }); + throw error; + } + } + + async resolve(onMissing?: (source: string) => Promise): Promise { + const accumulator = this.createAccumulator(); + const globalSettings = this.settingsManager.getGlobalSettings(); + const projectSettings = this.settingsManager.getProjectSettings(); + + // Collect all packages with scope (project first so cwd resources win collisions) + const allPackages: Array<{ pkg: PackageSource; scope: SourceScope }> = []; + for (const pkg of projectSettings.packages ?? []) { + allPackages.push({ pkg, scope: "project" }); + } + for (const pkg of globalSettings.packages ?? []) { + allPackages.push({ pkg, scope: "user" }); + } + + // Dedupe: project scope wins over global for same package identity + const packageSources = this.dedupePackages(allPackages); + await this.resolvePackageSources(packageSources, accumulator, onMissing); + + const globalBaseDir = this.agentDir; + const projectBaseDir = join(this.cwd, CONFIG_DIR_NAME); + + for (const resourceType of RESOURCE_TYPES) { + const target = this.getTargetMap(accumulator, resourceType); + const globalEntries = (globalSettings[resourceType] ?? []) as string[]; + const projectEntries = (projectSettings[resourceType] ?? []) as string[]; + this.resolveLocalEntries( + projectEntries, + resourceType, + target, + { + source: "local", + scope: "project", + origin: "top-level", + }, + projectBaseDir, + ); + this.resolveLocalEntries( + globalEntries, + resourceType, + target, + { + source: "local", + scope: "user", + origin: "top-level", + }, + globalBaseDir, + ); + } + + this.addAutoDiscoveredResources(accumulator, globalSettings, projectSettings, globalBaseDir, projectBaseDir); + + return this.toResolvedPaths(accumulator); + } + + async resolveExtensionSources( + sources: string[], + options?: { local?: boolean; temporary?: boolean }, + ): Promise { + const accumulator = this.createAccumulator(); + const scope: SourceScope = options?.temporary ? "temporary" : options?.local ? "project" : "user"; + const packageSources = sources.map((source) => ({ pkg: source as PackageSource, scope })); + await this.resolvePackageSources(packageSources, accumulator); + return this.toResolvedPaths(accumulator); + } + + listConfiguredPackages(): ConfiguredPackage[] { + const globalSettings = this.settingsManager.getGlobalSettings(); + const projectSettings = this.settingsManager.getProjectSettings(); + const configuredPackages: ConfiguredPackage[] = []; + + for (const pkg of globalSettings.packages ?? []) { + const source = typeof pkg === "string" ? pkg : pkg.source; + configuredPackages.push({ + source, + scope: "user", + filtered: typeof pkg === "object", + installedPath: this.getInstalledPath(source, "user"), + }); + } + + for (const pkg of projectSettings.packages ?? []) { + const source = typeof pkg === "string" ? pkg : pkg.source; + configuredPackages.push({ + source, + scope: "project", + filtered: typeof pkg === "object", + installedPath: this.getInstalledPath(source, "project"), + }); + } + + return configuredPackages; + } + + async install(source: string, options?: { local?: boolean }): Promise { + const parsed = this.parseSource(source); + const scope: SourceScope = options?.local ? "project" : "user"; + this.assertProjectTrustedForScope(scope); + await this.withProgress("install", source, `Installing ${source}...`, async () => { + if (parsed.type === "npm") { + await this.installNpm(parsed, scope, false); + return; + } + if (parsed.type === "git") { + await this.installGit(parsed, scope); + return; + } + if (parsed.type === "local") { + const resolved = this.resolvePath(parsed.path); + if (!existsSync(resolved)) { + throw new Error(`Path does not exist: ${resolved}`); + } + return; + } + throw new Error(`Unsupported install source: ${source}`); + }); + } + + async installAndPersist(source: string, options?: { local?: boolean }): Promise { + await this.install(source, options); + this.addSourceToSettings(source, options); + } + + async remove(source: string, options?: { local?: boolean }): Promise { + const parsed = this.parseSource(source); + const scope: SourceScope = options?.local ? "project" : "user"; + this.assertProjectTrustedForScope(scope); + await this.withProgress("remove", source, `Removing ${source}...`, async () => { + if (parsed.type === "npm") { + await this.uninstallNpm(parsed, scope); + return; + } + if (parsed.type === "git") { + await this.removeGit(parsed, scope); + return; + } + if (parsed.type === "local") { + return; + } + throw new Error(`Unsupported remove source: ${source}`); + }); + } + + async removeAndPersist(source: string, options?: { local?: boolean }): Promise { + await this.remove(source, options); + return this.removeSourceFromSettings(source, options); + } + + async update(source?: string): Promise { + const globalSettings = this.settingsManager.getGlobalSettings(); + const projectSettings = this.settingsManager.getProjectSettings(); + const identity = source ? this.getPackageIdentity(source) : undefined; + let matched = false; + const updateSources: ConfiguredUpdateSource[] = []; + + for (const pkg of globalSettings.packages ?? []) { + const sourceStr = typeof pkg === "string" ? pkg : pkg.source; + if (identity && this.getPackageIdentity(sourceStr, "user") !== identity) continue; + matched = true; + updateSources.push({ source: sourceStr, scope: "user" }); + } + for (const pkg of projectSettings.packages ?? []) { + const sourceStr = typeof pkg === "string" ? pkg : pkg.source; + if (identity && this.getPackageIdentity(sourceStr, "project") !== identity) continue; + matched = true; + updateSources.push({ source: sourceStr, scope: "project" }); + } + + if (source && !matched) { + throw new Error( + this.buildNoMatchingPackageMessage(source, [ + ...(globalSettings.packages ?? []), + ...(projectSettings.packages ?? []), + ]), + ); + } + + await this.updateConfiguredSources(updateSources); + } + + private async updateConfiguredSources(sources: ConfiguredUpdateSource[]): Promise { + if (isOfflineModeEnabled() || sources.length === 0) { + return; + } + + const npmCandidates: NpmUpdateTarget[] = []; + const gitCandidates: GitUpdateTarget[] = []; + + for (const entry of sources) { + const parsed = this.parseSource(entry.source); + // Pinned npm versions are fixed. Pinned git refs are configured checkout targets, + // so include them to reconcile an existing clone when the configured ref changes. + if (parsed.type === "npm") { + if (!parsed.pinned) { + npmCandidates.push({ ...entry, parsed }); + } + } else if (parsed.type === "git") { + gitCandidates.push({ ...entry, parsed }); + } + } + + const npmCheckTasks = npmCandidates.map((entry) => async () => ({ + entry, + shouldUpdate: await this.shouldUpdateNpmSource(entry.parsed, entry.scope), + })); + const npmCheckResults = await this.runWithConcurrency(npmCheckTasks, UPDATE_CHECK_CONCURRENCY); + const userNpmUpdates: NpmUpdateTarget[] = []; + const projectNpmUpdates: NpmUpdateTarget[] = []; + for (const result of npmCheckResults) { + if (!result.shouldUpdate) { + continue; + } + if (result.entry.scope === "user") { + userNpmUpdates.push(result.entry); + } else { + projectNpmUpdates.push(result.entry); + } + } + + const tasks: Promise[] = []; + if (userNpmUpdates.length > 0) { + tasks.push(this.updateNpmBatch(userNpmUpdates, "user")); + } + if (projectNpmUpdates.length > 0) { + tasks.push(this.updateNpmBatch(projectNpmUpdates, "project")); + } + if (gitCandidates.length > 0) { + const gitTasks = gitCandidates.map( + (entry) => async () => + this.withProgress("update", entry.source, `Updating ${entry.source}...`, async () => { + await this.updateGit(entry.parsed, entry.scope); + }), + ); + tasks.push(this.runWithConcurrency(gitTasks, GIT_UPDATE_CONCURRENCY).then(() => {})); + } + + await Promise.all(tasks); + } + + private async shouldUpdateNpmSource(source: NpmSource, scope: InstalledSourceScope): Promise { + const installedPath = this.getManagedNpmInstallPath(source, scope); + const installedVersion = existsSync(installedPath) ? this.getInstalledNpmVersion(installedPath) : undefined; + if (!installedVersion) { + return true; + } + + try { + const targetVersion = await this.getLatestNpmVersion(source.version ? source.spec : source.name, source.range); + return targetVersion !== installedVersion; + } catch { + // Preserve existing update behavior when version lookup fails. + return true; + } + } + + private async updateNpmBatch(sources: NpmUpdateTarget[], scope: InstalledSourceScope): Promise { + if (sources.length === 0) { + return; + } + + const sourceLabel = sources.length === 1 ? sources[0].source : `${scope} npm packages`; + const message = sources.length === 1 ? `Updating ${sources[0].source}...` : `Updating ${scope} npm packages...`; + const specs = sources.map((entry) => (entry.parsed.version ? entry.parsed.spec : `${entry.parsed.name}@latest`)); + + await this.withProgress("update", sourceLabel, message, async () => { + await this.installNpmBatch(specs, scope); + }); + } + + private async installNpmBatch(_specs: string[], _scope: InstalledSourceScope): Promise { + throw new Error("Package installation is disabled in this Cactus build."); + } + + async checkForAvailableUpdates(): Promise { + if (isOfflineModeEnabled()) { + return []; + } + + const globalSettings = this.settingsManager.getGlobalSettings(); + const projectSettings = this.settingsManager.getProjectSettings(); + const allPackages: Array<{ pkg: PackageSource; scope: SourceScope }> = []; + for (const pkg of projectSettings.packages ?? []) { + allPackages.push({ pkg, scope: "project" }); + } + for (const pkg of globalSettings.packages ?? []) { + allPackages.push({ pkg, scope: "user" }); + } + + const packageSources = this.dedupePackages(allPackages); + const checks = packageSources + .filter( + (entry): entry is { pkg: PackageSource; scope: Exclude } => + entry.scope !== "temporary", + ) + .map((entry) => async (): Promise => { + const source = typeof entry.pkg === "string" ? entry.pkg : entry.pkg.source; + const parsed = this.parseSource(source); + if (parsed.type === "local" || parsed.pinned) { + return undefined; + } + + if (parsed.type === "npm") { + const installedPath = this.getNpmInstallPath(parsed, entry.scope); + if (!existsSync(installedPath)) { + return undefined; + } + const hasUpdate = await this.npmHasAvailableUpdate(parsed, installedPath); + if (!hasUpdate) { + return undefined; + } + return { + source, + displayName: parsed.name, + type: "npm", + scope: entry.scope, + }; + } + + const installedPath = this.getGitInstallPath(parsed, entry.scope); + if (!existsSync(installedPath)) { + return undefined; + } + const hasUpdate = await this.gitHasAvailableUpdate(installedPath); + if (!hasUpdate) { + return undefined; + } + return { + source, + displayName: `${parsed.host}/${parsed.path}`, + type: "git", + scope: entry.scope, + }; + }); + + const results = await this.runWithConcurrency(checks, UPDATE_CHECK_CONCURRENCY); + return results.filter((result): result is PackageUpdate => result !== undefined); + } + + private async resolvePackageSources( + sources: Array<{ pkg: PackageSource; scope: SourceScope }>, + accumulator: ResourceAccumulator, + onMissing?: (source: string) => Promise, + ): Promise { + for (const { pkg, scope } of sources) { + const sourceStr = typeof pkg === "string" ? pkg : pkg.source; + const filter = typeof pkg === "object" ? pkg : undefined; + const parsed = this.parseSource(sourceStr); + const metadata: PathMetadata = { source: sourceStr, scope, origin: "package" }; + + if (parsed.type === "local") { + const baseDir = this.getBaseDirForScope(scope); + this.resolveLocalExtensionSource(parsed, accumulator, filter, metadata, baseDir); + continue; + } + + const installMissing = async (): Promise => { + if (isOfflineModeEnabled()) { + return false; + } + if (!onMissing) { + await this.installParsedSource(parsed, scope); + return true; + } + const action = await onMissing(sourceStr); + if (action === "skip") return false; + if (action === "error") throw new Error(`Missing source: ${sourceStr}`); + await this.installParsedSource(parsed, scope); + return true; + }; + + if (parsed.type === "npm") { + let installedPath = this.getNpmInstallPath(parsed, scope); + const needsInstall = + !existsSync(installedPath) || !(await this.installedNpmMatchesConfiguredVersion(parsed, installedPath)); + if (needsInstall) { + const installed = await installMissing(); + if (!installed) continue; + installedPath = this.getNpmInstallPath(parsed, scope); + } + metadata.baseDir = installedPath; + this.collectPackageResources(installedPath, accumulator, filter, metadata); + continue; + } + + if (parsed.type === "git") { + const installedPath = this.getGitInstallPath(parsed, scope); + if (!existsSync(installedPath)) { + const installed = await installMissing(); + if (!installed) continue; + } else if (scope === "temporary" && !parsed.pinned && !isOfflineModeEnabled()) { + await this.refreshTemporaryGitSource(parsed, sourceStr); + } + metadata.baseDir = installedPath; + this.collectPackageResources(installedPath, accumulator, filter, metadata); + } + } + } + + private resolveLocalExtensionSource( + source: LocalSource, + accumulator: ResourceAccumulator, + filter: PackageFilter | undefined, + metadata: PathMetadata, + baseDir: string, + ): void { + const resolved = this.resolvePathFromBase(source.path, baseDir); + if (!existsSync(resolved)) { + return; + } + + try { + const stats = statSync(resolved); + if (stats.isFile()) { + metadata.baseDir = dirname(resolved); + this.addResource(accumulator.extensions, resolved, metadata, true); + return; + } + if (stats.isDirectory()) { + metadata.baseDir = resolved; + const resources = this.collectPackageResources(resolved, accumulator, filter, metadata); + if (!resources) { + this.addResource(accumulator.extensions, resolved, metadata, true); + } + } + } catch { + return; + } + } + + private async installParsedSource(parsed: ParsedSource, scope: SourceScope): Promise { + if (parsed.type === "npm") { + await this.installNpm(parsed, scope, scope === "temporary"); + return; + } + if (parsed.type === "git") { + await this.installGit(parsed, scope); + return; + } + } + + private getPackageSourceString(pkg: PackageSource): string { + return typeof pkg === "string" ? pkg : pkg.source; + } + + private getSourceMatchKeyForInput(source: string): string { + const parsed = this.parseSource(source); + if (parsed.type === "npm") { + return `npm:${parsed.name}`; + } + if (parsed.type === "git") { + return `git:${parsed.host}/${parsed.path}`; + } + return `local:${this.resolvePath(parsed.path)}`; + } + + private getSourceMatchKeyForSettings(source: string, scope: SourceScope): string { + const parsed = this.parseSource(source); + if (parsed.type === "npm") { + return `npm:${parsed.name}`; + } + if (parsed.type === "git") { + return `git:${parsed.host}/${parsed.path}`; + } + const baseDir = this.getBaseDirForScope(scope); + return `local:${this.resolvePathFromBase(parsed.path, baseDir)}`; + } + + private buildNoMatchingPackageMessage(source: string, configuredPackages: PackageSource[]): string { + const suggestion = this.findSuggestedConfiguredSource(source, configuredPackages); + if (!suggestion) { + return `No matching package found for ${source}`; + } + return `No matching package found for ${source}. Did you mean ${suggestion}?`; + } + + private findSuggestedConfiguredSource(source: string, configuredPackages: PackageSource[]): string | undefined { + const trimmedSource = source.trim(); + const suggestions = new Set(); + + for (const pkg of configuredPackages) { + const sourceStr = this.getPackageSourceString(pkg); + const parsed = this.parseSource(sourceStr); + if (parsed.type === "npm") { + if (trimmedSource === parsed.name || trimmedSource === parsed.spec) { + suggestions.add(sourceStr); + } + continue; + } + if (parsed.type === "git") { + const shorthand = `${parsed.host}/${parsed.path}`; + const shorthandWithRef = parsed.ref ? `${shorthand}@${parsed.ref}` : undefined; + if (trimmedSource === shorthand || (shorthandWithRef && trimmedSource === shorthandWithRef)) { + suggestions.add(sourceStr); + } + } + } + + return suggestions.values().next().value; + } + + private packageSourcesMatch(existing: PackageSource, inputSource: string, scope: SourceScope): boolean { + const left = this.getSourceMatchKeyForSettings(this.getPackageSourceString(existing), scope); + const right = this.getSourceMatchKeyForInput(inputSource); + return left === right; + } + + private normalizePackageSourceForSettings(source: string, scope: SourceScope): string { + const parsed = this.parseSource(source); + if (parsed.type !== "local") { + return source; + } + const baseDir = this.getBaseDirForScope(scope); + const resolved = this.resolvePath(parsed.path); + const rel = relative(baseDir, resolved); + return rel || "."; + } + + private parseSource(source: string): ParsedSource { + if (source.startsWith("npm:")) { + const spec = source.slice("npm:".length).trim(); + const { name, version } = this.parseNpmSpec(spec); + return { + type: "npm", + spec, + name, + version, + range: getNpmVersionRange(version), + pinned: isExactNpmVersion(version), + }; + } + + if (isLocalPath(source)) { + return { type: "local", path: source }; + } + + // Try parsing as git URL + const gitParsed = parseGitUrl(source); + if (gitParsed) { + return gitParsed; + } + + return { type: "local", path: source }; + } + + private async installedNpmMatchesConfiguredVersion(source: NpmSource, installedPath: string): Promise { + const installedVersion = this.getInstalledNpmVersion(installedPath); + if (!installedVersion) { + return false; + } + return source.range ? satisfies(installedVersion, source.range) : true; + } + + private async npmHasAvailableUpdate(source: NpmSource, installedPath: string): Promise { + if (isOfflineModeEnabled()) { + return false; + } + + const installedVersion = this.getInstalledNpmVersion(installedPath); + if (!installedVersion) { + return false; + } + + try { + const targetVersion = await this.getLatestNpmVersion(source.version ? source.spec : source.name, source.range); + return targetVersion !== installedVersion; + } catch { + return false; + } + } + + private getInstalledNpmVersion(installedPath: string): string | undefined { + const packageJsonPath = join(installedPath, "package.json"); + if (!existsSync(packageJsonPath)) return undefined; + try { + const content = readFileSync(packageJsonPath, "utf-8"); + const pkg = JSON.parse(content) as { version?: string }; + return pkg.version; + } catch { + return undefined; + } + } + + private async getLatestNpmVersion(packageSpec: string, range?: string): Promise { + const npmCommand = this.getNpmCommand(); + const stdout = await this.runCommandCapture( + npmCommand.command, + [...npmCommand.args, "view", packageSpec, "version", "--json"], + { cwd: this.cwd, timeoutMs: NETWORK_TIMEOUT_MS }, + ); + const raw = stdout.trim(); + if (!raw) throw new Error("Empty response from npm view"); + const parsed = JSON.parse(raw) as unknown; + if (typeof parsed === "string") { + return parsed; + } + if (Array.isArray(parsed)) { + const versions = parsed.filter((value): value is string => typeof value === "string" && value.length > 0); + const latest = range ? maxSatisfying(versions, range) : [...versions].sort(rcompare)[0]; + if (latest) return latest; + } + throw new Error("Unexpected response from npm view"); + } + + private async gitHasAvailableUpdate(installedPath: string): Promise { + if (isOfflineModeEnabled()) { + return false; + } + + try { + const localHead = await this.runCommandCapture("git", ["rev-parse", "HEAD"], { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + const remoteHead = await this.getRemoteGitHead(installedPath); + return localHead.trim() !== remoteHead.trim(); + } catch { + return false; + } + } + + private async getRemoteGitHead(installedPath: string): Promise { + const upstreamRef = await this.getGitUpstreamRef(installedPath); + if (upstreamRef) { + const remoteHead = await this.runGitRemoteCommand(installedPath, ["ls-remote", "origin", upstreamRef]); + const match = remoteHead.match(/^([0-9a-f]{40})\s+/m); + if (match?.[1]) { + return match[1]; + } + } + + const remoteHead = await this.runGitRemoteCommand(installedPath, ["ls-remote", "origin", "HEAD"]); + const match = remoteHead.match(/^([0-9a-f]{40})\s+HEAD$/m); + if (!match?.[1]) { + throw new Error("Failed to determine remote HEAD"); + } + return match[1]; + } + + private async getLocalGitUpdateTarget( + installedPath: string, + ): Promise<{ ref: string; head: string; fetchArgs: string[] }> { + try { + const upstream = await this.runCommandCapture("git", ["rev-parse", "--abbrev-ref", "@{upstream}"], { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + const trimmedUpstream = upstream.trim(); + if (!trimmedUpstream.startsWith("origin/")) { + throw new Error(`Unsupported upstream remote: ${trimmedUpstream}`); + } + const branch = trimmedUpstream.slice("origin/".length); + if (!branch) { + throw new Error("Missing upstream branch name"); + } + const head = await this.runCommandCapture("git", ["rev-parse", "@{upstream}"], { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + return { + ref: "@{upstream}", + head, + fetchArgs: [ + "fetch", + "--prune", + "--no-tags", + "origin", + `+refs/heads/${branch}:refs/remotes/origin/${branch}`, + ], + }; + } catch { + await this.runCommand("git", ["remote", "set-head", "origin", "-a"], { cwd: installedPath }).catch(() => {}); + const head = await this.runCommandCapture("git", ["rev-parse", "origin/HEAD"], { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + const originHeadRef = await this.runCommandCapture("git", ["symbolic-ref", "refs/remotes/origin/HEAD"], { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + }).catch(() => ""); + const branch = originHeadRef.trim().replace(/^refs\/remotes\/origin\//, ""); + if (branch) { + return { + ref: "origin/HEAD", + head, + fetchArgs: [ + "fetch", + "--prune", + "--no-tags", + "origin", + `+refs/heads/${branch}:refs/remotes/origin/${branch}`, + ], + }; + } + return { + ref: "origin/HEAD", + head, + fetchArgs: ["fetch", "--prune", "--no-tags", "origin", "+HEAD:refs/remotes/origin/HEAD"], + }; + } + } + + private async getGitUpstreamRef(installedPath: string): Promise { + try { + const upstream = await this.runCommandCapture("git", ["rev-parse", "--abbrev-ref", "@{upstream}"], { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + const trimmed = upstream.trim(); + if (!trimmed.startsWith("origin/")) { + return undefined; + } + const branch = trimmed.slice("origin/".length); + return branch ? `refs/heads/${branch}` : undefined; + } catch { + return undefined; + } + } + + private runGitRemoteCommand(installedPath: string, args: string[]): Promise { + return this.runCommandCapture("git", args, { + cwd: installedPath, + timeoutMs: NETWORK_TIMEOUT_MS, + env: { + GIT_TERMINAL_PROMPT: "0", + }, + }); + } + + private async runWithConcurrency(tasks: Array<() => Promise>, limit: number): Promise { + if (tasks.length === 0) { + return []; + } + + const results: T[] = new Array(tasks.length); + let nextIndex = 0; + const workerCount = Math.max(1, Math.min(limit, tasks.length)); + + const worker = async () => { + while (true) { + const index = nextIndex; + nextIndex += 1; + if (index >= tasks.length) { + return; + } + results[index] = await tasks[index](); + } + }; + + await Promise.all(Array.from({ length: workerCount }, () => worker())); + return results; + } + + /** + * Get a unique identity for a package, ignoring version/ref. + * Used to detect when the same package is in both global and project settings. + * For git packages, uses normalized host/path to ensure SSH and HTTPS URLs + * for the same repository are treated as identical. + */ + private getPackageIdentity(source: string, scope?: SourceScope): string { + const parsed = this.parseSource(source); + if (parsed.type === "npm") { + return `npm:${parsed.name}`; + } + if (parsed.type === "git") { + // Use host/path for identity to normalize SSH and HTTPS + return `git:${parsed.host}/${parsed.path}`; + } + if (scope) { + const baseDir = this.getBaseDirForScope(scope); + return `local:${this.resolvePathFromBase(parsed.path, baseDir)}`; + } + return `local:${this.resolvePath(parsed.path)}`; + } + + /** + * Dedupe packages: if same package identity appears in both global and project, + * keep only the project one (project wins). + */ + private dedupePackages( + packages: Array<{ pkg: PackageSource; scope: SourceScope }>, + ): Array<{ pkg: PackageSource; scope: SourceScope }> { + const seen = new Map(); + + for (const entry of packages) { + const sourceStr = typeof entry.pkg === "string" ? entry.pkg : entry.pkg.source; + const identity = this.getPackageIdentity(sourceStr, entry.scope); + + const existing = seen.get(identity); + if (!existing) { + seen.set(identity, entry); + } else if (entry.scope === "project" && existing.scope === "user") { + // Project wins over user + seen.set(identity, entry); + } + // If existing is project and new is global, keep existing (project) + // If both are same scope, keep first one + } + + return Array.from(seen.values()); + } + + private parseNpmSpec(spec: string): { name: string; version?: string } { + const match = spec.match(/^(@?[^@]+(?:\/[^@]+)?)(?:@(.+))?$/); + if (!match) { + return { name: spec }; + } + const name = match[1] ?? spec; + const version = match[2]; + return { name, version }; + } + + private assertProjectTrustedForScope(scope: SourceScope): void { + if (scope === "project" && !this.settingsManager.isProjectTrusted()) { + throw new Error("Project is not trusted; refusing to access project package storage"); + } + } + + private getNpmCommand(): { command: string; args: string[] } { + const configuredCommand = this.settingsManager.getNpmCommand(); + if (!configuredCommand || configuredCommand.length === 0) { + return { command: "npm", args: [] }; + } + const [command, ...args] = configuredCommand; + if (!command) { + throw new Error("Invalid npmCommand: first array entry must be a non-empty command"); + } + return { command, args }; + } + + private getPackageManagerName(): string { + const npmCommand = this.getNpmCommand(); + const commandParts = [npmCommand.command, ...npmCommand.args]; + const separatorIndex = commandParts.lastIndexOf("--"); + const packageManagerCommand = separatorIndex >= 0 ? commandParts[separatorIndex + 1] : npmCommand.command; + return packageManagerCommand ? basename(packageManagerCommand).replace(/\.(cmd|exe)$/i, "") : ""; + } + + private async runNpmCommand(args: string[], options?: { cwd?: string }): Promise { + const npmCommand = this.getNpmCommand(); + await this.runCommand(npmCommand.command, [...npmCommand.args, ...args], options); + } + + private getGitDependencyInstallArgs(): string[] { + const configuredCommand = this.settingsManager.getNpmCommand(); + if (configuredCommand && configuredCommand.length > 0) { + return ["install"]; + } + return ["install", "--omit=dev"]; + } + + private runNpmCommandSync(args: string[]): string { + const npmCommand = this.getNpmCommand(); + return this.runCommandSync(npmCommand.command, [...npmCommand.args, ...args]); + } + + private getNpmInstallArgs(specs: string[], installRoot: string): string[] { + const packageManagerName = this.getPackageManagerName(); + // Extension packages run inside pi and resolve pi APIs through loader aliases/virtual modules. + // Disable peer dependency resolution for managed installs (npm's --legacy-peer-deps, and + // equivalent bun/pnpm settings) so package managers do not install or solve host-provided + // @earendil-works/pi-* peers. Stale auto-installed pi peers can otherwise block updates. + if (packageManagerName === "bun") { + return ["install", ...specs, "--cwd", installRoot, "--omit=peer"]; + } + if (packageManagerName === "pnpm") { + return [ + "install", + ...specs, + "--prefix", + installRoot, + "--config.auto-install-peers=false", + "--config.strict-peer-dependencies=false", + "--config.strict-dep-builds=false", + ]; + } + return ["install", ...specs, "--prefix", installRoot, "--legacy-peer-deps"]; + } + + private async installNpm(_source: NpmSource, _scope: SourceScope, _temporary: boolean): Promise { + throw new Error("Package installation is disabled in this Cactus build."); + } + + private async uninstallNpm(source: NpmSource, scope: SourceScope): Promise { + const installRoot = this.getNpmInstallRoot(scope, false); + if (!existsSync(installRoot)) { + return; + } + if (this.getPackageManagerName() === "bun") { + await this.runNpmCommand(["uninstall", source.name, "--cwd", installRoot]); + return; + } + await this.runNpmCommand(["uninstall", source.name, "--prefix", installRoot]); + } + + private async installGit(_source: GitSource, _scope: SourceScope): Promise { + throw new Error("Package installation is disabled in this Cactus build."); + } + + private async updateGit(source: GitSource, scope: SourceScope): Promise { + const targetDir = this.getGitInstallPath(source, scope); + if (!existsSync(targetDir)) { + await this.installGit(source, scope); + return; + } + + if (source.ref) { + await this.ensureGitRef(targetDir, ["fetch", "origin", source.ref], "FETCH_HEAD"); + return; + } + + const target = await this.getLocalGitUpdateTarget(targetDir); + await this.ensureGitRef(targetDir, target.fetchArgs, target.ref); + } + + private async ensureGitRef(targetDir: string, fetchArgs: string[], ref: string): Promise { + // Fetch only the ref we will reset to, avoiding unrelated branch/tag noise. + await this.runCommand("git", fetchArgs, { cwd: targetDir }); + + const localHead = await this.runCommandCapture("git", ["rev-parse", "HEAD"], { + cwd: targetDir, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + const commitRef = `${ref}^{commit}`; + const targetHead = await this.runCommandCapture("git", ["rev-parse", commitRef], { + cwd: targetDir, + timeoutMs: NETWORK_TIMEOUT_MS, + }); + if (localHead.trim() === targetHead.trim()) { + return; + } + + await this.runCommand("git", ["reset", "--hard", commitRef], { cwd: targetDir }); + + // Clean untracked files (extensions should be pristine) + await this.runCommand("git", ["clean", "-fdx"], { cwd: targetDir }); + + const packageJsonPath = join(targetDir, "package.json"); + if (existsSync(packageJsonPath)) { + await this.runNpmCommand(this.getGitDependencyInstallArgs(), { cwd: targetDir }); + } + } + + private async refreshTemporaryGitSource(source: GitSource, sourceStr: string): Promise { + if (isOfflineModeEnabled()) { + return; + } + try { + await this.withProgress("pull", sourceStr, `Refreshing ${sourceStr}...`, async () => { + await this.updateGit(source, "temporary"); + }); + } catch { + // Keep cached temporary checkout if refresh fails. + } + } + + private async removeGit(source: GitSource, scope: SourceScope): Promise { + const targetDir = this.getGitInstallPath(source, scope); + if (!existsSync(targetDir)) return; + rmSync(targetDir, { recursive: true, force: true }); + this.pruneEmptyGitParents(targetDir, this.getGitInstallRoot(scope)); + } + + private pruneEmptyGitParents(targetDir: string, installRoot: string | undefined): void { + if (!installRoot) return; + const resolvedRoot = resolve(installRoot); + let current = dirname(targetDir); + while (current.startsWith(resolvedRoot) && current !== resolvedRoot) { + if (!existsSync(current)) { + current = dirname(current); + continue; + } + const entries = readdirSync(current); + if (entries.length > 0) { + break; + } + try { + rmSync(current, { recursive: true, force: true }); + } catch { + break; + } + current = dirname(current); + } + } + + private ensureNpmProject(installRoot: string): void { + if (!existsSync(installRoot)) { + mkdirSync(installRoot, { recursive: true }); + } + markPathIgnoredByCloudSync(installRoot); + this.ensureGitIgnore(installRoot); + const packageJsonPath = join(installRoot, "package.json"); + if (!existsSync(packageJsonPath)) { + const pkgJson = { name: "pi-extensions", private: true }; + writeFileSync(packageJsonPath, JSON.stringify(pkgJson, null, 2), "utf-8"); + } + } + + private ensureGitIgnore(dir: string): void { + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + const ignorePath = join(dir, ".gitignore"); + if (!existsSync(ignorePath)) { + writeFileSync(ignorePath, "*\n!.gitignore\n", "utf-8"); + } + } + + private getNpmInstallRoot(scope: SourceScope, temporary: boolean): string { + if (temporary) { + return this.getTemporaryDir("npm"); + } + if (scope === "project") { + this.assertProjectTrustedForScope(scope); + return join(this.cwd, CONFIG_DIR_NAME, "npm"); + } + return join(this.agentDir, "npm"); + } + + private getGlobalNpmRoot(): string { + const npmCommand = this.getNpmCommand(); + const commandKey = [npmCommand.command, ...npmCommand.args].join("\0"); + if (this.globalNpmRoot && this.globalNpmRootCommandKey === commandKey) { + return this.globalNpmRoot; + } + if (this.getPackageManagerName() === "bun") { + const binDir = this.runNpmCommandSync(["pm", "bin", "-g"]).trim(); + this.globalNpmRoot = join(dirname(binDir), "install", "global", "node_modules"); + } else { + this.globalNpmRoot = this.runNpmCommandSync(["root", "-g"]).trim(); + } + this.globalNpmRootCommandKey = commandKey; + return this.globalNpmRoot; + } + + private getPnpmGlobalPackagePath(packageName: string): string | undefined { + if (this.getPackageManagerName() !== "pnpm") { + return undefined; + } + + const output = this.runNpmCommandSync(["list", "-g", "--depth", "0", "--json"]); + const entries = JSON.parse(output) as Array<{ dependencies?: Record }>; + for (const entry of entries) { + const path = entry.dependencies?.[packageName]?.path; + if (path) return path; + } + return undefined; + } + + private getManagedNpmInstallPath(source: NpmSource, scope: SourceScope): string { + if (scope === "temporary") { + return join(this.getTemporaryDir("npm"), "node_modules", source.name); + } + if (scope === "project") { + this.assertProjectTrustedForScope(scope); + return join(this.cwd, CONFIG_DIR_NAME, "npm", "node_modules", source.name); + } + return join(this.agentDir, "npm", "node_modules", source.name); + } + + private getLegacyGlobalNpmInstallPath(source: NpmSource): string | undefined { + try { + return this.getPnpmGlobalPackagePath(source.name) ?? join(this.getGlobalNpmRoot(), source.name); + } catch { + return undefined; + } + } + + private getNpmInstallPath(source: NpmSource, scope: SourceScope): string { + const managedPath = this.getManagedNpmInstallPath(source, scope); + if (scope !== "user" || existsSync(managedPath)) { + return managedPath; + } + const legacyPath = this.getLegacyGlobalNpmInstallPath(source); + return legacyPath && existsSync(legacyPath) ? legacyPath : managedPath; + } + + private getGitInstallPath(source: GitSource, scope: SourceScope): string { + if (scope === "temporary") { + return this.getTemporaryDir(`git-${source.host}`, source.path); + } + const installRoot = this.getGitInstallRoot(scope); + if (!installRoot) { + throw new Error("Missing git install root"); + } + return this.resolveManagedPath(installRoot, source.host, source.path); + } + + private getGitInstallRoot(scope: SourceScope): string | undefined { + if (scope === "temporary") { + return undefined; + } + if (scope === "project") { + this.assertProjectTrustedForScope(scope); + return join(this.cwd, CONFIG_DIR_NAME, "git"); + } + return join(this.agentDir, "git"); + } + + private getTemporaryDir(prefix: string, suffix?: string): string { + const root = this.resolveManagedPath(getExtensionTempFolder(this.agentDir), prefix); + const hash = createHash("sha256") + .update(`${prefix}-${suffix ?? ""}`) + .digest("hex") + .slice(0, 8); + return this.resolveManagedPath(root, hash, suffix ?? ""); + } + + private resolveManagedPath(root: string, ...parts: string[]): string { + const resolvedRoot = resolve(root); + const resolvedPath = resolve(resolvedRoot, ...parts); + if (resolvedPath !== resolvedRoot && !resolvedPath.startsWith(`${resolvedRoot}${sep}`)) { + throw new Error(`Refusing to use path outside package install root: ${resolvedPath}`); + } + return resolvedPath; + } + + private getBaseDirForScope(scope: SourceScope): string { + if (scope === "project") { + this.assertProjectTrustedForScope(scope); + return join(this.cwd, CONFIG_DIR_NAME); + } + if (scope === "user") { + return this.agentDir; + } + return this.cwd; + } + + private resolvePath(input: string): string { + return resolvePath(input, this.cwd, { homeDir: getHomeDir(), trim: true }); + } + + private resolvePathFromBase(input: string, baseDir: string): string { + return resolvePath(input, baseDir, { homeDir: getHomeDir(), trim: true }); + } + + private collectPackageResources( + packageRoot: string, + accumulator: ResourceAccumulator, + filter: PackageFilter | undefined, + metadata: PathMetadata, + ): boolean { + if (filter) { + for (const resourceType of RESOURCE_TYPES) { + const patterns = filter[resourceType as keyof PackageFilter]; + const target = this.getTargetMap(accumulator, resourceType); + if (patterns !== undefined) { + this.applyPackageFilter(packageRoot, patterns, resourceType, target, metadata); + } else { + this.collectDefaultResources(packageRoot, resourceType, target, metadata); + } + } + return true; + } + + const manifest = this.readPiManifest(packageRoot); + if (manifest) { + for (const resourceType of RESOURCE_TYPES) { + const entries = manifest[resourceType as keyof PiManifest]; + this.addManifestEntries( + entries, + packageRoot, + resourceType, + this.getTargetMap(accumulator, resourceType), + metadata, + ); + } + return true; + } + + let hasAnyDir = false; + for (const resourceType of RESOURCE_TYPES) { + const dir = join(packageRoot, resourceType); + if (existsSync(dir)) { + // Collect all files from the directory (all enabled by default) + const files = collectResourceFiles(dir, resourceType); + for (const f of files) { + this.addResource(this.getTargetMap(accumulator, resourceType), f, metadata, true); + } + hasAnyDir = true; + } + } + return hasAnyDir; + } + + private collectDefaultResources( + packageRoot: string, + resourceType: ResourceType, + target: Map, + metadata: PathMetadata, + ): void { + const manifest = this.readPiManifest(packageRoot); + const entries = manifest?.[resourceType as keyof PiManifest]; + if (entries) { + this.addManifestEntries(entries, packageRoot, resourceType, target, metadata); + return; + } + const dir = join(packageRoot, resourceType); + if (existsSync(dir)) { + // Collect all files from the directory (all enabled by default) + const files = collectResourceFiles(dir, resourceType); + for (const f of files) { + this.addResource(target, f, metadata, true); + } + } + } + + private applyPackageFilter( + packageRoot: string, + userPatterns: string[], + resourceType: ResourceType, + target: Map, + metadata: PathMetadata, + ): void { + const { allFiles } = this.collectManifestFiles(packageRoot, resourceType); + + if (userPatterns.length === 0) { + // Empty array explicitly disables all resources of this type + for (const f of allFiles) { + this.addResource(target, f, metadata, false); + } + return; + } + + // Apply user patterns + const enabledByUser = applyPatterns(allFiles, userPatterns, packageRoot); + + for (const f of allFiles) { + const enabled = enabledByUser.has(f); + this.addResource(target, f, metadata, enabled); + } + } + + /** + * Collect all files from a package for a resource type, applying manifest patterns. + * Returns { allFiles, enabledByManifest } where enabledByManifest is the set of files + * that pass the manifest's own patterns. + */ + private collectManifestFiles( + packageRoot: string, + resourceType: ResourceType, + ): { allFiles: string[]; enabledByManifest: Set } { + const manifest = this.readPiManifest(packageRoot); + const entries = manifest?.[resourceType as keyof PiManifest]; + if (entries && entries.length > 0) { + const allFiles = this.collectFilesFromManifestEntries(entries, packageRoot, resourceType); + const manifestPatterns = entries.filter(isOverridePattern); + const enabledByManifest = + manifestPatterns.length > 0 ? applyPatterns(allFiles, manifestPatterns, packageRoot) : new Set(allFiles); + return { allFiles: Array.from(enabledByManifest), enabledByManifest }; + } + + const conventionDir = join(packageRoot, resourceType); + if (!existsSync(conventionDir)) { + return { allFiles: [], enabledByManifest: new Set() }; + } + const allFiles = collectResourceFiles(conventionDir, resourceType); + return { allFiles, enabledByManifest: new Set(allFiles) }; + } + + private readPiManifest(packageRoot: string): PiManifest | null { + const packageJsonPath = join(packageRoot, "package.json"); + if (!existsSync(packageJsonPath)) { + return null; + } + + try { + const content = readFileSync(packageJsonPath, "utf-8"); + const pkg = JSON.parse(content) as { pi?: PiManifest }; + return pkg.pi ?? null; + } catch { + return null; + } + } + + private addManifestEntries( + entries: string[] | undefined, + root: string, + resourceType: ResourceType, + target: Map, + metadata: PathMetadata, + ): void { + if (!entries) return; + + const allFiles = this.collectFilesFromManifestEntries(entries, root, resourceType); + const patterns = entries.filter(isOverridePattern); + const enabledPaths = applyPatterns(allFiles, patterns, root); + + for (const f of allFiles) { + if (enabledPaths.has(f)) { + this.addResource(target, f, metadata, true); + } + } + } + + private collectFilesFromManifestEntries(entries: string[], root: string, resourceType: ResourceType): string[] { + const sourceEntries = entries.filter((entry) => !isOverridePattern(entry)); + const resolved = sourceEntries.flatMap((entry) => { + if (!hasGlobPattern(entry)) { + return [resolve(root, entry)]; + } + + return globSync(entry, { + cwd: root, + absolute: true, + dot: false, + nodir: false, + }).map((match) => resolve(match)); + }); + return this.collectFilesFromPaths(resolved, resourceType); + } + + private resolveLocalEntries( + entries: string[], + resourceType: ResourceType, + target: Map, + metadata: PathMetadata, + baseDir: string, + ): void { + if (entries.length === 0) return; + + // Collect all files from plain entries (non-pattern entries) + const { plain, patterns } = splitPatterns(entries); + const resolvedPlain = plain.map((p) => this.resolvePathFromBase(p, baseDir)); + const allFiles = this.collectFilesFromPaths(resolvedPlain, resourceType); + + // Determine which files are enabled based on patterns + const enabledPaths = applyPatterns(allFiles, patterns, baseDir); + + // Add all files with their enabled state + for (const f of allFiles) { + this.addResource(target, f, metadata, enabledPaths.has(f)); + } + } + + private addAutoDiscoveredResources( + accumulator: ResourceAccumulator, + globalSettings: ReturnType, + projectSettings: ReturnType, + globalBaseDir: string, + projectBaseDir: string, + ): void { + const userMetadata: PathMetadata = { + source: "auto", + scope: "user", + origin: "top-level", + baseDir: globalBaseDir, + }; + const projectMetadata: PathMetadata = { + source: "auto", + scope: "project", + origin: "top-level", + baseDir: projectBaseDir, + }; + + const userOverrides = { + extensions: (globalSettings.extensions ?? []) as string[], + skills: (globalSettings.skills ?? []) as string[], + prompts: (globalSettings.prompts ?? []) as string[], + themes: (globalSettings.themes ?? []) as string[], + }; + const projectOverrides = { + extensions: (projectSettings.extensions ?? []) as string[], + skills: (projectSettings.skills ?? []) as string[], + prompts: (projectSettings.prompts ?? []) as string[], + themes: (projectSettings.themes ?? []) as string[], + }; + + const userDirs = { + extensions: join(globalBaseDir, "extensions"), + skills: join(globalBaseDir, "skills"), + prompts: join(globalBaseDir, "prompts"), + themes: join(globalBaseDir, "themes"), + }; + const projectDirs = { + extensions: join(projectBaseDir, "extensions"), + skills: join(projectBaseDir, "skills"), + prompts: join(projectBaseDir, "prompts"), + themes: join(projectBaseDir, "themes"), + }; + const userAgentsSkillsDir = join(getHomeDir(), ".agents", "skills"); + const projectTrusted = this.settingsManager.isProjectTrusted(); + const projectAgentsSkillDirs = projectTrusted + ? collectAncestorAgentsSkillDirs(this.cwd).filter((dir) => resolve(dir) !== resolve(userAgentsSkillsDir)) + : []; + + const addResources = ( + resourceType: ResourceType, + paths: string[], + metadata: PathMetadata, + overrides: string[], + baseDir: string, + ) => { + const target = this.getTargetMap(accumulator, resourceType); + for (const path of paths) { + const enabled = isEnabledByOverrides(path, overrides, baseDir); + this.addResource(target, path, metadata, enabled); + } + }; + + if (projectTrusted) { + // Project extensions from .pi/ + addResources( + "extensions", + collectAutoExtensionEntries(projectDirs.extensions), + projectMetadata, + projectOverrides.extensions, + projectBaseDir, + ); + + // Project skills from .pi/ + addResources( + "skills", + collectAutoSkillEntries(projectDirs.skills, "pi"), + projectMetadata, + projectOverrides.skills, + projectBaseDir, + ); + } + + // Project skills from .agents/ (each with its own baseDir) + for (const agentsSkillsDir of projectAgentsSkillDirs) { + const agentsBaseDir = dirname(agentsSkillsDir); // the .agents directory + const agentsMetadata: PathMetadata = { + ...projectMetadata, + baseDir: agentsBaseDir, + }; + addResources( + "skills", + collectAutoSkillEntries(agentsSkillsDir, "agents"), + agentsMetadata, + projectOverrides.skills, + agentsBaseDir, + ); + } + + if (projectTrusted) { + addResources( + "prompts", + collectAutoPromptEntries(projectDirs.prompts), + projectMetadata, + projectOverrides.prompts, + projectBaseDir, + ); + addResources( + "themes", + collectAutoThemeEntries(projectDirs.themes), + projectMetadata, + projectOverrides.themes, + projectBaseDir, + ); + } + + // User extensions from ~/.pi/agent/ + addResources( + "extensions", + collectAutoExtensionEntries(userDirs.extensions), + userMetadata, + userOverrides.extensions, + globalBaseDir, + ); + + // User skills from ~/.pi/agent/ + addResources( + "skills", + collectAutoSkillEntries(userDirs.skills, "pi"), + userMetadata, + userOverrides.skills, + globalBaseDir, + ); + + // User skills from ~/.agents/ (with its own baseDir) + const userAgentsBaseDir = dirname(userAgentsSkillsDir); + const userAgentsMetadata: PathMetadata = { + ...userMetadata, + baseDir: userAgentsBaseDir, + }; + addResources( + "skills", + collectAutoSkillEntries(userAgentsSkillsDir, "agents"), + userAgentsMetadata, + userOverrides.skills, + userAgentsBaseDir, + ); + + addResources( + "prompts", + collectAutoPromptEntries(userDirs.prompts), + userMetadata, + userOverrides.prompts, + globalBaseDir, + ); + addResources( + "themes", + collectAutoThemeEntries(userDirs.themes), + userMetadata, + userOverrides.themes, + globalBaseDir, + ); + } + + private collectFilesFromPaths(paths: string[], resourceType: ResourceType): string[] { + const files: string[] = []; + for (const p of paths) { + if (!existsSync(p)) continue; + + try { + const stats = statSync(p); + if (stats.isFile()) { + files.push(p); + } else if (stats.isDirectory()) { + files.push(...collectResourceFiles(p, resourceType)); + } + } catch { + // Ignore errors + } + } + return files; + } + + private getTargetMap( + accumulator: ResourceAccumulator, + resourceType: ResourceType, + ): Map { + switch (resourceType) { + case "extensions": + return accumulator.extensions; + case "skills": + return accumulator.skills; + case "prompts": + return accumulator.prompts; + case "themes": + return accumulator.themes; + default: + throw new Error(`Unknown resource type: ${resourceType}`); + } + } + + private addResource( + map: Map, + path: string, + metadata: PathMetadata, + enabled: boolean, + ): void { + if (!path) return; + if (!map.has(path)) { + map.set(path, { metadata, enabled }); + } + } + + private createAccumulator(): ResourceAccumulator { + return { + extensions: new Map(), + skills: new Map(), + prompts: new Map(), + themes: new Map(), + }; + } + + private toResolvedPaths(accumulator: ResourceAccumulator): ResolvedPaths { + const mapToResolved = ( + entries: Map, + ): ResolvedResource[] => { + const resolved = Array.from(entries.entries()).map(([path, { metadata, enabled }]) => ({ + path, + enabled, + metadata, + })); + resolved.sort((a, b) => resourcePrecedenceRank(a.metadata) - resourcePrecedenceRank(b.metadata)); + + const seen = new Set(); + return resolved.filter((entry) => { + const canonicalPath = canonicalizePath(entry.path); + if (seen.has(canonicalPath)) return false; + seen.add(canonicalPath); + return true; + }); + }; + + return { + extensions: mapToResolved(accumulator.extensions), + skills: mapToResolved(accumulator.skills), + prompts: mapToResolved(accumulator.prompts), + themes: mapToResolved(accumulator.themes), + }; + } + + private spawnCommand(command: string, args: string[], options?: { cwd?: string }): ChildProcess { + const env = getEnv(); + return spawnProcess(command, args, { + cwd: options?.cwd, + stdio: isStdoutTakenOver() ? ["ignore", 2, 2] : "inherit", + env, + }); + } + + private spawnCaptureCommand( + command: string, + args: string[], + options?: { cwd?: string; env?: Record }, + ): ChildProcessByStdio { + const baseEnv = getEnv(); + const env = options?.env ? { ...baseEnv, ...options.env } : baseEnv; + return spawnProcess(command, args, { + cwd: options?.cwd, + stdio: ["ignore", "pipe", "pipe"], + env, + }); + } + + private runCommandCapture( + command: string, + args: string[], + options?: { cwd?: string; timeoutMs?: number; env?: Record }, + ): Promise { + return new Promise((resolvePromise, reject) => { + const child = this.spawnCaptureCommand(command, args, options); + let stdout = ""; + let stderr = ""; + let timedOut = false; + const timeout = + typeof options?.timeoutMs === "number" + ? setTimeout(() => { + timedOut = true; + child.kill(); + }, options.timeoutMs) + : undefined; + + child.stdout?.on("data", (data) => { + stdout += data.toString(); + }); + child.stderr?.on("data", (data) => { + stderr += data.toString(); + }); + child.once("error", (error) => { + if (timeout) clearTimeout(timeout); + reject(error); + }); + child.once("close", (code, signal) => { + if (timeout) clearTimeout(timeout); + if (timedOut) { + reject(new Error(`${command} ${args.join(" ")} timed out after ${options?.timeoutMs}ms`)); + return; + } + if (code === 0) { + resolvePromise(stdout.trim()); + return; + } + const exitStatus = code === null ? `signal ${signal ?? "unknown"}` : `code ${code}`; + reject(new Error(`${command} ${args.join(" ")} failed with ${exitStatus}: ${stderr || stdout}`)); + }); + }); + } + + private runCommand(command: string, args: string[], options?: { cwd?: string }): Promise { + return new Promise((resolvePromise, reject) => { + const child = this.spawnCommand(command, args, options); + child.on("error", reject); + child.on("exit", (code) => { + if (code === 0) { + resolvePromise(); + } else { + reject(new Error(`${command} ${args.join(" ")} failed with code ${code}`)); + } + }); + }); + } + + private runCommandSync(command: string, args: string[]): string { + const env = getEnv(); + const result = spawnProcessSync(command, args, { + stdio: ["ignore", "pipe", "pipe"], + encoding: "utf-8", + env, + }); + if (result.error || result.status !== 0) { + throw new Error( + `Failed to run ${command} ${args.join(" ")}: ${result.error?.message || result.stderr || result.stdout}`, + ); + } + return (result.stdout || result.stderr || "").trim(); + } +} diff --git a/cactus-code/packages/coding-agent/src/core/project-trust.ts b/cactus-code/packages/coding-agent/src/core/project-trust.ts new file mode 100644 index 000000000..4b0039774 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/project-trust.ts @@ -0,0 +1,96 @@ +import { CONFIG_DIR_NAME } from "../config.ts"; +import { emitProjectTrustEvent } from "./extensions/runner.ts"; +import type { LoadExtensionsResult, ProjectTrustContext } from "./extensions/types.ts"; +import type { DefaultProjectTrust } from "./settings-manager.ts"; +import { + getProjectTrustOptions, + hasTrustRequiringProjectResources, + type ProjectTrustOption, + type ProjectTrustStore, +} from "./trust-manager.ts"; + +export type AppMode = "interactive" | "print" | "json"; + +export interface ResolveProjectTrustedOptions { + cwd: string; + trustStore: ProjectTrustStore; + trustOverride?: boolean; + defaultProjectTrust?: DefaultProjectTrust; + extensionsResult?: LoadExtensionsResult; + projectTrustContext: ProjectTrustContext; + onExtensionError?: (message: string) => void; +} + +function formatProjectTrustPrompt(cwd: string): string { + return `Trust project folder?\n${cwd}\n\nThis allows Cactus to load ${CONFIG_DIR_NAME} settings and resources (skills, prompts, themes) from this project.`; +} + +async function selectProjectTrustOption( + cwd: string, + ctx: ProjectTrustContext, +): Promise { + const options = getProjectTrustOptions(cwd, { includeSessionOnly: true }); + const selected = await ctx.ui.select( + formatProjectTrustPrompt(cwd), + options.map((option) => option.label), + ); + return options.find((option) => option.label === selected); +} + +function saveProjectTrustPromptResult(trustStore: ProjectTrustStore, result: ProjectTrustOption): void { + if (result.updates.length > 0) { + trustStore.setMany(result.updates); + } +} + +export async function resolveProjectTrusted(options: ResolveProjectTrustedOptions): Promise { + if (options.trustOverride !== undefined) { + return options.trustOverride; + } + if (!hasTrustRequiringProjectResources(options.cwd)) { + return true; + } + + if (options.extensionsResult) { + const { result, errors } = await emitProjectTrustEvent( + options.extensionsResult, + { type: "project_trust", cwd: options.cwd }, + options.projectTrustContext, + ); + for (const error of errors) { + options.onExtensionError?.(`Extension "${error.extensionPath}" project_trust error: ${error.error}`); + } + if (result) { + const trusted = result.trusted === "yes"; + if (result.remember === true) { + options.trustStore.set(options.cwd, trusted); + } + return trusted; + } + } + + const decision = options.trustStore.get(options.cwd); + if (decision !== null) { + return decision; + } + + switch (options.defaultProjectTrust ?? "ask") { + case "always": + return true; + case "never": + return false; + case "ask": + break; + } + + if (!options.projectTrustContext.hasUI) { + return false; + } + + const selected = await selectProjectTrustOption(options.cwd, options.projectTrustContext); + if (selected !== undefined) { + saveProjectTrustPromptResult(options.trustStore, selected); + return selected.trusted; + } + return false; +} diff --git a/cactus-code/packages/coding-agent/src/core/prompt-templates.ts b/cactus-code/packages/coding-agent/src/core/prompt-templates.ts new file mode 100644 index 000000000..6b5b1e24a --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/prompt-templates.ts @@ -0,0 +1,284 @@ +import { existsSync, readdirSync, readFileSync, statSync } from "fs"; +import { basename, dirname, join, resolve, sep } from "path"; +import { CONFIG_DIR_NAME } from "../config.ts"; +import { parseFrontmatter } from "../utils/frontmatter.ts"; +import { resolvePath } from "../utils/paths.ts"; +import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.ts"; + +/** + * Represents a prompt template loaded from a markdown file + */ +export interface PromptTemplate { + name: string; + description: string; + argumentHint?: string; + content: string; + sourceInfo: SourceInfo; + filePath: string; // Absolute path to the template file +} + +/** + * Parse command arguments respecting quoted strings (bash-style) + * Returns array of arguments + */ +export function parseCommandArgs(argsString: string): string[] { + const args: string[] = []; + let current = ""; + let inQuote: string | null = null; + + for (let i = 0; i < argsString.length; i++) { + const char = argsString[i]; + + if (inQuote) { + if (char === inQuote) { + inQuote = null; + } else { + current += char; + } + } else if (char === '"' || char === "'") { + inQuote = char; + } else if (/\s/.test(char)) { + if (current) { + args.push(current); + current = ""; + } + } else { + current += char; + } + } + + if (current) { + args.push(current); + } + + return args; +} + +/** + * Substitute argument placeholders in template content + * Supports: + * - $1, $2, ... for positional args + * - $@ and $ARGUMENTS for all args + * - ${N:-default} for positional arg N with default when missing/empty + * - ${@:N} for args from Nth onwards (bash-style slicing) + * - ${@:N:L} for L args starting from Nth + * + * Note: Replacement happens on the template string only. Argument and default values + * containing patterns like $1, $@, or $ARGUMENTS are NOT recursively substituted. + */ +export function substituteArgs(content: string, args: string[]): string { + const allArgs = args.join(" "); + + return content.replace( + /\$\{(\d+):-([^}]*)\}|\$\{@:(\d+)(?::(\d+))?\}|\$(ARGUMENTS|@|\d+)/g, + (_match, defaultNum, defaultValue, sliceStart, sliceLength, simple) => { + if (defaultNum) { + const index = parseInt(defaultNum, 10) - 1; + const value = args[index]; + return value ? value : defaultValue; + } + + if (sliceStart) { + let start = parseInt(sliceStart, 10) - 1; // Convert to 0-indexed (user provides 1-indexed) + // Treat 0 as 1 (bash convention: args start at 1) + if (start < 0) start = 0; + + if (sliceLength) { + const length = parseInt(sliceLength, 10); + return args.slice(start, start + length).join(" "); + } + return args.slice(start).join(" "); + } + + if (simple === "ARGUMENTS" || simple === "@") { + return allArgs; + } + + const index = parseInt(simple, 10) - 1; + return args[index] ?? ""; + }, + ); +} + +function loadTemplateFromFile(filePath: string, sourceInfo: SourceInfo): PromptTemplate | null { + try { + const rawContent = readFileSync(filePath, "utf-8"); + const { frontmatter, body } = parseFrontmatter>(rawContent); + + const name = basename(filePath).replace(/\.md$/, ""); + + // Get description from frontmatter or first non-empty line + let description = frontmatter.description || ""; + if (!description) { + const firstLine = body.split("\n").find((line) => line.trim()); + if (firstLine) { + // Truncate if too long + description = firstLine.slice(0, 60); + if (firstLine.length > 60) description += "..."; + } + } + + return { + name, + description, + ...(frontmatter["argument-hint"] && { argumentHint: frontmatter["argument-hint"] }), + content: body, + sourceInfo, + filePath, + }; + } catch { + return null; + } +} + +/** + * Scan a directory for .md files (non-recursive) and load them as prompt templates. + */ +function loadTemplatesFromDir(dir: string, getSourceInfo: (filePath: string) => SourceInfo): PromptTemplate[] { + const templates: PromptTemplate[] = []; + + if (!existsSync(dir)) { + return templates; + } + + try { + const entries = readdirSync(dir, { withFileTypes: true }); + + for (const entry of entries) { + const fullPath = join(dir, entry.name); + + // For symlinks, check if they point to a file + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + const stats = statSync(fullPath); + isFile = stats.isFile(); + } catch { + // Broken symlink, skip it + continue; + } + } + + if (isFile && entry.name.endsWith(".md")) { + const template = loadTemplateFromFile(fullPath, getSourceInfo(fullPath)); + if (template) { + templates.push(template); + } + } + } + } catch { + return templates; + } + + return templates; +} + +export interface LoadPromptTemplatesOptions { + /** Working directory for project-local templates. */ + cwd: string; + /** Agent config directory for global templates. */ + agentDir: string; + /** Explicit prompt template paths (files or directories). */ + promptPaths: string[]; + /** Include default prompt directories. */ + includeDefaults: boolean; +} + +/** + * Load all prompt templates from: + * 1. Global: agentDir/prompts/ + * 2. Project: cwd/{CONFIG_DIR_NAME}/prompts/ + * 3. Explicit prompt paths + */ +export function loadPromptTemplates(options: LoadPromptTemplatesOptions): PromptTemplate[] { + const resolvedCwd = resolvePath(options.cwd); + const resolvedAgentDir = resolvePath(options.agentDir); + const promptPaths = options.promptPaths; + const includeDefaults = options.includeDefaults; + + const templates: PromptTemplate[] = []; + + const globalPromptsDir = join(resolvedAgentDir, "prompts"); + const projectPromptsDir = resolve(resolvedCwd, CONFIG_DIR_NAME, "prompts"); + + const isUnderPath = (target: string, root: string): boolean => { + const normalizedRoot = resolve(root); + if (target === normalizedRoot) { + return true; + } + const prefix = normalizedRoot.endsWith(sep) ? normalizedRoot : `${normalizedRoot}${sep}`; + return target.startsWith(prefix); + }; + + const getSourceInfo = (resolvedPath: string): SourceInfo => { + if (isUnderPath(resolvedPath, globalPromptsDir)) { + return createSyntheticSourceInfo(resolvedPath, { + source: "local", + scope: "user", + baseDir: globalPromptsDir, + }); + } + if (isUnderPath(resolvedPath, projectPromptsDir)) { + return createSyntheticSourceInfo(resolvedPath, { + source: "local", + scope: "project", + baseDir: projectPromptsDir, + }); + } + return createSyntheticSourceInfo(resolvedPath, { + source: "local", + baseDir: statSync(resolvedPath).isDirectory() ? resolvedPath : dirname(resolvedPath), + }); + }; + + if (includeDefaults) { + templates.push(...loadTemplatesFromDir(globalPromptsDir, getSourceInfo)); + templates.push(...loadTemplatesFromDir(projectPromptsDir, getSourceInfo)); + } + + // 3. Load explicit prompt paths + for (const rawPath of promptPaths) { + const resolvedPath = resolvePath(rawPath, resolvedCwd, { trim: true }); + if (!existsSync(resolvedPath)) { + continue; + } + + try { + const stats = statSync(resolvedPath); + if (stats.isDirectory()) { + templates.push(...loadTemplatesFromDir(resolvedPath, getSourceInfo)); + } else if (stats.isFile() && resolvedPath.endsWith(".md")) { + const template = loadTemplateFromFile(resolvedPath, getSourceInfo(resolvedPath)); + if (template) { + templates.push(template); + } + } + } catch { + // Ignore read failures + } + } + + return templates; +} + +/** + * Expand a prompt template if it matches a template name. + * Returns the expanded content or the original text if not a template. + */ +export function expandPromptTemplate(text: string, templates: PromptTemplate[]): string { + if (!text.startsWith("/")) return text; + + const match = text.match(/^\/([^\s]+)(?:\s+([\s\S]*))?$/); + if (!match) return text; + + const templateName = match[1]; + const argsString = match[2] ?? ""; + + const template = templates.find((t) => t.name === templateName); + if (template) { + const args = parseCommandArgs(argsString); + return substituteArgs(template.content, args); + } + + return text; +} diff --git a/cactus-code/packages/coding-agent/src/core/provider-attribution.ts b/cactus-code/packages/coding-agent/src/core/provider-attribution.ts new file mode 100644 index 000000000..b26f32927 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/provider-attribution.ts @@ -0,0 +1,17 @@ +import type { Api, Model, ProviderHeaders } from "@earendil-works/pi-ai"; +import type { SettingsManager } from "./settings-manager.ts"; + +export function mergeProviderAttributionHeaders( + _model: Model, + _settingsManager: SettingsManager, + _sessionId: string | undefined, + ...headerSources: Array +): ProviderHeaders | undefined { + const merged: ProviderHeaders = {}; + for (const headers of headerSources) { + if (headers) { + Object.assign(merged, headers); + } + } + return Object.keys(merged).length > 0 ? merged : undefined; +} diff --git a/cactus-code/packages/coding-agent/src/core/provider-display-names.ts b/cactus-code/packages/coding-agent/src/core/provider-display-names.ts new file mode 100644 index 000000000..d33c3d7d6 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/provider-display-names.ts @@ -0,0 +1,35 @@ +export const BUILT_IN_PROVIDER_DISPLAY_NAMES: Record = { + anthropic: "Anthropic", + "amazon-bedrock": "Amazon Bedrock", + "ant-ling": "Ant Ling", + "azure-openai-responses": "Azure OpenAI Responses", + cerebras: "Cerebras", + "cloudflare-ai-gateway": "Cloudflare AI Gateway", + "cloudflare-workers-ai": "Cloudflare Workers AI", + deepseek: "DeepSeek", + fireworks: "Fireworks", + google: "Google Gemini", + "google-vertex": "Google Vertex AI", + groq: "Groq", + huggingface: "Hugging Face", + "kimi-coding": "Kimi For Coding", + mistral: "Mistral", + minimax: "MiniMax", + "minimax-cn": "MiniMax (China)", + moonshotai: "Moonshot AI", + "moonshotai-cn": "Moonshot AI (China)", + nvidia: "NVIDIA NIM", + opencode: "OpenCode Zen", + "opencode-go": "OpenCode Go", + openai: "OpenAI", + openrouter: "OpenRouter", + together: "Together AI", + "vercel-ai-gateway": "Vercel AI Gateway", + xai: "xAI", + zai: "ZAI Coding Plan (Global)", + "zai-coding-cn": "ZAI Coding Plan (China)", + xiaomi: "Xiaomi MiMo", + "xiaomi-token-plan-cn": "Xiaomi MiMo Token Plan (China)", + "xiaomi-token-plan-ams": "Xiaomi MiMo Token Plan (Amsterdam)", + "xiaomi-token-plan-sgp": "Xiaomi MiMo Token Plan (Singapore)", +}; diff --git a/cactus-code/packages/coding-agent/src/core/resolve-config-value.ts b/cactus-code/packages/coding-agent/src/core/resolve-config-value.ts new file mode 100644 index 000000000..6d75b001e --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/resolve-config-value.ts @@ -0,0 +1,287 @@ +/** + * Resolve configuration values that may be shell commands, environment variables, or literals. + * Used by auth-storage.ts and model-registry.ts. + */ + +import { execSync, spawnSync } from "child_process"; +import { getShellConfig } from "../utils/shell.ts"; + +// Cache for shell command results (persists for process lifetime) +const commandResultCache = new Map(); +const ENV_VAR_NAME_RE = /^[A-Za-z_][A-Za-z0-9_]*$/; +const ENV_VAR_NAME_PREFIX_RE = /^[A-Za-z_][A-Za-z0-9_]*/; + +type TemplatePart = { type: "literal"; value: string } | { type: "env"; name: string }; + +type ConfigValueReference = { type: "command"; config: string } | { type: "template"; parts: TemplatePart[] }; + +function appendLiteral(parts: TemplatePart[], value: string): void { + if (!value) return; + const previousPart = parts[parts.length - 1]; + if (previousPart?.type === "literal") { + previousPart.value += value; + return; + } + parts.push({ type: "literal", value }); +} + +function parseConfigValueTemplate(config: string): TemplatePart[] { + const parts: TemplatePart[] = []; + let index = 0; + + while (index < config.length) { + const dollarIndex = config.indexOf("$", index); + if (dollarIndex < 0) { + appendLiteral(parts, config.slice(index)); + break; + } + + appendLiteral(parts, config.slice(index, dollarIndex)); + const nextChar = config[dollarIndex + 1]; + + if (nextChar === "$" || nextChar === "!") { + appendLiteral(parts, nextChar); + index = dollarIndex + 2; + continue; + } + + if (nextChar === "{") { + const endIndex = config.indexOf("}", dollarIndex + 2); + if (endIndex < 0) { + appendLiteral(parts, "$"); + index = dollarIndex + 1; + continue; + } + + const name = config.slice(dollarIndex + 2, endIndex); + if (ENV_VAR_NAME_RE.test(name)) { + parts.push({ type: "env", name }); + } else { + appendLiteral(parts, config.slice(dollarIndex, endIndex + 1)); + } + index = endIndex + 1; + continue; + } + + const match = config.slice(dollarIndex + 1).match(ENV_VAR_NAME_PREFIX_RE); + if (match) { + parts.push({ type: "env", name: match[0] }); + index = dollarIndex + 1 + match[0].length; + continue; + } + + appendLiteral(parts, "$"); + index = dollarIndex + 1; + } + + return parts; +} + +function parseConfigValueReference(config: string): ConfigValueReference { + if (config.startsWith("!")) { + return { type: "command", config }; + } + + return { type: "template", parts: parseConfigValueTemplate(config) }; +} + +function resolveEnvConfigValue(name: string, env?: Record): string | undefined { + return env?.[name] || process.env[name] || undefined; +} + +function getTemplateEnvVarNames(parts: TemplatePart[]): string[] { + const names: string[] = []; + for (const part of parts) { + if (part.type !== "env" || names.includes(part.name)) continue; + names.push(part.name); + } + return names; +} + +function resolveTemplate(parts: TemplatePart[], env?: Record): string | undefined { + let resolved = ""; + for (const part of parts) { + if (part.type === "literal") { + resolved += part.value; + continue; + } + const envValue = resolveEnvConfigValue(part.name, env); + if (envValue === undefined) return undefined; + resolved += envValue; + } + return resolved; +} + +export function getConfigValueEnvVarName(config: string): string | undefined { + const reference = parseConfigValueReference(config); + if (reference.type !== "template") return undefined; + return reference.parts.length === 1 && reference.parts[0]?.type === "env" ? reference.parts[0].name : undefined; +} + +export function getConfigValueEnvVarNames(config: string): string[] { + const reference = parseConfigValueReference(config); + return reference.type === "template" ? getTemplateEnvVarNames(reference.parts) : []; +} + +export function getMissingConfigValueEnvVarNames(config: string, env?: Record): string[] { + return getConfigValueEnvVarNames(config).filter((name) => resolveEnvConfigValue(name, env) === undefined); +} + +export function isCommandConfigValue(config: string): boolean { + return parseConfigValueReference(config).type === "command"; +} + +export function isConfigValueConfigured(config: string, env?: Record): boolean { + return getMissingConfigValueEnvVarNames(config, env).length === 0; +} + +/** + * Resolve a config value (API key, header value, etc.) to an actual value. + * - If starts with "!", executes the rest as a shell command and uses stdout (cached) + * - Interpolates "$ENV_VAR" or "${ENV_VAR}" references with the named environment variable + * - In non-command values, "$$" escapes a literal "$" and "$!" escapes a literal "!" + * - Otherwise treats the value as a literal + */ +export function resolveConfigValue(config: string, env?: Record): string | undefined { + const reference = parseConfigValueReference(config); + if (reference.type === "command") { + return executeCommand(reference.config); + } + return resolveTemplate(reference.parts, env); +} + +function executeWithConfiguredShell(command: string): { executed: boolean; value: string | undefined } { + try { + const { shell, args, commandTransport } = getShellConfig(); + const commandFromStdin = commandTransport === "stdin"; + const result = spawnSync(shell, commandFromStdin ? args : [...args, command], { + encoding: "utf-8", + input: commandFromStdin ? command : undefined, + timeout: 10000, + stdio: [commandFromStdin ? "pipe" : "ignore", "pipe", "ignore"], + shell: false, + windowsHide: true, + }); + + if (result.error) { + const error = result.error as NodeJS.ErrnoException; + if (error.code === "ENOENT") { + return { executed: false, value: undefined }; + } + return { executed: true, value: undefined }; + } + + if (result.status !== 0) { + return { executed: true, value: undefined }; + } + + const value = (result.stdout ?? "").trim(); + return { executed: true, value: value || undefined }; + } catch { + return { executed: false, value: undefined }; + } +} + +function executeWithDefaultShell(command: string): string | undefined { + try { + const output = execSync(command, { + encoding: "utf-8", + timeout: 10000, + stdio: ["ignore", "pipe", "ignore"], + }); + return output.trim() || undefined; + } catch { + return undefined; + } +} + +function executeCommandUncached(commandConfig: string): string | undefined { + const command = commandConfig.slice(1); + return process.platform === "win32" + ? (() => { + const configuredResult = executeWithConfiguredShell(command); + return configuredResult.executed ? configuredResult.value : executeWithDefaultShell(command); + })() + : executeWithDefaultShell(command); +} + +function executeCommand(commandConfig: string): string | undefined { + if (commandResultCache.has(commandConfig)) { + return commandResultCache.get(commandConfig); + } + + const result = executeCommandUncached(commandConfig); + commandResultCache.set(commandConfig, result); + return result; +} + +/** + * Resolve all header values using the same resolution logic as API keys. + */ +export function resolveConfigValueUncached(config: string, env?: Record): string | undefined { + const reference = parseConfigValueReference(config); + if (reference.type === "command") { + return executeCommandUncached(reference.config); + } + return resolveTemplate(reference.parts, env); +} + +export function resolveConfigValueOrThrow(config: string, description: string, env?: Record): string { + const resolvedValue = resolveConfigValueUncached(config, env); + if (resolvedValue !== undefined) { + return resolvedValue; + } + + const reference = parseConfigValueReference(config); + if (reference.type === "command") { + throw new Error(`Failed to resolve ${description} from shell command: ${reference.config.slice(1)}`); + } + + if (reference.type === "template") { + const missingEnvVars = getMissingConfigValueEnvVarNames(config, env); + if (missingEnvVars.length === 1) { + throw new Error(`Failed to resolve ${description} from environment variable: ${missingEnvVars[0]}`); + } + if (missingEnvVars.length > 1) { + throw new Error(`Failed to resolve ${description} from environment variables: ${missingEnvVars.join(", ")}`); + } + } + + throw new Error(`Failed to resolve ${description}`); +} + +/** + * Resolve all header values using the same resolution logic as API keys. + */ +export function resolveHeaders( + headers: Record | undefined, + env?: Record, +): Record | undefined { + if (!headers) return undefined; + const resolved: Record = {}; + for (const [key, value] of Object.entries(headers)) { + const resolvedValue = resolveConfigValue(value, env); + if (resolvedValue) { + resolved[key] = resolvedValue; + } + } + return Object.keys(resolved).length > 0 ? resolved : undefined; +} + +export function resolveHeadersOrThrow( + headers: Record | undefined, + description: string, + env?: Record, +): Record | undefined { + if (!headers) return undefined; + const resolved: Record = {}; + for (const [key, value] of Object.entries(headers)) { + resolved[key] = resolveConfigValueOrThrow(value, `${description} header "${key}"`, env); + } + return Object.keys(resolved).length > 0 ? resolved : undefined; +} + +/** Clear the config value command cache. Exported for testing. */ +export function clearConfigValueCache(): void { + commandResultCache.clear(); +} diff --git a/cactus-code/packages/coding-agent/src/core/resource-loader.ts b/cactus-code/packages/coding-agent/src/core/resource-loader.ts new file mode 100644 index 000000000..18486ead9 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/resource-loader.ts @@ -0,0 +1,1037 @@ +import { existsSync, readdirSync, readFileSync, statSync } from "node:fs"; +import { join, resolve, sep } from "node:path"; +import chalk from "chalk"; +import { CONFIG_DIR_NAME } from "../config.ts"; +import { loadThemeFromPath, type Theme } from "../modes/interactive/theme/theme.ts"; +import type { ResourceDiagnostic } from "./diagnostics.ts"; + +export type { ResourceCollision, ResourceDiagnostic } from "./diagnostics.ts"; + +import { canonicalizePath, isLocalPath, resolvePath } from "../utils/paths.ts"; +import { createEventBus, type EventBus } from "./event-bus.ts"; +import { + clearExtensionCache, + createExtensionRuntime, + loadExtensionFromFactory, + loadExtensionsCached, +} from "./extensions/loader.ts"; +import type { Extension, ExtensionFactory, ExtensionRuntime, LoadExtensionsResult } from "./extensions/types.ts"; +import { DefaultPackageManager, type PathMetadata, type ResolvedResource } from "./package-manager.ts"; +import type { PromptTemplate } from "./prompt-templates.ts"; +import { loadPromptTemplates } from "./prompt-templates.ts"; +import { SettingsManager } from "./settings-manager.ts"; +import type { Skill } from "./skills.ts"; +import { loadSkills } from "./skills.ts"; +import { createSourceInfo, type SourceInfo } from "./source-info.ts"; + +export interface ResourceExtensionPaths { + skillPaths?: Array<{ path: string; metadata: PathMetadata }>; + promptPaths?: Array<{ path: string; metadata: PathMetadata }>; + themePaths?: Array<{ path: string; metadata: PathMetadata }>; +} + +export interface ResourceLoaderReloadOptions { + resolveProjectTrust?: (input: { extensionsResult: LoadExtensionsResult }) => Promise; +} + +export interface ResourceLoader { + getExtensions(): LoadExtensionsResult; + getSkills(): { skills: Skill[]; diagnostics: ResourceDiagnostic[] }; + getPrompts(): { prompts: PromptTemplate[]; diagnostics: ResourceDiagnostic[] }; + getThemes(): { themes: Theme[]; diagnostics: ResourceDiagnostic[] }; + getAgentsFiles(): { agentsFiles: Array<{ path: string; content: string }> }; + getSystemPrompt(): string | undefined; + getAppendSystemPrompt(): string[]; + extendResources(paths: ResourceExtensionPaths): void; + reload(options?: ResourceLoaderReloadOptions): Promise; +} + +function resolvePromptInput(input: string | undefined, description: string): string | undefined { + if (!input) { + return undefined; + } + + if (existsSync(input)) { + try { + return readFileSync(input, "utf-8"); + } catch (error) { + console.error(chalk.yellow(`Warning: Could not read ${description} file ${input}: ${error}`)); + return input; + } + } + + return input; +} + +function loadContextFileFromDir(dir: string): { path: string; content: string } | null { + const candidates = ["AGENTS.md", "AGENTS.MD", "CLAUDE.md", "CLAUDE.MD"]; + for (const filename of candidates) { + const filePath = join(dir, filename); + if (existsSync(filePath)) { + try { + return { + path: filePath, + content: readFileSync(filePath, "utf-8"), + }; + } catch (error) { + console.error(chalk.yellow(`Warning: Could not read ${filePath}: ${error}`)); + } + } + } + return null; +} + +export function loadProjectContextFiles(options: { + cwd: string; + agentDir: string; +}): Array<{ path: string; content: string }> { + const resolvedCwd = resolvePath(options.cwd); + const resolvedAgentDir = resolvePath(options.agentDir); + + const contextFiles: Array<{ path: string; content: string }> = []; + const seenPaths = new Set(); + + const globalContext = loadContextFileFromDir(resolvedAgentDir); + if (globalContext) { + contextFiles.push(globalContext); + seenPaths.add(globalContext.path); + } + + const ancestorContextFiles: Array<{ path: string; content: string }> = []; + + let currentDir = resolvedCwd; + const root = resolve("/"); + + while (true) { + const contextFile = loadContextFileFromDir(currentDir); + if (contextFile && !seenPaths.has(contextFile.path)) { + ancestorContextFiles.unshift(contextFile); + seenPaths.add(contextFile.path); + } + + if (currentDir === root) break; + + const parentDir = resolve(currentDir, ".."); + if (parentDir === currentDir) break; + currentDir = parentDir; + } + + contextFiles.push(...ancestorContextFiles); + + return contextFiles; +} + +export interface DefaultResourceLoaderOptions { + cwd: string; + agentDir: string; + settingsManager?: SettingsManager; + eventBus?: EventBus; + additionalExtensionPaths?: string[]; + additionalSkillPaths?: string[]; + additionalPromptTemplatePaths?: string[]; + additionalThemePaths?: string[]; + extensionFactories?: ExtensionFactory[]; + noExtensions?: boolean; + noSkills?: boolean; + noPromptTemplates?: boolean; + noThemes?: boolean; + noContextFiles?: boolean; + systemPrompt?: string; + appendSystemPrompt?: string[]; + extensionsOverride?: (base: LoadExtensionsResult) => LoadExtensionsResult; + skillsOverride?: (base: { skills: Skill[]; diagnostics: ResourceDiagnostic[] }) => { + skills: Skill[]; + diagnostics: ResourceDiagnostic[]; + }; + promptsOverride?: (base: { prompts: PromptTemplate[]; diagnostics: ResourceDiagnostic[] }) => { + prompts: PromptTemplate[]; + diagnostics: ResourceDiagnostic[]; + }; + themesOverride?: (base: { themes: Theme[]; diagnostics: ResourceDiagnostic[] }) => { + themes: Theme[]; + diagnostics: ResourceDiagnostic[]; + }; + agentsFilesOverride?: (base: { agentsFiles: Array<{ path: string; content: string }> }) => { + agentsFiles: Array<{ path: string; content: string }>; + }; + systemPromptOverride?: (base: string | undefined) => string | undefined; + appendSystemPromptOverride?: (base: string[]) => string[]; +} + +export class DefaultResourceLoader implements ResourceLoader { + private cwd: string; + private agentDir: string; + private settingsManager: SettingsManager; + private eventBus: EventBus; + private packageManager: DefaultPackageManager; + private additionalExtensionPaths: string[]; + private additionalSkillPaths: string[]; + private additionalPromptTemplatePaths: string[]; + private additionalThemePaths: string[]; + private extensionFactories: ExtensionFactory[]; + private noExtensions: boolean; + private noSkills: boolean; + private noPromptTemplates: boolean; + private noThemes: boolean; + private noContextFiles: boolean; + private systemPromptSource?: string; + private appendSystemPromptSource?: string[]; + private extensionsOverride?: (base: LoadExtensionsResult) => LoadExtensionsResult; + private skillsOverride?: (base: { skills: Skill[]; diagnostics: ResourceDiagnostic[] }) => { + skills: Skill[]; + diagnostics: ResourceDiagnostic[]; + }; + private promptsOverride?: (base: { prompts: PromptTemplate[]; diagnostics: ResourceDiagnostic[] }) => { + prompts: PromptTemplate[]; + diagnostics: ResourceDiagnostic[]; + }; + private themesOverride?: (base: { themes: Theme[]; diagnostics: ResourceDiagnostic[] }) => { + themes: Theme[]; + diagnostics: ResourceDiagnostic[]; + }; + private agentsFilesOverride?: (base: { agentsFiles: Array<{ path: string; content: string }> }) => { + agentsFiles: Array<{ path: string; content: string }>; + }; + private systemPromptOverride?: (base: string | undefined) => string | undefined; + private appendSystemPromptOverride?: (base: string[]) => string[]; + + private extensionsResult: LoadExtensionsResult; + private skills: Skill[]; + private skillDiagnostics: ResourceDiagnostic[]; + private prompts: PromptTemplate[]; + private promptDiagnostics: ResourceDiagnostic[]; + private themes: Theme[]; + private themeDiagnostics: ResourceDiagnostic[]; + private agentsFiles: Array<{ path: string; content: string }>; + private systemPrompt?: string; + private appendSystemPrompt: string[]; + private lastSkillPaths: string[]; + private extensionSkillSourceInfos: Map; + private extensionPromptSourceInfos: Map; + private extensionThemeSourceInfos: Map; + private lastPromptPaths: string[]; + private lastThemePaths: string[]; + private loaded: boolean; + + constructor(options: DefaultResourceLoaderOptions) { + this.cwd = resolvePath(options.cwd); + this.agentDir = resolvePath(options.agentDir); + this.settingsManager = options.settingsManager ?? SettingsManager.create(this.cwd, this.agentDir); + this.eventBus = options.eventBus ?? createEventBus(); + this.packageManager = new DefaultPackageManager({ + cwd: this.cwd, + agentDir: this.agentDir, + settingsManager: this.settingsManager, + }); + this.additionalExtensionPaths = options.additionalExtensionPaths ?? []; + this.additionalSkillPaths = options.additionalSkillPaths ?? []; + this.additionalPromptTemplatePaths = options.additionalPromptTemplatePaths ?? []; + this.additionalThemePaths = options.additionalThemePaths ?? []; + this.extensionFactories = options.extensionFactories ?? []; + this.noExtensions = options.noExtensions ?? false; + this.noSkills = options.noSkills ?? false; + this.noPromptTemplates = options.noPromptTemplates ?? false; + this.noThemes = options.noThemes ?? false; + this.noContextFiles = options.noContextFiles ?? false; + this.systemPromptSource = options.systemPrompt; + this.appendSystemPromptSource = options.appendSystemPrompt; + this.extensionsOverride = options.extensionsOverride; + this.skillsOverride = options.skillsOverride; + this.promptsOverride = options.promptsOverride; + this.themesOverride = options.themesOverride; + this.agentsFilesOverride = options.agentsFilesOverride; + this.systemPromptOverride = options.systemPromptOverride; + this.appendSystemPromptOverride = options.appendSystemPromptOverride; + + this.extensionsResult = { extensions: [], errors: [], runtime: createExtensionRuntime() }; + this.skills = []; + this.skillDiagnostics = []; + this.prompts = []; + this.promptDiagnostics = []; + this.themes = []; + this.themeDiagnostics = []; + this.agentsFiles = []; + this.appendSystemPrompt = []; + this.lastSkillPaths = []; + this.extensionSkillSourceInfos = new Map(); + this.extensionPromptSourceInfos = new Map(); + this.extensionThemeSourceInfos = new Map(); + this.lastPromptPaths = []; + this.lastThemePaths = []; + this.loaded = false; + } + + getExtensions(): LoadExtensionsResult { + return this.extensionsResult; + } + + getSkills(): { skills: Skill[]; diagnostics: ResourceDiagnostic[] } { + return { skills: this.skills, diagnostics: this.skillDiagnostics }; + } + + getPrompts(): { prompts: PromptTemplate[]; diagnostics: ResourceDiagnostic[] } { + return { prompts: this.prompts, diagnostics: this.promptDiagnostics }; + } + + getThemes(): { themes: Theme[]; diagnostics: ResourceDiagnostic[] } { + return { themes: this.themes, diagnostics: this.themeDiagnostics }; + } + + getAgentsFiles(): { agentsFiles: Array<{ path: string; content: string }> } { + return { agentsFiles: this.agentsFiles }; + } + + getSystemPrompt(): string | undefined { + return this.systemPrompt; + } + + getAppendSystemPrompt(): string[] { + return this.appendSystemPrompt; + } + + extendResources(paths: ResourceExtensionPaths): void { + const skillPaths = this.normalizeExtensionPaths(paths.skillPaths ?? []); + const promptPaths = this.normalizeExtensionPaths(paths.promptPaths ?? []); + const themePaths = this.normalizeExtensionPaths(paths.themePaths ?? []); + + for (const entry of skillPaths) { + this.extensionSkillSourceInfos.set(entry.path, createSourceInfo(entry.path, entry.metadata)); + } + for (const entry of promptPaths) { + this.extensionPromptSourceInfos.set(entry.path, createSourceInfo(entry.path, entry.metadata)); + } + for (const entry of themePaths) { + this.extensionThemeSourceInfos.set(entry.path, createSourceInfo(entry.path, entry.metadata)); + } + + if (skillPaths.length > 0) { + this.lastSkillPaths = this.mergePaths( + this.lastSkillPaths, + skillPaths.map((entry) => entry.path), + ); + this.updateSkillsFromPaths(this.lastSkillPaths); + } + + if (promptPaths.length > 0) { + this.lastPromptPaths = this.mergePaths( + this.lastPromptPaths, + promptPaths.map((entry) => entry.path), + ); + this.updatePromptsFromPaths(this.lastPromptPaths); + } + + if (themePaths.length > 0) { + this.lastThemePaths = this.mergePaths( + this.lastThemePaths, + themePaths.map((entry) => entry.path), + ); + this.updateThemesFromPaths(this.lastThemePaths); + } + } + + async loadProjectTrustExtensions(): Promise { + // Force untrusted project settings for the bootstrap pass. This keeps project-local + // extensions/packages out while still loading user/global and temporary CLI extensions. + this.settingsManager.setProjectTrusted(false); + await this.settingsManager.reload(); + return this.loadCurrentExtensionSet({ includeInlineFactories: true }); + } + + async reload(options?: ResourceLoaderReloadOptions): Promise { + if (this.loaded) { + clearExtensionCache(); + } + + let preTrustExtensions: LoadExtensionsResult | undefined; + if (options?.resolveProjectTrust) { + preTrustExtensions = await this.loadProjectTrustExtensions(); + const projectTrusted = await options.resolveProjectTrust({ extensionsResult: preTrustExtensions }); + this.settingsManager.setProjectTrusted(projectTrusted); + } + + // reload() preserves SettingsManager.projectTrusted and reloads settings for that trust state. + await this.settingsManager.reload(); + const resolvedPaths = await this.packageManager.resolve(); + const cliExtensionPaths = await this.packageManager.resolveExtensionSources(this.additionalExtensionPaths, { + temporary: true, + }); + const metadataByPath = new Map(); + + this.extensionSkillSourceInfos = new Map(); + this.extensionPromptSourceInfos = new Map(); + this.extensionThemeSourceInfos = new Map(); + + // Helper to extract enabled paths and store metadata + const getEnabledResources = (resources: ResolvedResource[]): ResolvedResource[] => { + for (const r of resources) { + if (!metadataByPath.has(r.path)) { + metadataByPath.set(r.path, r.metadata); + } + } + return resources.filter((r) => r.enabled); + }; + + const getEnabledPaths = (resources: ResolvedResource[]): string[] => + getEnabledResources(resources).map((r) => r.path); + const enabledExtensions = getEnabledPaths(resolvedPaths.extensions); + const enabledSkillResources = getEnabledResources(resolvedPaths.skills); + const enabledPrompts = getEnabledPaths(resolvedPaths.prompts); + const enabledThemes = getEnabledPaths(resolvedPaths.themes); + + const enabledSkills = enabledSkillResources.map((resource) => this.mapSkillPath(resource, metadataByPath)); + + // Add CLI paths metadata + for (const r of cliExtensionPaths.extensions) { + if (!metadataByPath.has(r.path)) { + metadataByPath.set(r.path, { source: "cli", scope: "temporary", origin: "top-level" }); + } + } + for (const r of cliExtensionPaths.skills) { + if (!metadataByPath.has(r.path)) { + metadataByPath.set(r.path, { source: "cli", scope: "temporary", origin: "top-level" }); + } + } + + const cliEnabledExtensions = getEnabledPaths(cliExtensionPaths.extensions); + const cliEnabledSkills = getEnabledPaths(cliExtensionPaths.skills); + const cliEnabledPrompts = getEnabledPaths(cliExtensionPaths.prompts); + const cliEnabledThemes = getEnabledPaths(cliExtensionPaths.themes); + + const extensionPaths = this.noExtensions + ? cliEnabledExtensions + : this.mergePaths(cliEnabledExtensions, enabledExtensions); + + const extensionsResult = await this.loadFinalExtensionSet(extensionPaths, preTrustExtensions); + for (const p of this.additionalExtensionPaths) { + if (isLocalPath(p)) { + const resolved = this.resolveResourcePath(p); + if (!existsSync(resolved)) { + extensionsResult.errors.push({ path: resolved, error: `Extension path does not exist: ${resolved}` }); + } + } + } + this.extensionsResult = this.extensionsOverride ? this.extensionsOverride(extensionsResult) : extensionsResult; + this.applyExtensionSourceInfo(this.extensionsResult.extensions, metadataByPath); + + const skillPaths = this.noSkills + ? this.mergePaths(cliEnabledSkills, this.additionalSkillPaths) + : this.mergePaths([...cliEnabledSkills, ...enabledSkills], this.additionalSkillPaths); + + this.lastSkillPaths = skillPaths; + this.updateSkillsFromPaths(skillPaths, metadataByPath); + for (const p of this.additionalSkillPaths) { + if (isLocalPath(p)) { + const resolved = this.resolveResourcePath(p); + if (!existsSync(resolved) && !this.skillDiagnostics.some((d) => d.path === resolved)) { + this.skillDiagnostics.push({ type: "error", message: "Skill path does not exist", path: resolved }); + } + } + } + + const promptPaths = this.noPromptTemplates + ? this.mergePaths(cliEnabledPrompts, this.additionalPromptTemplatePaths) + : this.mergePaths([...cliEnabledPrompts, ...enabledPrompts], this.additionalPromptTemplatePaths); + + this.lastPromptPaths = promptPaths; + this.updatePromptsFromPaths(promptPaths, metadataByPath); + for (const p of this.additionalPromptTemplatePaths) { + if (isLocalPath(p)) { + const resolved = this.resolveResourcePath(p); + if (!existsSync(resolved) && !this.promptDiagnostics.some((d) => d.path === resolved)) { + this.promptDiagnostics.push({ + type: "error", + message: "Prompt template path does not exist", + path: resolved, + }); + } + } + } + + const themePaths = this.noThemes + ? this.mergePaths(cliEnabledThemes, this.additionalThemePaths) + : this.mergePaths([...cliEnabledThemes, ...enabledThemes], this.additionalThemePaths); + + this.lastThemePaths = themePaths; + this.updateThemesFromPaths(themePaths, metadataByPath); + for (const p of this.additionalThemePaths) { + const resolved = this.resolveResourcePath(p); + if (!existsSync(resolved) && !this.themeDiagnostics.some((d) => d.path === resolved)) { + this.themeDiagnostics.push({ type: "error", message: "Theme path does not exist", path: resolved }); + } + } + + const agentsFiles = { + agentsFiles: this.noContextFiles + ? [] + : loadProjectContextFiles({ + cwd: this.cwd, + agentDir: this.agentDir, + }), + }; + const resolvedAgentsFiles = this.agentsFilesOverride ? this.agentsFilesOverride(agentsFiles) : agentsFiles; + this.agentsFiles = resolvedAgentsFiles.agentsFiles; + + const baseSystemPrompt = resolvePromptInput( + this.systemPromptSource ?? this.discoverSystemPromptFile(), + "system prompt", + ); + this.systemPrompt = this.systemPromptOverride ? this.systemPromptOverride(baseSystemPrompt) : baseSystemPrompt; + + const appendSources = + this.appendSystemPromptSource ?? + (this.discoverAppendSystemPromptFile() ? [this.discoverAppendSystemPromptFile()!] : []); + const baseAppend = appendSources + .map((s) => resolvePromptInput(s, "append system prompt")) + .filter((s): s is string => s !== undefined); + this.appendSystemPrompt = this.appendSystemPromptOverride + ? this.appendSystemPromptOverride(baseAppend) + : baseAppend; + this.loaded = true; + } + + private async loadCurrentExtensionSet(options: { includeInlineFactories: boolean }): Promise { + const resolvedPaths = await this.packageManager.resolve(); + const cliExtensionPaths = await this.packageManager.resolveExtensionSources(this.additionalExtensionPaths, { + temporary: true, + }); + const enabledExtensions = resolvedPaths.extensions.filter((r) => r.enabled).map((r) => r.path); + const cliEnabledExtensions = cliExtensionPaths.extensions.filter((r) => r.enabled).map((r) => r.path); + const extensionPaths = this.noExtensions + ? cliEnabledExtensions + : this.mergePaths(cliEnabledExtensions, enabledExtensions); + const extensionsResult = await loadExtensionsCached(extensionPaths, this.cwd, this.eventBus); + if (!options.includeInlineFactories) { + return extensionsResult; + } + + const inlineExtensions = await this.loadExtensionFactories(extensionsResult.runtime); + extensionsResult.extensions.push(...inlineExtensions.extensions); + extensionsResult.errors.push(...inlineExtensions.errors); + return extensionsResult; + } + + private resolveExtensionLoadPath(path: string): string { + return resolvePath(path, this.cwd, { normalizeUnicodeSpaces: true }); + } + + private async loadFinalExtensionSet( + extensionPaths: string[], + preTrustExtensions: LoadExtensionsResult | undefined, + ): Promise { + if (!preTrustExtensions) { + const extensionsResult = await loadExtensionsCached(extensionPaths, this.cwd, this.eventBus); + const inlineExtensions = await this.loadExtensionFactories(extensionsResult.runtime); + extensionsResult.extensions.push(...inlineExtensions.extensions); + extensionsResult.errors.push(...inlineExtensions.errors); + this.addExtensionConflictDiagnostics(extensionsResult); + return extensionsResult; + } + + const preloadedByPath = new Map( + preTrustExtensions.extensions + .filter((extension) => !extension.path.startsWith(" [extension.resolvedPath, extension]), + ); + const failedPreloadPaths = new Set( + preTrustExtensions.errors.map((error) => this.resolveExtensionLoadPath(error.path)), + ); + const remainingPaths = extensionPaths.filter((path) => { + const resolvedPath = this.resolveExtensionLoadPath(path); + return !preloadedByPath.has(resolvedPath) && !failedPreloadPaths.has(resolvedPath); + }); + const remainingExtensions = await loadExtensionsCached( + remainingPaths, + this.cwd, + this.eventBus, + preTrustExtensions.runtime, + ); + const loadedByPath = new Map(preloadedByPath); + for (const extension of remainingExtensions.extensions) { + loadedByPath.set(extension.resolvedPath, extension); + } + + const inlineExtensions = preTrustExtensions.extensions.filter((extension) => + extension.path.startsWith(" loadedByPath.get(this.resolveExtensionLoadPath(path))) + .filter((extension): extension is Extension => extension !== undefined); + orderedExtensions.push(...inlineExtensions); + + const extensionsResult: LoadExtensionsResult = { + extensions: orderedExtensions, + errors: [...preTrustExtensions.errors, ...remainingExtensions.errors], + runtime: preTrustExtensions.runtime, + }; + this.addExtensionConflictDiagnostics(extensionsResult); + return extensionsResult; + } + + private addExtensionConflictDiagnostics(extensionsResult: LoadExtensionsResult): void { + // Detect extension conflicts (tools, commands, flags with same names from different extensions) + // Keep all extensions loaded. Conflicts are reported as diagnostics, and precedence is handled by load order. + const conflicts = this.detectExtensionConflicts(extensionsResult.extensions); + for (const conflict of conflicts) { + extensionsResult.errors.push({ path: conflict.path, error: conflict.message }); + } + } + + private mapSkillPath(resource: ResolvedResource, metadataByPath: Map): string { + if (resource.metadata.source !== "auto" && resource.metadata.origin !== "package") { + return resource.path; + } + try { + const stats = statSync(resource.path); + if (!stats.isDirectory()) { + return resource.path; + } + } catch { + return resource.path; + } + const skillFile = join(resource.path, "SKILL.md"); + if (existsSync(skillFile)) { + if (!metadataByPath.has(skillFile)) { + metadataByPath.set(skillFile, resource.metadata); + } + return skillFile; + } + return resource.path; + } + + private normalizeExtensionPaths( + entries: Array<{ path: string; metadata: PathMetadata }>, + ): Array<{ path: string; metadata: PathMetadata }> { + return entries.map((entry) => { + const metadata = entry.metadata.baseDir + ? { ...entry.metadata, baseDir: this.resolveResourcePath(entry.metadata.baseDir) } + : entry.metadata; + return { + path: this.resolveResourcePath(entry.path), + metadata, + }; + }); + } + + private updateSkillsFromPaths(skillPaths: string[], metadataByPath?: Map): void { + let skillsResult: { skills: Skill[]; diagnostics: ResourceDiagnostic[] }; + if (this.noSkills && skillPaths.length === 0) { + skillsResult = { skills: [], diagnostics: [] }; + } else { + skillsResult = loadSkills({ + cwd: this.cwd, + agentDir: this.agentDir, + skillPaths, + includeDefaults: false, + }); + } + const resolvedSkills = this.skillsOverride ? this.skillsOverride(skillsResult) : skillsResult; + this.skills = resolvedSkills.skills.map((skill) => ({ + ...skill, + sourceInfo: + this.findSourceInfoForPath(skill.filePath, this.extensionSkillSourceInfos, metadataByPath) ?? + skill.sourceInfo ?? + this.getDefaultSourceInfoForPath(skill.filePath), + })); + this.skillDiagnostics = resolvedSkills.diagnostics; + } + + private updatePromptsFromPaths(promptPaths: string[], metadataByPath?: Map): void { + let promptsResult: { prompts: PromptTemplate[]; diagnostics: ResourceDiagnostic[] }; + if (this.noPromptTemplates && promptPaths.length === 0) { + promptsResult = { prompts: [], diagnostics: [] }; + } else { + const allPrompts = loadPromptTemplates({ + cwd: this.cwd, + agentDir: this.agentDir, + promptPaths, + includeDefaults: false, + }); + promptsResult = this.dedupePrompts(allPrompts); + } + const resolvedPrompts = this.promptsOverride ? this.promptsOverride(promptsResult) : promptsResult; + this.prompts = resolvedPrompts.prompts.map((prompt) => ({ + ...prompt, + sourceInfo: + this.findSourceInfoForPath(prompt.filePath, this.extensionPromptSourceInfos, metadataByPath) ?? + prompt.sourceInfo ?? + this.getDefaultSourceInfoForPath(prompt.filePath), + })); + this.promptDiagnostics = resolvedPrompts.diagnostics; + } + + private updateThemesFromPaths(themePaths: string[], metadataByPath?: Map): void { + let themesResult: { themes: Theme[]; diagnostics: ResourceDiagnostic[] }; + if (this.noThemes && themePaths.length === 0) { + themesResult = { themes: [], diagnostics: [] }; + } else { + const loaded = this.loadThemes(themePaths, false); + const deduped = this.dedupeThemes(loaded.themes); + themesResult = { themes: deduped.themes, diagnostics: [...loaded.diagnostics, ...deduped.diagnostics] }; + } + const resolvedThemes = this.themesOverride ? this.themesOverride(themesResult) : themesResult; + this.themes = resolvedThemes.themes.map((theme) => { + const sourcePath = theme.sourcePath; + theme.sourceInfo = sourcePath + ? (this.findSourceInfoForPath(sourcePath, this.extensionThemeSourceInfos, metadataByPath) ?? + theme.sourceInfo ?? + this.getDefaultSourceInfoForPath(sourcePath)) + : theme.sourceInfo; + return theme; + }); + this.themeDiagnostics = resolvedThemes.diagnostics; + } + + private applyExtensionSourceInfo(extensions: Extension[], metadataByPath: Map): void { + for (const extension of extensions) { + extension.sourceInfo = + this.findSourceInfoForPath(extension.path, undefined, metadataByPath) ?? + this.getDefaultSourceInfoForPath(extension.path); + for (const command of extension.commands.values()) { + command.sourceInfo = extension.sourceInfo; + } + for (const tool of extension.tools.values()) { + tool.sourceInfo = extension.sourceInfo; + } + } + } + + private findSourceInfoForPath( + resourcePath: string, + extraSourceInfos?: Map, + metadataByPath?: Map, + ): SourceInfo | undefined { + if (!resourcePath) { + return undefined; + } + + if (resourcePath.startsWith("<")) { + return this.getDefaultSourceInfoForPath(resourcePath); + } + + const normalizedResourcePath = resolve(resourcePath); + if (extraSourceInfos) { + for (const [sourcePath, sourceInfo] of extraSourceInfos.entries()) { + const normalizedSourcePath = resolve(sourcePath); + if ( + normalizedResourcePath === normalizedSourcePath || + normalizedResourcePath.startsWith(`${normalizedSourcePath}${sep}`) + ) { + return { ...sourceInfo, path: resourcePath }; + } + } + } + + if (metadataByPath) { + const exact = metadataByPath.get(normalizedResourcePath) ?? metadataByPath.get(resourcePath); + if (exact) { + return createSourceInfo(resourcePath, exact); + } + + for (const [sourcePath, metadata] of metadataByPath.entries()) { + const normalizedSourcePath = resolve(sourcePath); + if ( + normalizedResourcePath === normalizedSourcePath || + normalizedResourcePath.startsWith(`${normalizedSourcePath}${sep}`) + ) { + return createSourceInfo(resourcePath, metadata); + } + } + } + + return undefined; + } + + private getDefaultSourceInfoForPath(filePath: string): SourceInfo { + if (filePath.startsWith("<") && filePath.endsWith(">")) { + return { + path: filePath, + source: filePath.slice(1, -1).split(":")[0] || "temporary", + scope: "temporary", + origin: "top-level", + }; + } + + const normalizedPath = resolve(filePath); + const agentRoots = [ + join(this.agentDir, "skills"), + join(this.agentDir, "prompts"), + join(this.agentDir, "themes"), + join(this.agentDir, "extensions"), + ]; + const projectRoots = [ + join(this.cwd, CONFIG_DIR_NAME, "skills"), + join(this.cwd, CONFIG_DIR_NAME, "prompts"), + join(this.cwd, CONFIG_DIR_NAME, "themes"), + join(this.cwd, CONFIG_DIR_NAME, "extensions"), + ]; + + for (const root of agentRoots) { + if (this.isUnderPath(normalizedPath, root)) { + return { path: filePath, source: "local", scope: "user", origin: "top-level", baseDir: root }; + } + } + + for (const root of projectRoots) { + if (this.isUnderPath(normalizedPath, root)) { + return { path: filePath, source: "local", scope: "project", origin: "top-level", baseDir: root }; + } + } + + return { + path: filePath, + source: "local", + scope: "temporary", + origin: "top-level", + baseDir: statSync(normalizedPath).isDirectory() ? normalizedPath : resolve(normalizedPath, ".."), + }; + } + + private mergePaths(primary: string[], additional: string[]): string[] { + const merged: string[] = []; + const seen = new Set(); + + for (const p of [...primary, ...additional]) { + const resolved = this.resolveResourcePath(p); + const canonicalPath = canonicalizePath(resolved); + if (seen.has(canonicalPath)) continue; + seen.add(canonicalPath); + merged.push(resolved); + } + + return merged; + } + + private resolveResourcePath(p: string): string { + return resolvePath(p, this.cwd, { trim: true }); + } + + private loadThemes( + paths: string[], + includeDefaults: boolean = true, + ): { + themes: Theme[]; + diagnostics: ResourceDiagnostic[]; + } { + const themes: Theme[] = []; + const diagnostics: ResourceDiagnostic[] = []; + if (includeDefaults) { + const defaultDirs = [join(this.agentDir, "themes"), join(this.cwd, CONFIG_DIR_NAME, "themes")]; + + for (const dir of defaultDirs) { + this.loadThemesFromDir(dir, themes, diagnostics); + } + } + + for (const p of paths) { + const resolved = this.resolveResourcePath(p); + if (!existsSync(resolved)) { + diagnostics.push({ type: "warning", message: "theme path does not exist", path: resolved }); + continue; + } + + try { + const stats = statSync(resolved); + if (stats.isDirectory()) { + this.loadThemesFromDir(resolved, themes, diagnostics); + } else if (stats.isFile() && resolved.endsWith(".json")) { + this.loadThemeFromFile(resolved, themes, diagnostics); + } else { + diagnostics.push({ type: "warning", message: "theme path is not a json file", path: resolved }); + } + } catch (error) { + const message = error instanceof Error ? error.message : "failed to read theme path"; + diagnostics.push({ type: "warning", message, path: resolved }); + } + } + + return { themes, diagnostics }; + } + + private loadThemesFromDir(dir: string, themes: Theme[], diagnostics: ResourceDiagnostic[]): void { + if (!existsSync(dir)) { + return; + } + + try { + const entries = readdirSync(dir, { withFileTypes: true }); + for (const entry of entries) { + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + isFile = statSync(join(dir, entry.name)).isFile(); + } catch { + continue; + } + } + if (!isFile) { + continue; + } + if (!entry.name.endsWith(".json")) { + continue; + } + this.loadThemeFromFile(join(dir, entry.name), themes, diagnostics); + } + } catch (error) { + const message = error instanceof Error ? error.message : "failed to read theme directory"; + diagnostics.push({ type: "warning", message, path: dir }); + } + } + + private loadThemeFromFile(filePath: string, themes: Theme[], diagnostics: ResourceDiagnostic[]): void { + try { + themes.push(loadThemeFromPath(filePath)); + } catch (error) { + const message = error instanceof Error ? error.message : "failed to load theme"; + diagnostics.push({ type: "warning", message, path: filePath }); + } + } + + private async loadExtensionFactories(runtime: ExtensionRuntime): Promise<{ + extensions: Extension[]; + errors: Array<{ path: string; error: string }>; + }> { + const extensions: Extension[] = []; + const errors: Array<{ path: string; error: string }> = []; + + for (const [index, factory] of this.extensionFactories.entries()) { + const extensionPath = ``; + try { + const extension = await loadExtensionFromFactory(factory, this.cwd, this.eventBus, runtime, extensionPath); + extensions.push(extension); + } catch (error) { + const message = error instanceof Error ? error.message : "failed to load extension"; + errors.push({ path: extensionPath, error: message }); + } + } + + return { extensions, errors }; + } + + private dedupePrompts(prompts: PromptTemplate[]): { prompts: PromptTemplate[]; diagnostics: ResourceDiagnostic[] } { + const seen = new Map(); + const diagnostics: ResourceDiagnostic[] = []; + + for (const prompt of prompts) { + const existing = seen.get(prompt.name); + if (existing) { + diagnostics.push({ + type: "collision", + message: `name "/${prompt.name}" collision`, + path: prompt.filePath, + collision: { + resourceType: "prompt", + name: prompt.name, + winnerPath: existing.filePath, + loserPath: prompt.filePath, + }, + }); + } else { + seen.set(prompt.name, prompt); + } + } + + return { prompts: Array.from(seen.values()), diagnostics }; + } + + private dedupeThemes(themes: Theme[]): { themes: Theme[]; diagnostics: ResourceDiagnostic[] } { + const seen = new Map(); + const diagnostics: ResourceDiagnostic[] = []; + + for (const t of themes) { + const name = t.name ?? "unnamed"; + const existing = seen.get(name); + if (existing) { + diagnostics.push({ + type: "collision", + message: `name "${name}" collision`, + path: t.sourcePath, + collision: { + resourceType: "theme", + name, + winnerPath: existing.sourcePath ?? "", + loserPath: t.sourcePath ?? "", + }, + }); + } else { + seen.set(name, t); + } + } + + return { themes: Array.from(seen.values()), diagnostics }; + } + + private discoverSystemPromptFile(): string | undefined { + const projectPath = join(this.cwd, CONFIG_DIR_NAME, "SYSTEM.md"); + if (this.settingsManager.isProjectTrusted() && existsSync(projectPath)) { + return projectPath; + } + + const globalPath = join(this.agentDir, "SYSTEM.md"); + if (existsSync(globalPath)) { + return globalPath; + } + + return undefined; + } + + private discoverAppendSystemPromptFile(): string | undefined { + const projectPath = join(this.cwd, CONFIG_DIR_NAME, "APPEND_SYSTEM.md"); + if (this.settingsManager.isProjectTrusted() && existsSync(projectPath)) { + return projectPath; + } + + const globalPath = join(this.agentDir, "APPEND_SYSTEM.md"); + if (existsSync(globalPath)) { + return globalPath; + } + + return undefined; + } + + private isUnderPath(target: string, root: string): boolean { + const normalizedRoot = resolve(root); + if (target === normalizedRoot) { + return true; + } + const prefix = normalizedRoot.endsWith(sep) ? normalizedRoot : `${normalizedRoot}${sep}`; + return target.startsWith(prefix); + } + + private detectExtensionConflicts(extensions: Extension[]): Array<{ path: string; message: string }> { + const conflicts: Array<{ path: string; message: string }> = []; + + // Track which extension registered each tool and flag + const toolOwners = new Map(); + const flagOwners = new Map(); + + for (const ext of extensions) { + // Check tools + for (const toolName of ext.tools.keys()) { + const existingOwner = toolOwners.get(toolName); + if (existingOwner && existingOwner !== ext.path) { + conflicts.push({ + path: ext.path, + message: `Tool "${toolName}" conflicts with ${existingOwner}`, + }); + } else { + toolOwners.set(toolName, ext.path); + } + } + + // Check flags + for (const flagName of ext.flags.keys()) { + const existingOwner = flagOwners.get(flagName); + if (existingOwner && existingOwner !== ext.path) { + conflicts.push({ + path: ext.path, + message: `Flag "--${flagName}" conflicts with ${existingOwner}`, + }); + } else { + flagOwners.set(flagName, ext.path); + } + } + } + + return conflicts; + } +} diff --git a/cactus-code/packages/coding-agent/src/core/sdk.ts b/cactus-code/packages/coding-agent/src/core/sdk.ts new file mode 100644 index 000000000..2175c60f8 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/sdk.ts @@ -0,0 +1,399 @@ +import { join } from "node:path"; +import { Agent, type AgentMessage, type ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { clampThinkingLevel, type Message, type Model, streamSimple } from "@earendil-works/pi-ai/compat"; +import { getAgentDir } from "../config.ts"; +import { resolvePath } from "../utils/paths.ts"; +import { AgentSession } from "./agent-session.ts"; +import { formatNoModelsAvailableMessage } from "./auth-guidance.ts"; +import { AuthStorage } from "./auth-storage.ts"; +import { DEFAULT_THINKING_LEVEL } from "./defaults.ts"; +import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.ts"; +import { convertToLlm } from "./messages.ts"; +import { ModelRegistry } from "./model-registry.ts"; +import { findInitialModel } from "./model-resolver.ts"; +import { mergeProviderAttributionHeaders } from "./provider-attribution.ts"; +import type { ResourceLoader } from "./resource-loader.ts"; +import { DefaultResourceLoader } from "./resource-loader.ts"; +import { getDefaultSessionDir, SessionManager } from "./session-manager.ts"; +import { SettingsManager } from "./settings-manager.ts"; +import { time } from "./timings.ts"; +import { + createBashTool, + createCodingTools, + createEditTool, + createFindTool, + createGrepTool, + createLsTool, + createReadOnlyTools, + createReadTool, + createWriteTool, + type ToolName, + withFileMutationQueue, +} from "./tools/index.ts"; + +export interface CreateAgentSessionOptions { + /** Working directory for project-local discovery. Default: process.cwd() */ + cwd?: string; + /** Global config directory. Default: ~/.pi/agent */ + agentDir?: string; + + /** Auth storage for credentials. Default: AuthStorage.create(agentDir/auth.json) */ + authStorage?: AuthStorage; + /** Model registry. Default: ModelRegistry.create(authStorage, agentDir/models.json) */ + modelRegistry?: ModelRegistry; + + /** Model to use. Default: from settings, else first available */ + model?: Model; + /** Thinking level. Default: from settings, else 'medium' (clamped to model capabilities) */ + thinkingLevel?: ThinkingLevel; + /** Models available for cycling (Ctrl+P in interactive mode) */ + scopedModels?: Array<{ model: Model; thinkingLevel?: ThinkingLevel }>; + + /** + * Optional default tool suppression mode when no explicit allowlist is provided. + * + * - "all": start with no tools enabled + * - "builtin": disable the default built-in tools (read, bash, edit, write) + * but keep extension/custom tools enabled + */ + noTools?: "all" | "builtin"; + /** + * Optional allowlist of tool names. + * + * When omitted, pi enables the default built-in tools (read, bash, edit, write) + * and leaves extension/custom tools enabled unless `noTools` changes that default. + * When provided, only the listed tool names are enabled. + */ + tools?: string[]; + /** Optional denylist of tool names to disable. Applies after `tools` when both are provided. */ + excludeTools?: string[]; + /** Custom tools to register (in addition to built-in tools). */ + customTools?: ToolDefinition[]; + + /** Resource loader. When omitted, DefaultResourceLoader is used. */ + resourceLoader?: ResourceLoader; + + /** Session manager. Default: SessionManager.create(cwd) */ + sessionManager?: SessionManager; + + /** Settings manager. Default: SettingsManager.create(cwd, agentDir) */ + settingsManager?: SettingsManager; + /** Session start event metadata for extension runtime startup. */ + sessionStartEvent?: SessionStartEvent; +} + +/** Result from createAgentSession */ +export interface CreateAgentSessionResult { + /** The created session */ + session: AgentSession; + /** Extensions result (for UI context setup in interactive mode) */ + extensionsResult: LoadExtensionsResult; + /** Warning if session was restored with a different model than saved */ + modelFallbackMessage?: string; +} + +// Re-exports + +export * from "./agent-session-runtime.ts"; +export type { + ExtensionAPI, + ExtensionCommandContext, + ExtensionContext, + ExtensionFactory, + SlashCommandInfo, + SlashCommandSource, + ToolDefinition, +} from "./extensions/index.ts"; +export type { PromptTemplate } from "./prompt-templates.ts"; +export type { Skill } from "./skills.ts"; +export type { Tool } from "./tools/index.ts"; + +export { + withFileMutationQueue, + // Tool factories (for custom cwd) + createCodingTools, + createReadOnlyTools, + createReadTool, + createBashTool, + createEditTool, + createWriteTool, + createGrepTool, + createFindTool, + createLsTool, +}; + +// Helper Functions + +function getDefaultAgentDir(): string { + return getAgentDir(); +} + +/** + * Create an AgentSession with the specified options. + * + * @example + * ```typescript + * // Minimal - uses defaults + * const { session } = await createAgentSession(); + * + * // With explicit model + * import { getModel } from '@earendil-works/pi-ai'; + * const { session } = await createAgentSession({ + * model: getModel('anthropic', 'claude-opus-4-5'), + * thinkingLevel: 'high', + * }); + * + * // Continue previous session + * const { session, modelFallbackMessage } = await createAgentSession({ + * continueSession: true, + * }); + * + * // Full control + * const loader = new DefaultResourceLoader({ + * cwd: process.cwd(), + * agentDir: getAgentDir(), + * settingsManager: SettingsManager.create(), + * }); + * await loader.reload(); + * const { session } = await createAgentSession({ + * model: myModel, + * tools: ["read", "bash"], + * resourceLoader: loader, + * sessionManager: SessionManager.inMemory(), + * }); + * ``` + */ +export async function createAgentSession(options: CreateAgentSessionOptions = {}): Promise { + const cwd = resolvePath(options.cwd ?? options.sessionManager?.getCwd() ?? process.cwd()); + const agentDir = options.agentDir ? resolvePath(options.agentDir) : getDefaultAgentDir(); + let resourceLoader = options.resourceLoader; + + // Use provided or create AuthStorage and ModelRegistry + const authPath = options.agentDir ? join(agentDir, "auth.json") : undefined; + const modelsPath = options.agentDir ? join(agentDir, "models.json") : undefined; + const authStorage = options.authStorage ?? AuthStorage.create(authPath); + const modelRegistry = options.modelRegistry ?? ModelRegistry.create(authStorage, modelsPath); + + const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir); + const sessionManager = options.sessionManager ?? SessionManager.create(cwd, getDefaultSessionDir(cwd, agentDir)); + + if (!resourceLoader) { + resourceLoader = new DefaultResourceLoader({ cwd, agentDir, settingsManager }); + await resourceLoader.reload(); + time("resourceLoader.reload"); + } + + // Check if session has existing data to restore + const existingSession = sessionManager.buildSessionContext(); + const hasExistingSession = existingSession.messages.length > 0; + const hasThinkingEntry = sessionManager.getBranch().some((entry) => entry.type === "thinking_level_change"); + + let model = options.model; + let modelFallbackMessage: string | undefined; + + // If session has data, try to restore model from it + if (!model && hasExistingSession && existingSession.model) { + const restoredModel = modelRegistry.find(existingSession.model.provider, existingSession.model.modelId); + if (restoredModel && modelRegistry.hasConfiguredAuth(restoredModel)) { + model = restoredModel; + } + if (!model) { + modelFallbackMessage = `Could not restore model ${existingSession.model.provider}/${existingSession.model.modelId}`; + } + } + + // If still no model, use findInitialModel (checks settings default, then provider defaults) + if (!model) { + const result = await findInitialModel({ + scopedModels: [], + isContinuing: hasExistingSession, + defaultProvider: settingsManager.getDefaultProvider(), + defaultModelId: settingsManager.getDefaultModel(), + defaultThinkingLevel: settingsManager.getDefaultThinkingLevel(), + modelRegistry, + }); + model = result.model; + if (!model) { + modelFallbackMessage = formatNoModelsAvailableMessage(); + } else if (modelFallbackMessage) { + modelFallbackMessage += `. Using ${model.provider}/${model.id}`; + } + } + + let thinkingLevel = options.thinkingLevel; + + // If session has data, restore thinking level from it + if (thinkingLevel === undefined && hasExistingSession) { + thinkingLevel = hasThinkingEntry + ? (existingSession.thinkingLevel as ThinkingLevel) + : (settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL); + } + + // Fall back to settings default + if (thinkingLevel === undefined) { + thinkingLevel = settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL; + } + + // Clamp to model capabilities + if (!model) { + thinkingLevel = "off"; + } else { + thinkingLevel = clampThinkingLevel(model, thinkingLevel) as ThinkingLevel; + } + + const defaultActiveToolNames: ToolName[] = ["read", "bash", "edit", "write", "grep", "find", "ls"]; + const allowedToolNames = options.tools ?? (options.noTools === "all" ? [] : undefined); + const excludedToolNames = options.excludeTools; + const excludedToolNameSet = excludedToolNames ? new Set(excludedToolNames) : undefined; + const initialActiveToolNames: string[] = ( + options.tools ? [...options.tools] : options.noTools ? [] : defaultActiveToolNames + ).filter((name) => !excludedToolNameSet?.has(name)); + + let agent: Agent; + + // Create convertToLlm wrapper that filters images if blockImages is enabled (defense-in-depth) + const convertToLlmWithBlockImages = (messages: AgentMessage[]): Message[] => { + const converted = convertToLlm(messages); + // Check setting dynamically so mid-session changes take effect + if (!settingsManager.getBlockImages()) { + return converted; + } + // Filter out ImageContent from all messages, replacing with text placeholder + return converted.map((msg) => { + if (msg.role === "user" || msg.role === "toolResult") { + const content = msg.content; + if (Array.isArray(content)) { + const hasImages = content.some((c) => c.type === "image"); + if (hasImages) { + const filteredContent = content + .map((c) => + c.type === "image" ? { type: "text" as const, text: "Image reading is disabled." } : c, + ) + .filter( + (c, i, arr) => + // Dedupe consecutive "Image reading is disabled." texts + !( + c.type === "text" && + c.text === "Image reading is disabled." && + i > 0 && + arr[i - 1].type === "text" && + (arr[i - 1] as { type: "text"; text: string }).text === "Image reading is disabled." + ), + ); + return { ...msg, content: filteredContent }; + } + } + } + return msg; + }); + }; + + const extensionRunnerRef: { current?: ExtensionRunner } = {}; + + agent = new Agent({ + initialState: { + systemPrompt: "", + model, + thinkingLevel, + tools: [], + }, + convertToLlm: convertToLlmWithBlockImages, + streamFn: async (model, context, options) => { + const auth = await modelRegistry.getApiKeyAndHeaders(model); + if (!auth.ok) { + throw new Error(auth.error); + } + const env = auth.env || options?.env ? { ...(auth.env ?? {}), ...(options?.env ?? {}) } : undefined; + const providerRetrySettings = settingsManager.getProviderRetrySettings(); + const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs(); + // SDKs treat timeout=0 as 0ms (immediate timeout), not "no timeout". + // Use max int32 to effectively disable the timeout. + const effectiveTimeoutMs = httpIdleTimeoutMs === 0 ? 2147483647 : httpIdleTimeoutMs; + const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs; + const websocketConnectTimeoutMs = + options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(); + return streamSimple(model, context, { + ...options, + apiKey: auth.apiKey, + env, + timeoutMs, + websocketConnectTimeoutMs, + maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries, + maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs, + headers: mergeProviderAttributionHeaders( + model, + settingsManager, + options?.sessionId, + auth.headers, + options?.headers, + ), + }); + }, + onPayload: async (payload, _model) => { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("before_provider_request")) { + return payload; + } + return runner.emitBeforeProviderRequest(payload); + }, + onResponse: async (response, _model) => { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("after_provider_response")) { + return; + } + await runner.emit({ + type: "after_provider_response", + status: response.status, + headers: response.headers, + }); + }, + sessionId: sessionManager.getSessionId(), + transformContext: async (messages) => { + const runner = extensionRunnerRef.current; + if (!runner) return messages; + return runner.emitContext(messages); + }, + steeringMode: settingsManager.getSteeringMode(), + followUpMode: settingsManager.getFollowUpMode(), + transport: settingsManager.getTransport(), + thinkingBudgets: settingsManager.getThinkingBudgets(), + maxRetryDelayMs: settingsManager.getProviderRetrySettings().maxRetryDelayMs, + }); + + // Restore messages if session has existing data + if (hasExistingSession) { + agent.state.messages = existingSession.messages; + if (!hasThinkingEntry) { + sessionManager.appendThinkingLevelChange(thinkingLevel); + } + } else { + // Save initial model and thinking level for new sessions so they can be restored on resume + if (model) { + sessionManager.appendModelChange(model.provider, model.id); + } + sessionManager.appendThinkingLevelChange(thinkingLevel); + } + + const session = new AgentSession({ + agent, + sessionManager, + settingsManager, + cwd, + scopedModels: options.scopedModels, + resourceLoader, + customTools: options.customTools, + modelRegistry, + initialActiveToolNames, + allowedToolNames, + excludedToolNames, + extensionRunnerRef, + sessionStartEvent: options.sessionStartEvent, + }); + const extensionsResult = resourceLoader.getExtensions(); + + return { + session, + extensionsResult, + modelFallbackMessage, + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/session-cwd.ts b/cactus-code/packages/coding-agent/src/core/session-cwd.ts new file mode 100644 index 000000000..79960df1f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/session-cwd.ts @@ -0,0 +1,59 @@ +import { existsSync } from "node:fs"; + +export interface SessionCwdIssue { + sessionFile?: string; + sessionCwd: string; + fallbackCwd: string; +} + +interface SessionCwdSource { + getCwd(): string; + getSessionFile(): string | undefined; +} + +export function getMissingSessionCwdIssue( + sessionManager: SessionCwdSource, + fallbackCwd: string, +): SessionCwdIssue | undefined { + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) { + return undefined; + } + + const sessionCwd = sessionManager.getCwd(); + if (!sessionCwd || existsSync(sessionCwd)) { + return undefined; + } + + return { + sessionFile, + sessionCwd, + fallbackCwd, + }; +} + +export function formatMissingSessionCwdError(issue: SessionCwdIssue): string { + const sessionFile = issue.sessionFile ? `\nSession file: ${issue.sessionFile}` : ""; + return `Stored session working directory does not exist: ${issue.sessionCwd}${sessionFile}\nCurrent working directory: ${issue.fallbackCwd}`; +} + +export function formatMissingSessionCwdPrompt(issue: SessionCwdIssue): string { + return `cwd from session file does not exist\n${issue.sessionCwd}\n\ncontinue in current cwd\n${issue.fallbackCwd}`; +} + +export class MissingSessionCwdError extends Error { + readonly issue: SessionCwdIssue; + + constructor(issue: SessionCwdIssue) { + super(formatMissingSessionCwdError(issue)); + this.name = "MissingSessionCwdError"; + this.issue = issue; + } +} + +export function assertSessionCwdExists(sessionManager: SessionCwdSource, fallbackCwd: string): void { + const issue = getMissingSessionCwdIssue(sessionManager, fallbackCwd); + if (issue) { + throw new MissingSessionCwdError(issue); + } +} diff --git a/cactus-code/packages/coding-agent/src/core/session-manager.ts b/cactus-code/packages/coding-agent/src/core/session-manager.ts new file mode 100644 index 000000000..e40ed0c47 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/session-manager.ts @@ -0,0 +1,1578 @@ +import { type AgentMessage, uuidv7 } from "@earendil-works/pi-agent-core"; +import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai"; +import { randomUUID } from "crypto"; +import { + appendFileSync, + closeSync, + createReadStream, + existsSync, + mkdirSync, + openSync, + readdirSync, + readSync, + statSync, + writeFileSync, +} from "fs"; +import { readdir, stat } from "fs/promises"; +import { join, resolve } from "path"; +import { createInterface } from "readline"; +import { StringDecoder } from "string_decoder"; +import { getAgentDir as getDefaultAgentDir, getSessionsDir } from "../config.ts"; +import { normalizePath, resolvePath } from "../utils/paths.ts"; +import { + type BashExecutionMessage, + type CustomMessage, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "./messages.ts"; + +export const CURRENT_SESSION_VERSION = 3; + +export interface SessionHeader { + type: "session"; + version?: number; // v1 sessions don't have this + id: string; + timestamp: string; + cwd: string; + parentSession?: string; +} + +export interface NewSessionOptions { + id?: string; + parentSession?: string; +} + +export interface SessionEntryBase { + type: string; + id: string; + parentId: string | null; + timestamp: string; +} + +export interface SessionMessageEntry extends SessionEntryBase { + type: "message"; + message: AgentMessage; +} + +export interface ThinkingLevelChangeEntry extends SessionEntryBase { + type: "thinking_level_change"; + thinkingLevel: string; +} + +export interface ModelChangeEntry extends SessionEntryBase { + type: "model_change"; + provider: string; + modelId: string; +} + +export interface CompactionEntry extends SessionEntryBase { + type: "compaction"; + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ + details?: T; + /** True if generated by an extension, undefined/false if pi-generated (backward compatible) */ + fromHook?: boolean; +} + +export interface BranchSummaryEntry extends SessionEntryBase { + type: "branch_summary"; + fromId: string; + summary: string; + /** Extension-specific data (not sent to LLM) */ + details?: T; + /** True if generated by an extension, false if pi-generated */ + fromHook?: boolean; +} + +/** + * Custom entry for extensions to store extension-specific data in the session. + * Use customType to identify your extension's entries. + * + * Purpose: Persist extension state across session reloads. On reload, extensions can + * scan entries for their customType and reconstruct internal state. + * + * Does NOT participate in LLM context (ignored by buildSessionContext). + * For injecting content into context, see CustomMessageEntry. + */ +export interface CustomEntry extends SessionEntryBase { + type: "custom"; + customType: string; + data?: T; +} + +/** Label entry for user-defined bookmarks/markers on entries. */ +export interface LabelEntry extends SessionEntryBase { + type: "label"; + targetId: string; + label: string | undefined; +} + +/** Session metadata entry (e.g., user-defined display name). */ +export interface SessionInfoEntry extends SessionEntryBase { + type: "session_info"; + name?: string; +} + +/** + * Custom message entry for extensions to inject messages into LLM context. + * Use customType to identify your extension's entries. + * + * Unlike CustomEntry, this DOES participate in LLM context. + * The content is converted to a user message in buildSessionContext(). + * Use details for extension-specific metadata (not sent to LLM). + * + * display controls TUI rendering: + * - false: hidden entirely + * - true: rendered with distinct styling (different from user messages) + */ +export interface CustomMessageEntry extends SessionEntryBase { + type: "custom_message"; + customType: string; + content: string | (TextContent | ImageContent)[]; + details?: T; + display: boolean; +} + +/** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */ +export type SessionEntry = + | SessionMessageEntry + | ThinkingLevelChangeEntry + | ModelChangeEntry + | CompactionEntry + | BranchSummaryEntry + | CustomEntry + | CustomMessageEntry + | LabelEntry + | SessionInfoEntry; + +/** Raw file entry (includes header) */ +export type FileEntry = SessionHeader | SessionEntry; + +/** Tree node for getTree() - defensive copy of session structure */ +export interface SessionTreeNode { + entry: SessionEntry; + children: SessionTreeNode[]; + /** Resolved label for this entry, if any */ + label?: string; + /** Timestamp of the latest label change for this entry, if any */ + labelTimestamp?: string; +} + +export interface SessionContext { + messages: AgentMessage[]; + thinkingLevel: string; + model: { provider: string; modelId: string } | null; +} + +export interface SessionInfo { + path: string; + id: string; + /** Working directory where the session was started. Empty string for old sessions. */ + cwd: string; + /** User-defined display name from session_info entries. */ + name?: string; + /** Path to the parent session (if this session was forked). */ + parentSessionPath?: string; + created: Date; + modified: Date; + messageCount: number; + firstMessage: string; + allMessagesText: string; +} + +export type ReadonlySessionManager = Pick< + SessionManager, + | "getCwd" + | "getSessionDir" + | "getSessionId" + | "getSessionFile" + | "getLeafId" + | "getLeafEntry" + | "getEntry" + | "getLabel" + | "getBranch" + | "getHeader" + | "getEntries" + | "getTree" + | "getSessionName" +>; + +function createSessionId(): string { + return uuidv7(); +} + +export function assertValidSessionId(id: string): void { + if (!/^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$/.test(id)) { + throw new Error( + "Session id must be non-empty, contain only alphanumeric characters, '-', '_', and '.', and start and end with an alphanumeric character", + ); + } +} + +/** Generate a unique short ID (8 hex chars, collision-checked) */ +function generateId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + const id = randomUUID().slice(0, 8); + if (!byId.has(id)) return id; + } + // Fallback to full UUID if somehow we have collisions + return randomUUID(); +} + +/** Migrate v1 → v2: add id/parentId tree structure. Mutates in place. */ +function migrateV1ToV2(entries: FileEntry[]): void { + const ids = new Set(); + let prevId: string | null = null; + + for (const entry of entries) { + if (entry.type === "session") { + entry.version = 2; + continue; + } + + entry.id = generateId(ids); + entry.parentId = prevId; + prevId = entry.id; + + // Convert firstKeptEntryIndex to firstKeptEntryId for compaction + if (entry.type === "compaction") { + const comp = entry as CompactionEntry & { firstKeptEntryIndex?: number }; + if (typeof comp.firstKeptEntryIndex === "number") { + const targetEntry = entries[comp.firstKeptEntryIndex]; + if (targetEntry && targetEntry.type !== "session") { + comp.firstKeptEntryId = targetEntry.id; + } + delete comp.firstKeptEntryIndex; + } + } + } +} + +/** Migrate v2 → v3: rename hookMessage role to custom. Mutates in place. */ +function migrateV2ToV3(entries: FileEntry[]): void { + for (const entry of entries) { + if (entry.type === "session") { + entry.version = 3; + continue; + } + + // Update message entries with hookMessage role + if (entry.type === "message") { + const msgEntry = entry as SessionMessageEntry; + if (msgEntry.message && (msgEntry.message as { role: string }).role === "hookMessage") { + (msgEntry.message as { role: string }).role = "custom"; + } + } + } +} + +/** + * Run all necessary migrations to bring entries to current version. + * Mutates entries in place. Returns true if any migration was applied. + */ +function migrateToCurrentVersion(entries: FileEntry[]): boolean { + const header = entries.find((e) => e.type === "session") as SessionHeader | undefined; + const version = header?.version ?? 1; + + if (version >= CURRENT_SESSION_VERSION) return false; + + if (version < 2) migrateV1ToV2(entries); + if (version < 3) migrateV2ToV3(entries); + + return true; +} + +/** Exported for testing */ +export function migrateSessionEntries(entries: FileEntry[]): void { + migrateToCurrentVersion(entries); +} + +/** Exported for compaction.test.ts */ +export function parseSessionEntries(content: string): FileEntry[] { + const entries: FileEntry[] = []; + const lines = content.trim().split("\n"); + + for (const line of lines) { + if (!line.trim()) continue; + try { + const entry = JSON.parse(line) as FileEntry; + entries.push(entry); + } catch { + // Skip malformed lines + } + } + + return entries; +} + +export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEntry | null { + for (let i = entries.length - 1; i >= 0; i--) { + if (entries[i].type === "compaction") { + return entries[i] as CompactionEntry; + } + } + return null; +} + +/** + * Build the session context from entries using tree traversal. + * If leafId is provided, walks from that entry to root. + * Handles compaction and branch summaries along the path. + */ +export function buildSessionContext( + entries: SessionEntry[], + leafId?: string | null, + byId?: Map, +): SessionContext { + // Build uuid index if not available + if (!byId) { + byId = new Map(); + for (const entry of entries) { + byId.set(entry.id, entry); + } + } + + // Find leaf + let leaf: SessionEntry | undefined; + if (leafId === null) { + // Explicitly null - return no messages (navigated to before first entry) + return { messages: [], thinkingLevel: "off", model: null }; + } + if (leafId) { + leaf = byId.get(leafId); + } + if (!leaf) { + // Fallback to last entry (when leafId is undefined) + leaf = entries[entries.length - 1]; + } + + if (!leaf) { + return { messages: [], thinkingLevel: "off", model: null }; + } + + // Walk from leaf to root, collecting path + const path: SessionEntry[] = []; + let current: SessionEntry | undefined = leaf; + while (current) { + path.push(current); + current = current.parentId ? byId.get(current.parentId) : undefined; + } + path.reverse(); + + // Extract settings and find compaction + let thinkingLevel = "off"; + let model: { provider: string; modelId: string } | null = null; + let compaction: CompactionEntry | null = null; + + for (const entry of path) { + if (entry.type === "thinking_level_change") { + thinkingLevel = entry.thinkingLevel; + } else if (entry.type === "model_change") { + model = { provider: entry.provider, modelId: entry.modelId }; + } else if (entry.type === "message" && entry.message.role === "assistant") { + model = { provider: entry.message.provider, modelId: entry.message.model }; + } else if (entry.type === "compaction") { + compaction = entry; + } + } + + // Build messages and collect corresponding entries + // When there's a compaction, we need to: + // 1. Emit summary first (entry = compaction) + // 2. Emit kept messages (from firstKeptEntryId up to compaction) + // 3. Emit messages after compaction + const messages: AgentMessage[] = []; + + const appendMessage = (entry: SessionEntry) => { + if (entry.type === "message") { + messages.push(entry.message); + } else if (entry.type === "custom_message") { + messages.push( + createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp), + ); + } else if (entry.type === "branch_summary" && entry.summary) { + messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)); + } + }; + + if (compaction) { + // Emit summary first + messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp)); + + // Find compaction index in path + const compactionIdx = path.findIndex((e) => e.type === "compaction" && e.id === compaction.id); + + // Emit kept messages (before compaction, starting from firstKeptEntryId) + let foundFirstKept = false; + for (let i = 0; i < compactionIdx; i++) { + const entry = path[i]; + if (entry.id === compaction.firstKeptEntryId) { + foundFirstKept = true; + } + if (foundFirstKept) { + appendMessage(entry); + } + } + + // Emit messages after compaction + for (let i = compactionIdx + 1; i < path.length; i++) { + const entry = path[i]; + appendMessage(entry); + } + } else { + // No compaction - emit all messages, handle branch summaries and custom messages + for (const entry of path) { + appendMessage(entry); + } + } + + return { messages, thinkingLevel, model }; +} + +/** + * Compute the default session directory for a cwd. + * Encodes cwd into a safe directory name under ~/.pi/agent/sessions/. + */ +function getDefaultSessionDirPath(cwd: string, agentDir: string = getDefaultAgentDir()): string { + const resolvedCwd = resolvePath(cwd); + const resolvedAgentDir = resolvePath(agentDir); + const safePath = `--${resolvedCwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`; + return join(resolvedAgentDir, "sessions", safePath); +} + +export function getDefaultSessionDir(cwd: string, agentDir: string = getDefaultAgentDir()): string { + const sessionDir = getDefaultSessionDirPath(cwd, agentDir); + if (!existsSync(sessionDir)) { + mkdirSync(sessionDir, { recursive: true }); + } + return sessionDir; +} + +const SESSION_READ_BUFFER_SIZE = 1024 * 1024; + +function parseSessionEntryLine(line: string): FileEntry | null { + if (!line.trim()) return null; + try { + return JSON.parse(line) as FileEntry; + } catch { + // Skip malformed lines + return null; + } +} + +/** Exported for testing */ +export function loadEntriesFromFile(filePath: string): FileEntry[] { + const resolvedFilePath = normalizePath(filePath); + if (!existsSync(resolvedFilePath)) return []; + + const entries: FileEntry[] = []; + const fd = openSync(resolvedFilePath, "r"); + try { + const decoder = new StringDecoder("utf8"); + const buffer = Buffer.allocUnsafe(SESSION_READ_BUFFER_SIZE); + let pending = ""; + + while (true) { + const bytesRead = readSync(fd, buffer, 0, buffer.length, null); + if (bytesRead === 0) break; + + pending += decoder.write(buffer.subarray(0, bytesRead)); + let lineStart = 0; + let newlineIndex = pending.indexOf("\n", lineStart); + while (newlineIndex !== -1) { + const entry = parseSessionEntryLine(pending.slice(lineStart, newlineIndex)); + if (entry) entries.push(entry); + lineStart = newlineIndex + 1; + newlineIndex = pending.indexOf("\n", lineStart); + } + pending = pending.slice(lineStart); + } + + pending += decoder.end(); + const finalEntry = parseSessionEntryLine(pending); + if (finalEntry) entries.push(finalEntry); + } finally { + closeSync(fd); + } + + // Validate session header + if (entries.length === 0) return entries; + const header = entries[0]; + if (header.type !== "session" || typeof (header as { id?: unknown }).id !== "string") { + return []; + } + + return entries; +} + +function readSessionHeader(filePath: string): SessionHeader | null { + try { + const fd = openSync(filePath, "r"); + const buffer = Buffer.alloc(512); + const bytesRead = readSync(fd, buffer, 0, 512, 0); + closeSync(fd); + const firstLine = buffer.toString("utf8", 0, bytesRead).split("\n")[0]; + if (!firstLine) return null; + const header = JSON.parse(firstLine) as Record; + if (header.type !== "session" || typeof header.id !== "string") { + return null; + } + return header as unknown as SessionHeader; + } catch { + return null; + } +} + +function getSessionHeaderCwd(header: SessionHeader): string | undefined { + const cwd = (header as { cwd?: unknown }).cwd; + return typeof cwd === "string" ? cwd : undefined; +} + +function sessionCwdMatches(cwd: string | undefined, resolvedCwd: string): boolean { + return cwd !== undefined && cwd !== "" && resolvePath(cwd) === resolvedCwd; +} + +/** Exported for testing */ +export function findMostRecentSession(sessionDir: string, cwd?: string): string | null { + const resolvedSessionDir = normalizePath(sessionDir); + const resolvedCwd = cwd ? resolvePath(cwd) : undefined; + try { + const files = readdirSync(resolvedSessionDir) + .filter((f) => f.endsWith(".jsonl")) + .map((f) => join(resolvedSessionDir, f)) + .map((path) => ({ path, header: readSessionHeader(path) })) + .filter( + (file): file is { path: string; header: SessionHeader } => + file.header !== null && + (!resolvedCwd || sessionCwdMatches(getSessionHeaderCwd(file.header), resolvedCwd)), + ) + .map(({ path }) => ({ path, mtime: statSync(path).mtime })) + .sort((a, b) => b.mtime.getTime() - a.mtime.getTime()); + + return files[0]?.path || null; + } catch { + return null; + } +} + +function isMessageWithContent(message: AgentMessage): message is Message { + return typeof (message as Message).role === "string" && "content" in message; +} + +function extractTextContent(message: Message): string { + const content = message.content; + if (typeof content === "string") { + return content; + } + return content + .filter((block): block is TextContent => block.type === "text") + .map((block) => block.text) + .join(" "); +} + +function getMessageActivityTime(entry: SessionMessageEntry): number | undefined { + const message = entry.message; + if (!isMessageWithContent(message)) return undefined; + if (message.role !== "user" && message.role !== "assistant") return undefined; + + const msgTimestamp = (message as { timestamp?: number }).timestamp; + if (typeof msgTimestamp === "number") { + return msgTimestamp; + } + + const t = new Date(entry.timestamp).getTime(); + return Number.isNaN(t) ? undefined : t; +} + +async function buildSessionInfo(filePath: string): Promise { + try { + const stats = await stat(filePath); + let header: SessionHeader | null = null; + let messageCount = 0; + let firstMessage = ""; + const allMessages: string[] = []; + let name: string | undefined; + let lastActivityTime: number | undefined; + + const rl = createInterface({ + input: createReadStream(filePath, { encoding: "utf8" }), + crlfDelay: Infinity, + }); + + for await (const line of rl) { + const entry = parseSessionEntryLine(line); + if (!entry) continue; + + if (!header) { + if (entry.type !== "session") return null; + header = entry; + continue; + } + + // Extract session name (use latest, including explicit clears) + if (entry.type === "session_info") { + name = entry.name?.trim() || undefined; + } + + if (entry.type !== "message") continue; + messageCount++; + + const activityTime = getMessageActivityTime(entry); + if (typeof activityTime === "number") { + lastActivityTime = Math.max(lastActivityTime ?? 0, activityTime); + } + + const message = entry.message; + if (!isMessageWithContent(message)) continue; + if (message.role !== "user" && message.role !== "assistant") continue; + + const textContent = extractTextContent(message); + if (!textContent) continue; + + allMessages.push(textContent); + if (!firstMessage && message.role === "user") { + firstMessage = textContent; + } + } + + if (!header) return null; + + const cwd = typeof header.cwd === "string" ? header.cwd : ""; + const parentSessionPath = header.parentSession; + const headerTime = typeof header.timestamp === "string" ? new Date(header.timestamp).getTime() : NaN; + const modified = + typeof lastActivityTime === "number" && lastActivityTime > 0 + ? new Date(lastActivityTime) + : !Number.isNaN(headerTime) + ? new Date(headerTime) + : stats.mtime; + + return { + path: filePath, + id: header.id, + cwd, + name, + parentSessionPath, + created: new Date(header.timestamp), + modified, + messageCount, + firstMessage: firstMessage || "(no messages)", + allMessagesText: allMessages.join(" "), + }; + } catch { + return null; + } +} + +export type SessionListProgress = (loaded: number, total: number) => void; + +const MAX_CONCURRENT_SESSION_INFO_LOADS = 10; + +async function buildSessionInfosWithConcurrency( + files: string[], + onLoaded: () => void, +): Promise<(SessionInfo | null)[]> { + const results: (SessionInfo | null)[] = new Array(files.length).fill(null); + const inFlight = new Set>(); + let nextIndex = 0; + + const startNext = (): void => { + const index = nextIndex++; + const file = files[index]; + if (!file) return; + + let task: Promise; + task = buildSessionInfo(file) + .then((info) => { + results[index] = info; + }) + .catch(() => { + results[index] = null; + }) + .finally(() => { + inFlight.delete(task); + onLoaded(); + }); + inFlight.add(task); + }; + + while (nextIndex < files.length || inFlight.size > 0) { + while (nextIndex < files.length && inFlight.size < MAX_CONCURRENT_SESSION_INFO_LOADS) { + startNext(); + } + if (inFlight.size > 0) { + await Promise.race(inFlight); + } + } + + return results; +} + +async function listSessionsFromDir( + dir: string, + onProgress?: SessionListProgress, + progressOffset = 0, + progressTotal?: number, +): Promise { + const sessions: SessionInfo[] = []; + if (!existsSync(dir)) { + return sessions; + } + + try { + const dirEntries = await readdir(dir); + const files = dirEntries.filter((f) => f.endsWith(".jsonl")).map((f) => join(dir, f)); + const total = progressTotal ?? files.length; + + let loaded = 0; + const results = await buildSessionInfosWithConcurrency(files, () => { + loaded++; + onProgress?.(progressOffset + loaded, total); + }); + for (const info of results) { + if (info) { + sessions.push(info); + } + } + } catch { + // Return empty list on error + } + + return sessions; +} + +/** + * Manages conversation sessions as append-only trees stored in JSONL files. + * + * Each session entry has an id and parentId forming a tree structure. The "leaf" + * pointer tracks the current position. Appending creates a child of the current leaf. + * Branching moves the leaf to an earlier entry, allowing new branches without + * modifying history. + * + * Use buildSessionContext() to get the resolved message list for the LLM, which + * handles compaction summaries and follows the path from root to current leaf. + */ +export class SessionManager { + private sessionId: string = ""; + private sessionFile: string | undefined; + private sessionDir: string; + private cwd: string; + private persist: boolean; + private flushed: boolean = false; + private fileEntries: FileEntry[] = []; + private byId: Map = new Map(); + private labelsById: Map = new Map(); + private labelTimestampsById: Map = new Map(); + private leafId: string | null = null; + + private constructor( + cwd: string, + sessionDir: string, + sessionFile: string | undefined, + persist: boolean, + newSessionOptions?: NewSessionOptions, + ) { + this.cwd = resolvePath(cwd); + this.sessionDir = normalizePath(sessionDir); + this.persist = persist; + if (persist && this.sessionDir && !existsSync(this.sessionDir)) { + mkdirSync(this.sessionDir, { recursive: true }); + } + + if (sessionFile) { + this.setSessionFile(sessionFile); + } else { + this.newSession(newSessionOptions); + } + } + + /** Switch to a different session file (used for resume and branching) */ + setSessionFile(sessionFile: string): void { + this.sessionFile = resolvePath(sessionFile); + if (existsSync(this.sessionFile)) { + this.fileEntries = loadEntriesFromFile(this.sessionFile); + + // If file was empty or corrupted (no valid header), truncate and start fresh + // to avoid appending messages without a session header (which breaks the session) + if (this.fileEntries.length === 0) { + const explicitPath = this.sessionFile; + this.newSession(); + this.sessionFile = explicitPath; + this._rewriteFile(); + this.flushed = true; + return; + } + + const header = this.fileEntries.find((e) => e.type === "session") as SessionHeader | undefined; + this.sessionId = header?.id ?? createSessionId(); + + if (migrateToCurrentVersion(this.fileEntries)) { + this._rewriteFile(); + } + + this._buildIndex(); + this.flushed = true; + } else { + const explicitPath = this.sessionFile; + this.newSession(); + this.sessionFile = explicitPath; // preserve explicit path from --session flag + } + } + + newSession(options?: NewSessionOptions): string | undefined { + if (options?.id !== undefined) { + assertValidSessionId(options.id); + } + this.sessionId = options?.id ?? createSessionId(); + const timestamp = new Date().toISOString(); + const header: SessionHeader = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: this.sessionId, + timestamp, + cwd: this.cwd, + parentSession: options?.parentSession, + }; + this.fileEntries = [header]; + this.byId.clear(); + this.labelsById.clear(); + this.leafId = null; + this.flushed = false; + + if (this.persist) { + const fileTimestamp = timestamp.replace(/[:.]/g, "-"); + this.sessionFile = join(this.getSessionDir(), `${fileTimestamp}_${this.sessionId}.jsonl`); + } + return this.sessionFile; + } + + private _buildIndex(): void { + this.byId.clear(); + this.labelsById.clear(); + this.labelTimestampsById.clear(); + this.leafId = null; + for (const entry of this.fileEntries) { + if (entry.type === "session") continue; + this.byId.set(entry.id, entry); + this.leafId = entry.id; + if (entry.type === "label") { + if (entry.label) { + this.labelsById.set(entry.targetId, entry.label); + this.labelTimestampsById.set(entry.targetId, entry.timestamp); + } else { + this.labelsById.delete(entry.targetId); + this.labelTimestampsById.delete(entry.targetId); + } + } + } + } + + private _rewriteFile(): void { + if (!this.persist || !this.sessionFile) return; + const fd = openSync(this.sessionFile, "w"); + try { + for (const entry of this.fileEntries) { + writeFileSync(fd, `${JSON.stringify(entry)}\n`); + } + } finally { + closeSync(fd); + } + } + + isPersisted(): boolean { + return this.persist; + } + + getCwd(): string { + return this.cwd; + } + + getSessionDir(): string { + return this.sessionDir; + } + + usesDefaultSessionDir(): boolean { + return this.sessionDir === getDefaultSessionDirPath(this.cwd); + } + + getSessionId(): string { + return this.sessionId; + } + + getSessionFile(): string | undefined { + return this.sessionFile; + } + + _persist(entry: SessionEntry): void { + if (!this.persist || !this.sessionFile) return; + + const hasAssistant = this.fileEntries.some((e) => e.type === "message" && e.message.role === "assistant"); + if (!hasAssistant) { + if (this.flushed) { + appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`); + } else { + // Mark as not flushed so when assistant arrives, all entries get written + this.flushed = false; + } + return; + } + + if (!this.flushed) { + const fd = openSync(this.sessionFile, "wx"); + try { + for (const e of this.fileEntries) { + writeFileSync(fd, `${JSON.stringify(e)}\n`); + } + } finally { + closeSync(fd); + } + this.flushed = true; + } else { + appendFileSync(this.sessionFile, `${JSON.stringify(entry)}\n`); + } + } + + private _appendEntry(entry: SessionEntry): void { + this.fileEntries.push(entry); + this.byId.set(entry.id, entry); + this.leafId = entry.id; + this._persist(entry); + } + + /** Append a message as child of current leaf, then advance leaf. Returns entry id. + * Does not allow writing CompactionSummaryMessage and BranchSummaryMessage directly. + * Reason: we want these to be top-level entries in the session, not message session entries, + * so it is easier to find them. + * These need to be appended via appendCompaction() and appendBranchSummary() methods. + */ + appendMessage(message: Message | CustomMessage | BashExecutionMessage): string { + const entry: SessionMessageEntry = { + type: "message", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + message, + }; + this._appendEntry(entry); + return entry.id; + } + + /** Append a thinking level change as child of current leaf, then advance leaf. Returns entry id. */ + appendThinkingLevelChange(thinkingLevel: string): string { + const entry: ThinkingLevelChangeEntry = { + type: "thinking_level_change", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + thinkingLevel, + }; + this._appendEntry(entry); + return entry.id; + } + + /** Append a model change as child of current leaf, then advance leaf. Returns entry id. */ + appendModelChange(provider: string, modelId: string): string { + const entry: ModelChangeEntry = { + type: "model_change", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + provider, + modelId, + }; + this._appendEntry(entry); + return entry.id; + } + + /** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */ + appendCompaction( + summary: string, + firstKeptEntryId: string, + tokensBefore: number, + details?: T, + fromHook?: boolean, + ): string { + const entry: CompactionEntry = { + type: "compaction", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + summary, + firstKeptEntryId, + tokensBefore, + details, + fromHook, + }; + this._appendEntry(entry); + return entry.id; + } + + /** Append a custom entry (for extensions) as child of current leaf, then advance leaf. Returns entry id. */ + appendCustomEntry(customType: string, data?: unknown): string { + const entry: CustomEntry = { + type: "custom", + customType, + data, + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + }; + this._appendEntry(entry); + return entry.id; + } + + /** Append a session info entry (e.g., display name). Returns entry id. */ + appendSessionInfo(name: string): string { + const sanitizedName = name.replace(/[\r\n]+/g, " ").trim(); + const entry: SessionInfoEntry = { + type: "session_info", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + name: sanitizedName, + }; + this._appendEntry(entry); + return entry.id; + } + + /** Get the current session name from the latest session_info entry, if any. */ + getSessionName(): string | undefined { + // Walk entries in reverse to find the latest session_info entry. + // Empty names explicitly clear the session title. + const entries = this.getEntries(); + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry.type === "session_info") { + return entry.name?.trim() || undefined; + } + } + return undefined; + } + + /** + * Append a custom message entry (for extensions) that participates in LLM context. + * @param customType Extension identifier for filtering on reload + * @param content Message content (string or TextContent/ImageContent array) + * @param display Whether to show in TUI (true = styled display, false = hidden) + * @param details Optional extension-specific metadata (not sent to LLM) + * @returns Entry id + */ + appendCustomMessageEntry( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details?: T, + ): string { + const entry: CustomMessageEntry = { + type: "custom_message", + customType, + content, + display, + details, + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + }; + this._appendEntry(entry); + return entry.id; + } + + // ========================================================================= + // Tree Traversal + // ========================================================================= + + getLeafId(): string | null { + return this.leafId; + } + + getLeafEntry(): SessionEntry | undefined { + return this.leafId ? this.byId.get(this.leafId) : undefined; + } + + getEntry(id: string): SessionEntry | undefined { + return this.byId.get(id); + } + + /** + * Get all direct children of an entry. + */ + getChildren(parentId: string): SessionEntry[] { + const children: SessionEntry[] = []; + for (const entry of this.byId.values()) { + if (entry.parentId === parentId) { + children.push(entry); + } + } + return children; + } + + /** + * Get the label for an entry, if any. + */ + getLabel(id: string): string | undefined { + return this.labelsById.get(id); + } + + /** + * Set or clear a label on an entry. + * Labels are user-defined markers for bookmarking/navigation. + * Pass undefined or empty string to clear the label. + */ + appendLabelChange(targetId: string, label: string | undefined): string { + if (!this.byId.has(targetId)) { + throw new Error(`Entry ${targetId} not found`); + } + const entry: LabelEntry = { + type: "label", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + targetId, + label, + }; + this._appendEntry(entry); + if (label) { + this.labelsById.set(targetId, label); + this.labelTimestampsById.set(targetId, entry.timestamp); + } else { + this.labelsById.delete(targetId); + this.labelTimestampsById.delete(targetId); + } + return entry.id; + } + + /** + * Walk from entry to root, returning all entries in path order. + * Includes all entry types (messages, compaction, model changes, etc.). + * Use buildSessionContext() to get the resolved messages for the LLM. + */ + getBranch(fromId?: string): SessionEntry[] { + const path: SessionEntry[] = []; + const startId = fromId ?? this.leafId; + let current = startId ? this.byId.get(startId) : undefined; + while (current) { + path.push(current); + current = current.parentId ? this.byId.get(current.parentId) : undefined; + } + path.reverse(); + return path; + } + + /** + * Build the session context (what gets sent to the LLM). + * Uses tree traversal from current leaf. + */ + buildSessionContext(): SessionContext { + return buildSessionContext(this.getEntries(), this.leafId, this.byId); + } + + /** + * Get session header. + */ + getHeader(): SessionHeader | null { + const h = this.fileEntries.find((e) => e.type === "session"); + return h ? (h as SessionHeader) : null; + } + + /** + * Get all session entries (excludes header). Returns a shallow copy. + * The session is append-only: use appendXXX() to add entries, branch() to + * change the leaf pointer. Entries cannot be modified or deleted. + */ + getEntries(): SessionEntry[] { + return this.fileEntries.filter((e): e is SessionEntry => e.type !== "session"); + } + + /** + * Get the session as a tree structure. Returns a shallow defensive copy of all entries. + * A well-formed session has exactly one root (first entry with parentId === null). + * Orphaned entries (broken parent chain) are also returned as roots. + */ + getTree(): SessionTreeNode[] { + const entries = this.getEntries(); + const nodeMap = new Map(); + const roots: SessionTreeNode[] = []; + + // Create nodes with resolved labels + for (const entry of entries) { + const label = this.labelsById.get(entry.id); + const labelTimestamp = this.labelTimestampsById.get(entry.id); + nodeMap.set(entry.id, { entry, children: [], label, labelTimestamp }); + } + + // Build tree + for (const entry of entries) { + const node = nodeMap.get(entry.id)!; + if (entry.parentId === null || entry.parentId === entry.id) { + roots.push(node); + } else { + const parent = nodeMap.get(entry.parentId); + if (parent) { + parent.children.push(node); + } else { + // Orphan - treat as root + roots.push(node); + } + } + } + + // Sort children by timestamp (oldest first, newest at bottom) + // Use iterative approach to avoid stack overflow on deep trees + const stack: SessionTreeNode[] = [...roots]; + while (stack.length > 0) { + const node = stack.pop()!; + node.children.sort((a, b) => new Date(a.entry.timestamp).getTime() - new Date(b.entry.timestamp).getTime()); + stack.push(...node.children); + } + + return roots; + } + + // ========================================================================= + // Branching + // ========================================================================= + + /** + * Start a new branch from an earlier entry. + * Moves the leaf pointer to the specified entry. The next appendXXX() call + * will create a child of that entry, forming a new branch. Existing entries + * are not modified or deleted. + */ + branch(branchFromId: string): void { + if (!this.byId.has(branchFromId)) { + throw new Error(`Entry ${branchFromId} not found`); + } + this.leafId = branchFromId; + } + + /** + * Reset the leaf pointer to null (before any entries). + * The next appendXXX() call will create a new root entry (parentId = null). + * Use this when navigating to re-edit the first user message. + */ + resetLeaf(): void { + this.leafId = null; + } + + /** + * Start a new branch with a summary of the abandoned path. + * Same as branch(), but also appends a branch_summary entry that captures + * context from the abandoned conversation path. + */ + branchWithSummary(branchFromId: string | null, summary: string, details?: unknown, fromHook?: boolean): string { + if (branchFromId !== null && !this.byId.has(branchFromId)) { + throw new Error(`Entry ${branchFromId} not found`); + } + this.leafId = branchFromId; + const entry: BranchSummaryEntry = { + type: "branch_summary", + id: generateId(this.byId), + parentId: branchFromId, + timestamp: new Date().toISOString(), + fromId: branchFromId ?? "root", + summary, + details, + fromHook, + }; + this._appendEntry(entry); + return entry.id; + } + + /** + * Create a new session file containing only the path from root to the specified leaf. + * Useful for extracting a single conversation path from a branched session. + * Returns the new session file path, or undefined if not persisting. + */ + createBranchedSession(leafId: string): string | undefined { + const previousSessionFile = this.sessionFile; + const path = this.getBranch(leafId); + if (path.length === 0) { + throw new Error(`Entry ${leafId} not found`); + } + + // Filter out LabelEntry from path - we'll recreate them from the resolved map. + // Because labels are real tree entries, later entries can be children of labels; + // removing labels requires re-chaining the retained path to avoid orphaned subtrees. + const pathWithoutLabels: SessionEntry[] = []; + let pathParentId: string | null = null; + for (const entry of path) { + if (entry.type === "label") continue; + pathWithoutLabels.push({ ...entry, parentId: pathParentId }); + pathParentId = entry.id; + } + + const newSessionId = createSessionId(); + const timestamp = new Date().toISOString(); + const fileTimestamp = timestamp.replace(/[:.]/g, "-"); + const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`); + + const header: SessionHeader = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: newSessionId, + timestamp, + cwd: this.cwd, + parentSession: this.persist ? previousSessionFile : undefined, + }; + + // Collect labels for entries in the path + const pathEntryIds = new Set(pathWithoutLabels.map((e) => e.id)); + const labelsToWrite: Array<{ targetId: string; label: string; timestamp: string }> = []; + for (const [targetId, label] of this.labelsById) { + if (pathEntryIds.has(targetId)) { + labelsToWrite.push({ targetId, label, timestamp: this.labelTimestampsById.get(targetId)! }); + } + } + + if (this.persist) { + // Build label entries + const lastEntryId = pathWithoutLabels[pathWithoutLabels.length - 1]?.id || null; + let parentId = lastEntryId; + const labelEntries: LabelEntry[] = []; + for (const { targetId, label, timestamp: labelTimestamp } of labelsToWrite) { + const labelEntry: LabelEntry = { + type: "label", + id: generateId(new Set(pathEntryIds)), + parentId, + timestamp: labelTimestamp, + targetId, + label, + }; + pathEntryIds.add(labelEntry.id); + labelEntries.push(labelEntry); + parentId = labelEntry.id; + } + + this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries]; + this.sessionId = newSessionId; + this.sessionFile = newSessionFile; + this._buildIndex(); + + // Only write the file now if it contains an assistant message. + // Otherwise defer to _persist(), which creates the file on the + // first assistant response, matching the newSession() contract + // and avoiding the duplicate-header bug when _persist()'s + // no-assistant guard later resets flushed to false. + const hasAssistant = this.fileEntries.some((e) => e.type === "message" && e.message.role === "assistant"); + if (hasAssistant) { + this._rewriteFile(); + this.flushed = true; + } else { + this.flushed = false; + } + + return newSessionFile; + } + + // In-memory mode: replace current session with the path + labels + const labelEntries: LabelEntry[] = []; + let parentId = pathWithoutLabels[pathWithoutLabels.length - 1]?.id || null; + for (const { targetId, label, timestamp: labelTimestamp } of labelsToWrite) { + const labelEntry: LabelEntry = { + type: "label", + id: generateId(new Set([...pathEntryIds, ...labelEntries.map((e) => e.id)])), + parentId, + timestamp: labelTimestamp, + targetId, + label, + }; + labelEntries.push(labelEntry); + parentId = labelEntry.id; + } + this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries]; + this.sessionId = newSessionId; + this._buildIndex(); + return undefined; + } + + /** + * Create a new session. + * @param cwd Working directory (stored in session header) + * @param sessionDir Optional session directory. If omitted, uses default (~/.pi/agent/sessions//). + */ + static create(cwd: string, sessionDir?: string, options?: NewSessionOptions): SessionManager { + const dir = sessionDir ? normalizePath(sessionDir) : getDefaultSessionDir(cwd); + return new SessionManager(cwd, dir, undefined, true, options); + } + + /** + * Open a specific session file. + * @param path Path to session file + * @param sessionDir Optional session directory for /new or /branch. If omitted, derives from file's parent. + * @param cwdOverride Optional cwd override instead of the session header cwd. + */ + static open(path: string, sessionDir?: string, cwdOverride?: string): SessionManager { + const resolvedPath = resolvePath(path); + // Extract cwd from session header if possible, otherwise use process.cwd() + const entries = loadEntriesFromFile(resolvedPath); + const header = entries.find((e) => e.type === "session") as SessionHeader | undefined; + const cwd = cwdOverride ?? header?.cwd ?? process.cwd(); + // If no sessionDir provided, derive from file's parent directory + const dir = sessionDir ? normalizePath(sessionDir) : resolve(resolvedPath, ".."); + return new SessionManager(cwd, dir, resolvedPath, true); + } + + /** + * Continue the most recent session, or create new if none. + * @param cwd Working directory + * @param sessionDir Optional session directory. If omitted, uses default (~/.pi/agent/sessions//). + */ + static continueRecent(cwd: string, sessionDir?: string): SessionManager { + const dir = sessionDir ? normalizePath(sessionDir) : getDefaultSessionDir(cwd); + const filterCwd = sessionDir !== undefined && dir !== getDefaultSessionDirPath(cwd); + const mostRecent = findMostRecentSession(dir, filterCwd ? cwd : undefined); + if (mostRecent) { + return new SessionManager(cwd, dir, mostRecent, true); + } + return new SessionManager(cwd, dir, undefined, true); + } + + /** Create an in-memory session (no file persistence) */ + static inMemory(cwd: string = process.cwd()): SessionManager { + return new SessionManager(cwd, "", undefined, false); + } + + /** + * Fork a session from another project directory into the current project. + * Creates a new session in the target cwd with the full history from the source session. + * @param sourcePath Path to the source session file + * @param targetCwd Target working directory (where the new session will be stored) + * @param sessionDir Optional session directory. If omitted, uses default for targetCwd. + */ + static forkFrom( + sourcePath: string, + targetCwd: string, + sessionDir?: string, + options?: NewSessionOptions, + ): SessionManager { + const resolvedSourcePath = resolvePath(sourcePath); + const resolvedTargetCwd = resolvePath(targetCwd); + const sourceEntries = loadEntriesFromFile(resolvedSourcePath); + if (sourceEntries.length === 0) { + throw new Error(`Cannot fork: source session file is empty or invalid: ${resolvedSourcePath}`); + } + + const sourceHeader = sourceEntries.find((e) => e.type === "session") as SessionHeader | undefined; + if (!sourceHeader) { + throw new Error(`Cannot fork: source session has no header: ${resolvedSourcePath}`); + } + + const dir = sessionDir ? normalizePath(sessionDir) : getDefaultSessionDir(resolvedTargetCwd); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + + // Create new session file with new ID but forked content + if (options?.id !== undefined) { + assertValidSessionId(options.id); + } + const newSessionId = options?.id ?? createSessionId(); + const timestamp = new Date().toISOString(); + const fileTimestamp = timestamp.replace(/[:.]/g, "-"); + const newSessionFile = join(dir, `${fileTimestamp}_${newSessionId}.jsonl`); + + // Write new header pointing to source as parent, with updated cwd + const newHeader: SessionHeader = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: newSessionId, + timestamp, + cwd: resolvedTargetCwd, + parentSession: resolvedSourcePath, + }; + writeFileSync(newSessionFile, `${JSON.stringify(newHeader)}\n`, { flag: "wx" }); + + // Copy all non-header entries from source + for (const entry of sourceEntries) { + if (entry.type !== "session") { + appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`); + } + } + + return new SessionManager(resolvedTargetCwd, dir, newSessionFile, true); + } + + /** + * List all sessions for a directory. + * @param cwd Working directory (used to compute default session directory) + * @param sessionDir Optional session directory. If omitted, uses default (~/.pi/agent/sessions//). + * @param onProgress Optional callback for progress updates (loaded, total) + */ + static async list(cwd: string, sessionDir?: string, onProgress?: SessionListProgress): Promise { + const dir = sessionDir ? normalizePath(sessionDir) : getDefaultSessionDir(cwd); + const filterCwd = sessionDir !== undefined && dir !== getDefaultSessionDirPath(cwd); + const resolvedCwd = resolvePath(cwd); + const sessions = (await listSessionsFromDir(dir, onProgress)).filter( + (session) => !filterCwd || sessionCwdMatches(session.cwd, resolvedCwd), + ); + sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime()); + return sessions; + } + + /** + * List all sessions across all project directories. + * @param onProgress Optional callback for progress updates (loaded, total) + */ + static async listAll(onProgress?: SessionListProgress): Promise; + static async listAll(sessionDir?: string, onProgress?: SessionListProgress): Promise; + static async listAll( + sessionDirOrOnProgress?: string | SessionListProgress, + onProgress?: SessionListProgress, + ): Promise { + const customSessionDir = + typeof sessionDirOrOnProgress === "string" ? normalizePath(sessionDirOrOnProgress) : undefined; + const progress = typeof sessionDirOrOnProgress === "function" ? sessionDirOrOnProgress : onProgress; + if (customSessionDir) { + const sessions = await listSessionsFromDir(customSessionDir, progress); + sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime()); + return sessions; + } + + const sessionsDir = getSessionsDir(); + + try { + if (!existsSync(sessionsDir)) { + return []; + } + const entries = await readdir(sessionsDir, { withFileTypes: true }); + const dirs = entries.filter((e) => e.isDirectory()).map((e) => join(sessionsDir, e.name)); + + // Count total files first for accurate progress + let totalFiles = 0; + const dirFiles: string[][] = []; + for (const dir of dirs) { + try { + const files = (await readdir(dir)).filter((f) => f.endsWith(".jsonl")); + dirFiles.push(files.map((f) => join(dir, f))); + totalFiles += files.length; + } catch { + dirFiles.push([]); + } + } + + // Process all files with progress tracking + let loaded = 0; + const sessions: SessionInfo[] = []; + const allFiles = dirFiles.flat(); + + const results = await buildSessionInfosWithConcurrency(allFiles, () => { + loaded++; + progress?.(loaded, totalFiles); + }); + + for (const info of results) { + if (info) { + sessions.push(info); + } + } + + sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime()); + return sessions; + } catch { + return []; + } + } +} diff --git a/cactus-code/packages/coding-agent/src/core/settings-manager.ts b/cactus-code/packages/coding-agent/src/core/settings-manager.ts new file mode 100644 index 000000000..f49fb7c19 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/settings-manager.ts @@ -0,0 +1,1175 @@ +import type { Transport } from "@earendil-works/pi-ai"; +import { randomUUID } from "crypto"; +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; +import { dirname, join } from "path"; +import lockfile from "proper-lockfile"; +import { CONFIG_DIR_NAME, getAgentDir } from "../config.ts"; +import { normalizePath, resolvePath } from "../utils/paths.ts"; +import { DEFAULT_HTTP_IDLE_TIMEOUT_MS, parseHttpIdleTimeoutMs } from "./http-dispatcher.ts"; + +export interface CompactionSettings { + enabled?: boolean; // default: true + reserveTokens?: number; // default: 16384 + keepRecentTokens?: number; // default: 20000 +} + +export interface BranchSummarySettings { + reserveTokens?: number; // default: 16384 (tokens reserved for prompt + LLM response) + skipPrompt?: boolean; // default: false - when true, skips "Summarize branch?" prompt and defaults to no summary +} + +export interface ProviderRetrySettings { + timeoutMs?: number; // SDK/provider request timeout in milliseconds + maxRetries?: number; // SDK/provider retry attempts + maxRetryDelayMs?: number; // default: 60000 (max server-requested delay before failing) +} + +export interface RetrySettings { + enabled?: boolean; // default: true + maxRetries?: number; // default: 3 + baseDelayMs?: number; // default: 2000 (exponential backoff: 2s, 4s, 8s) + provider?: ProviderRetrySettings; +} + +export interface TerminalSettings { + showImages?: boolean; // default: true (only relevant if terminal supports images) + imageWidthCells?: number; // default: 60 (preferred inline image width in terminal cells) + clearOnShrink?: boolean; // default: false (clear empty rows when content shrinks) + showTerminalProgress?: boolean; // default: false (OSC 9;4 terminal progress indicators) +} + +export interface ImageSettings { + autoResize?: boolean; // default: true (resize images to 2000x2000 max for better model compatibility) + blockImages?: boolean; // default: false - when true, prevents all images from being sent to LLM providers +} + +export interface ThinkingBudgetsSettings { + minimal?: number; + low?: number; + medium?: number; + high?: number; +} + +export interface MarkdownSettings { + codeBlockIndent?: string; // default: " " +} + +export interface WarningSettings { + anthropicExtraUsage?: boolean; // default: true +} + +export type DefaultProjectTrust = "ask" | "always" | "never"; + +export type TransportSetting = Transport; + +/** + * Package source for npm/git packages. + * - String form: load all resources from the package + * - Object form: filter which resources to load + */ +export type PackageSource = + | string + | { + source: string; + extensions?: string[]; + skills?: string[]; + prompts?: string[]; + themes?: string[]; + }; + +export interface Settings { + lastChangelogVersion?: string; + defaultProvider?: string; + defaultModel?: string; + defaultThinkingLevel?: "off" | "minimal" | "low" | "medium" | "high" | "xhigh"; + transport?: TransportSetting; // default: "auto" + steeringMode?: "all" | "one-at-a-time"; + followUpMode?: "all" | "one-at-a-time"; + theme?: string; + compaction?: CompactionSettings; + branchSummary?: BranchSummarySettings; + retry?: RetrySettings; + hideThinkingBlock?: boolean; + shellPath?: string; // Custom shell path (e.g., for Cygwin users on Windows) + quietStartup?: boolean; + defaultProjectTrust?: DefaultProjectTrust; // default: "ask"; global setting only + shellCommandPrefix?: string; // Prefix prepended to every bash command (e.g., "shopt -s expand_aliases" for alias support) + npmCommand?: string[]; // Command used for npm package lookup/install operations, argv-style (e.g., ["mise", "exec", "node@20", "--", "npm"]) + collapseChangelog?: boolean; // Show condensed changelog after update (use /changelog for full) + enableInstallTelemetry?: boolean; // default: true - anonymous version/update ping after changelog-detected updates + enableAnalytics?: boolean; // default: false - opt-in analytics data sharing + trackingId?: string; // analytics tracking identifier, generated when analytics is enabled + packages?: PackageSource[]; // Array of npm/git package sources (string or object with filtering) + extensions?: string[]; // Array of local extension file paths or directories + skills?: string[]; // Array of local skill file paths or directories + prompts?: string[]; // Array of local prompt template paths or directories + themes?: string[]; // Array of local theme file paths or directories + enableSkillCommands?: boolean; // default: true - register skills as /skill:name commands + terminal?: TerminalSettings; + images?: ImageSettings; + enabledModels?: string[]; // Model patterns for cycling (same format as --models CLI flag) + doubleEscapeAction?: "fork" | "tree" | "none"; // Action for double-escape with empty editor (default: "tree") + treeFilterMode?: "default" | "no-tools" | "user-only" | "labeled-only" | "all"; // Default filter when opening /tree + thinkingBudgets?: ThinkingBudgetsSettings; // Custom token budgets for thinking levels + editorPaddingX?: number; // Horizontal padding for input editor (default: 0) + autocompleteMaxVisible?: number; // Max visible items in autocomplete dropdown (default: 5) + showHardwareCursor?: boolean; // Show terminal cursor while still positioning it for IME + markdown?: MarkdownSettings; + warnings?: WarningSettings; + sessionDir?: string; // Custom session storage directory (same format as --session-dir CLI flag) + httpProxy?: string; // Proxy URL applied as HTTP_PROXY and HTTPS_PROXY for Pi-managed HTTP clients + httpIdleTimeoutMs?: number; // HTTP header/body idle timeout in milliseconds; 0 disables it + websocketConnectTimeoutMs?: number; // WebSocket connect/open handshake timeout in milliseconds; 0 disables it +} + +/** Deep merge settings: project/overrides take precedence, nested objects merge recursively */ +function deepMergeSettings(base: Settings, overrides: Settings): Settings { + const result: Settings = { ...base }; + + for (const key of Object.keys(overrides) as (keyof Settings)[]) { + const overrideValue = overrides[key]; + const baseValue = base[key]; + + if (overrideValue === undefined) { + continue; + } + + // For nested objects, merge recursively + if ( + typeof overrideValue === "object" && + overrideValue !== null && + !Array.isArray(overrideValue) && + typeof baseValue === "object" && + baseValue !== null && + !Array.isArray(baseValue) + ) { + (result as Record)[key] = { ...baseValue, ...overrideValue }; + } else { + // For primitives and arrays, override value wins + (result as Record)[key] = overrideValue; + } + } + + return result; +} + +function parseTimeoutSetting(value: unknown, settingName: string): number | undefined { + const timeoutMs = parseHttpIdleTimeoutMs(value); + if (timeoutMs !== undefined) { + return timeoutMs; + } + if (value !== undefined) { + throw new Error(`Invalid ${settingName} setting: ${String(value)}`); + } + return undefined; +} + +export type SettingsScope = "global" | "project"; + +export interface SettingsManagerCreateOptions { + projectTrusted?: boolean; +} + +export interface SettingsStorage { + withLock(scope: SettingsScope, fn: (current: string | undefined) => string | undefined): void; +} + +export interface SettingsError { + scope: SettingsScope; + error: Error; +} + +export class FileSettingsStorage implements SettingsStorage { + private globalSettingsPath: string; + private projectSettingsPath: string; + + constructor(cwd: string, agentDir: string) { + const resolvedCwd = resolvePath(cwd); + const resolvedAgentDir = resolvePath(agentDir); + this.globalSettingsPath = join(resolvedAgentDir, "settings.json"); + this.projectSettingsPath = join(resolvedCwd, CONFIG_DIR_NAME, "settings.json"); + } + + private acquireLockSyncWithRetry(path: string): () => void { + const maxAttempts = 10; + const delayMs = 20; + let lastError: unknown; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { + try { + return lockfile.lockSync(path, { realpath: false }); + } catch (error) { + const code = + typeof error === "object" && error !== null && "code" in error + ? String((error as { code?: unknown }).code) + : undefined; + if (code !== "ELOCKED" || attempt === maxAttempts) { + throw error; + } + lastError = error; + const start = Date.now(); + while (Date.now() - start < delayMs) { + // Sleep synchronously to avoid changing callers to async. + } + } + } + + throw (lastError as Error) ?? new Error("Failed to acquire settings lock"); + } + + withLock(scope: SettingsScope, fn: (current: string | undefined) => string | undefined): void { + const path = scope === "global" ? this.globalSettingsPath : this.projectSettingsPath; + const dir = dirname(path); + + let release: (() => void) | undefined; + try { + // Only create directory and lock if file exists or we need to write + const fileExists = existsSync(path); + if (fileExists) { + release = this.acquireLockSyncWithRetry(path); + } + const current = fileExists ? readFileSync(path, "utf-8") : undefined; + const next = fn(current); + if (next !== undefined) { + // Only create directory when we actually need to write + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + if (!release) { + release = this.acquireLockSyncWithRetry(path); + } + writeFileSync(path, next, "utf-8"); + } + } finally { + if (release) { + release(); + } + } + } +} + +export class InMemorySettingsStorage implements SettingsStorage { + private global: string | undefined; + private project: string | undefined; + + withLock(scope: SettingsScope, fn: (current: string | undefined) => string | undefined): void { + const current = scope === "global" ? this.global : this.project; + const next = fn(current); + if (next !== undefined) { + if (scope === "global") { + this.global = next; + } else { + this.project = next; + } + } + } +} + +export class SettingsManager { + private storage: SettingsStorage; + private globalSettings: Settings; + private projectSettings: Settings; + private settings: Settings; + private projectTrusted: boolean; + private modifiedFields = new Set(); // Track global fields modified during session + private modifiedNestedFields = new Map>(); // Track global nested field modifications + private modifiedProjectFields = new Set(); // Track project fields modified during session + private modifiedProjectNestedFields = new Map>(); // Track project nested field modifications + private globalSettingsLoadError: Error | null = null; // Track if global settings file had parse errors + private projectSettingsLoadError: Error | null = null; // Track if project settings file had parse errors + private writeQueue: Promise = Promise.resolve(); + private errors: SettingsError[]; + + private constructor( + storage: SettingsStorage, + initialGlobal: Settings, + initialProject: Settings, + globalLoadError: Error | null = null, + projectLoadError: Error | null = null, + initialErrors: SettingsError[] = [], + projectTrusted = true, + ) { + this.storage = storage; + this.globalSettings = initialGlobal; + this.projectSettings = initialProject; + this.projectTrusted = projectTrusted; + this.globalSettingsLoadError = globalLoadError; + this.projectSettingsLoadError = projectLoadError; + this.errors = [...initialErrors]; + this.settings = deepMergeSettings(this.globalSettings, this.projectSettings); + } + + /** Create a SettingsManager that loads from files */ + static create( + cwd: string, + agentDir: string = getAgentDir(), + options: SettingsManagerCreateOptions = {}, + ): SettingsManager { + const storage = new FileSettingsStorage(cwd, agentDir); + return SettingsManager.fromStorage(storage, options); + } + + /** Create a SettingsManager from an arbitrary storage backend */ + static fromStorage(storage: SettingsStorage, options: SettingsManagerCreateOptions = {}): SettingsManager { + const projectTrusted = options.projectTrusted ?? true; + const globalLoad = SettingsManager.tryLoadFromStorage(storage, "global"); + const projectLoad = SettingsManager.tryLoadFromStorage(storage, "project", projectTrusted); + const initialErrors: SettingsError[] = []; + if (globalLoad.error) { + initialErrors.push({ scope: "global", error: globalLoad.error }); + } + if (projectLoad.error) { + initialErrors.push({ scope: "project", error: projectLoad.error }); + } + + return new SettingsManager( + storage, + globalLoad.settings, + projectLoad.settings, + globalLoad.error, + projectLoad.error, + initialErrors, + projectTrusted, + ); + } + + /** Create an in-memory SettingsManager (no file I/O) */ + static inMemory(settings: Partial = {}, options: SettingsManagerCreateOptions = {}): SettingsManager { + const storage = new InMemorySettingsStorage(); + const initialSettings = SettingsManager.migrateSettings(structuredClone(settings) as Record); + storage.withLock("global", () => JSON.stringify(initialSettings, null, 2)); + return SettingsManager.fromStorage(storage, options); + } + + private static loadFromStorage(storage: SettingsStorage, scope: SettingsScope, projectTrusted = true): Settings { + if (scope === "project" && !projectTrusted) { + return {}; + } + + let content: string | undefined; + storage.withLock(scope, (current) => { + content = current; + return undefined; + }); + + if (!content) { + return {}; + } + const settings = JSON.parse(content); + return SettingsManager.migrateSettings(settings); + } + + private static tryLoadFromStorage( + storage: SettingsStorage, + scope: SettingsScope, + projectTrusted = true, + ): { settings: Settings; error: Error | null } { + try { + return { settings: SettingsManager.loadFromStorage(storage, scope, projectTrusted), error: null }; + } catch (error) { + return { settings: {}, error: error as Error }; + } + } + + /** Migrate old settings format to new format */ + private static migrateSettings(settings: Record): Settings { + // Migrate queueMode -> steeringMode + if ("queueMode" in settings && !("steeringMode" in settings)) { + settings.steeringMode = settings.queueMode; + delete settings.queueMode; + } + + // Migrate legacy websockets boolean -> transport enum + if (!("transport" in settings) && typeof settings.websockets === "boolean") { + settings.transport = settings.websockets ? "websocket" : "sse"; + delete settings.websockets; + } + + // Migrate old skills object format to new array format + if ( + "skills" in settings && + typeof settings.skills === "object" && + settings.skills !== null && + !Array.isArray(settings.skills) + ) { + const skillsSettings = settings.skills as { + enableSkillCommands?: boolean; + customDirectories?: unknown; + }; + if (skillsSettings.enableSkillCommands !== undefined && settings.enableSkillCommands === undefined) { + settings.enableSkillCommands = skillsSettings.enableSkillCommands; + } + if (Array.isArray(skillsSettings.customDirectories) && skillsSettings.customDirectories.length > 0) { + settings.skills = skillsSettings.customDirectories; + } else { + delete settings.skills; + } + } + + // Migrate retry.maxDelayMs -> retry.provider.maxRetryDelayMs + if ( + "retry" in settings && + typeof settings.retry === "object" && + settings.retry !== null && + !Array.isArray(settings.retry) + ) { + const retrySettings = settings.retry as Record; + const providerSettings = + typeof retrySettings.provider === "object" && retrySettings.provider !== null + ? (retrySettings.provider as Record) + : undefined; + if ( + typeof retrySettings.maxDelayMs === "number" && + (providerSettings?.maxRetryDelayMs === undefined || providerSettings?.maxRetryDelayMs === null) + ) { + retrySettings.provider = { + ...(providerSettings ?? {}), + maxRetryDelayMs: retrySettings.maxDelayMs, + }; + } + delete retrySettings.maxDelayMs; + } + + return settings as Settings; + } + + getGlobalSettings(): Settings { + return structuredClone(this.globalSettings); + } + + getProjectSettings(): Settings { + return structuredClone(this.projectSettings); + } + + isProjectTrusted(): boolean { + return this.projectTrusted; + } + + setProjectTrusted(trusted: boolean): void { + if (this.projectTrusted === trusted) { + return; + } + + this.projectTrusted = trusted; + this.modifiedProjectFields.clear(); + this.modifiedProjectNestedFields.clear(); + + if (!trusted) { + this.projectSettings = {}; + this.projectSettingsLoadError = null; + this.settings = deepMergeSettings(this.globalSettings, this.projectSettings); + return; + } + + const projectLoad = SettingsManager.tryLoadFromStorage(this.storage, "project", trusted); + this.projectSettings = projectLoad.settings; + this.projectSettingsLoadError = projectLoad.error; + if (projectLoad.error) { + this.recordError("project", projectLoad.error); + } + this.settings = deepMergeSettings(this.globalSettings, this.projectSettings); + } + + async reload(): Promise { + await this.writeQueue; + const globalLoad = SettingsManager.tryLoadFromStorage(this.storage, "global"); + if (!globalLoad.error) { + this.globalSettings = globalLoad.settings; + this.globalSettingsLoadError = null; + } else { + this.globalSettingsLoadError = globalLoad.error; + this.recordError("global", globalLoad.error); + } + + this.modifiedFields.clear(); + this.modifiedNestedFields.clear(); + this.modifiedProjectFields.clear(); + this.modifiedProjectNestedFields.clear(); + + const projectLoad = SettingsManager.tryLoadFromStorage(this.storage, "project", this.projectTrusted); + if (!projectLoad.error) { + this.projectSettings = projectLoad.settings; + this.projectSettingsLoadError = null; + } else { + this.projectSettingsLoadError = projectLoad.error; + this.recordError("project", projectLoad.error); + } + + this.settings = deepMergeSettings(this.globalSettings, this.projectSettings); + } + + /** Apply additional overrides on top of current settings */ + applyOverrides(overrides: Partial): void { + this.settings = deepMergeSettings(this.settings, overrides); + } + + /** Mark a global field as modified during this session */ + private markModified(field: keyof Settings, nestedKey?: string): void { + this.modifiedFields.add(field); + if (nestedKey) { + if (!this.modifiedNestedFields.has(field)) { + this.modifiedNestedFields.set(field, new Set()); + } + this.modifiedNestedFields.get(field)!.add(nestedKey); + } + } + + /** Mark a project field as modified during this session */ + private markProjectModified(field: keyof Settings, nestedKey?: string): void { + this.modifiedProjectFields.add(field); + if (nestedKey) { + if (!this.modifiedProjectNestedFields.has(field)) { + this.modifiedProjectNestedFields.set(field, new Set()); + } + this.modifiedProjectNestedFields.get(field)!.add(nestedKey); + } + } + + private assertProjectTrustedForWrite(): void { + if (!this.projectTrusted) { + throw new Error("Project is not trusted; refusing to write project settings"); + } + } + + private recordError(scope: SettingsScope, error: unknown): void { + const normalizedError = error instanceof Error ? error : new Error(String(error)); + this.errors.push({ scope, error: normalizedError }); + } + + private clearModifiedScope(scope: SettingsScope): void { + if (scope === "global") { + this.modifiedFields.clear(); + this.modifiedNestedFields.clear(); + return; + } + + this.modifiedProjectFields.clear(); + this.modifiedProjectNestedFields.clear(); + } + + private enqueueWrite(scope: SettingsScope, task: () => void): void { + this.writeQueue = this.writeQueue + .then(() => { + if (scope === "project") { + this.assertProjectTrustedForWrite(); + } + task(); + this.clearModifiedScope(scope); + }) + .catch((error) => { + this.recordError(scope, error); + }); + } + + private cloneModifiedNestedFields(source: Map>): Map> { + const snapshot = new Map>(); + for (const [key, value] of source.entries()) { + snapshot.set(key, new Set(value)); + } + return snapshot; + } + + private persistScopedSettings( + scope: SettingsScope, + snapshotSettings: Settings, + modifiedFields: Set, + modifiedNestedFields: Map>, + ): void { + this.storage.withLock(scope, (current) => { + const currentFileSettings = current + ? SettingsManager.migrateSettings(JSON.parse(current) as Record) + : {}; + const mergedSettings: Settings = { ...currentFileSettings }; + for (const field of modifiedFields) { + const value = snapshotSettings[field]; + if (modifiedNestedFields.has(field) && typeof value === "object" && value !== null) { + const nestedModified = modifiedNestedFields.get(field)!; + const baseNested = (currentFileSettings[field] as Record) ?? {}; + const inMemoryNested = value as Record; + const mergedNested = { ...baseNested }; + for (const nestedKey of nestedModified) { + mergedNested[nestedKey] = inMemoryNested[nestedKey]; + } + (mergedSettings as Record)[field] = mergedNested; + } else { + (mergedSettings as Record)[field] = value; + } + } + + return JSON.stringify(mergedSettings, null, 2); + }); + } + + private save(): void { + this.settings = deepMergeSettings(this.globalSettings, this.projectSettings); + + if (this.globalSettingsLoadError) { + return; + } + + const snapshotGlobalSettings = structuredClone(this.globalSettings); + const modifiedFields = new Set(this.modifiedFields); + const modifiedNestedFields = this.cloneModifiedNestedFields(this.modifiedNestedFields); + + this.enqueueWrite("global", () => { + this.persistScopedSettings("global", snapshotGlobalSettings, modifiedFields, modifiedNestedFields); + }); + } + + private saveProjectSettings(settings: Settings): void { + this.assertProjectTrustedForWrite(); + this.projectSettings = structuredClone(settings); + this.settings = deepMergeSettings(this.globalSettings, this.projectSettings); + + if (this.projectSettingsLoadError) { + return; + } + + const snapshotProjectSettings = structuredClone(this.projectSettings); + const modifiedFields = new Set(this.modifiedProjectFields); + const modifiedNestedFields = this.cloneModifiedNestedFields(this.modifiedProjectNestedFields); + this.enqueueWrite("project", () => { + this.persistScopedSettings("project", snapshotProjectSettings, modifiedFields, modifiedNestedFields); + }); + } + + private updateProjectSettings(field: keyof Settings, update: (settings: Settings) => void): void { + this.assertProjectTrustedForWrite(); + const projectSettings = structuredClone(this.projectSettings); + update(projectSettings); + this.markProjectModified(field); + this.saveProjectSettings(projectSettings); + } + + async flush(): Promise { + await this.writeQueue; + } + + drainErrors(): SettingsError[] { + const drained = [...this.errors]; + this.errors = []; + return drained; + } + + getLastChangelogVersion(): string | undefined { + return this.settings.lastChangelogVersion; + } + + setLastChangelogVersion(version: string): void { + this.globalSettings.lastChangelogVersion = version; + this.markModified("lastChangelogVersion"); + this.save(); + } + + getSessionDir(): string | undefined { + const sessionDir = this.settings.sessionDir; + return sessionDir ? normalizePath(sessionDir) : sessionDir; + } + + getDefaultProvider(): string | undefined { + return this.settings.defaultProvider; + } + + getDefaultModel(): string | undefined { + return this.settings.defaultModel; + } + + setDefaultProvider(provider: string): void { + this.globalSettings.defaultProvider = provider; + this.markModified("defaultProvider"); + this.save(); + } + + setDefaultModel(modelId: string): void { + this.globalSettings.defaultModel = modelId; + this.markModified("defaultModel"); + this.save(); + } + + setDefaultModelAndProvider(provider: string, modelId: string): void { + this.globalSettings.defaultProvider = provider; + this.globalSettings.defaultModel = modelId; + this.markModified("defaultProvider"); + this.markModified("defaultModel"); + this.save(); + } + + getSteeringMode(): "all" | "one-at-a-time" { + return this.settings.steeringMode || "one-at-a-time"; + } + + setSteeringMode(mode: "all" | "one-at-a-time"): void { + this.globalSettings.steeringMode = mode; + this.markModified("steeringMode"); + this.save(); + } + + getFollowUpMode(): "all" | "one-at-a-time" { + return this.settings.followUpMode || "one-at-a-time"; + } + + setFollowUpMode(mode: "all" | "one-at-a-time"): void { + this.globalSettings.followUpMode = mode; + this.markModified("followUpMode"); + this.save(); + } + + getThemeSetting(): string | undefined { + const value = this.settings.theme; + if (typeof value === "string") return value; + return undefined; + } + + getTheme(): string | undefined { + const theme = this.getThemeSetting(); + return theme?.includes("/") ? undefined : theme; + } + + setTheme(theme: string): void { + this.globalSettings.theme = theme; + this.markModified("theme"); + this.save(); + } + + getDefaultThinkingLevel(): "off" | "minimal" | "low" | "medium" | "high" | "xhigh" | undefined { + return this.settings.defaultThinkingLevel; + } + + setDefaultThinkingLevel(level: "off" | "minimal" | "low" | "medium" | "high" | "xhigh"): void { + this.globalSettings.defaultThinkingLevel = level; + this.markModified("defaultThinkingLevel"); + this.save(); + } + + getTransport(): TransportSetting { + return this.settings.transport ?? "auto"; + } + + setTransport(transport: TransportSetting): void { + this.globalSettings.transport = transport; + this.markModified("transport"); + this.save(); + } + + getCompactionEnabled(): boolean { + return this.settings.compaction?.enabled ?? true; + } + + setCompactionEnabled(enabled: boolean): void { + if (!this.globalSettings.compaction) { + this.globalSettings.compaction = {}; + } + this.globalSettings.compaction.enabled = enabled; + this.markModified("compaction", "enabled"); + this.save(); + } + + getCompactionReserveTokens(): number { + return this.settings.compaction?.reserveTokens ?? 16384; + } + + getCompactionKeepRecentTokens(): number { + return this.settings.compaction?.keepRecentTokens ?? 20000; + } + + getCompactionSettings(): { enabled: boolean; reserveTokens: number; keepRecentTokens: number } { + return { + enabled: this.getCompactionEnabled(), + reserveTokens: this.getCompactionReserveTokens(), + keepRecentTokens: this.getCompactionKeepRecentTokens(), + }; + } + + getBranchSummarySettings(): { reserveTokens: number; skipPrompt: boolean } { + return { + reserveTokens: this.settings.branchSummary?.reserveTokens ?? 16384, + skipPrompt: this.settings.branchSummary?.skipPrompt ?? false, + }; + } + + getBranchSummarySkipPrompt(): boolean { + return this.settings.branchSummary?.skipPrompt ?? false; + } + + getRetryEnabled(): boolean { + return this.settings.retry?.enabled ?? true; + } + + setRetryEnabled(enabled: boolean): void { + if (!this.globalSettings.retry) { + this.globalSettings.retry = {}; + } + this.globalSettings.retry.enabled = enabled; + this.markModified("retry", "enabled"); + this.save(); + } + + getRetrySettings(): { enabled: boolean; maxRetries: number; baseDelayMs: number } { + return { + enabled: this.getRetryEnabled(), + maxRetries: this.settings.retry?.maxRetries ?? 3, + baseDelayMs: this.settings.retry?.baseDelayMs ?? 2000, + }; + } + + getHttpIdleTimeoutMs(): number { + return parseTimeoutSetting(this.settings.httpIdleTimeoutMs, "httpIdleTimeoutMs") ?? DEFAULT_HTTP_IDLE_TIMEOUT_MS; + } + + setHttpIdleTimeoutMs(timeoutMs: number): void { + if (!Number.isFinite(timeoutMs) || timeoutMs < 0) { + throw new Error(`Invalid httpIdleTimeoutMs setting: ${String(timeoutMs)}`); + } + this.globalSettings.httpIdleTimeoutMs = Math.floor(timeoutMs); + this.markModified("httpIdleTimeoutMs"); + this.save(); + } + + getProviderRetrySettings(): { timeoutMs?: number; maxRetries?: number; maxRetryDelayMs: number } { + return { + timeoutMs: this.settings.retry?.provider?.timeoutMs, + maxRetries: this.settings.retry?.provider?.maxRetries, + maxRetryDelayMs: this.settings.retry?.provider?.maxRetryDelayMs ?? 60000, + }; + } + + getWebSocketConnectTimeoutMs(): number | undefined { + return parseTimeoutSetting(this.settings.websocketConnectTimeoutMs, "websocketConnectTimeoutMs"); + } + + getHideThinkingBlock(): boolean { + return this.settings.hideThinkingBlock ?? false; + } + + setHideThinkingBlock(hide: boolean): void { + this.globalSettings.hideThinkingBlock = hide; + this.markModified("hideThinkingBlock"); + this.save(); + } + + getShellPath(): string | undefined { + return this.settings.shellPath; + } + + setShellPath(path: string | undefined): void { + this.globalSettings.shellPath = path; + this.markModified("shellPath"); + this.save(); + } + + getQuietStartup(): boolean { + return this.settings.quietStartup ?? false; + } + + setQuietStartup(quiet: boolean): void { + this.globalSettings.quietStartup = quiet; + this.markModified("quietStartup"); + this.save(); + } + + getDefaultProjectTrust(): DefaultProjectTrust { + const value = this.globalSettings.defaultProjectTrust; + return value === "always" || value === "never" ? value : "ask"; + } + + setDefaultProjectTrust(defaultProjectTrust: DefaultProjectTrust): void { + this.globalSettings.defaultProjectTrust = defaultProjectTrust; + this.markModified("defaultProjectTrust"); + this.save(); + } + + getShellCommandPrefix(): string | undefined { + return this.settings.shellCommandPrefix; + } + + setShellCommandPrefix(prefix: string | undefined): void { + this.globalSettings.shellCommandPrefix = prefix; + this.markModified("shellCommandPrefix"); + this.save(); + } + + getNpmCommand(): string[] | undefined { + return this.settings.npmCommand ? [...this.settings.npmCommand] : undefined; + } + + setNpmCommand(command: string[] | undefined): void { + this.globalSettings.npmCommand = command ? [...command] : undefined; + this.markModified("npmCommand"); + this.save(); + } + + getEnableAnalytics(): boolean { + return this.settings.enableAnalytics ?? false; + } + + getTrackingId(): string | undefined { + return this.settings.trackingId; + } + + /** Set the analytics opt-in preference; generates a tracking identifier on first opt-in */ + setEnableAnalytics(enabled: boolean): void { + this.globalSettings.enableAnalytics = enabled; + this.markModified("enableAnalytics"); + if (enabled && !this.globalSettings.trackingId) { + this.globalSettings.trackingId = randomUUID(); + this.markModified("trackingId"); + } + this.save(); + } + + getPackages(): PackageSource[] { + return [...(this.settings.packages ?? [])]; + } + + setPackages(packages: PackageSource[]): void { + this.globalSettings.packages = packages; + this.markModified("packages"); + this.save(); + } + + setProjectPackages(packages: PackageSource[]): void { + this.updateProjectSettings("packages", (settings) => { + settings.packages = packages; + }); + } + + getExtensionPaths(): string[] { + return [...(this.settings.extensions ?? [])]; + } + + setExtensionPaths(paths: string[]): void { + this.globalSettings.extensions = paths; + this.markModified("extensions"); + this.save(); + } + + setProjectExtensionPaths(paths: string[]): void { + this.updateProjectSettings("extensions", (settings) => { + settings.extensions = paths; + }); + } + + getSkillPaths(): string[] { + return [...(this.settings.skills ?? [])]; + } + + setSkillPaths(paths: string[]): void { + this.globalSettings.skills = paths; + this.markModified("skills"); + this.save(); + } + + setProjectSkillPaths(paths: string[]): void { + this.updateProjectSettings("skills", (settings) => { + settings.skills = paths; + }); + } + + getPromptTemplatePaths(): string[] { + return [...(this.settings.prompts ?? [])]; + } + + setPromptTemplatePaths(paths: string[]): void { + this.globalSettings.prompts = paths; + this.markModified("prompts"); + this.save(); + } + + setProjectPromptTemplatePaths(paths: string[]): void { + this.updateProjectSettings("prompts", (settings) => { + settings.prompts = paths; + }); + } + + getThemePaths(): string[] { + return [...(this.settings.themes ?? [])]; + } + + setThemePaths(paths: string[]): void { + this.globalSettings.themes = paths; + this.markModified("themes"); + this.save(); + } + + setProjectThemePaths(paths: string[]): void { + this.updateProjectSettings("themes", (settings) => { + settings.themes = paths; + }); + } + + getEnableSkillCommands(): boolean { + return this.settings.enableSkillCommands ?? true; + } + + setEnableSkillCommands(enabled: boolean): void { + this.globalSettings.enableSkillCommands = enabled; + this.markModified("enableSkillCommands"); + this.save(); + } + + getThinkingBudgets(): ThinkingBudgetsSettings | undefined { + return this.settings.thinkingBudgets; + } + + getShowImages(): boolean { + return this.settings.terminal?.showImages ?? true; + } + + setShowImages(show: boolean): void { + if (!this.globalSettings.terminal) { + this.globalSettings.terminal = {}; + } + this.globalSettings.terminal.showImages = show; + this.markModified("terminal", "showImages"); + this.save(); + } + + getImageWidthCells(): number { + const width = this.settings.terminal?.imageWidthCells; + if (typeof width !== "number" || !Number.isFinite(width)) { + return 60; + } + return Math.max(1, Math.floor(width)); + } + + setImageWidthCells(width: number): void { + if (!this.globalSettings.terminal) { + this.globalSettings.terminal = {}; + } + this.globalSettings.terminal.imageWidthCells = Math.max(1, Math.floor(width)); + this.markModified("terminal", "imageWidthCells"); + this.save(); + } + + getClearOnShrink(): boolean { + // Settings takes precedence, then env var, then default false + if (this.settings.terminal?.clearOnShrink !== undefined) { + return this.settings.terminal.clearOnShrink; + } + return process.env.PI_CLEAR_ON_SHRINK === "1"; + } + + setClearOnShrink(enabled: boolean): void { + if (!this.globalSettings.terminal) { + this.globalSettings.terminal = {}; + } + this.globalSettings.terminal.clearOnShrink = enabled; + this.markModified("terminal", "clearOnShrink"); + this.save(); + } + + getShowTerminalProgress(): boolean { + return this.settings.terminal?.showTerminalProgress ?? false; + } + + setShowTerminalProgress(enabled: boolean): void { + if (!this.globalSettings.terminal) { + this.globalSettings.terminal = {}; + } + this.globalSettings.terminal.showTerminalProgress = enabled; + this.markModified("terminal", "showTerminalProgress"); + this.save(); + } + + getImageAutoResize(): boolean { + return this.settings.images?.autoResize ?? true; + } + + setImageAutoResize(enabled: boolean): void { + if (!this.globalSettings.images) { + this.globalSettings.images = {}; + } + this.globalSettings.images.autoResize = enabled; + this.markModified("images", "autoResize"); + this.save(); + } + + getBlockImages(): boolean { + return this.settings.images?.blockImages ?? false; + } + + setBlockImages(blocked: boolean): void { + if (!this.globalSettings.images) { + this.globalSettings.images = {}; + } + this.globalSettings.images.blockImages = blocked; + this.markModified("images", "blockImages"); + this.save(); + } + + getEnabledModels(): string[] | undefined { + return this.settings.enabledModels; + } + + setEnabledModels(patterns: string[] | undefined): void { + this.globalSettings.enabledModels = patterns; + this.markModified("enabledModels"); + this.save(); + } + + getDoubleEscapeAction(): "fork" | "tree" | "none" { + return this.settings.doubleEscapeAction ?? "tree"; + } + + setDoubleEscapeAction(action: "fork" | "tree" | "none"): void { + this.globalSettings.doubleEscapeAction = action; + this.markModified("doubleEscapeAction"); + this.save(); + } + + getTreeFilterMode(): "default" | "no-tools" | "user-only" | "labeled-only" | "all" { + const mode = this.settings.treeFilterMode; + const valid = ["default", "no-tools", "user-only", "labeled-only", "all"]; + return mode && valid.includes(mode) ? mode : "default"; + } + + setTreeFilterMode(mode: "default" | "no-tools" | "user-only" | "labeled-only" | "all"): void { + this.globalSettings.treeFilterMode = mode; + this.markModified("treeFilterMode"); + this.save(); + } + + getShowHardwareCursor(): boolean { + return this.settings.showHardwareCursor ?? process.env.PI_HARDWARE_CURSOR === "1"; + } + + setShowHardwareCursor(enabled: boolean): void { + this.globalSettings.showHardwareCursor = enabled; + this.markModified("showHardwareCursor"); + this.save(); + } + + getEditorPaddingX(): number { + return this.settings.editorPaddingX ?? 0; + } + + setEditorPaddingX(padding: number): void { + this.globalSettings.editorPaddingX = Math.max(0, Math.min(3, Math.floor(padding))); + this.markModified("editorPaddingX"); + this.save(); + } + + getAutocompleteMaxVisible(): number { + return this.settings.autocompleteMaxVisible ?? 5; + } + + setAutocompleteMaxVisible(maxVisible: number): void { + this.globalSettings.autocompleteMaxVisible = Math.max(3, Math.min(20, Math.floor(maxVisible))); + this.markModified("autocompleteMaxVisible"); + this.save(); + } + + getCodeBlockIndent(): string { + return this.settings.markdown?.codeBlockIndent ?? " "; + } + + getWarnings(): WarningSettings { + return { ...(this.settings.warnings ?? {}) }; + } + + setWarnings(warnings: WarningSettings): void { + this.globalSettings.warnings = { ...warnings }; + this.markModified("warnings"); + this.save(); + } +} diff --git a/cactus-code/packages/coding-agent/src/core/skills.ts b/cactus-code/packages/coding-agent/src/core/skills.ts new file mode 100644 index 000000000..c104c8564 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/skills.ts @@ -0,0 +1,487 @@ +import { existsSync, readdirSync, readFileSync, statSync } from "fs"; +import ignore from "ignore"; +import { basename, dirname, join, relative, resolve, sep } from "path"; +import { CONFIG_DIR_NAME, getAgentDir } from "../config.ts"; +import { parseFrontmatter } from "../utils/frontmatter.ts"; +import { canonicalizePath, resolvePath } from "../utils/paths.ts"; +import type { ResourceDiagnostic } from "./diagnostics.ts"; +import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.ts"; + +/** Max name length per spec */ +const MAX_NAME_LENGTH = 64; + +/** Max description length per spec */ +const MAX_DESCRIPTION_LENGTH = 1024; + +const IGNORE_FILE_NAMES = [".gitignore", ".ignore", ".fdignore"]; + +type IgnoreMatcher = ReturnType; + +function toPosixPath(p: string): string { + return p.split(sep).join("/"); +} + +function prefixIgnorePattern(line: string, prefix: string): string | null { + const trimmed = line.trim(); + if (!trimmed) return null; + if (trimmed.startsWith("#") && !trimmed.startsWith("\\#")) return null; + + let pattern = line; + let negated = false; + + if (pattern.startsWith("!")) { + negated = true; + pattern = pattern.slice(1); + } else if (pattern.startsWith("\\!")) { + pattern = pattern.slice(1); + } + + if (pattern.startsWith("/")) { + pattern = pattern.slice(1); + } + + const prefixed = prefix ? `${prefix}${pattern}` : pattern; + return negated ? `!${prefixed}` : prefixed; +} + +function addIgnoreRules(ig: IgnoreMatcher, dir: string, rootDir: string): void { + const relativeDir = relative(rootDir, dir); + const prefix = relativeDir ? `${toPosixPath(relativeDir)}/` : ""; + + for (const filename of IGNORE_FILE_NAMES) { + const ignorePath = join(dir, filename); + if (!existsSync(ignorePath)) continue; + try { + const content = readFileSync(ignorePath, "utf-8"); + const patterns = content + .split(/\r?\n/) + .map((line) => prefixIgnorePattern(line, prefix)) + .filter((line): line is string => Boolean(line)); + if (patterns.length > 0) { + ig.add(patterns); + } + } catch {} + } +} + +export interface SkillFrontmatter { + name?: string; + description?: string; + "disable-model-invocation"?: boolean; + [key: string]: unknown; +} + +export interface Skill { + name: string; + description: string; + filePath: string; + baseDir: string; + sourceInfo: SourceInfo; + disableModelInvocation: boolean; +} + +export interface LoadSkillsResult { + skills: Skill[]; + diagnostics: ResourceDiagnostic[]; +} + +/** + * Validate skill name per Agent Skills spec. + * Returns array of validation error messages (empty if valid). + */ +function validateName(name: string): string[] { + const errors: string[] = []; + + if (name.length > MAX_NAME_LENGTH) { + errors.push(`name exceeds ${MAX_NAME_LENGTH} characters (${name.length})`); + } + + if (!/^[a-z0-9-]+$/.test(name)) { + errors.push(`name contains invalid characters (must be lowercase a-z, 0-9, hyphens only)`); + } + + if (name.startsWith("-") || name.endsWith("-")) { + errors.push(`name must not start or end with a hyphen`); + } + + if (name.includes("--")) { + errors.push(`name must not contain consecutive hyphens`); + } + + return errors; +} + +/** + * Validate description per Agent Skills spec. + */ +function validateDescription(description: string | undefined): string[] { + const errors: string[] = []; + + if (!description || description.trim() === "") { + errors.push("description is required"); + } else if (description.length > MAX_DESCRIPTION_LENGTH) { + errors.push(`description exceeds ${MAX_DESCRIPTION_LENGTH} characters (${description.length})`); + } + + return errors; +} + +export interface LoadSkillsFromDirOptions { + /** Directory to scan for skills */ + dir: string; + /** Source identifier for these skills */ + source: string; +} + +function createSkillSourceInfo(filePath: string, baseDir: string, source: string): SourceInfo { + switch (source) { + case "user": + return createSyntheticSourceInfo(filePath, { + source: "local", + scope: "user", + baseDir, + }); + case "project": + return createSyntheticSourceInfo(filePath, { + source: "local", + scope: "project", + baseDir, + }); + case "path": + return createSyntheticSourceInfo(filePath, { + source: "local", + baseDir, + }); + default: + return createSyntheticSourceInfo(filePath, { source, baseDir }); + } +} + +/** + * Load skills from a directory. + * + * Discovery rules: + * - if a directory contains SKILL.md, treat it as a skill root and do not recurse further + * - otherwise, load direct .md children in the root + * - recurse into subdirectories to find SKILL.md + */ +export function loadSkillsFromDir(options: LoadSkillsFromDirOptions): LoadSkillsResult { + const { dir, source } = options; + return loadSkillsFromDirInternal(dir, source, true); +} + +function loadSkillsFromDirInternal( + dir: string, + source: string, + includeRootFiles: boolean, + ignoreMatcher?: IgnoreMatcher, + rootDir?: string, +): LoadSkillsResult { + const skills: Skill[] = []; + const diagnostics: ResourceDiagnostic[] = []; + + if (!existsSync(dir)) { + return { skills, diagnostics }; + } + + const root = rootDir ?? dir; + const ig = ignoreMatcher ?? ignore(); + addIgnoreRules(ig, dir, root); + + try { + const entries = readdirSync(dir, { withFileTypes: true }); + + for (const entry of entries) { + if (entry.name !== "SKILL.md") { + continue; + } + + const fullPath = join(dir, entry.name); + + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + isFile = statSync(fullPath).isFile(); + } catch { + continue; + } + } + + const relPath = toPosixPath(relative(root, fullPath)); + if (!isFile || ig.ignores(relPath)) { + continue; + } + + const result = loadSkillFromFile(fullPath, source); + if (result.skill) { + skills.push(result.skill); + } + diagnostics.push(...result.diagnostics); + return { skills, diagnostics }; + } + + for (const entry of entries) { + if (entry.name.startsWith(".")) { + continue; + } + + // Skip node_modules to avoid scanning dependencies + if (entry.name === "node_modules") { + continue; + } + + const fullPath = join(dir, entry.name); + + // For symlinks, check if they point to a directory and follow them + let isDirectory = entry.isDirectory(); + let isFile = entry.isFile(); + if (entry.isSymbolicLink()) { + try { + const stats = statSync(fullPath); + isDirectory = stats.isDirectory(); + isFile = stats.isFile(); + } catch { + // Broken symlink, skip it + continue; + } + } + + const relPath = toPosixPath(relative(root, fullPath)); + const ignorePath = isDirectory ? `${relPath}/` : relPath; + if (ig.ignores(ignorePath)) { + continue; + } + + if (isDirectory) { + const subResult = loadSkillsFromDirInternal(fullPath, source, false, ig, root); + skills.push(...subResult.skills); + diagnostics.push(...subResult.diagnostics); + continue; + } + + if (!isFile || !includeRootFiles || !entry.name.endsWith(".md")) { + continue; + } + + const result = loadSkillFromFile(fullPath, source); + if (result.skill) { + skills.push(result.skill); + } + diagnostics.push(...result.diagnostics); + } + } catch {} + + return { skills, diagnostics }; +} + +function loadSkillFromFile( + filePath: string, + source: string, +): { skill: Skill | null; diagnostics: ResourceDiagnostic[] } { + const diagnostics: ResourceDiagnostic[] = []; + + try { + const rawContent = readFileSync(filePath, "utf-8"); + const { frontmatter } = parseFrontmatter(rawContent); + const skillDir = dirname(filePath); + const parentDirName = basename(skillDir); + + // Validate description + const descErrors = validateDescription(frontmatter.description); + for (const error of descErrors) { + diagnostics.push({ type: "warning", message: error, path: filePath }); + } + + // Use name from frontmatter, or fall back to parent directory name + const name = frontmatter.name || parentDirName; + + // Validate name + const nameErrors = validateName(name); + for (const error of nameErrors) { + diagnostics.push({ type: "warning", message: error, path: filePath }); + } + + // Still load the skill even with warnings (unless description is completely missing) + if (!frontmatter.description || frontmatter.description.trim() === "") { + return { skill: null, diagnostics }; + } + + return { + skill: { + name, + description: frontmatter.description, + filePath, + baseDir: skillDir, + sourceInfo: createSkillSourceInfo(filePath, skillDir, source), + disableModelInvocation: frontmatter["disable-model-invocation"] === true, + }, + diagnostics, + }; + } catch (error) { + const message = error instanceof Error ? error.message : "failed to parse skill file"; + diagnostics.push({ type: "warning", message, path: filePath }); + return { skill: null, diagnostics }; + } +} + +/** + * Format skills for inclusion in a system prompt. + * Uses XML format per Agent Skills standard. + * See: https://agentskills.io/integrate-skills + * + * Skills with disableModelInvocation=true are excluded from the prompt + * (they can only be invoked explicitly via /skill:name commands). + */ +export function formatSkillsForPrompt(skills: Skill[]): string { + const visibleSkills = skills.filter((s) => !s.disableModelInvocation); + + if (visibleSkills.length === 0) { + return ""; + } + + const lines = [ + "\n\nThe following skills provide specialized instructions for specific tasks.", + "Use the read tool to load a skill's file when the task matches its description.", + "When a skill file references a relative path, resolve it against the skill directory (parent of SKILL.md / dirname of the path) and use that absolute path in tool commands.", + "", + "", + ]; + + for (const skill of visibleSkills) { + lines.push(" "); + lines.push(` ${escapeXml(skill.name)}`); + lines.push(` ${escapeXml(skill.description)}`); + lines.push(` ${escapeXml(skill.filePath)}`); + lines.push(" "); + } + + lines.push(""); + + return lines.join("\n"); +} + +function escapeXml(str: string): string { + return str + .replace(/&/g, "&") + .replace(//g, ">") + .replace(/"/g, """) + .replace(/'/g, "'"); +} + +export interface LoadSkillsOptions { + /** Working directory for project-local skills. */ + cwd: string; + /** Agent config directory for global skills. */ + agentDir: string; + /** Explicit skill paths (files or directories) */ + skillPaths: string[]; + /** Include default skills directories. */ + includeDefaults: boolean; +} + +/** + * Load skills from all configured locations. + * Returns skills and any validation diagnostics. + */ +export function loadSkills(options: LoadSkillsOptions): LoadSkillsResult { + const { agentDir, skillPaths, includeDefaults } = options; + + // Resolve agentDir - if not provided, use default from config + const resolvedCwd = resolvePath(options.cwd); + const resolvedAgentDir = resolvePath(agentDir ?? getAgentDir()); + + const skillMap = new Map(); + const realPathSet = new Set(); + const allDiagnostics: ResourceDiagnostic[] = []; + const collisionDiagnostics: ResourceDiagnostic[] = []; + + function addSkills(result: LoadSkillsResult) { + allDiagnostics.push(...result.diagnostics); + for (const skill of result.skills) { + // Resolve symlinks to detect duplicate files + const realPath = canonicalizePath(skill.filePath); + + // Skip silently if we've already loaded this exact file (via symlink) + if (realPathSet.has(realPath)) { + continue; + } + + const existing = skillMap.get(skill.name); + if (existing) { + collisionDiagnostics.push({ + type: "collision", + message: `name "${skill.name}" collision`, + path: skill.filePath, + collision: { + resourceType: "skill", + name: skill.name, + winnerPath: existing.filePath, + loserPath: skill.filePath, + }, + }); + } else { + skillMap.set(skill.name, skill); + realPathSet.add(realPath); + } + } + } + + if (includeDefaults) { + addSkills(loadSkillsFromDirInternal(join(resolvedAgentDir, "skills"), "user", true)); + addSkills(loadSkillsFromDirInternal(resolve(resolvedCwd, CONFIG_DIR_NAME, "skills"), "project", true)); + } + + const userSkillsDir = join(resolvedAgentDir, "skills"); + const projectSkillsDir = resolve(resolvedCwd, CONFIG_DIR_NAME, "skills"); + + const isUnderPath = (target: string, root: string): boolean => { + const normalizedRoot = resolve(root); + if (target === normalizedRoot) { + return true; + } + const prefix = normalizedRoot.endsWith(sep) ? normalizedRoot : `${normalizedRoot}${sep}`; + return target.startsWith(prefix); + }; + + const getSource = (resolvedPath: string): "user" | "project" | "path" => { + if (!includeDefaults) { + if (isUnderPath(resolvedPath, userSkillsDir)) return "user"; + if (isUnderPath(resolvedPath, projectSkillsDir)) return "project"; + } + return "path"; + }; + + for (const rawPath of skillPaths) { + const resolvedPath = resolvePath(rawPath, resolvedCwd, { trim: true }); + if (!existsSync(resolvedPath)) { + allDiagnostics.push({ type: "warning", message: "skill path does not exist", path: resolvedPath }); + continue; + } + + try { + const stats = statSync(resolvedPath); + const source = getSource(resolvedPath); + if (stats.isDirectory()) { + addSkills(loadSkillsFromDirInternal(resolvedPath, source, true)); + } else if (stats.isFile() && resolvedPath.endsWith(".md")) { + const result = loadSkillFromFile(resolvedPath, source); + if (result.skill) { + addSkills({ skills: [result.skill], diagnostics: result.diagnostics }); + } else { + allDiagnostics.push(...result.diagnostics); + } + } else { + allDiagnostics.push({ type: "warning", message: "skill path is not a markdown file", path: resolvedPath }); + } + } catch (error) { + const message = error instanceof Error ? error.message : "failed to read skill path"; + allDiagnostics.push({ type: "warning", message, path: resolvedPath }); + } + } + + return { + skills: Array.from(skillMap.values()), + diagnostics: [...allDiagnostics, ...collisionDiagnostics], + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/slash-commands.ts b/cactus-code/packages/coding-agent/src/core/slash-commands.ts new file mode 100644 index 000000000..80b5dff9d --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/slash-commands.ts @@ -0,0 +1,37 @@ +import { APP_NAME } from "../config.ts"; +import type { SourceInfo } from "./source-info.ts"; + +export type SlashCommandSource = "extension" | "prompt" | "skill"; + +export interface SlashCommandInfo { + name: string; + description?: string; + source: SlashCommandSource; + sourceInfo: SourceInfo; +} + +export interface BuiltinSlashCommand { + name: string; + description: string; +} + +export const BUILTIN_SLASH_COMMANDS: ReadonlyArray = [ + { name: "settings", description: "Open settings menu" }, + { name: "model", description: "Select model (opens selector UI)" }, + { name: "scoped-models", description: "Enable/disable models for Ctrl+P cycling" }, + { name: "export", description: "Export the current session to a .jsonl file" }, + { name: "import", description: "Import and resume a session from a JSONL file" }, + { name: "copy", description: "Copy last agent message to clipboard" }, + { name: "name", description: "Set session display name" }, + { name: "session", description: "Show session info and stats" }, + { name: "hotkeys", description: "Show all keyboard shortcuts" }, + { name: "fork", description: "Create a new fork from a previous user message" }, + { name: "clone", description: "Duplicate the current session at the current position" }, + { name: "tree", description: "Navigate session tree (switch branches)" }, + { name: "trust", description: "Save project trust decision for future sessions" }, + { name: "new", description: "Start a new session" }, + { name: "compact", description: "Manually compact the session context" }, + { name: "resume", description: "Resume a different session" }, + { name: "reload", description: "Reload keybindings, extensions, skills, prompts, and themes" }, + { name: "quit", description: `Quit ${APP_NAME}` }, +]; diff --git a/cactus-code/packages/coding-agent/src/core/source-info.ts b/cactus-code/packages/coding-agent/src/core/source-info.ts new file mode 100644 index 000000000..c8c9837d1 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/source-info.ts @@ -0,0 +1,40 @@ +import type { PathMetadata } from "./package-manager.ts"; + +export type SourceScope = "user" | "project" | "temporary"; +export type SourceOrigin = "package" | "top-level"; + +export interface SourceInfo { + path: string; + source: string; + scope: SourceScope; + origin: SourceOrigin; + baseDir?: string; +} + +export function createSourceInfo(path: string, metadata: PathMetadata): SourceInfo { + return { + path, + source: metadata.source, + scope: metadata.scope, + origin: metadata.origin, + baseDir: metadata.baseDir, + }; +} + +export function createSyntheticSourceInfo( + path: string, + options: { + source: string; + scope?: SourceScope; + origin?: SourceOrigin; + baseDir?: string; + }, +): SourceInfo { + return { + path, + source: options.source, + scope: options.scope ?? "temporary", + origin: options.origin ?? "top-level", + baseDir: options.baseDir, + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/system-prompt.ts b/cactus-code/packages/coding-agent/src/core/system-prompt.ts new file mode 100644 index 000000000..bff8d3c0b --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/system-prompt.ts @@ -0,0 +1,158 @@ +/** + * System prompt construction and project context loading + */ + +import { formatSkillsForPrompt, type Skill } from "./skills.ts"; + +export interface BuildSystemPromptOptions { + /** Custom system prompt (replaces default). */ + customPrompt?: string; + /** Tools to include in prompt. Default: [read, bash, edit, write] */ + selectedTools?: string[]; + /** Optional one-line tool snippets keyed by tool name. */ + toolSnippets?: Record; + /** Additional guideline bullets appended to the default system prompt guidelines. */ + promptGuidelines?: string[]; + /** Text to append to system prompt. */ + appendSystemPrompt?: string; + /** Working directory. */ + cwd: string; + /** Pre-loaded context files. */ + contextFiles?: Array<{ path: string; content: string }>; + /** Pre-loaded skills. */ + skills?: Skill[]; +} + +/** Build the system prompt with tools, guidelines, and context */ +export function buildSystemPrompt(options: BuildSystemPromptOptions): string { + const { + customPrompt, + selectedTools, + toolSnippets, + promptGuidelines, + appendSystemPrompt, + cwd, + contextFiles: providedContextFiles, + skills: providedSkills, + } = options; + const resolvedCwd = cwd; + const promptCwd = resolvedCwd.replace(/\\/g, "/"); + + const now = new Date(); + const year = now.getFullYear(); + const month = String(now.getMonth() + 1).padStart(2, "0"); + const day = String(now.getDate()).padStart(2, "0"); + const date = `${year}-${month}-${day}`; + + const appendSection = appendSystemPrompt ? `\n\n${appendSystemPrompt}` : ""; + + const contextFiles = providedContextFiles ?? []; + const skills = providedSkills ?? []; + + if (customPrompt) { + let prompt = customPrompt; + + if (appendSection) { + prompt += appendSection; + } + + // Append project context files + if (contextFiles.length > 0) { + prompt += "\n\n\n\n"; + prompt += "Project-specific instructions and guidelines:\n\n"; + for (const { path: filePath, content } of contextFiles) { + prompt += `\n${content}\n\n\n`; + } + prompt += "\n"; + } + + // Append skills section (only if read tool is available) + const customPromptHasRead = !selectedTools || selectedTools.includes("read"); + if (customPromptHasRead && skills.length > 0) { + prompt += formatSkillsForPrompt(skills); + } + + // Add date and working directory last + prompt += `\nCurrent date: ${date}`; + prompt += `\nCurrent working directory: ${promptCwd}`; + + return prompt; + } + + // Build tools list based on selected tools. + // A tool appears in Available tools only when the caller provides a one-line snippet. + const tools = selectedTools || ["read", "bash", "edit", "write"]; + const visibleTools = tools.filter((name) => !!toolSnippets?.[name]); + const toolsList = + visibleTools.length > 0 ? visibleTools.map((name) => `- ${name}: ${toolSnippets![name]}`).join("\n") : "(none)"; + + // Build guidelines based on which tools are actually available + const guidelinesList: string[] = []; + const guidelinesSet = new Set(); + const addGuideline = (guideline: string): void => { + if (guidelinesSet.has(guideline)) { + return; + } + guidelinesSet.add(guideline); + guidelinesList.push(guideline); + }; + + const hasBash = tools.includes("bash"); + const hasGrep = tools.includes("grep"); + const hasFind = tools.includes("find"); + const hasLs = tools.includes("ls"); + const hasRead = tools.includes("read"); + + // File exploration guidelines + if (hasBash && !hasGrep && !hasFind && !hasLs) { + addGuideline("Use bash for file operations like ls, rg, find"); + } + + for (const guideline of promptGuidelines ?? []) { + const normalized = guideline.trim(); + if (normalized.length > 0) { + addGuideline(normalized); + } + } + + // Always include these + addGuideline("Be concise in your responses"); + addGuideline("Show file paths clearly when working with files"); + + const guidelines = guidelinesList.map((g) => `- ${g}`).join("\n"); + + let prompt = `You are an expert coding assistant operating inside cactus, a coding agent harness. You help users by reading files, executing commands, editing code, and writing new files. + +Available tools: +${toolsList} + +In addition to the tools above, you may have access to other custom tools depending on the project. + +Guidelines: +${guidelines}`; + + if (appendSection) { + prompt += appendSection; + } + + // Append project context files + if (contextFiles.length > 0) { + prompt += "\n\n\n\n"; + prompt += "Project-specific instructions and guidelines:\n\n"; + for (const { path: filePath, content } of contextFiles) { + prompt += `\n${content}\n\n\n`; + } + prompt += "\n"; + } + + // Append skills section (only if read tool is available) + if (hasRead && skills.length > 0) { + prompt += formatSkillsForPrompt(skills); + } + + // Add date and working directory last + prompt += `\nCurrent date: ${date}`; + prompt += `\nCurrent working directory: ${promptCwd}`; + + return prompt; +} diff --git a/cactus-code/packages/coding-agent/src/core/timings.ts b/cactus-code/packages/coding-agent/src/core/timings.ts new file mode 100644 index 000000000..0b3e5cf48 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/timings.ts @@ -0,0 +1,31 @@ +/** + * Central timing instrumentation for startup profiling. + * Enable with PI_TIMING=1 environment variable. + */ + +const ENABLED = process.env.PI_TIMING === "1"; +const timings: Array<{ label: string; ms: number }> = []; +let lastTime = Date.now(); + +export function resetTimings(): void { + if (!ENABLED) return; + timings.length = 0; + lastTime = Date.now(); +} + +export function time(label: string): void { + if (!ENABLED) return; + const now = Date.now(); + timings.push({ label, ms: now - lastTime }); + lastTime = now; +} + +export function printTimings(): void { + if (!ENABLED || timings.length === 0) return; + console.error("\n--- Startup Timings ---"); + for (const t of timings) { + console.error(` ${t.label}: ${t.ms}ms`); + } + console.error(` TOTAL: ${timings.reduce((a, b) => a + b.ms, 0)}ms`); + console.error("------------------------\n"); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/bash.ts b/cactus-code/packages/coding-agent/src/core/tools/bash.ts new file mode 100644 index 000000000..5e29cccbb --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/bash.ts @@ -0,0 +1,453 @@ +import { constants } from "node:fs"; +import { access as fsAccess } from "node:fs/promises"; +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { Container, Text, truncateToWidth } from "@earendil-works/pi-tui"; +import { spawn } from "child_process"; +import { type Static, Type } from "typebox"; +import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts"; +import { truncateToVisualLines } from "../../modes/interactive/components/visual-truncate.ts"; +import { theme } from "../../modes/interactive/theme/theme.ts"; +import { waitForChildProcess } from "../../utils/child-process.ts"; +import { + getShellConfig, + getShellEnv, + killProcessTree, + trackDetachedChildPid, + untrackDetachedChildPid, +} from "../../utils/shell.ts"; +import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts"; +import { OutputAccumulator } from "./output-accumulator.ts"; +import { getTextOutput, invalidArgText, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; +import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult } from "./truncate.ts"; + +const bashSchema = Type.Object({ + command: Type.String({ description: "Bash command to execute" }), + timeout: Type.Optional(Type.Number({ description: "Timeout in seconds (optional, no default timeout)" })), +}); + +export type BashToolInput = Static; + +export interface BashToolDetails { + truncation?: TruncationResult; + fullOutputPath?: string; +} + +/** + * Pluggable operations for the bash tool. + * Override these to delegate command execution to remote systems (for example SSH). + */ +export interface BashOperations { + /** + * Execute a command and stream output. + * @param command The command to execute + * @param cwd Working directory + * @param options Execution options + * @returns Promise resolving to exit code (null if killed) + */ + exec: ( + command: string, + cwd: string, + options: { + onData: (data: Buffer) => void; + signal?: AbortSignal; + timeout?: number; + env?: NodeJS.ProcessEnv; + }, + ) => Promise<{ exitCode: number | null }>; +} + +/** + * Create bash operations using pi's built-in local shell execution backend. + * + * This is useful for extensions that intercept user_bash and still want pi's + * standard local shell behavior while wrapping or rewriting commands. + */ +export function createLocalBashOperations(options?: { shellPath?: string }): BashOperations { + return { + exec: async (command, cwd, { onData, signal, timeout, env }) => { + const shellConfig = getShellConfig(options?.shellPath); + try { + await fsAccess(cwd, constants.F_OK); + } catch { + throw new Error(`Working directory does not exist: ${cwd}\nCannot execute bash commands.`); + } + if (signal?.aborted) { + throw new Error("aborted"); + } + + const commandFromStdin = shellConfig.commandTransport === "stdin"; + const child = spawn(shellConfig.shell, commandFromStdin ? shellConfig.args : [...shellConfig.args, command], { + cwd, + detached: process.platform !== "win32", + env: env ?? getShellEnv(), + stdio: [commandFromStdin ? "pipe" : "ignore", "pipe", "pipe"], + windowsHide: true, + }); + if (commandFromStdin) { + child.stdin?.on("error", () => {}); + child.stdin?.end(command); + } + if (child.pid) trackDetachedChildPid(child.pid); + let timedOut = false; + let timeoutHandle: NodeJS.Timeout | undefined; + const onAbort = () => { + if (child.pid) killProcessTree(child.pid); + }; + + try { + // Set timeout if provided. + if (timeout !== undefined && timeout > 0) { + timeoutHandle = setTimeout(() => { + timedOut = true; + if (child.pid) killProcessTree(child.pid); + }, timeout * 1000); + } + // Stream stdout and stderr. + child.stdout?.on("data", onData); + child.stderr?.on("data", onData); + // Handle abort signal by killing the entire process tree. + if (signal) { + if (signal.aborted) onAbort(); + else signal.addEventListener("abort", onAbort, { once: true }); + } + // Handle shell spawn errors and wait for the process to terminate without hanging + // on inherited stdio handles held by detached descendants. + const exitCode = await waitForChildProcess(child); + if (signal?.aborted) { + throw new Error("aborted"); + } + if (timedOut) { + throw new Error(`timeout:${timeout}`); + } + return { exitCode }; + } finally { + if (child.pid) untrackDetachedChildPid(child.pid); + if (timeoutHandle) clearTimeout(timeoutHandle); + if (signal) signal.removeEventListener("abort", onAbort); + } + }, + }; +} + +export interface BashSpawnContext { + command: string; + cwd: string; + env: NodeJS.ProcessEnv; +} + +export type BashSpawnHook = (context: BashSpawnContext) => BashSpawnContext; + +function resolveSpawnContext(command: string, cwd: string, spawnHook?: BashSpawnHook): BashSpawnContext { + const baseContext: BashSpawnContext = { command, cwd, env: { ...getShellEnv() } }; + return spawnHook ? spawnHook(baseContext) : baseContext; +} + +export interface BashToolOptions { + /** Custom operations for command execution. Default: local shell */ + operations?: BashOperations; + /** Command prefix prepended to every command (for example shell setup commands) */ + commandPrefix?: string; + /** Optional explicit shell path from settings */ + shellPath?: string; + /** Hook to adjust command, cwd, or env before execution */ + spawnHook?: BashSpawnHook; +} + +const BASH_PREVIEW_LINES = 5; +const BASH_UPDATE_THROTTLE_MS = 100; + +type BashRenderState = { + startedAt: number | undefined; + endedAt: number | undefined; + interval: NodeJS.Timeout | undefined; +}; + +type BashResultRenderState = { + cachedWidth: number | undefined; + cachedLines: string[] | undefined; + cachedSkipped: number | undefined; +}; + +class BashResultRenderComponent extends Container { + state: BashResultRenderState = { + cachedWidth: undefined, + cachedLines: undefined, + cachedSkipped: undefined, + }; +} + +function formatDuration(ms: number): string { + return `${(ms / 1000).toFixed(1)}s`; +} + +function formatBashCall(args: { command?: string; timeout?: number } | undefined): string { + const command = str(args?.command); + const timeout = args?.timeout as number | undefined; + const timeoutSuffix = timeout ? theme.fg("muted", ` (timeout ${timeout}s)`) : ""; + const commandDisplay = command === null ? invalidArgText(theme) : command ? command : theme.fg("toolOutput", "..."); + return theme.fg("toolTitle", theme.bold(`$ ${commandDisplay}`)) + timeoutSuffix; +} + +function rebuildBashResultRenderComponent( + component: BashResultRenderComponent, + result: { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + details?: BashToolDetails; + }, + options: ToolRenderResultOptions, + showImages: boolean, + startedAt: number | undefined, + endedAt: number | undefined, +): void { + const state = component.state; + component.clear(); + + let output = getTextOutput(result as any, showImages).trim(); + const truncation = result.details?.truncation; + const fullOutputPath = result.details?.fullOutputPath; + if (!options.isPartial && truncation?.truncated && fullOutputPath && output.endsWith("]")) { + const footerStart = output.lastIndexOf("\n\n["); + if (footerStart !== -1 && output.slice(footerStart).includes(fullOutputPath)) { + output = output.slice(0, footerStart).trimEnd(); + } + } + + if (output) { + const styledOutput = output + .split("\n") + .map((line) => theme.fg("toolOutput", line)) + .join("\n"); + + if (options.expanded) { + component.addChild(new Text(`\n${styledOutput}`, 0, 0)); + } else { + component.addChild({ + render: (width: number) => { + if (state.cachedLines === undefined || state.cachedWidth !== width) { + const preview = truncateToVisualLines(styledOutput, BASH_PREVIEW_LINES, width); + state.cachedLines = preview.visualLines; + state.cachedSkipped = preview.skippedCount; + state.cachedWidth = width; + } + if (state.cachedSkipped && state.cachedSkipped > 0) { + const hint = + theme.fg("muted", `... (${state.cachedSkipped} earlier lines,`) + + ` ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`; + return ["", truncateToWidth(hint, width, "..."), ...(state.cachedLines ?? [])]; + } + return ["", ...(state.cachedLines ?? [])]; + }, + invalidate: () => { + state.cachedWidth = undefined; + state.cachedLines = undefined; + state.cachedSkipped = undefined; + }, + }); + } + } + + if (truncation?.truncated || fullOutputPath) { + const warnings: string[] = []; + if (fullOutputPath) { + warnings.push(`Full output: ${fullOutputPath}`); + } + if (truncation?.truncated) { + if (truncation.truncatedBy === "lines") { + warnings.push(`Truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines`); + } else { + warnings.push( + `Truncated: ${truncation.outputLines} lines shown (${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit)`, + ); + } + } + component.addChild(new Text(`\n${theme.fg("warning", `[${warnings.join(". ")}]`)}`, 0, 0)); + } + + if (startedAt !== undefined) { + const label = options.isPartial ? "Elapsed" : "Took"; + const endTime = endedAt ?? Date.now(); + component.addChild(new Text(`\n${theme.fg("muted", `${label} ${formatDuration(endTime - startedAt)}`)}`, 0, 0)); + } +} + +export function createBashToolDefinition( + cwd: string, + options?: BashToolOptions, +): ToolDefinition { + const ops = options?.operations ?? createLocalBashOperations({ shellPath: options?.shellPath }); + const commandPrefix = options?.commandPrefix; + const spawnHook = options?.spawnHook; + return { + name: "bash", + label: "bash", + description: `Execute a bash command in the current working directory. Returns stdout and stderr. Output is truncated to last ${DEFAULT_MAX_LINES} lines or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). If truncated, full output is saved to a temp file. Optionally provide a timeout in seconds.`, + promptSnippet: "Execute bash commands (ls, grep, find, etc.)", + parameters: bashSchema, + async execute( + _toolCallId, + { command, timeout }: { command: string; timeout?: number }, + signal?: AbortSignal, + onUpdate?, + _ctx?, + ) { + const resolvedCommand = commandPrefix ? `${commandPrefix}\n${command}` : command; + const spawnContext = resolveSpawnContext(resolvedCommand, cwd, spawnHook); + const output = new OutputAccumulator({ tempFilePrefix: "cactus-bash" }); + let acceptingOutput = true; + let updateTimer: NodeJS.Timeout | undefined; + let updateDirty = false; + let lastUpdateAt = 0; + + const emitOutputUpdate = () => { + if (!onUpdate || !updateDirty) return; + updateDirty = false; + lastUpdateAt = Date.now(); + const snapshot = output.snapshot({ persistIfTruncated: true }); + onUpdate({ + content: [{ type: "text", text: snapshot.content || "" }], + details: { + truncation: snapshot.truncation.truncated ? snapshot.truncation : undefined, + fullOutputPath: snapshot.fullOutputPath, + }, + }); + }; + + const clearUpdateTimer = () => { + if (updateTimer) { + clearTimeout(updateTimer); + updateTimer = undefined; + } + }; + + const scheduleOutputUpdate = () => { + if (!onUpdate) return; + updateDirty = true; + const delay = BASH_UPDATE_THROTTLE_MS - (Date.now() - lastUpdateAt); + if (delay <= 0) { + clearUpdateTimer(); + emitOutputUpdate(); + return; + } + updateTimer ??= setTimeout(() => { + updateTimer = undefined; + emitOutputUpdate(); + }, delay); + }; + + if (onUpdate) { + onUpdate({ content: [], details: undefined }); + } + + const handleData = (data: Buffer) => { + if (!acceptingOutput) return; + output.append(data); + scheduleOutputUpdate(); + }; + + const finishOutput = async () => { + acceptingOutput = false; + output.finish(); + clearUpdateTimer(); + emitOutputUpdate(); + const snapshot = output.snapshot({ persistIfTruncated: true }); + await output.closeTempFile(); + return snapshot; + }; + + const formatOutput = (snapshot: Awaited>, emptyText = "(no output)") => { + const truncation = snapshot.truncation; + let text = snapshot.content || emptyText; + let details: BashToolDetails | undefined; + if (truncation.truncated) { + details = { truncation, fullOutputPath: snapshot.fullOutputPath }; + const startLine = truncation.totalLines - truncation.outputLines + 1; + const endLine = truncation.totalLines; + if (truncation.lastLinePartial) { + const lastLineSize = formatSize(output.getLastLineBytes()); + text += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${snapshot.fullOutputPath}]`; + } else if (truncation.truncatedBy === "lines") { + text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${snapshot.fullOutputPath}]`; + } else { + text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${snapshot.fullOutputPath}]`; + } + } + return { text, details }; + }; + + const appendStatus = (text: string, status: string) => `${text ? `${text}\n\n` : ""}${status}`; + + try { + let exitCode: number | null; + try { + const result = await ops.exec(spawnContext.command, spawnContext.cwd, { + onData: handleData, + signal, + timeout, + env: spawnContext.env, + }); + exitCode = result.exitCode; + } catch (err) { + const snapshot = await finishOutput(); + const { text } = formatOutput(snapshot, ""); + if (err instanceof Error && err.message === "aborted") { + throw new Error(appendStatus(text, "Command aborted")); + } + if (err instanceof Error && err.message.startsWith("timeout:")) { + const timeoutSecs = err.message.split(":")[1]; + throw new Error(appendStatus(text, `Command timed out after ${timeoutSecs} seconds`)); + } + throw err; + } + + const snapshot = await finishOutput(); + const { text: outputText, details } = formatOutput(snapshot); + if (exitCode !== 0 && exitCode !== null) { + throw new Error(appendStatus(outputText, `Command exited with code ${exitCode}`)); + } + return { content: [{ type: "text", text: outputText }], details }; + } finally { + clearUpdateTimer(); + } + }, + renderCall(args, _theme, context) { + const state = context.state; + if (context.executionStarted && state.startedAt === undefined) { + state.startedAt = Date.now(); + state.endedAt = undefined; + } + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatBashCall(args)); + return text; + }, + renderResult(result, options, _theme, context) { + const state = context.state; + if (state.startedAt !== undefined && options.isPartial && !state.interval) { + state.interval = setInterval(() => context.invalidate(), 1000); + } + if (!options.isPartial || context.isError) { + state.endedAt ??= Date.now(); + if (state.interval) { + clearInterval(state.interval); + state.interval = undefined; + } + } + const component = + (context.lastComponent as BashResultRenderComponent | undefined) ?? new BashResultRenderComponent(); + rebuildBashResultRenderComponent( + component, + result as any, + options, + context.showImages, + state.startedAt, + state.endedAt, + ); + component.invalidate(); + return component; + }, + }; +} + +export function createBashTool(cwd: string, options?: BashToolOptions): AgentTool { + return wrapToolDefinition(createBashToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/edit-diff.ts b/cactus-code/packages/coding-agent/src/core/tools/edit-diff.ts new file mode 100644 index 000000000..5a4d966b0 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/edit-diff.ts @@ -0,0 +1,560 @@ +/** + * Shared diff computation utilities for the edit and similar tools. + */ + +import * as Diff from "diff"; +import { constants } from "fs"; +import { access, readFile } from "fs/promises"; +import { resolveToCwd } from "./path-utils.ts"; + +export function detectLineEnding(content: string): "\r\n" | "\n" { + const crlfIdx = content.indexOf("\r\n"); + const lfIdx = content.indexOf("\n"); + if (lfIdx === -1) return "\n"; + if (crlfIdx === -1) return "\n"; + return crlfIdx < lfIdx ? "\r\n" : "\n"; +} + +export function normalizeToLF(text: string): string { + return text.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); +} + +export function restoreLineEndings(text: string, ending: "\r\n" | "\n"): string { + return ending === "\r\n" ? text.replace(/\n/g, "\r\n") : text; +} + +/** + * Normalize text for fuzzy matching. Applies progressive transformations: + * - Strip trailing whitespace from each line + * - Normalize smart quotes to ASCII equivalents + * - Normalize Unicode dashes/hyphens to ASCII hyphen + * - Normalize special Unicode spaces to regular space + */ +export function normalizeForFuzzyMatch(text: string): string { + return ( + text + .normalize("NFKC") + // Strip trailing whitespace per line + .split("\n") + .map((line) => line.trimEnd()) + .join("\n") + // Smart single quotes → ' + .replace(/[\u2018\u2019\u201A\u201B]/g, "'") + // Smart double quotes → " + .replace(/[\u201C\u201D\u201E\u201F]/g, '"') + // Various dashes/hyphens → - + // U+2010 hyphen, U+2011 non-breaking hyphen, U+2012 figure dash, + // U+2013 en-dash, U+2014 em-dash, U+2015 horizontal bar, U+2212 minus + .replace(/[\u2010\u2011\u2012\u2013\u2014\u2015\u2212]/g, "-") + // Special spaces → regular space + // U+00A0 NBSP, U+2002-U+200A various spaces, U+202F narrow NBSP, + // U+205F medium math space, U+3000 ideographic space + .replace(/[\u00A0\u2002-\u200A\u202F\u205F\u3000]/g, " ") + ); +} + +function splitLinesWithEndings(content: string): string[] { + return content.match(/[^\n]*\n|[^\n]+/g) ?? []; +} + +interface LineSpan { + start: number; + end: number; +} + +interface MatchedEdit { + editIndex: number; + matchIndex: number; + matchLength: number; + newText: string; +} + +type TextReplacement = Pick; + +function getLineSpans(content: string): LineSpan[] { + let offset = 0; + return splitLinesWithEndings(content).map((line) => { + const span = { start: offset, end: offset + line.length }; + offset = span.end; + return span; + }); +} + +function getReplacementLineRange(lines: LineSpan[], replacement: TextReplacement) { + const replacementStart = replacement.matchIndex; + const replacementEnd = replacement.matchIndex + replacement.matchLength; + + let startLine = -1; + for (let i = 0; i < lines.length; i++) { + const line = lines[i]; + if (replacementStart >= line.start && replacementStart < line.end) { + startLine = i; + break; + } + } + if (startLine === -1) { + throw new Error("Replacement range is outside the base content."); + } + + let endLine = startLine; + while (endLine < lines.length && lines[endLine].end < replacementEnd) { + endLine++; + } + if (endLine >= lines.length) { + throw new Error("Replacement range is outside the base content."); + } + + return { startLine, endLine: endLine + 1 }; +} + +function applyReplacements(content: string, replacements: TextReplacement[], offset = 0): string { + let result = content; + for (let i = replacements.length - 1; i >= 0; i--) { + const replacement = replacements[i]; + const matchIndex = replacement.matchIndex - offset; + result = + result.substring(0, matchIndex) + replacement.newText + result.substring(matchIndex + replacement.matchLength); + } + return result; +} + +/** + * Apply replacements matched against `baseContent` to `originalContent` while + * preserving unchanged line blocks from the original. + * + * This is useful when `baseContent` is a normalized view of the original. Each + * replacement is widened to the lines it actually touches, those touched lines + * are rewritten from the normalized base, and all other lines are copied back + * from `originalContent`. The actual replacement ranges drive preservation so + * duplicate normalized lines cannot be aligned to the wrong occurrence. + */ +export function applyReplacementsPreservingUnchangedLines( + originalContent: string, + baseContent: string, + replacements: TextReplacement[], +): string { + const originalLines = splitLinesWithEndings(originalContent); + const baseLines = getLineSpans(baseContent); + if (originalLines.length !== baseLines.length) { + throw new Error("Cannot preserve unchanged lines because the base content has a different line count."); + } + + const groups: Array<{ startLine: number; endLine: number; replacements: TextReplacement[] }> = []; + const sortedReplacements = [...replacements].sort((a, b) => a.matchIndex - b.matchIndex); + for (const replacement of sortedReplacements) { + const range = getReplacementLineRange(baseLines, replacement); + const current = groups[groups.length - 1]; + if (current && range.startLine < current.endLine) { + current.endLine = Math.max(current.endLine, range.endLine); + current.replacements.push(replacement); + continue; + } + groups.push({ ...range, replacements: [replacement] }); + } + + let originalLineIndex = 0; + let result = ""; + for (const group of groups) { + result += originalLines.slice(originalLineIndex, group.startLine).join(""); + + const groupStartOffset = baseLines[group.startLine].start; + const groupEndOffset = baseLines[group.endLine - 1].end; + result += applyReplacements( + baseContent.slice(groupStartOffset, groupEndOffset), + group.replacements, + groupStartOffset, + ); + originalLineIndex = group.endLine; + } + result += originalLines.slice(originalLineIndex).join(""); + + return result; +} + +export interface FuzzyMatchResult { + /** Whether a match was found */ + found: boolean; + /** The index where the match starts (in the content that should be used for replacement) */ + index: number; + /** Length of the matched text */ + matchLength: number; + /** Whether fuzzy matching was used (false = exact match) */ + usedFuzzyMatch: boolean; + /** + * The content to use for replacement operations. + * When exact match: original content. When fuzzy match: normalized content. + */ + contentForReplacement: string; +} + +export interface Edit { + oldText: string; + newText: string; +} + +export interface AppliedEditsResult { + baseContent: string; + newContent: string; +} + +/** + * Find oldText in content, trying exact match first, then fuzzy match. + * When fuzzy matching is used, the returned contentForReplacement is the + * fuzzy-normalized version of the content (trailing whitespace stripped, + * Unicode quotes/dashes normalized to ASCII). + */ +export function fuzzyFindText(content: string, oldText: string): FuzzyMatchResult { + // Try exact match first + const exactIndex = content.indexOf(oldText); + if (exactIndex !== -1) { + return { + found: true, + index: exactIndex, + matchLength: oldText.length, + usedFuzzyMatch: false, + contentForReplacement: content, + }; + } + + // Try fuzzy match - work entirely in normalized space + const fuzzyContent = normalizeForFuzzyMatch(content); + const fuzzyOldText = normalizeForFuzzyMatch(oldText); + const fuzzyIndex = fuzzyContent.indexOf(fuzzyOldText); + + if (fuzzyIndex === -1) { + return { + found: false, + index: -1, + matchLength: 0, + usedFuzzyMatch: false, + contentForReplacement: content, + }; + } + + // When fuzzy matching, return offsets in normalized space. Callers can use + // the normalized content to compute replacements, then decide how much of + // that normalized output should be written back. + return { + found: true, + index: fuzzyIndex, + matchLength: fuzzyOldText.length, + usedFuzzyMatch: true, + contentForReplacement: fuzzyContent, + }; +} + +/** Strip UTF-8 BOM if present, return both the BOM (if any) and the text without it */ +export function stripBom(content: string): { bom: string; text: string } { + return content.startsWith("\uFEFF") ? { bom: "\uFEFF", text: content.slice(1) } : { bom: "", text: content }; +} + +function countOccurrences(content: string, oldText: string): number { + const fuzzyContent = normalizeForFuzzyMatch(content); + const fuzzyOldText = normalizeForFuzzyMatch(oldText); + return fuzzyContent.split(fuzzyOldText).length - 1; +} + +function getNotFoundError(path: string, editIndex: number, totalEdits: number): Error { + if (totalEdits === 1) { + return new Error( + `Could not find the exact text in ${path}. The old text must match exactly including all whitespace and newlines.`, + ); + } + return new Error( + `Could not find edits[${editIndex}] in ${path}. The oldText must match exactly including all whitespace and newlines.`, + ); +} + +function getDuplicateError(path: string, editIndex: number, totalEdits: number, occurrences: number): Error { + if (totalEdits === 1) { + return new Error( + `Found ${occurrences} occurrences of the text in ${path}. The text must be unique. Please provide more context to make it unique.`, + ); + } + return new Error( + `Found ${occurrences} occurrences of edits[${editIndex}] in ${path}. Each oldText must be unique. Please provide more context to make it unique.`, + ); +} + +function getEmptyOldTextError(path: string, editIndex: number, totalEdits: number): Error { + if (totalEdits === 1) { + return new Error(`oldText must not be empty in ${path}.`); + } + return new Error(`edits[${editIndex}].oldText must not be empty in ${path}.`); +} + +function getNoChangeError(path: string, totalEdits: number): Error { + if (totalEdits === 1) { + return new Error( + `No changes made to ${path}. The replacement produced identical content. This might indicate an issue with special characters or the text not existing as expected.`, + ); + } + return new Error(`No changes made to ${path}. The replacements produced identical content.`); +} + +/** + * Apply one or more exact-text replacements to LF-normalized content. + * + * All edits are matched against the same original content. Replacements are + * then applied in reverse order so offsets remain stable. If any edit needs + * fuzzy matching, the operation runs in fuzzy-normalized content space and then + * overlays those line-level changes onto the original content so unchanged line + * blocks keep their original bytes. + */ +export function applyEditsToNormalizedContent( + normalizedContent: string, + edits: Edit[], + path: string, +): AppliedEditsResult { + const normalizedEdits = edits.map((edit) => ({ + oldText: normalizeToLF(edit.oldText), + newText: normalizeToLF(edit.newText), + })); + + for (let i = 0; i < normalizedEdits.length; i++) { + if (normalizedEdits[i].oldText.length === 0) { + throw getEmptyOldTextError(path, i, normalizedEdits.length); + } + } + + const initialMatches = normalizedEdits.map((edit) => fuzzyFindText(normalizedContent, edit.oldText)); + const usedFuzzyMatch = initialMatches.some((match) => match.usedFuzzyMatch); + const replacementBaseContent = usedFuzzyMatch ? normalizeForFuzzyMatch(normalizedContent) : normalizedContent; + + const matchedEdits: MatchedEdit[] = []; + for (let i = 0; i < normalizedEdits.length; i++) { + const edit = normalizedEdits[i]; + const matchResult = fuzzyFindText(replacementBaseContent, edit.oldText); + if (!matchResult.found) { + throw getNotFoundError(path, i, normalizedEdits.length); + } + + const occurrences = countOccurrences(replacementBaseContent, edit.oldText); + if (occurrences > 1) { + throw getDuplicateError(path, i, normalizedEdits.length, occurrences); + } + + matchedEdits.push({ + editIndex: i, + matchIndex: matchResult.index, + matchLength: matchResult.matchLength, + newText: edit.newText, + }); + } + + matchedEdits.sort((a, b) => a.matchIndex - b.matchIndex); + for (let i = 1; i < matchedEdits.length; i++) { + const previous = matchedEdits[i - 1]; + const current = matchedEdits[i]; + if (previous.matchIndex + previous.matchLength > current.matchIndex) { + throw new Error( + `edits[${previous.editIndex}] and edits[${current.editIndex}] overlap in ${path}. Merge them into one edit or target disjoint regions.`, + ); + } + } + + const baseContent = normalizedContent; + const newContent = usedFuzzyMatch + ? applyReplacementsPreservingUnchangedLines(normalizedContent, replacementBaseContent, matchedEdits) + : applyReplacements(replacementBaseContent, matchedEdits); + + if (baseContent === newContent) { + throw getNoChangeError(path, normalizedEdits.length); + } + + return { baseContent, newContent }; +} + +/** Generate a standard unified patch. */ +export function generateUnifiedPatch(path: string, oldContent: string, newContent: string, contextLines = 4): string { + return Diff.createTwoFilesPatch(path, path, oldContent, newContent, undefined, undefined, { + context: contextLines, + headerOptions: Diff.FILE_HEADERS_ONLY, + }); +} + +/** + * Generate a display-oriented diff string with line numbers and context. + * Returns both the diff string and the first changed line number (in the new file). + */ +export function generateDiffString( + oldContent: string, + newContent: string, + contextLines = 4, +): { diff: string; firstChangedLine: number | undefined } { + const parts = Diff.diffLines(oldContent, newContent); + const output: string[] = []; + + const oldLines = oldContent.split("\n"); + const newLines = newContent.split("\n"); + const maxLineNum = Math.max(oldLines.length, newLines.length); + const lineNumWidth = String(maxLineNum).length; + + let oldLineNum = 1; + let newLineNum = 1; + let lastWasChange = false; + let firstChangedLine: number | undefined; + + for (let i = 0; i < parts.length; i++) { + const part = parts[i]; + const raw = part.value.split("\n"); + if (raw[raw.length - 1] === "") { + raw.pop(); + } + + if (part.added || part.removed) { + // Capture the first changed line (in the new file) + if (firstChangedLine === undefined) { + firstChangedLine = newLineNum; + } + + // Show the change + for (const line of raw) { + if (part.added) { + const lineNum = String(newLineNum).padStart(lineNumWidth, " "); + output.push(`+${lineNum} ${line}`); + newLineNum++; + } else { + // removed + const lineNum = String(oldLineNum).padStart(lineNumWidth, " "); + output.push(`-${lineNum} ${line}`); + oldLineNum++; + } + } + lastWasChange = true; + } else { + // Context lines - only show a few before/after changes + const nextPartIsChange = i < parts.length - 1 && (parts[i + 1].added || parts[i + 1].removed); + const hasLeadingChange = lastWasChange; + const hasTrailingChange = nextPartIsChange; + + if (hasLeadingChange && hasTrailingChange) { + if (raw.length <= contextLines * 2) { + for (const line of raw) { + const lineNum = String(oldLineNum).padStart(lineNumWidth, " "); + output.push(` ${lineNum} ${line}`); + oldLineNum++; + newLineNum++; + } + } else { + const leadingLines = raw.slice(0, contextLines); + const trailingLines = raw.slice(raw.length - contextLines); + const skippedLines = raw.length - leadingLines.length - trailingLines.length; + + for (const line of leadingLines) { + const lineNum = String(oldLineNum).padStart(lineNumWidth, " "); + output.push(` ${lineNum} ${line}`); + oldLineNum++; + newLineNum++; + } + + output.push(` ${"".padStart(lineNumWidth, " ")} ...`); + oldLineNum += skippedLines; + newLineNum += skippedLines; + + for (const line of trailingLines) { + const lineNum = String(oldLineNum).padStart(lineNumWidth, " "); + output.push(` ${lineNum} ${line}`); + oldLineNum++; + newLineNum++; + } + } + } else if (hasLeadingChange) { + const shownLines = raw.slice(0, contextLines); + const skippedLines = raw.length - shownLines.length; + + for (const line of shownLines) { + const lineNum = String(oldLineNum).padStart(lineNumWidth, " "); + output.push(` ${lineNum} ${line}`); + oldLineNum++; + newLineNum++; + } + + if (skippedLines > 0) { + output.push(` ${"".padStart(lineNumWidth, " ")} ...`); + oldLineNum += skippedLines; + newLineNum += skippedLines; + } + } else if (hasTrailingChange) { + const skippedLines = Math.max(0, raw.length - contextLines); + if (skippedLines > 0) { + output.push(` ${"".padStart(lineNumWidth, " ")} ...`); + oldLineNum += skippedLines; + newLineNum += skippedLines; + } + + for (const line of raw.slice(skippedLines)) { + const lineNum = String(oldLineNum).padStart(lineNumWidth, " "); + output.push(` ${lineNum} ${line}`); + oldLineNum++; + newLineNum++; + } + } else { + // Skip these context lines entirely + oldLineNum += raw.length; + newLineNum += raw.length; + } + + lastWasChange = false; + } + } + + return { diff: output.join("\n"), firstChangedLine }; +} + +export interface EditDiffResult { + diff: string; + firstChangedLine: number | undefined; +} + +export interface EditDiffError { + error: string; +} + +/** + * Compute the diff for one or more edit operations without applying them. + * Used for preview rendering in the TUI before the tool executes. + */ +export async function computeEditsDiff( + path: string, + edits: Edit[], + cwd: string, +): Promise { + const absolutePath = resolveToCwd(path, cwd); + + try { + // Check if file exists and is readable + try { + await access(absolutePath, constants.R_OK); + } catch (error: unknown) { + const errorMessage = error instanceof Error && "code" in error ? `Error code: ${error.code}` : String(error); + return { error: `Could not edit file: ${path}. ${errorMessage}.` }; + } + + // Read the file + const rawContent = await readFile(absolutePath, "utf-8"); + + // Strip BOM before matching (LLM won't include invisible BOM in oldText) + const { text: content } = stripBom(rawContent); + const normalizedContent = normalizeToLF(content); + const { baseContent, newContent } = applyEditsToNormalizedContent(normalizedContent, edits, path); + + // Generate the diff + return generateDiffString(baseContent, newContent); + } catch (err) { + return { error: err instanceof Error ? err.message : String(err) }; + } +} + +/** + * Compute the diff for a single edit operation without applying it. + * Kept as a convenience wrapper for single-edit callers. + */ +export async function computeEditDiff( + path: string, + oldText: string, + newText: string, + cwd: string, +): Promise { + return computeEditsDiff(path, [{ oldText, newText }], cwd); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/edit.ts b/cactus-code/packages/coding-agent/src/core/tools/edit.ts new file mode 100644 index 000000000..25aa1e17b --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/edit.ts @@ -0,0 +1,437 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { Box, Container, Spacer, Text } from "@earendil-works/pi-tui"; +import { constants } from "fs"; +import { access as fsAccess, readFile as fsReadFile, writeFile as fsWriteFile } from "fs/promises"; +import { type Static, Type } from "typebox"; +import { renderDiff } from "../../modes/interactive/components/diff.ts"; +import type { Theme } from "../../modes/interactive/theme/theme.ts"; +import type { ToolDefinition } from "../extensions/types.ts"; +import { + applyEditsToNormalizedContent, + computeEditsDiff, + detectLineEnding, + type Edit, + type EditDiffError, + type EditDiffResult, + generateDiffString, + generateUnifiedPatch, + normalizeToLF, + restoreLineEndings, + stripBom, +} from "./edit-diff.ts"; +import { withFileMutationQueue } from "./file-mutation-queue.ts"; +import { resolveToCwd } from "./path-utils.ts"; +import { renderToolPath, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; + +type EditPreview = EditDiffResult | EditDiffError; + +type EditRenderState = { + callComponent?: EditCallRenderComponent; +}; + +const replaceEditSchema = Type.Object( + { + oldText: Type.String({ + description: + "Exact text for one targeted replacement. It must be unique in the original file and must not overlap with any other edits[].oldText in the same call.", + }), + newText: Type.String({ description: "Replacement text for this targeted edit." }), + }, + { additionalProperties: false }, +); + +const editSchema = Type.Object( + { + path: Type.String({ description: "Path to the file to edit (relative or absolute)" }), + edits: Type.Array(replaceEditSchema, { + description: + "One or more targeted replacements. Each edit is matched against the original file, not incrementally. Do not include overlapping or nested edits. If two changes touch the same block or nearby lines, merge them into one edit instead.", + }), + }, + { additionalProperties: false }, +); + +export type EditToolInput = Static; +type LegacyEditToolInput = EditToolInput & { + oldText?: unknown; + newText?: unknown; +}; + +export interface EditToolDetails { + /** Display-oriented diff of the changes made */ + diff: string; + /** Standard unified patch of the changes made */ + patch: string; + /** Line number of the first change in the new file (for editor navigation) */ + firstChangedLine?: number; +} + +/** + * Pluggable operations for the edit tool. + * Override these to delegate file editing to remote systems (for example SSH). + */ +export interface EditOperations { + /** Read file contents as a Buffer */ + readFile: (absolutePath: string) => Promise; + /** Write content to a file */ + writeFile: (absolutePath: string, content: string) => Promise; + /** Check if file is readable and writable (throw if not) */ + access: (absolutePath: string) => Promise; +} + +const defaultEditOperations: EditOperations = { + readFile: (path) => fsReadFile(path), + writeFile: (path, content) => fsWriteFile(path, content, "utf-8"), + access: (path) => fsAccess(path, constants.R_OK | constants.W_OK), +}; + +export interface EditToolOptions { + /** Custom operations for file editing. Default: local filesystem */ + operations?: EditOperations; +} + +function prepareEditArguments(input: unknown): EditToolInput { + if (!input || typeof input !== "object") { + return input as EditToolInput; + } + + const args = input as Record; + + // Some models (Opus 4.6, GLM-5.1) send edits as a JSON string instead of an array + if (typeof args.edits === "string") { + try { + const parsed = JSON.parse(args.edits); + if (Array.isArray(parsed)) args.edits = parsed; + } catch {} + } + + const legacy = args as LegacyEditToolInput; + if (typeof legacy.oldText !== "string" || typeof legacy.newText !== "string") { + return args as EditToolInput; + } + + const edits = Array.isArray(legacy.edits) ? [...legacy.edits] : []; + edits.push({ oldText: legacy.oldText, newText: legacy.newText }); + const { oldText: _oldText, newText: _newText, ...rest } = legacy; + return { ...rest, edits } as EditToolInput; +} + +function validateEditInput(input: EditToolInput): { path: string; edits: Edit[] } { + if (!Array.isArray(input.edits) || input.edits.length === 0) { + throw new Error("Edit tool input is invalid. edits must contain at least one replacement."); + } + return { path: input.path, edits: input.edits }; +} + +type RenderableEditArgs = { + path?: string; + file_path?: string; + edits?: Edit[]; + oldText?: string; + newText?: string; +}; + +type EditToolResultLike = { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + details?: EditToolDetails; +}; + +type EditCallRenderComponent = Box & { + preview?: EditPreview; + previewArgsKey?: string; + previewPending?: boolean; + settledError?: boolean; +}; + +function createEditCallRenderComponent(): EditCallRenderComponent { + return Object.assign(new Box(1, 1, (text: string) => text), { + preview: undefined as EditPreview | undefined, + previewArgsKey: undefined as string | undefined, + previewPending: false, + settledError: false, + }); +} + +function getEditCallRenderComponent(state: EditRenderState, lastComponent: unknown): EditCallRenderComponent { + if (lastComponent instanceof Box) { + const component = lastComponent as EditCallRenderComponent; + state.callComponent = component; + return component; + } + if (state.callComponent) { + return state.callComponent; + } + const component = createEditCallRenderComponent(); + state.callComponent = component; + return component; +} + +function getRenderablePreviewInput(args: RenderableEditArgs | undefined): { path: string; edits: Edit[] } | null { + if (!args) { + return null; + } + + const path = typeof args.path === "string" ? args.path : typeof args.file_path === "string" ? args.file_path : null; + if (!path) { + return null; + } + + if ( + Array.isArray(args.edits) && + args.edits.length > 0 && + args.edits.every((edit) => typeof edit?.oldText === "string" && typeof edit?.newText === "string") + ) { + return { path, edits: args.edits }; + } + + if (typeof args.oldText === "string" && typeof args.newText === "string") { + return { path, edits: [{ oldText: args.oldText, newText: args.newText }] }; + } + + return null; +} + +function formatEditCall(args: RenderableEditArgs | undefined, theme: Theme, cwd: string): string { + const pathDisplay = renderToolPath(str(args?.file_path ?? args?.path), theme, cwd); + return `${theme.fg("toolTitle", theme.bold("edit"))} ${pathDisplay}`; +} + +function formatEditResult( + args: RenderableEditArgs | undefined, + preview: EditPreview | undefined, + result: EditToolResultLike, + theme: Theme, + isError: boolean, +): string | undefined { + const rawPath = str(args?.file_path ?? args?.path); + const previewDiff = preview && !("error" in preview) ? preview.diff : undefined; + const previewError = preview && "error" in preview ? preview.error : undefined; + if (isError) { + const errorText = result.content + .filter((c) => c.type === "text") + .map((c) => c.text || "") + .join("\n"); + if (!errorText || errorText === previewError) { + return undefined; + } + return theme.fg("error", errorText); + } + + const resultDiff = result.details?.diff; + if (resultDiff && resultDiff !== previewDiff) { + return renderDiff(resultDiff, { filePath: rawPath ?? undefined }); + } + + return undefined; +} + +function getEditHeaderBg( + preview: EditPreview | undefined, + settledError: boolean | undefined, + theme: Theme, +): (text: string) => string { + if (preview) { + if ("error" in preview) { + return (text: string) => theme.bg("toolErrorBg", text); + } + return (text: string) => theme.bg("toolSuccessBg", text); + } + if (settledError) { + return (text: string) => theme.bg("toolErrorBg", text); + } + return (text: string) => theme.bg("toolPendingBg", text); +} + +function buildEditCallComponent( + component: EditCallRenderComponent, + args: RenderableEditArgs | undefined, + theme: Theme, + cwd: string, +): EditCallRenderComponent { + component.setBgFn(getEditHeaderBg(component.preview, component.settledError, theme)); + component.clear(); + component.addChild(new Text(formatEditCall(args, theme, cwd), 0, 0)); + + if (!component.preview) { + return component; + } + + const body = + "error" in component.preview ? theme.fg("error", component.preview.error) : renderDiff(component.preview.diff); + component.addChild(new Spacer(1)); + component.addChild(new Text(body, 0, 0)); + return component; +} + +function setEditPreview( + component: EditCallRenderComponent, + preview: EditPreview, + argsKey: string | undefined, +): boolean { + const current = component.preview; + const changed = + current === undefined || + ("error" in current && "error" in preview + ? current.error !== preview.error + : "error" in current !== "error" in preview) || + (!("error" in current) && + !("error" in preview) && + (current.diff !== preview.diff || current.firstChangedLine !== preview.firstChangedLine)); + component.preview = preview; + component.previewArgsKey = argsKey; + component.previewPending = false; + return changed; +} + +export function createEditToolDefinition( + cwd: string, + options?: EditToolOptions, +): ToolDefinition { + const ops = options?.operations ?? defaultEditOperations; + return { + name: "edit", + label: "edit", + description: + "Edit a single file using exact text replacement. Every edits[].oldText must match a unique, non-overlapping region of the original file. If two changes affect the same block or nearby lines, merge them into one edit instead of emitting overlapping edits. Do not include large unchanged regions just to connect distant changes.", + promptSnippet: + "Make precise file edits with exact text replacement, including multiple disjoint edits in one call", + promptGuidelines: [ + "Use edit for precise changes (edits[].oldText must match exactly)", + "When changing multiple separate locations in one file, use one edit call with multiple entries in edits[] instead of multiple edit calls", + "Each edits[].oldText is matched against the original file, not after earlier edits are applied. Do not emit overlapping or nested edits. Merge nearby changes into one edit.", + "Keep edits[].oldText as small as possible while still being unique in the file. Do not pad with large unchanged regions.", + ], + parameters: editSchema, + renderShell: "self", + prepareArguments: prepareEditArguments, + async execute(_toolCallId, input: EditToolInput, signal?: AbortSignal, _onUpdate?, _ctx?) { + const { path, edits } = validateEditInput(input); + const absolutePath = resolveToCwd(path, cwd); + + return withFileMutationQueue(absolutePath, async () => { + // Do not reject from an abort event listener here: that would release the + // mutation queue while an in-flight filesystem operation may still finish. + // Checking signal.aborted after each await observes the same aborts while + // keeping the queue locked until the current operation has settled. + const throwIfAborted = (): void => { + if (signal?.aborted) throw new Error("Operation aborted"); + }; + + throwIfAborted(); + + // Check if file exists. + try { + await ops.access(absolutePath); + } catch (error: unknown) { + throwIfAborted(); + const errorMessage = + error instanceof Error && "code" in error ? `Error code: ${error.code}` : String(error); + throw new Error(`Could not edit file: ${path}. ${errorMessage}.`); + } + throwIfAborted(); + + // Read the file. + const buffer = await ops.readFile(absolutePath); + const rawContent = buffer.toString("utf-8"); + throwIfAborted(); + + // Strip BOM before matching. The model will not include an invisible BOM in oldText. + const { bom, text: content } = stripBom(rawContent); + const originalEnding = detectLineEnding(content); + const normalizedContent = normalizeToLF(content); + const { baseContent, newContent } = applyEditsToNormalizedContent(normalizedContent, edits, path); + throwIfAborted(); + + const finalContent = bom + restoreLineEndings(newContent, originalEnding); + await ops.writeFile(absolutePath, finalContent); + throwIfAborted(); + + const diffResult = generateDiffString(baseContent, newContent); + const patch = generateUnifiedPatch(path, baseContent, newContent); + return { + content: [ + { + type: "text", + text: `Successfully replaced ${edits.length} block(s) in ${path}.`, + }, + ], + details: { diff: diffResult.diff, patch, firstChangedLine: diffResult.firstChangedLine }, + }; + }); + }, + renderCall(args, theme, context) { + const component = getEditCallRenderComponent(context.state, context.lastComponent); + const previewInput = getRenderablePreviewInput(args as RenderableEditArgs | undefined); + const argsKey = previewInput + ? JSON.stringify({ path: previewInput.path, edits: previewInput.edits }) + : undefined; + + if (component.previewArgsKey !== argsKey) { + component.preview = undefined; + component.previewArgsKey = argsKey; + component.previewPending = false; + component.settledError = false; + } + + if (context.argsComplete && previewInput && !component.preview && !component.previewPending) { + component.previewPending = true; + const requestKey = argsKey; + void computeEditsDiff(previewInput.path, previewInput.edits, context.cwd).then((preview) => { + if (component.previewArgsKey === requestKey) { + setEditPreview(component, preview, requestKey); + context.invalidate(); + } + }); + } + + return buildEditCallComponent(component, args, theme, context.cwd); + }, + renderResult(result, _options, theme, context) { + const callComponent = context.state.callComponent; + const previewInput = getRenderablePreviewInput(context.args as RenderableEditArgs | undefined); + const argsKey = previewInput + ? JSON.stringify({ path: previewInput.path, edits: previewInput.edits }) + : undefined; + const typedResult = result as EditToolResultLike; + const resultDiff = !context.isError ? typedResult.details?.diff : undefined; + let changed = false; + if (callComponent) { + if (typeof resultDiff === "string") { + changed = + setEditPreview( + callComponent, + { diff: resultDiff, firstChangedLine: typedResult.details?.firstChangedLine }, + argsKey, + ) || changed; + } + if (callComponent.settledError !== context.isError) { + callComponent.settledError = context.isError; + changed = true; + } + if (changed) { + buildEditCallComponent( + callComponent, + context.args as RenderableEditArgs | undefined, + theme, + context.cwd, + ); + } + } + + const output = formatEditResult(context.args, callComponent?.preview, typedResult, theme, context.isError); + const component = (context.lastComponent as Container | undefined) ?? new Container(); + component.clear(); + if (!output) { + return component; + } + component.addChild(new Spacer(1)); + component.addChild(new Text(output, 1, 0)); + return component; + }, + }; +} + +export function createEditTool(cwd: string, options?: EditToolOptions): AgentTool { + return wrapToolDefinition(createEditToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/file-mutation-queue.ts b/cactus-code/packages/coding-agent/src/core/tools/file-mutation-queue.ts new file mode 100644 index 000000000..5505a7a27 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/file-mutation-queue.ts @@ -0,0 +1,61 @@ +import { realpath } from "node:fs/promises"; +import { resolve } from "node:path"; + +const fileMutationQueues = new Map>(); +let registrationQueue = Promise.resolve(); + +function isMissingPathError(error: unknown): boolean { + return ( + typeof error === "object" && + error !== null && + "code" in error && + (error.code === "ENOENT" || error.code === "ENOTDIR") + ); +} + +async function getMutationQueueKey(filePath: string): Promise { + const resolvedPath = resolve(filePath); + try { + return await realpath(resolvedPath); + } catch (error) { + if (isMissingPathError(error)) { + return resolvedPath; + } + throw error; + } +} + +/** + * Serialize file mutation operations targeting the same file. + * Operations for different files still run in parallel. + */ +export async function withFileMutationQueue(filePath: string, fn: () => Promise): Promise { + const registration = registrationQueue.then(async () => { + const key = await getMutationQueueKey(filePath); + const currentQueue = fileMutationQueues.get(key) ?? Promise.resolve(); + + let releaseNext!: () => void; + const nextQueue = new Promise((resolveQueue) => { + releaseNext = resolveQueue; + }); + const chainedQueue = currentQueue.then(() => nextQueue); + fileMutationQueues.set(key, chainedQueue); + + return { key, currentQueue, chainedQueue, releaseNext }; + }); + registrationQueue = registration.then( + () => undefined, + () => undefined, + ); + + const { key, currentQueue, chainedQueue, releaseNext } = await registration; + await currentQueue; + try { + return await fn(); + } finally { + releaseNext(); + if (fileMutationQueues.get(key) === chainedQueue) { + fileMutationQueues.delete(key); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/find.ts b/cactus-code/packages/coding-agent/src/core/tools/find.ts new file mode 100644 index 000000000..e03b728a9 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/find.ts @@ -0,0 +1,374 @@ +import { createInterface } from "node:readline"; +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { Text } from "@earendil-works/pi-tui"; +import { spawn } from "child_process"; +import path from "path"; +import { type Static, Type } from "typebox"; +import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts"; +import type { Theme } from "../../modes/interactive/theme/theme.ts"; +import { ensureTool } from "../../utils/tools-manager.ts"; +import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts"; +import { pathExists, resolveToCwd } from "./path-utils.ts"; +import { getTextOutput, invalidArgText, shortenPath, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; +import { DEFAULT_MAX_BYTES, formatSize, type TruncationResult, truncateHead } from "./truncate.ts"; + +function toPosixPath(value: string): string { + return value.split(path.sep).join("/"); +} + +const findSchema = Type.Object({ + pattern: Type.String({ + description: "Glob pattern to match files, e.g. '*.ts', '**/*.json', or 'src/**/*.spec.ts'", + }), + path: Type.Optional(Type.String({ description: "Directory to search in (default: current directory)" })), + limit: Type.Optional(Type.Number({ description: "Maximum number of results (default: 1000)" })), +}); + +export type FindToolInput = Static; + +const DEFAULT_LIMIT = 1000; + +export interface FindToolDetails { + truncation?: TruncationResult; + resultLimitReached?: number; +} + +/** + * Pluggable operations for the find tool. + * Override these to delegate file search to remote systems (for example SSH). + */ +export interface FindOperations { + /** Check if path exists */ + exists: (absolutePath: string) => Promise | boolean; + /** Find files matching glob pattern. Returns relative or absolute paths. */ + glob: (pattern: string, cwd: string, options: { ignore: string[]; limit: number }) => Promise | string[]; +} + +const defaultFindOperations: FindOperations = { + exists: pathExists, + // This is a placeholder. Actual fd execution happens in execute() when no custom glob is provided. + glob: () => [], +}; + +export interface FindToolOptions { + /** Custom operations for find. Default: local filesystem plus fd */ + operations?: FindOperations; +} + +function formatFindCall(args: { pattern: string; path?: string; limit?: number } | undefined, theme: Theme): string { + const pattern = str(args?.pattern); + const rawPath = str(args?.path); + const path = rawPath !== null ? shortenPath(rawPath || ".") : null; + const limit = args?.limit; + const invalidArg = invalidArgText(theme); + let text = + theme.fg("toolTitle", theme.bold("find")) + + " " + + (pattern === null ? invalidArg : theme.fg("accent", pattern || "")) + + theme.fg("toolOutput", ` in ${path === null ? invalidArg : path}`); + if (limit !== undefined) { + text += theme.fg("toolOutput", ` (limit ${limit})`); + } + return text; +} + +function formatFindResult( + result: { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + details?: FindToolDetails; + }, + options: ToolRenderResultOptions, + theme: Theme, + showImages: boolean, +): string { + const output = getTextOutput(result, showImages).trim(); + let text = ""; + if (output) { + const lines = output.split("\n"); + const maxLines = options.expanded ? lines.length : 20; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; + text += `\n${displayLines.map((line) => theme.fg("toolOutput", line)).join("\n")}`; + if (remaining > 0) { + text += `${theme.fg("muted", `\n... (${remaining} more lines,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`; + } + } + + const resultLimit = result.details?.resultLimitReached; + const truncation = result.details?.truncation; + if (resultLimit || truncation?.truncated) { + const warnings: string[] = []; + if (resultLimit) warnings.push(`${resultLimit} results limit`); + if (truncation?.truncated) warnings.push(`${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit`); + text += `\n${theme.fg("warning", `[Truncated: ${warnings.join(", ")}]`)}`; + } + return text; +} + +export function createFindToolDefinition( + cwd: string, + options?: FindToolOptions, +): ToolDefinition { + const customOps = options?.operations; + return { + name: "find", + label: "find", + description: `Search for files by glob pattern. Returns matching file paths relative to the search directory. Respects .gitignore. Output is truncated to ${DEFAULT_LIMIT} results or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first).`, + promptSnippet: "Find files by glob pattern (respects .gitignore)", + parameters: findSchema, + async execute( + _toolCallId, + { pattern, path: searchDir, limit }: { pattern: string; path?: string; limit?: number }, + signal?: AbortSignal, + _onUpdate?, + _ctx?, + ) { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error("Operation aborted")); + return; + } + + let settled = false; + let stopChild: (() => void) | undefined; + const settle = (fn: () => void) => { + if (settled) return; + settled = true; + signal?.removeEventListener("abort", onAbort); + stopChild = undefined; + fn(); + }; + const onAbort = () => { + stopChild?.(); + settle(() => reject(new Error("Operation aborted"))); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + + (async () => { + try { + const searchPath = resolveToCwd(searchDir || ".", cwd); + const effectiveLimit = limit ?? DEFAULT_LIMIT; + const ops = customOps ?? defaultFindOperations; + + // If custom operations provide glob(), use that instead of fd. + if (customOps?.glob) { + if (!(await ops.exists(searchPath))) { + settle(() => reject(new Error(`Path not found: ${searchPath}`))); + return; + } + if (signal?.aborted) { + settle(() => reject(new Error("Operation aborted"))); + return; + } + const results = await ops.glob(pattern, searchPath, { + ignore: ["**/node_modules/**", "**/.git/**"], + limit: effectiveLimit, + }); + if (signal?.aborted) { + settle(() => reject(new Error("Operation aborted"))); + return; + } + if (results.length === 0) { + settle(() => + resolve({ + content: [{ type: "text", text: "No files found matching pattern" }], + details: undefined, + }), + ); + return; + } + + // Relativize paths against the search root for stable output. + const relativized = results.map((p) => { + if (p.startsWith(searchPath)) return toPosixPath(p.slice(searchPath.length + 1)); + return toPosixPath(path.relative(searchPath, p)); + }); + const resultLimitReached = relativized.length >= effectiveLimit; + const rawOutput = relativized.join("\n"); + const truncation = truncateHead(rawOutput, { maxLines: Number.MAX_SAFE_INTEGER }); + let resultOutput = truncation.content; + const details: FindToolDetails = {}; + const notices: string[] = []; + if (resultLimitReached) { + notices.push(`${effectiveLimit} results limit reached`); + details.resultLimitReached = effectiveLimit; + } + if (truncation.truncated) { + notices.push(`${formatSize(DEFAULT_MAX_BYTES)} limit reached`); + details.truncation = truncation; + } + if (notices.length > 0) { + resultOutput += `\n\n[${notices.join(". ")}]`; + } + settle(() => + resolve({ + content: [{ type: "text", text: resultOutput }], + details: Object.keys(details).length > 0 ? details : undefined, + }), + ); + return; + } + + // Default implementation uses fd. + const fdPath = await ensureTool("fd", true); + if (signal?.aborted) { + settle(() => reject(new Error("Operation aborted"))); + return; + } + if (!fdPath) { + settle(() => reject(new Error("fd is not available and could not be downloaded"))); + return; + } + + const args: string[] = ["--glob", "--color=never", "--hidden"]; + + // fd normally ignores .gitignore outside git repos, so keep --no-require-git + // there. Inside repos, use fd's default git-aware behavior so parent + // .gitignore rules stop at nested repo boundaries: + // https://github.com/earendil-works/pi/issues/5960 + let insideGitRepo = false; + for (let current = searchPath; ; ) { + if (await pathExists(path.join(current, ".git"))) { + insideGitRepo = true; + break; + } + const parent = path.dirname(current); + if (parent === current) break; + current = parent; + } + if (!insideGitRepo) args.push("--no-require-git"); + args.push("--max-results", String(effectiveLimit)); + + // fd --glob matches against the basename unless --full-path is set; in --full-path + // mode it matches against the absolute candidate path, so a path-containing + // pattern like 'src/**/*.spec.ts' needs a leading '**/' to match anything. + let effectivePattern = pattern; + if (pattern.includes("/")) { + args.push("--full-path"); + if (!pattern.startsWith("/") && !pattern.startsWith("**/") && pattern !== "**") { + effectivePattern = `**/${pattern}`; + } + } + args.push("--", effectivePattern, searchPath); + + const child = spawn(fdPath, args, { stdio: ["ignore", "pipe", "pipe"] }); + const rl = createInterface({ input: child.stdout }); + let stderr = ""; + const lines: string[] = []; + + stopChild = () => { + if (!child.killed) { + child.kill(); + } + }; + + const cleanup = () => { + rl.close(); + }; + + child.stderr?.on("data", (chunk) => { + stderr += chunk.toString(); + }); + + rl.on("line", (line) => { + lines.push(line); + }); + + child.on("error", (error) => { + cleanup(); + settle(() => reject(new Error(`Failed to run fd: ${error.message}`))); + }); + + child.on("close", (code) => { + cleanup(); + if (signal?.aborted) { + settle(() => reject(new Error("Operation aborted"))); + return; + } + const output = lines.join("\n"); + if (code !== 0) { + const errorMsg = stderr.trim() || `fd exited with code ${code}`; + if (!output) { + settle(() => reject(new Error(errorMsg))); + return; + } + } + if (!output) { + settle(() => + resolve({ + content: [{ type: "text", text: "No files found matching pattern" }], + details: undefined, + }), + ); + return; + } + + const relativized: string[] = []; + for (const rawLine of lines) { + const line = rawLine.replace(/\r$/, "").trim(); + if (!line) continue; + const hadTrailingSlash = line.endsWith("/") || line.endsWith("\\"); + let relativePath = line; + if (line.startsWith(searchPath)) { + relativePath = line.slice(searchPath.length + 1); + } else { + relativePath = path.relative(searchPath, line); + } + if (hadTrailingSlash && !relativePath.endsWith("/")) relativePath += "/"; + relativized.push(toPosixPath(relativePath)); + } + + const resultLimitReached = relativized.length >= effectiveLimit; + const rawOutput = relativized.join("\n"); + const truncation = truncateHead(rawOutput, { maxLines: Number.MAX_SAFE_INTEGER }); + let resultOutput = truncation.content; + const details: FindToolDetails = {}; + const notices: string[] = []; + if (resultLimitReached) { + notices.push( + `${effectiveLimit} results limit reached. Use limit=${effectiveLimit * 2} for more, or refine pattern`, + ); + details.resultLimitReached = effectiveLimit; + } + if (truncation.truncated) { + notices.push(`${formatSize(DEFAULT_MAX_BYTES)} limit reached`); + details.truncation = truncation; + } + if (notices.length > 0) { + resultOutput += `\n\n[${notices.join(". ")}]`; + } + settle(() => + resolve({ + content: [{ type: "text", text: resultOutput }], + details: Object.keys(details).length > 0 ? details : undefined, + }), + ); + }); + } catch (e) { + if (signal?.aborted) { + settle(() => reject(new Error("Operation aborted"))); + return; + } + const error = e instanceof Error ? e : new Error(String(e)); + settle(() => reject(error)); + } + })(); + }); + }, + renderCall(args, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatFindCall(args, theme)); + return text; + }, + renderResult(result, options, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatFindResult(result as any, options, theme, context.showImages)); + return text; + }, + }; +} + +export function createFindTool(cwd: string, options?: FindToolOptions): AgentTool { + return wrapToolDefinition(createFindToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/grep.ts b/cactus-code/packages/coding-agent/src/core/tools/grep.ts new file mode 100644 index 000000000..e4ed36d1b --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/grep.ts @@ -0,0 +1,385 @@ +import { readFile as fsReadFile, stat as fsStat } from "node:fs/promises"; +import { createInterface } from "node:readline"; +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { Text } from "@earendil-works/pi-tui"; +import { spawn } from "child_process"; +import path from "path"; +import { type Static, Type } from "typebox"; +import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts"; +import type { Theme } from "../../modes/interactive/theme/theme.ts"; +import { ensureTool } from "../../utils/tools-manager.ts"; +import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts"; +import { resolveToCwd } from "./path-utils.ts"; +import { getTextOutput, invalidArgText, shortenPath, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; +import { + DEFAULT_MAX_BYTES, + formatSize, + GREP_MAX_LINE_LENGTH, + type TruncationResult, + truncateHead, + truncateLine, +} from "./truncate.ts"; + +const grepSchema = Type.Object({ + pattern: Type.String({ description: "Search pattern (regex or literal string)" }), + path: Type.Optional(Type.String({ description: "Directory or file to search (default: current directory)" })), + glob: Type.Optional(Type.String({ description: "Filter files by glob pattern, e.g. '*.ts' or '**/*.spec.ts'" })), + ignoreCase: Type.Optional(Type.Boolean({ description: "Case-insensitive search (default: false)" })), + literal: Type.Optional( + Type.Boolean({ description: "Treat pattern as literal string instead of regex (default: false)" }), + ), + context: Type.Optional( + Type.Number({ description: "Number of lines to show before and after each match (default: 0)" }), + ), + limit: Type.Optional(Type.Number({ description: "Maximum number of matches to return (default: 100)" })), +}); + +export type GrepToolInput = Static; +const DEFAULT_LIMIT = 100; + +export interface GrepToolDetails { + truncation?: TruncationResult; + matchLimitReached?: number; + linesTruncated?: boolean; +} + +/** + * Pluggable operations for the grep tool. + * Override these to delegate search to remote systems (for example SSH). + */ +export interface GrepOperations { + /** Check if path is a directory. Throws if path does not exist. */ + isDirectory: (absolutePath: string) => Promise | boolean; + /** Read file contents for context lines */ + readFile: (absolutePath: string) => Promise | string; +} + +const defaultGrepOperations: GrepOperations = { + isDirectory: async (p) => (await fsStat(p)).isDirectory(), + readFile: (p) => fsReadFile(p, "utf-8"), +}; + +export interface GrepToolOptions { + /** Custom operations for grep. Default: local filesystem plus ripgrep */ + operations?: GrepOperations; +} + +function formatGrepCall( + args: { pattern: string; path?: string; glob?: string; limit?: number } | undefined, + theme: Theme, +): string { + const pattern = str(args?.pattern); + const rawPath = str(args?.path); + const path = rawPath !== null ? shortenPath(rawPath || ".") : null; + const glob = str(args?.glob); + const limit = args?.limit; + const invalidArg = invalidArgText(theme); + let text = + theme.fg("toolTitle", theme.bold("grep")) + + " " + + (pattern === null ? invalidArg : theme.fg("accent", `/${pattern || ""}/`)) + + theme.fg("toolOutput", ` in ${path === null ? invalidArg : path}`); + if (glob) text += theme.fg("toolOutput", ` (${glob})`); + if (limit !== undefined) text += theme.fg("toolOutput", ` limit ${limit}`); + return text; +} + +function formatGrepResult( + result: { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + details?: GrepToolDetails; + }, + options: ToolRenderResultOptions, + theme: Theme, + showImages: boolean, +): string { + const output = getTextOutput(result, showImages).trim(); + let text = ""; + if (output) { + const lines = output.split("\n"); + const maxLines = options.expanded ? lines.length : 15; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; + text += `\n${displayLines.map((line) => theme.fg("toolOutput", line)).join("\n")}`; + if (remaining > 0) { + text += `${theme.fg("muted", `\n... (${remaining} more lines,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`; + } + } + + const matchLimit = result.details?.matchLimitReached; + const truncation = result.details?.truncation; + const linesTruncated = result.details?.linesTruncated; + if (matchLimit || truncation?.truncated || linesTruncated) { + const warnings: string[] = []; + if (matchLimit) warnings.push(`${matchLimit} matches limit`); + if (truncation?.truncated) warnings.push(`${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit`); + if (linesTruncated) warnings.push("some lines truncated"); + text += `\n${theme.fg("warning", `[Truncated: ${warnings.join(", ")}]`)}`; + } + return text; +} + +export function createGrepToolDefinition( + cwd: string, + options?: GrepToolOptions, +): ToolDefinition { + const customOps = options?.operations; + return { + name: "grep", + label: "grep", + description: `Search file contents for a pattern. Returns matching lines with file paths and line numbers. Respects .gitignore. Output is truncated to ${DEFAULT_LIMIT} matches or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). Long lines are truncated to ${GREP_MAX_LINE_LENGTH} chars.`, + promptSnippet: "Search file contents for patterns (respects .gitignore)", + parameters: grepSchema, + async execute( + _toolCallId, + { + pattern, + path: searchDir, + glob, + ignoreCase, + literal, + context, + limit, + }: { + pattern: string; + path?: string; + glob?: string; + ignoreCase?: boolean; + literal?: boolean; + context?: number; + limit?: number; + }, + signal?: AbortSignal, + _onUpdate?, + _ctx?, + ) { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error("Operation aborted")); + return; + } + let settled = false; + const settle = (fn: () => void) => { + if (!settled) { + settled = true; + fn(); + } + }; + + (async () => { + try { + const rgPath = await ensureTool("rg", true); + if (!rgPath) { + settle(() => reject(new Error("ripgrep (rg) is not available and could not be downloaded"))); + return; + } + + const searchPath = resolveToCwd(searchDir || ".", cwd); + const ops = customOps ?? defaultGrepOperations; + let isDirectory: boolean; + try { + isDirectory = await ops.isDirectory(searchPath); + } catch { + settle(() => reject(new Error(`Path not found: ${searchPath}`))); + return; + } + + const contextValue = context && context > 0 ? context : 0; + const effectiveLimit = Math.max(1, limit ?? DEFAULT_LIMIT); + const formatPath = (filePath: string): string => { + if (isDirectory) { + const relative = path.relative(searchPath, filePath); + if (relative && !relative.startsWith("..")) { + return relative.replace(/\\/g, "/"); + } + } + return path.basename(filePath); + }; + + const fileCache = new Map(); + const getFileLines = async (filePath: string): Promise => { + let lines = fileCache.get(filePath); + if (!lines) { + try { + const content = await ops.readFile(filePath); + lines = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n").split("\n"); + } catch { + lines = []; + } + fileCache.set(filePath, lines); + } + return lines; + }; + + const args: string[] = ["--json", "--line-number", "--color=never", "--hidden"]; + if (ignoreCase) args.push("--ignore-case"); + if (literal) args.push("--fixed-strings"); + if (glob) args.push("--glob", glob); + args.push("--", pattern, searchPath); + + const child = spawn(rgPath, args, { stdio: ["ignore", "pipe", "pipe"] }); + const rl = createInterface({ input: child.stdout }); + let stderr = ""; + let matchCount = 0; + let matchLimitReached = false; + let linesTruncated = false; + let aborted = false; + let killedDueToLimit = false; + const outputLines: string[] = []; + + const cleanup = () => { + rl.close(); + signal?.removeEventListener("abort", onAbort); + }; + const stopChild = (dueToLimit = false) => { + if (!child.killed) { + killedDueToLimit = dueToLimit; + child.kill(); + } + }; + const onAbort = () => { + aborted = true; + stopChild(); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + child.stderr?.on("data", (chunk) => { + stderr += chunk.toString(); + }); + + const formatBlock = async (filePath: string, lineNumber: number): Promise => { + const relativePath = formatPath(filePath); + const lines = await getFileLines(filePath); + if (!lines.length) return [`${relativePath}:${lineNumber}: (unable to read file)`]; + const block: string[] = []; + const start = contextValue > 0 ? Math.max(1, lineNumber - contextValue) : lineNumber; + const end = contextValue > 0 ? Math.min(lines.length, lineNumber + contextValue) : lineNumber; + for (let current = start; current <= end; current++) { + const lineText = lines[current - 1] ?? ""; + const sanitized = lineText.replace(/\r/g, ""); + const isMatchLine = current === lineNumber; + // Truncate long lines so grep output stays compact. + const { text: truncatedText, wasTruncated } = truncateLine(sanitized); + if (wasTruncated) linesTruncated = true; + if (isMatchLine) block.push(`${relativePath}:${current}: ${truncatedText}`); + else block.push(`${relativePath}-${current}- ${truncatedText}`); + } + return block; + }; + + // Collect matches during streaming, then format them after rg exits. + const matches: Array<{ filePath: string; lineNumber: number; lineText?: string }> = []; + rl.on("line", (line) => { + if (!line.trim() || matchCount >= effectiveLimit) return; + let event: any; + try { + event = JSON.parse(line); + } catch { + return; + } + if (event.type === "match") { + matchCount++; + const filePath = event.data?.path?.text; + const lineNumber = event.data?.line_number; + const lineText = event.data?.lines?.text; + if (filePath && typeof lineNumber === "number") + matches.push({ filePath, lineNumber, lineText }); + if (matchCount >= effectiveLimit) { + matchLimitReached = true; + stopChild(true); + } + } + }); + + child.on("error", (error) => { + cleanup(); + settle(() => reject(new Error(`Failed to run ripgrep: ${error.message}`))); + }); + child.on("close", async (code) => { + cleanup(); + if (aborted) { + settle(() => reject(new Error("Operation aborted"))); + return; + } + if (!killedDueToLimit && code !== 0 && code !== 1) { + const errorMsg = stderr.trim() || `ripgrep exited with code ${code}`; + settle(() => reject(new Error(errorMsg))); + return; + } + if (matchCount === 0) { + settle(() => + resolve({ content: [{ type: "text", text: "No matches found" }], details: undefined }), + ); + return; + } + + // Format matches after streaming finishes so custom readFile() backends can be async. + for (const match of matches) { + if (contextValue === 0 && match.lineText !== undefined) { + const relativePath = formatPath(match.filePath); + const sanitized = match.lineText + .replace(/\r\n/g, "\n") + .replace(/\r/g, "") + .replace(/\n$/, ""); + const { text: truncatedText, wasTruncated } = truncateLine(sanitized); + if (wasTruncated) linesTruncated = true; + outputLines.push(`${relativePath}:${match.lineNumber}: ${truncatedText}`); + } else { + const block = await formatBlock(match.filePath, match.lineNumber); + outputLines.push(...block); + } + } + + const rawOutput = outputLines.join("\n"); + // Apply byte truncation. There is no line limit here because the match limit already capped rows. + const truncation = truncateHead(rawOutput, { maxLines: Number.MAX_SAFE_INTEGER }); + let output = truncation.content; + const details: GrepToolDetails = {}; + // Build actionable notices for truncation and match limits. + const notices: string[] = []; + if (matchLimitReached) { + notices.push( + `${effectiveLimit} matches limit reached. Use limit=${effectiveLimit * 2} for more, or refine pattern`, + ); + details.matchLimitReached = effectiveLimit; + } + if (truncation.truncated) { + notices.push(`${formatSize(DEFAULT_MAX_BYTES)} limit reached`); + details.truncation = truncation; + } + if (linesTruncated) { + notices.push( + `Some lines truncated to ${GREP_MAX_LINE_LENGTH} chars. Use read tool to see full lines`, + ); + details.linesTruncated = true; + } + if (notices.length > 0) output += `\n\n[${notices.join(". ")}]`; + settle(() => + resolve({ + content: [{ type: "text", text: output }], + details: Object.keys(details).length > 0 ? details : undefined, + }), + ); + }); + } catch (err) { + settle(() => reject(err as Error)); + } + })(); + }); + }, + renderCall(args, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatGrepCall(args, theme)); + return text; + }, + renderResult(result, options, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatGrepResult(result as any, options, theme, context.showImages)); + return text; + }, + }; +} + +export function createGrepTool(cwd: string, options?: GrepToolOptions): AgentTool { + return wrapToolDefinition(createGrepToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/index.ts b/cactus-code/packages/coding-agent/src/core/tools/index.ts new file mode 100644 index 000000000..e55f91406 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/index.ts @@ -0,0 +1,196 @@ +export { + type BashOperations, + type BashSpawnContext, + type BashSpawnHook, + type BashToolDetails, + type BashToolInput, + type BashToolOptions, + createBashTool, + createBashToolDefinition, + createLocalBashOperations, +} from "./bash.ts"; +export { + createEditTool, + createEditToolDefinition, + type EditOperations, + type EditToolDetails, + type EditToolInput, + type EditToolOptions, +} from "./edit.ts"; +export { withFileMutationQueue } from "./file-mutation-queue.ts"; +export { + createFindTool, + createFindToolDefinition, + type FindOperations, + type FindToolDetails, + type FindToolInput, + type FindToolOptions, +} from "./find.ts"; +export { + createGrepTool, + createGrepToolDefinition, + type GrepOperations, + type GrepToolDetails, + type GrepToolInput, + type GrepToolOptions, +} from "./grep.ts"; +export { + createLsTool, + createLsToolDefinition, + type LsOperations, + type LsToolDetails, + type LsToolInput, + type LsToolOptions, +} from "./ls.ts"; +export { + createReadTool, + createReadToolDefinition, + type ReadOperations, + type ReadToolDetails, + type ReadToolInput, + type ReadToolOptions, +} from "./read.ts"; +export { + DEFAULT_MAX_BYTES, + DEFAULT_MAX_LINES, + formatSize, + type TruncationOptions, + type TruncationResult, + truncateHead, + truncateLine, + truncateTail, +} from "./truncate.ts"; +export { + createWriteTool, + createWriteToolDefinition, + type WriteOperations, + type WriteToolInput, + type WriteToolOptions, +} from "./write.ts"; + +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import type { ToolDefinition } from "../extensions/types.ts"; +import { type BashToolOptions, createBashTool, createBashToolDefinition } from "./bash.ts"; +import { createEditTool, createEditToolDefinition, type EditToolOptions } from "./edit.ts"; +import { createFindTool, createFindToolDefinition, type FindToolOptions } from "./find.ts"; +import { createGrepTool, createGrepToolDefinition, type GrepToolOptions } from "./grep.ts"; +import { createLsTool, createLsToolDefinition, type LsToolOptions } from "./ls.ts"; +import { createReadTool, createReadToolDefinition, type ReadToolOptions } from "./read.ts"; +import { createWriteTool, createWriteToolDefinition, type WriteToolOptions } from "./write.ts"; + +export type Tool = AgentTool; +export type ToolDef = ToolDefinition; +export type ToolName = "read" | "bash" | "edit" | "write" | "grep" | "find" | "ls"; +export const allToolNames: Set = new Set(["read", "bash", "edit", "write", "grep", "find", "ls"]); + +export interface ToolsOptions { + read?: ReadToolOptions; + bash?: BashToolOptions; + write?: WriteToolOptions; + edit?: EditToolOptions; + grep?: GrepToolOptions; + find?: FindToolOptions; + ls?: LsToolOptions; +} + +export function createToolDefinition(toolName: ToolName, cwd: string, options?: ToolsOptions): ToolDef { + switch (toolName) { + case "read": + return createReadToolDefinition(cwd, options?.read); + case "bash": + return createBashToolDefinition(cwd, options?.bash); + case "edit": + return createEditToolDefinition(cwd, options?.edit); + case "write": + return createWriteToolDefinition(cwd, options?.write); + case "grep": + return createGrepToolDefinition(cwd, options?.grep); + case "find": + return createFindToolDefinition(cwd, options?.find); + case "ls": + return createLsToolDefinition(cwd, options?.ls); + default: + throw new Error(`Unknown tool name: ${toolName}`); + } +} + +export function createTool(toolName: ToolName, cwd: string, options?: ToolsOptions): Tool { + switch (toolName) { + case "read": + return createReadTool(cwd, options?.read); + case "bash": + return createBashTool(cwd, options?.bash); + case "edit": + return createEditTool(cwd, options?.edit); + case "write": + return createWriteTool(cwd, options?.write); + case "grep": + return createGrepTool(cwd, options?.grep); + case "find": + return createFindTool(cwd, options?.find); + case "ls": + return createLsTool(cwd, options?.ls); + default: + throw new Error(`Unknown tool name: ${toolName}`); + } +} + +export function createCodingToolDefinitions(cwd: string, options?: ToolsOptions): ToolDef[] { + return [ + createReadToolDefinition(cwd, options?.read), + createBashToolDefinition(cwd, options?.bash), + createEditToolDefinition(cwd, options?.edit), + createWriteToolDefinition(cwd, options?.write), + ]; +} + +export function createReadOnlyToolDefinitions(cwd: string, options?: ToolsOptions): ToolDef[] { + return [ + createReadToolDefinition(cwd, options?.read), + createGrepToolDefinition(cwd, options?.grep), + createFindToolDefinition(cwd, options?.find), + createLsToolDefinition(cwd, options?.ls), + ]; +} + +export function createAllToolDefinitions(cwd: string, options?: ToolsOptions): Record { + return { + read: createReadToolDefinition(cwd, options?.read), + bash: createBashToolDefinition(cwd, options?.bash), + edit: createEditToolDefinition(cwd, options?.edit), + write: createWriteToolDefinition(cwd, options?.write), + grep: createGrepToolDefinition(cwd, options?.grep), + find: createFindToolDefinition(cwd, options?.find), + ls: createLsToolDefinition(cwd, options?.ls), + }; +} + +export function createCodingTools(cwd: string, options?: ToolsOptions): Tool[] { + return [ + createReadTool(cwd, options?.read), + createBashTool(cwd, options?.bash), + createEditTool(cwd, options?.edit), + createWriteTool(cwd, options?.write), + ]; +} + +export function createReadOnlyTools(cwd: string, options?: ToolsOptions): Tool[] { + return [ + createReadTool(cwd, options?.read), + createGrepTool(cwd, options?.grep), + createFindTool(cwd, options?.find), + createLsTool(cwd, options?.ls), + ]; +} + +export function createAllTools(cwd: string, options?: ToolsOptions): Record { + return { + read: createReadTool(cwd, options?.read), + bash: createBashTool(cwd, options?.bash), + edit: createEditTool(cwd, options?.edit), + write: createWriteTool(cwd, options?.write), + grep: createGrepTool(cwd, options?.grep), + find: createFindTool(cwd, options?.find), + ls: createLsTool(cwd, options?.ls), + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/ls.ts b/cactus-code/packages/coding-agent/src/core/tools/ls.ts new file mode 100644 index 000000000..8a689e8da --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/ls.ts @@ -0,0 +1,225 @@ +import { readdir as fsReaddir, stat as fsStat } from "node:fs/promises"; +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { Text } from "@earendil-works/pi-tui"; +import nodePath from "path"; +import { type Static, Type } from "typebox"; +import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts"; +import type { Theme } from "../../modes/interactive/theme/theme.ts"; +import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts"; +import { pathExists, resolveToCwd } from "./path-utils.ts"; +import { getTextOutput, renderToolPath, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; +import { DEFAULT_MAX_BYTES, formatSize, type TruncationResult, truncateHead } from "./truncate.ts"; + +const lsSchema = Type.Object({ + path: Type.Optional(Type.String({ description: "Directory to list (default: current directory)" })), + limit: Type.Optional(Type.Number({ description: "Maximum number of entries to return (default: 500)" })), +}); + +export type LsToolInput = Static; + +const DEFAULT_LIMIT = 500; + +export interface LsToolDetails { + truncation?: TruncationResult; + entryLimitReached?: number; +} + +/** + * Pluggable operations for the ls tool. + * Override these to delegate directory listing to remote systems (for example SSH). + */ +export interface LsOperations { + /** Check if path exists */ + exists: (absolutePath: string) => Promise | boolean; + /** Get file or directory stats. Throws if not found. */ + stat: (absolutePath: string) => Promise<{ isDirectory: () => boolean }> | { isDirectory: () => boolean }; + /** Read directory entries */ + readdir: (absolutePath: string) => Promise | string[]; +} + +const defaultLsOperations: LsOperations = { + exists: pathExists, + stat: fsStat, + readdir: fsReaddir, +}; + +export interface LsToolOptions { + /** Custom operations for directory listing. Default: local filesystem */ + operations?: LsOperations; +} + +function formatLsCall(args: { path?: string; limit?: number } | undefined, theme: Theme, cwd: string): string { + const limit = args?.limit; + const pathDisplay = renderToolPath(str(args?.path), theme, cwd, { emptyFallback: "." }); + let text = `${theme.fg("toolTitle", theme.bold("ls"))} ${pathDisplay}`; + if (limit !== undefined) { + text += theme.fg("toolOutput", ` (limit ${limit})`); + } + return text; +} + +function formatLsResult( + result: { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + details?: LsToolDetails; + }, + options: ToolRenderResultOptions, + theme: Theme, + showImages: boolean, +): string { + const output = getTextOutput(result, showImages).trim(); + let text = ""; + if (output) { + const lines = output.split("\n"); + const maxLines = options.expanded ? lines.length : 20; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; + text += `\n${displayLines.map((line) => theme.fg("toolOutput", line)).join("\n")}`; + if (remaining > 0) { + text += `${theme.fg("muted", `\n... (${remaining} more lines,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`; + } + } + + const entryLimit = result.details?.entryLimitReached; + const truncation = result.details?.truncation; + if (entryLimit || truncation?.truncated) { + const warnings: string[] = []; + if (entryLimit) warnings.push(`${entryLimit} entries limit`); + if (truncation?.truncated) warnings.push(`${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit`); + text += `\n${theme.fg("warning", `[Truncated: ${warnings.join(", ")}]`)}`; + } + return text; +} + +export function createLsToolDefinition( + cwd: string, + options?: LsToolOptions, +): ToolDefinition { + const ops = options?.operations ?? defaultLsOperations; + return { + name: "ls", + label: "ls", + description: `List directory contents. Returns entries sorted alphabetically, with '/' suffix for directories. Includes dotfiles. Output is truncated to ${DEFAULT_LIMIT} entries or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first).`, + promptSnippet: "List directory contents", + parameters: lsSchema, + async execute( + _toolCallId, + { path, limit }: { path?: string; limit?: number }, + signal?: AbortSignal, + _onUpdate?, + _ctx?, + ) { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error("Operation aborted")); + return; + } + + const onAbort = () => reject(new Error("Operation aborted")); + signal?.addEventListener("abort", onAbort, { once: true }); + + (async () => { + try { + const dirPath = resolveToCwd(path || ".", cwd); + const effectiveLimit = limit ?? DEFAULT_LIMIT; + + // Check if path exists. + if (!(await ops.exists(dirPath))) { + reject(new Error(`Path not found: ${dirPath}`)); + return; + } + + // Check if path is a directory. + const stat = await ops.stat(dirPath); + if (!stat.isDirectory()) { + reject(new Error(`Not a directory: ${dirPath}`)); + return; + } + + // Read directory entries. + let entries: string[]; + try { + entries = await ops.readdir(dirPath); + } catch (e: any) { + reject(new Error(`Cannot read directory: ${e.message}`)); + return; + } + + // Sort alphabetically, case-insensitive. + entries.sort((a, b) => a.toLowerCase().localeCompare(b.toLowerCase())); + + // Format entries with directory indicators. + const results: string[] = []; + let entryLimitReached = false; + for (const entry of entries) { + if (results.length >= effectiveLimit) { + entryLimitReached = true; + break; + } + + const fullPath = nodePath.join(dirPath, entry); + let suffix = ""; + try { + const entryStat = await ops.stat(fullPath); + if (entryStat.isDirectory()) suffix = "/"; + } catch { + // Skip entries we cannot stat. + continue; + } + results.push(entry + suffix); + } + + signal?.removeEventListener("abort", onAbort); + + if (results.length === 0) { + resolve({ content: [{ type: "text", text: "(empty directory)" }], details: undefined }); + return; + } + + const rawOutput = results.join("\n"); + // Apply byte truncation. There is no separate line limit because entry count is already capped. + const truncation = truncateHead(rawOutput, { maxLines: Number.MAX_SAFE_INTEGER }); + let output = truncation.content; + const details: LsToolDetails = {}; + // Build actionable notices for truncation and entry limits. + const notices: string[] = []; + if (entryLimitReached) { + notices.push(`${effectiveLimit} entries limit reached. Use limit=${effectiveLimit * 2} for more`); + details.entryLimitReached = effectiveLimit; + } + if (truncation.truncated) { + notices.push(`${formatSize(DEFAULT_MAX_BYTES)} limit reached`); + details.truncation = truncation; + } + if (notices.length > 0) { + output += `\n\n[${notices.join(". ")}]`; + } + + resolve({ + content: [{ type: "text", text: output }], + details: Object.keys(details).length > 0 ? details : undefined, + }); + } catch (e: any) { + signal?.removeEventListener("abort", onAbort); + reject(e); + } + })(); + }); + }, + renderCall(args, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatLsCall(args, theme, context.cwd)); + return text; + }, + renderResult(result, options, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(formatLsResult(result as any, options, theme, context.showImages)); + return text; + }, + }; +} + +export function createLsTool(cwd: string, options?: LsToolOptions): AgentTool { + return wrapToolDefinition(createLsToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/output-accumulator.ts b/cactus-code/packages/coding-agent/src/core/tools/output-accumulator.ts new file mode 100644 index 000000000..5c25b73fa --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/output-accumulator.ts @@ -0,0 +1,222 @@ +import { randomBytes } from "node:crypto"; +import { createWriteStream, type WriteStream } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, type TruncationResult, truncateTail } from "./truncate.ts"; + +export interface OutputAccumulatorOptions { + maxLines?: number; + maxBytes?: number; + tempFilePrefix?: string; +} + +export interface OutputSnapshot { + content: string; + truncation: TruncationResult; + fullOutputPath?: string; +} + +function defaultTempFilePath(prefix: string): string { + const id = randomBytes(8).toString("hex"); + return join(tmpdir(), `${prefix}-${id}.log`); +} + +function byteLength(text: string): number { + return Buffer.byteLength(text, "utf-8"); +} + +/** + * Incrementally tracks streaming output with bounded memory. + * + * Appends decode chunks with a streaming UTF-8 decoder, keeps only a decoded + * tail for display snapshots, and opens a temp file when the full output needs + * to be preserved. + */ +export class OutputAccumulator { + private readonly maxLines: number; + private readonly maxBytes: number; + private readonly maxRollingBytes: number; + private readonly tempFilePrefix: string; + private readonly decoder = new TextDecoder(); + + private rawChunks: Buffer[] = []; + private tailText = ""; + private tailBytes = 0; + private tailStartsAtLineBoundary = true; + private totalRawBytes = 0; + private totalDecodedBytes = 0; + private completedLines = 0; + private totalLines = 0; + private currentLineBytes = 0; + private hasOpenLine = false; + private finished = false; + + private tempFilePath: string | undefined; + private tempFileStream: WriteStream | undefined; + + constructor(options: OutputAccumulatorOptions = {}) { + this.maxLines = options.maxLines ?? DEFAULT_MAX_LINES; + this.maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + this.maxRollingBytes = Math.max(this.maxBytes * 2, 1); + this.tempFilePrefix = options.tempFilePrefix ?? "cactus-output"; + } + + append(data: Buffer): void { + if (this.finished) { + throw new Error("Cannot append to a finished output accumulator"); + } + + this.totalRawBytes += data.length; + this.appendDecodedText(this.decoder.decode(data, { stream: true })); + + if (this.tempFileStream || this.shouldUseTempFile()) { + this.ensureTempFile(); + this.tempFileStream?.write(data); + } else if (data.length > 0) { + this.rawChunks.push(data); + } + } + + finish(): void { + if (this.finished) { + return; + } + this.finished = true; + this.appendDecodedText(this.decoder.decode()); + if (this.shouldUseTempFile()) { + this.ensureTempFile(); + } + } + + snapshot(options: { persistIfTruncated?: boolean } = {}): OutputSnapshot { + const tailTruncation = truncateTail(this.getSnapshotText(), { + maxLines: this.maxLines, + maxBytes: this.maxBytes, + }); + const truncated = this.totalLines > this.maxLines || this.totalDecodedBytes > this.maxBytes; + const truncatedBy = truncated + ? (tailTruncation.truncatedBy ?? (this.totalDecodedBytes > this.maxBytes ? "bytes" : "lines")) + : null; + const truncation: TruncationResult = { + ...tailTruncation, + truncated, + truncatedBy, + totalLines: this.totalLines, + totalBytes: this.totalDecodedBytes, + maxLines: this.maxLines, + maxBytes: this.maxBytes, + }; + + if (options.persistIfTruncated && truncation.truncated) { + this.ensureTempFile(); + } + + return { + content: truncation.content, + truncation, + fullOutputPath: this.tempFilePath, + }; + } + + async closeTempFile(): Promise { + if (!this.tempFileStream) { + return; + } + + const stream = this.tempFileStream; + this.tempFileStream = undefined; + + await new Promise((resolve, reject) => { + const onError = (error: Error) => { + stream.off("finish", onFinish); + reject(error); + }; + const onFinish = () => { + stream.off("error", onError); + resolve(); + }; + stream.once("error", onError); + stream.once("finish", onFinish); + stream.end(); + }); + } + + getLastLineBytes(): number { + return this.currentLineBytes; + } + + private appendDecodedText(text: string): void { + if (text.length === 0) { + return; + } + + const bytes = byteLength(text); + this.totalDecodedBytes += bytes; + this.tailText += text; + this.tailBytes += bytes; + if (this.tailBytes > this.maxRollingBytes * 2) { + this.trimTail(); + } + + let newlines = 0; + let lastNewline = -1; + for (let i = text.indexOf("\n"); i !== -1; i = text.indexOf("\n", i + 1)) { + newlines++; + lastNewline = i; + } + if (newlines === 0) { + this.currentLineBytes += bytes; + this.hasOpenLine = true; + } else { + this.completedLines += newlines; + const tail = text.slice(lastNewline + 1); + this.currentLineBytes = byteLength(tail); + this.hasOpenLine = tail.length > 0; + } + this.totalLines = this.completedLines + (this.hasOpenLine ? 1 : 0); + } + + private trimTail(): void { + const buffer = Buffer.from(this.tailText, "utf-8"); + if (buffer.length <= this.maxRollingBytes) { + this.tailBytes = buffer.length; + return; + } + + let start = buffer.length - this.maxRollingBytes; + while (start < buffer.length && (buffer[start] & 0xc0) === 0x80) { + start++; + } + + this.tailStartsAtLineBoundary = start === 0 ? this.tailStartsAtLineBoundary : buffer[start - 1] === 0x0a; + this.tailText = buffer.subarray(start).toString("utf-8"); + this.tailBytes = byteLength(this.tailText); + } + + private getSnapshotText(): string { + if (this.tailStartsAtLineBoundary) { + return this.tailText; + } + + const firstNewline = this.tailText.indexOf("\n"); + return firstNewline === -1 ? this.tailText : this.tailText.slice(firstNewline + 1); + } + + private shouldUseTempFile(): boolean { + return ( + this.totalRawBytes > this.maxBytes || this.totalDecodedBytes > this.maxBytes || this.totalLines > this.maxLines + ); + } + + private ensureTempFile(): void { + if (this.tempFilePath) { + return; + } + this.tempFilePath = defaultTempFilePath(this.tempFilePrefix); + this.tempFileStream = createWriteStream(this.tempFilePath); + for (const chunk of this.rawChunks) { + this.tempFileStream.write(chunk); + } + this.rawChunks = []; + } +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/path-utils.ts b/cactus-code/packages/coding-agent/src/core/tools/path-utils.ts new file mode 100644 index 000000000..1f9ab4cc9 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/path-utils.ts @@ -0,0 +1,118 @@ +import { accessSync, constants } from "node:fs"; +import { access } from "node:fs/promises"; +import { normalizePath, resolvePath } from "../../utils/paths.ts"; + +const NARROW_NO_BREAK_SPACE = "\u202F"; + +function tryMacOSScreenshotPath(filePath: string): string { + return filePath.replace(/ (AM|PM)\./gi, `${NARROW_NO_BREAK_SPACE}$1.`); +} + +function tryNFDVariant(filePath: string): string { + // macOS stores filenames in NFD (decomposed) form, try converting user input to NFD + return filePath.normalize("NFD"); +} + +function tryCurlyQuoteVariant(filePath: string): string { + // macOS uses U+2019 (right single quotation mark) in screenshot names like "Capture d'écran" + // Users typically type U+0027 (straight apostrophe) + return filePath.replace(/'/g, "\u2019"); +} + +function fileExists(filePath: string): boolean { + try { + accessSync(filePath, constants.F_OK); + return true; + } catch { + return false; + } +} + +export async function pathExists(filePath: string): Promise { + try { + await access(filePath, constants.F_OK); + return true; + } catch { + return false; + } +} + +export function expandPath(filePath: string): string { + return normalizePath(filePath, { normalizeUnicodeSpaces: true, stripAtPrefix: true }); +} + +/** + * Resolve a path relative to the given cwd. + * Handles ~ expansion and absolute paths. + */ +export function resolveToCwd(filePath: string, cwd: string): string { + return resolvePath(filePath, cwd, { normalizeUnicodeSpaces: true, stripAtPrefix: true }); +} + +export function resolveReadPath(filePath: string, cwd: string): string { + const resolved = resolveToCwd(filePath, cwd); + + if (fileExists(resolved)) { + return resolved; + } + + // Try macOS AM/PM variant (narrow no-break space before AM/PM) + const amPmVariant = tryMacOSScreenshotPath(resolved); + if (amPmVariant !== resolved && fileExists(amPmVariant)) { + return amPmVariant; + } + + // Try NFD variant (macOS stores filenames in NFD form) + const nfdVariant = tryNFDVariant(resolved); + if (nfdVariant !== resolved && fileExists(nfdVariant)) { + return nfdVariant; + } + + // Try curly quote variant (macOS uses U+2019 in screenshot names) + const curlyVariant = tryCurlyQuoteVariant(resolved); + if (curlyVariant !== resolved && fileExists(curlyVariant)) { + return curlyVariant; + } + + // Try combined NFD + curly quote (for French macOS screenshots like "Capture d'écran") + const nfdCurlyVariant = tryCurlyQuoteVariant(nfdVariant); + if (nfdCurlyVariant !== resolved && fileExists(nfdCurlyVariant)) { + return nfdCurlyVariant; + } + + return resolved; +} + +export async function resolveReadPathAsync(filePath: string, cwd: string): Promise { + const resolved = resolveToCwd(filePath, cwd); + + if (await pathExists(resolved)) { + return resolved; + } + + // Try macOS AM/PM variant (narrow no-break space before AM/PM) + const amPmVariant = tryMacOSScreenshotPath(resolved); + if (amPmVariant !== resolved && (await pathExists(amPmVariant))) { + return amPmVariant; + } + + // Try NFD variant (macOS stores filenames in NFD form) + const nfdVariant = tryNFDVariant(resolved); + if (nfdVariant !== resolved && (await pathExists(nfdVariant))) { + return nfdVariant; + } + + // Try curly quote variant (macOS uses U+2019 in screenshot names) + const curlyVariant = tryCurlyQuoteVariant(resolved); + if (curlyVariant !== resolved && (await pathExists(curlyVariant))) { + return curlyVariant; + } + + // Try combined NFD + curly quote (for French macOS screenshots like "Capture d'écran") + const nfdCurlyVariant = tryCurlyQuoteVariant(nfdVariant); + if (nfdCurlyVariant !== resolved && (await pathExists(nfdCurlyVariant))) { + return nfdCurlyVariant; + } + + return resolved; +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/read.ts b/cactus-code/packages/coding-agent/src/core/tools/read.ts new file mode 100644 index 000000000..52e87b1df --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/read.ts @@ -0,0 +1,362 @@ +import { basename, dirname, isAbsolute, relative, resolve as resolvePath, sep } from "node:path"; +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import type { Api, ImageContent, Model, TextContent } from "@earendil-works/pi-ai"; +import { Text } from "@earendil-works/pi-tui"; +import { constants } from "fs"; +import { access as fsAccess, readFile as fsReadFile } from "fs/promises"; +import { type Static, Type } from "typebox"; +import { getReadmePath } from "../../config.ts"; +import { keyHint, keyText } from "../../modes/interactive/components/keybinding-hints.ts"; +import { getLanguageFromPath, highlightCode, type Theme } from "../../modes/interactive/theme/theme.ts"; +import { formatDimensionNote, resizeImage } from "../../utils/image-resize.ts"; +import { detectSupportedImageMimeTypeFromFile } from "../../utils/mime.ts"; +import { formatPathRelativeToCwdOrAbsolute } from "../../utils/paths.ts"; +import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts"; +import { resolveReadPathAsync, resolveToCwd } from "./path-utils.ts"; +import { getTextOutput, renderToolPath, replaceTabs, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; +import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult, truncateHead } from "./truncate.ts"; + +const readSchema = Type.Object({ + path: Type.String({ description: "Path to the file to read (relative or absolute)" }), + offset: Type.Optional(Type.Number({ description: "Line number to start reading from (1-indexed)" })), + limit: Type.Optional(Type.Number({ description: "Maximum number of lines to read" })), +}); + +export type ReadToolInput = Static; + +export interface ReadToolDetails { + truncation?: TruncationResult; +} + +interface CompactReadClassification { + kind: "docs" | "resource" | "skill"; + label: string; +} + +const COMPACT_RESOURCE_FILE_NAMES = new Set(["AGENTS.md", "AGENTS.MD", "CLAUDE.md", "CLAUDE.MD"]); + +/** + * Pluggable operations for the read tool. + * Override these to delegate file reading to remote systems (for example SSH). + */ +export interface ReadOperations { + /** Read file contents as a Buffer */ + readFile: (absolutePath: string) => Promise; + /** Check if file is readable (throw if not) */ + access: (absolutePath: string) => Promise; + /** Detect image MIME type, return null or undefined for non-images */ + detectImageMimeType?: (absolutePath: string) => Promise; +} + +const defaultReadOperations: ReadOperations = { + readFile: (path) => fsReadFile(path), + access: (path) => fsAccess(path, constants.R_OK), + detectImageMimeType: detectSupportedImageMimeTypeFromFile, +}; + +export interface ReadToolOptions { + /** Whether to auto-resize images to 2000x2000 max. Default: true */ + autoResizeImages?: boolean; + /** Custom operations for file reading. Default: local filesystem */ + operations?: ReadOperations; +} + +type ReadRenderArgs = { path?: string; file_path?: string; offset?: number; limit?: number }; + +function formatReadLineRange(args: ReadRenderArgs | undefined, theme: Theme): string { + if (args?.offset === undefined && args?.limit === undefined) return ""; + const startLine = args.offset ?? 1; + const endLine = args.limit !== undefined ? startLine + args.limit - 1 : ""; + return theme.fg("warning", `:${startLine}${endLine ? `-${endLine}` : ""}`); +} + +function formatReadCall(args: ReadRenderArgs | undefined, theme: Theme, cwd: string): string { + const pathDisplay = renderToolPath(str(args?.file_path ?? args?.path), theme, cwd); + return `${theme.fg("toolTitle", theme.bold("read"))} ${pathDisplay}${formatReadLineRange(args, theme)}`; +} + +function trimTrailingEmptyLines(lines: string[]): string[] { + let end = lines.length; + while (end > 0 && lines[end - 1] === "") { + end--; + } + return lines.slice(0, end); +} + +function getNonVisionImageNote(model: Model | undefined): string | undefined { + if (!model || model.input.includes("image")) { + return undefined; + } + return "[Current model does not support images. The image will be omitted from this request.]"; +} + +function toPosixPath(filePath: string): string { + return filePath.split(sep).join("/"); +} + +function getPiDocsClassification(absolutePath: string): CompactReadClassification | undefined { + const packageRoot = dirname(getReadmePath()); + const relativePath = relative(resolvePath(packageRoot), resolvePath(absolutePath)); + if ( + relativePath === "" || + relativePath === ".." || + relativePath.startsWith(`..${sep}`) || + isAbsolute(relativePath) + ) { + return undefined; + } + + const label = toPosixPath(relativePath); + if (label === "README.md" || label.startsWith("docs/") || label.startsWith("examples/")) { + return { kind: "docs", label }; + } + return undefined; +} + +function getCompactReadClassification( + args: ReadRenderArgs | undefined, + cwd: string, +): CompactReadClassification | undefined { + const rawPath = str(args?.file_path ?? args?.path); + if (!rawPath) return undefined; + + const absolutePath = resolveToCwd(rawPath, cwd); + const fileName = basename(absolutePath); + if (fileName === "SKILL.md") { + return { kind: "skill", label: basename(dirname(absolutePath)) || fileName }; + } + + const docsClassification = getPiDocsClassification(absolutePath); + if (docsClassification) return docsClassification; + + if (COMPACT_RESOURCE_FILE_NAMES.has(fileName)) { + return { kind: "resource", label: formatPathRelativeToCwdOrAbsolute(absolutePath, cwd) }; + } + + return undefined; +} + +function formatCompactReadCall( + classification: CompactReadClassification, + args: ReadRenderArgs | undefined, + theme: Theme, +): string { + const expandHint = theme.fg("dim", ` (${keyText("app.tools.expand")} to expand)`); + if (classification.kind === "skill") { + return ( + theme.fg("customMessageLabel", `\x1b[1m[skill]\x1b[22m `) + + theme.fg("customMessageText", classification.label) + + formatReadLineRange(args, theme) + + expandHint + ); + } + + return ( + theme.fg("toolTitle", theme.bold(`read ${classification.kind}`)) + + " " + + theme.fg("accent", classification.label) + + formatReadLineRange(args, theme) + + expandHint + ); +} + +function formatReadResult( + args: ReadRenderArgs | undefined, + result: { content: (TextContent | ImageContent)[]; details?: ReadToolDetails }, + options: ToolRenderResultOptions, + theme: Theme, + showImages: boolean, + _cwd: string, + isError: boolean, +): string { + if (!options.expanded && !isError) { + return ""; + } + + const rawPath = str(args?.file_path ?? args?.path); + const output = getTextOutput(result, showImages); + const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; + const renderedLines = lang ? highlightCode(replaceTabs(output), lang) : output.split("\n"); + const lines = trimTrailingEmptyLines(renderedLines); + const maxLines = options.expanded ? lines.length : 10; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; + let text = `\n${displayLines.map((line) => (lang ? replaceTabs(line) : theme.fg("toolOutput", replaceTabs(line)))).join("\n")}`; + if (remaining > 0) { + text += `${theme.fg("muted", `\n... (${remaining} more lines,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`; + } + + const truncation = result.details?.truncation; + if (truncation?.truncated) { + if (truncation.firstLineExceedsLimit) { + text += `\n${theme.fg("warning", `[First line exceeds ${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit]`)}`; + } else if (truncation.truncatedBy === "lines") { + text += `\n${theme.fg("warning", `[Truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines (${truncation.maxLines ?? DEFAULT_MAX_LINES} line limit)]`)}`; + } else { + text += `\n${theme.fg("warning", `[Truncated: ${truncation.outputLines} lines shown (${formatSize(truncation.maxBytes ?? DEFAULT_MAX_BYTES)} limit)]`)}`; + } + } + return text; +} + +export function createReadToolDefinition( + cwd: string, + options?: ReadToolOptions, +): ToolDefinition { + const autoResizeImages = options?.autoResizeImages ?? true; + const ops = options?.operations ?? defaultReadOperations; + return { + name: "read", + label: "read", + description: `Read the contents of a file. Supports text files and images (jpg, png, gif, webp). Images are sent as attachments. For text files, output is truncated to ${DEFAULT_MAX_LINES} lines or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). Use offset/limit for large files. When you need the full file, continue with offset until complete.`, + promptSnippet: "Read file contents", + promptGuidelines: ["Use read to examine files instead of cat or sed."], + parameters: readSchema, + async execute( + _toolCallId, + { path, offset, limit }: { path: string; offset?: number; limit?: number }, + signal?: AbortSignal, + _onUpdate?, + ctx?, + ) { + return new Promise<{ content: (TextContent | ImageContent)[]; details: ReadToolDetails | undefined }>( + (resolve, reject) => { + if (signal?.aborted) { + reject(new Error("Operation aborted")); + return; + } + let aborted = false; + const onAbort = () => { + aborted = true; + reject(new Error("Operation aborted")); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + + (async () => { + try { + const absolutePath = await resolveReadPathAsync(path, cwd); + if (aborted) return; + // Check if file exists and is readable. + await ops.access(absolutePath); + if (aborted) return; + const mimeType = ops.detectImageMimeType ? await ops.detectImageMimeType(absolutePath) : undefined; + let content: (TextContent | ImageContent)[]; + let details: ReadToolDetails | undefined; + const nonVisionImageNote = getNonVisionImageNote(ctx?.model); + if (mimeType) { + // Read image as binary. + const buffer = await ops.readFile(absolutePath); + if (autoResizeImages) { + // Resize image if needed before sending it back to the model. + const resized = await resizeImage(buffer, mimeType); + if (!resized) { + let textNote = `Read image file [${mimeType}]\n[Image omitted: could not be resized below the inline image size limit.]`; + if (nonVisionImageNote) textNote += `\n${nonVisionImageNote}`; + content = [{ type: "text", text: textNote }]; + } else { + const dimensionNote = formatDimensionNote(resized); + let textNote = `Read image file [${resized.mimeType}]`; + if (dimensionNote) textNote += `\n${dimensionNote}`; + if (nonVisionImageNote) textNote += `\n${nonVisionImageNote}`; + content = [ + { type: "text", text: textNote }, + { type: "image", data: resized.data, mimeType: resized.mimeType }, + ]; + } + } else { + let textNote = `Read image file [${mimeType}]`; + if (nonVisionImageNote) textNote += `\n${nonVisionImageNote}`; + content = [ + { type: "text", text: textNote }, + { type: "image", data: buffer.toString("base64"), mimeType }, + ]; + } + } else { + // Read text content. + const buffer = await ops.readFile(absolutePath); + const textContent = buffer.toString("utf-8"); + const allLines = textContent.split("\n"); + const totalFileLines = allLines.length; + // Apply offset if specified. Convert from 1-indexed input to 0-indexed array access. + const startLine = offset ? Math.max(0, offset - 1) : 0; + const startLineDisplay = startLine + 1; + // Check if offset is out of bounds. + if (startLine >= allLines.length) { + throw new Error(`Offset ${offset} is beyond end of file (${allLines.length} lines total)`); + } + let selectedContent: string; + let userLimitedLines: number | undefined; + // If limit is specified by the user, honor it first. Otherwise truncateHead decides. + if (limit !== undefined) { + const endLine = Math.min(startLine + limit, allLines.length); + selectedContent = allLines.slice(startLine, endLine).join("\n"); + userLimitedLines = endLine - startLine; + } else { + selectedContent = allLines.slice(startLine).join("\n"); + } + // Apply truncation, respecting both line and byte limits. + const truncation = truncateHead(selectedContent); + let outputText: string; + if (truncation.firstLineExceedsLimit) { + // First line alone exceeds the byte limit. Point the model at a bash fallback. + const firstLineSize = formatSize(Buffer.byteLength(allLines[startLine], "utf-8")); + outputText = `[Line ${startLineDisplay} is ${firstLineSize}, exceeds ${formatSize(DEFAULT_MAX_BYTES)} limit. Use bash: sed -n '${startLineDisplay}p' ${path} | head -c ${DEFAULT_MAX_BYTES}]`; + details = { truncation }; + } else if (truncation.truncated) { + // Truncation occurred. Build an actionable continuation notice. + const endLineDisplay = startLineDisplay + truncation.outputLines - 1; + const nextOffset = endLineDisplay + 1; + outputText = truncation.content; + if (truncation.truncatedBy === "lines") { + outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines}. Use offset=${nextOffset} to continue.]`; + } else { + outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Use offset=${nextOffset} to continue.]`; + } + details = { truncation }; + } else if (userLimitedLines !== undefined && startLine + userLimitedLines < allLines.length) { + // User-specified limit stopped early, but the file still has more content. + const remaining = allLines.length - (startLine + userLimitedLines); + const nextOffset = startLine + userLimitedLines + 1; + outputText = `${truncation.content}\n\n[${remaining} more lines in file. Use offset=${nextOffset} to continue.]`; + } else { + // No truncation and no remaining user-limited content. + outputText = truncation.content; + } + content = [{ type: "text", text: outputText }]; + } + + if (aborted) return; + signal?.removeEventListener("abort", onAbort); + resolve({ content, details }); + } catch (error: any) { + signal?.removeEventListener("abort", onAbort); + if (!aborted) reject(error); + } + })(); + }, + ); + }, + renderCall(args, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + const classification = !context.expanded ? getCompactReadClassification(args, context.cwd) : undefined; + text.setText( + classification + ? formatCompactReadCall(classification, args, theme) + : formatReadCall(args, theme, context.cwd), + ); + return text; + }, + renderResult(result, options, theme, context) { + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText( + formatReadResult(context.args, result, options, theme, context.showImages, context.cwd, context.isError), + ); + return text; + }, + }; +} + +export function createReadTool(cwd: string, options?: ReadToolOptions): AgentTool { + return wrapToolDefinition(createReadToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/render-utils.ts b/cactus-code/packages/coding-agent/src/core/tools/render-utils.ts new file mode 100644 index 000000000..48a5b0af0 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/render-utils.ts @@ -0,0 +1,85 @@ +import * as os from "node:os"; +import { pathToFileURL } from "node:url"; +import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import { getCapabilities, getImageDimensions, hyperlink, imageFallback } from "@earendil-works/pi-tui"; +import type { Theme } from "../../modes/interactive/theme/theme.ts"; +import { stripAnsi } from "../../utils/ansi.ts"; +import { resolvePath } from "../../utils/paths.ts"; +import { sanitizeBinaryOutput } from "../../utils/shell.ts"; + +export function shortenPath(path: unknown): string { + if (typeof path !== "string") return ""; + const home = os.homedir(); + if (path.startsWith(home)) { + return `~${path.slice(home.length)}`; + } + return path; +} + +export function linkPath(styledText: string, rawPath: string, cwd: string): string { + if (!getCapabilities().hyperlinks) return styledText; + const absolutePath = resolvePath(rawPath, cwd); + return hyperlink(styledText, pathToFileURL(absolutePath).href); +} + +export function str(value: unknown): string | null { + if (typeof value === "string") return value; + if (value == null) return ""; + return null; +} + +export function replaceTabs(text: string): string { + return text.replace(/\t/g, " "); +} + +export function normalizeDisplayText(text: string): string { + return text.replace(/\r/g, ""); +} + +export function getTextOutput( + result: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }> } | undefined, + showImages: boolean, +): string { + if (!result) return ""; + + const textBlocks = result.content.filter((c) => c.type === "text"); + const imageBlocks = result.content.filter((c) => c.type === "image"); + + let output = textBlocks.map((c) => sanitizeBinaryOutput(stripAnsi(c.text || "")).replace(/\r/g, "")).join("\n"); + + const caps = getCapabilities(); + if (imageBlocks.length > 0 && (!caps.images || !showImages)) { + const imageIndicators = imageBlocks + .map((img) => { + const mimeType = img.mimeType ?? "image/unknown"; + const dims = + img.data && img.mimeType ? (getImageDimensions(img.data, img.mimeType) ?? undefined) : undefined; + return imageFallback(mimeType, dims); + }) + .join("\n"); + output = output ? `${output}\n${imageIndicators}` : imageIndicators; + } + + return output; +} + +export type ToolRenderResultLike = { + content: (TextContent | ImageContent)[]; + details: TDetails; +}; + +export function invalidArgText(theme: Theme): string { + return theme.fg("error", "[invalid arg]"); +} + +export function renderToolPath( + rawPath: string | null, + theme: Theme, + cwd: string, + options?: { emptyFallback?: string }, +): string { + if (rawPath === null) return invalidArgText(theme); + const value = rawPath || options?.emptyFallback; + if (!value) return theme.fg("toolOutput", "..."); + return linkPath(theme.fg("accent", shortenPath(value)), value, cwd); +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts b/cactus-code/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts new file mode 100644 index 000000000..c9da525ad --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/tool-definition-wrapper.ts @@ -0,0 +1,45 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import type { ExtensionContext, ToolDefinition } from "../extensions/types.ts"; + +/** Wrap a ToolDefinition into an AgentTool for the core runtime. */ +export function wrapToolDefinition( + definition: ToolDefinition, + ctxFactory?: () => ExtensionContext, +): AgentTool { + return { + name: definition.name, + label: definition.label, + description: definition.description, + parameters: definition.parameters, + prepareArguments: definition.prepareArguments, + executionMode: definition.executionMode, + execute: (toolCallId, params, signal, onUpdate) => + definition.execute(toolCallId, params, signal, onUpdate, ctxFactory?.() as ExtensionContext), + }; +} + +/** Wrap multiple ToolDefinitions into AgentTools for the core runtime. */ +export function wrapToolDefinitions( + definitions: ToolDefinition[], + ctxFactory?: () => ExtensionContext, +): AgentTool[] { + return definitions.map((definition) => wrapToolDefinition(definition, ctxFactory)); +} + +/** + * Synthesize a minimal ToolDefinition from an AgentTool. + * + * This keeps AgentSession's internal registry definition-first even when a caller + * provides plain AgentTool overrides that do not include prompt metadata or renderers. + */ +export function createToolDefinitionFromAgentTool(tool: AgentTool): ToolDefinition { + return { + name: tool.name, + label: tool.label, + description: tool.description, + parameters: tool.parameters as any, + prepareArguments: tool.prepareArguments, + executionMode: tool.executionMode, + execute: async (toolCallId, params, signal, onUpdate) => tool.execute(toolCallId, params, signal, onUpdate), + }; +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/truncate.ts b/cactus-code/packages/coding-agent/src/core/tools/truncate.ts new file mode 100644 index 000000000..c638ee480 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/truncate.ts @@ -0,0 +1,276 @@ +/** + * Shared truncation utilities for tool outputs. + * + * Truncation is based on two independent limits - whichever is hit first wins: + * - Line limit (default: 2000 lines) + * - Byte limit (default: 50KB) + * + * Never returns partial lines (except bash tail truncation edge case). + */ + +export const DEFAULT_MAX_LINES = 2000; +export const DEFAULT_MAX_BYTES = 50 * 1024; // 50KB +export const GREP_MAX_LINE_LENGTH = 500; // Max chars per grep match line + +export interface TruncationResult { + /** The truncated content */ + content: string; + /** Whether truncation occurred */ + truncated: boolean; + /** Which limit was hit: "lines", "bytes", or null if not truncated */ + truncatedBy: "lines" | "bytes" | null; + /** Total number of lines in the original content */ + totalLines: number; + /** Total number of bytes in the original content */ + totalBytes: number; + /** Number of complete lines in the truncated output */ + outputLines: number; + /** Number of bytes in the truncated output */ + outputBytes: number; + /** Whether the last line was partially truncated (only for tail truncation edge case) */ + lastLinePartial: boolean; + /** Whether the first line exceeded the byte limit (for head truncation) */ + firstLineExceedsLimit: boolean; + /** The max lines limit that was applied */ + maxLines: number; + /** The max bytes limit that was applied */ + maxBytes: number; +} + +export interface TruncationOptions { + /** Maximum number of lines (default: 2000) */ + maxLines?: number; + /** Maximum number of bytes (default: 50KB) */ + maxBytes?: number; +} + +function splitLinesForCounting(content: string): string[] { + if (content.length === 0) { + return []; + } + const lines = content.split("\n"); + if (content.endsWith("\n")) { + lines.pop(); + } + return lines; +} + +/** + * Format bytes as human-readable size. + */ +export function formatSize(bytes: number): string { + if (bytes < 1024) { + return `${bytes}B`; + } else if (bytes < 1024 * 1024) { + return `${(bytes / 1024).toFixed(1)}KB`; + } else { + return `${(bytes / (1024 * 1024)).toFixed(1)}MB`; + } +} + +/** + * Truncate content from the head (keep first N lines/bytes). + * Suitable for file reads where you want to see the beginning. + * + * Never returns partial lines. If first line exceeds byte limit, + * returns empty content with firstLineExceedsLimit=true. + */ +export function truncateHead(content: string, options: TruncationOptions = {}): TruncationResult { + const maxLines = options.maxLines ?? DEFAULT_MAX_LINES; + const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + + const totalBytes = Buffer.byteLength(content, "utf-8"); + const lines = splitLinesForCounting(content); + const totalLines = lines.length; + + // Check if no truncation needed + if (totalLines <= maxLines && totalBytes <= maxBytes) { + return { + content, + truncated: false, + truncatedBy: null, + totalLines, + totalBytes, + outputLines: totalLines, + outputBytes: totalBytes, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; + } + + // Check if first line alone exceeds byte limit + const firstLineBytes = Buffer.byteLength(lines[0], "utf-8"); + if (firstLineBytes > maxBytes) { + return { + content: "", + truncated: true, + truncatedBy: "bytes", + totalLines, + totalBytes, + outputLines: 0, + outputBytes: 0, + lastLinePartial: false, + firstLineExceedsLimit: true, + maxLines, + maxBytes, + }; + } + + // Collect complete lines that fit + const outputLinesArr: string[] = []; + let outputBytesCount = 0; + let truncatedBy: "lines" | "bytes" = "lines"; + + for (let i = 0; i < lines.length && i < maxLines; i++) { + const line = lines[i]; + const lineBytes = Buffer.byteLength(line, "utf-8") + (i > 0 ? 1 : 0); // +1 for newline + + if (outputBytesCount + lineBytes > maxBytes) { + truncatedBy = "bytes"; + break; + } + + outputLinesArr.push(line); + outputBytesCount += lineBytes; + } + + // If we exited due to line limit + if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) { + truncatedBy = "lines"; + } + + const outputContent = outputLinesArr.join("\n"); + const finalOutputBytes = Buffer.byteLength(outputContent, "utf-8"); + + return { + content: outputContent, + truncated: true, + truncatedBy, + totalLines, + totalBytes, + outputLines: outputLinesArr.length, + outputBytes: finalOutputBytes, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; +} + +/** + * Truncate content from the tail (keep last N lines/bytes). + * Suitable for bash output where you want to see the end (errors, final results). + * + * May return partial first line if the last line of original content exceeds byte limit. + */ +export function truncateTail(content: string, options: TruncationOptions = {}): TruncationResult { + const maxLines = options.maxLines ?? DEFAULT_MAX_LINES; + const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + + const totalBytes = Buffer.byteLength(content, "utf-8"); + const lines = splitLinesForCounting(content); + const totalLines = lines.length; + + // Check if no truncation needed + if (totalLines <= maxLines && totalBytes <= maxBytes) { + return { + content, + truncated: false, + truncatedBy: null, + totalLines, + totalBytes, + outputLines: totalLines, + outputBytes: totalBytes, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; + } + + // Work backwards from the end + const outputLinesArr: string[] = []; + let outputBytesCount = 0; + let truncatedBy: "lines" | "bytes" = "lines"; + let lastLinePartial = false; + + for (let i = lines.length - 1; i >= 0 && outputLinesArr.length < maxLines; i--) { + const line = lines[i]; + const lineBytes = Buffer.byteLength(line, "utf-8") + (outputLinesArr.length > 0 ? 1 : 0); // +1 for newline + + if (outputBytesCount + lineBytes > maxBytes) { + truncatedBy = "bytes"; + // Edge case: if we haven't added ANY lines yet and this line exceeds maxBytes, + // take the end of the line (partial) + if (outputLinesArr.length === 0) { + const truncatedLine = truncateStringToBytesFromEnd(line, maxBytes); + outputLinesArr.unshift(truncatedLine); + outputBytesCount = Buffer.byteLength(truncatedLine, "utf-8"); + lastLinePartial = true; + } + break; + } + + outputLinesArr.unshift(line); + outputBytesCount += lineBytes; + } + + // If we exited due to line limit + if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) { + truncatedBy = "lines"; + } + + const outputContent = outputLinesArr.join("\n"); + const finalOutputBytes = Buffer.byteLength(outputContent, "utf-8"); + + return { + content: outputContent, + truncated: true, + truncatedBy, + totalLines, + totalBytes, + outputLines: outputLinesArr.length, + outputBytes: finalOutputBytes, + lastLinePartial, + firstLineExceedsLimit: false, + maxLines, + maxBytes, + }; +} + +/** + * Truncate a string to fit within a byte limit (from the end). + * Handles multi-byte UTF-8 characters correctly. + */ +function truncateStringToBytesFromEnd(str: string, maxBytes: number): string { + const buf = Buffer.from(str, "utf-8"); + if (buf.length <= maxBytes) { + return str; + } + + // Start from the end, skip maxBytes back + let start = buf.length - maxBytes; + + // Find a valid UTF-8 boundary (start of a character) + while (start < buf.length && (buf[start] & 0xc0) === 0x80) { + start++; + } + + return buf.slice(start).toString("utf-8"); +} + +/** + * Truncate a single line to max characters, adding [truncated] suffix. + * Used for grep match lines. + */ +export function truncateLine( + line: string, + maxChars: number = GREP_MAX_LINE_LENGTH, +): { text: string; wasTruncated: boolean } { + if (line.length <= maxChars) { + return { text: line, wasTruncated: false }; + } + return { text: `${line.slice(0, maxChars)}... [truncated]`, wasTruncated: true }; +} diff --git a/cactus-code/packages/coding-agent/src/core/tools/write.ts b/cactus-code/packages/coding-agent/src/core/tools/write.ts new file mode 100644 index 000000000..12668e61a --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/tools/write.ts @@ -0,0 +1,267 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { Container, Text } from "@earendil-works/pi-tui"; +import { mkdir as fsMkdir, writeFile as fsWriteFile } from "fs/promises"; +import { dirname } from "path"; +import { type Static, Type } from "typebox"; +import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts"; +import { getLanguageFromPath, highlightCode, type Theme } from "../../modes/interactive/theme/theme.ts"; +import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.ts"; +import { withFileMutationQueue } from "./file-mutation-queue.ts"; +import { resolveToCwd } from "./path-utils.ts"; +import { normalizeDisplayText, renderToolPath, replaceTabs, str } from "./render-utils.ts"; +import { wrapToolDefinition } from "./tool-definition-wrapper.ts"; + +const writeSchema = Type.Object({ + path: Type.String({ description: "Path to the file to write (relative or absolute)" }), + content: Type.String({ description: "Content to write to the file" }), +}); + +export type WriteToolInput = Static; + +/** + * Pluggable operations for the write tool. + * Override these to delegate file writing to remote systems (for example SSH). + */ +export interface WriteOperations { + /** Write content to a file */ + writeFile: (absolutePath: string, content: string) => Promise; + /** Create directory recursively */ + mkdir: (dir: string) => Promise; +} + +const defaultWriteOperations: WriteOperations = { + writeFile: (path, content) => fsWriteFile(path, content, "utf-8"), + mkdir: (dir) => fsMkdir(dir, { recursive: true }).then(() => {}), +}; + +export interface WriteToolOptions { + /** Custom operations for file writing. Default: local filesystem */ + operations?: WriteOperations; +} + +type WriteHighlightCache = { + rawPath: string | null; + lang: string; + rawContent: string; + normalizedLines: string[]; + highlightedLines: string[]; +}; + +class WriteCallRenderComponent extends Text { + cache?: WriteHighlightCache; + + constructor() { + super("", 0, 0); + } +} + +const WRITE_PARTIAL_FULL_HIGHLIGHT_LINES = 50; + +function highlightSingleLine(line: string, lang: string): string { + const highlighted = highlightCode(line, lang); + return highlighted[0] ?? ""; +} + +function refreshWriteHighlightPrefix(cache: WriteHighlightCache): void { + const prefixCount = Math.min(WRITE_PARTIAL_FULL_HIGHLIGHT_LINES, cache.normalizedLines.length); + if (prefixCount === 0) return; + const prefixSource = cache.normalizedLines.slice(0, prefixCount).join("\n"); + const prefixHighlighted = highlightCode(prefixSource, cache.lang); + for (let i = 0; i < prefixCount; i++) { + cache.highlightedLines[i] = + prefixHighlighted[i] ?? highlightSingleLine(cache.normalizedLines[i] ?? "", cache.lang); + } +} + +function rebuildWriteHighlightCacheFull(rawPath: string | null, fileContent: string): WriteHighlightCache | undefined { + const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; + if (!lang) return undefined; + const displayContent = normalizeDisplayText(fileContent); + const normalized = replaceTabs(displayContent); + return { + rawPath, + lang, + rawContent: fileContent, + normalizedLines: normalized.split("\n"), + highlightedLines: highlightCode(normalized, lang), + }; +} + +function updateWriteHighlightCacheIncremental( + cache: WriteHighlightCache | undefined, + rawPath: string | null, + fileContent: string, +): WriteHighlightCache | undefined { + const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; + if (!lang) return undefined; + if (!cache) return rebuildWriteHighlightCacheFull(rawPath, fileContent); + if (cache.lang !== lang || cache.rawPath !== rawPath) return rebuildWriteHighlightCacheFull(rawPath, fileContent); + if (!fileContent.startsWith(cache.rawContent)) return rebuildWriteHighlightCacheFull(rawPath, fileContent); + if (fileContent.length === cache.rawContent.length) return cache; + + const deltaRaw = fileContent.slice(cache.rawContent.length); + const deltaDisplay = normalizeDisplayText(deltaRaw); + const deltaNormalized = replaceTabs(deltaDisplay); + cache.rawContent = fileContent; + if (cache.normalizedLines.length === 0) { + cache.normalizedLines.push(""); + cache.highlightedLines.push(""); + } + + const segments = deltaNormalized.split("\n"); + const lastIndex = cache.normalizedLines.length - 1; + cache.normalizedLines[lastIndex] += segments[0]; + cache.highlightedLines[lastIndex] = highlightSingleLine(cache.normalizedLines[lastIndex], cache.lang); + for (let i = 1; i < segments.length; i++) { + cache.normalizedLines.push(segments[i]); + cache.highlightedLines.push(highlightSingleLine(segments[i], cache.lang)); + } + refreshWriteHighlightPrefix(cache); + return cache; +} + +function trimTrailingEmptyLines(lines: string[]): string[] { + let end = lines.length; + while (end > 0 && lines[end - 1] === "") { + end--; + } + return lines.slice(0, end); +} + +function formatWriteCall( + args: { path?: string; file_path?: string; content?: string } | undefined, + options: ToolRenderResultOptions, + theme: Theme, + cache: WriteHighlightCache | undefined, + cwd: string, +): string { + const rawPath = str(args?.file_path ?? args?.path); + const fileContent = str(args?.content); + const pathDisplay = renderToolPath(rawPath, theme, cwd); + let text = `${theme.fg("toolTitle", theme.bold("write"))} ${pathDisplay}`; + + if (fileContent === null) { + text += `\n\n${theme.fg("error", "[invalid content arg - expected string]")}`; + } else if (fileContent) { + const lang = rawPath ? getLanguageFromPath(rawPath) : undefined; + const renderedLines = lang + ? (cache?.highlightedLines ?? highlightCode(replaceTabs(normalizeDisplayText(fileContent)), lang)) + : normalizeDisplayText(fileContent).split("\n"); + const lines = trimTrailingEmptyLines(renderedLines); + const totalLines = lines.length; + const maxLines = options.expanded ? lines.length : 10; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; + text += `\n\n${displayLines.map((line) => (lang ? line : theme.fg("toolOutput", replaceTabs(line)))).join("\n")}`; + if (remaining > 0) { + text += `${theme.fg("muted", `\n... (${remaining} more lines, ${totalLines} total,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`; + } + } + + return text; +} + +function formatWriteResult( + result: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; isError?: boolean }, + theme: Theme, +): string | undefined { + if (!result.isError) { + return undefined; + } + const output = result.content + .filter((c) => c.type === "text") + .map((c) => c.text || "") + .join("\n"); + if (!output) { + return undefined; + } + return `\n${theme.fg("error", output)}`; +} + +export function createWriteToolDefinition( + cwd: string, + options?: WriteToolOptions, +): ToolDefinition { + const ops = options?.operations ?? defaultWriteOperations; + return { + name: "write", + label: "write", + description: + "Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Automatically creates parent directories.", + promptSnippet: "Create or overwrite files", + promptGuidelines: ["Use write only for new files or complete rewrites."], + parameters: writeSchema, + async execute( + _toolCallId, + { path, content }: { path: string; content: string }, + signal?: AbortSignal, + _onUpdate?, + _ctx?, + ) { + const absolutePath = resolveToCwd(path, cwd); + const dir = dirname(absolutePath); + return withFileMutationQueue(absolutePath, async () => { + // Do not reject from an abort event listener here: that would release the + // mutation queue while an in-flight filesystem operation may still finish. + // Checking signal.aborted after each await observes the same aborts while + // keeping the queue locked until the current operation has settled. + const throwIfAborted = (): void => { + if (signal?.aborted) throw new Error("Operation aborted"); + }; + + throwIfAborted(); + // Create parent directories if needed. + await ops.mkdir(dir); + throwIfAborted(); + + // Write the file contents. + await ops.writeFile(absolutePath, content); + throwIfAborted(); + + return { + content: [{ type: "text", text: `Successfully wrote ${content.length} bytes to ${path}` }], + details: undefined, + }; + }); + }, + renderCall(args, theme, context) { + const renderArgs = args as { path?: string; file_path?: string; content?: string } | undefined; + const rawPath = str(renderArgs?.file_path ?? renderArgs?.path); + const fileContent = str(renderArgs?.content); + const component = + (context.lastComponent as WriteCallRenderComponent | undefined) ?? new WriteCallRenderComponent(); + if (fileContent !== null) { + component.cache = context.argsComplete + ? rebuildWriteHighlightCacheFull(rawPath, fileContent) + : updateWriteHighlightCacheIncremental(component.cache, rawPath, fileContent); + } else { + component.cache = undefined; + } + component.setText( + formatWriteCall( + renderArgs, + { expanded: context.expanded, isPartial: context.isPartial }, + theme, + component.cache, + context.cwd, + ), + ); + return component; + }, + renderResult(result, _options, theme, context) { + const output = formatWriteResult({ ...result, isError: context.isError }, theme); + if (!output) { + const component = (context.lastComponent as Container | undefined) ?? new Container(); + component.clear(); + return component; + } + const text = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0); + text.setText(output); + return text; + }, + }; +} + +export function createWriteTool(cwd: string, options?: WriteToolOptions): AgentTool { + return wrapToolDefinition(createWriteToolDefinition(cwd, options)); +} diff --git a/cactus-code/packages/coding-agent/src/core/trust-manager.ts b/cactus-code/packages/coding-agent/src/core/trust-manager.ts new file mode 100644 index 000000000..9c494b47a --- /dev/null +++ b/cactus-code/packages/coding-agent/src/core/trust-manager.ts @@ -0,0 +1,244 @@ +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { homedir } from "node:os"; +import { dirname, join } from "node:path"; +import lockfile from "proper-lockfile"; +import { CONFIG_DIR_NAME } from "../config.ts"; +import { canonicalizePath, resolvePath } from "../utils/paths.ts"; + +export type ProjectTrustDecision = boolean | null; + +export interface ProjectTrustStoreEntry { + path: string; + decision: boolean; +} + +export interface ProjectTrustUpdate { + path: string; + decision: ProjectTrustDecision; +} + +export interface ProjectTrustOption { + label: string; + trusted: boolean; + updates: ProjectTrustUpdate[]; + savedPath?: string; +} + +type TrustFile = Record; + +const TRUST_REQUIRING_PROJECT_CONFIG_RESOURCES = [ + "settings.json", + "extensions", + "skills", + "prompts", + "themes", + "SYSTEM.md", + "APPEND_SYSTEM.md", +] as const; + +function normalizeCwd(cwd: string): string { + return canonicalizePath(resolvePath(cwd)); +} + +function findNearestTrustEntry(data: TrustFile, cwd: string): ProjectTrustStoreEntry | null { + let currentDir = normalizeCwd(cwd); + while (true) { + const value = data[currentDir]; + if (value === true || value === false) { + return { path: currentDir, decision: value }; + } + + const parentDir = dirname(currentDir); + if (parentDir === currentDir) { + return null; + } + currentDir = parentDir; + } +} + +export function getProjectTrustParentPath(cwd: string): string | undefined { + const trustPath = normalizeCwd(cwd); + const parentDir = dirname(trustPath); + return parentDir === trustPath ? undefined : parentDir; +} + +export function getProjectTrustOptions(cwd: string, options?: { includeSessionOnly?: boolean }): ProjectTrustOption[] { + const trustPath = normalizeCwd(cwd); + const trustOptions: ProjectTrustOption[] = [ + { label: "Trust", trusted: true, updates: [{ path: trustPath, decision: true }], savedPath: trustPath }, + ]; + const parentPath = getProjectTrustParentPath(cwd); + if (parentPath !== undefined) { + trustOptions.push({ + label: `Trust parent folder (${parentPath})`, + trusted: true, + updates: [ + { path: parentPath, decision: true }, + { path: trustPath, decision: null }, + ], + savedPath: parentPath, + }); + } + if (options?.includeSessionOnly) { + trustOptions.push({ label: "Trust (this session only)", trusted: true, updates: [] }); + } + trustOptions.push({ + label: "Do not trust", + trusted: false, + updates: [{ path: trustPath, decision: false }], + savedPath: trustPath, + }); + if (options?.includeSessionOnly) { + trustOptions.push({ label: "Do not trust (this session only)", trusted: false, updates: [] }); + } + return trustOptions; +} + +function readTrustFile(path: string): TrustFile { + if (!existsSync(path)) { + return {}; + } + + let parsed: unknown; + try { + parsed = JSON.parse(readFileSync(path, "utf-8")); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + throw new Error(`Failed to read trust store ${path}: ${message}`); + } + + if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) { + throw new Error(`Invalid trust store ${path}: expected an object`); + } + + const data: TrustFile = {}; + for (const [key, value] of Object.entries(parsed)) { + if (value !== true && value !== false && value !== null) { + throw new Error(`Invalid trust store ${path}: value for ${JSON.stringify(key)} must be true, false, or null`); + } + data[key] = value; + } + return data; +} + +function writeTrustFile(path: string, data: TrustFile): void { + const sorted: TrustFile = {}; + for (const key of Object.keys(data).sort()) { + const value = data[key]; + if (value === true || value === false || value === null) { + sorted[key] = value; + } + } + mkdirSync(dirname(path), { recursive: true }); + writeFileSync(path, `${JSON.stringify(sorted, null, 2)}\n`, "utf-8"); +} + +function acquireTrustLockSync(path: string): () => void { + const trustDir = dirname(path); + mkdirSync(trustDir, { recursive: true }); + const maxAttempts = 10; + const delayMs = 20; + let lastError: unknown; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { + try { + return lockfile.lockSync(trustDir, { realpath: false, lockfilePath: `${path}.lock` }); + } catch (error) { + const code = + typeof error === "object" && error !== null && "code" in error + ? String((error as { code?: unknown }).code) + : undefined; + if (code !== "ELOCKED" || attempt === maxAttempts) { + throw error; + } + lastError = error; + const start = Date.now(); + while (Date.now() - start < delayMs) { + // Sleep synchronously to avoid changing trust store callers to async. + } + } + } + + if (lastError instanceof Error) { + throw lastError; + } + throw new Error("Failed to acquire trust store lock"); +} + +function withTrustFileLock(path: string, fn: () => T): T { + const release = acquireTrustLockSync(path); + try { + return fn(); + } finally { + release(); + } +} + +/** + * Returns true when cwd has project-local resources that must be gated by + * project trust: trust-requiring entries under cwd/.pi, or .agents/skills in + * cwd or one of its ancestors. Returns false when no such project resources + * exist. The user/global ~/.agents/skills directory is always treated as a + * trusted user resource and is ignored here, even when cwd is $HOME. + */ +export function hasTrustRequiringProjectResources(cwd: string): boolean { + const homeDir = canonicalizePath(resolvePath(process.env.HOME || homedir())); + const userAgentsSkillsDir = join(homeDir, ".agents", "skills"); + let currentDir = canonicalizePath(resolvePath(cwd)); + + const configDir = join(currentDir, CONFIG_DIR_NAME); + if (TRUST_REQUIRING_PROJECT_CONFIG_RESOURCES.some((entry) => existsSync(join(configDir, entry)))) { + return true; + } + + while (true) { + const agentsSkillsDir = join(currentDir, ".agents", "skills"); + if (agentsSkillsDir !== userAgentsSkillsDir && existsSync(agentsSkillsDir)) { + return true; + } + + const parentDir = dirname(currentDir); + if (parentDir === currentDir) { + return false; + } + currentDir = parentDir; + } +} + +export class ProjectTrustStore { + private trustPath: string; + + constructor(agentDir: string) { + this.trustPath = join(resolvePath(agentDir), "trust.json"); + } + + get(cwd: string): ProjectTrustDecision { + return this.getEntry(cwd)?.decision ?? null; + } + + getEntry(cwd: string): ProjectTrustStoreEntry | null { + return withTrustFileLock(this.trustPath, () => { + const data = readTrustFile(this.trustPath); + return findNearestTrustEntry(data, cwd); + }); + } + + set(cwd: string, decision: ProjectTrustDecision): void { + this.setMany([{ path: cwd, decision }]); + } + + setMany(decisions: ProjectTrustUpdate[]): void { + withTrustFileLock(this.trustPath, () => { + const data = readTrustFile(this.trustPath); + for (const { path, decision } of decisions) { + const key = normalizeCwd(path); + if (decision === null) { + delete data[key]; + } else { + data[key] = decision; + } + } + writeTrustFile(this.trustPath, data); + }); + } +} diff --git a/cactus-code/packages/coding-agent/src/index.ts b/cactus-code/packages/coding-agent/src/index.ts new file mode 100644 index 000000000..2540dbcc3 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/index.ts @@ -0,0 +1,368 @@ +// Core session management + +export { type Args, parseArgs } from "./cli/args.ts"; + +// Config paths +export { + CONFIG_DIR_NAME, + getAgentDir, + getDocsPath, + getExamplesPath, + getPackageDir, + getReadmePath, + VERSION, +} from "./config.ts"; +export { + AgentSession, + type AgentSessionConfig, + type AgentSessionEvent, + type AgentSessionEventListener, + type ModelCycleResult, + type ParsedSkillBlock, + type PromptOptions, + parseSkillBlock, + type SessionStats, +} from "./core/agent-session.ts"; +// Auth and model registry +export { + type ApiKeyCredential, + type AuthCredential, + type AuthStatus, + AuthStorage, + type AuthStorageBackend, + FileAuthStorageBackend, + InMemoryAuthStorageBackend, + type OAuthCredential, +} from "./core/auth-storage.ts"; +// Compaction +export { + type BranchPreparation, + type BranchSummaryResult, + type CollectEntriesResult, + type CompactionResult, + type CutPointResult, + calculateContextTokens, + collectEntriesForBranchSummary, + compact, + DEFAULT_COMPACTION_SETTINGS, + estimateTokens, + type FileOperations, + findCutPoint, + findTurnStartIndex, + type GenerateBranchSummaryOptions, + generateBranchSummary, + generateSummary, + getLastAssistantUsage, + prepareBranchEntries, + serializeConversation, + shouldCompact, +} from "./core/compaction/index.ts"; +export { createEventBus, type EventBus, type EventBusController } from "./core/event-bus.ts"; +// Extension system +export type { + AgentEndEvent, + AgentStartEvent, + AgentToolResult, + AgentToolUpdateCallback, + AppKeybinding, + AutocompleteProviderFactory, + BashToolCallEvent, + BeforeAgentStartEvent, + BeforeAgentStartEventResult, + BeforeProviderRequestEvent, + BeforeProviderRequestEventResult, + BuildSystemPromptOptions, + CompactOptions, + ContextEvent, + ContextUsage, + CustomToolCallEvent, + EditToolCallEvent, + ExecOptions, + ExecResult, + Extension, + ExtensionActions, + ExtensionAPI, + ExtensionCommandContext, + ExtensionCommandContextActions, + ExtensionContext, + ExtensionContextActions, + ExtensionError, + ExtensionEvent, + ExtensionFactory, + ExtensionFlag, + ExtensionHandler, + ExtensionRuntime, + ExtensionShortcut, + ExtensionUIContext, + ExtensionUIDialogOptions, + ExtensionWidgetOptions, + FindToolCallEvent, + GrepToolCallEvent, + InputEvent, + InputEventResult, + InputSource, + KeybindingsManager, + LoadExtensionsResult, + LsToolCallEvent, + MessageRenderer, + MessageRenderOptions, + ProjectTrustContext, + ProjectTrustEvent, + ProjectTrustEventDecision, + ProjectTrustEventResult, + ProjectTrustHandler, + ProviderConfig, + ProviderModelConfig, + ReadToolCallEvent, + RegisteredCommand, + RegisteredTool, + ResolvedCommand, + SessionBeforeCompactEvent, + SessionBeforeForkEvent, + SessionBeforeSwitchEvent, + SessionBeforeTreeEvent, + SessionCompactEvent, + SessionShutdownEvent, + SessionStartEvent, + SessionTreeEvent, + SlashCommandInfo, + SlashCommandSource, + SourceInfo, + TerminalInputHandler, + ToolCallEvent, + ToolCallEventResult, + ToolDefinition, + ToolExecutionMode, + ToolInfo, + ToolRenderResultOptions, + ToolResultEvent, + TurnEndEvent, + TurnStartEvent, + UserBashEvent, + UserBashEventResult, + WidgetPlacement, + WorkingIndicatorOptions, + WriteToolCallEvent, +} from "./core/extensions/index.ts"; +export { + createExtensionRuntime, + defineTool, + discoverAndLoadExtensions, + ExtensionRunner, + isBashToolResult, + isEditToolResult, + isFindToolResult, + isGrepToolResult, + isLsToolResult, + isReadToolResult, + isToolCallEventType, + isWriteToolResult, + wrapRegisteredTool, + wrapRegisteredTools, +} from "./core/extensions/index.ts"; +// Footer data provider (git branch + extension statuses - data not otherwise available to extensions) +export type { ReadonlyFooterDataProvider } from "./core/footer-data-provider.ts"; +export { convertToLlm } from "./core/messages.ts"; +export { ModelRegistry } from "./core/model-registry.ts"; +export type { + PackageManager, + PathMetadata, + ProgressCallback, + ProgressEvent, + ResolvedPaths, + ResolvedResource, +} from "./core/package-manager.ts"; +export { DefaultPackageManager } from "./core/package-manager.ts"; +export type { ResourceCollision, ResourceDiagnostic, ResourceLoader } from "./core/resource-loader.ts"; +export { DefaultResourceLoader, loadProjectContextFiles } from "./core/resource-loader.ts"; +// SDK for programmatic usage +export { + AgentSessionRuntime, + type AgentSessionRuntimeDiagnostic, + type AgentSessionServices, + type CreateAgentSessionFromServicesOptions, + type CreateAgentSessionOptions, + type CreateAgentSessionResult, + type CreateAgentSessionRuntimeFactory, + type CreateAgentSessionRuntimeResult, + type CreateAgentSessionServicesOptions, + // Factory + createAgentSession, + createAgentSessionFromServices, + createAgentSessionRuntime, + createAgentSessionServices, + createBashTool, + // Tool factories (for custom cwd) + createCodingTools, + createEditTool, + createFindTool, + createGrepTool, + createLsTool, + createReadOnlyTools, + createReadTool, + createWriteTool, + type PromptTemplate, +} from "./core/sdk.ts"; +export { + type BranchSummaryEntry, + buildSessionContext, + type CompactionEntry, + CURRENT_SESSION_VERSION, + type CustomEntry, + type CustomMessageEntry, + type FileEntry, + getLatestCompactionEntry, + type ModelChangeEntry, + migrateSessionEntries, + type NewSessionOptions, + parseSessionEntries, + type SessionContext, + type SessionEntry, + type SessionEntryBase, + type SessionHeader, + type SessionInfo, + type SessionInfoEntry, + SessionManager, + type SessionMessageEntry, + type ThinkingLevelChangeEntry, +} from "./core/session-manager.ts"; +export { + type CompactionSettings, + type DefaultProjectTrust, + type ImageSettings, + type PackageSource, + type RetrySettings, + SettingsManager, + type SettingsManagerCreateOptions, +} from "./core/settings-manager.ts"; +// Skills +export { + formatSkillsForPrompt, + type LoadSkillsFromDirOptions, + type LoadSkillsResult, + loadSkills, + loadSkillsFromDir, + type Skill, + type SkillFrontmatter, +} from "./core/skills.ts"; +export { createSyntheticSourceInfo } from "./core/source-info.ts"; +export { type EditDiffResult, generateDiffString, generateUnifiedPatch } from "./core/tools/edit-diff.ts"; +// Tools +export { + type BashOperations, + type BashSpawnContext, + type BashSpawnHook, + type BashToolDetails, + type BashToolInput, + type BashToolOptions, + createBashToolDefinition, + createEditToolDefinition, + createFindToolDefinition, + createGrepToolDefinition, + createLocalBashOperations, + createLsToolDefinition, + createReadToolDefinition, + createWriteToolDefinition, + DEFAULT_MAX_BYTES, + DEFAULT_MAX_LINES, + type EditOperations, + type EditToolDetails, + type EditToolInput, + type EditToolOptions, + type FindOperations, + type FindToolDetails, + type FindToolInput, + type FindToolOptions, + formatSize, + type GrepOperations, + type GrepToolDetails, + type GrepToolInput, + type GrepToolOptions, + type LsOperations, + type LsToolDetails, + type LsToolInput, + type LsToolOptions, + type ReadOperations, + type ReadToolDetails, + type ReadToolInput, + type ReadToolOptions, + type ToolsOptions, + type TruncationOptions, + type TruncationResult, + truncateHead, + truncateLine, + truncateTail, + type WriteOperations, + type WriteToolInput, + type WriteToolOptions, + withFileMutationQueue, +} from "./core/tools/index.ts"; +export { + hasTrustRequiringProjectResources, + type ProjectTrustDecision, + ProjectTrustStore, + type ProjectTrustStoreEntry, + type ProjectTrustUpdate, +} from "./core/trust-manager.ts"; +// Main entry point +export { type MainOptions, main } from "./main.ts"; +// Run modes for programmatic SDK usage +export { + InteractiveMode, + type InteractiveModeOptions, + type PrintModeOptions, + runPrintMode, +} from "./modes/index.ts"; +// UI components for extensions +export { + AssistantMessageComponent, + BashExecutionComponent, + BorderedLoader, + BranchSummaryMessageComponent, + CompactionSummaryMessageComponent, + CustomEditor, + CustomMessageComponent, + DynamicBorder, + ExtensionEditorComponent, + ExtensionInputComponent, + ExtensionSelectorComponent, + FooterComponent, + keyHint, + keyText, + ModelSelectorComponent, + type RenderDiffOptions, + rawKeyHint, + renderDiff, + SessionSelectorComponent, + type SettingsCallbacks, + type SettingsConfig, + SettingsSelectorComponent, + SkillInvocationMessageComponent, + ThemeSelectorComponent, + ThinkingSelectorComponent, + ToolExecutionComponent, + type ToolExecutionOptions, + TreeSelectorComponent, + truncateToVisualLines, + UserMessageComponent, + UserMessageSelectorComponent, + type VisualTruncateResult, +} from "./modes/interactive/components/index.ts"; +// Theme utilities for custom tools and extensions +export { + getLanguageFromPath, + getMarkdownTheme, + getSelectListTheme, + getSettingsListTheme, + highlightCode, + initTheme, + Theme, + type ThemeColor, +} from "./modes/interactive/theme/theme.ts"; +// Clipboard utilities +export { copyToClipboard } from "./utils/clipboard.ts"; +export { parseFrontmatter, stripFrontmatter } from "./utils/frontmatter.ts"; +export { convertToPng } from "./utils/image-convert.ts"; +export { formatDimensionNote, type ResizedImage, resizeImage } from "./utils/image-resize.ts"; +// Shell utilities +export { getShellConfig } from "./utils/shell.ts"; diff --git a/cactus-code/packages/coding-agent/src/main.ts b/cactus-code/packages/coding-agent/src/main.ts new file mode 100644 index 000000000..acfb0abaa --- /dev/null +++ b/cactus-code/packages/coding-agent/src/main.ts @@ -0,0 +1,779 @@ +/** + * Main entry point for the coding agent CLI. + * + * This file handles CLI argument parsing and translates them into + * createAgentSession() options. The SDK does the heavy lifting. + */ + +import { createInterface } from "node:readline"; +import { type ImageContent, modelsAreEqual } from "@earendil-works/pi-ai"; +import chalk from "chalk"; +import { type Args, type Mode, parseArgs, printHelp } from "./cli/args.ts"; +import { processFileArguments } from "./cli/file-processor.ts"; +import { buildInitialMessage } from "./cli/initial-message.ts"; +import { listModels } from "./cli/list-models.ts"; +import { createProjectTrustContext } from "./cli/project-trust.ts"; +import { selectSession } from "./cli/session-picker.ts"; +import { shouldRunFirstTimeSetup, showFirstTimeSetup, showStartupSelector } from "./cli/startup-ui.ts"; +import { ENV_SESSION_DIR, expandTildePath, getAgentDir, VERSION } from "./config.ts"; +import { type CreateAgentSessionRuntimeFactory, createAgentSessionRuntime } from "./core/agent-session-runtime.ts"; +import { + type AgentSessionRuntimeDiagnostic, + createAgentSessionFromServices, + createAgentSessionServices, +} from "./core/agent-session-services.ts"; +import { formatNoModelsAvailableMessage } from "./core/auth-guidance.ts"; +import { AuthStorage } from "./core/auth-storage.ts"; +import type { ExtensionFactory } from "./core/extensions/types.ts"; +import { applyHttpProxySettings, configureHttpDispatcher } from "./core/http-dispatcher.ts"; +import type { ModelRegistry } from "./core/model-registry.ts"; +import { resolveCliModel, resolveModelScope, type ScopedModel } from "./core/model-resolver.ts"; +import { restoreStdout, takeOverStdout } from "./core/output-guard.ts"; +import { type AppMode, resolveProjectTrusted } from "./core/project-trust.ts"; +import type { CreateAgentSessionOptions } from "./core/sdk.ts"; +import { + formatMissingSessionCwdPrompt, + getMissingSessionCwdIssue, + MissingSessionCwdError, + type SessionCwdIssue, +} from "./core/session-cwd.ts"; +import { assertValidSessionId, SessionManager } from "./core/session-manager.ts"; +import { SettingsManager } from "./core/settings-manager.ts"; +import { printTimings, resetTimings, time } from "./core/timings.ts"; +import { hasTrustRequiringProjectResources, ProjectTrustStore } from "./core/trust-manager.ts"; +import { InteractiveMode, runPrintMode } from "./modes/index.ts"; +import { initTheme, stopThemeWatcher } from "./modes/interactive/theme/theme.ts"; +import { isLocalPath, normalizePath, resolvePath } from "./utils/paths.ts"; + + +/** + * Read all content from piped stdin. + * Returns undefined if stdin is a TTY (interactive terminal). + */ +async function readPipedStdin(): Promise { + // If stdin is a TTY, we're running interactively - don't read stdin + if (process.stdin.isTTY) { + return undefined; + } + + return new Promise((resolve) => { + let data = ""; + process.stdin.setEncoding("utf8"); + process.stdin.on("data", (chunk) => { + data += chunk; + }); + process.stdin.on("end", () => { + resolve(data.trim() || undefined); + }); + process.stdin.resume(); + }); +} + +function collectSettingsDiagnostics( + settingsManager: SettingsManager, + context: string, +): AgentSessionRuntimeDiagnostic[] { + return settingsManager.drainErrors().map(({ scope, error }) => ({ + type: "warning", + message: `(${context}, ${scope} settings) ${error.message}`, + })); +} + +function reportDiagnostics(diagnostics: readonly AgentSessionRuntimeDiagnostic[]): void { + for (const diagnostic of diagnostics) { + const color = diagnostic.type === "error" ? chalk.red : diagnostic.type === "warning" ? chalk.yellow : chalk.dim; + const prefix = diagnostic.type === "error" ? "Error: " : diagnostic.type === "warning" ? "Warning: " : ""; + console.error(color(`${prefix}${diagnostic.message}`)); + } +} + +function isTruthyEnvFlag(value: string | undefined): boolean { + if (!value) return false; + return value === "1" || value.toLowerCase() === "true" || value.toLowerCase() === "yes"; +} + +function resolveAppMode(parsed: Args, stdinIsTTY: boolean, stdoutIsTTY: boolean): AppMode { + if (parsed.mode === "json") { + return "json"; + } + if (parsed.print || !stdinIsTTY || !stdoutIsTTY) { + return "print"; + } + return "interactive"; +} + +function toPrintOutputMode(appMode: AppMode): Mode { + return appMode === "json" ? "json" : "text"; +} + +function isPlainRuntimeMetadataCommand(parsed: Args): boolean { + return !parsed.print && parsed.mode === undefined && (parsed.help === true || parsed.listModels !== undefined); +} + +async function prepareInitialMessage( + parsed: Args, + autoResizeImages: boolean, + stdinContent?: string, +): Promise<{ + initialMessage?: string; + initialImages?: ImageContent[]; +}> { + if (parsed.fileArgs.length === 0) { + return buildInitialMessage({ parsed, stdinContent }); + } + + const { text, images } = await processFileArguments(parsed.fileArgs, { autoResizeImages }); + return buildInitialMessage({ + parsed, + fileText: text, + fileImages: images, + stdinContent, + }); +} + +/** Result from resolving a session argument */ +type ResolvedSession = + | { type: "path"; path: string } // Direct file path + | { type: "local"; path: string } // Found in current project + | { type: "global"; path: string; cwd: string } // Found in different project + | { type: "not_found"; arg: string }; // Not found anywhere + +/** + * Resolve a session argument to a file path. + * If it looks like a path, use as-is. Otherwise try to match as session ID prefix. + */ +async function findLocalSessionByExactId( + sessionId: string, + cwd: string, + sessionDir?: string, +): Promise<{ type: "local"; path: string } | undefined> { + const localSessions = await SessionManager.list(cwd, sessionDir); + const localMatch = localSessions.find((s) => s.id === sessionId); + return localMatch ? { type: "local", path: localMatch.path } : undefined; +} + +async function resolveSessionPath(sessionArg: string, cwd: string, sessionDir?: string): Promise { + // If it looks like a file path, resolve it before handing it to the session manager. + if (sessionArg.includes("/") || sessionArg.includes("\\") || sessionArg.endsWith(".jsonl")) { + return { type: "path", path: resolvePath(sessionArg, cwd) }; + } + + // Try to match as session ID in current project first + const localSessions = await SessionManager.list(cwd, sessionDir); + const localMatch = + localSessions.find((s) => s.id === sessionArg) ?? localSessions.find((s) => s.id.startsWith(sessionArg)); + + if (localMatch) { + return { type: "local", path: localMatch.path }; + } + + // Try global search across all projects + const allSessions = await SessionManager.listAll(sessionDir); + const globalMatch = + allSessions.find((s) => s.id === sessionArg) ?? allSessions.find((s) => s.id.startsWith(sessionArg)); + + if (globalMatch) { + return { type: "global", path: globalMatch.path, cwd: globalMatch.cwd }; + } + + // Not found anywhere + return { type: "not_found", arg: sessionArg }; +} + +/** Prompt user for yes/no confirmation */ +async function promptConfirm(message: string): Promise { + return new Promise((resolve) => { + const rl = createInterface({ + input: process.stdin, + output: process.stdout, + }); + rl.question(`${message} [y/N] `, (answer) => { + rl.close(); + resolve(answer.toLowerCase() === "y" || answer.toLowerCase() === "yes"); + }); + }); +} + +function validateForkFlags(parsed: Args): void { + if (!parsed.fork) return; + + const conflictingFlags = [ + parsed.session ? "--session" : undefined, + parsed.continue ? "--continue" : undefined, + parsed.resume ? "--resume" : undefined, + parsed.noSession ? "--no-session" : undefined, + ].filter((flag): flag is string => flag !== undefined); + + if (conflictingFlags.length > 0) { + console.error(chalk.red(`Error: --fork cannot be combined with ${conflictingFlags.join(", ")}`)); + process.exit(1); + } +} + +function validateSessionIdFlags(parsed: Args): void { + if (parsed.sessionId === undefined) return; + + const conflictingFlags = [ + parsed.session ? "--session" : undefined, + parsed.continue ? "--continue" : undefined, + parsed.resume ? "--resume" : undefined, + parsed.noSession ? "--no-session" : undefined, + ].filter((flag): flag is string => flag !== undefined); + + if (conflictingFlags.length > 0) { + console.error(chalk.red(`Error: --session-id cannot be combined with ${conflictingFlags.join(", ")}`)); + process.exit(1); + } + + try { + assertValidSessionId(parsed.sessionId); + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + console.error(chalk.red(`Error: ${message}`)); + process.exit(1); + } +} + +function forkSessionOrExit(sourcePath: string, cwd: string, sessionDir?: string, sessionId?: string): SessionManager { + try { + return SessionManager.forkFrom(sourcePath, cwd, sessionDir, { id: sessionId }); + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + console.error(chalk.red(`Error: ${message}`)); + process.exit(1); + } +} + +async function createSessionManager( + parsed: Args, + cwd: string, + sessionDir: string | undefined, + settingsManager: SettingsManager, +): Promise { + if (parsed.noSession || parsed.help || parsed.listModels !== undefined) { + return SessionManager.inMemory(cwd); + } + + if (parsed.fork) { + if (parsed.sessionId) { + const existingTarget = await findLocalSessionByExactId(parsed.sessionId, cwd, sessionDir); + if (existingTarget) { + console.error(chalk.red(`Session already exists with id '${parsed.sessionId}'`)); + process.exit(1); + } + } + + const resolved = await resolveSessionPath(parsed.fork, cwd, sessionDir); + + switch (resolved.type) { + case "path": + case "local": + case "global": + return forkSessionOrExit(resolved.path, cwd, sessionDir, parsed.sessionId); + + case "not_found": + console.error(chalk.red(`No session found matching '${resolved.arg}'`)); + process.exit(1); + } + } + + if (parsed.session) { + const resolved = await resolveSessionPath(parsed.session, cwd, sessionDir); + + switch (resolved.type) { + case "path": + case "local": + return SessionManager.open(resolved.path, sessionDir); + + case "global": { + console.log(chalk.yellow(`Session found in different project: ${resolved.cwd}`)); + const shouldFork = await promptConfirm("Fork this session into current directory?"); + if (!shouldFork) { + console.log(chalk.dim("Aborted.")); + process.exit(0); + } + return forkSessionOrExit(resolved.path, cwd, sessionDir); + } + + case "not_found": + console.error(chalk.red(`No session found matching '${resolved.arg}'`)); + process.exit(1); + } + } + + if (parsed.resume) { + try { + const selectedPath = await selectSession( + (onProgress) => SessionManager.list(cwd, sessionDir, onProgress), + (onProgress) => SessionManager.listAll(sessionDir, onProgress), + settingsManager, + ); + if (!selectedPath) { + console.log(chalk.dim("No session selected")); + process.exit(0); + } + return SessionManager.open(selectedPath, sessionDir); + } finally { + stopThemeWatcher(); + } + } + + if (parsed.continue) { + return SessionManager.continueRecent(cwd, sessionDir); + } + + if (parsed.sessionId) { + const existingSession = await findLocalSessionByExactId(parsed.sessionId, cwd, sessionDir); + if (existingSession) { + return SessionManager.open(existingSession.path, sessionDir); + } + } + + return SessionManager.create(cwd, sessionDir, { id: parsed.sessionId }); +} + +function buildSessionOptions( + parsed: Args, + scopedModels: ScopedModel[], + hasExistingSession: boolean, + modelRegistry: ModelRegistry, + settingsManager: SettingsManager, +): { + options: CreateAgentSessionOptions; + cliThinkingFromModel: boolean; + diagnostics: AgentSessionRuntimeDiagnostic[]; +} { + const options: CreateAgentSessionOptions = {}; + const diagnostics: AgentSessionRuntimeDiagnostic[] = []; + let cliThinkingFromModel = false; + + // Model from CLI + // - supports --provider --model + // - supports --model / + if (parsed.model) { + const resolved = resolveCliModel({ + cliProvider: parsed.provider, + cliModel: parsed.model, + cliThinking: parsed.thinking, + modelRegistry, + }); + if (resolved.warning) { + diagnostics.push({ type: "warning", message: resolved.warning }); + } + if (resolved.error) { + diagnostics.push({ type: "error", message: resolved.error }); + } + if (resolved.model) { + options.model = resolved.model; + // Allow "--model :" as a shorthand. + // Explicit --thinking still takes precedence (applied later). + if (!parsed.thinking && resolved.thinkingLevel) { + options.thinkingLevel = resolved.thinkingLevel; + cliThinkingFromModel = true; + } + } + } + + if (!options.model && scopedModels.length > 0 && !hasExistingSession) { + // Check if saved default is in scoped models - use it if so, otherwise first scoped model + const savedProvider = settingsManager.getDefaultProvider(); + const savedModelId = settingsManager.getDefaultModel(); + const savedModel = savedProvider && savedModelId ? modelRegistry.find(savedProvider, savedModelId) : undefined; + const savedInScope = savedModel ? scopedModels.find((sm) => modelsAreEqual(sm.model, savedModel)) : undefined; + + if (savedInScope) { + options.model = savedInScope.model; + // Use thinking level from scoped model config if explicitly set + if (!parsed.thinking && savedInScope.thinkingLevel) { + options.thinkingLevel = savedInScope.thinkingLevel; + } + } else { + options.model = scopedModels[0].model; + // Use thinking level from first scoped model if explicitly set + if (!parsed.thinking && scopedModels[0].thinkingLevel) { + options.thinkingLevel = scopedModels[0].thinkingLevel; + } + } + } + + // Thinking level from CLI (takes precedence over scoped model thinking levels set above) + if (parsed.thinking) { + options.thinkingLevel = parsed.thinking; + } + + // Scoped models for Ctrl+P cycling + // Keep thinking level undefined when not explicitly set in the model pattern. + // Undefined means "inherit current session thinking level" during cycling. + if (scopedModels.length > 0) { + options.scopedModels = scopedModels.map((sm) => ({ + model: sm.model, + thinkingLevel: sm.thinkingLevel, + })); + } + + // API key from CLI - set in authStorage + // (handled by caller before createAgentSession) + + // Tools + if (parsed.noTools) { + options.noTools = "all"; + } else if (parsed.noBuiltinTools) { + options.noTools = "builtin"; + } + if (parsed.tools) { + options.tools = [...parsed.tools]; + } + if (parsed.excludeTools) { + options.excludeTools = [...parsed.excludeTools]; + } + + return { options, cliThinkingFromModel, diagnostics }; +} + +function resolveCliPaths(cwd: string, paths: string[] | undefined): string[] | undefined { + return paths?.map((value) => (isLocalPath(value) ? resolvePath(value, cwd) : value)); +} + +async function promptForMissingSessionCwd( + issue: SessionCwdIssue, + settingsManager: SettingsManager, +): Promise { + return showStartupSelector(settingsManager, formatMissingSessionCwdPrompt(issue), [ + { label: "Continue", value: issue.fallbackCwd }, + { label: "Cancel", value: undefined }, + ]); +} + +export interface MainOptions { + extensionFactories?: ExtensionFactory[]; +} + +export async function main(args: string[], options?: MainOptions) { + resetTimings(); + if (args[0] === "code") { + args = args.slice(1); + } + const offlineMode = args.includes("--offline") || isTruthyEnvFlag(process.env.PI_OFFLINE); + if (offlineMode) { + process.env.PI_OFFLINE = "1"; + process.env.PI_SKIP_VERSION_CHECK = "1"; + } + + const cwd = process.cwd(); + const agentDir = getAgentDir(); + const bootstrapSettingsManager = SettingsManager.create(cwd, agentDir, { projectTrusted: false }); + applyHttpProxySettings(bootstrapSettingsManager.getGlobalSettings().httpProxy); + configureHttpDispatcher(); + + const parsed = parseArgs(args); + if (parsed.diagnostics.length > 0) { + for (const d of parsed.diagnostics) { + const color = d.type === "error" ? chalk.red : chalk.yellow; + console.error(color(`${d.type === "error" ? "Error" : "Warning"}: ${d.message}`)); + } + if (parsed.diagnostics.some((d) => d.type === "error")) { + process.exit(1); + } + } + time("parseArgs"); + + if (parsed.version) { + console.log(VERSION); + process.exit(0); + } + + let appMode = resolveAppMode(parsed, process.stdin.isTTY, process.stdout.isTTY); + const shouldTakeOverStdout = appMode !== "interactive" && !isPlainRuntimeMetadataCommand(parsed); + if (shouldTakeOverStdout) { + takeOverStdout(); + } + + validateForkFlags(parsed); + validateSessionIdFlags(parsed); + + const migratedProviders: string[] = []; + + const startupSettingsManager = SettingsManager.create(cwd, agentDir); + reportDiagnostics(collectSettingsDiagnostics(startupSettingsManager, "startup session lookup")); + + // Experimental first-time setup: theme choice and analytics opt-in. + // Runs before any runtime services are created so the chosen settings apply everywhere. + if (appMode === "interactive" && !parsed.help && parsed.listModels === undefined && shouldRunFirstTimeSetup()) { + await showFirstTimeSetup(startupSettingsManager); + time("firstTimeSetup"); + } + + // Decide the final runtime cwd before creating cwd-bound runtime services. + // --session and --resume may select a session from another project, so project-local + // settings, resources, provider registrations, and models must be resolved only after + // the target session cwd is known. The startup-cwd settings manager is used only for + // sessionDir lookup during session selection. + const envSessionDir = process.env[ENV_SESSION_DIR]; + const sessionDir = + (parsed.sessionDir ? normalizePath(parsed.sessionDir) : undefined) ?? + (envSessionDir ? expandTildePath(envSessionDir) : undefined) ?? + startupSettingsManager.getSessionDir(); + let sessionManager = await createSessionManager(parsed, cwd, sessionDir, startupSettingsManager); + const missingSessionCwdIssue = getMissingSessionCwdIssue(sessionManager, cwd); + if (missingSessionCwdIssue) { + if (appMode === "interactive") { + const selectedCwd = await promptForMissingSessionCwd(missingSessionCwdIssue, startupSettingsManager); + if (!selectedCwd) { + process.exit(0); + } + sessionManager = SessionManager.open(missingSessionCwdIssue.sessionFile!, sessionDir, selectedCwd); + } else { + console.error(chalk.red(new MissingSessionCwdError(missingSessionCwdIssue).message)); + process.exit(1); + } + } + if (parsed.name !== undefined) { + const name = parsed.name.trim(); + if (!name) { + console.error(chalk.red("Error: --name requires a non-empty value")); + process.exit(1); + } + sessionManager.appendSessionInfo(name); + } + time("createSessionManager"); + + const trustStore = new ProjectTrustStore(agentDir); + const sessionCwd = sessionManager.getCwd(); + const autoTrustOnReloadCwd = + parsed.projectTrustOverride === undefined && !hasTrustRequiringProjectResources(sessionCwd) + ? sessionCwd + : undefined; + const trustPromptMode: AppMode = parsed.help || parsed.listModels !== undefined ? "print" : appMode; + const projectTrustByCwd = new Map(); + + const resolvedSkillPaths = resolveCliPaths(cwd, parsed.skills); + const resolvedPromptTemplatePaths = resolveCliPaths(cwd, parsed.promptTemplates); + const resolvedThemePaths = resolveCliPaths(cwd, parsed.themes); + const authStorage = AuthStorage.create(); + const createRuntime: CreateAgentSessionRuntimeFactory = async ({ + cwd, + agentDir, + sessionManager, + sessionStartEvent, + projectTrustContext, + }) => { + const isInitialRuntime = sessionStartEvent === undefined; + const projectTrustDiagnostics: AgentSessionRuntimeDiagnostic[] = []; + const cachedProjectTrust = projectTrustByCwd.get(cwd); + const hasTrustRequiringResources = hasTrustRequiringProjectResources(cwd); + const shouldResolveProjectTrust = + parsed.projectTrustOverride === undefined && cachedProjectTrust === undefined && hasTrustRequiringResources; + const projectTrusted = shouldResolveProjectTrust + ? false + : (cachedProjectTrust ?? + parsed.projectTrustOverride ?? + (!hasTrustRequiringResources || trustStore.get(cwd) === true)); + const runtimeSettingsManager = SettingsManager.create(cwd, agentDir, { projectTrusted }); + const services = await createAgentSessionServices({ + cwd, + agentDir, + authStorage, + settingsManager: runtimeSettingsManager, + extensionFlagValues: parsed.unknownFlags, + resourceLoaderReloadOptions: shouldResolveProjectTrust + ? { + resolveProjectTrust: async ({ extensionsResult }) => { + const trusted = await resolveProjectTrusted({ + cwd, + trustStore, + trustOverride: parsed.projectTrustOverride, + defaultProjectTrust: startupSettingsManager.getDefaultProjectTrust(), + extensionsResult, + projectTrustContext: + projectTrustContext ?? + createProjectTrustContext({ + cwd, + mode: isInitialRuntime ? trustPromptMode : appMode, + settingsManager: startupSettingsManager, + hasUI: isInitialRuntime && trustPromptMode === "interactive", + }), + onExtensionError: (message) => projectTrustDiagnostics.push({ type: "warning", message }), + }); + projectTrustByCwd.set(cwd, trusted); + return trusted; + }, + } + : undefined, + resourceLoaderOptions: { + additionalSkillPaths: resolvedSkillPaths, + additionalPromptTemplatePaths: resolvedPromptTemplatePaths, + additionalThemePaths: resolvedThemePaths, + noExtensions: true, + noSkills: parsed.noSkills, + noPromptTemplates: parsed.noPromptTemplates, + noThemes: parsed.noThemes, + noContextFiles: parsed.noContextFiles, + systemPrompt: parsed.systemPrompt, + appendSystemPrompt: parsed.appendSystemPrompt, + extensionFactories: options?.extensionFactories, + }, + }); + const { settingsManager, modelRegistry, resourceLoader } = services; + const diagnostics: AgentSessionRuntimeDiagnostic[] = [ + ...projectTrustDiagnostics, + ...services.diagnostics, + ...collectSettingsDiagnostics(settingsManager, "runtime creation"), + ...resourceLoader.getExtensions().errors.map(({ path, error }) => ({ + type: "error" as const, + message: `Failed to load extension "${path}": ${error}`, + })), + ]; + + const modelPatterns = parsed.models ?? settingsManager.getEnabledModels(); + const scopedModels = + modelPatterns && modelPatterns.length > 0 ? await resolveModelScope(modelPatterns, modelRegistry) : []; + const { + options: sessionOptions, + cliThinkingFromModel, + diagnostics: sessionOptionDiagnostics, + } = buildSessionOptions( + parsed, + scopedModels, + sessionManager.buildSessionContext().messages.length > 0, + modelRegistry, + settingsManager, + ); + diagnostics.push(...sessionOptionDiagnostics); + + if (parsed.apiKey) { + if (!sessionOptions.model) { + diagnostics.push({ + type: "error", + message: "--api-key requires a model to be specified via --model, --provider/--model, or --models", + }); + } else { + authStorage.setRuntimeApiKey(sessionOptions.model.provider, parsed.apiKey); + } + } + + const created = await createAgentSessionFromServices({ + services, + sessionManager, + sessionStartEvent, + model: sessionOptions.model, + thinkingLevel: sessionOptions.thinkingLevel, + scopedModels: sessionOptions.scopedModels, + tools: sessionOptions.tools, + excludeTools: sessionOptions.excludeTools, + noTools: sessionOptions.noTools, + customTools: sessionOptions.customTools, + }); + const cliThinkingOverride = parsed.thinking !== undefined || cliThinkingFromModel; + if (created.session.model && cliThinkingOverride) { + created.session.setThinkingLevel(created.session.thinkingLevel); + } + + return { + ...created, + services, + diagnostics, + }; + }; + time("createRuntime"); + const runtime = await createAgentSessionRuntime(createRuntime, { + cwd: sessionManager.getCwd(), + agentDir, + sessionManager, + }); + time("createAgentSessionRuntime"); + const { services, session, modelFallbackMessage } = runtime; + const { settingsManager, modelRegistry, resourceLoader } = services; + applyHttpProxySettings(settingsManager.getGlobalSettings().httpProxy); + configureHttpDispatcher(settingsManager.getHttpIdleTimeoutMs()); + + if (parsed.help) { + const extensionFlags = resourceLoader + .getExtensions() + .extensions.flatMap((extension) => Array.from(extension.flags.values())); + printHelp(extensionFlags); + process.exit(0); + } + + if (parsed.listModels !== undefined) { + const searchPattern = typeof parsed.listModels === "string" ? parsed.listModels : undefined; + await listModels(modelRegistry, searchPattern); + process.exit(0); + } + + // Read piped stdin content (if any) + let stdinContent: string | undefined = await readPipedStdin(); + if (stdinContent !== undefined && appMode === "interactive") { + appMode = "print"; + } + time("readPipedStdin"); + + const { initialMessage, initialImages } = await prepareInitialMessage( + parsed, + settingsManager.getImageAutoResize(), + stdinContent, + ); + time("prepareInitialMessage"); + initTheme(settingsManager.getTheme(), appMode === "interactive"); + time("initTheme"); + + time("resolveModelScope"); + reportDiagnostics(runtime.diagnostics); + if (runtime.diagnostics.some((diagnostic) => diagnostic.type === "error")) { + process.exit(1); + } + time("createAgentSession"); + + if (appMode !== "interactive" && !session.model) { + console.error(chalk.red(formatNoModelsAvailableMessage())); + process.exit(1); + } + + const startupBenchmark = isTruthyEnvFlag(process.env.PI_STARTUP_BENCHMARK); + if (startupBenchmark && appMode !== "interactive") { + console.error(chalk.red("Error: PI_STARTUP_BENCHMARK only supports interactive mode")); + process.exit(1); + } + + if (appMode === "interactive") { + const interactiveMode = new InteractiveMode(runtime, { + migratedProviders, + modelFallbackMessage, + autoTrustOnReloadCwd, + initialMessage, + initialImages, + initialMessages: parsed.messages, + verbose: parsed.verbose, + }); + if (startupBenchmark) { + await interactiveMode.init(); + time("interactiveMode.init"); + interactiveMode.stop(); + stopThemeWatcher(); + printTimings(); + if (process.stdout.writableLength > 0) { + await new Promise((resolve) => process.stdout.once("drain", resolve)); + } + if (process.stderr.writableLength > 0) { + await new Promise((resolve) => process.stderr.once("drain", resolve)); + } + return; + } + + printTimings(); + await interactiveMode.run(); + } else { + printTimings(); + const exitCode = await runPrintMode(runtime, { + mode: toPrintOutputMode(appMode), + messages: parsed.messages, + initialMessage, + initialImages, + }); + stopThemeWatcher(); + restoreStdout(); + if (exitCode !== 0) { + process.exitCode = exitCode; + } + return; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/index.ts b/cactus-code/packages/coding-agent/src/modes/index.ts new file mode 100644 index 000000000..ff9fb633e --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/index.ts @@ -0,0 +1,6 @@ +/** + * Run modes for the coding agent. + */ + +export { InteractiveMode, type InteractiveModeOptions } from "./interactive/interactive-mode.ts"; +export { type PrintModeOptions, runPrintMode } from "./print-mode.ts"; diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/assets/clankolas.png b/cactus-code/packages/coding-agent/src/modes/interactive/assets/clankolas.png new file mode 100644 index 000000000..d3dc65da9 Binary files /dev/null and b/cactus-code/packages/coding-agent/src/modes/interactive/assets/clankolas.png differ diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/assistant-message.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/assistant-message.ts new file mode 100644 index 000000000..b567cb886 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/assistant-message.ts @@ -0,0 +1,153 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { Container, Markdown, type MarkdownTheme, Spacer, Text } from "@earendil-works/pi-tui"; +import { getMarkdownTheme, theme } from "../theme/theme.ts"; + +const OSC133_ZONE_START = "\x1b]133;A\x07"; +const OSC133_ZONE_END = "\x1b]133;B\x07"; +const OSC133_ZONE_FINAL = "\x1b]133;C\x07"; + +/** + * Component that renders a complete assistant message + */ +export class AssistantMessageComponent extends Container { + private contentContainer: Container; + private hideThinkingBlock: boolean; + private markdownTheme: MarkdownTheme; + private hiddenThinkingLabel: string; + private lastMessage?: AssistantMessage; + private hasToolCalls = false; + + constructor( + message?: AssistantMessage, + hideThinkingBlock = false, + markdownTheme: MarkdownTheme = getMarkdownTheme(), + hiddenThinkingLabel = "Thinking...", + ) { + super(); + + this.hideThinkingBlock = hideThinkingBlock; + this.markdownTheme = markdownTheme; + this.hiddenThinkingLabel = hiddenThinkingLabel; + + // Container for text/thinking content + this.contentContainer = new Container(); + this.addChild(this.contentContainer); + + if (message) { + this.updateContent(message); + } + } + + override invalidate(): void { + super.invalidate(); + if (this.lastMessage) { + this.updateContent(this.lastMessage); + } + } + + setHideThinkingBlock(hide: boolean): void { + this.hideThinkingBlock = hide; + if (this.lastMessage) { + this.updateContent(this.lastMessage); + } + } + + setHiddenThinkingLabel(label: string): void { + this.hiddenThinkingLabel = label; + if (this.lastMessage) { + this.updateContent(this.lastMessage); + } + } + + override render(width: number): string[] { + const lines = super.render(width); + if (this.hasToolCalls || lines.length === 0) { + return lines; + } + + lines[0] = OSC133_ZONE_START + lines[0]; + lines[lines.length - 1] = OSC133_ZONE_END + OSC133_ZONE_FINAL + lines[lines.length - 1]; + return lines; + } + + updateContent(message: AssistantMessage): void { + this.lastMessage = message; + + // Clear content container + this.contentContainer.clear(); + + const hasVisibleContent = message.content.some( + (c) => (c.type === "text" && c.text.trim()) || (c.type === "thinking" && c.thinking.trim()), + ); + + if (message.cloudHandoff !== undefined) { + const origin = message.cloudHandoff ? "cloud" : "local"; + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild( + new Text(theme.fg(message.cloudHandoff ? "warning" : "muted", `● ${origin}`), 1, 0), + ); + } else if (hasVisibleContent) { + this.contentContainer.addChild(new Spacer(1)); + } + + // Render content in order + for (let i = 0; i < message.content.length; i++) { + const content = message.content[i]; + if (content.type === "text" && content.text.trim()) { + // Assistant text messages with no background - trim the text + // Set paddingY=0 to avoid extra spacing before tool executions + this.contentContainer.addChild(new Markdown(content.text.trim(), 1, 0, this.markdownTheme)); + } else if (content.type === "thinking" && content.thinking.trim()) { + // Add spacing only when another visible assistant content block follows. + // This avoids a superfluous blank line before separately-rendered tool execution blocks. + const hasVisibleContentAfter = message.content + .slice(i + 1) + .some((c) => (c.type === "text" && c.text.trim()) || (c.type === "thinking" && c.thinking.trim())); + + if (this.hideThinkingBlock) { + // Show static thinking label when hidden + this.contentContainer.addChild( + new Text(theme.italic(theme.fg("thinkingText", this.hiddenThinkingLabel)), 1, 0), + ); + if (hasVisibleContentAfter) { + this.contentContainer.addChild(new Spacer(1)); + } + } else { + // Thinking traces in thinkingText color, italic + this.contentContainer.addChild( + new Markdown(content.thinking.trim(), 1, 0, this.markdownTheme, { + color: (text: string) => theme.fg("thinkingText", text), + italic: true, + }), + ); + if (hasVisibleContentAfter) { + this.contentContainer.addChild(new Spacer(1)); + } + } + } + } + + // Check if aborted - show after partial content + // But only if there are no tool calls (tool execution components will show the error) + const hasToolCalls = message.content.some((c) => c.type === "toolCall"); + this.hasToolCalls = hasToolCalls; + if (!hasToolCalls) { + if (message.stopReason === "aborted") { + const abortMessage = + message.errorMessage && message.errorMessage !== "Request was aborted" + ? message.errorMessage + : "Operation aborted"; + if (hasVisibleContent) { + this.contentContainer.addChild(new Spacer(1)); + } else { + this.contentContainer.addChild(new Spacer(1)); + } + this.contentContainer.addChild(new Text(theme.fg("error", abortMessage), 1, 0)); + } else if (message.stopReason === "error") { + const errorMsg = message.errorMessage || "Unknown error"; + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild(new Text(theme.fg("error", `Error: ${errorMsg}`), 1, 0)); + } + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/bash-execution.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/bash-execution.ts new file mode 100644 index 000000000..b46a9d0b9 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/bash-execution.ts @@ -0,0 +1,220 @@ +/** + * Component for displaying bash command execution with streaming output. + */ + +import { Container, Loader, Spacer, Text, type TUI } from "@earendil-works/pi-tui"; +import { + DEFAULT_MAX_BYTES, + DEFAULT_MAX_LINES, + type TruncationResult, + truncateTail, +} from "../../../core/tools/truncate.ts"; +import { stripAnsi } from "../../../utils/ansi.ts"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint, keyText } from "./keybinding-hints.ts"; +import { truncateToVisualLines } from "./visual-truncate.ts"; + +// Preview line limit when not expanded (matches tool execution behavior) +const PREVIEW_LINES = 20; + +export class BashExecutionComponent extends Container { + private command: string; + private outputLines: string[] = []; + private status: "running" | "complete" | "cancelled" | "error" = "running"; + private exitCode: number | undefined = undefined; + private loader: Loader; + private truncationResult?: TruncationResult; + private fullOutputPath?: string; + private expanded = false; + private contentContainer: Container; + + constructor(command: string, ui: TUI, excludeFromContext = false) { + super(); + this.command = command; + + // Use dim border for excluded-from-context commands (!! prefix) + const colorKey = excludeFromContext ? "dim" : "bashMode"; + const borderColor = (str: string) => theme.fg(colorKey, str); + + // Add spacer + this.addChild(new Spacer(1)); + + // Top border + this.addChild(new DynamicBorder(borderColor)); + + // Content container (holds dynamic content between borders) + this.contentContainer = new Container(); + this.addChild(this.contentContainer); + + // Command header + const header = new Text(theme.fg(colorKey, theme.bold(`$ ${command}`)), 1, 0); + this.contentContainer.addChild(header); + + // Loader + this.loader = new Loader( + ui, + (spinner) => theme.fg(colorKey, spinner), + (text) => theme.fg("muted", text), + `Running... (${keyText("tui.select.cancel")} to cancel)`, // Plain text for loader + ); + this.contentContainer.addChild(this.loader); + + // Bottom border + this.addChild(new DynamicBorder(borderColor)); + } + + /** + * Set whether the output is expanded (shows full output) or collapsed (preview only). + */ + setExpanded(expanded: boolean): void { + this.expanded = expanded; + this.updateDisplay(); + } + + override invalidate(): void { + super.invalidate(); + this.updateDisplay(); + } + + appendOutput(chunk: string): void { + // Strip ANSI codes and normalize line endings + // Note: binary data is already sanitized in tui-renderer.ts executeBashCommand + const clean = stripAnsi(chunk).replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + + // Append to output lines + const newLines = clean.split("\n"); + if (this.outputLines.length > 0 && newLines.length > 0) { + // Append first chunk to last line (incomplete line continuation) + this.outputLines[this.outputLines.length - 1] += newLines[0]; + this.outputLines.push(...newLines.slice(1)); + } else { + this.outputLines.push(...newLines); + } + + this.updateDisplay(); + } + + setComplete( + exitCode: number | undefined, + cancelled: boolean, + truncationResult?: TruncationResult, + fullOutputPath?: string, + ): void { + this.exitCode = exitCode; + this.status = cancelled + ? "cancelled" + : exitCode !== 0 && exitCode !== undefined && exitCode !== null + ? "error" + : "complete"; + this.truncationResult = truncationResult; + this.fullOutputPath = fullOutputPath; + + // Stop loader + this.loader.stop(); + + this.updateDisplay(); + } + + private updateDisplay(): void { + // Apply truncation for LLM context limits (same limits as bash tool) + const fullOutput = this.outputLines.join("\n"); + const contextTruncation = truncateTail(fullOutput, { + maxLines: DEFAULT_MAX_LINES, + maxBytes: DEFAULT_MAX_BYTES, + }); + + // Get the lines to potentially display (after context truncation) + const availableLines = contextTruncation.content ? contextTruncation.content.split("\n") : []; + + // Apply preview truncation based on expanded state + const previewLogicalLines = availableLines.slice(-PREVIEW_LINES); + const hiddenLineCount = availableLines.length - previewLogicalLines.length; + + // Rebuild content container + this.contentContainer.clear(); + + // Command header + const header = new Text(theme.fg("bashMode", theme.bold(`$ ${this.command}`)), 1, 0); + this.contentContainer.addChild(header); + + // Output + if (availableLines.length > 0) { + if (this.expanded) { + // Show all lines + const displayText = availableLines.map((line) => theme.fg("muted", line)).join("\n"); + this.contentContainer.addChild(new Text(`\n${displayText}`, 1, 0)); + } else { + // Use shared visual truncation utility with width-aware caching + const styledOutput = previewLogicalLines.map((line) => theme.fg("muted", line)).join("\n"); + const styledInput = `\n${styledOutput}`; + let cachedWidth: number | undefined; + let cachedLines: string[] | undefined; + this.contentContainer.addChild({ + render: (width: number) => { + if (cachedLines === undefined || cachedWidth !== width) { + const result = truncateToVisualLines(styledInput, PREVIEW_LINES, width, 1); + cachedLines = result.visualLines; + cachedWidth = width; + } + return cachedLines ?? []; + }, + invalidate: () => { + cachedWidth = undefined; + cachedLines = undefined; + }, + }); + } + } + + // Loader or status + if (this.status === "running") { + this.contentContainer.addChild(this.loader); + } else { + const statusParts: string[] = []; + + // Show how many lines are hidden (collapsed preview) + if (hiddenLineCount > 0) { + if (this.expanded) { + statusParts.push( + `${theme.fg("muted", "(")}${keyHint("app.tools.expand", "to collapse")}${theme.fg("muted", ")")}`, + ); + } else { + statusParts.push( + `${theme.fg("muted", `... ${hiddenLineCount} more lines (`)}${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`, + ); + } + } + + if (this.status === "cancelled") { + statusParts.push(theme.fg("warning", "(cancelled)")); + } else if (this.status === "error") { + statusParts.push(theme.fg("error", `(exit ${this.exitCode})`)); + } + + // Add truncation warning (context truncation, not preview truncation) + const wasTruncated = this.truncationResult?.truncated || contextTruncation.truncated; + if (wasTruncated && this.fullOutputPath) { + statusParts.push(theme.fg("warning", `Output truncated. Full output: ${this.fullOutputPath}`)); + } + + if (statusParts.length > 0) { + this.contentContainer.addChild(new Text(`\n${statusParts.join("\n")}`, 1, 0)); + } + } + } + + /** + * Get the raw output for creating BashExecutionMessage. + */ + getOutput(): string { + return this.outputLines.join("\n"); + } + + /** + * Get the command that was executed. + */ + getCommand(): string { + return this.command; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/bordered-loader.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/bordered-loader.ts new file mode 100644 index 000000000..4d692ffc8 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/bordered-loader.ts @@ -0,0 +1,68 @@ +import { CancellableLoader, Container, Loader, Spacer, Text, type TUI } from "@earendil-works/pi-tui"; +import type { Theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint } from "./keybinding-hints.ts"; + +/** Loader wrapped with borders for extension UI */ +export class BorderedLoader extends Container { + private loader: CancellableLoader | Loader; + private cancellable: boolean; + private signalController?: AbortController; + + constructor(tui: TUI, theme: Theme, message: string, options?: { cancellable?: boolean }) { + super(); + this.cancellable = options?.cancellable ?? true; + const borderColor = (s: string) => theme.fg("border", s); + this.addChild(new DynamicBorder(borderColor)); + if (this.cancellable) { + this.loader = new CancellableLoader( + tui, + (s) => theme.fg("accent", s), + (s) => theme.fg("muted", s), + message, + ); + } else { + this.signalController = new AbortController(); + this.loader = new Loader( + tui, + (s) => theme.fg("accent", s), + (s) => theme.fg("muted", s), + message, + ); + } + this.addChild(this.loader); + if (this.cancellable) { + this.addChild(new Spacer(1)); + this.addChild(new Text(keyHint("tui.select.cancel", "cancel"), 1, 0)); + } + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder(borderColor)); + } + + get signal(): AbortSignal { + if (this.cancellable) { + return (this.loader as CancellableLoader).signal; + } + return this.signalController?.signal ?? new AbortController().signal; + } + + set onAbort(fn: (() => void) | undefined) { + if (this.cancellable) { + (this.loader as CancellableLoader).onAbort = fn; + } + } + + handleInput(data: string): void { + if (this.cancellable) { + (this.loader as CancellableLoader).handleInput(data); + } + } + + dispose(): void { + if ("dispose" in this.loader && typeof this.loader.dispose === "function") { + this.loader.dispose(); + } else if ("stop" in this.loader && typeof this.loader.stop === "function") { + this.loader.stop(); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/branch-summary-message.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/branch-summary-message.ts new file mode 100644 index 000000000..af85115f0 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/branch-summary-message.ts @@ -0,0 +1,58 @@ +import { Box, Markdown, type MarkdownTheme, Spacer, Text } from "@earendil-works/pi-tui"; +import type { BranchSummaryMessage } from "../../../core/messages.ts"; +import { getMarkdownTheme, theme } from "../theme/theme.ts"; +import { keyText } from "./keybinding-hints.ts"; + +/** + * Component that renders a branch summary message with collapsed/expanded state. + * Uses same background color as custom messages for visual consistency. + */ +export class BranchSummaryMessageComponent extends Box { + private expanded = false; + private message: BranchSummaryMessage; + private markdownTheme: MarkdownTheme; + + constructor(message: BranchSummaryMessage, markdownTheme: MarkdownTheme = getMarkdownTheme()) { + super(1, 1, (t) => theme.bg("customMessageBg", t)); + this.message = message; + this.markdownTheme = markdownTheme; + this.updateDisplay(); + } + + setExpanded(expanded: boolean): void { + this.expanded = expanded; + this.updateDisplay(); + } + + override invalidate(): void { + super.invalidate(); + this.updateDisplay(); + } + + private updateDisplay(): void { + this.clear(); + + const label = theme.fg("customMessageLabel", `\x1b[1m[branch]\x1b[22m`); + this.addChild(new Text(label, 0, 0)); + this.addChild(new Spacer(1)); + + if (this.expanded) { + const header = "**Branch Summary**\n\n"; + this.addChild( + new Markdown(header + this.message.summary, 0, 0, this.markdownTheme, { + color: (text: string) => theme.fg("customMessageText", text), + }), + ); + } else { + this.addChild( + new Text( + theme.fg("customMessageText", "Branch summary (") + + theme.fg("dim", keyText("app.tools.expand")) + + theme.fg("customMessageText", " to expand)"), + 0, + 0, + ), + ); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/compaction-summary-message.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/compaction-summary-message.ts new file mode 100644 index 000000000..cad613d76 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/compaction-summary-message.ts @@ -0,0 +1,59 @@ +import { Box, Markdown, type MarkdownTheme, Spacer, Text } from "@earendil-works/pi-tui"; +import type { CompactionSummaryMessage } from "../../../core/messages.ts"; +import { getMarkdownTheme, theme } from "../theme/theme.ts"; +import { keyText } from "./keybinding-hints.ts"; + +/** + * Component that renders a compaction message with collapsed/expanded state. + * Uses same background color as custom messages for visual consistency. + */ +export class CompactionSummaryMessageComponent extends Box { + private expanded = false; + private message: CompactionSummaryMessage; + private markdownTheme: MarkdownTheme; + + constructor(message: CompactionSummaryMessage, markdownTheme: MarkdownTheme = getMarkdownTheme()) { + super(1, 1, (t) => theme.bg("customMessageBg", t)); + this.message = message; + this.markdownTheme = markdownTheme; + this.updateDisplay(); + } + + setExpanded(expanded: boolean): void { + this.expanded = expanded; + this.updateDisplay(); + } + + override invalidate(): void { + super.invalidate(); + this.updateDisplay(); + } + + private updateDisplay(): void { + this.clear(); + + const tokenStr = this.message.tokensBefore.toLocaleString(); + const label = theme.fg("customMessageLabel", `\x1b[1m[compaction]\x1b[22m`); + this.addChild(new Text(label, 0, 0)); + this.addChild(new Spacer(1)); + + if (this.expanded) { + const header = `**Compacted from ${tokenStr} tokens**\n\n`; + this.addChild( + new Markdown(header + this.message.summary, 0, 0, this.markdownTheme, { + color: (text: string) => theme.fg("customMessageText", text), + }), + ); + } else { + this.addChild( + new Text( + theme.fg("customMessageText", `Compacted from ${tokenStr} tokens (`) + + theme.fg("dim", keyText("app.tools.expand")) + + theme.fg("customMessageText", " to expand)"), + 0, + 0, + ), + ); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/countdown-timer.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/countdown-timer.ts new file mode 100644 index 000000000..73ec2734d --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/countdown-timer.ts @@ -0,0 +1,39 @@ +/** + * Reusable countdown timer for dialog components. + */ + +import type { TUI } from "@earendil-works/pi-tui"; + +export class CountdownTimer { + private intervalId: ReturnType | undefined; + private remainingSeconds: number; + private tui: TUI | undefined; + private onTick: (seconds: number) => void; + private onExpire: () => void; + + constructor(timeoutMs: number, tui: TUI | undefined, onTick: (seconds: number) => void, onExpire: () => void) { + this.tui = tui; + this.onTick = onTick; + this.onExpire = onExpire; + this.remainingSeconds = Math.ceil(timeoutMs / 1000); + this.onTick(this.remainingSeconds); + + this.intervalId = setInterval(() => { + this.remainingSeconds--; + this.onTick(this.remainingSeconds); + this.tui?.requestRender(); + + if (this.remainingSeconds <= 0) { + this.dispose(); + this.onExpire(); + } + }, 1000); + } + + dispose(): void { + if (this.intervalId) { + clearInterval(this.intervalId); + this.intervalId = undefined; + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/custom-editor.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/custom-editor.ts new file mode 100644 index 000000000..a9d256222 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/custom-editor.ts @@ -0,0 +1,80 @@ +import { Editor, type EditorOptions, type EditorTheme, type TUI } from "@earendil-works/pi-tui"; +import type { AppKeybinding, KeybindingsManager } from "../../../core/keybindings.ts"; + +/** + * Custom editor that handles app-level keybindings for coding-agent. + */ +export class CustomEditor extends Editor { + private keybindings: KeybindingsManager; + public actionHandlers: Map void> = new Map(); + + // Special handlers that can be dynamically replaced + public onEscape?: () => void; + public onCtrlD?: () => void; + public onPasteImage?: () => void; + /** Handler for extension-registered shortcuts. Returns true if handled. */ + public onExtensionShortcut?: (data: string) => boolean; + + constructor(tui: TUI, theme: EditorTheme, keybindings: KeybindingsManager, options?: EditorOptions) { + super(tui, theme, options); + this.keybindings = keybindings; + } + + /** + * Register a handler for an app action. + */ + onAction(action: AppKeybinding, handler: () => void): void { + this.actionHandlers.set(action, handler); + } + + handleInput(data: string): void { + // Check extension-registered shortcuts first + if (this.onExtensionShortcut?.(data)) { + return; + } + + // Check for paste image keybinding + if (this.keybindings.matches(data, "app.clipboard.pasteImage")) { + this.onPasteImage?.(); + return; + } + + // Check app keybindings first + + // Escape/interrupt - only if autocomplete is NOT active + if (this.keybindings.matches(data, "app.interrupt")) { + if (!this.isShowingAutocomplete()) { + // Use dynamic onEscape if set, otherwise registered handler + const handler = this.onEscape ?? this.actionHandlers.get("app.interrupt"); + if (handler) { + handler(); + return; + } + } + // Let parent handle escape for autocomplete cancellation + super.handleInput(data); + return; + } + + // Exit (Ctrl+D) - only when editor is empty + if (this.keybindings.matches(data, "app.exit")) { + if (this.getText().length === 0) { + const handler = this.onCtrlD ?? this.actionHandlers.get("app.exit"); + if (handler) handler(); + return; + } + // Fall through to editor handling for delete-char-forward when not empty + } + + // Check all other app actions + for (const [action, handler] of this.actionHandlers) { + if (action !== "app.interrupt" && action !== "app.exit" && this.keybindings.matches(data, action)) { + handler(); + return; + } + } + + // Pass to parent for editor handling + super.handleInput(data); + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/custom-message.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/custom-message.ts new file mode 100644 index 000000000..a8fffc6f3 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/custom-message.ts @@ -0,0 +1,99 @@ +import type { TextContent } from "@earendil-works/pi-ai"; +import type { Component } from "@earendil-works/pi-tui"; +import { Box, Container, Markdown, type MarkdownTheme, Spacer, Text } from "@earendil-works/pi-tui"; +import type { MessageRenderer } from "../../../core/extensions/types.ts"; +import type { CustomMessage } from "../../../core/messages.ts"; +import { getMarkdownTheme, theme } from "../theme/theme.ts"; + +/** + * Component that renders a custom message entry from extensions. + * Uses distinct styling to differentiate from user messages. + */ +export class CustomMessageComponent extends Container { + private message: CustomMessage; + private customRenderer?: MessageRenderer; + private box: Box; + private customComponent?: Component; + private markdownTheme: MarkdownTheme; + private _expanded = false; + + constructor( + message: CustomMessage, + customRenderer?: MessageRenderer, + markdownTheme: MarkdownTheme = getMarkdownTheme(), + ) { + super(); + this.message = message; + this.customRenderer = customRenderer; + this.markdownTheme = markdownTheme; + + this.addChild(new Spacer(1)); + + // Create box with purple background (used for default rendering) + this.box = new Box(1, 1, (t) => theme.bg("customMessageBg", t)); + + this.rebuild(); + } + + setExpanded(expanded: boolean): void { + if (this._expanded !== expanded) { + this._expanded = expanded; + this.rebuild(); + } + } + + override invalidate(): void { + super.invalidate(); + this.rebuild(); + } + + private rebuild(): void { + // Remove previous content component + if (this.customComponent) { + this.removeChild(this.customComponent); + this.customComponent = undefined; + } + this.removeChild(this.box); + + // Try custom renderer first - it handles its own styling + if (this.customRenderer) { + try { + const component = this.customRenderer(this.message, { expanded: this._expanded }, theme); + if (component) { + // Custom renderer provides its own styled component + this.customComponent = component; + this.addChild(component); + return; + } + } catch { + // Fall through to default rendering + } + } + + // Default rendering uses our box + this.addChild(this.box); + this.box.clear(); + + // Default rendering: label + content + const label = theme.fg("customMessageLabel", `\x1b[1m[${this.message.customType}]\x1b[22m`); + this.box.addChild(new Text(label, 0, 0)); + this.box.addChild(new Spacer(1)); + + // Extract text content + let text: string; + if (typeof this.message.content === "string") { + text = this.message.content; + } else { + text = this.message.content + .filter((c): c is TextContent => c.type === "text") + .map((c) => c.text) + .join("\n"); + } + + this.box.addChild( + new Markdown(text, 0, 0, this.markdownTheme, { + color: (text: string) => theme.fg("customMessageText", text), + }), + ); + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/diff.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/diff.ts new file mode 100644 index 000000000..54e88273d --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/diff.ts @@ -0,0 +1,147 @@ +import * as Diff from "diff"; +import { theme } from "../theme/theme.ts"; + +/** + * Parse diff line to extract prefix, line number, and content. + * Format: "+123 content" or "-123 content" or " 123 content" or " ..." + */ +function parseDiffLine(line: string): { prefix: string; lineNum: string; content: string } | null { + const match = line.match(/^([+-\s])(\s*\d*)\s(.*)$/); + if (!match) return null; + return { prefix: match[1], lineNum: match[2], content: match[3] }; +} + +/** + * Replace tabs with spaces for consistent rendering. + */ +function replaceTabs(text: string): string { + return text.replace(/\t/g, " "); +} + +/** + * Compute word-level diff and render with inverse on changed parts. + * Uses diffWords which groups whitespace with adjacent words for cleaner highlighting. + * Strips leading whitespace from inverse to avoid highlighting indentation. + */ +function renderIntraLineDiff(oldContent: string, newContent: string): { removedLine: string; addedLine: string } { + const wordDiff = Diff.diffWords(oldContent, newContent); + + let removedLine = ""; + let addedLine = ""; + let isFirstRemoved = true; + let isFirstAdded = true; + + for (const part of wordDiff) { + if (part.removed) { + let value = part.value; + // Strip leading whitespace from the first removed part + if (isFirstRemoved) { + const leadingWs = value.match(/^(\s*)/)?.[1] || ""; + value = value.slice(leadingWs.length); + removedLine += leadingWs; + isFirstRemoved = false; + } + if (value) { + removedLine += theme.inverse(value); + } + } else if (part.added) { + let value = part.value; + // Strip leading whitespace from the first added part + if (isFirstAdded) { + const leadingWs = value.match(/^(\s*)/)?.[1] || ""; + value = value.slice(leadingWs.length); + addedLine += leadingWs; + isFirstAdded = false; + } + if (value) { + addedLine += theme.inverse(value); + } + } else { + removedLine += part.value; + addedLine += part.value; + } + } + + return { removedLine, addedLine }; +} + +export interface RenderDiffOptions { + /** File path (unused, kept for API compatibility) */ + filePath?: string; +} + +/** + * Render a diff string with colored lines and intra-line change highlighting. + * - Context lines: dim/gray + * - Removed lines: red, with inverse on changed tokens + * - Added lines: green, with inverse on changed tokens + */ +export function renderDiff(diffText: string, _options: RenderDiffOptions = {}): string { + const lines = diffText.split("\n"); + const result: string[] = []; + + let i = 0; + while (i < lines.length) { + const line = lines[i]; + const parsed = parseDiffLine(line); + + if (!parsed) { + result.push(theme.fg("toolDiffContext", line)); + i++; + continue; + } + + if (parsed.prefix === "-") { + // Collect consecutive removed lines + const removedLines: { lineNum: string; content: string }[] = []; + while (i < lines.length) { + const p = parseDiffLine(lines[i]); + if (!p || p.prefix !== "-") break; + removedLines.push({ lineNum: p.lineNum, content: p.content }); + i++; + } + + // Collect consecutive added lines + const addedLines: { lineNum: string; content: string }[] = []; + while (i < lines.length) { + const p = parseDiffLine(lines[i]); + if (!p || p.prefix !== "+") break; + addedLines.push({ lineNum: p.lineNum, content: p.content }); + i++; + } + + // Only do intra-line diffing when there's exactly one removed and one added line + // (indicating a single line modification). Otherwise, show lines as-is. + if (removedLines.length === 1 && addedLines.length === 1) { + const removed = removedLines[0]; + const added = addedLines[0]; + + const { removedLine, addedLine } = renderIntraLineDiff( + replaceTabs(removed.content), + replaceTabs(added.content), + ); + + result.push(theme.fg("toolDiffRemoved", `-${removed.lineNum} ${removedLine}`)); + result.push(theme.fg("toolDiffAdded", `+${added.lineNum} ${addedLine}`)); + } else { + // Show all removed lines first, then all added lines + for (const removed of removedLines) { + result.push(theme.fg("toolDiffRemoved", `-${removed.lineNum} ${replaceTabs(removed.content)}`)); + } + for (const added of addedLines) { + result.push(theme.fg("toolDiffAdded", `+${added.lineNum} ${replaceTabs(added.content)}`)); + } + } + } else if (parsed.prefix === "+") { + // Standalone added line + result.push(theme.fg("toolDiffAdded", `+${parsed.lineNum} ${replaceTabs(parsed.content)}`)); + i++; + } else { + // Context line + result.push(theme.fg("toolDiffContext", ` ${parsed.lineNum} ${replaceTabs(parsed.content)}`)); + i++; + } + } + + return result.join("\n"); +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/dynamic-border.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/dynamic-border.ts new file mode 100644 index 000000000..77342b25e --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/dynamic-border.ts @@ -0,0 +1,25 @@ +import type { Component } from "@earendil-works/pi-tui"; +import { theme } from "../theme/theme.ts"; + +/** + * Dynamic border component that adjusts to viewport width. + * + * Note: When used from extensions loaded via jiti, the global `theme` may be undefined + * because jiti creates a separate module cache. Always pass an explicit color + * function when using DynamicBorder in components exported for extension use. + */ +export class DynamicBorder implements Component { + private color: (str: string) => string; + + constructor(color: (str: string) => string = (str) => theme.fg("border", str)) { + this.color = color; + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(width: number): string[] { + return [this.color("─".repeat(Math.max(1, width)))]; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-editor.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-editor.ts new file mode 100644 index 000000000..96f081a1f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-editor.ts @@ -0,0 +1,156 @@ +/** + * Multi-line editor component for extensions. + * Supports Ctrl+G for external editor. + */ + +import { spawn } from "node:child_process"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { + Container, + Editor, + type EditorOptions, + type Focusable, + getKeybindings, + Spacer, + Text, + type TUI, +} from "@earendil-works/pi-tui"; +import type { KeybindingsManager } from "../../../core/keybindings.ts"; +import { getEditorTheme, theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint } from "./keybinding-hints.ts"; + +export class ExtensionEditorComponent extends Container implements Focusable { + private editor: Editor; + private onSubmitCallback: (value: string) => void; + private onCancelCallback: () => void; + private tui: TUI; + private keybindings: KeybindingsManager; + + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.editor.focused = value; + } + + constructor( + tui: TUI, + keybindings: KeybindingsManager, + title: string, + prefill: string | undefined, + onSubmit: (value: string) => void, + onCancel: () => void, + options?: EditorOptions, + ) { + super(); + + this.tui = tui; + this.keybindings = keybindings; + this.onSubmitCallback = onSubmit; + this.onCancelCallback = onCancel; + + // Add top border + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + + // Add title + this.addChild(new Text(theme.fg("accent", title), 1, 0)); + this.addChild(new Spacer(1)); + + // Create editor + this.editor = new Editor(tui, getEditorTheme(), options); + if (prefill) { + this.editor.setText(prefill); + } + // Wire up Enter to submit (Shift+Enter for newlines, like the main editor) + this.editor.onSubmit = (text: string) => { + this.onSubmitCallback(text); + }; + this.addChild(this.editor); + + this.addChild(new Spacer(1)); + + // Add hint + const hasExternalEditor = !!(process.env.VISUAL || process.env.EDITOR); + const hint = + keyHint("tui.select.confirm", "submit") + + " " + + keyHint("tui.input.newLine", "newline") + + " " + + keyHint("tui.select.cancel", "cancel") + + (hasExternalEditor ? ` ${keyHint("app.editor.external", "external editor")}` : ""); + this.addChild(new Text(hint, 1, 0)); + + this.addChild(new Spacer(1)); + + // Add bottom border + this.addChild(new DynamicBorder()); + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + // Escape or Ctrl+C to cancel + if (kb.matches(keyData, "tui.select.cancel")) { + this.onCancelCallback(); + return; + } + + // External editor (app keybinding) + if (this.keybindings.matches(keyData, "app.editor.external")) { + this.openExternalEditor(); + return; + } + + // Forward to editor + this.editor.handleInput(keyData); + } + + private async openExternalEditor(): Promise { + const editorCmd = process.env.VISUAL || process.env.EDITOR; + if (!editorCmd) { + return; + } + + const currentText = this.editor.getText(); + const tmpFile = path.join(os.tmpdir(), `cactus-extension-editor-${Date.now()}.md`); + + try { + fs.writeFileSync(tmpFile, currentText, "utf-8"); + this.tui.stop(); + + const [editor, ...editorArgs] = editorCmd.split(" "); + process.stdout.write(`Launching external editor: ${editorCmd}\nPi will resume when the editor exits.\n`); + + // Do not use spawnSync here. On Windows, synchronous child_process calls can keep + // Node/libuv's console input read active after tui.stop() pauses stdin, racing + // vim/nvim for the console input buffer until Ctrl+C cancels the pending read. + const status = await new Promise((resolve) => { + const child = spawn(editor, [...editorArgs, tmpFile], { + stdio: "inherit", + shell: process.platform === "win32", + }); + child.on("error", () => resolve(null)); + child.on("close", (code) => resolve(code)); + }); + + if (status === 0) { + const newContent = fs.readFileSync(tmpFile, "utf-8").replace(/\n$/, ""); + this.editor.setText(newContent); + } + } finally { + try { + fs.unlinkSync(tmpFile); + } catch { + // Ignore cleanup errors + } + this.tui.start(); + // Force full re-render since external editor uses alternate screen + this.tui.requestRender(true); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-input.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-input.ts new file mode 100644 index 000000000..581ef4f5c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-input.ts @@ -0,0 +1,87 @@ +/** + * Simple text input component for extensions. + */ + +import { Container, type Focusable, getKeybindings, Input, Spacer, Text, type TUI } from "@earendil-works/pi-tui"; +import { theme } from "../theme/theme.ts"; +import { CountdownTimer } from "./countdown-timer.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint } from "./keybinding-hints.ts"; + +export interface ExtensionInputOptions { + tui?: TUI; + timeout?: number; +} + +export class ExtensionInputComponent extends Container implements Focusable { + private input: Input; + private onSubmitCallback: (value: string) => void; + private onCancelCallback: () => void; + private titleText: Text; + private baseTitle: string; + private countdown: CountdownTimer | undefined; + + // Focusable implementation - propagate to input for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.input.focused = value; + } + + constructor( + title: string, + _placeholder: string | undefined, + onSubmit: (value: string) => void, + onCancel: () => void, + opts?: ExtensionInputOptions, + ) { + super(); + + this.onSubmitCallback = onSubmit; + this.onCancelCallback = onCancel; + this.baseTitle = title; + + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + + this.titleText = new Text(theme.fg("accent", title), 1, 0); + this.addChild(this.titleText); + this.addChild(new Spacer(1)); + + if (opts?.timeout && opts.timeout > 0 && opts.tui) { + this.countdown = new CountdownTimer( + opts.timeout, + opts.tui, + (s) => this.titleText.setText(theme.fg("accent", `${this.baseTitle} (${s}s)`)), + () => this.onCancelCallback(), + ); + } + + this.input = new Input(); + this.addChild(this.input); + this.addChild(new Spacer(1)); + this.addChild( + new Text(`${keyHint("tui.select.confirm", "submit")} ${keyHint("tui.select.cancel", "cancel")}`, 1, 0), + ); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "tui.select.confirm") || keyData === "\n") { + this.onSubmitCallback(this.input.getValue()); + } else if (kb.matches(keyData, "tui.select.cancel")) { + this.onCancelCallback(); + } else { + this.input.handleInput(keyData); + } + } + + dispose(): void { + this.countdown?.dispose(); + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-selector.ts new file mode 100644 index 000000000..d1dbdb4d0 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/extension-selector.ts @@ -0,0 +1,112 @@ +/** + * Generic selector component for extensions. + * Displays a list of string options with keyboard navigation. + */ + +import { Container, getKeybindings, Spacer, Text, type TUI } from "@earendil-works/pi-tui"; +import { theme } from "../theme/theme.ts"; +import { CountdownTimer } from "./countdown-timer.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint, rawKeyHint } from "./keybinding-hints.ts"; + +export interface ExtensionSelectorOptions { + tui?: TUI; + timeout?: number; + onToggleToolsExpanded?: () => void; +} + +export class ExtensionSelectorComponent extends Container { + private options: string[]; + private selectedIndex = 0; + private listContainer: Container; + private onSelectCallback: (option: string) => void; + private onCancelCallback: () => void; + private titleText: Text; + private baseTitle: string; + private countdown: CountdownTimer | undefined; + private onToggleToolsExpanded: (() => void) | undefined; + + constructor( + title: string, + options: string[], + onSelect: (option: string) => void, + onCancel: () => void, + opts?: ExtensionSelectorOptions, + ) { + super(); + + this.options = options; + this.onSelectCallback = onSelect; + this.onCancelCallback = onCancel; + this.onToggleToolsExpanded = opts?.onToggleToolsExpanded; + this.baseTitle = title; + + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + + this.titleText = new Text(theme.fg("accent", theme.bold(title)), 1, 0); + this.addChild(this.titleText); + this.addChild(new Spacer(1)); + + if (opts?.timeout && opts.timeout > 0 && opts.tui) { + this.countdown = new CountdownTimer( + opts.timeout, + opts.tui, + (s) => this.titleText.setText(theme.fg("accent", theme.bold(`${this.baseTitle} (${s}s)`))), + () => this.onCancelCallback(), + ); + } + + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild( + new Text( + rawKeyHint("↑↓", "navigate") + + " " + + keyHint("tui.select.confirm", "select") + + " " + + keyHint("tui.select.cancel", "cancel"), + 1, + 0, + ), + ); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + + this.updateList(); + } + + private updateList(): void { + this.listContainer.clear(); + for (let i = 0; i < this.options.length; i++) { + const isSelected = i === this.selectedIndex; + const text = isSelected + ? theme.fg("accent", "→ ") + theme.fg("accent", this.options[i]) + : ` ${theme.fg("text", this.options[i])}`; + this.listContainer.addChild(new Text(text, 1, 0)); + } + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "app.tools.expand")) { + this.onToggleToolsExpanded?.(); + } else if (kb.matches(keyData, "tui.select.up") || keyData === "k") { + this.selectedIndex = Math.max(0, this.selectedIndex - 1); + this.updateList(); + } else if (kb.matches(keyData, "tui.select.down") || keyData === "j") { + this.selectedIndex = Math.min(this.options.length - 1, this.selectedIndex + 1); + this.updateList(); + } else if (kb.matches(keyData, "tui.select.confirm") || keyData === "\n") { + const selected = this.options[this.selectedIndex]; + if (selected) this.onSelectCallback(selected); + } else if (kb.matches(keyData, "tui.select.cancel")) { + this.onCancelCallback(); + } + } + + dispose(): void { + this.countdown?.dispose(); + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/first-time-setup.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/first-time-setup.ts new file mode 100644 index 000000000..299f5b2ba --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/first-time-setup.ts @@ -0,0 +1,145 @@ +import { Container, getKeybindings, Spacer, Text } from "@earendil-works/pi-tui"; +import { APP_NAME } from "../../../config.ts"; +import { type TerminalTheme, theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint, rawKeyHint } from "./keybinding-hints.ts"; + +export interface FirstTimeSetupResult { + theme: TerminalTheme; + shareAnalytics: boolean; +} + +export interface FirstTimeSetupOptions { + detectedTheme: TerminalTheme; + onThemePreview: (themeName: TerminalTheme) => void; + onSubmit: (result: FirstTimeSetupResult) => void; + onCancel: () => void; +} + +const THEME_OPTIONS: Array<{ value: TerminalTheme; label: string }> = [ + { value: "dark", label: "Dark" }, + { value: "light", label: "Light" }, +]; + +const ANALYTICS_OPTIONS: Array<{ value: boolean; label: string }> = [ + { value: true, label: "Share anonymous usage data" }, + { value: false, label: "Don't share" }, +]; + +const SETUP_LOGO_LINES = ["██████", "██ ██", "████ ██", "██ ██"]; + +/** First-time setup dialog: theme choice and analytics opt-in. */ +export class FirstTimeSetupComponent extends Container { + private step: "theme" | "analytics" = "theme"; + private themeIndex: number; + private analyticsIndex = 0; + private readonly options: FirstTimeSetupOptions; + + constructor(options: FirstTimeSetupOptions) { + super(); + this.options = options; + this.themeIndex = Math.max( + 0, + THEME_OPTIONS.findIndex((option) => option.value === options.detectedTheme), + ); + this.update(); + } + + // Rebuild the whole dialog on every change so theme previews recolor all text. + private update(): void { + this.clear(); + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("accent", SETUP_LOGO_LINES.join("\n")), 1, 0)); + this.addChild(new Spacer(1)); + this.addChild( + new Text(theme.fg("accent", theme.bold(`Welcome to ${APP_NAME}, the minimal coding agent.`)), 1, 0), + ); + this.addChild(new Spacer(1)); + + if (this.step === "theme") { + this.addChild(new Text(theme.fg("text", "Pick a theme."), 1, 0)); + this.addChild(new Text(theme.fg("muted", `Detected system appearance: ${this.options.detectedTheme}`), 1, 0)); + this.addChild(new Spacer(1)); + this.addOptionList( + THEME_OPTIONS.map((option) => option.label), + this.themeIndex, + ); + } else { + this.addChild(new Text(theme.fg("text", "Opt-in to anonymous usage data sharing?"), 1, 0)); + this.addChild( + new Text( + theme.fg( + "muted", + "Opting in stores a tracking identifier in settings.json and enables anonymous\nusage analytics. This helps us to better debug, reproduce, and resolve issues\nand bugs within Cactus. You can observe what is shared using /privacy and make\nchanges anytime in settings.json.", + ), + 1, + 0, + ), + ); + this.addChild(new Spacer(1)); + this.addOptionList( + ANALYTICS_OPTIONS.map((option) => option.label), + this.analyticsIndex, + ); + } + + this.addChild(new Spacer(1)); + this.addChild( + new Text( + rawKeyHint("↑↓", "navigate") + + " " + + keyHint("tui.select.confirm", this.step === "theme" ? "continue" : "finish") + + " " + + keyHint("tui.select.cancel", "skip setup"), + 1, + 0, + ), + ); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + } + + private addOptionList(labels: string[], selectedIndex: number): void { + for (let i = 0; i < labels.length; i++) { + const isSelected = i === selectedIndex; + const prefix = isSelected ? theme.fg("accent", "→ ") : " "; + const label = isSelected ? theme.fg("accent", labels[i]) : theme.fg("text", labels[i]); + this.addChild(new Text(`${prefix}${label}`, 1, 0)); + } + } + + private moveSelection(delta: number): void { + if (this.step === "theme") { + const next = Math.max(0, Math.min(THEME_OPTIONS.length - 1, this.themeIndex + delta)); + if (next !== this.themeIndex) { + this.themeIndex = next; + this.options.onThemePreview(THEME_OPTIONS[this.themeIndex].value); + } + } else { + this.analyticsIndex = Math.max(0, Math.min(ANALYTICS_OPTIONS.length - 1, this.analyticsIndex + delta)); + } + this.update(); + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "tui.select.up") || keyData === "k") { + this.moveSelection(-1); + } else if (kb.matches(keyData, "tui.select.down") || keyData === "j") { + this.moveSelection(1); + } else if (kb.matches(keyData, "tui.select.confirm") || keyData === "\n") { + if (this.step === "theme") { + this.step = "analytics"; + this.update(); + } else { + this.options.onSubmit({ + theme: THEME_OPTIONS[this.themeIndex].value, + shareAnalytics: ANALYTICS_OPTIONS[this.analyticsIndex].value, + }); + } + } else if (kb.matches(keyData, "tui.select.cancel")) { + this.options.onCancel(); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/footer.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/footer.ts new file mode 100644 index 000000000..f58366109 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/footer.ts @@ -0,0 +1,247 @@ +import { isAbsolute, relative, resolve, sep } from "node:path"; +import { type Component, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui"; +import type { AgentSession } from "../../../core/agent-session.ts"; +import { areExperimentalFeaturesEnabled } from "../../../core/experimental.ts"; +import type { ReadonlyFooterDataProvider } from "../../../core/footer-data-provider.ts"; +import { theme } from "../theme/theme.ts"; + +/** + * Sanitize text for display in a single-line status. + * Removes newlines, tabs, carriage returns, and other control characters. + */ +function sanitizeStatusText(text: string): string { + // Replace newlines, tabs, carriage returns with space, then collapse multiple spaces + return text + .replace(/[\r\n\t]/g, " ") + .replace(/ +/g, " ") + .trim(); +} + +/** + * Format token counts for compact footer display. + */ +function formatTokens(count: number): string { + if (!Number.isFinite(count)) return "?"; + if (count < 1000) return count.toString(); + if (count < 10000) return `${(count / 1000).toFixed(1)}k`; + if (count < 1000000) return `${Math.round(count / 1000)}k`; + if (count < 10000000) return `${(count / 1000000).toFixed(1)}M`; + return `${Math.round(count / 1000000)}M`; +} + +export function formatCwdForFooter(cwd: string, home: string | undefined): string { + if (!home) return cwd; + + const resolvedCwd = resolve(cwd); + const resolvedHome = resolve(home); + const relativeToHome = relative(resolvedHome, resolvedCwd); + const isInsideHome = + relativeToHome === "" || + (relativeToHome !== ".." && !relativeToHome.startsWith(`..${sep}`) && !isAbsolute(relativeToHome)); + + if (!isInsideHome) return cwd; + return relativeToHome === "" ? "~" : `~${sep}${relativeToHome}`; +} + +/** + * Footer component that shows pwd, token stats, and context usage. + * Computes token/context stats from session, gets git branch and extension statuses from provider. + */ +export class FooterComponent implements Component { + private autoCompactEnabled = true; + private session: AgentSession; + private footerData: ReadonlyFooterDataProvider; + + constructor(session: AgentSession, footerData: ReadonlyFooterDataProvider) { + this.session = session; + this.footerData = footerData; + } + + setSession(session: AgentSession): void { + this.session = session; + } + + setAutoCompactEnabled(enabled: boolean): void { + this.autoCompactEnabled = enabled; + } + + /** + * No-op: git branch caching now handled by provider. + * Kept for compatibility with existing call sites in interactive-mode. + */ + invalidate(): void { + // No-op: git branch is cached/invalidated by provider + } + + /** + * Clean up resources. + * Git watcher cleanup now handled by provider. + */ + dispose(): void { + // Git watcher cleanup handled by provider + } + + render(width: number): string[] { + const state = this.session.state; + + // Calculate cumulative usage from ALL session entries (not just post-compaction messages) + let totalInput = 0; + let totalOutput = 0; + let totalCacheRead = 0; + let totalCacheWrite = 0; + let totalCost = 0; + let latestCacheHitRate: number | undefined; + + for (const entry of this.session.sessionManager.getEntries()) { + if (entry.type === "message" && entry.message.role === "assistant") { + totalInput += entry.message.usage.input; + totalOutput += entry.message.usage.output; + totalCacheRead += entry.message.usage.cacheRead; + totalCacheWrite += entry.message.usage.cacheWrite; + totalCost += entry.message.usage.cost.total; + + const latestPromptTokens = + entry.message.usage.input + entry.message.usage.cacheRead + entry.message.usage.cacheWrite; + latestCacheHitRate = + latestPromptTokens > 0 ? (entry.message.usage.cacheRead / latestPromptTokens) * 100 : undefined; + } + } + + // Calculate context usage from session (handles compaction correctly). + // After compaction, tokens are unknown until the next LLM response. + const contextUsage = this.session.getContextUsage(); + const contextWindow = contextUsage?.contextWindow ?? state.model?.contextWindow ?? 0; + const contextPercentValue = contextUsage?.percent ?? 0; + const contextPercent = contextUsage?.percent !== null ? contextPercentValue.toFixed(1) : "?"; + + // Replace home directory with ~ + let pwd = formatCwdForFooter(this.session.sessionManager.getCwd(), process.env.HOME || process.env.USERPROFILE); + + // Add git branch if available + const branch = this.footerData.getGitBranch(); + if (branch) { + pwd = `${pwd} (${branch})`; + } + + // Add session name if set + const sessionName = this.session.sessionManager.getSessionName(); + if (sessionName) { + pwd = `${pwd} • ${sessionName}`; + } + + // Build stats line + const statsParts = []; + if (totalInput) statsParts.push(`↑${formatTokens(totalInput)}`); + if (totalOutput) statsParts.push(`↓${formatTokens(totalOutput)}`); + if (totalCacheRead) statsParts.push(`R${formatTokens(totalCacheRead)}`); + if (totalCacheWrite) statsParts.push(`W${formatTokens(totalCacheWrite)}`); + if ((totalCacheRead > 0 || totalCacheWrite > 0) && latestCacheHitRate !== undefined) { + statsParts.push(`CH${latestCacheHitRate.toFixed(1)}%`); + } + + // Show cost with "(sub)" indicator if using OAuth subscription + const usingSubscription = state.model ? this.session.modelRegistry.isUsingOAuth(state.model) : false; + if (totalCost || usingSubscription) { + const costStr = `$${totalCost.toFixed(3)}${usingSubscription ? " (sub)" : ""}`; + statsParts.push(costStr); + } + + // Colorize context percentage based on usage + let contextPercentStr: string; + const autoIndicator = this.autoCompactEnabled ? " (auto)" : ""; + const contextPercentDisplay = + contextPercent === "?" + ? `?/${formatTokens(contextWindow)}${autoIndicator}` + : `${contextPercent}%/${formatTokens(contextWindow)}${autoIndicator}`; + if (contextPercentValue > 90) { + contextPercentStr = theme.fg("error", contextPercentDisplay); + } else if (contextPercentValue > 70) { + contextPercentStr = theme.fg("warning", contextPercentDisplay); + } else { + contextPercentStr = contextPercentDisplay; + } + statsParts.push(contextPercentStr); + if (areExperimentalFeaturesEnabled()) { + statsParts.push(`${theme.fg("dim", "•")} ${theme.bold(theme.fg("warning", "xp"))}`); + } + + let statsLeft = statsParts.join(" "); + + // Add model name on the right side, plus thinking level if model supports it + const modelName = state.model?.id || "no-model"; + + let statsLeftWidth = visibleWidth(statsLeft); + + // If statsLeft is too wide, truncate it + if (statsLeftWidth > width) { + statsLeft = truncateToWidth(statsLeft, width, "..."); + statsLeftWidth = visibleWidth(statsLeft); + } + + // Calculate available space for padding (minimum 2 spaces between stats and model) + const minPadding = 2; + + // Add thinking level indicator if model supports reasoning + let rightSideWithoutProvider = modelName; + if (state.model?.reasoning) { + const thinkingLevel = state.thinkingLevel || "off"; + rightSideWithoutProvider = + thinkingLevel === "off" ? `${modelName} • thinking off` : `${modelName} • ${thinkingLevel}`; + } + + // Prepend the provider in parentheses if there are multiple providers and there's enough room + let rightSide = rightSideWithoutProvider; + if (this.footerData.getAvailableProviderCount() > 1 && state.model) { + rightSide = `(${state.model!.provider}) ${rightSideWithoutProvider}`; + if (statsLeftWidth + minPadding + visibleWidth(rightSide) > width) { + // Too wide, fall back + rightSide = rightSideWithoutProvider; + } + } + + const rightSideWidth = visibleWidth(rightSide); + const totalNeeded = statsLeftWidth + minPadding + rightSideWidth; + + let statsLine: string; + if (totalNeeded <= width) { + // Both fit - add padding to right-align model + const padding = " ".repeat(width - statsLeftWidth - rightSideWidth); + statsLine = statsLeft + padding + rightSide; + } else { + // Need to truncate right side + const availableForRight = width - statsLeftWidth - minPadding; + if (availableForRight > 0) { + const truncatedRight = truncateToWidth(rightSide, availableForRight, ""); + const truncatedRightWidth = visibleWidth(truncatedRight); + const padding = " ".repeat(Math.max(0, width - statsLeftWidth - truncatedRightWidth)); + statsLine = statsLeft + padding + truncatedRight; + } else { + // Not enough space for right side at all + statsLine = statsLeft; + } + } + + // Apply dim to each part separately. statsLeft may contain color codes (for context %) + // that end with a reset, which would clear an outer dim wrapper. So we dim the parts + // before and after the colored section independently. + const dimStatsLeft = theme.fg("dim", statsLeft); + const remainder = statsLine.slice(statsLeft.length); // padding + rightSide + const dimRemainder = theme.fg("dim", remainder); + + const pwdLine = truncateToWidth(theme.fg("dim", pwd), width, theme.fg("dim", "...")); + const lines = [pwdLine, dimStatsLeft + dimRemainder]; + + // Add extension statuses on a single line, sorted by key alphabetically + const extensionStatuses = this.footerData.getExtensionStatuses(); + if (extensionStatuses.size > 0) { + const sortedStatuses = Array.from(extensionStatuses.entries()) + .sort(([a], [b]) => a.localeCompare(b)) + .map(([, text]) => sanitizeStatusText(text)); + const statusLine = sortedStatuses.join(" "); + // Truncate to terminal width with dim ellipsis for consistency with footer style + lines.push(truncateToWidth(statusLine, width, theme.fg("dim", "..."))); + } + + return lines; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/index.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/index.ts new file mode 100644 index 000000000..067ee8940 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/index.ts @@ -0,0 +1,33 @@ +// UI Components for extensions +export { AssistantMessageComponent } from "./assistant-message.ts"; +export { BashExecutionComponent } from "./bash-execution.ts"; +export { BorderedLoader } from "./bordered-loader.ts"; +export { BranchSummaryMessageComponent } from "./branch-summary-message.ts"; +export { CompactionSummaryMessageComponent } from "./compaction-summary-message.ts"; +export { CustomEditor } from "./custom-editor.ts"; +export { CustomMessageComponent } from "./custom-message.ts"; +export { type RenderDiffOptions, renderDiff } from "./diff.ts"; +export { DynamicBorder } from "./dynamic-border.ts"; +export { ExtensionEditorComponent } from "./extension-editor.ts"; +export { ExtensionInputComponent } from "./extension-input.ts"; +export { ExtensionSelectorComponent } from "./extension-selector.ts"; +export { + FirstTimeSetupComponent, + type FirstTimeSetupOptions, + type FirstTimeSetupResult, +} from "./first-time-setup.ts"; +export { FooterComponent } from "./footer.ts"; +export { keyHint, keyText, rawKeyHint } from "./keybinding-hints.ts"; +export { ModelSelectorComponent } from "./model-selector.ts"; +export { type ModelsCallbacks, type ModelsConfig, ScopedModelsSelectorComponent } from "./scoped-models-selector.ts"; +export { SessionSelectorComponent } from "./session-selector.ts"; +export { type SettingsCallbacks, type SettingsConfig, SettingsSelectorComponent } from "./settings-selector.ts"; +export { SkillInvocationMessageComponent } from "./skill-invocation-message.ts"; +export { ThemeSelectorComponent } from "./theme-selector.ts"; +export { ThinkingSelectorComponent } from "./thinking-selector.ts"; +export { ToolExecutionComponent, type ToolExecutionOptions } from "./tool-execution.ts"; +export { TreeSelectorComponent } from "./tree-selector.ts"; +export { TrustSelectorComponent } from "./trust-selector.ts"; +export { UserMessageComponent } from "./user-message.ts"; +export { UserMessageSelectorComponent } from "./user-message-selector.ts"; +export { truncateToVisualLines, type VisualTruncateResult } from "./visual-truncate.ts"; diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/keybinding-hints.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/keybinding-hints.ts new file mode 100644 index 000000000..22601abd9 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/keybinding-hints.ts @@ -0,0 +1,48 @@ +/** + * Utilities for formatting keybinding hints in the UI. + */ + +import { getKeybindings, type Keybinding, type KeyId } from "@earendil-works/pi-tui"; +import { theme } from "../theme/theme.ts"; + +export interface KeyTextFormatOptions { + capitalize?: boolean; +} + +function formatKeyPart(part: string, options: KeyTextFormatOptions): string { + const displayPart = process.platform === "darwin" && part.toLowerCase() === "alt" ? "option" : part; + return options.capitalize ? displayPart.charAt(0).toUpperCase() + displayPart.slice(1) : displayPart; +} + +export function formatKeyText(key: string, options: KeyTextFormatOptions = {}): string { + return key + .split("/") + .map((k) => + k + .split("+") + .map((part) => formatKeyPart(part, options)) + .join("+"), + ) + .join("/"); +} + +function formatKeys(keys: KeyId[], options: KeyTextFormatOptions = {}): string { + if (keys.length === 0) return ""; + return formatKeyText(keys.join("/"), options); +} + +export function keyText(keybinding: Keybinding): string { + return formatKeys(getKeybindings().getKeys(keybinding)); +} + +export function keyDisplayText(keybinding: Keybinding): string { + return formatKeys(getKeybindings().getKeys(keybinding), { capitalize: true }); +} + +export function keyHint(keybinding: Keybinding, description: string): string { + return theme.fg("dim", keyText(keybinding)) + theme.fg("muted", ` ${description}`); +} + +export function rawKeyHint(key: string, description: string): string { + return theme.fg("dim", formatKeyText(key)) + theme.fg("muted", ` ${description}`); +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/model-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/model-selector.ts new file mode 100644 index 000000000..327119293 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/model-selector.ts @@ -0,0 +1,337 @@ +import { type Model, modelsAreEqual } from "@earendil-works/pi-ai"; +import { + Container, + type Focusable, + fuzzyFilter, + getKeybindings, + Input, + Spacer, + Text, + type TUI, +} from "@earendil-works/pi-tui"; +import type { ModelRegistry } from "../../../core/model-registry.ts"; +import type { SettingsManager } from "../../../core/settings-manager.ts"; +import { getModelSelectorSearchText } from "../model-search.ts"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint } from "./keybinding-hints.ts"; + +interface ModelItem { + provider: string; + id: string; + model: Model; +} + +interface ScopedModelItem { + model: Model; + thinkingLevel?: string; +} + +type ModelScope = "all" | "scoped"; + +/** + * Component that renders a model selector with search + */ +export class ModelSelectorComponent extends Container implements Focusable { + private searchInput: Input; + + // Focusable implementation - propagate to searchInput for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.searchInput.focused = value; + } + private listContainer: Container; + private allModels: ModelItem[] = []; + private scopedModelItems: ModelItem[] = []; + private activeModels: ModelItem[] = []; + private filteredModels: ModelItem[] = []; + private selectedIndex: number = 0; + private currentModel?: Model; + private settingsManager: SettingsManager; + private modelRegistry: ModelRegistry; + private onSelectCallback: (model: Model) => void; + private onCancelCallback: () => void; + private errorMessage?: string; + private tui: TUI; + private scopedModels: ReadonlyArray; + private scope: ModelScope = "all"; + private scopeText?: Text; + private scopeHintText?: Text; + + constructor( + tui: TUI, + currentModel: Model | undefined, + settingsManager: SettingsManager, + modelRegistry: ModelRegistry, + scopedModels: ReadonlyArray, + onSelect: (model: Model) => void, + onCancel: () => void, + initialSearchInput?: string, + ) { + super(); + + this.tui = tui; + this.currentModel = currentModel; + this.settingsManager = settingsManager; + this.modelRegistry = modelRegistry; + this.scopedModels = scopedModels; + this.scope = scopedModels.length > 0 ? "scoped" : "all"; + this.onSelectCallback = onSelect; + this.onCancelCallback = onCancel; + + // Add top border + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + + // Add hint about model filtering + if (scopedModels.length > 0) { + this.scopeText = new Text(this.getScopeText(), 0, 0); + this.addChild(this.scopeText); + this.scopeHintText = new Text(this.getScopeHintText(), 0, 0); + this.addChild(this.scopeHintText); + } else { + const hintText = "Only showing models from configured providers. Use /login to add providers."; + this.addChild(new Text(theme.fg("warning", hintText), 0, 0)); + } + this.addChild(new Spacer(1)); + + // Create search input + this.searchInput = new Input(); + if (initialSearchInput) { + this.searchInput.setValue(initialSearchInput); + } + this.searchInput.onSubmit = () => { + // Enter on search input selects the first filtered item + if (this.filteredModels[this.selectedIndex]) { + this.handleSelect(this.filteredModels[this.selectedIndex].model); + } + }; + this.addChild(this.searchInput); + + this.addChild(new Spacer(1)); + + // Create list container + this.listContainer = new Container(); + this.addChild(this.listContainer); + + this.addChild(new Spacer(1)); + + // Add bottom border + this.addChild(new DynamicBorder()); + + // Load models and do initial render + this.loadModels().then(() => { + if (initialSearchInput) { + this.filterModels(initialSearchInput); + } else { + this.updateList(); + } + // Request re-render after models are loaded + this.tui.requestRender(); + }); + } + + private async loadModels(): Promise { + let models: ModelItem[]; + + // Refresh to pick up any changes to models.json + this.modelRegistry.refresh(); + + // Check for models.json errors + const loadError = this.modelRegistry.getError(); + if (loadError) { + this.errorMessage = loadError; + } + + // Load available models (built-in models still work even if models.json failed) + try { + const availableModels = await this.modelRegistry.getAvailable(); + models = availableModels.map((model: Model) => ({ + provider: model.provider, + id: model.id, + model, + })); + } catch (error) { + this.allModels = []; + this.scopedModelItems = []; + this.activeModels = []; + this.filteredModels = []; + this.errorMessage = error instanceof Error ? error.message : String(error); + return; + } + + this.allModels = this.sortModels(models); + this.scopedModels = this.scopedModels.map((scoped) => { + const refreshed = this.modelRegistry.find(scoped.model.provider, scoped.model.id); + return refreshed ? { ...scoped, model: refreshed } : scoped; + }); + this.scopedModelItems = this.scopedModels.map((scoped) => ({ + provider: scoped.model.provider, + id: scoped.model.id, + model: scoped.model, + })); + this.activeModels = this.scope === "scoped" ? this.scopedModelItems : this.allModels; + this.filteredModels = this.activeModels; + const currentIndex = this.filteredModels.findIndex((item) => modelsAreEqual(this.currentModel, item.model)); + this.selectedIndex = + currentIndex >= 0 ? currentIndex : Math.min(this.selectedIndex, Math.max(0, this.filteredModels.length - 1)); + } + + private sortModels(models: ModelItem[]): ModelItem[] { + const sorted = [...models]; + // Sort: current model first, then by provider + sorted.sort((a, b) => { + const aIsCurrent = modelsAreEqual(this.currentModel, a.model); + const bIsCurrent = modelsAreEqual(this.currentModel, b.model); + if (aIsCurrent && !bIsCurrent) return -1; + if (!aIsCurrent && bIsCurrent) return 1; + return a.provider.localeCompare(b.provider); + }); + return sorted; + } + + private getScopeText(): string { + const allText = this.scope === "all" ? theme.fg("accent", "all") : theme.fg("muted", "all"); + const scopedText = this.scope === "scoped" ? theme.fg("accent", "scoped") : theme.fg("muted", "scoped"); + return `${theme.fg("muted", "Scope: ")}${allText}${theme.fg("muted", " | ")}${scopedText}`; + } + + private getScopeHintText(): string { + return keyHint("tui.input.tab", "scope") + theme.fg("muted", " (all/scoped)"); + } + + private setScope(scope: ModelScope): void { + if (this.scope === scope) return; + this.scope = scope; + this.activeModels = this.scope === "scoped" ? this.scopedModelItems : this.allModels; + const currentIndex = this.activeModels.findIndex((item) => modelsAreEqual(this.currentModel, item.model)); + this.selectedIndex = currentIndex >= 0 ? currentIndex : 0; + this.filterModels(this.searchInput.getValue()); + if (this.scopeText) { + this.scopeText.setText(this.getScopeText()); + } + } + + private filterModels(query: string): void { + this.filteredModels = query + ? fuzzyFilter(this.activeModels, query, ({ id, provider, model }) => + getModelSelectorSearchText({ id, provider, name: model.name }), + ) + : this.activeModels; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredModels.length - 1)); + this.updateList(); + } + + private updateList(): void { + this.listContainer.clear(); + + const maxVisible = 10; + const startIndex = Math.max( + 0, + Math.min(this.selectedIndex - Math.floor(maxVisible / 2), this.filteredModels.length - maxVisible), + ); + const endIndex = Math.min(startIndex + maxVisible, this.filteredModels.length); + + // Show visible slice of filtered models + for (let i = startIndex; i < endIndex; i++) { + const item = this.filteredModels[i]; + if (!item) continue; + + const isSelected = i === this.selectedIndex; + const isCurrent = modelsAreEqual(this.currentModel, item.model); + + let line = ""; + if (isSelected) { + const prefix = theme.fg("accent", "→ "); + const modelText = `${item.id}`; + const providerBadge = theme.fg("muted", `[${item.provider}]`); + const checkmark = isCurrent ? theme.fg("success", " ✓") : ""; + line = `${prefix + theme.fg("accent", modelText)} ${providerBadge}${checkmark}`; + } else { + const modelText = ` ${item.id}`; + const providerBadge = theme.fg("muted", `[${item.provider}]`); + const checkmark = isCurrent ? theme.fg("success", " ✓") : ""; + line = `${modelText} ${providerBadge}${checkmark}`; + } + + this.listContainer.addChild(new Text(line, 0, 0)); + } + + // Add scroll indicator if needed + if (startIndex > 0 || endIndex < this.filteredModels.length) { + const scrollInfo = theme.fg("muted", ` (${this.selectedIndex + 1}/${this.filteredModels.length})`); + this.listContainer.addChild(new Text(scrollInfo, 0, 0)); + } + + // Show error message or "no results" if empty + if (this.errorMessage) { + // Show error in red + const errorLines = this.errorMessage.split("\n"); + for (const line of errorLines) { + this.listContainer.addChild(new Text(theme.fg("error", line), 0, 0)); + } + } else if (this.filteredModels.length === 0) { + this.listContainer.addChild(new Text(theme.fg("muted", " No matching models"), 0, 0)); + } else { + const selected = this.filteredModels[this.selectedIndex]; + this.listContainer.addChild(new Spacer(1)); + this.listContainer.addChild(new Text(theme.fg("muted", ` Model Name: ${selected.model.name}`), 0, 0)); + } + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "tui.input.tab")) { + if (this.scopedModelItems.length > 0) { + const nextScope: ModelScope = this.scope === "all" ? "scoped" : "all"; + this.setScope(nextScope); + if (this.scopeHintText) { + this.scopeHintText.setText(this.getScopeHintText()); + } + } + return; + } + // Up arrow - wrap to bottom when at top + if (kb.matches(keyData, "tui.select.up")) { + if (this.filteredModels.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredModels.length - 1 : this.selectedIndex - 1; + this.updateList(); + } + // Down arrow - wrap to top when at bottom + else if (kb.matches(keyData, "tui.select.down")) { + if (this.filteredModels.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredModels.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); + } + // Enter + else if (kb.matches(keyData, "tui.select.confirm")) { + const selectedModel = this.filteredModels[this.selectedIndex]; + if (selectedModel) { + this.handleSelect(selectedModel.model); + } + } + // Escape or Ctrl+C + else if (kb.matches(keyData, "tui.select.cancel")) { + this.onCancelCallback(); + } + // Pass everything else to search input + else { + this.searchInput.handleInput(keyData); + this.filterModels(this.searchInput.getValue()); + } + } + + private handleSelect(model: Model): void { + // Save as new default + this.settingsManager.setDefaultModelAndProvider(model.provider, model.id); + this.onSelectCallback(model); + } + + getSearchInput(): Input { + return this.searchInput; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts new file mode 100644 index 000000000..772e3af06 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/scoped-models-selector.ts @@ -0,0 +1,360 @@ +import type { Model } from "@earendil-works/pi-ai"; +import { + Container, + type Focusable, + fuzzyFilter, + getKeybindings, + Input, + Key, + matchesKey, + Spacer, + Text, +} from "@earendil-works/pi-tui"; +import { getModelSearchText } from "../model-search.ts"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyText } from "./keybinding-hints.ts"; + +// EnabledIds: null = all enabled (no filter), string[] = explicit ordered list +type EnabledIds = string[] | null; + +function isEnabled(enabledIds: EnabledIds, id: string): boolean { + return enabledIds === null || enabledIds.includes(id); +} + +function toggle(enabledIds: EnabledIds, id: string): EnabledIds { + if (enabledIds === null) return [id]; // First toggle: start with only this one + const index = enabledIds.indexOf(id); + if (index >= 0) return [...enabledIds.slice(0, index), ...enabledIds.slice(index + 1)]; + return [...enabledIds, id]; +} + +function enableAll(enabledIds: EnabledIds, allIds: string[], targetIds?: string[]): EnabledIds { + if (enabledIds === null) return null; // Already all enabled + const targets = targetIds ?? allIds; + const result = [...enabledIds]; + for (const id of targets) { + if (!result.includes(id)) result.push(id); + } + return result.length === allIds.length ? null : result; +} + +function clearAll(enabledIds: EnabledIds, allIds: string[], targetIds?: string[]): EnabledIds { + if (enabledIds === null) { + return targetIds ? allIds.filter((id) => !targetIds.includes(id)) : []; + } + const targets = new Set(targetIds ?? enabledIds); + return enabledIds.filter((id) => !targets.has(id)); +} + +function move(enabledIds: EnabledIds, id: string, delta: number): EnabledIds { + if (enabledIds === null) return null; + const list = [...enabledIds]; + const index = list.indexOf(id); + if (index < 0) return list; + const newIndex = index + delta; + if (newIndex < 0 || newIndex >= list.length) return list; + const result = [...list]; + [result[index], result[newIndex]] = [result[newIndex], result[index]]; + return result; +} + +function getSortedIds(enabledIds: EnabledIds, allIds: string[]): string[] { + if (enabledIds === null) return allIds; + const enabledSet = new Set(enabledIds); + return [...enabledIds, ...allIds.filter((id) => !enabledSet.has(id))]; +} + +interface ModelItem { + fullId: string; + model: Model; + enabled: boolean; +} + +export interface ModelsConfig { + allModels: Model[]; + enabledModelIds: string[] | null; +} + +export interface ModelsCallbacks { + /** Called whenever the enabled model set or order changes (session-only, no persist) */ + onChange: (enabledModelIds: string[] | null) => void | Promise; + /** Called when user wants to persist current selection to settings */ + onPersist: (enabledModelIds: string[] | null) => void | Promise; + onCancel: () => void; +} + +/** + * Component for enabling/disabling models for Ctrl+P cycling. + * Changes are session-only until explicitly persisted with Ctrl+S. + */ +export class ScopedModelsSelectorComponent extends Container implements Focusable { + private modelsById: Map> = new Map(); + private allIds: string[] = []; + private enabledIds: EnabledIds = null; + private filteredItems: ModelItem[] = []; + private selectedIndex = 0; + private searchInput: Input; + + // Focusable implementation - propagate to searchInput for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.searchInput.focused = value; + } + private listContainer: Container; + private footerText: Text; + private callbacks: ModelsCallbacks; + private maxVisible = 8; + private isDirty = false; + + constructor(config: ModelsConfig, callbacks: ModelsCallbacks) { + super(); + this.callbacks = callbacks; + + for (const model of config.allModels) { + const fullId = `${model.provider}/${model.id}`; + this.modelsById.set(fullId, model); + this.allIds.push(fullId); + } + + this.enabledIds = config.enabledModelIds === null ? null : [...config.enabledModelIds]; + this.filteredItems = this.buildItems(); + + // Header + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("accent", theme.bold("Model Configuration")), 0, 0)); + this.addChild( + new Text(theme.fg("muted", `Session-only. ${keyText("app.models.save")} to save to settings.`), 0, 0), + ); + this.addChild(new Spacer(1)); + + // Search input + this.searchInput = new Input(); + this.addChild(this.searchInput); + this.addChild(new Spacer(1)); + + // List container + this.listContainer = new Container(); + this.addChild(this.listContainer); + + // Footer hint + this.addChild(new Spacer(1)); + this.footerText = new Text(this.getFooterText(), 0, 0); + this.addChild(this.footerText); + + this.addChild(new DynamicBorder()); + this.updateList(); + } + + private buildItems(): ModelItem[] { + // Filter out IDs that no longer have a corresponding model (e.g., after logout) + return getSortedIds(this.enabledIds, this.allIds) + .filter((id) => this.modelsById.has(id)) + .map((id) => ({ + fullId: id, + model: this.modelsById.get(id)!, + enabled: isEnabled(this.enabledIds, id), + })); + } + + private getFooterText(): string { + const enabledCount = this.enabledIds?.length ?? this.allIds.length; + const allEnabled = this.enabledIds === null; + const countText = allEnabled ? "all enabled" : `${enabledCount}/${this.allIds.length} enabled`; + const parts = [ + `${keyText("tui.select.confirm")} toggle`, + `${keyText("app.models.enableAll")} all`, + `${keyText("app.models.clearAll")} clear`, + `${keyText("app.models.toggleProvider")} provider`, + `${keyText("app.models.reorderUp")}/${keyText("app.models.reorderDown")} reorder`, + `${keyText("app.models.save")} save`, + countText, + ]; + return this.isDirty + ? theme.fg("dim", ` ${parts.join(" · ")} `) + theme.fg("warning", "(unsaved)") + : theme.fg("dim", ` ${parts.join(" · ")}`); + } + + private refresh(): void { + const query = this.searchInput.getValue(); + const items = this.buildItems(); + this.filteredItems = query + ? fuzzyFilter(items, query, (i) => + getModelSearchText({ id: i.model.id, provider: i.model.provider, name: i.model.name }), + ) + : items; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + this.updateList(); + this.footerText.setText(this.getFooterText()); + } + + private notifyChange(): void { + this.callbacks.onChange(this.enabledIds === null ? null : [...this.enabledIds]); + } + + private updateList(): void { + this.listContainer.clear(); + + if (this.filteredItems.length === 0) { + this.listContainer.addChild(new Text(theme.fg("muted", " No matching models"), 0, 0)); + return; + } + + const startIndex = Math.max( + 0, + Math.min(this.selectedIndex - Math.floor(this.maxVisible / 2), this.filteredItems.length - this.maxVisible), + ); + const endIndex = Math.min(startIndex + this.maxVisible, this.filteredItems.length); + const allEnabled = this.enabledIds === null; + + for (let i = startIndex; i < endIndex; i++) { + const item = this.filteredItems[i]!; + const isSelected = i === this.selectedIndex; + const prefix = isSelected ? theme.fg("accent", "→ ") : " "; + const modelText = isSelected ? theme.fg("accent", item.model.id) : item.model.id; + const providerBadge = theme.fg("muted", ` [${item.model.provider}]`); + const status = allEnabled ? "" : item.enabled ? theme.fg("success", " ✓") : theme.fg("dim", " ✗"); + this.listContainer.addChild(new Text(`${prefix}${modelText}${providerBadge}${status}`, 0, 0)); + } + + // Add scroll indicator if needed + if (startIndex > 0 || endIndex < this.filteredItems.length) { + this.listContainer.addChild( + new Text(theme.fg("muted", ` (${this.selectedIndex + 1}/${this.filteredItems.length})`), 0, 0), + ); + } + + if (this.filteredItems.length > 0) { + const selected = this.filteredItems[this.selectedIndex]; + this.listContainer.addChild(new Spacer(1)); + this.listContainer.addChild(new Text(theme.fg("muted", ` Model Name: ${selected.model.name}`), 0, 0)); + } + } + + handleInput(data: string): void { + const kb = getKeybindings(); + + // Navigation + if (kb.matches(data, "tui.select.up")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.updateList(); + return; + } + if (kb.matches(data, "tui.select.down")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); + return; + } + + // Reorder enabled models + const reorderUp = kb.matches(data, "app.models.reorderUp"); + const reorderDown = kb.matches(data, "app.models.reorderDown"); + if (reorderUp || reorderDown) { + if (this.enabledIds === null) return; + const item = this.filteredItems[this.selectedIndex]; + if (item && isEnabled(this.enabledIds, item.fullId)) { + const delta = reorderUp ? -1 : 1; + const currentIndex = this.enabledIds.indexOf(item.fullId); + const newIndex = currentIndex + delta; + // Only move if within bounds + if (newIndex >= 0 && newIndex < this.enabledIds.length) { + this.enabledIds = move(this.enabledIds, item.fullId, delta); + this.isDirty = true; + this.selectedIndex += delta; + this.refresh(); + this.notifyChange(); + } + } + return; + } + + // Toggle on Enter + if (kb.matches(data, "tui.select.confirm")) { + const item = this.filteredItems[this.selectedIndex]; + if (item) { + this.enabledIds = toggle(this.enabledIds, item.fullId); + this.isDirty = true; + this.refresh(); + this.notifyChange(); + } + return; + } + + // Enable all (filtered if search active, otherwise all) + if (kb.matches(data, "app.models.enableAll")) { + const targetIds = this.searchInput.getValue() ? this.filteredItems.map((i) => i.fullId) : undefined; + this.enabledIds = enableAll(this.enabledIds, this.allIds, targetIds); + this.isDirty = true; + this.refresh(); + this.notifyChange(); + return; + } + + // Clear all (filtered if search active, otherwise all) + if (kb.matches(data, "app.models.clearAll")) { + const targetIds = this.searchInput.getValue() ? this.filteredItems.map((i) => i.fullId) : undefined; + this.enabledIds = clearAll(this.enabledIds, this.allIds, targetIds); + this.isDirty = true; + this.refresh(); + this.notifyChange(); + return; + } + + // Toggle provider of current item + if (kb.matches(data, "app.models.toggleProvider")) { + const item = this.filteredItems[this.selectedIndex]; + if (item) { + const provider = item.model.provider; + const providerIds = this.allIds.filter((id) => this.modelsById.get(id)!.provider === provider); + const allEnabled = providerIds.every((id) => isEnabled(this.enabledIds, id)); + this.enabledIds = allEnabled + ? clearAll(this.enabledIds, this.allIds, providerIds) + : enableAll(this.enabledIds, this.allIds, providerIds); + this.isDirty = true; + this.refresh(); + this.notifyChange(); + } + return; + } + + // Save/persist to settings + if (kb.matches(data, "app.models.save")) { + this.callbacks.onPersist(this.enabledIds === null ? null : [...this.enabledIds]); + this.isDirty = false; + this.footerText.setText(this.getFooterText()); + return; + } + + // Ctrl+C - clear search or cancel if empty + if (matchesKey(data, Key.ctrl("c"))) { + if (this.searchInput.getValue()) { + this.searchInput.setValue(""); + this.refresh(); + } else { + this.callbacks.onCancel(); + } + return; + } + + // Escape - cancel + if (matchesKey(data, Key.escape)) { + this.callbacks.onCancel(); + return; + } + + // Pass everything else to search input + this.searchInput.handleInput(data); + this.refresh(); + } + + getSearchInput(): Input { + return this.searchInput; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/session-selector-search.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/session-selector-search.ts new file mode 100644 index 000000000..9b5bf2327 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/session-selector-search.ts @@ -0,0 +1,194 @@ +import { fuzzyMatch } from "@earendil-works/pi-tui"; +import type { SessionInfo } from "../../../core/session-manager.ts"; + +export type SortMode = "threaded" | "recent" | "relevance"; + +export type NameFilter = "all" | "named"; + +export interface ParsedSearchQuery { + mode: "tokens" | "regex"; + tokens: { kind: "fuzzy" | "phrase"; value: string }[]; + regex: RegExp | null; + /** If set, parsing failed and we should treat query as non-matching. */ + error?: string; +} + +export interface MatchResult { + matches: boolean; + /** Lower is better; only meaningful when matches === true */ + score: number; +} + +function normalizeWhitespaceLower(text: string): string { + return text.toLowerCase().replace(/\s+/g, " ").trim(); +} + +function getSessionSearchText(session: SessionInfo): string { + return `${session.id} ${session.name ?? ""} ${session.allMessagesText} ${session.cwd}`; +} + +export function hasSessionName(session: SessionInfo): boolean { + return Boolean(session.name?.trim()); +} + +function matchesNameFilter(session: SessionInfo, filter: NameFilter): boolean { + if (filter === "all") return true; + return hasSessionName(session); +} + +export function parseSearchQuery(query: string): ParsedSearchQuery { + const trimmed = query.trim(); + if (!trimmed) { + return { mode: "tokens", tokens: [], regex: null }; + } + + // Regex mode: re: + if (trimmed.startsWith("re:")) { + const pattern = trimmed.slice(3).trim(); + if (!pattern) { + return { mode: "regex", tokens: [], regex: null, error: "Empty regex" }; + } + try { + return { mode: "regex", tokens: [], regex: new RegExp(pattern, "i") }; + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + return { mode: "regex", tokens: [], regex: null, error: msg }; + } + } + + // Token mode with quote support. + // Example: foo "node cve" bar + const tokens: { kind: "fuzzy" | "phrase"; value: string }[] = []; + let buf = ""; + let inQuote = false; + let hadUnclosedQuote = false; + + const flush = (kind: "fuzzy" | "phrase"): void => { + const v = buf.trim(); + buf = ""; + if (!v) return; + tokens.push({ kind, value: v }); + }; + + for (let i = 0; i < trimmed.length; i++) { + const ch = trimmed[i]!; + if (ch === '"') { + if (inQuote) { + flush("phrase"); + inQuote = false; + } else { + flush("fuzzy"); + inQuote = true; + } + continue; + } + + if (!inQuote && /\s/.test(ch)) { + flush("fuzzy"); + continue; + } + + buf += ch; + } + + if (inQuote) { + hadUnclosedQuote = true; + } + + // If quotes were unbalanced, fall back to plain whitespace tokenization. + if (hadUnclosedQuote) { + return { + mode: "tokens", + tokens: trimmed + .split(/\s+/) + .map((t) => t.trim()) + .filter((t) => t.length > 0) + .map((t) => ({ kind: "fuzzy" as const, value: t })), + regex: null, + }; + } + + flush(inQuote ? "phrase" : "fuzzy"); + + return { mode: "tokens", tokens, regex: null }; +} + +export function matchSession(session: SessionInfo, parsed: ParsedSearchQuery): MatchResult { + const text = getSessionSearchText(session); + + if (parsed.mode === "regex") { + if (!parsed.regex) { + return { matches: false, score: 0 }; + } + const idx = text.search(parsed.regex); + if (idx < 0) return { matches: false, score: 0 }; + return { matches: true, score: idx * 0.1 }; + } + + if (parsed.tokens.length === 0) { + return { matches: true, score: 0 }; + } + + let totalScore = 0; + let normalizedText: string | null = null; + + for (const token of parsed.tokens) { + if (token.kind === "phrase") { + if (normalizedText === null) { + normalizedText = normalizeWhitespaceLower(text); + } + const phrase = normalizeWhitespaceLower(token.value); + if (!phrase) continue; + const idx = normalizedText.indexOf(phrase); + if (idx < 0) return { matches: false, score: 0 }; + totalScore += idx * 0.1; + continue; + } + + const m = fuzzyMatch(token.value, text); + if (!m.matches) return { matches: false, score: 0 }; + totalScore += m.score; + } + + return { matches: true, score: totalScore }; +} + +export function filterAndSortSessions( + sessions: SessionInfo[], + query: string, + sortMode: SortMode, + nameFilter: NameFilter = "all", +): SessionInfo[] { + const nameFiltered = + nameFilter === "all" ? sessions : sessions.filter((session) => matchesNameFilter(session, nameFilter)); + const trimmed = query.trim(); + if (!trimmed) return nameFiltered; + + const parsed = parseSearchQuery(query); + if (parsed.error) return []; + + // Recent mode: filter only, keep incoming order. + if (sortMode === "recent") { + const filtered: SessionInfo[] = []; + for (const s of nameFiltered) { + const res = matchSession(s, parsed); + if (res.matches) filtered.push(s); + } + return filtered; + } + + // Relevance mode: sort by score, tie-break by modified desc. + const scored: { session: SessionInfo; score: number }[] = []; + for (const s of nameFiltered) { + const res = matchSession(s, parsed); + if (!res.matches) continue; + scored.push({ session: s, score: res.score }); + } + + scored.sort((a, b) => { + if (a.score !== b.score) return a.score - b.score; + return b.session.modified.getTime() - a.session.modified.getTime(); + }); + + return scored.map((r) => r.session); +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/session-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/session-selector.ts new file mode 100644 index 000000000..4949eee51 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/session-selector.ts @@ -0,0 +1,1031 @@ +import { spawnSync } from "node:child_process"; +import { existsSync } from "node:fs"; +import { unlink } from "node:fs/promises"; +import * as os from "node:os"; +import { + type Component, + Container, + type Focusable, + getKeybindings, + Input, + Spacer, + Text, + truncateToWidth, + visibleWidth, +} from "@earendil-works/pi-tui"; +import { KeybindingsManager } from "../../../core/keybindings.ts"; +import type { SessionInfo, SessionListProgress } from "../../../core/session-manager.ts"; +import { canonicalizePath as _canonicalizePath } from "../../../utils/paths.ts"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint, keyText } from "./keybinding-hints.ts"; +import { filterAndSortSessions, hasSessionName, type NameFilter, type SortMode } from "./session-selector-search.ts"; + +type SessionScope = "current" | "all"; + +function shortenPath(path: string): string { + const home = os.homedir(); + if (!path) return path; + if (path.startsWith(home)) { + return `~${path.slice(home.length)}`; + } + return path; +} + +function formatSessionDate(date: Date): string { + const now = new Date(); + const diffMs = now.getTime() - date.getTime(); + const diffMins = Math.floor(diffMs / 60000); + const diffHours = Math.floor(diffMs / 3600000); + const diffDays = Math.floor(diffMs / 86400000); + + if (diffMins < 1) return "now"; + if (diffMins < 60) return `${diffMins}m`; + if (diffHours < 24) return `${diffHours}h`; + if (diffDays < 7) return `${diffDays}d`; + if (diffDays < 30) return `${Math.floor(diffDays / 7)}w`; + if (diffDays < 365) return `${Math.floor(diffDays / 30)}mo`; + return `${Math.floor(diffDays / 365)}y`; +} + +function canonicalizePath(path: string | undefined): string | undefined { + if (!path) return path; + return _canonicalizePath(path); +} + +class SessionSelectorHeader implements Component { + private scope: SessionScope; + private sortMode: SortMode; + private nameFilter: NameFilter; + private requestRender: () => void; + private loading = false; + private loadProgress: { loaded: number; total: number } | null = null; + private showPath = false; + private confirmingDeletePath: string | null = null; + private statusMessage: { type: "info" | "error"; message: string } | null = null; + private statusTimeout: ReturnType | null = null; + private showRenameHint = false; + + constructor(scope: SessionScope, sortMode: SortMode, nameFilter: NameFilter, requestRender: () => void) { + this.scope = scope; + this.sortMode = sortMode; + this.nameFilter = nameFilter; + this.requestRender = requestRender; + } + + setScope(scope: SessionScope): void { + this.scope = scope; + } + + setSortMode(sortMode: SortMode): void { + this.sortMode = sortMode; + } + + setNameFilter(nameFilter: NameFilter): void { + this.nameFilter = nameFilter; + } + + setLoading(loading: boolean): void { + this.loading = loading; + // Progress is scoped to the current load; clear whenever the loading state is set + this.loadProgress = null; + } + + setProgress(loaded: number, total: number): void { + this.loadProgress = { loaded, total }; + } + + setShowPath(showPath: boolean): void { + this.showPath = showPath; + } + + setShowRenameHint(show: boolean): void { + this.showRenameHint = show; + } + + setConfirmingDeletePath(path: string | null): void { + this.confirmingDeletePath = path; + } + + private clearStatusTimeout(): void { + if (!this.statusTimeout) return; + clearTimeout(this.statusTimeout); + this.statusTimeout = null; + } + + setStatusMessage(msg: { type: "info" | "error"; message: string } | null, autoHideMs?: number): void { + this.clearStatusTimeout(); + this.statusMessage = msg; + if (!msg || !autoHideMs) return; + + this.statusTimeout = setTimeout(() => { + this.statusMessage = null; + this.statusTimeout = null; + this.requestRender(); + }, autoHideMs); + } + + invalidate(): void {} + + render(width: number): string[] { + const title = this.scope === "current" ? "Resume Session (Current Folder)" : "Resume Session (All)"; + const leftText = theme.bold(title); + + const sortLabel = this.sortMode === "threaded" ? "Threaded" : this.sortMode === "recent" ? "Recent" : "Fuzzy"; + const sortText = theme.fg("muted", "Sort: ") + theme.fg("accent", sortLabel); + + const nameLabel = this.nameFilter === "all" ? "All" : "Named"; + const nameText = theme.fg("muted", "Name: ") + theme.fg("accent", nameLabel); + + let scopeText: string; + if (this.loading) { + const progressText = this.loadProgress ? `${this.loadProgress.loaded}/${this.loadProgress.total}` : "..."; + scopeText = `${theme.fg("muted", "○ Current Folder | ")}${theme.fg("accent", `Loading ${progressText}`)}`; + } else if (this.scope === "current") { + scopeText = `${theme.fg("accent", "◉ Current Folder")}${theme.fg("muted", " | ○ All")}`; + } else { + scopeText = `${theme.fg("muted", "○ Current Folder | ")}${theme.fg("accent", "◉ All")}`; + } + + const rightText = truncateToWidth(`${scopeText} ${nameText} ${sortText}`, width, ""); + const availableLeft = Math.max(0, width - visibleWidth(rightText) - 1); + const left = truncateToWidth(leftText, availableLeft, ""); + const spacing = Math.max(0, width - visibleWidth(left) - visibleWidth(rightText)); + + // Build hint lines - changes based on state (all branches truncate to width) + let hintLine1: string; + let hintLine2: string; + if (this.confirmingDeletePath !== null) { + const confirmHint = `Delete session? ${keyHint("tui.select.confirm", "confirm")} · ${keyHint("tui.select.cancel", "cancel")}`; + hintLine1 = theme.fg("error", truncateToWidth(confirmHint, width, "…")); + hintLine2 = ""; + } else if (this.statusMessage) { + const color = this.statusMessage.type === "error" ? "error" : "accent"; + hintLine1 = theme.fg(color, truncateToWidth(this.statusMessage.message, width, "…")); + hintLine2 = ""; + } else { + const pathState = this.showPath ? "(on)" : "(off)"; + const sep = theme.fg("muted", " · "); + const hint1 = + keyHint("tui.input.tab", "scope") + sep + theme.fg("muted", 're: regex · "phrase" exact'); + const hint2Parts = [ + keyHint("app.session.toggleSort", "sort"), + keyHint("app.session.toggleNamedFilter", "named"), + keyHint("app.session.delete", "delete"), + keyHint("app.session.togglePath", `path ${pathState}`), + ]; + if (this.showRenameHint) { + hint2Parts.push(keyHint("app.session.rename", "rename")); + } + const hint2 = hint2Parts.join(sep); + hintLine1 = truncateToWidth(hint1, width, "…"); + hintLine2 = truncateToWidth(hint2, width, "…"); + } + + return [`${left}${" ".repeat(spacing)}${rightText}`, hintLine1, hintLine2]; + } +} + +/** A session tree node for hierarchical display */ +interface SessionTreeNode { + session: SessionInfo; + children: SessionTreeNode[]; + latestActivity: number; +} + +/** Flattened node for display with tree structure info */ +interface FlatSessionNode { + session: SessionInfo; + depth: number; + isLast: boolean; + /** For each ancestor level, whether there are more siblings after it */ + ancestorContinues: boolean[]; +} + +/** + * Build a tree structure from sessions based on parentSessionPath. + * Returns root nodes sorted by modified date (descending). + */ +function buildSessionTree(sessions: SessionInfo[]): SessionTreeNode[] { + const byPath = new Map(); + + for (const session of sessions) { + const sessionPath = canonicalizePath(session.path) ?? session.path; + byPath.set(sessionPath, { session, children: [], latestActivity: session.modified.getTime() }); + } + + const roots: SessionTreeNode[] = []; + + for (const session of sessions) { + const sessionPath = canonicalizePath(session.path) ?? session.path; + const node = byPath.get(sessionPath)!; + const parentPath = canonicalizePath(session.parentSessionPath); + + if (parentPath && byPath.has(parentPath)) { + byPath.get(parentPath)!.children.push(node); + } else { + roots.push(node); + } + } + + const updateLatestActivity = (node: SessionTreeNode): number => { + let latestActivity = node.session.modified.getTime(); + for (const child of node.children) { + latestActivity = Math.max(latestActivity, updateLatestActivity(child)); + } + node.latestActivity = latestActivity; + return latestActivity; + }; + + for (const root of roots) { + updateLatestActivity(root); + } + + // Sort children and roots by latest activity in each subtree (descending) + const sortNodes = (nodes: SessionTreeNode[]): void => { + nodes.sort((a, b) => b.latestActivity - a.latestActivity); + for (const node of nodes) { + sortNodes(node.children); + } + }; + sortNodes(roots); + + return roots; +} + +/** + * Flatten tree into display list with tree structure metadata. + */ +function flattenSessionTree(roots: SessionTreeNode[]): FlatSessionNode[] { + const result: FlatSessionNode[] = []; + + const walk = (node: SessionTreeNode, depth: number, ancestorContinues: boolean[], isLast: boolean): void => { + result.push({ session: node.session, depth, isLast, ancestorContinues }); + + for (let i = 0; i < node.children.length; i++) { + const childIsLast = i === node.children.length - 1; + // Only show continuation line for non-root ancestors + const continues = depth > 0 ? !isLast : false; + walk(node.children[i]!, depth + 1, [...ancestorContinues, continues], childIsLast); + } + }; + + for (let i = 0; i < roots.length; i++) { + walk(roots[i]!, 0, [], i === roots.length - 1); + } + + return result; +} + +/** + * Custom session list component with multi-line items and search + */ +class SessionList implements Component, Focusable { + public getSelectedSessionPath(): string | undefined { + const selected = this.filteredSessions[this.selectedIndex]; + return selected?.session.path; + } + private allSessions: SessionInfo[] = []; + private filteredSessions: FlatSessionNode[] = []; + private selectedIndex: number = 0; + private searchInput: Input; + private showCwd = false; + private sortMode: SortMode = "threaded"; + private nameFilter: NameFilter = "all"; + private keybindings: KeybindingsManager; + private showPath = false; + private confirmingDeletePath: string | null = null; + private currentSessionCanonicalPath?: string; + public onSelect?: (sessionPath: string) => void; + public onCancel?: () => void; + public onExit: () => void = () => {}; + public onToggleScope?: () => void; + public onToggleSort?: () => void; + public onToggleNameFilter?: () => void; + public onTogglePath?: (showPath: boolean) => void; + public onDeleteConfirmationChange?: (path: string | null) => void; + public onDeleteSession?: (sessionPath: string) => Promise; + public onRenameSession?: (sessionPath: string) => void; + public onError?: (message: string) => void; + private maxVisible: number = 10; // Max sessions visible (one line each) + + // Focusable implementation - propagate to searchInput for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.searchInput.focused = value; + } + + constructor( + sessions: SessionInfo[], + showCwd: boolean, + sortMode: SortMode, + nameFilter: NameFilter, + keybindings: KeybindingsManager, + currentSessionFilePath?: string, + ) { + this.allSessions = sessions; + this.filteredSessions = []; + this.searchInput = new Input(); + this.showCwd = showCwd; + this.sortMode = sortMode; + this.nameFilter = nameFilter; + this.keybindings = keybindings; + this.currentSessionCanonicalPath = canonicalizePath(currentSessionFilePath); + this.filterSessions(""); + + // Handle Enter in search input - select current item + this.searchInput.onSubmit = () => { + if (this.filteredSessions[this.selectedIndex]) { + const selected = this.filteredSessions[this.selectedIndex]; + if (this.onSelect) { + this.onSelect(selected.session.path); + } + } + }; + } + + setSortMode(sortMode: SortMode): void { + this.sortMode = sortMode; + this.filterSessions(this.searchInput.getValue()); + } + + setNameFilter(nameFilter: NameFilter): void { + this.nameFilter = nameFilter; + this.filterSessions(this.searchInput.getValue()); + } + + setSessions(sessions: SessionInfo[], showCwd: boolean): void { + this.allSessions = sessions; + this.showCwd = showCwd; + this.filterSessions(this.searchInput.getValue()); + } + + private filterSessions(query: string): void { + const trimmed = query.trim(); + const nameFiltered = + this.nameFilter === "all" ? this.allSessions : this.allSessions.filter((session) => hasSessionName(session)); + + if (this.sortMode === "threaded" && !trimmed) { + // Threaded mode without search: show tree structure + const roots = buildSessionTree(nameFiltered); + this.filteredSessions = flattenSessionTree(roots); + } else { + // Other modes or with search: flat list + const filtered = filterAndSortSessions(nameFiltered, query, this.sortMode, "all"); + this.filteredSessions = filtered.map((session) => ({ + session, + depth: 0, + isLast: true, + ancestorContinues: [], + })); + } + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredSessions.length - 1)); + } + + private setConfirmingDeletePath(path: string | null): void { + this.confirmingDeletePath = path; + this.onDeleteConfirmationChange?.(path); + } + + private startDeleteConfirmationForSelectedSession(): void { + const selected = this.filteredSessions[this.selectedIndex]; + if (!selected) return; + + // Prevent deleting current session + if (this.isCurrentSessionPath(selected.session.path)) { + this.onError?.("Cannot delete the currently active session"); + return; + } + + this.setConfirmingDeletePath(selected.session.path); + } + + private isCurrentSessionPath(path: string): boolean { + if (!this.currentSessionCanonicalPath) return false; + return (canonicalizePath(path) ?? path) === this.currentSessionCanonicalPath; + } + + invalidate(): void {} + + render(width: number): string[] { + const lines: string[] = []; + + // Render search input + lines.push(...this.searchInput.render(width)); + lines.push(""); // Blank line after search + + if (this.filteredSessions.length === 0) { + let emptyMessage: string; + if (this.nameFilter === "named") { + const toggleKey = keyText("app.session.toggleNamedFilter"); + if (this.showCwd) { + emptyMessage = ` No named sessions found. Press ${toggleKey} to show all.`; + } else { + emptyMessage = ` No named sessions in current folder. Press ${toggleKey} to show all, or Tab to view all.`; + } + } else if (this.showCwd) { + // "All" scope - no sessions anywhere that match filter + emptyMessage = " No sessions found"; + } else { + // "Current folder" scope - hint to try "all" + emptyMessage = " No sessions in current folder. Press Tab to view all."; + } + lines.push(theme.fg("muted", truncateToWidth(emptyMessage, width, "…"))); + return lines; + } + + // Calculate visible range with scrolling + const startIndex = Math.max( + 0, + Math.min(this.selectedIndex - Math.floor(this.maxVisible / 2), this.filteredSessions.length - this.maxVisible), + ); + const endIndex = Math.min(startIndex + this.maxVisible, this.filteredSessions.length); + + // Render visible sessions (one line each with tree structure) + for (let i = startIndex; i < endIndex; i++) { + const node = this.filteredSessions[i]!; + const session = node.session; + const isSelected = i === this.selectedIndex; + const isConfirmingDelete = session.path === this.confirmingDeletePath; + const isCurrent = this.isCurrentSessionPath(session.path); + + // Build tree prefix + const prefix = this.buildTreePrefix(node); + + // Session display text (name or first message) + const hasName = !!session.name; + const displayText = session.name ?? session.firstMessage; + const normalizedMessage = displayText.replace(/[\x00-\x1f\x7f]/g, " ").trim(); + + // Right side: message count and age + const age = formatSessionDate(session.modified); + const msgCount = String(session.messageCount); + let rightPart = `${msgCount} ${age}`; + if (this.showCwd && session.cwd) { + rightPart = `${shortenPath(session.cwd)} ${rightPart}`; + } + if (this.showPath) { + rightPart = `${shortenPath(session.path)} ${rightPart}`; + } + + // Cursor + const cursor = isSelected ? theme.fg("accent", "› ") : " "; + + // Calculate available width for message + const prefixWidth = visibleWidth(prefix); + const rightWidth = visibleWidth(rightPart) + 2; // +2 for spacing + const availableForMsg = width - 2 - prefixWidth - rightWidth; // -2 for cursor + + const truncatedMsg = truncateToWidth(normalizedMessage, Math.max(10, availableForMsg), "…"); + + // Style message + let messageColor: "error" | "warning" | "accent" | null = null; + if (isConfirmingDelete) { + messageColor = "error"; + } else if (isCurrent) { + messageColor = "accent"; + } else if (hasName) { + messageColor = "warning"; + } + let styledMsg = messageColor ? theme.fg(messageColor, truncatedMsg) : truncatedMsg; + if (isSelected) { + styledMsg = theme.bold(styledMsg); + } + + // Build line + const leftPart = cursor + theme.fg("dim", prefix) + styledMsg; + const leftWidth = visibleWidth(leftPart); + const spacing = Math.max(1, width - leftWidth - visibleWidth(rightPart)); + const styledRight = theme.fg(isConfirmingDelete ? "error" : "dim", rightPart); + + let line = leftPart + " ".repeat(spacing) + styledRight; + if (isSelected) { + line = theme.bg("selectedBg", line); + } + lines.push(truncateToWidth(line, width)); + } + + // Add scroll indicator if needed + if (startIndex > 0 || endIndex < this.filteredSessions.length) { + const scrollText = ` (${this.selectedIndex + 1}/${this.filteredSessions.length})`; + const scrollInfo = theme.fg("muted", truncateToWidth(scrollText, width, "")); + lines.push(scrollInfo); + } + + return lines; + } + + private buildTreePrefix(node: FlatSessionNode): string { + if (node.depth === 0) { + return ""; + } + + const parts = node.ancestorContinues.map((continues) => (continues ? "│ " : " ")); + const branch = node.isLast ? "└─ " : "├─ "; + return parts.join("") + branch; + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + + // Handle delete confirmation state first - intercept all keys + if (this.confirmingDeletePath !== null) { + if (kb.matches(keyData, "tui.select.confirm")) { + const pathToDelete = this.confirmingDeletePath; + this.setConfirmingDeletePath(null); + void this.onDeleteSession?.(pathToDelete); + return; + } + if (kb.matches(keyData, "tui.select.cancel")) { + this.setConfirmingDeletePath(null); + return; + } + // Ignore all other keys while confirming + return; + } + + if (kb.matches(keyData, "tui.input.tab")) { + if (this.onToggleScope) { + this.onToggleScope(); + } + return; + } + + if (kb.matches(keyData, "app.session.toggleSort")) { + this.onToggleSort?.(); + return; + } + + if (this.keybindings.matches(keyData, "app.session.toggleNamedFilter")) { + this.onToggleNameFilter?.(); + return; + } + + // Ctrl+P: toggle path display + if (kb.matches(keyData, "app.session.togglePath")) { + this.showPath = !this.showPath; + this.onTogglePath?.(this.showPath); + return; + } + + // Ctrl+D: initiate delete confirmation (useful on terminals that don't distinguish Ctrl+Backspace from Backspace) + if (kb.matches(keyData, "app.session.delete")) { + this.startDeleteConfirmationForSelectedSession(); + return; + } + + // Rename selected session + if (kb.matches(keyData, "app.session.rename")) { + const selected = this.filteredSessions[this.selectedIndex]; + if (selected) { + this.onRenameSession?.(selected.session.path); + } + return; + } + + // Ctrl+Backspace: non-invasive convenience alias for delete + // Only triggers deletion when the query is empty; otherwise it is forwarded to the input + if (kb.matches(keyData, "app.session.deleteNoninvasive")) { + if (this.searchInput.getValue().length > 0) { + this.searchInput.handleInput(keyData); + this.filterSessions(this.searchInput.getValue()); + return; + } + + this.startDeleteConfirmationForSelectedSession(); + return; + } + + // Up arrow + if (kb.matches(keyData, "tui.select.up")) { + this.selectedIndex = Math.max(0, this.selectedIndex - 1); + } + // Down arrow + else if (kb.matches(keyData, "tui.select.down")) { + this.selectedIndex = Math.min(this.filteredSessions.length - 1, this.selectedIndex + 1); + } + // Page up - jump up by maxVisible items + else if (kb.matches(keyData, "tui.select.pageUp")) { + this.selectedIndex = Math.max(0, this.selectedIndex - this.maxVisible); + } + // Page down - jump down by maxVisible items + else if (kb.matches(keyData, "tui.select.pageDown")) { + this.selectedIndex = Math.min(this.filteredSessions.length - 1, this.selectedIndex + this.maxVisible); + } + // Enter + else if (kb.matches(keyData, "tui.select.confirm")) { + const selected = this.filteredSessions[this.selectedIndex]; + if (selected && this.onSelect) { + this.onSelect(selected.session.path); + } + } + // Escape - cancel + else if (kb.matches(keyData, "tui.select.cancel")) { + if (this.onCancel) { + this.onCancel(); + } + } + // Pass everything else to search input + else { + this.searchInput.handleInput(keyData); + this.filterSessions(this.searchInput.getValue()); + } + } +} + +type SessionsLoader = (onProgress?: SessionListProgress) => Promise; + +/** + * Delete a session file, trying the `trash` CLI first, then falling back to unlink + */ +async function deleteSessionFile( + sessionPath: string, +): Promise<{ ok: boolean; method: "trash" | "unlink"; error?: string }> { + // Try `trash` first (if installed) + const trashArgs = sessionPath.startsWith("-") ? ["--", sessionPath] : [sessionPath]; + const trashResult = spawnSync("trash", trashArgs, { encoding: "utf-8" }); + + const getTrashErrorHint = (): string | null => { + const parts: string[] = []; + if (trashResult.error) { + parts.push(trashResult.error.message); + } + const stderr = trashResult.stderr?.trim(); + if (stderr) { + parts.push(stderr.split("\n")[0] ?? stderr); + } + if (parts.length === 0) return null; + return `trash: ${parts.join(" · ").slice(0, 200)}`; + }; + + // If trash reports success, or the file is gone afterwards, treat it as successful + if (trashResult.status === 0 || !existsSync(sessionPath)) { + return { ok: true, method: "trash" }; + } + + // Fallback to permanent deletion + try { + await unlink(sessionPath); + return { ok: true, method: "unlink" }; + } catch (err) { + const unlinkError = err instanceof Error ? err.message : String(err); + const trashErrorHint = getTrashErrorHint(); + const error = trashErrorHint ? `${unlinkError} (${trashErrorHint})` : unlinkError; + return { ok: false, method: "unlink", error }; + } +} + +/** + * Component that renders a session selector + */ +export class SessionSelectorComponent extends Container implements Focusable { + handleInput(data: string): void { + if (this.mode === "rename") { + const kb = getKeybindings(); + if (kb.matches(data, "tui.select.cancel")) { + this.exitRenameMode(); + return; + } + this.renameInput.handleInput(data); + return; + } + + this.sessionList.handleInput(data); + } + + private canRename = true; + private sessionList: SessionList; + private header: SessionSelectorHeader; + private keybindings: KeybindingsManager; + private scope: SessionScope = "current"; + private sortMode: SortMode = "threaded"; + private nameFilter: NameFilter = "all"; + private currentSessions: SessionInfo[] | null = null; + private allSessions: SessionInfo[] | null = null; + private currentSessionsLoader: SessionsLoader; + private allSessionsLoader: SessionsLoader; + private requestRender: () => void; + private renameSession?: (sessionPath: string, currentName: string | undefined) => Promise; + private currentLoading = false; + private allLoading = false; + private allLoadSeq = 0; + + private mode: "list" | "rename" = "list"; + private renameInput = new Input(); + private renameTargetPath: string | null = null; + + // Focusable implementation - propagate to sessionList for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.sessionList.focused = value; + this.renameInput.focused = value; + if (value && this.mode === "rename") { + this.renameInput.focused = true; + } + } + + private buildBaseLayout(content: Component, options?: { showHeader?: boolean }): void { + this.clear(); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s) => theme.fg("accent", s))); + this.addChild(new Spacer(1)); + if (options?.showHeader ?? true) { + this.addChild(this.header); + this.addChild(new Spacer(1)); + } + this.addChild(content); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s) => theme.fg("accent", s))); + } + + constructor( + currentSessionsLoader: SessionsLoader, + allSessionsLoader: SessionsLoader, + onSelect: (sessionPath: string) => void, + onCancel: () => void, + onExit: () => void, + requestRender: () => void, + options?: { + renameSession?: (sessionPath: string, currentName: string | undefined) => Promise; + showRenameHint?: boolean; + keybindings?: KeybindingsManager; + }, + currentSessionFilePath?: string, + ) { + super(); + this.keybindings = options?.keybindings ?? KeybindingsManager.create(); + this.currentSessionsLoader = currentSessionsLoader; + this.allSessionsLoader = allSessionsLoader; + this.requestRender = requestRender; + this.header = new SessionSelectorHeader(this.scope, this.sortMode, this.nameFilter, this.requestRender); + const renameSession = options?.renameSession; + this.renameSession = renameSession; + this.canRename = !!renameSession; + this.header.setShowRenameHint(options?.showRenameHint ?? this.canRename); + + // Create session list (starts empty, will be populated after load) + this.sessionList = new SessionList( + [], + false, + this.sortMode, + this.nameFilter, + this.keybindings, + currentSessionFilePath, + ); + + this.buildBaseLayout(this.sessionList); + + this.renameInput.onSubmit = (value) => { + void this.confirmRename(value); + }; + + // Ensure header status timeouts are cleared when leaving the selector + const clearStatusMessage = () => this.header.setStatusMessage(null); + this.sessionList.onSelect = (sessionPath) => { + clearStatusMessage(); + onSelect(sessionPath); + }; + this.sessionList.onCancel = () => { + clearStatusMessage(); + onCancel(); + }; + this.sessionList.onExit = () => { + clearStatusMessage(); + onExit(); + }; + this.sessionList.onToggleScope = () => this.toggleScope(); + this.sessionList.onToggleSort = () => this.toggleSortMode(); + this.sessionList.onToggleNameFilter = () => this.toggleNameFilter(); + this.sessionList.onRenameSession = (sessionPath) => { + if (!renameSession) return; + if (this.scope === "current" && this.currentLoading) return; + if (this.scope === "all" && this.allLoading) return; + + const sessions = this.scope === "all" ? (this.allSessions ?? []) : (this.currentSessions ?? []); + const session = sessions.find((s) => s.path === sessionPath); + this.enterRenameMode(sessionPath, session?.name); + }; + + // Sync list events to header + this.sessionList.onTogglePath = (showPath) => { + this.header.setShowPath(showPath); + this.requestRender(); + }; + this.sessionList.onDeleteConfirmationChange = (path) => { + this.header.setConfirmingDeletePath(path); + this.requestRender(); + }; + this.sessionList.onError = (msg) => { + this.header.setStatusMessage({ type: "error", message: msg }, 3000); + this.requestRender(); + }; + + // Handle session deletion + this.sessionList.onDeleteSession = async (sessionPath: string) => { + const result = await deleteSessionFile(sessionPath); + + if (result.ok) { + if (this.currentSessions) { + this.currentSessions = this.currentSessions.filter((s) => s.path !== sessionPath); + } + if (this.allSessions) { + this.allSessions = this.allSessions.filter((s) => s.path !== sessionPath); + } + + const sessions = this.scope === "all" ? (this.allSessions ?? []) : (this.currentSessions ?? []); + const showCwd = this.scope === "all"; + this.sessionList.setSessions(sessions, showCwd); + + const msg = result.method === "trash" ? "Session moved to trash" : "Session deleted"; + this.header.setStatusMessage({ type: "info", message: msg }, 2000); + await this.refreshSessionsAfterMutation(); + } else { + const errorMessage = result.error ?? "Unknown error"; + this.header.setStatusMessage({ type: "error", message: `Failed to delete: ${errorMessage}` }, 3000); + } + + this.requestRender(); + }; + + // Start loading current sessions immediately + this.loadCurrentSessions(); + } + + private loadCurrentSessions(): void { + void this.loadScope("current", "initial"); + } + + private enterRenameMode(sessionPath: string, currentName: string | undefined): void { + this.mode = "rename"; + this.renameTargetPath = sessionPath; + this.renameInput.setValue(currentName ?? ""); + this.renameInput.focused = true; + + const panel = new Container(); + panel.addChild(new Text(theme.bold("Rename Session"), 1, 0)); + panel.addChild(new Spacer(1)); + panel.addChild(this.renameInput); + panel.addChild(new Spacer(1)); + panel.addChild( + new Text( + theme.fg("muted", `${keyText("tui.select.confirm")} to save · ${keyText("tui.select.cancel")} to cancel`), + 1, + 0, + ), + ); + + this.buildBaseLayout(panel, { showHeader: false }); + this.requestRender(); + } + + private exitRenameMode(): void { + this.mode = "list"; + this.renameTargetPath = null; + + this.buildBaseLayout(this.sessionList); + + this.requestRender(); + } + + private async confirmRename(value: string): Promise { + const next = value.trim(); + if (!next) return; + const target = this.renameTargetPath; + if (!target) { + this.exitRenameMode(); + return; + } + + // Find current name for callback + const renameSession = this.renameSession; + if (!renameSession) { + this.exitRenameMode(); + return; + } + + try { + await renameSession(target, next); + await this.refreshSessionsAfterMutation(); + } finally { + this.exitRenameMode(); + } + } + + private async loadScope(scope: SessionScope, reason: "initial" | "refresh" | "toggle"): Promise { + const showCwd = scope === "all"; + + // Mark loading + if (scope === "current") { + this.currentLoading = true; + } else { + this.allLoading = true; + } + + const seq = scope === "all" ? ++this.allLoadSeq : undefined; + this.header.setScope(scope); + this.header.setLoading(true); + this.requestRender(); + + const onProgress = (loaded: number, total: number) => { + if (scope !== this.scope) return; + if (seq !== undefined && seq !== this.allLoadSeq) return; + this.header.setProgress(loaded, total); + this.requestRender(); + }; + + try { + const sessions = await (scope === "current" + ? this.currentSessionsLoader(onProgress) + : this.allSessionsLoader(onProgress)); + + if (scope === "current") { + this.currentSessions = sessions; + this.currentLoading = false; + } else { + this.allSessions = sessions; + this.allLoading = false; + } + + if (scope !== this.scope) return; + if (seq !== undefined && seq !== this.allLoadSeq) return; + + this.header.setLoading(false); + this.sessionList.setSessions(sessions, showCwd); + this.requestRender(); + } catch (err) { + if (scope === "current") { + this.currentLoading = false; + } else { + this.allLoading = false; + } + + if (scope !== this.scope) return; + if (seq !== undefined && seq !== this.allLoadSeq) return; + + const message = err instanceof Error ? err.message : String(err); + this.header.setLoading(false); + this.header.setStatusMessage({ type: "error", message: `Failed to load sessions: ${message}` }, 4000); + + if (reason === "initial") { + this.sessionList.setSessions([], showCwd); + } + this.requestRender(); + } + } + + private toggleSortMode(): void { + // Cycle: threaded -> recent -> relevance -> threaded + this.sortMode = this.sortMode === "threaded" ? "recent" : this.sortMode === "recent" ? "relevance" : "threaded"; + this.header.setSortMode(this.sortMode); + this.sessionList.setSortMode(this.sortMode); + this.requestRender(); + } + + private toggleNameFilter(): void { + this.nameFilter = this.nameFilter === "all" ? "named" : "all"; + this.header.setNameFilter(this.nameFilter); + this.sessionList.setNameFilter(this.nameFilter); + this.requestRender(); + } + + private async refreshSessionsAfterMutation(): Promise { + await this.loadScope(this.scope, "refresh"); + } + + private toggleScope(): void { + if (this.scope === "current") { + this.scope = "all"; + this.header.setScope(this.scope); + + if (this.allSessions !== null) { + this.header.setLoading(false); + this.sessionList.setSessions(this.allSessions, true); + this.requestRender(); + return; + } + + if (!this.allLoading) { + void this.loadScope("all", "toggle"); + } + return; + } + + this.scope = "current"; + this.header.setScope(this.scope); + this.header.setLoading(this.currentLoading); + this.sessionList.setSessions(this.currentSessions ?? [], false); + this.requestRender(); + } + + getSessionList(): SessionList { + return this.sessionList; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/settings-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/settings-selector.ts new file mode 100644 index 000000000..2a82f41b7 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/settings-selector.ts @@ -0,0 +1,786 @@ +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { Transport } from "@earendil-works/pi-ai"; +import { + type Component, + Container, + getCapabilities, + type SelectItem, + SelectList, + type SelectListLayoutOptions, + type SettingItem, + SettingsList, + Spacer, + Text, +} from "@earendil-works/pi-tui"; +import { formatHttpIdleTimeoutMs, HTTP_IDLE_TIMEOUT_CHOICES } from "../../../core/http-dispatcher.ts"; +import type { DefaultProjectTrust, WarningSettings } from "../../../core/settings-manager.ts"; +import { + getSelectListTheme, + getSettingsListTheme, + parseAutoThemeSetting, + type TerminalTheme, + theme, +} from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyDisplayText } from "./keybinding-hints.ts"; + +const SETTINGS_SUBMENU_SELECT_LIST_LAYOUT: SelectListLayoutOptions = { + minPrimaryColumnWidth: 12, + maxPrimaryColumnWidth: 32, +}; + +const THINKING_DESCRIPTIONS: Record = { + off: "No reasoning", + minimal: "Very brief reasoning (~1k tokens)", + low: "Light reasoning (~2k tokens)", + medium: "Moderate reasoning (~8k tokens)", + high: "Deep reasoning (~16k tokens)", + xhigh: "Maximum reasoning (~32k tokens)", +}; + +const DEFAULT_PROJECT_TRUST_LABELS: Record = { + ask: "Ask", + always: "Always trust", + never: "Never trust", +}; + +const DEFAULT_PROJECT_TRUST_BY_LABEL = new Map( + Object.entries(DEFAULT_PROJECT_TRUST_LABELS).map(([value, label]) => [label, value as DefaultProjectTrust]), +); + +export interface SettingsConfig { + autoCompact: boolean; + showImages: boolean; + imageWidthCells: number; + autoResizeImages: boolean; + blockImages: boolean; + enableSkillCommands: boolean; + steeringMode: "all" | "one-at-a-time"; + followUpMode: "all" | "one-at-a-time"; + transport: Transport; + httpIdleTimeoutMs: number; + thinkingLevel: ThinkingLevel; + availableThinkingLevels: ThinkingLevel[]; + currentTheme: string; + terminalTheme: TerminalTheme; + availableThemes: string[]; + hideThinkingBlock: boolean; + doubleEscapeAction: "fork" | "tree" | "none"; + treeFilterMode: "default" | "no-tools" | "user-only" | "labeled-only" | "all"; + showHardwareCursor: boolean; + editorPaddingX: number; + autocompleteMaxVisible: number; + quietStartup: boolean; + defaultProjectTrust: DefaultProjectTrust; + clearOnShrink: boolean; + showTerminalProgress: boolean; + warnings: WarningSettings; +} + +export interface SettingsCallbacks { + onAutoCompactChange: (enabled: boolean) => void; + onShowImagesChange: (enabled: boolean) => void; + onImageWidthCellsChange: (width: number) => void; + onAutoResizeImagesChange: (enabled: boolean) => void; + onBlockImagesChange: (blocked: boolean) => void; + onEnableSkillCommandsChange: (enabled: boolean) => void; + onSteeringModeChange: (mode: "all" | "one-at-a-time") => void; + onFollowUpModeChange: (mode: "all" | "one-at-a-time") => void; + onTransportChange: (transport: Transport) => void; + onHttpIdleTimeoutMsChange: (timeoutMs: number) => void; + onThinkingLevelChange: (level: ThinkingLevel) => void; + onThemeChange: (theme: string) => void; + onThemePreview?: (theme: string) => void; + onHideThinkingBlockChange: (hidden: boolean) => void; + onDoubleEscapeActionChange: (action: "fork" | "tree" | "none") => void; + onTreeFilterModeChange: (mode: "default" | "no-tools" | "user-only" | "labeled-only" | "all") => void; + onShowHardwareCursorChange: (enabled: boolean) => void; + onEditorPaddingXChange: (padding: number) => void; + onAutocompleteMaxVisibleChange: (maxVisible: number) => void; + onQuietStartupChange: (enabled: boolean) => void; + onDefaultProjectTrustChange: (defaultProjectTrust: DefaultProjectTrust) => void; + onClearOnShrinkChange: (enabled: boolean) => void; + onShowTerminalProgressChange: (enabled: boolean) => void; + onWarningsChange: (warnings: WarningSettings) => void; + onCancel: () => void; +} + +/** + * A submenu component for selecting from a list of options. + */ +class WarningSettingsSubmenu extends Container { + private settingsList: SettingsList; + private state: WarningSettings; + + constructor(warnings: WarningSettings, onChange: (warnings: WarningSettings) => void, onCancel: () => void) { + super(); + + this.state = { ...warnings }; + + const items: SettingItem[] = [ + { + id: "anthropic-extra-usage", + label: "Anthropic extra usage", + description: "Warn when Anthropic subscription auth may use paid extra usage", + currentValue: (this.state.anthropicExtraUsage ?? true) ? "true" : "false", + values: ["true", "false"], + }, + ]; + + this.settingsList = new SettingsList( + items, + Math.min(items.length, 10), + getSettingsListTheme(), + (id, newValue) => { + switch (id) { + case "anthropic-extra-usage": + this.state = { ...this.state, anthropicExtraUsage: newValue === "true" }; + onChange({ ...this.state }); + break; + } + }, + onCancel, + ); + + this.addChild(this.settingsList); + } + + handleInput(data: string): void { + this.settingsList.handleInput(data); + } +} + +class SelectSubmenu extends Container { + private selectList: SelectList; + + constructor( + title: string, + description: string, + options: SelectItem[], + currentValue: string, + onSelect: (value: string) => void, + onCancel: () => void, + onSelectionChange?: (value: string) => void, + ) { + super(); + + // Title + this.addChild(new Text(theme.bold(theme.fg("accent", title)), 0, 0)); + + // Description + if (description) { + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("muted", description), 0, 0)); + } + + // Spacer + this.addChild(new Spacer(1)); + + // Select list + this.selectList = new SelectList( + options, + Math.min(options.length, 10), + getSelectListTheme(), + SETTINGS_SUBMENU_SELECT_LIST_LAYOUT, + ); + + // Pre-select current value + const currentIndex = options.findIndex((o) => o.value === currentValue); + if (currentIndex !== -1) { + this.selectList.setSelectedIndex(currentIndex); + } + + this.selectList.onSelect = (item) => { + onSelect(item.value); + }; + + this.selectList.onCancel = onCancel; + + if (onSelectionChange) { + this.selectList.onSelectionChange = (item) => { + onSelectionChange(item.value); + }; + } + + this.addChild(this.selectList); + + // Hint + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("dim", " Enter to select · Esc to go back"), 0, 0)); + } + + handleInput(data: string): void { + this.selectList.handleInput(data); + } +} + +function themeItems(availableThemes: string[]): SelectItem[] { + return availableThemes.map((name) => ({ value: name, label: name })); +} + +const AUTOMATIC_THEME_VALUE = "/"; + +function singleModeThemeItems(availableThemes: string[]): SelectItem[] { + return [ + { + value: AUTOMATIC_THEME_VALUE, + label: "Automatic", + description: "Use separate themes for light and dark terminal appearance", + }, + ...themeItems(availableThemes), + ]; +} + +function preferredTheme(availableThemes: string[], preferred: string | undefined, fallback: string): string { + if (preferred && availableThemes.includes(preferred)) return preferred; + if (availableThemes.includes(fallback)) return fallback; + return availableThemes[0] ?? fallback; +} + +function defaultAutomaticThemes( + currentThemeSetting: string, + availableThemes: string[], +): { lightTheme: string; darkTheme: string } { + const autoTheme = parseAutoThemeSetting(currentThemeSetting); + if (autoTheme) return autoTheme; + + const currentFixedTheme = currentThemeSetting.includes("/") ? undefined : currentThemeSetting; + const themeName = preferredTheme(availableThemes, currentFixedTheme, "dark"); + return { lightTheme: themeName, darkTheme: themeName }; +} + +class ThemeSubmenu extends Container { + private inputComponent: Component | undefined; + private readonly callbacks: SettingsCallbacks; + private readonly availableThemes: string[]; + private readonly terminalTheme: TerminalTheme; + private readonly onDone: (selectedValue?: string) => void; + private readonly originalThemeSetting: string; + private mode: "single" | "automatic"; + private singleTheme: string; + private lightTheme: string; + private darkTheme: string; + + constructor( + currentThemeSetting: string, + terminalTheme: TerminalTheme, + availableThemes: string[], + callbacks: SettingsCallbacks, + onDone: (selectedValue?: string) => void, + ) { + super(); + this.callbacks = callbacks; + this.availableThemes = availableThemes; + this.terminalTheme = terminalTheme; + this.onDone = onDone; + this.originalThemeSetting = currentThemeSetting; + const autoTheme = parseAutoThemeSetting(currentThemeSetting); + const automaticThemes = defaultAutomaticThemes(currentThemeSetting, availableThemes); + const fixedTheme = autoTheme || currentThemeSetting.includes("/") ? undefined : currentThemeSetting; + this.mode = autoTheme ? "automatic" : "single"; + this.lightTheme = automaticThemes.lightTheme; + this.darkTheme = automaticThemes.darkTheme; + this.singleTheme = preferredTheme( + availableThemes, + fixedTheme ?? (autoTheme ? this.getActiveAutomaticTheme() : undefined), + "dark", + ); + + if (this.mode === "automatic") { + this.showAutomaticMenu(); + } else { + this.showSingleMenu(); + } + } + + handleInput(data: string): void { + this.inputComponent?.handleInput?.(data); + } + + private setContent(renderComponent: Component, inputComponent: Component = renderComponent): void { + this.clear(); + this.addChild(renderComponent); + this.inputComponent = inputComponent; + } + + private showSingleMenu(): void { + this.mode = "single"; + const menu = new SelectSubmenu( + "Theme", + "Select a theme, or choose Automatic to follow terminal appearance.", + singleModeThemeItems(this.availableThemes), + this.singleTheme, + (value) => { + if (value === AUTOMATIC_THEME_VALUE) { + this.mode = "automatic"; + this.callbacks.onThemePreview?.(this.getThemeSetting()); + this.showAutomaticMenu(); + return; + } + + this.singleTheme = value; + this.apply(value); + }, + () => this.cancel(), + (value) => { + this.callbacks.onThemePreview?.(value === AUTOMATIC_THEME_VALUE ? this.getAutomaticThemeSetting() : value); + }, + ); + this.setContent(menu); + } + + private showAutomaticMenu(): void { + this.mode = "automatic"; + const content = new Container(); + content.addChild(new Text(theme.bold(theme.fg("accent", "Automatic Theme")), 0, 0)); + content.addChild(new Spacer(1)); + content.addChild(new Text(theme.fg("muted", "Choose themes for terminal light and dark appearance."), 0, 0)); + content.addChild(new Text(theme.fg("muted", "Light/dark detection requires terminal support."), 0, 0)); + content.addChild(new Spacer(1)); + + const items: SettingItem[] = [ + { + id: "light-theme", + label: "Light theme", + description: "Theme to use in automatic mode when the terminal is light", + currentValue: this.lightTheme, + submenu: (currentValue, done) => + this.createThemeSelect( + "Light Theme", + "Select the theme to use for light terminal appearance", + currentValue, + done, + (value) => { + this.lightTheme = value; + this.callbacks.onThemePreview?.(this.getThemeSetting()); + done(value); + }, + ), + }, + { + id: "dark-theme", + label: "Dark theme", + description: "Theme to use in automatic mode when the terminal is dark", + currentValue: this.darkTheme, + submenu: (currentValue, done) => + this.createThemeSelect( + "Dark Theme", + "Select the theme to use for dark terminal appearance", + currentValue, + done, + (value) => { + this.darkTheme = value; + this.callbacks.onThemePreview?.(this.getThemeSetting()); + done(value); + }, + ), + }, + { + id: "apply", + label: "Apply", + description: "Save and go back", + currentValue: "save and go back", + values: ["save and go back"], + }, + { + id: "single-mode", + label: "Change mode", + description: "Switch to one theme for light and dark", + currentValue: "switch to single theme", + values: ["switch to single theme"], + }, + ]; + + const settingsList = new SettingsList( + items, + Math.min(items.length, 10), + getSettingsListTheme(), + (id) => { + switch (id) { + case "single-mode": + this.mode = "single"; + this.singleTheme = this.getActiveAutomaticTheme(); + this.callbacks.onThemePreview?.(this.singleTheme); + this.showSingleMenu(); + break; + case "apply": + this.apply(this.getAutomaticThemeSetting()); + break; + } + }, + () => this.cancel(), + ); + content.addChild(settingsList); + this.setContent(content, settingsList); + } + + private createThemeSelect( + title: string, + description: string, + currentValue: string, + done: (selectedValue?: string) => void, + onSelect: (value: string) => void, + ): SelectSubmenu { + return new SelectSubmenu( + title, + description, + themeItems(this.availableThemes), + currentValue, + onSelect, + () => { + this.callbacks.onThemePreview?.(this.getThemeSetting()); + done(); + }, + (value) => this.callbacks.onThemePreview?.(value), + ); + } + + private getThemeSetting(): string { + return this.mode === "automatic" ? this.getAutomaticThemeSetting() : this.singleTheme; + } + + private getActiveAutomaticTheme(): string { + return this.terminalTheme === "light" ? this.lightTheme : this.darkTheme; + } + + private getAutomaticThemeSetting(): string { + return `${this.lightTheme}/${this.darkTheme}`; + } + + private apply(themeSetting: string): void { + this.onDone(themeSetting); + } + + private cancel(): void { + this.callbacks.onThemePreview?.(this.originalThemeSetting); + this.onDone(); + } +} + +/** + * Main settings selector component. + */ +export class SettingsSelectorComponent extends Container { + private settingsList: SettingsList; + + constructor(config: SettingsConfig, callbacks: SettingsCallbacks) { + super(); + + const supportsImages = getCapabilities().images; + const followUpKey = keyDisplayText("app.message.followUp"); + let currentWarnings = { ...config.warnings }; + + const items: SettingItem[] = [ + { + id: "autocompact", + label: "Auto-compact", + description: "Automatically compact context when it gets too large", + currentValue: config.autoCompact ? "true" : "false", + values: ["true", "false"], + }, + { + id: "steering-mode", + label: "Steering mode", + description: + "Enter while streaming queues steering messages. 'one-at-a-time': deliver one, wait for response. 'all': deliver all at once.", + currentValue: config.steeringMode, + values: ["one-at-a-time", "all"], + }, + { + id: "follow-up-mode", + label: "Follow-up mode", + description: `${followUpKey} queues follow-up messages until agent stops. 'one-at-a-time': deliver one, wait for response. 'all': deliver all at once.`, + currentValue: config.followUpMode, + values: ["one-at-a-time", "all"], + }, + { + id: "transport", + label: "Transport", + description: "Preferred transport for providers that support multiple transports", + currentValue: config.transport, + values: ["sse", "websocket", "websocket-cached", "auto"], + }, + { + id: "http-idle-timeout", + label: "HTTP idle timeout", + description: + "Maximum idle gap while waiting for HTTP headers or body chunks. Disable for local models that pause longer than five minutes.", + currentValue: formatHttpIdleTimeoutMs(config.httpIdleTimeoutMs), + values: HTTP_IDLE_TIMEOUT_CHOICES.map((choice) => choice.label), + }, + { + id: "hide-thinking", + label: "Hide thinking", + description: "Hide thinking blocks in assistant responses", + currentValue: config.hideThinkingBlock ? "true" : "false", + values: ["true", "false"], + }, + { + id: "quiet-startup", + label: "Quiet startup", + description: "Disable verbose printing at startup", + currentValue: config.quietStartup ? "true" : "false", + values: ["true", "false"], + }, + { + id: "default-project-trust", + label: "Default project trust", + description: "Fallback behavior when no extension or saved trust decision decides project trust", + currentValue: DEFAULT_PROJECT_TRUST_LABELS[config.defaultProjectTrust], + values: Object.values(DEFAULT_PROJECT_TRUST_LABELS), + }, + { + id: "double-escape-action", + label: "Double-escape action", + description: "Action when pressing Escape twice with empty editor", + currentValue: config.doubleEscapeAction, + values: ["tree", "fork", "none"], + }, + { + id: "tree-filter-mode", + label: "Tree filter mode", + description: "Default filter when opening /tree", + currentValue: config.treeFilterMode, + values: ["default", "no-tools", "user-only", "labeled-only", "all"], + }, + { + id: "warnings", + label: "Warnings", + description: "Enable or disable individual warnings", + currentValue: "configure", + submenu: (_currentValue, done) => + new WarningSettingsSubmenu( + currentWarnings, + (warnings) => { + currentWarnings = warnings; + callbacks.onWarningsChange(warnings); + }, + () => done(), + ), + }, + { + id: "thinking", + label: "Thinking level", + description: "Reasoning depth for thinking-capable models", + currentValue: config.thinkingLevel, + submenu: (currentValue, done) => + new SelectSubmenu( + "Thinking Level", + "Select reasoning depth for thinking-capable models", + config.availableThinkingLevels.map((level) => ({ + value: level, + label: level, + description: THINKING_DESCRIPTIONS[level], + })), + currentValue, + (value) => { + callbacks.onThinkingLevelChange(value as ThinkingLevel); + done(value); + }, + () => done(), + ), + }, + { + id: "theme", + label: "Theme", + description: "Color theme for the interface", + currentValue: config.currentTheme, + submenu: (currentValue, done) => + new ThemeSubmenu(currentValue, config.terminalTheme, config.availableThemes, callbacks, done), + }, + ]; + + // Only show image toggle if terminal supports it + if (supportsImages) { + // Insert after autocompact + items.splice(1, 0, { + id: "show-images", + label: "Show images", + description: "Render images inline in terminal", + currentValue: config.showImages ? "true" : "false", + values: ["true", "false"], + }); + items.splice(2, 0, { + id: "image-width-cells", + label: "Image width", + description: "Preferred inline image width in terminal cells", + currentValue: String(config.imageWidthCells), + values: ["60", "80", "120"], + }); + } + + // Image auto-resize toggle (always available, affects both attached and read images) + items.splice(supportsImages ? 3 : 1, 0, { + id: "auto-resize-images", + label: "Auto-resize images", + description: "Resize large images to 2000x2000 max for better model compatibility", + currentValue: config.autoResizeImages ? "true" : "false", + values: ["true", "false"], + }); + + // Block images toggle (always available, insert after auto-resize-images) + const autoResizeIndex = items.findIndex((item) => item.id === "auto-resize-images"); + items.splice(autoResizeIndex + 1, 0, { + id: "block-images", + label: "Block images", + description: "Prevent images from being sent to LLM providers", + currentValue: config.blockImages ? "true" : "false", + values: ["true", "false"], + }); + + // Skill commands toggle (insert after block-images) + const blockImagesIndex = items.findIndex((item) => item.id === "block-images"); + items.splice(blockImagesIndex + 1, 0, { + id: "skill-commands", + label: "Skill commands", + description: "Register skills as /skill:name commands", + currentValue: config.enableSkillCommands ? "true" : "false", + values: ["true", "false"], + }); + + // Hardware cursor toggle (insert after skill-commands) + const skillCommandsIndex = items.findIndex((item) => item.id === "skill-commands"); + items.splice(skillCommandsIndex + 1, 0, { + id: "show-hardware-cursor", + label: "Show hardware cursor", + description: "Show the terminal cursor while still positioning it for IME support", + currentValue: config.showHardwareCursor ? "true" : "false", + values: ["true", "false"], + }); + + // Editor padding toggle (insert after show-hardware-cursor) + const hardwareCursorIndex = items.findIndex((item) => item.id === "show-hardware-cursor"); + items.splice(hardwareCursorIndex + 1, 0, { + id: "editor-padding", + label: "Editor padding", + description: "Horizontal padding for input editor (0-3)", + currentValue: String(config.editorPaddingX), + values: ["0", "1", "2", "3"], + }); + + // Autocomplete max visible toggle (insert after editor-padding) + const editorPaddingIndex = items.findIndex((item) => item.id === "editor-padding"); + items.splice(editorPaddingIndex + 1, 0, { + id: "autocomplete-max-visible", + label: "Autocomplete max items", + description: "Max visible items in autocomplete dropdown (3-20)", + currentValue: String(config.autocompleteMaxVisible), + values: ["3", "5", "7", "10", "15", "20"], + }); + + // Clear on shrink toggle (insert after autocomplete-max-visible) + const autocompleteIndex = items.findIndex((item) => item.id === "autocomplete-max-visible"); + items.splice(autocompleteIndex + 1, 0, { + id: "clear-on-shrink", + label: "Clear on shrink", + description: "Clear empty rows when content shrinks (may cause flicker)", + currentValue: config.clearOnShrink ? "true" : "false", + values: ["true", "false"], + }); + + // Terminal progress toggle (insert after clear-on-shrink) + const clearOnShrinkIndex = items.findIndex((item) => item.id === "clear-on-shrink"); + items.splice(clearOnShrinkIndex + 1, 0, { + id: "terminal-progress", + label: "Terminal progress", + description: "Show OSC 9;4 progress indicators in the terminal tab bar", + currentValue: config.showTerminalProgress ? "true" : "false", + values: ["true", "false"], + }); + + // Add borders + this.addChild(new DynamicBorder()); + + this.settingsList = new SettingsList( + items, + 10, + getSettingsListTheme(), + (id, newValue) => { + switch (id) { + case "autocompact": + callbacks.onAutoCompactChange(newValue === "true"); + break; + case "show-images": + callbacks.onShowImagesChange(newValue === "true"); + break; + case "image-width-cells": + callbacks.onImageWidthCellsChange(parseInt(newValue, 10)); + break; + case "auto-resize-images": + callbacks.onAutoResizeImagesChange(newValue === "true"); + break; + case "block-images": + callbacks.onBlockImagesChange(newValue === "true"); + break; + case "skill-commands": + callbacks.onEnableSkillCommandsChange(newValue === "true"); + break; + case "steering-mode": + callbacks.onSteeringModeChange(newValue as "all" | "one-at-a-time"); + break; + case "follow-up-mode": + callbacks.onFollowUpModeChange(newValue as "all" | "one-at-a-time"); + break; + case "transport": + callbacks.onTransportChange(newValue as Transport); + break; + case "http-idle-timeout": { + const choice = HTTP_IDLE_TIMEOUT_CHOICES.find((item) => item.label === newValue); + if (choice) { + callbacks.onHttpIdleTimeoutMsChange(choice.timeoutMs); + } + break; + } + case "hide-thinking": + callbacks.onHideThinkingBlockChange(newValue === "true"); + break; + case "quiet-startup": + callbacks.onQuietStartupChange(newValue === "true"); + break; + case "default-project-trust": { + const defaultProjectTrust = DEFAULT_PROJECT_TRUST_BY_LABEL.get(newValue); + if (defaultProjectTrust) { + callbacks.onDefaultProjectTrustChange(defaultProjectTrust); + } + break; + } + case "double-escape-action": + callbacks.onDoubleEscapeActionChange(newValue as "fork" | "tree"); + break; + case "tree-filter-mode": + callbacks.onTreeFilterModeChange( + newValue as "default" | "no-tools" | "user-only" | "labeled-only" | "all", + ); + break; + case "show-hardware-cursor": + callbacks.onShowHardwareCursorChange(newValue === "true"); + break; + case "editor-padding": + callbacks.onEditorPaddingXChange(parseInt(newValue, 10)); + break; + case "autocomplete-max-visible": + callbacks.onAutocompleteMaxVisibleChange(parseInt(newValue, 10)); + break; + case "clear-on-shrink": + callbacks.onClearOnShrinkChange(newValue === "true"); + break; + case "terminal-progress": + callbacks.onShowTerminalProgressChange(newValue === "true"); + break; + case "theme": + callbacks.onThemeChange(newValue); + break; + } + }, + callbacks.onCancel, + { enableSearch: true }, + ); + + this.addChild(this.settingsList); + this.addChild(new DynamicBorder()); + } + + getSettingsList(): SettingsList { + return this.settingsList; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/skill-invocation-message.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/skill-invocation-message.ts new file mode 100644 index 000000000..e09febc7c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/skill-invocation-message.ts @@ -0,0 +1,55 @@ +import { Box, Markdown, type MarkdownTheme, Text } from "@earendil-works/pi-tui"; +import type { ParsedSkillBlock } from "../../../core/agent-session.ts"; +import { getMarkdownTheme, theme } from "../theme/theme.ts"; +import { keyText } from "./keybinding-hints.ts"; + +/** + * Component that renders a skill invocation message with collapsed/expanded state. + * Uses same background color as custom messages for visual consistency. + * Only renders the skill block itself - user message is rendered separately. + */ +export class SkillInvocationMessageComponent extends Box { + private expanded = false; + private skillBlock: ParsedSkillBlock; + private markdownTheme: MarkdownTheme; + + constructor(skillBlock: ParsedSkillBlock, markdownTheme: MarkdownTheme = getMarkdownTheme()) { + super(1, 1, (t) => theme.bg("customMessageBg", t)); + this.skillBlock = skillBlock; + this.markdownTheme = markdownTheme; + this.updateDisplay(); + } + + setExpanded(expanded: boolean): void { + this.expanded = expanded; + this.updateDisplay(); + } + + override invalidate(): void { + super.invalidate(); + this.updateDisplay(); + } + + private updateDisplay(): void { + this.clear(); + + if (this.expanded) { + // Expanded: label + skill name header + full content + const label = theme.fg("customMessageLabel", `\x1b[1m[skill]\x1b[22m`); + this.addChild(new Text(label, 0, 0)); + const header = `**${this.skillBlock.name}**\n\n`; + this.addChild( + new Markdown(header + this.skillBlock.content, 0, 0, this.markdownTheme, { + color: (text: string) => theme.fg("customMessageText", text), + }), + ); + } else { + // Collapsed: single line - [skill] name (hint to expand) + const line = + theme.fg("customMessageLabel", `\x1b[1m[skill]\x1b[22m `) + + theme.fg("customMessageText", this.skillBlock.name) + + theme.fg("dim", ` (${keyText("app.tools.expand")} to expand)`); + this.addChild(new Text(line, 0, 0)); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/theme-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/theme-selector.ts new file mode 100644 index 000000000..07b377423 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/theme-selector.ts @@ -0,0 +1,67 @@ +import { Container, type SelectItem, SelectList, type SelectListLayoutOptions } from "@earendil-works/pi-tui"; +import { getAvailableThemes, getSelectListTheme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; + +const THEME_SELECT_LIST_LAYOUT: SelectListLayoutOptions = { + minPrimaryColumnWidth: 12, + maxPrimaryColumnWidth: 32, +}; + +/** + * Component that renders a theme selector + */ +export class ThemeSelectorComponent extends Container { + private selectList: SelectList; + private onPreview: (themeName: string) => void; + + constructor( + currentTheme: string, + onSelect: (themeName: string) => void, + onCancel: () => void, + onPreview: (themeName: string) => void, + ) { + super(); + this.onPreview = onPreview; + + // Get available themes and create select items + const themes = getAvailableThemes(); + const themeItems: SelectItem[] = themes.map((name) => ({ + value: name, + label: name, + description: name === currentTheme ? "(current)" : undefined, + })); + + // Add top border + this.addChild(new DynamicBorder()); + + // Create selector + this.selectList = new SelectList(themeItems, 10, getSelectListTheme(), THEME_SELECT_LIST_LAYOUT); + + // Preselect current theme + const currentIndex = themes.indexOf(currentTheme); + if (currentIndex !== -1) { + this.selectList.setSelectedIndex(currentIndex); + } + + this.selectList.onSelect = (item) => { + onSelect(item.value); + }; + + this.selectList.onCancel = () => { + onCancel(); + }; + + this.selectList.onSelectionChange = (item) => { + this.onPreview(item.value); + }; + + this.addChild(this.selectList); + + // Add bottom border + this.addChild(new DynamicBorder()); + } + + getSelectList(): SelectList { + return this.selectList; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/thinking-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/thinking-selector.ts new file mode 100644 index 000000000..91a3a6266 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/thinking-selector.ts @@ -0,0 +1,74 @@ +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { Container, type SelectItem, SelectList, type SelectListLayoutOptions } from "@earendil-works/pi-tui"; +import { getSelectListTheme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; + +const THINKING_SELECT_LIST_LAYOUT: SelectListLayoutOptions = { + minPrimaryColumnWidth: 12, + maxPrimaryColumnWidth: 32, +}; + +const LEVEL_DESCRIPTIONS: Record = { + off: "No reasoning", + minimal: "Very brief reasoning (~1k tokens)", + low: "Light reasoning (~2k tokens)", + medium: "Moderate reasoning (~8k tokens)", + high: "Deep reasoning (~16k tokens)", + xhigh: "Maximum reasoning (~32k tokens)", +}; + +/** + * Component that renders a thinking level selector with borders + */ +export class ThinkingSelectorComponent extends Container { + private selectList: SelectList; + + constructor( + currentLevel: ThinkingLevel, + availableLevels: ThinkingLevel[], + onSelect: (level: ThinkingLevel) => void, + onCancel: () => void, + ) { + super(); + + const thinkingLevels: SelectItem[] = availableLevels.map((level) => ({ + value: level, + label: level, + description: LEVEL_DESCRIPTIONS[level], + })); + + // Add top border + this.addChild(new DynamicBorder()); + + // Create selector + this.selectList = new SelectList( + thinkingLevels, + thinkingLevels.length, + getSelectListTheme(), + THINKING_SELECT_LIST_LAYOUT, + ); + + // Preselect current level + const currentIndex = thinkingLevels.findIndex((item) => item.value === currentLevel); + if (currentIndex !== -1) { + this.selectList.setSelectedIndex(currentIndex); + } + + this.selectList.onSelect = (item) => { + onSelect(item.value as ThinkingLevel); + }; + + this.selectList.onCancel = () => { + onCancel(); + }; + + this.addChild(this.selectList); + + // Add bottom border + this.addChild(new DynamicBorder()); + } + + getSelectList(): SelectList { + return this.selectList; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/tool-execution.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/tool-execution.ts new file mode 100644 index 000000000..ad84f4413 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/tool-execution.ts @@ -0,0 +1,377 @@ +import { Box, type Component, Container, getCapabilities, Image, Spacer, Text, type TUI } from "@earendil-works/pi-tui"; +import type { ToolDefinition, ToolRenderContext } from "../../../core/extensions/types.ts"; +import { createAllToolDefinitions, type ToolName } from "../../../core/tools/index.ts"; +import { getTextOutput as getRenderedTextOutput } from "../../../core/tools/render-utils.ts"; +import { convertToPng } from "../../../utils/image-convert.ts"; +import { theme } from "../theme/theme.ts"; + +export interface ToolExecutionOptions { + showImages?: boolean; + imageWidthCells?: number; +} + +export class ToolExecutionComponent extends Container { + private contentBox: Box; + private contentText: Text; + private selfRenderContainer: Container; + private callRendererComponent?: Component; + private resultRendererComponent?: Component; + private rendererState: any = {}; + private imageComponents: Image[] = []; + private imageSpacers: Spacer[] = []; + private toolName: string; + private toolCallId: string; + private args: any; + private expanded = false; + private showImages: boolean; + private imageWidthCells: number; + private isPartial = true; + private toolDefinition?: ToolDefinition; + private builtInToolDefinition?: ToolDefinition; + private ui: TUI; + private cwd: string; + private executionStarted = false; + private argsComplete = false; + private result?: { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + isError: boolean; + details?: any; + }; + private convertedImages: Map = new Map(); + private hideComponent = false; + + constructor( + toolName: string, + toolCallId: string, + args: any, + options: ToolExecutionOptions = {}, + toolDefinition: ToolDefinition | undefined, + ui: TUI, + cwd: string, + ) { + super(); + this.toolName = toolName; + this.toolCallId = toolCallId; + this.args = args; + this.toolDefinition = toolDefinition; + this.builtInToolDefinition = createAllToolDefinitions(cwd)[toolName as ToolName]; + this.showImages = options.showImages ?? true; + this.imageWidthCells = options.imageWidthCells ?? 60; + this.ui = ui; + this.cwd = cwd; + + this.addChild(new Spacer(1)); + + // Always create all shell variants. contentBox is used for default renderer-based composition. + // selfRenderContainer is used when the tool renders its own framing. + // contentText is reserved for generic fallback rendering when no tool definition exists. + this.contentBox = new Box(1, 1, (text: string) => theme.bg("toolPendingBg", text)); + this.contentText = new Text("", 1, 1, (text: string) => theme.bg("toolPendingBg", text)); + this.selfRenderContainer = new Container(); + + if (this.hasRendererDefinition()) { + this.addChild(this.getRenderShell() === "self" ? this.selfRenderContainer : this.contentBox); + } else { + this.addChild(this.contentText); + } + + this.updateDisplay(); + } + + private getCallRenderer(): ToolDefinition["renderCall"] | undefined { + if (!this.builtInToolDefinition) { + return this.toolDefinition?.renderCall; + } + if (!this.toolDefinition) { + return this.builtInToolDefinition.renderCall; + } + return this.toolDefinition.renderCall ?? this.builtInToolDefinition.renderCall; + } + + private getResultRenderer(): ToolDefinition["renderResult"] | undefined { + if (!this.builtInToolDefinition) { + return this.toolDefinition?.renderResult; + } + if (!this.toolDefinition) { + return this.builtInToolDefinition.renderResult; + } + return this.toolDefinition.renderResult ?? this.builtInToolDefinition.renderResult; + } + + private hasRendererDefinition(): boolean { + return this.builtInToolDefinition !== undefined || this.toolDefinition !== undefined; + } + + private getRenderShell(): "default" | "self" { + if (!this.builtInToolDefinition) { + return this.toolDefinition?.renderShell ?? "default"; + } + if (!this.toolDefinition) { + return this.builtInToolDefinition.renderShell ?? "default"; + } + return this.toolDefinition.renderShell ?? this.builtInToolDefinition.renderShell ?? "default"; + } + + private getRenderContext(lastComponent: Component | undefined): ToolRenderContext { + return { + args: this.args, + toolCallId: this.toolCallId, + invalidate: () => { + this.invalidate(); + this.ui.requestRender(); + }, + lastComponent, + state: this.rendererState, + cwd: this.cwd, + executionStarted: this.executionStarted, + argsComplete: this.argsComplete, + isPartial: this.isPartial, + expanded: this.expanded, + showImages: this.showImages, + isError: this.result?.isError ?? false, + }; + } + + private createCallFallback(): Component { + return new Text(theme.fg("toolTitle", theme.bold(this.toolName)), 0, 0); + } + + private createResultFallback(): Component | undefined { + const output = this.getTextOutput(); + if (!output) { + return undefined; + } + return new Text(theme.fg("toolOutput", output), 0, 0); + } + + updateArgs(args: any): void { + this.args = args; + this.updateDisplay(); + } + + markExecutionStarted(): void { + this.executionStarted = true; + this.updateDisplay(); + this.ui.requestRender(); + } + + setArgsComplete(): void { + this.argsComplete = true; + this.updateDisplay(); + this.ui.requestRender(); + } + + updateResult( + result: { + content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; + details?: any; + isError: boolean; + }, + isPartial = false, + ): void { + this.result = result; + this.isPartial = isPartial; + this.updateDisplay(); + this.maybeConvertImagesForKitty(); + } + + private maybeConvertImagesForKitty(): void { + const caps = getCapabilities(); + if (caps.images !== "kitty") return; + if (!this.result) return; + + const imageBlocks = this.result.content.filter((c) => c.type === "image"); + for (let i = 0; i < imageBlocks.length; i++) { + const img = imageBlocks[i]; + if (!img.data || !img.mimeType) continue; + if (img.mimeType === "image/png") continue; + if (this.convertedImages.has(i)) continue; + + const index = i; + convertToPng(img.data, img.mimeType).then((converted) => { + if (converted) { + this.convertedImages.set(index, converted); + this.updateDisplay(); + this.ui.requestRender(); + } + }); + } + } + + setExpanded(expanded: boolean): void { + this.expanded = expanded; + this.updateDisplay(); + } + + setShowImages(show: boolean): void { + this.showImages = show; + this.updateDisplay(); + } + + setImageWidthCells(width: number): void { + this.imageWidthCells = Math.max(1, Math.floor(width)); + this.updateDisplay(); + } + + override invalidate(): void { + super.invalidate(); + this.updateDisplay(); + } + + override render(width: number): string[] { + if (this.hideComponent) { + return []; + } + + if (this.hasRendererDefinition() && this.getRenderShell() === "self") { + const contentLines = this.selfRenderContainer.render(width); + if (contentLines.length === 0 && this.imageComponents.length === 0) { + return []; + } + + const lines: string[] = []; + if (contentLines.length > 0) { + lines.push(""); + lines.push(...contentLines); + } + for (let i = 0; i < this.imageComponents.length; i++) { + const spacer = this.imageSpacers[i]; + if (spacer) { + lines.push(...spacer.render(width)); + } + const imageComponent = this.imageComponents[i]; + if (imageComponent) { + lines.push(...imageComponent.render(width)); + } + } + return lines; + } + + return super.render(width); + } + + private updateDisplay(): void { + const bgFn = this.isPartial + ? (text: string) => theme.bg("toolPendingBg", text) + : this.result?.isError + ? (text: string) => theme.bg("toolErrorBg", text) + : (text: string) => theme.bg("toolSuccessBg", text); + + let hasContent = false; + this.hideComponent = false; + if (this.hasRendererDefinition()) { + const renderContainer = this.getRenderShell() === "self" ? this.selfRenderContainer : this.contentBox; + if (renderContainer instanceof Box) { + renderContainer.setBgFn(bgFn); + } + renderContainer.clear(); + + const callRenderer = this.getCallRenderer(); + if (!callRenderer) { + renderContainer.addChild(this.createCallFallback()); + hasContent = true; + } else { + try { + const component = callRenderer(this.args, theme, this.getRenderContext(this.callRendererComponent)); + this.callRendererComponent = component; + renderContainer.addChild(component); + hasContent = true; + } catch { + this.callRendererComponent = undefined; + renderContainer.addChild(this.createCallFallback()); + hasContent = true; + } + } + + if (this.result) { + const resultRenderer = this.getResultRenderer(); + if (!resultRenderer) { + const component = this.createResultFallback(); + if (component) { + renderContainer.addChild(component); + hasContent = true; + } + } else { + try { + const component = resultRenderer( + { content: this.result.content as any, details: this.result.details }, + { expanded: this.expanded, isPartial: this.isPartial }, + theme, + this.getRenderContext(this.resultRendererComponent), + ); + this.resultRendererComponent = component; + renderContainer.addChild(component); + hasContent = true; + } catch { + this.resultRendererComponent = undefined; + const component = this.createResultFallback(); + if (component) { + renderContainer.addChild(component); + hasContent = true; + } + } + } + } + } else { + this.contentText.setCustomBgFn(bgFn); + this.contentText.setText(this.formatToolExecution()); + hasContent = true; + } + + for (const img of this.imageComponents) { + this.removeChild(img); + } + this.imageComponents = []; + for (const spacer of this.imageSpacers) { + this.removeChild(spacer); + } + this.imageSpacers = []; + + if (this.result) { + const imageBlocks = this.result.content.filter((c) => c.type === "image"); + const caps = getCapabilities(); + for (let i = 0; i < imageBlocks.length; i++) { + const img = imageBlocks[i]; + if (caps.images && this.showImages && img.data && img.mimeType) { + const converted = this.convertedImages.get(i); + const imageData = converted?.data ?? img.data; + const imageMimeType = converted?.mimeType ?? img.mimeType; + if (caps.images === "kitty" && imageMimeType !== "image/png") continue; + + const spacer = new Spacer(1); + this.addChild(spacer); + this.imageSpacers.push(spacer); + const imageComponent = new Image( + imageData, + imageMimeType, + { fallbackColor: (s: string) => theme.fg("toolOutput", s) }, + { maxWidthCells: this.imageWidthCells }, + ); + this.imageComponents.push(imageComponent); + this.addChild(imageComponent); + } + } + } + + if (this.hasRendererDefinition() && !hasContent && this.imageComponents.length === 0) { + this.hideComponent = true; + } + } + + private getTextOutput(): string { + return getRenderedTextOutput(this.result, this.showImages); + } + + private formatToolExecution(): string { + let text = theme.fg("toolTitle", theme.bold(this.toolName)); + const content = JSON.stringify(this.args, null, 2); + if (content) { + text += `\n\n${content}`; + } + const output = this.getTextOutput(); + if (output) { + text += `\n${output}`; + } + return text; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/tree-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/tree-selector.ts new file mode 100644 index 000000000..8bcdc2fb8 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/tree-selector.ts @@ -0,0 +1,1386 @@ +import { + type Component, + Container, + type Focusable, + getKeybindings, + Input, + type Keybinding, + Spacer, + sliceByColumn, + Text, + truncateToWidth, + visibleWidth, + wrapTextWithAnsi, +} from "@earendil-works/pi-tui"; +import type { SessionTreeNode } from "../../../core/session-manager.ts"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { formatKeyText, keyHint } from "./keybinding-hints.ts"; + +/** Gutter info: position (displayIndent where connector was) and whether to show │ */ +interface GutterInfo { + position: number; // displayIndent level where the connector was shown + show: boolean; // true = show │, false = show spaces +} + +/** Flattened tree node for navigation */ +interface FlatNode { + node: SessionTreeNode; + /** Indentation level (each level = 3 chars) */ + indent: number; + /** Whether to show connector (├─ or └─) - true if parent has multiple children */ + showConnector: boolean; + /** If showConnector, true = last sibling (└─), false = not last (├─) */ + isLast: boolean; + /** Gutter info for each ancestor branch point */ + gutters: GutterInfo[]; + /** True if this node is a root under a virtual branching root (multiple roots) */ + isVirtualRootChild: boolean; +} + +interface HorizontalViewportRow { + gutter: string; + body: string; + anchorCol: number; + bodyWidth: number; + isSelected: boolean; +} + +const TREE_GUTTER_WIDTH = 2; +const MIN_VISIBLE_ANCHOR_CONTENT_WIDTH = 4; +const MAX_VISIBLE_ANCHOR_CONTENT_WIDTH = 20; +const MIN_ANCHOR_CONTEXT_WIDTH = 2; +const MAX_ANCHOR_CONTEXT_WIDTH = 12; + +/** + * Render tree rows into a horizontally clipped viewport. + * + * The tree gutter is always kept visible. The row bodies are shifted left only + * when the selected row's anchor (the start of its entry text after tree + * indentation/markers) would otherwise be too far right to see useful content. + */ +function renderHorizontalViewport(rows: HorizontalViewportRow[], width: number): string[] { + const viewportWidth = Math.max(0, width - TREE_GUTTER_WIDTH); + const maxBodyWidth = rows.reduce((max, row) => Math.max(max, row.bodyWidth), 0); + const maxHorizontalScroll = Math.max(0, maxBodyWidth - viewportWidth); + const selectedRow = rows.find((row) => row.isSelected); + + // Only pan horizontally when needed to keep enough selected-row content visible after its anchor. + let horizontalScroll = 0; + if (selectedRow && maxHorizontalScroll > 0) { + const minVisibleAnchorContentWidth = Math.min( + MAX_VISIBLE_ANCHOR_CONTENT_WIDTH, + Math.max(MIN_VISIBLE_ANCHOR_CONTENT_WIDTH, Math.floor(viewportWidth / 3)), + ); + if (selectedRow.anchorCol > viewportWidth - minVisibleAnchorContentWidth) { + const anchorContextWidth = Math.min( + MAX_ANCHOR_CONTEXT_WIDTH, + Math.max(MIN_ANCHOR_CONTEXT_WIDTH, Math.floor(viewportWidth / 4)), + ); + horizontalScroll = Math.min(maxHorizontalScroll, selectedRow.anchorCol - anchorContextWidth); + } + } + + // Clip only the body; the fixed-width gutter remains visible as navigation context. + return rows.map((row) => { + const line = + horizontalScroll > 0 + ? `${row.gutter}${sliceByColumn(row.body, horizontalScroll, viewportWidth, true)}\x1b[0m` + : row.gutter + row.body; + return truncateToWidth(line, width, ""); + }); +} + +/** Filter mode for tree display */ +export type FilterMode = "default" | "no-tools" | "user-only" | "labeled-only" | "all"; + +/** + * Tree list component with selection and ASCII art visualization + */ +/** Tool call info for lookup */ +interface ToolCallInfo { + name: string; + arguments: Record; +} + +class TreeList implements Component { + private flatNodes: FlatNode[] = []; + private filteredNodes: FlatNode[] = []; + private selectedIndex = 0; + private currentLeafId: string | null; + private maxVisibleLines: number; + private filterMode: FilterMode = "default"; + private searchQuery = ""; + private toolCallMap: Map = new Map(); + private multipleRoots = false; + private showLabelTimestamps = false; + private activePathIds: Set = new Set(); + private visibleParentMap: Map = new Map(); + private visibleChildrenMap: Map = new Map(); + private lastSelectedId: string | null = null; + private foldedNodes: Set = new Set(); + + public onSelect?: (entryId: string) => void; + public onCancel?: () => void; + public onLabelEdit?: (entryId: string, currentLabel: string | undefined) => void; + + constructor( + tree: SessionTreeNode[], + currentLeafId: string | null, + maxVisibleLines: number, + initialSelectedId?: string, + initialFilterMode?: FilterMode, + ) { + this.currentLeafId = currentLeafId; + this.maxVisibleLines = maxVisibleLines; + this.filterMode = initialFilterMode ?? "default"; + this.multipleRoots = tree.length > 1; + this.flatNodes = this.flattenTree(tree); + this.buildActivePath(); + this.applyFilter(); + + // Start with initialSelectedId if provided, otherwise current leaf + const targetId = initialSelectedId ?? currentLeafId; + this.selectedIndex = this.findNearestVisibleIndex(targetId); + this.lastSelectedId = this.filteredNodes[this.selectedIndex]?.node.entry.id ?? null; + } + + /** + * Find the index of the nearest visible entry, walking up the parent chain if needed. + * Returns the index in filteredNodes, or the last index as fallback. + */ + private findNearestVisibleIndex(entryId: string | null): number { + if (this.filteredNodes.length === 0) return 0; + + // Build a map for parent lookup + const entryMap = new Map(); + for (const flatNode of this.flatNodes) { + entryMap.set(flatNode.node.entry.id, flatNode); + } + + // Build a map of visible entry IDs to their indices in filteredNodes + const visibleIdToIndex = new Map(this.filteredNodes.map((node, i) => [node.node.entry.id, i])); + + // Walk from entryId up to root, looking for a visible entry + let currentId = entryId; + while (currentId !== null) { + const index = visibleIdToIndex.get(currentId); + if (index !== undefined) return index; + const node = entryMap.get(currentId); + if (!node) break; + currentId = node.node.entry.parentId ?? null; + } + + // Fallback: last visible entry + return this.filteredNodes.length - 1; + } + + /** Build the set of entry IDs on the path from root to current leaf */ + private buildActivePath(): void { + this.activePathIds.clear(); + if (!this.currentLeafId) return; + + // Build a map of id -> entry for parent lookup + const entryMap = new Map(); + for (const flatNode of this.flatNodes) { + entryMap.set(flatNode.node.entry.id, flatNode); + } + + // Walk from leaf to root + let currentId: string | null = this.currentLeafId; + while (currentId) { + this.activePathIds.add(currentId); + const node = entryMap.get(currentId); + if (!node) break; + currentId = node.node.entry.parentId ?? null; + } + } + + private flattenTree(roots: SessionTreeNode[]): FlatNode[] { + const result: FlatNode[] = []; + this.toolCallMap.clear(); + + // Indentation rules: + // - At indent 0: stay at 0 unless parent has >1 children (then +1) + // - At indent 1: children always go to indent 2 (visual grouping of subtree) + // - At indent 2+: stay flat for single-child chains, +1 only if parent branches + + // Stack items: [node, indent, justBranched, showConnector, isLast, gutters, isVirtualRootChild] + type StackItem = [SessionTreeNode, number, boolean, boolean, boolean, GutterInfo[], boolean]; + const stack: StackItem[] = []; + + // Determine which subtrees contain the active leaf (to sort current branch first) + // Use iterative post-order traversal to avoid stack overflow + const containsActive = new Map(); + const leafId = this.currentLeafId; + { + // Build list in pre-order, then process in reverse for post-order effect + const allNodes: SessionTreeNode[] = []; + const preOrderStack: SessionTreeNode[] = [...roots]; + while (preOrderStack.length > 0) { + const node = preOrderStack.pop()!; + allNodes.push(node); + // Push children in reverse so they're processed left-to-right + for (let i = node.children.length - 1; i >= 0; i--) { + preOrderStack.push(node.children[i]); + } + } + // Process in reverse (post-order): children before parents + for (let i = allNodes.length - 1; i >= 0; i--) { + const node = allNodes[i]; + let has = leafId !== null && node.entry.id === leafId; + for (const child of node.children) { + if (containsActive.get(child)) { + has = true; + } + } + containsActive.set(node, has); + } + } + + // Add roots in reverse order, prioritizing the one containing the active leaf + // If multiple roots, treat them as children of a virtual root that branches + const multipleRoots = roots.length > 1; + const orderedRoots = [...roots].sort((a, b) => Number(containsActive.get(b)) - Number(containsActive.get(a))); + for (let i = orderedRoots.length - 1; i >= 0; i--) { + const isLast = i === orderedRoots.length - 1; + stack.push([orderedRoots[i], multipleRoots ? 1 : 0, multipleRoots, multipleRoots, isLast, [], multipleRoots]); + } + + while (stack.length > 0) { + const [node, indent, justBranched, showConnector, isLast, gutters, isVirtualRootChild] = stack.pop()!; + + // Extract tool calls from assistant messages for later lookup + const entry = node.entry; + if (entry.type === "message" && entry.message.role === "assistant") { + const content = (entry.message as { content?: unknown }).content; + if (Array.isArray(content)) { + for (const block of content) { + if (typeof block === "object" && block !== null && "type" in block && block.type === "toolCall") { + const tc = block as { id: string; name: string; arguments: Record }; + this.toolCallMap.set(tc.id, { name: tc.name, arguments: tc.arguments }); + } + } + } + } + + result.push({ node, indent, showConnector, isLast, gutters, isVirtualRootChild }); + + const children = node.children; + const multipleChildren = children.length > 1; + + // Order children so the branch containing the active leaf comes first + const orderedChildren = (() => { + const prioritized: SessionTreeNode[] = []; + const rest: SessionTreeNode[] = []; + for (const child of children) { + if (containsActive.get(child)) { + prioritized.push(child); + } else { + rest.push(child); + } + } + return [...prioritized, ...rest]; + })(); + + // Calculate child indent + let childIndent: number; + if (multipleChildren) { + // Parent branches: children get +1 + childIndent = indent + 1; + } else if (justBranched && indent > 0) { + // First generation after a branch: +1 for visual grouping + childIndent = indent + 1; + } else { + // Single-child chain: stay flat + childIndent = indent; + } + + // Build gutters for children + // If this node showed a connector, add a gutter entry for descendants + // Only add gutter if connector is actually displayed (not suppressed for virtual root children) + const connectorDisplayed = showConnector && !isVirtualRootChild; + // When connector is displayed, add a gutter entry at the connector's position + // Connector is at position (displayIndent - 1), so gutter should be there too + const currentDisplayIndent = this.multipleRoots ? Math.max(0, indent - 1) : indent; + const connectorPosition = Math.max(0, currentDisplayIndent - 1); + const childGutters: GutterInfo[] = connectorDisplayed + ? [...gutters, { position: connectorPosition, show: !isLast }] + : gutters; + + // Add children in reverse order + for (let i = orderedChildren.length - 1; i >= 0; i--) { + const childIsLast = i === orderedChildren.length - 1; + stack.push([ + orderedChildren[i], + childIndent, + multipleChildren, + multipleChildren, + childIsLast, + childGutters, + false, + ]); + } + } + + return result; + } + + private applyFilter(): void { + // Update lastSelectedId only when we have a valid selection (non-empty list) + // This preserves the selection when switching through empty filter results + if (this.filteredNodes.length > 0) { + this.lastSelectedId = this.filteredNodes[this.selectedIndex]?.node.entry.id ?? this.lastSelectedId; + } + + const searchTokens = this.searchQuery.toLowerCase().split(/\s+/).filter(Boolean); + + this.filteredNodes = this.flatNodes.filter((flatNode) => { + const entry = flatNode.node.entry; + const isCurrentLeaf = entry.id === this.currentLeafId; + + // Skip assistant messages with only tool calls (no text) unless error/aborted + // Always show current leaf so active position is visible + if (entry.type === "message" && entry.message.role === "assistant" && !isCurrentLeaf) { + const msg = entry.message as { stopReason?: string; content?: unknown }; + const hasText = this.hasTextContent(msg.content); + const isErrorOrAborted = msg.stopReason && msg.stopReason !== "stop" && msg.stopReason !== "toolUse"; + // Only hide if no text AND not an error/aborted message + if (!hasText && !isErrorOrAborted) { + return false; + } + } + + // Apply filter mode + let passesFilter = true; + // Entry types hidden in default view (settings/bookkeeping) + const isSettingsEntry = + entry.type === "label" || + entry.type === "custom" || + entry.type === "model_change" || + entry.type === "thinking_level_change" || + entry.type === "session_info"; + + switch (this.filterMode) { + case "user-only": + // Just user messages + passesFilter = entry.type === "message" && entry.message.role === "user"; + break; + case "no-tools": + // Default minus tool results + passesFilter = !isSettingsEntry && !(entry.type === "message" && entry.message.role === "toolResult"); + break; + case "labeled-only": + // Just labeled entries + passesFilter = flatNode.node.label !== undefined; + break; + case "all": + // Show everything + passesFilter = true; + break; + default: + // Default mode: hide settings/bookkeeping entries + passesFilter = !isSettingsEntry; + break; + } + + if (!passesFilter) return false; + + // Apply search filter + if (searchTokens.length > 0) { + const nodeText = this.getSearchableText(flatNode.node).toLowerCase(); + return searchTokens.every((token) => nodeText.includes(token)); + } + + return true; + }); + + // Filter out descendants of folded nodes. + if (this.foldedNodes.size > 0) { + const skipSet = new Set(); + for (const flatNode of this.flatNodes) { + const { id, parentId } = flatNode.node.entry; + if (parentId != null && (this.foldedNodes.has(parentId) || skipSet.has(parentId))) { + skipSet.add(id); + } + } + this.filteredNodes = this.filteredNodes.filter((flatNode) => !skipSet.has(flatNode.node.entry.id)); + } + + // Recalculate visual structure (indent, connectors, gutters) based on visible tree + this.recalculateVisualStructure(); + + // Try to preserve cursor on the same node, or find nearest visible ancestor + if (this.lastSelectedId) { + this.selectedIndex = this.findNearestVisibleIndex(this.lastSelectedId); + } else if (this.selectedIndex >= this.filteredNodes.length) { + // Clamp index if out of bounds + this.selectedIndex = Math.max(0, this.filteredNodes.length - 1); + } + + // Update lastSelectedId to the actual selection (may have changed due to parent walk) + if (this.filteredNodes.length > 0) { + this.lastSelectedId = this.filteredNodes[this.selectedIndex]?.node.entry.id ?? this.lastSelectedId; + } + } + + /** + * Recompute indentation/connectors for the filtered view + * + * Filtering can hide intermediate entries; descendants attach to the nearest visible ancestor. + * Keep indentation semantics aligned with flattenTree() so single-child chains don't drift right. + */ + private recalculateVisualStructure(): void { + if (this.filteredNodes.length === 0) return; + + const visibleIds = new Set(this.filteredNodes.map((n) => n.node.entry.id)); + + // Build entry map for efficient parent lookup (using full tree) + const entryMap = new Map(); + for (const flatNode of this.flatNodes) { + entryMap.set(flatNode.node.entry.id, flatNode); + } + + // Find nearest visible ancestor for a node + const findVisibleAncestor = (nodeId: string): string | null => { + let currentId = entryMap.get(nodeId)?.node.entry.parentId ?? null; + while (currentId !== null) { + if (visibleIds.has(currentId)) { + return currentId; + } + currentId = entryMap.get(currentId)?.node.entry.parentId ?? null; + } + return null; + }; + + // Build visible tree structure: + // - visibleParent: nodeId → nearest visible ancestor (or null for roots) + // - visibleChildren: parentId → list of visible children (in filteredNodes order) + const visibleParent = new Map(); + const visibleChildren = new Map(); + visibleChildren.set(null, []); // root-level nodes + + for (const flatNode of this.filteredNodes) { + const nodeId = flatNode.node.entry.id; + const ancestorId = findVisibleAncestor(nodeId); + visibleParent.set(nodeId, ancestorId); + + if (!visibleChildren.has(ancestorId)) { + visibleChildren.set(ancestorId, []); + } + visibleChildren.get(ancestorId)!.push(nodeId); + } + + // Update multipleRoots based on visible roots + const visibleRootIds = visibleChildren.get(null)!; + this.multipleRoots = visibleRootIds.length > 1; + + // Build a map for quick lookup: nodeId → FlatNode + const filteredNodeMap = new Map(); + for (const flatNode of this.filteredNodes) { + filteredNodeMap.set(flatNode.node.entry.id, flatNode); + } + + // DFS over the visible tree using flattenTree() indentation semantics + // Stack items: [nodeId, indent, justBranched, showConnector, isLast, gutters, isVirtualRootChild] + type StackItem = [string, number, boolean, boolean, boolean, GutterInfo[], boolean]; + const stack: StackItem[] = []; + + // Add visible roots in reverse order (to process in forward order via stack) + for (let i = visibleRootIds.length - 1; i >= 0; i--) { + const isLast = i === visibleRootIds.length - 1; + stack.push([ + visibleRootIds[i], + this.multipleRoots ? 1 : 0, + this.multipleRoots, + this.multipleRoots, + isLast, + [], + this.multipleRoots, + ]); + } + + while (stack.length > 0) { + const [nodeId, indent, justBranched, showConnector, isLast, gutters, isVirtualRootChild] = stack.pop()!; + + const flatNode = filteredNodeMap.get(nodeId); + if (!flatNode) continue; + + // Update this node's visual properties + flatNode.indent = indent; + flatNode.showConnector = showConnector; + flatNode.isLast = isLast; + flatNode.gutters = gutters; + flatNode.isVirtualRootChild = isVirtualRootChild; + + // Get visible children of this node + const children = visibleChildren.get(nodeId) || []; + const multipleChildren = children.length > 1; + + // Child indent follows flattenTree(): branch points (and first generation after a branch) shift +1 + let childIndent: number; + if (multipleChildren) { + childIndent = indent + 1; + } else if (justBranched && indent > 0) { + childIndent = indent + 1; + } else { + childIndent = indent; + } + + // Child gutters follow flattenTree() connector/gutter rules + const connectorDisplayed = showConnector && !isVirtualRootChild; + const currentDisplayIndent = this.multipleRoots ? Math.max(0, indent - 1) : indent; + const connectorPosition = Math.max(0, currentDisplayIndent - 1); + const childGutters: GutterInfo[] = connectorDisplayed + ? [...gutters, { position: connectorPosition, show: !isLast }] + : gutters; + + // Add children in reverse order (to process in forward order via stack) + for (let i = children.length - 1; i >= 0; i--) { + const childIsLast = i === children.length - 1; + stack.push([ + children[i], + childIndent, + multipleChildren, + multipleChildren, + childIsLast, + childGutters, + false, + ]); + } + } + + // Store visible tree maps for ancestor/descendant lookups in navigation + this.visibleParentMap = visibleParent; + this.visibleChildrenMap = visibleChildren; + } + + /** Get searchable text content from a node */ + private getSearchableText(node: SessionTreeNode): string { + const entry = node.entry; + const parts: string[] = []; + + if (node.label) { + parts.push(node.label); + } + + switch (entry.type) { + case "message": { + const msg = entry.message; + parts.push(msg.role); + if ("content" in msg && msg.content) { + parts.push(this.extractContent(msg.content)); + } + if (msg.role === "bashExecution") { + const bashMsg = msg as { command?: string }; + if (bashMsg.command) parts.push(bashMsg.command); + } + break; + } + case "custom_message": { + parts.push(entry.customType); + if (typeof entry.content === "string") { + parts.push(entry.content); + } else { + parts.push(this.extractContent(entry.content)); + } + break; + } + case "compaction": + parts.push("compaction"); + break; + case "branch_summary": + parts.push("branch summary", entry.summary); + break; + case "session_info": + parts.push("title"); + if (entry.name) parts.push(entry.name); + break; + case "model_change": + parts.push("model", entry.modelId); + break; + case "thinking_level_change": + parts.push("thinking", entry.thinkingLevel); + break; + case "custom": + parts.push("custom", entry.customType); + break; + case "label": + parts.push("label", entry.label ?? ""); + break; + } + + return parts.join(" "); + } + + invalidate(): void {} + + getSearchQuery(): string { + return this.searchQuery; + } + + getSelectedNode(): SessionTreeNode | undefined { + return this.filteredNodes[this.selectedIndex]?.node; + } + + updateNodeLabel(entryId: string, label: string | undefined, labelTimestamp?: string): void { + for (const flatNode of this.flatNodes) { + if (flatNode.node.entry.id === entryId) { + flatNode.node.label = label; + flatNode.node.labelTimestamp = label ? (labelTimestamp ?? new Date().toISOString()) : undefined; + break; + } + } + } + + private getStatusLabels(): string { + let labels = ""; + switch (this.filterMode) { + case "no-tools": + labels += " [no-tools]"; + break; + case "user-only": + labels += " [user]"; + break; + case "labeled-only": + labels += " [labeled]"; + break; + case "all": + labels += " [all]"; + break; + } + if (this.showLabelTimestamps) { + labels += " [+label time]"; + } + return labels; + } + + render(width: number): string[] { + const lines: string[] = []; + + if (this.filteredNodes.length === 0) { + lines.push(truncateToWidth(theme.fg("muted", " No entries found"), width)); + lines.push(truncateToWidth(theme.fg("muted", ` (0/0)${this.getStatusLabels()}`), width)); + return lines; + } + + const startIndex = Math.max( + 0, + Math.min( + this.selectedIndex - Math.floor(this.maxVisibleLines / 2), + this.filteredNodes.length - this.maxVisibleLines, + ), + ); + const endIndex = Math.min(startIndex + this.maxVisibleLines, this.filteredNodes.length); + + const renderedRows: HorizontalViewportRow[] = []; + for (let i = startIndex; i < endIndex; i++) { + const flatNode = this.filteredNodes[i]; + const entry = flatNode.node.entry; + const isSelected = i === this.selectedIndex; + + // Build line: cursor + prefix + path marker + label + content + const cursor = isSelected ? theme.fg("accent", "› ") : " "; + + // If multiple roots, shift display (roots at 0, not 1) + const displayIndent = this.multipleRoots ? Math.max(0, flatNode.indent - 1) : flatNode.indent; + + // Build prefix with gutters at their correct positions + // Each gutter has a position (displayIndent where its connector was shown) + const connector = + flatNode.showConnector && !flatNode.isVirtualRootChild ? (flatNode.isLast ? "└─ " : "├─ ") : ""; + const connectorPosition = connector ? displayIndent - 1 : -1; + + // Build prefix char by char, placing gutters and connector at their positions + const totalChars = displayIndent * 3; + const prefixChars: string[] = []; + const isFolded = this.foldedNodes.has(entry.id); + for (let i = 0; i < totalChars; i++) { + const level = Math.floor(i / 3); + const posInLevel = i % 3; + + // Check if there's a gutter at this level + const gutter = flatNode.gutters.find((g) => g.position === level); + if (gutter) { + if (posInLevel === 0) { + prefixChars.push(gutter.show ? "│" : " "); + } else { + prefixChars.push(" "); + } + } else if (connector && level === connectorPosition) { + // Connector at this level, with fold indicator + if (posInLevel === 0) { + prefixChars.push(flatNode.isLast ? "└" : "├"); + } else if (posInLevel === 1) { + const foldable = this.isFoldable(entry.id); + prefixChars.push(isFolded ? "⊞" : foldable ? "⊟" : "─"); + } else { + prefixChars.push(" "); + } + } else { + prefixChars.push(" "); + } + } + const prefix = prefixChars.join(""); + + // Fold marker for nodes without connectors (roots) + const showsFoldInConnector = flatNode.showConnector && !flatNode.isVirtualRootChild; + const foldMarker = isFolded && !showsFoldInConnector ? theme.fg("accent", "⊞ ") : ""; + + // Active path marker - shown right before the entry text + const isOnActivePath = this.activePathIds.has(entry.id); + const pathMarker = isOnActivePath ? theme.fg("accent", "• ") : ""; + + const label = flatNode.node.label ? theme.fg("warning", `[${flatNode.node.label}] `) : ""; + const labelTimestamp = + this.showLabelTimestamps && flatNode.node.label && flatNode.node.labelTimestamp + ? theme.fg("muted", `${this.formatLabelTimestamp(flatNode.node.labelTimestamp)} `) + : ""; + const content = this.getEntryDisplayText(flatNode.node, isSelected); + const prefixPart = theme.fg("dim", prefix) + foldMarker + pathMarker; + const anchorCol = visibleWidth(prefixPart); + let gutter = cursor; + let body = prefixPart + label + labelTimestamp + content; + if (isSelected) { + gutter = theme.bg("selectedBg", gutter); + body = theme.bg("selectedBg", body); + } + renderedRows.push({ gutter, body, anchorCol, bodyWidth: visibleWidth(body), isSelected }); + } + + lines.push(...renderHorizontalViewport(renderedRows, width)); + lines.push( + truncateToWidth( + theme.fg("muted", ` (${this.selectedIndex + 1}/${this.filteredNodes.length})${this.getStatusLabels()}`), + width, + ), + ); + + return lines; + } + + private getEntryDisplayText(node: SessionTreeNode, isSelected: boolean): string { + const entry = node.entry; + let result: string; + + const normalize = (s: string) => s.replace(/[\n\t]/g, " ").trim(); + + switch (entry.type) { + case "message": { + const msg = entry.message; + const role = msg.role; + if (role === "user") { + const msgWithContent = msg as { content?: unknown }; + const content = normalize(this.extractContent(msgWithContent.content)); + result = theme.fg("accent", "user: ") + content; + } else if (role === "assistant") { + const msgWithContent = msg as { content?: unknown; stopReason?: string; errorMessage?: string }; + const textContent = normalize(this.extractContent(msgWithContent.content)); + if (textContent) { + result = theme.fg("success", "assistant: ") + textContent; + } else if (msgWithContent.stopReason === "aborted") { + result = theme.fg("success", "assistant: ") + theme.fg("muted", "(aborted)"); + } else if (msgWithContent.errorMessage) { + const errMsg = normalize(msgWithContent.errorMessage).slice(0, 80); + result = theme.fg("success", "assistant: ") + theme.fg("error", errMsg); + } else { + result = theme.fg("success", "assistant: ") + theme.fg("muted", "(no content)"); + } + } else if (role === "toolResult") { + const toolMsg = msg as { toolCallId?: string; toolName?: string }; + const toolCall = toolMsg.toolCallId ? this.toolCallMap.get(toolMsg.toolCallId) : undefined; + if (toolCall) { + result = theme.fg("muted", this.formatToolCall(toolCall.name, toolCall.arguments)); + } else { + result = theme.fg("muted", `[${toolMsg.toolName ?? "tool"}]`); + } + } else if (role === "bashExecution") { + const bashMsg = msg as { command?: string }; + result = theme.fg("dim", `[bash]: ${normalize(bashMsg.command ?? "")}`); + } else { + result = theme.fg("dim", `[${role}]`); + } + break; + } + case "custom_message": { + const content = + typeof entry.content === "string" + ? entry.content + : entry.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + result = theme.fg("customMessageLabel", `[${entry.customType}]: `) + normalize(content); + break; + } + case "compaction": { + const tokens = Math.round(entry.tokensBefore / 1000); + result = theme.fg("borderAccent", `[compaction: ${tokens}k tokens]`); + break; + } + case "branch_summary": + result = theme.fg("warning", `[branch summary]: `) + normalize(entry.summary); + break; + case "model_change": + result = theme.fg("dim", `[model: ${entry.modelId}]`); + break; + case "thinking_level_change": + result = theme.fg("dim", `[thinking: ${entry.thinkingLevel}]`); + break; + case "custom": + result = theme.fg("dim", `[custom: ${entry.customType}]`); + break; + case "label": + result = theme.fg("dim", `[label: ${entry.label ?? "(cleared)"}]`); + break; + case "session_info": + result = entry.name + ? [theme.fg("dim", "[title: "), theme.fg("dim", entry.name), theme.fg("dim", "]")].join("") + : [theme.fg("dim", "[title: "), theme.italic(theme.fg("dim", "empty")), theme.fg("dim", "]")].join(""); + break; + default: + result = ""; + } + + return isSelected ? theme.bold(result) : result; + } + + private formatLabelTimestamp(timestamp: string): string { + const date = new Date(timestamp); + const now = new Date(); + const hours = date.getHours().toString().padStart(2, "0"); + const minutes = date.getMinutes().toString().padStart(2, "0"); + const time = `${hours}:${minutes}`; + + if ( + date.getFullYear() === now.getFullYear() && + date.getMonth() === now.getMonth() && + date.getDate() === now.getDate() + ) { + return time; + } + + const month = date.getMonth() + 1; + const day = date.getDate(); + if (date.getFullYear() === now.getFullYear()) { + return `${month}/${day} ${time}`; + } + + const year = date.getFullYear().toString().slice(-2); + return `${year}/${month}/${day} ${time}`; + } + + private extractContent(content: unknown): string { + const maxLen = 200; + if (typeof content === "string") return content.slice(0, maxLen); + if (Array.isArray(content)) { + let result = ""; + for (const c of content) { + if (typeof c === "object" && c !== null && "type" in c && c.type === "text") { + result += (c as { text: string }).text; + if (result.length >= maxLen) return result.slice(0, maxLen); + } + } + return result; + } + return ""; + } + + private hasTextContent(content: unknown): boolean { + if (typeof content === "string") return content.trim().length > 0; + if (Array.isArray(content)) { + for (const c of content) { + if (typeof c === "object" && c !== null && "type" in c && c.type === "text") { + const text = (c as { text?: string }).text; + if (text && text.trim().length > 0) return true; + } + } + } + return false; + } + + private formatToolCall(name: string, args: Record): string { + const shortenPath = (p: string): string => { + const home = process.env.HOME || process.env.USERPROFILE || ""; + if (home && p.startsWith(home)) return `~${p.slice(home.length)}`; + return p; + }; + + switch (name) { + case "read": { + const path = shortenPath(String(args.path || args.file_path || "")); + const offset = args.offset as number | undefined; + const limit = args.limit as number | undefined; + let display = path; + if (offset !== undefined || limit !== undefined) { + const start = offset ?? 1; + const end = limit !== undefined ? start + limit - 1 : ""; + display += `:${start}${end ? `-${end}` : ""}`; + } + return `[read: ${display}]`; + } + case "write": { + const path = shortenPath(String(args.path || args.file_path || "")); + return `[write: ${path}]`; + } + case "edit": { + const path = shortenPath(String(args.path || args.file_path || "")); + return `[edit: ${path}]`; + } + case "bash": { + const rawCmd = String(args.command || ""); + const cmd = rawCmd + .replace(/[\n\t]/g, " ") + .trim() + .slice(0, 50); + return `[bash: ${cmd}${rawCmd.length > 50 ? "..." : ""}]`; + } + case "grep": { + const pattern = String(args.pattern || ""); + const path = shortenPath(String(args.path || ".")); + return `[grep: /${pattern}/ in ${path}]`; + } + case "find": { + const pattern = String(args.pattern || ""); + const path = shortenPath(String(args.path || ".")); + return `[find: ${pattern} in ${path}]`; + } + case "ls": { + const path = shortenPath(String(args.path || ".")); + return `[ls: ${path}]`; + } + default: { + // Custom tool - show name and truncated JSON args + const argsStr = JSON.stringify(args).slice(0, 40); + return `[${name}: ${argsStr}${JSON.stringify(args).length > 40 ? "..." : ""}]`; + } + } + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "tui.select.up")) { + this.selectedIndex = this.selectedIndex === 0 ? this.filteredNodes.length - 1 : this.selectedIndex - 1; + } else if (kb.matches(keyData, "tui.select.down")) { + this.selectedIndex = this.selectedIndex === this.filteredNodes.length - 1 ? 0 : this.selectedIndex + 1; + } else if (kb.matches(keyData, "app.tree.foldOrUp")) { + const currentId = this.filteredNodes[this.selectedIndex]?.node.entry.id; + if (currentId && this.isFoldable(currentId) && !this.foldedNodes.has(currentId)) { + this.foldedNodes.add(currentId); + this.applyFilter(); + } else { + this.selectedIndex = this.findBranchSegmentStart("up"); + } + } else if (kb.matches(keyData, "app.tree.unfoldOrDown")) { + const currentId = this.filteredNodes[this.selectedIndex]?.node.entry.id; + if (currentId && this.foldedNodes.has(currentId)) { + this.foldedNodes.delete(currentId); + this.applyFilter(); + } else { + this.selectedIndex = this.findBranchSegmentStart("down"); + } + } else if (kb.matches(keyData, "tui.editor.cursorLeft") || kb.matches(keyData, "tui.select.pageUp")) { + // Page up + this.selectedIndex = Math.max(0, this.selectedIndex - this.maxVisibleLines); + } else if (kb.matches(keyData, "tui.editor.cursorRight") || kb.matches(keyData, "tui.select.pageDown")) { + // Page down + this.selectedIndex = Math.min(this.filteredNodes.length - 1, this.selectedIndex + this.maxVisibleLines); + } else if (kb.matches(keyData, "tui.select.confirm")) { + const selected = this.filteredNodes[this.selectedIndex]; + if (selected && this.onSelect) { + this.onSelect(selected.node.entry.id); + } + } else if (kb.matches(keyData, "tui.select.cancel")) { + if (this.searchQuery) { + this.searchQuery = ""; + this.foldedNodes.clear(); + this.applyFilter(); + } else { + this.onCancel?.(); + } + } else if (kb.matches(keyData, "app.tree.filter.default")) { + // Direct filter: default + this.filterMode = "default"; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "app.tree.filter.noTools")) { + // Toggle filter: no-tools ↔ default + this.filterMode = this.filterMode === "no-tools" ? "default" : "no-tools"; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "app.tree.filter.userOnly")) { + // Toggle filter: user-only ↔ default + this.filterMode = this.filterMode === "user-only" ? "default" : "user-only"; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "app.tree.filter.labeledOnly")) { + // Toggle filter: labeled-only ↔ default + this.filterMode = this.filterMode === "labeled-only" ? "default" : "labeled-only"; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "app.tree.filter.all")) { + // Toggle filter: all ↔ default + this.filterMode = this.filterMode === "all" ? "default" : "all"; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "app.tree.filter.cycleBackward")) { + // Cycle filter backwards + const modes: FilterMode[] = ["default", "no-tools", "user-only", "labeled-only", "all"]; + const currentIndex = modes.indexOf(this.filterMode); + this.filterMode = modes[(currentIndex - 1 + modes.length) % modes.length]; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "app.tree.filter.cycleForward")) { + // Cycle filter forwards: default → no-tools → user-only → labeled-only → all → default + const modes: FilterMode[] = ["default", "no-tools", "user-only", "labeled-only", "all"]; + const currentIndex = modes.indexOf(this.filterMode); + this.filterMode = modes[(currentIndex + 1) % modes.length]; + this.foldedNodes.clear(); + this.applyFilter(); + } else if (kb.matches(keyData, "tui.editor.deleteCharBackward")) { + if (this.searchQuery.length > 0) { + this.searchQuery = this.searchQuery.slice(0, -1); + this.foldedNodes.clear(); + this.applyFilter(); + } + } else if (kb.matches(keyData, "app.tree.editLabel")) { + const selected = this.filteredNodes[this.selectedIndex]; + if (selected && this.onLabelEdit) { + this.onLabelEdit(selected.node.entry.id, selected.node.label); + } + } else if (kb.matches(keyData, "app.tree.toggleLabelTimestamp")) { + this.showLabelTimestamps = !this.showLabelTimestamps; + } else { + const hasControlChars = [...keyData].some((ch) => { + const code = ch.charCodeAt(0); + return code < 32 || code === 0x7f || (code >= 0x80 && code <= 0x9f); + }); + if (!hasControlChars && keyData.length > 0) { + this.searchQuery += keyData; + this.foldedNodes.clear(); + this.applyFilter(); + } + } + } + + /** + * Whether a node can be folded. A node is foldable if it has visible children + * and is either a root (no visible parent) or a segment start (visible parent + * has multiple visible children). + */ + private isFoldable(entryId: string): boolean { + const children = this.visibleChildrenMap.get(entryId); + if (!children || children.length === 0) return false; + const parentId = this.visibleParentMap.get(entryId); + if (parentId === null || parentId === undefined) return true; + const siblings = this.visibleChildrenMap.get(parentId); + return siblings !== undefined && siblings.length > 1; + } + + /** + * Find the index of the next branch segment start in the given direction. + * A segment start is the first child of a branch point. + * + * "up" walks the visible parent chain; "down" walks visible children + * (always following the first child). + */ + private findBranchSegmentStart(direction: "up" | "down"): number { + const selectedId = this.filteredNodes[this.selectedIndex]?.node.entry.id; + if (!selectedId) return this.selectedIndex; + + const indexByEntryId = new Map(this.filteredNodes.map((node, i) => [node.node.entry.id, i])); + let currentId: string = selectedId; + if (direction === "down") { + while (true) { + const children: string[] = this.visibleChildrenMap.get(currentId) ?? []; + if (children.length === 0) return indexByEntryId.get(currentId)!; + if (children.length > 1) return indexByEntryId.get(children[0])!; + currentId = children[0]; + } + } + + // direction === "up" + while (true) { + const parentId: string | null = this.visibleParentMap.get(currentId) ?? null; + if (parentId === null) return indexByEntryId.get(currentId)!; + const children = this.visibleChildrenMap.get(parentId) ?? []; + if (children.length > 1) { + const segmentStart = indexByEntryId.get(currentId)!; + if (segmentStart < this.selectedIndex) { + return segmentStart; + } + } + currentId = parentId; + } + } +} + +/** Component that displays the current search query */ +class SearchLine implements Component { + private treeList: TreeList; + + constructor(treeList: TreeList) { + this.treeList = treeList; + } + + invalidate(): void {} + + render(width: number): string[] { + const query = this.treeList.getSearchQuery(); + if (query) { + return [truncateToWidth(` ${theme.fg("muted", "Type to search:")} ${theme.fg("accent", query)}`, width)]; + } + return [truncateToWidth(` ${theme.fg("muted", "Type to search:")}`, width)]; + } + + handleInput(_keyData: string): void {} +} + +/** Component that renders tree help as semantic rows with chunk-aware wrapping */ +class TreeHelp implements Component { + invalidate(): void {} + + render(width: number): string[] { + const items = TREE_HELP_ITEMS.map(({ keys, label, labelFirst }) => { + const text = formatHelpKeys(keys); + if (!text) return label; + return labelFirst ? `${label} ${text}` : `${text} ${label}`; + }); + + const availableWidth = Math.max(1, width); + const indent = " "; + const separator = " · "; + const lines: string[] = []; + let currentLine = ""; + + for (const item of items) { + const candidate = currentLine + ? `${currentLine}${separator}${item}` + : visibleWidth(`${indent}${item}`) <= availableWidth + ? `${indent}${item}` + : item; + if (!currentLine || visibleWidth(candidate) <= availableWidth) { + currentLine = candidate; + continue; + } + + lines.push(...wrapTextWithAnsi(currentLine.trimEnd(), availableWidth)); + currentLine = visibleWidth(`${indent}${item}`) <= availableWidth ? `${indent}${item}` : item; + } + + if (currentLine) { + lines.push(...wrapTextWithAnsi(currentLine.trimEnd(), availableWidth)); + } + + return lines.map((line) => theme.fg("muted", line)); + } +} + +const TREE_HELP_ITEMS: Array<{ keys: Keybinding[]; label: string; labelFirst?: boolean }> = [ + { keys: ["tui.select.up", "tui.select.down"], label: "move" }, + { keys: ["tui.editor.cursorLeft", "tui.editor.cursorRight"], label: "page" }, + { keys: ["app.tree.foldOrUp", "app.tree.unfoldOrDown"], label: "branch" }, + { keys: ["app.tree.editLabel"], label: "label" }, + { keys: ["app.tree.toggleLabelTimestamp"], label: "label time" }, + { + keys: [ + "app.tree.filter.default", + "app.tree.filter.noTools", + "app.tree.filter.userOnly", + "app.tree.filter.labeledOnly", + "app.tree.filter.all", + ], + label: "filters", + labelFirst: true, + }, + { keys: ["app.tree.filter.cycleForward", "app.tree.filter.cycleBackward"], label: "cycle", labelFirst: true }, +]; + +function formatHelpKeys(keybindings: Keybinding[]): string { + const keys: string[] = []; + for (const keybinding of keybindings) { + const key = getKeybindings().getKeys(keybinding)[0]; + if (key !== undefined) keys.push(key); + } + if (keys.length === 0) return ""; + + return formatKeyText(compactRawKeys(keys)) + .replace(/\bpageUp\b/g, "pgup") + .replace(/\bpageDown\b/g, "pgdn") + .replace(/\bup\b/g, "↑") + .replace(/\bdown\b/g, "↓") + .replace(/\bleft\b/g, "←") + .replace(/\bright\b/g, "→"); +} + +function compactRawKeys(keys: string[]): string { + if (keys.length === 1) return keys[0]!; + + const parts = keys.map((key) => { + const separatorIndex = key.lastIndexOf("+"); + return separatorIndex === -1 + ? { prefix: "", suffix: key } + : { prefix: key.slice(0, separatorIndex + 1), suffix: key.slice(separatorIndex + 1) }; + }); + const prefix = parts[0]!.prefix; + return prefix && parts.every((part) => part.prefix === prefix) + ? `${prefix}${parts.map((part) => part.suffix).join("/")}` + : keys.join("/"); +} + +/** Label input component shown when editing a label */ +class LabelInput implements Component, Focusable { + private input: Input; + private entryId: string; + public onSubmit?: (entryId: string, label: string | undefined) => void; + public onCancel?: () => void; + + // Focusable implementation - propagate to input for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + this.input.focused = value; + } + + constructor(entryId: string, currentLabel: string | undefined) { + this.entryId = entryId; + this.input = new Input(); + if (currentLabel) { + this.input.setValue(currentLabel); + } + } + + invalidate(): void {} + + render(width: number): string[] { + const lines: string[] = []; + const indent = " "; + const availableWidth = width - indent.length; + lines.push(truncateToWidth(`${indent}${theme.fg("muted", "Label (empty to remove):")}`, width)); + lines.push(...this.input.render(availableWidth).map((line) => truncateToWidth(`${indent}${line}`, width))); + lines.push( + truncateToWidth( + `${indent}${keyHint("tui.select.confirm", "save")} ${keyHint("tui.select.cancel", "cancel")}`, + width, + ), + ); + return lines; + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "tui.select.confirm")) { + const value = this.input.getValue().trim(); + this.onSubmit?.(this.entryId, value || undefined); + } else if (kb.matches(keyData, "tui.select.cancel")) { + this.onCancel?.(); + } else { + this.input.handleInput(keyData); + } + } +} + +/** + * Component that renders a session tree selector for navigation + */ +export class TreeSelectorComponent extends Container implements Focusable { + private treeList: TreeList; + private labelInput: LabelInput | null = null; + private labelInputContainer: Container; + private treeContainer: Container; + private onLabelChangeCallback?: (entryId: string, label: string | undefined) => void; + + // Focusable implementation - propagate to labelInput when active for IME cursor positioning + private _focused = false; + get focused(): boolean { + return this._focused; + } + set focused(value: boolean) { + this._focused = value; + // Propagate to labelInput when it's active + if (this.labelInput) { + this.labelInput.focused = value; + } + } + + constructor( + tree: SessionTreeNode[], + currentLeafId: string | null, + terminalHeight: number, + onSelect: (entryId: string) => void, + onCancel: () => void, + onLabelChange?: (entryId: string, label: string | undefined) => void, + initialSelectedId?: string, + initialFilterMode?: FilterMode, + ) { + super(); + + this.onLabelChangeCallback = onLabelChange; + const maxVisibleLines = Math.max(5, Math.floor(terminalHeight / 2)); + + this.treeList = new TreeList(tree, currentLeafId, maxVisibleLines, initialSelectedId, initialFilterMode); + this.treeList.onSelect = onSelect; + this.treeList.onCancel = onCancel; + this.treeList.onLabelEdit = (entryId, currentLabel) => this.showLabelInput(entryId, currentLabel); + + this.treeContainer = new Container(); + this.treeContainer.addChild(this.treeList); + + this.labelInputContainer = new Container(); + + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + this.addChild(new Text(theme.bold(" Session Tree"), 1, 0)); + this.addChild(new TreeHelp()); + this.addChild(new SearchLine(this.treeList)); + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + this.addChild(this.treeContainer); + this.addChild(this.labelInputContainer); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + + if (tree.length === 0) { + setTimeout(() => onCancel(), 100); + } + } + + private showLabelInput(entryId: string, currentLabel: string | undefined): void { + this.labelInput = new LabelInput(entryId, currentLabel); + this.labelInput.onSubmit = (id, label) => { + this.treeList.updateNodeLabel(id, label); + this.onLabelChangeCallback?.(id, label); + this.hideLabelInput(); + }; + this.labelInput.onCancel = () => this.hideLabelInput(); + + // Propagate current focused state to the new labelInput + this.labelInput.focused = this._focused; + + this.treeContainer.clear(); + this.labelInputContainer.clear(); + this.labelInputContainer.addChild(this.labelInput); + } + + private hideLabelInput(): void { + this.labelInput = null; + this.labelInputContainer.clear(); + this.treeContainer.clear(); + this.treeContainer.addChild(this.treeList); + } + + handleInput(keyData: string): void { + if (this.labelInput) { + this.labelInput.handleInput(keyData); + } else { + this.treeList.handleInput(keyData); + } + } + + getTreeList(): TreeList { + return this.treeList; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/trust-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/trust-selector.ts new file mode 100644 index 000000000..92c232889 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/trust-selector.ts @@ -0,0 +1,134 @@ +import { Container, getKeybindings, Spacer, Text } from "@earendil-works/pi-tui"; +import { + getProjectTrustOptions, + type ProjectTrustOption, + type ProjectTrustStoreEntry, +} from "../../../core/trust-manager.ts"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; +import { keyHint, rawKeyHint } from "./keybinding-hints.ts"; + +export type TrustSelection = Pick; + +export interface TrustSelectorOptions { + cwd: string; + savedDecision: ProjectTrustStoreEntry | null; + projectTrusted: boolean; + onSelect: (selection: TrustSelection) => void; + onCancel: () => void; +} + +function formatDecision(trustPath: string | undefined, decision: ProjectTrustStoreEntry | null): string { + if (decision === null) { + return "none"; + } + const label = decision.decision ? "trusted" : "untrusted"; + if (trustPath !== undefined && decision.path !== trustPath) { + return `${label} (inherited from ${decision.path})`; + } + return `${label} (${decision.path})`; +} + +export class TrustSelectorComponent extends Container { + private selectedIndex: number; + private readonly listContainer: Container; + private readonly trustOptions: ProjectTrustOption[]; + private readonly savedDecision: ProjectTrustStoreEntry | null; + private readonly onSelectCallback: (selection: TrustSelection) => void; + private readonly onCancelCallback: () => void; + + constructor(options: TrustSelectorOptions) { + super(); + + this.savedDecision = options.savedDecision; + this.trustOptions = getProjectTrustOptions(options.cwd); + this.selectedIndex = Math.max( + 0, + this.trustOptions.findIndex((option) => this.isSavedOption(option)), + ); + this.onSelectCallback = options.onSelect; + this.onCancelCallback = options.onCancel; + + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("accent", theme.bold("Project trust")), 1, 0)); + this.addChild(new Text(theme.fg("muted", options.cwd), 1, 0)); + this.addChild(new Spacer(1)); + this.addChild( + new Text( + theme.fg( + "muted", + `Saved decision: ${formatDecision(this.trustOptions[0]?.savedPath, options.savedDecision)}`, + ), + 1, + 0, + ), + ); + this.addChild( + new Text(theme.fg("muted", `Current session: ${options.projectTrusted ? "trusted" : "untrusted"}`), 1, 0), + ); + this.addChild(new Spacer(1)); + + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild( + new Text( + rawKeyHint("↑↓", "navigate") + + " " + + keyHint("tui.select.confirm", "save") + + " " + + keyHint("tui.select.cancel", "cancel"), + 1, + 0, + ), + ); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + + this.updateList(); + } + + private isSavedOption(option: ProjectTrustOption): boolean { + return ( + option.savedPath !== undefined && + this.savedDecision?.decision === option.trusted && + this.savedDecision.path === option.savedPath + ); + } + + private updateList(): void { + this.listContainer.clear(); + for (let i = 0; i < this.trustOptions.length; i++) { + const option = this.trustOptions[i]; + if (!option) { + continue; + } + + const isSelected = i === this.selectedIndex; + const isCurrent = this.isSavedOption(option); + const checkmark = isCurrent ? theme.fg("success", " ✓") : ""; + const prefix = isSelected ? theme.fg("accent", "→ ") : " "; + const label = isSelected ? theme.fg("accent", option.label) : theme.fg("text", option.label); + this.listContainer.addChild(new Text(`${prefix}${label}${checkmark}`, 1, 0)); + } + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + if (kb.matches(keyData, "tui.select.up") || keyData === "k") { + this.selectedIndex = Math.max(0, this.selectedIndex - 1); + this.updateList(); + } else if (kb.matches(keyData, "tui.select.down") || keyData === "j") { + this.selectedIndex = Math.min(this.trustOptions.length - 1, this.selectedIndex + 1); + this.updateList(); + } else if (kb.matches(keyData, "tui.select.confirm") || keyData === "\n") { + const selected = this.trustOptions[this.selectedIndex]; + if (selected) { + this.onSelectCallback({ trusted: selected.trusted, updates: selected.updates }); + } + } else if (kb.matches(keyData, "tui.select.cancel")) { + this.onCancelCallback(); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts new file mode 100644 index 000000000..69fbbc82a --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts @@ -0,0 +1,155 @@ +import { type Component, Container, getKeybindings, Spacer, Text, truncateToWidth } from "@earendil-works/pi-tui"; +import { theme } from "../theme/theme.ts"; +import { DynamicBorder } from "./dynamic-border.ts"; + +interface UserMessageItem { + id: string; // Entry ID in the session + text: string; // The message text + timestamp?: string; // Optional timestamp if available +} + +/** + * Custom user message list component with selection + */ +class UserMessageList implements Component { + private messages: UserMessageItem[] = []; + private selectedIndex: number = 0; + public onSelect?: (entryId: string) => void; + public onCancel?: () => void; + private maxVisible: number = 10; // Max messages visible + + constructor(messages: UserMessageItem[], initialSelectedId?: string) { + // Store messages in chronological order (oldest to newest) + this.messages = messages; + const initialIndex = initialSelectedId ? messages.findIndex((message) => message.id === initialSelectedId) : -1; + // Start with selected message if provided, else default to the most recent + this.selectedIndex = initialIndex >= 0 ? initialIndex : Math.max(0, messages.length - 1); + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(width: number): string[] { + const lines: string[] = []; + + if (this.messages.length === 0) { + lines.push(theme.fg("muted", " No user messages found")); + return lines; + } + + // Calculate visible range with scrolling + const startIndex = Math.max( + 0, + Math.min(this.selectedIndex - Math.floor(this.maxVisible / 2), this.messages.length - this.maxVisible), + ); + const endIndex = Math.min(startIndex + this.maxVisible, this.messages.length); + + // Render visible messages (2 lines per message + blank line) + for (let i = startIndex; i < endIndex; i++) { + const message = this.messages[i]; + const isSelected = i === this.selectedIndex; + + // Normalize message to single line + const normalizedMessage = message.text.replace(/\n/g, " ").trim(); + + // First line: cursor + message + const cursor = isSelected ? theme.fg("accent", "› ") : " "; + const maxMsgWidth = width - 2; // Account for cursor (2 chars) + const truncatedMsg = truncateToWidth(normalizedMessage, maxMsgWidth); + const messageLine = cursor + (isSelected ? theme.bold(truncatedMsg) : truncatedMsg); + + lines.push(messageLine); + + // Second line: metadata (position in history) + const position = i + 1; + const metadata = ` Message ${position} of ${this.messages.length}`; + const metadataLine = theme.fg("muted", metadata); + lines.push(metadataLine); + lines.push(""); // Blank line between messages + } + + // Add scroll indicator if needed + if (startIndex > 0 || endIndex < this.messages.length) { + const scrollInfo = theme.fg("muted", ` (${this.selectedIndex + 1}/${this.messages.length})`); + lines.push(scrollInfo); + } + + return lines; + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + // Up arrow - go to previous (older) message, wrap to bottom when at top + if (kb.matches(keyData, "tui.select.up")) { + this.selectedIndex = this.selectedIndex === 0 ? this.messages.length - 1 : this.selectedIndex - 1; + } + // Down arrow - go to next (newer) message, wrap to top when at bottom + else if (kb.matches(keyData, "tui.select.down")) { + this.selectedIndex = this.selectedIndex === this.messages.length - 1 ? 0 : this.selectedIndex + 1; + } + // Enter - select message and branch + else if (kb.matches(keyData, "tui.select.confirm")) { + const selected = this.messages[this.selectedIndex]; + if (selected && this.onSelect) { + this.onSelect(selected.id); + } + } + // Escape - cancel + else if (kb.matches(keyData, "tui.select.cancel")) { + if (this.onCancel) { + this.onCancel(); + } + } + } +} + +/** + * Component that renders a user message selector for branching + */ +export class UserMessageSelectorComponent extends Container { + private messageList: UserMessageList; + + constructor( + messages: UserMessageItem[], + onSelect: (entryId: string) => void, + onCancel: () => void, + initialSelectedId?: string, + ) { + super(); + + // Add header + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.bold("Fork from Message"), 1, 0)); + this.addChild( + new Text( + theme.fg("muted", "Select a user message to copy the active path up to that point into a new session"), + 1, + 0, + ), + ); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + this.addChild(new Spacer(1)); + + // Create message list + this.messageList = new UserMessageList(messages, initialSelectedId); + this.messageList.onSelect = onSelect; + this.messageList.onCancel = onCancel; + + this.addChild(this.messageList); + + // Add bottom border + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder()); + + // Auto-cancel if no messages + if (messages.length === 0) { + setTimeout(() => onCancel(), 100); + } + } + + getMessageList(): UserMessageList { + return this.messageList; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/user-message.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/user-message.ts new file mode 100644 index 000000000..45d655776 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/user-message.ts @@ -0,0 +1,42 @@ +import { Box, Container, Markdown, type MarkdownTheme } from "@earendil-works/pi-tui"; +import { getMarkdownTheme, theme } from "../theme/theme.ts"; + +const OSC133_ZONE_START = "\x1b]133;A\x07"; +const OSC133_ZONE_END = "\x1b]133;B\x07"; +const OSC133_ZONE_FINAL = "\x1b]133;C\x07"; + +/** + * Component that renders a user message + */ +export class UserMessageComponent extends Container { + private contentBox: Box; + + constructor(text: string, markdownTheme: MarkdownTheme = getMarkdownTheme()) { + super(); + this.contentBox = new Box(1, 1, (content: string) => theme.bg("userMessageBg", content)); + this.contentBox.addChild( + new Markdown( + text, + 0, + 0, + markdownTheme, + { + color: (content: string) => theme.fg("userMessageText", content), + }, + { preserveOrderedListMarkers: true }, + ), + ); + this.addChild(this.contentBox); + } + + override render(width: number): string[] { + const lines = super.render(width); + if (lines.length === 0) { + return lines; + } + + lines[0] = OSC133_ZONE_START + lines[0]; + lines[lines.length - 1] = OSC133_ZONE_END + OSC133_ZONE_FINAL + lines[lines.length - 1]; + return lines; + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/components/visual-truncate.ts b/cactus-code/packages/coding-agent/src/modes/interactive/components/visual-truncate.ts new file mode 100644 index 000000000..827230416 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/components/visual-truncate.ts @@ -0,0 +1,50 @@ +/** + * Shared utility for truncating text to visual lines (accounting for line wrapping). + * Used by both tool-execution.ts and bash-execution.ts for consistent behavior. + */ + +import { Text } from "@earendil-works/pi-tui"; + +export interface VisualTruncateResult { + /** The visual lines to display */ + visualLines: string[]; + /** Number of visual lines that were skipped (hidden) */ + skippedCount: number; +} + +/** + * Truncate text to a maximum number of visual lines (from the end). + * This accounts for line wrapping based on terminal width. + * + * @param text - The text content (may contain newlines) + * @param maxVisualLines - Maximum number of visual lines to show + * @param width - Terminal/render width + * @param paddingX - Horizontal padding for Text component (default 0). + * Use 0 when result will be placed in a Box (Box adds its own padding). + * Use 1 when result will be placed in a plain Container. + * @returns The truncated visual lines and count of skipped lines + */ +export function truncateToVisualLines( + text: string, + maxVisualLines: number, + width: number, + paddingX: number = 0, +): VisualTruncateResult { + if (!text) { + return { visualLines: [], skippedCount: 0 }; + } + + // Create a temporary Text component to render and get visual lines + const tempText = new Text(text, paddingX, 0); + const allVisualLines = tempText.render(width); + + if (allVisualLines.length <= maxVisualLines) { + return { visualLines: allVisualLines, skippedCount: 0 }; + } + + // Take the last N visual lines + const truncatedLines = allVisualLines.slice(-maxVisualLines); + const skippedCount = allVisualLines.length - maxVisualLines; + + return { visualLines: truncatedLines, skippedCount }; +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/cactus-code/packages/coding-agent/src/modes/interactive/interactive-mode.ts new file mode 100644 index 000000000..7df647d94 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -0,0 +1,4966 @@ +/** + * Interactive mode for the coding agent. + * Handles TUI rendering and user interaction, delegating business logic to AgentSession. + */ + +import * as crypto from "node:crypto"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { + type AssistantMessage, + getProviders, + type ImageContent, + type Message, + type Model, +} from "@earendil-works/pi-ai/compat"; +import type { + AutocompleteItem, + AutocompleteProvider, + EditorComponent, + Keybinding, + KeyId, + MarkdownTheme, + OverlayHandle, + OverlayOptions, + SlashCommand, +} from "@earendil-works/pi-tui"; +import { + CombinedAutocompleteProvider, + type Component, + Container, + fuzzyFilter, + Loader, + type LoaderIndicatorOptions, + Markdown, + matchesKey, + ProcessTerminal, + Spacer, + setKeybindings, + Text, + TruncatedText, + TUI, + visibleWidth, +} from "@earendil-works/pi-tui"; +import chalk from "chalk"; +import { spawn } from "child_process"; +import { + APP_NAME, + APP_TITLE, + CONFIG_DIR_NAME, + getDebugLogPath, + VERSION, +} from "../../config.ts"; +import { type AgentSession, type AgentSessionEvent, parseSkillBlock } from "../../core/agent-session.ts"; +import { type AgentSessionRuntime, SessionImportFileNotFoundError } from "../../core/agent-session-runtime.ts"; +import type { + AutocompleteProviderFactory, + EditorFactory, + ExtensionCommandContext, + ExtensionContext, + ExtensionRunner, + ExtensionUIContext, + ExtensionUIDialogOptions, + ExtensionWidgetOptions, + ProjectTrustContext, +} from "../../core/extensions/index.ts"; +import { FooterDataProvider, type ReadonlyFooterDataProvider } from "../../core/footer-data-provider.ts"; +import { configureHttpDispatcher, formatHttpIdleTimeoutMs } from "../../core/http-dispatcher.ts"; +import { type AppKeybinding, KeybindingsManager } from "../../core/keybindings.ts"; +import { createCompactionSummaryMessage } from "../../core/messages.ts"; +import { defaultModelPerProvider, findExactModelReferenceMatch, resolveModelScope } from "../../core/model-resolver.ts"; +import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "../../core/provider-display-names.ts"; +import type { ResourceDiagnostic } from "../../core/resource-loader.ts"; +import { formatMissingSessionCwdPrompt, MissingSessionCwdError } from "../../core/session-cwd.ts"; +import { type SessionContext, SessionManager } from "../../core/session-manager.ts"; +import { BUILTIN_SLASH_COMMANDS } from "../../core/slash-commands.ts"; +import type { SourceInfo } from "../../core/source-info.ts"; +import type { TruncationResult } from "../../core/tools/truncate.ts"; +import { hasTrustRequiringProjectResources, ProjectTrustStore } from "../../core/trust-manager.ts"; +import { copyToClipboard } from "../../utils/clipboard.ts"; +import { extensionForImageMimeType, readClipboardImage } from "../../utils/clipboard-image.ts"; +import { parseGitUrl } from "../../utils/git.ts"; +import { getCwdRelativePath } from "../../utils/paths.ts"; +import { killTrackedDetachedChildren } from "../../utils/shell.ts"; +import { ensureTool } from "../../utils/tools-manager.ts"; +import { AssistantMessageComponent } from "./components/assistant-message.ts"; +import { BashExecutionComponent } from "./components/bash-execution.ts"; +import { BranchSummaryMessageComponent } from "./components/branch-summary-message.ts"; +import { CompactionSummaryMessageComponent } from "./components/compaction-summary-message.ts"; +import { CountdownTimer } from "./components/countdown-timer.ts"; +import { CustomEditor } from "./components/custom-editor.ts"; +import { CustomMessageComponent } from "./components/custom-message.ts"; +import { DynamicBorder } from "./components/dynamic-border.ts"; +import { ExtensionEditorComponent } from "./components/extension-editor.ts"; +import { ExtensionInputComponent } from "./components/extension-input.ts"; +import { ExtensionSelectorComponent } from "./components/extension-selector.ts"; +import { FooterComponent } from "./components/footer.ts"; +import { formatKeyText, keyDisplayText, keyHint, keyText, rawKeyHint } from "./components/keybinding-hints.ts"; +import { ModelSelectorComponent } from "./components/model-selector.ts"; +import { ScopedModelsSelectorComponent } from "./components/scoped-models-selector.ts"; +import { SessionSelectorComponent } from "./components/session-selector.ts"; +import { SettingsSelectorComponent } from "./components/settings-selector.ts"; +import { SkillInvocationMessageComponent } from "./components/skill-invocation-message.ts"; +import { ToolExecutionComponent } from "./components/tool-execution.ts"; +import { TreeSelectorComponent } from "./components/tree-selector.ts"; +import { TrustSelectorComponent } from "./components/trust-selector.ts"; +import { UserMessageComponent } from "./components/user-message.ts"; +import { UserMessageSelectorComponent } from "./components/user-message-selector.ts"; +import { getModelSearchText } from "./model-search.ts"; +import { + getAvailableThemes, + getAvailableThemesWithPaths, + getEditorTheme, + getMarkdownTheme, + getThemeByName, + onThemeChange, + setRegisteredThemes, + stopThemeWatcher, + Theme, + type ThemeColor, + theme, +} from "./theme/theme.ts"; +import { InteractiveThemeController } from "./theme/theme-controller.ts"; + +/** Interface for components that can be expanded/collapsed */ +interface Expandable { + setExpanded(expanded: boolean): void; +} + +function isExpandable(obj: unknown): obj is Expandable { + return typeof obj === "object" && obj !== null && "setExpanded" in obj && typeof obj.setExpanded === "function"; +} + +class ExpandableText extends Text implements Expandable { + private readonly getCollapsedText: () => string; + private readonly getExpandedText: () => string; + + constructor( + getCollapsedText: () => string, + getExpandedText: () => string, + expanded = false, + paddingX = 0, + paddingY = 0, + ) { + super(expanded ? getExpandedText() : getCollapsedText(), paddingX, paddingY); + this.getCollapsedText = getCollapsedText; + this.getExpandedText = getExpandedText; + } + + setExpanded(expanded: boolean): void { + this.setText(expanded ? this.getExpandedText() : this.getCollapsedText()); + } +} + +type CompactionQueuedMessage = { + text: string; + mode: "steer" | "followUp"; +}; + +const DEAD_TERMINAL_ERROR_CODES = new Set(["EIO", "EPIPE", "ENOTCONN"]); + +function isDeadTerminalError(error: unknown): boolean { + if (!error || typeof error !== "object" || !("code" in error)) { + return false; + } + const code = (error as NodeJS.ErrnoException).code; + return code !== undefined && DEAD_TERMINAL_ERROR_CODES.has(code); +} + +const ANTHROPIC_SUBSCRIPTION_AUTH_WARNING = + "Anthropic subscription auth is active. Third-party harness usage draws from extra usage and is billed per token, not your Claude plan limits. Manage extra usage at https://claude.ai/settings/usage."; + +function isAnthropicSubscriptionAuthKey(apiKey: string | undefined): boolean { + return typeof apiKey === "string" && apiKey.startsWith("sk-ant-oat"); +} + +function isUnknownModel(model: Model | undefined): boolean { + return !!model && model.provider === "unknown" && model.id === "unknown" && model.api === "unknown"; +} + +function quoteIfNeeded(value: string): string { + if (value.length > 0 && !/[^a-zA-Z0-9_\-./~:@]/.test(value)) { + return value; + } + return `'${value.replace(/'/g, `'\\''`)}'`; +} + +export function formatResumeCommand(sessionManager: SessionManager): string | undefined { + if (!process.stdout.isTTY) return undefined; + if (!sessionManager.isPersisted()) return undefined; + + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile || !fs.existsSync(sessionFile)) return undefined; + + const args = [APP_NAME]; + if (!sessionManager.usesDefaultSessionDir()) { + args.push("--session-dir", quoteIfNeeded(sessionManager.getSessionDir())); + } + args.push("--session", sessionManager.getSessionId()); + return args.join(" "); +} + +function hasDefaultModelProvider(providerId: string): providerId is keyof typeof defaultModelPerProvider { + return providerId in defaultModelPerProvider; +} + +const BUILT_IN_MODEL_PROVIDERS = new Set(getProviders()); + +export function isApiKeyLoginProvider( + providerId: string, + oauthProviderIds: ReadonlySet, + builtInProviderIds: ReadonlySet = BUILT_IN_MODEL_PROVIDERS, +): boolean { + if (BUILT_IN_PROVIDER_DISPLAY_NAMES[providerId]) { + return true; + } + if (builtInProviderIds.has(providerId)) { + return false; + } + return !oauthProviderIds.has(providerId); +} + +/** + * Options for InteractiveMode initialization. + */ +export interface InteractiveModeOptions { + /** Providers that were migrated to auth.json (shows warning) */ + migratedProviders?: string[]; + /** Warning message if session model couldn't be restored */ + modelFallbackMessage?: string; + /** Cwd to trust after reload if it gained a .pi directory during this implicitly trusted session. */ + autoTrustOnReloadCwd?: string; + /** Initial message to send on startup (can include @file content) */ + initialMessage?: string; + /** Images to attach to the initial message */ + initialImages?: ImageContent[]; + /** Additional messages to send after the initial message */ + initialMessages?: string[]; + /** Force verbose startup (overrides quietStartup setting) */ + verbose?: boolean; +} + +export class InteractiveMode { + private runtimeHost: AgentSessionRuntime; + private ui: TUI; + private chatContainer: Container; + private pendingMessagesContainer: Container; + private statusContainer: Container; + private defaultEditor: CustomEditor; + private editor: EditorComponent; + private editorComponentFactory: EditorFactory | undefined; + private autocompleteProvider: AutocompleteProvider | undefined; + private autocompleteProviderWrappers: AutocompleteProviderFactory[] = []; + private fdPath: string | undefined; + private editorContainer: Container; + private footer: FooterComponent; + private footerDataProvider: FooterDataProvider; + // Stored so the same manager can be injected into custom editors, selectors, and extension UI. + private keybindings: KeybindingsManager; + private version: string; + private isInitialized = false; + private onInputCallback?: (text: string) => void; + private pendingUserInputs: string[] = []; + private loadingAnimation: Loader | undefined = undefined; + private workingMessage: string | undefined = undefined; + private workingVisible = true; + private workingIndicatorOptions: LoaderIndicatorOptions | undefined = undefined; + private readonly defaultWorkingMessage = "Working..."; + private readonly defaultHiddenThinkingLabel = "Thinking..."; + private hiddenThinkingLabel = this.defaultHiddenThinkingLabel; + + private lastSigintTime = 0; + private lastEscapeTime = 0; + private anthropicSubscriptionWarningShown = false; + + // Status line tracking (for mutating immediately-sequential status updates) + private lastStatusSpacer: Spacer | undefined = undefined; + private lastStatusText: Text | undefined = undefined; + + // Streaming message tracking + private streamingComponent: AssistantMessageComponent | undefined = undefined; + private streamingMessage: AssistantMessage | undefined = undefined; + + // Tool execution tracking: toolCallId -> component + private pendingTools = new Map(); + + // Tool output expansion state + private toolOutputExpanded = false; + + // Thinking block visibility state + private hideThinkingBlock = false; + + // Skill commands: command name -> skill file path + private skillCommands = new Map(); + + // Agent subscription unsubscribe function + private unsubscribe?: () => void; + private signalCleanupHandlers: Array<() => void> = []; + + // Track if editor is in bash mode (text starts with !) + private isBashMode = false; + + // Track current bash execution component + private bashComponent: BashExecutionComponent | undefined = undefined; + + // Track pending bash components (shown in pending area, moved to chat on submit) + private pendingBashComponents: BashExecutionComponent[] = []; + + // Auto-compaction state + private autoCompactionLoader: Loader | undefined = undefined; + private autoCompactionEscapeHandler?: () => void; + + // Auto-retry state + private retryLoader: Loader | undefined = undefined; + private retryCountdown: CountdownTimer | undefined = undefined; + private retryEscapeHandler?: () => void; + + // Messages queued while compaction is running + private compactionQueuedMessages: CompactionQueuedMessage[] = []; + + // Shutdown state + private shutdownRequested = false; + + // Extension UI state + private extensionSelector: ExtensionSelectorComponent | undefined = undefined; + private extensionInput: ExtensionInputComponent | undefined = undefined; + private extensionEditor: ExtensionEditorComponent | undefined = undefined; + private extensionTerminalInputUnsubscribers = new Set<() => void>(); + + // Extension widgets (components rendered above/below the editor) + private extensionWidgetsAbove = new Map(); + private extensionWidgetsBelow = new Map(); + private widgetContainerAbove!: Container; + private widgetContainerBelow!: Container; + + // Custom footer from extension (undefined = use built-in footer) + private customFooter: (Component & { dispose?(): void }) | undefined = undefined; + + // Header container that holds the built-in or custom header + private headerContainer: Container; + + // Built-in header (logo + keybinding hints) + private builtInHeader: Component | undefined = undefined; + + // Custom header from extension (undefined = use built-in header) + private customHeader: (Component & { dispose?(): void }) | undefined = undefined; + + private options: InteractiveModeOptions; + private autoTrustOnReloadCwd: string | undefined; + private themeController: InteractiveThemeController; + + // Convenience accessors + private get session(): AgentSession { + return this.runtimeHost.session; + } + private get agent() { + return this.session.agent; + } + private get sessionManager() { + return this.session.sessionManager; + } + private get settingsManager() { + return this.session.settingsManager; + } + + constructor(runtimeHost: AgentSessionRuntime, options: InteractiveModeOptions = {}) { + this.runtimeHost = runtimeHost; + this.options = options; + this.autoTrustOnReloadCwd = options.autoTrustOnReloadCwd; + this.runtimeHost.setBeforeSessionInvalidate(() => { + this.resetExtensionUI(); + }); + this.runtimeHost.setRebindSession(async () => { + await this.rebindCurrentSession({ renderBeforeBind: true }); + }); + this.version = VERSION; + this.ui = new TUI(new ProcessTerminal(), this.settingsManager.getShowHardwareCursor()); + this.ui.setClearOnShrink(this.settingsManager.getClearOnShrink()); + this.headerContainer = new Container(); + this.chatContainer = new Container(); + this.pendingMessagesContainer = new Container(); + this.statusContainer = new Container(); + this.widgetContainerAbove = new Container(); + this.widgetContainerBelow = new Container(); + this.keybindings = KeybindingsManager.create(); + setKeybindings(this.keybindings); + const editorPaddingX = this.settingsManager.getEditorPaddingX(); + const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible(); + this.defaultEditor = new CustomEditor(this.ui, getEditorTheme(), this.keybindings, { + paddingX: editorPaddingX, + autocompleteMaxVisible, + }); + this.editor = this.defaultEditor; + this.editorContainer = new Container(); + this.editorContainer.addChild(this.editor as Component); + this.footerDataProvider = new FooterDataProvider(this.sessionManager.getCwd()); + this.footer = new FooterComponent(this.session, this.footerDataProvider); + this.footer.setAutoCompactEnabled(this.session.autoCompactionEnabled); + + // Load hide thinking block setting + this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock(); + + // Register themes from resource loader and initialize + setRegisteredThemes(this.session.resourceLoader.getThemes().themes); + this.themeController = new InteractiveThemeController( + this.ui, + this.settingsManager, + (message) => this.showError(message), + () => this.updateEditorBorderColor(), + ); + } + + private getAutocompleteSourceTag(sourceInfo?: SourceInfo): string | undefined { + if (!sourceInfo) { + return undefined; + } + + const scopePrefix = sourceInfo.scope === "user" ? "u" : sourceInfo.scope === "project" ? "p" : "t"; + const source = sourceInfo.source.trim(); + + if (source === "auto" || source === "local" || source === "cli") { + return scopePrefix; + } + + if (source.startsWith("npm:")) { + return `${scopePrefix}:${source}`; + } + + const gitSource = parseGitUrl(source); + if (gitSource) { + const ref = gitSource.ref ? `@${gitSource.ref}` : ""; + return `${scopePrefix}:git:${gitSource.host}/${gitSource.path}${ref}`; + } + + return scopePrefix; + } + + private prefixAutocompleteDescription(description: string | undefined, sourceInfo?: SourceInfo): string | undefined { + const sourceTag = this.getAutocompleteSourceTag(sourceInfo); + if (!sourceTag) { + return description; + } + return description ? `[${sourceTag}] ${description}` : `[${sourceTag}]`; + } + + private getBuiltInCommandConflictDiagnostics(extensionRunner: ExtensionRunner): ResourceDiagnostic[] { + const builtinNames = new Set(BUILTIN_SLASH_COMMANDS.map((command) => command.name)); + return extensionRunner + .getRegisteredCommands() + .filter((command) => builtinNames.has(command.name)) + .map((command) => ({ + type: "warning" as const, + message: + command.invocationName === command.name + ? `Extension command '/${command.name}' conflicts with built-in interactive command. Skipping in autocomplete.` + : `Extension command '/${command.name}' conflicts with built-in interactive command. Available as '/${command.invocationName}'.`, + path: command.sourceInfo.path, + })); + } + + private createBaseAutocompleteProvider(): AutocompleteProvider { + // Define commands for autocomplete + const slashCommands: SlashCommand[] = BUILTIN_SLASH_COMMANDS.map((command) => ({ + name: command.name, + description: command.description, + })); + + const modelCommand = slashCommands.find((command) => command.name === "model"); + if (modelCommand) { + modelCommand.getArgumentCompletions = (prefix: string): AutocompleteItem[] | null => { + // Get available models (scoped or from registry) + const models = + this.session.scopedModels.length > 0 + ? this.session.scopedModels.map((s) => s.model) + : this.session.modelRegistry.getAvailable(); + + if (models.length === 0) return null; + + // Create items with provider/id format + const items = models.map((m) => ({ + id: m.id, + provider: m.provider, + name: m.name, + label: `${m.provider}/${m.id}`, + })); + + // Fuzzy filter by model ID + provider in either order. + const filtered = fuzzyFilter(items, prefix, getModelSearchText); + + if (filtered.length === 0) return null; + + return filtered.map((item) => ({ + value: item.label, + label: item.id, + description: item.provider, + })); + }; + } + + // Convert prompt templates to SlashCommand format for autocomplete + const templateCommands: SlashCommand[] = this.session.promptTemplates.map((cmd) => ({ + name: cmd.name, + description: this.prefixAutocompleteDescription(cmd.description, cmd.sourceInfo), + ...(cmd.argumentHint && { argumentHint: cmd.argumentHint }), + })); + + // Convert extension commands to SlashCommand format + const builtinCommandNames = new Set(slashCommands.map((c) => c.name)); + const extensionCommands: SlashCommand[] = this.session.extensionRunner + .getRegisteredCommands() + .filter((cmd) => !builtinCommandNames.has(cmd.name)) + .map((cmd) => ({ + name: cmd.invocationName, + description: this.prefixAutocompleteDescription(cmd.description, cmd.sourceInfo), + getArgumentCompletions: cmd.getArgumentCompletions, + })); + + // Build skill commands from session.skills (if enabled) + this.skillCommands.clear(); + const skillCommandList: SlashCommand[] = []; + if (this.settingsManager.getEnableSkillCommands()) { + for (const skill of this.session.resourceLoader.getSkills().skills) { + const commandName = `skill:${skill.name}`; + this.skillCommands.set(commandName, skill.filePath); + skillCommandList.push({ + name: commandName, + description: this.prefixAutocompleteDescription(skill.description, skill.sourceInfo), + }); + } + } + + return new CombinedAutocompleteProvider( + [...slashCommands, ...templateCommands, ...extensionCommands, ...skillCommandList], + this.sessionManager.getCwd(), + this.fdPath, + ); + } + + private setupAutocompleteProvider(): void { + let provider = this.createBaseAutocompleteProvider(); + const triggerCharacters: string[] = []; + for (const wrapProvider of this.autocompleteProviderWrappers) { + provider = wrapProvider(provider); + triggerCharacters.push(...(provider.triggerCharacters ?? [])); + } + if (triggerCharacters.length > 0) { + provider.triggerCharacters = [...new Set(triggerCharacters)]; + } + + this.autocompleteProvider = provider; + this.defaultEditor.setAutocompleteProvider(provider); + if (this.editor !== this.defaultEditor) { + this.editor.setAutocompleteProvider?.(provider); + } + } + + async init(): Promise { + if (this.isInitialized) return; + + this.registerSignalHandlers(); + + // Ensure fd and rg are available (downloads if missing, adds to PATH via getBinDir) + // Both are needed: fd for autocomplete, rg for grep tool and bash commands + const [fdPath] = await Promise.all([ensureTool("fd"), ensureTool("rg")]); + this.fdPath = fdPath; + + if (this.session.scopedModels.length > 0 && (this.options.verbose || !this.settingsManager.getQuietStartup())) { + const modelList = this.session.scopedModels + .map((sm) => { + const thinkingStr = sm.thinkingLevel ? `:${sm.thinkingLevel}` : ""; + return `${sm.model.id}${thinkingStr}`; + }) + .join(", "); + const cycleKeys = this.keybindings.getKeys("app.model.cycleForward"); + const cycleHint = + cycleKeys.length > 0 + ? theme.fg("muted", ` (${formatKeyText(cycleKeys.join("/"), { capitalize: true })} to cycle)`) + : ""; + console.log(theme.fg("dim", `Model scope: ${modelList}${cycleHint}`)); + } + + // Add header container as first child. Populate it after detectThemeIfUnset. + this.ui.addChild(this.headerContainer); + + this.ui.addChild(this.chatContainer); + this.ui.addChild(this.pendingMessagesContainer); + this.ui.addChild(this.statusContainer); + this.renderWidgets(); // Initialize with default spacer + this.ui.addChild(this.widgetContainerAbove); + this.ui.addChild(this.editorContainer); + this.ui.addChild(this.widgetContainerBelow); + this.ui.addChild(this.footer); + this.ui.setFocus(this.editor); + + this.setupKeyHandlers(); + this.setupEditorSubmitHandler(); + + // Start the UI before initializing extensions so session_start handlers can use interactive dialogs + this.ui.start(); + this.isInitialized = true; + + await this.themeController.applyFromSettings(); + + // Add header with keybindings from config (unless silenced) + if (this.options.verbose || !this.settingsManager.getQuietStartup()) { + const logo = theme.bold(theme.fg("accent", APP_NAME)) + theme.fg("dim", ` v${this.version}`); + + // Build startup instructions using keybinding hint helpers + const hint = (keybinding: AppKeybinding, description: string) => keyHint(keybinding, description); + + const expandedInstructions = [ + hint("app.interrupt", "to interrupt"), + hint("app.clear", "to clear"), + rawKeyHint(`${keyText("app.clear")} twice`, "to exit"), + hint("app.exit", "to exit (empty)"), + hint("app.suspend", "to suspend"), + keyHint("tui.editor.deleteToLineEnd", "to delete to end"), + hint("app.thinking.cycle", "to cycle thinking level"), + rawKeyHint(`${keyText("app.model.cycleForward")}/${keyText("app.model.cycleBackward")}`, "to cycle models"), + hint("app.model.select", "to select model"), + hint("app.tools.expand", "to expand tools"), + hint("app.thinking.toggle", "to expand thinking"), + hint("app.editor.external", "for external editor"), + rawKeyHint("/", "for commands"), + rawKeyHint("!", "to run bash"), + rawKeyHint("!!", "to run bash (no context)"), + hint("app.message.followUp", "to queue follow-up"), + hint("app.message.dequeue", "to edit all queued messages"), + hint("app.clipboard.pasteImage", "to paste image"), + rawKeyHint("drop files", "to attach"), + ].join("\n"); + const compactInstructions = [ + hint("app.interrupt", "interrupt"), + rawKeyHint(`${keyText("app.clear")}/${keyText("app.exit")}`, "clear/exit"), + rawKeyHint("/", "commands"), + rawKeyHint("!", "bash"), + hint("app.tools.expand", "more"), + ].join(theme.fg("muted", " · ")); + const compactOnboarding = theme.fg( + "dim", + `Press ${keyText("app.tools.expand")} to show full startup help and loaded resources.`, + ); + const onboarding = theme.fg( + "dim", + `Cactus reads, edits, and runs code in your project. Describe a task to get started.`, + ); + this.builtInHeader = new ExpandableText( + () => `${logo}\n${compactInstructions}\n${compactOnboarding}\n\n${onboarding}`, + () => `${logo}\n${expandedInstructions}\n\n${onboarding}`, + this.getStartupExpansionState(), + 1, + 0, + ); + + // Setup UI layout + this.headerContainer.addChild(new Spacer(1)); + this.headerContainer.addChild(this.builtInHeader); + this.headerContainer.addChild(new Spacer(1)); + } else { + // Minimal header when silenced + this.builtInHeader = new Text("", 0, 0); + this.headerContainer.addChild(this.builtInHeader); + } + this.ui.requestRender(); + + // Initialize extensions first so resources are shown before messages + await this.rebindCurrentSession(); + + // Render initial messages AFTER showing loaded resources + this.renderInitialMessages(); + + // Set up theme file watcher + onThemeChange(() => { + this.ui.invalidate(); + this.updateEditorBorderColor(); + this.ui.requestRender(); + }); + + // Set up git branch watcher (uses provider instead of footer) + this.footerDataProvider.onBranchChange(() => { + this.ui.requestRender(); + }); + + // Initialize available provider count for footer display + await this.updateAvailableProviderCount(); + } + + /** + * Update terminal title with session name and cwd. + */ + private updateTerminalTitle(): void { + const cwdBasename = path.basename(this.sessionManager.getCwd()); + const sessionName = this.sessionManager.getSessionName(); + if (sessionName) { + this.ui.terminal.setTitle(`${APP_TITLE} - ${sessionName} - ${cwdBasename}`); + } else { + this.ui.terminal.setTitle(`${APP_TITLE} - ${cwdBasename}`); + } + } + + /** + * Run the interactive mode. This is the main entry point. + * Initializes the UI, shows warnings, processes initial messages, and starts the interactive loop. + */ + async run(): Promise { + await this.init(); + + // Check tmux keyboard setup asynchronously + this.checkTmuxKeyboardSetup().then((warning) => { + if (warning) { + this.showWarning(warning); + } + }); + + // Show startup warnings + const { migratedProviders, modelFallbackMessage, initialMessage, initialImages, initialMessages } = this.options; + + if (migratedProviders && migratedProviders.length > 0) { + this.showWarning(`Migrated credentials to auth.json: ${migratedProviders.join(", ")}`); + } + + const modelsJsonError = this.session.modelRegistry.getError(); + if (modelsJsonError) { + this.showError(`models.json error: ${modelsJsonError}`); + } + + if (modelFallbackMessage) { + this.showWarning(modelFallbackMessage); + } + + void this.maybeWarnAboutAnthropicSubscriptionAuth(); + + // Process initial messages + if (initialMessage) { + try { + await this.session.prompt(initialMessage, { images: initialImages }); + } catch (error: unknown) { + const errorMessage = error instanceof Error ? error.message : "Unknown error occurred"; + this.showError(errorMessage); + } + } + + if (initialMessages) { + for (const message of initialMessages) { + try { + await this.session.prompt(message); + } catch (error: unknown) { + const errorMessage = error instanceof Error ? error.message : "Unknown error occurred"; + this.showError(errorMessage); + } + } + } + + // Main interactive loop + while (true) { + const userInput = await this.getUserInput(); + try { + await this.session.prompt(userInput); + } catch (error: unknown) { + const errorMessage = error instanceof Error ? error.message : "Unknown error occurred"; + this.showError(errorMessage); + } + } + } + + private async checkTmuxKeyboardSetup(): Promise { + if (!process.env.TMUX) return undefined; + + const runTmuxShow = (option: string): Promise => { + return new Promise((resolve) => { + const proc = spawn("tmux", ["show", "-gv", option], { + stdio: ["ignore", "pipe", "ignore"], + }); + let stdout = ""; + const timer = setTimeout(() => { + proc.kill(); + resolve(undefined); + }, 2000); + + proc.stdout?.on("data", (data) => { + stdout += data.toString(); + }); + proc.on("error", () => { + clearTimeout(timer); + resolve(undefined); + }); + proc.on("close", (code) => { + clearTimeout(timer); + resolve(code === 0 ? stdout.trim() : undefined); + }); + }); + }; + + const [extendedKeys, extendedKeysFormat] = await Promise.all([ + runTmuxShow("extended-keys"), + runTmuxShow("extended-keys-format"), + ]); + + // If we couldn't query tmux (timeout, sandbox, etc.), don't warn + if (extendedKeys === undefined) return undefined; + + if (extendedKeys !== "on" && extendedKeys !== "always") { + return "tmux extended-keys is off. Modified Enter keys may not work. Add `set -g extended-keys on` to ~/.tmux.conf and restart tmux."; + } + + if (extendedKeysFormat === "xterm") { + return "tmux extended-keys-format is xterm. Cactus works best with csi-u. Add `set -g extended-keys-format csi-u` to ~/.tmux.conf and restart tmux."; + } + + return undefined; + } + + private getMarkdownThemeWithSettings(): MarkdownTheme { + return { + ...getMarkdownTheme(), + codeBlockIndent: this.settingsManager.getCodeBlockIndent(), + }; + } + + // ========================================================================= + // Extension System + // ========================================================================= + + private formatDisplayPath(p: string): string { + const home = os.homedir(); + let result = p; + + // Replace home directory with ~ + if (result.startsWith(home)) { + result = `~${result.slice(home.length)}`; + } + + return result; + } + + private formatExtensionDisplayPath(path: string): string { + let result = this.formatDisplayPath(path); + result = result.replace(/\/index\.ts$/, "").replace(/\/index\.js$/, ""); + return result; + } + + private formatContextPath(p: string): string { + const cwd = path.resolve(this.sessionManager.getCwd()); + const absolutePath = path.isAbsolute(p) ? path.resolve(p) : path.resolve(cwd, p); + const relativePath = getCwdRelativePath(absolutePath, cwd); + if (relativePath !== undefined) { + return relativePath; + } + + return this.formatDisplayPath(absolutePath); + } + + private getStartupExpansionState(): boolean { + return this.options.verbose || this.toolOutputExpanded; + } + + /** + * Get a short path relative to the package root for display. + */ + private getShortPath(fullPath: string, sourceInfo?: SourceInfo): string { + const baseDir = sourceInfo?.baseDir; + if (baseDir && this.isPackageSource(sourceInfo)) { + const relativePath = path.relative(path.resolve(baseDir), path.resolve(fullPath)); + if ( + relativePath && + relativePath !== "." && + !relativePath.startsWith("..") && + !relativePath.startsWith(`..${path.sep}`) && + !path.isAbsolute(relativePath) + ) { + return relativePath.replace(/\\/g, "/"); + } + } + + const source = sourceInfo?.source ?? ""; + const npmMatch = fullPath.match(/node_modules\/(@?[^/]+(?:\/[^/]+)?)\/(.*)/); + if (npmMatch && source.startsWith("npm:")) { + return npmMatch[2]; + } + + const gitMatch = fullPath.match(/git\/[^/]+\/[^/]+\/(.*)/); + if (gitMatch && source.startsWith("git:")) { + return gitMatch[1]; + } + + return this.formatDisplayPath(fullPath); + } + + private getCompactPathLabel(resourcePath: string, sourceInfo?: SourceInfo): string { + const shortPath = this.getShortPath(resourcePath, sourceInfo); + const normalizedPath = shortPath.replace(/\\/g, "/"); + const segments = normalizedPath.split("/").filter((segment) => segment.length > 0 && segment !== "~"); + if (segments.length > 0) { + return segments[segments.length - 1]!; + } + return shortPath; + } + + private getCompactPackageSourceLabel(sourceInfo?: SourceInfo): string { + const source = sourceInfo?.source ?? ""; + if (source.startsWith("npm:")) { + return source.slice("npm:".length) || source; + } + + const gitSource = parseGitUrl(source); + if (gitSource) { + return gitSource.path || source; + } + + return source; + } + + private getCompactExtensionLabel(resourcePath: string, sourceInfo?: SourceInfo): string { + if (!this.isPackageSource(sourceInfo)) { + return this.getCompactPathLabel(resourcePath, sourceInfo); + } + + const sourceLabel = this.getCompactPackageSourceLabel(sourceInfo); + if (!sourceLabel) { + return this.getCompactPathLabel(resourcePath, sourceInfo); + } + + const shortPath = this.getShortPath(resourcePath, sourceInfo).replace(/\\/g, "/"); + const packagePath = shortPath.startsWith("extensions/") ? shortPath.slice("extensions/".length) : shortPath; + const parsedPath = path.posix.parse(packagePath); + + if (parsedPath.name === "index") { + return !parsedPath.dir || parsedPath.dir === "." ? sourceLabel : `${sourceLabel}:${parsedPath.dir}`; + } + + return `${sourceLabel}:${packagePath}`; + } + + private getCompactDisplayPathSegments(resourcePath: string): string[] { + return this.formatDisplayPath(resourcePath) + .replace(/\\/g, "/") + .split("/") + .filter((segment) => segment.length > 0 && segment !== "~"); + } + + private getCompactNonPackageExtensionLabel( + resourcePath: string, + index: number, + allPaths: Array<{ path: string; segments: string[] }>, + ): string { + const segments = allPaths[index]?.segments; + if (!segments || segments.length === 0) { + return this.getCompactPathLabel(resourcePath); + } + + for (let segmentCount = 1; segmentCount <= segments.length; segmentCount += 1) { + const candidate = segments.slice(-segmentCount).join("/"); + const isUnique = allPaths.every((item, itemIndex) => { + if (itemIndex === index) { + return true; + } + return item.segments.slice(-segmentCount).join("/") !== candidate; + }); + + if (isUnique) { + return candidate; + } + } + + return segments.join("/"); + } + + private getCompactExtensionLabels(extensions: Array<{ path: string; sourceInfo?: SourceInfo }>): string[] { + const nonPackageExtensions = extensions + .map((extension) => { + const segments = this.getCompactDisplayPathSegments(extension.path); + const lastSegment = segments[segments.length - 1]; + if (segments.length > 1 && (lastSegment === "index.ts" || lastSegment === "index.js")) { + segments.pop(); + } + return { + path: extension.path, + sourceInfo: extension.sourceInfo, + segments, + }; + }) + .filter((extension) => !this.isPackageSource(extension.sourceInfo)); + + return extensions.map((extension) => { + if (this.isPackageSource(extension.sourceInfo)) { + return this.getCompactExtensionLabel(extension.path, extension.sourceInfo); + } + + const nonPackageIndex = nonPackageExtensions.findIndex((item) => item.path === extension.path); + if (nonPackageIndex === -1) { + return this.getCompactPathLabel(extension.path, extension.sourceInfo); + } + + return this.getCompactNonPackageExtensionLabel(extension.path, nonPackageIndex, nonPackageExtensions); + }); + } + + private getDisplaySourceInfo(sourceInfo?: SourceInfo): { + label: string; + scopeLabel?: string; + color: "accent" | "muted"; + } { + const source = sourceInfo?.source ?? "local"; + const scope = sourceInfo?.scope ?? "project"; + if (source === "local") { + if (scope === "user") { + return { label: "user", color: "muted" }; + } + if (scope === "project") { + return { label: "project", color: "muted" }; + } + if (scope === "temporary") { + return { label: "path", scopeLabel: "temp", color: "muted" }; + } + return { label: "path", color: "muted" }; + } + + if (source === "cli") { + return { label: "path", scopeLabel: scope === "temporary" ? "temp" : undefined, color: "muted" }; + } + + const scopeLabel = + scope === "user" ? "user" : scope === "project" ? "project" : scope === "temporary" ? "temp" : undefined; + return { label: source, scopeLabel, color: "accent" }; + } + + private getScopeGroup(sourceInfo?: SourceInfo): "user" | "project" | "path" { + const source = sourceInfo?.source ?? "local"; + const scope = sourceInfo?.scope ?? "project"; + if (source === "cli" || scope === "temporary") return "path"; + if (scope === "user") return "user"; + if (scope === "project") return "project"; + return "path"; + } + + private isPackageSource(sourceInfo?: SourceInfo): boolean { + const source = sourceInfo?.source ?? ""; + return source.startsWith("npm:") || source.startsWith("git:"); + } + + private buildScopeGroups(items: Array<{ path: string; sourceInfo?: SourceInfo }>): Array<{ + scope: "user" | "project" | "path"; + paths: Array<{ path: string; sourceInfo?: SourceInfo }>; + packages: Map>; + }> { + const groups: Record< + "user" | "project" | "path", + { + scope: "user" | "project" | "path"; + paths: Array<{ path: string; sourceInfo?: SourceInfo }>; + packages: Map>; + } + > = { + user: { scope: "user", paths: [], packages: new Map() }, + project: { scope: "project", paths: [], packages: new Map() }, + path: { scope: "path", paths: [], packages: new Map() }, + }; + + for (const item of items) { + const groupKey = this.getScopeGroup(item.sourceInfo); + const group = groups[groupKey]; + const source = item.sourceInfo?.source ?? "local"; + + if (this.isPackageSource(item.sourceInfo)) { + const list = group.packages.get(source) ?? []; + list.push(item); + group.packages.set(source, list); + } else { + group.paths.push(item); + } + } + + return [groups.project, groups.user, groups.path].filter( + (group) => group.paths.length > 0 || group.packages.size > 0, + ); + } + + private formatScopeGroups( + groups: Array<{ + scope: "user" | "project" | "path"; + paths: Array<{ path: string; sourceInfo?: SourceInfo }>; + packages: Map>; + }>, + options: { + formatPath: (item: { path: string; sourceInfo?: SourceInfo }) => string; + formatPackagePath: (item: { path: string; sourceInfo?: SourceInfo }, source: string) => string; + }, + ): string { + const lines: string[] = []; + + for (const group of groups) { + lines.push(` ${theme.fg("accent", group.scope)}`); + + const sortedPaths = [...group.paths].sort((a, b) => a.path.localeCompare(b.path)); + for (const item of sortedPaths) { + lines.push(theme.fg("dim", ` ${options.formatPath(item)}`)); + } + + const sortedPackages = Array.from(group.packages.entries()).sort(([a], [b]) => a.localeCompare(b)); + for (const [source, items] of sortedPackages) { + lines.push(` ${theme.fg("mdLink", source)}`); + const sortedPackagePaths = [...items].sort((a, b) => a.path.localeCompare(b.path)); + for (const item of sortedPackagePaths) { + lines.push(theme.fg("dim", ` ${options.formatPackagePath(item, source)}`)); + } + } + } + + return lines.join("\n"); + } + + private findSourceInfoForPath(p: string, sourceInfos: Map): SourceInfo | undefined { + const exact = sourceInfos.get(p); + if (exact) return exact; + + let current = p; + while (current.includes("/")) { + current = current.substring(0, current.lastIndexOf("/")); + const parent = sourceInfos.get(current); + if (parent) return parent; + } + + return undefined; + } + + private formatPathWithSource(p: string, sourceInfo?: SourceInfo): string { + if (sourceInfo) { + const shortPath = this.getShortPath(p, sourceInfo); + const { label, scopeLabel } = this.getDisplaySourceInfo(sourceInfo); + const labelText = scopeLabel ? `${label} (${scopeLabel})` : label; + return `${labelText} ${shortPath}`; + } + return this.formatDisplayPath(p); + } + + private formatDiagnostics(diagnostics: readonly ResourceDiagnostic[], sourceInfos: Map): string { + const lines: string[] = []; + + // Group collision diagnostics by name + const collisions = new Map(); + const otherDiagnostics: ResourceDiagnostic[] = []; + + for (const d of diagnostics) { + if (d.type === "collision" && d.collision) { + const list = collisions.get(d.collision.name) ?? []; + list.push(d); + collisions.set(d.collision.name, list); + } else { + otherDiagnostics.push(d); + } + } + + // Format collision diagnostics grouped by name + for (const [name, collisionList] of collisions) { + const first = collisionList[0]?.collision; + if (!first) continue; + lines.push(theme.fg("warning", ` "${name}" collision:`)); + lines.push( + theme.fg( + "dim", + ` ${theme.fg("success", "✓")} ${this.formatPathWithSource(first.winnerPath, this.findSourceInfoForPath(first.winnerPath, sourceInfos))}`, + ), + ); + for (const d of collisionList) { + if (d.collision) { + lines.push( + theme.fg( + "dim", + ` ${theme.fg("warning", "✗")} ${this.formatPathWithSource(d.collision.loserPath, this.findSourceInfoForPath(d.collision.loserPath, sourceInfos))} (skipped)`, + ), + ); + } + } + } + + for (const d of otherDiagnostics) { + if (d.path) { + const formattedPath = this.formatPathWithSource(d.path, this.findSourceInfoForPath(d.path, sourceInfos)); + lines.push(theme.fg(d.type === "error" ? "error" : "warning", ` ${formattedPath}`)); + lines.push(theme.fg(d.type === "error" ? "error" : "warning", ` ${d.message}`)); + } else { + lines.push(theme.fg(d.type === "error" ? "error" : "warning", ` ${d.message}`)); + } + } + + return lines.join("\n"); + } + + private showLoadedResources(options?: { + extensions?: Array<{ path: string; sourceInfo?: SourceInfo }>; + force?: boolean; + showDiagnosticsWhenQuiet?: boolean; + }): void { + const showListing = options?.force || this.options.verbose || !this.settingsManager.getQuietStartup(); + const showDiagnostics = showListing || options?.showDiagnosticsWhenQuiet === true; + if (!showListing && !showDiagnostics) { + return; + } + + const sectionHeader = (name: string, color: ThemeColor = "mdHeading") => theme.fg(color, `[${name}]`); + const formatCompactList = (items: string[], options?: { sort?: boolean }): string => { + const labels = items.map((item) => item.trim()).filter((item) => item.length > 0); + if (options?.sort !== false) { + labels.sort((a, b) => a.localeCompare(b)); + } + return theme.fg("dim", ` ${labels.join(", ")}`); + }; + const addLoadedSection = ( + name: string, + collapsedBody: string, + expandedBody = collapsedBody, + color: ThemeColor = "mdHeading", + ): void => { + const section = new ExpandableText( + () => `${sectionHeader(name, color)}\n${collapsedBody}`, + () => `${sectionHeader(name, color)}\n${expandedBody}`, + this.getStartupExpansionState(), + 0, + 0, + ); + this.chatContainer.addChild(section); + this.chatContainer.addChild(new Spacer(1)); + }; + + const skillsResult = this.session.resourceLoader.getSkills(); + const promptsResult = this.session.resourceLoader.getPrompts(); + const themesResult = this.session.resourceLoader.getThemes(); + const extensions = + options?.extensions ?? + this.session.resourceLoader.getExtensions().extensions.map((extension) => ({ + path: extension.path, + sourceInfo: extension.sourceInfo, + })); + const sourceInfos = new Map(); + for (const extension of extensions) { + if (extension.sourceInfo) { + sourceInfos.set(extension.path, extension.sourceInfo); + } + } + for (const skill of skillsResult.skills) { + if (skill.sourceInfo) { + sourceInfos.set(skill.filePath, skill.sourceInfo); + } + } + for (const prompt of promptsResult.prompts) { + if (prompt.sourceInfo) { + sourceInfos.set(prompt.filePath, prompt.sourceInfo); + } + } + for (const loadedTheme of themesResult.themes) { + if (loadedTheme.sourcePath && loadedTheme.sourceInfo) { + sourceInfos.set(loadedTheme.sourcePath, loadedTheme.sourceInfo); + } + } + + if (showListing) { + const contextFiles = this.session.resourceLoader.getAgentsFiles().agentsFiles; + if (contextFiles.length > 0) { + this.chatContainer.addChild(new Spacer(1)); + const contextList = contextFiles + .map((f) => theme.fg("dim", ` ${this.formatDisplayPath(f.path)}`)) + .join("\n"); + const contextCompactList = formatCompactList( + contextFiles.map((contextFile) => this.formatContextPath(contextFile.path)), + { sort: false }, + ); + addLoadedSection("Context", contextCompactList, contextList); + } + + const skills = skillsResult.skills; + if (skills.length > 0) { + const groups = this.buildScopeGroups( + skills.map((skill) => ({ path: skill.filePath, sourceInfo: skill.sourceInfo })), + ); + const skillList = this.formatScopeGroups(groups, { + formatPath: (item) => this.formatDisplayPath(item.path), + formatPackagePath: (item) => this.getShortPath(item.path, item.sourceInfo), + }); + const skillCompactList = formatCompactList(skills.map((skill) => skill.name)); + addLoadedSection("Skills", skillCompactList, skillList); + } + + const templates = this.session.promptTemplates; + if (templates.length > 0) { + const groups = this.buildScopeGroups( + templates.map((template) => ({ path: template.filePath, sourceInfo: template.sourceInfo })), + ); + const templateByPath = new Map(templates.map((t) => [t.filePath, t])); + const templateList = this.formatScopeGroups(groups, { + formatPath: (item) => { + const template = templateByPath.get(item.path); + return template ? `/${template.name}` : this.formatDisplayPath(item.path); + }, + formatPackagePath: (item) => { + const template = templateByPath.get(item.path); + return template ? `/${template.name}` : this.formatDisplayPath(item.path); + }, + }); + const promptCompactList = formatCompactList(templates.map((template) => `/${template.name}`)); + addLoadedSection("Prompts", promptCompactList, templateList); + } + + if (extensions.length > 0) { + const groups = this.buildScopeGroups(extensions); + const extList = this.formatScopeGroups(groups, { + formatPath: (item) => this.formatExtensionDisplayPath(item.path), + formatPackagePath: (item) => + this.formatExtensionDisplayPath(this.getShortPath(item.path, item.sourceInfo)), + }); + const extensionCompactList = formatCompactList(this.getCompactExtensionLabels(extensions)); + addLoadedSection("Extensions", extensionCompactList, extList, "mdHeading"); + } + + // Show loaded themes (excluding built-in) + const loadedThemes = themesResult.themes; + const customThemes = loadedThemes.filter((t) => t.sourcePath); + if (customThemes.length > 0) { + const groups = this.buildScopeGroups( + customThemes.map((loadedTheme) => ({ + path: loadedTheme.sourcePath!, + sourceInfo: loadedTheme.sourceInfo, + })), + ); + const themeList = this.formatScopeGroups(groups, { + formatPath: (item) => this.formatDisplayPath(item.path), + formatPackagePath: (item) => this.getShortPath(item.path, item.sourceInfo), + }); + const themeCompactList = formatCompactList( + customThemes.map( + (loadedTheme) => + loadedTheme.name ?? this.getCompactPathLabel(loadedTheme.sourcePath!, loadedTheme.sourceInfo), + ), + ); + addLoadedSection("Themes", themeCompactList, themeList); + } + } + + if (showDiagnostics) { + const skillDiagnostics = skillsResult.diagnostics; + if (skillDiagnostics.length > 0) { + const warningLines = this.formatDiagnostics(skillDiagnostics, sourceInfos); + this.chatContainer.addChild(new Text(`${theme.fg("warning", "[Skill conflicts]")}\n${warningLines}`, 0, 0)); + this.chatContainer.addChild(new Spacer(1)); + } + + const promptDiagnostics = promptsResult.diagnostics; + if (promptDiagnostics.length > 0) { + const warningLines = this.formatDiagnostics(promptDiagnostics, sourceInfos); + this.chatContainer.addChild( + new Text(`${theme.fg("warning", "[Prompt conflicts]")}\n${warningLines}`, 0, 0), + ); + this.chatContainer.addChild(new Spacer(1)); + } + + const extensionDiagnostics: ResourceDiagnostic[] = []; + const extensionErrors = this.session.resourceLoader.getExtensions().errors; + if (extensionErrors.length > 0) { + for (const error of extensionErrors) { + extensionDiagnostics.push({ type: "error", message: error.error, path: error.path }); + } + } + + const commandDiagnostics = this.session.extensionRunner.getCommandDiagnostics(); + extensionDiagnostics.push(...commandDiagnostics); + extensionDiagnostics.push(...this.getBuiltInCommandConflictDiagnostics(this.session.extensionRunner)); + + const shortcutDiagnostics = this.session.extensionRunner.getShortcutDiagnostics(); + extensionDiagnostics.push(...shortcutDiagnostics); + + if (extensionDiagnostics.length > 0) { + const warningLines = this.formatDiagnostics(extensionDiagnostics, sourceInfos); + this.chatContainer.addChild( + new Text(`${theme.fg("warning", "[Extension issues]")}\n${warningLines}`, 0, 0), + ); + this.chatContainer.addChild(new Spacer(1)); + } + + const themeDiagnostics = themesResult.diagnostics; + if (themeDiagnostics.length > 0) { + const warningLines = this.formatDiagnostics(themeDiagnostics, sourceInfos); + this.chatContainer.addChild(new Text(`${theme.fg("warning", "[Theme conflicts]")}\n${warningLines}`, 0, 0)); + this.chatContainer.addChild(new Spacer(1)); + } + } + } + + /** + * Initialize the extension system with TUI-based UI context. + */ + private async bindCurrentSessionExtensions(): Promise { + const uiContext = this.createExtensionUIContext(); + await this.session.bindExtensions({ + uiContext, + mode: "tui", + abortHandler: () => { + this.restoreQueuedMessagesToEditor({ abort: true }); + }, + commandContextActions: { + waitForIdle: () => this.session.agent.waitForIdle(), + newSession: async (options) => { + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.statusContainer.clear(); + try { + return await this.runtimeHost.newSession(options); + } catch (error: unknown) { + return this.handleFatalRuntimeError("Failed to create session", error); + } + }, + fork: async (entryId, options) => { + try { + const result = await this.runtimeHost.fork(entryId, options); + if (!result.cancelled) { + this.editor.setText(result.selectedText ?? ""); + this.showStatus("Forked to new session"); + } + return { cancelled: result.cancelled }; + } catch (error: unknown) { + return this.handleFatalRuntimeError("Failed to fork session", error); + } + }, + navigateTree: async (targetId, options) => { + const result = await this.session.navigateTree(targetId, { + summarize: options?.summarize, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }); + if (result.cancelled) { + return { cancelled: true }; + } + + this.chatContainer.clear(); + this.renderInitialMessages(); + if (result.editorText && !this.editor.getText().trim()) { + this.editor.setText(result.editorText); + } + this.showStatus("Navigated to selected point"); + void this.flushCompactionQueue({ willRetry: false }); + return { cancelled: false }; + }, + switchSession: async (sessionPath, options) => { + return this.handleResumeSession(sessionPath, options); + }, + reload: async () => { + await this.handleReloadCommand(); + }, + }, + shutdownHandler: () => { + this.shutdownRequested = true; + if (!this.session.isStreaming) { + void this.shutdown(); + } + }, + onError: (error) => { + this.showExtensionError(error.extensionPath, error.error, error.stack); + }, + }); + + setRegisteredThemes(this.session.resourceLoader.getThemes().themes); + this.setupAutocompleteProvider(); + + const extensionRunner = this.session.extensionRunner; + this.setupExtensionShortcuts(extensionRunner); + this.showLoadedResources({ force: false, showDiagnosticsWhenQuiet: true }); + } + + private applyRuntimeSettings(): void { + configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs()); + this.footer.setSession(this.session); + this.footer.setAutoCompactEnabled(this.session.autoCompactionEnabled); + this.footerDataProvider.setCwd(this.sessionManager.getCwd()); + this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock(); + this.ui.setShowHardwareCursor(this.settingsManager.getShowHardwareCursor()); + this.ui.setClearOnShrink(this.settingsManager.getClearOnShrink()); + const editorPaddingX = this.settingsManager.getEditorPaddingX(); + const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible(); + this.defaultEditor.setPaddingX(editorPaddingX); + this.defaultEditor.setAutocompleteMaxVisible(autocompleteMaxVisible); + if (this.editor !== this.defaultEditor) { + this.editor.setPaddingX?.(editorPaddingX); + this.editor.setAutocompleteMaxVisible?.(autocompleteMaxVisible); + } + } + + private async rebindCurrentSession(options: { renderBeforeBind?: boolean } = {}): Promise { + this.unsubscribe?.(); + this.unsubscribe = undefined; + this.applyRuntimeSettings(); + if (options.renderBeforeBind) { + this.renderCurrentSessionState(); + this.subscribeToAgent(); + await this.bindCurrentSessionExtensions(); + } else { + await this.bindCurrentSessionExtensions(); + this.subscribeToAgent(); + } + await this.updateAvailableProviderCount(); + this.updateEditorBorderColor(); + this.updateTerminalTitle(); + } + + private async handleFatalRuntimeError(prefix: string, error: unknown): Promise { + const message = error instanceof Error ? error.message : String(error); + this.showError(`${prefix}: ${message}`); + stopThemeWatcher(); + this.stop(); + process.exit(1); + } + + private renderCurrentSessionState(): void { + this.chatContainer.clear(); + this.pendingMessagesContainer.clear(); + this.compactionQueuedMessages = []; + this.streamingComponent = undefined; + this.streamingMessage = undefined; + this.pendingTools.clear(); + this.renderInitialMessages(); + } + + /** + * Get a registered tool definition by name (for custom rendering). + */ + private getRegisteredToolDefinition(toolName: string) { + return this.session.getToolDefinition(toolName); + } + + /** + * Set up keyboard shortcuts registered by extensions. + */ + private setupExtensionShortcuts(extensionRunner: ExtensionRunner): void { + const shortcuts = extensionRunner.getShortcuts(this.keybindings.getEffectiveConfig()); + if (shortcuts.size === 0) return; + + // Create a context for shortcut handlers + const createContext = (): ExtensionContext => ({ + ui: this.createExtensionUIContext(), + mode: "tui", + hasUI: true, + cwd: this.sessionManager.getCwd(), + sessionManager: this.sessionManager, + modelRegistry: this.session.modelRegistry, + model: this.session.model, + isIdle: () => !this.session.isStreaming, + isProjectTrusted: () => this.settingsManager.isProjectTrusted(), + signal: this.session.agent.signal, + abort: () => { + this.restoreQueuedMessagesToEditor({ abort: true }); + }, + hasPendingMessages: () => this.session.pendingMessageCount > 0, + shutdown: () => { + this.shutdownRequested = true; + }, + getContextUsage: () => this.session.getContextUsage(), + compact: (options) => { + void (async () => { + try { + const result = await this.session.compact(options?.customInstructions); + options?.onComplete?.(result); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + options?.onError?.(err); + } + })(); + }, + getSystemPrompt: () => this.session.systemPrompt, + }); + + // Set up the extension shortcut handler on the default editor + this.defaultEditor.onExtensionShortcut = (data: string) => { + for (const [shortcutStr, shortcut] of shortcuts) { + // Cast to KeyId - extension shortcuts use the same format + if (matchesKey(data, shortcutStr as KeyId)) { + // Run handler async, don't block input + Promise.resolve(shortcut.handler(createContext())).catch((err) => { + this.showError(`Shortcut handler error: ${err instanceof Error ? err.message : String(err)}`); + }); + return true; + } + } + return false; + }; + } + + /** + * Set extension status text in the footer. + */ + private setExtensionStatus(key: string, text: string | undefined): void { + this.footerDataProvider.setExtensionStatus(key, text); + this.ui.requestRender(); + } + + private getWorkingLoaderMessage(): string { + return this.workingMessage ?? this.defaultWorkingMessage; + } + + private createWorkingLoader(): Loader { + return new Loader( + this.ui, + (spinner) => theme.fg("accent", spinner), + (text) => theme.fg("muted", text), + this.getWorkingLoaderMessage(), + this.workingIndicatorOptions, + ); + } + + private stopWorkingLoader(): void { + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.statusContainer.clear(); + } + + private setWorkingVisible(visible: boolean): void { + this.workingVisible = visible; + if (!visible) { + this.stopWorkingLoader(); + this.ui.requestRender(); + return; + } + if (this.session.isStreaming && !this.loadingAnimation) { + this.statusContainer.clear(); + this.loadingAnimation = this.createWorkingLoader(); + this.statusContainer.addChild(this.loadingAnimation); + } + this.ui.requestRender(); + } + + private setWorkingIndicator(options?: LoaderIndicatorOptions): void { + this.workingIndicatorOptions = options; + this.loadingAnimation?.setIndicator(options); + this.ui.requestRender(); + } + + private setHiddenThinkingLabel(label?: string): void { + this.hiddenThinkingLabel = label ?? this.defaultHiddenThinkingLabel; + for (const child of this.chatContainer.children) { + if (child instanceof AssistantMessageComponent) { + child.setHiddenThinkingLabel(this.hiddenThinkingLabel); + } + } + if (this.streamingComponent) { + this.streamingComponent.setHiddenThinkingLabel(this.hiddenThinkingLabel); + } + this.ui.requestRender(); + } + + /** + * Set an extension widget (string array or custom component). + */ + private setExtensionWidget( + key: string, + content: string[] | ((tui: TUI, thm: Theme) => Component & { dispose?(): void }) | undefined, + options?: ExtensionWidgetOptions, + ): void { + const placement = options?.placement ?? "aboveEditor"; + const removeExisting = (map: Map) => { + const existing = map.get(key); + if (existing?.dispose) existing.dispose(); + map.delete(key); + }; + + removeExisting(this.extensionWidgetsAbove); + removeExisting(this.extensionWidgetsBelow); + + if (content === undefined) { + this.renderWidgets(); + return; + } + + let component: Component & { dispose?(): void }; + + if (Array.isArray(content)) { + // Wrap string array in a Container with Text components + const container = new Container(); + for (const line of content.slice(0, InteractiveMode.MAX_WIDGET_LINES)) { + container.addChild(new Text(line, 1, 0)); + } + if (content.length > InteractiveMode.MAX_WIDGET_LINES) { + container.addChild(new Text(theme.fg("muted", "... (widget truncated)"), 1, 0)); + } + component = container; + } else { + // Factory function - create component + component = content(this.ui, theme); + } + + const targetMap = placement === "belowEditor" ? this.extensionWidgetsBelow : this.extensionWidgetsAbove; + targetMap.set(key, component); + this.renderWidgets(); + } + + private clearExtensionWidgets(): void { + for (const widget of this.extensionWidgetsAbove.values()) { + widget.dispose?.(); + } + for (const widget of this.extensionWidgetsBelow.values()) { + widget.dispose?.(); + } + this.extensionWidgetsAbove.clear(); + this.extensionWidgetsBelow.clear(); + this.renderWidgets(); + } + + private resetExtensionUI(): void { + if (this.extensionSelector) { + this.hideExtensionSelector(); + } + if (this.extensionInput) { + this.hideExtensionInput(); + } + if (this.extensionEditor) { + this.hideExtensionEditor(); + } + this.ui.hideOverlay(); + this.clearExtensionTerminalInputListeners(); + this.setExtensionFooter(undefined); + this.setExtensionHeader(undefined); + this.clearExtensionWidgets(); + this.footerDataProvider.clearExtensionStatuses(); + this.footer.invalidate(); + this.autocompleteProviderWrappers = []; + this.setCustomEditorComponent(undefined); + this.setupAutocompleteProvider(); + this.defaultEditor.onExtensionShortcut = undefined; + this.updateTerminalTitle(); + this.workingMessage = undefined; + this.workingVisible = true; + this.setWorkingIndicator(); + if (this.loadingAnimation) { + this.loadingAnimation.setMessage(`${this.defaultWorkingMessage} (${keyText("app.interrupt")} to interrupt)`); + } + this.setHiddenThinkingLabel(); + } + + // Maximum total widget lines to prevent viewport overflow + private static readonly MAX_WIDGET_LINES = 10; + + /** + * Render all extension widgets to the widget container. + */ + private renderWidgets(): void { + if (!this.widgetContainerAbove || !this.widgetContainerBelow) return; + this.renderWidgetContainer(this.widgetContainerAbove, this.extensionWidgetsAbove, true, true); + this.renderWidgetContainer(this.widgetContainerBelow, this.extensionWidgetsBelow, false, false); + this.ui.requestRender(); + } + + private renderWidgetContainer( + container: Container, + widgets: Map, + spacerWhenEmpty: boolean, + leadingSpacer: boolean, + ): void { + container.clear(); + + if (widgets.size === 0) { + if (spacerWhenEmpty) { + container.addChild(new Spacer(1)); + } + return; + } + + if (leadingSpacer) { + container.addChild(new Spacer(1)); + } + for (const component of widgets.values()) { + container.addChild(component); + } + } + + /** + * Set a custom footer component, or restore the built-in footer. + */ + private setExtensionFooter( + factory: + | ((tui: TUI, thm: Theme, footerData: ReadonlyFooterDataProvider) => Component & { dispose?(): void }) + | undefined, + ): void { + // Dispose existing custom footer + if (this.customFooter?.dispose) { + this.customFooter.dispose(); + } + + // Remove current footer from UI + if (this.customFooter) { + this.ui.removeChild(this.customFooter); + } else { + this.ui.removeChild(this.footer); + } + + if (factory) { + // Create and add custom footer, passing the data provider + this.customFooter = factory(this.ui, theme, this.footerDataProvider); + this.ui.addChild(this.customFooter); + } else { + // Restore built-in footer + this.customFooter = undefined; + this.ui.addChild(this.footer); + } + + this.ui.requestRender(); + } + + /** + * Set a custom header component, or restore the built-in header. + */ + private setExtensionHeader(factory: ((tui: TUI, thm: Theme) => Component & { dispose?(): void }) | undefined): void { + // Header may not be initialized yet if called during early initialization + if (!this.builtInHeader) { + return; + } + + // Dispose existing custom header + if (this.customHeader?.dispose) { + this.customHeader.dispose(); + } + + // Find the index of the current header in the header container + const currentHeader = this.customHeader || this.builtInHeader; + const index = this.headerContainer.children.indexOf(currentHeader); + + if (factory) { + // Create and add custom header + this.customHeader = factory(this.ui, theme); + if (isExpandable(this.customHeader)) { + this.customHeader.setExpanded(this.toolOutputExpanded); + } + if (index !== -1) { + this.headerContainer.children[index] = this.customHeader; + } else { + // If not found (e.g. builtInHeader was never added), add at the top + this.headerContainer.children.unshift(this.customHeader); + } + } else { + // Restore built-in header + this.customHeader = undefined; + if (isExpandable(this.builtInHeader)) { + this.builtInHeader.setExpanded(this.toolOutputExpanded); + } + if (index !== -1) { + this.headerContainer.children[index] = this.builtInHeader; + } + } + + this.ui.requestRender(); + } + + private addExtensionTerminalInputListener( + handler: (data: string) => { consume?: boolean; data?: string } | undefined, + ): () => void { + const unsubscribe = this.ui.addInputListener(handler); + this.extensionTerminalInputUnsubscribers.add(unsubscribe); + return () => { + unsubscribe(); + this.extensionTerminalInputUnsubscribers.delete(unsubscribe); + }; + } + + private clearExtensionTerminalInputListeners(): void { + for (const unsubscribe of this.extensionTerminalInputUnsubscribers) { + unsubscribe(); + } + this.extensionTerminalInputUnsubscribers.clear(); + } + + /** + * Create the ExtensionUIContext for extensions. + */ + private createProjectTrustContext(cwd: string): ProjectTrustContext { + const ui = this.createExtensionUIContext(); + return { + cwd, + mode: "tui", + hasUI: true, + ui: { + select: ui.select, + confirm: ui.confirm, + input: ui.input, + notify: ui.notify, + }, + }; + } + + private createExtensionUIContext(): ExtensionUIContext { + return { + select: (title, options, opts) => this.showExtensionSelector(title, options, opts), + confirm: (title, message, opts) => this.showExtensionConfirm(title, message, opts), + input: (title, placeholder, opts) => this.showExtensionInput(title, placeholder, opts), + notify: (message, type) => this.showExtensionNotify(message, type), + onTerminalInput: (handler) => this.addExtensionTerminalInputListener(handler), + setStatus: (key, text) => this.setExtensionStatus(key, text), + setWorkingMessage: (message) => { + this.workingMessage = message; + if (this.loadingAnimation) { + this.loadingAnimation.setMessage(message ?? this.defaultWorkingMessage); + } + }, + setWorkingVisible: (visible) => this.setWorkingVisible(visible), + setWorkingIndicator: (options) => this.setWorkingIndicator(options), + setHiddenThinkingLabel: (label) => this.setHiddenThinkingLabel(label), + setWidget: (key, content, options) => this.setExtensionWidget(key, content, options), + setFooter: (factory) => this.setExtensionFooter(factory), + setHeader: (factory) => this.setExtensionHeader(factory), + setTitle: (title) => this.ui.terminal.setTitle(title), + custom: (factory, options) => this.showExtensionCustom(factory, options), + pasteToEditor: (text) => this.editor.handleInput(`\x1b[200~${text}\x1b[201~`), + setEditorText: (text) => this.editor.setText(text), + getEditorText: () => this.editor.getExpandedText?.() ?? this.editor.getText(), + editor: (title, prefill) => this.showExtensionEditor(title, prefill), + addAutocompleteProvider: (factory) => { + this.autocompleteProviderWrappers.push(factory); + this.setupAutocompleteProvider(); + }, + setEditorComponent: (factory) => this.setCustomEditorComponent(factory), + getEditorComponent: () => this.editorComponentFactory, + get theme() { + return theme; + }, + getAllThemes: () => getAvailableThemesWithPaths(), + getTheme: (name) => getThemeByName(name), + setTheme: (themeOrName) => { + if (themeOrName instanceof Theme) { + return this.themeController.setThemeInstance(themeOrName); + } + const result = this.themeController.setThemeName(themeOrName); + if (result.success) { + if (this.settingsManager.getTheme() !== themeOrName) { + this.settingsManager.setTheme(themeOrName); + } + } + return result; + }, + getToolsExpanded: () => this.toolOutputExpanded, + setToolsExpanded: (expanded) => this.setToolsExpanded(expanded), + }; + } + + /** + * Show a selector for extensions. + */ + private showExtensionSelector( + title: string, + options: string[], + opts?: ExtensionUIDialogOptions, + ): Promise { + return new Promise((resolve) => { + if (opts?.signal?.aborted) { + resolve(undefined); + return; + } + + const onAbort = () => { + this.hideExtensionSelector(); + resolve(undefined); + }; + opts?.signal?.addEventListener("abort", onAbort, { once: true }); + + this.extensionSelector = new ExtensionSelectorComponent( + title, + options, + (option) => { + opts?.signal?.removeEventListener("abort", onAbort); + this.hideExtensionSelector(); + resolve(option); + }, + () => { + opts?.signal?.removeEventListener("abort", onAbort); + this.hideExtensionSelector(); + resolve(undefined); + }, + { tui: this.ui, timeout: opts?.timeout, onToggleToolsExpanded: () => this.toggleToolOutputExpansion() }, + ); + + this.editorContainer.clear(); + this.editorContainer.addChild(this.extensionSelector); + this.ui.setFocus(this.extensionSelector); + this.ui.requestRender(); + }); + } + + /** + * Hide the extension selector. + */ + private hideExtensionSelector(): void { + this.extensionSelector?.dispose(); + this.editorContainer.clear(); + this.editorContainer.addChild(this.editor); + this.extensionSelector = undefined; + this.ui.setFocus(this.editor); + this.ui.requestRender(); + } + + /** + * Show a confirmation dialog for extensions. + */ + private async showExtensionConfirm( + title: string, + message: string, + opts?: ExtensionUIDialogOptions, + ): Promise { + const result = await this.showExtensionSelector(`${title}\n${message}`, ["Yes", "No"], opts); + return result === "Yes"; + } + + private async promptForMissingSessionCwd(error: MissingSessionCwdError): Promise { + const confirmed = await this.showExtensionConfirm( + "Session cwd not found", + formatMissingSessionCwdPrompt(error.issue), + ); + return confirmed ? error.issue.fallbackCwd : undefined; + } + + /** + * Show a text input for extensions. + */ + private showExtensionInput( + title: string, + placeholder?: string, + opts?: ExtensionUIDialogOptions, + ): Promise { + return new Promise((resolve) => { + if (opts?.signal?.aborted) { + resolve(undefined); + return; + } + + const onAbort = () => { + this.hideExtensionInput(); + resolve(undefined); + }; + opts?.signal?.addEventListener("abort", onAbort, { once: true }); + + this.extensionInput = new ExtensionInputComponent( + title, + placeholder, + (value) => { + opts?.signal?.removeEventListener("abort", onAbort); + this.hideExtensionInput(); + resolve(value); + }, + () => { + opts?.signal?.removeEventListener("abort", onAbort); + this.hideExtensionInput(); + resolve(undefined); + }, + { tui: this.ui, timeout: opts?.timeout }, + ); + + this.editorContainer.clear(); + this.editorContainer.addChild(this.extensionInput); + this.ui.setFocus(this.extensionInput); + this.ui.requestRender(); + }); + } + + /** + * Hide the extension input. + */ + private hideExtensionInput(): void { + this.extensionInput?.dispose(); + this.editorContainer.clear(); + this.editorContainer.addChild(this.editor); + this.extensionInput = undefined; + this.ui.setFocus(this.editor); + this.ui.requestRender(); + } + + /** + * Show a multi-line editor for extensions (with Ctrl+G support). + */ + private showExtensionEditor(title: string, prefill?: string): Promise { + return new Promise((resolve) => { + this.extensionEditor = new ExtensionEditorComponent( + this.ui, + this.keybindings, + title, + prefill, + (value) => { + this.hideExtensionEditor(); + resolve(value); + }, + () => { + this.hideExtensionEditor(); + resolve(undefined); + }, + ); + + this.editorContainer.clear(); + this.editorContainer.addChild(this.extensionEditor); + this.ui.setFocus(this.extensionEditor); + this.ui.requestRender(); + }); + } + + /** + * Hide the extension editor. + */ + private hideExtensionEditor(): void { + this.editorContainer.clear(); + this.editorContainer.addChild(this.editor); + this.extensionEditor = undefined; + this.ui.setFocus(this.editor); + this.ui.requestRender(); + } + + /** + * Set a custom editor component from an extension. + * Pass undefined to restore the default editor. + */ + private setCustomEditorComponent(factory: EditorFactory | undefined): void { + this.editorComponentFactory = factory; + + // Save text from current editor before switching + const currentText = this.editor.getText(); + + this.editorContainer.clear(); + + if (factory) { + // Create the custom editor with tui, theme, and keybindings + const newEditor = factory(this.ui, getEditorTheme(), this.keybindings); + + // Wire up callbacks from the default editor + newEditor.onSubmit = this.defaultEditor.onSubmit; + newEditor.onChange = this.defaultEditor.onChange; + + // Copy text from previous editor + newEditor.setText(currentText); + + // Copy appearance settings if supported + if (newEditor.borderColor !== undefined) { + newEditor.borderColor = this.defaultEditor.borderColor; + } + if (newEditor.setPaddingX !== undefined) { + newEditor.setPaddingX(this.defaultEditor.getPaddingX()); + } + + // Set autocomplete if supported + if (newEditor.setAutocompleteProvider && this.autocompleteProvider) { + newEditor.setAutocompleteProvider(this.autocompleteProvider); + } + + // If extending CustomEditor, copy app-level handlers + // Use duck typing since instanceof fails across jiti module boundaries + const customEditor = newEditor as unknown as Record; + if ("actionHandlers" in customEditor && customEditor.actionHandlers instanceof Map) { + if (!customEditor.onEscape) { + customEditor.onEscape = () => this.defaultEditor.onEscape?.(); + } + if (!customEditor.onCtrlD) { + customEditor.onCtrlD = () => this.defaultEditor.onCtrlD?.(); + } + if (!customEditor.onPasteImage) { + customEditor.onPasteImage = () => this.defaultEditor.onPasteImage?.(); + } + if (!customEditor.onExtensionShortcut) { + customEditor.onExtensionShortcut = (data: string) => this.defaultEditor.onExtensionShortcut?.(data); + } + // Copy action handlers (clear, suspend, model switching, etc.) + for (const [action, handler] of this.defaultEditor.actionHandlers) { + (customEditor.actionHandlers as Map void>).set(action, handler); + } + } + + this.editor = newEditor; + } else { + // Restore default editor with text from custom editor + this.defaultEditor.setText(currentText); + this.editor = this.defaultEditor; + } + + this.editorContainer.addChild(this.editor as Component); + this.ui.setFocus(this.editor as Component); + this.ui.requestRender(); + } + + /** + * Show a notification for extensions. + */ + private showExtensionNotify(message: string, type?: "info" | "warning" | "error"): void { + if (type === "error") { + this.showError(message); + } else if (type === "warning") { + this.showWarning(message); + } else { + this.showStatus(message); + } + } + + /** Show a custom component with keyboard focus. Overlay mode renders on top of existing content. */ + private async showExtensionCustom( + factory: ( + tui: TUI, + theme: Theme, + keybindings: KeybindingsManager, + done: (result: T) => void, + ) => (Component & { dispose?(): void }) | Promise, + options?: { + overlay?: boolean; + overlayOptions?: OverlayOptions | (() => OverlayOptions); + onHandle?: (handle: OverlayHandle) => void; + }, + ): Promise { + const savedText = this.editor.getText(); + const isOverlay = options?.overlay ?? false; + + const restoreEditor = () => { + this.editorContainer.clear(); + this.editorContainer.addChild(this.editor); + this.editor.setText(savedText); + this.ui.setFocus(this.editor); + this.ui.requestRender(); + }; + + return new Promise((resolve, reject) => { + let component: Component & { dispose?(): void }; + let closed = false; + + const close = (result: T) => { + if (closed) return; + closed = true; + if (isOverlay) this.ui.hideOverlay(); + else restoreEditor(); + // Note: both branches above already call requestRender + resolve(result); + try { + component?.dispose?.(); + } catch { + /* ignore dispose errors */ + } + }; + + Promise.resolve(factory(this.ui, theme, this.keybindings, close)) + .then((c) => { + if (closed) return; + component = c; + if (isOverlay) { + // Resolve overlay options - can be static or dynamic function + const resolveOptions = (): OverlayOptions | undefined => { + if (options?.overlayOptions) { + const opts = + typeof options.overlayOptions === "function" + ? options.overlayOptions() + : options.overlayOptions; + return opts; + } + // Fallback: use component's width property if available + const w = (component as { width?: number }).width; + return w ? { width: w } : undefined; + }; + const handle = this.ui.showOverlay(component, resolveOptions()); + // Expose handle to caller for visibility control + options?.onHandle?.(handle); + } else { + this.editorContainer.clear(); + this.editorContainer.addChild(component); + this.ui.setFocus(component); + this.ui.requestRender(); + } + }) + .catch((err) => { + if (closed) return; + if (!isOverlay) restoreEditor(); + reject(err); + }); + }); + } + + /** + * Show an extension error in the UI. + */ + private showExtensionError(extensionPath: string, error: string, stack?: string): void { + const errorMsg = `Extension "${extensionPath}" error: ${error}`; + const errorText = new Text(theme.fg("error", errorMsg), 1, 0); + this.chatContainer.addChild(errorText); + if (stack) { + // Show stack trace in dim color, indented + const stackLines = stack + .split("\n") + .slice(1) // Skip first line (duplicates error message) + .map((line) => theme.fg("dim", ` ${line.trim()}`)) + .join("\n"); + if (stackLines) { + this.chatContainer.addChild(new Text(stackLines, 1, 0)); + } + } + this.ui.requestRender(); + } + + // ========================================================================= + // Key Handlers + // ========================================================================= + + private setupKeyHandlers(): void { + // Set up handlers on defaultEditor - they use this.editor for text access + // so they work correctly regardless of which editor is active + this.defaultEditor.onEscape = () => { + if (this.session.isStreaming) { + this.restoreQueuedMessagesToEditor({ abort: true }); + } else if (this.session.isBashRunning) { + this.session.abortBash(); + } else if (this.isBashMode) { + this.editor.setText(""); + this.isBashMode = false; + this.updateEditorBorderColor(); + } else if (!this.editor.getText().trim()) { + // Double-escape with empty editor triggers /tree, /fork, or nothing based on setting + const action = this.settingsManager.getDoubleEscapeAction(); + if (action !== "none") { + const now = Date.now(); + if (now - this.lastEscapeTime < 500) { + if (action === "tree") { + this.showTreeSelector(); + } else { + this.showUserMessageSelector(); + } + this.lastEscapeTime = 0; + } else { + this.lastEscapeTime = now; + } + } + } + }; + + // Register app action handlers + this.defaultEditor.onAction("app.clear", () => this.handleCtrlC()); + this.defaultEditor.onCtrlD = () => this.handleCtrlD(); + this.defaultEditor.onAction("app.suspend", () => this.handleCtrlZ()); + this.defaultEditor.onAction("app.thinking.cycle", () => this.cycleThinkingLevel()); + this.defaultEditor.onAction("app.model.cycleForward", () => this.cycleModel("forward")); + this.defaultEditor.onAction("app.model.cycleBackward", () => this.cycleModel("backward")); + + // Global debug handler on TUI (works regardless of focus) + this.ui.onDebug = () => this.handleDebugCommand(); + this.defaultEditor.onAction("app.model.select", () => this.showModelSelector()); + this.defaultEditor.onAction("app.tools.expand", () => this.toggleToolOutputExpansion()); + this.defaultEditor.onAction("app.thinking.toggle", () => this.toggleThinkingBlockVisibility()); + this.defaultEditor.onAction("app.editor.external", () => this.openExternalEditor()); + this.defaultEditor.onAction("app.message.followUp", () => this.handleFollowUp()); + this.defaultEditor.onAction("app.message.dequeue", () => this.handleDequeue()); + this.defaultEditor.onAction("app.session.new", () => this.handleClearCommand()); + this.defaultEditor.onAction("app.session.tree", () => this.showTreeSelector()); + this.defaultEditor.onAction("app.session.fork", () => this.showUserMessageSelector()); + this.defaultEditor.onAction("app.session.resume", () => this.showSessionSelector()); + + this.defaultEditor.onChange = (text: string) => { + const wasBashMode = this.isBashMode; + this.isBashMode = text.trimStart().startsWith("!"); + if (wasBashMode !== this.isBashMode) { + this.updateEditorBorderColor(); + } + }; + + // Handle clipboard image paste (triggered on Ctrl+V) + this.defaultEditor.onPasteImage = () => { + this.handleClipboardImagePaste(); + }; + } + + private async handleClipboardImagePaste(): Promise { + try { + const image = await readClipboardImage(); + if (!image) { + return; + } + + // Write to temp file + const tmpDir = os.tmpdir(); + const ext = extensionForImageMimeType(image.mimeType) ?? "png"; + const fileName = `cactus-clipboard-${crypto.randomUUID()}.${ext}`; + const filePath = path.join(tmpDir, fileName); + fs.writeFileSync(filePath, Buffer.from(image.bytes)); + + // Insert file path directly + this.editor.insertTextAtCursor?.(filePath); + this.ui.requestRender(); + } catch { + // Silently ignore clipboard errors (may not have permission, etc.) + } + } + + private setupEditorSubmitHandler(): void { + this.defaultEditor.onSubmit = async (text: string) => { + text = text.trim(); + if (!text) return; + + // Handle commands + if (text === "/settings") { + this.showSettingsSelector(); + this.editor.setText(""); + return; + } + if (text === "/scoped-models") { + this.editor.setText(""); + await this.showModelsSelector(); + return; + } + if (text === "/model" || text.startsWith("/model ")) { + const searchTerm = text.startsWith("/model ") ? text.slice(7).trim() : undefined; + this.editor.setText(""); + await this.handleModelCommand(searchTerm); + return; + } + if (text === "/export" || text.startsWith("/export ")) { + await this.handleExportCommand(text); + this.editor.setText(""); + return; + } + if (text === "/import" || text.startsWith("/import ")) { + await this.handleImportCommand(text); + this.editor.setText(""); + return; + } + if (text === "/copy") { + await this.handleCopyCommand(); + this.editor.setText(""); + return; + } + if (text === "/name" || text.startsWith("/name ")) { + this.handleNameCommand(text); + this.editor.setText(""); + return; + } + if (text === "/session") { + this.handleSessionCommand(); + this.editor.setText(""); + return; + } + if (text === "/hotkeys") { + this.handleHotkeysCommand(); + this.editor.setText(""); + return; + } + if (text === "/fork") { + this.showUserMessageSelector(); + this.editor.setText(""); + return; + } + if (text === "/clone") { + this.editor.setText(""); + await this.handleCloneCommand(); + return; + } + if (text === "/tree") { + this.showTreeSelector(); + this.editor.setText(""); + return; + } + if (text === "/trust") { + this.showTrustSelector(); + this.editor.setText(""); + return; + } + if (text === "/new") { + this.editor.setText(""); + await this.handleClearCommand(); + return; + } + if (text === "/compact" || text.startsWith("/compact ")) { + const customInstructions = text.startsWith("/compact ") ? text.slice(9).trim() : undefined; + this.editor.setText(""); + await this.handleCompactCommand(customInstructions); + return; + } + if (text === "/reload") { + this.editor.setText(""); + await this.handleReloadCommand(); + return; + } + if (text === "/debug") { + this.handleDebugCommand(); + this.editor.setText(""); + return; + } + if (text === "/resume") { + this.showSessionSelector(); + this.editor.setText(""); + return; + } + if (text === "/quit") { + this.editor.setText(""); + await this.shutdown(); + return; + } + + // Handle bash command (! for normal, !! for excluded from context) + if (text.startsWith("!")) { + const isExcluded = text.startsWith("!!"); + const command = isExcluded ? text.slice(2).trim() : text.slice(1).trim(); + if (command) { + if (this.session.isBashRunning) { + this.showWarning("A bash command is already running. Press Esc to cancel it first."); + this.editor.setText(text); + return; + } + this.editor.addToHistory?.(text); + await this.handleBashCommand(command, isExcluded); + this.isBashMode = false; + this.updateEditorBorderColor(); + return; + } + } + + // Queue input during compaction (extension commands execute immediately) + if (this.session.isCompacting) { + if (this.isExtensionCommand(text)) { + this.editor.addToHistory?.(text); + this.editor.setText(""); + await this.session.prompt(text); + } else { + this.queueCompactionMessage(text, "steer"); + } + return; + } + + // If streaming, use prompt() with steer behavior + // This handles extension commands (execute immediately), prompt template expansion, and queueing + if (this.session.isStreaming) { + this.editor.addToHistory?.(text); + this.editor.setText(""); + await this.session.prompt(text, { streamingBehavior: "steer" }); + this.updatePendingMessagesDisplay(); + this.ui.requestRender(); + return; + } + + // Normal message submission + // First, move any pending bash components to chat + this.flushPendingBashComponents(); + + if (this.onInputCallback) { + this.onInputCallback(text); + } else { + this.pendingUserInputs.push(text); + } + this.editor.addToHistory?.(text); + }; + } + + private subscribeToAgent(): void { + this.unsubscribe = this.session.subscribe(async (event) => { + await this.handleEvent(event); + }); + } + + private async handleEvent(event: AgentSessionEvent): Promise { + if (!this.isInitialized) { + await this.init(); + } + + this.footer.invalidate(); + + switch (event.type) { + case "agent_start": + this.pendingTools.clear(); + if (this.settingsManager.getShowTerminalProgress()) { + this.ui.terminal.setProgress(true); + } + // Restore main escape handler if retry handler is still active + // (retry success event fires later, but we need main handler now) + if (this.retryEscapeHandler) { + this.defaultEditor.onEscape = this.retryEscapeHandler; + this.retryEscapeHandler = undefined; + } + if (this.retryCountdown) { + this.retryCountdown.dispose(); + this.retryCountdown = undefined; + } + if (this.retryLoader) { + this.retryLoader.stop(); + this.retryLoader = undefined; + } + this.stopWorkingLoader(); + if (this.workingVisible) { + this.loadingAnimation = this.createWorkingLoader(); + this.statusContainer.addChild(this.loadingAnimation); + } + this.ui.requestRender(); + break; + + case "queue_update": + this.updatePendingMessagesDisplay(); + this.ui.requestRender(); + break; + + case "session_info_changed": + this.updateTerminalTitle(); + this.footer.invalidate(); + this.ui.requestRender(); + break; + + case "thinking_level_changed": + this.footer.invalidate(); + this.updateEditorBorderColor(); + break; + + case "message_start": + if (event.message.role === "custom") { + this.addMessageToChat(event.message); + this.ui.requestRender(); + } else if (event.message.role === "user") { + this.addMessageToChat(event.message); + this.updatePendingMessagesDisplay(); + this.ui.requestRender(); + } else if (event.message.role === "assistant") { + this.streamingComponent = new AssistantMessageComponent( + undefined, + this.hideThinkingBlock, + this.getMarkdownThemeWithSettings(), + this.hiddenThinkingLabel, + ); + this.streamingMessage = event.message; + this.chatContainer.addChild(this.streamingComponent); + this.streamingComponent.updateContent(this.streamingMessage); + this.ui.requestRender(); + } + break; + + case "message_update": + if (this.streamingComponent && event.message.role === "assistant") { + this.streamingMessage = event.message; + this.streamingComponent.updateContent(this.streamingMessage); + + for (const content of this.streamingMessage.content) { + if (content.type === "toolCall") { + if (!this.pendingTools.has(content.id)) { + const component = new ToolExecutionComponent( + content.name, + content.id, + content.arguments, + { + showImages: this.settingsManager.getShowImages(), + imageWidthCells: this.settingsManager.getImageWidthCells(), + }, + this.getRegisteredToolDefinition(content.name), + this.ui, + this.sessionManager.getCwd(), + ); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + this.pendingTools.set(content.id, component); + } else { + const component = this.pendingTools.get(content.id); + if (component) { + component.updateArgs(content.arguments); + } + } + } + } + this.ui.requestRender(); + } + break; + + case "message_end": + if (event.message.role === "user") break; + if (this.streamingComponent && event.message.role === "assistant") { + this.streamingMessage = event.message; + let errorMessage: string | undefined; + if (this.streamingMessage.stopReason === "aborted") { + const retryAttempt = this.session.retryAttempt; + errorMessage = + retryAttempt > 0 + ? `Aborted after ${retryAttempt} retry attempt${retryAttempt > 1 ? "s" : ""}` + : "Operation aborted"; + this.streamingMessage.errorMessage = errorMessage; + } + this.streamingComponent.updateContent(this.streamingMessage); + + if (this.streamingMessage.stopReason === "aborted" || this.streamingMessage.stopReason === "error") { + if (!errorMessage) { + errorMessage = this.streamingMessage.errorMessage || "Error"; + } + for (const [, component] of this.pendingTools.entries()) { + component.updateResult({ + content: [{ type: "text", text: errorMessage }], + isError: true, + }); + } + this.pendingTools.clear(); + } else { + // Args are now complete - trigger diff computation for edit tools + for (const [, component] of this.pendingTools.entries()) { + component.setArgsComplete(); + } + } + this.streamingComponent = undefined; + this.streamingMessage = undefined; + this.footer.invalidate(); + } + this.ui.requestRender(); + break; + + case "tool_execution_start": { + let component = this.pendingTools.get(event.toolCallId); + if (!component) { + component = new ToolExecutionComponent( + event.toolName, + event.toolCallId, + event.args, + { + showImages: this.settingsManager.getShowImages(), + imageWidthCells: this.settingsManager.getImageWidthCells(), + }, + this.getRegisteredToolDefinition(event.toolName), + this.ui, + this.sessionManager.getCwd(), + ); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + this.pendingTools.set(event.toolCallId, component); + } + component.markExecutionStarted(); + this.ui.requestRender(); + break; + } + + case "tool_execution_update": { + const component = this.pendingTools.get(event.toolCallId); + if (component) { + component.updateResult({ ...event.partialResult, isError: false }, true); + this.ui.requestRender(); + } + break; + } + + case "tool_execution_end": { + const component = this.pendingTools.get(event.toolCallId); + if (component) { + component.updateResult({ ...event.result, isError: event.isError }); + this.pendingTools.delete(event.toolCallId); + this.ui.requestRender(); + } + break; + } + + case "agent_end": + if (this.settingsManager.getShowTerminalProgress()) { + this.ui.terminal.setProgress(false); + } + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + this.statusContainer.clear(); + } + if (this.streamingComponent) { + this.chatContainer.removeChild(this.streamingComponent); + this.streamingComponent = undefined; + this.streamingMessage = undefined; + } + this.pendingTools.clear(); + + await this.checkShutdownRequested(); + + this.ui.requestRender(); + break; + + case "compaction_start": { + if (this.settingsManager.getShowTerminalProgress()) { + this.ui.terminal.setProgress(true); + } + // Keep editor active; submissions are queued during compaction. + this.autoCompactionEscapeHandler = this.defaultEditor.onEscape; + this.defaultEditor.onEscape = () => { + this.session.abortCompaction(); + }; + this.statusContainer.clear(); + const cancelHint = `(${keyText("app.interrupt")} to cancel)`; + const label = + event.reason === "manual" + ? `Compacting context... ${cancelHint}` + : `${event.reason === "overflow" ? "Context overflow detected, " : ""}Auto-compacting... ${cancelHint}`; + this.autoCompactionLoader = new Loader( + this.ui, + (spinner) => theme.fg("accent", spinner), + (text) => theme.fg("muted", text), + label, + ); + this.statusContainer.addChild(this.autoCompactionLoader); + this.ui.requestRender(); + break; + } + + case "compaction_end": { + if (this.settingsManager.getShowTerminalProgress()) { + this.ui.terminal.setProgress(false); + } + if (this.autoCompactionEscapeHandler) { + this.defaultEditor.onEscape = this.autoCompactionEscapeHandler; + this.autoCompactionEscapeHandler = undefined; + } + if (this.autoCompactionLoader) { + this.autoCompactionLoader.stop(); + this.autoCompactionLoader = undefined; + this.statusContainer.clear(); + } + if (event.aborted) { + if (event.reason === "manual") { + this.showError("Compaction cancelled"); + } else { + this.showStatus("Auto-compaction cancelled"); + } + } else if (event.result) { + this.chatContainer.clear(); + this.rebuildChatFromMessages(); + this.addMessageToChat( + createCompactionSummaryMessage( + event.result.summary, + event.result.tokensBefore, + new Date().toISOString(), + ), + ); + this.footer.invalidate(); + } else if (event.errorMessage) { + if (event.reason === "manual") { + this.showError(event.errorMessage); + } else { + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("error", event.errorMessage), 1, 0)); + } + } + void this.flushCompactionQueue({ willRetry: event.willRetry }); + this.ui.requestRender(); + break; + } + + case "auto_retry_start": { + // Set up escape to abort retry + this.retryEscapeHandler = this.defaultEditor.onEscape; + this.defaultEditor.onEscape = () => { + this.session.abortRetry(); + }; + // Show retry indicator + this.statusContainer.clear(); + this.retryCountdown?.dispose(); + const retryMessage = (seconds: number) => + `Retrying (${event.attempt}/${event.maxAttempts}) in ${seconds}s... (${keyText("app.interrupt")} to cancel)`; + this.retryLoader = new Loader( + this.ui, + (spinner) => theme.fg("warning", spinner), + (text) => theme.fg("muted", text), + retryMessage(Math.ceil(event.delayMs / 1000)), + ); + this.retryCountdown = new CountdownTimer( + event.delayMs, + this.ui, + (seconds) => { + this.retryLoader?.setMessage(retryMessage(seconds)); + }, + () => { + this.retryCountdown = undefined; + }, + ); + this.statusContainer.addChild(this.retryLoader); + this.ui.requestRender(); + break; + } + + case "auto_retry_end": { + // Restore escape handler + if (this.retryEscapeHandler) { + this.defaultEditor.onEscape = this.retryEscapeHandler; + this.retryEscapeHandler = undefined; + } + if (this.retryCountdown) { + this.retryCountdown.dispose(); + this.retryCountdown = undefined; + } + // Stop loader + if (this.retryLoader) { + this.retryLoader.stop(); + this.retryLoader = undefined; + this.statusContainer.clear(); + } + // Show error only on final failure (success shows normal response) + if (!event.success) { + this.showError(`Retry failed after ${event.attempt} attempts: ${event.finalError || "Unknown error"}`); + } + this.ui.requestRender(); + break; + } + } + } + + /** Extract text content from a user message */ + private getUserMessageText(message: Message): string { + if (message.role !== "user") return ""; + const textBlocks = + typeof message.content === "string" + ? [{ type: "text", text: message.content }] + : message.content.filter((c: { type: string }) => c.type === "text"); + return textBlocks.map((c) => (c as { text: string }).text).join(""); + } + + /** + * Show a status message in the chat. + * + * If multiple status messages are emitted back-to-back (without anything else being added to the chat), + * we update the previous status line instead of appending new ones to avoid log spam. + */ + private showStatus(message: string): void { + const children = this.chatContainer.children; + const last = children.length > 0 ? children[children.length - 1] : undefined; + const secondLast = children.length > 1 ? children[children.length - 2] : undefined; + + if (last && secondLast && last === this.lastStatusText && secondLast === this.lastStatusSpacer) { + this.lastStatusText.setText(theme.fg("dim", message)); + this.ui.requestRender(); + return; + } + + const spacer = new Spacer(1); + const text = new Text(theme.fg("dim", message), 1, 0); + this.chatContainer.addChild(spacer); + this.chatContainer.addChild(text); + this.lastStatusSpacer = spacer; + this.lastStatusText = text; + this.ui.requestRender(); + } + + private addMessageToChat(message: AgentMessage, options?: { populateHistory?: boolean }): void { + switch (message.role) { + case "bashExecution": { + const component = new BashExecutionComponent(message.command, this.ui, message.excludeFromContext); + if (message.output) { + component.appendOutput(message.output); + } + component.setComplete( + message.exitCode, + message.cancelled, + message.truncated ? ({ truncated: true } as TruncationResult) : undefined, + message.fullOutputPath, + ); + this.chatContainer.addChild(component); + break; + } + case "custom": { + if (message.display) { + const renderer = this.session.extensionRunner.getMessageRenderer(message.customType); + const component = new CustomMessageComponent(message, renderer, this.getMarkdownThemeWithSettings()); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + } + break; + } + case "compactionSummary": { + this.chatContainer.addChild(new Spacer(1)); + const component = new CompactionSummaryMessageComponent(message, this.getMarkdownThemeWithSettings()); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + break; + } + case "branchSummary": { + this.chatContainer.addChild(new Spacer(1)); + const component = new BranchSummaryMessageComponent(message, this.getMarkdownThemeWithSettings()); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + break; + } + case "user": { + const textContent = this.getUserMessageText(message); + if (textContent) { + if (this.chatContainer.children.length > 0) { + this.chatContainer.addChild(new Spacer(1)); + } + const skillBlock = parseSkillBlock(textContent); + if (skillBlock) { + // Render skill block (collapsible) + const component = new SkillInvocationMessageComponent( + skillBlock, + this.getMarkdownThemeWithSettings(), + ); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + // Render user message separately if present + if (skillBlock.userMessage) { + this.chatContainer.addChild(new Spacer(1)); + const userComponent = new UserMessageComponent( + skillBlock.userMessage, + this.getMarkdownThemeWithSettings(), + ); + this.chatContainer.addChild(userComponent); + } + } else { + const userComponent = new UserMessageComponent(textContent, this.getMarkdownThemeWithSettings()); + this.chatContainer.addChild(userComponent); + } + if (options?.populateHistory) { + this.editor.addToHistory?.(textContent); + } + } + break; + } + case "assistant": { + const assistantComponent = new AssistantMessageComponent( + message, + this.hideThinkingBlock, + this.getMarkdownThemeWithSettings(), + this.hiddenThinkingLabel, + ); + this.chatContainer.addChild(assistantComponent); + break; + } + case "toolResult": { + // Tool results are rendered inline with tool calls, handled separately + break; + } + default: { + const _exhaustive: never = message; + } + } + } + + /** + * Render session context to chat. Used for initial load and rebuild after compaction. + * @param sessionContext Session context to render + * @param options.updateFooter Update footer state + * @param options.populateHistory Add user messages to editor history + */ + private renderSessionContext( + sessionContext: SessionContext, + options: { updateFooter?: boolean; populateHistory?: boolean } = {}, + ): void { + this.pendingTools.clear(); + const renderedPendingTools = new Map(); + + if (options.updateFooter) { + this.footer.invalidate(); + this.updateEditorBorderColor(); + } + + for (const message of sessionContext.messages) { + // Assistant messages need special handling for tool calls + if (message.role === "assistant") { + this.addMessageToChat(message); + // Render tool call components + for (const content of message.content) { + if (content.type === "toolCall") { + const component = new ToolExecutionComponent( + content.name, + content.id, + content.arguments, + { + showImages: this.settingsManager.getShowImages(), + imageWidthCells: this.settingsManager.getImageWidthCells(), + }, + this.getRegisteredToolDefinition(content.name), + this.ui, + this.sessionManager.getCwd(), + ); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + + if (message.stopReason === "aborted" || message.stopReason === "error") { + let errorMessage: string; + if (message.stopReason === "aborted") { + const retryAttempt = this.session.retryAttempt; + errorMessage = + retryAttempt > 0 + ? `Aborted after ${retryAttempt} retry attempt${retryAttempt > 1 ? "s" : ""}` + : "Operation aborted"; + } else { + errorMessage = message.errorMessage || "Error"; + } + component.updateResult({ content: [{ type: "text", text: errorMessage }], isError: true }); + } else { + renderedPendingTools.set(content.id, component); + } + } + } + } else if (message.role === "toolResult") { + // Match tool results to pending tool components + const component = renderedPendingTools.get(message.toolCallId); + if (component) { + component.updateResult(message); + renderedPendingTools.delete(message.toolCallId); + } + } else { + // All other messages use standard rendering + this.addMessageToChat(message, options); + } + } + + for (const [toolCallId, component] of renderedPendingTools) { + this.pendingTools.set(toolCallId, component); + } + this.ui.requestRender(); + } + + renderInitialMessages(): void { + // Get aligned messages and entries from session context + const context = this.sessionManager.buildSessionContext(); + this.renderSessionContext(context, { + updateFooter: true, + populateHistory: true, + }); + this.renderProjectTrustWarningIfNeeded(); + + // Show compaction info if session was compacted + const allEntries = this.sessionManager.getEntries(); + const compactionCount = allEntries.filter((e) => e.type === "compaction").length; + if (compactionCount > 0) { + const times = compactionCount === 1 ? "1 time" : `${compactionCount} times`; + this.showStatus(`Session compacted ${times}`); + } + } + + private renderProjectTrustWarningIfNeeded(): void { + if (this.settingsManager.isProjectTrusted() || !hasTrustRequiringProjectResources(this.sessionManager.getCwd())) { + return; + } + + if (this.chatContainer.children.length > 0) { + this.chatContainer.addChild(new Spacer(1)); + } + this.chatContainer.addChild( + new Text( + theme.fg( + "warning", + `This project is not trusted. Project ${CONFIG_DIR_NAME} resources are ignored. Use /trust to save a trust decision, then restart Cactus.`, + ), + 1, + 0, + ), + ); + } + + async getUserInput(): Promise { + const queuedInput = this.pendingUserInputs.shift(); + if (queuedInput !== undefined) { + return queuedInput; + } + + return new Promise((resolve) => { + this.onInputCallback = (text: string) => { + this.onInputCallback = undefined; + resolve(text); + }; + }); + } + + private rebuildChatFromMessages(): void { + this.chatContainer.clear(); + const context = this.sessionManager.buildSessionContext(); + this.renderSessionContext(context); + } + + // ========================================================================= + // Key handlers + // ========================================================================= + + private handleCtrlC(): void { + const now = Date.now(); + if (now - this.lastSigintTime < 500) { + void this.shutdown(); + } else { + this.clearEditor(); + this.lastSigintTime = now; + } + } + + private handleCtrlD(): void { + // Only called when editor is empty (enforced by CustomEditor) + void this.shutdown(); + } + + /** + * Gracefully shutdown the agent. + * Stops the TUI before emitting shutdown events so extension UI cleanup cannot + * repaint the final frame while the process is exiting. + */ + private isShuttingDown = false; + + private async shutdown(options?: { fromSignal?: boolean }): Promise { + if (this.isShuttingDown) return; + this.isShuttingDown = true; + // Keep signal handlers registered until terminal cleanup has completed. + // `signal-exit` checks the listener list during the same SIGTERM/SIGHUP + // dispatch and re-sends the signal if only its own listeners remain. + + if (options?.fromSignal) { + // Signal-triggered shutdown (SIGTERM/SIGHUP). Emit extension cleanup + // (session_shutdown) BEFORE touching the terminal. Extension teardown + // such as removing sockets does not write to the tty, so it must not be + // skipped if a later terminal-restore write fails on a dead or stalled + // terminal. If the terminal is gone, the restore writes below emit EIO, + // which the stdout/stderr error handler turns into emergencyTerminalExit; + // the render loop is already idle, so this cannot hot-spin (see #4144). + await this.runtimeHost.dispose(); + this.themeController.disableAutoSync(); + await this.ui.terminal.drainInput(1000); + this.stop(); + process.exit(0); + } + + // Interactive quit (Ctrl+D, Ctrl+C, /quit, extension shutdown()). Stop the + // TUI before emitting shutdown events so extension UI cleanup cannot repaint + // the final frame while the process is exiting. + // Drain any in-flight Kitty key release events before stopping. + // This prevents escape sequences from leaking to the parent shell over slow SSH. + this.themeController.disableAutoSync(); + await this.ui.terminal.drainInput(1000); + + this.stop(); + await this.runtimeHost.dispose(); + + const resumeCommand = formatResumeCommand(this.sessionManager); + if (resumeCommand) { + process.stdout.write(`${chalk.dim("To resume this session:")} ${resumeCommand}\n`); + } + + process.exit(0); + } + + private emergencyTerminalExit(): never { + this.isShuttingDown = true; + this.unregisterSignalHandlers(); + killTrackedDetachedChildren(); + // The terminal is gone. Do not run normal shutdown because TUI and + // extension cleanup can write restore sequences and re-trigger EIO. + process.exit(129); + } + + /** + * Last-resort handler for uncaught exceptions. The TUI puts stdin into raw + * mode and hides the cursor; without this handler, an uncaught throw from + * anywhere (e.g. an extension's async `ChildProcess.on("exit")` callback) + * tears down the process while leaving the terminal in raw mode with no + * cursor, requiring `stty sane && reset` to recover. + * + * Unlike emergencyTerminalExit, the terminal is still alive here, so we + * call ui.stop() to restore cooked mode, the cursor, and disable bracketed + * paste / Kitty / modifyOtherKeys sequences. + */ + private uncaughtCrash(error: Error): never { + if (this.isShuttingDown) { + process.exit(1); + } + this.isShuttingDown = true; + try { + this.unregisterSignalHandlers(); + } catch {} + try { + killTrackedDetachedChildren(); + } catch {} + try { + this.ui.stop(); + } catch {} + console.error("cactus exiting due to uncaughtException:"); + console.error(error); + process.exit(1); + } + + /** + * Check if shutdown was requested and perform shutdown if so. + */ + private async checkShutdownRequested(): Promise { + if (!this.shutdownRequested) return; + await this.shutdown(); + } + + private registerSignalHandlers(): void { + this.unregisterSignalHandlers(); + + const signals: NodeJS.Signals[] = ["SIGTERM"]; + if (process.platform !== "win32") { + signals.push("SIGHUP"); + } + + for (const signal of signals) { + const handler = () => { + // SIGHUP no longer hard-exits: graceful shutdown emits session_shutdown + // first, then attempts terminal restore. A genuinely dead terminal + // surfaces as an EIO on the restore writes, which the stdout/stderr + // error handler converts into emergencyTerminalExit (see #4144, #5080). + killTrackedDetachedChildren(); + void this.shutdown({ fromSignal: true }); + }; + process.prependListener(signal, handler); + this.signalCleanupHandlers.push(() => process.off(signal, handler)); + } + + const terminalErrorHandler = (error: Error) => { + if (isDeadTerminalError(error)) { + this.emergencyTerminalExit(); + } + throw error; + }; + process.stdout.on("error", terminalErrorHandler); + process.stderr.on("error", terminalErrorHandler); + this.signalCleanupHandlers.push(() => process.stdout.off("error", terminalErrorHandler)); + this.signalCleanupHandlers.push(() => process.stderr.off("error", terminalErrorHandler)); + + // Restore the terminal before the process dies on any uncaught throw. + // Without this, an unhandled exception from extension code (or anywhere + // in pi) leaves the terminal in raw mode with no cursor. + const uncaughtExceptionHandler = (error: Error) => this.uncaughtCrash(error); + process.prependListener("uncaughtException", uncaughtExceptionHandler); + this.signalCleanupHandlers.push(() => process.off("uncaughtException", uncaughtExceptionHandler)); + } + + private unregisterSignalHandlers(): void { + for (const cleanup of this.signalCleanupHandlers) { + cleanup(); + } + this.signalCleanupHandlers = []; + } + + private handleCtrlZ(): void { + if (process.platform === "win32") { + this.showStatus("Suspend to background is not supported on Windows"); + return; + } + + // Keep the event loop alive while suspended. Without this, stopping the TUI + // can leave Node with no ref'ed handles, causing the process to exit on fg + // before the SIGCONT handler gets a chance to restore the terminal. + const suspendKeepAlive = setInterval(() => {}, 2 ** 30); + + // Ignore SIGINT while suspended so Ctrl+C in the terminal does not + // kill the backgrounded process. The handler is removed on resume. + const ignoreSigint = () => {}; + process.on("SIGINT", ignoreSigint); + + // Set up handler to restore TUI when resumed + process.once("SIGCONT", () => { + clearInterval(suspendKeepAlive); + process.removeListener("SIGINT", ignoreSigint); + this.ui.start(); + this.ui.requestRender(true); + }); + + try { + // Stop the TUI (restore terminal to normal mode) + this.ui.stop(); + + // Send SIGTSTP to process group (pid=0 means all processes in group) + process.kill(0, "SIGTSTP"); + } catch (error) { + clearInterval(suspendKeepAlive); + process.removeListener("SIGINT", ignoreSigint); + throw error; + } + } + + private async handleFollowUp(): Promise { + const text = (this.editor.getExpandedText?.() ?? this.editor.getText()).trim(); + if (!text) return; + + // Queue input during compaction (extension commands execute immediately) + if (this.session.isCompacting) { + if (this.isExtensionCommand(text)) { + this.editor.addToHistory?.(text); + this.editor.setText(""); + await this.session.prompt(text); + } else { + this.queueCompactionMessage(text, "followUp"); + } + return; + } + + // Alt+Enter queues a follow-up message (waits until agent finishes) + // This handles extension commands (execute immediately), prompt template expansion, and queueing + if (this.session.isStreaming) { + this.editor.addToHistory?.(text); + this.editor.setText(""); + await this.session.prompt(text, { streamingBehavior: "followUp" }); + this.updatePendingMessagesDisplay(); + this.ui.requestRender(); + } + // If not streaming, Alt+Enter acts like regular Enter (trigger onSubmit) + else if (this.editor.onSubmit) { + this.editor.setText(""); + this.editor.onSubmit(text); + } + } + + private handleDequeue(): void { + const restored = this.restoreQueuedMessagesToEditor(); + if (restored === 0) { + this.showStatus("No queued messages to restore"); + } else { + this.showStatus(`Restored ${restored} queued message${restored > 1 ? "s" : ""} to editor`); + } + } + + private updateEditorBorderColor(): void { + if (this.isBashMode) { + this.editor.borderColor = theme.getBashModeBorderColor(); + } else { + const level = this.session.thinkingLevel || "off"; + this.editor.borderColor = theme.getThinkingBorderColor(level); + } + this.ui.requestRender(); + } + + private cycleThinkingLevel(): void { + const newLevel = this.session.cycleThinkingLevel(); + if (newLevel === undefined) { + this.showStatus("Current model does not support thinking"); + } else { + this.footer.invalidate(); + this.updateEditorBorderColor(); + this.showStatus(`Thinking level: ${newLevel}`); + } + } + + private async cycleModel(direction: "forward" | "backward"): Promise { + try { + const result = await this.session.cycleModel(direction); + if (result === undefined) { + const msg = this.session.scopedModels.length > 0 ? "Only one model in scope" : "Only one model available"; + this.showStatus(msg); + } else { + this.footer.invalidate(); + this.updateEditorBorderColor(); + const thinkingStr = + result.model.reasoning && result.thinkingLevel !== "off" ? ` (thinking: ${result.thinkingLevel})` : ""; + this.showStatus(`Switched to ${result.model.name || result.model.id}${thinkingStr}`); + void this.maybeWarnAboutAnthropicSubscriptionAuth(result.model); + } + } catch (error) { + this.showError(error instanceof Error ? error.message : String(error)); + } + } + + private toggleToolOutputExpansion(): void { + this.setToolsExpanded(!this.toolOutputExpanded); + } + + private setToolsExpanded(expanded: boolean): void { + this.toolOutputExpanded = expanded; + const activeHeader = this.customHeader ?? this.builtInHeader; + if (isExpandable(activeHeader)) { + activeHeader.setExpanded(expanded); + } + for (const child of this.chatContainer.children) { + if (isExpandable(child)) { + child.setExpanded(expanded); + } + } + this.ui.requestRender(); + } + + private toggleThinkingBlockVisibility(): void { + this.hideThinkingBlock = !this.hideThinkingBlock; + this.settingsManager.setHideThinkingBlock(this.hideThinkingBlock); + + // Rebuild chat from session messages + this.chatContainer.clear(); + this.rebuildChatFromMessages(); + + // If streaming, re-add the streaming component with updated visibility and re-render + if (this.streamingComponent && this.streamingMessage) { + this.streamingComponent.setHideThinkingBlock(this.hideThinkingBlock); + this.streamingComponent.updateContent(this.streamingMessage); + this.chatContainer.addChild(this.streamingComponent); + } + + this.showStatus(`Thinking blocks: ${this.hideThinkingBlock ? "hidden" : "visible"}`); + } + + private async openExternalEditor(): Promise { + // Determine editor (respect $VISUAL, then $EDITOR) + const editorCmd = process.env.VISUAL || process.env.EDITOR; + if (!editorCmd) { + this.showWarning("No editor configured. Set $VISUAL or $EDITOR environment variable."); + return; + } + + const currentText = this.editor.getExpandedText?.() ?? this.editor.getText(); + const tmpFile = path.join(os.tmpdir(), `cactus-editor-${Date.now()}.md`); + + try { + // Write current content to temp file + fs.writeFileSync(tmpFile, currentText, "utf-8"); + + // Stop TUI to release terminal + this.ui.stop(); + + // Split by space to support editor arguments (e.g., "code --wait") + const [editor, ...editorArgs] = editorCmd.split(" "); + + process.stdout.write(`Launching external editor: ${editorCmd}\nPi will resume when the editor exits.\n`); + + // Do not use spawnSync here. On Windows, synchronous child_process calls can keep + // Node/libuv's console input read active after ui.stop() pauses stdin, racing + // vim/nvim for the console input buffer until Ctrl+C cancels the pending read. + const status = await new Promise((resolve) => { + const child = spawn(editor, [...editorArgs, tmpFile], { + stdio: "inherit", + shell: process.platform === "win32", + }); + child.on("error", () => resolve(null)); + child.on("close", (code) => resolve(code)); + }); + + // On successful exit (status 0), replace editor content + if (status === 0) { + const newContent = fs.readFileSync(tmpFile, "utf-8").replace(/\n$/, ""); + this.editor.setText(newContent); + } + // On non-zero exit, keep original text (no action needed) + } finally { + // Clean up temp file + try { + fs.unlinkSync(tmpFile); + } catch { + // Ignore cleanup errors + } + + // Restart TUI + this.ui.start(); + // Force full re-render since external editor uses alternate screen + this.ui.requestRender(true); + } + } + + // ========================================================================= + // UI helpers + // ========================================================================= + + clearEditor(): void { + this.editor.setText(""); + this.ui.requestRender(); + } + + showError(errorMessage: string): void { + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("error", `Error: ${errorMessage}`), 1, 0)); + this.ui.requestRender(); + } + + showWarning(warningMessage: string): void { + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("warning", `Warning: ${warningMessage}`), 1, 0)); + this.ui.requestRender(); + } + + /** + * Get all queued messages (read-only). + * Combines session queue and compaction queue. + */ + private getAllQueuedMessages(): { steering: string[]; followUp: string[] } { + return { + steering: [ + ...this.session.getSteeringMessages(), + ...this.compactionQueuedMessages.filter((msg) => msg.mode === "steer").map((msg) => msg.text), + ], + followUp: [ + ...this.session.getFollowUpMessages(), + ...this.compactionQueuedMessages.filter((msg) => msg.mode === "followUp").map((msg) => msg.text), + ], + }; + } + + /** + * Clear all queued messages and return their contents. + * Clears both session queue and compaction queue. + */ + private clearAllQueues(): { steering: string[]; followUp: string[] } { + const { steering, followUp } = this.session.clearQueue(); + const compactionSteering = this.compactionQueuedMessages + .filter((msg) => msg.mode === "steer") + .map((msg) => msg.text); + const compactionFollowUp = this.compactionQueuedMessages + .filter((msg) => msg.mode === "followUp") + .map((msg) => msg.text); + this.compactionQueuedMessages = []; + return { + steering: [...steering, ...compactionSteering], + followUp: [...followUp, ...compactionFollowUp], + }; + } + + private updatePendingMessagesDisplay(): void { + this.pendingMessagesContainer.clear(); + const { steering: steeringMessages, followUp: followUpMessages } = this.getAllQueuedMessages(); + if (steeringMessages.length > 0 || followUpMessages.length > 0) { + this.pendingMessagesContainer.addChild(new Spacer(1)); + for (const message of steeringMessages) { + const text = theme.fg("dim", `Steering: ${message}`); + this.pendingMessagesContainer.addChild(new TruncatedText(text, 1, 0)); + } + for (const message of followUpMessages) { + const text = theme.fg("dim", `Follow-up: ${message}`); + this.pendingMessagesContainer.addChild(new TruncatedText(text, 1, 0)); + } + const dequeueHint = this.getAppKeyDisplay("app.message.dequeue"); + const hintText = theme.fg("dim", `↳ ${dequeueHint} to edit all queued messages`); + this.pendingMessagesContainer.addChild(new TruncatedText(hintText, 1, 0)); + } + } + + private restoreQueuedMessagesToEditor(options?: { abort?: boolean; currentText?: string }): number { + const { steering, followUp } = this.clearAllQueues(); + const allQueued = [...steering, ...followUp]; + if (allQueued.length === 0) { + this.updatePendingMessagesDisplay(); + if (options?.abort) { + this.agent.abort(); + } + return 0; + } + const queuedText = allQueued.join("\n\n"); + const currentText = options?.currentText ?? this.editor.getText(); + const combinedText = [queuedText, currentText].filter((t) => t.trim()).join("\n\n"); + this.editor.setText(combinedText); + this.updatePendingMessagesDisplay(); + if (options?.abort) { + this.agent.abort(); + } + return allQueued.length; + } + + private queueCompactionMessage(text: string, mode: "steer" | "followUp"): void { + this.compactionQueuedMessages.push({ text, mode }); + this.editor.addToHistory?.(text); + this.editor.setText(""); + this.updatePendingMessagesDisplay(); + this.showStatus("Queued message for after compaction"); + } + + private isExtensionCommand(text: string): boolean { + if (!text.startsWith("/")) return false; + + const extensionRunner = this.session.extensionRunner; + + const spaceIndex = text.indexOf(" "); + const commandName = spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex); + return !!extensionRunner.getCommand(commandName); + } + + private async flushCompactionQueue(options?: { willRetry?: boolean }): Promise { + if (this.compactionQueuedMessages.length === 0) { + return; + } + + const queuedMessages = [...this.compactionQueuedMessages]; + this.compactionQueuedMessages = []; + this.updatePendingMessagesDisplay(); + + const restoreQueue = (error: unknown) => { + this.session.clearQueue(); + this.compactionQueuedMessages = queuedMessages; + this.updatePendingMessagesDisplay(); + this.showError( + `Failed to send queued message${queuedMessages.length > 1 ? "s" : ""}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + }; + + try { + if (options?.willRetry) { + // When retry is pending, queue messages for the retry turn + for (const message of queuedMessages) { + if (this.isExtensionCommand(message.text)) { + await this.session.prompt(message.text); + } else if (message.mode === "followUp") { + await this.session.followUp(message.text); + } else { + await this.session.steer(message.text); + } + } + this.updatePendingMessagesDisplay(); + return; + } + + // Find first non-extension-command message to use as prompt + const firstPromptIndex = queuedMessages.findIndex((message) => !this.isExtensionCommand(message.text)); + if (firstPromptIndex === -1) { + // All extension commands - execute them all + for (const message of queuedMessages) { + await this.session.prompt(message.text); + } + return; + } + + // Execute any extension commands before the first prompt + const preCommands = queuedMessages.slice(0, firstPromptIndex); + const firstPrompt = queuedMessages[firstPromptIndex]; + const rest = queuedMessages.slice(firstPromptIndex + 1); + + for (const message of preCommands) { + await this.session.prompt(message.text); + } + + // Send first prompt (starts streaming) + const promptPromise = this.session.prompt(firstPrompt.text).catch((error) => { + restoreQueue(error); + }); + + // Queue remaining messages + for (const message of rest) { + if (this.isExtensionCommand(message.text)) { + await this.session.prompt(message.text); + } else if (message.mode === "followUp") { + await this.session.followUp(message.text); + } else { + await this.session.steer(message.text); + } + } + this.updatePendingMessagesDisplay(); + void promptPromise; + } catch (error) { + restoreQueue(error); + } + } + + /** Move pending bash components from pending area to chat */ + private flushPendingBashComponents(): void { + for (const component of this.pendingBashComponents) { + this.pendingMessagesContainer.removeChild(component); + this.chatContainer.addChild(component); + } + this.pendingBashComponents = []; + } + + // ========================================================================= + // Selectors + // ========================================================================= + + /** + * Shows a selector component in place of the editor. + * @param create Factory that receives a `done` callback and returns the component and focus target + */ + private showSelector(create: (done: () => void) => { component: Component; focus: Component }): void { + const done = () => { + this.editorContainer.clear(); + this.editorContainer.addChild(this.editor); + this.ui.setFocus(this.editor); + }; + const { component, focus } = create(done); + this.editorContainer.clear(); + this.editorContainer.addChild(component); + this.ui.setFocus(focus); + this.ui.requestRender(); + } + + private showSettingsSelector(): void { + this.showSelector((done) => { + const selector = new SettingsSelectorComponent( + { + autoCompact: this.session.autoCompactionEnabled, + showImages: this.settingsManager.getShowImages(), + imageWidthCells: this.settingsManager.getImageWidthCells(), + autoResizeImages: this.settingsManager.getImageAutoResize(), + blockImages: this.settingsManager.getBlockImages(), + enableSkillCommands: this.settingsManager.getEnableSkillCommands(), + steeringMode: this.session.steeringMode, + followUpMode: this.session.followUpMode, + transport: this.settingsManager.getTransport(), + httpIdleTimeoutMs: this.settingsManager.getHttpIdleTimeoutMs(), + thinkingLevel: this.session.thinkingLevel, + availableThinkingLevels: this.session.getAvailableThinkingLevels(), + currentTheme: this.settingsManager.getThemeSetting() || "dark", + terminalTheme: this.themeController.getTerminalTheme(), + availableThemes: getAvailableThemes(), + hideThinkingBlock: this.hideThinkingBlock, + doubleEscapeAction: this.settingsManager.getDoubleEscapeAction(), + treeFilterMode: this.settingsManager.getTreeFilterMode(), + showHardwareCursor: this.settingsManager.getShowHardwareCursor(), + defaultProjectTrust: this.settingsManager.getDefaultProjectTrust(), + editorPaddingX: this.settingsManager.getEditorPaddingX(), + autocompleteMaxVisible: this.settingsManager.getAutocompleteMaxVisible(), + quietStartup: this.settingsManager.getQuietStartup(), + clearOnShrink: this.settingsManager.getClearOnShrink(), + showTerminalProgress: this.settingsManager.getShowTerminalProgress(), + warnings: this.settingsManager.getWarnings(), + }, + { + onAutoCompactChange: (enabled) => { + this.session.setAutoCompactionEnabled(enabled); + this.footer.setAutoCompactEnabled(enabled); + }, + onShowImagesChange: (enabled) => { + this.settingsManager.setShowImages(enabled); + for (const child of this.chatContainer.children) { + if (child instanceof ToolExecutionComponent) { + child.setShowImages(enabled); + } + } + }, + onImageWidthCellsChange: (width) => { + this.settingsManager.setImageWidthCells(width); + for (const child of this.chatContainer.children) { + if (child instanceof ToolExecutionComponent) { + child.setImageWidthCells(width); + } + } + }, + onAutoResizeImagesChange: (enabled) => { + this.settingsManager.setImageAutoResize(enabled); + }, + onBlockImagesChange: (blocked) => { + this.settingsManager.setBlockImages(blocked); + }, + onEnableSkillCommandsChange: (enabled) => { + this.settingsManager.setEnableSkillCommands(enabled); + this.setupAutocompleteProvider(); + }, + onSteeringModeChange: (mode) => { + this.session.setSteeringMode(mode); + }, + onFollowUpModeChange: (mode) => { + this.session.setFollowUpMode(mode); + }, + onTransportChange: (transport) => { + this.settingsManager.setTransport(transport); + this.session.agent.transport = transport; + }, + onHttpIdleTimeoutMsChange: (timeoutMs) => { + this.settingsManager.setHttpIdleTimeoutMs(timeoutMs); + configureHttpDispatcher(timeoutMs); + this.showStatus(`HTTP idle timeout: ${formatHttpIdleTimeoutMs(timeoutMs)}`); + }, + onThinkingLevelChange: (level) => { + this.session.setThinkingLevel(level); + this.footer.invalidate(); + this.updateEditorBorderColor(); + }, + onThemeChange: (themeSetting) => { + this.settingsManager.setTheme(themeSetting); + void this.themeController.applyFromSettings(); + }, + onThemePreview: (themeName) => this.themeController.preview(themeName), + onHideThinkingBlockChange: (hidden) => { + this.hideThinkingBlock = hidden; + this.settingsManager.setHideThinkingBlock(hidden); + for (const child of this.chatContainer.children) { + if (child instanceof AssistantMessageComponent) { + child.setHideThinkingBlock(hidden); + } + } + this.chatContainer.clear(); + this.rebuildChatFromMessages(); + }, + onQuietStartupChange: (enabled) => { + this.settingsManager.setQuietStartup(enabled); + }, + onDefaultProjectTrustChange: (defaultProjectTrust) => { + this.settingsManager.setDefaultProjectTrust(defaultProjectTrust); + }, + onDoubleEscapeActionChange: (action) => { + this.settingsManager.setDoubleEscapeAction(action); + }, + onTreeFilterModeChange: (mode) => { + this.settingsManager.setTreeFilterMode(mode); + }, + onShowHardwareCursorChange: (enabled) => { + this.settingsManager.setShowHardwareCursor(enabled); + this.ui.setShowHardwareCursor(enabled); + }, + onEditorPaddingXChange: (padding) => { + this.settingsManager.setEditorPaddingX(padding); + this.defaultEditor.setPaddingX(padding); + if (this.editor !== this.defaultEditor && this.editor.setPaddingX !== undefined) { + this.editor.setPaddingX(padding); + } + }, + onAutocompleteMaxVisibleChange: (maxVisible) => { + this.settingsManager.setAutocompleteMaxVisible(maxVisible); + this.defaultEditor.setAutocompleteMaxVisible(maxVisible); + if (this.editor !== this.defaultEditor && this.editor.setAutocompleteMaxVisible !== undefined) { + this.editor.setAutocompleteMaxVisible(maxVisible); + } + }, + onClearOnShrinkChange: (enabled) => { + this.settingsManager.setClearOnShrink(enabled); + this.ui.setClearOnShrink(enabled); + }, + onShowTerminalProgressChange: (enabled) => { + this.settingsManager.setShowTerminalProgress(enabled); + }, + onWarningsChange: (warnings) => { + this.settingsManager.setWarnings(warnings); + }, + onCancel: () => { + done(); + this.ui.requestRender(); + }, + }, + ); + return { component: selector, focus: selector.getSettingsList() }; + }); + } + + private async handleModelCommand(searchTerm?: string): Promise { + if (!searchTerm) { + this.showModelSelector(); + return; + } + + const model = await this.findExactModelMatch(searchTerm); + if (model) { + try { + await this.session.setModel(model); + this.footer.invalidate(); + this.updateEditorBorderColor(); + this.showStatus(`Model: ${model.id}`); + void this.maybeWarnAboutAnthropicSubscriptionAuth(model); + } catch (error) { + this.showError(error instanceof Error ? error.message : String(error)); + } + return; + } + + this.showModelSelector(searchTerm); + } + + private async findExactModelMatch(searchTerm: string): Promise | undefined> { + const models = await this.getModelCandidates(); + return findExactModelReferenceMatch(searchTerm, models); + } + + private async getModelCandidates(): Promise[]> { + if (this.session.scopedModels.length > 0) { + return this.session.scopedModels.map((scoped) => scoped.model); + } + + this.session.modelRegistry.refresh(); + try { + return await this.session.modelRegistry.getAvailable(); + } catch { + return []; + } + } + + /** Update the footer's available provider count from current model candidates */ + private async updateAvailableProviderCount(): Promise { + const models = await this.getModelCandidates(); + const uniqueProviders = new Set(models.map((m) => m.provider)); + this.footerDataProvider.setAvailableProviderCount(uniqueProviders.size); + } + + private async maybeWarnAboutAnthropicSubscriptionAuth( + model: Model | undefined = this.session.model, + ): Promise { + if (this.settingsManager.getWarnings().anthropicExtraUsage === false) { + return; + } + if (this.anthropicSubscriptionWarningShown) { + return; + } + if (!model || model.provider !== "anthropic") { + return; + } + + const storedCredential = this.session.modelRegistry.authStorage.get("anthropic"); + if (storedCredential?.type === "oauth") { + this.anthropicSubscriptionWarningShown = true; + this.showWarning(ANTHROPIC_SUBSCRIPTION_AUTH_WARNING); + return; + } + + try { + const apiKey = await this.session.modelRegistry.getApiKeyForProvider(model.provider); + if (!isAnthropicSubscriptionAuthKey(apiKey)) { + return; + } + this.anthropicSubscriptionWarningShown = true; + this.showWarning(ANTHROPIC_SUBSCRIPTION_AUTH_WARNING); + } catch { + // Ignore auth lookup failures for warning-only checks. + } + } + + private maybeSaveImplicitProjectTrustAfterReload(): boolean { + const cwd = this.sessionManager.getCwd(); + if (this.autoTrustOnReloadCwd !== cwd) { + return false; + } + if (!this.settingsManager.isProjectTrusted() || !hasTrustRequiringProjectResources(cwd)) { + return false; + } + + const trustStore = new ProjectTrustStore(this.runtimeHost.services.agentDir); + try { + if (trustStore.get(cwd) !== null) { + this.autoTrustOnReloadCwd = undefined; + return false; + } + trustStore.set(cwd, true); + this.autoTrustOnReloadCwd = undefined; + return true; + } catch (error) { + this.showWarning( + `Could not save project trust after reload: ${error instanceof Error ? error.message : String(error)}`, + ); + return false; + } + } + + private showTrustSelector(): void { + const cwd = this.sessionManager.getCwd(); + const trustStore = new ProjectTrustStore(this.runtimeHost.services.agentDir); + const savedDecision = trustStore.getEntry(cwd); + this.showSelector((done) => { + const selector = new TrustSelectorComponent({ + cwd, + savedDecision, + projectTrusted: this.settingsManager.isProjectTrusted(), + onSelect: (selection) => { + trustStore.setMany(selection.updates); + done(); + this.showStatus( + `Saved trust decision: ${selection.trusted ? "trusted" : "untrusted"}. Restart Cactus for this to take effect.`, + ); + }, + onCancel: () => { + done(); + this.ui.requestRender(); + }, + }); + return { component: selector, focus: selector }; + }); + } + + private showModelSelector(initialSearchInput?: string): void { + this.showSelector((done) => { + const selector = new ModelSelectorComponent( + this.ui, + this.session.model, + this.settingsManager, + this.session.modelRegistry, + this.session.scopedModels, + async (model) => { + try { + await this.session.setModel(model); + this.footer.invalidate(); + this.updateEditorBorderColor(); + done(); + this.showStatus(`Model: ${model.id}`); + void this.maybeWarnAboutAnthropicSubscriptionAuth(model); + } catch (error) { + done(); + this.showError(error instanceof Error ? error.message : String(error)); + } + }, + () => { + done(); + this.ui.requestRender(); + }, + initialSearchInput, + ); + return { component: selector, focus: selector }; + }); + } + + private async showModelsSelector(): Promise { + // Get all available models + this.session.modelRegistry.refresh(); + const allModels = this.session.modelRegistry.getAvailable(); + + if (allModels.length === 0) { + this.showStatus("No models available"); + return; + } + + // Check if session has scoped models (from previous session-only changes or CLI --models) + const sessionScopedModels = this.session.scopedModels; + const hasSessionScope = sessionScopedModels.length > 0; + + // Build enabled model IDs from session state or settings + let currentEnabledIds: string[] | null = null; + + if (hasSessionScope) { + // Use current session's scoped models + currentEnabledIds = sessionScopedModels.map((scoped) => `${scoped.model.provider}/${scoped.model.id}`); + } else { + // Fall back to settings + const patterns = this.settingsManager.getEnabledModels(); + if (patterns !== undefined && patterns.length > 0) { + const scopedModels = await resolveModelScope(patterns, this.session.modelRegistry); + currentEnabledIds = scopedModels.map((scoped) => `${scoped.model.provider}/${scoped.model.id}`); + } + } + + // Helper to update session's scoped models (session-only, no persist) + const updateSessionModels = async (enabledIds: string[] | null) => { + currentEnabledIds = enabledIds === null ? null : [...enabledIds]; + if (enabledIds && enabledIds.length > 0 && enabledIds.length < allModels.length) { + const newScopedModels = await resolveModelScope(enabledIds, this.session.modelRegistry); + this.session.setScopedModels( + newScopedModels.map((sm) => ({ + model: sm.model, + thinkingLevel: sm.thinkingLevel, + })), + ); + } else { + // All enabled or none enabled = no filter + this.session.setScopedModels([]); + } + await this.updateAvailableProviderCount(); + this.ui.requestRender(); + }; + + this.showSelector((done) => { + const selector = new ScopedModelsSelectorComponent( + { + allModels, + enabledModelIds: currentEnabledIds, + }, + { + onChange: async (enabledIds) => { + await updateSessionModels(enabledIds); + }, + onPersist: (enabledIds) => { + // Persist to settings + const newPatterns = + enabledIds === null || enabledIds.length === allModels.length + ? undefined // All enabled = clear filter + : enabledIds; + this.settingsManager.setEnabledModels(newPatterns ? [...newPatterns] : undefined); + this.showStatus("Model selection saved to settings"); + }, + onCancel: () => { + done(); + this.ui.requestRender(); + }, + }, + ); + return { component: selector, focus: selector }; + }); + } + + private showUserMessageSelector(): void { + const userMessages = this.session.getUserMessagesForForking(); + + if (userMessages.length === 0) { + this.showStatus("No messages to fork from"); + return; + } + + const initialSelectedId = userMessages[userMessages.length - 1]?.entryId; + + this.showSelector((done) => { + const selector = new UserMessageSelectorComponent( + userMessages.map((m) => ({ id: m.entryId, text: m.text })), + async (entryId) => { + try { + const result = await this.runtimeHost.fork(entryId); + if (result.cancelled) { + done(); + this.ui.requestRender(); + return; + } + + this.editor.setText(result.selectedText ?? ""); + done(); + this.showStatus("Forked to new session"); + } catch (error: unknown) { + done(); + this.showError(error instanceof Error ? error.message : String(error)); + } + }, + () => { + done(); + this.ui.requestRender(); + }, + initialSelectedId, + ); + return { component: selector, focus: selector.getMessageList() }; + }); + } + + private async handleCloneCommand(): Promise { + const leafId = this.sessionManager.getLeafId(); + if (!leafId) { + this.showStatus("Nothing to clone yet"); + return; + } + + try { + const result = await this.runtimeHost.fork(leafId, { position: "at" }); + if (result.cancelled) { + this.ui.requestRender(); + return; + } + + this.editor.setText(""); + this.showStatus("Cloned to new session"); + } catch (error: unknown) { + this.showError(error instanceof Error ? error.message : String(error)); + } + } + + private showTreeSelector(initialSelectedId?: string): void { + const tree = this.sessionManager.getTree(); + const realLeafId = this.sessionManager.getLeafId(); + const initialFilterMode = this.settingsManager.getTreeFilterMode(); + + if (tree.length === 0) { + this.showStatus("No entries in session"); + return; + } + + this.showSelector((done) => { + const selector = new TreeSelectorComponent( + tree, + realLeafId, + this.ui.terminal.rows, + async (entryId) => { + // Selecting the current leaf is a no-op (already there) + if (entryId === realLeafId) { + done(); + this.showStatus("Already at this point"); + return; + } + + // Ask about summarization + done(); // Close selector first + + // Loop until user makes a complete choice or cancels to tree + let wantsSummary = false; + let customInstructions: string | undefined; + + // Check if we should skip the prompt (user preference to always default to no summary) + if (!this.settingsManager.getBranchSummarySkipPrompt()) { + while (true) { + const summaryChoice = await this.showExtensionSelector("Summarize branch?", [ + "No summary", + "Summarize", + "Summarize with custom prompt", + ]); + + if (summaryChoice === undefined) { + // User pressed escape - re-show tree selector with same selection + this.showTreeSelector(entryId); + return; + } + + wantsSummary = summaryChoice !== "No summary"; + + if (summaryChoice === "Summarize with custom prompt") { + customInstructions = await this.showExtensionEditor("Custom summarization instructions"); + if (customInstructions === undefined) { + // User cancelled - loop back to summary selector + continue; + } + } + + // User made a complete choice + break; + } + } + + // Set up escape handler and loader if summarizing + let summaryLoader: Loader | undefined; + const originalOnEscape = this.defaultEditor.onEscape; + + if (wantsSummary) { + this.defaultEditor.onEscape = () => { + this.session.abortBranchSummary(); + }; + this.chatContainer.addChild(new Spacer(1)); + summaryLoader = new Loader( + this.ui, + (spinner) => theme.fg("accent", spinner), + (text) => theme.fg("muted", text), + `Summarizing branch... (${keyText("app.interrupt")} to cancel)`, + ); + this.statusContainer.addChild(summaryLoader); + this.ui.requestRender(); + } + + try { + const result = await this.session.navigateTree(entryId, { + summarize: wantsSummary, + customInstructions, + }); + + if (result.aborted) { + // Summarization aborted - re-show tree selector with same selection + this.showStatus("Branch summarization cancelled"); + this.showTreeSelector(entryId); + return; + } + if (result.cancelled) { + this.showStatus("Navigation cancelled"); + return; + } + + // Update UI + this.chatContainer.clear(); + this.renderInitialMessages(); + if (result.editorText && !this.editor.getText().trim()) { + this.editor.setText(result.editorText); + } + this.showStatus("Navigated to selected point"); + void this.flushCompactionQueue({ willRetry: false }); + } catch (error) { + this.showError(error instanceof Error ? error.message : String(error)); + } finally { + if (summaryLoader) { + summaryLoader.stop(); + this.statusContainer.clear(); + } + this.defaultEditor.onEscape = originalOnEscape; + } + }, + () => { + done(); + this.ui.requestRender(); + }, + (entryId, label) => { + this.sessionManager.appendLabelChange(entryId, label); + this.ui.requestRender(); + }, + initialSelectedId, + initialFilterMode, + ); + return { component: selector, focus: selector }; + }); + } + + private showSessionSelector(): void { + this.showSelector((done) => { + const selector = new SessionSelectorComponent( + (onProgress) => + SessionManager.list(this.sessionManager.getCwd(), this.sessionManager.getSessionDir(), onProgress), + (onProgress) => + this.sessionManager.usesDefaultSessionDir() + ? SessionManager.listAll(onProgress) + : SessionManager.listAll(this.sessionManager.getSessionDir(), onProgress), + async (sessionPath) => { + done(); + await this.handleResumeSession(sessionPath); + }, + () => { + done(); + this.ui.requestRender(); + }, + () => { + void this.shutdown(); + }, + () => this.ui.requestRender(), + { + renameSession: async (sessionFilePath: string, nextName: string | undefined) => { + const next = (nextName ?? "").trim(); + if (!next) return; + const mgr = SessionManager.open(sessionFilePath); + mgr.appendSessionInfo(next); + }, + showRenameHint: true, + keybindings: this.keybindings, + }, + + this.sessionManager.getSessionFile(), + ); + return { component: selector, focus: selector }; + }); + } + + private async handleResumeSession( + sessionPath: string, + options?: Parameters[1], + ): Promise<{ cancelled: boolean }> { + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.statusContainer.clear(); + try { + const result = await this.runtimeHost.switchSession(sessionPath, { + withSession: options?.withSession, + projectTrustContextFactory: (cwd) => this.createProjectTrustContext(cwd), + }); + if (result.cancelled) { + return result; + } + this.showStatus("Resumed session"); + return result; + } catch (error: unknown) { + if (error instanceof MissingSessionCwdError) { + const selectedCwd = await this.promptForMissingSessionCwd(error); + if (!selectedCwd) { + this.showStatus("Resume cancelled"); + return { cancelled: true }; + } + const result = await this.runtimeHost.switchSession(sessionPath, { + cwdOverride: selectedCwd, + withSession: options?.withSession, + projectTrustContextFactory: (cwd) => this.createProjectTrustContext(cwd), + }); + if (result.cancelled) { + return result; + } + this.showStatus("Resumed session in current cwd"); + return result; + } + return this.handleFatalRuntimeError("Failed to resume session", error); + } + } + + // ========================================================================= + // Command handlers + // ========================================================================= + + private async handleReloadCommand(): Promise { + if (this.session.isStreaming) { + this.showWarning("Wait for the current response to finish before reloading."); + return; + } + if (this.session.isCompacting) { + this.showWarning("Wait for compaction to finish before reloading."); + return; + } + + this.resetExtensionUI(); + + const reloadBox = new Container(); + const borderColor = (s: string) => theme.fg("border", s); + reloadBox.addChild(new DynamicBorder(borderColor)); + reloadBox.addChild(new Spacer(1)); + reloadBox.addChild( + new Text(theme.fg("muted", "Reloading keybindings, extensions, skills, prompts, themes..."), 1, 0), + ); + reloadBox.addChild(new Spacer(1)); + reloadBox.addChild(new DynamicBorder(borderColor)); + + const previousEditor = this.editor; + this.editorContainer.clear(); + this.editorContainer.addChild(reloadBox); + this.ui.setFocus(reloadBox); + this.ui.requestRender(true); + await new Promise((resolve) => process.nextTick(resolve)); + + const dismissReloadBox = (editor: Component) => { + this.editorContainer.clear(); + this.editorContainer.addChild(editor); + this.ui.setFocus(editor); + this.ui.requestRender(); + }; + + let chatRestoredBeforeSessionStart = false; + let reloadBoxDismissed = false; + const restoreChatBeforeSessionStart = () => { + if (chatRestoredBeforeSessionStart) { + return; + } + this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock(); + this.rebuildChatFromMessages(); + chatRestoredBeforeSessionStart = true; + }; + + try { + await this.session.reload({ beforeSessionStart: restoreChatBeforeSessionStart }); + restoreChatBeforeSessionStart(); + configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs()); + this.keybindings.reload(); + const activeHeader = this.customHeader ?? this.builtInHeader; + if (isExpandable(activeHeader)) { + activeHeader.setExpanded(this.toolOutputExpanded); + } + setRegisteredThemes(this.session.resourceLoader.getThemes().themes); + await this.themeController.applyFromSettings(); + const editorPaddingX = this.settingsManager.getEditorPaddingX(); + const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible(); + this.defaultEditor.setPaddingX(editorPaddingX); + this.defaultEditor.setAutocompleteMaxVisible(autocompleteMaxVisible); + if (this.editor !== this.defaultEditor) { + this.editor.setPaddingX?.(editorPaddingX); + this.editor.setAutocompleteMaxVisible?.(autocompleteMaxVisible); + } + this.ui.setShowHardwareCursor(this.settingsManager.getShowHardwareCursor()); + this.ui.setClearOnShrink(this.settingsManager.getClearOnShrink()); + this.setupAutocompleteProvider(); + const runner = this.session.extensionRunner; + this.setupExtensionShortcuts(runner); + this.showLoadedResources({ + force: false, + showDiagnosticsWhenQuiet: true, + }); + const savedImplicitProjectTrust = this.maybeSaveImplicitProjectTrustAfterReload(); + const modelsJsonError = this.session.modelRegistry.getError(); + if (modelsJsonError) { + this.showError(`models.json error: ${modelsJsonError}`); + } + this.showStatus( + savedImplicitProjectTrust + ? "Reloaded keybindings, extensions, skills, prompts, themes; saved project trust" + : "Reloaded keybindings, extensions, skills, prompts, themes", + ); + dismissReloadBox(this.editor as Component); + reloadBoxDismissed = true; + } catch (error) { + if (!reloadBoxDismissed) { + dismissReloadBox(previousEditor as Component); + } + this.showError(`Reload failed: ${error instanceof Error ? error.message : String(error)}`); + } + } + + private async handleExportCommand(text: string): Promise { + const outputPath = this.getPathCommandArgument(text, "/export"); + + try { + const filePath = this.session.exportToJsonl(outputPath); + this.showStatus(`Session exported to: ${filePath}`); + } catch (error: unknown) { + this.showError(`Failed to export session: ${error instanceof Error ? error.message : "Unknown error"}`); + } + } + + private getPathCommandArgument(text: string, command: "/export" | "/import"): string | undefined { + if (text === command) { + return undefined; + } + if (!text.startsWith(`${command} `)) { + return undefined; + } + + const argsString = text.slice(command.length + 1).trimStart(); + if (!argsString) { + return undefined; + } + + const firstChar = argsString[0]; + if (firstChar === '"' || firstChar === "'") { + const closingQuoteIndex = argsString.indexOf(firstChar, 1); + if (closingQuoteIndex < 0) { + return undefined; + } + return argsString.slice(1, closingQuoteIndex); + } + + const firstWhitespaceIndex = argsString.search(/\s/); + if (firstWhitespaceIndex < 0) { + return argsString; + } + return argsString.slice(0, firstWhitespaceIndex); + } + + private async handleImportCommand(text: string): Promise { + const inputPath = this.getPathCommandArgument(text, "/import"); + if (!inputPath) { + this.showError("Usage: /import "); + return; + } + + const confirmed = await this.showExtensionConfirm("Import session", `Replace current session with ${inputPath}?`); + if (!confirmed) { + this.showStatus("Import cancelled"); + return; + } + + try { + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.statusContainer.clear(); + const result = await this.runtimeHost.importFromJsonl(inputPath); + if (result.cancelled) { + this.showStatus("Import cancelled"); + return; + } + this.showStatus(`Session imported from: ${inputPath}`); + } catch (error: unknown) { + if (error instanceof MissingSessionCwdError) { + const selectedCwd = await this.promptForMissingSessionCwd(error); + if (!selectedCwd) { + this.showStatus("Import cancelled"); + return; + } + const result = await this.runtimeHost.importFromJsonl(inputPath, selectedCwd); + if (result.cancelled) { + this.showStatus("Import cancelled"); + return; + } + this.showStatus(`Session imported from: ${inputPath}`); + return; + } + if (error instanceof SessionImportFileNotFoundError) { + this.showError(`Failed to import session: ${error.message}`); + return; + } + await this.handleFatalRuntimeError("Failed to import session", error); + } + } + + private async handleCopyCommand(): Promise { + const text = this.session.getLastAssistantText(); + if (!text) { + this.showError("No agent messages to copy yet."); + return; + } + + try { + await copyToClipboard(text); + this.showStatus("Copied last agent message to clipboard"); + } catch (error) { + this.showError(error instanceof Error ? error.message : String(error)); + } + } + + private handleNameCommand(text: string): void { + const name = text.replace(/^\/name\s*/, "").trim(); + if (!name) { + const currentName = this.sessionManager.getSessionName(); + if (currentName) { + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("dim", `Session name: ${currentName}`), 1, 0)); + } else { + this.showWarning("Usage: /name "); + } + this.ui.requestRender(); + return; + } + + this.session.setSessionName(name); + const sessionName = this.sessionManager.getSessionName(); + if (sessionName !== name) { + this.showWarning(`Session name was normalized from ${JSON.stringify(name)} to ${JSON.stringify(sessionName)}`); + } + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("dim", `Session name set: ${sessionName ?? name}`), 1, 0)); + this.ui.requestRender(); + } + + private handleSessionCommand(): void { + const stats = this.session.getSessionStats(); + const sessionName = this.sessionManager.getSessionName(); + + let info = `${theme.bold("Session Info")}\n\n`; + if (sessionName) { + info += `${theme.fg("dim", "Name:")} ${sessionName}\n`; + } + info += `${theme.fg("dim", "File:")} ${stats.sessionFile ?? "In-memory"}\n`; + info += `${theme.fg("dim", "ID:")} ${stats.sessionId}\n\n`; + info += `${theme.bold("Messages")}\n`; + info += `${theme.fg("dim", "User:")} ${stats.userMessages}\n`; + info += `${theme.fg("dim", "Assistant:")} ${stats.assistantMessages}\n`; + info += `${theme.fg("dim", "Tool Calls:")} ${stats.toolCalls}\n`; + info += `${theme.fg("dim", "Tool Results:")} ${stats.toolResults}\n`; + info += `${theme.fg("dim", "Total:")} ${stats.totalMessages}\n\n`; + info += `${theme.bold("Tokens")}\n`; + info += `${theme.fg("dim", "Input:")} ${stats.tokens.input.toLocaleString()}\n`; + info += `${theme.fg("dim", "Output:")} ${stats.tokens.output.toLocaleString()}\n`; + if (stats.tokens.cacheRead > 0) { + info += `${theme.fg("dim", "Cache Read:")} ${stats.tokens.cacheRead.toLocaleString()}\n`; + } + if (stats.tokens.cacheWrite > 0) { + info += `${theme.fg("dim", "Cache Write:")} ${stats.tokens.cacheWrite.toLocaleString()}\n`; + } + info += `${theme.fg("dim", "Total:")} ${stats.tokens.total.toLocaleString()}\n`; + + if (stats.cost > 0) { + info += `\n${theme.bold("Cost")}\n`; + info += `${theme.fg("dim", "Total:")} ${stats.cost.toFixed(4)}`; + } + + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(info, 1, 0)); + this.ui.requestRender(); + } + + /** + * Get capitalized display string for an app keybinding action. + */ + private getAppKeyDisplay(action: AppKeybinding): string { + return keyDisplayText(action); + } + + /** + * Get capitalized display string for an editor keybinding action. + */ + private getEditorKeyDisplay(action: Keybinding): string { + return keyDisplayText(action); + } + + private handleHotkeysCommand(): void { + // Navigation keybindings + const cursorUp = this.getEditorKeyDisplay("tui.editor.cursorUp"); + const cursorDown = this.getEditorKeyDisplay("tui.editor.cursorDown"); + const cursorLeft = this.getEditorKeyDisplay("tui.editor.cursorLeft"); + const cursorRight = this.getEditorKeyDisplay("tui.editor.cursorRight"); + const cursorWordLeft = this.getEditorKeyDisplay("tui.editor.cursorWordLeft"); + const cursorWordRight = this.getEditorKeyDisplay("tui.editor.cursorWordRight"); + const cursorLineStart = this.getEditorKeyDisplay("tui.editor.cursorLineStart"); + const cursorLineEnd = this.getEditorKeyDisplay("tui.editor.cursorLineEnd"); + const jumpForward = this.getEditorKeyDisplay("tui.editor.jumpForward"); + const jumpBackward = this.getEditorKeyDisplay("tui.editor.jumpBackward"); + const pageUp = this.getEditorKeyDisplay("tui.editor.pageUp"); + const pageDown = this.getEditorKeyDisplay("tui.editor.pageDown"); + + // Editing keybindings + const submit = this.getEditorKeyDisplay("tui.input.submit"); + const newLine = this.getEditorKeyDisplay("tui.input.newLine"); + const deleteWordBackward = this.getEditorKeyDisplay("tui.editor.deleteWordBackward"); + const deleteWordForward = this.getEditorKeyDisplay("tui.editor.deleteWordForward"); + const deleteToLineStart = this.getEditorKeyDisplay("tui.editor.deleteToLineStart"); + const deleteToLineEnd = this.getEditorKeyDisplay("tui.editor.deleteToLineEnd"); + const yank = this.getEditorKeyDisplay("tui.editor.yank"); + const yankPop = this.getEditorKeyDisplay("tui.editor.yankPop"); + const undo = this.getEditorKeyDisplay("tui.editor.undo"); + const tab = this.getEditorKeyDisplay("tui.input.tab"); + + // App keybindings + const interrupt = this.getAppKeyDisplay("app.interrupt"); + const clear = this.getAppKeyDisplay("app.clear"); + const exit = this.getAppKeyDisplay("app.exit"); + const suspend = this.getAppKeyDisplay("app.suspend"); + const cycleThinkingLevel = this.getAppKeyDisplay("app.thinking.cycle"); + const cycleModelForward = this.getAppKeyDisplay("app.model.cycleForward"); + const selectModel = this.getAppKeyDisplay("app.model.select"); + const expandTools = this.getAppKeyDisplay("app.tools.expand"); + const toggleThinking = this.getAppKeyDisplay("app.thinking.toggle"); + const externalEditor = this.getAppKeyDisplay("app.editor.external"); + const cycleModelBackward = this.getAppKeyDisplay("app.model.cycleBackward"); + const followUp = this.getAppKeyDisplay("app.message.followUp"); + const dequeue = this.getAppKeyDisplay("app.message.dequeue"); + const pasteImage = this.getAppKeyDisplay("app.clipboard.pasteImage"); + + let hotkeys = ` +**Navigation** +| Key | Action | +|-----|--------| +| \`${cursorUp}\` / \`${cursorDown}\` / \`${cursorLeft}\` / \`${cursorRight}\` | Move cursor / browse history | +| \`${cursorWordLeft}\` / \`${cursorWordRight}\` | Move by word | +| \`${cursorLineStart}\` | Start of line | +| \`${cursorLineEnd}\` | End of line | +| \`${jumpForward}\` | Jump forward to character | +| \`${jumpBackward}\` | Jump backward to character | +| \`${pageUp}\` / \`${pageDown}\` | Scroll by page | + +**Editing** +| Key | Action | +|-----|--------| +| \`${submit}\` | Send message | +| \`${newLine}\` | New line${process.platform === "win32" ? " (Ctrl+Enter on Windows Terminal)" : ""} | +| \`${deleteWordBackward}\` | Delete word backwards | +| \`${deleteWordForward}\` | Delete word forwards | +| \`${deleteToLineStart}\` | Delete to start of line | +| \`${deleteToLineEnd}\` | Delete to end of line | +| \`${yank}\` | Paste the most-recently-deleted text | +| \`${yankPop}\` | Cycle through the deleted text after pasting | +| \`${undo}\` | Undo | + +**Other** +| Key | Action | +|-----|--------| +| \`${tab}\` | Path completion / accept autocomplete | +| \`${interrupt}\` | Cancel autocomplete / abort streaming | +| \`${clear}\` | Clear editor (first) / exit (second) | +| \`${exit}\` | Exit (when editor is empty) | +| \`${suspend}\` | Suspend to background | +| \`${cycleThinkingLevel}\` | Cycle thinking level | +| \`${cycleModelForward}\` / \`${cycleModelBackward}\` | Cycle models | +| \`${selectModel}\` | Open model selector | +| \`${expandTools}\` | Toggle tool output expansion | +| \`${toggleThinking}\` | Toggle thinking block visibility | +| \`${externalEditor}\` | Edit message in external editor | +| \`${followUp}\` | Queue follow-up message | +| \`${dequeue}\` | Restore queued messages | +| \`${pasteImage}\` | Paste image from clipboard | +| \`/\` | Slash commands | +| \`!\` | Run bash command | +| \`!!\` | Run bash command (excluded from context) | +`; + + // Add extension-registered shortcuts + const extensionRunner = this.session.extensionRunner; + const shortcuts = extensionRunner.getShortcuts(this.keybindings.getEffectiveConfig()); + if (shortcuts.size > 0) { + hotkeys += ` +**Extensions** +| Key | Action | +|-----|--------| +`; + for (const [key, shortcut] of shortcuts) { + const description = shortcut.description ?? shortcut.extensionPath; + const keyDisplay = formatKeyText(key, { capitalize: true }); + hotkeys += `| \`${keyDisplay}\` | ${description} |\n`; + } + } + + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new DynamicBorder()); + this.chatContainer.addChild(new Text(theme.bold(theme.fg("accent", "Keyboard Shortcuts")), 1, 0)); + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Markdown(hotkeys.trim(), 1, 1, this.getMarkdownThemeWithSettings())); + this.chatContainer.addChild(new DynamicBorder()); + this.ui.requestRender(); + } + + private async handleClearCommand(): Promise { + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.statusContainer.clear(); + try { + const result = await this.runtimeHost.newSession(); + if (result.cancelled) { + return; + } + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1)); + this.ui.requestRender(); + } catch (error: unknown) { + await this.handleFatalRuntimeError("Failed to create session", error); + } + } + + private handleDebugCommand(): void { + const width = this.ui.terminal.columns; + const height = this.ui.terminal.rows; + const allLines = this.ui.render(width); + + const debugLogPath = getDebugLogPath(); + const debugData = [ + `Debug output at ${new Date().toISOString()}`, + `Terminal: ${width}x${height}`, + `Total lines: ${allLines.length}`, + "", + "=== All rendered lines with visible widths ===", + ...allLines.map((line, idx) => { + const vw = visibleWidth(line); + const escaped = JSON.stringify(line); + return `[${idx}] (w=${vw}) ${escaped}`; + }), + "", + "=== Agent messages (JSONL) ===", + ...this.session.messages.map((msg) => JSON.stringify(msg)), + "", + ].join("\n"); + + fs.mkdirSync(path.dirname(debugLogPath), { recursive: true }); + fs.writeFileSync(debugLogPath, debugData); + + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild( + new Text(`${theme.fg("accent", "✓ Debug log written")}\n${theme.fg("muted", debugLogPath)}`, 1, 1), + ); + this.ui.requestRender(); + } + + private async handleBashCommand(command: string, excludeFromContext = false): Promise { + const extensionRunner = this.session.extensionRunner; + + // Emit user_bash event to let extensions intercept + const eventResult = await extensionRunner.emitUserBash({ + type: "user_bash", + command, + excludeFromContext, + cwd: this.sessionManager.getCwd(), + }); + + // If extension returned a full result, use it directly + if (eventResult?.result) { + const result = eventResult.result; + + // Create UI component for display + this.bashComponent = new BashExecutionComponent(command, this.ui, excludeFromContext); + if (this.session.isStreaming) { + this.pendingMessagesContainer.addChild(this.bashComponent); + this.pendingBashComponents.push(this.bashComponent); + } else { + this.chatContainer.addChild(this.bashComponent); + } + + // Show output and complete + if (result.output) { + this.bashComponent.appendOutput(result.output); + } + this.bashComponent.setComplete( + result.exitCode, + result.cancelled, + result.truncated ? ({ truncated: true, content: result.output } as TruncationResult) : undefined, + result.fullOutputPath, + ); + + // Record the result in session + this.session.recordBashResult(command, result, { excludeFromContext }); + this.bashComponent = undefined; + this.ui.requestRender(); + return; + } + + // Normal execution path (possibly with custom operations) + const isDeferred = this.session.isStreaming; + this.bashComponent = new BashExecutionComponent(command, this.ui, excludeFromContext); + + if (isDeferred) { + // Show in pending area when agent is streaming + this.pendingMessagesContainer.addChild(this.bashComponent); + this.pendingBashComponents.push(this.bashComponent); + } else { + // Show in chat immediately when agent is idle + this.chatContainer.addChild(this.bashComponent); + } + this.ui.requestRender(); + + try { + const result = await this.session.executeBash( + command, + (chunk) => { + if (this.bashComponent) { + this.bashComponent.appendOutput(chunk); + this.ui.requestRender(); + } + }, + { excludeFromContext, operations: eventResult?.operations }, + ); + + if (this.bashComponent) { + this.bashComponent.setComplete( + result.exitCode, + result.cancelled, + result.truncated ? ({ truncated: true, content: result.output } as TruncationResult) : undefined, + result.fullOutputPath, + ); + } + } catch (error) { + if (this.bashComponent) { + this.bashComponent.setComplete(undefined, false); + } + this.showError(`Bash command failed: ${error instanceof Error ? error.message : "Unknown error"}`); + } + + this.bashComponent = undefined; + this.ui.requestRender(); + } + + private async handleCompactCommand(customInstructions?: string): Promise { + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.statusContainer.clear(); + + try { + await this.session.compact(customInstructions); + } catch { + // Ignore, will be emitted as an event + } + } + + stop(): void { + if (this.settingsManager.getShowTerminalProgress()) { + this.ui.terminal.setProgress(false); + } + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = undefined; + } + this.themeController.disableAutoSync(); + this.clearExtensionTerminalInputListeners(); + this.footer.dispose(); + this.footerDataProvider.dispose(); + if (this.unsubscribe) { + this.unsubscribe(); + } + if (this.isInitialized) { + this.ui.stop(); + this.isInitialized = false; + } + this.unregisterSignalHandlers(); + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/model-search.ts b/cactus-code/packages/coding-agent/src/modes/interactive/model-search.ts new file mode 100644 index 000000000..bab9c5a5b --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/model-search.ts @@ -0,0 +1,21 @@ +export interface ModelSearchItem { + id: string; + provider: string; + name?: string; +} + +export function getModelSearchText(item: ModelSearchItem): string { + const { id, provider } = item; + const name = item.name ? ` ${item.name}` : ""; + return `${id} ${provider} ${provider}/${id} ${provider} ${id}${name}`; +} + +/** + * The /model selector search should rank exact provider-prefixed queries before proxy-provider IDs + * like openrouter/openai/gpt-5, so keep the bare model ID out of the leading position. + */ +export function getModelSelectorSearchText(item: ModelSearchItem): string { + const { id, provider } = item; + const name = item.name ? ` ${item.name}` : ""; + return `${provider} ${provider}/${id} ${provider} ${id}${name}`; +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/theme/dark.json b/cactus-code/packages/coding-agent/src/modes/interactive/theme/dark.json new file mode 100644 index 000000000..c133e232f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/theme/dark.json @@ -0,0 +1,86 @@ +{ + "$schema": "https://raw.githubusercontent.com/earendil-works/pi/main/packages/coding-agent/src/modes/interactive/theme/theme-schema.json", + "name": "dark", + "vars": { + "cyan": "#00d7ff", + "blue": "#5f87ff", + "green": "#b5bd68", + "red": "#cc6666", + "yellow": "#ffff00", + "text": "#d4d4d4", + "gray": "#808080", + "dimGray": "#666666", + "darkGray": "#505050", + "accent": "#8abeb7", + "selectedBg": "#3a3a4a", + "userMsgBg": "#343541", + "toolPendingBg": "#282832", + "toolSuccessBg": "#283228", + "toolErrorBg": "#3c2828", + "customMsgBg": "#2d2838" + }, + "colors": { + "accent": "accent", + "border": "blue", + "borderAccent": "cyan", + "borderMuted": "darkGray", + "success": "green", + "error": "red", + "warning": "yellow", + "muted": "gray", + "dim": "dimGray", + "text": "text", + "thinkingText": "gray", + + "selectedBg": "selectedBg", + "userMessageBg": "userMsgBg", + "userMessageText": "text", + "customMessageBg": "customMsgBg", + "customMessageText": "text", + "customMessageLabel": "#9575cd", + "toolPendingBg": "toolPendingBg", + "toolSuccessBg": "toolSuccessBg", + "toolErrorBg": "toolErrorBg", + "toolTitle": "text", + "toolOutput": "gray", + + "mdHeading": "#f0c674", + "mdLink": "#81a2be", + "mdLinkUrl": "dimGray", + "mdCode": "accent", + "mdCodeBlock": "green", + "mdCodeBlockBorder": "gray", + "mdQuote": "gray", + "mdQuoteBorder": "gray", + "mdHr": "gray", + "mdListBullet": "accent", + + "toolDiffAdded": "green", + "toolDiffRemoved": "red", + "toolDiffContext": "gray", + + "syntaxComment": "#6A9955", + "syntaxKeyword": "#569CD6", + "syntaxFunction": "#DCDCAA", + "syntaxVariable": "#9CDCFE", + "syntaxString": "#CE9178", + "syntaxNumber": "#B5CEA8", + "syntaxType": "#4EC9B0", + "syntaxOperator": "#D4D4D4", + "syntaxPunctuation": "#D4D4D4", + + "thinkingOff": "darkGray", + "thinkingMinimal": "#6e6e6e", + "thinkingLow": "#5f87af", + "thinkingMedium": "#81a2be", + "thinkingHigh": "#b294bb", + "thinkingXhigh": "#d183e8", + + "bashMode": "green" + }, + "export": { + "pageBg": "#18181e", + "cardBg": "#1e1e24", + "infoBg": "#3c3728" + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/theme/light.json b/cactus-code/packages/coding-agent/src/modes/interactive/theme/light.json new file mode 100644 index 000000000..0a417b566 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/theme/light.json @@ -0,0 +1,85 @@ +{ + "$schema": "https://raw.githubusercontent.com/earendil-works/pi/main/packages/coding-agent/src/modes/interactive/theme/theme-schema.json", + "name": "light", + "vars": { + "teal": "#5a8080", + "blue": "#547da7", + "green": "#588458", + "red": "#aa5555", + "yellow": "#9a7326", + "text": "#1f2328", + "mediumGray": "#6c6c6c", + "dimGray": "#767676", + "lightGray": "#b0b0b0", + "selectedBg": "#d0d0e0", + "userMsgBg": "#e8e8e8", + "toolPendingBg": "#e8e8f0", + "toolSuccessBg": "#e8f0e8", + "toolErrorBg": "#f0e8e8", + "customMsgBg": "#ede7f6" + }, + "colors": { + "accent": "teal", + "border": "blue", + "borderAccent": "teal", + "borderMuted": "lightGray", + "success": "green", + "error": "red", + "warning": "yellow", + "muted": "mediumGray", + "dim": "dimGray", + "text": "text", + "thinkingText": "mediumGray", + + "selectedBg": "selectedBg", + "userMessageBg": "userMsgBg", + "userMessageText": "text", + "customMessageBg": "customMsgBg", + "customMessageText": "text", + "customMessageLabel": "#7e57c2", + "toolPendingBg": "toolPendingBg", + "toolSuccessBg": "toolSuccessBg", + "toolErrorBg": "toolErrorBg", + "toolTitle": "text", + "toolOutput": "mediumGray", + + "mdHeading": "yellow", + "mdLink": "blue", + "mdLinkUrl": "dimGray", + "mdCode": "teal", + "mdCodeBlock": "green", + "mdCodeBlockBorder": "mediumGray", + "mdQuote": "mediumGray", + "mdQuoteBorder": "mediumGray", + "mdHr": "mediumGray", + "mdListBullet": "green", + + "toolDiffAdded": "green", + "toolDiffRemoved": "red", + "toolDiffContext": "mediumGray", + + "syntaxComment": "#008000", + "syntaxKeyword": "#0000FF", + "syntaxFunction": "#795E26", + "syntaxVariable": "#001080", + "syntaxString": "#A31515", + "syntaxNumber": "#098658", + "syntaxType": "#267F99", + "syntaxOperator": "#000000", + "syntaxPunctuation": "#000000", + + "thinkingOff": "lightGray", + "thinkingMinimal": "#767676", + "thinkingLow": "blue", + "thinkingMedium": "teal", + "thinkingHigh": "#875f87", + "thinkingXhigh": "#8b008b", + + "bashMode": "green" + }, + "export": { + "pageBg": "#f8f8f8", + "cardBg": "#ffffff", + "infoBg": "#fffae6" + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme-controller.ts b/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme-controller.ts new file mode 100644 index 000000000..43ad620d3 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme-controller.ts @@ -0,0 +1,126 @@ +import type { TUI } from "@earendil-works/pi-tui"; +import type { SettingsManager } from "../../../core/settings-manager.ts"; +import { + detectTerminalBackgroundFromEnv, + detectTerminalBackgroundTheme, + detectTerminalThemeForAuto, + initTheme, + parseAutoThemeSetting, + resolveThemeSetting, + setTheme, + setThemeInstance, + type TerminalTheme, + type Theme, +} from "./theme.ts"; + +type ThemeResult = { success: boolean; error?: string }; + +export class InteractiveThemeController { + private readonly ui: TUI; + private readonly settingsManager: SettingsManager; + private readonly showError: (message: string) => void; + private readonly onChanged: () => void; + private terminalTheme: TerminalTheme = detectTerminalBackgroundFromEnv().theme; + private activeThemeName: string | undefined; + private autoSyncEnabled = false; + + constructor(ui: TUI, settingsManager: SettingsManager, showError: (message: string) => void, onChanged: () => void) { + this.ui = ui; + this.settingsManager = settingsManager; + this.showError = showError; + this.onChanged = onChanged; + this.activeThemeName = resolveThemeSetting(this.settingsManager.getThemeSetting(), this.terminalTheme); + initTheme(this.activeThemeName, true); + this.ui.onTerminalColorSchemeChange((terminalTheme) => this.applyTerminalTheme(terminalTheme)); + } + + async applyFromSettings(): Promise { + const themeSetting = this.settingsManager.getThemeSetting(); + const autoTheme = parseAutoThemeSetting(themeSetting); + if (autoTheme) { + this.terminalTheme = await detectTerminalThemeForAuto({ ui: this.ui, timeoutMs: 100 }); + this.setAutoSync(true); + this.applyThemeName(this.terminalTheme === "light" ? autoTheme.lightTheme : autoTheme.darkTheme, true); + return; + } + + this.setAutoSync(false); + if (themeSetting !== undefined) { + this.applyThemeName(themeSetting, true); + return; + } + + const detection = await detectTerminalBackgroundTheme({ ui: this.ui, timeoutMs: 100 }); + this.terminalTheme = detection.theme; + if (!this.applyThemeName(detection.theme).success) return; + if (detection.confidence === "high") { + this.settingsManager.setTheme(detection.theme); + await this.settingsManager.flush(); + } + } + + setThemeName(themeName: string, showError = false): ThemeResult { + this.setAutoSync(false); + return this.applyThemeName(themeName, showError); + } + + setThemeInstance(themeInstance: Theme): ThemeResult { + this.setAutoSync(false); + setThemeInstance(themeInstance); + this.activeThemeName = ""; + this.notifyChanged(); + return { success: true }; + } + + preview(themeSettingOrName: string): void { + const themeName = resolveThemeSetting(themeSettingOrName, this.terminalTheme) ?? this.activeThemeName; + if (!themeName) return; + if (setTheme(themeName, true).success) { + this.ui.invalidate(); + this.ui.requestRender(); + } + } + + disableAutoSync(): void { + this.setAutoSync(false); + } + + getTerminalTheme(): TerminalTheme { + return this.terminalTheme; + } + + private applyThemeName(themeName: string, showError = false): ThemeResult { + const result = setTheme(themeName, true); + this.activeThemeName = result.success ? themeName : "dark"; + this.notifyChanged(); + if (!result.success && showError) { + this.showError(`Failed to load theme "${themeName}": ${result.error}\nFell back to dark theme.`); + } + return result; + } + + private notifyChanged(): void { + this.ui.invalidate(); + this.onChanged(); + } + + private setAutoSync(enabled: boolean): void { + if (this.autoSyncEnabled === enabled) return; + this.autoSyncEnabled = enabled; + this.ui.setTerminalColorSchemeNotifications(enabled); + } + + private applyTerminalTheme(terminalTheme: TerminalTheme): void { + if (!this.autoSyncEnabled) return; + this.terminalTheme = terminalTheme; + const autoTheme = parseAutoThemeSetting(this.settingsManager.getThemeSetting()); + if (!autoTheme) { + this.setAutoSync(false); + return; + } + const themeName = terminalTheme === "light" ? autoTheme.lightTheme : autoTheme.darkTheme; + if (themeName !== this.activeThemeName) { + this.applyThemeName(themeName); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme-schema.json b/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme-schema.json new file mode 100644 index 000000000..9d94a12a7 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme-schema.json @@ -0,0 +1,336 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "Pi Coding Agent Theme", + "description": "Theme schema for Pi coding agent", + "type": "object", + "required": ["name", "colors"], + "properties": { + "$schema": { + "type": "string", + "description": "JSON schema reference" + }, + "name": { + "type": "string", + "pattern": "^[^/]+$", + "description": "Theme name. Must not contain '/' because it is reserved for automatic light/dark theme settings." + }, + "vars": { + "type": "object", + "description": "Reusable color variables", + "additionalProperties": { + "oneOf": [ + { + "type": "string", + "description": "Hex color (#RRGGBB), variable reference, or empty string for terminal default" + }, + { + "type": "integer", + "minimum": 0, + "maximum": 255, + "description": "256-color palette index (0-255)" + } + ] + } + }, + "colors": { + "type": "object", + "description": "Theme color definitions (all required)", + "required": [ + "accent", + "border", + "borderAccent", + "borderMuted", + "success", + "error", + "warning", + "muted", + "dim", + "text", + "thinkingText", + "selectedBg", + "userMessageBg", + "userMessageText", + "customMessageBg", + "customMessageText", + "customMessageLabel", + "toolPendingBg", + "toolSuccessBg", + "toolErrorBg", + "toolTitle", + "toolOutput", + "mdHeading", + "mdLink", + "mdLinkUrl", + "mdCode", + "mdCodeBlock", + "mdCodeBlockBorder", + "mdQuote", + "mdQuoteBorder", + "mdHr", + "mdListBullet", + "toolDiffAdded", + "toolDiffRemoved", + "toolDiffContext", + "syntaxComment", + "syntaxKeyword", + "syntaxFunction", + "syntaxVariable", + "syntaxString", + "syntaxNumber", + "syntaxType", + "syntaxOperator", + "syntaxPunctuation", + "thinkingOff", + "thinkingMinimal", + "thinkingLow", + "thinkingMedium", + "thinkingHigh", + "thinkingXhigh", + "bashMode" + ], + "properties": { + "accent": { + "$ref": "#/$defs/colorValue", + "description": "Primary accent color (logo, selected items, cursor)" + }, + "border": { + "$ref": "#/$defs/colorValue", + "description": "Normal borders" + }, + "borderAccent": { + "$ref": "#/$defs/colorValue", + "description": "Highlighted borders" + }, + "borderMuted": { + "$ref": "#/$defs/colorValue", + "description": "Subtle borders" + }, + "success": { + "$ref": "#/$defs/colorValue", + "description": "Success states" + }, + "error": { + "$ref": "#/$defs/colorValue", + "description": "Error states" + }, + "warning": { + "$ref": "#/$defs/colorValue", + "description": "Warning states" + }, + "muted": { + "$ref": "#/$defs/colorValue", + "description": "Secondary/dimmed text" + }, + "dim": { + "$ref": "#/$defs/colorValue", + "description": "Very dimmed text (more subtle than muted)" + }, + "text": { + "$ref": "#/$defs/colorValue", + "description": "Default text color (usually empty string)" + }, + "thinkingText": { + "$ref": "#/$defs/colorValue", + "description": "Thinking block text color" + }, + "selectedBg": { + "$ref": "#/$defs/colorValue", + "description": "Selected item background" + }, + "userMessageBg": { + "$ref": "#/$defs/colorValue", + "description": "User message background" + }, + "userMessageText": { + "$ref": "#/$defs/colorValue", + "description": "User message text color" + }, + "customMessageBg": { + "$ref": "#/$defs/colorValue", + "description": "Custom message background (hook-injected messages)" + }, + "customMessageText": { + "$ref": "#/$defs/colorValue", + "description": "Custom message text color" + }, + "customMessageLabel": { + "$ref": "#/$defs/colorValue", + "description": "Custom message type label color" + }, + "toolPendingBg": { + "$ref": "#/$defs/colorValue", + "description": "Tool execution box (pending state)" + }, + "toolSuccessBg": { + "$ref": "#/$defs/colorValue", + "description": "Tool execution box (success state)" + }, + "toolErrorBg": { + "$ref": "#/$defs/colorValue", + "description": "Tool execution box (error state)" + }, + "toolTitle": { + "$ref": "#/$defs/colorValue", + "description": "Tool execution box title color" + }, + "toolOutput": { + "$ref": "#/$defs/colorValue", + "description": "Tool execution box output text color" + }, + "mdHeading": { + "$ref": "#/$defs/colorValue", + "description": "Markdown heading text" + }, + "mdLink": { + "$ref": "#/$defs/colorValue", + "description": "Markdown link text" + }, + "mdLinkUrl": { + "$ref": "#/$defs/colorValue", + "description": "Markdown link URL" + }, + "mdCode": { + "$ref": "#/$defs/colorValue", + "description": "Markdown inline code" + }, + "mdCodeBlock": { + "$ref": "#/$defs/colorValue", + "description": "Markdown code block content" + }, + "mdCodeBlockBorder": { + "$ref": "#/$defs/colorValue", + "description": "Markdown code block fences" + }, + "mdQuote": { + "$ref": "#/$defs/colorValue", + "description": "Markdown blockquote text" + }, + "mdQuoteBorder": { + "$ref": "#/$defs/colorValue", + "description": "Markdown blockquote border" + }, + "mdHr": { + "$ref": "#/$defs/colorValue", + "description": "Markdown horizontal rule" + }, + "mdListBullet": { + "$ref": "#/$defs/colorValue", + "description": "Markdown list bullets/numbers" + }, + "toolDiffAdded": { + "$ref": "#/$defs/colorValue", + "description": "Added lines in tool diffs" + }, + "toolDiffRemoved": { + "$ref": "#/$defs/colorValue", + "description": "Removed lines in tool diffs" + }, + "toolDiffContext": { + "$ref": "#/$defs/colorValue", + "description": "Context lines in tool diffs" + }, + "syntaxComment": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: comments" + }, + "syntaxKeyword": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: keywords" + }, + "syntaxFunction": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: function names" + }, + "syntaxVariable": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: variable names" + }, + "syntaxString": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: string literals" + }, + "syntaxNumber": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: number literals" + }, + "syntaxType": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: type names" + }, + "syntaxOperator": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: operators" + }, + "syntaxPunctuation": { + "$ref": "#/$defs/colorValue", + "description": "Syntax highlighting: punctuation" + }, + "thinkingOff": { + "$ref": "#/$defs/colorValue", + "description": "Thinking level border: off" + }, + "thinkingMinimal": { + "$ref": "#/$defs/colorValue", + "description": "Thinking level border: minimal" + }, + "thinkingLow": { + "$ref": "#/$defs/colorValue", + "description": "Thinking level border: low" + }, + "thinkingMedium": { + "$ref": "#/$defs/colorValue", + "description": "Thinking level border: medium" + }, + "thinkingHigh": { + "$ref": "#/$defs/colorValue", + "description": "Thinking level border: high" + }, + "thinkingXhigh": { + "$ref": "#/$defs/colorValue", + "description": "Thinking level border: xhigh (OpenAI codex-max only)" + }, + "bashMode": { + "$ref": "#/$defs/colorValue", + "description": "Editor border color in bash mode" + } + }, + "additionalProperties": false + }, + "export": { + "type": "object", + "description": "Optional colors for HTML export (defaults derived from userMessageBg if not specified)", + "properties": { + "pageBg": { + "$ref": "#/$defs/colorValue", + "description": "Page background color" + }, + "cardBg": { + "$ref": "#/$defs/colorValue", + "description": "Card/container background color" + }, + "infoBg": { + "$ref": "#/$defs/colorValue", + "description": "Info sections background (system prompt, notices)" + } + }, + "additionalProperties": false + } + }, + "additionalProperties": false, + "$defs": { + "colorValue": { + "oneOf": [ + { + "type": "string", + "description": "Hex color (#RRGGBB), variable reference, or empty string for terminal default" + }, + { + "type": "integer", + "minimum": 0, + "maximum": 255, + "description": "256-color palette index (0-255)" + } + ] + } + } +} diff --git a/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme.ts b/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme.ts new file mode 100644 index 000000000..676bc5299 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/interactive/theme/theme.ts @@ -0,0 +1,1284 @@ +import * as fs from "node:fs"; +import * as path from "node:path"; +import { + type EditorTheme, + getCapabilities, + type MarkdownTheme, + type RgbColor, + type SelectListTheme, + type SettingsListTheme, +} from "@earendil-works/pi-tui"; +import chalk from "chalk"; +import { type Static, Type } from "typebox"; +import { Compile } from "typebox/compile"; +import { getCustomThemesDir, getThemesDir } from "../../../config.ts"; +import type { SourceInfo } from "../../../core/source-info.ts"; +import { closeWatcher, watchWithErrorHandler } from "../../../utils/fs-watch.ts"; +import { highlight, supportsLanguage } from "../../../utils/syntax-highlight.ts"; + +// ============================================================================ +// Types & Schema +// ============================================================================ + +const ColorValueSchema = Type.Union([ + Type.String(), // hex "#ff0000", var ref "primary", or empty "" + Type.Integer({ minimum: 0, maximum: 255 }), // 256-color index +]); + +type ColorValue = Static; + +const ThemeJsonSchema = Type.Object({ + $schema: Type.Optional(Type.String()), + name: Type.String(), + vars: Type.Optional(Type.Record(Type.String(), ColorValueSchema)), + colors: Type.Object({ + // Core UI (10 colors) + accent: ColorValueSchema, + border: ColorValueSchema, + borderAccent: ColorValueSchema, + borderMuted: ColorValueSchema, + success: ColorValueSchema, + error: ColorValueSchema, + warning: ColorValueSchema, + muted: ColorValueSchema, + dim: ColorValueSchema, + text: ColorValueSchema, + thinkingText: ColorValueSchema, + // Backgrounds & Content Text (11 colors) + selectedBg: ColorValueSchema, + userMessageBg: ColorValueSchema, + userMessageText: ColorValueSchema, + customMessageBg: ColorValueSchema, + customMessageText: ColorValueSchema, + customMessageLabel: ColorValueSchema, + toolPendingBg: ColorValueSchema, + toolSuccessBg: ColorValueSchema, + toolErrorBg: ColorValueSchema, + toolTitle: ColorValueSchema, + toolOutput: ColorValueSchema, + // Markdown (10 colors) + mdHeading: ColorValueSchema, + mdLink: ColorValueSchema, + mdLinkUrl: ColorValueSchema, + mdCode: ColorValueSchema, + mdCodeBlock: ColorValueSchema, + mdCodeBlockBorder: ColorValueSchema, + mdQuote: ColorValueSchema, + mdQuoteBorder: ColorValueSchema, + mdHr: ColorValueSchema, + mdListBullet: ColorValueSchema, + // Tool Diffs (3 colors) + toolDiffAdded: ColorValueSchema, + toolDiffRemoved: ColorValueSchema, + toolDiffContext: ColorValueSchema, + // Syntax Highlighting (9 colors) + syntaxComment: ColorValueSchema, + syntaxKeyword: ColorValueSchema, + syntaxFunction: ColorValueSchema, + syntaxVariable: ColorValueSchema, + syntaxString: ColorValueSchema, + syntaxNumber: ColorValueSchema, + syntaxType: ColorValueSchema, + syntaxOperator: ColorValueSchema, + syntaxPunctuation: ColorValueSchema, + // Thinking Level Borders (6 colors) + thinkingOff: ColorValueSchema, + thinkingMinimal: ColorValueSchema, + thinkingLow: ColorValueSchema, + thinkingMedium: ColorValueSchema, + thinkingHigh: ColorValueSchema, + thinkingXhigh: ColorValueSchema, + // Bash Mode (1 color) + bashMode: ColorValueSchema, + }), + export: Type.Optional( + Type.Object({ + pageBg: Type.Optional(ColorValueSchema), + cardBg: Type.Optional(ColorValueSchema), + infoBg: Type.Optional(ColorValueSchema), + }), + ), +}); + +type ThemeJson = Static; + +const validateThemeJson = Compile(ThemeJsonSchema); + +export type ThemeColor = + | "accent" + | "border" + | "borderAccent" + | "borderMuted" + | "success" + | "error" + | "warning" + | "muted" + | "dim" + | "text" + | "thinkingText" + | "userMessageText" + | "customMessageText" + | "customMessageLabel" + | "toolTitle" + | "toolOutput" + | "mdHeading" + | "mdLink" + | "mdLinkUrl" + | "mdCode" + | "mdCodeBlock" + | "mdCodeBlockBorder" + | "mdQuote" + | "mdQuoteBorder" + | "mdHr" + | "mdListBullet" + | "toolDiffAdded" + | "toolDiffRemoved" + | "toolDiffContext" + | "syntaxComment" + | "syntaxKeyword" + | "syntaxFunction" + | "syntaxVariable" + | "syntaxString" + | "syntaxNumber" + | "syntaxType" + | "syntaxOperator" + | "syntaxPunctuation" + | "thinkingOff" + | "thinkingMinimal" + | "thinkingLow" + | "thinkingMedium" + | "thinkingHigh" + | "thinkingXhigh" + | "bashMode"; + +export type ThemeBg = + | "selectedBg" + | "userMessageBg" + | "customMessageBg" + | "toolPendingBg" + | "toolSuccessBg" + | "toolErrorBg"; + +type ColorMode = "truecolor" | "256color"; + +// ============================================================================ +// Color Utilities +// ============================================================================ + +function hexToRgb(hex: string): { r: number; g: number; b: number } { + const cleaned = hex.replace("#", ""); + if (cleaned.length !== 6) { + throw new Error(`Invalid hex color: ${hex}`); + } + const r = parseInt(cleaned.substring(0, 2), 16); + const g = parseInt(cleaned.substring(2, 4), 16); + const b = parseInt(cleaned.substring(4, 6), 16); + if (Number.isNaN(r) || Number.isNaN(g) || Number.isNaN(b)) { + throw new Error(`Invalid hex color: ${hex}`); + } + return { r, g, b }; +} + +// The 6x6x6 color cube channel values (indices 0-5) +const CUBE_VALUES = [0, 95, 135, 175, 215, 255]; + +// Grayscale ramp values (indices 232-255, 24 grays from 8 to 238) +const GRAY_VALUES = Array.from({ length: 24 }, (_, i) => 8 + i * 10); + +function findClosestCubeIndex(value: number): number { + let minDist = Infinity; + let minIdx = 0; + for (let i = 0; i < CUBE_VALUES.length; i++) { + const dist = Math.abs(value - CUBE_VALUES[i]); + if (dist < minDist) { + minDist = dist; + minIdx = i; + } + } + return minIdx; +} + +function findClosestGrayIndex(gray: number): number { + let minDist = Infinity; + let minIdx = 0; + for (let i = 0; i < GRAY_VALUES.length; i++) { + const dist = Math.abs(gray - GRAY_VALUES[i]); + if (dist < minDist) { + minDist = dist; + minIdx = i; + } + } + return minIdx; +} + +function colorDistance(r1: number, g1: number, b1: number, r2: number, g2: number, b2: number): number { + // Weighted Euclidean distance (human eye is more sensitive to green) + const dr = r1 - r2; + const dg = g1 - g2; + const db = b1 - b2; + return dr * dr * 0.299 + dg * dg * 0.587 + db * db * 0.114; +} + +function rgbTo256(r: number, g: number, b: number): number { + // Find closest color in the 6x6x6 cube + const rIdx = findClosestCubeIndex(r); + const gIdx = findClosestCubeIndex(g); + const bIdx = findClosestCubeIndex(b); + const cubeR = CUBE_VALUES[rIdx]; + const cubeG = CUBE_VALUES[gIdx]; + const cubeB = CUBE_VALUES[bIdx]; + const cubeIndex = 16 + 36 * rIdx + 6 * gIdx + bIdx; + const cubeDist = colorDistance(r, g, b, cubeR, cubeG, cubeB); + + // Find closest grayscale + const gray = Math.round(0.299 * r + 0.587 * g + 0.114 * b); + const grayIdx = findClosestGrayIndex(gray); + const grayValue = GRAY_VALUES[grayIdx]; + const grayIndex = 232 + grayIdx; + const grayDist = colorDistance(r, g, b, grayValue, grayValue, grayValue); + + // Check if color has noticeable saturation (hue matters) + // If max-min spread is significant, prefer cube to preserve tint + const maxC = Math.max(r, g, b); + const minC = Math.min(r, g, b); + const spread = maxC - minC; + + // Only consider grayscale if color is nearly neutral (spread < 10) + // AND grayscale is actually closer + if (spread < 10 && grayDist < cubeDist) { + return grayIndex; + } + + return cubeIndex; +} + +function hexTo256(hex: string): number { + const { r, g, b } = hexToRgb(hex); + return rgbTo256(r, g, b); +} + +function fgAnsi(color: string | number, mode: ColorMode): string { + if (color === "") return "\x1b[39m"; + if (typeof color === "number") return `\x1b[38;5;${color}m`; + if (color.startsWith("#")) { + if (mode === "truecolor") { + const { r, g, b } = hexToRgb(color); + return `\x1b[38;2;${r};${g};${b}m`; + } else { + const index = hexTo256(color); + return `\x1b[38;5;${index}m`; + } + } + throw new Error(`Invalid color value: ${color}`); +} + +function bgAnsi(color: string | number, mode: ColorMode): string { + if (color === "") return "\x1b[49m"; + if (typeof color === "number") return `\x1b[48;5;${color}m`; + if (color.startsWith("#")) { + if (mode === "truecolor") { + const { r, g, b } = hexToRgb(color); + return `\x1b[48;2;${r};${g};${b}m`; + } else { + const index = hexTo256(color); + return `\x1b[48;5;${index}m`; + } + } + throw new Error(`Invalid color value: ${color}`); +} + +function resolveVarRefs( + value: ColorValue, + vars: Record, + visited = new Set(), +): string | number { + if (typeof value === "number" || value === "" || value.startsWith("#")) { + return value; + } + if (visited.has(value)) { + throw new Error(`Circular variable reference detected: ${value}`); + } + if (!(value in vars)) { + throw new Error(`Variable reference not found: ${value}`); + } + visited.add(value); + return resolveVarRefs(vars[value], vars, visited); +} + +function resolveThemeColors>( + colors: T, + vars: Record = {}, +): Record { + const resolved: Record = {}; + for (const [key, value] of Object.entries(colors)) { + resolved[key] = resolveVarRefs(value, vars); + } + return resolved as Record; +} + +// ============================================================================ +// Theme Class +// ============================================================================ + +export class Theme { + readonly name?: string; + readonly sourcePath?: string; + sourceInfo?: SourceInfo; + private fgColors: Map; + private bgColors: Map; + private mode: ColorMode; + + constructor( + fgColors: Record, + bgColors: Record, + mode: ColorMode, + options: { name?: string; sourcePath?: string; sourceInfo?: SourceInfo } = {}, + ) { + this.name = options.name; + this.sourcePath = options.sourcePath; + this.sourceInfo = options.sourceInfo; + this.mode = mode; + this.fgColors = new Map(); + for (const [key, value] of Object.entries(fgColors) as [ThemeColor, string | number][]) { + this.fgColors.set(key, fgAnsi(value, mode)); + } + this.bgColors = new Map(); + for (const [key, value] of Object.entries(bgColors) as [ThemeBg, string | number][]) { + this.bgColors.set(key, bgAnsi(value, mode)); + } + } + + fg(color: ThemeColor, text: string): string { + const ansi = this.fgColors.get(color); + if (!ansi) throw new Error(`Unknown theme color: ${color}`); + return `${ansi}${text}\x1b[39m`; // Reset only foreground color + } + + bg(color: ThemeBg, text: string): string { + const ansi = this.bgColors.get(color); + if (!ansi) throw new Error(`Unknown theme background color: ${color}`); + return `${ansi}${text}\x1b[49m`; // Reset only background color + } + + bold(text: string): string { + return chalk.bold(text); + } + + italic(text: string): string { + return chalk.italic(text); + } + + underline(text: string): string { + return chalk.underline(text); + } + + inverse(text: string): string { + return chalk.inverse(text); + } + + strikethrough(text: string): string { + return chalk.strikethrough(text); + } + + getFgAnsi(color: ThemeColor): string { + const ansi = this.fgColors.get(color); + if (!ansi) throw new Error(`Unknown theme color: ${color}`); + return ansi; + } + + getBgAnsi(color: ThemeBg): string { + const ansi = this.bgColors.get(color); + if (!ansi) throw new Error(`Unknown theme background color: ${color}`); + return ansi; + } + + getColorMode(): ColorMode { + return this.mode; + } + + getThinkingBorderColor(level: "off" | "minimal" | "low" | "medium" | "high" | "xhigh"): (str: string) => string { + // Map thinking levels to dedicated theme colors + switch (level) { + case "off": + return (str: string) => this.fg("thinkingOff", str); + case "minimal": + return (str: string) => this.fg("thinkingMinimal", str); + case "low": + return (str: string) => this.fg("thinkingLow", str); + case "medium": + return (str: string) => this.fg("thinkingMedium", str); + case "high": + return (str: string) => this.fg("thinkingHigh", str); + case "xhigh": + return (str: string) => this.fg("thinkingXhigh", str); + default: + return (str: string) => this.fg("thinkingOff", str); + } + } + + getBashModeBorderColor(): (str: string) => string { + return (str: string) => this.fg("bashMode", str); + } +} + +// ============================================================================ +// Theme Loading +// ============================================================================ + +let BUILTIN_THEMES: Record | undefined; + +function getBuiltinThemes(): Record { + if (!BUILTIN_THEMES) { + const themesDir = getThemesDir(); + const darkPath = path.join(themesDir, "dark.json"); + const lightPath = path.join(themesDir, "light.json"); + BUILTIN_THEMES = { + dark: JSON.parse(fs.readFileSync(darkPath, "utf-8")) as ThemeJson, + light: JSON.parse(fs.readFileSync(lightPath, "utf-8")) as ThemeJson, + }; + } + return BUILTIN_THEMES; +} + +export function getAvailableThemes(): string[] { + return getAvailableThemesWithPaths().map(({ name }) => name); +} + +export interface ThemeInfo { + name: string; + path: string | undefined; +} + +export function getAvailableThemesWithPaths(): ThemeInfo[] { + const themesDir = getThemesDir(); + const result: ThemeInfo[] = []; + const seen = new Set(); + const addTheme = (themeInfo: ThemeInfo) => { + if (seen.has(themeInfo.name)) { + return; + } + seen.add(themeInfo.name); + result.push(themeInfo); + }; + + // Built-in themes + for (const name of Object.keys(getBuiltinThemes())) { + addTheme({ name, path: path.join(themesDir, `${name}.json`) }); + } + + // Custom themes + for (const themeInfo of getCustomThemeInfos()) { + addTheme(themeInfo); + } + + for (const [name, theme] of registeredThemes.entries()) { + addTheme({ name, path: theme.sourcePath }); + } + + return result.sort((a, b) => a.name.localeCompare(b.name)); +} + +function getCustomThemeInfos(): ThemeInfo[] { + const customThemesDir = getCustomThemesDir(); + const result: ThemeInfo[] = []; + if (!fs.existsSync(customThemesDir)) { + return result; + } + + for (const file of fs.readdirSync(customThemesDir)) { + if (!file.endsWith(".json")) { + continue; + } + const themePath = path.join(customThemesDir, file); + try { + const customTheme = loadThemeFromPath(themePath); + if (customTheme.name) { + result.push({ name: customTheme.name, path: themePath }); + } + } catch { + // Invalid themes are ignored here; the resource loader reports them + // during normal startup/reload. + } + } + return result; +} + +function assertThemeNameIsValid(name: string): void { + if (name.includes("/")) { + throw new Error( + `Invalid theme name "${name}": theme names cannot contain "/" because it is reserved for automatic light/dark theme settings.`, + ); + } +} + +function parseThemeJson(label: string, json: unknown): ThemeJson { + if (!validateThemeJson.Check(json)) { + const errors = Array.from(validateThemeJson.Errors(json)); + const missingColors = new Set(); + const otherErrors: string[] = []; + + for (const error of errors) { + if (error.keyword === "required" && error.instancePath === "/colors") { + const requiredProperties = (error.params as { requiredProperties?: string[] }).requiredProperties; + for (const requiredProperty of requiredProperties ?? []) { + missingColors.add(requiredProperty); + } + continue; + } + + const path = error.instancePath || "/"; + otherErrors.push(` - ${path}: ${error.message}`); + } + + let errorMessage = `Invalid theme "${label}":\n`; + if (missingColors.size > 0) { + errorMessage += "\nMissing required color tokens:\n"; + errorMessage += Array.from(missingColors) + .sort() + .map((color) => ` - ${color}`) + .join("\n"); + errorMessage += '\n\nPlease add these colors to your theme\'s "colors" object.'; + errorMessage += "\nSee the built-in themes (dark.json, light.json) for reference values."; + } + if (otherErrors.length > 0) { + errorMessage += `\n\nOther errors:\n${otherErrors.join("\n")}`; + } + + throw new Error(errorMessage); + } + + const themeJson = json as ThemeJson; + assertThemeNameIsValid(themeJson.name); + return themeJson; +} + +function parseThemeJsonContent(label: string, content: string): ThemeJson { + let json: unknown; + try { + json = JSON.parse(content); + } catch (error) { + throw new Error(`Failed to parse theme ${label}: ${error}`); + } + return parseThemeJson(label, json); +} + +function loadThemeJson(name: string): ThemeJson { + const builtinThemes = getBuiltinThemes(); + if (name in builtinThemes) { + return builtinThemes[name]; + } + const registeredTheme = registeredThemes.get(name); + if (registeredTheme?.sourcePath) { + const content = fs.readFileSync(registeredTheme.sourcePath, "utf-8"); + return parseThemeJsonContent(registeredTheme.sourcePath, content); + } + if (registeredTheme) { + throw new Error(`Theme "${name}" does not have a source path for export`); + } + const customThemesDir = getCustomThemesDir(); + const themePath = path.join(customThemesDir, `${name}.json`); + if (!fs.existsSync(themePath)) { + throw new Error(`Theme not found: ${name}`); + } + const content = fs.readFileSync(themePath, "utf-8"); + return parseThemeJsonContent(name, content); +} + +function createTheme(themeJson: ThemeJson, mode?: ColorMode, sourcePath?: string): Theme { + const colorMode = mode ?? (getCapabilities().trueColor ? "truecolor" : "256color"); + const resolvedColors = resolveThemeColors(themeJson.colors, themeJson.vars); + const fgColors: Record = {} as Record; + const bgColors: Record = {} as Record; + const bgColorKeys: Set = new Set([ + "selectedBg", + "userMessageBg", + "customMessageBg", + "toolPendingBg", + "toolSuccessBg", + "toolErrorBg", + ]); + for (const [key, value] of Object.entries(resolvedColors)) { + if (bgColorKeys.has(key)) { + bgColors[key as ThemeBg] = value; + } else { + fgColors[key as ThemeColor] = value; + } + } + return new Theme(fgColors, bgColors, colorMode, { + name: themeJson.name, + sourcePath, + }); +} + +export function loadThemeFromPath(themePath: string, mode?: ColorMode): Theme { + const content = fs.readFileSync(themePath, "utf-8"); + const themeJson = parseThemeJsonContent(themePath, content); + return createTheme(themeJson, mode, themePath); +} + +function loadTheme(name: string, mode?: ColorMode): Theme { + const registeredTheme = registeredThemes.get(name); + if (registeredTheme) { + return registeredTheme; + } + const themeJson = loadThemeJson(name); + return createTheme(themeJson, mode); +} + +export function getThemeByName(name: string): Theme | undefined { + try { + return loadTheme(name); + } catch { + return undefined; + } +} + +export type TerminalTheme = "dark" | "light"; + +export function parseAutoThemeSetting( + themeSetting: string | undefined, +): { lightTheme: string; darkTheme: string } | undefined { + if (!themeSetting) return undefined; + const slashIndex = themeSetting.indexOf("/"); + if (slashIndex === -1 || themeSetting.indexOf("/", slashIndex + 1) !== -1) { + return undefined; + } + + const lightTheme = themeSetting.slice(0, slashIndex).trim(); + const darkTheme = themeSetting.slice(slashIndex + 1).trim(); + if (!lightTheme || !darkTheme) { + return undefined; + } + return { lightTheme, darkTheme }; +} + +export function resolveThemeSetting( + themeSetting: string | undefined, + terminalTheme: TerminalTheme, +): string | undefined { + const autoTheme = parseAutoThemeSetting(themeSetting); + if (autoTheme) { + return terminalTheme === "light" ? autoTheme.lightTheme : autoTheme.darkTheme; + } + if (themeSetting?.includes("/")) return undefined; + if (typeof themeSetting === "string") return themeSetting; + return undefined; +} + +export interface TerminalThemeDetection { + theme: TerminalTheme; + source: "terminal background" | "COLORFGBG" | "fallback"; + detail: string; + confidence: "high" | "low"; +} + +export interface TerminalThemeDetectionOptions { + env?: NodeJS.ProcessEnv; +} + +export interface TerminalBackgroundThemeDetector { + queryTerminalBackgroundColor({ timeoutMs }: { timeoutMs: number }): Promise; +} + +export interface TerminalAutoThemeDetector extends TerminalBackgroundThemeDetector { + queryTerminalColorScheme?({ timeoutMs }: { timeoutMs: number }): Promise; +} + +export interface TerminalBackgroundThemeDetectionOptions extends TerminalThemeDetectionOptions { + ui: TerminalBackgroundThemeDetector; + timeoutMs: number; +} + +export interface TerminalAutoThemeDetectionOptions extends TerminalThemeDetectionOptions { + ui: TerminalAutoThemeDetector; + timeoutMs: number; +} + +function getColorFgBgBackgroundIndex(colorfgbg: string): number | undefined { + const parts = colorfgbg.split(";"); + for (let i = parts.length - 1; i >= 0; i--) { + const bg = parseInt(parts[i].trim(), 10); + if (Number.isInteger(bg) && bg >= 0 && bg <= 255) { + return bg; + } + } + return undefined; +} + +function getRgbColorLuminance({ r, g, b }: RgbColor): number { + const toLinear = (channel: number) => { + const value = channel / 255; + return value <= 0.03928 ? value / 12.92 : ((value + 0.055) / 1.055) ** 2.4; + }; + return 0.2126 * toLinear(r) + 0.7152 * toLinear(g) + 0.0722 * toLinear(b); +} + +function getAnsiColorLuminance(index: number): number { + return getRgbColorLuminance(hexToRgb(ansi256ToHex(index))); +} + +export function getThemeForRgbColor(rgb: RgbColor): TerminalTheme { + return getRgbColorLuminance(rgb) >= 0.5 ? "light" : "dark"; +} + +export function detectTerminalBackgroundFromEnv(options: TerminalThemeDetectionOptions = {}): TerminalThemeDetection { + const env = options.env ?? process.env; + const colorfgbg = env.COLORFGBG || ""; + const bg = getColorFgBgBackgroundIndex(colorfgbg); + if (bg !== undefined) { + return { + theme: getAnsiColorLuminance(bg) >= 0.5 ? "light" : "dark", + source: "COLORFGBG", + detail: `background color index ${bg}`, + confidence: "high", + }; + } + + return { + theme: "dark", + source: "fallback", + detail: "no terminal background hint found", + confidence: "low", + }; +} + +export async function detectTerminalBackgroundTheme({ + ui, + timeoutMs, + env, +}: TerminalBackgroundThemeDetectionOptions): Promise { + try { + const rgb = await ui.queryTerminalBackgroundColor({ timeoutMs }); + if (rgb) { + return { + theme: getThemeForRgbColor(rgb), + source: "terminal background", + detail: `OSC 11 background rgb(${rgb.r}, ${rgb.g}, ${rgb.b})`, + confidence: "high", + }; + } + } catch { + // Fall back to environment-based detection when the terminal query fails. + } + + return detectTerminalBackgroundFromEnv({ env }); +} + +export async function detectTerminalThemeForAuto({ + ui, + timeoutMs, + env, +}: TerminalAutoThemeDetectionOptions): Promise { + try { + const colorScheme = await ui.queryTerminalColorScheme?.({ timeoutMs }); + if (colorScheme) return colorScheme; + } catch { + // Fall back to OSC 11 / COLORFGBG detection when color-scheme DSR is unsupported. + } + return (await detectTerminalBackgroundTheme({ ui, timeoutMs, env })).theme; +} + +export function getDefaultTheme(): string { + return detectTerminalBackgroundFromEnv().theme; +} + +// ============================================================================ +// Global Theme Instance +// ============================================================================ + +// Use globalThis to share theme across module loaders (tsx + jiti in dev mode) +const THEME_KEY = Symbol.for("@earendil-works/pi-coding-agent:theme"); +const THEME_KEY_OLD = Symbol.for("@mariozechner/pi-coding-agent:theme"); + +// Export theme as a getter that reads from globalThis +// This ensures all module instances (tsx, jiti) see the same theme +export const theme: Theme = new Proxy({} as Theme, { + get(_target, prop) { + const t = (globalThis as Record)[THEME_KEY]; + if (!t) throw new Error("Theme not initialized. Call initTheme() first."); + return (t as unknown as Record)[prop]; + }, +}); + +function setGlobalTheme(t: Theme): void { + (globalThis as Record)[THEME_KEY] = t; + (globalThis as Record)[THEME_KEY_OLD] = t; +} + +let currentThemeName: string | undefined; +let themeWatcher: fs.FSWatcher | undefined; +let themeReloadTimer: NodeJS.Timeout | undefined; +let onThemeChangeCallback: (() => void) | undefined; +const registeredThemes = new Map(); + +export function setRegisteredThemes(themes: Theme[]): void { + registeredThemes.clear(); + for (const theme of themes) { + if (theme.name) { + assertThemeNameIsValid(theme.name); + registeredThemes.set(theme.name, theme); + } + } +} + +export function initTheme(themeName?: string, enableWatcher: boolean = false): void { + const name = themeName ?? getDefaultTheme(); + currentThemeName = name; + try { + setGlobalTheme(loadTheme(name)); + if (enableWatcher) { + startThemeWatcher(); + } + } catch (_error) { + // Theme is invalid - fall back to dark theme silently + currentThemeName = "dark"; + setGlobalTheme(loadTheme("dark")); + // Don't start watcher for fallback theme + } +} + +export function setTheme(name: string, enableWatcher: boolean = false): { success: boolean; error?: string } { + currentThemeName = name; + try { + setGlobalTheme(loadTheme(name)); + if (enableWatcher) { + startThemeWatcher(); + } + if (onThemeChangeCallback) { + onThemeChangeCallback(); + } + return { success: true }; + } catch (error) { + // Theme is invalid - fall back to dark theme + currentThemeName = "dark"; + setGlobalTheme(loadTheme("dark")); + // Don't start watcher for fallback theme + return { + success: false, + error: error instanceof Error ? error.message : String(error), + }; + } +} + +export function setThemeInstance(themeInstance: Theme): void { + setGlobalTheme(themeInstance); + currentThemeName = ""; + stopThemeWatcher(); // Can't watch a direct instance + if (onThemeChangeCallback) { + onThemeChangeCallback(); + } +} + +export function onThemeChange(callback: () => void): void { + onThemeChangeCallback = callback; +} + +function startThemeWatcher(): void { + stopThemeWatcher(); + + // Only watch if it's a custom theme (not built-in) + if (!currentThemeName || currentThemeName === "dark" || currentThemeName === "light") { + return; + } + + const customThemesDir = getCustomThemesDir(); + const watchedThemeName = currentThemeName; + const watchedFileName = `${watchedThemeName}.json`; + const themeFile = path.join(customThemesDir, watchedFileName); + + // Only watch if the file exists + if (!fs.existsSync(themeFile)) { + return; + } + + const scheduleReload = () => { + if (themeReloadTimer) { + clearTimeout(themeReloadTimer); + } + themeReloadTimer = setTimeout(() => { + themeReloadTimer = undefined; + + // Ignore stale timers after switching themes or stopping the watcher + if (currentThemeName !== watchedThemeName) { + return; + } + + // Keep the last successfully loaded theme active if the file is temporarily missing + if (!fs.existsSync(themeFile)) { + return; + } + + try { + // Reload the theme from disk and refresh the registry cache + const reloadedTheme = loadThemeFromPath(themeFile); + registeredThemes.set(watchedThemeName, reloadedTheme); + setGlobalTheme(reloadedTheme); + // Notify callback (to invalidate UI) + if (onThemeChangeCallback) { + onThemeChangeCallback(); + } + } catch (_error) { + // Ignore errors (file might be in invalid state while being edited) + } + }, 100); + }; + + themeWatcher = + watchWithErrorHandler( + customThemesDir, + (_eventType, filename) => { + if (currentThemeName !== watchedThemeName) { + return; + } + if (!filename) { + scheduleReload(); + return; + } + if (filename !== watchedFileName) { + return; + } + scheduleReload(); + }, + () => { + closeWatcher(themeWatcher); + themeWatcher = undefined; + }, + ) ?? undefined; +} + +export function stopThemeWatcher(): void { + if (themeReloadTimer) { + clearTimeout(themeReloadTimer); + themeReloadTimer = undefined; + } + closeWatcher(themeWatcher); + themeWatcher = undefined; +} + +// ============================================================================ +// HTML Export Helpers +// ============================================================================ + +/** + * Convert a 256-color index to hex string. + * Indices 0-15: basic colors (approximate) + * Indices 16-231: 6x6x6 color cube + * Indices 232-255: grayscale ramp + */ +function ansi256ToHex(index: number): string { + // Basic colors (0-15) - approximate common terminal values + const basicColors = [ + "#000000", + "#800000", + "#008000", + "#808000", + "#000080", + "#800080", + "#008080", + "#c0c0c0", + "#808080", + "#ff0000", + "#00ff00", + "#ffff00", + "#0000ff", + "#ff00ff", + "#00ffff", + "#ffffff", + ]; + if (index < 16) { + return basicColors[index]; + } + + // Color cube (16-231): 6x6x6 = 216 colors + if (index < 232) { + const cubeIndex = index - 16; + const r = Math.floor(cubeIndex / 36); + const g = Math.floor((cubeIndex % 36) / 6); + const b = cubeIndex % 6; + const toHex = (n: number) => (n === 0 ? 0 : 55 + n * 40).toString(16).padStart(2, "0"); + return `#${toHex(r)}${toHex(g)}${toHex(b)}`; + } + + // Grayscale (232-255): 24 shades + const gray = 8 + (index - 232) * 10; + const grayHex = gray.toString(16).padStart(2, "0"); + return `#${grayHex}${grayHex}${grayHex}`; +} + +/** + * Get resolved theme colors as CSS-compatible hex strings. + * Used by HTML export to generate CSS custom properties. + */ +export function getResolvedThemeColors(themeName?: string): Record { + const name = themeName ?? currentThemeName ?? getDefaultTheme(); + const isLight = name === "light"; + const themeJson = loadThemeJson(name); + const resolved = resolveThemeColors(themeJson.colors, themeJson.vars); + + // Default text color for empty values (terminal uses default fg color) + const defaultText = isLight ? "#000000" : "#e5e5e7"; + + const cssColors: Record = {}; + for (const [key, value] of Object.entries(resolved)) { + if (typeof value === "number") { + cssColors[key] = ansi256ToHex(value); + } else if (value === "") { + // Empty means default terminal color - use sensible fallback for HTML + cssColors[key] = defaultText; + } else { + cssColors[key] = value; + } + } + return cssColors; +} + +/** + * Check if a theme is a "light" theme (for CSS that needs light/dark variants). + */ +export function isLightTheme(themeName?: string): boolean { + // Currently just check the name - could be extended to analyze colors + return themeName === "light"; +} + +/** + * Get explicit export colors from theme JSON, if specified. + * Returns undefined for each color that isn't explicitly set. + */ +export function getThemeExportColors(themeName?: string): { + pageBg?: string; + cardBg?: string; + infoBg?: string; +} { + const name = themeName ?? currentThemeName ?? getDefaultTheme(); + try { + const themeJson = loadThemeJson(name); + const exportSection = themeJson.export; + if (!exportSection) return {}; + + const vars = themeJson.vars ?? {}; + const resolve = (value: ColorValue | undefined): string | undefined => { + if (value === undefined) return undefined; + const resolved = resolveVarRefs(value, vars); + if (typeof resolved === "number") return ansi256ToHex(resolved); + if (resolved === "") return undefined; + return resolved; + }; + + return { + pageBg: resolve(exportSection.pageBg), + cardBg: resolve(exportSection.cardBg), + infoBg: resolve(exportSection.infoBg), + }; + } catch { + return {}; + } +} + +// ============================================================================ +// TUI Helpers +// ============================================================================ + +type CliHighlightTheme = Record string>; + +let cachedHighlightThemeFor: Theme | undefined; +let cachedCliHighlightTheme: CliHighlightTheme | undefined; + +function buildCliHighlightTheme(t: Theme): CliHighlightTheme { + return { + keyword: (s: string) => t.fg("syntaxKeyword", s), + built_in: (s: string) => t.fg("syntaxType", s), + literal: (s: string) => t.fg("syntaxNumber", s), + number: (s: string) => t.fg("syntaxNumber", s), + regexp: (s: string) => t.fg("syntaxString", s), + string: (s: string) => t.fg("syntaxString", s), + comment: (s: string) => t.fg("syntaxComment", s), + doctag: (s: string) => t.fg("syntaxComment", s), + meta: (s: string) => t.fg("muted", s), + function: (s: string) => t.fg("syntaxFunction", s), + title: (s: string) => t.fg("syntaxFunction", s), + class: (s: string) => t.fg("syntaxType", s), + type: (s: string) => t.fg("syntaxType", s), + tag: (s: string) => t.fg("syntaxPunctuation", s), + name: (s: string) => t.fg("syntaxKeyword", s), + attr: (s: string) => t.fg("syntaxVariable", s), + variable: (s: string) => t.fg("syntaxVariable", s), + params: (s: string) => t.fg("syntaxVariable", s), + operator: (s: string) => t.fg("syntaxOperator", s), + punctuation: (s: string) => t.fg("syntaxPunctuation", s), + emphasis: (s: string) => t.italic(s), + strong: (s: string) => t.bold(s), + link: (s: string) => t.underline(s), + addition: (s: string) => t.fg("toolDiffAdded", s), + deletion: (s: string) => t.fg("toolDiffRemoved", s), + }; +} + +function getCliHighlightTheme(t: Theme): CliHighlightTheme { + if (cachedHighlightThemeFor !== t || !cachedCliHighlightTheme) { + cachedHighlightThemeFor = t; + cachedCliHighlightTheme = buildCliHighlightTheme(t); + } + return cachedCliHighlightTheme; +} + +/** + * Highlight code with syntax coloring based on file extension or language. + * Returns array of highlighted lines. + */ +export function highlightCode(code: string, lang?: string): string[] { + // Validate language before highlighting to avoid stderr spam from cli-highlight + const validLang = lang && supportsLanguage(lang) ? lang : undefined; + // Skip highlighting when no valid language is specified. cli-highlight's + // auto-detection is unreliable and can misidentify prose as AppleScript, + // LiveCodeServer, etc., coloring random English words as keywords. + if (!validLang) { + return code.split("\n").map((line) => theme.fg("mdCodeBlock", line)); + } + const opts = { + language: validLang, + ignoreIllegals: true, + theme: getCliHighlightTheme(theme), + }; + try { + return highlight(code, opts).split("\n"); + } catch { + return code.split("\n"); + } +} + +/** + * Get language identifier from file path extension. + */ +export function getLanguageFromPath(filePath: string): string | undefined { + const ext = filePath.split(".").pop()?.toLowerCase(); + if (!ext) return undefined; + + const extToLang: Record = { + ts: "typescript", + tsx: "typescript", + js: "javascript", + jsx: "javascript", + mjs: "javascript", + cjs: "javascript", + py: "python", + rb: "ruby", + rs: "rust", + go: "go", + java: "java", + kt: "kotlin", + swift: "swift", + c: "c", + h: "c", + cpp: "cpp", + cc: "cpp", + cxx: "cpp", + hpp: "cpp", + cs: "csharp", + php: "php", + sh: "bash", + bash: "bash", + zsh: "bash", + fish: "fish", + ps1: "powershell", + sql: "sql", + html: "html", + htm: "html", + css: "css", + scss: "scss", + sass: "sass", + less: "less", + json: "json", + yaml: "yaml", + yml: "yaml", + toml: "toml", + xml: "xml", + md: "markdown", + markdown: "markdown", + dockerfile: "dockerfile", + makefile: "makefile", + cmake: "cmake", + lua: "lua", + perl: "perl", + r: "r", + scala: "scala", + clj: "clojure", + ex: "elixir", + exs: "elixir", + erl: "erlang", + hs: "haskell", + ml: "ocaml", + vim: "vim", + graphql: "graphql", + proto: "protobuf", + tf: "hcl", + hcl: "hcl", + }; + + return extToLang[ext]; +} + +export function getMarkdownTheme(): MarkdownTheme { + return { + heading: (text: string) => theme.fg("mdHeading", text), + link: (text: string) => theme.fg("mdLink", text), + linkUrl: (text: string) => theme.fg("mdLinkUrl", text), + code: (text: string) => theme.fg("mdCode", text), + codeBlock: (text: string) => theme.fg("mdCodeBlock", text), + codeBlockBorder: (text: string) => theme.fg("mdCodeBlockBorder", text), + quote: (text: string) => theme.fg("mdQuote", text), + quoteBorder: (text: string) => theme.fg("mdQuoteBorder", text), + hr: (text: string) => theme.fg("mdHr", text), + listBullet: (text: string) => theme.fg("mdListBullet", text), + bold: (text: string) => theme.bold(text), + italic: (text: string) => theme.italic(text), + underline: (text: string) => theme.underline(text), + strikethrough: (text: string) => chalk.strikethrough(text), + highlightCode: (code: string, lang?: string): string[] => { + // Validate language before highlighting to avoid stderr spam from cli-highlight + const validLang = lang && supportsLanguage(lang) ? lang : undefined; + // Skip highlighting when no valid language is specified. cli-highlight's + // auto-detection is unreliable and can misidentify prose as AppleScript, + // LiveCodeServer, etc., coloring random English words as keywords. + if (!validLang) { + return code.split("\n").map((line) => theme.fg("mdCodeBlock", line)); + } + const opts = { + language: validLang, + ignoreIllegals: true, + theme: getCliHighlightTheme(theme), + }; + try { + return highlight(code, opts).split("\n"); + } catch { + return code.split("\n").map((line) => theme.fg("mdCodeBlock", line)); + } + }, + }; +} + +export function getSelectListTheme(): SelectListTheme { + return { + selectedPrefix: (text: string) => theme.fg("accent", text), + selectedText: (text: string) => theme.fg("accent", text), + description: (text: string) => theme.fg("muted", text), + scrollInfo: (text: string) => theme.fg("muted", text), + noMatch: (text: string) => theme.fg("muted", text), + }; +} + +export function getEditorTheme(): EditorTheme { + return { + borderColor: (text: string) => theme.fg("borderMuted", text), + selectList: getSelectListTheme(), + }; +} + +export function getSettingsListTheme(): SettingsListTheme { + return { + label: (text: string, selected: boolean) => (selected ? theme.fg("accent", text) : text), + value: (text: string, selected: boolean) => (selected ? theme.fg("accent", text) : theme.fg("muted", text)), + description: (text: string) => theme.fg("dim", text), + cursor: theme.fg("accent", "→ "), + hint: (text: string) => theme.fg("dim", text), + }; +} diff --git a/cactus-code/packages/coding-agent/src/modes/print-mode.ts b/cactus-code/packages/coding-agent/src/modes/print-mode.ts new file mode 100644 index 000000000..a5ee02355 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/modes/print-mode.ts @@ -0,0 +1,159 @@ +/** + * Print mode (single-shot): Send prompts, output result, exit. + * + * Used for: + * - `pi -p "prompt"` - text output + * - `pi --mode json "prompt"` - JSON event stream + */ + +import type { AssistantMessage, ImageContent } from "@earendil-works/pi-ai"; +import type { AgentSessionRuntime } from "../core/agent-session-runtime.ts"; +import { flushRawStdout, writeRawStdout } from "../core/output-guard.ts"; +import { killTrackedDetachedChildren } from "../utils/shell.ts"; + +/** + * Options for print mode. + */ +export interface PrintModeOptions { + /** Output mode: "text" for final response only, "json" for all events */ + mode: "text" | "json"; + /** Array of additional prompts to send after initialMessage */ + messages?: string[]; + /** First message to send (may contain @file content) */ + initialMessage?: string; + /** Images to attach to the initial message */ + initialImages?: ImageContent[]; +} + +/** + * Run in print (single-shot) mode. + * Sends prompts to the agent and outputs the result. + */ +export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: PrintModeOptions): Promise { + const { mode, messages = [], initialMessage, initialImages } = options; + let exitCode = 0; + let session = runtimeHost.session; + let unsubscribe: (() => void) | undefined; + let disposed = false; + const signalCleanupHandlers: Array<() => void> = []; + + const disposeRuntime = async (): Promise => { + if (disposed) return; + disposed = true; + unsubscribe?.(); + await runtimeHost.dispose(); + }; + + const registerSignalHandlers = (): void => { + const signals: NodeJS.Signals[] = ["SIGTERM"]; + if (process.platform !== "win32") { + signals.push("SIGHUP"); + } + + for (const signal of signals) { + const handler = () => { + killTrackedDetachedChildren(); + void disposeRuntime().finally(() => { + process.exit(signal === "SIGHUP" ? 129 : 143); + }); + }; + process.on(signal, handler); + signalCleanupHandlers.push(() => process.off(signal, handler)); + } + }; + + registerSignalHandlers(); + + runtimeHost.setRebindSession(async () => { + await rebindSession(); + }); + + const rebindSession = async (): Promise => { + session = runtimeHost.session; + await session.bindExtensions({ + mode: mode === "json" ? "json" : "print", + commandContextActions: { + waitForIdle: () => session.agent.waitForIdle(), + newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions), + fork: async (entryId, forkOptions) => { + const result = await runtimeHost.fork(entryId, forkOptions); + return { cancelled: result.cancelled }; + }, + navigateTree: async (targetId, navigateOptions) => { + const result = await session.navigateTree(targetId, { + summarize: navigateOptions?.summarize, + customInstructions: navigateOptions?.customInstructions, + replaceInstructions: navigateOptions?.replaceInstructions, + label: navigateOptions?.label, + }); + return { cancelled: result.cancelled }; + }, + switchSession: async (sessionPath, switchOptions) => { + return runtimeHost.switchSession(sessionPath, switchOptions); + }, + reload: async () => { + await session.reload(); + }, + }, + onError: (err) => { + console.error(`Extension error (${err.extensionPath}): ${err.error}`); + }, + }); + + unsubscribe?.(); + unsubscribe = session.subscribe((event) => { + if (mode === "json") { + writeRawStdout(`${JSON.stringify(event)}\n`); + } + }); + }; + + try { + if (mode === "json") { + const header = session.sessionManager.getHeader(); + if (header) { + writeRawStdout(`${JSON.stringify(header)}\n`); + } + } + + await rebindSession(); + + if (initialMessage) { + await session.prompt(initialMessage, { images: initialImages }); + } + + for (const message of messages) { + await session.prompt(message); + } + + if (mode === "text") { + const state = session.state; + const lastMessage = state.messages[state.messages.length - 1]; + + if (lastMessage?.role === "assistant") { + const assistantMsg = lastMessage as AssistantMessage; + if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") { + console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`); + exitCode = 1; + } else { + for (const content of assistantMsg.content) { + if (content.type === "text") { + writeRawStdout(`${content.text}\n`); + } + } + } + } + } + + return exitCode; + } catch (error: unknown) { + console.error(error instanceof Error ? error.message : String(error)); + return 1; + } finally { + for (const cleanup of signalCleanupHandlers) { + cleanup(); + } + await disposeRuntime(); + await flushRawStdout(); + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/ansi.ts b/cactus-code/packages/coding-agent/src/utils/ansi.ts new file mode 100644 index 000000000..a95ded68f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/ansi.ts @@ -0,0 +1,60 @@ +/* + * Portions of this file are derived from: + * - ansi-regex (https://github.com/chalk/ansi-regex) + * - strip-ansi (https://github.com/chalk/strip-ansi) + * + * MIT License + * + * Copyright (c) Sindre Sorhus (https://sindresorhus.com) + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +function ansiRegex({ onlyFirst = false }: { onlyFirst?: boolean } = {}): RegExp { + // Valid string terminator sequences are BEL, ESC\, and 0x9c + const ST = "(?:\\u0007|\\u001B\\u005C|\\u009C)"; + + // OSC sequences only: ESC ] ... ST (non-greedy until the first ST) + const osc = `(?:\\u001B\\][\\s\\S]*?${ST})`; + + // CSI and related: ESC/C1, optional intermediates, optional params (supports ; and :) then final byte + const csi = "[\\u001B\\u009B][[\\]()#;?]*(?:\\d{1,4}(?:[;:]\\d{0,4})*)?[\\dA-PR-TZcf-nq-uy=><~]"; + + const pattern = `${osc}|${csi}`; + + return new RegExp(pattern, onlyFirst ? undefined : "g"); +} + +const regex = ansiRegex(); + +export function stripAnsi(value: string): string { + if (typeof value !== "string") { + throw new TypeError(`Expected a \`string\`, got \`${typeof value}\``); + } + + // Fast path: ANSI codes require ESC (7-bit) or CSI (8-bit) introducer + if (!value.includes("\u001B") && !value.includes("\u009B")) { + return value; + } + + // Even though the regex is global, we don't need to reset the `.lastIndex` + // because unlike `.exec()` and `.test()`, `.replace()` does it automatically + // and doing it manually has a performance penalty. + return value.replace(regex, ""); +} diff --git a/cactus-code/packages/coding-agent/src/utils/child-process.ts b/cactus-code/packages/coding-agent/src/utils/child-process.ts new file mode 100644 index 000000000..b152444d6 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/child-process.ts @@ -0,0 +1,137 @@ +import { + type ChildProcess, + type ChildProcessByStdio, + spawn as nodeSpawn, + spawnSync as nodeSpawnSync, + type SpawnOptions, + type SpawnOptionsWithStdioTuple, + type SpawnSyncOptionsWithStringEncoding, + type SpawnSyncReturns, + type StdioNull, + type StdioPipe, +} from "node:child_process"; +import type { Readable } from "node:stream"; +import crossSpawn from "cross-spawn"; + +const EXIT_STDIO_GRACE_MS = 100; + +export function spawnProcess( + command: string, + args: string[], + options: SpawnOptionsWithStdioTuple, +): ChildProcessByStdio; +export function spawnProcess(command: string, args: string[], options: SpawnOptions): ChildProcess; +export function spawnProcess(command: string, args: string[], options: SpawnOptions): ChildProcess { + return process.platform === "win32" ? crossSpawn(command, args, options) : nodeSpawn(command, args, options); +} + +export function spawnProcessSync( + command: string, + args: string[], + options: SpawnSyncOptionsWithStringEncoding, +): SpawnSyncReturns { + return process.platform === "win32" + ? crossSpawn.sync(command, args, options) + : nodeSpawnSync(command, args, options); +} + +/** + * Wait for a child process to terminate without hanging on inherited stdio handles. + * + * A short-lived child can `exit` while a detached descendant keeps its stdout/stderr + * pipe open. We must not resolve and destroy the streams on a fixed deadline measured + * from `exit`, or output still being written past that deadline is silently lost + * (earendil-works/pi#5303). Instead, after `exit` we wait for the pipes to fall idle: + * the grace timer is re-armed on every chunk, so an actively writing descendant keeps + * us reading, while a quiet inherited handle (e.g. a Windows daemonized descendant + * that never lets `close` fire) still releases us after the grace elapses. + */ +export function waitForChildProcess(child: ChildProcess): Promise { + return new Promise((resolve, reject) => { + let settled = false; + let exited = false; + let exitCode: number | null = null; + let postExitTimer: NodeJS.Timeout | undefined; + let stdoutEnded = child.stdout === null; + let stderrEnded = child.stderr === null; + + const cleanup = () => { + if (postExitTimer) { + clearTimeout(postExitTimer); + postExitTimer = undefined; + } + child.removeListener("error", onError); + child.removeListener("exit", onExit); + child.removeListener("close", onClose); + child.stdout?.removeListener("end", onStdoutEnd); + child.stderr?.removeListener("end", onStderrEnd); + child.stdout?.removeListener("data", onData); + child.stderr?.removeListener("data", onData); + }; + + const finalize = (code: number | null) => { + if (settled) return; + settled = true; + cleanup(); + child.stdout?.destroy(); + child.stderr?.destroy(); + resolve(code); + }; + + const maybeFinalizeAfterExit = () => { + if (!exited || settled) return; + if (stdoutEnded && stderrEnded) { + finalize(exitCode); + } + }; + + const armIdleTimer = () => { + if (postExitTimer) clearTimeout(postExitTimer); + postExitTimer = setTimeout(() => finalize(exitCode), EXIT_STDIO_GRACE_MS); + }; + + const onData = () => { + // Output is still arriving after exit; defer finalizing so we don't + // destroy the stream mid-write and truncate the tail. + if (exited && !settled) armIdleTimer(); + }; + + const onStdoutEnd = () => { + stdoutEnded = true; + maybeFinalizeAfterExit(); + }; + + const onStderrEnd = () => { + stderrEnded = true; + maybeFinalizeAfterExit(); + }; + + const onError = (err: Error) => { + if (settled) return; + settled = true; + cleanup(); + reject(err); + }; + + const onExit = (code: number | null) => { + exited = true; + exitCode = code; + maybeFinalizeAfterExit(); + if (!settled) { + armIdleTimer(); + } + }; + + const onClose = (code: number | null) => { + finalize(code); + }; + + child.stdout?.once("end", onStdoutEnd); + child.stderr?.once("end", onStderrEnd); + child.stdout?.on("data", onData); + child.stderr?.on("data", onData); + child.once("error", onError); + child.once("exit", onExit); + child.once("close", onClose); + }); +} diff --git a/cactus-code/packages/coding-agent/src/utils/clipboard-image.ts b/cactus-code/packages/coding-agent/src/utils/clipboard-image.ts new file mode 100644 index 000000000..b98d0e9bb --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/clipboard-image.ts @@ -0,0 +1,300 @@ +import { spawnSync } from "child_process"; +import { randomUUID } from "crypto"; +import { readFileSync, unlinkSync } from "fs"; +import { tmpdir } from "os"; +import { join } from "path"; + +import { clipboard } from "./clipboard-native.ts"; +import { loadPhoton } from "./photon.ts"; + +export type ClipboardImage = { + bytes: Uint8Array; + mimeType: string; +}; + +const SUPPORTED_IMAGE_MIME_TYPES = ["image/png", "image/jpeg", "image/webp", "image/gif"] as const; + +const DEFAULT_LIST_TIMEOUT_MS = 1000; +const DEFAULT_READ_TIMEOUT_MS = 3000; +const DEFAULT_POWERSHELL_TIMEOUT_MS = 5000; +const DEFAULT_MAX_BUFFER_BYTES = 50 * 1024 * 1024; + +export function isWaylandSession(env: NodeJS.ProcessEnv = process.env): boolean { + return Boolean(env.WAYLAND_DISPLAY) || env.XDG_SESSION_TYPE === "wayland"; +} + +function baseMimeType(mimeType: string): string { + return mimeType.split(";")[0]?.trim().toLowerCase() ?? mimeType.toLowerCase(); +} + +export function extensionForImageMimeType(mimeType: string): string | null { + switch (baseMimeType(mimeType)) { + case "image/png": + return "png"; + case "image/jpeg": + return "jpg"; + case "image/webp": + return "webp"; + case "image/gif": + return "gif"; + default: + return null; + } +} + +function selectPreferredImageMimeType(mimeTypes: string[]): string | null { + const normalized = mimeTypes + .map((t) => t.trim()) + .filter(Boolean) + .map((t) => ({ raw: t, base: baseMimeType(t) })); + + for (const preferred of SUPPORTED_IMAGE_MIME_TYPES) { + const match = normalized.find((t) => t.base === preferred); + if (match) { + return match.raw; + } + } + + const anyImage = normalized.find((t) => t.base.startsWith("image/")); + return anyImage?.raw ?? null; +} + +function isSupportedImageMimeType(mimeType: string): boolean { + const base = baseMimeType(mimeType); + return SUPPORTED_IMAGE_MIME_TYPES.some((t) => t === base); +} + +/** + * Convert unsupported image formats to PNG using Photon. + * Returns null if conversion is unavailable or fails. + */ +async function convertToPng(bytes: Uint8Array): Promise { + const photon = await loadPhoton(); + if (!photon) { + return null; + } + + try { + const image = photon.PhotonImage.new_from_byteslice(bytes); + try { + return image.get_bytes(); + } finally { + image.free(); + } + } catch { + return null; + } +} + +function runCommand( + command: string, + args: string[], + options?: { timeoutMs?: number; maxBufferBytes?: number; env?: NodeJS.ProcessEnv }, +): { stdout: Buffer; ok: boolean } { + const timeoutMs = options?.timeoutMs ?? DEFAULT_READ_TIMEOUT_MS; + const maxBufferBytes = options?.maxBufferBytes ?? DEFAULT_MAX_BUFFER_BYTES; + + const result = spawnSync(command, args, { + timeout: timeoutMs, + maxBuffer: maxBufferBytes, + env: options?.env, + }); + + if (result.error) { + return { ok: false, stdout: Buffer.alloc(0) }; + } + + if (result.status !== 0) { + return { ok: false, stdout: Buffer.alloc(0) }; + } + + const stdout = Buffer.isBuffer(result.stdout) + ? result.stdout + : Buffer.from(result.stdout ?? "", typeof result.stdout === "string" ? "utf-8" : undefined); + + return { ok: true, stdout }; +} + +function readClipboardImageViaWlPaste(): ClipboardImage | null { + const list = runCommand("wl-paste", ["--list-types"], { timeoutMs: DEFAULT_LIST_TIMEOUT_MS }); + if (!list.ok) { + return null; + } + + const types = list.stdout + .toString("utf-8") + .split(/\r?\n/) + .map((t) => t.trim()) + .filter(Boolean); + + const selectedType = selectPreferredImageMimeType(types); + if (!selectedType) { + return null; + } + + const data = runCommand("wl-paste", ["--type", selectedType, "--no-newline"]); + if (!data.ok || data.stdout.length === 0) { + return null; + } + + return { bytes: data.stdout, mimeType: baseMimeType(selectedType) }; +} + +function isWSL(env: NodeJS.ProcessEnv = process.env): boolean { + if (env.WSL_DISTRO_NAME || env.WSLENV) { + return true; + } + + try { + const release = readFileSync("/proc/version", "utf-8"); + return /microsoft|wsl/i.test(release); + } catch { + return false; + } +} + +/** + * On WSL, the Linux clipboard (Wayland/X11) does not receive image data from + * Windows screenshots (Win+Shift+S). PowerShell can access the Windows clipboard + * directly, so we use it as a fallback. + */ +function readClipboardImageViaPowerShell(): ClipboardImage | null { + const tmpFile = join(tmpdir(), `cactus-wsl-clip-${randomUUID()}.png`); + + try { + const winPathResult = runCommand("wslpath", ["-w", tmpFile], { timeoutMs: DEFAULT_LIST_TIMEOUT_MS }); + if (!winPathResult.ok) { + return null; + } + + const winPath = winPathResult.stdout.toString("utf-8").trim(); + if (!winPath) { + return null; + } + + const psQuotedWinPath = winPath.replaceAll("'", "''"); + const psScript = [ + "Add-Type -AssemblyName System.Windows.Forms", + "Add-Type -AssemblyName System.Drawing", + `$path = '${psQuotedWinPath}'`, + "$img = [System.Windows.Forms.Clipboard]::GetImage()", + "if ($img) { $img.Save($path, [System.Drawing.Imaging.ImageFormat]::Png); Write-Output 'ok' } else { Write-Output 'empty' }", + ].join("; "); + + const result = runCommand("powershell.exe", ["-NoProfile", "-Command", psScript], { + timeoutMs: DEFAULT_POWERSHELL_TIMEOUT_MS, + }); + if (!result.ok) { + return null; + } + + const output = result.stdout.toString("utf-8").trim(); + if (output !== "ok") { + return null; + } + + const bytes = readFileSync(tmpFile); + if (bytes.length === 0) { + return null; + } + + return { bytes: new Uint8Array(bytes), mimeType: "image/png" }; + } catch { + return null; + } finally { + try { + unlinkSync(tmpFile); + } catch { + // Ignore cleanup errors. + } + } +} + +function readClipboardImageViaXclip(): ClipboardImage | null { + const targets = runCommand("xclip", ["-selection", "clipboard", "-t", "TARGETS", "-o"], { + timeoutMs: DEFAULT_LIST_TIMEOUT_MS, + }); + + let candidateTypes: string[] = []; + if (targets.ok) { + candidateTypes = targets.stdout + .toString("utf-8") + .split(/\r?\n/) + .map((t) => t.trim()) + .filter(Boolean); + } + + const preferred = candidateTypes.length > 0 ? selectPreferredImageMimeType(candidateTypes) : null; + const tryTypes = preferred ? [preferred, ...SUPPORTED_IMAGE_MIME_TYPES] : [...SUPPORTED_IMAGE_MIME_TYPES]; + + for (const mimeType of tryTypes) { + const data = runCommand("xclip", ["-selection", "clipboard", "-t", mimeType, "-o"]); + if (data.ok && data.stdout.length > 0) { + return { bytes: data.stdout, mimeType: baseMimeType(mimeType) }; + } + } + + return null; +} + +async function readClipboardImageViaNativeClipboard(): Promise { + if (!clipboard || !clipboard.hasImage()) { + return null; + } + + const imageData = await clipboard.getImageBinary(); + if (!imageData || imageData.length === 0) { + return null; + } + + const bytes = imageData instanceof Uint8Array ? imageData : Uint8Array.from(imageData); + return { bytes, mimeType: "image/png" }; +} + +export async function readClipboardImage(options?: { + env?: NodeJS.ProcessEnv; + platform?: NodeJS.Platform; +}): Promise { + const env = options?.env ?? process.env; + const platform = options?.platform ?? process.platform; + + if (env.TERMUX_VERSION) { + return null; + } + + let image: ClipboardImage | null = null; + + if (platform === "linux") { + const wsl = isWSL(env); + const wayland = isWaylandSession(env); + + if (wayland || wsl) { + image = readClipboardImageViaWlPaste() ?? readClipboardImageViaXclip(); + } + + if (!image && wsl) { + image = readClipboardImageViaPowerShell(); + } + + if (!image && !wayland) { + image = await readClipboardImageViaNativeClipboard(); + } + } else { + image = await readClipboardImageViaNativeClipboard(); + } + + if (!image) { + return null; + } + + // Convert unsupported formats (e.g., BMP from WSLg) to PNG + if (!isSupportedImageMimeType(image.mimeType)) { + const pngBytes = await convertToPng(image.bytes); + if (!pngBytes) { + return null; + } + return { bytes: pngBytes, mimeType: "image/png" }; + } + + return image; +} diff --git a/cactus-code/packages/coding-agent/src/utils/clipboard-native.ts b/cactus-code/packages/coding-agent/src/utils/clipboard-native.ts new file mode 100644 index 000000000..f8aeb5472 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/clipboard-native.ts @@ -0,0 +1,32 @@ +import { createRequire } from "module"; +import { dirname, join } from "path"; +import { pathToFileURL } from "url"; + +export type ClipboardModule = { + setText: (text: string) => Promise; + hasImage: () => boolean; + getImageBinary: () => Promise>; +}; + +type ClipboardRequire = (id: string) => unknown; + +const moduleRequire = createRequire(import.meta.url); +const executableDirRequire = createRequire(pathToFileURL(join(dirname(process.execPath), "package.json")).href); +const hasDisplay = process.platform !== "linux" || Boolean(process.env.DISPLAY || process.env.WAYLAND_DISPLAY); + +export function loadClipboardNative( + requires: readonly ClipboardRequire[] = [moduleRequire, executableDirRequire], +): ClipboardModule | null { + for (const requireClipboard of requires) { + try { + return requireClipboard("@mariozechner/clipboard") as ClipboardModule; + } catch { + // Try the next resolution root. + } + } + return null; +} + +const clipboard = !process.env.TERMUX_VERSION && hasDisplay ? loadClipboardNative() : null; + +export { clipboard }; diff --git a/cactus-code/packages/coding-agent/src/utils/clipboard.ts b/cactus-code/packages/coding-agent/src/utils/clipboard.ts new file mode 100644 index 000000000..26e1d018d --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/clipboard.ts @@ -0,0 +1,127 @@ +import { execSync, spawn } from "child_process"; +import { platform } from "os"; +import { isWaylandSession } from "./clipboard-image.ts"; +import { clipboard } from "./clipboard-native.ts"; + +type NativeClipboardExecOptions = { + input: string; + timeout: number; + stdio: ["pipe", "ignore", "ignore"]; +}; + +function copyToX11Clipboard(options: NativeClipboardExecOptions): void { + try { + execSync("xclip -selection clipboard", options); + } catch { + execSync("xsel --clipboard --input", options); + } +} + +const MAX_OSC52_ENCODED_LENGTH = 100_000; + +function isRemoteSession(env: NodeJS.ProcessEnv = process.env): boolean { + return Boolean(env.SSH_CONNECTION || env.SSH_CLIENT || env.MOSH_CONNECTION); +} + +function emitOsc52(text: string): boolean { + const encoded = Buffer.from(text).toString("base64"); + if (encoded.length > MAX_OSC52_ENCODED_LENGTH) { + return false; + } + process.stdout.write(`\x1b]52;c;${encoded}\x07`); + return true; +} + +export async function copyToClipboard(text: string): Promise { + let copied = false; + + const p = platform(); + + // Prefer direct clipboard writes. Emitting OSC 52 first can make terminals + // write the same native clipboard concurrently with the addon, and very large + // OSC 52 payloads can desynchronize terminal rendering. + // + // On Linux, skip the native addon. The underlying `clipboard-rs` crate is + // X11-only and does not retain selection ownership after `set_text` + // resolves, so on Wayland-only compositors (Hyprland, Niri, ...) and even + // some X11 sessions the call resolves successfully without populating the + // clipboard. The platform tools below (wl-copy, xclip, xsel) properly + // daemonize and keep ownership. + try { + if (clipboard && p !== "linux") { + await clipboard.setText(text); + copied = true; + } + } catch { + // Fall through to platform-specific clipboard tools. + } + + const remote = isRemoteSession(); + if (copied && !remote) { + return; + } + + const options: NativeClipboardExecOptions = { input: text, timeout: 5000, stdio: ["pipe", "ignore", "ignore"] }; + + if (!copied) { + try { + if (p === "darwin") { + execSync("pbcopy", options); + copied = true; + } else if (p === "win32") { + execSync("clip", options); + copied = true; + } else { + // Linux. Try Termux, Wayland, or X11 clipboard tools. + if (process.env.TERMUX_VERSION) { + try { + execSync("termux-clipboard-set", options); + copied = true; + } catch { + // Fall back to Wayland or X11 tools. + } + } + + if (!copied) { + const hasWaylandDisplay = Boolean(process.env.WAYLAND_DISPLAY); + const hasX11Display = Boolean(process.env.DISPLAY); + const isWayland = isWaylandSession(); + if (isWayland && hasWaylandDisplay) { + try { + // Verify wl-copy exists (spawn errors are async and won't be caught) + execSync("which wl-copy", { stdio: "ignore" }); + // wl-copy with execSync hangs due to fork behavior; use spawn instead + const proc = spawn("wl-copy", [], { stdio: ["pipe", "ignore", "ignore"] }); + proc.stdin.on("error", () => { + // Ignore EPIPE errors if wl-copy exits early + }); + proc.stdin.write(text); + proc.stdin.end(); + proc.unref(); + copied = true; + } catch { + if (hasX11Display) { + copyToX11Clipboard(options); + copied = true; + } + } + } else if (hasX11Display) { + copyToX11Clipboard(options); + copied = true; + } + } + } + } catch { + // Fall through to OSC 52 fallback. + } + } + + if (remote || !copied) { + const osc52Copied = emitOsc52(text); + copied = copied || osc52Copied; + } + + if (!copied) { + throw new Error("Failed to copy to clipboard"); + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/deprecation.ts b/cactus-code/packages/coding-agent/src/utils/deprecation.ts new file mode 100644 index 000000000..78a2f1469 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/deprecation.ts @@ -0,0 +1,14 @@ +import chalk from "chalk"; + +const emittedDeprecationWarnings = new Set(); + +export function warnDeprecation(message: string): void { + if (emittedDeprecationWarnings.has(message)) return; + emittedDeprecationWarnings.add(message); + console.warn(chalk.yellow(`Deprecation warning: ${message}`)); +} + +/** Clear deprecation warning state. Exported for tests. */ +export function clearDeprecationWarningsForTests(): void { + emittedDeprecationWarnings.clear(); +} diff --git a/cactus-code/packages/coding-agent/src/utils/exif-orientation.ts b/cactus-code/packages/coding-agent/src/utils/exif-orientation.ts new file mode 100644 index 000000000..4b454afaa --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/exif-orientation.ts @@ -0,0 +1,183 @@ +import type { PhotonImageType } from "./photon.ts"; + +type Photon = typeof import("@silvia-odwyer/photon-node"); + +function readOrientationFromTiff(bytes: Uint8Array, tiffStart: number): number { + if (tiffStart + 8 > bytes.length) return 1; + + const byteOrder = (bytes[tiffStart] << 8) | bytes[tiffStart + 1]; + const le = byteOrder === 0x4949; + + const read16 = (pos: number): number => { + if (le) return bytes[pos] | (bytes[pos + 1] << 8); + return (bytes[pos] << 8) | bytes[pos + 1]; + }; + + const read32 = (pos: number): number => { + if (le) return bytes[pos] | (bytes[pos + 1] << 8) | (bytes[pos + 2] << 16) | (bytes[pos + 3] << 24); + return ((bytes[pos] << 24) | (bytes[pos + 1] << 16) | (bytes[pos + 2] << 8) | bytes[pos + 3]) >>> 0; + }; + + const ifdOffset = read32(tiffStart + 4); + const ifdStart = tiffStart + ifdOffset; + if (ifdStart + 2 > bytes.length) return 1; + + const entryCount = read16(ifdStart); + for (let i = 0; i < entryCount; i++) { + const entryPos = ifdStart + 2 + i * 12; + if (entryPos + 12 > bytes.length) return 1; + + if (read16(entryPos) === 0x0112) { + const value = read16(entryPos + 8); + return value >= 1 && value <= 8 ? value : 1; + } + } + + return 1; +} + +function findJpegTiffOffset(bytes: Uint8Array): number { + let offset = 2; + while (offset < bytes.length - 1) { + if (bytes[offset] !== 0xff) return -1; + const marker = bytes[offset + 1]; + if (marker === 0xff) { + offset++; + continue; + } + + if (marker === 0xe1) { + if (offset + 4 >= bytes.length) return -1; + const segmentStart = offset + 4; + if (segmentStart + 6 > bytes.length) return -1; + if (!hasExifHeader(bytes, segmentStart)) return -1; + return segmentStart + 6; + } + + if (offset + 4 > bytes.length) return -1; + const length = (bytes[offset + 2] << 8) | bytes[offset + 3]; + offset += 2 + length; + } + + return -1; +} + +function findWebpTiffOffset(bytes: Uint8Array): number { + let offset = 12; + while (offset + 8 <= bytes.length) { + const chunkId = String.fromCharCode(bytes[offset], bytes[offset + 1], bytes[offset + 2], bytes[offset + 3]); + const chunkSize = + bytes[offset + 4] | (bytes[offset + 5] << 8) | (bytes[offset + 6] << 16) | (bytes[offset + 7] << 24); + const dataStart = offset + 8; + + if (chunkId === "EXIF") { + if (dataStart + chunkSize > bytes.length) return -1; + // Some WebP files have "Exif\0\0" prefix before the TIFF header + const tiffStart = chunkSize >= 6 && hasExifHeader(bytes, dataStart) ? dataStart + 6 : dataStart; + return tiffStart; + } + + // RIFF chunks are padded to even size + offset = dataStart + chunkSize + (chunkSize % 2); + } + + return -1; +} + +function hasExifHeader(bytes: Uint8Array, offset: number): boolean { + return ( + bytes[offset] === 0x45 && + bytes[offset + 1] === 0x78 && + bytes[offset + 2] === 0x69 && + bytes[offset + 3] === 0x66 && + bytes[offset + 4] === 0x00 && + bytes[offset + 5] === 0x00 + ); +} + +function getExifOrientation(bytes: Uint8Array): number { + let tiffOffset = -1; + + // JPEG: starts with FF D8 + if (bytes.length >= 2 && bytes[0] === 0xff && bytes[1] === 0xd8) { + tiffOffset = findJpegTiffOffset(bytes); + } + // WebP: starts with RIFF....WEBP + else if ( + bytes.length >= 12 && + bytes[0] === 0x52 && + bytes[1] === 0x49 && + bytes[2] === 0x46 && + bytes[3] === 0x46 && + bytes[8] === 0x57 && + bytes[9] === 0x45 && + bytes[10] === 0x42 && + bytes[11] === 0x50 + ) { + tiffOffset = findWebpTiffOffset(bytes); + } + + if (tiffOffset === -1) return 1; + return readOrientationFromTiff(bytes, tiffOffset); +} + +type DstIndexFn = (x: number, y: number, w: number, h: number) => number; + +function rotate90(photon: Photon, image: PhotonImageType, dstIndex: DstIndexFn): PhotonImageType { + const w = image.get_width(); + const h = image.get_height(); + const src = image.get_raw_pixels(); + const dst = new Uint8Array(src.length); + + for (let y = 0; y < h; y++) { + for (let x = 0; x < w; x++) { + const srcIdx = (y * w + x) * 4; + const dstIdx = dstIndex(x, y, w, h) * 4; + dst[dstIdx] = src[srcIdx]; + dst[dstIdx + 1] = src[srcIdx + 1]; + dst[dstIdx + 2] = src[srcIdx + 2]; + dst[dstIdx + 3] = src[srcIdx + 3]; + } + } + + return new photon.PhotonImage(dst, h, w); +} + +// Flip orientations mutate in-place. Rotations return a new image (caller must free the old one if different). +export function applyExifOrientation( + photon: Photon, + image: PhotonImageType, + originalBytes: Uint8Array, +): PhotonImageType { + const orientation = getExifOrientation(originalBytes); + if (orientation === 1) return image; + + switch (orientation) { + case 2: + photon.fliph(image); + return image; + case 3: + photon.fliph(image); + photon.flipv(image); + return image; + case 4: + photon.flipv(image); + return image; + case 5: { + const rotated = rotate90(photon, image, (x, y, _w, h) => x * h + (h - 1 - y)); + photon.fliph(rotated); + return rotated; + } + case 6: + return rotate90(photon, image, (x, y, _w, h) => x * h + (h - 1 - y)); + case 7: { + const rotated = rotate90(photon, image, (x, y, w, h) => (w - 1 - x) * h + y); + photon.fliph(rotated); + return rotated; + } + case 8: + return rotate90(photon, image, (x, y, w, h) => (w - 1 - x) * h + y); + default: + return image; + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/frontmatter.ts b/cactus-code/packages/coding-agent/src/utils/frontmatter.ts new file mode 100644 index 000000000..847e2e539 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/frontmatter.ts @@ -0,0 +1,39 @@ +import { parse } from "yaml"; + +type ParsedFrontmatter> = { + frontmatter: T; + body: string; +}; + +const normalizeNewlines = (value: string): string => value.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + +const extractFrontmatter = (content: string): { yamlString: string | null; body: string } => { + const normalized = normalizeNewlines(content); + + if (!normalized.startsWith("---")) { + return { yamlString: null, body: normalized }; + } + + const endIndex = normalized.indexOf("\n---", 3); + if (endIndex === -1) { + return { yamlString: null, body: normalized }; + } + + return { + yamlString: normalized.slice(4, endIndex), + body: normalized.slice(endIndex + 4).trim(), + }; +}; + +export const parseFrontmatter = = Record>( + content: string, +): ParsedFrontmatter => { + const { yamlString, body } = extractFrontmatter(content); + if (!yamlString) { + return { frontmatter: {} as T, body }; + } + const parsed = parse(yamlString); + return { frontmatter: (parsed ?? {}) as T, body }; +}; + +export const stripFrontmatter = (content: string): string => parseFrontmatter(content).body; diff --git a/cactus-code/packages/coding-agent/src/utils/fs-watch.ts b/cactus-code/packages/coding-agent/src/utils/fs-watch.ts new file mode 100644 index 000000000..daaf80900 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/fs-watch.ts @@ -0,0 +1,30 @@ +import { type FSWatcher, type WatchListener, watch } from "node:fs"; + +export const FS_WATCH_RETRY_DELAY_MS = 5000; + +export function closeWatcher(watcher: FSWatcher | null | undefined): void { + if (!watcher) { + return; + } + + try { + watcher.close(); + } catch { + // Ignore watcher close errors + } +} + +export function watchWithErrorHandler( + path: string, + listener: WatchListener, + onError: () => void, +): FSWatcher | null { + try { + const watcher = watch(path, listener); + watcher.on("error", onError); + return watcher; + } catch { + onError(); + return null; + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/git.ts b/cactus-code/packages/coding-agent/src/utils/git.ts new file mode 100644 index 000000000..1314edea7 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/git.ts @@ -0,0 +1,226 @@ +import hostedGitInfo from "hosted-git-info"; + +/** + * Parsed git URL information. + */ +export type GitSource = { + /** Always "git" for git sources */ + type: "git"; + /** Clone URL (always valid for git clone, without ref suffix) */ + repo: string; + /** Git host domain (e.g., "github.com") */ + host: string; + /** Repository path (e.g., "user/repo") */ + path: string; + /** Git ref (branch, tag, commit) if specified */ + ref?: string; + /** True if ref was specified (package won't be auto-updated) */ + pinned: boolean; +}; + +function splitRef(url: string): { repo: string; ref?: string } { + const scpLikeMatch = url.match(/^git@([^:]+):(.+)$/); + if (scpLikeMatch) { + const pathWithMaybeRef = scpLikeMatch[2] ?? ""; + const refSeparator = pathWithMaybeRef.indexOf("@"); + if (refSeparator < 0) return { repo: url }; + const repoPath = pathWithMaybeRef.slice(0, refSeparator); + const ref = pathWithMaybeRef.slice(refSeparator + 1); + if (!repoPath || !ref) return { repo: url }; + return { + repo: `git@${scpLikeMatch[1] ?? ""}:${repoPath}`, + ref, + }; + } + + if (url.includes("://")) { + try { + const parsed = new URL(url); + const pathWithMaybeRef = parsed.pathname.replace(/^\/+/, ""); + const refSeparator = pathWithMaybeRef.indexOf("@"); + if (refSeparator < 0) return { repo: url }; + const repoPath = pathWithMaybeRef.slice(0, refSeparator); + const ref = pathWithMaybeRef.slice(refSeparator + 1); + if (!repoPath || !ref) return { repo: url }; + parsed.pathname = `/${repoPath}`; + return { + repo: parsed.toString().replace(/\/$/, ""), + ref, + }; + } catch { + return { repo: url }; + } + } + + const slashIndex = url.indexOf("/"); + if (slashIndex < 0) { + return { repo: url }; + } + const host = url.slice(0, slashIndex); + const pathWithMaybeRef = url.slice(slashIndex + 1); + const refSeparator = pathWithMaybeRef.indexOf("@"); + if (refSeparator < 0) { + return { repo: url }; + } + const repoPath = pathWithMaybeRef.slice(0, refSeparator); + const ref = pathWithMaybeRef.slice(refSeparator + 1); + if (!repoPath || !ref) { + return { repo: url }; + } + return { + repo: `${host}/${repoPath}`, + ref, + }; +} + +function decodeForValidation(value: string): string | null { + try { + return decodeURIComponent(value); + } catch { + return null; + } +} + +function hasUnsafeGitInstallPart(value: string, allowSlash: boolean): boolean { + const decoded = decodeForValidation(value); + if (decoded === null) { + return true; + } + const candidates = [value, decoded]; + for (const candidate of candidates) { + if (candidate.includes("\0") || candidate.includes("\\") || candidate.startsWith("/")) { + return true; + } + if (!allowSlash && candidate.includes("/")) { + return true; + } + if (candidate.split("/").includes("..")) { + return true; + } + } + return false; +} + +function buildGitSource(args: { repo: string; host: string; path: string; ref?: string }): GitSource | null { + if (args.path.startsWith("/")) { + return null; + } + const normalizedPath = args.path.replace(/\.git$/, "").replace(/^\/+/, ""); + if (!args.host || !normalizedPath || normalizedPath.split("/").length < 2) { + return null; + } + if (hasUnsafeGitInstallPart(args.host, false) || hasUnsafeGitInstallPart(normalizedPath, true)) { + return null; + } + + return { + type: "git", + repo: args.repo, + host: args.host, + path: normalizedPath, + ref: args.ref, + pinned: Boolean(args.ref), + }; +} + +function parseGenericGitUrl(url: string): GitSource | null { + const { repo: repoWithoutRef, ref } = splitRef(url); + let repo = repoWithoutRef; + let host = ""; + let path = ""; + + const scpLikeMatch = repoWithoutRef.match(/^git@([^:]+):(.+)$/); + if (scpLikeMatch) { + host = scpLikeMatch[1] ?? ""; + path = scpLikeMatch[2] ?? ""; + } else if ( + repoWithoutRef.startsWith("https://") || + repoWithoutRef.startsWith("http://") || + repoWithoutRef.startsWith("ssh://") || + repoWithoutRef.startsWith("git://") + ) { + try { + const parsed = new URL(repoWithoutRef); + host = parsed.hostname; + path = parsed.pathname.replace(/^\/+/, ""); + } catch { + return null; + } + } else { + const slashIndex = repoWithoutRef.indexOf("/"); + if (slashIndex < 0) { + return null; + } + host = repoWithoutRef.slice(0, slashIndex); + path = repoWithoutRef.slice(slashIndex + 1); + if (!host.includes(".") && host !== "localhost") { + return null; + } + repo = `https://${repoWithoutRef}`; + } + + return buildGitSource({ repo, host, path, ref }); +} + +/** + * Parse git source into a GitSource. + * + * Rules: + * - With git: prefix, accept all historical shorthand forms. + * - Without git: prefix, only accept explicit protocol URLs. + */ +export function parseGitUrl(source: string): GitSource | null { + const trimmed = source.trim(); + const hasGitPrefix = trimmed.startsWith("git:"); + const url = hasGitPrefix ? trimmed.slice(4).trim() : trimmed; + + if (!hasGitPrefix && !/^(https?|ssh|git):\/\//i.test(url)) { + return null; + } + + const split = splitRef(url); + + const hostedCandidates = [split.ref ? `${split.repo}#${split.ref}` : undefined, url].filter( + (value): value is string => Boolean(value), + ); + for (const candidate of hostedCandidates) { + const info = hostedGitInfo.fromUrl(candidate); + if (info) { + if (split.ref && info.project?.includes("@")) { + continue; + } + const useHttpsPrefix = + !split.repo.startsWith("http://") && + !split.repo.startsWith("https://") && + !split.repo.startsWith("ssh://") && + !split.repo.startsWith("git://") && + !split.repo.startsWith("git@"); + return buildGitSource({ + repo: useHttpsPrefix ? `https://${split.repo}` : split.repo, + host: info.domain || "", + path: `${info.user}/${info.project}`, + ref: info.committish || split.ref || undefined, + }); + } + } + + const httpsCandidates = [split.ref ? `https://${split.repo}#${split.ref}` : undefined, `https://${url}`].filter( + (value): value is string => Boolean(value), + ); + for (const candidate of httpsCandidates) { + const info = hostedGitInfo.fromUrl(candidate); + if (info) { + if (split.ref && info.project?.includes("@")) { + continue; + } + return buildGitSource({ + repo: `https://${split.repo}`, + host: info.domain || "", + path: `${info.user}/${info.project}`, + ref: info.committish || split.ref || undefined, + }); + } + } + + return parseGenericGitUrl(url); +} diff --git a/cactus-code/packages/coding-agent/src/utils/highlight-js-lib-index.d.ts b/cactus-code/packages/coding-agent/src/utils/highlight-js-lib-index.d.ts new file mode 100644 index 000000000..75e31da28 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/highlight-js-lib-index.d.ts @@ -0,0 +1,19 @@ +declare module "highlight.js/lib/index.js" { + interface HighlightResult { + value: string; + } + + interface HighlightOptions { + language: string; + ignoreIllegals?: boolean; + } + + interface HighlightJs { + highlight(code: string, options: HighlightOptions): HighlightResult; + highlightAuto(code: string, languageSubset?: string[]): HighlightResult; + getLanguage(name: string): unknown; + } + + const hljs: HighlightJs; + export default hljs; +} diff --git a/cactus-code/packages/coding-agent/src/utils/html.ts b/cactus-code/packages/coding-agent/src/utils/html.ts new file mode 100644 index 000000000..a13ad46f4 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/html.ts @@ -0,0 +1,51 @@ +export interface DecodedHtmlEntity { + text: string; + length: number; +} + +function decodeCodePoint(codePoint: number): string | undefined { + if (!Number.isInteger(codePoint) || codePoint < 0 || codePoint > 0x10ffff) { + return undefined; + } + return String.fromCodePoint(codePoint); +} + +export function decodeHtmlEntity(entity: string): string | undefined { + switch (entity) { + case "amp": + return "&"; + case "lt": + return "<"; + case "gt": + return ">"; + case "quot": + return '"'; + case "apos": + return "'"; + } + + if (entity.startsWith("#x") || entity.startsWith("#X")) { + return decodeCodePoint(Number.parseInt(entity.slice(2), 16)); + } + + if (entity.startsWith("#")) { + return decodeCodePoint(Number.parseInt(entity.slice(1), 10)); + } + + return undefined; +} + +export function decodeHtmlEntityAt(html: string, index: number): DecodedHtmlEntity | undefined { + const semicolonIndex = html.indexOf(";", index + 1); + if (semicolonIndex === -1 || semicolonIndex - index > 16) { + return undefined; + } + + const entity = html.slice(index + 1, semicolonIndex); + const decoded = decodeHtmlEntity(entity); + if (decoded === undefined) { + return undefined; + } + + return { text: decoded, length: semicolonIndex - index + 1 }; +} diff --git a/cactus-code/packages/coding-agent/src/utils/image-convert.ts b/cactus-code/packages/coding-agent/src/utils/image-convert.ts new file mode 100644 index 000000000..4d5f0204a --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/image-convert.ts @@ -0,0 +1,41 @@ +import { applyExifOrientation } from "./exif-orientation.ts"; +import { loadPhoton } from "./photon.ts"; + +/** + * Convert image to PNG format for terminal display. + * Kitty graphics protocol requires PNG format (f=100). + */ +export async function convertToPng( + base64Data: string, + mimeType: string, +): Promise<{ data: string; mimeType: string } | null> { + // Already PNG, no conversion needed + if (mimeType === "image/png") { + return { data: base64Data, mimeType }; + } + + const photon = await loadPhoton(); + if (!photon) { + // Photon not available, can't convert + return null; + } + + try { + const bytes = new Uint8Array(Buffer.from(base64Data, "base64")); + const rawImage = photon.PhotonImage.new_from_byteslice(bytes); + const image = applyExifOrientation(photon, rawImage, bytes); + if (image !== rawImage) rawImage.free(); + try { + const pngBuffer = image.get_bytes(); + return { + data: Buffer.from(pngBuffer).toString("base64"), + mimeType: "image/png", + }; + } finally { + image.free(); + } + } catch { + // Conversion failed + return null; + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/image-resize-core.ts b/cactus-code/packages/coding-agent/src/utils/image-resize-core.ts new file mode 100644 index 000000000..e60820c7b --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/image-resize-core.ts @@ -0,0 +1,164 @@ +import { applyExifOrientation } from "./exif-orientation.ts"; +import { loadPhoton } from "./photon.ts"; + +export interface ImageResizeOptions { + maxWidth?: number; // Default: 2000 + maxHeight?: number; // Default: 2000 + maxBytes?: number; // Default: 4.5MB of base64 payload (below Anthropic's 5MB limit) + jpegQuality?: number; // Default: 80 +} + +export interface ResizedImage { + data: string; // base64 + mimeType: string; + originalWidth: number; + originalHeight: number; + width: number; + height: number; + wasResized: boolean; +} + +// 4.5MB of base64 payload. Provides headroom below Anthropic's 5MB limit. +const DEFAULT_MAX_BYTES = 4.5 * 1024 * 1024; + +const DEFAULT_OPTIONS: Required = { + maxWidth: 2000, + maxHeight: 2000, + maxBytes: DEFAULT_MAX_BYTES, + jpegQuality: 80, +}; + +interface EncodedCandidate { + data: string; + encodedSize: number; + mimeType: string; +} + +function encodeCandidate(buffer: Uint8Array, mimeType: string): EncodedCandidate { + const data = Buffer.from(buffer).toString("base64"); + return { + data, + encodedSize: Buffer.byteLength(data, "utf-8"), + mimeType, + }; +} + +/** + * Resize an image to fit within the specified max dimensions and encoded file size. + * Returns null if the image cannot be resized below maxBytes. + * + * Uses Photon (Rust/WASM) for image processing. If Photon is not available, + * returns null. + * + * Strategy for staying under maxBytes: + * 1. First resize to maxWidth/maxHeight + * 2. Try both PNG and JPEG formats, pick the smaller one + * 3. If still too large, try JPEG with decreasing quality + * 4. If still too large, progressively reduce dimensions until 1x1 + */ +export async function resizeImageInProcess( + inputBytes: Uint8Array, + mimeType: string, + options?: ImageResizeOptions, +): Promise { + const opts = { ...DEFAULT_OPTIONS, ...options }; + const inputBase64Size = Math.ceil(inputBytes.byteLength / 3) * 4; + + const photon = await loadPhoton(); + if (!photon) { + return null; + } + + let image: ReturnType | undefined; + try { + const rawImage = photon.PhotonImage.new_from_byteslice(inputBytes); + image = applyExifOrientation(photon, rawImage, inputBytes); + if (image !== rawImage) rawImage.free(); + + const originalWidth = image.get_width(); + const originalHeight = image.get_height(); + const format = mimeType.split("/")[1] ?? "png"; + + // Check if already within all limits (dimensions AND encoded size) + if (originalWidth <= opts.maxWidth && originalHeight <= opts.maxHeight && inputBase64Size < opts.maxBytes) { + return { + data: Buffer.from(inputBytes).toString("base64"), + mimeType: mimeType || `image/${format}`, + originalWidth, + originalHeight, + width: originalWidth, + height: originalHeight, + wasResized: false, + }; + } + + // Calculate initial dimensions respecting max limits + let targetWidth = originalWidth; + let targetHeight = originalHeight; + + if (targetWidth > opts.maxWidth) { + targetHeight = Math.round((targetHeight * opts.maxWidth) / targetWidth); + targetWidth = opts.maxWidth; + } + if (targetHeight > opts.maxHeight) { + targetWidth = Math.round((targetWidth * opts.maxHeight) / targetHeight); + targetHeight = opts.maxHeight; + } + + function tryEncodings(width: number, height: number, jpegQualities: number[]): EncodedCandidate[] { + const resized = photon!.resize(image!, width, height, photon!.SamplingFilter.Lanczos3); + + try { + const candidates: EncodedCandidate[] = [encodeCandidate(resized.get_bytes(), "image/png")]; + for (const quality of jpegQualities) { + candidates.push(encodeCandidate(resized.get_bytes_jpeg(quality), "image/jpeg")); + } + return candidates; + } finally { + resized.free(); + } + } + + const qualitySteps = Array.from(new Set([opts.jpegQuality, 85, 70, 55, 40])); + let currentWidth = targetWidth; + let currentHeight = targetHeight; + + while (true) { + const candidates = tryEncodings(currentWidth, currentHeight, qualitySteps); + for (const candidate of candidates) { + if (candidate.encodedSize < opts.maxBytes) { + return { + data: candidate.data, + mimeType: candidate.mimeType, + originalWidth, + originalHeight, + width: currentWidth, + height: currentHeight, + wasResized: true, + }; + } + } + + if (currentWidth === 1 && currentHeight === 1) { + break; + } + + const nextWidth = currentWidth === 1 ? 1 : Math.max(1, Math.floor(currentWidth * 0.75)); + const nextHeight = currentHeight === 1 ? 1 : Math.max(1, Math.floor(currentHeight * 0.75)); + if (nextWidth === currentWidth && nextHeight === currentHeight) { + break; + } + + currentWidth = nextWidth; + currentHeight = nextHeight; + } + + return null; + } catch { + return null; + } finally { + if (image) { + image.free(); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/image-resize-worker.ts b/cactus-code/packages/coding-agent/src/utils/image-resize-worker.ts new file mode 100644 index 000000000..ee881b3d9 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/image-resize-worker.ts @@ -0,0 +1,42 @@ +import { parentPort } from "node:worker_threads"; +import { type ImageResizeOptions, type ResizedImage, resizeImageInProcess } from "./image-resize-core.ts"; + +interface ResizeImageWorkerRequest { + inputBytes: Uint8Array; + mimeType: string; + options?: ImageResizeOptions; +} + +interface ResizeImageWorkerResponse { + result?: ResizedImage | null; + error?: string; +} + +function isResizeImageWorkerRequest(value: unknown): value is ResizeImageWorkerRequest { + if (!value || typeof value !== "object") return false; + const record = value as Record; + return record.inputBytes instanceof Uint8Array && typeof record.mimeType === "string"; +} + +const port = parentPort; +if (!port) { + throw new Error("image resize worker requires parentPort"); +} + +port.once("message", (message: unknown) => { + void (async () => { + try { + if (!isResizeImageWorkerRequest(message)) { + throw new Error("Invalid image resize worker request"); + } + const result = await resizeImageInProcess(message.inputBytes, message.mimeType, message.options); + const response: ResizeImageWorkerResponse = { result }; + port.postMessage(response); + } catch (error) { + const response: ResizeImageWorkerResponse = { + error: error instanceof Error ? error.message : String(error), + }; + port.postMessage(response); + } + })(); +}); diff --git a/cactus-code/packages/coding-agent/src/utils/image-resize.ts b/cactus-code/packages/coding-agent/src/utils/image-resize.ts new file mode 100644 index 000000000..516a1e57a --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/image-resize.ts @@ -0,0 +1,123 @@ +import { Worker } from "node:worker_threads"; +import { type ImageResizeOptions, type ResizedImage, resizeImageInProcess } from "./image-resize-core.ts"; + +export type { ImageResizeOptions, ResizedImage } from "./image-resize-core.ts"; + +interface ResizeImageWorkerResponse { + result?: ResizedImage | null; + error?: string; +} + +function toTransferableBytes(input: Uint8Array): Uint8Array { + // Transfer detaches the buffer, so transfer a worker-owned copy and leave the + // caller's bytes intact. + return new Uint8Array(input); +} + +function isResizeImageWorkerResponse(value: unknown): value is ResizeImageWorkerResponse { + return value !== null && typeof value === "object"; +} + +function createResizeWorker(workerSpecifier: string | URL): Worker { + return new Worker(workerSpecifier); +} + +async function resizeImageInWorker( + workerSpecifier: string | URL, + inputBytes: Uint8Array, + mimeType: string, + options?: ImageResizeOptions, +): Promise { + const worker = createResizeWorker(workerSpecifier); + try { + const inputBytesForWorker = toTransferableBytes(inputBytes); + return await new Promise((resolve, reject) => { + let settled = false; + const settle = (result: ResizedImage | null): void => { + if (settled) return; + settled = true; + resolve(result); + }; + const fail = (error: Error): void => { + if (settled) return; + settled = true; + reject(error); + }; + + worker.once("message", (message: unknown) => { + if (!isResizeImageWorkerResponse(message)) { + fail(new Error("Invalid image resize worker response")); + return; + } + if (message.error) { + fail(new Error(message.error)); + return; + } + settle(message.result ?? null); + }); + worker.once("error", fail); + worker.once("exit", (code) => { + if (!settled) { + fail(new Error(`Image resize worker exited with code ${code}`)); + } + }); + worker.postMessage( + { + inputBytes: inputBytesForWorker, + mimeType, + options, + }, + [inputBytesForWorker.buffer], + ); + }); + } finally { + void worker.terminate().catch(() => undefined); + } +} + +/** + * Resize an image to fit within the specified max dimensions and encoded file size. + * Runs Photon in a worker thread so WASM decoding, resizing, and encoding do not + * block the TUI event loop. If the worker cannot be loaded (for example in some + * Bun compiled executable layouts), fall back to in-process resizing so image + * reads still work. + */ +export async function resizeImage( + inputBytes: Uint8Array, + mimeType: string, + options?: ImageResizeOptions, +): Promise { + const isTypeScriptRuntime = import.meta.url.endsWith(".ts"); + const workerUrl = new URL( + isTypeScriptRuntime ? "./image-resize-worker.ts" : "./image-resize-worker.js", + import.meta.url, + ); + + // Bun compiled executables resolve worker entrypoints by string path, not via + // new URL(..., import.meta.url). Try the string path first under Bun so the + // release binary uses the embedded worker instead of falling back in-process. + if (typeof process.versions.bun === "string") { + try { + return await resizeImageInWorker("./src/utils/image-resize-worker.ts", inputBytes, mimeType, options); + } catch {} + } + + try { + return await resizeImageInWorker(workerUrl, inputBytes, mimeType, options); + } catch { + return resizeImageInProcess(inputBytes, mimeType, options); + } +} + +/** + * Format a dimension note for resized images. + * This helps the model understand the coordinate mapping. + */ +export function formatDimensionNote(result: ResizedImage): string | undefined { + if (!result.wasResized) { + return undefined; + } + + const scale = result.originalWidth / result.width; + return `[Image: original ${result.originalWidth}x${result.originalHeight}, displayed at ${result.width}x${result.height}. Multiply coordinates by ${scale.toFixed(2)} to map to original image.]`; +} diff --git a/cactus-code/packages/coding-agent/src/utils/json.ts b/cactus-code/packages/coding-agent/src/utils/json.ts new file mode 100644 index 000000000..9ee7b7b1f --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/json.ts @@ -0,0 +1,6 @@ +/** Strip `//` line comments and trailing commas from JSON, leaving string literals untouched. */ +export function stripJsonComments(input: string): string { + return input + .replace(/"(?:\\.|[^"\\])*"|\/\/[^\n]*/g, (m) => (m[0] === '"' ? m : "")) + .replace(/"(?:\\.|[^"\\])*"|,(\s*[}\]])/g, (m, tail) => tail ?? (m[0] === '"' ? m : "")); +} diff --git a/cactus-code/packages/coding-agent/src/utils/mime.ts b/cactus-code/packages/coding-agent/src/utils/mime.ts new file mode 100644 index 000000000..7381d3703 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/mime.ts @@ -0,0 +1,74 @@ +import { open } from "node:fs/promises"; + +const IMAGE_TYPE_SNIFF_BYTES = 4100; +const PNG_SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]; + +export function detectSupportedImageMimeType(buffer: Uint8Array): string | null { + if (startsWith(buffer, [0xff, 0xd8, 0xff])) { + return buffer[3] === 0xf7 ? null : "image/jpeg"; + } + if (startsWith(buffer, PNG_SIGNATURE)) { + return isPng(buffer) && !isAnimatedPng(buffer) ? "image/png" : null; + } + if (startsWithAscii(buffer, 0, "GIF")) { + return "image/gif"; + } + if (startsWithAscii(buffer, 0, "RIFF") && startsWithAscii(buffer, 8, "WEBP")) { + return "image/webp"; + } + return null; +} + +export async function detectSupportedImageMimeTypeFromFile(filePath: string): Promise { + const fileHandle = await open(filePath, "r"); + try { + const buffer = Buffer.alloc(IMAGE_TYPE_SNIFF_BYTES); + const { bytesRead } = await fileHandle.read(buffer, 0, IMAGE_TYPE_SNIFF_BYTES, 0); + return detectSupportedImageMimeType(buffer.subarray(0, bytesRead)); + } finally { + await fileHandle.close(); + } +} + +function isPng(buffer: Uint8Array): boolean { + return ( + buffer.length >= 16 && readUint32BE(buffer, PNG_SIGNATURE.length) === 13 && startsWithAscii(buffer, 12, "IHDR") + ); +} + +function isAnimatedPng(buffer: Uint8Array): boolean { + let offset = PNG_SIGNATURE.length; + while (offset + 8 <= buffer.length) { + const chunkLength = readUint32BE(buffer, offset); + const chunkTypeOffset = offset + 4; + if (startsWithAscii(buffer, chunkTypeOffset, "acTL")) return true; + if (startsWithAscii(buffer, chunkTypeOffset, "IDAT")) return false; + + const nextOffset = offset + 8 + chunkLength + 4; + if (nextOffset <= offset || nextOffset > buffer.length) return false; + offset = nextOffset; + } + return false; +} + +function readUint32BE(buffer: Uint8Array, offset: number): number { + return ( + (buffer[offset] ?? 0) * 0x1000000 + + ((buffer[offset + 1] ?? 0) << 16) + + ((buffer[offset + 2] ?? 0) << 8) + + (buffer[offset + 3] ?? 0) + ); +} + +function startsWith(buffer: Uint8Array, bytes: number[]): boolean { + if (buffer.length < bytes.length) return false; + return bytes.every((byte, index) => buffer[index] === byte); +} + +function startsWithAscii(buffer: Uint8Array, offset: number, text: string): boolean { + if (buffer.length < offset + text.length) return false; + for (let index = 0; index < text.length; index++) { + if (buffer[offset + index] !== text.charCodeAt(index)) return false; + } + return true; +} diff --git a/cactus-code/packages/coding-agent/src/utils/open-browser.ts b/cactus-code/packages/coding-agent/src/utils/open-browser.ts new file mode 100644 index 000000000..435e23f97 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/open-browser.ts @@ -0,0 +1,24 @@ +import { spawn } from "node:child_process"; + +/** + * Open a URL or file in the platform browser/default handler. + * + * This intentionally never invokes a shell. On Windows, do not use + * `cmd /c start`: cmd.exe re-parses metacharacters (&, |, ^, ...) before + * `start` runs, which would make attacker-controlled URLs injectable. + */ +export function openBrowser(target: string): void { + const [cmd, args]: [string, string[]] = + process.platform === "darwin" + ? ["open", [target]] + : process.platform === "win32" + ? ["rundll32", ["url.dll,FileProtocolHandler", target]] + : ["xdg-open", [target]]; + + // spawn reports launcher failures (for example, missing xdg-open) via an + // error event. Browser launch is best-effort: callers still present the target + // to the user, so keep the launcher failure from becoming a process crash. + spawn(cmd, args, { stdio: "ignore", detached: true }) + .on("error", () => {}) + .unref(); +} diff --git a/cactus-code/packages/coding-agent/src/utils/paths.ts b/cactus-code/packages/coding-agent/src/utils/paths.ts new file mode 100644 index 000000000..7f1013446 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/paths.ts @@ -0,0 +1,118 @@ +import { realpathSync } from "node:fs"; +import { homedir } from "node:os"; +import { isAbsolute, join, resolve as nodeResolvePath, relative, sep } from "node:path"; +import { fileURLToPath } from "node:url"; +import { spawnProcessSync } from "./child-process.ts"; + +const UNICODE_SPACES = /[\u00A0\u2000-\u200A\u202F\u205F\u3000]/g; + +export interface PathInputOptions { + /** Trim leading/trailing whitespace before normalization. */ + trim?: boolean; + /** Expand leading `~` to a home directory. Defaults to true. */ + expandTilde?: boolean; + /** Home directory used for `~` expansion. Defaults to `os.homedir()`. */ + homeDir?: string; + /** Strip a leading `@`, used for CLI @file paths. */ + stripAtPrefix?: boolean; + /** Normalize unicode space variants to regular spaces. */ + normalizeUnicodeSpaces?: boolean; +} + +/** + * Resolve a path to its canonical (real) form, following symlinks. + * Falls back to the raw path if resolution fails (e.g. the target does + * not exist yet), so that callers never crash on missing filesystem + * entries. + */ +export function canonicalizePath(path: string): string { + try { + return realpathSync(path); + } catch { + return path; + } +} + +/** + * Returns true if the value is NOT a package source (npm:, git:, etc.) + * or a remote URL protocol. Bare names, relative paths, and file: URLs + * are considered local. + */ +export function isLocalPath(value: string): boolean { + const trimmed = value.trim(); + // Known non-local prefixes. file: URLs are local paths and are intentionally resolved by resolvePath(). + if ( + trimmed.startsWith("npm:") || + trimmed.startsWith("git:") || + trimmed.startsWith("github:") || + trimmed.startsWith("http:") || + trimmed.startsWith("https:") || + trimmed.startsWith("ssh:") + ) { + return false; + } + return true; +} + +export function normalizePath(input: string, options: PathInputOptions = {}): string { + let normalized = options.trim ? input.trim() : input; + if (options.normalizeUnicodeSpaces) { + normalized = normalized.replace(UNICODE_SPACES, " "); + } + if (options.stripAtPrefix && normalized.startsWith("@")) { + normalized = normalized.slice(1); + } + + if (options.expandTilde ?? true) { + const home = options.homeDir ?? homedir(); + if (normalized === "~") return home; + if (normalized.startsWith("~/") || (process.platform === "win32" && normalized.startsWith("~\\"))) { + return join(home, normalized.slice(2)); + } + } + + if (/^file:\/\//.test(normalized)) { + return fileURLToPath(normalized); + } + + return normalized; +} + +export function resolvePath(input: string, baseDir: string = process.cwd(), options: PathInputOptions = {}): string { + const normalized = normalizePath(input, options); + const normalizedBaseDir = normalizePath(baseDir); + return isAbsolute(normalized) ? nodeResolvePath(normalized) : nodeResolvePath(normalizedBaseDir, normalized); +} + +export function getCwdRelativePath(filePath: string, cwd: string): string | undefined { + const resolvedCwd = resolvePath(cwd); + const resolvedPath = resolvePath(filePath, resolvedCwd); + const relativePath = relative(resolvedCwd, resolvedPath); + const isInsideCwd = + relativePath === "" || + (relativePath !== ".." && !relativePath.startsWith(`..${sep}`) && !isAbsolute(relativePath)); + + return isInsideCwd ? relativePath || "." : undefined; +} + +export function formatPathRelativeToCwdOrAbsolute(filePath: string, cwd: string): string { + const absolutePath = resolvePath(filePath, cwd); + return (getCwdRelativePath(absolutePath, cwd) ?? absolutePath).split(sep).join("/"); +} + +export function markPathIgnoredByCloudSync(path: string): void { + const attrs = + process.platform === "darwin" + ? ["com.dropbox.ignored", "com.apple.fileprovider.ignore#P"] + : process.platform === "linux" + ? ["user.com.dropbox.ignored"] + : []; + + for (const attr of attrs) { + if (process.platform === "darwin") { + spawnProcessSync("xattr", ["-w", attr, "1", path], { encoding: "utf-8", stdio: "ignore" }); + } else { + spawnProcessSync("setfattr", ["-n", attr, "-v", "1", path], { encoding: "utf-8", stdio: "ignore" }); + } + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/photon.ts b/cactus-code/packages/coding-agent/src/utils/photon.ts new file mode 100644 index 000000000..6c320705e --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/photon.ts @@ -0,0 +1,139 @@ +/** + * Photon image processing wrapper. + * + * This module provides a unified interface to @silvia-odwyer/photon-node that works in: + * 1. Node.js (development, npm run build) + * 2. Bun compiled binaries (standalone distribution) + * + * The challenge: photon-node's CJS entry uses fs.readFileSync(__dirname + '/photon_rs_bg.wasm') + * which bakes the build machine's absolute path into Bun compiled binaries. + * + * Solution: + * 1. Patch fs.readFileSync to redirect missing photon_rs_bg.wasm reads + * 2. Copy photon_rs_bg.wasm next to the executable in build:binary + */ + +import type { PathOrFileDescriptor } from "fs"; +import { createRequire } from "module"; +import * as path from "path"; +import { fileURLToPath } from "url"; + +const require = createRequire(import.meta.url); +const fs = require("fs") as typeof import("fs"); + +// Re-export types from the main package +export type { PhotonImage as PhotonImageType } from "@silvia-odwyer/photon-node"; + +type ReadFileSync = typeof fs.readFileSync; + +const WASM_FILENAME = "photon_rs_bg.wasm"; + +// Lazy-loaded photon module +let photonModule: typeof import("@silvia-odwyer/photon-node") | null = null; +let loadPromise: Promise | null = null; + +function pathOrNull(file: PathOrFileDescriptor): string | null { + if (typeof file === "string") { + return file; + } + if (file instanceof URL) { + return fileURLToPath(file); + } + return null; +} + +function getFallbackWasmPaths(): string[] { + const execDir = path.dirname(process.execPath); + return [ + path.join(execDir, WASM_FILENAME), + path.join(execDir, "photon", WASM_FILENAME), + path.join(process.cwd(), WASM_FILENAME), + ]; +} + +function patchPhotonWasmRead(): () => void { + const originalReadFileSync: ReadFileSync = fs.readFileSync.bind(fs); + const fallbackPaths = getFallbackWasmPaths(); + const mutableFs = fs as { readFileSync: ReadFileSync }; + + const patchedReadFileSync: ReadFileSync = ((...args: Parameters) => { + const [file, options] = args; + const resolvedPath = pathOrNull(file); + + if (resolvedPath?.endsWith(WASM_FILENAME)) { + try { + return originalReadFileSync(...args); + } catch (error) { + const err = error as NodeJS.ErrnoException; + if (err?.code && err.code !== "ENOENT") { + throw error; + } + + for (const fallbackPath of fallbackPaths) { + if (!fs.existsSync(fallbackPath)) { + continue; + } + if (options === undefined) { + return originalReadFileSync(fallbackPath); + } + return originalReadFileSync(fallbackPath, options); + } + + throw error; + } + } + + return originalReadFileSync(...args); + }) as ReadFileSync; + + try { + mutableFs.readFileSync = patchedReadFileSync; + } catch { + Object.defineProperty(fs, "readFileSync", { + value: patchedReadFileSync, + writable: true, + configurable: true, + }); + } + + return () => { + try { + mutableFs.readFileSync = originalReadFileSync; + } catch { + Object.defineProperty(fs, "readFileSync", { + value: originalReadFileSync, + writable: true, + configurable: true, + }); + } + }; +} + +/** + * Load the photon module asynchronously. + * Returns cached module on subsequent calls. + */ +export async function loadPhoton(): Promise { + if (photonModule) { + return photonModule; + } + + if (loadPromise) { + return loadPromise; + } + + loadPromise = (async () => { + const restoreReadFileSync = patchPhotonWasmRead(); + try { + photonModule = await import("@silvia-odwyer/photon-node"); + return photonModule; + } catch { + photonModule = null; + return photonModule; + } finally { + restoreReadFileSync(); + } + })(); + + return loadPromise; +} diff --git a/cactus-code/packages/coding-agent/src/utils/shell.ts b/cactus-code/packages/coding-agent/src/utils/shell.ts new file mode 100644 index 000000000..2cafa595e --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/shell.ts @@ -0,0 +1,225 @@ +import { existsSync } from "node:fs"; +import { delimiter } from "node:path"; +import { spawn, spawnSync } from "child_process"; +import { getBinDir } from "../config.ts"; + +export interface ShellConfig { + shell: string; + args: string[]; + commandTransport?: "argv" | "stdin"; +} + +/** + * Find bash executable on PATH (cross-platform) + */ +function isLegacyWslBashPath(path: string): boolean { + const normalized = path.replace(/\//g, "\\").toLowerCase(); + return /^[a-z]:\\windows\\(?:system32|sysnative)\\bash\.exe$/.test(normalized); +} + +function getBashShellConfig(shell: string): ShellConfig { + return isLegacyWslBashPath(shell) ? { shell, args: ["-s"], commandTransport: "stdin" } : { shell, args: ["-c"] }; +} + +function findBashOnPath(): string | null { + if (process.platform === "win32") { + // Windows: Use 'where' and verify file exists (where can return non-existent paths) + try { + const result = spawnSync("where", ["bash.exe"], { + encoding: "utf-8", + timeout: 5000, + windowsHide: true, + }); + if (result.status === 0 && result.stdout) { + const firstMatch = result.stdout.trim().split(/\r?\n/)[0]; + if (firstMatch && existsSync(firstMatch)) { + return firstMatch; + } + } + } catch { + // Ignore errors + } + return null; + } + + // Unix: Use 'which' and trust its output (handles Termux and special filesystems) + try { + const result = spawnSync("which", ["bash"], { encoding: "utf-8", timeout: 5000 }); + if (result.status === 0 && result.stdout) { + const firstMatch = result.stdout.trim().split(/\r?\n/)[0]; + if (firstMatch) { + return firstMatch; + } + } + } catch { + // Ignore errors + } + return null; +} + +/** + * Resolve shell configuration based on platform and an optional explicit shell path. + * Resolution order: + * 1. User-specified shellPath + * 2. On Windows: Git Bash in known locations, then bash on PATH + * 3. On Unix: /bin/bash, then bash on PATH, then fallback to sh + */ +export function getShellConfig(customShellPath?: string): ShellConfig { + // 1. Check user-specified shell path + if (customShellPath) { + if (existsSync(customShellPath)) { + return getBashShellConfig(customShellPath); + } + throw new Error(`Custom shell path not found: ${customShellPath}`); + } + + if (process.platform === "win32") { + // 2. Try Git Bash in known locations + const paths: string[] = []; + const programFiles = process.env.ProgramFiles; + if (programFiles) { + paths.push(`${programFiles}\\Git\\bin\\bash.exe`); + } + const programFilesX86 = process.env["ProgramFiles(x86)"]; + if (programFilesX86) { + paths.push(`${programFilesX86}\\Git\\bin\\bash.exe`); + } + + for (const path of paths) { + if (existsSync(path)) { + return getBashShellConfig(path); + } + } + + // 3. Fallback: search bash.exe on PATH (Cygwin, MSYS2, WSL, etc.) + const bashOnPath = findBashOnPath(); + if (bashOnPath) { + return getBashShellConfig(bashOnPath); + } + + throw new Error( + `No bash shell found. Options:\n` + + ` 1. Install Git for Windows: https://git-scm.com/download/win\n` + + ` 2. Add your bash to PATH (Cygwin, MSYS2, etc.)\n` + + " 3. Set shellPath in settings.json\n\n" + + `Searched Git Bash in:\n${paths.map((p) => ` ${p}`).join("\n")}`, + ); + } + + // Unix: try /bin/bash, then bash on PATH, then fallback to sh + if (existsSync("/bin/bash")) { + return getBashShellConfig("/bin/bash"); + } + + const bashOnPath = findBashOnPath(); + if (bashOnPath) { + return getBashShellConfig(bashOnPath); + } + + return { shell: "sh", args: ["-c"] }; +} + +export function getShellEnv(): NodeJS.ProcessEnv { + const binDir = getBinDir(); + const pathKey = Object.keys(process.env).find((key) => key.toLowerCase() === "path") ?? "PATH"; + const currentPath = process.env[pathKey] ?? ""; + const pathEntries = currentPath.split(delimiter).filter(Boolean); + const hasBinDir = pathEntries.includes(binDir); + const updatedPath = hasBinDir ? currentPath : [binDir, currentPath].filter(Boolean).join(delimiter); + + return { + ...process.env, + [pathKey]: updatedPath, + }; +} + +/** + * Sanitize binary output for display/storage. + * Removes characters that crash string-width or cause display issues: + * - Control characters (except tab, newline, carriage return) + * - Lone surrogates + * - Unicode Format characters (crash string-width due to a bug) + * - Characters with undefined code points + */ +export function sanitizeBinaryOutput(str: string): string { + // Use Array.from to properly iterate over code points (not code units) + // This handles surrogate pairs correctly and catches edge cases where + // codePointAt() might return undefined + return Array.from(str) + .filter((char) => { + // Filter out characters that cause string-width to crash + // This includes: + // - Unicode format characters + // - Lone surrogates (already filtered by Array.from) + // - Control chars except \t \n \r + // - Characters with undefined code points + + const code = char.codePointAt(0); + + // Skip if code point is undefined (edge case with invalid strings) + if (code === undefined) return false; + + // Allow tab, newline, carriage return + if (code === 0x09 || code === 0x0a || code === 0x0d) return true; + + // Filter out control characters (0x00-0x1F, except 0x09, 0x0a, 0x0x0d) + if (code <= 0x1f) return false; + + // Filter out Unicode format characters + if (code >= 0xfff9 && code <= 0xfffb) return false; + + return true; + }) + .join(""); +} + +/** + * Detached child processes must be tracked so they can be killed on parent + * shutdown signals (SIGHUP/SIGTERM). + */ +const trackedDetachedChildPids = new Set(); + +export function trackDetachedChildPid(pid: number): void { + trackedDetachedChildPids.add(pid); +} + +export function untrackDetachedChildPid(pid: number): void { + trackedDetachedChildPids.delete(pid); +} + +export function killTrackedDetachedChildren(): void { + for (const pid of trackedDetachedChildPids) { + killProcessTree(pid); + } + trackedDetachedChildPids.clear(); +} + +/** + * Kill a process and all its children (cross-platform) + */ +export function killProcessTree(pid: number): void { + if (process.platform === "win32") { + // Use taskkill on Windows to kill process tree + try { + spawn("taskkill", ["/F", "/T", "/PID", String(pid)], { + stdio: "ignore", + detached: true, + windowsHide: true, + }); + } catch { + // Ignore errors if taskkill fails + } + } else { + // Use SIGKILL on Unix/Linux/Mac + try { + process.kill(-pid, "SIGKILL"); + } catch { + // Fallback to killing just the child if process group kill fails + try { + process.kill(pid, "SIGKILL"); + } catch { + // Process already dead + } + } + } +} diff --git a/cactus-code/packages/coding-agent/src/utils/sleep.ts b/cactus-code/packages/coding-agent/src/utils/sleep.ts new file mode 100644 index 000000000..948f93c47 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/sleep.ts @@ -0,0 +1,18 @@ +/** + * Sleep helper that respects abort signal. + */ +export function sleep(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(new Error("Aborted")); + return; + } + + const timeout = setTimeout(resolve, ms); + + signal?.addEventListener("abort", () => { + clearTimeout(timeout); + reject(new Error("Aborted")); + }); + }); +} diff --git a/cactus-code/packages/coding-agent/src/utils/syntax-highlight.ts b/cactus-code/packages/coding-agent/src/utils/syntax-highlight.ts new file mode 100644 index 000000000..bcc4add1c --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/syntax-highlight.ts @@ -0,0 +1,146 @@ +import hljs from "highlight.js/lib/index.js"; +import { decodeHtmlEntityAt } from "./html.ts"; + +export type HighlightFormatter = (text: string) => string; +export type HighlightTheme = Partial>; + +export interface HighlightOptions { + language?: string; + ignoreIllegals?: boolean; + languageSubset?: string[]; + theme?: HighlightTheme; +} + +const SPAN_CLOSE = ""; +const HIGHLIGHT_CLASS_PREFIX = "hljs-"; + +function getScopeFromSpanTag(tag: string): string | undefined { + const match = /\sclass\s*=\s*(?:"([^"]*)"|'([^']*)')/.exec(tag); + const classValue = match?.[1] ?? match?.[2]; + if (!classValue) { + return undefined; + } + + for (const className of classValue.split(/\s+/)) { + if (className.startsWith(HIGHLIGHT_CLASS_PREFIX)) { + return className.slice(HIGHLIGHT_CLASS_PREFIX.length); + } + } + + return undefined; +} + +function getScopeFormatter(scope: string, theme: HighlightTheme): HighlightFormatter | undefined { + const exact = theme[scope]; + if (exact) { + return exact; + } + + const dotIndex = scope.indexOf("."); + if (dotIndex !== -1) { + const prefixFormatter = theme[scope.slice(0, dotIndex)]; + if (prefixFormatter) { + return prefixFormatter; + } + } + + const dashIndex = scope.indexOf("-"); + if (dashIndex !== -1) { + const prefixFormatter = theme[scope.slice(0, dashIndex)]; + if (prefixFormatter) { + return prefixFormatter; + } + } + + return undefined; +} + +function getActiveFormatter(scopes: Array, theme: HighlightTheme): HighlightFormatter | undefined { + for (let i = scopes.length - 1; i >= 0; i--) { + const scope = scopes[i]; + if (!scope) { + continue; + } + const formatter = getScopeFormatter(scope, theme); + if (formatter) { + return formatter; + } + } + return theme.default; +} + +function isSpanOpenTagStart(html: string, index: number): boolean { + if (!html.startsWith("" || nextChar === " " || nextChar === "\t" || nextChar === "\n" || nextChar === "\r"; +} + +export function renderHighlightedHtml(html: string, theme: HighlightTheme = {}): string { + let output = ""; + let textBuffer = ""; + const scopes: Array = []; + + const flushText = () => { + if (!textBuffer) { + return; + } + const formatter = getActiveFormatter(scopes, theme); + output += formatter ? formatter(textBuffer) : textBuffer; + textBuffer = ""; + }; + + let index = 0; + while (index < html.length) { + if (isSpanOpenTagStart(html, index)) { + const tagEndIndex = html.indexOf(">", index + 5); + if (tagEndIndex !== -1) { + flushText(); + const tag = html.slice(index, tagEndIndex + 1); + const scope = getScopeFromSpanTag(tag); + scopes.push(scope); + index = tagEndIndex + 1; + continue; + } + } + + if (html.startsWith(SPAN_CLOSE, index)) { + flushText(); + if (scopes.length > 0) { + scopes.pop(); + } + index += SPAN_CLOSE.length; + continue; + } + + if (html[index] === "&") { + const decoded = decodeHtmlEntityAt(html, index); + if (decoded) { + textBuffer += decoded.text; + index += decoded.length; + continue; + } + } + + textBuffer += html[index]; + index++; + } + + flushText(); + return output; +} + +export function highlight(code: string, options: HighlightOptions = {}): string { + const html = options.language + ? hljs.highlight(code, { + language: options.language, + ignoreIllegals: options.ignoreIllegals, + }).value + : hljs.highlightAuto(code, options.languageSubset).value; + return renderHighlightedHtml(html, options.theme); +} + +export function supportsLanguage(name: string): boolean { + return hljs.getLanguage(name) !== undefined; +} diff --git a/cactus-code/packages/coding-agent/src/utils/tools-manager.ts b/cactus-code/packages/coding-agent/src/utils/tools-manager.ts new file mode 100644 index 000000000..09d9a1195 --- /dev/null +++ b/cactus-code/packages/coding-agent/src/utils/tools-manager.ts @@ -0,0 +1,369 @@ +import chalk from "chalk"; +import { type SpawnSyncReturns, spawnSync } from "child_process"; +import { chmodSync, createWriteStream, existsSync, mkdirSync, readdirSync, renameSync, rmSync } from "fs"; +import { arch, platform } from "os"; +import { join } from "path"; +import { Readable } from "stream"; +import { pipeline } from "stream/promises"; +import { APP_NAME, getBinDir } from "../config.ts"; + +const TOOLS_DIR = getBinDir(); +const NETWORK_TIMEOUT_MS = 10_000; +const DOWNLOAD_TIMEOUT_MS = 120_000; + +function isOfflineModeEnabled(): boolean { + const value = process.env.PI_OFFLINE; + if (!value) return false; + return value === "1" || value.toLowerCase() === "true" || value.toLowerCase() === "yes"; +} + +interface ToolConfig { + name: string; + repo: string; // GitHub repo (e.g., "sharkdp/fd") + binaryName: string; // Name of the binary inside the archive + systemBinaryNames?: string[]; // Alternative system command names to try before downloading + tagPrefix: string; // Prefix for tags (e.g., "v" for v1.0.0, "" for 1.0.0) + getAssetName: (version: string, plat: string, architecture: string) => string | null; +} + +const TOOLS: Record = { + fd: { + name: "fd", + repo: "sharkdp/fd", + binaryName: "fd", + systemBinaryNames: ["fd", "fdfind"], + tagPrefix: "v", + getAssetName: (version, plat, architecture) => { + if (plat === "darwin") { + const archStr = architecture === "arm64" ? "aarch64" : "x86_64"; + return `fd-v${version}-${archStr}-apple-darwin.tar.gz`; + } else if (plat === "linux") { + const archStr = architecture === "arm64" ? "aarch64" : "x86_64"; + return `fd-v${version}-${archStr}-unknown-linux-gnu.tar.gz`; + } else if (plat === "win32") { + const archStr = architecture === "arm64" ? "aarch64" : "x86_64"; + return `fd-v${version}-${archStr}-pc-windows-msvc.zip`; + } + return null; + }, + }, + rg: { + name: "ripgrep", + repo: "BurntSushi/ripgrep", + binaryName: "rg", + tagPrefix: "", + getAssetName: (version, plat, architecture) => { + if (plat === "darwin") { + const archStr = architecture === "arm64" ? "aarch64" : "x86_64"; + return `ripgrep-${version}-${archStr}-apple-darwin.tar.gz`; + } else if (plat === "linux") { + if (architecture === "arm64") { + return `ripgrep-${version}-aarch64-unknown-linux-gnu.tar.gz`; + } + return `ripgrep-${version}-x86_64-unknown-linux-musl.tar.gz`; + } else if (plat === "win32") { + const archStr = architecture === "arm64" ? "aarch64" : "x86_64"; + return `ripgrep-${version}-${archStr}-pc-windows-msvc.zip`; + } + return null; + }, + }, +}; + +// Check if a command exists in PATH by trying to run it +function commandExists(cmd: string): boolean { + try { + const result = spawnSync(cmd, ["--version"], { stdio: "pipe" }); + // Check for ENOENT error (command not found) + return result.error === undefined || result.error === null; + } catch { + return false; + } +} + +// Get the path to a tool (system-wide or in our tools dir) +export function getToolPath(tool: "fd" | "rg"): string | null { + const config = TOOLS[tool]; + if (!config) return null; + + // Check our tools directory first + const localPath = join(TOOLS_DIR, config.binaryName + (platform() === "win32" ? ".exe" : "")); + if (existsSync(localPath)) { + return localPath; + } + + // Check system PATH - if found, just return the command name (it's in PATH) + const systemBinaryNames = config.systemBinaryNames ?? [config.binaryName]; + for (const systemBinaryName of systemBinaryNames) { + if (commandExists(systemBinaryName)) { + return systemBinaryName; + } + } + + return null; +} + +// Fetch latest release version from GitHub +async function getLatestVersion(repo: string): Promise { + const response = await fetch(`https://api.github.com/repos/${repo}/releases/latest`, { + headers: { "User-Agent": `${APP_NAME}-coding-agent` }, + signal: AbortSignal.timeout(NETWORK_TIMEOUT_MS), + }); + + if (!response.ok) { + throw new Error(`GitHub API error: ${response.status}`); + } + + const data = (await response.json()) as { tag_name: string }; + return data.tag_name.replace(/^v/, ""); +} + +// Download a file from URL +async function downloadFile(url: string, dest: string): Promise { + const response = await fetch(url, { + signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS), + }); + + if (!response.ok) { + throw new Error(`Failed to download: ${response.status}`); + } + + if (!response.body) { + throw new Error("No response body"); + } + + const fileStream = createWriteStream(dest); + await pipeline(Readable.fromWeb(response.body as any), fileStream); +} + +function findBinaryRecursively(rootDir: string, binaryFileName: string): string | null { + const stack: string[] = [rootDir]; + + while (stack.length > 0) { + const currentDir = stack.pop(); + if (!currentDir) continue; + + const entries = readdirSync(currentDir, { withFileTypes: true }); + for (const entry of entries) { + const fullPath = join(currentDir, entry.name); + if (entry.isFile() && entry.name === binaryFileName) { + return fullPath; + } + if (entry.isDirectory()) { + stack.push(fullPath); + } + } + } + + return null; +} + +function formatSpawnFailure(result: SpawnSyncReturns): string { + if (result.error?.message) { + return result.error.message; + } + const stderr = result.stderr?.toString().trim(); + if (stderr) { + return stderr; + } + const stdout = result.stdout?.toString().trim(); + if (stdout) { + return stdout; + } + return `exit status ${result.status ?? "unknown"}`; +} + +function runExtractionCommand(command: string, args: string[]): string | null { + const result = spawnSync(command, args, { stdio: "pipe" }); + if (!result.error && result.status === 0) { + return null; + } + return `${command}: ${formatSpawnFailure(result)}`; +} + +function extractTarGzArchive(archivePath: string, extractDir: string, assetName: string): void { + const failure = runExtractionCommand("tar", ["xzf", archivePath, "-C", extractDir]); + if (failure) { + throw new Error(`Failed to extract ${assetName}: ${failure}`); + } +} + +function getWindowsTarCommand(): string { + const systemRoot = process.env.SystemRoot ?? process.env.WINDIR; + if (systemRoot) { + const systemTar = join(systemRoot, "System32", "tar.exe"); + if (existsSync(systemTar)) { + return systemTar; + } + } + return "tar.exe"; +} + +function extractZipArchive(archivePath: string, extractDir: string, assetName: string): void { + const failures: string[] = []; + + if (platform() === "win32") { + // Windows ships bsdtar as tar.exe, which supports zip files. Prefer the + // System32 binary over Git Bash's GNU tar, which does not handle zip archives. + const tarFailure = runExtractionCommand(getWindowsTarCommand(), ["xf", archivePath, "-C", extractDir]); + if (!tarFailure) return; + failures.push(tarFailure); + + const script = + "& { param($archive, $destination) $ErrorActionPreference = 'Stop'; Expand-Archive -LiteralPath $archive -DestinationPath $destination -Force }"; + const powershellFailure = runExtractionCommand("powershell.exe", [ + "-NoLogo", + "-NoProfile", + "-NonInteractive", + "-ExecutionPolicy", + "Bypass", + "-Command", + script, + archivePath, + extractDir, + ]); + if (!powershellFailure) return; + failures.push(powershellFailure); + } else { + const unzipFailure = runExtractionCommand("unzip", ["-q", archivePath, "-d", extractDir]); + if (!unzipFailure) return; + failures.push(unzipFailure); + + const tarFailure = runExtractionCommand("tar", ["xf", archivePath, "-C", extractDir]); + if (!tarFailure) return; + failures.push(tarFailure); + } + + throw new Error(`Failed to extract ${assetName}: ${failures.join("; ")}`); +} + +// Download and install a tool +async function downloadTool(tool: "fd" | "rg"): Promise { + const config = TOOLS[tool]; + if (!config) throw new Error(`Unknown tool: ${tool}`); + + const plat = platform(); + const architecture = arch(); + + // Get latest version + let version = await getLatestVersion(config.repo); + if (tool === "fd" && plat === "darwin" && architecture === "x64") { + version = "10.3.0"; + } + + // Get asset name for this platform + const assetName = config.getAssetName(version, plat, architecture); + if (!assetName) { + throw new Error(`Unsupported platform: ${plat}/${architecture}`); + } + + // Create tools directory + mkdirSync(TOOLS_DIR, { recursive: true }); + + const downloadUrl = `https://github.com/${config.repo}/releases/download/${config.tagPrefix}${version}/${assetName}`; + const archivePath = join(TOOLS_DIR, assetName); + const binaryExt = plat === "win32" ? ".exe" : ""; + const binaryPath = join(TOOLS_DIR, config.binaryName + binaryExt); + + // Download + await downloadFile(downloadUrl, archivePath); + + // Extract into a unique temp directory. fd and rg downloads can run concurrently + // during startup, so sharing a fixed directory causes races. + const extractDir = join( + TOOLS_DIR, + `extract_tmp_${config.binaryName}_${process.pid}_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`, + ); + mkdirSync(extractDir, { recursive: true }); + + try { + if (assetName.endsWith(".tar.gz")) { + extractTarGzArchive(archivePath, extractDir, assetName); + } else if (assetName.endsWith(".zip")) { + extractZipArchive(archivePath, extractDir, assetName); + } else { + throw new Error(`Unsupported archive format: ${assetName}`); + } + + // Find the binary in extracted files. Some archives contain files directly + // at root, others nest under a versioned subdirectory. + const binaryFileName = config.binaryName + binaryExt; + const extractedDir = join(extractDir, assetName.replace(/\.(tar\.gz|zip)$/, "")); + const extractedBinaryCandidates = [join(extractedDir, binaryFileName), join(extractDir, binaryFileName)]; + let extractedBinary = extractedBinaryCandidates.find((candidate) => existsSync(candidate)); + + if (!extractedBinary) { + extractedBinary = findBinaryRecursively(extractDir, binaryFileName) ?? undefined; + } + + if (extractedBinary) { + renameSync(extractedBinary, binaryPath); + } else { + throw new Error(`Binary not found in archive: expected ${binaryFileName} under ${extractDir}`); + } + + // Make executable (Unix only) + if (plat !== "win32") { + chmodSync(binaryPath, 0o755); + } + } finally { + // Cleanup + rmSync(archivePath, { force: true }); + rmSync(extractDir, { recursive: true, force: true }); + } + + return binaryPath; +} + +// Termux package names for tools +const TERMUX_PACKAGES: Record = { + fd: "fd", + rg: "ripgrep", +}; + +// Ensure a tool is available, downloading if necessary +// Returns the path to the tool, or null if unavailable +export async function ensureTool(tool: "fd" | "rg", silent: boolean = false): Promise { + const existingPath = getToolPath(tool); + if (existingPath) { + return existingPath; + } + + const config = TOOLS[tool]; + if (!config) return undefined; + + if (isOfflineModeEnabled()) { + if (!silent) { + console.log(chalk.yellow(`${config.name} not found. Offline mode enabled, skipping download.`)); + } + return undefined; + } + + // On Android/Termux, Linux binaries don't work due to Bionic libc incompatibility. + // Users must install via pkg. + if (platform() === "android") { + const pkgName = TERMUX_PACKAGES[tool] ?? tool; + if (!silent) { + console.log(chalk.yellow(`${config.name} not found. Install with: pkg install ${pkgName}`)); + } + return undefined; + } + + // Tool not found - download it + if (!silent) { + console.log(chalk.dim(`${config.name} not found. Downloading...`)); + } + + try { + const path = await downloadTool(tool); + if (!silent) { + console.log(chalk.dim(`${config.name} installed to ${path}`)); + } + return path; + } catch (e) { + if (!silent) { + console.log(chalk.yellow(`Failed to download ${config.name}: ${e instanceof Error ? e.message : e}`)); + } + return undefined; + } +} diff --git a/cactus-code/packages/coding-agent/tsconfig.build.json b/cactus-code/packages/coding-agent/tsconfig.build.json new file mode 100644 index 000000000..15a2436ce --- /dev/null +++ b/cactus-code/packages/coding-agent/tsconfig.build.json @@ -0,0 +1,17 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "paths": { + "@earendil-works/pi-agent-core": ["../agent/dist/index.d.ts"], + "@earendil-works/pi-agent-core/*": ["../agent/dist/*.d.ts"], + "@earendil-works/pi-ai": ["../ai/dist/index.d.ts"], + "@earendil-works/pi-ai/*": ["../ai/dist/*.d.ts", "../ai/dist/providers/*.d.ts"], + "@earendil-works/pi-tui": ["../tui/dist/index.d.ts"], + "@earendil-works/pi-tui/*": ["../tui/dist/*.d.ts", "../tui/dist/components/*.d.ts"] + }, + "rootDir": "./src" + }, + "include": ["src/**/*.ts", "src/**/*.d.ts"], + "exclude": ["node_modules", "dist"] +} diff --git a/cactus-code/packages/tui/native/darwin/prebuilds/darwin-arm64/darwin-modifiers.node b/cactus-code/packages/tui/native/darwin/prebuilds/darwin-arm64/darwin-modifiers.node new file mode 100755 index 000000000..461484269 Binary files /dev/null and b/cactus-code/packages/tui/native/darwin/prebuilds/darwin-arm64/darwin-modifiers.node differ diff --git a/cactus-code/packages/tui/native/darwin/prebuilds/darwin-x64/darwin-modifiers.node b/cactus-code/packages/tui/native/darwin/prebuilds/darwin-x64/darwin-modifiers.node new file mode 100755 index 000000000..0c82b3913 Binary files /dev/null and b/cactus-code/packages/tui/native/darwin/prebuilds/darwin-x64/darwin-modifiers.node differ diff --git a/cactus-code/packages/tui/native/darwin/src/darwin-modifiers.c b/cactus-code/packages/tui/native/darwin/src/darwin-modifiers.c new file mode 100644 index 000000000..7e612519d --- /dev/null +++ b/cactus-code/packages/tui/native/darwin/src/darwin-modifiers.c @@ -0,0 +1,70 @@ +#include +#include +#include +#include +#include + +#define NAPI_AUTO_LENGTH ((size_t)-1) + +typedef void* napi_env; +typedef void* napi_value; +typedef void* napi_callback_info; +typedef napi_value (*napi_callback)(napi_env, napi_callback_info); +typedef int (*napi_create_function_fn)(napi_env, const char*, size_t, napi_callback, void*, napi_value*); +typedef int (*napi_set_named_property_fn)(napi_env, napi_value, const char*, napi_value); +typedef int (*napi_get_boolean_fn)(napi_env, bool, napi_value*); +typedef int (*napi_get_cb_info_fn)(napi_env, napi_callback_info, size_t*, napi_value*, napi_value*, void**); +typedef int (*napi_get_value_string_utf8_fn)(napi_env, napi_value, char*, size_t, size_t*); + +static void* node_symbol(const char* name) { + return dlsym(RTLD_DEFAULT, name); +} + +static CGEventFlags modifier_mask_for_name(const char* name) { + if (strcmp(name, "shift") == 0) return kCGEventFlagMaskShift; + if (strcmp(name, "command") == 0) return kCGEventFlagMaskCommand; + if (strcmp(name, "control") == 0) return kCGEventFlagMaskControl; + if (strcmp(name, "option") == 0) return kCGEventFlagMaskAlternate; + return 0; +} + +static napi_value is_modifier_pressed(napi_env env, napi_callback_info info) { + napi_get_cb_info_fn napi_get_cb_info = (napi_get_cb_info_fn)node_symbol("napi_get_cb_info"); + napi_get_value_string_utf8_fn napi_get_value_string_utf8 = (napi_get_value_string_utf8_fn)node_symbol("napi_get_value_string_utf8"); + napi_get_boolean_fn napi_get_boolean = (napi_get_boolean_fn)node_symbol("napi_get_boolean"); + + bool pressed = false; + if (napi_get_cb_info && napi_get_value_string_utf8) { + size_t argc = 1; + napi_value args[1] = {0}; + if (napi_get_cb_info(env, info, &argc, args, 0, 0) == 0 && argc >= 1 && args[0]) { + char name[16] = {0}; + size_t copied = 0; + if (napi_get_value_string_utf8(env, args[0], name, sizeof(name), &copied) == 0) { + CGEventFlags mask = modifier_mask_for_name(name); + if (mask != 0) { + CGEventFlags flags = CGEventSourceFlagsState(kCGEventSourceStateCombinedSessionState); + pressed = (flags & mask) != 0; + } + } + } + } + + napi_value result = 0; + if (napi_get_boolean) napi_get_boolean(env, pressed, &result); + return result; +} + +__attribute__((visibility("default"))) napi_value napi_register_module_v1(napi_env env, napi_value exports) { + napi_create_function_fn napi_create_function = (napi_create_function_fn)node_symbol("napi_create_function"); + napi_set_named_property_fn napi_set_named_property = (napi_set_named_property_fn)node_symbol("napi_set_named_property"); + + napi_value fn = 0; + if (napi_create_function && + napi_set_named_property && + napi_create_function(env, "isModifierPressed", NAPI_AUTO_LENGTH, is_modifier_pressed, 0, &fn) == 0) { + napi_set_named_property(env, exports, "isModifierPressed", fn); + } + + return exports; +} diff --git a/cactus-code/packages/tui/native/win32/prebuilds/win32-arm64/win32-console-mode.node b/cactus-code/packages/tui/native/win32/prebuilds/win32-arm64/win32-console-mode.node new file mode 100644 index 000000000..42b2c77cc Binary files /dev/null and b/cactus-code/packages/tui/native/win32/prebuilds/win32-arm64/win32-console-mode.node differ diff --git a/cactus-code/packages/tui/native/win32/prebuilds/win32-x64/win32-console-mode.node b/cactus-code/packages/tui/native/win32/prebuilds/win32-x64/win32-console-mode.node new file mode 100644 index 000000000..2c6d86d87 Binary files /dev/null and b/cactus-code/packages/tui/native/win32/prebuilds/win32-x64/win32-console-mode.node differ diff --git a/cactus-code/packages/tui/native/win32/src/win32-console-mode.c b/cactus-code/packages/tui/native/win32/src/win32-console-mode.c new file mode 100644 index 000000000..d68810c70 --- /dev/null +++ b/cactus-code/packages/tui/native/win32/src/win32-console-mode.c @@ -0,0 +1,53 @@ +#include + +#ifndef ENABLE_VIRTUAL_TERMINAL_INPUT +#define ENABLE_VIRTUAL_TERMINAL_INPUT 0x0200 +#endif + +#define NAPI_AUTO_LENGTH ((unsigned long long)-1) + +typedef void* napi_env; +typedef void* napi_value; +typedef void* napi_callback_info; +typedef napi_value (__cdecl *napi_callback)(napi_env, napi_callback_info); +typedef int (__cdecl *napi_create_function_fn)(napi_env, const char*, unsigned long long, napi_callback, void*, napi_value*); +typedef int (__cdecl *napi_set_named_property_fn)(napi_env, napi_value, const char*, napi_value); +typedef int (__cdecl *napi_get_boolean_fn)(napi_env, int, napi_value*); + +static void* node_symbol(const char* name) { + HMODULE module = GetModuleHandleA(0); + void* proc = module ? (void*)GetProcAddress(module, name) : 0; + if (proc) return proc; + + module = GetModuleHandleA("node.dll"); + return module ? (void*)GetProcAddress(module, name) : 0; +} + +static napi_value __cdecl enable_virtual_terminal_input(napi_env env, napi_callback_info info) { + (void)info; + + HANDLE handle = GetStdHandle(STD_INPUT_HANDLE); + DWORD mode = 0; + int enabled = handle != INVALID_HANDLE_VALUE && + GetConsoleMode(handle, &mode) && + SetConsoleMode(handle, mode | ENABLE_VIRTUAL_TERMINAL_INPUT); + + napi_get_boolean_fn napi_get_boolean = (napi_get_boolean_fn)node_symbol("napi_get_boolean"); + napi_value result = 0; + if (napi_get_boolean) napi_get_boolean(env, enabled, &result); + return result; +} + +__declspec(dllexport) napi_value __cdecl napi_register_module_v1(napi_env env, napi_value exports) { + napi_create_function_fn napi_create_function = (napi_create_function_fn)node_symbol("napi_create_function"); + napi_set_named_property_fn napi_set_named_property = (napi_set_named_property_fn)node_symbol("napi_set_named_property"); + + napi_value fn = 0; + if (napi_create_function && + napi_set_named_property && + napi_create_function(env, "enableVirtualTerminalInput", NAPI_AUTO_LENGTH, enable_virtual_terminal_input, 0, &fn) == 0) { + napi_set_named_property(env, exports, "enableVirtualTerminalInput", fn); + } + + return exports; +} diff --git a/cactus-code/packages/tui/package.json b/cactus-code/packages/tui/package.json new file mode 100644 index 000000000..25786e929 --- /dev/null +++ b/cactus-code/packages/tui/package.json @@ -0,0 +1,45 @@ +{ + "name": "@earendil-works/pi-tui", + "version": "0.80.2", + "description": "Terminal User Interface library with differential rendering for efficient text-based applications", + "type": "module", + "main": "dist/index.js", + "scripts": { + "clean": "shx rm -rf dist", + "build": "tsgo -p tsconfig.build.json", + "prepublishOnly": "npm run clean && npm run build" + }, + "files": [ + "dist/**/*", + "native/win32/prebuilds/**/*.node", + "native/darwin/prebuilds/**/*.node", + "README.md" + ], + "keywords": [ + "tui", + "terminal", + "ui", + "text-editor", + "differential-rendering", + "typescript", + "cli" + ], + "author": "Mario Zechner", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/tui" + }, + "engines": { + "node": ">=22.19.0" + }, + "types": "./dist/index.d.ts", + "dependencies": { + "get-east-asian-width": "1.6.0", + "marked": "18.0.5" + }, + "devDependencies": { + "@types/node": "24.12.4" + } +} diff --git a/cactus-code/packages/tui/src/autocomplete.ts b/cactus-code/packages/tui/src/autocomplete.ts new file mode 100644 index 000000000..205748d88 --- /dev/null +++ b/cactus-code/packages/tui/src/autocomplete.ts @@ -0,0 +1,786 @@ +import { spawn } from "child_process"; +import { readdirSync, statSync } from "fs"; +import { homedir } from "os"; +import { basename, dirname, join } from "path"; +import { fuzzyFilter } from "./fuzzy.ts"; + +const PATH_DELIMITERS = new Set([" ", "\t", '"', "'", "="]); + +function toDisplayPath(value: string): string { + return value.replace(/\\/g, "/"); +} + +function escapeRegex(value: string): string { + return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); +} + +function buildFdPathQuery(query: string): string { + const normalized = toDisplayPath(query); + if (!normalized.includes("/")) { + return normalized; + } + + const hasTrailingSeparator = normalized.endsWith("/"); + const trimmed = normalized.replace(/^\/+|\/+$/g, ""); + if (!trimmed) { + return normalized; + } + + const separatorPattern = "[\\\\/]"; + const segments = trimmed + .split("/") + .filter(Boolean) + .map((segment) => escapeRegex(segment)); + if (segments.length === 0) { + return normalized; + } + + let pattern = segments.join(separatorPattern); + if (hasTrailingSeparator) { + pattern += separatorPattern; + } + return pattern; +} + +function findLastDelimiter(text: string): number { + for (let i = text.length - 1; i >= 0; i -= 1) { + if (PATH_DELIMITERS.has(text[i] ?? "")) { + return i; + } + } + return -1; +} + +function findUnclosedQuoteStart(text: string): number | null { + let inQuotes = false; + let quoteStart = -1; + + for (let i = 0; i < text.length; i += 1) { + if (text[i] === '"') { + inQuotes = !inQuotes; + if (inQuotes) { + quoteStart = i; + } + } + } + + return inQuotes ? quoteStart : null; +} + +function isTokenStart(text: string, index: number): boolean { + return index === 0 || PATH_DELIMITERS.has(text[index - 1] ?? ""); +} + +function extractQuotedPrefix(text: string): string | null { + const quoteStart = findUnclosedQuoteStart(text); + if (quoteStart === null) { + return null; + } + + if (quoteStart > 0 && text[quoteStart - 1] === "@") { + if (!isTokenStart(text, quoteStart - 1)) { + return null; + } + return text.slice(quoteStart - 1); + } + + if (!isTokenStart(text, quoteStart)) { + return null; + } + + return text.slice(quoteStart); +} + +function parsePathPrefix(prefix: string): { rawPrefix: string; isAtPrefix: boolean; isQuotedPrefix: boolean } { + if (prefix.startsWith('@"')) { + return { rawPrefix: prefix.slice(2), isAtPrefix: true, isQuotedPrefix: true }; + } + if (prefix.startsWith('"')) { + return { rawPrefix: prefix.slice(1), isAtPrefix: false, isQuotedPrefix: true }; + } + if (prefix.startsWith("@")) { + return { rawPrefix: prefix.slice(1), isAtPrefix: true, isQuotedPrefix: false }; + } + return { rawPrefix: prefix, isAtPrefix: false, isQuotedPrefix: false }; +} + +function buildCompletionValue( + path: string, + options: { isDirectory: boolean; isAtPrefix: boolean; isQuotedPrefix: boolean }, +): string { + const needsQuotes = options.isQuotedPrefix || path.includes(" "); + const prefix = options.isAtPrefix ? "@" : ""; + + if (!needsQuotes) { + return `${prefix}${path}`; + } + + const openQuote = `${prefix}"`; + const closeQuote = '"'; + return `${openQuote}${path}${closeQuote}`; +} + +// Use fd to walk directory tree (fast, respects .gitignore) +async function walkDirectoryWithFd( + baseDir: string, + fdPath: string, + query: string, + maxResults: number, + signal: AbortSignal, +): Promise> { + const args = [ + "--base-directory", + baseDir, + "--max-results", + String(maxResults), + "--type", + "f", + "--type", + "d", + "--follow", + "--hidden", + "--exclude", + ".git", + "--exclude", + ".git/*", + "--exclude", + ".git/**", + ]; + + if (toDisplayPath(query).includes("/")) { + args.push("--full-path"); + } + + if (query) { + args.push(buildFdPathQuery(query)); + } + + return await new Promise((resolve) => { + if (signal.aborted) { + resolve([]); + return; + } + + const child = spawn(fdPath, args, { + stdio: ["ignore", "pipe", "pipe"], + }); + let stdout = ""; + let resolved = false; + + const finish = (results: Array<{ path: string; isDirectory: boolean }>) => { + if (resolved) return; + resolved = true; + signal.removeEventListener("abort", onAbort); + resolve(results); + }; + + const onAbort = () => { + if (child.exitCode === null) { + child.kill("SIGKILL"); + } + }; + + signal.addEventListener("abort", onAbort, { once: true }); + child.stdout.setEncoding("utf-8"); + child.stdout.on("data", (chunk: string) => { + stdout += chunk; + }); + child.on("error", () => { + finish([]); + }); + child.on("close", (code) => { + if (signal.aborted || code !== 0 || !stdout) { + finish([]); + return; + } + + const lines = stdout.trim().split("\n").filter(Boolean); + const results: Array<{ path: string; isDirectory: boolean }> = []; + + for (const line of lines) { + const displayLine = toDisplayPath(line); + const hasTrailingSeparator = displayLine.endsWith("/"); + const normalizedPath = hasTrailingSeparator ? displayLine.slice(0, -1) : displayLine; + if (normalizedPath === ".git" || normalizedPath.startsWith(".git/") || normalizedPath.includes("/.git/")) { + continue; + } + + results.push({ + path: displayLine, + isDirectory: hasTrailingSeparator, + }); + } + + finish(results); + }); + }); +} + +export interface AutocompleteItem { + value: string; + label: string; + description?: string; +} + +type Awaitable = T | Promise; + +export interface SlashCommand { + name: string; + description?: string; + argumentHint?: string; + // Function to get argument completions for this command + // Returns null if no argument completion is available + getArgumentCompletions?(argumentPrefix: string): Awaitable; +} + +export interface AutocompleteSuggestions { + items: AutocompleteItem[]; + prefix: string; // What we're matching against (e.g., "/" or "src/") +} + +export interface AutocompleteProvider { + /** Characters that should naturally trigger this provider at token boundaries. */ + triggerCharacters?: string[]; + + // Get autocomplete suggestions for current text/cursor position + // Returns null if no suggestions available + getSuggestions( + lines: string[], + cursorLine: number, + cursorCol: number, + options: { signal: AbortSignal; force?: boolean }, + ): Promise; + + // Apply the selected item + // Returns the new text and cursor position + applyCompletion( + lines: string[], + cursorLine: number, + cursorCol: number, + item: AutocompleteItem, + prefix: string, + ): { + lines: string[]; + cursorLine: number; + cursorCol: number; + }; + + // Check if file completion should trigger for explicit Tab completion + shouldTriggerFileCompletion?(lines: string[], cursorLine: number, cursorCol: number): boolean; +} + +// Combined provider that handles both slash commands and file paths +export class CombinedAutocompleteProvider implements AutocompleteProvider { + private commands: (SlashCommand | AutocompleteItem)[]; + private basePath: string; + private fdPath: string | null; + + constructor(commands: (SlashCommand | AutocompleteItem)[] = [], basePath: string, fdPath: string | null = null) { + this.commands = commands; + this.basePath = basePath; + this.fdPath = fdPath; + } + + async getSuggestions( + lines: string[], + cursorLine: number, + cursorCol: number, + options: { signal: AbortSignal; force?: boolean }, + ): Promise { + const currentLine = lines[cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, cursorCol); + + const atPrefix = this.extractAtPrefix(textBeforeCursor); + if (atPrefix) { + const { rawPrefix, isQuotedPrefix } = parsePathPrefix(atPrefix); + const suggestions = await this.getFuzzyFileSuggestions(rawPrefix, { + isQuotedPrefix, + signal: options.signal, + }); + if (suggestions.length === 0) return null; + + return { + items: suggestions, + prefix: atPrefix, + }; + } + + if (!options.force && textBeforeCursor.startsWith("/")) { + const spaceIndex = textBeforeCursor.indexOf(" "); + + if (spaceIndex === -1) { + const prefix = textBeforeCursor.slice(1); + const commandItems = this.commands.map((cmd) => { + const name = "name" in cmd ? cmd.name : cmd.value; + const hint = "argumentHint" in cmd && cmd.argumentHint ? cmd.argumentHint : undefined; + const desc = cmd.description ?? ""; + const fullDesc = hint ? (desc ? `${hint} — ${desc}` : hint) : desc; + return { + name, + label: name, + description: fullDesc || undefined, + }; + }); + + const filtered = fuzzyFilter(commandItems, prefix, (item) => item.name).map((item) => ({ + value: item.name, + label: item.label, + ...(item.description && { description: item.description }), + })); + + if (filtered.length === 0) return null; + + return { + items: filtered, + prefix: textBeforeCursor, + }; + } + + const commandName = textBeforeCursor.slice(1, spaceIndex); + const argumentText = textBeforeCursor.slice(spaceIndex + 1); + + const command = this.commands.find((cmd) => { + const name = "name" in cmd ? cmd.name : cmd.value; + return name === commandName; + }); + if (!command || !("getArgumentCompletions" in command) || !command.getArgumentCompletions) { + return null; + } + + const argumentSuggestions = await command.getArgumentCompletions(argumentText); + if (!Array.isArray(argumentSuggestions) || argumentSuggestions.length === 0) { + return null; + } + + return { + items: argumentSuggestions, + prefix: argumentText, + }; + } + + const pathMatch = this.extractPathPrefix(textBeforeCursor, options.force ?? false); + if (pathMatch === null) { + return null; + } + + const suggestions = this.getFileSuggestions(pathMatch); + if (suggestions.length === 0) return null; + + return { + items: suggestions, + prefix: pathMatch, + }; + } + + applyCompletion( + lines: string[], + cursorLine: number, + cursorCol: number, + item: AutocompleteItem, + prefix: string, + ): { lines: string[]; cursorLine: number; cursorCol: number } { + const currentLine = lines[cursorLine] || ""; + const beforePrefix = currentLine.slice(0, cursorCol - prefix.length); + const afterCursor = currentLine.slice(cursorCol); + const isQuotedPrefix = prefix.startsWith('"') || prefix.startsWith('@"'); + const hasLeadingQuoteAfterCursor = afterCursor.startsWith('"'); + const hasTrailingQuoteInItem = item.value.endsWith('"'); + const adjustedAfterCursor = + isQuotedPrefix && hasTrailingQuoteInItem && hasLeadingQuoteAfterCursor ? afterCursor.slice(1) : afterCursor; + + // Check if we're completing a slash command (prefix starts with "/" but NOT a file path) + // Slash commands are at the start of the line and don't contain path separators after the first / + const isSlashCommand = prefix.startsWith("/") && beforePrefix.trim() === "" && !prefix.slice(1).includes("/"); + if (isSlashCommand) { + // This is a command name completion + const newLine = `${beforePrefix}/${item.value} ${adjustedAfterCursor}`; + const newLines = [...lines]; + newLines[cursorLine] = newLine; + + return { + lines: newLines, + cursorLine, + cursorCol: beforePrefix.length + item.value.length + 2, // +2 for "/" and space + }; + } + + // Check if we're completing a file attachment (prefix starts with "@") + if (prefix.startsWith("@")) { + // This is a file attachment completion + // Don't add space after directories so user can continue autocompleting + const isDirectory = item.label.endsWith("/"); + const suffix = isDirectory ? "" : " "; + const newLine = `${beforePrefix + item.value}${suffix}${adjustedAfterCursor}`; + const newLines = [...lines]; + newLines[cursorLine] = newLine; + + const hasTrailingQuote = item.value.endsWith('"'); + const cursorOffset = isDirectory && hasTrailingQuote ? item.value.length - 1 : item.value.length; + + return { + lines: newLines, + cursorLine, + cursorCol: beforePrefix.length + cursorOffset + suffix.length, + }; + } + + // Check if we're in a slash command context (beforePrefix contains "/command ") + const textBeforeCursor = currentLine.slice(0, cursorCol); + if (textBeforeCursor.includes("/") && textBeforeCursor.includes(" ")) { + // This is likely a command argument completion + const newLine = beforePrefix + item.value + adjustedAfterCursor; + const newLines = [...lines]; + newLines[cursorLine] = newLine; + + const isDirectory = item.label.endsWith("/"); + const hasTrailingQuote = item.value.endsWith('"'); + const cursorOffset = isDirectory && hasTrailingQuote ? item.value.length - 1 : item.value.length; + + return { + lines: newLines, + cursorLine, + cursorCol: beforePrefix.length + cursorOffset, + }; + } + + // For file paths, complete the path + const newLine = beforePrefix + item.value + adjustedAfterCursor; + const newLines = [...lines]; + newLines[cursorLine] = newLine; + + const isDirectory = item.label.endsWith("/"); + const hasTrailingQuote = item.value.endsWith('"'); + const cursorOffset = isDirectory && hasTrailingQuote ? item.value.length - 1 : item.value.length; + + return { + lines: newLines, + cursorLine, + cursorCol: beforePrefix.length + cursorOffset, + }; + } + + // Extract @ prefix for fuzzy file suggestions + private extractAtPrefix(text: string): string | null { + const quotedPrefix = extractQuotedPrefix(text); + if (quotedPrefix?.startsWith('@"')) { + return quotedPrefix; + } + + const lastDelimiterIndex = findLastDelimiter(text); + const tokenStart = lastDelimiterIndex === -1 ? 0 : lastDelimiterIndex + 1; + + if (text[tokenStart] === "@") { + return text.slice(tokenStart); + } + + return null; + } + + // Extract a path-like prefix from the text before cursor + private extractPathPrefix(text: string, forceExtract: boolean = false): string | null { + const quotedPrefix = extractQuotedPrefix(text); + if (quotedPrefix) { + return quotedPrefix; + } + + const lastDelimiterIndex = findLastDelimiter(text); + const pathPrefix = lastDelimiterIndex === -1 ? text : text.slice(lastDelimiterIndex + 1); + + // For forced extraction (Tab key), always return something + if (forceExtract) { + return pathPrefix; + } + + // For natural triggers, return if it looks like a path, ends with /, starts with ~/, . + // Only return empty string if the text looks like it's starting a path context + if (pathPrefix.includes("/") || pathPrefix.startsWith(".") || pathPrefix.startsWith("~/")) { + return pathPrefix; + } + + // Return empty string only after a space (not for completely empty text) + // Empty text should not trigger file suggestions - that's for forced Tab completion + if (pathPrefix === "" && text.endsWith(" ")) { + return pathPrefix; + } + + return null; + } + + // Expand home directory (~/) to actual home path + private expandHomePath(path: string): string { + if (path.startsWith("~/")) { + const expandedPath = join(homedir(), path.slice(2)); + // Preserve trailing slash if original path had one + return path.endsWith("/") && !expandedPath.endsWith("/") ? `${expandedPath}/` : expandedPath; + } else if (path === "~") { + return homedir(); + } + return path; + } + + private resolveScopedFuzzyQuery(rawQuery: string): { baseDir: string; query: string; displayBase: string } | null { + const normalizedQuery = toDisplayPath(rawQuery); + const slashIndex = normalizedQuery.lastIndexOf("/"); + if (slashIndex === -1) { + return null; + } + + const displayBase = normalizedQuery.slice(0, slashIndex + 1); + const query = normalizedQuery.slice(slashIndex + 1); + + let baseDir: string; + if (displayBase.startsWith("~/")) { + baseDir = this.expandHomePath(displayBase); + } else if (displayBase.startsWith("/")) { + baseDir = displayBase; + } else { + baseDir = join(this.basePath, displayBase); + } + + try { + if (!statSync(baseDir).isDirectory()) { + return null; + } + } catch { + return null; + } + + return { baseDir, query, displayBase }; + } + + private scopedPathForDisplay(displayBase: string, relativePath: string): string { + const normalizedRelativePath = toDisplayPath(relativePath); + if (displayBase === "/") { + return `/${normalizedRelativePath}`; + } + return `${toDisplayPath(displayBase)}${normalizedRelativePath}`; + } + + // Get file/directory suggestions for a given path prefix + private getFileSuggestions(prefix: string): AutocompleteItem[] { + try { + let searchDir: string; + let searchPrefix: string; + const { rawPrefix, isAtPrefix, isQuotedPrefix } = parsePathPrefix(prefix); + let expandedPrefix = rawPrefix; + + // Handle home directory expansion + if (expandedPrefix.startsWith("~")) { + expandedPrefix = this.expandHomePath(expandedPrefix); + } + + const isRootPrefix = + rawPrefix === "" || + rawPrefix === "./" || + rawPrefix === "../" || + rawPrefix === "~" || + rawPrefix === "~/" || + rawPrefix === "/" || + (isAtPrefix && rawPrefix === ""); + + if (isRootPrefix) { + // Complete from specified position + if (rawPrefix.startsWith("~") || expandedPrefix.startsWith("/")) { + searchDir = expandedPrefix; + } else { + searchDir = join(this.basePath, expandedPrefix); + } + searchPrefix = ""; + } else if (rawPrefix.endsWith("/")) { + // If prefix ends with /, show contents of that directory + if (rawPrefix.startsWith("~") || expandedPrefix.startsWith("/")) { + searchDir = expandedPrefix; + } else { + searchDir = join(this.basePath, expandedPrefix); + } + searchPrefix = ""; + } else { + // Split into directory and file prefix + const dir = dirname(expandedPrefix); + const file = basename(expandedPrefix); + if (rawPrefix.startsWith("~") || expandedPrefix.startsWith("/")) { + searchDir = dir; + } else { + searchDir = join(this.basePath, dir); + } + searchPrefix = file; + } + + const entries = readdirSync(searchDir, { withFileTypes: true }); + const suggestions: AutocompleteItem[] = []; + + for (const entry of entries) { + if (!entry.name.toLowerCase().startsWith(searchPrefix.toLowerCase())) { + continue; + } + + // Check if entry is a directory (or a symlink pointing to a directory) + let isDirectory = entry.isDirectory(); + if (!isDirectory && entry.isSymbolicLink()) { + try { + const fullPath = join(searchDir, entry.name); + isDirectory = statSync(fullPath).isDirectory(); + } catch { + // Broken symlink or permission error - treat as file + } + } + + let relativePath: string; + const name = entry.name; + const displayPrefix = rawPrefix; + + if (displayPrefix.endsWith("/")) { + // If prefix ends with /, append entry to the prefix + relativePath = displayPrefix + name; + } else if (displayPrefix.includes("/") || displayPrefix.includes("\\")) { + // Preserve ~/ format for home directory paths + if (displayPrefix.startsWith("~/")) { + const homeRelativeDir = displayPrefix.slice(2); // Remove ~/ + const dir = dirname(homeRelativeDir); + relativePath = `~/${dir === "." ? name : join(dir, name)}`; + } else if (displayPrefix.startsWith("/")) { + // Absolute path - construct properly + const dir = dirname(displayPrefix); + if (dir === "/") { + relativePath = `/${name}`; + } else { + relativePath = `${dir}/${name}`; + } + } else { + relativePath = join(dirname(displayPrefix), name); + // path.join normalizes away ./ prefix, preserve it + if (displayPrefix.startsWith("./") && !relativePath.startsWith("./")) { + relativePath = `./${relativePath}`; + } + } + } else { + // For standalone entries, preserve ~/ if original prefix was ~/ + if (displayPrefix.startsWith("~")) { + relativePath = `~/${name}`; + } else { + relativePath = name; + } + } + + relativePath = toDisplayPath(relativePath); + const pathValue = isDirectory ? `${relativePath}/` : relativePath; + const value = buildCompletionValue(pathValue, { + isDirectory, + isAtPrefix, + isQuotedPrefix, + }); + + suggestions.push({ + value, + label: name + (isDirectory ? "/" : ""), + }); + } + + // Sort directories first, then alphabetically + suggestions.sort((a, b) => { + const aIsDir = a.value.endsWith("/"); + const bIsDir = b.value.endsWith("/"); + if (aIsDir && !bIsDir) return -1; + if (!aIsDir && bIsDir) return 1; + return a.label.localeCompare(b.label); + }); + + return suggestions; + } catch (_e) { + // Directory doesn't exist or not accessible + return []; + } + } + + // Score an entry against the query (higher = better match) + // isDirectory adds bonus to prioritize folders + private scoreEntry(filePath: string, query: string, isDirectory: boolean): number { + const fileName = basename(filePath); + const lowerFileName = fileName.toLowerCase(); + const lowerQuery = query.toLowerCase(); + + let score = 0; + + // Exact filename match (highest) + if (lowerFileName === lowerQuery) score = 100; + // Filename starts with query + else if (lowerFileName.startsWith(lowerQuery)) score = 80; + // Substring match in filename + else if (lowerFileName.includes(lowerQuery)) score = 50; + // Substring match in full path + else if (filePath.toLowerCase().includes(lowerQuery)) score = 30; + + // Directories get a bonus to appear first + if (isDirectory && score > 0) score += 10; + + return score; + } + + // Fuzzy file search using fd (fast, respects .gitignore) + private async getFuzzyFileSuggestions( + query: string, + options: { isQuotedPrefix: boolean; signal: AbortSignal }, + ): Promise { + if (!this.fdPath || options.signal.aborted) { + return []; + } + + try { + const scopedQuery = this.resolveScopedFuzzyQuery(query); + const fdBaseDir = scopedQuery?.baseDir ?? this.basePath; + const fdQuery = scopedQuery?.query ?? query; + const entries = await walkDirectoryWithFd(fdBaseDir, this.fdPath, fdQuery, 100, options.signal); + if (options.signal.aborted) { + return []; + } + + const scoredEntries = entries + .map((entry) => ({ + ...entry, + score: fdQuery ? this.scoreEntry(entry.path, fdQuery, entry.isDirectory) : 1, + })) + .filter((entry) => entry.score > 0); + + scoredEntries.sort((a, b) => b.score - a.score); + const topEntries = scoredEntries.slice(0, 20); + + const suggestions: AutocompleteItem[] = []; + for (const { path: entryPath, isDirectory } of topEntries) { + const pathWithoutSlash = isDirectory ? entryPath.slice(0, -1) : entryPath; + const displayPath = scopedQuery + ? this.scopedPathForDisplay(scopedQuery.displayBase, pathWithoutSlash) + : pathWithoutSlash; + const entryName = basename(pathWithoutSlash); + const completionPath = isDirectory ? `${displayPath}/` : displayPath; + const value = buildCompletionValue(completionPath, { + isDirectory, + isAtPrefix: true, + isQuotedPrefix: options.isQuotedPrefix, + }); + + suggestions.push({ + value, + label: entryName + (isDirectory ? "/" : ""), + description: displayPath, + }); + } + + return suggestions; + } catch { + return []; + } + } + + // Check if we should trigger file completion (called on Tab key) + shouldTriggerFileCompletion(lines: string[], cursorLine: number, cursorCol: number): boolean { + const currentLine = lines[cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, cursorCol); + + // Don't trigger if we're typing a slash command at the start of the line + if (textBeforeCursor.trim().startsWith("/") && !textBeforeCursor.trim().includes(" ")) { + return false; + } + + return true; + } +} diff --git a/cactus-code/packages/tui/src/components/box.ts b/cactus-code/packages/tui/src/components/box.ts new file mode 100644 index 000000000..3573ab092 --- /dev/null +++ b/cactus-code/packages/tui/src/components/box.ts @@ -0,0 +1,137 @@ +import type { Component } from "../tui.ts"; +import { applyBackgroundToLine, visibleWidth } from "../utils.ts"; + +type RenderCache = { + childLines: string[]; + width: number; + bgSample: string | undefined; + lines: string[]; +}; + +/** + * Box component - a container that applies padding and background to all children + */ +export class Box implements Component { + children: Component[] = []; + private paddingX: number; + private paddingY: number; + private bgFn?: (text: string) => string; + + // Cache for rendered output + private cache?: RenderCache; + + constructor(paddingX = 1, paddingY = 1, bgFn?: (text: string) => string) { + this.paddingX = paddingX; + this.paddingY = paddingY; + this.bgFn = bgFn; + } + + addChild(component: Component): void { + this.children.push(component); + this.invalidateCache(); + } + + removeChild(component: Component): void { + const index = this.children.indexOf(component); + if (index !== -1) { + this.children.splice(index, 1); + this.invalidateCache(); + } + } + + clear(): void { + this.children = []; + this.invalidateCache(); + } + + setBgFn(bgFn?: (text: string) => string): void { + this.bgFn = bgFn; + // Don't invalidate here - we'll detect bgFn changes by sampling output + } + + private invalidateCache(): void { + this.cache = undefined; + } + + private matchCache(width: number, childLines: string[], bgSample: string | undefined): boolean { + const cache = this.cache; + return ( + !!cache && + cache.width === width && + cache.bgSample === bgSample && + cache.childLines.length === childLines.length && + cache.childLines.every((line, i) => line === childLines[i]) + ); + } + + invalidate(): void { + this.invalidateCache(); + for (const child of this.children) { + child.invalidate?.(); + } + } + + render(width: number): string[] { + if (this.children.length === 0) { + return []; + } + + const contentWidth = Math.max(1, width - this.paddingX * 2); + const leftPad = " ".repeat(this.paddingX); + + // Render all children + const childLines: string[] = []; + for (const child of this.children) { + const lines = child.render(contentWidth); + for (const line of lines) { + childLines.push(leftPad + line); + } + } + + if (childLines.length === 0) { + return []; + } + + // Check if bgFn output changed by sampling + const bgSample = this.bgFn ? this.bgFn("test") : undefined; + + // Check cache validity + if (this.matchCache(width, childLines, bgSample)) { + return this.cache!.lines; + } + + // Apply background and padding + const result: string[] = []; + + // Top padding + for (let i = 0; i < this.paddingY; i++) { + result.push(this.applyBg("", width)); + } + + // Content + for (const line of childLines) { + result.push(this.applyBg(line, width)); + } + + // Bottom padding + for (let i = 0; i < this.paddingY; i++) { + result.push(this.applyBg("", width)); + } + + // Update cache + this.cache = { childLines, width, bgSample, lines: result }; + + return result; + } + + private applyBg(line: string, width: number): string { + const visLen = visibleWidth(line); + const padNeeded = Math.max(0, width - visLen); + const padded = line + " ".repeat(padNeeded); + + if (this.bgFn) { + return applyBackgroundToLine(padded, width, this.bgFn); + } + return padded; + } +} diff --git a/cactus-code/packages/tui/src/components/cancellable-loader.ts b/cactus-code/packages/tui/src/components/cancellable-loader.ts new file mode 100644 index 000000000..7822cb042 --- /dev/null +++ b/cactus-code/packages/tui/src/components/cancellable-loader.ts @@ -0,0 +1,40 @@ +import { getKeybindings } from "../keybindings.ts"; +import { Loader } from "./loader.ts"; + +/** + * Loader that can be cancelled with Escape. + * Extends Loader with an AbortSignal for cancelling async operations. + * + * @example + * const loader = new CancellableLoader(tui, cyan, dim, "Working..."); + * loader.onAbort = () => done(null); + * doWork(loader.signal).then(done); + */ +export class CancellableLoader extends Loader { + private abortController = new AbortController(); + + /** Called when user presses Escape */ + onAbort?: () => void; + + /** AbortSignal that is aborted when user presses Escape */ + get signal(): AbortSignal { + return this.abortController.signal; + } + + /** Whether the loader was aborted */ + get aborted(): boolean { + return this.abortController.signal.aborted; + } + + handleInput(data: string): void { + const kb = getKeybindings(); + if (kb.matches(data, "tui.select.cancel")) { + this.abortController.abort(); + this.onAbort?.(); + } + } + + dispose(): void { + this.stop(); + } +} diff --git a/cactus-code/packages/tui/src/components/editor.ts b/cactus-code/packages/tui/src/components/editor.ts new file mode 100644 index 000000000..bedd01057 --- /dev/null +++ b/cactus-code/packages/tui/src/components/editor.ts @@ -0,0 +1,2307 @@ +import type { AutocompleteProvider, AutocompleteSuggestions } from "../autocomplete.ts"; +import { getKeybindings } from "../keybindings.ts"; +import { decodePrintableKey, matchesKey } from "../keys.ts"; +import { KillRing } from "../kill-ring.ts"; +import { type Component, CURSOR_MARKER, type Focusable, type TUI } from "../tui.ts"; +import { UndoStack } from "../undo-stack.ts"; +import { + cjkBreakRegex, + getGraphemeSegmenter, + getWordSegmenter, + isWhitespaceChar, + truncateToWidth, + visibleWidth, +} from "../utils.ts"; +import { findWordBackward, findWordForward } from "../word-navigation.ts"; +import { SelectList, type SelectListLayoutOptions, type SelectListTheme } from "./select-list.ts"; + +const graphemeSegmenter = getGraphemeSegmenter(); +const wordSegmenter = getWordSegmenter(); + +/** Regex matching paste markers like `[paste #1 +123 lines]` or `[paste #2 1234 chars]`. */ +const PASTE_MARKER_REGEX = /\[paste #(\d+)( (\+\d+ lines|\d+ chars))?\]/g; + +/** Non-global version for single-segment testing. */ +const PASTE_MARKER_SINGLE = /^\[paste #(\d+)( (\+\d+ lines|\d+ chars))?\]$/; + +/** Check if a segment is a paste marker (i.e. was merged by segmentWithMarkers). */ +function isPasteMarker(segment: string): boolean { + return segment.length >= 10 && PASTE_MARKER_SINGLE.test(segment); +} + +/** + * A segmenter that wraps Intl.Segmenter and merges graphemes that fall + * within paste markers into single atomic segments. This makes cursor + * movement, deletion, word-wrap, etc. treat paste markers as single units. + * + * Only markers whose numeric ID exists in `validIds` are merged. + */ +function segmentWithMarkers( + text: string, + baseSegmenter: Intl.Segmenter, + validIds: Set, +): Iterable { + // Fast path: no paste markers in the text or no valid IDs. + if (validIds.size === 0 || !text.includes("[paste #")) { + return baseSegmenter.segment(text); + } + + // Find all marker spans with valid IDs. + const markers: Array<{ start: number; end: number }> = []; + for (const m of text.matchAll(PASTE_MARKER_REGEX)) { + const id = Number.parseInt(m[1]!, 10); + if (!validIds.has(id)) continue; + markers.push({ start: m.index, end: m.index + m[0].length }); + } + if (markers.length === 0) { + return baseSegmenter.segment(text); + } + + // Build merged segment list. + const baseSegments = baseSegmenter.segment(text); + const result: Intl.SegmentData[] = []; + let markerIdx = 0; + + for (const seg of baseSegments) { + // Skip past markers that are entirely before this segment. + while (markerIdx < markers.length && markers[markerIdx]!.end <= seg.index) { + markerIdx++; + } + + const marker = markerIdx < markers.length ? markers[markerIdx]! : null; + + if (marker && seg.index >= marker.start && seg.index < marker.end) { + // This segment falls inside a marker. + // If this is the first segment of the marker, emit a merged segment. + if (seg.index === marker.start) { + const markerText = text.slice(marker.start, marker.end); + result.push({ + segment: markerText, + index: marker.start, + input: text, + }); + } + // Otherwise skip (already merged into the first segment). + } else { + result.push(seg); + } + } + + return result; +} + +/** + * Represents a chunk of text for word-wrap layout. + * Tracks both the text content and its position in the original line. + */ +export interface TextChunk { + text: string; + startIndex: number; + endIndex: number; +} + +/** + * Split a line into word-wrapped chunks. + * Wraps at word boundaries when possible, falling back to character-level + * wrapping for words longer than the available width. + * + * @param line - The text line to wrap + * @param maxWidth - Maximum visible width per chunk + * @param preSegmented - Optional pre-segmented graphemes (e.g. with paste-marker awareness). + * When omitted the default Intl.Segmenter is used. + * @returns Array of chunks with text and position information + */ +export function wordWrapLine(line: string, maxWidth: number, preSegmented?: Intl.SegmentData[]): TextChunk[] { + if (!line || maxWidth <= 0) { + return [{ text: "", startIndex: 0, endIndex: 0 }]; + } + + const lineWidth = visibleWidth(line); + if (lineWidth <= maxWidth) { + return [{ text: line, startIndex: 0, endIndex: line.length }]; + } + + const chunks: TextChunk[] = []; + const segments = preSegmented ?? [...graphemeSegmenter.segment(line)]; + + let currentWidth = 0; + let chunkStart = 0; + + // Wrap opportunity: the position after the last whitespace before a non-whitespace + // grapheme, i.e. where a line break is allowed. + let wrapOppIndex = -1; + let wrapOppWidth = 0; + + for (let i = 0; i < segments.length; i++) { + const seg = segments[i]!; + const grapheme = seg.segment; + const gWidth = visibleWidth(grapheme); + const charIndex = seg.index; + const isWs = !isPasteMarker(grapheme) && isWhitespaceChar(grapheme); + + // Overflow check before advancing. + if (currentWidth + gWidth > maxWidth) { + if (wrapOppIndex >= 0 && currentWidth - wrapOppWidth + gWidth <= maxWidth) { + // Backtrack to last wrap opportunity (the remaining content + // plus the current grapheme still fits within maxWidth). + chunks.push({ text: line.slice(chunkStart, wrapOppIndex), startIndex: chunkStart, endIndex: wrapOppIndex }); + chunkStart = wrapOppIndex; + currentWidth -= wrapOppWidth; + } else if (chunkStart < charIndex) { + // No viable wrap opportunity: force-break at current position. + // This also handles the case where backtracking to a word + // boundary wouldn't help because the remaining content plus + // the current grapheme (e.g. a wide character) still exceeds + // maxWidth. + chunks.push({ text: line.slice(chunkStart, charIndex), startIndex: chunkStart, endIndex: charIndex }); + chunkStart = charIndex; + currentWidth = 0; + } + wrapOppIndex = -1; + } + + if (gWidth > maxWidth) { + // Single atomic segment wider than maxWidth (e.g. paste marker + // in a narrow terminal). Re-wrap it at grapheme granularity. + + // The segment remains logically atomic for cursor + // movement / editing — the split is purely visual for word-wrap layout. + const subChunks = wordWrapLine(grapheme, maxWidth); + for (let j = 0; j < subChunks.length - 1; j++) { + const sc = subChunks[j]!; + chunks.push({ text: sc.text, startIndex: charIndex + sc.startIndex, endIndex: charIndex + sc.endIndex }); + } + const last = subChunks[subChunks.length - 1]!; + chunkStart = charIndex + last.startIndex; + currentWidth = visibleWidth(last.text); + wrapOppIndex = -1; + continue; + } + + // Advance. + currentWidth += gWidth; + + // Record wrap opportunity: whitespace followed by non-whitespace + // (multiple spaces join; the break point is after the last space), + // or at a boundary where either side is CJK (CJK allows breaking + // between any adjacent characters). + const next = segments[i + 1]; + if (isWs && next && (isPasteMarker(next.segment) || !isWhitespaceChar(next.segment))) { + wrapOppIndex = next.index; + wrapOppWidth = currentWidth; + } else if (!isWs && next && !isWhitespaceChar(next.segment)) { + const isCjk = !isPasteMarker(grapheme) && cjkBreakRegex.test(grapheme); + const nextIsCjk = !isPasteMarker(next.segment) && cjkBreakRegex.test(next.segment); + if (isCjk || nextIsCjk) { + wrapOppIndex = next.index; + wrapOppWidth = currentWidth; + } + } + } + + // Push final chunk. + chunks.push({ text: line.slice(chunkStart), startIndex: chunkStart, endIndex: line.length }); + + return chunks; +} + +// Kitty CSI-u sequences for printable keys, including optional shifted/base codepoints. +interface EditorState { + lines: string[]; + cursorLine: number; + cursorCol: number; +} + +interface LayoutLine { + text: string; + hasCursor: boolean; + cursorPos?: number; +} + +export interface EditorTheme { + borderColor: (str: string) => string; + selectList: SelectListTheme; +} + +export interface EditorOptions { + paddingX?: number; + autocompleteMaxVisible?: number; +} + +const SLASH_COMMAND_SELECT_LIST_LAYOUT: SelectListLayoutOptions = { + minPrimaryColumnWidth: 12, + maxPrimaryColumnWidth: 32, +}; + +const ATTACHMENT_AUTOCOMPLETE_DEBOUNCE_MS = 20; +const DEFAULT_AUTOCOMPLETE_TRIGGER_CHARACTERS = ["@", "#"]; + +function escapeCharacterClass(value: string): string { + return value.replace(/[\\^$.*+?()[\]{}|-]/g, "\\$&"); +} + +function buildTriggerPattern(triggerCharacters: string[]): RegExp { + return new RegExp(`(?:^|[\\s])[${triggerCharacters.map(escapeCharacterClass).join("")}][^\\s]*$`); +} + +function buildDebouncePattern(triggerCharacters: string[]): RegExp { + const escapedWithoutAt = triggerCharacters.filter((character) => character !== "@").map(escapeCharacterClass); + return new RegExp(`(?:^|[ \\t])(?:@(?:"[^"]*|[^\\s]*)|[${escapedWithoutAt.join("")}][^\\s]*)$`); +} + +export class Editor implements Component, Focusable { + private state: EditorState = { + lines: [""], + cursorLine: 0, + cursorCol: 0, + }; + + /** Focusable interface - set by TUI when focus changes */ + focused: boolean = false; + + protected tui: TUI; + private theme: EditorTheme; + private paddingX: number = 0; + + // Store last render width for cursor navigation + private lastWidth: number = 80; + + // Vertical scrolling support + private scrollOffset: number = 0; + + // Border color (can be changed dynamically) + public borderColor: (str: string) => string; + + // Autocomplete support + private autocompleteProvider?: AutocompleteProvider; + private autocompleteTriggerCharacters = [...DEFAULT_AUTOCOMPLETE_TRIGGER_CHARACTERS]; + private autocompleteTriggerPattern = buildTriggerPattern(this.autocompleteTriggerCharacters); + private autocompleteDebouncePattern = buildDebouncePattern(this.autocompleteTriggerCharacters); + private autocompleteList?: SelectList; + private autocompleteState: "regular" | "force" | null = null; + private autocompletePrefix: string = ""; + private autocompleteMaxVisible: number = 5; + private autocompleteAbort?: AbortController; + private autocompleteDebounceTimer?: ReturnType; + private autocompleteRequestTask: Promise = Promise.resolve(); + private autocompleteStartToken: number = 0; + private autocompleteRequestId: number = 0; + + // Paste tracking for large pastes + private pastes: Map = new Map(); + private pasteCounter: number = 0; + + // Bracketed paste mode buffering + private pasteBuffer: string = ""; + private isInPaste: boolean = false; + + // Prompt history for up/down navigation + private history: string[] = []; + private historyIndex: number = -1; // -1 = not browsing, 0 = most recent, 1 = older, etc. + private historyDraft: EditorState | null = null; + + // Kill ring for Emacs-style kill/yank operations + private killRing = new KillRing(); + private lastAction: "kill" | "yank" | "type-word" | null = null; + + // Character jump mode + private jumpMode: "forward" | "backward" | null = null; + + // Preferred visual column for vertical cursor movement (sticky column) + private preferredVisualCol: number | null = null; + + // When the cursor is snapped to the start of an atomic segment, e.g. a + // paste marker, cursorCol no longer reflects where the cursor would have + // landed. This field stores the pre-snap cursorCol so that the next + // vertical move can resolve it to a visual column on whatever VL it belongs + // to. + private snappedFromCursorCol: number | null = null; + + // Undo support + private undoStack = new UndoStack(); + + public onSubmit?: (text: string) => void; + public onChange?: (text: string) => void; + public disableSubmit: boolean = false; + + constructor(tui: TUI, theme: EditorTheme, options: EditorOptions = {}) { + this.tui = tui; + this.theme = theme; + this.borderColor = theme.borderColor; + const paddingX = options.paddingX ?? 0; + this.paddingX = Number.isFinite(paddingX) ? Math.max(0, Math.floor(paddingX)) : 0; + const maxVisible = options.autocompleteMaxVisible ?? 5; + this.autocompleteMaxVisible = Number.isFinite(maxVisible) ? Math.max(3, Math.min(20, Math.floor(maxVisible))) : 5; + } + + /** Set of currently valid paste IDs, for marker-aware segmentation. */ + private validPasteIds(): Set { + return new Set(this.pastes.keys()); + } + + /** Segment text with paste-marker awareness, only merging markers with valid IDs. */ + private segment(text: string, mode: "word" | "grapheme"): Iterable { + return segmentWithMarkers(text, mode === "word" ? wordSegmenter : graphemeSegmenter, this.validPasteIds()); + } + + getPaddingX(): number { + return this.paddingX; + } + + setPaddingX(padding: number): void { + const newPadding = Number.isFinite(padding) ? Math.max(0, Math.floor(padding)) : 0; + if (this.paddingX !== newPadding) { + this.paddingX = newPadding; + this.tui.requestRender(); + } + } + + getAutocompleteMaxVisible(): number { + return this.autocompleteMaxVisible; + } + + setAutocompleteMaxVisible(maxVisible: number): void { + const newMaxVisible = Number.isFinite(maxVisible) ? Math.max(3, Math.min(20, Math.floor(maxVisible))) : 5; + if (this.autocompleteMaxVisible !== newMaxVisible) { + this.autocompleteMaxVisible = newMaxVisible; + this.tui.requestRender(); + } + } + + setAutocompleteProvider(provider: AutocompleteProvider): void { + this.cancelAutocomplete(); + this.autocompleteProvider = provider; + this.setAutocompleteTriggerCharacters(provider.triggerCharacters ?? []); + } + + /** + * Add a prompt to history for up/down arrow navigation. + * Called after successful submission. + */ + addToHistory(text: string): void { + const trimmed = text.trim(); + if (!trimmed) return; + // Don't add consecutive duplicates + if (this.history.length > 0 && this.history[0] === trimmed) return; + this.history.unshift(trimmed); + // Limit history size + if (this.history.length > 100) { + this.history.pop(); + } + } + + private isEditorEmpty(): boolean { + return this.state.lines.length === 1 && this.state.lines[0] === ""; + } + + private isOnFirstVisualLine(): boolean { + const visualLines = this.buildVisualLineMap(this.lastWidth); + const currentVisualLine = this.findCurrentVisualLine(visualLines); + return currentVisualLine === 0; + } + + private isOnLastVisualLine(): boolean { + const visualLines = this.buildVisualLineMap(this.lastWidth); + const currentVisualLine = this.findCurrentVisualLine(visualLines); + return currentVisualLine === visualLines.length - 1; + } + + private navigateHistory(direction: 1 | -1): void { + this.lastAction = null; + if (this.history.length === 0) return; + + const newIndex = this.historyIndex - direction; // Up(-1) increases index, Down(1) decreases + if (newIndex < -1 || newIndex >= this.history.length) return; + + // Capture state when first entering history browsing mode + if (this.historyIndex === -1 && newIndex >= 0) { + this.pushUndoSnapshot(); + this.historyDraft = structuredClone(this.state); + } + + this.historyIndex = newIndex; + + if (this.historyIndex === -1) { + const draft = this.historyDraft; + this.historyDraft = null; + if (draft) { + this.state = draft; + this.preferredVisualCol = null; + this.snappedFromCursorCol = null; + this.scrollOffset = 0; + if (this.onChange) this.onChange(this.getText()); + } else { + this.setTextInternal(""); + } + } else { + this.setTextInternal(this.history[this.historyIndex] || "", direction === -1 ? "start" : "end"); + } + } + + private exitHistoryBrowsing(): void { + this.historyIndex = -1; + this.historyDraft = null; + } + + /** Internal setText that doesn't reset history state - used by navigateHistory */ + private setTextInternal(text: string, cursorPlacement: "start" | "end" = "end"): void { + const lines = text.split("\n"); + this.state.lines = lines.length === 0 ? [""] : lines; + this.state.cursorLine = cursorPlacement === "start" ? 0 : this.state.lines.length - 1; + this.setCursorCol(cursorPlacement === "start" ? 0 : this.state.lines[this.state.cursorLine]?.length || 0); + // Reset scroll - render() will adjust to show cursor + this.scrollOffset = 0; + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(width: number): string[] { + const maxPadding = Math.max(0, Math.floor((width - 1) / 2)); + const paddingX = Math.min(this.paddingX, maxPadding); + const contentWidth = Math.max(1, width - paddingX * 2); + + // Layout width: with padding the cursor can overflow into it, + // without padding we reserve 1 column for the cursor. + const layoutWidth = Math.max(1, contentWidth - (paddingX ? 0 : 1)); + + // Store for cursor navigation (must match wrapping width) + this.lastWidth = layoutWidth; + + const horizontal = this.borderColor("─"); + + // Layout the text + const layoutLines = this.layoutText(layoutWidth); + + // Calculate max visible lines: 30% of terminal height, minimum 5 lines + const terminalRows = this.tui.terminal.rows; + const maxVisibleLines = Math.max(5, Math.floor(terminalRows * 0.3)); + + // Find the cursor line index in layoutLines + let cursorLineIndex = layoutLines.findIndex((line) => line.hasCursor); + if (cursorLineIndex === -1) cursorLineIndex = 0; + + // Adjust scroll offset to keep cursor visible + if (cursorLineIndex < this.scrollOffset) { + this.scrollOffset = cursorLineIndex; + } else if (cursorLineIndex >= this.scrollOffset + maxVisibleLines) { + this.scrollOffset = cursorLineIndex - maxVisibleLines + 1; + } + + // Clamp scroll offset to valid range + const maxScrollOffset = Math.max(0, layoutLines.length - maxVisibleLines); + this.scrollOffset = Math.max(0, Math.min(this.scrollOffset, maxScrollOffset)); + + // Get visible lines slice + const visibleLines = layoutLines.slice(this.scrollOffset, this.scrollOffset + maxVisibleLines); + + const result: string[] = []; + const leftPadding = " ".repeat(paddingX); + const rightPadding = leftPadding; + + // Render top border (with scroll indicator if scrolled down) + if (this.scrollOffset > 0) { + const indicator = `─── ↑ ${this.scrollOffset} more `; + const remaining = width - visibleWidth(indicator); + if (remaining >= 0) { + result.push(this.borderColor(indicator + "─".repeat(remaining))); + } else { + result.push(this.borderColor(truncateToWidth(indicator, width))); + } + } else { + result.push(horizontal.repeat(width)); + } + + // Render each visible layout line + // Emit hardware cursor marker when focused so TUI can position the + // hardware cursor for IME candidate-window placement even while + // autocomplete (e.g. slash-command menu) is visible. + const emitCursorMarker = this.focused; + + for (const layoutLine of visibleLines) { + let displayText = layoutLine.text; + let lineVisibleWidth = visibleWidth(layoutLine.text); + let cursorInPadding = false; + + // Add cursor if this line has it + if (layoutLine.hasCursor && layoutLine.cursorPos !== undefined) { + const before = displayText.slice(0, layoutLine.cursorPos); + const after = displayText.slice(layoutLine.cursorPos); + + // Hardware cursor marker (zero-width, emitted before fake cursor for IME positioning) + const marker = emitCursorMarker ? CURSOR_MARKER : ""; + + if (after.length > 0) { + // Cursor is on a character (grapheme) - replace it with highlighted version + // Get the first grapheme from 'after' + const afterGraphemes = [...this.segment(after, "grapheme")]; + const firstGrapheme = afterGraphemes[0]?.segment || ""; + const restAfter = after.slice(firstGrapheme.length); + const cursor = `\x1b[7m${firstGrapheme}\x1b[0m`; + displayText = before + marker + cursor + restAfter; + // lineVisibleWidth stays the same - we're replacing, not adding + } else { + // Cursor is at the end - add highlighted space + const cursor = "\x1b[7m \x1b[0m"; + displayText = before + marker + cursor; + lineVisibleWidth = lineVisibleWidth + 1; + // If cursor overflows content width into the padding, flag it + if (lineVisibleWidth > contentWidth && paddingX > 0) { + cursorInPadding = true; + } + } + } + + // Calculate padding based on actual visible width + const padding = " ".repeat(Math.max(0, contentWidth - lineVisibleWidth)); + const lineRightPadding = cursorInPadding ? rightPadding.slice(1) : rightPadding; + + // Render the line (no side borders, just horizontal lines above and below) + result.push(`${leftPadding}${displayText}${padding}${lineRightPadding}`); + } + + // Render bottom border (with scroll indicator if more content below) + const linesBelow = layoutLines.length - (this.scrollOffset + visibleLines.length); + if (linesBelow > 0) { + const indicator = `─── ↓ ${linesBelow} more `; + const remaining = width - visibleWidth(indicator); + result.push(this.borderColor(indicator + "─".repeat(Math.max(0, remaining)))); + } else { + result.push(horizontal.repeat(width)); + } + + // Add autocomplete list if active + if (this.autocompleteState && this.autocompleteList) { + const autocompleteResult = this.autocompleteList.render(contentWidth); + for (const line of autocompleteResult) { + const lineWidth = visibleWidth(line); + const linePadding = " ".repeat(Math.max(0, contentWidth - lineWidth)); + result.push(`${leftPadding}${line}${linePadding}${rightPadding}`); + } + } + + return result; + } + + handleInput(data: string): void { + const kb = getKeybindings(); + + // Handle character jump mode (awaiting next character to jump to) + if (this.jumpMode !== null) { + // Cancel if the hotkey is pressed again + if (kb.matches(data, "tui.editor.jumpForward") || kb.matches(data, "tui.editor.jumpBackward")) { + this.jumpMode = null; + return; + } + + const printable = decodePrintableKey(data) ?? (data.charCodeAt(0) >= 32 ? data : undefined); + if (printable !== undefined) { + // Printable character - perform the jump + const direction = this.jumpMode; + this.jumpMode = null; + this.jumpToChar(printable, direction); + return; + } + + // Control character - cancel and fall through to normal handling + this.jumpMode = null; + } + + // Handle bracketed paste mode + if (data.includes("\x1b[200~")) { + this.isInPaste = true; + this.pasteBuffer = ""; + data = data.replace("\x1b[200~", ""); + } + + if (this.isInPaste) { + this.pasteBuffer += data; + const endIndex = this.pasteBuffer.indexOf("\x1b[201~"); + if (endIndex !== -1) { + const pasteContent = this.pasteBuffer.substring(0, endIndex); + if (pasteContent.length > 0) { + this.handlePaste(pasteContent); + } + this.isInPaste = false; + const remaining = this.pasteBuffer.substring(endIndex + 6); + this.pasteBuffer = ""; + if (remaining.length > 0) { + this.handleInput(remaining); + } + return; + } + return; + } + + // Ctrl+C - let parent handle (exit/clear) + if (kb.matches(data, "tui.input.copy")) { + return; + } + + // Undo + if (kb.matches(data, "tui.editor.undo")) { + this.undo(); + return; + } + + // Handle autocomplete mode + if (this.autocompleteState && this.autocompleteList) { + if (kb.matches(data, "tui.select.cancel")) { + this.cancelAutocomplete(); + return; + } + + if (kb.matches(data, "tui.select.up") || kb.matches(data, "tui.select.down")) { + this.autocompleteList.handleInput(data); + return; + } + + if (kb.matches(data, "tui.input.tab")) { + const selected = this.autocompleteList.getSelectedItem(); + if (selected && this.autocompleteProvider) { + this.pushUndoSnapshot(); + this.lastAction = null; + const result = this.autocompleteProvider.applyCompletion( + this.state.lines, + this.state.cursorLine, + this.state.cursorCol, + selected, + this.autocompletePrefix, + ); + this.state.lines = result.lines; + this.state.cursorLine = result.cursorLine; + this.setCursorCol(result.cursorCol); + this.cancelAutocomplete(); + if (this.onChange) this.onChange(this.getText()); + } + return; + } + + if (kb.matches(data, "tui.select.confirm")) { + const selected = this.autocompleteList.getSelectedItem(); + if (selected && this.autocompleteProvider) { + this.pushUndoSnapshot(); + this.lastAction = null; + const result = this.autocompleteProvider.applyCompletion( + this.state.lines, + this.state.cursorLine, + this.state.cursorCol, + selected, + this.autocompletePrefix, + ); + this.state.lines = result.lines; + this.state.cursorLine = result.cursorLine; + this.setCursorCol(result.cursorCol); + + if (this.autocompletePrefix.startsWith("/")) { + this.cancelAutocomplete(); + // Fall through to submit + } else { + this.cancelAutocomplete(); + if (this.onChange) this.onChange(this.getText()); + return; + } + } + } + } + + // Tab - trigger completion + if (kb.matches(data, "tui.input.tab") && !this.autocompleteState) { + this.handleTabCompletion(); + return; + } + + // Deletion actions + if (kb.matches(data, "tui.editor.deleteToLineEnd")) { + this.deleteToEndOfLine(); + return; + } + if (kb.matches(data, "tui.editor.deleteToLineStart")) { + this.deleteToStartOfLine(); + return; + } + if (kb.matches(data, "tui.editor.deleteWordBackward")) { + this.deleteWordBackwards(); + return; + } + if (kb.matches(data, "tui.editor.deleteWordForward")) { + this.deleteWordForward(); + return; + } + if (kb.matches(data, "tui.editor.deleteCharBackward") || matchesKey(data, "shift+backspace")) { + this.handleBackspace(); + return; + } + if (kb.matches(data, "tui.editor.deleteCharForward") || matchesKey(data, "shift+delete")) { + this.handleForwardDelete(); + return; + } + + // Kill ring actions + if (kb.matches(data, "tui.editor.yank")) { + this.yank(); + return; + } + if (kb.matches(data, "tui.editor.yankPop")) { + this.yankPop(); + return; + } + + // Cursor movement actions + if (kb.matches(data, "tui.editor.cursorLineStart")) { + this.moveToLineStart(); + return; + } + if (kb.matches(data, "tui.editor.cursorLineEnd")) { + this.moveToLineEnd(); + return; + } + if (kb.matches(data, "tui.editor.cursorWordLeft")) { + this.moveWordBackwards(); + return; + } + if (kb.matches(data, "tui.editor.cursorWordRight")) { + this.moveWordForwards(); + return; + } + + // New line + if ( + kb.matches(data, "tui.input.newLine") || + (data.charCodeAt(0) === 10 && data.length > 1) || + data === "\x1b\r" || + data === "\x1b[13;2~" || + (data.length > 1 && data.includes("\x1b") && data.includes("\r")) || + (data === "\n" && data.length === 1) + ) { + if (this.shouldSubmitOnBackslashEnter(data, kb)) { + this.handleBackspace(); + this.submitValue(); + return; + } + this.addNewLine(); + return; + } + + // Submit (Enter) + if (kb.matches(data, "tui.input.submit")) { + if (this.disableSubmit) return; + + // Workaround for terminals without Shift+Enter support: + // If char before cursor is \, delete it and insert newline instead of submitting. + const currentLine = this.state.lines[this.state.cursorLine] || ""; + if (this.state.cursorCol > 0 && currentLine[this.state.cursorCol - 1] === "\\") { + this.handleBackspace(); + this.addNewLine(); + return; + } + + this.submitValue(); + return; + } + + // Arrow key navigation (with history support) + if (kb.matches(data, "tui.editor.cursorUp")) { + if ( + this.isOnFirstVisualLine() && + (this.isEditorEmpty() || this.historyIndex > -1 || this.state.cursorCol === 0) + ) { + this.navigateHistory(-1); + } else if (this.isOnFirstVisualLine()) { + // Already at top - jump to start of line + this.moveToLineStart(); + } else { + this.moveCursor(-1, 0); + } + return; + } + if (kb.matches(data, "tui.editor.cursorDown")) { + if (this.historyIndex > -1 && this.isOnLastVisualLine()) { + this.navigateHistory(1); + } else if (this.isOnLastVisualLine()) { + // Already at bottom - jump to end of line + this.moveToLineEnd(); + } else { + this.moveCursor(1, 0); + } + return; + } + if (kb.matches(data, "tui.editor.cursorRight")) { + this.moveCursor(0, 1); + return; + } + if (kb.matches(data, "tui.editor.cursorLeft")) { + this.moveCursor(0, -1); + return; + } + + // Page up/down - scroll by page and move cursor + if (kb.matches(data, "tui.editor.pageUp")) { + this.pageScroll(-1); + return; + } + if (kb.matches(data, "tui.editor.pageDown")) { + this.pageScroll(1); + return; + } + + // Character jump mode triggers + if (kb.matches(data, "tui.editor.jumpForward")) { + this.jumpMode = "forward"; + return; + } + if (kb.matches(data, "tui.editor.jumpBackward")) { + this.jumpMode = "backward"; + return; + } + + // Shift+Space - insert regular space + if (matchesKey(data, "shift+space")) { + this.insertCharacter(" "); + return; + } + + const printable = decodePrintableKey(data); + if (printable !== undefined) { + this.insertCharacter(printable); + return; + } + + // Regular characters + if (data.charCodeAt(0) >= 32) { + this.insertCharacter(data); + } + } + + private layoutText(contentWidth: number): LayoutLine[] { + const layoutLines: LayoutLine[] = []; + + if (this.state.lines.length === 0 || (this.state.lines.length === 1 && this.state.lines[0] === "")) { + // Empty editor + layoutLines.push({ + text: "", + hasCursor: true, + cursorPos: 0, + }); + return layoutLines; + } + + // Process each logical line + for (let i = 0; i < this.state.lines.length; i++) { + const line = this.state.lines[i] || ""; + const isCurrentLine = i === this.state.cursorLine; + const lineVisibleWidth = visibleWidth(line); + + if (lineVisibleWidth <= contentWidth) { + // Line fits in one layout line + if (isCurrentLine) { + layoutLines.push({ + text: line, + hasCursor: true, + cursorPos: this.state.cursorCol, + }); + } else { + layoutLines.push({ + text: line, + hasCursor: false, + }); + } + } else { + // Line needs wrapping - use word-aware wrapping + const chunks = wordWrapLine(line, contentWidth, [...this.segment(line, "grapheme")]); + + for (let chunkIndex = 0; chunkIndex < chunks.length; chunkIndex++) { + const chunk = chunks[chunkIndex]; + if (!chunk) continue; + + const cursorPos = this.state.cursorCol; + const isLastChunk = chunkIndex === chunks.length - 1; + + // Determine if cursor is in this chunk + // For word-wrapped chunks, we need to handle the case where + // cursor might be in trimmed whitespace at end of chunk + let hasCursorInChunk = false; + let adjustedCursorPos = 0; + + if (isCurrentLine) { + if (isLastChunk) { + // Last chunk: cursor belongs here if >= startIndex + hasCursorInChunk = cursorPos >= chunk.startIndex; + adjustedCursorPos = cursorPos - chunk.startIndex; + } else { + // Non-last chunk: cursor belongs here if in range [startIndex, endIndex) + // But we need to handle the visual position in the trimmed text + hasCursorInChunk = cursorPos >= chunk.startIndex && cursorPos < chunk.endIndex; + if (hasCursorInChunk) { + adjustedCursorPos = cursorPos - chunk.startIndex; + // Clamp to text length (in case cursor was in trimmed whitespace) + if (adjustedCursorPos > chunk.text.length) { + adjustedCursorPos = chunk.text.length; + } + } + } + } + + if (hasCursorInChunk) { + layoutLines.push({ + text: chunk.text, + hasCursor: true, + cursorPos: adjustedCursorPos, + }); + } else { + layoutLines.push({ + text: chunk.text, + hasCursor: false, + }); + } + } + } + } + + return layoutLines; + } + + getText(): string { + return this.state.lines.join("\n"); + } + + private expandPasteMarkers(text: string): string { + let result = text; + for (const [pasteId, pasteContent] of this.pastes) { + const markerRegex = new RegExp(`\\[paste #${pasteId}( (\\+\\d+ lines|\\d+ chars))?\\]`, "g"); + result = result.replace(markerRegex, () => pasteContent); + } + return result; + } + + /** + * Get text with paste markers expanded to their actual content. + * Use this when you need the full content (e.g., for external editor). + */ + getExpandedText(): string { + return this.expandPasteMarkers(this.state.lines.join("\n")); + } + + getLines(): string[] { + return [...this.state.lines]; + } + + getCursor(): { line: number; col: number } { + return { line: this.state.cursorLine, col: this.state.cursorCol }; + } + + setText(text: string): void { + this.cancelAutocomplete(); + this.lastAction = null; + this.exitHistoryBrowsing(); + const normalized = this.normalizeText(text); + // Push undo snapshot if content differs (makes programmatic changes undoable) + if (this.getText() !== normalized) { + this.pushUndoSnapshot(); + } + this.setTextInternal(normalized); + } + + /** + * Insert text at the current cursor position. + * Used for programmatic insertion (e.g., clipboard image markers). + * This is atomic for undo - single undo restores entire pre-insert state. + */ + insertTextAtCursor(text: string): void { + if (!text) return; + this.cancelAutocomplete(); + this.pushUndoSnapshot(); + this.lastAction = null; + this.exitHistoryBrowsing(); + this.insertTextAtCursorInternal(text); + } + + /** + * Normalize text for editor storage: + * - Normalize line endings (\r\n and \r -> \n) + * - Expand tabs to 4 spaces + */ + private normalizeText(text: string): string { + return text.replace(/\r\n/g, "\n").replace(/\r/g, "\n").replace(/\t/g, " "); + } + + /** + * Internal text insertion at cursor. Handles single and multi-line text. + * Does not push undo snapshots or trigger autocomplete - caller is responsible. + * Normalizes line endings and calls onChange once at the end. + */ + private insertTextAtCursorInternal(text: string): void { + if (!text) return; + + // Normalize line endings and tabs + const normalized = this.normalizeText(text); + const insertedLines = normalized.split("\n"); + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const beforeCursor = currentLine.slice(0, this.state.cursorCol); + const afterCursor = currentLine.slice(this.state.cursorCol); + + if (insertedLines.length === 1) { + // Single line - insert at cursor position + this.state.lines[this.state.cursorLine] = beforeCursor + normalized + afterCursor; + this.setCursorCol(this.state.cursorCol + normalized.length); + } else { + // Multi-line insertion + this.state.lines = [ + // All lines before current line + ...this.state.lines.slice(0, this.state.cursorLine), + + // The first inserted line merged with text before cursor + beforeCursor + insertedLines[0], + + // All middle inserted lines + ...insertedLines.slice(1, -1), + + // The last inserted line with text after cursor + insertedLines[insertedLines.length - 1] + afterCursor, + + // All lines after current line + ...this.state.lines.slice(this.state.cursorLine + 1), + ]; + + this.state.cursorLine += insertedLines.length - 1; + this.setCursorCol((insertedLines[insertedLines.length - 1] || "").length); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + // All the editor methods from before... + private insertCharacter(char: string, skipUndoCoalescing?: boolean): void { + this.exitHistoryBrowsing(); + + // Undo coalescing (fish-style): + // - Consecutive word chars coalesce into one undo unit + // - Space captures state before itself (so undo removes space+following word together) + // - Each space is separately undoable + // Skip coalescing when called from atomic operations (e.g., handlePaste) + if (!skipUndoCoalescing) { + if (isWhitespaceChar(char) || this.lastAction !== "type-word") { + this.pushUndoSnapshot(); + } + this.lastAction = "type-word"; + } + + const line = this.state.lines[this.state.cursorLine] || ""; + + const before = line.slice(0, this.state.cursorCol); + const after = line.slice(this.state.cursorCol); + + this.state.lines[this.state.cursorLine] = before + char + after; + this.setCursorCol(this.state.cursorCol + char.length); + + if (this.onChange) { + this.onChange(this.getText()); + } + + // Check if we should trigger or update autocomplete + if (!this.autocompleteState) { + // Auto-trigger for "/" at the start of a line (slash commands) + if (char === "/" && this.isAtStartOfMessage()) { + this.tryTriggerAutocomplete(); + } + // Auto-trigger for symbol-based completion like @, #, or provider triggers at token boundaries + else if (this.autocompleteTriggerCharacters.includes(char)) { + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, this.state.cursorCol); + const charBeforeSymbol = textBeforeCursor[textBeforeCursor.length - 2]; + if (textBeforeCursor.length === 1 || charBeforeSymbol === " " || charBeforeSymbol === "\t") { + this.tryTriggerAutocomplete(); + } + } + // Also auto-trigger when typing letters in a slash command or symbol completion context + else if (/[a-zA-Z0-9.\-_]/.test(char)) { + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, this.state.cursorCol); + // Check if we're in a slash command (with or without space for arguments) + if (this.isInSlashCommandContext(textBeforeCursor)) { + this.tryTriggerAutocomplete(); + } + // Check if we're in a symbol-based completion context like @, #, or provider triggers + else if (this.autocompleteTriggerPattern.test(textBeforeCursor)) { + this.tryTriggerAutocomplete(); + } + } + } else { + this.updateAutocomplete(); + } + } + + private handlePaste(pastedText: string): void { + this.cancelAutocomplete(); + this.exitHistoryBrowsing(); + this.lastAction = null; + + this.pushUndoSnapshot(); + + // Some terminals (e.g. tmux popups with extended-keys-format=csi-u) re-encode + // control bytes inside bracketed paste as CSI-u Ctrl+ sequences + // (ESC [ ; 5 u). Decode those back to their literal byte so the + // per-char filter below preserves newlines instead of stripping ESC and + // leaking the printable tail (e.g. "[106;5u") into the editor. + const decodedText = pastedText.replace(/\x1b\[(\d+);5u/g, (match, code) => { + const cp = Number(code); + if (cp >= 97 && cp <= 122) return String.fromCharCode(cp - 96); + if (cp >= 65 && cp <= 90) return String.fromCharCode(cp - 64); + return match; + }); + + // Clean the pasted text: normalize line endings, expand tabs + const cleanText = this.normalizeText(decodedText); + + // Filter out non-printable characters except newlines + let filteredText = cleanText + .split("") + .filter((char) => char === "\n" || char.charCodeAt(0) >= 32) + .join(""); + + // If pasting a file path (starts with /, ~, or .) and the character before + // the cursor is a word character, prepend a space for better readability + if (/^[/~.]/.test(filteredText)) { + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const charBeforeCursor = this.state.cursorCol > 0 ? currentLine[this.state.cursorCol - 1] : ""; + if (charBeforeCursor && /\w/.test(charBeforeCursor)) { + filteredText = ` ${filteredText}`; + } + } + + // Split into lines to check for large paste + const pastedLines = filteredText.split("\n"); + + // Check if this is a large paste (> 10 lines or > 1000 characters) + const totalChars = filteredText.length; + if (pastedLines.length > 10 || totalChars > 1000) { + // Store the paste and insert a marker + this.pasteCounter++; + const pasteId = this.pasteCounter; + this.pastes.set(pasteId, filteredText); + + // Insert marker like "[paste #1 +123 lines]" or "[paste #1 1234 chars]" + const marker = + pastedLines.length > 10 + ? `[paste #${pasteId} +${pastedLines.length} lines]` + : `[paste #${pasteId} ${totalChars} chars]`; + this.insertTextAtCursorInternal(marker); + return; + } + + if (pastedLines.length === 1) { + // Single line - insert atomically (do not trigger autocomplete during paste) + this.insertTextAtCursorInternal(filteredText); + return; + } + + // Multi-line paste - use direct state manipulation + this.insertTextAtCursorInternal(filteredText); + } + + private addNewLine(): void { + this.cancelAutocomplete(); + this.exitHistoryBrowsing(); + this.lastAction = null; + + this.pushUndoSnapshot(); + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + const before = currentLine.slice(0, this.state.cursorCol); + const after = currentLine.slice(this.state.cursorCol); + + // Split current line + this.state.lines[this.state.cursorLine] = before; + this.state.lines.splice(this.state.cursorLine + 1, 0, after); + + // Move cursor to start of new line + this.state.cursorLine++; + this.setCursorCol(0); + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + private shouldSubmitOnBackslashEnter(data: string, kb: ReturnType): boolean { + if (this.disableSubmit) return false; + if (!matchesKey(data, "enter")) return false; + const submitKeys = kb.getKeys("tui.input.submit"); + const hasShiftEnter = submitKeys.includes("shift+enter") || submitKeys.includes("shift+return"); + if (!hasShiftEnter) return false; + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + return this.state.cursorCol > 0 && currentLine[this.state.cursorCol - 1] === "\\"; + } + + private submitValue(): void { + this.cancelAutocomplete(); + const result = this.expandPasteMarkers(this.state.lines.join("\n")).trim(); + + this.state = { lines: [""], cursorLine: 0, cursorCol: 0 }; + this.pastes.clear(); + this.pasteCounter = 0; + this.exitHistoryBrowsing(); + this.scrollOffset = 0; + this.undoStack.clear(); + this.lastAction = null; + + if (this.onChange) this.onChange(""); + if (this.onSubmit) this.onSubmit(result); + } + + private handleBackspace(): void { + this.exitHistoryBrowsing(); + this.lastAction = null; + + if (this.state.cursorCol > 0) { + this.pushUndoSnapshot(); + + // Delete grapheme before cursor (handles emojis, combining characters, etc.) + const line = this.state.lines[this.state.cursorLine] || ""; + const beforeCursor = line.slice(0, this.state.cursorCol); + + // Find the last grapheme in the text before cursor + const graphemes = [...this.segment(beforeCursor, "grapheme")]; + const lastGrapheme = graphemes[graphemes.length - 1]; + const graphemeLength = lastGrapheme ? lastGrapheme.segment.length : 1; + + const before = line.slice(0, this.state.cursorCol - graphemeLength); + const after = line.slice(this.state.cursorCol); + + this.state.lines[this.state.cursorLine] = before + after; + this.setCursorCol(this.state.cursorCol - graphemeLength); + } else if (this.state.cursorLine > 0) { + this.pushUndoSnapshot(); + + // Merge with previous line + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const previousLine = this.state.lines[this.state.cursorLine - 1] || ""; + + this.state.lines[this.state.cursorLine - 1] = previousLine + currentLine; + this.state.lines.splice(this.state.cursorLine, 1); + + this.state.cursorLine--; + this.setCursorCol(previousLine.length); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + + // Update or re-trigger autocomplete after backspace + if (this.autocompleteState) { + this.updateAutocomplete(); + } else { + // If autocomplete was cancelled (no matches), re-trigger if we're in a completable context + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, this.state.cursorCol); + // Slash command context + if (this.isInSlashCommandContext(textBeforeCursor)) { + this.tryTriggerAutocomplete(); + } + // Symbol-based completion context like @, #, or provider triggers + else if (this.autocompleteTriggerPattern.test(textBeforeCursor)) { + this.tryTriggerAutocomplete(); + } + } + } + + /** + * Set cursor column and clear preferredVisualCol. + * Use this for all non-vertical cursor movements to reset sticky column behavior. + */ + private setCursorCol(col: number): void { + this.state.cursorCol = col; + this.preferredVisualCol = null; + this.snappedFromCursorCol = null; + } + + /** + * Move cursor to a target visual line, applying sticky column logic. + * Shared by moveCursor() and pageScroll(). + */ + private moveToVisualLine( + visualLines: Array<{ logicalLine: number; startCol: number; length: number }>, + currentVisualLine: number, + targetVisualLine: number, + ): void { + const currentVL = visualLines[currentVisualLine]; + const targetVL = visualLines[targetVisualLine]; + if (!(currentVL && targetVL)) return; + + // When the cursor was snapped to a segment start, resolve the pre-snap + // position against the VL it belongs to. This gives the correct visual + // column even after a resize reshuffles VLs. + let currentVisualCol: number; + if (this.snappedFromCursorCol !== null) { + const vlIndex = this.findVisualLineAt(visualLines, currentVL.logicalLine, this.snappedFromCursorCol); + currentVisualCol = this.snappedFromCursorCol - visualLines[vlIndex].startCol; + } else { + currentVisualCol = this.state.cursorCol - currentVL.startCol; + } + + // For non-last segments, clamp to length-1 to stay within the segment + const isLastSourceSegment = + currentVisualLine === visualLines.length - 1 || + visualLines[currentVisualLine + 1]?.logicalLine !== currentVL.logicalLine; + const sourceMaxVisualCol = isLastSourceSegment ? currentVL.length : Math.max(0, currentVL.length - 1); + + const isLastTargetSegment = + targetVisualLine === visualLines.length - 1 || + visualLines[targetVisualLine + 1]?.logicalLine !== targetVL.logicalLine; + const targetMaxVisualCol = isLastTargetSegment ? targetVL.length : Math.max(0, targetVL.length - 1); + + const moveToVisualCol = this.computeVerticalMoveColumn(currentVisualCol, sourceMaxVisualCol, targetMaxVisualCol); + + // Set cursor position + this.state.cursorLine = targetVL.logicalLine; + const targetCol = targetVL.startCol + moveToVisualCol; + const logicalLine = this.state.lines[targetVL.logicalLine] || ""; + this.state.cursorCol = Math.min(targetCol, logicalLine.length); + + // Snap cursor to atomic segment boundary (e.g. paste markers) + // so the cursor never lands in the middle of a multi-grapheme unit. + // Single-grapheme segments don't need snapping. + const segments = [...this.segment(logicalLine, "grapheme")]; + for (const seg of segments) { + if (seg.index > this.state.cursorCol) break; + if (seg.segment.length <= 1) continue; + if (this.state.cursorCol < seg.index + seg.segment.length) { + const isContinuation = seg.index < targetVL.startCol; + const isMovingDown = targetVisualLine > currentVisualLine; + + if (isContinuation && isMovingDown) { + // The segment started on a previous visual line, and we + // already visited it on the way down. Skip all remaining + // continuation VLs and land on the first VL past it. + const segEnd = seg.index + seg.segment.length; + let next = targetVisualLine + 1; + while ( + next < visualLines.length && + visualLines[next].logicalLine === targetVL.logicalLine && + visualLines[next].startCol < segEnd + ) { + next++; + } + if (next < visualLines.length) { + this.moveToVisualLine(visualLines, currentVisualLine, next); + return; + } + } + + // Snap to the start of the segment so it gets highlighted. + // Store the pre-snap position so the next vertical move can + // resolve it to the correct visual column. + this.snappedFromCursorCol = this.state.cursorCol; + this.state.cursorCol = seg.index; + return; + } + } + + // No snap occurred – we moved out of the atomic segment. + this.snappedFromCursorCol = null; + } + + /** + * Compute the target visual column for vertical cursor movement. + * Implements the sticky column decision table: + * + * | P | S | T | U | Scenario | Set Preferred | Move To | + * |---|---|---|---| ---------------------------------------------------- |---------------|-------------| + * | 0 | * | 0 | - | Start nav, target fits | null | current | + * | 0 | * | 1 | - | Start nav, target shorter | current | target end | + * | 1 | 0 | 0 | 0 | Clamped, target fits preferred | null | preferred | + * | 1 | 0 | 0 | 1 | Clamped, target longer but still can't fit preferred | keep | target end | + * | 1 | 0 | 1 | - | Clamped, target even shorter | keep | target end | + * | 1 | 1 | 0 | - | Rewrapped, target fits current | null | current | + * | 1 | 1 | 1 | - | Rewrapped, target shorter than current | current | target end | + * + * Where: + * - P = preferred col is set + * - S = cursor in middle of source line (not clamped to end) + * - T = target line shorter than current visual col + * - U = target line shorter than preferred col + */ + private computeVerticalMoveColumn( + currentVisualCol: number, + sourceMaxVisualCol: number, + targetMaxVisualCol: number, + ): number { + const hasPreferred = this.preferredVisualCol !== null; // P + const cursorInMiddle = currentVisualCol < sourceMaxVisualCol; // S + const targetTooShort = targetMaxVisualCol < currentVisualCol; // T + + if (!hasPreferred || cursorInMiddle) { + if (targetTooShort) { + // Cases 2 and 7 + this.preferredVisualCol = currentVisualCol; + return targetMaxVisualCol; + } + + // Cases 1 and 6 + this.preferredVisualCol = null; + return currentVisualCol; + } + + const targetCantFitPreferred = targetMaxVisualCol < this.preferredVisualCol!; // U + if (targetTooShort || targetCantFitPreferred) { + // Cases 4 and 5 + return targetMaxVisualCol; + } + + // Case 3 + const result = this.preferredVisualCol!; + this.preferredVisualCol = null; + return result; + } + + private moveToLineStart(): void { + this.lastAction = null; + this.setCursorCol(0); + } + + private moveToLineEnd(): void { + this.lastAction = null; + const currentLine = this.state.lines[this.state.cursorLine] || ""; + this.setCursorCol(currentLine.length); + } + + private deleteToStartOfLine(): void { + this.exitHistoryBrowsing(); + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + if (this.state.cursorCol > 0) { + this.pushUndoSnapshot(); + + // Calculate text to be deleted and save to kill ring (backward deletion = prepend) + const deletedText = currentLine.slice(0, this.state.cursorCol); + this.killRing.push(deletedText, { prepend: true, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + + // Delete from start of line up to cursor + this.state.lines[this.state.cursorLine] = currentLine.slice(this.state.cursorCol); + this.setCursorCol(0); + } else if (this.state.cursorLine > 0) { + this.pushUndoSnapshot(); + + // At start of line - merge with previous line, treating newline as deleted text + this.killRing.push("\n", { prepend: true, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + + const previousLine = this.state.lines[this.state.cursorLine - 1] || ""; + this.state.lines[this.state.cursorLine - 1] = previousLine + currentLine; + this.state.lines.splice(this.state.cursorLine, 1); + this.state.cursorLine--; + this.setCursorCol(previousLine.length); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + private deleteToEndOfLine(): void { + this.exitHistoryBrowsing(); + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + if (this.state.cursorCol < currentLine.length) { + this.pushUndoSnapshot(); + + // Calculate text to be deleted and save to kill ring (forward deletion = append) + const deletedText = currentLine.slice(this.state.cursorCol); + this.killRing.push(deletedText, { prepend: false, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + + // Delete from cursor to end of line + this.state.lines[this.state.cursorLine] = currentLine.slice(0, this.state.cursorCol); + } else if (this.state.cursorLine < this.state.lines.length - 1) { + this.pushUndoSnapshot(); + + // At end of line - merge with next line, treating newline as deleted text + this.killRing.push("\n", { prepend: false, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + + const nextLine = this.state.lines[this.state.cursorLine + 1] || ""; + this.state.lines[this.state.cursorLine] = currentLine + nextLine; + this.state.lines.splice(this.state.cursorLine + 1, 1); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + private deleteWordBackwards(): void { + this.exitHistoryBrowsing(); + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + // If at start of line, behave like backspace at column 0 (merge with previous line) + if (this.state.cursorCol === 0) { + if (this.state.cursorLine > 0) { + this.pushUndoSnapshot(); + + // Treat newline as deleted text (backward deletion = prepend) + this.killRing.push("\n", { prepend: true, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + + const previousLine = this.state.lines[this.state.cursorLine - 1] || ""; + this.state.lines[this.state.cursorLine - 1] = previousLine + currentLine; + this.state.lines.splice(this.state.cursorLine, 1); + this.state.cursorLine--; + this.setCursorCol(previousLine.length); + } + } else { + this.pushUndoSnapshot(); + + // Save lastAction before cursor movement (moveWordBackwards resets it) + const wasKill = this.lastAction === "kill"; + + const oldCursorCol = this.state.cursorCol; + this.moveWordBackwards(); + const deleteFrom = this.state.cursorCol; + this.setCursorCol(oldCursorCol); + + const deletedText = currentLine.slice(deleteFrom, this.state.cursorCol); + this.killRing.push(deletedText, { prepend: true, accumulate: wasKill }); + this.lastAction = "kill"; + + this.state.lines[this.state.cursorLine] = + currentLine.slice(0, deleteFrom) + currentLine.slice(this.state.cursorCol); + this.setCursorCol(deleteFrom); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + private deleteWordForward(): void { + this.exitHistoryBrowsing(); + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + // If at end of line, merge with next line (delete the newline) + if (this.state.cursorCol >= currentLine.length) { + if (this.state.cursorLine < this.state.lines.length - 1) { + this.pushUndoSnapshot(); + + // Treat newline as deleted text (forward deletion = append) + this.killRing.push("\n", { prepend: false, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + + const nextLine = this.state.lines[this.state.cursorLine + 1] || ""; + this.state.lines[this.state.cursorLine] = currentLine + nextLine; + this.state.lines.splice(this.state.cursorLine + 1, 1); + } + } else { + this.pushUndoSnapshot(); + + // Save lastAction before cursor movement (moveWordForwards resets it) + const wasKill = this.lastAction === "kill"; + + const oldCursorCol = this.state.cursorCol; + this.moveWordForwards(); + const deleteTo = this.state.cursorCol; + this.setCursorCol(oldCursorCol); + + const deletedText = currentLine.slice(this.state.cursorCol, deleteTo); + this.killRing.push(deletedText, { prepend: false, accumulate: wasKill }); + this.lastAction = "kill"; + + this.state.lines[this.state.cursorLine] = + currentLine.slice(0, this.state.cursorCol) + currentLine.slice(deleteTo); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + private handleForwardDelete(): void { + this.exitHistoryBrowsing(); + this.lastAction = null; + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + if (this.state.cursorCol < currentLine.length) { + this.pushUndoSnapshot(); + + // Delete grapheme at cursor position (handles emojis, combining characters, etc.) + const afterCursor = currentLine.slice(this.state.cursorCol); + + // Find the first grapheme at cursor + const graphemes = [...this.segment(afterCursor, "grapheme")]; + const firstGrapheme = graphemes[0]; + const graphemeLength = firstGrapheme ? firstGrapheme.segment.length : 1; + + const before = currentLine.slice(0, this.state.cursorCol); + const after = currentLine.slice(this.state.cursorCol + graphemeLength); + this.state.lines[this.state.cursorLine] = before + after; + } else if (this.state.cursorLine < this.state.lines.length - 1) { + this.pushUndoSnapshot(); + + // At end of line - merge with next line + const nextLine = this.state.lines[this.state.cursorLine + 1] || ""; + this.state.lines[this.state.cursorLine] = currentLine + nextLine; + this.state.lines.splice(this.state.cursorLine + 1, 1); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + + // Update or re-trigger autocomplete after forward delete + if (this.autocompleteState) { + this.updateAutocomplete(); + } else { + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, this.state.cursorCol); + // Slash command context + if (this.isInSlashCommandContext(textBeforeCursor)) { + this.tryTriggerAutocomplete(); + } + // Symbol-based completion context like @, #, or provider triggers + else if (this.autocompleteTriggerPattern.test(textBeforeCursor)) { + this.tryTriggerAutocomplete(); + } + } + } + + /** + * Build a mapping from visual lines to logical positions. + * Returns an array where each element represents a visual line with: + * - logicalLine: index into this.state.lines + * - startCol: starting column in the logical line + * - length: length of this visual line segment + */ + private buildVisualLineMap(width: number): Array<{ logicalLine: number; startCol: number; length: number }> { + const visualLines: Array<{ logicalLine: number; startCol: number; length: number }> = []; + + for (let i = 0; i < this.state.lines.length; i++) { + const line = this.state.lines[i] || ""; + const lineVisWidth = visibleWidth(line); + if (line.length === 0) { + // Empty line still takes one visual line + visualLines.push({ logicalLine: i, startCol: 0, length: 0 }); + } else if (lineVisWidth <= width) { + visualLines.push({ logicalLine: i, startCol: 0, length: line.length }); + } else { + // Line needs wrapping - use word-aware wrapping + const chunks = wordWrapLine(line, width, [...this.segment(line, "grapheme")]); + for (const chunk of chunks) { + visualLines.push({ + logicalLine: i, + startCol: chunk.startIndex, + length: chunk.endIndex - chunk.startIndex, + }); + } + } + } + + return visualLines; + } + + /** + * Find the visual line index that contains the given logical position. + */ + private findVisualLineAt( + visualLines: Array<{ logicalLine: number; startCol: number; length: number }>, + line: number, + col: number, + ): number { + for (let i = 0; i < visualLines.length; i++) { + const vl = visualLines[i]; + if (!vl || vl.logicalLine !== line) continue; + const offset = col - vl.startCol; + // Cursor is in this segment if it's within range. For the last + // segment of a logical line, cursor can be at length (end position) + const isLastSegmentOfLine = i === visualLines.length - 1 || visualLines[i + 1]?.logicalLine !== vl.logicalLine; + if (offset >= 0 && (offset < vl.length || (isLastSegmentOfLine && offset === vl.length))) { + return i; + } + } + return visualLines.length - 1; + } + + /** + * Find the visual line index for the current cursor position. + */ + private findCurrentVisualLine( + visualLines: Array<{ logicalLine: number; startCol: number; length: number }>, + ): number { + return this.findVisualLineAt(visualLines, this.state.cursorLine, this.state.cursorCol); + } + + private moveCursor(deltaLine: number, deltaCol: number): void { + this.lastAction = null; + const visualLines = this.buildVisualLineMap(this.lastWidth); + const currentVisualLine = this.findCurrentVisualLine(visualLines); + + if (deltaLine !== 0) { + const targetVisualLine = currentVisualLine + deltaLine; + + if (targetVisualLine >= 0 && targetVisualLine < visualLines.length) { + this.moveToVisualLine(visualLines, currentVisualLine, targetVisualLine); + } + } + + if (deltaCol !== 0) { + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + if (deltaCol > 0) { + // Moving right - move by one grapheme (handles emojis, combining characters, etc.) + if (this.state.cursorCol < currentLine.length) { + const afterCursor = currentLine.slice(this.state.cursorCol); + const graphemes = [...this.segment(afterCursor, "grapheme")]; + const firstGrapheme = graphemes[0]; + this.setCursorCol(this.state.cursorCol + (firstGrapheme ? firstGrapheme.segment.length : 1)); + } else if (this.state.cursorLine < this.state.lines.length - 1) { + // Wrap to start of next logical line + this.state.cursorLine++; + this.setCursorCol(0); + } else { + // At end of last line - can't move, but set preferredVisualCol for up/down navigation + const currentVL = visualLines[currentVisualLine]; + if (currentVL) { + this.preferredVisualCol = this.state.cursorCol - currentVL.startCol; + } + } + } else { + // Moving left - move by one grapheme (handles emojis, combining characters, etc.) + if (this.state.cursorCol > 0) { + const beforeCursor = currentLine.slice(0, this.state.cursorCol); + const graphemes = [...this.segment(beforeCursor, "grapheme")]; + const lastGrapheme = graphemes[graphemes.length - 1]; + this.setCursorCol(this.state.cursorCol - (lastGrapheme ? lastGrapheme.segment.length : 1)); + } else if (this.state.cursorLine > 0) { + // Wrap to end of previous logical line + this.state.cursorLine--; + const prevLine = this.state.lines[this.state.cursorLine] || ""; + this.setCursorCol(prevLine.length); + } + } + } + + // Keep an open autocomplete picker in sync with the new cursor + // position: cursor movement changes the text before the cursor, so a + // picker computed for the old position is stale. Re-query so it + // refreshes — or closes when the new position yields no suggestions — + // mirroring insertCharacter()/handleBackspace(). Without this, arrowing + // left from `/cmd ` back into the command name leaves the argument + // picker showing against a `/cmd` prefix (and a Tab there would + // concatenate the stale suggestion onto the partial command name). + if (this.autocompleteState) { + this.updateAutocomplete(); + } + } + + /** + * Scroll by a page (direction: -1 for up, 1 for down). + * Moves cursor by the page size while keeping it in bounds. + */ + private pageScroll(direction: -1 | 1): void { + this.lastAction = null; + const terminalRows = this.tui.terminal.rows; + const pageSize = Math.max(5, Math.floor(terminalRows * 0.3)); + + const visualLines = this.buildVisualLineMap(this.lastWidth); + const currentVisualLine = this.findCurrentVisualLine(visualLines); + const targetVisualLine = Math.max(0, Math.min(visualLines.length - 1, currentVisualLine + direction * pageSize)); + + this.moveToVisualLine(visualLines, currentVisualLine, targetVisualLine); + } + + private moveWordBackwards(): void { + this.lastAction = null; + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + // If at start of line, move to end of previous line + if (this.state.cursorCol === 0) { + if (this.state.cursorLine > 0) { + this.state.cursorLine--; + const prevLine = this.state.lines[this.state.cursorLine] || ""; + this.setCursorCol(prevLine.length); + } + return; + } + + this.setCursorCol( + findWordBackward(currentLine, this.state.cursorCol, { + segment: (text) => this.segment(text, "word"), + isAtomicSegment: isPasteMarker, + }), + ); + } + + /** + * Yank (paste) the most recent kill ring entry at cursor position. + */ + private yank(): void { + if (this.killRing.length === 0) return; + + this.pushUndoSnapshot(); + + const text = this.killRing.peek()!; + this.insertYankedText(text); + + this.lastAction = "yank"; + } + + /** + * Cycle through kill ring (only works immediately after yank or yank-pop). + * Replaces the last yanked text with the previous entry in the ring. + */ + private yankPop(): void { + // Only works if we just yanked and have more than one entry + if (this.lastAction !== "yank" || this.killRing.length <= 1) return; + + this.pushUndoSnapshot(); + + // Delete the previously yanked text (still at end of ring before rotation) + this.deleteYankedText(); + + // Rotate the ring: move end to front + this.killRing.rotate(); + + // Insert the new most recent entry (now at end after rotation) + const text = this.killRing.peek()!; + this.insertYankedText(text); + + this.lastAction = "yank"; + } + + /** + * Insert text at cursor position (used by yank operations). + */ + private insertYankedText(text: string): void { + this.exitHistoryBrowsing(); + const lines = text.split("\n"); + + if (lines.length === 1) { + // Single line - insert at cursor + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const before = currentLine.slice(0, this.state.cursorCol); + const after = currentLine.slice(this.state.cursorCol); + this.state.lines[this.state.cursorLine] = before + text + after; + this.setCursorCol(this.state.cursorCol + text.length); + } else { + // Multi-line insert + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const before = currentLine.slice(0, this.state.cursorCol); + const after = currentLine.slice(this.state.cursorCol); + + // First line merges with text before cursor + this.state.lines[this.state.cursorLine] = before + (lines[0] || ""); + + // Insert middle lines + for (let i = 1; i < lines.length - 1; i++) { + this.state.lines.splice(this.state.cursorLine + i, 0, lines[i] || ""); + } + + // Last line merges with text after cursor + const lastLineIndex = this.state.cursorLine + lines.length - 1; + this.state.lines.splice(lastLineIndex, 0, (lines[lines.length - 1] || "") + after); + + // Update cursor position + this.state.cursorLine = lastLineIndex; + this.setCursorCol((lines[lines.length - 1] || "").length); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + /** + * Delete the previously yanked text (used by yank-pop). + * The yanked text is derived from killRing[end] since it hasn't been rotated yet. + */ + private deleteYankedText(): void { + const yankedText = this.killRing.peek(); + if (!yankedText) return; + + const yankLines = yankedText.split("\n"); + + if (yankLines.length === 1) { + // Single line - delete backward from cursor + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const deleteLen = yankedText.length; + const before = currentLine.slice(0, this.state.cursorCol - deleteLen); + const after = currentLine.slice(this.state.cursorCol); + this.state.lines[this.state.cursorLine] = before + after; + this.setCursorCol(this.state.cursorCol - deleteLen); + } else { + // Multi-line delete - cursor is at end of last yanked line + const startLine = this.state.cursorLine - (yankLines.length - 1); + const startCol = (this.state.lines[startLine] || "").length - (yankLines[0] || "").length; + + // Get text after cursor on current line + const afterCursor = (this.state.lines[this.state.cursorLine] || "").slice(this.state.cursorCol); + + // Get text before yank start position + const beforeYank = (this.state.lines[startLine] || "").slice(0, startCol); + + // Remove all lines from startLine to cursorLine and replace with merged line + this.state.lines.splice(startLine, yankLines.length, beforeYank + afterCursor); + + // Update cursor + this.state.cursorLine = startLine; + this.setCursorCol(startCol); + } + + if (this.onChange) { + this.onChange(this.getText()); + } + } + + private pushUndoSnapshot(): void { + this.undoStack.push(this.state); + } + + private undo(): void { + this.exitHistoryBrowsing(); + const snapshot = this.undoStack.pop(); + if (!snapshot) return; + Object.assign(this.state, snapshot); + this.lastAction = null; + this.preferredVisualCol = null; + if (this.onChange) { + this.onChange(this.getText()); + } + } + + /** + * Jump to the first occurrence of a character in the specified direction. + * Multi-line search. Case-sensitive. Skips the current cursor position. + */ + private jumpToChar(char: string, direction: "forward" | "backward"): void { + this.lastAction = null; + const isForward = direction === "forward"; + const lines = this.state.lines; + + const end = isForward ? lines.length : -1; + const step = isForward ? 1 : -1; + + for (let lineIdx = this.state.cursorLine; lineIdx !== end; lineIdx += step) { + const line = lines[lineIdx] || ""; + const isCurrentLine = lineIdx === this.state.cursorLine; + + // Current line: start after/before cursor; other lines: search full line + const searchFrom = isCurrentLine + ? isForward + ? this.state.cursorCol + 1 + : this.state.cursorCol - 1 + : undefined; + + const idx = isForward ? line.indexOf(char, searchFrom) : line.lastIndexOf(char, searchFrom); + + if (idx !== -1) { + this.state.cursorLine = lineIdx; + this.setCursorCol(idx); + return; + } + } + // No match found - cursor stays in place + } + + private moveWordForwards(): void { + this.lastAction = null; + const currentLine = this.state.lines[this.state.cursorLine] || ""; + + // If at end of line, move to start of next line + if (this.state.cursorCol >= currentLine.length) { + if (this.state.cursorLine < this.state.lines.length - 1) { + this.state.cursorLine++; + this.setCursorCol(0); + } + return; + } + + this.setCursorCol( + findWordForward(currentLine, this.state.cursorCol, { + segment: (text) => this.segment(text, "word"), + isAtomicSegment: isPasteMarker, + }), + ); + } + + // Slash menu only allowed on the first line of the editor + private isSlashMenuAllowed(): boolean { + return this.state.cursorLine === 0; + } + + // Helper method to check if cursor is at start of message (for slash command detection) + private isAtStartOfMessage(): boolean { + if (!this.isSlashMenuAllowed()) return false; + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const beforeCursor = currentLine.slice(0, this.state.cursorCol); + return beforeCursor.trim() === "" || beforeCursor.trim() === "/"; + } + + private isInSlashCommandContext(textBeforeCursor: string): boolean { + return this.isSlashMenuAllowed() && textBeforeCursor.trimStart().startsWith("/"); + } + + // Autocomplete methods + /** + * Find the best autocomplete item index for the given prefix. + * Returns -1 if no match is found. + * + * Match priority: + * 1. Exact match (prefix === item.value) -> always selected + * 2. Prefix match -> first item whose value starts with prefix + * 3. No match -> -1 (keep default highlight) + * + * Matching is case-sensitive and checks item.value only. + */ + private getBestAutocompleteMatchIndex(items: Array<{ value: string; label: string }>, prefix: string): number { + if (!prefix) return -1; + + let firstPrefixIndex = -1; + + for (let i = 0; i < items.length; i++) { + const value = items[i]!.value; + if (value === prefix) { + return i; // Exact match always wins + } + if (firstPrefixIndex === -1 && value.startsWith(prefix)) { + firstPrefixIndex = i; + } + } + + return firstPrefixIndex; + } + + private createAutocompleteList( + prefix: string, + items: Array<{ value: string; label: string; description?: string }>, + ): SelectList { + const layout = prefix.startsWith("/") ? SLASH_COMMAND_SELECT_LIST_LAYOUT : undefined; + return new SelectList(items, this.autocompleteMaxVisible, this.theme.selectList, layout); + } + + private tryTriggerAutocomplete(explicitTab: boolean = false): void { + this.requestAutocomplete({ force: false, explicitTab }); + } + + private handleTabCompletion(): void { + if (!this.autocompleteProvider) return; + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const beforeCursor = currentLine.slice(0, this.state.cursorCol); + + if (this.isInSlashCommandContext(beforeCursor) && !beforeCursor.trimStart().includes(" ")) { + this.handleSlashCommandCompletion(); + } else { + this.forceFileAutocomplete(true); + } + } + + private handleSlashCommandCompletion(): void { + this.requestAutocomplete({ force: false, explicitTab: true }); + } + + private forceFileAutocomplete(explicitTab: boolean = false): void { + this.requestAutocomplete({ force: true, explicitTab }); + } + + private requestAutocomplete(options: { force: boolean; explicitTab: boolean }): void { + if (!this.autocompleteProvider) return; + + if (options.force) { + const shouldTrigger = + !this.autocompleteProvider.shouldTriggerFileCompletion || + this.autocompleteProvider.shouldTriggerFileCompletion( + this.state.lines, + this.state.cursorLine, + this.state.cursorCol, + ); + if (!shouldTrigger) { + return; + } + } + + this.cancelAutocompleteRequest(); + const startToken = ++this.autocompleteStartToken; + + const debounceMs = this.getAutocompleteDebounceMs(options); + if (debounceMs > 0) { + this.autocompleteDebounceTimer = setTimeout(() => { + this.autocompleteDebounceTimer = undefined; + void this.startAutocompleteRequest(startToken, options); + }, debounceMs); + return; + } + + void this.startAutocompleteRequest(startToken, options); + } + + private async startAutocompleteRequest( + startToken: number, + options: { force: boolean; explicitTab: boolean }, + ): Promise { + const previousTask = this.autocompleteRequestTask; + this.autocompleteRequestTask = (async () => { + await previousTask; + if (startToken !== this.autocompleteStartToken || !this.autocompleteProvider) { + return; + } + + const controller = new AbortController(); + this.autocompleteAbort = controller; + const requestId = ++this.autocompleteRequestId; + const snapshotText = this.getText(); + const snapshotLine = this.state.cursorLine; + const snapshotCol = this.state.cursorCol; + + await this.runAutocompleteRequest(requestId, controller, snapshotText, snapshotLine, snapshotCol, options); + })(); + await this.autocompleteRequestTask; + } + + private setAutocompleteTriggerCharacters(triggerCharacters: string[]): void { + const next = [...DEFAULT_AUTOCOMPLETE_TRIGGER_CHARACTERS]; + for (const character of triggerCharacters) { + if (character.length !== 1 || character === "/" || isWhitespaceChar(character) || next.includes(character)) { + continue; + } + next.push(character); + } + this.autocompleteTriggerCharacters = next; + this.autocompleteTriggerPattern = buildTriggerPattern(next); + this.autocompleteDebouncePattern = buildDebouncePattern(next); + } + + private getAutocompleteDebounceMs(options: { force: boolean; explicitTab: boolean }): number { + if (options.explicitTab || options.force) { + return 0; + } + + const currentLine = this.state.lines[this.state.cursorLine] || ""; + const textBeforeCursor = currentLine.slice(0, this.state.cursorCol); + return this.autocompleteDebouncePattern.test(textBeforeCursor) ? ATTACHMENT_AUTOCOMPLETE_DEBOUNCE_MS : 0; + } + + private async runAutocompleteRequest( + requestId: number, + controller: AbortController, + snapshotText: string, + snapshotLine: number, + snapshotCol: number, + options: { force: boolean; explicitTab: boolean }, + ): Promise { + if (!this.autocompleteProvider) return; + + const suggestions = await this.autocompleteProvider.getSuggestions( + this.state.lines, + this.state.cursorLine, + this.state.cursorCol, + { signal: controller.signal, force: options.force }, + ); + + if (!this.isAutocompleteRequestCurrent(requestId, controller, snapshotText, snapshotLine, snapshotCol)) { + return; + } + + this.autocompleteAbort = undefined; + + if (!suggestions || !Array.isArray(suggestions.items) || suggestions.items.length === 0) { + this.cancelAutocomplete(); + this.tui.requestRender(); + return; + } + + if (options.force && options.explicitTab && suggestions.items.length === 1) { + const item = suggestions.items[0]!; + this.pushUndoSnapshot(); + this.lastAction = null; + const result = this.autocompleteProvider.applyCompletion( + this.state.lines, + this.state.cursorLine, + this.state.cursorCol, + item, + suggestions.prefix, + ); + this.state.lines = result.lines; + this.state.cursorLine = result.cursorLine; + this.setCursorCol(result.cursorCol); + if (this.onChange) this.onChange(this.getText()); + this.tui.requestRender(); + return; + } + + this.applyAutocompleteSuggestions(suggestions, options.force ? "force" : "regular"); + this.tui.requestRender(); + } + + private isAutocompleteRequestCurrent( + requestId: number, + controller: AbortController, + snapshotText: string, + snapshotLine: number, + snapshotCol: number, + ): boolean { + return ( + !controller.signal.aborted && + requestId === this.autocompleteRequestId && + this.getText() === snapshotText && + this.state.cursorLine === snapshotLine && + this.state.cursorCol === snapshotCol + ); + } + + private applyAutocompleteSuggestions(suggestions: AutocompleteSuggestions, state: "regular" | "force"): void { + this.autocompletePrefix = suggestions.prefix; + this.autocompleteList = this.createAutocompleteList(suggestions.prefix, suggestions.items); + + const bestMatchIndex = this.getBestAutocompleteMatchIndex(suggestions.items, suggestions.prefix); + if (bestMatchIndex >= 0) { + this.autocompleteList.setSelectedIndex(bestMatchIndex); + } + + this.autocompleteState = state; + } + + private cancelAutocompleteRequest(): void { + this.autocompleteStartToken += 1; + if (this.autocompleteDebounceTimer) { + clearTimeout(this.autocompleteDebounceTimer); + this.autocompleteDebounceTimer = undefined; + } + this.autocompleteAbort?.abort(); + this.autocompleteAbort = undefined; + } + + private clearAutocompleteUi(): void { + this.autocompleteState = null; + this.autocompleteList = undefined; + this.autocompletePrefix = ""; + } + + private cancelAutocomplete(): void { + this.cancelAutocompleteRequest(); + this.clearAutocompleteUi(); + } + + public isShowingAutocomplete(): boolean { + return this.autocompleteState !== null; + } + + private updateAutocomplete(): void { + if (!this.autocompleteState || !this.autocompleteProvider) return; + this.requestAutocomplete({ force: this.autocompleteState === "force", explicitTab: false }); + } +} diff --git a/cactus-code/packages/tui/src/components/image.ts b/cactus-code/packages/tui/src/components/image.ts new file mode 100644 index 000000000..50c1e0f78 --- /dev/null +++ b/cactus-code/packages/tui/src/components/image.ts @@ -0,0 +1,126 @@ +import { + allocateImageId, + getCapabilities, + getCellDimensions, + getImageDimensions, + type ImageDimensions, + imageFallback, + renderImage, +} from "../terminal-image.ts"; +import type { Component } from "../tui.ts"; + +export interface ImageTheme { + fallbackColor: (str: string) => string; +} + +export interface ImageOptions { + maxWidthCells?: number; + maxHeightCells?: number; + filename?: string; + /** Kitty image ID. If provided, reuses this ID (for animations/updates). */ + imageId?: number; +} + +export class Image implements Component { + private base64Data: string; + private mimeType: string; + private dimensions: ImageDimensions; + private theme: ImageTheme; + private options: ImageOptions; + private imageId?: number; + + private cachedLines?: string[]; + private cachedWidth?: number; + + constructor( + base64Data: string, + mimeType: string, + theme: ImageTheme, + options: ImageOptions = {}, + dimensions?: ImageDimensions, + ) { + this.base64Data = base64Data; + this.mimeType = mimeType; + this.theme = theme; + this.options = options; + this.dimensions = dimensions || getImageDimensions(base64Data, mimeType) || { widthPx: 800, heightPx: 600 }; + this.imageId = options.imageId; + } + + /** Get the Kitty image ID used by this image (if any). */ + getImageId(): number | undefined { + return this.imageId; + } + + invalidate(): void { + this.cachedLines = undefined; + this.cachedWidth = undefined; + } + + render(width: number): string[] { + if (this.cachedLines && this.cachedWidth === width) { + return this.cachedLines; + } + + const maxWidth = Math.max(1, Math.min(width - 2, this.options.maxWidthCells ?? 60)); + const cellDimensions = getCellDimensions(); + const defaultMaxHeight = Math.max(1, Math.ceil((maxWidth * cellDimensions.widthPx) / cellDimensions.heightPx)); + const maxHeight = this.options.maxHeightCells ?? defaultMaxHeight; + + const caps = getCapabilities(); + let lines: string[]; + + if (caps.images) { + if (caps.images === "kitty" && this.imageId === undefined) { + this.imageId = allocateImageId(); + } + const result = renderImage(this.base64Data, this.dimensions, { + maxWidthCells: maxWidth, + maxHeightCells: maxHeight, + imageId: this.imageId, + moveCursor: false, + }); + + if (result) { + // Store the image ID for later cleanup + if (result.imageId) { + this.imageId = result.imageId; + } + + if (caps.images === "kitty") { + // For Kitty: C=1 prevents cursor movement. + // Don't need the cursor movement. + lines = [result.sequence]; + + // Return `rows` lines so TUI accounts for image height. + for (let i = 0; i < result.rows - 1; i++) { + lines.push(""); + } + } else { + // Return `rows` lines so TUI accounts for image height. + // First (rows-1) lines are empty and cleared before the image is drawn. + // Last line: move cursor back up, draw the image, then move back down + // so TUI cursor accounting stays inside the scroll area. + lines = []; + for (let i = 0; i < result.rows - 1; i++) { + lines.push(""); + } + const rowOffset = result.rows - 1; + const moveUp = rowOffset > 0 ? `\x1b[${rowOffset}A` : ""; + lines.push(moveUp + result.sequence); + } + } else { + const fallback = imageFallback(this.mimeType, this.dimensions, this.options.filename); + lines = [this.theme.fallbackColor(fallback)]; + } + } else { + const fallback = imageFallback(this.mimeType, this.dimensions, this.options.filename); + lines = [this.theme.fallbackColor(fallback)]; + } + + this.cachedLines = lines; + this.cachedWidth = width; + + return lines; + } +} diff --git a/cactus-code/packages/tui/src/components/input.ts b/cactus-code/packages/tui/src/components/input.ts new file mode 100644 index 000000000..a054076d0 --- /dev/null +++ b/cactus-code/packages/tui/src/components/input.ts @@ -0,0 +1,447 @@ +import { getKeybindings } from "../keybindings.ts"; +import { decodeKittyPrintable } from "../keys.ts"; +import { KillRing } from "../kill-ring.ts"; +import { type Component, CURSOR_MARKER, type Focusable } from "../tui.ts"; +import { UndoStack } from "../undo-stack.ts"; +import { getGraphemeSegmenter, isWhitespaceChar, sliceByColumn, visibleWidth } from "../utils.ts"; +import { findWordBackward, findWordForward } from "../word-navigation.ts"; + +const segmenter = getGraphemeSegmenter(); + +interface InputState { + value: string; + cursor: number; +} + +/** + * Input component - single-line text input with horizontal scrolling + */ +export class Input implements Component, Focusable { + private value: string = ""; + private cursor: number = 0; // Cursor position in the value + public onSubmit?: (value: string) => void; + public onEscape?: () => void; + + /** Focusable interface - set by TUI when focus changes */ + focused: boolean = false; + + // Bracketed paste mode buffering + private pasteBuffer: string = ""; + private isInPaste: boolean = false; + + // Kill ring for Emacs-style kill/yank operations + private killRing = new KillRing(); + private lastAction: "kill" | "yank" | "type-word" | null = null; + + // Undo support + private undoStack = new UndoStack(); + + getValue(): string { + return this.value; + } + + setValue(value: string): void { + this.value = value; + this.cursor = Math.min(this.cursor, value.length); + } + + handleInput(data: string): void { + // Handle bracketed paste mode + // Start of paste: \x1b[200~ + // End of paste: \x1b[201~ + + // Check if we're starting a bracketed paste + if (data.includes("\x1b[200~")) { + this.isInPaste = true; + this.pasteBuffer = ""; + data = data.replace("\x1b[200~", ""); + } + + // If we're in a paste, buffer the data + if (this.isInPaste) { + // Check if this chunk contains the end marker + this.pasteBuffer += data; + + const endIndex = this.pasteBuffer.indexOf("\x1b[201~"); + if (endIndex !== -1) { + // Extract the pasted content + const pasteContent = this.pasteBuffer.substring(0, endIndex); + + // Process the complete paste + this.handlePaste(pasteContent); + + // Reset paste state + this.isInPaste = false; + + // Handle any remaining input after the paste marker + const remaining = this.pasteBuffer.substring(endIndex + 6); // 6 = length of \x1b[201~ + this.pasteBuffer = ""; + if (remaining) { + this.handleInput(remaining); + } + } + return; + } + + const kb = getKeybindings(); + + // Escape/Cancel + if (kb.matches(data, "tui.select.cancel")) { + if (this.onEscape) this.onEscape(); + return; + } + + // Undo + if (kb.matches(data, "tui.editor.undo")) { + this.undo(); + return; + } + + // Submit + if (kb.matches(data, "tui.input.submit") || data === "\n") { + if (this.onSubmit) this.onSubmit(this.value); + return; + } + + // Deletion + if (kb.matches(data, "tui.editor.deleteCharBackward")) { + this.handleBackspace(); + return; + } + + if (kb.matches(data, "tui.editor.deleteCharForward")) { + this.handleForwardDelete(); + return; + } + + if (kb.matches(data, "tui.editor.deleteWordBackward")) { + this.deleteWordBackwards(); + return; + } + + if (kb.matches(data, "tui.editor.deleteWordForward")) { + this.deleteWordForward(); + return; + } + + if (kb.matches(data, "tui.editor.deleteToLineStart")) { + this.deleteToLineStart(); + return; + } + + if (kb.matches(data, "tui.editor.deleteToLineEnd")) { + this.deleteToLineEnd(); + return; + } + + // Kill ring actions + if (kb.matches(data, "tui.editor.yank")) { + this.yank(); + return; + } + if (kb.matches(data, "tui.editor.yankPop")) { + this.yankPop(); + return; + } + + // Cursor movement + if (kb.matches(data, "tui.editor.cursorLeft")) { + this.lastAction = null; + if (this.cursor > 0) { + const beforeCursor = this.value.slice(0, this.cursor); + const graphemes = [...segmenter.segment(beforeCursor)]; + const lastGrapheme = graphemes[graphemes.length - 1]; + this.cursor -= lastGrapheme ? lastGrapheme.segment.length : 1; + } + return; + } + + if (kb.matches(data, "tui.editor.cursorRight")) { + this.lastAction = null; + if (this.cursor < this.value.length) { + const afterCursor = this.value.slice(this.cursor); + const graphemes = [...segmenter.segment(afterCursor)]; + const firstGrapheme = graphemes[0]; + this.cursor += firstGrapheme ? firstGrapheme.segment.length : 1; + } + return; + } + + if (kb.matches(data, "tui.editor.cursorLineStart")) { + this.lastAction = null; + this.cursor = 0; + return; + } + + if (kb.matches(data, "tui.editor.cursorLineEnd")) { + this.lastAction = null; + this.cursor = this.value.length; + return; + } + + if (kb.matches(data, "tui.editor.cursorWordLeft")) { + this.moveWordBackwards(); + return; + } + + if (kb.matches(data, "tui.editor.cursorWordRight")) { + this.moveWordForwards(); + return; + } + + // Kitty CSI-u printable character (e.g. \x1b[97u for 'a'). + // Terminals with Kitty protocol flag 1 (disambiguate) send CSI-u for all keys, + // including plain printable characters. Decode before the control-char check + // since CSI-u sequences contain \x1b which would be rejected. + const kittyPrintable = decodeKittyPrintable(data); + if (kittyPrintable !== undefined) { + this.insertCharacter(kittyPrintable); + return; + } + + // Regular character input - accept printable characters including Unicode, + // but reject control characters (C0: 0x00-0x1F, DEL: 0x7F, C1: 0x80-0x9F) + const hasControlChars = [...data].some((ch) => { + const code = ch.charCodeAt(0); + return code < 32 || code === 0x7f || (code >= 0x80 && code <= 0x9f); + }); + if (!hasControlChars) { + this.insertCharacter(data); + } + } + + private insertCharacter(char: string): void { + // Undo coalescing: consecutive word chars coalesce into one undo unit + if (isWhitespaceChar(char) || this.lastAction !== "type-word") { + this.pushUndo(); + } + this.lastAction = "type-word"; + + this.value = this.value.slice(0, this.cursor) + char + this.value.slice(this.cursor); + this.cursor += char.length; + } + + private handleBackspace(): void { + this.lastAction = null; + if (this.cursor > 0) { + this.pushUndo(); + const beforeCursor = this.value.slice(0, this.cursor); + const graphemes = [...segmenter.segment(beforeCursor)]; + const lastGrapheme = graphemes[graphemes.length - 1]; + const graphemeLength = lastGrapheme ? lastGrapheme.segment.length : 1; + this.value = this.value.slice(0, this.cursor - graphemeLength) + this.value.slice(this.cursor); + this.cursor -= graphemeLength; + } + } + + private handleForwardDelete(): void { + this.lastAction = null; + if (this.cursor < this.value.length) { + this.pushUndo(); + const afterCursor = this.value.slice(this.cursor); + const graphemes = [...segmenter.segment(afterCursor)]; + const firstGrapheme = graphemes[0]; + const graphemeLength = firstGrapheme ? firstGrapheme.segment.length : 1; + this.value = this.value.slice(0, this.cursor) + this.value.slice(this.cursor + graphemeLength); + } + } + + private deleteToLineStart(): void { + if (this.cursor === 0) return; + this.pushUndo(); + const deletedText = this.value.slice(0, this.cursor); + this.killRing.push(deletedText, { prepend: true, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + this.value = this.value.slice(this.cursor); + this.cursor = 0; + } + + private deleteToLineEnd(): void { + if (this.cursor >= this.value.length) return; + this.pushUndo(); + const deletedText = this.value.slice(this.cursor); + this.killRing.push(deletedText, { prepend: false, accumulate: this.lastAction === "kill" }); + this.lastAction = "kill"; + this.value = this.value.slice(0, this.cursor); + } + + private deleteWordBackwards(): void { + if (this.cursor === 0) return; + + // Save lastAction before cursor movement (moveWordBackwards resets it) + const wasKill = this.lastAction === "kill"; + + this.pushUndo(); + + const oldCursor = this.cursor; + this.moveWordBackwards(); + const deleteFrom = this.cursor; + this.cursor = oldCursor; + + const deletedText = this.value.slice(deleteFrom, this.cursor); + this.killRing.push(deletedText, { prepend: true, accumulate: wasKill }); + this.lastAction = "kill"; + + this.value = this.value.slice(0, deleteFrom) + this.value.slice(this.cursor); + this.cursor = deleteFrom; + } + + private deleteWordForward(): void { + if (this.cursor >= this.value.length) return; + + // Save lastAction before cursor movement (moveWordForwards resets it) + const wasKill = this.lastAction === "kill"; + + this.pushUndo(); + + const oldCursor = this.cursor; + this.moveWordForwards(); + const deleteTo = this.cursor; + this.cursor = oldCursor; + + const deletedText = this.value.slice(this.cursor, deleteTo); + this.killRing.push(deletedText, { prepend: false, accumulate: wasKill }); + this.lastAction = "kill"; + + this.value = this.value.slice(0, this.cursor) + this.value.slice(deleteTo); + } + + private yank(): void { + const text = this.killRing.peek(); + if (!text) return; + + this.pushUndo(); + + this.value = this.value.slice(0, this.cursor) + text + this.value.slice(this.cursor); + this.cursor += text.length; + this.lastAction = "yank"; + } + + private yankPop(): void { + if (this.lastAction !== "yank" || this.killRing.length <= 1) return; + + this.pushUndo(); + + // Delete the previously yanked text (still at end of ring before rotation) + const prevText = this.killRing.peek() || ""; + this.value = this.value.slice(0, this.cursor - prevText.length) + this.value.slice(this.cursor); + this.cursor -= prevText.length; + + // Rotate and insert new entry + this.killRing.rotate(); + const text = this.killRing.peek() || ""; + this.value = this.value.slice(0, this.cursor) + text + this.value.slice(this.cursor); + this.cursor += text.length; + this.lastAction = "yank"; + } + + private pushUndo(): void { + this.undoStack.push({ value: this.value, cursor: this.cursor }); + } + + private undo(): void { + const snapshot = this.undoStack.pop(); + if (!snapshot) return; + this.value = snapshot.value; + this.cursor = snapshot.cursor; + this.lastAction = null; + } + + private moveWordBackwards(): void { + if (this.cursor === 0) return; + this.lastAction = null; + this.cursor = findWordBackward(this.value, this.cursor); + } + + private moveWordForwards(): void { + if (this.cursor >= this.value.length) return; + this.lastAction = null; + this.cursor = findWordForward(this.value, this.cursor); + } + + private handlePaste(pastedText: string): void { + this.lastAction = null; + this.pushUndo(); + + // Clean the pasted text - remove newlines and carriage returns + const cleanText = pastedText.replace(/\r\n/g, "").replace(/\r/g, "").replace(/\n/g, "").replace(/\t/g, " "); + + // Insert at cursor position + this.value = this.value.slice(0, this.cursor) + cleanText + this.value.slice(this.cursor); + this.cursor += cleanText.length; + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(width: number): string[] { + // Calculate visible window + const prompt = "> "; + const availableWidth = width - prompt.length; + + if (availableWidth <= 0) { + return [prompt]; + } + + let visibleText = ""; + let cursorDisplay = this.cursor; + const totalWidth = visibleWidth(this.value); + + if (totalWidth < availableWidth) { + // Everything fits (leave room for cursor at end) + visibleText = this.value; + } else { + // Need horizontal scrolling + // Reserve one column for cursor if it's at the end + const scrollWidth = this.cursor === this.value.length ? availableWidth - 1 : availableWidth; + const cursorCol = visibleWidth(this.value.slice(0, this.cursor)); + + if (scrollWidth > 0) { + const halfWidth = Math.floor(scrollWidth / 2); + let startCol = 0; + + if (cursorCol < halfWidth) { + // Cursor near start + startCol = 0; + } else if (cursorCol > totalWidth - halfWidth) { + // Cursor near end + startCol = Math.max(0, totalWidth - scrollWidth); + } else { + // Cursor in middle + startCol = Math.max(0, cursorCol - halfWidth); + } + + visibleText = sliceByColumn(this.value, startCol, scrollWidth, true); + const beforeCursor = sliceByColumn(this.value, startCol, Math.max(0, cursorCol - startCol), true); + cursorDisplay = beforeCursor.length; + } else { + visibleText = ""; + cursorDisplay = 0; + } + } + + // Build line with fake cursor + // Insert cursor character at cursor position + const graphemes = [...segmenter.segment(visibleText.slice(cursorDisplay))]; + const cursorGrapheme = graphemes[0]; + + const beforeCursor = visibleText.slice(0, cursorDisplay); + const atCursor = cursorGrapheme?.segment ?? " "; // Character at cursor, or space if at end + const afterCursor = visibleText.slice(cursorDisplay + atCursor.length); + + // Hardware cursor marker (zero-width, emitted before fake cursor for IME positioning) + const marker = this.focused ? CURSOR_MARKER : ""; + + // Use inverse video to show cursor + const cursorChar = `\x1b[7m${atCursor}\x1b[27m`; // ESC[7m = reverse video, ESC[27m = normal + const textWithCursor = beforeCursor + marker + cursorChar + afterCursor; + + // Calculate visual width + const visualLength = visibleWidth(textWithCursor); + const padding = " ".repeat(Math.max(0, availableWidth - visualLength)); + const line = prompt + textWithCursor + padding; + + return [line]; + } +} diff --git a/cactus-code/packages/tui/src/components/loader.ts b/cactus-code/packages/tui/src/components/loader.ts new file mode 100644 index 000000000..4a93eb08f --- /dev/null +++ b/cactus-code/packages/tui/src/components/loader.ts @@ -0,0 +1,92 @@ +import type { TUI } from "../tui.ts"; +import { Text } from "./text.ts"; + +export interface LoaderIndicatorOptions { + /** Animation frames. Use an empty array to hide the indicator. */ + frames?: string[]; + /** Frame interval in milliseconds for animated indicators. */ + intervalMs?: number; +} + +const DEFAULT_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; +const DEFAULT_INTERVAL_MS = 80; + +/** + * Loader component that updates with an optional spinning animation. + */ +export class Loader extends Text { + private frames = [...DEFAULT_FRAMES]; + private intervalMs = DEFAULT_INTERVAL_MS; + private currentFrame = 0; + private intervalId: NodeJS.Timeout | null = null; + private ui: TUI | null = null; + private renderIndicatorVerbatim = false; + private spinnerColorFn: (str: string) => string; + private messageColorFn: (str: string) => string; + private message: string = "Loading..."; + + constructor( + ui: TUI, + spinnerColorFn: (str: string) => string, + messageColorFn: (str: string) => string, + message: string = "Loading...", + indicator?: LoaderIndicatorOptions, + ) { + super("", 1, 0); + this.ui = ui; + this.spinnerColorFn = spinnerColorFn; + this.messageColorFn = messageColorFn; + this.message = message; + this.setIndicator(indicator); + } + + render(width: number): string[] { + return ["", ...super.render(width)]; + } + + start(): void { + this.updateDisplay(); + this.restartAnimation(); + } + + stop(): void { + if (this.intervalId) { + clearInterval(this.intervalId); + this.intervalId = null; + } + } + + setMessage(message: string): void { + this.message = message; + this.updateDisplay(); + } + + setIndicator(indicator?: LoaderIndicatorOptions): void { + this.renderIndicatorVerbatim = indicator !== undefined; + this.frames = indicator?.frames !== undefined ? [...indicator.frames] : [...DEFAULT_FRAMES]; + this.intervalMs = indicator?.intervalMs && indicator.intervalMs > 0 ? indicator.intervalMs : DEFAULT_INTERVAL_MS; + this.currentFrame = 0; + this.start(); + } + + private restartAnimation(): void { + this.stop(); + if (this.frames.length <= 1) { + return; + } + this.intervalId = setInterval(() => { + this.currentFrame = (this.currentFrame + 1) % this.frames.length; + this.updateDisplay(); + }, this.intervalMs); + } + + private updateDisplay(): void { + const frame = this.frames[this.currentFrame] ?? ""; + const renderedFrame = this.renderIndicatorVerbatim ? frame : this.spinnerColorFn(frame); + const indicator = frame.length > 0 ? `${renderedFrame} ` : ""; + this.setText(`${indicator}${this.messageColorFn(this.message)}`); + if (this.ui) { + this.ui.requestRender(); + } + } +} diff --git a/cactus-code/packages/tui/src/components/markdown.ts b/cactus-code/packages/tui/src/components/markdown.ts new file mode 100644 index 000000000..1034cb476 --- /dev/null +++ b/cactus-code/packages/tui/src/components/markdown.ts @@ -0,0 +1,852 @@ +import { Marked, type Token, Tokenizer, type Tokens } from "marked"; +import { getCapabilities, hyperlink, isImageLine } from "../terminal-image.ts"; +import type { Component } from "../tui.ts"; +import { applyBackgroundToLine, visibleWidth, wrapTextWithAnsi } from "../utils.ts"; + +const STRICT_STRIKETHROUGH_REGEX = /^(~~)(?=[^\s~])((?:\\.|[^\\])*?(?:\\.|[^\s~\\]))\1(?=[^~]|$)/; + +class StrictStrikethroughTokenizer extends Tokenizer { + override del(src: string): Tokens.Del | undefined { + const match = STRICT_STRIKETHROUGH_REGEX.exec(src); + if (!match) { + return undefined; + } + + const text = match[2]; + return { + type: "del", + raw: match[0], + text, + tokens: this.lexer.inlineTokens(text), + }; + } +} + +function trimPartialClosingFences(tokens: readonly Token[]): void { + const token = tokens[tokens.length - 1]; + if (token?.type === "list") { + trimPartialClosingFences(token.items[token.items.length - 1]?.tokens ?? []); + return; + } + if (token?.type === "blockquote") { + trimPartialClosingFences(token.tokens ?? []); + return; + } + if (token?.type !== "code") { + return; + } + + // Trim streamed partial closing fences so code blocks do not shrink/flicker + // when the final fence character arrives. See https://github.com/earendil-works/pi/issues/5825. + const marker = /^(`{3,}|~{3,})/.exec(token.raw)?.[1]; + const lastLine = token.raw.split("\n").pop(); + if (!marker || !lastLine || lastLine.length >= marker.length || lastLine !== marker[0]?.repeat(lastLine.length)) { + return; + } + + token.text = token.text.slice(0, -lastLine.length).replace(/\n$/, ""); +} + +const markdownParser = new Marked(); +markdownParser.setOptions({ + tokenizer: new StrictStrikethroughTokenizer(), +}); + +/** + * Default text styling for markdown content. + * Applied to all text unless overridden by markdown formatting. + */ +export interface DefaultTextStyle { + /** Foreground color function */ + color?: (text: string) => string; + /** Background color function */ + bgColor?: (text: string) => string; + /** Bold text */ + bold?: boolean; + /** Italic text */ + italic?: boolean; + /** Strikethrough text */ + strikethrough?: boolean; + /** Underline text */ + underline?: boolean; +} + +/** + * Theme functions for markdown elements. + * Each function takes text and returns styled text with ANSI codes. + */ +export interface MarkdownTheme { + heading: (text: string) => string; + link: (text: string) => string; + linkUrl: (text: string) => string; + code: (text: string) => string; + codeBlock: (text: string) => string; + codeBlockBorder: (text: string) => string; + quote: (text: string) => string; + quoteBorder: (text: string) => string; + hr: (text: string) => string; + listBullet: (text: string) => string; + bold: (text: string) => string; + italic: (text: string) => string; + strikethrough: (text: string) => string; + underline: (text: string) => string; + highlightCode?: (code: string, lang?: string) => string[]; + /** Prefix applied to each rendered code block line (default: " ") */ + codeBlockIndent?: string; +} + +export interface MarkdownOptions { + /** Preserve source list markers instead of normalizing them. */ + preserveOrderedListMarkers?: boolean; +} + +interface InlineStyleContext { + applyText: (text: string) => string; + stylePrefix: string; +} + +export class Markdown implements Component { + private text: string; + private paddingX: number; // Left/right padding + private paddingY: number; // Top/bottom padding + private defaultTextStyle?: DefaultTextStyle; + private theme: MarkdownTheme; + private options: MarkdownOptions; + private defaultStylePrefix?: string; + + // Cache for rendered output + private cachedText?: string; + private cachedWidth?: number; + private cachedLines?: string[]; + + constructor( + text: string, + paddingX: number, + paddingY: number, + theme: MarkdownTheme, + defaultTextStyle?: DefaultTextStyle, + options?: MarkdownOptions, + ) { + this.text = text; + this.paddingX = paddingX; + this.paddingY = paddingY; + this.theme = theme; + this.defaultTextStyle = defaultTextStyle; + this.options = options ? { ...options } : {}; + } + + setText(text: string): void { + this.text = text; + this.invalidate(); + } + + invalidate(): void { + this.cachedText = undefined; + this.cachedWidth = undefined; + this.cachedLines = undefined; + } + + render(width: number): string[] { + // Check cache + if (this.cachedLines && this.cachedText === this.text && this.cachedWidth === width) { + return this.cachedLines; + } + + // Calculate available width for content (subtract horizontal padding) + const contentWidth = Math.max(1, width - this.paddingX * 2); + + // Don't render anything if there's no actual text + if (!this.text || this.text.trim() === "") { + const result: string[] = []; + // Update cache + this.cachedText = this.text; + this.cachedWidth = width; + this.cachedLines = result; + return result; + } + + // Replace tabs with 3 spaces for consistent rendering + const normalizedText = this.text.replace(/\t/g, " "); + + // Parse markdown to HTML-like tokens + const tokens = markdownParser.lexer(normalizedText); + trimPartialClosingFences(tokens); + + // Convert tokens to styled terminal output + const renderedLines: string[] = []; + + for (let i = 0; i < tokens.length; i++) { + const token = tokens[i]; + const nextToken = tokens[i + 1]; + const tokenLines = this.renderToken(token, contentWidth, nextToken?.type); + for (const tokenLine of tokenLines) { + renderedLines.push(tokenLine); + } + } + + // Wrap lines (NO padding, NO background yet) + const wrappedLines: string[] = []; + for (const line of renderedLines) { + if (isImageLine(line)) { + wrappedLines.push(line); + } else { + for (const wrappedLine of wrapTextWithAnsi(line, contentWidth)) { + wrappedLines.push(wrappedLine); + } + } + } + + // Add margins and background to each wrapped line + const leftMargin = " ".repeat(this.paddingX); + const rightMargin = " ".repeat(this.paddingX); + const bgFn = this.defaultTextStyle?.bgColor; + const contentLines: string[] = []; + + for (const line of wrappedLines) { + if (isImageLine(line)) { + contentLines.push(line); + continue; + } + + const lineWithMargins = leftMargin + line + rightMargin; + + if (bgFn) { + contentLines.push(applyBackgroundToLine(lineWithMargins, width, bgFn)); + } else { + // No background - just pad to width + const visibleLen = visibleWidth(lineWithMargins); + const paddingNeeded = Math.max(0, width - visibleLen); + contentLines.push(lineWithMargins + " ".repeat(paddingNeeded)); + } + } + + // Add top/bottom padding (empty lines) + const emptyLine = " ".repeat(width); + const emptyLines: string[] = []; + for (let i = 0; i < this.paddingY; i++) { + const line = bgFn ? applyBackgroundToLine(emptyLine, width, bgFn) : emptyLine; + emptyLines.push(line); + } + + // Combine top padding, content, and bottom padding + const result = emptyLines.concat(contentLines, emptyLines); + + // Update cache + this.cachedText = this.text; + this.cachedWidth = width; + this.cachedLines = result; + + return result.length > 0 ? result : [""]; + } + + /** + * Apply default text style to a string. + * This is the base styling applied to all text content. + * NOTE: Background color is NOT applied here - it's applied at the padding stage + * to ensure it extends to the full line width. + */ + private applyDefaultStyle(text: string): string { + if (!this.defaultTextStyle) { + return text; + } + + let styled = text; + + // Apply foreground color (NOT background - that's applied at padding stage) + if (this.defaultTextStyle.color) { + styled = this.defaultTextStyle.color(styled); + } + + // Apply text decorations using this.theme + if (this.defaultTextStyle.bold) { + styled = this.theme.bold(styled); + } + if (this.defaultTextStyle.italic) { + styled = this.theme.italic(styled); + } + if (this.defaultTextStyle.strikethrough) { + styled = this.theme.strikethrough(styled); + } + if (this.defaultTextStyle.underline) { + styled = this.theme.underline(styled); + } + + return styled; + } + + private getDefaultStylePrefix(): string { + if (!this.defaultTextStyle) { + return ""; + } + + if (this.defaultStylePrefix !== undefined) { + return this.defaultStylePrefix; + } + + const sentinel = "\u0000"; + let styled = sentinel; + + if (this.defaultTextStyle.color) { + styled = this.defaultTextStyle.color(styled); + } + + if (this.defaultTextStyle.bold) { + styled = this.theme.bold(styled); + } + if (this.defaultTextStyle.italic) { + styled = this.theme.italic(styled); + } + if (this.defaultTextStyle.strikethrough) { + styled = this.theme.strikethrough(styled); + } + if (this.defaultTextStyle.underline) { + styled = this.theme.underline(styled); + } + + const sentinelIndex = styled.indexOf(sentinel); + this.defaultStylePrefix = sentinelIndex >= 0 ? styled.slice(0, sentinelIndex) : ""; + return this.defaultStylePrefix; + } + + private getStylePrefix(styleFn: (text: string) => string): string { + const sentinel = "\u0000"; + const styled = styleFn(sentinel); + const sentinelIndex = styled.indexOf(sentinel); + return sentinelIndex >= 0 ? styled.slice(0, sentinelIndex) : ""; + } + + private getDefaultInlineStyleContext(): InlineStyleContext { + return { + applyText: (text: string) => this.applyDefaultStyle(text), + stylePrefix: this.getDefaultStylePrefix(), + }; + } + + private renderToken( + token: Token, + width: number, + nextTokenType?: string, + styleContext?: InlineStyleContext, + ): string[] { + const lines: string[] = []; + + switch (token.type) { + case "heading": { + const headingLevel = token.depth; + const headingPrefix = `${"#".repeat(headingLevel)} `; + + // Build a heading-specific style context so inline tokens (codespan, bold, etc.) + // restore heading styling after their own ANSI resets instead of falling back to + // the default text style. + let headingStyleFn: (text: string) => string; + if (headingLevel === 1) { + headingStyleFn = (text: string) => this.theme.heading(this.theme.bold(this.theme.underline(text))); + } else { + headingStyleFn = (text: string) => this.theme.heading(this.theme.bold(text)); + } + + const headingStyleContext: InlineStyleContext = { + applyText: headingStyleFn, + stylePrefix: this.getStylePrefix(headingStyleFn), + }; + + const headingText = this.renderInlineTokens(token.tokens || [], headingStyleContext); + const styledHeading = headingLevel >= 3 ? headingStyleFn(headingPrefix) + headingText : headingText; + lines.push(styledHeading); + if (nextTokenType && nextTokenType !== "space") { + lines.push(""); // Add spacing after headings (unless space token follows) + } + break; + } + + case "paragraph": { + const paragraphText = this.renderInlineTokens(token.tokens || [], styleContext); + lines.push(paragraphText); + // Don't add spacing if next token is space or list + if (nextTokenType && nextTokenType !== "list" && nextTokenType !== "space") { + lines.push(""); + } + break; + } + + case "text": + lines.push(this.renderInlineTokens([token], styleContext)); + break; + + case "code": { + const indent = this.theme.codeBlockIndent ?? " "; + lines.push(this.theme.codeBlockBorder(`\`\`\`${token.lang || ""}`)); + if (this.theme.highlightCode) { + const highlightedLines = this.theme.highlightCode(token.text, token.lang); + for (const hlLine of highlightedLines) { + lines.push(`${indent}${hlLine}`); + } + } else { + // Split code by newlines and style each line + const codeLines = token.text.split("\n"); + for (const codeLine of codeLines) { + lines.push(`${indent}${this.theme.codeBlock(codeLine)}`); + } + } + lines.push(this.theme.codeBlockBorder("```")); + if (nextTokenType && nextTokenType !== "space") { + lines.push(""); // Add spacing after code blocks (unless space token follows) + } + break; + } + + case "list": { + const listLines = this.renderList(token as Tokens.List, 0, width, styleContext); + lines.push(...listLines); + // Don't add spacing after lists if a space token follows + // (the space token will handle it) + break; + } + + case "table": { + const tableLines = this.renderTable(token as Tokens.Table, width, nextTokenType, styleContext); + lines.push(...tableLines); + break; + } + + case "blockquote": { + const quoteStyle = (text: string) => this.theme.quote(this.theme.italic(text)); + const quoteStylePrefix = this.getStylePrefix(quoteStyle); + const applyQuoteStyle = (line: string): string => { + if (!quoteStylePrefix) { + return quoteStyle(line); + } + const lineWithReappliedStyle = line.replace(/\x1b\[0m/g, `\x1b[0m${quoteStylePrefix}`); + return quoteStyle(lineWithReappliedStyle); + }; + + // Calculate available width for quote content (subtract border "│ " = 2 chars) + const quoteContentWidth = Math.max(1, width - 2); + + // Blockquotes contain block-level tokens (paragraph, list, code, etc.), so render + // children with renderToken() instead of renderInlineTokens(). + // Default message style should not apply inside blockquotes. + const quoteInlineStyleContext: InlineStyleContext = { + applyText: (text: string) => text, + stylePrefix: quoteStylePrefix, + }; + const quoteTokens = token.tokens || []; + const renderedQuoteLines: string[] = []; + for (let i = 0; i < quoteTokens.length; i++) { + const quoteToken = quoteTokens[i]; + const nextQuoteToken = quoteTokens[i + 1]; + renderedQuoteLines.push( + ...this.renderToken(quoteToken, quoteContentWidth, nextQuoteToken?.type, quoteInlineStyleContext), + ); + } + + // Avoid rendering an extra empty quote line before the outer blockquote spacing. + while (renderedQuoteLines.length > 0 && renderedQuoteLines[renderedQuoteLines.length - 1] === "") { + renderedQuoteLines.pop(); + } + + for (const quoteLine of renderedQuoteLines) { + const styledLine = applyQuoteStyle(quoteLine); + const wrappedLines = wrapTextWithAnsi(styledLine, quoteContentWidth); + for (const wrappedLine of wrappedLines) { + lines.push(this.theme.quoteBorder("│ ") + wrappedLine); + } + } + if (nextTokenType && nextTokenType !== "space") { + lines.push(""); // Add spacing after blockquotes (unless space token follows) + } + break; + } + + case "hr": + lines.push(this.theme.hr("─".repeat(Math.min(width, 80)))); + if (nextTokenType && nextTokenType !== "space") { + lines.push(""); // Add spacing after horizontal rules (unless space token follows) + } + break; + + case "html": + // Render HTML as plain text (escaped for terminal) + if ("raw" in token && typeof token.raw === "string") { + lines.push(this.applyDefaultStyle(token.raw.trim())); + } + break; + + case "space": + // Space tokens represent blank lines in markdown + lines.push(""); + break; + + default: + // Handle any other token types as plain text + if ("text" in token && typeof token.text === "string") { + lines.push(token.text); + } + } + + return lines; + } + + private renderInlineTokens(tokens: Token[], styleContext?: InlineStyleContext): string { + let result = ""; + const resolvedStyleContext = styleContext ?? this.getDefaultInlineStyleContext(); + const { applyText, stylePrefix } = resolvedStyleContext; + const applyTextWithNewlines = (text: string): string => { + const segments: string[] = text.split("\n"); + return segments.map((segment: string) => applyText(segment)).join("\n"); + }; + + for (const token of tokens) { + switch (token.type) { + case "text": + // Text tokens in list items can have nested tokens for inline formatting + if (token.tokens && token.tokens.length > 0) { + result += this.renderInlineTokens(token.tokens, resolvedStyleContext); + } else { + result += applyTextWithNewlines(token.text); + } + break; + + case "paragraph": + // Paragraph tokens contain nested inline tokens + result += this.renderInlineTokens(token.tokens || [], resolvedStyleContext); + break; + + case "strong": { + const boldContent = this.renderInlineTokens(token.tokens || [], resolvedStyleContext); + result += this.theme.bold(boldContent) + stylePrefix; + break; + } + + case "em": { + const italicContent = this.renderInlineTokens(token.tokens || [], resolvedStyleContext); + result += this.theme.italic(italicContent) + stylePrefix; + break; + } + + case "codespan": + result += this.theme.code(token.text) + stylePrefix; + break; + + case "link": { + const linkText = this.renderInlineTokens(token.tokens || [], resolvedStyleContext); + const styledLink = this.theme.link(this.theme.underline(linkText)); + if (getCapabilities().hyperlinks) { + // OSC 8: render as a clickable hyperlink. The URL is not printed inline, + // so we always show only the link text regardless of whether it matches href. + result += hyperlink(styledLink, token.href) + stylePrefix; + } else { + // Fallback: print URL in parentheses when text differs from href. + // Compare raw token.text (not styled) against href for the equality check. + // For mailto: links strip the prefix (autolinked emails use text="foo@bar.com" + // but href="mailto:foo@bar.com"). + const hrefForComparison = token.href.startsWith("mailto:") ? token.href.slice(7) : token.href; + if (token.text === token.href || token.text === hrefForComparison) { + result += styledLink + stylePrefix; + } else { + result += styledLink + this.theme.linkUrl(` (${token.href})`) + stylePrefix; + } + } + break; + } + + case "br": + result += "\n"; + break; + + case "del": { + const delContent = this.renderInlineTokens(token.tokens || [], resolvedStyleContext); + result += this.theme.strikethrough(delContent) + stylePrefix; + break; + } + + case "html": + // Render inline HTML as plain text + if ("raw" in token && typeof token.raw === "string") { + result += applyTextWithNewlines(token.raw); + } + break; + + default: + // Handle any other inline token types as plain text + if ("text" in token && typeof token.text === "string") { + result += applyTextWithNewlines(token.text); + } + } + } + + while (stylePrefix && result.endsWith(stylePrefix)) { + result = result.slice(0, -stylePrefix.length); + } + + return result; + } + + private getOrderedListMarker(item: Tokens.ListItem): string | undefined { + const match = /^(?: {0,3})(\d{1,9}[.)])[ \t]+/.exec(item.raw); + return match ? `${match[1]} ` : undefined; + } + + private getUnorderedListMarker(item: Tokens.ListItem): string | undefined { + const match = /^(?: {0,3})([-+*])(?:[ \t]+|(?=\r?\n|$))/.exec(item.raw); + return match ? `${match[1]} ` : undefined; + } + + /** + * Render a list with proper nesting support + */ + private renderList(token: Tokens.List, depth: number, width: number, styleContext?: InlineStyleContext): string[] { + const lines: string[] = []; + const indent = " ".repeat(depth); + // Use the list's start property (defaults to 1 for ordered lists) + const startNumber = typeof token.start === "number" ? token.start : 1; + + for (let i = 0; i < token.items.length; i++) { + const item = token.items[i]; + const isLastItem = i === token.items.length - 1; + const bullet = token.ordered + ? this.options.preserveOrderedListMarkers + ? (this.getOrderedListMarker(item) ?? `${startNumber + i}. `) + : `${startNumber + i}. ` + : this.options.preserveOrderedListMarkers + ? (this.getUnorderedListMarker(item) ?? "- ") + : "- "; + const taskMarker = item.task ? `[${item.checked ? "x" : " "}] ` : ""; + const marker = bullet + taskMarker; + const firstPrefix = indent + this.theme.listBullet(marker); + const continuationPrefix = indent + " ".repeat(visibleWidth(marker)); + const itemWidth = Math.max(1, width - visibleWidth(firstPrefix)); + let renderedAnyLine = false; + + for (const itemToken of item.tokens) { + if (itemToken.type === "list") { + lines.push(...this.renderList(itemToken as Tokens.List, depth + 1, width, styleContext)); + renderedAnyLine = true; + continue; + } + + const itemLines = this.renderToken(itemToken, itemWidth, undefined, styleContext); + for (const line of itemLines) { + for (const wrappedLine of wrapTextWithAnsi(line, itemWidth)) { + const linePrefix = renderedAnyLine ? continuationPrefix : firstPrefix; + lines.push(linePrefix + wrappedLine); + renderedAnyLine = true; + } + } + } + + if (!renderedAnyLine) { + lines.push(firstPrefix); + } + + if (token.loose && !isLastItem) { + lines.push(""); + } + } + + return lines; + } + + /** + * Get the visible width of the longest word in a string. + */ + private getLongestWordWidth(text: string, maxWidth?: number): number { + const words = text.split(/\s+/).filter((word) => word.length > 0); + let longest = 0; + for (const word of words) { + longest = Math.max(longest, visibleWidth(word)); + } + if (maxWidth === undefined) { + return longest; + } + return Math.min(longest, maxWidth); + } + + /** + * Wrap a table cell to fit into a column. + * + * Delegates to wrapTextWithAnsi() so ANSI codes + long tokens are handled + * consistently with the rest of the renderer. + */ + private wrapCellText(text: string, maxWidth: number): string[] { + return wrapTextWithAnsi(text, Math.max(1, maxWidth)); + } + + /** + * Render a table with width-aware cell wrapping. + * Cells that don't fit are wrapped to multiple lines. + */ + private renderTable( + token: Tokens.Table, + availableWidth: number, + nextTokenType?: string, + styleContext?: InlineStyleContext, + ): string[] { + const lines: string[] = []; + const numCols = token.header.length; + + if (numCols === 0) { + return lines; + } + + // Calculate border overhead: "│ " + (n-1) * " │ " + " │" + // = 2 + (n-1) * 3 + 2 = 3n + 1 + const borderOverhead = 3 * numCols + 1; + const availableForCells = availableWidth - borderOverhead; + if (availableForCells < numCols) { + // Too narrow to render a stable table. Fall back to raw markdown. + const fallbackLines = token.raw ? wrapTextWithAnsi(token.raw, availableWidth) : []; + if (nextTokenType && nextTokenType !== "space") { + fallbackLines.push(""); + } + return fallbackLines; + } + + const maxUnbrokenWordWidth = 30; + + // Calculate natural column widths (what each column needs without constraints) + const naturalWidths: number[] = []; + const minWordWidths: number[] = []; + for (let i = 0; i < numCols; i++) { + const headerText = this.renderInlineTokens(token.header[i].tokens || [], styleContext); + naturalWidths[i] = visibleWidth(headerText); + minWordWidths[i] = Math.max(1, this.getLongestWordWidth(headerText, maxUnbrokenWordWidth)); + } + for (const row of token.rows) { + for (let i = 0; i < row.length; i++) { + const cellText = this.renderInlineTokens(row[i].tokens || [], styleContext); + naturalWidths[i] = Math.max(naturalWidths[i] || 0, visibleWidth(cellText)); + minWordWidths[i] = Math.max( + minWordWidths[i] || 1, + this.getLongestWordWidth(cellText, maxUnbrokenWordWidth), + ); + } + } + + let minColumnWidths = minWordWidths; + let minCellsWidth = minColumnWidths.reduce((a, b) => a + b, 0); + + if (minCellsWidth > availableForCells) { + minColumnWidths = new Array(numCols).fill(1); + const remaining = availableForCells - numCols; + + if (remaining > 0) { + const totalWeight = minWordWidths.reduce((total, width) => total + Math.max(0, width - 1), 0); + const growth = minWordWidths.map((width) => { + const weight = Math.max(0, width - 1); + return totalWeight > 0 ? Math.floor((weight / totalWeight) * remaining) : 0; + }); + + for (let i = 0; i < numCols; i++) { + minColumnWidths[i] += growth[i] ?? 0; + } + + const allocated = growth.reduce((total, width) => total + width, 0); + let leftover = remaining - allocated; + for (let i = 0; leftover > 0 && i < numCols; i++) { + minColumnWidths[i]++; + leftover--; + } + } + + minCellsWidth = minColumnWidths.reduce((a, b) => a + b, 0); + } + + // Calculate column widths that fit within available width + const totalNaturalWidth = naturalWidths.reduce((a, b) => a + b, 0) + borderOverhead; + let columnWidths: number[]; + + if (totalNaturalWidth <= availableWidth) { + // Everything fits naturally + columnWidths = naturalWidths.map((width, index) => Math.max(width, minColumnWidths[index])); + } else { + // Need to shrink columns to fit + const totalGrowPotential = naturalWidths.reduce((total, width, index) => { + return total + Math.max(0, width - minColumnWidths[index]); + }, 0); + const extraWidth = Math.max(0, availableForCells - minCellsWidth); + columnWidths = minColumnWidths.map((minWidth, index) => { + const naturalWidth = naturalWidths[index]; + const minWidthDelta = Math.max(0, naturalWidth - minWidth); + let grow = 0; + if (totalGrowPotential > 0) { + grow = Math.floor((minWidthDelta / totalGrowPotential) * extraWidth); + } + return minWidth + grow; + }); + + // Adjust for rounding errors - distribute remaining space + const allocated = columnWidths.reduce((a, b) => a + b, 0); + let remaining = availableForCells - allocated; + while (remaining > 0) { + let grew = false; + for (let i = 0; i < numCols && remaining > 0; i++) { + if (columnWidths[i] < naturalWidths[i]) { + columnWidths[i]++; + remaining--; + grew = true; + } + } + if (!grew) { + break; + } + } + } + + // Render top border + const topBorderCells = columnWidths.map((w) => "─".repeat(w)); + lines.push(`┌─${topBorderCells.join("─┬─")}─┐`); + + // Render header with wrapping + const headerCellLines: string[][] = token.header.map((cell, i) => { + const text = this.renderInlineTokens(cell.tokens || [], styleContext); + return this.wrapCellText(text, columnWidths[i]); + }); + const headerLineCount = Math.max(...headerCellLines.map((c) => c.length)); + + for (let lineIdx = 0; lineIdx < headerLineCount; lineIdx++) { + const rowParts = headerCellLines.map((cellLines, colIdx) => { + const text = cellLines[lineIdx] || ""; + const padded = text + " ".repeat(Math.max(0, columnWidths[colIdx] - visibleWidth(text))); + return this.theme.bold(padded); + }); + lines.push(`│ ${rowParts.join(" │ ")} │`); + } + + // Render separator + const separatorCells = columnWidths.map((w) => "─".repeat(w)); + const separatorLine = `├─${separatorCells.join("─┼─")}─┤`; + lines.push(separatorLine); + + // Render rows with wrapping + for (let rowIndex = 0; rowIndex < token.rows.length; rowIndex++) { + const row = token.rows[rowIndex]; + const rowCellLines: string[][] = row.map((cell, i) => { + const text = this.renderInlineTokens(cell.tokens || [], styleContext); + return this.wrapCellText(text, columnWidths[i]); + }); + const rowLineCount = Math.max(...rowCellLines.map((c) => c.length)); + + for (let lineIdx = 0; lineIdx < rowLineCount; lineIdx++) { + const rowParts = rowCellLines.map((cellLines, colIdx) => { + const text = cellLines[lineIdx] || ""; + return text + " ".repeat(Math.max(0, columnWidths[colIdx] - visibleWidth(text))); + }); + lines.push(`│ ${rowParts.join(" │ ")} │`); + } + + if (rowIndex < token.rows.length - 1) { + lines.push(separatorLine); + } + } + + // Render bottom border + const bottomBorderCells = columnWidths.map((w) => "─".repeat(w)); + lines.push(`└─${bottomBorderCells.join("─┴─")}─┘`); + + if (nextTokenType && nextTokenType !== "space") { + lines.push(""); // Add spacing after table + } + return lines; + } +} diff --git a/cactus-code/packages/tui/src/components/select-list.ts b/cactus-code/packages/tui/src/components/select-list.ts new file mode 100644 index 000000000..26fdb685a --- /dev/null +++ b/cactus-code/packages/tui/src/components/select-list.ts @@ -0,0 +1,229 @@ +import { getKeybindings } from "../keybindings.ts"; +import type { Component } from "../tui.ts"; +import { truncateToWidth, visibleWidth } from "../utils.ts"; + +const DEFAULT_PRIMARY_COLUMN_WIDTH = 32; +const PRIMARY_COLUMN_GAP = 2; +const MIN_DESCRIPTION_WIDTH = 10; + +const normalizeToSingleLine = (text: string): string => text.replace(/[\r\n]+/g, " ").trim(); +const clamp = (value: number, min: number, max: number): number => Math.max(min, Math.min(value, max)); + +export interface SelectItem { + value: string; + label: string; + description?: string; +} + +export interface SelectListTheme { + selectedPrefix: (text: string) => string; + selectedText: (text: string) => string; + description: (text: string) => string; + scrollInfo: (text: string) => string; + noMatch: (text: string) => string; +} + +export interface SelectListTruncatePrimaryContext { + text: string; + maxWidth: number; + columnWidth: number; + item: SelectItem; + isSelected: boolean; +} + +export interface SelectListLayoutOptions { + minPrimaryColumnWidth?: number; + maxPrimaryColumnWidth?: number; + truncatePrimary?: (context: SelectListTruncatePrimaryContext) => string; +} + +export class SelectList implements Component { + private items: SelectItem[] = []; + private filteredItems: SelectItem[] = []; + private selectedIndex: number = 0; + private maxVisible: number = 5; + private theme: SelectListTheme; + private layout: SelectListLayoutOptions; + + public onSelect?: (item: SelectItem) => void; + public onCancel?: () => void; + public onSelectionChange?: (item: SelectItem) => void; + + constructor(items: SelectItem[], maxVisible: number, theme: SelectListTheme, layout: SelectListLayoutOptions = {}) { + this.items = items; + this.filteredItems = items; + this.maxVisible = maxVisible; + this.theme = theme; + this.layout = layout; + } + + setFilter(filter: string): void { + this.filteredItems = this.items.filter((item) => item.value.toLowerCase().startsWith(filter.toLowerCase())); + // Reset selection when filter changes + this.selectedIndex = 0; + } + + setSelectedIndex(index: number): void { + this.selectedIndex = Math.max(0, Math.min(index, this.filteredItems.length - 1)); + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(width: number): string[] { + const lines: string[] = []; + + // If no items match filter, show message + if (this.filteredItems.length === 0) { + lines.push(this.theme.noMatch(" No matching commands")); + return lines; + } + + const primaryColumnWidth = this.getPrimaryColumnWidth(); + + // Calculate visible range with scrolling + const startIndex = Math.max( + 0, + Math.min(this.selectedIndex - Math.floor(this.maxVisible / 2), this.filteredItems.length - this.maxVisible), + ); + const endIndex = Math.min(startIndex + this.maxVisible, this.filteredItems.length); + + // Render visible items + for (let i = startIndex; i < endIndex; i++) { + const item = this.filteredItems[i]; + if (!item) continue; + + const isSelected = i === this.selectedIndex; + const descriptionSingleLine = item.description ? normalizeToSingleLine(item.description) : undefined; + lines.push(this.renderItem(item, isSelected, width, descriptionSingleLine, primaryColumnWidth)); + } + + // Add scroll indicators if needed + if (startIndex > 0 || endIndex < this.filteredItems.length) { + const scrollText = ` (${this.selectedIndex + 1}/${this.filteredItems.length})`; + // Truncate if too long for terminal + lines.push(this.theme.scrollInfo(truncateToWidth(scrollText, width - 2, ""))); + } + + return lines; + } + + handleInput(keyData: string): void { + const kb = getKeybindings(); + // Up arrow - wrap to bottom when at top + if (kb.matches(keyData, "tui.select.up")) { + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.notifySelectionChange(); + } + // Down arrow - wrap to top when at bottom + else if (kb.matches(keyData, "tui.select.down")) { + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.notifySelectionChange(); + } + // Enter + else if (kb.matches(keyData, "tui.select.confirm")) { + const selectedItem = this.filteredItems[this.selectedIndex]; + if (selectedItem && this.onSelect) { + this.onSelect(selectedItem); + } + } + // Escape or Ctrl+C + else if (kb.matches(keyData, "tui.select.cancel")) { + if (this.onCancel) { + this.onCancel(); + } + } + } + + private renderItem( + item: SelectItem, + isSelected: boolean, + width: number, + descriptionSingleLine: string | undefined, + primaryColumnWidth: number, + ): string { + const prefix = isSelected ? "→ " : " "; + const prefixWidth = visibleWidth(prefix); + + if (descriptionSingleLine && width > 40) { + const effectivePrimaryColumnWidth = Math.max(1, Math.min(primaryColumnWidth, width - prefixWidth - 4)); + const maxPrimaryWidth = Math.max(1, effectivePrimaryColumnWidth - PRIMARY_COLUMN_GAP); + const truncatedValue = this.truncatePrimary(item, isSelected, maxPrimaryWidth, effectivePrimaryColumnWidth); + const truncatedValueWidth = visibleWidth(truncatedValue); + const spacing = " ".repeat(Math.max(1, effectivePrimaryColumnWidth - truncatedValueWidth)); + const descriptionStart = prefixWidth + truncatedValueWidth + spacing.length; + const remainingWidth = width - descriptionStart - 2; // -2 for safety + + if (remainingWidth > MIN_DESCRIPTION_WIDTH) { + const truncatedDesc = truncateToWidth(descriptionSingleLine, remainingWidth, ""); + if (isSelected) { + return this.theme.selectedText(`${prefix}${truncatedValue}${spacing}${truncatedDesc}`); + } + + const descText = this.theme.description(spacing + truncatedDesc); + return prefix + truncatedValue + descText; + } + } + + const maxWidth = width - prefixWidth - 2; + const truncatedValue = this.truncatePrimary(item, isSelected, maxWidth, maxWidth); + if (isSelected) { + return this.theme.selectedText(`${prefix}${truncatedValue}`); + } + + return prefix + truncatedValue; + } + + private getPrimaryColumnWidth(): number { + const { min, max } = this.getPrimaryColumnBounds(); + const widestPrimary = this.filteredItems.reduce((widest, item) => { + return Math.max(widest, visibleWidth(this.getDisplayValue(item)) + PRIMARY_COLUMN_GAP); + }, 0); + + return clamp(widestPrimary, min, max); + } + + private getPrimaryColumnBounds(): { min: number; max: number } { + const rawMin = + this.layout.minPrimaryColumnWidth ?? this.layout.maxPrimaryColumnWidth ?? DEFAULT_PRIMARY_COLUMN_WIDTH; + const rawMax = + this.layout.maxPrimaryColumnWidth ?? this.layout.minPrimaryColumnWidth ?? DEFAULT_PRIMARY_COLUMN_WIDTH; + + return { + min: Math.max(1, Math.min(rawMin, rawMax)), + max: Math.max(1, Math.max(rawMin, rawMax)), + }; + } + + private truncatePrimary(item: SelectItem, isSelected: boolean, maxWidth: number, columnWidth: number): string { + const displayValue = this.getDisplayValue(item); + const truncatedValue = this.layout.truncatePrimary + ? this.layout.truncatePrimary({ + text: displayValue, + maxWidth, + columnWidth, + item, + isSelected, + }) + : truncateToWidth(displayValue, maxWidth, ""); + + return truncateToWidth(truncatedValue, maxWidth, ""); + } + + private getDisplayValue(item: SelectItem): string { + return item.label || item.value; + } + + private notifySelectionChange(): void { + const selectedItem = this.filteredItems[this.selectedIndex]; + if (selectedItem && this.onSelectionChange) { + this.onSelectionChange(selectedItem); + } + } + + getSelectedItem(): SelectItem | null { + const item = this.filteredItems[this.selectedIndex]; + return item || null; + } +} diff --git a/cactus-code/packages/tui/src/components/settings-list.ts b/cactus-code/packages/tui/src/components/settings-list.ts new file mode 100644 index 000000000..88e9433f7 --- /dev/null +++ b/cactus-code/packages/tui/src/components/settings-list.ts @@ -0,0 +1,250 @@ +import { fuzzyFilter } from "../fuzzy.ts"; +import { getKeybindings } from "../keybindings.ts"; +import type { Component } from "../tui.ts"; +import { truncateToWidth, visibleWidth, wrapTextWithAnsi } from "../utils.ts"; +import { Input } from "./input.ts"; + +export interface SettingItem { + /** Unique identifier for this setting */ + id: string; + /** Display label (left side) */ + label: string; + /** Optional description shown when selected */ + description?: string; + /** Current value to display (right side) */ + currentValue: string; + /** If provided, Enter/Space cycles through these values */ + values?: string[]; + /** If provided, Enter opens this submenu. Receives current value and done callback. */ + submenu?: (currentValue: string, done: (selectedValue?: string) => void) => Component; +} + +export interface SettingsListTheme { + label: (text: string, selected: boolean) => string; + value: (text: string, selected: boolean) => string; + description: (text: string) => string; + cursor: string; + hint: (text: string) => string; +} + +export interface SettingsListOptions { + enableSearch?: boolean; +} + +export class SettingsList implements Component { + private items: SettingItem[]; + private filteredItems: SettingItem[]; + private theme: SettingsListTheme; + private selectedIndex = 0; + private maxVisible: number; + private onChange: (id: string, newValue: string) => void; + private onCancel: () => void; + private searchInput?: Input; + private searchEnabled: boolean; + + // Submenu state + private submenuComponent: Component | null = null; + private submenuItemIndex: number | null = null; + + constructor( + items: SettingItem[], + maxVisible: number, + theme: SettingsListTheme, + onChange: (id: string, newValue: string) => void, + onCancel: () => void, + options: SettingsListOptions = {}, + ) { + this.items = items; + this.filteredItems = items; + this.maxVisible = maxVisible; + this.theme = theme; + this.onChange = onChange; + this.onCancel = onCancel; + this.searchEnabled = options.enableSearch ?? false; + if (this.searchEnabled) { + this.searchInput = new Input(); + } + } + + /** Update an item's currentValue */ + updateValue(id: string, newValue: string): void { + const item = this.items.find((i) => i.id === id); + if (item) { + item.currentValue = newValue; + } + } + + invalidate(): void { + this.submenuComponent?.invalidate?.(); + } + + render(width: number): string[] { + // If submenu is active, render it instead + if (this.submenuComponent) { + return this.submenuComponent.render(width); + } + + return this.renderMainList(width); + } + + private renderMainList(width: number): string[] { + const lines: string[] = []; + + if (this.searchEnabled && this.searchInput) { + lines.push(...this.searchInput.render(width)); + lines.push(""); + } + + if (this.items.length === 0) { + lines.push(this.theme.hint(" No settings available")); + if (this.searchEnabled) { + this.addHintLine(lines, width); + } + return lines; + } + + const displayItems = this.searchEnabled ? this.filteredItems : this.items; + if (displayItems.length === 0) { + lines.push(truncateToWidth(this.theme.hint(" No matching settings"), width)); + this.addHintLine(lines, width); + return lines; + } + + // Calculate visible range with scrolling + const startIndex = Math.max( + 0, + Math.min(this.selectedIndex - Math.floor(this.maxVisible / 2), displayItems.length - this.maxVisible), + ); + const endIndex = Math.min(startIndex + this.maxVisible, displayItems.length); + + // Calculate max label width for alignment + const maxLabelWidth = Math.min(30, Math.max(...this.items.map((item) => visibleWidth(item.label)))); + + // Render visible items + for (let i = startIndex; i < endIndex; i++) { + const item = displayItems[i]; + if (!item) continue; + + const isSelected = i === this.selectedIndex; + const prefix = isSelected ? this.theme.cursor : " "; + const prefixWidth = visibleWidth(prefix); + + // Pad label to align values + const labelPadded = item.label + " ".repeat(Math.max(0, maxLabelWidth - visibleWidth(item.label))); + const labelText = this.theme.label(labelPadded, isSelected); + + // Calculate space for value + const separator = " "; + const usedWidth = prefixWidth + maxLabelWidth + visibleWidth(separator); + const valueMaxWidth = width - usedWidth - 2; + + const valueText = this.theme.value(truncateToWidth(item.currentValue, valueMaxWidth, ""), isSelected); + + lines.push(truncateToWidth(prefix + labelText + separator + valueText, width)); + } + + // Add scroll indicator if needed + if (startIndex > 0 || endIndex < displayItems.length) { + const scrollText = ` (${this.selectedIndex + 1}/${displayItems.length})`; + lines.push(this.theme.hint(truncateToWidth(scrollText, width - 2, ""))); + } + + // Add description for selected item + const selectedItem = displayItems[this.selectedIndex]; + if (selectedItem?.description) { + lines.push(""); + const wrappedDesc = wrapTextWithAnsi(selectedItem.description, width - 4); + for (const line of wrappedDesc) { + lines.push(this.theme.description(` ${line}`)); + } + } + + // Add hint + this.addHintLine(lines, width); + + return lines; + } + + handleInput(data: string): void { + // If submenu is active, delegate all input to it + // The submenu's onCancel (triggered by escape) will call done() which closes it + if (this.submenuComponent) { + this.submenuComponent.handleInput?.(data); + return; + } + + // Main list input handling + const kb = getKeybindings(); + const displayItems = this.searchEnabled ? this.filteredItems : this.items; + if (kb.matches(data, "tui.select.up")) { + if (displayItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? displayItems.length - 1 : this.selectedIndex - 1; + } else if (kb.matches(data, "tui.select.down")) { + if (displayItems.length === 0) return; + this.selectedIndex = this.selectedIndex === displayItems.length - 1 ? 0 : this.selectedIndex + 1; + } else if (kb.matches(data, "tui.select.confirm") || data === " ") { + this.activateItem(); + } else if (kb.matches(data, "tui.select.cancel")) { + this.onCancel(); + } else if (this.searchEnabled && this.searchInput) { + const sanitized = data.replace(/ /g, ""); + if (!sanitized) { + return; + } + this.searchInput.handleInput(sanitized); + this.applyFilter(this.searchInput.getValue()); + } + } + + private activateItem(): void { + const item = this.searchEnabled ? this.filteredItems[this.selectedIndex] : this.items[this.selectedIndex]; + if (!item) return; + + if (item.submenu) { + // Open submenu, passing current value so it can pre-select correctly + this.submenuItemIndex = this.selectedIndex; + this.submenuComponent = item.submenu(item.currentValue, (selectedValue?: string) => { + if (selectedValue !== undefined) { + item.currentValue = selectedValue; + this.onChange(item.id, selectedValue); + } + this.closeSubmenu(); + }); + } else if (item.values && item.values.length > 0) { + // Cycle through values + const currentIndex = item.values.indexOf(item.currentValue); + const nextIndex = (currentIndex + 1) % item.values.length; + const newValue = item.values[nextIndex]; + item.currentValue = newValue; + this.onChange(item.id, newValue); + } + } + + private closeSubmenu(): void { + this.submenuComponent = null; + // Restore selection to the item that opened the submenu + if (this.submenuItemIndex !== null) { + this.selectedIndex = this.submenuItemIndex; + this.submenuItemIndex = null; + } + } + + private applyFilter(query: string): void { + this.filteredItems = fuzzyFilter(this.items, query, (item) => item.label); + this.selectedIndex = 0; + } + + private addHintLine(lines: string[], width: number): void { + lines.push(""); + lines.push( + truncateToWidth( + this.theme.hint( + this.searchEnabled + ? " Type to search · Enter/Space to change · Esc to cancel" + : " Enter/Space to change · Esc to cancel", + ), + width, + ), + ); + } +} diff --git a/cactus-code/packages/tui/src/components/spacer.ts b/cactus-code/packages/tui/src/components/spacer.ts new file mode 100644 index 000000000..7abe1551c --- /dev/null +++ b/cactus-code/packages/tui/src/components/spacer.ts @@ -0,0 +1,28 @@ +import type { Component } from "../tui.ts"; + +/** + * Spacer component that renders empty lines + */ +export class Spacer implements Component { + private lines: number; + + constructor(lines: number = 1) { + this.lines = lines; + } + + setLines(lines: number): void { + this.lines = lines; + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(_width: number): string[] { + const result: string[] = []; + for (let i = 0; i < this.lines; i++) { + result.push(""); + } + return result; + } +} diff --git a/cactus-code/packages/tui/src/components/text.ts b/cactus-code/packages/tui/src/components/text.ts new file mode 100644 index 000000000..3809a48a8 --- /dev/null +++ b/cactus-code/packages/tui/src/components/text.ts @@ -0,0 +1,106 @@ +import type { Component } from "../tui.ts"; +import { applyBackgroundToLine, visibleWidth, wrapTextWithAnsi } from "../utils.ts"; + +/** + * Text component - displays multi-line text with word wrapping + */ +export class Text implements Component { + private text: string; + private paddingX: number; // Left/right padding + private paddingY: number; // Top/bottom padding + private customBgFn?: (text: string) => string; + + // Cache for rendered output + private cachedText?: string; + private cachedWidth?: number; + private cachedLines?: string[]; + + constructor(text: string = "", paddingX: number = 1, paddingY: number = 1, customBgFn?: (text: string) => string) { + this.text = text; + this.paddingX = paddingX; + this.paddingY = paddingY; + this.customBgFn = customBgFn; + } + + setText(text: string): void { + this.text = text; + this.cachedText = undefined; + this.cachedWidth = undefined; + this.cachedLines = undefined; + } + + setCustomBgFn(customBgFn?: (text: string) => string): void { + this.customBgFn = customBgFn; + this.cachedText = undefined; + this.cachedWidth = undefined; + this.cachedLines = undefined; + } + + invalidate(): void { + this.cachedText = undefined; + this.cachedWidth = undefined; + this.cachedLines = undefined; + } + + render(width: number): string[] { + // Check cache + if (this.cachedLines && this.cachedText === this.text && this.cachedWidth === width) { + return this.cachedLines; + } + + // Don't render anything if there's no actual text + if (!this.text || this.text.trim() === "") { + const result: string[] = []; + this.cachedText = this.text; + this.cachedWidth = width; + this.cachedLines = result; + return result; + } + + // Replace tabs with 3 spaces + const normalizedText = this.text.replace(/\t/g, " "); + + // Calculate content width (subtract left/right margins) + const contentWidth = Math.max(1, width - this.paddingX * 2); + + // Wrap text (this preserves ANSI codes but does NOT pad) + const wrappedLines = wrapTextWithAnsi(normalizedText, contentWidth); + + // Add margins and background to each line + const leftMargin = " ".repeat(this.paddingX); + const rightMargin = " ".repeat(this.paddingX); + const contentLines: string[] = []; + + for (const line of wrappedLines) { + // Add margins + const lineWithMargins = leftMargin + line + rightMargin; + + // Apply background if specified (this also pads to full width) + if (this.customBgFn) { + contentLines.push(applyBackgroundToLine(lineWithMargins, width, this.customBgFn)); + } else { + // No background - just pad to width with spaces + const visibleLen = visibleWidth(lineWithMargins); + const paddingNeeded = Math.max(0, width - visibleLen); + contentLines.push(lineWithMargins + " ".repeat(paddingNeeded)); + } + } + + // Add top/bottom padding (empty lines) + const emptyLine = " ".repeat(width); + const emptyLines: string[] = []; + for (let i = 0; i < this.paddingY; i++) { + const line = this.customBgFn ? applyBackgroundToLine(emptyLine, width, this.customBgFn) : emptyLine; + emptyLines.push(line); + } + + const result = [...emptyLines, ...contentLines, ...emptyLines]; + + // Update cache + this.cachedText = this.text; + this.cachedWidth = width; + this.cachedLines = result; + + return result.length > 0 ? result : [""]; + } +} diff --git a/cactus-code/packages/tui/src/components/truncated-text.ts b/cactus-code/packages/tui/src/components/truncated-text.ts new file mode 100644 index 000000000..c26b88929 --- /dev/null +++ b/cactus-code/packages/tui/src/components/truncated-text.ts @@ -0,0 +1,65 @@ +import type { Component } from "../tui.ts"; +import { truncateToWidth, visibleWidth } from "../utils.ts"; + +/** + * Text component that truncates to fit viewport width + */ +export class TruncatedText implements Component { + private text: string; + private paddingX: number; + private paddingY: number; + + constructor(text: string, paddingX: number = 0, paddingY: number = 0) { + this.text = text; + this.paddingX = paddingX; + this.paddingY = paddingY; + } + + invalidate(): void { + // No cached state to invalidate currently + } + + render(width: number): string[] { + const result: string[] = []; + + // Empty line padded to width + const emptyLine = " ".repeat(width); + + // Add vertical padding above + for (let i = 0; i < this.paddingY; i++) { + result.push(emptyLine); + } + + // Calculate available width after horizontal padding + const availableWidth = Math.max(1, width - this.paddingX * 2); + + // Take only the first line (stop at newline) + let singleLineText = this.text; + const newlineIndex = this.text.indexOf("\n"); + if (newlineIndex !== -1) { + singleLineText = this.text.substring(0, newlineIndex); + } + + // Truncate text if needed (accounting for ANSI codes) + const displayText = truncateToWidth(singleLineText, availableWidth); + + // Add horizontal padding + const leftPadding = " ".repeat(this.paddingX); + const rightPadding = " ".repeat(this.paddingX); + const lineWithPadding = leftPadding + displayText + rightPadding; + + // Pad line to exactly width characters + const lineVisibleWidth = visibleWidth(lineWithPadding); + const paddingNeeded = Math.max(0, width - lineVisibleWidth); + const finalLine = lineWithPadding + " ".repeat(paddingNeeded); + + result.push(finalLine); + + // Add vertical padding below + for (let i = 0; i < this.paddingY; i++) { + result.push(emptyLine); + } + + return result; + } +} diff --git a/cactus-code/packages/tui/src/editor-component.ts b/cactus-code/packages/tui/src/editor-component.ts new file mode 100644 index 000000000..595645033 --- /dev/null +++ b/cactus-code/packages/tui/src/editor-component.ts @@ -0,0 +1,74 @@ +import type { AutocompleteProvider } from "./autocomplete.ts"; +import type { Component } from "./tui.ts"; + +/** + * Interface for custom editor components. + * + * This allows extensions to provide their own editor implementation + * (e.g., vim mode, emacs mode, custom keybindings) while maintaining + * compatibility with the core application. + */ +export interface EditorComponent extends Component { + // ========================================================================= + // Core text access (required) + // ========================================================================= + + /** Get the current text content */ + getText(): string; + + /** Set the text content */ + setText(text: string): void; + + /** Handle raw terminal input (key presses, paste sequences, etc.) */ + handleInput(data: string): void; + + // ========================================================================= + // Callbacks (required) + // ========================================================================= + + /** Called when user submits (e.g., Enter key) */ + onSubmit?: (text: string) => void; + + /** Called when text changes */ + onChange?: (text: string) => void; + + // ========================================================================= + // History support (optional) + // ========================================================================= + + /** Add text to history for up/down navigation */ + addToHistory?(text: string): void; + + // ========================================================================= + // Advanced text manipulation (optional) + // ========================================================================= + + /** Insert text at current cursor position */ + insertTextAtCursor?(text: string): void; + + /** + * Get text with any markers expanded (e.g., paste markers). + * Falls back to getText() if not implemented. + */ + getExpandedText?(): string; + + // ========================================================================= + // Autocomplete support (optional) + // ========================================================================= + + /** Set the autocomplete provider */ + setAutocompleteProvider?(provider: AutocompleteProvider): void; + + // ========================================================================= + // Appearance (optional) + // ========================================================================= + + /** Border color function */ + borderColor?: (str: string) => string; + + /** Set horizontal padding */ + setPaddingX?(padding: number): void; + + /** Set max visible items in autocomplete dropdown */ + setAutocompleteMaxVisible?(maxVisible: number): void; +} diff --git a/cactus-code/packages/tui/src/fuzzy.ts b/cactus-code/packages/tui/src/fuzzy.ts new file mode 100644 index 000000000..73c10dcf4 --- /dev/null +++ b/cactus-code/packages/tui/src/fuzzy.ts @@ -0,0 +1,137 @@ +/** + * Fuzzy matching utilities. + * Matches if all query characters appear in order (not necessarily consecutive). + * Lower score = better match. + */ + +export interface FuzzyMatch { + matches: boolean; + score: number; +} + +export function fuzzyMatch(query: string, text: string): FuzzyMatch { + const queryLower = query.toLowerCase(); + const textLower = text.toLowerCase(); + + const matchQuery = (normalizedQuery: string): FuzzyMatch => { + if (normalizedQuery.length === 0) { + return { matches: true, score: 0 }; + } + + if (normalizedQuery.length > textLower.length) { + return { matches: false, score: 0 }; + } + + let queryIndex = 0; + let score = 0; + let lastMatchIndex = -1; + let consecutiveMatches = 0; + + for (let i = 0; i < textLower.length && queryIndex < normalizedQuery.length; i++) { + if (textLower[i] === normalizedQuery[queryIndex]) { + const isWordBoundary = i === 0 || /[\s\-_./:]/.test(textLower[i - 1]!); + + // Reward consecutive matches + if (lastMatchIndex === i - 1) { + consecutiveMatches++; + score -= consecutiveMatches * 5; + } else { + consecutiveMatches = 0; + // Penalize gaps + if (lastMatchIndex >= 0) { + score += (i - lastMatchIndex - 1) * 2; + } + } + + // Reward word boundary matches + if (isWordBoundary) { + score -= 10; + } + + // Slight penalty for later matches + score += i * 0.1; + + lastMatchIndex = i; + queryIndex++; + } + } + + if (queryIndex < normalizedQuery.length) { + return { matches: false, score: 0 }; + } + + if (normalizedQuery === textLower) { + score -= 100; + } + + return { matches: true, score }; + }; + + const primaryMatch = matchQuery(queryLower); + if (primaryMatch.matches) { + return primaryMatch; + } + + const alphaNumericMatch = queryLower.match(/^(?[a-z]+)(?[0-9]+)$/); + const numericAlphaMatch = queryLower.match(/^(?[0-9]+)(?[a-z]+)$/); + const swappedQuery = alphaNumericMatch + ? `${alphaNumericMatch.groups?.digits ?? ""}${alphaNumericMatch.groups?.letters ?? ""}` + : numericAlphaMatch + ? `${numericAlphaMatch.groups?.letters ?? ""}${numericAlphaMatch.groups?.digits ?? ""}` + : ""; + + if (!swappedQuery) { + return primaryMatch; + } + + const swappedMatch = matchQuery(swappedQuery); + if (!swappedMatch.matches) { + return primaryMatch; + } + + return { matches: true, score: swappedMatch.score + 5 }; +} + +/** + * Filter and sort items by fuzzy match quality (best matches first). + * Supports whitespace- and slash-separated tokens: all tokens must match. + */ +export function fuzzyFilter(items: T[], query: string, getText: (item: T) => string): T[] { + if (!query.trim()) { + return items; + } + + const tokens = query + .trim() + .split(/[\s/]+/) + .filter((t) => t.length > 0); + + if (tokens.length === 0) { + return items; + } + + const results: { item: T; totalScore: number }[] = []; + + for (const item of items) { + const text = getText(item); + let totalScore = 0; + let allMatch = true; + + for (const token of tokens) { + const match = fuzzyMatch(token, text); + if (match.matches) { + totalScore += match.score; + } else { + allMatch = false; + break; + } + } + + if (allMatch) { + results.push({ item, totalScore }); + } + } + + results.sort((a, b) => a.totalScore - b.totalScore); + return results.map((r) => r.item); +} diff --git a/cactus-code/packages/tui/src/index.ts b/cactus-code/packages/tui/src/index.ts new file mode 100644 index 000000000..4e76b1079 --- /dev/null +++ b/cactus-code/packages/tui/src/index.ts @@ -0,0 +1,114 @@ +// Core TUI interfaces and classes + +// Autocomplete support +export { + type AutocompleteItem, + type AutocompleteProvider, + type AutocompleteSuggestions, + CombinedAutocompleteProvider, + type SlashCommand, +} from "./autocomplete.ts"; +// Components +export { Box } from "./components/box.ts"; +export { CancellableLoader } from "./components/cancellable-loader.ts"; +export { Editor, type EditorOptions, type EditorTheme } from "./components/editor.ts"; +export { Image, type ImageOptions, type ImageTheme } from "./components/image.ts"; +export { Input } from "./components/input.ts"; +export { Loader, type LoaderIndicatorOptions } from "./components/loader.ts"; +export { type DefaultTextStyle, Markdown, type MarkdownOptions, type MarkdownTheme } from "./components/markdown.ts"; +export { + type SelectItem, + SelectList, + type SelectListLayoutOptions, + type SelectListTheme, + type SelectListTruncatePrimaryContext, +} from "./components/select-list.ts"; +export { type SettingItem, SettingsList, type SettingsListTheme } from "./components/settings-list.ts"; +export { Spacer } from "./components/spacer.ts"; +export { Text } from "./components/text.ts"; +export { TruncatedText } from "./components/truncated-text.ts"; +// Editor component interface (for custom editors) +export type { EditorComponent } from "./editor-component.ts"; +// Fuzzy matching +export { type FuzzyMatch, fuzzyFilter, fuzzyMatch } from "./fuzzy.ts"; +// Keybindings +export { + getKeybindings, + type Keybinding, + type KeybindingConflict, + type KeybindingDefinition, + type KeybindingDefinitions, + type Keybindings, + type KeybindingsConfig, + KeybindingsManager, + setKeybindings, + TUI_KEYBINDINGS, +} from "./keybindings.ts"; +// Keyboard input handling +export { + decodeKittyPrintable, + isKeyRelease, + isKeyRepeat, + isKittyProtocolActive, + Key, + type KeyEventType, + type KeyId, + matchesKey, + parseKey, + setKittyProtocolActive, +} from "./keys.ts"; +// Input buffering for batch splitting +export { StdinBuffer, type StdinBufferEventMap, type StdinBufferOptions } from "./stdin-buffer.ts"; +// Terminal interface and implementations +export { ProcessTerminal, type Terminal } from "./terminal.ts"; +// Terminal colors +export { + parseOsc11BackgroundColor, + parseTerminalColorSchemeReport, + type RgbColor, + type TerminalColorScheme, +} from "./terminal-colors.ts"; +// Terminal image support +export { + allocateImageId, + type CellDimensions, + calculateImageRows, + deleteAllKittyImages, + deleteKittyImage, + detectCapabilities, + encodeITerm2, + encodeKitty, + getCapabilities, + getCellDimensions, + getGifDimensions, + getImageDimensions, + getJpegDimensions, + getPngDimensions, + getWebpDimensions, + hyperlink, + type ImageDimensions, + type ImageProtocol, + type ImageRenderOptions, + imageFallback, + renderImage, + resetCapabilitiesCache, + setCapabilities, + setCellDimensions, + type TerminalCapabilities, +} from "./terminal-image.ts"; +export { + type Component, + Container, + CURSOR_MARKER, + type Focusable, + isFocusable, + type OverlayAnchor, + type OverlayHandle, + type OverlayMargin, + type OverlayOptions, + type OverlayUnfocusOptions, + type SizeValue, + TUI, +} from "./tui.ts"; +// Utilities +export { sliceByColumn, truncateToWidth, visibleWidth, wrapTextWithAnsi } from "./utils.ts"; diff --git a/cactus-code/packages/tui/src/keybindings.ts b/cactus-code/packages/tui/src/keybindings.ts new file mode 100644 index 000000000..eaf0838b4 --- /dev/null +++ b/cactus-code/packages/tui/src/keybindings.ts @@ -0,0 +1,244 @@ +import { type KeyId, matchesKey } from "./keys.ts"; + +/** + * Global keybinding registry. + * Downstream packages can add keybindings via declaration merging. + */ +export interface Keybindings { + // Editor navigation and editing + "tui.editor.cursorUp": true; + "tui.editor.cursorDown": true; + "tui.editor.cursorLeft": true; + "tui.editor.cursorRight": true; + "tui.editor.cursorWordLeft": true; + "tui.editor.cursorWordRight": true; + "tui.editor.cursorLineStart": true; + "tui.editor.cursorLineEnd": true; + "tui.editor.jumpForward": true; + "tui.editor.jumpBackward": true; + "tui.editor.pageUp": true; + "tui.editor.pageDown": true; + "tui.editor.deleteCharBackward": true; + "tui.editor.deleteCharForward": true; + "tui.editor.deleteWordBackward": true; + "tui.editor.deleteWordForward": true; + "tui.editor.deleteToLineStart": true; + "tui.editor.deleteToLineEnd": true; + "tui.editor.yank": true; + "tui.editor.yankPop": true; + "tui.editor.undo": true; + // Generic input actions + "tui.input.newLine": true; + "tui.input.submit": true; + "tui.input.tab": true; + "tui.input.copy": true; + // Generic selection actions + "tui.select.up": true; + "tui.select.down": true; + "tui.select.pageUp": true; + "tui.select.pageDown": true; + "tui.select.confirm": true; + "tui.select.cancel": true; +} + +export type Keybinding = keyof Keybindings; + +export interface KeybindingDefinition { + defaultKeys: KeyId | KeyId[]; + description?: string; +} + +export type KeybindingDefinitions = Record; +export type KeybindingsConfig = Record; + +export const TUI_KEYBINDINGS = { + "tui.editor.cursorUp": { defaultKeys: "up", description: "Move cursor up" }, + "tui.editor.cursorDown": { defaultKeys: "down", description: "Move cursor down" }, + "tui.editor.cursorLeft": { + defaultKeys: ["left", "ctrl+b"], + description: "Move cursor left", + }, + "tui.editor.cursorRight": { + defaultKeys: ["right", "ctrl+f"], + description: "Move cursor right", + }, + "tui.editor.cursorWordLeft": { + defaultKeys: ["alt+left", "ctrl+left", "alt+b"], + description: "Move cursor word left", + }, + "tui.editor.cursorWordRight": { + defaultKeys: ["alt+right", "ctrl+right", "alt+f"], + description: "Move cursor word right", + }, + "tui.editor.cursorLineStart": { + defaultKeys: ["home", "ctrl+a"], + description: "Move to line start", + }, + "tui.editor.cursorLineEnd": { + defaultKeys: ["end", "ctrl+e"], + description: "Move to line end", + }, + "tui.editor.jumpForward": { + defaultKeys: "ctrl+]", + description: "Jump forward to character", + }, + "tui.editor.jumpBackward": { + defaultKeys: "ctrl+alt+]", + description: "Jump backward to character", + }, + "tui.editor.pageUp": { defaultKeys: "pageUp", description: "Page up" }, + "tui.editor.pageDown": { defaultKeys: "pageDown", description: "Page down" }, + "tui.editor.deleteCharBackward": { + defaultKeys: "backspace", + description: "Delete character backward", + }, + "tui.editor.deleteCharForward": { + defaultKeys: ["delete", "ctrl+d"], + description: "Delete character forward", + }, + "tui.editor.deleteWordBackward": { + defaultKeys: ["ctrl+w", "alt+backspace"], + description: "Delete word backward", + }, + "tui.editor.deleteWordForward": { + defaultKeys: ["alt+d", "alt+delete"], + description: "Delete word forward", + }, + "tui.editor.deleteToLineStart": { + defaultKeys: "ctrl+u", + description: "Delete to line start", + }, + "tui.editor.deleteToLineEnd": { + defaultKeys: "ctrl+k", + description: "Delete to line end", + }, + "tui.editor.yank": { defaultKeys: "ctrl+y", description: "Yank" }, + "tui.editor.yankPop": { defaultKeys: "alt+y", description: "Yank pop" }, + "tui.editor.undo": { defaultKeys: "ctrl+-", description: "Undo" }, + "tui.input.newLine": { defaultKeys: ["shift+enter", "ctrl+j"], description: "Insert newline" }, + "tui.input.submit": { defaultKeys: "enter", description: "Submit input" }, + "tui.input.tab": { defaultKeys: "tab", description: "Tab / autocomplete" }, + "tui.input.copy": { defaultKeys: "ctrl+c", description: "Copy selection" }, + "tui.select.up": { defaultKeys: "up", description: "Move selection up" }, + "tui.select.down": { defaultKeys: "down", description: "Move selection down" }, + "tui.select.pageUp": { defaultKeys: "pageUp", description: "Selection page up" }, + "tui.select.pageDown": { + defaultKeys: "pageDown", + description: "Selection page down", + }, + "tui.select.confirm": { defaultKeys: "enter", description: "Confirm selection" }, + "tui.select.cancel": { + defaultKeys: ["escape", "ctrl+c"], + description: "Cancel selection", + }, +} as const satisfies KeybindingDefinitions; + +export interface KeybindingConflict { + key: KeyId; + keybindings: string[]; +} + +function normalizeKeys(keys: KeyId | KeyId[] | undefined): KeyId[] { + if (keys === undefined) return []; + const keyList = Array.isArray(keys) ? keys : [keys]; + const seen = new Set(); + const result: KeyId[] = []; + for (const key of keyList) { + if (!seen.has(key)) { + seen.add(key); + result.push(key); + } + } + return result; +} + +export class KeybindingsManager { + private definitions: KeybindingDefinitions; + private userBindings: KeybindingsConfig; + private keysById = new Map(); + private conflicts: KeybindingConflict[] = []; + + constructor(definitions: KeybindingDefinitions, userBindings: KeybindingsConfig = {}) { + this.definitions = definitions; + this.userBindings = userBindings; + this.rebuild(); + } + + private rebuild(): void { + this.keysById.clear(); + this.conflicts = []; + + const userClaims = new Map>(); + for (const [keybinding, keys] of Object.entries(this.userBindings)) { + if (!(keybinding in this.definitions)) continue; + for (const key of normalizeKeys(keys)) { + const claimants = userClaims.get(key) ?? new Set(); + claimants.add(keybinding as Keybinding); + userClaims.set(key, claimants); + } + } + + for (const [key, keybindings] of userClaims) { + if (keybindings.size > 1) { + this.conflicts.push({ key, keybindings: [...keybindings] }); + } + } + + for (const [id, definition] of Object.entries(this.definitions)) { + const userKeys = this.userBindings[id]; + const keys = userKeys === undefined ? normalizeKeys(definition.defaultKeys) : normalizeKeys(userKeys); + this.keysById.set(id as Keybinding, keys); + } + } + + matches(data: string, keybinding: Keybinding): boolean { + const keys = this.keysById.get(keybinding) ?? []; + for (const key of keys) { + if (matchesKey(data, key)) return true; + } + return false; + } + + getKeys(keybinding: Keybinding): KeyId[] { + return [...(this.keysById.get(keybinding) ?? [])]; + } + + getDefinition(keybinding: Keybinding): KeybindingDefinition { + return this.definitions[keybinding]; + } + + getConflicts(): KeybindingConflict[] { + return this.conflicts.map((conflict) => ({ ...conflict, keybindings: [...conflict.keybindings] })); + } + + setUserBindings(userBindings: KeybindingsConfig): void { + this.userBindings = userBindings; + this.rebuild(); + } + + getUserBindings(): KeybindingsConfig { + return { ...this.userBindings }; + } + + getResolvedBindings(): KeybindingsConfig { + const resolved: KeybindingsConfig = {}; + for (const id of Object.keys(this.definitions)) { + const keys = this.keysById.get(id as Keybinding) ?? []; + resolved[id] = keys.length === 1 ? keys[0]! : [...keys]; + } + return resolved; + } +} + +let globalKeybindings: KeybindingsManager | null = null; + +export function setKeybindings(keybindings: KeybindingsManager): void { + globalKeybindings = keybindings; +} + +export function getKeybindings(): KeybindingsManager { + if (!globalKeybindings) { + globalKeybindings = new KeybindingsManager(TUI_KEYBINDINGS); + } + return globalKeybindings; +} diff --git a/cactus-code/packages/tui/src/keys.ts b/cactus-code/packages/tui/src/keys.ts new file mode 100644 index 000000000..0d3c7085f --- /dev/null +++ b/cactus-code/packages/tui/src/keys.ts @@ -0,0 +1,1400 @@ +/** + * Keyboard input handling for terminal applications. + * + * Supports both legacy terminal sequences and Kitty keyboard protocol. + * See: https://sw.kovidgoyal.net/kitty/keyboard-protocol/ + * Reference: https://github.com/sst/opentui/blob/7da92b4088aebfe27b9f691c04163a48821e49fd/packages/core/src/lib/parse.keypress.ts + * + * Symbol keys are also supported, however some ctrl+symbol combos + * overlap with ASCII codes, e.g. ctrl+[ = ESC. + * See: https://sw.kovidgoyal.net/kitty/keyboard-protocol/#legacy-ctrl-mapping-of-ascii-keys + * Those can still be * used for ctrl+shift combos + * + * API: + * - matchesKey(data, keyId) - Check if input matches a key identifier + * - parseKey(data) - Parse input and return the key identifier + * - Key - Helper object for creating typed key identifiers + * - setKittyProtocolActive(active) - Set global Kitty protocol state + * - isKittyProtocolActive() - Query global Kitty protocol state + */ + +// ============================================================================= +// Global Kitty Protocol State +// ============================================================================= + +let _kittyProtocolActive = false; + +/** + * Set the global Kitty keyboard protocol state. + * Called by ProcessTerminal after detecting protocol support. + */ +export function setKittyProtocolActive(active: boolean): void { + _kittyProtocolActive = active; +} + +/** + * Query whether Kitty keyboard protocol is currently active. + */ +export function isKittyProtocolActive(): boolean { + return _kittyProtocolActive; +} + +// ============================================================================= +// Type-Safe Key Identifiers +// ============================================================================= + +type Letter = + | "a" + | "b" + | "c" + | "d" + | "e" + | "f" + | "g" + | "h" + | "i" + | "j" + | "k" + | "l" + | "m" + | "n" + | "o" + | "p" + | "q" + | "r" + | "s" + | "t" + | "u" + | "v" + | "w" + | "x" + | "y" + | "z"; + +type Digit = "0" | "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9"; + +type SymbolKey = + | "`" + | "-" + | "=" + | "[" + | "]" + | "\\" + | ";" + | "'" + | "," + | "." + | "/" + | "!" + | "@" + | "#" + | "$" + | "%" + | "^" + | "&" + | "*" + | "(" + | ")" + | "_" + | "+" + | "|" + | "~" + | "{" + | "}" + | ":" + | "<" + | ">" + | "?"; + +type SpecialKey = + | "escape" + | "esc" + | "enter" + | "return" + | "tab" + | "space" + | "backspace" + | "delete" + | "insert" + | "clear" + | "home" + | "end" + | "pageUp" + | "pageDown" + | "up" + | "down" + | "left" + | "right" + | "f1" + | "f2" + | "f3" + | "f4" + | "f5" + | "f6" + | "f7" + | "f8" + | "f9" + | "f10" + | "f11" + | "f12"; + +type BaseKey = Letter | Digit | SymbolKey | SpecialKey; +type ModifierName = "ctrl" | "shift" | "alt" | "super"; + +type ModifiedKeyId = { + [M in RemainingModifiers]: `${M}+${Key}` | `${M}+${ModifiedKeyId>}`; +}[RemainingModifiers]; + +/** + * Union type of all valid key identifiers. + * Provides autocomplete and catches typos at compile time. + */ +export type KeyId = BaseKey | ModifiedKeyId; + +/** + * Helper object for creating typed key identifiers with autocomplete. + * + * Usage: + * - Key.escape, Key.enter, Key.tab, etc. for special keys + * - Key.backtick, Key.comma, Key.period, etc. for symbol keys + * - Key.ctrl("c"), Key.alt("x"), Key.super("k") for single modifiers + * - Key.ctrlShift("p"), Key.ctrlAlt("x"), Key.ctrlSuper("k") for combined modifiers + */ +export const Key = { + // Special keys + escape: "escape" as const, + esc: "esc" as const, + enter: "enter" as const, + return: "return" as const, + tab: "tab" as const, + space: "space" as const, + backspace: "backspace" as const, + delete: "delete" as const, + insert: "insert" as const, + clear: "clear" as const, + home: "home" as const, + end: "end" as const, + pageUp: "pageUp" as const, + pageDown: "pageDown" as const, + up: "up" as const, + down: "down" as const, + left: "left" as const, + right: "right" as const, + f1: "f1" as const, + f2: "f2" as const, + f3: "f3" as const, + f4: "f4" as const, + f5: "f5" as const, + f6: "f6" as const, + f7: "f7" as const, + f8: "f8" as const, + f9: "f9" as const, + f10: "f10" as const, + f11: "f11" as const, + f12: "f12" as const, + + // Symbol keys + backtick: "`" as const, + hyphen: "-" as const, + equals: "=" as const, + leftbracket: "[" as const, + rightbracket: "]" as const, + backslash: "\\" as const, + semicolon: ";" as const, + quote: "'" as const, + comma: "," as const, + period: "." as const, + slash: "/" as const, + exclamation: "!" as const, + at: "@" as const, + hash: "#" as const, + dollar: "$" as const, + percent: "%" as const, + caret: "^" as const, + ampersand: "&" as const, + asterisk: "*" as const, + leftparen: "(" as const, + rightparen: ")" as const, + underscore: "_" as const, + plus: "+" as const, + pipe: "|" as const, + tilde: "~" as const, + leftbrace: "{" as const, + rightbrace: "}" as const, + colon: ":" as const, + lessthan: "<" as const, + greaterthan: ">" as const, + question: "?" as const, + + // Single modifiers + ctrl: (key: K): `ctrl+${K}` => `ctrl+${key}`, + shift: (key: K): `shift+${K}` => `shift+${key}`, + alt: (key: K): `alt+${K}` => `alt+${key}`, + super: (key: K): `super+${K}` => `super+${key}`, + + // Combined modifiers + ctrlShift: (key: K): `ctrl+shift+${K}` => `ctrl+shift+${key}`, + shiftCtrl: (key: K): `shift+ctrl+${K}` => `shift+ctrl+${key}`, + ctrlAlt: (key: K): `ctrl+alt+${K}` => `ctrl+alt+${key}`, + altCtrl: (key: K): `alt+ctrl+${K}` => `alt+ctrl+${key}`, + shiftAlt: (key: K): `shift+alt+${K}` => `shift+alt+${key}`, + altShift: (key: K): `alt+shift+${K}` => `alt+shift+${key}`, + ctrlSuper: (key: K): `ctrl+super+${K}` => `ctrl+super+${key}`, + superCtrl: (key: K): `super+ctrl+${K}` => `super+ctrl+${key}`, + shiftSuper: (key: K): `shift+super+${K}` => `shift+super+${key}`, + superShift: (key: K): `super+shift+${K}` => `super+shift+${key}`, + altSuper: (key: K): `alt+super+${K}` => `alt+super+${key}`, + superAlt: (key: K): `super+alt+${K}` => `super+alt+${key}`, + + // Triple modifiers + ctrlShiftAlt: (key: K): `ctrl+shift+alt+${K}` => `ctrl+shift+alt+${key}`, + ctrlShiftSuper: (key: K): `ctrl+shift+super+${K}` => `ctrl+shift+super+${key}`, +} as const; + +// ============================================================================= +// Constants +// ============================================================================= + +const SYMBOL_KEYS = new Set([ + "`", + "-", + "=", + "[", + "]", + "\\", + ";", + "'", + ",", + ".", + "/", + "!", + "@", + "#", + "$", + "%", + "^", + "&", + "*", + "(", + ")", + "_", + "+", + "|", + "~", + "{", + "}", + ":", + "<", + ">", + "?", +]); + +const MODIFIERS = { + shift: 1, + alt: 2, + ctrl: 4, + super: 8, +} as const; + +const LOCK_MASK = 64 + 128; // Caps Lock + Num Lock + +const CODEPOINTS = { + escape: 27, + tab: 9, + enter: 13, + space: 32, + backspace: 127, + kpEnter: 57414, // Numpad Enter (Kitty protocol) +} as const; + +const ARROW_CODEPOINTS = { + up: -1, + down: -2, + right: -3, + left: -4, +} as const; + +const FUNCTIONAL_CODEPOINTS = { + delete: -10, + insert: -11, + pageUp: -12, + pageDown: -13, + home: -14, + end: -15, +} as const; + +const KITTY_FUNCTIONAL_KEY_EQUIVALENTS = new Map([ + [57399, 48], // KP_0 -> 0 + [57400, 49], // KP_1 -> 1 + [57401, 50], // KP_2 -> 2 + [57402, 51], // KP_3 -> 3 + [57403, 52], // KP_4 -> 4 + [57404, 53], // KP_5 -> 5 + [57405, 54], // KP_6 -> 6 + [57406, 55], // KP_7 -> 7 + [57407, 56], // KP_8 -> 8 + [57408, 57], // KP_9 -> 9 + [57409, 46], // KP_DECIMAL -> . + [57410, 47], // KP_DIVIDE -> / + [57411, 42], // KP_MULTIPLY -> * + [57412, 45], // KP_SUBTRACT -> - + [57413, 43], // KP_ADD -> + + [57415, 61], // KP_EQUAL -> = + [57416, 44], // KP_SEPARATOR -> , + [57417, ARROW_CODEPOINTS.left], + [57418, ARROW_CODEPOINTS.right], + [57419, ARROW_CODEPOINTS.up], + [57420, ARROW_CODEPOINTS.down], + [57421, FUNCTIONAL_CODEPOINTS.pageUp], + [57422, FUNCTIONAL_CODEPOINTS.pageDown], + [57423, FUNCTIONAL_CODEPOINTS.home], + [57424, FUNCTIONAL_CODEPOINTS.end], + [57425, FUNCTIONAL_CODEPOINTS.insert], + [57426, FUNCTIONAL_CODEPOINTS.delete], +]); + +function normalizeKittyFunctionalCodepoint(codepoint: number): number { + return KITTY_FUNCTIONAL_KEY_EQUIVALENTS.get(codepoint) ?? codepoint; +} + +function normalizeShiftedLetterIdentityCodepoint(codepoint: number, modifier: number): number { + const effectiveModifier = modifier & ~LOCK_MASK; + if ((effectiveModifier & MODIFIERS.shift) !== 0 && codepoint >= 65 && codepoint <= 90) { + return codepoint + 32; + } + return codepoint; +} + +const LEGACY_KEY_SEQUENCES = { + up: ["\x1b[A", "\x1bOA"], + down: ["\x1b[B", "\x1bOB"], + right: ["\x1b[C", "\x1bOC"], + left: ["\x1b[D", "\x1bOD"], + home: ["\x1b[H", "\x1bOH", "\x1b[1~", "\x1b[7~"], + end: ["\x1b[F", "\x1bOF", "\x1b[4~", "\x1b[8~"], + insert: ["\x1b[2~"], + delete: ["\x1b[3~"], + pageUp: ["\x1b[5~", "\x1b[[5~"], + pageDown: ["\x1b[6~", "\x1b[[6~"], + clear: ["\x1b[E", "\x1bOE"], + f1: ["\x1bOP", "\x1b[11~", "\x1b[[A"], + f2: ["\x1bOQ", "\x1b[12~", "\x1b[[B"], + f3: ["\x1bOR", "\x1b[13~", "\x1b[[C"], + f4: ["\x1bOS", "\x1b[14~", "\x1b[[D"], + f5: ["\x1b[15~", "\x1b[[E"], + f6: ["\x1b[17~"], + f7: ["\x1b[18~"], + f8: ["\x1b[19~"], + f9: ["\x1b[20~"], + f10: ["\x1b[21~"], + f11: ["\x1b[23~"], + f12: ["\x1b[24~"], +} as const; + +const LEGACY_SHIFT_SEQUENCES = { + up: ["\x1b[a"], + down: ["\x1b[b"], + right: ["\x1b[c"], + left: ["\x1b[d"], + clear: ["\x1b[e"], + insert: ["\x1b[2$"], + delete: ["\x1b[3$"], + pageUp: ["\x1b[5$"], + pageDown: ["\x1b[6$"], + home: ["\x1b[7$"], + end: ["\x1b[8$"], +} as const; + +const LEGACY_CTRL_SEQUENCES = { + up: ["\x1bOa"], + down: ["\x1bOb"], + right: ["\x1bOc"], + left: ["\x1bOd"], + clear: ["\x1bOe"], + insert: ["\x1b[2^"], + delete: ["\x1b[3^"], + pageUp: ["\x1b[5^"], + pageDown: ["\x1b[6^"], + home: ["\x1b[7^"], + end: ["\x1b[8^"], +} as const; + +const LEGACY_SEQUENCE_KEY_IDS: Record = { + "\x1bOA": "up", + "\x1bOB": "down", + "\x1bOC": "right", + "\x1bOD": "left", + "\x1bOH": "home", + "\x1bOF": "end", + "\x1b[E": "clear", + "\x1bOE": "clear", + "\x1bOe": "ctrl+clear", + "\x1b[e": "shift+clear", + "\x1b[2~": "insert", + "\x1b[2$": "shift+insert", + "\x1b[2^": "ctrl+insert", + "\x1b[3$": "shift+delete", + "\x1b[3^": "ctrl+delete", + "\x1b[[5~": "pageUp", + "\x1b[[6~": "pageDown", + "\x1b[a": "shift+up", + "\x1b[b": "shift+down", + "\x1b[c": "shift+right", + "\x1b[d": "shift+left", + "\x1bOa": "ctrl+up", + "\x1bOb": "ctrl+down", + "\x1bOc": "ctrl+right", + "\x1bOd": "ctrl+left", + "\x1b[5$": "shift+pageUp", + "\x1b[6$": "shift+pageDown", + "\x1b[7$": "shift+home", + "\x1b[8$": "shift+end", + "\x1b[5^": "ctrl+pageUp", + "\x1b[6^": "ctrl+pageDown", + "\x1b[7^": "ctrl+home", + "\x1b[8^": "ctrl+end", + "\x1bOP": "f1", + "\x1bOQ": "f2", + "\x1bOR": "f3", + "\x1bOS": "f4", + "\x1b[11~": "f1", + "\x1b[12~": "f2", + "\x1b[13~": "f3", + "\x1b[14~": "f4", + "\x1b[[A": "f1", + "\x1b[[B": "f2", + "\x1b[[C": "f3", + "\x1b[[D": "f4", + "\x1b[[E": "f5", + "\x1b[15~": "f5", + "\x1b[17~": "f6", + "\x1b[18~": "f7", + "\x1b[19~": "f8", + "\x1b[20~": "f9", + "\x1b[21~": "f10", + "\x1b[23~": "f11", + "\x1b[24~": "f12", + "\x1bb": "alt+left", + "\x1bf": "alt+right", + "\x1bp": "alt+up", + "\x1bn": "alt+down", +} as const; + +type LegacyModifierKey = keyof typeof LEGACY_SHIFT_SEQUENCES; + +const matchesLegacySequence = (data: string, sequences: readonly string[]): boolean => sequences.includes(data); + +const matchesLegacyModifierSequence = (data: string, key: LegacyModifierKey, modifier: number): boolean => { + if (modifier === MODIFIERS.shift) { + return matchesLegacySequence(data, LEGACY_SHIFT_SEQUENCES[key]); + } + if (modifier === MODIFIERS.ctrl) { + return matchesLegacySequence(data, LEGACY_CTRL_SEQUENCES[key]); + } + return false; +}; + +// ============================================================================= +// Kitty Protocol Parsing +// ============================================================================= + +/** + * Event types from Kitty keyboard protocol (flag 2) + * 1 = key press, 2 = key repeat, 3 = key release + */ +export type KeyEventType = "press" | "repeat" | "release"; + +interface ParsedKittySequence { + codepoint: number; + shiftedKey?: number; // Shifted version of the key (when shift is pressed) + baseLayoutKey?: number; // Key in standard PC-101 layout (for non-Latin layouts) + modifier: number; + eventType: KeyEventType; +} + +interface ParsedModifyOtherKeysSequence { + codepoint: number; + modifier: number; +} + +// Store the last parsed event type for isKeyRelease() to query +let _lastEventType: KeyEventType = "press"; + +/** + * Check if the last parsed key event was a key release. + * Only meaningful when Kitty keyboard protocol with flag 2 is active. + */ +export function isKeyRelease(data: string): boolean { + // Don't treat bracketed paste content as key release, even if it contains + // patterns like ":3F" (e.g., bluetooth MAC addresses like "90:62:3F:A5"). + // Terminal.ts re-wraps paste content with bracketed paste markers before + // passing to TUI, so pasted data will always contain \x1b[200~. + if (data.includes("\x1b[200~")) { + return false; + } + + // Quick check: release events with flag 2 contain ":3" + // Format: \x1b[;:3u + if ( + data.includes(":3u") || + data.includes(":3~") || + data.includes(":3A") || + data.includes(":3B") || + data.includes(":3C") || + data.includes(":3D") || + data.includes(":3H") || + data.includes(":3F") + ) { + return true; + } + return false; +} + +/** + * Check if the last parsed key event was a key repeat. + * Only meaningful when Kitty keyboard protocol with flag 2 is active. + */ +export function isKeyRepeat(data: string): boolean { + // Don't treat bracketed paste content as key repeat, even if it contains + // patterns like ":2F". See isKeyRelease() for details. + if (data.includes("\x1b[200~")) { + return false; + } + + if ( + data.includes(":2u") || + data.includes(":2~") || + data.includes(":2A") || + data.includes(":2B") || + data.includes(":2C") || + data.includes(":2D") || + data.includes(":2H") || + data.includes(":2F") + ) { + return true; + } + return false; +} + +function parseEventType(eventTypeStr: string | undefined): KeyEventType { + if (!eventTypeStr) return "press"; + const eventType = parseInt(eventTypeStr, 10); + if (eventType === 2) return "repeat"; + if (eventType === 3) return "release"; + return "press"; +} + +function parseKittySequence(data: string): ParsedKittySequence | null { + // CSI u format with alternate keys (flag 4): + // \x1b[u + // \x1b[;u + // \x1b[;:u + // \x1b[:;u + // \x1b[::;u + // \x1b[::;u (no shifted key, only base) + // + // With flag 2, event type is appended after modifier colon: 1=press, 2=repeat, 3=release + // With flag 4, alternate keys are appended after codepoint with colons + const csiUMatch = data.match(/^\x1b\[(\d+)(?::(\d*))?(?::(\d+))?(?:;(\d+))?(?::(\d+))?u$/); + if (csiUMatch) { + const codepoint = parseInt(csiUMatch[1]!, 10); + const shiftedKey = csiUMatch[2] && csiUMatch[2].length > 0 ? parseInt(csiUMatch[2], 10) : undefined; + const baseLayoutKey = csiUMatch[3] ? parseInt(csiUMatch[3], 10) : undefined; + const modValue = csiUMatch[4] ? parseInt(csiUMatch[4], 10) : 1; + const eventType = parseEventType(csiUMatch[5]); + _lastEventType = eventType; + return { codepoint, shiftedKey, baseLayoutKey, modifier: modValue - 1, eventType }; + } + + // Arrow keys with modifier: \x1b[1;A/B/C/D or \x1b[1;:A/B/C/D + const arrowMatch = data.match(/^\x1b\[1;(\d+)(?::(\d+))?([ABCD])$/); + if (arrowMatch) { + const modValue = parseInt(arrowMatch[1]!, 10); + const eventType = parseEventType(arrowMatch[2]); + const arrowCodes: Record = { A: -1, B: -2, C: -3, D: -4 }; + _lastEventType = eventType; + return { codepoint: arrowCodes[arrowMatch[3]!]!, modifier: modValue - 1, eventType }; + } + + // Functional keys: \x1b[~ or \x1b[;~ or \x1b[;:~ + const funcMatch = data.match(/^\x1b\[(\d+)(?:;(\d+))?(?::(\d+))?~$/); + if (funcMatch) { + const keyNum = parseInt(funcMatch[1]!, 10); + const modValue = funcMatch[2] ? parseInt(funcMatch[2], 10) : 1; + const eventType = parseEventType(funcMatch[3]); + const funcCodes: Record = { + 2: FUNCTIONAL_CODEPOINTS.insert, + 3: FUNCTIONAL_CODEPOINTS.delete, + 5: FUNCTIONAL_CODEPOINTS.pageUp, + 6: FUNCTIONAL_CODEPOINTS.pageDown, + 7: FUNCTIONAL_CODEPOINTS.home, + 8: FUNCTIONAL_CODEPOINTS.end, + }; + const codepoint = funcCodes[keyNum]; + if (codepoint !== undefined) { + _lastEventType = eventType; + return { codepoint, modifier: modValue - 1, eventType }; + } + } + + // Home/End with modifier: \x1b[1;H/F or \x1b[1;:H/F + const homeEndMatch = data.match(/^\x1b\[1;(\d+)(?::(\d+))?([HF])$/); + if (homeEndMatch) { + const modValue = parseInt(homeEndMatch[1]!, 10); + const eventType = parseEventType(homeEndMatch[2]); + const codepoint = homeEndMatch[3] === "H" ? FUNCTIONAL_CODEPOINTS.home : FUNCTIONAL_CODEPOINTS.end; + _lastEventType = eventType; + return { codepoint, modifier: modValue - 1, eventType }; + } + + return null; +} + +function matchesKittySequence(data: string, expectedCodepoint: number, expectedModifier: number): boolean { + const parsed = parseKittySequence(data); + if (!parsed) return false; + const actualMod = parsed.modifier & ~LOCK_MASK; + const expectedMod = expectedModifier & ~LOCK_MASK; + + // Check if modifiers match + if (actualMod !== expectedMod) return false; + + const normalizedCodepoint = normalizeShiftedLetterIdentityCodepoint( + normalizeKittyFunctionalCodepoint(parsed.codepoint), + parsed.modifier, + ); + const normalizedExpectedCodepoint = normalizeShiftedLetterIdentityCodepoint( + normalizeKittyFunctionalCodepoint(expectedCodepoint), + expectedModifier, + ); + + // Primary match: codepoint matches directly after normalizing functional keys + if (normalizedCodepoint === normalizedExpectedCodepoint) return true; + + // Alternate match: use base layout key for non-Latin keyboard layouts. + // This allows Ctrl+С (Cyrillic) to match Ctrl+c (Latin) when terminal reports + // the base layout key (the key in standard PC-101 layout). + // + // Only fall back to base layout key when the codepoint is NOT already a + // recognized Latin letter (a-z) or symbol (e.g., /, -, [, ;, etc.). + // When the codepoint is a recognized key, it is authoritative regardless + // of physical key position. This prevents remapped layouts (Dvorak, Colemak, + // xremap, etc.) from causing false matches: both letters and symbols move + // to different physical positions, so Ctrl+K could falsely match Ctrl+V + // (letter remapping) and Ctrl+/ could falsely match Ctrl+[ (symbol remapping) + // if the base layout key were always considered. + if (parsed.baseLayoutKey !== undefined && parsed.baseLayoutKey === expectedCodepoint) { + const cp = normalizedCodepoint; + const isLatinLetter = cp >= 97 && cp <= 122; // a-z + const isKnownSymbol = SYMBOL_KEYS.has(String.fromCharCode(cp)); + if (!isLatinLetter && !isKnownSymbol) return true; + } + + return false; +} + +function parseModifyOtherKeysSequence(data: string): ParsedModifyOtherKeysSequence | null { + const match = data.match(/^\x1b\[27;(\d+);(\d+)~$/); + if (!match) return null; + const modValue = parseInt(match[1]!, 10); + const codepoint = parseInt(match[2]!, 10); + return { codepoint, modifier: modValue - 1 }; +} + +/** + * Match xterm modifyOtherKeys format: CSI 27 ; modifiers ; keycode ~ + * This is used by terminals when Kitty protocol is not enabled. + * Modifier values are 1-indexed: 2=shift, 3=alt, 5=ctrl, etc. + */ +function matchesModifyOtherKeys(data: string, expectedKeycode: number, expectedModifier: number): boolean { + const parsed = parseModifyOtherKeysSequence(data); + if (!parsed) return false; + return parsed.codepoint === expectedKeycode && parsed.modifier === expectedModifier; +} + +function isWindowsTerminalSession(): boolean { + return ( + Boolean(process.env.WT_SESSION) && !process.env.SSH_CONNECTION && !process.env.SSH_CLIENT && !process.env.SSH_TTY + ); +} + +/** + * Raw 0x08 (BS) is ambiguous in legacy terminals. + * + * - Windows Terminal uses it for Ctrl+Backspace. + * - Some legacy terminals and tmux setups send it for plain Backspace. + * + * Prefer explicit Kitty / CSI-u / modifyOtherKeys sequences whenever they are + * available. Fall back to a Windows Terminal heuristic only for raw BS bytes. + */ +function matchesRawBackspace(data: string, expectedModifier: number): boolean { + if (data === "\x7f") return expectedModifier === 0; + if (data !== "\x08") return false; + return isWindowsTerminalSession() ? expectedModifier === MODIFIERS.ctrl : expectedModifier === 0; +} + +// ============================================================================= +// Generic Key Matching +// ============================================================================= + +/** + * Get the control character for a key. + * Uses the universal formula: code & 0x1f (mask to lower 5 bits) + * + * Works for: + * - Letters a-z → 1-26 + * - Symbols [\]_ → 27, 28, 29, 31 + * - Also maps - to same as _ (same physical key on US keyboards) + */ +function rawCtrlChar(key: string): string | null { + const char = key.toLowerCase(); + const code = char.charCodeAt(0); + if ((code >= 97 && code <= 122) || char === "[" || char === "\\" || char === "]" || char === "_") { + return String.fromCharCode(code & 0x1f); + } + // Handle - as _ (same physical key on US keyboards) + if (char === "-") { + return String.fromCharCode(31); // Same as Ctrl+_ + } + return null; +} + +function isDigitKey(key: string): boolean { + return key >= "0" && key <= "9"; +} + +function matchesPrintableModifyOtherKeys(data: string, expectedKeycode: number, expectedModifier: number): boolean { + if (expectedModifier === 0) return false; + const parsed = parseModifyOtherKeysSequence(data); + if (!parsed || parsed.modifier !== expectedModifier) return false; + return ( + normalizeShiftedLetterIdentityCodepoint(parsed.codepoint, parsed.modifier) === + normalizeShiftedLetterIdentityCodepoint(expectedKeycode, expectedModifier) + ); +} + +function formatKeyNameWithModifiers(keyName: string, modifier: number): string | undefined { + const mods: string[] = []; + const effectiveMod = modifier & ~LOCK_MASK; + const supportedModifierMask = MODIFIERS.shift | MODIFIERS.ctrl | MODIFIERS.alt | MODIFIERS.super; + if ((effectiveMod & ~supportedModifierMask) !== 0) return undefined; + if (effectiveMod & MODIFIERS.shift) mods.push("shift"); + if (effectiveMod & MODIFIERS.ctrl) mods.push("ctrl"); + if (effectiveMod & MODIFIERS.alt) mods.push("alt"); + if (effectiveMod & MODIFIERS.super) mods.push("super"); + return mods.length > 0 ? `${mods.join("+")}+${keyName}` : keyName; +} + +function parseKeyId( + keyId: string, +): { key: string; ctrl: boolean; shift: boolean; alt: boolean; super: boolean } | null { + const parts = keyId.toLowerCase().split("+"); + const key = parts[parts.length - 1]; + if (!key) return null; + return { + key, + ctrl: parts.includes("ctrl"), + shift: parts.includes("shift"), + alt: parts.includes("alt"), + super: parts.includes("super"), + }; +} + +/** + * Match input data against a key identifier string. + * + * Supported key identifiers: + * - Single keys: "escape", "tab", "enter", "backspace", "delete", "home", "end", "space" + * - Arrow keys: "up", "down", "left", "right" + * - Ctrl combinations: "ctrl+c", "ctrl+z", etc. + * - Shift combinations: "shift+tab", "shift+enter" + * - Alt combinations: "alt+enter", "alt+backspace" + * - Super combinations: "super+k", "super+enter" + * - Combined modifiers: "shift+ctrl+p", "ctrl+alt+x", "ctrl+super+k" + * + * Use the Key helper for autocomplete: Key.ctrl("c"), Key.escape, Key.ctrlShift("p"), Key.super("k") + * + * @param data - Raw input data from terminal + * @param keyId - Key identifier (e.g., "ctrl+c", "escape", Key.ctrl("c")) + */ +export function matchesKey(data: string, keyId: KeyId): boolean { + const parsed = parseKeyId(keyId); + if (!parsed) return false; + + const { key, ctrl, shift, alt, super: superModifier } = parsed; + let modifier = 0; + if (shift) modifier |= MODIFIERS.shift; + if (alt) modifier |= MODIFIERS.alt; + if (ctrl) modifier |= MODIFIERS.ctrl; + if (superModifier) modifier |= MODIFIERS.super; + + switch (key) { + case "escape": + case "esc": + if (modifier !== 0) return false; + return ( + data === "\x1b" || + matchesKittySequence(data, CODEPOINTS.escape, 0) || + matchesModifyOtherKeys(data, CODEPOINTS.escape, 0) + ); + + case "space": + if (!_kittyProtocolActive) { + if (modifier === MODIFIERS.ctrl && data === "\x00") { + return true; + } + if (modifier === MODIFIERS.alt && data === "\x1b ") { + return true; + } + } + if (modifier === 0) { + return ( + data === " " || + matchesKittySequence(data, CODEPOINTS.space, 0) || + matchesModifyOtherKeys(data, CODEPOINTS.space, 0) + ); + } + return ( + matchesKittySequence(data, CODEPOINTS.space, modifier) || + matchesModifyOtherKeys(data, CODEPOINTS.space, modifier) + ); + + case "tab": + if (modifier === MODIFIERS.shift) { + return ( + data === "\x1b[Z" || + matchesKittySequence(data, CODEPOINTS.tab, MODIFIERS.shift) || + matchesModifyOtherKeys(data, CODEPOINTS.tab, MODIFIERS.shift) + ); + } + if (modifier === 0) { + return data === "\t" || matchesKittySequence(data, CODEPOINTS.tab, 0); + } + return ( + matchesKittySequence(data, CODEPOINTS.tab, modifier) || + matchesModifyOtherKeys(data, CODEPOINTS.tab, modifier) + ); + + case "enter": + case "return": + if (modifier === MODIFIERS.shift) { + // CSI u sequences (standard Kitty protocol) + if ( + matchesKittySequence(data, CODEPOINTS.enter, MODIFIERS.shift) || + matchesKittySequence(data, CODEPOINTS.kpEnter, MODIFIERS.shift) + ) { + return true; + } + // xterm modifyOtherKeys format (fallback when Kitty protocol not enabled) + if (matchesModifyOtherKeys(data, CODEPOINTS.enter, MODIFIERS.shift)) { + return true; + } + // When Kitty protocol is active, legacy sequences are custom terminal mappings + // \x1b\r = Kitty's "map shift+enter send_text all \e\r" + // \n = Ghostty's "keybind = shift+enter=text:\n" + if (_kittyProtocolActive) { + return data === "\x1b\r" || data === "\n"; + } + return false; + } + if (modifier === MODIFIERS.alt) { + // CSI u sequences (standard Kitty protocol) + if ( + matchesKittySequence(data, CODEPOINTS.enter, MODIFIERS.alt) || + matchesKittySequence(data, CODEPOINTS.kpEnter, MODIFIERS.alt) + ) { + return true; + } + // xterm modifyOtherKeys format (fallback when Kitty protocol not enabled) + if (matchesModifyOtherKeys(data, CODEPOINTS.enter, MODIFIERS.alt)) { + return true; + } + // \x1b\r is alt+enter only in legacy mode (no Kitty protocol) + // When Kitty protocol is active, alt+enter comes as CSI u sequence + if (!_kittyProtocolActive) { + return data === "\x1b\r"; + } + return false; + } + if (modifier === 0) { + return ( + data === "\r" || + (!_kittyProtocolActive && data === "\n") || + data === "\x1bOM" || // SS3 M (numpad enter in some terminals) + matchesKittySequence(data, CODEPOINTS.enter, 0) || + matchesKittySequence(data, CODEPOINTS.kpEnter, 0) + ); + } + return ( + matchesKittySequence(data, CODEPOINTS.enter, modifier) || + matchesKittySequence(data, CODEPOINTS.kpEnter, modifier) || + matchesModifyOtherKeys(data, CODEPOINTS.enter, modifier) + ); + + case "backspace": + if (modifier === MODIFIERS.alt) { + if (data === "\x1b\x7f" || data === "\x1b\b") { + return true; + } + return ( + matchesKittySequence(data, CODEPOINTS.backspace, MODIFIERS.alt) || + matchesModifyOtherKeys(data, CODEPOINTS.backspace, MODIFIERS.alt) + ); + } + if (modifier === MODIFIERS.ctrl) { + // Legacy raw 0x08 is ambiguous: it can be Ctrl+Backspace on Windows + // Terminal or plain Backspace on other terminals, while also + // overlapping with Ctrl+H. + if (matchesRawBackspace(data, MODIFIERS.ctrl)) return true; + return ( + matchesKittySequence(data, CODEPOINTS.backspace, MODIFIERS.ctrl) || + matchesModifyOtherKeys(data, CODEPOINTS.backspace, MODIFIERS.ctrl) + ); + } + if (modifier === 0) { + return ( + matchesRawBackspace(data, 0) || + matchesKittySequence(data, CODEPOINTS.backspace, 0) || + matchesModifyOtherKeys(data, CODEPOINTS.backspace, 0) + ); + } + return ( + matchesKittySequence(data, CODEPOINTS.backspace, modifier) || + matchesModifyOtherKeys(data, CODEPOINTS.backspace, modifier) + ); + + case "insert": + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.insert) || + matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.insert, 0) + ); + } + if (matchesLegacyModifierSequence(data, "insert", modifier)) { + return true; + } + return matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.insert, modifier); + + case "delete": + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.delete) || + matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.delete, 0) + ); + } + if (matchesLegacyModifierSequence(data, "delete", modifier)) { + return true; + } + return matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.delete, modifier); + + case "clear": + if (modifier === 0) { + return matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.clear); + } + return matchesLegacyModifierSequence(data, "clear", modifier); + + case "home": + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.home) || + matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.home, 0) + ); + } + if (matchesLegacyModifierSequence(data, "home", modifier)) { + return true; + } + return matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.home, modifier); + + case "end": + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.end) || + matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.end, 0) + ); + } + if (matchesLegacyModifierSequence(data, "end", modifier)) { + return true; + } + return matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.end, modifier); + + case "pageup": + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.pageUp) || + matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.pageUp, 0) + ); + } + if (matchesLegacyModifierSequence(data, "pageUp", modifier)) { + return true; + } + return matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.pageUp, modifier); + + case "pagedown": + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.pageDown) || + matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.pageDown, 0) + ); + } + if (matchesLegacyModifierSequence(data, "pageDown", modifier)) { + return true; + } + return matchesKittySequence(data, FUNCTIONAL_CODEPOINTS.pageDown, modifier); + + case "up": + if (modifier === MODIFIERS.alt) { + return data === "\x1bp" || matchesKittySequence(data, ARROW_CODEPOINTS.up, MODIFIERS.alt); + } + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.up) || + matchesKittySequence(data, ARROW_CODEPOINTS.up, 0) + ); + } + if (matchesLegacyModifierSequence(data, "up", modifier)) { + return true; + } + return matchesKittySequence(data, ARROW_CODEPOINTS.up, modifier); + + case "down": + if (modifier === MODIFIERS.alt) { + return data === "\x1bn" || matchesKittySequence(data, ARROW_CODEPOINTS.down, MODIFIERS.alt); + } + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.down) || + matchesKittySequence(data, ARROW_CODEPOINTS.down, 0) + ); + } + if (matchesLegacyModifierSequence(data, "down", modifier)) { + return true; + } + return matchesKittySequence(data, ARROW_CODEPOINTS.down, modifier); + + case "left": + if (modifier === MODIFIERS.alt) { + return ( + data === "\x1b[1;3D" || + (!_kittyProtocolActive && data === "\x1bB") || + data === "\x1bb" || + matchesKittySequence(data, ARROW_CODEPOINTS.left, MODIFIERS.alt) + ); + } + if (modifier === MODIFIERS.ctrl) { + return ( + data === "\x1b[1;5D" || + matchesLegacyModifierSequence(data, "left", MODIFIERS.ctrl) || + matchesKittySequence(data, ARROW_CODEPOINTS.left, MODIFIERS.ctrl) + ); + } + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.left) || + matchesKittySequence(data, ARROW_CODEPOINTS.left, 0) + ); + } + if (matchesLegacyModifierSequence(data, "left", modifier)) { + return true; + } + return matchesKittySequence(data, ARROW_CODEPOINTS.left, modifier); + + case "right": + if (modifier === MODIFIERS.alt) { + return ( + data === "\x1b[1;3C" || + (!_kittyProtocolActive && data === "\x1bF") || + data === "\x1bf" || + matchesKittySequence(data, ARROW_CODEPOINTS.right, MODIFIERS.alt) + ); + } + if (modifier === MODIFIERS.ctrl) { + return ( + data === "\x1b[1;5C" || + matchesLegacyModifierSequence(data, "right", MODIFIERS.ctrl) || + matchesKittySequence(data, ARROW_CODEPOINTS.right, MODIFIERS.ctrl) + ); + } + if (modifier === 0) { + return ( + matchesLegacySequence(data, LEGACY_KEY_SEQUENCES.right) || + matchesKittySequence(data, ARROW_CODEPOINTS.right, 0) + ); + } + if (matchesLegacyModifierSequence(data, "right", modifier)) { + return true; + } + return matchesKittySequence(data, ARROW_CODEPOINTS.right, modifier); + + case "f1": + case "f2": + case "f3": + case "f4": + case "f5": + case "f6": + case "f7": + case "f8": + case "f9": + case "f10": + case "f11": + case "f12": { + if (modifier !== 0) { + return false; + } + const functionKey = key as keyof typeof LEGACY_KEY_SEQUENCES; + return matchesLegacySequence(data, LEGACY_KEY_SEQUENCES[functionKey]); + } + } + + // Handle single letter/digit keys and symbols + if (key.length === 1 && ((key >= "a" && key <= "z") || isDigitKey(key) || SYMBOL_KEYS.has(key))) { + const codepoint = key.charCodeAt(0); + const rawCtrl = rawCtrlChar(key); + const isLetter = key >= "a" && key <= "z"; + const isDigit = isDigitKey(key); + + if (modifier === MODIFIERS.ctrl + MODIFIERS.alt && !_kittyProtocolActive && rawCtrl) { + // Legacy: ctrl+alt+key is ESC followed by the control character. + // If that legacy form does not match, continue so CSI-u and + // modifyOtherKeys sequences from tmux can still be recognized. + if (data === `\x1b${rawCtrl}`) return true; + } + + if (modifier === MODIFIERS.alt && !_kittyProtocolActive && (isLetter || isDigit)) { + // Legacy: alt+letter/digit is ESC followed by the key + if (data === `\x1b${key}`) return true; + } + + if (modifier === MODIFIERS.ctrl) { + // Legacy: ctrl+key sends the control character + if (rawCtrl && data === rawCtrl) return true; + return ( + matchesKittySequence(data, codepoint, MODIFIERS.ctrl) || + matchesPrintableModifyOtherKeys(data, codepoint, MODIFIERS.ctrl) + ); + } + + if (modifier === MODIFIERS.shift + MODIFIERS.ctrl) { + return ( + matchesKittySequence(data, codepoint, MODIFIERS.shift + MODIFIERS.ctrl) || + matchesPrintableModifyOtherKeys(data, codepoint, MODIFIERS.shift + MODIFIERS.ctrl) + ); + } + + if (modifier === MODIFIERS.shift) { + // Legacy: shift+letter produces uppercase + if (isLetter && data === key.toUpperCase()) return true; + return ( + matchesKittySequence(data, codepoint, MODIFIERS.shift) || + matchesPrintableModifyOtherKeys(data, codepoint, MODIFIERS.shift) + ); + } + + if (modifier !== 0) { + return ( + matchesKittySequence(data, codepoint, modifier) || + matchesPrintableModifyOtherKeys(data, codepoint, modifier) + ); + } + + // Check both raw char and Kitty sequence (needed for release events) + return data === key || matchesKittySequence(data, codepoint, 0); + } + + return false; +} + +/** + * Parse input data and return the key identifier if recognized. + * + * @param data - Raw input data from terminal + * @returns Key identifier string (e.g., "ctrl+c") or undefined + */ +function formatParsedKey(codepoint: number, modifier: number, baseLayoutKey?: number): string | undefined { + const normalizedCodepoint = normalizeKittyFunctionalCodepoint(codepoint); + const identityCodepoint = normalizeShiftedLetterIdentityCodepoint(normalizedCodepoint, modifier); + + // Use base layout key only when codepoint is not a recognized Latin + // letter (a-z), digit (0-9), or symbol (/, -, [, ;, etc.). For those, + // the codepoint is authoritative regardless of physical key position. + // This prevents remapped layouts (Dvorak, Colemak, xremap, etc.) from + // reporting the wrong key name based on the QWERTY physical position. + const isLatinLetter = identityCodepoint >= 97 && identityCodepoint <= 122; // a-z + const isDigit = identityCodepoint >= 48 && identityCodepoint <= 57; // 0-9 + const isKnownSymbol = SYMBOL_KEYS.has(String.fromCharCode(identityCodepoint)); + const effectiveCodepoint = + isLatinLetter || isDigit || isKnownSymbol ? identityCodepoint : (baseLayoutKey ?? identityCodepoint); + + let keyName: string | undefined; + if (effectiveCodepoint === CODEPOINTS.escape) keyName = "escape"; + else if (effectiveCodepoint === CODEPOINTS.tab) keyName = "tab"; + else if (effectiveCodepoint === CODEPOINTS.enter || effectiveCodepoint === CODEPOINTS.kpEnter) keyName = "enter"; + else if (effectiveCodepoint === CODEPOINTS.space) keyName = "space"; + else if (effectiveCodepoint === CODEPOINTS.backspace) keyName = "backspace"; + else if (effectiveCodepoint === FUNCTIONAL_CODEPOINTS.delete) keyName = "delete"; + else if (effectiveCodepoint === FUNCTIONAL_CODEPOINTS.insert) keyName = "insert"; + else if (effectiveCodepoint === FUNCTIONAL_CODEPOINTS.home) keyName = "home"; + else if (effectiveCodepoint === FUNCTIONAL_CODEPOINTS.end) keyName = "end"; + else if (effectiveCodepoint === FUNCTIONAL_CODEPOINTS.pageUp) keyName = "pageUp"; + else if (effectiveCodepoint === FUNCTIONAL_CODEPOINTS.pageDown) keyName = "pageDown"; + else if (effectiveCodepoint === ARROW_CODEPOINTS.up) keyName = "up"; + else if (effectiveCodepoint === ARROW_CODEPOINTS.down) keyName = "down"; + else if (effectiveCodepoint === ARROW_CODEPOINTS.left) keyName = "left"; + else if (effectiveCodepoint === ARROW_CODEPOINTS.right) keyName = "right"; + else if (effectiveCodepoint >= 48 && effectiveCodepoint <= 57) keyName = String.fromCharCode(effectiveCodepoint); + else if (effectiveCodepoint >= 97 && effectiveCodepoint <= 122) keyName = String.fromCharCode(effectiveCodepoint); + else if (SYMBOL_KEYS.has(String.fromCharCode(effectiveCodepoint))) keyName = String.fromCharCode(effectiveCodepoint); + + if (!keyName) return undefined; + return formatKeyNameWithModifiers(keyName, modifier); +} + +export function parseKey(data: string): string | undefined { + const kitty = parseKittySequence(data); + if (kitty) { + return formatParsedKey(kitty.codepoint, kitty.modifier, kitty.baseLayoutKey); + } + + const modifyOtherKeys = parseModifyOtherKeysSequence(data); + if (modifyOtherKeys) { + return formatParsedKey(modifyOtherKeys.codepoint, modifyOtherKeys.modifier); + } + + // Mode-aware legacy sequences + // When Kitty protocol is active, ambiguous sequences are interpreted as custom terminal mappings: + // - \x1b\r = shift+enter (Kitty mapping), not alt+enter + // - \n = shift+enter (Ghostty mapping) + if (_kittyProtocolActive) { + if (data === "\x1b\r" || data === "\n") return "shift+enter"; + } + + const legacySequenceKeyId = LEGACY_SEQUENCE_KEY_IDS[data]; + if (legacySequenceKeyId) return legacySequenceKeyId; + + // Legacy sequences (used when Kitty protocol is not active, or for unambiguous sequences) + if (data === "\x1b") return "escape"; + if (data === "\x1c") return "ctrl+\\"; + if (data === "\x1d") return "ctrl+]"; + if (data === "\x1f") return "ctrl+-"; + if (data === "\x1b\x1b") return "ctrl+alt+["; + if (data === "\x1b\x1c") return "ctrl+alt+\\"; + if (data === "\x1b\x1d") return "ctrl+alt+]"; + if (data === "\x1b\x1f") return "ctrl+alt+-"; + if (data === "\t") return "tab"; + if (data === "\r" || (!_kittyProtocolActive && data === "\n") || data === "\x1bOM") return "enter"; + if (data === "\x00") return "ctrl+space"; + if (data === " ") return "space"; + if (data === "\x7f") return "backspace"; + if (data === "\x08") return isWindowsTerminalSession() ? "ctrl+backspace" : "backspace"; + if (data === "\x1b[Z") return "shift+tab"; + if (!_kittyProtocolActive && data === "\x1b\r") return "alt+enter"; + if (!_kittyProtocolActive && data === "\x1b ") return "alt+space"; + if (data === "\x1b\x7f" || data === "\x1b\b") return "alt+backspace"; + if (!_kittyProtocolActive && data === "\x1bB") return "alt+left"; + if (!_kittyProtocolActive && data === "\x1bF") return "alt+right"; + if (!_kittyProtocolActive && data.length === 2 && data[0] === "\x1b") { + const code = data.charCodeAt(1); + if (code >= 1 && code <= 26) { + return `ctrl+alt+${String.fromCharCode(code + 96)}`; + } + // Legacy alt+letter/digit (ESC followed by the key) + if ((code >= 97 && code <= 122) || (code >= 48 && code <= 57)) { + return `alt+${String.fromCharCode(code)}`; + } + } + if (data === "\x1b[A") return "up"; + if (data === "\x1b[B") return "down"; + if (data === "\x1b[C") return "right"; + if (data === "\x1b[D") return "left"; + if (data === "\x1b[H" || data === "\x1bOH") return "home"; + if (data === "\x1b[F" || data === "\x1bOF") return "end"; + if (data === "\x1b[3~") return "delete"; + if (data === "\x1b[5~") return "pageUp"; + if (data === "\x1b[6~") return "pageDown"; + + // Raw Ctrl+letter + if (data.length === 1) { + const code = data.charCodeAt(0); + if (code >= 1 && code <= 26) { + return `ctrl+${String.fromCharCode(code + 96)}`; + } + if (code >= 32 && code <= 126) { + return data; + } + } + + return undefined; +} + +// ============================================================================= +// Kitty CSI-u Printable Decoding +// ============================================================================= + +const KITTY_CSI_U_REGEX = /^\x1b\[(\d+)(?::(\d*))?(?::(\d+))?(?:;(\d+))?(?::(\d+))?u$/; +const KITTY_PRINTABLE_ALLOWED_MODIFIERS = MODIFIERS.shift | LOCK_MASK; + +/** + * Decode a Kitty CSI-u sequence into a printable character, if applicable. + * + * When Kitty keyboard protocol flag 1 (disambiguate) is active, terminals send + * CSI-u sequences for all keys, including plain printable characters. This + * function extracts the printable character from such sequences. + * + * Only accepts plain or Shift-modified keys. Rejects Ctrl, Alt, and unsupported + * modifier combinations (those are handled by keybinding matching instead). + * Prefers the shifted keycode when Shift is held and a shifted key is reported. + * + * @param data - Raw input data from terminal + * @returns The printable character, or undefined if not a printable CSI-u sequence + */ +export function decodeKittyPrintable(data: string): string | undefined { + const match = data.match(KITTY_CSI_U_REGEX); + if (!match) return undefined; + + // CSI-u groups: [:[:]];[:]u + const codepoint = Number.parseInt(match[1] ?? "", 10); + if (!Number.isFinite(codepoint)) return undefined; + + const shiftedKey = match[2] && match[2].length > 0 ? Number.parseInt(match[2], 10) : undefined; + const modValue = match[4] ? Number.parseInt(match[4], 10) : 1; + // Modifiers are 1-indexed in CSI-u; normalize to our bitmask. + const modifier = Number.isFinite(modValue) ? modValue - 1 : 0; + + // Only accept printable CSI-u input for plain or Shift-modified text keys. + // Reject unsupported modifier bits (e.g. Super/Meta) to avoid inserting + // characters from modifier-only terminal events. + if ((modifier & ~KITTY_PRINTABLE_ALLOWED_MODIFIERS) !== 0) return undefined; + if (modifier & (MODIFIERS.alt | MODIFIERS.ctrl)) return undefined; + + // Prefer the shifted keycode when Shift is held. + let effectiveCodepoint = codepoint; + if (modifier & MODIFIERS.shift && typeof shiftedKey === "number") { + effectiveCodepoint = shiftedKey; + } + effectiveCodepoint = normalizeKittyFunctionalCodepoint(effectiveCodepoint); + // Drop control characters or invalid codepoints. + if (!Number.isFinite(effectiveCodepoint) || effectiveCodepoint < 32) return undefined; + + try { + return String.fromCodePoint(effectiveCodepoint); + } catch { + return undefined; + } +} + +function decodeModifyOtherKeysPrintable(data: string): string | undefined { + const parsed = parseModifyOtherKeysSequence(data); + if (!parsed) return undefined; + const modifier = parsed.modifier & ~LOCK_MASK; + if ((modifier & ~MODIFIERS.shift) !== 0) return undefined; + if (!Number.isFinite(parsed.codepoint) || parsed.codepoint < 32) return undefined; + + try { + return String.fromCodePoint(parsed.codepoint); + } catch { + return undefined; + } +} + +export function decodePrintableKey(data: string): string | undefined { + return decodeKittyPrintable(data) ?? decodeModifyOtherKeysPrintable(data); +} diff --git a/cactus-code/packages/tui/src/kill-ring.ts b/cactus-code/packages/tui/src/kill-ring.ts new file mode 100644 index 000000000..2292f91aa --- /dev/null +++ b/cactus-code/packages/tui/src/kill-ring.ts @@ -0,0 +1,46 @@ +/** + * Ring buffer for Emacs-style kill/yank operations. + * + * Tracks killed (deleted) text entries. Consecutive kills can accumulate + * into a single entry. Supports yank (paste most recent) and yank-pop + * (cycle through older entries). + */ +export class KillRing { + private ring: string[] = []; + + /** + * Add text to the kill ring. + * + * @param text - The killed text to add + * @param opts - Push options + * @param opts.prepend - If accumulating, prepend (backward deletion) or append (forward deletion) + * @param opts.accumulate - Merge with the most recent entry instead of creating a new one + */ + push(text: string, opts: { prepend: boolean; accumulate?: boolean }): void { + if (!text) return; + + if (opts.accumulate && this.ring.length > 0) { + const last = this.ring.pop()!; + this.ring.push(opts.prepend ? text + last : last + text); + } else { + this.ring.push(text); + } + } + + /** Get most recent entry without modifying the ring. */ + peek(): string | undefined { + return this.ring.length > 0 ? this.ring[this.ring.length - 1] : undefined; + } + + /** Move last entry to front (for yank-pop cycling). */ + rotate(): void { + if (this.ring.length > 1) { + const last = this.ring.pop()!; + this.ring.unshift(last); + } + } + + get length(): number { + return this.ring.length; + } +} diff --git a/cactus-code/packages/tui/src/native-modifiers.ts b/cactus-code/packages/tui/src/native-modifiers.ts new file mode 100644 index 000000000..e2cd631cb --- /dev/null +++ b/cactus-code/packages/tui/src/native-modifiers.ts @@ -0,0 +1,59 @@ +import { createRequire } from "node:module"; +import * as path from "node:path"; +import { fileURLToPath } from "node:url"; + +const cjsRequire = createRequire(import.meta.url); + +export type ModifierKey = "shift" | "command" | "control" | "option"; + +type NativeModifiersHelper = { + isModifierPressed: (name: ModifierKey) => boolean; +}; + +let nativeModifiersHelper: NativeModifiersHelper | null | undefined; + +function isNativeModifiersHelper(value: unknown): value is NativeModifiersHelper { + if (typeof value !== "object" || value === null) return false; + const candidate = (value as { isModifierPressed?: unknown }).isModifierPressed; + return typeof candidate === "function"; +} + +function loadNativeModifiersHelper(): NativeModifiersHelper | undefined { + if (nativeModifiersHelper !== undefined) return nativeModifiersHelper ?? undefined; + nativeModifiersHelper = null; + if (process.platform !== "darwin") return undefined; + const arch = process.arch; + if (arch !== "x64" && arch !== "arm64") return undefined; + + const moduleDir = path.dirname(fileURLToPath(import.meta.url)); + const nativePath = path.join("native", "darwin", "prebuilds", `darwin-${arch}`, "darwin-modifiers.node"); + const candidates = [ + path.join(moduleDir, "..", nativePath), + path.join(moduleDir, nativePath), + path.join(path.dirname(process.execPath), nativePath), + ]; + + for (const modulePath of candidates) { + try { + const helper = cjsRequire(modulePath) as unknown; + if (isNativeModifiersHelper(helper)) { + nativeModifiersHelper = helper; + return helper; + } + } catch { + // Try the next possible packaging location. + } + } + + return undefined; +} + +export function isNativeModifierPressed(key: ModifierKey): boolean { + const helper = loadNativeModifiersHelper(); + if (!helper) return false; + try { + return helper.isModifierPressed(key) === true; + } catch { + return false; + } +} diff --git a/cactus-code/packages/tui/src/stdin-buffer.ts b/cactus-code/packages/tui/src/stdin-buffer.ts new file mode 100644 index 000000000..a8b0b8478 --- /dev/null +++ b/cactus-code/packages/tui/src/stdin-buffer.ts @@ -0,0 +1,434 @@ +/** + * StdinBuffer buffers input and emits complete sequences. + * + * This is necessary because stdin data events can arrive in partial chunks, + * especially for escape sequences like mouse events. Without buffering, + * partial sequences can be misinterpreted as regular keypresses. + * + * For example, the mouse SGR sequence `\x1b[<35;20;5m` might arrive as: + * - Event 1: `\x1b` + * - Event 2: `[<35` + * - Event 3: `;20;5m` + * + * The buffer accumulates these until a complete sequence is detected. + * Call the `process()` method to feed input data. + * + * Based on code from OpenTUI (https://github.com/anomalyco/opentui) + * MIT License - Copyright (c) 2025 opentui + */ + +import { EventEmitter } from "events"; + +const ESC = "\x1b"; +const BRACKETED_PASTE_START = "\x1b[200~"; +const BRACKETED_PASTE_END = "\x1b[201~"; + +/** + * Check if a string is a complete escape sequence or needs more data + */ +function isCompleteSequence(data: string): "complete" | "incomplete" | "not-escape" { + if (!data.startsWith(ESC)) { + return "not-escape"; + } + + if (data.length === 1) { + return "incomplete"; + } + + const afterEsc = data.slice(1); + + // CSI sequences: ESC [ + if (afterEsc.startsWith("[")) { + // Check for old-style mouse sequence: ESC[M + 3 bytes + if (afterEsc.startsWith("[M")) { + // Old-style mouse needs ESC[M + 3 bytes = 6 total + return data.length >= 6 ? "complete" : "incomplete"; + } + return isCompleteCsiSequence(data); + } + + // OSC sequences: ESC ] + if (afterEsc.startsWith("]")) { + return isCompleteOscSequence(data); + } + + // DCS sequences: ESC P ... ESC \ (includes XTVersion responses) + if (afterEsc.startsWith("P")) { + return isCompleteDcsSequence(data); + } + + // APC sequences: ESC _ ... ESC \ (includes Kitty graphics responses) + if (afterEsc.startsWith("_")) { + return isCompleteApcSequence(data); + } + + // SS3 sequences: ESC O + if (afterEsc.startsWith("O")) { + // ESC O followed by a single character + return afterEsc.length >= 2 ? "complete" : "incomplete"; + } + + // Meta key sequences: ESC followed by a single character + if (afterEsc.length === 1) { + return "complete"; + } + + // Unknown escape sequence - treat as complete + return "complete"; +} + +/** + * Check if CSI sequence is complete + * CSI sequences: ESC [ ... followed by a final byte (0x40-0x7E) + */ +function isCompleteCsiSequence(data: string): "complete" | "incomplete" { + if (!data.startsWith(`${ESC}[`)) { + return "complete"; + } + + // Need at least ESC [ and one more character + if (data.length < 3) { + return "incomplete"; + } + + const payload = data.slice(2); + + // CSI sequences end with a byte in the range 0x40-0x7E (@-~) + // This includes all letters and several special characters + const lastChar = payload[payload.length - 1]; + const lastCharCode = lastChar.charCodeAt(0); + + if (lastCharCode >= 0x40 && lastCharCode <= 0x7e) { + // Special handling for SGR mouse sequences + // Format: ESC[ /^\d+$/.test(p))) { + return "complete"; + } + } + + return "incomplete"; + } + + return "complete"; + } + + return "incomplete"; +} + +/** + * Check if OSC sequence is complete + * OSC sequences: ESC ] ... ST (where ST is ESC \ or BEL) + */ +function isCompleteOscSequence(data: string): "complete" | "incomplete" { + if (!data.startsWith(`${ESC}]`)) { + return "complete"; + } + + // OSC sequences end with ST (ESC \) or BEL (\x07) + if (data.endsWith(`${ESC}\\`) || data.endsWith("\x07")) { + return "complete"; + } + + return "incomplete"; +} + +/** + * Check if DCS (Device Control String) sequence is complete + * DCS sequences: ESC P ... ST (where ST is ESC \) + * Used for XTVersion responses like ESC P >| ... ESC \ + */ +function isCompleteDcsSequence(data: string): "complete" | "incomplete" { + if (!data.startsWith(`${ESC}P`)) { + return "complete"; + } + + // DCS sequences end with ST (ESC \) + if (data.endsWith(`${ESC}\\`)) { + return "complete"; + } + + return "incomplete"; +} + +/** + * Check if APC (Application Program Command) sequence is complete + * APC sequences: ESC _ ... ST (where ST is ESC \) + * Used for Kitty graphics responses like ESC _ G ... ESC \ + */ +function isCompleteApcSequence(data: string): "complete" | "incomplete" { + if (!data.startsWith(`${ESC}_`)) { + return "complete"; + } + + // APC sequences end with ST (ESC \) + if (data.endsWith(`${ESC}\\`)) { + return "complete"; + } + + return "incomplete"; +} + +/** + * Split accumulated buffer into complete sequences + */ +function parseUnmodifiedKittyPrintableCodepoint(sequence: string): number | undefined { + const match = sequence.match(/^\x1b\[(\d+)(?::\d*)?(?::\d+)?u$/); + if (!match) return undefined; + + const codepoint = parseInt(match[1]!, 10); + return codepoint >= 32 ? codepoint : undefined; +} + +function extractCompleteSequences(buffer: string): { sequences: string[]; remainder: string } { + const sequences: string[] = []; + let pos = 0; + + while (pos < buffer.length) { + const remaining = buffer.slice(pos); + + // Try to extract a sequence starting at this position + if (remaining.startsWith(ESC)) { + // Find the end of this escape sequence + let seqEnd = 1; + while (seqEnd <= remaining.length) { + const candidate = remaining.slice(0, seqEnd); + const status = isCompleteSequence(candidate); + + if (status === "complete") { + // WezTerm with enable_kitty_keyboard sends the Escape key press as a + // raw '\x1b' byte (simple text path in encode_kitty, ignoring + // DISAMBIGUATE_ESCAPE_CODES) and the release as a full Kitty CSI-u + // sequence. These arrive concatenated as '\x1b\x1b[27;...u'. + // The buffer would normally treat '\x1b\x1b' as a complete meta-key + // sequence (ESC + single char), leaving '[27;...u' to be typed as + // plain text. If the character immediately following '\x1b\x1b' + // would begin a new escape sequence, emit only the first ESC and + // restart from the second. + if (candidate === "\x1b\x1b") { + const nextChar = remaining[seqEnd]; + if ( + nextChar === "[" || // CSI + nextChar === "]" || // OSC + nextChar === "O" || // SS3 + nextChar === "P" || // DCS + nextChar === "_" // APC + ) { + sequences.push(ESC); + pos += 1; + break; + } + } + sequences.push(candidate); + pos += seqEnd; + break; + } else if (status === "incomplete") { + seqEnd++; + } else { + // Should not happen when starting with ESC + sequences.push(candidate); + pos += seqEnd; + break; + } + } + + if (seqEnd > remaining.length) { + return { sequences, remainder: remaining }; + } + } else { + // Not an escape sequence - take a single character + sequences.push(remaining[0]!); + pos++; + } + } + + return { sequences, remainder: "" }; +} + +export type StdinBufferOptions = { + /** + * Maximum time to wait for sequence completion (default: 10ms) + * After this time, the buffer is flushed even if incomplete + */ + timeout?: number; +}; + +export type StdinBufferEventMap = { + data: [string]; + paste: [string]; +}; + +/** + * Buffers stdin input and emits complete sequences via the 'data' event. + * Handles partial escape sequences that arrive across multiple chunks. + */ +export class StdinBuffer extends EventEmitter { + private buffer: string = ""; + private timeout: ReturnType | null = null; + private readonly timeoutMs: number; + private pasteMode: boolean = false; + private pasteBuffer: string = ""; + private pendingKittyPrintableCodepoint: number | undefined; + + constructor(options: StdinBufferOptions = {}) { + super(); + this.timeoutMs = options.timeout ?? 10; + } + + public process(data: string | Buffer): void { + // Clear any pending timeout + if (this.timeout) { + clearTimeout(this.timeout); + this.timeout = null; + } + + // Handle high-byte conversion (for compatibility with parseKeypress) + // If buffer has single byte > 127, convert to ESC + (byte - 128) + let str: string; + if (Buffer.isBuffer(data)) { + if (data.length === 1 && data[0]! > 127) { + const byte = data[0]! - 128; + str = `\x1b${String.fromCharCode(byte)}`; + } else { + str = data.toString(); + } + } else { + str = data; + } + + if (str.length === 0 && this.buffer.length === 0) { + this.emitDataSequence(""); + return; + } + + this.buffer += str; + + if (this.pasteMode) { + this.pasteBuffer += this.buffer; + this.buffer = ""; + + const endIndex = this.pasteBuffer.indexOf(BRACKETED_PASTE_END); + if (endIndex !== -1) { + const pastedContent = this.pasteBuffer.slice(0, endIndex); + const remaining = this.pasteBuffer.slice(endIndex + BRACKETED_PASTE_END.length); + + this.pasteMode = false; + this.pasteBuffer = ""; + this.pendingKittyPrintableCodepoint = undefined; + + this.emit("paste", pastedContent); + + if (remaining.length > 0) { + this.process(remaining); + } + } + return; + } + + const startIndex = this.buffer.indexOf(BRACKETED_PASTE_START); + if (startIndex !== -1) { + if (startIndex > 0) { + const beforePaste = this.buffer.slice(0, startIndex); + const result = extractCompleteSequences(beforePaste); + for (const sequence of result.sequences) { + this.emitDataSequence(sequence); + } + } + + this.pendingKittyPrintableCodepoint = undefined; + this.buffer = this.buffer.slice(startIndex + BRACKETED_PASTE_START.length); + this.pasteMode = true; + this.pasteBuffer = this.buffer; + this.buffer = ""; + + const endIndex = this.pasteBuffer.indexOf(BRACKETED_PASTE_END); + if (endIndex !== -1) { + const pastedContent = this.pasteBuffer.slice(0, endIndex); + const remaining = this.pasteBuffer.slice(endIndex + BRACKETED_PASTE_END.length); + + this.pasteMode = false; + this.pasteBuffer = ""; + this.pendingKittyPrintableCodepoint = undefined; + + this.emit("paste", pastedContent); + + if (remaining.length > 0) { + this.process(remaining); + } + } + return; + } + + const result = extractCompleteSequences(this.buffer); + this.buffer = result.remainder; + + for (const sequence of result.sequences) { + this.emitDataSequence(sequence); + } + + if (this.buffer.length > 0) { + this.timeout = setTimeout(() => { + const flushed = this.flush(); + + for (const sequence of flushed) { + this.emitDataSequence(sequence); + } + }, this.timeoutMs); + } + } + + private emitDataSequence(sequence: string): void { + const rawCodepoint = sequence.length === 1 ? sequence.codePointAt(0) : undefined; + if (rawCodepoint !== undefined && rawCodepoint === this.pendingKittyPrintableCodepoint) { + this.pendingKittyPrintableCodepoint = undefined; + return; + } + + this.pendingKittyPrintableCodepoint = parseUnmodifiedKittyPrintableCodepoint(sequence); + this.emit("data", sequence); + } + + flush(): string[] { + if (this.timeout) { + clearTimeout(this.timeout); + this.timeout = null; + } + + if (this.buffer.length === 0) { + return []; + } + + const sequences = [this.buffer]; + this.buffer = ""; + this.pendingKittyPrintableCodepoint = undefined; + return sequences; + } + + clear(): void { + if (this.timeout) { + clearTimeout(this.timeout); + this.timeout = null; + } + this.buffer = ""; + this.pasteMode = false; + this.pasteBuffer = ""; + this.pendingKittyPrintableCodepoint = undefined; + } + + getBuffer(): string { + return this.buffer; + } + + destroy(): void { + this.clear(); + } +} diff --git a/cactus-code/packages/tui/src/terminal-colors.ts b/cactus-code/packages/tui/src/terminal-colors.ts new file mode 100644 index 000000000..cff6dc8e8 --- /dev/null +++ b/cactus-code/packages/tui/src/terminal-colors.ts @@ -0,0 +1,73 @@ +export interface RgbColor { + r: number; + g: number; + b: number; +} + +export type TerminalColorScheme = "dark" | "light"; + +function hexToRgb(hex: string): RgbColor { + const normalized = hex.startsWith("#") ? hex.slice(1) : hex; + const r = parseInt(normalized.slice(0, 2), 16); + const g = parseInt(normalized.slice(2, 4), 16); + const b = parseInt(normalized.slice(4, 6), 16); + return { r, g, b }; +} + +function parseOscHexChannel(channel: string): number | undefined { + if (!/^[0-9a-f]+$/i.test(channel)) { + return undefined; + } + const max = 16 ** channel.length - 1; + if (max <= 0) { + return undefined; + } + return Math.round((parseInt(channel, 16) / max) * 255); +} + +const OSC11_BACKGROUND_COLOR_RESPONSE_PATTERN = /^\x1b\]11;([^\x07\x1b]*)(?:\x07|\x1b\\)$/i; +const COLOR_SCHEME_REPORT_PATTERN = /^\x1b\[\?997;(1|2)n$/; + +export function isOsc11BackgroundColorResponse(data: string): boolean { + return OSC11_BACKGROUND_COLOR_RESPONSE_PATTERN.test(data); +} + +export function parseOsc11BackgroundColor(data: string): RgbColor | undefined { + const match = data.match(OSC11_BACKGROUND_COLOR_RESPONSE_PATTERN); + if (!match) { + return undefined; + } + + const value = match[1].trim(); + if (value.startsWith("#")) { + const hex = value.slice(1); + if (/^[0-9a-f]{6}$/i.test(hex)) { + return hexToRgb(value); + } + if (/^[0-9a-f]{12}$/i.test(hex)) { + const r = parseOscHexChannel(hex.slice(0, 4)); + const g = parseOscHexChannel(hex.slice(4, 8)); + const b = parseOscHexChannel(hex.slice(8, 12)); + return r !== undefined && g !== undefined && b !== undefined ? { r, g, b } : undefined; + } + return undefined; + } + + const rgbValue = value.replace(/^rgba?:/i, ""); + const [red, green, blue] = rgbValue.split("/"); + if (red === undefined || green === undefined || blue === undefined) { + return undefined; + } + const r = parseOscHexChannel(red); + const g = parseOscHexChannel(green); + const b = parseOscHexChannel(blue); + return r !== undefined && g !== undefined && b !== undefined ? { r, g, b } : undefined; +} + +export function parseTerminalColorSchemeReport(data: string): TerminalColorScheme | undefined { + const match = data.match(COLOR_SCHEME_REPORT_PATTERN); + if (!match) { + return undefined; + } + return match[1] === "2" ? "light" : "dark"; +} diff --git a/cactus-code/packages/tui/src/terminal-image.ts b/cactus-code/packages/tui/src/terminal-image.ts new file mode 100644 index 000000000..8854a86a7 --- /dev/null +++ b/cactus-code/packages/tui/src/terminal-image.ts @@ -0,0 +1,488 @@ +import { execSync } from "node:child_process"; + +export type ImageProtocol = "kitty" | "iterm2" | null; + +export interface TerminalCapabilities { + images: ImageProtocol; + trueColor: boolean; + hyperlinks: boolean; +} + +export interface CellDimensions { + widthPx: number; + heightPx: number; +} + +export interface ImageDimensions { + widthPx: number; + heightPx: number; +} + +export interface ImageRenderOptions { + maxWidthCells?: number; + maxHeightCells?: number; + preserveAspectRatio?: boolean; + /** Kitty image ID. If provided, reuses/replaces existing image with this ID. */ + imageId?: number; + /** Whether Kitty should apply its default cursor movement after placement. */ + moveCursor?: boolean; +} + +let cachedCapabilities: TerminalCapabilities | null = null; + +// Default cell dimensions - updated by TUI when terminal responds to query +let cellDimensions: CellDimensions = { widthPx: 9, heightPx: 18 }; + +export function getCellDimensions(): CellDimensions { + return cellDimensions; +} + +export function setCellDimensions(dims: CellDimensions): void { + cellDimensions = dims; +} + +/** + * Checks whether the attached tmux client forwards OSC 8 hyperlinks to the + * outer terminal. tmux only re-emits them when its `client_termfeatures` lists + * `hyperlinks`, and strips them otherwise. On any error fallbacks `false`. + */ +function probeTmuxHyperlinks(): boolean { + try { + const termfeatures = execSync("tmux display-message -p '#{client_termfeatures}'", { + encoding: "utf8", + timeout: 250, + stdio: ["ignore", "pipe", "ignore"], + }); + return termfeatures + .split(",") + .map((feature) => feature.trim()) + .includes("hyperlinks"); + } catch { + return false; + } +} + +export function detectCapabilities(tmuxForwardsHyperlink: () => boolean = probeTmuxHyperlinks): TerminalCapabilities { + const termProgram = process.env.TERM_PROGRAM?.toLowerCase() || ""; + const terminalEmulator = process.env.TERMINAL_EMULATOR?.toLowerCase() || ""; + const term = process.env.TERM?.toLowerCase() || ""; + const colorTerm = process.env.COLORTERM?.toLowerCase() || ""; + const hasTrueColorHint = colorTerm === "truecolor" || colorTerm === "24bit"; + + // Emit OSC 8 hyperlinks only when tmux confirms it forwards. + // Image protocols are unreliable under tmux, so leave `images: null`. + if (process.env.TMUX || term.startsWith("tmux")) { + return { images: null, trueColor: hasTrueColorHint, hyperlinks: tmuxForwardsHyperlink() }; + } + + // screen does not forward OSC 8 hyperlinks, so keep them off there. + if (term.startsWith("screen")) { + return { images: null, trueColor: hasTrueColorHint, hyperlinks: false }; + } + + if (process.env.KITTY_WINDOW_ID || termProgram === "kitty") { + return { images: "kitty", trueColor: true, hyperlinks: true }; + } + + if (termProgram === "ghostty" || term.includes("ghostty") || process.env.GHOSTTY_RESOURCES_DIR) { + return { images: "kitty", trueColor: true, hyperlinks: true }; + } + + if (process.env.WEZTERM_PANE || termProgram === "wezterm") { + return { images: "kitty", trueColor: true, hyperlinks: true }; + } + + // Warp supports the Kitty graphics protocol and OSC 8 hyperlinks. + if (termProgram === "warpterminal" || process.env.WARP_SESSION_ID || process.env.WARP_TERMINAL_SESSION_UUID) { + return { images: "kitty", trueColor: true, hyperlinks: true }; + } + + if (process.env.ITERM_SESSION_ID || termProgram === "iterm.app") { + return { images: "iterm2", trueColor: true, hyperlinks: true }; + } + + if (process.env.WT_SESSION) { + return { images: null, trueColor: true, hyperlinks: true }; + } + + if (termProgram === "vscode") { + return { images: null, trueColor: true, hyperlinks: true }; + } + + if (termProgram === "alacritty") { + return { images: null, trueColor: true, hyperlinks: true }; + } + + if (terminalEmulator === "jetbrains-jediterm") { + return { images: null, trueColor: true, hyperlinks: false }; + } + + // Unknown terminal: be conservative. OSC 8 is rendered invisibly as "just + // text" on terminals that swallow it, which means the URL disappears from + // the rendered output. Default to the legacy `text (url)` behavior unless we + // have positively identified a hyperlink-capable terminal above. + return { images: null, trueColor: hasTrueColorHint, hyperlinks: false }; +} + +export function getCapabilities(): TerminalCapabilities { + if (!cachedCapabilities) { + cachedCapabilities = detectCapabilities(); + } + return cachedCapabilities; +} + +export function resetCapabilitiesCache(): void { + cachedCapabilities = null; +} + +/** Override the cached capabilities. Useful in tests to exercise both code paths. */ +export function setCapabilities(caps: TerminalCapabilities): void { + cachedCapabilities = caps; +} + +const KITTY_PREFIX = "\x1b_G"; +const ITERM2_PREFIX = "\x1b]1337;File="; + +export function isImageLine(line: string): boolean { + // Fast path: sequence at line start (single-row images) + if (line.startsWith(KITTY_PREFIX) || line.startsWith(ITERM2_PREFIX)) { + return true; + } + // Slow path: sequence elsewhere (multi-row images have cursor-up prefix) + return line.includes(KITTY_PREFIX) || line.includes(ITERM2_PREFIX); +} + +/** + * Generate a random image ID for Kitty graphics protocol. + * Uses random IDs to avoid collisions between different module instances + * (e.g., main app vs extensions). + */ +export function allocateImageId(): number { + // Use random ID in range [1, 0xffffffff] to avoid collisions + return Math.floor(Math.random() * 0xfffffffe) + 1; +} + +export function encodeKitty( + base64Data: string, + options: { + columns?: number; + rows?: number; + imageId?: number; + /** Whether Kitty should apply its default cursor movement after placement. Default: true. */ + moveCursor?: boolean; + } = {}, +): string { + const CHUNK_SIZE = 4096; + + const params: string[] = ["a=T", "f=100", "q=2"]; + + if (options.moveCursor === false) params.push("C=1"); + if (options.columns) params.push(`c=${options.columns}`); + if (options.rows) params.push(`r=${options.rows}`); + if (options.imageId) params.push(`i=${options.imageId}`); + + if (base64Data.length <= CHUNK_SIZE) { + return `\x1b_G${params.join(",")};${base64Data}\x1b\\`; + } + + const chunks: string[] = []; + let offset = 0; + let isFirst = true; + + while (offset < base64Data.length) { + const chunk = base64Data.slice(offset, offset + CHUNK_SIZE); + const isLast = offset + CHUNK_SIZE >= base64Data.length; + + if (isFirst) { + chunks.push(`\x1b_G${params.join(",")},m=1;${chunk}\x1b\\`); + isFirst = false; + } else if (isLast) { + chunks.push(`\x1b_Gm=0;${chunk}\x1b\\`); + } else { + chunks.push(`\x1b_Gm=1;${chunk}\x1b\\`); + } + + offset += CHUNK_SIZE; + } + + return chunks.join(""); +} + +/** + * Delete a Kitty graphics image by ID. + * Uses uppercase 'I' to also free the image data. + */ +export function deleteKittyImage(imageId: number): string { + return `\x1b_Ga=d,d=I,i=${imageId},q=2\x1b\\`; +} + +/** + * Delete all visible Kitty graphics images. + * Uses uppercase 'A' to also free the image data. + */ +export function deleteAllKittyImages(): string { + return "\x1b_Ga=d,d=A,q=2\x1b\\"; +} + +export function encodeITerm2( + base64Data: string, + options: { + width?: number | string; + height?: number | string; + name?: string; + preserveAspectRatio?: boolean; + inline?: boolean; + } = {}, +): string { + const params: string[] = [`inline=${options.inline !== false ? 1 : 0}`]; + + if (options.width !== undefined) params.push(`width=${options.width}`); + if (options.height !== undefined) params.push(`height=${options.height}`); + if (options.name) { + const nameBase64 = Buffer.from(options.name).toString("base64"); + params.push(`name=${nameBase64}`); + } + if (options.preserveAspectRatio === false) { + params.push("preserveAspectRatio=0"); + } + + return `\x1b]1337;File=${params.join(";")}:${base64Data}\x07`; +} + +export interface ImageCellSize { + columns: number; + rows: number; +} + +export function calculateImageCellSize( + imageDimensions: ImageDimensions, + maxWidthCells: number, + maxHeightCells?: number, + cellDimensions: CellDimensions = { widthPx: 9, heightPx: 18 }, +): ImageCellSize { + const maxWidth = Math.max(1, Math.floor(maxWidthCells)); + const maxHeight = maxHeightCells === undefined ? undefined : Math.max(1, Math.floor(maxHeightCells)); + const imageWidth = Math.max(1, imageDimensions.widthPx); + const imageHeight = Math.max(1, imageDimensions.heightPx); + + const widthScale = (maxWidth * cellDimensions.widthPx) / imageWidth; + const heightScale = maxHeight === undefined ? widthScale : (maxHeight * cellDimensions.heightPx) / imageHeight; + const scale = Math.min(widthScale, heightScale); + + const scaledWidthPx = imageWidth * scale; + const scaledHeightPx = imageHeight * scale; + const columns = Math.ceil(scaledWidthPx / cellDimensions.widthPx); + const rows = Math.ceil(scaledHeightPx / cellDimensions.heightPx); + + return { + columns: Math.max(1, Math.min(maxWidth, columns)), + rows: Math.max(1, maxHeight === undefined ? rows : Math.min(maxHeight, rows)), + }; +} + +export function calculateImageRows( + imageDimensions: ImageDimensions, + targetWidthCells: number, + cellDimensions: CellDimensions = { widthPx: 9, heightPx: 18 }, +): number { + return calculateImageCellSize(imageDimensions, targetWidthCells, undefined, cellDimensions).rows; +} + +export function getPngDimensions(base64Data: string): ImageDimensions | null { + try { + const buffer = Buffer.from(base64Data, "base64"); + + if (buffer.length < 24) { + return null; + } + + if (buffer[0] !== 0x89 || buffer[1] !== 0x50 || buffer[2] !== 0x4e || buffer[3] !== 0x47) { + return null; + } + + const width = buffer.readUInt32BE(16); + const height = buffer.readUInt32BE(20); + + return { widthPx: width, heightPx: height }; + } catch { + return null; + } +} + +export function getJpegDimensions(base64Data: string): ImageDimensions | null { + try { + const buffer = Buffer.from(base64Data, "base64"); + + if (buffer.length < 2) { + return null; + } + + if (buffer[0] !== 0xff || buffer[1] !== 0xd8) { + return null; + } + + let offset = 2; + while (offset < buffer.length - 9) { + if (buffer[offset] !== 0xff) { + offset++; + continue; + } + + const marker = buffer[offset + 1]; + + if (marker >= 0xc0 && marker <= 0xc2) { + const height = buffer.readUInt16BE(offset + 5); + const width = buffer.readUInt16BE(offset + 7); + return { widthPx: width, heightPx: height }; + } + + if (offset + 3 >= buffer.length) { + return null; + } + const length = buffer.readUInt16BE(offset + 2); + if (length < 2) { + return null; + } + offset += 2 + length; + } + + return null; + } catch { + return null; + } +} + +export function getGifDimensions(base64Data: string): ImageDimensions | null { + try { + const buffer = Buffer.from(base64Data, "base64"); + + if (buffer.length < 10) { + return null; + } + + const sig = buffer.slice(0, 6).toString("ascii"); + if (sig !== "GIF87a" && sig !== "GIF89a") { + return null; + } + + const width = buffer.readUInt16LE(6); + const height = buffer.readUInt16LE(8); + + return { widthPx: width, heightPx: height }; + } catch { + return null; + } +} + +export function getWebpDimensions(base64Data: string): ImageDimensions | null { + try { + const buffer = Buffer.from(base64Data, "base64"); + + if (buffer.length < 30) { + return null; + } + + const riff = buffer.slice(0, 4).toString("ascii"); + const webp = buffer.slice(8, 12).toString("ascii"); + if (riff !== "RIFF" || webp !== "WEBP") { + return null; + } + + const chunk = buffer.slice(12, 16).toString("ascii"); + if (chunk === "VP8 ") { + if (buffer.length < 30) return null; + const width = buffer.readUInt16LE(26) & 0x3fff; + const height = buffer.readUInt16LE(28) & 0x3fff; + return { widthPx: width, heightPx: height }; + } else if (chunk === "VP8L") { + if (buffer.length < 25) return null; + const bits = buffer.readUInt32LE(21); + const width = (bits & 0x3fff) + 1; + const height = ((bits >> 14) & 0x3fff) + 1; + return { widthPx: width, heightPx: height }; + } else if (chunk === "VP8X") { + if (buffer.length < 30) return null; + const width = (buffer[24] | (buffer[25] << 8) | (buffer[26] << 16)) + 1; + const height = (buffer[27] | (buffer[28] << 8) | (buffer[29] << 16)) + 1; + return { widthPx: width, heightPx: height }; + } + + return null; + } catch { + return null; + } +} + +export function getImageDimensions(base64Data: string, mimeType: string): ImageDimensions | null { + if (mimeType === "image/png") { + return getPngDimensions(base64Data); + } + if (mimeType === "image/jpeg") { + return getJpegDimensions(base64Data); + } + if (mimeType === "image/gif") { + return getGifDimensions(base64Data); + } + if (mimeType === "image/webp") { + return getWebpDimensions(base64Data); + } + return null; +} + +export function renderImage( + base64Data: string, + imageDimensions: ImageDimensions, + options: ImageRenderOptions = {}, +): { sequence: string; rows: number; imageId?: number } | null { + const caps = getCapabilities(); + + if (!caps.images) { + return null; + } + + const maxWidth = options.maxWidthCells ?? 80; + const size = calculateImageCellSize(imageDimensions, maxWidth, options.maxHeightCells, getCellDimensions()); + + if (caps.images === "kitty") { + const sequence = encodeKitty(base64Data, { + columns: size.columns, + rows: size.rows, + imageId: options.imageId, + moveCursor: options.moveCursor, + }); + return { sequence, rows: size.rows, imageId: options.imageId }; + } + + if (caps.images === "iterm2") { + const sequence = encodeITerm2(base64Data, { + width: size.columns, + height: "auto", + preserveAspectRatio: options.preserveAspectRatio ?? true, + }); + return { sequence, rows: size.rows }; + } + + return null; +} + +/** + * Wrap text in an OSC 8 hyperlink sequence. + * The text is rendered as a clickable hyperlink in terminals that support OSC 8 + * (Ghostty, Kitty, WezTerm, iTerm2, VSCode, and others). + * In terminals that do not support OSC 8, the escape sequences are ignored + * and only the plain text is displayed. + * + * @param text - The visible text to display + * @param url - The URL to link to + */ +export function hyperlink(text: string, url: string): string { + return `\x1b]8;;${url}\x1b\\${text}\x1b]8;;\x1b\\`; +} + +export function imageFallback(mimeType: string, dimensions?: ImageDimensions, filename?: string): string { + const parts: string[] = []; + if (filename) parts.push(filename); + parts.push(`[${mimeType}]`); + if (dimensions) parts.push(`${dimensions.widthPx}x${dimensions.heightPx}`); + return `[Image: ${parts.join(" ")}]`; +} diff --git a/cactus-code/packages/tui/src/terminal.ts b/cactus-code/packages/tui/src/terminal.ts new file mode 100644 index 000000000..2542286ae --- /dev/null +++ b/cactus-code/packages/tui/src/terminal.ts @@ -0,0 +1,531 @@ +import * as fs from "node:fs"; +import { createRequire } from "node:module"; +import * as path from "node:path"; +import { fileURLToPath } from "node:url"; +import { setKittyProtocolActive } from "./keys.ts"; +import { isNativeModifierPressed } from "./native-modifiers.ts"; +import { StdinBuffer } from "./stdin-buffer.ts"; + +const cjsRequire = createRequire(import.meta.url); + +const TERMINAL_PROGRESS_KEEPALIVE_MS = 1000; +const TERMINAL_PROGRESS_ACTIVE_SEQUENCE = "\x1b]9;4;3\x07"; +const TERMINAL_PROGRESS_CLEAR_SEQUENCE = "\x1b]9;4;0;\x07"; +const APPLE_TERMINAL_SHIFT_ENTER_SEQUENCE = "\x1b[13;2u"; +const DESIRED_KITTY_KEYBOARD_PROTOCOL_FLAGS = 7; +const KEYBOARD_PROTOCOL_RESPONSE_FRAGMENT_TIMEOUT_MS = 150; +const KITTY_KEYBOARD_PROTOCOL_QUERY = `\x1b[>${DESIRED_KITTY_KEYBOARD_PROTOCOL_FLAGS}u\x1b[?u\x1b[c`; + +export type KeyboardProtocolNegotiationSequence = + | { type: "kitty-flags"; flags: number } + | { type: "device-attributes" }; + +export function parseKeyboardProtocolNegotiationSequence( + sequence: string, +): KeyboardProtocolNegotiationSequence | undefined { + const kittyFlags = sequence.match(/^\x1b\[\?(\d+)u$/); + if (kittyFlags) { + return { type: "kitty-flags", flags: Number.parseInt(kittyFlags[1]!, 10) }; + } + if (/^\x1b\[\?[\d;]*c$/.test(sequence)) { + return { type: "device-attributes" }; + } + return undefined; +} + +function isKeyboardProtocolNegotiationSequencePrefix(sequence: string): boolean { + return sequence === "\x1b[" || /^\x1b\[\?[\d;]*$/.test(sequence); +} + +export function isAppleTerminalSession(): boolean { + return process.platform === "darwin" && process.env.TERM_PROGRAM === "Apple_Terminal"; +} + +export function normalizeAppleTerminalInput(data: string, isAppleTerminal: boolean, isShiftPressed: boolean): string { + if (isAppleTerminal && data === "\r" && isShiftPressed) return APPLE_TERMINAL_SHIFT_ENTER_SEQUENCE; + return data; +} + +/** + * Minimal terminal interface for TUI + */ +export interface Terminal { + // Start the terminal with input and resize handlers + start(onInput: (data: string) => void, onResize: () => void): void; + + // Stop the terminal and restore state + stop(): void; + + /** + * Drain stdin before exiting to prevent Kitty key release events from + * leaking to the parent shell over slow SSH connections. + * @param maxMs - Maximum time to drain (default: 1000ms) + * @param idleMs - Exit early if no input arrives within this time (default: 50ms) + */ + drainInput(maxMs?: number, idleMs?: number): Promise; + + // Write output to terminal + write(data: string): void; + + // Get terminal dimensions + get columns(): number; + get rows(): number; + + // Whether Kitty keyboard protocol is active + get kittyProtocolActive(): boolean; + + // Cursor positioning (relative to current position) + moveBy(lines: number): void; // Move cursor up (negative) or down (positive) by N lines + + // Cursor visibility + hideCursor(): void; // Hide the cursor + showCursor(): void; // Show the cursor + + // Clear operations + clearLine(): void; // Clear current line + clearFromCursor(): void; // Clear from cursor to end of screen + clearScreen(): void; // Clear entire screen and move cursor to (0,0) + + // Title operations + setTitle(title: string): void; // Set terminal window title + + // Progress indicator (OSC 9;4) + setProgress(active: boolean): void; +} + +/** + * Real terminal using process.stdin/stdout + */ +export class ProcessTerminal implements Terminal { + private wasRaw = false; + private inputHandler?: (data: string) => void; + private resizeHandler?: () => void; + private _kittyProtocolActive = false; + private _modifyOtherKeysActive = false; + private keyboardProtocolPushed = false; + private keyboardProtocolNegotiationBuffer = ""; + private keyboardProtocolBufferFlushTimer?: ReturnType; + private stdinBuffer?: StdinBuffer; + private stdinDataHandler?: (data: string) => void; + private progressInterval?: ReturnType; + private writeLogPath = (() => { + const env = process.env.PI_TUI_WRITE_LOG || ""; + if (!env) return ""; + try { + if (fs.statSync(env).isDirectory()) { + const now = new Date(); + const ts = `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}-${String(now.getDate()).padStart(2, "0")}_${String(now.getHours()).padStart(2, "0")}-${String(now.getMinutes()).padStart(2, "0")}-${String(now.getSeconds()).padStart(2, "0")}`; + return path.join(env, `tui-${ts}-${process.pid}.log`); + } + } catch { + // Not an existing directory - use as-is (file path) + } + return env; + })(); + + get kittyProtocolActive(): boolean { + return this._kittyProtocolActive; + } + + get modifyOtherKeysActive(): boolean { + return this._modifyOtherKeysActive; + } + + start(onInput: (data: string) => void, onResize: () => void): void { + this.inputHandler = onInput; + this.resizeHandler = onResize; + + // Save previous state and enable raw mode + this.wasRaw = process.stdin.isRaw || false; + if (process.stdin.setRawMode) { + process.stdin.setRawMode(true); + } + process.stdin.setEncoding("utf8"); + process.stdin.resume(); + + // Enable bracketed paste mode - terminal will wrap pastes in \x1b[200~ ... \x1b[201~ + process.stdout.write("\x1b[?2004h"); + + // Set up resize handler immediately + process.stdout.on("resize", this.resizeHandler); + + // Refresh terminal dimensions - they may be stale after suspend/resume + // (SIGWINCH is lost while process is stopped). Unix only. + if (process.platform !== "win32") { + process.kill(process.pid, "SIGWINCH"); + } + + // On Windows, enable ENABLE_VIRTUAL_TERMINAL_INPUT so the console sends + // VT escape sequences (e.g. \x1b[Z for Shift+Tab) instead of raw console + // events that lose modifier information. Must run AFTER setRawMode(true) + // since that resets console mode flags. + this.enableWindowsVTInput(); + + // Query Kitty keyboard protocol and fall back to modifyOtherKeys when DA confirms no Kitty response. + // See: https://sw.kovidgoyal.net/kitty/keyboard-protocol/ + this.queryAndEnableKittyProtocol(); + } + + /** + * Set up StdinBuffer to split batched input into individual sequences. + * This ensures components receive single events, making matchesKey/isKeyRelease work correctly. + * + * Also watches for Kitty protocol response and enables it when detected. + * This is done here (after stdinBuffer parsing) rather than on raw stdin + * to handle the case where the response arrives split across multiple events. + */ + private setupStdinBuffer(): void { + this.stdinBuffer = new StdinBuffer({ timeout: 10 }); + + // Forward individual sequences to the input handler + this.stdinBuffer.on("data", (sequence) => { + const negotiationSequence = this.readKeyboardProtocolNegotiationSequence(sequence); + if (negotiationSequence === "pending") { + this.scheduleKeyboardProtocolNegotiationBufferFlush(); + return; // Wait briefly for the rest of a split Kitty response. + } + if (this.handleKeyboardProtocolNegotiationSequence(negotiationSequence)) { + return; + } + + this.forwardInputSequence(sequence); + }); + + // Re-wrap paste content with bracketed paste markers for existing editor handling + this.stdinBuffer.on("paste", (content) => { + if (this.inputHandler) { + this.inputHandler(`\x1b[200~${content}\x1b[201~`); + } + }); + + // Handler that pipes stdin data through the buffer + this.stdinDataHandler = (data: string) => { + this.stdinBuffer!.process(data); + }; + } + + /** + * Query terminal for Kitty keyboard protocol support and enable it if available. + * + * Kitty's progressive enhancement detection requires requesting the desired + * flags before querying them. The trailing DA query is a sentinel supported by + * terminals that do not know Kitty keyboard protocol; receiving DA before a + * Kitty response enables modifyOtherKeys fallback without a startup timeout. + * + * The requested flags are: + * - 1 = disambiguate escape codes + * - 2 = report event types (press/repeat/release) + * - 4 = report alternate keys (shifted key, base layout key) + */ + private queryAndEnableKittyProtocol(): void { + this.setupStdinBuffer(); + process.stdin.on("data", this.stdinDataHandler!); + this.keyboardProtocolPushed = true; + this.clearKeyboardProtocolNegotiationBuffer(); + process.stdout.write(KITTY_KEYBOARD_PROTOCOL_QUERY); + } + + private handleKeyboardProtocolNegotiationSequence( + negotiationSequence: KeyboardProtocolNegotiationSequence | undefined, + ): boolean { + if (!negotiationSequence) return false; + this.clearKeyboardProtocolNegotiationBuffer(); + if (negotiationSequence.type === "kitty-flags") { + if (negotiationSequence.flags !== 0) { + this.disableModifyOtherKeys(); + if (!this._kittyProtocolActive) { + this._kittyProtocolActive = true; + setKittyProtocolActive(true); + } + } else { + this.enableModifyOtherKeys(); + } + return true; + } + + if (!this._kittyProtocolActive) { + this.enableModifyOtherKeys(); + } + return true; + } + + private readKeyboardProtocolNegotiationSequence( + sequence: string, + ): KeyboardProtocolNegotiationSequence | "pending" | undefined { + if (this.keyboardProtocolNegotiationBuffer) { + const bufferedSequence = this.keyboardProtocolNegotiationBuffer + sequence; + const negotiationSequence = parseKeyboardProtocolNegotiationSequence(bufferedSequence); + if (negotiationSequence) { + this.clearKeyboardProtocolNegotiationBuffer(); + return negotiationSequence; + } + if (isKeyboardProtocolNegotiationSequencePrefix(bufferedSequence)) { + this.setKeyboardProtocolNegotiationBuffer(bufferedSequence); + return "pending"; + } + this.flushKeyboardProtocolNegotiationBufferAsInput(); + } + + const negotiationSequence = parseKeyboardProtocolNegotiationSequence(sequence); + if (negotiationSequence) return negotiationSequence; + if (isKeyboardProtocolNegotiationSequencePrefix(sequence)) { + this.setKeyboardProtocolNegotiationBuffer(sequence); + return "pending"; + } + return undefined; + } + + private setKeyboardProtocolNegotiationBuffer(sequence: string): void { + this.clearKeyboardProtocolNegotiationBufferFlushTimer(); + this.keyboardProtocolNegotiationBuffer = sequence; + } + + private clearKeyboardProtocolNegotiationBuffer(): void { + this.clearKeyboardProtocolNegotiationBufferFlushTimer(); + this.keyboardProtocolNegotiationBuffer = ""; + } + + private flushKeyboardProtocolNegotiationBufferAsInput(): void { + if (!this.keyboardProtocolNegotiationBuffer) return; + const sequence = this.keyboardProtocolNegotiationBuffer; + this.clearKeyboardProtocolNegotiationBuffer(); + this.forwardInputSequence(sequence); + } + + private scheduleKeyboardProtocolNegotiationBufferFlush(): void { + if (!this.keyboardProtocolNegotiationBuffer || this.keyboardProtocolBufferFlushTimer) return; + this.keyboardProtocolBufferFlushTimer = setTimeout(() => { + this.keyboardProtocolBufferFlushTimer = undefined; + this.flushKeyboardProtocolNegotiationBufferAsInput(); + }, KEYBOARD_PROTOCOL_RESPONSE_FRAGMENT_TIMEOUT_MS); + } + + private clearKeyboardProtocolNegotiationBufferFlushTimer(): void { + if (!this.keyboardProtocolBufferFlushTimer) return; + clearTimeout(this.keyboardProtocolBufferFlushTimer); + this.keyboardProtocolBufferFlushTimer = undefined; + } + + private forwardInputSequence(sequence: string): void { + if (!this.inputHandler) return; + const isAppleTerminal = sequence === "\r" && isAppleTerminalSession(); + const input = normalizeAppleTerminalInput( + sequence, + isAppleTerminal, + isAppleTerminal && isNativeModifierPressed("shift"), + ); + this.inputHandler(input); + } + + private enableModifyOtherKeys(): void { + if (this._kittyProtocolActive || this._modifyOtherKeysActive) return; + process.stdout.write("\x1b[>4;2m"); + this._modifyOtherKeysActive = true; + } + + private disableModifyOtherKeys(): void { + if (!this._modifyOtherKeysActive) return; + process.stdout.write("\x1b[>4;0m"); + this._modifyOtherKeysActive = false; + } + + /** + * On Windows, add ENABLE_VIRTUAL_TERMINAL_INPUT (0x0200) to the stdin + * console handle so the terminal sends VT sequences for modified keys + * (e.g. \x1b[Z for Shift+Tab). Without this, libuv's ReadConsoleInputW + * discards modifier state and Shift+Tab arrives as plain \t. + */ + private enableWindowsVTInput(): void { + if (process.platform !== "win32") return; + try { + const arch = process.arch; + if (arch !== "x64" && arch !== "arm64") return; + + // Dynamic require so non-Windows and bundled/browser paths never load the + // native helper. In the npm package native/ is next to dist/; in compiled + // binary archives native/ is copied next to the executable. + const moduleDir = path.dirname(fileURLToPath(import.meta.url)); + const nativePath = path.join("native", "win32", "prebuilds", `win32-${arch}`, "win32-console-mode.node"); + const candidates = [ + path.join(moduleDir, "..", nativePath), + path.join(moduleDir, nativePath), + path.join(path.dirname(process.execPath), nativePath), + ]; + for (const modulePath of candidates) { + try { + const helper = cjsRequire(modulePath) as { enableVirtualTerminalInput?: () => boolean }; + helper.enableVirtualTerminalInput?.(); + return; + } catch { + // Try the next possible packaging location. + } + } + } catch { + // Native helper not available — Shift+Tab won't be distinguishable from Tab. + } + } + + async drainInput(maxMs = 1000, idleMs = 50): Promise { + const shouldDisableKittyProtocol = this.keyboardProtocolPushed || this._kittyProtocolActive; + this.clearKeyboardProtocolNegotiationBuffer(); + if (shouldDisableKittyProtocol) { + // Disable Kitty keyboard protocol first so any late key releases + // do not generate new Kitty escape sequences. + process.stdout.write("\x1b[ { + lastDataTime = Date.now(); + }; + + process.stdin.on("data", onData); + const endTime = Date.now() + maxMs; + + try { + while (true) { + const now = Date.now(); + const timeLeft = endTime - now; + if (timeLeft <= 0) break; + if (now - lastDataTime >= idleMs) break; + await new Promise((resolve) => setTimeout(resolve, Math.min(idleMs, timeLeft))); + } + } finally { + process.stdin.removeListener("data", onData); + this.inputHandler = previousHandler; + } + } + + stop(): void { + if (this.clearProgressInterval()) { + process.stdout.write(TERMINAL_PROGRESS_CLEAR_SEQUENCE); + } + + // Disable bracketed paste mode + process.stdout.write("\x1b[?2004l"); + + const shouldDisableKittyProtocol = this.keyboardProtocolPushed || this._kittyProtocolActive; + this.clearKeyboardProtocolNegotiationBuffer(); + + // Disable Kitty keyboard protocol if not already done by drainInput() + if (shouldDisableKittyProtocol) { + process.stdout.write("\x1b[ 0) { + // Move down + process.stdout.write(`\x1b[${lines}B`); + } else if (lines < 0) { + // Move up + process.stdout.write(`\x1b[${-lines}A`); + } + // lines === 0: no movement + } + + hideCursor(): void { + process.stdout.write("\x1b[?25l"); + } + + showCursor(): void { + process.stdout.write("\x1b[?25h"); + } + + clearLine(): void { + process.stdout.write("\x1b[K"); + } + + clearFromCursor(): void { + process.stdout.write("\x1b[J"); + } + + clearScreen(): void { + process.stdout.write("\x1b[2J\x1b[H"); // Clear screen and move to home (1,1) + } + + setTitle(title: string): void { + // OSC 0;title BEL - set terminal window title + process.stdout.write(`\x1b]0;${title}\x07`); + } + + setProgress(active: boolean): void { + if (active) { + // OSC 9;4;3 - indeterminate progress + process.stdout.write(TERMINAL_PROGRESS_ACTIVE_SEQUENCE); + if (!this.progressInterval) { + this.progressInterval = setInterval(() => { + process.stdout.write(TERMINAL_PROGRESS_ACTIVE_SEQUENCE); + }, TERMINAL_PROGRESS_KEEPALIVE_MS); + } + } else { + this.clearProgressInterval(); + // OSC 9;4;0 - clear progress + process.stdout.write(TERMINAL_PROGRESS_CLEAR_SEQUENCE); + } + } + + private clearProgressInterval(): boolean { + if (!this.progressInterval) return false; + clearInterval(this.progressInterval); + this.progressInterval = undefined; + return true; + } +} diff --git a/cactus-code/packages/tui/src/tui.ts b/cactus-code/packages/tui/src/tui.ts new file mode 100644 index 000000000..c1c1cbbe0 --- /dev/null +++ b/cactus-code/packages/tui/src/tui.ts @@ -0,0 +1,1714 @@ +/** + * Minimal TUI implementation with differential rendering + */ + +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { performance } from "node:perf_hooks"; +import { isKeyRelease, matchesKey } from "./keys.ts"; +import type { Terminal } from "./terminal.ts"; +import { + isOsc11BackgroundColorResponse, + parseOsc11BackgroundColor, + parseTerminalColorSchemeReport, + type RgbColor, + type TerminalColorScheme, +} from "./terminal-colors.ts"; +import { deleteKittyImage, getCapabilities, isImageLine, setCellDimensions } from "./terminal-image.ts"; +import { extractSegments, normalizeTerminalOutput, sliceByColumn, sliceWithWidth, visibleWidth } from "./utils.ts"; + +const KITTY_SEQUENCE_PREFIX = "\x1b_G"; + +interface KittyImageHeader { + ids: number[]; + rows: number; +} + +function parseKittyImageHeader(line: string): KittyImageHeader | undefined { + const sequenceStart = line.indexOf(KITTY_SEQUENCE_PREFIX); + if (sequenceStart === -1) return undefined; + + const paramsStart = sequenceStart + KITTY_SEQUENCE_PREFIX.length; + const paramsEnd = line.indexOf(";", paramsStart); + if (paramsEnd === -1) return undefined; + + const ids: number[] = []; + let rows = 1; + const params = line.slice(paramsStart, paramsEnd); + for (const param of params.split(",")) { + const [key, value] = param.split("=", 2); + if (value === undefined) continue; + const numberValue = Number(value); + if (!Number.isInteger(numberValue) || numberValue <= 0 || numberValue > 0xffffffff) continue; + if (key === "i") { + ids.push(numberValue); + } else if (key === "r") { + rows = numberValue; + } + } + return { ids, rows }; +} + +function extractKittyImageIds(line: string): number[] { + return parseKittyImageHeader(line)?.ids ?? []; +} + +function extractKittyImageRows(line: string): number { + return parseKittyImageHeader(line)?.rows ?? 1; +} + +/** + * Component interface - all components must implement this + */ +export interface Component { + /** + * Render the component to lines for the given viewport width + * @param width - Current viewport width + * @returns Array of strings, each representing a line + */ + render(width: number): string[]; + + /** + * Optional handler for keyboard input when component has focus + */ + handleInput?(data: string): void; + + /** + * If true, component receives key release events (Kitty protocol). + * Default is false - release events are filtered out. + */ + wantsKeyRelease?: boolean; + + /** + * Invalidate any cached rendering state. + * Called when theme changes or when component needs to re-render from scratch. + */ + invalidate(): void; +} + +type InputListenerResult = { consume?: boolean; data?: string } | undefined; +type InputListener = (data: string) => InputListenerResult; +type PendingOsc11BackgroundQuery = { + settled: boolean; + resolve: ((rgb: RgbColor | undefined) => void) | undefined; + timer: NodeJS.Timeout | undefined; +}; + +/** + * Interface for components that can receive focus and display a hardware cursor. + * When focused, the component should emit CURSOR_MARKER at the cursor position + * in its render output. TUI will find this marker and position the hardware + * cursor there for proper IME candidate window positioning. + */ +export interface Focusable { + /** Set by TUI when focus changes. Component should emit CURSOR_MARKER when true. */ + focused: boolean; +} + +/** Type guard to check if a component implements Focusable */ +export function isFocusable(component: Component | null): component is Component & Focusable { + return component !== null && "focused" in component; +} + +/** + * Cursor position marker - APC (Application Program Command) sequence. + * This is a zero-width escape sequence that terminals ignore. + * Components emit this at the cursor position when focused. + * TUI finds and strips this marker, then positions the hardware cursor there. + */ +export const CURSOR_MARKER = "\x1b_pi:c\x07"; + +export { visibleWidth }; + +/** + * Anchor position for overlays + */ +export type OverlayAnchor = + | "center" + | "top-left" + | "top-right" + | "bottom-left" + | "bottom-right" + | "top-center" + | "bottom-center" + | "left-center" + | "right-center"; + +/** + * Margin configuration for overlays + */ +export interface OverlayMargin { + top?: number; + right?: number; + bottom?: number; + left?: number; +} + +/** Value that can be absolute (number) or percentage (string like "50%") */ +export type SizeValue = number | `${number}%`; + +/** Parse a SizeValue into absolute value given a reference size */ +function parseSizeValue(value: SizeValue | undefined, referenceSize: number): number | undefined { + if (value === undefined) return undefined; + if (typeof value === "number") return value; + // Parse percentage string like "50%" + const match = value.match(/^(\d+(?:\.\d+)?)%$/); + if (match) { + return Math.floor((referenceSize * parseFloat(match[1])) / 100); + } + return undefined; +} + +function isTermuxSession(): boolean { + return Boolean(process.env.TERMUX_VERSION); +} + +/** + * Options for overlay positioning and sizing. + * Values can be absolute numbers or percentage strings (e.g., "50%"). + */ +export interface OverlayOptions { + // === Sizing === + /** Width in columns, or percentage of terminal width (e.g., "50%") */ + width?: SizeValue; + /** Minimum width in columns */ + minWidth?: number; + /** Maximum height in rows, or percentage of terminal height (e.g., "50%") */ + maxHeight?: SizeValue; + + // === Positioning - anchor-based === + /** Anchor point for positioning (default: 'center') */ + anchor?: OverlayAnchor; + /** Horizontal offset from anchor position (positive = right) */ + offsetX?: number; + /** Vertical offset from anchor position (positive = down) */ + offsetY?: number; + + // === Positioning - percentage or absolute === + /** Row position: absolute number, or percentage (e.g., "25%" = 25% from top) */ + row?: SizeValue; + /** Column position: absolute number, or percentage (e.g., "50%" = centered horizontally) */ + col?: SizeValue; + + // === Margin from terminal edges === + /** Margin from terminal edges. Number applies to all sides. */ + margin?: OverlayMargin | number; + + // === Visibility === + /** + * Control overlay visibility based on terminal dimensions. + * If provided, overlay is only rendered when this returns true. + * Called each render cycle with current terminal dimensions. + */ + visible?: (termWidth: number, termHeight: number) => boolean; + /** If true, don't capture keyboard focus when shown */ + nonCapturing?: boolean; +} + +/** Options for {@link OverlayHandle.unfocus}. */ +export interface OverlayUnfocusOptions { + /** Explicit target to focus after releasing this overlay. */ + target: Component | null; +} + +/** + * Handle returned by showOverlay for controlling the overlay + */ +export interface OverlayHandle { + /** Permanently remove the overlay (cannot be shown again) */ + hide(): void; + /** Temporarily hide or show the overlay */ + setHidden(hidden: boolean): void; + /** Check if overlay is temporarily hidden */ + isHidden(): boolean; + /** Focus this overlay and bring it to the visual front */ + focus(): void; + /** Release focus to the next visible capturing overlay or previous target, or to an explicit target when provided */ + unfocus(options?: OverlayUnfocusOptions): void; + /** Check if this overlay currently has focus */ + isFocused(): boolean; +} + +type OverlayStackEntry = { + component: Component; + options?: OverlayOptions; + preFocus: Component | null; + hidden: boolean; + focusOrder: number; +}; + +type OverlayBlockedFocusResume = { status: "restore-overlay" } | { status: "focus-target"; target: Component | null }; +type EligibleOverlayFocusRestoreState = { status: "eligible"; overlay: OverlayStackEntry }; +type BlockedOverlayFocusRestoreState = { + status: "blocked"; + overlay: OverlayStackEntry; + blockedBy: Component; + resume: OverlayBlockedFocusResume; +}; +type ActiveOverlayFocusRestoreState = EligibleOverlayFocusRestoreState | BlockedOverlayFocusRestoreState; +type OverlayFocusRestoreState = { status: "inactive" } | ActiveOverlayFocusRestoreState; +type OverlayFocusRestorePolicy = "clear" | "preserve"; + +/** + * Container - a component that contains other components + */ +export class Container implements Component { + children: Component[] = []; + + addChild(component: Component): void { + this.children.push(component); + } + + removeChild(component: Component): void { + const index = this.children.indexOf(component); + if (index !== -1) { + this.children.splice(index, 1); + } + } + + clear(): void { + this.children = []; + } + + invalidate(): void { + for (const child of this.children) { + child.invalidate?.(); + } + } + + render(width: number): string[] { + const lines: string[] = []; + for (const child of this.children) { + const childLines = child.render(width); + for (const line of childLines) { + lines.push(line); + } + } + return lines; + } +} + +/** + * TUI - Main class for managing terminal UI with differential rendering + */ +export class TUI extends Container { + public terminal: Terminal; + private previousLines: string[] = []; + private previousKittyImageIds = new Set(); + private previousWidth = 0; + private previousHeight = 0; + private focusedComponent: Component | null = null; + private inputListeners = new Set(); + + /** Global callback for debug key (Shift+Ctrl+D). Called before input is forwarded to focused component. */ + public onDebug?: () => void; + private renderRequested = false; + private renderTimer: NodeJS.Timeout | undefined; + private lastRenderAt = 0; + private static readonly MIN_RENDER_INTERVAL_MS = 16; + private cursorRow = 0; // Logical cursor row (end of rendered content) + private hardwareCursorRow = 0; // Actual terminal cursor row (may differ due to IME positioning) + private showHardwareCursor = process.env.PI_HARDWARE_CURSOR === "1"; + private clearOnShrink = process.env.PI_CLEAR_ON_SHRINK === "1"; // Clear empty rows when content shrinks (default: off) + private maxLinesRendered = 0; // Track terminal's working area (max lines ever rendered) + private previousViewportTop = 0; // Track previous viewport top for resize-aware cursor moves + private fullRedrawCount = 0; + private stopped = false; + private pendingOsc11BackgroundReplies = 0; + private pendingOsc11BackgroundQueries: PendingOsc11BackgroundQuery[] = []; + private terminalColorSchemeListeners = new Set<(scheme: TerminalColorScheme) => void>(); + private terminalColorSchemeNotificationsEnabled = false; + + // Overlay stack for modal components rendered on top of base content + private focusOrderCounter = 0; + private overlayStack: OverlayStackEntry[] = []; + private overlayFocusRestore: OverlayFocusRestoreState = { status: "inactive" }; + + constructor(terminal: Terminal, showHardwareCursor?: boolean) { + super(); + this.terminal = terminal; + if (showHardwareCursor !== undefined) { + this.showHardwareCursor = showHardwareCursor; + } + } + + get fullRedraws(): number { + return this.fullRedrawCount; + } + + getShowHardwareCursor(): boolean { + return this.showHardwareCursor; + } + + setShowHardwareCursor(enabled: boolean): void { + if (this.showHardwareCursor === enabled) return; + this.showHardwareCursor = enabled; + if (!enabled) { + this.terminal.hideCursor(); + } + this.requestRender(); + } + + getClearOnShrink(): boolean { + return this.clearOnShrink; + } + + /** + * Set whether to trigger full re-render when content shrinks. + * When true (default), empty rows are cleared when content shrinks. + * When false, empty rows remain (reduces redraws on slower terminals). + */ + setClearOnShrink(enabled: boolean): void { + this.clearOnShrink = enabled; + } + + setFocus(component: Component | null): void { + this.setFocusInternal({ component, overlayFocusRestore: "clear" }); + } + + private setFocusInternal({ + component, + overlayFocusRestore, + }: { + component: Component | null; + overlayFocusRestore: OverlayFocusRestorePolicy; + }): void { + const previousFocus = this.focusedComponent; + let nextFocus = component; + const previousFocusedOverlay = previousFocus + ? this.overlayStack.find((entry) => entry.component === previousFocus && this.isOverlayVisible(entry)) + : undefined; + const nextFocusIsOverlay = nextFocus ? this.overlayStack.some((entry) => entry.component === nextFocus) : false; + const restoreState = this.getVisibleOverlayFocusRestore(); + if (nextFocus && !nextFocusIsOverlay) { + if (restoreState.status === "blocked" && restoreState.blockedBy === previousFocus) { + if (restoreState.resume.status === "focus-target" || !this.isComponentMounted(restoreState.blockedBy)) { + nextFocus = this.resolveBlockedOverlayFocusResume(restoreState); + } else { + this.overlayFocusRestore = { + status: "blocked", + overlay: restoreState.overlay, + blockedBy: nextFocus, + resume: restoreState.resume, + }; + } + } else if ( + previousFocusedOverlay && + restoreState.status !== "inactive" && + restoreState.overlay === previousFocusedOverlay && + !this.isOverlayFocusAncestor(previousFocusedOverlay, nextFocus) + ) { + this.overlayFocusRestore = { + status: "blocked", + overlay: previousFocusedOverlay, + blockedBy: nextFocus, + resume: { status: "restore-overlay" }, + }; + } + } else if (nextFocus === null) { + if (restoreState.status === "blocked" && restoreState.blockedBy === previousFocus) { + nextFocus = this.resolveBlockedOverlayFocusResume(restoreState); + } else if (overlayFocusRestore === "clear") { + this.clearOverlayFocusRestore(); + } + } + + if (isFocusable(this.focusedComponent)) { + this.focusedComponent.focused = false; + } + + this.focusedComponent = nextFocus; + + if (isFocusable(nextFocus)) { + nextFocus.focused = true; + } + + const focusedOverlay = nextFocus + ? this.overlayStack.find((entry) => entry.component === nextFocus && this.isOverlayVisible(entry)) + : undefined; + if (focusedOverlay) { + this.overlayFocusRestore = { status: "eligible", overlay: focusedOverlay }; + } + } + + private clearOverlayFocusRestore(): void { + this.overlayFocusRestore = { status: "inactive" }; + } + + private clearOverlayFocusRestoreFor(overlay: OverlayStackEntry): void { + if (this.overlayFocusRestore.status !== "inactive" && this.overlayFocusRestore.overlay === overlay) { + this.clearOverlayFocusRestore(); + } + } + + private resolveBlockedOverlayFocusResume(restoreState: BlockedOverlayFocusRestoreState): Component | null { + if (restoreState.resume.status === "restore-overlay") return restoreState.overlay.component; + this.clearOverlayFocusRestore(); + return restoreState.resume.target; + } + + private getVisibleOverlayFocusRestore(): OverlayFocusRestoreState { + const restoreState = this.overlayFocusRestore; + if (restoreState.status === "inactive") return restoreState; + if (!this.overlayStack.includes(restoreState.overlay) || !this.isOverlayVisible(restoreState.overlay)) { + return { status: "inactive" }; + } + return restoreState; + } + + private isOverlayFocusAncestor(entry: OverlayStackEntry, component: Component): boolean { + const visited = new Set(); + let current = entry.preFocus; + while (current && !visited.has(current)) { + visited.add(current); + if (current === component) return true; + current = this.overlayStack.find((overlay) => overlay.component === current)?.preFocus ?? null; + } + return false; + } + + private retargetOverlayPreFocus(removed: OverlayStackEntry): void { + for (const overlay of this.overlayStack) { + if (overlay !== removed && overlay.preFocus === removed.component) { + overlay.preFocus = removed.preFocus; + } + } + } + + private isComponentMounted(component: Component): boolean { + return this.children.some((child) => this.containsComponent(child, component)); + } + + private containsComponent(root: Component, target: Component): boolean { + if (root === target) return true; + if (!(root instanceof Container)) return false; + return root.children.some((child) => this.containsComponent(child, target)); + } + + /** + * Show an overlay component with configurable positioning and sizing. + * Returns a handle to control the overlay's visibility. + */ + showOverlay(component: Component, options?: OverlayOptions): OverlayHandle { + const entry: OverlayStackEntry = { + component, + ...(options === undefined ? {} : { options }), + preFocus: this.focusedComponent, + hidden: false, + focusOrder: ++this.focusOrderCounter, + }; + this.overlayStack.push(entry); + // Only focus if overlay is actually visible + if (!options?.nonCapturing && this.isOverlayVisible(entry)) { + this.setFocus(component); + } + this.terminal.hideCursor(); + this.requestRender(); + + // Return handle for controlling this overlay + return { + hide: () => { + const index = this.overlayStack.indexOf(entry); + if (index !== -1) { + this.clearOverlayFocusRestoreFor(entry); + this.retargetOverlayPreFocus(entry); + this.overlayStack.splice(index, 1); + // Restore focus if this overlay had focus + if (this.focusedComponent === component) { + const topVisible = this.getTopmostVisibleOverlay(); + this.setFocus(topVisible?.component ?? entry.preFocus); + } + if (this.overlayStack.length === 0) this.terminal.hideCursor(); + this.requestRender(); + } + }, + setHidden: (hidden: boolean) => { + if (entry.hidden === hidden) return; + entry.hidden = hidden; + // Update focus when hiding/showing + if (hidden) { + this.clearOverlayFocusRestoreFor(entry); + // If this overlay had focus, move focus to next visible or preFocus + if (this.focusedComponent === component) { + const topVisible = this.getTopmostVisibleOverlay(); + this.setFocus(topVisible?.component ?? entry.preFocus); + } + } else { + // Restore focus to this overlay when showing (if it's actually visible) + if (!options?.nonCapturing && this.isOverlayVisible(entry)) { + entry.focusOrder = ++this.focusOrderCounter; + this.setFocus(component); + } + } + this.requestRender(); + }, + isHidden: () => entry.hidden, + focus: () => { + if (!this.overlayStack.includes(entry) || !this.isOverlayVisible(entry)) return; + entry.focusOrder = ++this.focusOrderCounter; + this.setFocus(component); + this.requestRender(); + }, + unfocus: (unfocusOptions) => { + const isFocused = this.focusedComponent === component; + const restoreState = this.overlayFocusRestore; + const hasPendingRestore = restoreState.status !== "inactive" && restoreState.overlay === entry; + if (!isFocused && !hasPendingRestore) return; + if ( + restoreState.status === "blocked" && + restoreState.overlay === entry && + this.focusedComponent === restoreState.blockedBy + ) { + if (unfocusOptions) { + this.overlayFocusRestore = { + status: "blocked", + overlay: entry, + blockedBy: restoreState.blockedBy, + resume: { status: "focus-target", target: unfocusOptions.target }, + }; + } else { + this.clearOverlayFocusRestore(); + } + this.requestRender(); + return; + } + this.clearOverlayFocusRestoreFor(entry); + if (isFocused || unfocusOptions) { + const topVisible = this.getTopmostVisibleOverlay(); + const fallbackTarget = topVisible && topVisible !== entry ? topVisible.component : entry.preFocus; + this.setFocus(unfocusOptions ? unfocusOptions.target : fallbackTarget); + } + this.requestRender(); + }, + isFocused: () => this.focusedComponent === component, + }; + } + + /** Hide the topmost overlay and restore previous focus. */ + hideOverlay(): void { + const overlay = this.overlayStack[this.overlayStack.length - 1]; + if (!overlay) return; + this.clearOverlayFocusRestoreFor(overlay); + this.retargetOverlayPreFocus(overlay); + this.overlayStack.pop(); + if (this.focusedComponent === overlay.component) { + // Find topmost visible overlay, or fall back to preFocus + const topVisible = this.getTopmostVisibleOverlay(); + this.setFocus(topVisible?.component ?? overlay.preFocus); + } + if (this.overlayStack.length === 0) this.terminal.hideCursor(); + this.requestRender(); + } + + /** Check if there are any visible overlays */ + hasOverlay(): boolean { + return this.overlayStack.some((o) => this.isOverlayVisible(o)); + } + + /** Check if an overlay entry is currently visible */ + private isOverlayVisible(entry: OverlayStackEntry): boolean { + if (entry.hidden) return false; + if (entry.options?.visible) { + return entry.options.visible(this.terminal.columns, this.terminal.rows); + } + return true; + } + + /** Find the visual-frontmost visible capturing overlay, if any */ + private getTopmostVisibleOverlay(): OverlayStackEntry | undefined { + let topmost: OverlayStackEntry | undefined; + for (const overlay of this.overlayStack) { + if (overlay.options?.nonCapturing || !this.isOverlayVisible(overlay)) continue; + if (!topmost || overlay.focusOrder > topmost.focusOrder) { + topmost = overlay; + } + } + return topmost; + } + + override invalidate(): void { + super.invalidate(); + for (const overlay of this.overlayStack) overlay.component.invalidate?.(); + } + + start(): void { + this.stopped = false; + this.terminal.start( + (data) => this.handleInput(data), + () => this.requestRender(), + ); + this.terminal.hideCursor(); + if (this.terminalColorSchemeNotificationsEnabled) { + this.terminal.write("\x1b[?2031h"); + } + this.queryCellSize(); + this.requestRender(); + } + + addInputListener(listener: InputListener): () => void { + this.inputListeners.add(listener); + return () => { + this.inputListeners.delete(listener); + }; + } + + removeInputListener(listener: InputListener): void { + this.inputListeners.delete(listener); + } + + onTerminalColorSchemeChange(listener: (scheme: TerminalColorScheme) => void): () => void { + this.terminalColorSchemeListeners.add(listener); + return () => { + this.terminalColorSchemeListeners.delete(listener); + }; + } + + setTerminalColorSchemeNotifications(enabled: boolean): void { + if (this.terminalColorSchemeNotificationsEnabled === enabled) { + return; + } + this.terminalColorSchemeNotificationsEnabled = enabled; + if (!this.stopped) { + this.terminal.write(enabled ? "\x1b[?2031h" : "\x1b[?2031l"); + } + } + + private queryCellSize(): void { + // Only query if terminal supports images (cell size is only used for image rendering) + if (!getCapabilities().images) { + return; + } + // Query terminal for cell size in pixels: CSI 16 t + // Response format: CSI 6 ; height ; width t + this.terminal.write("\x1b[16t"); + } + + stop(): void { + this.stopped = true; + if (this.renderTimer) { + clearTimeout(this.renderTimer); + this.renderTimer = undefined; + } + if (this.terminalColorSchemeNotificationsEnabled) { + this.terminal.write("\x1b[?2031l"); + } + // Move cursor to the end of the content to prevent overwriting/artifacts on exit + if (this.previousLines.length > 0) { + const targetRow = this.previousLines.length; // Line after the last content + const lineDiff = targetRow - this.hardwareCursorRow; + if (lineDiff > 0) { + this.terminal.write(`\x1b[${lineDiff}B`); + } else if (lineDiff < 0) { + this.terminal.write(`\x1b[${-lineDiff}A`); + } + this.terminal.write("\r\n"); + } + + this.terminal.showCursor(); + this.terminal.stop(); + } + + requestRender(force = false): void { + if (force) { + this.previousLines = []; + this.previousWidth = -1; // -1 triggers widthChanged, forcing a full clear + this.previousHeight = -1; // -1 triggers heightChanged, forcing a full clear + this.cursorRow = 0; + this.hardwareCursorRow = 0; + this.maxLinesRendered = 0; + this.previousViewportTop = 0; + if (this.renderTimer) { + clearTimeout(this.renderTimer); + this.renderTimer = undefined; + } + this.renderRequested = true; + process.nextTick(() => { + if (this.stopped || !this.renderRequested) { + return; + } + this.renderRequested = false; + this.lastRenderAt = performance.now(); + this.doRender(); + }); + return; + } + if (this.renderRequested) return; + this.renderRequested = true; + process.nextTick(() => this.scheduleRender()); + } + + private scheduleRender(): void { + if (this.stopped || this.renderTimer || !this.renderRequested) { + return; + } + const elapsed = performance.now() - this.lastRenderAt; + const delay = Math.max(0, TUI.MIN_RENDER_INTERVAL_MS - elapsed); + this.renderTimer = setTimeout(() => { + this.renderTimer = undefined; + if (this.stopped || !this.renderRequested) { + return; + } + this.renderRequested = false; + this.lastRenderAt = performance.now(); + this.doRender(); + if (this.renderRequested) { + this.scheduleRender(); + } + }, delay); + } + + private handleInput(data: string): void { + if (this.consumeOsc11BackgroundResponse(data)) { + return; + } + if (this.consumeTerminalColorSchemeReport(data)) { + return; + } + + if (this.inputListeners.size > 0) { + let current = data; + for (const listener of this.inputListeners) { + const result = listener(current); + if (result?.consume) { + return; + } + if (result?.data !== undefined) { + current = result.data; + } + } + if (current.length === 0) { + return; + } + data = current; + } + + // Consume terminal cell size responses without blocking unrelated input. + if (this.consumeCellSizeResponse(data)) { + return; + } + + // Global debug key handler (Shift+Ctrl+D) + if (matchesKey(data, "shift+ctrl+d") && this.onDebug) { + this.onDebug(); + return; + } + + // If focused component is an overlay, verify it's still visible + // (visibility can change due to terminal resize or visible() callback) + const focusedOverlay = this.overlayStack.find((o) => o.component === this.focusedComponent); + if (focusedOverlay && !this.isOverlayVisible(focusedOverlay)) { + // Focused overlay is no longer visible, redirect to topmost visible overlay + const topVisible = this.getTopmostVisibleOverlay(); + if (topVisible) { + this.setFocus(topVisible.component); + } else { + this.setFocusInternal({ component: focusedOverlay.preFocus, overlayFocusRestore: "preserve" }); + } + } + + const focusIsOverlay = this.overlayStack.some((o) => o.component === this.focusedComponent); + if (!focusIsOverlay) { + const restoreState = this.getVisibleOverlayFocusRestore(); + if (restoreState.status === "eligible") { + this.setFocus(restoreState.overlay.component); + } else if (restoreState.status === "blocked" && restoreState.blockedBy !== this.focusedComponent) { + if (restoreState.resume.status === "restore-overlay") { + this.setFocus(restoreState.overlay.component); + } else { + this.clearOverlayFocusRestore(); + this.setFocus(restoreState.resume.target); + } + } + } + + // Pass input to focused component (including Ctrl+C) + // The focused component can decide how to handle Ctrl+C + if (this.focusedComponent?.handleInput) { + // Filter out key release events unless component opts in + if (isKeyRelease(data) && !this.focusedComponent.wantsKeyRelease) { + return; + } + this.focusedComponent.handleInput(data); + this.requestRender(); + } + } + + private consumeOsc11BackgroundResponse(data: string): boolean { + if (this.pendingOsc11BackgroundReplies <= 0) { + return false; + } + + if (!isOsc11BackgroundColorResponse(data)) { + return false; + } + + const rgb = parseOsc11BackgroundColor(data); + this.pendingOsc11BackgroundReplies -= 1; + const query = this.pendingOsc11BackgroundQueries.shift(); + if (query && !query.settled) { + query.settled = true; + if (query.timer) { + clearTimeout(query.timer); + query.timer = undefined; + } + query.resolve?.(rgb); + query.resolve = undefined; + } + return true; + } + + private consumeTerminalColorSchemeReport(data: string): boolean { + const scheme = parseTerminalColorSchemeReport(data); + if (!scheme) { + return false; + } + + for (const listener of this.terminalColorSchemeListeners) { + listener(scheme); + } + return true; + } + + private consumeCellSizeResponse(data: string): boolean { + // Response format: ESC [ 6 ; height ; width t + const match = data.match(/^\x1b\[6;(\d+);(\d+)t$/); + if (!match) { + return false; + } + + const heightPx = parseInt(match[1], 10); + const widthPx = parseInt(match[2], 10); + if (heightPx <= 0 || widthPx <= 0) { + return true; + } + + setCellDimensions({ widthPx, heightPx }); + // Invalidate all components so images re-render with correct dimensions. + this.invalidate(); + this.requestRender(); + return true; + } + + /** + * Resolve overlay layout from options. + * Returns { width, row, col, maxHeight } for rendering. + */ + private resolveOverlayLayout( + options: OverlayOptions | undefined, + overlayHeight: number, + termWidth: number, + termHeight: number, + ): { width: number; row: number; col: number; maxHeight: number | undefined } { + const opt = options ?? {}; + + // Parse margin (clamp to non-negative) + const margin = + typeof opt.margin === "number" + ? { top: opt.margin, right: opt.margin, bottom: opt.margin, left: opt.margin } + : (opt.margin ?? {}); + const marginTop = Math.max(0, margin.top ?? 0); + const marginRight = Math.max(0, margin.right ?? 0); + const marginBottom = Math.max(0, margin.bottom ?? 0); + const marginLeft = Math.max(0, margin.left ?? 0); + + // Available space after margins + const availWidth = Math.max(1, termWidth - marginLeft - marginRight); + const availHeight = Math.max(1, termHeight - marginTop - marginBottom); + + // === Resolve width === + let width = parseSizeValue(opt.width, termWidth) ?? Math.min(80, availWidth); + // Apply minWidth + if (opt.minWidth !== undefined) { + width = Math.max(width, opt.minWidth); + } + // Clamp to available space + width = Math.max(1, Math.min(width, availWidth)); + + // === Resolve maxHeight === + let maxHeight = parseSizeValue(opt.maxHeight, termHeight); + // Clamp to available space + if (maxHeight !== undefined) { + maxHeight = Math.max(1, Math.min(maxHeight, availHeight)); + } + + // Effective overlay height (may be clamped by maxHeight) + const effectiveHeight = maxHeight !== undefined ? Math.min(overlayHeight, maxHeight) : overlayHeight; + + // === Resolve position === + let row: number; + let col: number; + + if (opt.row !== undefined) { + if (typeof opt.row === "string") { + // Percentage: 0% = top, 100% = bottom (overlay stays within bounds) + const match = opt.row.match(/^(\d+(?:\.\d+)?)%$/); + if (match) { + const maxRow = Math.max(0, availHeight - effectiveHeight); + const percent = parseFloat(match[1]) / 100; + row = marginTop + Math.floor(maxRow * percent); + } else { + // Invalid format, fall back to center + row = this.resolveAnchorRow("center", effectiveHeight, availHeight, marginTop); + } + } else { + // Absolute row position + row = opt.row; + } + } else { + // Anchor-based (default: center) + const anchor = opt.anchor ?? "center"; + row = this.resolveAnchorRow(anchor, effectiveHeight, availHeight, marginTop); + } + + if (opt.col !== undefined) { + if (typeof opt.col === "string") { + // Percentage: 0% = left, 100% = right (overlay stays within bounds) + const match = opt.col.match(/^(\d+(?:\.\d+)?)%$/); + if (match) { + const maxCol = Math.max(0, availWidth - width); + const percent = parseFloat(match[1]) / 100; + col = marginLeft + Math.floor(maxCol * percent); + } else { + // Invalid format, fall back to center + col = this.resolveAnchorCol("center", width, availWidth, marginLeft); + } + } else { + // Absolute column position + col = opt.col; + } + } else { + // Anchor-based (default: center) + const anchor = opt.anchor ?? "center"; + col = this.resolveAnchorCol(anchor, width, availWidth, marginLeft); + } + + // Apply offsets + if (opt.offsetY !== undefined) row += opt.offsetY; + if (opt.offsetX !== undefined) col += opt.offsetX; + + // Clamp to terminal bounds (respecting margins) + row = Math.max(marginTop, Math.min(row, termHeight - marginBottom - effectiveHeight)); + col = Math.max(marginLeft, Math.min(col, termWidth - marginRight - width)); + + return { width, row, col, maxHeight }; + } + + private resolveAnchorRow(anchor: OverlayAnchor, height: number, availHeight: number, marginTop: number): number { + switch (anchor) { + case "top-left": + case "top-center": + case "top-right": + return marginTop; + case "bottom-left": + case "bottom-center": + case "bottom-right": + return marginTop + availHeight - height; + case "left-center": + case "center": + case "right-center": + return marginTop + Math.floor((availHeight - height) / 2); + } + } + + private resolveAnchorCol(anchor: OverlayAnchor, width: number, availWidth: number, marginLeft: number): number { + switch (anchor) { + case "top-left": + case "left-center": + case "bottom-left": + return marginLeft; + case "top-right": + case "right-center": + case "bottom-right": + return marginLeft + availWidth - width; + case "top-center": + case "center": + case "bottom-center": + return marginLeft + Math.floor((availWidth - width) / 2); + } + } + + /** Composite all overlays into content lines (sorted by focusOrder, higher = on top). */ + private compositeOverlays(lines: string[], termWidth: number, termHeight: number): string[] { + if (this.overlayStack.length === 0) return lines; + const result = [...lines]; + + // Pre-render all visible overlays and calculate positions + const rendered: { overlayLines: string[]; row: number; col: number; w: number }[] = []; + let minLinesNeeded = result.length; + + const visibleEntries = this.overlayStack.filter((e) => this.isOverlayVisible(e)); + visibleEntries.sort((a, b) => a.focusOrder - b.focusOrder); + for (const entry of visibleEntries) { + const { component, options } = entry; + + // Get layout with height=0 first to determine width and maxHeight + // (width and maxHeight don't depend on overlay height) + const { width, maxHeight } = this.resolveOverlayLayout(options, 0, termWidth, termHeight); + + // Render component at calculated width + let overlayLines = component.render(width); + + // Apply maxHeight if specified + if (maxHeight !== undefined && overlayLines.length > maxHeight) { + overlayLines = overlayLines.slice(0, maxHeight); + } + + // Get final row/col with actual overlay height + const { row, col } = this.resolveOverlayLayout(options, overlayLines.length, termWidth, termHeight); + + rendered.push({ overlayLines, row, col, w: width }); + minLinesNeeded = Math.max(minLinesNeeded, row + overlayLines.length); + } + + // Pad to at least terminal height so overlays have screen-relative positions. + // Excludes maxLinesRendered: the historical high-water mark caused self-reinforcing + // inflation that pushed content into scrollback on terminal widen. + const workingHeight = Math.max(result.length, termHeight, minLinesNeeded); + + // Extend result with empty lines if content is too short for overlay placement or working area + while (result.length < workingHeight) { + result.push(""); + } + + const viewportStart = Math.max(0, workingHeight - termHeight); + + // Composite each overlay + for (const { overlayLines, row, col, w } of rendered) { + for (let i = 0; i < overlayLines.length; i++) { + const idx = viewportStart + row + i; + if (idx >= 0 && idx < result.length) { + // Defensive: truncate overlay line to declared width before compositing + // (components should already respect width, but this ensures it) + const truncatedOverlayLine = + visibleWidth(overlayLines[i]) > w ? sliceByColumn(overlayLines[i], 0, w, true) : overlayLines[i]; + result[idx] = this.compositeLineAt(result[idx], truncatedOverlayLine, col, w, termWidth); + } + } + } + + return result; + } + + private static readonly SEGMENT_RESET = "\x1b[0m\x1b]8;;\x07"; + + private applyLineResets(lines: string[]): string[] { + const reset = TUI.SEGMENT_RESET; + for (let i = 0; i < lines.length; i++) { + const line = lines[i]; + if (!isImageLine(line)) { + lines[i] = normalizeTerminalOutput(line) + reset; + } + } + return lines; + } + + private collectKittyImageIds(lines: string[]): Set { + const ids = new Set(); + for (const line of lines) { + for (const id of extractKittyImageIds(line)) { + ids.add(id); + } + } + return ids; + } + + private deleteKittyImages(ids: Iterable): string { + let buffer = ""; + for (const id of ids) { + buffer += deleteKittyImage(id); + } + return buffer; + } + + private getKittyImageReservedRows(lines: string[], index: number, maxIndex = lines.length - 1): number { + const rows = extractKittyImageRows(lines[index] ?? ""); + if (rows <= 1) return 1; + + const maxRows = Math.min(rows, maxIndex - index + 1, lines.length - index); + let reservedRows = 1; + while (reservedRows < maxRows) { + const line = lines[index + reservedRows] ?? ""; + if (isImageLine(line) || visibleWidth(line) > 0) break; + reservedRows++; + } + return reservedRows; + } + + private expandChangedRangeForKittyImages( + firstChanged: number, + lastChanged: number, + newLines: string[], + ): { firstChanged: number; lastChanged: number } { + let expandedFirstChanged = firstChanged; + let expandedLastChanged = lastChanged; + const expandForLines = (lines: string[]): void => { + for (let i = 0; i < lines.length; i++) { + if (extractKittyImageIds(lines[i]).length === 0) continue; + const blockEnd = i + this.getKittyImageReservedRows(lines, i) - 1; + if (i >= firstChanged || (i <= lastChanged && blockEnd >= firstChanged)) { + expandedFirstChanged = Math.min(expandedFirstChanged, i); + expandedLastChanged = Math.max(expandedLastChanged, blockEnd); + } + } + }; + + expandForLines(this.previousLines); + expandForLines(newLines); + return { firstChanged: expandedFirstChanged, lastChanged: expandedLastChanged }; + } + + private deleteChangedKittyImages(firstChanged: number, lastChanged: number): string { + if (firstChanged < 0 || lastChanged < firstChanged) return ""; + + const ids = new Set(); + const maxLine = Math.min(lastChanged, this.previousLines.length - 1); + for (let i = firstChanged; i <= maxLine; i++) { + for (const id of extractKittyImageIds(this.previousLines[i] ?? "")) { + ids.add(id); + } + } + + return this.deleteKittyImages(ids); + } + + /** Splice overlay content into a base line at a specific column. Single-pass optimized. */ + private compositeLineAt( + baseLine: string, + overlayLine: string, + startCol: number, + overlayWidth: number, + totalWidth: number, + ): string { + if (isImageLine(baseLine)) return baseLine; + + // Single pass through baseLine extracts both before and after segments + const afterStart = startCol + overlayWidth; + const base = extractSegments(baseLine, startCol, afterStart, totalWidth - afterStart, true); + + // Extract overlay with width tracking (strict=true to exclude wide chars at boundary) + const overlay = sliceWithWidth(overlayLine, 0, overlayWidth, true); + + // Pad segments to target widths + const beforePad = Math.max(0, startCol - base.beforeWidth); + const overlayPad = Math.max(0, overlayWidth - overlay.width); + const actualBeforeWidth = Math.max(startCol, base.beforeWidth); + const actualOverlayWidth = Math.max(overlayWidth, overlay.width); + const afterTarget = Math.max(0, totalWidth - actualBeforeWidth - actualOverlayWidth); + const afterPad = Math.max(0, afterTarget - base.afterWidth); + + // Compose result + const r = TUI.SEGMENT_RESET; + const result = + base.before + + " ".repeat(beforePad) + + r + + overlay.text + + " ".repeat(overlayPad) + + r + + base.after + + " ".repeat(afterPad); + + // CRITICAL: Always verify and truncate to terminal width. + // This is the final safeguard against width overflow which would crash the TUI. + // Width tracking can drift from actual visible width due to: + // - Complex ANSI/OSC sequences (hyperlinks, colors) + // - Wide characters at segment boundaries + // - Edge cases in segment extraction + const resultWidth = visibleWidth(result); + if (resultWidth <= totalWidth) { + return result; + } + // Truncate with strict=true to ensure we don't exceed totalWidth + return sliceByColumn(result, 0, totalWidth, true); + } + + /** + * Find and extract cursor position from rendered lines. + * Searches for CURSOR_MARKER, calculates its position, and strips it from the output. + * Only scans the bottom terminal height lines (visible viewport). + * @param lines - Rendered lines to search + * @param height - Terminal height (visible viewport size) + * @returns Cursor position { row, col } or null if no marker found + */ + private extractCursorPosition(lines: string[], height: number): { row: number; col: number } | null { + // Only scan the bottom `height` lines (visible viewport) + const viewportTop = Math.max(0, lines.length - height); + for (let row = lines.length - 1; row >= viewportTop; row--) { + const line = lines[row]; + const markerIndex = line.indexOf(CURSOR_MARKER); + if (markerIndex !== -1) { + // Calculate visual column (width of text before marker) + const beforeMarker = line.slice(0, markerIndex); + const col = visibleWidth(beforeMarker); + + // Strip marker from the line + lines[row] = line.slice(0, markerIndex) + line.slice(markerIndex + CURSOR_MARKER.length); + + return { row, col }; + } + } + return null; + } + + private doRender(): void { + if (this.stopped) return; + const width = this.terminal.columns; + const height = this.terminal.rows; + const widthChanged = this.previousWidth !== 0 && this.previousWidth !== width; + const heightChanged = this.previousHeight !== 0 && this.previousHeight !== height; + const previousBufferLength = this.previousHeight > 0 ? this.previousViewportTop + this.previousHeight : height; + let prevViewportTop = heightChanged ? Math.max(0, previousBufferLength - height) : this.previousViewportTop; + let viewportTop = prevViewportTop; + let hardwareCursorRow = this.hardwareCursorRow; + const computeLineDiff = (targetRow: number): number => { + const currentScreenRow = hardwareCursorRow - prevViewportTop; + const targetScreenRow = targetRow - viewportTop; + return targetScreenRow - currentScreenRow; + }; + + // Render all components to get new lines + let newLines = this.render(width); + + // Composite overlays into the rendered lines (before differential compare) + if (this.overlayStack.length > 0) { + newLines = this.compositeOverlays(newLines, width, height); + } + + // Extract cursor position before applying line resets (marker must be found first) + const cursorPos = this.extractCursorPosition(newLines, height); + + newLines = this.applyLineResets(newLines); + + // Helper to clear scrollback and viewport and render all new lines + const fullRender = (clear: boolean): void => { + this.fullRedrawCount += 1; + let buffer = "\x1b[?2026h"; // Begin synchronized output + if (clear) { + buffer += this.deleteKittyImages(this.previousKittyImageIds); + buffer += "\x1b[2J\x1b[H\x1b[3J"; // Clear screen, home, then clear scrollback + } + for (let i = 0; i < newLines.length; i++) { + if (i > 0) buffer += "\r\n"; + const line = newLines[i]; + const isImage = isImageLine(line); + const imageReservedRows = isImage ? this.getKittyImageReservedRows(newLines, i) : 1; + if (imageReservedRows > 1 && imageReservedRows <= height) { + for (let row = 1; row < imageReservedRows; row++) { + buffer += "\r\n"; + } + buffer += `\x1b[${imageReservedRows - 1}A`; + buffer += line; + buffer += `\x1b[${imageReservedRows - 1}B`; + i += imageReservedRows - 1; + continue; + } + buffer += line; + } + buffer += "\x1b[?2026l"; // End synchronized output + this.terminal.write(buffer); + this.cursorRow = Math.max(0, newLines.length - 1); + this.hardwareCursorRow = this.cursorRow; + // Reset max lines when clearing, otherwise track growth + if (clear) { + this.maxLinesRendered = newLines.length; + } else { + this.maxLinesRendered = Math.max(this.maxLinesRendered, newLines.length); + } + const bufferLength = Math.max(height, newLines.length); + this.previousViewportTop = Math.max(0, bufferLength - height); + this.positionHardwareCursor(cursorPos, newLines.length); + this.previousLines = newLines; + this.previousKittyImageIds = this.collectKittyImageIds(newLines); + this.previousWidth = width; + this.previousHeight = height; + }; + + const debugRedraw = process.env.PI_DEBUG_REDRAW === "1"; + const logRedraw = (reason: string): void => { + if (!debugRedraw) return; + const logPath = path.join(os.homedir(), ".cactus", "agent", "cactus-debug.log"); + const msg = `[${new Date().toISOString()}] fullRender: ${reason} (prev=${this.previousLines.length}, new=${newLines.length}, height=${height})\n`; + fs.appendFileSync(logPath, msg); + }; + + // First render - just output everything without clearing (assumes clean screen) + if (this.previousLines.length === 0 && !widthChanged && !heightChanged) { + logRedraw("first render"); + fullRender(false); + return; + } + + // Width changes always need a full re-render because wrapping changes. + if (widthChanged) { + logRedraw(`terminal width changed (${this.previousWidth} -> ${width})`); + fullRender(true); + return; + } + + // Height changes normally need a full re-render to keep the visible viewport aligned, + // but Termux changes height when the software keyboard shows or hides. + // In that environment, a full redraw causes the entire history to replay on every toggle. + if (heightChanged && !isTermuxSession()) { + logRedraw(`terminal height changed (${this.previousHeight} -> ${height})`); + fullRender(true); + return; + } + + // Content shrunk below the working area and no overlays - re-render to clear empty rows + // (overlays need the padding, so only do this when no overlays are active) + // Configurable via setClearOnShrink() or PI_CLEAR_ON_SHRINK=0 env var + if (this.clearOnShrink && newLines.length < this.maxLinesRendered && this.overlayStack.length === 0) { + logRedraw(`clearOnShrink (maxLinesRendered=${this.maxLinesRendered})`); + fullRender(true); + return; + } + + // Find first and last changed lines + let firstChanged = -1; + let lastChanged = -1; + const maxLines = Math.max(newLines.length, this.previousLines.length); + for (let i = 0; i < maxLines; i++) { + const oldLine = i < this.previousLines.length ? this.previousLines[i] : ""; + const newLine = i < newLines.length ? newLines[i] : ""; + + if (oldLine !== newLine) { + if (firstChanged === -1) { + firstChanged = i; + } + lastChanged = i; + } + } + const appendedLines = newLines.length > this.previousLines.length; + if (appendedLines) { + if (firstChanged === -1) { + firstChanged = this.previousLines.length; + } + lastChanged = newLines.length - 1; + } + if (firstChanged !== -1) { + const expandedRange = this.expandChangedRangeForKittyImages(firstChanged, lastChanged, newLines); + firstChanged = expandedRange.firstChanged; + lastChanged = expandedRange.lastChanged; + } + const appendStart = appendedLines && firstChanged === this.previousLines.length && firstChanged > 0; + + // No changes - but still need to update hardware cursor position if it moved + if (firstChanged === -1) { + this.positionHardwareCursor(cursorPos, newLines.length); + this.previousViewportTop = prevViewportTop; + this.previousHeight = height; + return; + } + + // All changes are in deleted lines (nothing to render, just clear) + if (firstChanged >= newLines.length) { + if (this.previousLines.length > newLines.length) { + let buffer = "\x1b[?2026h"; + buffer += this.deleteChangedKittyImages(firstChanged, lastChanged); + // Move to end of new content (clamp to 0 for empty content) + const targetRow = Math.max(0, newLines.length - 1); + if (targetRow < prevViewportTop) { + logRedraw(`deleted lines moved viewport up (${targetRow} < ${prevViewportTop})`); + fullRender(true); + return; + } + const lineDiff = computeLineDiff(targetRow); + if (lineDiff > 0) buffer += `\x1b[${lineDiff}B`; + else if (lineDiff < 0) buffer += `\x1b[${-lineDiff}A`; + buffer += "\r"; + // Clear extra lines without scrolling + const extraLines = this.previousLines.length - newLines.length; + if (extraLines > height) { + logRedraw(`extraLines > height (${extraLines} > ${height})`); + fullRender(true); + return; + } + const clearStartOffset = newLines.length === 0 ? 0 : 1; + if (extraLines > 0 && clearStartOffset > 0) { + buffer += `\x1b[${clearStartOffset}B`; + } + for (let i = 0; i < extraLines; i++) { + buffer += "\r\x1b[2K"; + if (i < extraLines - 1) buffer += "\x1b[1B"; + } + const moveBack = Math.max(0, extraLines - 1 + clearStartOffset); + if (moveBack > 0) { + buffer += `\x1b[${moveBack}A`; + } + buffer += "\x1b[?2026l"; + this.terminal.write(buffer); + this.cursorRow = targetRow; + this.hardwareCursorRow = targetRow; + } + this.positionHardwareCursor(cursorPos, newLines.length); + this.previousLines = newLines; + this.previousKittyImageIds = this.collectKittyImageIds(newLines); + this.previousWidth = width; + this.previousHeight = height; + this.previousViewportTop = prevViewportTop; + return; + } + + // Differential rendering can only touch what was actually visible. + // If the first changed line is above the previous viewport, we need a full redraw. + if (firstChanged < prevViewportTop) { + logRedraw(`firstChanged < viewportTop (${firstChanged} < ${prevViewportTop})`); + fullRender(true); + return; + } + + // Render from first changed line to end + // Build buffer with all updates wrapped in synchronized output + let buffer = "\x1b[?2026h"; // Begin synchronized output + buffer += this.deleteChangedKittyImages(firstChanged, lastChanged); + const prevViewportBottom = prevViewportTop + height - 1; + const moveTargetRow = appendStart ? firstChanged - 1 : firstChanged; + if (moveTargetRow > prevViewportBottom) { + const currentScreenRow = Math.max(0, Math.min(height - 1, hardwareCursorRow - prevViewportTop)); + const moveToBottom = height - 1 - currentScreenRow; + if (moveToBottom > 0) { + buffer += `\x1b[${moveToBottom}B`; + } + const scroll = moveTargetRow - prevViewportBottom; + buffer += "\r\n".repeat(scroll); + prevViewportTop += scroll; + viewportTop += scroll; + hardwareCursorRow = moveTargetRow; + } + + // Move cursor to first changed line (use hardwareCursorRow for actual position) + const lineDiff = computeLineDiff(moveTargetRow); + if (lineDiff > 0) { + buffer += `\x1b[${lineDiff}B`; // Move down + } else if (lineDiff < 0) { + buffer += `\x1b[${-lineDiff}A`; // Move up + } + + buffer += appendStart ? "\r\n" : "\r"; // Move to column 0 + + // Only render changed lines (firstChanged to lastChanged), not all lines to end + // This reduces flicker when only a single line changes (e.g., spinner animation) + const renderEnd = Math.min(lastChanged, newLines.length - 1); + for (let i = firstChanged; i <= renderEnd; i++) { + if (i > firstChanged) buffer += "\r\n"; + const line = newLines[i]; + const isImage = isImageLine(line); + const imageReservedRows = isImage ? this.getKittyImageReservedRows(newLines, i, renderEnd) : 1; + if (imageReservedRows > 1) { + const imageStartScreenRow = i - viewportTop; + if (imageStartScreenRow < 0 || imageStartScreenRow + imageReservedRows > height) { + logRedraw( + `kitty image pre-clear would scroll (${imageStartScreenRow} + ${imageReservedRows} > ${height})`, + ); + fullRender(true); + return; + } + + buffer += "\x1b[2K"; + for (let row = 1; row < imageReservedRows; row++) { + buffer += "\r\n\x1b[2K"; + } + buffer += `\x1b[${imageReservedRows - 1}A`; + buffer += line; + buffer += `\x1b[${imageReservedRows - 1}B`; + i += imageReservedRows - 1; + continue; + } + + buffer += "\x1b[2K"; // Clear current line + if (!isImage && visibleWidth(line) > width) { + // Log all lines to crash file for debugging + const crashLogPath = path.join(os.homedir(), ".cactus", "agent", "cactus-crash.log"); + const crashData = [ + `Crash at ${new Date().toISOString()}`, + `Terminal width: ${width}`, + `Line ${i} visible width: ${visibleWidth(line)}`, + "", + "=== All rendered lines ===", + ...newLines.map((l, idx) => `[${idx}] (w=${visibleWidth(l)}) ${l}`), + "", + ].join("\n"); + fs.mkdirSync(path.dirname(crashLogPath), { recursive: true }); + fs.writeFileSync(crashLogPath, crashData); + + // Clean up terminal state before throwing + this.stop(); + + const errorMsg = [ + `Rendered line ${i} exceeds terminal width (${visibleWidth(line)} > ${width}).`, + "", + "This is likely caused by a custom TUI component not truncating its output.", + "Use visibleWidth() to measure and truncateToWidth() to truncate lines.", + "", + `Debug log written to: ${crashLogPath}`, + ].join("\n"); + throw new Error(errorMsg); + } + buffer += line; + } + + // Track where cursor ended up after rendering + let finalCursorRow = renderEnd; + + // If we had more lines before, clear them and move cursor back + if (this.previousLines.length > newLines.length) { + // Move to end of new content first if we stopped before it + if (renderEnd < newLines.length - 1) { + const moveDown = newLines.length - 1 - renderEnd; + buffer += `\x1b[${moveDown}B`; + finalCursorRow = newLines.length - 1; + } + const extraLines = this.previousLines.length - newLines.length; + for (let i = newLines.length; i < this.previousLines.length; i++) { + buffer += "\r\n\x1b[2K"; + } + // Move cursor back to end of new content + buffer += `\x1b[${extraLines}A`; + } + + buffer += "\x1b[?2026l"; // End synchronized output + + if (process.env.PI_TUI_DEBUG === "1") { + const debugDir = "/tmp/tui"; + fs.mkdirSync(debugDir, { recursive: true }); + const debugPath = path.join(debugDir, `render-${Date.now()}-${Math.random().toString(36).slice(2)}.log`); + const debugData = [ + `firstChanged: ${firstChanged}`, + `viewportTop: ${viewportTop}`, + `cursorRow: ${this.cursorRow}`, + `height: ${height}`, + `lineDiff: ${lineDiff}`, + `hardwareCursorRow: ${hardwareCursorRow}`, + `renderEnd: ${renderEnd}`, + `finalCursorRow: ${finalCursorRow}`, + `cursorPos: ${JSON.stringify(cursorPos)}`, + `newLines.length: ${newLines.length}`, + `previousLines.length: ${this.previousLines.length}`, + "", + "=== newLines ===", + JSON.stringify(newLines, null, 2), + "", + "=== previousLines ===", + JSON.stringify(this.previousLines, null, 2), + "", + "=== buffer ===", + JSON.stringify(buffer), + ].join("\n"); + fs.writeFileSync(debugPath, debugData); + } + + // Write entire buffer at once + this.terminal.write(buffer); + + // Track cursor position for next render + // cursorRow tracks end of content (for viewport calculation) + // hardwareCursorRow tracks actual terminal cursor position (for movement) + this.cursorRow = Math.max(0, newLines.length - 1); + this.hardwareCursorRow = finalCursorRow; + // Track terminal's working area (grows but doesn't shrink unless cleared) + this.maxLinesRendered = Math.max(this.maxLinesRendered, newLines.length); + this.previousViewportTop = Math.max(prevViewportTop, finalCursorRow - height + 1); + + // Position hardware cursor for IME + this.positionHardwareCursor(cursorPos, newLines.length); + + this.previousLines = newLines; + this.previousKittyImageIds = this.collectKittyImageIds(newLines); + this.previousWidth = width; + this.previousHeight = height; + } + + /** + * Position the hardware cursor for IME candidate window. + * @param cursorPos The cursor position extracted from rendered output, or null + * @param totalLines Total number of rendered lines + */ + private positionHardwareCursor(cursorPos: { row: number; col: number } | null, totalLines: number): void { + if (!cursorPos || totalLines <= 0) { + this.terminal.hideCursor(); + return; + } + + // Clamp cursor position to valid range + const targetRow = Math.max(0, Math.min(cursorPos.row, totalLines - 1)); + const targetCol = Math.max(0, cursorPos.col); + + // Move cursor from current position to target + const rowDelta = targetRow - this.hardwareCursorRow; + let buffer = ""; + if (rowDelta > 0) { + buffer += `\x1b[${rowDelta}B`; // Move down + } else if (rowDelta < 0) { + buffer += `\x1b[${-rowDelta}A`; // Move up + } + // Move to absolute column (1-indexed) + buffer += `\x1b[${targetCol + 1}G`; + + if (buffer) { + this.terminal.write(buffer); + } + + this.hardwareCursorRow = targetRow; + if (this.showHardwareCursor) { + this.terminal.showCursor(); + } else { + this.terminal.hideCursor(); + } + } + + /** + * Query the terminal's default background color with OSC 11 (`ESC ] 11 ; ? BEL`). + * @param timeoutMs Query timeout in milliseconds. + * @returns Promise containing the parsed RGB color, or undefined if it times out or fails to parse. + */ + queryTerminalBackgroundColor({ timeoutMs }: { timeoutMs: number }): Promise { + return new Promise((resolve) => { + const query: PendingOsc11BackgroundQuery = { + settled: false, + resolve, + timer: undefined, + }; + + query.timer = setTimeout(() => { + if (query.settled) { + return; + } + query.settled = true; + query.timer = undefined; + query.resolve?.(undefined); + query.resolve = undefined; + }, timeoutMs); + this.pendingOsc11BackgroundQueries.push(query); + this.pendingOsc11BackgroundReplies += 1; + this.terminal.write("\x1b]11;?\x07"); + }); + } + + /** + * Query the terminal's color-scheme preference with DSR (`CSI ? 996 n`). + * Terminals that support the color palette notification protocol reply with + * `CSI ? 997 ; 1 n` for dark or `CSI ? 997 ; 2 n` for light. + */ + queryTerminalColorScheme({ timeoutMs }: { timeoutMs: number }): Promise { + return new Promise((resolve) => { + let settled = false; + let timer: NodeJS.Timeout | undefined; + let unsubscribe: () => void = () => {}; + const settle = (scheme: TerminalColorScheme | undefined) => { + if (settled) return; + settled = true; + if (timer) { + clearTimeout(timer); + timer = undefined; + } + unsubscribe(); + resolve(scheme); + }; + + unsubscribe = this.onTerminalColorSchemeChange(settle); + timer = setTimeout(() => settle(undefined), timeoutMs); + this.terminal.write("\x1b[?996n"); + }); + } +} diff --git a/cactus-code/packages/tui/src/undo-stack.ts b/cactus-code/packages/tui/src/undo-stack.ts new file mode 100644 index 000000000..5b9a7e9ce --- /dev/null +++ b/cactus-code/packages/tui/src/undo-stack.ts @@ -0,0 +1,28 @@ +/** + * Generic undo stack with clone-on-push semantics. + * + * Stores deep clones of state snapshots. Popped snapshots are returned + * directly (no re-cloning) since they are already detached. + */ +export class UndoStack { + private stack: S[] = []; + + /** Push a deep clone of the given state onto the stack. */ + push(state: S): void { + this.stack.push(structuredClone(state)); + } + + /** Pop and return the most recent snapshot, or undefined if empty. */ + pop(): S | undefined { + return this.stack.pop(); + } + + /** Remove all snapshots. */ + clear(): void { + this.stack.length = 0; + } + + get length(): number { + return this.stack.length; + } +} diff --git a/cactus-code/packages/tui/src/utils.ts b/cactus-code/packages/tui/src/utils.ts new file mode 100644 index 000000000..b074a3a24 --- /dev/null +++ b/cactus-code/packages/tui/src/utils.ts @@ -0,0 +1,1188 @@ +import { eastAsianWidth } from "get-east-asian-width"; + +// segmenters (shared instance) +const graphemeSegmenter = new Intl.Segmenter(undefined, { granularity: "grapheme" }); +const wordSegmenter = new Intl.Segmenter(undefined, { granularity: "word" }); + +/** + * Get the shared grapheme segmenter instance. + */ +export function getGraphemeSegmenter(): Intl.Segmenter { + return graphemeSegmenter; +} + +/** + * Get the shared word segmenter instance. + */ +export function getWordSegmenter(): Intl.Segmenter { + return wordSegmenter; +} + +/** + * Check if a grapheme cluster (after segmentation) could possibly be an RGI emoji. + * This is a fast heuristic to avoid the expensive rgiEmojiRegex test. + * The tested Unicode blocks are deliberately broad to account for future + * Unicode additions. + */ +function couldBeEmoji(segment: string): boolean { + const cp = segment.codePointAt(0)!; + return ( + (cp >= 0x1f000 && cp <= 0x1fbff) || // Emoji and Pictograph + (cp >= 0x2300 && cp <= 0x23ff) || // Misc technical + (cp >= 0x2600 && cp <= 0x27bf) || // Misc symbols, dingbats + (cp >= 0x2b50 && cp <= 0x2b55) || // Specific stars/circles + segment.includes("\uFE0F") || // Contains VS16 (emoji presentation selector) + segment.length > 2 // Multi-codepoint sequences (ZWJ, skin tones, etc.) + ); +} + +// Regexes for character classification (same as string-width library) +const zeroWidthRegex = /^(?:\p{Default_Ignorable_Code_Point}|\p{Control}|\p{Mark}|\p{Surrogate})+$/v; +const leadingNonPrintingRegex = /^[\p{Default_Ignorable_Code_Point}\p{Control}\p{Format}\p{Mark}\p{Surrogate}]+/v; +const rgiEmojiRegex = /^\p{RGI_Emoji}$/v; + +// Cache for non-ASCII strings +const WIDTH_CACHE_SIZE = 512; +const widthCache = new Map(); + +export const cjkBreakRegex = + /[\p{Script_Extensions=Han}\p{Script_Extensions=Hiragana}\p{Script_Extensions=Katakana}\p{Script_Extensions=Hangul}\p{Script_Extensions=Bopomofo}]/u; + +function isPrintableAscii(str: string): boolean { + for (let i = 0; i < str.length; i++) { + const code = str.charCodeAt(i); + if (code < 0x20 || code > 0x7e) { + return false; + } + } + return true; +} + +function truncateFragmentToWidth(text: string, maxWidth: number): { text: string; width: number } { + if (maxWidth <= 0 || text.length === 0) { + return { text: "", width: 0 }; + } + + if (isPrintableAscii(text)) { + const clipped = text.slice(0, maxWidth); + return { text: clipped, width: clipped.length }; + } + + const hasAnsi = text.includes("\x1b"); + const hasTabs = text.includes("\t"); + if (!hasAnsi && !hasTabs) { + let result = ""; + let width = 0; + for (const { segment } of graphemeSegmenter.segment(text)) { + const w = graphemeWidth(segment); + if (width + w > maxWidth) { + break; + } + result += segment; + width += w; + } + return { text: result, width }; + } + + let result = ""; + let width = 0; + let i = 0; + let pendingAnsi = ""; + + while (i < text.length) { + const ansi = extractAnsiCode(text, i); + if (ansi) { + pendingAnsi += ansi.code; + i += ansi.length; + continue; + } + + if (text[i] === "\t") { + if (width + 3 > maxWidth) { + break; + } + if (pendingAnsi) { + result += pendingAnsi; + pendingAnsi = ""; + } + result += "\t"; + width += 3; + i++; + continue; + } + + let end = i; + while (end < text.length && text[end] !== "\t") { + const nextAnsi = extractAnsiCode(text, end); + if (nextAnsi) { + break; + } + end++; + } + + for (const { segment } of graphemeSegmenter.segment(text.slice(i, end))) { + const w = graphemeWidth(segment); + if (width + w > maxWidth) { + return { text: result, width }; + } + if (pendingAnsi) { + result += pendingAnsi; + pendingAnsi = ""; + } + result += segment; + width += w; + } + i = end; + } + + return { text: result, width }; +} + +function finalizeTruncatedResult( + prefix: string, + prefixWidth: number, + ellipsis: string, + ellipsisWidth: number, + maxWidth: number, + pad: boolean, +): string { + const reset = "\x1b[0m"; + const visibleWidth = prefixWidth + ellipsisWidth; + let result: string; + + if (ellipsis.length > 0) { + result = `${prefix}${reset}${ellipsis}${reset}`; + } else { + result = `${prefix}${reset}`; + } + + return pad ? result + " ".repeat(Math.max(0, maxWidth - visibleWidth)) : result; +} + +/** + * Calculate the terminal width of a single grapheme cluster. + * Based on code from the string-width library, but includes a possible-emoji + * check to avoid running the RGI_Emoji regex unnecessarily. + */ +function graphemeWidth(segment: string): number { + if (segment === "\t") { + return 3; + } + + // Zero-width clusters + if (zeroWidthRegex.test(segment)) { + return 0; + } + + // Emoji check with pre-filter + if (couldBeEmoji(segment) && rgiEmojiRegex.test(segment)) { + return 2; + } + + // Get base visible codepoint + const base = segment.replace(leadingNonPrintingRegex, ""); + const cp = base.codePointAt(0); + if (cp === undefined) { + return 0; + } + + // Regional indicator symbols (U+1F1E6..U+1F1FF) are often rendered as + // full-width emoji in terminals, even when isolated during streaming. + // Keep width conservative (2) to avoid terminal auto-wrap drift artifacts. + if (cp >= 0x1f1e6 && cp <= 0x1f1ff) { + return 2; + } + + let width = eastAsianWidth(cp); + + // Trailing halfwidth/fullwidth forms and AM vowels that segment with a base. + if (segment.length > 1) { + for (const char of segment.slice(1)) { + const c = char.codePointAt(0)!; + if (c >= 0xff00 && c <= 0xffef) { + width += eastAsianWidth(c); + } else if (c === 0x0e33 || c === 0x0eb3) { + width += 1; + } + } + } + + return width; +} + +/** + * Calculate the visible width of a string in terminal columns. + */ +export function visibleWidth(str: string): number { + if (str.length === 0) { + return 0; + } + + // Fast path: pure ASCII printable + if (isPrintableAscii(str)) { + return str.length; + } + + // Check cache + const cached = widthCache.get(str); + if (cached !== undefined) { + return cached; + } + + // Normalize: tabs to 3 spaces, strip ANSI escape codes + let clean = str; + if (str.includes("\t")) { + clean = clean.replace(/\t/g, " "); + } + if (clean.includes("\x1b")) { + // Strip supported ANSI/OSC/APC escape sequences in one pass. + // This covers CSI styling/cursor codes, OSC hyperlinks and prompt markers, + // and APC sequences like CURSOR_MARKER. + let stripped = ""; + let i = 0; + while (i < clean.length) { + const ansi = extractAnsiCode(clean, i); + if (ansi) { + i += ansi.length; + continue; + } + stripped += clean[i]; + i++; + } + clean = stripped; + } + + // Calculate width + let width = 0; + for (const { segment } of graphemeSegmenter.segment(clean)) { + width += graphemeWidth(segment); + } + + // Cache result + if (widthCache.size >= WIDTH_CACHE_SIZE) { + const firstKey = widthCache.keys().next().value; + if (firstKey !== undefined) { + widthCache.delete(firstKey); + } + } + widthCache.set(str, width); + + return width; +} + +/** + * Normalize text for terminal output without changing logical editor content. + * Some terminals render precomposed Thai/Lao AM vowels inconsistently during + * differential repaint. Their compatibility decompositions have the same cell + * width but avoid stale-cell artifacts in terminal renderers. + */ +const THAI_LAO_AM_REGEX = /[\u0e33\u0eb3]/; +const THAI_LAO_AM_GLOBAL_REGEX = /[\u0e33\u0eb3]/g; + +export function normalizeTerminalOutput(str: string): string { + if (!THAI_LAO_AM_REGEX.test(str)) return str; + return str.replace(THAI_LAO_AM_GLOBAL_REGEX, (char) => (char === "\u0e33" ? "\u0e4d\u0e32" : "\u0ecd\u0eb2")); +} + +/** + * Extract ANSI escape sequences from a string at the given position. + */ +export function extractAnsiCode(str: string, pos: number): { code: string; length: number } | null { + if (pos >= str.length || str[pos] !== "\x1b") return null; + + const next = str[pos + 1]; + + // CSI sequence: ESC [ ... m/G/K/H/J + if (next === "[") { + let j = pos + 2; + while (j < str.length && !/[mGKHJ]/.test(str[j]!)) j++; + if (j < str.length) return { code: str.substring(pos, j + 1), length: j + 1 - pos }; + return null; + } + + // OSC sequence: ESC ] ... BEL or ESC ] ... ST (ESC \) + // Used for hyperlinks (OSC 8), window titles, etc. + if (next === "]") { + let j = pos + 2; + while (j < str.length) { + if (str[j] === "\x07") return { code: str.substring(pos, j + 1), length: j + 1 - pos }; + if (str[j] === "\x1b" && str[j + 1] === "\\") return { code: str.substring(pos, j + 2), length: j + 2 - pos }; + j++; + } + return null; + } + + // APC sequence: ESC _ ... BEL or ESC _ ... ST (ESC \) + // Used for cursor marker and application-specific commands + if (next === "_") { + let j = pos + 2; + while (j < str.length) { + if (str[j] === "\x07") return { code: str.substring(pos, j + 1), length: j + 1 - pos }; + if (str[j] === "\x1b" && str[j + 1] === "\\") return { code: str.substring(pos, j + 2), length: j + 2 - pos }; + j++; + } + return null; + } + + return null; +} + +type Osc8Terminator = "\x07" | "\x1b\\"; + +interface ActiveHyperlink { + params: string; + url: string; + terminator: Osc8Terminator; +} + +function parseOsc8Hyperlink(ansiCode: string): ActiveHyperlink | null | undefined { + if (!ansiCode.startsWith("\x1b]8;")) { + return undefined; + } + + const terminator: Osc8Terminator = ansiCode.endsWith("\x07") ? "\x07" : "\x1b\\"; + const body = ansiCode.slice(4, terminator === "\x07" ? -1 : -2); + const separatorIndex = body.indexOf(";"); + if (separatorIndex === -1) { + return undefined; + } + + const params = body.slice(0, separatorIndex); + const url = body.slice(separatorIndex + 1); + if (!url) { + return null; + } + return { params, url, terminator }; +} + +function formatOsc8Hyperlink(hyperlink: ActiveHyperlink): string { + return `\x1b]8;${hyperlink.params};${hyperlink.url}${hyperlink.terminator}`; +} + +function formatOsc8Close(terminator: Osc8Terminator): string { + return `\x1b]8;;${terminator}`; +} + +/** + * Track active ANSI SGR codes to preserve styling across line breaks. + */ +class AnsiCodeTracker { + // Track individual attributes separately so we can reset them specifically + private bold = false; + private dim = false; + private italic = false; + private underline = false; + private blink = false; + private inverse = false; + private hidden = false; + private strikethrough = false; + private fgColor: string | null = null; // Stores the full code like "31" or "38;5;240" + private bgColor: string | null = null; // Stores the full code like "41" or "48;5;240" + private activeHyperlink: ActiveHyperlink | null = null; + + process(ansiCode: string): void { + // OSC 8 hyperlink: \x1b]8;;\x1b\\ (open) or \x1b]8;;\x1b\\ (close). + // Preserve the original terminator because some terminals only make BEL-terminated + // links clickable. OAuth login URLs use BEL, so reopening wrapped lines with ST + // made only the first physical line clickable in those terminals. + const hyperlink = parseOsc8Hyperlink(ansiCode); + if (hyperlink !== undefined) { + this.activeHyperlink = hyperlink; + return; + } + + if (!ansiCode.endsWith("m")) { + return; + } + + // Extract the parameters between \x1b[ and m + const match = ansiCode.match(/\x1b\[([\d;]*)m/); + if (!match) return; + + const params = match[1]; + if (params === "" || params === "0") { + // Full reset + this.reset(); + return; + } + + // Parse parameters (can be semicolon-separated) + const parts = params.split(";"); + let i = 0; + while (i < parts.length) { + const code = Number.parseInt(parts[i], 10); + + // Handle 256-color and RGB codes which consume multiple parameters + if (code === 38 || code === 48) { + // 38;5;N (256 color fg) or 38;2;R;G;B (RGB fg) + // 48;5;N (256 color bg) or 48;2;R;G;B (RGB bg) + if (parts[i + 1] === "5" && parts[i + 2] !== undefined) { + // 256 color: 38;5;N or 48;5;N + const colorCode = `${parts[i]};${parts[i + 1]};${parts[i + 2]}`; + if (code === 38) { + this.fgColor = colorCode; + } else { + this.bgColor = colorCode; + } + i += 3; + continue; + } else if (parts[i + 1] === "2" && parts[i + 4] !== undefined) { + // RGB color: 38;2;R;G;B or 48;2;R;G;B + const colorCode = `${parts[i]};${parts[i + 1]};${parts[i + 2]};${parts[i + 3]};${parts[i + 4]}`; + if (code === 38) { + this.fgColor = colorCode; + } else { + this.bgColor = colorCode; + } + i += 5; + continue; + } + } + + // Standard SGR codes + switch (code) { + case 0: + this.reset(); + break; + case 1: + this.bold = true; + break; + case 2: + this.dim = true; + break; + case 3: + this.italic = true; + break; + case 4: + this.underline = true; + break; + case 5: + this.blink = true; + break; + case 7: + this.inverse = true; + break; + case 8: + this.hidden = true; + break; + case 9: + this.strikethrough = true; + break; + case 21: + this.bold = false; + break; // Some terminals + case 22: + this.bold = false; + this.dim = false; + break; + case 23: + this.italic = false; + break; + case 24: + this.underline = false; + break; + case 25: + this.blink = false; + break; + case 27: + this.inverse = false; + break; + case 28: + this.hidden = false; + break; + case 29: + this.strikethrough = false; + break; + case 39: + this.fgColor = null; + break; // Default fg + case 49: + this.bgColor = null; + break; // Default bg + default: + // Standard foreground colors 30-37, 90-97 + if ((code >= 30 && code <= 37) || (code >= 90 && code <= 97)) { + this.fgColor = String(code); + } + // Standard background colors 40-47, 100-107 + else if ((code >= 40 && code <= 47) || (code >= 100 && code <= 107)) { + this.bgColor = String(code); + } + break; + } + i++; + } + } + + private reset(): void { + this.bold = false; + this.dim = false; + this.italic = false; + this.underline = false; + this.blink = false; + this.inverse = false; + this.hidden = false; + this.strikethrough = false; + this.fgColor = null; + this.bgColor = null; + // SGR reset does not affect OSC 8 hyperlink state + } + + /** Clear all state for reuse. */ + clear(): void { + this.reset(); + this.activeHyperlink = null; + } + + getActiveCodes(): string { + const codes: string[] = []; + if (this.bold) codes.push("1"); + if (this.dim) codes.push("2"); + if (this.italic) codes.push("3"); + if (this.underline) codes.push("4"); + if (this.blink) codes.push("5"); + if (this.inverse) codes.push("7"); + if (this.hidden) codes.push("8"); + if (this.strikethrough) codes.push("9"); + if (this.fgColor) codes.push(this.fgColor); + if (this.bgColor) codes.push(this.bgColor); + + let result = codes.length > 0 ? `\x1b[${codes.join(";")}m` : ""; + if (this.activeHyperlink) { + result += formatOsc8Hyperlink(this.activeHyperlink); + } + return result; + } + + hasActiveCodes(): boolean { + return ( + this.bold || + this.dim || + this.italic || + this.underline || + this.blink || + this.inverse || + this.hidden || + this.strikethrough || + this.fgColor !== null || + this.bgColor !== null || + this.activeHyperlink !== null + ); + } + + /** + * Get reset codes for attributes that need to be turned off at line end. + * Underline must be closed to prevent bleeding into padding. + * Active OSC 8 hyperlinks must be closed and re-opened on the next line. + * Returns empty string if no attributes need closing. + */ + getLineEndReset(): string { + let result = ""; + if (this.underline) { + result += "\x1b[24m"; // Underline off only + } + if (this.activeHyperlink) { + result += formatOsc8Close(this.activeHyperlink.terminator); // Re-opened at line start via getActiveCodes() + } + return result; + } +} + +function updateTrackerFromText(text: string, tracker: AnsiCodeTracker): void { + let i = 0; + while (i < text.length) { + const ansiResult = extractAnsiCode(text, i); + if (ansiResult) { + tracker.process(ansiResult.code); + i += ansiResult.length; + } else { + i++; + } + } +} + +/** + * Split text into words while keeping ANSI codes attached. + */ +function splitIntoTokensWithAnsi(text: string): string[] { + const tokens: string[] = []; + let current = ""; + let pendingAnsi = ""; // ANSI codes waiting to be attached to next visible content + let currentKind: "space" | "word" | null = null; + let i = 0; + + const flushCurrent = (): void => { + if (!current) { + return; + } + tokens.push(current); + current = ""; + currentKind = null; + }; + + while (i < text.length) { + const ansiResult = extractAnsiCode(text, i); + if (ansiResult) { + // Hold ANSI codes separately - they'll be attached to the next visible char + pendingAnsi += ansiResult.code; + i += ansiResult.length; + continue; + } + + let end = i; + while (end < text.length && !extractAnsiCode(text, end)) { + end++; + } + + for (const { segment } of graphemeSegmenter.segment(text.slice(i, end))) { + const segmentIsSpace = segment === " "; + if (!segmentIsSpace && cjkBreakRegex.test(segment)) { + flushCurrent(); + const token = pendingAnsi + segment; + pendingAnsi = ""; + tokens.push(token); + continue; + } + + const segmentKind = segmentIsSpace ? "space" : "word"; + if (current && currentKind !== segmentKind) { + flushCurrent(); + } + + // Attach any pending ANSI codes to this visible character + if (pendingAnsi) { + current += pendingAnsi; + pendingAnsi = ""; + } + + currentKind = segmentKind; + current += segment; + } + + i = end; + } + + // Handle any remaining pending ANSI codes (attach to last token) + if (pendingAnsi) { + if (current) { + current += pendingAnsi; + } else if (tokens.length > 0) { + tokens[tokens.length - 1] += pendingAnsi; + } else { + current = pendingAnsi; + } + } + + if (current) { + tokens.push(current); + } + + return tokens; +} + +/** + * Wrap text with ANSI codes preserved. + * + * ONLY does word wrapping - NO padding, NO background colors. + * Returns lines where each line is <= width visible chars. + * Active ANSI codes are preserved across line breaks. + * + * @param text - Text to wrap (may contain ANSI codes and newlines) + * @param width - Maximum visible width per line + * @returns Array of wrapped lines (NOT padded to width) + */ +export function wrapTextWithAnsi(text: string, width: number): string[] { + if (!text) { + return [""]; + } + + // Handle newlines by processing each line separately + // Track ANSI state across lines so styles carry over after literal newlines + const inputLines = text.split("\n"); + const result: string[] = []; + const tracker = new AnsiCodeTracker(); + + for (const inputLine of inputLines) { + // Prepend active ANSI codes from previous lines (except for first line) + const prefix = result.length > 0 ? tracker.getActiveCodes() : ""; + const wrappedLines = wrapSingleLine(prefix + inputLine, width); + for (const wrappedLine of wrappedLines) { + result.push(wrappedLine); + } + // Update tracker with codes from this line for next iteration + updateTrackerFromText(inputLine, tracker); + } + + return result.length > 0 ? result : [""]; +} + +function wrapSingleLine(line: string, width: number): string[] { + if (!line) { + return [""]; + } + + const visibleLength = visibleWidth(line); + if (visibleLength <= width) { + return [line]; + } + + const wrapped: string[] = []; + const tracker = new AnsiCodeTracker(); + const tokens = splitIntoTokensWithAnsi(line); + + let currentLine = ""; + let currentVisibleLength = 0; + + for (const token of tokens) { + const tokenVisibleLength = visibleWidth(token); + const isWhitespace = token.trim() === ""; + + // Token itself is too long - break it character by character + if (tokenVisibleLength > width && !isWhitespace) { + if (currentLine) { + // Add specific reset for underline only (preserves background) + const lineEndReset = tracker.getLineEndReset(); + if (lineEndReset) { + currentLine += lineEndReset; + } + wrapped.push(currentLine); + currentLine = ""; + currentVisibleLength = 0; + } + + // Break long token - breakLongWord handles its own resets + const broken = breakLongWord(token, width, tracker); + for (let i = 0; i < broken.length - 1; i++) { + wrapped.push(broken[i]!); + } + currentLine = broken[broken.length - 1]; + currentVisibleLength = visibleWidth(currentLine); + continue; + } + + // Check if adding this token would exceed width + const totalNeeded = currentVisibleLength + tokenVisibleLength; + + if (totalNeeded > width && currentVisibleLength > 0) { + // Trim trailing whitespace, then add underline reset (not full reset, to preserve background) + let lineToWrap = currentLine.trimEnd(); + const lineEndReset = tracker.getLineEndReset(); + if (lineEndReset) { + lineToWrap += lineEndReset; + } + wrapped.push(lineToWrap); + if (isWhitespace) { + // Don't start new line with whitespace + currentLine = tracker.getActiveCodes(); + currentVisibleLength = 0; + } else { + currentLine = tracker.getActiveCodes() + token; + currentVisibleLength = tokenVisibleLength; + } + } else { + // Add to current line + currentLine += token; + currentVisibleLength += tokenVisibleLength; + } + + updateTrackerFromText(token, tracker); + } + + if (currentLine) { + // No reset at end of final line - let caller handle it + wrapped.push(currentLine); + } + + // Trailing whitespace can cause lines to exceed the requested width + return wrapped.length > 0 ? wrapped.map((line) => line.trimEnd()) : [""]; +} + +export const PUNCTUATION_REGEX = /[(){}[\]<>.,;:'"!?+\-=*/\\|&%^$#@~`]/; + +/** + * Check if a character is whitespace. + */ +export function isWhitespaceChar(char: string): boolean { + return /\s/.test(char); +} + +/** + * Check if a character is punctuation. + */ +export function isPunctuationChar(char: string): boolean { + return PUNCTUATION_REGEX.test(char); +} + +function breakLongWord(word: string, width: number, tracker: AnsiCodeTracker): string[] { + const lines: string[] = []; + let currentLine = tracker.getActiveCodes(); + let currentWidth = 0; + + // First, separate ANSI codes from visible content + // We need to handle ANSI codes specially since they're not graphemes + let i = 0; + const segments: Array<{ type: "ansi" | "grapheme"; value: string }> = []; + + while (i < word.length) { + const ansiResult = extractAnsiCode(word, i); + if (ansiResult) { + segments.push({ type: "ansi", value: ansiResult.code }); + i += ansiResult.length; + } else { + // Find the next ANSI code or end of string + let end = i; + while (end < word.length) { + const nextAnsi = extractAnsiCode(word, end); + if (nextAnsi) break; + end++; + } + // Segment this non-ANSI portion into graphemes + const textPortion = word.slice(i, end); + for (const seg of graphemeSegmenter.segment(textPortion)) { + segments.push({ type: "grapheme", value: seg.segment }); + } + i = end; + } + } + + // Now process segments + for (const seg of segments) { + if (seg.type === "ansi") { + currentLine += seg.value; + tracker.process(seg.value); + continue; + } + + const grapheme = seg.value; + // Skip empty graphemes to avoid issues with string-width calculation + if (!grapheme) continue; + + const graphemeWidth = visibleWidth(grapheme); + + if (currentWidth + graphemeWidth > width) { + // Add specific reset for underline only (preserves background) + const lineEndReset = tracker.getLineEndReset(); + if (lineEndReset) { + currentLine += lineEndReset; + } + lines.push(currentLine); + currentLine = tracker.getActiveCodes(); + currentWidth = 0; + } + + currentLine += grapheme; + currentWidth += graphemeWidth; + } + + if (currentLine) { + // No reset at end of final segment - caller handles continuation + lines.push(currentLine); + } + + return lines.length > 0 ? lines : [""]; +} + +/** + * Apply background color to a line, padding to full width. + * + * @param line - Line of text (may contain ANSI codes) + * @param width - Total width to pad to + * @param bgFn - Background color function + * @returns Line with background applied and padded to width + */ +export function applyBackgroundToLine(line: string, width: number, bgFn: (text: string) => string): string { + // Calculate padding needed + const visibleLen = visibleWidth(line); + const paddingNeeded = Math.max(0, width - visibleLen); + const padding = " ".repeat(paddingNeeded); + + // Apply background to content + padding + const withPadding = line + padding; + return bgFn(withPadding); +} + +/** + * Truncate text to fit within a maximum visible width, adding ellipsis if needed. + * Optionally pad with spaces to reach exactly maxWidth. + * Properly handles ANSI escape codes (they don't count toward width). + * + * @param text - Text to truncate (may contain ANSI codes) + * @param maxWidth - Maximum visible width + * @param ellipsis - Ellipsis string to append when truncating (default: "...") + * @param pad - If true, pad result with spaces to exactly maxWidth (default: false) + * @returns Truncated text, optionally padded to exactly maxWidth + */ +export function truncateToWidth( + text: string, + maxWidth: number, + ellipsis: string = "...", + pad: boolean = false, +): string { + if (maxWidth <= 0) { + return ""; + } + + if (text.length === 0) { + return pad ? " ".repeat(maxWidth) : ""; + } + + const ellipsisWidth = visibleWidth(ellipsis); + if (ellipsisWidth >= maxWidth) { + const textWidth = visibleWidth(text); + if (textWidth <= maxWidth) { + return pad ? text + " ".repeat(maxWidth - textWidth) : text; + } + + const clippedEllipsis = truncateFragmentToWidth(ellipsis, maxWidth); + if (clippedEllipsis.width === 0) { + return pad ? " ".repeat(maxWidth) : ""; + } + return finalizeTruncatedResult("", 0, clippedEllipsis.text, clippedEllipsis.width, maxWidth, pad); + } + + if (isPrintableAscii(text)) { + if (text.length <= maxWidth) { + return pad ? text + " ".repeat(maxWidth - text.length) : text; + } + const targetWidth = maxWidth - ellipsisWidth; + return finalizeTruncatedResult(text.slice(0, targetWidth), targetWidth, ellipsis, ellipsisWidth, maxWidth, pad); + } + + const targetWidth = maxWidth - ellipsisWidth; + let result = ""; + let pendingAnsi = ""; + let visibleSoFar = 0; + let keptWidth = 0; + let keepContiguousPrefix = true; + let overflowed = false; + let exhaustedInput = false; + const hasAnsi = text.includes("\x1b"); + const hasTabs = text.includes("\t"); + + if (!hasAnsi && !hasTabs) { + for (const { segment } of graphemeSegmenter.segment(text)) { + const width = graphemeWidth(segment); + if (keepContiguousPrefix && keptWidth + width <= targetWidth) { + result += segment; + keptWidth += width; + } else { + keepContiguousPrefix = false; + } + visibleSoFar += width; + if (visibleSoFar > maxWidth) { + overflowed = true; + break; + } + } + exhaustedInput = !overflowed; + } else { + let i = 0; + while (i < text.length) { + const ansi = extractAnsiCode(text, i); + if (ansi) { + pendingAnsi += ansi.code; + i += ansi.length; + continue; + } + + if (text[i] === "\t") { + if (keepContiguousPrefix && keptWidth + 3 <= targetWidth) { + if (pendingAnsi) { + result += pendingAnsi; + pendingAnsi = ""; + } + result += "\t"; + keptWidth += 3; + } else { + keepContiguousPrefix = false; + pendingAnsi = ""; + } + visibleSoFar += 3; + if (visibleSoFar > maxWidth) { + overflowed = true; + break; + } + i++; + continue; + } + + let end = i; + while (end < text.length && text[end] !== "\t") { + const nextAnsi = extractAnsiCode(text, end); + if (nextAnsi) { + break; + } + end++; + } + + for (const { segment } of graphemeSegmenter.segment(text.slice(i, end))) { + const width = graphemeWidth(segment); + if (keepContiguousPrefix && keptWidth + width <= targetWidth) { + if (pendingAnsi) { + result += pendingAnsi; + pendingAnsi = ""; + } + result += segment; + keptWidth += width; + } else { + keepContiguousPrefix = false; + pendingAnsi = ""; + } + + visibleSoFar += width; + if (visibleSoFar > maxWidth) { + overflowed = true; + break; + } + } + if (overflowed) { + break; + } + i = end; + } + exhaustedInput = i >= text.length; + } + + if (!overflowed && exhaustedInput) { + return pad ? text + " ".repeat(Math.max(0, maxWidth - visibleSoFar)) : text; + } + + return finalizeTruncatedResult(result, keptWidth, ellipsis, ellipsisWidth, maxWidth, pad); +} + +/** + * Extract a range of visible columns from a line. Handles ANSI codes and wide chars. + * @param strict - If true, exclude wide chars at boundary that would extend past the range + */ +export function sliceByColumn(line: string, startCol: number, length: number, strict = false): string { + return sliceWithWidth(line, startCol, length, strict).text; +} + +/** Like sliceByColumn but also returns the actual visible width of the result. */ +export function sliceWithWidth( + line: string, + startCol: number, + length: number, + strict = false, +): { text: string; width: number } { + if (length <= 0) return { text: "", width: 0 }; + const endCol = startCol + length; + let result = "", + resultWidth = 0, + currentCol = 0, + i = 0, + pendingAnsi = ""; + + while (i < line.length) { + const ansi = extractAnsiCode(line, i); + if (ansi) { + if (currentCol >= startCol && currentCol < endCol) result += ansi.code; + else if (currentCol < startCol) pendingAnsi += ansi.code; + i += ansi.length; + continue; + } + + let textEnd = i; + while (textEnd < line.length && !extractAnsiCode(line, textEnd)) textEnd++; + + for (const { segment } of graphemeSegmenter.segment(line.slice(i, textEnd))) { + const w = graphemeWidth(segment); + const inRange = currentCol >= startCol && currentCol < endCol; + const fits = !strict || currentCol + w <= endCol; + if (inRange && fits) { + if (pendingAnsi) { + result += pendingAnsi; + pendingAnsi = ""; + } + result += segment; + resultWidth += w; + } + currentCol += w; + if (currentCol >= endCol) break; + } + i = textEnd; + if (currentCol >= endCol) break; + } + return { text: result, width: resultWidth }; +} + +// Pooled tracker instance for extractSegments (avoids allocation per call) +const pooledStyleTracker = new AnsiCodeTracker(); + +/** + * Extract "before" and "after" segments from a line in a single pass. + * Used for overlay compositing where we need content before and after the overlay region. + * Preserves styling from before the overlay that should affect content after it. + */ +export function extractSegments( + line: string, + beforeEnd: number, + afterStart: number, + afterLen: number, + strictAfter = false, +): { before: string; beforeWidth: number; after: string; afterWidth: number } { + let before = "", + beforeWidth = 0, + after = "", + afterWidth = 0; + let currentCol = 0, + i = 0; + let pendingAnsiBefore = ""; + let afterStarted = false; + const afterEnd = afterStart + afterLen; + + // Track styling state so "after" inherits styling from before the overlay + pooledStyleTracker.clear(); + + while (i < line.length) { + const ansi = extractAnsiCode(line, i); + if (ansi) { + // Track all SGR codes to know styling state at afterStart + pooledStyleTracker.process(ansi.code); + // Include ANSI codes in their respective segments + if (currentCol < beforeEnd) { + pendingAnsiBefore += ansi.code; + } else if (currentCol >= afterStart && currentCol < afterEnd && afterStarted) { + // Only include after we've started "after" (styling already prepended) + after += ansi.code; + } + i += ansi.length; + continue; + } + + let textEnd = i; + while (textEnd < line.length && !extractAnsiCode(line, textEnd)) textEnd++; + + for (const { segment } of graphemeSegmenter.segment(line.slice(i, textEnd))) { + const w = graphemeWidth(segment); + + if (currentCol < beforeEnd && currentCol + w <= beforeEnd) { + if (pendingAnsiBefore) { + before += pendingAnsiBefore; + pendingAnsiBefore = ""; + } + before += segment; + beforeWidth += w; + } else if (currentCol >= afterStart && currentCol < afterEnd) { + const fits = !strictAfter || currentCol + w <= afterEnd; + if (fits) { + // On first "after" grapheme, prepend inherited styling from before overlay + if (!afterStarted) { + after += pooledStyleTracker.getActiveCodes(); + afterStarted = true; + } + after += segment; + afterWidth += w; + } + } + + currentCol += w; + // Early exit: done with "before" only, or done with both segments + if (afterLen <= 0 ? currentCol >= beforeEnd : currentCol >= afterEnd) break; + } + i = textEnd; + if (afterLen <= 0 ? currentCol >= beforeEnd : currentCol >= afterEnd) break; + } + + return { before, beforeWidth, after, afterWidth }; +} diff --git a/cactus-code/packages/tui/src/word-navigation.ts b/cactus-code/packages/tui/src/word-navigation.ts new file mode 100644 index 000000000..7c7eced2b --- /dev/null +++ b/cactus-code/packages/tui/src/word-navigation.ts @@ -0,0 +1,117 @@ +import { getWordSegmenter, isWhitespaceChar, PUNCTUATION_REGEX } from "./utils.ts"; + +const wordSegmenter = getWordSegmenter(); + +/** + * Options for word navigation functions. + * When omitted, uses the default Intl.Segmenter word segmentation. + */ +export interface WordNavigationOptions { + /** Custom segmenter returning word segments for the given text. */ + segment?: (text: string) => Iterable; + /** Predicate identifying atomic segments that should be treated as single units (e.g. paste markers). */ + isAtomicSegment?: (segment: string) => boolean; +} + +/** + * Find the cursor position after moving one word backward from `cursor` in `text`. + * Skips trailing whitespace, then stops at the next word/punctuation boundary. + * + * Pure function - does not mutate any state. + */ +export function findWordBackward(text: string, cursor: number, options?: WordNavigationOptions): number { + if (cursor <= 0) return 0; + + const textBeforeCursor = text.slice(0, cursor); + const segmentFn = options?.segment; + const isAtomic = options?.isAtomicSegment; + const segments = segmentFn ? [...segmentFn(textBeforeCursor)] : [...wordSegmenter.segment(textBeforeCursor)]; + let newCursor = cursor; + + // Skip trailing whitespace + while ( + segments.length > 0 && + !isAtomic?.(segments[segments.length - 1]?.segment || "") && + isWhitespaceChar(segments[segments.length - 1]?.segment || "") + ) { + newCursor -= segments.pop()?.segment.length || 0; + } + + if (segments.length === 0) return newCursor; + + const last = segments[segments.length - 1]!; + + if (isAtomic?.(last.segment)) { + // Skip one atomic segment. + newCursor -= last.segment.length; + } else if (last.isWordLike) { + // Skip inside one word-like segment, preserving ASCII punctuation boundaries. + const segment = last.segment; + const matches = [...segment.matchAll(new RegExp(PUNCTUATION_REGEX, "g"))]; + if (matches.length <= 0) { + newCursor -= segment.length; + } else { + const lastMatch = matches[matches.length - 1]!; + newCursor -= segment.length - (lastMatch.index + lastMatch[0].length); + } + } else { + // Skip non-word non-whitespace run (punctuation) + while ( + segments.length > 0 && + !isAtomic?.(segments[segments.length - 1]?.segment || "") && + !segments[segments.length - 1]?.isWordLike && + !isWhitespaceChar(segments[segments.length - 1]?.segment || "") + ) { + newCursor -= segments.pop()?.segment.length || 0; + } + } + + return newCursor; +} + +/** + * Find the cursor position after moving one word forward from `cursor` in `text`. + * Skips leading whitespace, then stops at the next word/punctuation boundary. + * + * Pure function - does not mutate any state. + */ +export function findWordForward(text: string, cursor: number, options?: WordNavigationOptions): number { + if (cursor >= text.length) return text.length; + + const textAfterCursor = text.slice(cursor); + const segmentFn = options?.segment; + const isAtomic = options?.isAtomicSegment; + const segments = segmentFn ? segmentFn(textAfterCursor) : wordSegmenter.segment(textAfterCursor); + const iterator = segments[Symbol.iterator](); + let next = iterator.next(); + let newCursor = cursor; + + // Skip leading whitespace + while (!next.done && !isAtomic?.(next.value.segment) && isWhitespaceChar(next.value.segment)) { + newCursor += next.value.segment.length; + next = iterator.next(); + } + + if (next.done) return newCursor; + + if (isAtomic?.(next.value.segment)) { + // Skip one atomic segment. + newCursor += next.value.segment.length; + } else if (next.value.isWordLike) { + // Skip inside one word-like segment, preserving ASCII punctuation boundaries. + newCursor += PUNCTUATION_REGEX.exec(next.value.segment)?.index ?? next.value.segment.length; + } else { + // Skip non-word non-whitespace run (punctuation) + while ( + !next.done && + !isAtomic?.(next.value.segment) && + !next.value.isWordLike && + !isWhitespaceChar(next.value.segment) + ) { + newCursor += next.value.segment.length; + next = iterator.next(); + } + } + + return newCursor; +} diff --git a/cactus-code/packages/tui/tsconfig.build.json b/cactus-code/packages/tui/tsconfig.build.json new file mode 100644 index 000000000..5ce430296 --- /dev/null +++ b/cactus-code/packages/tui/tsconfig.build.json @@ -0,0 +1,9 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "rootDir": "./src" + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist"] +} \ No newline at end of file diff --git a/cactus-code/scripts/bundle-for-pypi.sh b/cactus-code/scripts/bundle-for-pypi.sh new file mode 100644 index 000000000..4455b2986 --- /dev/null +++ b/cactus-code/scripts/bundle-for-pypi.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# +# Bundle the built Cactus coding agent into the Python package so it ships in the +# pypi wheel and `cactus code` works after `pip install cactus-compute`. +# +# Produces a self-contained payload at python/cactus/code/ with the four built +# packages plus a production-only node_modules (workspace packages copied in via +# --install-links, so it's portable — no symlinks back into the repo). +# +# Requires: node >=22, npm. Run from anywhere. +set -euo pipefail + +CC_ROOT="$(cd "$(dirname "$0")/.." && pwd)" # cactus-code/ +OUT="$(cd "$CC_ROOT/.." && pwd)/python/cactus/code" # python/cactus/code/ + +echo "==> Building cactus-code" +cd "$CC_ROOT" +npm run build + +echo "==> Staging payload at $OUT" +rm -rf "$OUT" +mkdir -p "$OUT/packages" +for p in ai agent tui coding-agent; do + mkdir -p "$OUT/packages/$p" + cp -R "packages/$p/dist" "$OUT/packages/$p/dist" + cp "packages/$p/package.json" "$OUT/packages/$p/package.json" +done +# tui ships prebuilt native bindings (not part of dist); copy them so the TUI +# loads its key-modifier addon after a pip install. +if [ -d "packages/tui/native" ]; then + mkdir -p "$OUT/packages/tui/native" + cp -R "packages/tui/native" "$OUT/packages/tui/" +fi +cp package.json "$OUT/package.json" +[ -f package-lock.json ] && cp package-lock.json "$OUT/package-lock.json" || true +# Ship the Cactus version so the agent can display it after a pip install. +[ -f "$CC_ROOT/../CACTUS_VERSION" ] && cp "$CC_ROOT/../CACTUS_VERSION" "$OUT/CACTUS_VERSION" || true + +echo "==> Installing production dependencies" +cd "$OUT" +# --install-links copies workspace packages as real dirs; --omit=dev drops +# typescript/tsgo/etc. Optional deps (the native clipboard module used by /copy +# and image paste) are kept so those work; the wheel is platform-specific anyway +# since the Cactus engine ships native libraries per platform. +# --ignore-scripts: the bundled root package.json carries dev lifecycle hooks +# (e.g. `prepare: husky`) that must not run when installing the wheel payload. +npm install --omit=dev --install-links --no-audit --no-fund --ignore-scripts + +echo "==> Making payload wheel-safe (dereference symlinks, drop .bin)" +# Workspace packages get symlinked by npm; wheels don't carry symlinks reliably, +# so replace them with real directory copies. +if [ -d "$OUT/node_modules/@earendil-works" ]; then + for d in "$OUT/node_modules/@earendil-works"/*; do + if [ -L "$d" ]; then + real="$(cd "$d" && pwd -P)" + rm "$d" + cp -R "$real" "$d" + fi + done +fi +# .bin holds only symlinks to dev tools; not needed at runtime. +rm -rf "$OUT/node_modules/.bin" +# Drop any other stray symlinks so the wheel is fully self-contained. +find "$OUT" -type l -delete 2>/dev/null || true + +echo "==> Done. Bundled agent at: $OUT" +echo " entry: $OUT/packages/coding-agent/dist/cli.js" diff --git a/cactus-code/tsconfig.base.json b/cactus-code/tsconfig.base.json new file mode 100644 index 000000000..57e97d6e3 --- /dev/null +++ b/cactus-code/tsconfig.base.json @@ -0,0 +1,25 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "Node16", + "lib": ["ES2022"], + "strict": true, + "erasableSyntaxOnly": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "declaration": true, + "declarationMap": true, + "sourceMap": true, + "inlineSources": true, + "inlineSourceMap": false, + "moduleResolution": "Node16", + "resolveJsonModule": true, + "allowImportingTsExtensions": true, + "rewriteRelativeImportExtensions": true, + "experimentalDecorators": true, + "emitDecoratorMetadata": true, + "useDefineForClassFields": false, + "types": ["node"] + } +} diff --git a/cactus-engine/CMakeLists.txt b/cactus-engine/CMakeLists.txt new file mode 100644 index 000000000..ce09c48fe --- /dev/null +++ b/cactus-engine/CMakeLists.txt @@ -0,0 +1,162 @@ +cmake_minimum_required(VERSION 3.10) +project(CactusEngine LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT TARGET cactus_graph) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../cactus-graph ${CMAKE_CURRENT_BINARY_DIR}/cactus-graph) +endif() + +set(CACTUS_CURL_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/libs/curl" CACHE PATH "Path to vendored libcurl") + +set(ENGINE_SOURCES + src/bpe.cpp + src/sp.cpp + src/constraints.cpp + src/model.cpp + src/kv_compress.cpp + src/engine_image.cpp + src/index.cpp + src/index_ffi.cpp + src/log.cpp + src/rag.cpp + src/telemetry.cpp + src/telemetry_impl.cpp + src/cloud.cpp + src/init.cpp + src/embed.cpp + src/complete.cpp + src/transcribe.cpp + src/stream.cpp + src/tokenizer.cpp +) +add_library(cactus_engine STATIC ${ENGINE_SOURCES}) +set_target_properties(cactus_engine PROPERTIES + POSITION_INDEPENDENT_CODE ON + OUTPUT_NAME cactus_engine_core +) + +target_include_directories(cactus_engine + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/libs + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src +) + +target_link_libraries(cactus_engine PUBLIC cactus_graph) + +if(APPLE) + target_link_libraries(cactus_engine PUBLIC + "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" + ) + + if(CMAKE_SYSTEM_NAME STREQUAL "iOS") + string(TOLOWER "${CMAKE_OSX_SYSROOT}" _sysroot) + if(_sysroot MATCHES "simulator") + set(_curl "${CACTUS_CURL_ROOT}/ios/simulator/libcurl.a") + else() + set(_curl "${CACTUS_CURL_ROOT}/ios/device/libcurl.a") + endif() + else() + set(_curl "${CACTUS_CURL_ROOT}/macos/libcurl.a") + endif() + if(EXISTS "${_curl}") + target_link_libraries(cactus_engine PRIVATE "${_curl}") + target_include_directories(cactus_engine PRIVATE "${CACTUS_CURL_ROOT}/include") + target_compile_definitions(cactus_engine PRIVATE CACTUS_USE_CURL=1) + endif() +elseif(ANDROID) + set(_mbedtls ${CMAKE_CURRENT_SOURCE_DIR}/../android/mbedtls/${ANDROID_ABI}/lib) + target_link_libraries(cactus_engine PUBLIC + ${_mbedtls}/libmbedtls.a + ${_mbedtls}/libmbedx509.a + ${_mbedtls}/libmbedcrypto.a + log + android + ) + + set(_curl "${CACTUS_CURL_ROOT}/android/${ANDROID_ABI}/libcurl.a") + if(EXISTS "${_curl}") + target_link_libraries(cactus_engine PRIVATE "${_curl}") + target_include_directories(cactus_engine PRIVATE "${CACTUS_CURL_ROOT}/include") + target_compile_definitions(cactus_engine PRIVATE CACTUS_USE_CURL=1) + endif() +else() + find_package(CURL QUIET) + if(CURL_FOUND) + target_link_libraries(cactus_engine PRIVATE ${CURL_LIBRARIES}) + target_include_directories(cactus_engine PRIVATE ${CURL_INCLUDE_DIRS}) + target_compile_definitions(cactus_engine PRIVATE CACTUS_USE_CURL=1) + endif() +endif() + +set(_version_file "${CMAKE_CURRENT_SOURCE_DIR}/../CACTUS_VERSION") +if(EXISTS "${_version_file}") + file(READ "${_version_file}" _version) + string(STRIP "${_version}" _version) + target_compile_definitions(cactus_engine PRIVATE CACTUS_COMPILE_TIME_VERSION="${_version}") +endif() + +target_compile_options(cactus_engine PRIVATE + -Wall -Wextra -pedantic -Wno-missing-field-initializers + $<$:-Wno-c99-extensions> + $<$:-Wno-pedantic> +) + +set(_CACTUS_ENGINE_BUNDLE "${CMAKE_BINARY_DIR}/libcactus_engine.a") +set(_CACTUS_ENGINE_ARCHIVES + $ + $ + $ +) +set(_CACTUS_ENGINE_DEPS cactus_engine cactus_graph cactus_kernels) + +if(APPLE) + add_custom_target(cactus_engine_bundle ALL + COMMAND libtool -static -o ${_CACTUS_ENGINE_BUNDLE} ${_CACTUS_ENGINE_ARCHIVES} + DEPENDS ${_CACTUS_ENGINE_DEPS} + BYPRODUCTS ${_CACTUS_ENGINE_BUNDLE} + ) +else() + file(GENERATE OUTPUT ${CMAKE_BINARY_DIR}/cactus_engine_bundle.mri CONTENT + "create ${_CACTUS_ENGINE_BUNDLE}\naddlib $\naddlib $\naddlib $\nsave\nend\n" + ) + add_custom_target(cactus_engine_bundle ALL + COMMAND ${CMAKE_AR} -M < ${CMAKE_BINARY_DIR}/cactus_engine_bundle.mri + DEPENDS ${_CACTUS_ENGINE_DEPS} + BYPRODUCTS ${_CACTUS_ENGINE_BUNDLE} + ) +endif() + +if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + file(WRITE ${CMAKE_BINARY_DIR}/cactus_engine_shared_stub.cpp "") + add_library(cactus_engine_shared SHARED ${CMAKE_BINARY_DIR}/cactus_engine_shared_stub.cpp) + set_target_properties(cactus_engine_shared PROPERTIES OUTPUT_NAME cactus_engine) + add_dependencies(cactus_engine_shared ${_CACTUS_ENGINE_DEPS}) + target_link_libraries(cactus_engine_shared PRIVATE cactus_flags) + + if(APPLE) + target_link_options(cactus_engine_shared PRIVATE + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + ) + set(_curl "${CACTUS_CURL_ROOT}/macos/libcurl.a") + if(EXISTS "${_curl}") + target_link_libraries(cactus_engine_shared PRIVATE "${_curl}") + endif() + target_link_libraries(cactus_engine_shared PRIVATE + "-framework Accelerate" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" + ) + else() + target_link_options(cactus_engine_shared PRIVATE + -Wl,--whole-archive ${_CACTUS_ENGINE_ARCHIVES} -Wl,--no-whole-archive + ) + find_package(CURL QUIET) + if(CURL_FOUND) + target_link_libraries(cactus_engine_shared PRIVATE ${CURL_LIBRARIES}) + endif() + endif() +endif() diff --git a/cactus-engine/build.sh b/cactus-engine/build.sh new file mode 100755 index 000000000..68177f692 --- /dev/null +++ b/cactus-engine/build.sh @@ -0,0 +1,16 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")" + +echo "Building cactus-engine..." + +mkdir -p build +cd build + +cmake .. -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null +make -j$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) + +echo "cactus-engine built successfully!" +echo " Static: $(pwd)/libcactus_engine.a" +echo " Shared: $(pwd)/libcactus_engine.$([ "$(uname)" = Darwin ] && echo dylib || echo so)" diff --git a/cactus/ffi/cactus_ffi.h b/cactus-engine/cactus_engine.h similarity index 63% rename from cactus/ffi/cactus_ffi.h rename to cactus-engine/cactus_engine.h index e8a77f343..a60e4aec2 100644 --- a/cactus/ffi/cactus_ffi.h +++ b/cactus-engine/cactus_engine.h @@ -1,14 +1,12 @@ #ifndef CACTUS_FFI_H #define CACTUS_FFI_H +#include "cactus_graph.h" + #include #include #include -// ---------------------------------------------------------------------------- -// Export Macros -// ---------------------------------------------------------------------------- - #if __GNUC__ >= 4 #define CACTUS_FFI_EXPORT __attribute__((visibility("default"))) #define CACTUS_FFI_LOCAL __attribute__((visibility("hidden"))) @@ -21,32 +19,22 @@ extern "C" { #endif -// ---------------------------------------------------------------------------- -// Types -// ---------------------------------------------------------------------------- - typedef void* cactus_model_t; typedef void* cactus_index_t; typedef void* cactus_stream_transcribe_t; - typedef void (*cactus_token_callback)(const char* token, uint32_t token_id, void* user_data); -// ---------------------------------------------------------------------------- -// Model Lifecycle -// ---------------------------------------------------------------------------- - CACTUS_FFI_EXPORT cactus_model_t cactus_init( const char* model_path, - const char* corpus_dir // optional: NULL if no RAG corpus + const char* corpus_dir, // optional: NULL if no RAG corpus + bool cache_index // false = always rebuild index, true = load cached if available ); CACTUS_FFI_EXPORT void cactus_destroy(cactus_model_t model); CACTUS_FFI_EXPORT void cactus_reset(cactus_model_t model); CACTUS_FFI_EXPORT void cactus_stop(cactus_model_t model); -// ---------------------------------------------------------------------------- -// Text Completion -// ---------------------------------------------------------------------------- +CACTUS_FFI_EXPORT int cactus_set_backend(const char* backend); CACTUS_FFI_EXPORT int cactus_complete( cactus_model_t model, @@ -56,12 +44,21 @@ CACTUS_FFI_EXPORT int cactus_complete( const char* options_json, // optional const char* tools_json, // optional cactus_token_callback callback, // optional - void* user_data // optional + void* user_data, // optional + const uint8_t* pcm_buffer, // optional: NULL when not used + size_t pcm_buffer_size // optional: 0 when not used ); -// ---------------------------------------------------------------------------- -// Tokenization -// ---------------------------------------------------------------------------- +CACTUS_FFI_EXPORT int cactus_prefill( + cactus_model_t model, + const char* messages_json, + char* response_buffer, + size_t buffer_size, + const char* options_json, // optional + const char* tools_json, // optional + const uint8_t* pcm_buffer, // optional: NULL when not used + size_t pcm_buffer_size // optional: 0 when not used +); CACTUS_FFI_EXPORT int cactus_tokenize( cactus_model_t model, @@ -71,6 +68,15 @@ CACTUS_FFI_EXPORT int cactus_tokenize( size_t* out_token_len ); +CACTUS_FFI_EXPORT int cactus_render_prompt( + cactus_model_t model, + const char* messages_json, + const char* options_json, + const char* tools_json, + char* prompt_buffer, + size_t buffer_size +); + CACTUS_FFI_EXPORT int cactus_score_window( cactus_model_t model, const uint32_t* tokens, @@ -82,9 +88,14 @@ CACTUS_FFI_EXPORT int cactus_score_window( size_t buffer_size ); -// ---------------------------------------------------------------------------- -// Audio Transcription -// ---------------------------------------------------------------------------- +CACTUS_FFI_EXPORT int cactus_benchmark_tokens( + cactus_model_t model, + const uint32_t* prompt_tokens, + size_t prompt_token_len, + size_t decode_token_len, + char* response_buffer, + size_t buffer_size +); CACTUS_FFI_EXPORT int cactus_transcribe( cactus_model_t model, @@ -99,41 +110,36 @@ CACTUS_FFI_EXPORT int cactus_transcribe( size_t pcm_buffer_size ); -// ---------------------------------------------------------------------------- -// Streaming Transcription -// ---------------------------------------------------------------------------- - -CACTUS_FFI_EXPORT cactus_stream_transcribe_t cactus_stream_transcribe_init( - cactus_model_t model +CACTUS_FFI_EXPORT int cactus_preprocess_audio_features( + const char* audio_file_path, + const char* model_type, + size_t mel_bins, + float* features_buffer, + size_t buffer_size, + size_t* feature_count, + size_t* out_mel_bins, + size_t* out_frames ); -CACTUS_FFI_EXPORT int cactus_stream_transcribe_insert( - cactus_stream_transcribe_t stream, - const uint8_t* pcm_buffer, - size_t pcm_buffer_size +CACTUS_FFI_EXPORT cactus_stream_transcribe_t cactus_stream_transcribe_start( + cactus_model_t model, + const char* options_json ); CACTUS_FFI_EXPORT int cactus_stream_transcribe_process( cactus_stream_transcribe_t stream, + const uint8_t* pcm_buffer, + size_t pcm_buffer_size, char* response_buffer, - size_t buffer_size, - const char* options_json // optional + size_t buffer_size ); -CACTUS_FFI_EXPORT int cactus_stream_transcribe_finalize( +CACTUS_FFI_EXPORT int cactus_stream_transcribe_stop( cactus_stream_transcribe_t stream, char* response_buffer, size_t buffer_size ); -CACTUS_FFI_EXPORT void cactus_stream_transcribe_destroy( - cactus_stream_transcribe_t stream -); - -// ---------------------------------------------------------------------------- -// Embeddings -// ---------------------------------------------------------------------------- - CACTUS_FFI_EXPORT int cactus_embed( cactus_model_t model, const char* text, @@ -159,10 +165,6 @@ CACTUS_FFI_EXPORT int cactus_audio_embed( size_t* embedding_dim ); -// ---------------------------------------------------------------------------- -// RAG (Retrieval-Augmented Generation) -// ---------------------------------------------------------------------------- - CACTUS_FFI_EXPORT int cactus_rag_query( cactus_model_t model, const char* query, @@ -171,10 +173,6 @@ CACTUS_FFI_EXPORT int cactus_rag_query( size_t top_k ); -// ---------------------------------------------------------------------------- -// Vector Index -// ---------------------------------------------------------------------------- - CACTUS_FFI_EXPORT cactus_index_t cactus_index_init( const char* index_dir, size_t embedding_dim @@ -223,13 +221,18 @@ CACTUS_FFI_EXPORT int cactus_index_query( CACTUS_FFI_EXPORT int cactus_index_compact(cactus_index_t index); CACTUS_FFI_EXPORT void cactus_index_destroy(cactus_index_t index); -// ---------------------------------------------------------------------------- -// Utilities -// ---------------------------------------------------------------------------- - CACTUS_FFI_EXPORT const char* cactus_get_last_error(void); -CACTUS_FFI_EXPORT void cactus_set_telemetry_token(const char* token); -CACTUS_FFI_EXPORT void cactus_set_pro_key(const char* pro_key); + +// level: 0=DEBUG, 1=INFO, 2=WARN (default), 3=ERROR, 4=NONE +CACTUS_FFI_EXPORT void cactus_log_set_level(int level); + +typedef void (*cactus_log_callback_t)(int level, const char* component, const char* message, void* user_data); +CACTUS_FFI_EXPORT void cactus_log_set_callback(cactus_log_callback_t callback, void* user_data); + +CACTUS_FFI_EXPORT void cactus_set_telemetry_environment(const char* framework, const char* cache_location, const char* version); +CACTUS_FFI_EXPORT void cactus_set_app_id(const char* app_id); +CACTUS_FFI_EXPORT void cactus_telemetry_flush(void); +CACTUS_FFI_EXPORT void cactus_telemetry_shutdown(void); #ifdef __cplusplus } diff --git a/cactus-engine/libs/curl/android/arm64-v8a/libcurl.a b/cactus-engine/libs/curl/android/arm64-v8a/libcurl.a new file mode 100644 index 000000000..959af061a Binary files /dev/null and b/cactus-engine/libs/curl/android/arm64-v8a/libcurl.a differ diff --git a/cactus-engine/libs/curl/include/Makefile.am b/cactus-engine/libs/curl/include/Makefile.am new file mode 100644 index 000000000..d65bfeaa0 --- /dev/null +++ b/cactus-engine/libs/curl/include/Makefile.am @@ -0,0 +1,28 @@ +#*************************************************************************** +# _ _ ____ _ +# Project ___| | | | _ \| | +# / __| | | | |_) | | +# | (__| |_| | _ <| |___ +# \___|\___/|_| \_\_____| +# +# Copyright (C) Daniel Stenberg, , et al. +# +# This software is licensed as described in the file COPYING, which +# you should have received as part of this distribution. The terms +# are also available at https://curl.se/docs/copyright.html. +# +# You may opt to use, copy, modify, merge, publish, distribute and/or sell +# copies of the Software, and permit persons to whom the Software is +# furnished to do so, under the terms of the COPYING file. +# +# This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY +# KIND, either express or implied. +# +# SPDX-License-Identifier: curl +# +########################################################################### +SUBDIRS = curl + +EXTRA_DIST = README.md + +AUTOMAKE_OPTIONS = foreign no-dependencies diff --git a/cactus-engine/libs/curl/include/Makefile.in b/cactus-engine/libs/curl/include/Makefile.in new file mode 100644 index 000000000..e7e6e7582 --- /dev/null +++ b/cactus-engine/libs/curl/include/Makefile.in @@ -0,0 +1,713 @@ +# Makefile.in generated by automake 1.16.5 from Makefile.am. +# @configure_input@ + +# Copyright (C) 1994-2021 Free Software Foundation, Inc. + +# This Makefile.in is free software; the Free Software Foundation +# gives unlimited permission to copy and/or distribute it, +# with or without modifications, as long as this notice is preserved. + +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY, to the extent permitted by law; without +# even the implied warranty of MERCHANTABILITY or FITNESS FOR A +# PARTICULAR PURPOSE. + +@SET_MAKE@ +VPATH = @srcdir@ +am__is_gnu_make = { \ + if test -z '$(MAKELEVEL)'; then \ + false; \ + elif test -n '$(MAKE_HOST)'; then \ + true; \ + elif test -n '$(MAKE_VERSION)' && test -n '$(CURDIR)'; then \ + true; \ + else \ + false; \ + fi; \ +} +am__make_running_with_option = \ + case $${target_option-} in \ + ?) ;; \ + *) echo "am__make_running_with_option: internal error: invalid" \ + "target option '$${target_option-}' specified" >&2; \ + exit 1;; \ + esac; \ + has_opt=no; \ + sane_makeflags=$$MAKEFLAGS; \ + if $(am__is_gnu_make); then \ + sane_makeflags=$$MFLAGS; \ + else \ + case $$MAKEFLAGS in \ + *\\[\ \ ]*) \ + bs=\\; \ + sane_makeflags=`printf '%s\n' "$$MAKEFLAGS" \ + | sed "s/$$bs$$bs[$$bs $$bs ]*//g"`;; \ + esac; \ + fi; \ + skip_next=no; \ + strip_trailopt () \ + { \ + flg=`printf '%s\n' "$$flg" | sed "s/$$1.*$$//"`; \ + }; \ + for flg in $$sane_makeflags; do \ + test $$skip_next = yes && { skip_next=no; continue; }; \ + case $$flg in \ + *=*|--*) continue;; \ + -*I) strip_trailopt 'I'; skip_next=yes;; \ + -*I?*) strip_trailopt 'I';; \ + -*O) strip_trailopt 'O'; skip_next=yes;; \ + -*O?*) strip_trailopt 'O';; \ + -*l) strip_trailopt 'l'; skip_next=yes;; \ + -*l?*) strip_trailopt 'l';; \ + -[dEDm]) skip_next=yes;; \ + -[JT]) skip_next=yes;; \ + esac; \ + case $$flg in \ + *$$target_option*) has_opt=yes; break;; \ + esac; \ + done; \ + test $$has_opt = yes +am__make_dryrun = (target_option=n; $(am__make_running_with_option)) +am__make_keepgoing = (target_option=k; $(am__make_running_with_option)) +pkgdatadir = $(datadir)/@PACKAGE@ +pkgincludedir = $(includedir)/@PACKAGE@ +pkglibdir = $(libdir)/@PACKAGE@ +pkglibexecdir = $(libexecdir)/@PACKAGE@ +am__cd = CDPATH="$${ZSH_VERSION+.}$(PATH_SEPARATOR)" && cd +install_sh_DATA = $(install_sh) -c -m 644 +install_sh_PROGRAM = $(install_sh) -c +install_sh_SCRIPT = $(install_sh) -c +INSTALL_HEADER = $(INSTALL_DATA) +transform = $(program_transform_name) +NORMAL_INSTALL = : +PRE_INSTALL = : +POST_INSTALL = : +NORMAL_UNINSTALL = : +PRE_UNINSTALL = : +POST_UNINSTALL = : +build_triplet = @build@ +host_triplet = @host@ +subdir = include +ACLOCAL_M4 = $(top_srcdir)/aclocal.m4 +am__aclocal_m4_deps = $(top_srcdir)/m4/curl-amissl.m4 \ + $(top_srcdir)/m4/curl-apple-sectrust.m4 \ + $(top_srcdir)/m4/curl-compilers.m4 \ + $(top_srcdir)/m4/curl-confopts.m4 \ + $(top_srcdir)/m4/curl-functions.m4 \ + $(top_srcdir)/m4/curl-gnutls.m4 \ + $(top_srcdir)/m4/curl-mbedtls.m4 \ + $(top_srcdir)/m4/curl-openssl.m4 \ + $(top_srcdir)/m4/curl-override.m4 \ + $(top_srcdir)/m4/curl-reentrant.m4 \ + $(top_srcdir)/m4/curl-rustls.m4 \ + $(top_srcdir)/m4/curl-schannel.m4 \ + $(top_srcdir)/m4/curl-sysconfig.m4 \ + $(top_srcdir)/m4/curl-wolfssl.m4 $(top_srcdir)/m4/libtool.m4 \ + $(top_srcdir)/m4/ltoptions.m4 $(top_srcdir)/m4/ltsugar.m4 \ + $(top_srcdir)/m4/ltversion.m4 $(top_srcdir)/m4/lt~obsolete.m4 \ + $(top_srcdir)/m4/xc-am-iface.m4 \ + $(top_srcdir)/m4/xc-cc-check.m4 \ + $(top_srcdir)/m4/xc-lt-iface.m4 \ + $(top_srcdir)/m4/xc-val-flgs.m4 \ + $(top_srcdir)/m4/zz40-xc-ovr.m4 \ + $(top_srcdir)/m4/zz50-xc-ovr.m4 $(top_srcdir)/acinclude.m4 \ + $(top_srcdir)/configure.ac +am__configure_deps = $(am__aclocal_m4_deps) $(CONFIGURE_DEPENDENCIES) \ + $(ACLOCAL_M4) +DIST_COMMON = $(srcdir)/Makefile.am $(am__DIST_COMMON) +mkinstalldirs = $(install_sh) -d +CONFIG_HEADER = $(top_builddir)/lib/curl_config.h +CONFIG_CLEAN_FILES = +CONFIG_CLEAN_VPATH_FILES = +AM_V_P = $(am__v_P_@AM_V@) +am__v_P_ = $(am__v_P_@AM_DEFAULT_V@) +am__v_P_0 = false +am__v_P_1 = : +AM_V_GEN = $(am__v_GEN_@AM_V@) +am__v_GEN_ = $(am__v_GEN_@AM_DEFAULT_V@) +am__v_GEN_0 = @echo " GEN " $@; +am__v_GEN_1 = +AM_V_at = $(am__v_at_@AM_V@) +am__v_at_ = $(am__v_at_@AM_DEFAULT_V@) +am__v_at_0 = @ +am__v_at_1 = +depcomp = +am__maybe_remake_depfiles = +SOURCES = +DIST_SOURCES = +RECURSIVE_TARGETS = all-recursive check-recursive cscopelist-recursive \ + ctags-recursive dvi-recursive html-recursive info-recursive \ + install-data-recursive install-dvi-recursive \ + install-exec-recursive install-html-recursive \ + install-info-recursive install-pdf-recursive \ + install-ps-recursive install-recursive installcheck-recursive \ + installdirs-recursive pdf-recursive ps-recursive \ + tags-recursive uninstall-recursive +am__can_run_installinfo = \ + case $$AM_UPDATE_INFO_DIR in \ + n|no|NO) false;; \ + *) (install-info --version) >/dev/null 2>&1;; \ + esac +RECURSIVE_CLEAN_TARGETS = mostlyclean-recursive clean-recursive \ + distclean-recursive maintainer-clean-recursive +am__recursive_targets = \ + $(RECURSIVE_TARGETS) \ + $(RECURSIVE_CLEAN_TARGETS) \ + $(am__extra_recursive_targets) +AM_RECURSIVE_TARGETS = $(am__recursive_targets:-recursive=) TAGS CTAGS \ + distdir distdir-am +am__tagged_files = $(HEADERS) $(SOURCES) $(TAGS_FILES) $(LISP) +# Read a list of newline-separated strings from the standard input, +# and print each of them once, without duplicates. Input order is +# *not* preserved. +am__uniquify_input = $(AWK) '\ + BEGIN { nonempty = 0; } \ + { items[$$0] = 1; nonempty = 1; } \ + END { if (nonempty) { for (i in items) print i; }; } \ +' +# Make sure the list of sources is unique. This is necessary because, +# e.g., the same source file might be shared among _SOURCES variables +# for different programs/libraries. +am__define_uniq_tagged_files = \ + list='$(am__tagged_files)'; \ + unique=`for i in $$list; do \ + if test -f "$$i"; then echo $$i; else echo $(srcdir)/$$i; fi; \ + done | $(am__uniquify_input)` +DIST_SUBDIRS = $(SUBDIRS) +am__DIST_COMMON = $(srcdir)/Makefile.in README.md +DISTFILES = $(DIST_COMMON) $(DIST_SOURCES) $(TEXINFOS) $(EXTRA_DIST) +am__relativize = \ + dir0=`pwd`; \ + sed_first='s,^\([^/]*\)/.*$$,\1,'; \ + sed_rest='s,^[^/]*/*,,'; \ + sed_last='s,^.*/\([^/]*\)$$,\1,'; \ + sed_butlast='s,/*[^/]*$$,,'; \ + while test -n "$$dir1"; do \ + first=`echo "$$dir1" | sed -e "$$sed_first"`; \ + if test "$$first" != "."; then \ + if test "$$first" = ".."; then \ + dir2=`echo "$$dir0" | sed -e "$$sed_last"`/"$$dir2"; \ + dir0=`echo "$$dir0" | sed -e "$$sed_butlast"`; \ + else \ + first2=`echo "$$dir2" | sed -e "$$sed_first"`; \ + if test "$$first2" = "$$first"; then \ + dir2=`echo "$$dir2" | sed -e "$$sed_rest"`; \ + else \ + dir2="../$$dir2"; \ + fi; \ + dir0="$$dir0"/"$$first"; \ + fi; \ + fi; \ + dir1=`echo "$$dir1" | sed -e "$$sed_rest"`; \ + done; \ + reldir="$$dir2" +ACLOCAL = @ACLOCAL@ +AMTAR = @AMTAR@ +AM_DEFAULT_VERBOSITY = @AM_DEFAULT_VERBOSITY@ +APXS = @APXS@ +AR = @AR@ +AR_FLAGS = @AR_FLAGS@ +AS = @AS@ +AUTOCONF = @AUTOCONF@ +AUTOHEADER = @AUTOHEADER@ +AUTOMAKE = @AUTOMAKE@ +AWK = @AWK@ +BLANK_AT_MAKETIME = @BLANK_AT_MAKETIME@ +CADDY = @CADDY@ +CC = @CC@ +CCDEPMODE = @CCDEPMODE@ +CFLAGS = @CFLAGS@ +CFLAG_CURL_SYMBOL_HIDING = @CFLAG_CURL_SYMBOL_HIDING@ +CONFIGURE_OPTIONS = @CONFIGURE_OPTIONS@ +CPP = @CPP@ +CPPFLAGS = @CPPFLAGS@ +CSCOPE = @CSCOPE@ +CTAGS = @CTAGS@ +CURLVERSION = @CURLVERSION@ +CURL_CA_BUNDLE = @CURL_CA_BUNDLE@ +CURL_CA_EMBED = @CURL_CA_EMBED@ +CURL_CFLAG_EXTRAS = @CURL_CFLAG_EXTRAS@ +CURL_CPP = @CURL_CPP@ +CURL_LIBCURL_VERSIONED_SYMBOLS_PREFIX = @CURL_LIBCURL_VERSIONED_SYMBOLS_PREFIX@ +CURL_LIBCURL_VERSIONED_SYMBOLS_SONAME = @CURL_LIBCURL_VERSIONED_SYMBOLS_SONAME@ +CURL_NETWORK_AND_TIME_LIBS = @CURL_NETWORK_AND_TIME_LIBS@ +CYGPATH_W = @CYGPATH_W@ +DANTED = @DANTED@ +DEFS = @DEFS@ +DEPDIR = @DEPDIR@ +DLLTOOL = @DLLTOOL@ +DSYMUTIL = @DSYMUTIL@ +DUMPBIN = @DUMPBIN@ +ECHO_C = @ECHO_C@ +ECHO_N = @ECHO_N@ +ECHO_T = @ECHO_T@ +EGREP = @EGREP@ +ENABLE_SHARED = @ENABLE_SHARED@ +ENABLE_STATIC = @ENABLE_STATIC@ +ETAGS = @ETAGS@ +EXEEXT = @EXEEXT@ +FGREP = @FGREP@ +FILECMD = @FILECMD@ +FISH_FUNCTIONS_DIR = @FISH_FUNCTIONS_DIR@ +GCOV = @GCOV@ +GREP = @GREP@ +HAVE_LIBZ = @HAVE_LIBZ@ +HTTPD = @HTTPD@ +HTTPD_NGHTTPX = @HTTPD_NGHTTPX@ +INSTALL = @INSTALL@ +INSTALL_DATA = @INSTALL_DATA@ +INSTALL_PROGRAM = @INSTALL_PROGRAM@ +INSTALL_SCRIPT = @INSTALL_SCRIPT@ +INSTALL_STRIP_PROGRAM = @INSTALL_STRIP_PROGRAM@ +LCOV = @LCOV@ +LD = @LD@ +LDFLAGS = @LDFLAGS@ +LIBCURL_PC_CFLAGS = @LIBCURL_PC_CFLAGS@ +LIBCURL_PC_CFLAGS_PRIVATE = @LIBCURL_PC_CFLAGS_PRIVATE@ +LIBCURL_PC_LDFLAGS_PRIVATE = @LIBCURL_PC_LDFLAGS_PRIVATE@ +LIBCURL_PC_LIBS = @LIBCURL_PC_LIBS@ +LIBCURL_PC_LIBS_PRIVATE = @LIBCURL_PC_LIBS_PRIVATE@ +LIBCURL_PC_REQUIRES = @LIBCURL_PC_REQUIRES@ +LIBCURL_PC_REQUIRES_PRIVATE = @LIBCURL_PC_REQUIRES_PRIVATE@ +LIBOBJS = @LIBOBJS@ +LIBS = @LIBS@ +LIBTOOL = @LIBTOOL@ +LIPO = @LIPO@ +LN_S = @LN_S@ +LTLIBOBJS = @LTLIBOBJS@ +LT_SYS_LIBRARY_PATH = @LT_SYS_LIBRARY_PATH@ +MAINT = @MAINT@ +MAKEINFO = @MAKEINFO@ +MANIFEST_TOOL = @MANIFEST_TOOL@ +MKDIR_P = @MKDIR_P@ +NM = @NM@ +NMEDIT = @NMEDIT@ +OBJDUMP = @OBJDUMP@ +OBJEXT = @OBJEXT@ +OTOOL = @OTOOL@ +OTOOL64 = @OTOOL64@ +PACKAGE = @PACKAGE@ +PACKAGE_BUGREPORT = @PACKAGE_BUGREPORT@ +PACKAGE_NAME = @PACKAGE_NAME@ +PACKAGE_STRING = @PACKAGE_STRING@ +PACKAGE_TARNAME = @PACKAGE_TARNAME@ +PACKAGE_URL = @PACKAGE_URL@ +PACKAGE_VERSION = @PACKAGE_VERSION@ +PATH_SEPARATOR = @PATH_SEPARATOR@ +PERL = @PERL@ +PKGCONFIG = @PKGCONFIG@ +RANLIB = @RANLIB@ +RC = @RC@ +SED = @SED@ +SET_MAKE = @SET_MAKE@ +SFTPD = @SFTPD@ +SHELL = @SHELL@ +SSHD = @SSHD@ +SSL_BACKENDS = @SSL_BACKENDS@ +STRIP = @STRIP@ +SUPPORT_FEATURES = @SUPPORT_FEATURES@ +SUPPORT_PROTOCOLS = @SUPPORT_PROTOCOLS@ +TEST_NGHTTPX = @TEST_NGHTTPX@ +VERSION = @VERSION@ +VERSIONNUM = @VERSIONNUM@ +VSFTPD = @VSFTPD@ +ZLIB_LIBS = @ZLIB_LIBS@ +ZSH_FUNCTIONS_DIR = @ZSH_FUNCTIONS_DIR@ +abs_builddir = @abs_builddir@ +abs_srcdir = @abs_srcdir@ +abs_top_builddir = @abs_top_builddir@ +abs_top_srcdir = @abs_top_srcdir@ +ac_ct_AR = @ac_ct_AR@ +ac_ct_CC = @ac_ct_CC@ +ac_ct_DUMPBIN = @ac_ct_DUMPBIN@ +am__include = @am__include@ +am__leading_dot = @am__leading_dot@ +am__quote = @am__quote@ +am__tar = @am__tar@ +am__untar = @am__untar@ +bindir = @bindir@ +build = @build@ +build_alias = @build_alias@ +build_cpu = @build_cpu@ +build_os = @build_os@ +build_vendor = @build_vendor@ +builddir = @builddir@ +datadir = @datadir@ +datarootdir = @datarootdir@ +docdir = @docdir@ +dvidir = @dvidir@ +exec_prefix = @exec_prefix@ +host = @host@ +host_alias = @host_alias@ +host_cpu = @host_cpu@ +host_os = @host_os@ +host_vendor = @host_vendor@ +htmldir = @htmldir@ +includedir = @includedir@ +infodir = @infodir@ +install_sh = @install_sh@ +libdir = @libdir@ +libexecdir = @libexecdir@ +libext = @libext@ +localedir = @localedir@ +localstatedir = @localstatedir@ +mandir = @mandir@ +mkdir_p = @mkdir_p@ +oldincludedir = @oldincludedir@ +pdfdir = @pdfdir@ +prefix = @prefix@ +program_transform_name = @program_transform_name@ +psdir = @psdir@ +runstatedir = @runstatedir@ +sbindir = @sbindir@ +sharedstatedir = @sharedstatedir@ +srcdir = @srcdir@ +sysconfdir = @sysconfdir@ +target_alias = @target_alias@ +top_build_prefix = @top_build_prefix@ +top_builddir = @top_builddir@ +top_srcdir = @top_srcdir@ + +#*************************************************************************** +# _ _ ____ _ +# Project ___| | | | _ \| | +# / __| | | | |_) | | +# | (__| |_| | _ <| |___ +# \___|\___/|_| \_\_____| +# +# Copyright (C) Daniel Stenberg, , et al. +# +# This software is licensed as described in the file COPYING, which +# you should have received as part of this distribution. The terms +# are also available at https://curl.se/docs/copyright.html. +# +# You may opt to use, copy, modify, merge, publish, distribute and/or sell +# copies of the Software, and permit persons to whom the Software is +# furnished to do so, under the terms of the COPYING file. +# +# This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY +# KIND, either express or implied. +# +# SPDX-License-Identifier: curl +# +########################################################################### +SUBDIRS = curl +EXTRA_DIST = README.md +AUTOMAKE_OPTIONS = foreign no-dependencies +all: all-recursive + +.SUFFIXES: +$(srcdir)/Makefile.in: @MAINTAINER_MODE_TRUE@ $(srcdir)/Makefile.am $(am__configure_deps) + @for dep in $?; do \ + case '$(am__configure_deps)' in \ + *$$dep*) \ + ( cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh ) \ + && { if test -f $@; then exit 0; else break; fi; }; \ + exit 1;; \ + esac; \ + done; \ + echo ' cd $(top_srcdir) && $(AUTOMAKE) --foreign include/Makefile'; \ + $(am__cd) $(top_srcdir) && \ + $(AUTOMAKE) --foreign include/Makefile +Makefile: $(srcdir)/Makefile.in $(top_builddir)/config.status + @case '$?' in \ + *config.status*) \ + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh;; \ + *) \ + echo ' cd $(top_builddir) && $(SHELL) ./config.status $(subdir)/$@ $(am__maybe_remake_depfiles)'; \ + cd $(top_builddir) && $(SHELL) ./config.status $(subdir)/$@ $(am__maybe_remake_depfiles);; \ + esac; + +$(top_builddir)/config.status: $(top_srcdir)/configure $(CONFIG_STATUS_DEPENDENCIES) + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh + +$(top_srcdir)/configure: @MAINTAINER_MODE_TRUE@ $(am__configure_deps) + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh +$(ACLOCAL_M4): @MAINTAINER_MODE_TRUE@ $(am__aclocal_m4_deps) + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh +$(am__aclocal_m4_deps): + +mostlyclean-libtool: + -rm -f *.lo + +clean-libtool: + -rm -rf .libs _libs + +# This directory's subdirectories are mostly independent; you can cd +# into them and run 'make' without going through this Makefile. +# To change the values of 'make' variables: instead of editing Makefiles, +# (1) if the variable is set in 'config.status', edit 'config.status' +# (which will cause the Makefiles to be regenerated when you run 'make'); +# (2) otherwise, pass the desired values on the 'make' command line. +$(am__recursive_targets): + @fail=; \ + if $(am__make_keepgoing); then \ + failcom='fail=yes'; \ + else \ + failcom='exit 1'; \ + fi; \ + dot_seen=no; \ + target=`echo $@ | sed s/-recursive//`; \ + case "$@" in \ + distclean-* | maintainer-clean-*) list='$(DIST_SUBDIRS)' ;; \ + *) list='$(SUBDIRS)' ;; \ + esac; \ + for subdir in $$list; do \ + echo "Making $$target in $$subdir"; \ + if test "$$subdir" = "."; then \ + dot_seen=yes; \ + local_target="$$target-am"; \ + else \ + local_target="$$target"; \ + fi; \ + ($(am__cd) $$subdir && $(MAKE) $(AM_MAKEFLAGS) $$local_target) \ + || eval $$failcom; \ + done; \ + if test "$$dot_seen" = "no"; then \ + $(MAKE) $(AM_MAKEFLAGS) "$$target-am" || exit 1; \ + fi; test -z "$$fail" + +ID: $(am__tagged_files) + $(am__define_uniq_tagged_files); mkid -fID $$unique +tags: tags-recursive +TAGS: tags + +tags-am: $(TAGS_DEPENDENCIES) $(am__tagged_files) + set x; \ + here=`pwd`; \ + if ($(ETAGS) --etags-include --version) >/dev/null 2>&1; then \ + include_option=--etags-include; \ + empty_fix=.; \ + else \ + include_option=--include; \ + empty_fix=; \ + fi; \ + list='$(SUBDIRS)'; for subdir in $$list; do \ + if test "$$subdir" = .; then :; else \ + test ! -f $$subdir/TAGS || \ + set "$$@" "$$include_option=$$here/$$subdir/TAGS"; \ + fi; \ + done; \ + $(am__define_uniq_tagged_files); \ + shift; \ + if test -z "$(ETAGS_ARGS)$$*$$unique"; then :; else \ + test -n "$$unique" || unique=$$empty_fix; \ + if test $$# -gt 0; then \ + $(ETAGS) $(ETAGSFLAGS) $(AM_ETAGSFLAGS) $(ETAGS_ARGS) \ + "$$@" $$unique; \ + else \ + $(ETAGS) $(ETAGSFLAGS) $(AM_ETAGSFLAGS) $(ETAGS_ARGS) \ + $$unique; \ + fi; \ + fi +ctags: ctags-recursive + +CTAGS: ctags +ctags-am: $(TAGS_DEPENDENCIES) $(am__tagged_files) + $(am__define_uniq_tagged_files); \ + test -z "$(CTAGS_ARGS)$$unique" \ + || $(CTAGS) $(CTAGSFLAGS) $(AM_CTAGSFLAGS) $(CTAGS_ARGS) \ + $$unique + +GTAGS: + here=`$(am__cd) $(top_builddir) && pwd` \ + && $(am__cd) $(top_srcdir) \ + && gtags -i $(GTAGS_ARGS) "$$here" +cscopelist: cscopelist-recursive + +cscopelist-am: $(am__tagged_files) + list='$(am__tagged_files)'; \ + case "$(srcdir)" in \ + [\\/]* | ?:[\\/]*) sdir="$(srcdir)" ;; \ + *) sdir=$(subdir)/$(srcdir) ;; \ + esac; \ + for i in $$list; do \ + if test -f "$$i"; then \ + echo "$(subdir)/$$i"; \ + else \ + echo "$$sdir/$$i"; \ + fi; \ + done >> $(top_builddir)/cscope.files + +distclean-tags: + -rm -f TAGS ID GTAGS GRTAGS GSYMS GPATH tags +distdir: $(BUILT_SOURCES) + $(MAKE) $(AM_MAKEFLAGS) distdir-am + +distdir-am: $(DISTFILES) + @srcdirstrip=`echo "$(srcdir)" | sed 's/[].[^$$\\*]/\\\\&/g'`; \ + topsrcdirstrip=`echo "$(top_srcdir)" | sed 's/[].[^$$\\*]/\\\\&/g'`; \ + list='$(DISTFILES)'; \ + dist_files=`for file in $$list; do echo $$file; done | \ + sed -e "s|^$$srcdirstrip/||;t" \ + -e "s|^$$topsrcdirstrip/|$(top_builddir)/|;t"`; \ + case $$dist_files in \ + */*) $(MKDIR_P) `echo "$$dist_files" | \ + sed '/\//!d;s|^|$(distdir)/|;s,/[^/]*$$,,' | \ + sort -u` ;; \ + esac; \ + for file in $$dist_files; do \ + if test -f $$file || test -d $$file; then d=.; else d=$(srcdir); fi; \ + if test -d $$d/$$file; then \ + dir=`echo "/$$file" | sed -e 's,/[^/]*$$,,'`; \ + if test -d "$(distdir)/$$file"; then \ + find "$(distdir)/$$file" -type d ! -perm -700 -exec chmod u+rwx {} \;; \ + fi; \ + if test -d $(srcdir)/$$file && test $$d != $(srcdir); then \ + cp -fpR $(srcdir)/$$file "$(distdir)$$dir" || exit 1; \ + find "$(distdir)/$$file" -type d ! -perm -700 -exec chmod u+rwx {} \;; \ + fi; \ + cp -fpR $$d/$$file "$(distdir)$$dir" || exit 1; \ + else \ + test -f "$(distdir)/$$file" \ + || cp -p $$d/$$file "$(distdir)/$$file" \ + || exit 1; \ + fi; \ + done + @list='$(DIST_SUBDIRS)'; for subdir in $$list; do \ + if test "$$subdir" = .; then :; else \ + $(am__make_dryrun) \ + || test -d "$(distdir)/$$subdir" \ + || $(MKDIR_P) "$(distdir)/$$subdir" \ + || exit 1; \ + dir1=$$subdir; dir2="$(distdir)/$$subdir"; \ + $(am__relativize); \ + new_distdir=$$reldir; \ + dir1=$$subdir; dir2="$(top_distdir)"; \ + $(am__relativize); \ + new_top_distdir=$$reldir; \ + echo " (cd $$subdir && $(MAKE) $(AM_MAKEFLAGS) top_distdir="$$new_top_distdir" distdir="$$new_distdir" \\"; \ + echo " am__remove_distdir=: am__skip_length_check=: am__skip_mode_fix=: distdir)"; \ + ($(am__cd) $$subdir && \ + $(MAKE) $(AM_MAKEFLAGS) \ + top_distdir="$$new_top_distdir" \ + distdir="$$new_distdir" \ + am__remove_distdir=: \ + am__skip_length_check=: \ + am__skip_mode_fix=: \ + distdir) \ + || exit 1; \ + fi; \ + done +check-am: all-am +check: check-recursive +all-am: Makefile +installdirs: installdirs-recursive +installdirs-am: +install: install-recursive +install-exec: install-exec-recursive +install-data: install-data-recursive +uninstall: uninstall-recursive + +install-am: all-am + @$(MAKE) $(AM_MAKEFLAGS) install-exec-am install-data-am + +installcheck: installcheck-recursive +install-strip: + if test -z '$(STRIP)'; then \ + $(MAKE) $(AM_MAKEFLAGS) INSTALL_PROGRAM="$(INSTALL_STRIP_PROGRAM)" \ + install_sh_PROGRAM="$(INSTALL_STRIP_PROGRAM)" INSTALL_STRIP_FLAG=-s \ + install; \ + else \ + $(MAKE) $(AM_MAKEFLAGS) INSTALL_PROGRAM="$(INSTALL_STRIP_PROGRAM)" \ + install_sh_PROGRAM="$(INSTALL_STRIP_PROGRAM)" INSTALL_STRIP_FLAG=-s \ + "INSTALL_PROGRAM_ENV=STRIPPROG='$(STRIP)'" install; \ + fi +mostlyclean-generic: + +clean-generic: + +distclean-generic: + -test -z "$(CONFIG_CLEAN_FILES)" || rm -f $(CONFIG_CLEAN_FILES) + -test . = "$(srcdir)" || test -z "$(CONFIG_CLEAN_VPATH_FILES)" || rm -f $(CONFIG_CLEAN_VPATH_FILES) + +maintainer-clean-generic: + @echo "This command is intended for maintainers to use" + @echo "it deletes files that may require special tools to rebuild." +clean: clean-recursive + +clean-am: clean-generic clean-libtool mostlyclean-am + +distclean: distclean-recursive + -rm -f Makefile +distclean-am: clean-am distclean-generic distclean-tags + +dvi: dvi-recursive + +dvi-am: + +html: html-recursive + +html-am: + +info: info-recursive + +info-am: + +install-data-am: + +install-dvi: install-dvi-recursive + +install-dvi-am: + +install-exec-am: + +install-html: install-html-recursive + +install-html-am: + +install-info: install-info-recursive + +install-info-am: + +install-man: + +install-pdf: install-pdf-recursive + +install-pdf-am: + +install-ps: install-ps-recursive + +install-ps-am: + +installcheck-am: + +maintainer-clean: maintainer-clean-recursive + -rm -f Makefile +maintainer-clean-am: distclean-am maintainer-clean-generic + +mostlyclean: mostlyclean-recursive + +mostlyclean-am: mostlyclean-generic mostlyclean-libtool + +pdf: pdf-recursive + +pdf-am: + +ps: ps-recursive + +ps-am: + +uninstall-am: + +.MAKE: $(am__recursive_targets) install-am install-strip + +.PHONY: $(am__recursive_targets) CTAGS GTAGS TAGS all all-am check \ + check-am clean clean-generic clean-libtool cscopelist-am ctags \ + ctags-am distclean distclean-generic distclean-libtool \ + distclean-tags distdir dvi dvi-am html html-am info info-am \ + install install-am install-data install-data-am install-dvi \ + install-dvi-am install-exec install-exec-am install-html \ + install-html-am install-info install-info-am install-man \ + install-pdf install-pdf-am install-ps install-ps-am \ + install-strip installcheck installcheck-am installdirs \ + installdirs-am maintainer-clean maintainer-clean-generic \ + mostlyclean mostlyclean-generic mostlyclean-libtool pdf pdf-am \ + ps ps-am tags tags-am uninstall uninstall-am + +.PRECIOUS: Makefile + + +# Tell versions [3.59,3.63) of GNU make to not export all variables. +# Otherwise a system limit (for SysV at least) may be exceeded. +.NOEXPORT: diff --git a/cactus-engine/libs/curl/include/curl/Makefile.am b/cactus-engine/libs/curl/include/curl/Makefile.am new file mode 100644 index 000000000..f664f6ab7 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/Makefile.am @@ -0,0 +1,43 @@ +#*************************************************************************** +# _ _ ____ _ +# Project ___| | | | _ \| | +# / __| | | | |_) | | +# | (__| |_| | _ <| |___ +# \___|\___/|_| \_\_____| +# +# Copyright (C) Daniel Stenberg, , et al. +# +# This software is licensed as described in the file COPYING, which +# you should have received as part of this distribution. The terms +# are also available at https://curl.se/docs/copyright.html. +# +# You may opt to use, copy, modify, merge, publish, distribute and/or sell +# copies of the Software, and permit persons to whom the Software is +# furnished to do so, under the terms of the COPYING file. +# +# This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY +# KIND, either express or implied. +# +# SPDX-License-Identifier: curl +# +########################################################################### +pkginclude_HEADERS = \ + curl.h curlver.h easy.h mprintf.h stdcheaders.h multi.h \ + typecheck-gcc.h system.h urlapi.h options.h header.h websockets.h + +pkgincludedir= $(includedir)/curl + +CHECKSRC = $(CS_$(V)) +CS_0 = @echo " RUN " $@; +CS_1 = +CS_ = $(CS_0) + +checksrc: + $(CHECKSRC)@PERL@ $(top_srcdir)/scripts/checksrc.pl -D$(top_srcdir)/include/curl $(pkginclude_HEADERS) + +if NOT_CURL_CI +if DEBUGBUILD +# for debug builds, we scan the sources on all regular make invokes +all-local: checksrc +endif +endif diff --git a/cactus-engine/libs/curl/include/curl/Makefile.in b/cactus-engine/libs/curl/include/curl/Makefile.in new file mode 100644 index 000000000..ed4f28818 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/Makefile.in @@ -0,0 +1,665 @@ +# Makefile.in generated by automake 1.16.5 from Makefile.am. +# @configure_input@ + +# Copyright (C) 1994-2021 Free Software Foundation, Inc. + +# This Makefile.in is free software; the Free Software Foundation +# gives unlimited permission to copy and/or distribute it, +# with or without modifications, as long as this notice is preserved. + +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY, to the extent permitted by law; without +# even the implied warranty of MERCHANTABILITY or FITNESS FOR A +# PARTICULAR PURPOSE. + +@SET_MAKE@ + +VPATH = @srcdir@ +am__is_gnu_make = { \ + if test -z '$(MAKELEVEL)'; then \ + false; \ + elif test -n '$(MAKE_HOST)'; then \ + true; \ + elif test -n '$(MAKE_VERSION)' && test -n '$(CURDIR)'; then \ + true; \ + else \ + false; \ + fi; \ +} +am__make_running_with_option = \ + case $${target_option-} in \ + ?) ;; \ + *) echo "am__make_running_with_option: internal error: invalid" \ + "target option '$${target_option-}' specified" >&2; \ + exit 1;; \ + esac; \ + has_opt=no; \ + sane_makeflags=$$MAKEFLAGS; \ + if $(am__is_gnu_make); then \ + sane_makeflags=$$MFLAGS; \ + else \ + case $$MAKEFLAGS in \ + *\\[\ \ ]*) \ + bs=\\; \ + sane_makeflags=`printf '%s\n' "$$MAKEFLAGS" \ + | sed "s/$$bs$$bs[$$bs $$bs ]*//g"`;; \ + esac; \ + fi; \ + skip_next=no; \ + strip_trailopt () \ + { \ + flg=`printf '%s\n' "$$flg" | sed "s/$$1.*$$//"`; \ + }; \ + for flg in $$sane_makeflags; do \ + test $$skip_next = yes && { skip_next=no; continue; }; \ + case $$flg in \ + *=*|--*) continue;; \ + -*I) strip_trailopt 'I'; skip_next=yes;; \ + -*I?*) strip_trailopt 'I';; \ + -*O) strip_trailopt 'O'; skip_next=yes;; \ + -*O?*) strip_trailopt 'O';; \ + -*l) strip_trailopt 'l'; skip_next=yes;; \ + -*l?*) strip_trailopt 'l';; \ + -[dEDm]) skip_next=yes;; \ + -[JT]) skip_next=yes;; \ + esac; \ + case $$flg in \ + *$$target_option*) has_opt=yes; break;; \ + esac; \ + done; \ + test $$has_opt = yes +am__make_dryrun = (target_option=n; $(am__make_running_with_option)) +am__make_keepgoing = (target_option=k; $(am__make_running_with_option)) +pkgdatadir = $(datadir)/@PACKAGE@ +pkglibdir = $(libdir)/@PACKAGE@ +pkglibexecdir = $(libexecdir)/@PACKAGE@ +am__cd = CDPATH="$${ZSH_VERSION+.}$(PATH_SEPARATOR)" && cd +install_sh_DATA = $(install_sh) -c -m 644 +install_sh_PROGRAM = $(install_sh) -c +install_sh_SCRIPT = $(install_sh) -c +INSTALL_HEADER = $(INSTALL_DATA) +transform = $(program_transform_name) +NORMAL_INSTALL = : +PRE_INSTALL = : +POST_INSTALL = : +NORMAL_UNINSTALL = : +PRE_UNINSTALL = : +POST_UNINSTALL = : +build_triplet = @build@ +host_triplet = @host@ +subdir = include/curl +ACLOCAL_M4 = $(top_srcdir)/aclocal.m4 +am__aclocal_m4_deps = $(top_srcdir)/m4/curl-amissl.m4 \ + $(top_srcdir)/m4/curl-apple-sectrust.m4 \ + $(top_srcdir)/m4/curl-compilers.m4 \ + $(top_srcdir)/m4/curl-confopts.m4 \ + $(top_srcdir)/m4/curl-functions.m4 \ + $(top_srcdir)/m4/curl-gnutls.m4 \ + $(top_srcdir)/m4/curl-mbedtls.m4 \ + $(top_srcdir)/m4/curl-openssl.m4 \ + $(top_srcdir)/m4/curl-override.m4 \ + $(top_srcdir)/m4/curl-reentrant.m4 \ + $(top_srcdir)/m4/curl-rustls.m4 \ + $(top_srcdir)/m4/curl-schannel.m4 \ + $(top_srcdir)/m4/curl-sysconfig.m4 \ + $(top_srcdir)/m4/curl-wolfssl.m4 $(top_srcdir)/m4/libtool.m4 \ + $(top_srcdir)/m4/ltoptions.m4 $(top_srcdir)/m4/ltsugar.m4 \ + $(top_srcdir)/m4/ltversion.m4 $(top_srcdir)/m4/lt~obsolete.m4 \ + $(top_srcdir)/m4/xc-am-iface.m4 \ + $(top_srcdir)/m4/xc-cc-check.m4 \ + $(top_srcdir)/m4/xc-lt-iface.m4 \ + $(top_srcdir)/m4/xc-val-flgs.m4 \ + $(top_srcdir)/m4/zz40-xc-ovr.m4 \ + $(top_srcdir)/m4/zz50-xc-ovr.m4 $(top_srcdir)/acinclude.m4 \ + $(top_srcdir)/configure.ac +am__configure_deps = $(am__aclocal_m4_deps) $(CONFIGURE_DEPENDENCIES) \ + $(ACLOCAL_M4) +DIST_COMMON = $(srcdir)/Makefile.am $(pkginclude_HEADERS) \ + $(am__DIST_COMMON) +mkinstalldirs = $(install_sh) -d +CONFIG_HEADER = $(top_builddir)/lib/curl_config.h +CONFIG_CLEAN_FILES = +CONFIG_CLEAN_VPATH_FILES = +AM_V_P = $(am__v_P_@AM_V@) +am__v_P_ = $(am__v_P_@AM_DEFAULT_V@) +am__v_P_0 = false +am__v_P_1 = : +AM_V_GEN = $(am__v_GEN_@AM_V@) +am__v_GEN_ = $(am__v_GEN_@AM_DEFAULT_V@) +am__v_GEN_0 = @echo " GEN " $@; +am__v_GEN_1 = +AM_V_at = $(am__v_at_@AM_V@) +am__v_at_ = $(am__v_at_@AM_DEFAULT_V@) +am__v_at_0 = @ +am__v_at_1 = +SOURCES = +DIST_SOURCES = +am__can_run_installinfo = \ + case $$AM_UPDATE_INFO_DIR in \ + n|no|NO) false;; \ + *) (install-info --version) >/dev/null 2>&1;; \ + esac +am__vpath_adj_setup = srcdirstrip=`echo "$(srcdir)" | sed 's|.|.|g'`; +am__vpath_adj = case $$p in \ + $(srcdir)/*) f=`echo "$$p" | sed "s|^$$srcdirstrip/||"`;; \ + *) f=$$p;; \ + esac; +am__strip_dir = f=`echo $$p | sed -e 's|^.*/||'`; +am__install_max = 40 +am__nobase_strip_setup = \ + srcdirstrip=`echo "$(srcdir)" | sed 's/[].[^$$\\*|]/\\\\&/g'` +am__nobase_strip = \ + for p in $$list; do echo "$$p"; done | sed -e "s|$$srcdirstrip/||" +am__nobase_list = $(am__nobase_strip_setup); \ + for p in $$list; do echo "$$p $$p"; done | \ + sed "s| $$srcdirstrip/| |;"' / .*\//!s/ .*/ ./; s,\( .*\)/[^/]*$$,\1,' | \ + $(AWK) 'BEGIN { files["."] = "" } { files[$$2] = files[$$2] " " $$1; \ + if (++n[$$2] == $(am__install_max)) \ + { print $$2, files[$$2]; n[$$2] = 0; files[$$2] = "" } } \ + END { for (dir in files) print dir, files[dir] }' +am__base_list = \ + sed '$$!N;$$!N;$$!N;$$!N;$$!N;$$!N;$$!N;s/\n/ /g' | \ + sed '$$!N;$$!N;$$!N;$$!N;s/\n/ /g' +am__uninstall_files_from_dir = { \ + test -z "$$files" \ + || { test ! -d "$$dir" && test ! -f "$$dir" && test ! -r "$$dir"; } \ + || { echo " ( cd '$$dir' && rm -f" $$files ")"; \ + $(am__cd) "$$dir" && rm -f $$files; }; \ + } +am__installdirs = "$(DESTDIR)$(pkgincludedir)" +HEADERS = $(pkginclude_HEADERS) +am__tagged_files = $(HEADERS) $(SOURCES) $(TAGS_FILES) $(LISP) +# Read a list of newline-separated strings from the standard input, +# and print each of them once, without duplicates. Input order is +# *not* preserved. +am__uniquify_input = $(AWK) '\ + BEGIN { nonempty = 0; } \ + { items[$$0] = 1; nonempty = 1; } \ + END { if (nonempty) { for (i in items) print i; }; } \ +' +# Make sure the list of sources is unique. This is necessary because, +# e.g., the same source file might be shared among _SOURCES variables +# for different programs/libraries. +am__define_uniq_tagged_files = \ + list='$(am__tagged_files)'; \ + unique=`for i in $$list; do \ + if test -f "$$i"; then echo $$i; else echo $(srcdir)/$$i; fi; \ + done | $(am__uniquify_input)` +am__DIST_COMMON = $(srcdir)/Makefile.in +DISTFILES = $(DIST_COMMON) $(DIST_SOURCES) $(TEXINFOS) $(EXTRA_DIST) +pkgincludedir = $(includedir)/curl +ACLOCAL = @ACLOCAL@ +AMTAR = @AMTAR@ +AM_DEFAULT_VERBOSITY = @AM_DEFAULT_VERBOSITY@ +APXS = @APXS@ +AR = @AR@ +AR_FLAGS = @AR_FLAGS@ +AS = @AS@ +AUTOCONF = @AUTOCONF@ +AUTOHEADER = @AUTOHEADER@ +AUTOMAKE = @AUTOMAKE@ +AWK = @AWK@ +BLANK_AT_MAKETIME = @BLANK_AT_MAKETIME@ +CADDY = @CADDY@ +CC = @CC@ +CCDEPMODE = @CCDEPMODE@ +CFLAGS = @CFLAGS@ +CFLAG_CURL_SYMBOL_HIDING = @CFLAG_CURL_SYMBOL_HIDING@ +CONFIGURE_OPTIONS = @CONFIGURE_OPTIONS@ +CPP = @CPP@ +CPPFLAGS = @CPPFLAGS@ +CSCOPE = @CSCOPE@ +CTAGS = @CTAGS@ +CURLVERSION = @CURLVERSION@ +CURL_CA_BUNDLE = @CURL_CA_BUNDLE@ +CURL_CA_EMBED = @CURL_CA_EMBED@ +CURL_CFLAG_EXTRAS = @CURL_CFLAG_EXTRAS@ +CURL_CPP = @CURL_CPP@ +CURL_LIBCURL_VERSIONED_SYMBOLS_PREFIX = @CURL_LIBCURL_VERSIONED_SYMBOLS_PREFIX@ +CURL_LIBCURL_VERSIONED_SYMBOLS_SONAME = @CURL_LIBCURL_VERSIONED_SYMBOLS_SONAME@ +CURL_NETWORK_AND_TIME_LIBS = @CURL_NETWORK_AND_TIME_LIBS@ +CYGPATH_W = @CYGPATH_W@ +DANTED = @DANTED@ +DEFS = @DEFS@ +DEPDIR = @DEPDIR@ +DLLTOOL = @DLLTOOL@ +DSYMUTIL = @DSYMUTIL@ +DUMPBIN = @DUMPBIN@ +ECHO_C = @ECHO_C@ +ECHO_N = @ECHO_N@ +ECHO_T = @ECHO_T@ +EGREP = @EGREP@ +ENABLE_SHARED = @ENABLE_SHARED@ +ENABLE_STATIC = @ENABLE_STATIC@ +ETAGS = @ETAGS@ +EXEEXT = @EXEEXT@ +FGREP = @FGREP@ +FILECMD = @FILECMD@ +FISH_FUNCTIONS_DIR = @FISH_FUNCTIONS_DIR@ +GCOV = @GCOV@ +GREP = @GREP@ +HAVE_LIBZ = @HAVE_LIBZ@ +HTTPD = @HTTPD@ +HTTPD_NGHTTPX = @HTTPD_NGHTTPX@ +INSTALL = @INSTALL@ +INSTALL_DATA = @INSTALL_DATA@ +INSTALL_PROGRAM = @INSTALL_PROGRAM@ +INSTALL_SCRIPT = @INSTALL_SCRIPT@ +INSTALL_STRIP_PROGRAM = @INSTALL_STRIP_PROGRAM@ +LCOV = @LCOV@ +LD = @LD@ +LDFLAGS = @LDFLAGS@ +LIBCURL_PC_CFLAGS = @LIBCURL_PC_CFLAGS@ +LIBCURL_PC_CFLAGS_PRIVATE = @LIBCURL_PC_CFLAGS_PRIVATE@ +LIBCURL_PC_LDFLAGS_PRIVATE = @LIBCURL_PC_LDFLAGS_PRIVATE@ +LIBCURL_PC_LIBS = @LIBCURL_PC_LIBS@ +LIBCURL_PC_LIBS_PRIVATE = @LIBCURL_PC_LIBS_PRIVATE@ +LIBCURL_PC_REQUIRES = @LIBCURL_PC_REQUIRES@ +LIBCURL_PC_REQUIRES_PRIVATE = @LIBCURL_PC_REQUIRES_PRIVATE@ +LIBOBJS = @LIBOBJS@ +LIBS = @LIBS@ +LIBTOOL = @LIBTOOL@ +LIPO = @LIPO@ +LN_S = @LN_S@ +LTLIBOBJS = @LTLIBOBJS@ +LT_SYS_LIBRARY_PATH = @LT_SYS_LIBRARY_PATH@ +MAINT = @MAINT@ +MAKEINFO = @MAKEINFO@ +MANIFEST_TOOL = @MANIFEST_TOOL@ +MKDIR_P = @MKDIR_P@ +NM = @NM@ +NMEDIT = @NMEDIT@ +OBJDUMP = @OBJDUMP@ +OBJEXT = @OBJEXT@ +OTOOL = @OTOOL@ +OTOOL64 = @OTOOL64@ +PACKAGE = @PACKAGE@ +PACKAGE_BUGREPORT = @PACKAGE_BUGREPORT@ +PACKAGE_NAME = @PACKAGE_NAME@ +PACKAGE_STRING = @PACKAGE_STRING@ +PACKAGE_TARNAME = @PACKAGE_TARNAME@ +PACKAGE_URL = @PACKAGE_URL@ +PACKAGE_VERSION = @PACKAGE_VERSION@ +PATH_SEPARATOR = @PATH_SEPARATOR@ +PERL = @PERL@ +PKGCONFIG = @PKGCONFIG@ +RANLIB = @RANLIB@ +RC = @RC@ +SED = @SED@ +SET_MAKE = @SET_MAKE@ +SFTPD = @SFTPD@ +SHELL = @SHELL@ +SSHD = @SSHD@ +SSL_BACKENDS = @SSL_BACKENDS@ +STRIP = @STRIP@ +SUPPORT_FEATURES = @SUPPORT_FEATURES@ +SUPPORT_PROTOCOLS = @SUPPORT_PROTOCOLS@ +TEST_NGHTTPX = @TEST_NGHTTPX@ +VERSION = @VERSION@ +VERSIONNUM = @VERSIONNUM@ +VSFTPD = @VSFTPD@ +ZLIB_LIBS = @ZLIB_LIBS@ +ZSH_FUNCTIONS_DIR = @ZSH_FUNCTIONS_DIR@ +abs_builddir = @abs_builddir@ +abs_srcdir = @abs_srcdir@ +abs_top_builddir = @abs_top_builddir@ +abs_top_srcdir = @abs_top_srcdir@ +ac_ct_AR = @ac_ct_AR@ +ac_ct_CC = @ac_ct_CC@ +ac_ct_DUMPBIN = @ac_ct_DUMPBIN@ +am__include = @am__include@ +am__leading_dot = @am__leading_dot@ +am__quote = @am__quote@ +am__tar = @am__tar@ +am__untar = @am__untar@ +bindir = @bindir@ +build = @build@ +build_alias = @build_alias@ +build_cpu = @build_cpu@ +build_os = @build_os@ +build_vendor = @build_vendor@ +builddir = @builddir@ +datadir = @datadir@ +datarootdir = @datarootdir@ +docdir = @docdir@ +dvidir = @dvidir@ +exec_prefix = @exec_prefix@ +host = @host@ +host_alias = @host_alias@ +host_cpu = @host_cpu@ +host_os = @host_os@ +host_vendor = @host_vendor@ +htmldir = @htmldir@ +includedir = @includedir@ +infodir = @infodir@ +install_sh = @install_sh@ +libdir = @libdir@ +libexecdir = @libexecdir@ +libext = @libext@ +localedir = @localedir@ +localstatedir = @localstatedir@ +mandir = @mandir@ +mkdir_p = @mkdir_p@ +oldincludedir = @oldincludedir@ +pdfdir = @pdfdir@ +prefix = @prefix@ +program_transform_name = @program_transform_name@ +psdir = @psdir@ +runstatedir = @runstatedir@ +sbindir = @sbindir@ +sharedstatedir = @sharedstatedir@ +srcdir = @srcdir@ +sysconfdir = @sysconfdir@ +target_alias = @target_alias@ +top_build_prefix = @top_build_prefix@ +top_builddir = @top_builddir@ +top_srcdir = @top_srcdir@ + +#*************************************************************************** +# _ _ ____ _ +# Project ___| | | | _ \| | +# / __| | | | |_) | | +# | (__| |_| | _ <| |___ +# \___|\___/|_| \_\_____| +# +# Copyright (C) Daniel Stenberg, , et al. +# +# This software is licensed as described in the file COPYING, which +# you should have received as part of this distribution. The terms +# are also available at https://curl.se/docs/copyright.html. +# +# You may opt to use, copy, modify, merge, publish, distribute and/or sell +# copies of the Software, and permit persons to whom the Software is +# furnished to do so, under the terms of the COPYING file. +# +# This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY +# KIND, either express or implied. +# +# SPDX-License-Identifier: curl +# +########################################################################### +pkginclude_HEADERS = \ + curl.h curlver.h easy.h mprintf.h stdcheaders.h multi.h \ + typecheck-gcc.h system.h urlapi.h options.h header.h websockets.h + +CHECKSRC = $(CS_$(V)) +CS_0 = @echo " RUN " $@; +CS_1 = +CS_ = $(CS_0) +all: all-am + +.SUFFIXES: +$(srcdir)/Makefile.in: @MAINTAINER_MODE_TRUE@ $(srcdir)/Makefile.am $(am__configure_deps) + @for dep in $?; do \ + case '$(am__configure_deps)' in \ + *$$dep*) \ + ( cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh ) \ + && { if test -f $@; then exit 0; else break; fi; }; \ + exit 1;; \ + esac; \ + done; \ + echo ' cd $(top_srcdir) && $(AUTOMAKE) --gnu include/curl/Makefile'; \ + $(am__cd) $(top_srcdir) && \ + $(AUTOMAKE) --gnu include/curl/Makefile +Makefile: $(srcdir)/Makefile.in $(top_builddir)/config.status + @case '$?' in \ + *config.status*) \ + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh;; \ + *) \ + echo ' cd $(top_builddir) && $(SHELL) ./config.status $(subdir)/$@ $(am__maybe_remake_depfiles)'; \ + cd $(top_builddir) && $(SHELL) ./config.status $(subdir)/$@ $(am__maybe_remake_depfiles);; \ + esac; + +$(top_builddir)/config.status: $(top_srcdir)/configure $(CONFIG_STATUS_DEPENDENCIES) + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh + +$(top_srcdir)/configure: @MAINTAINER_MODE_TRUE@ $(am__configure_deps) + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh +$(ACLOCAL_M4): @MAINTAINER_MODE_TRUE@ $(am__aclocal_m4_deps) + cd $(top_builddir) && $(MAKE) $(AM_MAKEFLAGS) am--refresh +$(am__aclocal_m4_deps): + +mostlyclean-libtool: + -rm -f *.lo + +clean-libtool: + -rm -rf .libs _libs +install-pkgincludeHEADERS: $(pkginclude_HEADERS) + @$(NORMAL_INSTALL) + @list='$(pkginclude_HEADERS)'; test -n "$(pkgincludedir)" || list=; \ + if test -n "$$list"; then \ + echo " $(MKDIR_P) '$(DESTDIR)$(pkgincludedir)'"; \ + $(MKDIR_P) "$(DESTDIR)$(pkgincludedir)" || exit 1; \ + fi; \ + for p in $$list; do \ + if test -f "$$p"; then d=; else d="$(srcdir)/"; fi; \ + echo "$$d$$p"; \ + done | $(am__base_list) | \ + while read files; do \ + echo " $(INSTALL_HEADER) $$files '$(DESTDIR)$(pkgincludedir)'"; \ + $(INSTALL_HEADER) $$files "$(DESTDIR)$(pkgincludedir)" || exit $$?; \ + done + +uninstall-pkgincludeHEADERS: + @$(NORMAL_UNINSTALL) + @list='$(pkginclude_HEADERS)'; test -n "$(pkgincludedir)" || list=; \ + files=`for p in $$list; do echo $$p; done | sed -e 's|^.*/||'`; \ + dir='$(DESTDIR)$(pkgincludedir)'; $(am__uninstall_files_from_dir) + +ID: $(am__tagged_files) + $(am__define_uniq_tagged_files); mkid -fID $$unique +tags: tags-am +TAGS: tags + +tags-am: $(TAGS_DEPENDENCIES) $(am__tagged_files) + set x; \ + here=`pwd`; \ + $(am__define_uniq_tagged_files); \ + shift; \ + if test -z "$(ETAGS_ARGS)$$*$$unique"; then :; else \ + test -n "$$unique" || unique=$$empty_fix; \ + if test $$# -gt 0; then \ + $(ETAGS) $(ETAGSFLAGS) $(AM_ETAGSFLAGS) $(ETAGS_ARGS) \ + "$$@" $$unique; \ + else \ + $(ETAGS) $(ETAGSFLAGS) $(AM_ETAGSFLAGS) $(ETAGS_ARGS) \ + $$unique; \ + fi; \ + fi +ctags: ctags-am + +CTAGS: ctags +ctags-am: $(TAGS_DEPENDENCIES) $(am__tagged_files) + $(am__define_uniq_tagged_files); \ + test -z "$(CTAGS_ARGS)$$unique" \ + || $(CTAGS) $(CTAGSFLAGS) $(AM_CTAGSFLAGS) $(CTAGS_ARGS) \ + $$unique + +GTAGS: + here=`$(am__cd) $(top_builddir) && pwd` \ + && $(am__cd) $(top_srcdir) \ + && gtags -i $(GTAGS_ARGS) "$$here" +cscopelist: cscopelist-am + +cscopelist-am: $(am__tagged_files) + list='$(am__tagged_files)'; \ + case "$(srcdir)" in \ + [\\/]* | ?:[\\/]*) sdir="$(srcdir)" ;; \ + *) sdir=$(subdir)/$(srcdir) ;; \ + esac; \ + for i in $$list; do \ + if test -f "$$i"; then \ + echo "$(subdir)/$$i"; \ + else \ + echo "$$sdir/$$i"; \ + fi; \ + done >> $(top_builddir)/cscope.files + +distclean-tags: + -rm -f TAGS ID GTAGS GRTAGS GSYMS GPATH tags +distdir: $(BUILT_SOURCES) + $(MAKE) $(AM_MAKEFLAGS) distdir-am + +distdir-am: $(DISTFILES) + @srcdirstrip=`echo "$(srcdir)" | sed 's/[].[^$$\\*]/\\\\&/g'`; \ + topsrcdirstrip=`echo "$(top_srcdir)" | sed 's/[].[^$$\\*]/\\\\&/g'`; \ + list='$(DISTFILES)'; \ + dist_files=`for file in $$list; do echo $$file; done | \ + sed -e "s|^$$srcdirstrip/||;t" \ + -e "s|^$$topsrcdirstrip/|$(top_builddir)/|;t"`; \ + case $$dist_files in \ + */*) $(MKDIR_P) `echo "$$dist_files" | \ + sed '/\//!d;s|^|$(distdir)/|;s,/[^/]*$$,,' | \ + sort -u` ;; \ + esac; \ + for file in $$dist_files; do \ + if test -f $$file || test -d $$file; then d=.; else d=$(srcdir); fi; \ + if test -d $$d/$$file; then \ + dir=`echo "/$$file" | sed -e 's,/[^/]*$$,,'`; \ + if test -d "$(distdir)/$$file"; then \ + find "$(distdir)/$$file" -type d ! -perm -700 -exec chmod u+rwx {} \;; \ + fi; \ + if test -d $(srcdir)/$$file && test $$d != $(srcdir); then \ + cp -fpR $(srcdir)/$$file "$(distdir)$$dir" || exit 1; \ + find "$(distdir)/$$file" -type d ! -perm -700 -exec chmod u+rwx {} \;; \ + fi; \ + cp -fpR $$d/$$file "$(distdir)$$dir" || exit 1; \ + else \ + test -f "$(distdir)/$$file" \ + || cp -p $$d/$$file "$(distdir)/$$file" \ + || exit 1; \ + fi; \ + done +check-am: all-am +check: check-am +@DEBUGBUILD_FALSE@all-local: +@NOT_CURL_CI_FALSE@all-local: +all-am: Makefile $(HEADERS) all-local +installdirs: + for dir in "$(DESTDIR)$(pkgincludedir)"; do \ + test -z "$$dir" || $(MKDIR_P) "$$dir"; \ + done +install: install-am +install-exec: install-exec-am +install-data: install-data-am +uninstall: uninstall-am + +install-am: all-am + @$(MAKE) $(AM_MAKEFLAGS) install-exec-am install-data-am + +installcheck: installcheck-am +install-strip: + if test -z '$(STRIP)'; then \ + $(MAKE) $(AM_MAKEFLAGS) INSTALL_PROGRAM="$(INSTALL_STRIP_PROGRAM)" \ + install_sh_PROGRAM="$(INSTALL_STRIP_PROGRAM)" INSTALL_STRIP_FLAG=-s \ + install; \ + else \ + $(MAKE) $(AM_MAKEFLAGS) INSTALL_PROGRAM="$(INSTALL_STRIP_PROGRAM)" \ + install_sh_PROGRAM="$(INSTALL_STRIP_PROGRAM)" INSTALL_STRIP_FLAG=-s \ + "INSTALL_PROGRAM_ENV=STRIPPROG='$(STRIP)'" install; \ + fi +mostlyclean-generic: + +clean-generic: + +distclean-generic: + -test -z "$(CONFIG_CLEAN_FILES)" || rm -f $(CONFIG_CLEAN_FILES) + -test . = "$(srcdir)" || test -z "$(CONFIG_CLEAN_VPATH_FILES)" || rm -f $(CONFIG_CLEAN_VPATH_FILES) + +maintainer-clean-generic: + @echo "This command is intended for maintainers to use" + @echo "it deletes files that may require special tools to rebuild." +clean: clean-am + +clean-am: clean-generic clean-libtool mostlyclean-am + +distclean: distclean-am + -rm -f Makefile +distclean-am: clean-am distclean-generic distclean-tags + +dvi: dvi-am + +dvi-am: + +html: html-am + +html-am: + +info: info-am + +info-am: + +install-data-am: install-pkgincludeHEADERS + +install-dvi: install-dvi-am + +install-dvi-am: + +install-exec-am: + +install-html: install-html-am + +install-html-am: + +install-info: install-info-am + +install-info-am: + +install-man: + +install-pdf: install-pdf-am + +install-pdf-am: + +install-ps: install-ps-am + +install-ps-am: + +installcheck-am: + +maintainer-clean: maintainer-clean-am + -rm -f Makefile +maintainer-clean-am: distclean-am maintainer-clean-generic + +mostlyclean: mostlyclean-am + +mostlyclean-am: mostlyclean-generic mostlyclean-libtool + +pdf: pdf-am + +pdf-am: + +ps: ps-am + +ps-am: + +uninstall-am: uninstall-pkgincludeHEADERS + +.MAKE: install-am install-strip + +.PHONY: CTAGS GTAGS TAGS all all-am all-local check check-am clean \ + clean-generic clean-libtool cscopelist-am ctags ctags-am \ + distclean distclean-generic distclean-libtool distclean-tags \ + distdir dvi dvi-am html html-am info info-am install \ + install-am install-data install-data-am install-dvi \ + install-dvi-am install-exec install-exec-am install-html \ + install-html-am install-info install-info-am install-man \ + install-pdf install-pdf-am install-pkgincludeHEADERS \ + install-ps install-ps-am install-strip installcheck \ + installcheck-am installdirs maintainer-clean \ + maintainer-clean-generic mostlyclean mostlyclean-generic \ + mostlyclean-libtool pdf pdf-am ps ps-am tags tags-am uninstall \ + uninstall-am uninstall-pkgincludeHEADERS + +.PRECIOUS: Makefile + + +checksrc: + $(CHECKSRC)@PERL@ $(top_srcdir)/scripts/checksrc.pl -D$(top_srcdir)/include/curl $(pkginclude_HEADERS) + +# for debug builds, we scan the sources on all regular make invokes +@DEBUGBUILD_TRUE@@NOT_CURL_CI_TRUE@all-local: checksrc + +# Tell versions [3.59,3.63) of GNU make to not export all variables. +# Otherwise a system limit (for SysV at least) may be exceeded. +.NOEXPORT: diff --git a/cactus-engine/libs/curl/include/curl/curl.h b/cactus-engine/libs/curl/include/curl/curl.h new file mode 100644 index 000000000..e755f098f --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/curl.h @@ -0,0 +1,3334 @@ +#ifndef CURLINC_CURL_H +#define CURLINC_CURL_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +/* + * If you have libcurl problems, all docs and details are found here: + * https://curl.se/libcurl/ + */ + +#ifdef CURL_NO_OLDIES +#define CURL_STRICTER /* not used since 8.11.0 */ +#endif + +/* Compile-time deprecation macros. */ +#if (defined(__GNUC__) && \ + ((__GNUC__ > 12) || ((__GNUC__ == 12) && (__GNUC_MINOR__ >= 1))) || \ + (defined(__clang__) && __clang_major__ >= 3) || \ + defined(__IAR_SYSTEMS_ICC__)) && \ + !defined(__INTEL_COMPILER) && \ + !defined(CURL_DISABLE_DEPRECATION) && !defined(BUILDING_LIBCURL) +#define CURL_DEPRECATED(version, message) \ + __attribute__((deprecated("since " # version ". " message))) +#ifdef __IAR_SYSTEMS_ICC__ +#define CURL_IGNORE_DEPRECATION(statements) \ + _Pragma("diag_suppress=Pe1444") \ + statements \ + _Pragma("diag_default=Pe1444") +#else +#define CURL_IGNORE_DEPRECATION(statements) \ + _Pragma("GCC diagnostic push") \ + _Pragma("GCC diagnostic ignored \"-Wdeprecated-declarations\"") \ + statements \ + _Pragma("GCC diagnostic pop") +#endif +#else +#define CURL_DEPRECATED(version, message) +#define CURL_IGNORE_DEPRECATION(statements) statements +#endif + +#include "curlver.h" /* libcurl version defines */ +#include "system.h" /* determine things runtime */ + +#include +#include + +#if defined(__FreeBSD__) || defined(__MidnightBSD__) +/* Needed for __FreeBSD_version or __MidnightBSD_version symbol definition */ +#include +#endif + +/* The include stuff here below is mainly for time_t! */ +#include +#include + +#if defined(_WIN32) && !defined(_WIN32_WCE) && !defined(__CYGWIN__) +#if !(defined(_WINSOCKAPI_) || defined(_WINSOCK_H) || \ + defined(__LWIP_OPT_H__) || defined(LWIP_HDR_OPT_H)) +/* The check above prevents the winsock2.h inclusion if winsock.h already was + included, since they cannot co-exist without problems */ +#include +#include +#endif +#endif + +/* HP-UX systems version 9, 10 and 11 lack sys/select.h and so does oldish + libc5-based Linux systems. Only include it on systems that are known to + require it! */ +#if defined(_AIX) || defined(__NOVELL_LIBC__) || defined(__NetBSD__) || \ + defined(__minix) || defined(__INTEGRITY) || \ + defined(ANDROID) || defined(__ANDROID__) || defined(__OpenBSD__) || \ + defined(__CYGWIN__) || defined(AMIGA) || defined(__NuttX__) || \ + (defined(__FreeBSD_version) && (__FreeBSD_version < 800000)) || \ + (defined(__MidnightBSD_version) && (__MidnightBSD_version < 100000)) || \ + defined(__sun__) || defined(__serenity__) || defined(__vxworks__) +#include +#endif + +#ifndef _WIN32 +#include +#include +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +typedef void CURL; +typedef void CURLSH; + +/* + * libcurl external API function linkage decorations. + */ + +#ifdef __has_declspec_attribute +#define CURL_HAS_DECLSPEC_ATTRIBUTE(x) __has_declspec_attribute(x) +#else +#define CURL_HAS_DECLSPEC_ATTRIBUTE(x) 0 +#endif + +#ifdef CURL_STATICLIB +# define CURL_EXTERN +#elif defined(_WIN32) || \ + (CURL_HAS_DECLSPEC_ATTRIBUTE(dllexport) && \ + CURL_HAS_DECLSPEC_ATTRIBUTE(dllimport)) +# ifdef BUILDING_LIBCURL +# define CURL_EXTERN __declspec(dllexport) +# else +# define CURL_EXTERN __declspec(dllimport) +# endif +#elif defined(BUILDING_LIBCURL) && defined(CURL_HIDDEN_SYMBOLS) +# define CURL_EXTERN CURL_EXTERN_SYMBOL +#else +# define CURL_EXTERN +#endif + +#ifndef curl_socket_typedef +/* socket typedef */ +#if defined(_WIN32) && !defined(__LWIP_OPT_H__) && !defined(LWIP_HDR_OPT_H) +typedef SOCKET curl_socket_t; +#define CURL_SOCKET_BAD INVALID_SOCKET +#else +typedef int curl_socket_t; +#define CURL_SOCKET_BAD -1 +#endif +#define curl_socket_typedef +#endif /* curl_socket_typedef */ + +/* enum for the different supported SSL backends */ +typedef enum { + CURLSSLBACKEND_NONE = 0, + CURLSSLBACKEND_OPENSSL = 1, + CURLSSLBACKEND_GNUTLS = 2, + CURLSSLBACKEND_NSS CURL_DEPRECATED(8.3.0, "") = 3, + CURLSSLBACKEND_OBSOLETE4 = 4, /* Was QSOSSL. */ + CURLSSLBACKEND_GSKIT CURL_DEPRECATED(8.3.0, "") = 5, + CURLSSLBACKEND_POLARSSL CURL_DEPRECATED(7.69.0, "") = 6, + CURLSSLBACKEND_WOLFSSL = 7, + CURLSSLBACKEND_SCHANNEL = 8, + CURLSSLBACKEND_SECURETRANSPORT CURL_DEPRECATED(8.15.0, "") = 9, + CURLSSLBACKEND_AXTLS CURL_DEPRECATED(7.61.0, "") = 10, + CURLSSLBACKEND_MBEDTLS = 11, + CURLSSLBACKEND_MESALINK CURL_DEPRECATED(7.82.0, "") = 12, + CURLSSLBACKEND_BEARSSL CURL_DEPRECATED(8.15.0, "") = 13, + CURLSSLBACKEND_RUSTLS = 14 +} curl_sslbackend; + +/* aliases for library clones and renames */ +#define CURLSSLBACKEND_AWSLC CURLSSLBACKEND_OPENSSL +#define CURLSSLBACKEND_BORINGSSL CURLSSLBACKEND_OPENSSL +#define CURLSSLBACKEND_LIBRESSL CURLSSLBACKEND_OPENSSL + +/* deprecated names: */ +#define CURLSSLBACKEND_CYASSL CURLSSLBACKEND_WOLFSSL +#define CURLSSLBACKEND_DARWINSSL CURLSSLBACKEND_SECURETRANSPORT + +/* bits for the CURLOPT_FOLLOWLOCATION option */ +#define CURLFOLLOW_ALL 1L /* generic follow redirects */ + +/* Do not use the custom method in the follow-up request if the HTTP code + instructs so (301, 302, 303). */ +#define CURLFOLLOW_OBEYCODE 2L + +/* Only use the custom method in the first request, always reset in the next */ +#define CURLFOLLOW_FIRSTONLY 3L + +struct curl_httppost { + struct curl_httppost *next; /* next entry in the list */ + char *name; /* pointer to allocated name */ + long namelength; /* length of name length */ + char *contents; /* pointer to allocated data contents */ + long contentslength; /* length of contents field, see also + CURL_HTTPPOST_LARGE */ + char *buffer; /* pointer to allocated buffer contents */ + long bufferlength; /* length of buffer field */ + char *contenttype; /* Content-Type */ + struct curl_slist *contentheader; /* list of extra headers for this form */ + struct curl_httppost *more; /* if one field name has more than one + file, this link should link to following + files */ + long flags; /* as defined below */ + +/* specified content is a filename */ +#define CURL_HTTPPOST_FILENAME (1 << 0) +/* specified content is a filename */ +#define CURL_HTTPPOST_READFILE (1 << 1) +/* name is only stored pointer do not free in formfree */ +#define CURL_HTTPPOST_PTRNAME (1 << 2) +/* contents is only stored pointer do not free in formfree */ +#define CURL_HTTPPOST_PTRCONTENTS (1 << 3) +/* upload file from buffer */ +#define CURL_HTTPPOST_BUFFER (1 << 4) +/* upload file from pointer contents */ +#define CURL_HTTPPOST_PTRBUFFER (1 << 5) +/* upload file contents by using the regular read callback to get the data and + pass the given pointer as custom pointer */ +#define CURL_HTTPPOST_CALLBACK (1 << 6) +/* use size in 'contentlen', added in 7.46.0 */ +#define CURL_HTTPPOST_LARGE (1 << 7) + + char *showfilename; /* The filename to show. If not set, the + actual filename will be used (if this + is a file part) */ + void *userp; /* custom pointer used for + HTTPPOST_CALLBACK posts */ + curl_off_t contentlen; /* alternative length of contents + field. Used if CURL_HTTPPOST_LARGE is + set. Added in 7.46.0 */ +}; + +/* This is a return code for the progress callback that, when returned, will + signal libcurl to continue executing the default progress function */ +#define CURL_PROGRESSFUNC_CONTINUE 0x10000001 + +/* This is the CURLOPT_PROGRESSFUNCTION callback prototype. It is now + considered deprecated but was the only choice up until 7.31.0 */ +typedef int (*curl_progress_callback)(void *clientp, + double dltotal, + double dlnow, + double ultotal, + double ulnow); + +/* This is the CURLOPT_XFERINFOFUNCTION callback prototype. It was introduced + in 7.32.0, avoids the use of floating point numbers and provides more + detailed information. */ +typedef int (*curl_xferinfo_callback)(void *clientp, + curl_off_t dltotal, + curl_off_t dlnow, + curl_off_t ultotal, + curl_off_t ulnow); + +#ifndef CURL_MAX_READ_SIZE + /* The maximum receive buffer size configurable via CURLOPT_BUFFERSIZE. */ +#define CURL_MAX_READ_SIZE (10*1024*1024) +#endif + +#ifndef CURL_MAX_WRITE_SIZE + /* Tests have proven that 20K is a bad buffer size for uploads on Windows, + while 16K for some odd reason performed a lot better. We do the ifndef + check to allow this value to easier be changed at build time for those + who feel adventurous. The practical minimum is about 400 bytes since + libcurl uses a buffer of this size as a scratch area (unrelated to + network send operations). */ +#define CURL_MAX_WRITE_SIZE 16384 +#endif + +#ifndef CURL_MAX_HTTP_HEADER +/* The only reason to have a max limit for this is to avoid the risk of a bad + server feeding libcurl with a never-ending header that will cause reallocs + infinitely */ +#define CURL_MAX_HTTP_HEADER (100*1024) +#endif + +/* This is a magic return code for the write callback that, when returned, + will signal libcurl to pause receiving on the current transfer. */ +#define CURL_WRITEFUNC_PAUSE 0x10000001 + +/* This is a magic return code for the write callback that, when returned, + will signal an error from the callback. */ +#define CURL_WRITEFUNC_ERROR 0xFFFFFFFF + +typedef size_t (*curl_write_callback)(char *buffer, + size_t size, + size_t nitems, + void *outstream); + +/* This callback will be called when a new resolver request is made */ +typedef int (*curl_resolver_start_callback)(void *resolver_state, + void *reserved, void *userdata); + +/* enumeration of file types */ +typedef enum { + CURLFILETYPE_FILE = 0, + CURLFILETYPE_DIRECTORY, + CURLFILETYPE_SYMLINK, + CURLFILETYPE_DEVICE_BLOCK, + CURLFILETYPE_DEVICE_CHAR, + CURLFILETYPE_NAMEDPIPE, + CURLFILETYPE_SOCKET, + CURLFILETYPE_DOOR, /* is possible only on Sun Solaris now */ + + CURLFILETYPE_UNKNOWN /* should never occur */ +} curlfiletype; + +#define CURLFINFOFLAG_KNOWN_FILENAME (1 << 0) +#define CURLFINFOFLAG_KNOWN_FILETYPE (1 << 1) +#define CURLFINFOFLAG_KNOWN_TIME (1 << 2) +#define CURLFINFOFLAG_KNOWN_PERM (1 << 3) +#define CURLFINFOFLAG_KNOWN_UID (1 << 4) +#define CURLFINFOFLAG_KNOWN_GID (1 << 5) +#define CURLFINFOFLAG_KNOWN_SIZE (1 << 6) +#define CURLFINFOFLAG_KNOWN_HLINKCOUNT (1 << 7) + +/* Information about a single file, used when doing FTP wildcard matching */ +struct curl_fileinfo { + char *filename; + curlfiletype filetype; + time_t time; /* always zero! */ + unsigned int perm; + int uid; + int gid; + curl_off_t size; + long int hardlinks; + + struct { + /* If some of these fields is not NULL, it is a pointer to b_data. */ + char *time; + char *perm; + char *user; + char *group; + char *target; /* pointer to the target filename of a symlink */ + } strings; + + unsigned int flags; + + /* These are libcurl private struct fields. Previously used by libcurl, so + they must never be interfered with. */ + char *b_data; + size_t b_size; + size_t b_used; +}; + +/* return codes for CURLOPT_CHUNK_BGN_FUNCTION */ +#define CURL_CHUNK_BGN_FUNC_OK 0 +#define CURL_CHUNK_BGN_FUNC_FAIL 1 /* tell the lib to end the task */ +#define CURL_CHUNK_BGN_FUNC_SKIP 2 /* skip this chunk over */ + +/* if splitting of data transfer is enabled, this callback is called before + download of an individual chunk started. Note that parameter "remains" works + only for FTP wildcard downloading (for now), otherwise is not used */ +typedef long (*curl_chunk_bgn_callback)(const void *transfer_info, + void *ptr, + int remains); + +/* return codes for CURLOPT_CHUNK_END_FUNCTION */ +#define CURL_CHUNK_END_FUNC_OK 0 +#define CURL_CHUNK_END_FUNC_FAIL 1 /* tell the lib to end the task */ + +/* If splitting of data transfer is enabled this callback is called after + download of an individual chunk finished. + Note! After this callback was set then it have to be called FOR ALL chunks. + Even if downloading of this chunk was skipped in CHUNK_BGN_FUNC. + This is the reason why we do not need "transfer_info" parameter in this + callback and we are not interested in "remains" parameter too. */ +typedef long (*curl_chunk_end_callback)(void *ptr); + +/* return codes for FNMATCHFUNCTION */ +#define CURL_FNMATCHFUNC_MATCH 0 /* string corresponds to the pattern */ +#define CURL_FNMATCHFUNC_NOMATCH 1 /* pattern does not match the string */ +#define CURL_FNMATCHFUNC_FAIL 2 /* an error occurred */ + +/* callback type for wildcard downloading pattern matching. If the + string matches the pattern, return CURL_FNMATCHFUNC_MATCH value, etc. */ +typedef int (*curl_fnmatch_callback)(void *ptr, + const char *pattern, + const char *string); + +/* These are the return codes for the seek callbacks */ +#define CURL_SEEKFUNC_OK 0 +#define CURL_SEEKFUNC_FAIL 1 /* fail the entire transfer */ +#define CURL_SEEKFUNC_CANTSEEK 2 /* tell libcurl seeking cannot be done, so + libcurl might try other means instead */ +typedef int (*curl_seek_callback)(void *instream, + curl_off_t offset, + int origin); /* 'whence' */ + +/* This is a return code for the read callback that, when returned, will + signal libcurl to immediately abort the current transfer. */ +#define CURL_READFUNC_ABORT 0x10000000 +/* This is a return code for the read callback that, when returned, will + signal libcurl to pause sending data on the current transfer. */ +#define CURL_READFUNC_PAUSE 0x10000001 + +/* Return code for when the trailing headers' callback has terminated + without any errors */ +#define CURL_TRAILERFUNC_OK 0 +/* Return code for when was an error in the trailing header's list and we + want to abort the request */ +#define CURL_TRAILERFUNC_ABORT 1 + +typedef size_t (*curl_read_callback)(char *buffer, + size_t size, + size_t nitems, + void *instream); + +typedef int (*curl_trailer_callback)(struct curl_slist **list, + void *userdata); + +typedef enum { + CURLSOCKTYPE_IPCXN, /* socket created for a specific IP connection */ + CURLSOCKTYPE_ACCEPT, /* socket created by accept() call */ + CURLSOCKTYPE_LAST /* never use */ +} curlsocktype; + +/* The return code from the sockopt_callback can signal information back + to libcurl: */ +#define CURL_SOCKOPT_OK 0 +#define CURL_SOCKOPT_ERROR 1 /* causes libcurl to abort and return + CURLE_ABORTED_BY_CALLBACK */ +#define CURL_SOCKOPT_ALREADY_CONNECTED 2 + +typedef int (*curl_sockopt_callback)(void *clientp, + curl_socket_t curlfd, + curlsocktype purpose); + +struct curl_sockaddr { + int family; + int socktype; + int protocol; + unsigned int addrlen; /* addrlen was a socklen_t type before 7.18.0 but it + turned really ugly and painful on the systems that + lack this type */ + struct sockaddr addr; +}; + +typedef curl_socket_t +(*curl_opensocket_callback)(void *clientp, + curlsocktype purpose, + struct curl_sockaddr *address); + +typedef int +(*curl_closesocket_callback)(void *clientp, curl_socket_t item); + +typedef enum { + CURLIOE_OK, /* I/O operation successful */ + CURLIOE_UNKNOWNCMD, /* command was unknown to callback */ + CURLIOE_FAILRESTART, /* failed to restart the read */ + CURLIOE_LAST /* never use */ +} curlioerr; + +typedef enum { + CURLIOCMD_NOP, /* no operation */ + CURLIOCMD_RESTARTREAD, /* restart the read stream from start */ + CURLIOCMD_LAST /* never use */ +} curliocmd; + +typedef curlioerr (*curl_ioctl_callback)(CURL *handle, + int cmd, + void *clientp); + +#ifndef CURL_DID_MEMORY_FUNC_TYPEDEFS +/* + * The following typedef's are signatures of malloc, free, realloc, strdup and + * calloc respectively. Function pointers of these types can be passed to the + * curl_global_init_mem() function to set user defined memory management + * callback routines. + */ +typedef void *(*curl_malloc_callback)(size_t size); +typedef void (*curl_free_callback)(void *ptr); +typedef void *(*curl_realloc_callback)(void *ptr, size_t size); +typedef char *(*curl_strdup_callback)(const char *str); +typedef void *(*curl_calloc_callback)(size_t nmemb, size_t size); + +#define CURL_DID_MEMORY_FUNC_TYPEDEFS +#endif + +/* the kind of data that is passed to information_callback */ +typedef enum { + CURLINFO_TEXT = 0, + CURLINFO_HEADER_IN, /* 1 */ + CURLINFO_HEADER_OUT, /* 2 */ + CURLINFO_DATA_IN, /* 3 */ + CURLINFO_DATA_OUT, /* 4 */ + CURLINFO_SSL_DATA_IN, /* 5 */ + CURLINFO_SSL_DATA_OUT, /* 6 */ + CURLINFO_END +} curl_infotype; + +typedef int (*curl_debug_callback) + (CURL *handle, /* the handle/transfer this concerns */ + curl_infotype type, /* what kind of data */ + char *data, /* points to the data */ + size_t size, /* size of the data pointed to */ + void *userptr); /* whatever the user please */ + +/* This is the CURLOPT_PREREQFUNCTION callback prototype. */ +typedef int (*curl_prereq_callback)(void *clientp, + char *conn_primary_ip, + char *conn_local_ip, + int conn_primary_port, + int conn_local_port); + +/* Return code for when the pre-request callback has terminated without + any errors */ +#define CURL_PREREQFUNC_OK 0 +/* Return code for when the pre-request callback wants to abort the + request */ +#define CURL_PREREQFUNC_ABORT 1 + +/* All possible error codes from all sorts of curl functions. Future versions + may return other values, stay prepared. + + Always add new return codes last. Never *EVER* remove any. The return + codes must remain the same! + */ + +typedef enum { + CURLE_OK = 0, + CURLE_UNSUPPORTED_PROTOCOL, /* 1 */ + CURLE_FAILED_INIT, /* 2 */ + CURLE_URL_MALFORMAT, /* 3 */ + CURLE_NOT_BUILT_IN, /* 4 - [was obsoleted in August 2007 for + 7.17.0, reused in April 2011 for 7.21.5] */ + CURLE_COULDNT_RESOLVE_PROXY, /* 5 */ + CURLE_COULDNT_RESOLVE_HOST, /* 6 */ + CURLE_COULDNT_CONNECT, /* 7 */ + CURLE_WEIRD_SERVER_REPLY, /* 8 */ + CURLE_REMOTE_ACCESS_DENIED, /* 9 a service was denied by the server + due to lack of access - when login fails + this is not returned. */ + CURLE_FTP_ACCEPT_FAILED, /* 10 - [was obsoleted in April 2006 for + 7.15.4, reused in Dec 2011 for 7.24.0]*/ + CURLE_FTP_WEIRD_PASS_REPLY, /* 11 */ + CURLE_FTP_ACCEPT_TIMEOUT, /* 12 - timeout occurred accepting server + [was obsoleted in August 2007 for 7.17.0, + reused in Dec 2011 for 7.24.0]*/ + CURLE_FTP_WEIRD_PASV_REPLY, /* 13 */ + CURLE_FTP_WEIRD_227_FORMAT, /* 14 */ + CURLE_FTP_CANT_GET_HOST, /* 15 */ + CURLE_HTTP2, /* 16 - A problem in the http2 framing layer. + [was obsoleted in August 2007 for 7.17.0, + reused in July 2014 for 7.38.0] */ + CURLE_FTP_COULDNT_SET_TYPE, /* 17 */ + CURLE_PARTIAL_FILE, /* 18 */ + CURLE_FTP_COULDNT_RETR_FILE, /* 19 */ + CURLE_OBSOLETE20, /* 20 - NOT USED */ + CURLE_QUOTE_ERROR, /* 21 - quote command failure */ + CURLE_HTTP_RETURNED_ERROR, /* 22 */ + CURLE_WRITE_ERROR, /* 23 */ + CURLE_OBSOLETE24, /* 24 - NOT USED */ + CURLE_UPLOAD_FAILED, /* 25 - failed upload "command" */ + CURLE_READ_ERROR, /* 26 - could not open/read from file */ + CURLE_OUT_OF_MEMORY, /* 27 */ + CURLE_OPERATION_TIMEDOUT, /* 28 - the timeout time was reached */ + CURLE_OBSOLETE29, /* 29 - NOT USED */ + CURLE_FTP_PORT_FAILED, /* 30 - FTP PORT operation failed */ + CURLE_FTP_COULDNT_USE_REST, /* 31 - the REST command failed */ + CURLE_OBSOLETE32, /* 32 - NOT USED */ + CURLE_RANGE_ERROR, /* 33 - RANGE "command" did not work */ + CURLE_OBSOLETE34, /* 34 */ + CURLE_SSL_CONNECT_ERROR, /* 35 - wrong when connecting with SSL */ + CURLE_BAD_DOWNLOAD_RESUME, /* 36 - could not resume download */ + CURLE_FILE_COULDNT_READ_FILE, /* 37 */ + CURLE_LDAP_CANNOT_BIND, /* 38 */ + CURLE_LDAP_SEARCH_FAILED, /* 39 */ + CURLE_OBSOLETE40, /* 40 - NOT USED */ + CURLE_OBSOLETE41, /* 41 - NOT USED starting with 7.53.0 */ + CURLE_ABORTED_BY_CALLBACK, /* 42 */ + CURLE_BAD_FUNCTION_ARGUMENT, /* 43 */ + CURLE_OBSOLETE44, /* 44 - NOT USED */ + CURLE_INTERFACE_FAILED, /* 45 - CURLOPT_INTERFACE failed */ + CURLE_OBSOLETE46, /* 46 - NOT USED */ + CURLE_TOO_MANY_REDIRECTS, /* 47 - catch endless re-direct loops */ + CURLE_UNKNOWN_OPTION, /* 48 - User specified an unknown option */ + CURLE_SETOPT_OPTION_SYNTAX, /* 49 - Malformed setopt option */ + CURLE_OBSOLETE50, /* 50 - NOT USED */ + CURLE_OBSOLETE51, /* 51 - NOT USED */ + CURLE_GOT_NOTHING, /* 52 - when this is a specific error */ + CURLE_SSL_ENGINE_NOTFOUND, /* 53 - SSL crypto engine not found */ + CURLE_SSL_ENGINE_SETFAILED, /* 54 - can not set SSL crypto engine as + default */ + CURLE_SEND_ERROR, /* 55 - failed sending network data */ + CURLE_RECV_ERROR, /* 56 - failure in receiving network data */ + CURLE_OBSOLETE57, /* 57 - NOT IN USE */ + CURLE_SSL_CERTPROBLEM, /* 58 - problem with the local certificate */ + CURLE_SSL_CIPHER, /* 59 - could not use specified cipher */ + CURLE_PEER_FAILED_VERIFICATION, /* 60 - peer's certificate or fingerprint + was not verified fine */ + CURLE_BAD_CONTENT_ENCODING, /* 61 - Unrecognized/bad encoding */ + CURLE_OBSOLETE62, /* 62 - NOT IN USE since 7.82.0 */ + CURLE_FILESIZE_EXCEEDED, /* 63 - Maximum file size exceeded */ + CURLE_USE_SSL_FAILED, /* 64 - Requested FTP SSL level failed */ + CURLE_SEND_FAIL_REWIND, /* 65 - Sending the data requires a rewind + that failed */ + CURLE_SSL_ENGINE_INITFAILED, /* 66 - failed to initialise ENGINE */ + CURLE_LOGIN_DENIED, /* 67 - user, password or similar was not + accepted and we failed to login */ + CURLE_TFTP_NOTFOUND, /* 68 - file not found on server */ + CURLE_TFTP_PERM, /* 69 - permission problem on server */ + CURLE_REMOTE_DISK_FULL, /* 70 - out of disk space on server */ + CURLE_TFTP_ILLEGAL, /* 71 - Illegal TFTP operation */ + CURLE_TFTP_UNKNOWNID, /* 72 - Unknown transfer ID */ + CURLE_REMOTE_FILE_EXISTS, /* 73 - File already exists */ + CURLE_TFTP_NOSUCHUSER, /* 74 - No such user */ + CURLE_OBSOLETE75, /* 75 - NOT IN USE since 7.82.0 */ + CURLE_OBSOLETE76, /* 76 - NOT IN USE since 7.82.0 */ + CURLE_SSL_CACERT_BADFILE, /* 77 - could not load CACERT file, missing + or wrong format */ + CURLE_REMOTE_FILE_NOT_FOUND, /* 78 - remote file not found */ + CURLE_SSH, /* 79 - error from the SSH layer, somewhat + generic so the error message will be of + interest when this has happened */ + + CURLE_SSL_SHUTDOWN_FAILED, /* 80 - Failed to shut down the SSL + connection */ + CURLE_AGAIN, /* 81 - socket is not ready for send/recv, + wait till it is ready and try again (Added + in 7.18.2) */ + CURLE_SSL_CRL_BADFILE, /* 82 - could not load CRL file, missing or + wrong format (Added in 7.19.0) */ + CURLE_SSL_ISSUER_ERROR, /* 83 - Issuer check failed. (Added in + 7.19.0) */ + CURLE_FTP_PRET_FAILED, /* 84 - a PRET command failed */ + CURLE_RTSP_CSEQ_ERROR, /* 85 - mismatch of RTSP CSeq numbers */ + CURLE_RTSP_SESSION_ERROR, /* 86 - mismatch of RTSP Session Ids */ + CURLE_FTP_BAD_FILE_LIST, /* 87 - unable to parse FTP file list */ + CURLE_CHUNK_FAILED, /* 88 - chunk callback reported error */ + CURLE_NO_CONNECTION_AVAILABLE, /* 89 - No connection available, the + session will be queued */ + CURLE_SSL_PINNEDPUBKEYNOTMATCH, /* 90 - specified pinned public key did not + match */ + CURLE_SSL_INVALIDCERTSTATUS, /* 91 - invalid certificate status */ + CURLE_HTTP2_STREAM, /* 92 - stream error in HTTP/2 framing layer + */ + CURLE_RECURSIVE_API_CALL, /* 93 - an api function was called from + inside a callback */ + CURLE_AUTH_ERROR, /* 94 - an authentication function returned an + error */ + CURLE_HTTP3, /* 95 - An HTTP/3 layer problem */ + CURLE_QUIC_CONNECT_ERROR, /* 96 - QUIC connection error */ + CURLE_PROXY, /* 97 - proxy handshake error */ + CURLE_SSL_CLIENTCERT, /* 98 - client-side certificate required */ + CURLE_UNRECOVERABLE_POLL, /* 99 - poll/select returned fatal error */ + CURLE_TOO_LARGE, /* 100 - a value/data met its maximum */ + CURLE_ECH_REQUIRED, /* 101 - ECH tried but failed */ + CURL_LAST /* never use! */ +} CURLcode; + +#ifndef CURL_NO_OLDIES /* define this to test if your app builds with all + the obsolete stuff removed! */ + +/* removed in 7.53.0 */ +#define CURLE_FUNCTION_NOT_FOUND CURLE_OBSOLETE41 + +/* removed in 7.56.0 */ +#define CURLE_HTTP_POST_ERROR CURLE_OBSOLETE34 + +/* Previously obsolete error code reused in 7.38.0 */ +#define CURLE_OBSOLETE16 CURLE_HTTP2 + +/* Previously obsolete error codes reused in 7.24.0 */ +#define CURLE_OBSOLETE10 CURLE_FTP_ACCEPT_FAILED +#define CURLE_OBSOLETE12 CURLE_FTP_ACCEPT_TIMEOUT + +/* compatibility with older names */ +#define CURLOPT_ENCODING CURLOPT_ACCEPT_ENCODING +#define CURLE_FTP_WEIRD_SERVER_REPLY CURLE_WEIRD_SERVER_REPLY + +/* The following were added in 7.62.0 */ +#define CURLE_SSL_CACERT CURLE_PEER_FAILED_VERIFICATION + +/* The following were added in 7.21.5, April 2011 */ +#define CURLE_UNKNOWN_TELNET_OPTION CURLE_UNKNOWN_OPTION + +/* Added for 7.78.0 */ +#define CURLE_TELNET_OPTION_SYNTAX CURLE_SETOPT_OPTION_SYNTAX + +/* The following were added in 7.17.1 */ +/* These are scheduled to disappear by 2009 */ +#define CURLE_SSL_PEER_CERTIFICATE CURLE_PEER_FAILED_VERIFICATION + +/* The following were added in 7.17.0 */ +/* These are scheduled to disappear by 2009 */ +#define CURLE_OBSOLETE CURLE_OBSOLETE50 /* no one should be using this! */ +#define CURLE_BAD_PASSWORD_ENTERED CURLE_OBSOLETE46 +#define CURLE_BAD_CALLING_ORDER CURLE_OBSOLETE44 +#define CURLE_FTP_USER_PASSWORD_INCORRECT CURLE_OBSOLETE10 +#define CURLE_FTP_CANT_RECONNECT CURLE_OBSOLETE16 +#define CURLE_FTP_COULDNT_GET_SIZE CURLE_OBSOLETE32 +#define CURLE_FTP_COULDNT_SET_ASCII CURLE_OBSOLETE29 +#define CURLE_FTP_WEIRD_USER_REPLY CURLE_OBSOLETE12 +#define CURLE_FTP_WRITE_ERROR CURLE_OBSOLETE20 +#define CURLE_LIBRARY_NOT_FOUND CURLE_OBSOLETE40 +#define CURLE_MALFORMAT_USER CURLE_OBSOLETE24 +#define CURLE_SHARE_IN_USE CURLE_OBSOLETE57 +#define CURLE_URL_MALFORMAT_USER CURLE_NOT_BUILT_IN + +#define CURLE_FTP_ACCESS_DENIED CURLE_REMOTE_ACCESS_DENIED +#define CURLE_FTP_COULDNT_SET_BINARY CURLE_FTP_COULDNT_SET_TYPE +#define CURLE_FTP_QUOTE_ERROR CURLE_QUOTE_ERROR +#define CURLE_TFTP_DISKFULL CURLE_REMOTE_DISK_FULL +#define CURLE_TFTP_EXISTS CURLE_REMOTE_FILE_EXISTS +#define CURLE_HTTP_RANGE_ERROR CURLE_RANGE_ERROR +#define CURLE_FTP_SSL_FAILED CURLE_USE_SSL_FAILED + +/* The following were added earlier */ + +#define CURLE_OPERATION_TIMEOUTED CURLE_OPERATION_TIMEDOUT +#define CURLE_HTTP_NOT_FOUND CURLE_HTTP_RETURNED_ERROR +#define CURLE_HTTP_PORT_FAILED CURLE_INTERFACE_FAILED +#define CURLE_FTP_COULDNT_STOR_FILE CURLE_UPLOAD_FAILED +#define CURLE_FTP_PARTIAL_FILE CURLE_PARTIAL_FILE +#define CURLE_FTP_BAD_DOWNLOAD_RESUME CURLE_BAD_DOWNLOAD_RESUME +#define CURLE_LDAP_INVALID_URL CURLE_OBSOLETE62 +#define CURLE_CONV_REQD CURLE_OBSOLETE76 +#define CURLE_CONV_FAILED CURLE_OBSOLETE75 + +/* This was the error code 50 in 7.7.3 and a few earlier versions, this + is no longer used by libcurl but is instead #defined here only to not + make programs break */ +#define CURLE_ALREADY_COMPLETE 99999 + +/* Provide defines for really old option names */ +#define CURLOPT_FILE CURLOPT_WRITEDATA /* name changed in 7.9.7 */ +#define CURLOPT_INFILE CURLOPT_READDATA /* name changed in 7.9.7 */ +#define CURLOPT_WRITEHEADER CURLOPT_HEADERDATA + +/* Since long deprecated options with no code in the lib that does anything + with them. */ +#define CURLOPT_WRITEINFO CURLOPT_OBSOLETE40 +#define CURLOPT_CLOSEPOLICY CURLOPT_OBSOLETE72 +#define CURLOPT_OBSOLETE72 9999 +#define CURLOPT_OBSOLETE40 9999 + +#endif /* !CURL_NO_OLDIES */ + +/* + * Proxy error codes. Returned in CURLINFO_PROXY_ERROR if CURLE_PROXY was + * return for the transfers. + */ +typedef enum { + CURLPX_OK, + CURLPX_BAD_ADDRESS_TYPE, + CURLPX_BAD_VERSION, + CURLPX_CLOSED, + CURLPX_GSSAPI, + CURLPX_GSSAPI_PERMSG, + CURLPX_GSSAPI_PROTECTION, + CURLPX_IDENTD, + CURLPX_IDENTD_DIFFER, + CURLPX_LONG_HOSTNAME, + CURLPX_LONG_PASSWD, + CURLPX_LONG_USER, + CURLPX_NO_AUTH, + CURLPX_RECV_ADDRESS, + CURLPX_RECV_AUTH, + CURLPX_RECV_CONNECT, + CURLPX_RECV_REQACK, + CURLPX_REPLY_ADDRESS_TYPE_NOT_SUPPORTED, + CURLPX_REPLY_COMMAND_NOT_SUPPORTED, + CURLPX_REPLY_CONNECTION_REFUSED, + CURLPX_REPLY_GENERAL_SERVER_FAILURE, + CURLPX_REPLY_HOST_UNREACHABLE, + CURLPX_REPLY_NETWORK_UNREACHABLE, + CURLPX_REPLY_NOT_ALLOWED, + CURLPX_REPLY_TTL_EXPIRED, + CURLPX_REPLY_UNASSIGNED, + CURLPX_REQUEST_FAILED, + CURLPX_RESOLVE_HOST, + CURLPX_SEND_AUTH, + CURLPX_SEND_CONNECT, + CURLPX_SEND_REQUEST, + CURLPX_UNKNOWN_FAIL, + CURLPX_UNKNOWN_MODE, + CURLPX_USER_REJECTED, + CURLPX_LAST /* never use */ +} CURLproxycode; + +/* This prototype applies to all conversion callbacks */ +typedef CURLcode (*curl_conv_callback)(char *buffer, size_t length); + +typedef CURLcode (*curl_ssl_ctx_callback)(CURL *curl, /* easy handle */ + void *ssl_ctx, /* actually an OpenSSL + or wolfSSL SSL_CTX, + or an mbedTLS + mbedtls_ssl_config */ + void *userptr); + +#define CURLPROXY_HTTP 0L /* added in 7.10, new in 7.19.4 default is + to use CONNECT HTTP/1.1 */ +#define CURLPROXY_HTTP_1_0 1L /* force to use CONNECT HTTP/1.0 + added in 7.19.4 */ +#define CURLPROXY_HTTPS 2L /* HTTPS but stick to HTTP/1 + added in 7.52.0 */ +#define CURLPROXY_HTTPS2 3L /* HTTPS and attempt HTTP/2 + added in 8.2.0 */ +#define CURLPROXY_SOCKS4 4L /* support added in 7.15.2, enum existed + already in 7.10 */ +#define CURLPROXY_SOCKS5 5L /* added in 7.10 */ +#define CURLPROXY_SOCKS4A 6L /* added in 7.18.0 */ +#define CURLPROXY_SOCKS5_HOSTNAME 7L /* Use the SOCKS5 protocol but pass along + the hostname rather than the IP + address. added in 7.18.0 */ + +typedef enum { + CURLPROXY_LAST = 8 /* never use */ +} curl_proxytype; /* this enum was added in 7.10 */ + +/* + * Bitmasks for CURLOPT_HTTPAUTH and CURLOPT_PROXYAUTH options: + * + * CURLAUTH_NONE - No HTTP authentication + * CURLAUTH_BASIC - HTTP Basic authentication (default) + * CURLAUTH_DIGEST - HTTP Digest authentication + * CURLAUTH_NEGOTIATE - HTTP Negotiate (SPNEGO) authentication + * CURLAUTH_GSSNEGOTIATE - Alias for CURLAUTH_NEGOTIATE (deprecated) + * CURLAUTH_NTLM - HTTP NTLM authentication + * CURLAUTH_DIGEST_IE - HTTP Digest authentication with IE flavour + * CURLAUTH_NTLM_WB - HTTP NTLM authentication delegated to winbind helper + * CURLAUTH_BEARER - HTTP Bearer token authentication + * CURLAUTH_ONLY - Use together with a single other type to force no + * authentication or just that single type + * CURLAUTH_ANY - All fine types set + * CURLAUTH_ANYSAFE - All fine types except Basic + */ + +#define CURLAUTH_NONE ((unsigned long)0) +#define CURLAUTH_BASIC (((unsigned long)1) << 0) +#define CURLAUTH_DIGEST (((unsigned long)1) << 1) +#define CURLAUTH_NEGOTIATE (((unsigned long)1) << 2) +/* Deprecated since the advent of CURLAUTH_NEGOTIATE */ +#define CURLAUTH_GSSNEGOTIATE CURLAUTH_NEGOTIATE +/* Used for CURLOPT_SOCKS5_AUTH to stay terminologically correct */ +#define CURLAUTH_GSSAPI CURLAUTH_NEGOTIATE +#define CURLAUTH_NTLM (((unsigned long)1) << 3) +#define CURLAUTH_DIGEST_IE (((unsigned long)1) << 4) +#ifndef CURL_NO_OLDIES + /* functionality removed since 8.8.0 */ +#define CURLAUTH_NTLM_WB (((unsigned long)1) << 5) +#endif +#define CURLAUTH_BEARER (((unsigned long)1) << 6) +#define CURLAUTH_AWS_SIGV4 (((unsigned long)1) << 7) +#define CURLAUTH_ONLY (((unsigned long)1) << 31) +#define CURLAUTH_ANY (~CURLAUTH_DIGEST_IE) +#define CURLAUTH_ANYSAFE (~(CURLAUTH_BASIC | CURLAUTH_DIGEST_IE)) + +#define CURLSSH_AUTH_ANY ~0L /* all types supported by server */ +#define CURLSSH_AUTH_NONE 0L /* none allowed, silly but complete */ +#define CURLSSH_AUTH_PUBLICKEY (1L << 0) /* public/private key files */ +#define CURLSSH_AUTH_PASSWORD (1L << 1) /* password */ +#define CURLSSH_AUTH_HOST (1L << 2) /* host key files */ +#define CURLSSH_AUTH_KEYBOARD (1L << 3) /* keyboard interactive */ +#define CURLSSH_AUTH_AGENT (1L << 4) /* agent (ssh-agent, pageant...) */ +#define CURLSSH_AUTH_GSSAPI (1L << 5) /* gssapi (kerberos, ...) */ +#define CURLSSH_AUTH_DEFAULT CURLSSH_AUTH_ANY + +#define CURLGSSAPI_DELEGATION_NONE 0L /* no delegation (default) */ +#define CURLGSSAPI_DELEGATION_POLICY_FLAG (1L<<0) /* if permitted by policy */ +#define CURLGSSAPI_DELEGATION_FLAG (1L<<1) /* delegate always */ + +#define CURL_ERROR_SIZE 256 + +enum curl_khtype { + CURLKHTYPE_UNKNOWN, + CURLKHTYPE_RSA1, + CURLKHTYPE_RSA, + CURLKHTYPE_DSS, + CURLKHTYPE_ECDSA, + CURLKHTYPE_ED25519 +}; + +struct curl_khkey { + const char *key; /* points to a null-terminated string encoded with base64 + if len is zero, otherwise to the "raw" data */ + size_t len; + enum curl_khtype keytype; +}; + +/* this is the set of return values expected from the curl_sshkeycallback + callback */ +enum curl_khstat { + CURLKHSTAT_FINE_ADD_TO_FILE, + CURLKHSTAT_FINE, + CURLKHSTAT_REJECT, /* reject the connection, return an error */ + CURLKHSTAT_DEFER, /* do not accept it, but we cannot answer right now. + Causes a CURLE_PEER_FAILED_VERIFICATION error but the + connection will be left intact etc */ + CURLKHSTAT_FINE_REPLACE, /* accept and replace the wrong key */ + CURLKHSTAT_LAST /* not for use, only a marker for last-in-list */ +}; + +/* this is the set of status codes pass in to the callback */ +enum curl_khmatch { + CURLKHMATCH_OK, /* match */ + CURLKHMATCH_MISMATCH, /* host found, key mismatch! */ + CURLKHMATCH_MISSING, /* no matching host/key found */ + CURLKHMATCH_LAST /* not for use, only a marker for last-in-list */ +}; + +typedef int + (*curl_sshkeycallback) (CURL *easy, /* easy handle */ + const struct curl_khkey *knownkey, /* known */ + const struct curl_khkey *foundkey, /* found */ + enum curl_khmatch, /* libcurl's view on the keys */ + void *clientp); /* custom pointer passed with */ + /* CURLOPT_SSH_KEYDATA */ + +typedef int + (*curl_sshhostkeycallback) (void *clientp,/* custom pointer passed */ + /* with CURLOPT_SSH_HOSTKEYDATA */ + int keytype, /* CURLKHTYPE */ + const char *key, /* hostkey to check */ + size_t keylen); /* length of the key */ + /* return CURLE_OK to accept */ + /* or something else to refuse */ + +/* parameter for the CURLOPT_USE_SSL option */ +#define CURLUSESSL_NONE 0L /* do not attempt to use SSL */ +#define CURLUSESSL_TRY 1L /* try using SSL, proceed anyway otherwise */ +#define CURLUSESSL_CONTROL 2L /* SSL for the control connection or fail */ +#define CURLUSESSL_ALL 3L /* SSL for all communication or fail */ + +typedef enum { + CURLUSESSL_LAST = 4 /* not an option, never use */ +} curl_usessl; + +/* Definition of bits for the CURLOPT_SSL_OPTIONS argument: */ + +/* - ALLOW_BEAST tells libcurl to allow the BEAST SSL vulnerability in the + name of improving interoperability with older servers. Some SSL libraries + have introduced work-arounds for this flaw but those work-arounds sometimes + make the SSL communication fail. To regain functionality with those broken + servers, a user can this way allow the vulnerability back. */ +#define CURLSSLOPT_ALLOW_BEAST (1L << 0) + +/* - NO_REVOKE tells libcurl to disable certificate revocation checks for those + SSL backends where such behavior is present. */ +#define CURLSSLOPT_NO_REVOKE (1L << 1) + +/* - NO_PARTIALCHAIN tells libcurl to *NOT* accept a partial certificate chain + if possible. The OpenSSL backend has this ability. */ +#define CURLSSLOPT_NO_PARTIALCHAIN (1L << 2) + +/* - REVOKE_BEST_EFFORT tells libcurl to ignore certificate revocation offline + checks and ignore missing revocation list for those SSL backends where such + behavior is present. */ +#define CURLSSLOPT_REVOKE_BEST_EFFORT (1L << 3) + +/* - CURLSSLOPT_NATIVE_CA tells libcurl to use standard certificate store of + operating system. Currently implemented under MS-Windows. */ +#define CURLSSLOPT_NATIVE_CA (1L << 4) + +/* - CURLSSLOPT_AUTO_CLIENT_CERT tells libcurl to automatically locate and use + a client certificate for authentication. (Schannel) */ +#define CURLSSLOPT_AUTO_CLIENT_CERT (1L << 5) + +/* If possible, send data using TLS 1.3 early data */ +#define CURLSSLOPT_EARLYDATA (1L << 6) + +/* The default connection attempt delay in milliseconds for happy eyeballs. + CURLOPT_HAPPY_EYEBALLS_TIMEOUT_MS.3 and happy-eyeballs-timeout-ms.d document + this value, keep them in sync. */ +#define CURL_HET_DEFAULT 200L + +/* The default connection upkeep interval in milliseconds. */ +#define CURL_UPKEEP_INTERVAL_DEFAULT 60000L + +#ifndef CURL_NO_OLDIES /* define this to test if your app builds with all + the obsolete stuff removed! */ + +/* Backwards compatibility with older names */ +/* These are scheduled to disappear by 2009 */ + +#define CURLFTPSSL_NONE CURLUSESSL_NONE +#define CURLFTPSSL_TRY CURLUSESSL_TRY +#define CURLFTPSSL_CONTROL CURLUSESSL_CONTROL +#define CURLFTPSSL_ALL CURLUSESSL_ALL +#define CURLFTPSSL_LAST CURLUSESSL_LAST +#define curl_ftpssl curl_usessl +#endif /* !CURL_NO_OLDIES */ + +/* parameter for the CURLOPT_FTP_SSL_CCC option */ +#define CURLFTPSSL_CCC_NONE 0L /* do not send CCC */ +#define CURLFTPSSL_CCC_PASSIVE 1L /* Let the server initiate the shutdown */ +#define CURLFTPSSL_CCC_ACTIVE 2L /* Initiate the shutdown */ + +typedef enum { + CURLFTPSSL_CCC_LAST = 3 /* not an option, never use */ +} curl_ftpccc; + +/* parameter for the CURLOPT_FTPSSLAUTH option */ +#define CURLFTPAUTH_DEFAULT 0L /* let libcurl decide */ +#define CURLFTPAUTH_SSL 1L /* use "AUTH SSL" */ +#define CURLFTPAUTH_TLS 2L /* use "AUTH TLS" */ + +typedef enum { + CURLFTPAUTH_LAST = 3 /* not an option, never use */ +} curl_ftpauth; + +/* parameter for the CURLOPT_FTP_CREATE_MISSING_DIRS option */ +#define CURLFTP_CREATE_DIR_NONE 0L /* do NOT create missing dirs! */ +#define CURLFTP_CREATE_DIR 1L /* (FTP/SFTP) if CWD fails, try MKD and + then CWD again if MKD succeeded, for + SFTP this does similar magic */ +#define CURLFTP_CREATE_DIR_RETRY 2L /* (FTP only) if CWD fails, try MKD and + then CWD again even if MKD failed! */ + +typedef enum { + CURLFTP_CREATE_DIR_LAST = 3 /* not an option, never use */ +} curl_ftpcreatedir; + +/* parameter for the CURLOPT_FTP_FILEMETHOD option */ +#define CURLFTPMETHOD_DEFAULT 0L /* let libcurl pick */ +#define CURLFTPMETHOD_MULTICWD 1L /* single CWD operation for each path + part */ +#define CURLFTPMETHOD_NOCWD 2L /* no CWD at all */ +#define CURLFTPMETHOD_SINGLECWD 3L /* one CWD to full dir, then work on file */ + +typedef enum { + CURLFTPMETHOD_LAST = 4 /* not an option, never use */ +} curl_ftpmethod; + +/* bitmask defines for CURLOPT_HEADEROPT */ +#define CURLHEADER_UNIFIED 0L +#define CURLHEADER_SEPARATE (1L << 0) + +/* CURLALTSVC_* are bits for the CURLOPT_ALTSVC_CTRL option */ +#define CURLALTSVC_READONLYFILE (1L << 2) +#define CURLALTSVC_H1 (1L << 3) +#define CURLALTSVC_H2 (1L << 4) +#define CURLALTSVC_H3 (1L << 5) + +/* bitmask values for CURLOPT_UPLOAD_FLAGS */ +#define CURLULFLAG_ANSWERED (1L << 0) +#define CURLULFLAG_DELETED (1L << 1) +#define CURLULFLAG_DRAFT (1L << 2) +#define CURLULFLAG_FLAGGED (1L << 3) +#define CURLULFLAG_SEEN (1L << 4) + +struct curl_hstsentry { + char *name; + size_t namelen; + unsigned int includeSubDomains:1; + char expire[18]; /* YYYYMMDD HH:MM:SS [null-terminated] */ +}; + +struct curl_index { + size_t index; /* the provided entry's "index" or count */ + size_t total; /* total number of entries to save */ +}; + +typedef enum { + CURLSTS_OK, + CURLSTS_DONE, + CURLSTS_FAIL +} CURLSTScode; + +typedef CURLSTScode (*curl_hstsread_callback)(CURL *easy, + struct curl_hstsentry *e, + void *userp); +typedef CURLSTScode (*curl_hstswrite_callback)(CURL *easy, + struct curl_hstsentry *e, + struct curl_index *i, + void *userp); + +/* CURLHSTS_* are bits for the CURLOPT_HSTS option */ +#define CURLHSTS_ENABLE (1L << 0) +#define CURLHSTS_READONLYFILE (1L << 1) + +/* The CURLPROTO_ defines below are for the **deprecated** CURLOPT_*PROTOCOLS + options. Do not use. */ +#define CURLPROTO_HTTP (1L << 0) +#define CURLPROTO_HTTPS (1L << 1) +#define CURLPROTO_FTP (1L << 2) +#define CURLPROTO_FTPS (1L << 3) +#define CURLPROTO_SCP (1L << 4) +#define CURLPROTO_SFTP (1L << 5) +#define CURLPROTO_TELNET (1L << 6) +#define CURLPROTO_LDAP (1L << 7) +#define CURLPROTO_LDAPS (1L << 8) +#define CURLPROTO_DICT (1L << 9) +#define CURLPROTO_FILE (1L << 10) +#define CURLPROTO_TFTP (1L << 11) +#define CURLPROTO_IMAP (1L << 12) +#define CURLPROTO_IMAPS (1L << 13) +#define CURLPROTO_POP3 (1L << 14) +#define CURLPROTO_POP3S (1L << 15) +#define CURLPROTO_SMTP (1L << 16) +#define CURLPROTO_SMTPS (1L << 17) +#define CURLPROTO_RTSP (1L << 18) +#define CURLPROTO_RTMP (1L << 19) +#define CURLPROTO_RTMPT (1L << 20) +#define CURLPROTO_RTMPE (1L << 21) +#define CURLPROTO_RTMPTE (1L << 22) +#define CURLPROTO_RTMPS (1L << 23) +#define CURLPROTO_RTMPTS (1L << 24) +#define CURLPROTO_GOPHER (1L << 25) +#define CURLPROTO_SMB (1L << 26) +#define CURLPROTO_SMBS (1L << 27) +#define CURLPROTO_MQTT (1L << 28) +#define CURLPROTO_GOPHERS (1L << 29) +#define CURLPROTO_ALL (~0L) /* enable everything */ + +/* long may be 32 or 64 bits, but we should never depend on anything else + but 32 */ +#define CURLOPTTYPE_LONG 0 +#define CURLOPTTYPE_OBJECTPOINT 10000 +#define CURLOPTTYPE_FUNCTIONPOINT 20000 +#define CURLOPTTYPE_OFF_T 30000 +#define CURLOPTTYPE_BLOB 40000 + +/* *STRINGPOINT is an alias for OBJECTPOINT to allow tools to extract the + string options from the header file */ + +#define CURLOPT(na,t,nu) na = t + nu +#define CURLOPTDEPRECATED(na,t,nu,v,m) na CURL_DEPRECATED(v,m) = t + nu + +/* CURLOPT aliases that make no runtime difference */ + +/* 'char *' argument to a string with a trailing zero */ +#define CURLOPTTYPE_STRINGPOINT CURLOPTTYPE_OBJECTPOINT + +/* 'struct curl_slist *' argument */ +#define CURLOPTTYPE_SLISTPOINT CURLOPTTYPE_OBJECTPOINT + +/* 'void *' argument passed untouched to callback */ +#define CURLOPTTYPE_CBPOINT CURLOPTTYPE_OBJECTPOINT + +/* 'long' argument with a set of values/bitmask */ +#define CURLOPTTYPE_VALUES CURLOPTTYPE_LONG + +/* + * All CURLOPT_* values. + */ + +typedef enum { + /* This is the FILE * or void * the regular output should be written to. */ + CURLOPT(CURLOPT_WRITEDATA, CURLOPTTYPE_CBPOINT, 1), + + /* The full URL to get/put */ + CURLOPT(CURLOPT_URL, CURLOPTTYPE_STRINGPOINT, 2), + + /* Port number to connect to, if other than default. */ + CURLOPT(CURLOPT_PORT, CURLOPTTYPE_LONG, 3), + + /* Name of proxy to use. */ + CURLOPT(CURLOPT_PROXY, CURLOPTTYPE_STRINGPOINT, 4), + + /* "user:password;options" to use when fetching. */ + CURLOPT(CURLOPT_USERPWD, CURLOPTTYPE_STRINGPOINT, 5), + + /* "user:password" to use with proxy. */ + CURLOPT(CURLOPT_PROXYUSERPWD, CURLOPTTYPE_STRINGPOINT, 6), + + /* Range to get, specified as an ASCII string. */ + CURLOPT(CURLOPT_RANGE, CURLOPTTYPE_STRINGPOINT, 7), + + /* not used */ + + /* Specified file stream to upload from (use as input): */ + CURLOPT(CURLOPT_READDATA, CURLOPTTYPE_CBPOINT, 9), + + /* Buffer to receive error messages in, must be at least CURL_ERROR_SIZE + * bytes big. */ + CURLOPT(CURLOPT_ERRORBUFFER, CURLOPTTYPE_OBJECTPOINT, 10), + + /* Function that will be called to store the output (instead of fwrite). The + * parameters will use fwrite() syntax, make sure to follow them. */ + CURLOPT(CURLOPT_WRITEFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 11), + + /* Function that will be called to read the input (instead of fread). The + * parameters will use fread() syntax, make sure to follow them. */ + CURLOPT(CURLOPT_READFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 12), + + /* Time-out the read operation after this amount of seconds */ + CURLOPT(CURLOPT_TIMEOUT, CURLOPTTYPE_LONG, 13), + + /* If CURLOPT_READDATA is used, this can be used to inform libcurl about + * how large the file being sent really is. That allows better error + * checking and better verifies that the upload was successful. -1 means + * unknown size. + * + * For large file support, there is also a _LARGE version of the key + * which takes an off_t type, allowing platforms with larger off_t + * sizes to handle larger files. See below for INFILESIZE_LARGE. + */ + CURLOPT(CURLOPT_INFILESIZE, CURLOPTTYPE_LONG, 14), + + /* POST static input fields. */ + CURLOPT(CURLOPT_POSTFIELDS, CURLOPTTYPE_OBJECTPOINT, 15), + + /* Set the referrer page (needed by some CGIs) */ + CURLOPT(CURLOPT_REFERER, CURLOPTTYPE_STRINGPOINT, 16), + + /* Set the FTP PORT string (interface name, named or numerical IP address) + Use i.e '-' to use default address. */ + CURLOPT(CURLOPT_FTPPORT, CURLOPTTYPE_STRINGPOINT, 17), + + /* Set the User-Agent string (examined by some CGIs) */ + CURLOPT(CURLOPT_USERAGENT, CURLOPTTYPE_STRINGPOINT, 18), + + /* If the download receives less than "low speed limit" bytes/second + * during "low speed time" seconds, the operations is aborted. + * You could i.e if you have a pretty high speed connection, abort if + * it is less than 2000 bytes/sec during 20 seconds. + */ + + /* Set the "low speed limit" */ + CURLOPT(CURLOPT_LOW_SPEED_LIMIT, CURLOPTTYPE_LONG, 19), + + /* Set the "low speed time" */ + CURLOPT(CURLOPT_LOW_SPEED_TIME, CURLOPTTYPE_LONG, 20), + + /* Set the continuation offset. + * + * Note there is also a _LARGE version of this key which uses + * off_t types, allowing for large file offsets on platforms which + * use larger-than-32-bit off_t's. Look below for RESUME_FROM_LARGE. + */ + CURLOPT(CURLOPT_RESUME_FROM, CURLOPTTYPE_LONG, 21), + + /* Set cookie in request: */ + CURLOPT(CURLOPT_COOKIE, CURLOPTTYPE_STRINGPOINT, 22), + + /* This points to a linked list of headers, struct curl_slist kind. This + list is also used for RTSP (in spite of its name) */ + CURLOPT(CURLOPT_HTTPHEADER, CURLOPTTYPE_SLISTPOINT, 23), + + /* This points to a linked list of post entries, struct curl_httppost */ + CURLOPTDEPRECATED(CURLOPT_HTTPPOST, CURLOPTTYPE_OBJECTPOINT, 24, + 7.56.0, "Use CURLOPT_MIMEPOST"), + + /* name of the file keeping your private SSL-certificate */ + CURLOPT(CURLOPT_SSLCERT, CURLOPTTYPE_STRINGPOINT, 25), + + /* password for the SSL or SSH private key */ + CURLOPT(CURLOPT_KEYPASSWD, CURLOPTTYPE_STRINGPOINT, 26), + + /* send TYPE parameter? */ + CURLOPT(CURLOPT_CRLF, CURLOPTTYPE_LONG, 27), + + /* send linked-list of QUOTE commands */ + CURLOPT(CURLOPT_QUOTE, CURLOPTTYPE_SLISTPOINT, 28), + + /* send FILE * or void * to store headers to, if you use a callback it + is simply passed to the callback unmodified */ + CURLOPT(CURLOPT_HEADERDATA, CURLOPTTYPE_CBPOINT, 29), + + /* point to a file to read the initial cookies from, also enables + "cookie awareness" */ + CURLOPT(CURLOPT_COOKIEFILE, CURLOPTTYPE_STRINGPOINT, 31), + + /* What version to specifically try to use. + See CURL_SSLVERSION defines below. */ + CURLOPT(CURLOPT_SSLVERSION, CURLOPTTYPE_VALUES, 32), + + /* What kind of HTTP time condition to use, see defines */ + CURLOPT(CURLOPT_TIMECONDITION, CURLOPTTYPE_VALUES, 33), + + /* Time to use with the above condition. Specified in number of seconds + since 1 Jan 1970 */ + CURLOPT(CURLOPT_TIMEVALUE, CURLOPTTYPE_LONG, 34), + + /* 35 = OBSOLETE */ + + /* Custom request, for customizing the get command like + HTTP: DELETE, TRACE and others + FTP: to use a different list command + */ + CURLOPT(CURLOPT_CUSTOMREQUEST, CURLOPTTYPE_STRINGPOINT, 36), + + /* FILE handle to use instead of stderr */ + CURLOPT(CURLOPT_STDERR, CURLOPTTYPE_OBJECTPOINT, 37), + + /* 38 is not used */ + + /* send linked-list of post-transfer QUOTE commands */ + CURLOPT(CURLOPT_POSTQUOTE, CURLOPTTYPE_SLISTPOINT, 39), + + /* 40 is not used */ + + /* talk a lot */ + CURLOPT(CURLOPT_VERBOSE, CURLOPTTYPE_LONG, 41), + + /* throw the header out too */ + CURLOPT(CURLOPT_HEADER, CURLOPTTYPE_LONG, 42), + + /* shut off the progress meter */ + CURLOPT(CURLOPT_NOPROGRESS, CURLOPTTYPE_LONG, 43), + + /* use HEAD to get http document */ + CURLOPT(CURLOPT_NOBODY, CURLOPTTYPE_LONG, 44), + + /* no output on http error codes >= 400 */ + CURLOPT(CURLOPT_FAILONERROR, CURLOPTTYPE_LONG, 45), + + /* this is an upload */ + CURLOPT(CURLOPT_UPLOAD, CURLOPTTYPE_LONG, 46), + + /* HTTP POST method */ + CURLOPT(CURLOPT_POST, CURLOPTTYPE_LONG, 47), + + /* bare names when listing directories */ + CURLOPT(CURLOPT_DIRLISTONLY, CURLOPTTYPE_LONG, 48), + + /* Append instead of overwrite on upload! */ + CURLOPT(CURLOPT_APPEND, CURLOPTTYPE_LONG, 50), + + /* Specify whether to read the user+password from the .netrc or the URL. + * This must be one of the CURL_NETRC_* enums below. */ + CURLOPT(CURLOPT_NETRC, CURLOPTTYPE_VALUES, 51), + + /* use Location: Luke! */ + CURLOPT(CURLOPT_FOLLOWLOCATION, CURLOPTTYPE_LONG, 52), + + /* transfer data in text/ASCII format */ + CURLOPT(CURLOPT_TRANSFERTEXT, CURLOPTTYPE_LONG, 53), + + /* HTTP PUT */ + CURLOPTDEPRECATED(CURLOPT_PUT, CURLOPTTYPE_LONG, 54, + 7.12.1, "Use CURLOPT_UPLOAD"), + + /* 55 = OBSOLETE */ + + /* DEPRECATED + * Function that will be called instead of the internal progress display + * function. This function should be defined as the curl_progress_callback + * prototype defines. */ + CURLOPTDEPRECATED(CURLOPT_PROGRESSFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 56, + 7.32.0, "Use CURLOPT_XFERINFOFUNCTION"), + + /* Data passed to the CURLOPT_PROGRESSFUNCTION and CURLOPT_XFERINFOFUNCTION + callbacks */ + CURLOPT(CURLOPT_XFERINFODATA, CURLOPTTYPE_CBPOINT, 57), +#define CURLOPT_PROGRESSDATA CURLOPT_XFERINFODATA + + /* We want the referrer field set automatically when following locations */ + CURLOPT(CURLOPT_AUTOREFERER, CURLOPTTYPE_LONG, 58), + + /* Port of the proxy, can be set in the proxy string as well with: + "[host]:[port]" */ + CURLOPT(CURLOPT_PROXYPORT, CURLOPTTYPE_LONG, 59), + + /* size of the POST input data, if strlen() is not good to use */ + CURLOPT(CURLOPT_POSTFIELDSIZE, CURLOPTTYPE_LONG, 60), + + /* tunnel non-http operations through an HTTP proxy */ + CURLOPT(CURLOPT_HTTPPROXYTUNNEL, CURLOPTTYPE_LONG, 61), + + /* Set the interface string to use as outgoing network interface */ + CURLOPT(CURLOPT_INTERFACE, CURLOPTTYPE_STRINGPOINT, 62), + + /* Set the krb4/5 security level, this also enables krb4/5 awareness. This + * is a string, 'clear', 'safe', 'confidential' or 'private'. If the string + * is set but does not match one of these, 'private' will be used. */ + CURLOPTDEPRECATED(CURLOPT_KRBLEVEL, CURLOPTTYPE_STRINGPOINT, 63, + 8.17.0, "removed"), + + /* Set if we should verify the peer in ssl handshake, set 1 to verify. */ + CURLOPT(CURLOPT_SSL_VERIFYPEER, CURLOPTTYPE_LONG, 64), + + /* The CApath or CAfile used to validate the peer certificate + this option is used only if SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_CAINFO, CURLOPTTYPE_STRINGPOINT, 65), + + /* 66 = OBSOLETE */ + /* 67 = OBSOLETE */ + + /* Maximum number of http redirects to follow */ + CURLOPT(CURLOPT_MAXREDIRS, CURLOPTTYPE_LONG, 68), + + /* Pass a long set to 1 to get the date of the requested document (if + possible)! Pass a zero to shut it off. */ + CURLOPT(CURLOPT_FILETIME, CURLOPTTYPE_LONG, 69), + + /* This points to a linked list of telnet options */ + CURLOPT(CURLOPT_TELNETOPTIONS, CURLOPTTYPE_SLISTPOINT, 70), + + /* Max amount of cached alive connections */ + CURLOPT(CURLOPT_MAXCONNECTS, CURLOPTTYPE_LONG, 71), + + /* 72 = OBSOLETE */ + /* 73 = OBSOLETE */ + + /* Set to explicitly use a new connection for the upcoming transfer. + Do not use this unless you are absolutely sure of this, as it makes the + operation slower and is less friendly for the network. */ + CURLOPT(CURLOPT_FRESH_CONNECT, CURLOPTTYPE_LONG, 74), + + /* Set to explicitly forbid the upcoming transfer's connection to be reused + when done. Do not use this unless you are absolutely sure of this, as it + makes the operation slower and is less friendly for the network. */ + CURLOPT(CURLOPT_FORBID_REUSE, CURLOPTTYPE_LONG, 75), + + /* Set to a filename that contains random data for libcurl to use to + seed the random engine when doing SSL connects. */ + CURLOPTDEPRECATED(CURLOPT_RANDOM_FILE, CURLOPTTYPE_STRINGPOINT, 76, + 7.84.0, "Serves no purpose anymore"), + + /* Set to the Entropy Gathering Daemon socket pathname */ + CURLOPTDEPRECATED(CURLOPT_EGDSOCKET, CURLOPTTYPE_STRINGPOINT, 77, + 7.84.0, "Serves no purpose anymore"), + + /* Time-out connect operations after this amount of seconds, if connects are + OK within this time, then fine... This only aborts the connect phase. */ + CURLOPT(CURLOPT_CONNECTTIMEOUT, CURLOPTTYPE_LONG, 78), + + /* Function that will be called to store headers (instead of fwrite). The + * parameters will use fwrite() syntax, make sure to follow them. */ + CURLOPT(CURLOPT_HEADERFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 79), + + /* Set this to force the HTTP request to get back to GET. Only really usable + if POST, PUT or a custom request have been used first. + */ + CURLOPT(CURLOPT_HTTPGET, CURLOPTTYPE_LONG, 80), + + /* Set if we should verify the Common name from the peer certificate in ssl + * handshake, set 1 to check existence, 2 to ensure that it matches the + * provided hostname. */ + CURLOPT(CURLOPT_SSL_VERIFYHOST, CURLOPTTYPE_LONG, 81), + + /* Specify which filename to write all known cookies in after completed + operation. Set filename to "-" (dash) to make it go to stdout. */ + CURLOPT(CURLOPT_COOKIEJAR, CURLOPTTYPE_STRINGPOINT, 82), + + /* Specify which TLS 1.2 (1.1, 1.0) ciphers to use */ + CURLOPT(CURLOPT_SSL_CIPHER_LIST, CURLOPTTYPE_STRINGPOINT, 83), + + /* Specify which HTTP version to use! This must be set to one of the + CURL_HTTP_VERSION* enums set below. */ + CURLOPT(CURLOPT_HTTP_VERSION, CURLOPTTYPE_VALUES, 84), + + /* Specifically switch on or off the FTP engine's use of the EPSV command. By + default, that one will always be attempted before the more traditional + PASV command. */ + CURLOPT(CURLOPT_FTP_USE_EPSV, CURLOPTTYPE_LONG, 85), + + /* type of the file keeping your SSL-certificate ("DER", "PEM", "ENG") */ + CURLOPT(CURLOPT_SSLCERTTYPE, CURLOPTTYPE_STRINGPOINT, 86), + + /* name of the file keeping your private SSL-key */ + CURLOPT(CURLOPT_SSLKEY, CURLOPTTYPE_STRINGPOINT, 87), + + /* type of the file keeping your private SSL-key ("DER", "PEM", "ENG") */ + CURLOPT(CURLOPT_SSLKEYTYPE, CURLOPTTYPE_STRINGPOINT, 88), + + /* crypto engine for the SSL-sub system */ + CURLOPT(CURLOPT_SSLENGINE, CURLOPTTYPE_STRINGPOINT, 89), + + /* set the crypto engine for the SSL-sub system as default + the param has no meaning... + */ + CURLOPT(CURLOPT_SSLENGINE_DEFAULT, CURLOPTTYPE_LONG, 90), + + /* Non-zero value means to use the global dns cache */ + /* DEPRECATED, do not use! */ + CURLOPTDEPRECATED(CURLOPT_DNS_USE_GLOBAL_CACHE, CURLOPTTYPE_LONG, 91, + 7.11.1, "Use CURLOPT_SHARE"), + + /* DNS cache timeout */ + CURLOPT(CURLOPT_DNS_CACHE_TIMEOUT, CURLOPTTYPE_LONG, 92), + + /* send linked-list of pre-transfer QUOTE commands */ + CURLOPT(CURLOPT_PREQUOTE, CURLOPTTYPE_SLISTPOINT, 93), + + /* set the debug function */ + CURLOPT(CURLOPT_DEBUGFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 94), + + /* set the data for the debug function */ + CURLOPT(CURLOPT_DEBUGDATA, CURLOPTTYPE_CBPOINT, 95), + + /* mark this as start of a cookie session */ + CURLOPT(CURLOPT_COOKIESESSION, CURLOPTTYPE_LONG, 96), + + /* The CApath directory used to validate the peer certificate + this option is used only if SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_CAPATH, CURLOPTTYPE_STRINGPOINT, 97), + + /* Instruct libcurl to use a smaller receive buffer */ + CURLOPT(CURLOPT_BUFFERSIZE, CURLOPTTYPE_LONG, 98), + + /* Instruct libcurl to not use any signal/alarm handlers, even when using + timeouts. This option is useful for multi-threaded applications. + See libcurl-the-guide for more background information. */ + CURLOPT(CURLOPT_NOSIGNAL, CURLOPTTYPE_LONG, 99), + + /* Provide a CURLShare for mutexing non-ts data */ + CURLOPT(CURLOPT_SHARE, CURLOPTTYPE_OBJECTPOINT, 100), + + /* indicates type of proxy. accepted values are CURLPROXY_HTTP (default), + CURLPROXY_HTTPS, CURLPROXY_SOCKS4, CURLPROXY_SOCKS4A and + CURLPROXY_SOCKS5. */ + CURLOPT(CURLOPT_PROXYTYPE, CURLOPTTYPE_VALUES, 101), + + /* Set the Accept-Encoding string. Use this to tell a server you would like + the response to be compressed. Before 7.21.6, this was known as + CURLOPT_ENCODING */ + CURLOPT(CURLOPT_ACCEPT_ENCODING, CURLOPTTYPE_STRINGPOINT, 102), + + /* Set pointer to private data */ + CURLOPT(CURLOPT_PRIVATE, CURLOPTTYPE_OBJECTPOINT, 103), + + /* Set aliases for HTTP 200 in the HTTP Response header */ + CURLOPT(CURLOPT_HTTP200ALIASES, CURLOPTTYPE_SLISTPOINT, 104), + + /* Continue to send authentication (user+password) when following locations, + even when hostname changed. This can potentially send off the name + and password to whatever host the server decides. */ + CURLOPT(CURLOPT_UNRESTRICTED_AUTH, CURLOPTTYPE_LONG, 105), + + /* Specifically switch on or off the FTP engine's use of the EPRT command ( + it also disables the LPRT attempt). By default, those ones will always be + attempted before the good old traditional PORT command. */ + CURLOPT(CURLOPT_FTP_USE_EPRT, CURLOPTTYPE_LONG, 106), + + /* Set this to a bitmask value to enable the particular authentications + methods you like. Use this in combination with CURLOPT_USERPWD. + Note that setting multiple bits may cause extra network round-trips. */ + CURLOPT(CURLOPT_HTTPAUTH, CURLOPTTYPE_VALUES, 107), + + /* Set the ssl context callback function, currently only for OpenSSL or + wolfSSL ssl_ctx, or mbedTLS mbedtls_ssl_config in the second argument. + The function must match the curl_ssl_ctx_callback prototype. */ + CURLOPT(CURLOPT_SSL_CTX_FUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 108), + + /* Set the userdata for the ssl context callback function's third + argument */ + CURLOPT(CURLOPT_SSL_CTX_DATA, CURLOPTTYPE_CBPOINT, 109), + + /* FTP Option that causes missing dirs to be created on the remote server. + In 7.19.4 we introduced the convenience enums for this option using the + CURLFTP_CREATE_DIR prefix. + */ + CURLOPT(CURLOPT_FTP_CREATE_MISSING_DIRS, CURLOPTTYPE_LONG, 110), + + /* Set this to a bitmask value to enable the particular authentications + methods you like. Use this in combination with CURLOPT_PROXYUSERPWD. + Note that setting multiple bits may cause extra network round-trips. */ + CURLOPT(CURLOPT_PROXYAUTH, CURLOPTTYPE_VALUES, 111), + + /* Option that changes the timeout, in seconds, associated with getting a + response. This is different from transfer timeout time and essentially + places a demand on the server to acknowledge commands in a timely + manner. For FTP, SMTP, IMAP and POP3. */ + CURLOPT(CURLOPT_SERVER_RESPONSE_TIMEOUT, CURLOPTTYPE_LONG, 112), + + /* Set this option to one of the CURL_IPRESOLVE_* defines (see below) to + tell libcurl to use those IP versions only. This only has effect on + systems with support for more than one, i.e IPv4 _and_ IPv6. */ + CURLOPT(CURLOPT_IPRESOLVE, CURLOPTTYPE_VALUES, 113), + + /* Set this option to limit the size of a file that will be downloaded from + an HTTP or FTP server. + + Note there is also _LARGE version which adds large file support for + platforms which have larger off_t sizes. See MAXFILESIZE_LARGE below. */ + CURLOPT(CURLOPT_MAXFILESIZE, CURLOPTTYPE_LONG, 114), + + /* See the comment for INFILESIZE above, but in short, specifies + * the size of the file being uploaded. -1 means unknown. + */ + CURLOPT(CURLOPT_INFILESIZE_LARGE, CURLOPTTYPE_OFF_T, 115), + + /* Sets the continuation offset. There is also a CURLOPTTYPE_LONG version + * of this; look above for RESUME_FROM. + */ + CURLOPT(CURLOPT_RESUME_FROM_LARGE, CURLOPTTYPE_OFF_T, 116), + + /* Sets the maximum size of data that will be downloaded from + * an HTTP or FTP server. See MAXFILESIZE above for the LONG version. + */ + CURLOPT(CURLOPT_MAXFILESIZE_LARGE, CURLOPTTYPE_OFF_T, 117), + + /* Set this option to the filename of your .netrc file you want libcurl + to parse (using the CURLOPT_NETRC option). If not set, libcurl will do + a poor attempt to find the user's home directory and check for a .netrc + file in there. */ + CURLOPT(CURLOPT_NETRC_FILE, CURLOPTTYPE_STRINGPOINT, 118), + + /* Enable SSL/TLS for FTP, pick one of: + CURLUSESSL_TRY - try using SSL, proceed anyway otherwise + CURLUSESSL_CONTROL - SSL for the control connection or fail + CURLUSESSL_ALL - SSL for all communication or fail + */ + CURLOPT(CURLOPT_USE_SSL, CURLOPTTYPE_VALUES, 119), + + /* The _LARGE version of the standard POSTFIELDSIZE option */ + CURLOPT(CURLOPT_POSTFIELDSIZE_LARGE, CURLOPTTYPE_OFF_T, 120), + + /* Enable/disable the TCP Nagle algorithm */ + CURLOPT(CURLOPT_TCP_NODELAY, CURLOPTTYPE_LONG, 121), + + /* 122 OBSOLETE, used in 7.12.3. Gone in 7.13.0 */ + /* 123 OBSOLETE. Gone in 7.16.0 */ + /* 124 OBSOLETE, used in 7.12.3. Gone in 7.13.0 */ + /* 125 OBSOLETE, used in 7.12.3. Gone in 7.13.0 */ + /* 126 OBSOLETE, used in 7.12.3. Gone in 7.13.0 */ + /* 127 OBSOLETE. Gone in 7.16.0 */ + /* 128 OBSOLETE. Gone in 7.16.0 */ + + /* When FTP over SSL/TLS is selected (with CURLOPT_USE_SSL), this option + can be used to change libcurl's default action which is to first try + "AUTH SSL" and then "AUTH TLS" in this order, and proceed when a OK + response has been received. + + Available parameters are: + CURLFTPAUTH_DEFAULT - let libcurl decide + CURLFTPAUTH_SSL - try "AUTH SSL" first, then TLS + CURLFTPAUTH_TLS - try "AUTH TLS" first, then SSL + */ + CURLOPT(CURLOPT_FTPSSLAUTH, CURLOPTTYPE_VALUES, 129), + + CURLOPTDEPRECATED(CURLOPT_IOCTLFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 130, + 7.18.0, "Use CURLOPT_SEEKFUNCTION"), + CURLOPTDEPRECATED(CURLOPT_IOCTLDATA, CURLOPTTYPE_CBPOINT, 131, + 7.18.0, "Use CURLOPT_SEEKDATA"), + + /* 132 OBSOLETE. Gone in 7.16.0 */ + /* 133 OBSOLETE. Gone in 7.16.0 */ + + /* null-terminated string for pass on to the FTP server when asked for + "account" info */ + CURLOPT(CURLOPT_FTP_ACCOUNT, CURLOPTTYPE_STRINGPOINT, 134), + + /* feed cookie into cookie engine */ + CURLOPT(CURLOPT_COOKIELIST, CURLOPTTYPE_STRINGPOINT, 135), + + /* ignore Content-Length */ + CURLOPT(CURLOPT_IGNORE_CONTENT_LENGTH, CURLOPTTYPE_LONG, 136), + + /* Set to non-zero to skip the IP address received in a 227 PASV FTP server + response. Typically used for FTP-SSL purposes but is not restricted to + that. libcurl will then instead use the same IP address it used for the + control connection. */ + CURLOPT(CURLOPT_FTP_SKIP_PASV_IP, CURLOPTTYPE_LONG, 137), + + /* Select "file method" to use when doing FTP, see the curl_ftpmethod + above. */ + CURLOPT(CURLOPT_FTP_FILEMETHOD, CURLOPTTYPE_VALUES, 138), + + /* Local port number to bind the socket to */ + CURLOPT(CURLOPT_LOCALPORT, CURLOPTTYPE_LONG, 139), + + /* Number of ports to try, including the first one set with LOCALPORT. + Thus, setting it to 1 will make no additional attempts but the first. + */ + CURLOPT(CURLOPT_LOCALPORTRANGE, CURLOPTTYPE_LONG, 140), + + /* no transfer, set up connection and let application use the socket by + extracting it with CURLINFO_LASTSOCKET */ + CURLOPT(CURLOPT_CONNECT_ONLY, CURLOPTTYPE_LONG, 141), + + /* Function that will be called to convert from the + network encoding (instead of using the iconv calls in libcurl) */ + CURLOPTDEPRECATED(CURLOPT_CONV_FROM_NETWORK_FUNCTION, + CURLOPTTYPE_FUNCTIONPOINT, 142, + 7.82.0, "Serves no purpose anymore"), + + /* Function that will be called to convert to the + network encoding (instead of using the iconv calls in libcurl) */ + CURLOPTDEPRECATED(CURLOPT_CONV_TO_NETWORK_FUNCTION, + CURLOPTTYPE_FUNCTIONPOINT, 143, + 7.82.0, "Serves no purpose anymore"), + + /* Function that will be called to convert from UTF8 + (instead of using the iconv calls in libcurl) + Note that this is used only for SSL certificate processing */ + CURLOPTDEPRECATED(CURLOPT_CONV_FROM_UTF8_FUNCTION, + CURLOPTTYPE_FUNCTIONPOINT, 144, + 7.82.0, "Serves no purpose anymore"), + + /* if the connection proceeds too quickly then need to slow it down */ + /* limit-rate: maximum number of bytes per second to send or receive */ + CURLOPT(CURLOPT_MAX_SEND_SPEED_LARGE, CURLOPTTYPE_OFF_T, 145), + CURLOPT(CURLOPT_MAX_RECV_SPEED_LARGE, CURLOPTTYPE_OFF_T, 146), + + /* Pointer to command string to send if USER/PASS fails. */ + CURLOPT(CURLOPT_FTP_ALTERNATIVE_TO_USER, CURLOPTTYPE_STRINGPOINT, 147), + + /* callback function for setting socket options */ + CURLOPT(CURLOPT_SOCKOPTFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 148), + CURLOPT(CURLOPT_SOCKOPTDATA, CURLOPTTYPE_CBPOINT, 149), + + /* set to 0 to disable session ID reuse for this transfer, default is + enabled (== 1) */ + CURLOPT(CURLOPT_SSL_SESSIONID_CACHE, CURLOPTTYPE_LONG, 150), + + /* allowed SSH authentication methods */ + CURLOPT(CURLOPT_SSH_AUTH_TYPES, CURLOPTTYPE_VALUES, 151), + + /* Used by scp/sftp to do public/private key authentication */ + CURLOPT(CURLOPT_SSH_PUBLIC_KEYFILE, CURLOPTTYPE_STRINGPOINT, 152), + CURLOPT(CURLOPT_SSH_PRIVATE_KEYFILE, CURLOPTTYPE_STRINGPOINT, 153), + + /* Send CCC (Clear Command Channel) after authentication */ + CURLOPT(CURLOPT_FTP_SSL_CCC, CURLOPTTYPE_LONG, 154), + + /* Same as TIMEOUT and CONNECTTIMEOUT, but with ms resolution */ + CURLOPT(CURLOPT_TIMEOUT_MS, CURLOPTTYPE_LONG, 155), + CURLOPT(CURLOPT_CONNECTTIMEOUT_MS, CURLOPTTYPE_LONG, 156), + + /* set to zero to disable the libcurl's decoding and thus pass the raw body + data to the application even when it is encoded/compressed */ + CURLOPT(CURLOPT_HTTP_TRANSFER_DECODING, CURLOPTTYPE_LONG, 157), + CURLOPT(CURLOPT_HTTP_CONTENT_DECODING, CURLOPTTYPE_LONG, 158), + + /* Permission used when creating new files and directories on the remote + server for protocols that support it, SFTP/SCP/FILE */ + CURLOPT(CURLOPT_NEW_FILE_PERMS, CURLOPTTYPE_LONG, 159), + CURLOPT(CURLOPT_NEW_DIRECTORY_PERMS, CURLOPTTYPE_LONG, 160), + + /* Set the behavior of POST when redirecting. Values must be set to one + of CURL_REDIR* defines below. This used to be called CURLOPT_POST301 */ + CURLOPT(CURLOPT_POSTREDIR, CURLOPTTYPE_VALUES, 161), + + /* used by scp/sftp to verify the host's public key */ + CURLOPT(CURLOPT_SSH_HOST_PUBLIC_KEY_MD5, CURLOPTTYPE_STRINGPOINT, 162), + + /* Callback function for opening socket (instead of socket(2)). Optionally, + callback is able change the address or refuse to connect returning + CURL_SOCKET_BAD. The callback should have type + curl_opensocket_callback */ + CURLOPT(CURLOPT_OPENSOCKETFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 163), + CURLOPT(CURLOPT_OPENSOCKETDATA, CURLOPTTYPE_CBPOINT, 164), + + /* POST volatile input fields. */ + CURLOPT(CURLOPT_COPYPOSTFIELDS, CURLOPTTYPE_OBJECTPOINT, 165), + + /* set transfer mode (;type=) when doing FTP via an HTTP proxy */ + CURLOPT(CURLOPT_PROXY_TRANSFER_MODE, CURLOPTTYPE_LONG, 166), + + /* Callback function for seeking in the input stream */ + CURLOPT(CURLOPT_SEEKFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 167), + CURLOPT(CURLOPT_SEEKDATA, CURLOPTTYPE_CBPOINT, 168), + + /* CRL file */ + CURLOPT(CURLOPT_CRLFILE, CURLOPTTYPE_STRINGPOINT, 169), + + /* Issuer certificate */ + CURLOPT(CURLOPT_ISSUERCERT, CURLOPTTYPE_STRINGPOINT, 170), + + /* (IPv6) Address scope */ + CURLOPT(CURLOPT_ADDRESS_SCOPE, CURLOPTTYPE_LONG, 171), + + /* Collect certificate chain info and allow it to get retrievable with + CURLINFO_CERTINFO after the transfer is complete. */ + CURLOPT(CURLOPT_CERTINFO, CURLOPTTYPE_LONG, 172), + + /* "name" and "pwd" to use when fetching. */ + CURLOPT(CURLOPT_USERNAME, CURLOPTTYPE_STRINGPOINT, 173), + CURLOPT(CURLOPT_PASSWORD, CURLOPTTYPE_STRINGPOINT, 174), + + /* "name" and "pwd" to use with Proxy when fetching. */ + CURLOPT(CURLOPT_PROXYUSERNAME, CURLOPTTYPE_STRINGPOINT, 175), + CURLOPT(CURLOPT_PROXYPASSWORD, CURLOPTTYPE_STRINGPOINT, 176), + + /* Comma separated list of hostnames defining no-proxy zones. These should + match both hostnames directly, and hostnames within a domain. For + example, local.com will match local.com and www.local.com, but NOT + notlocal.com or www.notlocal.com. For compatibility with other + implementations of this, .local.com will be considered to be the same as + local.com. A single * is the only valid wildcard, and effectively + disables the use of proxy. */ + CURLOPT(CURLOPT_NOPROXY, CURLOPTTYPE_STRINGPOINT, 177), + + /* block size for TFTP transfers */ + CURLOPT(CURLOPT_TFTP_BLKSIZE, CURLOPTTYPE_LONG, 178), + + /* Socks Service */ + /* DEPRECATED, do not use! */ + CURLOPTDEPRECATED(CURLOPT_SOCKS5_GSSAPI_SERVICE, + CURLOPTTYPE_STRINGPOINT, 179, + 7.49.0, "Use CURLOPT_PROXY_SERVICE_NAME"), + + /* Socks Service */ + CURLOPT(CURLOPT_SOCKS5_GSSAPI_NEC, CURLOPTTYPE_LONG, 180), + + /* set the bitmask for the protocols that are allowed to be used for the + transfer, which thus helps the app which takes URLs from users or other + external inputs and want to restrict what protocol(s) to deal + with. Defaults to CURLPROTO_ALL. */ + CURLOPTDEPRECATED(CURLOPT_PROTOCOLS, CURLOPTTYPE_LONG, 181, + 7.85.0, "Use CURLOPT_PROTOCOLS_STR"), + + /* set the bitmask for the protocols that libcurl is allowed to follow to, + as a subset of the CURLOPT_PROTOCOLS ones. That means the protocol needs + to be set in both bitmasks to be allowed to get redirected to. */ + CURLOPTDEPRECATED(CURLOPT_REDIR_PROTOCOLS, CURLOPTTYPE_LONG, 182, + 7.85.0, "Use CURLOPT_REDIR_PROTOCOLS_STR"), + + /* set the SSH knownhost filename to use */ + CURLOPT(CURLOPT_SSH_KNOWNHOSTS, CURLOPTTYPE_STRINGPOINT, 183), + + /* set the SSH host key callback, must point to a curl_sshkeycallback + function */ + CURLOPT(CURLOPT_SSH_KEYFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 184), + + /* set the SSH host key callback custom pointer */ + CURLOPT(CURLOPT_SSH_KEYDATA, CURLOPTTYPE_CBPOINT, 185), + + /* set the SMTP mail originator */ + CURLOPT(CURLOPT_MAIL_FROM, CURLOPTTYPE_STRINGPOINT, 186), + + /* set the list of SMTP mail receiver(s) */ + CURLOPT(CURLOPT_MAIL_RCPT, CURLOPTTYPE_SLISTPOINT, 187), + + /* FTP: send PRET before PASV */ + CURLOPT(CURLOPT_FTP_USE_PRET, CURLOPTTYPE_LONG, 188), + + /* RTSP request method (OPTIONS, SETUP, PLAY, etc...) */ + CURLOPT(CURLOPT_RTSP_REQUEST, CURLOPTTYPE_VALUES, 189), + + /* The RTSP session identifier */ + CURLOPT(CURLOPT_RTSP_SESSION_ID, CURLOPTTYPE_STRINGPOINT, 190), + + /* The RTSP stream URI */ + CURLOPT(CURLOPT_RTSP_STREAM_URI, CURLOPTTYPE_STRINGPOINT, 191), + + /* The Transport: header to use in RTSP requests */ + CURLOPT(CURLOPT_RTSP_TRANSPORT, CURLOPTTYPE_STRINGPOINT, 192), + + /* Manually initialize the client RTSP CSeq for this handle */ + CURLOPT(CURLOPT_RTSP_CLIENT_CSEQ, CURLOPTTYPE_LONG, 193), + + /* Manually initialize the server RTSP CSeq for this handle */ + CURLOPT(CURLOPT_RTSP_SERVER_CSEQ, CURLOPTTYPE_LONG, 194), + + /* The stream to pass to INTERLEAVEFUNCTION. */ + CURLOPT(CURLOPT_INTERLEAVEDATA, CURLOPTTYPE_CBPOINT, 195), + + /* Let the application define a custom write method for RTP data */ + CURLOPT(CURLOPT_INTERLEAVEFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 196), + + /* Turn on wildcard matching */ + CURLOPT(CURLOPT_WILDCARDMATCH, CURLOPTTYPE_LONG, 197), + + /* Directory matching callback called before downloading of an + individual file (chunk) started */ + CURLOPT(CURLOPT_CHUNK_BGN_FUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 198), + + /* Directory matching callback called after the file (chunk) + was downloaded, or skipped */ + CURLOPT(CURLOPT_CHUNK_END_FUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 199), + + /* Change match (fnmatch-like) callback for wildcard matching */ + CURLOPT(CURLOPT_FNMATCH_FUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 200), + + /* Let the application define custom chunk data pointer */ + CURLOPT(CURLOPT_CHUNK_DATA, CURLOPTTYPE_CBPOINT, 201), + + /* FNMATCH_FUNCTION user pointer */ + CURLOPT(CURLOPT_FNMATCH_DATA, CURLOPTTYPE_CBPOINT, 202), + + /* send linked-list of name:port:address sets */ + CURLOPT(CURLOPT_RESOLVE, CURLOPTTYPE_SLISTPOINT, 203), + + /* Set a username for authenticated TLS */ + CURLOPT(CURLOPT_TLSAUTH_USERNAME, CURLOPTTYPE_STRINGPOINT, 204), + + /* Set a password for authenticated TLS */ + CURLOPT(CURLOPT_TLSAUTH_PASSWORD, CURLOPTTYPE_STRINGPOINT, 205), + + /* Set authentication type for authenticated TLS */ + CURLOPT(CURLOPT_TLSAUTH_TYPE, CURLOPTTYPE_STRINGPOINT, 206), + + /* Set to 1 to enable the "TE:" header in HTTP requests to ask for + compressed transfer-encoded responses. Set to 0 to disable the use of TE: + in outgoing requests. The current default is 0, but it might change in a + future libcurl release. + + libcurl will ask for the compressed methods it knows of, and if that + is not any, it will not ask for transfer-encoding at all even if this + option is set to 1. */ + CURLOPT(CURLOPT_TRANSFER_ENCODING, CURLOPTTYPE_LONG, 207), + + /* Callback function for closing socket (instead of close(2)). The callback + should have type curl_closesocket_callback */ + CURLOPT(CURLOPT_CLOSESOCKETFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 208), + CURLOPT(CURLOPT_CLOSESOCKETDATA, CURLOPTTYPE_CBPOINT, 209), + + /* allow GSSAPI credential delegation */ + CURLOPT(CURLOPT_GSSAPI_DELEGATION, CURLOPTTYPE_VALUES, 210), + + /* Set the name servers to use for DNS resolution. + * Only supported by the c-ares DNS backend */ + CURLOPT(CURLOPT_DNS_SERVERS, CURLOPTTYPE_STRINGPOINT, 211), + + /* Time-out accept operations (currently for FTP only) after this amount + of milliseconds. */ + CURLOPT(CURLOPT_ACCEPTTIMEOUT_MS, CURLOPTTYPE_LONG, 212), + + /* Set TCP keepalive */ + CURLOPT(CURLOPT_TCP_KEEPALIVE, CURLOPTTYPE_LONG, 213), + + /* non-universal keepalive knobs (Linux, AIX, HP-UX, more) */ + CURLOPT(CURLOPT_TCP_KEEPIDLE, CURLOPTTYPE_LONG, 214), + CURLOPT(CURLOPT_TCP_KEEPINTVL, CURLOPTTYPE_LONG, 215), + + /* Enable/disable specific SSL features with a bitmask, see CURLSSLOPT_* */ + CURLOPT(CURLOPT_SSL_OPTIONS, CURLOPTTYPE_VALUES, 216), + + /* Set the SMTP auth originator */ + CURLOPT(CURLOPT_MAIL_AUTH, CURLOPTTYPE_STRINGPOINT, 217), + + /* Enable/disable SASL initial response */ + CURLOPT(CURLOPT_SASL_IR, CURLOPTTYPE_LONG, 218), + + /* Function that will be called instead of the internal progress display + * function. This function should be defined as the curl_xferinfo_callback + * prototype defines. (Deprecates CURLOPT_PROGRESSFUNCTION) */ + CURLOPT(CURLOPT_XFERINFOFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 219), + + /* The XOAUTH2 bearer token */ + CURLOPT(CURLOPT_XOAUTH2_BEARER, CURLOPTTYPE_STRINGPOINT, 220), + + /* Set the interface string to use as outgoing network + * interface for DNS requests. + * Only supported by the c-ares DNS backend */ + CURLOPT(CURLOPT_DNS_INTERFACE, CURLOPTTYPE_STRINGPOINT, 221), + + /* Set the local IPv4 address to use for outgoing DNS requests. + * Only supported by the c-ares DNS backend */ + CURLOPT(CURLOPT_DNS_LOCAL_IP4, CURLOPTTYPE_STRINGPOINT, 222), + + /* Set the local IPv6 address to use for outgoing DNS requests. + * Only supported by the c-ares DNS backend */ + CURLOPT(CURLOPT_DNS_LOCAL_IP6, CURLOPTTYPE_STRINGPOINT, 223), + + /* Set authentication options directly */ + CURLOPT(CURLOPT_LOGIN_OPTIONS, CURLOPTTYPE_STRINGPOINT, 224), + + /* Enable/disable TLS NPN extension (http2 over ssl might fail without) */ + CURLOPTDEPRECATED(CURLOPT_SSL_ENABLE_NPN, CURLOPTTYPE_LONG, 225, + 7.86.0, "Has no function"), + + /* Enable/disable TLS ALPN extension (http2 over ssl might fail without) */ + CURLOPT(CURLOPT_SSL_ENABLE_ALPN, CURLOPTTYPE_LONG, 226), + + /* Time to wait for a response to an HTTP request containing an + * Expect: 100-continue header before sending the data anyway. */ + CURLOPT(CURLOPT_EXPECT_100_TIMEOUT_MS, CURLOPTTYPE_LONG, 227), + + /* This points to a linked list of headers used for proxy requests only, + struct curl_slist kind */ + CURLOPT(CURLOPT_PROXYHEADER, CURLOPTTYPE_SLISTPOINT, 228), + + /* Pass in a bitmask of "header options" */ + CURLOPT(CURLOPT_HEADEROPT, CURLOPTTYPE_VALUES, 229), + + /* The public key used to validate the peer public key */ + CURLOPT(CURLOPT_PINNEDPUBLICKEY, CURLOPTTYPE_STRINGPOINT, 230), + + /* Path to Unix domain socket */ + CURLOPT(CURLOPT_UNIX_SOCKET_PATH, CURLOPTTYPE_STRINGPOINT, 231), + + /* Set if we should verify the certificate status. */ + CURLOPT(CURLOPT_SSL_VERIFYSTATUS, CURLOPTTYPE_LONG, 232), + + /* Set if we should enable TLS false start. */ + CURLOPTDEPRECATED(CURLOPT_SSL_FALSESTART, CURLOPTTYPE_LONG, 233, + 8.15.0, "Has no function"), + + /* Do not squash dot-dot sequences */ + CURLOPT(CURLOPT_PATH_AS_IS, CURLOPTTYPE_LONG, 234), + + /* Proxy Service Name */ + CURLOPT(CURLOPT_PROXY_SERVICE_NAME, CURLOPTTYPE_STRINGPOINT, 235), + + /* Service Name */ + CURLOPT(CURLOPT_SERVICE_NAME, CURLOPTTYPE_STRINGPOINT, 236), + + /* Wait/do not wait for pipe/mutex to clarify */ + CURLOPT(CURLOPT_PIPEWAIT, CURLOPTTYPE_LONG, 237), + + /* Set the protocol used when curl is given a URL without a protocol */ + CURLOPT(CURLOPT_DEFAULT_PROTOCOL, CURLOPTTYPE_STRINGPOINT, 238), + + /* Set stream weight, 1 - 256 (default is 16) */ + CURLOPT(CURLOPT_STREAM_WEIGHT, CURLOPTTYPE_LONG, 239), + + /* Set stream dependency on another curl handle */ + CURLOPT(CURLOPT_STREAM_DEPENDS, CURLOPTTYPE_OBJECTPOINT, 240), + + /* Set E-xclusive stream dependency on another curl handle */ + CURLOPT(CURLOPT_STREAM_DEPENDS_E, CURLOPTTYPE_OBJECTPOINT, 241), + + /* Do not send any tftp option requests to the server */ + CURLOPT(CURLOPT_TFTP_NO_OPTIONS, CURLOPTTYPE_LONG, 242), + + /* Linked-list of host:port:connect-to-host:connect-to-port, + overrides the URL's host:port (only for the network layer) */ + CURLOPT(CURLOPT_CONNECT_TO, CURLOPTTYPE_SLISTPOINT, 243), + + /* Set TCP Fast Open */ + CURLOPT(CURLOPT_TCP_FASTOPEN, CURLOPTTYPE_LONG, 244), + + /* Continue to send data if the server responds early with an + * HTTP status code >= 300 */ + CURLOPT(CURLOPT_KEEP_SENDING_ON_ERROR, CURLOPTTYPE_LONG, 245), + + /* The CApath or CAfile used to validate the proxy certificate + this option is used only if PROXY_SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_PROXY_CAINFO, CURLOPTTYPE_STRINGPOINT, 246), + + /* The CApath directory used to validate the proxy certificate + this option is used only if PROXY_SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_PROXY_CAPATH, CURLOPTTYPE_STRINGPOINT, 247), + + /* Set if we should verify the proxy in ssl handshake, + set 1 to verify. */ + CURLOPT(CURLOPT_PROXY_SSL_VERIFYPEER, CURLOPTTYPE_LONG, 248), + + /* Set if we should verify the Common name from the proxy certificate in ssl + * handshake, set 1 to check existence, 2 to ensure that it matches + * the provided hostname. */ + CURLOPT(CURLOPT_PROXY_SSL_VERIFYHOST, CURLOPTTYPE_LONG, 249), + + /* What version to specifically try to use for proxy. + See CURL_SSLVERSION defines below. */ + CURLOPT(CURLOPT_PROXY_SSLVERSION, CURLOPTTYPE_VALUES, 250), + + /* Set a username for authenticated TLS for proxy */ + CURLOPT(CURLOPT_PROXY_TLSAUTH_USERNAME, CURLOPTTYPE_STRINGPOINT, 251), + + /* Set a password for authenticated TLS for proxy */ + CURLOPT(CURLOPT_PROXY_TLSAUTH_PASSWORD, CURLOPTTYPE_STRINGPOINT, 252), + + /* Set authentication type for authenticated TLS for proxy */ + CURLOPT(CURLOPT_PROXY_TLSAUTH_TYPE, CURLOPTTYPE_STRINGPOINT, 253), + + /* name of the file keeping your private SSL-certificate for proxy */ + CURLOPT(CURLOPT_PROXY_SSLCERT, CURLOPTTYPE_STRINGPOINT, 254), + + /* type of the file keeping your SSL-certificate ("DER", "PEM", "ENG") for + proxy */ + CURLOPT(CURLOPT_PROXY_SSLCERTTYPE, CURLOPTTYPE_STRINGPOINT, 255), + + /* name of the file keeping your private SSL-key for proxy */ + CURLOPT(CURLOPT_PROXY_SSLKEY, CURLOPTTYPE_STRINGPOINT, 256), + + /* type of the file keeping your private SSL-key ("DER", "PEM", "ENG") for + proxy */ + CURLOPT(CURLOPT_PROXY_SSLKEYTYPE, CURLOPTTYPE_STRINGPOINT, 257), + + /* password for the SSL private key for proxy */ + CURLOPT(CURLOPT_PROXY_KEYPASSWD, CURLOPTTYPE_STRINGPOINT, 258), + + /* Specify which TLS 1.2 (1.1, 1.0) ciphers to use for proxy */ + CURLOPT(CURLOPT_PROXY_SSL_CIPHER_LIST, CURLOPTTYPE_STRINGPOINT, 259), + + /* CRL file for proxy */ + CURLOPT(CURLOPT_PROXY_CRLFILE, CURLOPTTYPE_STRINGPOINT, 260), + + /* Enable/disable specific SSL features with a bitmask for proxy, see + CURLSSLOPT_* */ + CURLOPT(CURLOPT_PROXY_SSL_OPTIONS, CURLOPTTYPE_LONG, 261), + + /* Name of pre proxy to use. */ + CURLOPT(CURLOPT_PRE_PROXY, CURLOPTTYPE_STRINGPOINT, 262), + + /* The public key in DER form used to validate the proxy public key + this option is used only if PROXY_SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_PROXY_PINNEDPUBLICKEY, CURLOPTTYPE_STRINGPOINT, 263), + + /* Path to an abstract Unix domain socket */ + CURLOPT(CURLOPT_ABSTRACT_UNIX_SOCKET, CURLOPTTYPE_STRINGPOINT, 264), + + /* Suppress proxy CONNECT response headers from user callbacks */ + CURLOPT(CURLOPT_SUPPRESS_CONNECT_HEADERS, CURLOPTTYPE_LONG, 265), + + /* The request target, instead of extracted from the URL */ + CURLOPT(CURLOPT_REQUEST_TARGET, CURLOPTTYPE_STRINGPOINT, 266), + + /* bitmask of allowed auth methods for connections to SOCKS5 proxies */ + CURLOPT(CURLOPT_SOCKS5_AUTH, CURLOPTTYPE_LONG, 267), + + /* Enable/disable SSH compression */ + CURLOPT(CURLOPT_SSH_COMPRESSION, CURLOPTTYPE_LONG, 268), + + /* Post MIME data. */ + CURLOPT(CURLOPT_MIMEPOST, CURLOPTTYPE_OBJECTPOINT, 269), + + /* Time to use with the CURLOPT_TIMECONDITION. Specified in number of + seconds since 1 Jan 1970. */ + CURLOPT(CURLOPT_TIMEVALUE_LARGE, CURLOPTTYPE_OFF_T, 270), + + /* Head start in milliseconds to give happy eyeballs. */ + CURLOPT(CURLOPT_HAPPY_EYEBALLS_TIMEOUT_MS, CURLOPTTYPE_LONG, 271), + + /* Function that will be called before a resolver request is made */ + CURLOPT(CURLOPT_RESOLVER_START_FUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 272), + + /* User data to pass to the resolver start callback. */ + CURLOPT(CURLOPT_RESOLVER_START_DATA, CURLOPTTYPE_CBPOINT, 273), + + /* send HAProxy PROXY protocol header? */ + CURLOPT(CURLOPT_HAPROXYPROTOCOL, CURLOPTTYPE_LONG, 274), + + /* shuffle addresses before use when DNS returns multiple */ + CURLOPT(CURLOPT_DNS_SHUFFLE_ADDRESSES, CURLOPTTYPE_LONG, 275), + + /* Specify which TLS 1.3 ciphers suites to use */ + CURLOPT(CURLOPT_TLS13_CIPHERS, CURLOPTTYPE_STRINGPOINT, 276), + CURLOPT(CURLOPT_PROXY_TLS13_CIPHERS, CURLOPTTYPE_STRINGPOINT, 277), + + /* Disallow specifying username/login in URL. */ + CURLOPT(CURLOPT_DISALLOW_USERNAME_IN_URL, CURLOPTTYPE_LONG, 278), + + /* DNS-over-HTTPS URL */ + CURLOPT(CURLOPT_DOH_URL, CURLOPTTYPE_STRINGPOINT, 279), + + /* Preferred buffer size to use for uploads */ + CURLOPT(CURLOPT_UPLOAD_BUFFERSIZE, CURLOPTTYPE_LONG, 280), + + /* Time in ms between connection upkeep calls for long-lived connections. */ + CURLOPT(CURLOPT_UPKEEP_INTERVAL_MS, CURLOPTTYPE_LONG, 281), + + /* Specify URL using CURL URL API. */ + CURLOPT(CURLOPT_CURLU, CURLOPTTYPE_OBJECTPOINT, 282), + + /* add trailing data just after no more data is available */ + CURLOPT(CURLOPT_TRAILERFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 283), + + /* pointer to be passed to HTTP_TRAILER_FUNCTION */ + CURLOPT(CURLOPT_TRAILERDATA, CURLOPTTYPE_CBPOINT, 284), + + /* set this to 1L to allow HTTP/0.9 responses or 0L to disallow */ + CURLOPT(CURLOPT_HTTP09_ALLOWED, CURLOPTTYPE_LONG, 285), + + /* alt-svc control bitmask */ + CURLOPT(CURLOPT_ALTSVC_CTRL, CURLOPTTYPE_LONG, 286), + + /* alt-svc cache filename to possibly read from/write to */ + CURLOPT(CURLOPT_ALTSVC, CURLOPTTYPE_STRINGPOINT, 287), + + /* maximum age (idle time) of a connection to consider it for reuse + * (in seconds) */ + CURLOPT(CURLOPT_MAXAGE_CONN, CURLOPTTYPE_LONG, 288), + + /* SASL authorization identity */ + CURLOPT(CURLOPT_SASL_AUTHZID, CURLOPTTYPE_STRINGPOINT, 289), + + /* allow RCPT TO command to fail for some recipients */ + CURLOPT(CURLOPT_MAIL_RCPT_ALLOWFAILS, CURLOPTTYPE_LONG, 290), + + /* the private SSL-certificate as a "blob" */ + CURLOPT(CURLOPT_SSLCERT_BLOB, CURLOPTTYPE_BLOB, 291), + CURLOPT(CURLOPT_SSLKEY_BLOB, CURLOPTTYPE_BLOB, 292), + CURLOPT(CURLOPT_PROXY_SSLCERT_BLOB, CURLOPTTYPE_BLOB, 293), + CURLOPT(CURLOPT_PROXY_SSLKEY_BLOB, CURLOPTTYPE_BLOB, 294), + CURLOPT(CURLOPT_ISSUERCERT_BLOB, CURLOPTTYPE_BLOB, 295), + + /* Issuer certificate for proxy */ + CURLOPT(CURLOPT_PROXY_ISSUERCERT, CURLOPTTYPE_STRINGPOINT, 296), + CURLOPT(CURLOPT_PROXY_ISSUERCERT_BLOB, CURLOPTTYPE_BLOB, 297), + + /* the EC curves requested by the TLS client (RFC 8422, 5.1); + * OpenSSL support via 'set_groups'/'set_curves': + * https://docs.openssl.org/master/man3/SSL_CTX_set1_curves/ + */ + CURLOPT(CURLOPT_SSL_EC_CURVES, CURLOPTTYPE_STRINGPOINT, 298), + + /* HSTS bitmask */ + CURLOPT(CURLOPT_HSTS_CTRL, CURLOPTTYPE_LONG, 299), + /* HSTS filename */ + CURLOPT(CURLOPT_HSTS, CURLOPTTYPE_STRINGPOINT, 300), + + /* HSTS read callback */ + CURLOPT(CURLOPT_HSTSREADFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 301), + CURLOPT(CURLOPT_HSTSREADDATA, CURLOPTTYPE_CBPOINT, 302), + + /* HSTS write callback */ + CURLOPT(CURLOPT_HSTSWRITEFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 303), + CURLOPT(CURLOPT_HSTSWRITEDATA, CURLOPTTYPE_CBPOINT, 304), + + /* Parameters for V4 signature */ + CURLOPT(CURLOPT_AWS_SIGV4, CURLOPTTYPE_STRINGPOINT, 305), + + /* Same as CURLOPT_SSL_VERIFYPEER but for DoH (DNS-over-HTTPS) servers. */ + CURLOPT(CURLOPT_DOH_SSL_VERIFYPEER, CURLOPTTYPE_LONG, 306), + + /* Same as CURLOPT_SSL_VERIFYHOST but for DoH (DNS-over-HTTPS) servers. */ + CURLOPT(CURLOPT_DOH_SSL_VERIFYHOST, CURLOPTTYPE_LONG, 307), + + /* Same as CURLOPT_SSL_VERIFYSTATUS but for DoH (DNS-over-HTTPS) servers. */ + CURLOPT(CURLOPT_DOH_SSL_VERIFYSTATUS, CURLOPTTYPE_LONG, 308), + + /* The CA certificates as "blob" used to validate the peer certificate + this option is used only if SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_CAINFO_BLOB, CURLOPTTYPE_BLOB, 309), + + /* The CA certificates as "blob" used to validate the proxy certificate + this option is used only if PROXY_SSL_VERIFYPEER is true */ + CURLOPT(CURLOPT_PROXY_CAINFO_BLOB, CURLOPTTYPE_BLOB, 310), + + /* used by scp/sftp to verify the host's public key */ + CURLOPT(CURLOPT_SSH_HOST_PUBLIC_KEY_SHA256, CURLOPTTYPE_STRINGPOINT, 311), + + /* Function that will be called immediately before the initial request + is made on a connection (after any protocol negotiation step). */ + CURLOPT(CURLOPT_PREREQFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 312), + + /* Data passed to the CURLOPT_PREREQFUNCTION callback */ + CURLOPT(CURLOPT_PREREQDATA, CURLOPTTYPE_CBPOINT, 313), + + /* maximum age (since creation) of a connection to consider it for reuse + * (in seconds) */ + CURLOPT(CURLOPT_MAXLIFETIME_CONN, CURLOPTTYPE_LONG, 314), + + /* Set MIME option flags. */ + CURLOPT(CURLOPT_MIME_OPTIONS, CURLOPTTYPE_LONG, 315), + + /* set the SSH host key callback, must point to a curl_sshkeycallback + function */ + CURLOPT(CURLOPT_SSH_HOSTKEYFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 316), + + /* set the SSH host key callback custom pointer */ + CURLOPT(CURLOPT_SSH_HOSTKEYDATA, CURLOPTTYPE_CBPOINT, 317), + + /* specify which protocols that are allowed to be used for the transfer, + which thus helps the app which takes URLs from users or other external + inputs and want to restrict what protocol(s) to deal with. Defaults to + all built-in protocols. */ + CURLOPT(CURLOPT_PROTOCOLS_STR, CURLOPTTYPE_STRINGPOINT, 318), + + /* specify which protocols that libcurl is allowed to follow directs to */ + CURLOPT(CURLOPT_REDIR_PROTOCOLS_STR, CURLOPTTYPE_STRINGPOINT, 319), + + /* WebSockets options */ + CURLOPT(CURLOPT_WS_OPTIONS, CURLOPTTYPE_LONG, 320), + + /* CA cache timeout */ + CURLOPT(CURLOPT_CA_CACHE_TIMEOUT, CURLOPTTYPE_LONG, 321), + + /* Can leak things, gonna exit() soon */ + CURLOPT(CURLOPT_QUICK_EXIT, CURLOPTTYPE_LONG, 322), + + /* set a specific client IP for HAProxy PROXY protocol header? */ + CURLOPT(CURLOPT_HAPROXY_CLIENT_IP, CURLOPTTYPE_STRINGPOINT, 323), + + /* millisecond version */ + CURLOPT(CURLOPT_SERVER_RESPONSE_TIMEOUT_MS, CURLOPTTYPE_LONG, 324), + + /* set ECH configuration */ + CURLOPT(CURLOPT_ECH, CURLOPTTYPE_STRINGPOINT, 325), + + /* maximum number of keepalive probes (Linux, *BSD, macOS, etc.) */ + CURLOPT(CURLOPT_TCP_KEEPCNT, CURLOPTTYPE_LONG, 326), + + CURLOPT(CURLOPT_UPLOAD_FLAGS, CURLOPTTYPE_LONG, 327), + + /* set TLS supported signature algorithms */ + CURLOPT(CURLOPT_SSL_SIGNATURE_ALGORITHMS, CURLOPTTYPE_STRINGPOINT, 328), + + CURLOPT_LASTENTRY /* the last unused */ +} CURLoption; + +#ifndef CURL_NO_OLDIES /* define this to test if your app builds with all + the obsolete stuff removed! */ + +/* Backwards compatibility with older names */ +/* These are scheduled to disappear by 2011 */ + +/* This was added in version 7.19.1 */ +#define CURLOPT_POST301 CURLOPT_POSTREDIR + +/* These are scheduled to disappear by 2009 */ + +/* The following were added in 7.17.0 */ +#define CURLOPT_SSLKEYPASSWD CURLOPT_KEYPASSWD +#define CURLOPT_FTPAPPEND CURLOPT_APPEND +#define CURLOPT_FTPLISTONLY CURLOPT_DIRLISTONLY +#define CURLOPT_FTP_SSL CURLOPT_USE_SSL + +/* The following were added earlier */ + +#define CURLOPT_SSLCERTPASSWD CURLOPT_KEYPASSWD +#define CURLOPT_KRB4LEVEL CURLOPT_KRBLEVEL + +/* */ +#define CURLOPT_FTP_RESPONSE_TIMEOUT CURLOPT_SERVER_RESPONSE_TIMEOUT + +/* Added in 8.2.0 */ +#define CURLOPT_MAIL_RCPT_ALLLOWFAILS CURLOPT_MAIL_RCPT_ALLOWFAILS + +#else +/* This is set if CURL_NO_OLDIES is defined at compile-time */ +#undef CURLOPT_DNS_USE_GLOBAL_CACHE /* soon obsolete */ +#endif + +/* Below here follows defines for the CURLOPT_IPRESOLVE option. If a host + name resolves addresses using more than one IP protocol version, this + option might be handy to force libcurl to use a specific IP version. */ +#define CURL_IPRESOLVE_WHATEVER 0L /* default, uses addresses to all IP + versions that your system allows */ +#define CURL_IPRESOLVE_V4 1L /* uses only IPv4 addresses/connections */ +#define CURL_IPRESOLVE_V6 2L /* uses only IPv6 addresses/connections */ + + /* Convenient "aliases" */ +#define CURLOPT_RTSPHEADER CURLOPT_HTTPHEADER + +/* These constants are for use with the CURLOPT_HTTP_VERSION option. */ +#define CURL_HTTP_VERSION_NONE 0L /* setting this means we do not care, and + that we would like the library to choose + the best possible for us! */ +#define CURL_HTTP_VERSION_1_0 1L /* please use HTTP 1.0 in the request */ +#define CURL_HTTP_VERSION_1_1 2L /* please use HTTP 1.1 in the request */ +#define CURL_HTTP_VERSION_2_0 3L /* please use HTTP 2 in the request */ +#define CURL_HTTP_VERSION_2TLS 4L /* use version 2 for HTTPS, version 1.1 for + HTTP */ +#define CURL_HTTP_VERSION_2_PRIOR_KNOWLEDGE 5L /* please use HTTP 2 without + HTTP/1.1 Upgrade */ +#define CURL_HTTP_VERSION_3 30L /* Use HTTP/3, fallback to HTTP/2 or + HTTP/1 if needed. For HTTPS only. For + HTTP, this option makes libcurl + return error. */ +#define CURL_HTTP_VERSION_3ONLY 31L /* Use HTTP/3 without fallback. For + HTTPS only. For HTTP, this makes + libcurl return error. */ +#define CURL_HTTP_VERSION_LAST 32L /* *ILLEGAL* http version */ + +/* Convenience definition simple because the name of the version is HTTP/2 and + not 2.0. The 2_0 version of the enum name was set while the version was + still planned to be 2.0 and we stick to it for compatibility. */ +#define CURL_HTTP_VERSION_2 CURL_HTTP_VERSION_2_0 + +/* + * Public API enums for RTSP requests + */ + +#define CURL_RTSPREQ_NONE 0L +#define CURL_RTSPREQ_OPTIONS 1L +#define CURL_RTSPREQ_DESCRIBE 2L +#define CURL_RTSPREQ_ANNOUNCE 3L +#define CURL_RTSPREQ_SETUP 4L +#define CURL_RTSPREQ_PLAY 5L +#define CURL_RTSPREQ_PAUSE 6L +#define CURL_RTSPREQ_TEARDOWN 7L +#define CURL_RTSPREQ_GET_PARAMETER 8L +#define CURL_RTSPREQ_SET_PARAMETER 9L +#define CURL_RTSPREQ_RECORD 10L +#define CURL_RTSPREQ_RECEIVE 11L +#define CURL_RTSPREQ_LAST 12L /* not used */ + + /* These enums are for use with the CURLOPT_NETRC option. */ +#define CURL_NETRC_IGNORED 0L /* The .netrc will never be read. + This is the default. */ +#define CURL_NETRC_OPTIONAL 1L /* A user:password in the URL will be preferred + to one in the .netrc. */ +#define CURL_NETRC_REQUIRED 2L /* A user:password in the URL will be ignored. + Unless one is set programmatically, the + .netrc will be queried. */ +enum CURL_NETRC_OPTION { + /* we set a single member here, just to make sure we still provide the enum, + but the values to use are defined above with L suffixes */ + CURL_NETRC_LAST = 3 +}; + +#define CURL_SSLVERSION_DEFAULT 0L +#define CURL_SSLVERSION_TLSv1 1L /* TLS 1.x */ +#define CURL_SSLVERSION_SSLv2 2L +#define CURL_SSLVERSION_SSLv3 3L +#define CURL_SSLVERSION_TLSv1_0 4L +#define CURL_SSLVERSION_TLSv1_1 5L +#define CURL_SSLVERSION_TLSv1_2 6L +#define CURL_SSLVERSION_TLSv1_3 7L + +#define CURL_SSLVERSION_LAST 8L /* never use, keep last */ + +#define CURL_SSLVERSION_MAX_NONE 0L +#define CURL_SSLVERSION_MAX_DEFAULT (CURL_SSLVERSION_TLSv1 << 16) +#define CURL_SSLVERSION_MAX_TLSv1_0 (CURL_SSLVERSION_TLSv1_0 << 16) +#define CURL_SSLVERSION_MAX_TLSv1_1 (CURL_SSLVERSION_TLSv1_1 << 16) +#define CURL_SSLVERSION_MAX_TLSv1_2 (CURL_SSLVERSION_TLSv1_2 << 16) +#define CURL_SSLVERSION_MAX_TLSv1_3 (CURL_SSLVERSION_TLSv1_3 << 16) + +/* never use, keep last */ +#define CURL_SSLVERSION_MAX_LAST (CURL_SSLVERSION_LAST << 16) + +#define CURL_TLSAUTH_NONE 0L +#define CURL_TLSAUTH_SRP 1L + +enum CURL_TLSAUTH { + /* we set a single member here, just to make sure we still provide the enum, + but the values to use are defined above with L suffixes */ + CURL_TLSAUTH_LAST = 2 +}; + +/* symbols to use with CURLOPT_POSTREDIR. + CURL_REDIR_POST_301, CURL_REDIR_POST_302 and CURL_REDIR_POST_303 + can be bitwise ORed so that CURL_REDIR_POST_301 | CURL_REDIR_POST_302 + | CURL_REDIR_POST_303 == CURL_REDIR_POST_ALL */ + +#define CURL_REDIR_GET_ALL 0L +#define CURL_REDIR_POST_301 1L +#define CURL_REDIR_POST_302 2L +#define CURL_REDIR_POST_303 4L +#define CURL_REDIR_POST_ALL \ + (CURL_REDIR_POST_301 | CURL_REDIR_POST_302 | CURL_REDIR_POST_303) + +#define CURL_TIMECOND_NONE 0L +#define CURL_TIMECOND_IFMODSINCE 1L +#define CURL_TIMECOND_IFUNMODSINCE 2L +#define CURL_TIMECOND_LASTMOD 3L + +typedef enum { + /* we set a single member here, just to make sure we still provide + the enum typedef, but the values to use are defined above with L + suffixes */ + CURL_TIMECOND_LAST = 4 +} curl_TimeCond; + +/* Special size_t value signaling a null-terminated string. */ +#define CURL_ZERO_TERMINATED ((size_t)-1) + +/* curl_strequal() and curl_strnequal() are subject for removal in a future + release */ +CURL_EXTERN int curl_strequal(const char *s1, const char *s2); +CURL_EXTERN int curl_strnequal(const char *s1, const char *s2, size_t n); + +/* Mime/form handling support. */ +typedef struct curl_mime curl_mime; /* Mime context. */ +typedef struct curl_mimepart curl_mimepart; /* Mime part context. */ + +/* CURLMIMEOPT_ defines are for the CURLOPT_MIME_OPTIONS option. */ +#define CURLMIMEOPT_FORMESCAPE (1L << 0) /* Use backslash-escaping for forms */ + +/* + * NAME curl_mime_init() + * + * DESCRIPTION + * + * Create a mime context and return its handle. The easy parameter is the + * target handle. + */ +CURL_EXTERN curl_mime *curl_mime_init(CURL *easy); + +/* + * NAME curl_mime_free() + * + * DESCRIPTION + * + * release a mime handle and its substructures. + */ +CURL_EXTERN void curl_mime_free(curl_mime *mime); + +/* + * NAME curl_mime_addpart() + * + * DESCRIPTION + * + * Append a new empty part to the given mime context and return a handle to + * the created part. + */ +CURL_EXTERN curl_mimepart *curl_mime_addpart(curl_mime *mime); + +/* + * NAME curl_mime_name() + * + * DESCRIPTION + * + * Set mime/form part name. + */ +CURL_EXTERN CURLcode curl_mime_name(curl_mimepart *part, const char *name); + +/* + * NAME curl_mime_filename() + * + * DESCRIPTION + * + * Set mime part remote filename. + */ +CURL_EXTERN CURLcode curl_mime_filename(curl_mimepart *part, + const char *filename); + +/* + * NAME curl_mime_type() + * + * DESCRIPTION + * + * Set mime part type. + */ +CURL_EXTERN CURLcode curl_mime_type(curl_mimepart *part, const char *mimetype); + +/* + * NAME curl_mime_encoder() + * + * DESCRIPTION + * + * Set mime data transfer encoder. + */ +CURL_EXTERN CURLcode curl_mime_encoder(curl_mimepart *part, + const char *encoding); + +/* + * NAME curl_mime_data() + * + * DESCRIPTION + * + * Set mime part data source from memory data, + */ +CURL_EXTERN CURLcode curl_mime_data(curl_mimepart *part, + const char *data, size_t datasize); + +/* + * NAME curl_mime_filedata() + * + * DESCRIPTION + * + * Set mime part data source from named file. + */ +CURL_EXTERN CURLcode curl_mime_filedata(curl_mimepart *part, + const char *filename); + +/* + * NAME curl_mime_data_cb() + * + * DESCRIPTION + * + * Set mime part data source from callback function. + */ +CURL_EXTERN CURLcode curl_mime_data_cb(curl_mimepart *part, + curl_off_t datasize, + curl_read_callback readfunc, + curl_seek_callback seekfunc, + curl_free_callback freefunc, + void *arg); + +/* + * NAME curl_mime_subparts() + * + * DESCRIPTION + * + * Set mime part data source from subparts. + */ +CURL_EXTERN CURLcode curl_mime_subparts(curl_mimepart *part, + curl_mime *subparts); +/* + * NAME curl_mime_headers() + * + * DESCRIPTION + * + * Set mime part headers. + */ +CURL_EXTERN CURLcode curl_mime_headers(curl_mimepart *part, + struct curl_slist *headers, + int take_ownership); + +typedef enum { + /********* the first one is unused ************/ + CURLFORM_NOTHING CURL_DEPRECATED(7.56.0, ""), + CURLFORM_COPYNAME CURL_DEPRECATED(7.56.0, "Use curl_mime_name()"), + CURLFORM_PTRNAME CURL_DEPRECATED(7.56.0, "Use curl_mime_name()"), + CURLFORM_NAMELENGTH CURL_DEPRECATED(7.56.0, ""), + CURLFORM_COPYCONTENTS CURL_DEPRECATED(7.56.0, "Use curl_mime_data()"), + CURLFORM_PTRCONTENTS CURL_DEPRECATED(7.56.0, "Use curl_mime_data()"), + CURLFORM_CONTENTSLENGTH CURL_DEPRECATED(7.56.0, "Use curl_mime_data()"), + CURLFORM_FILECONTENT CURL_DEPRECATED(7.56.0, "Use curl_mime_data_cb()"), + CURLFORM_ARRAY CURL_DEPRECATED(7.56.0, ""), + CURLFORM_OBSOLETE, + CURLFORM_FILE CURL_DEPRECATED(7.56.0, "Use curl_mime_filedata()"), + + CURLFORM_BUFFER CURL_DEPRECATED(7.56.0, "Use curl_mime_filename()"), + CURLFORM_BUFFERPTR CURL_DEPRECATED(7.56.0, "Use curl_mime_data()"), + CURLFORM_BUFFERLENGTH CURL_DEPRECATED(7.56.0, "Use curl_mime_data()"), + + CURLFORM_CONTENTTYPE CURL_DEPRECATED(7.56.0, "Use curl_mime_type()"), + CURLFORM_CONTENTHEADER CURL_DEPRECATED(7.56.0, "Use curl_mime_headers()"), + CURLFORM_FILENAME CURL_DEPRECATED(7.56.0, "Use curl_mime_filename()"), + CURLFORM_END, + CURLFORM_OBSOLETE2, + + CURLFORM_STREAM CURL_DEPRECATED(7.56.0, "Use curl_mime_data_cb()"), + CURLFORM_CONTENTLEN /* added in 7.46.0, provide a curl_off_t length */ + CURL_DEPRECATED(7.56.0, "Use curl_mime_data()"), + + CURLFORM_LASTENTRY /* the last unused */ +} CURLformoption; + +/* structure to be used as parameter for CURLFORM_ARRAY */ +struct curl_forms { + CURLformoption option; + const char *value; +}; + +/* use this for multipart formpost building */ +/* Returns code for curl_formadd() + * + * Returns: + * CURL_FORMADD_OK on success + * CURL_FORMADD_MEMORY if the FormInfo allocation fails + * CURL_FORMADD_OPTION_TWICE if one option is given twice for one Form + * CURL_FORMADD_NULL if a null pointer was given for a char + * CURL_FORMADD_MEMORY if the allocation of a FormInfo struct failed + * CURL_FORMADD_UNKNOWN_OPTION if an unknown option was used + * CURL_FORMADD_INCOMPLETE if the some FormInfo is not complete (or error) + * CURL_FORMADD_MEMORY if a curl_httppost struct cannot be allocated + * CURL_FORMADD_MEMORY if some allocation for string copying failed. + * CURL_FORMADD_ILLEGAL_ARRAY if an illegal option is used in an array + * + ***************************************************************************/ +typedef enum { + CURL_FORMADD_OK CURL_DEPRECATED(7.56.0, ""), /* 1st, no error */ + + CURL_FORMADD_MEMORY CURL_DEPRECATED(7.56.0, ""), + CURL_FORMADD_OPTION_TWICE CURL_DEPRECATED(7.56.0, ""), + CURL_FORMADD_NULL CURL_DEPRECATED(7.56.0, ""), + CURL_FORMADD_UNKNOWN_OPTION CURL_DEPRECATED(7.56.0, ""), + CURL_FORMADD_INCOMPLETE CURL_DEPRECATED(7.56.0, ""), + CURL_FORMADD_ILLEGAL_ARRAY CURL_DEPRECATED(7.56.0, ""), + /* libcurl was built with form api disabled */ + CURL_FORMADD_DISABLED CURL_DEPRECATED(7.56.0, ""), + + CURL_FORMADD_LAST /* last */ +} CURLFORMcode; + +/* + * NAME curl_formadd() + * + * DESCRIPTION + * + * Pretty advanced function for building multi-part formposts. Each invoke + * adds one part that together construct a full post. Then use + * CURLOPT_HTTPPOST to send it off to libcurl. + */ +CURL_EXTERN CURLFORMcode CURL_DEPRECATED(7.56.0, "Use curl_mime_init()") +curl_formadd(struct curl_httppost **httppost, + struct curl_httppost **last_post, + ...); + +/* + * callback function for curl_formget() + * The void *arg pointer will be the one passed as second argument to + * curl_formget(). + * The character buffer passed to it must not be freed. + * Should return the buffer length passed to it as the argument "len" on + * success. + */ +typedef size_t (*curl_formget_callback)(void *arg, const char *buf, + size_t len); + +/* + * NAME curl_formget() + * + * DESCRIPTION + * + * Serialize a curl_httppost struct built with curl_formadd(). + * Accepts a void pointer as second argument which will be passed to + * the curl_formget_callback function. + * Returns 0 on success. + */ +CURL_EXTERN int CURL_DEPRECATED(7.56.0, "") +curl_formget(struct curl_httppost *form, void *arg, + curl_formget_callback append); +/* + * NAME curl_formfree() + * + * DESCRIPTION + * + * Free a multipart formpost previously built with curl_formadd(). + */ +CURL_EXTERN void CURL_DEPRECATED(7.56.0, "Use curl_mime_free()") +curl_formfree(struct curl_httppost *form); + +/* + * NAME curl_getenv() + * + * DESCRIPTION + * + * Returns a malloc()'ed string that MUST be curl_free()ed after usage is + * complete. DEPRECATED - see lib/README.curlx + */ +CURL_EXTERN char *curl_getenv(const char *variable); + +/* + * NAME curl_version() + * + * DESCRIPTION + * + * Returns a static ASCII string of the libcurl version. + */ +CURL_EXTERN char *curl_version(void); + +/* + * NAME curl_easy_escape() + * + * DESCRIPTION + * + * Escapes URL strings (converts all letters consider illegal in URLs to their + * %XX versions). This function returns a new allocated string or NULL if an + * error occurred. + */ +CURL_EXTERN char *curl_easy_escape(CURL *handle, + const char *string, + int length); + +/* the previous version: */ +CURL_EXTERN char *curl_escape(const char *string, + int length); + +/* + * NAME curl_easy_unescape() + * + * DESCRIPTION + * + * Unescapes URL encoding in strings (converts all %XX codes to their 8bit + * versions). This function returns a new allocated string or NULL if an error + * occurred. + * Conversion Note: On non-ASCII platforms the ASCII %XX codes are + * converted into the host encoding. + */ +CURL_EXTERN char *curl_easy_unescape(CURL *handle, + const char *string, + int length, + int *outlength); + +/* the previous version */ +CURL_EXTERN char *curl_unescape(const char *string, + int length); + +/* + * NAME curl_free() + * + * DESCRIPTION + * + * Provided for de-allocation in the same translation unit that did the + * allocation. Added in libcurl 7.10 + */ +CURL_EXTERN void curl_free(void *p); + +/* + * NAME curl_global_init() + * + * DESCRIPTION + * + * curl_global_init() should be invoked exactly once for each application that + * uses libcurl and before any call of other libcurl functions. + * + * This function is thread-safe if CURL_VERSION_THREADSAFE is set in the + * curl_version_info_data.features flag (fetch by curl_version_info()). + */ +CURL_EXTERN CURLcode curl_global_init(long flags); + +/* + * NAME curl_global_init_mem() + * + * DESCRIPTION + * + * curl_global_init() or curl_global_init_mem() should be invoked exactly once + * for each application that uses libcurl. This function can be used to + * initialize libcurl and set user defined memory management callback + * functions. Users can implement memory management routines to check for + * memory leaks, check for misuse of the curl library etc. User registered + * callback routines will be invoked by this library instead of the system + * memory management routines like malloc, free etc. + */ +CURL_EXTERN CURLcode curl_global_init_mem(long flags, + curl_malloc_callback m, + curl_free_callback f, + curl_realloc_callback r, + curl_strdup_callback s, + curl_calloc_callback c); + +/* + * NAME curl_global_cleanup() + * + * DESCRIPTION + * + * curl_global_cleanup() should be invoked exactly once for each application + * that uses libcurl + */ +CURL_EXTERN void curl_global_cleanup(void); + +/* + * NAME curl_global_trace() + * + * DESCRIPTION + * + * curl_global_trace() can be invoked at application start to + * configure which components in curl should participate in tracing. + * + * This function is thread-safe if CURL_VERSION_THREADSAFE is set in the + * curl_version_info_data.features flag (fetch by curl_version_info()). + */ +CURL_EXTERN CURLcode curl_global_trace(const char *config); + +/* linked-list structure for the CURLOPT_QUOTE option (and other) */ +struct curl_slist { + char *data; + struct curl_slist *next; +}; + +/* + * NAME curl_global_sslset() + * + * DESCRIPTION + * + * When built with multiple SSL backends, curl_global_sslset() allows to + * choose one. This function can only be called once, and it must be called + * *before* curl_global_init(). + * + * The backend can be identified by the id (e.g. CURLSSLBACKEND_OPENSSL). The + * backend can also be specified via the name parameter (passing -1 as id). If + * both id and name are specified, the name will be ignored. If neither id nor + * name are specified, the function will fail with CURLSSLSET_UNKNOWN_BACKEND + * and set the "avail" pointer to the NULL-terminated list of available + * backends. + * + * Upon success, the function returns CURLSSLSET_OK. + * + * If the specified SSL backend is not available, the function returns + * CURLSSLSET_UNKNOWN_BACKEND and sets the "avail" pointer to a + * NULL-terminated list of available SSL backends. + * + * The SSL backend can be set only once. If it has already been set, a + * subsequent attempt to change it will result in a CURLSSLSET_TOO_LATE. + */ + +struct curl_ssl_backend { + curl_sslbackend id; + const char *name; +}; +typedef struct curl_ssl_backend curl_ssl_backend; + +typedef enum { + CURLSSLSET_OK = 0, + CURLSSLSET_UNKNOWN_BACKEND, + CURLSSLSET_TOO_LATE, + CURLSSLSET_NO_BACKENDS /* libcurl was built without any SSL support */ +} CURLsslset; + +CURL_EXTERN CURLsslset curl_global_sslset(curl_sslbackend id, const char *name, + const curl_ssl_backend ***avail); + +/* + * NAME curl_slist_append() + * + * DESCRIPTION + * + * Appends a string to a linked list. If no list exists, it will be created + * first. Returns the new list, after appending. + */ +CURL_EXTERN struct curl_slist *curl_slist_append(struct curl_slist *list, + const char *data); + +/* + * NAME curl_slist_free_all() + * + * DESCRIPTION + * + * free a previously built curl_slist. + */ +CURL_EXTERN void curl_slist_free_all(struct curl_slist *list); + +/* + * NAME curl_getdate() + * + * DESCRIPTION + * + * Returns the time, in seconds since 1 Jan 1970 of the time string given in + * the first argument. The time argument in the second parameter is unused + * and should be set to NULL. + */ +CURL_EXTERN time_t curl_getdate(const char *p, const time_t *unused); + +/* info about the certificate chain, for SSL backends that support it. Asked + for with CURLOPT_CERTINFO / CURLINFO_CERTINFO */ +struct curl_certinfo { + int num_of_certs; /* number of certificates with information */ + struct curl_slist **certinfo; /* for each index in this array, there is a + linked list with textual information for a + certificate in the format "name:content". + eg "Subject:foo", "Issuer:bar", etc. */ +}; + +/* Information about the SSL library used and the respective internal SSL + handle, which can be used to obtain further information regarding the + connection. Asked for with CURLINFO_TLS_SSL_PTR or CURLINFO_TLS_SESSION. */ +struct curl_tlssessioninfo { + curl_sslbackend backend; + void *internals; +}; + +#define CURLINFO_STRING 0x100000 +#define CURLINFO_LONG 0x200000 +#define CURLINFO_DOUBLE 0x300000 +#define CURLINFO_SLIST 0x400000 +#define CURLINFO_PTR 0x400000 /* same as SLIST */ +#define CURLINFO_SOCKET 0x500000 +#define CURLINFO_OFF_T 0x600000 +#define CURLINFO_MASK 0x0fffff +#define CURLINFO_TYPEMASK 0xf00000 + +typedef enum { + CURLINFO_NONE, /* first, never use this */ + CURLINFO_EFFECTIVE_URL = CURLINFO_STRING + 1, + CURLINFO_RESPONSE_CODE = CURLINFO_LONG + 2, + CURLINFO_TOTAL_TIME = CURLINFO_DOUBLE + 3, + CURLINFO_NAMELOOKUP_TIME = CURLINFO_DOUBLE + 4, + CURLINFO_CONNECT_TIME = CURLINFO_DOUBLE + 5, + CURLINFO_PRETRANSFER_TIME = CURLINFO_DOUBLE + 6, + CURLINFO_SIZE_UPLOAD CURL_DEPRECATED(7.55.0, "Use CURLINFO_SIZE_UPLOAD_T") + = CURLINFO_DOUBLE + 7, + CURLINFO_SIZE_UPLOAD_T = CURLINFO_OFF_T + 7, + CURLINFO_SIZE_DOWNLOAD + CURL_DEPRECATED(7.55.0, "Use CURLINFO_SIZE_DOWNLOAD_T") + = CURLINFO_DOUBLE + 8, + CURLINFO_SIZE_DOWNLOAD_T = CURLINFO_OFF_T + 8, + CURLINFO_SPEED_DOWNLOAD + CURL_DEPRECATED(7.55.0, "Use CURLINFO_SPEED_DOWNLOAD_T") + = CURLINFO_DOUBLE + 9, + CURLINFO_SPEED_DOWNLOAD_T = CURLINFO_OFF_T + 9, + CURLINFO_SPEED_UPLOAD + CURL_DEPRECATED(7.55.0, "Use CURLINFO_SPEED_UPLOAD_T") + = CURLINFO_DOUBLE + 10, + CURLINFO_SPEED_UPLOAD_T = CURLINFO_OFF_T + 10, + CURLINFO_HEADER_SIZE = CURLINFO_LONG + 11, + CURLINFO_REQUEST_SIZE = CURLINFO_LONG + 12, + CURLINFO_SSL_VERIFYRESULT = CURLINFO_LONG + 13, + CURLINFO_FILETIME = CURLINFO_LONG + 14, + CURLINFO_FILETIME_T = CURLINFO_OFF_T + 14, + CURLINFO_CONTENT_LENGTH_DOWNLOAD + CURL_DEPRECATED(7.55.0, + "Use CURLINFO_CONTENT_LENGTH_DOWNLOAD_T") + = CURLINFO_DOUBLE + 15, + CURLINFO_CONTENT_LENGTH_DOWNLOAD_T = CURLINFO_OFF_T + 15, + CURLINFO_CONTENT_LENGTH_UPLOAD + CURL_DEPRECATED(7.55.0, + "Use CURLINFO_CONTENT_LENGTH_UPLOAD_T") + = CURLINFO_DOUBLE + 16, + CURLINFO_CONTENT_LENGTH_UPLOAD_T = CURLINFO_OFF_T + 16, + CURLINFO_STARTTRANSFER_TIME = CURLINFO_DOUBLE + 17, + CURLINFO_CONTENT_TYPE = CURLINFO_STRING + 18, + CURLINFO_REDIRECT_TIME = CURLINFO_DOUBLE + 19, + CURLINFO_REDIRECT_COUNT = CURLINFO_LONG + 20, + CURLINFO_PRIVATE = CURLINFO_STRING + 21, + CURLINFO_HTTP_CONNECTCODE = CURLINFO_LONG + 22, + CURLINFO_HTTPAUTH_AVAIL = CURLINFO_LONG + 23, + CURLINFO_PROXYAUTH_AVAIL = CURLINFO_LONG + 24, + CURLINFO_OS_ERRNO = CURLINFO_LONG + 25, + CURLINFO_NUM_CONNECTS = CURLINFO_LONG + 26, + CURLINFO_SSL_ENGINES = CURLINFO_SLIST + 27, + CURLINFO_COOKIELIST = CURLINFO_SLIST + 28, + CURLINFO_LASTSOCKET CURL_DEPRECATED(7.45.0, "Use CURLINFO_ACTIVESOCKET") + = CURLINFO_LONG + 29, + CURLINFO_FTP_ENTRY_PATH = CURLINFO_STRING + 30, + CURLINFO_REDIRECT_URL = CURLINFO_STRING + 31, + CURLINFO_PRIMARY_IP = CURLINFO_STRING + 32, + CURLINFO_APPCONNECT_TIME = CURLINFO_DOUBLE + 33, + CURLINFO_CERTINFO = CURLINFO_PTR + 34, + CURLINFO_CONDITION_UNMET = CURLINFO_LONG + 35, + CURLINFO_RTSP_SESSION_ID = CURLINFO_STRING + 36, + CURLINFO_RTSP_CLIENT_CSEQ = CURLINFO_LONG + 37, + CURLINFO_RTSP_SERVER_CSEQ = CURLINFO_LONG + 38, + CURLINFO_RTSP_CSEQ_RECV = CURLINFO_LONG + 39, + CURLINFO_PRIMARY_PORT = CURLINFO_LONG + 40, + CURLINFO_LOCAL_IP = CURLINFO_STRING + 41, + CURLINFO_LOCAL_PORT = CURLINFO_LONG + 42, + CURLINFO_TLS_SESSION CURL_DEPRECATED(7.48.0, "Use CURLINFO_TLS_SSL_PTR") + = CURLINFO_PTR + 43, + CURLINFO_ACTIVESOCKET = CURLINFO_SOCKET + 44, + CURLINFO_TLS_SSL_PTR = CURLINFO_PTR + 45, + CURLINFO_HTTP_VERSION = CURLINFO_LONG + 46, + CURLINFO_PROXY_SSL_VERIFYRESULT = CURLINFO_LONG + 47, + CURLINFO_PROTOCOL CURL_DEPRECATED(7.85.0, "Use CURLINFO_SCHEME") + = CURLINFO_LONG + 48, + CURLINFO_SCHEME = CURLINFO_STRING + 49, + CURLINFO_TOTAL_TIME_T = CURLINFO_OFF_T + 50, + CURLINFO_NAMELOOKUP_TIME_T = CURLINFO_OFF_T + 51, + CURLINFO_CONNECT_TIME_T = CURLINFO_OFF_T + 52, + CURLINFO_PRETRANSFER_TIME_T = CURLINFO_OFF_T + 53, + CURLINFO_STARTTRANSFER_TIME_T = CURLINFO_OFF_T + 54, + CURLINFO_REDIRECT_TIME_T = CURLINFO_OFF_T + 55, + CURLINFO_APPCONNECT_TIME_T = CURLINFO_OFF_T + 56, + CURLINFO_RETRY_AFTER = CURLINFO_OFF_T + 57, + CURLINFO_EFFECTIVE_METHOD = CURLINFO_STRING + 58, + CURLINFO_PROXY_ERROR = CURLINFO_LONG + 59, + CURLINFO_REFERER = CURLINFO_STRING + 60, + CURLINFO_CAINFO = CURLINFO_STRING + 61, + CURLINFO_CAPATH = CURLINFO_STRING + 62, + CURLINFO_XFER_ID = CURLINFO_OFF_T + 63, + CURLINFO_CONN_ID = CURLINFO_OFF_T + 64, + CURLINFO_QUEUE_TIME_T = CURLINFO_OFF_T + 65, + CURLINFO_USED_PROXY = CURLINFO_LONG + 66, + CURLINFO_POSTTRANSFER_TIME_T = CURLINFO_OFF_T + 67, + CURLINFO_EARLYDATA_SENT_T = CURLINFO_OFF_T + 68, + CURLINFO_HTTPAUTH_USED = CURLINFO_LONG + 69, + CURLINFO_PROXYAUTH_USED = CURLINFO_LONG + 70, + CURLINFO_LASTONE = 70 +} CURLINFO; + +/* CURLINFO_RESPONSE_CODE is the new name for the option previously known as + CURLINFO_HTTP_CODE */ +#define CURLINFO_HTTP_CODE CURLINFO_RESPONSE_CODE + +typedef enum { + CURLCLOSEPOLICY_NONE, /* first, never use this */ + + CURLCLOSEPOLICY_OLDEST, + CURLCLOSEPOLICY_LEAST_RECENTLY_USED, + CURLCLOSEPOLICY_LEAST_TRAFFIC, + CURLCLOSEPOLICY_SLOWEST, + CURLCLOSEPOLICY_CALLBACK, + + CURLCLOSEPOLICY_LAST /* last, never use this */ +} curl_closepolicy; + +#define CURL_GLOBAL_SSL (1 << 0) /* no purpose since 7.57.0 */ +#define CURL_GLOBAL_WIN32 (1 << 1) +#define CURL_GLOBAL_ALL (CURL_GLOBAL_SSL | CURL_GLOBAL_WIN32) +#define CURL_GLOBAL_NOTHING 0 +#define CURL_GLOBAL_DEFAULT CURL_GLOBAL_ALL +#define CURL_GLOBAL_ACK_EINTR (1 << 2) + +/***************************************************************************** + * Setup defines, protos etc for the sharing stuff. + */ + +/* Different data locks for a single share */ +typedef enum { + CURL_LOCK_DATA_NONE = 0, + /* CURL_LOCK_DATA_SHARE is used internally to say that + * the locking is just made to change the internal state of the share + * itself. + */ + CURL_LOCK_DATA_SHARE, + CURL_LOCK_DATA_COOKIE, + CURL_LOCK_DATA_DNS, + CURL_LOCK_DATA_SSL_SESSION, + CURL_LOCK_DATA_CONNECT, + CURL_LOCK_DATA_PSL, + CURL_LOCK_DATA_HSTS, + CURL_LOCK_DATA_LAST +} curl_lock_data; + +/* Different lock access types */ +typedef enum { + CURL_LOCK_ACCESS_NONE = 0, /* unspecified action */ + CURL_LOCK_ACCESS_SHARED = 1, /* for read perhaps */ + CURL_LOCK_ACCESS_SINGLE = 2, /* for write perhaps */ + CURL_LOCK_ACCESS_LAST /* never use */ +} curl_lock_access; + +typedef void (*curl_lock_function)(CURL *handle, + curl_lock_data data, + curl_lock_access locktype, + void *userptr); +typedef void (*curl_unlock_function)(CURL *handle, + curl_lock_data data, + void *userptr); + +typedef enum { + CURLSHE_OK, /* all is fine */ + CURLSHE_BAD_OPTION, /* 1 */ + CURLSHE_IN_USE, /* 2 */ + CURLSHE_INVALID, /* 3 */ + CURLSHE_NOMEM, /* 4 out of memory */ + CURLSHE_NOT_BUILT_IN, /* 5 feature not present in lib */ + CURLSHE_LAST /* never use */ +} CURLSHcode; + +typedef enum { + CURLSHOPT_NONE, /* do not use */ + CURLSHOPT_SHARE, /* specify a data type to share */ + CURLSHOPT_UNSHARE, /* specify which data type to stop sharing */ + CURLSHOPT_LOCKFUNC, /* pass in a 'curl_lock_function' pointer */ + CURLSHOPT_UNLOCKFUNC, /* pass in a 'curl_unlock_function' pointer */ + CURLSHOPT_USERDATA, /* pass in a user data pointer used in the lock/unlock + callback functions */ + CURLSHOPT_LAST /* never use */ +} CURLSHoption; + +CURL_EXTERN CURLSH *curl_share_init(void); +CURL_EXTERN CURLSHcode curl_share_setopt(CURLSH *share, CURLSHoption option, + ...); +CURL_EXTERN CURLSHcode curl_share_cleanup(CURLSH *share); + +/**************************************************************************** + * Structures for querying information about the curl library at runtime. + */ + +typedef enum { + CURLVERSION_FIRST, /* 7.10 */ + CURLVERSION_SECOND, /* 7.11.1 */ + CURLVERSION_THIRD, /* 7.12.0 */ + CURLVERSION_FOURTH, /* 7.16.1 */ + CURLVERSION_FIFTH, /* 7.57.0 */ + CURLVERSION_SIXTH, /* 7.66.0 */ + CURLVERSION_SEVENTH, /* 7.70.0 */ + CURLVERSION_EIGHTH, /* 7.72.0 */ + CURLVERSION_NINTH, /* 7.75.0 */ + CURLVERSION_TENTH, /* 7.77.0 */ + CURLVERSION_ELEVENTH, /* 7.87.0 */ + CURLVERSION_TWELFTH, /* 8.8.0 */ + CURLVERSION_LAST /* never actually use this */ +} CURLversion; + +/* The 'CURLVERSION_NOW' is the symbolic name meant to be used by + basically all programs ever that want to get version information. It is + meant to be a built-in version number for what kind of struct the caller + expects. If the struct ever changes, we redefine the NOW to another enum + from above. */ +#define CURLVERSION_NOW CURLVERSION_TWELFTH + +struct curl_version_info_data { + CURLversion age; /* age of the returned struct */ + const char *version; /* LIBCURL_VERSION */ + unsigned int version_num; /* LIBCURL_VERSION_NUM */ + const char *host; /* OS/host/cpu/machine when configured */ + int features; /* bitmask, see defines below */ + const char *ssl_version; /* human readable string */ + long ssl_version_num; /* not used anymore, always 0 */ + const char *libz_version; /* human readable string */ + /* protocols is terminated by an entry with a NULL protoname */ + const char * const *protocols; + + /* The fields below this were added in CURLVERSION_SECOND */ + const char *ares; + int ares_num; + + /* This field was added in CURLVERSION_THIRD */ + const char *libidn; + + /* These field were added in CURLVERSION_FOURTH */ + + /* Same as '_libiconv_version' if built with HAVE_ICONV */ + int iconv_ver_num; + + const char *libssh_version; /* human readable string */ + + /* These fields were added in CURLVERSION_FIFTH */ + unsigned int brotli_ver_num; /* Numeric Brotli version + (MAJOR << 24) | (MINOR << 12) | PATCH */ + const char *brotli_version; /* human readable string. */ + + /* These fields were added in CURLVERSION_SIXTH */ + unsigned int nghttp2_ver_num; /* Numeric nghttp2 version + (MAJOR << 16) | (MINOR << 8) | PATCH */ + const char *nghttp2_version; /* human readable string. */ + const char *quic_version; /* human readable quic (+ HTTP/3) library + + version or NULL */ + + /* These fields were added in CURLVERSION_SEVENTH */ + const char *cainfo; /* the built-in default CURLOPT_CAINFO, might + be NULL */ + const char *capath; /* the built-in default CURLOPT_CAPATH, might + be NULL */ + + /* These fields were added in CURLVERSION_EIGHTH */ + unsigned int zstd_ver_num; /* Numeric Zstd version + (MAJOR << 24) | (MINOR << 12) | PATCH */ + const char *zstd_version; /* human readable string. */ + + /* These fields were added in CURLVERSION_NINTH */ + const char *hyper_version; /* human readable string. */ + + /* These fields were added in CURLVERSION_TENTH */ + const char *gsasl_version; /* human readable string. */ + + /* These fields were added in CURLVERSION_ELEVENTH */ + /* feature_names is terminated by an entry with a NULL feature name */ + const char * const *feature_names; + + /* These fields were added in CURLVERSION_TWELFTH */ + const char *rtmp_version; /* human readable string. */ +}; +typedef struct curl_version_info_data curl_version_info_data; + +#define CURL_VERSION_IPV6 (1<<0) /* IPv6-enabled */ +#define CURL_VERSION_KERBEROS4 (1<<1) /* Kerberos V4 auth is supported + (deprecated) */ +#define CURL_VERSION_SSL (1<<2) /* SSL options are present */ +#define CURL_VERSION_LIBZ (1<<3) /* libz features are present */ +#define CURL_VERSION_NTLM (1<<4) /* NTLM auth is supported */ +#define CURL_VERSION_GSSNEGOTIATE (1<<5) /* Negotiate auth is supported + (deprecated) */ +#define CURL_VERSION_DEBUG (1<<6) /* Built with debug capabilities */ +#define CURL_VERSION_ASYNCHDNS (1<<7) /* Asynchronous DNS resolves */ +#define CURL_VERSION_SPNEGO (1<<8) /* SPNEGO auth is supported */ +#define CURL_VERSION_LARGEFILE (1<<9) /* Supports files larger than 2GB */ +#define CURL_VERSION_IDN (1<<10) /* Internationized Domain Names are + supported */ +#define CURL_VERSION_SSPI (1<<11) /* Built against Windows SSPI */ +#define CURL_VERSION_CONV (1<<12) /* Character conversions supported */ +#define CURL_VERSION_CURLDEBUG (1<<13) /* Debug memory tracking supported */ +#define CURL_VERSION_TLSAUTH_SRP (1<<14) /* TLS-SRP auth is supported */ +#define CURL_VERSION_NTLM_WB (1<<15) /* NTLM delegation to winbind helper + is supported */ +#define CURL_VERSION_HTTP2 (1<<16) /* HTTP2 support built-in */ +#define CURL_VERSION_GSSAPI (1<<17) /* Built against a GSS-API library */ +#define CURL_VERSION_KERBEROS5 (1<<18) /* Kerberos V5 auth is supported */ +#define CURL_VERSION_UNIX_SOCKETS (1<<19) /* Unix domain sockets support */ +#define CURL_VERSION_PSL (1<<20) /* Mozilla's Public Suffix List, used + for cookie domain verification */ +#define CURL_VERSION_HTTPS_PROXY (1<<21) /* HTTPS-proxy support built-in */ +#define CURL_VERSION_MULTI_SSL (1<<22) /* Multiple SSL backends available */ +#define CURL_VERSION_BROTLI (1<<23) /* Brotli features are present. */ +#define CURL_VERSION_ALTSVC (1<<24) /* Alt-Svc handling built-in */ +#define CURL_VERSION_HTTP3 (1<<25) /* HTTP3 support built-in */ +#define CURL_VERSION_ZSTD (1<<26) /* zstd features are present */ +#define CURL_VERSION_UNICODE (1<<27) /* Unicode support on Windows */ +#define CURL_VERSION_HSTS (1<<28) /* HSTS is supported */ +#define CURL_VERSION_GSASL (1<<29) /* libgsasl is supported */ +#define CURL_VERSION_THREADSAFE (1<<30) /* libcurl API is thread-safe */ + +/* + * NAME curl_version_info() + * + * DESCRIPTION + * + * This function returns a pointer to a static copy of the version info + * struct. See above. + */ +CURL_EXTERN curl_version_info_data *curl_version_info(CURLversion); + +/* + * NAME curl_easy_strerror() + * + * DESCRIPTION + * + * The curl_easy_strerror function may be used to turn a CURLcode value + * into the equivalent human readable error string. This is useful + * for printing meaningful error messages. + */ +CURL_EXTERN const char *curl_easy_strerror(CURLcode); + +/* + * NAME curl_share_strerror() + * + * DESCRIPTION + * + * The curl_share_strerror function may be used to turn a CURLSHcode value + * into the equivalent human readable error string. This is useful + * for printing meaningful error messages. + */ +CURL_EXTERN const char *curl_share_strerror(CURLSHcode); + +/* + * NAME curl_easy_pause() + * + * DESCRIPTION + * + * The curl_easy_pause function pauses or unpauses transfers. Select the new + * state by setting the bitmask, use the convenience defines below. + * + */ +CURL_EXTERN CURLcode curl_easy_pause(CURL *handle, int bitmask); + +#define CURLPAUSE_RECV (1 << 0) +#define CURLPAUSE_RECV_CONT (0) + +#define CURLPAUSE_SEND (1 << 2) +#define CURLPAUSE_SEND_CONT (0) + +#define CURLPAUSE_ALL (CURLPAUSE_RECV | CURLPAUSE_SEND) +#define CURLPAUSE_CONT (CURLPAUSE_RECV_CONT | CURLPAUSE_SEND_CONT) + +/* + * NAME curl_easy_ssls_import() + * + * DESCRIPTION + * + * The curl_easy_ssls_import function adds a previously exported SSL session + * to the SSL session cache of the easy handle (or the underlying share). + */ +CURL_EXTERN CURLcode curl_easy_ssls_import(CURL *handle, + const char *session_key, + const unsigned char *shmac, + size_t shmac_len, + const unsigned char *sdata, + size_t sdata_len); + +/* This is the curl_ssls_export_cb callback prototype. It + * is passed to curl_easy_ssls_export() to extract SSL sessions/tickets. */ +typedef CURLcode curl_ssls_export_cb(CURL *handle, + void *userptr, + const char *session_key, + const unsigned char *shmac, + size_t shmac_len, + const unsigned char *sdata, + size_t sdata_len, + curl_off_t valid_until, + int ietf_tls_id, + const char *alpn, + size_t earlydata_max); + +/* + * NAME curl_easy_ssls_export() + * + * DESCRIPTION + * + * The curl_easy_ssls_export function iterates over all SSL sessions stored + * in the easy handle (or underlying share) and invokes the passed + * callback. + * + */ +CURL_EXTERN CURLcode curl_easy_ssls_export(CURL *handle, + curl_ssls_export_cb *export_fn, + void *userptr); + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif + +/* unfortunately, the easy.h and multi.h include files need options and info + stuff before they can be included! */ +#include "easy.h" /* nothing in curl is fun without the easy stuff */ +#include "multi.h" +#include "urlapi.h" +#include "options.h" +#include "header.h" +#include "websockets.h" +#include "mprintf.h" + +/* the typechecker does not work in C++ (yet) */ +#if ((defined(__GNUC__) && defined(__GNUC_MINOR__) && \ + ((__GNUC__ > 4) || (__GNUC__ == 4 && __GNUC_MINOR__ >= 3))) || \ + (defined(__clang__) && __clang_major__ >= 14)) && \ + !defined(__cplusplus) && !defined(CURL_DISABLE_TYPECHECK) +#include "typecheck-gcc.h" +#else +#if defined(__STDC__) && (__STDC__ >= 1) +/* This preprocessor magic that replaces a call with the exact same call is + only done to make sure application authors pass exactly three arguments + to these functions. */ +#define curl_easy_setopt(handle,opt,param) curl_easy_setopt(handle,opt,param) +#define curl_easy_getinfo(handle,info,arg) curl_easy_getinfo(handle,info,arg) +#define curl_share_setopt(share,opt,param) curl_share_setopt(share,opt,param) +#define curl_multi_setopt(handle,opt,param) curl_multi_setopt(handle,opt,param) +#endif /* __STDC__ >= 1 */ +#endif /* gcc >= 4.3 && !__cplusplus && !CURL_DISABLE_TYPECHECK */ + +#endif /* CURLINC_CURL_H */ diff --git a/cactus-engine/libs/curl/include/curl/curlver.h b/cactus-engine/libs/curl/include/curl/curlver.h new file mode 100644 index 000000000..d1608c9ce --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/curlver.h @@ -0,0 +1,78 @@ +#ifndef CURLINC_CURLVER_H +#define CURLINC_CURLVER_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +/* This header file contains nothing but libcurl version info, generated by + a script at release-time. This was made its own header file in 7.11.2 */ + +/* This is the global package copyright */ +#define LIBCURL_COPYRIGHT "Daniel Stenberg, ." + +/* This is the version number of the libcurl package from which this header + file origins: */ +#define LIBCURL_VERSION "8.18.0" + +/* The numeric version number is also available "in parts" by using these + defines: */ +#define LIBCURL_VERSION_MAJOR 8 +#define LIBCURL_VERSION_MINOR 18 +#define LIBCURL_VERSION_PATCH 0 +/* This is the numeric version of the libcurl version number, meant for easier + parsing and comparisons by programs. The LIBCURL_VERSION_NUM define will + always follow this syntax: + + 0xXXYYZZ + + Where XX, YY and ZZ are the main version, release and patch numbers in + hexadecimal (using 8 bits each). All three numbers are always represented + using two digits. 1.2 would appear as "0x010200" while version 9.11.7 + appears as "0x090b07". + + This 6-digit (24 bits) hexadecimal number does not show pre-release number, + and it is always a greater number in a more recent release. It makes + comparisons with greater than and less than work. + + Note: This define is the full hex number and _does not_ use the + CURL_VERSION_BITS() macro since curl's own configure script greps for it + and needs it to contain the full number. +*/ +#define LIBCURL_VERSION_NUM 0x081200 + +/* + * This is the date and time when the full source package was created. The + * timestamp is not stored in git, as the timestamp is properly set in the + * tarballs by the maketgz script. + * + * The format of the date follows this template: + * + * "2007-11-23" + */ +#define LIBCURL_TIMESTAMP "2026-01-07" + +#define CURL_VERSION_BITS(x, y, z) ((x) << 16 | (y) << 8 | (z)) +#define CURL_AT_LEAST_VERSION(x, y, z) \ + (LIBCURL_VERSION_NUM >= CURL_VERSION_BITS(x, y, z)) + +#endif /* CURLINC_CURLVER_H */ diff --git a/cactus-engine/libs/curl/include/curl/easy.h b/cactus-engine/libs/curl/include/curl/easy.h new file mode 100644 index 000000000..5b3cdbd64 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/easy.h @@ -0,0 +1,123 @@ +#ifndef CURLINC_EASY_H +#define CURLINC_EASY_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ +#ifdef __cplusplus +extern "C" { +#endif + +/* Flag bits in the curl_blob struct: */ +#define CURL_BLOB_COPY 1 /* tell libcurl to copy the data */ +#define CURL_BLOB_NOCOPY 0 /* tell libcurl to NOT copy the data */ + +struct curl_blob { + void *data; + size_t len; + unsigned int flags; /* bit 0 is defined, the rest are reserved and should be + left zeroes */ +}; + +CURL_EXTERN CURL *curl_easy_init(void); +CURL_EXTERN CURLcode curl_easy_setopt(CURL *curl, CURLoption option, ...); +CURL_EXTERN CURLcode curl_easy_perform(CURL *curl); +CURL_EXTERN void curl_easy_cleanup(CURL *curl); + +/* + * NAME curl_easy_getinfo() + * + * DESCRIPTION + * + * Request internal information from the curl session with this function. + * The third argument MUST be pointing to the specific type of the used option + * which is documented in each man page of the option. The data pointed to + * will be filled in accordingly and can be relied upon only if the function + * returns CURLE_OK. This function is intended to get used *AFTER* a performed + * transfer, all results from this function are undefined until the transfer + * is completed. + */ +CURL_EXTERN CURLcode curl_easy_getinfo(CURL *curl, CURLINFO info, ...); + +/* + * NAME curl_easy_duphandle() + * + * DESCRIPTION + * + * Creates a new curl session handle with the same options set for the handle + * passed in. Duplicating a handle could only be a matter of cloning data and + * options, internal state info and things like persistent connections cannot + * be transferred. It is useful in multi-threaded applications when you can run + * curl_easy_duphandle() for each new thread to avoid a series of identical + * curl_easy_setopt() invokes in every thread. + */ +CURL_EXTERN CURL *curl_easy_duphandle(CURL *curl); + +/* + * NAME curl_easy_reset() + * + * DESCRIPTION + * + * Re-initializes a curl handle to the default values. This puts back the + * handle to the same state as it was in when it was just created. + * + * It does keep: live connections, the Session ID cache, the DNS cache and the + * cookies. + */ +CURL_EXTERN void curl_easy_reset(CURL *curl); + +/* + * NAME curl_easy_recv() + * + * DESCRIPTION + * + * Receives data from the connected socket. Use after successful + * curl_easy_perform() with CURLOPT_CONNECT_ONLY option. + */ +CURL_EXTERN CURLcode curl_easy_recv(CURL *curl, void *buffer, size_t buflen, + size_t *n); + +/* + * NAME curl_easy_send() + * + * DESCRIPTION + * + * Sends data over the connected socket. Use after successful + * curl_easy_perform() with CURLOPT_CONNECT_ONLY option. + */ +CURL_EXTERN CURLcode curl_easy_send(CURL *curl, const void *buffer, + size_t buflen, size_t *n); + +/* + * NAME curl_easy_upkeep() + * + * DESCRIPTION + * + * Performs connection upkeep for the given session handle. + */ +CURL_EXTERN CURLcode curl_easy_upkeep(CURL *curl); + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif + +#endif diff --git a/cactus-engine/libs/curl/include/curl/header.h b/cactus-engine/libs/curl/include/curl/header.h new file mode 100644 index 000000000..e7334b5a3 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/header.h @@ -0,0 +1,74 @@ +#ifndef CURLINC_HEADER_H +#define CURLINC_HEADER_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +#ifdef __cplusplus +extern "C" { +#endif + +struct curl_header { + char *name; /* this might not use the same case */ + char *value; + size_t amount; /* number of headers using this name */ + size_t index; /* ... of this instance, 0 or higher */ + unsigned int origin; /* see bits below */ + void *anchor; /* handle privately used by libcurl */ +}; + +/* 'origin' bits */ +#define CURLH_HEADER (1 << 0) /* plain server header */ +#define CURLH_TRAILER (1 << 1) /* trailers */ +#define CURLH_CONNECT (1 << 2) /* CONNECT headers */ +#define CURLH_1XX (1 << 3) /* 1xx headers */ +#define CURLH_PSEUDO (1 << 4) /* pseudo headers */ + +typedef enum { + CURLHE_OK, + CURLHE_BADINDEX, /* header exists but not with this index */ + CURLHE_MISSING, /* no such header exists */ + CURLHE_NOHEADERS, /* no headers at all exist (yet) */ + CURLHE_NOREQUEST, /* no request with this number was used */ + CURLHE_OUT_OF_MEMORY, /* out of memory while processing */ + CURLHE_BAD_ARGUMENT, /* a function argument was not okay */ + CURLHE_NOT_BUILT_IN /* if API was disabled in the build */ +} CURLHcode; + +CURL_EXTERN CURLHcode curl_easy_header(CURL *easy, + const char *name, + size_t index, + unsigned int origin, + int request, + struct curl_header **hout); + +CURL_EXTERN struct curl_header *curl_easy_nextheader(CURL *easy, + unsigned int origin, + int request, + struct curl_header *prev); + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif + +#endif /* CURLINC_HEADER_H */ diff --git a/cactus-engine/libs/curl/include/curl/mprintf.h b/cactus-engine/libs/curl/include/curl/mprintf.h new file mode 100644 index 000000000..9272e7489 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/mprintf.h @@ -0,0 +1,85 @@ +#ifndef CURLINC_MPRINTF_H +#define CURLINC_MPRINTF_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +#include +#include /* needed for FILE */ +#include "curl.h" /* for CURL_EXTERN */ + +#ifdef __cplusplus +extern "C" { +#endif + +#ifndef CURL_TEMP_PRINTF +#if (defined(__GNUC__) || defined(__clang__) || \ + defined(__IAR_SYSTEMS_ICC__)) && \ + defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) && \ + !defined(CURL_NO_FMT_CHECKS) +#if defined(__MINGW32__) && !defined(__clang__) +#ifdef __MINGW_PRINTF_FORMAT /* mingw-w64 3.0.0+. Needs stdio.h. */ +#define CURL_TEMP_PRINTF(fmt, arg) \ + __attribute__((format(__MINGW_PRINTF_FORMAT, fmt, arg))) +#else +#define CURL_TEMP_PRINTF(fmt, arg) +#endif +#else +#define CURL_TEMP_PRINTF(fmt, arg) \ + __attribute__((format(printf, fmt, arg))) +#endif +#else +#define CURL_TEMP_PRINTF(fmt, arg) +#endif +#endif + +CURL_EXTERN int curl_mprintf(const char *format, ...) + CURL_TEMP_PRINTF(1, 2); +CURL_EXTERN int curl_mfprintf(FILE *fd, const char *format, ...) + CURL_TEMP_PRINTF(2, 3); +CURL_EXTERN int curl_msprintf(char *buffer, const char *format, ...) + CURL_TEMP_PRINTF(2, 3); +CURL_EXTERN int curl_msnprintf(char *buffer, size_t maxlength, + const char *format, ...) + CURL_TEMP_PRINTF(3, 4); +CURL_EXTERN int curl_mvprintf(const char *format, va_list args) + CURL_TEMP_PRINTF(1, 0); +CURL_EXTERN int curl_mvfprintf(FILE *fd, const char *format, va_list args) + CURL_TEMP_PRINTF(2, 0); +CURL_EXTERN int curl_mvsprintf(char *buffer, const char *format, va_list args) + CURL_TEMP_PRINTF(2, 0); +CURL_EXTERN int curl_mvsnprintf(char *buffer, size_t maxlength, + const char *format, va_list args) + CURL_TEMP_PRINTF(3, 0); +CURL_EXTERN char *curl_maprintf(const char *format, ...) + CURL_TEMP_PRINTF(1, 2); +CURL_EXTERN char *curl_mvaprintf(const char *format, va_list args) + CURL_TEMP_PRINTF(1, 0); + +#undef CURL_TEMP_PRINTF + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif + +#endif /* CURLINC_MPRINTF_H */ diff --git a/cactus-engine/libs/curl/include/curl/multi.h b/cactus-engine/libs/curl/include/curl/multi.h new file mode 100644 index 000000000..531c1a954 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/multi.h @@ -0,0 +1,551 @@ +#ifndef CURLINC_MULTI_H +#define CURLINC_MULTI_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ +/* + This is an "external" header file. Do not give away any internals here! + + GOALS + + o Enable a "pull" interface. The application that uses libcurl decides where + and when to ask libcurl to get/send data. + + o Enable multiple simultaneous transfers in the same thread without making it + complicated for the application. + + o Enable the application to select() on its own file descriptors and curl's + file descriptors simultaneous easily. + +*/ + +/* + * This header file should not really need to include "curl.h" since curl.h + * itself includes this file and we expect user applications to do #include + * without the need for especially including multi.h. + * + * For some reason we added this include here at one point, and rather than to + * break existing (wrongly written) libcurl applications, we leave it as-is + * but with this warning attached. + */ +#include "curl.h" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef void CURLM; + +typedef enum { + CURLM_CALL_MULTI_PERFORM = -1, /* please call curl_multi_perform() or + curl_multi_socket*() soon */ + CURLM_OK, + CURLM_BAD_HANDLE, /* the passed-in handle is not a valid CURLM handle */ + CURLM_BAD_EASY_HANDLE, /* an easy handle was not good/valid */ + CURLM_OUT_OF_MEMORY, /* if you ever get this, you are in deep sh*t */ + CURLM_INTERNAL_ERROR, /* this is a libcurl bug */ + CURLM_BAD_SOCKET, /* the passed in socket argument did not match */ + CURLM_UNKNOWN_OPTION, /* curl_multi_setopt() with unsupported option */ + CURLM_ADDED_ALREADY, /* an easy handle already added to a multi handle was + attempted to get added - again */ + CURLM_RECURSIVE_API_CALL, /* an api function was called from inside a + callback */ + CURLM_WAKEUP_FAILURE, /* wakeup is unavailable or failed */ + CURLM_BAD_FUNCTION_ARGUMENT, /* function called with a bad parameter */ + CURLM_ABORTED_BY_CALLBACK, + CURLM_UNRECOVERABLE_POLL, + CURLM_LAST +} CURLMcode; + +/* just to make code nicer when using curl_multi_socket() you can now check + for CURLM_CALL_MULTI_SOCKET too in the same style it works for + curl_multi_perform() and CURLM_CALL_MULTI_PERFORM */ +#define CURLM_CALL_MULTI_SOCKET CURLM_CALL_MULTI_PERFORM + +/* bitmask bits for CURLMOPT_PIPELINING */ +#define CURLPIPE_NOTHING 0L +#define CURLPIPE_HTTP1 1L +#define CURLPIPE_MULTIPLEX 2L + +typedef enum { + CURLMSG_NONE, /* first, not used */ + CURLMSG_DONE, /* This easy handle has completed. 'result' contains + the CURLcode of the transfer */ + CURLMSG_LAST /* last, not used */ +} CURLMSG; + +struct CURLMsg { + CURLMSG msg; /* what this message means */ + CURL *easy_handle; /* the handle it concerns */ + union { + void *whatever; /* message-specific data */ + CURLcode result; /* return code for transfer */ + } data; +}; +typedef struct CURLMsg CURLMsg; + +/* Based on poll(2) structure and values. + * We do not use pollfd and POLL* constants explicitly + * to cover platforms without poll(). */ +#define CURL_WAIT_POLLIN 0x0001 +#define CURL_WAIT_POLLPRI 0x0002 +#define CURL_WAIT_POLLOUT 0x0004 + +struct curl_waitfd { + curl_socket_t fd; + short events; + short revents; +}; + +/* + * Name: curl_multi_init() + * + * Desc: initialize multi-style curl usage + * + * Returns: a new CURLM handle to use in all 'curl_multi' functions. + */ +CURL_EXTERN CURLM *curl_multi_init(void); + +/* + * Name: curl_multi_add_handle() + * + * Desc: add a standard curl handle to the multi stack + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_add_handle(CURLM *multi_handle, + CURL *curl_handle); + +/* + * Name: curl_multi_remove_handle() + * + * Desc: removes a curl handle from the multi stack again + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_remove_handle(CURLM *multi_handle, + CURL *curl_handle); + +/* + * Name: curl_multi_fdset() + * + * Desc: Ask curl for its fd_set sets. The app can use these to select() or + * poll() on. We want curl_multi_perform() called as soon as one of + * them are ready. + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_fdset(CURLM *multi_handle, + fd_set *read_fd_set, + fd_set *write_fd_set, + fd_set *exc_fd_set, + int *max_fd); + +/* + * Name: curl_multi_wait() + * + * Desc: Poll on all fds within a CURLM set as well as any + * additional fds passed to the function. + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_wait(CURLM *multi_handle, + struct curl_waitfd extra_fds[], + unsigned int extra_nfds, + int timeout_ms, + int *ret); + +/* + * Name: curl_multi_poll() + * + * Desc: Poll on all fds within a CURLM set as well as any + * additional fds passed to the function. + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_poll(CURLM *multi_handle, + struct curl_waitfd extra_fds[], + unsigned int extra_nfds, + int timeout_ms, + int *ret); + +/* + * Name: curl_multi_wakeup() + * + * Desc: wakes up a sleeping curl_multi_poll call. + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_wakeup(CURLM *multi_handle); + +/* + * Name: curl_multi_perform() + * + * Desc: When the app thinks there is data available for curl it calls this + * function to read/write whatever there is right now. This returns + * as soon as the reads and writes are done. This function does not + * require that there actually is data available for reading or that + * data can be written, it can be called just in case. It returns + * the number of handles that still transfer data in the second + * argument's integer-pointer. + * + * Returns: CURLMcode type, general multi error code. *NOTE* that this only + * returns errors etc regarding the whole multi stack. There might + * still have occurred problems on individual transfers even when + * this returns OK. + */ +CURL_EXTERN CURLMcode curl_multi_perform(CURLM *multi_handle, + int *running_handles); + +/* + * Name: curl_multi_cleanup() + * + * Desc: Cleans up and removes a whole multi stack. It does not free or + * touch any individual easy handles in any way. We need to define + * in what state those handles will be if this function is called + * in the middle of a transfer. + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_cleanup(CURLM *multi_handle); + +/* + * Name: curl_multi_info_read() + * + * Desc: Ask the multi handle if there is any messages/informationals from + * the individual transfers. Messages include informationals such as + * error code from the transfer or just the fact that a transfer is + * completed. More details on these should be written down as well. + * + * Repeated calls to this function will return a new struct each + * time, until a special "end of msgs" struct is returned as a signal + * that there is no more to get at this point. + * + * The data the returned pointer points to will not survive calling + * curl_multi_cleanup(). + * + * The 'CURLMsg' struct is meant to be simple and only contain basic + * information. If more involved information is wanted, we will + * provide the particular "transfer handle" in that struct and that + * should/could/would be used in subsequent curl_easy_getinfo() calls + * (or similar). The point being that we must never expose complex + * structs to applications, as then we will undoubtably get backwards + * compatibility problems in the future. + * + * Returns: A pointer to a filled-in struct, or NULL if it failed or ran out + * of structs. It also writes the number of messages left in the + * queue (after this read) in the integer the second argument points + * to. + */ +CURL_EXTERN CURLMsg *curl_multi_info_read(CURLM *multi_handle, + int *msgs_in_queue); + +/* + * Name: curl_multi_strerror() + * + * Desc: The curl_multi_strerror function may be used to turn a CURLMcode + * value into the equivalent human readable error string. This is + * useful for printing meaningful error messages. + * + * Returns: A pointer to a null-terminated error message. + */ +CURL_EXTERN const char *curl_multi_strerror(CURLMcode); + +/* + * Name: curl_multi_socket() and + * curl_multi_socket_all() + * + * Desc: An alternative version of curl_multi_perform() that allows the + * application to pass in one of the file descriptors that have been + * detected to have "action" on them and let libcurl perform. + * See man page for details. + */ +#define CURL_POLL_NONE 0 +#define CURL_POLL_IN 1 +#define CURL_POLL_OUT 2 +#define CURL_POLL_INOUT 3 +#define CURL_POLL_REMOVE 4 + +#define CURL_SOCKET_TIMEOUT CURL_SOCKET_BAD + +#define CURL_CSELECT_IN 0x01 +#define CURL_CSELECT_OUT 0x02 +#define CURL_CSELECT_ERR 0x04 + +typedef int (*curl_socket_callback)(CURL *easy, /* easy handle */ + curl_socket_t s, /* socket */ + int what, /* see above */ + void *userp, /* private callback + pointer */ + void *socketp); /* private socket + pointer */ +/* + * Name: curl_multi_timer_callback + * + * Desc: Called by libcurl whenever the library detects a change in the + * maximum number of milliseconds the app is allowed to wait before + * curl_multi_socket() or curl_multi_perform() must be called + * (to allow libcurl's timed events to take place). + * + * Returns: The callback should return zero. + */ +typedef int (*curl_multi_timer_callback)(CURLM *multi, /* multi handle */ + long timeout_ms, /* see above */ + void *userp); /* private callback + pointer */ + +CURL_EXTERN CURLMcode CURL_DEPRECATED(7.19.5, "Use curl_multi_socket_action()") +curl_multi_socket(CURLM *multi_handle, curl_socket_t s, int *running_handles); + +CURL_EXTERN CURLMcode curl_multi_socket_action(CURLM *multi_handle, + curl_socket_t s, + int ev_bitmask, + int *running_handles); + +CURL_EXTERN CURLMcode CURL_DEPRECATED(7.19.5, "Use curl_multi_socket_action()") +curl_multi_socket_all(CURLM *multi_handle, int *running_handles); + +#ifndef CURL_ALLOW_OLD_MULTI_SOCKET +/* This macro below was added in 7.16.3 to push users who recompile to use + * the new curl_multi_socket_action() instead of the old curl_multi_socket() + */ +#define curl_multi_socket(x,y,z) curl_multi_socket_action(x,y,0,z) +#endif + +/* + * Name: curl_multi_timeout() + * + * Desc: Returns the maximum number of milliseconds the app is allowed to + * wait before curl_multi_socket() or curl_multi_perform() must be + * called (to allow libcurl's timed events to take place). + * + * Returns: CURLM error code. + */ +CURL_EXTERN CURLMcode curl_multi_timeout(CURLM *multi_handle, + long *milliseconds); + +typedef enum { + /* This is the socket callback function pointer */ + CURLOPT(CURLMOPT_SOCKETFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 1), + + /* This is the argument passed to the socket callback */ + CURLOPT(CURLMOPT_SOCKETDATA, CURLOPTTYPE_OBJECTPOINT, 2), + + /* set to 1 to enable pipelining for this multi handle */ + CURLOPT(CURLMOPT_PIPELINING, CURLOPTTYPE_LONG, 3), + + /* This is the timer callback function pointer */ + CURLOPT(CURLMOPT_TIMERFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 4), + + /* This is the argument passed to the timer callback */ + CURLOPT(CURLMOPT_TIMERDATA, CURLOPTTYPE_OBJECTPOINT, 5), + + /* maximum number of entries in the connection cache */ + CURLOPT(CURLMOPT_MAXCONNECTS, CURLOPTTYPE_LONG, 6), + + /* maximum number of (pipelining) connections to one host */ + CURLOPT(CURLMOPT_MAX_HOST_CONNECTIONS, CURLOPTTYPE_LONG, 7), + + /* maximum number of requests in a pipeline */ + CURLOPT(CURLMOPT_MAX_PIPELINE_LENGTH, CURLOPTTYPE_LONG, 8), + + /* a connection with a content-length longer than this + will not be considered for pipelining */ + CURLOPT(CURLMOPT_CONTENT_LENGTH_PENALTY_SIZE, CURLOPTTYPE_OFF_T, 9), + + /* a connection with a chunk length longer than this + will not be considered for pipelining */ + CURLOPT(CURLMOPT_CHUNK_LENGTH_PENALTY_SIZE, CURLOPTTYPE_OFF_T, 10), + + /* a list of site names(+port) that are blocked from pipelining */ + CURLOPT(CURLMOPT_PIPELINING_SITE_BL, CURLOPTTYPE_OBJECTPOINT, 11), + + /* a list of server types that are blocked from pipelining */ + CURLOPT(CURLMOPT_PIPELINING_SERVER_BL, CURLOPTTYPE_OBJECTPOINT, 12), + + /* maximum number of open connections in total */ + CURLOPT(CURLMOPT_MAX_TOTAL_CONNECTIONS, CURLOPTTYPE_LONG, 13), + + /* This is the server push callback function pointer */ + CURLOPT(CURLMOPT_PUSHFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 14), + + /* This is the argument passed to the server push callback */ + CURLOPT(CURLMOPT_PUSHDATA, CURLOPTTYPE_OBJECTPOINT, 15), + + /* maximum number of concurrent streams to support on a connection */ + CURLOPT(CURLMOPT_MAX_CONCURRENT_STREAMS, CURLOPTTYPE_LONG, 16), + + /* network has changed, adjust caches/connection reuse */ + CURLOPT(CURLMOPT_NETWORK_CHANGED, CURLOPTTYPE_LONG, 17), + + /* This is the notify callback function pointer */ + CURLOPT(CURLMOPT_NOTIFYFUNCTION, CURLOPTTYPE_FUNCTIONPOINT, 18), + + /* This is the argument passed to the notify callback */ + CURLOPT(CURLMOPT_NOTIFYDATA, CURLOPTTYPE_OBJECTPOINT, 19), + + CURLMOPT_LASTENTRY /* the last unused */ +} CURLMoption; + +/* Definition of bits for the CURLMOPT_NETWORK_CHANGED argument: */ + +/* - CURLMNWC_CLEAR_CONNS tells libcurl to prevent further reuse of existing + connections. Connections that are idle will be closed. Ongoing transfers + will continue with the connection they have. */ +#define CURLMNWC_CLEAR_CONNS (1L << 0) + +/* - CURLMNWC_CLEAR_DNS tells libcurl to prevent further reuse of existing + connections. Connections that are idle will be closed. Ongoing transfers + will continue with the connection they have. */ +#define CURLMNWC_CLEAR_DNS (1L << 0) + +/* + * Name: curl_multi_setopt() + * + * Desc: Sets options for the multi handle. + * + * Returns: CURLM error code. + */ +CURL_EXTERN CURLMcode curl_multi_setopt(CURLM *multi_handle, + CURLMoption option, ...); + +/* + * Name: curl_multi_assign() + * + * Desc: This function sets an association in the multi handle between the + * given socket and a private pointer of the application. This is + * (only) useful for curl_multi_socket uses. + * + * Returns: CURLM error code. + */ +CURL_EXTERN CURLMcode curl_multi_assign(CURLM *multi_handle, + curl_socket_t sockfd, void *sockp); + +/* + * Name: curl_multi_get_handles() + * + * Desc: Returns an allocated array holding all handles currently added to + * the multi handle. Marks the final entry with a NULL pointer. If + * there is no easy handle added to the multi handle, this function + * returns an array with the first entry as a NULL pointer. + * + * Returns: NULL on failure, otherwise a CURL **array pointer + */ +CURL_EXTERN CURL **curl_multi_get_handles(CURLM *multi_handle); + +typedef enum { + CURLMINFO_NONE, /* first, never use this */ + /* The number of easy handles currently managed by the multi handle, + * e.g. have been added but not yet removed. */ + CURLMINFO_XFERS_CURRENT = 1, + /* The number of easy handles running, e.g. not done and not queueing. */ + CURLMINFO_XFERS_RUNNING = 2, + /* The number of easy handles waiting to start, e.g. for a connection + * to become available due to limits on parallelism, max connections + * or other factors. */ + CURLMINFO_XFERS_PENDING = 3, + /* The number of easy handles finished, waiting for their results to + * be read via `curl_multi_info_read()`. */ + CURLMINFO_XFERS_DONE = 4, + /* The total number of easy handles added to the multi handle, ever. */ + CURLMINFO_XFERS_ADDED = 5, + + CURLMINFO_LASTENTRY /* the last unused */ +} CURLMinfo_offt; + +/* + * Name: curl_multi_get_offt() + * + * Desc: Retrieves a numeric value for the `CURLMINFO_*` enums. + * + * Returns: CULRM_OK or error when value could not be obtained. + */ +CURL_EXTERN CURLMcode curl_multi_get_offt(CURLM *multi_handle, + CURLMinfo_offt info, + curl_off_t *pvalue); + +/* + * Name: curl_push_callback + * + * Desc: This callback gets called when a new stream is being pushed by the + * server. It approves or denies the new stream. It can also decide + * to completely fail the connection. + * + * Returns: CURL_PUSH_OK, CURL_PUSH_DENY or CURL_PUSH_ERROROUT + */ +#define CURL_PUSH_OK 0 +#define CURL_PUSH_DENY 1 +#define CURL_PUSH_ERROROUT 2 /* added in 7.72.0 */ + +struct curl_pushheaders; /* forward declaration only */ + +CURL_EXTERN char *curl_pushheader_bynum(struct curl_pushheaders *h, + size_t num); +CURL_EXTERN char *curl_pushheader_byname(struct curl_pushheaders *h, + const char *name); + +typedef int (*curl_push_callback)(CURL *parent, + CURL *easy, + size_t num_headers, + struct curl_pushheaders *headers, + void *userp); + +/* + * Name: curl_multi_waitfds() + * + * Desc: Ask curl for fds for polling. The app can use these to poll on. + * We want curl_multi_perform() called as soon as one of them are + * ready. Passing zero size allows to get just a number of fds. + * + * Returns: CURLMcode type, general multi error code. + */ +CURL_EXTERN CURLMcode curl_multi_waitfds(CURLM *multi, + struct curl_waitfd *ufds, + unsigned int size, + unsigned int *fd_count); + +/* + * Notifications dispatched by a multi handle, when enabled. + */ +#define CURLMNOTIFY_INFO_READ 0 +#define CURLMNOTIFY_EASY_DONE 1 + +/* + * Callback to install via CURLMOPT_NOTIFYFUNCTION. + */ +typedef void (*curl_notify_callback)(CURLM *multi, + unsigned int notification, + CURL *easy, + void *user_data); + +CURL_EXTERN CURLMcode curl_multi_notify_disable(CURLM *multi, + unsigned int notification); + +CURL_EXTERN CURLMcode curl_multi_notify_enable(CURLM *multi, + unsigned int notification); + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif + +#endif diff --git a/cactus-engine/libs/curl/include/curl/options.h b/cactus-engine/libs/curl/include/curl/options.h new file mode 100644 index 000000000..835a722e9 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/options.h @@ -0,0 +1,70 @@ +#ifndef CURLINC_OPTIONS_H +#define CURLINC_OPTIONS_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +#ifdef __cplusplus +extern "C" { +#endif + +typedef enum { + CURLOT_LONG, /* long (a range of values) */ + CURLOT_VALUES, /* (a defined set or bitmask) */ + CURLOT_OFF_T, /* curl_off_t (a range of values) */ + CURLOT_OBJECT, /* pointer (void *) */ + CURLOT_STRING, /* (char * to null-terminated buffer) */ + CURLOT_SLIST, /* (struct curl_slist *) */ + CURLOT_CBPTR, /* (void * passed as-is to a callback) */ + CURLOT_BLOB, /* blob (struct curl_blob *) */ + CURLOT_FUNCTION /* function pointer */ +} curl_easytype; + +/* Flag bits */ + +/* "alias" means it is provided for old programs to remain functional, + we prefer another name */ +#define CURLOT_FLAG_ALIAS (1 << 0) + +/* The CURLOPTTYPE_* id ranges can still be used to figure out what type/size + to use for curl_easy_setopt() for the given id */ +struct curl_easyoption { + const char *name; + CURLoption id; + curl_easytype type; + unsigned int flags; +}; + +CURL_EXTERN const struct curl_easyoption * +curl_easy_option_by_name(const char *name); + +CURL_EXTERN const struct curl_easyoption * +curl_easy_option_by_id(CURLoption id); + +CURL_EXTERN const struct curl_easyoption * +curl_easy_option_next(const struct curl_easyoption *prev); + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif +#endif /* CURLINC_OPTIONS_H */ diff --git a/cactus-engine/libs/curl/include/curl/stdcheaders.h b/cactus-engine/libs/curl/include/curl/stdcheaders.h new file mode 100644 index 000000000..7451aa305 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/stdcheaders.h @@ -0,0 +1,35 @@ +#ifndef CURLINC_STDCHEADERS_H +#define CURLINC_STDCHEADERS_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +#include + +size_t fread(void *, size_t, size_t, FILE *); +size_t fwrite(const void *, size_t, size_t, FILE *); + +int strcasecmp(const char *, const char *); +int strncasecmp(const char *, const char *, size_t); + +#endif /* CURLINC_STDCHEADERS_H */ diff --git a/cactus-engine/libs/curl/include/curl/system.h b/cactus-engine/libs/curl/include/curl/system.h new file mode 100644 index 000000000..a5b3e9eba --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/system.h @@ -0,0 +1,399 @@ +#ifndef CURLINC_SYSTEM_H +#define CURLINC_SYSTEM_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +/* + * Try to keep one section per platform, compiler and architecture, otherwise, + * if an existing section is reused for a different one and later on the + * original is adjusted, probably the piggybacking one can be adversely + * changed. + * + * In order to differentiate between platforms/compilers/architectures use + * only compiler built-in predefined preprocessor symbols. + * + * curl_off_t + * ---------- + * + * For any given platform/compiler curl_off_t MUST be typedef'ed to a 64-bit + * wide signed integral data type. The width of this data type must remain + * constant and independent of any possible large file support settings. + * + * As a general rule, curl_off_t shall not be mapped to off_t. This rule shall + * only be violated if off_t is the only 64-bit data type available and the + * size of off_t is independent of large file support settings. Keep your + * build on the safe side avoiding an off_t gating. If you have a 64-bit + * off_t then take for sure that another 64-bit data type exists, dig deeper + * and you will find it. + * + */ + +#ifdef __DJGPP__ +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# define CURL_TYPEOF_CURL_SOCKLEN_T int + +#elif defined(__BORLANDC__) +# define CURL_TYPEOF_CURL_OFF_T __int64 +# define CURL_FORMAT_CURL_OFF_T "I64d" +# define CURL_FORMAT_CURL_OFF_TU "I64u" +# define CURL_SUFFIX_CURL_OFF_T i64 +# define CURL_SUFFIX_CURL_OFF_TU ui64 +# define CURL_TYPEOF_CURL_SOCKLEN_T int + +#elif defined(__POCC__) +# ifdef _MSC_VER +# define CURL_TYPEOF_CURL_OFF_T __int64 +# define CURL_FORMAT_CURL_OFF_T "I64d" +# define CURL_FORMAT_CURL_OFF_TU "I64u" +# define CURL_SUFFIX_CURL_OFF_T i64 +# define CURL_SUFFIX_CURL_OFF_TU ui64 +# else +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T int + +#elif defined(__LCC__) +# ifdef __MCST__ /* MCST eLbrus Compiler Collection */ +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 +# else /* Local (or Little) C Compiler */ +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# define CURL_TYPEOF_CURL_SOCKLEN_T int +# endif + +#elif defined(macintosh) +# include +# if TYPE_LONGLONG +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# else +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T unsigned int + +#elif defined(__TANDEM) +# ifndef __LP64 +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# define CURL_TYPEOF_CURL_SOCKLEN_T int +# else +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# define CURL_TYPEOF_CURL_SOCKLEN_T unsigned int +# endif + +#elif defined(UNDER_CE) +# define CURL_TYPEOF_CURL_OFF_T __int64 +# define CURL_FORMAT_CURL_OFF_T "I64d" +# define CURL_FORMAT_CURL_OFF_TU "I64u" +# define CURL_SUFFIX_CURL_OFF_T i64 +# define CURL_SUFFIX_CURL_OFF_TU ui64 +# define CURL_TYPEOF_CURL_SOCKLEN_T int + +#elif defined(__MINGW32__) +# include +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T PRId64 +# define CURL_FORMAT_CURL_OFF_TU PRIu64 +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# define CURL_TYPEOF_CURL_SOCKLEN_T int +# define CURL_PULL_SYS_TYPES_H 1 + +#elif defined(__VMS) +# ifdef __VAX +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# else +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T unsigned int + +#elif defined(__OS400__) +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +#elif defined(__MVS__) +# ifdef _LONG_LONG +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# else /* _LP64 and default */ +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +#elif defined(__370__) +# if defined(__IBMC__) || defined(__IBMCPP__) +# ifdef _LONG_LONG +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# else /* _LP64 and default */ +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 +# endif + +#elif defined(TPF) +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# define CURL_TYPEOF_CURL_SOCKLEN_T int + +#elif defined(__TINYC__) /* also known as tcc */ +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +#elif defined(__SUNPRO_C) || defined(__SUNPRO_CC) /* Oracle Solaris Studio */ +# if !defined(__LP64) && (defined(__ILP32) || \ + defined(__i386) || \ + defined(__sparcv8) || \ + defined(__sparcv8plus)) +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# elif defined(__LP64) || \ + defined(__amd64) || defined(__sparcv9) +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +#elif defined(__xlc__) /* IBM xlc compiler */ +# ifndef _LP64 +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# else +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +#elif defined(__hpux) /* HP aCC compiler */ +# ifndef _LP64 +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# else +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +/* ===================================== */ +/* KEEP MSVC THE PENULTIMATE ENTRY */ +/* ===================================== */ + +#elif defined(_MSC_VER) +# if (_MSC_VER >= 1800) +# include +# define CURL_FORMAT_CURL_OFF_T PRId64 +# define CURL_FORMAT_CURL_OFF_TU PRIu64 +# else +# define CURL_FORMAT_CURL_OFF_T "I64d" +# define CURL_FORMAT_CURL_OFF_TU "I64u" +# endif +# define CURL_TYPEOF_CURL_OFF_T __int64 +# define CURL_SUFFIX_CURL_OFF_T i64 +# define CURL_SUFFIX_CURL_OFF_TU ui64 +# define CURL_TYPEOF_CURL_SOCKLEN_T int + +/* ===================================== */ +/* KEEP GENERIC GCC THE LAST ENTRY */ +/* ===================================== */ + +#elif defined(__GNUC__) && !defined(_SCO_DS) +# if !defined(__LP64__) && \ + (defined(__ILP32__) || defined(__i386__) || defined(__hppa__) || \ + defined(__ppc__) || defined(__powerpc__) || defined(__arm__) || \ + defined(__sparc__) || defined(__mips__) || defined(__sh__) || \ + defined(__XTENSA__) || \ + (defined(__SIZEOF_LONG__) && __SIZEOF_LONG__ == 4) || \ + (defined(__LONG_MAX__) && __LONG_MAX__ == 2147483647L)) +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# if (__GNUC__ >= 4) || \ + ((__GNUC__ == 3) && defined(__GNUC_MINOR__) && (__GNUC_MINOR__ >= 4)) +# define CURL_POPCOUNT64(x) __builtin_popcountll(x) +# define CURL_CTZ64(x) __builtin_ctzll(x) +# endif +# elif defined(__LP64__) || \ + defined(__x86_64__) || defined(__ppc64__) || defined(__sparc64__) || \ + defined(__e2k__) || \ + (defined(__SIZEOF_LONG__) && __SIZEOF_LONG__ == 8) || \ + (defined(__LONG_MAX__) && __LONG_MAX__ == 9223372036854775807L) +# define CURL_TYPEOF_CURL_OFF_T long +# define CURL_FORMAT_CURL_OFF_T "ld" +# define CURL_FORMAT_CURL_OFF_TU "lu" +# define CURL_SUFFIX_CURL_OFF_T L +# define CURL_SUFFIX_CURL_OFF_TU UL +# if (__GNUC__ >= 4) || \ + ((__GNUC__ == 3) && defined(__GNUC_MINOR__) && (__GNUC_MINOR__ >= 4)) +# define CURL_POPCOUNT64(x) __builtin_popcountl(x) +# define CURL_CTZ64(x) __builtin_ctzl(x) +# endif +# endif +# define CURL_TYPEOF_CURL_SOCKLEN_T socklen_t +# define CURL_PULL_SYS_TYPES_H 1 +# define CURL_PULL_SYS_SOCKET_H 1 + +#else +/* generic "safe guess" on old 32-bit style */ +# define CURL_TYPEOF_CURL_OFF_T long long +# define CURL_FORMAT_CURL_OFF_T "lld" +# define CURL_FORMAT_CURL_OFF_TU "llu" +# define CURL_SUFFIX_CURL_OFF_T LL +# define CURL_SUFFIX_CURL_OFF_TU ULL +# define CURL_TYPEOF_CURL_SOCKLEN_T int +#endif + +#ifdef _AIX +/* AIX needs */ +#define CURL_PULL_SYS_POLL_H +#endif + +/* CURL_PULL_SYS_TYPES_H is defined above when inclusion of header file */ +/* sys/types.h is required here to properly make type definitions below. */ +#ifdef CURL_PULL_SYS_TYPES_H +# include +#endif + +/* CURL_PULL_SYS_SOCKET_H is defined above when inclusion of header file */ +/* sys/socket.h is required here to properly make type definitions below. */ +#ifdef CURL_PULL_SYS_SOCKET_H +# include +#endif + +/* CURL_PULL_SYS_POLL_H is defined above when inclusion of header file */ +/* sys/poll.h is required here to properly make type definitions below. */ +#ifdef CURL_PULL_SYS_POLL_H +# include +#endif + +/* Data type definition of curl_socklen_t. */ +#ifdef CURL_TYPEOF_CURL_SOCKLEN_T + typedef CURL_TYPEOF_CURL_SOCKLEN_T curl_socklen_t; +#endif + +/* Data type definition of curl_off_t. */ + +#ifdef CURL_TYPEOF_CURL_OFF_T + typedef CURL_TYPEOF_CURL_OFF_T curl_off_t; +#endif + +#endif /* CURLINC_SYSTEM_H */ diff --git a/cactus-engine/libs/curl/include/curl/typecheck-gcc.h b/cactus-engine/libs/curl/include/curl/typecheck-gcc.h new file mode 100644 index 000000000..0642afd91 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/typecheck-gcc.h @@ -0,0 +1,957 @@ +#ifndef CURLINC_TYPECHECK_GCC_H +#define CURLINC_TYPECHECK_GCC_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +/* wraps curl_easy_setopt() with type checking */ + +/* To add a new kind of warning, add an + * if(curlcheck_sometype_option(_curl_opt)) + * if(!curlcheck_sometype(value)) + * Wcurl_easy_setopt_err_sometype(); + * block and define curlcheck_sometype_option, curlcheck_sometype and + * Wcurl_easy_setopt_err_sometype below + * + * NOTE: We use two nested 'if' statements here instead of the && operator, in + * order to work around gcc bug #32061. It affects only gcc 4.3.x/4.4.x + * when compiling with -Wlogical-op. + * + * To add an option that uses the same type as an existing option, you will + * just need to extend the appropriate _curl_*_option macro + */ + +#define curl_easy_setopt(handle, option, value) \ + __extension__({ \ + if(__builtin_constant_p(option)) { \ + CURL_IGNORE_DEPRECATION( \ + if(curlcheck_long_option(option)) \ + if(!curlcheck_long(value)) \ + Wcurl_easy_setopt_err_long(); \ + if(curlcheck_off_t_option(option)) \ + if(!curlcheck_off_t(value)) \ + Wcurl_easy_setopt_err_curl_off_t(); \ + if(curlcheck_string_option(option)) \ + if(!curlcheck_string(value)) \ + Wcurl_easy_setopt_err_string(); \ + if((option) == CURLOPT_PRIVATE) { } \ + if(curlcheck_write_cb_option(option)) \ + if(!curlcheck_write_cb(value)) \ + Wcurl_easy_setopt_err_write_callback(); \ + if(curlcheck_curl_option(option)) \ + if(!curlcheck_curl(value)) \ + Wcurl_easy_setopt_err_curl(); \ + if((option) == CURLOPT_RESOLVER_START_FUNCTION) \ + if(!curlcheck_resolver_start_callback(value)) \ + Wcurl_easy_setopt_err_resolver_start_callback(); \ + if((option) == CURLOPT_READFUNCTION) \ + if(!curlcheck_read_cb(value)) \ + Wcurl_easy_setopt_err_read_cb(); \ + if((option) == CURLOPT_IOCTLFUNCTION) \ + if(!curlcheck_ioctl_cb(value)) \ + Wcurl_easy_setopt_err_ioctl_cb(); \ + if((option) == CURLOPT_SOCKOPTFUNCTION) \ + if(!curlcheck_sockopt_cb(value)) \ + Wcurl_easy_setopt_err_sockopt_cb(); \ + if((option) == CURLOPT_OPENSOCKETFUNCTION) \ + if(!curlcheck_opensocket_cb(value)) \ + Wcurl_easy_setopt_err_opensocket_cb(); \ + if((option) == CURLOPT_PROGRESSFUNCTION) \ + if(!curlcheck_progress_cb(value)) \ + Wcurl_easy_setopt_err_progress_cb(); \ + if((option) == CURLOPT_XFERINFOFUNCTION) \ + if(!curlcheck_xferinfo_cb(value)) \ + Wcurl_easy_setopt_err_xferinfo_cb(); \ + if((option) == CURLOPT_DEBUGFUNCTION) \ + if(!curlcheck_debug_cb(value)) \ + Wcurl_easy_setopt_err_debug_cb(); \ + if((option) == CURLOPT_SSL_CTX_FUNCTION) \ + if(!curlcheck_ssl_ctx_cb(value)) \ + Wcurl_easy_setopt_err_ssl_ctx_cb(); \ + if(curlcheck_conv_cb_option(option)) \ + if(!curlcheck_conv_cb(value)) \ + Wcurl_easy_setopt_err_conv_cb(); \ + if((option) == CURLOPT_SEEKFUNCTION) \ + if(!curlcheck_seek_cb(value)) \ + Wcurl_easy_setopt_err_seek_cb(); \ + if((option) == CURLOPT_CHUNK_BGN_FUNCTION) \ + if(!curlcheck_chunk_bgn_cb(value)) \ + Wcurl_easy_setopt_err_chunk_bgn_cb(); \ + if((option) == CURLOPT_CHUNK_END_FUNCTION) \ + if(!curlcheck_chunk_end_cb(value)) \ + Wcurl_easy_setopt_err_chunk_end_cb(); \ + if((option) == CURLOPT_CLOSESOCKETFUNCTION) \ + if(!curlcheck_close_socket_cb(value)) \ + Wcurl_easy_setopt_err_close_socket_cb(); \ + if((option) == CURLOPT_FNMATCH_FUNCTION) \ + if(!curlcheck_fnmatch_cb(value)) \ + Wcurl_easy_setopt_err_fnmatch_cb(); \ + if((option) == CURLOPT_HSTSREADFUNCTION) \ + if(!curlcheck_hstsread_cb(value)) \ + Wcurl_easy_setopt_err_hstsread_cb(); \ + if((option) == CURLOPT_HSTSWRITEFUNCTION) \ + if(!curlcheck_hstswrite_cb(value)) \ + Wcurl_easy_setopt_err_hstswrite_cb(); \ + if((option) == CURLOPT_SSH_HOSTKEYFUNCTION) \ + if(!curlcheck_ssh_hostkey_cb(value)) \ + Wcurl_easy_setopt_err_ssh_hostkey_cb(); \ + if((option) == CURLOPT_SSH_KEYFUNCTION) \ + if(!curlcheck_ssh_key_cb(value)) \ + Wcurl_easy_setopt_err_ssh_key_cb(); \ + if((option) == CURLOPT_INTERLEAVEFUNCTION) \ + if(!curlcheck_interleave_cb(value)) \ + Wcurl_easy_setopt_err_interleave_cb(); \ + if((option) == CURLOPT_PREREQFUNCTION) \ + if(!curlcheck_prereq_cb(value)) \ + Wcurl_easy_setopt_err_prereq_cb(); \ + if((option) == CURLOPT_TRAILERFUNCTION) \ + if(!curlcheck_trailer_cb(value)) \ + Wcurl_easy_setopt_err_trailer_cb(); \ + if(curlcheck_cb_data_option(option)) \ + if(!curlcheck_cb_data(value)) \ + Wcurl_easy_setopt_err_cb_data(); \ + if((option) == CURLOPT_ERRORBUFFER) \ + if(!curlcheck_error_buffer(value)) \ + Wcurl_easy_setopt_err_error_buffer(); \ + if((option) == CURLOPT_CURLU) \ + if(!curlcheck_ptr((value), CURLU)) \ + Wcurl_easy_setopt_err_curlu(); \ + if((option) == CURLOPT_STDERR) \ + if(!curlcheck_FILE(value)) \ + Wcurl_easy_setopt_err_FILE(); \ + if(curlcheck_postfields_option(option)) \ + if(!curlcheck_postfields(value)) \ + Wcurl_easy_setopt_err_postfields(); \ + if((option) == CURLOPT_HTTPPOST) \ + if(!curlcheck_arr((value), struct curl_httppost)) \ + Wcurl_easy_setopt_err_curl_httpost(); \ + if((option) == CURLOPT_MIMEPOST) \ + if(!curlcheck_ptr((value), curl_mime)) \ + Wcurl_easy_setopt_err_curl_mimepost(); \ + if(curlcheck_slist_option(option)) \ + if(!curlcheck_arr((value), struct curl_slist)) \ + Wcurl_easy_setopt_err_curl_slist(); \ + if((option) == CURLOPT_SHARE) \ + if(!curlcheck_ptr((value), CURLSH)) \ + Wcurl_easy_setopt_err_CURLSH(); \ + ) \ + } \ + curl_easy_setopt(handle, option, value); \ + }) + +/* wraps curl_easy_getinfo() with type checking */ +#define curl_easy_getinfo(handle, info, arg) \ + __extension__({ \ + if(__builtin_constant_p(info)) { \ + CURL_IGNORE_DEPRECATION( \ + if(curlcheck_string_info(info)) \ + if(!curlcheck_arr((arg), char *)) \ + Wcurl_easy_getinfo_err_string(); \ + if(curlcheck_long_info(info)) \ + if(!curlcheck_arr((arg), long)) \ + Wcurl_easy_getinfo_err_long(); \ + if(curlcheck_double_info(info)) \ + if(!curlcheck_arr((arg), double)) \ + Wcurl_easy_getinfo_err_double(); \ + if(curlcheck_slist_info(info)) \ + if(!curlcheck_arr((arg), struct curl_slist *)) \ + Wcurl_easy_getinfo_err_curl_slist(); \ + if(curlcheck_tlssessioninfo_info(info)) \ + if(!curlcheck_arr((arg), struct curl_tlssessioninfo *)) \ + Wcurl_easy_getinfo_err_curl_tlssessioninfo(); \ + if(curlcheck_certinfo_info(info)) \ + if(!curlcheck_arr((arg), struct curl_certinfo *)) \ + Wcurl_easy_getinfo_err_curl_certinfo(); \ + if(curlcheck_socket_info(info)) \ + if(!curlcheck_arr((arg), curl_socket_t)) \ + Wcurl_easy_getinfo_err_curl_socket(); \ + if(curlcheck_off_t_info(info)) \ + if(!curlcheck_arr((arg), curl_off_t)) \ + Wcurl_easy_getinfo_err_curl_off_t(); \ + ) \ + } \ + curl_easy_getinfo(handle, info, arg); \ + }) + +#define curl_multi_setopt(handle, option, value) \ + __extension__({ \ + if(__builtin_constant_p(option)) { \ + if(curlcheck_long_option(option)) \ + if(!curlcheck_long(value)) \ + Wcurl_multi_setopt_err_long(); \ + if(curlcheck_off_t_option(option)) \ + if(!curlcheck_off_t(value)) \ + Wcurl_multi_setopt_err_curl_off_t(); \ + if(curlcheck_multicb_data_option(option)) \ + if(!curlcheck_cb_data(value)) \ + Wcurl_multi_setopt_err_cb_data(); \ + if(curlcheck_charpp_option(option)) \ + if(!curlcheck_ptrptr(value, char)) \ + Wcurl_multi_setopt_err_charpp(); \ + if((option) == CURLMOPT_NOTIFYFUNCTION) \ + if(!curlcheck_multinotify_cb(value)) \ + Wcurl_multi_setopt_err_notifycb(); \ + if((option) == CURLMOPT_PUSHFUNCTION) \ + if(!curlcheck_multipush_cb(value)) \ + Wcurl_multi_setopt_err_pushcb(); \ + if((option) == CURLMOPT_SOCKETFUNCTION) \ + if(!curlcheck_multisocket_cb(value)) \ + Wcurl_multi_setopt_err_socketcb(); \ + if((option) == CURLMOPT_TIMERFUNCTION) \ + if(!curlcheck_multitimer_cb(value)) \ + Wcurl_multi_setopt_err_timercb(); \ + } \ + curl_multi_setopt(handle, option, value); \ + }) + +/* evaluates to true if the option takes a data argument to pass to a + callback */ +#define curlcheck_multicb_data_option(option) \ + ((option) == CURLMOPT_NOTIFYDATA || \ + (option) == CURLMOPT_PUSHDATA || \ + (option) == CURLMOPT_SOCKETDATA || \ + (option) == CURLMOPT_TIMERDATA || \ + 0) + +/* evaluates to true if the option takes a char ** argument */ +#define curlcheck_charpp_option(option) \ + ((option) == CURLMOPT_PIPELINING_SERVER_BL || \ + (option) == CURLMOPT_PIPELINING_SITE_BL || \ + 0) + +/* evaluates to true if expr is of type curl_multi_timer_callback */ +#define curlcheck_multitimer_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_multi_timer_callback)) + +/* evaluates to true if expr is of type curl_socket_callback */ +#define curlcheck_multisocket_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_socket_callback)) + +/* evaluates to true if expr is of type curl_push_callback */ +#define curlcheck_multipush_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_push_callback)) + +/* evaluates to true if expr is of type curl_push_callback */ +#define curlcheck_multinotify_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_notify_callback)) + +/* + * For now, just make sure that the functions are called with three arguments + */ +#define curl_share_setopt(share,opt,param) curl_share_setopt(share,opt,param) + +/* the actual warnings, triggered by calling the Wcurl_easy_setopt_err* + * functions */ + +/* To define a new warning, use _CURL_WARNING(identifier, "message") */ +#define CURLWARNING(id, message) \ + static void __attribute__((__warning__(message))) \ + __attribute__((__unused__)) __attribute__((__noinline__)) \ + id(void) { __asm__(""); } + +CURLWARNING(Wcurl_multi_setopt_err_long, + "curl_multi_setopt expects a long argument") +CURLWARNING(Wcurl_multi_setopt_err_curl_off_t, + "curl_multi_setopt expects a curl_off_t argument") +CURLWARNING(Wcurl_multi_setopt_err_cb_data, + "curl_multi_setopt expects a 'void *' argument") +CURLWARNING(Wcurl_multi_setopt_err_charpp, + "curl_multi_setopt expects a 'char **' argument") +CURLWARNING(Wcurl_multi_setopt_err_pushcb, + "curl_multi_setopt expects a curl_push_callback argument") +CURLWARNING(Wcurl_multi_setopt_err_notifycb, + "curl_multi_setopt expects a curl_notify_callback argument") +CURLWARNING(Wcurl_multi_setopt_err_socketcb, + "curl_multi_setopt expects a curl_socket_callback argument") +CURLWARNING(Wcurl_multi_setopt_err_timercb, + "curl_multi_setopt expects a curl_multi_timer_callback argument") + +CURLWARNING(Wcurl_easy_setopt_err_long, + "curl_easy_setopt expects a long argument") +CURLWARNING(Wcurl_easy_setopt_err_curl_off_t, + "curl_easy_setopt expects a curl_off_t argument") +CURLWARNING(Wcurl_easy_setopt_err_string, + "curl_easy_setopt expects a " + "string ('char *' or char[]) argument") +CURLWARNING(Wcurl_easy_setopt_err_write_callback, + "curl_easy_setopt expects a curl_write_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_resolver_start_callback, + "curl_easy_setopt expects a " + "curl_resolver_start_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_read_cb, + "curl_easy_setopt expects a curl_read_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_ioctl_cb, + "curl_easy_setopt expects a curl_ioctl_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_sockopt_cb, + "curl_easy_setopt expects a curl_sockopt_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_opensocket_cb, + "curl_easy_setopt expects a " + "curl_opensocket_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_progress_cb, + "curl_easy_setopt expects a curl_progress_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_xferinfo_cb, + "curl_easy_setopt expects a curl_xferinfo_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_debug_cb, + "curl_easy_setopt expects a curl_debug_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_ssl_ctx_cb, + "curl_easy_setopt expects a curl_ssl_ctx_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_conv_cb, + "curl_easy_setopt expects a curl_conv_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_seek_cb, + "curl_easy_setopt expects a curl_seek_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_cb_data, + "curl_easy_setopt expects a " + "private data pointer as argument") +CURLWARNING(Wcurl_easy_setopt_err_chunk_bgn_cb, + "curl_easy_setopt expects a curl_chunk_bgn_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_chunk_end_cb, + "curl_easy_setopt expects a curl_chunk_end_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_close_socket_cb, + "curl_easy_setopt expects a curl_closesocket_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_fnmatch_cb, + "curl_easy_setopt expects a curl_fnmatch_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_hstsread_cb, + "curl_easy_setopt expects a curl_hstsread_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_hstswrite_cb, + "curl_easy_setopt expects a curl_hstswrite_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_ssh_key_cb, + "curl_easy_setopt expects a curl_sshkeycallback argument") +CURLWARNING(Wcurl_easy_setopt_err_ssh_hostkey_cb, + "curl_easy_setopt expects a curl_sshhostkeycallback argument") +CURLWARNING(Wcurl_easy_setopt_err_interleave_cb, + "curl_easy_setopt expects a curl_interleave_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_prereq_cb, + "curl_easy_setopt expects a curl_prereq_callback argument") +CURLWARNING(Wcurl_easy_setopt_err_trailer_cb, + "curl_easy_setopt expects a curl_trailerfunc_ok argument") +CURLWARNING(Wcurl_easy_setopt_err_error_buffer, + "curl_easy_setopt expects a " + "char buffer of CURL_ERROR_SIZE as argument") +CURLWARNING(Wcurl_easy_setopt_err_curlu, + "curl_easy_setopt expects a 'CURLU *' argument") +CURLWARNING(Wcurl_easy_setopt_err_curl, + "curl_easy_setopt expects a 'CURL *' argument") +CURLWARNING(Wcurl_easy_setopt_err_FILE, + "curl_easy_setopt expects a 'FILE *' argument") +CURLWARNING(Wcurl_easy_setopt_err_postfields, + "curl_easy_setopt expects a 'void *' or 'char *' argument") +CURLWARNING(Wcurl_easy_setopt_err_curl_httpost, + "curl_easy_setopt expects a 'struct curl_httppost *' " + "argument") +CURLWARNING(Wcurl_easy_setopt_err_curl_mimepost, + "curl_easy_setopt expects a 'curl_mime *' " + "argument") +CURLWARNING(Wcurl_easy_setopt_err_curl_slist, + "curl_easy_setopt expects a 'struct curl_slist *' argument") +CURLWARNING(Wcurl_easy_setopt_err_CURLSH, + "curl_easy_setopt expects a CURLSH* argument") +CURLWARNING(Wcurl_easy_getinfo_err_string, + "curl_easy_getinfo expects a pointer to 'char *'") +CURLWARNING(Wcurl_easy_getinfo_err_long, + "curl_easy_getinfo expects a pointer to long") +CURLWARNING(Wcurl_easy_getinfo_err_double, + "curl_easy_getinfo expects a pointer to double") +CURLWARNING(Wcurl_easy_getinfo_err_curl_slist, + "curl_easy_getinfo expects a pointer to 'struct curl_slist *'") +CURLWARNING(Wcurl_easy_getinfo_err_curl_tlssessioninfo, + "curl_easy_getinfo expects a pointer to " + "'struct curl_tlssessioninfo *'") +CURLWARNING(Wcurl_easy_getinfo_err_curl_certinfo, + "curl_easy_getinfo expects a pointer to " + "'struct curl_certinfo *'") +CURLWARNING(Wcurl_easy_getinfo_err_curl_socket, + "curl_easy_getinfo expects a pointer to curl_socket_t") +CURLWARNING(Wcurl_easy_getinfo_err_curl_off_t, + "curl_easy_getinfo expects a pointer to curl_off_t") + +/* groups of curl_easy_setops options that take the same type of argument */ + +/* evaluates to true if option takes a long argument */ +#define curlcheck_long_option(option) \ + (0 < (option) && (option) < CURLOPTTYPE_OBJECTPOINT) + +#define curlcheck_off_t_option(option) \ + (((option) > CURLOPTTYPE_OFF_T) && ((option) < CURLOPTTYPE_BLOB)) + +/* option takes a CURL * argument */ +#define curlcheck_curl_option(option) \ + ((option) == CURLOPT_STREAM_DEPENDS || \ + (option) == CURLOPT_STREAM_DEPENDS_E || \ + 0) + +/* evaluates to true if option takes a char* argument */ +#define curlcheck_string_option(option) \ + ((option) == CURLOPT_ABSTRACT_UNIX_SOCKET || \ + (option) == CURLOPT_ACCEPT_ENCODING || \ + (option) == CURLOPT_ALTSVC || \ + (option) == CURLOPT_CAINFO || \ + (option) == CURLOPT_CAPATH || \ + (option) == CURLOPT_COOKIE || \ + (option) == CURLOPT_COOKIEFILE || \ + (option) == CURLOPT_COOKIEJAR || \ + (option) == CURLOPT_COOKIELIST || \ + (option) == CURLOPT_CRLFILE || \ + (option) == CURLOPT_CUSTOMREQUEST || \ + (option) == CURLOPT_DEFAULT_PROTOCOL || \ + (option) == CURLOPT_DNS_INTERFACE || \ + (option) == CURLOPT_DNS_LOCAL_IP4 || \ + (option) == CURLOPT_DNS_LOCAL_IP6 || \ + (option) == CURLOPT_DNS_SERVERS || \ + (option) == CURLOPT_DOH_URL || \ + (option) == CURLOPT_ECH || \ + (option) == CURLOPT_EGDSOCKET || \ + (option) == CURLOPT_FTP_ACCOUNT || \ + (option) == CURLOPT_FTP_ALTERNATIVE_TO_USER || \ + (option) == CURLOPT_FTPPORT || \ + (option) == CURLOPT_HAPROXY_CLIENT_IP || \ + (option) == CURLOPT_HSTS || \ + (option) == CURLOPT_INTERFACE || \ + (option) == CURLOPT_ISSUERCERT || \ + (option) == CURLOPT_KEYPASSWD || \ + (option) == CURLOPT_KRBLEVEL || \ + (option) == CURLOPT_LOGIN_OPTIONS || \ + (option) == CURLOPT_MAIL_AUTH || \ + (option) == CURLOPT_MAIL_FROM || \ + (option) == CURLOPT_NETRC_FILE || \ + (option) == CURLOPT_NOPROXY || \ + (option) == CURLOPT_PASSWORD || \ + (option) == CURLOPT_PINNEDPUBLICKEY || \ + (option) == CURLOPT_PRE_PROXY || \ + (option) == CURLOPT_PROTOCOLS_STR || \ + (option) == CURLOPT_PROXY || \ + (option) == CURLOPT_PROXY_CAINFO || \ + (option) == CURLOPT_PROXY_CAPATH || \ + (option) == CURLOPT_PROXY_CRLFILE || \ + (option) == CURLOPT_PROXY_ISSUERCERT || \ + (option) == CURLOPT_PROXY_KEYPASSWD || \ + (option) == CURLOPT_PROXY_PINNEDPUBLICKEY || \ + (option) == CURLOPT_PROXY_SERVICE_NAME || \ + (option) == CURLOPT_PROXY_SSL_CIPHER_LIST || \ + (option) == CURLOPT_PROXY_SSLCERT || \ + (option) == CURLOPT_PROXY_SSLCERTTYPE || \ + (option) == CURLOPT_PROXY_SSLKEY || \ + (option) == CURLOPT_PROXY_SSLKEYTYPE || \ + (option) == CURLOPT_PROXY_TLS13_CIPHERS || \ + (option) == CURLOPT_PROXY_TLSAUTH_PASSWORD || \ + (option) == CURLOPT_PROXY_TLSAUTH_TYPE || \ + (option) == CURLOPT_PROXY_TLSAUTH_USERNAME || \ + (option) == CURLOPT_PROXYPASSWORD || \ + (option) == CURLOPT_PROXYUSERNAME || \ + (option) == CURLOPT_PROXYUSERPWD || \ + (option) == CURLOPT_RANDOM_FILE || \ + (option) == CURLOPT_RANGE || \ + (option) == CURLOPT_REDIR_PROTOCOLS_STR || \ + (option) == CURLOPT_REFERER || \ + (option) == CURLOPT_REQUEST_TARGET || \ + (option) == CURLOPT_RTSP_SESSION_ID || \ + (option) == CURLOPT_RTSP_STREAM_URI || \ + (option) == CURLOPT_RTSP_TRANSPORT || \ + (option) == CURLOPT_SASL_AUTHZID || \ + (option) == CURLOPT_SERVICE_NAME || \ + (option) == CURLOPT_SOCKS5_GSSAPI_SERVICE || \ + (option) == CURLOPT_SSH_HOST_PUBLIC_KEY_MD5 || \ + (option) == CURLOPT_SSH_HOST_PUBLIC_KEY_SHA256 || \ + (option) == CURLOPT_SSH_KNOWNHOSTS || \ + (option) == CURLOPT_SSH_PRIVATE_KEYFILE || \ + (option) == CURLOPT_SSH_PUBLIC_KEYFILE || \ + (option) == CURLOPT_SSLCERT || \ + (option) == CURLOPT_SSLCERTTYPE || \ + (option) == CURLOPT_SSLENGINE || \ + (option) == CURLOPT_SSLKEY || \ + (option) == CURLOPT_SSLKEYTYPE || \ + (option) == CURLOPT_SSL_CIPHER_LIST || \ + (option) == CURLOPT_SSL_EC_CURVES || \ + (option) == CURLOPT_SSL_SIGNATURE_ALGORITHMS || \ + (option) == CURLOPT_TLS13_CIPHERS || \ + (option) == CURLOPT_TLSAUTH_PASSWORD || \ + (option) == CURLOPT_TLSAUTH_TYPE || \ + (option) == CURLOPT_TLSAUTH_USERNAME || \ + (option) == CURLOPT_UNIX_SOCKET_PATH || \ + (option) == CURLOPT_URL || \ + (option) == CURLOPT_USERAGENT || \ + (option) == CURLOPT_USERNAME || \ + (option) == CURLOPT_AWS_SIGV4 || \ + (option) == CURLOPT_USERPWD || \ + (option) == CURLOPT_XOAUTH2_BEARER || \ + 0) + +/* evaluates to true if option takes a curl_write_callback argument */ +#define curlcheck_write_cb_option(option) \ + ((option) == CURLOPT_HEADERFUNCTION || \ + (option) == CURLOPT_WRITEFUNCTION) + +/* evaluates to true if option takes a curl_conv_callback argument */ +#define curlcheck_conv_cb_option(option) \ + ((option) == CURLOPT_CONV_TO_NETWORK_FUNCTION || \ + (option) == CURLOPT_CONV_FROM_NETWORK_FUNCTION || \ + (option) == CURLOPT_CONV_FROM_UTF8_FUNCTION) + +/* evaluates to true if option takes a data argument to pass to a callback */ +#define curlcheck_cb_data_option(option) \ + ((option) == CURLOPT_CHUNK_DATA || \ + (option) == CURLOPT_CLOSESOCKETDATA || \ + (option) == CURLOPT_DEBUGDATA || \ + (option) == CURLOPT_FNMATCH_DATA || \ + (option) == CURLOPT_HEADERDATA || \ + (option) == CURLOPT_HSTSREADDATA || \ + (option) == CURLOPT_HSTSWRITEDATA || \ + (option) == CURLOPT_INTERLEAVEDATA || \ + (option) == CURLOPT_IOCTLDATA || \ + (option) == CURLOPT_OPENSOCKETDATA || \ + (option) == CURLOPT_PREREQDATA || \ + (option) == CURLOPT_XFERINFODATA || \ + (option) == CURLOPT_READDATA || \ + (option) == CURLOPT_SEEKDATA || \ + (option) == CURLOPT_SOCKOPTDATA || \ + (option) == CURLOPT_SSH_KEYDATA || \ + (option) == CURLOPT_SSL_CTX_DATA || \ + (option) == CURLOPT_WRITEDATA || \ + (option) == CURLOPT_RESOLVER_START_DATA || \ + (option) == CURLOPT_TRAILERDATA || \ + (option) == CURLOPT_SSH_HOSTKEYDATA || \ + 0) + +/* evaluates to true if option takes a POST data argument (void* or char*) */ +#define curlcheck_postfields_option(option) \ + ((option) == CURLOPT_POSTFIELDS || \ + (option) == CURLOPT_COPYPOSTFIELDS || \ + 0) + +/* evaluates to true if option takes a struct curl_slist * argument */ +#define curlcheck_slist_option(option) \ + ((option) == CURLOPT_HTTP200ALIASES || \ + (option) == CURLOPT_HTTPHEADER || \ + (option) == CURLOPT_MAIL_RCPT || \ + (option) == CURLOPT_POSTQUOTE || \ + (option) == CURLOPT_PREQUOTE || \ + (option) == CURLOPT_PROXYHEADER || \ + (option) == CURLOPT_QUOTE || \ + (option) == CURLOPT_RESOLVE || \ + (option) == CURLOPT_TELNETOPTIONS || \ + (option) == CURLOPT_CONNECT_TO || \ + 0) + +/* groups of curl_easy_getinfo infos that take the same type of argument */ + +/* evaluates to true if info expects a pointer to char * argument */ +#define curlcheck_string_info(info) \ + (CURLINFO_STRING < (info) && (info) < CURLINFO_LONG && \ + (info) != CURLINFO_PRIVATE) + +/* evaluates to true if info expects a pointer to long argument */ +#define curlcheck_long_info(info) \ + (CURLINFO_LONG < (info) && (info) < CURLINFO_DOUBLE) + +/* evaluates to true if info expects a pointer to double argument */ +#define curlcheck_double_info(info) \ + (CURLINFO_DOUBLE < (info) && (info) < CURLINFO_SLIST) + +/* true if info expects a pointer to struct curl_slist * argument */ +#define curlcheck_slist_info(info) \ + (((info) == CURLINFO_SSL_ENGINES) || ((info) == CURLINFO_COOKIELIST)) + +/* true if info expects a pointer to struct curl_tlssessioninfo * argument */ +#define curlcheck_tlssessioninfo_info(info) \ + (((info) == CURLINFO_TLS_SSL_PTR) || ((info) == CURLINFO_TLS_SESSION)) + +/* true if info expects a pointer to struct curl_certinfo * argument */ +#define curlcheck_certinfo_info(info) ((info) == CURLINFO_CERTINFO) + +/* true if info expects a pointer to struct curl_socket_t argument */ +#define curlcheck_socket_info(info) \ + (CURLINFO_SOCKET < (info) && (info) < CURLINFO_OFF_T) + +/* true if info expects a pointer to curl_off_t argument */ +#define curlcheck_off_t_info(info) \ + (CURLINFO_OFF_T < (info)) + +/* + * typecheck helpers -- check whether given expression has requested type + */ + +/* For pointers, you can use the curlcheck_ptr/curlcheck_arr macros, + * otherwise define a new macro. Search for __builtin_types_compatible_p + * in the GCC manual. + * NOTE: these macros MUST NOT EVALUATE their arguments! The argument is + * the actual expression passed to the curl_easy_setopt macro. This + * means that you can only apply the sizeof and __typeof__ operators, no + * == or whatsoever. + */ + +/* XXX: should evaluate to true if expr is a pointer */ +#define curlcheck_any_ptr(expr) \ + (sizeof(expr) == sizeof(void *)) + +/* evaluates to true if expr is NULL */ +/* XXX: must not evaluate expr, so this check is not accurate */ +#define curlcheck_NULL(expr) \ + (__builtin_types_compatible_p(__typeof__(expr), __typeof__(NULL))) + +/* evaluates to true if expr is type*, const type* or NULL */ +#define curlcheck_ptr(expr, type) \ + (curlcheck_NULL(expr) || \ + __builtin_types_compatible_p(__typeof__(expr), type *) || \ + __builtin_types_compatible_p(__typeof__(expr), const type *)) + +/* evaluates to true if expr is type**, const type** or NULL */ +#define curlcheck_ptrptr(expr, type) \ + (curlcheck_NULL(expr) || \ + __builtin_types_compatible_p(__typeof__(expr), type **) || \ + __builtin_types_compatible_p(__typeof__(expr), type *[]) || \ + __builtin_types_compatible_p(__typeof__(expr), const type *[]) || \ + __builtin_types_compatible_p(__typeof__(expr), const type **)) + +/* evaluates to true if expr is one of type[], type*, NULL or const type* */ +#define curlcheck_arr(expr, type) \ + (curlcheck_ptr((expr), type) || \ + __builtin_types_compatible_p(__typeof__(expr), type [])) + +/* evaluates to true if expr is a string */ +#define curlcheck_string(expr) \ + (curlcheck_arr((expr), char) || \ + curlcheck_arr((expr), signed char) || \ + curlcheck_arr((expr), unsigned char)) + +/* evaluates to true if expr is a CURL * */ +#define curlcheck_curl(expr) \ + (curlcheck_NULL(expr) || \ + __builtin_types_compatible_p(__typeof__(expr), CURL *)) + +/* evaluates to true if expr is a long (no matter the signedness) + * XXX: for now, int is also accepted (and therefore short and char, which + * are promoted to int when passed to a variadic function) */ +#define curlcheck_long(expr) \ + ( \ + ((sizeof(long) != sizeof(int)) && \ + (__builtin_types_compatible_p(__typeof__(expr), long) || \ + __builtin_types_compatible_p(__typeof__(expr), signed long) || \ + __builtin_types_compatible_p(__typeof__(expr), unsigned long))) \ + || \ + ((sizeof(long) == sizeof(int)) && \ + (__builtin_types_compatible_p(__typeof__(expr), long) || \ + __builtin_types_compatible_p(__typeof__(expr), signed long) || \ + __builtin_types_compatible_p(__typeof__(expr), unsigned long) || \ + __builtin_types_compatible_p(__typeof__(expr), int) || \ + __builtin_types_compatible_p(__typeof__(expr), signed int) || \ + __builtin_types_compatible_p(__typeof__(expr), unsigned int) || \ + __builtin_types_compatible_p(__typeof__(expr), short) || \ + __builtin_types_compatible_p(__typeof__(expr), signed short) || \ + __builtin_types_compatible_p(__typeof__(expr), unsigned short) || \ + __builtin_types_compatible_p(__typeof__(expr), char) || \ + __builtin_types_compatible_p(__typeof__(expr), signed char) || \ + __builtin_types_compatible_p(__typeof__(expr), unsigned char))) \ + ) + +/* evaluates to true if expr is of type curl_off_t */ +#define curlcheck_off_t(expr) \ + (__builtin_types_compatible_p(__typeof__(expr), curl_off_t)) + +/* evaluates to true if expr is abuffer suitable for CURLOPT_ERRORBUFFER */ +/* XXX: also check size of an char[] array? */ +#define curlcheck_error_buffer(expr) \ + (curlcheck_NULL(expr) || \ + __builtin_types_compatible_p(__typeof__(expr), char *) || \ + __builtin_types_compatible_p(__typeof__(expr), char[])) + +/* evaluates to true if expr is of type (const) void* or (const) FILE* */ +#if 0 +#define curlcheck_cb_data(expr) \ + (curlcheck_ptr((expr), void) || \ + curlcheck_ptr((expr), FILE)) +#else /* be less strict */ +#define curlcheck_cb_data(expr) \ + curlcheck_any_ptr(expr) +#endif + +/* evaluates to true if expr is of type FILE* */ +#define curlcheck_FILE(expr) \ + (curlcheck_NULL(expr) || \ + (__builtin_types_compatible_p(__typeof__(expr), FILE *))) + +/* evaluates to true if expr can be passed as POST data (void* or char*) */ +#define curlcheck_postfields(expr) \ + (curlcheck_ptr((expr), void) || \ + curlcheck_arr((expr), char) || \ + curlcheck_arr((expr), unsigned char)) + +/* helper: __builtin_types_compatible_p distinguishes between functions and + * function pointers, hide it */ +#define curlcheck_cb_compatible(func, type) \ + (__builtin_types_compatible_p(__typeof__(func), type) || \ + __builtin_types_compatible_p(__typeof__(func) *, type)) + +/* evaluates to true if expr is of type curl_resolver_start_callback */ +#define curlcheck_resolver_start_callback(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_resolver_start_callback)) + +/* evaluates to true if expr is of type curl_read_callback or "similar" */ +#define curlcheck_read_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), __typeof__(fread) *) || \ + curlcheck_cb_compatible((expr), curl_read_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_read_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_read_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_read_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_read_callback4) || \ + curlcheck_cb_compatible((expr), Wcurl_read_callback5) || \ + curlcheck_cb_compatible((expr), Wcurl_read_callback6)) +typedef size_t (*Wcurl_read_callback1)(char *, size_t, size_t, void *); +typedef size_t (*Wcurl_read_callback2)(char *, size_t, size_t, const void *); +typedef size_t (*Wcurl_read_callback3)(char *, size_t, size_t, FILE *); +typedef size_t (*Wcurl_read_callback4)(void *, size_t, size_t, void *); +typedef size_t (*Wcurl_read_callback5)(void *, size_t, size_t, const void *); +typedef size_t (*Wcurl_read_callback6)(void *, size_t, size_t, FILE *); + +/* evaluates to true if expr is of type curl_write_callback or "similar" */ +#define curlcheck_write_cb(expr) \ + (curlcheck_read_cb(expr) || \ + curlcheck_cb_compatible((expr), __typeof__(fwrite) *) || \ + curlcheck_cb_compatible((expr), curl_write_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_write_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_write_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_write_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_write_callback4) || \ + curlcheck_cb_compatible((expr), Wcurl_write_callback5) || \ + curlcheck_cb_compatible((expr), Wcurl_write_callback6)) +typedef size_t (*Wcurl_write_callback1)(const char *, size_t, size_t, void *); +typedef size_t (*Wcurl_write_callback2)(const char *, size_t, size_t, + const void *); +typedef size_t (*Wcurl_write_callback3)(const char *, size_t, size_t, FILE *); +typedef size_t (*Wcurl_write_callback4)(const void *, size_t, size_t, void *); +typedef size_t (*Wcurl_write_callback5)(const void *, size_t, size_t, + const void *); +typedef size_t (*Wcurl_write_callback6)(const void *, size_t, size_t, FILE *); + +/* evaluates to true if expr is of type curl_ioctl_callback or "similar" */ +#define curlcheck_ioctl_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_ioctl_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_ioctl_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_ioctl_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_ioctl_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_ioctl_callback4)) +typedef curlioerr (*Wcurl_ioctl_callback1)(CURL *, int, void *); +typedef curlioerr (*Wcurl_ioctl_callback2)(CURL *, int, const void *); +typedef curlioerr (*Wcurl_ioctl_callback3)(CURL *, curliocmd, void *); +typedef curlioerr (*Wcurl_ioctl_callback4)(CURL *, curliocmd, const void *); + +/* evaluates to true if expr is of type curl_sockopt_callback or "similar" */ +#define curlcheck_sockopt_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_sockopt_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_sockopt_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_sockopt_callback2)) +typedef int (*Wcurl_sockopt_callback1)(void *, curl_socket_t, curlsocktype); +typedef int (*Wcurl_sockopt_callback2)(const void *, curl_socket_t, + curlsocktype); + +/* evaluates to true if expr is of type curl_opensocket_callback or + "similar" */ +#define curlcheck_opensocket_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_opensocket_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_opensocket_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_opensocket_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_opensocket_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_opensocket_callback4)) +typedef curl_socket_t (*Wcurl_opensocket_callback1) + (void *, curlsocktype, struct curl_sockaddr *); +typedef curl_socket_t (*Wcurl_opensocket_callback2) + (void *, curlsocktype, const struct curl_sockaddr *); +typedef curl_socket_t (*Wcurl_opensocket_callback3) + (const void *, curlsocktype, struct curl_sockaddr *); +typedef curl_socket_t (*Wcurl_opensocket_callback4) + (const void *, curlsocktype, const struct curl_sockaddr *); + +/* evaluates to true if expr is of type curl_progress_callback or "similar" */ +#define curlcheck_progress_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_progress_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_progress_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_progress_callback2)) +typedef int (*Wcurl_progress_callback1)(void *, + double, double, double, double); +typedef int (*Wcurl_progress_callback2)(const void *, + double, double, double, double); + +/* evaluates to true if expr is of type curl_xferinfo_callback */ +#define curlcheck_xferinfo_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_xferinfo_callback)) + +/* evaluates to true if expr is of type curl_debug_callback or "similar" */ +#define curlcheck_debug_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_debug_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback4) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback5) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback6) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback7) || \ + curlcheck_cb_compatible((expr), Wcurl_debug_callback8)) +typedef int (*Wcurl_debug_callback1) (CURL *, + curl_infotype, char *, size_t, void *); +typedef int (*Wcurl_debug_callback2) (CURL *, + curl_infotype, char *, size_t, const void *); +typedef int (*Wcurl_debug_callback3) (CURL *, + curl_infotype, const char *, size_t, void *); +typedef int (*Wcurl_debug_callback4) (CURL *, + curl_infotype, const char *, size_t, const void *); +typedef int (*Wcurl_debug_callback5) (CURL *, + curl_infotype, unsigned char *, size_t, void *); +typedef int (*Wcurl_debug_callback6) (CURL *, + curl_infotype, unsigned char *, size_t, const void *); +typedef int (*Wcurl_debug_callback7) (CURL *, + curl_infotype, const unsigned char *, size_t, void *); +typedef int (*Wcurl_debug_callback8) (CURL *, + curl_infotype, const unsigned char *, size_t, const void *); + +/* evaluates to true if expr is of type curl_ssl_ctx_callback or "similar" */ +/* this is getting even messier... */ +#define curlcheck_ssl_ctx_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_ssl_ctx_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback4) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback5) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback6) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback7) || \ + curlcheck_cb_compatible((expr), Wcurl_ssl_ctx_callback8)) +typedef CURLcode (*Wcurl_ssl_ctx_callback1)(CURL *, void *, void *); +typedef CURLcode (*Wcurl_ssl_ctx_callback2)(CURL *, void *, const void *); +typedef CURLcode (*Wcurl_ssl_ctx_callback3)(CURL *, const void *, void *); +typedef CURLcode (*Wcurl_ssl_ctx_callback4)(CURL *, const void *, + const void *); +#ifdef HEADER_SSL_H +/* hack: if we included OpenSSL's ssl.h, we know about SSL_CTX + * this will of course break if we are included before OpenSSL headers... + */ +typedef CURLcode (*Wcurl_ssl_ctx_callback5)(CURL *, SSL_CTX *, void *); +typedef CURLcode (*Wcurl_ssl_ctx_callback6)(CURL *, SSL_CTX *, const void *); +typedef CURLcode (*Wcurl_ssl_ctx_callback7)(CURL *, const SSL_CTX *, void *); +typedef CURLcode (*Wcurl_ssl_ctx_callback8)(CURL *, const SSL_CTX *, + const void *); +#else +typedef Wcurl_ssl_ctx_callback1 Wcurl_ssl_ctx_callback5; +typedef Wcurl_ssl_ctx_callback1 Wcurl_ssl_ctx_callback6; +typedef Wcurl_ssl_ctx_callback1 Wcurl_ssl_ctx_callback7; +typedef Wcurl_ssl_ctx_callback1 Wcurl_ssl_ctx_callback8; +#endif + +/* evaluates to true if expr is of type curl_conv_callback or "similar" */ +#define curlcheck_conv_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_conv_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_conv_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_conv_callback2) || \ + curlcheck_cb_compatible((expr), Wcurl_conv_callback3) || \ + curlcheck_cb_compatible((expr), Wcurl_conv_callback4)) +typedef CURLcode (*Wcurl_conv_callback1)(char *, size_t length); +typedef CURLcode (*Wcurl_conv_callback2)(const char *, size_t length); +typedef CURLcode (*Wcurl_conv_callback3)(void *, size_t length); +typedef CURLcode (*Wcurl_conv_callback4)(const void *, size_t length); + +/* evaluates to true if expr is of type curl_seek_callback or "similar" */ +#define curlcheck_seek_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_seek_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_seek_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_seek_callback2)) +typedef CURLcode (*Wcurl_seek_callback1)(void *, curl_off_t, int); +typedef CURLcode (*Wcurl_seek_callback2)(const void *, curl_off_t, int); + +/* evaluates to true if expr is of type curl_chunk_bgn_callback */ +#define curlcheck_chunk_bgn_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_chunk_bgn_callback) || \ + curlcheck_cb_compatible((expr), Wcurl_chunk_bgn_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_chunk_bgn_callback2)) +typedef long (*Wcurl_chunk_bgn_callback1)(struct curl_fileinfo *, + void *, int); +typedef long (*Wcurl_chunk_bgn_callback2)(void *, void *, int); + +/* evaluates to true if expr is of type curl_chunk_end_callback */ +#define curlcheck_chunk_end_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_chunk_end_callback)) + +/* evaluates to true if expr is of type curl_closesocket_callback */ +#define curlcheck_close_socket_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_closesocket_callback)) + +/* evaluates to true if expr is of type curl_fnmatch_callback */ +#define curlcheck_fnmatch_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_fnmatch_callback)) + +/* evaluates to true if expr is of type curl_hstsread_callback */ +#define curlcheck_hstsread_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_hstsread_callback)) + +/* evaluates to true if expr is of type curl_hstswrite_callback */ +#define curlcheck_hstswrite_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_hstswrite_callback)) + +/* evaluates to true if expr is of type curl_sshhostkeycallback */ +#define curlcheck_ssh_hostkey_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_sshhostkeycallback)) + +/* evaluates to true if expr is of type curl_sshkeycallback */ +#define curlcheck_ssh_key_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_sshkeycallback)) + +/* evaluates to true if expr is of type curl_interleave_callback */ +#define curlcheck_interleave_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), Wcurl_interleave_callback1) || \ + curlcheck_cb_compatible((expr), Wcurl_interleave_callback2)) +typedef size_t (*Wcurl_interleave_callback1)(void *p, size_t s, + size_t n, void *u); +typedef size_t (*Wcurl_interleave_callback2)(char *p, size_t s, + size_t n, void *u); + +/* evaluates to true if expr is of type curl_prereq_callback */ +#define curlcheck_prereq_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_prereq_callback)) + +/* evaluates to true if expr is of type curl_trailer_callback */ +#define curlcheck_trailer_cb(expr) \ + (curlcheck_NULL(expr) || \ + curlcheck_cb_compatible((expr), curl_trailer_callback)) + +#endif /* CURLINC_TYPECHECK_GCC_H */ diff --git a/cactus-engine/libs/curl/include/curl/urlapi.h b/cactus-engine/libs/curl/include/curl/urlapi.h new file mode 100644 index 000000000..dabe44a47 --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/urlapi.h @@ -0,0 +1,155 @@ +#ifndef CURLINC_URLAPI_H +#define CURLINC_URLAPI_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +#include "curl.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* the error codes for the URL API */ +typedef enum { + CURLUE_OK, + CURLUE_BAD_HANDLE, /* 1 */ + CURLUE_BAD_PARTPOINTER, /* 2 */ + CURLUE_MALFORMED_INPUT, /* 3 */ + CURLUE_BAD_PORT_NUMBER, /* 4 */ + CURLUE_UNSUPPORTED_SCHEME, /* 5 */ + CURLUE_URLDECODE, /* 6 */ + CURLUE_OUT_OF_MEMORY, /* 7 */ + CURLUE_USER_NOT_ALLOWED, /* 8 */ + CURLUE_UNKNOWN_PART, /* 9 */ + CURLUE_NO_SCHEME, /* 10 */ + CURLUE_NO_USER, /* 11 */ + CURLUE_NO_PASSWORD, /* 12 */ + CURLUE_NO_OPTIONS, /* 13 */ + CURLUE_NO_HOST, /* 14 */ + CURLUE_NO_PORT, /* 15 */ + CURLUE_NO_QUERY, /* 16 */ + CURLUE_NO_FRAGMENT, /* 17 */ + CURLUE_NO_ZONEID, /* 18 */ + CURLUE_BAD_FILE_URL, /* 19 */ + CURLUE_BAD_FRAGMENT, /* 20 */ + CURLUE_BAD_HOSTNAME, /* 21 */ + CURLUE_BAD_IPV6, /* 22 */ + CURLUE_BAD_LOGIN, /* 23 */ + CURLUE_BAD_PASSWORD, /* 24 */ + CURLUE_BAD_PATH, /* 25 */ + CURLUE_BAD_QUERY, /* 26 */ + CURLUE_BAD_SCHEME, /* 27 */ + CURLUE_BAD_SLASHES, /* 28 */ + CURLUE_BAD_USER, /* 29 */ + CURLUE_LACKS_IDN, /* 30 */ + CURLUE_TOO_LARGE, /* 31 */ + CURLUE_LAST +} CURLUcode; + +typedef enum { + CURLUPART_URL, + CURLUPART_SCHEME, + CURLUPART_USER, + CURLUPART_PASSWORD, + CURLUPART_OPTIONS, + CURLUPART_HOST, + CURLUPART_PORT, + CURLUPART_PATH, + CURLUPART_QUERY, + CURLUPART_FRAGMENT, + CURLUPART_ZONEID /* added in 7.65.0 */ +} CURLUPart; + +#define CURLU_DEFAULT_PORT (1 << 0) /* return default port number */ +#define CURLU_NO_DEFAULT_PORT (1 << 1) /* act as if no port number was set, + if the port number matches the + default for the scheme */ +#define CURLU_DEFAULT_SCHEME (1 << 2) /* return default scheme if + missing */ +#define CURLU_NON_SUPPORT_SCHEME (1 << 3) /* allow non-supported scheme */ +#define CURLU_PATH_AS_IS (1 << 4) /* leave dot sequences */ +#define CURLU_DISALLOW_USER (1 << 5) /* no user+password allowed */ +#define CURLU_URLDECODE (1 << 6) /* URL decode on get */ +#define CURLU_URLENCODE (1 << 7) /* URL encode on set */ +#define CURLU_APPENDQUERY (1 << 8) /* append a form style part */ +#define CURLU_GUESS_SCHEME (1 << 9) /* legacy curl-style guessing */ +#define CURLU_NO_AUTHORITY (1 << 10) /* Allow empty authority when the + scheme is unknown. */ +#define CURLU_ALLOW_SPACE (1 << 11) /* Allow spaces in the URL */ +#define CURLU_PUNYCODE (1 << 12) /* get the hostname in punycode */ +#define CURLU_PUNY2IDN (1 << 13) /* punycode => IDN conversion */ +#define CURLU_GET_EMPTY (1 << 14) /* allow empty queries and fragments + when extracting the URL or the + components */ +#define CURLU_NO_GUESS_SCHEME (1 << 15) /* for get, do not accept a guess */ + +typedef struct Curl_URL CURLU; + +/* + * curl_url() creates a new CURLU handle and returns a pointer to it. + * Must be freed with curl_url_cleanup(). + */ +CURL_EXTERN CURLU *curl_url(void); + +/* + * curl_url_cleanup() frees the CURLU handle and related resources used for + * the URL parsing. It will not free strings previously returned with the URL + * API. + */ +CURL_EXTERN void curl_url_cleanup(CURLU *handle); + +/* + * curl_url_dup() duplicates a CURLU handle and returns a new copy. The new + * handle must also be freed with curl_url_cleanup(). + */ +CURL_EXTERN CURLU *curl_url_dup(const CURLU *in); + +/* + * curl_url_get() extracts a specific part of the URL from a CURLU + * handle. Returns error code. The returned pointer MUST be freed with + * curl_free() afterwards. + */ +CURL_EXTERN CURLUcode curl_url_get(const CURLU *handle, CURLUPart what, + char **part, unsigned int flags); + +/* + * curl_url_set() sets a specific part of the URL in a CURLU handle. Returns + * error code. The passed in string will be copied. Passing a NULL instead of + * a part string, clears that part. + */ +CURL_EXTERN CURLUcode curl_url_set(CURLU *handle, CURLUPart what, + const char *part, unsigned int flags); + +/* + * curl_url_strerror() turns a CURLUcode value into the equivalent human + * readable error string. This is useful for printing meaningful error + * messages. + */ +CURL_EXTERN const char *curl_url_strerror(CURLUcode); + +#ifdef __cplusplus +} /* end of extern "C" */ +#endif + +#endif /* CURLINC_URLAPI_H */ diff --git a/cactus-engine/libs/curl/include/curl/websockets.h b/cactus-engine/libs/curl/include/curl/websockets.h new file mode 100644 index 000000000..402a2ca0f --- /dev/null +++ b/cactus-engine/libs/curl/include/curl/websockets.h @@ -0,0 +1,98 @@ +#ifndef CURLINC_WEBSOCKETS_H +#define CURLINC_WEBSOCKETS_H +/*************************************************************************** + * _ _ ____ _ + * Project ___| | | | _ \| | + * / __| | | | |_) | | + * | (__| |_| | _ <| |___ + * \___|\___/|_| \_\_____| + * + * Copyright (C) Daniel Stenberg, , et al. + * + * This software is licensed as described in the file COPYING, which + * you should have received as part of this distribution. The terms + * are also available at https://curl.se/docs/copyright.html. + * + * You may opt to use, copy, modify, merge, publish, distribute and/or sell + * copies of the Software, and permit persons to whom the Software is + * furnished to do so, under the terms of the COPYING file. + * + * This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY + * KIND, either express or implied. + * + * SPDX-License-Identifier: curl + * + ***************************************************************************/ + +#ifdef __cplusplus +extern "C" { +#endif + +struct curl_ws_frame { + int age; /* zero */ + int flags; /* See the CURLWS_* defines */ + curl_off_t offset; /* the offset of this data into the frame */ + curl_off_t bytesleft; /* number of pending bytes left of the payload */ + size_t len; /* size of the current data chunk */ +}; + +/* flag bits */ +#define CURLWS_TEXT (1 << 0) +#define CURLWS_BINARY (1 << 1) +#define CURLWS_CONT (1 << 2) +#define CURLWS_CLOSE (1 << 3) +#define CURLWS_PING (1 << 4) +#define CURLWS_OFFSET (1 << 5) + +/* + * NAME curl_ws_recv() + * + * DESCRIPTION + * + * Receives data from the websocket connection. Use after successful + * curl_easy_perform() with CURLOPT_CONNECT_ONLY option. + */ +CURL_EXTERN CURLcode curl_ws_recv(CURL *curl, void *buffer, size_t buflen, + size_t *recv, + const struct curl_ws_frame **metap); + +/* flags for curl_ws_send() */ +#define CURLWS_PONG (1 << 6) + +/* + * NAME curl_ws_send() + * + * DESCRIPTION + * + * Sends data over the websocket connection. Use after successful + * curl_easy_perform() with CURLOPT_CONNECT_ONLY option. + */ +CURL_EXTERN CURLcode curl_ws_send(CURL *curl, const void *buffer, + size_t buflen, size_t *sent, + curl_off_t fragsize, + unsigned int flags); + +/* + * NAME curl_ws_start_frame() + * + * DESCRIPTION + * + * Buffers a websocket frame header with the given flags and length. + * Errors when a previous frame is not complete, e.g. not all its + * payload has been added. + */ +CURL_EXTERN CURLcode curl_ws_start_frame(CURL *curl, + unsigned int flags, + curl_off_t frame_len); + +/* bits for the CURLOPT_WS_OPTIONS bitmask: */ +#define CURLWS_RAW_MODE (1L << 0) +#define CURLWS_NOAUTOPONG (1L << 1) + +CURL_EXTERN const struct curl_ws_frame *curl_ws_meta(CURL *curl); + +#ifdef __cplusplus +} +#endif + +#endif /* CURLINC_WEBSOCKETS_H */ diff --git a/cactus-engine/libs/curl/ios/device/libcurl.a b/cactus-engine/libs/curl/ios/device/libcurl.a new file mode 100644 index 000000000..ee93a86b0 Binary files /dev/null and b/cactus-engine/libs/curl/ios/device/libcurl.a differ diff --git a/cactus-engine/libs/curl/ios/simulator/libcurl.a b/cactus-engine/libs/curl/ios/simulator/libcurl.a new file mode 100644 index 000000000..13bde09d7 Binary files /dev/null and b/cactus-engine/libs/curl/ios/simulator/libcurl.a differ diff --git a/cactus-engine/libs/curl/macos/libcurl.a b/cactus-engine/libs/curl/macos/libcurl.a new file mode 100644 index 000000000..0d2583b2e Binary files /dev/null and b/cactus-engine/libs/curl/macos/libcurl.a differ diff --git a/cactus-engine/libs/picojson.h b/cactus-engine/libs/picojson.h new file mode 100644 index 000000000..76742fe06 --- /dev/null +++ b/cactus-engine/libs/picojson.h @@ -0,0 +1,1200 @@ +/* + * Copyright 2009-2010 Cybozu Labs, Inc. + * Copyright 2011-2014 Kazuho Oku + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + */ +#ifndef picojson_h +#define picojson_h + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// for isnan/isinf +#if __cplusplus >= 201103L +#include +#else +extern "C" { +#ifdef _MSC_VER +#include +#elif defined(__INTEL_COMPILER) +#include +#else +#include +#endif +} +#endif + +#ifndef PICOJSON_USE_RVALUE_REFERENCE +#if (defined(__cpp_rvalue_references) && __cpp_rvalue_references >= 200610) || (defined(_MSC_VER) && _MSC_VER >= 1600) +#define PICOJSON_USE_RVALUE_REFERENCE 1 +#else +#define PICOJSON_USE_RVALUE_REFERENCE 0 +#endif +#endif // PICOJSON_USE_RVALUE_REFERENCE + +#ifndef PICOJSON_NOEXCEPT +#if PICOJSON_USE_RVALUE_REFERENCE +#define PICOJSON_NOEXCEPT noexcept +#else +#define PICOJSON_NOEXCEPT throw() +#endif +#endif + +// experimental support for int64_t (see README.mkdn for detail) +#ifdef PICOJSON_USE_INT64 +#define __STDC_FORMAT_MACROS +#include +#if __cplusplus >= 201103L +#include +#else +extern "C" { +#include +} +#endif +#endif + +// to disable the use of localeconv(3), set PICOJSON_USE_LOCALE to 0 +#ifndef PICOJSON_USE_LOCALE +#define PICOJSON_USE_LOCALE 1 +#endif +#if PICOJSON_USE_LOCALE +extern "C" { +#include +} +#endif + +#ifndef PICOJSON_ASSERT +#define PICOJSON_ASSERT(e) \ + do { \ + if (!(e)) \ + throw std::runtime_error(#e); \ + } while (0) +#endif + +#ifdef _MSC_VER +#define SNPRINTF _snprintf_s +#pragma warning(push) +#pragma warning(disable : 4244) // conversion from int to char +#pragma warning(disable : 4127) // conditional expression is constant +#pragma warning(disable : 4702) // unreachable code +#pragma warning(disable : 4706) // assignment within conditional expression +#else +#define SNPRINTF snprintf +#endif + +namespace picojson { + +enum { + null_type, + boolean_type, + number_type, + string_type, + array_type, + object_type +#ifdef PICOJSON_USE_INT64 + , + int64_type +#endif +}; + +enum { INDENT_WIDTH = 2, DEFAULT_MAX_DEPTHS = 100 }; + +struct null {}; + +class value { +public: + typedef std::vector array; + typedef std::map object; + union _storage { + bool boolean_; + double number_; +#ifdef PICOJSON_USE_INT64 + int64_t int64_; +#endif + std::string *string_; + array *array_; + object *object_; + }; + +protected: + int type_; + _storage u_; + +public: + value(); + value(int type, bool); + explicit value(bool b); +#ifdef PICOJSON_USE_INT64 + explicit value(int64_t i); +#endif + explicit value(double n); + explicit value(const std::string &s); + explicit value(const array &a); + explicit value(const object &o); +#if PICOJSON_USE_RVALUE_REFERENCE + explicit value(std::string &&s); + explicit value(array &&a); + explicit value(object &&o); +#endif + explicit value(const char *s); + value(const char *s, size_t len); + ~value(); + value(const value &x); + value &operator=(const value &x); +#if PICOJSON_USE_RVALUE_REFERENCE + value(value &&x) PICOJSON_NOEXCEPT; + value &operator=(value &&x) PICOJSON_NOEXCEPT; +#endif + void swap(value &x) PICOJSON_NOEXCEPT; + template bool is() const; + template const T &get() const; + template T &get(); + template void set(const T &); +#if PICOJSON_USE_RVALUE_REFERENCE + template void set(T &&); +#endif + bool evaluate_as_boolean() const; + const value &get(const size_t idx) const; + const value &get(const std::string &key) const; + value &get(const size_t idx); + value &get(const std::string &key); + + bool contains(const size_t idx) const; + bool contains(const std::string &key) const; + std::string to_str() const; + template void serialize(Iter os, bool prettify = false) const; + std::string serialize(bool prettify = false) const; + +private: + template value(const T *); // intentionally defined to block implicit conversion of pointer to bool + template static void _indent(Iter os, int indent); + template void _serialize(Iter os, int indent) const; + std::string _serialize(int indent) const; + void clear(); +}; + +typedef value::array array; +typedef value::object object; + +inline value::value() : type_(null_type), u_() { +} + +inline value::value(int type, bool) : type_(type), u_() { + switch (type) { +#define INIT(p, v) \ + case p##type: \ + u_.p = v; \ + break + INIT(boolean_, false); + INIT(number_, 0.0); +#ifdef PICOJSON_USE_INT64 + INIT(int64_, 0); +#endif + INIT(string_, new std::string()); + INIT(array_, new array()); + INIT(object_, new object()); +#undef INIT + default: + break; + } +} + +inline value::value(bool b) : type_(boolean_type), u_() { + u_.boolean_ = b; +} + +#ifdef PICOJSON_USE_INT64 +inline value::value(int64_t i) : type_(int64_type), u_() { + u_.int64_ = i; +} +#endif + +inline value::value(double n) : type_(number_type), u_() { + if ( +#ifdef _MSC_VER + !_finite(n) +#elif __cplusplus >= 201103L + std::isnan(n) || std::isinf(n) +#else + isnan(n) || isinf(n) +#endif + ) { + throw std::overflow_error(""); + } + u_.number_ = n; +} + +inline value::value(const std::string &s) : type_(string_type), u_() { + u_.string_ = new std::string(s); +} + +inline value::value(const array &a) : type_(array_type), u_() { + u_.array_ = new array(a); +} + +inline value::value(const object &o) : type_(object_type), u_() { + u_.object_ = new object(o); +} + +#if PICOJSON_USE_RVALUE_REFERENCE +inline value::value(std::string &&s) : type_(string_type), u_() { + u_.string_ = new std::string(std::move(s)); +} + +inline value::value(array &&a) : type_(array_type), u_() { + u_.array_ = new array(std::move(a)); +} + +inline value::value(object &&o) : type_(object_type), u_() { + u_.object_ = new object(std::move(o)); +} +#endif + +inline value::value(const char *s) : type_(string_type), u_() { + u_.string_ = new std::string(s); +} + +inline value::value(const char *s, size_t len) : type_(string_type), u_() { + u_.string_ = new std::string(s, len); +} + +inline void value::clear() { + switch (type_) { +#define DEINIT(p) \ + case p##type: \ + delete u_.p; \ + break + DEINIT(string_); + DEINIT(array_); + DEINIT(object_); +#undef DEINIT + default: + break; + } +} + +inline value::~value() { + clear(); +} + +inline value::value(const value &x) : type_(x.type_), u_() { + switch (type_) { +#define INIT(p, v) \ + case p##type: \ + u_.p = v; \ + break + INIT(string_, new std::string(*x.u_.string_)); + INIT(array_, new array(*x.u_.array_)); + INIT(object_, new object(*x.u_.object_)); +#undef INIT + default: + u_ = x.u_; + break; + } +} + +inline value &value::operator=(const value &x) { + if (this != &x) { + value t(x); + swap(t); + } + return *this; +} + +#if PICOJSON_USE_RVALUE_REFERENCE +inline value::value(value &&x) PICOJSON_NOEXCEPT : type_(null_type), u_() { + swap(x); +} +inline value &value::operator=(value &&x) PICOJSON_NOEXCEPT { + swap(x); + return *this; +} +#endif +inline void value::swap(value &x) PICOJSON_NOEXCEPT { + std::swap(type_, x.type_); + std::swap(u_, x.u_); +} + +#define IS(ctype, jtype) \ + template <> inline bool value::is() const { \ + return type_ == jtype##_type; \ + } +IS(null, null) +IS(bool, boolean) +#ifdef PICOJSON_USE_INT64 +IS(int64_t, int64) +#endif +IS(std::string, string) +IS(array, array) +IS(object, object) +#undef IS +template <> inline bool value::is() const { + return type_ == number_type +#ifdef PICOJSON_USE_INT64 + || type_ == int64_type +#endif + ; +} + +#define GET(ctype, var) \ + template <> inline const ctype &value::get() const { \ + PICOJSON_ASSERT("type mismatch! call is() before get()" && is()); \ + return var; \ + } \ + template <> inline ctype &value::get() { \ + PICOJSON_ASSERT("type mismatch! call is() before get()" && is()); \ + return var; \ + } +GET(bool, u_.boolean_) +GET(std::string, *u_.string_) +GET(array, *u_.array_) +GET(object, *u_.object_) +#ifdef PICOJSON_USE_INT64 +GET(double, + (type_ == int64_type && (const_cast(this)->type_ = number_type, (const_cast(this)->u_.number_ = u_.int64_)), + u_.number_)) +GET(int64_t, u_.int64_) +#else +GET(double, u_.number_) +#endif +#undef GET + +#define SET(ctype, jtype, setter) \ + template <> inline void value::set(const ctype &_val) { \ + clear(); \ + type_ = jtype##_type; \ + setter \ + } +SET(bool, boolean, u_.boolean_ = _val;) +SET(std::string, string, u_.string_ = new std::string(_val);) +SET(array, array, u_.array_ = new array(_val);) +SET(object, object, u_.object_ = new object(_val);) +SET(double, number, u_.number_ = _val;) +#ifdef PICOJSON_USE_INT64 +SET(int64_t, int64, u_.int64_ = _val;) +#endif +#undef SET + +#if PICOJSON_USE_RVALUE_REFERENCE +#define MOVESET(ctype, jtype, setter) \ + template <> inline void value::set(ctype && _val) { \ + clear(); \ + type_ = jtype##_type; \ + setter \ + } +MOVESET(std::string, string, u_.string_ = new std::string(std::move(_val));) +MOVESET(array, array, u_.array_ = new array(std::move(_val));) +MOVESET(object, object, u_.object_ = new object(std::move(_val));) +#undef MOVESET +#endif + +inline bool value::evaluate_as_boolean() const { + switch (type_) { + case null_type: + return false; + case boolean_type: + return u_.boolean_; + case number_type: + return u_.number_ != 0; +#ifdef PICOJSON_USE_INT64 + case int64_type: + return u_.int64_ != 0; +#endif + case string_type: + return !u_.string_->empty(); + default: + return true; + } +} + +inline const value &value::get(const size_t idx) const { + static value s_null; + PICOJSON_ASSERT(is()); + return idx < u_.array_->size() ? (*u_.array_)[idx] : s_null; +} + +inline value &value::get(const size_t idx) { + static value s_null; + PICOJSON_ASSERT(is()); + return idx < u_.array_->size() ? (*u_.array_)[idx] : s_null; +} + +inline const value &value::get(const std::string &key) const { + static value s_null; + PICOJSON_ASSERT(is()); + object::const_iterator i = u_.object_->find(key); + return i != u_.object_->end() ? i->second : s_null; +} + +inline value &value::get(const std::string &key) { + static value s_null; + PICOJSON_ASSERT(is()); + object::iterator i = u_.object_->find(key); + return i != u_.object_->end() ? i->second : s_null; +} + +inline bool value::contains(const size_t idx) const { + PICOJSON_ASSERT(is()); + return idx < u_.array_->size(); +} + +inline bool value::contains(const std::string &key) const { + PICOJSON_ASSERT(is()); + object::const_iterator i = u_.object_->find(key); + return i != u_.object_->end(); +} + +inline std::string value::to_str() const { + switch (type_) { + case null_type: + return "null"; + case boolean_type: + return u_.boolean_ ? "true" : "false"; +#ifdef PICOJSON_USE_INT64 + case int64_type: { + char buf[sizeof("-9223372036854775808")]; + SNPRINTF(buf, sizeof(buf), "%" PRId64, u_.int64_); + return buf; + } +#endif + case number_type: { + char buf[256]; + double tmp; + SNPRINTF(buf, sizeof(buf), fabs(u_.number_) < (1ULL << 53) && modf(u_.number_, &tmp) == 0 ? "%.f" : "%.17g", u_.number_); +#if PICOJSON_USE_LOCALE + char *decimal_point = localeconv()->decimal_point; + if (strcmp(decimal_point, ".") != 0) { + size_t decimal_point_len = strlen(decimal_point); + for (char *p = buf; *p != '\0'; ++p) { + if (strncmp(p, decimal_point, decimal_point_len) == 0) { + return std::string(buf, p) + "." + (p + decimal_point_len); + } + } + } +#endif + return buf; + } + case string_type: + return *u_.string_; + case array_type: + return "array"; + case object_type: + return "object"; + default: + PICOJSON_ASSERT(0); +#ifdef _MSC_VER + __assume(0); +#endif + } + return std::string(); +} + +template void copy(const std::string &s, Iter oi) { + std::copy(s.begin(), s.end(), oi); +} + +template struct serialize_str_char { + Iter oi; + void operator()(char c) { + switch (c) { +#define MAP(val, sym) \ + case val: \ + copy(sym, oi); \ + break + MAP('"', "\\\""); + MAP('\\', "\\\\"); + MAP('/', "\\/"); + MAP('\b', "\\b"); + MAP('\f', "\\f"); + MAP('\n', "\\n"); + MAP('\r', "\\r"); + MAP('\t', "\\t"); +#undef MAP + default: + if (static_cast(c) < 0x20 || c == 0x7f) { + char buf[7]; + SNPRINTF(buf, sizeof(buf), "\\u%04x", c & 0xff); + copy(buf, buf + 6, oi); + } else { + *oi++ = c; + } + break; + } + } +}; + +template void serialize_str(const std::string &s, Iter oi) { + *oi++ = '"'; + serialize_str_char process_char = {oi}; + std::for_each(s.begin(), s.end(), process_char); + *oi++ = '"'; +} + +template void value::serialize(Iter oi, bool prettify) const { + return _serialize(oi, prettify ? 0 : -1); +} + +inline std::string value::serialize(bool prettify) const { + return _serialize(prettify ? 0 : -1); +} + +template void value::_indent(Iter oi, int indent) { + *oi++ = '\n'; + for (int i = 0; i < indent * INDENT_WIDTH; ++i) { + *oi++ = ' '; + } +} + +template void value::_serialize(Iter oi, int indent) const { + switch (type_) { + case string_type: + serialize_str(*u_.string_, oi); + break; + case array_type: { + *oi++ = '['; + if (indent != -1) { + ++indent; + } + for (array::const_iterator i = u_.array_->begin(); i != u_.array_->end(); ++i) { + if (i != u_.array_->begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + i->_serialize(oi, indent); + } + if (indent != -1) { + --indent; + if (!u_.array_->empty()) { + _indent(oi, indent); + } + } + *oi++ = ']'; + break; + } + case object_type: { + *oi++ = '{'; + if (indent != -1) { + ++indent; + } + for (object::const_iterator i = u_.object_->begin(); i != u_.object_->end(); ++i) { + if (i != u_.object_->begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + serialize_str(i->first, oi); + *oi++ = ':'; + if (indent != -1) { + *oi++ = ' '; + } + i->second._serialize(oi, indent); + } + if (indent != -1) { + --indent; + if (!u_.object_->empty()) { + _indent(oi, indent); + } + } + *oi++ = '}'; + break; + } + default: + copy(to_str(), oi); + break; + } + if (indent == 0) { + *oi++ = '\n'; + } +} + +inline std::string value::_serialize(int indent) const { + std::string s; + _serialize(std::back_inserter(s), indent); + return s; +} + +template class input { +protected: + Iter cur_, end_; + bool consumed_; + int line_; + +public: + input(const Iter &first, const Iter &last) : cur_(first), end_(last), consumed_(false), line_(1) { + } + int getc() { + if (consumed_) { + if (*cur_ == '\n') { + ++line_; + } + ++cur_; + } + if (cur_ == end_) { + consumed_ = false; + return -1; + } + consumed_ = true; + return *cur_ & 0xff; + } + void ungetc() { + consumed_ = false; + } + Iter cur() const { + if (consumed_) { + input *self = const_cast *>(this); + self->consumed_ = false; + ++self->cur_; + } + return cur_; + } + int line() const { + return line_; + } + void skip_ws() { + while (1) { + int ch = getc(); + if (!(ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r')) { + ungetc(); + break; + } + } + } + bool expect(const int expected) { + skip_ws(); + if (getc() != expected) { + ungetc(); + return false; + } + return true; + } + bool match(const std::string &pattern) { + for (std::string::const_iterator pi(pattern.begin()); pi != pattern.end(); ++pi) { + if (getc() != *pi) { + ungetc(); + return false; + } + } + return true; + } +}; + +template inline int _parse_quadhex(input &in) { + int uni_ch = 0, hex; + for (int i = 0; i < 4; i++) { + if ((hex = in.getc()) == -1) { + return -1; + } + if ('0' <= hex && hex <= '9') { + hex -= '0'; + } else if ('A' <= hex && hex <= 'F') { + hex -= 'A' - 0xa; + } else if ('a' <= hex && hex <= 'f') { + hex -= 'a' - 0xa; + } else { + in.ungetc(); + return -1; + } + uni_ch = uni_ch * 16 + hex; + } + return uni_ch; +} + +template inline bool _parse_codepoint(String &out, input &in) { + int uni_ch; + if ((uni_ch = _parse_quadhex(in)) == -1) { + return false; + } + if (0xd800 <= uni_ch && uni_ch <= 0xdfff) { + if (0xdc00 <= uni_ch) { + // a second 16-bit of a surrogate pair appeared + return false; + } + // first 16-bit of surrogate pair, get the next one + if (in.getc() != '\\' || in.getc() != 'u') { + in.ungetc(); + return false; + } + int second = _parse_quadhex(in); + if (!(0xdc00 <= second && second <= 0xdfff)) { + return false; + } + uni_ch = ((uni_ch - 0xd800) << 10) | ((second - 0xdc00) & 0x3ff); + uni_ch += 0x10000; + } + if (uni_ch < 0x80) { + out.push_back(static_cast(uni_ch)); + } else { + if (uni_ch < 0x800) { + out.push_back(static_cast(0xc0 | (uni_ch >> 6))); + } else { + if (uni_ch < 0x10000) { + out.push_back(static_cast(0xe0 | (uni_ch >> 12))); + } else { + out.push_back(static_cast(0xf0 | (uni_ch >> 18))); + out.push_back(static_cast(0x80 | ((uni_ch >> 12) & 0x3f))); + } + out.push_back(static_cast(0x80 | ((uni_ch >> 6) & 0x3f))); + } + out.push_back(static_cast(0x80 | (uni_ch & 0x3f))); + } + return true; +} + +template inline bool _parse_string(String &out, input &in) { + while (1) { + int ch = in.getc(); + if (ch < ' ') { + in.ungetc(); + return false; + } else if (ch == '"') { + return true; + } else if (ch == '\\') { + if ((ch = in.getc()) == -1) { + return false; + } + switch (ch) { +#define MAP(sym, val) \ + case sym: \ + out.push_back(val); \ + break + MAP('"', '\"'); + MAP('\\', '\\'); + MAP('/', '/'); + MAP('b', '\b'); + MAP('f', '\f'); + MAP('n', '\n'); + MAP('r', '\r'); + MAP('t', '\t'); +#undef MAP + case 'u': + if (!_parse_codepoint(out, in)) { + return false; + } + break; + default: + return false; + } + } else { + out.push_back(static_cast(ch)); + } + } + return false; +} + +template inline bool _parse_array(Context &ctx, input &in) { + if (!ctx.parse_array_start()) { + return false; + } + size_t idx = 0; + if (in.expect(']')) { + return ctx.parse_array_stop(idx); + } + do { + if (!ctx.parse_array_item(in, idx)) { + return false; + } + idx++; + } while (in.expect(',')); + return in.expect(']') && ctx.parse_array_stop(idx); +} + +template inline bool _parse_object(Context &ctx, input &in) { + if (!ctx.parse_object_start()) { + return false; + } + if (in.expect('}')) { + return ctx.parse_object_stop(); + } + do { + std::string key; + if (!in.expect('"') || !_parse_string(key, in) || !in.expect(':')) { + return false; + } + if (!ctx.parse_object_item(in, key)) { + return false; + } + } while (in.expect(',')); + return in.expect('}') && ctx.parse_object_stop(); +} + +template inline std::string _parse_number(input &in) { + std::string num_str; + while (1) { + int ch = in.getc(); + if (('0' <= ch && ch <= '9') || ch == '+' || ch == '-' || ch == 'e' || ch == 'E') { + num_str.push_back(static_cast(ch)); + } else if (ch == '.') { +#if PICOJSON_USE_LOCALE + num_str += localeconv()->decimal_point; +#else + num_str.push_back('.'); +#endif + } else { + in.ungetc(); + break; + } + } + return num_str; +} + +template inline bool _parse(Context &ctx, input &in) { + in.skip_ws(); + int ch = in.getc(); + switch (ch) { +#define IS(ch, text, op) \ + case ch: \ + if (in.match(text) && op) { \ + return true; \ + } else { \ + return false; \ + } + IS('n', "ull", ctx.set_null()); + IS('f', "alse", ctx.set_bool(false)); + IS('t', "rue", ctx.set_bool(true)); +#undef IS + case '"': + return ctx.parse_string(in); + case '[': + return _parse_array(ctx, in); + case '{': + return _parse_object(ctx, in); + default: + if (('0' <= ch && ch <= '9') || ch == '-') { + double f; + char *endp; + in.ungetc(); + std::string num_str(_parse_number(in)); + if (num_str.empty()) { + return false; + } +#ifdef PICOJSON_USE_INT64 + { + errno = 0; + intmax_t ival = strtoimax(num_str.c_str(), &endp, 10); + if (errno == 0 && std::numeric_limits::min() <= ival && ival <= std::numeric_limits::max() && + endp == num_str.c_str() + num_str.size()) { + ctx.set_int64(ival); + return true; + } + } +#endif + f = strtod(num_str.c_str(), &endp); + if (endp == num_str.c_str() + num_str.size()) { + ctx.set_number(f); + return true; + } + return false; + } + break; + } + in.ungetc(); + return false; +} + +class deny_parse_context { +public: + bool set_null() { + return false; + } + bool set_bool(bool) { + return false; + } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t) { + return false; + } +#endif + bool set_number(double) { + return false; + } + template bool parse_string(input &) { + return false; + } + bool parse_array_start() { + return false; + } + template bool parse_array_item(input &, size_t) { + return false; + } + bool parse_array_stop(size_t) { + return false; + } + bool parse_object_start() { + return false; + } + template bool parse_object_item(input &, const std::string &) { + return false; + } +}; + +class default_parse_context { +protected: + value *out_; + size_t depths_; + +public: + default_parse_context(value *out, size_t depths = DEFAULT_MAX_DEPTHS) : out_(out), depths_(depths) { + } + bool set_null() { + *out_ = value(); + return true; + } + bool set_bool(bool b) { + *out_ = value(b); + return true; + } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t i) { + *out_ = value(i); + return true; + } +#endif + bool set_number(double f) { + *out_ = value(f); + return true; + } + template bool parse_string(input &in) { + *out_ = value(string_type, false); + return _parse_string(out_->get(), in); + } + bool parse_array_start() { + if (depths_ == 0) + return false; + --depths_; + *out_ = value(array_type, false); + return true; + } + template bool parse_array_item(input &in, size_t) { + array &a = out_->get(); + a.push_back(value()); + default_parse_context ctx(&a.back(), depths_); + return _parse(ctx, in); + } + bool parse_array_stop(size_t) { + ++depths_; + return true; + } + bool parse_object_start() { + if (depths_ == 0) + return false; + *out_ = value(object_type, false); + return true; + } + template bool parse_object_item(input &in, const std::string &key) { + object &o = out_->get(); + default_parse_context ctx(&o[key], depths_); + return _parse(ctx, in); + } + bool parse_object_stop() { + ++depths_; + return true; + } + +private: + default_parse_context(const default_parse_context &); + default_parse_context &operator=(const default_parse_context &); +}; + +class null_parse_context { +protected: + size_t depths_; + +public: + struct dummy_str { + void push_back(int) { + } + }; + +public: + null_parse_context(size_t depths = DEFAULT_MAX_DEPTHS) : depths_(depths) { + } + bool set_null() { + return true; + } + bool set_bool(bool) { + return true; + } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t) { + return true; + } +#endif + bool set_number(double) { + return true; + } + template bool parse_string(input &in) { + dummy_str s; + return _parse_string(s, in); + } + bool parse_array_start() { + if (depths_ == 0) + return false; + --depths_; + return true; + } + template bool parse_array_item(input &in, size_t) { + return _parse(*this, in); + } + bool parse_array_stop(size_t) { + ++depths_; + return true; + } + bool parse_object_start() { + if (depths_ == 0) + return false; + --depths_; + return true; + } + template bool parse_object_item(input &in, const std::string &) { + ++depths_; + return _parse(*this, in); + } + bool parse_object_stop() { + return true; + } + +private: + null_parse_context(const null_parse_context &); + null_parse_context &operator=(const null_parse_context &); +}; + +// obsolete, use the version below +template inline std::string parse(value &out, Iter &pos, const Iter &last) { + std::string err; + pos = parse(out, pos, last, &err); + return err; +} + +template inline Iter _parse(Context &ctx, const Iter &first, const Iter &last, std::string *err) { + input in(first, last); + if (!_parse(ctx, in) && err != NULL) { + char buf[64]; + SNPRINTF(buf, sizeof(buf), "syntax error at line %d near: ", in.line()); + *err = buf; + while (1) { + int ch = in.getc(); + if (ch == -1 || ch == '\n') { + break; + } else if (ch >= ' ') { + err->push_back(static_cast(ch)); + } + } + } + return in.cur(); +} + +template inline Iter parse(value &out, const Iter &first, const Iter &last, std::string *err) { + default_parse_context ctx(&out); + return _parse(ctx, first, last, err); +} + +inline std::string parse(value &out, const std::string &s) { + std::string err; + parse(out, s.begin(), s.end(), &err); + return err; +} + +inline std::string parse(value &out, std::istream &is) { + std::string err; + parse(out, std::istreambuf_iterator(is.rdbuf()), std::istreambuf_iterator(), &err); + return err; +} + +template struct last_error_t { static std::string s; }; +template std::string last_error_t::s; + +inline void set_last_error(const std::string &s) { + last_error_t::s = s; +} + +inline const std::string &get_last_error() { + return last_error_t::s; +} + +inline bool operator==(const value &x, const value &y) { + if (x.is()) + return y.is(); +#define PICOJSON_CMP(type) \ + if (x.is()) \ + return y.is() && x.get() == y.get() + PICOJSON_CMP(bool); + PICOJSON_CMP(double); + PICOJSON_CMP(std::string); + PICOJSON_CMP(array); + PICOJSON_CMP(object); +#undef PICOJSON_CMP + PICOJSON_ASSERT(0); +#ifdef _MSC_VER + __assume(0); +#endif + return false; +} + +inline bool operator!=(const value &x, const value &y) { + return !(x == y); +} +} + +#if !PICOJSON_USE_RVALUE_REFERENCE +namespace std { +template <> inline void swap(picojson::value &x, picojson::value &y) { + x.swap(y); +} +} +#endif + +inline std::istream &operator>>(std::istream &is, picojson::value &x) { + picojson::set_last_error(std::string()); + const std::string err(picojson::parse(x, is)); + if (!err.empty()) { + picojson::set_last_error(err); + is.setstate(std::ios::failbit); + } + return is; +} + +inline std::ostream &operator<<(std::ostream &os, const picojson::value &x) { + x.serialize(std::ostream_iterator(os)); + return os; +} +#ifdef _MSC_VER +#pragma warning(pop) +#endif + +#endif diff --git a/cactus/engine/engine_bpe.cpp b/cactus-engine/src/bpe.cpp similarity index 63% rename from cactus/engine/engine_bpe.cpp rename to cactus-engine/src/bpe.cpp index f29b1a4db..41d49c795 100644 --- a/cactus/engine/engine_bpe.cpp +++ b/cactus-engine/src/bpe.cpp @@ -2,15 +2,19 @@ #include #include #include +#include #include #include #include #include -#include namespace cactus { namespace engine { +namespace { +constexpr const char* kMetaspace = "\xE2\x96\x81"; +} // namespace + BPETokenizer::BPETokenizer() : vocab_size_(0), unk_token_id_(0), bos_token_id_(1), eos_token_id_(2), vocab_mmap_ptr_(nullptr), vocab_mmap_size_(0), @@ -57,23 +61,58 @@ bool BPETokenizer::load_vocabulary_mmap(const std::string& vocab_file, const std }; std::string line; - uint32_t id = 0; token_to_id_.clear(); id_to_token_.clear(); special_tokens_.clear(); + const bool use_id_tab_vocab = + runtime_config_.vocab_format == TokenizerRuntimeConfig::VocabFormat::ID_TAB_TOKEN; - while (std::getline(vocab_stream, line)) { - rtrim_cr(line); - if (line.empty()) continue; - token_to_id_[line] = id; - id_to_token_.push_back(line); + if (use_id_tab_vocab) { + while (std::getline(vocab_stream, line)) { + rtrim_cr(line); - if (!line.empty() && line.front() == '<' && line.back() == '>') { - special_tokens_[line] = id; + std::string token; + uint32_t id = UINT32_MAX; + + std::istringstream iss(line); + if (iss >> id) { + if (std::getline(iss, token) && !token.empty() && token[0] == '\t') { + token = token.substr(1); + } + + if (token.empty()) { + auto last_pos = vocab_stream.tellg(); + while (std::getline(vocab_stream, line)) { + rtrim_cr(line); + if (!line.empty()) { + break; + } + token += '\n'; + last_pos = vocab_stream.tellg(); + } + vocab_stream.seekg(last_pos); + } + } + + if (!token.empty() && id != UINT32_MAX) { + token_to_id_[token] = id; + if (id >= id_to_token_.size()) { + id_to_token_.resize(id + 1); + } + id_to_token_[id] = token; + } + } + vocab_size_ = static_cast(id_to_token_.size()); + } else { + uint32_t id = 0; + while (std::getline(vocab_stream, line)) { + rtrim_cr(line); + token_to_id_[line] = id; + id_to_token_.push_back(line); + ++id; } - id++; + vocab_size_ = id; } - vocab_size_ = id; int merges_fd = open(merges_file.c_str(), O_RDONLY); if (merges_fd == -1) return false; @@ -109,7 +148,7 @@ bool BPETokenizer::load_vocabulary_mmap(const std::string& vocab_file, const std std::string merged = first + second; merge_rules_.emplace_back(first, second, merged, priority); - std::string key = first + "\x00" + second; + std::string key = first + '\0' + second; auto it = merge_map_.find(key); if (it == merge_map_.end() || priority < it->second) { merge_map_[key] = priority; @@ -117,6 +156,34 @@ bool BPETokenizer::load_vocabulary_mmap(const std::string& vocab_file, const std priority++; } } + + auto is_whitespace_only = [](const std::string& s) { + if (s.empty()) return false; + for (char c : s) { + if (c != ' ' && c != '\n' && c != '\t' && c != '\r') return false; + } + return true; + }; + std::vector ws_tokens; + for (const auto& tok : id_to_token_) { + if (tok.size() >= 2 && is_whitespace_only(tok)) ws_tokens.push_back(tok); + } + std::sort(ws_tokens.begin(), ws_tokens.end(), + [](const std::string& a, const std::string& b) { return a.size() < b.size(); }); + for (const auto& tok : ws_tokens) { + for (size_t i = 1; i < tok.size(); i++) { + std::string first = tok.substr(0, i); + std::string second = tok.substr(i); + if (!token_to_id_.count(first) || !token_to_id_.count(second)) continue; + std::string key = first + '\0' + second; + if (merge_map_.find(key) == merge_map_.end()) { + merge_rules_.emplace_back(first, second, tok, priority); + merge_map_[key] = priority; + priority++; + } + break; + } + } std::sort(merge_rules_.begin(), merge_rules_.end(), [](const MergeRule& a, const MergeRule& b) { @@ -127,6 +194,7 @@ bool BPETokenizer::load_vocabulary_mmap(const std::string& vocab_file, const std } bool BPETokenizer::load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) { + runtime_config_ = load_tokenizer_runtime_config(config_file); if (!load_vocabulary_mmap(vocab_file, merges_file)) { return false; } @@ -171,55 +239,6 @@ bool BPETokenizer::load_vocabulary_with_config(const std::string& vocab_file, co std::string special_tokens_path = dir + "/special_tokens.json"; load_special_tokens(special_tokens_path); - try { - std::ifstream tok_json(dir + "/tokenizer.json"); - if (tok_json.is_open()) { - std::string content((std::istreambuf_iterator(tok_json)), std::istreambuf_iterator()); - size_t pos = 0; - while (true) { - size_t id_key = content.find("\"id\"", pos); - if (id_key == std::string::npos) break; - size_t id_colon = content.find(':', id_key); - if (id_colon == std::string::npos) break; - size_t id_end = content.find_first_of(",}\n", id_colon + 1); - if (id_end == std::string::npos) break; - std::string id_str = content.substr(id_colon + 1, id_end - id_colon - 1); - id_str.erase(0, id_str.find_first_not_of(" \t\n\r")); - id_str.erase(id_str.find_last_not_of(" \t\n\r") + 1); - - size_t content_key = content.find("\"content\"", id_end); - if (content_key == std::string::npos) { pos = id_end; continue; } - size_t cont_quote1 = content.find('"', content_key + 9); - if (cont_quote1 == std::string::npos) { pos = id_end; continue; } - size_t cont_quote2 = content.find('"', cont_quote1 + 1); - if (cont_quote2 == std::string::npos) { pos = id_end; continue; } - std::string token_content = content.substr(cont_quote1 + 1, cont_quote2 - cont_quote1 - 1); - - size_t special_key = content.find("\"special\"", cont_quote2); - if (special_key == std::string::npos) { pos = cont_quote2; continue; } - size_t special_colon = content.find(':', special_key); - if (special_colon == std::string::npos) { pos = cont_quote2; continue; } - size_t special_val_start = content.find_first_not_of(" \t\n\r", special_colon + 1); - if (special_val_start == std::string::npos) { pos = cont_quote2; continue; } - bool is_special = false; - if (content.compare(special_val_start, 4, "true") == 0) { - is_special = true; - } - if (is_special) { - try { - uint32_t token_id = static_cast(std::stoul(id_str)); - special_tokens_[token_content] = token_id; - } catch (...) { - - } - } - pos = cont_quote2 + 1; - } - } - } catch (...) { - - } - std::string template_path = dir + "/chat_template.jinja2"; load_chat_template(template_path); @@ -230,87 +249,11 @@ bool BPETokenizer::load_vocabulary_with_config(const std::string& vocab_file, co } void BPETokenizer::load_special_tokens(const std::string& config_file) { - std::ifstream file(config_file); - if (!file.is_open()) { - return; - } - - std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - - size_t pos = content.find("\"special_tokens\""); - if (pos == std::string::npos) return; - - pos = content.find("{", pos); - if (pos == std::string::npos) return; - - size_t end_pos = content.find("}", pos); - if (end_pos == std::string::npos) return; - - std::string special_tokens_section = content.substr(pos + 1, end_pos - pos - 1); - - std::istringstream iss(special_tokens_section); - std::string line; - - while (std::getline(iss, line)) { - size_t colon_pos = line.find(":"); - if (colon_pos == std::string::npos) continue; - - std::string id_part = line.substr(0, colon_pos); - std::string token_part = line.substr(colon_pos + 1); - - size_t id_start = id_part.find("\""); - size_t id_end = id_part.find("\"", id_start + 1); - if (id_start == std::string::npos || id_end == std::string::npos) continue; - - std::string id_str = id_part.substr(id_start + 1, id_end - id_start - 1); - uint32_t token_id = std::stoul(id_str); - - size_t token_start = token_part.find("\""); - size_t token_end = token_part.rfind("\""); - if (token_start == std::string::npos || token_end == std::string::npos || token_start >= token_end) continue; - - std::string token_content = token_part.substr(token_start + 1, token_end - token_start - 1); - - special_tokens_[token_content] = token_id; - } - + load_special_tokens_map(config_file, special_tokens_); } std::vector BPETokenizer::split_with_special_tokens(const std::string& text) const { - std::vector result; - - size_t start = 0; - while (start < text.size()) { - size_t best_match_pos = text.size(); - size_t best_match_len = 0; - std::string best_special_token; - - for (const auto& [special_token, token_id] : special_tokens_) { - size_t pos = text.find(special_token, start); - if (pos != std::string::npos && pos < best_match_pos) { - best_match_pos = pos; - best_match_len = special_token.length(); - best_special_token = special_token; - } - } - - if (best_match_pos < text.size()) { - if (best_match_pos > start) { - std::string before = text.substr(start, best_match_pos - start); - result.push_back(before); - } - - result.push_back(best_special_token); - start = best_match_pos + best_match_len; - } else { - if (start < text.size()) { - result.push_back(text.substr(start)); - } - break; - } - } - - return result; + return cactus::engine::split_with_special_tokens(text, special_tokens_); } void BPETokenizer::init_byte_mappings() const { @@ -402,6 +345,8 @@ std::string BPETokenizer::unicode_to_bytes(const std::string& text) const { auto it = unicode_to_byte_.find(unicode_char); if (it != unicode_to_byte_.end()) { result += static_cast(it->second); + } else if (unicode_char.size() == 1 && (unsigned char)unicode_char[0] < 128) { + result += unicode_char[0]; } else { result += '?'; } @@ -436,13 +381,39 @@ std::vector BPETokenizer::byte_level_split(const std::string& text) return chars; } +std::vector BPETokenizer::utf8_split(const std::string& text) const { + std::vector chars; + size_t i = 0; + while (i < text.length()) { + size_t char_len = 1; + const unsigned char byte = static_cast(text[i]); + + if ((byte & 0x80) == 0) { + char_len = 1; + } else if ((byte & 0xE0) == 0xC0) { + char_len = 2; + } else if ((byte & 0xF0) == 0xE0) { + char_len = 3; + } else if ((byte & 0xF8) == 0xF0) { + char_len = 4; + } + + if (i + char_len <= text.length()) { + chars.push_back(text.substr(i, char_len)); + } + i += char_len; + } + + return chars; +} + std::pair BPETokenizer::find_best_merge_fast(const std::vector& tokens) const { int best_pos = -1; uint32_t best_priority = UINT32_MAX; for (size_t i = 0; i < tokens.size() - 1; ++i) { - std::string key = tokens[i] + "\x00" + tokens[i + 1]; + std::string key = tokens[i] + '\0' + tokens[i + 1]; auto it = merge_map_.find(key); if (it != merge_map_.end()) { if (it->second < best_priority) { @@ -498,16 +469,48 @@ std::vector BPETokenizer::encode(const std::string& text) const { if (special_it != special_tokens_.end()) { token_ids.push_back(special_it->second); } else { - auto chars = byte_level_split(segment); + std::string normalized_segment = segment; + std::vector chars; + if (runtime_config_.normalizer == TokenizerRuntimeConfig::Normalizer::METASPACE) { + size_t pos = 0; + while ((pos = normalized_segment.find(' ', pos)) != std::string::npos) { + normalized_segment.replace(pos, 1, kMetaspace); + pos += 3; + } + chars = utf8_split(normalized_segment); + } else { + chars = byte_level_split(segment); + } auto bpe_tokens = apply_bpe(chars); - for (const auto& token : bpe_tokens) { auto it = token_to_id_.find(token); if (it != token_to_id_.end()) { token_ids.push_back(it->second); } else { - token_ids.push_back(unk_token_id_); + if (!runtime_config_.byte_fallback) { + token_ids.push_back(unk_token_id_); + continue; + } + + std::vector fallback_ids; + fallback_ids.reserve(token.size()); + for (unsigned char byte : token) { + char fallback_token[7]; + std::snprintf(fallback_token, sizeof(fallback_token), "<0x%02X>", byte); + auto fallback_it = token_to_id_.find(fallback_token); + if (fallback_it == token_to_id_.end()) { + fallback_ids.clear(); + break; + } + fallback_ids.push_back(fallback_it->second); + } + + if (fallback_ids.empty()) { + token_ids.push_back(unk_token_id_); + } else { + token_ids.insert(token_ids.end(), fallback_ids.begin(), fallback_ids.end()); + } } } } @@ -517,28 +520,47 @@ std::vector BPETokenizer::encode(const std::string& text) const { } std::string BPETokenizer::decode(const std::vector& tokens) const { + if (runtime_config_.decoder == TokenizerRuntimeConfig::Decoder::REPLACE_METASPACE) { + std::string result; + result.reserve(tokens.size() * 4); + + for (uint32_t token_id : tokens) { + if (token_id >= id_to_token_.size()) continue; + result += id_to_token_[token_id]; + } + + size_t pos = 0; + while ((pos = result.find(kMetaspace, pos)) != std::string::npos) { + result.replace(pos, 3, " "); + pos += 1; + } + return result; + } + std::string unicode_result; + unicode_result.reserve(tokens.size() * 4); + for (uint32_t token_id : tokens) { - if (token_id < id_to_token_.size()) { - unicode_result += id_to_token_[token_id]; + if (token_id >= id_to_token_.size()) continue; + const std::string& tok = id_to_token_[token_id]; + + size_t pos = 0; + while (pos < tok.size()) { + if (pos + 3 <= tok.size() && + (unsigned char)tok[pos] == 0xE2 && + (unsigned char)tok[pos+1] == 0x96 && + (unsigned char)tok[pos+2] == 0x81) { + unicode_result.push_back(' '); + pos += 3; + } else { + unicode_result.push_back(tok[pos++]); + } } } - std::string result = unicode_to_bytes(unicode_result); - - return result; + return unicode_to_bytes(unicode_result); } -void BPETokenizer::load_chat_template(const std::string& template_file) { - std::ifstream file(template_file); - if (!file.is_open()) { - has_chat_template_ = false; - return; - } - chat_template_ = std::string((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - has_chat_template_ = !chat_template_.empty(); } - } -} \ No newline at end of file diff --git a/cactus-engine/src/chat_tools.h b/cactus-engine/src/chat_tools.h new file mode 100644 index 000000000..f81159911 --- /dev/null +++ b/cactus-engine/src/chat_tools.h @@ -0,0 +1,285 @@ +#pragma once + +#include +#include +#include + +namespace chat_tools { + +inline std::string json_escape(const std::string& s) { + std::string out; + out.reserve(s.size() + 2); + for (char c : s) { + switch (c) { + case '"': out += "\\\""; break; + case '\\': out += "\\\\"; break; + case '\n': out += "\\n"; break; + case '\r': out += "\\r"; break; + case '\t': out += "\\t"; break; + default: + if (static_cast(c) < 0x20) { + static const char* hex = "0123456789abcdef"; + out += "\\u00"; + out += hex[(c >> 4) & 0xF]; + out += hex[c & 0xF]; + } else { + out += c; + } + } + } + return out; +} + +inline std::string respace_json(const std::string& s) { + std::string out; + out.reserve(s.size() + s.size() / 8); + bool in_str = false, esc = false; + for (char c : s) { + if (in_str) { + out += c; + if (esc) esc = false; + else if (c == '\\') esc = true; + else if (c == '"') in_str = false; + continue; + } + if (c == '"') { in_str = true; out += c; continue; } + if (std::isspace(static_cast(c))) continue; + out += c; + if (c == ',' || c == ':') out += ' '; + } + return out; +} + +template +inline std::string tool_object_json(const ToolFunction& tool) { + std::string schema; + auto it = tool.parameters.find("schema"); + if (it != tool.parameters.end()) schema = it->second; + std::string obj = "{\"type\":\"function\",\"function\":{\"name\":\"" + json_escape(tool.name) + + "\",\"description\":\"" + json_escape(tool.description) + "\""; + if (!schema.empty()) obj += ",\"parameters\":" + schema; + obj += "}}"; + return respace_json(obj); +} + +template +inline std::string serialize_tools_qwen(const std::vector& tools) { + if (tools.empty()) return ""; + std::string r = + "# Tools\n\nYou may call one or more functions to assist with the user query.\n\n" + "You are provided with function signatures within XML tags:\n"; + for (const auto& t : tools) r += "\n" + tool_object_json(t); + r += "\n\n\nFor each function call, return a json object with function name and " + "arguments within XML tags:\n\n" + "{\"name\": , \"arguments\": }\n"; + return r; +} + +template +inline std::string serialize_tools_lfm2(const std::vector& tools) { + if (tools.empty()) return ""; + std::string r = "List of tools: <|tool_list_start|>["; + for (size_t i = 0; i < tools.size(); ++i) { + if (i) r += ", "; + r += tool_object_json(tools[i]); + } + r += "]<|tool_list_end|>"; + return r; +} + +inline void skip_ws(const std::string& s, size_t& p) { + while (p < s.size() && std::isspace(static_cast(s[p]))) ++p; +} + +inline std::string py_value(const std::string& s, size_t& p); + +inline std::string py_string(const std::string& s, size_t& p) { + std::string out = "\""; + ++p; + while (p < s.size() && s[p] != '"') { + if (s[p] == '\\' && p + 1 < s.size()) { out += s[p]; out += s[p + 1]; p += 2; } + else { out += s[p++]; } + } + if (p < s.size()) ++p; + out += "\""; + return out; +} + +inline std::string py_value(const std::string& s, size_t& p) { + skip_ws(s, p); + if (p >= s.size()) return ""; + char c = s[p]; + if (c == '"') return py_string(s, p); + if (c == '{') { + std::string out = "{"; + ++p; bool first = true; + while (true) { + skip_ws(s, p); + if (p >= s.size() || s[p] == '}') { if (p < s.size()) ++p; break; } + if (s[p] == ',') { ++p; continue; } + std::string key = py_string(s, p); + skip_ws(s, p); + if (p < s.size() && s[p] == ':') ++p; + std::string val = py_value(s, p); + if (!first) out += ", "; + first = false; + out += key + ": " + val; + } + return out + "}"; + } + if (c == '[') { + std::string out = "["; + ++p; bool first = true; + while (true) { + skip_ws(s, p); + if (p >= s.size() || s[p] == ']') { if (p < s.size()) ++p; break; } + if (s[p] == ',') { ++p; continue; } + std::string val = py_value(s, p); + if (!first) out += ", "; + first = false; + out += val; + } + return out + "]"; + } + if (s.compare(p, 4, "true") == 0) { p += 4; return "True"; } + if (s.compare(p, 5, "false") == 0) { p += 5; return "False"; } + if (s.compare(p, 4, "null") == 0) { p += 4; return "None"; } + size_t start = p; + while (p < s.size() && (std::isdigit(static_cast(s[p])) || s[p] == '.' || + s[p] == '-' || s[p] == '+' || s[p] == 'e' || s[p] == 'E')) ++p; + return s.substr(start, p - start); +} + +inline std::string trim(const std::string& s) { + size_t a = s.find_first_not_of(" \t\r\n"); + if (a == std::string::npos) return ""; + size_t b = s.find_last_not_of(" \t\r\n"); + return s.substr(a, b - a + 1); +} + +inline void extract_qwen_tool_calls(std::string& text, std::vector& calls) { + const std::string OPEN = "", CLOSE = ""; + size_t pos; + while ((pos = text.find(OPEN)) != std::string::npos) { + size_t cs = pos + OPEN.size(); + size_t ce = text.find(CLOSE, cs); + std::string inner = trim(text.substr(cs, (ce == std::string::npos ? text.size() : ce) - cs)); + if (!inner.empty() && inner.front() == '{') calls.push_back(inner); + text.erase(pos, (ce == std::string::npos ? text.size() : ce + CLOSE.size()) - pos); + } +} + +inline std::string py_literal_to_json(const std::string& s, size_t& p) { + skip_ws(s, p); + if (p >= s.size()) return "null"; + char c = s[p]; + if (c == '"' || c == '\'') { + char q = c; ++p; + std::string out = "\""; + while (p < s.size() && s[p] != q) { + if (s[p] == '\\' && p + 1 < s.size()) { out += s[p]; out += s[p + 1]; p += 2; continue; } + if (s[p] == '"') out += '\\'; + out += s[p++]; + } + if (p < s.size()) ++p; + return out + "\""; + } + if (c == '[' || c == '{') { + char close = (c == '[') ? ']' : '}'; + std::string out(1, c == '[' ? '[' : '{'); + ++p; bool first = true; + while (true) { + skip_ws(s, p); + if (p >= s.size() || s[p] == close) { if (p < s.size()) ++p; break; } + if (s[p] == ',') { ++p; continue; } + if (!first) out += ", "; + first = false; + std::string key_or_val = py_literal_to_json(s, p); + skip_ws(s, p); + if (close == '}' && p < s.size() && s[p] == ':') { + ++p; + out += key_or_val + ": " + py_literal_to_json(s, p); + } else { + out += key_or_val; + } + } + return out + close; + } + if (s.compare(p, 4, "True") == 0) { p += 4; return "true"; } + if (s.compare(p, 5, "False") == 0) { p += 5; return "false"; } + if (s.compare(p, 4, "None") == 0) { p += 4; return "null"; } + size_t start = p; + while (p < s.size() && s[p] != ',' && s[p] != ')' && s[p] != ']' && s[p] != '}') ++p; + return trim(s.substr(start, p - start)); +} + +inline void extract_lfm2_tool_calls(std::string& text, std::vector& calls) { + const std::string OPEN = "<|tool_call_start|>", CLOSE = "<|tool_call_end|>"; + size_t pos; + while ((pos = text.find(OPEN)) != std::string::npos) { + size_t cs = pos + OPEN.size(); + size_t ce = text.find(CLOSE, cs); + std::string inner = text.substr(cs, (ce == std::string::npos ? text.size() : ce) - cs); + size_t p = 0; + skip_ws(inner, p); + if (p < inner.size() && inner[p] == '[') ++p; + while (p < inner.size()) { + skip_ws(inner, p); + if (p >= inner.size() || inner[p] == ']') break; + if (inner[p] == ',') { ++p; continue; } + size_t name_start = p; + while (p < inner.size() && (std::isalnum(static_cast(inner[p])) || + inner[p] == '_' || inner[p] == '.')) ++p; + std::string name = inner.substr(name_start, p - name_start); + skip_ws(inner, p); + if (name.empty() || p >= inner.size() || inner[p] != '(') { ++p; continue; } + ++p; // ( + std::string args = "{"; + bool first = true; + while (true) { + skip_ws(inner, p); + if (p >= inner.size() || inner[p] == ')') { if (p < inner.size()) ++p; break; } + if (inner[p] == ',') { ++p; continue; } + size_t ks = p; + while (p < inner.size() && (std::isalnum(static_cast(inner[p])) || inner[p] == '_')) ++p; + std::string key = inner.substr(ks, p - ks); + skip_ws(inner, p); + if (p < inner.size() && inner[p] == '=') ++p; + std::string val = py_literal_to_json(inner, p); + if (!first) args += ", "; + first = false; + args += "\"" + key + "\": " + val; + } + args += "}"; + calls.push_back("{\"name\": \"" + name + "\", \"arguments\": " + args + "}"); + } + text.erase(pos, (ce == std::string::npos ? text.size() : ce + CLOSE.size()) - pos); + } +} + +inline std::string pythonic_call(const std::string& name, const std::string& args_json) { + std::string out = name + "("; + size_t p = 0; + skip_ws(args_json, p); + bool first = true; + if (p < args_json.size() && args_json[p] == '{') { + ++p; + while (true) { + skip_ws(args_json, p); + if (p >= args_json.size() || args_json[p] == '}') break; + if (args_json[p] == ',') { ++p; continue; } + if (args_json[p] != '"') break; + std::string key; + { std::string q = py_string(args_json, p); key = q.substr(1, q.size() - 2); } + skip_ws(args_json, p); + if (p < args_json.size() && args_json[p] == ':') ++p; + std::string val = py_value(args_json, p); + if (!first) out += ", "; + first = false; + out += key + "=" + val; + } + } + return out + ")"; +} + +} // namespace chat_tools diff --git a/cactus-engine/src/cloud.cpp b/cactus-engine/src/cloud.cpp new file mode 100644 index 000000000..3ef099947 --- /dev/null +++ b/cactus-engine/src/cloud.cpp @@ -0,0 +1,568 @@ +#include "cloud.h" +#include "telemetry.h" + +#include +#include +#include +#include +#include +#include +#include + +#ifdef CACTUS_USE_CURL +#include +#endif + +namespace cactus { +namespace ffi { + +namespace { + +#ifdef CACTUS_USE_CURL +std::string compact_error_detail(std::string detail) { + detail = trim_string(detail); + if (detail.empty()) return {}; + + for (char& c : detail) { + if (c == '\n' || c == '\r' || c == '\t') c = ' '; + } + while (detail.find(" ") != std::string::npos) { + detail.erase(detail.find(" "), 1); + } + + constexpr size_t kMaxLen = 160; + if (detail.size() > kMaxLen) { + detail = detail.substr(0, kMaxLen) + "..."; + } + return detail; +} +#endif + +} // namespace + +std::string resolve_cloud_api_key(const char* cloud_key_param) { + const char* env_cloud = std::getenv("CACTUS_CLOUD_KEY"); + const char* env_cloud_legacy = std::getenv("CACTUS_CLOUD_API_KEY"); // keeping this cuz Matthew Hayes had some code reliant on Cactus_cloud_api_key + std::string resolved_cloud_key; + bool key_from_param_or_env = false; + + if (cloud_key_param && *cloud_key_param) { + resolved_cloud_key = cloud_key_param; + key_from_param_or_env = true; + } else if (env_cloud && *env_cloud) { + resolved_cloud_key = env_cloud; + key_from_param_or_env = true; + } else if (env_cloud_legacy && *env_cloud_legacy) { + resolved_cloud_key = env_cloud_legacy; + key_from_param_or_env = true; + } else { + resolved_cloud_key = cactus::telemetry::loadCachedCloudApiKey(); + } + + resolved_cloud_key = trim_string(resolved_cloud_key); + + const bool should_cache = + !resolved_cloud_key.empty() && + key_from_param_or_env; + if (should_cache) { + cactus::telemetry::cacheCloudApiKey(resolved_cloud_key.c_str()); + } + + return resolved_cloud_key; +} + +namespace { + +#ifdef CACTUS_USE_CURL +static std::atomic g_warned_missing_cloud_api_key{false}; + +static void apply_curl_tls_trust(CURL* curl) { + if (!curl) return; + const char* ca_bundle = std::getenv("CACTUS_CA_BUNDLE"); + if (ca_bundle && ca_bundle[0] != '\0') { + curl_easy_setopt(curl, CURLOPT_CAINFO, ca_bundle); + } +#if defined(__ANDROID__) + const char* ca_path = std::getenv("CACTUS_CA_PATH"); + if (ca_path && ca_path[0] != '\0') { + curl_easy_setopt(curl, CURLOPT_CAPATH, ca_path); + } else { + curl_easy_setopt(curl, CURLOPT_CAPATH, "/system/etc/security/cacerts"); + } +#endif +} + +static size_t curl_write_cb(void* ptr, size_t size, size_t nmemb, void* userdata) { + auto* s = static_cast(userdata); + s->append(static_cast(ptr), size * nmemb); + return size * nmemb; +} + +static bool read_file_bytes(const std::string& path, std::vector& out) { + std::ifstream in(path, std::ios::binary); + if (!in.is_open()) return false; + in.seekg(0, std::ios::end); + std::streamsize size = in.tellg(); + if (size <= 0) return false; + in.seekg(0, std::ios::beg); + + out.resize(static_cast(size)); + if (!in.read(reinterpret_cast(out.data()), size)) return false; + return true; +} + +static std::string infer_mime_type(const std::string& path) { + auto dot = path.find_last_of('.'); + if (dot == std::string::npos) return "image/jpeg"; + std::string ext = path.substr(dot + 1); + std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); }); + + if (ext == "png") return "image/png"; + if (ext == "webp") return "image/webp"; + if (ext == "gif") return "image/gif"; + return "image/jpeg"; +} + +static std::string build_cloud_text_prompt(const CloudCompletionRequest& request) { + std::ostringstream oss; + oss << "You are continuing an existing assistant conversation.\\n"; + oss << "Output contract:\\n"; + oss << "1) Never include role prefixes like 'assistant:'.\\n"; + oss << "2) Never include markdown/code fences/backticks.\\n"; + oss << "3) Return only the final assistant answer text unless a tool call is required.\\n"; + oss << "4) If a tool call is required, return ONLY JSON with this exact shape:\\n"; + oss << "[{\"name\":\"tool_name\",\"arguments\":{\"arg\":\"value\"}}]\\n"; + oss << "5) Do not include any prose before or after that JSON tool-call output.\\n"; + oss << "6) Never write a tool call as a function-call literal such as " + "tool_name({...}) or tool_name(arg=value); only the JSON array shape above is accepted.\\n"; + oss << "\\nConversation:\\n"; + for (const auto& m : request.messages) { + if (m.role == "tool") { + oss << "[tool:" << (m.name.empty() ? "result" : m.name) << "] " << m.content << "\\n"; + continue; + } + oss << "[" << m.role << "] " << m.content; + for (const auto& tc : m.tool_calls) { + oss << " (called tool " << tc.name << " with arguments " << tc.arguments << ")"; + } + oss << "\\n"; + } + + if (!request.tools.empty()) { + oss << "\\nAvailable tools JSON (use only these tool names and arguments):\\n"; + oss << serialize_tools_json(request.tools) << "\\n"; + oss << "If tools are relevant, prefer the strict JSON tool-call output contract above.\\n"; + } + + if (!request.local_output.empty()) { + oss << "\\nLocal model draft (useful fallback reference, may be low confidence):\\n"; + oss << request.local_output << "\\n"; + } + + return oss.str(); +} + +static void extract_textual_tool_calls(std::string& response, + const std::vector& tools, + std::vector& function_calls) { + if (tools.empty()) return; + for (const auto& tool : tools) { + const std::string& name = tool.name; + if (name.empty()) continue; + size_t search = 0; + while (true) { + size_t pos = response.find(name, search); + if (pos == std::string::npos) break; + bool left_ok = (pos == 0) || + !(std::isalnum(static_cast(response[pos - 1])) || response[pos - 1] == '_'); + size_t paren = pos + name.size(); + while (paren < response.size() && + std::isspace(static_cast(response[paren]))) { + ++paren; + } + if (!left_ok || paren >= response.size() || response[paren] != '(') { + search = pos + name.size(); + continue; + } + size_t args_end = find_matching_delimiter(response, paren, '(', ')'); + if (args_end > response.size()) { + search = pos + name.size(); + continue; + } + std::string inner = trim_string(response.substr(paren + 1, args_end - paren - 2)); + std::string args_json; + if (inner.empty()) { + args_json = "{}"; + } else if (inner.front() == '{' && inner.back() == '}') { + args_json = inner; + } else { + search = pos + name.size(); + continue; + } + function_calls.push_back("{\"name\":\"" + name + "\",\"arguments\":" + args_json + "}"); + response.erase(pos, args_end - pos); + search = pos; + } + } + if (!function_calls.empty()) response = trim_string(response); +} + +static void extract_called_tool_phrases(std::string& response, + const std::vector& tools, + std::vector& function_calls) { + const std::string marker = "called tool "; + size_t search = 0; + while (true) { + size_t pos = response.find(marker, search); + if (pos == std::string::npos) break; + size_t name_start = pos + marker.size(); + size_t name_end = name_start; + while (name_end < response.size() && + (std::isalnum(static_cast(response[name_end])) || response[name_end] == '_')) { + ++name_end; + } + std::string name = response.substr(name_start, name_end - name_start); + bool known = false; + for (const auto& t : tools) { + if (t.name == name) { known = true; break; } + } + if (name.empty() || !known) { search = pos + marker.size(); continue; } + size_t with = response.find("with arguments", name_end); + size_t brace = (with == std::string::npos) ? std::string::npos : response.find('{', with); + if (brace == std::string::npos) { search = pos + marker.size(); continue; } + size_t end = find_matching_delimiter(response, brace, '{', '}'); + if (end > response.size()) { search = pos + marker.size(); continue; } + std::string args = trim_string(response.substr(brace, end - brace)); + function_calls.push_back("{\"name\":\"" + name + "\",\"arguments\":" + args + "}"); + size_t erase_end = end; + while (erase_end < response.size() && + std::isspace(static_cast(response[erase_end]))) { + ++erase_end; + } + if (erase_end < response.size() && response[erase_end] == ')') ++erase_end; + size_t erase_start = (pos > 0 && response[pos - 1] == '(') ? pos - 1 : pos; + response.erase(erase_start, erase_end - erase_start); + search = erase_start; + } + if (!function_calls.empty()) response = trim_string(response); +} + +static std::string call_cloud_endpoint(const std::string& url, + const std::string& payload, + long timeout_ms, + const char* cloud_key_param, + std::string& err_out) { + std::string api_key = resolve_cloud_api_key(cloud_key_param); + if (api_key.empty()) { + if (!g_warned_missing_cloud_api_key.exchange(true)) { + CACTUS_LOG_WARN("cloud_handoff", "No cloud key found (cloud_key param, CACTUS_CLOUD_KEY, or CACTUS_CLOUD_API_KEY env); cloud handoff will fall back to local output"); + } + err_out = "missing_api_key"; + return {}; + } + + CURL* curl = curl_easy_init(); + if (!curl) { + err_out = "curl_init_failed"; + return {}; + } + + std::string response_body; + struct curl_slist* headers = nullptr; + headers = curl_slist_append(headers, ("X-API-Key: " + api_key).c_str()); + headers = curl_slist_append(headers, "Content-Type: application/json"); + + const char* extra_hdrs = std::getenv("CACTUS_CLOUD_HEADERS"); + if (extra_hdrs && extra_hdrs[0] != '\0') { + std::istringstream stream(extra_hdrs); + std::string pair; + while (std::getline(stream, pair, ',')) { + auto eq = pair.find('='); + if (eq != std::string::npos && eq > 0) { + std::string header = pair.substr(0, eq) + ": " + pair.substr(eq + 1); + headers = curl_slist_append(headers, header.c_str()); + } + } + } + + curl_easy_setopt(curl, CURLOPT_URL, url.c_str()); + curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers); + curl_easy_setopt(curl, CURLOPT_POSTFIELDS, payload.c_str()); + curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, static_cast(payload.size())); + curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, curl_write_cb); + curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response_body); + curl_easy_setopt(curl, CURLOPT_TIMEOUT_MS, timeout_ms); + curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT_MS, std::min(timeout_ms, 2000L)); + curl_easy_setopt(curl, CURLOPT_NOSIGNAL, 1L); + + if (!env_flag_enabled("CACTUS_CLOUD_STRICT_SSL")) { + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 0L); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0L); + } else { + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L); + apply_curl_tls_trust(curl); + } + + CURLcode res = curl_easy_perform(curl); + long http_status = 0; + curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &http_status); + curl_slist_free_all(headers); + curl_easy_cleanup(curl); + + if (res != CURLE_OK) { + err_out = curl_easy_strerror(res); + return {}; + } + + if (http_status >= 400) { + std::string detail = json_string_field(response_body, "error"); + if (detail.empty()) detail = json_string_field(response_body, "message"); + if (detail.empty()) detail = json_string_field(response_body, "detail"); + detail = compact_error_detail(detail); + + err_out = "http_" + std::to_string(http_status); + if (!detail.empty()) { + err_out += ":" + detail; + } + return {}; + } + + return response_body; +} +#endif + +} // namespace + +std::string cloud_base64_encode(const uint8_t* data, size_t len) { + static const char table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + std::string out; + out.reserve(((len + 2) / 3) * 4); + for (size_t i = 0; i < len; i += 3) { + uint32_t n = static_cast(data[i]) << 16; + if (i + 1 < len) n |= static_cast(data[i + 1]) << 8; + if (i + 2 < len) n |= static_cast(data[i + 2]); + out += table[(n >> 18) & 0x3F]; + out += table[(n >> 12) & 0x3F]; + out += (i + 1 < len) ? table[(n >> 6) & 0x3F] : '='; + out += (i + 2 < len) ? table[n & 0x3F] : '='; + } + return out; +} + +std::vector cloud_build_wav(const uint8_t* pcm, size_t pcm_bytes) { + constexpr uint32_t sample_rate = 16000; + constexpr uint16_t channels = 1; + constexpr uint16_t bits = 16; + const uint32_t byte_rate = sample_rate * channels * bits / 8; + const uint16_t block_align = channels * bits / 8; + const uint32_t data_size = static_cast(pcm_bytes); + const uint32_t file_size = 36 + data_size; + + std::vector wav(44 + pcm_bytes); + auto w16 = [&](size_t off, uint16_t v) { + wav[off] = v & 0xFF; + wav[off + 1] = v >> 8; + }; + auto w32 = [&](size_t off, uint32_t v) { + wav[off] = v & 0xFF; + wav[off + 1] = (v >> 8) & 0xFF; + wav[off + 2] = (v >> 16) & 0xFF; + wav[off + 3] = (v >> 24) & 0xFF; + }; + + std::memcpy(wav.data(), "RIFF", 4); + w32(4, file_size); + std::memcpy(wav.data() + 8, "WAVE", 4); + std::memcpy(wav.data() + 12, "fmt ", 4); + w32(16, 16); + w16(20, 1); + w16(22, channels); + w32(24, sample_rate); + w32(28, byte_rate); + w16(32, block_align); + w16(34, bits); + std::memcpy(wav.data() + 36, "data", 4); + w32(40, data_size); + std::memcpy(wav.data() + 44, pcm, pcm_bytes); + return wav; +} + +CloudResponse cloud_transcribe_request(const std::string& audio_b64, + const std::string& fallback_text, + long timeout_seconds, + const char* cloud_key) { + if (env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF")) { + return {fallback_text, "", false, "cloud_handoff_disabled"}; + } +#ifdef CACTUS_USE_CURL + std::string base = env_or_default("CACTUS_CLOUD_API_BASE", "https://104.198.76.3/api/v1"); + std::string endpoint = base + "/transcribe"; + + std::string payload = "{\"audio\":\"" + audio_b64 + "\",\"mime_type\":\"audio/wav\",\"language\":\"en-US\"}"; + + std::string err; + std::string body = call_cloud_endpoint(endpoint, payload, timeout_seconds * 1000L, cloud_key, err); + if (body.empty()) { + return {fallback_text, "", false, err.empty() ? "request_failed" : err}; + } + + std::string transcript = json_string_field(body, "transcript"); + if (transcript.empty()) { + transcript = json_string_field(body, "text"); + } + if (transcript.empty()) { + transcript = json_string_field(body, "response"); + } + if (transcript.empty()) { + transcript = json_string_field(body, "analysis"); + } + if (transcript.empty()) { + std::string detail = json_string_field(body, "error"); + if (detail.empty()) detail = json_string_field(body, "message"); + if (detail.empty()) detail = json_string_field(body, "detail"); + detail = compact_error_detail(detail); + if (!detail.empty()) { + return {fallback_text, "", false, "missing_transcript:" + detail}; + } + return {fallback_text, "", false, "missing_transcript"}; + } + + std::string api_key_hash = json_string_field(body, "api_key_hash"); + return {transcript, api_key_hash, true, ""}; +#else + (void)audio_b64; + (void)timeout_seconds; + (void)cloud_key; + return {fallback_text, "", false, "curl_not_enabled"}; +#endif +} + +CloudCompletionResult cloud_complete_request(const CloudCompletionRequest& request, + long timeout_ms) { + if (env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF")) { + CloudCompletionResult result; + result.ok = false; + result.error = "cloud_handoff_disabled"; + return result; + } +#ifdef CACTUS_USE_CURL + std::string base = env_or_default("CACTUS_CLOUD_API_BASE", "https://104.198.76.3/api/v1"); + std::string model = env_or_default("CACTUS_CLOUD_MODEL", "gemini-2.5-flash"); + + std::string endpoint; + std::string payload; + + std::string img_b64; + std::string img_mime; + if (request.has_images) { + for (const auto& message : request.messages) { + for (const auto& image : message.images) { + if (!image.empty()) { + std::vector img_bytes; + if (read_file_bytes(image, img_bytes)) { + img_b64 = cloud_base64_encode(img_bytes.data(), img_bytes.size()); + img_mime = infer_mime_type(image); + } + break; + } + } + if (!img_b64.empty()) break; + } + } + + std::string audio_b64; + if (request.has_audio && !request.audio_pcm.empty()) { + auto wav = cloud_build_wav(request.audio_pcm.data(), request.audio_pcm.size()); + audio_b64 = cloud_base64_encode(wav.data(), wav.size()); + } + + std::string text_prompt = build_cloud_text_prompt(request); + bool use_omni = !audio_b64.empty() || (request.has_images && request.has_audio); + + if (use_omni) { + endpoint = base + "/omni"; + payload = "{"; + payload += "\"text\":\"" + escape_json_string(text_prompt) + "\""; + if (!img_b64.empty()) { + payload += ",\"image\":\"" + img_b64 + "\""; + payload += ",\"image_mime_type\":\"" + img_mime + "\""; + } + if (!audio_b64.empty()) { + payload += ",\"audio\":\"" + audio_b64 + "\""; + payload += ",\"audio_mime_type\":\"audio/wav\""; + } + payload += ",\"model\":\"" + escape_json_string(model) + "\""; + payload += ",\"language\":\"en-US\""; + payload += "}"; + } else if (request.has_images && !img_b64.empty()) { + endpoint = base + "/vlm"; + payload = "{" + "\"image\":\"" + img_b64 + "\"," + "\"mime_type\":\"" + img_mime + "\"," + "\"prompt\":\"" + escape_json_string(text_prompt) + "\"," + "\"language\":\"en-US\"," + "\"model\":\"" + escape_json_string(model) + "\"" + "}"; + } else { + endpoint = base + "/text"; + payload = "{" + "\"text\":\"" + escape_json_string(text_prompt) + "\"," + "\"language\":\"en-US\"," + "\"model\":\"" + escape_json_string(model) + "\"" + "}"; + } + + std::string err; + const char* cloud_key = request.cloud_key.empty() ? nullptr : request.cloud_key.c_str(); + std::string body = call_cloud_endpoint(endpoint, payload, timeout_ms, cloud_key, err); + if (body.empty()) { + return {false, false, "", {}, err.empty() ? "request_failed" : err}; + } + + std::vector function_calls; + std::string calls_json = json_array_field(body, "function_calls"); + if (calls_json != "[]") { + function_calls = split_json_array(calls_json); + } + + std::string response = json_string_field(body, "text"); + if (response.empty()) { + response = json_string_field(body, "analysis"); + } + + if (!response.empty() && function_calls.empty()) { + std::string regular_response; + parse_function_calls_from_response(response, regular_response, function_calls); + response = regular_response; + } + + if (!response.empty() && function_calls.empty()) { + extract_textual_tool_calls(response, request.tools, function_calls); + } + + if (!response.empty() && function_calls.empty()) { + extract_called_tool_phrases(response, request.tools, function_calls); + } + + sanitize_function_calls(function_calls); + + if (response.empty() && function_calls.empty()) { + return {false, false, "", {}, "missing_text"}; + } + + CloudCompletionResult out; + out.ok = true; + out.used_cloud = true; + out.response = response; + out.function_calls = std::move(function_calls); + return out; +#else + (void)request; + (void)timeout_ms; + return {false, false, "", {}, "curl_not_enabled"}; +#endif +} + +} // namespace ffi +} // namespace cactus diff --git a/cactus-engine/src/cloud.h b/cactus-engine/src/cloud.h new file mode 100644 index 000000000..ed1b7ac05 --- /dev/null +++ b/cactus-engine/src/cloud.h @@ -0,0 +1,50 @@ +#ifndef CACTUS_CLOUD_H +#define CACTUS_CLOUD_H + +#include "utils.h" +#include +#include + +namespace cactus { +namespace ffi { + +struct CloudResponse { + std::string transcript; + std::string api_key_hash; + bool used_cloud = false; + std::string error; +}; + +struct CloudCompletionRequest { + std::vector messages; + std::vector tools; + std::string local_output; + std::vector local_function_calls; + bool has_images = false; + bool has_audio = false; + std::vector audio_pcm; + std::string cloud_key; +}; + +struct CloudCompletionResult { + bool ok = false; + bool used_cloud = false; + std::string response; + std::vector function_calls; + std::string error; +}; + +std::string cloud_base64_encode(const uint8_t* data, size_t len); +std::vector cloud_build_wav(const uint8_t* pcm, size_t pcm_bytes); +std::string resolve_cloud_api_key(const char* cloud_key_param); +CloudResponse cloud_transcribe_request(const std::string& audio_b64, + const std::string& fallback_text, + long timeout_seconds = 15L, + const char* cloud_key = nullptr); +CloudCompletionResult cloud_complete_request(const CloudCompletionRequest& request, + long timeout_ms); + +} // namespace ffi +} // namespace cactus + +#endif // CACTUS_CLOUD_H diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp new file mode 100644 index 000000000..0821cf9bb --- /dev/null +++ b/cactus-engine/src/complete.cpp @@ -0,0 +1,1557 @@ +#include "../cactus_engine.h" +#include "cloud.h" +#include "utils.h" +#include "chat_tools.h" +#include "telemetry.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include "wav.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace cactus::engine; +using namespace cactus::ffi; + +namespace { + +std::vector> extract_schema_property_types(const std::string& schema); +std::vector extract_schema_required(const std::string& schema); +std::vector extract_string_array(const std::string& json, const std::string& key); +std::string extract_json_object_field(const std::string& json, const std::string& key); + +std::string extract_last_user_query(const std::vector& messages) { + for (auto it = messages.rbegin(); it != messages.rend(); ++it) { + if (it->role == "user") { + return it->content; + } + } + return {}; +} + +void inject_rag_context(CactusModelHandle* handle, std::vector& messages) { + if (!handle->corpus_index) return; + + std::string query = extract_last_user_query(messages); + if (query.empty()) return; + + std::string rag_context = retrieve_rag_context(handle, query); + if (rag_context.empty()) return; + + if (!messages.empty() && messages[0].role == "system") { + messages[0].content = rag_context + messages[0].content; + } else { + ChatMessage system_msg; + system_msg.role = "system"; + system_msg.content = rag_context + "Answer the user's question using ONLY the context above. Do not use any prior knowledge. If the answer cannot be found in the context, respond with \"I don't have enough information to answer that.\""; + messages.insert(messages.begin(), system_msg); + } +} + +std::vector build_tool_constraint_specs(const std::vector& tools) { + std::vector specs; + specs.reserve(tools.size()); + + for (const auto& tool : tools) { + ToolConstraintSpec spec; + spec.name = tool.name; + + auto schema_it = tool.parameters.find("schema"); + if (schema_it != tool.parameters.end()) { + auto properties = extract_schema_property_types(schema_it->second); + std::string properties_object = extract_json_object_field(schema_it->second, "properties"); + spec.parameter_names.reserve(properties.size()); + spec.parameter_enums.reserve(properties.size()); + for (const auto& [name, _] : properties) { + spec.parameter_names.push_back(name); + spec.parameter_enums.push_back( + extract_string_array(extract_json_object_field(properties_object, name), "enum")); + } + spec.required_parameter_names = extract_schema_required(schema_it->second); + } + + specs.push_back(std::move(spec)); + } + + return specs; +} + +void setup_tool_constraints(CactusModelHandle* handle, const std::vector& tools, + bool force_tools, float& temperature) { + if (!force_tools || tools.empty()) return; + + handle->model->set_tool_constraints(build_tool_constraint_specs(tools)); + + if (temperature == 0.0f) { + temperature = 0.01f; + } +} + +size_t find_json_block_end(const std::string& json, size_t start) { + if (start >= json.size() || json[start] != '{') { + return std::string::npos; + } + + int depth = 1; + bool in_string = false; + bool escaped = false; + size_t pos = start + 1; + while (pos < json.size() && depth > 0) { + char c = json[pos]; + if (in_string) { + if (escaped) { + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + in_string = false; + } + } else { + if (c == '"') { + in_string = true; + } else if (c == '{') { + depth++; + } else if (c == '}') { + depth--; + } + } + ++pos; + } + + return depth == 0 ? pos : std::string::npos; +} + +std::string extract_json_object_field(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":"; + size_t key_pos = json.find(pattern); + if (key_pos == std::string::npos) { + return {}; + } + + size_t object_start = json.find('{', key_pos + pattern.size()); + if (object_start == std::string::npos) { + return {}; + } + + size_t object_end = find_json_block_end(json, object_start); + if (object_end == std::string::npos) { + return {}; + } + + return json.substr(object_start, object_end - object_start); +} + +std::vector> extract_schema_property_types(const std::string& schema) { + std::vector> properties; + std::string properties_object = extract_json_object_field(schema, "properties"); + if (properties_object.empty() || properties_object.size() < 2) { + return properties; + } + + size_t pos = 1; + while (pos + 1 < properties_object.size()) { + size_t key_start = properties_object.find('"', pos); + if (key_start == std::string::npos || key_start + 1 >= properties_object.size()) { + break; + } + size_t key_end = properties_object.find('"', key_start + 1); + if (key_end == std::string::npos) { + break; + } + + std::string name = properties_object.substr(key_start + 1, key_end - key_start - 1); + size_t value_start = properties_object.find('{', key_end); + if (value_start == std::string::npos) { + break; + } + size_t value_end = find_json_block_end(properties_object, value_start); + if (value_end == std::string::npos) { + break; + } + + std::string value = properties_object.substr(value_start, value_end - value_start); + std::string type = "string"; + std::string type_pattern = "\"type\":\""; + size_t type_pos = value.find(type_pattern); + if (type_pos != std::string::npos) { + size_t type_start = type_pos + type_pattern.size(); + size_t type_end = value.find('"', type_start); + if (type_end != std::string::npos) { + type = value.substr(type_start, type_end - type_start); + } + } else if (value.find("\"enum\"") != std::string::npos) { + type = "string"; + } else if (value.find("\"properties\"") != std::string::npos) { + type = "object"; + } + + properties.emplace_back(std::move(name), std::move(type)); + pos = value_end; + } + + return properties; +} + +std::vector extract_string_array(const std::string& json, const std::string& key) { + std::vector values; + if (json.empty()) return values; + std::string pattern = "\"" + key + "\""; + size_t key_pos = json.find(pattern); + if (key_pos == std::string::npos) return values; + size_t arr_start = json.find('[', key_pos + pattern.size()); + if (arr_start == std::string::npos) return values; + size_t arr_end = json.find(']', arr_start); + if (arr_end == std::string::npos) return values; + size_t pos = arr_start + 1; + while (pos < arr_end) { + size_t qs = json.find('"', pos); + if (qs == std::string::npos || qs >= arr_end) break; + size_t qe = json.find('"', qs + 1); + if (qe == std::string::npos || qe > arr_end) break; + values.push_back(json.substr(qs + 1, qe - qs - 1)); + pos = qe + 1; + } + return values; +} + +std::vector extract_schema_required(const std::string& schema) { + return extract_string_array(schema, "required"); +} + +static bool json_object_has_key(const std::string& obj, const std::string& key) { + const std::string target = "\"" + key + "\""; + int depth = 0; + bool in_str = false, esc = false; + for (size_t i = 0; i < obj.size(); ++i) { + char c = obj[i]; + if (in_str) { + if (esc) esc = false; + else if (c == '\\') esc = true; + else if (c == '"') in_str = false; + continue; + } + if (c == '"') { + if (depth == 1 && obj.compare(i, target.size(), target) == 0) { + size_t j = i + target.size(); + while (j < obj.size() && std::isspace(static_cast(obj[j]))) ++j; + if (j < obj.size() && obj[j] == ':') return true; + } + in_str = true; + continue; + } + if (c == '{' || c == '[') ++depth; + else if (c == '}' || c == ']') --depth; + } + return false; +} + +static bool function_calls_missing_required(const std::vector& calls, + const std::vector& tools) { + if (calls.empty() || tools.empty()) return false; + for (const auto& call : calls) { + std::string name = json_string_field(call, "name"); + const ToolFunction* tool = nullptr; + for (const auto& t : tools) { + if (t.name == name) { tool = &t; break; } + } + if (!tool) continue; + auto schema_it = tool->parameters.find("schema"); + if (schema_it == tool->parameters.end()) continue; + std::vector required = extract_schema_required(schema_it->second); + if (required.empty()) continue; + std::string args = extract_json_object_field(call, "arguments"); + for (const auto& req : required) { + if (!json_object_has_key(args, req)) return true; + } + } + return false; +} + +std::string serialize_needle_tools(const std::vector& tools) { + if (tools.empty()) return "[]"; + + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < tools.size(); ++i) { + if (i > 0) oss << ","; + + oss << "{\"name\":\"" << escape_json_string(tools[i].name) << "\""; + if (!tools[i].description.empty()) { + oss << ",\"description\":\"" << escape_json_string(tools[i].description) << "\""; + } + + oss << ",\"parameters\":"; + auto schema_it = tools[i].parameters.find("schema"); + if (schema_it == tools[i].parameters.end()) { + oss << "{}"; + } else { + auto properties = extract_schema_property_types(schema_it->second); + auto required = extract_schema_required(schema_it->second); + std::unordered_set required_set(required.begin(), required.end()); + + if (properties.empty()) { + oss << "{}"; + } else { + std::string properties_object = extract_json_object_field(schema_it->second, "properties"); + oss << "{"; + for (size_t p = 0; p < properties.size(); ++p) { + if (p > 0) oss << ","; + const std::string& param_name = properties[p].first; + const std::string& param_type = properties[p].second; + bool is_required = required_set.count(param_name) > 0; + + std::string description; + std::string prop_obj = extract_json_object_field(properties_object, param_name); + if (!prop_obj.empty()) { + std::string desc_pattern = "\"description\":\""; + size_t desc_pos = prop_obj.find(desc_pattern); + if (desc_pos != std::string::npos) { + size_t desc_start = desc_pos + desc_pattern.size(); + size_t desc_end = desc_start; + bool escaped = false; + while (desc_end < prop_obj.size()) { + if (escaped) { + escaped = false; + desc_end++; + continue; + } + if (prop_obj[desc_end] == '\\') { + escaped = true; + desc_end++; + continue; + } + if (prop_obj[desc_end] == '"') break; + desc_end++; + } + description = prop_obj.substr(desc_start, desc_end - desc_start); + } + } + + oss << "\"" << escape_json_string(param_name) << "\":{"; + oss << "\"type\":\"" << escape_json_string(param_type) << "\""; + if (!description.empty()) { + oss << ",\"description\":\"" << description << "\""; + } + oss << ",\"required\":" << (is_required ? "true" : "false") << "}"; + } + oss << "}"; + } + } + oss << "}"; + } + oss << "]"; + return oss.str(); +} + +std::vector> build_stop_sequences( + Tokenizer* tokenizer, + const std::vector& stop_sequences, + Config::ModelType model_type, + bool has_tools +) { + std::vector> stop_token_sequences; + stop_token_sequences.push_back({tokenizer->get_eos_token()}); + + std::vector sequences = stop_sequences; + if (sequences.empty()) { + std::string default_stop = tokenizer->get_default_stop_sequence(); + if (!default_stop.empty()) { + sequences.push_back(default_stop); + } + } + for (const auto& stop_seq : sequences) { + stop_token_sequences.push_back(tokenizer->encode(stop_seq)); + } + + if (model_type == Config::ModelType::GEMMA4) { + stop_token_sequences.push_back(tokenizer->encode("")); + if (has_tools) { + stop_token_sequences.push_back(tokenizer->encode("<|tool_response>")); + } + } + + return stop_token_sequences; +} + +void trim_stop_suffix(std::vector& generated_tokens, + const std::vector>& stop_token_sequences, + bool include_stop_sequences) { + if (include_stop_sequences) return; + for (const auto& stop_seq : stop_token_sequences) { + if (stop_seq.empty()) continue; + if (generated_tokens.size() >= stop_seq.size() && + std::equal(stop_seq.rbegin(), stop_seq.rend(), generated_tokens.rbegin())) { + generated_tokens.resize(generated_tokens.size() - stop_seq.size()); + break; + } + } +} + +void reset_cache(CactusModelHandle* handle) { + handle->model->reset_cache(); + handle->processed_tokens.clear(); + handle->processed_images.clear(); + handle->user_audio_counts.clear(); +} + +struct PrefillResult { + std::vector remaining_tokens; + size_t prefilled_count = 0; + bool was_prefix = false; + bool was_exact_match = false; +}; + +struct EntropyState { + float total_sum = 0.0f; + size_t total_count = 0; + + void add(float entropy) { + total_sum += entropy; + total_count++; + } + + float mean_confidence() const { + return 1.0f - (total_sum / static_cast(total_count)); + } +}; + +struct PreparedPrompt { + InferenceOptions options; + Config::ModelType model_type = Config::ModelType::GEMMA4; + std::vector image_paths; + std::vector audio_paths; + std::vector messages; + std::vector tools; + std::string rendered; + std::vector tokens; + size_t context_token_count = 0; + std::vector> images; + + std::vector> audio_features; + size_t audio_num_frames = 0; + + bool has_images() const { + return std::any_of(images.begin(), images.end(), + [](const auto& msg_imgs) { return !msg_imgs.empty(); }); + } + + bool has_audio() const { + return std::any_of(audio_features.begin(), audio_features.end(), + [](const auto& mel) { return !mel.empty(); }); + } +}; + +CactusModelHandle::ProcessedImage image_signature(const std::string& image_path) { + std::filesystem::path normalized_path(image_path); + std::error_code ec; + + auto absolute_path = std::filesystem::absolute(normalized_path, ec); + if (!ec) { + normalized_path = absolute_path; + } + + CactusModelHandle::ProcessedImage image; + image.path = normalized_path.string(); + + ec.clear(); + auto status = std::filesystem::status(normalized_path, ec); + if (!ec && std::filesystem::is_regular_file(status)) { + std::error_code time_ec; + auto mtime = std::filesystem::last_write_time(normalized_path, time_ec); + if (!time_ec) { + image.last_modified_timestamp = static_cast(mtime.time_since_epoch().count()); + } + } + + return image; +} + +std::vector> images_from_message(const std::vector& messages) { + std::vector> message_signatures; + message_signatures.reserve(messages.size()); + + for (const auto& message : messages) { + std::vector image_signatures; + image_signatures.reserve(message.images.size()); + for (const auto& image_path : message.images) { + image_signatures.push_back(image_signature(image_path)); + } + message_signatures.push_back(std::move(image_signatures)); + } + + return message_signatures; +} + +bool image_context_prefix_matches( + const std::vector>& prefix, + const std::vector>& full +) { + return prefix.size() <= full.size() && + std::equal(prefix.begin(), prefix.end(), full.begin()); +} + +bool prompt_context_matches( + const CactusModelHandle* handle, + const PreparedPrompt& prompt +) { + if (handle->processed_tokens.empty()) { + return false; + } + if (handle->model && + handle->model->get_config().model_type == Config::ModelType::NEEDLE && + prompt.tokens.size() != handle->processed_tokens.size() + 1) { + return false; + } + if (prompt.context_token_count < handle->processed_tokens.size()) { + return false; + } + if (!std::equal(handle->processed_tokens.begin(), handle->processed_tokens.end(), prompt.tokens.begin())) { + return false; + } + if (prompt.has_images()) { + return image_context_prefix_matches(handle->processed_images, prompt.images); + } + return !prompt.has_images(); +} + +PreparedPrompt prepare_prompt( + CactusModelHandle* handle, + const char* messages_json, + const char* options_json, + const char* tools_json, + bool apply_tool_constraints, + bool add_generation_prompt, + const uint8_t* pcm_buffer = nullptr, + size_t pcm_buffer_size = 0 +) { + if (!handle || !handle->model) { + throw std::runtime_error("Invalid model handle"); + } + + PreparedPrompt prompt; + prompt.options = parse_inference_options_json(options_json ? options_json : ""); + prompt.messages = parse_messages_json(messages_json, prompt.image_paths, &prompt.audio_paths); + if (prompt.messages.empty()) { + throw std::runtime_error("No messages provided"); + } + + inject_rag_context(handle, prompt.messages); + + if (tools_json && std::strlen(tools_json) > 0) { + prompt.tools = parse_tools_json(tools_json); + } + + if (prompt.options.tool_rag_top_k > 0 && prompt.tools.size() > prompt.options.tool_rag_top_k) { + std::string query = extract_last_user_query(prompt.messages); + if (!query.empty()) { + prompt.tools = select_relevant_tools(handle, query, prompt.tools, prompt.options.tool_rag_top_k); + } + } + + auto* tokenizer = handle->model->get_tokenizer(); + if (!tokenizer) { + throw std::runtime_error("Tokenizer unavailable"); + } + + prompt.model_type = handle->model->get_config().model_type; + + if (prompt.options.confidence_threshold < 0.0f) { + if (handle->model->has_handoff_probe()) { + // The Gemma4 probe returns p_wrong; confidence is 1 - p_wrong. + // Route when the probe is less than 50% confident in local output. + prompt.options.confidence_threshold = 0.50f; + } else { + float model_default = handle->model->get_config().default_cloud_handoff_threshold; + prompt.options.confidence_threshold = (model_default > 0.0f) ? model_default : 0.7f; + } + } + + if (prompt.model_type == Config::ModelType::GEMMA4) { + std::vector user_indices; + for (size_t i = 0; i < prompt.messages.size(); i++) { + if (prompt.messages[i].role == "user") user_indices.push_back(i); + } + auto& counts = handle->user_audio_counts; + if (counts.size() < user_indices.size()) counts.resize(user_indices.size(), 0); + + if (pcm_buffer != nullptr && pcm_buffer_size > 1 && !user_indices.empty()) { + auto waveform_fp32 = cactus::audio::pcm_buffer_to_float_samples(pcm_buffer, pcm_buffer_size); + auto samples_16k = resample_to_16k_fp32(waveform_fp32, 16000); + if (!samples_16k.empty()) { + auto audio_prep = cactus::audio::preprocess_audio_for_gemma4(samples_16k, handle->model->get_config()); + prompt.audio_features.push_back(std::move(audio_prep.features)); + size_t u = user_indices.size() - 1; + prompt.messages[user_indices[u]].audio_soft_token_count = audio_prep.num_soft_tokens; + counts[u] = audio_prep.num_soft_tokens; + } + } else if (!prompt.audio_paths.empty()) { + for (size_t u = 0; u < user_indices.size(); u++) { + const auto& msg = prompt.messages[user_indices[u]]; + if (msg.audio.empty()) continue; + const std::string& audio_path = msg.audio.back(); + AudioFP32 wav = load_wav(audio_path); + auto samples_16k = resample_to_16k_fp32(wav.samples, wav.sample_rate); + if (samples_16k.empty()) continue; + auto audio_prep = cactus::audio::preprocess_audio_for_gemma4(samples_16k, handle->model->get_config()); + prompt.audio_features.push_back(std::move(audio_prep.features)); + prompt.messages[user_indices[u]].audio_soft_token_count = audio_prep.num_soft_tokens; + counts[u] = audio_prep.num_soft_tokens; + } + } + } + + if (prompt.model_type == Config::ModelType::NEEDLE && !prompt.tools.empty()) { + prompt.options.force_tools = true; + } + + // Build the tool-definition block in the format each model family was trained on. + std::string formatted_tools; + if (prompt.model_type == Config::ModelType::NEEDLE) { + formatted_tools = serialize_needle_tools(prompt.tools); + } else if (tokenizer->is_qwen_family()) { + formatted_tools = chat_tools::serialize_tools_qwen(prompt.tools); + } else if (tokenizer->is_lfm2_family()) { + formatted_tools = chat_tools::serialize_tools_lfm2(prompt.tools); + } else { + formatted_tools = gemma::format_tools(prompt.tools, !Config::is_gemma3_family(prompt.model_type)); + } + + if (apply_tool_constraints) { + setup_tool_constraints(handle, prompt.tools, prompt.options.force_tools, prompt.options.temperature); + } + + prompt.rendered = tokenizer->format_chat_prompt( + prompt.messages, + add_generation_prompt, + formatted_tools, + prompt.options.enable_thinking_if_supported + ); + if (prompt.rendered.find("ERROR:") == 0) { + throw std::runtime_error(prompt.rendered.substr(6)); + } + prompt.tokens = tokenizer->encode(prompt.rendered); + prompt.context_token_count = prompt.tokens.size(); + prompt.images = images_from_message(prompt.messages); + return prompt; +} + +PrefillResult do_prefill( + CactusModelHandle* handle, + const PreparedPrompt& prompt, + const std::vector& target_tokens +) { + PrefillResult result = {}; + bool has_images = prompt.has_images(); + bool has_audio = prompt.has_audio(); + + result.was_prefix = prompt_context_matches(handle, prompt); + result.was_exact_match = result.was_prefix && + target_tokens.size() == handle->processed_tokens.size(); + + if (result.was_exact_match) { + return result; + } + + std::vector tokens_to_process; + if (!result.was_prefix) { + reset_cache(handle); + tokens_to_process = target_tokens; + } else { + tokens_to_process.assign( + target_tokens.begin() + handle->processed_tokens.size(), + target_tokens.end() + ); + } + + if (tokens_to_process.size() > 1) { + std::vector prefill_tokens(tokens_to_process.begin(), tokens_to_process.end() - 1); + result.prefilled_count = prefill_tokens.size(); + + auto slice_delta_audio = [&]() -> std::vector> { + if (!result.was_prefix) return prompt.audio_features; + const size_t cached_msg_count = handle->processed_images.size(); + size_t cached_audio_count = 0; + for (size_t i = 0; i < cached_msg_count && i < prompt.messages.size(); ++i) { + const auto& msg = prompt.messages[i]; + if (msg.role == "user" && !msg.audio.empty()) cached_audio_count++; + } + cached_audio_count = std::min(cached_audio_count, prompt.audio_features.size()); + return std::vector>( + prompt.audio_features.begin() + cached_audio_count, + prompt.audio_features.end()); + }; + + if (has_images && has_audio) { + std::vector delta_image_paths; + if (result.was_prefix) { + size_t cached_image_count = 0; + for (const auto& msg_imgs : handle->processed_images) { + cached_image_count += msg_imgs.size(); + } + delta_image_paths.assign( + prompt.image_paths.begin() + cached_image_count, + prompt.image_paths.end() + ); + } else { + delta_image_paths = prompt.image_paths; + } + handle->model->prefill_with_media(prefill_tokens, delta_image_paths, slice_delta_audio()); + } else if (has_images) { + std::vector delta_image_paths; + if (result.was_prefix) { + size_t cached_image_count = 0; + for (const auto& msg_imgs : handle->processed_images) { + cached_image_count += msg_imgs.size(); + } + delta_image_paths.assign( + prompt.image_paths.begin() + cached_image_count, + prompt.image_paths.end() + ); + } else { + delta_image_paths = prompt.image_paths; + } + handle->model->prefill_with_images(prefill_tokens, delta_image_paths); + } else if (has_audio) { + handle->model->prefill_with_audio(prefill_tokens, slice_delta_audio()); + } else { + handle->model->prefill(prefill_tokens, handle->model->get_prefill_chunk_size()); + } + result.remaining_tokens = {tokens_to_process.back()}; + } else { + result.remaining_tokens = tokens_to_process; + } + + return result; +} + +uint32_t decode( + std::unique_ptr& model, + const std::vector& tokens, + const InferenceOptions& options, + float* out_entropy +) { + return model->decode(tokens, options.temperature, options.top_p, options.top_k, + "", out_entropy, options.min_p, options.repetition_penalty); +} + +uint32_t generate_first_token( + CactusModelHandle* handle, + const PrefillResult& prefill_result, + const PreparedPrompt& prompt, + float* first_token_entropy +) { + if (prefill_result.was_exact_match || prefill_result.remaining_tokens.empty()) { + if (handle->processed_tokens.empty()) { + throw std::runtime_error("Cannot generate from empty prompt"); + } + return decode(handle->model, {handle->processed_tokens.back()}, prompt.options, first_token_entropy); + } + return decode(handle->model, prefill_result.remaining_tokens, prompt.options, first_token_entropy); +} + +std::string construct_prefill_response_json( + bool success, + const std::string* error, + size_t prefill_tokens, + double prefill_tps, + double total_time_ms +) { + std::ostringstream json; + json << "{"; + json << "\"success\":" << (success ? "true" : "false") << ","; + if (error) { + json << "\"error\":\"" << escape_json_string(*error) << "\","; + } else { + json << "\"error\":null,"; + } + json << "\"prefill_tokens\":" << prefill_tokens << ","; + json << "\"prefill_tps\":" << std::fixed << std::setprecision(2) << prefill_tps << ","; + json << "\"total_time_ms\":" << std::fixed << std::setprecision(2) << total_time_ms << ","; + json << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb(); + json << "}"; + return json.str(); +} + +} // anonymous namespace + +extern "C" { + +int cactus_complete( + cactus_model_t model, + const char* messages_json, + char* response_buffer, + size_t buffer_size, + const char* options_json, + const char* tools_json, + cactus_token_callback callback, + void* user_data, + const uint8_t* pcm_buffer, + size_t pcm_buffer_size +) { + struct MetalTrimGuard { + ~MetalTrimGuard() { cactus_metal_trim_prefill_cache(); } + } metal_trim_guard; + + if (!model) { + std::string error_msg = last_error_message.empty() ? + "Model not initialized. Check model path and files." : last_error_message; + CACTUS_LOG_ERROR("complete", error_msg); + handle_error_response(error_msg, response_buffer, buffer_size); + return -1; + } + + if (!messages_json || !response_buffer || buffer_size == 0) { + CACTUS_LOG_ERROR("complete", "Invalid parameters: messages_json, response_buffer, or buffer_size"); + handle_error_response("Invalid parameters", response_buffer, buffer_size); + return -1; + } + + try { + CactusThreading::prepare_current_thread_for_cactus_work(); + auto start_time = std::chrono::high_resolution_clock::now(); + + auto* handle = static_cast(model); + handle->should_stop = false; + if (!handle->model->can_generate()) { + const std::string err = "Model has no language-model decoder; cannot generate (embedding/encoder-only model)"; + CACTUS_LOG_ERROR("complete", err); + handle_error_response(err, response_buffer, buffer_size); + return -1; + } + auto* tokenizer = handle->model->get_tokenizer(); + auto prompt = prepare_prompt(handle, messages_json, options_json, tools_json, true, true, pcm_buffer, pcm_buffer_size); + if (prompt.options.sample_seed != 0) { + handle->model->set_sample_seed(prompt.options.sample_seed); + } + + CACTUS_LOG_DEBUG("complete", "Prompt tokens: " << prompt.tokens.size() + << ", max_tokens: " << prompt.options.max_tokens); + + bool has_images = prompt.has_images(); + bool has_audio = prompt.has_audio(); + const bool cloud_disabled = env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF"); + const bool cloud_eligible = !cloud_disabled && !handle->cloud_handoff_disabled && + prompt.options.auto_handoff && (!has_images || prompt.options.handoff_with_images); + handle->model->reset_handoff_probe_rollout(); + const bool defer_local_stream_until_probe = cloud_eligible && handle->model->has_handoff_probe(); + bool pre_generation_cloud_attempted = false; + + std::string handoff_reason; + if (!prompt.options.auto_handoff) { + handoff_reason = "handoff off"; + } else if (cloud_disabled) { + handoff_reason = "disabled (env)"; + } else if (handle->cloud_handoff_disabled) { + handoff_reason = "disabled (auth failed earlier)"; + } else if (has_images && !prompt.options.handoff_with_images) { + handoff_reason = "images kept local"; + } + auto friendly_cloud_error = [](const std::string& e) -> std::string { + if (e == "missing_api_key") return "no API key"; + if (e.rfind("http_401", 0) == 0) return "auth failed (401)"; + if (e.rfind("http_403", 0) == 0) return "auth failed (403)"; + if (e == "cloud_handoff_disabled") return "disabled"; + if (e == "curl_init_failed") return "curl init failed"; + if (e.empty()) return "cloud error"; + return e; + }; + + auto make_cloud_request = [&](const std::string& local_output_hint, + const std::vector& local_calls_hint) { + CloudCompletionRequest request; + request.messages = prompt.messages; + request.tools = prompt.tools; + request.local_output = local_output_hint; + request.local_function_calls = local_calls_hint; + request.has_images = has_images; + request.has_audio = has_audio; + if (has_audio && pcm_buffer != nullptr && pcm_buffer_size > 0) { + request.audio_pcm.assign(pcm_buffer, pcm_buffer + pcm_buffer_size); + } + request.cloud_key = resolve_cloud_api_key(nullptr); + return request; + }; + + auto disable_handoff_on_auth_failure = [&](const std::string& error) { + if (error.rfind("http_401", 0) == 0 || error.rfind("http_403", 0) == 0) { + handle->cloud_handoff_disabled = true; + CACTUS_LOG_WARN("cloud_handoff", "Cloud auth failed (" << error + << "); disabling cloud handoff for this session"); + } + }; + + auto return_cloud_completion = [&](const CloudCompletionResult& cloud_result, + double ttft_ms, + double total_ms, + float confidence, + size_t prompt_token_count, + const std::string& reason) { + std::string cloud_response = cloud_result.response; + std::vector cloud_calls = cloud_result.function_calls; + if (callback && !cloud_response.empty()) { + callback(cloud_response.c_str(), 0, user_data); + } + std::string result = construct_response_json(cloud_response, cloud_calls, ttft_ms, + total_ms, 0.0, 0.0, prompt_token_count, + 0, confidence, true, "", {}, "", + prompt.options.confidence_threshold, reason); + if (result.length() >= buffer_size) { + handle_error_response("Response buffer too small", response_buffer, buffer_size); + return -1; + } + std::strcpy(response_buffer, result.c_str()); + + cactus::telemetry::CompletionMetrics metrics{}; + metrics.success = true; + metrics.cloud_handoff = true; + metrics.ttft_ms = ttft_ms; + metrics.prefill_tps = 0.0; + metrics.decode_tps = 0.0; + metrics.response_time_ms = total_ms; + metrics.confidence = confidence; + metrics.ram_usage_mb = get_ram_usage_mb(); + metrics.prefill_tokens = prompt_token_count; + metrics.decode_tokens = 0; + metrics.error_message = nullptr; + metrics.function_calls_json = nullptr; + cactus::telemetry::recordCompletion(handle->model_name.c_str(), metrics); + return static_cast(result.length()); + }; + + if (cloud_eligible && prompt.options.confidence_threshold >= 1.0f) { + pre_generation_cloud_attempted = true; + CACTUS_LOG_INFO("cloud_handoff", "Cloud handoff triggered before local generation; waiting up to " + << prompt.options.cloud_timeout_ms << " ms before falling back"); + auto cloud_result = cloud_complete_request( + make_cloud_request("", {}), + static_cast(prompt.options.cloud_timeout_ms)); + auto now = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration_cast(now - start_time).count() / 1000.0; + if (cloud_result.ok && (!cloud_result.response.empty() || !cloud_result.function_calls.empty())) { + return return_cloud_completion(cloud_result, elapsed_ms, elapsed_ms, 0.0f, prompt.tokens.size(), + "forced (threshold 100%)"); + } + std::string cloud_error = cloud_result.error.empty() ? "cloud completion failed" : cloud_result.error; + CACTUS_LOG_WARN("cloud_handoff", "Cloud completion failed before local generation: " << cloud_error); + disable_handoff_on_auth_failure(cloud_error); + handle_error_response(("cloud handoff failed before local generation: " + cloud_error).c_str(), + response_buffer, buffer_size); + return -1; + } + + auto stop_token_sequences = build_stop_sequences(tokenizer, prompt.options.stop_sequences, prompt.model_type, !prompt.tools.empty()); + + std::vector generated_tokens; + double time_to_first_token = 0.0; + float first_token_entropy = 0.0f; + uint32_t next_token; + size_t prompt_tokens; + + if (has_audio && !handle->processed_tokens.empty()) { + auto& cache = handle->processed_tokens; + size_t common = 0; + size_t limit = std::min(cache.size(), prompt.tokens.size()); + while (common < limit && cache[common] == prompt.tokens[common]) common++; + if (common < cache.size()) { + CACTUS_LOG_WARN("complete", "KV cache diverges from new prompt at position " << common + << "/" << cache.size() << "; resetting cache for clean audio re-prefill"); + reset_cache(handle); + } + } + + bool first_token_from_prefill = false; + if (!has_images && !has_audio && handle->processed_tokens.empty()) { + reset_cache(handle); + first_token_from_prefill = handle->model->prefill_and_sample_first_token(prompt.tokens, next_token, &first_token_entropy); + if (first_token_from_prefill) { + prompt_tokens = prompt.tokens.size(); + } + } + if (!first_token_from_prefill) { + auto prefill_result = do_prefill(handle, prompt, prompt.tokens); + prompt_tokens = prefill_result.prefilled_count + prefill_result.remaining_tokens.size(); + next_token = generate_first_token(handle, prefill_result, prompt, &first_token_entropy); + } + + handle->processed_tokens = prompt.tokens; + handle->processed_images = prompt.images; + + auto token_end = std::chrono::high_resolution_clock::now(); + time_to_first_token = std::chrono::duration_cast(token_end - start_time).count() / 1000.0; + + float confidence = 1.0f - first_token_entropy; + std::string cloud_error; + + generated_tokens.push_back(next_token); + handle->processed_tokens.push_back(next_token); + + if (prompt.options.force_tools && !prompt.tools.empty()) { + handle->model->update_tool_constraints(next_token); + } + + EntropyState entropy; + entropy.add(first_token_entropy); + + if (!matches_stop_sequence(generated_tokens, stop_token_sequences)) { + if (cloud_eligible + && !defer_local_stream_until_probe + && !pre_generation_cloud_attempted + && confidence < prompt.options.confidence_threshold) { + CACTUS_LOG_INFO("cloud_handoff", "Cloud handoff triggered before local streaming; waiting up to " + << prompt.options.cloud_timeout_ms << " ms before falling back"); + CloudCompletionResult cloud_result = cloud_complete_request( + make_cloud_request("", {}), + static_cast(prompt.options.cloud_timeout_ms)); + auto now = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration_cast(now - start_time).count() / 1000.0; + if (cloud_result.ok && (!cloud_result.response.empty() || !cloud_result.function_calls.empty())) { + if (prompt.options.force_tools && !prompt.tools.empty()) { + handle->model->clear_tool_constraints(); + } + return return_cloud_completion(cloud_result, elapsed_ms, elapsed_ms, confidence, prompt_tokens, + "low confidence"); + } + cloud_error = cloud_result.error.empty() ? "cloud completion failed" : cloud_result.error; + CACTUS_LOG_WARN("cloud_handoff", "Cloud completion failed before local streaming, falling back to local output: " << cloud_error); + disable_handoff_on_auth_failure(cloud_error); + handoff_reason = "handoff failed: " + friendly_cloud_error(cloud_error); + } + + if (callback && !defer_local_stream_until_probe) { + std::string new_text = tokenizer->decode({next_token}); + callback(new_text.c_str(), next_token, user_data); + } + + for (size_t i = 1; i < prompt.options.max_tokens; i++) { + if (handle->should_stop) break; + + float token_entropy = 0.0f; + if (has_audio) { + uint32_t last_token = handle->processed_tokens.empty() ? next_token : handle->processed_tokens.back(); + next_token = handle->model->decode_with_audio( + {last_token}, prompt.audio_features, + prompt.options.temperature, prompt.options.top_p, prompt.options.top_k, + "", &token_entropy, + prompt.options.min_p, prompt.options.repetition_penalty); + } else { + next_token = decode(handle->model, {next_token}, prompt.options, &token_entropy); + } + handle->processed_tokens.push_back(next_token); + generated_tokens.push_back(next_token); + + entropy.add(token_entropy); + + if (prompt.options.force_tools && !prompt.tools.empty()) { + handle->model->update_tool_constraints(next_token); + } + + if (matches_stop_sequence(generated_tokens, stop_token_sequences)) { + trim_stop_suffix(generated_tokens, stop_token_sequences, prompt.options.include_stop_sequences); + break; + } + + if (callback && !defer_local_stream_until_probe) { + std::string new_text = tokenizer->decode({next_token}); + callback(new_text.c_str(), next_token, user_data); + } + } + } else { + trim_stop_suffix(generated_tokens, stop_token_sequences, prompt.options.include_stop_sequences); + } + + if (defer_local_stream_until_probe) { + confidence = entropy.mean_confidence(); + if (handle->model->has_handoff_probe_rollout()) { + float wrong_probability = handle->model->handoff_probe_wrong_probability(); + if (std::isfinite(wrong_probability)) { + confidence = std::max(0.0f, std::min(1.0f, 1.0f - wrong_probability)); + CACTUS_LOG_DEBUG("cloud_handoff", "Gemma4 handoff probe p_wrong=" + << wrong_probability << " confidence=" << confidence); + } + } + } + + if (prompt.options.force_tools && !prompt.tools.empty()) { + handle->model->clear_tool_constraints(); + } + + auto end_time = std::chrono::high_resolution_clock::now(); + double total_time_ms = std::chrono::duration_cast(end_time - start_time).count() / 1000.0; + + size_t completion_tokens = generated_tokens.size(); + double decode_time_ms = total_time_ms - time_to_first_token; + double prefill_tps = time_to_first_token > 0 ? (prompt_tokens * 1000.0) / time_to_first_token : 0.0; + double decode_tps = (completion_tokens > 1 && decode_time_ms > 0) ? ((completion_tokens - 1) * 1000.0) / decode_time_ms : 0.0; + + std::string response_text = tokenizer->decode(generated_tokens); + + std::string regular_response; + std::vector function_calls; + parse_function_calls_from_response(response_text, regular_response, function_calls, prompt.tools); + + std::string thinking_text; + if (prompt.model_type == Config::ModelType::GEMMA4 || prompt.options.enable_thinking_if_supported) { + std::string stripped_content; + partition_thinking_response(regular_response, thinking_text, stripped_content); + regular_response = stripped_content; + if (!prompt.options.enable_thinking_if_supported) { + thinking_text.clear(); + } + } + + std::string local_completion = regular_response; + const bool parsed_empty_needle_tool_call = + prompt.model_type == Config::ModelType::NEEDLE && + response_text.find("") != std::string::npos; + if (local_completion.empty() && function_calls.empty() && !parsed_empty_needle_tool_call) { + local_completion = response_text; + } + std::string primary_response = local_completion; + std::vector primary_function_calls = function_calls; + + bool handoff_succeeded = false; + const bool low_confidence = defer_local_stream_until_probe + && confidence < prompt.options.confidence_threshold; + const bool invalid_local_tool_call = cloud_eligible && defer_local_stream_until_probe + && function_calls_missing_required(function_calls, prompt.tools); + if (!pre_generation_cloud_attempted && (low_confidence || invalid_local_tool_call)) { + const char* trigger_reason = invalid_local_tool_call && !low_confidence + ? "local tool call missing required args" : "low confidence (probe)"; + CACTUS_LOG_INFO("cloud_handoff", "Cloud handoff triggered (" << trigger_reason + << "): p_wrong=" << (1.0f - confidence) << " confidence=" << confidence + << "; waiting up to " << prompt.options.cloud_timeout_ms << " ms"); + CloudCompletionResult cloud_result = cloud_complete_request( + make_cloud_request(local_completion, function_calls), + static_cast(prompt.options.cloud_timeout_ms)); + auto now = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration_cast(now - start_time).count() / 1000.0; + if (cloud_result.ok && (!cloud_result.response.empty() || !cloud_result.function_calls.empty())) { + if (prompt.options.force_tools && !prompt.tools.empty()) { + handle->model->clear_tool_constraints(); + } + return return_cloud_completion(cloud_result, elapsed_ms, elapsed_ms, confidence, prompt_tokens, + trigger_reason); + } + cloud_error = cloud_result.error.empty() ? "cloud completion failed" : cloud_result.error; + CACTUS_LOG_WARN("cloud_handoff", "Cloud completion failed after probe handoff, falling back to local output: " + << cloud_error); + disable_handoff_on_auth_failure(cloud_error); + handoff_reason = "handoff failed: " + friendly_cloud_error(cloud_error); + } + + if (callback && defer_local_stream_until_probe && !primary_response.empty()) { + callback(primary_response.c_str(), 0, user_data); + } + + // No cloud handoff happened (success paths returned early above). If we never recorded a + // reason, the request was eligible and the local model cleared the threshold. + if (handoff_reason.empty()) { + handoff_reason = (confidence >= prompt.options.confidence_threshold) + ? "above threshold" + : "kept local"; + } + + if (prompt.options.force_tools && !prompt.tools.empty() && primary_function_calls.empty()) { + CACTUS_LOG_WARN("force_tools", "force_tools was set but no tool call was parsed from the constrained output"); + } + + std::string result = construct_response_json(primary_response, primary_function_calls, time_to_first_token, + total_time_ms, prefill_tps, decode_tps, prompt_tokens, + completion_tokens, confidence, handoff_succeeded, + thinking_text, {}, response_text, + prompt.options.confidence_threshold, handoff_reason); + + if (result.length() >= buffer_size) { + handle_error_response("Response buffer too small", response_buffer, buffer_size); + return -1; + } + + std::strcpy(response_buffer, result.c_str()); + + std::string function_calls_json = serialize_function_calls(primary_function_calls); + cactus::telemetry::CompletionMetrics metrics{}; + metrics.success = true; + metrics.cloud_handoff = handoff_succeeded; + metrics.ttft_ms = time_to_first_token; + metrics.prefill_tps = prefill_tps; + metrics.decode_tps = decode_tps; + metrics.response_time_ms = total_time_ms; + metrics.confidence = confidence; + metrics.ram_usage_mb = get_ram_usage_mb(); + metrics.prefill_tokens = prompt_tokens; + metrics.decode_tokens = completion_tokens; + metrics.error_message = nullptr; + metrics.function_calls_json = nullptr; + cactus::telemetry::recordCompletion(handle->model_name.c_str(), metrics); + + return static_cast(result.length()); + + } catch (const std::exception& e) { + CACTUS_LOG_ERROR("complete", "Exception: " << e.what()); + handle_error_response(e.what(), response_buffer, buffer_size); + + cactus::telemetry::CompletionMetrics metrics{}; + metrics.success = false; + metrics.cloud_handoff = false; + metrics.ttft_ms = 0.0; + metrics.prefill_tps = 0.0; + metrics.decode_tps = 0.0; + metrics.response_time_ms = 0.0; + metrics.confidence = 0.0; + metrics.ram_usage_mb = get_ram_usage_mb(); + metrics.prefill_tokens = 0; + metrics.decode_tokens = 0; + metrics.error_message = e.what(); + metrics.function_calls_json = nullptr; + auto* h = static_cast(model); + cactus::telemetry::recordCompletion(h ? h->model_name.c_str() : "unknown", metrics); + + return -1; + } catch (...) { + CACTUS_LOG_ERROR("complete", "Unknown exception during completion"); + handle_error_response("Unknown error during completion", response_buffer, buffer_size); + + cactus::telemetry::CompletionMetrics metrics{}; + metrics.success = false; + metrics.cloud_handoff = false; + metrics.ttft_ms = 0.0; + metrics.prefill_tps = 0.0; + metrics.decode_tps = 0.0; + metrics.response_time_ms = 0.0; + metrics.confidence = 0.0; + metrics.ram_usage_mb = get_ram_usage_mb(); + metrics.prefill_tokens = 0; + metrics.decode_tokens = 0; + metrics.error_message = "Unknown error during completion"; + metrics.function_calls_json = nullptr; + auto* h = static_cast(model); + cactus::telemetry::recordCompletion(h ? h->model_name.c_str() : "unknown", metrics); + + return -1; + } +} + +int cactus_prefill( + cactus_model_t model, + const char* messages_json, + char* response_buffer, + size_t buffer_size, + const char* options_json, + const char* tools_json, + const uint8_t* pcm_buffer, + size_t pcm_buffer_size +) { + if (!model) { + std::string error_msg = last_error_message.empty() + ? "Model not initialized. Check model path and files." + : last_error_message; + if (response_buffer && buffer_size > 0) { + std::string result = construct_prefill_response_json(false, &error_msg, 0, 0.0, 0.0); + if (result.size() < buffer_size) { + std::strcpy(response_buffer, result.c_str()); + } + } + return -1; + } + + if (!messages_json || !response_buffer || buffer_size == 0) { + std::string error_msg = "Invalid parameters"; + if (response_buffer && buffer_size > 0) { + std::string result = construct_prefill_response_json(false, &error_msg, 0, 0.0, 0.0); + if (result.size() < buffer_size) { + std::strcpy(response_buffer, result.c_str()); + } + } + return -1; + } + + try { + CactusThreading::prepare_current_thread_for_cactus_work(); + auto start_time = std::chrono::high_resolution_clock::now(); + + auto* handle = static_cast(model); + if (!handle->model->can_generate()) { + std::string err = "Model has no language-model decoder; cannot generate (embedding/encoder-only model)"; + CACTUS_LOG_ERROR("prefill", err); + std::string result = construct_prefill_response_json(false, &err, 0, 0.0, 0.0); + if (result.size() < buffer_size) std::strcpy(response_buffer, result.c_str()); + return -1; + } + auto prompt = prepare_prompt(handle, messages_json, options_json, tools_json, false, false, pcm_buffer, pcm_buffer_size); + + std::vector context_tokens(prompt.tokens.begin(), prompt.tokens.begin() + prompt.context_token_count); + auto prefill_result = do_prefill(handle, prompt, context_tokens); + + if (!prefill_result.was_exact_match) { + handle->processed_tokens = context_tokens; + if (!handle->processed_tokens.empty()) { + handle->processed_tokens.pop_back(); + } + } + handle->processed_images = prompt.images; + + auto end_time = std::chrono::high_resolution_clock::now(); + double elapsed_ms = std::chrono::duration_cast(end_time - start_time).count() / 1000.0; + double prefill_tps = (prefill_result.prefilled_count > 0 && elapsed_ms > 0.0) + ? (static_cast(prefill_result.prefilled_count) * 1000.0) / elapsed_ms + : 0.0; + + std::string result = construct_prefill_response_json(true, nullptr, prefill_result.prefilled_count, prefill_tps, elapsed_ms); + if (result.size() >= buffer_size) { + std::string error_msg = "Response buffer too small"; + std::string error_json = construct_prefill_response_json(false, &error_msg, 0, 0.0, 0.0); + if (error_json.size() < buffer_size) { + std::strcpy(response_buffer, error_json.c_str()); + } + return -1; + } + + std::strcpy(response_buffer, result.c_str()); + return static_cast(result.size()); + } catch (const std::exception& e) { + std::string error_msg = e.what(); + std::string result = construct_prefill_response_json(false, &error_msg, 0, 0.0, 0.0); + if (result.size() < buffer_size) { + std::strcpy(response_buffer, result.c_str()); + } + return -1; + } catch (...) { + std::string error_msg = "Unknown error during prefill"; + std::string result = construct_prefill_response_json(false, &error_msg, 0, 0.0, 0.0); + if (result.size() < buffer_size) { + std::strcpy(response_buffer, result.c_str()); + } + return -1; + } +} + +int cactus_tokenize( + cactus_model_t model, + const char* text, + uint32_t* token_buffer, + size_t token_buffer_len, + size_t* out_token_len +) { + if (!model || !text || !out_token_len) return -1; + + try { + auto* handle = static_cast(model); + auto* tokenizer = handle->model->get_tokenizer(); + + std::vector toks = tokenizer->encode(std::string(text)); + *out_token_len = toks.size(); + + if (!token_buffer || token_buffer_len == 0) return 0; + if (token_buffer_len < toks.size()) return -2; + + std::memcpy(token_buffer, toks.data(), toks.size() * sizeof(uint32_t)); + return 0; + } catch (...) { + return -1; + } +} + +int cactus_render_prompt( + cactus_model_t model, + const char* messages_json, + const char* options_json, + const char* tools_json, + char* prompt_buffer, + size_t buffer_size +) { + if (!model || !messages_json || !prompt_buffer || buffer_size == 0) return -1; + + try { + auto* handle = static_cast(model); + auto prompt = prepare_prompt(handle, messages_json, options_json, tools_json, false, true); + if (prompt.rendered.size() >= buffer_size) return -2; + std::strcpy(prompt_buffer, prompt.rendered.c_str()); + return static_cast(prompt.rendered.size()); + } catch (...) { + return -1; + } +} + +int cactus_score_window( + cactus_model_t model, + const uint32_t* tokens, + size_t token_len, + size_t start, + size_t end, + size_t context, + char* response_buffer, + size_t buffer_size +) { + if (!model || !tokens || token_len == 0 || !response_buffer || buffer_size == 0) { + handle_error_response("Invalid parameters", response_buffer, buffer_size); + return -1; + } + + try { + CactusThreading::prepare_current_thread_for_cactus_work(); + auto* handle = static_cast(model); + + std::vector vec(tokens, tokens + token_len); + + size_t scored = 0; + double logprob = handle->model->score_tokens_window_logprob(vec, start, end, context, &scored); + + std::ostringstream oss; + oss << "{" + << "\"success\":true," + << "\"logprob\":" << std::setprecision(10) << logprob << "," + << "\"tokens\":" << scored + << "}"; + + std::string result = oss.str(); + if (result.size() >= buffer_size) { + handle_error_response("Response buffer too small", response_buffer, buffer_size); + return -1; + } + + std::strcpy(response_buffer, result.c_str()); + return (int)result.size(); + + } catch (const std::exception& e) { + handle_error_response(e.what(), response_buffer, buffer_size); + return -1; + } +} + +int cactus_benchmark_tokens( + cactus_model_t model, + const uint32_t* prompt_tokens, + size_t prompt_token_len, + size_t decode_token_len, + char* response_buffer, + size_t buffer_size +) { + if (!model || !prompt_tokens || prompt_token_len == 0 || !response_buffer || buffer_size == 0) { + handle_error_response("Invalid parameters", response_buffer, buffer_size); + return -1; + } + + try { + CactusThreading::prepare_current_thread_for_cactus_work(); + auto* handle = static_cast(model); + std::vector prompt(prompt_tokens, prompt_tokens + prompt_token_len); + + auto start_time = std::chrono::high_resolution_clock::now(); + double peak_ram_usage_mb = get_ram_usage_mb(); + auto sample_peak_ram = [&]() { + peak_ram_usage_mb = std::max(peak_ram_usage_mb, get_ram_usage_mb()); + }; + handle->model->reset_cache(); + auto prefill_start_time = std::chrono::high_resolution_clock::now(); + size_t cache_prime_tokens = 0; + uint32_t next_token = 0; + bool first_token_from_prefill = false; + if (decode_token_len == 0) { + handle->model->prefill(prompt, handle->model->get_prefill_chunk_size(), "", false); + cache_prime_tokens = prompt.size(); + } else { + first_token_from_prefill = handle->model->prefill_and_sample_first_token(prompt, next_token); + if (first_token_from_prefill) { + cache_prime_tokens = prompt.size(); + } else if (prompt.size() > 1) { + handle->model->prefill( + std::vector(prompt.begin(), prompt.end() - 1), + handle->model->get_prefill_chunk_size() + ); + cache_prime_tokens = prompt.size() - 1; + } + } + auto prefill_end_time = std::chrono::high_resolution_clock::now(); + sample_peak_ram(); + auto first_decode_start_time = std::chrono::high_resolution_clock::now(); + if (decode_token_len > 0 && !first_token_from_prefill) { + next_token = handle->model->decode({prompt.back()}, 0.0f, 1.0f, 1); + } + sample_peak_ram(); + auto first_token_time = std::chrono::high_resolution_clock::now(); + + std::vector completion_tokens; + if (decode_token_len > 0) completion_tokens.push_back(next_token); + size_t generated = decode_token_len > 0 ? 1 : 0; + while (generated < decode_token_len) { + next_token = handle->model->decode({next_token}, 0.0f, 1.0f, 1); + completion_tokens.push_back(next_token); + sample_peak_ram(); + ++generated; + } + + auto end_time = std::chrono::high_resolution_clock::now(); + double ttft_ms = std::chrono::duration_cast(first_token_time - start_time).count() / 1000.0; + double total_ms = std::chrono::duration_cast(end_time - start_time).count() / 1000.0; + double cache_prime_ms = std::chrono::duration_cast(prefill_end_time - prefill_start_time).count() / 1000.0; + double first_decode_ms = std::chrono::duration_cast(first_token_time - first_decode_start_time).count() / 1000.0; + double cache_state_copy_ms = handle->model->last_prefill_cache_copy_ms(); + double cache_prime_compute_ms = std::max(0.0, cache_prime_ms - cache_state_copy_ms); + double decode_ms = std::max(0.0, total_ms - ttft_ms); + double prefill_prepare_tps = (cache_prime_tokens > 0 && cache_prime_ms > 0.0) + ? (static_cast(cache_prime_tokens) * 1000.0) / cache_prime_ms + : 0.0; + double prefill_tps = (cache_prime_tokens > 0 && cache_prime_compute_ms > 0.0) + ? (static_cast(cache_prime_tokens) * 1000.0) / cache_prime_compute_ms + : 0.0; + double ttft_prompt_tps = ttft_ms > 0.0 ? (static_cast(prompt_token_len) * 1000.0) / ttft_ms : 0.0; + double decode_tps = (generated > 1 && decode_ms > 0.0) ? (static_cast(generated - 1) * 1000.0) / decode_ms : 0.0; + + std::ostringstream oss; + oss << "{" + << "\"success\":true," + << "\"time_to_first_token_ms\":" << std::fixed << std::setprecision(2) << ttft_ms << "," + << "\"total_time_ms\":" << std::fixed << std::setprecision(2) << total_ms << "," + << "\"prefill_tps\":" << std::fixed << std::setprecision(2) << prefill_tps << "," + << "\"prefill_compute_tps\":" << std::fixed << std::setprecision(2) << prefill_tps << "," + << "\"prefill_prepare_tps\":" << std::fixed << std::setprecision(2) << prefill_prepare_tps << "," + << "\"ttft_prompt_tps\":" << std::fixed << std::setprecision(2) << ttft_prompt_tps << "," + << "\"cache_prime_ms\":" << std::fixed << std::setprecision(2) << cache_prime_ms << "," + << "\"cache_prime_compute_ms\":" << std::fixed << std::setprecision(2) << cache_prime_compute_ms << "," + << "\"cache_state_copy_ms\":" << std::fixed << std::setprecision(2) << cache_state_copy_ms << "," + << "\"cache_prime_tokens\":" << cache_prime_tokens << "," + << "\"first_decode_ms\":" << std::fixed << std::setprecision(2) << first_decode_ms << "," + << "\"first_token_from_prefill\":" << (first_token_from_prefill ? "true" : "false") << "," + << "\"prefill_padding_tokens\":" << handle->model->last_prefill_padding_tokens() << "," + << "\"prefill_scalar_tail_tokens\":" << handle->model->last_prefill_scalar_tail_tokens() << "," + << "\"prefill_tail_chunk_tokens\":" << handle->model->last_prefill_tail_chunk_tokens() << "," + << "\"prefill_tail_padding_tokens\":" << handle->model->last_prefill_tail_padding_tokens() << "," + << "\"decode_tps\":" << std::fixed << std::setprecision(2) << decode_tps << "," + << "\"prompt_tokens\":" << prompt_token_len << "," + << "\"completion_tokens\":" << generated << "," + << "\"completion_token_ids\":["; + for (size_t i = 0; i < completion_tokens.size(); ++i) { + if (i) oss << ","; + oss << completion_tokens[i]; + } + oss << "]," + << "\"peak_ram_usage_mb\":" << std::fixed << std::setprecision(2) << peak_ram_usage_mb << "," + << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() + << "}"; + + std::string result = oss.str(); + if (result.size() >= buffer_size) { + handle_error_response("Response buffer too small", response_buffer, buffer_size); + return -1; + } + std::strcpy(response_buffer, result.c_str()); + return static_cast(result.size()); + } catch (const std::exception& e) { + handle_error_response(e.what(), response_buffer, buffer_size); + return -1; + } catch (...) { + handle_error_response("Unknown error during token benchmark", response_buffer, buffer_size); + return -1; + } +} + +} diff --git a/cactus-engine/src/constraints.cpp b/cactus-engine/src/constraints.cpp new file mode 100644 index 000000000..c966b7000 --- /dev/null +++ b/cactus-engine/src/constraints.cpp @@ -0,0 +1,648 @@ +#include "engine.h" + +namespace cactus { +namespace engine { + +constexpr float FORCE_BIAS = 500.0f; +constexpr float BLOCK_BIAS = -500.0f; +constexpr float TRIE_BLOCK_BIAS = -1e9f; + +static constexpr const char* kEscapeTag = "<|\"|>"; +static constexpr size_t kEscapeTagLen = 5; + +static bool ends_with_escape_tag(const std::string& text) { + return text.size() >= kEscapeTagLen && + text.compare(text.size() - kEscapeTagLen, kEscapeTagLen, kEscapeTag) == 0; +} + +void ToolCallConstrainer::add_tokens_for_string(const std::string& str, std::unordered_set& token_set) { + if (!tokenizer_) return; + auto tokens = tokenizer_->encode(str); + for (uint32_t t : tokens) { + token_set.insert(t); + } +} + +void ToolCallConstrainer::init_common_tokens() { + backtick_tokens_.clear(); + add_tokens_for_string("`", backtick_tokens_); + add_tokens_for_string("``", backtick_tokens_); + add_tokens_for_string("```", backtick_tokens_); + add_tokens_for_string("````", backtick_tokens_); + add_tokens_for_string("```json", backtick_tokens_); + add_tokens_for_string("```JSON", backtick_tokens_); + add_tokens_for_string("``` json", backtick_tokens_); + add_tokens_for_string("```\n", backtick_tokens_); + add_tokens_for_string("` ", backtick_tokens_); +} + +void ToolCallConstrainer::trie_insert(TrieNode* root, const std::string& word) { + if (!root) return; + TrieNode* node = root; + for (char ch : word) { + auto& child = node->children[ch]; + if (!child) child = std::make_unique(); + node = child.get(); + } + node->is_terminal = true; +} + +const ToolCallConstrainer::TrieNode* ToolCallConstrainer::trie_seek( + const TrieNode* root, + const std::string& prefix +) const { + const TrieNode* node = root; + for (char ch : prefix) { + if (!node) return nullptr; + auto it = node->children.find(ch); + if (it == node->children.end()) return nullptr; + node = it->second.get(); + } + return node; +} + +bool ToolCallConstrainer::trie_token_valid(const std::string& token_text, + const TrieNode* node, + char terminator) const { + if (!node) return false; + for (char ch : token_text) { + if (ch == terminator) return node->is_terminal; + auto it = node->children.find(ch); + if (it == node->children.end()) return false; + node = it->second.get(); + } + return true; +} + +const ToolCallConstrainer::TrieNode* ToolCallConstrainer::current_param_trie() const { + auto it = param_tries_.find(current_function_); + return it != param_tries_.end() ? it->second.get() : nullptr; +} + +void ToolCallConstrainer::mark_trie_bias(const TrieNode* node, char terminator, + const std::vector& extra_allowed, + int32_t terminal_token, bool exact_terminator) { + if (!node) return; + + std::vector valid(token_strings_.size(), false); + bool has_valid = false; + auto mark = [&](char first_char, bool unconditional) { + auto idx_it = token_index_.find(first_char); + if (idx_it == token_index_.end()) return; + for (uint32_t token_id : idx_it->second) { + if (valid[token_id]) continue; + if (unconditional || trie_token_valid(token_strings_[token_id], node, terminator)) { + valid[token_id] = true; + has_valid = true; + } + } + }; + + for (const auto& [first_char, _] : node->children) mark(first_char, false); + if (node->is_terminal) { + if (exact_terminator) { + auto idx_it = token_index_.find(terminator); + if (idx_it != token_index_.end()) { + for (uint32_t token_id : idx_it->second) { + if (!valid[token_id] && token_strings_[token_id].size() == 1) { + valid[token_id] = true; + has_valid = true; + } + } + } + } else { + mark(terminator, false); + } + } + for (char c : extra_allowed) mark(c, true); + if (node->is_terminal && terminal_token >= 0) { + uint32_t tid = static_cast(terminal_token); + if (tid < valid.size() && !valid[tid]) { + valid[tid] = true; + has_valid = true; + } + } + if (!has_valid) return; + + dense_bias_.assign(valid.size(), TRIE_BLOCK_BIAS); + for (size_t token_id = 0; token_id < valid.size(); ++token_id) { + if (valid[token_id]) dense_bias_[token_id] = 0.0f; + } + dense_ready_ = true; +} + +void ToolCallConstrainer::reset_constraint_state() { + region_ = Region::FREE; + constraint_buffer_.clear(); + constrained_buf_.clear(); + current_function_.clear(); + in_arguments_ = false; + arguments_depth_ = 0; + nesting_depth_ = 0; + in_string_value_ = false; + prev_char_escape_ = false; + at_key_start_ = false; + await_enum_open_ = false; + active_enum_trie_ = nullptr; + emitted_keys_.clear(); + remaining_key_trie_.reset(); +} + +void ToolCallConstrainer::rebuild_remaining_keys() { + remaining_key_trie_ = std::make_unique(); + for (const auto& tool : tool_specs_) { + if (tool.name != current_function_) continue; + for (const auto& param_name : tool.parameter_names) { + if (!param_name.empty() && !emitted_keys_.count(param_name)) { + trie_insert(remaining_key_trie_.get(), param_name); + } + } + break; + } +} + +bool ToolCallConstrainer::required_satisfied() const { + for (const auto& tool : tool_specs_) { + if (tool.name != current_function_) continue; + for (const auto& req : tool.required_parameter_names) { + if (!emitted_keys_.count(req)) return false; + } + return true; + } + return true; +} + +void ToolCallConstrainer::build_constraint_tables() { + name_trie_ = std::make_unique(); + param_tries_.clear(); + enum_tries_.clear(); + for (const auto& tool : tool_specs_) { + if (!tool.name.empty()) trie_insert(name_trie_.get(), tool.name); + + if (is_needle()) { + auto param_root = std::make_unique(); + for (const auto& param_name : tool.parameter_names) { + if (!param_name.empty()) trie_insert(param_root.get(), param_name); + } + param_tries_[tool.name] = std::move(param_root); + } + + for (size_t i = 0; i < tool.parameter_names.size() && i < tool.parameter_enums.size(); ++i) { + if (tool.parameter_enums[i].empty()) continue; + auto enum_root = std::make_unique(); + for (const auto& value : tool.parameter_enums[i]) { + if (!value.empty()) trie_insert(enum_root.get(), value); + } + enum_tries_[tool.name][tool.parameter_names[i]] = std::move(enum_root); + } + } + + const uint32_t vocab_size = tokenizer_ ? tokenizer_->get_vocab_size() : 0; + if (token_strings_.size() == vocab_size) return; + + token_strings_.assign(vocab_size, ""); + token_index_.clear(); + + const uint32_t eos = tokenizer_->get_eos_token(); + const uint32_t bos = tokenizer_->get_bos_token(); + const uint32_t unk = tokenizer_->get_unk_token(); + constexpr uint32_t pad = 0; + + for (uint32_t token_id = 0; token_id < vocab_size; ++token_id) { + if (token_id == pad || token_id == eos || token_id == bos || token_id == unk) continue; + std::string decoded = tokenizer_->decode({token_id}); + if (decoded.size() == 1) { + unsigned char b = static_cast(decoded[0]); + if (b >= 0x80) { + decoded.assign({static_cast(0xC0 | (b >> 6)), + static_cast(0x80 | (b & 0x3F))}); + } + } + token_strings_[token_id] = decoded; + if (!decoded.empty()) token_index_[decoded.front()].push_back(token_id); + } +} + +bool ToolCallConstrainer::needle_at_arg_key_start() const { + if (constraint_buffer_.size() < 2) return false; + return constraint_buffer_.compare(constraint_buffer_.size() - 2, 2, "{\"") == 0 || + constraint_buffer_.compare(constraint_buffer_.size() - 2, 2, ",\"") == 0; +} + +bool ToolCallConstrainer::needle_is_value_string_start() const { + if (constraint_buffer_.empty()) return false; + for (size_t i = constraint_buffer_.size() - 1; i-- > 0;) { + char ch = constraint_buffer_[i]; + if (ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r') continue; + return ch == ':'; + } + return false; +} + +void ToolCallConstrainer::feed_needle_char(char ch) { + if (region_ == Region::IN_NAME || region_ == Region::IN_ARG_KEY) { + if (ch == '"') { + if (region_ == Region::IN_NAME) current_function_ = constrained_buf_; + constrained_buf_.clear(); + region_ = Region::FREE; + } else { + constrained_buf_.push_back(ch); + } + constraint_buffer_.push_back(ch); + return; + } + + constraint_buffer_.push_back(ch); + + if (in_string_value_) { + if (prev_char_escape_) { + prev_char_escape_ = false; + return; + } + if (ch == '\\') { + prev_char_escape_ = true; + return; + } + if (ch == '"') in_string_value_ = false; + return; + } + + if (ch == '{' || ch == '[') { + nesting_depth_++; + } else if (ch == '}' || ch == ']') { + nesting_depth_ = std::max(0, nesting_depth_ - 1); + if (ch == '}' && in_arguments_ && nesting_depth_ < arguments_depth_) { + in_arguments_ = false; + } + } + + if (constraint_buffer_.size() >= 8 && !in_arguments_ && + constraint_buffer_.compare(constraint_buffer_.size() - 8, 8, "\"name\":\"") == 0) { + region_ = Region::IN_NAME; + constrained_buf_.clear(); + return; + } + + if (constraint_buffer_.size() >= 13 && + constraint_buffer_.compare(constraint_buffer_.size() - 13, 13, "\"arguments\":{") == 0) { + in_arguments_ = true; + arguments_depth_ = nesting_depth_; + return; + } + + if (in_arguments_ && + nesting_depth_ == arguments_depth_ && + needle_at_arg_key_start()) { + region_ = Region::IN_ARG_KEY; + constrained_buf_.clear(); + return; + } + + if (ch == '"' && needle_is_value_string_start()) { + in_string_value_ = true; + } +} + +void ToolCallConstrainer::feed_needle_text(const std::string& text) { + for (char ch : text) feed_needle_char(ch); +} + +void ToolCallConstrainer::feed_gemma_char(char ch) { + if (state_ != State::GEMMA_NAME && state_ != State::GEMMA_ARGS) return; + + constraint_buffer_.push_back(ch); + + if (region_ == Region::IN_NAME) { + if (ch == '{') { + current_function_ = constrained_buf_; + constrained_buf_.clear(); + region_ = Region::FREE; + state_ = State::GEMMA_ARGS; + in_arguments_ = true; + nesting_depth_ = 1; + arguments_depth_ = nesting_depth_; + at_key_start_ = true; + rebuild_remaining_keys(); + } else { + constrained_buf_.push_back(ch); + } + return; + } + + if (region_ == Region::IN_ARG_KEY) { + if (ch == ':') { + emitted_keys_.insert(constrained_buf_); + active_enum_trie_ = nullptr; + if (escape_tag_valid_) { + auto fit = enum_tries_.find(current_function_); + if (fit != enum_tries_.end()) { + auto kit = fit->second.find(constrained_buf_); + if (kit != fit->second.end()) active_enum_trie_ = kit->second.get(); + } + } + await_enum_open_ = active_enum_trie_ != nullptr; + rebuild_remaining_keys(); + constrained_buf_.clear(); + region_ = Region::FREE; + } else { + constrained_buf_.push_back(ch); + } + return; + } + + if (in_string_value_) { + if (ends_with_escape_tag(constraint_buffer_)) { + in_string_value_ = false; + active_enum_trie_ = nullptr; + constrained_buf_.clear(); + } else if (active_enum_trie_) { + constrained_buf_.push_back(ch); + } + return; + } + if (ends_with_escape_tag(constraint_buffer_)) { + in_string_value_ = true; + at_key_start_ = false; + await_enum_open_ = false; + constrained_buf_.clear(); + return; + } + if (ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r') return; + + if (at_key_start_ && ch != '}' && ch != ',') { + region_ = Region::IN_ARG_KEY; + constrained_buf_.clear(); + constrained_buf_.push_back(ch); + at_key_start_ = false; + return; + } + + if (ch == '{' || ch == '[') { + nesting_depth_++; + } else if (ch == '}' || ch == ']') { + nesting_depth_--; + if (ch == '}' && nesting_depth_ < arguments_depth_) { + in_arguments_ = false; + state_ = State::GEMMA_EXPECT_END; + forced_tag_progress_ = 0; + } + } else if (ch == ',' && nesting_depth_ == arguments_depth_) { + at_key_start_ = true; + } +} + +void ToolCallConstrainer::feed_gemma_text(const std::string& text) { + for (char ch : text) feed_gemma_char(ch); +} + +void ToolCallConstrainer::tokenize_grammar_elements() { + if (!tokenizer_) return; + + open_brace_tokens_.clear(); + close_brace_tokens_.clear(); + + init_common_tokens(); + + gemma_call_start_tokens_.clear(); + gemma_call_end_tokens_.clear(); + gemma_call_prefix_tokens_.clear(); + gemma_response_start_tokens_.clear(); + + call_start_sequence_ = tokenizer_->encode(call_start_tag_); + call_end_sequence_ = tokenizer_->encode(call_end_tag_); + std::vector response_sequence = tokenizer_->encode(response_start_tag_); + if (call_start_sequence_.size() == 1) gemma_call_start_tokens_.insert(call_start_sequence_[0]); + if (call_end_sequence_.size() == 1) gemma_call_end_tokens_.insert(call_end_sequence_[0]); + if (response_sequence.size() == 1) gemma_response_start_tokens_.insert(response_sequence[0]); + add_tokens_for_string("call:", gemma_call_prefix_tokens_); + + add_tokens_for_string("{", open_brace_tokens_); + add_tokens_for_string("}", close_brace_tokens_); + + std::vector escape_sequence = tokenizer_->encode(kEscapeTag); + escape_tag_valid_ = escape_sequence.size() == 1; + escape_tag_token_ = escape_tag_valid_ ? escape_sequence[0] : 0; +} + +void ToolCallConstrainer::init(Config::ModelType model_type, + const std::vector& tools, + Tokenizer* tokenizer) { + model_type_ = model_type; + tool_specs_ = tools; + tokenizer_ = tokenizer; + generated_text_.clear(); + current_bias_.clear(); + dense_ready_ = false; + forced_tag_progress_ = 0; + const bool model_supported = is_needle() || Config::is_gemma_family(model_type_); + active_ = model_supported && !tool_specs_.empty() && tokenizer != nullptr; + + reset_constraint_state(); + state_ = is_needle() ? State::NEEDLE_START : State::GEMMA_START; + if (Config::is_gemma3_family(model_type_)) { + call_start_tag_ = ""; + call_end_tag_ = ""; + response_start_tag_ = ""; + } else { + call_start_tag_ = "<|tool_call>"; + call_end_tag_ = ""; + response_start_tag_ = "<|tool_response>"; + } + + if (!active_) { + return; + } + + tokenize_grammar_elements(); + build_constraint_tables(); + compute_bias(); +} + +void ToolCallConstrainer::advance_forced_tag(const std::vector& sequence, uint32_t token_id) { + if (forced_tag_progress_ < sequence.size() && token_id == sequence[forced_tag_progress_]) { + forced_tag_progress_++; + } else { + forced_tag_progress_ = (!sequence.empty() && token_id == sequence[0]) ? 1 : 0; + } +} + +void ToolCallConstrainer::update(uint32_t token_id, const std::string& decoded_text) { + if (!active_) return; + + generated_text_ += decoded_text; + + if (is_needle()) { + feed_needle_text(decoded_text); + generated_text_.clear(); + compute_bias(); + return; + } + + switch (state_) { + case State::GEMMA_START: + advance_forced_tag(call_start_sequence_, token_id); + if (generated_text_.find(call_start_tag_) != std::string::npos) { + state_ = State::GEMMA_EXPECT_CALL; + generated_text_.clear(); + forced_tag_progress_ = 0; + } + break; + + case State::GEMMA_EXPECT_CALL: + if (generated_text_.find("call:") != std::string::npos) { + state_ = State::GEMMA_NAME; + generated_text_.clear(); + region_ = Region::IN_NAME; + constrained_buf_.clear(); + constraint_buffer_.clear(); + } + break; + + case State::GEMMA_NAME: + case State::GEMMA_ARGS: + feed_gemma_text(decoded_text); + generated_text_.clear(); + break; + + case State::GEMMA_EXPECT_END: + advance_forced_tag(call_end_sequence_, token_id); + if (generated_text_.find(call_end_tag_) != std::string::npos) { + state_ = State::DONE; + generated_text_.clear(); + } + break; + + default: + break; + } + + compute_bias(); +} + +void ToolCallConstrainer::compute_bias() { + current_bias_.clear(); + dense_ready_ = false; + + if (!active_) return; + + if (!is_needle()) { + for (uint32_t t : backtick_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + for (uint32_t t : gemma_response_start_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + } + + if (region_ != Region::FREE) { + const TrieNode* node = nullptr; + char terminator = '"'; + if (region_ == Region::IN_NAME) { + node = trie_seek(name_trie_.get(), constrained_buf_); + if (!is_needle()) terminator = '{'; + } else { + const TrieNode* root = is_needle() ? current_param_trie() : remaining_key_trie_.get(); + node = trie_seek(root, constrained_buf_); + if (!is_needle()) terminator = ':'; + } + mark_trie_bias(node, terminator, {}, -1, !is_needle()); + return; + } + + if (is_needle()) return; + + switch (state_) { + case State::GEMMA_START: + if (forced_tag_progress_ < call_start_sequence_.size()) { + current_bias_[call_start_sequence_[forced_tag_progress_]] = FORCE_BIAS; + } + for (uint32_t t : open_brace_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + for (uint32_t t : close_brace_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + break; + + case State::GEMMA_EXPECT_CALL: + for (uint32_t t : gemma_call_prefix_tokens_) { + current_bias_[t] = FORCE_BIAS; + } + for (uint32_t t : open_brace_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + for (uint32_t t : gemma_call_end_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + break; + + case State::GEMMA_ARGS: + if (await_enum_open_) { + current_bias_[escape_tag_token_] = FORCE_BIAS; + for (uint32_t t : gemma_call_end_tokens_) current_bias_[t] = BLOCK_BIAS; + for (uint32_t t : gemma_call_start_tokens_) current_bias_[t] = BLOCK_BIAS; + } else if (active_enum_trie_ && in_string_value_) { + mark_trie_bias(trie_seek(active_enum_trie_, constrained_buf_), '\0', {}, + static_cast(escape_tag_token_)); + } else if (at_key_start_) { + std::vector extra; + if (required_satisfied()) extra.push_back('}'); + mark_trie_bias(remaining_key_trie_.get(), ':', extra); + } else { + for (uint32_t t : gemma_call_end_tokens_) current_bias_[t] = BLOCK_BIAS; + for (uint32_t t : gemma_call_start_tokens_) current_bias_[t] = BLOCK_BIAS; + if (nesting_depth_ == arguments_depth_ && !in_string_value_ && !required_satisfied()) { + for (uint32_t t : close_brace_tokens_) current_bias_[t] = BLOCK_BIAS; + } + } + break; + + case State::GEMMA_EXPECT_END: + if (forced_tag_progress_ < call_end_sequence_.size()) { + current_bias_[call_end_sequence_[forced_tag_progress_]] = FORCE_BIAS; + } + for (uint32_t t : open_brace_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + for (uint32_t t : gemma_call_start_tokens_) { + current_bias_[t] = BLOCK_BIAS; + } + break; + + default: + break; + } +} + +void ToolCallConstrainer::reset() { + generated_text_.clear(); + current_bias_.clear(); + dense_ready_ = false; + forced_tag_progress_ = 0; + reset_constraint_state(); + + state_ = is_needle() ? State::NEEDLE_START : State::GEMMA_START; + + if (active_) { + compute_bias(); + } +} + + +void Model::set_tool_constraints(const std::vector& tools) { + tool_constrainer_.init(config_.model_type, tools, tokenizer_.get()); +} + +void Model::clear_tool_constraints() { + tool_constrainer_.init(config_.model_type, {}, tokenizer_.get()); +} + +void Model::update_tool_constraints(uint32_t token_id) { + if (tool_constrainer_.is_active() && tokenizer_) { + std::string decoded = tokenizer_->decode({token_id}); + tool_constrainer_.update(token_id, decoded); + } +} + +} // namespace engine +} // namespace cactus diff --git a/cactus/ffi/cactus_embed.cpp b/cactus-engine/src/embed.cpp similarity index 53% rename from cactus/ffi/cactus_embed.cpp rename to cactus-engine/src/embed.cpp index f1085a7d6..33e9daf7f 100644 --- a/cactus/ffi/cactus_embed.cpp +++ b/cactus-engine/src/embed.cpp @@ -1,72 +1,31 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" -#include "cactus_telemetry.h" -#include "../../libs/audio/wav.h" +#include "../cactus_engine.h" +#include "utils.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include "wav.h" #include -#include #include using namespace cactus::engine; using namespace cactus::ffi; -static constexpr size_t WHISPER_TARGET_FRAMES = 3000; -static constexpr int WHISPER_SAMPLE_RATE = 16000; - -static AudioProcessor::SpectrogramConfig get_whisper_spectrogram_config() { - AudioProcessor::SpectrogramConfig cfg{}; - cfg.n_fft = 400; - cfg.frame_length = 400; - cfg.hop_length = 160; - cfg.power = 2.0f; - cfg.center = true; - cfg.pad_mode = "reflect"; - cfg.onesided = true; - cfg.dither = 0.0f; - cfg.mel_floor = 1e-10f; - cfg.log_mel = "log10"; - cfg.reference = 1.0f; - cfg.min_value = 1e-10f; - cfg.remove_dc_offset = true; - return cfg; -} - -static std::vector compute_mel_from_wav(const std::string& wav_path) { +static std::vector compute_mel_from_wav(const std::string& wav_path, size_t mel_bins) { AudioFP32 audio = load_wav(wav_path); std::vector waveform_16k = resample_to_16k_fp32(audio.samples, audio.sample_rate); - auto cfg = get_whisper_spectrogram_config(); - const size_t num_mel_filters = 80; - const size_t num_frequency_bins = cfg.n_fft / 2 + 1; - - AudioProcessor ap; - ap.init_mel_filters(num_frequency_bins, num_mel_filters, 0.0f, 8000.0f, WHISPER_SAMPLE_RATE); - std::vector mel = ap.compute_spectrogram(waveform_16k, cfg); + auto cfg = cactus::audio::get_whisper_spectrogram_config(); + const size_t num_mel_filters = std::max(1, mel_bins); + const bool is_v3 = num_mel_filters > 80; + if (is_v3) cactus::audio::apply_whisper_v3_overrides(cfg); + int norm_type = is_v3 ? 1 : 1; + int scale_type = is_v3 ? 2 : 2; + cactus::audio::pad_or_trim_whisper_waveform(waveform_16k); + std::vector mel = cactus::audio::compute_spectrogram_graph( + waveform_16k, cfg, num_mel_filters, 0.0f, 8000.0f, 16000, norm_type, scale_type); if (mel.empty()) return mel; - size_t n_mels = num_mel_filters; - size_t n_frames = mel.size() / n_mels; - - float max_val = -std::numeric_limits::infinity(); - for (float v : mel) if (v > max_val) max_val = v; - - float min_allowed = max_val - 8.0f; - for (float& v : mel) { - if (v < min_allowed) v = min_allowed; - v = (v + 4.0f) / 4.0f; - } - - if (n_frames != WHISPER_TARGET_FRAMES) { - std::vector fixed(n_mels * WHISPER_TARGET_FRAMES, 0.0f); - size_t copy_frames = std::min(n_frames, WHISPER_TARGET_FRAMES); - for (size_t m = 0; m < n_mels; ++m) { - const float* src = &mel[m * n_frames]; - float* dst = &fixed[m * WHISPER_TARGET_FRAMES]; - std::copy(src, src + copy_frames, dst); - } - return fixed; - } - return mel; + return cactus::audio::normalize_whisper_mel(mel, num_mel_filters, is_v3); } extern "C" { @@ -79,6 +38,9 @@ int cactus_embed( size_t* embedding_dim, bool normalize ) { + struct MetalTrimGuard { + ~MetalTrimGuard() { cactus_metal_trim_prefill_cache(); } + } metal_trim_guard; if (!model || !text || !embeddings_buffer || buffer_size == 0) { CACTUS_LOG_ERROR("embed", "Invalid parameters for text embedding"); return -1; @@ -87,53 +49,40 @@ int cactus_embed( try { auto* handle = static_cast(model); auto* tokenizer = handle->model->get_tokenizer(); + if (!tokenizer) { + CACTUS_LOG_ERROR("embed", "Model tokenizer not available"); + return -1; + } - std::vector tokens = tokenizer->encode(text); + auto tokens = tokenizer->encode(text); if (tokens.empty()) { - CACTUS_LOG_ERROR("embed", "Tokenization produced empty result"); + CACTUS_LOG_ERROR("embed", "Failed to tokenize input text"); + return -1; + } + + std::vector embeddings = handle->model->get_text_embeddings(tokens, normalize); + if (embeddings.empty()) { + CACTUS_LOG_ERROR("embed", "Embedding returned empty result"); return -1; } - std::vector embeddings = handle->model->get_embeddings(tokens, true, normalize); if (embeddings.size() * sizeof(float) > buffer_size) { - CACTUS_LOG_ERROR("embed", "Buffer too small: need " << embeddings.size() * sizeof(float) << " bytes, got " << buffer_size); + CACTUS_LOG_ERROR("embed", "Buffer too small: need " << embeddings.size() * sizeof(float) << " bytes"); return -2; } std::memcpy(embeddings_buffer, embeddings.data(), embeddings.size() * sizeof(float)); if (embedding_dim) *embedding_dim = embeddings.size(); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - true, - "" - ); - return static_cast(embeddings.size()); } catch (const std::exception& e) { last_error_message = e.what(); CACTUS_LOG_ERROR("embed", "Exception: " << e.what()); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - false, - last_error_message - ); - return -1; } catch (...) { last_error_message = "Unknown error during embedding"; CACTUS_LOG_ERROR("embed", last_error_message); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - false, - last_error_message - ); - return -1; } } @@ -145,6 +94,9 @@ int cactus_image_embed( size_t buffer_size, size_t* embedding_dim ) { + struct MetalTrimGuard { + ~MetalTrimGuard() { cactus_metal_trim_prefill_cache(); } + } metal_trim_guard; if (!model || !image_path || !embeddings_buffer || buffer_size == 0) { CACTUS_LOG_ERROR("image_embed", "Invalid parameters for image embedding"); return -1; @@ -152,8 +104,6 @@ int cactus_image_embed( try { auto* handle = static_cast(model); - - CACTUS_LOG_DEBUG("image_embed", "Processing image: " << image_path); std::vector embeddings = handle->model->get_image_embeddings(image_path); if (embeddings.empty()) { CACTUS_LOG_ERROR("image_embed", "Image embedding returned empty result"); @@ -167,37 +117,15 @@ int cactus_image_embed( std::memcpy(embeddings_buffer, embeddings.data(), embeddings.size() * sizeof(float)); if (embedding_dim) *embedding_dim = embeddings.size(); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - true, - "" - ); - return static_cast(embeddings.size()); } catch (const std::exception& e) { last_error_message = e.what(); CACTUS_LOG_ERROR("image_embed", "Exception: " << e.what()); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - false, - last_error_message - ); - return -1; } catch (...) { last_error_message = "Unknown error during image embedding"; CACTUS_LOG_ERROR("image_embed", last_error_message); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - false, - last_error_message - ); - return -1; } } @@ -209,6 +137,9 @@ int cactus_audio_embed( size_t buffer_size, size_t* embedding_dim ) { + struct MetalTrimGuard { + ~MetalTrimGuard() { cactus_metal_trim_prefill_cache(); } + } metal_trim_guard; if (!model || !audio_path || !embeddings_buffer || buffer_size == 0) { CACTUS_LOG_ERROR("audio_embed", "Invalid parameters for audio embedding"); return -1; @@ -216,9 +147,7 @@ int cactus_audio_embed( try { auto* handle = static_cast(model); - - CACTUS_LOG_DEBUG("audio_embed", "Processing audio: " << audio_path); - auto mel_bins = compute_mel_from_wav(audio_path); + auto mel_bins = compute_mel_from_wav(audio_path, handle->model->get_config().num_mel_bins); if (mel_bins.empty()) { last_error_message = "Failed to compute mel spectrogram"; CACTUS_LOG_ERROR("audio_embed", last_error_message << " for: " << audio_path); @@ -238,37 +167,15 @@ int cactus_audio_embed( std::memcpy(embeddings_buffer, embeddings.data(), embeddings.size() * sizeof(float)); if (embedding_dim) *embedding_dim = embeddings.size(); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - true, - "" - ); - return static_cast(embeddings.size()); } catch (const std::exception& e) { last_error_message = e.what(); CACTUS_LOG_ERROR("audio_embed", "Exception: " << e.what()); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - false, - last_error_message - ); - return -1; } catch (...) { last_error_message = "Unknown error during audio embedding"; CACTUS_LOG_ERROR("audio_embed", last_error_message); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordEmbedding( - handle->model_name, - false, - last_error_message - ); - return -1; } } diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h new file mode 100644 index 000000000..73fc6d47a --- /dev/null +++ b/cactus-engine/src/engine.h @@ -0,0 +1,1202 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cactus_graph.h" +#include "kv_compress.h" + +class CactusGraph; + +namespace cactus { +namespace engine { + +struct Config { + uint32_t vocab_size = 151936; + uint32_t bos_token_id = 151643; + uint32_t eos_token_id = 151645; + uint32_t decoder_start_token_id = 151643; + std::vector decoder_prompt_token_ids; + uint32_t num_layers = 28; + uint32_t hidden_dim = 1024; + uint32_t ffn_intermediate_dim = 3072; + uint32_t attention_heads = 16; + uint32_t attention_kv_heads = 8; + uint32_t attention_head_dim = 128; + float layer_norm_eps = 1e-6f; + float rope_theta = 1000000.0f; + uint32_t num_experts = 0; + uint32_t num_shared_experts = 0; + uint32_t num_top_experts = 0; + uint32_t moe_every_n_layers = 0; + uint32_t moe_intermediate_dim = 0; + uint32_t num_dense_layers = 0; + uint32_t num_experts_per_tok = 0; + bool norm_topk_prob = false; + bool use_expert_bias = false; + float routed_scaling_factor = 1.0f; + bool tie_word_embeddings = true; + + uint32_t vision_hidden_dim = 0; + uint32_t vision_num_layers = 0; + uint32_t vision_attention_heads = 0; + uint32_t vision_image_size = 0; + uint32_t vision_patch_size = 0; + uint32_t vision_num_channels = 3; + uint32_t vision_embed_dim = 0; + uint32_t visual_tokens_per_img = 0; + bool use_pixel_shuffle = false; + uint32_t pixel_shuffle_factor = 1; + bool use_image_tokens = false; + uint32_t image_token_id = 0; + bool use_layout_tags = false; + uint32_t image_seq_len = 64; + + uint32_t global_image_size = 2048; + uint32_t max_tile_size = 512; + float rescale_factor = 0.00392156862745098f; + float image_mean = 0.5f; + float image_std = 0.5f; + + uint32_t downsample_factor = 2; + uint32_t min_tiles = 2; + uint32_t max_tiles = 10; + bool use_thumbnail = true; + uint32_t min_image_tokens = 64; + uint32_t max_image_tokens = 256; + uint32_t max_num_patches = 1024; + uint32_t tile_size = 512; + float max_pixels_tolerance = 2.0f; + bool do_image_splitting = true; + bool encoder_act_gelu = false; + bool decoder_act_gelu = false; + uint32_t num_encoder_layers = 0; + uint32_t num_decoder_layers = 0; + float partial_rotary_factor = 0.0f; + uint32_t pad_token_id = 0; + uint32_t conv_kernel_size = 0; + uint32_t subsampling_conv_kernel_size = 0; + uint32_t subsampling_conv_stride = 0; + uint32_t subsampling_conv_channels = 0; + uint32_t subsampling_factor = 0; + uint32_t num_mel_bins = 80; + std::string encoder_hidden_act = "silu"; + uint32_t linear_num_key_heads = 0; + uint32_t linear_key_head_dim = 0; + uint32_t linear_num_value_heads = 0; + uint32_t linear_value_head_dim = 0; + uint32_t linear_q_proj_dim = 0; + uint32_t linear_k_proj_dim = 0; + uint32_t linear_v_proj_dim = 0; + + uint32_t kv_lora_rank = 0; + uint32_t q_lora_rank = 0; + uint32_t qk_head_dim = 0; + uint32_t qk_nope_head_dim = 0; + uint32_t qk_rope_head_dim = 0; + uint32_t v_head_dim = 0; + uint32_t rope_interleave = 0; + bool attention_bias = false; + float rope_scaling_factor = 1.0f; + float rope_mscale_all_dim = 0.0f; + + enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18}; + uint32_t predictor_hidden_dim = 0; + uint32_t predictor_num_layers = 0; + uint32_t tdt_joint_dim = 0; + uint32_t tdt_num_durations = 0; + uint32_t tdt_blank_id = 0; + std::vector tdt_durations; + + ModelType model_type = ModelType::GEMMA4; + + enum class ModelVariant {DEFAULT = 0, VLM = 1, EXTRACT = 2, RAG = 3}; + ModelVariant model_variant = ModelVariant::DEFAULT; + + enum class Activation {GELU = 0, SILU = 1}; + Activation activation = Activation::SILU; + + enum class Precision {INT8 = 0, FP16 = 1, FP32 = 2}; + Precision precision = Precision::FP32; + + float default_temperature = 0.6f; + float default_top_p = 0.95f; + size_t default_top_k = 20; + float default_max_tps = -1.0f; + float default_cloud_handoff_threshold = 0.0f; + + std::vector layer_types; + size_t conv_L_cache = 0; + + // Rolling bounded KV compaction (default ON, 4096 -> 2048). Override at runtime with + // CACTUS_KV_COMPRESS_AT (trigger) / CACTUS_KV_COMPRESS_TO (target); CACTUS_KV_COMPRESS_AT=0 disables. + bool kv_compress = true; + float kv_compress_recent_frac = 0.30f; + uint32_t kv_compress_sink = 4; + int32_t kv_compress_trigger_len = 4096; + int32_t kv_compress_target_len = 2048; + bool kv_compress_preserve_special = true; + + uint32_t altup_num_inputs = 4; + uint32_t laurel_rank = 64; + static constexpr uint32_t UNSET_U32 = UINT32_MAX; + static constexpr float UNSET_F32 = -1e30f; + uint32_t hidden_size_per_layer_input = UNSET_U32; + uint32_t num_kv_shared_layers = UNSET_U32; + uint32_t sliding_window = UNSET_U32; + float rope_local_base_freq = UNSET_F32; + float final_logit_softcapping = UNSET_F32; + float global_partial_rotary_factor = UNSET_F32; + uint32_t expert_intermediate_size = 0; + uint32_t global_head_dim = UNSET_U32; + uint32_t num_global_kv_heads = 0; + bool attention_k_eq_v = false; + bool enable_moe_block = false; + std::vector activation_sparsity_ppf; + + uint32_t vision_head_dim = 64; + uint32_t vision_kv_heads = 12; + uint32_t vision_intermediate_size = 3072; + uint32_t vision_position_embedding_size = 10240; + uint32_t vision_pooling_kernel_size = 3; + uint32_t vision_default_output_length = 280; + float vision_rope_theta = 100.0f; + + uint32_t audio_hidden_dim = 0; + uint32_t audio_num_layers = 0; + uint32_t audio_num_heads = 0; + uint32_t audio_head_dim = 0; + uint32_t audio_input_feat_size = 128; + uint32_t audio_conf_conv_kernel_size = 5; + uint32_t audio_chunk_size = 12; + uint32_t audio_context_left = 13; + uint32_t audio_context_right = 0; + float audio_logit_cap = 50.0f; + float audio_residual_weight = 0.5f; + uint32_t audio_output_proj_dims = 0; + uint32_t audio_vocab_size = 128; + uint32_t audio_vocab_offset = 0; + uint32_t audio_soft_tokens = 188; + uint32_t audio_sscp_conv0_channels = 128; + uint32_t audio_sscp_conv1_channels = 32; + float audio_sscp_conv_eps = 1e-3f; + float audio_rms_norm_eps = 1e-6f; + uint32_t audio_fft_length = 1024; + uint32_t audio_token_id = 0; + bool audio_fft_overdrive = false; + + static bool is_gemma_family(ModelType t) { + return t == ModelType::GEMMA || t == ModelType::GEMMA3N || t == ModelType::GEMMA4; + } + + static bool is_gemma3_family(ModelType t) { + return t == ModelType::GEMMA || t == ModelType::GEMMA3N; + } + + bool from_json(const std::string& json_path); + std::string to_json() const; + // Disable rolling unless 0 < target < trigger (when trigger > 0). + void validate_kv_compress(); + bool parse_kv_compress_override(const char* trigger_env, const char* target_env); +}; + + + +struct MergeRule { + std::string first; + std::string second; + std::string merged; + uint32_t priority; + + MergeRule(const std::string& f, const std::string& s, const std::string& m, uint32_t p) + : first(f), second(s), merged(m), priority(p) {} +}; + + +struct ToolCallInfo { + std::string name; + std::string arguments; +}; + +struct ChatMessage { + std::string role; + std::string content; + std::string name; + std::vector images; + std::vector audio; + size_t audio_soft_token_count = 0; + std::vector tool_calls; +}; + +struct ToolConstraintSpec { + std::string name; + std::vector parameter_names; + std::vector> parameter_enums; + std::vector required_parameter_names; +}; + +struct TokenizerRuntimeConfig { + enum class TokenizerType { UNKNOWN, BPE, SENTENCEPIECE }; + enum class VocabFormat { UNKNOWN, ID_TAB_TOKEN, LINE_TOKEN }; + enum class Normalizer { NONE, METASPACE, BYTE_LEVEL }; + enum class Decoder { NONE, REPLACE_METASPACE, BYTE_LEVEL }; + + TokenizerType tokenizer_type = TokenizerType::UNKNOWN; + VocabFormat vocab_format = VocabFormat::UNKNOWN; + Normalizer normalizer = Normalizer::NONE; + Decoder decoder = Decoder::NONE; + bool byte_fallback = false; + bool has_chat_template = false; +}; + +TokenizerRuntimeConfig load_tokenizer_runtime_config(const std::string& config_file); +void load_special_tokens_map(const std::string& config_file, std::unordered_map& special_tokens); +std::vector split_with_special_tokens(const std::string& text, const std::unordered_map& special_tokens); + +inline std::string extract_json_string(const std::string& json, size_t& pos) { + std::string value; + while (pos < json.size() && json[pos] != '"') { + if (json[pos] == '\\' && pos + 1 < json.size()) { + pos++; + if (json[pos] == 'n') value += '\n'; + else if (json[pos] == 't') value += '\t'; + else if (json[pos] == 'r') value += '\r'; + else if (json[pos] == '"') value += '"'; + else if (json[pos] == '\\') value += '\\'; + else value += json[pos]; + } else { + value += json[pos]; + } + pos++; + } + if (pos < json.size()) pos++; + return value; +} + +class Tokenizer { +public: + virtual ~Tokenizer() = default; + + virtual std::vector encode(const std::string& text) const = 0; + virtual std::string decode(const std::vector& tokens) const = 0; + + virtual std::vector apply_chat_template(const std::vector& messages, bool add_generation_prompt = true) const; + virtual std::string format_chat_prompt(const std::vector& messages, bool add_generation_prompt = true, const std::string& tools_json = "", bool enable_thinking_if_supported = false) const; + + virtual uint32_t get_vocab_size() const = 0; + virtual uint32_t get_unk_token() const = 0; + virtual uint32_t get_bos_token() const = 0; + virtual uint32_t get_eos_token() const = 0; + virtual std::unordered_set special_token_ids() const { return {}; } + virtual bool has_chat_template() const { return has_chat_template_; } + bool is_qwen_family() const { return model_type_ == ModelType::QWEN; } + bool is_lfm2_family() const { return model_type_ == ModelType::LFM2; } + bool has_function_call_tokens() const { return encode("").size() == 1; } + std::string get_default_stop_sequence() const; + + virtual bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) = 0; + + uint32_t get_image_token_id() const { return image_token_id_; } + uint32_t get_fake_token_id() const { return fake_token_id_; } + uint32_t get_global_img_token_id() const { return global_img_token_id_; } + + void set_image_soft_token_count(size_t n) { image_soft_token_count_ = n; } + size_t get_image_soft_token_count() const { return image_soft_token_count_; } + + void set_lfm2_vision_config(const Config& cfg) { lfm2_vision_config_ = cfg; has_lfm2_vision_config_ = true; } + +protected: + enum class ModelType { UNKNOWN, GEMMA4, GEMMA, QWEN, LFM2, NEEDLE }; + ModelType model_type_ = ModelType::UNKNOWN; + enum class ModelVariant { DEFAULT, VLM, EXTRACT, RAG}; + ModelVariant model_variant_ = ModelVariant::DEFAULT; + bool has_chat_template_ = false; + std::string chat_template_; + + uint32_t image_token_id_ = 396; + uint32_t fake_token_id_ = 49189; + uint32_t global_img_token_id_ = 49152; + + + uint32_t vision_patch_size_ = 16; + uint32_t vision_pooling_kernel_size_ = 3; + uint32_t vision_default_output_length_ = 280; + uint32_t vision_image_size_ = 768; + size_t image_soft_token_count_ = 0; + Config lfm2_vision_config_{}; + bool has_lfm2_vision_config_ = false; + TokenizerRuntimeConfig runtime_config_; + + void detect_model_type(const std::string& config_path); + void load_chat_template(const std::string& template_file); + std::string format_gemma4_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; + std::string format_qwen_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; + std::string format_lfm2_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; + std::string format_needle_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; + std::string format_gemma_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; +}; + +class BPETokenizer : public Tokenizer { +public: + BPETokenizer(); + ~BPETokenizer(); + + bool load_vocabulary_mmap(const std::string& vocab_file, const std::string& merges_file); + bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) override; + + std::vector encode(const std::string& text) const override; + std::string decode(const std::vector& tokens) const override; + + uint32_t get_vocab_size() const override { return vocab_size_; } + uint32_t get_unk_token() const override { return unk_token_id_; } + uint32_t get_bos_token() const override { return bos_token_id_; } + uint32_t get_eos_token() const override { return eos_token_id_; } + std::unordered_set special_token_ids() const override { + std::unordered_set ids; + for (const auto& kv : special_tokens_) ids.insert(kv.second); + return ids; + } + +private: + std::unordered_map token_to_id_; + std::vector id_to_token_; + std::vector merge_rules_; + std::unordered_map merge_map_; + + uint32_t vocab_size_; + uint32_t unk_token_id_; + uint32_t bos_token_id_; + uint32_t eos_token_id_; + + void* vocab_mmap_ptr_; + size_t vocab_mmap_size_; + + void* merges_mmap_ptr_; + size_t merges_mmap_size_; + + std::vector apply_bpe(const std::vector& tokens) const; + std::pair find_best_merge_fast(const std::vector& tokens) const; + + std::string bytes_to_unicode(const std::string& text) const; + std::string unicode_to_bytes(const std::string& text) const; + std::vector byte_level_split(const std::string& text) const; + std::vector utf8_split(const std::string& text) const; + + void cleanup_mmap(); + +private: + mutable std::unordered_map byte_to_unicode_; + mutable std::unordered_map unicode_to_byte_; + void init_byte_mappings() const; + + std::unordered_map special_tokens_; + std::vector split_with_special_tokens(const std::string& text) const; + void load_special_tokens(const std::string& config_file); +}; + +class SPTokenizer : public Tokenizer { +public: + SPTokenizer(); + ~SPTokenizer(); + + bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) override; + + std::vector encode(const std::string& text) const override; + std::string decode(const std::vector& tokens) const override; + + uint32_t get_vocab_size() const override { return vocab_size_; } + uint32_t get_unk_token() const override { return unk_token_id_; } + uint32_t get_bos_token() const override { return bos_token_id_; } + uint32_t get_eos_token() const override { return eos_token_id_; } + std::unordered_set special_token_ids() const override { + std::unordered_set ids; + for (const auto& kv : special_tokens_) ids.insert(kv.second); + return ids; + } + +private: + struct TrieNode { + std::unordered_map> children; + int32_t token_id = -1; + float score = 0.0f; + }; + + std::unique_ptr trie_root_; + std::unordered_map token_to_id_; + std::vector id_to_token_; + std::vector token_scores_; + + uint32_t vocab_size_; + uint32_t unk_token_id_; + uint32_t bos_token_id_; + uint32_t eos_token_id_; + uint32_t pad_token_id_; + + bool sp_bpe_mode_ = false; + bool sp_add_dummy_prefix_ = false; + bool sp_byte_fallback_ = false; + + void* vocab_mmap_ptr_; + size_t vocab_mmap_size_; + + void build_trie(); + std::vector> tokenize_with_trie(const std::string& text) const; + std::vector tokenize_with_bpe(const std::string& text) const; + std::string preprocess_text(const std::string& text) const; + std::string postprocess_text(const std::string& text) const; + std::vector split_by_unicode_spaces(const std::string& text) const; + + void cleanup_mmap(); + + std::unordered_map special_tokens_; + std::vector split_with_special_tokens(const std::string& text) const; + void load_special_tokens(const std::string& config_file); +}; + + + +class ToolCallConstrainer { +public: + enum class State { + DONE, + GEMMA_START, + GEMMA_EXPECT_CALL, + GEMMA_NAME, + GEMMA_ARGS, + GEMMA_EXPECT_END, + NEEDLE_START + }; + + void init(Config::ModelType model_type, + const std::vector& tools, + Tokenizer* tokenizer); + + const std::unordered_map& get_bias() const { return current_bias_; } + const std::vector* get_dense_bias() const { return dense_ready_ ? &dense_bias_ : nullptr; } + + void update(uint32_t token_id, const std::string& decoded_text); + + void reset(); + + bool is_active() const { return active_; } + +private: + enum class Region { FREE, IN_NAME, IN_ARG_KEY }; + struct TrieNode { + std::unordered_map> children; + bool is_terminal = false; + }; + + bool active_ = false; + State state_ = State::GEMMA_START; + Config::ModelType model_type_ = Config::ModelType::GEMMA4; + Tokenizer* tokenizer_ = nullptr; + + std::vector tool_specs_; + std::string generated_text_; + + std::string call_start_tag_; + std::string call_end_tag_; + std::string response_start_tag_; + std::vector call_start_sequence_; + std::vector call_end_sequence_; + size_t forced_tag_progress_ = 0; + + std::unordered_set gemma_call_start_tokens_; + std::unordered_set gemma_call_end_tokens_; + std::unordered_set gemma_response_start_tokens_; + std::unordered_set gemma_call_prefix_tokens_; + + std::unordered_set backtick_tokens_; + std::unordered_set open_brace_tokens_; + std::unordered_set close_brace_tokens_; + + std::unordered_map current_bias_; + std::vector dense_bias_; + bool dense_ready_ = false; + + void compute_bias(); + void tokenize_grammar_elements(); + void add_tokens_for_string(const std::string& str, std::unordered_set& token_set); + void advance_forced_tag(const std::vector& sequence, uint32_t token_id); + void init_common_tokens(); + + bool is_needle() const { return model_type_ == Config::ModelType::NEEDLE; } + + Region region_ = Region::FREE; + std::string constraint_buffer_; + std::string constrained_buf_; + std::string current_function_; + bool in_arguments_ = false; + int arguments_depth_ = 0; + int nesting_depth_ = 0; + bool in_string_value_ = false; + bool prev_char_escape_ = false; + bool at_key_start_ = false; + bool await_enum_open_ = false; + const TrieNode* active_enum_trie_ = nullptr; + std::unordered_set emitted_keys_; + std::unique_ptr name_trie_; + std::unordered_map> param_tries_; + std::unordered_map>> enum_tries_; + std::unique_ptr remaining_key_trie_; + std::vector token_strings_; + std::unordered_map> token_index_; + uint32_t escape_tag_token_ = 0; + bool escape_tag_valid_ = false; + + void build_constraint_tables(); + void reset_constraint_state(); + void trie_insert(TrieNode* root, const std::string& word); + const TrieNode* trie_seek(const TrieNode* root, const std::string& prefix) const; + bool trie_token_valid(const std::string& token_text, const TrieNode* node, char terminator) const; + void mark_trie_bias(const TrieNode* node, char terminator, const std::vector& extra_allowed, + int32_t terminal_token = -1, bool exact_terminator = false); + const TrieNode* current_param_trie() const; + void rebuild_remaining_keys(); + bool required_satisfied() const; + + void feed_needle_text(const std::string& text); + void feed_needle_char(char ch); + bool needle_at_arg_key_start() const; + bool needle_is_value_string_start() const; + + void feed_gemma_text(const std::string& text); + void feed_gemma_char(char ch); +}; + +class Model { +public: + struct DebugNode { + uint32_t layer_idx; + std::string name; + size_t node_id; + }; + + Model(); + explicit Model(const Config& config); + ~Model(); + + const Config& get_config() const { return config_; } + Tokenizer* get_tokenizer() const { return tokenizer_.get(); } + const std::vector& get_debug_nodes() const; + + bool init(const std::string& bundle_dir, size_t context_size, + const std::string& system_prompt = "", bool do_warmup = true); + + uint32_t decode(const std::vector& tokens, float temperature = -1.0f, float top_p = -1.0f, + size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr, + float min_p = 0.15f, float repetition_penalty = 1.1f); + void set_sample_seed(uint64_t seed) { sample_rng_.seed(static_cast(seed)); } + bool prefill_and_sample_first_token(const std::vector& tokens, uint32_t& out_token, + float* out_uncertainty = nullptr); + + std::vector> decode_batch(const std::vector& seed_tokens, + size_t max_new_tokens); + std::vector> generate_batch(const std::vector>& prompts, + size_t max_new_tokens, bool stop_on_eos = false); + + bool supports_dynamic_batch(); + void set_decode_slots(size_t num_slots); + std::vector batch_stop_token_ids() const; + + void prefill(const std::vector& tokens, size_t chunk_size = 128, const std::string& profile_file = "", + bool prepare_decode = true); + + void prefill_with_images(const std::vector& tokens, const std::vector& image_paths, + const std::string& profile_file = ""); + + void prefill_with_audio(const std::vector& tokens, + const std::vector>& audio_features_per_message, + const std::string& profile_file = ""); + + void prefill_with_media(const std::vector& tokens, + const std::vector& image_paths, + const std::vector>& audio_features_per_message, + const std::string& profile_file = ""); + + uint32_t decode_with_images(const std::vector& tokens, const std::vector& image_paths, + float temperature = -1.0f, float top_p = -1.0f, + size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr, + float min_p = 0.15f, float repetition_penalty = 1.1f); + + uint32_t decode_with_audio(const std::vector& tokens, + const std::vector>& audio_features_per_message, + float temperature = 0.0f, float top_p = 0.0f, + size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr, + float min_p = 0.15f, float repetition_penalty = 1.1f, + float* out_token_time_start = nullptr, float* out_token_time_end = nullptr); + + struct ParakeetTdtStreamState { + bool initialized = false; + uint32_t last_token = 0; + size_t time_index = 0; + std::vector> dec_state; + std::vector pending; + float confirmed_sec = 0.0f; + size_t decoded_tokens = 0; + double raw_decode_ms = 0.0; + }; + + std::vector transcribe_parakeet_tdt(const std::vector& audio_features, + ParakeetTdtStreamState* stream = nullptr, + bool is_final = true, + size_t end_frame = 0, + const std::atomic* should_stop = nullptr); + std::vector transcribe_whisper_seq2seq(const std::vector& audio_features, + const std::vector& decoder_prompt_tokens, + size_t max_tokens, + const std::vector>& stop_token_sequences, + const std::atomic* should_stop = nullptr, + int64_t suppress_token_id = -1); + + std::vector get_embeddings(const std::vector& tokens, bool pooled = true, + bool normalize = false, const std::string& profile_file = ""); + + std::vector get_text_embeddings(const std::vector& tokens, bool normalize = false); + std::vector get_lm_embeddings(const std::vector& tokens, bool normalize = false); + bool has_lm_embedding() const { return decoder_embed_ != nullptr; } + bool has_text_embedding() const { return components_.count("text_embedding") > 0; } + bool supports_warm_media_injection() const { + if (lm_encoder_media_step_ != nullptr) return true; + return encoder_ != nullptr && decoder_ != nullptr + && input_index(*decoder_, "inputs_embeds") >= 0; + } + + std::vector get_image_embeddings(const std::string& image_path); + + std::vector get_audio_embeddings(const std::vector& audio_features); + + void reset_cache(); + void record_sampled_token(uint32_t token) { + if (token_history_.size() >= MAX_TOKEN_HISTORY) { + token_history_.erase(token_history_.begin(), token_history_.begin() + (MAX_TOKEN_HISTORY / 2)); + } + token_history_.push_back(token); + } + + double score_tokens_window_logprob(const std::vector& tokens, size_t start, size_t end, + size_t context, size_t* tokens_scored); + + void set_cache_window(size_t window_size, size_t sink_size = 4); + size_t get_cache_size() const { return cache_total_seq_len_; } + + size_t get_prefill_chunk_size() const { return 128; } + bool can_generate() const { return decoder_ != nullptr; } // false for embedding/encoder-only models + double last_prefill_cache_copy_ms() const { return last_prefill_cache_copy_ms_; } + size_t last_prefill_padding_tokens() const { return last_prefill_padding_tokens_; } + size_t last_prefill_scalar_tail_tokens() const { return last_prefill_scalar_tail_tokens_; } + size_t last_prefill_tail_chunk_tokens() const { return last_prefill_tail_chunk_tokens_; } + size_t last_prefill_tail_padding_tokens() const { return last_prefill_tail_padding_tokens_; } + + void remove_thinking_tokens(const std::vector>& ranges); + void compact_kv_cache() {} + + void compress_kv_cache_keydiff(const cactus::kvcompress::Params& params); + void maybe_roll_compact(); + std::vector compressible_layers() const; + void apply_kv_compress_env_override(); + + void set_tool_constraints(const std::vector& tools); + void clear_tool_constraints(); + void update_tool_constraints(uint32_t token_id); + + void set_vocab_bias(const std::unordered_map& bias) { vocab_bias_ = bias; } + void clear_vocab_bias() { vocab_bias_.clear(); } + bool has_vocab_bias() const { return !vocab_bias_.empty(); } + const std::unordered_map& get_vocab_bias() const { return vocab_bias_; } + + bool has_handoff_probe() const { return handoff_probe_loaded_; } + void reset_handoff_probe_rollout() { handoff_probe_hidden_.clear(); } + bool has_handoff_probe_rollout() const; + float handoff_probe_wrong_probability() const; + +private: + struct Binding { + int node_id = -1; + std::string path; + }; + + struct CacheStateBinding { + std::string layer_key; + int key_node_id = -1; + int value_node_id = -1; + }; + + struct Component { + std::string name; + std::string graph_path; + std::vector runtime_input_node_ids; + std::vector logical_inputs; + std::vector output_node_ids; + std::vector logical_outputs; + std::vector bindings; + std::vector cache_states; + std::map metadata; + std::unique_ptr graph; + std::vector> input_buffers; + }; + + void copy_cache_state(const Component& src, Component& dst); + + struct ChunkedPrefillResult { + size_t logical_tokens = 0; + size_t executed_tokens = 0; + size_t padding_tokens = 0; + size_t scalar_tail_tokens = 0; + size_t last_logit_row = 0; + }; + + bool load_manifest(); + bool setup_tokenizer(); + bool load_components(const std::unordered_set& required_components); + bool load_component_graph(Component& comp); + void unload_component_graph(Component& comp); + bool bind_runtime_buffers(Component& comp); + void run_step(uint32_t token_id, size_t position, bool read_logits, bool use_fused = true); + void run_step_batch(const std::vector& token_ids, const std::vector& positions); + void set_component_batch(Component& comp, size_t batch); + size_t decoder_cache_num_slots(); + void run_encoder_step(uint32_t token_id, size_t position); + void run_media_step(size_t position, const uint8_t* feature_row, size_t feature_row_bytes, + Precision feature_precision); + void write_media_embeds_row(Component& comp, int embeds_idx, const uint8_t* feature_row, + size_t feature_row_bytes, Precision feature_precision); + void reset_encoder_cross_kv_route_state(); + bool finish_encoder_cross_kv_prepare(); + bool prepare_encoder_cross_kv_from_text(const std::vector& tokens); + bool prepare_encoder_cross_kv_from_audio(const std::vector& audio_features); + bool run_encoder_cross_kv_decoder_step(uint32_t token_id, size_t position); + std::vector run_encoder_cross_kv_decode_loop( + const std::vector& decoder_prompt_tokens, + size_t max_tokens, + const std::vector>& stop_token_sequences, + const std::atomic* should_stop); + void copy_component_outputs_to_inputs(const Component& source, Component& target); + bool copy_cross_kv_outputs_to_decoder_cache_inputs(const Component& source, Component& target, size_t source_len); + void copy_encoder_outputs_to_decoder(const Component& enc); + void copy_component_outputs_to_chunk_inputs(const Component& source, Component& target, size_t token_index); + void copy_component_outputs_to_chunk_inputs_range(const Component& source, Component& target, size_t token_offset); + bool cache_states_compatible(const Component& source, const Component& target) const; + void move_cache_states(Component& source, Component& target, size_t logical_current = std::numeric_limits::max()); + void set_cache_current_len(Component& comp, size_t len); + void reset_component_cache_states(Component& comp); + void reset_prefill_stats(); + size_t component_chunk_tokens(const Component& comp, const std::string& input_name) const; + size_t component_output_tokens(const Component& comp, const std::string& output_name) const; + ChunkedPrefillResult run_chunked_prefill(const std::vector& tokens, size_t start_position, + size_t chunk_size, bool prepare_decode); + void execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, size_t encoder_chunk, + size_t chunk_tokens, const std::vector& tokens, + size_t processed, size_t start_position); + void run_full_context_text(); + void prepare_sampling_context(float repetition_penalty); + uint32_t sample_component_logits(Component& comp, float temperature, float top_p, size_t top_k, + float min_p, bool greedy, float* out_uncertainty); + uint32_t argmax_component_logits(Component& comp, size_t logit_row = std::numeric_limits::max(), + float* out_uncertainty = nullptr); + uint32_t argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty); + std::vector argmax_component_logits_batch(Component& comp, size_t batch); + void write_int_input(Component& comp, const std::string& name, int64_t value); + void write_int_input_at(Component& comp, const std::string& name, size_t index, int64_t value); + void write_bytes_input(Component& comp, const std::string& name, const void* data, size_t byte_size); + int input_index(const Component& comp, const std::string& name) const; + int output_index(const Component& comp, const std::string& name) const; + uint32_t argmax_last_logits(float* out_uncertainty = nullptr); + bool load_handoff_probe(); + void maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name = "probe_hidden"); + void run_vision_encoder(const std::string& image_path); + void run_vision_encoder_lfm2_vl(const std::string& image_path); + void encode_lfm2_vl_image_into_features(const std::string& image_path); + bool load_lfm2_vl_position_grid(); + void run_audio_encoder(const std::vector& audio_features); + void run_audio_encoder_messages(const std::vector>& audio_features_per_message); + bool run_chunk_prefill_path(const std::vector& tokens, + const std::vector& image_paths, + const std::vector>& audio_features_per_message); + bool build_lm_encoder_outputs_dynamic_gemma4( + const std::vector& tokens, + std::map>& store_bytes, + std::map& store_prec, + std::map>& store_shape); + + std::string bundle_dir_; + std::map components_; + Component* encoder_ = nullptr; + Component* decoder_ = nullptr; + Component* decoder_prefill_ = nullptr; + Component* decoder_embed_ = nullptr; + Component* prefill_encoder_ = nullptr; + enum class DecodeRoute { CACHED_STEP, DIRECT_DECODER_STEP, FULL_CONTEXT_TEXT, ENCODER_CROSS_KV_STEP }; + DecodeRoute decode_route_ = DecodeRoute::CACHED_STEP; + Component* source_encoder_ = nullptr; + Component* decoder_cross_kv_ = nullptr; + Component* vision_encoder_ = nullptr; + Component* vision_projector_ = nullptr; + Component* audio_encoder_ = nullptr; + Component* lm_encoder_media_step_ = nullptr; + Component* decoder_prefill_chunk_ = nullptr; + Component* lm_encoder_ = nullptr; + Component* lm_encoder_text_chunk_ = nullptr; + Component* lm_encoder_media_chunk_ = nullptr; + std::string encoder_cross_kv_source_kind_; + bool encoder_cross_kv_ready_ = false; + size_t encoder_cross_kv_source_len_ = 0; + bool prefill_tail_pad_disabled_ = false; + + std::string family_; + + std::map> media_features_; + std::map> media_feature_shapes_; + std::map media_feature_precisions_; + + std::vector lfm2_pos_grid_; + int lfm2_pos_grid_h_ = 0; + int lfm2_pos_grid_w_ = 0; + int lfm2_pos_grid_dim_ = 0; + bool lfm2_pos_grid_loaded_ = false; + + Config config_; + std::unique_ptr tokenizer_; + bool initialized_ = false; + size_t cache_total_seq_len_ = 0; + std::vector cache_token_ids_; // token id per cache row (canonical head-0 view) + std::unordered_set special_ids_; // special-token ids force-kept during compaction + cactus::kvcompress::SpecialRowTracker special_rows_; // per-(layer,head) special rows for compaction protect + size_t cache_max_seq_len_ = 4096; + size_t last_logit_position_ = 0; + double last_prefill_cache_copy_ms_ = 0.0; + size_t last_prefill_padding_tokens_ = 0; + size_t last_prefill_scalar_tail_tokens_ = 0; + size_t last_prefill_tail_chunk_tokens_ = 0; + size_t last_prefill_tail_padding_tokens_ = 0; + std::vector context_tokens_; + + static constexpr size_t MAX_TOKEN_HISTORY = 128; + std::vector token_history_; + std::vector samp_recent_; + std::vector samp_bias_dense_; + bool samp_has_bias_ = false; + float samp_penalty_ = 1.0f; + bool samp_ctx_active_ = false; + std::mt19937 sample_rng_{std::random_device{}()}; + FusedEmbedCtx fused_embed_ctx_; + int ple_probe_state_ = 0; + + ToolCallConstrainer tool_constrainer_; + std::unordered_map vocab_bias_; + int64_t suppressed_token_id_ = -1; + + bool handoff_probe_loaded_ = false; + uint32_t handoff_probe_feat_dim_ = 0; + uint32_t handoff_probe_t_h_ = 0; + uint32_t handoff_probe_h1_ = 0; + uint32_t handoff_probe_h2_ = 0; + std::vector handoff_probe_norm_weight_; + std::vector handoff_probe_norm_bias_; + std::vector handoff_probe_proj_weight_; + std::vector handoff_probe_proj_bias_; + std::vector handoff_probe_attn_query_; + std::vector handoff_probe_head0_weight_; + std::vector handoff_probe_head0_bias_; + std::vector handoff_probe_head2_weight_; + std::vector handoff_probe_head2_bias_; + std::vector handoff_probe_head4_weight_; + std::vector handoff_probe_head4_bias_; + std::vector handoff_probe_hidden_; + + mutable std::vector debug_nodes_; +}; + +class ConvCache { +public: + struct CircularView { + const void* ptr1; + size_t len1; + const void* ptr2; + size_t len2; + size_t total_len; + }; + + void init(size_t layers, size_t hidden_dim, size_t window_len, Precision model_precision); + CircularView get_window(size_t layer) const; + void update(CactusGraph* gb, size_t layer, const size_t latest_token); + void reset(); + + bool is_empty() const { return num_layers == 0; } + + size_t num_layers = 0; + size_t hidden_size = 0; + size_t window_size = 0; + Precision precision = Precision::FP32; + size_t element_size = 4; + +private: + struct LayerState { + std::vector data; + size_t head = 0; + size_t count = 0; + }; + + std::vector layer_states; +}; + +class Siglip2Preprocessor { +public: + struct Config { + int patch_size = 16; + int downsample_factor = 2; + int min_tiles = 2; + int max_tiles = 10; + bool use_thumbnail = true; + int min_image_tokens = 64; + int max_image_tokens = 256; + int max_num_patches = 1024; + int tile_size = 512; + float max_pixels_tolerance = 2.0f; + bool do_resize = true; + bool do_rescale = true; + bool do_normalize = true; + bool do_convert_rgb = true; + bool do_image_splitting = true; + float rescale_factor = 1.0f / 255.0f; + float image_mean[3] = {0.5f, 0.5f, 0.5f}; + float image_std[3] = {0.5f, 0.5f, 0.5f}; + }; + + struct PreprocessedImage { + std::vector pixel_values; + std::vector pixel_attention_mask; + std::vector> spatial_shapes; + std::vector pixel_values_shape; + std::vector pixel_attention_mask_shape; + std::vector spatial_shapes_shape; + int num_patches_height; + int num_patches_width; + int actual_num_patches; + int num_tiles; + int patch_dim; + int max_patches_per_tile; + int image_rows; + int image_cols; + int image_height; + int image_width; + int tokens_per_tile; + int thumbnail_tokens; + + ~PreprocessedImage(); + }; + + struct SpatialShapeResult { + std::vector> shapes; + int grid_rows; + int grid_cols; + }; + + explicit Siglip2Preprocessor(const Config& config); + Siglip2Preprocessor(); + ~Siglip2Preprocessor(); + + PreprocessedImage preprocess_from_file(const std::string& image_path); + PreprocessedImage preprocess_from_memory(const unsigned char* img_data, int width, int height, int channels); + SpatialShapeResult compute_spatial_shapes(int height, int width); + +private: + Config config_; + + std::pair compute_pixel_limits() const; + std::vector convert_to_rgb(const unsigned char* img_data, int width, int height, int channels); + std::pair smart_resize(int height, int width); + bool is_image_too_large(int height, int width); + std::pair get_grid_layout(int height, int width); + std::pair find_closest_aspect_ratio(float aspect_ratio, int width, int height); + std::vector resize_image(const unsigned char* img_data, int src_width, int src_height, + int dst_width, int dst_height, int channels); + std::vector normalize_image(const float* img_data, int width, int height, int channels); + std::vector> convert_image_to_patches( + const std::vector& image, int width, int height, int channels, int patch_size); + PreprocessedImage pad_patches(const std::vector>& tile_patches, + const std::vector>& spatial_shapes, + int patch_dim, + int max_patches_per_tile); + int round_by_factor(int number, int factor); +}; + +std::unique_ptr create_model(const std::string& model_folder); + +struct Gemma4ImagePreprocessed { + std::vector pixel_values; + std::vector pixel_position_ids; + size_t num_patches = 0; + size_t max_patches = 0; + size_t patch_dim = 0; +}; + +Gemma4ImagePreprocessed preprocess_gemma4_image(const std::string& image_path, const Config& config); + +struct Lfm2VlImagePreprocessed { + std::vector pixel_values; + std::vector pixel_attention_mask; + std::vector> spatial_shapes; + size_t max_num_patches = 0; + size_t patch_dim = 0; +}; + +Lfm2VlImagePreprocessed preprocess_lfm2_vl_image(const std::string& image_path, const Config& config); + +void interpolate_position_embeddings(const float* grid, int grid_h, int grid_w, int dim, + int out_h, int out_w, float* out); + +void pixel_unshuffle(const float* feature, int h, int w, int dim, int factor, float* out); + +struct Lfm2VlTokenLayout { + int grid_rows = 1; + int grid_cols = 1; + int tokens_per_tile = 0; + int thumbnail_tokens = 0; + bool has_thumbnail = false; +}; + +Lfm2VlTokenLayout lfm2_vl_token_layout(int image_height, int image_width, const Config& config); + +struct Qwen3VlImagePreprocessed { + std::vector pixel_values; + size_t grid_t = 1; + size_t grid_h = 0; + size_t grid_w = 0; + size_t patch_dim = 0; +}; + +Qwen3VlImagePreprocessed preprocess_qwen3_vl_image(const std::string& image_path, const Config& config); + + +struct SpectrogramConfig { + size_t n_fft = 400; + size_t hop_length = 160; + size_t frame_length = 400; + float power = 2.0f; + bool center = true; + const char* pad_mode = "reflect"; + bool onesided = true; + float dither = 0.0f; + float mel_floor = 1e-10f; + const char* log_mel = nullptr; + float reference = 1.0f; + float min_value = 1e-10f; + bool remove_dc_offset = false; + float preemphasis = 0.0f; + bool hann_periodic = true; + float window_a0 = 0.5f; + size_t fft_override = 0; + bool mel_floor_additive = false; +}; + +namespace index { + constexpr uint32_t MAGIC = 0x43414354; + constexpr uint32_t VERSION = 1; + + struct Document { + int id; + std::vector embedding; + std::string content; + std::string metadata; + }; + + struct QueryResult { + int doc_id; + float score; + + QueryResult(int doc_id, float score) : doc_id(doc_id), score(score) {} + }; + + struct QueryOptions { + size_t top_k = 10; + float score_threshold = -1.0f; + }; + + class Index { + public: + Index(const std::string& index_path, const std::string& data_path, size_t embedding_dim); + ~Index(); + + Index(const Index&) = delete; + Index& operator=(const Index&) = delete; + Index(Index&&) = delete; + Index& operator=(Index&&) = delete; + + void add_documents(const std::vector& documents); + void delete_documents(const std::vector& doc_ids); + std::vector get_documents(const std::vector& doc_ids); + std::vector> query(const std::vector>& embeddings, const QueryOptions& options); + void compact(); + + private: + struct IndexHeader { + uint32_t magic; + uint32_t version; + uint32_t embedding_dim; + uint32_t num_documents; + }; + + struct IndexEntry { + int32_t doc_id; + uint64_t data_offset; + uint8_t flags; // bit 0: tombstone + + const __fp16* embedding() const { + return reinterpret_cast(this + 1); + } + + static size_t size(size_t embedding_dim) { + return sizeof(IndexEntry) + embedding_dim * sizeof(__fp16); + } + }; + + struct DataHeader { + uint32_t magic; + uint32_t version; + }; + + struct DataEntry { + uint16_t content_len; + uint16_t metadata_len; + + const char* content() const { + return reinterpret_cast(this + 1); + } + + const char* metadata() const { + return content() + content_len; + } + }; + + void parse_index_header(); + void parse_data_header(); + void build_doc_id_map(); + void validate_documents(const std::vector& documents); + void validate_doc_ids(const std::vector& doc_ids); + ssize_t write_full(int fd, const void* buf, size_t count); + + std::unordered_map doc_id_map_; + + std::string index_path_, data_path_; + size_t embedding_dim_; + size_t index_entry_size_; + uint32_t num_documents_; + + int index_fd_, data_fd_; + void *mapped_index_, *mapped_data_; + size_t index_file_size_, data_file_size_; + }; +} // namespace index + +} +} diff --git a/cactus-engine/src/engine_image.cpp b/cactus-engine/src/engine_image.cpp new file mode 100644 index 000000000..3fdb92e66 --- /dev/null +++ b/cactus-engine/src/engine_image.cpp @@ -0,0 +1,1055 @@ +#define STBI_NO_BMP +#define STBI_NO_PSD +#define STBI_NO_HDR +#define STBI_NO_PIC +#define STBI_NO_PNM +#define STBI_NO_TGA + +#include "engine.h" +#include "stb_image.h" +#include "stb_image_resize2.h" +#include +#include +#include +#include +#include +#include +#include + +namespace cactus { +namespace engine { + +namespace { + +constexpr int kPillowResizePrecisionBits = 22; + +static unsigned char pillow_clip8(int64_t value) { + int shifted = static_cast(value >> kPillowResizePrecisionBits); + if (shifted < 0) return 0; + if (shifted > 255) return 255; + return static_cast(shifted); +} + +static void precompute_pillow_bilinear_coeffs(int in_size, int out_size, + std::vector& bounds, + std::vector& coeffs, + int& ksize) { + const double scale = static_cast(in_size) / static_cast(out_size); + const double filterscale = std::max(scale, 1.0); + const double support = filterscale; + const double inv_filterscale = 1.0 / filterscale; + ksize = static_cast(std::ceil(support)) * 2 + 1; + bounds.assign(static_cast(out_size) * 2, 0); + coeffs.assign(static_cast(out_size) * ksize, 0); + + for (int out = 0; out < out_size; ++out) { + const double center = (static_cast(out) + 0.5) * scale; + int xmin = static_cast(center - support + 0.5); + if (xmin < 0) xmin = 0; + int xmax = static_cast(center + support + 0.5); + if (xmax > in_size) xmax = in_size; + xmax -= xmin; + + bounds[static_cast(out) * 2] = xmin; + bounds[static_cast(out) * 2 + 1] = xmax; + + std::vector weights(static_cast(ksize), 0.0); + double total = 0.0; + for (int i = 0; i < xmax; ++i) { + double x = (static_cast(i + xmin) - center + 0.5) * inv_filterscale; + if (x < 0.0) x = -x; + double weight = x < 1.0 ? 1.0 - x : 0.0; + weights[static_cast(i)] = weight; + total += weight; + } + for (int i = 0; i < ksize; ++i) { + double weight = total != 0.0 ? weights[static_cast(i)] / total : 0.0; + coeffs[static_cast(out) * ksize + i] = static_cast( + weight < 0.0 + ? -0.5 + weight * static_cast(1 << kPillowResizePrecisionBits) + : 0.5 + weight * static_cast(1 << kPillowResizePrecisionBits)); + } + } +} + +static std::vector resize_rgb_uint8_pillow_bilinear(const unsigned char* src, + int src_w, + int src_h, + int dst_w, + int dst_h) { + std::vector x_bounds; + std::vector y_bounds; + std::vector x_coeffs; + std::vector y_coeffs; + int x_ksize = 0; + int y_ksize = 0; + precompute_pillow_bilinear_coeffs(src_w, dst_w, x_bounds, x_coeffs, x_ksize); + precompute_pillow_bilinear_coeffs(src_h, dst_h, y_bounds, y_coeffs, y_ksize); + + const int y_first = y_bounds[0]; + const int y_last = y_bounds[static_cast(dst_h - 1) * 2] + + y_bounds[static_cast(dst_h - 1) * 2 + 1]; + const int temp_h = y_last - y_first; + std::vector temp(static_cast(dst_w) * temp_h * 3); + + for (int ty = 0; ty < temp_h; ++ty) { + const int sy = y_first + ty; + for (int dx = 0; dx < dst_w; ++dx) { + const int xmin = x_bounds[static_cast(dx) * 2]; + const int count = x_bounds[static_cast(dx) * 2 + 1]; + const int32_t* k = x_coeffs.data() + static_cast(dx) * x_ksize; + for (int c = 0; c < 3; ++c) { + int64_t sum = static_cast(1) << (kPillowResizePrecisionBits - 1); + for (int x = 0; x < count; ++x) { + sum += static_cast(src[(static_cast(sy) * src_w + xmin + x) * 3 + c]) * k[x]; + } + temp[(static_cast(ty) * dst_w + dx) * 3 + c] = pillow_clip8(sum); + } + } + } + + std::vector dst(static_cast(dst_w) * dst_h * 3); + for (int dy = 0; dy < dst_h; ++dy) { + const int ymin = y_bounds[static_cast(dy) * 2] - y_first; + const int count = y_bounds[static_cast(dy) * 2 + 1]; + const int32_t* k = y_coeffs.data() + static_cast(dy) * y_ksize; + for (int dx = 0; dx < dst_w; ++dx) { + for (int c = 0; c < 3; ++c) { + int64_t sum = static_cast(1) << (kPillowResizePrecisionBits - 1); + for (int y = 0; y < count; ++y) { + sum += static_cast(temp[(static_cast(ymin + y) * dst_w + dx) * 3 + c]) * k[y]; + } + dst[(static_cast(dy) * dst_w + dx) * 3 + c] = pillow_clip8(sum); + } + } + } + return dst; +} + +} // namespace + +Siglip2Preprocessor::PreprocessedImage::~PreprocessedImage() { + pixel_values.clear(); + pixel_attention_mask.clear(); +} + +Siglip2Preprocessor::Siglip2Preprocessor(const Config& config) + : config_(config) {} + +Siglip2Preprocessor::Siglip2Preprocessor() : config_() {} + +Siglip2Preprocessor::~Siglip2Preprocessor() = default; + +std::pair Siglip2Preprocessor::compute_pixel_limits() const { + int64_t min_pixels = static_cast(config_.min_image_tokens) * + config_.patch_size * config_.patch_size * + config_.downsample_factor * config_.downsample_factor; + int64_t max_pixels = static_cast(config_.max_image_tokens) * + config_.patch_size * config_.patch_size * + config_.downsample_factor * config_.downsample_factor; + return {min_pixels, max_pixels}; +} + +int Siglip2Preprocessor::round_by_factor(int number, int factor) { + if (factor == 0) { + return number; + } + double scaled = static_cast(number) / static_cast(factor); + long long rounded = static_cast(std::nearbyint(scaled)); + return static_cast(rounded * factor); +} + +std::pair Siglip2Preprocessor::smart_resize(int height, int width) { + const int total_factor = config_.patch_size * config_.downsample_factor; + auto [min_pixels, max_pixels] = compute_pixel_limits(); + + int h_bar = std::max(total_factor, round_by_factor(height, total_factor)); + int w_bar = std::max(total_factor, round_by_factor(width, total_factor)); + + + if (static_cast(h_bar) * static_cast(w_bar) > max_pixels) { + double beta = std::sqrt((static_cast(height) * static_cast(width)) / + static_cast(max_pixels)); + h_bar = std::max(total_factor, + static_cast(std::floor(height / beta / total_factor)) * total_factor); + w_bar = std::max(total_factor, + static_cast(std::floor(width / beta / total_factor)) * total_factor); + } else if (static_cast(h_bar) * static_cast(w_bar) < min_pixels) { + double beta = std::sqrt(static_cast(min_pixels) / + (static_cast(height) * static_cast(width))); + h_bar = std::max(total_factor, + static_cast(std::ceil(height * beta / total_factor)) * total_factor); + w_bar = std::max(total_factor, + static_cast(std::ceil(width * beta / total_factor)) * total_factor); + } + + return {w_bar, h_bar}; +} + + +bool Siglip2Preprocessor::is_image_too_large(int height, int width) { + const int total_factor = config_.patch_size * config_.downsample_factor; + const int h_bar = std::max(config_.patch_size, round_by_factor(height, total_factor)); + const int w_bar = std::max(config_.patch_size, round_by_factor(width, total_factor)); + + const int64_t pixels = static_cast(h_bar) * static_cast(w_bar); + auto [min_pixels, max_pixels] = compute_pixel_limits(); + const double max_pixels_with_tolerance = static_cast(max_pixels) * config_.max_pixels_tolerance; + + bool result = static_cast(pixels) > max_pixels_with_tolerance; + + return result; +} + + +std::pair Siglip2Preprocessor::find_closest_aspect_ratio(float aspect_ratio, int width, int height) { + float best_ratio_diff = std::numeric_limits::infinity(); + std::pair best_ratio = {1, 1}; + int area = width * height; + + std::vector> target_ratios; + for (int n = config_.min_tiles; n <= config_.max_tiles; ++n) { + for (int w = 1; w <= n; ++w) { + for (int h = 1; h <= n; ++h) { + int total_tiles = w * h; + if (total_tiles >= config_.min_tiles && total_tiles <= config_.max_tiles) { + target_ratios.push_back({w, h}); + } + } + } + } + + std::sort(target_ratios.begin(), target_ratios.end()); + target_ratios.erase(std::unique(target_ratios.begin(), target_ratios.end()), target_ratios.end()); + std::sort(target_ratios.begin(), target_ratios.end(), [](const auto& a, const auto& b) { + return (a.first * a.second) < (b.first * b.second); + }); + + for (const auto& ratio : target_ratios) { + float target_aspect_ratio = static_cast(ratio.first) / ratio.second; + float ratio_diff = std::abs(aspect_ratio - target_aspect_ratio); + + if (ratio_diff < best_ratio_diff) { + best_ratio_diff = ratio_diff; + best_ratio = ratio; + } else if (ratio_diff == best_ratio_diff) { + int target_area = config_.tile_size * config_.tile_size * ratio.first * ratio.second; + if (area > 0.5f * target_area) { + best_ratio = ratio; + } + } + } + + + return best_ratio; +} + +std::pair Siglip2Preprocessor::get_grid_layout(int height, int width) { + float aspect_ratio = (float)width / (float)height; + auto [grid_width, grid_height] = find_closest_aspect_ratio(aspect_ratio, width, height); + + int target_width = config_.tile_size * grid_width; + int target_height = config_.tile_size * grid_height; + + return {target_width, target_height}; +} + +Siglip2Preprocessor::SpatialShapeResult Siglip2Preprocessor::compute_spatial_shapes(int height, int width) { + if (height <= 0 || width <= 0) { + throw std::runtime_error("Image dimensions must be positive"); + } + + if (config_.patch_size <= 0) { + throw std::runtime_error("Patch size must be positive"); + } + + const int patch = config_.patch_size; + auto [resized_width, resized_height] = smart_resize(height, width); + const bool should_split = config_.do_image_splitting && is_image_too_large(height, width); + + + SpatialShapeResult result; + result.grid_rows = 1; + result.grid_cols = 1; + + if (should_split) { + if (config_.tile_size % patch != 0) { + throw std::runtime_error("Tile size must be divisible by patch size"); + } + + auto [grid_target_width, grid_target_height] = get_grid_layout(height, width); + result.grid_cols = grid_target_width / config_.tile_size; + result.grid_rows = grid_target_height / config_.tile_size; + + + const int patches_per_tile_side = config_.tile_size / patch; + const auto tile_shape = std::make_pair(patches_per_tile_side, patches_per_tile_side); + + result.shapes.reserve(static_cast(result.grid_rows) * result.grid_cols + 1); + for (int row = 0; row < result.grid_rows; ++row) { + for (int col = 0; col < result.grid_cols; ++col) { + result.shapes.push_back(tile_shape); + + } + } + + if (config_.use_thumbnail && result.grid_rows * result.grid_cols != 1) { + if (resized_height % patch != 0 || resized_width % patch != 0) { + throw std::runtime_error("Resized thumbnail dimensions must be divisible by patch size"); + } + result.shapes.emplace_back(resized_height / patch, resized_width / patch); + + } + } else { + int target_width = resized_width; + int target_height = resized_height; + if (!config_.do_resize) { + target_width = width; + target_height = height; + } + + if (target_height % patch != 0 || target_width % patch != 0) { + throw std::runtime_error("Target dimensions must be divisible by patch size"); + } + + result.shapes.emplace_back(target_height / patch, target_width / patch); + + } + + return result; +} + + +Siglip2Preprocessor::PreprocessedImage Siglip2Preprocessor::preprocess_from_file(const std::string& image_path) { + int width, height, channels; + unsigned char* img_data = stbi_load(image_path.c_str(), &width, &height, &channels, 0); + + if (!img_data) { + throw std::runtime_error("Failed to load image: " + image_path + " - " + std::string(stbi_failure_reason())); + } + + PreprocessedImage result; + result = preprocess_from_memory(img_data, width, height, channels); + + stbi_image_free(img_data); + return result; +} + +Siglip2Preprocessor::PreprocessedImage Siglip2Preprocessor::preprocess_from_memory( + const unsigned char* img_data, int width, int height, int channels) { + if (!img_data) { + throw std::runtime_error("Invalid image data pointer"); + } + if (width <= 0 || height <= 0) { + throw std::runtime_error("Image dimensions must be positive"); + } + + const int expected_channels = 3; + std::vector rgb_data; + const unsigned char* source_data = img_data; + int source_channels = channels; + + if (config_.do_convert_rgb && channels != expected_channels) { + rgb_data = convert_to_rgb(img_data, width, height, channels); + source_data = rgb_data.data(); + source_channels = expected_channels; + + } + + if (source_channels != expected_channels) { + throw std::runtime_error("Image must have 3 channels (RGB)"); + } + + const int patch = config_.patch_size; + const int downsample = config_.downsample_factor; + const int patch_dim = patch * patch * expected_channels; + + if (patch <= 0) { + throw std::runtime_error("Patch size must be positive"); + } + if (config_.tile_size % patch != 0) { + throw std::runtime_error("Tile size must be divisible by patch size"); + } + + auto [resized_width, resized_height] = smart_resize(height, width); + + + const int patches_per_tile_side = config_.tile_size / patch; + const int tile_patch_count = patches_per_tile_side * patches_per_tile_side; + const int thumbnail_patch_cap = config_.max_image_tokens * downsample * downsample; + int max_patches_per_tile = config_.max_num_patches; + max_patches_per_tile = std::max(max_patches_per_tile, tile_patch_count); + max_patches_per_tile = std::max(max_patches_per_tile, thumbnail_patch_cap); + + const bool allow_splitting = config_.do_image_splitting; + const bool should_split = allow_splitting && is_image_too_large(height, width); + + + size_t expected_tiles = should_split ? static_cast(config_.max_tiles) : 1; + if (should_split && config_.use_thumbnail) { + expected_tiles += 1; + } + + std::vector> tile_patches; + tile_patches.reserve(expected_tiles); + std::vector> spatial_shapes; + spatial_shapes.reserve(expected_tiles); + + auto normalize_and_patchify = [&](const float* data_ptr, int img_width, int img_height) { + if (img_height % patch != 0 || img_width % patch != 0) { + throw std::runtime_error("Image dimensions must be divisible by patch size"); + } + std::vector normalized = normalize_image(data_ptr, img_width, img_height, expected_channels); + auto patches = convert_image_to_patches(normalized, img_width, img_height, expected_channels, patch); + std::vector flattened(patches.size() * patch_dim); + for (size_t idx = 0; idx < patches.size(); ++idx) { + std::copy(patches[idx].begin(), patches[idx].end(), flattened.begin() + idx * patch_dim); + } + tile_patches.push_back(std::move(flattened)); + spatial_shapes.emplace_back(img_height / patch, img_width / patch); + }; + + int grid_rows = 1; + int grid_cols = 1; + bool thumbnail_added = false; + + if (should_split) { + auto [grid_target_width, grid_target_height] = get_grid_layout(height, width); + grid_cols = grid_target_width / config_.tile_size; + grid_rows = grid_target_height / config_.tile_size; + + std::vector resized_grid = resize_image( + source_data, width, height, grid_target_width, grid_target_height, expected_channels); + + std::vector tile_buffer( + static_cast(config_.tile_size) * config_.tile_size * expected_channels); + + for (int row = 0; row < grid_rows; ++row) { + for (int col = 0; col < grid_cols; ++col) { + for (int y = 0; y < config_.tile_size; ++y) { + const float* src_row = resized_grid.data() + + ((row * config_.tile_size + y) * grid_target_width + col * config_.tile_size) * + expected_channels; + float* dst_row = tile_buffer.data() + (y * config_.tile_size) * expected_channels; + std::copy_n(src_row, + static_cast(config_.tile_size) * expected_channels, + dst_row); + } + normalize_and_patchify(tile_buffer.data(), config_.tile_size, config_.tile_size); + + } + } + + if (config_.use_thumbnail && grid_rows * grid_cols != 1) { + std::vector thumbnail_bytes = resize_image( + source_data, width, height, resized_width, resized_height, expected_channels); + normalize_and_patchify(thumbnail_bytes.data(), resized_width, resized_height); + thumbnail_added = true; + + } + } else { + int target_width = resized_width; + int target_height = resized_height; + + const bool needs_resize = config_.do_resize && (width != target_width || height != target_height); + + std::vector resized_image; + if (needs_resize) { + resized_image = resize_image(source_data, width, height, target_width, target_height, expected_channels); + normalize_and_patchify(resized_image.data(), target_width, target_height); + + } else { + resized_image.resize(static_cast(width) * height * expected_channels); + for (size_t idx = 0; idx < resized_image.size(); ++idx) { + resized_image[idx] = static_cast(source_data[idx]); + } + normalize_and_patchify(resized_image.data(), width, height); + target_width = width; + target_height = height; + resized_width = target_width; + resized_height = target_height; + + } + + grid_rows = 1; + grid_cols = 1; + } + + PreprocessedImage result = pad_patches(tile_patches, spatial_shapes, patch_dim, max_patches_per_tile); + + + result.image_rows = grid_rows; + result.image_cols = grid_cols; + result.image_width = resized_width; + result.image_height = resized_height; + + auto compute_tokens = [&](int patches_h, int patches_w) -> int { + int tokens_h = (patches_h + downsample - 1) / downsample; + int tokens_w = (patches_w + downsample - 1) / downsample; + return tokens_h * tokens_w; + }; + + result.tokens_per_tile = spatial_shapes.empty() + ? 0 + : compute_tokens(spatial_shapes.front().first, spatial_shapes.front().second); + + + if (thumbnail_added && !spatial_shapes.empty()) { + const auto& thumb_shape = spatial_shapes.back(); + result.thumbnail_tokens = compute_tokens(thumb_shape.first, thumb_shape.second); + + } else { + result.thumbnail_tokens = 0; + } + + return result; +} + +std::vector Siglip2Preprocessor::convert_to_rgb( + const unsigned char* img_data, int width, int height, int channels) { + + std::vector rgb_data(width * height * 3); + + if (channels == 1) { + for (int i = 0; i < width * height; ++i) { + rgb_data[i * 3 + 0] = img_data[i]; + rgb_data[i * 3 + 1] = img_data[i]; + rgb_data[i * 3 + 2] = img_data[i]; + } + } else if (channels == 4) { + for (int i = 0; i < width * height; ++i) { + rgb_data[i * 3 + 0] = img_data[i * 4 + 0]; + rgb_data[i * 3 + 1] = img_data[i * 4 + 1]; + rgb_data[i * 3 + 2] = img_data[i * 4 + 2]; + } + } else if (channels == 2) { + for (int i = 0; i < width * height; ++i) { + rgb_data[i * 3 + 0] = img_data[i * 2 + 0]; + rgb_data[i * 3 + 1] = img_data[i * 2 + 0]; + rgb_data[i * 3 + 2] = img_data[i * 2 + 0]; + } + } else { + throw std::runtime_error("Unsupported number of channels: " + std::to_string(channels)); + } + + return rgb_data; +} + +std::vector Siglip2Preprocessor::resize_image( + const unsigned char* img_data, int src_width, int src_height, + int dst_width, int dst_height, int channels) { + + const size_t src_elements = static_cast(src_width) * src_height * channels; + std::vector src_float(src_elements); + for (size_t idx = 0; idx < src_elements; ++idx) { + src_float[idx] = static_cast(img_data[idx]); + } + + + std::vector resized_data(static_cast(dst_width) * dst_height * channels); + + stbir_pixel_layout layout = (channels == 1) ? STBIR_1CHANNEL : + (channels == 3) ? STBIR_RGB : STBIR_RGBA; + + float* result = stbir_resize_float_linear( + src_float.data(), src_width, src_height, 0, + resized_data.data(), dst_width, dst_height, 0, + layout + ); + + if (!result) { + throw std::runtime_error("Failed to resize image"); + } + + + return resized_data; +} + +std::vector Siglip2Preprocessor::normalize_image( + const float* img_data, int width, int height, int channels) { + + size_t total_pixels = width * height * channels; + std::vector normalized(total_pixels); + + for (size_t i = 0; i < static_cast(width * height); ++i) { + for (int c = 0; c < channels; ++c) { + size_t idx = i * channels + c; + float pixel = img_data[idx]; + + if (config_.do_rescale) { + pixel *= config_.rescale_factor; + } + + if (config_.do_normalize) { + pixel = (pixel - config_.image_mean[c]) / config_.image_std[c]; + } + + normalized[idx] = pixel; + } + } + + + return normalized; +} + +std::vector> Siglip2Preprocessor::convert_image_to_patches( + const std::vector& image, int width, int height, int channels, int patch_size) { + + int num_patches_height = height / patch_size; + int num_patches_width = width / patch_size; + int num_patches = num_patches_height * num_patches_width; + int patch_elements = patch_size * patch_size * channels; + + std::vector> patches(num_patches, std::vector(patch_elements)); + + for (int ph = 0; ph < num_patches_height; ++ph) { + for (int pw = 0; pw < num_patches_width; ++pw) { + int patch_idx = ph * num_patches_width + pw; + + for (int y = 0; y < patch_size; ++y) { + for (int x = 0; x < patch_size; ++x) { + int img_y = ph * patch_size + y; + int img_x = pw * patch_size + x; + int img_idx = (img_y * width + img_x) * channels; + int patch_offset = (y * patch_size + x) * channels; + + for (int c = 0; c < channels; ++c) { + patches[patch_idx][patch_offset + c] = image[img_idx + c]; + } + } + } + } + } + + return patches; +} + +Siglip2Preprocessor::PreprocessedImage Siglip2Preprocessor::pad_patches( + const std::vector>& tile_patches, + const std::vector>& spatial_shapes, + int patch_dim, + int max_patches_per_tile) { + + if (tile_patches.size() != spatial_shapes.size()) { + throw std::runtime_error("Mismatch between tile data and spatial shapes"); + } + + PreprocessedImage result; + + const int num_tiles = static_cast(tile_patches.size()); + result.num_tiles = num_tiles; + result.patch_dim = patch_dim; + result.max_patches_per_tile = max_patches_per_tile; + result.spatial_shapes = spatial_shapes; + result.actual_num_patches = 0; + + const size_t total_values = static_cast(num_tiles) * max_patches_per_tile * patch_dim; + result.pixel_values.assign(total_values, 0.0f); + result.pixel_attention_mask.assign(static_cast(num_tiles) * max_patches_per_tile, 0); + + + for (int tile_idx = 0; tile_idx < num_tiles; ++tile_idx) { + const auto& [patches_h, patches_w] = spatial_shapes[tile_idx]; + const int actual_patches = patches_h * patches_w; + + if (actual_patches > max_patches_per_tile) { + throw std::runtime_error("Actual patches exceed max_patches_per_tile"); + } + + const auto& flattened = tile_patches[tile_idx]; + const size_t expected_size = static_cast(actual_patches) * patch_dim; + if (flattened.size() != expected_size) { + throw std::runtime_error("Tile patch data has unexpected size"); + } + + float* destination = result.pixel_values.data() + + static_cast(tile_idx) * max_patches_per_tile * patch_dim; + std::memcpy(destination, flattened.data(), expected_size * sizeof(float)); + + + int mask_offset = tile_idx * max_patches_per_tile; + for (int p = 0; p < actual_patches; ++p) { + result.pixel_attention_mask[mask_offset + p] = 1; + } + + result.actual_num_patches += actual_patches; + } + + if (!spatial_shapes.empty()) { + result.num_patches_height = spatial_shapes.front().first; + result.num_patches_width = spatial_shapes.front().second; + } else { + result.num_patches_height = 0; + result.num_patches_width = 0; + } + + result.pixel_values_shape = { + static_cast(num_tiles), + static_cast(max_patches_per_tile), + static_cast(patch_dim) + }; + + + result.pixel_attention_mask_shape = { + static_cast(num_tiles), + static_cast(max_patches_per_tile) + }; + + result.spatial_shapes_shape = { + static_cast(num_tiles), + static_cast(2) + }; + + + return result; +} + +Gemma4ImagePreprocessed preprocess_gemma4_image(const std::string& image_path, const Config& config) { + Gemma4ImagePreprocessed result; + + uint32_t patch_size_u = config.vision_patch_size ? config.vision_patch_size : 16; + uint32_t pooling_k_u = config.vision_pooling_kernel_size ? config.vision_pooling_kernel_size : 3; + uint32_t max_soft_tokens_u = config.vision_default_output_length ? config.vision_default_output_length : 280; + if (patch_size_u == 0 || pooling_k_u == 0 || max_soft_tokens_u == 0) { + throw std::runtime_error("Gemma4 image config has invalid patch/pooling/soft-token values"); + } + + const int patch_size = static_cast(patch_size_u); + const int pooling_k = static_cast(pooling_k_u); + const size_t max_patches = static_cast(max_soft_tokens_u) * pooling_k * pooling_k; + const int side_multiple = pooling_k * patch_size; + const size_t patch_dim = static_cast(3) * patch_size * patch_size; + + int width = 0, height = 0, channels = 0; + unsigned char* raw = stbi_load(image_path.c_str(), &width, &height, &channels, 3); + if (!raw) { + throw std::runtime_error("Failed to load image: " + image_path + " - " + std::string(stbi_failure_reason())); + } + if (width <= 0 || height <= 0) { + stbi_image_free(raw); + throw std::runtime_error("Loaded image has invalid dimensions: " + image_path); + } + + const double target_pixels = static_cast(max_patches) * patch_size * patch_size; + const double pixel_count = std::max(1.0, static_cast(width) * static_cast(height)); + const double factor = std::sqrt(target_pixels / pixel_count); + int target_h = static_cast(std::floor(factor * height / side_multiple)) * side_multiple; + int target_w = static_cast(std::floor(factor * width / side_multiple)) * side_multiple; + if (target_h == 0) target_h = side_multiple; + if (target_w == 0) target_w = side_multiple; + + std::vector resized(static_cast(target_w) * target_h * 3); + + if (target_w == width && target_h == height) { + for (size_t i = 0; i < resized.size(); ++i) { + resized[i] = static_cast(raw[i]); + } + } else { + std::vector resized_u8 = resize_rgb_uint8_pillow_bilinear(raw, width, height, target_w, target_h); + for (size_t i = 0; i < resized.size(); ++i) { + resized[i] = static_cast(resized_u8[i]); + } + } + stbi_image_free(raw); + + const bool do_rescale = true; + const float rescale_factor = config.rescale_factor > 0.0f ? config.rescale_factor : (1.0f / 255.0f); + const bool do_normalize = false; + const float mean_v = do_normalize ? config.image_mean : 0.0f; + const float std_v = (do_normalize && config.image_std != 0.0f) ? config.image_std : 1.0f; + + if (do_rescale || do_normalize) { + const float inv_std = 1.0f / std_v; + for (size_t i = 0; i < resized.size(); ++i) { + float v = resized[i]; + if (do_rescale) v *= rescale_factor; + if (do_normalize) v = (v - mean_v) * inv_std; + resized[i] = v; + } + } + + const int patch_h = target_h / patch_size; + const int patch_w = target_w / patch_size; + const size_t num_patches = static_cast(patch_h) * static_cast(patch_w); + if (num_patches > max_patches) { + throw std::runtime_error("Gemma4 native image preprocessing produced too many patches"); + } + + result.pixel_values.assign(max_patches * patch_dim, 0.0f); + result.pixel_position_ids.assign(max_patches * 2, -1); + + for (int py = 0; py < patch_h; ++py) { + for (int px = 0; px < patch_w; ++px) { + const size_t patch_idx = static_cast(py) * patch_w + px; + float* dst = result.pixel_values.data() + patch_idx * patch_dim; + for (int y = 0; y < patch_size; ++y) { + const int img_y = py * patch_size + y; + for (int x = 0; x < patch_size; ++x) { + const int img_x = px * patch_size + x; + const size_t src_off = (static_cast(img_y) * target_w + img_x) * 3; + const size_t dst_off = (static_cast(y) * patch_size + x) * 3; + dst[dst_off + 0] = resized[src_off + 0]; + dst[dst_off + 1] = resized[src_off + 1]; + dst[dst_off + 2] = resized[src_off + 2]; + } + } + result.pixel_position_ids[patch_idx * 2 + 0] = static_cast(px); + result.pixel_position_ids[patch_idx * 2 + 1] = static_cast(py); + } + } + + result.num_patches = num_patches; + result.max_patches = max_patches; + result.patch_dim = patch_dim; + return result; +} + +Qwen3VlImagePreprocessed preprocess_qwen3_vl_image(const std::string& image_path, const Config& config) { + Qwen3VlImagePreprocessed result; + const int patch_size = static_cast(config.vision_patch_size ? config.vision_patch_size : 16); + const int temporal_patch_size = 2; + const int merge_size = 2; + const size_t image_tokens = config.image_seq_len ? static_cast(config.image_seq_len) : 64; + const size_t merged_patches = image_tokens * static_cast(merge_size * merge_size); + const int grid_side = static_cast(std::sqrt(static_cast(merged_patches))); + if (patch_size <= 0 || grid_side <= 0 || static_cast(grid_side * grid_side) != merged_patches) { + throw std::runtime_error("Qwen3-VL native image preprocessing requires a square static image grid"); + } + + int width = 0, height = 0, channels = 0; + unsigned char* raw = stbi_load(image_path.c_str(), &width, &height, &channels, 3); + if (!raw) { + throw std::runtime_error("Failed to load image: " + image_path + " - " + std::string(stbi_failure_reason())); + } + if (width <= 0 || height <= 0) { + stbi_image_free(raw); + throw std::runtime_error("Loaded image has invalid dimensions: " + image_path); + } + + const int target_h = grid_side * patch_size; + const int target_w = grid_side * patch_size; + std::vector resized(static_cast(target_w) * target_h * 3); + if (target_w == width && target_h == height) { + for (size_t i = 0; i < resized.size(); ++i) { + resized[i] = static_cast(raw[i]); + } + } else { + std::vector src_float(static_cast(width) * height * 3); + for (size_t i = 0; i < src_float.size(); ++i) { + src_float[i] = static_cast(raw[i]); + } + if (!stbir_resize_float_linear( + src_float.data(), width, height, 0, + resized.data(), target_w, target_h, 0, + STBIR_RGB)) { + stbi_image_free(raw); + throw std::runtime_error("Failed to resize image: " + image_path); + } + } + stbi_image_free(raw); + + const float rescale_factor = config.rescale_factor > 0.0f ? config.rescale_factor : (1.0f / 255.0f); + const float mean_v = config.image_mean; + const float std_v = config.image_std != 0.0f ? config.image_std : 1.0f; + for (float& v : resized) { + v = (v * rescale_factor - mean_v) / std_v; + } + + const size_t channel = 3; + const size_t grid_t = 1; + const size_t grid_h = static_cast(grid_side); + const size_t grid_w = static_cast(grid_side); + const size_t patch_dim = channel * static_cast(temporal_patch_size) * patch_size * patch_size; + result.pixel_values.assign(grid_t * grid_h * grid_w * patch_dim, 0.0f); + + for (size_t gh_major = 0; gh_major < grid_h / merge_size; ++gh_major) { + for (size_t gw_major = 0; gw_major < grid_w / merge_size; ++gw_major) { + for (size_t mh = 0; mh < static_cast(merge_size); ++mh) { + for (size_t mw = 0; mw < static_cast(merge_size); ++mw) { + const size_t patch_y = gh_major * merge_size + mh; + const size_t patch_x = gw_major * merge_size + mw; + const size_t patch_index = + (((gh_major * (grid_w / merge_size) + gw_major) * merge_size + mh) * merge_size + mw); + float* dst = result.pixel_values.data() + patch_index * patch_dim; + size_t dst_i = 0; + for (size_t c = 0; c < channel; ++c) { + for (size_t t = 0; t < static_cast(temporal_patch_size); ++t) { + (void)t; + for (int y = 0; y < patch_size; ++y) { + const size_t img_y = patch_y * patch_size + static_cast(y); + for (int x = 0; x < patch_size; ++x) { + const size_t img_x = patch_x * patch_size + static_cast(x); + const size_t src_off = (img_y * static_cast(target_w) + img_x) * channel + c; + dst[dst_i++] = resized[src_off]; + } + } + } + } + } + } + } + } + + result.grid_t = grid_t; + result.grid_h = grid_h; + result.grid_w = grid_w; + result.patch_dim = patch_dim; + return result; +} + +static Siglip2Preprocessor::Config make_lfm2_vl_siglip_config(const Config& config) { + Siglip2Preprocessor::Config sp_cfg; + sp_cfg.patch_size = config.vision_patch_size ? static_cast(config.vision_patch_size) : 16; + sp_cfg.downsample_factor = config.downsample_factor ? static_cast(config.downsample_factor) : 2; + sp_cfg.min_tiles = config.min_tiles ? static_cast(config.min_tiles) : 2; + sp_cfg.max_tiles = config.max_tiles ? static_cast(config.max_tiles) : 10; + sp_cfg.use_thumbnail = config.use_thumbnail; + sp_cfg.min_image_tokens = config.min_image_tokens ? static_cast(config.min_image_tokens) : 64; + sp_cfg.max_image_tokens = config.max_image_tokens ? static_cast(config.max_image_tokens) : 256; + sp_cfg.max_num_patches = config.max_num_patches ? static_cast(config.max_num_patches) : 1024; + sp_cfg.tile_size = config.tile_size ? static_cast(config.tile_size) : 512; + sp_cfg.max_pixels_tolerance = config.max_pixels_tolerance > 0.0f ? config.max_pixels_tolerance : 2.0f; + sp_cfg.rescale_factor = config.rescale_factor > 0.0f ? config.rescale_factor : (1.0f / 255.0f); + sp_cfg.image_mean[0] = sp_cfg.image_mean[1] = sp_cfg.image_mean[2] = config.image_mean; + sp_cfg.image_std[0] = sp_cfg.image_std[1] = sp_cfg.image_std[2] = config.image_std; + sp_cfg.do_image_splitting = true; + sp_cfg.do_resize = true; + sp_cfg.do_rescale = true; + sp_cfg.do_normalize = true; + sp_cfg.do_convert_rgb = true; + return sp_cfg; +} + +Lfm2VlImagePreprocessed preprocess_lfm2_vl_image(const std::string& image_path, const Config& config) { + Siglip2Preprocessor::Config sp_cfg = make_lfm2_vl_siglip_config(config); + Siglip2Preprocessor pre(sp_cfg); + auto sp_out = pre.preprocess_from_file(image_path); + + Lfm2VlImagePreprocessed result; + result.pixel_values = std::move(sp_out.pixel_values); + result.pixel_attention_mask.assign(sp_out.pixel_attention_mask.begin(), + sp_out.pixel_attention_mask.end()); + result.spatial_shapes = sp_out.spatial_shapes; + result.patch_dim = static_cast(sp_out.patch_dim); + result.max_num_patches = static_cast(sp_out.max_patches_per_tile); + return result; +} + +Lfm2VlTokenLayout lfm2_vl_token_layout(int image_height, int image_width, const Config& config) { + Siglip2Preprocessor::Config sp_cfg = make_lfm2_vl_siglip_config(config); + Siglip2Preprocessor pre(sp_cfg); + auto shapes = pre.compute_spatial_shapes(image_height, image_width); + const int df = sp_cfg.downsample_factor > 0 ? sp_cfg.downsample_factor : 1; + auto tokens_for = [df](const std::pair& s) { + return ((s.first + df - 1) / df) * ((s.second + df - 1) / df); + }; + + Lfm2VlTokenLayout layout; + layout.grid_rows = shapes.grid_rows; + layout.grid_cols = shapes.grid_cols; + const int tile_count = shapes.grid_rows * shapes.grid_cols; + const bool split = tile_count > 1; + if (!split || shapes.shapes.empty()) { + layout.tokens_per_tile = shapes.shapes.empty() ? 0 : tokens_for(shapes.shapes.front()); + return layout; + } + + layout.tokens_per_tile = tokens_for(shapes.shapes.front()); + if (sp_cfg.use_thumbnail && static_cast(shapes.shapes.size()) == tile_count + 1) { + layout.has_thumbnail = true; + layout.thumbnail_tokens = tokens_for(shapes.shapes.back()); + } + return layout; +} + +namespace { +void compute_bilinear_antialias_coeffs(int in_size, int out_size, + std::vector& starts, + std::vector>& weights) { + const double scale = static_cast(in_size) / static_cast(out_size); + const double filterscale = std::max(scale, 1.0); + const double support = filterscale; + const double inv_filterscale = 1.0 / filterscale; + starts.assign(out_size, 0); + weights.assign(out_size, {}); + for (int o = 0; o < out_size; ++o) { + const double center = (static_cast(o) + 0.5) * scale; + int xmin = static_cast(center - support + 0.5); + if (xmin < 0) xmin = 0; + int xmax = static_cast(center + support + 0.5); + if (xmax > in_size) xmax = in_size; + std::vector w; + w.reserve(static_cast(xmax - xmin)); + double total = 0.0; + for (int x = xmin; x < xmax; ++x) { + double t = (static_cast(x) - center + 0.5) * inv_filterscale; + if (t < 0.0) t = -t; + const double weight = t < 1.0 ? 1.0 - t : 0.0; + w.push_back(static_cast(weight)); + total += weight; + } + if (total > 0.0) for (float& v : w) v = static_cast(v / total); + starts[o] = xmin; + weights[o] = std::move(w); + } +} +} // namespace + +void interpolate_position_embeddings(const float* grid, int grid_h, int grid_w, int dim, + int out_h, int out_w, float* out) { + std::vector x_start, y_start; + std::vector> x_w, y_w; + compute_bilinear_antialias_coeffs(grid_w, out_w, x_start, x_w); + compute_bilinear_antialias_coeffs(grid_h, out_h, y_start, y_w); + + std::vector tmp(static_cast(grid_h) * out_w * dim, 0.0f); + for (int gy = 0; gy < grid_h; ++gy) { + for (int ox = 0; ox < out_w; ++ox) { + float* dst = tmp.data() + (static_cast(gy) * out_w + ox) * dim; + const auto& w = x_w[ox]; + const int xs = x_start[ox]; + for (size_t k = 0; k < w.size(); ++k) { + const float wk = w[k]; + const float* src = grid + (static_cast(gy) * grid_w + (xs + static_cast(k))) * dim; + for (int d = 0; d < dim; ++d) dst[d] += wk * src[d]; + } + } + } + for (int oy = 0; oy < out_h; ++oy) { + const auto& w = y_w[oy]; + const int ys = y_start[oy]; + for (int ox = 0; ox < out_w; ++ox) { + float* dst = out + (static_cast(oy) * out_w + ox) * dim; + for (int d = 0; d < dim; ++d) dst[d] = 0.0f; + for (size_t k = 0; k < w.size(); ++k) { + const float wk = w[k]; + const float* src = tmp.data() + (static_cast(ys + static_cast(k)) * out_w + ox) * dim; + for (int d = 0; d < dim; ++d) dst[d] += wk * src[d]; + } + } + } +} + +void pixel_unshuffle(const float* feature, int h, int w, int dim, int factor, float* out) { + const int hf = h / factor; + const int wf = w / factor; + const int cf = dim * factor; + const int cff = dim * factor * factor; + for (int yq = 0; yq < hf; ++yq) { + for (int xq = 0; xq < wf; ++xq) { + float* token = out + (static_cast(yq) * wf + xq) * cff; + for (int a = 0; a < factor; ++a) { + const int y = yq * factor + a; + for (int xr = 0; xr < factor; ++xr) { + const int x = xq * factor + xr; + const float* src = feature + (static_cast(y) * w + x) * dim; + float* dst = token + a * cf + xr * dim; + for (int c = 0; c < dim; ++c) dst[c] = src[c]; + } + } + } + } +} + +} // namespace engine +} // namespace cactus diff --git a/cactus-engine/src/gemma_tools.h b/cactus-engine/src/gemma_tools.h new file mode 100644 index 000000000..6b10a74a0 --- /dev/null +++ b/cactus-engine/src/gemma_tools.h @@ -0,0 +1,860 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "json_escape.h" +#include "picojson.h" + +namespace gemma { + +inline std::string to_upper(const std::string& s) { + std::string result = s; + for (auto& c : result) c = std::toupper(c); + return result; +} + +inline std::string escape(const std::string& s) { + return "<|\"|>" + s + "<|\"|>"; +} + +inline std::string use_escape_tags(std::string s) { + size_t pos = 0; + while ((pos = s.find("<|\"|>", pos)) != std::string::npos) { + s.replace(pos, 5, ""); + pos += 8; + } + return s; +} + +inline void skip_whitespace(const std::string& json, size_t& pos) { + while (pos < json.length() && std::isspace(json[pos])) pos++; +} + +inline std::string extract_json_string(const std::string& json, size_t& pos) { + std::string value; + while (pos < json.length() && json[pos] != '"') { + if (json[pos] == '\\' && pos + 1 < json.length()) { + pos++; + if (json[pos] == 'n') value += '\n'; + else if (json[pos] == 't') value += '\t'; + else if (json[pos] == 'r') value += '\r'; + else if (json[pos] == '"') value += '"'; + else if (json[pos] == '\\') value += '\\'; + else if (json[pos] == 'u' && pos + 4 < json.length()) { + // Decode \uXXXX (BMP) -> UTF-8. Without this, JSON-escaped + // chars like < leak literally as "u003c" in tool args. + unsigned int cp = 0; + bool ok = true; + for (int k = 1; k <= 4; k++) { + char h = json[pos + k]; + cp <<= 4; + if (h >= '0' && h <= '9') cp |= static_cast(h - '0'); + else if (h >= 'a' && h <= 'f') cp |= static_cast(h - 'a' + 10); + else if (h >= 'A' && h <= 'F') cp |= static_cast(h - 'A' + 10); + else { ok = false; break; } + } + if (ok) { + pos += 4; + if (cp < 0x80) { + value += static_cast(cp); + } else if (cp < 0x800) { + value += static_cast(0xC0 | (cp >> 6)); + value += static_cast(0x80 | (cp & 0x3F)); + } else { + value += static_cast(0xE0 | (cp >> 12)); + value += static_cast(0x80 | ((cp >> 6) & 0x3F)); + value += static_cast(0x80 | (cp & 0x3F)); + } + } else { + value += json[pos]; + } + } + else value += json[pos]; + } else { + value += json[pos]; + } + pos++; + } + if (pos < json.length()) pos++; + return value; +} + +std::string format_argument(const std::string& json, size_t& pos, bool escape_keys); +std::string format_parameters(const std::string& properties_json, const std::string& /*required_json*/); + +inline std::string format_argument(const std::string& json, size_t& pos, bool escape_keys = true) { + skip_whitespace(json, pos); + if (pos >= json.length()) return ""; + + char c = json[pos]; + + if (c == '"') { + pos++; + std::string value = extract_json_string(json, pos); + return escape(value); + } else if (c == '{') { + std::string result = "{"; + pos++; + bool first = true; + + while (pos < json.length()) { + skip_whitespace(json, pos); + if (pos >= json.length() || json[pos] == '}') { pos++; break; } + if (json[pos] == ',') { pos++; continue; } + + if (json[pos] != '"') break; + pos++; + std::string key = extract_json_string(json, pos); + + skip_whitespace(json, pos); + if (pos < json.length() && json[pos] == ':') pos++; + + std::string value = format_argument(json, pos, escape_keys); + + if (!first) result += ","; + first = false; + if (escape_keys) { + result += escape(key) + ":" + value; + } else { + result += key + ":" + value; + } + } + result += "}"; + return result; + } else if (c == '[') { + std::string result = "["; + pos++; + bool first = true; + + while (pos < json.length()) { + skip_whitespace(json, pos); + if (pos >= json.length() || json[pos] == ']') { pos++; break; } + if (json[pos] == ',') { pos++; continue; } + + std::string value = format_argument(json, pos, escape_keys); + + if (!first) result += ","; + first = false; + result += value; + } + result += "]"; + return result; + } else if (json.compare(pos, 4, "true") == 0) { + pos += 4; + return "true"; + } else if (json.compare(pos, 5, "false") == 0) { + pos += 5; + return "false"; + } else if (json.compare(pos, 4, "null") == 0) { + pos += 4; + return "null"; + } else { + size_t start = pos; + while (pos < json.length() && (std::isdigit(json[pos]) || json[pos] == '.' || + json[pos] == '-' || json[pos] == '+' || json[pos] == 'e' || json[pos] == 'E')) { + pos++; + } + if (pos == start && pos < json.length()) pos++; + return json.substr(start, pos - start); + } +} + +inline std::map parse_json_object_raw(const std::string& json, size_t& pos) { + std::map result; + skip_whitespace(json, pos); + if (pos >= json.length() || json[pos] != '{') return result; + pos++; + + while (pos < json.length()) { + skip_whitespace(json, pos); + if (pos >= json.length() || json[pos] == '}') { pos++; break; } + if (json[pos] == ',') { pos++; continue; } + + if (json[pos] != '"') break; + pos++; + std::string key = extract_json_string(json, pos); + + skip_whitespace(json, pos); + if (pos < json.length() && json[pos] == ':') pos++; + skip_whitespace(json, pos); + + size_t value_start = pos; + if (json[pos] == '"') { + pos++; + while (pos < json.length() && json[pos] != '"') { + if (json[pos] == '\\') pos++; + pos++; + } + pos++; + } else if (json[pos] == '{') { + int depth = 1; + pos++; + while (pos < json.length() && depth > 0) { + if (json[pos] == '{') depth++; + else if (json[pos] == '}') depth--; + else if (json[pos] == '"') { + pos++; + while (pos < json.length() && json[pos] != '"') { + if (json[pos] == '\\') pos++; + pos++; + } + } + pos++; + } + } else if (json[pos] == '[') { + int depth = 1; + pos++; + while (pos < json.length() && depth > 0) { + if (json[pos] == '[') depth++; + else if (json[pos] == ']') depth--; + else if (json[pos] == '"') { + pos++; + while (pos < json.length() && json[pos] != '"') { + if (json[pos] == '\\') pos++; + pos++; + } + } + pos++; + } + } else { + while (pos < json.length() && json[pos] != ',' && json[pos] != '}') pos++; + } + result[key] = json.substr(value_start, pos - value_start); + } + return result; +} + +inline std::string get_json_string_value(const std::string& json, size_t pos) { + skip_whitespace(json, pos); + if (pos < json.length() && json[pos] == '"') { + pos++; + return extract_json_string(json, pos); + } + return ""; +} + +inline std::string format_parameters(const std::string& properties_json, const std::string& /*required_json*/) { + static const std::set standard_keys = {"description", "type", "properties", "required", "nullable"}; + + size_t pos = 0; + auto properties = parse_json_object_raw(properties_json, pos); + + std::string result; + bool first = true; + + for (const auto& [key, value_json] : properties) { + if (standard_keys.count(key)) continue; + + if (!first) result += ","; + first = false; + + size_t prop_pos = 0; + auto prop_obj = parse_json_object_raw(value_json, prop_pos); + + result += key + ":{"; + + if (prop_obj.count("description")) { + std::string desc = get_json_string_value(prop_obj["description"], 0); + result += "description:" + escape(desc); + } + + std::string type_val; + if (prop_obj.count("type")) { + type_val = get_json_string_value(prop_obj["type"], 0); + } + + if (to_upper(type_val) == "STRING") { + if (prop_obj.count("enum")) { + size_t enum_pos = 0; + std::string enum_formatted = format_argument(prop_obj["enum"], enum_pos, true); + result += ",enum:" + enum_formatted; + } + } else if (to_upper(type_val) == "OBJECT") { + if (prop_obj.count("properties")) { + std::string nested_required; + if (prop_obj.count("required")) { + nested_required = prop_obj["required"]; + } + result += ",properties:{" + format_parameters(prop_obj["properties"], nested_required) + "}"; + } + if (prop_obj.count("required")) { + std::string req_items; + size_t req_pos = 0; + skip_whitespace(prop_obj["required"], req_pos); + if (req_pos < prop_obj["required"].length() && prop_obj["required"][req_pos] == '[') { + req_pos++; + bool req_first = true; + while (req_pos < prop_obj["required"].length()) { + skip_whitespace(prop_obj["required"], req_pos); + if (prop_obj["required"][req_pos] == ']') break; + if (prop_obj["required"][req_pos] == ',') { req_pos++; continue; } + if (prop_obj["required"][req_pos] == '"') { + req_pos++; + std::string req_item = extract_json_string(prop_obj["required"], req_pos); + if (!req_first) req_items += ","; + req_first = false; + req_items += escape(req_item); + } + } + } + if (!req_items.empty()) { + result += ",required:[" + req_items + "]"; + } + } + } else if (to_upper(type_val) == "ARRAY") { + if (prop_obj.count("items")) { + result += ",items:{"; + size_t items_pos = 0; + auto items_obj = parse_json_object_raw(prop_obj["items"], items_pos); + bool items_first = true; + + for (const auto& [item_key, item_value] : items_obj) { + if (!items_first) result += ","; + items_first = false; + + if (item_key == "properties") { + std::string items_required; + if (items_obj.count("required")) { + items_required = items_obj["required"]; + } + result += "properties:{" + format_parameters(item_value, items_required) + "}"; + } else if (item_key == "required") { + result += "required:["; + size_t req_pos = 0; + skip_whitespace(item_value, req_pos); + if (req_pos < item_value.length() && item_value[req_pos] == '[') { + req_pos++; + bool req_first = true; + while (req_pos < item_value.length()) { + skip_whitespace(item_value, req_pos); + if (item_value[req_pos] == ']') break; + if (item_value[req_pos] == ',') { req_pos++; continue; } + if (item_value[req_pos] == '"') { + req_pos++; + std::string req_item = extract_json_string(item_value, req_pos); + if (!req_first) result += ","; + req_first = false; + result += escape(req_item); + } + } + } + result += "]"; + } else if (item_key == "type") { + std::string item_type = get_json_string_value(item_value, 0); + result += "type:" + escape(to_upper(item_type)); + } else { + size_t val_pos = 0; + result += item_key + ":" + format_argument(item_value, val_pos, true); + } + } + result += "}"; + } + } + + if (!type_val.empty()) { + result += ",type:" + escape(to_upper(type_val)); + } + + result += "}"; + } + + return result; +} + +inline std::string format_function_declaration(const std::string& name, + const std::string& description, + const std::string& params_json) { + std::string result = "declaration:" + name + "{"; + result += "description:" + escape(description); + + if (!params_json.empty()) { + result += ",parameters:{"; + + size_t pos = 0; + auto params = parse_json_object_raw(params_json, pos); + + if (params.count("properties")) { + std::string required_json; + if (params.count("required")) { + required_json = params["required"]; + } + result += "properties:{" + format_parameters(params["properties"], required_json) + "}"; + } + + if (params.count("required")) { + std::string req_items; + size_t req_pos = 0; + skip_whitespace(params["required"], req_pos); + if (req_pos < params["required"].length() && params["required"][req_pos] == '[') { + req_pos++; + bool first = true; + while (req_pos < params["required"].length()) { + skip_whitespace(params["required"], req_pos); + if (params["required"][req_pos] == ']') break; + if (params["required"][req_pos] == ',') { req_pos++; continue; } + if (params["required"][req_pos] == '"') { + req_pos++; + std::string item = extract_json_string(params["required"], req_pos); + if (!first) req_items += ","; + first = false; + req_items += escape(item); + } + } + } + if (!req_items.empty()) { + result += ",required:[" + req_items + "]"; + } + } + + if (params.count("type")) { + std::string type_val = get_json_string_value(params["type"], 0); + result += ",type:" + escape(to_upper(type_val)); + } + + result += "}"; + } + + result += "}"; + return result; +} + +template +inline std::string format_tools(const std::vector& tools, bool use_pipe_tags = false) { + if (tools.empty()) return ""; + + const char* decl_start = use_pipe_tags ? "<|tool>" : ""; + const char* decl_end = use_pipe_tags ? "" : ""; + + std::string result; + for (const auto& tool : tools) { + result += decl_start; + std::string params_json; + auto it = tool.parameters.find("schema"); + if (it != tool.parameters.end()) { + params_json = it->second; + } + + std::string declaration = format_function_declaration(tool.name, tool.description, params_json); + result += use_pipe_tags ? declaration : use_escape_tags(declaration); + result += decl_end; + } + return result; +} + + +inline size_t match_quote_tag(const std::string& s, size_t pos) { + if (s.compare(pos, 8, "") == 0) return 8; + if (s.compare(pos, 5, "<|\"|>") == 0) return 5; + return 0; +} + +inline size_t find_quote_tag(const std::string& s, size_t pos) { + size_t e = s.find("", pos); + size_t t = s.find("<|\"|>", pos); + if (e == std::string::npos) return t; + if (t == std::string::npos) return e; + return std::min(e, t); +} + +inline std::string unescape(const std::string& s) { + const std::string ESCAPE_TAG = ""; + std::string result = s; + size_t pos = 0; + while ((pos = result.find(ESCAPE_TAG, pos)) != std::string::npos) { + result.erase(pos, ESCAPE_TAG.length()); + } + return result; +} + +inline std::string to_json_value(const std::string& v) { + if (v == "true" || v == "false" || v == "null") return v; + if (!v.empty()) { + char* end = nullptr; + std::strtod(v.c_str(), &end); + if (end == v.c_str() + v.size()) return v; + } + return "\"" + escape_json_string(v) + "\""; +} + +inline void append_utf8(std::string& out, unsigned cp) { + if (cp < 0x80) { out += static_cast(cp); } + else if (cp < 0x800) { out += static_cast(0xC0 | (cp >> 6)); out += static_cast(0x80 | (cp & 0x3F)); } + else if (cp < 0x10000) { out += static_cast(0xE0 | (cp >> 12)); out += static_cast(0x80 | ((cp >> 6) & 0x3F)); out += static_cast(0x80 | (cp & 0x3F)); } + else { out += static_cast(0xF0 | (cp >> 18)); out += static_cast(0x80 | ((cp >> 12) & 0x3F)); out += static_cast(0x80 | ((cp >> 6) & 0x3F)); out += static_cast(0x80 | (cp & 0x3F)); } +} + +inline std::string read_quoted_json_string(const std::string& s, size_t& pos, char quote = '"') { + std::string out; + if (pos < s.size() && s[pos] == quote) pos++; + auto hex4 = [&](size_t p, unsigned& v) -> bool { + if (p + 4 > s.size()) return false; + v = 0; + for (int k = 0; k < 4; ++k) { + char h = s[p + k]; v <<= 4; + if (h >= '0' && h <= '9') v |= static_cast(h - '0'); + else if (h >= 'a' && h <= 'f') v |= static_cast(h - 'a' + 10); + else if (h >= 'A' && h <= 'F') v |= static_cast(h - 'A' + 10); + else return false; + } + return true; + }; + while (pos < s.size()) { + char c = s[pos]; + if (c == quote) { pos++; break; } + if (c == '\\' && pos + 1 < s.size()) { + char n = s[pos + 1]; + switch (n) { + case '\'': out += '\''; pos += 2; break; + case 'n': out += '\n'; pos += 2; break; + case 't': out += '\t'; pos += 2; break; + case 'r': out += '\r'; pos += 2; break; + case 'b': out += '\b'; pos += 2; break; + case 'f': out += '\f'; pos += 2; break; + case '"': out += '"'; pos += 2; break; + case '\\': out += '\\'; pos += 2; break; + case '/': out += '/'; pos += 2; break; + case 'u': { + unsigned cp; + if (hex4(pos + 2, cp)) { + size_t adv = pos + 6; + if (cp >= 0xD800 && cp <= 0xDBFF) { + unsigned lo; + if (adv + 6 <= s.size() && s[adv] == '\\' && s[adv + 1] == 'u' && + hex4(adv + 2, lo) && lo >= 0xDC00 && lo <= 0xDFFF) { + cp = 0x10000 + ((cp - 0xD800) << 10) + (lo - 0xDC00); adv += 6; + } else { cp = 0xFFFD; } + } else if (cp >= 0xDC00 && cp <= 0xDFFF) { cp = 0xFFFD; } + append_utf8(out, cp); pos = adv; + } else { out += '\\'; out += 'u'; pos += 2; } + break; + } + default: out += n; pos += 2; break; + } + } else { + out += c; pos++; + } + } + return out; +} + +inline std::string repair_json(const std::string& s) { + std::string out; + out.reserve(s.size() + 16); + bool in_str = false, esc = false; + for (size_t i = 0; i < s.size(); ++i) { + char c = s[i]; + if (in_str) { + if (esc) { out += c; esc = false; continue; } + if (c == '\\') { out += c; esc = true; continue; } + if (c == '"') { out += c; in_str = false; continue; } + unsigned char uc = static_cast(c); + if (uc == '\n') out += "\\n"; + else if (uc == '\t') out += "\\t"; + else if (uc == '\r') out += "\\r"; + else if (uc < 0x20) { char b[8]; std::snprintf(b, sizeof(b), "\\u%04x", uc); out += b; } + else out += c; + } else { + if (c == '"') { in_str = true; out += c; } + else if (c == ',') { + size_t j = i + 1; + while (j < s.size() && std::isspace(static_cast(s[j]))) j++; + bool trailing = (j < s.size() && (s[j] == '}' || s[j] == ']')); + if (!trailing) out += c; + } else out += c; + } + } + return out; +} + +inline size_t find_raw_value_end(const std::string& s, size_t pos, + const std::vector& known_keys) { + int brace = 0, brack = 0; + for (size_t i = pos; i < s.size(); ++i) { + char c = s[i]; + if (c == '{') brace++; + else if (c == '}') { if (brace == 0) return i; brace--; } + else if (c == '[') brack++; + else if (c == ']') { if (brack > 0) brack--; } + else if (c == ',' && brace == 0 && brack == 0) { + size_t j = i + 1; + while (j < s.size() && std::isspace(static_cast(s[j]))) j++; + bool quoted = (j < s.size() && s[j] == '"'); + size_t ks = quoted ? j + 1 : j; + for (const auto& k : known_keys) { + if (k.empty() || s.compare(ks, k.size(), k) != 0) continue; + size_t m = ks + k.size(); + if (quoted) { if (m < s.size() && s[m] == '"') m++; else continue; } + while (m < s.size() && std::isspace(static_cast(s[m]))) m++; + if (m < s.size() && s[m] == ':') return i; + } + } + } + return s.size(); +} + +inline std::string args_to_json(const std::string& args_content, + const std::vector& known_keys = {}, + const std::set& string_keys = {}); +inline std::string parse_array_items(const std::string& inner); + +inline std::string try_json_object(const std::string& s) { + auto attempt = [](const std::string& c) -> std::string { + picojson::value v; + const char* b = c.data(); + const char* e = b + c.size(); + std::string err; + const char* end = picojson::parse(v, b, e, &err); + while (end < e && std::isspace(static_cast(*end))) ++end; + return (err.empty() && end == e && v.is()) ? v.serialize() : std::string(); + }; + std::string r = attempt(s); + if (!r.empty()) return r; + std::string repaired = repair_json(s); + return (repaired != s) ? attempt(repaired) : std::string(); +} + +inline bool skip_delimited_string(const std::string& s, size_t& i) { + if (size_t qtag = match_quote_tag(s, i)) { + size_t end = find_quote_tag(s, i + qtag); + i = (end == std::string::npos) ? s.size() : end + match_quote_tag(s, end); + return true; + } + char q = (i < s.size()) ? s[i] : '\0'; + if (q == '"' || q == '\'') { + for (i++; i < s.size(); ++i) { + if (s[i] == '\\') { i++; continue; } + if (s[i] == q) { i++; break; } + } + return true; + } + return false; +} + +inline std::string read_balanced(const std::string& s, size_t& pos, char open, char close) { + size_t start = pos; + int depth = 0; + while (pos < s.size()) { + if (skip_delimited_string(s, pos)) continue; + char c = s[pos++]; + if (c == open) depth++; + else if (c == close && --depth == 0) break; + } + return s.substr(start, pos - start); +} + +inline bool read_delimited_string(const std::string& s, size_t& pos, std::string& out) { + if (size_t qtag = match_quote_tag(s, pos)) { + pos += qtag; + size_t end = find_quote_tag(s, pos); + size_t stop = (end == std::string::npos) ? s.size() : end; + out = "\"" + escape_json_string(s.substr(pos, stop - pos)) + "\""; + pos = (end == std::string::npos) ? s.size() : end + match_quote_tag(s, end); + return true; + } + if (pos < s.size() && (s[pos] == '"' || s[pos] == '\'')) { + out = "\"" + escape_json_string(read_quoted_json_string(s, pos, s[pos])) + "\""; + return true; + } + return false; +} + +inline std::string read_raw_value(const std::string& s, size_t& pos, + const std::vector& known_keys, bool want_string) { + size_t start = pos; + size_t end = known_keys.empty() ? s.find_first_of(",}", pos) + : find_raw_value_end(s, pos, known_keys); + if (end == std::string::npos) end = s.size(); + std::string raw = s.substr(start, end - start); + pos = end; + + std::string trimmed = raw; + while (!trimmed.empty() && std::isspace(static_cast(trimmed.back()))) trimmed.pop_back(); + bool scalar = !want_string && (trimmed == "true" || trimmed == "false" || trimmed == "null"); + if (!want_string && !scalar && !trimmed.empty()) { + char* ep = nullptr; + std::strtod(trimmed.c_str(), &ep); + scalar = (ep == trimmed.c_str() + trimmed.size()); + } + return scalar ? trimmed : ("\"" + escape_json_string(raw) + "\""); +} + +inline std::string parse_array_items(const std::string& inner) { + std::string out = "["; + size_t pos = 0; + bool first = true; + while (pos < inner.size()) { + while (pos < inner.size() && (std::isspace(static_cast(inner[pos])) || inner[pos] == ',')) pos++; + if (pos >= inner.size()) break; + if (!first) out += ","; + first = false; + + std::string item; + if (read_delimited_string(inner, pos, item)) { + out += item; + } else if (inner[pos] == '{') { + out += args_to_json(read_balanced(inner, pos, '{', '}')); + } else if (inner[pos] == '[') { + std::string nested = read_balanced(inner, pos, '[', ']'); + out += parse_array_items(nested.substr(1, nested.size() - 2)); + } else { + out += read_raw_value(inner, pos, {}, false); + } + } + out += "]"; + return out; +} + +inline std::string args_to_json(const std::string& args_content, + const std::vector& known_keys, + const std::set& string_keys) { + if (std::string json = try_json_object(args_content); !json.empty()) return json; + + std::string result = "{"; + size_t pos = (!args_content.empty() && args_content[0] == '{') ? 1 : 0; + bool first = true; + + while (pos < args_content.size()) { + while (pos < args_content.size() && std::isspace(static_cast(args_content[pos]))) pos++; + if (pos >= args_content.size() || args_content[pos] == '}') break; + if (args_content[pos] == ',') { pos++; continue; } + + std::string key; + if (args_content[pos] == '"') { + key = read_quoted_json_string(args_content, pos); + } else { + size_t ks = pos; + while (pos < args_content.size() && args_content[pos] != ':') pos++; + key = args_content.substr(ks, pos - ks); + while (!key.empty() && std::isspace(static_cast(key.back()))) key.pop_back(); + } + while (pos < args_content.size() && args_content[pos] != ':') pos++; + if (pos < args_content.size()) pos++; + while (pos < args_content.size() && std::isspace(static_cast(args_content[pos]))) pos++; + + const bool want_string = string_keys.count(key) > 0; + std::string value; + if (pos < args_content.size()) { + if (read_delimited_string(args_content, pos, value)) { + } else if (!want_string && args_content[pos] == '{') { + value = args_to_json(read_balanced(args_content, pos, '{', '}')); + } else if (!want_string && args_content[pos] == '[') { + std::string arr = read_balanced(args_content, pos, '[', ']'); + value = parse_array_items(arr.substr(1, arr.size() - 2)); + } else { + value = read_raw_value(args_content, pos, known_keys, want_string); + } + } + + if (!first) result += ","; + first = false; + result += "\"" + escape_json_string(key) + "\":" + (value.empty() ? "\"\"" : value); + } + + result += "}"; + return result; +} + +inline void parse_function_calls_core(std::string& response, std::vector& function_calls, + const std::map>& fn_keys, + const std::map>& fn_strkeys = {}) { + const std::string CALL_START = (response.find("<|tool_call>") != std::string::npos) + ? "<|tool_call>" : ""; + const std::string CALL_END = (CALL_START == "<|tool_call>") + ? "" : ""; + size_t pos = 0; + + auto emit_call = [&](const std::string& fn, std::string args_content) { + if (args_content.empty() || args_content.back() != '}') args_content += "}"; + std::vector keys; + std::set strkeys; + if (auto it = fn_keys.find(fn); it != fn_keys.end()) keys = it->second; + if (auto it = fn_strkeys.find(fn); it != fn_strkeys.end()) strkeys = it->second; + std::string args_json = args_to_json(args_content, keys, strkeys); + function_calls.push_back("{\"name\":\"" + escape_json_string(fn) + "\",\"arguments\":" + args_json + "}"); + }; + + while ((pos = response.find(CALL_START, pos)) != std::string::npos) { + size_t content_start = pos + CALL_START.length(); + while (content_start < response.length() && + std::isspace(static_cast(response[content_start]))) { + content_start++; + } + + const bool is_call = response.compare(content_start, 5, "call:") == 0; + size_t scan_from = content_start; + if (is_call) { + size_t brace_pos = response.find('{', content_start + 5); + if (brace_pos != std::string::npos) { + size_t bp = brace_pos; + read_balanced(response, bp, '{', '}'); + scan_from = bp; + } + } + size_t call_end_pos = response.find(CALL_END, scan_from); + + size_t content_end = (call_end_pos != std::string::npos) ? call_end_pos : response.length(); + std::string call_content = response.substr(content_start, content_end - content_start); + + if (is_call) { + size_t brace_pos = call_content.find('{'); + if (brace_pos != std::string::npos) { + emit_call(call_content.substr(5, brace_pos - 5), call_content.substr(brace_pos)); + } else { + size_t sep_pos = call_content.find_first_of(", ", 5); + if (sep_pos != std::string::npos) { + size_t args_start = call_content.find_first_not_of(", ", sep_pos); + emit_call(call_content.substr(5, sep_pos - 5), + "{" + (args_start == std::string::npos ? std::string() : call_content.substr(args_start))); + } + } + } + + size_t erase_end = (call_end_pos != std::string::npos) ? + call_end_pos + CALL_END.length() : response.length(); + response.erase(pos, erase_end - pos); + } +} + +inline void parse_function_calls(std::string& response, std::vector& function_calls) { + parse_function_calls_core(response, function_calls, {}); +} + +template +inline void parse_function_calls(std::string& response, std::vector& function_calls, + const std::vector& tools) { + std::map> fn_keys; + std::map> fn_strkeys; + for (const auto& t : tools) { + auto it = t.parameters.find("schema"); + if (it == t.parameters.end()) continue; + picojson::value v; + if (!picojson::parse(v, it->second).empty() || !v.is()) continue; + const auto& obj = v.get(); + auto pit = obj.find("properties"); + if (pit == obj.end() || !pit->second.is()) continue; + std::vector names; + std::set strkeys; + for (const auto& kv : pit->second.get()) { + names.push_back(kv.first); + if (kv.second.is()) { + const auto& prop = kv.second.get(); + auto tit = prop.find("type"); + if (tit != prop.end() && tit->second.is() && + tit->second.get() == "string") { + strkeys.insert(kv.first); + } + } + } + fn_keys[t.name] = std::move(names); + fn_strkeys[t.name] = std::move(strkeys); + } + parse_function_calls_core(response, function_calls, fn_keys, fn_strkeys); +} + +} // namespace gemma \ No newline at end of file diff --git a/cactus/engine/engine_index.cpp b/cactus-engine/src/index.cpp similarity index 95% rename from cactus/engine/engine_index.cpp rename to cactus-engine/src/index.cpp index 105a915f9..940a476e4 100644 --- a/cactus/engine/engine_index.cpp +++ b/cactus-engine/src/index.cpp @@ -1,5 +1,5 @@ #include "engine.h" -#include "kernel/kernel.h" +#include "cactus_kernels.h" #include #include #include @@ -8,6 +8,9 @@ #include #include #include +#include +#include +#include namespace cactus { namespace engine { @@ -140,6 +143,7 @@ namespace index { } void Index::add_documents(const std::vector& documents) { + if (mapped_index_ == MAP_FAILED || mapped_data_ == MAP_FAILED) throw std::runtime_error("Index is in a failed memory-mapped state"); validate_documents(documents); size_t added_data_size = 0; @@ -161,20 +165,24 @@ namespace index { munmap(mapped_index_, index_file_size_); munmap(mapped_data_, data_file_size_); - mapped_index_ = mmap(nullptr, new_index_size, PROT_READ | PROT_WRITE, MAP_SHARED, index_fd_, 0); - if (mapped_index_ == MAP_FAILED) { + void* new_index_map = mmap(nullptr, new_index_size, PROT_READ | PROT_WRITE, MAP_SHARED, index_fd_, 0); + if (new_index_map == MAP_FAILED) { + mapped_index_ = mapped_data_ = MAP_FAILED; throw std::runtime_error("Failed to remap index file"); } - mapped_data_ = mmap(nullptr, new_data_size, PROT_READ | PROT_WRITE, MAP_SHARED, data_fd_, 0); - if (mapped_data_ == MAP_FAILED) { - munmap(mapped_index_, new_index_size); + void* new_data_map = mmap(nullptr, new_data_size, PROT_READ | PROT_WRITE, MAP_SHARED, data_fd_, 0); + if (new_data_map == MAP_FAILED) { + munmap(new_index_map, new_index_size); + mapped_index_ = mapped_data_ = MAP_FAILED; throw std::runtime_error("Failed to remap data file"); } size_t old_index_size = index_file_size_; size_t old_data_size = data_file_size_; + mapped_index_ = new_index_map; + mapped_data_ = new_data_map; index_file_size_ = new_index_size; data_file_size_ = new_data_size; @@ -229,6 +237,7 @@ namespace index { } void Index::delete_documents(const std::vector& doc_ids) { + if (mapped_index_ == MAP_FAILED || mapped_data_ == MAP_FAILED) throw std::runtime_error("Index is in a failed memory-mapped state"); validate_doc_ids(doc_ids); char* index_ptr = static_cast(mapped_index_); @@ -255,6 +264,7 @@ namespace index { } std::vector Index::get_documents(const std::vector& doc_ids) { + if (mapped_index_ == MAP_FAILED || mapped_data_ == MAP_FAILED) throw std::runtime_error("Index is in a failed memory-mapped state"); validate_doc_ids(doc_ids); const char* index_ptr = static_cast(mapped_index_); @@ -304,11 +314,18 @@ namespace index { if (embeddings.empty()) { return {}; } + if (mapped_index_ == MAP_FAILED || mapped_data_ == MAP_FAILED) throw std::runtime_error("Index is in a failed memory-mapped state"); if (options.top_k == 0) { return std::vector>(embeddings.size()); } + for (const auto& embedding : embeddings) { + if (embedding.size() != embedding_dim_) { + throw std::runtime_error("Query embedding dimension mismatch"); + } + } + if (num_documents_ > 0) { size_t last_entry_offset = (num_documents_ - 1) * index_entry_size_; if (sizeof(IndexHeader) + last_entry_offset + index_entry_size_ > index_file_size_) { @@ -378,6 +395,7 @@ namespace index { } void Index::compact() { + if (mapped_index_ == MAP_FAILED || mapped_data_ == MAP_FAILED) throw std::runtime_error("Index is in a failed memory-mapped state"); std::string temp_index_path = index_path_ + ".tmp"; std::string temp_data_path = data_path_ + ".tmp"; diff --git a/cactus/ffi/cactus_index.cpp b/cactus-engine/src/index_ffi.cpp similarity index 99% rename from cactus/ffi/cactus_index.cpp rename to cactus-engine/src/index_ffi.cpp index a81f62377..537d817a7 100644 --- a/cactus/ffi/cactus_index.cpp +++ b/cactus-engine/src/index_ffi.cpp @@ -1,5 +1,5 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" +#include "../cactus_engine.h" +#include "utils.h" #include using namespace cactus::ffi; diff --git a/cactus/ffi/cactus_init.cpp b/cactus-engine/src/init.cpp similarity index 82% rename from cactus/ffi/cactus_init.cpp rename to cactus-engine/src/init.cpp index 64df60b63..ca11df7fa 100644 --- a/cactus/ffi/cactus_init.cpp +++ b/cactus-engine/src/init.cpp @@ -1,6 +1,6 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" -#include "cactus_telemetry.h" +#include "../cactus_engine.h" +#include "utils.h" +#include "telemetry.h" #include #include #include @@ -8,6 +8,8 @@ #include #include #include +#include +#include using namespace cactus::engine; using namespace cactus::ffi; @@ -16,6 +18,12 @@ static constexpr size_t RAG_MAX_CHUNK_TOKENS = 128; static constexpr size_t RAG_MIN_CHUNK_TOKENS = 24; static constexpr size_t RAG_CHUNK_OVERLAP = 32; +static void apply_no_cloud_telemetry_env() { + if (cactus::ffi::env_flag_enabled("CACTUS_NO_CLOUD_TELE")) { + cactus::telemetry::setCloudDisabled(true); + } +} + static time_t get_file_mtime(const std::string& path) { struct stat st; if (stat(path.c_str(), &st) == 0) { @@ -229,7 +237,7 @@ static bool build_corpus_index(CactusModelHandle* handle, const std::string& cor CACTUS_LOG_INFO("init", "Generated " << chunks.size() << " chunks from corpus"); std::vector test_tokens = tokenizer->encode("test"); - std::vector test_embedding = handle->model->get_embeddings(test_tokens, true, true); + std::vector test_embedding = handle->model->get_text_embeddings(test_tokens, true); if (test_embedding.empty()) { CACTUS_LOG_ERROR("init", "Failed to get embedding dimension"); return false; @@ -242,6 +250,9 @@ static bool build_corpus_index(CactusModelHandle* handle, const std::string& cor std::string index_path = corpus_dir + "/index.bin"; std::string data_path = corpus_dir + "/data.bin"; + std::remove(index_path.c_str()); + std::remove(data_path.c_str()); + try { handle->corpus_index = std::make_unique(index_path, data_path, embedding_dim); } catch (const std::exception& e) { @@ -256,7 +267,7 @@ static bool build_corpus_index(CactusModelHandle* handle, const std::string& cor const auto& [chunk_text, source_file] = chunks[i]; std::vector tokens = tokenizer->encode(chunk_text); - std::vector embedding = handle->model->get_embeddings(tokens, true, true); + std::vector embedding = handle->model->get_text_embeddings(tokens, true); if (embedding.size() != embedding_dim) { CACTUS_LOG_WARN("init", "Skipping chunk " << i << " - embedding dimension mismatch"); @@ -306,7 +317,7 @@ static bool load_corpus_index(CactusModelHandle* handle, const std::string& corp auto* tokenizer = handle->model->get_tokenizer(); std::vector test_tokens = tokenizer->encode("test"); - std::vector test_embedding = handle->model->get_embeddings(test_tokens, true, true); + std::vector test_embedding = handle->model->get_text_embeddings(test_tokens, true); if (test_embedding.empty()) { CACTUS_LOG_ERROR("init", "Failed to get embedding dimension for index loading"); return false; @@ -324,8 +335,6 @@ static bool load_corpus_index(CactusModelHandle* handle, const std::string& corp } } -std::string last_error_message; - bool matches_stop_sequence(const std::vector& generated_tokens, const std::vector>& stop_sequences) { for (const auto& stop_seq : stop_sequences) { @@ -340,13 +349,9 @@ bool matches_stop_sequence(const std::vector& generated_tokens, extern "C" { -const char* cactus_get_last_error() { - return last_error_message.c_str(); -} - -cactus_model_t cactus_init(const char* model_path, const char* corpus_dir) { - CactusTelemetry::getInstance().ensureInitialized(); +int cactus_set_backend(const char* backend) { return cactus_backend_select(backend); } +cactus_model_t cactus_init(const char* model_path, const char* corpus_dir, bool cache_index) { constexpr size_t DEFAULT_CONTEXT_SIZE = 512; // matches default sliding window size std::string model_path_str = model_path ? std::string(model_path) : "unknown"; @@ -359,19 +364,23 @@ cactus_model_t cactus_init(const char* model_path, const char* corpus_dir) { CACTUS_LOG_INFO("init", "Loading model: " << model_name << " from " << model_path_str); + apply_no_cloud_telemetry_env(); + cactus::telemetry::init(nullptr, model_path_str.c_str(), nullptr); + + auto __cactus_init_start = std::chrono::steady_clock::now(); + try { auto* handle = new CactusModelHandle(); - handle->model = create_model(model_path); + handle->model = create_model(model_path_str); handle->model_name = model_name; if (!handle->model) { last_error_message = "Failed to create model - check config.txt exists at: " + model_path_str; CACTUS_LOG_ERROR("init", last_error_message); - - CactusTelemetry::getInstance().recordInit( - model_name, false, last_error_message - ); - + { + auto __cactus_init_err_dur = std::chrono::duration_cast(std::chrono::steady_clock::now() - __cactus_init_start).count(); + cactus::telemetry::recordInit(model_name.c_str(), false, static_cast(__cactus_init_err_dur), last_error_message.c_str()); + } delete handle; return nullptr; } @@ -379,11 +388,10 @@ cactus_model_t cactus_init(const char* model_path, const char* corpus_dir) { if (!handle->model->init(model_path, DEFAULT_CONTEXT_SIZE)) { last_error_message = "Failed to initialize model - check weight files at: " + model_path_str; CACTUS_LOG_ERROR("init", last_error_message); - - CactusTelemetry::getInstance().recordInit( - model_name, false, last_error_message - ); - + { + auto __cactus_init_err_dur = std::chrono::duration_cast(std::chrono::steady_clock::now() - __cactus_init_start).count(); + cactus::telemetry::recordInit(model_name.c_str(), false, static_cast(__cactus_init_err_dur), last_error_message.c_str()); + } delete handle; return nullptr; } @@ -391,8 +399,13 @@ cactus_model_t cactus_init(const char* model_path, const char* corpus_dir) { if (corpus_dir != nullptr && strlen(corpus_dir) > 0) { handle->corpus_dir = std::string(corpus_dir); - if (!load_corpus_index(handle, handle->corpus_dir)) { - CACTUS_LOG_INFO("init", "No existing index found, building new corpus index"); + bool loaded = false; + if (cache_index) { + loaded = load_corpus_index(handle, handle->corpus_dir); + } + + if (!loaded) { + CACTUS_LOG_INFO("init", (cache_index ? "No existing index found, building new corpus index" : "Building fresh corpus index (caching disabled)")); if (!build_corpus_index(handle, handle->corpus_dir)) { CACTUS_LOG_WARN("init", "Failed to build corpus index - RAG disabled"); } @@ -400,29 +413,27 @@ cactus_model_t cactus_init(const char* model_path, const char* corpus_dir) { } CACTUS_LOG_INFO("init", "Model loaded successfully: " << model_name); - - CactusTelemetry::getInstance().recordInit( - model_name, true, "Model initialized successfully" - ); + { + auto __cactus_init_ok_dur = std::chrono::duration_cast(std::chrono::steady_clock::now() - __cactus_init_start).count(); + cactus::telemetry::recordInit(model_name.c_str(), true, static_cast(__cactus_init_ok_dur), ""); + } return handle; } catch (const std::exception& e) { last_error_message = "Exception during init: " + std::string(e.what()); CACTUS_LOG_ERROR("init", last_error_message); - - CactusTelemetry::getInstance().recordInit( - model_name, false, last_error_message - ); - + { + auto __cactus_init_err_dur = std::chrono::duration_cast(std::chrono::steady_clock::now() - __cactus_init_start).count(); + cactus::telemetry::recordInit(model_name.c_str(), false, static_cast(__cactus_init_err_dur), last_error_message.c_str()); + } return nullptr; } catch (...) { last_error_message = "Unknown exception during model initialization"; CACTUS_LOG_ERROR("init", last_error_message); - - CactusTelemetry::getInstance().recordInit( - model_name, false, last_error_message - ); - + { + auto __cactus_init_err_dur = std::chrono::duration_cast(std::chrono::steady_clock::now() - __cactus_init_start).count(); + cactus::telemetry::recordInit(model_name.c_str(), false, static_cast(__cactus_init_err_dur), last_error_message.c_str()); + } return nullptr; } } @@ -436,6 +447,8 @@ void cactus_reset(cactus_model_t model) { auto* handle = static_cast(model); handle->model->reset_cache(); handle->processed_tokens.clear(); + handle->processed_images.clear(); + handle->user_audio_counts.clear(); } void cactus_stop(cactus_model_t model) { diff --git a/cactus-engine/src/json_escape.h b/cactus-engine/src/json_escape.h new file mode 100644 index 000000000..c650fa8a3 --- /dev/null +++ b/cactus-engine/src/json_escape.h @@ -0,0 +1,37 @@ +#ifndef CACTUS_JSON_ESCAPE_H +#define CACTUS_JSON_ESCAPE_H + +#include +#include + +inline std::string escape_json_string(const std::string& s) { + std::string result; + result.reserve(s.size()); + for (unsigned char c : s) { + switch (c) { + case '"': result += "\\\""; break; + case '\\': result += "\\\\"; break; + case '\b': result += "\\b"; break; + case '\f': result += "\\f"; break; + case '\n': result += "\\n"; break; + case '\r': result += "\\r"; break; + case '\t': result += "\\t"; break; + default: + if (c < 0x20) { + char buf[7]; + std::snprintf(buf, sizeof(buf), "\\u%04x", c); + result += buf; + } else { + result += static_cast(c); + } + break; + } + } + return result; +} + +inline std::string escape_json_string(const char* str) { + return str ? escape_json_string(std::string(str)) : std::string(); +} + +#endif diff --git a/cactus-engine/src/kv_compress.cpp b/cactus-engine/src/kv_compress.cpp new file mode 100644 index 000000000..eff9c4eb6 --- /dev/null +++ b/cactus-engine/src/kv_compress.cpp @@ -0,0 +1,482 @@ +#include "kv_compress.h" + +#include +#include +#include +#include +#include + +#if defined(__ARM_NEON) +#include +#define CACTUS_KV_NEON 1 +#else +#define CACTUS_KV_NEON 0 +#endif + +namespace cactus { +namespace kvcompress { + +namespace { +// Toggled off by tests to exercise/compare the scalar fallback on a NEON build. +bool g_use_simd = true; +} + +void kv_set_simd(bool on) { g_use_simd = on; } + +namespace { + +// Round half-to-even like numpy. recent_frac widens from float to a slightly-off double, so snap +// near-halves first to avoid flipping an exact .5 boundary the Python reference keeps. +long py_round(double x) { + double twice = x * 2.0; + double twice_rounded = std::nearbyint(twice); + if (std::fabs(twice - twice_rounded) < 1e-6) x = twice_rounded * 0.5; + double r = std::nearbyint(x); // FE_TONEAREST default == round-half-to-even + return static_cast(r); +} + +} // namespace + +void keydiff_score(const float* keys, size_t n, size_t head_dim, float* out) { + // double accumulation matches the float64 reference. + std::vector mu(head_dim, 0.0); + for (size_t i = 0; i < n; ++i) { + const float* k = keys + i * head_dim; + for (size_t d = 0; d < head_dim; ++d) mu[d] += static_cast(k[d]); + } + double mu_norm = 0.0; + for (size_t d = 0; d < head_dim; ++d) { + mu[d] /= static_cast(n); + mu_norm += mu[d] * mu[d]; + } + mu_norm = std::sqrt(mu_norm) + 1e-8; + for (size_t d = 0; d < head_dim; ++d) mu[d] /= mu_norm; + + for (size_t i = 0; i < n; ++i) { + const float* k = keys + i * head_dim; + double knorm = 0.0, dot = 0.0; + for (size_t d = 0; d < head_dim; ++d) { + double v = static_cast(k[d]); + knorm += v * v; + dot += v * mu[d]; + } + knorm = std::sqrt(knorm) + 1e-8; + out[i] = static_cast(-(dot / knorm)); + } +} + +std::vector keepset_for_head(const float* scores, size_t n, const Params& p) { + long budget = std::max(1, static_cast(p.abs_budget)); + long B = std::min(budget, static_cast(n)); + long sink = std::min(std::max(static_cast(p.sink), 0), static_cast(n)); + long n_recent = std::min(py_round(static_cast(p.recent_frac) * static_cast(B)), + static_cast(n)); + + auto in_range = [&](int idx) { return idx >= 0 && idx < static_cast(n); }; + std::set reserved; + for (long i = 0; i < sink; ++i) reserved.insert(i); + for (int idx : p.protect) if (in_range(idx)) reserved.insert(static_cast(idx)); + for (long i = static_cast(n) - n_recent; i < static_cast(n); ++i) + if (i >= 0) reserved.insert(i); + + // Reserved may exceed B: keep sink, protected, then recent, until B. + if (static_cast(reserved.size()) > B) { + std::vector ordered; + for (long i = 0; i < sink; ++i) ordered.push_back(i); + for (int idx : p.protect) if (in_range(idx)) ordered.push_back(static_cast(idx)); + for (long i = static_cast(n) - 1; i >= sink; --i) ordered.push_back(i); + reserved.clear(); + for (long i : ordered) { + if (i >= 0 && i < static_cast(n)) reserved.insert(i); + if (static_cast(reserved.size()) == B) break; + } + } + + std::set keep(reserved); + long remaining = B - static_cast(keep.size()); + if (remaining > 0) { + std::vector order(n); + std::iota(order.begin(), order.end(), 0L); + std::stable_sort(order.begin(), order.end(), [&](long a, long b) { + return scores[a] > scores[b]; + }); + for (long i : order) { + if (keep.count(i)) continue; + keep.insert(i); + if (--remaining == 0) break; + } + } + + std::vector result; + result.reserve(keep.size()); + for (long i : keep) result.push_back(static_cast(i)); + return result; +} + +namespace { + +std::vector rope_inv_freq(size_t head_dim, double rope_theta) { + std::vector inv(head_dim / 2); + for (size_t i = 0; i < inv.size(); ++i) + inv[i] = std::pow(rope_theta, -(2.0 * static_cast(i)) / static_cast(head_dim)); + return inv; +} + +RopeRotation rope_rot(const std::vector& inv_freq, double delta_pos) { + RopeRotation r; + r.cos.resize(inv_freq.size()); + r.sin.resize(inv_freq.size()); + for (size_t i = 0; i < inv_freq.size(); ++i) { + double a = delta_pos * inv_freq[i]; + r.cos[i] = std::cos(a); + r.sin[i] = std::sin(a); + } + return r; +} + +void rope_apply(float* row, const RopeRotation& r) { + size_t half = r.cos.size(); + for (size_t i = 0; i < half; ++i) { + double x1 = row[i], x2 = row[i + half]; + row[i] = static_cast(x1 * r.cos[i] - x2 * r.sin[i]); + row[i + half] = static_cast(x2 * r.cos[i] + x1 * r.sin[i]); + } +} + +// Apply the inverse rotation (negate the table angle): re-rope by +t given the un-rope table[t]. +void rope_apply_conj(float* row, const RopeRotation& r) { + size_t half = r.cos.size(); + for (size_t i = 0; i < half; ++i) { + double x1 = row[i], x2 = row[i + half]; + row[i] = static_cast(x1 * r.cos[i] + x2 * r.sin[i]); + row[i + half] = static_cast(x2 * r.cos[i] - x1 * r.sin[i]); + } +} + +size_t max_kept_index(const std::vector>& kept) { + size_t m = 0; + for (const auto& kh : kept) for (int idx : kh) m = std::max(m, static_cast(idx)); + return m; +} + +} // namespace + +void rope_rotate_row(float* row, size_t head_dim, double rope_theta, double delta_pos) { + rope_apply(row, rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos)); +} + +std::vector unrope_table(size_t n, size_t head_dim, double rope_theta) { + auto inv = rope_inv_freq(head_dim, rope_theta); + std::vector table(n); + for (size_t t = 0; t < n; ++t) table[t] = rope_rot(inv, -static_cast(t)); + return table; +} + +namespace { + +// Matches the engine's KV quant convention (quantize_group_fp16_to_int8 in +// cactus-kernels/src/quants.cpp): scale floored at 1e-10, roundf, clamp to [-128, 127]. +void requant_row(const float* row, int8_t* dst, float* dsc, size_t head_dim, size_t group_size) { + size_t groups = (head_dim + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + size_t lo = g * group_size, hi = std::min(head_dim, lo + group_size); + float amax = 0.0f; + for (size_t d = lo; d < hi; ++d) amax = std::max(amax, std::fabs(row[d])); + float scale = amax / 127.0f; + if (scale < 1e-10f) scale = 1e-10f; + dsc[g] = scale; + float inv = 1.0f / scale; + for (size_t d = lo; d < hi; ++d) { + int32_t q = static_cast(std::roundf(row[d] * inv)); + q = std::max(-128, std::min(127, q)); + dst[d] = static_cast(q); + } + } +} + +void dequant_row(const int8_t* src, const float* scale, size_t head_dim, size_t group_size, + float* out) { +#if CACTUS_KV_NEON + if (g_use_simd) { + size_t groups = (head_dim + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + size_t lo = g * group_size, hi = std::min(head_dim, lo + group_size); + float32x4_t scv = vdupq_n_f32(scale[g]); + size_t d = lo; + for (; d + 4 <= hi; d += 4) { + int32_t four; + std::memcpy(&four, src + d, 4); + int16x8_t i16 = vmovl_s8(vreinterpret_s8_s32(vdup_n_s32(four))); + vst1q_f32(out + d, vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(i16))), scv)); + } + for (; d < hi; ++d) out[d] = static_cast(src[d]) * scale[g]; + } + return; + } +#endif + for (size_t d = 0; d < head_dim; ++d) out[d] = static_cast(src[d]) * scale[d / group_size]; +} + +void rotate_int8_row_rot(int8_t* int8, float* scale, size_t head_dim, size_t group_size, + const RopeRotation& rot) { + std::vector row(head_dim); + dequant_row(int8, scale, head_dim, group_size, row.data()); + rope_apply(row.data(), rot); + requant_row(row.data(), int8, scale, head_dim, group_size); +} + +// fill_post(h, post) gathers head h's post-RoPE rows; scoring is shared across fp16/int8. +template +std::vector> keepsets_per_head(size_t n, size_t kv_heads, size_t head_dim, + const std::vector& unrope, + const Params& p, FillPost fill_post, + const std::vector>& protect_per_head) { + const bool per_head_protect = protect_per_head.size() == kv_heads; + std::vector post(n * head_dim), pre(n * head_dim), scores(n); + std::vector> out; + out.reserve(kv_heads); + Params ph = p; + for (size_t h = 0; h < kv_heads; ++h) { + fill_post(h, post.data()); + for (size_t t = 0; t < n; ++t) { + const float* src = post.data() + t * head_dim; + float* dst = pre.data() + t * head_dim; + for (size_t d = 0; d < head_dim; ++d) dst[d] = src[d]; + rope_apply(dst, unrope[t]); + } + keydiff_score(pre.data(), n, head_dim, scores.data()); + if (per_head_protect) ph.protect = protect_per_head[h]; + out.push_back(keepset_for_head(scores.data(), n, per_head_protect ? ph : p)); + } + return out; +} + +} // namespace + +void compact_fp16(uint16_t* key_rows_u, uint16_t* val_rows_u, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, + const std::vector& unrope) { + __fp16* key_rows = reinterpret_cast<__fp16*>(key_rows_u); + __fp16* val_rows = reinterpret_cast<__fp16*>(val_rows_u); + std::vector krow(head_dim), vrow(head_dim); + for (size_t h = 0; h < kv_heads; ++h) { + const std::vector& kept = kept_per_head[h]; + for (size_t rank = 0; rank < kept.size(); ++rank) { + int abs_pos = kept[rank]; + const __fp16* ksrc = key_rows + (static_cast(abs_pos) * kv_heads + h) * head_dim; + const __fp16* vsrc = val_rows + (static_cast(abs_pos) * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { krow[d] = ksrc[d]; vrow[d] = vsrc[d]; } + // renumber abs_pos -> rank: un-rope by -abs_pos then re-rope by +rank (no per-row trig) + rope_apply(krow.data(), unrope[static_cast(abs_pos)]); + rope_apply_conj(krow.data(), unrope[rank]); + __fp16* kdst = key_rows + (rank * kv_heads + h) * head_dim; + __fp16* vdst = val_rows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + kdst[d] = static_cast<__fp16>(krow[d]); + vdst[d] = static_cast<__fp16>(vrow[d]); + } + } + } +} + +void compact_fp16(uint16_t* key_rows, uint16_t* val_rows, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, double rope_theta) { + compact_fp16(key_rows, val_rows, kv_heads, head_dim, kept_per_head, + unrope_table(max_kept_index(kept_per_head) + 1, head_dim, rope_theta)); +} + +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, + const std::vector& unrope, bool renumber) { + size_t groups = (head_dim + group_size - 1) / group_size; + size_t int8_stride = kv_heads * head_dim; + size_t scale_stride = kv_heads * groups; + std::vector row(head_dim); + for (size_t h = 0; h < kv_heads; ++h) { + const std::vector& kept = kept_per_head[h]; + for (size_t rank = 0; rank < kept.size(); ++rank) { + int abs_pos = kept[rank]; + size_t src_t = static_cast(abs_pos); + const int8_t* src = int8_rows + src_t * int8_stride + h * head_dim; + const float* ssc = scale_rows + src_t * scale_stride + h * groups; + dequant_row(src, ssc, head_dim, group_size, row.data()); + if (renumber) { // un-rope by -abs_pos then re-rope by +rank (no per-row trig) + rope_apply(row.data(), unrope[static_cast(abs_pos)]); + rope_apply_conj(row.data(), unrope[rank]); + } + int8_t* dst = int8_rows + rank * int8_stride + h * head_dim; + float* dsc = scale_rows + rank * scale_stride + h * groups; + requant_row(row.data(), dst, dsc, head_dim, group_size); + } + } +} + +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, double rope_theta, + bool renumber) { + std::vector unrope; + if (renumber) unrope = unrope_table(max_kept_index(kept_per_head) + 1, head_dim, rope_theta); + compact_int8(int8_rows, scale_rows, kv_heads, head_dim, group_size, kept_per_head, unrope, renumber); +} + +void rotate_int8_row(int8_t* int8, float* scale, size_t head_dim, size_t group_size, + double rope_theta, double delta_pos) { + rotate_int8_row_rot(int8, scale, head_dim, group_size, + rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos)); +} + +void rerope_recent_fp16(uint16_t* key_rows_u, size_t kv_heads, size_t head_dim, + size_t lo, size_t hi, double rope_theta, double delta_pos) { + if (delta_pos == 0.0 || hi <= lo) return; + __fp16* key_rows = reinterpret_cast<__fp16*>(key_rows_u); + RopeRotation rot = rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos); + std::vector row(head_dim); + for (size_t t = lo; t < hi; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + __fp16* r = key_rows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) row[d] = r[d]; + rope_apply(row.data(), rot); + for (size_t d = 0; d < head_dim; ++d) r[d] = static_cast<__fp16>(row[d]); + } +} + +void rerope_recent_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, size_t head_dim, + size_t group_size, size_t lo, size_t hi, double rope_theta, double delta_pos) { + if (delta_pos == 0.0 || hi <= lo) return; + size_t groups = (head_dim + group_size - 1) / group_size; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + RopeRotation rot = rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos); + for (size_t t = lo; t < hi; ++t) + for (size_t h = 0; h < kv_heads; ++h) + rotate_int8_row_rot(int8_rows + t * int8_stride + h * head_dim, + scale_rows + t * scale_stride + h * groups, + head_dim, group_size, rot); +} + +std::vector> keepsets_from_fp16(const uint16_t* key_rows_u, size_t n, + size_t kv_heads, size_t head_dim, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head) { + const __fp16* key_rows = reinterpret_cast(key_rows_u); + return keepsets_per_head(n, kv_heads, head_dim, unrope, p, [&](size_t h, float* post) { + for (size_t t = 0; t < n; ++t) { + const __fp16* src = key_rows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) post[t * head_dim + d] = src[d]; + } + }, protect_per_head); +} + +std::vector> keepsets_from_fp16(const uint16_t* key_rows, size_t n, + size_t kv_heads, size_t head_dim, + double rope_theta, const Params& p) { + return keepsets_from_fp16(key_rows, n, kv_heads, head_dim, + unrope_table(n, head_dim, rope_theta), p); +} + +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head) { + size_t groups = (head_dim + group_size - 1) / group_size; + size_t int8_stride = kv_heads * head_dim; + size_t scale_stride = kv_heads * groups; + return keepsets_per_head(n, kv_heads, head_dim, unrope, p, [&](size_t h, float* post) { + for (size_t t = 0; t < n; ++t) { + const int8_t* src = int8_rows + t * int8_stride + h * head_dim; + const float* ssc = scale_rows + t * scale_stride + h * groups; + dequant_row(src, ssc, head_dim, group_size, post + t * head_dim); + } + }, protect_per_head); +} + +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, double rope_theta, + const Params& p) { + return keepsets_from_int8(int8_rows, scale_rows, n, kv_heads, head_dim, group_size, + unrope_table(n, head_dim, rope_theta), p); +} + +std::vector remap_rows_through_kept(const std::vector& rows, const std::vector& kept) { + std::vector out; + size_t i = 0; + for (size_t rank = 0; rank < kept.size(); ++rank) { + while (i < rows.size() && rows[i] < kept[rank]) ++i; + if (i < rows.size() && rows[i] == kept[rank]) out.push_back(static_cast(rank)); + } + return out; +} + +void SpecialRowTracker::add_appended(size_t layer, size_t kv_heads, + const std::vector& appended_rows) { + if (layer_rows_.size() <= layer) layer_rows_.resize(layer + 1); + auto& heads = layer_rows_[layer]; + if (heads.empty()) heads.resize(kv_heads); + for (auto& rows : heads) + rows.insert(rows.end(), appended_rows.begin(), appended_rows.end()); // appended >= existing, stays sorted +} + +const std::vector>& SpecialRowTracker::protect(size_t layer) const { + static const std::vector> kEmpty; + return layer < layer_rows_.size() ? layer_rows_[layer] : kEmpty; +} + +size_t SpecialRowTracker::max_reserved(size_t layer, size_t sink, + const std::vector& appended) const { + const auto& heads = protect(layer); + if (heads.empty()) { + size_t extra = 0; + for (int r : appended) if (static_cast(r) >= sink) ++extra; + return sink + extra; + } + size_t worst = 0; + for (const auto& rows : heads) { + size_t count = sink; + for (int r : rows) if (static_cast(r) >= sink) ++count; + for (int r : appended) if (static_cast(r) >= sink) ++count; + worst = std::max(worst, count); + } + return worst; +} + +void SpecialRowTracker::remap(size_t layer, const std::vector>& kept_per_head) { + if (layer >= layer_rows_.size()) return; + auto& heads = layer_rows_[layer]; + for (size_t h = 0; h < heads.size() && h < kept_per_head.size(); ++h) + heads[h] = remap_rows_through_kept(heads[h], kept_per_head[h]); +} + +bool is_sliding_layer(const std::vector& layer_types, size_t li) { + return li < layer_types.size() && layer_types[li].find("sliding") != std::string::npos; +} + +std::vector physical_compressible_layers(const std::vector& layer_types, + size_t num_layers, size_t num_kv_shared) { + auto is_full = [&](size_t i) { return !is_sliding_layer(layer_types, i); }; + std::vector full; + for (size_t i = 0; i < num_layers; ++i) if (is_full(i)) full.push_back(i); + + if (num_kv_shared == 0 || num_kv_shared >= num_layers) return full; + size_t first_shared = num_layers - num_kv_shared; + + std::set sources; + std::vector> last_of_type; + for (size_t i = 0; i < first_shared && i < layer_types.size(); ++i) { + bool found = false; + for (auto& e : last_of_type) if (e.first == layer_types[i]) { e.second = i; found = true; break; } + if (!found) last_of_type.emplace_back(layer_types[i], i); + } + for (auto& e : last_of_type) sources.insert(e.second); + + std::vector out; + for (size_t i : full) if (i < first_shared && !sources.count(i)) out.push_back(i); + return out; +} + +} // namespace kvcompress +} // namespace cactus diff --git a/cactus-engine/src/kv_compress.h b/cactus-engine/src/kv_compress.h new file mode 100644 index 000000000..c33a5e659 --- /dev/null +++ b/cactus-engine/src/kv_compress.h @@ -0,0 +1,127 @@ +#ifndef CACTUS_KV_COMPRESS_H +#define CACTUS_KV_COMPRESS_H + +#include +#include +#include +#include +#include + +namespace cactus { +namespace kvcompress { + +struct Params { + float recent_frac = 0.30f; + size_t sink = 4; + int abs_budget = 0; // per (layer, kv-head) keep budget, clamped to [1, n] + std::vector protect; // positions always kept (special tokens) +}; + +// Mirrors cactus-graph CacheMetadata. +struct CacheHeader { + uint64_t current_seq_len; + uint64_t max_seq_len; + uint64_t num_kv_heads; + uint64_t head_dim; + uint64_t sink_size; + uint64_t reserved[3]; +}; +static_assert(sizeof(CacheHeader) == 64, "CacheHeader must be 64 bytes"); + +void kv_set_simd(bool on); + +// score s_i = -cos(k_i, mean(k)); keys pre-RoPE. +void keydiff_score(const float* keys, size_t n, size_t head_dim, float* out); + +// Keep-set: sink + recent + top-score middle; ties by ascending index. +std::vector keepset_for_head(const float* scores, size_t n, const Params& p); + +// Keys stored POST-RoPE; rotating by delta_pos re-RoPEs to orig+delta_pos. +void rope_rotate_row(float* row, size_t head_dim, double rope_theta, double delta_pos); + +void rotate_int8_row(int8_t* int8, float* scale, size_t head_dim, size_t group_size, + double rope_theta, double delta_pos); + +struct RopeRotation { std::vector cos, sin; }; + +// Un-rope rotations for [0, n): row t rotates by -t. +std::vector unrope_table(size_t n, size_t head_dim, double rope_theta); + +// Gather survivors in rank order; re-rope K to 0..B-1, gather V unchanged. +void compact_fp16(uint16_t* key_rows, uint16_t* val_rows, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, + const std::vector& unrope); +void compact_fp16(uint16_t* key_rows, uint16_t* val_rows, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, double rope_theta); + +// renumber=true (K) re-ropes each row to its new rank; false (V) gathers as-is. +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, + const std::vector& unrope, bool renumber); +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, double rope_theta, + bool renumber); + +// Rotate recent K rows [lo, hi) by delta_pos; sink and V stay fixed. +void rerope_recent_fp16(uint16_t* key_rows, size_t kv_heads, size_t head_dim, + size_t lo, size_t hi, double rope_theta, double delta_pos); +void rerope_recent_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, size_t head_dim, + size_t group_size, size_t lo, size_t hi, double rope_theta, double delta_pos); + +// Per-head keep-sets from POST-RoPE keys. protect_per_head overrides Params::protect; empty -> shared. +std::vector> keepsets_from_fp16(const uint16_t* key_rows, size_t n, + size_t kv_heads, size_t head_dim, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head = {}); +std::vector> keepsets_from_fp16(const uint16_t* key_rows, size_t n, + size_t kv_heads, size_t head_dim, + double rope_theta, const Params& p); + +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head = {}); +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, double rope_theta, + const Params& p); + +std::vector remap_rows_through_kept(const std::vector& rows, const std::vector& kept); + +// Per-(compressible layer, head) special rows; lets compaction re-protect each head's own specials. +class SpecialRowTracker { +public: + void clear() { tracked_len_ = 0; layer_rows_.clear(); valid_ = true; } + size_t tracked_len() const { return tracked_len_; } + void set_tracked_len(size_t len) { tracked_len_ = len; } + bool valid() const { return valid_; } + void invalidate() { valid_ = false; layer_rows_.clear(); } + + // appended_rows must be >= tracked_len (the still head-aligned region). + void add_appended(size_t layer, size_t kv_heads, const std::vector& appended_rows); + const std::vector>& protect(size_t layer) const; + size_t max_reserved(size_t layer, size_t sink, const std::vector& appended) const; + void remap(size_t layer, const std::vector>& kept_per_head); + +private: + size_t tracked_len_ = 0; + bool valid_ = true; + std::vector>> layer_rows_; // [layer][head] -> sorted rows +}; + +// Sliding (local) vs full-attention; selects local vs global theta. +bool is_sliding_layer(const std::vector& layer_types, size_t li); + +// Compressible layers: empty -> all (Qwen); num_kv_shared drops shared consumers+sources (Gemma -> {4,9}). +std::vector physical_compressible_layers(const std::vector& layer_types, + size_t num_layers, size_t num_kv_shared); + +} // namespace kvcompress +} // namespace cactus + +#endif // CACTUS_KV_COMPRESS_H diff --git a/cactus-engine/src/log.cpp b/cactus-engine/src/log.cpp new file mode 100644 index 000000000..0c6228022 --- /dev/null +++ b/cactus-engine/src/log.cpp @@ -0,0 +1,23 @@ +#include "../cactus_engine.h" +#include "cactus_graph.h" + +extern "C" { + +void cactus_log_set_level(int level) { + cactus::Logger::instance().set_level(static_cast(level)); +} + +void cactus_log_set_callback(cactus_log_callback_t callback, void* user_data) { + if (callback) { + cactus::Logger::instance().set_callback( + [callback, user_data](cactus::LogLevel level, + const std::string& component, + const std::string& message) { + callback(static_cast(level), component.c_str(), message.c_str(), user_data); + }); + } else { + cactus::Logger::instance().set_callback(nullptr); + } +} + +} diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp new file mode 100644 index 000000000..8cbb7c1dc --- /dev/null +++ b/cactus-engine/src/model.cpp @@ -0,0 +1,4397 @@ +#include "engine.h" +#include "cactus_graph.h" +#include "cactus_kernels.h" + +#define PICOJSON_USE_INT64 +#include "picojson.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cactus { +namespace engine { + +std::vector parse_config_uint_list(const std::string& value) { + std::vector numbers; + size_t pos = 0; + while (pos < value.size()) { + while (pos < value.size() && !std::isdigit(static_cast(value[pos]))) { + ++pos; + } + if (pos >= value.size()) break; + size_t end = pos; + while (end < value.size() && std::isdigit(static_cast(value[end]))) { + ++end; + } + numbers.push_back(static_cast(std::stoul(value.substr(pos, end - pos)))); + pos = end; + } + return numbers; +} + +float read_scalar_value(Precision precision, const uint8_t* data, size_t index) { + const uint8_t* ptr = data + PrecisionTraits::byte_offset_of(precision, index); + switch (precision) { + case Precision::FP32: + return *reinterpret_cast(ptr); + case Precision::FP16: + return static_cast(*reinterpret_cast(ptr)); + case Precision::INT8: + return static_cast(*reinterpret_cast(ptr)); + default: + return 0.0f; + } +} + +void write_scalar_value(Precision precision, uint8_t* data, size_t index, float value) { + uint8_t* ptr = data + PrecisionTraits::byte_offset_of(precision, index); + switch (precision) { + case Precision::FP32: + *reinterpret_cast(ptr) = value; + break; + case Precision::FP16: + *reinterpret_cast<__fp16*>(ptr) = static_cast<__fp16>(value); + break; + case Precision::INT8: + *reinterpret_cast(ptr) = static_cast(value); + break; + default: + break; + } +} + +bool copy_component_tensor(CactusGraph& source_graph, + const BufferDesc& src_desc, + size_t src_node, + const BufferDesc& dst_desc, + std::vector& dst_buffer, + size_t dst_element_offset, + size_t element_count, + const std::string& name) { + const auto* src_ptr = static_cast(source_graph.get_output(src_node)); + if (src_desc.precision == dst_desc.precision) { + size_t dst_offset = PrecisionTraits::byte_offset_of(dst_desc.precision, dst_element_offset); + std::memcpy( + dst_buffer.data() + dst_offset, + src_ptr, + PrecisionTraits::packed_size_of(src_desc.precision, element_count)); + return true; + } + if (name != "position_ids" && name != "attention_mask") return false; + for (size_t i = 0; i < element_count; ++i) { + write_scalar_value( + dst_desc.precision, + dst_buffer.data(), + dst_element_offset + i, + read_scalar_value(src_desc.precision, src_ptr, i)); + } + return true; +} + +struct CrossKVCacheMetadata { + uint64_t current_seq_len; + uint64_t max_seq_len; + uint64_t num_kv_heads; + uint64_t head_dim; + uint64_t sink_size; + uint64_t reserved[3]; +}; + +static_assert(sizeof(CrossKVCacheMetadata) == 64, "CrossKVCacheMetadata must be 64 bytes"); + +size_t cross_kv_cache_buffer_size(size_t max_seq, size_t kv_heads, size_t head_dim) { + size_t num_groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + return sizeof(CrossKVCacheMetadata) + max_seq * kv_heads * head_dim + + max_seq * kv_heads * num_groups * sizeof(float); +} + +bool write_cross_kv_cache_buffer(const BufferDesc& src_desc, + const void* src_ptr, + const BufferDesc& dst_desc, + std::vector& dst_buffer, + size_t source_len, + const std::string& name) { + if (src_desc.precision != Precision::FP16 || dst_desc.precision != Precision::INT8) return false; + if (src_desc.shape.size() != 4) return false; + + const size_t max_seq = src_desc.shape[1]; + const size_t kv_heads = src_desc.shape[2]; + const size_t head_dim = src_desc.shape[3]; + const size_t expected_bytes = cross_kv_cache_buffer_size(max_seq, kv_heads, head_dim); + if (dst_buffer.size() < expected_bytes || dst_desc.byte_size < expected_bytes) { + CACTUS_LOG_ERROR("model", "cross-KV cache input buffer is too small for " << name); + return false; + } + + source_len = std::min(source_len, max_seq); + std::fill(dst_buffer.begin(), dst_buffer.end(), 0); + auto* meta = reinterpret_cast(dst_buffer.data()); + meta->current_seq_len = static_cast(source_len); + meta->max_seq_len = static_cast(max_seq); + meta->num_kv_heads = static_cast(kv_heads); + meta->head_dim = static_cast(head_dim); + meta->sink_size = 0; + + if (source_len == 0) return true; + auto* int8_base = reinterpret_cast(dst_buffer.data() + sizeof(CrossKVCacheMetadata)); + const size_t int8_bytes = max_seq * kv_heads * head_dim; + auto* scale_base = reinterpret_cast(dst_buffer.data() + sizeof(CrossKVCacheMetadata) + int8_bytes); + cactus_quantize_kv_fp16_to_int8( + static_cast(src_ptr), + int8_base, + scale_base, + source_len, + kv_heads, + head_dim); + return true; +} + +void ConvCache::init(size_t layers, size_t hidden_dim, size_t window_len, Precision model_precision) { + num_layers = layers; + hidden_size = hidden_dim; + window_size = window_len; + precision = model_precision; + element_size = PrecisionTraits::size_of(precision); + + size_t state_bytes = window_size * hidden_size * element_size; + layer_states.resize(num_layers); + for (auto& state : layer_states) { + state.data.resize(state_bytes); + std::memset(state.data.data(), 0, state_bytes); + state.head = 0; + state.count = 0; + } +} + +ConvCache::CircularView ConvCache::get_window(size_t layer) const { + CircularView view{}; + if (layer >= num_layers) { + return view; + } + + const auto& state = layer_states[layer]; + if (state.count == 0) { + return view; + } + + size_t stride = hidden_size * element_size; + if (state.count < window_size) { + view.ptr1 = state.data.data(); + view.len1 = state.count; + view.total_len = state.count; + return view; + } + + view.ptr1 = state.data.data(); + view.len1 = state.head; + view.ptr2 = state.data.data() + state.head * stride; + view.len2 = window_size - state.head; + view.total_len = window_size; + return view; +} + +void ConvCache::update(CactusGraph* gb, size_t layer, const size_t bx_node) { + if (layer >= num_layers || !bx_node || window_size == 0 || hidden_size == 0) { + return; + } + + auto& state = layer_states[layer]; + const void* output_ptr = gb->get_output(bx_node); + if (!output_ptr) { + return; + } + + const auto& buffer = gb->get_output_buffer(bx_node); + const size_t stride_bytes = hidden_size * element_size; + + size_t rows = 1; + if (!buffer.shape.empty()) { + rows = buffer.shape.size() == 1 ? 1 : buffer.shape[0]; + } + + if (buffer.total_size > 0 && hidden_size > 0) { + size_t inferred = buffer.total_size / hidden_size; + if (inferred > 0) { + rows = inferred; + } + } + + if (rows == 0) { + return; + } + + size_t copy_rows = std::min(rows, window_size); + size_t start_row = rows > window_size ? rows - window_size : 0; + const auto* src = static_cast(output_ptr) + start_row * stride_bytes; + + for (size_t i = 0; i < copy_rows; ++i) { + std::memcpy(state.data.data() + state.head * stride_bytes, src + i * stride_bytes, stride_bytes); + state.head = (state.head + 1) % window_size; + if (state.count < window_size) { + ++state.count; + } + } +} + +void ConvCache::reset() { + for (auto& state : layer_states) { + std::fill(state.data.begin(), state.data.end(), 0); + state.head = 0; + state.count = 0; + } +} + + +namespace fs = std::filesystem; + +Model::Model() : config_() {} + +Model::Model(const Config& config) : config_(config) {} + +Model::~Model() = default; + +namespace { + +bool read_exact(std::ifstream& in, void* data, size_t bytes) { + in.read(static_cast(data), static_cast(bytes)); + return static_cast(in.gcount()) == bytes; +} + +bool read_float_vector(std::ifstream& in, std::vector& out, size_t count) { + out.resize(count); + return read_exact(in, out.data(), count * sizeof(float)); +} + +float relu(float x) { + return x > 0.0f ? x : 0.0f; +} + +} // namespace + +bool Model::load_handoff_probe() { + fs::path path = fs::path(bundle_dir_) / "handoff_probe.bin"; + if (!fs::exists(path)) return false; + + std::ifstream in(path, std::ios::binary); + if (!in.is_open()) return false; + + char magic[8] = {}; + uint32_t version = 0; + if (!read_exact(in, magic, sizeof(magic)) || std::string(magic, sizeof(magic)) != std::string("CHP10P6\0", 8)) { + CACTUS_LOG_WARN("cloud_handoff", "Ignoring invalid handoff probe header at " << path); + return false; + } + if (!read_exact(in, &version, sizeof(version)) + || !read_exact(in, &handoff_probe_feat_dim_, sizeof(handoff_probe_feat_dim_)) + || !read_exact(in, &handoff_probe_t_h_, sizeof(handoff_probe_t_h_)) + || !read_exact(in, &handoff_probe_h1_, sizeof(handoff_probe_h1_)) + || !read_exact(in, &handoff_probe_h2_, sizeof(handoff_probe_h2_))) { + CACTUS_LOG_WARN("cloud_handoff", "Ignoring truncated handoff probe header at " << path); + return false; + } + if (version != 1 || handoff_probe_feat_dim_ == 0 || handoff_probe_t_h_ == 0 + || handoff_probe_h1_ == 0 || handoff_probe_h2_ == 0) { + CACTUS_LOG_WARN("cloud_handoff", "Ignoring unsupported handoff probe metadata at " << path); + return false; + } + + const size_t feat = handoff_probe_feat_dim_; + const size_t th = handoff_probe_t_h_; + const size_t h1 = handoff_probe_h1_; + const size_t h2 = handoff_probe_h2_; + bool ok = true; + ok = ok && read_float_vector(in, handoff_probe_norm_weight_, feat); + ok = ok && read_float_vector(in, handoff_probe_norm_bias_, feat); + ok = ok && read_float_vector(in, handoff_probe_proj_weight_, th * feat); + ok = ok && read_float_vector(in, handoff_probe_proj_bias_, th); + ok = ok && read_float_vector(in, handoff_probe_attn_query_, th); + ok = ok && read_float_vector(in, handoff_probe_head0_weight_, h1 * th); + ok = ok && read_float_vector(in, handoff_probe_head0_bias_, h1); + ok = ok && read_float_vector(in, handoff_probe_head2_weight_, h2 * h1); + ok = ok && read_float_vector(in, handoff_probe_head2_bias_, h2); + ok = ok && read_float_vector(in, handoff_probe_head4_weight_, h2); + ok = ok && read_float_vector(in, handoff_probe_head4_bias_, 1); + if (!ok) { + CACTUS_LOG_WARN("cloud_handoff", "Ignoring truncated handoff probe weights at " << path); + return false; + } + + handoff_probe_hidden_.clear(); + handoff_probe_loaded_ = true; + CACTUS_LOG_INFO("cloud_handoff", "Loaded handoff probe from " << path); + return true; +} + +bool Model::has_handoff_probe_rollout() const { + return handoff_probe_loaded_ + && handoff_probe_feat_dim_ > 0 + && handoff_probe_hidden_.size() >= static_cast(handoff_probe_feat_dim_); +} + +void Model::maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name) { + if (!handoff_probe_loaded_ || handoff_probe_feat_dim_ == 0) return; + int idx = output_index(comp, output_name); + if (idx < 0 || static_cast(idx) >= comp.output_node_ids.size()) return; + size_t node = static_cast(comp.output_node_ids[idx]); + const auto& desc = comp.graph->get_output_buffer(node); + if (desc.total_size < handoff_probe_feat_dim_) return; + if (!desc.shape.empty() && + static_cast(desc.shape.back()) != static_cast(handoff_probe_feat_dim_)) return; + size_t rows = desc.total_size / handoff_probe_feat_dim_; + if (rows == 0) return; + const auto* data = static_cast(comp.graph->get_output(node)); + if (!data) return; + for (size_t row = 0; row < rows; ++row) { + size_t base = row * static_cast(handoff_probe_feat_dim_); + for (size_t i = 0; i < handoff_probe_feat_dim_; ++i) { + handoff_probe_hidden_.push_back(read_scalar_value(desc.precision, data, base + i)); + } + } +} + +float Model::handoff_probe_wrong_probability() const { + if (!has_handoff_probe_rollout()) return std::numeric_limits::quiet_NaN(); + + const size_t feat = handoff_probe_feat_dim_; + const size_t th = handoff_probe_t_h_; + const size_t h1 = handoff_probe_h1_; + const size_t h2 = handoff_probe_h2_; + const size_t tokens = std::min(handoff_probe_hidden_.size() / feat, 1024); + if (tokens == 0) return std::numeric_limits::quiet_NaN(); + + std::vector projected(tokens * th); + std::vector scores(tokens); + for (size_t t = 0; t < tokens; ++t) { + const float* x = handoff_probe_hidden_.data() + t * feat; + double mean = 0.0; + for (size_t i = 0; i < feat; ++i) mean += x[i]; + mean /= static_cast(feat); + double var = 0.0; + for (size_t i = 0; i < feat; ++i) { + double d = static_cast(x[i]) - mean; + var += d * d; + } + var /= static_cast(feat); + float inv_std = static_cast(1.0 / std::sqrt(var + 1e-5)); + + for (size_t j = 0; j < th; ++j) { + double acc = handoff_probe_proj_bias_[j]; + const float* w = handoff_probe_proj_weight_.data() + j * feat; + for (size_t i = 0; i < feat; ++i) { + float xn = (x[i] - static_cast(mean)) * inv_std; + xn = xn * handoff_probe_norm_weight_[i] + handoff_probe_norm_bias_[i]; + acc += static_cast(w[i]) * xn; + } + float u = relu(static_cast(acc)); + projected[t * th + j] = u; + scores[t] += u * handoff_probe_attn_query_[j]; + } + scores[t] /= std::sqrt(static_cast(th)); + } + + float max_score = *std::max_element(scores.begin(), scores.end()); + double denom = 0.0; + for (float s : scores) denom += std::exp(static_cast(s - max_score)); + std::vector pooled(th, 0.0f); + for (size_t t = 0; t < tokens; ++t) { + float alpha = static_cast(std::exp(static_cast(scores[t] - max_score)) / denom); + const float* u = projected.data() + t * th; + for (size_t j = 0; j < th; ++j) pooled[j] += alpha * u[j]; + } + + std::vector y1(h1); + for (size_t i = 0; i < h1; ++i) { + double acc = handoff_probe_head0_bias_[i]; + const float* w = handoff_probe_head0_weight_.data() + i * th; + for (size_t j = 0; j < th; ++j) acc += static_cast(w[j]) * pooled[j]; + y1[i] = relu(static_cast(acc)); + } + std::vector y2(h2); + for (size_t i = 0; i < h2; ++i) { + double acc = handoff_probe_head2_bias_[i]; + const float* w = handoff_probe_head2_weight_.data() + i * h1; + for (size_t j = 0; j < h1; ++j) acc += static_cast(w[j]) * y1[j]; + y2[i] = relu(static_cast(acc)); + } + double logit = handoff_probe_head4_bias_[0]; + for (size_t j = 0; j < h2; ++j) logit += static_cast(handoff_probe_head4_weight_[j]) * y2[j]; + return static_cast(1.0 / (1.0 + std::exp(-logit))); +} + +bool Model::init(const std::string& bundle_dir, size_t context_size, + const std::string& /*system_prompt*/, bool /*do_warmup*/) { + if (initialized_) return true; + bundle_dir_ = bundle_dir; + + if (!config_.from_json(bundle_dir + "/config.txt")) { + CACTUS_LOG_ERROR("model", "Failed to load config.txt from: " << bundle_dir); + return false; + } + apply_kv_compress_env_override(); + if (!load_manifest()) { + CACTUS_LOG_ERROR("model", "Failed to load bundle manifest from: " << bundle_dir); + return false; + } + if (!setup_tokenizer()) { + CACTUS_LOG_ERROR("model", "Tokenizer init failed for bundle: " << bundle_dir); + return false; + } + const bool is_text_embedding = + components_.count("text_embedding") + && !components_.count("decoder") + && !components_.count("decoder_step") + && !components_.count("lm_encoder_step"); + if (is_text_embedding) { + cache_max_seq_len_ = context_size; + initialized_ = true; + return true; + } + std::string encoder_name; + std::string decoder_name; + std::string source_encoder_name; + std::string decoder_cross_kv_name; + std::unordered_set required_components; + for (const auto& [name, comp] : components_) { + auto route_it = comp.metadata.find("runtime_route"); + if (route_it == comp.metadata.end() || route_it->second != "encoder_cross_kv_decoder_step") { + continue; + } + auto role_it = comp.metadata.find("runtime_role"); + if (role_it == comp.metadata.end()) { + continue; + } + if (role_it->second == "source_encoder") { + source_encoder_name = name; + auto source_kind_it = comp.metadata.find("source_kind"); + if (source_kind_it != comp.metadata.end()) { + encoder_cross_kv_source_kind_ = source_kind_it->second; + } + } else if (role_it->second == "decoder_cross_kv") { + decoder_cross_kv_name = name; + } else if (role_it->second == "decoder_step") { + decoder_name = name; + } + } + const bool has_metadata_encoder_cross_kv_route = + !source_encoder_name.empty() && + !decoder_cross_kv_name.empty() && + !decoder_name.empty(); + bool has_chunked_prefill = components_.count("lm_encoder_step") + && components_.count("decoder_media_step") + && components_.count("lm_encoder_text_chunk") + && components_.count("decoder_prefill_chunk"); + if (has_metadata_encoder_cross_kv_route) { + decode_route_ = DecodeRoute::ENCODER_CROSS_KV_STEP; + required_components = {source_encoder_name, decoder_cross_kv_name, decoder_name}; + } else if (has_chunked_prefill) { + encoder_name = "lm_encoder_step"; + decoder_name = "decoder_media_step"; + decode_route_ = DecodeRoute::CACHED_STEP; + required_components = { + encoder_name, + decoder_name, + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + }; + } else if (components_.count("decoder_step") + && input_index(components_.at("decoder_step"), "input_ids") >= 0 + && input_index(components_.at("decoder_step"), "position_ids") >= 0) { + decoder_name = "decoder_step"; + decode_route_ = DecodeRoute::DIRECT_DECODER_STEP; + required_components = {decoder_name}; + } else if (components_.count("lm_encoder_step") && components_.count("decoder_step")) { + encoder_name = "lm_encoder_step"; + decoder_name = "decoder_step"; + decode_route_ = DecodeRoute::CACHED_STEP; + required_components = {encoder_name, decoder_name}; + if (components_.count("decoder_prefill_chunk")) { + required_components.insert("decoder_prefill_chunk"); + } + if (components_.count("lm_encoder_text_chunk")) { + required_components.insert("lm_encoder_text_chunk"); + } + } else if (components_.count("text_lm_encoder") && components_.count("decoder")) { + encoder_name = "text_lm_encoder"; + decoder_name = "decoder"; + decode_route_ = DecodeRoute::FULL_CONTEXT_TEXT; + required_components = {encoder_name, decoder_name}; + } else if (components_.count("audio_encoder") && + (components_.count("decoder") || components_.count("decoder_joint"))) { + decoder_name = components_.count("decoder") ? "decoder" : "decoder_joint"; + decode_route_ = DecodeRoute::DIRECT_DECODER_STEP; + required_components = {"audio_encoder", decoder_name}; + } else { + CACTUS_LOG_ERROR("model", "Bundle missing required components: need lm_encoder_step+decoder_step (LM), text_lm_encoder+decoder, audio_encoder+decoder (transcription), or source/audio_encoder+decoder_cross_kv+decoder_step"); + return false; + } + for (const auto& optional : { + "vision_encoder", + "audio_encoder", + "lm_encoder_media_step", + "decoder_prefill_chunk", + "decoder_embed_chunk", + "lm_encoder", + }) { + if (components_.count(optional)) { + required_components.insert(optional); + } + } + if (!load_components(required_components)) return false; + if (!encoder_name.empty()) encoder_ = &components_.at(encoder_name); + if (!decoder_name.empty()) decoder_ = &components_.at(decoder_name); + if (!source_encoder_name.empty()) source_encoder_ = &components_.at(source_encoder_name); + if (!decoder_cross_kv_name.empty()) decoder_cross_kv_ = &components_.at(decoder_cross_kv_name); + if (components_.count("decoder_prefill_chunk") && components_.at("decoder_prefill_chunk").graph) { + decoder_prefill_ = &components_.at("decoder_prefill_chunk"); + decoder_prefill_chunk_ = decoder_prefill_; + } else if (decode_route_ != DecodeRoute::ENCODER_CROSS_KV_STEP && !components_.count("audio_encoder")) { + CACTUS_LOG_WARN("model", "Bundle has no decoder_prefill_chunk component; prompts will prefill token-by-token (prefill speed ~= decode speed)."); + } + if (decoder_ && decoder_->graph) decoder_->graph->prewarm_metal_quant_weights(); + if (decoder_prefill_ && decoder_prefill_->graph) decoder_prefill_->graph->prewarm_metal_quant_weights(); + if (components_.count("decoder_embed_chunk") && components_.at("decoder_embed_chunk").graph) { + decoder_embed_ = &components_.at("decoder_embed_chunk"); + } + if (components_.count("lm_encoder_text_chunk") && components_.at("lm_encoder_text_chunk").graph) { + prefill_encoder_ = &components_.at("lm_encoder_text_chunk"); + } + if (const char* env = std::getenv("CACTUS_DISABLE_PREFILL_TAIL_PAD")) { + prefill_tail_pad_disabled_ = std::atoi(env) != 0; + } + vision_encoder_ = components_.count("vision_encoder") ? &components_.at("vision_encoder") : nullptr; + vision_projector_ = components_.count("vision_projector") ? &components_.at("vision_projector") : nullptr; + audio_encoder_ = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; + lm_encoder_media_step_ = components_.count("lm_encoder_media_step") ? &components_.at("lm_encoder_media_step") : nullptr; + lm_encoder_ = components_.count("lm_encoder") ? &components_.at("lm_encoder") : nullptr; + lm_encoder_text_chunk_ = components_.count("lm_encoder_text_chunk") ? &components_.at("lm_encoder_text_chunk") : nullptr; + lm_encoder_media_chunk_ = components_.count("lm_encoder_media_chunk") ? &components_.at("lm_encoder_media_chunk") : nullptr; + std::vector to_bind = { + encoder_, + source_encoder_, + prefill_encoder_, + decoder_, + decoder_cross_kv_, + decoder_prefill_, + vision_encoder_, + vision_projector_, + audio_encoder_, + lm_encoder_media_step_, + lm_encoder_, + lm_encoder_text_chunk_, + lm_encoder_media_chunk_, + }; + std::unordered_set bound; + for (Component* comp : to_bind) { + if (!comp || !comp->graph || bound.count(comp)) continue; + if (!bind_runtime_buffers(*comp)) return false; + bound.insert(comp); + } + + if (decoder_prefill_ && decoder_prefill_->graph) { + try { + if (prefill_encoder_ && prefill_encoder_->graph) { + for (auto& b : prefill_encoder_->input_buffers) std::fill(b.begin(), b.end(), 0); + prefill_encoder_->graph->execute(); + reset_component_cache_states(*prefill_encoder_); + } + for (auto& b : decoder_prefill_->input_buffers) std::fill(b.begin(), b.end(), 0); + decoder_prefill_->graph->execute(); + reset_component_cache_states(*decoder_prefill_); + } catch (const std::exception& e) { + CACTUS_LOG_WARN("model", std::string("prefill warmup skipped: ") + e.what()); + } + } + + if (vision_encoder_ && tokenizer_ && !vision_encoder_->output_node_ids.empty()) { + size_t out_node = static_cast(vision_encoder_->output_node_ids[0]); + const auto& desc = vision_encoder_->graph->get_output_buffer(out_node); + size_t n = 0; + if (desc.shape.size() >= 3) n = desc.shape[desc.shape.size() - 2]; + else if (desc.shape.size() >= 2) n = desc.shape[0]; + if (n > 0) tokenizer_->set_image_soft_token_count(n); + } + + if (family_ == "lfm2_vl" && tokenizer_) { + tokenizer_->set_lfm2_vision_config(config_); + } + + cache_max_seq_len_ = context_size; + + if (load_handoff_probe()) { + const bool is_parakeet = config_.model_type == Config::ModelType::PARAKEET_TDT; + const Component* probe_comp = is_parakeet ? audio_encoder_ : decoder_; + const char* probe_output = is_parakeet ? "encoder_hidden_states" : "probe_hidden"; + if (!probe_comp || output_index(*probe_comp, probe_output) < 0) { + CACTUS_LOG_WARN("cloud_handoff", "Handoff probe is packaged, but its probe input is not exposed; " + "reconvert to enable probe-based handoff"); + handoff_probe_loaded_ = false; + } + } + + initialized_ = true; + return true; +} + +bool Model::load_manifest() { + std::ifstream in(fs::path(bundle_dir_) / "components" / "manifest.json"); + if (!in.is_open()) return false; + picojson::value root; + std::string err = picojson::parse(root, in); + if (!err.empty() || !root.is()) { + CACTUS_LOG_ERROR("model", "manifest parse: " << err); + return false; + } + const auto& obj = root.get(); + if (obj.count("family") && obj.at("family").is()) { + family_ = obj.at("family").get(); + } + if (!obj.count("components")) return false; + for (const auto& cv : obj.at("components").get()) { + const auto& c = cv.get(); + Component comp; + comp.name = c.at("component").get(); + comp.graph_path = c.count("graph") ? c.at("graph").get() : ""; + if (c.count("runtime_input_node_ids")) { + for (const auto& v : c.at("runtime_input_node_ids").get()) + comp.runtime_input_node_ids.push_back(static_cast(v.get())); + } + if (c.count("logical_inputs")) { + for (const auto& v : c.at("logical_inputs").get()) + comp.logical_inputs.push_back(v.get()); + } + if (c.count("output_node_ids")) { + for (const auto& v : c.at("output_node_ids").get()) + comp.output_node_ids.push_back(static_cast(v.get())); + } + if (c.count("logical_outputs")) { + for (const auto& v : c.at("logical_outputs").get()) + comp.logical_outputs.push_back(v.get()); + } + if (c.count("metadata") && c.at("metadata").is()) { + for (const auto& mv : c.at("metadata").get()) { + if (mv.second.is()) { + comp.metadata[mv.first] = mv.second.get(); + } + } + } + if (c.count("bound_constant_bindings")) { + for (const auto& bv : c.at("bound_constant_bindings").get()) { + const auto& b = bv.get(); + Binding bd; + bd.node_id = static_cast(b.at("node_id").get()); + bd.path = b.at("path").get(); + comp.bindings.push_back(std::move(bd)); + } + } + if (c.count("cache_state_node_ids")) { + for (const auto& sv : c.at("cache_state_node_ids").get()) { + if (!sv.is()) continue; + const auto& s = sv.get(); + CacheStateBinding cs; + if (s.count("layer_key")) cs.layer_key = s.at("layer_key").get(); + if (s.count("key") && s.at("key").is()) + cs.key_node_id = static_cast(s.at("key").get()); + if (s.count("value") && s.at("value").is()) + cs.value_node_id = static_cast(s.at("value").get()); + if (cs.key_node_id >= 0 && cs.value_node_id >= 0) { + comp.cache_states.push_back(std::move(cs)); + } + } + } + components_[comp.name] = std::move(comp); + } + return true; +} + +bool Model::setup_tokenizer() { + std::string vocab = bundle_dir_ + "/vocab.txt"; + std::string merges = bundle_dir_ + "/merges.txt"; + std::string cfg = bundle_dir_ + "/tokenizer_config.txt"; + if (!fs::exists(vocab)) return false; + auto rt = load_tokenizer_runtime_config(cfg); + bool use_bpe = rt.tokenizer_type == TokenizerRuntimeConfig::TokenizerType::BPE + || (rt.tokenizer_type == TokenizerRuntimeConfig::TokenizerType::UNKNOWN + && fs::exists(merges)); + if (use_bpe) tokenizer_ = std::make_unique(); + else tokenizer_ = std::make_unique(); + return tokenizer_->load_vocabulary_with_config(vocab, merges, cfg); +} + +bool Model::load_components(const std::unordered_set& required_components) { + for (auto& [name, comp] : components_) { + if (!required_components.empty() && !required_components.count(name)) continue; + if (!load_component_graph(comp)) return false; + } + return true; +} + +bool Model::load_component_graph(Component& comp) { + if (comp.graph) return true; + if (comp.graph_path.empty()) return true; + fs::path full = fs::path(bundle_dir_) / comp.graph_path; + try { + comp.graph = std::make_unique(CactusGraph::load(full.string())); + comp.graph->retain_outputs(comp.output_node_ids); + } catch (const std::exception& e) { + CACTUS_LOG_ERROR("model", "load " << comp.graph_path << ": " << e.what()); + return false; + } + for (const auto& b : comp.bindings) { + if (b.node_id < 0) continue; + try { + fs::path weight_path(b.path); + if (weight_path.is_absolute()) { + fs::path local = fs::path(bundle_dir_) / weight_path.filename(); + if (fs::exists(local)) weight_path = local; + } else { + weight_path = fs::path(bundle_dir_) / weight_path; + } + comp.graph->bind_mmap_weights(static_cast(b.node_id), weight_path.string()); + } catch (const std::exception& e) { + CACTUS_LOG_ERROR("model", "bind " << b.path << ": " << e.what()); + return false; + } + } + return bind_runtime_buffers(comp); +} + +void Model::unload_component_graph(Component& comp) { + if (cactus_default_backend() == ComputeBackend::METAL) return; + if (comp.graph) { + comp.graph->release_runtime_buffers(); + comp.graph->release_all_weight_pages(); + } + comp.input_buffers.clear(); + comp.graph.reset(); +} + +bool Model::bind_runtime_buffers(Component& comp) { + comp.input_buffers.resize(comp.runtime_input_node_ids.size()); + for (size_t i = 0; i < comp.runtime_input_node_ids.size(); ++i) { + size_t node_id = static_cast(comp.runtime_input_node_ids[i]); + const auto& desc = comp.graph->get_output_buffer(node_id); + comp.input_buffers[i].assign(desc.byte_size, 0); + comp.graph->set_external_input(node_id, comp.input_buffers[i].data(), desc.precision); + } + return true; +} + +int Model::input_index(const Component& comp, const std::string& name) const { + for (size_t i = 0; i < comp.logical_inputs.size(); ++i) { + if (comp.logical_inputs[i] == name) return static_cast(i); + } + return -1; +} + +void Model::write_int_input(Component& comp, const std::string& name, int64_t value) { + write_int_input_at(comp, name, 0, value); +} + +void Model::write_int_input_at(Component& comp, const std::string& name, size_t index, int64_t value) { + int idx = input_index(comp, name); + if (idx < 0) return; + size_t node_id = static_cast(comp.runtime_input_node_ids[idx]); + const auto& desc = comp.graph->get_output_buffer(node_id); + auto& buf = comp.input_buffers[idx]; + if (index >= desc.total_size) return; + size_t offset = PrecisionTraits::byte_offset_of(desc.precision, index); + auto* dst = buf.data() + offset; + switch (desc.precision) { + case Precision::FP32: + *reinterpret_cast(dst) = static_cast(value); + break; + case Precision::FP16: + *reinterpret_cast<__fp16*>(dst) = static_cast<__fp16>(value); + break; + case Precision::INT8: + *reinterpret_cast(dst) = static_cast(value); + break; + default: + *reinterpret_cast(dst) = static_cast(value); + break; + } +} + +void Model::write_bytes_input(Component& comp, const std::string& name, const void* data, size_t byte_size) { + int idx = input_index(comp, name); + if (idx < 0) return; + auto& buf = comp.input_buffers[idx]; + size_t to_copy = std::min(byte_size, buf.size()); + std::memcpy(buf.data(), data, to_copy); + if (to_copy < buf.size()) { + std::memset(buf.data() + to_copy, 0, buf.size() - to_copy); + } +} + +int Model::output_index(const Component& comp, const std::string& name) const { + for (size_t i = 0; i < comp.logical_outputs.size(); ++i) { + if (comp.logical_outputs[i] == name) return static_cast(i); + } + return -1; +} + +void Model::copy_encoder_outputs_to_decoder(const Component& enc) { + for (size_t i = 0; i < enc.output_node_ids.size() && i < enc.logical_outputs.size(); ++i) { + const std::string& out_name = enc.logical_outputs[i]; + int dst_idx = input_index(*decoder_, out_name); + if (dst_idx < 0) continue; + size_t src_node = static_cast(enc.output_node_ids[i]); + const auto& src_desc = enc.graph->get_output_buffer(src_node); + void* src_ptr = enc.graph->get_output(src_node); + auto& dst_buf = decoder_->input_buffers[dst_idx]; + size_t to_copy = std::min(src_desc.byte_size, dst_buf.size()); + std::memcpy(dst_buf.data(), src_ptr, to_copy); + } +} + +void Model::run_step(uint32_t token_id, size_t position, bool /*read_logits*/, bool use_fused) { + if (decode_route_ == DecodeRoute::DIRECT_DECODER_STEP) { + write_int_input(*decoder_, "input_ids", static_cast(token_id)); + write_int_input(*decoder_, "position_ids", static_cast(position)); + if (!use_fused || cactus_default_backend() != ComputeBackend::METAL) cactus_graph_set_prefill_consistent(true); + decoder_->graph->execute(); + cactus_graph_set_prefill_consistent(false); + maybe_capture_handoff_probe_hidden(*decoder_); + return; + } + if (!use_fused) cactus_graph_set_prefill_consistent(true); + if (use_fused && cactus_default_backend() == ComputeBackend::METAL) { + if (ple_probe_state_ == 0) + ple_probe_state_ = encoder_->graph->extract_ple_pathway(fused_embed_ctx_) ? 1 : 2; + if (ple_probe_state_ == 1) { + fused_embed_ctx_.token_id = static_cast(token_id); + fused_embed_ctx_.position = static_cast(position); + cactus_graph_set_fused_embed(&fused_embed_ctx_); + write_int_input(*decoder_, "position_ids", static_cast(position)); + decoder_->graph->execute(); + cactus_graph_set_fused_embed(nullptr); + cactus_graph_set_prefill_consistent(false); + maybe_capture_handoff_probe_hidden(*decoder_); + return; + } + } + run_encoder_step(token_id, position); + copy_component_outputs_to_inputs(*encoder_, *decoder_); + decoder_->graph->execute(); + cactus_graph_set_prefill_consistent(false); + maybe_capture_handoff_probe_hidden(*decoder_); +} + +void Model::run_encoder_step(uint32_t token_id, size_t position) { + write_int_input(*encoder_, "input_ids", static_cast(token_id)); + write_int_input(*encoder_, "position_ids", static_cast(position)); + encoder_->graph->execute(); +} + +void Model::set_component_batch(Component& comp, size_t batch) { + for (size_t i = 0; i < comp.runtime_input_node_ids.size(); ++i) { + size_t node_id = static_cast(comp.runtime_input_node_ids[i]); + const auto& desc = comp.graph->get_output_buffer(node_id); + if (!desc.has_dynamic_dims() || desc.shape.empty() || desc.shape[0] == batch) continue; + std::vector shape = desc.shape; + shape[0] = batch; + comp.graph->set_runtime_input_shape(node_id, shape); + const auto& resized = comp.graph->get_output_buffer(node_id); + comp.input_buffers[i].assign(resized.byte_size, 0); + comp.graph->set_external_input(node_id, comp.input_buffers[i].data(), resized.precision); + } +} + +size_t Model::decoder_cache_num_slots() { + if (!decoder_ || !decoder_->graph) return 1; + for (const auto& state : decoder_->cache_states) { + int node_id = state.key_node_id; + if (node_id < 0) continue; + if (decoder_->graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; + size_t num_slots = decoder_->graph->get_node_cache_num_slots(static_cast(node_id)); + return num_slots ? num_slots : 1; + } + return 1; +} + +void Model::run_step_batch(const std::vector& token_ids, const std::vector& positions) { + if (decode_route_ == DecodeRoute::DIRECT_DECODER_STEP) { + for (size_t b = 0; b < token_ids.size(); ++b) { + write_int_input_at(*decoder_, "input_ids", b, static_cast(token_ids[b])); + write_int_input_at(*decoder_, "position_ids", b, static_cast(positions[b])); + } + decoder_->graph->execute(); + maybe_capture_handoff_probe_hidden(*decoder_); + return; + } + for (size_t b = 0; b < token_ids.size(); ++b) { + write_int_input_at(*encoder_, "input_ids", b, static_cast(token_ids[b])); + write_int_input_at(*encoder_, "position_ids", b, static_cast(positions[b])); + } + encoder_->graph->execute(); + copy_component_outputs_to_inputs(*encoder_, *decoder_); + decoder_->graph->execute(); + maybe_capture_handoff_probe_hidden(*decoder_); +} + +std::vector Model::batch_stop_token_ids() const { + std::vector stops; + stops.push_back(config_.eos_token_id); + Tokenizer* tk = get_tokenizer(); + auto add_if_single = [&](const std::string& s) { + if (!tk || s.empty()) return; + std::vector t = tk->encode(s); + if (t.size() == 1) stops.push_back(t[0]); + }; + if (tk) add_if_single(tk->get_default_stop_sequence()); + if (config_.model_type == Config::ModelType::GEMMA4) add_if_single(""); + return stops; +} + +std::vector> Model::generate_batch(const std::vector>& prompts, + size_t max_new_tokens, bool stop_on_eos) { + size_t batch = prompts.size(); + const bool cached = decode_route_ == DecodeRoute::CACHED_STEP; + const bool direct = decode_route_ == DecodeRoute::DIRECT_DECODER_STEP; + if (batch == 0 || !decoder_ || (!cached && !direct)) return {}; + if (cached && !encoder_) return {}; + for (const auto& p : prompts) if (p.empty()) return {}; + if (cached && !load_component_graph(*encoder_)) return {}; + if (!load_component_graph(*decoder_)) return {}; + if (batch > decoder_cache_num_slots()) return {}; + auto has_dynamic_input = [](Component& c) { + for (int node_id : c.runtime_input_node_ids) { + if (node_id < 0) continue; + if (c.graph->get_output_buffer(static_cast(node_id)).has_dynamic_dims()) return true; + } + return false; + }; + if (batch > 1) { + if (!has_dynamic_input(*decoder_)) return {}; + if (cached && !has_dynamic_input(*encoder_)) return {}; + for (const auto& np : decoder_->graph->nodes_) { + if (np->op_type == OpType::CONV_CACHE_STATE || np->op_type == OpType::RECURRENT_CACHE_STATE) return {}; + } + } + + size_t exec_batch = batch; + if (batch > 1 || cactus_default_backend() == ComputeBackend::METAL) { + exec_batch = std::max(batch, decoder_cache_num_slots()); + } + reset_component_cache_states(*decoder_); + if (cached) set_component_batch(*encoder_, exec_batch); + set_component_batch(*decoder_, exec_batch); + + std::vector> out(batch); + std::vector fed(batch, 0); + std::vector last(batch, 0); + std::vector done(batch, false); + size_t remaining = batch; + + const std::vector stop_ids = stop_on_eos ? batch_stop_token_ids() : std::vector{}; + auto is_stop = [&](uint32_t t) { + for (uint32_t s : stop_ids) if (s == t) return true; + return false; + }; + + std::vector tokens(exec_batch); + std::vector positions(exec_batch); + while (remaining > 0) { + for (size_t b = 0; b < batch; ++b) { + positions[b] = fed[b]; + tokens[b] = (fed[b] < prompts[b].size()) ? prompts[b][fed[b]] : last[b]; + } + for (size_t b = batch; b < exec_batch; ++b) { + positions[b] = positions[0]; + tokens[b] = tokens[0]; + } + run_step_batch(tokens, positions); + std::vector sampled = argmax_component_logits_batch(*decoder_, batch); + for (size_t b = 0; b < batch; ++b) { + ++fed[b]; + if (done[b]) continue; + if (fed[b] >= prompts[b].size()) { + last[b] = sampled[b]; + if (is_stop(sampled[b])) { + done[b] = true; + --remaining; + continue; + } + out[b].push_back(sampled[b]); + if (out[b].size() >= max_new_tokens) { + done[b] = true; + --remaining; + } + } + } + } + return out; +} + +std::vector> Model::decode_batch(const std::vector& seed_tokens, + size_t max_new_tokens) { + std::vector> prompts; + prompts.reserve(seed_tokens.size()); + for (uint32_t seed : seed_tokens) prompts.push_back({seed}); + return generate_batch(prompts, max_new_tokens); +} + +bool Model::supports_dynamic_batch() { + if (!decoder_) return false; + if (!decoder_->graph && !load_component_graph(*decoder_)) return false; + return decoder_->graph->has_dynamic_shapes(); +} + +void Model::set_decode_slots(size_t num_slots) { + if (num_slots == 0) num_slots = 1; + if (!decoder_) return; + if (!decoder_->graph && !load_component_graph(*decoder_)) return; + for (const auto& state : decoder_->cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + if (decoder_->graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; + decoder_->graph->resize_cache_slots(static_cast(node_id), num_slots); + } + } +} + +#define FOR_EACH_MATCHED_OUTPUT(source, target, body) \ + for (size_t _i = 0; _i < (source).output_node_ids.size() && _i < (source).logical_outputs.size(); ++_i) { \ + const std::string& out_name = (source).logical_outputs[_i]; \ + int dst_idx = input_index((target), out_name); \ + if (dst_idx < 0) continue; \ + size_t src_node = static_cast((source).output_node_ids[_i]); \ + const auto& src_desc = (source).graph->get_output_buffer(src_node); \ + size_t dst_node = static_cast((target).runtime_input_node_ids[dst_idx]); \ + const auto& dst_desc = (target).graph->get_output_buffer(dst_node); \ + auto& dst_buf = (target).input_buffers[dst_idx]; \ + body \ + } + +void Model::copy_component_outputs_to_inputs(const Component& source, Component& target) { + FOR_EACH_MATCHED_OUTPUT(source, target, { + std::fill(dst_buf.begin(), dst_buf.end(), 0); + size_t elements = std::min(src_desc.total_size, dst_desc.total_size); + if (!copy_component_tensor(*source.graph, src_desc, src_node, dst_desc, dst_buf, 0, elements, out_name)) + throw std::runtime_error("component output/input precision mismatch for " + out_name); + }) +} + +bool Model::copy_cross_kv_outputs_to_decoder_cache_inputs(const Component& source, Component& target, size_t source_len) { + bool copied_any = false; + for (size_t i = 0; i < source.output_node_ids.size() && i < source.logical_outputs.size(); ++i) { + const std::string& out_name = source.logical_outputs[i]; + if (out_name.rfind("cross_k_", 0) != 0 && out_name.rfind("cross_v_", 0) != 0) { + continue; + } + int dst_idx = input_index(target, out_name); + if (dst_idx < 0) continue; + size_t src_node = static_cast(source.output_node_ids[i]); + const auto& src_desc = source.graph->get_output_buffer(src_node); + const void* src_ptr = source.graph->get_output(src_node); + size_t dst_node = static_cast(target.runtime_input_node_ids[dst_idx]); + const auto& dst_desc = target.graph->get_output_buffer(dst_node); + if (!write_cross_kv_cache_buffer(src_desc, src_ptr, dst_desc, target.input_buffers[dst_idx], source_len, out_name)) { + return false; + } + copied_any = true; + } + return copied_any; +} + +void Model::copy_component_outputs_to_chunk_inputs(const Component& source, Component& target, size_t token_index) { + FOR_EACH_MATCHED_OUTPUT(source, target, { + size_t chunk_tokens = component_chunk_tokens(target, out_name); + if (chunk_tokens <= token_index || chunk_tokens == 0) + throw std::runtime_error("chunk prefill token index exceeds input capacity for " + out_name); + if (dst_desc.total_size % chunk_tokens != 0) + throw std::runtime_error("chunk prefill input shape is not token-aligned for " + out_name); + size_t elements_per_token = dst_desc.total_size / chunk_tokens; + if (src_desc.total_size != elements_per_token) + throw std::runtime_error("component output/input token shape mismatch for " + out_name); + if (!copy_component_tensor(*source.graph, src_desc, src_node, dst_desc, + dst_buf, token_index * elements_per_token, src_desc.total_size, out_name)) + throw std::runtime_error("component output/input precision mismatch for " + out_name); + }) +} + +void Model::copy_component_outputs_to_chunk_inputs_range(const Component& source, Component& target, size_t token_offset) { + FOR_EACH_MATCHED_OUTPUT(source, target, { + size_t src_tokens = component_output_tokens(source, out_name); + size_t dst_tokens = component_chunk_tokens(target, out_name); + if (src_tokens == 0 || dst_tokens == 0 || token_offset + src_tokens > dst_tokens) + throw std::runtime_error("chunk prefill output range exceeds input capacity for " + out_name); + if (src_desc.total_size % src_tokens != 0 || dst_desc.total_size % dst_tokens != 0) + throw std::runtime_error("chunk prefill output/input shape is not token-aligned for " + out_name); + size_t src_elements_per_token = src_desc.total_size / src_tokens; + size_t dst_elements_per_token = dst_desc.total_size / dst_tokens; + if (src_elements_per_token != dst_elements_per_token) + throw std::runtime_error("component output/input token shape mismatch for " + out_name); + if (!copy_component_tensor(*source.graph, src_desc, src_node, dst_desc, + dst_buf, token_offset * dst_elements_per_token, src_desc.total_size, out_name)) + throw std::runtime_error("component output/input precision mismatch for " + out_name); + }) +} + +#undef FOR_EACH_MATCHED_OUTPUT + +bool Model::cache_states_compatible(const Component& source, const Component& target) const { + if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) return false; + for (size_t i = 0; i < source.cache_states.size(); ++i) { + const auto& src = source.cache_states[i]; + const auto& dst = target.cache_states[i]; + if (src.layer_key != dst.layer_key) return false; + if (src.key_node_id < 0 || src.value_node_id < 0 || dst.key_node_id < 0 || dst.value_node_id < 0) return false; + } + return true; +} + +void Model::move_cache_states(Component& source, Component& target, size_t logical_current) { + if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) { + throw std::runtime_error("prefill and step cache states are not compatible"); + } + for (size_t i = 0; i < source.cache_states.size(); ++i) { + const auto& src = source.cache_states[i]; + const auto& dst = target.cache_states[i]; + if (src.layer_key != dst.layer_key) { + throw std::runtime_error("prefill and step cache layer mismatch: " + src.layer_key + " != " + dst.layer_key); + } + int moved_src = -1, moved_dst = -1; + for (auto [src_node, dst_node] : {std::pair{src.key_node_id, dst.key_node_id}, std::pair{src.value_node_id, dst.value_node_id}}) { + if (src_node < 0 || dst_node < 0) continue; + // Conv/recurrent caches share one node for key and value; move it only once. + if (src_node == moved_src && dst_node == moved_dst) continue; + moved_src = src_node; + moved_dst = dst_node; + target.graph->steal_cache_buffer(static_cast(dst_node), *source.graph, static_cast(src_node)); + if (target.graph->get_node_op_type(static_cast(dst_node)) == OpType::KV_CACHE_STATE && + logical_current != std::numeric_limits::max()) { + auto* meta = static_cast(target.graph->get_output(static_cast(dst_node))); + if (meta && logical_current < meta[0]) { + meta[0] = logical_current; + } + } + } + } +} + +void Model::set_cache_current_len(Component& comp, size_t len) { + for (const auto& state : comp.cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; + auto* meta = static_cast(comp.graph->get_output(static_cast(node_id))); + if (!meta || len >= meta[0]) continue; + meta[0] = len; + } + } +} + +void Model::reset_component_cache_states(Component& comp) { + for (const auto& state : comp.cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + const auto& desc = comp.graph->get_output_buffer(static_cast(node_id)); + if (desc.byte_size == 0 || !desc.get_data()) continue; + void* ptr = comp.graph->get_output(static_cast(node_id)); + if (!ptr) continue; + const OpType op_type = comp.graph->get_node_op_type(static_cast(node_id)); + switch (op_type) { + case OpType::KV_CACHE_STATE: + if (desc.byte_size >= sizeof(uint64_t)) { + auto* meta = static_cast(ptr); + uint64_t num_slots = (desc.byte_size >= 6 * sizeof(uint64_t) && meta[5] > 0) ? meta[5] : 1; + size_t slot_stride = num_slots ? desc.byte_size / num_slots : desc.byte_size; + for (uint64_t s = 0; s < num_slots; ++s) { + *reinterpret_cast(static_cast(ptr) + s * slot_stride) = 0; + } + } + break; + case OpType::CONV_CACHE_STATE: + if (desc.byte_size >= 2 * sizeof(uint64_t)) { + auto* meta = static_cast(ptr); + meta[0] = 0; // head + meta[1] = 0; // count + } + break; + case OpType::RECURRENT_CACHE_STATE: + std::memset(ptr, 0, desc.byte_size); + break; + default: + break; + } + } + } +} + +size_t Model::component_chunk_tokens(const Component& comp, const std::string& input_name) const { + int idx = input_index(comp, input_name); + if (idx < 0) return 0; + const auto& desc = comp.graph->get_output_buffer(static_cast(comp.runtime_input_node_ids[idx])); + if (desc.shape.size() >= 2 && desc.shape[0] == 1) return desc.shape[1]; + return desc.shape.empty() ? 0 : desc.shape[0]; +} + +size_t Model::component_output_tokens(const Component& comp, const std::string& output_name) const { + for (size_t i = 0; i < comp.logical_outputs.size() && i < comp.output_node_ids.size(); ++i) { + if (comp.logical_outputs[i] != output_name) continue; + const auto& desc = comp.graph->get_output_buffer(static_cast(comp.output_node_ids[i])); + if (desc.shape.size() >= 2 && desc.shape[0] == 1) return desc.shape[1]; + return desc.shape.empty() ? 0 : desc.shape[0]; + } + return 0; +} + +void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, size_t encoder_chunk, + size_t chunk_tokens, const std::vector& tokens, + size_t processed, size_t start_position) { + for (size_t i = 0; i < chunk_comp.input_buffers.size(); ++i) { + std::fill(chunk_comp.input_buffers[i].begin(), chunk_comp.input_buffers[i].end(), 0); + } + if (enc_comp && encoder_chunk > 0) { + for (size_t chunk_offset = 0; chunk_offset < chunk_tokens; chunk_offset += encoder_chunk) { + for (size_t i = 0; i < enc_comp->input_buffers.size(); ++i) { + std::fill(enc_comp->input_buffers[i].begin(), enc_comp->input_buffers[i].end(), 0); + } + for (size_t i = 0; i < encoder_chunk; ++i) { + size_t index = processed + chunk_offset + i; + uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); + write_int_input_at(*enc_comp, "input_ids", i, static_cast(token)); + write_int_input_at(*enc_comp, "position_ids", i, static_cast(start_position + processed + chunk_offset + i)); + } + enc_comp->graph->execute(); + copy_component_outputs_to_chunk_inputs_range(*enc_comp, chunk_comp, chunk_offset); + } + } else { + for (size_t i = 0; i < chunk_tokens; ++i) { + size_t index = processed + i; + uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); + run_encoder_step(token, start_position + processed + i); + copy_component_outputs_to_chunk_inputs(*encoder_, chunk_comp, i); + } + } + chunk_comp.graph->execute(); +} + +void Model::reset_prefill_stats() { + last_prefill_cache_copy_ms_ = 0.0; + last_prefill_padding_tokens_ = 0; + last_prefill_scalar_tail_tokens_ = 0; + last_prefill_tail_chunk_tokens_ = 0; + last_prefill_tail_padding_tokens_ = 0; +} + +Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector& tokens, size_t start_position, size_t chunk_size, bool prepare_decode) { + ChunkedPrefillResult result; + reset_prefill_stats(); + if (decode_route_ != DecodeRoute::CACHED_STEP || !encoder_ || !decoder_ || !decoder_prefill_) return result; + if (start_position != 0) return result; + if (!load_component_graph(*decoder_prefill_)) return result; + if (prefill_encoder_ && !load_component_graph(*prefill_encoder_)) return result; + if (!cache_states_compatible(*decoder_prefill_, *decoder_)) return result; + size_t component_tokens = component_chunk_tokens(*decoder_prefill_, "inputs_embeds"); + if (component_tokens <= 1) return result; + size_t effective_chunk = chunk_size > 0 ? std::min(chunk_size, component_tokens) : component_tokens; + if (effective_chunk != component_tokens) effective_chunk = component_tokens; + size_t whole_chunks_end = (tokens.size() / effective_chunk) * effective_chunk; + auto any_cache_node = [&](auto predicate) { + if (!decoder_prefill_->graph) return false; + for (const auto& state : decoder_prefill_->cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + if (predicate(static_cast(node_id))) return true; + } + } + return false; + }; + const bool has_recurrent_state = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::RECURRENT_CACHE_STATE; + }); + if (has_recurrent_state && whole_chunks_end > effective_chunk) { + whole_chunks_end = effective_chunk; + } + const bool has_sliding_window_cache = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::KV_CACHE_STATE + && decoder_prefill_->graph->get_node_window_size(id) > 0; + }); + const size_t tail_tokens = tokens.size() - whole_chunks_end; + const size_t padding_cutoff = std::max(1, effective_chunk / 16); + const bool has_conv_state = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::CONV_CACHE_STATE; + }); + const bool pad_tail = !has_conv_state + && !has_recurrent_state + && !has_sliding_window_cache + && tail_tokens >= padding_cutoff; + const bool padded_window_too_small = any_cache_node([&](size_t id) { + if (decoder_prefill_->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) return false; + size_t window = decoder_prefill_->graph->get_node_window_size(id); + size_t sink = decoder_prefill_->graph->get_node_sink_size(id); + return window > 0 && (window < effective_chunk * 4 || window <= effective_chunk + sink); + }); + const bool use_padded_tail = !pad_tail && !prefill_tail_pad_disabled_ + && has_sliding_window_cache && !has_recurrent_state && !has_conv_state + && tail_tokens > 8 && !padded_window_too_small; + const size_t executable_tokens = whole_chunks_end + (pad_tail ? effective_chunk : 0); + if (executable_tokens == 0 && !use_padded_tail) { + result.scalar_tail_tokens = tail_tokens; + last_prefill_scalar_tail_tokens_ = tail_tokens; + return result; + } + + size_t encoder_chunk = 0; + if (prefill_encoder_ && input_index(*prefill_encoder_, "input_ids") >= 0 && input_index(*prefill_encoder_, "position_ids") >= 0) { + encoder_chunk = component_chunk_tokens(*prefill_encoder_, "input_ids"); + if (encoder_chunk == 0 || effective_chunk % encoder_chunk != 0) { + encoder_chunk = 0; + } + } + + size_t processed = 0; + while (processed + effective_chunk <= executable_tokens) { + execute_prefill_chunk(*decoder_prefill_, prefill_encoder_, encoder_chunk, + effective_chunk, tokens, processed, start_position); + processed += effective_chunk; + } + + size_t tail_executed = 0; + size_t tail_padding = 0; + if (use_padded_tail) { + const size_t pads = effective_chunk - tail_tokens; + const size_t kept_real = tail_tokens - 1; + std::vector>> backups; + for (const auto& state : decoder_prefill_->cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + size_t id = static_cast(node_id); + if (decoder_prefill_->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) continue; + backups.emplace_back(id, decoder_prefill_->graph->snapshot_cache_padded_append(id, kept_real, pads + 1)); + } + } + execute_prefill_chunk(*decoder_prefill_, prefill_encoder_, encoder_chunk, + effective_chunk, tokens, processed, start_position); + for (auto& [id, backup] : backups) { + decoder_prefill_->graph->rollback_cache_padded_append(id, kept_real, pads + 1, backup); + } + processed += kept_real; + tail_executed = kept_real; + tail_padding = pads; + } + + result.executed_tokens = processed; + result.logical_tokens = std::min(tokens.size(), processed); + if (result.logical_tokens > 0) { + result.last_logit_row = (result.logical_tokens - 1) % effective_chunk; + } + result.padding_tokens = processed > tokens.size() ? processed - tokens.size() : 0; + result.scalar_tail_tokens = tokens.size() - result.logical_tokens; + last_prefill_padding_tokens_ = result.padding_tokens; + last_prefill_scalar_tail_tokens_ = result.scalar_tail_tokens; + last_prefill_tail_chunk_tokens_ = tail_executed; + last_prefill_tail_padding_tokens_ = tail_padding; + + if (processed > 0 && prepare_decode) { + for (size_t i = 0; i < decoder_->input_buffers.size(); ++i) { + std::fill(decoder_->input_buffers[i].begin(), decoder_->input_buffers[i].end(), 0); + } + auto copy_start = std::chrono::high_resolution_clock::now(); + move_cache_states(*decoder_prefill_, *decoder_, start_position + result.logical_tokens); + auto copy_end = std::chrono::high_resolution_clock::now(); + last_prefill_cache_copy_ms_ = std::chrono::duration_cast(copy_end - copy_start).count() / 1000.0; + } + return result; +} + +void Model::run_full_context_text() { + if (!encoder_ || !decoder_ || context_tokens_.empty()) return; + int input_ids_idx = input_index(*encoder_, "input_ids"); + int attention_mask_idx = input_index(*encoder_, "attention_mask"); + if (input_ids_idx < 0 || attention_mask_idx < 0) { + throw std::runtime_error("text_lm_encoder requires input_ids and attention_mask inputs"); + } + size_t input_node = static_cast(encoder_->runtime_input_node_ids[input_ids_idx]); + const auto& input_desc = encoder_->graph->get_output_buffer(input_node); + if (context_tokens_.size() > input_desc.total_size) { + throw std::runtime_error("context exceeds transpiled text_lm_encoder capacity"); + } + std::fill(encoder_->input_buffers[input_ids_idx].begin(), encoder_->input_buffers[input_ids_idx].end(), 0); + std::fill(encoder_->input_buffers[attention_mask_idx].begin(), encoder_->input_buffers[attention_mask_idx].end(), 0); + for (size_t i = 0; i < context_tokens_.size(); ++i) { + write_int_input_at(*encoder_, "input_ids", i, static_cast(context_tokens_[i])); + write_int_input_at(*encoder_, "attention_mask", i, 1); + } + encoder_->graph->execute(); + for (size_t i = 0; i < encoder_->output_node_ids.size() && i < encoder_->logical_outputs.size(); ++i) { + const std::string& out_name = encoder_->logical_outputs[i]; + int dst_idx = input_index(*decoder_, out_name); + if (dst_idx < 0) continue; + size_t src_node = static_cast(encoder_->output_node_ids[i]); + const auto& src_desc = encoder_->graph->get_output_buffer(src_node); + void* src_ptr = encoder_->graph->get_output(src_node); + std::memcpy(decoder_->input_buffers[dst_idx].data(), src_ptr, src_desc.byte_size); + } + last_logit_position_ = context_tokens_.empty() ? 0 : context_tokens_.size() - 1; + decoder_->graph->execute(); +} + +void Model::write_media_embeds_row(Component& comp, int embeds_idx, const uint8_t* feature_row, + size_t feature_row_bytes, Precision feature_precision) { + auto& embeds_buf = comp.input_buffers[embeds_idx]; + size_t node_id = static_cast(comp.runtime_input_node_ids[embeds_idx]); + const auto& desc = comp.graph->get_output_buffer(node_id); + if (desc.precision == feature_precision) { + size_t to_copy = std::min(feature_row_bytes, embeds_buf.size()); + std::memcpy(embeds_buf.data(), feature_row, to_copy); + if (to_copy < embeds_buf.size()) { + std::memset(embeds_buf.data() + to_copy, 0, embeds_buf.size() - to_copy); + } + return; + } + size_t src_elem = PrecisionTraits::size_of(feature_precision); + size_t dst_elem = PrecisionTraits::size_of(desc.precision); + size_t src_count = src_elem ? feature_row_bytes / src_elem : 0; + size_t dst_count = dst_elem ? embeds_buf.size() / dst_elem : 0; + size_t n = std::min(src_count, dst_count); + auto load_float = [&](size_t i) -> float { + if (feature_precision == Precision::FP16) return static_cast(reinterpret_cast(feature_row)[i]); + if (feature_precision == Precision::FP32) return reinterpret_cast(feature_row)[i]; + return static_cast(reinterpret_cast(feature_row)[i]); + }; + for (size_t i = 0; i < n; ++i) { + float v = load_float(i); + if (desc.precision == Precision::FP16) reinterpret_cast<__fp16*>(embeds_buf.data())[i] = static_cast<__fp16>(v); + else if (desc.precision == Precision::FP32) reinterpret_cast(embeds_buf.data())[i] = v; + else reinterpret_cast(embeds_buf.data())[i] = static_cast(v); + } + if (n < dst_count) { + std::memset(embeds_buf.data() + n * dst_elem, 0, (dst_count - n) * dst_elem); + } +} + +void Model::run_media_step(size_t position, const uint8_t* feature_row, size_t feature_row_bytes, + Precision feature_precision) { + if (lm_encoder_media_step_) { + int embeds_idx = input_index(*lm_encoder_media_step_, "inputs_embeds"); + if (embeds_idx >= 0) { + write_media_embeds_row(*lm_encoder_media_step_, embeds_idx, feature_row, feature_row_bytes, feature_precision); + write_int_input(*lm_encoder_media_step_, "input_ids", 0); + write_int_input(*lm_encoder_media_step_, "position_ids", static_cast(position)); + lm_encoder_media_step_->graph->execute(); + copy_encoder_outputs_to_decoder(*lm_encoder_media_step_); + decoder_->graph->execute(); + return; + } + } else if (encoder_ != nullptr && decoder_ != nullptr) { + int dec_embeds_idx = input_index(*decoder_, "inputs_embeds"); + if (dec_embeds_idx >= 0) { + run_encoder_step(static_cast(config_.pad_token_id), position); + copy_component_outputs_to_inputs(*encoder_, *decoder_); + write_media_embeds_row(*decoder_, dec_embeds_idx, feature_row, feature_row_bytes, feature_precision); + decoder_->graph->execute(); + return; + } + } + run_step(static_cast(config_.pad_token_id), position, false); +} + +namespace { +void write_typed_buffer(std::vector& buf, Precision dst_prec, + const void* src_data, size_t src_bytes, Precision src_prec); +} // namespace + +void Model::run_vision_encoder(const std::string& image_path) { + if (!vision_encoder_) return; + if (family_ == "lfm2_vl") { + run_vision_encoder_lfm2_vl(image_path); + return; + } + if (!load_component_graph(*vision_encoder_)) { + throw std::runtime_error("failed to load vision_encoder"); + } + + auto write_int_buffer_typed = [&](int idx, const int64_t* src, size_t src_count) { + auto& buf = vision_encoder_->input_buffers[idx]; + size_t node = static_cast(vision_encoder_->runtime_input_node_ids[idx]); + const auto& desc = vision_encoder_->graph->get_output_buffer(node); + const size_t elem = PrecisionTraits::size_of(desc.precision); + const size_t cap = elem ? buf.size() / elem : 0; + const size_t n = std::min(cap, src_count); + for (size_t i = 0; i < n; ++i) { + int64_t v = src[i]; + switch (desc.precision) { + case Precision::FP32: reinterpret_cast(buf.data())[i] = static_cast(v); break; + case Precision::FP16: reinterpret_cast<__fp16*>(buf.data())[i] = static_cast<__fp16>(v); break; + case Precision::INT8: reinterpret_cast(buf.data())[i] = static_cast(v); break; + default: + if (elem == 8) reinterpret_cast(buf.data())[i] = v; + else if (elem == 4) reinterpret_cast(buf.data())[i] = static_cast(v); + break; + } + } + if (n < cap) std::memset(buf.data() + n * elem, 0, (cap - n) * elem); + }; + + if (family_ == "qwen3_5" || family_ == "qwen3_vl" || config_.model_type == Config::ModelType::QWEN) { + Qwen3VlImagePreprocessed prep = preprocess_qwen3_vl_image(image_path, config_); + int pv_idx = input_index(*vision_encoder_, "pixel_values"); + if (pv_idx < 0) { + throw std::runtime_error("Qwen3-VL vision_encoder missing pixel_values input"); + } + auto& pv_buf = vision_encoder_->input_buffers[pv_idx]; + size_t pv_node = static_cast(vision_encoder_->runtime_input_node_ids[pv_idx]); + const auto& pv_desc = vision_encoder_->graph->get_output_buffer(pv_node); + write_typed_buffer(pv_buf, pv_desc.precision, + prep.pixel_values.data(), + prep.pixel_values.size() * sizeof(float), + Precision::FP32); + } else { + Gemma4ImagePreprocessed prep = preprocess_gemma4_image(image_path, config_); + int pv_idx = input_index(*vision_encoder_, "pixel_values"); + if (pv_idx >= 0) { + auto& pv_buf = vision_encoder_->input_buffers[pv_idx]; + size_t pv_node = static_cast(vision_encoder_->runtime_input_node_ids[pv_idx]); + const auto& pv_desc = vision_encoder_->graph->get_output_buffer(pv_node); + write_typed_buffer(pv_buf, pv_desc.precision, + prep.pixel_values.data(), + prep.pixel_values.size() * sizeof(float), + Precision::FP32); + } + int pp_idx = input_index(*vision_encoder_, "pixel_position_ids"); + if (pp_idx >= 0) { + write_int_buffer_typed(pp_idx, prep.pixel_position_ids.data(), + prep.pixel_position_ids.size()); + } + } + + vision_encoder_->graph->execute(); + for (size_t i = 0; i < vision_encoder_->output_node_ids.size() && i < vision_encoder_->logical_outputs.size(); ++i) { + const std::string& name = vision_encoder_->logical_outputs[i]; + size_t node_id = static_cast(vision_encoder_->output_node_ids[i]); + const auto& desc = vision_encoder_->graph->get_output_buffer(node_id); + void* ptr = vision_encoder_->graph->get_output(node_id); + auto& slot = media_features_[name]; + slot.assign(desc.byte_size, 0); + std::memcpy(slot.data(), ptr, desc.byte_size); + media_feature_shapes_[name] = desc.shape; + media_feature_precisions_[name] = desc.precision; + } + if (cactus_default_backend() != ComputeBackend::METAL) { + vision_encoder_->graph->release_runtime_buffers(); + vision_encoder_->graph->release_all_weight_pages(); + } + unload_component_graph(*vision_encoder_); +} + +namespace { +void typed_buffer_to_float(const void* src, Precision prec, float* dst, size_t n) { + switch (prec) { + case Precision::FP32: std::memcpy(dst, src, n * sizeof(float)); break; + case Precision::FP16: { + const __fp16* s = reinterpret_cast(src); + for (size_t i = 0; i < n; ++i) dst[i] = static_cast(s[i]); + break; + } + case Precision::INT8: { + const int8_t* s = reinterpret_cast(src); + for (size_t i = 0; i < n; ++i) dst[i] = static_cast(s[i]); + break; + } + default: std::memset(dst, 0, n * sizeof(float)); break; + } +} +} // namespace + +bool Model::load_lfm2_vl_position_grid() { + if (lfm2_pos_grid_loaded_) return !lfm2_pos_grid_.empty(); + lfm2_pos_grid_loaded_ = true; + if (!vision_encoder_) return false; + auto path_it = vision_encoder_->metadata.find("position_embedding_grid_path"); + auto shape_it = vision_encoder_->metadata.find("position_embedding_grid_shape"); + if (path_it == vision_encoder_->metadata.end() || shape_it == vision_encoder_->metadata.end()) return false; + int gh = 0, gw = 0, gd = 0; + if (std::sscanf(shape_it->second.c_str(), "%d,%d,%d", &gh, &gw, &gd) != 3) return false; + if (gh <= 0 || gw <= 0 || gd <= 0) return false; + fs::path full = fs::path(bundle_dir_) / path_it->second; + std::ifstream f(full, std::ios::binary); + if (!f.is_open()) return false; + const size_t count = static_cast(gh) * gw * gd; + lfm2_pos_grid_.resize(count); + f.read(reinterpret_cast(lfm2_pos_grid_.data()), static_cast(count * sizeof(float))); + if (!f) { lfm2_pos_grid_.clear(); return false; } + lfm2_pos_grid_h_ = gh; + lfm2_pos_grid_w_ = gw; + lfm2_pos_grid_dim_ = gd; + return true; +} + +void Model::run_vision_encoder_lfm2_vl(const std::string& image_path) { + if (!vision_encoder_ || !vision_projector_) { + throw std::runtime_error("lfm2_vl requires vision_encoder and vision_projector components"); + } + if (!load_lfm2_vl_position_grid()) { + throw std::runtime_error("lfm2_vl vision position-embedding grid is missing from the bundle"); + } + if (!load_component_graph(*vision_encoder_)) throw std::runtime_error("failed to load vision_encoder"); + if (!load_component_graph(*vision_projector_)) throw std::runtime_error("failed to load vision_projector"); + + media_features_.erase("image_features"); + media_feature_shapes_.erase("image_features"); + media_feature_precisions_.erase("image_features"); + + encode_lfm2_vl_image_into_features(image_path); + + unload_component_graph(*vision_encoder_); + unload_component_graph(*vision_projector_); +} + +void Model::encode_lfm2_vl_image_into_features(const std::string& image_path) { + Lfm2VlImagePreprocessed prep = preprocess_lfm2_vl_image(image_path, config_); + const int dim = lfm2_pos_grid_dim_; + const int factor = config_.downsample_factor ? static_cast(config_.downsample_factor) : 2; + const size_t max_patches = prep.max_num_patches; + const size_t patch_dim = prep.patch_dim; + const int cff = dim * factor * factor; + + const int pv_idx = input_index(*vision_encoder_, "pixel_values"); + const int pm_idx = input_index(*vision_encoder_, "pixel_attention_mask"); + const int pe_idx = input_index(*vision_encoder_, "positional_embeddings"); + const int vf_idx = input_index(*vision_projector_, "vision_features"); + if (pv_idx < 0 || pe_idx < 0 || vf_idx < 0) { + throw std::runtime_error("lfm2_vl vision components are missing expected inputs"); + } + + const size_t enc_out_node = static_cast(vision_encoder_->output_node_ids[0]); + const size_t proj_out_node = static_cast(vision_projector_->output_node_ids[0]); + const auto& proj_out_desc = vision_projector_->graph->get_output_buffer(proj_out_node); + const size_t proj_rows = proj_out_desc.shape.size() >= 2 + ? proj_out_desc.shape[proj_out_desc.shape.size() - 2] : 0; + const size_t text_hidden = proj_out_desc.shape.empty() ? 0 : proj_out_desc.shape.back(); + const Precision proj_prec = proj_out_desc.precision; + const size_t proj_elem = PrecisionTraits::size_of(proj_prec); + if (proj_rows == 0 || text_hidden == 0) { + throw std::runtime_error("lfm2_vl vision_projector has an invalid output shape"); + } + const auto& enc_out_desc0 = vision_encoder_->graph->get_output_buffer(enc_out_node); + const size_t enc_elem = PrecisionTraits::size_of(enc_out_desc0.precision); + const size_t enc_out_elems = enc_elem ? enc_out_desc0.byte_size / enc_elem : 0; + if (enc_out_elems < max_patches * static_cast(dim)) { + throw std::runtime_error("lfm2_vl vision_encoder output is smaller than the preprocessed patch grid"); + } + + auto write_mask = [&](int comp_idx, const int64_t* src, size_t count) { + auto& buf = vision_encoder_->input_buffers[comp_idx]; + size_t node = static_cast(vision_encoder_->runtime_input_node_ids[comp_idx]); + const auto& d = vision_encoder_->graph->get_output_buffer(node); + const size_t elem = PrecisionTraits::size_of(d.precision); + const size_t cap = elem ? buf.size() / elem : 0; + const size_t n = std::min(cap, count); + for (size_t i = 0; i < n; ++i) { + int64_t v = src[i]; + switch (d.precision) { + case Precision::FP32: reinterpret_cast(buf.data())[i] = static_cast(v); break; + case Precision::FP16: reinterpret_cast<__fp16*>(buf.data())[i] = static_cast<__fp16>(v); break; + case Precision::INT8: reinterpret_cast(buf.data())[i] = static_cast(v); break; + default: + if (elem == 8) reinterpret_cast(buf.data())[i] = v; + else if (elem == 4) reinterpret_cast(buf.data())[i] = static_cast(v); + break; + } + } + if (n < cap) std::memset(buf.data() + n * elem, 0, (cap - n) * elem); + }; + + auto write_float_input = [](Component& comp, int idx, const float* src, size_t count) { + auto& buf = comp.input_buffers[idx]; + size_t node = static_cast(comp.runtime_input_node_ids[idx]); + const auto& d = comp.graph->get_output_buffer(node); + write_typed_buffer(buf, d.precision, src, count * sizeof(float), Precision::FP32); + }; + + std::vector image_features; + size_t total_tokens = 0; + + std::vector pos_buf(max_patches * static_cast(dim)); + std::vector enc_out_f(max_patches * static_cast(dim)); + std::vector unshuf; + std::vector proj_in; + + for (size_t t = 0; t < prep.spatial_shapes.size(); ++t) { + const int h = prep.spatial_shapes[t].first; + const int w = prep.spatial_shapes[t].second; + const int num_tokens = (h / factor) * (w / factor); + if (static_cast(num_tokens) > proj_rows) { + throw std::runtime_error("lfm2_vl sub-image exceeds the traced projector capacity"); + } + + std::fill(pos_buf.begin(), pos_buf.end(), 0.0f); + interpolate_position_embeddings( + lfm2_pos_grid_.data(), lfm2_pos_grid_h_, lfm2_pos_grid_w_, dim, h, w, pos_buf.data()); + + write_float_input(*vision_encoder_, pv_idx, + prep.pixel_values.data() + t * max_patches * patch_dim, max_patches * patch_dim); + if (pm_idx >= 0) { + write_mask(pm_idx, prep.pixel_attention_mask.data() + t * max_patches, max_patches); + } + write_float_input(*vision_encoder_, pe_idx, pos_buf.data(), max_patches * static_cast(dim)); + vision_encoder_->graph->execute(); + const auto& enc_desc = vision_encoder_->graph->get_output_buffer(enc_out_node); + const void* enc_ptr = vision_encoder_->graph->get_output(enc_out_node); + typed_buffer_to_float(enc_ptr, enc_desc.precision, enc_out_f.data(), + max_patches * static_cast(dim)); + + unshuf.assign(static_cast(num_tokens) * cff, 0.0f); + pixel_unshuffle(enc_out_f.data(), h, w, dim, factor, unshuf.data()); + + proj_in.assign(proj_rows * static_cast(cff), 0.0f); + std::copy(unshuf.begin(), unshuf.end(), proj_in.begin()); + write_float_input(*vision_projector_, vf_idx, proj_in.data(), proj_in.size()); + + vision_projector_->graph->execute(); + const void* proj_ptr = vision_projector_->graph->get_output(proj_out_node); + const size_t row_bytes = text_hidden * proj_elem; + const size_t append_bytes = static_cast(num_tokens) * row_bytes; + const size_t base = image_features.size(); + image_features.resize(base + append_bytes); + std::memcpy(image_features.data() + base, proj_ptr, append_bytes); + total_tokens += static_cast(num_tokens); + } + + auto& slot = media_features_["image_features"]; + const size_t prev_bytes = slot.size(); + slot.resize(prev_bytes + image_features.size()); + std::memcpy(slot.data() + prev_bytes, image_features.data(), image_features.size()); + auto& shape = media_feature_shapes_["image_features"]; + if (shape.size() != 2) shape = { total_tokens, text_hidden }; + else shape[0] += total_tokens; + media_feature_precisions_["image_features"] = proj_prec; +} + +void Model::run_audio_encoder_messages(const std::vector>& audio_features_per_message) { + if (!audio_encoder_) return; + if (audio_features_per_message.empty()) return; + if (!load_component_graph(*audio_encoder_)) { + throw std::runtime_error("failed to load audio_encoder"); + } + for (const std::string& logical : audio_encoder_->logical_outputs) { + media_features_.erase(logical); + media_feature_shapes_.erase(logical); + media_feature_precisions_.erase(logical); + } + for (const auto& mel : audio_features_per_message) { + if (mel.empty()) continue; + run_audio_encoder(mel); + } + if (cactus_default_backend() != ComputeBackend::METAL) { + audio_encoder_->graph->release_runtime_buffers(); + audio_encoder_->graph->release_all_weight_pages(); + } + unload_component_graph(*audio_encoder_); +} + +void Model::run_audio_encoder(const std::vector& audio_features) { + if (!audio_encoder_) return; + const std::vector candidate_input_names = {"input_features", "audio_features"}; + int feature_idx = -1; + for (const auto& name : candidate_input_names) { + int idx = input_index(*audio_encoder_, name); + if (idx >= 0) { feature_idx = idx; break; } + } + if (feature_idx < 0) { + CACTUS_LOG_WARN("model", "audio_encoder has no input named input_features/audio_features; skipping"); + return; + } + const size_t feature_node = static_cast(audio_encoder_->runtime_input_node_ids[feature_idx]); + const auto& feature_desc = audio_encoder_->graph->get_output_buffer(feature_node); + const size_t mel_bins = feature_desc.shape.size() >= 3 + ? feature_desc.shape[2] + : static_cast(config_.audio_input_feat_size); + const size_t max_frames_per_chunk = feature_desc.shape.size() >= 2 ? feature_desc.shape[1] : 0; + if (max_frames_per_chunk == 0 || mel_bins == 0) { + CACTUS_LOG_WARN("model", "audio_encoder feature input has unexpected shape; skipping"); + return; + } + const size_t total_frames = audio_features.size() / mel_bins; + const size_t num_chunks = total_frames == 0 + ? 1 + : (total_frames + max_frames_per_chunk - 1) / max_frames_per_chunk; + + const int mask_idx = input_index(*audio_encoder_, "input_features_mask"); + + for (size_t chunk_idx = 0; chunk_idx < num_chunks; ++chunk_idx) { + const size_t frame_start = chunk_idx * max_frames_per_chunk; + const size_t frames_in_chunk = frame_start >= total_frames + ? 0 + : std::min(max_frames_per_chunk, total_frames - frame_start); + const size_t chunk_feature_elems = frames_in_chunk * mel_bins; + const float* chunk_src = audio_features.data() + frame_start * mel_bins; + + auto& buf = audio_encoder_->input_buffers[feature_idx]; + if (feature_desc.precision == Precision::FP32) { + const size_t n_bytes = std::min(chunk_feature_elems * sizeof(float), buf.size()); + std::memcpy(buf.data(), chunk_src, n_bytes); + if (n_bytes < buf.size()) std::memset(buf.data() + n_bytes, 0, buf.size() - n_bytes); + } else if (feature_desc.precision == Precision::FP16) { + const size_t cap_elems = buf.size() / sizeof(__fp16); + const size_t n_elems = std::min(chunk_feature_elems, cap_elems); + __fp16* dst = reinterpret_cast<__fp16*>(buf.data()); + for (size_t i = 0; i < n_elems; ++i) dst[i] = static_cast<__fp16>(chunk_src[i]); + if (n_elems < cap_elems) { + std::memset(buf.data() + n_elems * sizeof(__fp16), 0, (cap_elems - n_elems) * sizeof(__fp16)); + } + } else { + const size_t n_elems = std::min(chunk_feature_elems, buf.size()); + int8_t* dst = reinterpret_cast(buf.data()); + for (size_t i = 0; i < n_elems; ++i) dst[i] = static_cast(chunk_src[i]); + if (n_elems < buf.size()) std::memset(buf.data() + n_elems, 0, buf.size() - n_elems); + } + + if (mask_idx >= 0) { + auto& mb = audio_encoder_->input_buffers[mask_idx]; + const size_t mask_node = static_cast(audio_encoder_->runtime_input_node_ids[mask_idx]); + const auto& mask_desc = audio_encoder_->graph->get_output_buffer(mask_node); + const size_t elem = PrecisionTraits::size_of(mask_desc.precision); + const size_t cap = elem ? mb.size() / elem : 0; + const size_t n = std::min(cap, frames_in_chunk); + for (size_t i = 0; i < n; ++i) { + switch (mask_desc.precision) { + case Precision::FP32: reinterpret_cast(mb.data())[i] = 1.0f; break; + case Precision::FP16: reinterpret_cast<__fp16*>(mb.data())[i] = static_cast<__fp16>(1.0f); break; + case Precision::INT8: reinterpret_cast(mb.data())[i] = 1; break; + default: reinterpret_cast(mb.data())[i] = 1; break; + } + } + if (n < cap) std::memset(mb.data() + n * elem, 0, (cap - n) * elem); + } + + audio_encoder_->graph->execute(); + + for (size_t i = 0; i < audio_encoder_->output_node_ids.size() && i < audio_encoder_->logical_outputs.size(); ++i) { + const std::string& name = audio_encoder_->logical_outputs[i]; + const size_t node_id = static_cast(audio_encoder_->output_node_ids[i]); + const auto& desc = audio_encoder_->graph->get_output_buffer(node_id); + void* ptr = audio_encoder_->graph->get_output(node_id); + auto& slot = media_features_[name]; + const size_t prev_bytes = slot.size(); + slot.resize(prev_bytes + desc.byte_size); + std::memcpy(slot.data() + prev_bytes, ptr, desc.byte_size); + auto shape_it = media_feature_shapes_.find(name); + if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { + media_feature_shapes_[name] = desc.shape; + } else if (desc.shape.size() >= 2 && shape_it->second.size() == desc.shape.size()) { + shape_it->second[shape_it->second.size() - 2] += desc.shape[desc.shape.size() - 2]; + } + media_feature_precisions_[name] = desc.precision; + } + } +} + +static float uncertainty_from_margin(float best, float second) { + float confidence = 1.0f; + if (std::isfinite(best) && std::isfinite(second)) { + float margin = std::max(-60.0f, std::min(60.0f, best - second)); + confidence = 1.0f / (1.0f + std::exp(-margin)); + } + return std::max(0.0f, std::min(1.0f, 1.0f - confidence)); +} + +uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty) { + size_t vocab = desc.shape.empty() ? 0 : desc.shape.back(); + uint32_t best = 0; + float best_v = -std::numeric_limits::infinity(); + float second_v = -std::numeric_limits::infinity(); + const auto& tool_bias = tool_constrainer_.get_bias(); + const std::vector* tool_dense = tool_constrainer_.get_dense_bias(); + { + uint32_t gidx; float gbest, gsecond; + if (!tool_dense && tool_bias.empty() && vocab_bias_.empty() && suppressed_token_id_ < 0 && + cactus_graph_metal_argmax(&gidx, &gbest, &gsecond)) { + if (out_uncertainty) *out_uncertainty = uncertainty_from_margin(gbest, gsecond); + return gidx; + } + } + auto score_with_bias = [&](size_t token_id, float value) { + if (tool_dense && token_id < tool_dense->size()) value += (*tool_dense)[token_id]; + auto tool_it = tool_bias.find(static_cast(token_id)); + if (tool_it != tool_bias.end()) value += tool_it->second; + auto vocab_it = vocab_bias_.find(static_cast(token_id)); + if (vocab_it != vocab_bias_.end()) value += vocab_it->second; + return value; + }; + auto observe_logit = [&](size_t i, float v) { + if (static_cast(i) == suppressed_token_id_) return; + v = score_with_bias(i, v); + if (v > best_v) { + second_v = best_v; + best_v = v; + best = static_cast(i); + } else if (v > second_v) { + second_v = v; + } + }; + if (desc.precision == Precision::FP32) { + float* p = static_cast(ptr) + row_off; + for (size_t i = 0; i < vocab; ++i) observe_logit(i, p[i]); + } else if (desc.precision == Precision::FP16) { + __fp16* p = static_cast<__fp16*>(ptr) + row_off; + for (size_t i = 0; i < vocab; ++i) observe_logit(i, static_cast(p[i])); + } else { + int8_t* p = static_cast(ptr) + row_off; + for (size_t i = 0; i < vocab; ++i) observe_logit(i, static_cast(p[i])); + } + if (out_uncertainty) *out_uncertainty = uncertainty_from_margin(best_v, second_v); + return best; +} + +uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float* out_uncertainty) { + size_t out_node = static_cast(comp.output_node_ids.empty() ? 0 : comp.output_node_ids[0]); + const auto& desc = comp.graph->get_output_buffer(out_node); + void* ptr = comp.graph->get_output(out_node); + size_t vocab = desc.shape.empty() ? 0 : desc.shape.back(); + size_t seq = desc.shape.size() >= 2 ? desc.shape[desc.shape.size() - 2] : 1; + size_t row = seq > 0 ? seq - 1 : 0; + if (logit_row != std::numeric_limits::max()) { + row = std::min(logit_row, seq > 0 ? seq - 1 : 0); + } else if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { + row = std::min(last_logit_position_, seq > 0 ? seq - 1 : 0); + } + return argmax_logits_at(desc, ptr, row * vocab, out_uncertainty); +} + +std::vector Model::argmax_component_logits_batch(Component& comp, size_t batch) { + std::vector out(batch, 0); + if (batch == 0) return out; + size_t out_node = static_cast(comp.output_node_ids.empty() ? 0 : comp.output_node_ids[0]); + const auto& desc = comp.graph->get_output_buffer(out_node); + void* ptr = comp.graph->get_output(out_node); + size_t vocab = desc.shape.empty() ? 0 : desc.shape.back(); + if (vocab == 0) return out; + size_t total_rows = desc.total_size / vocab; + size_t seq = total_rows / batch; + for (size_t b = 0; b < batch; ++b) { + size_t row = b * seq + (seq > 0 ? seq - 1 : 0); + out[b] = argmax_logits_at(desc, ptr, row * vocab, nullptr); + } + return out; +} + +uint32_t Model::argmax_last_logits(float* out_uncertainty) { + return argmax_component_logits(*decoder_, std::numeric_limits::max(), out_uncertainty); +} + +void Model::prepare_sampling_context(float repetition_penalty) { + samp_recent_.clear(); + samp_has_bias_ = false; + samp_penalty_ = repetition_penalty; + constexpr size_t kPenaltyWindow = 64; + size_t start = token_history_.size() > kPenaltyWindow ? token_history_.size() - kPenaltyWindow : 0; + for (size_t i = start; i < token_history_.size(); ++i) { + uint32_t id = token_history_[i]; + bool dup = false; + for (uint32_t r : samp_recent_) if (r == id) { dup = true; break; } + if (!dup) samp_recent_.push_back(id); + } + const std::vector* tool_dense = tool_constrainer_.get_dense_bias(); + const auto& tool_bias = tool_constrainer_.get_bias(); + if (tool_dense || !tool_bias.empty() || !vocab_bias_.empty()) { + size_t n = tokenizer_ ? tokenizer_->get_vocab_size() : 0; + if (tool_dense && tool_dense->size() > n) n = tool_dense->size(); + if (n > 0) { + if (tool_dense) samp_bias_dense_.assign(tool_dense->begin(), tool_dense->end()); + else samp_bias_dense_.assign(n, 0.0f); + if (samp_bias_dense_.size() < n) samp_bias_dense_.resize(n, 0.0f); + for (const auto& kv : tool_bias) + if (kv.first < samp_bias_dense_.size()) samp_bias_dense_[kv.first] += kv.second; + for (const auto& kv : vocab_bias_) + if (kv.first < samp_bias_dense_.size()) samp_bias_dense_[kv.first] += kv.second; + samp_has_bias_ = true; + } + } + samp_ctx_active_ = (repetition_penalty != 1.0f && !samp_recent_.empty()) || + samp_has_bias_ || suppressed_token_id_ >= 0; + if (samp_ctx_active_) { + cactus_graph_set_sampling(samp_recent_.data(), (int)samp_recent_.size(), repetition_penalty, + samp_has_bias_ ? samp_bias_dense_.data() : nullptr, + samp_has_bias_ ? samp_bias_dense_.size() : 0, + suppressed_token_id_); + } else { + cactus_graph_clear_sampling(); + } +} + +uint32_t Model::sample_component_logits(Component& comp, float temperature, float top_p, size_t top_k, + float min_p, bool greedy, float* out_uncertainty) { + size_t out_node = static_cast(comp.output_node_ids.empty() ? 0 : comp.output_node_ids[0]); + const auto& desc = comp.graph->get_output_buffer(out_node); + void* ptr = comp.graph->get_output(out_node); + size_t vocab = desc.shape.empty() ? 0 : desc.shape.back(); + size_t seq = desc.shape.size() >= 2 ? desc.shape[desc.shape.size() - 2] : 1; + size_t row = seq > 0 ? seq - 1 : 0; + if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { + row = std::min(last_logit_position_, seq > 0 ? seq - 1 : 0); + } + size_t off = row * vocab; + if (vocab == 0 || (desc.precision != Precision::FP16 && desc.precision != Precision::FP32)) { + return argmax_component_logits(comp, std::numeric_limits::max(), out_uncertainty); + } + const bool fp16 = desc.precision == Precision::FP16; + __fp16* h = fp16 ? static_cast<__fp16*>(ptr) + off : nullptr; + float* f = fp16 ? nullptr : static_cast(ptr) + off; + auto get = [&](size_t i) -> float { return fp16 ? (float)h[i] : f[i]; }; + auto put = [&](size_t i, float v) { if (fp16) h[i] = (__fp16)v; else f[i] = v; }; + + if (samp_ctx_active_ && !cactus_graph_metal_adjusted()) { + if (samp_penalty_ != 1.0f) { + for (uint32_t id : samp_recent_) { + if (id >= vocab) continue; + float v = get(id); + put(id, v > 0.0f ? v / samp_penalty_ : v * samp_penalty_); + } + } + if (suppressed_token_id_ >= 0 && (size_t)suppressed_token_id_ < vocab) + put((size_t)suppressed_token_id_, -65504.0f); + } + if (greedy) { + uint32_t gidx; float gbest, gsecond; + if ((!samp_has_bias_ || cactus_graph_metal_argmax_biased()) && + cactus_graph_metal_argmax(&gidx, &gbest, &gsecond)) { + if (out_uncertainty) *out_uncertainty = uncertainty_from_margin(gbest, gsecond); + return gidx; + } + } + const float* bd = samp_has_bias_ ? samp_bias_dense_.data() : nullptr; + const size_t bn = samp_has_bias_ ? samp_bias_dense_.size() : 0; + auto biased = [&](size_t i) -> float { + float v = get(i); + if (i < bn) v += bd[i]; + return v; + }; + + if (greedy) { + uint32_t best = 0; + float bv = -std::numeric_limits::infinity(), sv = bv; + for (size_t i = 0; i < vocab; ++i) { + float v = biased(i); + if (v > bv) { sv = bv; bv = v; best = (uint32_t)i; } + else if (v > sv) sv = v; + } + if (out_uncertainty) *out_uncertainty = uncertainty_from_margin(bv, sv); + return best; + } + + size_t K = std::min(std::max(top_k, 1), 512); + std::vector> cand; + cand.reserve(2 * K + 16); + float kmin = -std::numeric_limits::infinity(); + for (size_t i = 0; i < vocab; ++i) { + float v = biased(i); + if (v <= kmin) continue; + cand.emplace_back(v, (uint32_t)i); + if (cand.size() >= 2 * K) { + std::nth_element(cand.begin(), cand.begin() + (K - 1), cand.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + cand.resize(K); + kmin = cand.back().first; + for (const auto& c : cand) if (c.first < kmin) kmin = c.first; + } + } + std::sort(cand.begin(), cand.end(), [](const auto& a, const auto& b) { return a.first > b.first; }); + if (cand.size() > K) cand.resize(K); + if (cand.empty()) return 0; + if (out_uncertainty) { + float sv = cand.size() > 1 ? cand[1].first : -std::numeric_limits::infinity(); + *out_uncertainty = uncertainty_from_margin(cand[0].first, sv); + } + if (cand.size() == 1) return cand[0].second; + + float t = std::max(temperature, 1e-3f); + float maxl = cand[0].first; + std::vector probs(cand.size()); + double denom = 0.0; + for (size_t i = 0; i < cand.size(); ++i) { + probs[i] = std::exp((cand[i].first - maxl) / t); + denom += probs[i]; + } + size_t keep = cand.size(); + if (top_p > 0.0f && top_p < 1.0f) { + double cum = 0.0; + for (size_t i = 0; i < cand.size(); ++i) { + cum += probs[i] / denom; + if (cum >= top_p) { keep = i + 1; break; } + } + } + if (min_p > 0.0f) { + float cut = min_p * probs[0]; + while (keep > 1 && probs[keep - 1] < cut) --keep; + } + double mass = 0.0; + for (size_t i = 0; i < keep; ++i) mass += probs[i]; + double r = std::uniform_real_distribution(0.0, mass)(sample_rng_); + double cum = 0.0; + for (size_t i = 0; i < keep; ++i) { + cum += probs[i]; + if (r <= cum) return cand[i].second; + } + return cand[keep - 1].second; +} + +bool Model::prefill_and_sample_first_token(const std::vector& tokens, uint32_t& out_token, + float* out_uncertainty) { + reset_prefill_stats(); + if (out_uncertainty) *out_uncertainty = 0.0f; + if (tokens.empty() || !decoder_ || cache_total_seq_len_ != 0) { + return false; + } + if (decode_route_ == DecodeRoute::ENCODER_CROSS_KV_STEP && encoder_cross_kv_source_kind_ == "text_tokens") { + std::vector source_tokens = tokens; + std::vector decoder_seed = {config_.decoder_start_token_id}; + if (!source_tokens.empty() && source_tokens.back() == config_.decoder_start_token_id) { + decoder_seed = {source_tokens.back()}; + source_tokens.pop_back(); + } + if (!prepare_encoder_cross_kv_from_text(source_tokens)) { + return false; + } + std::vector emitted = run_encoder_cross_kv_decode_loop( + decoder_seed, + 1, + {}, + nullptr); + if (emitted.empty()) return false; + out_token = emitted.front(); + return true; + } + cache_token_ids_ = tokens; + if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { + context_tokens_.insert(context_tokens_.end(), tokens.begin(), tokens.end()); + run_full_context_text(); + cache_total_seq_len_ = context_tokens_.size(); + out_token = argmax_last_logits(out_uncertainty); + record_sampled_token(out_token); + return true; + } + if (decode_route_ == DecodeRoute::DIRECT_DECODER_STEP) { + for (size_t i = 0; i < tokens.size(); ++i) { + run_step(tokens[i], i, i + 1 == tokens.size()); + } + cache_total_seq_len_ = tokens.size(); + out_token = argmax_last_logits(out_uncertainty); + record_sampled_token(out_token); + return true; + } + if (!encoder_) { + return false; + } + ChunkedPrefillResult chunked; + if (decoder_prefill_) { + chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), true); + if (chunked.logical_tokens == tokens.size() && chunked.padding_tokens > 0 && tokens.size() > 0) { + // Cache already moved into the step component; drop the padded last row and re-run it for logits. + set_cache_current_len(*decoder_, tokens.size() - 1); + cache_total_seq_len_ = tokens.size() - 1; + run_step(tokens.back(), cache_total_seq_len_, true, /*use_fused=*/false); + ++cache_total_seq_len_; + out_token = argmax_last_logits(out_uncertainty); + record_sampled_token(out_token); + last_prefill_scalar_tail_tokens_ = 1; + maybe_roll_compact(); + return true; + } + cache_total_seq_len_ += chunked.logical_tokens; + } + for (size_t i = chunked.logical_tokens; i < tokens.size(); ++i) { + run_step(tokens[i], cache_total_seq_len_, i + 1 == tokens.size(), /*use_fused=*/false); + ++cache_total_seq_len_; + } + last_prefill_scalar_tail_tokens_ = tokens.size() - chunked.logical_tokens; + if (chunked.logical_tokens == tokens.size() && chunked.logical_tokens > 0 && decoder_prefill_) { + out_token = argmax_component_logits(*decoder_prefill_, chunked.last_logit_row, out_uncertainty); + } else { + out_token = argmax_last_logits(out_uncertainty); + } + record_sampled_token(out_token); + maybe_roll_compact(); + return true; +} + +void Model::prefill(const std::vector& tokens, size_t /*chunk_size*/, const std::string& /*profile_file*/, bool prepare_decode) { + reset_prefill_stats(); + if (decode_route_ == DecodeRoute::ENCODER_CROSS_KV_STEP && encoder_cross_kv_source_kind_ == "text_tokens") { + (void)prepare_decode; + prepare_encoder_cross_kv_from_text(tokens); + return; + } + if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { + context_tokens_.insert(context_tokens_.end(), tokens.begin(), tokens.end()); + if (!context_tokens_.empty()) run_full_context_text(); + cache_total_seq_len_ = context_tokens_.size(); + cache_token_ids_ = context_tokens_; + return; + } + ChunkedPrefillResult chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), prepare_decode); + cache_total_seq_len_ += chunked.logical_tokens; + for (size_t i = chunked.logical_tokens; i < tokens.size(); ++i) { + run_step(tokens[i], cache_total_seq_len_, /*read_logits=*/false, /*use_fused=*/false); + ++cache_total_seq_len_; + } + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); + last_prefill_scalar_tail_tokens_ = tokens.size() - chunked.logical_tokens; + + if (prepare_decode) { + // After the prompt reaches full length here -- never mid-chunk -- bound it to target_len. + maybe_roll_compact(); + } +} + +void Model::prefill_with_images(const std::vector& tokens, + const std::vector& image_paths, + const std::string& profile_file) { + prefill_with_media(tokens, image_paths, {}, profile_file); +} + +void Model::prefill_with_audio(const std::vector& tokens, + const std::vector>& audio_features_per_message, + const std::string& profile_file) { + prefill_with_media(tokens, {}, audio_features_per_message, profile_file); +} + +namespace { + +void write_typed_buffer(std::vector& buf, Precision dst_prec, + const void* src_data, size_t src_bytes, Precision src_prec) { + if (dst_prec == src_prec) { + size_t to_copy = std::min(src_bytes, buf.size()); + std::memcpy(buf.data(), src_data, to_copy); + if (to_copy < buf.size()) std::memset(buf.data() + to_copy, 0, buf.size() - to_copy); + return; + } + const size_t src_elem = PrecisionTraits::size_of(src_prec); + const size_t dst_elem = PrecisionTraits::size_of(dst_prec); + const size_t src_count = src_elem ? src_bytes / src_elem : 0; + const size_t dst_count = dst_elem ? buf.size() / dst_elem : 0; + const size_t n = std::min(src_count, dst_count); + auto load_float = [&](size_t i) -> float { + if (src_prec == Precision::FP16) return static_cast(reinterpret_cast(src_data)[i]); + if (src_prec == Precision::FP32) return reinterpret_cast(src_data)[i]; + return static_cast(reinterpret_cast(src_data)[i]); + }; + for (size_t i = 0; i < n; ++i) { + float v = load_float(i); + if (dst_prec == Precision::FP16) reinterpret_cast<__fp16*>(buf.data())[i] = static_cast<__fp16>(v); + else if (dst_prec == Precision::FP32) reinterpret_cast(buf.data())[i] = v; + else reinterpret_cast(buf.data())[i] = static_cast(v); + } + if (n < dst_count) { + std::memset(buf.data() + n * dst_elem, 0, (dst_count - n) * dst_elem); + } +} + +static inline void write_int_element(uint8_t* buf, Precision prec, size_t index, int64_t v) { + switch (prec) { + case Precision::FP32: reinterpret_cast(buf)[index] = static_cast(v); break; + case Precision::FP16: reinterpret_cast<__fp16*>(buf)[index] = static_cast<__fp16>(v); break; + case Precision::INT8: reinterpret_cast(buf)[index] = static_cast(v); break; + default: { + size_t elem = PrecisionTraits::size_of(prec); + if (elem == 8) reinterpret_cast(buf)[index] = v; + else if (elem == 4) reinterpret_cast(buf)[index] = static_cast(v); + break; + } + } +} + +static inline size_t typed_buf_capacity(const std::vector& buf, Precision prec) { + size_t elem = PrecisionTraits::size_of(prec); + return elem ? buf.size() / elem : 0; +} + +static inline void zero_fill_remainder(std::vector& buf, Precision prec, size_t written, size_t cap) { + if (written < cap) { + size_t elem = PrecisionTraits::size_of(prec); + std::memset(buf.data() + written * elem, 0, (cap - written) * elem); + } +} + +void fill_int_buffer(std::vector& buf, Precision prec, int64_t value, size_t count) { + const size_t cap = typed_buf_capacity(buf, prec); + const size_t n = std::min(cap, count); + for (size_t i = 0; i < n; ++i) write_int_element(buf.data(), prec, i, value); + zero_fill_remainder(buf, prec, n, cap); +} + +void write_tokens_buffer(std::vector& buf, Precision prec, + const std::vector& tokens, size_t offset) { + const size_t cap = typed_buf_capacity(buf, prec); + const size_t avail = (offset < tokens.size()) ? (tokens.size() - offset) : 0; + const size_t n = std::min(cap, avail); + for (size_t i = 0; i < n; ++i) write_int_element(buf.data(), prec, i, static_cast(tokens[offset + i])); + zero_fill_remainder(buf, prec, n, cap); +} + +void write_int_vector_buffer(std::vector& buf, Precision prec, const std::vector& values) { + const size_t cap = typed_buf_capacity(buf, prec); + const size_t n = std::min(cap, values.size()); + for (size_t i = 0; i < n; ++i) write_int_element(buf.data(), prec, i, values[i]); + zero_fill_remainder(buf, prec, n, cap); +} + +std::vector qwen3_vl_position_ids(const std::vector& tokens, + size_t capacity, + const std::vector& images, + uint32_t image_token_id) { + std::vector positions(3 * capacity, 0); + size_t token_index = 0; + size_t image_index = 0; + int64_t current_pos = 0; + while (token_index < tokens.size() && token_index < capacity) { + if (image_token_id != 0 && tokens[token_index] == image_token_id) { + if (image_index >= images.size()) { + throw std::runtime_error("Qwen3-VL prompt contains more image token groups than image inputs"); + } + const auto& image = images[image_index++]; + const size_t merge_size = 2; + const size_t grid_t = image.grid_t; + const size_t llm_grid_h = image.grid_h / merge_size; + const size_t llm_grid_w = image.grid_w / merge_size; + const size_t image_seq = grid_t * llm_grid_h * llm_grid_w; + size_t count = 0; + while (token_index + count < tokens.size() + && token_index + count < capacity + && tokens[token_index + count] == image_token_id) { + ++count; + } + if (count != image_seq) { + throw std::runtime_error("Qwen3-VL image token count does not match vision feature grid"); + } + size_t local = 0; + for (size_t t = 0; t < grid_t; ++t) { + (void)t; + for (size_t h = 0; h < llm_grid_h; ++h) { + for (size_t w = 0; w < llm_grid_w; ++w) { + size_t pos = token_index + local++; + positions[pos] = current_pos; + positions[capacity + pos] = current_pos + static_cast(h); + positions[2 * capacity + pos] = current_pos + static_cast(w); + } + } + } + current_pos += static_cast(std::max(image.grid_h, image.grid_w) / merge_size); + token_index += count; + continue; + } + + size_t text_count = 0; + while (token_index + text_count < tokens.size() + && token_index + text_count < capacity + && (image_token_id == 0 || tokens[token_index + text_count] != image_token_id)) { + size_t pos = token_index + text_count; + int64_t value = current_pos + static_cast(text_count); + positions[pos] = value; + positions[capacity + pos] = value; + positions[2 * capacity + pos] = value; + ++text_count; + } + current_pos += static_cast(text_count); + token_index += text_count; + } + return positions; +} + +} // namespace + +void Model::reset_encoder_cross_kv_route_state() { + if (!decoder_) return; + reset_component_cache_states(*decoder_); + cache_total_seq_len_ = 0; + token_history_.clear(); + encoder_cross_kv_ready_ = false; + encoder_cross_kv_source_len_ = 0; +} + +bool Model::finish_encoder_cross_kv_prepare() { + if (!source_encoder_ || !decoder_cross_kv_ || !decoder_) return false; + source_encoder_->graph->execute(); + copy_component_outputs_to_inputs(*source_encoder_, *decoder_cross_kv_); + decoder_cross_kv_->graph->execute(); + copy_component_outputs_to_inputs(*source_encoder_, *decoder_); + if (!copy_cross_kv_outputs_to_decoder_cache_inputs(*decoder_cross_kv_, *decoder_, encoder_cross_kv_source_len_)) { + copy_component_outputs_to_inputs(*decoder_cross_kv_, *decoder_); + } + encoder_cross_kv_ready_ = true; + return true; +} + +bool Model::prepare_encoder_cross_kv_from_text(const std::vector& tokens) { + if (!source_encoder_ || !decoder_cross_kv_ || !decoder_ || tokens.empty()) return false; + if (encoder_cross_kv_source_kind_ != "text_tokens") return false; + + reset_encoder_cross_kv_route_state(); + for (auto& buf : source_encoder_->input_buffers) { + std::fill(buf.begin(), buf.end(), 0); + } + + int ids_idx = input_index(*source_encoder_, "input_ids"); + if (ids_idx < 0) { + CACTUS_LOG_ERROR("model", "source encoder missing input_ids input"); + return false; + } + size_t ids_node = static_cast(source_encoder_->runtime_input_node_ids[ids_idx]); + const auto& ids_desc = source_encoder_->graph->get_output_buffer(ids_node); + if (tokens.size() > ids_desc.total_size) { + CACTUS_LOG_ERROR("model", "source token count exceeds source encoder capacity"); + return false; + } + encoder_cross_kv_source_len_ = tokens.size(); + for (size_t i = 0; i < tokens.size(); ++i) { + write_int_input_at(*source_encoder_, "input_ids", i, static_cast(tokens[i])); + } + + int mask_idx = input_index(*source_encoder_, "attention_mask"); + if (mask_idx >= 0) { + for (size_t i = 0; i < tokens.size(); ++i) { + write_int_input_at(*source_encoder_, "attention_mask", i, 1); + } + } + + return finish_encoder_cross_kv_prepare(); +} + +bool Model::prepare_encoder_cross_kv_from_audio(const std::vector& audio_features) { + if (!source_encoder_ || !decoder_cross_kv_ || !decoder_) return false; + int feature_idx = input_index(*source_encoder_, "input_features"); + if (feature_idx < 0) { + feature_idx = input_index(*source_encoder_, "audio_features"); + } + if (feature_idx < 0) { + CACTUS_LOG_ERROR("model", "audio source encoder missing input_features/audio_features input"); + return false; + } + + reset_encoder_cross_kv_route_state(); + + auto& feature_buf = source_encoder_->input_buffers[feature_idx]; + size_t feature_node = static_cast(source_encoder_->runtime_input_node_ids[feature_idx]); + const auto& feature_desc = source_encoder_->graph->get_output_buffer(feature_node); + write_typed_buffer( + feature_buf, + feature_desc.precision, + audio_features.data(), + audio_features.size() * sizeof(float), + Precision::FP32); + + return finish_encoder_cross_kv_prepare(); +} + +bool Model::run_encoder_cross_kv_decoder_step(uint32_t token_id, size_t position) { + if (!encoder_cross_kv_ready_ || !decoder_) return false; + int ids_idx = input_index(*decoder_, "decoder_input_ids"); + const char* ids_name = "decoder_input_ids"; + if (ids_idx < 0) { + ids_idx = input_index(*decoder_, "input_ids"); + ids_name = "input_ids"; + } + int pos_idx = input_index(*decoder_, "position_ids"); + if (ids_idx < 0 || pos_idx < 0) { + CACTUS_LOG_ERROR("model", "decoder_step missing decoder_input_ids/input_ids or position_ids input"); + return false; + } + write_int_input(*decoder_, ids_name, static_cast(token_id)); + write_int_input(*decoder_, "position_ids", static_cast(position)); + decoder_->graph->execute(); + return true; +} + +std::vector Model::run_encoder_cross_kv_decode_loop( + const std::vector& decoder_prompt_tokens, + size_t max_tokens, + const std::vector>& stop_token_sequences, + const std::atomic* should_stop) { + std::vector emitted; + if (!encoder_cross_kv_ready_ || decoder_prompt_tokens.empty() || max_tokens == 0) return emitted; + + std::vector tokens = decoder_prompt_tokens; + auto stopped = [&]() { + for (const auto& stop_seq : stop_token_sequences) { + if (stop_seq.empty() || emitted.size() < stop_seq.size()) continue; + if (std::equal(stop_seq.rbegin(), stop_seq.rend(), emitted.rbegin())) return true; + } + return false; + }; + + for (size_t i = 0; i < max_tokens; ++i) { + if (should_stop && should_stop->load()) break; + const size_t start = cache_total_seq_len_ < tokens.size() ? cache_total_seq_len_ : tokens.size() - 1; + for (size_t pos = start; pos < tokens.size(); ++pos) { + if (!run_encoder_cross_kv_decoder_step(tokens[pos], pos)) return emitted; + } + cache_total_seq_len_ = tokens.size(); + + uint32_t next_token = argmax_last_logits(); + record_sampled_token(next_token); + emitted.push_back(next_token); + if (stopped()) break; + tokens.push_back(next_token); + } + + return emitted; +} + +bool Model::build_lm_encoder_outputs_dynamic_gemma4( + const std::vector& tokens, + std::map>& store_bytes, + std::map& store_prec, + std::map>& store_shape) { + if (!encoder_ || !lm_encoder_media_step_ || tokens.empty()) return false; + + const uint32_t image_tok = config_.image_token_id; + const uint32_t audio_tok = config_.audio_token_id; + + auto audio_it = media_features_.find("audio_features"); + const bool have_audio_features = audio_it != media_features_.end() && !audio_it->second.empty(); + size_t audio_rows = 0; + size_t audio_row_bytes = 0; + Precision audio_prec = Precision::FP16; + if (have_audio_features) { + const auto& shape = media_feature_shapes_["audio_features"]; + audio_rows = shape.size() >= 2 ? shape[shape.size() - 2] : 0; + audio_prec = media_feature_precisions_["audio_features"]; + audio_row_bytes = audio_rows > 0 ? audio_it->second.size() / audio_rows : audio_it->second.size(); + } + + auto image_it = media_features_.find("image_features"); + const bool have_image_features = image_it != media_features_.end() && !image_it->second.empty(); + size_t image_rows = 0; + size_t image_row_bytes = 0; + Precision image_prec = Precision::FP16; + if (have_image_features) { + const auto& shape = media_feature_shapes_["image_features"]; + image_rows = shape.size() >= 2 ? shape[shape.size() - 2] : 0; + image_prec = media_feature_precisions_["image_features"]; + image_row_bytes = image_rows > 0 ? image_it->second.size() / image_rows : image_it->second.size(); + } + + struct OutputInfo { + std::string name; + int text_idx = -1; + int media_idx = -1; + size_t per_token_bytes = 0; + Precision precision = Precision::FP16; + std::vector shape_template; + }; + + std::vector outputs; + for (size_t i = 0; i < encoder_->logical_outputs.size() && i < encoder_->output_node_ids.size(); ++i) { + OutputInfo info; + info.name = encoder_->logical_outputs[i]; + info.text_idx = static_cast(i); + info.media_idx = output_index(*lm_encoder_media_step_, info.name); + if (info.media_idx < 0) { + throw std::runtime_error("lm_encoder_media_step missing output " + info.name); + } + size_t node_id = static_cast(encoder_->output_node_ids[i]); + const auto& desc = encoder_->graph->get_output_buffer(node_id); + info.per_token_bytes = desc.byte_size; + info.precision = desc.precision; + info.shape_template = desc.shape; + outputs.push_back(std::move(info)); + } + if (outputs.empty()) return false; + + const size_t token_count = tokens.size(); + for (const auto& info : outputs) { + store_bytes[info.name].assign(token_count * info.per_token_bytes, 0); + store_prec[info.name] = info.precision; + std::vector shape = info.shape_template; + if (shape.size() >= 2 && shape[shape.size() - 2] == 1) { + shape[shape.size() - 2] = token_count; + } else if (shape.size() == 1) { + shape[0] = token_count; + } + store_shape[info.name] = std::move(shape); + } + + size_t audio_idx = 0; + size_t image_idx = 0; + for (size_t pos = 0; pos < token_count; ++pos) { + const uint32_t token = tokens[pos]; + Component* component = encoder_; + const uint8_t* media_row = nullptr; + size_t media_row_bytes = 0; + Precision media_prec = Precision::FP16; + + if (audio_tok != 0 && token == audio_tok && have_audio_features) { + if (audio_idx >= audio_rows) { + throw std::runtime_error("Gemma4 prompt contains more audio tokens than audio feature rows"); + } + component = lm_encoder_media_step_; + media_row = audio_it->second.data() + audio_idx * audio_row_bytes; + media_row_bytes = audio_row_bytes; + media_prec = audio_prec; + ++audio_idx; + } else if (image_tok != 0 && token == image_tok && have_image_features) { + if (image_idx >= image_rows) { + throw std::runtime_error("Gemma4 prompt contains more image tokens than image feature rows"); + } + component = lm_encoder_media_step_; + media_row = image_it->second.data() + image_idx * image_row_bytes; + media_row_bytes = image_row_bytes; + media_prec = image_prec; + ++image_idx; + } + + if (component == lm_encoder_media_step_) { + int embeds_idx = input_index(*component, "inputs_embeds"); + if (embeds_idx < 0) { + throw std::runtime_error("lm_encoder_media_step missing inputs_embeds input"); + } + auto& buf = component->input_buffers[embeds_idx]; + size_t node_id = static_cast(component->runtime_input_node_ids[embeds_idx]); + const auto& desc = component->graph->get_output_buffer(node_id); + write_typed_buffer(buf, desc.precision, media_row, media_row_bytes, media_prec); + write_int_input(*component, "input_ids", 0); + write_int_input(*component, "position_ids", static_cast(pos)); + } else { + write_int_input(*component, "input_ids", static_cast(token)); + write_int_input(*component, "position_ids", static_cast(pos)); + } + component->graph->execute(); + + for (const auto& info : outputs) { + int out_idx = component == encoder_ ? info.text_idx : info.media_idx; + size_t node_id = static_cast(component->output_node_ids[out_idx]); + const auto& desc = component->graph->get_output_buffer(node_id); + if (desc.byte_size != info.per_token_bytes || desc.precision != info.precision) { + throw std::runtime_error("Gemma4 dynamic output shape mismatch for " + info.name); + } + const void* ptr = component->graph->get_output(node_id); + std::memcpy(store_bytes[info.name].data() + pos * info.per_token_bytes, ptr, info.per_token_bytes); + } + component->graph->release_runtime_buffers(); + } + encoder_->graph->release_all_weight_pages(); + lm_encoder_media_step_->graph->release_all_weight_pages(); + return true; +} + + +bool Model::run_chunk_prefill_path(const std::vector& tokens, + const std::vector& image_paths, + const std::vector>& audio_features_per_message) { + if (cache_total_seq_len_ > 0) return false; + const bool have_images = !image_paths.empty() && vision_encoder_ != nullptr; + bool any_audio = false; + for (const auto& mel : audio_features_per_message) { if (!mel.empty()) { any_audio = true; break; } } + const bool have_audio = any_audio && audio_encoder_ != nullptr; + std::vector qwen_images; + + if (have_images) { + if (!load_component_graph(*vision_encoder_)) { + throw std::runtime_error("failed to load vision_encoder"); + } + for (const std::string& logical : vision_encoder_->logical_outputs) { + media_features_.erase(logical); + media_feature_shapes_.erase(logical); + media_feature_precisions_.erase(logical); + } + const bool lfm2_vision = family_ == "lfm2_vl"; + if (lfm2_vision) { + if (!vision_projector_) throw std::runtime_error("lfm2_vl requires a vision_projector component"); + if (!load_lfm2_vl_position_grid()) { + throw std::runtime_error("lfm2_vl vision position-embedding grid is missing from the bundle"); + } + if (!load_component_graph(*vision_projector_)) throw std::runtime_error("failed to load vision_projector"); + media_features_.erase("image_features"); + media_feature_shapes_.erase("image_features"); + media_feature_precisions_.erase("image_features"); + } + for (const auto& path : image_paths) { + if (family_ == "lfm2_vl") { + encode_lfm2_vl_image_into_features(path); + continue; + } else if (family_ == "qwen3_5" || family_ == "qwen3_vl" || config_.model_type == Config::ModelType::QWEN) { + Qwen3VlImagePreprocessed prep = preprocess_qwen3_vl_image(path, config_); + int pv_idx = input_index(*vision_encoder_, "pixel_values"); + if (pv_idx < 0) { + throw std::runtime_error("Qwen3-VL vision_encoder missing pixel_values input"); + } + auto& pv_buf = vision_encoder_->input_buffers[pv_idx]; + size_t pv_node = static_cast(vision_encoder_->runtime_input_node_ids[pv_idx]); + const auto& pv_desc = vision_encoder_->graph->get_output_buffer(pv_node); + write_typed_buffer(pv_buf, pv_desc.precision, + prep.pixel_values.data(), + prep.pixel_values.size() * sizeof(float), + Precision::FP32); + qwen_images.push_back(std::move(prep)); + } else { + Gemma4ImagePreprocessed prep = preprocess_gemma4_image(path, config_); + int pv_idx = input_index(*vision_encoder_, "pixel_values"); + if (pv_idx >= 0) { + auto& pv_buf = vision_encoder_->input_buffers[pv_idx]; + size_t pv_node = static_cast(vision_encoder_->runtime_input_node_ids[pv_idx]); + const auto& pv_desc = vision_encoder_->graph->get_output_buffer(pv_node); + write_typed_buffer(pv_buf, pv_desc.precision, + prep.pixel_values.data(), + prep.pixel_values.size() * sizeof(float), + Precision::FP32); + } + int pp_idx = input_index(*vision_encoder_, "pixel_position_ids"); + if (pp_idx >= 0) { + auto& pp_buf = vision_encoder_->input_buffers[pp_idx]; + size_t pp_node = static_cast(vision_encoder_->runtime_input_node_ids[pp_idx]); + const auto& pp_desc = vision_encoder_->graph->get_output_buffer(pp_node); + const size_t elem = PrecisionTraits::size_of(pp_desc.precision); + const size_t cap = elem ? pp_buf.size() / elem : 0; + const size_t n = std::min(cap, prep.pixel_position_ids.size()); + for (size_t i = 0; i < n; ++i) { + int64_t v = prep.pixel_position_ids[i]; + switch (pp_desc.precision) { + case Precision::FP32: reinterpret_cast(pp_buf.data())[i] = static_cast(v); break; + case Precision::FP16: reinterpret_cast<__fp16*>(pp_buf.data())[i] = static_cast<__fp16>(v); break; + case Precision::INT8: reinterpret_cast(pp_buf.data())[i] = static_cast(v); break; + default: + if (elem == 8) reinterpret_cast(pp_buf.data())[i] = v; + else if (elem == 4) reinterpret_cast(pp_buf.data())[i] = static_cast(v); + break; + } + } + if (n < cap) std::memset(pp_buf.data() + n * elem, 0, (cap - n) * elem); + } + } + vision_encoder_->graph->execute(); + for (size_t i = 0; i < vision_encoder_->output_node_ids.size() + && i < vision_encoder_->logical_outputs.size(); ++i) { + const std::string& name = vision_encoder_->logical_outputs[i]; + size_t node_id = static_cast(vision_encoder_->output_node_ids[i]); + const auto& desc = vision_encoder_->graph->get_output_buffer(node_id); + void* ptr = vision_encoder_->graph->get_output(node_id); + auto& slot = media_features_[name]; + const size_t prev_bytes = slot.size(); + slot.resize(prev_bytes + desc.byte_size); + std::memcpy(slot.data() + prev_bytes, ptr, desc.byte_size); + auto shape_it = media_feature_shapes_.find(name); + if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { + media_feature_shapes_[name] = desc.shape; + } else if (desc.shape.size() >= 2 && shape_it->second.size() == desc.shape.size()) { + shape_it->second[shape_it->second.size() - 2] += desc.shape[desc.shape.size() - 2]; + } + media_feature_precisions_[name] = desc.precision; + } + vision_encoder_->graph->release_runtime_buffers(); + vision_encoder_->graph->release_all_weight_pages(); + } + if (lfm2_vision) { + unload_component_graph(*vision_projector_); + } + unload_component_graph(*vision_encoder_); + } + + if (have_audio) { + run_audio_encoder_messages(audio_features_per_message); + } + + std::map> store_bytes; + std::map store_prec; + std::map> store_shape; + const bool needs_dynamic_walk = family_ == "gemma4" && (have_images || have_audio) && lm_encoder_media_step_ != nullptr; + + if (needs_dynamic_walk) { + if (!build_lm_encoder_outputs_dynamic_gemma4(tokens, store_bytes, store_prec, store_shape)) { + return false; + } + } else { + if (!load_component_graph(*lm_encoder_)) { + throw std::runtime_error("failed to load lm_encoder"); + } + int ids_idx = input_index(*lm_encoder_, "input_ids"); + if (ids_idx >= 0) { + auto& ids_buf = lm_encoder_->input_buffers[ids_idx]; + size_t ids_node = static_cast(lm_encoder_->runtime_input_node_ids[ids_idx]); + const auto& ids_desc = lm_encoder_->graph->get_output_buffer(ids_node); + write_tokens_buffer(ids_buf, ids_desc.precision, tokens, 0); + } + + int mask_idx = input_index(*lm_encoder_, "attention_mask"); + if (mask_idx >= 0) { + auto& mb = lm_encoder_->input_buffers[mask_idx]; + size_t mnode = static_cast(lm_encoder_->runtime_input_node_ids[mask_idx]); + const auto& mdesc = lm_encoder_->graph->get_output_buffer(mnode); + fill_int_buffer(mb, mdesc.precision, 1, tokens.size()); + } + + int pos_idx = input_index(*lm_encoder_, "position_ids"); + if (pos_idx >= 0) { + auto& pos_buf = lm_encoder_->input_buffers[pos_idx]; + size_t pos_node = static_cast(lm_encoder_->runtime_input_node_ids[pos_idx]); + const auto& pos_desc = lm_encoder_->graph->get_output_buffer(pos_node); + if (pos_desc.shape.size() >= 3 && pos_desc.shape[0] == 3 && !qwen_images.empty()) { + size_t capacity = pos_desc.shape[pos_desc.shape.size() - 1]; + auto positions = qwen3_vl_position_ids(tokens, capacity, qwen_images, config_.image_token_id); + write_int_vector_buffer(pos_buf, pos_desc.precision, positions); + } + } + + for (const auto& kv : media_features_) { + const std::string& name = kv.first; + int idx = input_index(*lm_encoder_, name); + if (idx < 0) continue; + auto& dst_buf = lm_encoder_->input_buffers[idx]; + size_t node_id = static_cast(lm_encoder_->runtime_input_node_ids[idx]); + const auto& desc = lm_encoder_->graph->get_output_buffer(node_id); + Precision src_prec = media_feature_precisions_[name]; + write_typed_buffer(dst_buf, desc.precision, + kv.second.data(), kv.second.size(), src_prec); + } + lm_encoder_->graph->execute(); + + for (size_t i = 0; i < lm_encoder_->output_node_ids.size() + && i < lm_encoder_->logical_outputs.size(); ++i) { + const std::string& name = lm_encoder_->logical_outputs[i]; + size_t node_id = static_cast(lm_encoder_->output_node_ids[i]); + const auto& desc = lm_encoder_->graph->get_output_buffer(node_id); + void* ptr = lm_encoder_->graph->get_output(node_id); + auto& slot = store_bytes[name]; + slot.assign(desc.byte_size, 0); + std::memcpy(slot.data(), ptr, desc.byte_size); + store_prec[name] = desc.precision; + store_shape[name] = desc.shape; + } + + if (have_images && family_ == "lfm2_vl") { + auto feat_it = media_features_.find("image_features"); + auto emb_it = store_bytes.find("inputs_embeds"); + if (feat_it != media_features_.end() && emb_it != store_bytes.end() && tokenizer_) { + const auto& emb_shape = store_shape["inputs_embeds"]; + const auto& feat_shape = media_feature_shapes_["image_features"]; + const size_t hidden = emb_shape.empty() ? 0 : emb_shape.back(); + const size_t emb_seq = emb_shape.size() >= 2 ? emb_shape[emb_shape.size() - 2] : 0; + const size_t feat_rows = feat_shape.empty() ? 0 : feat_shape[0]; + const Precision emb_prec = store_prec["inputs_embeds"]; + const Precision feat_prec = media_feature_precisions_["image_features"]; + const size_t emb_elem = PrecisionTraits::size_of(emb_prec); + const size_t feat_elem = PrecisionTraits::size_of(feat_prec); + const uint32_t image_tok = tokenizer_->get_image_token_id(); + uint8_t* emb_data = emb_it->second.data(); + const uint8_t* feat_data = feat_it->second.data(); + std::vector frow(hidden); + size_t fcur = 0; + const size_t limit = std::min(emb_seq, tokens.size()); + for (size_t p = 0; p < limit && fcur < feat_rows; ++p) { + if (tokens[p] != image_tok) continue; + typed_buffer_to_float(feat_data + fcur * hidden * feat_elem, feat_prec, frow.data(), hidden); + uint8_t* dst = emb_data + p * hidden * emb_elem; + for (size_t d = 0; d < hidden; ++d) { + switch (emb_prec) { + case Precision::FP32: reinterpret_cast(dst)[d] = frow[d]; break; + case Precision::FP16: reinterpret_cast<__fp16*>(dst)[d] = static_cast<__fp16>(frow[d]); break; + default: reinterpret_cast(dst)[d] = frow[d]; break; + } + } + ++fcur; + } + } + } + + lm_encoder_->graph->release_runtime_buffers(); + lm_encoder_->graph->release_all_weight_pages(); + unload_component_graph(*lm_encoder_); + } + + auto embeds_shape_it = store_shape.find("inputs_embeds"); + if (embeds_shape_it == store_shape.end()) { + return false; + } + size_t full_seq = 0; + { + const auto& sh = embeds_shape_it->second; + if (sh.size() >= 3) full_seq = sh[sh.size() - 2]; + else if (!sh.empty()) full_seq = sh[0]; + } + if (full_seq == 0) return false; + + size_t chunk_seq = 0; + { + if (!load_component_graph(*decoder_prefill_chunk_)) { + throw std::runtime_error("failed to load decoder_prefill_chunk"); + } + int idx = input_index(*decoder_prefill_chunk_, "inputs_embeds"); + if (idx < 0) return false; + size_t node_id = static_cast(decoder_prefill_chunk_->runtime_input_node_ids[idx]); + const auto& desc = decoder_prefill_chunk_->graph->get_output_buffer(node_id); + const auto& sh = desc.shape; + if (sh.size() >= 3) chunk_seq = sh[sh.size() - 2]; + else if (!sh.empty()) chunk_seq = sh[0]; + } + if (chunk_seq == 0) return false; + + std::map per_pos_bytes; + for (const auto& kv : store_bytes) { + per_pos_bytes[kv.first] = kv.second.size() / full_seq; + } + + size_t valid_seq = tokens.size(); + auto mask_it = store_bytes.find("attention_mask"); + if (mask_it != store_bytes.end() && per_pos_bytes.count("attention_mask")) { + Precision mp = store_prec["attention_mask"]; + size_t per = per_pos_bytes["attention_mask"]; + const uint8_t* mp_data = mask_it->second.data(); + size_t count = 0; + for (size_t i = 0; i < full_seq; ++i) { + const uint8_t* pos = mp_data + i * per; + bool nonzero = false; + switch (mp) { + case Precision::INT8: + nonzero = (*reinterpret_cast(pos) != 0); break; + case Precision::FP16: + nonzero = (static_cast(*reinterpret_cast(pos)) != 0.0f); break; + case Precision::FP32: + nonzero = (*reinterpret_cast(pos) != 0.0f); break; + default: + if (per == 8) nonzero = (*reinterpret_cast(pos) != 0); + else if (per == 4) nonzero = (*reinterpret_cast(pos) != 0); + else nonzero = (*pos != 0); + break; + } + if (nonzero) ++count; + } + if (count > 0) valid_seq = count; + } + valid_seq = std::min(valid_seq, full_seq); + const size_t whole_chunks_end = (valid_seq / chunk_seq) * chunk_seq; + for (size_t chunk_start = 0; chunk_start < whole_chunks_end; chunk_start += chunk_seq) { + for (const auto& kv : store_bytes) { + const std::string& name = kv.first; + int idx = input_index(*decoder_prefill_chunk_, name); + if (idx < 0) continue; + auto& dst_buf = decoder_prefill_chunk_->input_buffers[idx]; + size_t node_id = static_cast(decoder_prefill_chunk_->runtime_input_node_ids[idx]); + const auto& desc = decoder_prefill_chunk_->graph->get_output_buffer(node_id); + Precision src_prec = store_prec[name]; + size_t src_per_pos = per_pos_bytes[name]; + const uint8_t* src_ptr = kv.second.data() + chunk_start * src_per_pos; + size_t src_slice_bytes = chunk_seq * src_per_pos; + write_typed_buffer(dst_buf, desc.precision, src_ptr, src_slice_bytes, src_prec); + } + decoder_prefill_chunk_->graph->execute(); + } + if (whole_chunks_end > 0 && decoder_ != nullptr) { + move_cache_states(*decoder_prefill_chunk_, *decoder_); + decoder_prefill_chunk_->graph->release_runtime_buffers(); + unload_component_graph(*decoder_prefill_chunk_); + } + for (size_t pos = whole_chunks_end; pos < valid_seq; ++pos) { + for (const auto& kv : store_bytes) { + const std::string& name = kv.first; + int idx = input_index(*decoder_, name); + if (idx < 0) continue; + auto& dst_buf = decoder_->input_buffers[idx]; + size_t node_id = static_cast(decoder_->runtime_input_node_ids[idx]); + const auto& desc = decoder_->graph->get_output_buffer(node_id); + Precision src_prec = store_prec[name]; + size_t src_per_pos = per_pos_bytes[name]; + const uint8_t* src_ptr = kv.second.data() + pos * src_per_pos; + write_typed_buffer(dst_buf, desc.precision, src_ptr, src_per_pos, src_prec); + } + decoder_->graph->execute(); + } + cache_total_seq_len_ += valid_seq; + return true; +} + +void Model::prefill_with_media(const std::vector& tokens, + const std::vector& image_paths, + const std::vector>& audio_features_per_message, + const std::string& profile_file) { + if (tokens.empty()) return; + if (!image_paths.empty() && vision_encoder_ == nullptr) { + throw std::runtime_error("Model bundle does not include a vision_encoder for image input"); + } + bool any_audio = false; + for (const auto& mel : audio_features_per_message) { if (!mel.empty()) { any_audio = true; break; } } + if (any_audio && audio_encoder_ == nullptr) { + throw std::runtime_error("Model bundle does not include an audio_encoder for audio input"); + } + const bool have_images = !image_paths.empty(); + const bool have_audio = any_audio; + if (!have_images && !have_audio) { + prefill(tokens, get_prefill_chunk_size(), profile_file); + return; + } + + const bool can_chunk_prefill = + lm_encoder_ != nullptr && decoder_prefill_chunk_ != nullptr && + (vision_encoder_ != nullptr || audio_encoder_ != nullptr); + if (can_chunk_prefill) { + if (run_chunk_prefill_path(tokens, image_paths, audio_features_per_message)) { + (void)profile_file; + return; + } + } + if (!supports_warm_media_injection()) { + CACTUS_LOG_WARN("model", "Bundle supports neither chunk-prefill nor warm media injection; falling back to text-only prefill"); + prefill(tokens, get_prefill_chunk_size(), profile_file); + return; + } + + if (have_images) { + for (const auto& path : image_paths) { + run_vision_encoder(path); + } + } + if (have_audio) { + run_audio_encoder_messages(audio_features_per_message); + } + + std::string image_feature_name; + Precision image_feature_prec = Precision::FP16; + size_t image_row_bytes = 0; + if (have_images) { + const std::vector candidates = {"image_features", "image_embeddings", "vision_features", "inputs_embeds"}; + for (const auto& name : candidates) { + if (media_features_.count(name)) { image_feature_name = name; break; } + } + if (image_feature_name.empty() && !media_features_.empty()) { + image_feature_name = media_features_.begin()->first; + } + if (!image_feature_name.empty()) { + const auto& shape = media_feature_shapes_[image_feature_name]; + image_feature_prec = media_feature_precisions_[image_feature_name]; + if (shape.size() >= 2) { + size_t rows = (shape.size() >= 3) ? shape[shape.size() - 2] : shape[0]; + size_t total = media_features_[image_feature_name].size(); + image_row_bytes = rows > 0 ? total / rows : total; + } else { + image_row_bytes = media_features_[image_feature_name].size(); + } + } + } + + std::string audio_feature_name; + Precision audio_feature_prec = Precision::FP16; + size_t audio_row_bytes = 0; + if (have_audio) { + const std::vector candidates = {"audio_features", "audio_embeddings", "encoder_hidden_states", "inputs_embeds"}; + for (const auto& name : candidates) { + if (media_features_.count(name) && name != image_feature_name) { audio_feature_name = name; break; } + } + if (audio_feature_name.empty()) { + for (const auto& kv : media_features_) { + if (kv.first != image_feature_name) { audio_feature_name = kv.first; break; } + } + } + if (!audio_feature_name.empty()) { + const auto& shape = media_feature_shapes_[audio_feature_name]; + audio_feature_prec = media_feature_precisions_[audio_feature_name]; + if (shape.size() >= 2) { + size_t rows = (shape.size() >= 3) ? shape[shape.size() - 2] : shape[0]; + size_t total = media_features_[audio_feature_name].size(); + audio_row_bytes = rows > 0 ? total / rows : total; + } else { + audio_row_bytes = media_features_[audio_feature_name].size(); + } + } + } + + size_t image_consumed = 0; + size_t audio_consumed = 0; + const uint32_t image_tok = config_.image_token_id; + const uint32_t audio_tok = config_.audio_token_id; + const bool warm_media = supports_warm_media_injection(); + + for (size_t i = 0; i < tokens.size(); ++i) { + uint32_t t = tokens[i]; + size_t pos = cache_total_seq_len_ + i; + if (image_tok != 0 && t == image_tok && !image_feature_name.empty() && warm_media) { + const auto& feat = media_features_[image_feature_name]; + const uint8_t* row = feat.data() + image_consumed * image_row_bytes; + if (image_consumed * image_row_bytes + image_row_bytes <= feat.size()) { + run_media_step(pos, row, image_row_bytes, image_feature_prec); + ++image_consumed; + continue; + } + } + if (audio_tok != 0 && t == audio_tok && !audio_feature_name.empty() && warm_media) { + const auto& feat = media_features_[audio_feature_name]; + const uint8_t* row = feat.data() + audio_consumed * audio_row_bytes; + if (audio_consumed * audio_row_bytes + audio_row_bytes <= feat.size()) { + run_media_step(pos, row, audio_row_bytes, audio_feature_prec); + ++audio_consumed; + continue; + } + } + run_step(t, pos, false); + } + cache_total_seq_len_ += tokens.size(); + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); + (void)profile_file; +} + +uint32_t Model::decode(const std::vector& tokens, float temperature, float top_p, + size_t top_k, const std::string& /*profile_file*/, float* out_entropy, + float min_p, float repetition_penalty) { + if (tokens.empty()) return 0; + float temp = temperature < 0.0f ? config_.default_temperature : temperature; + float tp = (top_p <= 0.0f || top_p > 1.0f) ? config_.default_top_p : top_p; + size_t tk = top_k == 0 ? config_.default_top_k : top_k; + const bool greedy = temp <= 0.011f; + prepare_sampling_context(repetition_penalty); + struct SampClearGuard { + Model* m; + ~SampClearGuard() { cactus_graph_clear_sampling(); m->samp_ctx_active_ = false; } + } samp_guard{this}; + if (decode_route_ == DecodeRoute::ENCODER_CROSS_KV_STEP) { + if (!encoder_cross_kv_ready_ && encoder_cross_kv_source_kind_ == "text_tokens") { + std::vector source_tokens = tokens; + std::vector decoder_seed = {config_.decoder_start_token_id}; + if (!source_tokens.empty() && source_tokens.back() == config_.decoder_start_token_id) { + decoder_seed = {source_tokens.back()}; + source_tokens.pop_back(); + } + if (!prepare_encoder_cross_kv_from_text(source_tokens)) return 0; + std::vector emitted = run_encoder_cross_kv_decode_loop( + decoder_seed, + 1, + {}, + nullptr); + if (out_entropy) *out_entropy = 0.0f; + return emitted.empty() ? 0 : emitted.front(); + } + for (size_t i = 0; i < tokens.size(); ++i) { + if (!run_encoder_cross_kv_decoder_step(tokens[i], cache_total_seq_len_ + i)) return 0; + } + cache_total_seq_len_ += tokens.size(); + if (out_entropy) *out_entropy = 0.0f; + uint32_t result = argmax_last_logits(); + record_sampled_token(result); + return result; + } + if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { + context_tokens_.insert(context_tokens_.end(), tokens.begin(), tokens.end()); + run_full_context_text(); + cache_total_seq_len_ = context_tokens_.size(); + cache_token_ids_ = context_tokens_; + uint32_t result = (greedy && !samp_ctx_active_) + ? argmax_last_logits(out_entropy) + : sample_component_logits(*decoder_, temp, tp, tk, min_p, greedy, out_entropy); + record_sampled_token(result); + return result; + } + for (size_t i = 0; i + 1 < tokens.size(); ++i) { + run_step(tokens[i], cache_total_seq_len_ + i, /*read_logits=*/false); + } + run_step(tokens.back(), cache_total_seq_len_ + tokens.size() - 1, /*read_logits=*/true); + cache_total_seq_len_ += tokens.size(); + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); + maybe_roll_compact(); + uint32_t result = (greedy && !samp_ctx_active_) + ? argmax_last_logits(out_entropy) + : sample_component_logits(*decoder_, temp, tp, tk, min_p, greedy, out_entropy); + record_sampled_token(result); + return result; +} + +uint32_t Model::decode_with_audio(const std::vector& tokens, + const std::vector>& /*audio_features_per_message*/, + float temperature, float top_p, size_t top_k, const std::string& profile_file, + float* out_entropy, float min_p, float repetition_penalty, + float* /*out_token_time_start*/, float* /*out_token_time_end*/) { + return decode(tokens, temperature, top_p, top_k, profile_file, out_entropy, min_p, repetition_penalty); +} + +std::vector Model::transcribe_whisper_seq2seq( + const std::vector& audio_features, + const std::vector& decoder_prompt_tokens, + size_t max_tokens, + const std::vector>& stop_token_sequences, + const std::atomic* should_stop, + int64_t suppress_token_id) { + if (decoder_prompt_tokens.empty() || max_tokens == 0) return {}; + if (decode_route_ != DecodeRoute::ENCODER_CROSS_KV_STEP || encoder_cross_kv_source_kind_ != "audio_features") { + CACTUS_LOG_ERROR("model", "Whisper bundle missing encoder_cross_kv_decoder_step route metadata"); + return {}; + } + if (!prepare_encoder_cross_kv_from_audio(audio_features)) return {}; + struct SuppressGuard { int64_t& id; ~SuppressGuard() { id = -1; } } guard{suppressed_token_id_}; + suppressed_token_id_ = suppress_token_id; + return run_encoder_cross_kv_decode_loop( + decoder_prompt_tokens, + max_tokens, + stop_token_sequences, + should_stop); +} + +std::vector Model::transcribe_parakeet_tdt(const std::vector& audio_features, + ParakeetTdtStreamState* stream, bool is_final, + size_t end_frame, + const std::atomic* should_stop) { + std::vector emitted; + double raw_decode_ms = 0.0; + + reset_handoff_probe_rollout(); + + Component* audio_enc = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; + Component* dec = components_.count("decoder") ? &components_.at("decoder") : nullptr; + if (!audio_enc || !dec) { + CACTUS_LOG_ERROR("model", "Parakeet TDT bundle missing audio_encoder or decoder component"); + return emitted; + } + if (!bind_runtime_buffers(*audio_enc)) return emitted; + if (!bind_runtime_buffers(*dec)) return emitted; + + int feat_idx = input_index(*audio_enc, "input_features"); + if (feat_idx < 0) { + CACTUS_LOG_ERROR("model", "audio_encoder has no input_features input"); + return emitted; + } + auto& feat_buf = audio_enc->input_buffers[feat_idx]; + size_t feat_node = static_cast(audio_enc->runtime_input_node_ids[feat_idx]); + const auto& feat_desc = audio_enc->graph->get_output_buffer(feat_node); + if (feat_desc.shape.size() != 3) { + CACTUS_LOG_ERROR("model", "audio_encoder expects [1, frames, mels] input shape"); + return emitted; + } + const size_t expected_frames = feat_desc.shape[1]; + const size_t expected_mels = feat_desc.shape[2]; + const size_t source_frames = expected_mels > 0 ? audio_features.size() / expected_mels : 0; + const size_t copy_frames = std::min(source_frames, expected_frames); + std::vector transposed(expected_frames * expected_mels, 0.0f); + for (size_t t = 0; t < copy_frames; ++t) { + for (size_t m = 0; m < expected_mels; ++m) { + transposed[t * expected_mels + m] = audio_features[m * source_frames + t]; + } + } + write_typed_buffer(feat_buf, feat_desc.precision, transposed.data(), + transposed.size() * sizeof(float), Precision::FP32); + + if (should_stop && should_stop->load()) return emitted; + audio_enc->graph->execute(); + maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states"); + + int hidden_idx = output_index(*audio_enc, "encoder_hidden_states"); + if (hidden_idx < 0) { + CACTUS_LOG_ERROR("model", "audio_encoder has no encoder_hidden_states output"); + return emitted; + } + size_t hidden_node = static_cast(audio_enc->output_node_ids[hidden_idx]); + const auto& hidden_desc = audio_enc->graph->get_output_buffer(hidden_node); + const uint8_t* hidden_ptr = static_cast(audio_enc->graph->get_output(hidden_node)); + if (hidden_desc.shape.size() < 3 || hidden_ptr == nullptr) { + CACTUS_LOG_ERROR("model", "encoder_hidden_states must be 3D [B, T, D]"); + return emitted; + } + const size_t T = hidden_desc.shape[1]; + const size_t D = hidden_desc.shape[2]; + const Precision hidden_precision = hidden_desc.precision; + const size_t hidden_elem = PrecisionTraits::size_of(hidden_precision); + const size_t frame_bytes = D * hidden_elem; + + auto zero_state = [&](const std::string& name) { + int idx = input_index(*dec, name); + if (idx < 0) return; + auto& buf = dec->input_buffers[idx]; + std::memset(buf.data(), 0, buf.size()); + }; + if (stream && stream->initialized && stream->dec_state.size() == 4) { + const char* sn[4] = {"state_h_0", "state_c_0", "state_h_1", "state_c_1"}; + for (size_t i = 0; i < 4; ++i) { + int idx = input_index(*dec, sn[i]); + if (idx < 0) continue; + auto& buf = dec->input_buffers[idx]; + std::memcpy(buf.data(), stream->dec_state[i].data(), + std::min(buf.size(), stream->dec_state[i].size())); + } + } else { + zero_state("state_h_0"); + zero_state("state_c_0"); + zero_state("state_h_1"); + zero_state("state_c_1"); + } + + std::vector durations = config_.tdt_durations; + if (durations.empty()) { + for (uint32_t i = 0; i < config_.tdt_num_durations; ++i) durations.push_back(i); + } + if (durations.empty()) durations.push_back(1); + + const uint32_t configured_blank = config_.tdt_blank_id; + uint32_t last_token = (stream && stream->initialized) ? stream->last_token : configured_blank; + size_t time_index = (stream && stream->initialized) ? stream->time_index : 0; + + const int ef_idx = input_index(*dec, "encoder_frame"); + const int tok_in_idx = input_index(*dec, "token_ids"); + const int logits_idx = output_index(*dec, "step_logits"); + if (ef_idx < 0 || tok_in_idx < 0 || logits_idx < 0) { + CACTUS_LOG_ERROR("model", "decoder missing encoder_frame / token_ids / step_logits ports"); + return emitted; + } + auto& ef_buf = dec->input_buffers[ef_idx]; + const auto& ef_desc = dec->graph->get_output_buffer(static_cast(dec->runtime_input_node_ids[ef_idx])); + auto& tok_buf = dec->input_buffers[tok_in_idx]; + const Precision tok_prec = dec->graph->get_output_buffer(static_cast(dec->runtime_input_node_ids[tok_in_idx])).precision; + void* tok_data = tok_buf.data(); + const size_t logits_node = static_cast(dec->output_node_ids[logits_idx]); + const auto& logits_desc = dec->graph->get_output_buffer(logits_node); + const Precision logits_prec = logits_desc.precision; + const size_t total_classes = logits_desc.shape.empty() ? 0 : logits_desc.shape.back(); + const size_t num_durations = durations.size(); + const size_t token_class_count = (total_classes > num_durations) ? (total_classes - num_durations) : total_classes; + if (token_class_count == 0) return emitted; + uint32_t effective_blank = configured_blank; + if (effective_blank >= token_class_count) effective_blank = static_cast(token_class_count - 1); + + const std::array state_names = {"state_h_0", "state_c_0", "state_h_1", "state_c_1"}; + struct StateCopy { void* in_data; const void* out_ptr; size_t bytes; }; + std::array state_copies{}; + size_t state_copy_count = 0; + for (const char* state_name : state_names) { + int out_idx = output_index(*dec, state_name); + int in_idx = input_index(*dec, state_name); + if (out_idx < 0 || in_idx < 0) continue; + size_t out_node = static_cast(dec->output_node_ids[out_idx]); + const auto& out_desc = dec->graph->get_output_buffer(out_node); + auto& in_buf = dec->input_buffers[in_idx]; + state_copies[state_copy_count++] = { + in_buf.data(), + dec->graph->get_output(out_node), + std::min(out_desc.byte_size, in_buf.size()) + }; + } + + size_t commit_to = T; + if (stream) { + size_t valid_hidden = T; + if (expected_frames > 0) + valid_hidden = std::min(T, (copy_frames * T) / expected_frames); + commit_to = (end_frame > 0) ? std::min(end_frame, valid_hidden) : valid_hidden; + } + Tokenizer* stream_tok = stream ? get_tokenizer() : nullptr; + const auto& tdt_vocab_bias = get_vocab_bias(); + const float frame_sec = (160.0f / 16000.0f) * + static_cast(std::max(1, config_.subsampling_factor)); + constexpr uint32_t kMaxStreamDurationSkipFrames = 2; + + auto snapshot_state = [&]() { + std::vector> snap(state_copy_count); + for (size_t s = 0; s < state_copy_count; ++s) { + const uint8_t* p = static_cast(state_copies[s].in_data); + snap[s].assign(p, p + state_copies[s].bytes); + } + return snap; + }; + + std::vector> snap_state; + uint32_t snap_last_token = last_token; + size_t snap_time = time_index; + size_t confirmed_count = 0; + if (stream) snap_state = snapshot_state(); + + while (time_index < commit_to) { + if (should_stop && should_stop->load()) break; + const uint8_t* frame_ptr = hidden_ptr + time_index * frame_bytes; + write_typed_buffer(ef_buf, ef_desc.precision, frame_ptr, frame_bytes, hidden_precision); + + size_t symbols_added = 0; + bool advanced = false; + while (symbols_added < 10) { + switch (tok_prec) { + case Precision::FP32: *reinterpret_cast(tok_data) = static_cast(last_token); break; + case Precision::FP16: *reinterpret_cast<__fp16*>(tok_data) = static_cast<__fp16>(last_token); break; + case Precision::INT8: *reinterpret_cast(tok_data) = static_cast(last_token); break; + default: *reinterpret_cast(tok_data) = static_cast(last_token); break; + } + const auto exec_t0 = std::chrono::steady_clock::now(); + dec->graph->execute(); + raw_decode_ms += std::chrono::duration( + std::chrono::steady_clock::now() - exec_t0).count(); + + const void* logits_ptr = dec->graph->get_output(logits_node); + auto get_logit = [&](size_t i) -> float { + if (logits_prec == Precision::FP32) return reinterpret_cast(logits_ptr)[i]; + if (logits_prec == Precision::FP16) return static_cast(reinterpret_cast(logits_ptr)[i]); + return static_cast(reinterpret_cast(logits_ptr)[i]); + }; + + uint32_t next_token = 0; + float best_token_score = -std::numeric_limits::infinity(); + for (size_t i = 0; i < token_class_count; ++i) { + float v = get_logit(i); + if (stream && !tdt_vocab_bias.empty()) { + auto it = tdt_vocab_bias.find(static_cast(i)); + if (it != tdt_vocab_bias.end()) v += it->second; + } + if (v > best_token_score) { best_token_score = v; next_token = static_cast(i); } + } + uint32_t best_duration_idx = 0; + float best_duration_score = -std::numeric_limits::infinity(); + for (size_t i = 0; i < num_durations; ++i) { + float v = get_logit(token_class_count + i); + if (v > best_duration_score) { best_duration_score = v; best_duration_idx = static_cast(i); } + } + + uint32_t skip = durations[std::min(best_duration_idx, static_cast(durations.size() - 1))]; + if (stream && skip > kMaxStreamDurationSkipFrames) skip = kMaxStreamDurationSkipFrames; + + if (next_token != effective_blank) { + if (stream && !is_final && stream_tok && !emitted.empty()) { + std::string piece = stream_tok->decode({next_token}); + if (!piece.empty() && piece[0] == ' ') { + snap_state = snapshot_state(); + snap_last_token = last_token; + snap_time = time_index; + confirmed_count = emitted.size(); + } + } + emitted.push_back(next_token); + last_token = next_token; + for (size_t s = 0; s < state_copy_count; ++s) { + std::memcpy(state_copies[s].in_data, state_copies[s].out_ptr, state_copies[s].bytes); + } + } + + ++symbols_added; + + if (skip > 0) { + time_index += skip; + advanced = true; + break; + } + if (next_token == effective_blank) { + time_index += 1; + advanced = true; + break; + } + } + + if (!advanced) time_index += 1; + } + + if (stream) { + stream->initialized = true; + stream->decoded_tokens = emitted.size(); + stream->raw_decode_ms = raw_decode_ms; + stream->pending.clear(); + if (is_final) { + stream->dec_state = snapshot_state(); + stream->last_token = last_token; + stream->time_index = time_index; + stream->confirmed_sec = static_cast(time_index) * frame_sec; + } else { + stream->dec_state = std::move(snap_state); + stream->last_token = snap_last_token; + stream->time_index = snap_time; + stream->confirmed_sec = static_cast(snap_time) * frame_sec; + stream->pending.assign(emitted.begin() + confirmed_count, emitted.end()); + emitted.resize(confirmed_count); + } + } + + return emitted; +} + +uint32_t Model::decode_with_images(const std::vector& tokens, const std::vector& /*image_paths*/, + float temperature, float top_p, size_t top_k, const std::string& profile_file, + float* out_entropy, float min_p, float repetition_penalty) { + return decode(tokens, temperature, top_p, top_k, profile_file, out_entropy, min_p, repetition_penalty); +} + +namespace { + +std::vector pool_and_normalize_media_feature( + const std::vector& bytes, + const std::vector& shape, + Precision precision, + const std::string& source +) { + const size_t elem_size = PrecisionTraits::size_of(precision); + if (elem_size == 0 || bytes.empty() || shape.empty()) { + throw std::runtime_error(source + " produced empty feature output"); + } + const size_t total_elems = bytes.size() / elem_size; + const size_t hidden_dim = shape.back(); + if (hidden_dim == 0 || total_elems == 0 || total_elems % hidden_dim != 0) { + throw std::runtime_error(source + " feature shape inconsistent with hidden_dim"); + } + + std::vector fp32(total_elems); + switch (precision) { + case Precision::FP32: + std::memcpy(fp32.data(), bytes.data(), total_elems * sizeof(float)); + break; + case Precision::FP16: + Quantization::fp16_to_fp32(reinterpret_cast(bytes.data()), fp32.data(), total_elems); + break; + case Precision::INT8: + Quantization::int8_to_fp32(reinterpret_cast(bytes.data()), fp32.data(), total_elems, 1.0f); + break; + default: + throw std::runtime_error(source + " feature precision not supported for embeddings"); + } + + const size_t n_rows = total_elems / hidden_dim; + std::vector pooled(hidden_dim, 0.0f); + for (size_t r = 0; r < n_rows; ++r) { + const float* src = fp32.data() + r * hidden_dim; + for (size_t d = 0; d < hidden_dim; ++d) pooled[d] += src[d]; + } + const float inv = 1.0f / static_cast(n_rows); + for (float& v : pooled) v *= inv; + + float norm_sq = 0.0f; + for (float v : pooled) norm_sq += v * v; + if (norm_sq > 1e-12f) { + const float inv_norm = 1.0f / std::sqrt(norm_sq); + for (float& v : pooled) v *= inv_norm; + } + return pooled; +} + +} // namespace + +std::vector Model::get_image_embeddings(const std::string& image_path) { + if (!vision_encoder_) { + throw std::runtime_error("Model has no vision_encoder component"); + } + if (vision_encoder_->logical_outputs.empty()) { + throw std::runtime_error("vision_encoder has no logical outputs"); + } + std::string output_name = vision_encoder_->logical_outputs[0]; + + run_vision_encoder(image_path); + + if (!media_features_.count(output_name)) { + for (const char* name : {"image_features", "image_embeddings", "vision_features"}) { + if (media_features_.count(name)) { output_name = name; break; } + } + } + auto bytes_it = media_features_.find(output_name); + auto shape_it = media_feature_shapes_.find(output_name); + auto prec_it = media_feature_precisions_.find(output_name); + if (bytes_it == media_features_.end() || shape_it == media_feature_shapes_.end() + || prec_it == media_feature_precisions_.end()) { + throw std::runtime_error("vision_encoder produced no output for '" + output_name + "'"); + } + + std::vector embedding = pool_and_normalize_media_feature( + bytes_it->second, shape_it->second, prec_it->second, "vision_encoder"); + + for (const std::string& name : vision_encoder_->logical_outputs) { + media_features_.erase(name); + media_feature_shapes_.erase(name); + media_feature_precisions_.erase(name); + } + media_features_.erase(output_name); + media_feature_shapes_.erase(output_name); + media_feature_precisions_.erase(output_name); + // run_vision_encoder unloads the graph; restore so subsequent paths that + // assume the encoder is loaded (e.g. transcribe_*) keep working. + load_component_graph(*vision_encoder_); + return embedding; +} + +std::vector Model::get_audio_embeddings(const std::vector& mel_bins) { + if (!audio_encoder_) { + throw std::runtime_error("Model has no audio_encoder component"); + } + if (mel_bins.empty()) { + throw std::runtime_error("Empty audio features"); + } + if (audio_encoder_->logical_outputs.empty()) { + throw std::runtime_error("audio_encoder has no logical outputs"); + } + const std::string output_name = audio_encoder_->logical_outputs[0]; + + run_audio_encoder_messages({mel_bins}); + + auto bytes_it = media_features_.find(output_name); + auto shape_it = media_feature_shapes_.find(output_name); + auto prec_it = media_feature_precisions_.find(output_name); + if (bytes_it == media_features_.end() || shape_it == media_feature_shapes_.end() + || prec_it == media_feature_precisions_.end()) { + throw std::runtime_error("audio_encoder produced no output for '" + output_name + "'"); + } + + std::vector embedding = pool_and_normalize_media_feature( + bytes_it->second, shape_it->second, prec_it->second, "audio_encoder"); + + for (const std::string& name : audio_encoder_->logical_outputs) { + media_features_.erase(name); + media_feature_shapes_.erase(name); + media_feature_precisions_.erase(name); + } + load_component_graph(*audio_encoder_); + return embedding; +} + +void Model::reset_cache() { + cache_total_seq_len_ = 0; + last_logit_position_ = 0; + encoder_cross_kv_ready_ = false; + context_tokens_.clear(); + cache_token_ids_.clear(); + special_rows_.clear(); + token_history_.clear(); + media_features_.clear(); + media_feature_shapes_.clear(); + media_feature_precisions_.clear(); + for (auto& kv : components_) { + Component& comp = kv.second; + if (!comp.graph) continue; + reset_component_cache_states(comp); + } +} + +void Model::set_cache_window(size_t /*window_size*/, size_t /*sink_size*/) {} + +void Model::apply_kv_compress_env_override() { + config_.parse_kv_compress_override(std::getenv("CACTUS_KV_COMPRESS_AT"), + std::getenv("CACTUS_KV_COMPRESS_TO")); +} + +std::vector Model::compressible_layers() const { + size_t shared = (config_.num_kv_shared_layers == Config::UNSET_U32) + ? 0 : config_.num_kv_shared_layers; + return cactus::kvcompress::physical_compressible_layers( + config_.layer_types, config_.num_layers, shared); +} + +void Model::compress_kv_cache_keydiff(const cactus::kvcompress::Params& params) { + if (!decoder_) return; + using cactus::kvcompress::CacheHeader; + constexpr size_t kHeaderBytes = sizeof(CacheHeader); + + std::vector layers = compressible_layers(); + if (layers.empty()) return; + std::set compressible(layers.begin(), layers.end()); + const double rope_theta = static_cast(config_.rope_theta); + const double rope_local_theta = (config_.rope_local_base_freq == Config::UNSET_F32) + ? rope_theta : static_cast(config_.rope_local_base_freq); + const size_t old_total = cache_total_seq_len_; + + cactus::kvcompress::Params params_local = params; + bool preserve = config_.kv_compress_preserve_special; + if (const char* e = std::getenv("CACTUS_KV_PRESERVE_SPECIAL")) preserve = (std::atoi(e) != 0); + const bool map_valid = cache_token_ids_.size() == old_total && media_features_.empty(); + const bool per_head_protect = preserve && map_valid && special_rows_.valid(); + // cache_token_ids_ past tracked_len is still head-aligned, so specials there apply to every head. + std::vector appended_special; + if (per_head_protect) { + if (special_ids_.empty() && tokenizer_) special_ids_ = tokenizer_->special_token_ids(); + for (size_t r = special_rows_.tracked_len(); r < old_total && r < cache_token_ids_.size(); ++r) + if (special_ids_.count(cache_token_ids_[r])) appended_special.push_back(static_cast(r)); + } + + Component& comp = *decoder_; + if (!comp.graph) return; + // Skip the whole pass (before mutating) if any layer's V dim differs from K (MLA), or a head + // can't fit sink + all its specials in the budget. + const size_t protect_budget = params.abs_budget > 0 ? static_cast(params.abs_budget) : 0; + for (size_t li = 0; li < comp.cache_states.size(); ++li) { + if (!compressible.count(li)) continue; + const auto& cs = comp.cache_states[li]; + if (cs.key_node_id < 0 || cs.value_node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(cs.key_node_id)) != OpType::KV_CACHE_STATE) continue; + void* kraw = comp.graph->get_output(static_cast(cs.key_node_id)); + void* vraw = comp.graph->get_output(static_cast(cs.value_node_id)); + if (!kraw || !vraw) continue; + if (static_cast(vraw)->head_dim != static_cast(kraw)->head_dim) return; + if (per_head_protect && protect_budget > 0 && + special_rows_.max_reserved(li, params.sink, appended_special) > protect_budget) return; + } + size_t new_seq_len = 0; + bool have_new_seq_len = false; + std::vector canonical_keep; + bool canonical_captured = false; + std::vector unrope; + size_t shrink_cap = 1; + while (shrink_cap < static_cast(config_.kv_compress_trigger_len)) shrink_cap <<= 1; + for (size_t li = 0; li < comp.cache_states.size(); ++li) { + if (!compressible.count(li)) continue; + const auto& cs = comp.cache_states[li]; + if (cs.key_node_id < 0 || cs.value_node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(cs.key_node_id)) != OpType::KV_CACHE_STATE) continue; + + const auto& kdesc = comp.graph->get_output_buffer(static_cast(cs.key_node_id)); + const auto& vdesc = comp.graph->get_output_buffer(static_cast(cs.value_node_id)); + if (kdesc.byte_size <= kHeaderBytes || vdesc.byte_size <= kHeaderBytes) continue; + void* kraw = comp.graph->get_output(static_cast(cs.key_node_id)); + void* vraw = comp.graph->get_output(static_cast(cs.value_node_id)); + if (!kraw || !vraw) continue; + + auto* khdr = static_cast(kraw); + auto* vhdr = static_cast(vraw); + size_t n = khdr->current_seq_len; + size_t kv_heads = khdr->num_kv_heads; + size_t head_dim = khdr->head_dim; + if (kv_heads == 0 || head_dim == 0) continue; + if (unrope.empty()) unrope = cactus::kvcompress::unrope_table(n, head_dim, rope_theta); + + static const std::vector> kNoProtect; + if (per_head_protect) special_rows_.add_appended(li, kv_heads, appended_special); + const std::vector>& pph = per_head_protect ? special_rows_.protect(li) : kNoProtect; + + if (kdesc.precision == Precision::FP16) { + auto* kbase = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + auto* vbase = reinterpret_cast(static_cast(vraw) + kHeaderBytes); + auto kept = cactus::kvcompress::keepsets_from_fp16( + kbase, n, kv_heads, head_dim, unrope, params_local, pph); + if (!canonical_captured) { canonical_keep = kept.empty() ? std::vector{} : kept[0]; canonical_captured = true; } + cactus::kvcompress::compact_fp16(kbase, vbase, kv_heads, head_dim, kept, unrope); + if (per_head_protect) special_rows_.remap(li, kept); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; + vhdr->current_seq_len = B; + new_seq_len = B; + have_new_seq_len = true; + } else if (kdesc.precision == Precision::INT8) { + size_t max_seq = khdr->max_seq_len; + auto* k_i8 = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + auto* k_sc = reinterpret_cast(static_cast(kraw) + kHeaderBytes + + max_seq * kv_heads * head_dim); + auto* v_i8 = reinterpret_cast(static_cast(vraw) + kHeaderBytes); + auto* v_sc = reinterpret_cast(static_cast(vraw) + kHeaderBytes + + max_seq * kv_heads * head_dim); + auto kept = cactus::kvcompress::keepsets_from_int8( + k_i8, k_sc, n, kv_heads, head_dim, KV_QUANT_GROUP_SIZE, unrope, params_local, pph); + if (!canonical_captured) { canonical_keep = kept.empty() ? std::vector{} : kept[0]; canonical_captured = true; } + cactus::kvcompress::compact_int8(k_i8, k_sc, kv_heads, head_dim, KV_QUANT_GROUP_SIZE, + kept, unrope, /*renumber=*/true); + cactus::kvcompress::compact_int8(v_i8, v_sc, kv_heads, head_dim, KV_QUANT_GROUP_SIZE, + kept, unrope, /*renumber=*/false); + if (per_head_protect) special_rows_.remap(li, kept); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; + vhdr->current_seq_len = B; + new_seq_len = B; + have_new_seq_len = true; + } + comp.graph->shrink_cache_buffer(static_cast(cs.key_node_id), shrink_cap); + comp.graph->shrink_cache_buffer(static_cast(cs.value_node_id), shrink_cap); + } + + const size_t Delta = (have_new_seq_len && old_total >= new_seq_len) ? old_total - new_seq_len : 0; + if (Delta > 0) { + const double dpos = -static_cast(Delta); + for (size_t li = 0; li < comp.cache_states.size(); ++li) { + if (compressible.count(li)) continue; + const auto& cs = comp.cache_states[li]; + if (cs.key_node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(cs.key_node_id)) != OpType::KV_CACHE_STATE) continue; + const auto& kdesc = comp.graph->get_output_buffer(static_cast(cs.key_node_id)); + if (kdesc.byte_size <= kHeaderBytes) continue; + void* kraw = comp.graph->get_output(static_cast(cs.key_node_id)); + if (!kraw) continue; + auto* khdr = static_cast(kraw); + size_t kv_heads = khdr->num_kv_heads, head_dim = khdr->head_dim; + if (kv_heads == 0 || head_dim == 0) continue; + size_t hi = khdr->current_seq_len; + size_t lo = std::min(khdr->sink_size, hi); + const double layer_theta = cactus::kvcompress::is_sliding_layer(config_.layer_types, li) + ? rope_local_theta : rope_theta; + if (kdesc.precision == Precision::FP16) { + auto* kbase = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + cactus::kvcompress::rerope_recent_fp16(kbase, kv_heads, head_dim, lo, hi, + layer_theta, dpos); + } else if (kdesc.precision == Precision::INT8) { + size_t max_seq = khdr->max_seq_len; + auto* k_i8 = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + auto* k_sc = reinterpret_cast(static_cast(kraw) + kHeaderBytes + + max_seq * kv_heads * head_dim); + cactus::kvcompress::rerope_recent_int8(k_i8, k_sc, kv_heads, head_dim, + KV_QUANT_GROUP_SIZE, lo, hi, layer_theta, dpos); + } + } + } + + if (have_new_seq_len) { + cache_total_seq_len_ = new_seq_len; + if (per_head_protect) special_rows_.set_tracked_len(new_seq_len); + else special_rows_.invalidate(); + if (map_valid && canonical_captured) { + std::vector compacted; + compacted.reserve(canonical_keep.size()); + for (int idx : canonical_keep) + if (idx >= 0 && idx < static_cast(cache_token_ids_.size())) compacted.push_back(cache_token_ids_[idx]); + cache_token_ids_ = std::move(compacted); + } else { + cache_token_ids_.clear(); + } + } +} + +void Model::maybe_roll_compact() { + if (!config_.kv_compress || config_.kv_compress_trigger_len <= 0) return; + if (cache_total_seq_len_ < static_cast(config_.kv_compress_trigger_len)) return; + + cactus::kvcompress::Params p; + p.recent_frac = config_.kv_compress_recent_frac; + p.sink = config_.kv_compress_sink; + p.abs_budget = config_.kv_compress_target_len; + compress_kv_cache_keydiff(p); +} + +static void l2_normalize_inplace(std::vector& v) { + double norm = 0.0; + for (float x : v) norm += static_cast(x) * x; + float inv = static_cast(1.0 / std::max(std::sqrt(norm), 1e-12)); + for (float& x : v) x *= inv; +} + +static std::vector finalize_pooled_embedding(const std::vector& sum, size_t count, bool normalize) { + std::vector out(sum.size(), 0.0f); + if (count > 0) { + for (size_t h = 0; h < sum.size(); ++h) out[h] = static_cast(sum[h] / static_cast(count)); + } + if (normalize) l2_normalize_inplace(out); + return out; +} + +std::vector Model::get_text_embeddings(const std::vector& tokens, bool normalize) { + if (has_text_embedding()) { + return get_embeddings(tokens, /*pooled=*/true, normalize); + } + if (has_lm_embedding()) { + return get_lm_embeddings(tokens, normalize); + } + throw std::runtime_error("get_text_embeddings: bundle has neither a text_embedding nor a decoder_embed_chunk component"); +} + +std::vector Model::get_lm_embeddings(const std::vector& tokens, bool normalize) { + if (!decoder_embed_) { + throw std::runtime_error("get_lm_embeddings: bundle has no decoder_embed_chunk component"); + } + if (tokens.empty()) { + throw std::runtime_error("get_lm_embeddings: empty token sequence"); + } + if (decode_route_ != DecodeRoute::CACHED_STEP || !encoder_) { + throw std::runtime_error("get_lm_embeddings: model does not support chunked LM embeddings"); + } + if (!load_component_graph(*decoder_embed_)) { + throw std::runtime_error("get_lm_embeddings: failed to load decoder_embed_chunk graph"); + } + if (prefill_encoder_ && !load_component_graph(*prefill_encoder_)) { + throw std::runtime_error("get_lm_embeddings: failed to load prefill encoder graph"); + } + reset_component_cache_states(*decoder_embed_); + + const size_t component_tokens = component_chunk_tokens(*decoder_embed_, "inputs_embeds"); + if (component_tokens <= 1) { + throw std::runtime_error("get_lm_embeddings: decoder_embed_chunk is not chunk-shaped"); + } + const size_t effective_chunk = component_tokens; + + bool recurrent_state = false; + if (decoder_embed_->graph) { + for (const auto& state : decoder_embed_->cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + if (decoder_embed_->graph->get_node_op_type(static_cast(node_id)) == OpType::RECURRENT_CACHE_STATE) { + recurrent_state = true; + } + } + } + } + const size_t embed_limit = recurrent_state ? std::min(tokens.size(), effective_chunk) : tokens.size(); + + size_t encoder_chunk = 0; + if (prefill_encoder_ && input_index(*prefill_encoder_, "input_ids") >= 0 && input_index(*prefill_encoder_, "position_ids") >= 0) { + encoder_chunk = component_chunk_tokens(*prefill_encoder_, "input_ids"); + if (encoder_chunk == 0 || effective_chunk % encoder_chunk != 0) encoder_chunk = 0; + } + + const int out_idx = output_index(*decoder_embed_, "last_hidden_state"); + if (out_idx < 0 || static_cast(out_idx) >= decoder_embed_->output_node_ids.size()) { + throw std::runtime_error("get_lm_embeddings: decoder_embed_chunk missing last_hidden_state output"); + } + const size_t out_node = static_cast(decoder_embed_->output_node_ids[out_idx]); + + std::vector sum; + size_t count = 0; + + size_t processed = 0; + while (processed < embed_limit) { + for (auto& buf : decoder_embed_->input_buffers) std::fill(buf.begin(), buf.end(), 0); + if (encoder_chunk > 0) { + for (size_t chunk_offset = 0; chunk_offset < effective_chunk; chunk_offset += encoder_chunk) { + for (auto& buf : prefill_encoder_->input_buffers) std::fill(buf.begin(), buf.end(), 0); + for (size_t i = 0; i < encoder_chunk; ++i) { + const size_t index = processed + chunk_offset + i; + const uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); + write_int_input_at(*prefill_encoder_, "input_ids", i, static_cast(token)); + write_int_input_at(*prefill_encoder_, "position_ids", i, static_cast(processed + chunk_offset + i)); + } + prefill_encoder_->graph->execute(); + copy_component_outputs_to_chunk_inputs_range(*prefill_encoder_, *decoder_embed_, chunk_offset); + } + } else { + for (size_t i = 0; i < effective_chunk; ++i) { + const size_t index = processed + i; + const uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); + run_encoder_step(token, processed + i); + copy_component_outputs_to_chunk_inputs(*encoder_, *decoder_embed_, i); + } + } + decoder_embed_->graph->execute(); + + const auto& desc = decoder_embed_->graph->get_output_buffer(out_node); + void* ptr = decoder_embed_->graph->get_output(out_node); + const size_t hidden = desc.shape.empty() ? 0 : desc.shape.back(); + const size_t seq = (desc.shape.size() >= 2) ? desc.shape[desc.shape.size() - 2] : 1; + if (hidden == 0 || !ptr) { + throw std::runtime_error("get_lm_embeddings: decoder_embed_chunk produced no last_hidden_state"); + } + if (sum.empty()) sum.assign(hidden, 0.0); + const bool is_fp16 = desc.precision == Precision::FP16; + auto read_at = [&](size_t i) -> float { + return is_fp16 ? static_cast(reinterpret_cast(ptr)[i]) + : reinterpret_cast(ptr)[i]; + }; + size_t real_rows = std::min(effective_chunk, tokens.size() - processed); + real_rows = std::min(real_rows, seq); + for (size_t t = 0; t < real_rows; ++t) { + for (size_t h = 0; h < hidden; ++h) sum[h] += static_cast(read_at(t * hidden + h)); + } + count += real_rows; + processed += effective_chunk; + } + + reset_component_cache_states(*decoder_embed_); + decoder_embed_->graph->release_runtime_buffers(); + decoder_embed_->graph->release_all_weight_pages(); + unload_component_graph(*decoder_embed_); + return finalize_pooled_embedding(sum, count, normalize); +} + +std::vector Model::get_embeddings(const std::vector& tokens, bool pooled, + bool normalize, const std::string& /*profile_file*/) { + if (!components_.count("text_embedding")) { + throw std::runtime_error("get_embeddings: bundle has no text_embedding component"); + } + if (tokens.empty()) { + throw std::runtime_error("get_embeddings: empty token sequence"); + } + Component* comp = &components_.at("text_embedding"); + if (!load_component_graph(*comp)) { + throw std::runtime_error("get_embeddings: failed to load embedding component graph"); + } + + std::vector wrapped; + wrapped.reserve(tokens.size() + 2); + if (tokenizer_) wrapped.push_back(tokenizer_->get_bos_token()); + wrapped.insert(wrapped.end(), tokens.begin(), tokens.end()); + if (tokenizer_) wrapped.push_back(tokenizer_->get_eos_token()); + + int ids_idx = input_index(*comp, "input_ids"); + if (ids_idx < 0) { + throw std::runtime_error("get_embeddings: embedding component missing input_ids"); + } + auto& ids_buf = comp->input_buffers[ids_idx]; + size_t ids_node = static_cast(comp->runtime_input_node_ids[ids_idx]); + const auto& ids_desc = comp->graph->get_output_buffer(ids_node); + size_t capacity = PrecisionTraits::size_of(ids_desc.precision) + ? ids_buf.size() / PrecisionTraits::size_of(ids_desc.precision) : wrapped.size(); + size_t n_real = std::min(capacity, wrapped.size()); + write_tokens_buffer(ids_buf, ids_desc.precision, wrapped, 0); + + int mask_idx = input_index(*comp, "attention_mask"); + if (mask_idx >= 0) { + auto& mb = comp->input_buffers[mask_idx]; + size_t mnode = static_cast(comp->runtime_input_node_ids[mask_idx]); + const auto& mdesc = comp->graph->get_output_buffer(mnode); + fill_int_buffer(mb, mdesc.precision, 1, n_real); + } + + comp->graph->execute(); + + if (comp->output_node_ids.empty()) { + throw std::runtime_error("get_embeddings: embedding component produced no outputs"); + } + size_t out_node = static_cast(comp->output_node_ids[0]); + const auto& desc = comp->graph->get_output_buffer(out_node); + void* ptr = comp->graph->get_output(out_node); + size_t hidden = desc.shape.empty() ? 0 : desc.shape.back(); + size_t seq = (desc.shape.size() >= 2) ? desc.shape[desc.shape.size() - 2] : 1; + if (hidden == 0) { + throw std::runtime_error("get_embeddings: embedding output has zero hidden dim"); + } + + const bool is_fp16 = desc.precision == Precision::FP16; + auto read_at = [&](size_t i) -> float { + return is_fp16 ? static_cast(reinterpret_cast(ptr)[i]) + : reinterpret_cast(ptr)[i]; + }; + + std::vector result(hidden, 0.0f); + if (pooled) { + size_t pool_rows = std::min(seq, std::max(1, n_real)); + for (size_t t = 0; t < pool_rows; ++t) { + for (size_t h = 0; h < hidden; ++h) result[h] += read_at(t * hidden + h); + } + for (size_t h = 0; h < hidden; ++h) result[h] /= static_cast(pool_rows); + } else { + for (size_t h = 0; h < hidden; ++h) result[h] = read_at(h); + } + + if (normalize) l2_normalize_inplace(result); + + comp->graph->release_runtime_buffers(); + comp->graph->release_all_weight_pages(); + unload_component_graph(*comp); + return result; +} + +bool Config::from_json(const std::string& config_path) { + std::ifstream file(config_path); + if (!file) { + CACTUS_LOG_ERROR("config", "Failed to open config file: " << config_path); + return false; + } + + std::string line; + bool decoder_start_token_seen = false; + while (std::getline(file, line)) { + if (line.empty() || line[0] == '#') continue; + + size_t eq_pos = line.find('='); + if (eq_pos == std::string::npos) continue; + + std::string key = line.substr(0, eq_pos); + std::string value = line.substr(eq_pos + 1); + + key.erase(0, key.find_first_not_of(" \t")); + key.erase(key.find_last_not_of(" \t") + 1); + value.erase(0, value.find_first_not_of(" \t")); + value.erase(value.find_last_not_of(" \t") + 1); + + if (key == "vocab_size") vocab_size = static_cast(std::stoul(value)); + else if (key == "bos_token_id") bos_token_id = static_cast(std::stoul(value)); + else if (key == "eos_token_id") eos_token_id = static_cast(std::stoul(value)); + else if (key == "decoder_start_token_id") { + decoder_start_token_id = static_cast(std::stoul(value)); + decoder_start_token_seen = true; + } + else if (key == "decoder_prompt_token_ids") decoder_prompt_token_ids = parse_config_uint_list(value); + else if (key == "num_layers") num_layers = static_cast(std::stoul(value)); + else if (key == "hidden_dim") hidden_dim = static_cast(std::stoul(value)); + else if (key == "ffn_intermediate_dim") ffn_intermediate_dim = static_cast(std::stoul(value)); + else if (key == "attention_heads") attention_heads = static_cast(std::stoul(value)); + else if (key == "attention_kv_heads") attention_kv_heads = static_cast(std::stoul(value)); + else if (key == "attention_head_dim") attention_head_dim = static_cast(std::stoul(value)); + else if (key == "layer_norm_eps") layer_norm_eps = std::stof(value); + else if (key == "rope_theta") rope_theta = std::stof(value); + else if (key == "num_experts") num_experts = static_cast(std::stoul(value)); + else if (key == "num_shared_experts") num_shared_experts = static_cast(std::stoul(value)); + else if (key == "num_top_experts") num_top_experts = static_cast(std::stoul(value)); + else if (key == "moe_every_n_layers") moe_every_n_layers = static_cast(std::stoul(value)); + else if (key == "moe_intermediate_dim" || key == "moe_intermediate_size") moe_intermediate_dim = static_cast(std::stoul(value)); + else if (key == "num_dense_layers") num_dense_layers = static_cast(std::stoul(value)); + else if (key == "num_experts_per_tok") num_experts_per_tok = static_cast(std::stoul(value)); + else if (key == "norm_topk_prob") norm_topk_prob = (value == "true" || value == "1"); + else if (key == "use_expert_bias") use_expert_bias = (value == "true" || value == "1"); + else if (key == "routed_scaling_factor") routed_scaling_factor = std::stof(value); + else if (key == "tie_word_embeddings") tie_word_embeddings = (value == "true" || value == "1"); + else if (key == "vision_hidden_dim" || key == "vision_hidden_size") vision_hidden_dim = static_cast(std::stoul(value)); + else if (key == "vision_num_layers") vision_num_layers = static_cast(std::stoul(value)); + else if (key == "vision_attention_heads") vision_attention_heads = static_cast(std::stoul(value)); + else if (key == "vision_image_size") vision_image_size = static_cast(std::stoul(value)); + else if (key == "vision_patch_size") vision_patch_size = static_cast(std::stoul(value)); + else if (key == "vision_num_channels") vision_num_channels = static_cast(std::stoul(value)); + else if (key == "vision_embed_dim") vision_embed_dim = static_cast(std::stoul(value)); + else if (key == "visual_tokens_per_img") visual_tokens_per_img = static_cast(std::stoul(value)); + else if (key == "use_pixel_shuffle") use_pixel_shuffle = (value == "true" || value == "1"); + else if (key == "pixel_shuffle_factor") pixel_shuffle_factor = static_cast(std::stoul(value)); + else if (key == "use_image_tokens") use_image_tokens = (value == "true" || value == "1"); + else if (key == "image_token_id") image_token_id = static_cast(std::stoul(value)); + else if (key == "use_layout_tags") use_layout_tags = (value == "true" || value == "1"); + else if (key == "image_seq_len") image_seq_len = static_cast(std::stoul(value)); + else if (key == "global_image_size") global_image_size = static_cast(std::stoul(value)); + else if (key == "max_tile_size") max_tile_size = static_cast(std::stoul(value)); + else if (key == "rescale_factor") rescale_factor = std::stof(value); + else if (key == "image_mean") image_mean = std::stof(value); + else if (key == "image_std") image_std = std::stof(value); + else if (key == "downsample_factor") downsample_factor = static_cast(std::stoul(value)); + else if (key == "min_tiles") min_tiles = static_cast(std::stoul(value)); + else if (key == "max_tiles") max_tiles = static_cast(std::stoul(value)); + else if (key == "use_thumbnail") use_thumbnail = (value == "true" || value == "1"); + else if (key == "min_image_tokens") min_image_tokens = static_cast(std::stoul(value)); + else if (key == "max_image_tokens") max_image_tokens = static_cast(std::stoul(value)); + else if (key == "tile_size") tile_size = static_cast(std::stoul(value)); + else if (key == "max_pixels_tolerance") max_pixels_tolerance = std::stof(value); + else if (key == "do_image_splitting") do_image_splitting = (value == "true" || value == "1"); + else if (key == "precision") { + if (value == "INT8") precision = Precision::INT8; + else if (value == "FP16") precision = Precision::FP16; + else precision = Precision::FP32; + } + else if (key == "model_type") { + std::string mt = value; + std::transform(mt.begin(), mt.end(), mt.begin(), ::tolower); + if (mt == "qwen") model_type = ModelType::QWEN; + else if (mt == "qwen3p5" || mt == "qwen3_5") model_type = ModelType::QWEN3P5; + else if (mt == "gemma" || mt == "gemma3") model_type = ModelType::GEMMA; + else if (mt == "gemma3n") model_type = ModelType::GEMMA3N; + else if (mt == "lfm2") model_type = ModelType::LFM2; + else if (mt == "whisper") model_type = ModelType::WHISPER; + else if (mt == "parakeet_tdt" || mt == "parakeet-tdt") model_type = ModelType::PARAKEET_TDT; + else if (mt == "youtu") model_type = ModelType::YOUTU; + else if (mt == "needle") model_type = ModelType::NEEDLE; + else if (mt == "bert" || mt == "nomic") model_type = ModelType::NOMIC; + else model_type = ModelType::GEMMA4; + } + else if (key == "model_variant") { + std::string v = value; + std::transform(v.begin(), v.end(), v.begin(), ::tolower); + if (v == "vlm") model_variant = ModelVariant::VLM; + else if (v == "extract") model_variant = ModelVariant::EXTRACT; + else if (v == "rag") model_variant = ModelVariant::RAG; + else model_variant = ModelVariant::DEFAULT; + } + else if (key == "conv_L_cache") conv_L_cache = static_cast(std::stoul(value)); + else if (key == "kv_compress") kv_compress = (value == "true" || value == "1"); + else if (key == "kv_compress_recent_frac") kv_compress_recent_frac = std::stof(value); + else if (key == "kv_compress_sink") kv_compress_sink = static_cast(std::stoul(value)); + else if (key == "kv_compress_trigger_len") kv_compress_trigger_len = static_cast(std::stol(value)); + else if (key == "kv_compress_target_len") kv_compress_target_len = static_cast(std::stol(value)); + else if (key == "kv_compress_preserve_special") kv_compress_preserve_special = (value == "true" || value == "1"); + else if (key == "layer_types") { + layer_types.clear(); + std::string sanitized; + sanitized.reserve(value.size()); + for (char c : value) { + if (c == '[' || c == ']' || c == '\'' || c == '"') { + continue; + } + sanitized.push_back(c); + } + std::stringstream ss(sanitized); + std::string item; + while (std::getline(ss, item, ',')) { + if (!item.empty()) { + item.erase(0, item.find_first_not_of(" \t")); + item.erase(item.find_last_not_of(" \t") + 1); + if (!item.empty()) layer_types.push_back(item); + } + } + } + else if (key == "enc_hidden_act") encoder_act_gelu = (value == "gelu"); + else if (key == "dec_hidden_act") decoder_act_gelu = (value == "gelu"); + else if (key == "num_encoder_layers") num_encoder_layers = static_cast(std::stoul(value)); + else if (key == "num_decoder_layers") num_decoder_layers = static_cast(std::stoul(value)); + else if (key == "partial_rotary_factor") partial_rotary_factor = std::stof(value); + else if (key == "pad_token_id") pad_token_id = static_cast(std::stoul(value)); + else if (key == "conv_kernel_size") conv_kernel_size = static_cast(std::stoul(value)); + else if (key == "subsampling_conv_kernel_size") subsampling_conv_kernel_size = static_cast(std::stoul(value)); + else if (key == "subsampling_conv_stride") subsampling_conv_stride = static_cast(std::stoul(value)); + else if (key == "subsampling_conv_channels") subsampling_conv_channels = static_cast(std::stoul(value)); + else if (key == "subsampling_factor") subsampling_factor = static_cast(std::stoul(value)); + else if (key == "num_mel_bins") num_mel_bins = static_cast(std::stoul(value)); + else if (key == "encoder_hidden_act") encoder_hidden_act = value; + else if (key == "linear_num_key_heads") linear_num_key_heads = static_cast(std::stoul(value)); + else if (key == "linear_key_head_dim") linear_key_head_dim = static_cast(std::stoul(value)); + else if (key == "linear_num_value_heads") linear_num_value_heads = static_cast(std::stoul(value)); + else if (key == "linear_value_head_dim") linear_value_head_dim = static_cast(std::stoul(value)); + else if (key == "linear_q_proj_dim") linear_q_proj_dim = static_cast(std::stoul(value)); + else if (key == "kv_lora_rank") kv_lora_rank = static_cast(std::stoul(value)); + else if (key == "q_lora_rank") q_lora_rank = static_cast(std::stoul(value)); + else if (key == "qk_head_dim") qk_head_dim = static_cast(std::stoul(value)); + else if (key == "qk_nope_head_dim") qk_nope_head_dim = static_cast(std::stoul(value)); + else if (key == "qk_rope_head_dim") qk_rope_head_dim = static_cast(std::stoul(value)); + else if (key == "v_head_dim") v_head_dim = static_cast(std::stoul(value)); + else if (key == "rope_interleave") rope_interleave = (value == "true" || value == "1"); + else if (key == "attention_bias") attention_bias = (value == "true" || value == "1"); + else if (key == "rope_scaling_factor") rope_scaling_factor = std::stof(value); + else if (key == "rope_mscale_all_dim") rope_mscale_all_dim = std::stof(value); + else if (key == "linear_k_proj_dim") linear_k_proj_dim = static_cast(std::stoul(value)); + else if (key == "linear_v_proj_dim") linear_v_proj_dim = static_cast(std::stoul(value)); + else if (key == "predictor_hidden_dim") predictor_hidden_dim = static_cast(std::stoul(value)); + else if (key == "predictor_num_layers") predictor_num_layers = static_cast(std::stoul(value)); + else if (key == "tdt_joint_dim") tdt_joint_dim = static_cast(std::stoul(value)); + else if (key == "tdt_num_durations") tdt_num_durations = static_cast(std::stoul(value)); + else if (key == "tdt_blank_id") tdt_blank_id = static_cast(std::stoul(value)); + else if (key == "tdt_durations") { + tdt_durations.clear(); + std::stringstream ss(value); + std::string item; + while (std::getline(ss, item, ',')) { + size_t first = item.find_first_not_of(" \t"); + if (first == std::string::npos) continue; + size_t last = item.find_last_not_of(" \t"); + item = item.substr(first, last - first + 1); + tdt_durations.push_back(static_cast(std::stoul(item))); + } + } + else if (key == "altup_num_inputs") altup_num_inputs = static_cast(std::stoul(value)); + else if (key == "laurel_rank") laurel_rank = static_cast(std::stoul(value)); + else if (key == "hidden_size_per_layer_input") hidden_size_per_layer_input = static_cast(std::stoul(value)); + else if (key == "num_kv_shared_layers") num_kv_shared_layers = static_cast(std::stoul(value)); + else if (key == "sliding_window") sliding_window = static_cast(std::stoul(value)); + else if (key == "rope_local_base_freq") rope_local_base_freq = std::stof(value); + else if (key == "final_logit_softcapping") final_logit_softcapping = std::stof(value); + else if (key == "global_partial_rotary_factor") global_partial_rotary_factor = std::stof(value); + else if (key == "expert_intermediate_size") expert_intermediate_size = static_cast(std::stoul(value)); + else if (key == "global_head_dim") global_head_dim = static_cast(std::stoul(value)); + else if (key == "num_global_kv_heads" || key == "num_global_key_value_heads") num_global_kv_heads = static_cast(std::stoul(value)); + else if (key == "attention_k_eq_v") attention_k_eq_v = (value == "true" || value == "1"); + else if (key == "enable_moe_block") enable_moe_block = (value == "true" || value == "1"); + else if (key == "vision_head_dim") vision_head_dim = static_cast(std::stoul(value)); + else if (key == "vision_kv_heads") vision_kv_heads = static_cast(std::stoul(value)); + else if (key == "vision_intermediate_size") vision_intermediate_size = static_cast(std::stoul(value)); + else if (key == "vision_position_embedding_size") vision_position_embedding_size = static_cast(std::stoul(value)); + else if (key == "vision_pooling_kernel_size") vision_pooling_kernel_size = static_cast(std::stoul(value)); + else if (key == "vision_default_output_length") vision_default_output_length = static_cast(std::stoul(value)); + else if (key == "vision_rope_theta") vision_rope_theta = std::stof(value); + else if (key == "audio_hidden_dim") audio_hidden_dim = static_cast(std::stoul(value)); + else if (key == "audio_num_layers") audio_num_layers = static_cast(std::stoul(value)); + else if (key == "audio_num_heads") audio_num_heads = static_cast(std::stoul(value)); + else if (key == "audio_head_dim") audio_head_dim = static_cast(std::stoul(value)); + else if (key == "audio_input_feat_size") audio_input_feat_size = static_cast(std::stoul(value)); + else if (key == "audio_conf_conv_kernel_size") audio_conf_conv_kernel_size = static_cast(std::stoul(value)); + else if (key == "audio_chunk_size") audio_chunk_size = static_cast(std::stoul(value)); + else if (key == "audio_context_left") audio_context_left = static_cast(std::stoul(value)); + else if (key == "audio_context_right") audio_context_right = static_cast(std::stoul(value)); + else if (key == "audio_logit_cap") audio_logit_cap = std::stof(value); + else if (key == "audio_residual_weight") audio_residual_weight = std::stof(value); + else if (key == "audio_output_proj_dims") audio_output_proj_dims = static_cast(std::stoul(value)); + else if (key == "audio_vocab_size") audio_vocab_size = static_cast(std::stoul(value)); + else if (key == "audio_vocab_offset") audio_vocab_offset = static_cast(std::stoul(value)); + else if (key == "audio_soft_tokens") audio_soft_tokens = static_cast(std::stoul(value)); + else if (key == "audio_sscp_conv0_channels") audio_sscp_conv0_channels = static_cast(std::stoul(value)); + else if (key == "audio_sscp_conv1_channels") audio_sscp_conv1_channels = static_cast(std::stoul(value)); + else if (key == "audio_sscp_conv_eps") audio_sscp_conv_eps = std::stof(value); + else if (key == "audio_rms_norm_eps") audio_rms_norm_eps = std::stof(value); + else if (key == "audio_fft_length") audio_fft_length = static_cast(std::stoul(value)); + else if (key == "audio_fft_overdrive") { + audio_fft_overdrive = (value == "true" || value == "1"); + audio_fft_length = audio_fft_overdrive ? 1024u : 512u; + } + else if (key == "audio_token_id") audio_token_id = static_cast(std::stoul(value)); + else if (key == "activation_sparsity_ppf") { + activation_sparsity_ppf.clear(); + std::stringstream ss(value); + std::string item; + while (std::getline(ss, item, ',')) { + size_t first = item.find_first_not_of(" \t"); + if (first == std::string::npos) continue; + size_t last = item.find_last_not_of(" \t"); + item = item.substr(first, last - first + 1); + activation_sparsity_ppf.push_back(std::stof(item)); + } + } + } + + if (is_gemma_family(model_type)) { + default_temperature = 1.0f; + default_top_p = 0.95f; + default_top_k = 64; + if (model_type == ModelType::GEMMA4) { + default_cloud_handoff_threshold = 0.81f; + } + } else if (model_type == ModelType::LFM2) { + default_temperature = 0.3f; + default_top_p = 0.95f; + default_top_k = 20; + } else if (model_type == ModelType::QWEN) { + default_temperature = 0.6f; + default_top_p = 0.95f; + default_top_k = 20; + } else if (model_type == ModelType::QWEN3P5) { + default_temperature = 0.7f; + default_top_p = 0.8f; + default_top_k = 20; + } + + if (model_type == ModelType::GEMMA4) { + auto missing_u32 = [](uint32_t v) { return v == UNSET_U32; }; + auto missing_f32 = [](float v) { return v == UNSET_F32; }; + std::string missing; + if (missing_u32(hidden_size_per_layer_input)) missing += " hidden_size_per_layer_input"; + if (missing_u32(num_kv_shared_layers)) missing += " num_kv_shared_layers"; + if (missing_u32(sliding_window)) missing += " sliding_window"; + if (missing_u32(global_head_dim)) missing += " global_head_dim"; + if (missing_f32(rope_local_base_freq)) missing += " rope_local_base_freq"; + if (missing_f32(final_logit_softcapping)) missing += " final_logit_softcapping"; + if (missing_f32(global_partial_rotary_factor)) missing += " global_partial_rotary_factor"; + if (layer_types.empty()) missing += " layer_types"; + if (!missing.empty()) { + CACTUS_LOG_ERROR("config", "Gemma4 config missing required fields:" << missing); + return false; + } + } + + if (!decoder_start_token_seen) { + decoder_start_token_id = bos_token_id; + } + + validate_kv_compress(); + return true; +} + +bool Config::parse_kv_compress_override(const char* trigger_env, const char* target_env) { + const bool has_trigger = trigger_env && *trigger_env; + const bool has_target = target_env && *target_env; + if (!has_trigger && !has_target) return false; + if (has_trigger) kv_compress_trigger_len = static_cast(std::stol(trigger_env)); + if (has_target) kv_compress_target_len = static_cast(std::stol(target_env)); + if (kv_compress_trigger_len <= 0) { + kv_compress = false; + kv_compress_trigger_len = 0; + kv_compress_target_len = 0; + CACTUS_LOG_INFO("kv_compress", "rolling compaction disabled (CACTUS_KV_COMPRESS_AT <= 0)"); + return true; + } + kv_compress = true; + validate_kv_compress(); + CACTUS_LOG_INFO("kv_compress", "rolling override: trigger_len=" << kv_compress_trigger_len + << " target_len=" << kv_compress_target_len); + return true; +} + +void Config::validate_kv_compress() { + if (kv_compress_trigger_len <= 0) return; + if (kv_compress_target_len <= 0 || kv_compress_target_len >= kv_compress_trigger_len) { + CACTUS_LOG_WARN("kv_compress", "invalid rolling config (target_len=" << kv_compress_target_len + << ", trigger_len=" << kv_compress_trigger_len + << "): require 0 < target_len < trigger_len; disabling rolling compaction"); + kv_compress = false; + kv_compress_trigger_len = 0; + kv_compress_target_len = 0; + } +} + +std::string Config::to_json() const { + return "{}"; +} + +std::unique_ptr create_model(const std::string& bundle_dir) { + CACTUS_LOG_DEBUG("model", "Creating model from: " << bundle_dir); + fs::path manifest = fs::path(bundle_dir) / "components" / "manifest.json"; + if (!fs::exists(manifest)) { + CACTUS_LOG_ERROR("model", + "Not a transpiled bundle (no components/manifest.json at " << bundle_dir << "). " + "Run `cactus convert ` to produce one."); + return nullptr; + } + return std::make_unique(); +} + +const std::vector& Model::get_debug_nodes() const { + debug_nodes_.clear(); + return debug_nodes_; +} + + +double Model::score_tokens_window_logprob(const std::vector& /*tokens*/, size_t /*start*/, + size_t /*end*/, size_t /*context*/, size_t* tokens_scored) { + if (tokens_scored) *tokens_scored = 0; + return 0.0; +} + +} +} diff --git a/cactus/ffi/cactus_rag.cpp b/cactus-engine/src/rag.cpp similarity index 74% rename from cactus/ffi/cactus_rag.cpp rename to cactus-engine/src/rag.cpp index ad949ce13..81300fba2 100644 --- a/cactus/ffi/cactus_rag.cpp +++ b/cactus-engine/src/rag.cpp @@ -1,5 +1,5 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" +#include "../cactus_engine.h" +#include "utils.h" #include #include #include @@ -19,6 +19,41 @@ static constexpr float RRF_BM25_WEIGHT = 0.2f; static constexpr float BM25_K1 = 1.5f; static constexpr float BM25_B = 0.75f; +static std::vector> compute_rrf_scores( + const std::vector>& emb_ranked, + const std::vector>& bm25_ranked, + size_t num_items +) { + std::vector> emb_sorted = emb_ranked; + std::vector> bm25_sorted = bm25_ranked; + + std::sort(emb_sorted.begin(), emb_sorted.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + std::sort(bm25_sorted.begin(), bm25_sorted.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + + std::unordered_map emb_rank_map, bm25_rank_map; + for (size_t r = 0; r < emb_sorted.size(); ++r) { + emb_rank_map[emb_sorted[r].second] = r + 1; + } + for (size_t r = 0; r < bm25_sorted.size(); ++r) { + bm25_rank_map[bm25_sorted[r].second] = r + 1; + } + + std::vector> rrf_scored; + rrf_scored.reserve(num_items); + for (size_t i = 0; i < num_items; ++i) { + float rrf = RRF_EMB_WEIGHT / (RRF_K + emb_rank_map[i]) + + RRF_BM25_WEIGHT / (RRF_K + bm25_rank_map[i]); + rrf_scored.emplace_back(rrf, i); + } + + std::sort(rrf_scored.begin(), rrf_scored.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + + return rrf_scored; +} + static std::vector tokenize_words(const std::string& text) { std::vector words; std::string current; @@ -89,7 +124,13 @@ std::string retrieve_rag_context(CactusModelHandle* handle, const std::string& q std::vector query_tokens = tokenizer->encode(query); if (query_tokens.empty()) return ""; - std::vector query_embedding = handle->model->get_embeddings(query_tokens, true, true); + std::vector query_embedding; + try { + query_embedding = handle->model->get_text_embeddings(query_tokens, true); + } catch (const std::exception& e) { + CACTUS_LOG_WARN("rag", "get_embeddings unavailable, skipping RAG context: " << e.what()); + return ""; + } if (query_embedding.size() != handle->corpus_embedding_dim) { CACTUS_LOG_WARN("rag", "Query embedding dimension mismatch"); return ""; @@ -143,28 +184,7 @@ std::string retrieve_rag_context(CactusModelHandle* handle, const std::string& q bm25_ranked.emplace_back(bm25, i); } - std::sort(emb_ranked.begin(), emb_ranked.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - std::sort(bm25_ranked.begin(), bm25_ranked.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - - std::unordered_map emb_rank_map, bm25_rank_map; - for (size_t r = 0; r < emb_ranked.size(); ++r) { - emb_rank_map[emb_ranked[r].second] = r + 1; - } - for (size_t r = 0; r < bm25_ranked.size(); ++r) { - bm25_rank_map[bm25_ranked[r].second] = r + 1; - } - - std::vector> rrf_scored; - for (size_t i = 0; i < docs.size(); ++i) { - float rrf = RRF_EMB_WEIGHT / (RRF_K + emb_rank_map[i]) + - RRF_BM25_WEIGHT / (RRF_K + bm25_rank_map[i]); - rrf_scored.emplace_back(rrf, i); - } - - std::sort(rrf_scored.begin(), rrf_scored.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); + auto rrf_scored = compute_rrf_scores(emb_ranked, bm25_ranked, docs.size()); std::string context = "[Retrieved Context - Use ONLY this information to answer. If the answer is not in the context, say \"I don't have enough information to answer that.\"]\n"; size_t count = std::min(RAG_TOP_K, rrf_scored.size()); @@ -248,8 +268,15 @@ std::vector select_relevant_tools( std::vector tokens = tokenizer->encode(text); if (!tokens.empty()) { - std::vector emb = handle->model->get_embeddings(tokens, true, true); - handle->tool_embeddings.push_back(std::move(emb)); + try { + std::vector emb = handle->model->get_text_embeddings(tokens, true); + handle->tool_embeddings.push_back(std::move(emb)); + } catch (const std::exception& e) { + CACTUS_LOG_WARN("tool_rag", "get_embeddings unavailable, returning all tools: " << e.what()); + handle->tool_texts.clear(); + handle->tool_embeddings.clear(); + return all_tools; + } } else { handle->tool_embeddings.push_back({}); } @@ -263,7 +290,13 @@ std::vector select_relevant_tools( return all_tools; } - std::vector query_embedding = handle->model->get_embeddings(query_tokens, true, true); + std::vector query_embedding; + try { + query_embedding = handle->model->get_text_embeddings(query_tokens, true); + } catch (const std::exception& e) { + CACTUS_LOG_WARN("tool_rag", "get_embeddings unavailable, returning all tools: " << e.what()); + return all_tools; + } if (query_embedding.empty()) { CACTUS_LOG_WARN("tool_rag", "Failed to get query embedding, returning all tools"); return all_tools; @@ -271,7 +304,6 @@ std::vector select_relevant_tools( auto query_words = tokenize_words(query); - // Compute BM25 stats float total_len = 0.0f; std::unordered_map doc_freqs; for (const auto& text : handle->tool_texts) { @@ -284,14 +316,12 @@ std::vector select_relevant_tools( } float avg_doc_len = total_len / handle->tool_texts.size(); - // Compute embedding similarity scores std::vector> emb_ranked; for (size_t i = 0; i < handle->tool_embeddings.size(); ++i) { float sim = cosine_similarity(query_embedding, handle->tool_embeddings[i]); emb_ranked.emplace_back(sim, i); } - // Compute BM25 scores std::vector> bm25_ranked; for (size_t i = 0; i < handle->tool_texts.size(); ++i) { float bm25 = compute_bm25_score( @@ -300,32 +330,7 @@ std::vector select_relevant_tools( bm25_ranked.emplace_back(bm25, i); } - // Sort by scores (descending) - std::sort(emb_ranked.begin(), emb_ranked.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - std::sort(bm25_ranked.begin(), bm25_ranked.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - - // Build rank maps - std::unordered_map emb_rank_map, bm25_rank_map; - for (size_t r = 0; r < emb_ranked.size(); ++r) { - emb_rank_map[emb_ranked[r].second] = r + 1; - } - for (size_t r = 0; r < bm25_ranked.size(); ++r) { - bm25_rank_map[bm25_ranked[r].second] = r + 1; - } - - // Compute RRF scores - std::vector> rrf_scored; - for (size_t i = 0; i < all_tools.size(); ++i) { - float rrf = RRF_EMB_WEIGHT / (RRF_K + emb_rank_map[i]) + - RRF_BM25_WEIGHT / (RRF_K + bm25_rank_map[i]); - rrf_scored.emplace_back(rrf, i); - } - - // Sort by RRF score (descending) - std::sort(rrf_scored.begin(), rrf_scored.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); + auto rrf_scored = compute_rrf_scores(emb_ranked, bm25_ranked, all_tools.size()); // Select top-k tools std::vector selected; @@ -340,6 +345,16 @@ std::vector select_relevant_tools( return selected; } +static int rag_write_response(char* buffer, size_t buffer_size, const char* str, int success_code) { + size_t len = std::strlen(str); + if (len >= buffer_size) { + buffer[0] = '\0'; + return -1; + } + std::memcpy(buffer, str, len + 1); + return success_code; +} + extern "C" { int cactus_rag_query( @@ -356,27 +371,23 @@ int cactus_rag_query( auto* handle = static_cast(model); if (!handle->corpus_index || handle->corpus_embedding_dim == 0) { - std::strcpy(response_buffer, "{\"chunks\":[],\"error\":\"No corpus index loaded\"}"); - return 0; + return rag_write_response(response_buffer, buffer_size, "{\"chunks\":[],\"error\":\"No corpus index loaded\"}", 0); } try { auto* tokenizer = handle->model->get_tokenizer(); if (!tokenizer) { - std::strcpy(response_buffer, "{\"chunks\":[],\"error\":\"No tokenizer\"}"); - return 0; + return rag_write_response(response_buffer, buffer_size, "{\"chunks\":[],\"error\":\"No tokenizer\"}", 0); } std::vector query_tokens = tokenizer->encode(query); if (query_tokens.empty()) { - std::strcpy(response_buffer, "{\"chunks\":[],\"error\":\"Empty query\"}"); - return 0; + return rag_write_response(response_buffer, buffer_size, "{\"chunks\":[],\"error\":\"Empty query\"}", 0); } - std::vector query_embedding = handle->model->get_embeddings(query_tokens, true, true); + std::vector query_embedding = handle->model->get_text_embeddings(query_tokens, true); if (query_embedding.size() != handle->corpus_embedding_dim) { - std::strcpy(response_buffer, "{\"chunks\":[],\"error\":\"Embedding dimension mismatch\"}"); - return 0; + return rag_write_response(response_buffer, buffer_size, "{\"chunks\":[],\"error\":\"Embedding dimension mismatch\"}", 0); } index::QueryOptions options; @@ -387,8 +398,7 @@ int cactus_rag_query( auto results = handle->corpus_index->query(query_embeddings, options); if (results.empty() || results[0].empty()) { - std::strcpy(response_buffer, "{\"chunks\":[]}"); - return 0; + return rag_write_response(response_buffer, buffer_size, "{\"chunks\":[]}", 0); } std::vector doc_ids; @@ -424,28 +434,7 @@ int cactus_rag_query( bm25_ranked.emplace_back(bm25, i); } - std::sort(emb_ranked.begin(), emb_ranked.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - std::sort(bm25_ranked.begin(), bm25_ranked.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - - std::unordered_map emb_rank_map, bm25_rank_map; - for (size_t r = 0; r < emb_ranked.size(); ++r) { - emb_rank_map[emb_ranked[r].second] = r + 1; - } - for (size_t r = 0; r < bm25_ranked.size(); ++r) { - bm25_rank_map[bm25_ranked[r].second] = r + 1; - } - - std::vector> rrf_scored; - for (size_t i = 0; i < docs.size(); ++i) { - float rrf = RRF_EMB_WEIGHT / (RRF_K + emb_rank_map[i]) + - RRF_BM25_WEIGHT / (RRF_K + bm25_rank_map[i]); - rrf_scored.emplace_back(rrf, i); - } - - std::sort(rrf_scored.begin(), rrf_scored.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); + auto rrf_scored = compute_rrf_scores(emb_ranked, bm25_ranked, docs.size()); size_t result_count = std::min(top_k > 0 ? top_k : RAG_TOP_K, rrf_scored.size()); @@ -457,36 +446,21 @@ int cactus_rag_query( if (i > 0) oss << ","; oss << "{\"score\":" << std::setprecision(4) << final_score - << ",\"source\":\"" << docs[idx].metadata << "\"" - << ",\"content\":\""; - for (char c : docs[idx].content) { - switch (c) { - case '"': oss << "\\\""; break; - case '\\': oss << "\\\\"; break; - case '\n': oss << "\\n"; break; - case '\r': oss << "\\r"; break; - case '\t': oss << "\\t"; break; - default: oss << c; - } - } - oss << "\"}"; + << ",\"source\":\"" << escape_json_string(docs[idx].metadata) << "\"" + << ",\"content\":\"" << escape_json_string(docs[idx].content) << "\"}"; } oss << "]}"; std::string result = oss.str(); if (result.size() >= buffer_size) { - std::strcpy(response_buffer, "{\"chunks\":[],\"error\":\"Buffer too small\"}"); - return -1; + return rag_write_response(response_buffer, buffer_size, "{\"chunks\":[],\"error\":\"Buffer too small\"}", -1); } - std::strcpy(response_buffer, result.c_str()); - return static_cast(result.size()); + return rag_write_response(response_buffer, buffer_size, result.c_str(), static_cast(result.size())); } catch (const std::exception& e) { - std::string error = "{\"chunks\":[],\"error\":\"" + std::string(e.what()) + "\"}"; - if (error.size() < buffer_size) { - std::strcpy(response_buffer, error.c_str()); - } + std::string error = "{\"chunks\":[],\"error\":\"" + escape_json_string(e.what()) + "\"}"; + rag_write_response(response_buffer, buffer_size, error.c_str(), -1); return -1; } } diff --git a/cactus/engine/engine_sp.cpp b/cactus-engine/src/sp.cpp similarity index 72% rename from cactus/engine/engine_sp.cpp rename to cactus-engine/src/sp.cpp index 7f1e2bbf7..89e536ece 100644 --- a/cactus/engine/engine_sp.cpp +++ b/cactus-engine/src/sp.cpp @@ -1,6 +1,8 @@ #include "engine.h" #include +#include #include +#include #include #include #include @@ -35,6 +37,7 @@ void SPTokenizer::cleanup_mmap() { } bool SPTokenizer::load_vocabulary_with_config(const std::string& vocab_file, const std::string& /*merges_file*/, const std::string& config_file) { + runtime_config_ = load_tokenizer_runtime_config(config_file); std::string config_path = config_file.substr(0, config_file.find_last_of("/\\")) + "/config.txt"; detect_model_type(config_path); @@ -83,13 +86,20 @@ bool SPTokenizer::load_vocabulary_with_config(const std::string& vocab_file, con } if (!token.empty() && id != UINT32_MAX) { + float score = -static_cast(id); + size_t tab_in_token = token.find('\t'); + if (tab_in_token != std::string::npos) { + std::string score_str = token.substr(tab_in_token + 1); + if (!score_str.empty()) score = std::stof(score_str); + token = token.substr(0, tab_in_token); + } token_to_id_[token] = id; if (id >= id_to_token_.size()) { id_to_token_.resize(id + 1); token_scores_.resize(id + 1, 0.0f); } id_to_token_[id] = token; - token_scores_[id] = 0.0f; + token_scores_[id] = score; } } vocab_size_ = id_to_token_.size(); @@ -136,6 +146,12 @@ bool SPTokenizer::load_vocabulary_with_config(const std::string& vocab_file, con unk_token_id_ = std::stoul(value); } else if (key == "bos_token_id") { bos_token_id_ = std::stoul(value); + } else if (key == "sp_model_type") { + sp_bpe_mode_ = (value == "bpe" || value == "BPE"); + } else if (key == "sp_add_dummy_prefix") { + sp_add_dummy_prefix_ = (value == "true" || value == "1"); + } else if (key == "sp_byte_fallback") { + sp_byte_fallback_ = (value == "true" || value == "1"); } } } @@ -208,12 +224,22 @@ void SPTokenizer::build_trie() { std::string SPTokenizer::preprocess_text(const std::string& text) const { if (text.empty()) return text; - std::string processed = ""; - if (model_type_ == ModelType::BERT) { - processed = "▁"; + if (sp_bpe_mode_) { + std::string processed; + if (sp_add_dummy_prefix_) processed += "\xE2\x96\x81"; + for (char c : text) { + if (c == ' ') processed += "\xE2\x96\x81"; + else processed += c; + } + return processed; } - for (size_t i = text.find_first_not_of(" "); i < text.length(); i++) { + std::string processed = ""; + if (sp_add_dummy_prefix_) processed += "\xE2\x96\x81"; + + size_t start = text.find_first_not_of(" "); + if (start == std::string::npos) return processed; + for (size_t i = start; i < text.length(); i++) { char c = text[i]; if (c == ' ') { processed += "▁"; @@ -358,42 +384,55 @@ std::vector> SPTokenizer::tokenize_with_trie(co return result; } -std::vector SPTokenizer::split_with_special_tokens(const std::string& text) const { - std::vector result; - - size_t start = 0; - while (start < text.size()) { - size_t best_match_pos = text.size(); - size_t best_match_len = 0; - std::string best_special_token; - - for (const auto& [special_token, token_id] : special_tokens_) { - size_t pos = text.find(special_token, start); - if (pos != std::string::npos && pos < best_match_pos) { - best_match_pos = pos; - best_match_len = special_token.length(); - best_special_token = special_token; - } +std::vector SPTokenizer::tokenize_with_bpe(const std::string& text) const { + std::vector symbols; + for (size_t i = 0; i < text.size(); ) { + size_t char_len = 1; + unsigned char byte = static_cast(text[i]); + if ((byte & 0xE0) == 0xC0) char_len = 2; + else if ((byte & 0xF0) == 0xE0) char_len = 3; + else if ((byte & 0xF8) == 0xF0) char_len = 4; + if (i + char_len <= text.size()) symbols.push_back(text.substr(i, char_len)); + i += char_len; + } + + while (symbols.size() > 1) { + int best_index = -1; + float best_score = -std::numeric_limits::infinity(); + for (size_t i = 0; i + 1 < symbols.size(); ++i) { + auto it = token_to_id_.find(symbols[i] + symbols[i + 1]); + if (it == token_to_id_.end()) continue; + float score = it->second < token_scores_.size() + ? token_scores_[it->second] + : -static_cast(it->second); + if (score > best_score) { best_score = score; best_index = static_cast(i); } } - - if (best_match_pos < text.size()) { - if (best_match_pos > start) { - std::string before = text.substr(start, best_match_pos - start); - result.push_back(before); - } - result.push_back(best_special_token); - start = best_match_pos + best_match_len; - } else { - if (start < text.size()) { - result.push_back(text.substr(start)); + if (best_index < 0) break; + symbols[best_index] += symbols[best_index + 1]; + symbols.erase(symbols.begin() + best_index + 1); + } + + std::vector result; + for (const auto& symbol : symbols) { + auto it = token_to_id_.find(symbol); + if (it != token_to_id_.end()) { result.push_back(it->second); continue; } + if (sp_byte_fallback_) { + for (unsigned char b : symbol) { + char buf[7]; std::snprintf(buf, sizeof(buf), "<0x%02X>", b); + auto byte_it = token_to_id_.find(buf); + result.push_back(byte_it != token_to_id_.end() ? byte_it->second : unk_token_id_); } - break; + continue; } + result.push_back(unk_token_id_); } - return result; } +std::vector SPTokenizer::split_with_special_tokens(const std::string& text) const { + return cactus::engine::split_with_special_tokens(text, special_tokens_); +} + std::vector SPTokenizer::encode(const std::string& text) const { if (text.empty()) return {}; @@ -406,11 +445,16 @@ std::vector SPTokenizer::encode(const std::string& text) const { token_ids.push_back(special_it->second); } else { std::string processed = preprocess_text(segment); + if (processed.empty()) continue; - auto token_pairs = tokenize_with_trie(processed); - - for (const auto& [token, id] : token_pairs) { - token_ids.push_back(id); + if (sp_bpe_mode_) { + auto ids = tokenize_with_bpe(processed); + token_ids.insert(token_ids.end(), ids.begin(), ids.end()); + } else { + auto token_pairs = tokenize_with_trie(processed); + for (const auto& [token, id] : token_pairs) { + token_ids.push_back(id); + } } } } @@ -419,113 +463,45 @@ std::vector SPTokenizer::encode(const std::string& text) const { } std::string SPTokenizer::decode(const std::vector& tokens) const { - std::string result; - if (tokens.size() == 1) { - uint32_t token_id = tokens[0]; - if (token_id < id_to_token_.size()) { - std::string token = id_to_token_[token_id]; - - size_t pos = 0; - while (pos < token.length()) { - if (pos + 2 < token.length() && - static_cast(token[pos]) == 0xE2 && - static_cast(token[pos+1]) == 0x96 && - static_cast(token[pos+2]) == 0x81) { - result += ' '; - pos += 3; - } else { - result += token[pos]; - pos++; - } + if (tokens[0] >= id_to_token_.size()) return {}; + const std::string& piece = id_to_token_[tokens[0]]; + unsigned int byte_val; + if (piece.size() == 6 && std::sscanf(piece.c_str(), "<0x%02X>", &byte_val) == 1) { + return std::string(1, static_cast(byte_val)); + } + std::string result; + for (size_t i = 0; i < piece.length(); ) { + if (i + 2 < piece.length() && + static_cast(piece[i]) == 0xE2 && + static_cast(piece[i+1]) == 0x96 && + static_cast(piece[i+2]) == 0x81) { + result += ' '; + i += 3; + } else { + result += piece[i++]; } } return result; } - for (size_t i = 0; i < tokens.size(); i++) { - uint32_t token_id = tokens[i]; - if (token_id < id_to_token_.size()) { - std::string token = id_to_token_[token_id]; - - size_t pos = 0; - while (pos < token.length()) { - if (pos + 2 < token.length() && - static_cast(token[pos]) == 0xE2 && - static_cast(token[pos+1]) == 0x96 && - static_cast(token[pos+2]) == 0x81) { - if (!result.empty()) { - result += ' '; - } - pos += 3; - } else { - result += token[pos]; - pos++; - } - } + std::string raw; + for (uint32_t token_id : tokens) { + if (token_id >= id_to_token_.size()) continue; + const std::string& piece = id_to_token_[token_id]; + unsigned int byte_val; + if (piece.size() == 6 && std::sscanf(piece.c_str(), "<0x%02X>", &byte_val) == 1) { + raw.push_back(static_cast(byte_val)); + } else { + raw += piece; } } - - return result; + return postprocess_text(raw); } void SPTokenizer::load_special_tokens(const std::string& config_file) { - std::ifstream file(config_file); - if (!file.is_open()) { - return; - } - - std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - - size_t pos = content.find("\"special_tokens\""); - if (pos == std::string::npos) return; - - pos = content.find("{", pos); - if (pos == std::string::npos) return; - - size_t end_pos = content.find("}", pos); - if (end_pos == std::string::npos) return; - - std::string special_tokens_section = content.substr(pos + 1, end_pos - pos - 1); - - std::istringstream iss(special_tokens_section); - std::string line; - - while (std::getline(iss, line)) { - size_t colon_pos = line.find(":"); - if (colon_pos == std::string::npos) continue; - - std::string id_part = line.substr(0, colon_pos); - std::string token_part = line.substr(colon_pos + 1); - - size_t id_start = id_part.find("\""); - size_t id_end = id_part.find("\"", id_start + 1); - if (id_start == std::string::npos || id_end == std::string::npos) continue; - - std::string id_str = id_part.substr(id_start + 1, id_end - id_start - 1); - uint32_t token_id = std::stoul(id_str); - - size_t token_start = token_part.find("\""); - size_t token_end = token_part.rfind("\""); - if (token_start == std::string::npos || token_end == std::string::npos || token_start >= token_end) continue; - - std::string token_content = token_part.substr(token_start + 1, token_end - token_start - 1); - - special_tokens_[token_content] = token_id; - } + load_special_tokens_map(config_file, special_tokens_); } -void SPTokenizer::load_chat_template(const std::string& template_file) { - std::ifstream file(template_file); - if (!file.is_open()) { - has_chat_template_ = false; - return; - } - - chat_template_ = std::string((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - has_chat_template_ = !chat_template_.empty(); -} - - } // namespace engine -} // namespace cactus \ No newline at end of file +} // namespace cactus diff --git a/cactus-engine/src/stream.cpp b/cactus-engine/src/stream.cpp new file mode 100644 index 000000000..182d81e4e --- /dev/null +++ b/cactus-engine/src/stream.cpp @@ -0,0 +1,407 @@ +#include "../cactus_engine.h" +#include "utils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace cactus::ffi; + +namespace { + +constexpr size_t kScratchSize = 1u << 16; +constexpr size_t kSampleRate = 16000; +constexpr float kSampleRateF = 16000.0f; +constexpr size_t kLeftContextSamples = 8 * kSampleRate; // left context re-fed each window +constexpr size_t kRightContextSamples = 1 * kSampleRate; // look-ahead before confirming +constexpr size_t kChunkSamples = 1 * kSampleRate; // warm decode step +constexpr size_t kColdStartSamples = 6 * kSampleRate; // first decode from empty state +constexpr size_t kResumeColdSamples = 2 * kSampleRate; // first decode after silence +constexpr size_t kSilenceResetSamples = 3 * kSampleRate; // silence run that triggers a cold restart +constexpr size_t kMaxDecodeSamples = 10 * kSampleRate; + +struct StreamStats { + size_t decode_tokens = 0; + double total_time_ms = 0.0; + double decode_tps = 0.0; + double time_to_first_token_ms = 0.0; + + void finalize() { + if (total_time_ms > 0.0) decode_tps = decode_tokens * 1000.0 / total_time_ms; + } +}; + +struct StreamTranscribe { + CactusModelHandle* model = nullptr; + std::string options_json; + bool is_parakeet = false; + + std::vector samples; + size_t samples_decoded_up_to = 0; + size_t silence_run = 0; + bool cold_restart = false; + + cactus::engine::Model::ParakeetTdtStreamState pstate; + std::vector committed_tokens; + std::string emitted_text; + std::string previous_pending; + + std::vector whisper_prev_segments; + size_t whisper_last_len = 0; +}; + +double json_num(const std::string& json, const std::string& key) { + float v = 0.0f; + return try_parse_json_float(json, key, v) ? static_cast(v) : 0.0; +} + +std::vector to_pcm16(const std::vector& samples) { + std::vector pcm(samples.size()); + for (size_t i = 0; i < samples.size(); ++i) { + float x = std::max(-32768.0f, std::min(32767.0f, samples[i] * 32768.0f)); + pcm[i] = static_cast(x); + } + return pcm; +} + +std::string with_timestamps_option(std::string opts) { + if (opts.find("\"timestamps\"") != std::string::npos) return opts; + if (opts.empty() || opts == "{}") return "{\"timestamps\":true}"; + return opts.substr(0, opts.find_last_of('}')) + ",\"timestamps\":true}"; +} + +std::string trim(const std::string& s) { + size_t b = s.find_first_not_of(' '), e = s.find_last_not_of(' '); + return b == std::string::npos ? std::string() : s.substr(b, e - b + 1); +} + +std::string strip_annotations(const std::string& text) { + static const std::regex pattern(R"(\([^)]*\)|\[[^\]]*\]|\.\.\.)"); + std::string out = std::regex_replace(text, pattern, ""); + out = std::regex_replace(out, std::regex(R"(\s+)"), " "); + return trim(out); +} + +std::vector parse_segments(const std::string& json) { + std::vector segs; + for (const std::string& obj : split_json_array(json_array_field(json, "segments"))) { + float start = 0.0f, end = 0.0f; + try_parse_json_float(obj, "start", start); + try_parse_json_float(obj, "end", end); + segs.push_back({start, end, strip_annotations(json_string_field(obj, "text"))}); + } + return segs; +} + +std::vector whisper_transcribe(StreamTranscribe* s, const std::vector& samples, + StreamStats& stats) { + std::vector pcm = to_pcm16(samples); + std::string scratch(kScratchSize, '\0'); + const std::string opts = with_timestamps_option(s->options_json); + const bool prev_should_stop = s->model->should_stop.load(); + const int rc = cactus_transcribe( + static_cast(s->model), nullptr, nullptr, + scratch.data(), scratch.size(), opts.c_str(), + nullptr, nullptr, + reinterpret_cast(pcm.data()), pcm.size() * sizeof(int16_t)); + if (prev_should_stop) s->model->should_stop.store(true); + if (rc <= 0) return {}; + const std::string json(scratch.c_str()); + stats.decode_tps = json_num(json, "decode_tps"); + stats.total_time_ms = json_num(json, "total_time_ms"); + stats.time_to_first_token_ms = json_num(json, "time_to_first_token_ms"); + stats.decode_tokens = static_cast(json_num(json, "decode_tokens")); + return parse_segments(json); +} + +std::vector window_features(std::vector window, size_t mel_bins) { + if (window.empty()) return {}; + auto cfg = cactus::audio::get_parakeet_spectrogram_config(); + const size_t waveform_samples = window.size(); + cactus::audio::apply_preemphasis(window, 0.97f); + std::vector features = cactus::audio::compute_spectrogram_graph( + window, cfg, mel_bins, 0.0f, 8000.0f, cactus::audio::WHISPER_SAMPLE_RATE, 0, 0); + cactus::audio::normalize_parakeet_log_mel(features, mel_bins); + size_t valid_frames = waveform_samples / cfg.hop_length; + if (valid_frames == 0) valid_frames = 1; + cactus::audio::trim_mel_frames(features, mel_bins, valid_frames); + return features; +} + +std::string parakeet_decode_window(StreamTranscribe* s, size_t window_start, size_t window_end, + size_t decode_start_frame, size_t decode_end_frame, + bool is_final, std::string* pending_text, StreamStats& stats) { + if (pending_text) pending_text->clear(); + auto* model = s->model->model.get(); + const size_t mel_bins = std::max(1, static_cast(model->get_config().num_mel_bins)); + std::vector features = window_features( + std::vector(s->samples.begin() + window_start, s->samples.begin() + window_end), mel_bins); + if (features.empty()) return ""; + + s->pstate.time_index = decode_start_frame; + const auto t0 = std::chrono::steady_clock::now(); + std::vector tokens = model->transcribe_parakeet_tdt( + features, &s->pstate, is_final, is_final ? 0 : decode_end_frame); + stats.total_time_ms += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + stats.decode_tokens += s->pstate.decoded_tokens; + + auto* tokenizer = model->get_tokenizer(); + if (!tokenizer) return ""; + + s->committed_tokens.insert(s->committed_tokens.end(), tokens.begin(), tokens.end()); + std::string full = tokenizer->decode(s->committed_tokens); + std::string delta = full.size() > s->emitted_text.size() ? full.substr(s->emitted_text.size()) : std::string(); + s->emitted_text = full; + + if (pending_text && !s->pstate.pending.empty()) { + std::vector combined = s->committed_tokens; + combined.insert(combined.end(), s->pstate.pending.begin(), s->pstate.pending.end()); + std::string with_pending = tokenizer->decode(combined); + if (with_pending.size() > full.size()) *pending_text = with_pending.substr(full.size()); + } + + if (!is_final && !tokens.empty() && s->pstate.confirmed_sec > 0.0f) { + s->samples_decoded_up_to = window_start + static_cast(s->pstate.confirmed_sec * kSampleRateF); + } + return delta; +} + +size_t parakeet_spf(StreamTranscribe* s) { + const uint32_t subsampling = std::max(1, s->model->model->get_config().subsampling_factor); + return cactus::audio::get_parakeet_spectrogram_config().hop_length * subsampling; +} + +std::string parakeet_process(StreamTranscribe* s, std::string& pending, StreamStats& stats) { + std::lock_guard lock(s->model->model_mutex); + const size_t spf = parakeet_spf(s); + + std::string confirmed; + for (;;) { + const size_t total = s->samples.size(); + const size_t decodable = total > kRightContextSamples ? total - kRightContextSamples : 0; + const bool cold = s->samples_decoded_up_to == 0 || s->cold_restart; + const size_t cold_min = s->cold_restart ? kResumeColdSamples : kColdStartSamples; + const size_t min_chunk = cold ? cold_min : kChunkSamples; + if (decodable <= s->samples_decoded_up_to || + decodable - s->samples_decoded_up_to < min_chunk) { + break; + } + const size_t decode_to = std::min(decodable, s->samples_decoded_up_to + kMaxDecodeSamples); + const size_t window_start = cold ? s->samples_decoded_up_to + : (s->samples_decoded_up_to > kLeftContextSamples ? s->samples_decoded_up_to - kLeftContextSamples : 0); + const size_t window_end = decode_to + kRightContextSamples; + const size_t decode_start_frame = (s->samples_decoded_up_to - window_start) / spf; + const size_t decode_end_frame = decode_start_frame + (decode_to - s->samples_decoded_up_to) / spf; + if (cold) s->pstate = {}; + + std::string pend; + const size_t prev_cursor = s->samples_decoded_up_to; + const size_t tokens_before = stats.decode_tokens; + confirmed += parakeet_decode_window(s, window_start, window_end, + decode_start_frame, decode_end_frame, false, &pend, stats); + pending = pend; + if (s->pstate.decoded_tokens > 0) { s->cold_restart = false; s->silence_run = 0; } + if (s->samples_decoded_up_to > prev_cursor) continue; + if (s->pstate.decoded_tokens == 0) { + s->silence_run += decode_to - s->samples_decoded_up_to; + s->samples_decoded_up_to = decode_to; + if (s->silence_run >= kSilenceResetSamples) s->cold_restart = true; + continue; + } + if (decode_to - s->samples_decoded_up_to < cold_min) break; + stats.decode_tokens = tokens_before; + confirmed += parakeet_decode_window(s, window_start, window_end, + decode_start_frame, decode_end_frame, true, &pend, stats); + pending = pend; + s->samples_decoded_up_to = decode_to; + } + + const size_t keep_from = s->samples_decoded_up_to > kLeftContextSamples + ? s->samples_decoded_up_to - kLeftContextSamples : 0; + if (keep_from > kLeftContextSamples) { + s->samples.erase(s->samples.begin(), s->samples.begin() + keep_from); + s->samples_decoded_up_to -= keep_from; + } + + stats.finalize(); + if (confirmed.empty() && pending.empty()) pending = s->previous_pending; + else s->previous_pending = pending; + return confirmed; +} + +std::string parakeet_flush(StreamTranscribe* s, StreamStats& stats) { + std::lock_guard lock(s->model->model_mutex); + const size_t total = s->samples.size(); + if (total <= s->samples_decoded_up_to) return ""; + const size_t spf = parakeet_spf(s); + const bool cold = s->samples_decoded_up_to == 0 || s->cold_restart; + if (cold) s->pstate = {}; + const size_t window_start = cold + ? s->samples_decoded_up_to + : (s->samples_decoded_up_to > kLeftContextSamples ? s->samples_decoded_up_to - kLeftContextSamples : 0); + const size_t decode_start_frame = (s->samples_decoded_up_to - window_start) / spf; + std::string confirmed = parakeet_decode_window(s, window_start, total, decode_start_frame, 0, true, nullptr, stats); + stats.finalize(); + return confirmed; +} + +std::string segments_text(const std::vector& segs, size_t from, size_t to) { + std::string out; + for (size_t i = from; i < to && i < segs.size(); ++i) { + if (segs[i].text.empty()) continue; + if (!out.empty()) out += ' '; + out += segs[i].text; + } + return out; +} + +std::string whisper_process(StreamTranscribe* s, std::string& pending, StreamStats& stats) { + if (s->samples.size() < s->whisper_last_len + kChunkSamples) { + pending = s->previous_pending; + return ""; + } + s->whisper_last_len = s->samples.size(); + std::vector segs = whisper_transcribe(s, s->samples, stats); + + size_t agree = 0; + float confirmed_end_sec = 0.0f; + while (agree < segs.size() && agree < s->whisper_prev_segments.size() && + segs[agree].text == s->whisper_prev_segments[agree].text) { + confirmed_end_sec = std::min(segs[agree].end, s->whisper_prev_segments[agree].end); + ++agree; + } + + std::string confirmed; + pending = segments_text(segs, agree, segs.size()); + if (agree > 0 && confirmed_end_sec > 0.0f) { + confirmed = segments_text(segs, 0, agree); + const size_t cut = std::min(static_cast(confirmed_end_sec * kSampleRateF), s->samples.size()); + s->samples.erase(s->samples.begin(), s->samples.begin() + cut); + s->whisper_last_len -= cut; + std::vector tail; + for (size_t i = agree; i < segs.size(); ++i) + tail.push_back({segs[i].start - confirmed_end_sec, segs[i].end - confirmed_end_sec, segs[i].text}); + s->whisper_prev_segments = std::move(tail); + } else { + s->whisper_prev_segments = segs; + } + + s->previous_pending = pending; + return confirmed; +} + +std::string whisper_flush(StreamTranscribe* s, StreamStats& stats) { + std::vector segs = whisper_transcribe(s, s->samples, stats); + return segments_text(segs, 0, segs.size()); +} + +int write_result(char* buffer, size_t size, const std::string& confirmed, + const std::string& pending, const StreamStats& stats) { + std::ostringstream os; + os << "{\"success\":true,\"confirmed\":\"" << escape_json_string(confirmed) + << "\",\"pending\":\"" << escape_json_string(pending) + << "\",\"decode_tps\":" << stats.decode_tps + << ",\"total_time_ms\":" << stats.total_time_ms + << ",\"time_to_first_token_ms\":" << stats.time_to_first_token_ms + << ",\"decode_tokens\":" << stats.decode_tokens << "}"; + const std::string json = os.str(); + if (!buffer || size == 0) return 0; + if (json.size() >= size) { + handle_error_response("Stream response buffer too small", buffer, size); + return -1; + } + std::memcpy(buffer, json.c_str(), json.size() + 1); + return static_cast(json.size()); +} + +} // namespace + +extern "C" { + +cactus_stream_transcribe_t cactus_stream_transcribe_start(cactus_model_t model, const char* options_json) { + if (!model) { + last_error_message = "stream_transcribe_start: model is null"; + CACTUS_LOG_ERROR("stream_transcribe_start", last_error_message); + return nullptr; + } + try { + auto* handle = static_cast(model); + const auto type = handle->model->get_config().model_type; + const bool is_whisper = type == cactus::engine::Config::ModelType::WHISPER; + const bool is_parakeet = type == cactus::engine::Config::ModelType::PARAKEET_TDT; + if (!is_whisper && !is_parakeet) { + last_error_message = "stream_transcribe_start: only Whisper and Parakeet models support streaming"; + CACTUS_LOG_ERROR("stream_transcribe_start", last_error_message); + return nullptr; + } + auto s = std::make_unique(); + s->model = handle; + s->is_parakeet = is_parakeet; + if (options_json && options_json[0] != '\0') s->options_json = options_json; + CACTUS_LOG_INFO("stream_transcribe_start", "streaming session opened"); + return static_cast(s.release()); + } catch (const std::exception& e) { + last_error_message = std::string("stream_transcribe_start: ") + e.what(); + CACTUS_LOG_ERROR("stream_transcribe_start", last_error_message); + return nullptr; + } +} + +int cactus_stream_transcribe_process(cactus_stream_transcribe_t stream, + const uint8_t* pcm_buffer, size_t pcm_buffer_size, + char* response_buffer, size_t buffer_size) { + if (!stream) { + last_error_message = "stream_transcribe_process: stream is null"; + CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); + handle_error_response(last_error_message, response_buffer, buffer_size); + return -1; + } + auto* s = static_cast(stream); + try { + if (pcm_buffer && pcm_buffer_size >= sizeof(int16_t)) { + std::vector news = cactus::audio::pcm_buffer_to_float_samples(pcm_buffer, pcm_buffer_size); + s->samples.insert(s->samples.end(), news.begin(), news.end()); + } + StreamStats stats; + std::string pending; + std::string confirmed = s->is_parakeet ? parakeet_process(s, pending, stats) + : whisper_process(s, pending, stats); + return write_result(response_buffer, buffer_size, confirmed, pending, stats); + } catch (const std::exception& e) { + last_error_message = std::string("stream_transcribe_process: ") + e.what(); + CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); + handle_error_response(last_error_message, response_buffer, buffer_size); + return -1; + } +} + +int cactus_stream_transcribe_stop(cactus_stream_transcribe_t stream, + char* response_buffer, size_t buffer_size) { + if (!stream) { + last_error_message = "stream_transcribe_stop: stream is null"; + CACTUS_LOG_ERROR("stream_transcribe_stop", last_error_message); + handle_error_response(last_error_message, response_buffer, buffer_size); + return -1; + } + auto* s = static_cast(stream); + int result = 0; + try { + StreamStats stats; + std::string confirmed = s->is_parakeet ? parakeet_flush(s, stats) : whisper_flush(s, stats); + result = write_result(response_buffer, buffer_size, confirmed, "", stats); + } catch (const std::exception& e) { + last_error_message = std::string("stream_transcribe_stop: ") + e.what(); + CACTUS_LOG_ERROR("stream_transcribe_stop", last_error_message); + handle_error_response(last_error_message, response_buffer, buffer_size); + result = -1; + } + delete s; + return result; +} + +} // extern "C" diff --git a/cactus-engine/src/telemetry.cpp b/cactus-engine/src/telemetry.cpp new file mode 100644 index 000000000..360f53973 --- /dev/null +++ b/cactus-engine/src/telemetry.cpp @@ -0,0 +1,22 @@ +#include "../cactus_engine.h" +#include "telemetry.h" + +extern "C" { + +void cactus_set_telemetry_environment(const char* framework, const char* cache_location, const char* version) { + cactus::telemetry::setTelemetryEnvironment(framework, cache_location, version); +} + +void cactus_set_app_id(const char* app_id) { + cactus::telemetry::setAppId(app_id); +} + +void cactus_telemetry_flush(void) { + cactus::telemetry::flush(); +} + +void cactus_telemetry_shutdown(void) { + cactus::telemetry::shutdown(); +} + +} diff --git a/cactus-engine/src/telemetry.h b/cactus-engine/src/telemetry.h new file mode 100644 index 000000000..cd7b49f45 --- /dev/null +++ b/cactus-engine/src/telemetry.h @@ -0,0 +1,44 @@ +#pragma once +#include +#include + +namespace cactus { +namespace telemetry { + +struct CompletionMetrics { + bool success; + bool cloud_handoff; + double ttft_ms; + double prefill_tps; + double decode_tps; + double response_time_ms; + double confidence; + double ram_usage_mb; + size_t prefill_tokens; + size_t decode_tokens; + const char* error_message; + const char* function_calls_json; +}; + +void init(const char* project_id = nullptr, const char* project_scope = nullptr, const char* cloud_key = nullptr); +void setEnabled(bool enabled); +void setCloudDisabled(bool disabled); +void setAppId(const char* app_id); +void setTelemetryEnvironment(const char* framework, const char* cache_location, const char* version = nullptr); +void setCloudKey(const char* key); +void cacheCloudApiKey(const char* key); +std::string loadCachedCloudApiKey(); +void recordInit(const char* model, bool success, double response_time_ms, const char* message); +void recordCompletion(const char* model, const CompletionMetrics& metrics); +void recordCompletion(const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, const char* message); +void recordEmbedding(const char* model, bool success, const char* message); +void recordTranscription(const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, double ram_usage_mb, const char* message); +void recordStreamTranscription(const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, double session_ttft_ms, double session_tps, double session_time_ms, int session_tokens, const char* message); +void setStreamMode(bool in_stream); +bool isStreamMode(); +void markInference(bool active); +void flush(); +void shutdown(); + +} // namespace telemetry +} // namespace cactus diff --git a/cactus-engine/src/telemetry_impl.cpp b/cactus-engine/src/telemetry_impl.cpp new file mode 100644 index 000000000..150440b53 --- /dev/null +++ b/cactus-engine/src/telemetry_impl.cpp @@ -0,0 +1,1539 @@ +#include "telemetry.h" +#include "json_escape.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if defined(__APPLE__) +#include +#include +#endif +#include +#include +#include +#include +#include + +namespace cactus { +namespace telemetry { + +enum EventType { INIT = 0, COMPLETION = 1, EMBEDDING = 2, TRANSCRIPTION = 3, STREAM_TRANSCRIBE = 4 }; + +struct Event { + EventType type; + char model[128]; + bool success; + bool cloud_handoff; + double ttft_ms; + double prefill_tps; + double decode_tps; + double tps; + double response_time_ms; + double confidence; + double ram_usage_mb; + int tokens; + int prefill_tokens; + int decode_tokens; + double session_ttft_ms; + double session_tps; + double session_time_ms; + int session_tokens; + char message[256]; + char error[256]; + char function_calls[1024]; + std::chrono::system_clock::time_point timestamp; +}; +static bool enabled = false; +static int inference_active = 0; +static bool shutdown_called = false; +static bool atexit_registered = false; +static bool curl_initialized = false; +static bool cloud_disabled = false; +static std::string supabase_url = "https://vlqqczxwyaodtcdmdmlw.supabase.co"; +static std::string supabase_key = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InZscXFjenh3eWFvZHRjZG1kbWx3Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3NTE1MTg2MzIsImV4cCI6MjA2NzA5NDYzMn0.nBzqGuK9j6RZ6mOPWU2boAC_5H9XDs-fPpo5P3WZYbI"; +static std::string device_id; +static std::string project_id; +static std::string cloud_key; +static std::string project_scope; +static std::string device_model; +static std::string device_os; +static std::string device_os_version; +static std::string device_brand; +static std::string cactus_version; +#ifndef CACTUS_DEFAULT_FRAMEWORK +#define CACTUS_DEFAULT_FRAMEWORK "cpp" +#endif +static std::string framework = CACTUS_DEFAULT_FRAMEWORK; +static std::string custom_cache_location; +static std::string app_id; +static std::string device_registered_file; +static std::string project_registered_file; +static bool device_registered = false; +static bool project_registered = false; +static bool ids_ready = false; +static bool in_stream_mode = false; + +static std::mutex telemetry_mutex; +static std::condition_variable telemetry_lifecycle_cv; + +enum class TelemetryLifecycleState { + Stopped, + Running, + ShuttingDown, +}; + +static TelemetryLifecycleState lifecycle_state = TelemetryLifecycleState::Stopped; + +static bool can_record_event_locked() { + return enabled && ids_ready && lifecycle_state == TelemetryLifecycleState::Running; +} + +struct CloudSendResult { + bool payload_ok = false; + bool project_registered_ok = false; + bool device_registered_ok = false; +}; + +struct CloudConfigurationStateSnapshot { + std::string supabase_url; + std::string supabase_key; + std::string device_id; + std::string project_id; + std::string cloud_key; + std::string project_scope; + std::string framework; + std::string cactus_version; + std::string device_model; + std::string device_os; + std::string device_os_version; + std::string device_brand; + std::string app_id; + std::string device_registered_file; + std::string project_registered_file; + bool enabled = false; + bool cloud_disabled = false; + bool device_registered = false; + bool project_registered = false; +}; + +static std::string new_uuid(); +static std::string format_timestamp(const std::chrono::system_clock::time_point& tp); +static std::string model_basename(const char* model_path) { + if (!model_path) return {}; + std::string m(model_path); + size_t pos = m.find_last_of("/\\"); + if (pos != std::string::npos && pos + 1 < m.size()) { + m = m.substr(pos + 1); + } + return m; +} + +// Forward declarations for helpers used before definition +static std::string event_type_to_string(EventType t); +static bool event_type_from_string(const std::string& s, EventType& t); +static bool extract_string_field(const std::string& line, const std::string& key, std::string& out); +static bool extract_json_field(const std::string& line, const std::string& key, std::string& out); +static bool extract_bool_field(const std::string& line, const std::string& key, bool& out); +static bool extract_double_field(const std::string& line, const std::string& key, double& out); +static bool extract_int_field(const std::string& line, const std::string& key, int& out); +static bool extract_double_field_raw(const std::string& line, const std::string& key, double& out); +static void process_events(const std::vector& fresh_events); +static std::string get_telemetry_dir_locked(); +static CloudConfigurationStateSnapshot capture_cloud_configuration_state_snapshot_locked(); + +class TelemetryDispatcher { +public: + TelemetryDispatcher() = default; + + static TelemetryDispatcher& instance() { + static std::once_flag once; + static TelemetryDispatcher* singleton = nullptr; + std::call_once(once, [] { + singleton = new TelemetryDispatcher(); + }); + return *singleton; + } + + void start() { + std::lock_guard lock(queue_mutex_); + if (running_) return; + stop_ = false; + processing_ = false; + running_ = true; + worker_thread_ = std::thread([this] { worker_loop(); }); + } + + bool enqueue(const Event& event) { + std::lock_guard lock(queue_mutex_); + if (!running_) return false; + pending_events_.push_back(event); + enqueue_seq_ += 1; + queue_cv_.notify_one(); + return true; + } + + void flush() { + uint64_t target = 0; + { + std::unique_lock lock(queue_mutex_); + if (running_) { + target = enqueue_seq_; + queue_cv_.notify_one(); + flush_cv_.wait(lock, [this, target] { + return processed_seq_ >= target && !processing_; + }); + } + } + process_with_io_lock({}); + } + + void stop() { + { + std::lock_guard lock(queue_mutex_); + stop_ = true; + queue_cv_.notify_all(); + } + + if (worker_thread_.joinable()) { + worker_thread_.join(); + } + + std::lock_guard lock(queue_mutex_); + running_ = false; + stop_ = false; + processing_ = false; + pending_events_.clear(); + enqueue_seq_ = 0; + processed_seq_ = 0; + } + +private: + void process_with_io_lock(const std::vector& batch) { + std::lock_guard io_guard(io_mutex_); + process_events(batch); + } + + void worker_loop() { + while (true) { + std::vector batch; + uint64_t processed_count = 0; + { + std::unique_lock lock(queue_mutex_); + queue_cv_.wait(lock, [this] { + return stop_ || !pending_events_.empty(); + }); + if (stop_ && pending_events_.empty()) { + break; + } + batch.assign(pending_events_.begin(), pending_events_.end()); + processed_count = static_cast(batch.size()); + pending_events_.clear(); + processing_ = true; + } + + process_with_io_lock(batch); + + { + std::lock_guard lock(queue_mutex_); + processed_seq_ += processed_count; + processing_ = false; + flush_cv_.notify_all(); + } + } + + std::lock_guard lock(queue_mutex_); + running_ = false; + processing_ = false; + flush_cv_.notify_all(); + } + std::mutex queue_mutex_; + std::mutex io_mutex_; + std::condition_variable queue_cv_; + std::condition_variable flush_cv_; + std::deque pending_events_; + std::thread worker_thread_; + bool stop_ = false; + bool running_ = false; + bool processing_ = false; + uint64_t enqueue_seq_ = 0; + uint64_t processed_seq_ = 0; +}; + +static void mkdir_p(const std::string& path) { + if (path.empty()) return; + + std::string current; + for (size_t i = 0; i < path.size(); ++i) { + current += path[i]; + if (path[i] == '/' && i > 0) { + mkdir(current.c_str(), 0755); + } + } + mkdir(path.c_str(), 0755); +} + +static std::string get_telemetry_dir_locked() { + const std::string& cache_location = custom_cache_location; + if (!cache_location.empty()) { + mkdir_p(cache_location); + return cache_location; + } + +#if defined(__ANDROID__) + { + std::ifstream f("/proc/self/cmdline"); + if (f.is_open()) { + std::string pkg; + std::getline(f, pkg, '\0'); + size_t colon = pkg.find(':'); + if (colon != std::string::npos) pkg = pkg.substr(0, colon); + if (!pkg.empty()) { + uid_t uid = getuid(); + int user_id = static_cast((uid - 10000) / 100000); + std::string dir = "/data/user/" + std::to_string(user_id) + "/" + pkg + "/cache/cactus/telemetry"; + mkdir_p(dir); + struct stat st; + if (stat(dir.c_str(), &st) == 0) return dir; + dir = "/data/data/" + pkg + "/cache/cactus/telemetry"; + mkdir_p(dir); + if (stat(dir.c_str(), &st) == 0) return dir; + } + } + } +#endif + + const char* home = getenv("HOME"); + if (!home) home = "/tmp"; + std::string dir = std::string(home) + "/Library/Caches/cactus/telemetry"; + mkdir((std::string(home) + "/Library").c_str(), 0755); + mkdir((std::string(home) + "/Library/Caches").c_str(), 0755); + mkdir((std::string(home) + "/Library/Caches/cactus").c_str(), 0755); + mkdir(dir.c_str(), 0755); + return dir; +} + +static std::string cloud_api_key_cache_file_locked() { + return get_telemetry_dir_locked() + "/cloud_api_key"; +} + +static std::string load_or_create_id(const std::string& file) { + std::ifstream in(file); + if (in.is_open()) { + std::string line; + if (std::getline(in, line) && !line.empty()) { + return line; + } + } + std::string id = new_uuid(); + std::ofstream out(file, std::ios::trunc); + if (out.is_open()) { + out << id; + } + return id; +} + +static bool load_registered_flag(const std::string& file) { + std::ifstream in(file); + if (!in.is_open()) return false; + std::string line; + if (std::getline(in, line) && !line.empty()) { + return line[0] == '1'; + } + return false; +} + +static void persist_registered_flag(const std::string& file) { + std::ofstream out(file, std::ios::trunc); + if (out.is_open()) { + out << "1"; + } +} + +#if defined(__APPLE__) +static std::string sysctl_string(const char* key) { + size_t size = 0; + if (sysctlbyname(key, nullptr, &size, nullptr, 0) != 0 || size == 0) return {}; + std::string out(size, '\0'); + if (sysctlbyname(key, out.data(), &size, nullptr, 0) != 0) return {}; + if (!out.empty() && out.back() == '\0') out.pop_back(); + return out; +} +#endif + +static Event make_event(EventType type, const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, const char* message) { + Event e; + e.type = type; + e.success = success; + e.cloud_handoff = false; + e.ttft_ms = ttft_ms; + e.prefill_tps = 0.0; + e.decode_tps = tps; + e.tps = tps; + e.response_time_ms = response_time_ms; + e.confidence = 0.0; + e.ram_usage_mb = 0.0; + e.tokens = tokens; + e.prefill_tokens = 0; + e.decode_tokens = tokens; + e.session_ttft_ms = 0.0; + e.session_tps = 0.0; + e.session_time_ms = 0.0; + e.session_tokens = 0; + e.timestamp = std::chrono::system_clock::now(); + std::memset(e.model, 0, sizeof(e.model)); + std::memset(e.message, 0, sizeof(e.message)); + std::memset(e.error, 0, sizeof(e.error)); + std::memset(e.function_calls, 0, sizeof(e.function_calls)); + std::string safe_model = model_basename(model); + if (!safe_model.empty()) std::strncpy(e.model, safe_model.c_str(), sizeof(e.model)-1); + if (message) std::strncpy(e.message, message, sizeof(e.message)-1); + return e; +} + +static Event make_event_extended(EventType type, const char* model, const CompletionMetrics& metrics) { + Event e; + e.type = type; + e.success = metrics.success; + e.cloud_handoff = metrics.cloud_handoff; + e.ttft_ms = metrics.ttft_ms; + e.prefill_tps = metrics.prefill_tps; + e.decode_tps = metrics.decode_tps; + e.tps = metrics.decode_tps; + e.response_time_ms = metrics.response_time_ms; + e.confidence = metrics.confidence; + e.ram_usage_mb = metrics.ram_usage_mb; + e.tokens = static_cast(metrics.prefill_tokens + metrics.decode_tokens); + e.prefill_tokens = static_cast(metrics.prefill_tokens); + e.decode_tokens = static_cast(metrics.decode_tokens); + e.session_ttft_ms = 0.0; + e.session_tps = 0.0; + e.session_time_ms = 0.0; + e.session_tokens = 0; + e.timestamp = std::chrono::system_clock::now(); + std::memset(e.model, 0, sizeof(e.model)); + std::memset(e.message, 0, sizeof(e.message)); + std::memset(e.error, 0, sizeof(e.error)); + std::memset(e.function_calls, 0, sizeof(e.function_calls)); + std::string safe_model = model_basename(model); + if (!safe_model.empty()) std::strncpy(e.model, safe_model.c_str(), sizeof(e.model)-1); + if (!metrics.success && metrics.error_message) std::strncpy(e.error, metrics.error_message, sizeof(e.error)-1); + if (metrics.function_calls_json) std::strncpy(e.function_calls, metrics.function_calls_json, sizeof(e.function_calls)-1); + return e; +} + +static bool parse_event_line(const std::string& line, Event& out) { + std::string type_str; + if (!extract_string_field(line, "event_type", type_str)) return false; + EventType et; + if (!event_type_from_string(type_str, et)) return false; + + std::string model; + extract_string_field(line, "model", model); + bool success = false; + extract_bool_field(line, "success", success); + bool cloud_handoff = false; + extract_bool_field(line, "cloud_handoff", cloud_handoff); + double ttft = 0.0; + extract_double_field(line, "ttft", ttft); + double prefill_tps = 0.0; + extract_double_field(line, "prefill_tps", prefill_tps); + double decode_tps = 0.0; + extract_double_field(line, "decode_tps", decode_tps); + double tps = 0.0; + extract_double_field(line, "tps", tps); + if (decode_tps == 0.0 && tps != 0.0) decode_tps = tps; + double response_time = 0.0; + extract_double_field(line, "response_time", response_time); + double confidence = 0.0; + extract_double_field(line, "confidence", confidence); + double ram_usage_mb = 0.0; + extract_double_field(line, "ram_usage_mb", ram_usage_mb); + int tokens = 0; + extract_int_field(line, "tokens", tokens); + int prefill_tokens = 0; + extract_int_field(line, "prefill_tokens", prefill_tokens); + int decode_tokens = 0; + extract_int_field(line, "decode_tokens", decode_tokens); + double session_ttft_ms = 0.0; + extract_double_field(line, "session_ttft", session_ttft_ms); + double session_tps = 0.0; + extract_double_field(line, "session_tps", session_tps); + double session_time_ms = 0.0; + extract_double_field(line, "session_time_ms", session_time_ms); + int session_tokens = 0; + extract_int_field(line, "session_tokens", session_tokens); + std::string message; + extract_string_field(line, "message", message); + std::string error; + extract_string_field(line, "error", error); + std::string function_calls; + extract_json_field(line, "function_calls", function_calls); + double ts_ms = 0.0; + if (!extract_double_field_raw(line, "ts_ms", ts_ms)) { + extract_double_field(line, "ts_ms", ts_ms); + } + auto ts_point = std::chrono::system_clock::now(); + if (ts_ms > 0.0) { + ts_point = std::chrono::system_clock::time_point(std::chrono::milliseconds(static_cast(ts_ms))); + } + + out = make_event(et, + model.empty() ? nullptr : model.c_str(), + success, + ttft, + tps, + response_time, + tokens, + message.empty() ? nullptr : message.c_str()); + out.timestamp = ts_point; + out.cloud_handoff = cloud_handoff; + out.prefill_tps = prefill_tps; + out.decode_tps = decode_tps; + out.confidence = confidence; + out.ram_usage_mb = ram_usage_mb; + out.prefill_tokens = prefill_tokens; + out.decode_tokens = decode_tokens; + out.session_ttft_ms = session_ttft_ms; + out.session_tps = session_tps; + out.session_time_ms = session_time_ms; + out.session_tokens = session_tokens; + if (!error.empty()) std::strncpy(out.error, error.c_str(), sizeof(out.error)-1); + if (!function_calls.empty()) std::strncpy(out.function_calls, function_calls.c_str(), sizeof(out.function_calls)-1); + return true; +} + +static std::string event_type_to_string(EventType t) { + switch (t) { + case INIT: return "init"; + case COMPLETION: return "completion"; + case EMBEDDING: return "embedding"; + case TRANSCRIPTION: return "transcription"; + case STREAM_TRANSCRIBE: return "stream_transcribe"; + default: return "unknown"; + } +} + +static bool event_type_from_string(const std::string& s, EventType& t) { + if (s == "init") { t = INIT; return true; } + if (s == "completion") { t = COMPLETION; return true; } + if (s == "embedding") { t = EMBEDDING; return true; } + if (s == "transcription") { t = TRANSCRIPTION; return true; } + if (s == "stream_transcribe") { t = STREAM_TRANSCRIBE; return true; } + return false; +} + +static bool extract_string_field(const std::string& line, const std::string& key, std::string& out) { + std::string needle = "\"" + key + "\":"; + size_t pos = line.find(needle); + if (pos == std::string::npos) return false; + pos += needle.size(); + while (pos < line.size() && line[pos] == ' ') pos++; + if (pos >= line.size()) return false; + if (line[pos] == '"') { + pos++; + size_t end = line.find('"', pos); + if (end == std::string::npos) return false; + out = line.substr(pos, end - pos); + return true; + } + size_t end = line.find_first_of(",}", pos); + if (end == std::string::npos) end = line.size(); + out = line.substr(pos, end - pos); + if (out == "null") out.clear(); + return true; +} + +static bool extract_json_field(const std::string& line, const std::string& key, std::string& out) { + std::string needle = "\"" + key + "\":"; + size_t pos = line.find(needle); + if (pos == std::string::npos) return false; + pos += needle.size(); + while (pos < line.size() && line[pos] == ' ') pos++; + if (pos >= line.size()) return false; + + char start_char = line[pos]; + if (start_char == '[' || start_char == '{') { + char end_char = (start_char == '[') ? ']' : '}'; + int depth = 0; + size_t end = pos; + bool in_string = false; + bool escape_next = false; + + for (; end < line.size(); end++) { + if (escape_next) { + escape_next = false; + continue; + } + if (line[end] == '\\') { + escape_next = true; + continue; + } + if (line[end] == '"') { + in_string = !in_string; + continue; + } + if (in_string) continue; + + if (line[end] == start_char) { + depth++; + } else if (line[end] == end_char) { + depth--; + if (depth == 0) { + out = line.substr(pos, end - pos + 1); + return true; + } + } + } + return false; + } + + size_t end = line.find_first_of(",}", pos); + if (end == std::string::npos) end = line.size(); + out = line.substr(pos, end - pos); + if (out == "null") out.clear(); + return true; +} + +static bool extract_bool_field(const std::string& line, const std::string& key, bool& out) { + std::string raw; + if (!extract_string_field(line, key, raw)) return false; + if (raw == "true") { out = true; return true; } + if (raw == "false") { out = false; return true; } + return false; +} + +static bool extract_double_field(const std::string& line, const std::string& key, double& out) { + std::string raw; + if (!extract_string_field(line, key, raw)) return false; + try { + out = std::stod(raw); + return true; + } catch (...) { + return false; + } +} + +// Helper used when the field is already numeric (no quotes) in the log cache. +static bool extract_double_field_raw(const std::string& line, const std::string& key, double& out) { + std::string needle = "\"" + key + "\":"; + size_t pos = line.find(needle); + if (pos == std::string::npos) return false; + pos += needle.size(); + while (pos < line.size() && line[pos] == ' ') pos++; + if (pos >= line.size()) return false; + size_t end = line.find_first_of(",}", pos); + if (end == std::string::npos) end = line.size(); + std::string raw = line.substr(pos, end - pos); + try { + out = std::stod(raw); + return true; + } catch (...) { + return false; + } +} + +static bool extract_int_field(const std::string& line, const std::string& key, int& out) { + std::string raw; + if (!extract_string_field(line, key, raw)) return false; + try { + out = std::stoi(raw); + return true; + } catch (...) { + return false; + } +} + +static std::string new_uuid() { + static thread_local std::mt19937_64 rng(std::random_device{}()); + uint64_t a = rng(); + uint64_t b = rng(); + a = (a & 0xffffffffffff0fffULL) | 0x0000000000004000ULL; + b = (b & 0x3fffffffffffffffULL) | 0x8000000000000000ULL; + std::ostringstream oss; + oss << std::hex << std::setfill('0'); + oss << std::setw(8) << ((a >> 32) & 0xffffffffULL); + oss << "-" << std::setw(4) << ((a >> 16) & 0xffffULL); + oss << "-" << std::setw(4) << (a & 0xffffULL); + oss << "-" << std::setw(4) << ((b >> 48) & 0xffffULL); + oss << "-" << std::setw(12) << (b & 0xffffffffffffULL); + return oss.str(); +} + +static uint64_t fnv1a_64(const std::string& s) { + uint64_t h = 14695981039346656037ULL; + for (unsigned char c : s) { h ^= c; h *= 1099511628211ULL; } + return h; +} + +static std::string uuid_from_string(const std::string& input) { + uint64_t a = fnv1a_64(input); + uint64_t b = fnv1a_64(input + "\xff"); + a = (a & 0xffffffffffff0fffULL) | 0x0000000000004000ULL; + b = (b & 0x3fffffffffffffffULL) | 0x8000000000000000ULL; + std::ostringstream oss; + oss << std::hex << std::setfill('0'); + oss << std::setw(8) << ((a >> 32) & 0xffffffffULL); + oss << "-" << std::setw(4) << ((a >> 16) & 0xffffULL); + oss << "-" << std::setw(4) << (a & 0xffffULL); + oss << "-" << std::setw(4) << ((b >> 48) & 0xffffULL); + oss << "-" << std::setw(12) << (b & 0xffffffffffffULL); + return oss.str(); +} + +static std::string read_git_remote_url() { + const char* cwd_env = std::getenv("PWD"); + if (!cwd_env) return ""; + std::string dir = cwd_env; + while (!dir.empty()) { + std::ifstream f(dir + "/.git/config"); + if (f.is_open()) { + std::string line; + bool in_origin = false; + while (std::getline(f, line)) { + if (line.find("[remote \"origin\"]") != std::string::npos) { + in_origin = true; + } else if (in_origin) { + size_t pos = line.find("url = "); + if (pos != std::string::npos) return line.substr(pos + 6); + if (!line.empty() && line[0] == '[') break; + } + } + return ""; + } + size_t pos = dir.rfind('/'); + if (pos == 0 || pos == std::string::npos) break; + dir = dir.substr(0, pos); + } + return ""; +} + +#if defined(__APPLE__) +static std::string detect_app_id() { + CFBundleRef bundle = CFBundleGetMainBundle(); + if (!bundle) return ""; + CFStringRef id = (CFStringRef)CFBundleGetValueForInfoDictionaryKey(bundle, CFSTR("CFBundleIdentifier")); + if (!id) return ""; + char buf[256]; + return CFStringGetCString(id, buf, sizeof(buf), kCFStringEncodingUTF8) ? buf : ""; +} +#elif defined(__linux__) || defined(__ANDROID__) +static std::string detect_app_id() { + std::ifstream f("/proc/self/cmdline"); + if (!f.is_open()) return ""; + std::string s; + std::getline(f, s, '\0'); + size_t colon = s.find(':'); + if (colon != std::string::npos) s = s.substr(0, colon); + return s; +} +#else +static std::string detect_app_id() { return ""; } +#endif + +static std::string extract_android_package_from_path(const std::string& path) { + for (const char* prefix : {"/data/user/", "/data/user_de/", "/data/data/"}) { + size_t pos = path.find(prefix); + if (pos == std::string::npos) continue; + size_t start = pos + strlen(prefix); + if (strcmp(prefix, "/data/data/") != 0) { + start = path.find('/', start); + if (start == std::string::npos) continue; + start++; + } + size_t end = path.find('/', start); + if (end == std::string::npos) end = path.size(); + std::string pkg = path.substr(start, end - start); + if (!pkg.empty()) return pkg; + } + return ""; +} + +static std::string format_timestamp(const std::chrono::system_clock::time_point& tp) { + using namespace std::chrono; + auto secs = time_point_cast(tp); + auto ms = duration_cast(tp - secs).count(); + std::time_t t = system_clock::to_time_t(secs); + std::tm tm{}; +#if defined(_WIN32) + gmtime_s(&tm, &t); +#else + gmtime_r(&t, &tm); +#endif + char buf[32]; + std::snprintf(buf, sizeof(buf), "%04d-%02d-%02dT%02d:%02d:%02d.%03lldZ", + tm.tm_year + 1900, tm.tm_mon + 1, tm.tm_mday, + tm.tm_hour, tm.tm_min, tm.tm_sec, + static_cast(ms)); + return std::string(buf); +} + +static void collect_device_info() { + struct utsname u; + if (uname(&u) == 0) { + const std::string sysname = u.sysname; + device_os = (sysname == "Darwin") ? "macos" : sysname; + device_model = u.machine; + device_os_version = u.release; + device_brand = "unknown"; +#if defined(__APPLE__) + // Prefer user-facing macOS info over Darwin kernel version. + std::string hw_model = sysctl_string("hw.model"); + std::string os_product = sysctl_string("kern.osproductversion"); + if (!hw_model.empty()) device_model = hw_model; + if (!os_product.empty()) device_os_version = os_product; + device_brand = "apple"; +#elif defined(__ANDROID__) + device_brand = "android"; +#endif + } +} + +static bool parse_version_line(const std::string& line) { + size_t start = line.find_first_not_of(" \t\r\n"); + size_t end = line.find_last_not_of(" \t\r\n"); + if (start != std::string::npos && end != std::string::npos) { + cactus_version = line.substr(start, end - start + 1); + return true; + } + return false; +} + +static void read_cactus_version() { + const char* version_paths[] = { + "CACTUS_VERSION", + "../CACTUS_VERSION", + "../../CACTUS_VERSION", + "../../../CACTUS_VERSION" + }; + + for (const char* path : version_paths) { + std::ifstream file(path); + if (file.is_open()) { + std::string line; + if (std::getline(file, line) && parse_version_line(line)) return; + } + } + +#if defined(__APPLE__) + CFBundleRef bundle = CFBundleGetMainBundle(); + if (bundle) { + CFURLRef url = CFBundleCopyResourceURL(bundle, CFSTR("CACTUS_VERSION"), NULL, NULL); + if (url) { + char fspath[4096]; + if (CFURLGetFileSystemRepresentation(url, true, (UInt8*)fspath, sizeof(fspath))) { + std::ifstream file(fspath); + if (file.is_open()) { + std::string line; + if (std::getline(file, line) && parse_version_line(line)) { + CFRelease(url); + return; + } + } + } + CFRelease(url); + } + } +#endif + +#ifdef CACTUS_COMPILE_TIME_VERSION + cactus_version = CACTUS_COMPILE_TIME_VERSION; + return; +#endif + + cactus_version = ""; +} + +static void apply_curl_tls_trust(CURL* curl) { + if (!curl) return; + const char* ca_bundle = std::getenv("CACTUS_CA_BUNDLE"); + if (ca_bundle && ca_bundle[0] != '\0') { + curl_easy_setopt(curl, CURLOPT_CAINFO, ca_bundle); + } +#if defined(__ANDROID__) + const char* ca_path = std::getenv("CACTUS_CA_PATH"); + if (ca_path && ca_path[0] != '\0') { + curl_easy_setopt(curl, CURLOPT_CAPATH, ca_path); + } else { + curl_easy_setopt(curl, CURLOPT_CAPATH, "/system/etc/security/cacerts"); + } +#endif +} + +static bool ensure_project_row_remote(CURL* curl, const CloudConfigurationStateSnapshot& snapshot) { + if (snapshot.project_id.empty() || snapshot.project_registered) return true; + std::string url = snapshot.supabase_url + "/rest/v1/projects"; + std::ostringstream payload; + payload << "[{"; + payload << "\"project_key\":\"" << escape_json_string(snapshot.project_id) << "\""; + if (!snapshot.project_scope.empty()) { + payload << ",\"name\":\"" << escape_json_string(snapshot.project_scope) << "\""; + } + payload << "}]"; + struct curl_slist* headers = nullptr; + headers = curl_slist_append(headers, "Content-Type: application/json"); + std::string apikey_hdr = std::string("apikey: ") + snapshot.supabase_key; + std::string auth_hdr = std::string("Authorization: Bearer ") + snapshot.supabase_key; + headers = curl_slist_append(headers, apikey_hdr.c_str()); + headers = curl_slist_append(headers, auth_hdr.c_str()); + headers = curl_slist_append(headers, "Prefer: resolution=ignore-duplicates"); + headers = curl_slist_append(headers, "Content-Profile: cactus"); + headers = curl_slist_append(headers, "Accept-Profile: cactus"); + curl_easy_setopt(curl, CURLOPT_URL, url.c_str()); + curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers); + curl_easy_setopt(curl, CURLOPT_POST, 1L); + std::string body = payload.str(); + curl_easy_setopt(curl, CURLOPT_POSTFIELDS, body.c_str()); + curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, body.size()); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L); + curl_easy_setopt(curl, CURLOPT_TIMEOUT, 5L); + apply_curl_tls_trust(curl); + CURLcode res = curl_easy_perform(curl); + long code = 0; + curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &code); + if (headers) curl_slist_free_all(headers); + bool ok = (res == CURLE_OK) && ((code >= 200 && code < 300) || code == 409); + return ok; +} + +static bool ensure_device_row_remote(CURL* curl, const CloudConfigurationStateSnapshot& snapshot) { + if (snapshot.device_registered) return true; + if (snapshot.device_id.empty()) return false; + std::string url = snapshot.supabase_url + "/rest/v1/devices"; + std::ostringstream payload; + payload << "[{"; + payload << "\"id\":\"" << escape_json_string(snapshot.device_id) << "\""; + payload << ",\"device_id\":\"" << escape_json_string(snapshot.device_id) << "\""; + if (!snapshot.device_model.empty()) payload << ",\"model\":\"" << escape_json_string(snapshot.device_model) << "\""; + if (!snapshot.device_os.empty()) payload << ",\"os\":\"" << escape_json_string(snapshot.device_os) << "\""; + if (!snapshot.device_os_version.empty()) payload << ",\"os_version\":\"" << escape_json_string(snapshot.device_os_version) << "\""; + if (!snapshot.device_brand.empty()) payload << ",\"brand\":\"" << escape_json_string(snapshot.device_brand) << "\""; + payload << "}]"; + struct curl_slist* headers = nullptr; + headers = curl_slist_append(headers, "Content-Type: application/json"); + std::string apikey_hdr = std::string("apikey: ") + snapshot.supabase_key; + std::string auth_hdr = std::string("Authorization: Bearer ") + snapshot.supabase_key; + headers = curl_slist_append(headers, apikey_hdr.c_str()); + headers = curl_slist_append(headers, auth_hdr.c_str()); + headers = curl_slist_append(headers, "Prefer: resolution=ignore-duplicates"); + headers = curl_slist_append(headers, "Content-Profile: cactus"); + headers = curl_slist_append(headers, "Accept-Profile: cactus"); + curl_easy_setopt(curl, CURLOPT_URL, url.c_str()); + curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers); + curl_easy_setopt(curl, CURLOPT_POST, 1L); + std::string body = payload.str(); + curl_easy_setopt(curl, CURLOPT_POSTFIELDS, body.c_str()); + curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, body.size()); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L); + curl_easy_setopt(curl, CURLOPT_TIMEOUT, 5L); + apply_curl_tls_trust(curl); + CURLcode res = curl_easy_perform(curl); + long code = 0; + curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &code); + if (headers) curl_slist_free_all(headers); + bool ok = (res == CURLE_OK) && ((code >= 200 && code < 300) || code == 409); + return ok; +} + +static bool send_payload(CURL* curl, const std::string& url, const std::string& body, const CloudConfigurationStateSnapshot& snapshot) { + struct curl_slist* headers = nullptr; + headers = curl_slist_append(headers, "Content-Type: application/json"); + std::string apikey_hdr = std::string("apikey: ") + snapshot.supabase_key; + std::string auth_hdr = std::string("Authorization: Bearer ") + snapshot.supabase_key; + headers = curl_slist_append(headers, apikey_hdr.c_str()); + headers = curl_slist_append(headers, auth_hdr.c_str()); + headers = curl_slist_append(headers, "Prefer: return=minimal"); + headers = curl_slist_append(headers, "Content-Profile: cactus"); + headers = curl_slist_append(headers, "Accept-Profile: cactus"); + curl_easy_setopt(curl, CURLOPT_URL, url.c_str()); + curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers); + curl_easy_setopt(curl, CURLOPT_POST, 1L); + curl_easy_setopt(curl, CURLOPT_POSTFIELDS, body.c_str()); + curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, body.size()); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L); + curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L); + curl_easy_setopt(curl, CURLOPT_TIMEOUT, 5L); + apply_curl_tls_trust(curl); + CURLcode res = curl_easy_perform(curl); + long code = 0; + curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &code); + if (headers) curl_slist_free_all(headers); + return (res == CURLE_OK) && (code >= 200 && code < 300); +} + +static CloudSendResult send_batch_to_cloud(const std::vector& local, const CloudConfigurationStateSnapshot& snapshot) { + CloudSendResult result; + if (!snapshot.enabled) return result; + if (local.empty()) { + result.payload_ok = true; + return result; + } + if (snapshot.cloud_disabled) return result; + if (snapshot.device_id.empty()) return result; + CURL* curl = curl_easy_init(); + if (!curl) return result; + result.project_registered_ok = ensure_project_row_remote(curl, snapshot); + if (!snapshot.device_registered) { + result.device_registered_ok = ensure_device_row_remote(curl, snapshot); + } else { + result.device_registered_ok = true; + } + std::string url = snapshot.supabase_url + "/rest/v1/logs"; + std::ostringstream payload; + payload << "["; + for (size_t i = 0; i < local.size(); ++i) { + const Event& e = local[i]; + payload << "{"; + payload << "\"event_type\":\"" << event_type_to_string(e.type) << "\","; + payload << "\"model\":\"" << escape_json_string(e.model) << "\","; + payload << "\"success\":" << (e.success ? "true" : "false") << ","; + payload << "\"cloud_handoff\":" << (e.cloud_handoff ? "true" : "false") << ","; + if (e.type == INIT || !e.success) { + payload << "\"ttft\":null,"; + payload << "\"prefill_tps\":null,"; + payload << "\"decode_tps\":null,"; + payload << "\"tps\":null,"; + } else { + payload << "\"ttft\":" << e.ttft_ms << ","; + payload << "\"prefill_tps\":" << e.prefill_tps << ","; + payload << "\"decode_tps\":" << e.decode_tps << ","; + payload << "\"tps\":" << e.tps << ","; + } + if (!e.success) { + payload << "\"response_time\":null,"; + payload << "\"ram_usage_mb\":null,"; + } else { + payload << "\"response_time\":" << e.response_time_ms << ","; + payload << "\"ram_usage_mb\":" << e.ram_usage_mb << ","; + } + if (e.type == INIT || !e.success) { + payload << "\"confidence\":null,"; + payload << "\"tokens\":null,"; + payload << "\"prefill_tokens\":null,"; + payload << "\"decode_tokens\":null,"; + } else { + payload << "\"confidence\":" << e.confidence << ","; + payload << "\"tokens\":" << e.tokens << ","; + payload << "\"prefill_tokens\":" << e.prefill_tokens << ","; + payload << "\"decode_tokens\":" << e.decode_tokens << ","; + } + if (e.type == STREAM_TRANSCRIBE) { + payload << "\"session_ttft\":" << e.session_ttft_ms << ","; + payload << "\"session_tps\":" << e.session_tps << ","; + payload << "\"session_time_ms\":" << e.session_time_ms << ","; + payload << "\"session_tokens\":" << e.session_tokens << ","; + } else { + payload << "\"session_ttft\":null,"; + payload << "\"session_tps\":null,"; + payload << "\"session_time_ms\":null,"; + payload << "\"session_tokens\":null,"; + } + payload << "\"created_at\":\"" << format_timestamp(e.timestamp) << "\","; + if (!snapshot.project_id.empty()) { + payload << "\"project_id\":\"" << escape_json_string(snapshot.project_id) << "\","; + } + if (!snapshot.cloud_key.empty()) { + payload << "\"key_hash\":\"" << escape_json_string(snapshot.cloud_key) << "\","; + } + payload << "\"framework\":\"" << escape_json_string(snapshot.framework) << "\","; + if (!snapshot.cactus_version.empty()) { + payload << "\"framework_version\":\"" << escape_json_string(snapshot.cactus_version) << "\","; + } + payload << "\"device_id\":\"" << escape_json_string(snapshot.device_id) << "\""; + if (!snapshot.app_id.empty()) { + payload << ",\"app_id\":\"" << escape_json_string(snapshot.app_id.c_str()) << "\""; + } else { + payload << ",\"app_id\":null"; + } + if (e.message[0] != '\0') { + payload << ",\"message\":\"" << escape_json_string(e.message) << "\""; + } else { + payload << ",\"message\":null"; + } + if (e.error[0] != '\0') { + payload << ",\"error\":\"" << escape_json_string(e.error) << "\""; + } else { + payload << ",\"error\":null"; + } + payload << ",\"function_calls\":null"; + payload << "}"; + if (i + 1 < local.size()) payload << ","; + } + payload << "]"; + result.payload_ok = send_payload(curl, url, payload.str(), snapshot); + curl_easy_cleanup(curl); + return result; +} + +static const size_t kMaxCachedEventsPerFile = 1000; + +static void trim_cache_file(const std::string& file, size_t max_lines) { + // Stream the file keeping only the last max_lines, so a pre-existing oversized cache cannot blow up memory. + std::deque lines; + bool trimmed = false; + { + std::ifstream in(file); + if (!in.is_open()) return; + std::string line; + while (std::getline(in, line)) { + lines.push_back(std::move(line)); + if (lines.size() > max_lines) { + lines.pop_front(); + trimmed = true; + } + } + } + if (!trimmed) return; + std::ofstream out(file, std::ios::trunc); + if (!out.is_open()) return; + for (const auto& line : lines) out << line << "\n"; +} + +static void write_events_to_cache_in_dir(const std::vector& local, const std::string& dir) { + std::vector touched_files; + for (const auto &e : local) { + std::ostringstream oss; + oss << "{\"event_type\":\"" << event_type_to_string(e.type) << "\","; + oss << "\"model\":\"" << escape_json_string(e.model) << "\","; + oss << "\"success\":" << (e.success ? "true" : "false") << ","; + oss << "\"cloud_handoff\":" << (e.cloud_handoff ? "true" : "false") << ","; + if (e.type == INIT || !e.success) { + oss << "\"ttft\":null,"; + oss << "\"prefill_tps\":null,"; + oss << "\"decode_tps\":null,"; + oss << "\"tps\":null,"; + } else { + oss << "\"ttft\":" << e.ttft_ms << ","; + oss << "\"prefill_tps\":" << e.prefill_tps << ","; + oss << "\"decode_tps\":" << e.decode_tps << ","; + oss << "\"tps\":" << e.tps << ","; + } + if (!e.success) { + oss << "\"response_time\":null,"; + oss << "\"ram_usage_mb\":null,"; + } else { + oss << "\"response_time\":" << e.response_time_ms << ","; + oss << "\"ram_usage_mb\":" << e.ram_usage_mb << ","; + } + if (e.type == INIT || !e.success) { + oss << "\"confidence\":null,"; + oss << "\"tokens\":null,"; + oss << "\"prefill_tokens\":null,"; + oss << "\"decode_tokens\":null,"; + } else { + oss << "\"confidence\":" << e.confidence << ","; + oss << "\"tokens\":" << e.tokens << ","; + oss << "\"prefill_tokens\":" << e.prefill_tokens << ","; + oss << "\"decode_tokens\":" << e.decode_tokens << ","; + } + if (e.type == STREAM_TRANSCRIBE) { + oss << "\"session_ttft\":" << e.session_ttft_ms << ","; + oss << "\"session_tps\":" << e.session_tps << ","; + oss << "\"session_time_ms\":" << e.session_time_ms << ","; + oss << "\"session_tokens\":" << e.session_tokens; + } else { + oss << "\"session_ttft\":null,"; + oss << "\"session_tps\":null,"; + oss << "\"session_time_ms\":null,"; + oss << "\"session_tokens\":null"; + } + oss << ",\"ts_ms\":" << std::chrono::duration_cast(e.timestamp.time_since_epoch()).count(); + if (e.message[0] != '\0') { + oss << ",\"message\":\"" << escape_json_string(e.message) << "\""; + } else { + oss << ",\"message\":null"; + } + if (e.error[0] != '\0') { + oss << ",\"error\":\"" << escape_json_string(e.error) << "\""; + } else { + oss << ",\"error\":null"; + } + oss << ",\"function_calls\":null"; + if (!app_id.empty()) { + oss << ",\"app_id\":\"" << escape_json_string(app_id.c_str()) << "\""; + } else { + oss << ",\"app_id\":null"; + } + oss << "}"; + std::string file = dir + "/" + event_type_to_string(e.type) + ".log"; + std::ofstream out(file, std::ios::app); + if (out.is_open()) { + out << oss.str() << "\n"; + out.close(); + } + bool seen = false; + for (const auto& f : touched_files) { if (f == file) { seen = true; break; } } + if (!seen) touched_files.push_back(file); + } + for (const auto& file : touched_files) { + trim_cache_file(file, kMaxCachedEventsPerFile); + } +} + +static std::vector load_cached_events_in_dir(const std::string& dir) { + std::vector events; + DIR* d = opendir(dir.c_str()); + if (!d) return events; + struct dirent* ent; + while ((ent = readdir(d)) != nullptr) { + std::string name = ent->d_name; + if (name.size() < 5 || name.substr(name.size() - 4) != ".log") continue; + std::ifstream in(dir + "/" + name); + std::string line; + // Load only the last kMaxCachedEventsPerFile lines, matching trim, so an oversized file is bounded here too. + std::deque tail; + while (std::getline(in, line)) { + tail.push_back(std::move(line)); + if (tail.size() > kMaxCachedEventsPerFile) tail.pop_front(); + } + for (const auto& cached : tail) { + Event e; + if (parse_event_line(cached, e)) { + events.push_back(e); + } + } + } + closedir(d); + return events; +} + +static void clear_cache_files_in_dir(const std::string& dir) { + DIR* d = opendir(dir.c_str()); + if (!d) return; + struct dirent* ent; + while ((ent = readdir(d)) != nullptr) { + std::string name = ent->d_name; + if (name.size() < 5 || name.substr(name.size() - 4) != ".log") continue; + std::remove((dir + "/" + name).c_str()); + } + closedir(d); +} + +static CloudConfigurationStateSnapshot capture_cloud_configuration_state_snapshot_locked() { + CloudConfigurationStateSnapshot snapshot; + snapshot.supabase_url = supabase_url; + snapshot.supabase_key = supabase_key; + snapshot.device_id = device_id; + snapshot.project_id = project_id; + snapshot.cloud_key = cloud_key; + snapshot.project_scope = project_scope; + snapshot.framework = framework; + snapshot.cactus_version = cactus_version; + snapshot.device_model = device_model; + snapshot.device_os = device_os; + snapshot.device_os_version = device_os_version; + snapshot.device_brand = device_brand; + snapshot.app_id = app_id; + snapshot.device_registered_file = device_registered_file; + snapshot.project_registered_file = project_registered_file; + snapshot.enabled = enabled; + snapshot.cloud_disabled = cloud_disabled; + snapshot.device_registered = device_registered; + snapshot.project_registered = project_registered; + return snapshot; +} + +static void process_events(const std::vector& fresh_events) { + std::vector events; + CloudConfigurationStateSnapshot snapshot; + std::string telemetry_dir; + + { + std::lock_guard guard(telemetry_mutex); + telemetry_dir = get_telemetry_dir_locked(); + events = load_cached_events_in_dir(telemetry_dir); + if (!fresh_events.empty()) { + events.insert(events.end(), fresh_events.begin(), fresh_events.end()); + } + if (events.empty()) return; + snapshot = capture_cloud_configuration_state_snapshot_locked(); + } + + CloudSendResult send_result = send_batch_to_cloud(events, snapshot); + + { + std::lock_guard guard(telemetry_mutex); + if (send_result.project_registered_ok && !project_registered) { + if (!project_registered_file.empty()) { + persist_registered_flag(project_registered_file); + } + project_registered = true; + } + if (send_result.device_registered_ok && !device_registered) { + if (!device_registered_file.empty()) { + persist_registered_flag(device_registered_file); + } + device_registered = true; + } + + if (send_result.payload_ok) { + clear_cache_files_in_dir(telemetry_dir); + } else if (!fresh_events.empty()) { + write_events_to_cache_in_dir(fresh_events, telemetry_dir); + } + } +} + +void init(const char* project_id_param, const char* project_scope_param, const char* cloud_key_param) { + std::unique_lock lifecycle_guard(telemetry_mutex); + telemetry_lifecycle_cv.wait(lifecycle_guard, [] { + return lifecycle_state != TelemetryLifecycleState::ShuttingDown; + }); + + std::string scope = project_scope_param ? project_scope_param : "default"; + + bool cloud_disabled_from_env = false; + if (const char* env = std::getenv("CACTUS_NO_CLOUD_TELE")) { + if (env[0] != '\0' && !(env[0] == '0' && env[1] == '\0')) { + cloud_disabled_from_env = true; + } + } + const char* env_url = std::getenv("CACTUS_SUPABASE_URL"); + const char* env_key = std::getenv("CACTUS_SUPABASE_KEY"); + const char* env_project = std::getenv("CACTUS_PROJECT_ID"); + const char* env_cloud = std::getenv("CACTUS_CLOUD_KEY"); + + std::string dir = get_telemetry_dir_locked(); + + std::string device_file = dir + "/device_id"; + std::string device_flag_file = dir + "/device_registered"; + bool device_was_registered = load_registered_flag(device_flag_file); + std::string resolved_device_id = load_or_create_id(device_file); + + std::string resolved_project_id; + if (project_id_param && *project_id_param) { + resolved_project_id = project_id_param; + } else if (env_project && *env_project) { + resolved_project_id = env_project; + } else { + std::string remote_url = read_git_remote_url(); + resolved_project_id = remote_url.empty() ? resolved_device_id : uuid_from_string(remote_url); + } + + std::string project_flag_file = dir + "/project_reg_" + resolved_project_id.substr(0, 8); + bool project_was_registered = load_registered_flag(project_flag_file); + + std::string resolved_cloud_key; + if (cloud_key_param && *cloud_key_param) { + resolved_cloud_key = cloud_key_param; + } else if (env_cloud && *env_cloud) { + resolved_cloud_key = env_cloud; + } + + collect_device_info(); + if (cactus_version.empty()) read_cactus_version(); + + if (env_url && *env_url) supabase_url = env_url; + if (env_key && *env_key) supabase_key = env_key; + project_id = resolved_project_id; + project_scope = scope; + project_registered_file = project_flag_file; + cloud_key = resolved_cloud_key; + device_registered_file = device_flag_file; + device_id = resolved_device_id; + if (app_id.empty()) app_id = detect_app_id(); + + project_registered = project_was_registered; + device_registered = device_was_registered; + if (cloud_disabled_from_env) { + cloud_disabled = true; + } + + if (!curl_initialized) { + if (curl_global_init(CURL_GLOBAL_DEFAULT) == CURLE_OK) { + curl_initialized = true; + } + } + + if (!atexit_registered) { + atexit_registered = true; + std::atexit([](){ shutdown(); }); + } + + shutdown_called = false; + ids_ready = true; + enabled = true; + lifecycle_state = TelemetryLifecycleState::Running; + TelemetryDispatcher::instance().start(); +} + +void setEnabled(bool en) { + std::lock_guard guard(telemetry_mutex); + enabled = en; +} + +void setCloudDisabled(bool disabled) { + std::lock_guard guard(telemetry_mutex); + cloud_disabled = disabled; +} + +void setAppId(const char* id) { + std::lock_guard guard(telemetry_mutex); + if (id && id[0] != '\0') app_id = id; +} + +void setTelemetryEnvironment(const char* framework_str, const char* cache_location_str, const char* version_str) { + std::lock_guard guard(telemetry_mutex); + if (framework_str && framework_str[0] != '\0') { + framework = framework_str; + } + if (cache_location_str && cache_location_str[0] != '\0') { + custom_cache_location = cache_location_str; + if (app_id.empty()) { + std::string pkg = extract_android_package_from_path(custom_cache_location); + if (!pkg.empty()) app_id = pkg; + } + } + if (version_str && version_str[0] != '\0') { + cactus_version = version_str; + } +} + +void setCloudKey(const char* key) { + if (!key || key[0] == '\0') return; + std::lock_guard guard(telemetry_mutex); + cloud_key = key; +} + +void cacheCloudApiKey(const char* key) { + if (!key || key[0] == '\0') return; + std::lock_guard guard(telemetry_mutex); + const std::string file = cloud_api_key_cache_file_locked(); + std::ofstream out(file, std::ios::trunc); + if (out.is_open()) { + out << key; + } +} + +std::string loadCachedCloudApiKey() { + std::lock_guard guard(telemetry_mutex); + const std::string file = cloud_api_key_cache_file_locked(); + std::ifstream in(file); + if (!in.is_open()) return {}; + std::string line; + if (std::getline(in, line) && !line.empty()) { + return line; + } + return {}; +} + +void recordInit(const char* model, bool success, double response_time_ms, const char* message) { + double nan = std::numeric_limits::quiet_NaN(); + Event e = make_event(INIT, model, success, nan, nan, response_time_ms, 0, message); + std::lock_guard guard(telemetry_mutex); + if (!can_record_event_locked()) return; + TelemetryDispatcher::instance().enqueue(e); +} + +void recordCompletion(const char* model, const CompletionMetrics& metrics) { + Event e = make_event_extended(COMPLETION, model, metrics); + std::lock_guard guard(telemetry_mutex); + if (!can_record_event_locked()) return; + TelemetryDispatcher::instance().enqueue(e); +} + +void recordCompletion(const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, const char* message) { + Event e = make_event(COMPLETION, model, success, ttft_ms, tps, response_time_ms, tokens, message); + std::lock_guard guard(telemetry_mutex); + if (!can_record_event_locked()) return; + TelemetryDispatcher::instance().enqueue(e); +} + +void recordEmbedding(const char* model, bool success, const char* message) { + Event e = make_event(EMBEDDING, model, success, 0.0, 0.0, 0.0, 0, message); + std::lock_guard guard(telemetry_mutex); + if (!can_record_event_locked()) return; + TelemetryDispatcher::instance().enqueue(e); +} + +void recordTranscription(const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, double ram_usage_mb, const char* message) { + std::lock_guard guard(telemetry_mutex); + if (in_stream_mode) return; + Event e = make_event(TRANSCRIPTION, model, success, ttft_ms, tps, response_time_ms, tokens, message); + e.ram_usage_mb = ram_usage_mb; + if (!can_record_event_locked()) return; + TelemetryDispatcher::instance().enqueue(e); +} + +void recordStreamTranscription(const char* model, bool success, double ttft_ms, double tps, double response_time_ms, int tokens, double session_ttft_ms, double session_tps, double session_time_ms, int session_tokens, const char* message) { + Event e = make_event(STREAM_TRANSCRIBE, model, success, ttft_ms, tps, response_time_ms, tokens, message); + e.session_ttft_ms = session_ttft_ms; + e.session_tps = session_tps; + e.session_time_ms = session_time_ms; + e.session_tokens = session_tokens; + std::lock_guard guard(telemetry_mutex); + if (!can_record_event_locked()) return; + TelemetryDispatcher::instance().enqueue(e); +} + +void setStreamMode(bool in_stream) { + std::lock_guard guard(telemetry_mutex); + in_stream_mode = in_stream; +} + +bool isStreamMode() { + std::lock_guard guard(telemetry_mutex); + return in_stream_mode; +} + +void markInference(bool active) { + std::lock_guard guard(telemetry_mutex); + if (active) { + inference_active += 1; + } else if (inference_active > 0) { + inference_active -= 1; + } +} + +void flush() { + TelemetryDispatcher::instance().flush(); +} + +void shutdown() { + { + std::lock_guard lifecycle_guard(telemetry_mutex); + if (shutdown_called) { + return; + } + + shutdown_called = true; + lifecycle_state = TelemetryLifecycleState::ShuttingDown; + } + + flush(); + + { + std::lock_guard lifecycle_guard(telemetry_mutex); + enabled = false; + ids_ready = false; + } + + TelemetryDispatcher::instance().stop(); + + { + std::lock_guard lifecycle_guard(telemetry_mutex); + if (curl_initialized) { + curl_initialized = false; + curl_global_cleanup(); + } + lifecycle_state = TelemetryLifecycleState::Stopped; + telemetry_lifecycle_cv.notify_all(); + } +} + +} // namespace telemetry +} // namespace cactus diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp new file mode 100644 index 000000000..d792578db --- /dev/null +++ b/cactus-engine/src/tokenizer.cpp @@ -0,0 +1,778 @@ +#include "engine.h" +#include "cactus_kernels.h" +#include "gemma_tools.h" +#include "chat_tools.h" +#include +#include +#include +#include +#include + +namespace cactus { +namespace engine { + +namespace { + +std::string format_needle_query_text(const std::vector& messages) { + std::string system_text; + std::string user_query; + + for (const auto& msg : messages) { + if (msg.role == "system" || msg.role == "developer") { + if (!system_text.empty()) system_text += "\n"; + system_text += msg.content; + } else if (msg.role == "user") { + user_query = msg.content; + } + } + + if (user_query.empty() && !messages.empty()) user_query = messages.back().content; + if (system_text.empty()) return user_query; + if (user_query.empty()) return system_text; + return system_text + "\n\n" + user_query; +} + +std::string format_tool_call_for_prompt(const std::string& name, const std::string& arguments, bool gemma4) { + std::string args = arguments.empty() ? "{}" : arguments; + size_t pos = 0; + std::string dsl = gemma::format_argument(args, pos, false); + if (dsl.empty() || dsl.front() != '{') dsl = "{" + dsl + "}"; + if (gemma4) { + return "<|tool_call>call:" + name + dsl + ""; + } + return "call:" + name + gemma::use_escape_tags(dsl) + ""; +} + +std::string format_tool_response_for_prompt(const std::string& name, const std::string& content, bool gemma4) { + if (gemma4) { + return "<|tool_response>response:" + name + "{value:<|\"|>" + content + "<|\"|>}"; + } + return "response:" + name + "{value:" + content + "}"; +} + +std::string trim_copy(const std::string& value) { + size_t start = value.find_first_not_of(" \t\r\n"); + if (start == std::string::npos) return ""; + size_t end = value.find_last_not_of(" \t\r\n"); + return value.substr(start, end - start + 1); +} + +TokenizerRuntimeConfig::TokenizerType parse_tokenizer_type(const std::string& value) { + if (value == "bpe") return TokenizerRuntimeConfig::TokenizerType::BPE; + if (value == "sentencepiece") return TokenizerRuntimeConfig::TokenizerType::SENTENCEPIECE; + return TokenizerRuntimeConfig::TokenizerType::UNKNOWN; +} + +TokenizerRuntimeConfig::VocabFormat parse_vocab_format(const std::string& value) { + if (value == "id_tab_token") return TokenizerRuntimeConfig::VocabFormat::ID_TAB_TOKEN; + if (value == "line_token") return TokenizerRuntimeConfig::VocabFormat::LINE_TOKEN; + return TokenizerRuntimeConfig::VocabFormat::UNKNOWN; +} + +TokenizerRuntimeConfig::Normalizer parse_normalizer(const std::string& value) { + if (value == "metaspace") return TokenizerRuntimeConfig::Normalizer::METASPACE; + if (value == "byte_level") return TokenizerRuntimeConfig::Normalizer::BYTE_LEVEL; + return TokenizerRuntimeConfig::Normalizer::NONE; +} + +TokenizerRuntimeConfig::Decoder parse_decoder(const std::string& value) { + if (value == "replace_metaspace") return TokenizerRuntimeConfig::Decoder::REPLACE_METASPACE; + if (value == "byte_level") return TokenizerRuntimeConfig::Decoder::BYTE_LEVEL; + return TokenizerRuntimeConfig::Decoder::NONE; +} + +void skip_json_whitespace(const std::string& json, size_t& pos) { + while (pos < json.size() && std::isspace(static_cast(json[pos]))) { + ++pos; + } +} + +bool extract_added_token_object(const std::string& json, size_t& pos, std::string& out_object) { + skip_json_whitespace(json, pos); + if (pos >= json.size() || json[pos] != '{') { + return false; + } + + size_t start = pos; + size_t depth = 0; + bool in_string = false; + bool escaped = false; + + while (pos < json.size()) { + char c = json[pos++]; + if (escaped) { + escaped = false; + continue; + } + if (c == '\\') { + escaped = true; + continue; + } + if (c == '"') { + in_string = !in_string; + continue; + } + if (in_string) { + continue; + } + if (c == '{') { + ++depth; + } else if (c == '}') { + if (depth == 0) { + return false; + } + --depth; + if (depth == 0) { + out_object = json.substr(start, pos - start); + return true; + } + } + } + + return false; +} + +bool parse_added_token_entry(const std::string& object, std::string& token_content, uint32_t& token_id, + bool& is_special) { + token_content.clear(); + token_id = 0; + is_special = false; + + size_t id_key = object.find("\"id\""); + if (id_key == std::string::npos) { + return false; + } + size_t id_colon = object.find(':', id_key); + if (id_colon == std::string::npos) { + return false; + } + size_t id_pos = id_colon + 1; + skip_json_whitespace(object, id_pos); + size_t id_end = id_pos; + while (id_end < object.size() && std::isdigit(static_cast(object[id_end]))) { + ++id_end; + } + if (id_end == id_pos) { + return false; + } + token_id = static_cast(std::stoul(object.substr(id_pos, id_end - id_pos))); + + size_t content_key = object.find("\"content\""); + if (content_key == std::string::npos) { + return false; + } + size_t content_colon = object.find(':', content_key); + if (content_colon == std::string::npos) { + return false; + } + size_t content_pos = object.find('"', content_colon + 1); + if (content_pos == std::string::npos) { + return false; + } + ++content_pos; + token_content = extract_json_string(object, content_pos); + + size_t special_key = object.find("\"special\""); + if (special_key != std::string::npos) { + size_t special_colon = object.find(':', special_key); + if (special_colon != std::string::npos) { + size_t special_pos = special_colon + 1; + skip_json_whitespace(object, special_pos); + is_special = object.compare(special_pos, 4, "true") == 0; + } + } + + return true; +} + +void load_tokenizer_json_added_special_tokens( + const std::string& tokenizer_json_path, + std::unordered_map& special_tokens) { + std::ifstream file(tokenizer_json_path); + if (!file.is_open()) { + return; + } + + std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + size_t pos = content.find("\"added_tokens\""); + if (pos == std::string::npos) { + return; + } + + pos = content.find('[', pos); + if (pos == std::string::npos) { + return; + } + ++pos; + + while (pos < content.size()) { + skip_json_whitespace(content, pos); + if (pos >= content.size() || content[pos] == ']') { + break; + } + + std::string object; + if (!extract_added_token_object(content, pos, object)) { + ++pos; + continue; + } + + std::string token_content; + uint32_t token_id = 0; + bool is_special = false; + if (parse_added_token_entry(object, token_content, token_id, is_special) && is_special) { + special_tokens[token_content] = token_id; + } + + skip_json_whitespace(content, pos); + if (pos < content.size() && content[pos] == ',') { + ++pos; + } + } +} + +} // namespace + +TokenizerRuntimeConfig load_tokenizer_runtime_config(const std::string& config_file) { + TokenizerRuntimeConfig config; + + std::ifstream file(config_file); + if (!file.is_open()) { + return config; + } + + std::string line; + while (std::getline(file, line)) { + if (line.empty() || line[0] == '#') continue; + + size_t eq_pos = line.find('='); + if (eq_pos == std::string::npos) continue; + + const std::string key = trim_copy(line.substr(0, eq_pos)); + const std::string value = trim_copy(line.substr(eq_pos + 1)); + + if (key == "tokenizer_type") { + config.tokenizer_type = parse_tokenizer_type(value); + } else if (key == "vocab_format") { + config.vocab_format = parse_vocab_format(value); + } else if (key == "normalizer") { + config.normalizer = parse_normalizer(value); + } else if (key == "decoder") { + config.decoder = parse_decoder(value); + } else if (key == "byte_fallback") { + config.byte_fallback = (value == "true" || value == "1"); + } else if (key == "has_chat_template") { + config.has_chat_template = (value == "true" || value == "1"); + } + } + + return config; +} + +void load_special_tokens_map(const std::string& config_file, std::unordered_map& special_tokens) { + special_tokens.clear(); + + std::ifstream file(config_file); + if (file.is_open()) { + std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + + size_t pos = content.find("\"special_tokens\""); + if (pos != std::string::npos) { + pos = content.find("{", pos); + if (pos != std::string::npos) { + size_t end_pos = content.find("}", pos); + if (end_pos != std::string::npos) { + std::string special_tokens_section = content.substr(pos + 1, end_pos - pos - 1); + std::istringstream iss(special_tokens_section); + std::string line; + + while (std::getline(iss, line)) { + size_t colon_pos = line.find(":"); + if (colon_pos == std::string::npos) continue; + + std::string id_part = line.substr(0, colon_pos); + std::string token_part = line.substr(colon_pos + 1); + + size_t id_start = id_part.find("\""); + size_t id_end = id_part.find("\"", id_start + 1); + if (id_start == std::string::npos || id_end == std::string::npos) continue; + + uint32_t token_id = + static_cast(std::stoul(id_part.substr(id_start + 1, id_end - id_start - 1))); + + size_t token_start = token_part.find("\""); + if (token_start == std::string::npos) continue; + size_t value_pos = token_start + 1; + std::string token_content = extract_json_string(token_part, value_pos); + special_tokens[token_content] = token_id; + } + } + } + } + } + + size_t slash_pos = config_file.find_last_of("/\\"); + std::string dir = (slash_pos == std::string::npos) ? "." : config_file.substr(0, slash_pos); + load_tokenizer_json_added_special_tokens(dir + "/tokenizer.json", special_tokens); +} + +std::vector split_with_special_tokens(const std::string& text, + const std::unordered_map& special_tokens) { + std::vector result; + size_t start = 0; + while (start < text.size()) { + size_t best_match_pos = text.size(); + size_t best_match_len = 0; + std::string best_special_token; + + for (const auto& [special_token, token_id] : special_tokens) { + if (special_token.empty()) continue; + size_t pos = text.find(special_token, start); + if (pos != std::string::npos && + (pos < best_match_pos || (pos == best_match_pos && special_token.length() > best_match_len))) { + best_match_pos = pos; + best_match_len = special_token.length(); + best_special_token = special_token; + } + } + + if (best_match_pos < text.size()) { + if (best_match_pos > start) + result.push_back(text.substr(start, best_match_pos - start)); + result.push_back(best_special_token); + start = best_match_pos + best_match_len; + } else { + if (start < text.size()) + result.push_back(text.substr(start)); + break; + } + } + return result; +} + +void Tokenizer::load_chat_template(const std::string& template_file) { + std::ifstream file(template_file); + if (!file.is_open()) { + has_chat_template_ = false; + return; + } + chat_template_ = std::string((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + has_chat_template_ = !chat_template_.empty(); +} + +void Tokenizer::detect_model_type(const std::string& config_path) { + model_type_ = ModelType::GEMMA4; + + std::ifstream file(config_path); + if (!file.is_open()) return; + + std::string line; + while (std::getline(file, line)) { + std::string lower_line = line; + std::transform(lower_line.begin(), lower_line.end(), lower_line.begin(), ::tolower); + + if (lower_line.find("model_type") != std::string::npos) { + if (lower_line.find("needle") != std::string::npos) { + model_type_ = ModelType::NEEDLE; + } else if (lower_line.find("qwen") != std::string::npos) { + model_type_ = ModelType::QWEN; + } else if (lower_line.find("lfm2") != std::string::npos) { + model_type_ = ModelType::LFM2; + } else if (lower_line.find("gemma4") != std::string::npos) { + model_type_ = ModelType::GEMMA4; + } else if (lower_line.find("gemma") != std::string::npos) { + model_type_ = ModelType::GEMMA; + } + } else if (lower_line.find("model_variant") != std::string::npos) { + if (lower_line.find("vlm") != std::string::npos) { model_variant_ = ModelVariant::VLM; } + else if (lower_line.find("extract") != std::string::npos) { model_variant_ = ModelVariant::EXTRACT; } + else if (lower_line.find("rag") != std::string::npos) { model_variant_ = ModelVariant::RAG; } + } + } + + file.clear(); + file.seekg(0); + while (std::getline(file, line)) { + auto parse_uint = [&](const std::string& key, uint32_t& out) { + size_t p = line.find(key + "="); + if (p != std::string::npos) { + out = static_cast(std::stoul(line.substr(p + key.size() + 1))); + } + }; + parse_uint("vision_patch_size", vision_patch_size_); + parse_uint("vision_pooling_kernel_size", vision_pooling_kernel_size_); + parse_uint("vision_default_output_length", vision_default_output_length_); + parse_uint("vision_image_size", vision_image_size_); + } +} + +std::string Tokenizer::get_default_stop_sequence() const { + if (model_type_ == ModelType::NEEDLE) { + return ""; + } + if (model_type_ == ModelType::QWEN || model_type_ == ModelType::LFM2) { + return "<|im_end|>"; + } + if (model_type_ == ModelType::GEMMA) { + return ""; + } + return ""; +} + +std::vector Tokenizer::apply_chat_template(const std::vector& messages, bool add_generation_prompt) const { + return encode(format_chat_prompt(messages, add_generation_prompt)); +} + +std::string Tokenizer::format_chat_prompt(const std::vector& messages, bool add_generation_prompt, + const std::string& tools_json, bool enable_thinking_if_supported) const { + if (model_type_ == ModelType::QWEN) { + return format_qwen_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); + } + if (model_type_ == ModelType::LFM2) { + return format_lfm2_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); + } + if (model_type_ == ModelType::NEEDLE) { + return format_needle_style(messages, add_generation_prompt, tools_json); + } + if (model_type_ == ModelType::GEMMA) { + return format_gemma_style(messages, add_generation_prompt, tools_json); + } + return format_gemma4_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); +} + +std::string Tokenizer::format_gemma_style(const std::vector& messages, bool add_generation_prompt, + const std::string& tools_json) const { + std::string result = ""; + if (tools_json.empty() && !has_function_call_tokens()) { + std::string pending_context; + for (const auto& msg : messages) { + if (msg.role == "system" || msg.role == "developer") { + pending_context += msg.content + "\n\n"; + continue; + } + std::string role = msg.role == "assistant" ? "model" : "user"; + std::string content = msg.content; + for (const auto& tc : msg.tool_calls) { + content += "\ncall:" + tc.name + "(" + tc.arguments + ")"; + } + if (role == "user" && !pending_context.empty()) { + content = pending_context + content; + pending_context.clear(); + } + result += "" + role + "\n" + content + "\n"; + } + if (add_generation_prompt) { + result += "model\n"; + } + return result; + } + + size_t first = 0; + std::string sys; + if (!messages.empty() && (messages[0].role == "system" || messages[0].role == "developer")) { + sys = messages[0].content; + first = 1; + } + if (first > 0 || !tools_json.empty()) { + result += "developer\n" + sys + tools_json + "\n"; + } + bool after_tool_response = false; + for (size_t i = first; i < messages.size(); i++) { + const auto& msg = messages[i]; + if (msg.role == "tool") { + result += format_tool_response_for_prompt(msg.name, msg.content, false); + after_tool_response = true; + continue; + } + std::string role = msg.role == "assistant" ? "model" : msg.role; + if (!after_tool_response) { + result += "" + role + "\n"; + } + after_tool_response = false; + result += msg.content; + for (const auto& tc : msg.tool_calls) { + result += format_tool_call_for_prompt(tc.name, tc.arguments, false); + } + if (msg.tool_calls.empty()) { + result += "\n"; + } else if (i + 1 == messages.size()) { + result += ""; + } + } + if (add_generation_prompt && !after_tool_response) { + result += "model\n"; + } + return result; +} + +namespace { +std::string strip_newlines(const std::string& s) { + size_t a = s.find_first_not_of('\n'); + if (a == std::string::npos) return ""; + size_t b = s.find_last_not_of('\n'); + return s.substr(a, b - a + 1); +} +std::string lstrip_newlines(const std::string& s) { + size_t a = s.find_first_not_of('\n'); + return a == std::string::npos ? "" : s.substr(a); +} +} // namespace + +std::string Tokenizer::format_qwen_style(const std::vector& messages, bool add_generation_prompt, + const std::string& tools_json, bool enable_thinking_if_supported) const { + std::string result; + const size_t n = messages.size(); + const bool template_has_thinking = !has_chat_template_ || chat_template_.find("") != std::string::npos; + + size_t first = 0; + std::string sys; + const bool has_sys = n > 0 && (messages[0].role == "system" || messages[0].role == "developer"); + if (has_sys) { sys = messages[0].content; first = 1; } + if (!tools_json.empty()) { + result += "<|im_start|>system\n"; + if (has_sys) result += sys + "\n\n"; + result += tools_json; + result += "<|im_end|>\n"; + } else if (has_sys) { + result += "<|im_start|>system\n" + sys + "<|im_end|>\n"; + } + + long last_query_index = static_cast(n) - 1; + for (long i = static_cast(n) - 1; i >= 0; --i) { + if (messages[i].role == "user") { last_query_index = i; break; } + } + + for (size_t i = first; i < n; i++) { + const auto& msg = messages[i]; + std::string role = msg.role; + if (role == "developer") role = "system"; + else if (role != "system" && role != "assistant" && role != "tool") role = "user"; + + if (role == "user" || role == "system") { + result += "<|im_start|>" + role + "\n"; + if (role == "user") { + const size_t soft_n = image_soft_token_count_ > 0 ? image_soft_token_count_ : 1; + for (const auto& image_path : msg.images) { + (void)image_path; + result += "<|vision_start|>"; + for (size_t k = 0; k < soft_n; ++k) result += "<|image_pad|>"; + result += "<|vision_end|>"; + } + } + result += msg.content + "<|im_end|>\n"; + } else if (role == "assistant") { + std::string content = msg.content; + std::string reasoning; + size_t tpos = content.find(""); + if (tpos != std::string::npos) { + std::string head = content.substr(0, tpos); + size_t ts = head.rfind(""); + reasoning = strip_newlines(ts != std::string::npos ? head.substr(ts + 7) : head); + content = lstrip_newlines(content.substr(tpos + 8)); + } + result += "<|im_start|>assistant\n"; + if (template_has_thinking && static_cast(i) > last_query_index && (i == n - 1 || !reasoning.empty())) { + result += "\n" + reasoning + "\n\n\n" + lstrip_newlines(content); + } else { + result += content; + } + for (size_t k = 0; k < msg.tool_calls.size(); ++k) { + const auto& tc = msg.tool_calls[k]; + if ((k == 0 && !content.empty()) || k > 0) result += "\n"; + result += "\n{\"name\": \"" + tc.name + "\", \"arguments\": " + + chat_tools::respace_json(tc.arguments.empty() ? "{}" : tc.arguments) + "}\n"; + } + result += "<|im_end|>\n"; + } else { // tool + if (i == 0 || messages[i - 1].role != "tool") result += "<|im_start|>user"; + result += "\n\n" + msg.content + "\n"; + if (i + 1 >= n || messages[i + 1].role != "tool") result += "<|im_end|>\n"; + } + } + + if (add_generation_prompt) { + result += "<|im_start|>assistant\n"; + if (!enable_thinking_if_supported && template_has_thinking) result += "\n\n\n\n"; + } + return result; +} + +std::string Tokenizer::format_lfm2_style(const std::vector& messages, bool add_generation_prompt, + const std::string& tools_json, bool /*enable_thinking_if_supported*/) const { + std::string result = "<|startoftext|>"; + const size_t n = messages.size(); + + size_t first = 0; + std::string sys; + const bool has_sys = n > 0 && (messages[0].role == "system" || messages[0].role == "developer"); + if (has_sys) { sys = messages[0].content; first = 1; } + if (!tools_json.empty() || has_sys) { + result += "<|im_start|>system\n"; + if (has_sys) result += sys; + if (!tools_json.empty()) { if (has_sys) result += "\n"; result += tools_json; } + result += "<|im_end|>\n"; + } + + for (size_t i = first; i < n; i++) { + const auto& msg = messages[i]; + std::string role = msg.role; + if (role == "developer") role = "system"; + else if (role != "system" && role != "assistant" && role != "tool") role = "user"; + + result += "<|im_start|>" + role + "\n"; + if (role == "user") { + for (const auto& image_path : msg.images) { + int iw = 0, ih = 0, ic = 0; + Lfm2VlTokenLayout layout; + if (has_lfm2_vision_config_ && cactus_image_info(image_path.c_str(), &iw, &ih, &ic)) { + layout = lfm2_vl_token_layout(ih, iw, lfm2_vision_config_); + } else { + layout.grid_rows = 1; + layout.grid_cols = 1; + layout.tokens_per_tile = image_soft_token_count_ > 0 ? static_cast(image_soft_token_count_) : 1; + } + result += "<|image_start|>"; + const bool multi_tile = layout.grid_rows > 1 || layout.grid_cols > 1; + if (multi_tile) { + for (int r = 0; r < layout.grid_rows; ++r) { + for (int c = 0; c < layout.grid_cols; ++c) { + result += "<|img_row_" + std::to_string(r + 1) + "_col_" + std::to_string(c + 1) + "|>"; + for (int k = 0; k < layout.tokens_per_tile; ++k) result += ""; + } + } + if (layout.has_thumbnail) { + result += "<|img_thumbnail|>"; + for (int k = 0; k < layout.thumbnail_tokens; ++k) result += ""; + } + } else { + for (int k = 0; k < layout.tokens_per_tile; ++k) result += ""; + } + result += "<|image_end|>\n"; + } + } + if (role == "tool") { + result += "<|tool_response_start|>" + msg.content + "<|tool_response_end|>"; + } else { + result += msg.content; + } + if (role == "assistant" && !msg.tool_calls.empty()) { + result += "<|tool_call_start|>["; + for (size_t k = 0; k < msg.tool_calls.size(); ++k) { + if (k) result += ", "; + result += chat_tools::pythonic_call(msg.tool_calls[k].name, msg.tool_calls[k].arguments); + } + result += "]<|tool_call_end|>"; + } + result += "<|im_end|>\n"; + } + + if (add_generation_prompt) result += "<|im_start|>assistant\n"; + return result; +} + +std::string Tokenizer::format_needle_style(const std::vector& messages, bool /*add_generation_prompt*/, + const std::string& tools_json) const { + std::string serialized_tools = tools_json.empty() ? "[]" : tools_json; + return format_needle_query_text(messages) + "" + serialized_tools + ""; +} + +std::string Tokenizer::format_gemma4_style(const std::vector& messages, bool add_generation_prompt, + const std::string& tools_json, bool enable_thinking_if_supported) const { + std::string result = ""; + + std::string sys_content; + size_t first_msg = 0; + if (!messages.empty() && (messages[0].role == "system" || messages[0].role == "developer")) { + sys_content = messages[0].content; + first_msg = 1; + } + + if (enable_thinking_if_supported || !sys_content.empty() || !tools_json.empty()) { + result += "<|turn>system\n"; + if (enable_thinking_if_supported) { + result += "<|think|>\n"; + } + result += sys_content; + result += tools_json; + result += "\n"; + } + + auto compute_soft_tokens = [&](const std::string& image_path) -> size_t { + if (image_soft_token_count_ > 0) return image_soft_token_count_; + + int w = 0, h = 0, c = 0; + if (!cactus_image_info(image_path.c_str(), &w, &h, &c)) return 0; + + uint32_t p = vision_patch_size_ ? vision_patch_size_ : 16; + uint32_t k = vision_pooling_kernel_size_ ? vision_pooling_kernel_size_ : 3; + uint32_t out_len = vision_default_output_length_ ? vision_default_output_length_ : 280; + uint32_t side = k * p; + uint32_t max_patches = out_len * k * k; + float factor = std::sqrt(static_cast(max_patches) * p * p / + (static_cast(h) * w)); + int th = static_cast(std::floor(factor * h / side)) * side; + int tw = static_cast(std::floor(factor * w / side)) * side; + if (th == 0) th = side; + if (tw == 0) tw = side; + return static_cast((th / p / k) * (tw / p / k)); + }; + bool in_model_turn = false; + std::vector pending_call_names; + auto close_model_turn = [&]() { + if (in_model_turn) { result += "\n"; in_model_turn = false; } + }; + + for (size_t i = first_msg; i < messages.size(); i++) { + const auto& msg = messages[i]; + const std::string role = (msg.role == "assistant") ? "model" + : (msg.role == "developer") ? "system" : msg.role; + + if (role == "model") { + if (!in_model_turn) { result += "<|turn>model\n"; in_model_turn = true; } + result += msg.content; + for (const auto& tc : msg.tool_calls) { + result += format_tool_call_for_prompt(tc.name, tc.arguments, true); + pending_call_names.push_back(tc.name); + } + if (msg.tool_calls.empty()) close_model_turn(); + } else if (role == "tool") { + if (!in_model_turn) { result += "<|turn>model\n"; in_model_turn = true; } + std::string fn = !msg.name.empty() ? msg.name + : (!pending_call_names.empty() ? pending_call_names.front() + : std::string("unknown")); + if (!pending_call_names.empty()) pending_call_names.erase(pending_call_names.begin()); + result += format_tool_response_for_prompt(fn, msg.content, true); + } else { + close_model_turn(); + result += "<|turn>" + role + "\n"; + for (const auto& image_path : msg.images) { + size_t n = compute_soft_tokens(image_path); + if (n > 0) { + result += "<|image>"; + for (size_t j = 0; j < n; j++) + result += "<|image|>"; + result += ""; + } + } + result += msg.content; + if (msg.audio_soft_token_count > 0) { + result += "<|audio>"; + for (size_t j = 0; j < msg.audio_soft_token_count; j++) + result += "<|audio|>"; + result += ""; + } + result += "\n"; + } + } + + if (add_generation_prompt) { + if (!in_model_turn) result += "<|turn>model\n"; + } else { + close_model_turn(); + } + + return result; +} + +} // namespace engine +} // namespace cactus diff --git a/cactus-engine/src/transcribe.cpp b/cactus-engine/src/transcribe.cpp new file mode 100644 index 000000000..9f2ebe928 --- /dev/null +++ b/cactus-engine/src/transcribe.cpp @@ -0,0 +1,529 @@ +#include "../cactus_engine.h" +#include "utils.h" +#include "cloud.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include "wav.h" +#include +#include +#include +#include +#include +#include + +using namespace cactus::ffi; + +namespace { + +constexpr float kWhisperTimestampStepSec = 0.02f; + +uint32_t whisper_timestamp_begin(cactus::engine::Tokenizer* tokenizer) { + return tokenizer->encode("<|notimestamps|>")[0] + 1; +} + +std::vector parse_whisper_timestamp_segments( + cactus::engine::Tokenizer* tokenizer, const std::vector& tokens, + uint32_t timestamp_begin, std::string& clean_text) { + std::vector segments; + std::vector text_tokens; + std::vector flat; + float seg_start = -1.0f; + auto flush = [&](float end_sec) { + if (!text_tokens.empty()) { + std::string text = tokenizer->decode(text_tokens); + if (!text.empty() && text[0] == ' ') text.erase(0, 1); + segments.push_back({seg_start, end_sec, text}); + } + text_tokens.clear(); + }; + for (uint32_t tok : tokens) { + if (tok >= timestamp_begin) { + const float t = static_cast(tok - timestamp_begin) * kWhisperTimestampStepSec; + if (seg_start < 0.0f) seg_start = t; + else { flush(t); seg_start = t; } + } else { + text_tokens.push_back(tok); + flat.push_back(tok); + } + } + flush(seg_start); + clean_text = tokenizer->decode(flat); + return segments; +} + +} // namespace + +extern "C" { + +CACTUS_FFI_EXPORT int cactus_preprocess_audio_features( + const char* audio_file_path, + const char* model_type, + size_t mel_bins, + float* features_buffer, + size_t buffer_size, + size_t* feature_count, + size_t* out_mel_bins, + size_t* out_frames +) { + if (!audio_file_path || !model_type || !features_buffer || !feature_count || !out_mel_bins || !out_frames) { + last_error_message = "Invalid parameters for audio feature preprocessing"; + CACTUS_LOG_ERROR("audio_preprocess", last_error_message); + return -1; + } + + try { + AudioFP32 audio = load_wav(audio_file_path); + std::vector audio_samples = resample_to_16k_fp32(audio.samples, audio.sample_rate); + if (audio_samples.empty()) { + last_error_message = "No audio samples available for preprocessing"; + CACTUS_LOG_ERROR("audio_preprocess", last_error_message); + return -1; + } + + const size_t bins = std::max(1, mel_bins); + std::string lowered(model_type); + std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + + std::vector features; + size_t frames = 0; + + if (lowered.find("gemma4") != std::string::npos || lowered.find("gemma-4") != std::string::npos) { + cactus::engine::Config cfg; + cfg.model_type = cactus::engine::Config::ModelType::GEMMA4; + cfg.audio_input_feat_size = static_cast(bins); + // The Python transpiled runtime shape-specializes Gemma4 audio to + // its requested capacity (currently capped at 30s). Do not apply + // the native interactive model's shorter audio_soft_tokens limit + // here, otherwise long-file multimodal prompts are silently + // truncated before they reach the graph bundle. + cfg.audio_soft_tokens = 0; + cfg.audio_fft_length = 512; + auto prepared = cactus::audio::preprocess_audio_for_gemma4(std::move(audio_samples), cfg); + features = std::move(prepared.features); + frames = prepared.num_frames; + } else if (lowered.find("parakeet") != std::string::npos) { + auto cfg = cactus::audio::get_parakeet_spectrogram_config(); + const size_t waveform_samples = audio_samples.size(); + cactus::audio::apply_preemphasis(audio_samples, 0.97f); + features = cactus::audio::compute_spectrogram_graph( + audio_samples, cfg, bins, 0.0f, 8000.0f, + cactus::audio::WHISPER_SAMPLE_RATE, 0, 0); + cactus::audio::normalize_parakeet_log_mel(features, bins); + size_t valid_frames = waveform_samples / cfg.hop_length; + if (valid_frames == 0) valid_frames = 1; + cactus::audio::trim_mel_frames(features, bins, valid_frames); + frames = features.size() / bins; + } else { + cactus::audio::pad_or_trim_whisper_waveform(audio_samples); + auto cfg = cactus::audio::get_whisper_spectrogram_config(); + const bool is_whisper_v3 = bins > 80; + if (is_whisper_v3) cactus::audio::apply_whisper_v3_overrides(cfg); + int norm_type = 1; // Whisper HF feature extractor uses Slaney-normalized mel filters. + int scale_type = 2; // Whisper HF feature extractor uses the Slaney mel scale. + std::vector mel = cactus::audio::compute_spectrogram_graph( + audio_samples, cfg, bins, 0.0f, 8000.0f, + cactus::audio::WHISPER_SAMPLE_RATE, norm_type, scale_type); + features = cactus::audio::normalize_whisper_mel(mel, bins, true); + frames = features.size() / bins; + } + + const size_t bytes_needed = features.size() * sizeof(float); + if (bytes_needed > buffer_size) { + last_error_message = "Audio feature output buffer too small"; + CACTUS_LOG_ERROR("audio_preprocess", last_error_message); + return -2; + } + + std::memcpy(features_buffer, features.data(), bytes_needed); + *feature_count = features.size(); + *out_mel_bins = bins; + *out_frames = frames; + return static_cast(features.size()); + } catch (const std::exception& e) { + last_error_message = e.what(); + CACTUS_LOG_ERROR("audio_preprocess", last_error_message); + return -1; + } catch (...) { + last_error_message = "Unknown error during audio feature preprocessing"; + CACTUS_LOG_ERROR("audio_preprocess", last_error_message); + return -1; + } +} + +int cactus_transcribe( + cactus_model_t model, + const char* audio_file_path, + const char* prompt, + char* response_buffer, + size_t buffer_size, + const char* options_json, + cactus_token_callback callback, + void* user_data, + const uint8_t* pcm_buffer, + size_t pcm_buffer_size +) { + struct MetalTrimGuard { + ~MetalTrimGuard() { cactus_metal_trim_prefill_cache(); } + } metal_trim_guard; + if (validate_audio_params("transcribe", model, response_buffer, buffer_size, + audio_file_path, pcm_buffer, pcm_buffer_size) != 0) { + return -1; + } + + try { + auto* handle = static_cast(model); + const auto model_type = handle->model->get_config().model_type; + const bool is_whisper = model_type == cactus::engine::Config::ModelType::WHISPER; + const bool is_parakeet = model_type == cactus::engine::Config::ModelType::PARAKEET_TDT; + + if (!is_whisper && !is_parakeet) { + const uint8_t* pcm_data = pcm_buffer; + size_t pcm_size = pcm_buffer_size; + std::vector file_pcm; + + if (audio_file_path && (!pcm_buffer || pcm_buffer_size == 0)) { + AudioFP32 audio = load_wav(audio_file_path); + if (audio.samples.empty()) { + CACTUS_LOG_ERROR("transcribe", "Failed to load audio file: " << audio_file_path); + handle_error_response("Failed to load audio file", response_buffer, buffer_size); + return -1; + } + file_pcm.resize(audio.samples.size() * sizeof(int16_t)); + int16_t* out = reinterpret_cast(file_pcm.data()); + for (size_t i = 0; i < audio.samples.size(); i++) { + float clamped = std::max(-1.0f, std::min(1.0f, audio.samples[i])); + out[i] = static_cast(clamped * 32767.0f); + } + pcm_data = file_pcm.data(); + pcm_size = file_pcm.size(); + } + + std::string user_content = "Transcribe the following audio."; + if (prompt && prompt[0] != '\0') { + user_content = prompt; + } + + std::string messages_json = "[{\"role\": \"user\", \"content\": \"" + + escape_json_string(user_content) + "\"}]"; + + return cactus_complete( + model, + messages_json.c_str(), + response_buffer, + buffer_size, + options_json, + nullptr, + callback, + user_data, + pcm_data, + pcm_size + ); + } + + std::lock_guard lock(handle->model_mutex); + handle->should_stop = false; + auto start_time = std::chrono::high_resolution_clock::now(); + InferenceOptions options = parse_inference_options_json(options_json ? options_json : ""); + + std::vector audio_samples; + if (audio_file_path == nullptr) { + auto waveform_fp32 = cactus::audio::pcm_buffer_to_float_samples(pcm_buffer, pcm_buffer_size); + audio_samples = resample_to_16k_fp32(waveform_fp32, cactus::audio::WHISPER_SAMPLE_RATE); + } else { + AudioFP32 audio = load_wav(audio_file_path); + audio_samples = resample_to_16k_fp32(audio.samples, audio.sample_rate); + } + if (audio_samples.empty()) { + handle_error_response("No audio input provided", response_buffer, buffer_size); + return -1; + } + const size_t original_audio_sample_count = audio_samples.size(); + + const size_t mel_bins = std::max(1, static_cast(handle->model->get_config().num_mel_bins)); + std::vector audio_features; + std::vector handoff_audio_samples; + if (is_parakeet) { + auto cfg = cactus::audio::get_parakeet_spectrogram_config(); + const size_t waveform_samples = audio_samples.size(); + handoff_audio_samples = audio_samples; + cactus::audio::apply_preemphasis(audio_samples, 0.97f); + audio_features = cactus::audio::compute_spectrogram_graph( + audio_samples, cfg, mel_bins, 0.0f, 8000.0f, + cactus::audio::WHISPER_SAMPLE_RATE, 0, 0); + cactus::audio::normalize_parakeet_log_mel(audio_features, mel_bins); + size_t valid_frames = waveform_samples / cfg.hop_length; + if (valid_frames == 0) valid_frames = 1; + cactus::audio::trim_mel_frames(audio_features, mel_bins, valid_frames); + } else { + cactus::audio::pad_or_trim_whisper_waveform(audio_samples); + auto cfg = cactus::audio::get_whisper_spectrogram_config(); + const bool is_whisper_v3 = mel_bins > 80; + if (is_whisper_v3) cactus::audio::apply_whisper_v3_overrides(cfg); + int norm_type = 1; // Whisper HF feature extractor uses Slaney-normalized mel filters. + int scale_type = 2; // Whisper HF feature extractor uses the Slaney mel scale. + std::vector mel = cactus::audio::compute_spectrogram_graph( + audio_samples, cfg, mel_bins, 0.0f, 8000.0f, + cactus::audio::WHISPER_SAMPLE_RATE, norm_type, scale_type); + audio_features = cactus::audio::normalize_whisper_mel(mel, mel_bins, true); + } + if (audio_features.empty()) { + handle_error_response("Computed audio features are empty", response_buffer, buffer_size); + return -1; + } + + auto* tokenizer = handle->model->get_tokenizer(); + if (!tokenizer) { + handle_error_response("Tokenizer unavailable", response_buffer, buffer_size); + return -1; + } + + const bool want_timestamps = is_whisper && options.timestamps; + const uint32_t ts_begin = want_timestamps ? whisper_timestamp_begin(tokenizer) : 0; + + std::vector tokens; + if (prompt && prompt[0] != '\0') { + tokens = tokenizer->encode(prompt); + } else if (is_whisper) { + tokens = handle->model->get_config().decoder_prompt_token_ids; + std::string language = json_string_field(options_json ? options_json : "", "language"); + if (!language.empty()) { + std::vector language_token = tokenizer->encode("<|" + language + "|>"); + if (language_token.size() == 1) { + for (uint32_t& token : tokens) { + std::string piece = tokenizer->decode({token}); + if (piece.size() == 6 && piece[0] == '<' && piece[1] == '|' && + piece[4] == '|' && piece[5] == '>' && + std::islower((unsigned char)piece[2]) && std::islower((unsigned char)piece[3])) { + token = language_token[0]; + break; + } + } + } + } + if (want_timestamps) + tokens.erase(std::remove(tokens.begin(), tokens.end(), ts_begin - 1), tokens.end()); + } + + if (tokens.empty() && is_whisper) { + handle_error_response("Decoder input tokens empty", response_buffer, buffer_size); + return -1; + } + + std::vector> stop_token_sequences = {{tokenizer->get_eos_token()}}; + auto append_exact_stop_sequence = [&](const char* stop_text) { + std::vector seq = tokenizer->encode(stop_text); + if (!seq.empty() && tokenizer->decode(seq) == stop_text) { + stop_token_sequences.push_back(std::move(seq)); + } + }; + append_exact_stop_sequence("<|endoftext|>"); + append_exact_stop_sequence("<|endoftranscript|>"); + append_exact_stop_sequence(""); + append_exact_stop_sequence(""); + + const float audio_length_sec = + static_cast(original_audio_sample_count) / static_cast(cactus::audio::WHISPER_SAMPLE_RATE); + if (options.max_tokens == 100 && (!options_json || std::string(options_json).find("\"max_tokens\"") == std::string::npos)) { + options.max_tokens = std::max(100, static_cast(audio_length_sec * (is_parakeet ? 30.0f : 20.0f))); + } + + constexpr size_t WHISPER_MAX_DECODER_POSITIONS = 448; + if (is_whisper && tokens.size() < WHISPER_MAX_DECODER_POSITIONS) { + options.max_tokens = std::min(options.max_tokens, WHISPER_MAX_DECODER_POSITIONS - tokens.size()); + } + + std::string final_text; + std::vector segments; + std::vector generated_tokens; + generated_tokens.reserve(options.max_tokens); + const size_t prompt_tokens = tokens.size(); + double time_to_first_token = 0.0; + float total_entropy_sum = 0.0f; + + if (is_parakeet) { + generated_tokens = handle->model->transcribe_parakeet_tdt( + audio_features, nullptr, true, 0, &handle->should_stop); + auto t_first = std::chrono::high_resolution_clock::now(); + time_to_first_token = + std::chrono::duration_cast(t_first - start_time).count() / 1000.0; + final_text = tokenizer->decode(generated_tokens); + if (callback) { + for (uint32_t tok : generated_tokens) { + std::string piece = tokenizer->decode({tok}); + callback(piece.c_str(), tok, user_data); + } + } + } else if (is_whisper) { + generated_tokens = handle->model->transcribe_whisper_seq2seq( + audio_features, + tokens, + options.max_tokens, + stop_token_sequences, + &handle->should_stop, + want_timestamps ? static_cast(ts_begin) - 1 : -1); + auto t_first = std::chrono::high_resolution_clock::now(); + time_to_first_token = + std::chrono::duration_cast(t_first - start_time).count() / 1000.0; + for (const auto& stop_seq : stop_token_sequences) { + if (stop_seq.empty() || generated_tokens.size() < stop_seq.size()) continue; + if (std::equal(stop_seq.rbegin(), stop_seq.rend(), generated_tokens.rbegin())) { + generated_tokens.resize(generated_tokens.size() - stop_seq.size()); + break; + } + } + if (want_timestamps) { + segments = parse_whisper_timestamp_segments(tokenizer, generated_tokens, ts_begin, final_text); + } else { + final_text = tokenizer->decode(generated_tokens); + } + if (callback) { + for (uint32_t tok : generated_tokens) { + if (want_timestamps && tok >= ts_begin) continue; + std::string piece = tokenizer->decode({tok}); + callback(piece.c_str(), tok, user_data); + } + } + } else { + for (size_t i = 0; i < options.max_tokens; ++i) { + if (handle->should_stop) break; + + float token_entropy = 0.0f; + uint32_t next_token = handle->model->decode_with_audio( + tokens, std::vector>{audio_features}, + options.temperature, options.top_p, options.top_k, + "", &token_entropy, + options.min_p, options.repetition_penalty, + nullptr, nullptr); + + if (generated_tokens.empty()) { + auto t_first = std::chrono::high_resolution_clock::now(); + time_to_first_token = + std::chrono::duration_cast(t_first - start_time).count() / 1000.0; + } + + total_entropy_sum += token_entropy; + generated_tokens.push_back(next_token); + if (matches_stop_sequence(generated_tokens, stop_token_sequences)) { + break; + } + + std::string piece = tokenizer->decode({next_token}); + if (piece == "<|endoftext|>" || piece == "<|endoftranscript|>" || piece == "" || piece == "") { + break; + } + + tokens.push_back(next_token); + final_text += piece; + if (callback) callback(piece.c_str(), next_token, user_data); + } + } + + handle->model->reset_cache(); + + if (!final_text.empty() && final_text[0] == ' ') { + final_text.erase(0, 1); + } + + const size_t completion_tokens = generated_tokens.size(); + float mean_entropy = completion_tokens > 0 + ? total_entropy_sum / static_cast(completion_tokens) + : 0.0f; + float confidence = 1.0f - mean_entropy; + + bool cloud_handoff_used = false; + std::string handoff_reason; + float reported_threshold = -1.0f; + if (is_parakeet) { + const bool cloud_disabled = env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF"); + const bool cloud_eligible = + !cloud_disabled && !handle->cloud_handoff_disabled && options.auto_handoff; + float threshold = options.confidence_threshold; + if (threshold < 0.0f) { + const float model_default = handle->model->get_config().default_cloud_handoff_threshold; + threshold = handle->model->has_handoff_probe() + ? 0.50f + : (model_default > 0.0f ? model_default : 0.7f); + } + reported_threshold = threshold; + + if (!options.auto_handoff) { + handoff_reason = "handoff off"; + } else if (cloud_disabled) { + handoff_reason = "disabled (env)"; + } else if (handle->cloud_handoff_disabled) { + handoff_reason = "disabled (auth failed earlier)"; + } + + if (cloud_eligible && handle->model->has_handoff_probe_rollout()) { + float p_wrong = handle->model->handoff_probe_wrong_probability(); + if (std::isfinite(p_wrong)) { + confidence = std::max(0.0f, std::min(1.0f, 1.0f - p_wrong)); + CACTUS_LOG_DEBUG("cloud_handoff", "Parakeet handoff probe p_wrong=" + << p_wrong << " confidence=" << confidence); + if (confidence < threshold) { + CACTUS_LOG_WARN("cloud_handoff", "Cloud transcription handoff triggered: p_wrong=" + << p_wrong << " confidence=" << confidence << " threshold=" << threshold); + std::vector pcm16(handoff_audio_samples.size()); + for (size_t i = 0; i < handoff_audio_samples.size(); ++i) { + float clamped = std::max(-1.0f, std::min(1.0f, handoff_audio_samples[i])); + pcm16[i] = static_cast(std::lround(clamped * 32767.0f)); + } + std::vector wav = cactus::ffi::cloud_build_wav( + reinterpret_cast(pcm16.data()), pcm16.size() * sizeof(int16_t)); + std::string audio_b64 = cactus::ffi::cloud_base64_encode(wav.data(), wav.size()); + std::string cloud_key = cactus::ffi::resolve_cloud_api_key(nullptr); + CloudResponse cloud = cactus::ffi::cloud_transcribe_request( + audio_b64, final_text, + static_cast(options.cloud_timeout_ms / 1000), + cloud_key.empty() ? nullptr : cloud_key.c_str()); + if (cloud.used_cloud && !cloud.transcript.empty()) { + final_text = cloud.transcript; + cloud_handoff_used = true; + handoff_reason = "low confidence (probe)"; + } else { + if (cloud.error.rfind("http_401", 0) == 0 || cloud.error.rfind("http_403", 0) == 0) { + handle->cloud_handoff_disabled = true; + } + handoff_reason = "handoff failed: " + cloud.error; + CACTUS_LOG_WARN("cloud_handoff", "Cloud transcription failed, keeping local output: " + << cloud.error); + } + } + } + } + if (handoff_reason.empty()) { + handoff_reason = (confidence >= threshold) ? "above threshold" : "kept local"; + } + } + + auto end_time = std::chrono::high_resolution_clock::now(); + double total_time_ms = + std::chrono::duration_cast(end_time - start_time).count() / 1000.0; + double prefill_tps = time_to_first_token > 0 ? (prompt_tokens * 1000.0) / time_to_first_token : 0.0; + double decode_time_ms = std::max(0.0, total_time_ms - time_to_first_token); + double decode_tps = + (completion_tokens > 1 && decode_time_ms > 0.0) + ? ((completion_tokens - 1) * 1000.0) / decode_time_ms + : 0.0; + + std::string json = construct_response_json( + final_text, {}, time_to_first_token, total_time_ms, + prefill_tps, decode_tps, prompt_tokens, completion_tokens, + confidence, cloud_handoff_used, "", segments, "", + reported_threshold, handoff_reason); + if (json.size() >= buffer_size) { + handle_error_response("Response buffer too small", response_buffer, buffer_size); + return -1; + } + std::strcpy(response_buffer, json.c_str()); + return static_cast(json.size()); + } catch (const std::exception& e) { + CACTUS_LOG_ERROR("transcribe", "Exception: " << e.what()); + handle_error_response(e.what(), response_buffer, buffer_size); + return -1; + } +} + +} // extern "C" diff --git a/cactus-engine/src/utils.h b/cactus-engine/src/utils.h new file mode 100644 index 000000000..20062eb68 --- /dev/null +++ b/cactus-engine/src/utils.h @@ -0,0 +1,2031 @@ +#ifndef CACTUS_UTILS_H +#define CACTUS_UTILS_H + +#include "engine.h" +#include "cactus_kernels.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +#include +#elif defined(_WIN32) +#include +#include +#elif defined(__linux__) || defined(__ANDROID__) +#include +#endif + +inline size_t get_memory_footprint_bytes() { +#ifdef __APPLE__ + task_vm_info_data_t vm_info; + mach_msg_type_number_t count = TASK_VM_INFO_COUNT; + if (task_info(mach_task_self(), TASK_VM_INFO, (task_info_t)&vm_info, &count) == KERN_SUCCESS) + return vm_info.phys_footprint; + +#elif defined(_WIN32) + PROCESS_MEMORY_COUNTERS_EX pmc; + if (GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS*)&pmc, sizeof(pmc))) + return pmc.PrivateUsage; + +#elif defined(__linux__) || defined(__ANDROID__) + std::ifstream statm("/proc/self/statm"); + if (statm.is_open()) { + size_t size, resident; + statm >> size >> resident; + return resident * sysconf(_SC_PAGESIZE); + } +#endif + return 0; +} + +inline double get_ram_usage_mb() { + return get_memory_footprint_bytes() / (1024.0 * 1024.0); +} + +struct CactusModelHandle { + std::unique_ptr model; + std::unique_ptr vad_model; + std::atomic should_stop; + std::vector processed_tokens; + struct ProcessedImage { + std::string path; + long long last_modified_timestamp = 0; + + bool operator==(const ProcessedImage& other) const { + return path == other.path && last_modified_timestamp == other.last_modified_timestamp; + } + }; + + std::vector> processed_images; + std::vector user_audio_counts; + std::mutex model_mutex; + std::string model_name; + std::unique_ptr corpus_index; + std::string corpus_dir; + size_t corpus_embedding_dim = 0; + std::vector> tool_embeddings; + std::vector tool_texts; + bool cloud_handoff_disabled = false; + + CactusModelHandle() : should_stop(false) {} +}; + +extern thread_local std::string last_error_message; + +bool matches_stop_sequence(const std::vector& generated_tokens, + const std::vector>& stop_sequences); + +std::string retrieve_rag_context(CactusModelHandle* handle, const std::string& query); + +namespace cactus { +namespace audio { + +static constexpr size_t WHISPER_TARGET_FRAMES = 3000; +static constexpr int WHISPER_SAMPLE_RATE = 16000; +static constexpr size_t WHISPER_TARGET_SAMPLES = WHISPER_SAMPLE_RATE * 30; + +inline void pad_or_trim_whisper_waveform(std::vector& samples) { + if (samples.size() > WHISPER_TARGET_SAMPLES) { + samples.resize(WHISPER_TARGET_SAMPLES); + } else if (samples.size() < WHISPER_TARGET_SAMPLES) { + samples.resize(WHISPER_TARGET_SAMPLES, 0.0f); + } +} + +inline cactus::engine::SpectrogramConfig get_whisper_spectrogram_config() { + cactus::engine::SpectrogramConfig cfg{}; + cfg.n_fft = 400; + cfg.frame_length = 400; + cfg.hop_length = 160; + cfg.power = 2.0f; + cfg.center = true; + cfg.pad_mode = "reflect"; + cfg.onesided = true; + cfg.dither = 0.0f; + cfg.mel_floor = 1e-10f; + cfg.log_mel = "log10"; + cfg.reference = 1.0f; + cfg.min_value = 1e-10f; + cfg.remove_dc_offset = false; + return cfg; +} + +inline cactus::engine::SpectrogramConfig get_parakeet_spectrogram_config() { + cactus::engine::SpectrogramConfig cfg{}; + cfg.n_fft = 512; + cfg.frame_length = 400; + cfg.hop_length = 160; + cfg.power = 2.0f; + cfg.center = true; + cfg.pad_mode = "constant"; + cfg.onesided = true; + cfg.dither = 0.0f; + cfg.mel_floor = 5.960464477539063e-08f; // 2^-24 guard value used by HF Parakeet. + cfg.log_mel = "log"; + cfg.reference = 1.0f; + cfg.min_value = 1e-10f; + cfg.remove_dc_offset = false; + cfg.hann_periodic = false; + return cfg; +} + +inline cactus::engine::SpectrogramConfig get_htk_spectrogram_config() { + cactus::engine::SpectrogramConfig cfg{}; + cfg.n_fft = 321; + cfg.frame_length = 320; + cfg.fft_override = 1024; + cfg.hop_length = 160; + cfg.power = 1.0f; + cfg.center = false; + cfg.pad_mode = "constant"; + cfg.onesided = true; + cfg.dither = 0.0f; + cfg.mel_floor = 0.001f; + cfg.log_mel = "log"; + cfg.reference = 1.0f; + cfg.min_value = 0.001f; + cfg.remove_dc_offset = false; + cfg.hann_periodic = true; + return cfg; +} + +inline cactus::engine::SpectrogramConfig get_gemma4_audio_spectrogram_config( + const cactus::engine::Config& model_config) { + auto cfg = get_htk_spectrogram_config(); + cfg.fft_override = model_config.audio_fft_length; + cfg.mel_floor_additive = true; + return cfg; +} + +// Whisper v1/v2: 80 mel bins, HTK. Whisper v3: 128 mel bins, Slaney, 512-FFT, no DC removal. +inline void apply_whisper_v3_overrides(cactus::engine::SpectrogramConfig& cfg) { + cfg.fft_override = 512; + cfg.remove_dc_offset = false; +} + +inline std::vector compute_spectrogram_graph( + const std::vector& waveform, + const cactus::engine::SpectrogramConfig& cfg, + size_t mel_bins, + float min_freq, float max_freq, size_t sampling_rate, + int norm_type = 1, int scale_type = 2) { + + size_t fft_len = cfg.fft_override > 0 ? cfg.fft_override : cfg.n_fft; + size_t num_freq_bins = fft_len / 2 + 1; + + int log_mode = 0; + if (cfg.log_mel) { + if (std::strcmp(cfg.log_mel, "log") == 0) log_mode = 1; + else if (std::strcmp(cfg.log_mel, "log10") == 0) log_mode = 2; + else if (std::strcmp(cfg.log_mel, "dB") == 0) log_mode = 3; + } + int pad = (cfg.pad_mode && std::strcmp(cfg.pad_mode, "constant") == 0) ? 1 : 0; + + float effective_floor = cfg.mel_floor_additive ? 0.0f : cfg.mel_floor; + int effective_log = cfg.mel_floor_additive ? 0 : log_mode; + + CactusGraph g; + size_t mel_node = g.mel_filter_bank(num_freq_bins, mel_bins, + min_freq, max_freq, sampling_rate, + norm_type, scale_type); + + size_t wav_input = g.input({waveform.size()}, Precision::FP32); + g.set_input(wav_input, waveform.data(), Precision::FP32); + + size_t spec = g.spectrogram(wav_input, mel_node, + cfg.n_fft, cfg.hop_length, fft_len, + cfg.power, cfg.center, pad, + effective_floor, effective_log, + cfg.dither, cfg.preemphasis, + cfg.remove_dc_offset); + g.execute(); + + const auto& buf = g.get_output_buffer(spec); + float* out_ptr = static_cast(g.get_output(spec)); + std::vector result(out_ptr, out_ptr + buf.total_size); + + if (cfg.mel_floor_additive) { + for (size_t i = 0; i < result.size(); i++) + result[i] = std::log(result[i] + cfg.mel_floor); + } + + return result; +} + +inline std::vector compute_gemma4_audio_spectrogram( + const std::vector& waveform, + const cactus::engine::SpectrogramConfig& cfg, + size_t mel_bins, + float min_freq, + float max_freq, + size_t sampling_rate) { + + size_t fft_len = cfg.fft_override > 0 ? cfg.fft_override : cfg.n_fft; + size_t num_freq_bins = fft_len / 2 + 1; + std::vector mel_filters(mel_bins * num_freq_bins); + cactus_generate_mel_filter_bank( + mel_filters.data(), + static_cast(num_freq_bins), + static_cast(mel_bins), + min_freq, + max_freq, + static_cast(sampling_rate), + nullptr, + "htk", + false); + + // Gemma4's HF feature extractor unfolds 321 samples for frame validity but + // applies a 320-sample periodic Hann window and implicitly zeroes sample 321 + // before the 512-point FFT. Keep that exact split while still using Cactus + // DSP kernels for the frontend. + const size_t analysis_frame_length = cfg.n_fft; + const size_t windowed_frame_length = cfg.frame_length; + std::vector window(analysis_frame_length, 0.0f); + for (size_t i = 0; i < std::min(windowed_frame_length, analysis_frame_length); ++i) { + window[i] = 0.5f * (1.0f - std::cos(2.0f * static_cast(M_PI) * i / + static_cast(windowed_frame_length))); + } + + if (waveform.size() < analysis_frame_length) return {}; + size_t num_frames = 1 + (waveform.size() - analysis_frame_length) / cfg.hop_length; + std::vector result(mel_bins * num_frames); + cactus_compute_spectrogram_f32( + waveform.data(), + waveform.size(), + window.data(), + window.size(), + analysis_frame_length, + cfg.hop_length, + &fft_len, + result.data(), + cfg.power, + cfg.center, + "constant", + true, + cfg.dither, + nullptr, + mel_filters.data(), + mel_filters.size(), + 0.0f, + nullptr, + cfg.reference, + cfg.min_value, + nullptr, + cfg.remove_dc_offset); + + for (float& value : result) { + value = std::log(value + cfg.mel_floor); + } + return result; +} + +// use_mel_floor_padding=true pads short audio with the normalized mel floor (required for v3). +inline std::vector normalize_whisper_mel(std::vector& mel, size_t n_mels, + bool use_mel_floor_padding = false) { + if (mel.empty() || n_mels == 0) return mel; + size_t n_frames = mel.size() / n_mels; + + if (n_frames > WHISPER_TARGET_FRAMES) { + std::vector trimmed(n_mels * WHISPER_TARGET_FRAMES); + for (size_t m = 0; m < n_mels; ++m) { + const float* src = &mel[m * n_frames]; + float* dst = &trimmed[m * WHISPER_TARGET_FRAMES]; + std::copy(src, src + WHISPER_TARGET_FRAMES, dst); + } + mel = std::move(trimmed); + n_frames = WHISPER_TARGET_FRAMES; + } + + float max_val = -std::numeric_limits::infinity(); + for (float v : mel) if (v > max_val) max_val = v; + + float min_allowed = max_val - 8.0f; + for (float& v : mel) { + if (v < min_allowed) v = min_allowed; + v = (v + 4.0f) * 0.25f; + } + + if (n_frames != WHISPER_TARGET_FRAMES) { + float pad_val = use_mel_floor_padding ? (min_allowed + 4.0f) * 0.25f : 0.0f; + std::vector fixed(n_mels * WHISPER_TARGET_FRAMES, pad_val); + size_t copy_frames = std::min(n_frames, WHISPER_TARGET_FRAMES); + for (size_t m = 0; m < n_mels; ++m) { + const float* src = &mel[m * n_frames]; + float* dst = &fixed[m * WHISPER_TARGET_FRAMES]; + std::copy(src, src + copy_frames, dst); + } + return fixed; + } + return std::move(mel); +} + +inline std::vector transpose_mel_to_frame_major(const std::vector& mel, + size_t num_mels, size_t num_frames) { + std::vector transposed(num_frames * num_mels); + for (size_t m = 0; m < num_mels; m++) { + for (size_t t = 0; t < num_frames; t++) { + transposed[t * num_mels + m] = mel[m * num_frames + t]; + } + } + return transposed; +} + +inline void apply_preemphasis(std::vector& waveform, float coefficient = 0.97f) { + if (waveform.size() < 2 || coefficient == 0.0f) { + return; + } + for (size_t i = waveform.size() - 1; i > 0; --i) { + waveform[i] -= coefficient * waveform[i - 1]; + } +} + +inline void normalize_parakeet_log_mel(std::vector& mel, size_t num_mels, float epsilon = 1e-5f) { + if (mel.empty() || num_mels == 0 || (mel.size() % num_mels) != 0) { + return; + } + const size_t num_frames = mel.size() / num_mels; + if (num_frames == 0) { + return; + } + + for (size_t m = 0; m < num_mels; ++m) { + const size_t base = m * num_frames; + float mean = 0.0f; + for (size_t t = 0; t < num_frames; ++t) { + mean += mel[base + t]; + } + mean /= static_cast(num_frames); + + float variance = 0.0f; + for (size_t t = 0; t < num_frames; ++t) { + const float d = mel[base + t] - mean; + variance += d * d; + } + const float denom = static_cast(std::max(1, num_frames - 1)); + const float inv_std = 1.0f / std::sqrt((variance / denom) + epsilon); + for (size_t t = 0; t < num_frames; ++t) { + mel[base + t] = (mel[base + t] - mean) * inv_std; + } + } +} + +inline void trim_mel_frames(std::vector& mel, size_t num_mels, size_t valid_frames) { + if (mel.empty() || num_mels == 0 || (mel.size() % num_mels) != 0) { + return; + } + size_t total_frames = mel.size() / num_mels; + if (valid_frames == 0 || valid_frames >= total_frames) { + return; + } + std::vector trimmed(num_mels * valid_frames); + for (size_t m = 0; m < num_mels; ++m) { + const float* src = &mel[m * total_frames]; + float* dst = &trimmed[m * valid_frames]; + std::copy(src, src + valid_frames, dst); + } + mel.swap(trimmed); +} + +struct AudioPreprocessResult { + std::vector features; + size_t num_frames = 0; + size_t num_soft_tokens = 0; +}; + +inline AudioPreprocessResult preprocess_audio_for_gemma4( + std::vector audio_samples, + const cactus::engine::Config& model_config +) { + AudioPreprocessResult result; + if (audio_samples.empty()) return result; + + size_t pad_amt = 320 - (audio_samples.size() % 320); + if (pad_amt < 320) + audio_samples.resize(audio_samples.size() + pad_amt, 0.0f); + + size_t mel_bins = model_config.audio_input_feat_size; + auto cfg = get_gemma4_audio_spectrogram_config(model_config); + + size_t semicausal_pad = cfg.frame_length / 2; + audio_samples.insert(audio_samples.begin(), semicausal_pad, 0.0f); + + std::vector mel = compute_gemma4_audio_spectrogram( + audio_samples, cfg, mel_bins, 0.0f, 8000.0f, 16000); + + result.num_frames = mel.size() / mel_bins; + result.features = transpose_mel_to_frame_major(mel, mel_bins, result.num_frames); + + size_t max_soft = model_config.audio_soft_tokens; + size_t max_frames_per_chunk = max_soft > 0 ? (max_soft * 2 - 1) * 2 - 1 : 0; + if (max_soft > 0 && result.num_frames > max_frames_per_chunk) { + size_t num_chunks = (result.num_frames + max_frames_per_chunk - 1) / max_frames_per_chunk; + result.num_soft_tokens = num_chunks * max_soft; + } else { + size_t after_stage1 = (result.num_frames + 1) / 2; + result.num_soft_tokens = (after_stage1 + 1) / 2; + } + + return result; +} + +inline std::vector pcm_buffer_to_float_samples( + const uint8_t* pcm_buffer, size_t pcm_buffer_size +) { + const int16_t* pcm_samples = reinterpret_cast(pcm_buffer); + size_t num_samples = pcm_buffer_size / 2; + std::vector waveform_fp32(num_samples); + constexpr float inv_32768 = 1.0f / 32768.0f; + for (size_t i = 0; i < num_samples; i++) + waveform_fp32[i] = static_cast(pcm_samples[i]) * inv_32768; + return waveform_fp32; +} + +} // namespace audio +} // namespace cactus + +namespace cactus { +namespace ffi { + +inline bool env_flag_enabled(const char* key) { + const char* value = std::getenv(key); + return value && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0'); +} + +inline std::string generateUUID() { +#ifdef __APPLE__ + uuid_t uuid; + uuid_generate_random(uuid); + char uuid_str[37]; + uuid_unparse_lower(uuid, uuid_str); + return std::string(uuid_str); +#else + static std::random_device rd; + static std::mt19937 gen(rd()); + static std::uniform_int_distribution<> dis(0, 15); + static std::uniform_int_distribution<> dis2(8, 11); + + std::stringstream ss; + ss << std::hex; + for (int i = 0; i < 8; i++) ss << dis(gen); + ss << "-"; + for (int i = 0; i < 4; i++) ss << dis(gen); + ss << "-4"; + for (int i = 0; i < 3; i++) ss << dis(gen); + ss << "-"; + ss << dis2(gen); + for (int i = 0; i < 3; i++) ss << dis(gen); + ss << "-"; + for (int i = 0; i < 12; i++) ss << dis(gen); + return ss.str(); +#endif +} + +struct ToolFunction { + std::string name; + std::string description; + std::unordered_map parameters; +}; + +struct InferenceOptions { + float temperature = 0.0f; + float top_p = 0.0f; + float min_p = 0.15f; + float repetition_penalty = 1.1f; + uint64_t sample_seed = 0; + float confidence_threshold = -1.0f; + size_t top_k = 0; + size_t max_tokens = 100; + size_t tool_rag_top_k = 2; + size_t cloud_timeout_ms = 15000; + std::vector stop_sequences; + bool force_tools = false; + bool include_stop_sequences = false; + bool use_vad = true; + bool timestamps = false; + bool telemetry_enabled = true; + bool auto_handoff = true; + bool handoff_with_images = true; + bool enable_thinking_if_supported = false; +}; + +} // namespace ffi +} // namespace cactus + +std::vector select_relevant_tools( + CactusModelHandle* handle, + const std::string& query, + const std::vector& all_tools, + size_t top_k); + +#include "gemma_tools.h" +#include "chat_tools.h" + +namespace cactus { +namespace ffi { + +// Defined in json_escape.h — included here for backward compatibility +#include "json_escape.h" + + +inline std::string trim_string(const std::string& s) { + size_t start = 0; + while (start < s.size() && std::isspace(static_cast(s[start]))) ++start; + size_t end = s.size(); + while (end > start && std::isspace(static_cast(s[end - 1]))) --end; + return s.substr(start, end - start); +} + +inline size_t find_matching_delimiter(const std::string& s, size_t pos, char open, char close) { + int depth = 1; + pos++; + while (pos < s.length() && depth > 0) { + if (s[pos] == open) depth++; + else if (s[pos] == close) depth--; + else if (s[pos] == '"') { + pos++; + while (pos < s.length() && s[pos] != '"') { + if (s[pos] == '\\') pos++; + pos++; + } + } + pos++; + } + return pos; +} + +inline std::string env_or_default(const char* key, const char* fallback) { + const char* v = std::getenv(key); + if (v && v[0] != '\0') return std::string(v); + return std::string(fallback); +} + +inline std::string json_string_field(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return {}; + + size_t i = pos + pattern.size(); + while (i < json.size() && std::isspace(static_cast(json[i]))) i++; + if (i >= json.size() || json[i] != '"') return {}; + ++i; + + std::string out; + out.reserve(128); + while (i < json.size()) { + char c = json[i++]; + if (c == '"') return out; + if (c == '\\' && i < json.size()) { + char e = json[i++]; + switch (e) { + case '"': out.push_back('"'); break; + case '\\': out.push_back('\\'); break; + case '/': out.push_back('/'); break; + case 'b': out.push_back('\b'); break; + case 'f': out.push_back('\f'); break; + case 'n': out.push_back('\n'); break; + case 'r': out.push_back('\r'); break; + case 't': out.push_back('\t'); break; + case 'u': { + unsigned int cp = 0; + bool ok = i + 4 <= json.size(); + for (int k = 0; ok && k < 4; k++) { + char h = json[i + k]; + cp <<= 4; + if (h >= '0' && h <= '9') cp |= static_cast(h - '0'); + else if (h >= 'a' && h <= 'f') cp |= static_cast(h - 'a' + 10); + else if (h >= 'A' && h <= 'F') cp |= static_cast(h - 'A' + 10); + else ok = false; + } + if (ok) { + i += 4; + if (cp < 0x80) { + out.push_back(static_cast(cp)); + } else if (cp < 0x800) { + out.push_back(static_cast(0xC0 | (cp >> 6))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else { + out.push_back(static_cast(0xE0 | (cp >> 12))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } + } else { + out.push_back('u'); + } + break; + } + default: out.push_back(e); break; + } + continue; + } + out.push_back(c); + } + return {}; +} + +inline std::string json_array_field(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return "[]"; + size_t start = pos + pattern.size(); + while (start < json.size() && std::isspace(static_cast(json[start]))) ++start; + if (start >= json.size() || json[start] != '[') return "[]"; + + int depth = 1; + size_t end = start + 1; + bool in_string = false, escape = false; + while (end < json.size() && depth > 0) { + const char c = json[end]; + if (in_string) { + if (escape) escape = false; + else if (c == '\\') escape = true; + else if (c == '"') in_string = false; + } else if (c == '"') { + in_string = true; + } else if (c == '[') { + depth++; + } else if (c == ']') { + depth--; + } + end++; + } + return json.substr(start, end - start); +} + +inline std::vector split_json_array(const std::string& array_json) { + std::vector out; + if (array_json.size() < 2 || array_json.front() != '[' || array_json.back() != ']') return out; + + size_t i = 1; + while (i + 1 < array_json.size()) { + while (i + 1 < array_json.size() && + (std::isspace(static_cast(array_json[i])) || array_json[i] == ',')) i++; + if (i + 1 >= array_json.size() || array_json[i] != '{') break; + + size_t start = i; + int depth = 0; + bool in_str = false; + bool esc = false; + for (; i < array_json.size(); ++i) { + char c = array_json[i]; + if (in_str) { + if (esc) esc = false; + else if (c == '\\') esc = true; + else if (c == '"') in_str = false; + continue; + } + if (c == '"') { in_str = true; continue; } + if (c == '{') depth++; + if (c == '}') { + depth--; + if (depth == 0) { + out.push_back(array_json.substr(start, i - start + 1)); + i++; + break; + } + } + } + } + return out; +} + +namespace jsonval { + +inline void skip_ws(const std::string& s, size_t& p) { + while (p < s.size() && (s[p] == ' ' || s[p] == '\t' || s[p] == '\n' || s[p] == '\r')) p++; +} + +inline bool parse_value(const std::string& s, size_t& p); + +inline bool parse_string(const std::string& s, size_t& p) { + if (p >= s.size() || s[p] != '"') return false; + p++; + while (p < s.size()) { + char c = s[p++]; + if (c == '"') return true; + if (c == '\\') { + if (p >= s.size()) return false; + char e = s[p++]; + if (e == 'u') { + if (p + 4 > s.size()) return false; + for (int i = 0; i < 4; i++) { + if (!std::isxdigit(static_cast(s[p++]))) return false; + } + } else if (e != '"' && e != '\\' && e != '/' && e != 'b' && + e != 'f' && e != 'n' && e != 'r' && e != 't') { + return false; + } + } else if (static_cast(c) < 0x20) { + return false; + } + } + return false; +} + +inline bool parse_number(const std::string& s, size_t& p) { + size_t start = p; + if (p < s.size() && s[p] == '-') p++; + if (p < s.size() && s[p] == '0') { + p++; + } else if (p < s.size() && s[p] >= '1' && s[p] <= '9') { + while (p < s.size() && std::isdigit(static_cast(s[p]))) p++; + } else { + return false; + } + if (p < s.size() && s[p] == '.') { + p++; + if (p >= s.size() || !std::isdigit(static_cast(s[p]))) return false; + while (p < s.size() && std::isdigit(static_cast(s[p]))) p++; + } + if (p < s.size() && (s[p] == 'e' || s[p] == 'E')) { + p++; + if (p < s.size() && (s[p] == '+' || s[p] == '-')) p++; + if (p >= s.size() || !std::isdigit(static_cast(s[p]))) return false; + while (p < s.size() && std::isdigit(static_cast(s[p]))) p++; + } + return p > start; +} + +inline bool parse_literal(const std::string& s, size_t& p, const std::string& lit) { + if (s.compare(p, lit.size(), lit) == 0) { + p += lit.size(); + return true; + } + return false; +} + +inline bool parse_object(const std::string& s, size_t& p) { + p++; + skip_ws(s, p); + if (p < s.size() && s[p] == '}') { p++; return true; } + while (true) { + skip_ws(s, p); + if (!parse_string(s, p)) return false; + skip_ws(s, p); + if (p >= s.size() || s[p] != ':') return false; + p++; + if (!parse_value(s, p)) return false; + skip_ws(s, p); + if (p >= s.size()) return false; + if (s[p] == ',') { p++; continue; } + if (s[p] == '}') { p++; return true; } + return false; + } +} + +inline bool parse_array(const std::string& s, size_t& p) { + p++; + skip_ws(s, p); + if (p < s.size() && s[p] == ']') { p++; return true; } + while (true) { + if (!parse_value(s, p)) return false; + skip_ws(s, p); + if (p >= s.size()) return false; + if (s[p] == ',') { p++; continue; } + if (s[p] == ']') { p++; return true; } + return false; + } +} + +inline bool parse_value(const std::string& s, size_t& p) { + skip_ws(s, p); + if (p >= s.size()) return false; + char c = s[p]; + if (c == '{') return parse_object(s, p); + if (c == '[') return parse_array(s, p); + if (c == '"') return parse_string(s, p); + if (c == 't') return parse_literal(s, p, "true"); + if (c == 'f') return parse_literal(s, p, "false"); + if (c == 'n') return parse_literal(s, p, "null"); + return parse_number(s, p); +} + +} // namespace jsonval + +inline bool is_valid_json(const std::string& s) { + size_t p = 0; + if (!jsonval::parse_value(s, p)) return false; + jsonval::skip_ws(s, p); + return p == s.size(); +} + +inline void sanitize_function_calls(std::vector& calls) { + std::vector valid; + valid.reserve(calls.size()); + for (const auto& call : calls) { + std::string trimmed = trim_string(call); + if (!trimmed.empty() && trimmed.front() == '{' && + trimmed.find("\"name\"") != std::string::npos && is_valid_json(trimmed)) { + valid.push_back(trimmed); + } + } + calls.swap(valid); +} + +inline std::string serialize_tools_json(const std::vector& tools) { + if (tools.empty()) return ""; + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < tools.size(); ++i) { + if (i > 0) oss << ","; + oss << "{\"type\":\"function\",\"function\":{"; + oss << "\"name\":\"" << escape_json_string(tools[i].name) << "\","; + oss << "\"description\":\"" << escape_json_string(tools[i].description) << "\""; + auto it = tools[i].parameters.find("schema"); + if (it != tools[i].parameters.end()) { + oss << ",\"parameters\":" << it->second; + } + oss << "}}"; + } + oss << "]"; + return oss.str(); +} + +namespace json_sorted { + +inline void skip_ws(const std::string& s, size_t& p) { + while (p < s.size() && std::isspace(static_cast(s[p]))) p++; +} + +inline std::string parse_string(const std::string& s, size_t& p) { + std::string r = "\""; + p++; + while (p < s.size()) { + if (s[p] == '\\') { + r += s[p++]; + if (p < s.size()) r += s[p++]; + } else if (s[p] == '"') { + r += '"'; + p++; + return r; + } else { + r += s[p++]; + } + } + return r; +} + +inline std::string parse_value(const std::string& s, size_t& p); + +inline std::string parse_object(const std::string& s, size_t& p) { + p++; + std::map entries; + skip_ws(s, p); + while (p < s.size() && s[p] != '}') { + if (s[p] == ',') { p++; skip_ws(s, p); continue; } + std::string key = parse_string(s, p); + skip_ws(s, p); + if (p < s.size() && s[p] == ':') p++; + skip_ws(s, p); + std::string val = parse_value(s, p); + entries[key] = val; + skip_ws(s, p); + } + if (p < s.size()) p++; + std::string r = "{"; + bool first = true; + for (const auto& kv : entries) { + if (!first) r += ", "; + r += kv.first + ": " + kv.second; + first = false; + } + r += "}"; + return r; +} + +inline std::string parse_array(const std::string& s, size_t& p) { + p++; + std::vector items; + skip_ws(s, p); + while (p < s.size() && s[p] != ']') { + if (s[p] == ',') { p++; skip_ws(s, p); continue; } + items.push_back(parse_value(s, p)); + skip_ws(s, p); + } + if (p < s.size()) p++; + std::string r = "["; + for (size_t i = 0; i < items.size(); i++) { + if (i > 0) r += ", "; + r += items[i]; + } + r += "]"; + return r; +} + +inline std::string parse_value(const std::string& s, size_t& p) { + skip_ws(s, p); + if (p >= s.size()) return ""; + if (s[p] == '"') return parse_string(s, p); + if (s[p] == '{') return parse_object(s, p); + if (s[p] == '[') return parse_array(s, p); + size_t start = p; + while (p < s.size() && s[p] != ',' && s[p] != '}' && s[p] != ']' && !std::isspace(static_cast(s[p]))) p++; + return s.substr(start, p - start); +} + +inline std::string reformat(const std::string& json) { + size_t p = 0; + return parse_value(json, p); +} + +} // namespace json_sorted + +inline std::string serialize_tools_for_template(const std::vector& tools) { + if (tools.empty()) return ""; + std::string result; + for (const auto& tool : tools) { + std::map func_fields; + func_fields["\"description\""] = "\"" + escape_json_string(tool.description) + "\""; + func_fields["\"name\""] = "\"" + escape_json_string(tool.name) + "\""; + auto it = tool.parameters.find("schema"); + if (it != tool.parameters.end()) { + func_fields["\"parameters\""] = json_sorted::reformat(it->second); + } + std::string func_json = "{"; + bool first = true; + for (const auto& kv : func_fields) { + if (!first) func_json += ", "; + func_json += kv.first + ": " + kv.second; + first = false; + } + func_json += "}"; + result += "\n{\"function\": " + func_json + ", \"type\": \"function\"}"; + } + return result; +} + +inline void handle_error_response(const std::string& error_message, char* response_buffer, size_t buffer_size) { + std::ostringstream json; + json << "{"; + json << "\"success\":false,"; + json << "\"error\":\"" << escape_json_string(error_message) << "\","; + json << "\"cloud_handoff\":false,"; + json << "\"response\":null,"; + json << "\"function_calls\":[],"; + json << "\"confidence\":null,"; + json << "\"time_to_first_token_ms\":0.0,"; + json << "\"total_time_ms\":0.0,"; + json << "\"prefill_tps\":0.0,"; + json << "\"decode_tps\":0.0,"; + json << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() << ","; + json << "\"prefill_tokens\":0,"; + json << "\"decode_tokens\":0,"; + json << "\"total_tokens\":0"; + json << "}"; + std::string error_json = json.str(); + if (response_buffer && error_json.length() < buffer_size) { + std::strcpy(response_buffer, error_json.c_str()); + } +} + +inline std::vector parse_messages_json(const std::string& json, + std::vector& out_image_paths, + std::vector* out_audio_paths = nullptr) { + std::vector messages; + out_image_paths.clear(); + if (out_audio_paths) out_audio_paths->clear(); + + size_t pos = json.find('['); + if (pos == std::string::npos) { + throw std::runtime_error("Invalid JSON: expected array"); + } + + pos = json.find('{', pos); + while (pos != std::string::npos) { + cactus::engine::ChatMessage msg; + + size_t obj_start = pos; + int brace_count = 1; + size_t obj_end = obj_start + 1; + while (obj_end < json.length() && brace_count > 0) { + if (json[obj_end] == '{') brace_count++; + else if (json[obj_end] == '}') brace_count--; + obj_end++; + } + + size_t role_pos = json.find("\"role\"", pos); + if (role_pos == std::string::npos || role_pos >= obj_end) break; + + size_t role_start = json.find('"', role_pos + 6) + 1; + size_t role_end = json.find('"', role_start); + msg.role = json.substr(role_start, role_end - role_start); + + size_t content_pos = json.find("\"content\"", obj_start); + if (content_pos != std::string::npos && content_pos < obj_end) { + size_t content_start = json.find('"', content_pos + 9) + 1; + size_t content_end = content_start; + + while (content_end < json.length()) { + content_end = json.find('"', content_end); + if (content_end == std::string::npos) break; + if (json[content_end - 1] != '\\') break; + content_end++; + } + + msg.content = json.substr(content_start, content_end - content_start); + + size_t escape_pos = 0; + while ((escape_pos = msg.content.find("\\n", escape_pos)) != std::string::npos) { + msg.content.replace(escape_pos, 2, "\n"); + escape_pos += 1; + } + escape_pos = 0; + while ((escape_pos = msg.content.find("\\\"", escape_pos)) != std::string::npos) { + msg.content.replace(escape_pos, 2, "\""); + escape_pos += 1; + } + } + + auto parse_path_array = [&](const char* key, std::vector& dest, + std::vector* out_paths) { + size_t key_pos = json.find(key, pos); + if (key_pos == std::string::npos || key_pos >= obj_end) return; + size_t array_start = json.find('[', key_pos); + if (array_start == std::string::npos || array_start >= obj_end) return; + size_t array_end = json.find(']', array_start); + if (array_end == std::string::npos || array_end >= obj_end) return; + size_t cur = array_start; + while (true) { + cur = json.find('"', cur + 1); + if (cur == std::string::npos || cur >= array_end) break; + size_t str_start = cur + 1; + size_t str_end = json.find('"', str_start); + if (str_end == std::string::npos || str_end > array_end) break; + std::string path = std::filesystem::absolute( + std::filesystem::path(json.substr(str_start, str_end - str_start))).string(); + dest.push_back(path); + if (out_paths) out_paths->push_back(path); + cur = str_end; + } + }; + + parse_path_array("\"images\"", msg.images, &out_image_paths); + parse_path_array("\"audio\"", msg.audio, out_audio_paths); + + if (msg.role == "tool") { + size_t name_pos = json.find("\"name\"", obj_start); + if (name_pos != std::string::npos && name_pos < obj_end) { + size_t name_quote = json.find('"', name_pos + 6); + if (name_quote != std::string::npos && name_quote < obj_end) { + size_t name_start = name_quote + 1; + size_t name_end = json.find('"', name_start); + if (name_end != std::string::npos && name_end < obj_end) { + msg.name = json.substr(name_start, name_end - name_start); + } + } + } + } + + size_t tool_calls_pos = json.find("\"tool_calls\"", obj_start); + if (tool_calls_pos != std::string::npos && tool_calls_pos < obj_end) { + size_t tool_calls_arr_start = json.find('[', tool_calls_pos); + if (tool_calls_arr_start != std::string::npos && tool_calls_arr_start < obj_end) { + size_t tool_calls_arr_end = find_matching_delimiter(json, tool_calls_arr_start, '[', ']'); + + size_t search_pos = tool_calls_arr_start; + while (true) { + size_t func_pos = json.find("\"function\"", search_pos); + if (func_pos == std::string::npos || func_pos >= tool_calls_arr_end) break; + + size_t func_obj_start = json.find('{', func_pos + 10); + if (func_obj_start == std::string::npos || func_obj_start >= tool_calls_arr_end) break; + + size_t func_obj_end = find_matching_delimiter(json, func_obj_start, '{', '}'); + + cactus::engine::ToolCallInfo tool_call; + + size_t fn_name_pos = json.find("\"name\"", func_obj_start); + if (fn_name_pos != std::string::npos && fn_name_pos < func_obj_end) { + size_t fn_name_quote = json.find('"', fn_name_pos + 6); + if (fn_name_quote != std::string::npos && fn_name_quote < func_obj_end) { + size_t fn_name_start = fn_name_quote + 1; + size_t fn_name_end = json.find('"', fn_name_start); + if (fn_name_end != std::string::npos && fn_name_end < func_obj_end) { + tool_call.name = json.substr(fn_name_start, fn_name_end - fn_name_start); + } + } + } + + size_t args_pos = json.find("\"arguments\"", func_obj_start); + if (args_pos != std::string::npos && args_pos < func_obj_end) { + size_t colon_pos = json.find(':', args_pos + 11); + if (colon_pos != std::string::npos && colon_pos < func_obj_end) { + size_t args_start = colon_pos + 1; + while (args_start < json.length() && std::isspace(static_cast(json[args_start]))) args_start++; + + if (args_start < func_obj_end && json[args_start] == '{') { + size_t args_end = find_matching_delimiter(json, args_start, '{', '}'); + tool_call.arguments = json.substr(args_start, args_end - args_start); + } else if (args_start < func_obj_end && json[args_start] == '"') { + size_t str_start = args_start + 1; + size_t str_end = str_start; + while (str_end < json.length() && json[str_end] != '"') { + if (json[str_end] == '\\') str_end++; + str_end++; + } + tool_call.arguments = json.substr(str_start, str_end - str_start); + } + } + } + + if (!tool_call.name.empty()) { + msg.tool_calls.push_back(tool_call); + } + search_pos = func_obj_end; + } + } + } + + messages.push_back(msg); + + pos = json.find('{', obj_end); + } + + return messages; +} + +inline std::vector parse_tools_json(const std::string& json) { + std::vector tools; + + if (json.empty()) return tools; + + size_t pos = json.find('['); + if (pos == std::string::npos) return tools; + + pos = json.find("\"function\"", pos); + while (pos != std::string::npos) { + ToolFunction tool; + size_t next_search = pos + 1; + + size_t name_pos = json.find("\"name\"", pos); + if (name_pos != std::string::npos) { + size_t name_start = json.find('"', name_pos + 6) + 1; + size_t name_end = json.find('"', name_start); + tool.name = json.substr(name_start, name_end - name_start); + } + + size_t desc_pos = json.find("\"description\"", pos); + if (desc_pos != std::string::npos) { + size_t desc_start = json.find('"', desc_pos + 13) + 1; + size_t desc_end = json.find('"', desc_start); + tool.description = json.substr(desc_start, desc_end - desc_start); + } + + size_t params_pos = json.find("\"parameters\"", pos); + if (params_pos != std::string::npos) { + size_t params_start = json.find('{', params_pos); + if (params_start != std::string::npos) { + int brace_count = 1; + size_t params_end = params_start + 1; + while (params_end < json.length() && brace_count > 0) { + if (json[params_end] == '{') brace_count++; + else if (json[params_end] == '}') brace_count--; + params_end++; + } + tool.parameters["schema"] = json.substr(params_start, params_end - params_start); + next_search = params_end; + } + } + + if (!tool.name.empty()) { + tools.push_back(tool); + } + + pos = json.find("\"function\"", next_search); + } + + return tools; +} + +inline bool try_parse_json_float(const std::string& json, const std::string& key, float& out_value) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return false; + + size_t start = pos + pattern.size(); + while (start < json.size() && std::isspace(static_cast(json[start]))) ++start; + + size_t end = start; + while (end < json.size() && std::string(",}] \t\n\r").find(json[end]) == std::string::npos) ++end; + + try { + out_value = std::stof(json.substr(start, end - start)); + return true; + } catch (...) { + return false; + } +} + +inline bool try_parse_json_uint(const std::string& json, const std::string& key, size_t& out_value) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return false; + + size_t start = pos + pattern.size(); + while (start < json.size() && std::isspace(static_cast(json[start]))) ++start; + + size_t end = start; + while (end < json.size() && std::string(",}] \t\n\r").find(json[end]) == std::string::npos) ++end; + + try { + long long parsed = std::stoll(json.substr(start, end - start)); + if (parsed < 0) return false; + out_value = static_cast(parsed); + return true; + } catch (...) { + return false; + } +} + +inline std::vector parse_json_string_array_field(const std::string& json, const std::string& key) { + std::vector out; + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return out; + + size_t start = pos + pattern.size(); + while (start < json.size() && std::isspace(static_cast(json[start]))) ++start; + if (start >= json.size() || json[start] != '[') return out; + + int depth = 1; + bool in_string = false; + bool escaped = false; + size_t end = start + 1; + + while (end < json.size() && depth > 0) { + char c = json[end]; + if (in_string) { + if (escaped) escaped = false; + else if (c == '\\') escaped = true; + else if (c == '"') in_string = false; + } else { + if (c == '"') in_string = true; + else if (c == '[') depth++; + else if (c == ']') depth--; + } + ++end; + } + + if (depth != 0) return out; + const std::string array_json = json.substr(start, end - start); + if (array_json.size() < 2 || array_json.front() != '[' || array_json.back() != ']') return out; + + size_t i = 1; + while (i + 1 < array_json.size()) { + while (i + 1 < array_json.size() && + (std::isspace(static_cast(array_json[i])) || array_json[i] == ',')) { + ++i; + } + if (i + 1 >= array_json.size() || array_json[i] == ']') break; + if (array_json[i] != '"') break; + + ++i; + std::string value; + bool escaped = false; + while (i < array_json.size()) { + char c = array_json[i++]; + if (escaped) { + switch (c) { + case '"': value.push_back('"'); break; + case '\\': value.push_back('\\'); break; + case '/': value.push_back('/'); break; + case 'b': value.push_back('\b'); break; + case 'f': value.push_back('\f'); break; + case 'n': value.push_back('\n'); break; + case 'r': value.push_back('\r'); break; + case 't': value.push_back('\t'); break; + case 'u': { + unsigned int cp = 0; + bool ok = i + 4 <= array_json.size(); + for (int k = 0; ok && k < 4; k++) { + char h = array_json[i + k]; + cp <<= 4; + if (h >= '0' && h <= '9') cp |= static_cast(h - '0'); + else if (h >= 'a' && h <= 'f') cp |= static_cast(h - 'a' + 10); + else if (h >= 'A' && h <= 'F') cp |= static_cast(h - 'A' + 10); + else ok = false; + } + if (ok) { + i += 4; + if (cp < 0x80) { + value.push_back(static_cast(cp)); + } else if (cp < 0x800) { + value.push_back(static_cast(0xC0 | (cp >> 6))); + value.push_back(static_cast(0x80 | (cp & 0x3F))); + } else { + value.push_back(static_cast(0xE0 | (cp >> 12))); + value.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + value.push_back(static_cast(0x80 | (cp & 0x3F))); + } + } else { + value.push_back('u'); + } + break; + } + default: value.push_back(c); break; + } + escaped = false; + continue; + } + if (c == '\\') { + escaped = true; + continue; + } + if (c == '"') { + out.push_back(value); + break; + } + value.push_back(c); + } + } + + return out; +} + +inline void parse_custom_vocabulary_options(const std::string& json, + std::vector& custom_vocabulary, + float& vocabulary_boost) { + custom_vocabulary.clear(); + vocabulary_boost = 5.0f; + if (json.empty()) return; + + float parsed_boost = vocabulary_boost; + if (try_parse_json_float(json, "vocabulary_boost", parsed_boost)) { + vocabulary_boost = std::clamp(parsed_boost, 0.0f, 20.0f); + } + + custom_vocabulary = parse_json_string_array_field(json, "custom_vocabulary"); +} + +inline std::unordered_map build_token_bias_map(const std::vector>& tokenized_entries, + float vocabulary_boost) { + std::unordered_map vocab_bias; + const float clamped_boost = std::clamp(vocabulary_boost, 0.0f, 20.0f); + if (clamped_boost == 0.0f) return vocab_bias; + + for (const auto& token_ids : tokenized_entries) { + for (uint32_t token_id : token_ids) { + float& entry = vocab_bias[token_id]; + if (entry < clamped_boost) { + entry = clamped_boost; + } + } + } + + return vocab_bias; +} + +inline std::unordered_map build_custom_vocabulary_bias(cactus::engine::Tokenizer* tokenizer, + const std::vector& custom_vocabulary, + float vocabulary_boost) { + if (!tokenizer || custom_vocabulary.empty()) return {}; + std::vector> tokenized_entries; + tokenized_entries.reserve(custom_vocabulary.size()); + + for (const auto& word : custom_vocabulary) { + if (word.empty()) continue; + tokenized_entries.push_back(tokenizer->encode(word)); + } + + return build_token_bias_map(tokenized_entries, vocabulary_boost); +} + +inline void apply_custom_vocabulary_options(cactus::engine::Model* model, const std::string& json) { + if (!model) return; + + std::vector custom_vocabulary; + float vocabulary_boost = 5.0f; + parse_custom_vocabulary_options(json, custom_vocabulary, vocabulary_boost); + model->set_vocab_bias(build_custom_vocabulary_bias(model->get_tokenizer(), custom_vocabulary, vocabulary_boost)); +} + +inline size_t levenshtein_ci(const std::string& a, const std::string& b) { + const size_t m = a.size(), n = b.size(); + std::vector prev(n + 1), curr(n + 1); + for (size_t j = 0; j <= n; ++j) prev[j] = j; + for (size_t i = 1; i <= m; ++i) { + curr[0] = i; + for (size_t j = 1; j <= n; ++j) { + const bool match = std::tolower(static_cast(a[i - 1])) == + std::tolower(static_cast(b[j - 1])); + curr[j] = std::min({prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + (match ? 0 : 1)}); + } + std::swap(prev, curr); + } + return prev[n]; +} + +inline std::string collapse_spaces(const std::string& s) { + std::string out; + out.reserve(s.size()); + for (char c : s) { + if (c != ' ') out += c; + } + return out; +} + +inline void apply_vocabulary_spelling_correction( + std::string& text, + const std::vector& custom_vocabulary) +{ + if (custom_vocabulary.empty() || text.empty()) return; + + struct VocabEntry { + const std::string* original; + std::string collapsed; + }; + std::vector vocab_entries; + vocab_entries.reserve(custom_vocabulary.size()); + for (const auto& v : custom_vocabulary) { + vocab_entries.push_back({&v, collapse_spaces(v)}); + } + + struct Token { std::string text; bool is_word; }; + std::vector tokens; + size_t pos = 0; + while (pos < text.size()) { + if (std::isalnum(static_cast(text[pos])) || + text[pos] == '\'' || text[pos] == '-') { + size_t start = pos; + while (pos < text.size() && (std::isalnum(static_cast(text[pos])) || + text[pos] == '\'' || text[pos] == '-')) { + ++pos; + } + tokens.push_back({text.substr(start, pos - start), true}); + } else { + size_t start = pos; + while (pos < text.size() && !std::isalnum(static_cast(text[pos])) && + text[pos] != '\'' && text[pos] != '-') { + ++pos; + } + tokens.push_back({text.substr(start, pos - start), false}); + } + } + + std::vector word_indices; + for (size_t i = 0; i < tokens.size(); ++i) { + if (tokens[i].is_word) word_indices.push_back(i); + } + + std::vector consumed(tokens.size(), false); + + auto strip_suffix = [](const std::string& word) -> std::pair { + if (word.size() >= 3 && word.substr(word.size() - 2) == "'s") { + return {word.substr(0, word.size() - 2), "'s"}; + } + if (word.size() >= 3 && word.substr(word.size() - 2) == "'t") { + return {word.substr(0, word.size() - 2), "'t"}; + } + if (word.size() >= 4 && word.back() == 's' && + word[word.size() - 2] != 's' && // avoid stripping from "boss", "class" + std::isalpha(static_cast(word[word.size() - 2]))) { + return {word.substr(0, word.size() - 1), "s"}; + } + return {word, ""}; + }; + + size_t wi = 0; + while (wi < word_indices.size()) { + size_t best_dist = std::numeric_limits::max(); + const std::string* best_match = nullptr; + size_t best_window = 0; + size_t best_first_token = 0; + size_t best_last_token = 0; + std::string best_suffix; + + for (size_t window = std::min(3, word_indices.size() - wi); window >= 1; --window) { + std::string window_collapsed; + const size_t first_tok = word_indices[wi]; + const size_t last_tok = word_indices[wi + window - 1]; + for (size_t w = 0; w < window; ++w) { + window_collapsed += tokens[word_indices[wi + w]].text; + } + + if (window == 1 && window_collapsed.size() < 3) break; + + auto [stem, suffix] = strip_suffix(window_collapsed); + const std::string* candidates[] = {&window_collapsed, &stem}; + const std::string suffixes[] = {"", suffix}; + const size_t num_candidates = suffix.empty() ? 1 : 2; + + for (size_t ci = 0; ci < num_candidates; ++ci) { + const std::string& candidate = *candidates[ci]; + if (candidate.empty()) continue; + + for (const auto& entry : vocab_entries) { + const size_t wlen = candidate.size(); + const size_t vlen = entry.collapsed.size(); + + const size_t len_diff = wlen > vlen ? wlen - vlen : vlen - wlen; + const size_t max_dist = std::max(1, std::min(wlen, vlen) / 3); + if (len_diff > max_dist) continue; + + const size_t dist = levenshtein_ci(candidate, entry.collapsed); + + // For single-edit corrections, require first char match to prevent + // false positives like "vortex" → "Cortex". + if (dist == 1 && window == 1) { + const bool first_char_match = + std::tolower(static_cast(candidate[0])) == + std::tolower(static_cast(entry.collapsed[0])); + if (!first_char_match) continue; + } + + if (dist <= max_dist && dist < best_dist) { + best_dist = dist; + best_match = entry.original; + best_window = window; + best_first_token = first_tok; + best_last_token = last_tok; + best_suffix = suffixes[ci]; + } + } + } + + if (best_dist == 0) break; + } + + // Allow dist==0 for multi-word merges where word boundaries changed. + const bool should_replace = best_match && + best_dist != std::numeric_limits::max() && + (best_dist > 0 || best_window > 1); + + if (should_replace) { + tokens[best_first_token].text = *best_match + best_suffix; + for (size_t t = best_first_token + 1; t <= best_last_token; ++t) { + consumed[t] = true; + } + for (size_t t = best_first_token + 1; t <= best_last_token; ++t) { + if (t > 0) consumed[t - 1] = consumed[t - 1] || !tokens[t - 1].is_word; + } + wi += best_window; + } else { + ++wi; + } + } + + std::string result; + result.reserve(text.size()); + for (size_t i = 0; i < tokens.size(); ++i) { + if (!consumed[i]) { + result += tokens[i].text; + } + } + + text = std::move(result); +} + +inline InferenceOptions parse_inference_options_json(const std::string& json) { + InferenceOptions options; + + if (json.empty()) return options; + + size_t pos = json.find("\"temperature\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + options.temperature = std::stof(json.substr(pos)); + } + + pos = json.find("\"top_p\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + options.top_p = std::stof(json.substr(pos)); + } + + float parsed_min_p = options.min_p; + if (try_parse_json_float(json, "min_p", parsed_min_p)) { + options.min_p = std::clamp(parsed_min_p, 0.0f, 1.0f); + } + + float parsed_rep_penalty = options.repetition_penalty; + if (try_parse_json_float(json, "repetition_penalty", parsed_rep_penalty)) { + if (std::isfinite(parsed_rep_penalty) && parsed_rep_penalty > 0.0f) { + options.repetition_penalty = parsed_rep_penalty; + } + } + + size_t parsed_seed = 0; + if (try_parse_json_uint(json, "seed", parsed_seed)) { + options.sample_seed = parsed_seed; + } + + size_t parsed_top_k = options.top_k; + if (try_parse_json_uint(json, "top_k", parsed_top_k)) { + options.top_k = parsed_top_k; + } + + size_t parsed_max_tokens = options.max_tokens; + if (try_parse_json_uint(json, "max_tokens", parsed_max_tokens)) { + options.max_tokens = parsed_max_tokens; + } + + pos = json.find("\"force_tools\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.force_tools = (json.substr(pos, 4) == "true"); + } + + size_t parsed_tool_rag_top_k = options.tool_rag_top_k; + if (try_parse_json_uint(json, "tool_rag_top_k", parsed_tool_rag_top_k)) { + options.tool_rag_top_k = parsed_tool_rag_top_k; + } + + pos = json.find("\"confidence_threshold\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + options.confidence_threshold = std::stof(json.substr(pos)); + } + + pos = json.find("\"include_stop_sequences\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.include_stop_sequences = (json.substr(pos, 4) == "true"); + } + + pos = json.find("\"use_vad\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.use_vad = (json.substr(pos, 4) == "true"); + } + + pos = json.find("\"timestamps\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.timestamps = (json.substr(pos, 4) == "true"); + } + + pos = json.find("\"telemetry_enabled\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.telemetry_enabled = (json.substr(pos, 4) == "true"); + } + + pos = json.find("\"auto_handoff\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.auto_handoff = (json.substr(pos, 4) == "true"); + } + + size_t parsed_cloud_timeout_ms = options.cloud_timeout_ms; + if (try_parse_json_uint(json, "cloud_timeout_ms", parsed_cloud_timeout_ms)) { + options.cloud_timeout_ms = parsed_cloud_timeout_ms; + } + + pos = json.find("\"handoff_with_images\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.handoff_with_images = (json.substr(pos, 4) == "true"); + } + + pos = json.find("\"enable_thinking_if_supported\""); + if (pos != std::string::npos) { + pos = json.find(':', pos) + 1; + while (pos < json.length() && std::isspace(static_cast(json[pos]))) pos++; + options.enable_thinking_if_supported = (json.substr(pos, 4) == "true"); + } + + pos = json.find("\"stop_sequences\""); + if (pos != std::string::npos) { + pos = json.find('[', pos); + if (pos != std::string::npos) { + size_t end_pos = json.find(']', pos); + size_t seq_pos = json.find('"', pos); + + while (seq_pos != std::string::npos && seq_pos < end_pos) { + size_t seq_start = seq_pos + 1; + size_t seq_end = json.find('"', seq_start); + if (seq_end != std::string::npos) { + options.stop_sequences.push_back(json.substr(seq_start, seq_end - seq_start)); + } + seq_pos = json.find('"', seq_end + 1); + } + } + } + + return options; +} + +inline void parse_function_calls_from_response(const std::string& response_text, + std::string& regular_response, + std::vector& function_calls, + const std::vector& tools = {}) { + regular_response = response_text; + function_calls.clear(); + + const std::string needle_marker = ""; + size_t needle_pos = regular_response.find(needle_marker); + if (needle_pos != std::string::npos) { + size_t scan = needle_pos + needle_marker.size(); + while (scan < regular_response.size() && + std::isspace(static_cast(regular_response[scan]))) { + ++scan; + } + if (scan < regular_response.size() && regular_response[scan] == '[') { + size_t array_start = scan; + size_t array_end = find_matching_delimiter(regular_response, array_start, '[', ']'); + if (array_end <= regular_response.size()) { + size_t pos = array_start + 1; + while (pos < array_end) { + size_t obj_start = regular_response.find('{', pos); + if (obj_start == std::string::npos || obj_start >= array_end) break; + size_t obj_end = find_matching_delimiter(regular_response, obj_start, '{', '}'); + if (obj_end > regular_response.size() || obj_end > array_end) break; + function_calls.push_back(regular_response.substr(obj_start, obj_end - obj_start)); + pos = obj_end; + } + regular_response = trim_string( + regular_response.substr(0, needle_pos) + regular_response.substr(array_end)); + sanitize_function_calls(function_calls); + return; + } + } + } + + gemma::parse_function_calls(regular_response, function_calls, tools); + chat_tools::extract_qwen_tool_calls(regular_response, function_calls); + chat_tools::extract_lfm2_tool_calls(regular_response, function_calls); + + const char* FUNCTION_CALL_MARKER = "\"function_call\""; + size_t search_pos = 0; + const size_t text_len = regular_response.length(); + + while (search_pos < text_len) { + size_t marker_pos = regular_response.find(FUNCTION_CALL_MARKER, search_pos); + if (marker_pos == std::string::npos) break; + + size_t json_start = regular_response.find('{', marker_pos); + if (json_start == std::string::npos) break; + + int brace_count = 1; + size_t json_end = json_start + 1; + while (json_end < text_len && brace_count > 0) { + char c = regular_response[json_end]; + brace_count += (c == '{') - (c == '}'); + json_end++; + } + + if (brace_count == 0) { + function_calls.push_back(regular_response.substr(json_start, json_end - json_start)); + regular_response = regular_response.substr(0, marker_pos); + size_t last_bracket = regular_response.rfind('{'); + if(last_bracket != std::string::npos) { + regular_response = regular_response.substr(0, last_bracket); + } + } + search_pos = json_end; + } + + if (function_calls.empty()) { + size_t pos = 0; + while (pos < regular_response.size()) { + char open = regular_response[pos]; + if (open != '[' && open != '{') { ++pos; continue; } + size_t end = find_matching_delimiter(regular_response, pos, open, open == '[' ? ']' : '}'); + if (end > regular_response.size()) { ++pos; continue; } + std::string candidate = regular_response.substr(pos, end - pos); + bool call_shape = candidate.find("\"name\"") != std::string::npos && + (open == '[' || candidate.find("\"arguments\"") != std::string::npos); + if (call_shape && is_valid_json(candidate)) { + if (open == '[') { + for (const auto& item : split_json_array(candidate)) function_calls.push_back(item); + } else { + function_calls.push_back(candidate); + } + regular_response = trim_string( + regular_response.substr(0, pos) + regular_response.substr(end)); + pos = 0; + continue; + } + ++pos; + } + } + + if (function_calls.empty()) { + std::string remainder = trim_string(regular_response); + if (!remainder.empty() && (remainder.front() == '{' || remainder.front() == '[')) { + size_t span_end = find_matching_delimiter(remainder, 0, remainder.front(), + remainder.front() == '[' ? ']' : '}'); + bool whole_span = span_end == remainder.size(); + bool call_shape = remainder.find("\"name\"") != std::string::npos && + (remainder.front() == '[' || remainder.find("\"arguments\"") != std::string::npos); + if (whole_span && call_shape) regular_response.clear(); + } + } + + sanitize_function_calls(function_calls); +} + +inline void strip_tag_blocks(std::string& text, std::string& extracted, + const std::string& open_tag, const std::string& close_tag) { + std::string result; + size_t pos = 0; + + size_t first_close = text.find(close_tag); + size_t first_open = text.find(open_tag); + if (first_close != std::string::npos && + (first_open == std::string::npos || first_close < first_open)) { + extracted += text.substr(0, first_close); + pos = first_close + close_tag.size(); + } + + while (pos < text.size()) { + size_t open_pos = text.find(open_tag, pos); + if (open_pos == std::string::npos) { + result += text.substr(pos); + break; + } + result += text.substr(pos, open_pos - pos); + size_t content_start = open_pos + open_tag.size(); + size_t close_pos = text.find(close_tag, content_start); + if (close_pos == std::string::npos) { + if (!extracted.empty()) extracted += "\n"; + extracted += text.substr(content_start); + break; + } + if (!extracted.empty()) extracted += "\n"; + extracted += text.substr(content_start, close_pos - content_start); + pos = close_pos + close_tag.size(); + } + text = result; +} + +inline void partition_thinking_response(const std::string& input, std::string& thinking, std::string& content) { + thinking.clear(); + content = input; + + auto trim = [](std::string& s) { + size_t first = s.find_first_not_of(" \t\n\r"); + size_t last = s.find_last_not_of(" \t\n\r"); + if (first != std::string::npos && last != std::string::npos) + s = s.substr(first, last - first + 1); + else + s.clear(); + }; + + if (content.find("<|channel>") != std::string::npos || content.find("") != std::string::npos) { + strip_tag_blocks(content, thinking, "<|channel>", ""); + } else if (content.find("") != std::string::npos || content.find("") != std::string::npos) { + strip_tag_blocks(content, thinking, "", ""); + } else { + return; + } + + trim(thinking); + trim(content); +} + +struct TranscriptSegment { + float start = 0.0f; + float end = 0.0f; + std::string text; +}; + +inline std::string construct_response_json(const std::string& regular_response, + const std::vector& function_calls, + double time_to_first_token, + double total_time_ms, + double prefill_tps, + double decode_tps, + size_t prompt_tokens, + size_t completion_tokens, + float confidence = 0.0f, + bool cloud_handoff = false, + const std::string& thinking = "", + const std::vector& segments = {}, + const std::string& context_response = "", + float confidence_threshold = -1.0f, + const std::string& cloud_handoff_reason = "") { + std::ostringstream json; + json << "{"; + json << "\"success\":true,"; + json << "\"error\":null,"; + json << "\"cloud_handoff\":" << (cloud_handoff ? "true" : "false") << ","; + json << "\"cloud_handoff_reason\":\"" << escape_json_string(cloud_handoff_reason) << "\","; + json << "\"response\":\"" << escape_json_string(regular_response) << "\","; + if (!context_response.empty()) { + json << "\"context_response\":\"" << escape_json_string(context_response) << "\","; + } + if (!thinking.empty()) { + json << "\"thinking\":\"" << escape_json_string(thinking) << "\","; + } + json << "\"function_calls\":["; + for (size_t i = 0; i < function_calls.size(); ++i) { + if (i > 0) json << ","; + json << function_calls[i]; + } + json << "],"; + json << "\"segments\":["; + for (size_t i = 0; i < segments.size(); ++i) { + if (i > 0) json << ","; + json << "{\"start\":" << std::fixed << std::setprecision(3) << segments[i].start + << ",\"end\":" << std::fixed << std::setprecision(3) << segments[i].end + << ",\"text\":\"" << escape_json_string(segments[i].text) << "\"}"; + } + json << "],"; + json << "\"confidence\":" << std::fixed << std::setprecision(4) << confidence << ","; + json << "\"confidence_threshold\":" << std::fixed << std::setprecision(4) << confidence_threshold << ","; + json << "\"time_to_first_token_ms\":" << std::fixed << std::setprecision(2) << time_to_first_token << ","; + json << "\"total_time_ms\":" << std::fixed << std::setprecision(2) << total_time_ms << ","; + json << "\"prefill_tps\":" << std::fixed << std::setprecision(2) << prefill_tps << ","; + json << "\"decode_tps\":" << std::fixed << std::setprecision(2) << decode_tps << ","; + json << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() << ","; + json << "\"prefill_tokens\":" << prompt_tokens << ","; + json << "\"decode_tokens\":" << completion_tokens << ","; + json << "\"total_tokens\":" << (prompt_tokens + completion_tokens); + json << "}"; + return json.str(); +} + +inline std::string serialize_function_calls(const std::vector& calls) { + if (calls.empty()) return "[]"; + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < calls.size(); ++i) { + if (i > 0) oss << ","; + oss << calls[i]; + } + oss << "]"; + return oss.str(); +} + +inline int validate_audio_params( + const char* component, + void* model, + char* response_buffer, size_t buffer_size, + const char* audio_file_path, + const uint8_t* pcm_buffer, size_t pcm_buffer_size) { + if (!model) { + std::string err = last_error_message.empty() ? "Model not initialized." : last_error_message; + CACTUS_LOG_ERROR(component, err); + handle_error_response(err, response_buffer, buffer_size); + return -1; + } + if (!response_buffer || buffer_size == 0) { + CACTUS_LOG_ERROR(component, "Invalid parameters: response_buffer or buffer_size"); + handle_error_response("Invalid parameters", response_buffer, buffer_size); + return -1; + } + if (!audio_file_path && (!pcm_buffer || pcm_buffer_size == 0)) { + CACTUS_LOG_ERROR(component, "No audio input provided"); + handle_error_response("Either audio_file_path or pcm_buffer must be provided", response_buffer, buffer_size); + return -1; + } + if (audio_file_path && pcm_buffer && pcm_buffer_size > 0) { + CACTUS_LOG_ERROR(component, "Both audio_file_path and pcm_buffer provided"); + handle_error_response("Cannot provide both audio_file_path and pcm_buffer", response_buffer, buffer_size); + return -1; + } + if (pcm_buffer && pcm_buffer_size > 0 && (pcm_buffer_size < 2 || pcm_buffer_size % 2 != 0)) { + CACTUS_LOG_ERROR(component, "Invalid pcm_buffer_size"); + handle_error_response("pcm_buffer_size must be even and at least 2 bytes", response_buffer, buffer_size); + return -1; + } + return 0; +} + +inline std::vector pcm_to_float(const uint8_t* pcm_buffer, size_t pcm_buffer_size) { + const int16_t* samples = reinterpret_cast(pcm_buffer); + size_t n = pcm_buffer_size / 2; + std::vector out(n); + for (size_t i = 0; i < n; ++i) + out[i] = static_cast(samples[i]) / 32768.0f; + return out; +} + +} // namespace ffi +} // namespace cactus + +#ifdef __cplusplus +extern "C" { +#endif + +const char* cactus_get_last_error(); + +#ifdef __cplusplus +} +#endif + +#endif // CACTUS_UTILS_H diff --git a/cactus-engine/test.sh b/cactus-engine/test.sh new file mode 100755 index 000000000..aa793087a --- /dev/null +++ b/cactus-engine/test.sh @@ -0,0 +1,78 @@ +#!/bin/bash +set -e +cd "$(dirname "$0")" + +PROJECT_ROOT="$(pwd)/.." +ASSETS_DIR="$(pwd)/tests/assets" + +IOS_MODE=false +ANDROID_MODE=false +SUITE="" +BACKEND="" + +while [[ $# -gt 0 ]]; do + case $1 in + --ios) IOS_MODE=true; shift ;; + --android) ANDROID_MODE=true; shift ;; + --suite) SUITE="${2:?--suite needs an argument}"; shift 2 ;; + --model) CACTUS_TEST_MODEL="${2:?--model needs an argument}"; shift 2 ;; + --transcription-model) CACTUS_TEST_TRANSCRIPTION_MODEL="${2:?--transcription-model needs an argument}"; shift 2 ;; + --backend) BACKEND="${2:?--backend needs an argument}"; shift 2 ;; + *) echo "Unknown arg: $1" >&2; exit 2 ;; + esac +done + +require_bundle() { + local dir="$1" label="$2" + if [ -z "$dir" ] || [ ! -f "$dir/components/manifest.json" ]; then + echo "Error: --$label must point to a prepared bundle. Run engine tests via 'cactus test', which prepares them." >&2 + exit 2 + fi + echo "$dir" +} + +BUNDLE_DIR="$(require_bundle "$CACTUS_TEST_MODEL" "model")" +TRANSCRIPTION_BUNDLE_DIR="$(require_bundle "$CACTUS_TEST_TRANSCRIPTION_MODEL" "transcription-model")" + +if [ "$IOS_MODE" = true ]; then + export CACTUS_TEST_MODEL="$BUNDLE_DIR" CACTUS_TEST_TRANSCRIPTION_MODEL="$TRANSCRIPTION_BUNDLE_DIR" CACTUS_TEST_SUITE="$SUITE" CACTUS_TEST_BACKEND="$BACKEND" + exec "$(pwd)/tests/ios/run.sh" +fi +if [ "$ANDROID_MODE" = true ]; then + export CACTUS_TEST_MODEL="$BUNDLE_DIR" CACTUS_TEST_TRANSCRIPTION_MODEL="$TRANSCRIPTION_BUNDLE_DIR" CACTUS_TEST_SUITE="$SUITE" CACTUS_TEST_BACKEND="$BACKEND" + exec "$(pwd)/tests/android/run.sh" +fi + +echo "Model: $CACTUS_TEST_MODEL" +echo "Bundle: $BUNDLE_DIR" +echo "Transcription model: $CACTUS_TEST_TRANSCRIPTION_MODEL" +echo "Transcription bundle: $TRANSCRIPTION_BUNDLE_DIR" +echo "Backend: $BACKEND" + +cd "$PROJECT_ROOT/cactus-engine/tests" +rm -rf build && mkdir build && cd build +cmake .. -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null +make -j"$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)" + +export CACTUS_TEST_MODEL="$BUNDLE_DIR" +export CACTUS_TEST_TRANSCRIPTION_MODEL="$TRANSCRIPTION_BUNDLE_DIR" +export CACTUS_TEST_ASSETS="$ASSETS_DIR" +export CACTUS_INDEX_PATH="$ASSETS_DIR" +export CACTUS_TEST_BACKEND="$BACKEND" + +FAILED=0 +if [ -n "$SUITE" ]; then + target="./test_$SUITE" + if [ -x "$target" ]; then + "$target" || FAILED=1 + else + echo "Test not found: $target" >&2 + FAILED=1 + fi +else + for t in ./test_*; do + [ -x "$t" ] || continue + "$t" || FAILED=1 + done +fi +exit $FAILED diff --git a/cactus-engine/tests/CMakeLists.txt b/cactus-engine/tests/CMakeLists.txt new file mode 100644 index 000000000..e3dd2800d --- /dev/null +++ b/cactus-engine/tests/CMakeLists.txt @@ -0,0 +1,57 @@ +cmake_minimum_required(VERSION 3.10) +project(CactusEngineTests LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(APPLE) + set(CMAKE_OSX_ARCHITECTURES "arm64") +endif() + +if(NOT TARGET cactus_engine) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/.. ${CMAKE_CURRENT_BINARY_DIR}/cactus-engine) +endif() + +file(GLOB TEST_SOURCES "test_*.cpp") +list(REMOVE_ITEM TEST_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/test_utils.cpp") + +foreach(TEST_FILE ${TEST_SOURCES}) + get_filename_component(TEST_NAME ${TEST_FILE} NAME_WE) + add_executable(${TEST_NAME} ${TEST_FILE} test_utils.cpp) + target_link_libraries(${TEST_NAME} PRIVATE cactus_engine) +endforeach() + +if(TARGET test_curl AND EXISTS "${CACTUS_CURL_ROOT}/include/curl/curl.h") + target_include_directories(test_curl PRIVATE "${CACTUS_CURL_ROOT}/include") +endif() + +find_package(SDL2 QUIET CONFIG) +if(NOT SDL2_FOUND) + find_library(SDL2_LIBRARY NAMES SDL2 PATHS /opt/homebrew/lib /usr/local/lib) + find_path(SDL2_INCLUDE_DIR SDL.h PATHS /opt/homebrew/include/SDL2 /usr/local/include/SDL2) + if(SDL2_LIBRARY AND SDL2_INCLUDE_DIR) + set(SDL2_FOUND TRUE) + endif() +endif() +if(SDL2_FOUND) + add_executable(run run.cpp test_utils.cpp) + target_link_libraries(run PRIVATE cactus_engine) + if(TARGET SDL2::SDL2) + target_link_libraries(run PRIVATE SDL2::SDL2) + else() + target_link_libraries(run PRIVATE ${SDL2_LIBRARY}) + target_include_directories(run PRIVATE ${SDL2_INCLUDE_DIR}) + endif() +endif() + +add_executable(transcribe transcribe.cpp) +target_link_libraries(transcribe PRIVATE cactus_engine) +if(SDL2_FOUND) + target_compile_definitions(transcribe PRIVATE HAVE_SDL2) + if(TARGET SDL2::SDL2) + target_link_libraries(transcribe PRIVATE SDL2::SDL2) + else() + target_link_libraries(transcribe PRIVATE ${SDL2_LIBRARY}) + target_include_directories(transcribe PRIVATE ${SDL2_INCLUDE_DIR}) + endif() +endif() diff --git a/cactus-engine/tests/android/CMakeLists.txt b/cactus-engine/tests/android/CMakeLists.txt new file mode 100644 index 000000000..9ff1fc7c2 --- /dev/null +++ b/cactus-engine/tests/android/CMakeLists.txt @@ -0,0 +1,23 @@ +cmake_minimum_required(VERSION 3.10) +project(CactusAndroidTests LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT TARGET cactus_engine) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../.. ${CMAKE_CURRENT_BINARY_DIR}/cactus-engine) +endif() + +set(TESTS_DIR ${CMAKE_CURRENT_SOURCE_DIR}/..) +file(GLOB TEST_SOURCES "${TESTS_DIR}/test_*.cpp") +list(REMOVE_ITEM TEST_SOURCES "${TESTS_DIR}/test_utils.cpp") + +foreach(TEST_FILE ${TEST_SOURCES}) + get_filename_component(TEST_NAME ${TEST_FILE} NAME_WE) + add_executable(${TEST_NAME} ${TEST_FILE} ${TESTS_DIR}/test_utils.cpp) + target_link_libraries(${TEST_NAME} PRIVATE cactus_engine) +endforeach() + +if(TARGET test_curl AND EXISTS "${CACTUS_CURL_ROOT}/include/curl/curl.h") + target_include_directories(test_curl PRIVATE "${CACTUS_CURL_ROOT}/include") +endif() diff --git a/tests/android/run.sh b/cactus-engine/tests/android/run.sh similarity index 71% rename from tests/android/run.sh rename to cactus-engine/tests/android/run.sh index 3396561ed..1f5eddee5 100755 --- a/tests/android/run.sh +++ b/cactus-engine/tests/android/run.sh @@ -1,15 +1,21 @@ #!/bin/bash SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" - -MODEL_NAME="$1" -TRANSCRIBE_MODEL_NAME="$2" +PROJECT_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" +CACTUS_CURL_ROOT="${CACTUS_CURL_ROOT:-$PROJECT_ROOT/cactus-engine/libs/curl}" +export CACTUS_CURL_ROOT + +for var in CACTUS_TEST_MODEL CACTUS_TEST_TRANSCRIPTION_MODEL; do + if [ -z "${!var:-}" ] || [ ! -d "${!var}" ]; then + echo "Error: $var must point to a prepared bundle (got: '${!var:-}')" >&2 + exit 1 + fi +done echo "Running Cactus tests on Android..." echo "============================" -if ! command -v adb; then +if ! command -v adb &>/dev/null; then echo "adb not found" echo "Install Android SDK Platform Tools and ensure it's in your PATH" echo "Installation:" @@ -24,7 +30,7 @@ adb start-server connected_devices=$(adb devices | grep -E "device$|emulator" | grep -v "^List" | awk '{print $1}') -if command -v emulator; then +if command -v emulator &>/dev/null; then available_emulators=$(emulator -list-avds | grep -v "^INFO" || true) else available_emulators="" @@ -132,10 +138,9 @@ if [ "$device_type" = "emulator" ]; then echo "Selected: $device_name (Emulator - Not Running)" echo "Starting emulator..." - if ! command -v emulator; then + if ! command -v emulator &>/dev/null; then echo "Emulator command not found" echo "Add Android SDK emulator to your PATH" - echo "Location: \$ANDROID_HOME/emulator or \$ANDROID_SDK_ROOT/emulator" exit 1 fi @@ -171,31 +176,17 @@ else echo "Selected: $device_name (Device)" fi -if ! adb -s "$DEVICE_ID" shell echo "test"; then +if ! adb -s "$DEVICE_ID" shell echo "test" &>/dev/null; then echo "Device not responding" echo "Ensure the device is connected and USB debugging is authorized" exit 1 fi -echo "" -echo "Step 2: Building Cactus library for Android..." - -if ! "$PROJECT_ROOT/android/build.sh"; then - echo "Failed to build Cactus library" - exit 1 -fi - -echo "" -echo "Step 3: Building Android tests..." - -android_test_dir="$SCRIPT_DIR" -android_build_dir="$android_test_dir/build" - if [ -z "$ANDROID_NDK_HOME" ]; then if [ -n "$ANDROID_HOME" ]; then - ANDROID_NDK_HOME=$(ls -d "$ANDROID_HOME/ndk/"* | sort -V | tail -1) + ANDROID_NDK_HOME=$(ls -d "$ANDROID_HOME/ndk/"* 2>/dev/null | sort -V | tail -1) elif [ -d "$HOME/Library/Android/sdk" ]; then - ANDROID_NDK_HOME=$(ls -d "$HOME/Library/Android/sdk/ndk/"* | sort -V | tail -1) + ANDROID_NDK_HOME=$(ls -d "$HOME/Library/Android/sdk/ndk/"* 2>/dev/null | sort -V | tail -1) fi fi @@ -205,86 +196,141 @@ if [ -z "$ANDROID_NDK_HOME" ] || [ ! -d "$ANDROID_NDK_HOME" ]; then exit 1 fi -cmake_toolchain_file="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" +cmake_toolchain="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" android_platform=${ANDROID_PLATFORM:-android-21} -android_abi="arm64-v8a" +n_jobs=$(nproc 2>/dev/null || sysctl -n hw.logicalcpu 2>/dev/null || echo 4) + +echo "" +echo "Step 2: Building Cactus library for Android..." + +if ! "$PROJECT_ROOT/android/build.sh"; then + echo "Failed to build Cactus library" + exit 1 +fi + +echo "" +echo "Step 3: Building Android tests..." + +android_test_dir="$SCRIPT_DIR" +android_build_dir="$android_test_dir/build" rm -rf "$android_build_dir" mkdir -p "$android_build_dir" if ! cmake -S "$android_test_dir" -B "$android_build_dir" \ - -DCMAKE_TOOLCHAIN_FILE="$cmake_toolchain_file" \ - -DANDROID_ABI="$android_abi" \ + -DCMAKE_TOOLCHAIN_FILE="$cmake_toolchain" \ + -DANDROID_ABI="arm64-v8a" \ -DANDROID_PLATFORM="$android_platform" \ -DCMAKE_BUILD_TYPE=Release \ + -DCACTUS_CURL_ROOT="$CACTUS_CURL_ROOT" \ -DCMAKE_RULE_MESSAGES=OFF \ -DCMAKE_VERBOSE_MAKEFILE=OFF; then echo "Failed to configure tests" exit 1 fi -n_jobs=$(nproc || sysctl -n hw.logicalcpu || echo 4) if ! cmake --build "$android_build_dir" -j "$n_jobs"; then echo "Failed to build tests" exit 1 fi echo "Discovering test executables..." -test_executables=($(find "$android_build_dir" -maxdepth 1 -name "test_*" -type f | sort)) +all_test_executables=($(find "$android_build_dir" -maxdepth 1 -name "test_*" -type f | sort)) -if [ ${#test_executables[@]} -eq 0 ]; then +if [ ${#all_test_executables[@]} -eq 0 ]; then echo "No test executables found" exit 1 fi +test_executables=() +for test_exe in "${all_test_executables[@]}"; do + test_name=$(basename "$test_exe" | sed 's/^test_//') + if [ -z "${CACTUS_TEST_SUITE:-}" ] || [ "$test_name" = "$CACTUS_TEST_SUITE" ]; then + test_executables+=("$test_exe") + fi +done + +if [ ${#test_executables[@]} -eq 0 ]; then + echo "No test executables match filter: $CACTUS_TEST_SUITE" + exit 1 +fi + echo "Found ${#test_executables[@]} test executable(s)" echo "" echo "Step 4: Deploying to device..." -model_dir=$(echo "$MODEL_NAME" | sed 's|.*/||' | tr '[:upper:]' '[:lower:]') -transcribe_model_dir=$(echo "$TRANSCRIBE_MODEL_NAME" | sed 's|.*/||' | tr '[:upper:]' '[:lower:]') -model_src="$PROJECT_ROOT/weights/$model_dir" -transcribe_model_src="$PROJECT_ROOT/weights/$transcribe_model_dir" -assets_src="$PROJECT_ROOT/tests/assets" +model_dir=$(basename "$CACTUS_TEST_MODEL") +transcription_dir=$(basename "$CACTUS_TEST_TRANSCRIPTION_MODEL") +assets_src="$PROJECT_ROOT/cactus-engine/tests/assets" device_test_dir="/data/local/tmp/cactus_tests" device_model_dir="/data/local/tmp/cactus_models" device_assets_dir="/data/local/tmp/cactus_assets" +adb -s "$DEVICE_ID" shell "rm -rf $device_test_dir $device_model_dir $device_assets_dir" adb -s "$DEVICE_ID" shell "mkdir -p $device_test_dir $device_model_dir $device_assets_dir" echo "Pushing model weights..." -adb -s "$DEVICE_ID" push "$model_src" "$device_model_dir/" -adb -s "$DEVICE_ID" push "$transcribe_model_src" "$device_model_dir/" +if ! adb -s "$DEVICE_ID" push "$CACTUS_TEST_MODEL" "$device_model_dir/" >/dev/null; then + echo "Failed to push model weights" + exit 1 +fi -echo "Pushing test assets..." -adb -s "$DEVICE_ID" push "$assets_src" "$device_assets_dir/" +echo "Pushing transcription model..." +if ! adb -s "$DEVICE_ID" push "$CACTUS_TEST_TRANSCRIPTION_MODEL" "$device_model_dir/" >/dev/null; then + echo "Failed to push transcription model" + exit 1 +fi + +if [ -d "$assets_src" ]; then + echo "Pushing test assets..." + if ! adb -s "$DEVICE_ID" push "$assets_src" "$device_assets_dir/" >/dev/null; then + echo "Failed to push test assets" + exit 1 + fi +fi echo "Pushing test executables..." for test_exe in "${test_executables[@]}"; do test_name=$(basename "$test_exe") - adb -s "$DEVICE_ID" push "$test_exe" "$device_test_dir/" + if ! adb -s "$DEVICE_ID" push "$test_exe" "$device_test_dir/" >/dev/null; then + echo "Failed to push $test_name" + exit 1 + fi adb -s "$DEVICE_ID" shell "chmod +x $device_test_dir/$test_name" done echo "" echo "Step 5: Running tests..." echo "------------------------" -echo "Using model path: $device_model_dir/$model_dir" -echo "Using transcribe model path: $device_model_dir/$transcribe_model_dir" -echo "Using assets path: $device_assets_dir/assets" -echo "Using index path: $device_assets_dir/assets" +echo "Using model path: $device_model_dir/$model_dir" +echo "Using transcription model path: $device_model_dir/$transcription_dir" +echo "Using assets path: $device_assets_dir/assets" +FAILED=0 for test_exe in "${test_executables[@]}"; do test_name=$(basename "$test_exe") + echo "" + echo "Running $test_name..." - adb -s "$DEVICE_ID" shell "cd $device_test_dir && \ + if ! adb -s "$DEVICE_ID" shell "cd $device_test_dir && \ export CACTUS_TEST_MODEL=$device_model_dir/$model_dir && \ - export CACTUS_TEST_TRANSCRIBE_MODEL=$device_model_dir/$transcribe_model_dir && \ + export CACTUS_TEST_TRANSCRIPTION_MODEL=$device_model_dir/$transcription_dir && \ export CACTUS_TEST_ASSETS=$device_assets_dir/assets && \ export CACTUS_INDEX_PATH=$device_assets_dir/assets && \ - ./$test_name" + export CACTUS_NO_CLOUD_TELE=${CACTUS_NO_CLOUD_TELE:-1} && \ + export CACTUS_TEST_BACKEND=${CACTUS_TEST_BACKEND:-auto} && \ + ./$test_name"; then + FAILED=1 + fi done echo "" +if [ $FAILED -eq 0 ]; then + echo "All tests passed." +else + echo "Some tests failed." +fi + +exit $FAILED diff --git a/cactus-engine/tests/assets/hotword.wav b/cactus-engine/tests/assets/hotword.wav new file mode 100644 index 000000000..ad2dd47cd Binary files /dev/null and b/cactus-engine/tests/assets/hotword.wav differ diff --git a/tests/assets/rag_corpus/rag_sample1.txt b/cactus-engine/tests/assets/rag_corpus/rag_sample1.txt similarity index 100% rename from tests/assets/rag_corpus/rag_sample1.txt rename to cactus-engine/tests/assets/rag_corpus/rag_sample1.txt diff --git a/cactus-engine/tests/assets/record.wav b/cactus-engine/tests/assets/record.wav new file mode 100644 index 000000000..a5f238194 Binary files /dev/null and b/cactus-engine/tests/assets/record.wav differ diff --git a/cactus-engine/tests/assets/test.wav b/cactus-engine/tests/assets/test.wav new file mode 100644 index 000000000..99638e372 Binary files /dev/null and b/cactus-engine/tests/assets/test.wav differ diff --git a/cactus-engine/tests/assets/test_long.wav b/cactus-engine/tests/assets/test_long.wav new file mode 100644 index 000000000..788a9ace9 Binary files /dev/null and b/cactus-engine/tests/assets/test_long.wav differ diff --git a/tests/assets/test_monkey.png b/cactus-engine/tests/assets/test_monkey.png similarity index 100% rename from tests/assets/test_monkey.png rename to cactus-engine/tests/assets/test_monkey.png diff --git a/cactus-engine/tests/assets/test_thing.png b/cactus-engine/tests/assets/test_thing.png new file mode 100644 index 000000000..395edada6 Binary files /dev/null and b/cactus-engine/tests/assets/test_thing.png differ diff --git a/cactus-engine/tests/ios/.gitignore b/cactus-engine/tests/ios/.gitignore new file mode 100644 index 000000000..143f8b56e --- /dev/null +++ b/cactus-engine/tests/ios/.gitignore @@ -0,0 +1,4 @@ +CactusTest/CactusTest.xcodeproj/* +!CactusTest/CactusTest.xcodeproj/project.pbxproj.template +CactusTest/CactusTest/AppDelegate.mm +build/ diff --git a/tests/ios/CactusTest/CactusTest.xcodeproj/project.pbxproj.template b/cactus-engine/tests/ios/CactusTest/CactusTest.xcodeproj/project.pbxproj.template similarity index 100% rename from tests/ios/CactusTest/CactusTest.xcodeproj/project.pbxproj.template rename to cactus-engine/tests/ios/CactusTest/CactusTest.xcodeproj/project.pbxproj.template diff --git a/tests/ios/CactusTest/CactusTest/AppDelegate.h b/cactus-engine/tests/ios/CactusTest/CactusTest/AppDelegate.h similarity index 100% rename from tests/ios/CactusTest/CactusTest/AppDelegate.h rename to cactus-engine/tests/ios/CactusTest/CactusTest/AppDelegate.h diff --git a/cactus-engine/tests/ios/CactusTest/CactusTest/CactusTest.entitlements b/cactus-engine/tests/ios/CactusTest/CactusTest/CactusTest.entitlements new file mode 100644 index 000000000..475132a2d --- /dev/null +++ b/cactus-engine/tests/ios/CactusTest/CactusTest/CactusTest.entitlements @@ -0,0 +1,8 @@ + + + + + com.apple.developer.kernel.increased-memory-limit + + + diff --git a/tests/ios/CactusTest/CactusTest/Info.plist b/cactus-engine/tests/ios/CactusTest/CactusTest/Info.plist similarity index 100% rename from tests/ios/CactusTest/CactusTest/Info.plist rename to cactus-engine/tests/ios/CactusTest/CactusTest/Info.plist diff --git a/tests/ios/CactusTest/CactusTest/main.m b/cactus-engine/tests/ios/CactusTest/CactusTest/main.m similarity index 100% rename from tests/ios/CactusTest/CactusTest/main.m rename to cactus-engine/tests/ios/CactusTest/CactusTest/main.m diff --git a/tests/ios/configure_xcode.rb b/cactus-engine/tests/ios/configure_xcode.rb similarity index 50% rename from tests/ios/configure_xcode.rb rename to cactus-engine/tests/ios/configure_xcode.rb index 921774b35..6e8379500 100755 --- a/tests/ios/configure_xcode.rb +++ b/cactus-engine/tests/ios/configure_xcode.rb @@ -7,59 +7,96 @@ def fail_with(message) end def generate_app_delegate(output_path, test_files) + raw_level = ENV.fetch('CACTUS_TEST_LOG_LEVEL', 'WARN').to_s.upcase + level_int = case raw_level + when 'DEBUG' then 0 + when 'INFO' then 1 + when 'WARN' then 2 + when 'ERROR' then 3 + when 'NONE' then 4 + else 2 + end + test_names = test_files.map { |f| File.basename(f, '.cpp') } extern_declarations = test_names.map { |name| "extern int #{name}_main();" }.join("\n") - test_calls = test_names.map { |name| " #{name}_main();" }.join("\n") + test_calls = test_names.map { |name| + filter_name = name.sub(/^test_/, '') + " if (should_run(\"#{filter_name}\")) failed |= (#{name}_main() != 0);" + }.join("\n") app_delegate_content = <<~OBJC - // AUTO-GENERATED by configure_xcode.rb - DO NOT EDIT MANUALLY - // This file is regenerated on each test run with discovered test files - #import "AppDelegate.h" - #import +#import "AppDelegate.h" +#import +#import +#include "cactus_engine.h" +#include - #{extern_declarations} +#{extern_declarations} - @implementation AppDelegate +@implementation AppDelegate - - (void)copyFromBundle:(NSString *)bundlePath toDocuments:(const char *)name { - if (!name) return; - NSFileManager *fileManager = [NSFileManager defaultManager]; - NSString *itemName = [NSString stringWithUTF8String:name]; - NSString *sourceItemPath = [NSString stringWithFormat:@"%@/%@", bundlePath, itemName]; - if ([fileManager fileExistsAtPath:itemName]) { - [fileManager removeItemAtPath:itemName error:nil]; - } - [fileManager copyItemAtPath:sourceItemPath toPath:itemName error:nil]; +- (void)copyFromBundle:(NSString *)bundlePath toDocuments:(const char *)name { + if (!name) return; + NSFileManager *fileManager = [NSFileManager defaultManager]; + NSString *itemName = [NSString stringWithUTF8String:name]; + NSString *sourceItemPath = [NSString stringWithFormat:@"%@/%@", bundlePath, itemName]; + if (![fileManager fileExistsAtPath:sourceItemPath]) { + fprintf(stderr, "[CactusTest] copyFromBundle: source not found: %s\\n", [sourceItemPath UTF8String]); + return; } - - - (BOOL)application:(UIApplication *)application didFinishLaunchingWithOptions:(NSDictionary *)launchOptions { - NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES); - NSString *documentsDirectory = paths[0]; - chdir([documentsDirectory UTF8String]); - - #if !TARGET_OS_SIMULATOR - freopen("cactus_test.log", "w", stdout); - freopen("cactus_test.log", "a", stderr); - setbuf(stdout, NULL); - setbuf(stderr, NULL); - #endif - - NSString *bundlePath = [[NSBundle mainBundle] bundlePath]; - [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_MODEL")]; - [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_TRANSCRIBE_MODEL")]; - [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_ASSETS")]; - - dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(NSEC_PER_SEC)), dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ - #{test_calls} - exit(0); - }); - - return YES; + if ([fileManager fileExistsAtPath:itemName]) { + NSError *removeError = nil; + [fileManager removeItemAtPath:itemName error:&removeError]; + } + NSError *copyError = nil; + [fileManager copyItemAtPath:sourceItemPath toPath:itemName error:©Error]; + if (copyError) { + fprintf(stderr, "[CactusTest] copyFromBundle: failed to copy %s -> %s: %s\\n", + [sourceItemPath UTF8String], [itemName UTF8String], [[copyError localizedDescription] UTF8String]); } +} + +- (BOOL)application:(UIApplication *)application didFinishLaunchingWithOptions:(NSDictionary *)launchOptions { + NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES); + NSString *documentsDirectory = paths[0]; + chdir([documentsDirectory UTF8String]); + +#if !TARGET_OS_SIMULATOR + freopen("cactus_test.log", "w", stdout); + dup2(fileno(stdout), fileno(stderr)); + setbuf(stdout, NULL); + setbuf(stderr, NULL); +#endif + + cactus_log_set_level(#{level_int}); + + NSString *bundlePath = [[NSBundle mainBundle] bundlePath]; + [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_MODEL")]; + [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_TRANSCRIPTION_MODEL")]; + [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_ASSETS")]; + + dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(NSEC_PER_SEC)), dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + const char* only = getenv("CACTUS_TEST_ONLY"); + std::string filter = (only && only[0]) ? std::string(only) : ""; + auto should_run = [&](const char* name) { + return filter.empty() || filter == name; + }; + + int failed = 0; + #{test_calls} - @end - OBJC + FILE* f = fopen("cactus_test.exitcode", "w"); + if (f) { fprintf(f, "%d\\n", failed); fclose(f); } + exit(failed); + }); + + return YES; +} + +@end + +OBJC File.write(output_path, app_delegate_content) puts "Generated AppDelegate.mm with #{test_names.length} test(s)" @@ -67,7 +104,6 @@ def generate_app_delegate(output_path, test_files) project_root = ENV['PROJECT_ROOT'] tests_root = ENV['TESTS_ROOT'] -cactus_root = ENV['CACTUS_ROOT'] project_path = ENV['XCODEPROJ_PATH'] bundle_id = ENV['BUNDLE_ID'] team_id = ENV['DEVELOPMENT_TEAM'] @@ -75,7 +111,6 @@ def generate_app_delegate(output_path, test_files) fail_with("PROJECT_ROOT not set") unless project_root fail_with("TESTS_ROOT not set") unless tests_root -fail_with("CACTUS_ROOT not set") unless cactus_root fail_with("XCODEPROJ_PATH not set") unless project_path fail_with("DEVICE_TYPE not set (should be 'device' or 'simulator')") unless device_type fail_with("Xcode project not found") unless File.exist?(project_path) @@ -132,12 +167,30 @@ def generate_app_delegate(output_path, test_files) apple_dir = File.join(project_root, 'apple') static_lib_path = device_type == 'simulator' ? - File.join(apple_dir, 'libcactus-simulator.a') : - File.join(apple_dir, 'libcactus-device.a') + File.join(apple_dir, 'libcactus_engine-simulator.a') : + File.join(apple_dir, 'libcactus_engine-device.a') fail_with("Static library not found at: #{static_lib_path}") unless File.exist?(static_lib_path) puts "Using static library: #{static_lib_path}" +cactus_engine_dir = File.join(project_root, 'cactus-engine') +cactus_graph_dir = File.join(project_root, 'cactus-graph') +cactus_kernels_dir = File.join(project_root, 'cactus-kernels') + +curl_root = ENV['CACTUS_CURL_ROOT'] +vendored_curl_lib = nil +if curl_root && !curl_root.empty? + vendored_curl_lib = device_type == 'simulator' ? + File.join(curl_root, 'ios', 'simulator', 'libcurl.a') : + File.join(curl_root, 'ios', 'device', 'libcurl.a') + if File.exist?(vendored_curl_lib) + puts "Using vendored iOS libcurl: #{vendored_curl_lib}" + else + vendored_curl_lib = nil + puts "Vendored iOS libcurl not found under CACTUS_CURL_ROOT=#{curl_root}; continuing without explicit curl link" + end +end + target.frameworks_build_phase.files.to_a.each do |build_file| if build_file.file_ref&.path&.to_s&.include?('libcactus') target.frameworks_build_phase.files.delete(build_file) @@ -153,12 +206,19 @@ def generate_app_delegate(output_path, test_files) target.build_configurations.each do |config| config.build_settings['HEADER_SEARCH_PATHS'] ||= ['$(inherited)'] - [tests_root, cactus_root].each do |path| + [tests_root, cactus_engine_dir, File.join(cactus_engine_dir, 'libs'), cactus_graph_dir, cactus_kernels_dir, File.join(cactus_kernels_dir, 'src')].each do |path| config.build_settings['HEADER_SEARCH_PATHS'] << path unless config.build_settings['HEADER_SEARCH_PATHS'].include?(path) end + if curl_root && !curl_root.empty? + curl_include = File.join(curl_root, 'include') + if File.exist?(File.join(curl_include, 'curl', 'curl.h')) + config.build_settings['HEADER_SEARCH_PATHS'] << curl_include unless config.build_settings['HEADER_SEARCH_PATHS'].include?(curl_include) + end + end config.build_settings['CLANG_CXX_LANGUAGE_STANDARD'] = 'c++20' config.build_settings['CLANG_CXX_LIBRARY'] = 'libc++' + config.build_settings['GCC_ENABLE_CPP_RTTI'] = 'NO' config.build_settings['IPHONEOS_DEPLOYMENT_TARGET'] = '13.0' config.build_settings['CODE_SIGN_STYLE'] = 'Automatic' config.build_settings['PRODUCT_BUNDLE_IDENTIFIER'] = bundle_id if bundle_id @@ -172,9 +232,12 @@ def generate_app_delegate(output_path, test_files) config.build_settings['OTHER_LDFLAGS'].reject! { |flag| flag.to_s.include?('libcactus') } config.build_settings['OTHER_LDFLAGS'] << static_lib_path - ['-framework CoreML', '-framework Foundation'].each do |framework| + ['-framework Foundation', '-framework Accelerate', '-framework Security', '-framework SystemConfiguration', '-framework CFNetwork', '-framework Metal', '-framework MetalPerformanceShaders'].each do |framework| config.build_settings['OTHER_LDFLAGS'] << framework unless config.build_settings['OTHER_LDFLAGS'].include?(framework) end + if vendored_curl_lib + config.build_settings['OTHER_LDFLAGS'] << vendored_curl_lib unless config.build_settings['OTHER_LDFLAGS'].include?(vendored_curl_lib) + end end project.save rescue fail_with("Failed to save Xcode project") diff --git a/tests/ios/run.sh b/cactus-engine/tests/ios/run.sh similarity index 54% rename from tests/ios/run.sh rename to cactus-engine/tests/ios/run.sh index a8d5ac730..5db0401fa 100755 --- a/tests/ios/run.sh +++ b/cactus-engine/tests/ios/run.sh @@ -1,10 +1,16 @@ #!/bin/bash SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +PROJECT_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" +CACTUS_CURL_ROOT="${CACTUS_CURL_ROOT:-$PROJECT_ROOT/cactus-engine/libs/curl}" +export CACTUS_CURL_ROOT -MODEL_NAME="$1" -TRANSCRIBE_MODEL_NAME="$2" +for var in CACTUS_TEST_MODEL CACTUS_TEST_TRANSCRIPTION_MODEL; do + if [ -z "${!var:-}" ] || [ ! -d "${!var}" ]; then + echo "Error: $var must point to a prepared bundle (got: '${!var:-}')" >&2 + exit 1 + fi +done echo "Running Cactus tests on iOS..." echo "============================" @@ -15,19 +21,19 @@ if [ ! -d "/Applications/Xcode.app" ]; then exit 1 fi -if ! xcode-select -p; then +if ! xcode-select -p &>/dev/null; then echo "Xcode Command Line Tools not installed" echo "Install with: xcode-select --install" exit 1 fi -if ! /usr/bin/xcrun --version; then +if ! /usr/bin/xcrun --version &>/dev/null; then echo "Xcode license not accepted" echo "Accept the license with: sudo xcodebuild -license accept" exit 1 fi -if ! command -v xcodebuild; then +if ! command -v xcodebuild &>/dev/null; then echo "xcodebuild not found" exit 1 fi @@ -143,7 +149,10 @@ else echo "Selected: $device_name (Device)" fi - if ! security find-identity -v -p codesigning | grep -q "Apple Development"; then + matching_identity=$(security find-identity -p codesigning 2>/dev/null | grep "Apple Development" | head -1 || true) + valid_identity=$(security find-identity -v -p codesigning 2>/dev/null | grep "Apple Development" | head -1 || true) + + if [ -z "$matching_identity" ]; then echo "No development certificates found" echo "To fix this:" echo " 1. Open Xcode > Settings > Accounts" @@ -151,6 +160,11 @@ else echo " 3. Download development certificates" exit 1 fi + + if [ -z "$valid_identity" ]; then + echo "Warning: Apple Development identity found, but macOS did not report it as a valid identity." + echo "Proceeding and letting xcodebuild attempt automatic signing anyway." + fi fi echo "" @@ -165,8 +179,7 @@ echo "" echo "Step 3: Configuring Xcode project..." xcodeproj_path="$SCRIPT_DIR/CactusTest/CactusTest.xcodeproj" -tests_root="$(cd "$SCRIPT_DIR/.." && pwd)" -cactus_root="$PROJECT_ROOT/cactus" +tests_root="$PROJECT_ROOT/cactus-engine/tests" project_file="$xcodeproj_path/project.pbxproj" template_file="$xcodeproj_path/project.pbxproj.template" @@ -177,20 +190,56 @@ bundle_id="com.cactus.test.${USER}" echo "Using Bundle ID: $bundle_id" if [ "$device_type" = "device" ]; then - development_team=$(security find-certificate -a -c "Apple Development" -p | openssl x509 -noout -subject | grep -oE 'OU=[A-Z0-9]{10}' | head -1 | cut -d= -f2) + if [ -n "$DEVELOPMENT_TEAM" ]; then + development_team="$DEVELOPMENT_TEAM" + else + teams="" + while IFS= read -r line; do + team_id=$(echo "$line" | grep -oE 'OU=[A-Z0-9]+' | cut -d= -f2) + cn=$(echo "$line" | grep -oE 'CN=[^,]+' | cut -d= -f2) + [ -n "$team_id" ] && teams=$(printf "%s\n%s|%s" "$teams" "$team_id" "$cn") + done <<< "$(security find-identity -v -p codesigning 2>/dev/null | grep "Apple Development" | sed 's/.*"\(.*\)".*/\1/' | while read name; do security find-certificate -c "$name" -p 2>/dev/null | openssl x509 -noout -subject 2>/dev/null; done | sort -u)" + teams=$(echo "$teams" | grep -v '^$') + + if [ -z "$teams" ]; then + echo "No Apple Development certificates found" + echo "Add one in Xcode > Settings > Accounts" + exit 1 + fi + + team_count=$(echo "$teams" | wc -l | tr -d ' ') + if [ "$team_count" -eq 1 ]; then + development_team=$(echo "$teams" | head -1 | cut -d'|' -f1) + else + echo "" + echo "Development Teams:" + team_num=0 + while IFS='|' read -r tid tname; do + team_num=$((team_num + 1)) + printf " %2d. %s (%s)\n" "$team_num" "$tname" "$tid" + done <<< "$teams" + echo "" + read -p "Select team (1-$team_num): " team_selection + if ! [[ "$team_selection" =~ ^[0-9]+$ ]] || [ "$team_selection" -lt 1 ] || [ "$team_selection" -gt "$team_num" ]; then + echo "Invalid selection" + exit 1 + fi + development_team=$(echo "$teams" | sed -n "${team_selection}p" | cut -d'|' -f1) + fi + fi if [ -z "$development_team" ]; then - echo "Could not extract Team ID from certificate" + echo "Could not determine Team ID" exit 1 fi echo "Using Team ID: $development_team" fi -if ! command -v ruby; then +if ! command -v ruby &>/dev/null; then echo "Ruby not found" exit 1 fi -if ! gem list xcodeproj -i; then +if ! gem list xcodeproj -i &>/dev/null; then echo "Installing xcodeproj gem..." if ! gem install xcodeproj; then echo "Failed to install xcodeproj gem" @@ -198,7 +247,7 @@ if ! gem list xcodeproj -i; then fi fi -export PROJECT_ROOT TESTS_ROOT="$tests_root" CACTUS_ROOT="$cactus_root" XCODEPROJ_PATH="$xcodeproj_path" BUNDLE_ID="$bundle_id" DEVELOPMENT_TEAM="$development_team" DEVICE_TYPE="$device_type" +export PROJECT_ROOT TESTS_ROOT="$tests_root" XCODEPROJ_PATH="$xcodeproj_path" BUNDLE_ID="$bundle_id" DEVELOPMENT_TEAM="$development_team" DEVICE_TYPE="$device_type" CACTUS_CURL_ROOT="$CACTUS_CURL_ROOT" if ! ruby "$SCRIPT_DIR/configure_xcode.rb"; then echo "Failed to configure Xcode project" exit 1 @@ -249,6 +298,7 @@ else SDKROOT="$ios_sdk_path" \ PRODUCT_BUNDLE_IDENTIFIER="$bundle_id" \ CODE_SIGN_STYLE="Automatic" \ + CODE_SIGN_ENTITLEMENTS="$SCRIPT_DIR/CactusTest/CactusTest/CactusTest.entitlements" \ build; then echo "Build failed" exit 1 @@ -257,27 +307,38 @@ else app_path="$SCRIPT_DIR/build/Build/Products/Release-iphoneos/CactusTest.app" fi -model_dir=$(echo "$MODEL_NAME" | sed 's|.*/||' | tr '[:upper:]' '[:lower:]') -transcribe_model_dir=$(echo "$TRANSCRIBE_MODEL_NAME" | sed 's|.*/||' | tr '[:upper:]' '[:lower:]') -model_src="$PROJECT_ROOT/weights/$model_dir" -transcribe_model_src="$PROJECT_ROOT/weights/$transcribe_model_dir" -assets_src="$PROJECT_ROOT/tests/assets" +echo "" +echo "Step 5: Bundling model weights and assets..." + +model_dir=$(basename "$CACTUS_TEST_MODEL") +transcription_dir=$(basename "$CACTUS_TEST_TRANSCRIPTION_MODEL") +assets_src="$PROJECT_ROOT/cactus-engine/tests/assets" echo "Copying model weights to app bundle..." -if ! cp -R "$model_src" "$app_path/"; then - echo "Warning: Could not copy model weights" +rm -rf "$app_path/$model_dir" +if ! cp -R "$CACTUS_TEST_MODEL" "$app_path/"; then + echo "Error: Could not copy model weights from $CACTUS_TEST_MODEL" + exit 1 fi -if ! cp -R "$transcribe_model_src" "$app_path/"; then - echo "Warning: Could not copy transcribe model weights" + +echo "Copying transcription model to app bundle..." +rm -rf "$app_path/$transcription_dir" +if ! cp -R "$CACTUS_TEST_TRANSCRIPTION_MODEL" "$app_path/"; then + echo "Error: Could not copy transcription model from $CACTUS_TEST_TRANSCRIPTION_MODEL" + exit 1 fi -echo "Copying test assets to app bundle..." -if ! cp -R "$assets_src" "$app_path/"; then - echo "Warning: Could not copy test assets" +if [ -d "$assets_src" ]; then + echo "Copying test assets to app bundle..." + rm -rf "$app_path/assets" + if ! cp -R "$assets_src" "$app_path/"; then + echo "Error: Could not copy test assets from $assets_src" + exit 1 + fi fi echo "" -echo "Step 5: Running tests..." +echo "Step 6: Running tests..." echo "------------------------" if [ "$device_type" = "simulator" ]; then @@ -291,16 +352,31 @@ if [ "$device_type" = "simulator" ]; then fi echo "Launching tests..." - echo "Using model path: $model_dir" - echo "Using transcribe model path: $transcribe_model_dir" - echo "Using assets path: assets" - echo "Using index path: assets" - - SIMCTL_CHILD_CACTUS_TEST_MODEL="$model_dir" \ - SIMCTL_CHILD_CACTUS_TEST_TRANSCRIBE_MODEL="$transcribe_model_dir" \ - SIMCTL_CHILD_CACTUS_TEST_ASSETS="assets" \ - SIMCTL_CHILD_CACTUS_INDEX_PATH="assets" \ - xcrun simctl launch --console-pty "$device_uuid" "$bundle_id" + echo "Using model path: $model_dir" + echo "Using transcription model path: $transcription_dir" + echo "Using assets path: assets" + + sim_env=( + "SIMCTL_CHILD_CACTUS_TEST_MODEL=$model_dir" + "SIMCTL_CHILD_CACTUS_TEST_TRANSCRIPTION_MODEL=$transcription_dir" + "SIMCTL_CHILD_CACTUS_TEST_ASSETS=assets" + "SIMCTL_CHILD_CACTUS_INDEX_PATH=assets" + "SIMCTL_CHILD_CACTUS_NO_CLOUD_TELE=${CACTUS_NO_CLOUD_TELE:-1}" + "SIMCTL_CHILD_CACTUS_TEST_ONLY=${CACTUS_TEST_SUITE:-}" + "SIMCTL_CHILD_CACTUS_TEST_BACKEND=${CACTUS_TEST_BACKEND:-auto}" + ) + + env "${sim_env[@]}" xcrun simctl launch --console-pty --terminate-running-process "$device_uuid" "$bundle_id" + + data_container=$(xcrun simctl get_app_container "$device_uuid" "$bundle_id" data 2>/dev/null || true) + exitcode_file="$data_container/Documents/cactus_test.exitcode" + if [ -n "$data_container" ] && [ -f "$exitcode_file" ]; then + SUITE_EXIT=$(tr -d '[:space:]' < "$exitcode_file") + [[ "$SUITE_EXIT" =~ ^[0-9]+$ ]] || SUITE_EXIT=1 + else + echo "Could not retrieve exit-code marker from simulator" + SUITE_EXIT=1 + fi else echo "Installing on: $device_name" @@ -315,23 +391,29 @@ else echo "Launching tests..." echo "(Logs will be fetched from device after completion)" - echo "Using model path: $model_dir" - echo "Using transcribe model path: $transcribe_model_dir" - echo "Using assets path: assets" - echo "Using index path: assets" - - launch_output=$(DEVICECTL_CHILD_CACTUS_TEST_MODEL="$model_dir" \ - DEVICECTL_CHILD_CACTUS_TEST_TRANSCRIBE_MODEL="$transcribe_model_dir" \ - DEVICECTL_CHILD_CACTUS_TEST_ASSETS="assets" \ - DEVICECTL_CHILD_CACTUS_INDEX_PATH="assets" \ - xcrun devicectl device process launch --device "$device_uuid" "$bundle_id" 2>&1) || true - - echo "$launch_output" - + echo "Using model path: $model_dir" + echo "Using transcription model path: $transcription_dir" + echo "Using assets path: assets" + + device_env=( + "DEVICECTL_CHILD_CACTUS_TEST_MODEL=$model_dir" + "DEVICECTL_CHILD_CACTUS_TEST_TRANSCRIPTION_MODEL=$transcription_dir" + "DEVICECTL_CHILD_CACTUS_TEST_ASSETS=assets" + "DEVICECTL_CHILD_CACTUS_INDEX_PATH=assets" + "DEVICECTL_CHILD_CACTUS_NO_CLOUD_TELE=${CACTUS_NO_CLOUD_TELE:-1}" + "DEVICECTL_CHILD_CACTUS_TEST_ONLY=${CACTUS_TEST_SUITE:-}" + "DEVICECTL_CHILD_CACTUS_TEST_BACKEND=${CACTUS_TEST_BACKEND:-auto}" + ) + + env "${device_env[@]}" \ + xcrun devicectl device process launch --device "$device_uuid" "$bundle_id" 2>&1 + SUITE_EXIT=$? + + echo "Waiting for tests to complete..." max_wait=300 elapsed=0 while [ $elapsed -lt $max_wait ]; do - if xcrun devicectl device info processes --device "$device_uuid" | grep -q "CactusTest.app/CactusTest"; then + if xcrun devicectl device info processes --device "$device_uuid" 2>/dev/null | grep -q "CactusTest.app/CactusTest"; then sleep 2 elapsed=$((elapsed + 2)) else @@ -345,38 +427,42 @@ else sleep 1 + echo "" echo "Fetching logs from device..." - - temp_log_dir=$(mktemp -d) - temp_log_file="$temp_log_dir/cactus_test.log" - - if xcrun devicectl device copy from \ - --device "$device_uuid" \ - --source "Documents/cactus_test.log" \ - --destination "$temp_log_file" \ + log_dir=$(mktemp -d) + xcrun devicectl device copy from --device "$device_uuid" \ + --domain-identifier "$bundle_id" \ --domain-type appDataContainer \ - --domain-identifier "$bundle_id"; then - - if [ -f "$temp_log_file" ]; then - echo "" - cat "$temp_log_file" - else - echo "Warning: Could not find downloaded log file" - fi + --source Documents/cactus_test.log \ + --destination "$log_dir/cactus_test.log" 2>/dev/null || true + xcrun devicectl device copy from --device "$device_uuid" \ + --domain-identifier "$bundle_id" \ + --domain-type appDataContainer \ + --source Documents/cactus_test.exitcode \ + --destination "$log_dir/exitcode" 2>/dev/null || true - rm -rf "$temp_log_dir" + if [ -f "$log_dir/cactus_test.log" ]; then + echo "=== Device Test Output ===" + cat "$log_dir/cactus_test.log" + echo "=== End Device Test Output ===" else - echo "Warning: Could not fetch log file from device" + echo "Could not retrieve test logs from device" fi - if echo "$launch_output" | grep -q "FBSOpenApplicationErrorDomain error 3"; then - echo "" - echo "App launch failed: Untrusted Developer" - echo "To trust the developer profile on your device:" - echo " 1. Open Settings > General > VPN & Device Management" - echo " 2. Under Developer App, tap your Apple ID" - echo " 3. Tap Trust and confirm" - echo "Then run this script again." - echo "" + if [ -f "$log_dir/exitcode" ]; then + SUITE_EXIT=$(tr -d '[:space:]' < "$log_dir/exitcode") + [[ "$SUITE_EXIT" =~ ^[0-9]+$ ]] || SUITE_EXIT=1 + else + echo "Could not retrieve exit-code marker from device" + SUITE_EXIT=1 fi + rm -rf "$log_dir" +fi + +echo "" +if [ "${SUITE_EXIT:-1}" -eq 0 ]; then + echo "All tests passed." +else + echo "Some tests failed." fi +exit "${SUITE_EXIT:-1}" diff --git a/cactus-engine/tests/run.cpp b/cactus-engine/tests/run.cpp new file mode 100644 index 000000000..6e8f520b2 --- /dev/null +++ b/cactus-engine/tests/run.cpp @@ -0,0 +1,1106 @@ +#include "../cactus_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef HAVE_SDL2 +#include +#include +#endif + +namespace { + +constexpr int kMaxTokens = 1024; +constexpr size_t kResponseBufferSize = kMaxTokens * 128; +constexpr int kAudioSampleRate = 16000; + +#ifdef HAVE_SDL2 +constexpr int kRecordSampleRate = kAudioSampleRate; + +struct RecordState { + std::mutex mutex; + std::vector buffer; + std::atomic recording{false}; + int actual_sample_rate = kRecordSampleRate; + SDL_AudioFormat actual_format = AUDIO_S16LSB; + int actual_channels = 1; +}; + +RecordState g_record; + +void record_callback(void*, Uint8* stream, int len) { + if (!g_record.recording) return; + std::lock_guard lock(g_record.mutex); + g_record.buffer.insert(g_record.buffer.end(), stream, stream + len); +} + +std::vector decode_sdl_audio_to_mono_f32(const std::vector& input, + SDL_AudioFormat format, + int channels) { + if (input.empty() || channels <= 0) return {}; + + size_t bytes_per_sample = SDL_AUDIO_BITSIZE(format) / 8; + if (bytes_per_sample == 0) return {}; + size_t frame_count = input.size() / (bytes_per_sample * static_cast(channels)); + std::vector mono(frame_count); + + auto sample_at = [&](size_t sample_index) -> float { + const uint8_t* p = input.data() + sample_index * bytes_per_sample; + switch (format) { + case AUDIO_S16LSB: { + int16_t v; + std::memcpy(&v, p, sizeof(v)); + return static_cast(v) / 32768.0f; + } + case AUDIO_U16LSB: { + uint16_t v; + std::memcpy(&v, p, sizeof(v)); + return (static_cast(v) - 32768.0f) / 32768.0f; + } + case AUDIO_S16MSB: { + int16_t v = static_cast((p[0] << 8) | p[1]); + return static_cast(v) / 32768.0f; + } + case AUDIO_U16MSB: { + uint16_t v = static_cast((p[0] << 8) | p[1]); + return (static_cast(v) - 32768.0f) / 32768.0f; + } + case AUDIO_S8: + return static_cast(*reinterpret_cast(p)) / 128.0f; + case AUDIO_U8: + return (static_cast(*p) - 128.0f) / 128.0f; + case AUDIO_F32LSB: { + float v; + std::memcpy(&v, p, sizeof(v)); + return std::clamp(v, -1.0f, 1.0f); + } + default: + return 0.0f; + } + }; + + for (size_t frame = 0; frame < frame_count; ++frame) { + float sum = 0.0f; + for (int ch = 0; ch < channels; ++ch) { + sum += sample_at(frame * static_cast(channels) + static_cast(ch)); + } + mono[frame] = sum / static_cast(channels); + } + return mono; +} + +std::vector resample_f32_to_s16_pcm(const std::vector& input, int source_rate, int target_rate) { + if (input.empty()) return {}; + double ratio = static_cast(target_rate) / static_cast(source_rate); + size_t out_count = static_cast(static_cast(input.size()) * ratio); + if (out_count == 0) return {}; + + std::vector out(out_count); + for (size_t i = 0; i < out_count; ++i) { + double src_pos = static_cast(i) / ratio; + size_t i0 = static_cast(src_pos); + size_t i1 = std::min(i0 + 1, input.size() - 1); + double frac = src_pos - static_cast(i0); + double sample = static_cast(input[i0]) * (1.0 - frac) + static_cast(input[i1]) * frac; + sample = std::clamp(sample, -1.0, 1.0); + out[i] = static_cast(std::lrint(sample * 32767.0)); + } + + std::vector result(out.size() * sizeof(int16_t)); + std::memcpy(result.data(), out.data(), result.size()); + return result; +} + +bool record_audio(std::vector& pcm_out) { + if (SDL_Init(SDL_INIT_AUDIO) < 0) { + std::cerr << "Failed to init SDL audio: " << SDL_GetError() << "\n"; + return false; + } + + SDL_AudioSpec want; + SDL_AudioSpec have; + SDL_zero(want); + want.freq = kRecordSampleRate; + want.format = AUDIO_S16LSB; + want.channels = 1; + want.samples = static_cast((kRecordSampleRate * 100) / 1000); + want.callback = record_callback; + + SDL_AudioDeviceID device = SDL_OpenAudioDevice(nullptr, 1, &want, &have, + SDL_AUDIO_ALLOW_FREQUENCY_CHANGE | + SDL_AUDIO_ALLOW_FORMAT_CHANGE | + SDL_AUDIO_ALLOW_CHANNELS_CHANGE); + if (device == 0) { + std::cerr << "Failed to open microphone: " << SDL_GetError() << "\n"; + SDL_QuitSubSystem(SDL_INIT_AUDIO); + return false; + } + + { + std::lock_guard lock(g_record.mutex); + g_record.buffer.clear(); + } + g_record.actual_sample_rate = have.freq; + g_record.actual_format = have.format; + g_record.actual_channels = have.channels; + g_record.recording = true; + SDL_PauseAudioDevice(device, 0); + + std::cout << "Recording... press Enter to stop.\n" << std::flush; + std::string line; + std::getline(std::cin, line); + + g_record.recording = false; + SDL_PauseAudioDevice(device, 1); + + { + std::lock_guard lock(g_record.mutex); + auto mono = decode_sdl_audio_to_mono_f32(g_record.buffer, + g_record.actual_format, + g_record.actual_channels); + pcm_out = resample_f32_to_s16_pcm(mono, g_record.actual_sample_rate, kRecordSampleRate); + } + + SDL_CloseAudioDevice(device); + SDL_QuitSubSystem(SDL_INIT_AUDIO); + + double seconds = static_cast(pcm_out.size() / sizeof(int16_t)) / kRecordSampleRate; + std::cout << "Recorded " << std::fixed << std::setprecision(1) << seconds << "s of audio.\n"; + return !pcm_out.empty(); +} +#endif + +namespace ansi { +constexpr const char* reset = "\033[0m"; +constexpr const char* bold = "\033[1m"; +constexpr const char* dim = "\033[2m"; +constexpr const char* italic = "\033[3m"; +constexpr const char* underline = "\033[4m"; +constexpr const char* cyan = "\033[36m"; +constexpr const char* green = "\033[32m"; +constexpr const char* yellow = "\033[33m"; +constexpr const char* blue = "\033[34m"; +} // namespace ansi + +bool stdout_is_terminal() { + return isatty(STDOUT_FILENO) != 0 && std::getenv("NO_COLOR") == nullptr; +} + +int terminal_width() { + struct winsize ws; + if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_col > 0) return ws.ws_col; + if (const char* cols = std::getenv("COLUMNS")) { + int c = std::atoi(cols); + if (c > 0) return c; + } + return 80; +} + +void print_turn_separator() { + std::cout << "\n"; + if (stdout_is_terminal()) { + std::string rule; + for (int i = 0, w = terminal_width(); i < w; ++i) rule += "─"; + std::cout << ansi::dim << rule << ansi::reset << "\n"; + } +} + +std::atomic g_cloud_active{false}; + +void print_banner() { + if (!stdout_is_terminal()) { + std::cout << "Cactus Hybrid Chat\n\n"; + return; + } + std::cout << ansi::bold << ansi::green + << " ██████╗ █████╗ ██████╗████████╗██╗ ██╗███████╗\n" + << "██╔════╝██╔══██╗██╔════╝╚══██╔══╝██║ ██║██╔════╝\n" + << "██║ ███████║██║ ██║ ██║ ██║███████╗\n" + << "██║ ██╔══██║██║ ██║ ██║ ██║╚════██║\n" + << "╚██████╗██║ ██║╚██████╗ ██║ ╚██████╔╝███████║\n" + << " ╚═════╝╚═╝ ╚═╝ ╚═════╝ ╚═╝ ╚═════╝ ╚══════╝\n" + << ansi::reset + << ansi::dim << " Hybrid Chat" << ansi::reset << "\n\n"; +} + +struct TokenPrinter { + std::chrono::steady_clock::time_point start; + std::chrono::steady_clock::time_point first; + bool saw_first = false; + int count = 0; + std::string pending; + bool in_code_fence = false; + bool color = false; + bool suppress_thinking_stream = false; + bool show_thinking = false; + std::string stream_carry; + bool label_printed = false; + std::thread spinner_thread; + std::atomic spinner_running{false}; + + ~TokenPrinter() { stop_thinking(); } + + void reset() { + start = std::chrono::steady_clock::now(); + saw_first = false; + count = 0; + pending.clear(); + in_code_fence = false; + label_printed = false; + color = stdout_is_terminal(); + suppress_thinking_stream = false; + stream_carry.clear(); + } + + void start_thinking() { + if (!color || spinner_running.load()) return; + spinner_running = true; + spinner_thread = std::thread([this]() { + const char* frames[] = {"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}; + for (int idx = 0; spinner_running.load(); ++idx) { + std::cout << "\r" << ansi::dim << frames[idx % 10] << " thinking…" << ansi::reset << std::flush; + for (int k = 0; k < 5 && spinner_running.load(); ++k) { + std::this_thread::sleep_for(std::chrono::milliseconds(16)); + } + } + }); + } + + void stop_thinking() { + if (spinner_running.exchange(false)) { + if (spinner_thread.joinable()) spinner_thread.join(); + std::cout << "\r\033[K" << std::flush; + } + } + + void print_label() { + if (label_printed) return; + label_printed = true; + stop_thinking(); + const char* origin = g_cloud_active.load() ? "cloud" : "local"; + if (color) { + const char* origin_color = g_cloud_active.load() ? ansi::yellow : ansi::cyan; + std::cout << ansi::bold << ansi::green << "Assistant" << ansi::reset + << " " << origin_color << "(" << origin << ")" << ansi::reset + << ansi::bold << ansi::green << ":" << ansi::reset << " "; + } else { + std::cout << "Assistant (" << origin << "): "; + } + std::cout << std::flush; + } + + // Render the markdown inline spans (**bold**, *italic*, `code`) within one line. + void render_inline(const std::string& s) const { + if (!color) { std::cout << s; return; } + size_t i = 0, n = s.size(); + while (i < n) { + char c = s[i]; + if (c == '*' && i + 1 < n && s[i + 1] == '*') { + size_t end = s.find("**", i + 2); + if (end != std::string::npos && end > i + 2) { + std::cout << ansi::bold << s.substr(i + 2, end - (i + 2)) << ansi::reset; + i = end + 2; + continue; + } + } + if (c == '`') { + size_t end = s.find('`', i + 1); + if (end != std::string::npos) { + std::cout << ansi::cyan << s.substr(i + 1, end - (i + 1)) << ansi::reset; + i = end + 1; + continue; + } + } + + if (c == '[') { + size_t rb = s.find(']', i + 1); + if (rb != std::string::npos && rb + 1 < n && s[rb + 1] == '(') { + size_t rp = s.find(')', rb + 2); + if (rp != std::string::npos) { + std::cout << ansi::underline << ansi::cyan + << s.substr(i + 1, rb - (i + 1)) << ansi::reset; + i = rp + 1; + continue; + } + } + } + if (c == '*' && i + 1 < n && s[i + 1] != '*' && s[i + 1] != ' ') { + size_t end = s.find('*', i + 1); + if (end != std::string::npos && s[end - 1] != ' ') { + std::cout << ansi::italic << s.substr(i + 1, end - (i + 1)) << ansi::reset; + i = end + 1; + continue; + } + } + std::cout << c; + ++i; + } + } + + // Render a single complete line, handling block-level markdown (headings, lists, code fences). + void render_line(const std::string& raw) { + if (!color) { std::cout << raw; return; } + + size_t s0 = raw.find_first_not_of(" \t"); + std::string lead = (s0 == std::string::npos) ? raw : raw.substr(0, s0); + std::string body = (s0 == std::string::npos) ? "" : raw.substr(s0); + + if (body.rfind("```", 0) == 0) { + if (!in_code_fence) { + in_code_fence = true; + std::string lang = body.substr(3); + size_t a = lang.find_first_not_of(" \t"); + size_t b = lang.find_last_not_of(" \t"); + lang = (a == std::string::npos) ? "" : lang.substr(a, b - a + 1); + std::cout << ansi::dim << " ┌── " << (lang.empty() ? "code" : lang) << ansi::reset; + } else { + in_code_fence = false; + std::cout << ansi::dim << " └──" << ansi::reset; + } + return; + } + if (in_code_fence) { + std::cout << ansi::dim << " │ " << ansi::reset << ansi::green << raw << ansi::reset; + return; + } + + if (!body.empty() && body[0] == '#') { + size_t h = 0; + while (h < body.size() && body[h] == '#') ++h; + if (h >= 1 && h <= 6 && h < body.size() && body[h] == ' ') { + std::cout << lead << ansi::bold << ansi::cyan << body.substr(h + 1) << ansi::reset; + return; + } + } + + if (body[0] == '-' || body[0] == '*' || body[0] == '_') { + char hc = body[0]; + size_t marks = 0; + bool only = true; + for (char ch : body) { + if (ch == hc) ++marks; + else if (ch != ' ') { only = false; break; } + } + if (only && marks >= 3) { + std::cout << lead << ansi::dim + << "────────────────────────────" + << ansi::reset; + return; + } + } + + if (body[0] == '>') { + std::string rest = body.substr(1); + if (!rest.empty() && rest[0] == ' ') rest.erase(0, 1); + std::cout << lead << ansi::dim << "▏ " << ansi::reset; + render_inline(rest); + return; + } + + // Bullet list: -, *, + then space(s) + if (body.size() >= 2 && (body[0] == '-' || body[0] == '*' || body[0] == '+') && body[1] == ' ') { + std::string rest = body.substr(2); + size_t r0 = rest.find_first_not_of(" \t"); + rest = (r0 == std::string::npos) ? "" : rest.substr(r0); + std::cout << lead << ansi::yellow << "• " << ansi::reset; + render_inline(rest); + return; + } + + size_t d = 0; + while (d < body.size() && std::isdigit(static_cast(body[d]))) ++d; + if (d > 0 && d + 1 < body.size() && body[d] == '.' && body[d + 1] == ' ') { + std::string rest = body.substr(d + 1); + size_t r0 = rest.find_first_not_of(" \t"); + rest = (r0 == std::string::npos) ? "" : rest.substr(r0); + std::cout << lead << ansi::yellow << body.substr(0, d + 1) << ansi::reset << " "; + render_inline(rest); + return; + } + + std::cout << lead; + render_inline(body); + } + + void on_token(const char* text) { + if (!saw_first) { + first = std::chrono::steady_clock::now(); + saw_first = true; + } + ++count; + if (!text) return; + feed(text); + std::cout << std::flush; + } + + void feed(const std::string& chunk) { + static const std::vector markers = { + "<|channel>", "", "<|turn>", "", "<|think|>", "", + }; + stream_carry += chunk; + size_t i = 0; + while (i < stream_carry.size()) { + const std::string* matched = nullptr; + bool maybe_partial = false; + for (const auto& m : markers) { + if (stream_carry.compare(i, m.size(), m) == 0) { matched = &m; break; } + size_t avail = stream_carry.size() - i; + if (avail < m.size() && stream_carry.compare(i, avail, m, 0, avail) == 0) maybe_partial = true; + } + if (matched) { + if (*matched == "<|channel>") { suppress_thinking_stream = true; if (!show_thinking) start_thinking(); } + else if (*matched == "") { suppress_thinking_stream = false; stop_thinking(); } + i += matched->size(); + continue; + } + if (maybe_partial) break; + if (!suppress_thinking_stream || show_thinking) emit_visible(stream_carry[i]); + ++i; + } + stream_carry.erase(0, i); + } + + void emit_visible(char c) { + print_label(); + pending += c; + size_t nl; + while ((nl = pending.find('\n')) != std::string::npos) { + render_line(pending.substr(0, nl)); + std::cout << "\n"; + pending.erase(0, nl + 1); + } + } + + // Flush the trailing partial line once generation completes (no trailing newline). + void finish() { + stop_thinking(); + if (!stream_carry.empty() && (!suppress_thinking_stream || show_thinking)) { + for (char c : stream_carry) emit_visible(c); + } + stream_carry.clear(); + print_label(); + if (!pending.empty()) { + render_line(pending); + pending.clear(); + } + if (color && in_code_fence) std::cout << ansi::reset; + in_code_fence = false; + std::cout << std::flush; + } + + void print_stats(double ram_mb, double confidence, bool cloud_handoff, double threshold, + int reported_tokens = -1, double reported_decode_tps = -1.0, + double reported_ttft_ms = -1.0, double reported_total_ms = -1.0, + const std::string& cloud_reason = "") const { + auto end = std::chrono::steady_clock::now(); + double total_s = std::chrono::duration(end - start).count(); + double ttft_s = saw_first ? std::chrono::duration(first - start).count() : 0.0; + double decode_s = saw_first ? std::chrono::duration(end - first).count() : total_s; + int display_tokens = reported_tokens >= 0 ? reported_tokens : count; + double tps = (count > 1 && decode_s > 0.0) ? (count - 1) / decode_s : (total_s > 0.0 ? count / total_s : 0.0); + if (reported_decode_tps >= 0.0) tps = reported_decode_tps; + if (reported_ttft_ms >= 0.0) ttft_s = reported_ttft_ms / 1000.0; + if (reported_total_ms >= 0.0) total_s = reported_total_ms / 1000.0; + std::cout << "\n"; + if (color) std::cout << ansi::dim; + std::cout << "[" << display_tokens << " tokens | latency: " + << std::fixed << std::setprecision(3) << ttft_s + << "s | total: " << total_s + << "s | " << std::setprecision(1) << tps << " tok/s"; + if (confidence >= 0.0) { + std::cout << " | confidence: " << std::max(0.0, std::min(100.0, confidence * 100.0)) << "%"; + } else { + std::cout << " | confidence: n/a"; + } + if (threshold >= 0.0) { + std::cout << " | threshold: " << std::max(0.0, std::min(100.0, threshold * 100.0)) << "%"; + } + std::cout << " | cloud: " << (cloud_handoff ? "yes" : "no"); + if (!cloud_reason.empty()) { + std::cout << " (" << cloud_reason << ")"; + } + if (ram_mb > 0.0) { + std::cout << " | RAM: " << ram_mb << " MB"; + } + std::cout << "]"; + if (color) std::cout << ansi::reset; + std::cout << "\n"; + } +}; + +TokenPrinter* g_printer = nullptr; + +void token_callback(const char* text, uint32_t, void*) { + if (g_printer) { + g_printer->on_token(text); + } +} + +void log_callback(int level, const char* component, const char* message, void*) { + const std::string comp = component ? component : ""; + const std::string msg = message ? message : ""; + if (comp == "cloud_handoff") { + if (msg.find("triggered") != std::string::npos) { + g_cloud_active = true; + } else if (msg.find("falling back") != std::string::npos || + msg.find("failed") != std::string::npos || + msg.find("disabling") != std::string::npos) { + g_cloud_active = false; + } + return; + } + const bool color = isatty(STDERR_FILENO) != 0 && std::getenv("NO_COLOR") == nullptr; + const char* names[] = {"DEBUG", "INFO", "WARN", "ERROR"}; + const char* name = (level >= 0 && level <= 3) ? names[level] : "LOG"; + std::cerr << "\n"; + if (color) std::cerr << ansi::dim; + std::cerr << "[" << name << "] " << comp << ": " << msg; + if (color) std::cerr << ansi::reset; + std::cerr << "\n"; +} + +std::string escape_json(const std::string& s) { + std::ostringstream out; + for (unsigned char c : s) { + switch (c) { + case '"': out << "\\\""; break; + case '\\': out << "\\\\"; break; + case '\b': out << "\\b"; break; + case '\f': out << "\\f"; break; + case '\n': out << "\\n"; break; + case '\r': out << "\\r"; break; + case '\t': out << "\\t"; break; + default: + if (c < 0x20) { + out << "\\u" << std::hex << std::setw(4) << std::setfill('0') << static_cast(c); + } else { + out << c; + } + } + } + return out.str(); +} + +std::string unescape_json(const std::string& s) { + std::string out; + out.reserve(s.size()); + for (size_t i = 0; i < s.size(); ++i) { + if (s[i] != '\\' || i + 1 >= s.size()) { + out.push_back(s[i]); + continue; + } + char n = s[++i]; + switch (n) { + case '"': out.push_back('"'); break; + case '\\': out.push_back('\\'); break; + case 'b': out.push_back('\b'); break; + case 'f': out.push_back('\f'); break; + case 'n': out.push_back('\n'); break; + case 'r': out.push_back('\r'); break; + case 't': out.push_back('\t'); break; + default: out.push_back(n); break; + } + } + return out; +} + +std::string expand_tilde(const std::string& path) { + if (path.size() < 2 || path[0] != '~' || path[1] != '/') return path; + const char* home = std::getenv("HOME"); + return home ? std::string(home) + path.substr(1) : path; +} + +bool file_exists(const std::string& path) { + std::ifstream f(path); + return f.good(); +} + +std::vector parse_token_ids(const std::string& text) { + std::vector tokens; + std::string normalized = text; + for (char& ch : normalized) { + if (ch == ' ' || ch == '\t' || ch == '\r' || ch == '\n') ch = ','; + } + std::stringstream ss(normalized); + std::string part; + while (std::getline(ss, part, ',')) { + size_t start = part.find_first_not_of(" \t\r\n"); + if (start == std::string::npos) continue; + size_t end = part.find_last_not_of(" \t\r\n"); + std::string value = part.substr(start, end - start + 1); + char* parsed_end = nullptr; + unsigned long parsed = std::strtoul(value.c_str(), &parsed_end, 10); + if (!parsed_end || *parsed_end != '\0') { + throw std::runtime_error("Invalid token id: " + value); + } + tokens.push_back(static_cast(parsed)); + } + return tokens; +} + +bool read_text_file(const std::string& path, std::string& text) { + std::ifstream in(path); + if (!in) return false; + std::ostringstream buffer; + buffer << in.rdbuf(); + if (!in.good() && !in.eof()) return false; + text = buffer.str(); + return true; +} + +bool write_text_file(const std::string& path, const std::string& text) { + std::ofstream out(path); + if (!out) return false; + out << text; + return out.good(); +} + +std::string json_string_value(const std::string& json, const std::string& key) { + std::string needle = "\"" + key + "\":\""; + size_t start = json.find(needle); + if (start == std::string::npos) return {}; + start += needle.size(); + size_t end = start; + while ((end = json.find('"', end)) != std::string::npos) { + size_t slashes = 0; + for (size_t i = end; i > start && json[i - 1] == '\\'; --i) ++slashes; + if ((slashes % 2) == 0) break; + ++end; + } + if (end == std::string::npos) return {}; + return unescape_json(json.substr(start, end - start)); +} + +double json_number_value(const std::string& json, const std::string& key, double fallback = 0.0) { + std::string needle = "\"" + key + "\":"; + size_t start = json.find(needle); + if (start == std::string::npos) return fallback; + start += needle.size(); + while (start < json.size() && std::isspace(static_cast(json[start]))) { + ++start; + } + if (json.compare(start, 4, "null") == 0) return fallback; + char* end = nullptr; + return std::strtod(json.c_str() + start, &end); +} + +bool json_bool_value(const std::string& json, const std::string& key) { + std::string needle = "\"" + key + "\":"; + size_t start = json.find(needle); + if (start == std::string::npos) return false; + start += needle.size(); + while (start < json.size() && std::isspace(static_cast(json[start]))) { + ++start; + } + return json.compare(start, 4, "true") == 0; +} + +struct ChatTurn { + std::string role; + std::string content; + std::string image; + std::string audio; +}; + +std::string build_messages(const std::string& system_prompt, + const std::vector& history) { + std::ostringstream msg; + msg << "["; + bool need_comma = false; + if (!system_prompt.empty()) { + msg << "{\"role\":\"system\",\"content\":\"" << escape_json(system_prompt) << "\"}"; + need_comma = true; + } + for (const auto& turn : history) { + if (need_comma) msg << ","; + need_comma = true; + msg << "{\"role\":\"" << turn.role << "\",\"content\":\"" + << escape_json(turn.content) << "\""; + if (!turn.image.empty()) msg << ",\"images\":[\"" << escape_json(turn.image) << "\"]"; + if (!turn.audio.empty()) msg << ",\"audio\":[\"" << escape_json(turn.audio) << "\"]"; + msg << "}"; + } + msg << "]"; + return msg.str(); +} + +void print_usage(const char* argv0) { + std::cerr << "Usage: " << argv0 + << " [--backend cpu|metal] [--system ] [--image ] [--audio ]" + << " [--prompt ] [--input-ids ] [--input-ids-file ] [--max-new-tokens ]" + << " [--result-json ] [--thinking] [--no-cloud-handoff]" + << " [--confidence-threshold ] [--cloud-timeout-ms ] [-h|--help]\n"; +} + +} // namespace + +int main(int argc, char** argv) { + std::cout << std::unitbuf; + std::cerr << std::unitbuf; + + if (argc < 2) { + print_usage(argv[0]); + return 1; + } + + std::string model_path = argv[1]; + std::string system_prompt; + std::string current_image; + std::string current_audio; + std::string initial_prompt; + std::string input_ids; + std::string input_ids_file; + std::string result_json; + int max_new_tokens = kMaxTokens; + bool thinking = false; + bool auto_handoff = true; + double confidence_threshold = -1.0; + int cloud_timeout_ms = 15000; + std::string backend; + + int i = 2; + auto need_value = [&](const char* flag) -> bool { + if (i + 1 >= argc) { + std::cerr << "Error: " << flag << " requires a value\n"; + print_usage(argv[0]); + return false; + } + return true; + }; + for (; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "-h" || arg == "--help") { + print_usage(argv[0]); + return 0; + } else if (arg == "--system") { + if (!need_value("--system")) return 1; + system_prompt = argv[++i]; + } else if (arg == "--image") { + if (!need_value("--image")) return 1; + current_image = expand_tilde(argv[++i]); + } else if (arg == "--audio") { + if (!need_value("--audio")) return 1; + current_audio = expand_tilde(argv[++i]); + } else if (arg == "--prompt") { + if (!need_value("--prompt")) return 1; + initial_prompt = argv[++i]; + } else if (arg == "--input-ids") { + if (!need_value("--input-ids")) return 1; + input_ids = argv[++i]; + } else if (arg == "--input-ids-file") { + if (!need_value("--input-ids-file")) return 1; + input_ids_file = expand_tilde(argv[++i]); + } else if (arg == "--max-new-tokens") { + if (!need_value("--max-new-tokens")) return 1; + max_new_tokens = std::max(1, std::atoi(argv[++i])); + } else if (arg == "--result-json") { + if (!need_value("--result-json")) return 1; + result_json = argv[++i]; + } else if (arg == "--thinking") { + thinking = true; + } else if (arg == "--no-cloud-handoff") { + auto_handoff = false; + } else if (arg == "--confidence-threshold") { + if (!need_value("--confidence-threshold")) return 1; + confidence_threshold = std::atof(argv[++i]); + } else if (arg == "--cloud-timeout-ms") { + if (!need_value("--cloud-timeout-ms")) return 1; + cloud_timeout_ms = std::max(0, std::atoi(argv[++i])); + } else if (arg == "--backend") { + if (!need_value("--backend")) return 1; + backend = argv[++i]; + } else { + std::cerr << "Error: unknown option '" << arg << "'\n"; + print_usage(argv[0]); + return 1; + } + } + + if (!current_image.empty() && !file_exists(current_image)) { + std::cerr << "Image not found: " << current_image << "\n"; + return 1; + } + if (!current_audio.empty() && !file_exists(current_audio)) { + std::cerr << "Audio file not found: " << current_audio << "\n"; + return 1; + } + if (!input_ids.empty() && !input_ids_file.empty()) { + std::cerr << "Use either --input-ids or --input-ids-file, not both\n"; + return 1; + } + if (!input_ids_file.empty()) { + if (!file_exists(input_ids_file)) { + std::cerr << "Input ids file not found: " << input_ids_file << "\n"; + return 1; + } + if (!read_text_file(input_ids_file, input_ids)) { + std::cerr << "Failed to read input ids file: " << input_ids_file << "\n"; + return 1; + } + } + + if (backend == "cpu" || backend == "metal") { + if (cactus_set_backend(backend.c_str()) == 0) + std::cout << "Backend: " << (backend == "metal" ? "Metal GPU" : "CPU") << "\n"; + else + std::cout << "Metal not available; using CPU\n"; + } else if (!backend.empty()) { + std::cerr << "Error: unknown backend '" << backend << "' (expected 'cpu' or 'metal')\n"; + return 1; + } + + cactus_log_set_callback(log_callback, nullptr); + + std::cout << "Loading model from " << model_path << "...\n"; + cactus_model_t model = cactus_init(model_path.c_str(), nullptr, false); + if (!model) { + std::cerr << "Failed to initialize model\n"; + return 1; + } + + if (!input_ids.empty()) { + std::vector tokens; + try { + tokens = parse_token_ids(input_ids); + } catch (const std::exception& e) { + std::cerr << e.what() << "\n"; + cactus_destroy(model); + return 1; + } + if (tokens.empty()) { + std::cerr << "--input-ids did not contain any token ids\n"; + cactus_destroy(model); + return 1; + } + std::vector response(kResponseBufferSize, 0); + int rc = cactus_benchmark_tokens( + model, + tokens.data(), + tokens.size(), + static_cast(max_new_tokens), + response.data(), + response.size()); + std::string response_json(response.data()); + if (!result_json.empty() && !write_text_file(result_json, response_json)) { + std::cerr << "Failed to write result JSON: " << result_json << "\n"; + cactus_destroy(model); + return 1; + } + std::cout << response_json << "\n"; + cactus_destroy(model); + return rc < 0 ? 1 : 0; + } + + if (stdout_is_terminal()) std::cout << "\033[2J\033[3J\033[H"; + print_banner(); + + { + const bool tty = stdout_is_terminal(); + const char* d = tty ? ansi::dim : ""; + const char* r = tty ? ansi::reset : ""; + if (!auto_handoff) { + std::cout << d << "Hybrid handoff off (--no-cloud-handoff): answering fully on-device." << r << "\n\n"; + } else { + std::cout << d + << "Hybrid mode: answers on-device, handing off to the cloud whenever the local\n" + << "model's confidence drops below its default threshold. Set your own with\n" + << "--confidence-threshold <0-1>, or pass --no-cloud-handoff to stay fully local.\n" + << "Each reply's confidence and the active threshold are shown in its stats line." << r << "\n\n"; + } + } + + std::cout << (stdout_is_terminal() ? "\033[1mCommands:\033[0m\n" : "Commands:\n"); + auto print_command = [](const char* command, const char* description) { + if (stdout_is_terminal()) { + std::cout << " " << ansi::cyan << std::left << std::setw(24) << command << ansi::reset + << ansi::dim << description << ansi::reset << "\n"; + } else { + std::cout << " " << std::left << std::setw(24) << command << description << "\n"; + } + }; + print_command("/image [prompt]", "attach an image and chat about it"); + print_command("/audio [prompt]", "attach an audio file as a spoken prompt"); +#ifdef HAVE_SDL2 + print_command("/record [prompt]", "record from the mic as a spoken prompt"); +#endif + print_command("/clear", "drop the attached image/audio"); + print_command("reset", "start a new conversation"); + print_command("exit", "quit"); + std::cout << std::right << "\n"; + + std::vector history; + std::vector current_pcm; + TokenPrinter printer; + g_printer = &printer; + printer.show_thinking = thinking; + bool auto_send = !initial_prompt.empty() || !current_audio.empty() || !current_image.empty(); + + while (true) { + std::string input; + if (auto_send) { + auto_send = false; + if (!initial_prompt.empty()) input = initial_prompt; + else if (!current_image.empty()) input = "Describe this image."; + else input = "Respond to the spoken request in this audio."; + const char* you = stdout_is_terminal() ? "\033[1;34mYou:\033[0m " : "You: "; + std::cout << you << input << "\n"; + } else { + const char* you = stdout_is_terminal() ? "\033[1;34mYou:\033[0m " : "You: "; + std::cout << you << std::flush; + if (!std::getline(std::cin, input)) break; + } + + while (!input.empty() && (input.back() == ' ' || input.back() == '\t')) input.pop_back(); + if (input.empty()) continue; + if (input == "exit" || input == "quit") break; + if (input == "reset") { + history.clear(); + current_image.clear(); + current_audio.clear(); + current_pcm.clear(); + cactus_reset(model); + std::cout << "Conversation reset.\n"; + continue; + } + if (input == "/clear") { + current_image.clear(); + current_audio.clear(); + current_pcm.clear(); + std::cout << "Attachments cleared.\n"; + continue; + } + + auto parse_attachment = [&](const std::string& prefix, std::string& target) -> bool { + if (input.rfind(prefix, 0) != 0) return false; + std::string rest = input.substr(prefix.size()); + size_t split = rest.find(' '); + std::string path = expand_tilde(split == std::string::npos ? rest : rest.substr(0, split)); + if (!file_exists(path)) { + std::cerr << "File not found: " << path << "\n"; + input.clear(); + return true; + } + target = path; + input = split == std::string::npos ? "" : rest.substr(split + 1); + return true; + }; + + if (parse_attachment("/image ", current_image) && input.empty()) { + std::cout << "Image attached: " << current_image << "\n"; + continue; + } + if (parse_attachment("/audio ", current_audio) && input.empty()) { + std::cout << "Audio attached: " << current_audio << "\n"; + continue; + } + + static const std::string kRecordCmd = "/record"; + if (input == kRecordCmd || input.rfind(kRecordCmd + " ", 0) == 0) { +#ifdef HAVE_SDL2 + std::string record_prompt; + if (input.size() > kRecordCmd.size() + 1) { + record_prompt = input.substr(kRecordCmd.size() + 1); + while (!record_prompt.empty() && (record_prompt.front() == ' ' || record_prompt.front() == '\t')) { + record_prompt.erase(record_prompt.begin()); + } + } + current_pcm.clear(); + current_audio.clear(); + if (!record_audio(current_pcm)) { + std::cerr << "Recording failed.\n"; + continue; + } + input = record_prompt.empty() ? "Respond to the spoken request in this audio." : record_prompt; +#else + std::cerr << "Recording requires SDL2, but this binary was built without SDL2.\n"; + continue; +#endif + } + if (input.empty()) continue; + + bool attach_media = !current_image.empty() || !current_audio.empty() || !current_pcm.empty(); + if (attach_media) { + cactus_reset(model); + } + history.push_back({"user", input, current_image, current_audio}); + std::string messages = build_messages(system_prompt, history); + std::string options = "{\"temperature\":0.7,\"top_p\":0.95,\"top_k\":40,\"max_tokens\":" + + std::to_string(max_new_tokens) + + ",\"enable_thinking_if_supported\":" + (thinking ? "true" : "false") + + ",\"auto_handoff\":" + (auto_handoff ? "true" : "false") + + ",\"confidence_threshold\":" + std::to_string(confidence_threshold) + + ",\"cloud_timeout_ms\":" + std::to_string(cloud_timeout_ms) + + ",\"stop_sequences\":[\"<|im_end|>\",\"\",\"\"]}"; + + if (!current_image.empty()) std::cout << "[image: " << current_image << "]\n"; + if (!current_audio.empty()) std::cout << "[audio: " << current_audio << "]\n"; + if (!current_pcm.empty()) { + double seconds = static_cast(current_pcm.size() / sizeof(int16_t)) / static_cast(kAudioSampleRate); + std::cout << "[recorded audio: " << std::fixed << std::setprecision(1) << seconds << "s]\n"; + } + std::vector response(kResponseBufferSize, 0); + g_cloud_active = false; + printer.reset(); + printer.start_thinking(); + int rc = cactus_complete(model, + messages.c_str(), + response.data(), + response.size(), + options.c_str(), + nullptr, + token_callback, + nullptr, + current_pcm.empty() ? nullptr : current_pcm.data(), + current_pcm.size()); + printer.finish(); + + std::string response_json(response.data()); + if (!result_json.empty() && !write_text_file(result_json, response_json)) { + std::cerr << "Failed to write result JSON: " << result_json << "\n"; + } + bool cloud_handoff = json_bool_value(response_json, "cloud_handoff"); + double confidence = json_number_value(response_json, "confidence", -1.0); + double threshold = json_number_value(response_json, "confidence_threshold", -1.0); + double ram_mb = json_number_value(response_json, "ram_usage_mb"); + int decode_tokens = static_cast(json_number_value(response_json, "decode_tokens", -1.0)); + double decode_tps = json_number_value(response_json, "decode_tps", -1.0); + double ttft_ms = json_number_value(response_json, "time_to_first_token_ms", -1.0); + double total_ms = json_number_value(response_json, "total_time_ms", -1.0); + std::string cloud_reason = json_string_value(response_json, "cloud_handoff_reason"); + printer.print_stats(ram_mb, confidence, cloud_handoff, threshold, decode_tokens, decode_tps, ttft_ms, total_ms, cloud_reason); + + if (rc < 0) { + std::cout << "Error: " << response.data() << "\n"; + history.pop_back(); + continue; + } + + std::string assistant = json_string_value(response_json, "context_response"); + if (assistant.empty()) assistant = json_string_value(response_json, "response"); + history.push_back({"assistant", assistant, "", ""}); + current_image.clear(); + current_audio.clear(); + current_pcm.clear(); + print_turn_separator(); + } + + cactus_destroy(model); + std::cout << "Goodbye.\n"; + return 0; +} diff --git a/cactus-engine/tests/test_benchmark.cpp b/cactus-engine/tests/test_benchmark.cpp new file mode 100644 index 000000000..a24144ecc --- /dev/null +++ b/cactus-engine/tests/test_benchmark.cpp @@ -0,0 +1,292 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +using namespace EngineTestUtils; + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); +static const char* g_transcription_path = std::getenv("CACTUS_TEST_TRANSCRIPTION_MODEL"); +static const char* g_assets_path = std::getenv("CACTUS_TEST_ASSETS"); + +static constexpr size_t kPrefillTokens = 1000; +static constexpr size_t kDecodeTokens = 100; +static constexpr int kMeasuredRuns = 3; + +static void print_mean(const char* label, const char* unit, const std::vector& samples) { + double mean = std::accumulate(samples.begin(), samples.end(), 0.0) / samples.size(); + std::cout << "├─ " << std::left << std::setw(17) << label << std::right + << std::fixed << std::setprecision(2) << std::setw(12) << mean + << " " << unit << "\n"; +} + +static void print_header(const char* title) { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << std::string(" ") + title << "║\n" + << "╚══════════════════════════════════════════╝\n"; +} + +static std::string build_context_messages(int filler) { + std::string msg = "[{\"role\": \"system\", \"content\": \"/no_think You are helpful. "; + for (int i = 0; i < filler; i++) { + msg += "Context " + std::to_string(i) + ": Background knowledge. "; + } + msg += "\"}, {\"role\": \"user\", \"content\": \""; + for (int i = 0; i < filler; i++) { + msg += "Data " + std::to_string(i) + " = " + std::to_string(i * 3.14159) + ". "; + } + msg += "Explain the data.\"}]"; + return msg; +} + +static bool build_exact_prompt_tokens(cactus_model_t model, std::vector& tokens) { + for (int filler = 50; filler <= 3200 && tokens.size() < kPrefillTokens; filler *= 2) { + std::string msg = build_context_messages(filler); + std::vector prompt(msg.size() + 4096); + if (cactus_render_prompt(model, msg.c_str(), nullptr, nullptr, + prompt.data(), prompt.size()) < 0) { + std::cerr << "[✗] Failed to render prompt\n"; + return false; + } + size_t token_len = 0; + std::vector buffer(prompt.size()); + if (cactus_tokenize(model, prompt.data(), buffer.data(), buffer.size(), &token_len) != 0) { + std::cerr << "[✗] Failed to tokenize prompt\n"; + return false; + } + buffer.resize(token_len); + tokens = std::move(buffer); + } + if (tokens.size() < kPrefillTokens) { + std::cerr << "[✗] Could not build a " << kPrefillTokens << "-token prompt\n"; + return false; + } + std::cout << "├─ Prompt tokens: " << kPrefillTokens + << " (rendered " << tokens.size() << ", truncated)\n" + << "├─ Decode tokens: " << kDecodeTokens << "\n"; + tokens.resize(kPrefillTokens); + return true; +} + +static bool llm_benchmark() { + print_header("LLM BENCHMARK (1k prefill / 100 decode)"); + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + std::vector tokens; + if (!build_exact_prompt_tokens(model, tokens)) { + cactus_destroy(model); + return false; + } + + std::vector prefill_tps, decode_tps, peak_ram; + for (int run = 0; run <= kMeasuredRuns; run++) { + char response[16384]; + int rc = cactus_benchmark_tokens(model, tokens.data(), tokens.size(), kDecodeTokens, + response, sizeof(response)); + std::string json(response); + bool exact = rc >= 0 && + json.find("\"success\":true") != std::string::npos && + json_number(json, "prompt_tokens") == static_cast(kPrefillTokens) && + json_number(json, "completion_tokens") == static_cast(kDecodeTokens); + if (!exact) { + std::cerr << "[✗] Benchmark run failed: " << json << "\n"; + cactus_destroy(model); + return false; + } + if (run == 0) { + std::cout << "├─ Warmup: prefill_tps=" << std::fixed << std::setprecision(2) + << json_number(json, "prefill_tps") + << " decode_tps=" << json_number(json, "decode_tps") << "\n"; + continue; + } + prefill_tps.push_back(json_number(json, "prefill_tps")); + decode_tps.push_back(json_number(json, "decode_tps")); + peak_ram.push_back(json_number(json, "peak_ram_usage_mb")); + } + cactus_destroy(model); + + std::cout << "\n[Mean of " << kMeasuredRuns << " runs]\n"; + print_mean("prefill_tps", "tok/s", prefill_tps); + print_mean("decode_tps", "tok/s", decode_tps); + print_mean("peak_ram", "MB", peak_ram); + std::cout << "└─ Status: PASSED ✓\n"; + return true; +} + +static bool vlm_benchmark() { + print_header("VLM BENCHMARK (image embed / decode)"); + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + std::string image_path = std::string(g_assets_path) + "/test_monkey.png"; + std::string audio_path = std::string(g_assets_path) + "/test.wav"; + const size_t embed_buffer_size = 16 * 1024 * 1024; + std::vector embeddings(embed_buffer_size / sizeof(float)); + + std::vector embed_ms; + for (int run = 0; run <= kMeasuredRuns; run++) { + size_t dim = 0; + Timer t; + int rc = cactus_image_embed(model, image_path.c_str(), embeddings.data(), + embed_buffer_size, &dim); + double elapsed = t.elapsed_ms(); + if (rc <= 0 || dim == 0) { + std::cerr << "[✗] Image embedding failed\n"; + cactus_destroy(model); + return false; + } + if (run == 0) { + std::cout << "├─ Warmup: image_embed=" << std::fixed << std::setprecision(2) + << elapsed << "ms (dim " << dim << ")\n"; + continue; + } + embed_ms.push_back(elapsed); + } + + std::vector audio_ms; + for (int run = 0; run <= kMeasuredRuns; run++) { + size_t dim = 0; + Timer t; + int rc = cactus_audio_embed(model, audio_path.c_str(), embeddings.data(), + embed_buffer_size, &dim); + double elapsed = t.elapsed_ms(); + if (rc <= 0 || dim == 0) { + std::cerr << "[✗] Audio embedding failed\n"; + cactus_destroy(model); + return false; + } + if (run == 0) { + std::cout << "├─ Warmup: audio_embed=" << std::fixed << std::setprecision(2) + << elapsed << "ms (dim " << dim << ")\n"; + continue; + } + audio_ms.push_back(elapsed); + } + + std::string messages = "[{\"role\": \"user\", " + "\"content\": \"Describe this image in detail.\", " + "\"images\": [\"" + image_path + "\"]}]"; + const char* options = R"({ + "max_tokens": 100, + "temperature": 0, + "top_k": 1, + "stop_sequences": ["<|im_end|>", ""], + "telemetry_enabled": false, + "auto_handoff": false + })"; + + std::vector decode_tps; + for (int run = 0; run <= kMeasuredRuns; run++) { + cactus_reset(model); + char response[8192]; + int rc = cactus_complete(model, messages.c_str(), response, sizeof(response), + options, nullptr, nullptr, nullptr, nullptr, 0); + std::string json(response); + double tps = json_number(json, "decode_tps"); + if (rc <= 0 || tps <= 0.0) { + std::cerr << "[✗] VLM completion failed: " << json << "\n"; + cactus_destroy(model); + return false; + } + if (run == 0) { + std::cout << "├─ Warmup: decode_tps=" << tps << " (decode_tokens=" + << static_cast(json_number(json, "decode_tokens")) << ")\n"; + continue; + } + decode_tps.push_back(tps); + } + cactus_destroy(model); + + std::cout << "\n[Mean of " << kMeasuredRuns << " runs]\n"; + print_mean("image_embed", "ms", embed_ms); + print_mean("audio_embed", "ms", audio_ms); + print_mean("decode_tps", "tok/s", decode_tps); + std::cout << "└─ Status: PASSED ✓\n"; + return true; +} + +static bool transcribe_benchmark() { + print_header("TRANSCRIBE BENCHMARK (audio embed / stt)"); + + cactus_model_t model = cactus_init(g_transcription_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize transcription model\n"; + return false; + } + + std::string audio_path = std::string(g_assets_path) + "/test.wav"; + const size_t embed_buffer_size = 1024 * 1024; + std::vector embeddings(embed_buffer_size / sizeof(float)); + + std::vector embed_ms; + for (int run = 0; run <= kMeasuredRuns; run++) { + size_t dim = 0; + Timer t; + int rc = cactus_audio_embed(model, audio_path.c_str(), embeddings.data(), + embed_buffer_size, &dim); + double elapsed = t.elapsed_ms(); + if (rc <= 0 || dim == 0) { + std::cerr << "[✗] Audio embedding failed\n"; + cactus_destroy(model); + return false; + } + if (run == 0) { + std::cout << "├─ Warmup: audio_embed=" << std::fixed << std::setprecision(2) + << elapsed << "ms (dim " << dim << ")\n"; + continue; + } + embed_ms.push_back(elapsed); + } + + const char* options = R"({"max_tokens": 500, "telemetry_enabled": false, "auto_handoff": false})"; + std::vector decode_tps, transcribe_ms; + for (int run = 0; run <= kMeasuredRuns; run++) { + char response[1 << 15] = {0}; + Timer t; + int rc = cactus_transcribe(model, audio_path.c_str(), nullptr, + response, sizeof(response), options, + nullptr, nullptr, nullptr, 0); + double elapsed = t.elapsed_ms(); + std::string json(response); + if (rc <= 0 || json_string(json, "response").length() <= 5) { + std::cerr << "[✗] Transcription failed: " << json << "\n"; + cactus_destroy(model); + return false; + } + if (run == 0) { + std::cout << "├─ Warmup: transcribe=" << std::fixed << std::setprecision(2) + << elapsed << "ms decode_tps=" << json_number(json, "decode_tps") << "\n"; + continue; + } + decode_tps.push_back(json_number(json, "decode_tps")); + transcribe_ms.push_back(elapsed); + } + cactus_destroy(model); + + std::cout << "\n[Mean of " << kMeasuredRuns << " runs]\n"; + print_mean("audio_embed", "ms", embed_ms); + print_mean("transcribe_total", "ms", transcribe_ms); + print_mean("decode_tps", "tok/s", decode_tps); + std::cout << "└─ Status: PASSED ✓\n"; + return true; +} + +int main() { + TestUtils::TestRunner runner("Benchmark"); + runner.run_test("llm_benchmark", llm_benchmark()); + runner.run_test("vlm_benchmark", vlm_benchmark()); + runner.run_test("transcribe_benchmark", transcribe_benchmark()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_curl.cpp b/cactus-engine/tests/test_curl.cpp new file mode 100644 index 000000000..e73296c1c --- /dev/null +++ b/cactus-engine/tests/test_curl.cpp @@ -0,0 +1,140 @@ +#include "test_utils.h" + +#include +#include +#include +#include +#include + +#include + +bool test_curl_version_info() { + curl_version_info_data* info = curl_version_info(CURLVERSION_NOW); + if (!info) return false; + if (!info->version || std::string(info->version).empty()) return false; + if (!info->host || std::string(info->host).empty()) return false; + return true; +} + +bool test_curl_easy_init() { + CURL* handle = curl_easy_init(); + if (!handle) return false; + curl_easy_cleanup(handle); + return true; +} + +bool test_curl_url_api() { + CURL* handle = curl_easy_init(); + if (!handle) return false; + + bool ok = true; + ok = ok && (curl_easy_setopt(handle, CURLOPT_URL, "https://example.com/api/v1/ping?x=1") == CURLE_OK); + ok = ok && (curl_easy_setopt(handle, CURLOPT_NOBODY, 1L) == CURLE_OK); + ok = ok && (curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, 200L) == CURLE_OK); + ok = ok && (curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, 200L) == CURLE_OK); + + curl_easy_cleanup(handle); + return ok; +} + +struct CurlHttpCheck { + bool pass = false; + std::string reason; + long http_code = 0; + std::string body_preview; +}; + +static size_t curl_write_to_string(char* ptr, size_t size, size_t nmemb, void* userdata) { + if (!userdata) return 0; + std::string* out = static_cast(userdata); + out->append(ptr, size * nmemb); + return size * nmemb; +} + +static void configure_curl_tls(CURL* handle) { + if (!handle) return; + curl_easy_setopt(handle, CURLOPT_SSL_VERIFYPEER, 1L); + curl_easy_setopt(handle, CURLOPT_SSL_VERIFYHOST, 2L); + const char* ca_bundle = std::getenv("CACTUS_CA_BUNDLE"); + if (ca_bundle && ca_bundle[0] != '\0') { + curl_easy_setopt(handle, CURLOPT_CAINFO, ca_bundle); + } +#if defined(__ANDROID__) + const char* ca_path = std::getenv("CACTUS_CA_PATH"); + if (ca_path && ca_path[0] != '\0') { + curl_easy_setopt(handle, CURLOPT_CAPATH, ca_path); + } else { + curl_easy_setopt(handle, CURLOPT_CAPATH, "/system/etc/security/cacerts"); + } +#endif +} + +CurlHttpCheck run_curl_http_request_check() { + CURL* handle = curl_easy_init(); + if (!handle) return {false, "curl_easy_init returned null", 0, ""}; + + bool opts_ok = true; + std::string body; + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_URL, "https://sha256.badssl.com/") == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_HTTPGET, 1L) == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 1L) == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_USERAGENT, "cactus-curl-test/1.0") == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_WRITEFUNCTION, curl_write_to_string) == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_WRITEDATA, &body) == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, 3000L) == CURLE_OK); + opts_ok = opts_ok && (curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT_MS, 2000L) == CURLE_OK); + configure_curl_tls(handle); + if (!opts_ok) { + curl_easy_cleanup(handle); + return {false, "failed to configure curl options", 0, ""}; + } + + CURLcode rc = curl_easy_perform(handle); + long http_code = 0; + curl_easy_getinfo(handle, CURLINFO_RESPONSE_CODE, &http_code); + curl_easy_cleanup(handle); + std::string body_preview = body.substr(0, std::min(body.size(), 400)); + + if (rc == CURLE_OK && http_code >= 200 && http_code < 400) { + if (body.find("badssl.com") != std::string::npos) { + return {true, "", http_code, body_preview}; + } + return {false, "https request succeeded but response body did not contain expected marker", + http_code, body_preview}; + } + + std::ostringstream oss; + oss << "request failed rc=" << static_cast(rc) + << " (" << curl_easy_strerror(rc) << "), http_code=" << http_code; + return {false, oss.str(), http_code, body_preview}; +} + +int main() { + TestUtils::TestRunner runner("Curl Tests"); + + CURLcode init_rc = curl_global_init(CURL_GLOBAL_DEFAULT); + if (init_rc != CURLE_OK) { + runner.run_test("global_init", false); + runner.print_summary(); + return 1; + } + + runner.run_test("version_info", test_curl_version_info()); + runner.run_test("easy_init", test_curl_easy_init()); + runner.run_test("url_api", test_curl_url_api()); + CurlHttpCheck http_check = run_curl_http_request_check(); + runner.run_test("http_request", http_check.pass); + if (http_check.http_code > 0) { + std::cout << " http_code: " << http_check.http_code << "\n"; + } + if (!http_check.body_preview.empty()) { + std::cout << " body_preview: " << http_check.body_preview << "\n"; + } + if (!http_check.pass && !http_check.reason.empty()) { + std::cout << " reason: " << http_check.reason << "\n"; + } + + curl_global_cleanup(); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_embed.cpp b/cactus-engine/tests/test_embed.cpp new file mode 100644 index 000000000..13078c473 --- /dev/null +++ b/cactus-engine/tests/test_embed.cpp @@ -0,0 +1,95 @@ +#include "test_utils.h" +#include +#include + +using namespace EngineTestUtils; + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); +static const char* g_transcription_model_path = std::getenv("CACTUS_TEST_TRANSCRIPTION_MODEL"); +static const char* g_assets_path = std::getenv("CACTUS_TEST_ASSETS"); + +bool test_embeddings() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ EMBEDDINGS TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) return false; + + const char* texts[] = {"My name is Henry Ndubuaku", "Your name is Henry Ndubuaku"}; + std::vector emb1(2048), emb2(2048); + size_t dim1 = 0, dim2 = 0; + + int rc1 = cactus_embed(model, texts[0], emb1.data(), emb1.size() * sizeof(float), &dim1, true); + int rc2 = cactus_embed(model, texts[1], emb2.data(), emb2.size() * sizeof(float), &dim2, true); + + if (rc1 < 0 || rc2 < 0 || dim1 == 0 || dim1 != dim2 || dim1 > emb1.size()) { + cactus_destroy(model); + return false; + } + + float similarity = 0; + for (size_t i = 0; i < dim1; ++i) { + similarity += emb1[i] * emb2[i]; + } + + std::cout << "\n[Results]\n" + << "├─ Embedding dim: " << dim1 << "\n" + << "└─ Similarity: " << std::fixed << std::setprecision(4) << similarity << std::endl; + + cactus_destroy(model); + return similarity > 0.5f; +} + +static bool test_image_embeddings() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ IMAGE EMBEDDING TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + + std::string image_path = std::string(g_assets_path) + "/test_monkey.png"; + const size_t buffer_size = 1024 * 1024 * 4; + std::vector embeddings(buffer_size / sizeof(float)); + size_t embedding_dim = 0; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) return false; + + int result = cactus_image_embed(model, image_path.c_str(), embeddings.data(), buffer_size, &embedding_dim); + + cactus_destroy(model); + + std::cout << "└─ Embedding dim: " << embedding_dim << std::endl; + + return result > 0 && embedding_dim > 0; +} + +static bool test_audio_embeddings() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ AUDIO EMBEDDING TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + + const size_t buffer_size = 1024 * 1024; + std::vector embeddings(buffer_size / sizeof(float)); + size_t embedding_dim = 0; + + cactus_model_t model = cactus_init(g_transcription_model_path, nullptr, false); + if (!model) return false; + + std::string audio_path = std::string(g_assets_path) + "/test.wav"; + int result = cactus_audio_embed(model, audio_path.c_str(), embeddings.data(), buffer_size, &embedding_dim); + + cactus_destroy(model); + + std::cout << "└─ Embedding dim: " << embedding_dim << std::endl; + + return result > 0 && embedding_dim > 0; +} + +int main() { + TestUtils::TestRunner runner("Embedding Tests"); + runner.run_test("embeddings", test_embeddings()); + runner.run_test("image_embeddings", test_image_embeddings()); + runner.run_test("audio_embeddings", test_audio_embeddings()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/tests/test_index.cpp b/cactus-engine/tests/test_index.cpp similarity index 71% rename from tests/test_index.cpp rename to cactus-engine/tests/test_index.cpp index d6459c9bb..27eebcbe5 100644 --- a/tests/test_index.cpp +++ b/cactus-engine/tests/test_index.cpp @@ -1,12 +1,10 @@ -#include "../cactus/ffi/cactus_ffi.h" +#include "../cactus_engine.h" #include "test_utils.h" #include #include #include #include #include -#include -#include #include #include #include @@ -267,83 +265,6 @@ bool test_constructor() { return true; } -void run_benchmarks(TestUtils::TestRunner& runner, uint32_t num_docs) { - IndexFixture f("bench", DIM); - f.init(); - - std::vector ids(num_docs); - std::vector docs(num_docs); - std::vector doc_ptrs(num_docs), meta_ptrs(num_docs); - std::vector> embs(num_docs); - std::vector emb_ptrs(num_docs); - - for (uint32_t i = 0; i < num_docs; ++i) { - ids[i] = i; - docs[i] = "doc" + std::to_string(i); - doc_ptrs[i] = docs[i].c_str(); - meta_ptrs[i] = "meta"; - embs[i] = random_embedding(); - emb_ptrs[i] = embs[i].data(); - } - - auto t0 = std::chrono::high_resolution_clock::now(); - for (size_t i = 0; i < num_docs; i += 1000) { - size_t n = std::min(size_t(1000), num_docs - i); - cactus_index_add(f.get_idx(), ids.data() + i, doc_ptrs.data() + i, meta_ptrs.data() + i, - emb_ptrs.data() + i, n, DIM); - } - auto t1 = std::chrono::high_resolution_clock::now(); - auto add_ms = std::chrono::duration_cast(t1 - t0).count() / 1000.0; - - t0 = std::chrono::high_resolution_clock::now(); - f.reopen(); - t1 = std::chrono::high_resolution_clock::now(); - auto load_ms = std::chrono::duration_cast(t1 - t0).count() / 1000.0; - - t0 = std::chrono::high_resolution_clock::now(); - f.query(random_embedding(), 10); - t1 = std::chrono::high_resolution_clock::now(); - auto query_ms = std::chrono::duration_cast(t1 - t0).count() / 1000.0; - - t0 = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < 1000; ++i) f.get(i * (num_docs / 1000)); - t1 = std::chrono::high_resolution_clock::now(); - auto get_ms = std::chrono::duration_cast(t1 - t0).count() / 1000.0; - - std::vector del_ids(1000); - for (int i = 0; i < 1000; ++i) del_ids[i] = i; - t0 = std::chrono::high_resolution_clock::now(); - f.del(del_ids); - t1 = std::chrono::high_resolution_clock::now(); - auto del_ms = std::chrono::duration_cast(t1 - t0).count() / 1000.0; - - t0 = std::chrono::high_resolution_clock::now(); - f.compact(); - t1 = std::chrono::high_resolution_clock::now(); - auto compact_ms = std::chrono::duration_cast(t1 - t0).count() / 1000.0; - - std::ostringstream ss; - ss << std::fixed << std::setprecision(2); - - ss.str(""); ss << add_ms << "ms"; - runner.log_performance("Add 100k docs", ss.str()); - - ss.str(""); ss << load_ms << "ms"; - runner.log_performance("Load 100k docs", ss.str()); - - ss.str(""); ss << query_ms << "ms"; - runner.log_performance("Query top-10", ss.str()); - - ss.str(""); ss << get_ms << "ms"; - runner.log_performance("Get 1k docs", ss.str()); - - ss.str(""); ss << del_ms << "ms"; - runner.log_performance("Delete 1k docs", ss.str()); - - ss.str(""); ss << compact_ms << "ms"; - runner.log_performance("Compact", ss.str()); -} - int main() { TestUtils::TestRunner runner("Index Tests"); @@ -356,8 +277,6 @@ int main() { runner.run_test("unicode", test_unicode()); runner.run_test("constructor", test_constructor()); - run_benchmarks(runner, 100000); - runner.print_summary(); return runner.all_passed() ? 0 : 1; diff --git a/cactus-engine/tests/test_json_parsing.cpp b/cactus-engine/tests/test_json_parsing.cpp new file mode 100644 index 000000000..7dc8fb95e --- /dev/null +++ b/cactus-engine/tests/test_json_parsing.cpp @@ -0,0 +1,60 @@ +#include "test_utils.h" +#include "../src/utils.h" +#include +#include + +bool test_role_before_content() { + std::vector images; + auto messages = cactus::ffi::parse_messages_json( + R"([{"role":"user","content":"hello"}])", images); + + bool passed = messages.size() == 1 && + messages[0].role == "user" && + messages[0].content == "hello"; + if (!passed) { + std::cerr << "[✗] role-before-content: role=[" << (messages.empty() ? "" : messages[0].role) + << "] content=[" << (messages.empty() ? "" : messages[0].content) << "]\n"; + } + return passed; +} + +bool test_content_before_role() { + std::vector images; + auto messages = cactus::ffi::parse_messages_json( + R"([{"content":"hello","role":"user"}])", images); + + bool passed = messages.size() == 1 && + messages[0].role == "user" && + messages[0].content == "hello"; + if (!passed) { + std::cerr << "[✗] content-before-role: role=[" << (messages.empty() ? "" : messages[0].role) + << "] content=[" << (messages.empty() ? "" : messages[0].content) << "]\n"; + } + return passed; +} + +bool test_multi_message_mixed_order() { + std::vector images; + auto messages = cactus::ffi::parse_messages_json( + R"([{"role":"user","content":"first"},{"content":"second","role":"assistant"}])", images); + + bool passed = messages.size() == 2 && + messages[0].role == "user" && messages[0].content == "first" && + messages[1].role == "assistant" && messages[1].content == "second"; + if (!passed) { + std::cerr << "[✗] multi-message-mixed-order: got " << messages.size() << " messages\n"; + for (const auto& m : messages) { + std::cerr << " role=[" << m.role << "] content=[" << m.content << "]\n"; + } + } + return passed; +} + +int main() { + TestUtils::TestRunner runner("JSON Message Parsing Tests"); + runner.run_test("role before content", test_role_before_content()); + runner.run_test("content before role", test_content_before_role()); + runner.run_test("multi-message mixed order", test_multi_message_mixed_order()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_kv_compress.cpp b/cactus-engine/tests/test_kv_compress.cpp new file mode 100644 index 000000000..07365db95 --- /dev/null +++ b/cactus-engine/tests/test_kv_compress.cpp @@ -0,0 +1,1010 @@ +#include "test_utils.h" +#include "../src/kv_compress.h" +#include "../src/engine.h" +#include "cactus_graph.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace TestUtils; +using namespace cactus::kvcompress; + +namespace { + +using Header = cactus::kvcompress::CacheHeader; +constexpr size_t kHeaderBytes = sizeof(Header); +constexpr size_t kGroupSize = 32; // KV_QUANT_GROUP_SIZE + +float f16_to_f32(uint16_t u) { __fp16 v; std::memcpy(&v, &u, 2); return static_cast(v); } +uint16_t f32_to_f16(float f) { __fp16 v = static_cast<__fp16>(f); uint16_t u; std::memcpy(&u, &v, 2); return u; } + +using EngineTestUtils::rope_reference; + +// `pre_rope[h][t][d]` is the planted pre-RoPE key, stored post-RoPE at position t. +std::vector make_fp16_cache(size_t n, size_t max_seq, size_t kv_heads, size_t head_dim, + double theta, std::vector>>& pre_rope, + std::vector>>& values) { + std::vector buf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + auto* h = reinterpret_cast(buf.data()); + h->current_seq_len = n; h->max_seq_len = max_seq; h->num_kv_heads = kv_heads; + h->head_dim = head_dim; h->sink_size = 4; + auto* rows = reinterpret_cast(buf.data() + kHeaderBytes); + pre_rope.assign(kv_heads, std::vector>(n)); + values.assign(kv_heads, std::vector>(n)); + for (size_t hh = 0; hh < kv_heads; ++hh) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim), v(head_dim); + for (size_t d = 0; d < head_dim; ++d) { + k[d] = std::sin(0.1f * (hh + 1) * (t + 1) + 0.3f * d) + 0.05f * d; + v[d] = std::cos(0.07f * (hh + 1) * (t + 2) + 0.2f * d); + } + pre_rope[hh][t] = k; values[hh][t] = v; + std::vector kr = rope_reference(k, (double)t, theta); + for (size_t d = 0; d < head_dim; ++d) + rows[(t * kv_heads + hh) * head_dim + d] = f32_to_f16(kr[d]); + } + return buf; +} + +std::vector make_fp16_value_cache(size_t n, size_t max_seq, size_t kv_heads, size_t head_dim, + const std::vector>>& values) { + std::vector buf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + auto* h = reinterpret_cast(buf.data()); + h->current_seq_len = n; h->max_seq_len = max_seq; h->num_kv_heads = kv_heads; + h->head_dim = head_dim; h->sink_size = 4; + auto* rows = reinterpret_cast(buf.data() + kHeaderBytes); + for (size_t hh = 0; hh < kv_heads; ++hh) + for (size_t t = 0; t < n; ++t) + for (size_t d = 0; d < head_dim; ++d) + rows[(t * kv_heads + hh) * head_dim + d] = f32_to_f16(values[hh][t][d]); + return buf; +} + +// Independent oracle quantizer, kept separate from production requant so tests don't self-validate. +void quantize_row(const std::vector& row, int8_t* dst, float* sc, + size_t head_dim, size_t group_size) { + size_t groups = (head_dim + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + size_t lo = g * group_size, hi = std::min(head_dim, lo + group_size); + float amax = 0.f; for (size_t d = lo; d < hi; ++d) amax = std::max(amax, std::fabs(row[d])); + float scale = amax / 127.f; if (scale < 1e-10f) scale = 1e-10f; sc[g] = scale; + float inv = 1.f / scale; + for (size_t d = lo; d < hi; ++d) { int32_t q = (int32_t)std::roundf(row[d] * inv); q = std::max(-128, std::min(127, q)); dst[d] = (int8_t)q; } + } +} + +} // namespace + +bool test_compact_fp16_cache() { + const size_t n = 60, max_seq = 128, kv_heads = 3, head_dim = 16; + const double theta = 1000000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + std::vector> kept = { + {0, 5, 10, 20, 30, 59}, {1, 4, 12, 25, 40, 58}, {0, 2, 15, 33, 50, 55}}; + size_t B = kept[0].size(); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + + bool ok = true; + for (size_t h = 0; h < kv_heads; ++h) { + for (size_t rank = 0; rank < B; ++rank) { + int abs_pos = kept[h][rank]; + std::vector expK = rope_reference(pre[h][abs_pos], (double)rank, theta); + const uint16_t* kdst = krows + (rank * kv_heads + h) * head_dim; + const uint16_t* vdst = vrows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + if (std::abs(f16_to_f32(kdst[d]) - expK[d]) > 5e-2f) ok = false; + if (std::abs(f16_to_f32(vdst[d]) - vals[h][abs_pos][d]) > 5e-2f) ok = false; + } + } + } + return ok && B == 6; +} + +bool test_dense_check_full_budget() { + // abs_budget == n keeps every index in order, so renumber delta is 0 and the buffer is unchanged. + const size_t n = 40, max_seq = 64, kv_heads = 2, head_dim = 16; + const double theta = 1000000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + std::vector kbuf0 = kbuf, vbuf0 = vbuf; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + Params p; p.recent_frac = 0.3f; p.sink = 4; p.abs_budget = (int)n; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, theta, p); + for (auto& k : kept) if (k.size() != n) return false; + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + + size_t live = n * kv_heads * head_dim * sizeof(uint16_t); + if (std::memcmp(kbuf.data() + kHeaderBytes, kbuf0.data() + kHeaderBytes, live) != 0) return false; + if (std::memcmp(vbuf.data() + kHeaderBytes, vbuf0.data() + kHeaderBytes, live) != 0) return false; + return true; +} + +bool test_rope_renumber_contiguous() { + // rope(stored_post_rope_at_abs, rank - abs) == rope(pre, rank). + const size_t head_dim = 16; const double theta = 1000000.0; + std::vector pre(head_dim); + for (size_t d = 0; d < head_dim; ++d) pre[d] = std::sin(0.4 * d) + 0.1 * d; + int abs_pos = 137; size_t rank = 5; + std::vector at_abs = rope_reference(pre, (double)abs_pos, theta); + std::vector at_rank = rope_reference(pre, (double)rank, theta); + std::vector delta = at_abs; + rope_rotate_row(delta.data(), head_dim, theta, (double)rank - (double)abs_pos); + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(delta[d] - at_rank[d]) > 1e-3f) return false; + return true; +} + +bool test_fp16_storage_round_trip() { + // Each survivor must read back as rope(pre, rank) within fp16 tolerance. + const size_t n = 64, max_seq = 128, kv_heads = 2, head_dim = 16; + const double theta = 1000000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + Params p; p.recent_frac = 0.3f; p.sink = 4; p.abs_budget = 16; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, theta, p); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + + bool ok = true; + for (size_t h = 0; h < kv_heads; ++h) { + for (size_t rank = 0; rank < kept[h].size(); ++rank) { + int abs_pos = kept[h][rank]; + std::vector expK = rope_reference(pre[h][abs_pos], (double)rank, theta); + const uint16_t* kdst = krows + (rank * kv_heads + h) * head_dim; + const uint16_t* vdst = vrows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + if (std::abs(f16_to_f32(kdst[d]) - expK[d]) > 5e-2f) ok = false; + if (std::abs(f16_to_f32(vdst[d]) - vals[h][abs_pos][d]) > 5e-2f) ok = false; + } + } + } + return ok; +} + +bool test_gemma_layer_selection() { + std::vector qwen(28, "full_attention"); + auto q = physical_compressible_layers(qwen, 28, 0); + if (q.size() != 28) return false; + for (size_t i = 0; i < 28; ++i) if (q[i] != i) return false; + + std::vector g; + for (int i = 0; i < 35; ++i) g.push_back(((i + 1) % 5 == 0) ? "global" : "sliding"); + auto gg = physical_compressible_layers(g, 35, 20); + std::vector expect = {4, 9}; + if (gg != expect) return false; + + // Pass-2 theta follows is_sliding_layer, NOT compressibility: layer 14 is a global KV-shared + // *source* excluded from compaction (not in {4,9}) yet must still re-rope with the global theta. + if (is_sliding_layer(g, 14)) return false; + if (is_sliding_layer(g, 4)) return false; + if (!is_sliding_layer(g, 3)) return false; + for (size_t i = 0; i < 28; ++i) if (is_sliding_layer(qwen, i)) return false; + return true; +} + +bool test_compact_int8_cache() { + const size_t n = 50, max_seq = 96, kv_heads = 2, head_dim = 32; + const double theta = 1000000.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + + std::vector>> pre(kv_heads, std::vector>(n)); + std::vector>> vals(kv_heads, std::vector>(n)); + std::vector k_i8(max_seq * int8_stride, 0), v_i8(max_seq * int8_stride, 0); + std::vector k_sc(max_seq * scale_stride, 0.f), v_sc(max_seq * scale_stride, 0.f); + + auto quant_row = [&](const std::vector& row, int8_t* dst, float* sc) { + quantize_row(row, dst, sc, head_dim, kGroupSize); + }; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim), v(head_dim); + for (size_t d = 0; d < head_dim; ++d) { k[d] = std::sin(0.1 * (h + 1) * (t + 1) + 0.2 * d); v[d] = std::cos(0.05 * (t + 1) + 0.1 * d); } + pre[h][t] = k; vals[h][t] = v; + std::vector kr = rope_reference(k, (double)t, theta); + quant_row(kr, k_i8.data() + t * int8_stride + h * head_dim, k_sc.data() + t * scale_stride + h * groups); + quant_row(v, v_i8.data() + t * int8_stride + h * head_dim, v_sc.data() + t * scale_stride + h * groups); + } + + std::vector> kept = {{0, 5, 11, 22, 33, 49}, {1, 4, 13, 26, 40, 48}}; + size_t B = kept[0].size(); + compact_int8(k_i8.data(), k_sc.data(), kv_heads, head_dim, kGroupSize, kept, theta, true); + compact_int8(v_i8.data(), v_sc.data(), kv_heads, head_dim, kGroupSize, kept, theta, false); + + bool ok = true; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t rank = 0; rank < B; ++rank) { + int abs_pos = kept[h][rank]; + std::vector expK = rope_reference(pre[h][abs_pos], (double)rank, theta); + const int8_t* kd = k_i8.data() + rank * int8_stride + h * head_dim; + const float* ks = k_sc.data() + rank * scale_stride + h * groups; + const int8_t* vd = v_i8.data() + rank * int8_stride + h * head_dim; + const float* vs = v_sc.data() + rank * scale_stride + h * groups; + for (size_t d = 0; d < head_dim; ++d) { + float dq = (float)kd[d] * ks[d / kGroupSize]; + if (std::abs(dq - expK[d]) > 0.05f) ok = false; + float dqv = (float)vd[d] * vs[d / kGroupSize]; + if (std::abs(dqv - vals[h][abs_pos][d]) > 0.05f) ok = false; + } + } + return ok; +} + +// Mirrors Model::maybe_roll_compact on an in-place FP16 K/V pair; returns the new bounded length B. +size_t roll_compact_once(uint16_t* krows, uint16_t* vrows, Header* khdr, Header* vhdr, + size_t kv_heads, size_t head_dim, double theta, int target_len) { + size_t n = khdr->current_seq_len; + Params p; p.recent_frac = 0.30f; p.sink = 4; + p.abs_budget = target_len; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, theta, p); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; + vhdr->current_seq_len = B; + return B; +} + +bool test_rolling_bounded_compaction() { + // Grow to trigger_len -> compact to target_len -> repeat, asserting the cache stays bounded and + // a planted distinctive mid-token survives every compaction. + const int trigger_len = 4096, target_len = 2048; + const size_t kv_heads = 2, head_dim = 16, max_seq = trigger_len + 8; + const double theta = 1000000.0; + + std::vector kbuf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + std::vector vbuf(kbuf.size(), 0); + auto* khdr = reinterpret_cast(kbuf.data()); + auto* vhdr = reinterpret_cast(vbuf.data()); + *khdr = Header{0, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + *vhdr = *khdr; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + // A centroid-opposite direction that KeyDiff must keep. + std::vector distinctive(head_dim); + for (size_t d = 0; d < head_dim; ++d) distinctive[d] = (d % 2 == 0 ? -1.0f : 1.0f) * 4.0f; + + // The token at `plant_abs` is the distinctive one; the rest cluster on a common centroid. + auto append = [&](size_t from, size_t to, long plant_abs) { + for (size_t t = from; t < to; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector k(head_dim); + if ((long)t == plant_abs) { + k = distinctive; + } else { + for (size_t d = 0; d < head_dim; ++d) + k[d] = std::sin(0.05 * (h + 1) + 0.3 * d) + 0.01f * std::sin(0.01 * t); + } + std::vector kr = rope_reference(k, (double)t, theta); + std::vector vr(head_dim); + for (size_t d = 0; d < head_dim; ++d) vr[d] = std::cos(0.07 * (t + 1) + 0.2 * d); + for (size_t d = 0; d < head_dim; ++d) { + krows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(kr[d]); + vrows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(vr[d]); + } + } + }; + + auto distinctive_survives = [&](size_t B) -> bool { + for (size_t h = 0; h < kv_heads; ++h) + for (size_t rank = 0; rank < B; ++rank) { + std::vector pre(head_dim); + const uint16_t* src = krows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) pre[d] = f16_to_f32(src[d]); + rope_rotate_row(pre.data(), head_dim, theta, -(double)rank); + double dot = 0, na = 0, nb = 0; + for (size_t d = 0; d < head_dim; ++d) { + dot += pre[d] * distinctive[d]; na += pre[d] * pre[d]; nb += distinctive[d] * distinctive[d]; + } + if (dot / (std::sqrt(na) * std::sqrt(nb) + 1e-9) > 0.99) return true; + } + return false; + }; + + bool ok = true; + const int cycles = 3; + for (int c = 0; c < cycles; ++c) { + size_t start = khdr->current_seq_len; + long plant_abs = (c == 0) ? (trigger_len / 2) : -1; + append(start, (size_t)trigger_len, plant_abs); + khdr->current_seq_len = trigger_len; + vhdr->current_seq_len = trigger_len; + + if (c == 0 && !distinctive_survives(trigger_len)) ok = false; + + size_t B = roll_compact_once(krows, vrows, khdr, vhdr, kv_heads, head_dim, theta, target_len); + if (B != (size_t)target_len) ok = false; + if (khdr->current_seq_len > (size_t)trigger_len) ok = false; + if (khdr->current_seq_len != (size_t)target_len) ok = false; + if (!distinctive_survives(B)) ok = false; + } + if (khdr->current_seq_len != (size_t)target_len) ok = false; + return ok; +} + +bool test_config_parse_rolling_fields() { + cactus::engine::Config def; + if (!def.kv_compress) return false; + if (def.kv_compress_trigger_len != 4096 || def.kv_compress_target_len != 2048) return false; + + std::string tmpl_str = (std::filesystem::temp_directory_path() / "cactus_kvcfg_XXXXXX").string(); + std::vector tmpl(tmpl_str.begin(), tmpl_str.end()); + tmpl.push_back('\0'); + int fd = mkstemp(tmpl.data()); + if (fd < 0) return false; + ::close(fd); + { + std::ofstream f(tmpl.data()); + // model_type=qwen avoids the Gemma4-only required-field validation. + f << "model_type=qwen\n" + << "kv_compress=true\n" + << "kv_compress_trigger_len=4096\n" + << "kv_compress_target_len=2048\n"; + } + cactus::engine::Config cfg; + bool parsed = cfg.from_json(tmpl.data()); + std::remove(tmpl.data()); + if (!parsed) return false; + return cfg.kv_compress && cfg.kv_compress_trigger_len == 4096 && cfg.kv_compress_target_len == 2048; +} + +bool test_trigger_zero_gates_rolling() { + // The maybe_roll_compact gate fires iff kv_compress && trigger_len > 0 && seq_len >= trigger_len. + const size_t n = 5000, kv_heads = 2, head_dim = 16, max_seq = 5008; + const double theta = 1000000.0; + + auto fresh_cache = [&]() { + std::vector buf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + *reinterpret_cast(buf.data()) = Header{n, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + auto* rows = reinterpret_cast(buf.data() + kHeaderBytes); + for (size_t i = 0; i < n * kv_heads * head_dim; ++i) rows[i] = f32_to_f16(0.01f * (i % 97)); + return buf; + }; + auto roll_if_gated = [&](cactus::engine::Config& cfg, std::vector& kbuf, std::vector& vbuf) { + auto* khdr = reinterpret_cast(kbuf.data()); + bool fire = cfg.kv_compress && cfg.kv_compress_trigger_len > 0 && + khdr->current_seq_len >= (size_t)cfg.kv_compress_trigger_len; + if (fire) + roll_compact_once(reinterpret_cast(kbuf.data() + kHeaderBytes), + reinterpret_cast(vbuf.data() + kHeaderBytes), khdr, + reinterpret_cast(vbuf.data()), kv_heads, head_dim, theta, + cfg.kv_compress_target_len); + }; + + // Gate closed: cache byte-identical, length unchanged. + cactus::engine::Config off; off.kv_compress_trigger_len = 0; off.kv_compress_target_len = 0; + std::vector k_off = fresh_cache(), v_off(k_off.size(), 0), k_off0 = k_off, v_off0 = v_off; + roll_if_gated(off, k_off, v_off); + if (reinterpret_cast(k_off.data())->current_seq_len != n) return false; + if (std::memcmp(k_off.data(), k_off0.data(), k_off.size()) != 0) return false; + if (std::memcmp(v_off.data(), v_off0.data(), v_off.size()) != 0) return false; + + // Gate open: compacts to target_len and rewrites bytes. + cactus::engine::Config on; on.kv_compress_trigger_len = 4096; on.kv_compress_target_len = 2048; + std::vector k_on = fresh_cache(), v_on(k_on.size(), 0), k_on0 = k_on; + roll_if_gated(on, k_on, v_on); + if (reinterpret_cast(k_on.data())->current_seq_len != 2048) return false; + if (std::memcmp(k_on.data(), k_on0.data(), k_on.size()) == 0) return false; + return true; +} + +bool test_degenerate_rolling_config_disabled() { + // validate_kv_compress requires 0 < target_len < trigger_len; bad configs reset both to 0. + auto disabled = [](int trig, int targ) { + cactus::engine::Config c; + c.kv_compress = true; + c.kv_compress_trigger_len = trig; + c.kv_compress_target_len = targ; + c.validate_kv_compress(); + return c.kv_compress_trigger_len == 0 && c.kv_compress_target_len == 0; + }; + if (!disabled(4096, 0)) return false; + if (!disabled(4096, 4096)) return false; + if (!disabled(2048, 4096)) return false; + if (!disabled(4096, -1)) return false; + cactus::engine::Config ok; + ok.kv_compress = true; + ok.kv_compress_trigger_len = 4096; + ok.kv_compress_target_len = 2048; + ok.validate_kv_compress(); + return ok.kv_compress_trigger_len == 4096 && ok.kv_compress_target_len == 2048; +} + +bool test_env_override_parse() { + cactus::engine::Config off; + if (off.parse_kv_compress_override(nullptr, nullptr)) return false; + if (!off.kv_compress || off.kv_compress_trigger_len != 4096 || off.kv_compress_target_len != 2048) return false; + + cactus::engine::Config both; + if (!both.parse_kv_compress_override("3000", "1000")) return false; + if (!both.kv_compress || both.kv_compress_trigger_len != 3000 || both.kv_compress_target_len != 1000) return false; + + // Setting only one keeps the other's default. + cactus::engine::Config one; + if (!one.parse_kv_compress_override(nullptr, "1500")) return false; + if (one.kv_compress_trigger_len != 4096 || one.kv_compress_target_len != 1500) return false; + + cactus::engine::Config disabled; + if (!disabled.parse_kv_compress_override("0", nullptr)) return false; + if (disabled.kv_compress) return false; + + cactus::engine::Config bad; + if (!bad.parse_kv_compress_override("2048", "4096")) return false; + if (bad.kv_compress_trigger_len != 0 || bad.kv_compress_target_len != 0) return false; + return true; +} + +bool test_rerope_recent_fp16_uniform() { + // After a delta shift each recent row reads as rope(pre, t+delta); sink rows stay byte-identical. + const size_t n = 40, max_seq = 64, kv_heads = 2, head_dim = 16, sink = 4; + const double theta = 10000.0, delta = -7.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + std::vector sink_before(krows, krows + sink * kv_heads * head_dim); + + rerope_recent_fp16(krows, kv_heads, head_dim, sink, n, theta, delta); + + if (std::memcmp(krows, sink_before.data(), sink_before.size() * sizeof(uint16_t)) != 0) return false; + for (size_t t = sink; t < n; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector exp = rope_reference(pre[h][t], (double)t + delta, theta); + const uint16_t* r = krows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(f16_to_f32(r[d]) - exp[d]) > 5e-2f) return false; + } + return true; +} + +namespace { +void build_int8_key_cache(size_t n, size_t kv_heads, size_t head_dim, double theta, + std::vector>>& pre, + std::vector& k_i8, std::vector& k_sc) { + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + pre.assign(kv_heads, std::vector>(n)); + k_i8.assign(n * int8_stride, 0); k_sc.assign(n * scale_stride, 0.f); + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim); + for (size_t d = 0; d < head_dim; ++d) k[d] = std::sin(0.1 * (h + 1) * (t + 1) + 0.2 * d); + pre[h][t] = k; + std::vector kr = rope_reference(k, (double)t, theta); + quantize_row(kr, k_i8.data() + t * int8_stride + h * head_dim, + k_sc.data() + t * scale_stride + h * groups, head_dim, kGroupSize); + } +} +} // namespace + +bool test_rerope_recent_int8() { + const size_t n = 40, kv_heads = 2, head_dim = 32, sink = 4; + const double theta = 10000.0, delta = -9.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + std::vector>> pre; + std::vector k_i8; std::vector k_sc; + build_int8_key_cache(n, kv_heads, head_dim, theta, pre, k_i8, k_sc); + + std::vector sink_i8(k_i8.begin(), k_i8.begin() + sink * int8_stride); + std::vector sink_sc(k_sc.begin(), k_sc.begin() + sink * scale_stride); + rerope_recent_int8(k_i8.data(), k_sc.data(), kv_heads, head_dim, kGroupSize, sink, n, theta, delta); + + if (std::memcmp(k_i8.data(), sink_i8.data(), sink_i8.size()) != 0) return false; + if (std::memcmp(k_sc.data(), sink_sc.data(), sink_sc.size() * sizeof(float)) != 0) return false; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = sink; t < n; ++t) { + std::vector exp = rope_reference(pre[h][t], (double)t + delta, theta); + const int8_t* kd = k_i8.data() + t * int8_stride + h * head_dim; + const float* ks = k_sc.data() + t * scale_stride + h * groups; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs((float)kd[d] * ks[d / kGroupSize] - exp[d]) > 0.05f) return false; + } + return true; +} + +bool test_rotate_int8_row_matches_inline() { + // rotate_int8_row must be bitwise-identical to the inline dequant -> rotate -> requant sequence. + const size_t head_dim = 32; const double theta = 10000.0, delta = -9.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + std::vector pre(head_dim); + for (size_t d = 0; d < head_dim; ++d) pre[d] = std::sin(0.3 * d) + 0.1 * d; + std::vector stored = rope_reference(pre, 21.0, theta); + std::vector a_i8(head_dim), b_i8(head_dim); + std::vector a_sc(groups), b_sc(groups); + quantize_row(stored, a_i8.data(), a_sc.data(), head_dim, kGroupSize); + b_i8 = a_i8; b_sc = a_sc; + + std::vector row(head_dim); + for (size_t d = 0; d < head_dim; ++d) row[d] = (float)a_i8[d] * a_sc[d / kGroupSize]; + rope_rotate_row(row.data(), head_dim, theta, delta); + quantize_row(row, a_i8.data(), a_sc.data(), head_dim, kGroupSize); + rotate_int8_row(b_i8.data(), b_sc.data(), head_dim, kGroupSize, theta, delta); + + return std::memcmp(a_i8.data(), b_i8.data(), head_dim) == 0 + && std::memcmp(a_sc.data(), b_sc.data(), groups * sizeof(float)) == 0; +} + +bool test_sliding_plus_global_rolling() { + // A global cache renumbers survivors to 0..B-1 while a sliding cache re-ropes its recent rows by + // -Δ (Δ = old_total - B), keeping both in the same compressed frame across 3 cycles. + const size_t kv_heads = 2, head_dim = 16, sink = 4; + const double gtheta = 1000000.0, ltheta = 10000.0; + + { + const size_t n = 60, max_seq = 96; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, gtheta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + Params p; p.recent_frac = 0.3f; p.sink = sink; p.abs_budget = 16; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, gtheta, p); + size_t B = kept[0].size(); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, gtheta); + for (auto& k : kept) if (k.size() != B) return false; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t rank = 0; rank < B; ++rank) { + std::vector exp = rope_reference(pre[h][kept[h][rank]], (double)rank, gtheta); + const uint16_t* r = krows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(f16_to_f32(r[d]) - exp[d]) > 5e-2f) return false; + } + } + + const size_t n = 16, max_seq = 32; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, ltheta, pre, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + std::vector sink_before(krows, krows + sink * kv_heads * head_dim); + + const size_t old_total = 30, B = 12; + const double Delta = (double)(old_total - B); + double accum = 0.0; + for (int cycle = 0; cycle < 3; ++cycle) { + rerope_recent_fp16(krows, kv_heads, head_dim, sink, n, ltheta, -Delta); + accum -= Delta; + } + if (std::memcmp(krows, sink_before.data(), sink_before.size() * sizeof(uint16_t)) != 0) return false; + for (size_t t = sink; t < n; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector exp = rope_reference(pre[h][t], (double)t + accum, ltheta); + const uint16_t* r = krows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(f16_to_f32(r[d]) - exp[d]) > 8e-2f) return false; // looser for 3-cycle drift + } + return true; +} + +bool test_rerope_local_theta_used() { + // The sliding re-rope must use the LOCAL theta; the wrong (global) theta misses the position. + const size_t n = 24, max_seq = 32, kv_heads = 1, head_dim = 16, sink = 4; + const double ltheta = 10000.0, gtheta = 1000000.0, delta = -5.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, ltheta, pre, vals); + std::vector kbuf2 = kbuf; + auto* a = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* b = reinterpret_cast(kbuf2.data() + kHeaderBytes); + + rerope_recent_fp16(a, kv_heads, head_dim, sink, n, ltheta, delta); // correct theta + rerope_recent_fp16(b, kv_heads, head_dim, sink, n, gtheta, delta); // wrong theta + + bool correct_ok = true, wrong_differs = false; + for (size_t t = sink; t < n; ++t) { + std::vector exp = rope_reference(pre[0][t], (double)t + delta, ltheta); + const uint16_t* ra = a + (t * kv_heads) * head_dim; + const uint16_t* rb = b + (t * kv_heads) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + if (std::abs(f16_to_f32(ra[d]) - exp[d]) > 5e-2f) correct_ok = false; + if (std::abs(f16_to_f32(rb[d]) - exp[d]) > 5e-2f) wrong_differs = true; + } + } + return correct_ok && wrong_differs; +} + +bool test_rerope_zero_delta_noop() { + const size_t n = 20, max_seq = 32, kv_heads = 2, head_dim = 16, sink = 4; + const double theta = 10000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + std::vector before = kbuf; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + rerope_recent_fp16(krows, kv_heads, head_dim, sink, n, theta, 0.0); // zero delta + rerope_recent_fp16(krows, kv_heads, head_dim, n, n, theta, -3.0); // empty range hi<=lo + return kbuf == before; +} + +// dequant_row's NEON path must match the scalar fallback, via int8 compact and the int8 keep-set fill. +bool test_dequant_simd_matches_scalar() { + const size_t n = 64, max_seq = 96, kv_heads = 4, head_dim = 128; + const double theta = 1000000.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + Params p; p.recent_frac = 0.30f; p.sink = 4; p.abs_budget = 32; + + std::vector k8(max_seq * int8_stride, 0), v8(max_seq * int8_stride, 0); + std::vector ks(max_seq * scale_stride, 0.f), vs(max_seq * scale_stride, 0.f); + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim), v(head_dim); + for (size_t d = 0; d < head_dim; ++d) { k[d] = std::sin(0.1 * (h + 1) * (t + 1) + 0.2 * d); v[d] = std::cos(0.05 * (t + 1) + 0.1 * d); } + std::vector kr = rope_reference(k, (double)t, theta); + quantize_row(kr, k8.data() + t * int8_stride + h * head_dim, ks.data() + t * scale_stride + h * groups, head_dim, kGroupSize); + quantize_row(v, v8.data() + t * int8_stride + h * head_dim, vs.data() + t * scale_stride + h * groups, head_dim, kGroupSize); + } + std::vector> kept = {{0, 5, 11, 22}, {1, 4, 13, 26}, {2, 8, 17, 30}, {3, 9, 19, 40}}; + + auto compact = [&](bool simd) { + kv_set_simd(simd); + std::vector a = k8; + std::vector b = ks; + compact_int8(a.data(), b.data(), kv_heads, head_dim, kGroupSize, kept, theta, true); + a.insert(a.end(), reinterpret_cast(b.data()), reinterpret_cast(b.data() + b.size())); + return a; + }; + auto keepsets = [&](bool simd) { + kv_set_simd(simd); + return keepsets_from_int8(k8.data(), ks.data(), n, kv_heads, head_dim, kGroupSize, theta, p); + }; + bool ok = compact(true) == compact(false) && keepsets(true) == keepsets(false); + kv_set_simd(true); + return ok; +} + +bool test_keepset_protect_union() { + // A low-score position KeyDiff would evict is force-kept via Params::protect. + const size_t n = 64; + std::vector scores(n); + for (size_t i = 0; i < n; ++i) scores[i] = 0.01f * static_cast(i); + Params p; p.sink = 4; p.recent_frac = 0.25f; p.abs_budget = 24; + const int special = 9; + + auto base = keepset_for_head(scores.data(), n, p); + if (std::find(base.begin(), base.end(), special) != base.end()) return false; + + p.protect = {special}; + auto prot = keepset_for_head(scores.data(), n, p); + if (std::find(prot.begin(), prot.end(), special) == prot.end()) return false; + if (prot.size() != 24) return false; + for (int i = 0; i < 4; ++i) + if (std::find(prot.begin(), prot.end(), i) == prot.end()) return false; + return true; +} + +bool test_rolling_protect_survives() { + // A token KeyDiff would evict survives each compaction cycle when re-protected. + const int trigger_len = 256, target_len = 128; + const size_t kv_heads = 2, head_dim = 16, max_seq = trigger_len + 8; + const double theta = 1000000.0; + + std::vector kbuf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + std::vector vbuf(kbuf.size(), 0); + auto* khdr = reinterpret_cast(kbuf.data()); + auto* vhdr = reinterpret_cast(vbuf.data()); + *khdr = Header{0, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + *vhdr = *khdr; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + std::vector special_v(head_dim); + for (size_t d = 0; d < head_dim; ++d) special_v[d] = (d % 2 ? -3.0f : 3.0f); + + auto append = [&](size_t from, size_t to, long special_abs) { + for (size_t t = from; t < to; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector k(head_dim); + for (size_t d = 0; d < head_dim; ++d) k[d] = std::sin(0.05 * (h + 1) + 0.3 * d); + std::vector kr = rope_reference(k, (double)t, theta); + std::vector vr(head_dim); + if ((long)t == special_abs) vr = special_v; + else for (size_t d = 0; d < head_dim; ++d) vr[d] = std::cos(0.07 * (t + 1) + 0.2 * d); + for (size_t d = 0; d < head_dim; ++d) { + krows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(kr[d]); + vrows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(vr[d]); + } + } + }; + + auto special_rank = [&]() -> long { + for (size_t r = 0; r < khdr->current_seq_len; ++r) { + const uint16_t* v = vrows + (r * kv_heads + 0) * head_dim; + double dot = 0, na = 0, nb = 0; + for (size_t d = 0; d < head_dim; ++d) { + float x = f16_to_f32(v[d]); dot += x * special_v[d]; na += x * x; nb += special_v[d] * special_v[d]; + } + if (dot / (std::sqrt(na) * std::sqrt(nb) + 1e-9) > 0.99) return (long)r; + } + return -1; + }; + + long special_pos = trigger_len / 2; + for (int c = 0; c < 3; ++c) { + size_t start = khdr->current_seq_len; + append(start, (size_t)trigger_len, c == 0 ? special_pos : -1); + khdr->current_seq_len = trigger_len; + vhdr->current_seq_len = trigger_len; + + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = target_len; + p.protect = {(int)special_pos}; + auto kept = keepsets_from_fp16(krows, (size_t)trigger_len, kv_heads, head_dim, theta, p); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + khdr->current_seq_len = kept[0].size(); + vhdr->current_seq_len = kept[0].size(); + + special_pos = special_rank(); + if (special_pos < 0) return false; + } + return true; +} + +bool test_cache_starts_small_and_grows() { + CactusGraph gb; + const size_t kv_heads = 2, head_dim = 64, ceiling = 100000, chunk = 100; + const size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + const size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + size_t new_kv = gb.input({chunk, kv_heads, head_dim}, Precision::FP16); + size_t state = gb.kv_cache_state(ceiling, kv_heads, head_dim, /*window*/0, /*sink*/4); + size_t append = gb.kv_cache_append(new_kv, state, /*window*/0, /*sink*/4); + gb.retain_outputs({static_cast(state), static_cast(append)}); + + bool ok = true; + std::vector expect; + auto val = [](size_t row, size_t c) { return std::sin(0.013 * row + 0.07 * c); }; + for (size_t step = 0; step < 3; ++step) { + std::vector in(chunk * int8_stride); + for (size_t r = 0; r < chunk; ++r) + for (size_t c = 0; c < int8_stride; ++c) { + float v = static_cast(val(step * chunk + r, c)); + in[r * int8_stride + c] = f32_to_f16(v); + expect.push_back(v); + } + gb.set_input(new_kv, in.data(), Precision::FP16); + gb.execute(); + if (step == 0 && reinterpret_cast(gb.get_output(state))->max_seq_len != 256) ok = false; + } + + const auto* hdr = reinterpret_cast(gb.get_output(state)); + if (hdr->max_seq_len != 512) ok = false; + if (hdr->current_seq_len != 300) ok = false; + + const auto* base = static_cast(gb.get_output(state)); + const int8_t* i8 = reinterpret_cast(base + kHeaderBytes); + const float* sc = reinterpret_cast(base + kHeaderBytes + 512 * int8_stride); + for (size_t r = 0; r < 300 && ok; ++r) + for (size_t h = 0; h < kv_heads; ++h) + for (size_t d = 0; d < head_dim; ++d) { + float dq = static_cast(i8[r * int8_stride + h * head_dim + d]) * + sc[r * scale_stride + h * groups + d / kGroupSize]; + if (std::abs(dq - expect[r * int8_stride + h * head_dim + d]) > 0.05f) ok = false; + } + return ok; +} + +bool test_sliding_layer_fixed_capacity() { + CactusGraph gb; + const size_t kv_heads = 2, head_dim = 64, ceiling = 100000, window = 512, sink = 4; + size_t state = gb.kv_cache_state(ceiling, kv_heads, head_dim, window, sink); + gb.retain_outputs({static_cast(state)}); + gb.execute(); + auto* hdr = reinterpret_cast(gb.get_output(state)); + return hdr->max_seq_len == window + sink + 1; +} + +bool test_special_rows_remap_through_kept() { + std::vector rows = {3, 7, 10}; + std::vector kept = {0, 3, 5, 7, 9, 10, 12}; + return remap_rows_through_kept(rows, kept) == std::vector{1, 3, 5}; +} + +bool test_per_head_protect_keeps_specials() { + const double theta = 1000000.0; + const size_t n = 64, kv_heads = 3, head_dim = 16; + std::vector>> pre, val; + auto kbuf = make_fp16_cache(n, n, kv_heads, head_dim, theta, pre, val); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = 24; + auto unrope = unrope_table(n, head_dim, theta); + std::vector> protect(kv_heads); + for (size_t h = 0; h < kv_heads; ++h) protect[h] = {static_cast(20 + h)}; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, unrope, p, protect); + if (kept.size() != kv_heads) return false; + for (size_t h = 0; h < kv_heads; ++h) + if (std::find(kept[h].begin(), kept[h].end(), static_cast(20 + h)) == kept[h].end()) return false; + return true; +} + +bool test_preflight_specials_exceed_budget() { + SpecialRowTracker tracked; + std::vector rows; + for (int r = 0; r < 50; ++r) rows.push_back(r); + tracked.add_appended(0, 2, rows); + SpecialRowTracker untracked; + return tracked.max_reserved(0, 4, {}) == 50 && tracked.max_reserved(0, 4, {}) > 40 && + untracked.max_reserved(0, 4, rows) == 50 && untracked.max_reserved(0, 4, rows) > 40; +} + +bool test_empty_protect_per_head_uses_params_fallback() { + const double theta = 1000000.0; + const size_t n = 64, kv_heads = 3, head_dim = 16; + std::vector>> pre, val; + auto kbuf = make_fp16_cache(n, n, kv_heads, head_dim, theta, pre, val); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = 24; p.protect = {25}; + auto unrope = unrope_table(n, head_dim, theta); + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, unrope, p, /*protect_per_head=*/{}); + if (kept.size() != kv_heads) return false; + for (size_t h = 0; h < kv_heads; ++h) + if (std::find(kept[h].begin(), kept[h].end(), 25) == kept[h].end()) return false; + return true; +} + +bool test_all_heads_keep_special_across_cycles() { + const int trigger_len = 256, target_len = 128; + const size_t kv_heads = 4, head_dim = 16, max_seq = trigger_len + 8; + const double theta = 1000000.0; + std::vector kbuf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + std::vector vbuf(kbuf.size(), 0); + auto* khdr = reinterpret_cast(kbuf.data()); + auto* vhdr = reinterpret_cast(vbuf.data()); + *khdr = Header{0, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + *vhdr = *khdr; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + auto special_val = [&](size_t h) { + std::vector v(head_dim); + for (size_t d = 0; d < head_dim; ++d) v[d] = (d % 2 ? -2.0f : 2.0f) * static_cast(h + 1); + return v; + }; + const long special_abs = trigger_len / 2; // a MIDDLE position (not sink, not recent) + auto append = [&](size_t from, size_t to) { + for (size_t t = from; t < to; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector k(head_dim); + // Special key is non-distinctive (KeyDiff would evict it); other rows differ per (h,t). + for (size_t d = 0; d < head_dim; ++d) + k[d] = (static_cast(t) == special_abs) ? 0.01f + : std::sin(0.11 * (h + 1) * (t + 1) + 0.3 * d) + 0.05 * d; + std::vector kr = rope_reference(k, static_cast(t), theta); + std::vector vr(head_dim); + if (static_cast(t) == special_abs) vr = special_val(h); + else for (size_t d = 0; d < head_dim; ++d) vr[d] = std::cos(0.07 * (t + 1) + 0.2 * d); + for (size_t d = 0; d < head_dim; ++d) { + krows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(kr[d]); + vrows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(vr[d]); + } + } + }; + auto special_rank = [&](size_t h) -> long { + std::vector sv = special_val(h); + for (size_t r = 0; r < khdr->current_seq_len; ++r) { + const uint16_t* v = vrows + (r * kv_heads + h) * head_dim; + double dot = 0, na = 0, nb = 0; + for (size_t d = 0; d < head_dim; ++d) { float x = f16_to_f32(v[d]); dot += x * sv[d]; na += x * x; nb += sv[d] * sv[d]; } + if (dot / (std::sqrt(na) * std::sqrt(nb) + 1e-9) > 0.99) return static_cast(r); + } + return -1; + }; + + SpecialRowTracker tracker; + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = target_len; + auto unrope = unrope_table(trigger_len, head_dim, theta); + for (int c = 0; c < 3; ++c) { + size_t start = khdr->current_seq_len; + append(start, static_cast(trigger_len)); + khdr->current_seq_len = trigger_len; vhdr->current_seq_len = trigger_len; + std::vector appended; + if (c == 0) appended.push_back(static_cast(special_abs)); + tracker.add_appended(0, kv_heads, appended); + auto kept = keepsets_from_fp16(krows, static_cast(trigger_len), kv_heads, head_dim, + unrope, p, tracker.protect(0)); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + tracker.remap(0, kept); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; vhdr->current_seq_len = B; + tracker.set_tracked_len(B); + } + long r0 = special_rank(0); + if (r0 < 0) return false; + bool diverged = false; + for (size_t h = 0; h < kv_heads; ++h) { + long r = special_rank(h); + if (r < 0) return false; + if (r != r0) diverged = true; // ranks differ across heads -> a real per-head scenario, not trivially aligned + } + return diverged; +} + +bool test_shrink_cache_buffer_preserves_rows() { + CactusGraph gb; + const size_t kv_heads = 2, head_dim = 64, ceiling = 100000, chunk = 300; + const size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + const size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + size_t new_kv = gb.input({chunk, kv_heads, head_dim}, Precision::FP16); + size_t state = gb.kv_cache_state(ceiling, kv_heads, head_dim, /*window*/0, /*sink*/4); + size_t append = gb.kv_cache_append(new_kv, state, /*window*/0, /*sink*/4); + gb.retain_outputs({static_cast(state), static_cast(append)}); + + std::vector expect; + auto val = [](size_t row, size_t c) { return std::sin(0.013 * row + 0.07 * c); }; + std::vector in(chunk * int8_stride); + for (size_t r = 0; r < chunk; ++r) + for (size_t c = 0; c < int8_stride; ++c) { + float v = static_cast(val(r, c)); + in[r * int8_stride + c] = f32_to_f16(v); + expect.push_back(v); + } + gb.set_input(new_kv, in.data(), Precision::FP16); + gb.execute(); // appends 300 -> grows 256->512, occupancy 300 + + auto* hdr = reinterpret_cast(gb.get_output(state)); + if (hdr->max_seq_len != 512 || hdr->current_seq_len != 300) return false; + + gb.shrink_cache_buffer(state, 256); // clamped up to occupancy 300 + hdr = reinterpret_cast(gb.get_output(state)); + if (hdr->max_seq_len != 300 || hdr->current_seq_len != 300) return false; + + const auto* base = static_cast(gb.get_output(state)); + const int8_t* i8 = reinterpret_cast(base + kHeaderBytes); + const float* sc = reinterpret_cast(base + kHeaderBytes + 300 * int8_stride); + bool ok = true; + for (size_t r = 0; r < 300 && ok; ++r) + for (size_t h = 0; h < kv_heads; ++h) + for (size_t d = 0; d < head_dim; ++d) { + float dq = static_cast(i8[r * int8_stride + h * head_dim + d]) * + sc[r * scale_stride + h * groups + d / kGroupSize]; + if (std::abs(dq - expect[r * int8_stride + h * head_dim + d]) > 0.05f) ok = false; + } + return ok; +} + +int main() { + TestUtils::TestRunner runner("KV Compress Free-Function Tests"); + runner.run_test("cache_starts_small_and_grows", test_cache_starts_small_and_grows()); + runner.run_test("sliding_layer_fixed_capacity", test_sliding_layer_fixed_capacity()); + runner.run_test("compact_fp16_cache", test_compact_fp16_cache()); + runner.run_test("dense_check_full_budget", test_dense_check_full_budget()); + runner.run_test("rope_renumber_contiguous", test_rope_renumber_contiguous()); + runner.run_test("fp16_storage_round_trip", test_fp16_storage_round_trip()); + runner.run_test("gemma_layer_selection", test_gemma_layer_selection()); + runner.run_test("compact_int8_cache", test_compact_int8_cache()); + runner.run_test("rerope_recent_fp16_uniform", test_rerope_recent_fp16_uniform()); + runner.run_test("rerope_recent_int8", test_rerope_recent_int8()); + runner.run_test("rotate_int8_row_matches_inline", test_rotate_int8_row_matches_inline()); + runner.run_test("sliding_plus_global_rolling", test_sliding_plus_global_rolling()); + runner.run_test("rerope_local_theta_used", test_rerope_local_theta_used()); + runner.run_test("rerope_zero_delta_noop", test_rerope_zero_delta_noop()); + runner.run_test("rolling_bounded_compaction", test_rolling_bounded_compaction()); + runner.run_test("config_parse_rolling_fields", test_config_parse_rolling_fields()); + runner.run_test("trigger_zero_gates_rolling", test_trigger_zero_gates_rolling()); + runner.run_test("degenerate_rolling_config_disabled", test_degenerate_rolling_config_disabled()); + runner.run_test("env_override_parse", test_env_override_parse()); + runner.run_test("dequant_simd_matches_scalar", test_dequant_simd_matches_scalar()); + runner.run_test("keepset_protect_union", test_keepset_protect_union()); + runner.run_test("rolling_protect_survives", test_rolling_protect_survives()); + runner.run_test("special_rows_remap_through_kept", test_special_rows_remap_through_kept()); + runner.run_test("per_head_protect_keeps_specials", test_per_head_protect_keeps_specials()); + runner.run_test("preflight_specials_exceed_budget", test_preflight_specials_exceed_budget()); + runner.run_test("empty_protect_per_head_uses_params_fallback", test_empty_protect_per_head_uses_params_fallback()); + runner.run_test("all_heads_keep_special_across_cycles", test_all_heads_keep_special_across_cycles()); + runner.run_test("shrink_cache_buffer_preserves_rows", test_shrink_cache_buffer_preserves_rows()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp new file mode 100644 index 000000000..8bcff7526 --- /dev/null +++ b/cactus-engine/tests/test_llm.cpp @@ -0,0 +1,678 @@ +#include "test_utils.h" +#include "../src/utils.h" +#include +#include +#include +#include + +#if __has_include() +#include +#define CACTUS_ENGINE_TEST_HAS_CURL 1 +#else +#define CACTUS_ENGINE_TEST_HAS_CURL 0 +#endif + +using namespace EngineTestUtils; +using namespace cactus::engine; +using cactus::ffi::partition_thinking_response; + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); + +static bool check_partition(const std::string& input, + const std::string& expected_thinking, + const std::string& expected_content) { + std::string thinking, content; + partition_thinking_response(input, thinking, content); + if (thinking != expected_thinking) { + std::cerr << " thinking: '" << thinking << "' != '" << expected_thinking << "'\n"; + return false; + } + if (content != expected_content) { + std::cerr << " content: '" << content << "' != '" << expected_content << "'\n"; + return false; + } + return true; +} + +static cactus_model_t load_gemma4_or_skip() { + if (!g_model_path) { std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; return nullptr; } + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { std::cout << " [WARN] Could not load model; skipping\n"; return nullptr; } + if (static_cast(model)->model->get_config().model_type != Config::ModelType::GEMMA4) { + std::cout << " [WARN] chosen model is not Gemma4; skipping\n"; + cactus_destroy(model); + return nullptr; + } + return model; +} + +static cactus_model_t load_dynamic_batch_model_or_skip() { + if (!g_model_path) { std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; return nullptr; } + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { std::cout << " [WARN] Could not load model; skipping\n"; return nullptr; } + if (!static_cast(model)->model->supports_dynamic_batch()) { + std::cout << " [WARN] model is not dynamic-batch capable (reconvert with `cactus convert`); skipping\n"; + cactus_destroy(model); + return nullptr; + } + return model; +} + +static const char* g_options = R"({ + "max_tokens": 256, + "stop_sequences": ["<|im_end|>", ""], + "telemetry_enabled": false + })"; + +bool test_streaming() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " STREAMING & FOLLOW-UP TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + const char* messages1 = R"([ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "My name is Henry Ndubuaku, how are you?"} + ])"; + + StreamingData data1; + data1.model = model; + char response1[4096]; + + std::cout << "\n[Turn 1]\n"; + std::cout << "User: My name is Henry Ndubuaku, how are you?\n"; + std::cout << "Assistant: "; + + int result1 = cactus_complete(model, messages1, response1, sizeof(response1), + g_options, nullptr, stream_callback, &data1, nullptr, 0); + + std::cout << "\n\n[Results - Turn 1]\n"; + Metrics metrics1; + metrics1.parse(response1); + metrics1.print_json(); + + bool success1 = result1 > 0 && data1.token_count > 0; + + if (!success1) { + std::cout << "└─ Status: FAILED ✗\n"; + cactus_destroy(model); + return false; + } + + std::string assistant_response; + for(const auto& token : data1.tokens) { + assistant_response += token; + } + + std::string messages2_str = R"([ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "My name is Henry Ndubuaku, how are you?"}, + {"role": "assistant", "content": ")" + escape_json(assistant_response) + R"("}, + {"role": "user", "content": "What is my name?"} + ])"; + + StreamingData data2; + data2.model = model; + char response2[4096]; + + std::cout << "\n[Turn 2]\n"; + std::cout << "User: What is my name?\n"; + std::cout << "Assistant: "; + + int result2 = cactus_complete(model, messages2_str.c_str(), response2, sizeof(response2), + g_options, nullptr, stream_callback, &data2, nullptr, 0); + + std::cout << "\n\n[Results - Turn 2]\n"; + Metrics metrics2; + metrics2.parse(response2); + metrics2.print_json(); + + bool success2 = result2 > 0 && data2.token_count > 0; + bool recalled_name = metrics2.response.find("Henry") != std::string::npos + || metrics2.response.find("henry") != std::string::npos; + std::cout << "├─ Turn 2 recalls name: " << (recalled_name ? "YES" : "NO") << "\n"; + + cactus_destroy(model); + return success1 && success2 && recalled_name; +} + +bool test_prefill_invalidated_on_message_change() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " PREFILL INVALIDATION (LLM) TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + const char* prefill_messages = R"([ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "Summarize the phrase 'brainrot' in one sentence."} + ])"; + + const char* complete_messages = R"([ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "Give one sentence about the power of the 'brainrot'."} + ])"; + + const char* options = R"({ + "max_tokens": 128, + "stop_sequences": ["<|im_end|>", ""], + "confidence_threshold": 0.0, + "telemetry_enabled": false + })"; + + char prefill_response[2048] = {0}; + int prefill_result = cactus_prefill(model, prefill_messages, prefill_response, sizeof(prefill_response), nullptr, nullptr, nullptr, 0); + PrefillMetrics prefill_metrics; + prefill_metrics.parse(prefill_response); + + char complete_response_warm[4096] = {0}; + int complete_result_warm = cactus_complete(model, complete_messages, complete_response_warm, sizeof(complete_response_warm), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics warm_metrics; + warm_metrics.parse(complete_response_warm); + + cactus_reset(model); + + char complete_response_cold[4096] = {0}; + int complete_result_cold = cactus_complete(model, complete_messages, complete_response_cold, sizeof(complete_response_cold), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics cold_metrics; + cold_metrics.parse(complete_response_cold); + + std::cout << "\n\n[Results]\n"; + std::cout << "├─ Prefill success: " << ((prefill_result > 0 && prefill_metrics.success) ? "YES" : "NO") << "\n" + << "├─ Complete(warm mismatched) prefill_tokens: " << warm_metrics.prefill_tokens << "\n" + << "├─ Complete(cold) prefill_tokens: " << cold_metrics.prefill_tokens << "\n"; + + bool all_success = prefill_result > 0 && prefill_metrics.success + && complete_result_warm > 0 && warm_metrics.success + && complete_result_cold > 0 && cold_metrics.success; + bool invalidated = warm_metrics.prefill_tokens == cold_metrics.prefill_tokens; + + std::cout << "├─ Calls successful: " << (all_success ? "YES" : "NO") << "\n" + << "└─ Mismatch invalidated cache: " << (invalidated ? "YES" : "NO") << std::endl; + + cactus_destroy(model); + return all_success && invalidated; +} + +bool test_prefill() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " PREFILL API TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + const char* prefill_messages = R"([ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "Explain what brainrot means in one short sentence."}, + {"role": "assistant", "content": "Brainrot is internet slang for obsessive, meme-heavy online fixation."} + ])"; + + const char* complete_messages = R"([ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "Explain what brainrot means in one short sentence."}, + {"role": "assistant", "content": "Brainrot is internet slang for obsessive, meme-heavy online fixation."}, + {"role": "user", "content": "Now rewrite that in six words."} + ])"; + + const char* options = R"({ + "max_tokens": 128, + "stop_sequences": ["<|im_end|>", ""], + "confidence_threshold": 0.0, + "telemetry_enabled": false + })"; + + char prefill_response[2048] = {0}; + int prefill_result = cactus_prefill(model, prefill_messages, prefill_response, sizeof(prefill_response), nullptr, nullptr, nullptr, 0); + PrefillMetrics prefill_metrics; + prefill_metrics.parse(prefill_response); + + char complete_response_warm[4096] = {0}; + int complete_result_warm = cactus_complete(model, complete_messages, complete_response_warm, sizeof(complete_response_warm), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics warm_metrics; + warm_metrics.parse(complete_response_warm); + + cactus_reset(model); + + char complete_response_cold[4096] = {0}; + int complete_result_cold = cactus_complete(model, complete_messages, complete_response_cold, sizeof(complete_response_cold), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics cold_metrics; + cold_metrics.parse(complete_response_cold); + + std::cout << "\n\n[Results]\n"; + std::cout << "├─ Prefill success: " << ((prefill_result > 0 && prefill_metrics.success) ? "YES" : "NO") << "\n" + << "├─ Prefill metrics: "; + prefill_metrics.print_line(); + std::cout << "\n"; + std::cout << "├─ Complete warm metrics:\n"; + warm_metrics.print_json(); + std::cout << "├─ Complete cold metrics:\n"; + cold_metrics.print_json(); + + bool all_success = prefill_result > 0 && prefill_metrics.success + && complete_result_warm > 0 && warm_metrics.success + && complete_result_cold > 0 && cold_metrics.success; + bool warm_prefilled_less = warm_metrics.prefill_tokens < cold_metrics.prefill_tokens; + + std::cout << "├─ Calls successful: " << (all_success ? "YES" : "NO") << "\n" + << "└─ Warm prefilled less than cold: " << (warm_prefilled_less ? "YES" : "NO") << std::endl; + + cactus_destroy(model); + return all_success && warm_prefilled_less; +} + +struct LengthTokenizer : public Tokenizer { + std::vector encode(const std::string& text) const override { + return {static_cast(text.size())}; + } + std::string decode(const std::vector&) const override { return std::string(); } + uint32_t get_vocab_size() const override { return 0; } + uint32_t get_unk_token() const override { return 0; } + uint32_t get_bos_token() const override { return 0; } + uint32_t get_eos_token() const override { return 0; } + bool load_vocabulary_with_config(const std::string&, const std::string&, const std::string&) override { return true; } +}; + +bool test_tool_constraint_clear_releases_bias() { + LengthTokenizer tok; + ToolCallConstrainer constrainer; + constrainer.init(Config::ModelType::GEMMA4, {{"get_weather", {"location"}, {}, {"location"}}}, &tok); + if (constrainer.get_bias().empty()) { + std::cerr << " expected bias after activating init\n"; + return false; + } + constrainer.reset(); + constrainer.init(Config::ModelType::GEMMA4, {}, &tok); + if (!constrainer.get_bias().empty()) { + std::cerr << " stale bias survived deactivating init\n"; + return false; + } + return true; +} + +bool test_tool_call() { + const char* messages = R"([ + {"role": "system", "content": "You are a helpful assistant that can use tools."}, + {"role": "user", "content": "What's the weather in San Francisco?"} + ])"; + + const char* tools = R"([{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City, State, Country"} + }, + "required": ["location"] + } + } + }])"; + + const char* options_with_force_tools = R"({ + "max_tokens": 256, + "stop_sequences": ["<|im_end|>", ""], + "force_tools": true + })"; + + return EngineTestUtils::run_test("TOOL CALL TEST", g_model_path, messages, options_with_force_tools, + [](int result, const StreamingData&, const std::string& response, const Metrics& m) { + bool has_function = response.find("\"function_calls\":[") != std::string::npos; + bool has_tool = has_function && response.find("get_weather") != std::string::npos; + std::cout << "├─ Function call: " << (has_function ? "YES" : "NO") << "\n" + << "├─ Correct tool: " << (has_tool ? "YES" : "NO") << "\n"; + m.print_json(); + return result > 0 && has_function && has_tool; + }, tools, -1, "What's the weather in San Francisco?"); +} + +bool test_multiple_tool_call_invocations() { + const char* messages = R"([ + {"role": "system", "content": "You are a helpful assistant that can use tools."}, + {"role": "user", "content": "Send a message to Bob and get the weather for San Francisco."} + ])"; + + const char* tools = R"([{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City, State, Country"} + }, + "required": ["location"] + } + } + }, { + "type": "function", + "function": { + "name": "send_message", + "description": "Send a message to a contact", + "parameters": { + "type": "object", + "properties": { + "recipient": {"type": "string", "description": "Name of the person to send the message to"}, + "message": {"type": "string", "description": "The message content to send"} + }, + "required": ["recipient", "message"] + } + } + }])"; + + const char* options_with_force_tools = R"({ + "max_tokens": 256, + "stop_sequences": ["<|im_end|>", ""], + "force_tools": true + })"; + + return EngineTestUtils::run_test("MULTIPLE TOOLS TEST", g_model_path, messages, options_with_force_tools, + [](int result, const StreamingData&, const std::string& response, const Metrics& m) { + bool has_function = response.find("\"function_calls\":[") != std::string::npos; + bool has_weather_tool = has_function + && (response.find("\"name\":\"get_weather\"") != std::string::npos + || response.find("\"name\": \"get_weather\"") != std::string::npos); + bool has_message_tool = has_function + && (response.find("\"name\":\"send_message\"") != std::string::npos + || response.find("\"name\": \"send_message\"") != std::string::npos); + std::cout << "├─ Function call: " << (has_function ? "YES" : "NO") << "\n" + << "├─ Correct tool: " << (has_weather_tool && has_message_tool ? "YES" : "NO") << "\n"; + m.print_json(); + return result > 0 && has_function && has_weather_tool && has_message_tool; + }, tools, -1, "Send a message to Bob and get the weather for San Francisco."); +} + +bool test_partition_thinking_response() { + return check_partition("<|channel>reasonanswer", "reason", "answer") + && check_partition("<|channel>\n reason\n\n\nanswer", "reason", "answer") + && check_partition("no tags here", "", "no tags here") + && check_partition("reasonanswer", "reason", "answer") + && check_partition("<|channel>thought1text1<|channel>thought2text2", + "thought1\nthought2", "text1text2"); +} + +bool test_prompt_gemma4_retains_thinking() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tok = handle->model->get_tokenizer(); + + std::vector msgs = { + {"user", "hello", "", {}, {}, 0, {}}, + {"assistant", "<|channel>internal reasoningvisible response", "", {}, {}, 0, {}}, + {"user", "followup", "", {}, {}, 0, {}} + }; + + std::string prompt = tok->format_chat_prompt(msgs, true, "", true); + cactus_destroy(model); + + bool has_visible = prompt.find("visible response") != std::string::npos; + bool has_reasoning = prompt.find("internal reasoning") != std::string::npos; + bool has_channel_tags = prompt.find("<|channel>") != std::string::npos + && prompt.find("") != std::string::npos; + + if (!has_visible) std::cerr << " missing visible response in prompt\n"; + if (!has_reasoning) std::cerr << " thinking content not retained in assistant turn\n"; + if (!has_channel_tags) std::cerr << " channel tags not retained in prompt\n"; + + return has_visible && has_reasoning && has_channel_tags; +} + +bool test_complete_gemma4_thinking_api_clean() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + const char* msgs = R"([{"role": "user", "content": "What is 2+2?"}])"; + char buf[8192]; + + int r = cactus_complete(model, msgs, buf, sizeof(buf), + R"({"max_tokens":128,"enable_thinking_if_supported":true,"telemetry_enabled":false})", + nullptr, nullptr, nullptr, nullptr, 0); + std::string resp(buf); + cactus_destroy(model); + + std::string response = EngineTestUtils::json_string(resp, "response"); + bool ok = r > 0 + && resp.find("\"success\":true") != std::string::npos + && response.find("<|channel>") == std::string::npos + && response.find("") == std::string::npos; + if (!ok) std::cerr << " thinking-enabled completion api not clean: " << resp << "\n"; + return ok; +} + +bool test_multiturn_thinking_persist() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tokenizer = handle->model->get_tokenizer(); + const char* options = R"({"max_tokens":128,"temperature":0,"top_k":1,"enable_thinking_if_supported":true,"telemetry_enabled":false,"auto_handoff":false})"; + const char* turn1_msgs = R"([{"role": "user", "content": "My name is Alice. Please remember this."}])"; + char buf[16384]; + + int r1 = cactus_complete(model, turn1_msgs, buf, sizeof(buf), options, nullptr, nullptr, nullptr, nullptr, 0); + if (r1 <= 0) { std::cerr << " Turn 1 failed\n"; cactus_destroy(model); return false; } + + std::vector processed_after_t1 = handle->processed_tokens; + std::string turn1_json(buf); + std::string context_response = EngineTestUtils::json_string(turn1_json, "context_response"); + if (context_response.empty()) { + std::cerr << " context_response missing from turn 1 result\n"; + cactus_destroy(model); + return false; + } + + std::vector t2_chat = { + {"user", "My name is Alice. Please remember this.", "", {}, {}, 0, {}}, + {"assistant", context_response, "", {}, {}, 0, {}}, + {"user", "What is my name?", "", {}, {}, 0, {}} + }; + std::vector t2_prompt_tokens = tokenizer->encode(tokenizer->format_chat_prompt(t2_chat, true, "", true)); + + bool prefix_ok = (t2_prompt_tokens.size() >= processed_after_t1.size()) && + std::equal(processed_after_t1.begin(), processed_after_t1.end(), t2_prompt_tokens.begin()); + std::cout << " Prefix match (cache reuse): " << (prefix_ok ? "YES" : "NO") << "\n"; + + std::string escaped = EngineTestUtils::escape_json(context_response); + std::string turn2_json = R"([{"role": "user", "content": "My name is Alice. Please remember this."},{"role": "assistant", "content": ")" + + escaped + R"("},{"role": "user", "content": "What is my name?"}])"; + + int r2 = cactus_complete(model, turn2_json.c_str(), buf, sizeof(buf), options, nullptr, nullptr, nullptr, nullptr, 0); + if (r2 <= 0) { std::cerr << " Turn 2 failed\n"; cactus_destroy(model); return false; } + + std::string turn2_result(buf); + std::string turn2_response = EngineTestUtils::json_string(turn2_result, "response"); + bool mentions_alice = turn2_response.find("Alice") != std::string::npos + || turn2_response.find("alice") != std::string::npos; + std::cout << " Turn 2 mentions Alice: " << (mentions_alice ? "YES" : "NO") << "\n"; + + cactus_destroy(model); + + if (!prefix_ok) std::cerr << " FAIL: re-rendered history did not prefix-match cache\n"; + if (!mentions_alice) std::cerr << " FAIL: turn 2 did not recall the name\n"; + return prefix_ok && mentions_alice; +} + +static std::string benchmark_tokens_json(cactus_model_t model, const std::vector& ids, size_t max_new) { + std::vector response(1 << 16, 0); + int rc = cactus_benchmark_tokens(model, ids.data(), ids.size(), max_new, response.data(), response.size()); + return rc < 0 ? std::string() : std::string(response.data()); +} + +static std::string completion_ids_field(const std::string& json) { + size_t start = json.find("\"completion_token_ids\":["); + if (start == std::string::npos) return std::string(); + size_t end = json.find(']', start); + return end == std::string::npos ? std::string() : json.substr(start, end - start + 1); +} + +static std::string first_completion_id(const std::string& json) { + std::string ids = completion_ids_field(json); + size_t open = ids.find('['); + if (open == std::string::npos) return std::string(); + size_t end = ids.find_first_of(",]", open + 1); + return end == std::string::npos ? std::string() : ids.substr(open + 1, end - open - 1); +} + +bool test_chunked_prefill_padding() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " CHUNKED PREFILL PADDING TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + if (!g_model_path) { std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; return true; } + + std::vector sentence; + { + cactus_model_t tok_model = cactus_init(g_model_path, nullptr, false); + if (!tok_model) { std::cerr << " [✗] model init failed\n"; return false; } + auto* tok = static_cast(tok_model)->model->get_tokenizer(); + sentence = tok->encode( + "The quick brown fox jumps over the lazy dog. Paris is the capital of France. " + "Water is composed of hydrogen and oxygen. The Earth orbits the Sun once every year. " + "Photosynthesis converts sunlight into chemical energy in plants."); + cactus_destroy(tok_model); + if (sentence.empty()) { std::cerr << " [✗] tokenizer produced no ids\n"; return false; } + } + + for (size_t prompt_len : {size_t(95), size_t(600)}) { + std::vector ids(prompt_len); + for (size_t i = 0; i < ids.size(); i++) ids[i] = sentence[i % sentence.size()]; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { std::cerr << " [✗] model init failed\n"; return false; } + std::string padded = benchmark_tokens_json(model, ids, 4); + std::string padded_again = benchmark_tokens_json(model, ids, 4); + cactus_destroy(model); + if (padded.empty() || padded.find("\"success\":true") == std::string::npos + || padded_again.empty() || padded_again.find("\"success\":true") == std::string::npos) { + std::cerr << " [✗] repeated chunked benchmark failed at len " << prompt_len << "\n"; + return false; + } + long tail_chunk = static_cast(EngineTestUtils::json_number(padded, "prefill_tail_chunk_tokens", -1)); + long tail_pads = static_cast(EngineTestUtils::json_number(padded, "prefill_tail_padding_tokens", -1)); + long scalar = static_cast(EngineTestUtils::json_number(padded, "prefill_scalar_tail_tokens", -1)); + std::cout << " len " << prompt_len << ": tail_chunk=" << tail_chunk + << " pads=" << tail_pads << " scalar=" << scalar << "\n"; + if (tail_chunk <= 0) { + std::cout << " [WARN] padded tail did not engage (no sliding caches?); skipping\n"; + continue; + } + if (tail_pads <= 0 || scalar > 1) { + std::cerr << " [✗] unexpected padding telemetry\n"; + return false; + } + if (completion_ids_field(padded).empty() + || completion_ids_field(padded) != completion_ids_field(padded_again)) { + std::cerr << " [✗] padded prefill is not deterministic\n"; + return false; + } + + setenv("CACTUS_DISABLE_PREFILL_TAIL_PAD", "1", 1); + model = cactus_init(g_model_path, nullptr, false); + std::string scalar_run = model ? benchmark_tokens_json(model, ids, 4) : std::string(); + if (model) cactus_destroy(model); + unsetenv("CACTUS_DISABLE_PREFILL_TAIL_PAD"); + if (scalar_run.empty() || scalar_run.find("\"success\":true") == std::string::npos) { + std::cerr << " [✗] kill-switch benchmark failed\n"; + return false; + } + long off_tail = static_cast(EngineTestUtils::json_number(scalar_run, "prefill_tail_chunk_tokens", -1)); + long off_scalar = static_cast(EngineTestUtils::json_number(scalar_run, "prefill_scalar_tail_tokens", -1)); + const size_t chunk_size = static_cast(tail_chunk + tail_pads + 1); + if (off_tail != 0 || off_scalar != static_cast(prompt_len % chunk_size)) { + std::cerr << " [✗] kill switch did not restore the scalar tail (tail=" << off_tail + << " scalar=" << off_scalar << ")\n"; + return false; + } + if (first_completion_id(padded).empty() + || first_completion_id(padded) != first_completion_id(scalar_run)) { + std::cerr << " [✗] padded prefill diverged from the scalar tail\n" + << " padded: " << completion_ids_field(padded) << "\n" + << " scalar: " << completion_ids_field(scalar_run) << "\n"; + return false; + } + } + return true; +} + +bool test_batch_distinct4_matches_single() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " BATCH DISTINCT-4 VS SINGLE TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + cactus_model_t model = load_dynamic_batch_model_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tok = handle->model->get_tokenizer(); + if (!handle->model->supports_dynamic_batch()) { + std::cout << " [WARN] bundle not dynamic-batch capable (reconvert with `cactus convert`); skipping\n"; + cactus_destroy(model); + return true; + } + + const std::vector questions = { + "In one short sentence, what is Paris famous for?", + "List three primary colors, comma separated.", + "Count from one to five.", + "Name two planets in our solar system." + }; + std::vector> prompts; + for (const auto& q : questions) { + ChatMessage msg; msg.role = "user"; msg.content = q; + prompts.push_back(tok->encode(tok->format_chat_prompt({msg}, true, "", false))); + } + const size_t M = 24; + const size_t K = prompts.size(); + + handle->model->set_decode_slots(K); + auto batched = handle->model->generate_batch(prompts, M, true); + if (batched.size() != K) { std::cerr << " batched returned " << batched.size() << " streams\n"; cactus_destroy(model); return false; } + + bool ok = true; + for (size_t i = 0; i < K; ++i) { + auto ref = handle->model->generate_batch({prompts[i]}, M, true); + bool match = !ref.empty() && ref[0] == batched[i]; + std::cout << " row " << i << " (\"" << questions[i] << "\") " << (match ? "MATCH" : "DIVERGED") + << ": " << tok->decode(batched[i]) << "\n"; + if (!match) ok = false; + } + cactus_destroy(model); + return ok; +} + +int main() { + TestUtils::TestRunner runner("LLM Tests"); + runner.run_test("streaming", test_streaming()); + runner.run_test("prefill", test_prefill()); + runner.run_test("prefill_invalidated_on_message_change", test_prefill_invalidated_on_message_change()); + runner.run_test("chunked_prefill_padding", test_chunked_prefill_padding()); + runner.run_test("tool_calls", test_tool_call()); + runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations()); + runner.run_test("tool_constraint_clear_releases_bias", test_tool_constraint_clear_releases_bias()); + runner.run_test("partition_thinking_response", test_partition_thinking_response()); + runner.run_test("prompt_retains_thinking", test_prompt_gemma4_retains_thinking()); + runner.run_test("complete_thinking_api_clean", test_complete_gemma4_thinking_api_clean()); + runner.run_test("multiturn_thinking_persist", test_multiturn_thinking_persist()); + runner.run_test("batch_distinct4_matches_single", test_batch_distinct4_matches_single()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_model_loading.cpp b/cactus-engine/tests/test_model_loading.cpp new file mode 100644 index 000000000..648ea7de0 --- /dev/null +++ b/cactus-engine/tests/test_model_loading.cpp @@ -0,0 +1,74 @@ +#include "test_utils.h" +#include +#include + +namespace fs = std::filesystem; + +static std::string make_temp_dir(const std::string& suffix) { + std::string dir = fs::temp_directory_path().string() + "/cactus_test_" + suffix; + fs::create_directories(dir); + return dir; +} + +static void write_file(const std::string& path, const std::string& content) { + std::ofstream(path, std::ios::binary) << content; +} + +static bool expect_init_fails(const std::string& path) { + cactus_model_t model = cactus_init(path.c_str(), nullptr, false); + if (model) { cactus_destroy(model); return false; } + return true; +} + +static const char* MINIMAL_CONFIG = R"({"model_type":"qwen","model_variant":"default","precision":"INT8","num_layers":2,"hidden_dim":64,"ffn_intermediate_dim":128,"attention_heads":2,"attention_kv_heads":2,"attention_head_dim":32,"vocab_size":100,"context_length":512})"; + +static bool test_missing_directory() { + return expect_init_fails("/nonexistent/path/to/model"); +} + +static bool test_missing_config() { + std::string dir = make_temp_dir("missing_config"); + write_file(dir + "/dummy.bin", "placeholder"); + bool ok = expect_init_fails(dir); + fs::remove_all(dir); + return ok; +} + +static bool test_corrupt_weights() { + std::string dir = make_temp_dir("corrupt_weights"); + write_file(dir + "/config.txt", MINIMAL_CONFIG); + write_file(dir + "/vocab.txt", "hello\nworld\n"); + write_file(dir + "/weights.bin", std::string("\xDE\xAD\xBE\xEF", 4) + std::string(124, '\xDE')); + bool ok = expect_init_fails(dir); + fs::remove_all(dir); + return ok; +} + +static bool test_empty_weight_file() { + std::string dir = make_temp_dir("empty_weights"); + write_file(dir + "/config.txt", MINIMAL_CONFIG); + write_file(dir + "/vocab.txt", "hello\nworld\n"); + write_file(dir + "/weights.bin", ""); + bool ok = expect_init_fails(dir); + fs::remove_all(dir); + return ok; +} + +static bool test_missing_vocab() { + std::string dir = make_temp_dir("missing_vocab"); + write_file(dir + "/config.txt", MINIMAL_CONFIG); + bool ok = expect_init_fails(dir); + fs::remove_all(dir); + return ok; +} + +int main() { + TestUtils::TestRunner runner("Model Loading Failure Tests"); + runner.run_test("missing_directory", test_missing_directory()); + runner.run_test("missing_config", test_missing_config()); + runner.run_test("corrupt_weights", test_corrupt_weights()); + runner.run_test("empty_weight_file", test_empty_weight_file()); + runner.run_test("missing_vocab", test_missing_vocab()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_needle.cpp b/cactus-engine/tests/test_needle.cpp new file mode 100644 index 000000000..54fb1728b --- /dev/null +++ b/cactus-engine/tests/test_needle.cpp @@ -0,0 +1,93 @@ +#include "test_utils.h" + +#include +#include +#include +#include +#include + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); + +static bool model_is_needle(const std::string& bundle_path) { + std::ifstream f(bundle_path + "/config.txt"); + if (!f) return false; + std::string line; + while (std::getline(f, line)) { + while (!line.empty() && (line.back() == '\r' || line.back() == ' ')) line.pop_back(); + if (line.rfind("model_type=", 0) == 0) { + return line.substr(11) == "needle"; + } + } + return false; +} + +bool test_needle_tool_call() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ NEEDLE TOOL CALL TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + if (!g_model_path) { + std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; + return true; + } + if (!model_is_needle(g_model_path)) { + std::cout << " [SKIP] model under test is not needle; skipping needle tool-call test\n"; + return true; + } + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << " [✗] Failed to initialize needle model at " << g_model_path << "\n"; + return false; + } + + const char* messages = R"([ + {"role": "user", "content": "What's the weather in San Francisco?"} + ])"; + + const char* tools = R"([{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name."} + }, + "required": ["location"] + } + } + }])"; + + const char* options = R"({"max_tokens": 64, "force_tools": true, "telemetry_enabled": false, "auto_handoff": false})"; + + char response[1 << 15] = {0}; + int rc = cactus_complete(model, messages, response, sizeof(response), options, tools, + nullptr, nullptr, nullptr, 0); + std::string r(response); + + std::cout << "├─ rc: " << rc << "\n"; + std::cout << "├─ Response: " << r << "\n"; + + const bool has_function_calls = + r.find("\"function_calls\":[{") != std::string::npos; + const bool calls_tool = r.find("get_weather") != std::string::npos; + const bool has_arg = r.find("San Francisco") != std::string::npos; + std::cout << "├─ function_calls present: " << (has_function_calls ? "YES" : "NO") << "\n"; + std::cout << "├─ Calls get_weather: " << (calls_tool ? "YES" : "NO") << "\n"; + std::cout << "├─ Location arg present: " << (has_arg ? "YES" : "NO") << "\n"; + + const bool ok = rc > 0 && has_function_calls && calls_tool && has_arg; + std::cout << "└─ Status: " << (ok ? "PASSED ✓" : "FAILED ✗") << "\n"; + + cactus_destroy(model); + return ok; +} + +int main() { + TestUtils::apply_backend(); + bool ok = true; + ok &= test_needle_tool_call(); + std::cout << "\n" << (ok ? "✓ needle tool-call test passed" : "✗ needle tool-call test failed") << "\n"; + return ok ? 0 : 1; +} diff --git a/cactus-engine/tests/test_rag.cpp b/cactus-engine/tests/test_rag.cpp new file mode 100644 index 000000000..bf9f75b95 --- /dev/null +++ b/cactus-engine/tests/test_rag.cpp @@ -0,0 +1,144 @@ +#include "test_utils.h" +#include +#include +#include +#include + +using namespace EngineTestUtils; + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); +static const char* g_assets_path = std::getenv("CACTUS_TEST_ASSETS"); + +static const char* g_options = R"({ + "max_tokens": 256, + "stop_sequences": ["<|im_end|>", ""], + "telemetry_enabled": false + })"; + +static bool check_correctness(const std::string& response) { + const std::vector required_keywords = { + "henry", "ndubuaku", "roman", "shemet", "founder" + }; + std::string lower = response; + std::transform(lower.begin(), lower.end(), lower.begin(), + [](unsigned char c) { return std::tolower(c); }); + for (const auto& kw : required_keywords) { + if (lower.find(kw) == std::string::npos) { + std::cerr << "[✗] Response missing expected keyword: \"" << kw << "\"\n"; + return false; + } + } + return true; +} + +bool test_rag() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ RAG TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + + std::string corpus_dir = std::string(g_assets_path) + "/rag_corpus"; + + std::cout << "├─ Corpus dir: " << corpus_dir << "\n"; + std::cout << "├─ Initializing model with RAG...\n"; + + cactus_model_t model = cactus_init(g_model_path, corpus_dir.c_str(), false); + if (!model) { + std::cerr << "[✗] Failed to initialize model with corpus dir\n"; + return false; + } + + auto print_chunks = [](cactus_model_t m, const char* query) -> bool { + char chunks_buf[16384]; + int rc = cactus_rag_query(m, query, chunks_buf, sizeof(chunks_buf), 5); + if (rc > 0) { + std::cout << "Retrieved chunks:\n"; + std::string chunks_str(chunks_buf); + size_t pos = 0; + int chunk_num = 1; + while ((pos = chunks_str.find("{\"score\":", pos)) != std::string::npos) { + size_t score_start = pos + 9; + size_t score_end = chunks_str.find(",", score_start); + std::string score = chunks_str.substr(score_start, score_end - score_start); + + size_t source_pos = chunks_str.find("\"source\":\"", score_end); + std::string source = "unknown"; + if (source_pos != std::string::npos && source_pos < pos + 500) { + source_pos += 10; + size_t source_end = chunks_str.find("\"", source_pos); + source = chunks_str.substr(source_pos, source_end - source_pos); + } + + size_t content_pos = chunks_str.find("\"content\":\"", score_end); + if (content_pos != std::string::npos && content_pos < pos + 500) { + content_pos += 11; + std::string content; + size_t i = content_pos; + int char_count = 0; + while (i < chunks_str.size() && char_count < 80) { + if (chunks_str[i] == '\\' && i + 1 < chunks_str.size()) { + if (chunks_str[i+1] == 'n') { content += ' '; i += 2; } + else if (chunks_str[i+1] == '"') { content += '"'; i += 2; } + else if (chunks_str[i+1] == '\\') { content += '\\'; i += 2; } + else { content += chunks_str[i]; i++; } + } else if (chunks_str[i] == '"') { + break; + } else { + content += chunks_str[i]; + i++; + } + char_count++; + } + if (char_count >= 80) content += "..."; + std::cout << " [" << chunk_num++ << "] " << source << " (score: " << score << ")\n" + << " \"" << content << "\"\n"; + } + pos = score_end; + } + return true; + } else { + std::cerr << "[✗] RAG retrieval failed (rc=" << rc << "): " << chunks_buf << "\n"; + return false; + } + }; + + const char* query = "Who are the founders of Cactus and what are their roles?"; + const char* messages = R"([ + {"role": "system", "content": "You are a helpful assistant. Answer based on the context provided."}, + {"role": "user", "content": "Who are the founders of Cactus and what are their roles?"} + ])"; + + StreamingData data; + data.model = model; + char response[4096]; + + std::cout << "\n[Query] " << query << "\n"; + + if (!print_chunks(model, query)) { + cactus_destroy(model); + return false; + } + + std::cout << "Response: "; + + int result = cactus_complete(model, messages, response, sizeof(response), + g_options, nullptr, stream_callback, &data, nullptr, 0); + + std::cout << "\n"; + + Metrics metrics; + metrics.parse(response); + metrics.print_json(); + + cactus_destroy(model); + + if (!(result > 0 && data.token_count > 0)) return false; + + return check_correctness(metrics.response); +} + +int main() { + TestUtils::TestRunner runner("RAG Tests"); + runner.run_test("rag_preprocessing", test_rag()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_stream_transcribe.cpp b/cactus-engine/tests/test_stream_transcribe.cpp new file mode 100644 index 000000000..70ddc1c63 --- /dev/null +++ b/cactus-engine/tests/test_stream_transcribe.cpp @@ -0,0 +1,153 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace EngineTestUtils; + +static const char* g_model = std::getenv("CACTUS_TEST_TRANSCRIPTION_MODEL"); +static const char* g_assets = std::getenv("CACTUS_TEST_ASSETS"); + +static std::vector words_of(const std::string& text) { + std::vector words; + std::istringstream iss(text); + std::string w; + while (iss >> w) { + std::string n; + for (char c : w) if (std::isalnum((unsigned char)c)) n += (char)std::tolower((unsigned char)c); + if (!n.empty()) words.push_back(n); + } + return words; +} + +static double recall(const std::vector& ref, const std::vector& hyp) { + if (ref.empty()) return 1.0; + size_t found = 0; + for (const auto& w : ref) if (std::find(hyp.begin(), hyp.end(), w) != hyp.end()) ++found; + return (double)found / (double)ref.size(); +} + +static std::vector load_wav(const std::string& path) { + FILE* f = fopen(path.c_str(), "rb"); + if (!f) return {}; + char tag[4]; + uint32_t sz; + if (fread(tag, 1, 4, f) != 4 || std::strncmp(tag, "RIFF", 4)) { fclose(f); return {}; } + fread(&sz, 4, 1, f); + if (fread(tag, 1, 4, f) != 4 || std::strncmp(tag, "WAVE", 4)) { fclose(f); return {}; } + uint16_t ch = 1, bits = 16, fmt = 0; + uint32_t rate = 16000; + bool have_fmt = false; + std::vector mono; + while (fread(tag, 1, 4, f) == 4 && fread(&sz, 4, 1, f) == 1) { + if (!std::strncmp(tag, "fmt ", 4)) { + uint16_t af, c, al, bps; + uint32_t sr, br; + fread(&af, 2, 1, f); fread(&c, 2, 1, f); fread(&sr, 4, 1, f); + fread(&br, 4, 1, f); fread(&al, 2, 1, f); fread(&bps, 2, 1, f); + fmt = af; ch = c ? c : 1; bits = bps; rate = sr; have_fmt = true; + if (sz > 16) fseek(f, (long)sz - 16, SEEK_CUR); + } else if (!std::strncmp(tag, "data", 4)) { + if (!have_fmt || fmt != 1 || bits != 16) { fclose(f); return {}; } + size_t frames = sz / (2u * ch); + std::vector raw(sz / 2); + fread(raw.data(), 2, raw.size(), f); + mono.resize(frames); + for (size_t i = 0; i < frames; ++i) { + int acc = 0; + for (uint16_t k = 0; k < ch; ++k) acc += raw[i * ch + k]; + mono[i] = (int16_t)(acc / ch); + } + break; + } else { + fseek(f, (long)(sz + (sz & 1)), SEEK_CUR); + } + } + fclose(f); + if (mono.empty() || rate == 16000) return mono; + double ratio = 16000.0 / (double)rate; + size_t on = (size_t)(mono.size() * ratio); + std::vector out(on); + for (size_t i = 0; i < on; ++i) { + double pos = (double)i / ratio; + size_t i0 = (size_t)pos; + double fr = pos - (double)i0; + int16_t a = mono[std::min(i0, mono.size() - 1)]; + int16_t b = mono[std::min(i0 + 1, mono.size() - 1)]; + out[i] = (int16_t)((double)a + ((double)b - (double)a) * fr); + } + return out; +} + +static std::string transcribe_oneshot(cactus_model_t model, const std::vector& pcm) { + char resp[1 << 16] = {0}; + int rc = cactus_transcribe(model, nullptr, nullptr, resp, sizeof(resp), + R"({"telemetry_enabled": false, "auto_handoff": false, "timestamps": true})", + nullptr, nullptr, + reinterpret_cast(pcm.data()), pcm.size() * sizeof(int16_t)); + return rc <= 0 ? "" : json_string(std::string(resp), "response"); +} + +static bool test_stream_matches_oneshot() { + std::cout << "\n=== STREAM vs ONE-SHOT ===\n"; + if (!g_model || !g_assets) { std::cout << "SKIP (model/assets not set)\n"; return true; } + cactus_model_t model = cactus_init(g_model, nullptr, false); + if (!model) { std::cerr << "[x] model init failed\n"; return false; } + + std::vector pcm = load_wav(std::string(g_assets) + "/test.wav"); + if (pcm.size() > 20 * 16000) pcm.resize(20 * 16000); + std::string oneshot_text = transcribe_oneshot(model, pcm); + auto golden = words_of(oneshot_text); + if (golden.size() < 4) { cactus_destroy(model); std::cout << " SKIP (test.wav has no usable speech)\n"; return true; } + + cactus_stream_transcribe_t stream = cactus_stream_transcribe_start(model, nullptr); + if (!stream) { cactus_destroy(model); std::cerr << "[x] stream start failed\n"; return false; } + std::string streamed; + std::vector resp(1 << 16); + auto append = [&]() { + std::string c = json_string(std::string(resp.data()), "confirmed"); + if (c.empty()) return; + if (!streamed.empty() && streamed.back() != ' ' && c.front() != ' ') streamed += ' '; + streamed += c; + }; + for (size_t off = 0; off < pcm.size(); off += 16000) { + size_t n = std::min(16000, pcm.size() - off); + resp[0] = '\0'; + if (cactus_stream_transcribe_process(stream, reinterpret_cast(pcm.data() + off), + n * sizeof(int16_t), resp.data(), resp.size()) < 0) { + cactus_stream_transcribe_stop(stream, nullptr, 0); + cactus_destroy(model); + std::cerr << "[x] process failed\n"; + return false; + } + append(); + } + resp[0] = '\0'; + cactus_stream_transcribe_stop(stream, resp.data(), resp.size()); + append(); + cactus_destroy(model); + + auto sw = words_of(streamed); + std::cout << " ONE-SHOT: " << oneshot_text << "\n"; + std::cout << " STREAM: " << streamed << "\n"; + double rc = recall(golden, sw), pr = recall(sw, golden); + bool ok = rc >= 0.85 && pr >= 0.85; + std::cout << " recall=" << std::fixed << std::setprecision(3) << rc << " precision=" << pr + << " " << (ok ? "OK" : "FAIL") << "\n Status: " << (ok ? "PASSED" : "FAILED") << "\n"; + return ok; +} + +int main() { + TestUtils::TestRunner runner("Stream Transcribe Tests"); + runner.run_test("stream_transcribe_matches_oneshot", test_stream_matches_oneshot()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_stt.cpp b/cactus-engine/tests/test_stt.cpp new file mode 100644 index 000000000..d66ef894d --- /dev/null +++ b/cactus-engine/tests/test_stt.cpp @@ -0,0 +1,72 @@ +#include "test_utils.h" +#include +#include +#include + +using namespace EngineTestUtils; + +static const char* g_transcription_model_path = std::getenv("CACTUS_TEST_TRANSCRIPTION_MODEL"); +static const char* g_assets_path = std::getenv("CACTUS_TEST_ASSETS"); + +static const char* g_options = R"({"max_tokens": 200, "telemetry_enabled": false, "auto_handoff": false})"; + +bool test_transcription() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ TRANSCRIPTION TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + + cactus_model_t model = cactus_init(g_transcription_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + std::string audio_path = std::string(g_assets_path) + "/test.wav"; + char response[1 << 15] = {0}; + + int rc = cactus_transcribe(model, audio_path.c_str(), nullptr, + response, sizeof(response), g_options, + nullptr, nullptr, nullptr, 0); + std::string file_transcript = json_string(std::string(response), "response"); + std::cout << "├─ File transcript: " << file_transcript << "\n"; + bool file_ok = rc > 0 && file_transcript.length() > 5; + if (!file_ok) std::cerr << "[✗] File transcription failed: " << response << "\n"; + + FILE* wav_file = fopen(audio_path.c_str(), "rb"); + if (!wav_file) { + std::cerr << "[✗] Failed to open audio file\n"; + cactus_destroy(model); + return false; + } + fseek(wav_file, 44, SEEK_SET); + std::vector pcm_data; + uint8_t buf[4096]; + size_t n; + while ((n = fread(buf, 1, sizeof(buf), wav_file)) > 0) { + pcm_data.insert(pcm_data.end(), buf, buf + n); + } + fclose(wav_file); + + char pcm_response[1 << 15] = {0}; + int pcm_rc = cactus_transcribe(model, nullptr, nullptr, + pcm_response, sizeof(pcm_response), g_options, + nullptr, nullptr, + pcm_data.data(), pcm_data.size()); + std::string pcm_transcript = json_string(std::string(pcm_response), "response"); + std::cout << "├─ PCM transcript: " << pcm_transcript << "\n"; + bool pcm_ok = pcm_rc > 0 && pcm_transcript.length() > 5; + if (!pcm_ok) std::cerr << "[✗] PCM transcription failed: " << pcm_response << "\n"; + + cactus_destroy(model); + + bool passed = file_ok && pcm_ok; + std::cout << "└─ Status: " << (passed ? "PASSED ✓" : "FAILED ✗") << "\n"; + return passed; +} + +int main() { + TestUtils::TestRunner runner("STT Tests"); + runner.run_test("transcription", test_transcription()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_telemetry.cpp b/cactus-engine/tests/test_telemetry.cpp new file mode 100644 index 000000000..8ba571e4e --- /dev/null +++ b/cactus-engine/tests/test_telemetry.cpp @@ -0,0 +1,206 @@ +#include "test_utils.h" +#include "../src/cloud.h" +#include "../src/telemetry.h" +#include "../../cactus-kernels/src/threading.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +std::string make_temp_dir(const char* prefix) { + const std::string base = std::filesystem::temp_directory_path().string(); + char pattern[512] = {0}; + std::snprintf(pattern, sizeof(pattern), "%s/%s_XXXXXX", base.c_str(), prefix); + const char* dir = mkdtemp(pattern); + return dir ? std::string(dir) : std::string(); +} + +int count_events(const std::string& file_path) { + std::ifstream in(file_path); + if (!in.is_open()) return 0; + int count = 0; + std::string line; + while (std::getline(in, line)) { + ++count; + } + return count; +} + +std::string random_project_id() { + static thread_local std::mt19937_64 rng(std::random_device{}()); + uint64_t a = rng(); + uint64_t b = rng(); + a = (a & 0xffffffffffff0fffULL) | 0x0000000000004000ULL; + b = (b & 0x3fffffffffffffffULL) | 0x8000000000000000ULL; + char buf[37]; + std::snprintf(buf, sizeof(buf), + "%08llx-%04llx-%04llx-%04llx-%012llx", + (unsigned long long)((a >> 32) & 0xffffffffULL), + (unsigned long long)((a >> 16) & 0xffffULL), + (unsigned long long)(a & 0xffffULL), + (unsigned long long)((b >> 48) & 0xffffULL), + (unsigned long long)(b & 0xffffffffffffULL)); + return std::string(buf); +} + +bool test_record_many_then_flush() { + const std::string cache_dir = make_temp_dir("cactus_record_many_flush"); + + cactus::telemetry::setTelemetryEnvironment("cpp-test", cache_dir.c_str()); + cactus::telemetry::setCloudDisabled(true); + cactus::telemetry::init("telemetry-test-project", "record-many", nullptr); + + constexpr int expected_event_count = 200; + for (int i = 0; i < expected_event_count; ++i) { + cactus::telemetry::recordCompletion("test-model", true, 10.0, 25.0, 30.0, 32, "ok"); + } + + cactus::telemetry::flush(); + + const std::string completion_log = cache_dir + "/completion.log"; + const int event_count = count_events(completion_log); + + cactus::telemetry::shutdown(); + rmdir(cache_dir.c_str()); + return event_count == expected_event_count; +} + +bool test_shutdown_then_reinit_then_record() { + const std::string cache_dir = make_temp_dir("cactus_shutdown_reinit"); + + cactus::telemetry::setTelemetryEnvironment("cpp-test", cache_dir.c_str()); + cactus::telemetry::setCloudDisabled(true); + cactus::telemetry::init("telemetry-test-project", "shutdown-reinit", nullptr); + + cactus::telemetry::recordCompletion("test-model", true, 5.0, 20.0, 18.0, 16, "before-shutdown"); + cactus::telemetry::flush(); + + const std::string completion_log = cache_dir + "/completion.log"; + const int lines_before_shutdown = count_events(completion_log); + + cactus::telemetry::shutdown(); + + cactus::telemetry::init("telemetry-test-project", "shutdown-reinit", nullptr); + cactus::telemetry::recordCompletion("test-model", true, 6.0, 21.0, 19.0, 17, "after-reinit"); + cactus::telemetry::flush(); + + const int lines_after_reinit = count_events(completion_log); + + cactus::telemetry::shutdown(); + rmdir(cache_dir.c_str()); + return lines_after_reinit > lines_before_shutdown; +} + +bool test_record_and_flush_race_no_deadlock() { + const std::string cache_dir = make_temp_dir("cactus_telemetry_race"); + + cactus::telemetry::setTelemetryEnvironment("cpp-test", cache_dir.c_str()); + cactus::telemetry::setCloudDisabled(true); + cactus::telemetry::init("telemetry-test-project", "race-test", nullptr); + + constexpr int producer_tasks = 8; + constexpr int records_per_task = 120; // keep total (960) under the 1000-event offline-cache cap so the no-loss check stays exact + constexpr int flusher_tasks = 4; + constexpr int flushes_per_task = 40; + constexpr int expected_event_count = producer_tasks * records_per_task; + + auto& pool = CactusThreading::get_thread_pool(); + std::vector> futures; + futures.reserve(producer_tasks + flusher_tasks); + + for (int i = 0; i < producer_tasks; ++i) { + futures.push_back(pool.enqueue([]() { + for (int j = 0; j < records_per_task; ++j) { + cactus::telemetry::recordCompletion("race-model", true, 1.0, 1.0, 2.0, 1, "race"); + } + })); + } + + for (int i = 0; i < flusher_tasks; ++i) { + futures.push_back(pool.enqueue([]() { + for (int j = 0; j < flushes_per_task; ++j) { + cactus::telemetry::flush(); + } + })); + } + + bool all_completed = true; + for (auto& future : futures) { + if (future.wait_for(std::chrono::seconds(10)) != std::future_status::ready) { + all_completed = false; + break; + } + future.get(); + } + + cactus::telemetry::flush(); + + const std::string completion_log = cache_dir + "/completion.log"; + const int event_count = count_events(completion_log); + + cactus::telemetry::shutdown(); + std::remove(completion_log.c_str()); + rmdir(cache_dir.c_str()); + + return all_completed && event_count == expected_event_count; +} + +enum class CloudTelemetryTestResult { + Passed, + Failed, + Skipped, +}; + +CloudTelemetryTestResult test_cloud_upload_record_then_flush() { + const char* no_cloud_tele = std::getenv("CACTUS_NO_CLOUD_TELE"); + if (no_cloud_tele && no_cloud_tele[0] != '\0') { + return CloudTelemetryTestResult::Skipped; + } + + const std::string telemetry_key = cactus::ffi::resolve_cloud_api_key(nullptr); + if (telemetry_key.empty()) { + return CloudTelemetryTestResult::Skipped; + } + + const std::string cache_dir = make_temp_dir("cactus_cloud_telemetry"); + const std::string completion_log = cache_dir + "/completion.log"; + const std::string project_id = random_project_id(); + + cactus::telemetry::setTelemetryEnvironment("cpp", cache_dir.c_str()); + cactus::telemetry::setCloudDisabled(false); + cactus::telemetry::init(project_id.c_str(), "cloud-upload", telemetry_key.c_str()); + + cactus::telemetry::recordCompletion("cloud-model", true, 7.0, 19.0, 15.0, 9, "cloud-ok"); + cactus::telemetry::flush(); + + const int cached_count = count_events(completion_log); + + cactus::telemetry::shutdown(); + std::remove(completion_log.c_str()); + rmdir(cache_dir.c_str()); + + return cached_count == 0 ? CloudTelemetryTestResult::Passed : CloudTelemetryTestResult::Failed; +} + +int main() { + TestUtils::TestRunner runner("Telemetry Tests"); + runner.run_test("Record many then Flush", test_record_many_then_flush()); + runner.run_test("Shutdown then Reinit", test_shutdown_then_reinit_then_record()); + runner.run_test("Record and Flush Race", test_record_and_flush_race_no_deadlock()); + CloudTelemetryTestResult cloud_result = test_cloud_upload_record_then_flush(); + if (cloud_result == CloudTelemetryTestResult::Skipped) { + runner.log_skip("Cloud record + Flush", "--enable-telemetry and resolved cloud key (env/cache) required"); + } else { + runner.run_test("Cloud record + Flush", cloud_result == CloudTelemetryTestResult::Passed); + } + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_utils.cpp b/cactus-engine/tests/test_utils.cpp new file mode 100644 index 000000000..76c3ddd97 --- /dev/null +++ b/cactus-engine/tests/test_utils.cpp @@ -0,0 +1,230 @@ +#include "test_utils.h" +#include +#include +#include + +namespace TestUtils { + +void apply_backend() { + static bool applied = false; + if (applied) return; + applied = true; + const char* b = std::getenv("CACTUS_TEST_BACKEND"); + if (!b || !*b) return; + if (cactus_set_backend(b) == 0) + std::cout << "Backend: " << (std::strcmp(b, "metal") == 0 ? "Metal GPU" : "CPU") << "\n"; + else + std::cout << "Backend '" << b << "' unavailable; using default\n"; +} + +TestRunner::TestRunner(const std::string& suite_name) + : suite_name_(suite_name), passed_count_(0), total_count_(0) { + apply_backend(); + std::cout << "\n╔══════════════════════════════════════════════════════════════════════════════════════╗\n" + << "║ Running " << std::left << std::setw(76) << suite_name_ << " ║\n" + << "╚══════════════════════════════════════════════════════════════════════════════════════╝\n"; +} + +void TestRunner::run_test(const std::string& test_name, bool result) { + total_count_++; + if (result) { + passed_count_++; + std::cout << "✓ PASS │ " << std::left << std::setw(25) << test_name << "\n"; + } else { + std::cout << "✗ FAIL │ " << std::left << std::setw(25) << test_name << "\n"; + } +} + +void TestRunner::log_skip(const std::string& test_name, const std::string& reason) { + std::cout << "⊘ SKIP │ " << std::left << std::setw(25) << test_name << " │ " << reason << "\n"; +} + +void TestRunner::print_summary() { + std::cout << "────────────────────────────────────────────────────────────────────────────────────────\n"; + if (all_passed()) + std::cout << "✓ All " << total_count_ << " tests passed!\n"; + else + std::cout << "✗ " << (total_count_ - passed_count_) << " of " << total_count_ << " tests failed!\n"; + std::cout << "\n"; +} + +bool TestRunner::all_passed() const { + return passed_count_ == total_count_; +} + +} + +namespace EngineTestUtils { + +Timer::Timer() : start(std::chrono::high_resolution_clock::now()) {} + +double Timer::elapsed_ms() const { + auto end = std::chrono::high_resolution_clock::now(); + return std::chrono::duration_cast(end - start).count() / 1000.0; +} + +double json_number(const std::string& json, const std::string& key, double def) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return def; + size_t start = pos + pattern.size(); + while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; + size_t end = start; + while (end < json.size() && std::string(",}] \t\n\r").find(json[end]) == std::string::npos) ++end; + try { return std::stod(json.substr(start, end - start)); } + catch (...) { return def; } +} + +std::string json_string(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return {}; + size_t start = pos + pattern.size(); + while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; + if (start >= json.size() || json[start] != '"') return {}; + ++start; + + std::string out; + out.reserve(128); + bool escaped = false; + for (size_t i = start; i < json.size(); ++i) { + char c = json[i]; + if (escaped) { + switch (c) { + case '"': out.push_back('"'); break; + case '\\': out.push_back('\\'); break; + case 'n': out.push_back('\n'); break; + case 'r': out.push_back('\r'); break; + case 't': out.push_back('\t'); break; + default: out.push_back(c); break; + } + escaped = false; + continue; + } + if (c == '\\') { escaped = true; continue; } + if (c == '"') return out; + out.push_back(c); + } + return {}; +} + +std::string escape_json(const std::string& s) { + std::ostringstream o; + for (auto c : s) { + switch (c) { + case '"': o << "\\\""; break; + case '\\': o << "\\\\"; break; + case '\n': o << "\\n"; break; + case '\r': o << "\\r"; break; + default: o << c; break; + } + } + return o.str(); +} + +void stream_callback(const char* token, uint32_t token_id, void* user_data) { + auto* data = static_cast(user_data); + data->tokens.push_back(token ? token : ""); + data->token_ids.push_back(token_id); + data->token_count++; + + std::string out = token ? token : ""; + for (char& c : out) if (c == '\n') c = ' '; + std::cout << out << std::flush; + + if (data->stop_at > 0 && data->token_count >= data->stop_at) { + std::cout << " [-> stopped]" << std::flush; + cactus_stop(data->model); + } +} + +static bool json_bool(const std::string& json, const std::string& key, bool def = false) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return def; + size_t start = pos + pattern.size(); + while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; + if (start + 4 <= json.size() && json.substr(start, 4) == "true") return true; + if (start + 5 <= json.size() && json.substr(start, 5) == "false") return false; + return def; +} + +static std::string json_array(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":"; + size_t pos = json.find(pattern); + if (pos == std::string::npos) return "[]"; + size_t start = pos + pattern.size(); + while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; + if (start >= json.size() || json[start] != '[') return "[]"; + int depth = 1; + size_t end = start + 1; + while (end < json.size() && depth > 0) { + if (json[end] == '[') depth++; + else if (json[end] == ']') depth--; + end++; + } + return json.substr(start, end - start); +} + +void Metrics::parse(const std::string& json) { + success = json_bool(json, "success", false); + error = json_string(json, "error"); + cloud_handoff = json_bool(json, "cloud_handoff", false); + response = json_string(json, "response"); + thinking = json_string(json, "thinking"); + function_calls = json_array(json, "function_calls"); + confidence = json_number(json, "confidence", -1.0); + ttft = json_number(json, "time_to_first_token_ms"); + total_ms = json_number(json, "total_time_ms"); + prefill_tps = json_number(json, "prefill_tps"); + decode_tps = json_number(json, "decode_tps"); + ram_mb = json_number(json, "ram_usage_mb"); + prefill_tokens = json_number(json, "prefill_tokens"); + completion_tokens = json_number(json, "decode_tokens"); + total_tokens = json_number(json, "total_tokens"); + segments = json_array(json, "segments"); +} + +void Metrics::print_json() const { + std::cout << " \"success\": " << (success ? "true" : "false") << ",\n" + << " \"error\": " << (error.empty() ? "null" : "\"" + error + "\"") << ",\n" + << " \"cloud_handoff\": " << (cloud_handoff ? "true" : "false") << ",\n" + << " \"response\": \"" << response << "\",\n" + << " \"thinking\": " << (thinking.empty() ? "null" : "\"" + thinking + "\"") << ",\n" + << " \"function_calls\": " << function_calls << ",\n" + << " \"segments\": " << segments << ",\n" + << " \"confidence\": " << std::fixed << std::setprecision(4) << confidence << ",\n" + << " \"time_to_first_token_ms\": " << std::setprecision(2) << ttft << ",\n" + << " \"total_time_ms\": " << total_ms << ",\n" + << " \"prefill_tps\": " << prefill_tps << ",\n" + << " \"decode_tps\": " << decode_tps << ",\n" + << " \"ram_usage_mb\": " << ram_mb << ",\n" + << " \"prefill_tokens\": " << std::setprecision(0) << prefill_tokens << ",\n" + << " \"decode_tokens\": " << completion_tokens << ",\n" + << " \"total_tokens\": " << total_tokens << std::endl; +} + +void PrefillMetrics::parse(const std::string& json) { + success = json_bool(json, "success", false); + error = json_string(json, "error"); + prefill_tokens = json_number(json, "prefill_tokens"); + prefill_tps = json_number(json, "prefill_tps"); + total_ms = json_number(json, "total_time_ms"); + ram_mb = json_number(json, "ram_usage_mb"); +} + +std::string PrefillMetrics::line() const { + std::ostringstream oss; + oss << std::fixed << std::setprecision(2) + << "prefill_tokens=" << std::setprecision(0) << prefill_tokens + << ", prefill_tps=" << std::setprecision(2) << prefill_tps + << ", total_time_ms=" << std::setprecision(2) << total_ms + << ", ram_usage_mb=" << std::setprecision(2) << ram_mb; + return oss.str(); +} + +void PrefillMetrics::print_line() const { + std::cout << line(); +} + +} diff --git a/cactus-engine/tests/test_utils.h b/cactus-engine/tests/test_utils.h new file mode 100644 index 000000000..b9fe96c64 --- /dev/null +++ b/cactus-engine/tests/test_utils.h @@ -0,0 +1,146 @@ +#ifndef TEST_UTILS_H +#define TEST_UTILS_H + +#include "../cactus_engine.h" +#include +#include +#include +#include +#include +#include +#include + +namespace TestUtils { + +class TestRunner { +public: + TestRunner(const std::string& suite_name); + void run_test(const std::string& test_name, bool result); + void log_skip(const std::string& test_name, const std::string& reason); + void print_summary(); + bool all_passed() const; + +private: + std::string suite_name_; + int passed_count_; + int total_count_; +}; + +void apply_backend(); + +} + +namespace EngineTestUtils { + +// Double-precision rotate_half RoPE oracle for kv_compress tests (mirrors norms_rope.cpp). +inline std::vector rope_reference(const std::vector& v, double pos, double theta) { + size_t d = v.size(), half = d / 2; + std::vector o(d); + for (size_t i = 0; i < half; ++i) { + double inv = std::pow(theta, -(2.0 * (double)i) / (double)d); + double a = pos * inv, c = std::cos(a), s = std::sin(a); + o[i] = (float)(v[i] * c - v[i + half] * s); + o[i + half] = (float)(v[i + half] * c + v[i] * s); + } + return o; +} + +struct Timer { + std::chrono::high_resolution_clock::time_point start; + Timer(); + double elapsed_ms() const; +}; + +double json_number(const std::string& json, const std::string& key, double def = 0.0); +std::string json_string(const std::string& json, const std::string& key); +std::string escape_json(const std::string& s); + +struct StreamingData { + std::vector tokens; + std::vector token_ids; + int token_count = 0; + cactus_model_t model = nullptr; + int stop_at = -1; +}; + +void stream_callback(const char* token, uint32_t token_id, void* user_data); + +struct Metrics { + bool success = false; + std::string error; + bool cloud_handoff = false; + std::string response; + std::string thinking; + std::string function_calls; + double confidence = -1.0; + double ttft = 0.0; + double total_ms = 0.0; + double prefill_tps = 0.0; + double decode_tps = 0.0; + double ram_mb = 0.0; + double prefill_tokens = 0.0; + double completion_tokens = 0.0; + double total_tokens = 0.0; + std::string segments; + + void parse(const std::string& json); + void print_json() const; +}; + +struct PrefillMetrics { + bool success = false; + std::string error; + double prefill_tokens = 0.0; + double prefill_tps = 0.0; + double total_ms = 0.0; + double ram_mb = 0.0; + + void parse(const std::string& json); + std::string line() const; + void print_line() const; +}; + +template +bool run_test(const char* title, const char* model_path, const char* messages, + const char* options, TestFunc test_logic, + const char* tools = nullptr, int stop_at = -1, + const char* user_prompt = nullptr) { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << std::string(" ") + title << "║\n" + << "╚══════════════════════════════════════════╝\n"; + + if (user_prompt) { + std::cout << "├─ User prompt: " << user_prompt << "\n"; + } + + cactus_model_t model = cactus_init(model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + StreamingData data; + data.model = model; + data.stop_at = stop_at; + + char response[4096]; + std::cout << "Response: "; + + int result = cactus_complete(model, messages, response, sizeof(response), + options, tools, stream_callback, &data, nullptr, 0); + + std::cout << "\n\n[Results]\n"; + + Metrics metrics; + metrics.parse(response); + + bool success = test_logic(result, data, response, metrics); + std::cout << "└─ Status: " << (success ? "PASSED ✓" : "FAILED ✗") << std::endl; + + cactus_destroy(model); + return success; +} + +} + +#endif diff --git a/cactus-engine/tests/test_vlm.cpp b/cactus-engine/tests/test_vlm.cpp new file mode 100644 index 000000000..85f0ae223 --- /dev/null +++ b/cactus-engine/tests/test_vlm.cpp @@ -0,0 +1,242 @@ +#include "test_utils.h" +#include +#include + +using namespace EngineTestUtils; + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); +static const char* g_assets_path = std::getenv("CACTUS_TEST_ASSETS"); + +static const char* g_options = R"({ + "max_tokens": 256, + "stop_sequences": ["<|im_end|>", ""], + "telemetry_enabled": false + })"; + +bool test_vlm_multiturn() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ VLM MULTI-TURN TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "Failed to initialize model for VLM multi-turn test" << std::endl; + return false; + } + + std::string img_path = std::string(g_assets_path) + "/test_monkey.png"; + + std::string messages1 = "[{\"role\": \"user\", " + "\"content\": \"Describe what is happening in this image in two sentences.\", " + "\"images\": [\"" + img_path + "\"]}]"; + + StreamingData stream_data1; + stream_data1.model = model; + char response1[4096]; + + std::cout << "\n[Turn 1]\n"; + std::cout << "User: Describe what is happening in this image in two sentences.\n"; + std::cout << "Assistant: "; + + int result1 = cactus_complete(model, messages1.c_str(), response1, sizeof(response1), + g_options, nullptr, stream_callback, &stream_data1, nullptr, 0); + + std::cout << "\n\n[Results - Turn 1]\n"; + Metrics metrics1; + metrics1.parse(response1); + metrics1.print_json(); + + bool success1 = result1 > 0 && stream_data1.token_count > 0; + + if (!success1) { + std::cout << "└─ Status: FAILED ✗\n"; + cactus_destroy(model); + return false; + } + + std::string assistant_response; + for (const auto& token : stream_data1.tokens) { + assistant_response += token; + } + + std::string messages2 = "[{\"role\": \"user\", " + "\"content\": \"Describe what is happening in this image in two sentences.\", " + "\"images\": [\"" + img_path + "\"]}, " + "{\"role\": \"assistant\", \"content\": \"" + escape_json(assistant_response) + "\"}, " + "{\"role\": \"user\", \"content\": \"Describe the image once again.\"}]"; + + StreamingData stream_data2; + stream_data2.model = model; + char response2[4096]; + + std::cout << "\n[Turn 2]\n"; + std::cout << "User: Describe the image once again.\n"; + std::cout << "Assistant: "; + + int result2 = cactus_complete(model, messages2.c_str(), response2, sizeof(response2), + g_options, nullptr, stream_callback, &stream_data2, nullptr, 0); + + std::cout << "\n\n[Results - Turn 2]\n"; + Metrics metrics2; + metrics2.parse(response2); + metrics2.print_json(); + + bool success2 = result2 > 0 && stream_data2.token_count > 0; + + if (!success2) { + std::cout << "└─ Status: FAILED ✗ (Follow-up message failed)\n"; + } + + cactus_destroy(model); + return success1 && success2; +} + +bool test_prefill_invalidated_on_message_change_vlm() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " PREFILL INVALIDATION (VLM) TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + + std::string prefill_img_path = std::string(g_assets_path) + "/test_monkey.png"; + std::string complete_img_path = std::string(g_assets_path) + "/test_thing.png"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + std::string prefill_messages = "[{\"role\": \"user\", \"content\": \"Describe this image in one short sentence.\", \"images\": [\"" + + prefill_img_path + "\"]}]"; + + std::string complete_messages = "[{\"role\": \"user\", \"content\": \"Describe this image in one short sentence.\", \"images\": [\"" + + complete_img_path + "\"]}]"; + + const char* options = R"({ + "max_tokens": 128, + "stop_sequences": ["<|im_end|>", ""], + "confidence_threshold": 0.0, + "telemetry_enabled": false + })"; + + char prefill_response[2048] = {0}; + int prefill_result = cactus_prefill(model, prefill_messages.c_str(), prefill_response, sizeof(prefill_response), nullptr, nullptr, nullptr, 0); + PrefillMetrics prefill_metrics; + prefill_metrics.parse(prefill_response); + + char complete_response_warm[4096] = {0}; + int complete_result_warm = cactus_complete(model, complete_messages.c_str(), complete_response_warm, sizeof(complete_response_warm), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics warm_metrics; + warm_metrics.parse(complete_response_warm); + + cactus_reset(model); + + char complete_response_cold[4096] = {0}; + int complete_result_cold = cactus_complete(model, complete_messages.c_str(), complete_response_cold, sizeof(complete_response_cold), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics cold_metrics; + cold_metrics.parse(complete_response_cold); + + std::cout << "\n\n[Results]\n"; + std::cout << "├─ Prefill success: " << ((prefill_result > 0 && prefill_metrics.success) ? "YES" : "NO") << "\n" + << "├─ Complete(warm mismatched) prefill_tokens: " << warm_metrics.prefill_tokens << "\n" + << "├─ Complete(cold) prefill_tokens: " << cold_metrics.prefill_tokens << "\n"; + + bool all_success = prefill_result > 0 && prefill_metrics.success + && complete_result_warm > 0 && warm_metrics.success + && complete_result_cold > 0 && cold_metrics.success; + bool invalidated = warm_metrics.prefill_tokens == cold_metrics.prefill_tokens; + + std::cout << "├─ Calls successful: " << (all_success ? "YES" : "NO") << "\n" + << "└─ Mismatch invalidated cache: " << (invalidated ? "YES" : "NO") << std::endl; + + cactus_destroy(model); + return all_success && invalidated; +} + +bool test_prefill_with_images() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " PREFILL WITH IMAGES TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + + std::string prefill_img_path = std::string(g_assets_path) + "/test_monkey.png"; + std::string extension_img_path = std::string(g_assets_path) + "/test_thing.png"; + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << "[✗] Failed to initialize model\n"; + return false; + } + + std::string prefill_messages = "[" + "{\"role\": \"system\", \"content\": \"You are a helpful assistant. Be concise.\"}," + "{\"role\": \"user\", \"content\": \"Describe this image in one short sentence.\", \"images\": [\"" + + prefill_img_path + "\"]}," + "{\"role\": \"assistant\", \"content\": \"This image shows a close-up of a monkey face.\"}" + "]"; + + std::string complete_messages = "[" + "{\"role\": \"system\", \"content\": \"You are a helpful assistant. Be concise.\"}," + "{\"role\": \"user\", \"content\": \"Describe this image in one short sentence.\", \"images\": [\"" + + prefill_img_path + "\"]}," + "{\"role\": \"assistant\", \"content\": \"This image shows a close-up of a monkey face.\"}," + "{\"role\": \"user\", \"content\": \"Now describe this second image in one short sentence.\", \"images\": [\"" + + extension_img_path + "\"]}" + "]"; + + const char* options = R"({ + "max_tokens": 128, + "stop_sequences": ["<|im_end|>", ""], + "confidence_threshold": 0.0, + "telemetry_enabled": false + })"; + + char prefill_response[2048] = {0}; + int prefill_result = cactus_prefill(model, prefill_messages.c_str(), prefill_response, sizeof(prefill_response), nullptr, nullptr, nullptr, 0); + PrefillMetrics prefill_metrics; + prefill_metrics.parse(prefill_response); + + char complete_response_warm[4096] = {0}; + int complete_result_warm = cactus_complete(model, complete_messages.c_str(), complete_response_warm, sizeof(complete_response_warm), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics warm_metrics; + warm_metrics.parse(complete_response_warm); + + cactus_reset(model); + + char complete_response_cold[4096] = {0}; + int complete_result_cold = cactus_complete(model, complete_messages.c_str(), complete_response_cold, sizeof(complete_response_cold), + options, nullptr, nullptr, nullptr, nullptr, 0); + Metrics cold_metrics; + cold_metrics.parse(complete_response_cold); + + std::cout << "\n\n[Results]\n"; + std::cout << "├─ Prefill success: " << ((prefill_result > 0 && prefill_metrics.success) ? "YES" : "NO") << "\n" + << "├─ Prefill metrics: "; + prefill_metrics.print_line(); + std::cout << "\n"; + std::cout << "├─ Complete warm metrics:\n"; + warm_metrics.print_json(); + std::cout << "├─ Complete cold metrics:\n"; + cold_metrics.print_json(); + + bool all_success = prefill_result > 0 && prefill_metrics.success + && complete_result_warm > 0 && warm_metrics.success + && complete_result_cold > 0 && cold_metrics.success; + bool warm_prefilled_less = warm_metrics.prefill_tokens < cold_metrics.prefill_tokens; + + std::cout << "├─ Calls successful: " << (all_success ? "YES" : "NO") << "\n" + << "└─ Warm prefilled less than cold: " << (warm_prefilled_less ? "YES" : "NO") << std::endl; + + cactus_destroy(model); + return all_success && warm_prefilled_less; +} + +int main() { + TestUtils::TestRunner runner("VLM Tests"); + runner.run_test("prefill_with_images", test_prefill_with_images()); + runner.run_test("prefill_invalidated_on_message_change", test_prefill_invalidated_on_message_change_vlm()); + runner.run_test("vlm_multiturn", test_vlm_multiturn()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/transcribe.cpp b/cactus-engine/tests/transcribe.cpp new file mode 100644 index 000000000..52aac854e --- /dev/null +++ b/cactus-engine/tests/transcribe.cpp @@ -0,0 +1,406 @@ +#include "../cactus_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef HAVE_SDL2 +#include +#include +#include +#endif + +constexpr size_t RESPONSE_BUFFER_SIZE = 65536; + +namespace ansi { +constexpr const char* reset = "\033[0m"; +constexpr const char* bold = "\033[1m"; +constexpr const char* dim = "\033[2m"; +constexpr const char* green = "\033[32m"; +constexpr const char* yellow = "\033[33m"; +constexpr const char* cyan = "\033[36m"; +} + +static bool stdout_is_terminal() { + return isatty(STDOUT_FILENO) != 0 && std::getenv("NO_COLOR") == nullptr; +} +static const bool g_color = stdout_is_terminal(); + +static std::string colored(const std::string& text, const char* color) { + return g_color ? color + text + ansi::reset : text; +} + +static int terminal_width() { + struct winsize w; + if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &w) == -1 || w.ws_col == 0) return 80; + return w.ws_col; +} + +static void print_rule() { + std::string rule(terminal_width(), '-'); + std::cout << colored(rule, ansi::dim) << "\n"; +} + +static void print_banner() { + if (!g_color) { + std::cout << "Cactus Transcription\n\n"; + return; + } + std::cout << ansi::bold << ansi::green + << " ██████╗ █████╗ ██████╗████████╗██╗ ██╗███████╗\n" + << "██╔════╝██╔══██╗██╔════╝╚══██╔══╝██║ ██║██╔════╝\n" + << "██║ ███████║██║ ██║ ██║ ██║███████╗\n" + << "██║ ██╔══██║██║ ██║ ██║ ██║╚════██║\n" + << "╚██████╗██║ ██║╚██████╗ ██║ ╚██████╔╝███████║\n" + << " ╚═════╝╚═╝ ╚═╝ ╚═════╝ ╚═╝ ╚═════╝ ╚══════╝\n" + << ansi::reset << ansi::dim << " Transcription" << ansi::reset << "\n\n"; +} + +static std::string json_string_value(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":\""; + size_t start = json.find(pattern); + if (start == std::string::npos) return ""; + start += pattern.length(); + size_t end = start; + while (end < json.length() && json[end] != '"') { + if (json[end] == '\\' && end + 1 < json.length()) end++; + end++; + } + return json.substr(start, end - start); +} + +static double json_number_value(const std::string& json, const std::string& key) { + std::string pattern = "\"" + key + "\":"; + size_t start = json.find(pattern); + if (start == std::string::npos) return 0.0; + start += pattern.length(); + while (start < json.length() && std::isspace((unsigned char)json[start])) start++; + return std::strtod(json.c_str() + start, nullptr); +} + +static std::string whisper_prompt(const std::string& model_path, const std::string& language) { + std::string p = model_path; + std::transform(p.begin(), p.end(), p.begin(), [](unsigned char c) { return (char)std::tolower(c); }); + if (p.find("whisper") != std::string::npos) + return "<|startoftranscript|><|" + language + "|><|transcribe|>"; + return ""; +} + +static void print_token(const char* token, uint32_t, void*) { + std::cout << token << std::flush; +} + +static int transcribe_file(cactus_model_t model, const std::string& audio_path, + const std::string& model_path, const std::string& language) { + std::string prompt = whisper_prompt(model_path, language); + std::vector response_buffer(RESPONSE_BUFFER_SIZE, 0); + std::string options = "{\"max_tokens\":500,\"auto_handoff\":false,\"language\":\"" + language + "\"}"; + + auto start = std::chrono::steady_clock::now(); + int rc = cactus_transcribe( + model, audio_path.c_str(), prompt.empty() ? nullptr : prompt.c_str(), + response_buffer.data(), response_buffer.size(), + options.c_str(), print_token, nullptr, nullptr, 0); + double total_s = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + if (rc < 0) { + std::cerr << "\nTranscription failed: " << cactus_get_last_error() << "\n"; + return -1; + } + + std::string json(response_buffer.data()); + double model_ms = json_number_value(json, "total_time_ms"); + std::ostringstream stats; + stats << std::fixed << std::setprecision(2) << "[processed: " << total_s << "s"; + if (model_ms > 0.0) stats << " | model: " << model_ms / 1000.0 << "s"; + stats << "]"; + std::cout << "\n\n" << colored(stats.str(), ansi::dim) << "\n"; + return 0; +} + +#ifdef HAVE_SDL2 + +constexpr int TARGET_SAMPLE_RATE = 16000; +constexpr int AUDIO_BUFFER_MS = 100; +constexpr int PROCESS_INTERVAL_MS = 250; + +struct AudioState { + std::mutex mutex; + std::vector buffer; + std::atomic recording{false}; + int actual_sample_rate{TARGET_SAMPLE_RATE}; +}; +static AudioState g_audio; + +static std::vector resample_to_16k(const std::vector& input, int source_rate) { + if (source_rate == TARGET_SAMPLE_RATE) return input; + size_t num_input = input.size() / 2; + const int16_t* in = reinterpret_cast(input.data()); + double ratio = static_cast(TARGET_SAMPLE_RATE) / source_rate; + size_t num_output = static_cast(num_input * ratio); + std::vector out(num_output); + for (size_t i = 0; i < num_output; i++) { + double src = i / ratio; + size_t i0 = static_cast(src); + size_t i1 = std::min(i0 + 1, num_input - 1); + double frac = src - i0; + out[i] = static_cast(in[i0] * (1.0 - frac) + in[i1] * frac); + } + std::vector result(num_output * 2); + std::memcpy(result.data(), out.data(), result.size()); + return result; +} + +static void audio_callback(void*, Uint8* stream, int len) { + if (!g_audio.recording) return; + std::lock_guard lock(g_audio.mutex); + g_audio.buffer.insert(g_audio.buffer.end(), stream, stream + len); +} + +static size_t wrap_index(const std::string& s, size_t limit) { + size_t len = 0; + bool in_esc = false; + for (size_t i = 0; i < s.length(); ++i) { + char c = s[i]; + if (c == '\033') { in_esc = true; continue; } + if (in_esc) { if (c == 'm') in_esc = false; continue; } + if ((c & 0xC0) != 0x80) len++; + if (len >= limit && c == ' ') return i; + } + return std::string::npos; +} + +static int run_live_transcription(cactus_model_t model, const std::string& language) { + if (SDL_Init(SDL_INIT_AUDIO) < 0) { + std::cerr << "Failed to init SDL audio: " << SDL_GetError() << "\n"; + return 1; + } + + SDL_AudioSpec want, have; + SDL_zero(want); + want.freq = TARGET_SAMPLE_RATE; + want.format = AUDIO_S16LSB; + want.channels = 1; + want.samples = (TARGET_SAMPLE_RATE * AUDIO_BUFFER_MS) / 1000; + want.callback = audio_callback; + SDL_AudioDeviceID device = SDL_OpenAudioDevice(nullptr, 1, &want, &have, SDL_AUDIO_ALLOW_FREQUENCY_CHANGE); + if (device == 0) { + std::cerr << "Failed to open microphone: " << SDL_GetError() << "\n"; + SDL_Quit(); + return 1; + } + g_audio.actual_sample_rate = have.freq; + + std::string options = "{\"language\":\"" + language + "\"}"; + cactus_stream_transcribe_t stream = cactus_stream_transcribe_start(model, options.c_str()); + if (!stream) { + std::cerr << "Failed to start streaming (live mode needs a Whisper or Parakeet TDT model): " + << cactus_get_last_error() << "\n"; + SDL_CloseAudioDevice(device); + SDL_Quit(); + return 1; + } + + std::cout << colored("Listening…", ansi::yellow) << " press " << colored("Enter", ansi::cyan) << " to stop\n"; + print_rule(); + std::cout << "\n"; + + g_audio.buffer.clear(); + g_audio.recording = true; + SDL_PauseAudioDevice(device, 0); + + std::atomic should_stop{false}; + std::thread input_thread([&should_stop]() { + std::string line; + std::getline(std::cin, line); + should_stop = true; + }); + + std::string confirmed_text; + std::string current_line; + int pending_lines = 0; + std::string stats; + std::vector response_buffer(RESPONSE_BUFFER_SIZE, 0); + auto last_process = std::chrono::steady_clock::now(); + + auto step = [&](std::vector chunk, bool render) { + if (chunk.empty()) return; + std::vector pcm = resample_to_16k(chunk, g_audio.actual_sample_rate); + response_buffer[0] = '\0'; + auto t0 = std::chrono::steady_clock::now(); + if (cactus_stream_transcribe_process(stream, pcm.data(), pcm.size(), + response_buffer.data(), response_buffer.size()) < 0) return; + + std::string json(response_buffer.data()); + std::string confirmed = json_string_value(json, "confirmed"); + std::string pending = json_string_value(json, "pending"); + confirmed_text += confirmed; + if (!render) return; + + double latency_ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + if (!confirmed.empty() || !pending.empty()) { + std::ostringstream st; + st << "[latency: " << (int)latency_ms << "ms]"; + stats = colored(st.str(), ansi::dim); + } + + int limit = (int)(terminal_width() * 0.7); + + if (pending_lines > 0) { + std::cout << "\r\033[2K"; + for (int i = 0; i < pending_lines; ++i) std::cout << "\033[1A\033[2K"; + } else { + std::cout << "\r"; + } + + if (!confirmed.empty()) current_line += colored(confirmed, ansi::green); + + while (true) { + size_t idx = wrap_index(current_line, (size_t)limit); + if (idx == std::string::npos) break; + std::cout << "\r\033[K" << current_line.substr(0, idx) << ansi::reset << "\n"; + current_line = std::string(ansi::green) + current_line.substr(idx + 1); + } + std::cout << "\r\033[K" << current_line; + + std::string ghost = stats; + if (!pending.empty()) { + if (!ghost.empty()) ghost += "\n"; + ghost += colored("[pending] " + pending, ansi::yellow); + } + pending_lines = 0; + if (!ghost.empty()) { + std::cout << "\n"; + std::stringstream ss(ghost); + std::string line; + bool first = true; + while (std::getline(ss, line)) { + while (true) { + size_t idx = wrap_index(line, (size_t)limit); + if (idx == std::string::npos) break; + if (!first) std::cout << "\n"; + std::cout << line.substr(0, idx); + line = line.substr(idx + 1); + pending_lines++; + first = false; + } + if (!first) std::cout << "\n"; + std::cout << line; + pending_lines++; + first = false; + } + } + std::cout << std::flush; + }; + + while (!should_stop) { + auto now = std::chrono::steady_clock::now(); + if (now - last_process >= std::chrono::milliseconds(PROCESS_INTERVAL_MS)) { + last_process = now; + std::vector chunk; + { std::lock_guard lock(g_audio.mutex); chunk.swap(g_audio.buffer); } + step(std::move(chunk), true); + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + + g_audio.recording = false; + SDL_PauseAudioDevice(device, 1); + std::vector tail; + { std::lock_guard lock(g_audio.mutex); tail.swap(g_audio.buffer); } + step(std::move(tail), false); + + response_buffer[0] = '\0'; + cactus_stream_transcribe_stop(stream, response_buffer.data(), response_buffer.size()); + confirmed_text += json_string_value(std::string(response_buffer.data()), "confirmed"); + + std::cout << "\n\n"; + print_rule(); + std::cout << colored("Final transcript:", ansi::bold) << "\n" << confirmed_text << "\n"; + print_rule(); + + if (input_thread.joinable()) input_thread.detach(); + SDL_CloseAudioDevice(device); + SDL_Quit(); + return 0; +} + +#else + +static int run_live_transcription(cactus_model_t, const std::string&) { + std::cerr << "Live microphone transcription is unavailable in this build.\n" + "Transcribe an audio file instead: cactus transcribe --file audio.wav\n"; + return 1; +} + +#endif + +static void print_usage(const char* argv0) { + std::cerr << "Usage: " << argv0 << " [audio_file] [--language ]\n" + << " With an audio file: one-shot transcription. Without: live from the microphone.\n"; +} + +int main(int argc, char** argv) { + std::cout << std::unitbuf; + + if (argc < 2 || std::string(argv[1]) == "-h" || std::string(argv[1]) == "--help") { + print_usage(argv[0]); + return argc < 2 ? 1 : 0; + } + + std::string model_path = argv[1]; + std::string audio_file; + std::string language = "en"; + for (int i = 2; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "-h" || arg == "--help") { + print_usage(argv[0]); + return 0; + } else if (arg == "--language") { + if (i + 1 >= argc) { + std::cerr << "Error: --language requires a value\n"; + return 1; + } + language = argv[++i]; + } else if (!arg.empty() && arg[0] == '-') { + std::cerr << "Error: unknown option '" << arg << "'\n"; + print_usage(argv[0]); + return 1; + } else { + audio_file = arg; + } + } + + std::cout << "Loading model from " << model_path << "...\n"; + cactus_model_t model = cactus_init(model_path.c_str(), nullptr, false); + if (!model) { + std::cerr << "Failed to initialize model: " << cactus_get_last_error() << "\n"; + return 1; + } + + if (g_color) std::cout << "\033[2J\033[3J\033[H"; + print_banner(); + + int result; + if (!audio_file.empty()) { + std::cout << colored("Transcribing ", ansi::bold) << audio_file << "\n"; + print_rule(); + std::cout << "\n"; + result = transcribe_file(model, audio_file, model_path, language); + } else { + result = run_live_transcription(model, language); + } + + cactus_destroy(model); + std::cout << "\nGoodbye.\n"; + return result >= 0 ? 0 : 1; +} diff --git a/cactus-graph/CMakeLists.txt b/cactus-graph/CMakeLists.txt new file mode 100644 index 000000000..9c431a855 --- /dev/null +++ b/cactus-graph/CMakeLists.txt @@ -0,0 +1,54 @@ +cmake_minimum_required(VERSION 3.10) +project(CactusGraph LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT TARGET cactus_kernels) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../cactus-kernels ${CMAKE_CURRENT_BINARY_DIR}/cactus-kernels) +endif() + +add_library(cactus_graph STATIC + src/core.cpp + src/builder.cpp + src/execute.cpp + src/metal_runtime.cpp + src/metal_plan.cpp + src/io.cpp + src/param_io.cpp + src/ops_math.cpp + src/ops_nn.cpp + src/ops_conv.cpp + src/ops_recurrent.cpp + src/ops_tensor.cpp + src/ops_sample.cpp + src/ops_cache.cpp + src/ops_dsp.cpp + src/ops_image.cpp + src/graph_ffi.cpp + src/last_error.cpp +) + +set_target_properties(cactus_graph PROPERTIES POSITION_INDEPENDENT_CODE ON) + +target_include_directories(cactus_graph + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src +) + +target_link_libraries(cactus_graph PUBLIC cactus_kernels) + +target_compile_options(cactus_graph PRIVATE + -Wall -Wextra -pedantic -Wno-missing-field-initializers + $<$:-Wno-c99-extensions> + $<$:-Wno-pedantic> +) + +if((BUILD_TESTING OR CACTUS_BUILD_TESTS) AND CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + file(GLOB TEST_SOURCES "tests/test_*.cpp") + foreach(TEST_SRC ${TEST_SOURCES}) + get_filename_component(TEST_NAME ${TEST_SRC} NAME_WE) + add_executable(${TEST_NAME} ${TEST_SRC}) + target_link_libraries(${TEST_NAME} PRIVATE cactus_graph) + endforeach() +endif() diff --git a/cactus/build.sh b/cactus-graph/build.sh similarity index 57% rename from cactus/build.sh rename to cactus-graph/build.sh index bbdaeef82..c711c750a 100755 --- a/cactus/build.sh +++ b/cactus-graph/build.sh @@ -1,18 +1,16 @@ #!/bin/bash - set -e -echo "Building Cactus library..." +cd "$(dirname "$0")" -cd "$(dirname "$0")/../cactus" +echo "Building cactus-graph..." rm -rf build - mkdir -p build cd build cmake .. -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null 2>&1 make -j$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) -echo "Cactus library built successfully!" -echo "Library location: $(pwd)/lib/libcactus.a" +echo "cactus-graph built successfully!" +echo "Library: $(pwd)/libcactus_graph.a" diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h new file mode 100644 index 000000000..279d1be02 --- /dev/null +++ b/cactus-graph/cactus_graph.h @@ -0,0 +1,1168 @@ +#ifndef CACTUS_GRAPH_H +#define CACTUS_GRAPH_H + +#include "cactus_kernels.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +int cactus_backend_select(const char* backend); + +namespace cactus { + +enum class LogLevel { + DEBUG = 0, + INFO = 1, + WARN = 2, + ERROR = 3, + NONE = 4 +}; + +class Logger { +public: + static Logger& instance() { + static Logger logger; + return logger; + } + + void set_level(LogLevel level) { min_level_ = level; } + LogLevel get_level() const { return min_level_; } + + void set_callback(std::function cb) { + std::lock_guard lock(mutex_); + callback_ = cb; + } + + void log(LogLevel level, const std::string& component, const std::string& message) { + if (level < min_level_) return; + std::lock_guard lock(mutex_); + if (callback_) { + callback_(level, component, message); + } else { + const char* names[] = {"DEBUG", "INFO", "WARN", "ERROR"}; + std::cerr << "[" << names[static_cast(level)] << "] [" << component << "] " << message << std::endl; + } + if (level == LogLevel::ERROR) last_error_ = "[" + component + "] " + message; + } + + const std::string& last_error() const { return last_error_; } + void clear_error() { last_error_.clear(); } + +private: + Logger() : min_level_(LogLevel::WARN) {} + LogLevel min_level_; + std::mutex mutex_; + std::string last_error_; + std::function callback_; +}; + +} // namespace cactus + +#define CACTUS_LOG(level, component, msg) \ + do { \ + if (static_cast(level) >= static_cast(cactus::Logger::instance().get_level())) { \ + std::ostringstream _cactus_log_ss; \ + _cactus_log_ss << msg; \ + cactus::Logger::instance().log(level, component, _cactus_log_ss.str()); \ + } \ + } while(0) + +#define CACTUS_LOG_DEBUG(component, msg) CACTUS_LOG(cactus::LogLevel::DEBUG, component, msg) +#define CACTUS_LOG_INFO(component, msg) CACTUS_LOG(cactus::LogLevel::INFO, component, msg) +#define CACTUS_LOG_WARN(component, msg) CACTUS_LOG(cactus::LogLevel::WARN, component, msg) +#define CACTUS_LOG_ERROR(component, msg) CACTUS_LOG(cactus::LogLevel::ERROR, component, msg) + +enum class ComputeBackend { CPU = 0, METAL = 1 }; + +ComputeBackend cactus_default_backend(); + +enum class Activation { SILU, GELU, GELU_ERF, RELU, SIGMOID, TANH }; + +enum class OpType { + INPUT, PRECISION_CAST, + ADD, ADD_CLIPPED, SUBTRACT, MULTIPLY, DIVIDE, + ABS, POW, FLATTEN, VIEW, + MATMUL, TRANSPOSE, RESHAPE, SLICE, GATHER, EMBEDDING, + BILINEAR_INTERPOLATION, + SUM, MEAN, VARIANCE, MIN, MAX, CUMSUM, + RMS_NORM, ROPE, ROPE_GPTJ, SOFTMAX, + ATTENTION, ATTENTION_INT8_HYBRID, REL_POS_BIAS, + CONV1D_CAUSAL, CONV1D_K3, CONV1D_K7S3, CONV1D, + CONV1D_SAME_DEPTHWISE_K9, CONV1D_POINTWISE, + CONV2D_K3S2P1, CONV2D_DEPTHWISE_K3S2P1, CONV2D_POINTWISE_1X1, + GLU, BATCHNORM, + SCALAR_ADD, SCALAR_SUBTRACT, SCALAR_MULTIPLY, SCALAR_DIVIDE, + SCALAR_EXP, SCALAR_SQRT, SCALAR_COS, SCALAR_SIN, SCALAR_LOG, + RELU, SILU, GELU, GELU_ERF, SIGMOID, TANH, + SAMPLE, CONCAT, CAT, + SCATTER_TOPK, TOPK, LAYERNORM, GROUPNORM, + MOE_LAYER, INDEX, PERSISTENT, + LSTM_CELL, GATED_DELTANET_DECODE, GATED_DELTANET_PREFILL, + STFT, ALTUP_PREDICT, ALTUP_CORRECT, GAUSSIAN_TOPK, + MAXPOOL1D, BILSTM_SEQUENCE, LEAKY_RELU, + CONV2D_K3S1P1, STATS_POOL, WEIGHTED_STATS_POOL, + KV_CACHE_STATE, KV_CACHE_APPEND, ATTENTION_CACHED, + CONV_CACHE_STATE, CONV_CACHE_APPEND, + RFFT, IRFFT, MEL_FILTER_BANK, SPECTROGRAM, + IMAGE_PREPROCESS, + CLAMP, + DENSE_MLP_TQ_FUSED, + NOT_EQUAL, + SCALAR_NOT_EQUAL, + RECURRENT_CACHE_STATE, + RECURRENT_CACHE_WRITE, + CONV_CACHE_INITIALIZE +}; + +struct PrecisionTraits { + static constexpr size_t size_of(Precision prec) { + switch (prec) { + case Precision::INT8: return 1; + case Precision::FP16: return 2; + case Precision::FP32: return 4; + case Precision::CQ1: + case Precision::CQ2: + case Precision::CQ3: + case Precision::CQ4: return 1; // packed, not element-sized + } + return 1; + } + + static constexpr bool is_cq(Precision prec) { + return prec == Precision::CQ1 || prec == Precision::CQ2 || + prec == Precision::CQ3 || prec == Precision::CQ4; + } + + static constexpr uint32_t cq_bits(Precision prec) { + switch (prec) { + case Precision::CQ1: return 1; + case Precision::CQ2: return 2; + case Precision::CQ3: return 3; + case Precision::CQ4: return 4; + default: return 0; + } + } + + static constexpr size_t packed_size_of(Precision prec, size_t count) { + if (is_cq(prec)) { + uint32_t bits = cq_bits(prec); + return (count * bits + 7) / 8; + } + return count * size_of(prec); + } + + static size_t byte_offset_of(Precision prec, size_t element_offset) { + if (is_cq(prec)) { + uint32_t bits = cq_bits(prec); + return (element_offset * bits) / 8; + } + return element_offset * size_of(prec); + } + + static constexpr bool is_quantized(Precision prec) { + return is_cq(prec); + } + + static constexpr bool is_floating_point(Precision prec) { + return prec == Precision::FP16 || prec == Precision::FP32; + } +}; + +namespace Quantization { + void int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale = 1.0f); + void fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale = 1.0f); + void fp16_to_fp32(const __fp16* src, float* dst, size_t count); + void fp32_to_fp16(const float* src, __fp16* dst, size_t count); + void int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale = 1.0f); + void fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale = 1.0f); +} + +struct BroadcastInfo { + std::vector output_shape; + bool needs_broadcasting; + static BroadcastInfo compute(const std::vector& lhs, const std::vector& rhs); +}; + +struct TensorConfig { + Precision default_precision = Precision::FP16; + Precision compute_precision = Precision::FP16; + Precision output_precision = Precision::FP16; + bool auto_mixed_precision = false; + static TensorConfig& global(); +}; + +class BufferPool { +public: + BufferPool() = default; + ~BufferPool() = default; + BufferPool(const BufferPool&) = delete; + BufferPool& operator=(const BufferPool&) = delete; + BufferPool(BufferPool&&) noexcept = default; + BufferPool& operator=(BufferPool&&) noexcept = default; + + char* acquire(size_t byte_size); + void release(char* ptr, size_t byte_size); + void clear(); + + size_t active_bytes() const { return active_bytes_; } + size_t pool_bytes() const { return pool_bytes_; } + size_t peak_bytes() const { return peak_bytes_; } + +private: + std::unordered_map>> free_buffers_; + size_t active_bytes_ = 0; + size_t pool_bytes_ = 0; + size_t peak_bytes_ = 0; + size_t round_up_size(size_t size) const; +}; + +struct BufferDesc { + std::vector shape; + size_t total_size; + size_t byte_size; + std::unique_ptr data; + void* external_data; + char* pooled_data; + Precision precision; + + std::vector dynamic_dims; + size_t pooled_byte_size = 0; + + size_t group_size = 0; + size_t num_groups = 0; + + void* activation_scales_data = nullptr; + std::unique_ptr owned_activation_scales; + size_t num_rows_for_activation_scales = 0; + + BufferDesc(); + BufferDesc(const std::vector& s, Precision prec = Precision::FP16); + ~BufferDesc(); + BufferDesc(BufferDesc&& other) noexcept; + BufferDesc& operator=(BufferDesc&& other) noexcept; + BufferDesc(const BufferDesc&) = delete; + BufferDesc& operator=(const BufferDesc&) = delete; + + void* get_data(); + const void* get_data() const; + + template T* data_as() { return static_cast(get_data()); } + template const T* data_as() const { return static_cast(get_data()); } + + bool is_cq() const { return PrecisionTraits::is_cq(precision) && group_size > 0; } + + const __fp16* cq_codebook = nullptr; + const __fp16* cq_input_scale = nullptr; + const __fp16* cq_input_scale_recip = nullptr; + const __fp16* cq_norms = nullptr; + const int8_t* cq_left_signs = nullptr; + const int8_t* cq_right_signs = nullptr; + const uint32_t* cq_permutation = nullptr; + const __fp16* cq_rotation = nullptr; + uint32_t cq_flags = 0; + + CactusQuantMatrix to_cq_matrix() const { + return CactusQuantMatrix{ + .bits = PrecisionTraits::cq_bits(precision), + .K = static_cast(shape.size() >= 2 ? shape[1] : shape[0]), + .N = static_cast(shape.size() >= 2 ? shape[0] : 1), + .group_size = static_cast(group_size), + .num_groups = static_cast(num_groups), + .flags = cq_flags, + .codebook = cq_codebook, + .input_scale = cq_input_scale, + .input_scale_recip = cq_input_scale_recip, + .norms = cq_norms, + .packed_indices = static_cast(get_data()), + .left_signs = cq_left_signs, + .right_signs = cq_right_signs, + .permutation = cq_permutation, + .rotation = cq_rotation, + .expanded = nullptr, + .norm_f32 = nullptr, + }; + } + + bool has_activation_scales() const { return activation_scales_data != nullptr && num_rows_for_activation_scales > 0; } + const float* activation_scales_as_float() const { return reinterpret_cast(activation_scales_data); } + float* activation_scales_as_float() { return reinterpret_cast(activation_scales_data); } + void allocate_activation_scales(size_t num_rows) { + num_rows_for_activation_scales = num_rows; + owned_activation_scales = std::make_unique(num_rows * sizeof(float)); + activation_scales_data = owned_activation_scales.get(); + } + void set_activation_scales(void* scales_ptr, size_t num_rows) { + activation_scales_data = scales_ptr; num_rows_for_activation_scales = num_rows; + } + + void allocate(); + void allocate_from_pool(BufferPool& pool); + void release_to_pool(BufferPool& pool); + void release_memory(BufferPool& pool); + void set_external(void* ptr); + + bool has_dynamic_dims() const { return !dynamic_dims.empty(); } + void set_shape(const std::vector& new_shape); + void resize_from_pool(BufferPool& pool); +}; + +struct OpParams { + float scalar = 0.0f; + float scale = 1.0f; + float theta = 10000.0f; + float epsilon = 1e-6f; + int axis = -1; + bool pretransposed_rhs = false; + size_t position_offset = 0; + size_t slice_start = 0; + size_t slice_length = 0; + size_t window_size = 0; + bool is_causal = true; + bool attention_mask_is_additive = false; + float logit_cap = 0.0f; + std::vector new_shape; + std::vector permutation; + Precision output_precision = Precision::FP16; + BroadcastInfo broadcast_info; + ComputeBackend backend = cactus_default_backend(); + + size_t dilation = 1; + size_t stride = 1; + float temperature = 1.0f; + float top_p = 1.0f; + float min_p = 0.15f; + float repetition_penalty = 1.1f; + size_t top_k = 0; + size_t random_seed = 0; + + size_t index_value = 0; + size_t num_classes = 0; + size_t num_groups = 0; + size_t dst_height = 0; + size_t dst_width = 0; + bool align_corners = true; + bool normalize_routing = false; + size_t num_experts = 0; + size_t num_experts_per_tok = 0; + bool moe_gated = true; + Activation activation = Activation::SILU; + + std::vector bias_values; + std::vector bias_indices; + + const int8_t* cached_keys_int8 = nullptr; + const int8_t* cached_values_int8 = nullptr; + const float* cached_k_scales = nullptr; + const float* cached_v_scales = nullptr; + size_t cache_seq_len = 0; + size_t num_kv_heads = 0; + size_t head_dim = 0; + size_t num_fft_bins = 0; + size_t chunk_size = 0; + size_t num_altup_inputs = 0; + size_t v_head_dim = 0; + size_t kernel_size = 0; + size_t max_cache_seq_len = 0; + size_t cache_sink_size = 0; + size_t cache_slot = 0; + size_t cache_num_slots = 1; + + size_t hop_length = 0; + float power = 2.0f; + bool center = true; + float mel_floor = 1e-10f; + float dither = 0.0f; + float preemphasis_coef = 0.0f; + bool remove_dc_offset = false; + int log_mel_mode = 0; + int pad_mode_type = 0; + size_t num_mel_filters = 0; + size_t sampling_rate = 16000; + float min_frequency = 0.0f; + float max_frequency = 8000.0f; + int mel_norm_type = 0; + int mel_scale_type = 0; + + int patch_size = 16; + float rescale_factor = 1.0f / 255.0f; + float image_mean[3] = {0.5f, 0.5f, 0.5f}; + float image_std[3] = {0.5f, 0.5f, 0.5f}; + int target_width = 0; + int target_height = 0; + int image_channels = 3; +}; + +struct GraphNode { + size_t id; + OpType op_type; + std::vector input_ids; + BufferDesc output_buffer; + OpParams params; + GraphNode(size_t node_id, OpType type); +}; + +using nodes_vector = std::vector>; +using node_index_map_t = std::unordered_map; + +inline const BufferDesc& get_input( + const GraphNode& node, + size_t idx, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + return nodes[node_index_map.at(node.input_ids[idx])]->output_buffer; +} + +struct AxisDims { + size_t outer, axis_size, inner; + static AxisDims from_shape(const std::vector& shape, size_t axis) { + AxisDims d; + d.outer = 1; + for (size_t i = 0; i < axis; i++) d.outer *= shape[i]; + d.axis_size = shape[axis]; + d.inner = 1; + for (size_t i = axis + 1; i < shape.size(); i++) d.inner *= shape[i]; + return d; + } +}; + +template +void dispatch_binary_op(OpType op, const T* lhs, const T* rhs, T* output, size_t count); + +template +void dispatch_unary_op(OpType op, const T* input, T* output, size_t count, float param = 0.0f); + +void compute_node_optimized(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& node_index_map); +void shrink_thread_local_buffers(); + +namespace ValidationUtils { + void validate_tensor_dims(const std::vector& shape, size_t required_dims, const std::string& op_name); + void validate_precision(Precision actual, Precision required, const std::string& op_name); + void validate_input_count(size_t actual, size_t required, const std::string& op_name); +} + +namespace GraphFile { + class MappedFile; + struct SerializedGraph; +} + +struct FusedEmbedCtx { + int token_id = -1, position = 0; + CactusQuantMatrix ple{}; + CactusQuantMatrix proj{}; + const void* rms_weight = nullptr; + float emb_scale = 0.0f, ple_scale = 0.0f, proj_scale = 0.0f, final_scale = 0.0f, rms_eps = 1e-6f; + bool ok = false; +}; +void cactus_graph_set_fused_embed(const FusedEmbedCtx* ctx); +const FusedEmbedCtx* cactus_graph_fused_embed(); +bool cactus_graph_metal_fold_prologue(void* h_buf, void* ple_buf, void* pos_buf, + const CactusQuantMatrix* lm_head, size_t nl, size_t ple_dim); +bool cactus_graph_metal_tail(void* logits, size_t vocab); +void cactus_graph_metal_tail_commit(); + +bool cactus_graph_metal_argmax(uint32_t* idx, float* best, float* second); +void cactus_graph_mark_unadjusted(); +void cactus_graph_set_prefill_consistent(bool on); +bool cactus_graph_prefill_consistent(); +void cactus_graph_on_destroy(const void* graph); +void cactus_graph_set_sampling(const uint32_t* recent, int n_recent, float rep_penalty, + const float* bias_dense, size_t bias_len, + long long suppressed); +void cactus_graph_clear_sampling(); +bool cactus_graph_metal_adjusted(); +bool cactus_graph_metal_argmax_biased(); + +class CactusGraph { +public: + CactusGraph(); + ~CactusGraph(); + CactusGraph(const CactusGraph&) = delete; + CactusGraph& operator=(const CactusGraph&) = delete; + CactusGraph(CactusGraph&&) noexcept = default; + CactusGraph& operator=(CactusGraph&&) noexcept = default; + + struct DebugNodeEntry { + uint32_t layer_idx; + std::string name; + size_t node_id; + }; + + void save(const std::string& path); + static CactusGraph load(const std::string& path); + + size_t input(const std::vector& shape, Precision precision = Precision::FP16); + void set_input(size_t node_id, const void* data, Precision precision); + void set_external_input(size_t node_id, void* data, Precision precision); + void* get_output(size_t node_id); + + size_t add(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t add_clipped(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t subtract(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t multiply(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t divide(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t not_equal(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + + + size_t scalar_add(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_subtract(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_multiply(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_divide(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_not_equal(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_exp(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t scalar_sqrt(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t scalar_cos(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t scalar_sin(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t scalar_log(size_t input, ComputeBackend backend = cactus_default_backend()); + + size_t abs(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t pow(size_t input, float exponent, ComputeBackend backend = cactus_default_backend()); + size_t precision_cast(size_t input, Precision target_precision, ComputeBackend backend = cactus_default_backend()); + + size_t relu(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t leaky_relu(size_t input, float negative_slope = 0.01f, ComputeBackend backend = cactus_default_backend()); + size_t clamp(size_t input, float lo, float hi, ComputeBackend backend = cactus_default_backend()); + size_t silu(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t gelu(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t gelu_erf(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t sigmoid(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t tanh(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t glu(size_t input, int axis = -1, ComputeBackend backend = cactus_default_backend()); + + + void set_node_backend(size_t node_id, ComputeBackend backend); + + size_t sum(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t mean(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t variance(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t min(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t max(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t cumsum(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t softmax(size_t input, int axis = -1, ComputeBackend backend = cactus_default_backend()); + size_t topk(size_t input, size_t k, ComputeBackend backend = cactus_default_backend()); + + size_t reshape(size_t input, const std::vector& new_shape, ComputeBackend backend = cactus_default_backend()); + size_t view(size_t input, const std::vector& new_shape, ComputeBackend backend = cactus_default_backend()); + size_t flatten(size_t input, int start_dim = 0, int end_dim = -1, ComputeBackend backend = cactus_default_backend()); + size_t transpose(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t transposeN(size_t input, const std::vector& permutation, ComputeBackend backend = cactus_default_backend()); + size_t slice(size_t input, int axis, size_t start, size_t length, ComputeBackend backend = cactus_default_backend()); + size_t index(size_t input, size_t index_value, int dim, ComputeBackend backend = cactus_default_backend()); + size_t concat(size_t input1, size_t input2, int axis = 0, ComputeBackend backend = cactus_default_backend()); + size_t cat(const std::vector& inputs, int axis, ComputeBackend backend = cactus_default_backend()); + + size_t matmul( + size_t input1, + size_t input2, + bool pretransposed_rhs = false, + ComputeBackend backend = cactus_default_backend()); + + size_t rms_norm(size_t input, size_t weight, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t layernorm(size_t input, size_t weight, size_t bias, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t layernorm(size_t input, size_t weight, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t groupnorm(size_t input, size_t weight, size_t bias, size_t num_groups = 32, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t batchnorm(size_t input, size_t weight, size_t bias, size_t running_mean, size_t running_var, int axis = 1, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + + + size_t rope(size_t input, float theta, size_t position_offset = 0, ComputeBackend backend = cactus_default_backend()); + size_t rope_gptj(size_t input, float theta, size_t position_offset = 0, size_t rot_dim = 0, ComputeBackend backend = cactus_default_backend()); + + + size_t attention(size_t query, size_t key, size_t value, float scale, + bool is_causal = true, ComputeBackend backend = cactus_default_backend()); + size_t attention(size_t query, size_t key, size_t value, float scale, + size_t position_offset, ComputeBackend backend = cactus_default_backend()); + size_t attention(size_t query, size_t key, size_t value, float scale, + size_t position_offset, size_t window_size, ComputeBackend backend = cactus_default_backend()); + size_t attention_masked( + size_t query, size_t key, size_t value, size_t mask, float scale, + bool is_causal = true, ComputeBackend backend = cactus_default_backend(), + bool additive_mask = false, size_t position_offset = 0, size_t window_size = 0, + float logit_cap = 0.0f); + size_t rel_pos_bias(size_t query, size_t relative_key, float scale, ComputeBackend backend = cactus_default_backend()); + size_t attention_int8_hybrid( + size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset, + const int8_t* cached_keys, const int8_t* cached_values, + const float* k_scales, const float* v_scales, + size_t cache_len, size_t num_kv_heads, size_t head_dim, + size_t window_size = 0, size_t v_head_dim = 0, ComputeBackend backend = cactus_default_backend()); + + size_t kv_cache_state( + size_t max_seq_len, + size_t num_kv_heads, + size_t head_dim, + size_t window_size = 0, + size_t sink_size = 4, + size_t num_slots = 1, + ComputeBackend backend = cactus_default_backend()); + + size_t kv_cache_append( + size_t new_kv, + size_t cache_state_node, + size_t window_size = 0, + size_t sink_size = 4, + size_t cache_slot = 0, + ComputeBackend backend = cactus_default_backend()); + + size_t attention_cached( + size_t query, + size_t key_new, + size_t value_new, + size_t k_cache_state, + size_t v_cache_state, + float scale, + size_t position_offset = 0, + size_t window_size = 0, + size_t v_head_dim = 0, + size_t cache_slot = 0, + ComputeBackend backend = cactus_default_backend()); + + size_t conv_cache_state(size_t window_size, size_t hidden_dim, ComputeBackend backend = cactus_default_backend()); + size_t conv_cache_append(size_t new_data, size_t cache_state_node, ComputeBackend backend = cactus_default_backend()); + size_t conv_cache_initialize(size_t rows, size_t cache_state_node, ComputeBackend backend = cactus_default_backend()); + + size_t recurrent_cache_state(const std::vector& shape, Precision precision, ComputeBackend backend = cactus_default_backend()); + size_t recurrent_cache_write(size_t new_value, size_t cache_state, ComputeBackend backend = cactus_default_backend()); + + size_t conv1d_causal(size_t input, size_t weight, size_t kernel_size, size_t dilation = 1, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_k3(size_t input, size_t weight, size_t stride, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_k7s3(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv1d(size_t input, size_t weight, size_t stride, ComputeBackend backend = cactus_default_backend()); + size_t conv1d(size_t input, size_t weight, size_t bias, size_t stride, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_same_depthwise_k9(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_pointwise(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_pointwise(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s2p1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_depthwise_k3s2p1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_depthwise_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_pointwise_1x1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_pointwise_1x1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s1p1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s1p1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t stft(size_t input, size_t weight, size_t stride, size_t num_fft_bins, ComputeBackend backend = cactus_default_backend()); + + size_t rfft(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t irfft(size_t input, size_t output_length, ComputeBackend backend = cactus_default_backend()); + size_t mel_filter_bank( + size_t num_frequency_bins, size_t num_mel_filters, + float min_frequency, float max_frequency, size_t sampling_rate, + int norm_type = 1, int scale_type = 2, ComputeBackend backend = cactus_default_backend()); + size_t spectrogram( + size_t waveform, size_t mel_filters_node, + size_t frame_length, size_t hop_length, size_t fft_length, + float power = 2.0f, bool center = true, int pad_mode = 0, + float mel_floor = 1e-10f, int log_mel_mode = 0, + float dither = 0.0f, float preemphasis = 0.0f, + bool remove_dc_offset = false, ComputeBackend backend = cactus_default_backend()); + + size_t image_preprocess( + size_t pixel_input, + int src_width, int src_height, + int target_width, int target_height, + int patch_size, int channels = 3, + float rescale_factor = 1.0f / 255.0f, + const float* mean = nullptr, const float* std_dev = nullptr, + ComputeBackend backend = cactus_default_backend()); + + size_t bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners = true, ComputeBackend backend = cactus_default_backend()); + size_t maxpool1d(size_t input, size_t kernel_size, size_t stride, ComputeBackend backend = cactus_default_backend()); + + size_t lstm_cell( + size_t input, size_t h_prev, size_t c_prev, + size_t weight_ih, size_t weight_hh, size_t bias_ih, size_t bias_hh, + ComputeBackend backend = cactus_default_backend()); + size_t bilstm_sequence( + size_t input, + size_t w_ih_fwd, size_t w_hh_fwd, size_t b_ih_fwd, size_t b_hh_fwd, + size_t w_ih_bwd, size_t w_hh_bwd, size_t b_ih_bwd, size_t b_hh_bwd, + ComputeBackend backend = cactus_default_backend()); + size_t gated_deltanet_decode( + size_t query, size_t key, size_t value, + size_t gate_log, size_t beta, size_t initial_state, float scale = 0.0f, + ComputeBackend backend = cactus_default_backend()); + size_t gated_deltanet_prefill( + size_t query, size_t key, size_t value, + size_t gate_log, size_t beta, size_t initial_state, + size_t chunk_size = 64, float scale = 0.0f, + ComputeBackend backend = cactus_default_backend()); + + size_t altup_predict(size_t coefs, const size_t* streams, size_t num_streams, ComputeBackend backend = cactus_default_backend()); + size_t altup_correct(size_t coefs, size_t innovation, const size_t* predictions, size_t num_predictions, ComputeBackend backend = cactus_default_backend()); + size_t gaussian_topk(size_t input, float ppf, ComputeBackend backend = cactus_default_backend()); + size_t moe_layer( + size_t hidden, size_t routing_probs, size_t topk_indices, + const std::vector& w1_weights, const std::vector& w3_weights, + const std::vector& w2_weights, + size_t num_experts, size_t num_experts_per_tok, + bool normalize_routing, float epsilon, float routed_scaling_factor, + Activation activation = Activation::SILU, size_t per_expert_scale = 0, + ComputeBackend backend = cactus_default_backend()); + size_t moe_layer( + size_t hidden, size_t routing_probs, size_t topk_indices, + const std::vector& w1_weights, const std::vector& w2_weights, + size_t num_experts, size_t num_experts_per_tok, + bool normalize_routing, float epsilon, float routed_scaling_factor, + Activation activation, ComputeBackend backend = cactus_default_backend()); + size_t dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t up_weight, size_t down_weight, float product_scale = 1.0f, ComputeBackend backend = cactus_default_backend()); + size_t stats_pool(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t weighted_stats_pool(size_t input, size_t weights, ComputeBackend backend = cactus_default_backend()); + + size_t sample( + size_t logits, float temperature = 0.6f, float top_p = 0.95f, size_t top_k = 20, + const std::unordered_map& logit_bias = {}, + ComputeBackend backend = cactus_default_backend()); + size_t sample_with_options( + size_t logits, float temperature, float top_p, float min_p, float repetition_penalty, + size_t top_k, const std::unordered_map& logit_bias = {}, + ComputeBackend backend = cactus_default_backend()); + size_t scatter_topk(size_t indices, size_t values, size_t num_classes, ComputeBackend backend = cactus_default_backend()); + + size_t gather(size_t embeddings, size_t indices, ComputeBackend backend = cactus_default_backend()); + size_t embedding(const std::string& filename, size_t indices, ComputeBackend backend = cactus_default_backend()); + size_t embedding(size_t embedding_tensor, size_t indices, ComputeBackend backend = cactus_default_backend()); + size_t mmap_embeddings(const std::string& filename); + size_t mmap_weights(const std::string& filename); + void bind_mmap_weights(size_t node_id, const std::string& filename); + void mark_embedded_input(size_t node_id); + bool is_embedded_input(size_t node_id) const; + void release_weight_pages(size_t node_id); + void prefetch_weight_pages(size_t node_id); + void release_all_weight_pages(); + void release_runtime_buffers(); + void clear_buffer_pool(); + void retain_outputs(const std::vector& node_ids); + + size_t persistent(size_t source_node, ComputeBackend backend = cactus_default_backend()); + bool is_populated(size_t persistent_node_id) const; + void invalidate_persistent(size_t persistent_node_id); + + void execute(const std::string& profile_file = ""); + bool extract_ple_pathway(FusedEmbedCtx& ctx) const; + void hard_reset(); + void soft_reset(); + void soft_reset_keep_pool(); + void set_prefill_mode(bool enabled) { prefill_mode_ = enabled; } + + void register_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id); + void capture_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id); + const std::vector& get_debug_nodes() const; + void clear_debug_nodes(); + + size_t add_node(OpType op_type, const std::vector& inputs, + const std::vector& output_shape, const OpParams& params = {}); + const BufferDesc& get_output_buffer(size_t node_id) const; + OpType get_node_op_type(size_t node_id) const; + size_t get_node_window_size(size_t node_id) const; + size_t get_node_sink_size(size_t node_id) const; + size_t get_node_cache_num_slots(size_t node_id) const; + void resize_cache_slots(size_t node_id, size_t num_slots); + void steal_cache_buffer(size_t dst_node, CactusGraph& src, size_t src_node); + void shrink_cache_buffer(size_t node_id, size_t new_capacity); + std::vector snapshot_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens) const; + void rollback_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens, + const std::vector& backup); + void allocate_buffers(); + size_t get_node_count() const; + void prewarm_metal_quant_weights(); + void set_runtime_input_shape(size_t node_id, const std::vector& shape); + void set_input_dynamic_dims(size_t node_id, const std::vector& dynamic_dims); + bool has_dynamic_shapes() const { return has_dynamic_shapes_; } + + std::vector> nodes_; + std::unordered_map node_index_map_; + +private: + size_t binary_broadcast_op(OpType op, size_t input1, size_t input2); + size_t tag_backend(size_t node_id, ComputeBackend backend); + void infer_shapes(); + size_t reduction_op(OpType op, size_t input, int axis); + size_t attach_conv_bias(size_t node, size_t bias, size_t expected_size, const char* op_name); + static CactusGraph from_serialized(const GraphFile::SerializedGraph& serialized); + size_t next_node_id_; + std::vector> mapped_files_; + std::unordered_map weight_cache_; + std::unordered_map node_to_mapped_file_; + std::vector debug_nodes_; + BufferPool buffer_pool_; + bool prefill_mode_ = false; + bool has_dynamic_shapes_ = false; + bool runtime_shapes_dirty_ = false; + std::unordered_set persistent_node_ids_; + std::unordered_set populated_node_ids_; + std::unordered_set embedded_input_node_ids_; + std::unordered_set retained_output_node_ids_; + void build_metal_retype_plan(); + void invalidate_metal_state(); + std::unordered_map metal_plans_; + std::unordered_map> metal_plan_banned_; + uint64_t metal_plan_sig_ = 0; + std::vector metal_retype_plan_; + bool metal_retype_built_ = false; + bool metal_retype_disabled_ = false; + std::unordered_map> metal_persistent_acts_; +}; + +namespace GraphFile { + struct GraphHeader { + uint32_t magic; + uint32_t version; + uint32_t node_count; + uint32_t flags = 0; + }; + + struct NodeEntry { + uint32_t index; + OpType op_type; + std::vector inputs; + std::vector output_shape; + Precision precision; + OpParams params; + bool has_embedded_data = false; + std::vector embedded_data; + std::vector dynamic_mask; + }; + + struct SerializedGraph { + GraphHeader header; + std::vector nodes; + std::vector graph_inputs; + std::vector graph_outputs; + }; + + SerializedGraph load_graph(const std::string& filename); + void save_graph(const CactusGraph& graph, const std::string& filename); + void save_node(CactusGraph& graph, size_t node_id, const std::string& filename); + + class MappedFile { + public: + MappedFile(const std::string& filename); + ~MappedFile(); + MappedFile(const MappedFile&) = delete; + MappedFile& operator=(const MappedFile&) = delete; + MappedFile(MappedFile&& other) noexcept; + MappedFile& operator=(MappedFile&& other) noexcept; + + const std::vector& shape() const; + Precision precision() const; + size_t byte_size() const; + size_t group_size() const { return group_size_; } + size_t num_groups() const { return num_groups_; } + const void* scales_data() const; + bool is_orthogonal_rotation() const { return is_orthogonal_rotation_; } + bool is_interleaved_4row() const { return is_interleaved_4row_; } + size_t original_N() const { return original_N_; } + void* data(); + const void* data() const; + template const T* typed_data() const; + void release_pages(); + void prefetch_pages(); + + private: + int fd_; + void* mapped_data_; + size_t file_size_, data_offset_; + std::vector shape_; + Precision precision_; + size_t byte_size_; + size_t group_size_ = 0; + size_t num_groups_ = 0; + size_t scales_offset_ = 0; + size_t scales_bytes_ = 0; + uint32_t alignment_ = 32; + bool is_orthogonal_rotation_ = false; + bool is_interleaved_4row_ = false; + size_t original_N_ = 0; + + void parse_header(); + void apply_madvise_hints(); + }; +} + +#if __GNUC__ >= 4 + #define CACTUS_FFI_EXPORT __attribute__((visibility("default"))) +#else + #define CACTUS_FFI_EXPORT +#endif + +extern thread_local std::string last_error_message; + +#ifdef __cplusplus +extern "C" { +#endif + +typedef void* cactus_graph_t; +typedef uint64_t cactus_node_t; + +typedef struct { + int32_t precision; + size_t rank; + size_t shape[8]; + size_t num_elements; + size_t byte_size; +} cactus_tensor_info_t; + +CACTUS_FFI_EXPORT const char* cactus_get_last_error(void); + +CACTUS_FFI_EXPORT cactus_graph_t cactus_graph_create(void); +CACTUS_FFI_EXPORT void cactus_graph_destroy(cactus_graph_t graph); +CACTUS_FFI_EXPORT int cactus_graph_hard_reset(cactus_graph_t graph); + +CACTUS_FFI_EXPORT int cactus_graph_set_node_backend( + cactus_graph_t graph, cactus_node_t node, int32_t backend); + +CACTUS_FFI_EXPORT int cactus_graph_save(cactus_graph_t graph, const char* filename); +CACTUS_FFI_EXPORT cactus_graph_t cactus_graph_load(const char* filename); + +CACTUS_FFI_EXPORT int cactus_graph_input( + cactus_graph_t graph, const size_t* shape, size_t rank, int32_t precision, +cactus_node_t* out_node); + +CACTUS_FFI_EXPORT int cactus_graph_set_input( + cactus_graph_t graph, cactus_node_t node, const void* data, int32_t +precision); +CACTUS_FFI_EXPORT int cactus_graph_set_external_input( + cactus_graph_t graph, cactus_node_t node, void* data, int32_t precision); +CACTUS_FFI_EXPORT int cactus_graph_mark_embedded_input( + cactus_graph_t graph, cactus_node_t node); +CACTUS_FFI_EXPORT int cactus_graph_set_runtime_input_shape( + cactus_graph_t graph, cactus_node_t node, const size_t* shape, size_t rank); +CACTUS_FFI_EXPORT int cactus_graph_set_input_dynamic_dims( + cactus_graph_t graph, cactus_node_t node, const uint8_t* mask, size_t rank); + +CACTUS_FFI_EXPORT int cactus_graph_precision_cast( + cactus_graph_t graph, cactus_node_t input, int32_t target_precision, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_add(cactus_graph_t graph, cactus_node_t a, +cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_add_clipped(cactus_graph_t graph, cactus_node_t a, +cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_subtract(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_multiply(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_divide(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_not_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_subtract(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_multiply(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_exp(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_sqrt(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_cos(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_sin(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_log(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_abs(cactus_graph_t graph, cactus_node_t x, +cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_pow(cactus_graph_t graph, cactus_node_t x, +float exponent, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_view( + cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, +cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_flatten( + cactus_graph_t graph, cactus_node_t x, int32_t start_dim, int32_t end_dim, +cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_reshape( + cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_transpose( + cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_transpose_n( + cactus_graph_t graph, cactus_node_t x, const size_t* permutation, size_t rank, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_slice( + cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_index( + cactus_graph_t graph, cactus_node_t x, size_t index_value, int32_t dim, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_sum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_mean(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_variance(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_min(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_max(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_cumsum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_concat( + cactus_graph_t graph, cactus_node_t a, cactus_node_t b, int32_t axis, +cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_cat( + cactus_graph_t graph, const cactus_node_t* nodes, size_t count, int32_t +axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_matmul( + cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gather( + cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_embedding_from_tensor( + cactus_graph_t graph, cactus_node_t embedding_tensor, cactus_node_t indices, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_embedding_from_file( + cactus_graph_t graph, const char* filename, cactus_node_t indices, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_mmap_embeddings( + cactus_graph_t graph, const char* filename, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_mmap_weights( + cactus_graph_t graph, const char* filename, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bind_mmap_weights( + cactus_graph_t graph, cactus_node_t node, const char* filename); +CACTUS_FFI_EXPORT int cactus_graph_bilinear_interpolation( + cactus_graph_t graph, cactus_node_t pos_embeds, size_t dst_height, size_t dst_width, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_release_weight_pages(cactus_graph_t graph, cactus_node_t node); +CACTUS_FFI_EXPORT int cactus_graph_prefetch_weight_pages(cactus_graph_t graph, cactus_node_t node); +CACTUS_FFI_EXPORT int cactus_graph_release_all_weight_pages(cactus_graph_t graph); + +CACTUS_FFI_EXPORT int cactus_graph_relu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_silu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gelu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gelu_erf(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_sigmoid(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_tanh(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_glu(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_clamp(cactus_graph_t graph, cactus_node_t input, float lo, float hi, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused( + cactus_graph_t graph, cactus_node_t hidden, cactus_node_t gate_weight, cactus_node_t up_weight, + cactus_node_t down_weight, float product_scale, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_layernorm( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, float epsilon, bool has_bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_groupnorm( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, size_t num_groups, float epsilon, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_batchnorm( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, cactus_node_t running_mean, cactus_node_t running_var, int32_t axis, float epsilon, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_topk(cactus_graph_t graph, cactus_node_t input, size_t k, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_rms_norm( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, float epsilon, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_rope( + cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_rope_gptj( + cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, size_t rot_dim, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_softmax(cactus_graph_t graph, cactus_node_t input, int32_t axis, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_attention( + cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_rel_pos_bias( + cactus_graph_t graph, cactus_node_t query, cactus_node_t relative_key, float scale, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_attention_int8_hybrid( + cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, float scale, size_t position_offset, + const int8_t* cached_keys, const int8_t* cached_values, const float* k_scales, const float* v_scales, + size_t cache_len, size_t num_kv_heads, size_t head_dim, size_t window_size, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_kv_cache_state( + cactus_graph_t graph, size_t max_seq_len, size_t num_kv_heads, size_t head_dim, size_t window_size, size_t sink_size, size_t num_slots, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_kv_cache_append( + cactus_graph_t graph, cactus_node_t new_kv, cactus_node_t cache_state, size_t window_size, size_t sink_size, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_attention_cached( + cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, + cactus_node_t k_cache_state, cactus_node_t v_cache_state, + float scale, size_t position_offset, size_t window_size, size_t v_head_dim, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_conv_cache_state( + cactus_graph_t graph, size_t window_size, size_t hidden_dim, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv_cache_append( + cactus_graph_t graph, cactus_node_t new_data, cactus_node_t cache_state, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv_cache_initialize( + cactus_graph_t graph, cactus_node_t rows, cactus_node_t cache_state, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_recurrent_cache_state( + cactus_graph_t graph, const size_t* shape, size_t shape_len, int precision, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_recurrent_cache_write( + cactus_graph_t graph, cactus_node_t new_value, cactus_node_t cache_input, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_rfft( + cactus_graph_t graph, cactus_node_t input, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_irfft( + cactus_graph_t graph, cactus_node_t input, size_t output_length, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_mel_filter_bank( + cactus_graph_t graph, size_t num_frequency_bins, size_t num_mel_filters, + float min_frequency, float max_frequency, size_t sampling_rate, + int norm_type, int scale_type, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_spectrogram( + cactus_graph_t graph, cactus_node_t waveform, cactus_node_t mel_filters, + size_t frame_length, size_t hop_length, size_t fft_length, + float power, bool center, int pad_mode, + float mel_floor, int log_mel_mode, + float dither, float preemphasis, bool remove_dc_offset, + cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_image_preprocess( + cactus_graph_t graph, cactus_node_t pixel_input, + int src_width, int src_height, int target_width, int target_height, + int patch_size, int channels, float rescale_factor, + const float* mean, const float* std_dev, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_conv1d_causal( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t kernel_size, size_t dilation, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_k3( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t stride, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_k7s3( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_same_depthwise_k9( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_pointwise( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv2d_k3s2p1( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv2d_depthwise_k3s2p1( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv2d_pointwise_1x1( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_lstm_cell( + cactus_graph_t graph, cactus_node_t input, cactus_node_t h_prev, cactus_node_t c_prev, cactus_node_t weight_ih, cactus_node_t weight_hh, cactus_node_t bias_ih, cactus_node_t bias_hh, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gated_deltanet_decode( + cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, cactus_node_t gate_log, cactus_node_t beta, cactus_node_t initial_state, float scale, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gated_deltanet_prefill( + cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, cactus_node_t gate_log, cactus_node_t beta, cactus_node_t initial_state, size_t chunk_size, float scale, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_stft( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t stride, size_t num_fft_bins, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_altup_predict( + cactus_graph_t graph, cactus_node_t coefs, const cactus_node_t* streams, size_t num_streams, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_altup_correct( + cactus_graph_t graph, cactus_node_t coefs, cactus_node_t innovation, const cactus_node_t* predictions, size_t num_predictions, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gaussian_topk( + cactus_graph_t graph, cactus_node_t input, float ppf, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_moe_layer_gated( + cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, + const cactus_node_t* w1_weights, const cactus_node_t* w3_weights, const cactus_node_t* w2_weights, + size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_moe_layer_ungated( + cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, + const cactus_node_t* w1_weights, const cactus_node_t* w2_weights, + size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out); + +CACTUS_FFI_EXPORT int cactus_graph_sample( + cactus_graph_t graph, cactus_node_t logits, float temperature, float top_p, size_t top_k, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scatter_topk( + cactus_graph_t graph, cactus_node_t indices, cactus_node_t values, size_t num_classes, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_persistent( + cactus_graph_t graph, cactus_node_t source_node, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_is_populated( + cactus_graph_t graph, cactus_node_t persistent_node, int32_t* out_is_populated); +CACTUS_FFI_EXPORT int cactus_graph_invalidate_persistent( + cactus_graph_t graph, cactus_node_t persistent_node); + +CACTUS_FFI_EXPORT int cactus_graph_execute(cactus_graph_t graph); +CACTUS_FFI_EXPORT int cactus_graph_get_output_ptr(cactus_graph_t graph, +cactus_node_t node, void** out_ptr); +CACTUS_FFI_EXPORT int cactus_graph_get_output_info(cactus_graph_t graph, +cactus_node_t node, cactus_tensor_info_t* out_info); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp new file mode 100644 index 000000000..cc52bda29 --- /dev/null +++ b/cactus-graph/src/builder.cpp @@ -0,0 +1,1677 @@ +#include "../cactus_graph.h" +#include +#include +#include +#include +#include +#include + +namespace { + +bool use_fp16_kv_cache_for_builder() { + const char* value = std::getenv("CACTUS_KV_CACHE_FP16"); + return value != nullptr && std::strcmp(value, "1") == 0; +} + +} // namespace + +size_t CactusGraph::attach_conv_bias(size_t node, size_t bias, size_t expected_size, const char* op_name) { + const auto& b = get_output_buffer(bias); + if (b.total_size != expected_size) + throw std::runtime_error(std::string(op_name) + " bias size mismatch"); + nodes_[node_index_map_[node]]->input_ids.push_back(bias); + return node; +} + +size_t CactusGraph::input(const std::vector& shape, Precision precision) { + return add_node(OpType::INPUT, {}, shape, {.output_precision = precision}); +} + +void CactusGraph::set_node_backend(size_t node_id, ComputeBackend backend) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) throw std::runtime_error("set_node_backend: unknown node id"); + nodes_[it->second]->params.backend = backend; +} + +size_t CactusGraph::tag_backend(size_t node_id, ComputeBackend backend) { + nodes_[node_index_map_[node_id]]->params.backend = backend; + return node_id; +} + +size_t CactusGraph::binary_broadcast_op(OpType op, size_t input1, size_t input2) { + BroadcastInfo bi = BroadcastInfo::compute( + get_output_buffer(input1).shape, get_output_buffer(input2).shape); + return add_node(op, {input1, input2}, bi.output_shape, {.broadcast_info = bi}); +} + +size_t CactusGraph::add(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::ADD, a, b), backend); } +size_t CactusGraph::add_clipped(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::ADD_CLIPPED, a, b), backend); } +size_t CactusGraph::subtract(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::SUBTRACT, a, b), backend); } +size_t CactusGraph::multiply(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::MULTIPLY, a, b), backend); } +size_t CactusGraph::divide(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::DIVIDE, a, b), backend); } +size_t CactusGraph::not_equal(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::NOT_EQUAL, a, b), backend); } + +size_t CactusGraph::abs(size_t input, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + OpParams params{.output_precision = input_buffer.precision}; + return tag_backend(add_node(OpType::ABS, {input}, input_buffer.shape, params), backend); +} + +size_t CactusGraph::pow(size_t input, float exponent, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + OpParams params{.scalar = exponent, .output_precision = input_buffer.precision}; + return tag_backend(add_node(OpType::POW, {input}, input_buffer.shape, params), backend); +} + +size_t CactusGraph::view(size_t input, const std::vector& new_shape, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + + size_t input_elements = 1; + for (size_t dim : input_buffer.shape) { + input_elements *= dim; + } + + size_t new_elements = 1; + for (size_t dim : new_shape) { + new_elements *= dim; + } + + if (input_elements != new_elements) { + throw std::runtime_error("View operation requires total number of elements to remain the same"); + } + + OpParams params{.new_shape = new_shape}; + return tag_backend(add_node(OpType::VIEW, {input}, new_shape, params), backend); +} + +size_t CactusGraph::flatten(size_t input, int start_dim, int end_dim, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + const auto& shape = input_buffer.shape; + size_t rank = shape.size(); + + if (start_dim < 0) start_dim += rank; + if (end_dim < 0) end_dim += rank; + + if (start_dim < 0 || static_cast(start_dim) >= rank || + end_dim < 0 || static_cast(end_dim) >= rank || + start_dim > end_dim) { + throw std::runtime_error("Invalid start_dim or end_dim for flatten operation"); + } + + std::vector output_shape; + + for (int i = 0; i < start_dim; ++i) { + output_shape.push_back(shape[i]); + } + + size_t flattened_dim = 1; + for (int i = start_dim; i <= end_dim; ++i) { + flattened_dim *= shape[i]; + } + output_shape.push_back(flattened_dim); + + for (size_t i = end_dim + 1; i < rank; ++i) { + output_shape.push_back(shape[i]); + } + + OpParams params{.new_shape = output_shape}; + return tag_backend(add_node(OpType::FLATTEN, {input}, output_shape, params), backend); +} + +size_t CactusGraph::matmul(size_t input1, size_t input2, bool pretransposed_rhs, ComputeBackend backend) { + const auto& lhs_buffer = get_output_buffer(input1); + const auto& rhs_buffer = get_output_buffer(input2); + + if (lhs_buffer.shape.size() != 2 || rhs_buffer.shape.size() != 2) { + throw std::invalid_argument("Matrix multiplication requires 2D tensors"); + } + + size_t M = lhs_buffer.shape[0]; + size_t K = lhs_buffer.shape[1]; + size_t rhs_K = pretransposed_rhs ? rhs_buffer.shape[1] : rhs_buffer.shape[0]; + + size_t N = pretransposed_rhs ? rhs_buffer.shape[0] : rhs_buffer.shape[1]; + + if (K != rhs_K) { + std::cout << "Matrix dimensions incompatible for multiplication: " + << "lhs=[" << M << "," << K << "] rhs=[" << rhs_buffer.shape[0] << "," << rhs_buffer.shape[1] << "]" + << " pretransposed=" << pretransposed_rhs + << " (K=" << K << " != rhs_K=" << rhs_K << ")" << std::endl; + throw std::invalid_argument("Matrix dimensions incompatible for multiplication"); + } + + std::vector output_shape = {M, N}; + OpParams params{.pretransposed_rhs = pretransposed_rhs, .backend = backend}; + return add_node(OpType::MATMUL, {input1, input2}, output_shape, params); +} + +size_t CactusGraph::transpose(size_t input, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + std::vector output_shape = input_buffer.shape; + + if (output_shape.size() >= 2) { + std::swap(output_shape[output_shape.size()-2], output_shape[output_shape.size()-1]); + } + + std::vector permutation; + for (size_t i = 0; i < input_buffer.shape.size(); ++i) { + permutation.push_back(i); + } + if (permutation.size() >= 2) { + std::swap(permutation[permutation.size()-2], permutation[permutation.size()-1]); + } + + OpParams params{.permutation = permutation, .backend = backend}; + return add_node(OpType::TRANSPOSE, {input}, output_shape, params); +} + +size_t CactusGraph::transposeN(size_t input, const std::vector& permutation, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + if (permutation.size() != input_buffer.shape.size()) { + throw std::runtime_error("transposeN permutation size must match tensor rank"); + } + std::vector output_shape(permutation.size()); + for (size_t i = 0; i < permutation.size(); ++i) { + size_t p = permutation[i]; + if (p >= input_buffer.shape.size()) { + throw std::runtime_error("transposeN permutation index out of range"); + } + output_shape[i] = input_buffer.shape[p]; + } + OpParams params{.permutation = permutation, .backend = backend}; + return add_node(OpType::TRANSPOSE, {input}, output_shape, params); +} + +size_t CactusGraph::reshape(size_t input, const std::vector& new_shape, ComputeBackend backend) { + OpParams params{.new_shape = new_shape}; + return tag_backend(add_node(OpType::RESHAPE, {input}, new_shape, params), backend); +} + +size_t CactusGraph::index(size_t input, size_t index_value, int dim, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + const auto& shape = input_buffer.shape; + + if (shape.empty()) { + throw std::invalid_argument("Cannot index a scalar tensor"); + } + + int actual_dim = dim; + if (actual_dim < 0) { + actual_dim += static_cast(shape.size()); + } + + if (actual_dim < 0 || static_cast(actual_dim) >= shape.size()) { + throw std::invalid_argument("Index dimension out of bounds"); + } + + if (index_value >= shape[actual_dim]) { + throw std::invalid_argument("Index value " + std::to_string(index_value) + + " out of bounds for dimension " + std::to_string(actual_dim) + + " with size " + std::to_string(shape[actual_dim])); + } + + std::vector output_shape; + for (size_t i = 0; i < shape.size(); ++i) { + if (static_cast(i) != actual_dim) { + output_shape.push_back(shape[i]); + } + } + + if (output_shape.empty()) { + output_shape = {1}; + } + + OpParams params{.axis = actual_dim, .output_precision = input_buffer.precision, .index_value = index_value}; + return tag_backend(add_node(OpType::INDEX, {input}, output_shape, params), backend); +} + +size_t CactusGraph::reduction_op(OpType op, size_t input, int axis) { + const auto& buf = get_output_buffer(input); + std::vector out; + if (axis == -1) { + out = {1}; + } else { + out = buf.shape; + out.erase(out.begin() + axis); + if (out.empty()) out = {1}; + } + return add_node(op, {input}, out, {.axis = axis, .output_precision = buf.precision}); +} + +size_t CactusGraph::sum(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::SUM, input, axis), backend); } +size_t CactusGraph::mean(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::MEAN, input, axis), backend); } +size_t CactusGraph::variance(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::VARIANCE, input, axis), backend); } +size_t CactusGraph::min(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::MIN, input, axis), backend); } +size_t CactusGraph::max(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::MAX, input, axis), backend); } + +size_t CactusGraph::cumsum(size_t input, int axis, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + if (input_buffer.shape.empty()) { + throw std::runtime_error("Cumsum requires at least one dimension"); + } + + int actual_axis = axis; + if (actual_axis < 0) { + actual_axis += static_cast(input_buffer.shape.size()); + } + if (actual_axis < 0 || static_cast(actual_axis) >= input_buffer.shape.size()) { + throw std::runtime_error("Invalid axis for cumsum operation"); + } + + return tag_backend(add_node(OpType::CUMSUM, {input}, input_buffer.shape, {.axis = actual_axis, .output_precision = input_buffer.precision}), backend); +} + +size_t CactusGraph::rms_norm(size_t input, size_t weight, float epsilon, ComputeBackend backend) { + OpParams params{.epsilon = epsilon}; + return tag_backend(add_node(OpType::RMS_NORM, {input, weight}, {}, params), backend); +} + +size_t CactusGraph::rope(size_t input, float theta, size_t position_offset, ComputeBackend backend) { + OpParams params{.theta = theta, .position_offset = position_offset, .backend = backend}; + return add_node(OpType::ROPE, {input}, {}, params); +} + +size_t CactusGraph::softmax(size_t input, int axis, ComputeBackend backend) { + OpParams params{.axis = axis}; + return tag_backend(add_node(OpType::SOFTMAX, {input}, {}, params), backend); +} + +size_t CactusGraph::topk(size_t input, size_t k, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + + if (input_buffer.shape.empty()) { + throw std::runtime_error("TopK requires non-empty input tensor"); + } + + std::vector output_shape = {2, input_buffer.shape[0], k}; + OpParams params{.output_precision = Precision::FP32, .top_k = k}; + + return tag_backend(add_node(OpType::TOPK, {input}, output_shape, params), backend); +} + +size_t CactusGraph::moe_layer(size_t hidden, + size_t routing_probs, + size_t topk_indices, + const std::vector& w1_weights, + const std::vector& w3_weights, + const std::vector& w2_weights, + size_t num_experts, + size_t num_experts_per_tok, + bool normalize_routing, + float epsilon, + float routed_scaling_factor, + Activation activation, + size_t per_expert_scale, + ComputeBackend backend) { + const auto& hidden_buffer = get_output_buffer(hidden); + const auto& routing_buffer = get_output_buffer(routing_probs); + const auto& topk_buffer = get_output_buffer(topk_indices); + + if (hidden_buffer.shape.size() != 2) { + throw std::runtime_error("moe_layer expects [tokens, hidden_dim] for hidden"); + } + if (routing_buffer.shape.size() != 2 || topk_buffer.shape.size() != 2) { + throw std::runtime_error("moe_layer expects 2D routing_probs and topk_indices"); + } + if (routing_buffer.shape[0] != hidden_buffer.shape[0] || topk_buffer.shape[0] != hidden_buffer.shape[0]) { + throw std::runtime_error("moe_layer token dimension mismatch across inputs"); + } + if (w1_weights.size() != num_experts || w3_weights.size() != num_experts || w2_weights.size() != num_experts) { + throw std::runtime_error("moe_layer expects num_experts weight tensors for each of w1, w3, w2"); + } + + std::vector input_ids; + input_ids.reserve(3 + 3 * num_experts + 1); + input_ids.push_back(hidden); + input_ids.push_back(routing_probs); + input_ids.push_back(topk_indices); + for (size_t i = 0; i < num_experts; ++i) input_ids.push_back(w1_weights[i]); + for (size_t i = 0; i < num_experts; ++i) input_ids.push_back(w3_weights[i]); + for (size_t i = 0; i < num_experts; ++i) input_ids.push_back(w2_weights[i]); + if (per_expert_scale != 0) input_ids.push_back(per_expert_scale); + + OpParams params; + params.num_experts = num_experts; + params.num_experts_per_tok = num_experts_per_tok; + params.normalize_routing = normalize_routing; + params.epsilon = epsilon; + params.scalar = routed_scaling_factor; + params.output_precision = hidden_buffer.precision; + params.activation = activation; + params.moe_gated = true; + + return tag_backend(add_node(OpType::MOE_LAYER, input_ids, hidden_buffer.shape, params), backend); +} + +size_t CactusGraph::dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t up_weight, size_t down_weight, float product_scale, ComputeBackend backend) { + const auto& hidden_buffer = get_output_buffer(hidden); + const auto& down_buffer = get_output_buffer(down_weight); + if (hidden_buffer.shape.empty()) { + throw std::runtime_error("dense_mlp_tq_fused expects non-scalar hidden"); + } + if (!PrecisionTraits::is_cq(down_buffer.precision) || down_buffer.shape.size() != 2) { + throw std::runtime_error("dense_mlp_tq_fused expects 2D TQ down weight"); + } + + std::vector output_shape = hidden_buffer.shape; + output_shape.back() = down_buffer.shape[0]; + + OpParams params; + params.output_precision = Precision::FP16; + params.scalar = product_scale; + return tag_backend(add_node(OpType::DENSE_MLP_TQ_FUSED, + {hidden, gate_weight, up_weight, down_weight}, + output_shape, params), backend); +} + +size_t CactusGraph::moe_layer(size_t hidden, + size_t routing_probs, + size_t topk_indices, + const std::vector& w1_weights, + const std::vector& w2_weights, + size_t num_experts, + size_t num_experts_per_tok, + bool normalize_routing, + float epsilon, + float routed_scaling_factor, + Activation activation, + ComputeBackend backend) { + const auto& hidden_buffer = get_output_buffer(hidden); + const auto& routing_buffer = get_output_buffer(routing_probs); + const auto& topk_buffer = get_output_buffer(topk_indices); + + if (hidden_buffer.shape.size() != 2) { + throw std::runtime_error("moe_layer expects [tokens, hidden_dim] for hidden"); + } + if (routing_buffer.shape.size() != 2 || topk_buffer.shape.size() != 2) { + throw std::runtime_error("moe_layer expects 2D routing_probs and topk_indices"); + } + if (routing_buffer.shape[0] != hidden_buffer.shape[0] || topk_buffer.shape[0] != hidden_buffer.shape[0]) { + throw std::runtime_error("moe_layer token dimension mismatch across inputs"); + } + if (w1_weights.size() != num_experts || w2_weights.size() != num_experts) { + throw std::runtime_error("moe_layer expects num_experts weight tensors for each of w1, w2"); + } + + std::vector input_ids; + input_ids.reserve(3 + 2 * num_experts); + input_ids.push_back(hidden); + input_ids.push_back(routing_probs); + input_ids.push_back(topk_indices); + for (size_t i = 0; i < num_experts; ++i) input_ids.push_back(w1_weights[i]); + for (size_t i = 0; i < num_experts; ++i) input_ids.push_back(w2_weights[i]); + + OpParams params; + params.num_experts = num_experts; + params.num_experts_per_tok = num_experts_per_tok; + params.normalize_routing = normalize_routing; + params.epsilon = epsilon; + params.scalar = routed_scaling_factor; + params.output_precision = hidden_buffer.precision; + params.moe_gated = false; + params.activation = activation; + + return tag_backend(add_node(OpType::MOE_LAYER, input_ids, hidden_buffer.shape, params), backend); +} + +size_t CactusGraph::layernorm(size_t input, size_t weight, size_t bias, float epsilon, ComputeBackend backend) { + OpParams params{.epsilon = epsilon}; + return tag_backend(add_node(OpType::LAYERNORM, {input, weight, bias}, {}, params), backend); +} + +size_t CactusGraph::layernorm(size_t input, size_t weight, float epsilon, ComputeBackend backend) { + OpParams params{.epsilon = epsilon}; + return tag_backend(add_node(OpType::LAYERNORM, {input, weight}, {}, params), backend); +} + +size_t CactusGraph::groupnorm(size_t input, size_t weight, size_t bias, size_t num_groups, float epsilon, ComputeBackend backend) { + OpParams params{.epsilon = epsilon, .num_groups = num_groups}; + return tag_backend(add_node(OpType::GROUPNORM, {input, weight, bias}, {}, params), backend); +} + +size_t CactusGraph::batchnorm(size_t input, size_t weight, size_t bias, size_t running_mean, size_t running_var, int axis, float epsilon, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + const auto& b = get_output_buffer(bias); + const auto& rm = get_output_buffer(running_mean); + const auto& rv = get_output_buffer(running_var); + + if (xin.shape.empty()) { + throw std::runtime_error("batchnorm expects non-scalar input"); + } + + int actual_axis = axis; + if (actual_axis < 0) { + actual_axis += static_cast(xin.shape.size()); + } + if (actual_axis < 0 || static_cast(actual_axis) >= xin.shape.size()) { + throw std::runtime_error("batchnorm axis out of range"); + } + + const size_t C = xin.shape[static_cast(actual_axis)]; + if (w.total_size != C || b.total_size != C || rm.total_size != C || rv.total_size != C) { + throw std::runtime_error("batchnorm parameter size mismatch"); + } + + OpParams params{.epsilon = epsilon, .axis = actual_axis}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::BATCHNORM, {input, weight, bias, running_mean, running_var}, xin.shape, params), backend); +} + +size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, bool is_causal, ComputeBackend backend) { + OpParams params{.scale = scale, .is_causal = is_causal, .backend = backend}; + const auto& qs = get_output_buffer(query).shape; + const auto& vs = get_output_buffer(value).shape; + return add_node(OpType::ATTENTION, {query, key, value}, {qs[0], qs[1], qs[2], vs[3]}, params); +} + +size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, ComputeBackend backend) { + OpParams params{.scale = scale, .position_offset = position_offset, .backend = backend}; + const auto& qs = get_output_buffer(query).shape; + const auto& vs = get_output_buffer(value).shape; + return add_node(OpType::ATTENTION, {query, key, value}, {qs[0], qs[1], qs[2], vs[3]}, params); +} + +size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, size_t window_size, ComputeBackend backend) { + OpParams params{.scale = scale, .position_offset = position_offset, .window_size = window_size, .backend = backend}; + const auto& qs = get_output_buffer(query).shape; + const auto& vs = get_output_buffer(value).shape; + return add_node(OpType::ATTENTION, {query, key, value}, {qs[0], qs[1], qs[2], vs[3]}, params); +} + +size_t CactusGraph::attention_masked(size_t query, size_t key, size_t value, size_t mask, float scale, + bool is_causal, ComputeBackend backend, bool additive_mask, + size_t position_offset, size_t window_size, float logit_cap) { + OpParams params{ + .scale = scale, + .position_offset = position_offset, + .window_size = window_size, + .is_causal = is_causal, + .backend = backend + }; + params.attention_mask_is_additive = additive_mask; + params.logit_cap = logit_cap; + const auto& qs = get_output_buffer(query).shape; + const auto& vs = get_output_buffer(value).shape; + return add_node(OpType::ATTENTION, {query, key, value, mask}, {qs[0], qs[1], qs[2], vs[3]}, params); +} + +size_t CactusGraph::rel_pos_bias(size_t query, size_t relative_key, float scale, ComputeBackend backend) { + const auto& q = get_output_buffer(query); + const auto& r = get_output_buffer(relative_key); + + if (q.shape.size() != 4) { + throw std::runtime_error("rel_pos_bias expects query [B, T, H, D]"); + } + if (r.shape.size() != 4) { + throw std::runtime_error("rel_pos_bias expects relative_key [B, R, H, D]"); + } + + const size_t B = q.shape[0]; + const size_t T = q.shape[1]; + const size_t H = q.shape[2]; + const size_t D = q.shape[3]; + const size_t R = r.shape[1]; + + if (r.shape[0] != 1 && r.shape[0] != B) { + throw std::runtime_error("rel_pos_bias relative_key batch must be 1 or match query batch"); + } + if (r.shape[2] != H || r.shape[3] != D) { + throw std::runtime_error("rel_pos_bias expects matching [H, D] between query and relative_key"); + } + if (R < (2 * T - 1)) { + throw std::runtime_error("rel_pos_bias requires relative_key length >= 2*T-1"); + } + + OpParams params{.scale = scale, .output_precision = q.precision}; + return tag_backend(add_node(OpType::REL_POS_BIAS, {query, relative_key}, {B, H, T, T}, params), backend); +} + +size_t CactusGraph::attention_int8_hybrid(size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset, + const int8_t* cached_keys, const int8_t* cached_values, + const float* k_scales, const float* v_scales, + size_t cache_len, size_t num_kv_heads, size_t head_dim, + size_t window_size, size_t v_head_dim, ComputeBackend backend) { + OpParams params; + params.scale = scale; + params.position_offset = position_offset; + params.window_size = window_size; + params.cached_keys_int8 = cached_keys; + params.cached_values_int8 = cached_values; + params.cached_k_scales = k_scales; + params.cached_v_scales = v_scales; + params.cache_seq_len = cache_len; + params.num_kv_heads = num_kv_heads; + params.head_dim = head_dim; + params.v_head_dim = v_head_dim; + std::vector out_shape; + if (v_head_dim != 0 && v_head_dim != head_dim) { + const auto& q_buf = get_output_buffer(query); + out_shape = {q_buf.shape[0], q_buf.shape[1], q_buf.shape[2], v_head_dim}; + } + return tag_backend(add_node(OpType::ATTENTION_INT8_HYBRID, {query, key_new, value_new}, out_shape, params), backend); +} + +size_t CactusGraph::conv1d_causal(size_t input, size_t weight, size_t, size_t dilation, ComputeBackend backend) { + OpParams params{.dilation = dilation}; + return tag_backend(add_node(OpType::CONV1D_CAUSAL, {input, weight}, {}, params), backend); +} + +size_t CactusGraph::conv1d_k3(size_t input, size_t weight, size_t stride, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + + if (xin.shape.size() != 3) throw std::runtime_error("conv1d_k3 expects N,C,L"); + if (w.shape.size() != 3) throw std::runtime_error("weight must be [C_out,C_in,3]"); + if (w.shape[1] != xin.shape[1]) throw std::runtime_error("C_in mismatch in conv1d_k3"); + if (w.shape[2] != 3) throw std::runtime_error("K=3 expected in conv1d_k3"); + + const size_t N = xin.shape[0]; + const size_t L = xin.shape[2]; + const size_t C_out= w.shape[0]; + const size_t K = w.shape[2]; + + const size_t pad = 1; + const size_t L_out = (L + 2 * pad - K) / stride + 1; + + OpParams params{}; + params.stride = stride; + params.output_precision = xin.precision; + + std::vector out_shape{N, C_out, L_out}; + return tag_backend(add_node(OpType::CONV1D_K3, {input, weight}, out_shape, params), backend); +} + +size_t CactusGraph::conv1d_k7s3(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + const auto& b = get_output_buffer(bias); + + if (xin.shape.size() != 3) throw std::runtime_error("conv1d_k7s3 expects N,C,L"); + if (w.shape.size() != 3) throw std::runtime_error("weight must be [C_in, 7, C_out]"); + if (w.shape[0] != xin.shape[1]) throw std::runtime_error("C_in mismatch in conv1d_k7s3"); + if (w.shape[1] != 7) throw std::runtime_error("K=7 expected in conv1d_k7s3"); + + size_t C_out = w.shape[2]; + if (b.total_size != C_out) throw std::runtime_error("Bias size mismatch"); + + const size_t N = xin.shape[0]; + const size_t L = xin.shape[2]; + const size_t K = 7; + const size_t stride = 3; + + const size_t L_out = (L < K) ? 0 : (L - K) / stride + 1; + + OpParams params{}; + params.stride = stride; + params.output_precision = xin.precision; + + std::vector out_shape{N, C_out, L_out}; + return tag_backend(add_node(OpType::CONV1D_K7S3, {input, weight, bias}, out_shape, params), backend); +} + +size_t CactusGraph::conv1d(size_t input, size_t weight, size_t stride, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + if (xin.shape.size() != 3) throw std::runtime_error("conv1d expects N,C,L"); + if (w.shape.size() != 3) throw std::runtime_error("conv1d weight expects [C_out, C_in, K]"); + size_t L_out = (xin.shape[2] - w.shape[2]) / stride + 1; + OpParams params{.stride = stride}; + return tag_backend(add_node(OpType::CONV1D, {input, weight}, {xin.shape[0], w.shape[0], L_out}, params), backend); +} + +size_t CactusGraph::conv1d(size_t input, size_t weight, size_t bias, size_t stride, ComputeBackend backend) { + const auto& b = get_output_buffer(bias); + const auto& w = get_output_buffer(weight); + if (b.total_size != w.shape[0]) throw std::runtime_error("conv1d bias size mismatch"); + // Reuse without-bias validation via conv1d(input, weight, stride), but need bias in inputs + const auto& xin = get_output_buffer(input); + if (xin.shape.size() != 3) throw std::runtime_error("conv1d expects N,C,L"); + if (w.shape.size() != 3) throw std::runtime_error("conv1d weight expects [C_out, C_in, K]"); + size_t L_out = (xin.shape[2] - w.shape[2]) / stride + 1; + OpParams params{.output_precision = xin.precision, .stride = stride}; + return tag_backend(add_node(OpType::CONV1D, {input, weight, bias}, {xin.shape[0], w.shape[0], L_out}, params), backend); +} + +size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + if (xin.shape.size() != 3) throw std::runtime_error("conv1d_same_depthwise_k9 expects input [N, L, C]"); + const size_t C = xin.shape[2]; + if (w.shape.size() == 2) { + if (w.shape[0] != C || w.shape[1] != 9) throw std::runtime_error("conv1d_same_depthwise_k9 weight must be [C, 9]"); + } else if (w.shape.size() == 3) { + if (w.shape[0] != C || w.shape[1] != 1 || w.shape[2] != 9) throw std::runtime_error("conv1d_same_depthwise_k9 weight must be [C, 1, 9]"); + } else { + throw std::runtime_error("conv1d_same_depthwise_k9 weight must be rank 2 or 3"); + } + OpParams params{}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::CONV1D_SAME_DEPTHWISE_K9, {input, weight}, {xin.shape[0], xin.shape[1], C}, params), backend); +} + +size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv1d_same_depthwise_k9(input, weight, backend); + return attach_conv_bias(node, bias, get_output_buffer(input).shape[2], "conv1d_same_depthwise_k9"); +} + +size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + if (xin.shape.size() != 3) throw std::runtime_error("conv1d_pointwise expects input [N, L, C_in]"); + const size_t C_in = xin.shape[2]; + size_t C_out = w.shape[0]; + if (w.shape.size() == 2) { + if (w.shape[1] != C_in) throw std::runtime_error("conv1d_pointwise weight must be [C_out, C_in]"); + } else if (w.shape.size() == 3) { + if (w.shape[1] != C_in || w.shape[2] != 1) throw std::runtime_error("conv1d_pointwise weight must be [C_out, C_in, 1]"); + } else { + throw std::runtime_error("conv1d_pointwise weight must be rank 2 or 3"); + } + OpParams params{}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::CONV1D_POINTWISE, {input, weight}, {xin.shape[0], xin.shape[1], C_out}, params), backend); +} + +size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv1d_pointwise(input, weight, backend); + return attach_conv_bias(node, bias, get_output_buffer(node).shape[2], "conv1d_pointwise"); +} + +size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + + if (xin.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s2p1 expects input [N, C_in, H, W]"); + } + if (w.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s2p1 weight must be [C_out, C_in, 3, 3]"); + } + + const size_t N = xin.shape[0]; + const size_t C_in = xin.shape[1]; + const size_t H = xin.shape[2]; + const size_t W = xin.shape[3]; + const size_t C_out = w.shape[0]; + + if (w.shape[1] != C_in || w.shape[2] != 3 || w.shape[3] != 3) { + throw std::runtime_error("conv2d_k3s2p1 weight must match [C_out, C_in, 3, 3]"); + } + if (H == 0 || W == 0) { + throw std::runtime_error("conv2d_k3s2p1 input spatial dimensions must be > 0"); + } + + const size_t H_out = (H - 1) / 2 + 1; + const size_t W_out = (W - 1) / 2 + 1; + + OpParams params{}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::CONV2D_K3S2P1, {input, weight}, {N, C_out, H_out, W_out}, params), backend); +} + +size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_k3s2p1(input, weight, backend); + return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_k3s2p1"); +} + +size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + + if (xin.shape.size() != 4) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 expects input [N, C, H, W]"); + } + + const size_t N = xin.shape[0]; + const size_t C = xin.shape[1]; + const size_t H = xin.shape[2]; + const size_t W = xin.shape[3]; + + if (w.shape.size() == 3) { + if (w.shape[0] != C || w.shape[1] != 3 || w.shape[2] != 3) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 weight must be [C, 3, 3]"); + } + } else if (w.shape.size() == 4) { + if (w.shape[0] != C || w.shape[1] != 1 || w.shape[2] != 3 || w.shape[3] != 3) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 weight must be [C, 1, 3, 3]"); + } + } else { + throw std::runtime_error("conv2d_depthwise_k3s2p1 weight must be rank 3 or 4"); + } + + if (H == 0 || W == 0) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 input spatial dimensions must be > 0"); + } + + const size_t H_out = (H - 1) / 2 + 1; + const size_t W_out = (W - 1) / 2 + 1; + OpParams params{}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::CONV2D_DEPTHWISE_K3S2P1, {input, weight}, {N, C, H_out, W_out}, params), backend); +} + +size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_depthwise_k3s2p1(input, weight, backend); + return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_depthwise_k3s2p1"); +} + +size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + + if (xin.shape.size() != 4) { + throw std::runtime_error("conv2d_pointwise_1x1 expects input [N, C_in, H, W]"); + } + + const size_t N = xin.shape[0]; + const size_t C_in = xin.shape[1]; + const size_t H = xin.shape[2]; + const size_t W = xin.shape[3]; + size_t C_out = 0; + + if (w.shape.size() == 2) { + C_out = w.shape[0]; + if (w.shape[1] != C_in) { + throw std::runtime_error("conv2d_pointwise_1x1 weight must be [C_out, C_in]"); + } + } else if (w.shape.size() == 4) { + C_out = w.shape[0]; + if (w.shape[1] != C_in || w.shape[2] != 1 || w.shape[3] != 1) { + throw std::runtime_error("conv2d_pointwise_1x1 weight must be [C_out, C_in, 1, 1]"); + } + } else { + throw std::runtime_error("conv2d_pointwise_1x1 weight must be rank 2 or 4"); + } + + if (H == 0 || W == 0) { + throw std::runtime_error("conv2d_pointwise_1x1 input spatial dimensions must be > 0"); + } + + OpParams params{}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::CONV2D_POINTWISE_1X1, {input, weight}, {N, C_out, H, W}, params), backend); +} + +size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_pointwise_1x1(input, weight, backend); + return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_pointwise_1x1"); +} + +size_t CactusGraph::lstm_cell(size_t input, size_t h_prev, size_t c_prev, size_t weight_ih, size_t weight_hh, size_t bias_ih, size_t bias_hh, ComputeBackend backend) { + const auto& h_buffer = get_output_buffer(h_prev); + std::vector output_shape = {h_buffer.shape[0], h_buffer.shape[1], 2}; + return tag_backend(add_node(OpType::LSTM_CELL, {input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh}, output_shape, {}), backend); +} + +size_t CactusGraph::gated_deltanet_decode(size_t query, size_t key, size_t value, size_t gate_log, size_t beta, + size_t initial_state, float scale, ComputeBackend backend) { + const auto& q = get_output_buffer(query); + const auto& k = get_output_buffer(key); + const auto& v = get_output_buffer(value); + const auto& g = get_output_buffer(gate_log); + const auto& b = get_output_buffer(beta); + const auto& s = get_output_buffer(initial_state); + + if (q.shape.size() != 4 || k.shape.size() != 4 || v.shape.size() != 4) { + throw std::runtime_error("gated_deltanet_decode expects query/key/value rank 4 [B, T, H, D]"); + } + if (g.shape.size() != 3 || b.shape.size() != 3) { + throw std::runtime_error("gated_deltanet_decode expects gate_log/beta rank 3 [B, T, H]"); + } + if (s.shape.size() != 4) { + throw std::runtime_error("gated_deltanet_decode expects initial_state rank 4 [B, K, H, V]"); + } + + const size_t B = q.shape[0]; + const size_t T = q.shape[1]; + const size_t Hq = q.shape[2]; + const size_t K = q.shape[3]; + const size_t Hv = v.shape[2]; + const size_t V = v.shape[3]; + if (T != 1) { + throw std::runtime_error("gated_deltanet_decode expects sequence length T=1"); + } + auto is_supported_precision = [](Precision p) { + return p == Precision::FP16 || p == Precision::FP32; + }; + if (!is_supported_precision(q.precision) || !is_supported_precision(k.precision) || + !is_supported_precision(v.precision) || !is_supported_precision(g.precision) || + !is_supported_precision(b.precision) || !is_supported_precision(s.precision)) { + throw std::runtime_error("gated_deltanet_decode requires FP16/FP32 inputs"); + } + + float op_scale = scale; + if (op_scale == 0.0f) { + op_scale = 1.0f / std::sqrt(static_cast(K)); + } + + OpParams params; + params.scale = op_scale; + params.num_kv_heads = Hq; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::GATED_DELTANET_DECODE, + {query, key, value, gate_log, beta, initial_state}, + {B, static_cast(1 + K), Hv, V}, + params), backend); +} + +size_t CactusGraph::gated_deltanet_prefill(size_t query, size_t key, size_t value, size_t gate_log, size_t beta, + size_t initial_state, size_t chunk_size, float scale, ComputeBackend backend) { + const auto& q = get_output_buffer(query); + const auto& k = get_output_buffer(key); + const auto& v = get_output_buffer(value); + const auto& g = get_output_buffer(gate_log); + const auto& b = get_output_buffer(beta); + const auto& s = get_output_buffer(initial_state); + + if (q.shape.size() != 4 || k.shape.size() != 4 || v.shape.size() != 4) { + throw std::runtime_error("gated_deltanet_prefill expects query/key/value rank 4 [B, T, H, D]"); + } + if (g.shape.size() != 3 || b.shape.size() != 3) { + throw std::runtime_error("gated_deltanet_prefill expects gate_log/beta rank 3 [B, T, H]"); + } + if (s.shape.size() != 4) { + throw std::runtime_error("gated_deltanet_prefill expects initial_state rank 4 [B, K, H, V]"); + } + + const size_t B = q.shape[0]; + const size_t T = q.shape[1]; + const size_t Hq = q.shape[2]; + const size_t K = q.shape[3]; + const size_t Hv = v.shape[2]; + const size_t V = v.shape[3]; + auto is_supported_precision = [](Precision p) { + return p == Precision::FP16 || p == Precision::FP32; + }; + if (!is_supported_precision(q.precision) || !is_supported_precision(k.precision) || + !is_supported_precision(v.precision) || !is_supported_precision(g.precision) || + !is_supported_precision(b.precision) || !is_supported_precision(s.precision)) { + throw std::runtime_error("gated_deltanet_prefill requires FP16/FP32 inputs"); + } + + if (chunk_size == 0) { + chunk_size = 64; + } + + float op_scale = scale; + if (op_scale == 0.0f) { + op_scale = 1.0f / std::sqrt(static_cast(K)); + } + + OpParams params; + params.scale = op_scale; + params.chunk_size = chunk_size; + params.num_kv_heads = Hq; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::GATED_DELTANET_PREFILL, + {query, key, value, gate_log, beta, initial_state}, + {B, T + K, Hv, V}, + params), backend); +} + +size_t CactusGraph::stft(size_t input, size_t weight, size_t stride, size_t num_fft_bins, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + + if (xin.shape.size() != 3) throw std::runtime_error("stft expects N,C,L input"); + if (w.shape.size() != 3) throw std::runtime_error("stft weight expects [C_out, C_in, K]"); + + size_t N = xin.shape[0]; + size_t L = xin.shape[2]; + size_t K = w.shape[2]; + size_t L_out = (L - K) / stride + 1; + + OpParams params{}; + params.stride = stride; + params.num_fft_bins = num_fft_bins; + + return tag_backend(add_node(OpType::STFT, {input, weight}, {N, 2 * num_fft_bins, L_out}, params), backend); +} + +size_t CactusGraph::concat(size_t input1, size_t input2, int axis, ComputeBackend backend) { + const auto& buffer1 = get_output_buffer(input1); + const auto& buffer2 = get_output_buffer(input2); + + if (buffer1.shape.size() != buffer2.shape.size()) { + throw std::runtime_error("Concat requires inputs with same number of dimensions"); + } + + std::vector output_shape = buffer1.shape; + size_t ndims = output_shape.size(); + + if (axis < 0) axis += ndims; + if (axis < 0 || static_cast(axis) >= ndims) { + throw std::runtime_error("Invalid axis for concat operation"); + } + + for (size_t i = 0; i < ndims; ++i) { + if (i != static_cast(axis) && buffer1.shape[i] != buffer2.shape[i]) { + throw std::runtime_error("Concat inputs must have same shape except on concat axis"); + } + } + + output_shape[axis] = buffer1.shape[axis] + buffer2.shape[axis]; + + OpParams params; + params.axis = axis; + return tag_backend(add_node(OpType::CONCAT, {input1, input2}, output_shape, params), backend); +} + +size_t CactusGraph::cat(const std::vector& inputs, int axis, ComputeBackend backend) { + if (inputs.empty()) { + throw std::runtime_error("Cat requires at least one input"); + } + + const auto& first_buffer = get_output_buffer(inputs[0]); + std::vector output_shape = first_buffer.shape; + size_t ndims = output_shape.size(); + + if (axis < 0) axis += ndims; + if (axis < 0 || static_cast(axis) >= ndims) { + throw std::runtime_error("Invalid axis for cat operation"); + } + + for (size_t i = 1; i < inputs.size(); ++i) { + const auto& buffer = get_output_buffer(inputs[i]); + if (buffer.shape.size() != ndims) { + throw std::runtime_error("All inputs to cat must have same number of dimensions"); + } + for (size_t d = 0; d < ndims; ++d) { + if (d != static_cast(axis) && buffer.shape[d] != output_shape[d]) { + throw std::runtime_error("All inputs to cat must have same shape except on cat axis"); + } + } + output_shape[axis] += buffer.shape[axis]; + } + + OpParams params; + params.axis = axis; + return tag_backend(add_node(OpType::CAT, inputs, output_shape, params), backend); +} + +size_t CactusGraph::scatter_topk(size_t indices, size_t values, size_t num_classes, ComputeBackend backend) { + const auto& indices_buffer = get_output_buffer(indices); + const auto& values_buffer = get_output_buffer(values); + + if (indices_buffer.shape != values_buffer.shape) { + throw std::runtime_error("ScatterTopK requires indices and values with identical shapes"); + } + if (indices_buffer.shape.size() != 2) { + throw std::runtime_error("ScatterTopK currently supports 2D tensors [batch, top_k]"); + } + if (indices_buffer.precision != Precision::FP32 || values_buffer.precision != Precision::FP32) { + throw std::runtime_error("ScatterTopK expects FP32 indices and values"); + } + + std::vector output_shape = {num_classes, indices_buffer.shape[0]}; + OpParams params{.output_precision = Precision::FP32, .num_classes = num_classes}; + return tag_backend(add_node(OpType::SCATTER_TOPK, {indices, values}, output_shape, params), backend); +} + +size_t CactusGraph::sample(size_t logits, float temperature, float top_p, size_t top_k, + const std::unordered_map& logit_bias, ComputeBackend backend) { + return this->sample_with_options(logits, temperature, top_p, 0.15f, 1.1f, top_k, logit_bias, backend); +} + +size_t CactusGraph::sample_with_options(size_t logits, float temperature, float top_p, + float min_p, float repetition_penalty, size_t top_k, + const std::unordered_map& logit_bias, ComputeBackend backend) { + const auto& logits_buffer = get_output_buffer(logits); + + if (logits_buffer.shape.empty()) { + throw std::runtime_error("Sample requires non-empty logits tensor"); + } + + OpParams params; + params.temperature = temperature; + params.top_p = top_p; + params.min_p = min_p; + params.repetition_penalty = repetition_penalty; + params.top_k = top_k; + params.random_seed = std::chrono::high_resolution_clock::now().time_since_epoch().count(); + params.output_precision = Precision::FP32; + + if (!logit_bias.empty()) { + params.bias_indices.reserve(logit_bias.size()); + params.bias_values.reserve(logit_bias.size()); + for (const auto& [idx, val] : logit_bias) { + params.bias_indices.push_back(idx); + params.bias_values.push_back(val); + } + } + + std::vector output_shape = {1}; + return tag_backend(add_node(OpType::SAMPLE, {logits}, output_shape, params), backend); +} + +static size_t scalar_val_op(CactusGraph& g, OpType op, size_t input, float value) { + OpParams params{}; + params.scalar = value; + params.output_precision = g.get_output_buffer(input).precision; + return g.add_node(op, {input}, {}, params); +} + +size_t CactusGraph::scalar_add(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_ADD, input, value), backend); } +size_t CactusGraph::scalar_subtract(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_SUBTRACT, input, value), backend); } +size_t CactusGraph::scalar_multiply(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_MULTIPLY, input, value), backend); } +size_t CactusGraph::scalar_divide(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_DIVIDE, input, value), backend); } +size_t CactusGraph::scalar_not_equal(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_NOT_EQUAL, input, value), backend); } + +size_t CactusGraph::scalar_exp(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_EXP, {input}, {}), backend); +} + +size_t CactusGraph::scalar_sqrt(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_SQRT, {input}, {}), backend); +} + +size_t CactusGraph::scalar_cos(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_COS, {input}, {}), backend); +} + +size_t CactusGraph::scalar_sin(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_SIN, {input}, {}), backend); +} + +size_t CactusGraph::scalar_log(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_LOG, {input}, {}), backend); +} + +size_t CactusGraph::relu(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::RELU, {input}, {}), backend); +} + +size_t CactusGraph::silu(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SILU, {input}, {}), backend); +} + +size_t CactusGraph::gelu(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::GELU, {input}, {}), backend); +} + +size_t CactusGraph::gelu_erf(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::GELU_ERF, {input}, {}), backend); +} + +size_t CactusGraph::sigmoid(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SIGMOID, {input}, {}), backend); +} + +size_t CactusGraph::tanh(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::TANH, {input}, {}), backend); +} + +size_t CactusGraph::glu(size_t input, int axis, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + if (xin.shape.empty()) { + throw std::runtime_error("glu expects non-scalar input"); + } + + int actual_axis = axis; + if (actual_axis < 0) { + actual_axis += static_cast(xin.shape.size()); + } + if (actual_axis < 0 || static_cast(actual_axis) >= xin.shape.size()) { + throw std::runtime_error("glu axis out of range"); + } + + std::vector out_shape = xin.shape; + const size_t axis_size = out_shape[static_cast(actual_axis)]; + if ((axis_size % 2) != 0) { + throw std::runtime_error("glu split dimension must be even"); + } + out_shape[static_cast(actual_axis)] = axis_size / 2; + + OpParams params{}; + params.axis = actual_axis; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::GLU, {input}, out_shape, params), backend); +} + +size_t CactusGraph::rope_gptj(size_t input, float theta, size_t position_offset, size_t rot_dim, ComputeBackend backend) { + OpParams params; + params.theta = theta; + params.position_offset = position_offset; + params.scalar = static_cast(rot_dim); + params.backend = backend; + return add_node(OpType::ROPE_GPTJ, {input}, {}, params); +} + +size_t CactusGraph::gather(size_t tensor, size_t indices, ComputeBackend backend) { + const auto& tensor_buffer = get_output_buffer(tensor); + const auto& idx_shape = get_output_buffer(indices).shape; + + if (tensor_buffer.shape.empty()) { + throw std::runtime_error("Cannot gather from scalar tensor"); + } + + std::vector output_shape = idx_shape; + for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { + output_shape.push_back(tensor_buffer.shape[i]); + } + + OpParams params; + params.output_precision = tensor_buffer.precision; + + return tag_backend(add_node(OpType::GATHER, {tensor, indices}, output_shape, params), backend); +} + +size_t CactusGraph::slice(size_t input, int axis, size_t start, size_t length, ComputeBackend backend) { + const auto& input_buffer = get_output_buffer(input); + if (input_buffer.shape.empty()) { + throw std::runtime_error("Cannot slice a scalar tensor"); + } + + size_t axis_index = static_cast(axis); + size_t axis_size = input_buffer.shape[axis_index]; + + if (start + length > axis_size) { + throw std::runtime_error("Slice range extends beyond axis size"); + } + + std::vector output_shape = input_buffer.shape; + output_shape[axis_index] = length; + + OpParams params; + params.axis = axis_index; + params.slice_start = start; + params.slice_length = length; + params.output_precision = input_buffer.precision; + + return tag_backend(add_node(OpType::SLICE, {input}, output_shape, params), backend); +} + +size_t CactusGraph::embedding(size_t embedding_tensor, size_t indices, ComputeBackend backend) { + const auto& emb_buffer = get_output_buffer(embedding_tensor); + const auto& idx_shape = get_output_buffer(indices).shape; + + if (emb_buffer.shape.size() != 2) { + std::cerr << "Error: Embedding tensor " << embedding_tensor << " has invalid shape: ["; + for(auto d : emb_buffer.shape) std::cerr << d << ","; + std::cerr << "]. OpType=" << (int)nodes_[node_index_map_[embedding_tensor]]->op_type + << " ExtData=" << emb_buffer.external_data << std::endl; + throw std::runtime_error("Embedding tensor must be 2D [vocab_size, hidden_dim]"); + } + + std::vector output_shape = idx_shape; + output_shape.push_back(emb_buffer.shape[1]); + + OpParams params; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::EMBEDDING, {embedding_tensor, indices}, output_shape, params), backend); +} + +size_t CactusGraph::bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners, ComputeBackend backend) { + const auto& pos_embeds_buffer = get_output_buffer(pos_embeds); + size_t embed_dim = pos_embeds_buffer.shape[1]; + std::vector output_shape = {dst_height * dst_width, embed_dim}; + + OpParams params; + params.dst_height = dst_height; + params.dst_width = dst_width; + params.align_corners = align_corners; + params.output_precision = Precision::FP16; + + return tag_backend(add_node(OpType::BILINEAR_INTERPOLATION, {pos_embeds}, output_shape, params), backend); +} + +size_t CactusGraph::precision_cast(size_t input, Precision target_precision, ComputeBackend backend) { + OpParams params{}; + params.output_precision = target_precision; + return tag_backend(add_node(OpType::PRECISION_CAST, {input}, {}, params), backend); +} + +size_t CactusGraph::add_node(OpType op_type, const std::vector& inputs, const std::vector& output_shape, const OpParams& params) { + auto node = std::make_unique(next_node_id_, op_type); + node->input_ids = inputs; + node->params = params; + + std::vector result_shape = output_shape; + if (result_shape.empty() && !inputs.empty()) { + result_shape = nodes_[node_index_map_[inputs[0]]]->output_buffer.shape; + } + + Precision result_precision = params.output_precision; + if (op_type == OpType::PRECISION_CAST || + op_type == OpType::EMBEDDING || + op_type == OpType::TOPK || + op_type == OpType::SCATTER_TOPK || + op_type == OpType::SAMPLE || + op_type == OpType::GAUSSIAN_TOPK) { + result_precision = params.output_precision; + } else if (!inputs.empty()) { + result_precision = nodes_[node_index_map_[inputs[0]]]->output_buffer.precision; + } + + node->output_buffer = BufferDesc(result_shape, result_precision); + + size_t node_id = next_node_id_++; + node_index_map_[node_id] = nodes_.size(); + nodes_.push_back(std::move(node)); + + return node_id; +} + +const BufferDesc& CactusGraph::get_output_buffer(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->output_buffer; +} + +OpType CactusGraph::get_node_op_type(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->op_type; +} + +size_t CactusGraph::get_node_window_size(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->params.window_size; +} + +size_t CactusGraph::get_node_sink_size(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->params.cache_sink_size; +} + +size_t CactusGraph::get_node_cache_num_slots(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->params.cache_num_slots; +} + +void CactusGraph::resize_cache_slots(size_t node_id, size_t num_slots) { + GraphNode& node = *nodes_[node_index_map_.at(node_id)]; + node.params.cache_num_slots = num_slots; + node.output_buffer.data.reset(); +} + +size_t CactusGraph::persistent(size_t source_node, ComputeBackend backend) { + const auto& source_buffer = get_output_buffer(source_node); + OpParams params; + params.output_precision = source_buffer.precision; + size_t node_id = add_node(OpType::PERSISTENT, {source_node}, source_buffer.shape, params); + persistent_node_ids_.insert(node_id); + return tag_backend(node_id, backend); +} + +bool CactusGraph::is_populated(size_t persistent_node_id) const { + return populated_node_ids_.count(persistent_node_id) > 0; +} + +void CactusGraph::invalidate_persistent(size_t persistent_node_id) { + populated_node_ids_.erase(persistent_node_id); + persistent_node_ids_.erase(persistent_node_id); +} + +size_t CactusGraph::altup_predict(size_t coefs, const size_t* streams, size_t num_streams, ComputeBackend backend) { + const auto& stream0_buf = get_output_buffer(streams[0]); + + size_t seq_len = stream0_buf.shape[0]; + size_t hidden_dim = stream0_buf.shape[1]; + + std::vector input_ids = {coefs}; + for (size_t i = 0; i < num_streams; i++) + input_ids.push_back(streams[i]); + + OpParams params; + params.num_altup_inputs = num_streams; + return tag_backend(add_node(OpType::ALTUP_PREDICT, input_ids, {num_streams * seq_len, hidden_dim}, params), backend); +} + +size_t CactusGraph::altup_correct(size_t coefs, size_t innovation, const size_t* predictions, size_t num_predictions, ComputeBackend backend) { + const auto& pred0_buf = get_output_buffer(predictions[0]); + + size_t seq_len = pred0_buf.shape[0]; + size_t hidden_dim = pred0_buf.shape[1]; + + std::vector input_ids = {coefs, innovation}; + for (size_t i = 0; i < num_predictions; i++) + input_ids.push_back(predictions[i]); + + OpParams params; + params.num_altup_inputs = num_predictions; + return tag_backend(add_node(OpType::ALTUP_CORRECT, input_ids, {num_predictions * seq_len, hidden_dim}, params), backend); +} + +size_t CactusGraph::gaussian_topk(size_t input, float ppf, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + if (in_buf.precision != Precision::FP16) { + throw std::runtime_error("gaussian_topk only supports FP16 input"); + } + OpParams params; + params.scalar = ppf; + params.output_precision = in_buf.precision; + return tag_backend(add_node(OpType::GAUSSIAN_TOPK, {input}, in_buf.shape, params), backend); +} + +size_t CactusGraph::leaky_relu(size_t input, float negative_slope, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + OpParams params; + params.scalar = negative_slope; + return tag_backend(add_node(OpType::LEAKY_RELU, {input}, in_buf.shape, params), backend); +} + +size_t CactusGraph::clamp(size_t input, float lo, float hi, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + OpParams params; + params.scalar = lo; + params.scale = hi; + return tag_backend(add_node(OpType::CLAMP, {input}, in_buf.shape, params), backend); +} + +size_t CactusGraph::bilstm_sequence(size_t input, size_t w_ih_fwd, size_t w_hh_fwd, size_t b_ih_fwd, size_t b_hh_fwd, + size_t w_ih_bwd, size_t w_hh_bwd, size_t b_ih_bwd, size_t b_hh_bwd, + ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + const auto& w_ih = get_output_buffer(w_ih_fwd); + size_t batch = in_buf.shape[0]; + size_t seq_len = in_buf.shape[1]; + size_t hidden_size = w_ih.shape[0] / 4; + return tag_backend(add_node(OpType::BILSTM_SEQUENCE, + {input, w_ih_fwd, w_hh_fwd, b_ih_fwd, b_hh_fwd, w_ih_bwd, w_hh_bwd, b_ih_bwd, b_hh_bwd}, + {batch, seq_len, 2 * hidden_size}, {}), backend); +} + +size_t CactusGraph::maxpool1d(size_t input, size_t kernel_size, size_t stride, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + size_t batch = in_buf.shape[0]; + size_t channels = in_buf.shape[1]; + size_t input_length = in_buf.shape[2]; + size_t output_length = (input_length - kernel_size) / stride + 1; + + OpParams params; + params.kernel_size = kernel_size; + params.stride = stride; + return tag_backend(add_node(OpType::MAXPOOL1D, {input}, {batch, channels, output_length}, params), backend); +} + +size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + const auto& w = get_output_buffer(weight); + + if (xin.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s1p1 expects input [N, C_in, H, W]"); + } + if (w.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s1p1 weight must be [C_out, C_in, 3, 3]"); + } + + const size_t N = xin.shape[0]; + const size_t C_in = xin.shape[1]; + const size_t H = xin.shape[2]; + const size_t W = xin.shape[3]; + const size_t C_out = w.shape[0]; + + if (w.shape[1] != C_in || w.shape[2] != 3 || w.shape[3] != 3) { + throw std::runtime_error("conv2d_k3s1p1 weight must match [C_out, C_in, 3, 3]"); + } + if (H == 0 || W == 0) { + throw std::runtime_error("conv2d_k3s1p1 input spatial dimensions must be > 0"); + } + + OpParams params{}; + params.output_precision = xin.precision; + return tag_backend(add_node(OpType::CONV2D_K3S1P1, {input, weight}, {N, C_out, H, W}, params), backend); +} + +size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_k3s1p1(input, weight, backend); + return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_k3s1p1"); +} + +size_t CactusGraph::stats_pool(size_t input, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + size_t batch = xin.shape[0]; + size_t features = 1; + for (size_t i = 1; i < xin.shape.size() - 1; ++i) features *= xin.shape[i]; + return tag_backend(add_node(OpType::STATS_POOL, {input}, {batch, features * 2}), backend); +} + +size_t CactusGraph::weighted_stats_pool(size_t input, size_t weights, ComputeBackend backend) { + const auto& xin = get_output_buffer(input); + size_t batch = xin.shape[0]; + size_t features = 1; + for (size_t i = 1; i < xin.shape.size() - 1; ++i) features *= xin.shape[i]; + return tag_backend(add_node(OpType::WEIGHTED_STATS_POOL, {input, weights}, {batch, features * 2}), backend); +} + +size_t CactusGraph::kv_cache_state(size_t max_seq_len, size_t num_kv_heads, size_t head_dim, + size_t window_size, size_t sink_size, size_t num_slots, ComputeBackend backend) { + if (num_slots == 0) num_slots = 1; + bool fp16_cache = use_fp16_kv_cache_for_builder(); + size_t total_elements = 0; + Precision precision = Precision::INT8; + if (fp16_cache) { + total_elements = (64 / sizeof(__fp16)) + max_seq_len * num_kv_heads * head_dim; + precision = Precision::FP16; + } else { + size_t num_groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + total_elements = 64 + max_seq_len * num_kv_heads * head_dim + + max_seq_len * num_kv_heads * num_groups * sizeof(float); + } + total_elements *= num_slots; + OpParams params{}; + params.max_cache_seq_len = max_seq_len; + params.num_kv_heads = num_kv_heads; + params.head_dim = head_dim; + params.window_size = window_size; + params.cache_sink_size = sink_size; + params.cache_num_slots = num_slots; + params.output_precision = precision; + size_t node_id = add_node(OpType::KV_CACHE_STATE, {}, {total_elements}, params); + persistent_node_ids_.insert(node_id); + return tag_backend(node_id, backend); +} + +size_t CactusGraph::kv_cache_append(size_t new_kv, size_t cache_state_node, + size_t window_size, size_t sink_size, size_t cache_slot, ComputeBackend backend) { + OpParams params{}; + params.window_size = window_size; + params.cache_sink_size = sink_size; + params.cache_slot = cache_slot; + params.output_precision = Precision::FP32; + return tag_backend(add_node(OpType::KV_CACHE_APPEND, {new_kv, cache_state_node}, {1}, params), backend); +} + +size_t CactusGraph::conv_cache_state(size_t ws, size_t hidden_dim, ComputeBackend backend) { + size_t total_bytes = sizeof(uint64_t) * 8 + ws * hidden_dim * sizeof(__fp16); + OpParams params{}; + params.window_size = ws; + params.head_dim = hidden_dim; + params.output_precision = Precision::INT8; + size_t node_id = add_node(OpType::CONV_CACHE_STATE, {}, {total_bytes}, params); + persistent_node_ids_.insert(node_id); + return tag_backend(node_id, backend); +} + +size_t CactusGraph::conv_cache_append(size_t new_data, size_t cache_state_node, ComputeBackend backend) { + const auto& cache_buf = get_output_buffer(cache_state_node); + auto* raw = static_cast(cache_buf.get_data()); + size_t ws, hd; + if (raw) { + ws = *reinterpret_cast(raw + 16); + hd = *reinterpret_cast(raw + 24); + } else { + auto it = node_index_map_.find(cache_state_node); + ws = nodes_[it->second]->params.window_size; + hd = nodes_[it->second]->params.head_dim; + } + OpParams params{}; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::CONV_CACHE_APPEND, {new_data, cache_state_node}, {ws, hd}, params), backend); +} + +size_t CactusGraph::conv_cache_initialize(size_t rows, size_t cache_state_node, ComputeBackend backend) { + if (get_node_op_type(cache_state_node) != OpType::CONV_CACHE_STATE) { + throw std::invalid_argument( + "conv_cache_initialize target must be a CONV_CACHE_STATE node"); + } + const auto& rows_buf = get_output_buffer(rows); + if (rows_buf.precision != Precision::FP16 && rows_buf.precision != Precision::FP32) { + throw std::invalid_argument( + "conv_cache_initialize requires FP16/FP32 input rows"); + } + const size_t hidden_dim = nodes_[node_index_map_.at(cache_state_node)]->params.head_dim; + if (hidden_dim == 0) { + throw std::invalid_argument( + "conv_cache_initialize: cache hidden_dim is zero"); + } + if (rows_buf.shape.empty() || rows_buf.shape.back() != hidden_dim) { + throw std::invalid_argument( + "conv_cache_initialize: rows last dim must equal cache hidden_dim"); + } + if (rows_buf.total_size % hidden_dim != 0) { + throw std::invalid_argument( + "conv_cache_initialize: rows total size must be a multiple of cache hidden_dim"); + } + OpParams params{}; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::CONV_CACHE_INITIALIZE, {rows, cache_state_node}, {0}, params), backend); +} + +size_t CactusGraph::recurrent_cache_state(const std::vector& shape, Precision precision, ComputeBackend backend) { + OpParams params{}; + params.output_precision = precision; + size_t node_id = add_node(OpType::RECURRENT_CACHE_STATE, {}, shape, params); + persistent_node_ids_.insert(node_id); + return tag_backend(node_id, backend); +} + +size_t CactusGraph::recurrent_cache_write(size_t new_value, size_t cache_state, ComputeBackend backend) { + const auto& src_buf = get_output_buffer(new_value); + const auto& dst_buf = get_output_buffer(cache_state); + if (get_node_op_type(cache_state) != OpType::RECURRENT_CACHE_STATE) { + throw std::invalid_argument( + "recurrent_cache_write target must be a RECURRENT_CACHE_STATE node"); + } + if (src_buf.shape != dst_buf.shape) { + throw std::invalid_argument("recurrent_cache_write requires matching shapes"); + } + if (src_buf.precision != dst_buf.precision) { + throw std::invalid_argument("recurrent_cache_write requires matching precision"); + } + OpParams params{}; + params.output_precision = dst_buf.precision; + return tag_backend(add_node(OpType::RECURRENT_CACHE_WRITE, {new_value, cache_state}, {0}, params), backend); +} + +size_t CactusGraph::image_preprocess( + size_t pixel_input, + int src_width, int src_height, + int target_width, int target_height, + int patch_size, int channels, + float rescale_factor, + const float* mean, const float* std_dev, + ComputeBackend backend) { + + int tw = target_width > 0 ? target_width : src_width; + int th = target_height > 0 ? target_height : src_height; + int ph = th / patch_size; + int pw = tw / patch_size; + size_t num_patches = static_cast(ph) * pw; + size_t patch_dim = static_cast(patch_size) * patch_size * channels; + + OpParams params{}; + params.patch_size = patch_size; + params.rescale_factor = rescale_factor; + params.target_width = tw; + params.target_height = th; + params.image_channels = channels; + params.dst_width = static_cast(src_width); + params.dst_height = static_cast(src_height); + float default_mean[3] = {0.5f, 0.5f, 0.5f}; + float default_std[3] = {0.5f, 0.5f, 0.5f}; + const float* m = mean ? mean : default_mean; + const float* s = std_dev ? std_dev : default_std; + for (int i = 0; i < 3; i++) { + params.image_mean[i] = m[i]; + params.image_std[i] = s[i]; + } + params.output_precision = Precision::FP32; + + return tag_backend(add_node(OpType::IMAGE_PREPROCESS, {pixel_input}, {num_patches, patch_dim}, params), backend); +} + +size_t CactusGraph::rfft(size_t input, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + size_t n = in_buf.total_size; + size_t out_len = (n / 2 + 1) * 2; + OpParams params{}; + params.output_precision = in_buf.precision; + return tag_backend(add_node(OpType::RFFT, {input}, {out_len}, params), backend); +} + +size_t CactusGraph::irfft(size_t input, size_t output_length, ComputeBackend backend) { + const auto& in_buf = get_output_buffer(input); + OpParams params{}; + params.output_precision = in_buf.precision; + return tag_backend(add_node(OpType::IRFFT, {input}, {output_length}, params), backend); +} + +size_t CactusGraph::mel_filter_bank( + size_t num_frequency_bins, size_t num_mel_filters, + float min_frequency, float max_frequency, size_t sampling_rate, + int norm_type, int scale_type, ComputeBackend backend) { + OpParams params{}; + params.num_mel_filters = num_mel_filters; + params.min_frequency = min_frequency; + params.max_frequency = max_frequency; + params.sampling_rate = sampling_rate; + params.mel_norm_type = norm_type; + params.mel_scale_type = scale_type; + params.output_precision = Precision::FP32; + return tag_backend(add_node(OpType::MEL_FILTER_BANK, {}, {num_frequency_bins, num_mel_filters}, params), backend); +} + +size_t CactusGraph::spectrogram( + size_t waveform, size_t mel_filters_node, + size_t frame_length, size_t hop_length, size_t fft_length, + float power, bool center, int pad_mode, + float mel_floor, int log_mel_mode, + float dither_val, float preemphasis, + bool remove_dc_offset, ComputeBackend backend) { + + const auto& wav_buf = get_output_buffer(waveform); + const auto& mel_buf = get_output_buffer(mel_filters_node); + + size_t waveform_length = wav_buf.total_size; + size_t num_frequency_bins = fft_length / 2 + 1; + size_t num_mel_bins = mel_buf.total_size / num_frequency_bins; + size_t pad_len = center ? frame_length / 2 : 0; + size_t padded_len = waveform_length + 2 * pad_len; + size_t num_frames = 1 + (padded_len - frame_length) / hop_length; + + OpParams params{}; + params.num_fft_bins = frame_length; + params.hop_length = hop_length; + params.stride = fft_length; + params.power = power; + params.center = center; + params.pad_mode_type = pad_mode; + params.mel_floor = mel_floor; + params.log_mel_mode = log_mel_mode; + params.dither = dither_val; + params.preemphasis_coef = preemphasis; + params.remove_dc_offset = remove_dc_offset; + params.output_precision = Precision::FP32; + + return tag_backend(add_node(OpType::SPECTROGRAM, {waveform, mel_filters_node}, {num_mel_bins, num_frames}, params), backend); +} + +size_t CactusGraph::attention_cached(size_t query, size_t key_new, size_t value_new, + size_t k_cache_state, size_t v_cache_state, + float scale, size_t position_offset, + size_t window_size, size_t v_head_dim, size_t cache_slot, + ComputeBackend backend) { + const auto& q_shape = get_output_buffer(query).shape; + size_t batch = q_shape[0]; + size_t seq_len = q_shape[1]; + size_t num_q_heads = q_shape[2]; + size_t head_dim = q_shape[3]; + size_t out_v_dim = v_head_dim > 0 ? v_head_dim : head_dim; + + OpParams params{}; + params.scale = scale; + params.position_offset = position_offset; + params.window_size = window_size; + params.v_head_dim = v_head_dim; + params.cache_slot = cache_slot; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::ATTENTION_CACHED, + {query, key_new, value_new, k_cache_state, v_cache_state}, + {batch, seq_len, num_q_heads, out_v_dim}, params), backend); +} diff --git a/cactus/graph/graph_core.cpp b/cactus-graph/src/core.cpp similarity index 65% rename from cactus/graph/graph_core.cpp rename to cactus-graph/src/core.cpp index 4ba6d54e4..f850c99bc 100644 --- a/cactus/graph/graph_core.cpp +++ b/cactus-graph/src/core.cpp @@ -1,4 +1,4 @@ -#include "graph.h" +#include "../cactus_graph.h" #include #include #include @@ -65,14 +65,13 @@ BufferDesc::BufferDesc(const std::vector& s, Precision prec) : shape(s), external_data(nullptr), pooled_data(nullptr), precision(prec) { total_size = 1; for (size_t dim : shape) total_size *= dim; - byte_size = total_size * PrecisionTraits::size_of(prec); + byte_size = PrecisionTraits::packed_size_of(prec, total_size); } BufferDesc::~BufferDesc() { - if (pooled_data) { delete[] pooled_data; - pooled_data = nullptr; - } + + } BufferDesc::BufferDesc(BufferDesc&& other) noexcept @@ -83,33 +82,45 @@ BufferDesc::BufferDesc(BufferDesc&& other) noexcept external_data(other.external_data), pooled_data(other.pooled_data), precision(other.precision), + dynamic_dims(std::move(other.dynamic_dims)), + pooled_byte_size(other.pooled_byte_size), group_size(other.group_size), num_groups(other.num_groups), - scales_data(other.scales_data), - owned_scales(std::move(other.owned_scales)), - packed_int4_data(other.packed_int4_data), - packed_int4_size(other.packed_int4_size), activation_scales_data(other.activation_scales_data), owned_activation_scales(std::move(other.owned_activation_scales)), - num_rows_for_activation_scales(other.num_rows_for_activation_scales) { + num_rows_for_activation_scales(other.num_rows_for_activation_scales), + cq_codebook(other.cq_codebook), + cq_input_scale(other.cq_input_scale), + cq_input_scale_recip(other.cq_input_scale_recip), + cq_norms(other.cq_norms), + cq_left_signs(other.cq_left_signs), + cq_right_signs(other.cq_right_signs), + cq_permutation(other.cq_permutation), + cq_rotation(other.cq_rotation), + cq_flags(other.cq_flags) { other.total_size = 0; other.byte_size = 0; other.external_data = nullptr; other.pooled_data = nullptr; + other.pooled_byte_size = 0; other.group_size = 0; other.num_groups = 0; - other.scales_data = nullptr; - other.packed_int4_data = nullptr; - other.packed_int4_size = 0; other.activation_scales_data = nullptr; other.num_rows_for_activation_scales = 0; + other.cq_codebook = nullptr; + other.cq_input_scale = nullptr; + other.cq_input_scale_recip = nullptr; + other.cq_norms = nullptr; + other.cq_left_signs = nullptr; + other.cq_right_signs = nullptr; + other.cq_permutation = nullptr; + other.cq_rotation = nullptr; + other.cq_flags = 0; } BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept { if (this != &other) { - if (pooled_data) { - delete[] pooled_data; - } + delete[] pooled_data; shape = std::move(other.shape); total_size = other.total_size; @@ -118,27 +129,41 @@ BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept { external_data = other.external_data; pooled_data = other.pooled_data; precision = other.precision; + dynamic_dims = std::move(other.dynamic_dims); + pooled_byte_size = other.pooled_byte_size; group_size = other.group_size; num_groups = other.num_groups; - scales_data = other.scales_data; - owned_scales = std::move(other.owned_scales); - packed_int4_data = other.packed_int4_data; - packed_int4_size = other.packed_int4_size; activation_scales_data = other.activation_scales_data; owned_activation_scales = std::move(other.owned_activation_scales); num_rows_for_activation_scales = other.num_rows_for_activation_scales; + cq_codebook = other.cq_codebook; + cq_input_scale = other.cq_input_scale; + cq_input_scale_recip = other.cq_input_scale_recip; + cq_norms = other.cq_norms; + cq_left_signs = other.cq_left_signs; + cq_right_signs = other.cq_right_signs; + cq_permutation = other.cq_permutation; + cq_rotation = other.cq_rotation; + cq_flags = other.cq_flags; other.total_size = 0; other.byte_size = 0; other.external_data = nullptr; other.pooled_data = nullptr; + other.pooled_byte_size = 0; other.group_size = 0; other.num_groups = 0; - other.scales_data = nullptr; - other.packed_int4_data = nullptr; - other.packed_int4_size = 0; other.activation_scales_data = nullptr; other.num_rows_for_activation_scales = 0; + other.cq_codebook = nullptr; + other.cq_input_scale = nullptr; + other.cq_input_scale_recip = nullptr; + other.cq_norms = nullptr; + other.cq_left_signs = nullptr; + other.cq_right_signs = nullptr; + other.cq_permutation = nullptr; + other.cq_rotation = nullptr; + other.cq_flags = 0; } return *this; } @@ -164,32 +189,57 @@ void BufferDesc::allocate() { void BufferDesc::allocate_from_pool(BufferPool& pool) { if (!data && !external_data && !pooled_data && byte_size > 0) { pooled_data = pool.acquire(byte_size); + pooled_byte_size = byte_size; } } void BufferDesc::release_to_pool(BufferPool& pool) { - if (pooled_data && byte_size > 0) { - pool.release(pooled_data, byte_size); + if (pooled_data) { + pool.release(pooled_data, pooled_byte_size > 0 ? pooled_byte_size : byte_size); pooled_data = nullptr; + pooled_byte_size = 0; + } +} + +void BufferDesc::set_shape(const std::vector& new_shape) { + shape = new_shape; + total_size = 1; + for (size_t dim : shape) total_size *= dim; + byte_size = PrecisionTraits::packed_size_of(precision, total_size); +} + +void BufferDesc::resize_from_pool(BufferPool& pool) { + if (data || external_data) return; + if (pooled_data && byte_size == pooled_byte_size) return; + if (pooled_data) release_to_pool(pool); + if (byte_size > 0) { + pooled_data = pool.acquire(byte_size); + pooled_byte_size = byte_size; } } +void BufferDesc::release_memory(BufferPool& pool) { + release_to_pool(pool); + data.reset(); + external_data = nullptr; +} + void BufferDesc::set_external(void* ptr) { external_data = ptr; data.reset(); + if (pooled_data) { + delete[] pooled_data; + } pooled_data = nullptr; } -// GraphNode implementation GraphNode::GraphNode(size_t node_id, OpType type) : id(node_id), op_type(type) {} -// TensorConfig implementation TensorConfig& TensorConfig::global() { static TensorConfig instance; return instance; } -// BroadcastInfo implementation BroadcastInfo BroadcastInfo::compute(const std::vector& lhs, const std::vector& rhs) { BroadcastInfo info; size_t max_dims = std::max(lhs.size(), rhs.size()); @@ -234,7 +284,7 @@ namespace ValidationUtils { } } -CactusGraph::CactusGraph() : next_node_id_(0) {} +CactusGraph::CactusGraph() : next_node_id_(1) {} size_t CactusGraph::get_node_count() const { return nodes_.size(); @@ -258,8 +308,35 @@ void CactusGraph::clear_debug_nodes() { void CactusGraph::allocate_buffers() { for (auto& node : nodes_) { - if (node->op_type != OpType::INPUT) { - node->output_buffer.allocate(); + if (node->op_type == OpType::INPUT) continue; + if (node->op_type == OpType::KV_CACHE_STATE + || node->op_type == OpType::CONV_CACHE_STATE + || node->op_type == OpType::RECURRENT_CACHE_STATE) continue; + node->output_buffer.allocate(); + } +} + +void CactusGraph::release_runtime_buffers() { + for (auto& node : nodes_) { + if (node->op_type == OpType::INPUT) continue; + if (node->op_type == OpType::KV_CACHE_STATE + || node->op_type == OpType::CONV_CACHE_STATE + || node->op_type == OpType::RECURRENT_CACHE_STATE) continue; + if (persistent_node_ids_.count(node->id)) continue; + node->output_buffer.release_memory(buffer_pool_); + } + buffer_pool_.clear(); + shrink_thread_local_buffers(); +} + +void CactusGraph::clear_buffer_pool() { + buffer_pool_.clear(); +} + +void CactusGraph::retain_outputs(const std::vector& node_ids) { + for (int node_id : node_ids) { + if (node_id >= 0) { + retained_output_node_ids_.insert(static_cast(node_id)); } } } diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp new file mode 100644 index 000000000..fc19b2ccd --- /dev/null +++ b/cactus-graph/src/execute.cpp @@ -0,0 +1,2264 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static int g_selected_backend = -1; + +ComputeBackend cactus_default_backend() { + if (g_selected_backend >= 0) return static_cast(g_selected_backend); + if (cactus_metal_available()) return ComputeBackend::METAL; + return ComputeBackend::CPU; +} + +int cactus_backend_select(const char* backend) { + if (!backend) return -1; + ComputeBackend selected; + if (std::strcmp(backend, "cpu") == 0) selected = ComputeBackend::CPU; + else if (std::strcmp(backend, "metal") == 0 && cactus_metal_available()) selected = ComputeBackend::METAL; + else return -1; + g_selected_backend = static_cast(selected); + return 0; +} + + +using ComputeFn = void(*)(GraphNode&, const nodes_vector&, const node_index_map_t&); + +#define DECLARE_COMPUTE(name) \ + extern void name(GraphNode&, const nodes_vector&, const node_index_map_t&) + +DECLARE_COMPUTE(compute_binary_op_node); +DECLARE_COMPUTE(compute_unary_op_node); +DECLARE_COMPUTE(compute_activation_node); +DECLARE_COMPUTE(compute_reduce_node); +DECLARE_COMPUTE(compute_reshape_node); +DECLARE_COMPUTE(compute_precision_cast_node); +DECLARE_COMPUTE(compute_matmul_node); +DECLARE_COMPUTE(compute_rms_norm_node); +DECLARE_COMPUTE(compute_rope_node); +DECLARE_COMPUTE(compute_softmax_node); +DECLARE_COMPUTE(compute_attention_node); +DECLARE_COMPUTE(compute_attention_int8_hybrid_node); +DECLARE_COMPUTE(compute_rel_pos_bias_node); +DECLARE_COMPUTE(compute_layernorm_node); +DECLARE_COMPUTE(compute_conv1d_causal_node); +DECLARE_COMPUTE(compute_conv1d_k3_node); +DECLARE_COMPUTE(compute_conv1d_k7s3_node); +DECLARE_COMPUTE(compute_conv1d_node); +DECLARE_COMPUTE(compute_conv1d_same_depthwise_k9_node); +DECLARE_COMPUTE(compute_conv1d_pointwise_node); +DECLARE_COMPUTE(compute_conv2d_k3s2p1_node); +DECLARE_COMPUTE(compute_conv2d_depthwise_k3s2p1_node); +DECLARE_COMPUTE(compute_conv2d_pointwise_1x1_node); +DECLARE_COMPUTE(compute_glu_node); +DECLARE_COMPUTE(compute_batchnorm_node); +DECLARE_COMPUTE(compute_groupnorm_node); +DECLARE_COMPUTE(compute_rope_gptj_node); +DECLARE_COMPUTE(compute_lstm_cell_node); +DECLARE_COMPUTE(compute_gated_deltanet_decode_node); +DECLARE_COMPUTE(compute_gated_deltanet_prefill_node); +DECLARE_COMPUTE(compute_stft_node); +DECLARE_COMPUTE(compute_altup_predict_node); +DECLARE_COMPUTE(compute_altup_correct_node); +DECLARE_COMPUTE(compute_gaussian_topk_node); +DECLARE_COMPUTE(compute_maxpool1d_node); +DECLARE_COMPUTE(compute_bilstm_sequence_node); +DECLARE_COMPUTE(compute_conv2d_k3s1p1_node); +DECLARE_COMPUTE(compute_stats_pool_node); +DECLARE_COMPUTE(compute_weighted_stats_pool_node); +DECLARE_COMPUTE(compute_transpose_node); +DECLARE_COMPUTE(compute_gather_node); +DECLARE_COMPUTE(compute_slice_node); +DECLARE_COMPUTE(compute_embedding_node); +DECLARE_COMPUTE(compute_concat_node); +DECLARE_COMPUTE(compute_cat_node); +DECLARE_COMPUTE(compute_index_node); +DECLARE_COMPUTE(compute_bilinear_interpolation_node); +DECLARE_COMPUTE(compute_sample_node); +DECLARE_COMPUTE(compute_topk_node); +DECLARE_COMPUTE(compute_scatter_topk_node); +DECLARE_COMPUTE(compute_moe_layer_node); +DECLARE_COMPUTE(compute_dense_mlp_tq_fused_node); +DECLARE_COMPUTE(compute_persistent_node); +DECLARE_COMPUTE(compute_kv_cache_state_node); +DECLARE_COMPUTE(compute_kv_cache_append_node); +DECLARE_COMPUTE(compute_attention_cached_node); +DECLARE_COMPUTE(compute_conv_cache_state_node); +DECLARE_COMPUTE(compute_conv_cache_append_node); +DECLARE_COMPUTE(compute_recurrent_cache_state_node); +DECLARE_COMPUTE(compute_recurrent_cache_write_node); +DECLARE_COMPUTE(compute_conv_cache_initialize_node); +DECLARE_COMPUTE(compute_image_preprocess_node); +DECLARE_COMPUTE(compute_rfft_node); +DECLARE_COMPUTE(compute_irfft_node); +DECLARE_COMPUTE(compute_mel_filter_bank_node); +DECLARE_COMPUTE(compute_spectrogram_node); +extern void shrink_thread_local_buffers(); +#undef DECLARE_COMPUTE + +static constexpr int OP_TYPE_COUNT = static_cast(OpType::CONV_CACHE_INITIALIZE) + 1; +static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow"); +static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {}; + +static bool init_dispatch() { + dispatch_flat[static_cast(OpType::ADD)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::ADD_CLIPPED)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::SUBTRACT)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::MULTIPLY)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::DIVIDE)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::NOT_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_ADD)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_SUBTRACT)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_MULTIPLY)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_DIVIDE)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_NOT_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_EXP)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_SQRT)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_COS)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_SIN)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_LOG)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::ABS)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::POW)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::RELU)] = compute_activation_node; + dispatch_flat[static_cast(OpType::SILU)] = compute_activation_node; + dispatch_flat[static_cast(OpType::GELU)] = compute_activation_node; + dispatch_flat[static_cast(OpType::GELU_ERF)] = compute_activation_node; + dispatch_flat[static_cast(OpType::SIGMOID)] = compute_activation_node; + dispatch_flat[static_cast(OpType::TANH)] = compute_activation_node; + dispatch_flat[static_cast(OpType::LEAKY_RELU)] = compute_activation_node; + dispatch_flat[static_cast(OpType::CLAMP)] = compute_activation_node; + dispatch_flat[static_cast(OpType::SUM)] = compute_reduce_node; + dispatch_flat[static_cast(OpType::MEAN)] = compute_reduce_node; + dispatch_flat[static_cast(OpType::VARIANCE)] = compute_reduce_node; + dispatch_flat[static_cast(OpType::MIN)] = compute_reduce_node; + dispatch_flat[static_cast(OpType::MAX)] = compute_reduce_node; + dispatch_flat[static_cast(OpType::CUMSUM)] = compute_reduce_node; + dispatch_flat[static_cast(OpType::FLATTEN)] = compute_reshape_node; + dispatch_flat[static_cast(OpType::VIEW)] = compute_reshape_node; + dispatch_flat[static_cast(OpType::RESHAPE)] = compute_reshape_node; + dispatch_flat[static_cast(OpType::PRECISION_CAST)] = compute_precision_cast_node; + dispatch_flat[static_cast(OpType::MATMUL)] = compute_matmul_node; + dispatch_flat[static_cast(OpType::RMS_NORM)] = compute_rms_norm_node; + dispatch_flat[static_cast(OpType::LAYERNORM)] = compute_layernorm_node; + dispatch_flat[static_cast(OpType::GROUPNORM)] = compute_groupnorm_node; + dispatch_flat[static_cast(OpType::BATCHNORM)] = compute_batchnorm_node; + dispatch_flat[static_cast(OpType::ROPE)] = compute_rope_node; + dispatch_flat[static_cast(OpType::ROPE_GPTJ)] = compute_rope_gptj_node; + dispatch_flat[static_cast(OpType::SOFTMAX)] = compute_softmax_node; + dispatch_flat[static_cast(OpType::ATTENTION)] = compute_attention_node; + dispatch_flat[static_cast(OpType::ATTENTION_INT8_HYBRID)] = compute_attention_int8_hybrid_node; + dispatch_flat[static_cast(OpType::REL_POS_BIAS)] = compute_rel_pos_bias_node; + dispatch_flat[static_cast(OpType::CONV1D_CAUSAL)] = compute_conv1d_causal_node; + dispatch_flat[static_cast(OpType::CONV1D_K3)] = compute_conv1d_k3_node; + dispatch_flat[static_cast(OpType::CONV1D_K7S3)] = compute_conv1d_k7s3_node; + dispatch_flat[static_cast(OpType::CONV1D)] = compute_conv1d_node; + dispatch_flat[static_cast(OpType::CONV1D_SAME_DEPTHWISE_K9)] = compute_conv1d_same_depthwise_k9_node; + dispatch_flat[static_cast(OpType::CONV1D_POINTWISE)] = compute_conv1d_pointwise_node; + dispatch_flat[static_cast(OpType::CONV2D_K3S2P1)] = compute_conv2d_k3s2p1_node; + dispatch_flat[static_cast(OpType::CONV2D_DEPTHWISE_K3S2P1)] = compute_conv2d_depthwise_k3s2p1_node; + dispatch_flat[static_cast(OpType::CONV2D_POINTWISE_1X1)] = compute_conv2d_pointwise_1x1_node; + dispatch_flat[static_cast(OpType::CONV2D_K3S1P1)] = compute_conv2d_k3s1p1_node; + dispatch_flat[static_cast(OpType::GLU)] = compute_glu_node; + dispatch_flat[static_cast(OpType::TRANSPOSE)] = compute_transpose_node; + dispatch_flat[static_cast(OpType::GATHER)] = compute_gather_node; + dispatch_flat[static_cast(OpType::SLICE)] = compute_slice_node; + dispatch_flat[static_cast(OpType::EMBEDDING)] = compute_embedding_node; + dispatch_flat[static_cast(OpType::CONCAT)] = compute_concat_node; + dispatch_flat[static_cast(OpType::CAT)] = compute_cat_node; + dispatch_flat[static_cast(OpType::INDEX)] = compute_index_node; + dispatch_flat[static_cast(OpType::BILINEAR_INTERPOLATION)] = compute_bilinear_interpolation_node; + dispatch_flat[static_cast(OpType::SAMPLE)] = compute_sample_node; + dispatch_flat[static_cast(OpType::TOPK)] = compute_topk_node; + dispatch_flat[static_cast(OpType::SCATTER_TOPK)] = compute_scatter_topk_node; + dispatch_flat[static_cast(OpType::MOE_LAYER)] = compute_moe_layer_node; + dispatch_flat[static_cast(OpType::DENSE_MLP_TQ_FUSED)] = compute_dense_mlp_tq_fused_node; + dispatch_flat[static_cast(OpType::PERSISTENT)] = compute_persistent_node; + dispatch_flat[static_cast(OpType::LSTM_CELL)] = compute_lstm_cell_node; + dispatch_flat[static_cast(OpType::GATED_DELTANET_DECODE)] = compute_gated_deltanet_decode_node; + dispatch_flat[static_cast(OpType::GATED_DELTANET_PREFILL)] = compute_gated_deltanet_prefill_node; + dispatch_flat[static_cast(OpType::STFT)] = compute_stft_node; + dispatch_flat[static_cast(OpType::ALTUP_PREDICT)] = compute_altup_predict_node; + dispatch_flat[static_cast(OpType::ALTUP_CORRECT)] = compute_altup_correct_node; + dispatch_flat[static_cast(OpType::GAUSSIAN_TOPK)] = compute_gaussian_topk_node; + dispatch_flat[static_cast(OpType::MAXPOOL1D)] = compute_maxpool1d_node; + dispatch_flat[static_cast(OpType::BILSTM_SEQUENCE)] = compute_bilstm_sequence_node; + dispatch_flat[static_cast(OpType::STATS_POOL)] = compute_stats_pool_node; + dispatch_flat[static_cast(OpType::WEIGHTED_STATS_POOL)] = compute_weighted_stats_pool_node; + dispatch_flat[static_cast(OpType::KV_CACHE_STATE)] = compute_kv_cache_state_node; + dispatch_flat[static_cast(OpType::KV_CACHE_APPEND)] = compute_kv_cache_append_node; + dispatch_flat[static_cast(OpType::ATTENTION_CACHED)] = compute_attention_cached_node; + dispatch_flat[static_cast(OpType::CONV_CACHE_STATE)] = compute_conv_cache_state_node; + dispatch_flat[static_cast(OpType::CONV_CACHE_APPEND)] = compute_conv_cache_append_node; + dispatch_flat[static_cast(OpType::RECURRENT_CACHE_STATE)] = compute_recurrent_cache_state_node; + dispatch_flat[static_cast(OpType::RECURRENT_CACHE_WRITE)] = compute_recurrent_cache_write_node; + dispatch_flat[static_cast(OpType::CONV_CACHE_INITIALIZE)] = compute_conv_cache_initialize_node; + dispatch_flat[static_cast(OpType::IMAGE_PREPROCESS)] = compute_image_preprocess_node; + dispatch_flat[static_cast(OpType::RFFT)] = compute_rfft_node; + dispatch_flat[static_cast(OpType::IRFFT)] = compute_irfft_node; + dispatch_flat[static_cast(OpType::MEL_FILTER_BANK)] = compute_mel_filter_bank_node; + dispatch_flat[static_cast(OpType::SPECTROGRAM)] = compute_spectrogram_node; + return true; +} + +static const bool dispatch_initialized = init_dispatch(); + +static inline void dispatch_node(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& node_index_map) { + int op = static_cast(node.op_type); + ComputeFn fn = dispatch_flat[op]; + if (fn) { + fn(node, nodes, node_index_map); + } else { + throw std::runtime_error("Unknown operation type: " + std::to_string(op)); + } +} + +static const char* op_type_names[] = { + "INPUT", "PRECISION_CAST", + "ADD", "ADD_CLIPPED", "SUBTRACT", "MULTIPLY", "DIVIDE", + "ABS", "POW", "FLATTEN", "VIEW", + "MATMUL", "TRANSPOSE", "RESHAPE", "SLICE", "GATHER", "EMBEDDING", + "BILINEAR_INTERPOLATION", + "SUM", "MEAN", "VARIANCE", "MIN", "MAX", "CUMSUM", + "RMS_NORM", "ROPE", "ROPE_GPTJ", "SOFTMAX", + "ATTENTION", "ATTENTION_INT8_HYBRID", "REL_POS_BIAS", + "CONV1D_CAUSAL", "CONV1D_K3", "CONV1D_K7S3", "CONV1D", + "CONV1D_SAME_DEPTHWISE_K9", "CONV1D_POINTWISE", + "CONV2D_K3S2P1", "CONV2D_DEPTHWISE_K3S2P1", "CONV2D_POINTWISE_1X1", + "GLU", "BATCHNORM", + "SCALAR_ADD", "SCALAR_SUBTRACT", "SCALAR_MULTIPLY", "SCALAR_DIVIDE", + "SCALAR_EXP", "SCALAR_SQRT", "SCALAR_COS", "SCALAR_SIN", "SCALAR_LOG", + "RELU", "SILU", "GELU", "GELU_ERF", "SIGMOID", "TANH", + "SAMPLE", "CONCAT", "CAT", + "SCATTER_TOPK", "TOPK", "LAYERNORM", "GROUPNORM", + "MOE_LAYER", "INDEX", "PERSISTENT", + "LSTM_CELL", "GATED_DELTANET_DECODE", "GATED_DELTANET_PREFILL", + "STFT", "ALTUP_PREDICT", "ALTUP_CORRECT", "GAUSSIAN_TOPK", + "MAXPOOL1D", "BILSTM_SEQUENCE", "LEAKY_RELU", + "CONV2D_K3S1P1", "STATS_POOL", "WEIGHTED_STATS_POOL", + "KV_CACHE_STATE", "KV_CACHE_APPEND", "ATTENTION_CACHED", + "CONV_CACHE_STATE", "CONV_CACHE_APPEND", + "RFFT", "IRFFT", "MEL_FILTER_BANK", "SPECTROGRAM", + "IMAGE_PREPROCESS", "CLAMP", "DENSE_MLP_TQ_FUSED", + "NOT_EQUAL", "SCALAR_NOT_EQUAL", + "RECURRENT_CACHE_STATE", + "RECURRENT_CACHE_WRITE", + "CONV_CACHE_INITIALIZE" +}; + +static const char* get_op_name(OpType op) { + return op_type_names[static_cast(op)]; +} + +void compute_node_optimized(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& node_index_map) { + if (node.op_type == OpType::INPUT) return; + dispatch_node(node, nodes, node_index_map); +} + +void CactusGraph::set_input(size_t node_id, const void* data, Precision) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) { + throw std::out_of_range("Unknown input node id: " + std::to_string(node_id)); + } + + auto& node = *nodes_[it->second]; + if (node.op_type != OpType::INPUT) { + throw std::invalid_argument("Can only set data on input nodes"); + } + + if (!node.output_buffer.data && !node.output_buffer.external_data) { + node.output_buffer.allocate(); + } + + if (node.output_buffer.external_data) { + node.output_buffer.external_data = nullptr; + node.output_buffer.allocate(); + } + + std::memcpy(node.output_buffer.get_data(), data, node.output_buffer.byte_size); +} + +void CactusGraph::set_external_input(size_t node_id, void* data, Precision) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) { + throw std::out_of_range("Unknown input node id: " + std::to_string(node_id)); + } + + auto& node = *nodes_[it->second]; + if (node.op_type != OpType::INPUT) { + throw std::invalid_argument("Can only set data on input nodes"); + } + + node.output_buffer.set_external(data); + embedded_input_node_ids_.erase(node_id); +} + +void* CactusGraph::get_output(size_t node_id) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) { + throw std::out_of_range("Unknown output node id: " + std::to_string(node_id)); + } + + auto& buffer = nodes_[it->second]->output_buffer; + if (!buffer.get_data()) { + buffer.allocate(); + } + return buffer.get_data(); +} + +static bool check_debug_env() { + const char* v1 = std::getenv("CACTUS_CAPTURE_ENABLE"); + const char* v2 = std::getenv("CACTUS_CAPTURE_STDOUT"); + const char* v3 = std::getenv("CACTUS_CAPTURE_FILE"); + const char* v4 = std::getenv("CACTUS_CAPTURE_DIR"); + const char* v5 = std::getenv("CACTUS_PROFILE_FILE"); + const char* v6 = std::getenv("CACTUS_PROFILE"); + return (v1 && v1[0] != '0') || (v2 && v2[0] != '0') || + (v3 && v3[0]) || (v4 && v4[0]) || (v5 && v5[0]) || (v6 && v6[0]); +} + +namespace { +std::vector infer_output_shape(const GraphNode& node, const nodes_vector& nodes, const node_index_map_t& idx) { + auto in = [&](size_t i) -> const std::vector& { return get_input(node, i, nodes, idx).shape; }; + switch (node.op_type) { + case OpType::MATMUL: { + const auto& lhs = in(0); + const auto& rhs = in(1); + std::vector out = lhs; + out.back() = node.params.pretransposed_rhs ? rhs[rhs.size() - 2] : rhs[rhs.size() - 1]; + return out; + } + case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: + case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + return BroadcastInfo::compute(in(0), in(1)).output_shape; + case OpType::ATTENTION: case OpType::ATTENTION_CACHED: case OpType::ATTENTION_INT8_HYBRID: { + std::vector out = in(0); + if (node.params.v_head_dim > 0) out.back() = node.params.v_head_dim; + return out; + } + case OpType::TRANSPOSE: { + const auto& x = in(0); + const auto& perm = node.params.permutation; + if (perm.size() != x.size()) return x; + std::vector out(x.size()); + for (size_t i = 0; i < perm.size(); ++i) out[i] = x[perm[i]]; + return out; + } + case OpType::RESHAPE: case OpType::VIEW: case OpType::FLATTEN: { + const auto& x = in(0); + std::vector out = node.params.new_shape; + if (out.empty()) return x; + size_t in_total = 1; + for (size_t d : x) in_total *= d; + size_t rest = 1; + for (size_t i = 1; i < out.size(); ++i) rest *= out[i]; + if (rest > 0 && in_total % rest == 0) out[0] = in_total / rest; + return out; + } + case OpType::CONCAT: case OpType::CAT: { + std::vector out = in(0); + size_t axis = node.params.axis < 0 ? out.size() + static_cast(node.params.axis) + : static_cast(node.params.axis); + size_t sum = 0; + for (size_t i = 0; i < node.input_ids.size(); ++i) sum += in(i)[axis]; + out[axis] = sum; + return out; + } + case OpType::SLICE: { + std::vector out = in(0); + size_t axis = node.params.axis < 0 ? out.size() + static_cast(node.params.axis) + : static_cast(node.params.axis); + out[axis] = node.params.slice_length; + return out; + } + default: { + std::vector out = node.output_buffer.shape; + if (out.empty()) return in(0); + for (size_t i = 0; i < node.input_ids.size(); ++i) { + const auto& inp = get_input(node, i, nodes, idx); + if (inp.has_dynamic_dims() && !inp.shape.empty()) { + out[0] = inp.shape[0]; + return out; + } + } + return out; + } + } +} + +bool skip_shape_infer(OpType op) { + switch (op) { + case OpType::KV_CACHE_STATE: case OpType::KV_CACHE_APPEND: + case OpType::CONV_CACHE_STATE: case OpType::CONV_CACHE_APPEND: case OpType::CONV_CACHE_INITIALIZE: + case OpType::RECURRENT_CACHE_STATE: case OpType::RECURRENT_CACHE_WRITE: case OpType::PERSISTENT: + return true; + default: + return false; + } +} +} + +void CactusGraph::set_runtime_input_shape(size_t node_id, const std::vector& shape) { + GraphNode& node = *nodes_[node_index_map_.at(node_id)]; + node.output_buffer.set_shape(shape); + node.output_buffer.dynamic_dims.assign(shape.size(), 1); + node.output_buffer.data.reset(); + has_dynamic_shapes_ = true; + runtime_shapes_dirty_ = true; +} + +void CactusGraph::set_input_dynamic_dims(size_t node_id, const std::vector& dynamic_dims) { + GraphNode& node = *nodes_[node_index_map_.at(node_id)]; + node.output_buffer.dynamic_dims = dynamic_dims; + if (!dynamic_dims.empty()) has_dynamic_shapes_ = true; +} + +void CactusGraph::infer_shapes() { + if (!has_dynamic_shapes_ || !runtime_shapes_dirty_) return; + for (auto& np : nodes_) { + GraphNode& node = *np; + if (node.op_type == OpType::INPUT || persistent_node_ids_.count(node.id) || skip_shape_infer(node.op_type)) continue; + bool dyn = false; + for (size_t i = 0; i < node.input_ids.size() && !dyn; ++i) { + dyn = get_input(node, i, nodes_, node_index_map_).has_dynamic_dims(); + } + if (!dyn) continue; + node.output_buffer.set_shape(infer_output_shape(node, nodes_, node_index_map_)); + switch (node.op_type) { + case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: + case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + node.params.broadcast_info = BroadcastInfo::compute( + get_input(node, 0, nodes_, node_index_map_).shape, + get_input(node, 1, nodes_, node_index_map_).shape); + break; + default: + break; + } + node.output_buffer.dynamic_dims.assign(node.output_buffer.shape.size(), 1); + } + runtime_shapes_dirty_ = false; +} + +static void row_strides(const std::vector& shape, size_t* out) { + size_t s=1; for (int k=(int)shape.size()-1; k>=0; --k){ out[k]=s; s*=shape[k]; } +} +static void bcast_strides(const std::vector& in_shape, const std::vector& out_shape, uint32_t* out) { + size_t off = out_shape.size() - in_shape.size(); + size_t istr[8]; row_strides(in_shape, istr); + for (size_t d=0; d dequantize_int8_weights_to_fp16(const BufferDesc& W, size_t rows, size_t cols, + const char* op_name); + +static const __fp16* metal_conv_weight_f16(const BufferDesc& w, size_t rows, size_t cols) { + if (w.precision == Precision::FP16) return w.data_as<__fp16>(); + if (w.precision != Precision::INT8) return nullptr; + const int8_t* p = w.data_as(); + size_t n = rows * cols; + if (!p || n == 0 || w.total_size < n) return nullptr; + uint64_t h = 1469598103934665603ull; + size_t take = n < 64 ? n : 64; + for (size_t i = 0; i < take; ++i) { h ^= (uint8_t)p[i]; h *= 1099511628211ull; } + for (size_t i = n - take; i < n; ++i) { h ^= (uint8_t)p[i]; h *= 1099511628211ull; } + h ^= n; h *= 1099511628211ull; + struct Entry { std::vector<__fp16> data; uint64_t fp; }; + static std::mutex mu; + static std::unordered_map cache; + std::lock_guard lk(mu); + Entry& e = cache[p]; + if (e.data.size() != n || e.fp != h) { + e.data = dequantize_int8_weights_to_fp16(w, rows, cols, "conv_metal"); + e.fp = h; + } + return e.data.data(); +} + +static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& map) { + if (node.params.backend != ComputeBackend::METAL) return false; + BufferDesc& out = node.output_buffer; + auto fp16 = [](const BufferDesc& b){ return b.precision == Precision::FP16; }; + switch (node.op_type) { + case OpType::MATMUL: { + const auto& lhs = get_input(node, 0, nodes, map); + const auto& rhs = get_input(node, 1, nodes, map); + size_t M = lhs.shape[lhs.shape.size() - 2]; + if (!fp16(lhs)) return false; + if (PrecisionTraits::is_cq(rhs.precision) && rhs.group_size > 0) { + CactusQuantMatrix mat = rhs.to_cq_matrix(); + if (M == 1 && !cactus_graph_prefill_consistent()) + return cactus_metal_encode_quant_matmul(out.get_data(), lhs.get_data(), &mat); + return cactus_metal_encode_quant_matmul_m(out.get_data(), lhs.get_data(), &mat, (uint32_t)M); + } + if (!fp16(rhs) || !fp16(out) || rhs.shape.size() != 2) return false; + size_t K = lhs.shape.back(); + size_t N = node.params.pretransposed_rhs ? rhs.shape[0] : rhs.shape[1]; + size_t rhs_k = node.params.pretransposed_rhs ? rhs.shape[1] : rhs.shape[0]; + if (K == 0 || rhs_k != K || lhs.total_size % K != 0) return false; + size_t rows = lhs.total_size / K; + if (rows != M) { + size_t batch = rows / M; + if (batch * M != rows) return false; + } + return cactus_metal_encode_gemm_f16(out.get_data(), lhs.get_data(), rhs.get_data(), + (uint32_t)rows, (uint32_t)K, (uint32_t)N, node.params.pretransposed_rhs ? 1 : 0); + } + case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: + case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: { + const auto& a = get_input(node, 0, nodes, map); + const auto& b = get_input(node, 1, nodes, map); + int code = node.op_type==OpType::ADD?0: node.op_type==OpType::ADD_CLIPPED?1: + node.op_type==OpType::SUBTRACT?2: node.op_type==OpType::MULTIPLY?3: + node.op_type==OpType::DIVIDE?4:5; + bool f32 = a.precision == Precision::FP32 && b.precision == Precision::FP32 + && out.precision == Precision::FP32; + if (f32 && a.total_size == out.total_size && b.total_size == out.total_size) + return cactus_metal_encode_binary_f32(code, out.get_data(), a.get_data(), b.get_data(), out.total_size); + if (!fp16(a) || !fp16(b) || !fp16(out)) return false; + if (a.total_size == out.total_size && b.total_size == out.total_size) + return cactus_metal_encode_binary(code, out.get_data(), a.get_data(), b.get_data(), out.total_size); + const auto& osh = out.shape; + uint32_t nd = (uint32_t)osh.size(); + if (nd == 0 || nd > 8 || a.shape.size() > nd || b.shape.size() > nd) return false; + uint32_t oshape[8], astr[8], bstr[8]; + for (uint32_t d=0; d 8 || in.shape.size() != nd || out.shape.size() != nd) return false; + bool tailswap = nd >= 2; + for (uint32_t d = 0; d + 2 < nd && tailswap; ++d) if (perm[d] != d) tailswap = false; + if (tailswap && (perm[nd-2] != nd-1 || perm[nd-1] != nd-2)) tailswap = false; + if (tailswap) { + size_t batch = 1; + for (uint32_t d = 0; d + 2 < nd; ++d) batch *= in.shape[d]; + return cactus_metal_encode_transpose2d(out.get_data(), in.get_data(), + (uint32_t)batch, (uint32_t)in.shape[nd-2], (uint32_t)in.shape[nd-1]); + } + size_t istr[8]; row_strides(in.shape, istr); + uint32_t oshape[8], sstride[8]; + for (uint32_t d=0; d 8 || axis >= nd || out.shape.size() != nd) return false; + size_t istr[8]; row_strides(in.shape, istr); + uint32_t oshape[8], sstride[8]; + for (uint32_t d=0; d= in.shape.size()) return false; + size_t istr[8]; row_strides(in.shape, istr); + size_t slice = istr[axis], block = istr[axis-1]; + uint32_t oshape[2] = { (uint32_t)(in.total_size/block), (uint32_t)slice }; + uint32_t sstride[2] = { (uint32_t)block, 1u }; + return cactus_metal_encode_strided_copy(out.get_data(), in.get_data(), oshape, sstride, 2, + (uint32_t)out.total_size, (uint32_t)(node.params.index_value * slice), in.byte_size, out.byte_size); + } + case OpType::CAT: { + if (node.input_ids.size() < 2 || out.shape.empty()) return false; + size_t axis = (size_t)node.params.axis; + uint32_t nd = (uint32_t)out.shape.size(); + if (axis >= nd || nd > 8) return false; + for (size_t ii = 0; ii < node.input_ids.size(); ++ii) { + const auto& cin = get_input(node, ii, nodes, map); + if (!fp16(cin) || cin.shape.size() != nd) return false; + } + size_t ostr[8]; row_strides(out.shape, ostr); + size_t axis_off = 0; + for (size_t ii = 0; ii < node.input_ids.size(); ++ii) { + const auto& cin = get_input(node, ii, nodes, map); + uint32_t ishape[8], ostride[8]; + for (uint32_t d=0; d 0 && cin.shape[0] == 1 && out.shape[0] > 1) ? out.shape[0] : 1; + for (size_t r = 0; r < bcast; ++r) { + if (!cactus_metal_encode_strided_scatter(out.get_data(), cin.get_data(), ishape, ostride, nd, + (uint32_t)cin.total_size, (uint32_t)(axis_off * ostr[axis] + r * ostr[0]), cin.byte_size, out.byte_size)) + return false; + } + axis_off += cin.shape[axis]; + } + return true; + } + case OpType::SCALAR_ADD: case OpType::SCALAR_SUBTRACT: + case OpType::SCALAR_MULTIPLY: case OpType::SCALAR_DIVIDE: + case OpType::SCALAR_EXP: case OpType::SCALAR_SQRT: case OpType::SCALAR_COS: + case OpType::SCALAR_SIN: case OpType::SCALAR_LOG: case OpType::ABS: + case OpType::POW: case OpType::SCALAR_NOT_EQUAL: case OpType::LEAKY_RELU: { + const auto& in = get_input(node, 0, nodes, map); + int code; + switch (node.op_type) { + case OpType::SCALAR_ADD: code = 0; break; + case OpType::SCALAR_SUBTRACT: code = 1; break; + case OpType::SCALAR_MULTIPLY: code = 2; break; + case OpType::SCALAR_DIVIDE: code = 3; break; + case OpType::SCALAR_EXP: code = 4; break; + case OpType::SCALAR_SQRT: code = 5; break; + case OpType::SCALAR_COS: code = 6; break; + case OpType::SCALAR_SIN: code = 7; break; + case OpType::SCALAR_LOG: code = 8; break; + case OpType::ABS: code = 9; break; + case OpType::POW: code = 10; break; + case OpType::SCALAR_NOT_EQUAL: code = 11; break; + default: code = 12; break; + } + if (in.precision == Precision::FP32 && out.precision == Precision::FP32) + return cactus_metal_encode_scalar_f32(code, out.get_data(), in.get_data(), out.total_size, node.params.scalar); + if (!fp16(in) || !fp16(out)) return false; + return cactus_metal_encode_scalar(code, out.get_data(), in.get_data(), out.total_size, node.params.scalar); + } + case OpType::GELU: case OpType::TANH: case OpType::SILU: case OpType::RELU: + case OpType::GELU_ERF: case OpType::SIGMOID: { + const auto& in = get_input(node, 0, nodes, map); + int code = node.op_type==OpType::GELU?0: node.op_type==OpType::TANH?1: node.op_type==OpType::SILU?2: + node.op_type==OpType::GELU_ERF?4: node.op_type==OpType::SIGMOID?5:3; + if (in.precision == Precision::FP32 && out.precision == Precision::FP32) + return cactus_metal_encode_unary_f32(code, out.get_data(), in.get_data(), out.total_size); + if (!fp16(in) || !fp16(out)) return false; + return cactus_metal_encode_unary(code, out.get_data(), in.get_data(), out.total_size); + } + case OpType::SUM: case OpType::MEAN: case OpType::VARIANCE: + case OpType::MIN: case OpType::MAX: { + const auto& in = get_input(node, 0, nodes, map); + int axis = node.params.axis; + if (axis < 0 || (size_t)axis >= in.shape.size()) return false; + bool f32 = in.precision == Precision::FP32 && out.precision == Precision::FP32; + if (!f32 && (!fp16(in) || !fp16(out))) return false; + int code = node.op_type==OpType::SUM?0: node.op_type==OpType::MEAN?1: + node.op_type==OpType::VARIANCE?2: node.op_type==OpType::MIN?3:4; + size_t outer = 1, inner = 1; + for (size_t d = 0; d < (size_t)axis; ++d) outer *= in.shape[d]; + for (size_t d = (size_t)axis + 1; d < in.shape.size(); ++d) inner *= in.shape[d]; + return cactus_metal_encode_reduce_axis(code, out.get_data(), in.get_data(), + (uint32_t)outer, (uint32_t)in.shape[(size_t)axis], (uint32_t)inner, f32 ? 1 : 0); + } + case OpType::CUMSUM: { + const auto& in = get_input(node, 0, nodes, map); + int axis = node.params.axis; + if (axis < 0 || (size_t)axis >= in.shape.size()) return false; + bool f32 = in.precision == Precision::FP32 && out.precision == Precision::FP32; + if (!f32 && (!fp16(in) || !fp16(out))) return false; + size_t outer = 1, inner = 1; + for (size_t d = 0; d < (size_t)axis; ++d) outer *= in.shape[d]; + for (size_t d = (size_t)axis + 1; d < in.shape.size(); ++d) inner *= in.shape[d]; + return cactus_metal_encode_cumsum(out.get_data(), in.get_data(), + (uint32_t)outer, (uint32_t)in.shape[(size_t)axis], (uint32_t)inner, f32 ? 1 : 0); + } + case OpType::CONCAT: { + if (node.input_ids.size() != 2) return false; + const auto& a = get_input(node, 0, nodes, map); + const auto& b = get_input(node, 1, nodes, map); + if (!fp16(a) || !fp16(b) || !fp16(out)) return false; + int axis = node.params.axis; + if (axis < 0) axis += (int)a.shape.size(); + if (axis < 0 || (size_t)axis >= a.shape.size() || a.shape.size() != b.shape.size()) return false; + size_t a_outer = 1, b_outer = 1, inner = 1; + for (size_t d = 0; d < (size_t)axis; ++d) { a_outer *= a.shape[d]; b_outer *= b.shape[d]; } + for (size_t d = (size_t)axis + 1; d < a.shape.size(); ++d) inner *= a.shape[d]; + return cactus_metal_encode_concat2(out.get_data(), a.get_data(), b.get_data(), + (uint32_t)a_outer, (uint32_t)b_outer, + (uint32_t)a.shape[(size_t)axis], (uint32_t)b.shape[(size_t)axis], (uint32_t)inner); + } + case OpType::GATHER: { + const auto& t = get_input(node, 0, nodes, map); + const auto& idx = get_input(node, 1, nodes, map); + if (!fp16(t) || !fp16(out) || idx.precision != Precision::FP32) return false; + if (t.shape.size() < 2) return false; + size_t D = t.total_size / t.shape[0]; + return cactus_metal_encode_gather_f32idx(out.get_data(), t.get_data(), idx.get_data(), + (uint32_t)idx.total_size, (uint32_t)D, t.byte_size); + } + case OpType::ROPE: case OpType::ROPE_GPTJ: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || !fp16(out) || in.shape.size() < 4) return false; + size_t S = in.shape[1], H = in.shape[2], D = in.shape[3]; + bool gptj = node.op_type == OpType::ROPE_GPTJ; + size_t rot = gptj ? (size_t)node.params.scalar : D; + if (rot == 0 || rot > D || (rot % 2) != 0) return false; + size_t tokens = in.total_size / D; + return cactus_metal_encode_rope_full(out.get_data(), in.get_data(), + (uint32_t)tokens, (uint32_t)S, (uint32_t)H, (uint32_t)D, (uint32_t)rot, + (uint32_t)node.params.position_offset, node.params.theta, gptj ? 1 : 0); + } + case OpType::MAXPOOL1D: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || !fp16(out) || in.shape.size() != 3 || out.shape.size() != 3) return false; + return cactus_metal_encode_maxpool1d(out.get_data(), in.get_data(), + (uint32_t)(in.shape[0] * in.shape[1]), (uint32_t)in.shape[2], + (uint32_t)out.shape[2], (uint32_t)node.params.kernel_size, (uint32_t)node.params.stride); + } + case OpType::BILINEAR_INTERPOLATION: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || !fp16(out) || in.shape.size() != 2) return false; + size_t src = (size_t)std::sqrt((double)in.shape[0]); + if (src * src != in.shape[0]) return false; + return cactus_metal_encode_bilinear(out.get_data(), in.get_data(), + (uint32_t)src, (uint32_t)src, (uint32_t)node.params.dst_height, + (uint32_t)node.params.dst_width, (uint32_t)in.shape[1], + node.params.align_corners ? 1 : 0); + } + case OpType::CONV1D: case OpType::CONV1D_K7S3: { + if (node.input_ids.size() < 2) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + const BufferDesc* b = node.input_ids.size() > 2 ? &get_input(node, 2, nodes, map) : nullptr; + if (!fp16(x) || !fp16(w) || !fp16(out)) return false; + if (b && b->precision != Precision::FP16) return false; + if (x.shape.size() != 3 || w.shape.size() != 3 || out.shape.size() != 3) return false; + bool ck_co = node.op_type == OpType::CONV1D_K7S3; + size_t K = ck_co ? w.shape[1] : w.shape[2]; + size_t Cout = ck_co ? w.shape[2] : w.shape[0]; + return cactus_metal_encode_conv1d_gen(out.get_data(), x.get_data(), w.get_data(), + b ? b->get_data() : nullptr, (uint32_t)x.shape[0], (uint32_t)x.shape[1], + (uint32_t)x.shape[2], (uint32_t)Cout, (uint32_t)out.shape[2], (uint32_t)K, + (uint32_t)(node.params.stride ? node.params.stride : 1), ck_co ? 1 : 0); + } + case OpType::CONV1D_CAUSAL: case OpType::CONV1D_SAME_DEPTHWISE_K9: { + if (node.input_ids.size() < 2) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + bool causal = node.op_type == OpType::CONV1D_CAUSAL; + const BufferDesc* b = (!causal && node.input_ids.size() > 2) ? &get_input(node, 2, nodes, map) : nullptr; + if (!fp16(x) || !fp16(out)) return false; + if (b && b->precision != Precision::FP16) return false; + if (x.shape.size() != 3) return false; + size_t C = x.shape[2]; + size_t K; + if (causal) { + if (w.shape.size() != 3 || w.shape[0] != C || w.shape[1] != 1) return false; + K = w.shape[2]; + } else { + K = w.shape.back(); + if (w.total_size != C * K) return false; + } + const __fp16* wptr = metal_conv_weight_f16(w, C, K); + if (!wptr) return false; + size_t dil = causal ? (node.params.dilation ? node.params.dilation : 1) : 1; + size_t pad = causal ? (K - 1) * dil : K / 2; + return cactus_metal_encode_conv1d_nlc_dw(out.get_data(), x.get_data(), wptr, + b ? b->get_data() : nullptr, (uint32_t)x.shape[0], (uint32_t)x.shape[1], + (uint32_t)C, (uint32_t)K, (uint32_t)dil, (uint32_t)pad); + } + case OpType::CONV2D_K3S2P1: case OpType::CONV2D_K3S1P1: + case OpType::CONV2D_DEPTHWISE_K3S2P1: case OpType::CONV2D_POINTWISE_1X1: { + if (node.input_ids.size() < 2) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + const BufferDesc* b = node.input_ids.size() > 2 ? &get_input(node, 2, nodes, map) : nullptr; + if (!fp16(x) || !fp16(w) || !fp16(out)) return false; + if (b && b->precision != Precision::FP16) return false; + if (x.shape.size() != 4 || out.shape.size() != 4) return false; + bool dw = node.op_type == OpType::CONV2D_DEPTHWISE_K3S2P1; + bool pw = node.op_type == OpType::CONV2D_POINTWISE_1X1; + uint32_t K = pw ? 1 : 3; + uint32_t stride = (node.op_type == OpType::CONV2D_K3S1P1 || pw) ? 1 : 2; + uint32_t pad = pw ? 0 : 1; + return cactus_metal_encode_conv2d(out.get_data(), x.get_data(), w.get_data(), + b ? b->get_data() : nullptr, (uint32_t)x.shape[0], (uint32_t)x.shape[1], + (uint32_t)x.shape[2], (uint32_t)x.shape[3], (uint32_t)out.shape[1], + (uint32_t)out.shape[2], (uint32_t)out.shape[3], K, stride, pad, dw ? 1 : 0); + } + case OpType::CONV1D_POINTWISE: { + if (node.input_ids.size() < 2) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + const BufferDesc* b = node.input_ids.size() > 2 ? &get_input(node, 2, nodes, map) : nullptr; + if (!fp16(x) || !fp16(out)) return false; + if (b && b->precision != Precision::FP16) return false; + if (x.shape.size() != 3 || out.shape.size() != 3) return false; + size_t Cin = x.shape[2], Cout = out.shape[2]; + if (w.total_size != Cin * Cout) return false; + const __fp16* wptr = metal_conv_weight_f16(w, Cout, Cin); + if (!wptr) return false; + size_t rows = x.shape[0] * x.shape[1]; + if (!cactus_metal_encode_gemm_f16(out.get_data(), x.get_data(), wptr, + (uint32_t)rows, (uint32_t)Cin, (uint32_t)Cout, 1)) return false; + if (b) return cactus_metal_encode_bias_add_rows(out.get_data(), b->get_data(), + (uint32_t)Cout, (uint32_t)out.total_size); + return true; + } + case OpType::REL_POS_BIAS: { + if (node.input_ids.size() != 2) return false; + const auto& q = get_input(node, 0, nodes, map); + const auto& r = get_input(node, 1, nodes, map); + if (!fp16(q) || !fp16(r) || !fp16(out)) return false; + if (q.shape.size() != 4 || r.shape.size() != 4) return false; + size_t B = q.shape[0], T = q.shape[1], H = q.shape[2], D = q.shape[3]; + size_t Rb = r.shape[0], R = r.shape[1]; + if (B == 0 || T == 0 || (Rb != 1 && Rb != B)) return false; + if (r.shape[2] != H || r.shape[3] != D) return false; + if (R < 2 * T - 1 || out.total_size != B * H * T * T) return false; + return cactus_metal_encode_rel_pos_bias(out.get_data(), q.get_data(), r.get_data(), + (uint32_t)B, (uint32_t)T, (uint32_t)H, (uint32_t)D, (uint32_t)R, + Rb == 1 ? 0 : 1, node.params.scale); + } + case OpType::BATCHNORM: { + if (node.input_ids.size() != 5) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + const auto& b = get_input(node, 2, nodes, map); + const auto& rm = get_input(node, 3, nodes, map); + const auto& rv = get_input(node, 4, nodes, map); + if (!fp16(x) || !fp16(out) || !fp16(w) || !fp16(b) || !fp16(rm) || !fp16(rv)) return false; + int axis = node.params.axis; + if (axis < 0) axis += (int)x.shape.size(); + if (axis < 0 || (size_t)axis >= x.shape.size()) return false; + size_t inner = 1; + for (size_t d = (size_t)axis + 1; d < x.shape.size(); ++d) inner *= x.shape[d]; + return cactus_metal_encode_batchnorm(out.get_data(), x.get_data(), w.get_data(), + b.get_data(), rm.get_data(), rv.get_data(), (uint32_t)x.shape[(size_t)axis], + (uint32_t)inner, (uint32_t)x.total_size, node.params.epsilon); + } + case OpType::GROUPNORM: { + if (node.input_ids.size() < 3) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + const auto& b = get_input(node, 2, nodes, map); + if (!fp16(x) || !fp16(out) || !fp16(w) || !fp16(b) || x.shape.size() < 2) return false; + size_t spatial = 1; + for (size_t d = 2; d < x.shape.size(); ++d) spatial *= x.shape[d]; + size_t groups = node.params.num_groups ? node.params.num_groups : 32; + return cactus_metal_encode_groupnorm(out.get_data(), x.get_data(), w.get_data(), + b.get_data(), (uint32_t)x.shape[0], (uint32_t)x.shape[1], (uint32_t)spatial, + (uint32_t)groups, node.params.epsilon); + } + case OpType::CLAMP: { + const auto& in = get_input(node, 0, nodes, map); + bool f32 = in.precision == Precision::FP32 && out.precision == Precision::FP32; + if (!f32 && (!fp16(in) || !fp16(out))) return false; + return cactus_metal_encode_clamp(out.get_data(), in.get_data(), out.total_size, + node.params.scalar, node.params.scale, f32 ? 1 : 0); + } + case OpType::GLU: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || !fp16(out) || out.shape.empty()) return false; + int axis = node.params.axis; + if (axis < 0) axis += (int)in.shape.size(); + if (axis < 0 || (size_t)axis >= in.shape.size()) return false; + size_t split = in.shape[(size_t)axis] / 2; + size_t inner = 1; + for (size_t i = (size_t)axis + 1; i < in.shape.size(); ++i) inner *= in.shape[i]; + return cactus_metal_encode_glu(out.get_data(), in.get_data(), split, inner, out.total_size); + } + case OpType::SOFTMAX: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || !fp16(out) || out.shape.empty()) return false; + size_t cols = out.shape.back(); + if (cols == 0) return false; + return cactus_metal_encode_softmax_rows(out.get_data(), in.get_data(), out.total_size / cols, cols); + } + case OpType::GATED_DELTANET_DECODE: { + if (node.input_ids.size() != 6) return false; + const auto& q = get_input(node, 0, nodes, map); + const auto& k = get_input(node, 1, nodes, map); + const auto& v = get_input(node, 2, nodes, map); + const auto& g = get_input(node, 3, nodes, map); + const auto& b = get_input(node, 4, nodes, map); + const auto& st = get_input(node, 5, nodes, map); + if (!fp16(q) || !fp16(k) || !fp16(v) || !fp16(g) || !fp16(b) || !fp16(st) || !fp16(out)) return false; + if (q.shape.size() != 4 || v.shape.size() != 4 || st.shape.size() != 4) return false; + if (q.shape[1] != 1) return false; + size_t B = q.shape[0], Hq = q.shape[2], K = q.shape[3]; + size_t Hv = v.shape[2], V = v.shape[3]; + if (Hq == 0 || (Hv % Hq) != 0) return false; + if (st.shape[0] != B || st.shape[1] != K || st.shape[2] != Hv || st.shape[3] != V) return false; + if (out.total_size != B * (1 + K) * Hv * V) return false; + return cactus_metal_encode_deltanet_decode(out.get_data(), q.get_data(), k.get_data(), + v.get_data(), g.get_data(), b.get_data(), st.get_data(), + (uint32_t)B, (uint32_t)Hq, (uint32_t)Hv, (uint32_t)K, (uint32_t)V, node.params.scale); + } + case OpType::GATED_DELTANET_PREFILL: { + if (node.input_ids.size() != 6) return false; + const auto& q = get_input(node, 0, nodes, map); + const auto& k = get_input(node, 1, nodes, map); + const auto& v = get_input(node, 2, nodes, map); + const auto& g = get_input(node, 3, nodes, map); + const auto& b = get_input(node, 4, nodes, map); + const auto& st = get_input(node, 5, nodes, map); + if (!fp16(q) || !fp16(k) || !fp16(v) || !fp16(g) || !fp16(b) || !fp16(st) || !fp16(out)) return false; + if (q.shape.size() != 4 || v.shape.size() != 4 || st.shape.size() != 4) return false; + size_t B = q.shape[0], T = q.shape[1], Hq = q.shape[2], K = q.shape[3]; + size_t Hv = v.shape[2], V = v.shape[3]; + if (Hq == 0 || (Hv % Hq) != 0) return false; + if (st.shape[0] != B || st.shape[1] != K || st.shape[2] != Hv || st.shape[3] != V) return false; + if (out.total_size != B * (T + K) * Hv * V) return false; + return cactus_metal_encode_deltanet_prefill(out.get_data(), q.get_data(), k.get_data(), + v.get_data(), g.get_data(), b.get_data(), st.get_data(), + (uint32_t)B, (uint32_t)T, (uint32_t)Hq, (uint32_t)Hv, (uint32_t)K, (uint32_t)V, node.params.scale); + } + case OpType::RECURRENT_CACHE_WRITE: { + if (node.input_ids.size() < 2) return false; + const auto& src = get_input(node, 0, nodes, map); + BufferDesc& cache = nodes[map.at(node.input_ids[1])]->output_buffer; + if (!src.get_data() || !cache.get_data() || src.byte_size == 0) return false; + if (src.byte_size > cache.byte_size) return false; + return cactus_metal_encode_copy(cache.get_data(), src.get_data(), src.byte_size); + } + case OpType::TOPK: { + const auto& in = get_input(node, 0, nodes, map); + if (!fp16(in) || out.precision != Precision::FP32) return false; + if (in.shape.size() != 2 || node.params.top_k == 0 || node.params.top_k > 16) return false; + return cactus_metal_encode_topk_rows(out.get_data(), in.get_data(), + in.shape[0], in.shape[1], node.params.top_k); + } + case OpType::MOE_LAYER: { + const size_t num_experts = node.params.num_experts; + const size_t top_k = node.params.num_experts_per_tok; + if (!node.params.moe_gated || num_experts == 0 || top_k == 0 || top_k > 16) return false; + if (node.input_ids.size() != 3 + 3 * num_experts) return false; + const auto& hidden = get_input(node, 0, nodes, map); + const auto& probs = get_input(node, 1, nodes, map); + const auto& topk = get_input(node, 2, nodes, map); + if (!fp16(hidden) || !fp16(out) || !fp16(probs)) return false; + if (topk.precision != Precision::FP32) return false; + if (hidden.shape.size() != 2 || probs.shape.size() != 2 || probs.shape[1] != num_experts) return false; + uint32_t act; + switch (node.params.activation) { + case Activation::SILU: act = 0u; break; + case Activation::GELU: act = 1u; break; + default: return false; + } + const auto& w1_0 = get_input(node, 3, nodes, map); + if (!PrecisionTraits::is_cq(w1_0.precision) || w1_0.group_size == 0) return false; + CactusQuantMatrix key = w1_0.to_cq_matrix(); + if (!cactus_metal_moe_cq4_ready(&key)) { + std::vector w1s(num_experts), w3s(num_experts), w2s(num_experts); + for (size_t e = 0; e < num_experts; ++e) { + const auto& b1 = get_input(node, 3 + e, nodes, map); + const auto& b3 = get_input(node, 3 + num_experts + e, nodes, map); + const auto& b2 = get_input(node, 3 + 2 * num_experts + e, nodes, map); + if (!PrecisionTraits::is_cq(b1.precision) || b1.group_size == 0 || + !PrecisionTraits::is_cq(b3.precision) || b3.group_size == 0 || + !PrecisionTraits::is_cq(b2.precision) || b2.group_size == 0) return false; + w1s[e] = b1.to_cq_matrix(); + w3s[e] = b3.to_cq_matrix(); + w2s[e] = b2.to_cq_matrix(); + } + if (!cactus_metal_moe_cq4_build(w1s.data(), w3s.data(), w2s.data(), (uint32_t)num_experts)) + return false; + } + return cactus_metal_encode_moe_gated_cq4(out.get_data(), hidden.get_data(), probs.get_data(), + topk.get_data(), &key, (uint32_t)num_experts, (uint32_t)top_k, (uint32_t)hidden.shape[0], + act, node.params.normalize_routing ? 1u : 0u, node.params.epsilon, node.params.scalar); + } + case OpType::LAYERNORM: { + if (node.input_ids.size() < 2 || out.shape.empty()) return false; + const auto& in = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + const BufferDesc* b = node.input_ids.size() > 2 ? &get_input(node, 2, nodes, map) : nullptr; + if (!fp16(in) || !fp16(out) || !fp16(w) || (b && b->precision != Precision::FP16)) return false; + size_t dim = out.shape.back(); + if (dim == 0 || w.total_size != dim) return false; + return cactus_metal_encode_layer_norm(out.get_data(), in.get_data(), w.get_data(), + b ? b->get_data() : nullptr, out.total_size / dim, dim, node.params.epsilon); + } + case OpType::CONV1D_K3: { + if (node.input_ids.size() < 2) return false; + const auto& x = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + if (!fp16(x) || !fp16(out)) return false; + if (x.shape.size() != 3 || w.shape.size() != 3 || x.shape[0] != 1 || w.shape[2] != 3) return false; + size_t Cin = x.shape[1], L = x.shape[2], Cout = w.shape[0]; + size_t stride = node.params.stride ? node.params.stride : 1; + size_t Lout = ((L - 1) / stride) + 1; + int w_int8 = w.precision == Precision::INT8 ? 1 : 0; + if (!w_int8 && w.precision != Precision::FP16) return false; + if (out.total_size != Cout * Lout) return false; + return cactus_metal_encode_conv1d_k3(out.get_data(), x.get_data(), w.get_data(), w_int8, + w.activation_scales_data, (uint32_t)w.group_size, + (uint32_t)Cin, (uint32_t)L, (uint32_t)Cout, (uint32_t)Lout, (uint32_t)stride); + } + case OpType::ATTENTION: { + if (node.input_ids.size() < 3) return false; + const auto& q = get_input(node, 0, nodes, map); + const auto& k = get_input(node, 1, nodes, map); + const auto& v = get_input(node, 2, nodes, map); + const BufferDesc* mk = node.input_ids.size() > 3 ? &get_input(node, 3, nodes, map) : nullptr; + if (!fp16(q) || !fp16(k) || !fp16(v) || !fp16(out)) return false; + if (q.shape.size() != 4 || k.shape.size() != 4 || v.shape.size() != 4) return false; + size_t B = q.shape[0], T = q.shape[1], HQ = q.shape[2], D = q.shape[3]; + size_t S = k.shape[1], HKV = k.shape[2], DV = v.shape[3]; + if (HKV == 0 || HQ % HKV != 0) return false; + uint32_t mask_mode = 0; + if (mk) { + if (mk->precision != Precision::FP16) return false; + bool per_head = mk->shape.size() == 4; + if (!per_head && mk->shape.size() != 3) return false; + mask_mode = (node.params.attention_mask_is_additive ? 1u : 2u) + (per_head ? 2u : 0u); + } + float scale = node.params.scale != 0.0f ? node.params.scale : 1.0f/std::sqrt((float)D); + return cactus_metal_encode_attention_f16(out.get_data(), q.get_data(), k.get_data(), v.get_data(), + mk ? mk->get_data() : nullptr, + (uint32_t)B, (uint32_t)T, (uint32_t)S, (uint32_t)HQ, (uint32_t)HKV, (uint32_t)D, (uint32_t)DV, + scale, node.params.is_causal ? 1u : 0u, (uint32_t)node.params.position_offset, + (uint32_t)node.params.window_size, node.params.logit_cap, mask_mode); + } + case OpType::RMS_NORM: { + if (node.input_ids.size() < 2 || out.shape.empty()) return false; + const auto& in = get_input(node, 0, nodes, map); + const auto& w = get_input(node, 1, nodes, map); + if (!fp16(in) || !fp16(out) || !fp16(w)) return false; + size_t dim = out.shape.back(); + if (dim == 0) return false; + return cactus_metal_encode_rms_norm(out.get_data(), in.get_data(), w.get_data(), + out.total_size / dim, dim, node.params.epsilon); + } + case OpType::PRECISION_CAST: { + const auto& in = get_input(node, 0, nodes, map); + return cactus_metal_encode_cast(out.get_data(), static_cast(out.precision), + in.get_data(), static_cast(in.precision), in.total_size); + } + case OpType::ATTENTION_CACHED: { + if (node.input_ids.size() < 5) return false; + const auto& qb = get_input(node, 0, nodes, map); + const auto& knew = get_input(node, 1, nodes, map); + const auto& vnew = get_input(node, 2, nodes, map); + const auto& kcache = get_input(node, 3, nodes, map); + const auto& vcache = get_input(node, 4, nodes, map); + if (qb.shape.size() < 3) return false; + size_t batch=qb.shape[0], seq=qb.shape[1], nqh=qb.shape[2]; + if (batch != 1) return false; + if (kcache.precision == Precision::FP16 || vcache.precision == Precision::FP16) return false; + const uint64_t* km = reinterpret_cast(kcache.get_data()); + const uint64_t* vm = reinterpret_cast(vcache.get_data()); + if (!km || !vm) return false; + size_t cache_len=km[0], max_seq=km[1], kv_heads=km[2], hdim=km[3], v_max=vm[1]; + if (kv_heads == 0 || hdim == 0 || nqh % kv_heads != 0) return false; + size_t v_hdim = node.params.v_head_dim > 0 ? node.params.v_head_dim : hdim; + size_t new_seq_len = knew.total_size / (kv_heads * hdim); + size_t history_len = cache_len >= new_seq_len ? cache_len - new_seq_len : 0; + size_t po = node.params.position_offset, pos, new_len = seq; + if (po == std::numeric_limits::max()) pos = history_len; + else if (po == std::numeric_limits::max() - 1) { + pos = (cache_len >= seq) ? cache_len - seq : 0; history_len = cache_len; new_len = 0; + } else pos = po; + size_t total_keys = history_len + new_len; + size_t win = node.params.window_size; + size_t kv_start = (win > 0 && pos > win) ? pos - win : 0; + if (pos > history_len) kv_start = 0; + size_t kv_end = node.params.is_causal ? std::min(total_keys, pos + 1) : total_keys; + float scale = node.params.scale != 0.0f ? node.params.scale : 1.0f/std::sqrt((float)hdim); + const char* bk = static_cast(kcache.get_data()); + const char* bv = static_cast(vcache.get_data()); + size_t ngK=(hdim+31)/32, ngV=(v_hdim+31)/32; + if (seq > 1) { + size_t sink = km[4]; + uint32_t ringv = (win > 0 && max_seq > 2*sink + 1) ? (uint32_t)(max_seq - 2*sink - 1) : 0u; + return cactus_metal_encode_attention_i8_prefill( + out.get_data(), qb.get_data(), knew.get_data(), vnew.get_data(), + bk + 64, bv + 64, + bk + 64 + max_seq*kv_heads*hdim, bv + 64 + v_max*kv_heads*v_hdim, + (uint32_t)nqh, (uint32_t)kv_heads, (uint32_t)hdim, (uint32_t)v_hdim, + (uint32_t)history_len, (uint32_t)new_len, (uint32_t)pos, + (uint32_t)win, node.params.is_causal ? 1u : 0u, (uint32_t)seq, scale, + max_seq*kv_heads*hdim, v_max*kv_heads*v_hdim, + max_seq*kv_heads*ngK*sizeof(float), v_max*kv_heads*ngV*sizeof(float), + (uint32_t)sink, ringv); + } + size_t cache_ceiling = nodes[map.at(node.input_ids[3])]->params.max_cache_seq_len; + if (win > 0 && win < cache_ceiling) { + size_t sink = km[4]; + size_t ringW = max_seq > sink + 1 ? max_seq - sink - 1 : 0; + if (ringW == 0) return false; + if (total_keys > ringW) { + history_len = ringW; + total_keys = ringW; + kv_start = 0; + kv_end = ringW; + } + } + bool ok = cactus_metal_encode_attention_i8( + out.get_data(), qb.get_data(), knew.get_data(), vnew.get_data(), + bk + 64, bv + 64, + bk + 64 + max_seq*kv_heads*hdim, bv + 64 + v_max*kv_heads*v_hdim, + (uint32_t)nqh, (uint32_t)kv_heads, (uint32_t)hdim, (uint32_t)v_hdim, + (uint32_t)history_len, (uint32_t)total_keys, (uint32_t)kv_start, (uint32_t)kv_end, scale, + history_len*kv_heads*hdim, history_len*kv_heads*v_hdim, + history_len*kv_heads*ngK*sizeof(float), history_len*kv_heads*ngV*sizeof(float)); + return ok; + } + case OpType::KV_CACHE_APPEND: { + if (node.input_ids.size() < 2) return false; + const auto& new_kv = get_input(node, 0, nodes, map); + GraphNode& cache_node = *nodes[map.at(node.input_ids[1])]; + BufferDesc& cache = cache_node.output_buffer; + if (!fp16(new_kv) || cache.precision == Precision::FP16 || !cache.get_data()) return false; + uint64_t* km = reinterpret_cast(cache.get_data()); + size_t current_len=km[0], max_len=km[1], kv_heads=km[2], hdim=km[3], sink=km[4], num_slots=km[5]; + if (num_slots != 1 || kv_heads == 0 || hdim == 0) return false; + size_t new_seq_len = new_kv.total_size / (kv_heads * hdim); + size_t ceiling = cache_node.params.max_cache_seq_len, ws = node.params.window_size; + bool sliding = ws > 0 && ws < ceiling; + size_t new_total = current_len + new_seq_len; + char* base = static_cast(cache.get_data()); + size_t ngK = (hdim + 31)/32; + if (new_seq_len > 1) { + if (sliding) { + if (max_len <= sink + 1) return false; + uint32_t W = (uint32_t)(max_len - sink - 1); + if (!cactus_metal_encode_kv_append_ring_i8_m(new_kv.get_data(), base + 64, + base + 64 + max_len*kv_heads*hdim, (uint32_t)kv_heads, (uint32_t)hdim, + (uint32_t)current_len, 32, (uint32_t)new_seq_len, (uint32_t)sink, W, + new_kv.byte_size, max_len*kv_heads*hdim, max_len*kv_heads*ngK*sizeof(float))) + return false; + km[0] = new_total; + } else if (ws > 0 && new_total > max_len) { + size_t window = max_len; + size_t keep_sink = std::min({sink, current_len, window}); + size_t tail_capacity = window - keep_sink; + if (new_seq_len >= tail_capacity) return false; + size_t remaining = std::min(tail_capacity - new_seq_len, current_len - keep_sink); + size_t shift_src = current_len - remaining; + if (!cactus_metal_encode_kv_append_sliding_i8_m(new_kv.get_data(), base + 64, + base + 64 + max_len*kv_heads*hdim, (uint32_t)kv_heads, (uint32_t)hdim, + (uint32_t)keep_sink, (uint32_t)remaining, (uint32_t)shift_src, 32, (uint32_t)new_seq_len, + new_kv.byte_size, max_len*kv_heads*hdim, max_len*kv_heads*ngK*sizeof(float))) + return false; + km[0] = keep_sink + remaining + new_seq_len; + } else { + if (new_total > max_len) { + if (!cactus_kv_cache_grow(cache, new_total, ceiling)) return false; + km = reinterpret_cast(cache.get_data()); + max_len = km[1]; base = static_cast(cache.get_data()); + if (new_total > max_len) return false; + } + if (!cactus_metal_encode_kv_append_i8_m(new_kv.get_data(), base + 64, + base + 64 + max_len*kv_heads*hdim, (uint32_t)kv_heads, (uint32_t)hdim, + (uint32_t)current_len, 32, (uint32_t)new_seq_len, + new_kv.byte_size, max_len*kv_heads*hdim, max_len*kv_heads*ngK*sizeof(float))) + return false; + km[0] = new_total; + } + if (out.get_data()) *out.data_as() = static_cast(km[0]); + return true; + } + if (sliding) { + if (max_len <= sink + 1) return false; + uint32_t W = (uint32_t)(max_len - sink - 1); + if (!cactus_metal_encode_kv_append_ring_i8_m(new_kv.get_data(), base + 64, + base + 64 + max_len*kv_heads*hdim, (uint32_t)kv_heads, (uint32_t)hdim, + (uint32_t)current_len, 32, 1u, (uint32_t)sink, W, + new_kv.byte_size, max_len*kv_heads*hdim, max_len*kv_heads*ngK*sizeof(float))) + return false; + km[0] = new_total; + if (out.get_data()) *out.data_as() = static_cast(new_total); + return true; + } + if (new_total > max_len) { + if (!cactus_kv_cache_grow(cache, new_total, ceiling)) return false; + km = reinterpret_cast(cache.get_data()); + max_len = km[1]; base = static_cast(cache.get_data()); + if (new_total > max_len) return false; + } + size_t window = max_len; + if (new_total > window) { + size_t keep_sink = std::min({sink, current_len, window}); + size_t tail_capacity = window - keep_sink; + if (new_seq_len >= tail_capacity) return false; + size_t remaining = std::min(tail_capacity - new_seq_len, current_len - keep_sink); + size_t shift_src = current_len - remaining; + if (!cactus_metal_encode_kv_append_sliding_i8(new_kv.get_data(), base + 64, + base + 64 + max_len*kv_heads*hdim, (uint32_t)kv_heads, (uint32_t)hdim, + (uint32_t)keep_sink, (uint32_t)remaining, (uint32_t)shift_src, 32, new_kv.byte_size, + max_len*kv_heads*hdim, max_len*kv_heads*ngK*sizeof(float))) + return false; + km[0] = keep_sink + remaining + new_seq_len; + if (out.get_data()) *out.data_as() = static_cast(km[0]); + return true; + } + if (!cactus_metal_encode_kv_append_i8(new_kv.get_data(), base + 64, + base + 64 + max_len*kv_heads*hdim, (uint32_t)kv_heads, (uint32_t)hdim, + (uint32_t)current_len, 32, new_kv.byte_size, + max_len*kv_heads*hdim, max_len*kv_heads*ngK*sizeof(float))) + return false; + km[0] = new_total; + if (out.get_data()) *out.data_as() = static_cast(new_total); + return true; + } + case OpType::CONV_CACHE_APPEND: { + if (node.input_ids.size() < 2) return false; + const auto& new_data = get_input(node, 0, nodes, map); + BufferDesc& cache = nodes[map.at(node.input_ids[1])]->output_buffer; + if (!fp16(out) || !cache.get_data() || !out.get_data()) return false; + if (new_data.precision != Precision::FP16 && new_data.precision != Precision::FP32) return false; + uint64_t* cm = reinterpret_cast(cache.get_data()); + uint64_t head = cm[0], count = cm[1], ws = cm[2], hd = cm[3]; + if (ws == 0 || hd == 0 || head >= ws) return false; + size_t num_rows = new_data.total_size / hd; + if (num_rows == 0 || out.total_size != ws * hd) return false; + uint32_t nnew = (uint32_t)std::min(num_rows, ws); + uint64_t count_new = std::min(ws, count + num_rows); + if (!cactus_metal_encode_conv_cache_append(out.get_data(), new_data.get_data(), + static_cast(cache.get_data()) + 64, + (uint32_t)hd, (uint32_t)ws, nnew, (uint32_t)head, (uint32_t)count_new, + (uint32_t)num_rows, new_data.precision == Precision::FP32 ? 1 : 0)) + return false; + cm[0] = (head + nnew) % ws; + cm[1] = count_new; + return true; + } + case OpType::EMBEDDING: { + if (node.input_ids.size() < 2) return false; + const auto& emb = get_input(node, 0, nodes, map); + const auto& idxb = get_input(node, 1, nodes, map); + if (!fp16(out)) return false; + size_t M = idxb.total_size; + if (M == 0 || M > 4096) return false; + std::vector rows(M); + if (idxb.precision == Precision::FP32) { const float* p = idxb.data_as(); for (size_t i=0;i(); for (size_t i=0;i(); for (size_t i=0;i 0) { + CactusQuantMatrix W = emb.to_cq_matrix(); + if (W.flags & CACTUS_QUANT_FLAG_ORTHOGONAL) + return cactus_metal_encode_embedding_ortho_m(out.get_data(), &W, rows.data(), (uint32_t)M); + return cactus_metal_encode_embedding_hadamard_m(out.get_data(), &W, rows.data(), (uint32_t)M); + } + if (emb.precision == Precision::FP16 && emb.shape.size() >= 2) { + uint32_t D = (uint32_t)emb.shape[1]; + return cactus_metal_encode_gather_f16(out.get_data(), emb.get_data(), + emb.total_size * sizeof(__fp16), rows.data(), (uint32_t)M, D); + } + return false; + } + default: return false; + } +} + +struct MetalFusePlan; +MetalFusePlan* cactus_metal_plan_build(const std::vector>& nodes, + const std::unordered_map& map, + const std::unordered_set& pinned_ids, + const std::vector& retype, + const std::unordered_set* banned = nullptr); +void cactus_metal_plan_free(MetalFusePlan* p); +bool cactus_metal_plan_fold(MetalFusePlan* p, const std::vector>& nodes); +const std::vector* cactus_metal_plan_exec_list(const MetalFusePlan* p); +void* cactus_metal_plan_arena_ptr(const MetalFusePlan* p, size_t i); +bool cactus_metal_plan_has_arena(const MetalFusePlan* p); +void cactus_metal_plan_extend_last_use(const MetalFusePlan* p, std::vector& last_use); +int32_t cactus_metal_plan_action(const MetalFusePlan* p, size_t i); +bool cactus_metal_plan_encode(MetalFusePlan* p, int32_t cid, + const std::vector>& nodes, + const std::unordered_map& map); + +void CactusGraph::build_metal_retype_plan() { + metal_retype_built_ = true; + const size_t n = nodes_.size(); + metal_retype_plan_.assign(n, 0); + auto idxof = [&](size_t id) -> long { + auto it = node_index_map_.find(id); + return it == node_index_map_.end() ? -1 : (long)it->second; + }; + std::vector> cons(n); + for (size_t i = 0; i < n; ++i) + for (size_t id : nodes_[i]->input_ids) { long j = idxof(id); if (j >= 0) cons[(size_t)j].push_back(i); } + auto interior = [&](const GraphNode& nd) { + if (nd.params.backend != ComputeBackend::METAL) return false; + if (nd.output_buffer.precision != Precision::FP32) return false; + if (nd.output_buffer.shape.size() > 8) return false; + switch (nd.op_type) { + case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: + case OpType::MULTIPLY: case OpType::DIVIDE: + case OpType::SCALAR_ADD: case OpType::SCALAR_SUBTRACT: + case OpType::SCALAR_MULTIPLY: case OpType::SCALAR_DIVIDE: + case OpType::CLAMP: case OpType::GELU: case OpType::TANH: + case OpType::SILU: case OpType::RELU: + case OpType::VIEW: case OpType::RESHAPE: case OpType::FLATTEN: + case OpType::TRANSPOSE: case OpType::CAT: + return std::fabs(nd.params.scalar) < 60000.0f && std::fabs(nd.params.scale) < 60000.0f; + default: return false; + } + }; + std::vector cand(n, 0); + for (size_t i = 0; i < n; ++i) { + const GraphNode& nd = *nodes_[i]; + if (nd.op_type == OpType::PRECISION_CAST && nd.output_buffer.precision == Precision::FP32 + && !nd.input_ids.empty()) { + long j = idxof(nd.input_ids[0]); + if (j >= 0 && nodes_[(size_t)j]->output_buffer.precision == Precision::FP16) cand[i] = 1; + continue; + } + if (!interior(nd) || nd.input_ids.empty()) continue; + bool all = true; + for (size_t id : nd.input_ids) { + long j = idxof(id); + if (j < 0 || !cand[(size_t)j]) { all = false; break; } + } + cand[i] = all ? 1 : 0; + } + std::vector ok = cand; + for (size_t ri = n; ri-- > 0;) { + if (!ok[ri]) continue; + const GraphNode& nd = *nodes_[ri]; + if (retained_output_node_ids_.count(nd.id) || persistent_node_ids_.count(nd.id)) { ok[ri] = 0; continue; } + if (cons[ri].empty()) { ok[ri] = 0; continue; } + for (size_t c : cons[ri]) { + const GraphNode& cn = *nodes_[c]; + bool absorbs = (cn.op_type == OpType::PRECISION_CAST + && cn.output_buffer.precision == Precision::FP16) + || (cand[c] && ok[c]); + if (!absorbs) { ok[ri] = 0; break; } + } + } + for (size_t i = 0; i < n; ++i) { + if (!ok[i]) continue; + metal_retype_plan_[i] = nodes_[i]->op_type == OpType::PRECISION_CAST ? 1 : 2; + } + for (size_t i = 0; i < n; ++i) { + const GraphNode& nd = *nodes_[i]; + if (nd.op_type != OpType::PRECISION_CAST || nd.output_buffer.precision != Precision::FP16 + || nd.input_ids.empty()) continue; + long j = idxof(nd.input_ids[0]); + if (j >= 0 && metal_retype_plan_[(size_t)j] != 0) metal_retype_plan_[i] = 1; + } +} + +void CactusGraph::execute(const std::string& profile_file) { + cactus_graph_mark_unadjusted(); + BufferPool& pool = buffer_pool_; + const size_t n = nodes_.size(); + infer_shapes(); + + auto get_env_int = [](const char* name, int fallback) -> int { + const char* val = std::getenv(name); + return val ? std::atoi(val) : fallback; + }; + + bool trace_execution = get_env_int("CACTUS_TRACE_EXECUTE", 0) != 0; + static const size_t flush_cadence = (size_t)get_env_int("CACTUS_FLUSH_CADENCE", 48); + bool trace_nan = get_env_int("CACTUS_TRACE_NAN", 0) != 0; + bool need_debug = !profile_file.empty(); + if (!need_debug) { + static const bool env_debug = check_debug_env(); + need_debug = env_debug; + } + if (trace_execution) { + need_debug = true; + } + + auto trace_nonfinite = [&](size_t node_idx, const GraphNode& node) { + if (!trace_nan) return; + const BufferDesc& buffer = node.output_buffer; + const void* data = buffer.get_data(); + if (!data || buffer.total_size == 0) return; + + auto report = [&](size_t elem_idx, float value) { + std::cerr << "[cactus:nan] idx=" << node_idx + << " id=" << node.id + << " op=" << get_op_name(node.op_type) + << " elem=" << elem_idx + << " value=" << value + << " shape=["; + for (size_t dim_idx = 0; dim_idx < buffer.shape.size(); ++dim_idx) { + if (dim_idx > 0) std::cerr << ","; + std::cerr << buffer.shape[dim_idx]; + } + std::cerr << "]" << std::endl; + }; + + if (buffer.precision == Precision::FP16) { + const __fp16* values = buffer.data_as<__fp16>(); + for (size_t i = 0; i < buffer.total_size; ++i) { + float value = static_cast(values[i]); + if (!std::isfinite(value)) { + report(i, value); + return; + } + } + } else if (buffer.precision == Precision::FP32) { + const float* values = buffer.data_as(); + for (size_t i = 0; i < buffer.total_size; ++i) { + float value = values[i]; + if (!std::isfinite(value)) { + report(i, value); + return; + } + } + } + }; + + auto can_release_node = [&](size_t node_idx) { + const auto& node = nodes_[node_idx]; + if (node->op_type == OpType::INPUT) return false; + if (node->op_type == OpType::KV_CACHE_STATE + || node->op_type == OpType::CONV_CACHE_STATE + || node->op_type == OpType::RECURRENT_CACHE_STATE) return false; + if (persistent_node_ids_.count(node->id)) return false; + if (retained_output_node_ids_.count(node->id)) return false; + return true; + }; + + bool metal_mode = false; + if (cactus_metal_available()) { + for (size_t i = 0; i < n; ++i) { + if (nodes_[i]->op_type == OpType::INPUT) continue; + if (nodes_[i]->params.backend == ComputeBackend::METAL) { metal_mode = true; break; } + } + } + if (metal_mode && n < 100) { + for (size_t i = 0; i < n; ++i) { + OpType t = nodes_[i]->op_type; + if (t == OpType::LSTM_CELL || t == OpType::BILSTM_SEQUENCE) { + metal_mode = false; + break; + } + } + } + if (metal_mode && !need_debug && !metal_retype_built_) build_metal_retype_plan(); + MetalFusePlan* fplan = nullptr; + if (metal_mode && !need_debug) { + uint64_t sig = 1469598103934665603ull; + sig ^= (uint64_t)n; sig *= 1099511628211ull; + for (auto& np : nodes_) { + sig ^= (uint64_t)np->op_type + 0x9e3779b97f4a7c15ull; + sig *= 1099511628211ull; + sig ^= (uint64_t)np->params.backend + 0x9e3779b97f4a7c15ull; + sig *= 1099511628211ull; + if (np->op_type != OpType::INPUT) continue; + for (size_t d : np->output_buffer.shape) { + sig ^= (uint64_t)d + 0x9e3779b97f4a7c15ull; + sig *= 1099511628211ull; + } + } + auto it = metal_plans_.find(sig); + if (it == metal_plans_.end()) { + std::unordered_set pinned(retained_output_node_ids_.begin(), retained_output_node_ids_.end()); + pinned.insert(persistent_node_ids_.begin(), persistent_node_ids_.end()); + static const std::vector no_retype; + const std::vector& rt = (!metal_retype_disabled_ && metal_retype_plan_.size() == n) + ? metal_retype_plan_ : no_retype; + auto bit = metal_plan_banned_.find(sig); + const std::unordered_set* banned = + bit == metal_plan_banned_.end() ? nullptr : &bit->second; + it = metal_plans_.emplace(sig, cactus_metal_plan_build(nodes_, node_index_map_, pinned, rt, banned)).first; + } + fplan = it->second; + metal_plan_sig_ = sig; + } + const uint8_t* rplan = (metal_mode && !need_debug && !metal_retype_disabled_ + && metal_retype_plan_.size() == n) ? metal_retype_plan_.data() : nullptr; + auto plan_of = [&](const GraphNode& node) -> uint8_t { + if (!rplan) return 0; + auto it = node_index_map_.find(node.id); + return it == node_index_map_.end() ? 0 : rplan[it->second]; + }; + auto aliases_input = [&](const GraphNode& node) -> bool { + if (plan_of(node) == 1) return true; + if (node.op_type == OpType::SLICE && !node.input_ids.empty()) { + auto it = node_index_map_.find(node.input_ids[0]); + if (it != node_index_map_.end()) { + const auto& ish = nodes_[it->second]->output_buffer.shape; + size_t ax = static_cast(node.params.axis); + if (ax < ish.size()) { + size_t outer = 1; + for (size_t d = 0; d < ax; ++d) outer *= ish[d]; + if (outer == 1) return true; + } + } + } + if (node.op_type == OpType::INDEX && node.params.axis == 0) return true; + if (!metal_mode) return false; + if (node.op_type == OpType::VIEW || node.op_type == OpType::RESHAPE || node.op_type == OpType::FLATTEN) + return true; + if (node.op_type == OpType::PRECISION_CAST && !node.input_ids.empty()) { + auto it = node_index_map_.find(node.input_ids[0]); + if (it != node_index_map_.end() && + nodes_[it->second]->output_buffer.precision == node.output_buffer.precision) return true; + } + return false; + }; + auto preallocates_output = [&](const GraphNode& node) { return !aliases_input(node); }; + + std::vector last_use(n, 0); + std::vector use_count(n, 0); + for (size_t i = 0; i < n; ++i) { + for (size_t input_id : nodes_[i]->input_ids) { + auto it = node_index_map_.find(input_id); + if (it != node_index_map_.end()) { + last_use[it->second] = std::max(last_use[it->second], i); + ++use_count[it->second]; + } + } + } + + std::vector keep_until_graph_cleanup(n, false); + for (size_t i = n; i-- > 0;) { + const auto& node = *nodes_[i]; + if (!aliases_input(node) || node.input_ids.empty()) continue; + auto it = node_index_map_.find(node.input_ids[0]); + if (it == node_index_map_.end()) continue; + size_t base_idx = it->second; + if (use_count[i] == 0 || keep_until_graph_cleanup[i]) { + keep_until_graph_cleanup[base_idx] = true; + } else { + last_use[base_idx] = std::max(last_use[base_idx], last_use[i]); + } + } + + if (fplan) cactus_metal_plan_extend_last_use(fplan, last_use); + + std::vector> release_after(n); + for (size_t i = 0; i < n; ++i) { + if (!can_release_node(i) || use_count[i] == 0 || keep_until_graph_cleanup[i]) continue; + release_after[last_use[i]].push_back(i); + } + + if (metal_mode && !need_debug) { + const bool transient_acts = n >= 1500 && !cactus_metal_plan_has_arena(fplan); + std::vector transient_ptr(transient_acts ? n : 0, nullptr); + std::vector transient_dead; + std::vector transient_ok(transient_acts ? n : 0, 0); + auto release_transients = [&]() { + for (void* p : transient_dead) cactus_metal_free_shared(p); + transient_dead.clear(); + for (size_t i = 0; i < transient_ptr.size(); ++i) + if (transient_ptr[i]) { cactus_metal_free_shared(transient_ptr[i]); transient_ptr[i] = nullptr; } + }; + struct CacheWordSnap { size_t idx; size_t word; uint64_t value; }; + std::vector kv_snapshot; + auto kv_restore = [&]() { + for (const auto& s : kv_snapshot) { + uint64_t* km = reinterpret_cast(nodes_[s.idx]->output_buffer.get_data()); + if (km) km[s.word] = s.value; + } + }; + auto metal_abort_cleanup = [&]() { + cactus_metal_session_sync(); + kv_restore(); + release_transients(); + cactus_metal_set_active(false); + cactus_metal_session_end(); + }; + struct MetalExecGuard { + decltype(metal_abort_cleanup)& cleanup; + bool armed = true; + ~MetalExecGuard() { if (armed) cleanup(); } + } metal_guard{metal_abort_cleanup}; + cactus_metal_session_begin(); + cactus_metal_set_active(true); + if (transient_acts) { + for (size_t i = 0; i < n; ++i) { + size_t id = nodes_[i]->id; + transient_ok[i] = !retained_output_node_ids_.count(id) + && !persistent_node_ids_.count(id) + && !keep_until_graph_cleanup[i] + && nodes_[i]->output_buffer.byte_size >= (256u << 10); + } + } + auto metal_release = [&](size_t idx) { + GraphNode& nd = *nodes_[idx]; + if (aliases_input(nd)) nd.output_buffer.external_data = nullptr; + if (transient_acts && transient_ptr[idx]) { + transient_dead.push_back(transient_ptr[idx]); + transient_ptr[idx] = nullptr; + } + }; + auto assign_persistent_act = [&](GraphNode& nd) { + size_t need = nd.output_buffer.byte_size; + auto pit = metal_persistent_acts_.find(nd.id); + void* p = nullptr; + if (pit != metal_persistent_acts_.end() && pit->second.second >= need) { + p = pit->second.first; + } else { + if (pit != metal_persistent_acts_.end()) cactus_metal_free_shared(pit->second.first); + p = cactus_metal_alloc_pooled(need); + if (p) metal_persistent_acts_[nd.id] = { p, need }; + } + if (p) { nd.output_buffer.release_to_pool(pool); nd.output_buffer.set_external(p); } + else nd.output_buffer.resize_from_pool(pool); + }; + size_t since_flush = 0, since_recycle = 0; + if (fplan) cactus_metal_plan_fold(fplan, nodes_); + std::vector metal_live(n, 0); + auto maybe_recycle = [&]() { + if (!transient_acts || ++since_recycle < 256 || transient_dead.empty()) return; + for (void* p : transient_dead) cactus_metal_free_shared(p); + transient_dead.clear(); + cactus_metal_session_sync(); + std::fill(metal_live.begin(), metal_live.end(), 0); + since_flush = 0; + since_recycle = 0; + }; + for (size_t i = 0; i < n; ++i) { + const GraphNode& nd = *nodes_[i]; + bool kv = nd.op_type == OpType::KV_CACHE_APPEND; + bool conv = nd.op_type == OpType::CONV_CACHE_APPEND; + if ((!kv && !conv) || nd.input_ids.size() < 2) continue; + auto ci = node_index_map_.find(nd.input_ids[1]); + if (ci == node_index_map_.end()) continue; + const uint64_t* km = reinterpret_cast(nodes_[ci->second]->output_buffer.get_data()); + if (!km) continue; + kv_snapshot.push_back({ci->second, 0, km[0]}); + if (conv) kv_snapshot.push_back({ci->second, 1, km[1]}); + } + auto input_metal_live = [&](const GraphNode& nd) -> bool { + for (size_t id : nd.input_ids) { + auto it = node_index_map_.find(id); + if (it != node_index_map_.end() && metal_live[it->second]) return true; + } + return false; + }; + const std::vector* elist = fplan ? cactus_metal_plan_exec_list(fplan) : nullptr; + size_t iter_count = elist ? elist->size() : n; + if (fplan) { + for (size_t ii = 0; ii < iter_count; ++ii) { + size_t i = elist ? (size_t)(*elist)[ii] : ii; + if (cactus_metal_plan_action(fplan, i) == -3) assign_persistent_act(*nodes_[i]); + } + } + for (size_t ii = 0; ii < iter_count; ++ii) { + size_t i = elist ? (size_t)(*elist)[ii] : ii; + auto& node = nodes_[i]; + const OpType ot = node->op_type; + if (ot == OpType::INPUT) continue; + if (ot == OpType::KV_CACHE_STATE || ot == OpType::CONV_CACHE_STATE + || ot == OpType::RECURRENT_CACHE_STATE) { + dispatch_node(*node, nodes_, node_index_map_); + populated_node_ids_.insert(node->id); + for (size_t r : release_after[i]) metal_release(r); + continue; + } + int32_t fact = fplan ? cactus_metal_plan_action(fplan, i) : -1; + if (fact == -3) { + assign_persistent_act(*node); + metal_live[i] = 1; + for (size_t r : release_after[i]) metal_release(r); + continue; + } + if (aliases_input(*node)) { + if (rplan && rplan[i] == 1) { + const auto& src = get_input(*node, 0, nodes_, node_index_map_); + node->output_buffer.set_external(const_cast(static_cast(src.get_data()))); + } else { + dispatch_node(*node, nodes_, node_index_map_); + } + metal_live[i] = input_metal_live(*node) ? 1 : 0; + for (size_t r : release_after[i]) metal_release(r); + continue; + } + if (preallocates_output(*node)) { + void* ap = fplan ? cactus_metal_plan_arena_ptr(fplan, i) : nullptr; + if (ap) { + node->output_buffer.release_to_pool(pool); + node->output_buffer.set_external(ap); + } else if (transient_acts && transient_ok[i]) { + void* tp = cactus_metal_alloc_shared(node->output_buffer.byte_size); + if (tp) { + node->output_buffer.release_to_pool(pool); + node->output_buffer.set_external(tp); + transient_ptr[i] = tp; + } else { + assign_persistent_act(*node); + } + } else { + assign_persistent_act(*node); + } + } + if (fact >= 0) { + if (cactus_metal_plan_encode(fplan, fact, nodes_, node_index_map_)) { + metal_live[i] = 1; + for (size_t r : release_after[i]) metal_release(r); + if (++since_flush >= flush_cadence) { cactus_metal_session_flush(); since_flush = 0; } + maybe_recycle(); + continue; + } + metal_guard.armed = false; + metal_abort_cleanup(); + metal_plan_banned_[metal_plan_sig_].insert(i); + metal_plans_.erase(metal_plan_sig_); + cactus_metal_plan_free(fplan); + execute(profile_file); + return; + } + std::vector> flipped; + if (rplan && rplan[i] == 2) { + auto flip = [&](BufferDesc& b) { + if (b.precision == Precision::FP32) { flipped.push_back({&b, b.precision}); b.precision = Precision::FP16; } + }; + flip(node->output_buffer); + for (size_t id : node->input_ids) { + auto it = node_index_map_.find(id); + if (it != node_index_map_.end() && rplan[it->second] != 0) flip(nodes_[it->second]->output_buffer); + } + } + bool force_cpu = (ot == OpType::PRECISION_CAST + && node->output_buffer.total_size <= 8 && !input_metal_live(*node)) + || (ot == OpType::EMBEDDING && input_metal_live(*node)); + bool encoded = !force_cpu && try_encode_metal(*node, nodes_, node_index_map_); + for (auto& fp : flipped) fp.first->precision = fp.second; + if (!encoded && rplan && rplan[i] == 2) { + metal_retype_disabled_ = true; + metal_guard.armed = false; + metal_abort_cleanup(); + execute(profile_file); + return; + } + if (!encoded) { + if (input_metal_live(*node)) { + cactus_metal_session_sync(); + std::fill(metal_live.begin(), metal_live.end(), 0); + } + dispatch_node(*node, nodes_, node_index_map_); + } else { + metal_live[i] = 1; + } + if (ot == OpType::PERSISTENT) populated_node_ids_.insert(node->id); + for (size_t r : release_after[i]) metal_release(r); + if (encoded && ++since_flush >= flush_cadence) { cactus_metal_session_flush(); since_flush = 0; } + if (encoded) maybe_recycle(); + } + metal_guard.armed = false; + release_transients(); + cactus_metal_set_active(false); + cactus_metal_session_end(); + cactus_graph_metal_tail_commit(); + return; + } + + if (!need_debug) { + auto run = [&](auto& nd) { dispatch_node(nd, nodes_, node_index_map_); }; + for (size_t i = 0; i < n; ++i) { + auto& node = nodes_[i]; + if (node->op_type == OpType::INPUT) continue; + if (node->op_type == OpType::KV_CACHE_STATE + || node->op_type == OpType::CONV_CACHE_STATE + || node->op_type == OpType::RECURRENT_CACHE_STATE) { + run(*node); + populated_node_ids_.insert(node->id); + for (size_t release_idx : release_after[i]) { + nodes_[release_idx]->output_buffer.release_memory(pool); + } + continue; + } + if (preallocates_output(*node)) { + node->output_buffer.resize_from_pool(pool); + } + run(*node); + trace_nonfinite(i, *node); + if (node->op_type == OpType::PERSISTENT) { + populated_node_ids_.insert(node->id); + } + for (size_t release_idx : release_after[i]) { + nodes_[release_idx]->output_buffer.release_memory(pool); + } + } + return; + } + + auto get_env_str = [](const char* name) -> std::string { + const char* val = std::getenv(name); + return val ? std::string(val) : std::string(); + }; + + bool capture_to_stdout = get_env_int("CACTUS_CAPTURE_STDOUT", 0) != 0; + std::string capture_file_path = get_env_str("CACTUS_CAPTURE_FILE"); + bool capture_requested = get_env_int("CACTUS_CAPTURE_ENABLE", 0) != 0; + std::string capture_dir = get_env_str("CACTUS_CAPTURE_DIR"); + + if (!capture_requested) { + capture_requested = capture_to_stdout || !capture_file_path.empty() || !capture_dir.empty(); + } else if (capture_file_path.empty() && !capture_to_stdout && capture_dir.empty()) { + capture_to_stdout = true; + } + + size_t capture_preview_count = static_cast(get_env_int("CACTUS_CAPTURE_PREVIEW_COUNT", 8)); + size_t capture_max_elements = static_cast(get_env_int("CACTUS_CAPTURE_MAX_ELEMENTS", 65536)); + + std::string env_profile = get_env_str("CACTUS_PROFILE_FILE"); + if (env_profile.empty()) env_profile = get_env_str("CACTUS_PROFILE"); + + std::string target_profile = profile_file; + if (target_profile.empty() && !env_profile.empty()) { + target_profile = env_profile; + } + + bool enable_profiling = !target_profile.empty(); + bool to_stdout = (target_profile == "stdout" || target_profile == "-"); + + std::ofstream profile_out; + std::ostream* out = &std::cout; + + if (enable_profiling && !to_stdout) { + profile_out.open(target_profile, std::ios::app); + if (profile_out.is_open()) { + out = &profile_out; + } + } + + auto total_start = std::chrono::high_resolution_clock::now(); + + if (enable_profiling) { + *out << "=== Graph Execution Profile ===" << std::endl; + *out << std::left << std::setw(24) << "Operation" + << std::setw(12) << "Time (ms)" + << std::setw(20) << "Output Shape" + << "Backend" << std::endl; + *out << std::string(72, '-') << std::endl; + } + + for (size_t node_idx = 0; node_idx < n; ++node_idx) { + auto& node = nodes_[node_idx]; + + if (node->op_type == OpType::INPUT) { + continue; + } + + if (node->op_type == OpType::KV_CACHE_STATE + || node->op_type == OpType::CONV_CACHE_STATE + || node->op_type == OpType::RECURRENT_CACHE_STATE) { + dispatch_node(*node, nodes_, node_index_map_); + if (trace_execution) { + std::cerr << "[cactus:execute] cache-state idx=" << node_idx + << " id=" << node->id + << " op=" << get_op_name(node->op_type) + << std::endl; + } + trace_nonfinite(node_idx, *node); + populated_node_ids_.insert(node->id); + continue; + } + + node->output_buffer.allocate_from_pool(pool); + + if (trace_execution) { + std::cerr << "[cactus:execute] begin idx=" << node_idx + << " id=" << node->id + << " op=" << get_op_name(node->op_type) + << " shape=["; + for (size_t dim_idx = 0; dim_idx < node->output_buffer.shape.size(); ++dim_idx) { + if (dim_idx > 0) std::cerr << ","; + std::cerr << node->output_buffer.shape[dim_idx]; + } + std::cerr << "]" << std::endl; + } + + if (enable_profiling) { + auto start = std::chrono::high_resolution_clock::now(); + dispatch_node(*node, nodes_, node_index_map_); + trace_nonfinite(node_idx, *node); + if (node->op_type == OpType::PERSISTENT) { + populated_node_ids_.insert(node->id); + } + auto end = std::chrono::high_resolution_clock::now(); + double ms = std::chrono::duration_cast(end - start).count() / 1000.0; + + std::string shape_str = "["; + for (size_t i = 0; i < node->output_buffer.shape.size(); ++i) { + if (i > 0) shape_str += ","; + shape_str += std::to_string(node->output_buffer.shape[i]); + } + shape_str += "]"; + + *out << std::left << std::setw(24) << get_op_name(node->op_type) + << std::setw(12) << std::fixed << std::setprecision(3) << ms + << std::setw(20) << shape_str << std::endl; + } else { + dispatch_node(*node, nodes_, node_index_map_); + trace_nonfinite(node_idx, *node); + if (node->op_type == OpType::PERSISTENT) { + populated_node_ids_.insert(node->id); + } + } + + if (trace_execution) { + std::cerr << "[cactus:execute] done idx=" << node_idx + << " id=" << node->id + << " op=" << get_op_name(node->op_type) + << std::endl; + } + } + + std::unique_ptr capture_file_stream; + std::vector capture_outputs; + + if (capture_requested) { + if (capture_to_stdout) { + capture_outputs.push_back(&std::cout); + } + + if (!capture_file_path.empty()) { + std::filesystem::path capture_path(capture_file_path); + if (capture_path.has_parent_path()) { + std::error_code ec; + std::filesystem::create_directories(capture_path.parent_path(), ec); + } + + auto stream_ptr = std::make_unique(capture_path, std::ios::out | std::ios::app); + if (stream_ptr->is_open()) { + capture_outputs.push_back(stream_ptr.get()); + capture_file_stream = std::move(stream_ptr); + } else { + std::cerr << "Failed to open capture file: " << capture_path << std::endl; + } + } + + if (!capture_dir.empty()) { + std::filesystem::path dir_path(capture_dir); + std::error_code ec; + std::filesystem::create_directories(dir_path, ec); + } + + if (capture_outputs.empty() && capture_dir.empty()) { + capture_requested = false; + } + } + + if (capture_requested) { + auto precision_to_string = [](Precision p) -> const char* { + switch (p) { + case Precision::FP32: return "FP32"; + case Precision::FP16: return "FP16"; + case Precision::INT8: return "INT8"; + default: return "UNKNOWN"; + } + }; + + auto format_double = [](double value) { + std::ostringstream oss; + oss << std::fixed << std::setprecision(6) << value; + return oss.str(); + }; + + auto now = std::chrono::system_clock::now(); + std::time_t now_time = std::chrono::system_clock::to_time_t(now); + std::tm time_info{}; +#if defined(_WIN32) + localtime_s(&time_info, &now_time); +#else + localtime_r(&now_time, &time_info); +#endif + + auto write_header = [&](std::ostream& stream) { + stream << "=== Graph Debug Capture ===" << std::endl; + stream << "Timestamp: " << std::put_time(&time_info, "%Y-%m-%d %H:%M:%S") << std::endl; + stream << "Captured nodes: " << debug_nodes_.size() << std::endl; + stream << std::string(60, '-') << std::endl; + }; + + auto write_separator = [](std::ostream& stream) { + stream << std::string(60, '-') << std::endl; + }; + + if (debug_nodes_.empty()) { + for (auto* stream : capture_outputs) { + write_header(*stream); + *stream << "No debug nodes registered on this graph." << std::endl; + write_separator(*stream); + stream->flush(); + } + } else { + for (auto* stream : capture_outputs) { + write_header(*stream); + } + + for (const auto& entry : debug_nodes_) { + auto node_it = node_index_map_.find(entry.node_id); + const GraphNode* node_ptr = nullptr; + if (node_it != node_index_map_.end()) { + node_ptr = nodes_[node_it->second].get(); + } + + if (!node_ptr) { + for (auto* stream : capture_outputs) { + *stream << "Layer " << entry.layer_idx << " - " << entry.name + << " (node " << entry.node_id << ")" << std::endl; + *stream << " Data: " << std::endl; + write_separator(*stream); + } + continue; + } + + const BufferDesc& buffer = node_ptr->output_buffer; + const void* data_ptr = buffer.get_data(); + size_t total_size = buffer.total_size; + + std::ostringstream shape_ss; + shape_ss << "["; + for (size_t i = 0; i < buffer.shape.size(); ++i) { + if (i > 0) { + shape_ss << ","; + } + shape_ss << buffer.shape[i]; + } + shape_ss << "]"; + std::string shape_str = shape_ss.str(); + + bool has_data = data_ptr != nullptr && total_size > 0; + size_t elements_to_process = total_size; + bool truncated = false; + if (has_data && elements_to_process > capture_max_elements && capture_max_elements > 0) { + elements_to_process = capture_max_elements; + truncated = true; + } + + std::vector preview_values; + if (capture_preview_count > 0) { + preview_values.reserve(std::min(capture_preview_count, elements_to_process)); + } + + double min_val = std::numeric_limits::infinity(); + double max_val = -std::numeric_limits::infinity(); + long double sum = 0.0L; + long double sum_sq = 0.0L; + + if (has_data && elements_to_process > 0) { + auto accumulate = [&](float value, size_t index) { + double v = static_cast(value); + min_val = std::min(min_val, v); + max_val = std::max(max_val, v); + sum += static_cast(value); + sum_sq += static_cast(value) * static_cast(value); + if (capture_preview_count > 0 && index < capture_preview_count) { + preview_values.push_back(value); + } + }; + + if (buffer.precision == Precision::FP32) { + const float* typed = static_cast(data_ptr); + for (size_t i = 0; i < elements_to_process; ++i) { + accumulate(typed[i], i); + } + } else if (buffer.precision == Precision::FP16) { + const __fp16* typed = reinterpret_cast(data_ptr); + for (size_t i = 0; i < elements_to_process; ++i) { + accumulate(static_cast(typed[i]), i); + } + } else if (buffer.precision == Precision::INT8) { + const int8_t* typed = reinterpret_cast(data_ptr); + for (size_t i = 0; i < elements_to_process; ++i) { + accumulate(static_cast(typed[i]), i); + } + } else { + has_data = false; + } + } else { + has_data = false; + } + + if (!capture_dir.empty() && has_data) { + std::string safe_name = entry.name; + std::string filename = capture_dir + "/" + safe_name + ".bin"; + std::ofstream bin_file(filename, std::ios::binary); + if (bin_file.is_open()) { + size_t bytes_to_write = buffer.byte_size; + if (truncated) { + bytes_to_write = PrecisionTraits::packed_size_of(buffer.precision, elements_to_process); + } + bin_file.write(reinterpret_cast(data_ptr), bytes_to_write); + } + } + + size_t processed_count = has_data ? elements_to_process : 0; + long double mean_ld = processed_count > 0 ? sum / processed_count : 0.0L; + long double variance_ld = processed_count > 0 ? (sum_sq / processed_count) - (mean_ld * mean_ld) : 0.0L; + if (variance_ld < 0.0L) { + variance_ld = 0.0L; + } + double mean_val = static_cast(mean_ld); + double stddev_val = processed_count > 0 ? std::sqrt(static_cast(variance_ld)) : 0.0; + + std::ostringstream preview_ss; + if (capture_preview_count > 0 && !preview_values.empty()) { + preview_ss << "["; + for (size_t i = 0; i < preview_values.size(); ++i) { + if (i > 0) { + preview_ss << ", "; + } + preview_ss << format_double(static_cast(preview_values[i])); + } + if (processed_count > preview_values.size()) { + if (!preview_values.empty()) { + preview_ss << ", ..."; + } else { + preview_ss << "..."; + } + } + preview_ss << "]"; + } + + for (auto* stream : capture_outputs) { + *stream << "Layer " << entry.layer_idx << " - " << entry.name + << " (node " << entry.node_id << ")" << std::endl; + *stream << " Shape: " << shape_str << " Precision: " << precision_to_string(buffer.precision) << std::endl; + if (!has_data) { + *stream << " Data: " << std::endl; + } else { + *stream << " Stats: min=" << format_double(min_val) + << " max=" << format_double(max_val) + << " mean=" << format_double(mean_val) + << " std=" << format_double(stddev_val) << std::endl; + if (truncated || processed_count < total_size) { + *stream << " Note: stats computed on first " << processed_count + << " of " << total_size << " values" << std::endl; + } + if (capture_preview_count > 0 && !preview_values.empty()) { + *stream << " Preview: " << preview_ss.str() << std::endl; + } + } + write_separator(*stream); + } + } + + for (auto* stream : capture_outputs) { + stream->flush(); + } + } + } + + if (enable_profiling) { + auto total_end = std::chrono::high_resolution_clock::now(); + auto total_duration = std::chrono::duration_cast(total_end - total_start); + double total_ms = total_duration.count() / 1000.0; + + *out << std::string(72, '-') << std::endl; + *out << "Total execution time: " << std::fixed << std::setprecision(3) << total_ms << " ms" << std::endl; + *out << "================================" << std::endl; + + if (profile_out.is_open()) { + profile_out.close(); + } + } +} + +void CactusGraph::invalidate_metal_state() { + for (auto& kv : metal_plans_) cactus_metal_plan_free(kv.second); + metal_plans_.clear(); + for (auto& kv : metal_persistent_acts_) cactus_metal_free_shared(kv.second.first); + metal_persistent_acts_.clear(); + metal_retype_plan_.clear(); + metal_retype_built_ = false; +} + +void CactusGraph::hard_reset() { + invalidate_metal_state(); + nodes_.clear(); + node_index_map_.clear(); + mapped_files_.clear(); + weight_cache_.clear(); + next_node_id_ = 1; + debug_nodes_.clear(); + buffer_pool_.clear(); +} + +void CactusGraph::soft_reset() { + invalidate_metal_state(); + std::set cached_node_ids; + for (const auto& cache_entry : weight_cache_) { + cached_node_ids.insert(cache_entry.second); + } + + for (size_t pid : persistent_node_ids_) { + cached_node_ids.insert(pid); + } + + size_t max_preserved_id = 0; + for (const auto& node : nodes_) { + if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || + cached_node_ids.count(node->id)) { + max_preserved_id = std::max(max_preserved_id, node->id); + } + } + + auto preserved_nodes = std::move(nodes_); + auto preserved_index_map = std::move(node_index_map_); + + nodes_.clear(); + node_index_map_.clear(); + + for (auto& node : preserved_nodes) { + if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || + cached_node_ids.count(node->id)) { + size_t index = nodes_.size(); + node_index_map_[node->id] = index; + nodes_.push_back(std::move(node)); + } + } + + next_node_id_ = max_preserved_id + 1; + debug_nodes_.clear(); + if (!prefill_mode_) { + buffer_pool_.clear(); + shrink_thread_local_buffers(); + } +} + +void CactusGraph::soft_reset_keep_pool() { + invalidate_metal_state(); + std::set cached_node_ids; + for (const auto& cache_entry : weight_cache_) { + cached_node_ids.insert(cache_entry.second); + } + + for (size_t pid : persistent_node_ids_) { + cached_node_ids.insert(pid); + } + + size_t max_preserved_id = 0; + for (const auto& node : nodes_) { + if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || + cached_node_ids.count(node->id)) { + max_preserved_id = std::max(max_preserved_id, node->id); + } + } + + auto preserved_nodes = std::move(nodes_); + + nodes_.clear(); + node_index_map_.clear(); + + for (auto& node : preserved_nodes) { + if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || + cached_node_ids.count(node->id)) { + size_t index = nodes_.size(); + node_index_map_[node->id] = index; + nodes_.push_back(std::move(node)); + } + } + + next_node_id_ = max_preserved_id + 1; + debug_nodes_.clear(); +} + +void CactusGraph::prewarm_metal_quant_weights() { + if (!cactus_metal_available()) return; + for (auto& np : nodes_) { + GraphNode& node = *np; + if (node.params.backend != ComputeBackend::METAL) continue; + if (node.op_type != OpType::MATMUL || node.input_ids.size() < 2) continue; + auto it = node_index_map_.find(node.input_ids[1]); + if (it == node_index_map_.end()) continue; + const BufferDesc& rhs = nodes_[it->second]->output_buffer; + if (!(PrecisionTraits::is_cq(rhs.precision) && rhs.group_size > 0)) continue; + CactusQuantMatrix mat = rhs.to_cq_matrix(); + cactus_metal_prewarm_quant(&mat); + } +} diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp new file mode 100644 index 000000000..99cc069bb --- /dev/null +++ b/cactus-graph/src/graph_ffi.cpp @@ -0,0 +1,1532 @@ +#include "../cactus_graph.h" +#include +#include + +struct GraphHandle { + CactusGraph graph; +}; + +static GraphHandle* as_graph(cactus_graph_t g) { + return reinterpret_cast(g); +} + +static int fail_invalid(const char* msg) { + last_error_message = msg; + return -1; +} + +#define GRAPH_WRAP_UNARY(fn_name, method_name) \ +int fn_name(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { \ + if (!graph || !out) return fail_invalid("Invalid args to " #fn_name); \ + try { \ + *out = static_cast(as_graph(graph)->graph.method_name(static_cast(x))); \ + return 0; \ + } catch (const std::exception& e) { \ + last_error_message = e.what(); \ + return -1; \ + } \ +} + +#define GRAPH_WRAP_BINARY(fn_name, method_name) \ +int fn_name(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { \ + if (!graph || !out) return fail_invalid("Invalid args to " #fn_name); \ + try { \ + *out = static_cast(as_graph(graph)->graph.method_name(static_cast(a), static_cast(b))); \ + return 0; \ + } catch (const std::exception& e) { \ + last_error_message = e.what(); \ + return -1; \ + } \ +} + +extern "C" { + +cactus_graph_t cactus_graph_create(void) { + try { + return reinterpret_cast(new GraphHandle()); + } catch (const std::exception& e) { + last_error_message = e.what(); + return nullptr; + } catch (...) { + last_error_message = "Unknown error creating graph"; + return nullptr; + } +} + +void cactus_graph_destroy(cactus_graph_t graph) { + delete as_graph(graph); +} + +int cactus_graph_hard_reset(cactus_graph_t graph) { + if (!graph) { + last_error_message = "Invalid args to cactus_graph_hard_reset"; + return -1; + } + try { + as_graph(graph)->graph.hard_reset(); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_save(cactus_graph_t graph, const char* filename) { + if (!graph || !filename) { + last_error_message = "Invalid args to cactus_graph_save"; + return -1; + } + try { + as_graph(graph)->graph.save(filename); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + + +cactus_graph_t cactus_graph_load(const char* filename) { + if (!filename) { + last_error_message = "Invalid args to cactus_graph_load"; + return nullptr; + } + try { + auto handle = std::make_unique(); + handle->graph = CactusGraph::load(std::string(filename)); + return reinterpret_cast(handle.release()); + } catch (const std::exception& e) { + last_error_message = e.what(); + return nullptr; + } catch (...) { + last_error_message = "Unknown error loading graph"; + return nullptr; + } +} + +int cactus_graph_input(cactus_graph_t graph, const size_t* shape, size_t rank, int32_t precision, cactus_node_t* out_node) { + if (!graph || !shape || rank == 0 || !out_node) { + last_error_message = "Invalid args to cactus_graph_input"; + return -1; + } + try { + std::vector s(shape, shape + rank); + auto id = as_graph(graph)->graph.input(s, static_cast(precision)); + *out_node = static_cast(id); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_set_input(cactus_graph_t graph, cactus_node_t node, const void* data, int32_t precision) { + if (!graph || !data) { + last_error_message = "Invalid args to cactus_graph_set_input"; + return -1; + } + try { + as_graph(graph)->graph.set_input(static_cast(node), data, static_cast(precision)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_set_external_input(cactus_graph_t graph, cactus_node_t node, void* data, int32_t precision) { + if (!graph || !data) { + return fail_invalid("Invalid args to cactus_graph_set_external_input"); + } + try { + as_graph(graph)->graph.set_external_input(static_cast(node), data, static_cast(precision)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_set_runtime_input_shape(cactus_graph_t graph, cactus_node_t node, const size_t* shape, size_t rank) { + if (!graph || !shape) { + return fail_invalid("Invalid args to cactus_graph_set_runtime_input_shape"); + } + try { + as_graph(graph)->graph.set_runtime_input_shape(static_cast(node), std::vector(shape, shape + rank)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_set_input_dynamic_dims(cactus_graph_t graph, cactus_node_t node, const uint8_t* mask, size_t rank) { + if (!graph || !mask) { + return fail_invalid("Invalid args to cactus_graph_set_input_dynamic_dims"); + } + try { + as_graph(graph)->graph.set_input_dynamic_dims(static_cast(node), std::vector(mask, mask + rank)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_mark_embedded_input(cactus_graph_t graph, cactus_node_t node) { + if (!graph) { + return fail_invalid("Invalid args to cactus_graph_mark_embedded_input"); + } + try { + as_graph(graph)->graph.mark_embedded_input(static_cast(node)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_precision_cast(cactus_graph_t graph, cactus_node_t input, int32_t target_precision, cactus_node_t* out) { + if (!graph || !out) { + return fail_invalid("Invalid args to cactus_graph_precision_cast"); + } + try { + *out = static_cast(as_graph(graph)->graph.precision_cast(static_cast(input), static_cast(target_precision))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + + +int cactus_graph_add(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_add"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.add(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_add_clipped(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_add_clipped"); + try { + *out = static_cast(as_graph(graph)->graph.add_clipped(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_subtract(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_subtract"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.subtract(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_multiply(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_multiply"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.multiply(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_divide(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_divide"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.divide(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_not_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_not_equal"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.not_equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_scalar_add"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.scalar_add(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_subtract(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_subtract"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_subtract(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_multiply(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_scalar_multiply"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.scalar_multiply(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_scalar_divide"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.scalar_divide(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_scalar_not_equal"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.scalar_not_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_exp(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_scalar_exp"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.scalar_exp(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_sqrt(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_scalar_sqrt"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.scalar_sqrt(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_cos(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_cos"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_cos(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_sin(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_sin"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_sin(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scalar_log(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_log"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_log(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_abs(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_abs"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.abs(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_pow(cactus_graph_t graph, cactus_node_t x, float exponent, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_pow"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.pow(static_cast(x), exponent)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +static bool valid_backend(int32_t backend) { + switch (static_cast(backend)) { + case ComputeBackend::CPU: + case ComputeBackend::METAL: + return true; + } + return false; +} + +int cactus_graph_set_node_backend(cactus_graph_t graph, cactus_node_t node, int32_t backend) { + if (!graph || !valid_backend(backend)) { + last_error_message = "Invalid args to cactus_graph_set_node_backend"; + return -1; + } + try { + as_graph(graph)->graph.set_node_backend(static_cast(node), static_cast(backend)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_view(cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out) { + if (!graph || !shape || rank == 0 || !out) { + last_error_message = "Invalid args to cactus_graph_view"; + return -1; + } + try { + std::vector s(shape, shape + rank); + *out = static_cast(as_graph(graph)->graph.view(static_cast(x), s)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_flatten(cactus_graph_t graph, cactus_node_t x, int32_t start_dim, int32_t end_dim, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_flatten"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.flatten(static_cast(x), start_dim, end_dim)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_reshape(cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out) { + if (!graph || !shape || rank == 0 || !out) return fail_invalid("Invalid args to cactus_graph_reshape"); + try { + std::vector s(shape, shape + rank); + *out = static_cast(as_graph(graph)->graph.reshape(static_cast(x), s)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_transpose(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_transpose"); + try { + *out = static_cast(as_graph(graph)->graph.transpose(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_transpose_n(cactus_graph_t graph, cactus_node_t x, const size_t* permutation, size_t rank, cactus_node_t* out) { + if (!graph || !permutation || rank == 0 || !out) return fail_invalid("Invalid args to cactus_graph_transpose_n"); + try { + std::vector p(permutation, permutation + rank); + *out = static_cast(as_graph(graph)->graph.transposeN(static_cast(x), p)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_slice(cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_slice"); + try { + *out = static_cast(as_graph(graph)->graph.slice(static_cast(x), axis, start, length)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_index(cactus_graph_t graph, cactus_node_t x, size_t index_value, int32_t dim, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_index"); + try { + *out = static_cast(as_graph(graph)->graph.index(static_cast(x), index_value, dim)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_sum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_sum"); + try { + *out = static_cast(as_graph(graph)->graph.sum(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_mean(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_mean"); + try { + *out = static_cast(as_graph(graph)->graph.mean(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_variance(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_variance"); + try { + *out = static_cast(as_graph(graph)->graph.variance(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_min(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_min"); + try { + *out = static_cast(as_graph(graph)->graph.min(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_max(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_max"); + try { + *out = static_cast(as_graph(graph)->graph.max(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_cumsum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_cumsum"); + try { + *out = static_cast(as_graph(graph)->graph.cumsum(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_concat(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, int32_t axis, cactus_node_t* out) { + if (!graph || !out) { + last_error_message = "Invalid args to cactus_graph_concat"; + return -1; + } + try { + *out = static_cast(as_graph(graph)->graph.concat(static_cast(a), static_cast(b), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_cat(cactus_graph_t graph, const cactus_node_t* nodes, size_t count, int32_t axis, cactus_node_t* out) { + if (!graph || !nodes || !out || count == 0) { + last_error_message = "Invalid args to cactus_graph_cat"; + return -1; + } + try { + std::vector ns; + for (size_t i = 0; i < count; ++i) { + ns.push_back(static_cast(nodes[i])); + } + *out = static_cast(as_graph(graph)->graph.cat(ns, axis)); + return 0; + + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_matmul(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_matmul"); + try { + *out = static_cast(as_graph(graph)->graph.matmul(static_cast(a), static_cast(b), pretransposed_rhs)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_gather(cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gather"); + try { + *out = static_cast(as_graph(graph)->graph.gather(static_cast(tensor), static_cast(indices))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_embedding_from_tensor(cactus_graph_t graph, cactus_node_t embedding_tensor, cactus_node_t indices, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_embedding_from_tensor"); + try { + *out = static_cast(as_graph(graph)->graph.embedding(static_cast(embedding_tensor), static_cast(indices))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_embedding_from_file(cactus_graph_t graph, const char* filename, cactus_node_t indices, cactus_node_t* out) { + if (!graph || !filename || !out) return fail_invalid("Invalid args to cactus_graph_embedding_from_file"); + try { + *out = static_cast(as_graph(graph)->graph.embedding(std::string(filename), static_cast(indices))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_mmap_embeddings(cactus_graph_t graph, const char* filename, cactus_node_t* out) { + if (!graph || !filename || !out) return fail_invalid("Invalid args to cactus_graph_mmap_embeddings"); + try { + *out = static_cast(as_graph(graph)->graph.mmap_embeddings(std::string(filename))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_mmap_weights(cactus_graph_t graph, const char* filename, cactus_node_t* out) { + if (!graph || !filename || !out) return fail_invalid("Invalid args to cactus_graph_mmap_weights"); + try { + *out = static_cast(as_graph(graph)->graph.mmap_weights(std::string(filename))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_bind_mmap_weights(cactus_graph_t graph, cactus_node_t node, const char* filename) { + if (!graph || !filename) return fail_invalid("Invalid args to cactus_graph_bind_mmap_weights"); + try { + as_graph(graph)->graph.bind_mmap_weights(static_cast(node), std::string(filename)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_bilinear_interpolation(cactus_graph_t graph, cactus_node_t pos_embeds, size_t dst_height, size_t dst_width, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bilinear_interpolation"); + try { + *out = static_cast(as_graph(graph)->graph.bilinear_interpolation(static_cast(pos_embeds), dst_height, dst_width)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + + + +int cactus_graph_release_weight_pages(cactus_graph_t graph, cactus_node_t node) { + if (!graph) return fail_invalid("Invalid args to cactus_graph_release_weight_pages"); + try { + as_graph(graph)->graph.release_weight_pages(static_cast(node)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_prefetch_weight_pages(cactus_graph_t graph, cactus_node_t node) { + if (!graph) return fail_invalid("Invalid args to cactus_graph_prefetch_weight_pages"); + try { + as_graph(graph)->graph.prefetch_weight_pages(static_cast(node)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_release_all_weight_pages(cactus_graph_t graph) { + if (!graph) return fail_invalid("Invalid args to cactus_graph_release_all_weight_pages"); + try { + as_graph(graph)->graph.release_all_weight_pages(); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_relu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_relu"); + try { + *out = static_cast(as_graph(graph)->graph.relu(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_silu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_silu"); + try { + *out = static_cast(as_graph(graph)->graph.silu(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_gelu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gelu"); + try { + *out = static_cast(as_graph(graph)->graph.gelu(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_gelu_erf(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gelu_erf"); + try { + *out = static_cast(as_graph(graph)->graph.gelu_erf(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_sigmoid(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_sigmoid"); + try { + *out = static_cast(as_graph(graph)->graph.sigmoid(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_tanh(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_tanh"); + try { + *out = static_cast(as_graph(graph)->graph.tanh(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_glu(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_glu"); + try { + *out = static_cast(as_graph(graph)->graph.glu(static_cast(x), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_layernorm(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, float epsilon, bool has_bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_layernorm"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.layernorm(static_cast(input), static_cast(weight), static_cast(bias), epsilon)); + } else { + *out = static_cast(as_graph(graph)->graph.layernorm(static_cast(input), static_cast(weight), epsilon)); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_groupnorm(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, size_t num_groups, float epsilon, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_groupnorm"); + try { + *out = static_cast(as_graph(graph)->graph.groupnorm(static_cast(input), static_cast(weight), static_cast(bias), num_groups, epsilon)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_batchnorm(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, cactus_node_t running_mean, cactus_node_t running_var, int32_t axis, float epsilon, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_batchnorm"); + try { + *out = static_cast(as_graph(graph)->graph.batchnorm(static_cast(input), static_cast(weight), static_cast(bias), static_cast(running_mean), static_cast(running_var), axis, epsilon)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_topk(cactus_graph_t graph, cactus_node_t input, size_t k, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_topk"); + try { + *out = static_cast(as_graph(graph)->graph.topk(static_cast(input), k)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_rms_norm(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, float epsilon, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rms_norm"); + try { + *out = static_cast(as_graph(graph)->graph.rms_norm(static_cast(input), static_cast(weight), epsilon)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_rope(cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rope"); + try { + *out = static_cast(as_graph(graph)->graph.rope(static_cast(input), theta, position_offset)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_rope_gptj(cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, size_t rot_dim, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rope_gptj"); + try { + *out = static_cast(as_graph(graph)->graph.rope_gptj(static_cast(input), theta, position_offset, rot_dim)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_softmax(cactus_graph_t graph, cactus_node_t input, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_softmax"); + try { + *out = static_cast(as_graph(graph)->graph.softmax(static_cast(input), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_attention(cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_attention"); + try { + if (use_mask) { + *out = static_cast(as_graph(graph)->graph.attention_masked(static_cast(query), static_cast(key), static_cast(value), static_cast(mask), scale, is_causal, cactus_default_backend(), additive_mask, position_offset, window_size)); + } else if (window_size > 0 || position_offset > 0) { + *out = static_cast(as_graph(graph)->graph.attention(static_cast(query), static_cast(key), static_cast(value), scale, position_offset, window_size)); + } else { + *out = static_cast(as_graph(graph)->graph.attention(static_cast(query), static_cast(key), static_cast(value), scale, is_causal)); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_rel_pos_bias(cactus_graph_t graph, cactus_node_t query, cactus_node_t relative_key, float scale, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rel_pos_bias"); + try { + *out = static_cast(as_graph(graph)->graph.rel_pos_bias(static_cast(query), static_cast(relative_key), scale)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_attention_int8_hybrid(cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, float scale, size_t position_offset, + const int8_t* cached_keys, const int8_t* cached_values, const float* k_scales, const float* v_scales, + size_t cache_len, size_t num_kv_heads, size_t head_dim, size_t window_size, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_attention_int8_hybrid"); + try { + *out = static_cast(as_graph(graph)->graph.attention_int8_hybrid(static_cast(query), static_cast(key_new), static_cast(value_new), scale, position_offset, cached_keys, cached_values, k_scales, v_scales, cache_len, num_kv_heads, head_dim, window_size)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_kv_cache_state(cactus_graph_t graph, size_t max_seq_len, size_t num_kv_heads, size_t head_dim, size_t window_size, size_t sink_size, size_t num_slots, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_kv_cache_state"); + try { + *out = static_cast(as_graph(graph)->graph.kv_cache_state(max_seq_len, num_kv_heads, head_dim, window_size, sink_size, num_slots)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_kv_cache_append(cactus_graph_t graph, cactus_node_t new_kv, cactus_node_t cache_state, size_t window_size, size_t sink_size, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_kv_cache_append"); + try { + *out = static_cast(as_graph(graph)->graph.kv_cache_append(static_cast(new_kv), static_cast(cache_state), window_size, sink_size)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_attention_cached(cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, + cactus_node_t k_cache_state, cactus_node_t v_cache_state, + float scale, size_t position_offset, size_t window_size, size_t v_head_dim, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_attention_cached"); + try { + *out = static_cast(as_graph(graph)->graph.attention_cached( + static_cast(query), static_cast(key_new), static_cast(value_new), + static_cast(k_cache_state), static_cast(v_cache_state), + scale, position_offset, window_size, v_head_dim)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv_cache_state(cactus_graph_t graph, size_t window_size, size_t hidden_dim, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv_cache_state"); + try { + *out = static_cast(as_graph(graph)->graph.conv_cache_state(window_size, hidden_dim)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv_cache_append(cactus_graph_t graph, cactus_node_t new_data, cactus_node_t cache_state, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv_cache_append"); + try { + *out = static_cast(as_graph(graph)->graph.conv_cache_append(static_cast(new_data), static_cast(cache_state))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv_cache_initialize(cactus_graph_t graph, cactus_node_t rows, cactus_node_t cache_state, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv_cache_initialize"); + try { + *out = static_cast(as_graph(graph)->graph.conv_cache_initialize(static_cast(rows), static_cast(cache_state))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_recurrent_cache_state(cactus_graph_t graph, const size_t* shape, size_t shape_len, int precision, cactus_node_t* out) { + if (!graph || !out || (!shape && shape_len > 0)) { + return fail_invalid("Invalid args to cactus_graph_recurrent_cache_state"); + } + try { + std::vector shape_vec(shape, shape + shape_len); + *out = static_cast(as_graph(graph)->graph.recurrent_cache_state(shape_vec, static_cast(precision))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_recurrent_cache_write(cactus_graph_t graph, cactus_node_t new_value, cactus_node_t cache_input, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_recurrent_cache_write"); + try { + *out = static_cast(as_graph(graph)->graph.recurrent_cache_write(static_cast(new_value), static_cast(cache_input))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_rfft(cactus_graph_t graph, cactus_node_t input, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rfft"); + try { + *out = static_cast(as_graph(graph)->graph.rfft(static_cast(input))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_irfft(cactus_graph_t graph, cactus_node_t input, size_t output_length, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_irfft"); + try { + *out = static_cast(as_graph(graph)->graph.irfft(static_cast(input), output_length)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_mel_filter_bank(cactus_graph_t graph, size_t num_frequency_bins, size_t num_mel_filters, + float min_frequency, float max_frequency, size_t sampling_rate, + int norm_type, int scale_type, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_mel_filter_bank"); + try { + *out = static_cast(as_graph(graph)->graph.mel_filter_bank( + num_frequency_bins, num_mel_filters, min_frequency, max_frequency, sampling_rate, norm_type, scale_type)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_spectrogram(cactus_graph_t graph, cactus_node_t waveform, cactus_node_t mel_filters, + size_t frame_length, size_t hop_length, size_t fft_length, + float power, bool center, int pad_mode, + float mel_floor, int log_mel_mode, + float dither, float preemphasis, bool remove_dc_offset, + cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_spectrogram"); + try { + *out = static_cast(as_graph(graph)->graph.spectrogram( + static_cast(waveform), static_cast(mel_filters), + frame_length, hop_length, fft_length, + power, center, pad_mode, mel_floor, log_mel_mode, + dither, preemphasis, remove_dc_offset)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_image_preprocess(cactus_graph_t graph, cactus_node_t pixel_input, + int src_width, int src_height, int target_width, int target_height, + int patch_size, int channels, float rescale_factor, + const float* mean, const float* std_dev, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_image_preprocess"); + try { + *out = static_cast(as_graph(graph)->graph.image_preprocess( + static_cast(pixel_input), + src_width, src_height, target_width, target_height, + patch_size, channels, rescale_factor, mean, std_dev)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv1d_causal(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t kernel_size, size_t dilation, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_causal"); + try { + *out = static_cast(as_graph(graph)->graph.conv1d_causal(static_cast(input), static_cast(weight), kernel_size, dilation)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv1d_k3(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t stride, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_k3"); + try { + *out = static_cast(as_graph(graph)->graph.conv1d_k3(static_cast(input), static_cast(weight), stride)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv1d_k7s3(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_k7s3"); + try { + *out = static_cast(as_graph(graph)->graph.conv1d_k7s3(static_cast(input), static_cast(weight), static_cast(bias))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv1d(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv1d(static_cast(input), static_cast(weight), static_cast(bias), stride)); + } else { + *out = static_cast(as_graph(graph)->graph.conv1d(static_cast(input), static_cast(weight), stride)); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv1d_same_depthwise_k9(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_same_depthwise_k9"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv1d_same_depthwise_k9(static_cast(input), static_cast(weight), static_cast(bias))); + } else { + *out = static_cast(as_graph(graph)->graph.conv1d_same_depthwise_k9(static_cast(input), static_cast(weight))); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv1d_pointwise(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_pointwise"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv1d_pointwise(static_cast(input), static_cast(weight), static_cast(bias))); + } else { + *out = static_cast(as_graph(graph)->graph.conv1d_pointwise(static_cast(input), static_cast(weight))); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_clamp(cactus_graph_t graph, cactus_node_t input, float lo, float hi, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_clamp"); + try { + *out = static_cast(as_graph(graph)->graph.clamp(static_cast(input), lo, hi)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv2d_k3s2p1(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv2d_k3s2p1"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv2d_k3s2p1(static_cast(input), static_cast(weight), static_cast(bias))); + } else { + *out = static_cast(as_graph(graph)->graph.conv2d_k3s2p1(static_cast(input), static_cast(weight))); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv2d_depthwise_k3s2p1(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv2d_depthwise_k3s2p1"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv2d_depthwise_k3s2p1(static_cast(input), static_cast(weight), static_cast(bias))); + } else { + *out = static_cast(as_graph(graph)->graph.conv2d_depthwise_k3s2p1(static_cast(input), static_cast(weight))); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_conv2d_pointwise_1x1(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv2d_pointwise_1x1"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv2d_pointwise_1x1(static_cast(input), static_cast(weight), static_cast(bias))); + } else { + *out = static_cast(as_graph(graph)->graph.conv2d_pointwise_1x1(static_cast(input), static_cast(weight))); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_lstm_cell(cactus_graph_t graph, cactus_node_t input, cactus_node_t h_prev, cactus_node_t c_prev, cactus_node_t weight_ih, cactus_node_t weight_hh, cactus_node_t bias_ih, cactus_node_t bias_hh, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_lstm_cell"); + try { + *out = static_cast(as_graph(graph)->graph.lstm_cell(static_cast(input), static_cast(h_prev), static_cast(c_prev), static_cast(weight_ih), static_cast(weight_hh), static_cast(bias_ih), static_cast(bias_hh))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_gated_deltanet_decode(cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, cactus_node_t gate_log, cactus_node_t beta, cactus_node_t initial_state, float scale, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gated_deltanet_decode"); + try { + *out = static_cast(as_graph(graph)->graph.gated_deltanet_decode(static_cast(query), static_cast(key), static_cast(value), static_cast(gate_log), static_cast(beta), static_cast(initial_state), scale)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_gated_deltanet_prefill(cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, cactus_node_t gate_log, cactus_node_t beta, cactus_node_t initial_state, size_t chunk_size, float scale, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gated_deltanet_prefill"); + try { + *out = static_cast(as_graph(graph)->graph.gated_deltanet_prefill(static_cast(query), static_cast(key), static_cast(value), static_cast(gate_log), static_cast(beta), static_cast(initial_state), chunk_size, scale)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_stft(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t stride, size_t num_fft_bins, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_stft"); + try { + *out = static_cast(as_graph(graph)->graph.stft(static_cast(input), static_cast(weight), stride, num_fft_bins)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_altup_predict(cactus_graph_t graph, cactus_node_t coefs, const cactus_node_t* streams, size_t num_streams, cactus_node_t* out) { + if (!graph || !streams || num_streams == 0 || !out) return fail_invalid("Invalid args to cactus_graph_altup_predict"); + try { + std::vector s(num_streams); + for (size_t i = 0; i < num_streams; ++i) s[i] = static_cast(streams[i]); + *out = static_cast(as_graph(graph)->graph.altup_predict(static_cast(coefs), s.data(), num_streams)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_altup_correct(cactus_graph_t graph, cactus_node_t coefs, cactus_node_t innovation, const cactus_node_t* predictions, size_t num_predictions, cactus_node_t* out) { + if (!graph || !predictions || num_predictions == 0 || !out) return fail_invalid("Invalid args to cactus_graph_altup_correct"); + try { + std::vector p(num_predictions); + for (size_t i = 0; i < num_predictions; ++i) p[i] = static_cast(predictions[i]); + *out = static_cast(as_graph(graph)->graph.altup_correct(static_cast(coefs), static_cast(innovation), p.data(), num_predictions)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_gaussian_topk(cactus_graph_t graph, cactus_node_t input, float ppf, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gaussian_topk"); + try { + *out = static_cast(as_graph(graph)->graph.gaussian_topk(static_cast(input), ppf)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_moe_layer_gated(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, + const cactus_node_t* w1_weights, const cactus_node_t* w3_weights, const cactus_node_t* w2_weights, + size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out) { + if (!graph || !w1_weights || !w3_weights || !w2_weights || !out) return fail_invalid("Invalid args to cactus_graph_moe_layer_gated"); + try { + std::vector w1(num_experts), w3(num_experts), w2(num_experts); + for (size_t i = 0; i < num_experts; ++i) { + w1[i] = static_cast(w1_weights[i]); + w3[i] = static_cast(w3_weights[i]); + w2[i] = static_cast(w2_weights[i]); + } + *out = static_cast(as_graph(graph)->graph.moe_layer(static_cast(hidden), static_cast(routing_probs), static_cast(topk_indices), w1, w3, w2, num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor, static_cast(activation))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t gate_weight, cactus_node_t up_weight, cactus_node_t down_weight, float product_scale, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_dense_mlp_tq_fused"); + try { + *out = static_cast( + as_graph(graph)->graph.dense_mlp_tq_fused( + static_cast(hidden), + static_cast(gate_weight), + static_cast(up_weight), + static_cast(down_weight), + product_scale + ) + ); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_moe_layer_ungated(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, + const cactus_node_t* w1_weights, const cactus_node_t* w2_weights, + size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out) { + if (!graph || !w1_weights || !w2_weights || !out) return fail_invalid("Invalid args to cactus_graph_moe_layer_ungated"); + try { + std::vector w1(num_experts), w2(num_experts); + for (size_t i = 0; i < num_experts; ++i) { + w1[i] = static_cast(w1_weights[i]); + w2[i] = static_cast(w2_weights[i]); + } + *out = static_cast(as_graph(graph)->graph.moe_layer(static_cast(hidden), static_cast(routing_probs), static_cast(topk_indices), w1, w2, num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor, static_cast(activation))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_sample(cactus_graph_t graph, cactus_node_t logits, float temperature, float top_p, size_t top_k, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_sample"); + try { + std::unordered_map empty_bias; + *out = static_cast(as_graph(graph)->graph.sample(static_cast(logits), temperature, top_p, top_k, empty_bias)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_scatter_topk(cactus_graph_t graph, cactus_node_t indices, cactus_node_t values, size_t num_classes, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scatter_topk"); + try { + *out = static_cast(as_graph(graph)->graph.scatter_topk(static_cast(indices), static_cast(values), num_classes)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_persistent(cactus_graph_t graph, cactus_node_t source_node, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_persistent"); + try { + *out = static_cast(as_graph(graph)->graph.persistent(static_cast(source_node))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_is_populated(cactus_graph_t graph, cactus_node_t persistent_node, int32_t* out_is_populated) { + if (!graph || !out_is_populated) return fail_invalid("Invalid args to cactus_graph_is_populated"); + try { + *out_is_populated = as_graph(graph)->graph.is_populated(static_cast(persistent_node)) ? 1 : 0; + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_invalidate_persistent(cactus_graph_t graph, cactus_node_t persistent_node) { + if (!graph) return fail_invalid("Invalid args to cactus_graph_invalidate_persistent"); + try { + as_graph(graph)->graph.invalidate_persistent(static_cast(persistent_node)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_execute(cactus_graph_t graph) { + if (!graph) { + last_error_message = "Graph is null"; + return -1; + } + try { + as_graph(graph)->graph.execute(); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_get_output_ptr(cactus_graph_t graph, cactus_node_t node, void** out_ptr) { + if (!graph || !out_ptr) { + last_error_message = "Invalid args to cactus_graph_get_output_ptr"; + return -1; + } + try { + *out_ptr = as_graph(graph)->graph.get_output(static_cast(node)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_get_output_info(cactus_graph_t graph, cactus_node_t node, cactus_tensor_info_t* out_info) { + if (!graph || !out_info) { + last_error_message = "Invalid args to cactus_graph_get_output_info"; + return -1; + } + try { + const auto& buf = as_graph(graph)->graph.get_output_buffer(static_cast(node)); + out_info->precision = static_cast(buf.precision); + out_info->rank = buf.shape.size(); + if (out_info->rank > 8) { + last_error_message = "Rank exceeds cactus_tensor_info_t shape capacity"; + return -1; + } + for (size_t i = 0; i < out_info->rank; ++i) { + out_info->shape[i] = buf.shape[i]; + } + for (size_t i = out_info->rank; i < 8; ++i) { + out_info->shape[i] = 0; + } + out_info->num_elements = buf.total_size; + out_info->byte_size = buf.byte_size; + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} +} diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp new file mode 100644 index 000000000..147eeb98b --- /dev/null +++ b/cactus-graph/src/io.cpp @@ -0,0 +1,1104 @@ +#include "../cactus_graph.h" +#include "param_io.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + constexpr uint32_t fourcc(char a, char b, char c, char d) { + return static_cast(static_cast(a)) | + (static_cast(static_cast(b)) << 8) | + (static_cast(static_cast(c)) << 16) | + (static_cast(static_cast(d)) << 24); + } + + constexpr uint32_t CACTUS_MAGIC = 0x54434143; + constexpr uint32_t CACTUS_GRAPH_MAGIC = fourcc('C', 'G', 'R', 'F'); + constexpr uint32_t CACTUS_GRAPH_VERSION_LEGACY = 4; + constexpr uint32_t CACTUS_GRAPH_VERSION_EMBEDDED_INPUTS = 5; + constexpr uint32_t CACTUS_GRAPH_VERSION_DYNAMIC_SHAPES = 6; + constexpr uint32_t FLAG_ORTHOGONAL_ROTATION = 1 << 1; + constexpr uint32_t FLAG_INTERLEAVED_4ROW = 1 << 2; + constexpr uint32_t FLAG_EXTENDED_SHAPE = 1 << 4; + constexpr size_t HEADER_SIZE = 84; + + inline size_t align_offset(size_t offset, size_t alignment) { + size_t remainder = offset % alignment; + if (remainder == 0) return offset; + return offset + (alignment - remainder); + } + + std::string resolve_quantized_weight_file(const std::string& filename) { + constexpr const char* suffix = ".weights"; + if (filename.size() <= std::strlen(suffix) || + filename.compare(filename.size() - std::strlen(suffix), std::strlen(suffix), suffix) != 0) { + return filename; + } + + const std::string stem = filename.substr(0, filename.size() - std::strlen(suffix)); + const std::string cq4 = stem + ".cq4.weights"; + if (std::filesystem::exists(cq4)) { + return cq4; + } + const std::string cq2 = stem + ".cq2.weights"; + if (std::filesystem::exists(cq2)) { + return cq2; + } + return filename; + } + + inline void write_u32(std::ostream& out, uint32_t v) { + out.write(reinterpret_cast(&v), sizeof(v)); + } + + inline void write_u64(std::ostream& out, uint64_t v) { + out.write(reinterpret_cast(&v), sizeof(v)); + } + + void write_size_vector(std::ostream& out, const std::vector& values) { + uint32_t size = static_cast(values.size()); + write_u32(out, size); + for (size_t v : values) { + write_u64(out, static_cast(v)); + } + } + + void write_u32_vector(std::ostream& out, const std::vector& values) { + uint32_t size = static_cast(values.size()); + write_u32(out, size); + for (uint32_t v : values) { + write_u32(out, v); + } + } + + std::vector compute_leaf_outputs(const GraphFile::SerializedGraph& sg) { + std::unordered_set node_ids; + std::unordered_set referenced; + node_ids.reserve(sg.nodes.size()); + referenced.reserve(sg.nodes.size()); + + for (const auto& node : sg.nodes) { + node_ids.insert(node.index); + } + + for (const auto& node : sg.nodes) { + for (uint32_t input_idx : node.inputs) { + if (node_ids.find(input_idx) == node_ids.end()) { + throw std::runtime_error("Graph save failed: input node id not found"); + } + referenced.insert(input_idx); + } + } + + std::vector outputs; + outputs.reserve(sg.nodes.size()); + + for (const auto& node : sg.nodes) { + if (referenced.find(node.index) == referenced.end()) { + outputs.push_back(node.index); + } + } + + return outputs; + } + + void write_serialized_graph(std::ostream& out, const GraphFile::SerializedGraph& sg) { + write_u32(out, sg.header.magic); + write_u32(out, sg.header.version); + write_u32(out, sg.header.node_count); + write_u32(out, sg.header.flags); + + write_u32_vector(out, sg.graph_inputs); + write_u32_vector(out, sg.graph_outputs); + + for (const auto& node : sg.nodes) { + write_u32(out, node.index); + write_u32(out, static_cast(node.op_type)); + write_u32_vector(out, node.inputs); + write_size_vector(out, node.output_shape); + write_u32(out, static_cast(node.precision)); + GraphParamIO::write_op_params(out, node.op_type, node.params); + if (sg.header.version >= CACTUS_GRAPH_VERSION_EMBEDDED_INPUTS) { + write_u32(out, node.has_embedded_data ? 1u : 0u); + write_u64(out, static_cast(node.embedded_data.size())); + if (!node.embedded_data.empty()) { + out.write( + reinterpret_cast(node.embedded_data.data()), + static_cast(node.embedded_data.size()) + ); + } + } + if (sg.header.version >= CACTUS_GRAPH_VERSION_DYNAMIC_SHAPES) { + write_u32(out, static_cast(node.dynamic_mask.size())); + if (!node.dynamic_mask.empty()) { + out.write( + reinterpret_cast(node.dynamic_mask.data()), + static_cast(node.dynamic_mask.size()) + ); + } + } + } + + if (!out) { + throw std::runtime_error("Error writing serialized graph"); + } + } + + // read helpers + uint32_t read_u32(std::istream& in) { + uint32_t v = 0; + in.read(reinterpret_cast(&v), sizeof(v)); + if (!in) { + throw std::runtime_error("Unexpected EOF while reading uint32"); + } + return v; + } + + uint64_t read_u64(std::istream& in) { + uint64_t v = 0; + in.read(reinterpret_cast(&v), sizeof(v)); + if (!in) { + throw std::runtime_error("Unexpected EOF while reading uint64"); + } + return v; + } + + std::vector read_u32_vector(std::istream& in) { + uint32_t count = read_u32(in); + std::vector values; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + values.push_back(read_u32(in)); + } + return values; + } + + std::vector read_size_vector(std::istream& in) { + uint32_t count = read_u32(in); + std::vector values; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + values.push_back(static_cast(read_u64(in))); + } + return values; + } + + bool is_binary_broadcast_op(OpType op_type) { + switch (op_type) { + case OpType::ADD: + case OpType::ADD_CLIPPED: + case OpType::SUBTRACT: + case OpType::MULTIPLY: + case OpType::DIVIDE: + case OpType::NOT_EQUAL: + return true; + default: + return false; + } + } + + void populate_derived_params(CactusGraph& graph, + const GraphFile::NodeEntry& node_entry, + const std::vector& runtime_inputs, + OpParams& params) { + if (is_binary_broadcast_op(node_entry.op_type)) { + if (runtime_inputs.size() != 2) { + throw std::runtime_error("Graph file corrupted: binary op missing inputs"); + } + const auto& lhs = graph.get_output_buffer(runtime_inputs[0]); + const auto& rhs = graph.get_output_buffer(runtime_inputs[1]); + params.broadcast_info = BroadcastInfo::compute(lhs.shape, rhs.shape); + } + } + + GraphFile::GraphHeader read_graph_header(std::istream& in) { + GraphFile::GraphHeader header; + header.magic = read_u32(in); + header.version = read_u32(in); + header.node_count = read_u32(in); + header.flags = read_u32(in); + + if (header.magic != CACTUS_GRAPH_MAGIC) { + throw std::runtime_error("Invalid graph file: bad magic"); + } + if (header.version != CACTUS_GRAPH_VERSION_LEGACY && + header.version != CACTUS_GRAPH_VERSION_EMBEDDED_INPUTS && + header.version != CACTUS_GRAPH_VERSION_DYNAMIC_SHAPES) { + throw std::runtime_error("Unsupported graph file version: " + + std::to_string(header.version)); + } + + return header; + } + + GraphFile::NodeEntry read_node_entry(std::istream& in, uint32_t graph_version) { + GraphFile::NodeEntry node; + node.index = read_u32(in); + uint32_t op_type_val = read_u32(in); + if (op_type_val > static_cast(OpType::CONV_CACHE_INITIALIZE)) { + throw std::runtime_error("Graph file corrupted: invalid op type"); + } + node.op_type = static_cast(op_type_val); + node.inputs = read_u32_vector(in); + node.output_shape = read_size_vector(in); + uint32_t precision_val = read_u32(in); + if (precision_val > static_cast(Precision::CQ4)) { + throw std::runtime_error("Graph file corrupted: invalid precision"); + } + node.precision = static_cast(precision_val); + GraphParamIO::read_op_params(in, node.op_type, node.params); + if (graph_version >= CACTUS_GRAPH_VERSION_EMBEDDED_INPUTS) { + node.has_embedded_data = read_u32(in) != 0; + uint64_t data_size = read_u64(in); + if (data_size > 0) { + node.embedded_data.resize(static_cast(data_size)); + in.read( + reinterpret_cast(node.embedded_data.data()), + static_cast(node.embedded_data.size()) + ); + if (!in) { + throw std::runtime_error("Unexpected EOF while reading embedded graph input data"); + } + } + } + if (graph_version >= CACTUS_GRAPH_VERSION_DYNAMIC_SHAPES) { + uint32_t mask_size = read_u32(in); + if (mask_size > 0) { + node.dynamic_mask.resize(mask_size); + in.read(reinterpret_cast(node.dynamic_mask.data()), + static_cast(mask_size)); + if (!in) { + throw std::runtime_error("Unexpected EOF while reading dynamic mask"); + } + } + } + return node; + } + + +} // namespace + +void CactusGraph::save(const std::string& path) { + GraphFile::save_graph(*this, path); +} + +void CactusGraph::mark_embedded_input(size_t node_id) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) { + throw std::out_of_range("Unknown input node id: " + std::to_string(node_id)); + } + const auto& node = *nodes_[it->second]; + if (node.op_type != OpType::INPUT) { + throw std::invalid_argument("Can only embed input nodes"); + } + embedded_input_node_ids_.insert(node_id); +} + +bool CactusGraph::is_embedded_input(size_t node_id) const { + return embedded_input_node_ids_.find(node_id) != embedded_input_node_ids_.end(); +} + +CactusGraph CactusGraph::from_serialized(const GraphFile::SerializedGraph& sg) { + CactusGraph graph; + std::vector runtime_ids; + runtime_ids.reserve(sg.nodes.size()); + std::unordered_map serialized_id_to_runtime_id; + serialized_id_to_runtime_id.reserve(sg.nodes.size()); + + if (sg.nodes.size() != sg.header.node_count) { + throw std::runtime_error("Graph file corrupted: node count mismatch"); + } + + bool dense_legacy_indices = true; + for (size_t i = 0; i < sg.nodes.size(); ++i) { + if (sg.nodes[i].index != i) { + dense_legacy_indices = false; + break; + } + } + + for (size_t i = 0; i < sg.nodes.size(); ++i) { + const auto& node_entry = sg.nodes[i]; + + std::vector runtime_inputs; + runtime_inputs.reserve(node_entry.inputs.size()); + + for (uint32_t serialized_input_id : node_entry.inputs) { + if (dense_legacy_indices) { + if (serialized_input_id >= runtime_ids.size()) { + throw std::runtime_error( + "Graph file corrupted: input refers to a node that has not been reconstructed yet" + ); + } + runtime_inputs.push_back(runtime_ids[serialized_input_id]); + } else { + auto it = serialized_id_to_runtime_id.find(serialized_input_id); + if (it == serialized_id_to_runtime_id.end()) { + throw std::runtime_error( + "Graph file corrupted: input refers to a node that has not been reconstructed yet" + ); + } + runtime_inputs.push_back(it->second); + } + } + + size_t new_node_id = 0; + + if (!dense_legacy_indices) { + graph.next_node_id_ = static_cast(node_entry.index); + } + + if (node_entry.op_type == OpType::INPUT) { + new_node_id = graph.input(node_entry.output_shape, node_entry.precision); + if (node_entry.has_embedded_data) { + const auto& buffer = graph.get_output_buffer(new_node_id); + if (node_entry.embedded_data.size() != buffer.byte_size) { + throw std::runtime_error( + "Graph file corrupted: embedded input byte-size mismatch for node " + + std::to_string(node_entry.index) + ); + } + graph.set_input(new_node_id, node_entry.embedded_data.data(), node_entry.precision); + graph.mark_embedded_input(new_node_id); + } + } + else { + OpParams params = node_entry.params; + params.output_precision = node_entry.precision; + populate_derived_params(graph, node_entry, runtime_inputs, params); + new_node_id = graph.add_node(node_entry.op_type, runtime_inputs, node_entry.output_shape, params); + + if (node_entry.op_type == OpType::PERSISTENT + || node_entry.op_type == OpType::KV_CACHE_STATE + || node_entry.op_type == OpType::CONV_CACHE_STATE + || node_entry.op_type == OpType::RECURRENT_CACHE_STATE) { + graph.persistent_node_ids_.insert(new_node_id); + } + } + if (!node_entry.dynamic_mask.empty()) { + graph.nodes_[graph.node_index_map_.at(new_node_id)]->output_buffer.dynamic_dims = node_entry.dynamic_mask; + graph.has_dynamic_shapes_ = true; + } + runtime_ids.push_back(new_node_id); + serialized_id_to_runtime_id[node_entry.index] = new_node_id; + } + return graph; +} + +CactusGraph CactusGraph::load(const std::string& path) { + GraphFile::SerializedGraph sg = GraphFile::load_graph(path); + return from_serialized(sg); +} + +size_t CactusGraph::mmap_embeddings(const std::string& filename) { + const std::string resolved_filename = resolve_quantized_weight_file(filename); + auto mapped_file = std::make_unique(resolved_filename); + + const auto& shape = mapped_file->shape(); + if (shape.size() != 2) { + throw std::runtime_error("Memory-mapped embeddings must be 2D [vocab_size, embedding_dim]"); + } + + Precision precision = mapped_file->precision(); + + size_t node_id = input(shape, precision); + set_external_input(node_id, const_cast(mapped_file->data()), precision); + + auto& buffer = nodes_[node_index_map_.at(node_id)]->output_buffer; + if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { + buffer.group_size = mapped_file->group_size(); + buffer.num_groups = mapped_file->num_groups(); + + const char* scales_base = static_cast(mapped_file->scales_data()); + uint32_t bits = PrecisionTraits::cq_bits(precision); + uint32_t cb_size = 1u << bits; + uint32_t gs = static_cast(mapped_file->group_size()); + uint32_t K = gs * static_cast(mapped_file->num_groups()); + uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); + + size_t off = 0; + buffer.cq_codebook = reinterpret_cast(scales_base + off); + off += cb_size * sizeof(__fp16); + buffer.cq_input_scale = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_norms = reinterpret_cast(scales_base + off); + off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); + + if (mapped_file->is_orthogonal_rotation()) { + buffer.cq_rotation = reinterpret_cast(scales_base + off); + buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else { + buffer.cq_left_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_right_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_permutation = reinterpret_cast(scales_base + off); + buffer.cq_flags = 0; + } + if (mapped_file->is_interleaved_4row()) { + buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + } + } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) { + buffer.group_size = mapped_file->group_size(); + buffer.num_groups = mapped_file->num_groups(); + buffer.set_activation_scales(const_cast(mapped_file->scales_data()), shape.empty() ? 0 : shape[0]); + } + + size_t file_idx = mapped_files_.size(); + mapped_files_.push_back(std::move(mapped_file)); + node_to_mapped_file_[node_id] = file_idx; + weight_cache_[filename] = node_id; + return node_id; +} + +size_t CactusGraph::mmap_weights(const std::string& filename) { + const std::string resolved_filename = resolve_quantized_weight_file(filename); + auto it = weight_cache_.find(resolved_filename); + if (it != weight_cache_.end()) { + return it->second; + } + + auto mapped_file = std::make_unique(resolved_filename); + + const auto& shape = mapped_file->shape(); + Precision precision = mapped_file->precision(); + + size_t node_id = input(shape, precision); + set_external_input(node_id, const_cast(mapped_file->data()), precision); + + if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { + auto& buffer = nodes_[node_index_map_.at(node_id)]->output_buffer; + buffer.group_size = mapped_file->group_size(); + buffer.num_groups = mapped_file->num_groups(); + + + const char* scales_base = static_cast(mapped_file->scales_data()); + uint32_t bits = PrecisionTraits::cq_bits(precision); + uint32_t cb_size = 1u << bits; + uint32_t gs = static_cast(mapped_file->group_size()); + uint32_t K = gs * static_cast(mapped_file->num_groups()); + uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); + + size_t off = 0; + buffer.cq_codebook = reinterpret_cast(scales_base + off); + off += cb_size * sizeof(__fp16); + buffer.cq_input_scale = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_norms = reinterpret_cast(scales_base + off); + off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); + + if (mapped_file->is_orthogonal_rotation()) { + buffer.cq_rotation = reinterpret_cast(scales_base + off); + buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else { + buffer.cq_left_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_right_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_permutation = reinterpret_cast(scales_base + off); + buffer.cq_flags = 0; + } + if (mapped_file->is_interleaved_4row()) { + buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + } + } + + size_t file_idx = mapped_files_.size(); + mapped_files_.push_back(std::move(mapped_file)); + node_to_mapped_file_[node_id] = file_idx; + weight_cache_[resolved_filename] = node_id; + return node_id; +} + +void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) { + const std::string resolved_filename = resolve_quantized_weight_file(filename); + auto node_it = node_index_map_.find(node_id); + if (node_it == node_index_map_.end()) { + throw std::out_of_range("Unknown input node id: " + std::to_string(node_id)); + } + + auto& node = *nodes_[node_it->second]; + if (node.op_type != OpType::INPUT) { + throw std::invalid_argument("Can only bind mmap weights to input nodes"); + } + + auto mapped_file = std::make_unique(resolved_filename); + const auto& shape = mapped_file->shape(); + Precision precision = mapped_file->precision(); + auto& buffer = node.output_buffer; + if (buffer.shape != shape) { + throw std::runtime_error("mmap weight shape mismatch for node " + std::to_string(node_id)); + } + if (buffer.precision != precision) { + throw std::runtime_error("mmap weight precision mismatch for node " + std::to_string(node_id)); + } + + set_external_input(node_id, const_cast(mapped_file->data()), precision); + buffer.group_size = 0; + buffer.num_groups = 0; + buffer.activation_scales_data = nullptr; + buffer.owned_activation_scales.reset(); + buffer.num_rows_for_activation_scales = 0; + buffer.cq_codebook = nullptr; + buffer.cq_input_scale = nullptr; + buffer.cq_input_scale_recip = nullptr; + buffer.cq_norms = nullptr; + buffer.cq_left_signs = nullptr; + buffer.cq_right_signs = nullptr; + buffer.cq_permutation = nullptr; + buffer.cq_rotation = nullptr; + buffer.cq_flags = 0; + + if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { + buffer.group_size = mapped_file->group_size(); + buffer.num_groups = mapped_file->num_groups(); + + const char* scales_base = static_cast(mapped_file->scales_data()); + uint32_t bits = PrecisionTraits::cq_bits(precision); + uint32_t cb_size = 1u << bits; + uint32_t gs = static_cast(mapped_file->group_size()); + uint32_t K = gs * static_cast(mapped_file->num_groups()); + uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); + + size_t off = 0; + buffer.cq_codebook = reinterpret_cast(scales_base + off); + off += cb_size * sizeof(__fp16); + buffer.cq_input_scale = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_norms = reinterpret_cast(scales_base + off); + off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); + + if (mapped_file->is_orthogonal_rotation()) { + buffer.cq_rotation = reinterpret_cast(scales_base + off); + buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else { + buffer.cq_left_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_right_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_permutation = reinterpret_cast(scales_base + off); + buffer.cq_flags = 0; + } + if (mapped_file->is_interleaved_4row()) { + buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + } + } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) { + buffer.group_size = mapped_file->group_size(); + buffer.num_groups = mapped_file->num_groups(); + buffer.set_activation_scales(const_cast(mapped_file->scales_data()), shape.empty() ? 0 : shape[0]); + } + + size_t file_idx = mapped_files_.size(); + mapped_files_.push_back(std::move(mapped_file)); + node_to_mapped_file_[node_id] = file_idx; + weight_cache_[resolved_filename] = node_id; +} + +void CactusGraph::release_weight_pages(size_t node_id) { + auto it = node_to_mapped_file_.find(node_id); + if (it != node_to_mapped_file_.end() && it->second < mapped_files_.size()) { + mapped_files_[it->second]->release_pages(); + } +} + +void CactusGraph::prefetch_weight_pages(size_t node_id) { + auto it = node_to_mapped_file_.find(node_id); + if (it != node_to_mapped_file_.end() && it->second < mapped_files_.size()) { + mapped_files_[it->second]->prefetch_pages(); + } +} + +void CactusGraph::release_all_weight_pages() { + for (auto& mf : mapped_files_) { + if (mf) mf->release_pages(); + } +} + +size_t CactusGraph::embedding(const std::string& filename, size_t indices, ComputeBackend backend) { + auto mapped_file = std::make_unique(filename); + + const auto& shape = mapped_file->shape(); + if (shape.size() != 2) { + throw std::runtime_error("Embedding file must contain 2D tensor [vocab_size, hidden_dim]"); + } + + Precision precision = mapped_file->precision(); + size_t embeddings_node = input(shape, precision); + set_external_input(embeddings_node, const_cast(mapped_file->data()), precision); + + mapped_files_.push_back(std::move(mapped_file)); + + const auto& idx_shape = get_output_buffer(indices).shape; + std::vector output_shape = idx_shape; + output_shape.push_back(shape[1]); + + OpParams params; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::EMBEDDING, {embeddings_node, indices}, output_shape, params), backend); +} + + +namespace GraphFile { + +void save_graph(const CactusGraph& graph, + const std::string& filename) { + + std::ofstream out(filename, std::ios::binary); + if (!out) { + throw std::runtime_error("Cannot open file for writing: " + filename); + } + + SerializedGraph sg; + sg.header.magic = CACTUS_GRAPH_MAGIC; + bool any_dynamic = false; + for (const auto& node : graph.nodes_) { + if (!node->output_buffer.dynamic_dims.empty()) { any_dynamic = true; break; } + } + sg.header.version = any_dynamic ? CACTUS_GRAPH_VERSION_DYNAMIC_SHAPES + : CACTUS_GRAPH_VERSION_EMBEDDED_INPUTS; + sg.header.node_count = static_cast(graph.nodes_.size()); + sg.header.flags = 0; + + sg.nodes.reserve(graph.nodes_.size()); + + for (uint32_t i = 0; i < graph.nodes_.size(); ++i) { + const auto& node = graph.nodes_[i]; + + NodeEntry entry; + entry.index = static_cast(node->id); + entry.op_type = node->op_type; + entry.output_shape = node->output_buffer.shape; + entry.precision = node->output_buffer.precision; + entry.params = node->params; + entry.dynamic_mask = node->output_buffer.dynamic_dims; + + entry.inputs.reserve(node->input_ids.size()); + for (size_t input_id : node->input_ids) { + entry.inputs.push_back(static_cast(input_id)); + } + + if (node->op_type == OpType::INPUT && graph.is_embedded_input(node->id)) { + if (node->output_buffer.external_data != nullptr) { + throw std::runtime_error( + "Graph save failed: embedded input node " + std::to_string(node->id) + + " is backed by external data" + ); + } + const void* data = node->output_buffer.get_data(); + if (data == nullptr && node->output_buffer.byte_size > 0) { + throw std::runtime_error( + "Graph save failed: embedded input node " + std::to_string(node->id) + + " has no materialized data" + ); + } + entry.has_embedded_data = true; + if (node->output_buffer.byte_size > 0) { + const auto* bytes = static_cast(data); + entry.embedded_data.assign(bytes, bytes + node->output_buffer.byte_size); + } + } + + if (node->op_type == OpType::INPUT && !entry.has_embedded_data) { + sg.graph_inputs.push_back(entry.index); + } + + sg.nodes.push_back(std::move(entry)); + } + + sg.graph_outputs = compute_leaf_outputs(sg); + + write_serialized_graph(out, sg); + + if (!out) { + throw std::runtime_error("Error writing graph data to file: " + filename); + } +} + +SerializedGraph load_graph(const std::string& filename) { + std::ifstream in(filename, std::ios::binary); + if (!in) { + throw std::runtime_error("Cannot open file for reading: " + filename); + } + SerializedGraph sg; + sg.header = read_graph_header(in); + sg.graph_inputs = read_u32_vector(in); + sg.graph_outputs = read_u32_vector(in); + sg.nodes.reserve(sg.header.node_count); + for (uint32_t i = 0; i < sg.header.node_count; ++i) { + sg.nodes.push_back(read_node_entry(in, sg.header.version)); + } + + if (sg.nodes.size() != sg.header.node_count) { + throw std::runtime_error("Graph file corrupted: node count mismatch"); + } + + bool dense_legacy_indices = true; + std::unordered_set node_ids; + node_ids.reserve(sg.nodes.size()); + for (uint32_t i = 0; i < sg.nodes.size(); ++i) { + const auto& node = sg.nodes[i]; + if (node.index != i) { + dense_legacy_indices = false; + } + node_ids.insert(node.index); + } + + for (const auto& node : sg.nodes) { + for (uint32_t input_idx : node.inputs) { + if (dense_legacy_indices) { + if (input_idx >= sg.nodes.size()) { + throw std::runtime_error("Graph file corrupted: input index out of range"); + } + } else if (node_ids.find(input_idx) == node_ids.end()) { + throw std::runtime_error("Graph file corrupted: input node id out of range"); + } + } + } + + for (uint32_t input_idx : sg.graph_inputs) { + if (dense_legacy_indices) { + if (input_idx >= sg.nodes.size()) { + throw std::runtime_error("Graph file corrupted: graph input index out of range"); + } + } else if (node_ids.find(input_idx) == node_ids.end()) { + throw std::runtime_error("Graph file corrupted: graph input node id out of range"); + } + } + + for (uint32_t output_idx : sg.graph_outputs) { + if (dense_legacy_indices) { + if (output_idx >= sg.nodes.size()) { + throw std::runtime_error("Graph file corrupted: graph output index out of range"); + } + } else if (node_ids.find(output_idx) == node_ids.end()) { + throw std::runtime_error("Graph file corrupted: graph output node id out of range"); + } + } + + return sg; +} + +void save_node(CactusGraph& graph, size_t node_id, const std::string& filename) { + graph.execute(); + void* data = graph.get_output(node_id); + + const auto& buffer = graph.get_output_buffer(node_id); + const auto& shape = buffer.shape; + Precision precision = buffer.precision; + + std::ofstream file(filename, std::ios::binary); + if (!file) { + throw std::runtime_error("Cannot open file for writing: " + filename); + } + + size_t total_elements = 1; + for (size_t dim : shape) { + total_elements *= dim; + } + + size_t byte_size = PrecisionTraits::packed_size_of(precision, total_elements); + + size_t N = shape.size() >= 1 ? shape[0] : 1; + + uint32_t ndim = static_cast(shape.size()); + uint32_t flags = 0; + uint32_t alignment = 32; + + uint32_t magic = CACTUS_MAGIC; + file.write(reinterpret_cast(&magic), sizeof(magic)); + file.write(reinterpret_cast(&flags), sizeof(flags)); + file.write(reinterpret_cast(&alignment), sizeof(alignment)); + file.write(reinterpret_cast(&ndim), sizeof(ndim)); + + for (uint32_t i = 0; i < 4; i++) { + uint64_t dim_val = (i < shape.size()) ? static_cast(shape[i]) : 0; + file.write(reinterpret_cast(&dim_val), sizeof(dim_val)); + } + + uint32_t prec_val = static_cast(precision); + file.write(reinterpret_cast(&prec_val), sizeof(prec_val)); + + uint64_t data_bytes = static_cast(byte_size); + uint64_t zero64 = 0; + uint32_t zero32 = 0; + file.write(reinterpret_cast(&data_bytes), sizeof(data_bytes)); + file.write(reinterpret_cast(&zero64), sizeof(zero64)); // scales_bytes = 0 + file.write(reinterpret_cast(&zero32), sizeof(zero32)); // group_size = 0 + file.write(reinterpret_cast(&zero32), sizeof(zero32)); // num_groups = 0 + uint64_t original_N_val = static_cast(N); + file.write(reinterpret_cast(&original_N_val), sizeof(original_N_val)); + + size_t header_end = HEADER_SIZE; + size_t aligned_header = align_offset(header_end, alignment); + size_t header_padding = aligned_header - header_end; + for (size_t i = 0; i < header_padding; i++) { + char zero = 0; + file.write(&zero, 1); + } + + file.write(static_cast(data), byte_size); + + if (!file) { + throw std::runtime_error("Error writing node data to file: " + filename); + } +} + +// MappedFile implementation + +MappedFile::MappedFile(const std::string& filename) + : fd_(-1), mapped_data_(nullptr), file_size_(0), data_offset_(0) { + fd_ = open(filename.c_str(), O_RDONLY); + if (fd_ == -1) { + throw std::runtime_error("Cannot open file for mapping: " + filename); + } + + struct stat st; + if (fstat(fd_, &st) == -1) { + close(fd_); + throw std::runtime_error("Cannot get file size: " + filename); + } + file_size_ = static_cast(st.st_size); + + mapped_data_ = mmap(nullptr, file_size_, PROT_READ, MAP_SHARED, fd_, 0); + if (mapped_data_ == MAP_FAILED) { + close(fd_); + throw std::runtime_error("Cannot map file: " + filename); + } + + close(fd_); + fd_ = -1; + + parse_header(); + apply_madvise_hints(); +} + +MappedFile::~MappedFile() { + if (mapped_data_ != nullptr && mapped_data_ != MAP_FAILED) { + madvise(mapped_data_, file_size_, MADV_DONTNEED); + munmap(mapped_data_, file_size_); + mapped_data_ = nullptr; + } + if (fd_ != -1) { + close(fd_); + fd_ = -1; + } +} + +MappedFile::MappedFile(MappedFile&& other) noexcept + : fd_(other.fd_), mapped_data_(other.mapped_data_), file_size_(other.file_size_), + data_offset_(other.data_offset_), shape_(std::move(other.shape_)), + precision_(other.precision_), byte_size_(other.byte_size_), + group_size_(other.group_size_), num_groups_(other.num_groups_), + scales_offset_(other.scales_offset_), scales_bytes_(other.scales_bytes_), + alignment_(other.alignment_), + is_orthogonal_rotation_(other.is_orthogonal_rotation_), + is_interleaved_4row_(other.is_interleaved_4row_), + original_N_(other.original_N_) { + other.fd_ = -1; + other.mapped_data_ = nullptr; + other.file_size_ = 0; + other.is_orthogonal_rotation_ = false; + other.is_interleaved_4row_ = false; + other.original_N_ = 0; +} + +MappedFile& MappedFile::operator=(MappedFile&& other) noexcept { + if (this != &other) { + if (mapped_data_ != nullptr && mapped_data_ != MAP_FAILED) { + munmap(mapped_data_, file_size_); + } + if (fd_ != -1) { + close(fd_); + } + + fd_ = other.fd_; + mapped_data_ = other.mapped_data_; + file_size_ = other.file_size_; + data_offset_ = other.data_offset_; + shape_ = std::move(other.shape_); + precision_ = other.precision_; + byte_size_ = other.byte_size_; + group_size_ = other.group_size_; + num_groups_ = other.num_groups_; + scales_offset_ = other.scales_offset_; + scales_bytes_ = other.scales_bytes_; + alignment_ = other.alignment_; + is_orthogonal_rotation_ = other.is_orthogonal_rotation_; + is_interleaved_4row_ = other.is_interleaved_4row_; + original_N_ = other.original_N_; + other.fd_ = -1; + other.mapped_data_ = nullptr; + other.file_size_ = 0; + other.is_orthogonal_rotation_ = false; + other.is_interleaved_4row_ = false; + other.original_N_ = 0; + } + return *this; +} + +const std::vector& MappedFile::shape() const { + return shape_; +} + +Precision MappedFile::precision() const { + return precision_; +} + +size_t MappedFile::byte_size() const { + return byte_size_; +} + +const void* MappedFile::scales_data() const { + return static_cast(mapped_data_) + scales_offset_; +} + +void* MappedFile::data() { + return static_cast(mapped_data_) + data_offset_; +} + +const void* MappedFile::data() const { + return static_cast(mapped_data_) + data_offset_; +} + +template +const T* MappedFile::typed_data() const { + return static_cast(data()); +} + +void MappedFile::parse_header() { + if (file_size_ < HEADER_SIZE) { + throw std::runtime_error("File too small: insufficient data for header"); + } + + const char* ptr = static_cast(mapped_data_); + size_t offset = 0; + + uint32_t magic = *reinterpret_cast(ptr + offset); + offset += sizeof(uint32_t); + if (magic != CACTUS_MAGIC) { + throw std::runtime_error("Invalid tensor file: missing CACT magic number"); + } + + uint32_t flags = *reinterpret_cast(ptr + offset); + offset += sizeof(uint32_t); + is_orthogonal_rotation_ = (flags & FLAG_ORTHOGONAL_ROTATION) != 0; + is_interleaved_4row_ = (flags & FLAG_INTERLEAVED_4ROW) != 0; + + alignment_ = *reinterpret_cast(ptr + offset); + offset += sizeof(uint32_t); + if (alignment_ == 0) alignment_ = 1; + + uint32_t ndim = *reinterpret_cast(ptr + offset); + offset += sizeof(uint32_t); + + std::vector dims; + dims.reserve(8); + for (uint32_t i = 0; i < 4; i++) { + uint64_t dim_val = *reinterpret_cast(ptr + offset); + offset += sizeof(uint64_t); + dims.push_back(dim_val); + } + + uint32_t prec_val = *reinterpret_cast(ptr + offset); + precision_ = static_cast(prec_val); + offset += sizeof(uint32_t); + + byte_size_ = *reinterpret_cast(ptr + offset); + offset += sizeof(uint64_t); + + scales_bytes_ = *reinterpret_cast(ptr + offset); + offset += sizeof(uint64_t); + + group_size_ = *reinterpret_cast(ptr + offset); + offset += sizeof(uint32_t); + + num_groups_ = *reinterpret_cast(ptr + offset); + offset += sizeof(uint32_t); + + original_N_ = *reinterpret_cast(ptr + offset); + offset += sizeof(uint64_t); + + size_t header_size = HEADER_SIZE; + if ((flags & FLAG_EXTENDED_SHAPE) != 0) { + header_size += 4 * sizeof(uint64_t); + if (file_size_ < header_size) { + throw std::runtime_error("File too small: insufficient data for extended shape header"); + } + for (uint32_t i = 0; i < 4; i++) { + uint64_t dim_val = *reinterpret_cast(ptr + offset); + offset += sizeof(uint64_t); + dims.push_back(dim_val); + } + } + + shape_.clear(); + if (ndim > dims.size()) { + throw std::runtime_error("Invalid tensor file: ndim exceeds encoded shape rank"); + } + for (uint32_t i = 0; i < ndim; i++) { + if (dims[i] > 0) { + shape_.push_back(static_cast(dims[i])); + } + } + + size_t aligned_header = align_offset(header_size, alignment_); + + if (scales_bytes_ > 0) { + scales_offset_ = aligned_header; + size_t scales_end = scales_offset_ + scales_bytes_; + data_offset_ = align_offset(scales_end, alignment_); + } else { + scales_offset_ = 0; + data_offset_ = aligned_header; + } + + if (data_offset_ + byte_size_ > file_size_) { + throw std::runtime_error("File corrupted: data extends beyond file size"); + } + +} + +void MappedFile::apply_madvise_hints() { + if (scales_bytes_ > 0 && scales_offset_ > 0) { + madvise(static_cast(mapped_data_) + scales_offset_, scales_bytes_, MADV_WILLNEED); + } + + madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_SEQUENTIAL); +} + +void MappedFile::release_pages() { + if (mapped_data_ == nullptr || mapped_data_ == MAP_FAILED) return; + + if (scales_bytes_ > 0 && scales_offset_ > 0) { + madvise(static_cast(mapped_data_) + scales_offset_, scales_bytes_, MADV_DONTNEED); + } + madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_DONTNEED); +} + +void MappedFile::prefetch_pages() { + if (mapped_data_ == nullptr || mapped_data_ == MAP_FAILED) return; + + if (scales_bytes_ > 0 && scales_offset_ > 0) { + madvise(static_cast(mapped_data_) + scales_offset_, scales_bytes_, MADV_WILLNEED); + } + madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_WILLNEED); +} + +template const int8_t* MappedFile::typed_data() const; +template const float* MappedFile::typed_data() const; +template const uint16_t* MappedFile::typed_data() const; +template const uint8_t* MappedFile::typed_data() const; + +} // namespace GraphFile diff --git a/cactus-graph/src/last_error.cpp b/cactus-graph/src/last_error.cpp new file mode 100644 index 000000000..3175cec97 --- /dev/null +++ b/cactus-graph/src/last_error.cpp @@ -0,0 +1,8 @@ +#include "../cactus_graph.h" +#include + +thread_local std::string last_error_message; + +extern "C" const char* cactus_get_last_error() { + return last_error_message.c_str(); +} diff --git a/cactus-graph/src/metal_plan.cpp b/cactus-graph/src/metal_plan.cpp new file mode 100644 index 000000000..e0fb923ae --- /dev/null +++ b/cactus-graph/src/metal_plan.cpp @@ -0,0 +1,1825 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include +#include +#include + +bool cactus_kv_cache_grow(BufferDesc&, size_t, size_t); + +namespace { + +bool is_alias_op(OpType op) { + return op == OpType::VIEW || op == OpType::RESHAPE || op == OpType::FLATTEN; +} + +bool is_noop_transpose(const GraphNode& nd) { + if (nd.op_type != OpType::TRANSPOSE) return false; + size_t real = 0; + for (size_t d : nd.output_buffer.shape) if (d != 1) ++real; + return real <= 1; +} + +bool is_same_cast(const GraphNode& nd, const std::vector>& nodes, + const std::unordered_map& map) { + if (nd.op_type != OpType::PRECISION_CAST || nd.input_ids.empty()) return false; + auto it = map.find(nd.input_ids[0]); + return it != map.end() && nodes[it->second]->output_buffer.precision == nd.output_buffer.precision; +} + +} + +struct MetalCluster { + int rule = 0; + size_t a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0; + size_t b0 = 0, b1 = 0, b2 = 0, b3 = 0, b4 = 0; + float f0 = 0.0f, f1 = 0.0f; + uint32_t u0 = 0, u1 = 0; + void* s0 = nullptr; + void* s1 = nullptr; + void* sc[3] = {nullptr, nullptr, nullptr}; +}; + +struct MetalFusePlan { + std::vector action; + std::vector clusters; + long fold_h = -1, fold_ple = -1, fold_pos = -1, fold_w = -1; + size_t fold_nl = 0, fold_pd = 0; + std::vector exec_list; + std::vector arena_off; + char* arena_base = nullptr; + std::vector> blobs; +}; + +struct EwChainStep { + int32_t kind; + int32_t code; + float p0; + float p1; +}; + +MetalFusePlan* cactus_metal_plan_build( + const std::vector>& nodes, + const std::unordered_map& map, + const std::unordered_set& pinned_ids, + const std::vector& retype, + const std::unordered_set* banned) { + const size_t n = nodes.size(); + auto plan = new MetalFusePlan(); + plan->action.assign(n, -1); + std::vector pinned(n, 0); + for (size_t i = 0; i < n; ++i) + if (pinned_ids.count(nodes[i]->id)) pinned[i] = 1; + std::vector cpu_only(n, 0); + for (size_t i = 0; i < n; ++i) + if (nodes[i]->params.backend != ComputeBackend::METAL) { cpu_only[i] = 1; pinned[i] = 1; } + + auto idxof = [&](size_t id) -> long { + auto it = map.find(id); + return it == map.end() ? -1 : (long)it->second; + }; + std::vector> cons(n); + for (size_t i = 0; i < n; ++i) + for (size_t id : nodes[i]->input_ids) { long j = idxof(id); if (j >= 0) cons[(size_t)j].push_back(i); } + + auto retyped = [&](size_t idx) { return idx < retype.size() && retype[idx] != 0; }; + auto passthrough = [&](const GraphNode& nd) { + return is_alias_op(nd.op_type) || is_same_cast(nd, nodes, map) || is_noop_transpose(nd); + }; + auto up = [&](long i) -> long { + while (i >= 0) { + const GraphNode& nd = *nodes[(size_t)i]; + if (passthrough(nd) && !nd.input_ids.empty()) { + i = idxof(nd.input_ids[0]); + continue; + } + return i; + } + return -1; + }; + auto up_in = [&](size_t i, size_t k) -> long { + if (k >= nodes[i]->input_ids.size()) return -1; + return up(idxof(nodes[i]->input_ids[k])); + }; + auto deep_f16_source = [&](long i) -> long { + long cur = i; + while (cur >= 0) { + const GraphNode& nd = *nodes[(size_t)cur]; + if ((passthrough(nd) || nd.op_type == OpType::PRECISION_CAST) && !nd.input_ids.empty()) { + cur = idxof(nd.input_ids[0]); + continue; + } + break; + } + if (cur >= 0 && nodes[(size_t)cur]->output_buffer.precision == Precision::FP16) return cur; + return i; + }; + auto collect_chain = [&](size_t from, long to, std::vector& out_chain) { + long i = (long)from; + while (i >= 0 && i != to) { + const GraphNode& nd = *nodes[(size_t)i]; + if (!passthrough(nd)) break; + out_chain.push_back((size_t)i); + i = nd.input_ids.empty() ? -1 : idxof(nd.input_ids[0]); + } + }; + auto sole_use = [&](long i, const std::vector& allowed) -> bool { + if (i < 0) return false; + for (size_t c : cons[(size_t)i]) { + bool ok = false; + for (size_t a : allowed) if (a == c) { ok = true; break; } + if (!ok) { + const GraphNode& cn = *nodes[c]; + if (is_alias_op(cn.op_type) || is_same_cast(cn, nodes, map) || is_noop_transpose(cn)) { + bool subok = true; + for (size_t cc : cons[c]) { + bool found = false; + for (size_t a : allowed) if (a == cc) { found = true; break; } + if (!found) { subok = false; break; } + } + if (!subok) return false; + } else return false; + } + } + return true; + }; + + auto match_rope = [&](long add_i, long* x_out, long* cos_out, long* sin_out, + std::vector& cover) -> bool { + if (add_i < 0 || nodes[(size_t)add_i]->op_type != OpType::ADD) return false; + long m0 = up_in((size_t)add_i, 0), m1 = up_in((size_t)add_i, 1); + if (m0 < 0 || m1 < 0) return false; + if (nodes[(size_t)m0]->op_type != OpType::MULTIPLY) std::swap(m0, m1); + if (m0 < 0 || m1 < 0) return false; + if (nodes[(size_t)m0]->op_type != OpType::MULTIPLY || + nodes[(size_t)m1]->op_type != OpType::MULTIPLY) return false; + long rot = up_in((size_t)m1, 0); + long sinb = up_in((size_t)m1, 1); + if (rot < 0 || nodes[(size_t)rot]->op_type != OpType::CAT) { std::swap(rot, sinb); } + if (rot < 0 || sinb < 0 || nodes[(size_t)rot]->op_type != OpType::CAT) { + std::swap(m0, m1); + rot = up_in((size_t)m1, 0); sinb = up_in((size_t)m1, 1); + if (rot < 0 || nodes[(size_t)rot]->op_type != OpType::CAT) { std::swap(rot, sinb); } + if (rot < 0 || sinb < 0 || nodes[(size_t)rot]->op_type != OpType::CAT) return false; + } + long x = up_in((size_t)m0, 0), cosb = up_in((size_t)m0, 1); + if (x < 0 || cosb < 0) return false; + if (nodes[(size_t)rot]->input_ids.size() != 2) return false; + long neg = up_in((size_t)rot, 0), lo = up_in((size_t)rot, 1); + if (neg < 0 || lo < 0) return false; + if (nodes[(size_t)neg]->op_type != OpType::SCALAR_MULTIPLY || + nodes[(size_t)neg]->params.scalar != -1.0f) return false; + long hi = up_in((size_t)neg, 0); + if (hi < 0 || nodes[(size_t)hi]->op_type != OpType::SLICE || + nodes[(size_t)lo]->op_type != OpType::SLICE) return false; + long xs1 = up_in((size_t)hi, 0), xs2 = up_in((size_t)lo, 0); + if (xs1 != x || xs2 != x) { + long xa = up_in((size_t)m0, 1); + if (xs1 == xa && xs2 == xa) { x = xa; cosb = up_in((size_t)m0, 0); } + else return false; + } + *x_out = x; *cos_out = cosb; *sin_out = sinb; + cover.push_back((size_t)add_i); cover.push_back((size_t)m0); cover.push_back((size_t)m1); + cover.push_back((size_t)rot); cover.push_back((size_t)neg); + cover.push_back((size_t)hi); cover.push_back((size_t)lo); + return true; + }; + + auto alias_base = [&](long i, long* off_out) -> long { + long off = 0; + while (i >= 0) { + const GraphNode& nd = *nodes[(size_t)i]; + if (passthrough(nd) && !nd.input_ids.empty()) { + i = idxof(nd.input_ids[0]); + continue; + } + if (nd.op_type == OpType::INDEX && nd.params.axis == 0 && !nd.input_ids.empty()) { + long blk = 1; + for (size_t d : nd.output_buffer.shape) blk *= (long)d; + off += (long)nd.params.index_value * blk; + i = idxof(nd.input_ids[0]); + continue; + } + if (nd.op_type == OpType::SLICE && !nd.input_ids.empty()) { + long j = idxof(nd.input_ids[0]); + if (j < 0) break; + const auto& ish = nodes[(size_t)j]->output_buffer.shape; + size_t ax = (size_t)nd.params.axis; + if (ax >= ish.size()) break; + size_t outer = 1; + for (size_t d = 0; d < ax; ++d) outer *= ish[d]; + if (outer != 1) break; + long inner = 1; + for (size_t d = ax + 1; d < ish.size(); ++d) inner *= (long)ish[d]; + off += (long)nd.params.slice_start * inner; + i = j; + continue; + } + break; + } + *off_out = off; + return i; + }; + + auto release_scratch = [&](MetalCluster& c) { + if (c.s0) { cactus_metal_free_shared(c.s0); c.s0 = nullptr; } + if (c.s1) { cactus_metal_free_shared(c.s1); c.s1 = nullptr; } + for (auto& sp : c.sc) if (sp) { cactus_metal_free_shared(sp); sp = nullptr; } + }; + auto add_cluster = [&](MetalCluster c, size_t anchor, const std::vector& cover) -> bool { + if (banned && banned->count(anchor)) { release_scratch(c); return false; } + if (cpu_only[anchor]) { release_scratch(c); return false; } + for (size_t v : cover) if (v != anchor && pinned[v]) { release_scratch(c); return false; } + int32_t cid = (int32_t)plan->clusters.size(); + plan->clusters.push_back(c); + for (size_t v : cover) if (v != anchor) plan->action[v] = -2; + plan->action[anchor] = cid; + return true; + }; + + struct AttnCand { + MetalCluster c; + size_t anchor; + std::vector cover; + size_t kcache; + long kapp, vapp; + }; + std::vector cands; + + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& nd = *nodes[i]; + + if (nd.op_type == OpType::ATTENTION_CACHED && nd.input_ids.size() >= 5) { + long qidx = idxof(nd.input_ids[0]); + if (qidx < 0) continue; + const BufferDesc& qb = nodes[(size_t)qidx]->output_buffer; + if (qb.shape.size() < 4 || qb.shape[0] != 1 || qb.shape[1] != 1) continue; + uint32_t nqh = (uint32_t)qb.shape[2], hd = (uint32_t)qb.shape[3]; + long qk = idxof(nd.input_ids[1]); + long vk = idxof(nd.input_ids[2]); + long kcache = idxof(nd.input_ids[3]); + long vcache = idxof(nd.input_ids[4]); + if (qk < 0 || vk < 0 || kcache < 0 || vcache < 0) continue; + if (nodes[(size_t)kcache]->params.num_kv_heads != 1) continue; + + std::vector cover; + long q_rope = up_in(i, 0); + long qx = -1, qcos = -1, qsin = -1; + if (!match_rope(q_rope, &qx, &qcos, &qsin, cover)) continue; + long qnorm = qx; + if (qnorm < 0 || nodes[(size_t)qnorm]->op_type != OpType::RMS_NORM) continue; + long qraw = up_in((size_t)qnorm, 0); + long qw = up_in((size_t)qnorm, 1); + if (qraw < 0 || qw < 0) continue; + cover.push_back((size_t)qnorm); + + long k_rope = up(qk); + long kx = -1, kcos = -1, ksin = -1; + std::vector kcover; + long knorm = -1, kraw = -1, kw = -1, vnorm = -1, vraw = -1, vw = -1; + long kapp = -1, vapp = -1; + bool has_new = false; + if (match_rope(k_rope, &kx, &kcos, &ksin, kcover)) { + knorm = kx; + if (knorm >= 0 && nodes[(size_t)knorm]->op_type == OpType::RMS_NORM) { + kraw = up_in((size_t)knorm, 0); + kw = up_in((size_t)knorm, 1); + long vn = up(vk); + if (kraw >= 0 && kw >= 0 && vn >= 0 && nodes[(size_t)vn]->op_type == OpType::RMS_NORM) { + vnorm = vn; + vraw = up_in((size_t)vn, 0); + vw = up_in((size_t)vn, 1); + for (size_t c = 0; c < n; ++c) { + if (nodes[c]->op_type != OpType::KV_CACHE_APPEND || nodes[c]->input_ids.size() < 2) continue; + long src = up(idxof(nodes[c]->input_ids[0])); + long cache = idxof(nodes[c]->input_ids[1]); + if (src == k_rope && cache == kcache) kapp = (long)c; + if (src == (long)vnorm && cache == vcache) vapp = (long)c; + } + if (kapp >= 0 && vapp >= 0 && vraw >= 0 && vw >= 0) has_new = true; + } + } + } + bool shared = false; + if (has_new) { + for (auto& pc : cands) + if (pc.kapp == kapp || pc.kcache == (size_t)kcache) { shared = true; break; } + } + + AttnCand cand; + cand.anchor = i; + cand.kcache = (size_t)kcache; + cand.kapp = kapp; cand.vapp = vapp; + MetalCluster& c = cand.c; + c.rule = 1; + c.b0 = (size_t)qcos; c.b1 = (size_t)qsin; + c.b2 = (size_t)kcache; c.b3 = (size_t)vcache; + c.b4 = i; + c.a0 = (size_t)qraw; c.a3 = (size_t)deep_f16_source(qw); + c.f0 = nodes[(size_t)qnorm]->params.epsilon; + c.f1 = nd.params.scale != 0.0f ? nd.params.scale : 1.0f / std::sqrt((float)hd); + c.u0 = nqh; c.u1 = hd; + cand.cover = cover; + collect_chain((size_t)idxof(nd.input_ids[0]), q_rope, cand.cover); + if (has_new && !shared) { + c.a1 = (size_t)kraw; c.a2 = (size_t)vraw; + c.a4 = (size_t)deep_f16_source(kw); c.a5 = (size_t)deep_f16_source(vw); + c.u1 |= 0x10000u; + cand.cover.insert(cand.cover.end(), kcover.begin(), kcover.end()); + cand.cover.push_back((size_t)knorm); + cand.cover.push_back((size_t)vnorm); + cand.cover.push_back((size_t)kapp); + cand.cover.push_back((size_t)vapp); + collect_chain((size_t)idxof(nd.input_ids[1]), k_rope, cand.cover); + collect_chain((size_t)idxof(nd.input_ids[2]), vnorm, cand.cover); + collect_chain((size_t)idxof(nodes[(size_t)kapp]->input_ids[0]), k_rope, cand.cover); + collect_chain((size_t)idxof(nodes[(size_t)vapp]->input_ids[0]), vnorm, cand.cover); + } else if (!shared) { + continue; + } + cands.push_back(std::move(cand)); + continue; + } + + if (nd.op_type == OpType::MATMUL && nd.input_ids.size() >= 2) { + long r = idxof(nd.input_ids[1]); + const BufferDesc* rb = r >= 0 ? &nodes[(size_t)r]->output_buffer : nullptr; + size_t M = 1; + for (size_t d = 0; d + 1 < nd.output_buffer.shape.size(); ++d) M *= nd.output_buffer.shape[d]; + if (rb && M == 1 && PrecisionTraits::is_cq(rb->precision) && rb->group_size > 0 + && PrecisionTraits::cq_bits(rb->precision) == 4 + && !(rb->cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL)) { + long mul = up_in(i, 0); + if (mul >= 0 && nodes[(size_t)mul]->op_type == OpType::MULTIPLY && sole_use(mul, {i})) { + long g0 = up_in((size_t)mul, 0), u0v = up_in((size_t)mul, 1); + for (int side = 0; side < 2; ++side) { + long gel = side == 0 ? g0 : u0v; + long other = side == 0 ? u0v : g0; + long sm = -1; + float sc_val = 1.0f; + if (gel >= 0 && nodes[(size_t)gel]->op_type == OpType::SCALAR_MULTIPLY + && sole_use(gel, {(size_t)mul})) { + sm = gel; + sc_val = nodes[(size_t)gel]->params.scalar; + gel = up_in((size_t)gel, 0); + } + if (gel >= 0 && other >= 0 && nodes[(size_t)gel]->op_type == OpType::GELU + && sole_use(gel, {sm >= 0 ? (size_t)sm : (size_t)mul})) { + long gate = up_in((size_t)gel, 0); + if (gate < 0) break; + if (sc_val == 1.0f && nodes[(size_t)gate]->op_type == OpType::MATMUL + && nodes[(size_t)gate]->input_ids.size() >= 2 + && sole_use(gate, {(size_t)gel})) { + long gr = idxof(nodes[(size_t)gate]->input_ids[1]); + const BufferDesc* grb = gr >= 0 ? &nodes[(size_t)gr]->output_buffer : nullptr; + size_t gM = 1; + for (size_t d = 0; d + 1 < nodes[(size_t)gate]->output_buffer.shape.size(); ++d) + gM *= nodes[(size_t)gate]->output_buffer.shape[d]; + if (grb && gM == 1 && PrecisionTraits::is_cq(grb->precision) + && grb->group_size == 128 && PrecisionTraits::cq_bits(grb->precision) == 4 + && !(grb->cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) + && cactus_metal_transform_gemv_fits((uint32_t)( + grb->shape.size() >= 2 ? grb->shape[1] : grb->shape[0]))) { + long gx = up_in((size_t)gate, 0); + if (gx >= 0) { + MetalCluster c7; + c7.rule = 7; + c7.a0 = (size_t)gx; + c7.a1 = (size_t)other; + c7.b0 = (size_t)gate; + c7.b4 = i; + c7.s0 = cactus_metal_alloc_shared(nodes[(size_t)gate]->output_buffer.total_size * 2); + std::vector cover; + cover.push_back((size_t)mul); + cover.push_back((size_t)gel); + cover.push_back((size_t)gate); + collect_chain((size_t)idxof(nd.input_ids[0]), mul, cover); + collect_chain((size_t)idxof(nodes[(size_t)gel]->input_ids[0]), gate, cover); + const GraphNode& rn = *nodes[(size_t)other]; + if (rn.op_type == OpType::INDEX && !rn.input_ids.empty() + && sole_use(other, {(size_t)mul})) { + long ib = idxof(rn.input_ids[0]); + if (ib >= 0) { + size_t rowlen = rn.output_buffer.total_size; + c7.a1 = (size_t)ib; + c7.u1 = (uint32_t)(rn.params.index_value * rowlen * 2); + c7.u0 = 1u; + cover.push_back((size_t)other); + } + } + add_cluster(c7, i, cover); + break; + } + } + } + MetalCluster c; + c.rule = 4; + c.a0 = (size_t)gate; + c.a1 = (size_t)other; + c.b4 = i; + c.f0 = sc_val; + const BufferDesc& gb = nodes[(size_t)gate]->output_buffer; + c.s0 = cactus_metal_alloc_shared(gb.total_size * 2); + std::vector cover; + cover.push_back((size_t)mul); + cover.push_back((size_t)gel); + if (sm >= 0) cover.push_back((size_t)sm); + collect_chain((size_t)idxof(nd.input_ids[0]), mul, cover); + add_cluster(c, i, cover); + break; + } + } + if (plan->action[i] >= 0) continue; + } + } + } + + if (nd.op_type == OpType::SCALAR_MULTIPLY && !nd.output_buffer.shape.empty() + && nd.output_buffer.shape.back() >= 32768 + && nd.output_buffer.total_size == nd.output_buffer.shape.back()) { + long th = up_in(i, 0); + if (th >= 0 && nodes[(size_t)th]->op_type == OpType::TANH) { + long dv = up_in((size_t)th, 0); + if (dv >= 0 && nodes[(size_t)dv]->op_type == OpType::SCALAR_DIVIDE + && nodes[(size_t)dv]->params.scalar == nd.params.scalar) { + long mm = up_in((size_t)dv, 0); + if (mm >= 0 && nodes[(size_t)mm]->op_type == OpType::MATMUL + && nodes[(size_t)mm]->input_ids.size() >= 2) { + long wnode = idxof(nodes[(size_t)mm]->input_ids[1]); + long xn = up_in((size_t)mm, 0); + if (wnode >= 0 && xn >= 0) { + const BufferDesc& wb = nodes[(size_t)wnode]->output_buffer; + if (PrecisionTraits::is_cq(wb.precision) && wb.group_size > 0 + && (wb.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) + && sole_use(mm, {(size_t)dv}) && sole_use(dv, {(size_t)th}) + && sole_use(th, {i})) { + MetalCluster c; + c.rule = 5; + c.a0 = (size_t)xn; c.a1 = (size_t)wnode; c.b4 = i; + c.f0 = nd.params.scalar; + c.s0 = cactus_metal_alloc_shared(wb.shape.empty() ? 0 : wb.to_cq_matrix().K * 2); + std::vector cover; + cover.push_back((size_t)mm); + cover.push_back((size_t)dv); + cover.push_back((size_t)th); + collect_chain((size_t)idxof(nodes[(size_t)dv]->input_ids[0]), mm, cover); + collect_chain((size_t)idxof(nd.input_ids[0]), th, cover); + add_cluster(c, i, cover); + continue; + } + } + } + } + } + } + + if (nd.op_type == OpType::CAT && nd.input_ids.size() >= 4 && !nd.output_buffer.shape.empty()) { + size_t cnt = nd.input_ids.size(); + bool blocky = true; + { + size_t ax = (size_t)nd.params.axis; + const auto& osh = nd.output_buffer.shape; + if (retyped(i)) blocky = false; + if (ax >= osh.size()) blocky = false; + for (size_t d = 0; blocky && d < ax; ++d) if (osh[d] != 1) blocky = false; + long fi = idxof(nd.input_ids[0]); + if (blocky && fi >= 0) { + const auto& ish = nodes[(size_t)fi]->output_buffer.shape; + size_t ip = 1, op2 = 1; + for (size_t d : ish) ip *= d; + for (size_t d : osh) op2 *= d; + if (ip * cnt != op2) blocky = false; + } else blocky = false; + } + if (blocky) { + long m0 = -1; + { + long t = up_in(i, 0); + while (t >= 0 && nodes[(size_t)t]->op_type == OpType::PRECISION_CAST + && !nodes[(size_t)t]->input_ids.empty()) + t = up(idxof(nodes[(size_t)t]->input_ids[0])); + m0 = t; + } + if (m0 >= 0 && nodes[(size_t)m0]->op_type == OpType::MATMUL + && nodes[(size_t)m0]->output_buffer.shape.size() == 2 + && nodes[(size_t)m0]->output_buffer.precision == Precision::FP16) { + size_t M = nodes[(size_t)m0]->output_buffer.shape[0]; + size_t N = nodes[(size_t)m0]->output_buffer.shape[1]; + long ka = idxof(nodes[(size_t)m0]->input_ids[0]); + size_t K = ka >= 0 ? nodes[(size_t)ka]->output_buffer.shape.back() : 0; + bool f32out = nd.output_buffer.precision == Precision::FP32; + long parentA = -1, parentB = -1, offA0 = 0, offB0 = 0; + bool ok = K > 0; + bool f32a = false; + std::vector cover; + for (size_t j = 0; j < cnt && ok; ++j) { + long t = idxof(nd.input_ids[j]); + std::vector chain; + while (t >= 0 && (passthrough(*nodes[(size_t)t]) + || (nodes[(size_t)t]->op_type == OpType::PRECISION_CAST))) { + if (cons[(size_t)t].size() != 1) { t = -1; break; } + chain.push_back((size_t)t); + t = idxof(nodes[(size_t)t]->input_ids[0]); + } + if (t < 0 || nodes[(size_t)t]->op_type != OpType::MATMUL + || plan->action[(size_t)t] != -1 + || cons[(size_t)t].size() != 1 + || nodes[(size_t)t]->output_buffer.shape.size() != 2 + || nodes[(size_t)t]->output_buffer.shape[0] != M + || nodes[(size_t)t]->output_buffer.shape[1] != N + || nodes[(size_t)t]->params.pretransposed_rhs) { ok = false; break; } + long offA, offB; + long lin = idxof(nodes[(size_t)t]->input_ids[0]); + bool jf32 = false; + long lcast = -1; + if (lin >= 0 && nodes[(size_t)lin]->op_type == OpType::PRECISION_CAST + && nodes[(size_t)lin]->output_buffer.precision == Precision::FP16 + && cons[(size_t)lin].size() == 1 + && !nodes[(size_t)lin]->input_ids.empty()) { + long ci = idxof(nodes[(size_t)lin]->input_ids[0]); + if (ci >= 0 && nodes[(size_t)ci]->output_buffer.precision == Precision::FP32) { + jf32 = true; lcast = lin; lin = ci; + } + } + long pa = alias_base(lin, &offA); + long pb = alias_base(idxof(nodes[(size_t)t]->input_ids[1]), &offB); + Precision wantA = jf32 ? Precision::FP32 : Precision::FP16; + if (pa < 0 || pb < 0 + || nodes[(size_t)pa]->output_buffer.precision != wantA + || (jf32 && retyped((size_t)pa)) + || nodes[(size_t)pb]->output_buffer.precision != Precision::FP16) { ok = false; break; } + if (j == 0) { parentA = pa; parentB = pb; offA0 = offA; offB0 = offB; f32a = jf32; } + else if (jf32 != f32a) { ok = false; break; } + else if (pa != parentA || pb != parentB + || offA != offA0 + (long)(j * M * K) + || offB != offB0 + (long)(j * K * N)) { ok = false; break; } + cover.push_back((size_t)t); + if (lcast >= 0) cover.push_back((size_t)lcast); + for (size_t cnode : chain) cover.push_back(cnode); + } + if (ok) { + MetalCluster c; + c.rule = 8; + c.a0 = (size_t)parentA; c.a1 = (size_t)parentB; + c.a2 = (size_t)offA0; c.a3 = (size_t)offB0; + c.u0 = (uint32_t)cnt; c.u1 = (uint32_t)M; + c.b0 = K; c.b1 = N; c.b2 = f32out ? 1 : 0; + c.b3 = f32a ? 1 : 0; + c.b4 = i; + add_cluster(c, i, cover); + continue; + } + } + } + bool samesz = true; + { + long fi = idxof(nd.input_ids[0]); + if (fi < 0) samesz = false; + else { + size_t ip = nodes[(size_t)fi]->output_buffer.total_size; + samesz = ip * cnt == nd.output_buffer.total_size; + } + } + long tp0 = samesz ? up_in(i, 0) : -1; + if (tp0 >= 0 && nodes[(size_t)tp0]->op_type == OpType::TRANSPOSE + && nodes[(size_t)tp0]->output_buffer.precision == Precision::FP16 + && nd.output_buffer.precision == Precision::FP16 + && nodes[(size_t)tp0]->output_buffer.shape.size() + 1 <= 8 + && !is_noop_transpose(*nodes[(size_t)tp0])) { + const auto& perm0 = nodes[(size_t)tp0]->params.permutation; + const auto& tsh0 = nodes[(size_t)tp0]->output_buffer.shape; + long parentP = -1, offP0 = 0, strideP = -1; + size_t blk = nodes[(size_t)tp0]->output_buffer.total_size; + bool ok = !perm0.empty(); + std::vector cover; + for (size_t j = 0; j < cnt && ok; ++j) { + long t = idxof(nd.input_ids[j]); + std::vector chain; + while (t >= 0 && passthrough(*nodes[(size_t)t])) { + if (cons[(size_t)t].size() != 1) { t = -1; break; } + chain.push_back((size_t)t); + t = idxof(nodes[(size_t)t]->input_ids[0]); + } + if (t < 0 || nodes[(size_t)t]->op_type != OpType::TRANSPOSE + || plan->action[(size_t)t] != -1 + || cons[(size_t)t].size() != 1 + || nodes[(size_t)t]->params.permutation != perm0 + || nodes[(size_t)t]->output_buffer.shape != tsh0 + || nodes[(size_t)t]->output_buffer.precision != Precision::FP16) { ok = false; break; } + long offP; + long pp = alias_base(idxof(nodes[(size_t)t]->input_ids[0]), &offP); + if (pp < 0 || nodes[(size_t)pp]->output_buffer.precision != Precision::FP16) { ok = false; break; } + long ii = idxof(nodes[(size_t)t]->input_ids[0]); + if (ii < 0 || nodes[(size_t)ii]->output_buffer.total_size != blk) { ok = false; break; } + if (j == 0) { parentP = pp; offP0 = offP; } + else if (j == 1 && pp == parentP && offP > offP0) { strideP = offP - offP0; } + else if (j == 1 || pp != parentP || offP != offP0 + (long)j * strideP) { ok = false; break; } + cover.push_back((size_t)t); + for (size_t cnode : chain) cover.push_back(cnode); + } + if (ok && (strideP < 0 || cnt < 2)) ok = false; + if (ok) { + size_t ax = (size_t)nd.params.axis; + long fi = idxof(nd.input_ids[0]); + const auto& fsh = nodes[(size_t)fi]->output_buffer.shape; + if (ax >= fsh.size()) ok = false; + else { + size_t pre = 1; + for (size_t d = 0; d < ax && d < fsh.size(); ++d) pre *= fsh[d]; + size_t fpre = 1; + const auto& csh = nd.output_buffer.shape; + for (size_t d = 0; d < ax && d < csh.size(); ++d) fpre *= csh[d]; + if (pre != fpre) ok = false; + } + } + if (ok) { + MetalCluster c; + c.rule = 10; + c.a0 = (size_t)parentP; c.a1 = (size_t)up_in(i, 0); + c.a2 = (size_t)offP0; + c.u0 = (uint32_t)cnt; + c.u1 = (uint32_t)nd.params.axis; + c.b0 = (size_t)strideP; + c.b4 = i; + add_cluster(c, i, cover); + continue; + } + } + + long c0 = up_in(i, 0); + if (!blocky) c0 = -1; + if (c0 >= 0 && nodes[(size_t)c0]->op_type == OpType::CONV1D + && nodes[(size_t)c0]->output_buffer.precision == Precision::FP16) { + const auto& osh = nodes[(size_t)c0]->output_buffer.shape; + size_t Lout = osh.empty() ? 0 : osh.back(); + size_t stride = nodes[(size_t)c0]->params.stride ? nodes[(size_t)c0]->params.stride : 1; + long parentX = -1, parentW = -1, offX0 = 0, offW0 = 0; + size_t L = 0, Kk = 0; + bool ok = Lout > 0; + std::vector cover; + for (size_t j = 0; j < cnt && ok; ++j) { + long t = up(idxof(nd.input_ids[j])); + if (t < 0 || nodes[(size_t)t]->op_type != OpType::CONV1D + || plan->action[(size_t)t] != -1 + || cons[(size_t)t].size() != 1 + || nodes[(size_t)t]->input_ids.size() != 2 + || nodes[(size_t)t]->output_buffer.shape.back() != Lout) { ok = false; break; } + long offX, offW; + long px = alias_base(idxof(nodes[(size_t)t]->input_ids[0]), &offX); + long pw = alias_base(idxof(nodes[(size_t)t]->input_ids[1]), &offW); + if (px < 0 || pw < 0 + || nodes[(size_t)px]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)pw]->output_buffer.precision != Precision::FP16) { ok = false; break; } + long xin = idxof(nodes[(size_t)t]->input_ids[0]); + long win = idxof(nodes[(size_t)t]->input_ids[1]); + size_t Lj = nodes[(size_t)xin]->output_buffer.shape.back(); + size_t Kj = nodes[(size_t)win]->output_buffer.shape.back(); + if (j == 0) { + parentX = px; parentW = pw; offX0 = offX; offW0 = offW; L = Lj; Kk = Kj; + if ((L - Kk) / stride + 1 != Lout) { ok = false; break; } + } else if (px != parentX || pw != parentW || Lj != L || Kj != Kk + || offX != offX0 + (long)(j * L) + || offW != offW0 + (long)(j * Kk)) { ok = false; break; } + cover.push_back((size_t)t); + } + if (ok) { + MetalCluster c; + c.rule = 9; + c.a0 = (size_t)parentX; c.a1 = (size_t)parentW; + c.a2 = (size_t)offX0; c.a3 = (size_t)offW0; + c.u0 = (uint32_t)cnt; c.u1 = (uint32_t)L; + c.b0 = Lout; c.b1 = Kk; c.b2 = stride; + c.b4 = i; + add_cluster(c, i, cover); + continue; + } + } + } + + if (nd.op_type == OpType::ADD_CLIPPED && nd.input_ids.size() == 2) { + long n0 = up_in(i, 0), n1 = up_in(i, 1); + auto rms2_side = [&](long nrm, long* src, long* w) -> bool { + if (nrm < 0 || nodes[(size_t)nrm]->op_type != OpType::RMS_NORM + || nodes[(size_t)nrm]->input_ids.size() < 2 + || plan->action[(size_t)nrm] != -1 || pinned[(size_t)nrm] || retyped((size_t)nrm) + || nodes[(size_t)nrm]->output_buffer.precision != Precision::FP16 + || !sole_use(nrm, {i})) return false; + long s0 = up_in((size_t)nrm, 0); + long w0 = up_in((size_t)nrm, 1); + if (w0 >= 0) w0 = deep_f16_source(w0); + if (s0 < 0 || w0 < 0 + || nodes[(size_t)s0]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)w0]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)s0]->output_buffer.shape.empty() + || nodes[(size_t)s0]->output_buffer.total_size != + nodes[(size_t)s0]->output_buffer.shape.back()) return false; + *src = s0; *w = w0; + return true; + }; + size_t D2 = nd.output_buffer.shape.empty() ? 0 : nd.output_buffer.shape.back(); + long srcA = -1, wA = -1, srcB = -1, wB = -1; + if (n0 != n1 && D2 > 0 && nd.output_buffer.total_size == D2 + && nd.output_buffer.precision == Precision::FP16 && !retyped(i) + && rms2_side(n0, &srcA, &wA) && rms2_side(n1, &srcB, &wB) + && nodes[(size_t)srcA]->output_buffer.total_size == D2 + && nodes[(size_t)srcB]->output_buffer.total_size == D2 + && nodes[(size_t)wA]->output_buffer.total_size == D2 + && nodes[(size_t)wB]->output_buffer.total_size == D2) { + MetalCluster c; + c.rule = 18; + c.a0 = (size_t)srcA; c.a1 = (size_t)wA; + c.a2 = (size_t)srcB; c.a3 = (size_t)wB; + c.u1 = (uint32_t)D2; + c.f0 = nodes[(size_t)n0]->params.epsilon; + c.f1 = nodes[(size_t)n1]->params.epsilon; + c.b4 = i; + std::vector cover{(size_t)n0, (size_t)n1}; + if (add_cluster(c, i, cover)) continue; + } + long norm = -1, res = -1; + if (n0 >= 0 && nodes[(size_t)n0]->op_type == OpType::RMS_NORM) { norm = n0; res = n1; } + else if (n1 >= 0 && nodes[(size_t)n1]->op_type == OpType::RMS_NORM) { norm = n1; res = n0; } + if (norm >= 0 && res >= 0 && sole_use(norm, {i})) { + long src = up_in((size_t)norm, 0); + long w = up_in((size_t)norm, 1); + auto f16 = [&](long j) { + return j >= 0 && nodes[(size_t)j]->output_buffer.precision == Precision::FP16; + }; + if (w >= 0) w = deep_f16_source(w); + bool one_row = src >= 0 && !nodes[(size_t)src]->output_buffer.shape.empty() + && nodes[(size_t)src]->output_buffer.total_size == + nodes[(size_t)src]->output_buffer.shape.back(); + if (src >= 0 && w >= 0 && one_row && f16((long)i) && f16(src) && f16(res) && f16(w) && f16(norm)) { + MetalCluster c; + c.rule = 3; + c.a0 = (size_t)src; c.a1 = (size_t)w; c.a2 = (size_t)res; + c.b4 = i; + c.f0 = nodes[(size_t)norm]->params.epsilon; + std::vector cover; + cover.push_back((size_t)norm); + size_t h_node = i; + if (cons[i].size() == 1) { + long mt = (long)cons[i][0]; + while (mt >= 0 && passthrough(*nodes[(size_t)mt]) && cons[(size_t)mt].size() == 1) + mt = (long)cons[(size_t)mt][0]; + if (mt >= 0 && nodes[(size_t)mt]->op_type == OpType::MULTIPLY + && nodes[(size_t)mt]->input_ids.size() == 2 + && nodes[(size_t)mt]->output_buffer.precision == Precision::FP16) { + long o0 = up_in((size_t)mt, 0), o1 = up_in((size_t)mt, 1); + long scn = (o0 == (long)i || (o0 >= 0 && up(o0) == (long)i)) ? o1 : o0; + long base = scn == o1 ? o0 : o1; + if (base >= 0 && scn >= 0 + && nodes[(size_t)scn]->op_type == OpType::INPUT + && nodes[(size_t)scn]->output_buffer.total_size == 1 + && nodes[(size_t)scn]->output_buffer.precision == Precision::FP16) { + c.b1 = (size_t)scn; + c.u1 = 1u; + collect_chain(cons[i][0], (long)i, cover); + cover.push_back((size_t)mt); + h_node = (size_t)mt; + c.b4 = h_node; + cover.push_back(i); + } + } + } + long next_norm = -1; + for (size_t cc : cons[h_node]) { + long t = (long)cc; + while (t >= 0 && passthrough(*nodes[(size_t)t])) { + if (cons[(size_t)t].size() != 1) { t = -1; break; } + t = (long)cons[(size_t)t][0]; + } + if (t >= 0 && nodes[(size_t)t]->op_type == OpType::RMS_NORM + && up_in((size_t)t, 0) == (long)h_node + && nodes[(size_t)t]->output_buffer.precision == Precision::FP16) { + next_norm = t; + break; + } + } + if (next_norm >= 0) { + long nw = deep_f16_source(up_in((size_t)next_norm, 1)); + size_t ndim = nodes[(size_t)next_norm]->output_buffer.shape.empty() ? 0 + : nodes[(size_t)next_norm]->output_buffer.shape.back(); + size_t adim = nodes[i]->output_buffer.shape.empty() ? 0 + : nodes[i]->output_buffer.shape.back(); + if (nw >= 0 && ndim == adim && f16(nw)) { + c.u0 = 1u; + c.a4 = (size_t)nw; + c.a5 = (size_t)next_norm; + c.s1 = cactus_metal_alloc_shared(nodes[(size_t)next_norm]->output_buffer.total_size * 2); + cover.push_back((size_t)next_norm); + long ch = idxof(nodes[(size_t)next_norm]->input_ids[0]); + collect_chain((size_t)ch, (long)h_node, cover); + } + } + add_cluster(c, h_node, cover); + continue; + } + } + } + } + { + std::unordered_map> groups; + std::vector order; + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + const GraphNode& mn = *nodes[i]; + if (mn.op_type != OpType::MATMUL || mn.input_ids.size() < 2) continue; + long r = idxof(mn.input_ids[1]); + if (r < 0) continue; + const BufferDesc& rb = nodes[(size_t)r]->output_buffer; + size_t M = 1; + for (size_t d = 0; d + 1 < mn.output_buffer.shape.size(); ++d) M *= mn.output_buffer.shape[d]; + if (M != 1 || !PrecisionTraits::is_cq(rb.precision) || rb.group_size != 128 + || PrecisionTraits::cq_bits(rb.precision) != 4 + || (rb.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL)) continue; + long src = up_in(i, 0); + if (src < 0) continue; + auto it = groups.find(src); + if (it == groups.end()) { groups[src] = {i}; order.push_back(src); } + else it->second.push_back(i); + } + for (long src : order) { + auto& mms = groups[src]; + if (mms.size() < 2) continue; + for (size_t base = 0; base + 1 < mms.size(); base += 3) { + size_t cnt = std::min(3, mms.size() - base); + if (cnt < 2) break; + size_t anchor_i = mms[base]; + MetalCluster c; + c.rule = 2; + c.a0 = (size_t)src; + c.b0 = mms[base]; + c.b1 = mms[base + 1]; + c.b2 = cnt > 2 ? mms[base + 2] : 0; + c.u0 = (uint32_t)cnt; + c.b4 = anchor_i; + size_t maxK = 0; + for (size_t mi = 0; mi < cnt; ++mi) { + const BufferDesc& rb = nodes[(size_t)idxof(nodes[mms[base + mi]]->input_ids[1])]->output_buffer; + CactusQuantMatrix W = rb.to_cq_matrix(); + if (W.K > maxK) maxK = W.K; + } + for (uint32_t bi = 0; bi < c.u0; ++bi) + if (!c.sc[bi]) c.sc[bi] = cactus_metal_alloc_shared(maxK * 2); + std::vector cover; + if (add_cluster(c, anchor_i, cover)) + for (size_t mi = 1; mi < cnt; ++mi) plan->action[mms[base + mi]] = -3; + } + } + } + + cands.erase(std::remove_if(cands.begin(), cands.end(), + [&](const AttnCand& cd) { return cpu_only[cd.anchor] != 0; }), cands.end()); + if (!cands.empty()) { + std::vector mark(n, 0); + for (auto& cd : cands) { + mark[cd.anchor] = 2; + for (size_t v : cd.cover) if (mark[v] == 0) mark[v] = 1; + } + for (size_t ri = n; ri-- > 0;) { + if (mark[ri]) continue; + const GraphNode& ndp = *nodes[ri]; + if (!(is_alias_op(ndp.op_type) || is_same_cast(ndp, nodes, map) || is_noop_transpose(ndp))) continue; + if (cons[ri].empty()) { mark[ri] = 3; continue; } + bool all_in = true; + for (size_t cc : cons[ri]) if (!mark[cc]) { all_in = false; break; } + if (all_in) mark[ri] = 3; + } + bool valid = true; + for (size_t v = 0; v < n && valid; ++v) { + if (mark[v] != 1 && mark[v] != 3) continue; + for (size_t cc : cons[v]) { + if (!mark[cc]) { + valid = false; + break; + } + } + } + for (size_t v = 0; v < n && valid; ++v) + if ((mark[v] == 1 || mark[v] == 3) && pinned[v]) valid = false; + if (valid) { + for (size_t v = 0; v < n; ++v) + if ((mark[v] == 1 || mark[v] == 3) && plan->action[v] == -1) plan->action[v] = -2; + for (auto& cd : cands) { + int32_t cid = (int32_t)plan->clusters.size(); + plan->clusters.push_back(cd.c); + plan->action[cd.anchor] = cid; + } + } + } + + for (auto& cl : plan->clusters) { + if (cl.rule != 5) continue; + const auto& sh0 = nodes[0]->output_buffer; + const auto& sh1 = n > 1 ? nodes[1]->output_buffer : nodes[0]->output_buffer; + const auto& sh2 = n > 2 ? nodes[2]->output_buffer : nodes[0]->output_buffer; + if (n > 2 && nodes[0]->op_type == OpType::INPUT && nodes[1]->op_type == OpType::INPUT + && nodes[2]->op_type == OpType::INPUT + && sh0.precision == Precision::FP16 && sh0.total_size > 0 + && sh1.precision == Precision::FP16 && sh1.shape.size() >= 2 + && sh2.total_size == 1) { + plan->fold_h = 0; + plan->fold_ple = 1; + plan->fold_pos = 2; + plan->fold_w = (long)cl.a1; + plan->fold_nl = sh1.shape[sh1.shape.size() - 2]; + plan->fold_pd = sh1.shape.back(); + } + break; + } + + { + auto ew_kind = [&](const GraphNode& nd, int* kind, int* code, float* p0, float* p1) -> bool { + *p0 = 0; *p1 = 0; + switch (nd.op_type) { + case OpType::GELU: *kind = 0; *code = 0; return true; + case OpType::TANH: *kind = 0; *code = 1; return true; + case OpType::SILU: *kind = 0; *code = 2; return true; + case OpType::RELU: *kind = 0; *code = 3; return true; + case OpType::GELU_ERF: *kind = 0; *code = 4; return true; + case OpType::SIGMOID: *kind = 0; *code = 5; return true; + case OpType::SCALAR_ADD: *kind = 1; *code = 0; *p0 = nd.params.scalar; return true; + case OpType::SCALAR_SUBTRACT: *kind = 1; *code = 1; *p0 = nd.params.scalar; return true; + case OpType::SCALAR_MULTIPLY: *kind = 1; *code = 2; *p0 = nd.params.scalar; return true; + case OpType::SCALAR_DIVIDE: *kind = 1; *code = 3; *p0 = nd.params.scalar; return true; + case OpType::SCALAR_EXP: *kind = 1; *code = 4; return true; + case OpType::SCALAR_SQRT: *kind = 1; *code = 5; return true; + case OpType::SCALAR_COS: *kind = 1; *code = 6; return true; + case OpType::SCALAR_SIN: *kind = 1; *code = 7; return true; + case OpType::SCALAR_LOG: *kind = 1; *code = 8; return true; + case OpType::ABS: *kind = 1; *code = 9; return true; + case OpType::POW: *kind = 1; *code = 10; *p0 = nd.params.scalar; return true; + case OpType::SCALAR_NOT_EQUAL: *kind = 1; *code = 11; *p0 = nd.params.scalar; return true; + case OpType::LEAKY_RELU: *kind = 1; *code = 12; *p0 = nd.params.scalar; return true; + case OpType::ADD: *kind = 2; *code = 0; return true; + case OpType::ADD_CLIPPED: *kind = 2; *code = 1; return true; + case OpType::SUBTRACT: *kind = 2; *code = 2; return true; + case OpType::MULTIPLY: *kind = 2; *code = 3; return true; + case OpType::DIVIDE: *kind = 2; *code = 4; return true; + case OpType::NOT_EQUAL: *kind = 2; *code = 5; return true; + case OpType::CLAMP: *kind = 3; *code = 0; *p0 = nd.params.scalar; *p1 = nd.params.scale; return true; + case OpType::PRECISION_CAST: *kind = 4; + *code = nd.output_buffer.precision == Precision::FP32 ? 1 : 0; + return nd.output_buffer.precision == Precision::FP16 + || nd.output_buffer.precision == Precision::FP32; + default: return false; + } + }; + auto prec_ok = [](Precision p) { return p == Precision::FP16 || p == Precision::FP32; }; + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& nd0 = *nodes[i]; + if (nd0.op_type != OpType::RMS_NORM || nd0.input_ids.size() < 2) continue; + if (nd0.output_buffer.precision != Precision::FP16 || retyped(i)) continue; + const auto& osh = nd0.output_buffer.shape; + if (osh.empty()) continue; + size_t dim = osh.back(); + size_t rows = nd0.output_buffer.total_size / dim; + if (rows < 2 || dim < 32) continue; + long ai = idxof(nd0.input_ids[0]); + long wi = idxof(nd0.input_ids[1]); + if (ai < 0 || wi < 0 || plan->action[(size_t)ai] != -1 || retyped((size_t)ai)) continue; + const GraphNode& an = *nodes[(size_t)ai]; + if ((an.op_type != OpType::ADD && an.op_type != OpType::ADD_CLIPPED) + || an.input_ids.size() != 2 + || an.output_buffer.precision != Precision::FP16 + || an.output_buffer.total_size != nd0.output_buffer.total_size + || pinned[(size_t)ai]) continue; + long x0 = idxof(an.input_ids[0]); + long x1 = idxof(an.input_ids[1]); + if (x0 < 0 || x1 < 0 + || nodes[(size_t)x0]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)x1]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)x0]->output_buffer.total_size != an.output_buffer.total_size + || nodes[(size_t)x1]->output_buffer.total_size != an.output_buffer.total_size) continue; + long ww = deep_f16_source(wi); + if (ww < 0 || nodes[(size_t)ww]->output_buffer.total_size != dim) continue; + bool ordered = true; + for (size_t cc : cons[(size_t)ai]) + if (cc < i) { ordered = false; break; } + if (!ordered) continue; + MetalCluster c; + c.rule = 12; + c.a0 = (size_t)x0; c.a1 = (size_t)x1; c.a2 = (size_t)ai; c.a3 = (size_t)ww; + c.u0 = (uint32_t)rows; c.u1 = (uint32_t)dim; + c.f0 = nd0.params.epsilon; + c.b0 = an.op_type == OpType::ADD_CLIPPED ? 1 : 0; + c.b4 = i; + std::vector cover; + if (add_cluster(c, i, cover)) plan->action[(size_t)ai] = -3; + } + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& nd0 = *nodes[i]; + if (nd0.op_type != OpType::ADD || nd0.input_ids.size() != 2) continue; + if (nd0.output_buffer.precision != Precision::FP16 || retyped(i)) continue; + long i0 = idxof(nd0.input_ids[0]); + long i1 = idxof(nd0.input_ids[1]); + if (i0 < 0 || i1 < 0 || i0 == i1) continue; + long mm = -1, bs = -1; + if (nodes[(size_t)i0]->op_type == OpType::MATMUL) { mm = i0; bs = i1; } + else if (nodes[(size_t)i1]->op_type == OpType::MATMUL) { mm = i1; bs = i0; } + if (mm < 0 || plan->action[(size_t)mm] != -1 || pinned[(size_t)mm]) continue; + if (plan->action[(size_t)bs] == -2) continue; + if (retyped((size_t)mm) || retyped((size_t)bs)) continue; + if (cons[(size_t)mm].size() != 1) continue; + GraphNode& m = *nodes[(size_t)mm]; + if (m.input_ids.size() < 2) continue; + long li = idxof(m.input_ids[0]); + long ri = idxof(m.input_ids[1]); + if (li < 0 || ri < 0 || retyped((size_t)li) || retyped((size_t)ri)) continue; + if (plan->action[(size_t)li] == -2 || plan->action[(size_t)ri] == -2) continue; + const BufferDesc& lb = nodes[(size_t)li]->output_buffer; + const BufferDesc& rb = nodes[(size_t)ri]->output_buffer; + const BufferDesc& bb = nodes[(size_t)bs]->output_buffer; + if (lb.precision != Precision::FP16 || rb.precision != Precision::FP16 + || bb.precision != Precision::FP16) continue; + if (rb.shape.size() != 2 || lb.shape.empty()) continue; + size_t K = lb.shape.back(); + if (K == 0 || lb.total_size != K) continue; + size_t N = m.params.pretransposed_rhs ? rb.shape[0] : rb.shape[1]; + size_t rk = m.params.pretransposed_rhs ? rb.shape[1] : rb.shape[0]; + if (rk != K || N == 0 || m.output_buffer.total_size != N) continue; + if (bb.total_size != N || nd0.output_buffer.total_size != N) continue; + MetalCluster c; + c.rule = 13; + c.a0 = (size_t)li; c.a1 = (size_t)ri; c.a2 = (size_t)bs; + c.u0 = (uint32_t)K; c.u1 = (uint32_t)N; + c.b0 = m.params.pretransposed_rhs ? 1 : 0; + c.b4 = i; + std::vector cover; + cover.push_back((size_t)mm); + add_cluster(c, i, cover); + } + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& nd0 = *nodes[i]; + if (nd0.op_type != OpType::ADD || nd0.input_ids.size() != 2) continue; + if (nd0.output_buffer.precision != Precision::FP16 || retyped(i)) continue; + const auto& osh = nd0.output_buffer.shape; + if (osh.size() < 2) continue; + size_t D = osh.back(); + if (D == 0 || (D % 2) != 0) continue; + size_t H = nd0.output_buffer.total_size / D; + if (H == 0 || H * D != nd0.output_buffer.total_size) continue; + long mA = up_in(i, 0), mB = up_in(i, 1); + if (mA < 0 || mB < 0 || mA == mB) continue; + auto is_mul = [&](long j) { + return nodes[(size_t)j]->op_type == OpType::MULTIPLY + && nodes[(size_t)j]->input_ids.size() == 2 + && plan->action[(size_t)j] == -1 && !pinned[(size_t)j] && !retyped((size_t)j) + && nodes[(size_t)j]->output_buffer.precision == Precision::FP16 + && nodes[(size_t)j]->output_buffer.total_size == H * D; + }; + if (!is_mul(mA) || !is_mul(mB)) continue; + auto split_mul = [&](long m, long* big, long* small) -> bool { + long i0 = up_in((size_t)m, 0), i1 = up_in((size_t)m, 1); + if (i0 < 0 || i1 < 0) return false; + size_t t0 = nodes[(size_t)i0]->output_buffer.total_size; + size_t t1 = nodes[(size_t)i1]->output_buffer.total_size; + if (t0 == H * D && t1 == D) { *big = i0; *small = i1; return true; } + if (t1 == H * D && t0 == D) { *big = i1; *small = i0; return true; } + return false; + }; + long xA = -1, cA = -1, xB = -1, sB = -1; + if (!split_mul(mA, &xA, &cA) || !split_mul(mB, &xB, &sB)) continue; + long mulCos = -1, mulSin = -1, xi = -1, ci = -1, si = -1, cat = -1; + if (nodes[(size_t)xB]->op_type == OpType::CAT) { + mulCos = mA; mulSin = mB; xi = xA; ci = cA; si = sB; cat = xB; + } else if (nodes[(size_t)xA]->op_type == OpType::CAT) { + mulCos = mB; mulSin = mA; xi = xB; ci = sB; si = cA; cat = xA; + } else continue; + GraphNode& cn = *nodes[(size_t)cat]; + if (cn.input_ids.size() != 2 || plan->action[(size_t)cat] != -1 + || pinned[(size_t)cat] || retyped((size_t)cat)) continue; + if (cn.output_buffer.total_size != H * D) continue; + size_t cat_ax = (size_t)cn.params.axis; + if (cat_ax != cn.output_buffer.shape.size() - 1) continue; + long neg = up_in((size_t)cat, 0); + long lo = up_in((size_t)cat, 1); + if (neg < 0 || lo < 0) continue; + GraphNode& ng = *nodes[(size_t)neg]; + if (ng.op_type != OpType::SCALAR_MULTIPLY || ng.params.scalar != -1.0f + || plan->action[(size_t)neg] != -1 || pinned[(size_t)neg] || retyped((size_t)neg)) continue; + long hi = up_in((size_t)neg, 0); + if (hi < 0) continue; + GraphNode& hs = *nodes[(size_t)hi]; + GraphNode& ls = *nodes[(size_t)lo]; + if (hs.op_type != OpType::SLICE || ls.op_type != OpType::SLICE) continue; + if (plan->action[(size_t)hi] != -1 || plan->action[(size_t)lo] != -1 + || pinned[(size_t)hi] || pinned[(size_t)lo]) continue; + if (hs.params.slice_start != D / 2 || ls.params.slice_start != 0) continue; + size_t hlen = hs.params.slice_length ? hs.params.slice_length : D / 2; + size_t llen = ls.params.slice_length ? ls.params.slice_length : D / 2; + if (hlen != D / 2 || llen != D / 2) continue; + if ((size_t)hs.params.axis != hs.output_buffer.shape.size() - 1 + || (size_t)ls.params.axis != ls.output_buffer.shape.size() - 1) continue; + long xh = up_in((size_t)hi, 0), xl = up_in((size_t)lo, 0); + if (xh != xi || xl != xi) continue; + if (nodes[(size_t)xi]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)ci]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)si]->output_buffer.precision != Precision::FP16) continue; + if (!sole_use(mulCos, {i}) || !sole_use(mulSin, {i}) + || !sole_use(cat, {(size_t)mulSin}) || !sole_use(neg, {(size_t)cat}) + || !sole_use(hi, {(size_t)neg}) || !sole_use(lo, {(size_t)cat})) continue; + MetalCluster c; + c.rule = 14; + c.a0 = (size_t)xi; c.a1 = (size_t)ci; c.a2 = (size_t)si; + c.u0 = (uint32_t)H; c.u1 = (uint32_t)D; + c.b4 = i; + std::vector cover{(size_t)mulCos, (size_t)mulSin, (size_t)cat, + (size_t)neg, (size_t)hi, (size_t)lo}; + collect_chain((size_t)idxof(nd0.input_ids[0]), mulCos, cover); + collect_chain((size_t)idxof(nd0.input_ids[1]), mulSin, cover); + GraphNode& xn = *nodes[(size_t)xi]; + if (D <= 1024 && xn.op_type == OpType::RMS_NORM && xn.input_ids.size() >= 2 + && plan->action[(size_t)xi] == -1 && !pinned[(size_t)xi] && !retyped((size_t)xi) + && sole_use(xi, {(size_t)mulCos, (size_t)hi, (size_t)lo})) { + long rsrc = up_in((size_t)xi, 0); + long rw = up_in((size_t)xi, 1); + if (rw >= 0) rw = deep_f16_source(rw); + if (rsrc >= 0 && rw >= 0 + && nodes[(size_t)rsrc]->output_buffer.precision == Precision::FP16 + && nodes[(size_t)rsrc]->output_buffer.total_size == H * D + && nodes[(size_t)rw]->output_buffer.precision == Precision::FP16 + && nodes[(size_t)rw]->output_buffer.total_size == D) { + c.rule = 17; + c.a0 = (size_t)rsrc; c.a3 = (size_t)rw; + c.f0 = xn.params.epsilon; + cover.push_back((size_t)xi); + } + } + add_cluster(c, i, cover); + } + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& nd0 = *nodes[i]; + if (nd0.op_type != OpType::SCALAR_MULTIPLY || nd0.input_ids.empty()) continue; + if (nd0.output_buffer.precision != Precision::FP16 || retyped(i)) continue; + const auto& osh = nd0.output_buffer.shape; + if (osh.empty()) continue; + size_t D = osh.back(); + size_t rows = nd0.output_buffer.total_size / D; + if (D < 32 || rows == 0 || rows * D != nd0.output_buffer.total_size) continue; + long rn = up_in(i, 0); + if (rn < 0 || nodes[(size_t)rn]->op_type != OpType::RMS_NORM + || nodes[(size_t)rn]->input_ids.size() < 2 + || plan->action[(size_t)rn] != -1 || pinned[(size_t)rn] || retyped((size_t)rn) + || nodes[(size_t)rn]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)rn]->output_buffer.total_size != rows * D) continue; + if (!sole_use(rn, {i})) continue; + long src = up_in((size_t)rn, 0); + long w = up_in((size_t)rn, 1); + if (w >= 0) w = deep_f16_source(w); + if (src < 0 || w < 0 + || nodes[(size_t)src]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)w]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)w]->output_buffer.total_size != D) continue; + MetalCluster c; + c.rule = 15; + c.a0 = (size_t)src; c.a1 = (size_t)w; + c.u0 = (uint32_t)rows; c.u1 = (uint32_t)D; + c.f0 = nodes[(size_t)rn]->params.epsilon; + c.f1 = nd0.params.scalar; + c.b4 = i; + std::vector cover{(size_t)rn}; + collect_chain((size_t)idxof(nd0.input_ids[0]), rn, cover); + add_cluster(c, i, cover); + } + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& nd0 = *nodes[i]; + if (nd0.op_type != OpType::TOPK || nd0.input_ids.empty()) continue; + if (nd0.output_buffer.precision != Precision::FP32) continue; + size_t k = nd0.params.top_k; + if (k == 0 || k > 16) continue; + long sm = up_in(i, 0); + if (sm < 0 || nodes[(size_t)sm]->op_type != OpType::SCALAR_MULTIPLY + || nodes[(size_t)sm]->input_ids.empty() + || plan->action[(size_t)sm] != -1 || pinned[(size_t)sm] || retyped((size_t)sm) + || nodes[(size_t)sm]->output_buffer.precision != Precision::FP16) continue; + const auto& ssh = nodes[(size_t)sm]->output_buffer.shape; + if (ssh.size() != 2) continue; + size_t rows = ssh[0], E = ssh[1]; + if (rows == 0 || E == 0 || E > 4096) continue; + if (nd0.output_buffer.total_size != 2 * rows * k) continue; + long sfm = -1; + for (size_t cc : cons[(size_t)sm]) { + if (cc == i) continue; + const GraphNode& cn = *nodes[cc]; + if (cn.op_type == OpType::SOFTMAX && plan->action[cc] == -1 + && !pinned[cc] && !retyped(cc) + && cn.output_buffer.precision == Precision::FP16 + && cn.output_buffer.total_size == rows * E) { + if (sfm < 0) sfm = (long)cc; + else { sfm = -1; break; } + } else if (!is_alias_op(cn.op_type)) { sfm = -1; break; } + } + if (sfm < 0 || !sole_use(sm, {i, (size_t)sfm})) continue; + if ((size_t)sfm > i) continue; + bool ordered = true; + for (size_t cc : cons[(size_t)sfm]) + if (cc < i) { ordered = false; break; } + if (!ordered) continue; + long lg = up_in((size_t)sm, 0); + if (lg < 0 || nodes[(size_t)lg]->output_buffer.precision != Precision::FP16 + || nodes[(size_t)lg]->output_buffer.total_size != rows * E) continue; + MetalCluster c; + c.rule = 16; + c.a0 = (size_t)lg; c.a1 = (size_t)sfm; + c.u0 = (uint32_t)rows; c.u1 = (uint32_t)E; + c.b0 = k; + c.f0 = nodes[(size_t)sm]->params.scalar; + c.b4 = i; + std::vector cover{(size_t)sm}; + collect_chain((size_t)idxof(nd0.input_ids[0]), sm, cover); + if (add_cluster(c, i, cover)) plan->action[(size_t)sfm] = -3; + } + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + int kind, code; float p0, p1; + if (!ew_kind(*nodes[i], &kind, &code, &p0, &p1)) continue; + if (!prec_ok(nodes[i]->output_buffer.precision) || retyped(i)) continue; + size_t total = nodes[i]->output_buffer.total_size; + std::vector chain{i}; + size_t cur = i; + while (chain.size() < 12) { + if (cons[cur].size() != 1) break; + size_t nxt = cons[cur][0]; + if (plan->action[nxt] != -1) break; + int k2, c2; float q0, q1; + if (!ew_kind(*nodes[nxt], &k2, &c2, &q0, &q1)) break; + if (!prec_ok(nodes[nxt]->output_buffer.precision) || retyped(nxt)) break; + if (nodes[nxt]->output_buffer.total_size != total) break; + bool feeds = false; + for (size_t a = 0; a < nodes[nxt]->input_ids.size() && a < 2; ++a) + if (idxof(nodes[nxt]->input_ids[a]) == (long)cur) feeds = true; + if (!feeds) break; + chain.push_back(nxt); + cur = nxt; + } + size_t real_ops = 0; + for (size_t ci = 0; ci < chain.size(); ++ci) { + int k2, c2; float q0, q1; + ew_kind(*nodes[chain[ci]], &k2, &c2, &q0, &q1); + if (k2 != 4) ++real_ops; + } + if (chain.size() < 2 || real_ops == 0) continue; + std::vector blob; + std::vector sides; + std::vector side_f32; + bool ok = true; + long head_in = -1; + Precision run = Precision::FP16; + for (size_t ci = 0; ci < chain.size() && ok; ++ci) { + const GraphNode& cn = *nodes[chain[ci]]; + int k2, c2; float q0, q1; + ew_kind(cn, &k2, &c2, &q0, &q1); + long prev = ci == 0 ? -1 : (long)chain[ci - 1]; + if (ci == 0) { + long hi2 = idxof(cn.input_ids[0]); + if (hi2 < 0 || !prec_ok(nodes[(size_t)hi2]->output_buffer.precision) + || retyped((size_t)hi2) + || nodes[(size_t)hi2]->output_buffer.total_size != total) { ok = false; break; } + head_in = hi2; + run = nodes[(size_t)hi2]->output_buffer.precision; + } + if (k2 == 2) { + long in0 = idxof(cn.input_ids[0]); + long in1 = idxof(cn.input_ids[1]); + long side, chain_in; + bool rhs; + if (ci > 0 && in0 == prev) { chain_in = in0; side = in1; rhs = false; } + else if (ci > 0 && in1 == prev) { chain_in = in1; side = in0; rhs = true; } + else { chain_in = in0; side = in1; rhs = false; } + if (side < 0 || chain_in < 0 + || nodes[(size_t)side]->output_buffer.precision != run + || retyped((size_t)side) + || plan->action[(size_t)side] == -2) { ok = false; break; } + size_t inner = nodes[i]->output_buffer.shape.empty() + ? 1 : nodes[i]->output_buffer.shape.back(); + size_t stotal = nodes[(size_t)side]->output_buffer.total_size; + int bmode; + if (stotal == total) bmode = 0; + else if (inner > 1 && stotal * inner == total + && !nodes[(size_t)side]->output_buffer.shape.empty() + && nodes[(size_t)side]->output_buffer.shape.back() == 1) bmode = 1; + else if (inner > 1 && stotal == inner) bmode = 2; + else { ok = false; break; } + size_t slot = sides.size(); + for (size_t si = 0; si < sides.size(); ++si) + if (sides[si] == (size_t)side) slot = si; + if (slot == sides.size()) { + if (sides.size() >= 3) { ok = false; break; } + sides.push_back((size_t)side); + side_f32.push_back(run == Precision::FP32 ? 1 : 0); + } else if (side_f32[slot] != (run == Precision::FP32 ? 1 : 0)) { ok = false; break; } + c2 |= (rhs ? 16 : 0) | (run == Precision::FP32 ? 32 : 0) | ((int)slot << 6) | (bmode << 8); + } + if (k2 == 4) { + Precision from = ci == 0 ? run : nodes[(size_t)prev]->output_buffer.precision; + if (from == cn.output_buffer.precision) { ok = false; break; } + run = cn.output_buffer.precision; + } else if (cn.output_buffer.precision != run) { ok = false; break; } + EwChainStep st{k2, c2, q0, q1}; + float packed[4]; + std::memcpy(packed, &st, sizeof(st)); + blob.insert(blob.end(), packed, packed + 4); + } + if (!ok || head_in < 0) continue; + Precision head_prec = nodes[(size_t)head_in]->output_buffer.precision; + uint32_t flags = (head_prec == Precision::FP32 ? 1u : 0u) + | (nodes[chain.back()]->output_buffer.precision == Precision::FP32 ? 2u : 0u); + for (size_t si = 0; si < sides.size(); ++si) + if (side_f32[si]) flags |= (4u << si); + MetalCluster c; + c.rule = 11; + c.a0 = (size_t)head_in; + c.a1 = sides.size() > 0 ? sides[0] : (size_t)0; + c.a2 = sides.size() > 1 ? sides[1] : (size_t)0; + c.a3 = sides.size() > 2 ? sides[2] : (size_t)0; + c.u0 = (uint32_t)chain.size(); + c.u1 = (uint32_t)sides.size(); + c.b0 = plan->blobs.size(); + c.b1 = flags; + c.b4 = chain.back(); + plan->blobs.push_back(std::move(blob)); + std::vector cover(chain.begin(), chain.end() - 1); + add_cluster(c, chain.back(), cover); + } + } + + plan->exec_list.reserve(n); + for (size_t i = 0; i < n; ++i) + if (plan->action[i] != -2 && nodes[i]->op_type != OpType::INPUT) + plan->exec_list.push_back((uint32_t)i); + + bool decode_like = false; + for (auto& cl : plan->clusters) if (cl.rule == 1) { decode_like = true; break; } + if (decode_like) { + std::vector last_read(n, 0); + auto base_of = [&](size_t i) -> size_t { + size_t cur = i; + while (passthrough(*nodes[cur]) || plan->action[cur] == -2 || is_noop_transpose(*nodes[cur]) + || nodes[cur]->op_type == OpType::SLICE || nodes[cur]->op_type == OpType::INDEX) { + if (nodes[cur]->input_ids.empty()) break; + long j = idxof(nodes[cur]->input_ids[0]); + if (j < 0) break; + if (plan->action[cur] != -2 && !passthrough(*nodes[cur]) + && !is_noop_transpose(*nodes[cur])) { + bool aliasing = false; + const GraphNode& sn = *nodes[cur]; + if (sn.op_type == OpType::INDEX && sn.params.axis == 0) aliasing = true; + if (sn.op_type == OpType::SLICE) { + size_t ax = (size_t)sn.params.axis; + long jj = idxof(sn.input_ids[0]); + if (jj >= 0 && ax < nodes[(size_t)jj]->output_buffer.shape.size()) { + size_t outer = 1; + for (size_t d = 0; d < ax; ++d) outer *= nodes[(size_t)jj]->output_buffer.shape[d]; + if (outer == 1) aliasing = true; + } + } + if (!aliasing) break; + } + cur = (size_t)j; + } + return cur; + }; + for (size_t i = 0; i < n; ++i) { + for (size_t id : nodes[i]->input_ids) { + long j = idxof(id); + if (j >= 0) { + size_t b = base_of((size_t)j); + if (i > last_read[b]) last_read[b] = i; + } + } + } + for (auto& cl : plan->clusters) { + size_t anchor = cl.b4 ? cl.b4 : (cl.rule == 2 ? std::max(cl.b0, std::max(cl.b1, cl.b2)) : 0); + if (!anchor) continue; + for (size_t arg : {cl.a0, cl.a1, cl.a2, cl.a3, cl.a4, cl.a5, cl.b0, cl.b1, cl.b2, cl.b3}) { + if (arg == 0 || arg >= n) continue; + size_t b = base_of(arg); + if (anchor > last_read[b]) last_read[b] = anchor; + } + } + plan->arena_off.assign(n, -1); + struct Range { size_t off, size, free_at; }; + std::vector live; + size_t cursor = 0, peak = 0; + for (uint32_t ui : plan->exec_list) { + size_t i = ui; + const GraphNode& nd = *nodes[i]; + if (plan->action[i] == -2) continue; + if (pinned[i]) continue; + if (nd.op_type == OpType::KV_CACHE_STATE || nd.op_type == OpType::CONV_CACHE_STATE + || nd.op_type == OpType::RECURRENT_CACHE_STATE || nd.op_type == OpType::KV_CACHE_APPEND + || nd.op_type == OpType::PERSISTENT) continue; + if (passthrough(nd) || is_noop_transpose(nd)) continue; + if (nd.op_type == OpType::SLICE || nd.op_type == OpType::INDEX) continue; + if (last_read[i] == 0) continue; + size_t need = (nd.output_buffer.byte_size + 255) & ~size_t(255); + if (need == 0 || need > (8u << 20)) continue; + bool placed = false; + for (auto& r : live) { + if (r.free_at < i && r.size >= need) { + plan->arena_off[i] = (long)r.off; + r.free_at = last_read[i]; + placed = true; + break; + } + } + if (!placed) { + plan->arena_off[i] = (long)cursor; + live.push_back({cursor, need, last_read[i]}); + cursor += need; + if (cursor > peak) peak = cursor; + } + } + if (peak > 0 && peak < (256u << 20)) + plan->arena_base = (char*)cactus_metal_alloc_shared(peak); + if (!plan->arena_base) plan->arena_off.assign(n, -1); + } + + return plan; +} + +static void release_cluster_buffers(MetalFusePlan* p) { + for (auto& c : p->clusters) { + if (c.s0) { cactus_metal_free_shared(c.s0); c.s0 = nullptr; } + if (c.s1) { cactus_metal_free_shared(c.s1); c.s1 = nullptr; } + for (auto& s : c.sc) if (s) { cactus_metal_free_shared(s); s = nullptr; } + } +} + +void cactus_metal_plan_free(MetalFusePlan* p) { + if (!p) return; + release_cluster_buffers(p); + if (p->arena_base) { cactus_metal_free_shared(p->arena_base); p->arena_base = nullptr; } + delete p; +} + +bool cactus_metal_plan_has_arena(const MetalFusePlan* p) { + if (!p) return false; + if (p->arena_base) return true; + for (const auto& c : p->clusters) if (c.rule == 1) return true; + return false; +} + +void cactus_metal_plan_extend_last_use(const MetalFusePlan* p, std::vector& last_use) { + if (!p) return; + const size_t n = last_use.size(); + for (const auto& c : p->clusters) { + size_t ai = c.b4; + if (ai >= n) continue; + for (size_t v : {c.a0, c.a1, c.a2, c.a3, c.a4, c.a5, c.b0, c.b1, c.b2, c.b3}) + if (v < n && last_use[v] < ai) last_use[v] = ai; + } +} + +int32_t cactus_metal_plan_action(const MetalFusePlan* p, size_t i) { + return p && i < p->action.size() ? p->action[i] : -1; +} + + +const std::vector* cactus_metal_plan_exec_list(const MetalFusePlan* p) { + return p ? &p->exec_list : nullptr; +} + +void* cactus_metal_plan_arena_ptr(const MetalFusePlan* p, size_t i) { + if (!p || !p->arena_base || i >= p->arena_off.size() || p->arena_off[i] < 0) return nullptr; + return p->arena_base + p->arena_off[i]; +} + +bool cactus_metal_plan_fold(MetalFusePlan* p, + const std::vector>& nodes) { + if (!p || p->fold_h < 0 || !cactus_graph_fused_embed()) return false; + void* h = nodes[(size_t)p->fold_h]->output_buffer.get_data(); + void* ple = nodes[(size_t)p->fold_ple]->output_buffer.get_data(); + void* pos = nodes[(size_t)p->fold_pos]->output_buffer.get_data(); + if (!h || !ple) return false; + CactusQuantMatrix W = nodes[(size_t)p->fold_w]->output_buffer.to_cq_matrix(); + return cactus_graph_metal_fold_prologue(h, ple, pos, &W, p->fold_nl, p->fold_pd); +} + +bool cactus_metal_plan_encode(MetalFusePlan* p, int32_t cid, + const std::vector>& nodes, + const std::unordered_map& map) { + MetalCluster& c = p->clusters[(size_t)cid]; + auto data = [&](size_t i) { return nodes[i]->output_buffer.get_data(); }; + switch (c.rule) { + case 1: { + GraphNode& anchor = *nodes[c.b4]; + BufferDesc& kc = nodes[c.b2]->output_buffer; + BufferDesc& vc = nodes[c.b3]->output_buffer; + uint64_t* km = reinterpret_cast(kc.get_data()); + uint64_t* vm = reinterpret_cast(vc.get_data()); + if (!km || !vm) return false; + size_t clen = km[0], mx = km[1], sink = km[4]; + uint32_t hd = c.u1 & 0xFFFFu, nqh = c.u0; + if (km[2] != 1 || vm[3] != km[3] || hd == 0 || hd > 512u) return false; + bool has_new = (c.u1 & 0x10000u) != 0; + size_t ng = (hd + 31) / 32; + size_t win = anchor.params.window_size; + bool sliding = win > 0; + if (has_new && !sliding && clen >= mx) { + size_t ceiling = nodes[c.b2]->params.max_cache_seq_len; + if (!cactus_kv_cache_grow(kc, clen + 1, ceiling) + || !cactus_kv_cache_grow(vc, clen + 1, ceiling)) return false; + km = reinterpret_cast(kc.get_data()); + vm = reinterpret_cast(vc.get_data()); + if (!km || !vm) return false; + mx = km[1]; + } + uint32_t Wn = sliding ? (uint32_t)(mx - sink - 1) : 0u; + uint32_t Sn = sliding ? (uint32_t)sink : 0u; + uint32_t Rn = (Wn > Sn) ? (Wn - Sn) : 1u; + char* kb = (char*)kc.get_data(); + char* vb = (char*)vc.get_data(); + size_t slot = 0, kv_end; + if (has_new) { + bool wrap = sliding && clen >= (size_t)Wn; + slot = wrap ? (size_t)(Sn + ((clen - Sn) % Rn)) : clen; + kv_end = wrap ? (size_t)Wn : clen + 1; + } else { + kv_end = (sliding && clen > (size_t)Wn) ? (size_t)Wn : clen; + } + bool ok = cactus_metal_encode_attention_fused_i8( + anchor.output_buffer.get_data(), data(c.a0), + has_new ? data(c.a1) : nullptr, has_new ? data(c.a2) : nullptr, + kb + 64, vb + 64, kb + 64 + mx * hd, vb + 64 + mx * hd, + data(c.a3), has_new ? data(c.a4) : nullptr, has_new ? data(c.a5) : nullptr, + data(c.b0), data(c.b1), + nqh, hd, hd, + 0u, (uint32_t)kv_end, (uint32_t)slot, has_new ? 1u : 0u, + c.f0, c.f1, + mx * hd, mx * hd, mx * ng * sizeof(float), mx * ng * sizeof(float)); + if (!ok) return false; + if (has_new) { km[0] = clen + 1; vm[0] = clen + 1; } + return true; + } + case 2: { + size_t mm[3] = { c.b0, c.b1, c.b2 }; + CactusQuantMatrix W[3]; + void* outs[3]; + const CactusQuantMatrix* Wp[3]; + for (uint32_t bi = 0; bi < c.u0; ++bi) { + GraphNode& mn = *nodes[mm[bi]]; + auto it = map.find(mn.input_ids[1]); + if (it == map.end()) return false; + W[bi] = nodes[it->second]->output_buffer.to_cq_matrix(); + Wp[bi] = &W[bi]; + outs[bi] = mn.output_buffer.get_data(); + if (!outs[bi] || !c.sc[bi]) return false; + } + const void* x = nodes[c.a0]->output_buffer.get_data(); + if (!x) return false; + void* codes[3] = { c.sc[0], c.sc[1], c.sc[2] }; + if (cactus_metal_encode_transform_batch(x, Wp, (int)c.u0, codes)) { + const void* ccodes[3] = { c.sc[0], c.sc[1], c.sc[2] }; + if (!cactus_metal_encode_gemv_cat(outs, ccodes, Wp, (int)c.u0)) { + for (uint32_t bi = 0; bi < c.u0; ++bi) + if (!cactus_metal_encode_gemv_precoded(outs[bi], c.sc[bi], Wp[bi])) return false; + } + return true; + } + for (uint32_t bi = 0; bi < c.u0; ++bi) + if (!cactus_metal_encode_quant_matmul(outs[bi], x, Wp[bi])) return false; + return true; + } + case 7: { + GraphNode& anchor = *nodes[c.b4]; + GraphNode& down = *nodes[c.b0]; + auto itd = map.find(down.input_ids[1]); + auto itu = map.find(anchor.input_ids[1]); + if (itd == map.end() || itu == map.end() || !c.s0) return false; + CactusQuantMatrix Wd = nodes[itd->second]->output_buffer.to_cq_matrix(); + CactusQuantMatrix Wu = nodes[itu->second]->output_buffer.to_cq_matrix(); + const void* x = nodes[c.a0]->output_buffer.get_data(); + const void* row = nodes[c.a1]->output_buffer.get_data(); + if (row && c.u0 == 1u) row = (const char*)row + c.u1; + void* out = anchor.output_buffer.get_data(); + if (!x || !row || !out) return false; + if (!cactus_metal_encode_transform_gemv(c.s0, x, &Wd, row)) return false; + if (!cactus_metal_encode_transform_gemv(out, c.s0, &Wu, nullptr) + && !cactus_metal_encode_quant_matmul(out, c.s0, &Wu)) return false; + return true; + } + case 4: { + GraphNode& anchor = *nodes[c.b4]; + auto it = map.find(anchor.input_ids[1]); + if (it == map.end() || !c.s0) return false; + CactusQuantMatrix W = nodes[it->second]->output_buffer.to_cq_matrix(); + void* out = anchor.output_buffer.get_data(); + const void* gate = nodes[c.a0]->output_buffer.get_data(); + const void* up = nodes[c.a1]->output_buffer.get_data(); + if (!out || !gate || !up) return false; + if (cactus_metal_encode_swiglu_transform(c.s0, gate, up, &W, c.f0) + && cactus_metal_encode_gemv_precoded(out, c.s0, &W)) return true; + size_t M = W.K; + if (!cactus_metal_encode_swiglu(c.s0, gate, up, M, c.f0)) return false; + return cactus_metal_encode_quant_matmul(out, c.s0, &W); + } + case 3: { + GraphNode& anchor = *nodes[c.b4]; + const BufferDesc& src = nodes[c.a0]->output_buffer; + size_t dim = anchor.output_buffer.shape.empty() ? 0 : anchor.output_buffer.shape.back(); + if (dim == 0) return false; + size_t rows = src.total_size / dim; + float out_scale = 1.0f; + if (c.u1 == 1u) { + const void* sp = nodes[c.b1]->output_buffer.get_data(); + if (!sp) return false; + out_scale = (float)*(const __fp16*)sp; + } + if (c.u0 == 1u && rows == 1) { + GraphNode& nn = *nodes[c.a5]; + if (!nn.output_buffer.get_data()) { + if (!c.s1) return false; + nn.output_buffer.set_external(c.s1); + } + if (cactus_metal_encode_rms_norm_add_rms( + anchor.output_buffer.get_data(), nn.output_buffer.get_data(), + src.get_data(), data(c.a1), data(c.a2), data(c.a4), + rows, dim, c.f0, out_scale)) return true; + } + if (c.u1 == 1u && c.u0 != 1u && rows == 1) { + if (cactus_metal_encode_rms_norm_add_scale( + anchor.output_buffer.get_data(), src.get_data(), data(c.a1), data(c.a2), + rows, dim, c.f0, out_scale)) return true; + return false; + } + if (c.u1 == 1u) return false; + if (c.u0 == 1u) { + GraphNode& nn = *nodes[c.a5]; + if (!nn.output_buffer.get_data()) { + if (!c.s1) return false; + nn.output_buffer.set_external(c.s1); + } + if (!cactus_metal_encode_rms_norm_add( + anchor.output_buffer.get_data(), src.get_data(), data(c.a1), data(c.a2), + rows, dim, c.f0)) return false; + return cactus_metal_encode_rms_norm( + nn.output_buffer.get_data(), anchor.output_buffer.get_data(), data(c.a4), + rows, dim, c.f0); + } + return cactus_metal_encode_rms_norm_add( + anchor.output_buffer.get_data(), src.get_data(), data(c.a1), data(c.a2), + rows, dim, c.f0); + } + case 5: { + GraphNode& anchor = *nodes[c.b4]; + const BufferDesc& wb = nodes[c.a1]->output_buffer; + CactusQuantMatrix W = wb.to_cq_matrix(); + void* lg = anchor.output_buffer.get_data(); + if (!c.s0 || !cactus_metal_encode_quant_matmul_ortho(lg, data(c.a0), c.s0, &W)) return false; + size_t V = anchor.output_buffer.total_size; + if (!cactus_metal_encode_softcap(lg, lg, V, c.f0)) { + cactus_metal_encode_scalar(3, lg, lg, V, c.f0); + cactus_metal_encode_unary(1, lg, lg, V); + cactus_metal_encode_scalar(2, lg, lg, V, c.f0); + } + return cactus_graph_metal_tail(lg, V); + } + case 8: { + GraphNode& anchor = *nodes[c.b4]; + const char* pa = (const char*)nodes[c.a0]->output_buffer.get_data(); + const char* pb = (const char*)nodes[c.a1]->output_buffer.get_data(); + void* out = anchor.output_buffer.get_data(); + if (!pa || !pb || !out) return false; + return cactus_metal_encode_gemm_batch(out, pa + c.a2 * (c.b3 ? 4 : 2), pb + c.a3 * 2, + c.u1, (uint32_t)c.b0, (uint32_t)c.b1, c.u0, (int)c.b2, (int)c.b3); + } + case 9: { + GraphNode& anchor = *nodes[c.b4]; + const char* px = (const char*)nodes[c.a0]->output_buffer.get_data(); + const char* pw = (const char*)nodes[c.a1]->output_buffer.get_data(); + void* out = anchor.output_buffer.get_data(); + if (!px || !pw || !out) return false; + return cactus_metal_encode_conv1d_dw(out, px + c.a2 * 2, pw + c.a3 * 2, + c.u0, c.u1, (uint32_t)c.b0, (uint32_t)c.b1, (uint32_t)c.b2); + } + case 10: { + GraphNode& anchor = *nodes[c.b4]; + GraphNode& tn = *nodes[c.a1]; + const char* pp = (const char*)nodes[c.a0]->output_buffer.get_data(); + void* out = anchor.output_buffer.get_data(); + if (!pp || !out) return false; + long ii = -1; + { + auto it = map.find(tn.input_ids[0]); + if (it == map.end()) return false; + ii = (long)it->second; + } + const auto& ish = nodes[(size_t)ii]->output_buffer.shape; + const auto& osh = tn.output_buffer.shape; + const auto& perm = tn.params.permutation; + uint32_t ndim = (uint32_t)perm.size(); + uint32_t ax = c.u1; + if (ndim == 0 || ndim + 1 > 8 || ish.size() != ndim || osh.size() != ndim || ax > ndim) return false; + size_t istr[8]; + istr[ndim - 1] = 1; + for (size_t d = ndim - 1; d > 0; --d) istr[d - 1] = istr[d] * ish[d]; + uint32_t oshape[8], sstride[8]; + for (uint32_t d = 0; d < ax; ++d) { + oshape[d] = (uint32_t)osh[d]; + sstride[d] = (uint32_t)istr[perm[d]]; + } + oshape[ax] = c.u0; sstride[ax] = (uint32_t)c.b0; + for (uint32_t d = ax; d < ndim; ++d) { + oshape[d + 1] = (uint32_t)osh[d]; + sstride[d + 1] = (uint32_t)istr[perm[d]]; + } + size_t total = (size_t)c.u0 * tn.output_buffer.total_size; + return cactus_metal_encode_strided_copy(out, pp + c.a2 * 2, oshape, sstride, ndim + 1, + (uint32_t)total, 0, nodes[c.a0]->output_buffer.byte_size - c.a2 * 2, anchor.output_buffer.byte_size); + } + case 11: { + GraphNode& anchor = *nodes[c.b4]; + const void* in = nodes[c.a0]->output_buffer.get_data(); + void* out = anchor.output_buffer.get_data(); + if (!in || !out) return false; + const void* sides[3] = { nullptr, nullptr, nullptr }; + size_t side_elems[3] = { 0, 0, 0 }; + size_t side_idx[3] = { c.a1, c.a2, c.a3 }; + for (uint32_t si = 0; si < c.u1; ++si) { + sides[si] = nodes[side_idx[si]]->output_buffer.get_data(); + side_elems[si] = nodes[side_idx[si]]->output_buffer.total_size; + if (!sides[si]) return false; + } + const auto& blob = p->blobs[c.b0]; + size_t inner = anchor.output_buffer.shape.empty() ? 1 : anchor.output_buffer.shape.back(); + return cactus_metal_encode_elemwise_chain(out, in, blob.data(), c.u0, + sides[0], sides[1], sides[2], side_elems, anchor.output_buffer.total_size, + (uint32_t)c.b1, (uint32_t)inner); + } + case 12: { + GraphNode& anchor = *nodes[c.b4]; + const void* x = nodes[c.a0]->output_buffer.get_data(); + const void* res = nodes[c.a1]->output_buffer.get_data(); + void* ysum = nodes[c.a2]->output_buffer.get_data(); + const void* w = nodes[c.a3]->output_buffer.get_data(); + void* ynorm = anchor.output_buffer.get_data(); + if (!x || !res || !ysum || !w || !ynorm) return false; + return cactus_metal_encode_rms_norm_add_rows(ysum, ynorm, x, res, w, + c.u0, c.u1, c.f0, (int)c.b0); + } + case 13: { + GraphNode& anchor = *nodes[c.b4]; + const void* x = nodes[c.a0]->output_buffer.get_data(); + const void* w = nodes[c.a1]->output_buffer.get_data(); + const void* b = nodes[c.a2]->output_buffer.get_data(); + void* y = anchor.output_buffer.get_data(); + if (!x || !w || !b || !y) return false; + return cactus_metal_encode_gemv_bias(y, x, w, b, c.u0, c.u1, (int)c.b0); + } + case 14: { + GraphNode& anchor = *nodes[c.b4]; + const void* x = nodes[c.a0]->output_buffer.get_data(); + const void* cs = nodes[c.a1]->output_buffer.get_data(); + const void* sn = nodes[c.a2]->output_buffer.get_data(); + void* y = anchor.output_buffer.get_data(); + if (!x || !cs || !sn || !y) return false; + return cactus_metal_encode_rope_pair(y, x, cs, sn, c.u0, c.u1); + } + case 15: { + GraphNode& anchor = *nodes[c.b4]; + const void* x = nodes[c.a0]->output_buffer.get_data(); + const void* w = nodes[c.a1]->output_buffer.get_data(); + void* y = anchor.output_buffer.get_data(); + if (!x || !w || !y) return false; + return cactus_metal_encode_rms_norm_scale(y, x, w, c.u0, c.u1, c.f0, c.f1); + } + case 16: { + GraphNode& anchor = *nodes[c.b4]; + const void* lg = nodes[c.a0]->output_buffer.get_data(); + void* probs = nodes[c.a1]->output_buffer.get_data(); + void* tk = anchor.output_buffer.get_data(); + if (!lg || !probs || !tk) return false; + return cactus_metal_encode_softmax_topk(probs, tk, lg, c.u0, c.u1, c.b0, c.f0); + } + case 18: { + GraphNode& anchor = *nodes[c.b4]; + const void* a = nodes[c.a0]->output_buffer.get_data(); + const void* wa = nodes[c.a1]->output_buffer.get_data(); + const void* b = nodes[c.a2]->output_buffer.get_data(); + const void* wb = nodes[c.a3]->output_buffer.get_data(); + void* y = anchor.output_buffer.get_data(); + if (!a || !wa || !b || !wb || !y) return false; + return cactus_metal_encode_rms2_add_clip(y, a, wa, b, wb, c.u1, c.f0, c.f1); + } + case 17: { + GraphNode& anchor = *nodes[c.b4]; + const void* x = nodes[c.a0]->output_buffer.get_data(); + const void* w = nodes[c.a3]->output_buffer.get_data(); + const void* cs = nodes[c.a1]->output_buffer.get_data(); + const void* sn = nodes[c.a2]->output_buffer.get_data(); + void* y = anchor.output_buffer.get_data(); + if (!x || !w || !cs || !sn || !y) return false; + return cactus_metal_encode_rope_pair_rms(y, x, w, cs, sn, c.u0, c.u1, c.f0); + } + default: return false; + } +} diff --git a/cactus-graph/src/metal_runtime.cpp b/cactus-graph/src/metal_runtime.cpp new file mode 100644 index 000000000..311dec03a --- /dev/null +++ b/cactus-graph/src/metal_runtime.cpp @@ -0,0 +1,196 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include "metal_backend.h" + +#include +#include +#include + +namespace { + +static bool g_metal_argmax_valid = false; +static const float* g_metal_argmax_buf = nullptr; + +} + +static bool g_last_step_adjusted = false; +static bool g_last_amax_biased = false; +static bool g_prefill_consistent = false; + +void cactus_graph_mark_unadjusted() { + g_last_step_adjusted = false; + g_last_amax_biased = false; + g_metal_argmax_valid = false; +} + +void cactus_graph_set_prefill_consistent(bool on) { g_prefill_consistent = on; } + +bool cactus_graph_prefill_consistent() { return g_prefill_consistent; } + +struct GSampling { + bool active = false; + float rep_penalty = 1.0f; + std::vector recent; + const float* bias_dense = nullptr; + size_t bias_len = 0; + long long suppressed = -1; +}; +static GSampling g_samp; +static float* g_bias_dense = nullptr; +static size_t g_bias_dense_n = 0; + +void cactus_graph_set_sampling(const uint32_t* recent, int n_recent, float rep_penalty, + const float* bias_dense, size_t bias_len, + long long suppressed) { + g_samp.active = true; + g_samp.rep_penalty = rep_penalty; + g_samp.recent.assign(recent, recent + (n_recent > 0 ? n_recent : 0)); + g_samp.bias_dense = bias_dense; + g_samp.bias_len = bias_dense ? bias_len : 0; + g_samp.suppressed = suppressed; +} +void cactus_graph_clear_sampling() { g_samp = GSampling(); } +bool cactus_graph_metal_adjusted() { return g_last_step_adjusted; } +bool cactus_graph_metal_argmax_biased() { return g_last_amax_biased; } + +static void* g_plan_amax = nullptr; +static bool g_plan_tail_pending = false; +static bool g_plan_tail_adjusted = false; +static bool g_plan_tail_biased = false; + + +static const float* update_bias_dense(size_t vocab) { + if (!g_samp.bias_dense || g_samp.bias_len == 0) return nullptr; + if (!g_bias_dense || g_bias_dense_n < vocab) { + if (g_bias_dense) cactus_metal_free_shared(g_bias_dense); + g_bias_dense = (float*)cactus_metal_alloc_shared(vocab * sizeof(float)); + if (!g_bias_dense) { g_bias_dense_n = 0; return nullptr; } + g_bias_dense_n = vocab; + } + size_t n = g_samp.bias_len < vocab ? g_samp.bias_len : vocab; + std::memcpy(g_bias_dense, g_samp.bias_dense, n * sizeof(float)); + if (n < vocab) std::memset(g_bias_dense + n, 0, (vocab - n) * sizeof(float)); + return g_bias_dense; +} + +bool cactus_graph_metal_tail(void* logits, size_t vocab) { + g_plan_tail_pending = false; + bool adjusted = false; + if (g_samp.active && + ((g_samp.rep_penalty != 1.0f && !g_samp.recent.empty()) || g_samp.suppressed >= 0)) { + adjusted = cactus_metal_encode_adjust_logits(logits, vocab, + g_samp.recent.data(), (uint32_t)g_samp.recent.size(), + g_samp.suppressed, g_samp.rep_penalty); + if (!adjusted) return false; + } + if (!g_plan_amax) { + g_plan_amax = cactus_metal_alloc_shared(3 * sizeof(float)); + if (!g_plan_amax) return false; + } + const float* bias = g_samp.active ? update_bias_dense(vocab) : nullptr; + if (!cactus_metal_encode_argmax(logits, (uint32_t)vocab, g_plan_amax, bias)) return false; + g_plan_tail_pending = true; + g_plan_tail_adjusted = adjusted; + g_plan_tail_biased = (bias != nullptr); + return true; +} + +void cactus_graph_metal_tail_commit() { + if (!g_plan_tail_pending) return; + g_plan_tail_pending = false; + g_metal_argmax_buf = (const float*)g_plan_amax; + g_metal_argmax_valid = true; + g_last_step_adjusted = g_plan_tail_adjusted; + g_last_amax_biased = g_plan_tail_biased; +} + + +bool cactus_graph_metal_argmax(uint32_t* idx, float* best, float* second) { + if (!g_metal_argmax_valid || !g_metal_argmax_buf) return false; + g_metal_argmax_valid = false; + *best = g_metal_argmax_buf[0]; *second = g_metal_argmax_buf[1]; + *idx = (uint32_t)g_metal_argmax_buf[2]; + return true; +} + +static FusedEmbedCtx g_fe; +void cactus_graph_set_fused_embed(const FusedEmbedCtx* ctx) { + if (ctx && ctx->ok) g_fe = *ctx; else g_fe.ok = false; +} +const FusedEmbedCtx* cactus_graph_fused_embed() { return g_fe.ok ? &g_fe : nullptr; } + +bool cactus_graph_metal_fold_prologue(void* h_buf, void* ple_buf, void* pos_buf, + const CactusQuantMatrix* lm_head, size_t nl, size_t ple_dim) { + if (!g_fe.ok || !lm_head) return false; + static void* pe = nullptr; + static void* pa = nullptr; + static void* pj = nullptr; + static void* pjs = nullptr; + static size_t cap = 0; + static size_t cap_pe = 0; + const size_t PK = g_fe.proj.N; + const size_t EK = g_fe.ple.K; + if (cap_pe < EK) { + if (pe) cactus_metal_free_shared(pe); + pe = cactus_metal_alloc_shared(EK * 2); + cap_pe = pe ? EK : 0; + } + if (cap < PK) { + if (pa) cactus_metal_free_shared(pa); + if (pj) cactus_metal_free_shared(pj); + if (pjs) cactus_metal_free_shared(pjs); + pa = cactus_metal_alloc_shared(PK * 2); + pj = cactus_metal_alloc_shared(PK * 2); + pjs = cactus_metal_alloc_shared(PK * 2); + cap = (pa && pj && pjs) ? PK : 0; + } + if (!pe || !pa || !pj || !pjs) return false; + uint32_t tok = (uint32_t)g_fe.token_id; + if (!cactus_metal_encode_embedding_ortho(h_buf, tok, lm_head, g_fe.emb_scale)) return false; + if (!cactus_metal_encode_embedding_hadamard(pe, tok, &g_fe.ple)) return false; + cactus_metal_encode_scalar(2, pa, pe, PK, g_fe.ple_scale); + if (!cactus_metal_encode_transform_gemv(pj, h_buf, &g_fe.proj, nullptr) + && !cactus_metal_encode_quant_matmul(pj, h_buf, &g_fe.proj)) return false; + cactus_metal_encode_scalar(2, pjs, pj, PK, g_fe.proj_scale); + if (!cactus_metal_encode_rms_norm_add_scale(ple_buf, pjs, g_fe.rms_weight, pa, nl, ple_dim, g_fe.rms_eps, g_fe.final_scale)) { + if (!cactus_metal_encode_rms_norm_add(pa, pjs, g_fe.rms_weight, pa, nl, ple_dim, g_fe.rms_eps)) return false; + cactus_metal_encode_scalar(2, ple_buf, pa, PK, g_fe.final_scale); + } + if (pos_buf) *(float*)pos_buf = (float)g_fe.position; + return true; +} + +bool CactusGraph::extract_ple_pathway(FusedEmbedCtx& ctx) const { + + if (nodes_.size() != 23) return false; + auto op = [&](size_t i){ return nodes_[i]->op_type; }; + if (op(3) != OpType::INPUT || op(4) != OpType::INPUT || op(5) != OpType::INPUT) return false; + if (op(8) != OpType::EMBEDDING || op(12) != OpType::MATMUL || op(19) != OpType::RMS_NORM) return false; + if (op(7) != OpType::SCALAR_MULTIPLY || op(9) != OpType::SCALAR_MULTIPLY || + op(14) != OpType::SCALAR_MULTIPLY || op(22) != OpType::SCALAR_MULTIPLY) return false; + const BufferDesc& pleW = nodes_[3]->output_buffer; + const BufferDesc& projW = nodes_[5]->output_buffer; + if (!pleW.is_cq() || (pleW.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) || !projW.is_cq()) return false; + ctx.ple = pleW.to_cq_matrix(); + ctx.proj = projW.to_cq_matrix(); + ctx.rms_weight = nodes_[4]->output_buffer.get_data(); + ctx.emb_scale = nodes_[7]->params.scalar; + ctx.ple_scale = nodes_[9]->params.scalar; + ctx.proj_scale = nodes_[14]->params.scalar; + ctx.final_scale = nodes_[22]->params.scalar; + ctx.rms_eps = nodes_[19]->params.epsilon; + if (!ctx.rms_weight || ctx.emb_scale == 0.0f || ctx.proj.N == 0) return false; + ctx.ok = true; + return true; +} + +void cactus_graph_on_destroy(const void* graph) { + (void)graph; + g_samp = GSampling(); + cactus_metal_invalidate_host_wraps(); +} + +CactusGraph::~CactusGraph() { + cactus_graph_on_destroy(this); + invalidate_metal_state(); +} diff --git a/cactus-graph/src/ops_cache.cpp b/cactus-graph/src/ops_cache.cpp new file mode 100644 index 000000000..ed9db8edd --- /dev/null +++ b/cactus-graph/src/ops_cache.cpp @@ -0,0 +1,757 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include +#include +#include +#include +#include +#include + +namespace { + +struct ConvCacheMetadata { + uint64_t head; + uint64_t count; + uint64_t window_size; + uint64_t hidden_dim; + uint64_t reserved[4]; +}; + +static_assert(sizeof(ConvCacheMetadata) == 64, "ConvCacheMetadata must be 64 bytes"); + +inline ConvCacheMetadata* get_conv_meta(BufferDesc& buf) { + return static_cast(buf.get_data()); +} + +inline __fp16* get_conv_data(BufferDesc& buf) { + return reinterpret_cast<__fp16*>(static_cast(buf.get_data()) + sizeof(ConvCacheMetadata)); +} + +struct CacheMetadata { + uint64_t current_seq_len; + uint64_t max_seq_len; + uint64_t num_kv_heads; + uint64_t head_dim; + uint64_t sink_size; + uint64_t num_slots; + uint64_t reserved[2]; +}; + +static_assert(sizeof(CacheMetadata) == 64, "CacheMetadata must be 64 bytes"); + +inline size_t cache_buffer_size(size_t max_seq, size_t kv_heads, size_t head_dim) { + size_t num_groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + return sizeof(CacheMetadata) + max_seq * kv_heads * head_dim + max_seq * kv_heads * num_groups * sizeof(float); +} + +inline size_t fp16_cache_elements(size_t max_seq, size_t kv_heads, size_t head_dim) { + return (sizeof(CacheMetadata) / sizeof(__fp16)) + max_seq * kv_heads * head_dim; +} + +inline size_t kv_slot_off(const BufferDesc& buf, size_t slot) { + if (slot == 0) return 0; + const CacheMetadata* m = reinterpret_cast(buf.get_data()); + size_t per_slot = (buf.precision == Precision::FP16) + ? fp16_cache_elements(m->max_seq_len, m->num_kv_heads, m->head_dim) * sizeof(__fp16) + : cache_buffer_size(m->max_seq_len, m->num_kv_heads, m->head_dim); + return slot * per_slot; +} + +inline CacheMetadata* get_meta(BufferDesc& buf, size_t slot = 0) { + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot)); +} + +inline const CacheMetadata* get_meta(const BufferDesc& buf, size_t slot = 0) { + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot)); +} + +inline int8_t* get_int8_data(BufferDesc& buf, size_t slot = 0) { + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot) + sizeof(CacheMetadata)); +} + +inline const int8_t* get_int8_data(const BufferDesc& buf, size_t slot = 0) { + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot) + sizeof(CacheMetadata)); +} + +inline float* get_scales(BufferDesc& buf, size_t max_seq, size_t kv_heads, size_t head_dim, size_t slot = 0) { + size_t int8_bytes = max_seq * kv_heads * head_dim; + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot) + sizeof(CacheMetadata) + int8_bytes); +} + +inline const float* get_scales(const BufferDesc& buf, size_t max_seq, size_t kv_heads, size_t head_dim, size_t slot = 0) { + size_t int8_bytes = max_seq * kv_heads * head_dim; + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot) + sizeof(CacheMetadata) + int8_bytes); +} + +inline __fp16* get_fp16_data(BufferDesc& buf, size_t slot = 0) { + return reinterpret_cast<__fp16*>(static_cast(buf.get_data()) + kv_slot_off(buf, slot) + sizeof(CacheMetadata)); +} + +inline const __fp16* get_fp16_data(const BufferDesc& buf, size_t slot = 0) { + return reinterpret_cast(static_cast(buf.get_data()) + kv_slot_off(buf, slot) + sizeof(CacheMetadata)); +} + +inline bool use_fp16_kv_cache() { + static const bool cached = [] { + const char* value = std::getenv("CACTUS_KV_CACHE_FP16"); + return value != nullptr && std::strcmp(value, "1") == 0; + }(); + return cached; +} + +constexpr size_t kInitialCacheEntries = 256; + +inline bool kv_cache_resident() { + return cactus_default_backend() == ComputeBackend::METAL; +} + +inline bool resize_cache_buffer(BufferDesc& buf, size_t new_max) { + auto* meta = get_meta(buf); + size_t cur = meta->max_seq_len; + const size_t current_seq = meta->current_seq_len; + if (new_max == cur || new_max < current_seq) return false; + + const size_t kv_heads = meta->num_kv_heads; + const size_t hdim = meta->head_dim; + const bool fp16_cache = buf.precision == Precision::FP16; + + size_t total = fp16_cache ? fp16_cache_elements(new_max, kv_heads, hdim) + : cache_buffer_size(new_max, kv_heads, hdim); + BufferDesc resized({total}, fp16_cache ? Precision::FP16 : Precision::INT8); + const bool metal = kv_cache_resident(); + void* old_data = buf.get_data(); + if (metal) { + cactus_metal_session_sync(); + void* p = cactus_metal_alloc_shared(resized.byte_size); + if (p) resized.set_external(p); else resized.allocate(); + } else { + resized.allocate(); + } + std::memset(resized.get_data(), 0, resized.byte_size); + + std::memcpy(resized.get_data(), buf.get_data(), sizeof(CacheMetadata)); + if (fp16_cache) { + std::memcpy(get_fp16_data(resized), get_fp16_data(buf), + current_seq * kv_heads * hdim * sizeof(__fp16)); + } else { + std::memcpy(get_int8_data(resized), get_int8_data(buf), current_seq * kv_heads * hdim); + const size_t groups = (hdim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + std::memcpy(get_scales(resized, new_max, kv_heads, hdim), + get_scales(buf, cur, kv_heads, hdim), + current_seq * kv_heads * groups * sizeof(float)); + } + get_meta(resized)->max_seq_len = new_max; + buf = std::move(resized); + cactus_metal_free_shared(old_data); + return true; +} + +inline bool grow_cache_buffer(BufferDesc& buf, size_t needed, size_t ceiling) { + size_t cur = get_meta(buf)->max_seq_len; + if (needed <= cur || cur >= ceiling) return false; + size_t new_max = cur; + while (new_max < needed) new_max <<= 1; + if (new_max > ceiling) new_max = ceiling; + if (new_max <= cur) return false; + return resize_cache_buffer(buf, new_max); +} + +} // namespace + +bool cactus_kv_cache_grow(BufferDesc& buf, size_t needed, size_t ceiling) { + return grow_cache_buffer(buf, needed, ceiling); +} + +void compute_kv_cache_state_node( + GraphNode& node, + const nodes_vector&, + const node_index_map_t&) { + + if (node.output_buffer.get_data()) return; + + size_t ceiling = node.params.max_cache_seq_len; + size_t window = node.params.window_size; + size_t num_slots = node.params.cache_num_slots > 0 ? node.params.cache_num_slots : 1; + bool sliding = window > 0 && window < ceiling; + size_t max_seq; + if (sliding) max_seq = std::min(ceiling, window + node.params.cache_sink_size + 1); + else if (num_slots > 1 || window > 0) max_seq = ceiling; + else max_seq = std::min(ceiling, kInitialCacheEntries); + size_t kv_heads = node.params.num_kv_heads; + size_t hdim = node.params.head_dim; + const bool fp16_cache = use_fp16_kv_cache(); + size_t per_slot = fp16_cache + ? fp16_cache_elements(max_seq, kv_heads, hdim) + : cache_buffer_size(max_seq, kv_heads, hdim); + + node.output_buffer = BufferDesc({num_slots * per_slot}, fp16_cache ? Precision::FP16 : Precision::INT8); + if (kv_cache_resident()) { + void* p = cactus_metal_alloc_shared(node.output_buffer.byte_size); + if (p) node.output_buffer.set_external(p); else node.output_buffer.allocate(); + } else { + node.output_buffer.allocate(); + } + std::memset(node.output_buffer.get_data(), 0, node.output_buffer.byte_size); + + auto* meta0 = get_meta(node.output_buffer, 0); + meta0->current_seq_len = 0; + meta0->max_seq_len = max_seq; + meta0->num_kv_heads = kv_heads; + meta0->head_dim = hdim; + meta0->sink_size = node.params.cache_sink_size; + meta0->num_slots = num_slots; + for (size_t s = 1; s < num_slots; ++s) { + *get_meta(node.output_buffer, s) = *meta0; + } +} + +void kv_append_one_slot(BufferDesc& cache_buf, size_t slot, size_t num_slots, + const __fp16* source, size_t new_seq_len, + size_t window_size, size_t ceiling) { + auto* meta = get_meta(cache_buf, slot); + size_t current_len = meta->current_seq_len; + size_t max_len = meta->max_seq_len; + size_t kv_heads = meta->num_kv_heads; + size_t hdim = meta->head_dim; + size_t sink = meta->sink_size; + size_t num_groups = (hdim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + size_t int8_stride = kv_heads * hdim; + size_t scale_stride = kv_heads * num_groups; + bool sliding = window_size > 0 && window_size < ceiling; + + if (num_slots <= 1 && !sliding && current_len + new_seq_len > max_len) { + if (grow_cache_buffer(cache_buf, current_len + new_seq_len, ceiling)) { + meta = get_meta(cache_buf, slot); + max_len = meta->max_seq_len; + } + } + + if (cache_buf.precision == Precision::FP16) { + size_t stride = kv_heads * hdim; + __fp16* fp16_base = get_fp16_data(cache_buf, slot); + size_t window = sliding ? window_size : max_len; + size_t new_total = current_len + new_seq_len; + if (new_total > window) { + size_t keep_sink = std::min({sink, current_len, window}); + size_t tail_capacity = window - keep_sink; + if (new_seq_len >= tail_capacity) { + if (tail_capacity > 0) { + size_t source_offset = new_seq_len - tail_capacity; + std::memcpy(fp16_base + keep_sink * stride, source + source_offset * stride, + tail_capacity * stride * sizeof(__fp16)); + } + meta->current_seq_len = keep_sink + tail_capacity; + return; + } + size_t remaining = std::min(tail_capacity - new_seq_len, current_len - keep_sink); + size_t shift_src = current_len - remaining; + if (remaining > 0 && shift_src > keep_sink) { + std::memmove(fp16_base + keep_sink * stride, fp16_base + shift_src * stride, + remaining * stride * sizeof(__fp16)); + } + size_t append_offset = keep_sink + remaining; + std::memcpy(fp16_base + append_offset * stride, source, new_seq_len * stride * sizeof(__fp16)); + meta->current_seq_len = append_offset + new_seq_len; + } else { + std::memcpy(fp16_base + current_len * stride, source, new_seq_len * stride * sizeof(__fp16)); + meta->current_seq_len = new_total; + } + return; + } + + int8_t* int8_base = get_int8_data(cache_buf, slot); + float* scale_base = get_scales(cache_buf, max_len, kv_heads, hdim, slot); + size_t window = sliding ? window_size : max_len; + size_t new_total = current_len + new_seq_len; + if (new_total > window) { + size_t keep_sink = std::min({sink, current_len, window}); + size_t tail_capacity = window - keep_sink; + if (new_seq_len >= tail_capacity) { + if (tail_capacity > 0) { + size_t source_offset = new_seq_len - tail_capacity; + cactus_quantize_kv_fp16_to_int8(source + source_offset * int8_stride, + int8_base + keep_sink * int8_stride, scale_base + keep_sink * scale_stride, + tail_capacity, kv_heads, hdim); + } + meta->current_seq_len = keep_sink + tail_capacity; + return; + } + size_t remaining = std::min(tail_capacity - new_seq_len, current_len - keep_sink); + size_t shift_src = current_len - remaining; + if (remaining > 0 && shift_src > keep_sink) { + std::memmove(int8_base + keep_sink * int8_stride, int8_base + shift_src * int8_stride, + remaining * int8_stride); + std::memmove(scale_base + keep_sink * scale_stride, scale_base + shift_src * scale_stride, + remaining * scale_stride * sizeof(float)); + } + size_t append_offset = keep_sink + remaining; + cactus_quantize_kv_fp16_to_int8(source, int8_base + append_offset * int8_stride, + scale_base + append_offset * scale_stride, new_seq_len, kv_heads, hdim); + meta->current_seq_len = append_offset + new_seq_len; + } else { + cactus_quantize_kv_fp16_to_int8(source, int8_base + current_len * int8_stride, + scale_base + current_len * scale_stride, new_seq_len, kv_heads, hdim); + meta->current_seq_len = new_total; + } +} + +void compute_kv_cache_append_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& new_kv = get_input(node, 0, nodes, node_index_map); + auto& cache_buf = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; + size_t slot = node.params.cache_slot; + const auto* m0 = get_meta(cache_buf, 0); + size_t num_slots = m0->num_slots ? m0->num_slots : 1; + size_t kv_heads = m0->num_kv_heads; + size_t hdim = m0->head_dim; + size_t ceiling = nodes[node_index_map.at(node.input_ids[1])]->params.max_cache_seq_len; + size_t window_size = node.params.window_size; + const __fp16* src = new_kv.data_as<__fp16>(); + + size_t batch = new_kv.shape.empty() ? 1 : new_kv.shape[0]; + if (num_slots > 1 && batch > 1 && batch <= num_slots) { + size_t row_elems = new_kv.total_size / batch; + size_t per_row_tokens = row_elems / (kv_heads * hdim); + for (size_t i = 0; i < batch; ++i) { + kv_append_one_slot(cache_buf, i, num_slots, src + i * row_elems, per_row_tokens, window_size, ceiling); + } + *node.output_buffer.data_as() = static_cast(get_meta(cache_buf, batch - 1)->current_seq_len); + return; + } + + size_t new_seq_len = new_kv.total_size / (kv_heads * hdim); + kv_append_one_slot(cache_buf, slot, num_slots, src, new_seq_len, window_size, ceiling); + *node.output_buffer.data_as() = static_cast(get_meta(cache_buf, slot)->current_seq_len); +} + +void compute_attention_cached_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& query_buf = get_input(node, 0, nodes, node_index_map); + const auto& key_new_buf = get_input(node, 1, nodes, node_index_map); + const auto& val_new_buf = get_input(node, 2, nodes, node_index_map); + const auto& k_cache_buf = get_input(node, 3, nodes, node_index_map); + const auto& v_cache_buf = get_input(node, 4, nodes, node_index_map); + + size_t slot = node.params.cache_slot; + const auto* k_meta = get_meta(k_cache_buf, slot); + size_t num_slots = k_meta->num_slots ? k_meta->num_slots : 1; + size_t cache_len = k_meta->current_seq_len; + size_t k_max = k_meta->max_seq_len; + size_t kv_heads = k_meta->num_kv_heads; + size_t hdim = k_meta->head_dim; + + const auto* v_meta = get_meta(v_cache_buf, slot); + size_t v_hdim = node.params.v_head_dim > 0 ? node.params.v_head_dim : hdim; + size_t v_max = v_meta->max_seq_len; + + const int8_t* cached_keys = get_int8_data(k_cache_buf, slot); + const float* k_scales = get_scales(k_cache_buf, k_max, kv_heads, hdim, slot); + const int8_t* cached_values = get_int8_data(v_cache_buf, slot); + const float* v_scales = get_scales(v_cache_buf, v_max, kv_heads, v_hdim, slot); + + const auto& q_shape = query_buf.shape; + size_t batch_size = q_shape[0]; + size_t seq_len = q_shape[1]; + size_t num_q_heads = q_shape[2]; + + size_t new_seq_len = key_new_buf.total_size / (kv_heads * hdim); + size_t history_len = (cache_len >= new_seq_len) ? cache_len - new_seq_len : 0; + bool cache_only_attention = false; + size_t position_offset = node.params.position_offset; + if (position_offset == std::numeric_limits::max()) { + position_offset = history_len; + } else if (position_offset == std::numeric_limits::max() - 1) { + position_offset = (cache_len >= seq_len) ? cache_len - seq_len : 0; + history_len = cache_len; + new_seq_len = 0; + cache_only_attention = true; + } + + if (batch_size > 1 && num_slots > 1) { + bool fp16_cache = (k_cache_buf.precision == Precision::FP16 || v_cache_buf.precision == Precision::FP16); + size_t q_stride = query_buf.total_size / batch_size; + size_t knew_stride = key_new_buf.total_size / batch_size; + size_t vnew_stride = val_new_buf.total_size / batch_size; + size_t out_stride = node.output_buffer.total_size / batch_size; + const __fp16* q_all = query_buf.data_as<__fp16>(); + const __fp16* knew_all = key_new_buf.data_as<__fp16>(); + const __fp16* vnew_all = val_new_buf.data_as<__fp16>(); + __fp16* out_all = node.output_buffer.data_as<__fp16>(); + for (size_t i = 0; i < batch_size; ++i) { + size_t ci = get_meta(k_cache_buf, i)->current_seq_len; + size_t hist_i = (ci >= seq_len) ? ci - seq_len : 0; + if (fp16_cache) { + cactus_attention_f16( + q_all + i * q_stride, + get_fp16_data(k_cache_buf, i), + get_fp16_data(v_cache_buf, i), + out_all + i * out_stride, + 1, seq_len, ci, + num_q_heads, kv_heads, hdim, + node.params.scale, nullptr, hist_i, node.params.window_size, + true, false, false, v_hdim); + } else { + cactus_attention_hybrid_int8_fp16( + q_all + i * q_stride, + get_int8_data(k_cache_buf, i), + get_int8_data(v_cache_buf, i), + get_scales(k_cache_buf, k_max, kv_heads, hdim, i), + get_scales(v_cache_buf, v_max, kv_heads, v_hdim, i), + knew_all + i * knew_stride, + vnew_all + i * vnew_stride, + out_all + i * out_stride, + 1, seq_len, hist_i, seq_len, + num_q_heads, kv_heads, hdim, + node.params.scale, hist_i, true, node.params.window_size, + KV_QUANT_GROUP_SIZE, v_hdim); + } + } + return; + } + + if (k_cache_buf.precision == Precision::FP16 || v_cache_buf.precision == Precision::FP16) { + cactus_attention_f16( + query_buf.data_as<__fp16>(), + get_fp16_data(k_cache_buf, slot), + get_fp16_data(v_cache_buf, slot), + node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, cache_len, + num_q_heads, kv_heads, hdim, + node.params.scale, + nullptr, + position_offset, + node.params.window_size, + true, + false, + false, + v_hdim); + return; + } + + cactus_attention_hybrid_int8_fp16( + query_buf.data_as<__fp16>(), + cached_keys, + cached_values, + k_scales, + v_scales, + key_new_buf.data_as<__fp16>(), + val_new_buf.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, history_len, cache_only_attention ? 0 : seq_len, + num_q_heads, kv_heads, hdim, + node.params.scale, + position_offset, + true, + node.params.window_size, + KV_QUANT_GROUP_SIZE, + v_hdim); +} + +void compute_conv_cache_state_node( + GraphNode& node, + const nodes_vector&, + const node_index_map_t&) { + + if (node.output_buffer.get_data()) return; + + size_t ws = node.params.window_size; + size_t hd = node.params.head_dim; + size_t total = sizeof(ConvCacheMetadata) + ws * hd * sizeof(__fp16); + + node.output_buffer = BufferDesc({total}, Precision::INT8); + node.output_buffer.allocate(); + std::memset(node.output_buffer.get_data(), 0, total); + + auto* meta = get_conv_meta(node.output_buffer); + meta->head = 0; + meta->count = 0; + meta->window_size = ws; + meta->hidden_dim = hd; +} + +void compute_conv_cache_append_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& new_data = get_input(node, 0, nodes, node_index_map); + auto& cache_buf = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; + auto* meta = get_conv_meta(cache_buf); + + size_t ws = meta->window_size; + size_t hd = meta->hidden_dim; + size_t head = meta->head; + size_t count = meta->count; + + size_t num_rows = new_data.total_size / hd; + if (num_rows == 0) return; + + __fp16* cache_data = get_conv_data(cache_buf); + + const __fp16* src; + std::vector<__fp16> converted; + if (new_data.precision == Precision::FP16) { + src = new_data.data_as<__fp16>(); + } else if (new_data.precision == Precision::FP32) { + converted.resize(new_data.total_size); + Quantization::fp32_to_fp16(new_data.data_as(), converted.data(), new_data.total_size); + src = converted.data(); + } else { + converted.resize(new_data.total_size); + Quantization::int8_to_fp16(new_data.data_as(), converted.data(), new_data.total_size); + src = converted.data(); + } + + size_t copy_rows = std::min(num_rows, ws); + size_t start_row = num_rows > ws ? num_rows - ws : 0; + + for (size_t i = 0; i < copy_rows; ++i) { + std::memcpy(cache_data + head * hd, src + (start_row + i) * hd, hd * sizeof(__fp16)); + head = (head + 1) % ws; + if (count < ws) ++count; + } + + meta->head = head; + meta->count = count; + + __fp16* out = node.output_buffer.data_as<__fp16>(); + if (count < ws) { + const size_t pad_rows = ws - count; + std::memset(out, 0, pad_rows * hd * sizeof(__fp16)); + std::memcpy(out + pad_rows * hd, cache_data, count * hd * sizeof(__fp16)); + } else { + size_t tail_rows = ws - head; + if (tail_rows > 0) { + std::memcpy(out, cache_data + head * hd, tail_rows * hd * sizeof(__fp16)); + } + if (head > 0) { + std::memcpy(out + tail_rows * hd, cache_data, head * hd * sizeof(__fp16)); + } + } +} + +void compute_conv_cache_initialize_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& new_data = get_input(node, 0, nodes, node_index_map); + auto& cache_buf = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; + auto* meta = get_conv_meta(cache_buf); + + const size_t ws = meta->window_size; + const size_t hd = meta->hidden_dim; + + __fp16* cache_data = get_conv_data(cache_buf); + std::memset(cache_data, 0, ws * hd * sizeof(__fp16)); + meta->head = 0; + meta->count = 0; + + const size_t num_rows = new_data.total_size / hd; + if (num_rows == 0) return; + + const __fp16* src; + std::vector<__fp16> converted; + if (new_data.precision == Precision::FP16) { + src = new_data.data_as<__fp16>(); + } else if (new_data.precision == Precision::FP32) { + converted.resize(new_data.total_size); + Quantization::fp32_to_fp16(new_data.data_as(), converted.data(), new_data.total_size); + src = converted.data(); + } else { + converted.resize(new_data.total_size); + Quantization::int8_to_fp16(new_data.data_as(), converted.data(), new_data.total_size); + src = converted.data(); + } + + const size_t copy_rows = std::min(num_rows, ws); + const size_t start_row = num_rows - copy_rows; + std::memcpy(cache_data, src + start_row * hd, copy_rows * hd * sizeof(__fp16)); + meta->head = copy_rows % ws; + meta->count = copy_rows; +} + +void compute_recurrent_cache_state_node( + GraphNode& node, + const nodes_vector&, + const node_index_map_t&) { + + if (node.output_buffer.get_data()) return; + node.output_buffer.allocate(); + std::memset(node.output_buffer.get_data(), 0, node.output_buffer.byte_size); +} + +void compute_recurrent_cache_write_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& src = get_input(node, 0, nodes, node_index_map); + auto& cache_buf = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; + std::memcpy(cache_buf.get_data(), src.get_data(), src.byte_size); +} + +void CactusGraph::steal_cache_buffer(size_t dst_node, CactusGraph& src, size_t src_node) { + auto& dst = nodes_[node_index_map_.at(dst_node)]; + auto& s = src.nodes_[src.node_index_map_.at(src_node)]; + // Buffer carries its own runtime precision (may be fp16); only op_type is invariant pre-move. + assert(dst->op_type == s->op_type); + auto shape = s->output_buffer.shape; + auto prec = s->output_buffer.precision; + dst->output_buffer = std::move(s->output_buffer); + s->output_buffer = BufferDesc(shape, prec); +} + +namespace { + +struct PaddedAppendBackup { + uint64_t overshoot; + uint64_t keep_sink; + uint64_t kept_padded; + uint64_t ring_rows; +}; + +size_t ring_row_for_pos(size_t pos, size_t sink, size_t ring_capacity) { + if (pos < ring_capacity) return pos; + size_t recent = ring_capacity > sink ? ring_capacity - sink : 1; + return sink + (pos - sink) % recent; +} + +struct CacheRowRegion { + size_t offset; + size_t row_bytes; +}; + +std::vector cache_row_regions(const BufferDesc& buf, const CacheMetadata* meta) { + const auto* base = static_cast(buf.get_data()); + const size_t stride = meta->num_kv_heads * meta->head_dim; + if (buf.precision == Precision::FP16) { + return {{static_cast(reinterpret_cast(get_fp16_data(buf)) - base), stride * sizeof(__fp16)}}; + } + const size_t num_groups = (meta->head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + return { + {static_cast(reinterpret_cast(get_int8_data(buf)) - base), stride}, + {static_cast(reinterpret_cast(get_scales(buf, meta->max_seq_len, meta->num_kv_heads, meta->head_dim)) - base), + meta->num_kv_heads * num_groups * sizeof(float)}, + }; +} + +} // namespace + +std::vector CactusGraph::snapshot_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens) const { + const auto& node = *nodes_[node_index_map_.at(node_id)]; + const auto& buf = node.output_buffer; + if (node.op_type != OpType::KV_CACHE_STATE || !buf.get_data() || pad_tokens == 0) return {}; + const size_t window = node.params.window_size; + const size_t ceiling = node.params.max_cache_seq_len; + if (window == 0 || window >= ceiling) return {}; + const auto* meta = get_meta(buf); + const size_t len0 = meta->current_seq_len; + const size_t appended = real_tokens + pad_tokens; + if (kv_cache_resident()) { + const size_t ring_window = meta->max_seq_len - meta->sink_size - 1; + if (ring_window <= meta->sink_size || appended >= ring_window - meta->sink_size) { + throw std::runtime_error("padded cache append larger than the attention window is not supported"); + } + std::vector backup(sizeof(PaddedAppendBackup)); + auto* header = reinterpret_cast(backup.data()); + header->overshoot = 0; + header->keep_sink = meta->sink_size; + header->kept_padded = ring_window; + header->ring_rows = pad_tokens; + const auto* base = static_cast(buf.get_data()); + for (const auto& region : cache_row_regions(buf, meta)) { + for (size_t t = 0; t < pad_tokens; ++t) { + size_t row = ring_row_for_pos(len0 + real_tokens + t, meta->sink_size, ring_window); + const uint8_t* src = base + region.offset + row * region.row_bytes; + backup.insert(backup.end(), src, src + region.row_bytes); + } + } + return backup; + } + if (len0 == 0 || len0 + appended <= window) return {}; + const size_t keep_sink = std::min({static_cast(meta->sink_size), len0, window}); + const size_t tail_capacity = window - keep_sink; + if (appended >= tail_capacity) { + throw std::runtime_error("padded cache append larger than the attention window is not supported"); + } + auto kept_after = [&](size_t n) { return std::min(tail_capacity - n, len0 - keep_sink); }; + const size_t kept_real = kept_after(real_tokens); + const size_t kept_padded = kept_after(appended); + const size_t overshoot = kept_real - kept_padded; + if (overshoot == 0) return {}; + + const size_t first_saved_row = len0 - kept_real; + std::vector backup(sizeof(PaddedAppendBackup)); + auto* header = reinterpret_cast(backup.data()); + header->overshoot = overshoot; + header->keep_sink = keep_sink; + header->kept_padded = kept_padded; + const auto* base = static_cast(buf.get_data()); + for (const auto& region : cache_row_regions(buf, meta)) { + const uint8_t* rows = base + region.offset + first_saved_row * region.row_bytes; + backup.insert(backup.end(), rows, rows + overshoot * region.row_bytes); + } + return backup; +} + +void CactusGraph::rollback_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens, + const std::vector& backup) { + auto& node = *nodes_[node_index_map_.at(node_id)]; + auto& buf = node.output_buffer; + if (node.op_type != OpType::KV_CACHE_STATE || !buf.get_data() || pad_tokens == 0) return; + auto* meta = get_meta(buf); + if (backup.empty()) { + meta->current_seq_len = meta->current_seq_len >= pad_tokens ? meta->current_seq_len - pad_tokens : 0; + return; + } + const auto* header = reinterpret_cast(backup.data()); + if (header->ring_rows > 0) { + const size_t ring_window = header->kept_padded; + const size_t sink = header->keep_sink; + const size_t len_after = meta->current_seq_len; + const size_t len0 = len_after >= real_tokens + pad_tokens ? len_after - real_tokens - pad_tokens : 0; + auto* base = static_cast(buf.get_data()); + const uint8_t* saved = backup.data() + sizeof(PaddedAppendBackup); + for (const auto& region : cache_row_regions(buf, meta)) { + for (size_t t = 0; t < header->ring_rows; ++t) { + size_t row = ring_row_for_pos(len0 + real_tokens + t, sink, ring_window); + std::memcpy(base + region.offset + row * region.row_bytes, saved, region.row_bytes); + saved += region.row_bytes; + } + } + meta->current_seq_len = len0 + real_tokens; + return; + } + const size_t overshoot = header->overshoot; + const size_t keep_sink = header->keep_sink; + const size_t kept_padded = header->kept_padded; + const size_t move_rows = kept_padded + real_tokens; + auto* base = static_cast(buf.get_data()); + const uint8_t* saved = backup.data() + sizeof(PaddedAppendBackup); + for (const auto& region : cache_row_regions(buf, meta)) { + uint8_t* rows = base + region.offset; + std::memmove(rows + (keep_sink + overshoot) * region.row_bytes, + rows + keep_sink * region.row_bytes, move_rows * region.row_bytes); + std::memcpy(rows + keep_sink * region.row_bytes, saved, overshoot * region.row_bytes); + saved += overshoot * region.row_bytes; + } + meta->current_seq_len = keep_sink + overshoot + kept_padded + real_tokens; +} + +void CactusGraph::shrink_cache_buffer(size_t node_id, size_t new_capacity) { + auto& buf = nodes_[node_index_map_.at(node_id)]->output_buffer; + if (!buf.get_data()) return; + auto* meta = get_meta(buf); + size_t target = std::max(new_capacity, meta->current_seq_len); + if (target >= meta->max_seq_len) return; + resize_cache_buffer(buf, target); +} + diff --git a/cactus-graph/src/ops_conv.cpp b/cactus-graph/src/ops_conv.cpp new file mode 100644 index 000000000..20cf09030 --- /dev/null +++ b/cactus-graph/src/ops_conv.cpp @@ -0,0 +1,883 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include +#include +#include +#include +#include +#include + +std::vector<__fp16> dequantize_int8_weights_to_fp16( + const BufferDesc& W, + size_t rows, + size_t cols, + const char* op_name) { + const int8_t* W_int8 = W.data_as(); + std::vector<__fp16> W_fp16(rows * cols); + + if (W.group_size == 0 || W.activation_scales_data == nullptr) { + for (size_t i = 0; i < W_fp16.size(); ++i) { + W_fp16[i] = static_cast<__fp16>(W_int8[i]); + } + return W_fp16; + } + + const size_t group_size = W.group_size; + if ((cols % group_size) != 0) { + throw std::runtime_error(std::string(op_name) + " grouped INT8 weight columns must be divisible by group size"); + } + + const __fp16* scales = reinterpret_cast(W.activation_scales_data); + const size_t num_groups = cols / group_size; + for (size_t row = 0; row < rows; ++row) { + for (size_t col = 0; col < cols; ++col) { + const size_t idx = row * cols + col; + const size_t group_idx = col / group_size; + const float scale = static_cast(scales[row * num_groups + group_idx]); + W_fp16[idx] = static_cast<__fp16>(W_int8[idx] * scale); + } + } + return W_fp16; +} + +void compute_conv1d_causal_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + auto& Y = node.output_buffer; + + if (X.shape.size() != 3) { + throw std::runtime_error("Causal conv requires 3D input [batch, seq_len, in_channels]"); + } + if (W.shape.size() != 3) { + throw std::runtime_error("Weight must be 3D"); + } + + const size_t N = X.shape[0]; + const size_t L = X.shape[1]; + const size_t C_in = X.shape[2]; + const size_t W0 = W.shape[0]; + const size_t W1 = W.shape[1]; + const size_t W2 = W.shape[2]; + const size_t dil = node.params.dilation; + if (dil < 1) throw std::runtime_error("dilation must be >= 1"); + + size_t M = 1; + size_t C_out = 0; + const bool standard_layout = (W1 == 1); + const bool transposed_layout = (W2 == 1); + if ((!standard_layout && !transposed_layout) || (W0 % C_in != 0)) { + throw std::runtime_error("Only depthwise causal convolution is supported currently"); + } + const size_t K = standard_layout ? W2 : W1; + M = W0 / C_in; + C_out = C_in * M; + + Y.shape = { N, L, C_out }; + Y.precision = X.precision; + + auto transpose_depthwise_weights_fp16 = [&](const __fp16* src) { + std::vector<__fp16> transposed(W0 * K); + for (size_t oc = 0; oc < W0; ++oc) { + for (size_t k = 0; k < K; ++k) { + transposed[oc * K + k] = src[(oc * W1 + k) * W2]; + } + } + return transposed; + }; + + if (W.precision == Precision::INT8) { + auto W_fp16 = dequantize_int8_weights_to_fp16(W, W0, K, "conv1d_causal"); + if (transposed_layout && !standard_layout) { + auto fixed = transpose_depthwise_weights_fp16(W_fp16.data()); + cactus_conv1d_causal_depthwise_f16( + X.data_as<__fp16>(), fixed.data(), Y.data_as<__fp16>(), + N, L, C_in, K, dil); + } else { + cactus_conv1d_causal_depthwise_f16( + X.data_as<__fp16>(), W_fp16.data(), Y.data_as<__fp16>(), + N, L, C_in, K, dil); + } + } else if (W.precision == Precision::FP16) { + if (transposed_layout && !standard_layout) { + auto fixed = transpose_depthwise_weights_fp16(W.data_as<__fp16>()); + cactus_conv1d_causal_depthwise_f16( + X.data_as<__fp16>(), fixed.data(), Y.data_as<__fp16>(), + N, L, C_in, K, dil); + } else { + cactus_conv1d_causal_depthwise_f16( + X.data_as<__fp16>(), W.data_as<__fp16>(), Y.data_as<__fp16>(), + N, L, C_in, K, dil); + } + } else { + throw std::runtime_error("Conv requires FP16 weights"); + } +} + +void compute_conv1d_k3_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + auto& Y = node.output_buffer; + + if (X.shape.size() != 3) + throw std::runtime_error("Conv requires 3D input [N, C_in, L]!"); + + if (W.shape.size() != 3) + throw std::runtime_error("Weight must be [C_out, C_in, 3]!"); + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t L = X.shape[2]; + + const size_t C_out = W.shape[0]; + const size_t K = W.shape[2]; + const size_t stride = node.params.stride; + + if (K != 3) + throw std::runtime_error("Conv1d_k3 only supports K=3!"); + + size_t L_out = ((L - 1) / stride) + 1; + Y.shape = { N, C_out, L_out }; + Y.precision = X.precision; + + if (X.precision != Precision::FP16) { + throw std::runtime_error("Conv1d_k3 only supports FP16 activations"); + } + + if (W.precision == Precision::INT8) { + auto W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in * K, "conv1d_k3"); + cactus_conv1d_f16_k3( + X.data_as<__fp16>(), + W_fp16.data(), + Y.data_as<__fp16>(), + N, L, C_in, C_out, stride + ); + } else if (W.precision == Precision::FP16) { + cactus_conv1d_f16_k3( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + Y.data_as<__fp16>(), + N, L, C_in, C_out, stride + ); + } else { + throw std::runtime_error("Conv requires FP16 weights"); + } +} + +void compute_conv1d_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + + auto& Y = node.output_buffer; + + if (X.shape.size() != 3) { + throw std::runtime_error("conv1d expects input [N, C_in, L]"); + } + if (W.shape.size() != 3) { + throw std::runtime_error("conv1d weight must be [C_out, C_in, K]"); + } + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t L = X.shape[2]; + const size_t C_out = W.shape[0]; + const size_t K = W.shape[2]; + const size_t stride = node.params.stride; + + if (W.shape[1] != C_in) { + throw std::runtime_error("conv1d weight C_in mismatch"); + } + + if (X.precision != Precision::FP16) { + throw std::runtime_error("Conv1d only supports FP16 activations"); + } + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->total_size != C_out) { + throw std::runtime_error("conv1d bias size mismatch"); + } + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C_out); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C_out); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv1d bias only supports FP16/FP32"); + } + } + + std::vector<__fp16> W_fp16; + const __fp16* W_ptr = nullptr; + if (W.precision == Precision::FP16) { + W_ptr = W.data_as<__fp16>(); + } else if (W.precision == Precision::INT8) { + W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in * K, "conv1d"); + W_ptr = W_fp16.data(); + } else { + throw std::runtime_error("Conv1d only supports FP16/INT8 weights"); + } + + cactus_conv1d_f16(X.data_as<__fp16>(), W_ptr, bias_ptr, + Y.data_as<__fp16>(), N, L, C_in, C_out, K, stride); +} + +void compute_conv1d_same_depthwise_k9_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + auto& Y = node.output_buffer; + + if (X.shape.size() != 3) { + throw std::runtime_error("conv1d_same_depthwise_k9 expects input [N, L, C]"); + } + if (X.precision != Precision::FP16) { + throw std::runtime_error("conv1d_same_depthwise_k9 only supports FP16 activations"); + } + + const size_t N = X.shape[0]; + const size_t L = X.shape[1]; + const size_t C = X.shape[2]; + const size_t K = 9; + + if (W.shape.size() == 2) { + if (W.shape[0] != C || W.shape[1] != K) { + throw std::runtime_error("conv1d_same_depthwise_k9 weight must be [C, 9]"); + } + } else if (W.shape.size() == 3) { + if (W.shape[0] != C || W.shape[1] != 1 || W.shape[2] != K) { + throw std::runtime_error("conv1d_same_depthwise_k9 weight must be [C, 1, 9]"); + } + } else { + throw std::runtime_error("conv1d_same_depthwise_k9 weight must be rank 2 or 3"); + } + + Y.shape = {N, L, C}; + Y.precision = Precision::FP16; + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->total_size != C) { + throw std::runtime_error("conv1d_same_depthwise_k9 bias size mismatch"); + } + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv1d_same_depthwise_k9 bias only supports FP16/FP32"); + } + } + + if (W.precision == Precision::FP16) { + cactus_conv1d_same_depthwise_f16_k9( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + bias_ptr, + Y.data_as<__fp16>(), + N, L, C + ); + return; + } + + if (W.precision == Precision::INT8) { + auto W_fp16 = dequantize_int8_weights_to_fp16(W, C, K, "conv1d_same_depthwise_k9"); + cactus_conv1d_same_depthwise_f16_k9( + X.data_as<__fp16>(), + W_fp16.data(), + bias_ptr, + Y.data_as<__fp16>(), + N, L, C + ); + return; + } + + throw std::runtime_error("Conv requires FP16/INT8 weights"); +} + +void compute_conv2d_k3s2p1_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + auto& Y = node.output_buffer; + + if (X.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s2p1 expects input [N, C_in, H, W]"); + } + if (X.precision != Precision::FP16) { + throw std::runtime_error("conv2d_k3s2p1 only supports FP16 activations"); + } + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t H = X.shape[2]; + const size_t W_in = X.shape[3]; + const size_t C_out = Y.shape[1]; + + if (H == 0 || W_in == 0) { + throw std::runtime_error("conv2d_k3s2p1 input spatial dimensions must be > 0"); + } + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C_out); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C_out); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv2d_k3s2p1 bias only supports FP16/FP32"); + } + } + + if (W.precision == Precision::FP16) { + if (W.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s2p1 FP16 weight must be [C_out, C_in, 3, 3]"); + } + cactus_conv2d_f16_k3s2p1_nchw( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + bias_ptr, + Y.data_as<__fp16>(), + N, C_in, H, W_in, C_out + ); + return; + } + + if (W.precision == Precision::INT8) { + if (W.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s2p1 INT8 weight must be [C_out, C_in, 3, 3]"); + } + auto W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in * 3 * 3, "conv2d_k3s2p1"); + cactus_conv2d_f16_k3s2p1_nchw( + X.data_as<__fp16>(), + W_fp16.data(), + bias_ptr, + Y.data_as<__fp16>(), + N, C_in, H, W_in, C_out + ); + return; + } + + throw std::runtime_error("Conv requires FP16/INT8 weights"); +} + +void compute_conv2d_depthwise_k3s2p1_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + auto& Y = node.output_buffer; + + if (X.shape.size() != 4) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 expects input [N, C, H, W]"); + } + if (X.precision != Precision::FP16) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 only supports FP16 activations"); + } + + const size_t N = X.shape[0]; + const size_t C = X.shape[1]; + const size_t H = X.shape[2]; + const size_t W_in = X.shape[3]; + if (H == 0 || W_in == 0) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 input spatial dimensions must be > 0"); + } + + if (W.shape.size() == 3) { + if (W.shape[0] != C || W.shape[1] != 3 || W.shape[2] != 3) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 weight must be [C, 3, 3]"); + } + } else if (W.shape.size() == 4) { + if (W.shape[0] != C || W.shape[1] != 1 || W.shape[2] != 3 || W.shape[3] != 3) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 weight must be [C, 1, 3, 3]"); + } + } else { + throw std::runtime_error("conv2d_depthwise_k3s2p1 weight must be rank 3 or 4"); + } + + const size_t H_out = (H - 1) / 2 + 1; + const size_t W_out = (W_in - 1) / 2 + 1; + Y.shape = {N, C, H_out, W_out}; + Y.precision = Precision::FP16; + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->total_size != C) { + throw std::runtime_error("conv2d_depthwise_k3s2p1 bias size mismatch"); + } + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv2d_depthwise_k3s2p1 bias only supports FP16/FP32"); + } + } + + if (W.precision == Precision::FP16) { + cactus_conv2d_depthwise_f16_k3s2p1_nchw( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + bias_ptr, + Y.data_as<__fp16>(), + N, C, H, W_in + ); + return; + } + + if (W.precision == Precision::INT8) { + auto W_fp16 = dequantize_int8_weights_to_fp16(W, C, 3 * 3, "conv2d_depthwise_k3s2p1"); + cactus_conv2d_depthwise_f16_k3s2p1_nchw( + X.data_as<__fp16>(), + W_fp16.data(), + bias_ptr, + Y.data_as<__fp16>(), + N, C, H, W_in + ); + return; + } + + throw std::runtime_error("Conv requires FP16/INT8 weights"); +} + +void compute_conv2d_pointwise_1x1_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + auto& Y = node.output_buffer; + + if (X.shape.size() != 4) { + throw std::runtime_error("conv2d_pointwise_1x1 expects input [N, C_in, H, W]"); + } + if (X.precision != Precision::FP16) { + throw std::runtime_error("conv2d_pointwise_1x1 only supports FP16 activations"); + } + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t H = X.shape[2]; + const size_t W_in = X.shape[3]; + if (H == 0 || W_in == 0) { + throw std::runtime_error("conv2d_pointwise_1x1 input spatial dimensions must be > 0"); + } + + size_t C_out = 0; + if (W.shape.size() == 2) { + C_out = W.shape[0]; + if (W.shape[1] != C_in) { + throw std::runtime_error("conv2d_pointwise_1x1 weight must be [C_out, C_in]"); + } + } else if (W.shape.size() == 4) { + C_out = W.shape[0]; + if (W.shape[1] != C_in || W.shape[2] != 1 || W.shape[3] != 1) { + throw std::runtime_error("conv2d_pointwise_1x1 weight must be [C_out, C_in, 1, 1]"); + } + } else { + throw std::runtime_error("conv2d_pointwise_1x1 weight must be rank 2 or 4"); + } + + Y.shape = {N, C_out, H, W_in}; + Y.precision = Precision::FP16; + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->total_size != C_out) { + throw std::runtime_error("conv2d_pointwise_1x1 bias size mismatch"); + } + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C_out); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C_out); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv2d_pointwise_1x1 bias only supports FP16/FP32"); + } + } + + if (W.precision == Precision::FP16) { + cactus_conv2d_pointwise_f16_1x1_nchw_gemm( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + bias_ptr, + Y.data_as<__fp16>(), + N, C_in, H, W_in, C_out + ); + return; + } + + if (W.precision == Precision::INT8) { + auto W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in, "conv2d_pointwise_1x1"); + cactus_conv2d_pointwise_f16_1x1_nchw_gemm( + X.data_as<__fp16>(), + W_fp16.data(), + bias_ptr, + Y.data_as<__fp16>(), + N, C_in, H, W_in, C_out + ); + return; + } + + throw std::runtime_error("Conv requires FP16/INT8 weights"); +} + +void compute_conv1d_pointwise_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + auto& Y = node.output_buffer; + + if (X.shape.size() != 3) { + throw std::runtime_error("conv1d_pointwise expects input [N, L, C_in]"); + } + if (X.precision != Precision::FP16) { + throw std::runtime_error("conv1d_pointwise only supports FP16 activations"); + } + + const size_t N = X.shape[0]; + const size_t L = X.shape[1]; + const size_t C_in = X.shape[2]; + + size_t C_out = 0; + if (W.shape.size() == 2) { + C_out = W.shape[0]; + if (W.shape[1] != C_in) { + throw std::runtime_error("conv1d_pointwise weight must be [C_out, C_in]"); + } + } else if (W.shape.size() == 3) { + C_out = W.shape[0]; + if (W.shape[1] != C_in || W.shape[2] != 1) { + throw std::runtime_error("conv1d_pointwise weight must be [C_out, C_in, 1]"); + } + } else { + throw std::runtime_error("conv1d_pointwise weight must be rank 2 or 3"); + } + + Y.shape = {N, L, C_out}; + Y.precision = Precision::FP16; + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->total_size != C_out) { + throw std::runtime_error("conv1d_pointwise bias size mismatch"); + } + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C_out); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C_out); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv1d_pointwise bias only supports FP16/FP32"); + } + } + + if (W.precision == Precision::FP16) { + cactus_conv1d_pointwise_f16_gemm( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + bias_ptr, + Y.data_as<__fp16>(), + N, L, C_in, C_out + ); + return; + } + + if (W.precision == Precision::INT8) { + auto W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in, "conv1d_pointwise"); + cactus_conv1d_pointwise_f16_gemm( + X.data_as<__fp16>(), + W_fp16.data(), + bias_ptr, + Y.data_as<__fp16>(), + N, L, C_in, C_out + ); + return; + } + + throw std::runtime_error("Conv requires FP16/INT8 weights"); +} + +void compute_batchnorm_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + if (node.input_ids.size() != 5) { + throw std::runtime_error("BatchNorm expects 5 inputs: input, weight, bias, running_mean, running_var"); + } + + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const auto& B = get_input(node, 2, nodes, node_index_map); + const auto& RM = get_input(node, 3, nodes, node_index_map); + const auto& RV = get_input(node, 4, nodes, node_index_map); + auto& Y = node.output_buffer; + + if (X.shape.empty()) { + throw std::runtime_error("BatchNorm expects non-scalar input"); + } + + int axis = node.params.axis; + if (axis < 0) axis += static_cast(X.shape.size()); + if (axis < 0 || static_cast(axis) >= X.shape.size()) { + throw std::runtime_error("BatchNorm axis out of range"); + } + + const size_t C = X.shape[static_cast(axis)]; + if (W.total_size != C || B.total_size != C || RM.total_size != C || RV.total_size != C) { + throw std::runtime_error("BatchNorm parameter size mismatch"); + } + + auto load_1d_float = [C](const BufferDesc& buf, const char* name) -> std::vector { + if (buf.total_size != C) { + throw std::runtime_error(std::string("BatchNorm parameter size mismatch for ") + name); + } + std::vector out(C); + if (buf.precision == Precision::FP16) { + const __fp16* p = buf.data_as<__fp16>(); + for (size_t i = 0; i < C; ++i) out[i] = static_cast(p[i]); + } else if (buf.precision == Precision::FP32) { + std::memcpy(out.data(), buf.data_as(), C * sizeof(float)); + } else { + throw std::runtime_error(std::string("BatchNorm parameter ") + name + " must be FP16 or FP32"); + } + return out; + }; + + const std::vector gamma = load_1d_float(W, "weight"); + const std::vector beta = load_1d_float(B, "bias"); + const std::vector mean = load_1d_float(RM, "running_mean"); + const std::vector var = load_1d_float(RV, "running_var"); + + size_t outer = 1; + for (int i = 0; i < axis; ++i) { + outer *= X.shape[static_cast(i)]; + } + size_t inner = 1; + for (size_t i = static_cast(axis) + 1; i < X.shape.size(); ++i) { + inner *= X.shape[i]; + } + + Y.shape = X.shape; + Y.precision = X.precision; + + if (X.precision == Precision::FP16) { + cactus_batchnorm_f16( + X.data_as<__fp16>(), + gamma.data(), + beta.data(), + mean.data(), + var.data(), + Y.data_as<__fp16>(), + outer, + C, + inner, + node.params.epsilon + ); + return; + } + + if (X.precision == Precision::FP32) { + cactus_batchnorm_f32( + X.data_as(), + gamma.data(), + beta.data(), + mean.data(), + var.data(), + Y.data_as(), + outer, + C, + inner, + node.params.epsilon + ); + return; + } + + throw std::runtime_error("BatchNorm only supports FP16/FP32 activations"); +} + +void compute_stft_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + auto& Y = node.output_buffer; + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t L = X.shape[2]; + const size_t C_out = W.shape[0]; + const size_t K = W.shape[2]; + const size_t stride = node.params.stride; + const size_t num_fft_bins = node.params.num_fft_bins; + + if (X.precision != Precision::FP16 || W.precision != Precision::FP16) { + throw std::runtime_error("stft only supports FP16"); + } + + cactus_stft_f16(X.data_as<__fp16>(), W.data_as<__fp16>(), + Y.data_as<__fp16>(), N, L, C_in, C_out, K, stride, num_fft_bins); +} + +void compute_conv1d_k7s3_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + + auto& Y = node.output_buffer; + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t L = X.shape[2]; + + if (W.shape.size() != 3) throw std::runtime_error("Weight must be 3D"); + const size_t C_in_W = W.shape[0]; + const size_t K = W.shape[1]; + const size_t C_out = W.shape[2]; + const size_t stride = node.params.stride; + + if (C_in != C_in_W) throw std::runtime_error("Channel mismatch in conv1d_k7s3"); + if (K != 7 || stride != 3) throw std::runtime_error("conv1d_k7s3 requires K=7, stride=3"); + + if (X.precision != Precision::FP16 || W.precision != Precision::FP16) { + throw std::runtime_error("Conv1d specialized only supports FP16"); + } + + size_t L_out = (L < 7) ? 0 : (L - 7) / 3 + 1; + Y.shape = {N, C_out, L_out}; + Y.precision = Precision::FP16; + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->total_size != C_out) { + throw std::runtime_error("conv1d_k7s3 bias size mismatch"); + } + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C_out); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C_out); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv1d_k7s3 bias only supports FP16/FP32"); + } + } + + cactus_conv1d_f16_k7s3_oc8( + X.data_as<__fp16>(), + W.data_as<__fp16>(), + bias_ptr, + Y.data_as<__fp16>(), + N, L, C_in, C_out + ); +} + +void compute_maxpool1d_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + + size_t batch_size = input.shape[0]; + size_t channels = input.shape[1]; + size_t input_length = input.shape[2]; + size_t kernel_size = node.params.kernel_size; + size_t stride = node.params.stride; + + cactus_maxpool1d_f16( + input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + batch_size, channels, input_length, + kernel_size, stride); +} + +void compute_conv2d_k3s1p1_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + const BufferDesc* B = nullptr; + if (node.input_ids.size() >= 3) { + B = &get_input(node, 2, nodes, node_index_map); + } + auto& Y = node.output_buffer; + + if (X.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s1p1 expects input [N, C_in, H, W]"); + } + if (X.precision != Precision::FP16) { + throw std::runtime_error("conv2d_k3s1p1 only supports FP16 activations"); + } + + const size_t N = X.shape[0]; + const size_t C_in = X.shape[1]; + const size_t H = X.shape[2]; + const size_t W_in = X.shape[3]; + const size_t C_out = Y.shape[1]; + + if (H == 0 || W_in == 0) { + throw std::runtime_error("conv2d_k3s1p1 input spatial dimensions must be > 0"); + } + + const __fp16* bias_ptr = nullptr; + std::vector<__fp16> bias_fp16; + if (B) { + if (B->precision == Precision::FP16) { + bias_ptr = B->data_as<__fp16>(); + } else if (B->precision == Precision::FP32) { + bias_fp16.resize(C_out); + cactus_fp32_to_fp16(B->data_as(), bias_fp16.data(), C_out); + bias_ptr = bias_fp16.data(); + } else { + throw std::runtime_error("conv2d_k3s1p1 bias only supports FP16/FP32"); + } + } + + if (W.precision == Precision::FP16) { + if (W.shape.size() != 4) { + throw std::runtime_error("conv2d_k3s1p1 FP16 weight must be [C_out, C_in, 3, 3]"); + } + cactus_conv2d_f16_k3s1p1_nchw( + X.data_as<__fp16>(), W.data_as<__fp16>(), bias_ptr, + Y.data_as<__fp16>(), + N, C_in, H, W_in, C_out); + return; + } + + throw std::runtime_error("conv2d_k3s1p1 only supports FP16 weights"); +} diff --git a/cactus-graph/src/ops_dsp.cpp b/cactus-graph/src/ops_dsp.cpp new file mode 100644 index 000000000..a1c2a00d1 --- /dev/null +++ b/cactus-graph/src/ops_dsp.cpp @@ -0,0 +1,155 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include +#include +#include + +void compute_rfft_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& input = get_input(node, 0, nodes, node_index_map); + size_t n = input.total_size; + + if (input.precision == Precision::FP32) { + cactus_rfft_f32_1d( + input.data_as(), + node.output_buffer.data_as(), + n, "backward"); + } else if (input.precision == Precision::FP16) { + std::vector in_f32(n); + Quantization::fp16_to_fp32(input.data_as<__fp16>(), in_f32.data(), n); + std::vector out_f32((n / 2 + 1) * 2); + cactus_rfft_f32_1d(in_f32.data(), out_f32.data(), n, "backward"); + Quantization::fp32_to_fp16(out_f32.data(), node.output_buffer.data_as<__fp16>(), out_f32.size()); + } else { + throw std::runtime_error("RFFT requires FP32 or FP16 input"); + } +} + +void compute_irfft_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& input = get_input(node, 0, nodes, node_index_map); + size_t n = node.output_buffer.total_size; + + if (input.precision == Precision::FP32) { + cactus_irfft_f32_1d( + input.data_as(), + node.output_buffer.data_as(), + n, "backward"); + } else if (input.precision == Precision::FP16) { + size_t in_len = input.total_size; + std::vector in_f32(in_len); + Quantization::fp16_to_fp32(input.data_as<__fp16>(), in_f32.data(), in_len); + std::vector out_f32(n); + cactus_irfft_f32_1d(in_f32.data(), out_f32.data(), n, "backward"); + Quantization::fp32_to_fp16(out_f32.data(), node.output_buffer.data_as<__fp16>(), n); + } else { + throw std::runtime_error("IRFFT requires FP32 or FP16 input"); + } +} + +void compute_mel_filter_bank_node( + GraphNode& node, + const nodes_vector&, + const node_index_map_t&) { + + size_t num_freq_bins = node.output_buffer.shape[0]; + size_t num_mel = node.output_buffer.shape[1]; + + const char* norm_strs[] = {nullptr, "slaney"}; + const char* scale_strs[] = {"htk", "kaldi", "slaney"}; + + int norm_idx = node.params.mel_norm_type; + int scale_idx = node.params.mel_scale_type; + const char* norm = (norm_idx >= 0 && norm_idx <= 1) ? norm_strs[norm_idx] : "slaney"; + const char* scale = (scale_idx >= 0 && scale_idx <= 2) ? scale_strs[scale_idx] : "slaney"; + + std::vector filters(num_mel * num_freq_bins); + cactus_generate_mel_filter_bank( + filters.data(), + static_cast(num_freq_bins), + static_cast(num_mel), + node.params.min_frequency, + node.params.max_frequency, + static_cast(node.params.sampling_rate), + norm, scale, false); + + if (node.output_buffer.precision == Precision::FP32) { + std::memcpy(node.output_buffer.get_data(), filters.data(), filters.size() * sizeof(float)); + } else if (node.output_buffer.precision == Precision::FP16) { + Quantization::fp32_to_fp16(filters.data(), node.output_buffer.data_as<__fp16>(), filters.size()); + } +} + +void compute_spectrogram_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& waveform_buf = get_input(node, 0, nodes, node_index_map); + const auto& mel_buf = get_input(node, 1, nodes, node_index_map); + + size_t waveform_length = waveform_buf.total_size; + size_t frame_length = node.params.num_fft_bins; + size_t hop_length = node.params.hop_length; + size_t fft_length = node.params.stride; + float power = node.params.power; + bool center = node.params.center; + float mel_floor = node.params.mel_floor; + float dither_val = node.params.dither; + float preemph = node.params.preemphasis_coef; + bool rm_dc = node.params.remove_dc_offset; + + const char* pad_strs[] = {"reflect", "constant"}; + const char* pad_mode = (node.params.pad_mode_type >= 0 && node.params.pad_mode_type <= 1) + ? pad_strs[node.params.pad_mode_type] : "reflect"; + + const char* log_strs[] = {nullptr, "log", "log10", "dB"}; + const char* log_mel = (node.params.log_mel_mode >= 0 && node.params.log_mel_mode <= 3) + ? log_strs[node.params.log_mel_mode] : nullptr; + + std::vector waveform_f32(waveform_length); + if (waveform_buf.precision == Precision::FP32) { + std::memcpy(waveform_f32.data(), waveform_buf.get_data(), waveform_length * sizeof(float)); + } else if (waveform_buf.precision == Precision::FP16) { + Quantization::fp16_to_fp32(waveform_buf.data_as<__fp16>(), waveform_f32.data(), waveform_length); + } + + std::vector mel_f32(mel_buf.total_size); + if (mel_buf.precision == Precision::FP32) { + std::memcpy(mel_f32.data(), mel_buf.get_data(), mel_buf.total_size * sizeof(float)); + } else if (mel_buf.precision == Precision::FP16) { + Quantization::fp16_to_fp32(mel_buf.data_as<__fp16>(), mel_f32.data(), mel_buf.total_size); + } + + size_t num_frequency_bins = fft_length / 2 + 1; + size_t num_mel_bins = mel_buf.total_size / num_frequency_bins; + size_t pad_len = center ? frame_length / 2 : 0; + size_t padded_len = waveform_length + 2 * pad_len; + size_t num_frames = 1 + (padded_len - frame_length) / hop_length; + + std::vector output_f32(num_mel_bins * num_frames); + + cactus_compute_spectrogram_f32( + waveform_f32.data(), waveform_length, + nullptr, 0, + frame_length, hop_length, &fft_length, + output_f32.data(), power, + center, pad_mode, true, + dither_val, + preemph != 0.0f ? &preemph : nullptr, + mel_f32.data(), mel_f32.size(), + mel_floor, log_mel, + 1.0f, 1e-10f, nullptr, rm_dc); + + if (node.output_buffer.precision == Precision::FP32) { + std::memcpy(node.output_buffer.get_data(), output_f32.data(), output_f32.size() * sizeof(float)); + } else if (node.output_buffer.precision == Precision::FP16) { + Quantization::fp32_to_fp16(output_f32.data(), node.output_buffer.data_as<__fp16>(), output_f32.size()); + } +} diff --git a/cactus-graph/src/ops_image.cpp b/cactus-graph/src/ops_image.cpp new file mode 100644 index 000000000..81b645841 --- /dev/null +++ b/cactus-graph/src/ops_image.cpp @@ -0,0 +1,52 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include +#include + +void compute_image_preprocess_node( + GraphNode& node, + const nodes_vector& nodes, + const node_index_map_t& node_index_map) { + + const auto& input = get_input(node, 0, nodes, node_index_map); + + int src_w = static_cast(node.params.dst_width); + int src_h = static_cast(node.params.dst_height); + int dst_w = node.params.target_width; + int dst_h = node.params.target_height; + int ps = node.params.patch_size; + int ch = node.params.image_channels; + float rescale = node.params.rescale_factor; + + std::vector src_float(static_cast(src_w) * src_h * ch); + if (input.precision == Precision::FP32) { + std::memcpy(src_float.data(), input.get_data(), src_float.size() * sizeof(float)); + } else if (input.precision == Precision::INT8) { + const uint8_t* raw = static_cast(input.get_data()); + for (size_t i = 0; i < src_float.size(); i++) + src_float[i] = static_cast(raw[i]); + } else if (input.precision == Precision::FP16) { + Quantization::fp16_to_fp32(input.data_as<__fp16>(), src_float.data(), src_float.size()); + } + + std::vector resized; + const float* norm_input; + + if (dst_w != src_w || dst_h != src_h) { + resized.resize(static_cast(dst_w) * dst_h * ch); + cactus_image_resize_float(src_float.data(), src_w, src_h, + resized.data(), dst_w, dst_h, ch); + norm_input = resized.data(); + } else { + norm_input = src_float.data(); + } + + std::vector normalized(static_cast(dst_w) * dst_h * ch); + cactus_image_normalize(norm_input, normalized.data(), + dst_w, dst_h, ch, + rescale, node.params.image_mean, node.params.image_std); + + cactus_image_to_patches(normalized.data(), + node.output_buffer.data_as(), + dst_w, dst_h, ch, ps); +} diff --git a/cactus-graph/src/ops_math.cpp b/cactus-graph/src/ops_math.cpp new file mode 100644 index 000000000..e12661268 --- /dev/null +++ b/cactus-graph/src/ops_math.cpp @@ -0,0 +1,636 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include +#include +#include + +namespace Quantization { + void int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale) { + cactus_int8_to_fp32(src, dst, count, scale); + } + + void fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale) { + cactus_fp32_to_int8(src, dst, count, scale); + } + + void fp16_to_fp32(const __fp16* src, float* dst, size_t count) { + cactus_fp16_to_fp32(src, dst, count); + } + + void fp32_to_fp16(const float* src, __fp16* dst, size_t count) { + cactus_fp32_to_fp16(src, dst, count); + } + + void int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale) { + cactus_int8_to_fp16(src, dst, count, scale); + } + + void fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale) { + cactus_fp16_to_int8(src, dst, count, scale); + } +} + +static std::vector compute_strides(const std::vector& shape, const std::vector& target_shape) { + std::vector strides(target_shape.size()); + + size_t shape_offset = target_shape.size() - shape.size(); + + for (size_t i = 0; i < target_shape.size(); ++i) { + if (i < shape_offset) { + strides[i] = 0; + } else { + size_t dim_idx = i - shape_offset; + if (shape[dim_idx] == 1) { + strides[i] = 0; + } else { + strides[i] = 1; + for (size_t j = dim_idx + 1; j < shape.size(); ++j) { + strides[i] *= shape[j]; + } + } + } + } + + return strides; +} + +static size_t broadcast_linear_index(const std::vector& coords, const std::vector& strides) { + size_t index = 0; + for (size_t i = 0; i < coords.size(); ++i) { + index += coords[i] * strides[i]; + } + return index; +} + +static void increment_broadcast_coords(std::vector& coords, const std::vector& shape) { + for (int axis = static_cast(coords.size()) - 1; axis >= 0; --axis) { + size_t idx = static_cast(axis); + coords[idx]++; + if (coords[idx] < shape[idx]) { + break; + } + coords[idx] = 0; + } +} + +void dispatch_binary_op_f16(OpType op, const __fp16* lhs, const __fp16* rhs, __fp16* output, size_t count) { + switch (op) { + case OpType::ADD: + cactus_add_f16(lhs, rhs, output, count); + break; + case OpType::ADD_CLIPPED: + cactus_add_f16_clipped(lhs, rhs, output, count); + break; + case OpType::SUBTRACT: + cactus_subtract_f16(lhs, rhs, output, count); + break; + case OpType::MULTIPLY: + cactus_multiply_f16(lhs, rhs, output, count); + break; + case OpType::DIVIDE: + cactus_divide_f16(lhs, rhs, output, count); + break; + case OpType::NOT_EQUAL: + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output[i] = static_cast<__fp16>( + static_cast(lhs[i]) != static_cast(rhs[i]) ? 1.0f : 0.0f); + } + }); + break; + default: + break; + } +} + +static float apply_binary_op_f32(OpType op, float lhs, float rhs) { + switch (op) { + case OpType::ADD: + case OpType::ADD_CLIPPED: + return lhs + rhs; + case OpType::SUBTRACT: + return lhs - rhs; + case OpType::MULTIPLY: + return lhs * rhs; + case OpType::DIVIDE: + return lhs / rhs; + case OpType::NOT_EQUAL: + return lhs != rhs ? 1.0f : 0.0f; + default: + return lhs; + } +} + +void dispatch_binary_op_f32(OpType op, const float* lhs, const float* rhs, float* output, size_t count) { + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output[i] = apply_binary_op_f32(op, lhs[i], rhs[i]); + } + }); +} + +void dispatch_unary_op_f16(OpType op, const __fp16* input, __fp16* output, size_t count, float param) { + ScalarOpType scalar_op; + switch (op) { + case OpType::SCALAR_ADD: scalar_op = ScalarOpType::ADD; break; + case OpType::SCALAR_SUBTRACT: scalar_op = ScalarOpType::SUBTRACT; break; + case OpType::SCALAR_MULTIPLY: scalar_op = ScalarOpType::MULTIPLY; break; + case OpType::SCALAR_DIVIDE: scalar_op = ScalarOpType::DIVIDE; break; + case OpType::SCALAR_EXP: scalar_op = ScalarOpType::EXP; break; + case OpType::SCALAR_SQRT: scalar_op = ScalarOpType::SQRT; break; + case OpType::SCALAR_COS: scalar_op = ScalarOpType::COS; break; + case OpType::SCALAR_SIN: scalar_op = ScalarOpType::SIN; break; + case OpType::SCALAR_LOG: scalar_op = ScalarOpType::LOG; break; + case OpType::ABS: scalar_op = ScalarOpType::ABS; break; + case OpType::POW: scalar_op = ScalarOpType::POW; break; + default: return; + } + + cactus_scalar_op_f16(input, output, count, param, scalar_op); +} + +static float apply_scalar_op_f32(OpType op, float value, float param) { + switch (op) { + case OpType::SCALAR_ADD: return value + param; + case OpType::SCALAR_SUBTRACT: return value - param; + case OpType::SCALAR_MULTIPLY: return value * param; + case OpType::SCALAR_DIVIDE: return value / param; + case OpType::SCALAR_EXP: return std::exp(value); + case OpType::SCALAR_SQRT: return std::sqrt(value); + case OpType::SCALAR_COS: return std::cos(value); + case OpType::SCALAR_SIN: return std::sin(value); + case OpType::SCALAR_LOG: return std::log(value); + case OpType::ABS: return std::fabs(value); + case OpType::POW: return std::pow(value, param); + default: return value; + } +} + +void dispatch_unary_op_f32(OpType op, const float* input, float* output, size_t count, float param) { + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output[i] = apply_scalar_op_f32(op, input[i], param); + } + }); +} + +void compute_binary_op_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& lhs = get_input(node, 0, nodes, node_index_map); + const auto& rhs = get_input(node, 1, nodes, node_index_map); + + if (lhs.precision != rhs.precision) { + throw std::runtime_error("Binary operations require matching precision"); + } + if (lhs.precision != Precision::FP16 && lhs.precision != Precision::FP32) { + throw std::runtime_error("Binary operations only support FP16/FP32 precision"); + } + + if (node.params.broadcast_info.needs_broadcasting) { + std::vector lhs_strides = compute_strides(lhs.shape, node.params.broadcast_info.output_shape); + std::vector rhs_strides = compute_strides(rhs.shape, node.params.broadcast_info.output_shape); + const auto& output_shape = node.params.broadcast_info.output_shape; + size_t total_elements = 1; + for (size_t dim : output_shape) { + total_elements *= dim; + } + if (lhs.precision == Precision::FP32) { + const float* lhs_data = lhs.data_as(); + const float* rhs_data = rhs.data_as(); + float* out_data = node.output_buffer.data_as(); + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + std::vector coords(output_shape.size()); + size_t tmp = start_idx; + for (int axis = static_cast(output_shape.size()) - 1; axis >= 0; --axis) { + coords[static_cast(axis)] = tmp % output_shape[static_cast(axis)]; + tmp /= output_shape[static_cast(axis)]; + } + for (size_t linear_idx = start_idx; linear_idx < end_idx; ++linear_idx) { + size_t lhs_idx = broadcast_linear_index(coords, lhs_strides); + size_t rhs_idx = broadcast_linear_index(coords, rhs_strides); + out_data[linear_idx] = apply_binary_op_f32(node.op_type, lhs_data[lhs_idx], rhs_data[rhs_idx]); + increment_broadcast_coords(coords, output_shape); + } + }); + return; + } + + switch (node.op_type) { + case OpType::ADD: + case OpType::ADD_CLIPPED: + cactus_add_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + lhs_strides.data(), rhs_strides.data(), + node.params.broadcast_info.output_shape.data(), + node.params.broadcast_info.output_shape.size()); + break; + case OpType::SUBTRACT: + cactus_subtract_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + lhs_strides.data(), rhs_strides.data(), + node.params.broadcast_info.output_shape.data(), + node.params.broadcast_info.output_shape.size()); + break; + case OpType::MULTIPLY: + cactus_multiply_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + lhs_strides.data(), rhs_strides.data(), + node.params.broadcast_info.output_shape.data(), + node.params.broadcast_info.output_shape.size()); + break; + case OpType::DIVIDE: + cactus_divide_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + lhs_strides.data(), rhs_strides.data(), + node.params.broadcast_info.output_shape.data(), + node.params.broadcast_info.output_shape.size()); + break; + case OpType::NOT_EQUAL: { + const __fp16* lhs_data = lhs.data_as<__fp16>(); + const __fp16* rhs_data = rhs.data_as<__fp16>(); + __fp16* out_data = node.output_buffer.data_as<__fp16>(); + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + std::vector coords(output_shape.size()); + size_t tmp = start_idx; + for (int axis = static_cast(output_shape.size()) - 1; axis >= 0; --axis) { + coords[static_cast(axis)] = tmp % output_shape[static_cast(axis)]; + tmp /= output_shape[static_cast(axis)]; + } + for (size_t linear_idx = start_idx; linear_idx < end_idx; ++linear_idx) { + size_t lhs_idx = broadcast_linear_index(coords, lhs_strides); + size_t rhs_idx = broadcast_linear_index(coords, rhs_strides); + out_data[linear_idx] = static_cast<__fp16>( + static_cast(lhs_data[lhs_idx]) != static_cast(rhs_data[rhs_idx]) ? 1.0f : 0.0f); + increment_broadcast_coords(coords, output_shape); + } + }); + break; + } + default: break; + } + } else { + if (lhs.precision == Precision::FP16) { + dispatch_binary_op_f16(node.op_type, lhs.data_as<__fp16>(), + rhs.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + } else { + dispatch_binary_op_f32(node.op_type, lhs.data_as(), + rhs.data_as(), node.output_buffer.data_as(), + node.output_buffer.total_size); + } + } +} + +void compute_unary_op_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + + if (input.precision != Precision::FP16 && input.precision != Precision::FP32) { + throw std::runtime_error("Scalar operations only support FP16/FP32 precision"); + } + + if (node.op_type == OpType::SCALAR_NOT_EQUAL) { + const float scalar = node.params.scalar; + if (input.precision == Precision::FP16) { + const __fp16* input_data = input.data_as<__fp16>(); + __fp16* output_data = node.output_buffer.data_as<__fp16>(); + CactusThreading::parallel_for(node.output_buffer.total_size, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output_data[i] = static_cast<__fp16>( + static_cast(input_data[i]) != scalar ? 1.0f : 0.0f); + } + }); + } else { + const float* input_data = input.data_as(); + float* output_data = node.output_buffer.data_as(); + CactusThreading::parallel_for(node.output_buffer.total_size, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output_data[i] = input_data[i] != scalar ? 1.0f : 0.0f; + } + }); + } + return; + } + + if (input.precision == Precision::FP16) { + dispatch_unary_op_f16(node.op_type, input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size, node.params.scalar); + } else { + dispatch_unary_op_f32(node.op_type, input.data_as(), + node.output_buffer.data_as(), + node.output_buffer.total_size, node.params.scalar); + } +} + +void compute_activation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + + if (input.precision != Precision::FP16) { + throw std::runtime_error("Activation operations only support FP16 precision"); + } + + switch (node.op_type) { + case OpType::RELU: + cactus_relu_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + break; + case OpType::SILU: + cactus_silu_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + break; + case OpType::GELU: + cactus_gelu_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + break; + case OpType::GELU_ERF: + cactus_gelu_f16_erf(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + break; + case OpType::SIGMOID: + cactus_sigmoid_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + break; + case OpType::TANH: + cactus_tanh_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size); + break; + case OpType::LEAKY_RELU: + cactus_leaky_relu_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size, node.params.scalar); + break; + case OpType::CLAMP: + cactus_clamp_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size, + node.params.scalar, node.params.scale); + break; + default: + break; + } +} + +void compute_reduce_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + int axis = node.params.axis; + + if (input_buffer.precision != Precision::FP16 && input_buffer.precision != Precision::FP32) { + throw std::runtime_error("Reduction operations only support FP16/FP32 precision"); + } + + if (node.op_type == OpType::CUMSUM) { + if (axis < 0 || static_cast(axis) >= input_buffer.shape.size()) { + throw std::runtime_error("Cumsum axis out of range"); + } + + auto dims = AxisDims::from_shape(input_buffer.shape, static_cast(axis)); + if (input_buffer.precision == Precision::FP32) { + const float* input = input_buffer.data_as(); + float* output = node.output_buffer.data_as(); + for (size_t outer = 0; outer < dims.outer; ++outer) { + for (size_t inner = 0; inner < dims.inner; ++inner) { + float running = 0.0f; + for (size_t axis_index = 0; axis_index < dims.axis_size; ++axis_index) { + size_t flat_index = ((outer * dims.axis_size) + axis_index) * dims.inner + inner; + running += input[flat_index]; + output[flat_index] = running; + } + } + } + return; + } + + const __fp16* input = input_buffer.data_as<__fp16>(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + + for (size_t outer = 0; outer < dims.outer; ++outer) { + for (size_t inner = 0; inner < dims.inner; ++inner) { + float running = 0.0f; + for (size_t axis_index = 0; axis_index < dims.axis_size; ++axis_index) { + size_t flat_index = ((outer * dims.axis_size) + axis_index) * dims.inner + inner; + running += static_cast(input[flat_index]); + output[flat_index] = static_cast<__fp16>(running); + } + } + } + return; + } + + if (input_buffer.precision == Precision::FP32) { + const float* input = input_buffer.data_as(); + float* output = node.output_buffer.data_as(); + if (axis == -1) { + float result = 0.0f; + switch (node.op_type) { + case OpType::SUM: + case OpType::MEAN: + for (size_t i = 0; i < input_buffer.total_size; ++i) result += input[i]; + if (node.op_type == OpType::MEAN && input_buffer.total_size > 0) { + result /= static_cast(input_buffer.total_size); + } + output[0] = result; + break; + case OpType::VARIANCE: { + float mean = 0.0f; + for (size_t i = 0; i < input_buffer.total_size; ++i) mean += input[i]; + mean /= static_cast(std::max(input_buffer.total_size, 1)); + float variance = 0.0f; + for (size_t i = 0; i < input_buffer.total_size; ++i) { + float diff = input[i] - mean; + variance += diff * diff; + } + output[0] = variance / static_cast(std::max(input_buffer.total_size, 1)); + break; + } + case OpType::MIN: + result = input_buffer.total_size == 0 ? 0.0f : input[0]; + for (size_t i = 1; i < input_buffer.total_size; ++i) result = std::min(result, input[i]); + output[0] = result; + break; + case OpType::MAX: + result = input_buffer.total_size == 0 ? 0.0f : input[0]; + for (size_t i = 1; i < input_buffer.total_size; ++i) result = std::max(result, input[i]); + output[0] = result; + break; + default: + break; + } + return; + } + + auto dims = AxisDims::from_shape(input_buffer.shape, static_cast(axis)); + CactusThreading::parallel_for(dims.outer * dims.inner, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t out_idx = start_idx; out_idx < end_idx; ++out_idx) { + size_t outer = out_idx / dims.inner; + size_t inner = out_idx % dims.inner; + size_t base = outer * dims.axis_size * dims.inner + inner; + float result = 0.0f; + switch (node.op_type) { + case OpType::SUM: + case OpType::MEAN: + for (size_t axis_index = 0; axis_index < dims.axis_size; ++axis_index) { + result += input[base + axis_index * dims.inner]; + } + if (node.op_type == OpType::MEAN && dims.axis_size > 0) { + result /= static_cast(dims.axis_size); + } + output[out_idx] = result; + break; + case OpType::VARIANCE: { + float mean = 0.0f; + for (size_t axis_index = 0; axis_index < dims.axis_size; ++axis_index) { + mean += input[base + axis_index * dims.inner]; + } + mean /= static_cast(std::max(dims.axis_size, 1)); + float variance = 0.0f; + for (size_t axis_index = 0; axis_index < dims.axis_size; ++axis_index) { + float diff = input[base + axis_index * dims.inner] - mean; + variance += diff * diff; + } + output[out_idx] = variance / static_cast(std::max(dims.axis_size, 1)); + break; + } + case OpType::MIN: + result = dims.axis_size == 0 ? 0.0f : input[base]; + for (size_t axis_index = 1; axis_index < dims.axis_size; ++axis_index) { + result = std::min(result, input[base + axis_index * dims.inner]); + } + output[out_idx] = result; + break; + case OpType::MAX: + result = dims.axis_size == 0 ? 0.0f : input[base]; + for (size_t axis_index = 1; axis_index < dims.axis_size; ++axis_index) { + result = std::max(result, input[base + axis_index * dims.inner]); + } + output[out_idx] = result; + break; + default: + break; + } + } + }); + return; + } + + if (axis == -1) { + switch (node.op_type) { + case OpType::SUM: { + double result = cactus_sum_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); + node.output_buffer.data_as<__fp16>()[0] = static_cast<__fp16>(result); + break; + } + case OpType::MEAN: { + double result = cactus_mean_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); + node.output_buffer.data_as<__fp16>()[0] = static_cast<__fp16>(result); + break; + } + case OpType::VARIANCE: { + double result = cactus_variance_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); + node.output_buffer.data_as<__fp16>()[0] = static_cast<__fp16>(result); + break; + } + case OpType::MIN: { + __fp16 result = cactus_min_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); + node.output_buffer.data_as<__fp16>()[0] = result; + break; + } + case OpType::MAX: { + __fp16 result = cactus_max_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); + node.output_buffer.data_as<__fp16>()[0] = result; + break; + } + default: break; + } + } else { + auto dims = AxisDims::from_shape(input_buffer.shape, static_cast(axis)); + + switch (node.op_type) { + case OpType::SUM: + cactus_sum_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + dims.outer, dims.axis_size, dims.inner); + break; + case OpType::MEAN: + cactus_mean_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + dims.outer, dims.axis_size, dims.inner); + break; + case OpType::VARIANCE: + cactus_variance_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + dims.outer, dims.axis_size, dims.inner); + break; + case OpType::MIN: + cactus_min_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + dims.outer, dims.axis_size, dims.inner); + break; + case OpType::MAX: + cactus_max_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + dims.outer, dims.axis_size, dims.inner); + break; + default: break; + } + } +} + +void compute_reshape_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + + size_t input_total_elements = input_buffer.total_size; + size_t output_total_elements = node.output_buffer.total_size; + + if (input_total_elements != output_total_elements) { + throw std::runtime_error("Reshape operation: input elements (" + std::to_string(input_total_elements) + + ") must match output elements (" + std::to_string(output_total_elements) + ")"); + } + + if (cactus_metal_active_mode()) { + node.output_buffer.set_external(const_cast(input_buffer.get_data())); + } else { + std::memcpy(node.output_buffer.get_data(), input_buffer.get_data(), input_buffer.byte_size); + } +} + +void compute_precision_cast_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buf = get_input(node, 0, nodes, node_index_map); + + if (input_buf.precision == node.output_buffer.precision) { + if (cactus_metal_active_mode()) { + node.output_buffer.set_external(const_cast(input_buf.get_data())); + } else { + std::memcpy(node.output_buffer.get_data(), input_buf.get_data(), input_buf.byte_size); + } + return; + } + + size_t count = input_buf.total_size; + + if (input_buf.precision == Precision::INT8 && node.output_buffer.precision == Precision::FP32) { + Quantization::int8_to_fp32(input_buf.data_as(), node.output_buffer.data_as(), count, 1.0f); + } else if (input_buf.precision == Precision::FP32 && node.output_buffer.precision == Precision::INT8) { + Quantization::fp32_to_int8(input_buf.data_as(), node.output_buffer.data_as(), count, 1.0f); + } else if (input_buf.precision == Precision::FP16 && node.output_buffer.precision == Precision::FP32) { + Quantization::fp16_to_fp32(input_buf.data_as<__fp16>(), node.output_buffer.data_as(), count); + } else if (input_buf.precision == Precision::FP32 && node.output_buffer.precision == Precision::FP16) { + Quantization::fp32_to_fp16(input_buf.data_as(), node.output_buffer.data_as<__fp16>(), count); + } else if (input_buf.precision == Precision::INT8 && node.output_buffer.precision == Precision::FP16) { + Quantization::int8_to_fp16(input_buf.data_as(), node.output_buffer.data_as<__fp16>(), count, 1.0f); + } else if (input_buf.precision == Precision::FP16 && node.output_buffer.precision == Precision::INT8) { + Quantization::fp16_to_int8(input_buf.data_as<__fp16>(), node.output_buffer.data_as(), count, 1.0f); + } else { + throw std::runtime_error("Unsupported precision conversion from " + + std::to_string(static_cast(input_buf.precision)) + + " to " + std::to_string(static_cast(node.output_buffer.precision))); + } +} diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp new file mode 100644 index 000000000..3d192f52e --- /dev/null +++ b/cactus-graph/src/ops_nn.cpp @@ -0,0 +1,945 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include "metal_backend.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + thread_local std::vector<__fp16> transpose_buffer_fp16; + + void ensure_transpose_buffer_fp16(size_t required_size) { + if (transpose_buffer_fp16.size() < required_size) { + transpose_buffer_fp16.resize(required_size); + } + } +} + +void shrink_thread_local_buffers() { + std::vector<__fp16>().swap(transpose_buffer_fp16); +} + +void compute_matmul_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& lhs_buffer = get_input(node, 0, nodes, node_index_map); + const auto& rhs_buffer = get_input(node, 1, nodes, node_index_map); + const auto& lhs_shape = lhs_buffer.shape; + const auto& rhs_shape = rhs_buffer.shape; + + size_t M = lhs_shape[lhs_shape.size() - 2]; + size_t K = lhs_shape[lhs_shape.size() - 1]; + size_t N; + N = node.params.pretransposed_rhs ? + rhs_shape[rhs_shape.size() - 2] : rhs_shape[rhs_shape.size() - 1]; + + bool pretransposed_rhs = node.params.pretransposed_rhs; + + if (PrecisionTraits::is_cq(rhs_buffer.precision) && rhs_buffer.group_size > 0) { + if (lhs_buffer.precision != Precision::FP16) { + throw std::runtime_error("TQ matmul requires FP16 activations"); + } + + const __fp16* lhs = lhs_buffer.data_as<__fp16>(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + + CactusQuantMatrix mat = rhs_buffer.to_cq_matrix(); + if (rhs_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) + cactus_quant_orthogonal_matmul(&mat, lhs, static_cast(M), output); + else if (node.params.backend == ComputeBackend::METAL && cactus_metal_available()) + cactus_metal_quant_matmul(&mat, lhs, static_cast(M), output); + else + cactus_quant_matmul(&mat, lhs, static_cast(M), output); + } else { + if (lhs_buffer.precision != Precision::FP16) { + throw std::runtime_error("FP16 matmul requires FP16 activations (got precision " + std::to_string(static_cast(lhs_buffer.precision)) + ")"); + } + + const __fp16* lhs = lhs_buffer.data_as<__fp16>(); + const __fp16* rhs = rhs_buffer.data_as<__fp16>(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + + if (pretransposed_rhs) { + cactus_matmul_f16(lhs, rhs, output, M, K, N); + } else { + size_t transpose_size = rhs_shape[0] * rhs_shape[1]; + ensure_transpose_buffer_fp16(transpose_size); + + cactus_transpose_2d_f16(rhs, transpose_buffer_fp16.data(), + rhs_shape[0], rhs_shape[1], 0, rhs_shape[0]); + cactus_matmul_f16(lhs, transpose_buffer_fp16.data(), output, M, K, N); + } + } +} + +namespace { + thread_local std::vector<__fp16> moe_compact_hidden_buf; + thread_local std::vector<__fp16> moe_gate_buf; + thread_local std::vector<__fp16> moe_gate_pad_buf; + thread_local std::vector<__fp16> moe_up_buf; + thread_local std::vector<__fp16> moe_expert_out_buf; + thread_local std::vector moe_expert_offsets_buf; + thread_local std::vector moe_expert_tokens_buf; + thread_local std::vector moe_routing_denom_buf; + + void ensure_moe_buffers(size_t max_tokens, size_t hidden_dim, size_t intermediate_dim, + size_t num_experts, size_t top_k) { + size_t hidden_size = max_tokens * hidden_dim; + size_t inter_size = max_tokens * intermediate_dim; + if (moe_compact_hidden_buf.size() < hidden_size) moe_compact_hidden_buf.resize(hidden_size); + if (moe_gate_buf.size() < inter_size) moe_gate_buf.resize(inter_size); + if (moe_up_buf.size() < inter_size) moe_up_buf.resize(inter_size); + if (moe_expert_out_buf.size() < hidden_size) moe_expert_out_buf.resize(hidden_size); + size_t total_assignments = max_tokens * top_k; + if (moe_expert_offsets_buf.size() < num_experts + 1) moe_expert_offsets_buf.resize(num_experts + 1); + if (moe_expert_tokens_buf.size() < total_assignments) moe_expert_tokens_buf.resize(total_assignments); + if (moe_routing_denom_buf.size() < max_tokens) moe_routing_denom_buf.resize(max_tokens); + } + + void moe_matmul(const __fp16* lhs, + size_t M, + size_t K, + const BufferDesc& rhs_buffer, + __fp16* output, + size_t N) { + if (rhs_buffer.precision == Precision::FP16) { + cactus_matmul_f16(lhs, rhs_buffer.data_as<__fp16>(), output, M, K, N); + return; + } + + if (PrecisionTraits::is_cq(rhs_buffer.precision) && rhs_buffer.group_size > 0) { + CactusQuantMatrix mat = rhs_buffer.to_cq_matrix(); + if (rhs_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) + cactus_quant_orthogonal_matmul(&mat, lhs, static_cast(M), output); + else + cactus_quant_matmul(&mat, lhs, static_cast(M), output); + return; + } + + throw std::runtime_error("moe_layer only supports FP16 or TQ expert weights"); + } +} + +void compute_moe_layer_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const size_t num_experts = node.params.num_experts; + const size_t top_k = node.params.num_experts_per_tok; + const bool normalize_routing = node.params.normalize_routing; + const float eps = node.params.epsilon; + const float routed_scaling_factor = node.params.scalar; + const bool gated = node.params.moe_gated; + const Activation activation = node.params.activation; + const size_t base_inputs = gated ? (3 + 3 * num_experts) : (3 + 2 * num_experts); + bool has_per_expert_scale = node.input_ids.size() == base_inputs + 1; + if (node.input_ids.size() != base_inputs && node.input_ids.size() != base_inputs + 1) { + throw std::runtime_error("moe_layer expects " + std::to_string(base_inputs) + " or " + std::to_string(base_inputs + 1) + " inputs, got " + std::to_string(node.input_ids.size())); + } + + const auto& hidden_buffer = get_input(node, 0, nodes, node_index_map); + const auto& routing_buffer = get_input(node, 1, nodes, node_index_map); + const auto& topk_idx_buffer = get_input(node, 2, nodes, node_index_map); + + if (hidden_buffer.precision != Precision::FP16 || node.output_buffer.precision != Precision::FP16) { + throw std::runtime_error("moe_layer expects FP16 hidden/output"); + } + if (topk_idx_buffer.precision != Precision::FP32) { + throw std::runtime_error("moe_layer expects FP32 topk indices"); + } + + const __fp16* expert_scales_fp16 = nullptr; + if (has_per_expert_scale) { + const auto& scale_buffer = get_input(node, base_inputs, nodes, node_index_map); + if (scale_buffer.precision != Precision::FP16) { + throw std::runtime_error("moe_layer expects FP16 per_expert_scale"); + } + expert_scales_fp16 = scale_buffer.data_as<__fp16>(); + } + + const size_t token_count = hidden_buffer.shape[0]; + const size_t hidden_dim = hidden_buffer.shape[1]; + const size_t total_num_experts = routing_buffer.shape[1]; + + const auto& w1_0_buffer = get_input(node, 3, nodes, node_index_map); + const size_t expert_intermediate_dim = w1_0_buffer.shape[0]; + + const auto* hidden = hidden_buffer.data_as<__fp16>(); + auto* output = node.output_buffer.data_as<__fp16>(); + const auto* topk_idx = topk_idx_buffer.data_as(); + const auto* routing_fp16 = routing_buffer.precision == Precision::FP16 ? routing_buffer.data_as<__fp16>() : nullptr; + const auto* routing_fp32 = routing_buffer.precision == Precision::FP32 ? routing_buffer.data_as() : nullptr; + + auto routing_prob = [&](size_t tok, size_t exp) -> float { + const size_t offset = tok * total_num_experts + exp; + if (routing_fp16) return static_cast(routing_fp16[offset]); + return routing_fp32[offset]; + }; + + ensure_moe_buffers(token_count, hidden_dim, expert_intermediate_dim, num_experts, top_k); + + size_t* expert_offsets = moe_expert_offsets_buf.data(); + size_t* expert_tokens_flat = moe_expert_tokens_buf.data(); + + auto expert_index = [&](float raw_idx) -> size_t { + if (!std::isfinite(raw_idx) || raw_idx < 0.0f) { + throw std::runtime_error("moe_layer got invalid expert index"); + } + size_t idx = static_cast(raw_idx + 0.5f); + if (idx >= num_experts) { + throw std::runtime_error("moe_layer got expert index out of range"); + } + return idx; + }; + + std::memset(expert_offsets, 0, (num_experts + 1) * sizeof(size_t)); + for (size_t tok = 0; tok < token_count; ++tok) { + for (size_t k = 0; k < top_k; ++k) { + expert_offsets[expert_index(topk_idx[tok * top_k + k]) + 1]++; + } + } + + for (size_t e = 0; e < num_experts; ++e) { + expert_offsets[e + 1] += expert_offsets[e]; + } + + thread_local std::vector moe_write_cursors; + if (moe_write_cursors.size() < num_experts) moe_write_cursors.resize(num_experts); + std::memcpy(moe_write_cursors.data(), expert_offsets, num_experts * sizeof(size_t)); + + for (size_t tok = 0; tok < token_count; ++tok) { + for (size_t k = 0; k < top_k; ++k) { + size_t idx = expert_index(topk_idx[tok * top_k + k]); + expert_tokens_flat[moe_write_cursors[idx]++] = tok; + } + } + + float* routing_denom = moe_routing_denom_buf.data(); + if (normalize_routing) { + for (size_t tok = 0; tok < token_count; ++tok) { + float sum_probs = 0.0f; + for (size_t k = 0; k < top_k; ++k) { + size_t idx = expert_index(topk_idx[tok * top_k + k]); + sum_probs += routing_prob(tok, idx); + } + routing_denom[tok] = sum_probs + eps; + } + } + + std::memset(output, 0, token_count * hidden_dim * sizeof(__fp16)); + + for (size_t expert_idx = 0; expert_idx < num_experts; ++expert_idx) { + const size_t start = expert_offsets[expert_idx]; + const size_t end = expert_offsets[expert_idx + 1]; + if (start == end) continue; + + const size_t selected_count = end - start; + const size_t* selected_tokens = expert_tokens_flat + start; + + const auto& w1_buffer = get_input(node, 3 + expert_idx, nodes, node_index_map); + const auto& w2_buffer = gated + ? get_input(node, 3 + 2 * num_experts + expert_idx, nodes, node_index_map) + : get_input(node, 3 + num_experts + expert_idx, nodes, node_index_map); + + __fp16* compact_hidden = moe_compact_hidden_buf.data(); + for (size_t i = 0; i < selected_count; ++i) { + std::memcpy(compact_hidden + i * hidden_dim, + hidden + selected_tokens[i] * hidden_dim, + hidden_dim * sizeof(__fp16)); + } + + __fp16* gate = moe_gate_buf.data(); + __fp16* up = moe_up_buf.data(); + __fp16* expert_out = moe_expert_out_buf.data(); + + moe_matmul(compact_hidden, selected_count, hidden_dim, w1_buffer, gate, expert_intermediate_dim); + + switch (activation) { + case Activation::GELU: + cactus_gelu_f16(gate, gate, selected_count * expert_intermediate_dim); + break; + case Activation::GELU_ERF: + cactus_gelu_f16_erf(gate, gate, selected_count * expert_intermediate_dim); + break; + case Activation::RELU: + cactus_relu_f16(gate, gate, selected_count * expert_intermediate_dim); + break; + case Activation::SIGMOID: + cactus_sigmoid_f16(gate, gate, selected_count * expert_intermediate_dim); + break; + case Activation::TANH: + cactus_tanh_f16(gate, gate, selected_count * expert_intermediate_dim); + break; + case Activation::SILU: + default: + cactus_silu_f16(gate, gate, selected_count * expert_intermediate_dim); + break; + } + + if (gated) { + const auto& w3_buffer = get_input(node, 3 + num_experts + expert_idx, nodes, node_index_map); + moe_matmul(compact_hidden, selected_count, hidden_dim, w3_buffer, up, expert_intermediate_dim); + cactus_multiply_f16(gate, up, gate, selected_count * expert_intermediate_dim); + } + + const size_t w2_k = w2_buffer.shape.size() == 2 ? w2_buffer.shape[1] : 0; + if (w2_k < expert_intermediate_dim) { + throw std::runtime_error("moe_layer down-proj weight K smaller than expert intermediate dim"); + } + const __fp16* w2_input = gate; + if (w2_k != expert_intermediate_dim) { + if (moe_gate_pad_buf.size() < selected_count * w2_k) moe_gate_pad_buf.resize(selected_count * w2_k); + std::memset(moe_gate_pad_buf.data(), 0, selected_count * w2_k * sizeof(__fp16)); + for (size_t i = 0; i < selected_count; ++i) { + std::memcpy(moe_gate_pad_buf.data() + i * w2_k, + gate + i * expert_intermediate_dim, + expert_intermediate_dim * sizeof(__fp16)); + } + w2_input = moe_gate_pad_buf.data(); + } + + moe_matmul(w2_input, selected_count, w2_k, w2_buffer, expert_out, hidden_dim); + + for (size_t i = 0; i < selected_count; ++i) { + const size_t tok = selected_tokens[i]; + float expert_prob = routing_prob(tok, expert_idx); + if (expert_prob <= 0.0f) continue; + + float route_weight = expert_prob; + if (normalize_routing) { + route_weight = expert_prob / routing_denom[tok]; + } + route_weight *= routed_scaling_factor; + if (expert_scales_fp16) { + route_weight *= static_cast(expert_scales_fp16[expert_idx]); + } + + auto* out_row = output + tok * hidden_dim; + const auto* expert_row = expert_out + i * hidden_dim; + cactus_add_scaled_f16(out_row, expert_row, out_row, hidden_dim, route_weight); + } + } +} + +void compute_dense_mlp_tq_fused_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& hidden_buffer = get_input(node, 0, nodes, node_index_map); + const auto& gate_buffer = get_input(node, 1, nodes, node_index_map); + const auto& up_buffer = get_input(node, 2, nodes, node_index_map); + const auto& down_buffer = get_input(node, 3, nodes, node_index_map); + + if (hidden_buffer.precision != Precision::FP16 || node.output_buffer.precision != Precision::FP16) { + throw std::runtime_error("dense_mlp_tq_fused expects FP16 hidden/output"); + } + if (!PrecisionTraits::is_cq(gate_buffer.precision) || !PrecisionTraits::is_cq(up_buffer.precision) || + !PrecisionTraits::is_cq(down_buffer.precision) || + gate_buffer.group_size == 0 || up_buffer.group_size == 0 || down_buffer.group_size == 0) { + throw std::runtime_error("dense_mlp_tq_fused expects TQ gate/up/down weights"); + } + if (hidden_buffer.shape.empty() || gate_buffer.shape.size() != 2 || up_buffer.shape.size() != 2 || down_buffer.shape.size() != 2) { + throw std::runtime_error("dense_mlp_tq_fused expects rank >=1 hidden and rank-2 weights"); + } + + const size_t hidden_dim = hidden_buffer.shape.back(); + size_t M = 1; + for (size_t i = 0; i + 1 < hidden_buffer.shape.size(); ++i) M *= hidden_buffer.shape[i]; + + const size_t d_ffn = gate_buffer.shape[0]; + if (up_buffer.shape[0] != d_ffn || gate_buffer.shape[1] != hidden_dim || up_buffer.shape[1] != hidden_dim || + down_buffer.shape[1] != d_ffn) { + throw std::runtime_error("dense_mlp_tq_fused weight dimensions do not match hidden"); + } + + thread_local std::vector<__fp16> gate_buf; + thread_local std::vector<__fp16> up_buf; + thread_local std::vector<__fp16> prod_buf; + const size_t inter_size = M * d_ffn; + if (gate_buf.size() < inter_size) gate_buf.resize(inter_size); + if (up_buf.size() < inter_size) up_buf.resize(inter_size); + if (prod_buf.size() < inter_size) prod_buf.resize(inter_size); + + const __fp16* hidden = hidden_buffer.data_as<__fp16>(); + __fp16* gate = gate_buf.data(); + __fp16* up = up_buf.data(); + __fp16* prod = prod_buf.data(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + + CactusQuantMatrix gate_mat = gate_buffer.to_cq_matrix(); + CactusQuantMatrix up_mat = up_buffer.to_cq_matrix(); + CactusQuantMatrix down_mat = down_buffer.to_cq_matrix(); + const bool use_safe_product_scale = node.params.scalar != 0.0f && node.params.scalar != 1.0f; + const bool trace_dense_mlp = std::getenv("CACTUS_TRACE_DENSE_MLP") != nullptr; + + cactus_quant_matmul(&gate_mat, hidden, static_cast(M), gate); + cactus_gelu_f16(gate, gate, inter_size); + float max_gate_abs = 0.0f; + size_t gate_nonfinite = 0; + if (use_safe_product_scale) { + const __fp16 scale = static_cast<__fp16>(node.params.scalar); + for (size_t i = 0; i < inter_size; ++i) { + float value = static_cast(gate[i]); + if (!std::isfinite(value)) { + gate_nonfinite += 1; + value = std::copysign(65504.0f, value); + } + value *= static_cast(scale); + max_gate_abs = std::max(max_gate_abs, std::abs(value)); + gate[i] = static_cast<__fp16>(value); + } + } + cactus_quant_matmul(&up_mat, hidden, static_cast(M), up); + float max_up_abs = 0.0f; + size_t up_nonfinite = 0; + if (use_safe_product_scale) { + for (size_t i = 0; i < inter_size; ++i) { + float value = static_cast(up[i]); + if (!std::isfinite(value)) { + up_nonfinite += 1; + value = std::copysign(65504.0f, value); + } + max_up_abs = std::max(max_up_abs, std::abs(value)); + up[i] = static_cast<__fp16>(value); + } + const float product_bound = max_gate_abs * max_up_abs; + constexpr float kSafeHadamardProductBound = 1024.0f; + if (std::isfinite(product_bound) && product_bound > kSafeHadamardProductBound) { + const __fp16 extra_scale = static_cast<__fp16>(kSafeHadamardProductBound / product_bound); + for (size_t i = 0; i < inter_size; ++i) { + gate[i] = static_cast<__fp16>(gate[i] * extra_scale); + } + max_gate_abs *= static_cast(extra_scale); + } + } + cactus_multiply_f16(gate, up, prod, inter_size); + if (trace_dense_mlp) { + size_t prod_nonfinite = 0; + float max_prod_abs = 0.0f; + for (size_t i = 0; i < inter_size; ++i) { + float value = static_cast(prod[i]); + if (!std::isfinite(value)) { + prod_nonfinite += 1; + continue; + } + max_prod_abs = std::max(max_prod_abs, std::abs(value)); + } + std::cerr << "[cactus:dense_mlp] id=" << node.id + << " scale=" << node.params.scalar + << " gate_nonfinite=" << gate_nonfinite + << " up_nonfinite=" << up_nonfinite + << " prod_nonfinite=" << prod_nonfinite + << " max_gate=" << max_gate_abs + << " max_up=" << max_up_abs + << " max_prod=" << max_prod_abs + << " shape=["; + for (size_t i = 0; i < node.output_buffer.shape.size(); ++i) { + if (i) std::cerr << ","; + std::cerr << node.output_buffer.shape[i]; + } + std::cerr << "]" << std::endl; + } + cactus_quant_matmul(&down_mat, prod, static_cast(M), output); + if (trace_dense_mlp) { + size_t output_nonfinite = 0; + float max_output_abs = 0.0f; + const size_t output_count = node.output_buffer.total_size; + for (size_t i = 0; i < output_count; ++i) { + float value = static_cast(output[i]); + if (!std::isfinite(value)) { + output_nonfinite += 1; + continue; + } + max_output_abs = std::max(max_output_abs, std::abs(value)); + } + std::cerr << "[cactus:dense_mlp_out] id=" << node.id + << " output_nonfinite=" << output_nonfinite + << " max_output=" << max_output_abs + << std::endl; + } +} + +void compute_rms_norm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& weight_buffer = get_input(node, 1, nodes, node_index_map); + + if (input_buffer.shape.size() != 2) { + throw std::runtime_error("RMS normalization requires 2D input tensor [batch_size, dims], got " + + std::to_string(input_buffer.shape.size()) + "D tensor"); + } + + size_t batch_size = input_buffer.shape[0]; + size_t dims = input_buffer.shape[1]; + + if (input_buffer.precision != Precision::FP16) { + throw std::runtime_error("RMS normalization only supports FP16 precision"); + } + + cactus_rms_norm_f16(input_buffer.data_as<__fp16>(), weight_buffer.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), batch_size, dims, node.params.epsilon); +} + +void compute_rope_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& shape = input_buffer.shape; + + if (shape.size() < 4) { + throw std::runtime_error("RoPE operation requires 4D tensor with shape [batch, seq_len, num_heads, head_dim], got " + + std::to_string(shape.size()) + "D tensor"); + } + + if (input_buffer.precision != Precision::FP16 || node.output_buffer.precision != Precision::FP16) { + throw std::runtime_error("RoPE operation only supports FP16 precision"); + } + + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + size_t num_heads = shape[2]; + size_t head_dim = shape[3]; + + cactus_rope_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, num_heads, head_dim, node.params.position_offset, node.params.theta); +} + +void compute_softmax_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& shape = input_buffer.shape; + + if (shape.size() < 2) { + throw std::runtime_error("Softmax operation requires at least 2D tensor, got " + + std::to_string(shape.size()) + "D tensor"); + } + + if (input_buffer.precision != Precision::FP16) { + throw std::runtime_error("Softmax operation only supports FP16 precision"); + } + + size_t batch_size = 1; + for (size_t i = 0; i < shape.size() - 1; i++) { + batch_size *= shape[i]; + } + size_t vocab_size = shape[shape.size() - 1]; + + cactus_softmax_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + batch_size, 1, vocab_size); +} + +void compute_rel_pos_bias_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + if (node.input_ids.size() != 2) { + throw std::runtime_error("REL_POS_BIAS requires 2 inputs (query, relative_key)"); + } + + const auto& q_buffer = get_input(node, 0, nodes, node_index_map); + const auto& r_buffer = get_input(node, 1, nodes, node_index_map); + auto& y_buffer = node.output_buffer; + + if (q_buffer.shape.size() != 4) { + throw std::runtime_error("REL_POS_BIAS query must be [B, T, H, D]"); + } + if (r_buffer.shape.size() != 4) { + throw std::runtime_error("REL_POS_BIAS relative_key must be [B, R, H, D]"); + } + if (q_buffer.precision != Precision::FP16 || r_buffer.precision != Precision::FP16) { + throw std::runtime_error("REL_POS_BIAS currently only supports FP16 tensors"); + } + + const size_t B = q_buffer.shape[0]; + const size_t T = q_buffer.shape[1]; + const size_t H = q_buffer.shape[2]; + const size_t D = q_buffer.shape[3]; + const size_t Rb = r_buffer.shape[0]; + const size_t R = r_buffer.shape[1]; + + if (Rb != 1 && Rb != B) { + throw std::runtime_error("REL_POS_BIAS relative_key batch must be 1 or match query batch"); + } + if (r_buffer.shape[2] != H || r_buffer.shape[3] != D) { + throw std::runtime_error("REL_POS_BIAS expects matching [H, D] between query and relative_key"); + } + if (R < (2 * T - 1)) { + throw std::runtime_error("REL_POS_BIAS requires relative_key length >= 2*T-1"); + } + + const __fp16* q = q_buffer.data_as<__fp16>(); + const __fp16* r = r_buffer.data_as<__fp16>(); + __fp16* y = y_buffer.data_as<__fp16>(); + + const float scale = node.params.scale; + + const size_t q_batch_stride = T * H * D; + const size_t r_batch_stride = R * H * D; + const size_t y_batch_stride = H * T * T; + const size_t q_head_stride = D; + const size_t r_head_stride = D; + const size_t q_time_stride = H * D; + const size_t r_time_stride = H * D; + + CactusThreading::parallel_for(B * H * T, CactusThreading::Thresholds::ATTENTION, + [&](size_t start_idx, size_t end_idx) { + for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { + const size_t b = work_idx / (H * T); + const size_t rem = work_idx % (H * T); + const size_t h = rem / T; + const size_t t = rem % T; + + const size_t rb = (Rb == 1) ? 0 : b; + const __fp16* q_vec = q + b * q_batch_stride + t * q_time_stride + h * q_head_stride; + const __fp16* r_base = r + rb * r_batch_stride + h * r_head_stride; + __fp16* y_row = y + b * y_batch_stride + h * (T * T) + t * T; + + for (size_t j = 0; j < T; ++j) { + const size_t rel_idx = (T - 1) - t + j; + const __fp16* r_vec = r_base + rel_idx * r_time_stride; + + float acc = 0.0f; + for (size_t d = 0; d < D; ++d) { + acc += static_cast(q_vec[d]) * static_cast(r_vec[d]); + } + y_row[j] = static_cast<__fp16>(acc * scale); + } + } + }); +} + +void compute_attention_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + if (node.input_ids.size() < 3 || node.input_ids.size() > 4) { + throw std::runtime_error("Attention operation requires 3 or 4 inputs (query, key, value[, mask]), got " + + std::to_string(node.input_ids.size()) + " inputs"); + } + + const auto& query_buffer = get_input(node, 0, nodes, node_index_map); + const auto& key_buffer = get_input(node, 1, nodes, node_index_map); + const auto& value_buffer = get_input(node, 2, nodes, node_index_map); + const BufferDesc* mask_buffer = nullptr; + if (node.input_ids.size() == 4) { + mask_buffer = &get_input(node, 3, nodes, node_index_map); + } + const auto& q_shape = query_buffer.shape; + const auto& k_shape = key_buffer.shape; + + if (q_shape.size() < 4) { + throw std::runtime_error("Attention operation requires 4D tensors [batch, seq_len, num_heads, head_dim], got " + + std::to_string(q_shape.size()) + "D tensor"); + } + + if (query_buffer.precision != Precision::FP16) { + throw std::runtime_error("Attention operation only supports FP16 precision"); + } + + size_t batch_size = q_shape[0]; + size_t seq_len = q_shape[1]; + size_t num_q_heads = q_shape[2]; + size_t head_dim = q_shape[3]; + size_t num_kv_heads = k_shape[2]; + size_t kv_seq_len = key_buffer.shape[1]; + size_t v_head_dim = value_buffer.shape[3]; + bool mask_per_head = false; + const __fp16* mask_ptr = nullptr; + + if (mask_buffer) { + if (mask_buffer->precision != Precision::FP16) { + throw std::runtime_error("Attention mask tensor must be FP16"); + } + + if (mask_buffer->shape.size() == 3) { + if (mask_buffer->shape[0] != batch_size || + mask_buffer->shape[1] != seq_len || + mask_buffer->shape[2] != kv_seq_len) { + throw std::runtime_error("Attention mask [B, T, S] shape mismatch"); + } + mask_per_head = false; + } else if (mask_buffer->shape.size() == 4) { + if (mask_buffer->shape[0] != batch_size || + mask_buffer->shape[1] != num_q_heads || + mask_buffer->shape[2] != seq_len || + mask_buffer->shape[3] != kv_seq_len) { + throw std::runtime_error("Attention mask [B, H, T, S] shape mismatch"); + } + mask_per_head = true; + } else { + throw std::runtime_error("Attention mask must be rank 3 or 4"); + } + + mask_ptr = mask_buffer->data_as<__fp16>(); + } + + cactus_attention_f16(query_buffer.data_as<__fp16>(), key_buffer.data_as<__fp16>(), + value_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, kv_seq_len, num_q_heads, num_kv_heads, head_dim, node.params.scale, mask_ptr, + node.params.position_offset, node.params.window_size, node.params.is_causal, + node.params.attention_mask_is_additive, mask_per_head, v_head_dim, node.params.logit_cap); +} + +void compute_attention_int8_hybrid_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& query_buffer = get_input(node, 0, nodes, node_index_map); + const auto& key_new_buffer = get_input(node, 1, nodes, node_index_map); + const auto& value_new_buffer = get_input(node, 2, nodes, node_index_map); + const auto& q_shape = query_buffer.shape; + + if (q_shape.size() < 4) { + throw std::runtime_error("ATTENTION_INT8_HYBRID requires 4D query tensor"); + } + + size_t batch_size = q_shape[0]; + size_t seq_len = q_shape[1]; + size_t num_q_heads = q_shape[2]; + size_t head_dim = node.params.head_dim; + size_t v_head_dim = node.params.v_head_dim; + size_t num_kv_heads = node.params.num_kv_heads; + size_t cache_len = node.params.cache_seq_len; + size_t new_len = key_new_buffer.shape[1]; + + cactus_attention_hybrid_int8_fp16( + query_buffer.data_as<__fp16>(), + node.params.cached_keys_int8, + node.params.cached_values_int8, + node.params.cached_k_scales, + node.params.cached_v_scales, + key_new_buffer.data_as<__fp16>(), + value_new_buffer.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, cache_len, new_len, + num_q_heads, num_kv_heads, head_dim, + node.params.scale, node.params.position_offset, true, + node.params.window_size, KV_QUANT_GROUP_SIZE, v_head_dim + ); +} + +void compute_layernorm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& weight_buffer = get_input(node, 1, nodes, node_index_map); + bool has_bias = node.input_ids.size() > 2; + float epsilon = node.params.epsilon; + + if (input_buffer.shape.empty()) { + throw std::runtime_error("LayerNorm requires non-empty input tensor"); + } + + size_t feature_size = input_buffer.shape.back(); + size_t batch_size = input_buffer.total_size / feature_size; + + if (weight_buffer.total_size != feature_size) { + throw std::runtime_error("LayerNorm weight size mismatch with input feature dimension"); + } + + using BufferDesc = std::remove_reference_t; + const BufferDesc* bias_buffer_ptr = nullptr; + if (has_bias) { + const auto& bias_buffer = get_input(node, 2, nodes, node_index_map); + if (bias_buffer.total_size != feature_size) { + throw std::runtime_error("LayerNorm bias size mismatch with input feature dimension"); + } + bias_buffer_ptr = &bias_buffer; + } + + if (input_buffer.precision == Precision::FP16 && + weight_buffer.precision == Precision::FP16 && + node.output_buffer.precision == Precision::FP16 && + (!has_bias || bias_buffer_ptr->precision == Precision::FP16)) { + cactus_layer_norm_f16( + input_buffer.data_as<__fp16>(), + weight_buffer.data_as<__fp16>(), + has_bias ? bias_buffer_ptr->data_as<__fp16>() : nullptr, + node.output_buffer.data_as<__fp16>(), + batch_size, + feature_size, + epsilon); + return; + } + + std::vector input_float(input_buffer.total_size); + std::vector weight_float(feature_size); + std::vector bias_float(feature_size, 0.0f); + + if (input_buffer.precision == Precision::INT8) { + throw std::runtime_error("LayerNorm currently does not support INT8 input"); + } else if (input_buffer.precision == Precision::FP16) { + const __fp16* input_fp16 = input_buffer.data_as<__fp16>(); + for (size_t i = 0; i < input_buffer.total_size; ++i) { + input_float[i] = static_cast(input_fp16[i]); + } + } else { + std::memcpy(input_float.data(), input_buffer.data_as(), input_buffer.total_size * sizeof(float)); + } + + if (weight_buffer.precision == Precision::INT8) { + throw std::runtime_error("LayerNorm currently does not support INT8 weight"); + } else if (weight_buffer.precision == Precision::FP16) { + const __fp16* weight_fp16 = weight_buffer.data_as<__fp16>(); + for (size_t i = 0; i < feature_size; ++i) { + weight_float[i] = static_cast(weight_fp16[i]); + } + } else { + std::memcpy(weight_float.data(), weight_buffer.data_as(), feature_size * sizeof(float)); + } + + if (has_bias) { + const auto& bias_buffer = *bias_buffer_ptr; + if (bias_buffer.precision == Precision::INT8) { + throw std::runtime_error("LayerNorm currently does not support INT8 bias"); + } else if (bias_buffer.precision == Precision::FP16) { + const __fp16* bias_fp16 = bias_buffer.data_as<__fp16>(); + for (size_t i = 0; i < feature_size; ++i) { + bias_float[i] = static_cast(bias_fp16[i]); + } + } else { + std::memcpy(bias_float.data(), bias_buffer.data_as(), feature_size * sizeof(float)); + } + } + + std::vector output_float(input_buffer.total_size); + for (size_t b = 0; b < batch_size; ++b) { + const float* input_row = input_float.data() + b * feature_size; + float* output_row = output_float.data() + b * feature_size; + + float mean = 0.0f; + for (size_t i = 0; i < feature_size; ++i) { + mean += input_row[i]; + } + mean /= feature_size; + + float variance = 0.0f; + for (size_t i = 0; i < feature_size; ++i) { + float diff = input_row[i] - mean; + variance += diff * diff; + } + variance /= feature_size; + + float std_inv = 1.0f / std::sqrt(variance + epsilon); + for (size_t i = 0; i < feature_size; ++i) { + output_row[i] = (input_row[i] - mean) * std_inv * weight_float[i] + bias_float[i]; + } + } + + if (node.output_buffer.precision == Precision::INT8) { + throw std::runtime_error("LayerNorm currently does not support INT8 output"); + } else if (node.output_buffer.precision == Precision::FP16) { + __fp16* output_fp16 = node.output_buffer.data_as<__fp16>(); + for (size_t i = 0; i < input_buffer.total_size; ++i) { + output_fp16[i] = static_cast<__fp16>(output_float[i]); + } + } else { + std::memcpy(node.output_buffer.data_as(), output_float.data(), input_buffer.total_size * sizeof(float)); + } +} + + + + + + + + + +void compute_glu_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + auto& Y = node.output_buffer; + + if (X.shape.empty()) { + throw std::runtime_error("GLU expects non-scalar input"); + } + + int axis = node.params.axis; + if (axis < 0) axis += static_cast(X.shape.size()); + if (axis < 0 || static_cast(axis) >= X.shape.size()) { + throw std::runtime_error("GLU axis out of range"); + } + + const size_t axis_size = X.shape[static_cast(axis)]; + if ((axis_size % 2) != 0) { + throw std::runtime_error("GLU split dimension must be even"); + } + const size_t split = axis_size / 2; + + size_t outer = 1; + for (int i = 0; i < axis; ++i) { + outer *= X.shape[static_cast(i)]; + } + size_t inner = 1; + for (size_t i = static_cast(axis) + 1; i < X.shape.size(); ++i) { + inner *= X.shape[i]; + } + + std::vector out_shape = X.shape; + out_shape[static_cast(axis)] = split; + Y.shape = out_shape; + Y.precision = X.precision; + + if (X.precision == Precision::FP16) { + cactus_glu_f16(X.data_as<__fp16>(), Y.data_as<__fp16>(), outer, split, inner); + return; + } + + if (X.precision == Precision::FP32) { + cactus_glu_f32(X.data_as(), Y.data_as(), outer, split, inner); + return; + } + + throw std::runtime_error("GLU only supports FP16/FP32"); +} + + + + + +void compute_groupnorm_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + const auto& weight = get_input(node, 1, nodes, node_index_map); + const auto& bias = get_input(node, 2, nodes, node_index_map); + float epsilon = node.params.epsilon; + + size_t batch_size = input.shape[0]; + size_t channels = input.shape[1]; + size_t spatial_size = 1; + for (size_t i = 2; i < input.shape.size(); ++i) spatial_size *= input.shape[i]; + + size_t num_groups = node.params.num_groups; + if (num_groups == 0) num_groups = 32; + + if (channels % num_groups != 0) { + throw std::runtime_error("GroupNorm: channels must be divisible by num_groups"); + } + + size_t channels_per_group = channels / num_groups; + + const __fp16* src = input.data_as<__fp16>(); + const __fp16* w = weight.data_as<__fp16>(); + const __fp16* b = bias.data_as<__fp16>(); + __fp16* dst = node.output_buffer.data_as<__fp16>(); + + for (size_t n = 0; n < batch_size; ++n) { + for (size_t g = 0; g < num_groups; ++g) { + float sum = 0.0f, sum_sq = 0.0f; + size_t count = 0; + + for (size_t c = 0; c < channels_per_group; ++c) { + size_t ch = g * channels_per_group + c; + for (size_t s = 0; s < spatial_size; ++s) { + size_t idx = n * channels * spatial_size + ch * spatial_size + s; + float val = static_cast(src[idx]); + sum += val; + sum_sq += val * val; + count++; + } + } + + float mean = sum / count; + float var = (sum_sq / count) - (mean * mean); + float inv_std = 1.0f / std::sqrt(var + epsilon); + + for (size_t c = 0; c < channels_per_group; ++c) { + size_t ch = g * channels_per_group + c; + float wt = static_cast(w[ch]); + float bi = static_cast(b[ch]); + + for (size_t s = 0; s < spatial_size; ++s) { + size_t idx = n * channels * spatial_size + ch * spatial_size + s; + float val = static_cast(src[idx]); + dst[idx] = static_cast<__fp16>((val - mean) * inv_std * wt + bi); + } + } + } + } +} diff --git a/cactus-graph/src/ops_recurrent.cpp b/cactus-graph/src/ops_recurrent.cpp new file mode 100644 index 000000000..ee5ecf0cc --- /dev/null +++ b/cactus-graph/src/ops_recurrent.cpp @@ -0,0 +1,403 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include +#include +#include +#include +#include + +namespace { + const __fp16* as_fp16_ptr(const BufferDesc& buffer, std::vector<__fp16>& scratch) { + if (buffer.precision == Precision::FP16) { + return buffer.data_as<__fp16>(); + } + if (buffer.precision == Precision::FP32) { + scratch.resize(buffer.total_size); + cactus_fp32_to_fp16(buffer.data_as(), scratch.data(), buffer.total_size); + return scratch.data(); + } + throw std::runtime_error("GATED_DELTANET unsupported precision (expected FP16/FP32)"); + } + + void validate_gated_deltanet_inputs( + const BufferDesc& q, + const BufferDesc& k, + const BufferDesc& v, + const BufferDesc& g, + const BufferDesc& b, + const BufferDesc& s) { + auto is_supported_precision = [](Precision p) { + return p == Precision::FP16 || p == Precision::FP32; + }; + if (!is_supported_precision(q.precision) || !is_supported_precision(k.precision) || + !is_supported_precision(v.precision) || !is_supported_precision(g.precision) || + !is_supported_precision(b.precision) || !is_supported_precision(s.precision)) { + throw std::runtime_error("GATED_DELTANET requires FP16/FP32 inputs"); + } + + if (q.shape.size() != 4 || k.shape.size() != 4 || v.shape.size() != 4) { + throw std::runtime_error("GATED_DELTANET expects query/key/value rank 4 [B, T, H, D]"); + } + if (g.shape.size() != 3 || b.shape.size() != 3) { + throw std::runtime_error("GATED_DELTANET expects gate_log/beta rank 3 [B, T, H]"); + } + if (s.shape.size() != 4) { + throw std::runtime_error("GATED_DELTANET expects state rank 4 [B, K, H, V]"); + } + + const size_t B = q.shape[0]; + const size_t T = q.shape[1]; + const size_t Hq = q.shape[2]; + const size_t K = q.shape[3]; + + if (k.shape[0] != B || k.shape[1] != T || k.shape[2] != Hq || k.shape[3] != K) { + throw std::runtime_error("GATED_DELTANET query/key shape mismatch"); + } + if (v.shape[0] != B || v.shape[1] != T) { + throw std::runtime_error("GATED_DELTANET value shape mismatch"); + } + const size_t Hv = v.shape[2]; + if (g.shape[0] != B || g.shape[1] != T || g.shape[2] != Hv || + b.shape[0] != B || b.shape[1] != T || b.shape[2] != Hv) { + throw std::runtime_error("GATED_DELTANET gate_log/beta shape mismatch"); + } + if (Hq == 0 || Hv == 0 || (Hv % Hq) != 0) { + throw std::runtime_error("GATED_DELTANET expects value heads divisible by q/k heads"); + } + const size_t V = v.shape[3]; + if (s.shape[0] != B || s.shape[1] != K || s.shape[2] != Hv || s.shape[3] != V) { + throw std::runtime_error("GATED_DELTANET state shape mismatch"); + } + } + + +} + +void compute_gated_deltanet_decode_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + if (node.input_ids.size() != 6) { + throw std::runtime_error("GATED_DELTANET_DECODE expects 6 inputs"); + } + + const auto& q = get_input(node, 0, nodes, node_index_map); + const auto& k = get_input(node, 1, nodes, node_index_map); + const auto& v = get_input(node, 2, nodes, node_index_map); + const auto& g = get_input(node, 3, nodes, node_index_map); + const auto& b = get_input(node, 4, nodes, node_index_map); + const auto& s = get_input(node, 5, nodes, node_index_map); + + validate_gated_deltanet_inputs(q, k, v, g, b, s); + if (q.shape[1] != 1) { + throw std::runtime_error("GATED_DELTANET_DECODE expects T=1"); + } + + const size_t B = q.shape[0]; + const size_t Hq = q.shape[2]; + const size_t K = q.shape[3]; + const size_t Hv = v.shape[2]; + const size_t V = v.shape[3]; + const size_t qk_heads_from_params = node.params.num_kv_heads; + if (qk_heads_from_params != 0 && qk_heads_from_params != Hq) { + throw std::runtime_error("GATED_DELTANET_DECODE num_qk_heads param mismatch"); + } + + std::vector<__fp16> q_cast; + std::vector<__fp16> k_cast; + std::vector<__fp16> v_cast; + std::vector<__fp16> g_cast; + std::vector<__fp16> b_cast; + std::vector<__fp16> s_cast; + const __fp16* q_data = as_fp16_ptr(q, q_cast); + const __fp16* k_data = as_fp16_ptr(k, k_cast); + const __fp16* v_data = as_fp16_ptr(v, v_cast); + const __fp16* g_data = as_fp16_ptr(g, g_cast); + const __fp16* b_data = as_fp16_ptr(b, b_cast); + const __fp16* s_data = as_fp16_ptr(s, s_cast); + __fp16* out = node.output_buffer.data_as<__fp16>(); + + cactus_gated_deltanet_decode_f16( + q_data, k_data, v_data, g_data, b_data, s_data, out, + B, Hq, Hv, K, V, node.params.scale); +} + + +void compute_gated_deltanet_prefill_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + if (node.input_ids.size() != 6) { + throw std::runtime_error("GATED_DELTANET_PREFILL expects 6 inputs"); + } + + const auto& q = get_input(node, 0, nodes, node_index_map); + const auto& k = get_input(node, 1, nodes, node_index_map); + const auto& v = get_input(node, 2, nodes, node_index_map); + const auto& g = get_input(node, 3, nodes, node_index_map); + const auto& b = get_input(node, 4, nodes, node_index_map); + const auto& s = get_input(node, 5, nodes, node_index_map); + + validate_gated_deltanet_inputs(q, k, v, g, b, s); + + const size_t B = q.shape[0]; + const size_t T = q.shape[1]; + const size_t Hq = q.shape[2]; + const size_t K = q.shape[3]; + const size_t Hv = v.shape[2]; + const size_t V = v.shape[3]; + const size_t qk_heads_from_params = node.params.num_kv_heads; + if (qk_heads_from_params != 0 && qk_heads_from_params != Hq) { + throw std::runtime_error("GATED_DELTANET_PREFILL num_qk_heads param mismatch"); + } + + std::vector<__fp16> q_cast; + std::vector<__fp16> k_cast; + std::vector<__fp16> v_cast; + std::vector<__fp16> g_cast; + std::vector<__fp16> b_cast; + std::vector<__fp16> s_cast; + const __fp16* q_data = as_fp16_ptr(q, q_cast); + const __fp16* k_data = as_fp16_ptr(k, k_cast); + const __fp16* v_data = as_fp16_ptr(v, v_cast); + const __fp16* g_data = as_fp16_ptr(g, g_cast); + const __fp16* b_data = as_fp16_ptr(b, b_cast); + const __fp16* s_data = as_fp16_ptr(s, s_cast); + __fp16* out = node.output_buffer.data_as<__fp16>(); + + cactus_gated_deltanet_prefill_f16( + q_data, k_data, v_data, g_data, b_data, s_data, out, + B, T, Hq, Hv, K, V, node.params.chunk_size, node.params.scale); +} + + +void compute_rope_gptj_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& shape = input_buffer.shape; + + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + size_t num_heads = shape[2]; + size_t head_dim = shape[3]; + size_t rot_dim = static_cast(node.params.scalar); + + cactus_gpt_j_rope_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, num_heads, head_dim, rot_dim, + node.params.position_offset, node.params.theta); +} + +void compute_lstm_cell_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& h_prev_buffer = get_input(node, 1, nodes, node_index_map); + const auto& c_prev_buffer = get_input(node, 2, nodes, node_index_map); + const auto& weight_ih_buffer = get_input(node, 3, nodes, node_index_map); + const auto& weight_hh_buffer = get_input(node, 4, nodes, node_index_map); + const auto& bias_ih_buffer = get_input(node, 5, nodes, node_index_map); + const auto& bias_hh_buffer = get_input(node, 6, nodes, node_index_map); + + if (input_buffer.precision != Precision::FP16 || h_prev_buffer.precision != Precision::FP16 || + c_prev_buffer.precision != Precision::FP16 || weight_ih_buffer.precision != Precision::FP16 || + weight_hh_buffer.precision != Precision::FP16 || bias_ih_buffer.precision != Precision::FP16 || + bias_hh_buffer.precision != Precision::FP16) { + throw std::runtime_error("LSTM cell requires all inputs to be FP16"); + } + + if (input_buffer.shape.size() != 2 || h_prev_buffer.shape.size() != 2 || c_prev_buffer.shape.size() != 2) { + throw std::runtime_error("LSTM cell input/state shapes must be 2D [batch, features]"); + } + + const size_t batch_size = input_buffer.shape[0]; + const size_t input_size = input_buffer.shape[1]; + const size_t hidden_size = h_prev_buffer.shape[1]; + + const __fp16* x_input = input_buffer.data_as<__fp16>(); + const __fp16* h_prev = h_prev_buffer.data_as<__fp16>(); + const __fp16* c_prev = c_prev_buffer.data_as<__fp16>(); + const __fp16* weight_ih = weight_ih_buffer.data_as<__fp16>(); + const __fp16* weight_hh = weight_hh_buffer.data_as<__fp16>(); + const __fp16* bias_ih = bias_ih_buffer.data_as<__fp16>(); + const __fp16* bias_hh = bias_hh_buffer.data_as<__fp16>(); + + node.output_buffer.shape = {batch_size, hidden_size, 2}; + node.output_buffer.total_size = batch_size * hidden_size * 2; + node.output_buffer.precision = Precision::FP16; + node.output_buffer.allocate(); + + std::vector<__fp16> h_new_temp(batch_size * hidden_size); + std::vector<__fp16> c_new_temp(batch_size * hidden_size); + + cactus_lstm_cell_f16( + x_input, h_prev, c_prev, + weight_ih, weight_hh, + bias_ih, bias_hh, + h_new_temp.data(), c_new_temp.data(), + batch_size, input_size, hidden_size + ); + + __fp16* output = node.output_buffer.data_as<__fp16>(); + for (size_t b = 0; b < batch_size; ++b) { + for (size_t i = 0; i < hidden_size; ++i) { + const size_t idx = b * hidden_size + i; + output[b * hidden_size * 2 + i * 2] = h_new_temp[idx]; + output[b * hidden_size * 2 + i * 2 + 1] = c_new_temp[idx]; + } + } +} + +void compute_altup_predict_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + size_t n = node.params.num_altup_inputs; + const auto& coefs_buf = get_input(node, 0, nodes, node_index_map); + + std::vector stream_ptrs(n); + for (size_t i = 0; i < n; i++) { + stream_ptrs[i] = get_input(node, 1 + i, nodes, node_index_map).data_as<__fp16>(); + } + + const auto& stream0_buf = get_input(node, 1, nodes, node_index_map); + size_t seq_len = stream0_buf.shape[0]; + size_t hidden_dim = stream0_buf.shape[1]; + + cactus_altup_predict_f16( + coefs_buf.data_as<__fp16>(), + stream_ptrs.data(), + node.output_buffer.data_as<__fp16>(), + n, seq_len, hidden_dim); +} + +void compute_gaussian_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buf = get_input(node, 0, nodes, node_index_map); + const __fp16* input = input_buf.data_as<__fp16>(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + + size_t rows = input_buf.shape[0]; + size_t cols = input_buf.shape[1]; + float ppf = node.params.scalar; + + cactus_gaussian_topk_f16(input, output, rows, cols, ppf); +} + +void compute_altup_correct_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + size_t n = node.params.num_altup_inputs; + const auto& coefs_buf = get_input(node, 0, nodes, node_index_map); + const auto& innov_buf = get_input(node, 1, nodes, node_index_map); + + std::vector pred_ptrs(n); + for (size_t i = 0; i < n; i++) { + pred_ptrs[i] = get_input(node, 2 + i, nodes, node_index_map).data_as<__fp16>(); + } + + size_t seq_len = innov_buf.shape[0]; + size_t hidden_dim = innov_buf.shape[1]; + + cactus_altup_correct_f16( + coefs_buf.data_as<__fp16>(), + innov_buf.data_as<__fp16>(), + pred_ptrs.data(), + node.output_buffer.data_as<__fp16>(), + n, seq_len, hidden_dim); +} + +void compute_bilstm_sequence_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + const auto& w_ih_fwd = get_input(node, 1, nodes, node_index_map); + const auto& w_hh_fwd = get_input(node, 2, nodes, node_index_map); + const auto& b_ih_fwd = get_input(node, 3, nodes, node_index_map); + const auto& b_hh_fwd = get_input(node, 4, nodes, node_index_map); + const auto& w_ih_bwd = get_input(node, 5, nodes, node_index_map); + const auto& w_hh_bwd = get_input(node, 6, nodes, node_index_map); + const auto& b_ih_bwd = get_input(node, 7, nodes, node_index_map); + const auto& b_hh_bwd = get_input(node, 8, nodes, node_index_map); + + size_t batch_size = input.shape[0]; + size_t seq_len = input.shape[1]; + size_t input_size = input.shape[2]; + size_t hidden_size = w_ih_fwd.shape[0] / 4; + + cactus_bilstm_sequence_f16( + input.data_as<__fp16>(), + w_ih_fwd.data_as<__fp16>(), w_hh_fwd.data_as<__fp16>(), + b_ih_fwd.data_as<__fp16>(), b_hh_fwd.data_as<__fp16>(), + w_ih_bwd.data_as<__fp16>(), w_hh_bwd.data_as<__fp16>(), + b_ih_bwd.data_as<__fp16>(), b_hh_bwd.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + batch_size, seq_len, input_size, hidden_size); +} + +void compute_stats_pool_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + const __fp16* src = input.data_as<__fp16>(); + __fp16* dst = node.output_buffer.data_as<__fp16>(); + + size_t batch = input.shape[0]; + size_t total_per_batch = input.total_size / batch; + size_t T = input.shape.back(); + size_t features = total_per_batch / T; + + for (size_t b = 0; b < batch; ++b) { + const __fp16* batch_src = src + b * total_per_batch; + __fp16* batch_dst = dst + b * features * 2; + + for (size_t f = 0; f < features; ++f) { + float sum = 0.0f, sum_sq = 0.0f; + for (size_t t = 0; t < T; ++t) { + float v = static_cast(batch_src[f * T + t]); + sum += v; + sum_sq += v * v; + } + float mean = sum / static_cast(T); + float var = T > 1 ? (sum_sq - static_cast(T) * mean * mean) / static_cast(T - 1) : 0.0f; + float std_val = sqrtf(fmaxf(var, 0.0f)); + batch_dst[f] = static_cast<__fp16>(mean); + batch_dst[features + f] = static_cast<__fp16>(std_val); + } + } +} + +void compute_weighted_stats_pool_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& input = get_input(node, 0, nodes, node_index_map); + const auto& weight_buf = get_input(node, 1, nodes, node_index_map); + const __fp16* src = input.data_as<__fp16>(); + const float* weights = weight_buf.data_as(); + __fp16* dst = node.output_buffer.data_as<__fp16>(); + + size_t batch = input.shape[0]; + size_t total_per_batch = input.total_size / batch; + size_t T = input.shape.back(); + size_t features = total_per_batch / T; + + constexpr float eps = 1e-8f; + + for (size_t b = 0; b < batch; ++b) { + const __fp16* batch_src = src + b * total_per_batch; + const float* batch_w = weights + b * T; + __fp16* batch_dst = dst + b * features * 2; + + float v1 = 0.0f, v2 = 0.0f; + for (size_t t = 0; t < T; ++t) { + float w = batch_w[t]; + v1 += w; + v2 += w * w; + } + float v1_safe = v1 + eps; + float var_denom = v1_safe - v2 / v1_safe + eps; + + for (size_t f = 0; f < features; ++f) { + float wsum = 0.0f; + for (size_t t = 0; t < T; ++t) { + wsum += static_cast(batch_src[f * T + t]) * batch_w[t]; + } + float mean = wsum / v1_safe; + + float wvar = 0.0f; + for (size_t t = 0; t < T; ++t) { + float dx = static_cast(batch_src[f * T + t]) - mean; + wvar += batch_w[t] * dx * dx; + } + float std_val = sqrtf(fmaxf(wvar / var_denom, 0.0f)); + + batch_dst[f] = static_cast<__fp16>(mean); + batch_dst[features + f] = static_cast<__fp16>(std_val); + } + } +} + diff --git a/cactus/graph/graph_ops_sample.cpp b/cactus-graph/src/ops_sample.cpp similarity index 83% rename from cactus/graph/graph_ops_sample.cpp rename to cactus-graph/src/ops_sample.cpp index 44194713b..1b2ace329 100644 --- a/cactus/graph/graph_ops_sample.cpp +++ b/cactus-graph/src/ops_sample.cpp @@ -1,15 +1,17 @@ -#include "graph.h" -#include "../kernel/kernel.h" +#include "../cactus_graph.h" +#include "cactus_kernels.h" #include #include #include #include void compute_sample_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& logits_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; + const auto& logits_buffer = get_input(node, 0, nodes, node_index_map); float temperature = node.params.temperature; float top_p = node.params.top_p; + float min_p = node.params.min_p; + float repetition_penalty = node.params.repetition_penalty; size_t top_k = node.params.top_k; size_t random_seed = node.params.random_seed; @@ -27,19 +29,19 @@ void compute_sample_node(GraphNode& node, const std::vector(); - cactus_sample_f16(logits_fp16 + last_token_offset, node.output_buffer.data_as(), - vocab_size, temperature, top_p, top_k, random_seed, - bias_values, bias_indices, bias_count); + cactus_sample_f16_ex(logits_fp16 + last_token_offset, node.output_buffer.data_as(), + vocab_size, temperature, top_p, min_p, repetition_penalty, top_k, random_seed, + bias_values, bias_indices, bias_count); } else { const float* logits_fp32 = logits_buffer.data_as(); - cactus_sample_f32(logits_fp32 + last_token_offset, node.output_buffer.data_as(), - vocab_size, temperature, top_p, top_k, random_seed, - bias_values, bias_indices, bias_count); + cactus_sample_f32_ex(logits_fp32 + last_token_offset, node.output_buffer.data_as(), + vocab_size, temperature, top_p, min_p, repetition_penalty, top_k, random_seed, + bias_values, bias_indices, bias_count); } } void compute_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); if (input_buffer.shape.size() != 2) { throw std::runtime_error("TopK currently only supports 2D tensors [batch, features]"); } @@ -87,8 +89,8 @@ void compute_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& indices_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& values_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; + const auto& indices_buffer = get_input(node, 0, nodes, node_index_map); + const auto& values_buffer = get_input(node, 1, nodes, node_index_map); if (indices_buffer.shape != values_buffer.shape) { throw std::runtime_error("ScatterTopK requires indices and values with identical shapes"); diff --git a/cactus-graph/src/ops_tensor.cpp b/cactus-graph/src/ops_tensor.cpp new file mode 100644 index 000000000..1afb75563 --- /dev/null +++ b/cactus-graph/src/ops_tensor.cpp @@ -0,0 +1,375 @@ +#include "../cactus_graph.h" +#include "cactus_kernels.h" +#include +#include +#include +#include +#include + +void compute_transpose_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + + if (input_buffer.precision != Precision::FP16) { + throw std::runtime_error("Transpose only supports FP16 precision"); + } + + const auto& permutation = node.params.permutation; + + const __fp16* input = input_buffer.data_as<__fp16>(); + __fp16* output = node.output_buffer.data_as<__fp16>(); + cactus_transpose_f16(input, output, input_buffer.shape.data(), permutation.data(), permutation.size(), 0, input_buffer.total_size); +} + +void compute_gather_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& tensor_buffer = get_input(node, 0, nodes, node_index_map); + const auto& indices_buffer = get_input(node, 1, nodes, node_index_map); + + size_t first_dim = tensor_buffer.shape[0]; + size_t element_size = 1; + for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { + element_size *= tensor_buffer.shape[i]; + } + + size_t num_indices = indices_buffer.total_size; + size_t bytes_per_element = PrecisionTraits::packed_size_of(tensor_buffer.precision, element_size); + + { + const char* tensor_data = static_cast(tensor_buffer.get_data()); + char* output = static_cast(node.output_buffer.get_data()); + + auto gather_loop = [&](auto get_idx) { + for (size_t i = 0; i < num_indices; i++) { + size_t idx = get_idx(i); + if (idx >= first_dim) + throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); + std::memcpy(output + PrecisionTraits::byte_offset_of(tensor_buffer.precision, i * element_size), + tensor_data + PrecisionTraits::byte_offset_of(tensor_buffer.precision, idx * element_size), + bytes_per_element); + } + }; + + if (indices_buffer.precision == Precision::INT8) { + const int8_t* indices = indices_buffer.data_as(); + gather_loop([&](size_t i) { return static_cast(indices[i]); }); + } else { + const float* indices = indices_buffer.data_as(); + gather_loop([&](size_t i) { return static_cast(indices[i]); }); + } + } +} + +void compute_slice_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + auto* input_node = nodes[node_index_map.at(node.input_ids[0])].get(); + auto& input_buffer = input_node->output_buffer; + + const size_t axis_index = static_cast(node.params.axis); + + const size_t axis_size = input_buffer.shape[axis_index]; + const size_t slice_start = node.params.slice_start; + size_t slice_length = node.params.slice_length; + + if (slice_length == 0) { + slice_length = axis_size - slice_start; + } + + auto dims = AxisDims::from_shape(input_buffer.shape, axis_index); + + if (dims.outer == 1) { + auto* base_ptr = static_cast(input_buffer.get_data()); + if (!base_ptr) { + throw std::runtime_error("Slice input buffer is not available"); + } + + const size_t byte_offset = PrecisionTraits::byte_offset_of(input_buffer.precision, slice_start * dims.inner); + + node.output_buffer.set_external(base_ptr + byte_offset); + node.output_buffer.precision = input_buffer.precision; + + if (input_buffer.is_cq()) { + node.output_buffer.group_size = input_buffer.group_size; + node.output_buffer.num_groups = input_buffer.num_groups; + } + return; + } + + const char* input_ptr = static_cast(input_buffer.get_data()); + if (!input_ptr) { + throw std::runtime_error("Slice input buffer is not available"); + } + + node.output_buffer.external_data = nullptr; + node.output_buffer.allocate(); + node.output_buffer.precision = input_buffer.precision; + + auto* output_ptr = static_cast(node.output_buffer.get_data()); + if (!output_ptr) { + throw std::runtime_error("Slice output buffer could not be allocated"); + } + + const size_t copy_block_elements = slice_length * dims.inner; + const size_t axis_stride_elements = axis_size * dims.inner; + const size_t copy_block_bytes = PrecisionTraits::byte_offset_of(input_buffer.precision, copy_block_elements); + const size_t axis_stride_bytes = PrecisionTraits::byte_offset_of(input_buffer.precision, axis_stride_elements); + + for (size_t outer = 0; outer < dims.outer; ++outer) { + const char* src = input_ptr + outer * axis_stride_bytes + PrecisionTraits::byte_offset_of(input_buffer.precision, slice_start * dims.inner); + char* dst = output_ptr + outer * copy_block_bytes; + std::memcpy(dst, src, copy_block_bytes); + } +} + +void compute_embedding_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& embeddings_buffer = get_input(node, 0, nodes, node_index_map); + const auto& indices_buffer = get_input(node, 1, nodes, node_index_map); + + size_t hidden_dim = embeddings_buffer.shape[1]; + size_t num_indices = indices_buffer.total_size; + size_t vocab_size = embeddings_buffer.shape[0]; + + std::vector indices_float; + const float* indices_ptr; + if (indices_buffer.precision == Precision::FP32) { + indices_ptr = indices_buffer.data_as(); + } else if (indices_buffer.precision == Precision::FP16) { + indices_float.resize(num_indices); + const __fp16* half_indices = indices_buffer.data_as<__fp16>(); + for (size_t i = 0; i < num_indices; i++) { + indices_float[i] = static_cast(half_indices[i]); + } + indices_ptr = indices_float.data(); + } else { + indices_float.resize(num_indices); + const int8_t* int_indices = indices_buffer.data_as(); + for (size_t i = 0; i < num_indices; i++) { + indices_float[i] = static_cast(int_indices[i]); + } + indices_ptr = indices_float.data(); + } + + __fp16* output = node.output_buffer.data_as<__fp16>(); + + Precision emb_prec = embeddings_buffer.precision; + if (PrecisionTraits::is_cq(emb_prec) && embeddings_buffer.group_size > 0) { + bool orthogonal = (embeddings_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) != 0; + std::unordered_map row_cache; + std::vector<__fp16> cached_rows; + if (num_indices > 16) { + row_cache.reserve(std::min(num_indices, 256)); + } + for (size_t i = 0; i < num_indices; i++) { + size_t idx = static_cast(indices_ptr[i]); + if (idx >= vocab_size) { + throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); + } + if (num_indices > 16) { + auto it = row_cache.find(idx); + if (it != row_cache.end()) { + std::memcpy(output + i * hidden_dim, + cached_rows.data() + it->second * hidden_dim, + hidden_dim * sizeof(__fp16)); + continue; + } + } + if (orthogonal) { + cactus_quant_dequantize_orthogonal_embedding_row( + PrecisionTraits::cq_bits(emb_prec), + static_cast(hidden_dim), + idx, + embeddings_buffer.data_as(), + embeddings_buffer.cq_codebook, + embeddings_buffer.cq_norms, + embeddings_buffer.cq_input_scale_recip, + embeddings_buffer.cq_rotation, + embeddings_buffer.cq_flags, + output + i * hidden_dim); + } else { + cactus_quant_dequantize_hadamard_embedding_row( + PrecisionTraits::cq_bits(emb_prec), + static_cast(hidden_dim), + static_cast(embeddings_buffer.group_size), + static_cast(embeddings_buffer.num_groups), + idx, + embeddings_buffer.data_as(), + embeddings_buffer.cq_codebook, + embeddings_buffer.cq_norms, + embeddings_buffer.cq_input_scale_recip, + embeddings_buffer.cq_left_signs, + embeddings_buffer.cq_right_signs, + embeddings_buffer.cq_permutation, + output + i * hidden_dim); + } + if (num_indices > 16) { + const size_t cache_slot = cached_rows.size() / hidden_dim; + row_cache.emplace(idx, cache_slot); + cached_rows.insert(cached_rows.end(), output + i * hidden_dim, output + (i + 1) * hidden_dim); + } + } + } else if (embeddings_buffer.precision == Precision::FP16) { + const __fp16* embeddings = embeddings_buffer.data_as<__fp16>(); + for (size_t i = 0; i < num_indices; i++) { + size_t idx = static_cast(indices_ptr[i]); + if (idx >= vocab_size) { + throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); + } + std::memcpy(output + i * hidden_dim, embeddings + idx * hidden_dim, hidden_dim * sizeof(__fp16)); + } + } else { + throw std::runtime_error("Embedding requires CQ quantized or FP16 data"); + } +} + +void compute_concat_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input1_buffer = get_input(node, 0, nodes, node_index_map); + const auto& input2_buffer = get_input(node, 1, nodes, node_index_map); + + std::vector shape1 = input1_buffer.shape; + std::vector shape2 = input2_buffer.shape; + std::vector output_shape = node.output_buffer.shape; + + if (input1_buffer.precision != Precision::FP16) { + throw std::runtime_error("Concat operation only supports FP16 precision"); + } + cactus_concat_f16(input1_buffer.data_as<__fp16>(), input2_buffer.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + shape1.data(), shape2.data(), output_shape.data(), + shape1.size(), node.params.axis); +} + +void compute_cat_node( + GraphNode& node, + const std::vector>& nodes, + const std::unordered_map& node_index_map +) { + if (node.params.axis < 0) { + throw std::runtime_error("Cat operation does not support negative axis"); + } + if (node.input_ids.size() < 2) { + throw std::runtime_error("Cat operation requires at least 2 input tensors"); + } + + const auto& first_buffer = get_input(node, 0, nodes, node_index_map); + + if (first_buffer.precision != Precision::FP16) { + throw std::runtime_error("Cat operation only supports FP16 precision"); + } + + std::vector input_data_ptrs(node.input_ids.size()); + std::vector input_shape_ptrs(node.input_ids.size()); + + for (size_t i = 0; i < node.input_ids.size(); i++) { + const auto& buffer = get_input(node, i, nodes, node_index_map); + + if (buffer.precision != Precision::FP16) { + throw std::runtime_error("Cat operation only supports FP16 precision"); + } + + input_data_ptrs[i] = buffer.data_as<__fp16>(); + input_shape_ptrs[i] = buffer.shape.data(); + } + + cactus_cat_f16(input_data_ptrs.data(), + node.output_buffer.data_as<__fp16>(), + input_shape_ptrs.data(), + node.output_buffer.shape.data(), + node.input_ids.size(), + node.output_buffer.shape.size(), + node.params.axis); +} + +void compute_index_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& input_buffer = get_input(node, 0, nodes, node_index_map); + const auto& input_shape = input_buffer.shape; + + int dim = node.params.axis; + size_t index_value = node.params.index_value; + + const char* input_data = static_cast(input_buffer.get_data()); + char* output_data = static_cast(node.output_buffer.get_data()); + + if (dim == 0) { + size_t slice_size = input_buffer.total_size / input_shape[0]; + size_t offset_bytes = PrecisionTraits::byte_offset_of(input_buffer.precision, index_value * slice_size); + node.output_buffer.set_external(const_cast(input_data) + offset_bytes); + return; + } + + std::vector input_strides(input_shape.size()); + input_strides[input_shape.size() - 1] = 1; + for (int i = static_cast(input_shape.size()) - 2; i >= 0; --i) { + input_strides[i] = input_strides[i + 1] * input_shape[i + 1]; + } + + size_t slice_size = input_strides[dim]; + size_t outer_size = input_buffer.total_size / input_strides[dim - 1]; + size_t dim_stride = input_strides[dim]; + size_t block_size = dim_stride * input_shape[dim]; + + size_t output_idx = 0; + for (size_t outer_idx = 0; outer_idx < outer_size; ++outer_idx) { + size_t input_base = outer_idx * block_size + index_value * dim_stride; + + char* output_offset_bytes = output_data + PrecisionTraits::byte_offset_of(input_buffer.precision, output_idx); + const char* input_offset_bytes = input_data + PrecisionTraits::byte_offset_of(input_buffer.precision, input_base); + size_t length = PrecisionTraits::byte_offset_of(input_buffer.precision, slice_size); + std::memcpy(output_offset_bytes, input_offset_bytes, length); + + output_idx += slice_size; + } +} + +void compute_bilinear_interpolation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& pos_embeds_buffer = get_input(node, 0, nodes, node_index_map); + + size_t total_pos_embeds = pos_embeds_buffer.shape[0]; + size_t embed_dim = pos_embeds_buffer.shape[1]; + + size_t src_height = static_cast(std::sqrt(total_pos_embeds)); + size_t src_width = src_height; + + size_t dst_height = node.params.dst_height; + size_t dst_width = node.params.dst_width; + bool align_corners = node.params.align_corners; + + __fp16* output = node.output_buffer.data_as<__fp16>(); + + if (pos_embeds_buffer.precision == Precision::FP16) { + const __fp16* input = pos_embeds_buffer.data_as<__fp16>(); + cactus_bilinear_interpolation_f16(input, output, src_height, src_width, embed_dim, + dst_height, dst_width, align_corners); + } + else if (pos_embeds_buffer.precision == Precision::INT8) { + std::vector<__fp16> input_fp16(total_pos_embeds * embed_dim); + cactus_int8_to_fp16(pos_embeds_buffer.data_as(), input_fp16.data(), + total_pos_embeds * embed_dim); + cactus_bilinear_interpolation_f16(input_fp16.data(), output, src_height, src_width, embed_dim, + dst_height, dst_width, align_corners); + } + else { + throw std::runtime_error("BILINEAR_INTERPOLATION only supports INT8 and FP16 input precision"); + } +} + +void compute_persistent_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + if (node.input_ids.empty()) { + return; + } + + auto it = node_index_map.find(node.input_ids[0]); + + if (it != node_index_map.end()) { + const auto& input_buffer = nodes[it->second]->output_buffer; + + if (!node.output_buffer.get_data()) { + node.output_buffer.allocate(); + } + + std::memcpy(node.output_buffer.get_data(), + input_buffer.get_data(), + input_buffer.byte_size); + } else { + if (node.output_buffer.get_data()) { + return; + } + throw std::runtime_error("PERSISTENT node input not found and not populated - this should not happen"); + } +} diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp new file mode 100644 index 000000000..e698638f7 --- /dev/null +++ b/cactus-graph/src/param_io.cpp @@ -0,0 +1,409 @@ +#include "param_io.h" + +#include +#include +#include +#include + +namespace { + +void write_u32(std::ostream& out, uint32_t v) { + out.write(reinterpret_cast(&v), sizeof(v)); +} + +void write_u64(std::ostream& out, uint64_t v) { + out.write(reinterpret_cast(&v), sizeof(v)); +} + +void write_i32(std::ostream& out, int32_t v) { + out.write(reinterpret_cast(&v), sizeof(v)); +} + +void write_f32(std::ostream& out, float v) { + out.write(reinterpret_cast(&v), sizeof(v)); +} + +uint32_t read_u32(std::istream& in) { + uint32_t v = 0; + in.read(reinterpret_cast(&v), sizeof(v)); + if (!in) throw std::runtime_error("Unexpected EOF while reading uint32"); + return v; +} + +uint64_t read_u64(std::istream& in) { + uint64_t v = 0; + in.read(reinterpret_cast(&v), sizeof(v)); + if (!in) throw std::runtime_error("Unexpected EOF while reading uint64"); + return v; +} + +int32_t read_i32(std::istream& in) { + int32_t v = 0; + in.read(reinterpret_cast(&v), sizeof(v)); + if (!in) throw std::runtime_error("Unexpected EOF while reading int32"); + return v; +} + +float read_f32(std::istream& in) { + float v = 0.0f; + in.read(reinterpret_cast(&v), sizeof(v)); + if (!in) throw std::runtime_error("Unexpected EOF while reading float"); + return v; +} + +void write_size_vector(std::ostream& out, const std::vector& values) { + write_u32(out, static_cast(values.size())); + for (size_t v : values) write_u64(out, static_cast(v)); +} + +std::vector read_size_vector(std::istream& in) { + uint32_t count = read_u32(in); + std::vector values; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) values.push_back(static_cast(read_u64(in))); + return values; +} + +void write_u32_vector(std::ostream& out, const std::vector& values) { + write_u32(out, static_cast(values.size())); + for (uint32_t v : values) write_u32(out, v); +} + +std::vector read_u32_vector(std::istream& in) { + uint32_t count = read_u32(in); + std::vector values; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) values.push_back(read_u32(in)); + return values; +} + +void write_f32_vector(std::ostream& out, const std::vector& values) { + write_u32(out, static_cast(values.size())); + for (float v : values) write_f32(out, v); +} + +std::vector read_f32_vector(std::istream& in) { + uint32_t count = read_u32(in); + std::vector values; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) values.push_back(read_f32(in)); + return values; +} + +enum class ParamField : uint32_t { + Scalar = 1, + Axis, + NewShape, + PretransposedRhs, + Backend, + SliceStart, + SliceLength, + Epsilon, + NumGroups, + IndexValue, + Permutation, + Scale, + Theta, + PositionOffset, + WindowSize, + IsCausal, + AttentionMaskIsAdditive, + LogitCap, + TopK, + Temperature, + TopP, + MinP, + RepetitionPenalty, + RandomSeed, + BiasIndices, + BiasValues, + NormalizeRouting, + NumExperts, + NumExpertsPerTok, + MoeGated, + Activation, + NumKvHeads, + ChunkSize, + DstHeight, + DstWidth, + AlignCorners, + NumClasses, + NumAltupInputs, + KernelSize, + Stride, + Dilation, + NumFftBins, + CacheSeqLen, + HeadDim, + VHeadDim, + CachedKeysInt8Ptr, + CachedValuesInt8Ptr, + CachedKScalesPtr, + CachedVScalesPtr, + MaxCacheSeqLen, + CacheSinkSize, + CacheNumSlots, +}; + +enum class FieldPersistence { + Persistent, + RuntimeOnly, + Derived, +}; + +struct FieldSpec { + ParamField field; + FieldPersistence persistence; +}; + +using Schema = std::vector; + +const Schema& op_schema(OpType op_type) { + static const Schema kEmpty{}; + static const std::unordered_map kSchemas = { + {OpType::POW, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_ADD, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_SUBTRACT, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_MULTIPLY, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_DIVIDE, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::SCALAR_NOT_EQUAL, {{ParamField::Scalar, FieldPersistence::Persistent}}}, + {OpType::CLAMP, {{ParamField::Scalar, FieldPersistence::Persistent}, {ParamField::Scale, FieldPersistence::Persistent}}}, + {OpType::SOFTMAX, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::SUM, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::MEAN, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::VARIANCE, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::MIN, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::MAX, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::CUMSUM, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::INDEX, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::IndexValue, FieldPersistence::Persistent}}}, + {OpType::CONCAT, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::CAT, {{ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::VIEW, {{ParamField::NewShape, FieldPersistence::Persistent}}}, + {OpType::RESHAPE, {{ParamField::NewShape, FieldPersistence::Persistent}}}, + {OpType::FLATTEN, {{ParamField::NewShape, FieldPersistence::Persistent}}}, + {OpType::MATMUL, {{ParamField::PretransposedRhs, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, + {OpType::TRANSPOSE, {{ParamField::Permutation, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, + {OpType::SLICE, {{ParamField::Axis, FieldPersistence::Persistent}, {ParamField::SliceStart, FieldPersistence::Persistent}, {ParamField::SliceLength, FieldPersistence::Persistent}}}, + {OpType::RMS_NORM, {{ParamField::Epsilon, FieldPersistence::Persistent}}}, + {OpType::LAYERNORM, {{ParamField::Epsilon, FieldPersistence::Persistent}}}, + {OpType::GROUPNORM, {{ParamField::Epsilon, FieldPersistence::Persistent}, {ParamField::NumGroups, FieldPersistence::Persistent}}}, + {OpType::BATCHNORM, {{ParamField::Epsilon, FieldPersistence::Persistent}, {ParamField::Axis, FieldPersistence::Persistent}}}, + {OpType::ROPE, {{ParamField::Theta, FieldPersistence::Persistent}, {ParamField::PositionOffset, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, + {OpType::ROPE_GPTJ, {{ParamField::Theta, FieldPersistence::Persistent}, {ParamField::PositionOffset, FieldPersistence::Persistent}, {ParamField::Scalar, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, + {OpType::TOPK, {{ParamField::TopK, FieldPersistence::Persistent}}}, + {OpType::ATTENTION, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::PositionOffset, FieldPersistence::Persistent}, {ParamField::WindowSize, FieldPersistence::Persistent}, {ParamField::IsCausal, FieldPersistence::Persistent}, {ParamField::AttentionMaskIsAdditive, FieldPersistence::Persistent}, {ParamField::LogitCap, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}}, + {OpType::REL_POS_BIAS, {{ParamField::Scale, FieldPersistence::Persistent}}}, + {OpType::ATTENTION_INT8_HYBRID, { + {ParamField::Scale, FieldPersistence::Persistent}, + {ParamField::PositionOffset, FieldPersistence::Persistent}, + {ParamField::WindowSize, FieldPersistence::Persistent}, + {ParamField::NumKvHeads, FieldPersistence::Persistent}, + {ParamField::CacheSeqLen, FieldPersistence::Persistent}, + {ParamField::HeadDim, FieldPersistence::Persistent}, + {ParamField::VHeadDim, FieldPersistence::Persistent}, + {ParamField::CachedKeysInt8Ptr, FieldPersistence::RuntimeOnly}, + {ParamField::CachedValuesInt8Ptr, FieldPersistence::RuntimeOnly}, + {ParamField::CachedKScalesPtr, FieldPersistence::RuntimeOnly}, + {ParamField::CachedVScalesPtr, FieldPersistence::RuntimeOnly}, + }}, + {OpType::KV_CACHE_STATE, { + {ParamField::MaxCacheSeqLen, FieldPersistence::Persistent}, + {ParamField::NumKvHeads, FieldPersistence::Persistent}, + {ParamField::HeadDim, FieldPersistence::Persistent}, + {ParamField::WindowSize, FieldPersistence::Persistent}, + {ParamField::CacheSinkSize, FieldPersistence::Persistent}, + {ParamField::CacheNumSlots, FieldPersistence::Persistent}, + }}, + {OpType::KV_CACHE_APPEND, { + {ParamField::WindowSize, FieldPersistence::Persistent}, + {ParamField::CacheSinkSize, FieldPersistence::Persistent}, + }}, + {OpType::ATTENTION_CACHED, { + {ParamField::Scale, FieldPersistence::Persistent}, + {ParamField::PositionOffset, FieldPersistence::Persistent}, + {ParamField::WindowSize, FieldPersistence::Persistent}, + {ParamField::VHeadDim, FieldPersistence::Persistent}, + }}, + {OpType::CONV_CACHE_STATE, { + {ParamField::WindowSize, FieldPersistence::Persistent}, + {ParamField::HeadDim, FieldPersistence::Persistent}, + }}, + {OpType::MOE_LAYER, {{ParamField::Scalar, FieldPersistence::Persistent}, {ParamField::Epsilon, FieldPersistence::Persistent}, {ParamField::NormalizeRouting, FieldPersistence::Persistent}, {ParamField::NumExperts, FieldPersistence::Persistent}, {ParamField::NumExpertsPerTok, FieldPersistence::Persistent}, {ParamField::MoeGated, FieldPersistence::Persistent}, {ParamField::Activation, FieldPersistence::Persistent}}}, + {OpType::GATED_DELTANET_DECODE, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::NumKvHeads, FieldPersistence::Persistent}}}, + {OpType::GATED_DELTANET_PREFILL, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::NumKvHeads, FieldPersistence::Persistent}, {ParamField::ChunkSize, FieldPersistence::Persistent}}}, + {OpType::STFT, {{ParamField::Stride, FieldPersistence::Persistent}, {ParamField::NumFftBins, FieldPersistence::Persistent}}}, + {OpType::BILINEAR_INTERPOLATION, {{ParamField::DstHeight, FieldPersistence::Persistent}, {ParamField::DstWidth, FieldPersistence::Persistent}, {ParamField::AlignCorners, FieldPersistence::Persistent}}}, + {OpType::SCATTER_TOPK, {{ParamField::NumClasses, FieldPersistence::Persistent}}}, + {OpType::ALTUP_PREDICT, {{ParamField::NumAltupInputs, FieldPersistence::Persistent}}}, + {OpType::ALTUP_CORRECT, {{ParamField::NumAltupInputs, FieldPersistence::Persistent}}}, + {OpType::MAXPOOL1D, {{ParamField::KernelSize, FieldPersistence::Persistent}, {ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::CONV1D_CAUSAL, {{ParamField::Dilation, FieldPersistence::Persistent}}}, + {OpType::CONV1D_K3, {{ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::CONV1D_K7S3, {{ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::CONV1D, {{ParamField::Stride, FieldPersistence::Persistent}}}, + {OpType::SAMPLE, {{ParamField::Temperature, FieldPersistence::Persistent}, {ParamField::TopP, FieldPersistence::Persistent}, {ParamField::MinP, FieldPersistence::Persistent}, {ParamField::RepetitionPenalty, FieldPersistence::Persistent}, {ParamField::TopK, FieldPersistence::Persistent}, {ParamField::RandomSeed, FieldPersistence::Persistent}, {ParamField::BiasIndices, FieldPersistence::Persistent}, {ParamField::BiasValues, FieldPersistence::Persistent}}}, + }; + + auto it = kSchemas.find(op_type); + return it == kSchemas.end() ? kEmpty : it->second; +} + +void write_field(std::ostream& out, ParamField field, const OpParams& params) { + switch (field) { + case ParamField::Scalar: write_f32(out, params.scalar); break; + case ParamField::Axis: write_i32(out, static_cast(params.axis)); break; + case ParamField::NewShape: write_size_vector(out, params.new_shape); break; + case ParamField::PretransposedRhs: write_u32(out, params.pretransposed_rhs ? 1u : 0u); break; + case ParamField::SliceStart: write_u64(out, static_cast(params.slice_start)); break; + case ParamField::SliceLength: write_u64(out, static_cast(params.slice_length)); break; + case ParamField::Epsilon: write_f32(out, params.epsilon); break; + case ParamField::NumGroups: write_u64(out, static_cast(params.num_groups)); break; + case ParamField::IndexValue: write_u64(out, static_cast(params.index_value)); break; + case ParamField::Permutation: write_size_vector(out, params.permutation); break; + case ParamField::Scale: write_f32(out, params.scale); break; + case ParamField::Theta: write_f32(out, params.theta); break; + case ParamField::PositionOffset: write_u64(out, static_cast(params.position_offset)); break; + case ParamField::WindowSize: write_u64(out, static_cast(params.window_size)); break; + case ParamField::IsCausal: write_u32(out, params.is_causal ? 1u : 0u); break; + case ParamField::AttentionMaskIsAdditive: write_u32(out, params.attention_mask_is_additive ? 1u : 0u); break; + case ParamField::LogitCap: write_f32(out, params.logit_cap); break; + case ParamField::TopK: write_u64(out, static_cast(params.top_k)); break; + case ParamField::Temperature: write_f32(out, params.temperature); break; + case ParamField::TopP: write_f32(out, params.top_p); break; + case ParamField::MinP: write_f32(out, params.min_p); break; + case ParamField::RepetitionPenalty: write_f32(out, params.repetition_penalty); break; + case ParamField::RandomSeed: write_u64(out, static_cast(params.random_seed)); break; + case ParamField::BiasIndices: write_u32_vector(out, params.bias_indices); break; + case ParamField::BiasValues: write_f32_vector(out, params.bias_values); break; + case ParamField::NormalizeRouting: write_u32(out, params.normalize_routing ? 1u : 0u); break; + case ParamField::NumExperts: write_u64(out, static_cast(params.num_experts)); break; + case ParamField::NumExpertsPerTok: write_u64(out, static_cast(params.num_experts_per_tok)); break; + case ParamField::MoeGated: write_u32(out, params.moe_gated ? 1u : 0u); break; + case ParamField::Activation: write_u32(out, static_cast(params.activation)); break; + case ParamField::NumKvHeads: write_u64(out, static_cast(params.num_kv_heads)); break; + case ParamField::ChunkSize: write_u64(out, static_cast(params.chunk_size)); break; + case ParamField::DstHeight: write_u64(out, static_cast(params.dst_height)); break; + case ParamField::DstWidth: write_u64(out, static_cast(params.dst_width)); break; + case ParamField::AlignCorners: write_u32(out, params.align_corners ? 1u : 0u); break; + case ParamField::NumClasses: write_u64(out, static_cast(params.num_classes)); break; + case ParamField::NumAltupInputs: write_u64(out, static_cast(params.num_altup_inputs)); break; + case ParamField::KernelSize: write_u64(out, static_cast(params.kernel_size)); break; + case ParamField::Stride: write_u64(out, static_cast(params.stride)); break; + case ParamField::Dilation: write_u64(out, static_cast(params.dilation)); break; + case ParamField::NumFftBins: write_u64(out, static_cast(params.num_fft_bins)); break; + case ParamField::CacheSeqLen: write_u64(out, static_cast(params.cache_seq_len)); break; + case ParamField::HeadDim: write_u64(out, static_cast(params.head_dim)); break; + case ParamField::VHeadDim: write_u64(out, static_cast(params.v_head_dim)); break; + case ParamField::Backend: + case ParamField::CachedKeysInt8Ptr: + case ParamField::CachedValuesInt8Ptr: + case ParamField::CachedKScalesPtr: + case ParamField::CachedVScalesPtr: + throw std::runtime_error("Attempted to serialize runtime-only field"); + case ParamField::MaxCacheSeqLen: write_u64(out, static_cast(params.max_cache_seq_len)); break; + case ParamField::CacheSinkSize: write_u64(out, static_cast(params.cache_sink_size)); break; + case ParamField::CacheNumSlots: write_u64(out, static_cast(params.cache_num_slots)); break; + } +} + +void read_field(std::istream& in, ParamField field, OpParams& params) { + switch (field) { + case ParamField::Scalar: params.scalar = read_f32(in); break; + case ParamField::Axis: params.axis = static_cast(read_i32(in)); break; + case ParamField::NewShape: params.new_shape = read_size_vector(in); break; + case ParamField::PretransposedRhs: params.pretransposed_rhs = (read_u32(in) != 0); break; + case ParamField::Backend: + read_u32(in); + break; + case ParamField::SliceStart: params.slice_start = static_cast(read_u64(in)); break; + case ParamField::SliceLength: params.slice_length = static_cast(read_u64(in)); break; + case ParamField::Epsilon: params.epsilon = read_f32(in); break; + case ParamField::NumGroups: params.num_groups = static_cast(read_u64(in)); break; + case ParamField::IndexValue: params.index_value = static_cast(read_u64(in)); break; + case ParamField::Permutation: params.permutation = read_size_vector(in); break; + case ParamField::Scale: params.scale = read_f32(in); break; + case ParamField::Theta: params.theta = read_f32(in); break; + case ParamField::PositionOffset: params.position_offset = static_cast(read_u64(in)); break; + case ParamField::WindowSize: params.window_size = static_cast(read_u64(in)); break; + case ParamField::IsCausal: params.is_causal = (read_u32(in) != 0); break; + case ParamField::AttentionMaskIsAdditive: params.attention_mask_is_additive = (read_u32(in) != 0); break; + case ParamField::LogitCap: params.logit_cap = read_f32(in); break; + case ParamField::TopK: params.top_k = static_cast(read_u64(in)); break; + case ParamField::Temperature: params.temperature = read_f32(in); break; + case ParamField::TopP: params.top_p = read_f32(in); break; + case ParamField::MinP: params.min_p = read_f32(in); break; + case ParamField::RepetitionPenalty: params.repetition_penalty = read_f32(in); break; + case ParamField::RandomSeed: params.random_seed = static_cast(read_u64(in)); break; + case ParamField::BiasIndices: params.bias_indices = read_u32_vector(in); break; + case ParamField::BiasValues: params.bias_values = read_f32_vector(in); break; + case ParamField::NormalizeRouting: params.normalize_routing = (read_u32(in) != 0); break; + case ParamField::NumExperts: params.num_experts = static_cast(read_u64(in)); break; + case ParamField::NumExpertsPerTok: params.num_experts_per_tok = static_cast(read_u64(in)); break; + case ParamField::MoeGated: params.moe_gated = (read_u32(in) != 0); break; + case ParamField::Activation: params.activation = static_cast(read_u32(in)); break; + case ParamField::NumKvHeads: params.num_kv_heads = static_cast(read_u64(in)); break; + case ParamField::ChunkSize: params.chunk_size = static_cast(read_u64(in)); break; + case ParamField::DstHeight: params.dst_height = static_cast(read_u64(in)); break; + case ParamField::DstWidth: params.dst_width = static_cast(read_u64(in)); break; + case ParamField::AlignCorners: params.align_corners = (read_u32(in) != 0); break; + case ParamField::NumClasses: params.num_classes = static_cast(read_u64(in)); break; + case ParamField::NumAltupInputs: params.num_altup_inputs = static_cast(read_u64(in)); break; + case ParamField::KernelSize: params.kernel_size = static_cast(read_u64(in)); break; + case ParamField::Stride: params.stride = static_cast(read_u64(in)); break; + case ParamField::Dilation: params.dilation = static_cast(read_u64(in)); break; + case ParamField::NumFftBins: params.num_fft_bins = static_cast(read_u64(in)); break; + case ParamField::CacheSeqLen: params.cache_seq_len = static_cast(read_u64(in)); break; + case ParamField::HeadDim: params.head_dim = static_cast(read_u64(in)); break; + case ParamField::VHeadDim: params.v_head_dim = static_cast(read_u64(in)); break; + case ParamField::CachedKeysInt8Ptr: + case ParamField::CachedValuesInt8Ptr: + case ParamField::CachedKScalesPtr: + case ParamField::CachedVScalesPtr: + throw std::runtime_error("Graph file corrupted: runtime-only field serialized"); + case ParamField::MaxCacheSeqLen: params.max_cache_seq_len = static_cast(read_u64(in)); break; + case ParamField::CacheSinkSize: params.cache_sink_size = static_cast(read_u64(in)); break; + case ParamField::CacheNumSlots: params.cache_num_slots = static_cast(read_u64(in)); break; + } +} + +} // namespace + +namespace GraphParamIO { + +void write_op_params(std::ostream& out, OpType op_type, const OpParams& params) { + const auto& schema = op_schema(op_type); + std::vector fields; + fields.reserve(schema.size()); + for (const auto& spec : schema) { + if (spec.persistence == FieldPersistence::Persistent && spec.field != ParamField::Backend) { + fields.push_back(spec.field); + } + } + + write_u32(out, static_cast(fields.size())); + for (ParamField field : fields) { + write_u32(out, static_cast(field)); + write_field(out, field, params); + } +} + +void read_op_params(std::istream& in, OpType op_type, OpParams& params) { + uint32_t field_count = read_u32(in); + const auto& schema = op_schema(op_type); + std::unordered_set allowed; + for (const auto& spec : schema) { + if (spec.persistence == FieldPersistence::Persistent) { + allowed.insert(static_cast(spec.field)); + } + } + + for (uint32_t i = 0; i < field_count; ++i) { + uint32_t raw_field = read_u32(in); + if (allowed.count(raw_field) == 0) { + throw std::runtime_error("Graph file corrupted: field not allowed for op"); + } + read_field(in, static_cast(raw_field), params); + } +} + +} // namespace GraphParamIO diff --git a/cactus-graph/src/param_io.h b/cactus-graph/src/param_io.h new file mode 100644 index 000000000..93fb575ea --- /dev/null +++ b/cactus-graph/src/param_io.h @@ -0,0 +1,11 @@ +#pragma once + +#include "../cactus_graph.h" +#include + +namespace GraphParamIO { + +void write_op_params(std::ostream& out, OpType op_type, const OpParams& params); +void read_op_params(std::istream& in, OpType op_type, OpParams& params); + +} // namespace GraphParamIO diff --git a/cactus-graph/test.sh b/cactus-graph/test.sh new file mode 100755 index 000000000..4bf78b6d8 --- /dev/null +++ b/cactus-graph/test.sh @@ -0,0 +1,40 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")" + +SUITE="" +while [[ $# -gt 0 ]]; do + case $1 in + --suite) SUITE="${2:?--suite needs an argument}"; shift 2 ;; + *) echo "Unknown arg: $1" >&2; exit 2 ;; + esac +done + +echo "Building and testing cactus-graph..." + +rm -rf build +mkdir -p build +cd build + +cmake .. -DCACTUS_BUILD_TESTS=ON -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null +make -j$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) + +echo "" +FAILED=0 +if [ -n "$SUITE" ]; then + target="./test_${SUITE}" + if [ -x "$target" ]; then + "$target" || FAILED=1 + else + echo "Test not found: $target" >&2 + FAILED=1 + fi +else + for t in ./test_*; do + [ -x "$t" ] || continue + "$t" || FAILED=1 + done +fi + +exit $FAILED diff --git a/cactus-graph/tests/test_cache.cpp b/cactus-graph/tests/test_cache.cpp new file mode 100644 index 000000000..1fe297041 --- /dev/null +++ b/cactus-graph/tests/test_cache.cpp @@ -0,0 +1,1137 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +using namespace TestUtils; + +bool test_kv_cache_state_init() { + CactusGraph g; + + const size_t max_seq = 64, kv_heads = 4, head_dim = 16; + size_t cache_node = g.kv_cache_state(max_seq, kv_heads, head_dim); + g.execute(); + + auto* raw = static_cast(g.get_output(cache_node)); + if (!raw) return false; + + uint64_t current_seq = *reinterpret_cast(raw + 0); + uint64_t stored_max = *reinterpret_cast(raw + 8); + uint64_t stored_kv = *reinterpret_cast(raw + 16); + uint64_t stored_hdim = *reinterpret_cast(raw + 24); + + if (current_seq != 0) return false; + if (stored_max != max_seq) return false; + if (stored_kv != kv_heads) return false; + if (stored_hdim != head_dim) return false; + + return true; +} + +bool test_kv_cache_state_persistent() { + CactusGraph g; + + size_t cache_node = g.kv_cache_state(32, 2, 16); + g.execute(); + + auto* raw = static_cast(g.get_output(cache_node)); + if (!raw) return false; + + g.soft_reset(); + + if (!g.is_populated(cache_node)) return false; + + return true; +} + +bool test_kv_cache_append_basic() { + CactusGraph g; + + const size_t max_seq = 64, kv_heads = 2, head_dim = 16; + const size_t new_tokens = 3; + const size_t kv_elements = new_tokens * kv_heads * head_dim; + + size_t cache_node = g.kv_cache_state(max_seq, kv_heads, head_dim); + + size_t kv_input = g.input({kv_elements}, Precision::FP16); + std::vector<__fp16> kv_data(kv_elements); + for (size_t i = 0; i < kv_elements; i++) { + kv_data[i] = static_cast<__fp16>(static_cast(i) * 0.1f); + } + g.set_input(kv_input, kv_data.data(), Precision::FP16); + + size_t append_result = g.kv_cache_append(kv_input, cache_node); + g.execute(); + + float* result = static_cast(g.get_output(append_result)); + if (!result) return false; + if (static_cast(*result) != new_tokens) return false; + + auto* raw = static_cast(g.get_output(cache_node)); + uint64_t current_seq = *reinterpret_cast(raw + 0); + if (current_seq != new_tokens) return false; + + return true; +} + +bool test_kv_cache_append_multiple() { + CactusGraph g; + + const size_t max_seq = 64, kv_heads = 2, head_dim = 16; + size_t cache_node = g.kv_cache_state(max_seq, kv_heads, head_dim); + + { + const size_t tokens = 2; + const size_t elements = tokens * kv_heads * head_dim; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements, static_cast<__fp16>(1.0f)); + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node); + g.execute(); + } + + auto* raw = static_cast(g.get_output(cache_node)); + if (*reinterpret_cast(raw) != 2) return false; + + g.soft_reset(); + { + const size_t tokens = 3; + const size_t elements = tokens * kv_heads * head_dim; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements, static_cast<__fp16>(2.0f)); + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node); + g.execute(); + } + + raw = static_cast(g.get_output(cache_node)); + uint64_t total_seq = *reinterpret_cast(raw); + if (total_seq != 5) return false; + + return true; +} + +bool test_kv_cache_append_eviction() { + CactusGraph g; + + const size_t kv_heads = 1, head_dim = 16; + const size_t window = 8, sink = 2; + + size_t cache_node = g.kv_cache_state(window, kv_heads, head_dim, window, sink); + + { + const size_t tokens = 8; + const size_t elements = tokens * kv_heads * head_dim; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements); + for (size_t t = 0; t < tokens; t++) { + for (size_t j = 0; j < kv_heads * head_dim; j++) { + data[t * kv_heads * head_dim + j] = static_cast<__fp16>(static_cast(t)); + } + } + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node, window, sink); + g.execute(); + } + + auto* raw = static_cast(g.get_output(cache_node)); + uint64_t seq = *reinterpret_cast(raw); + if (seq != 8) return false; + + g.soft_reset(); + { + const size_t tokens = 2; + const size_t elements = tokens * kv_heads * head_dim; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements); + for (size_t t = 0; t < tokens; t++) { + for (size_t j = 0; j < kv_heads * head_dim; j++) { + data[t * kv_heads * head_dim + j] = static_cast<__fp16>(100.0f + static_cast(t)); + } + } + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node, window, sink); + g.execute(); + } + + raw = static_cast(g.get_output(cache_node)); + seq = *reinterpret_cast(raw); + if (seq != window) return false; + + return true; +} + +bool test_kv_cache_append_full_window_eviction() { + CactusGraph g; + + const size_t kv_heads = 1, head_dim = 16; + const size_t window = 8, sink = 2; + + size_t cache_node = g.kv_cache_state(window, kv_heads, head_dim, window, sink); + + for (int pass = 0; pass < 2; pass++) { + if (pass > 0) g.soft_reset(); + const size_t tokens = window; + const size_t elements = tokens * kv_heads * head_dim; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements, static_cast<__fp16>(static_cast(pass + 1))); + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node, window, sink); + g.execute(); + + auto* raw = static_cast(g.get_output(cache_node)); + uint64_t seq = *reinterpret_cast(raw); + if (seq != window) return false; + } + + return true; +} + +static bool padded_rollback_matches_exact_append(size_t prefix, size_t real, size_t pads) { + const size_t kv_heads = 1, head_dim = 16, max_seq = 64, window = 16, sink = 2; + const size_t stride = kv_heads * head_dim; + + auto append = [&](CactusGraph& g, size_t cache_node, size_t real_tokens, size_t pad_tokens, float base) { + const size_t tokens = real_tokens + pad_tokens; + const size_t elements = tokens * stride; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements); + for (size_t t = 0; t < tokens; t++) { + float value = t < real_tokens ? base + static_cast(t) : 9999.0f; + for (size_t j = 0; j < stride; j++) data[t * stride + j] = static_cast<__fp16>(value); + } + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node, window, sink); + g.execute(); + g.soft_reset(); + }; + + CactusGraph g_exact, g_padded; + size_t exact_node = g_exact.kv_cache_state(max_seq, kv_heads, head_dim, window, sink); + size_t padded_node = g_padded.kv_cache_state(max_seq, kv_heads, head_dim, window, sink); + if (prefix > 0) { + append(g_exact, exact_node, prefix, 0, 1.0f); + append(g_padded, padded_node, prefix, 0, 1.0f); + } + append(g_exact, exact_node, real, 0, 100.0f); + auto backup = g_padded.snapshot_cache_padded_append(padded_node, real, pads); + append(g_padded, padded_node, real, pads, 100.0f); + g_padded.rollback_cache_padded_append(padded_node, real, pads, backup); + + auto* exact_raw = static_cast(g_exact.get_output(exact_node)); + auto* padded_raw = static_cast(g_padded.get_output(padded_node)); + uint64_t exact_len = *reinterpret_cast(exact_raw); + uint64_t padded_len = *reinterpret_cast(padded_raw); + if (exact_len != padded_len) { + std::cerr << " current_seq_len mismatch: " << exact_len << " != " << padded_len << "\n"; + return false; + } + uint64_t buffer_max_seq = *reinterpret_cast(exact_raw + 8); + const size_t meta_bytes = 64; + if (std::memcmp(exact_raw + meta_bytes, padded_raw + meta_bytes, exact_len * stride) != 0) { + std::cerr << " cache data rows differ after rollback\n"; + return false; + } + const size_t num_groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + const size_t scales_offset = meta_bytes + buffer_max_seq * stride; + if (std::memcmp(exact_raw + scales_offset, padded_raw + scales_offset, + exact_len * kv_heads * num_groups * sizeof(float)) != 0) { + std::cerr << " cache scale rows differ after rollback\n"; + return false; + } + return true; +} + +bool test_padded_rollback_no_eviction() { + return padded_rollback_matches_exact_append(4, 3, 5); +} + +bool test_padded_rollback_eviction_overshoot() { + return padded_rollback_matches_exact_append(16, 3, 6); +} + +bool test_padded_rollback_only_pads_evict() { + return padded_rollback_matches_exact_append(10, 2, 8); +} + +bool test_padded_rollback_empty_cache() { + return padded_rollback_matches_exact_append(0, 5, 11); +} + +bool test_attention_cached_basic() { + const size_t b = 1, s = 1, h = 2, kv = 2, d = 16; + const size_t max_seq = 64; + + CactusGraph g; + + size_t k_cache = g.kv_cache_state(max_seq, kv, d); + size_t v_cache = g.kv_cache_state(max_seq, kv, d); + + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + + std::vector<__fp16> q(b * s * h * d), k_new(b * s * kv * d), v_new(b * s * kv * d); + fill_random_fp16(q); + fill_random_fp16(k_new); + fill_random_fp16(v_new); + + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k_new.data(), Precision::FP16); + g.set_input(iv, v_new.data(), Precision::FP16); + + g.kv_cache_append(ik, k_cache); + g.kv_cache_append(iv, v_cache); + + float scale = 1.0f / std::sqrt(static_cast(d)); + size_t attn = g.attention_cached(iq, ik, iv, k_cache, v_cache, scale, 0); + + g.execute(); + + __fp16* result = static_cast<__fp16*>(g.get_output(attn)); + size_t out_size = b * s * h * d; + + bool has_nonzero = false; + for (size_t i = 0; i < out_size; i++) { + if (!std::isfinite(static_cast(result[i]))) return false; + if (std::abs(static_cast(result[i])) > 1e-6f) has_nonzero = true; + } + return has_nonzero; +} + +bool test_attention_cached_multistep() { + const size_t b = 1, h = 2, kv = 2, d = 16; + const size_t max_seq = 64; + + CactusGraph g; + + size_t k_cache = g.kv_cache_state(max_seq, kv, d); + size_t v_cache = g.kv_cache_state(max_seq, kv, d); + + { + const size_t s = 4; + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + + std::vector<__fp16> q(b*s*h*d), k(b*s*kv*d), v(b*s*kv*d); + fill_random_fp16(q); + fill_random_fp16(k); + fill_random_fp16(v); + + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k.data(), Precision::FP16); + g.set_input(iv, v.data(), Precision::FP16); + + g.kv_cache_append(ik, k_cache); + g.kv_cache_append(iv, v_cache); + g.attention_cached(iq, ik, iv, k_cache, v_cache, + 1.0f / std::sqrt(static_cast(d)), 0); + g.execute(); + } + + auto* raw = static_cast(g.get_output(k_cache)); + if (*reinterpret_cast(raw) != 4) return false; + + g.soft_reset(); + { + const size_t s = 1; + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + + std::vector<__fp16> q(b*s*h*d), k(b*s*kv*d), v(b*s*kv*d); + fill_random_fp16(q); + fill_random_fp16(k); + fill_random_fp16(v); + + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k.data(), Precision::FP16); + g.set_input(iv, v.data(), Precision::FP16); + + g.kv_cache_append(ik, k_cache); + g.kv_cache_append(iv, v_cache); + size_t attn = g.attention_cached(iq, ik, iv, k_cache, v_cache, + 1.0f / std::sqrt(static_cast(d)), 4); + g.execute(); + + __fp16* result = static_cast<__fp16*>(g.get_output(attn)); + for (size_t i = 0; i < b*s*h*d; i++) { + if (!std::isfinite(static_cast(result[i]))) return false; + } + } + + raw = static_cast(g.get_output(k_cache)); + if (*reinterpret_cast(raw) != 5) return false; + + return true; +} + +bool test_kv_cache_invalidate() { + CactusGraph g; + + const size_t max_seq = 32, kv_heads = 2, head_dim = 16; + size_t cache_node = g.kv_cache_state(max_seq, kv_heads, head_dim); + + { + const size_t tokens = 4; + const size_t elements = tokens * kv_heads * head_dim; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements, static_cast<__fp16>(1.0f)); + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node); + g.execute(); + } + + if (!g.is_populated(cache_node)) return false; + + g.invalidate_persistent(cache_node); + if (g.is_populated(cache_node)) return false; + + g.soft_reset(); + size_t new_cache = g.kv_cache_state(max_seq, kv_heads, head_dim); + g.execute(); + + auto* raw = static_cast(g.get_output(new_cache)); + uint64_t seq = *reinterpret_cast(raw); + if (seq != 0) return false; + + return true; +} + +bool test_conv_cache_state_init() { + CactusGraph g; + + const size_t ws = 8, hd = 32; + size_t cache = g.conv_cache_state(ws, hd); + g.execute(); + + auto* raw = static_cast(g.get_output(cache)); + if (!raw) return false; + + uint64_t head = *reinterpret_cast(raw + 0); + uint64_t count = *reinterpret_cast(raw + 8); + uint64_t stored_ws = *reinterpret_cast(raw + 16); + uint64_t stored_hd = *reinterpret_cast(raw + 24); + + if (head != 0 || count != 0) return false; + if (stored_ws != ws || stored_hd != hd) return false; + + return true; +} + +bool test_conv_cache_append_basic() { + CactusGraph g; + + const size_t ws = 4, hd = 8; + size_t cache = g.conv_cache_state(ws, hd); + + size_t inp = g.input({2, hd}, Precision::FP16); + std::vector<__fp16> data(2 * hd); + for (size_t i = 0; i < 2 * hd; i++) data[i] = static_cast<__fp16>(static_cast(i)); + g.set_input(inp, data.data(), Precision::FP16); + + size_t window_out = g.conv_cache_append(inp, cache); + g.execute(); + + const auto& buf = g.get_output_buffer(window_out); + if (buf.shape[0] != ws || buf.shape[1] != hd) return false; + + __fp16* out = static_cast<__fp16*>(g.get_output(window_out)); + + const size_t pad_rows = ws - 2; + for (size_t i = 0; i < pad_rows * hd; i++) { + if (std::abs(static_cast(out[i])) > 1e-3f) return false; + } + for (size_t i = 0; i < 2 * hd; i++) { + if (std::abs(static_cast(out[pad_rows * hd + i]) - static_cast(data[i])) > 1e-3f) return false; + } + + return true; +} + +bool test_conv_cache_append_circular() { + CactusGraph g; + + const size_t ws = 3, hd = 4; + size_t cache = g.conv_cache_state(ws, hd); + + for (int step = 0; step < 5; step++) { + if (step > 0) g.soft_reset(); + size_t inp = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> data(hd); + for (size_t j = 0; j < hd; j++) data[j] = static_cast<__fp16>(static_cast(step * 10 + j)); + g.set_input(inp, data.data(), Precision::FP16); + g.conv_cache_append(inp, cache); + g.execute(); + } + + g.soft_reset(); + size_t inp = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> data(hd, static_cast<__fp16>(99.0f)); + g.set_input(inp, data.data(), Precision::FP16); + size_t window_out = g.conv_cache_append(inp, cache); + g.execute(); + + __fp16* out = static_cast<__fp16*>(g.get_output(window_out)); + const auto& buf = g.get_output_buffer(window_out); + if (buf.shape[0] != ws) return false; + + for (size_t j = 0; j < hd; j++) { + if (std::abs(static_cast(out[(ws - 1) * hd + j]) - 99.0f) > 1e-2f) return false; + } + + return true; +} + +bool test_conv_cache_persistent() { + CactusGraph g; + + size_t cache = g.conv_cache_state(4, 8); + g.execute(); + + g.soft_reset(); + if (!g.is_populated(cache)) return false; + + g.invalidate_persistent(cache); + if (g.is_populated(cache)) return false; + + return true; +} + +bool test_conv_cache_initialize_populates_trailing_rows() { + CactusGraph g; + + constexpr size_t ws = 4; + constexpr size_t hd = 8; + constexpr size_t num_rows = 6; + + size_t cache = g.conv_cache_state(ws, hd); + + size_t rows_input = g.input({num_rows, hd}, Precision::FP16); + std::vector<__fp16> data(num_rows * hd); + for (size_t r = 0; r < num_rows; ++r) { + for (size_t c = 0; c < hd; ++c) { + data[r * hd + c] = static_cast<__fp16>(static_cast(r + 1)); + } + } + g.set_input(rows_input, data.data(), Precision::FP16); + + g.conv_cache_initialize(rows_input, cache); + g.execute(); + + g.soft_reset_keep_pool(); + + size_t append_input = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> append_data(hd, static_cast<__fp16>(99.0f)); + g.set_input(append_input, append_data.data(), Precision::FP16); + + size_t window = g.conv_cache_append(append_input, cache); + g.execute(); + + const __fp16* w = static_cast(g.get_output(window)); + if (!w) return false; + const float expected[ws] = {4.0f, 5.0f, 6.0f, 99.0f}; + for (size_t r = 0; r < ws; ++r) { + for (size_t c = 0; c < hd; ++c) { + if (static_cast(w[r * hd + c]) != expected[r]) return false; + } + } + return true; +} + +bool test_conv_cache_initialize_rejects_non_state_input() { + CactusGraph g; + size_t plain_input = g.input({4, 8}, Precision::FP16); + std::vector<__fp16> zeros(32, static_cast<__fp16>(0.0f)); + g.set_input(plain_input, zeros.data(), Precision::FP16); + size_t rows_input = g.input({2, 8}, Precision::FP16); + g.set_input(rows_input, zeros.data(), Precision::FP16); + try { + g.conv_cache_initialize(rows_input, plain_input); + } catch (const std::invalid_argument&) { + return true; + } + return false; +} + +bool test_conv_cache_initialize_output_is_zero_byte() { + CactusGraph g; + size_t cache = g.conv_cache_state(/*window=*/4, /*hidden_dim=*/8); + size_t rows = g.input({2, 8}, Precision::FP16); + std::vector<__fp16> data(16, static_cast<__fp16>(1.0f)); + g.set_input(rows, data.data(), Precision::FP16); + size_t out_node = g.conv_cache_initialize(rows, cache); + const auto& desc = g.get_output_buffer(out_node); + return desc.byte_size == 0; +} + +bool test_conv_cache_initialize_resets_dirty_cache() { + CactusGraph g; + constexpr size_t ws = 4; + constexpr size_t hd = 8; + size_t cache = g.conv_cache_state(ws, hd); + + size_t dirty_input = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> dirty(hd, static_cast<__fp16>(77.0f)); + g.set_input(dirty_input, dirty.data(), Precision::FP16); + g.conv_cache_append(dirty_input, cache); + g.execute(); + + g.soft_reset_keep_pool(); + size_t rows_input = g.input({ws, hd}, Precision::FP16); + std::vector<__fp16> data(ws * hd); + for (size_t r = 0; r < ws; ++r) { + for (size_t c = 0; c < hd; ++c) { + data[r * hd + c] = static_cast<__fp16>(static_cast(r + 1)); + } + } + g.set_input(rows_input, data.data(), Precision::FP16); + g.conv_cache_initialize(rows_input, cache); + g.execute(); + + g.soft_reset_keep_pool(); + size_t append_input = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> append_data(hd, static_cast<__fp16>(99.0f)); + g.set_input(append_input, append_data.data(), Precision::FP16); + size_t window = g.conv_cache_append(append_input, cache); + g.execute(); + + const __fp16* w = static_cast(g.get_output(window)); + if (!w) return false; + const float expected[ws] = {2.0f, 3.0f, 4.0f, 99.0f}; + for (size_t r = 0; r < ws; ++r) { + for (size_t c = 0; c < hd; ++c) { + if (static_cast(w[r * hd + c]) != expected[r]) return false; + } + } + return true; +} + +bool test_recurrent_cache_write_carries_state_across_executions() { + CactusGraph g; + + const std::vector shape{2, 4}; + const size_t elements = 2 * 4; + + size_t cache_state = g.recurrent_cache_state(shape, Precision::FP16); + + size_t new_value = g.input(shape, Precision::FP16); + std::vector<__fp16> first(elements); + for (size_t i = 0; i < elements; ++i) first[i] = static_cast<__fp16>(static_cast(i + 1)); + g.set_input(new_value, first.data(), Precision::FP16); + + g.recurrent_cache_write(new_value, cache_state); + g.execute(); + + const __fp16* cache_after = static_cast(g.get_output(cache_state)); + if (!cache_after) return false; + for (size_t i = 0; i < elements; ++i) { + if (static_cast(cache_after[i]) != static_cast(first[i])) return false; + } + + g.soft_reset_keep_pool(); + size_t second_input = g.input(shape, Precision::FP16); + std::vector<__fp16> second(elements); + for (size_t i = 0; i < elements; ++i) second[i] = static_cast<__fp16>(static_cast(i + 1) * 10.0f); + g.set_input(second_input, second.data(), Precision::FP16); + g.recurrent_cache_write(second_input, cache_state); + g.execute(); + + const __fp16* cache_after2 = static_cast(g.get_output(cache_state)); + if (!cache_after2) return false; + for (size_t i = 0; i < elements; ++i) { + if (static_cast(cache_after2[i]) != static_cast(second[i])) return false; + } + + return true; +} + +bool test_recurrent_cache_state_initializes_to_zero() { + CactusGraph g; + const std::vector shape{3, 5}; + const size_t elements = 3 * 5; + size_t cache_state = g.recurrent_cache_state(shape, Precision::FP16); + g.execute(); + const __fp16* data = static_cast(g.get_output(cache_state)); + if (!data) return false; + for (size_t i = 0; i < elements; ++i) { + if (static_cast(data[i]) != 0.0f) return false; + } + return true; +} + +bool test_steal_cache_buffer_preserves_source_descriptor() { + CactusGraph src; + CactusGraph dst; + const std::vector shape{1, 16, 4, 32}; + const size_t elements = 1 * 16 * 4 * 32; + size_t src_state = src.recurrent_cache_state(shape, Precision::FP16); + size_t dst_state = dst.recurrent_cache_state(shape, Precision::FP16); + src.execute(); + dst.steal_cache_buffer(dst_state, src, src_state); + if (!dst.get_output(dst_state)) return false; + if (src.get_output_buffer(src_state).shape != shape) return false; + src.execute(); + const BufferDesc& buf = src.get_output_buffer(src_state); + const __fp16* data = static_cast(buf.get_data()); + if (!data) return false; + if (buf.byte_size != elements * sizeof(__fp16)) return false; + for (size_t i = 0; i < elements; ++i) { + if (static_cast(data[i]) != 0.0f) return false; + } + return true; +} + +bool test_recurrent_cache_write_rejects_non_state_input() { + CactusGraph g; + size_t plain_input = g.input({2, 4}, Precision::FP16); + std::vector<__fp16> zeros(8, static_cast<__fp16>(0.0f)); + g.set_input(plain_input, zeros.data(), Precision::FP16); + size_t new_value = g.input({2, 4}, Precision::FP16); + g.set_input(new_value, zeros.data(), Precision::FP16); + try { + g.recurrent_cache_write(new_value, plain_input); + } catch (const std::invalid_argument&) { + return true; + } + return false; +} + +bool test_recurrent_cache_write_rejects_shape_mismatch() { + CactusGraph g; + + size_t cache_state = g.recurrent_cache_state({2, 4}, Precision::FP16); + + size_t new_value = g.input({2, 8}, Precision::FP16); + std::vector<__fp16> data(16, static_cast<__fp16>(0.0f)); + g.set_input(new_value, data.data(), Precision::FP16); + + try { + g.recurrent_cache_write(new_value, cache_state); + } catch (const std::invalid_argument&) { + return true; + } + return false; +} + +bool test_allocate_buffers_preserves_recurrent_state_init() { + CactusGraph g; + const std::vector shape{3, 5}; + const size_t elements = 3 * 5; + size_t cache_state = g.recurrent_cache_state(shape, Precision::FP16); + g.allocate_buffers(); + g.execute(); + const __fp16* data = static_cast(g.get_output(cache_state)); + if (!data) return false; + for (size_t i = 0; i < elements; ++i) { + if (static_cast(data[i]) != 0.0f) return false; + } + return true; +} + +bool test_conv_cache_initialize_rejects_hidden_dim_mismatch() { + CactusGraph g; + size_t cache = g.conv_cache_state(/*window=*/4, /*hidden_dim=*/8); + size_t rows = g.input({2, 7}, Precision::FP16); + std::vector<__fp16> data(14, static_cast<__fp16>(0.0f)); + g.set_input(rows, data.data(), Precision::FP16); + try { + g.conv_cache_initialize(rows, cache); + } catch (const std::invalid_argument&) { + return true; + } + return false; +} + + +bool test_conv_cache_initialize_single_row() { + CactusGraph g; + constexpr size_t ws = 4, hd = 8; + size_t cache = g.conv_cache_state(ws, hd); + + size_t rows = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> data(hd, static_cast<__fp16>(42.0f)); + g.set_input(rows, data.data(), Precision::FP16); + g.conv_cache_initialize(rows, cache); + g.execute(); + + g.soft_reset_keep_pool(); + size_t append_input = g.input({1, hd}, Precision::FP16); + std::vector<__fp16> append_data(hd, static_cast<__fp16>(7.0f)); + g.set_input(append_input, append_data.data(), Precision::FP16); + size_t window = g.conv_cache_append(append_input, cache); + g.execute(); + + const __fp16* w = static_cast(g.get_output(window)); + if (!w) return false; + const float expected[ws] = {0.0f, 0.0f, 42.0f, 7.0f}; + for (size_t r = 0; r < ws; ++r) { + for (size_t c = 0; c < hd; ++c) { + if (static_cast(w[r * hd + c]) != expected[r]) return false; + } + } + return true; +} + +bool test_allocate_buffers_idempotent() { + CactusGraph g; + const std::vector shape{2, 4}; + size_t cache_state = g.recurrent_cache_state(shape, Precision::FP16); + g.allocate_buffers(); + g.allocate_buffers(); + g.execute(); + const __fp16* data = static_cast(g.get_output(cache_state)); + if (!data) return false; + for (size_t i = 0; i < 8; ++i) { + if (static_cast(data[i]) != 0.0f) return false; + } + return true; +} + +bool test_recurrent_cache_write_idempotent() { + CactusGraph g; + const std::vector shape{2, 4}; + size_t cache_state = g.recurrent_cache_state(shape, Precision::FP16); + + size_t new_value = g.input(shape, Precision::FP16); + std::vector<__fp16> data(8); + for (size_t i = 0; i < 8; ++i) data[i] = static_cast<__fp16>(static_cast(i + 1)); + g.set_input(new_value, data.data(), Precision::FP16); + + g.recurrent_cache_write(new_value, cache_state); + g.execute(); + + g.soft_reset_keep_pool(); + size_t new_value2 = g.input(shape, Precision::FP16); + g.set_input(new_value2, data.data(), Precision::FP16); + g.recurrent_cache_write(new_value2, cache_state); + g.execute(); + + const __fp16* read = static_cast(g.get_output(cache_state)); + if (!read) return false; + for (size_t i = 0; i < 8; ++i) { + if (static_cast(read[i]) != static_cast(data[i])) return false; + } + return true; +} + +bool run_benchmarks() { + auto bench = [](const char* label, auto setup, auto run) { + setup(); + run(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) run(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + { + const size_t kv = 8, d = 128, max_seq = 1024; + + CactusGraph g; + size_t k_cache = g.kv_cache_state(max_seq, kv, d); + + size_t prefill_elements = 512 * kv * d; + size_t prefill_input = g.input({prefill_elements}, Precision::FP16); + std::vector<__fp16> prefill_data(prefill_elements); + fill_random_fp16(prefill_data); + g.set_input(prefill_input, prefill_data.data(), Precision::FP16); + g.kv_cache_append(prefill_input, k_cache); + g.execute(); + + size_t append_elements = 1 * kv * d; + std::vector<__fp16> append_data(append_elements); + fill_random_fp16(append_data); + + bench("kv_cache_append 1tok@512", []{}, [&]{ + g.soft_reset_keep_pool(); + size_t inp = g.input({append_elements}, Precision::FP16); + g.set_input(inp, append_data.data(), Precision::FP16); + g.kv_cache_append(inp, k_cache); + g.execute(); + }); + } + + { + const size_t b = 1, s = 1, h = 16, kv = 8, d = 128, max_seq = 1024; + float scale = 1.0f / std::sqrt(static_cast(d)); + + CactusGraph g; + size_t k_cache = g.kv_cache_state(max_seq, kv, d); + size_t v_cache = g.kv_cache_state(max_seq, kv, d); + + size_t prefill_elements = 512 * kv * d; + std::vector<__fp16> prefill_data(prefill_elements); + fill_random_fp16(prefill_data); + + size_t pk = g.input({prefill_elements}, Precision::FP16); + size_t pv = g.input({prefill_elements}, Precision::FP16); + g.set_input(pk, prefill_data.data(), Precision::FP16); + g.set_input(pv, prefill_data.data(), Precision::FP16); + g.kv_cache_append(pk, k_cache); + g.kv_cache_append(pv, v_cache); + g.execute(); + + std::vector<__fp16> q(b*s*h*d), k_new(b*s*kv*d), v_new(b*s*kv*d); + fill_random_fp16(q); + fill_random_fp16(k_new); + fill_random_fp16(v_new); + + bench("attention_cached 1tok@512", []{}, [&]{ + g.soft_reset_keep_pool(); + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k_new.data(), Precision::FP16); + g.set_input(iv, v_new.data(), Precision::FP16); + g.kv_cache_append(ik, k_cache); + g.kv_cache_append(iv, v_cache); + g.attention_cached(iq, ik, iv, k_cache, v_cache, scale, 512); + g.execute(); + }); + } + + return true; +} + +bool test_kv_cache_slots_independent() { + const size_t b = 1, s = 1, h = 2, kv = 2, d = 16, max_seq = 64; + const float scale = 1.0f / std::sqrt(static_cast(d)); + + std::vector<__fp16> q(b * s * h * d), k0(b * s * kv * d), v0(b * s * kv * d), + k1(b * s * kv * d), v1(b * s * kv * d); + fill_random_fp16(q); + fill_random_fp16(k0); + fill_random_fp16(v0); + fill_random_fp16(k1); + fill_random_fp16(v1); + + std::vector ref; + { + CactusGraph g; + size_t kc = g.kv_cache_state(max_seq, kv, d); + size_t vc = g.kv_cache_state(max_seq, kv, d); + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k0.data(), Precision::FP16); + g.set_input(iv, v0.data(), Precision::FP16); + g.kv_cache_append(ik, kc); + g.kv_cache_append(iv, vc); + size_t attn = g.attention_cached(iq, ik, iv, kc, vc, scale, 0); + g.execute(); + __fp16* r = static_cast<__fp16*>(g.get_output(attn)); + ref.assign(r, r + b * s * h * d); + g.hard_reset(); + } + + CactusGraph g; + size_t kc = g.kv_cache_state(max_seq, kv, d, 0, 4, 2); + size_t vc = g.kv_cache_state(max_seq, kv, d, 0, 4, 2); + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik0 = g.input({b, s, kv, d}, Precision::FP16); + size_t iv0 = g.input({b, s, kv, d}, Precision::FP16); + size_t ik1 = g.input({b, s, kv, d}, Precision::FP16); + size_t iv1 = g.input({b, s, kv, d}, Precision::FP16); + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik0, k0.data(), Precision::FP16); + g.set_input(iv0, v0.data(), Precision::FP16); + g.set_input(ik1, k1.data(), Precision::FP16); + g.set_input(iv1, v1.data(), Precision::FP16); + g.kv_cache_append(ik0, kc, 0, 4, 0); + g.kv_cache_append(iv0, vc, 0, 4, 0); + g.kv_cache_append(ik1, kc, 0, 4, 1); + g.kv_cache_append(iv1, vc, 0, 4, 1); + size_t attn0 = g.attention_cached(iq, ik0, iv0, kc, vc, scale, 0, 0, 0, 0); + size_t attn1 = g.attention_cached(iq, ik1, iv1, kc, vc, scale, 0, 0, 0, 1); + g.execute(); + + __fp16* r0 = static_cast<__fp16*>(g.get_output(attn0)); + __fp16* r1 = static_cast<__fp16*>(g.get_output(attn1)); + size_t n = b * s * h * d; + bool slot0_matches_ref = true, slots_differ = false; + for (size_t i = 0; i < n; i++) { + if (!std::isfinite(static_cast(r0[i])) || !std::isfinite(static_cast(r1[i]))) { + g.hard_reset(); + return false; + } + if (std::abs(static_cast(r0[i]) - ref[i]) > 1e-2f) slot0_matches_ref = false; + if (std::abs(static_cast(r0[i]) - static_cast(r1[i])) > 1e-3f) slots_differ = true; + } + g.hard_reset(); + return slot0_matches_ref && slots_differ; +} + +bool test_batched_per_slot_attention() { + const size_t h = 2, kv = 2, d = 16, max_seq = 64; + const float scale = 1.0f / std::sqrt(static_cast(d)); + + std::vector<__fp16> q0(h * d), q1(h * d), k0(kv * d), v0(kv * d), k1(kv * d), v1(kv * d); + fill_random_fp16(q0); + fill_random_fp16(q1); + fill_random_fp16(k0); + fill_random_fp16(v0); + fill_random_fp16(k1); + fill_random_fp16(v1); + + CactusGraph g; + size_t kc = g.kv_cache_state(max_seq, kv, d, 0, 4, 2); + size_t vc = g.kv_cache_state(max_seq, kv, d, 0, 4, 2); + + size_t iq0 = g.input({1, 1, h, d}, Precision::FP16); + size_t ik0 = g.input({1, 1, kv, d}, Precision::FP16); + size_t iv0 = g.input({1, 1, kv, d}, Precision::FP16); + size_t iq1 = g.input({1, 1, h, d}, Precision::FP16); + size_t ik1 = g.input({1, 1, kv, d}, Precision::FP16); + size_t iv1 = g.input({1, 1, kv, d}, Precision::FP16); + g.set_input(iq0, q0.data(), Precision::FP16); + g.set_input(ik0, k0.data(), Precision::FP16); + g.set_input(iv0, v0.data(), Precision::FP16); + g.set_input(iq1, q1.data(), Precision::FP16); + g.set_input(ik1, k1.data(), Precision::FP16); + g.set_input(iv1, v1.data(), Precision::FP16); + + g.kv_cache_append(ik0, kc, 0, 4, 0); + g.kv_cache_append(iv0, vc, 0, 4, 0); + g.kv_cache_append(ik1, kc, 0, 4, 1); + g.kv_cache_append(iv1, vc, 0, 4, 1); + + size_t a0 = g.attention_cached(iq0, ik0, iv0, kc, vc, scale, 0, 0, 0, 0); + size_t a1 = g.attention_cached(iq1, ik1, iv1, kc, vc, scale, 0, 0, 0, 1); + + std::vector<__fp16> qb(q0); qb.insert(qb.end(), q1.begin(), q1.end()); + std::vector<__fp16> kb(k0); kb.insert(kb.end(), k1.begin(), k1.end()); + std::vector<__fp16> vb(v0); vb.insert(vb.end(), v1.begin(), v1.end()); + size_t iqb = g.input({2, 1, h, d}, Precision::FP16); + size_t ikb = g.input({2, 1, kv, d}, Precision::FP16); + size_t ivb = g.input({2, 1, kv, d}, Precision::FP16); + g.set_input(iqb, qb.data(), Precision::FP16); + g.set_input(ikb, kb.data(), Precision::FP16); + g.set_input(ivb, vb.data(), Precision::FP16); + size_t ab = g.attention_cached(iqb, ikb, ivb, kc, vc, scale, 0); + + g.execute(); + + __fp16* r0 = static_cast<__fp16*>(g.get_output(a0)); + __fp16* r1 = static_cast<__fp16*>(g.get_output(a1)); + __fp16* rb = static_cast<__fp16*>(g.get_output(ab)); + size_t n = h * d; + for (size_t i = 0; i < n; i++) { + if (std::abs(static_cast(rb[i]) - static_cast(r0[i])) > 1e-3f) { g.hard_reset(); return false; } + if (std::abs(static_cast(rb[n + i]) - static_cast(r1[i])) > 1e-3f) { g.hard_reset(); return false; } + } + g.hard_reset(); + return true; +} + +bool test_batched_kv_append() { + const size_t h = 2, kv = 2, d = 16, max_seq = 64; + const float scale = 1.0f / std::sqrt(static_cast(d)); + + std::vector<__fp16> q0(h * d), q1(h * d), k0(kv * d), v0(kv * d), k1(kv * d), v1(kv * d); + fill_random_fp16(q0); + fill_random_fp16(q1); + fill_random_fp16(k0); + fill_random_fp16(v0); + fill_random_fp16(k1); + fill_random_fp16(v1); + + auto run = [&](bool batched_append, std::vector& out0, std::vector& out1) { + CactusGraph g; + size_t kc = g.kv_cache_state(max_seq, kv, d, 0, 4, 2); + size_t vc = g.kv_cache_state(max_seq, kv, d, 0, 4, 2); + size_t iq0 = g.input({1, 1, h, d}, Precision::FP16); + size_t ik0 = g.input({1, 1, kv, d}, Precision::FP16); + size_t iv0 = g.input({1, 1, kv, d}, Precision::FP16); + size_t iq1 = g.input({1, 1, h, d}, Precision::FP16); + size_t ik1 = g.input({1, 1, kv, d}, Precision::FP16); + size_t iv1 = g.input({1, 1, kv, d}, Precision::FP16); + g.set_input(iq0, q0.data(), Precision::FP16); + g.set_input(ik0, k0.data(), Precision::FP16); + g.set_input(iv0, v0.data(), Precision::FP16); + g.set_input(iq1, q1.data(), Precision::FP16); + g.set_input(ik1, k1.data(), Precision::FP16); + g.set_input(iv1, v1.data(), Precision::FP16); + if (batched_append) { + std::vector<__fp16> kb(k0); kb.insert(kb.end(), k1.begin(), k1.end()); + std::vector<__fp16> vb(v0); vb.insert(vb.end(), v1.begin(), v1.end()); + size_t ikb = g.input({2, 1, kv, d}, Precision::FP16); + size_t ivb = g.input({2, 1, kv, d}, Precision::FP16); + g.set_input(ikb, kb.data(), Precision::FP16); + g.set_input(ivb, vb.data(), Precision::FP16); + g.kv_cache_append(ikb, kc); + g.kv_cache_append(ivb, vc); + } else { + g.kv_cache_append(ik0, kc, 0, 4, 0); + g.kv_cache_append(iv0, vc, 0, 4, 0); + g.kv_cache_append(ik1, kc, 0, 4, 1); + g.kv_cache_append(iv1, vc, 0, 4, 1); + } + size_t a0 = g.attention_cached(iq0, ik0, iv0, kc, vc, scale, 0, 0, 0, 0); + size_t a1 = g.attention_cached(iq1, ik1, iv1, kc, vc, scale, 0, 0, 0, 1); + g.execute(); + __fp16* r0 = static_cast<__fp16*>(g.get_output(a0)); + __fp16* r1 = static_cast<__fp16*>(g.get_output(a1)); + out0.assign(r0, r0 + h * d); + out1.assign(r1, r1 + h * d); + g.hard_reset(); + }; + + std::vector ref0, ref1, b0, b1; + run(false, ref0, ref1); + run(true, b0, b1); + for (size_t i = 0; i < h * d; i++) { + if (std::abs(b0[i] - ref0[i]) > 1e-3f) return false; + if (std::abs(b1[i] - ref1[i]) > 1e-3f) return false; + } + return true; +} + +int main() { + TestUtils::TestRunner runner("Cache Tests"); + + runner.run_test("KV Cache State Init", test_kv_cache_state_init()); + runner.run_test("KV Cache Persistent", test_kv_cache_state_persistent()); + runner.run_test("KV Cache Append Basic", test_kv_cache_append_basic()); + runner.run_test("KV Cache Append Multiple", test_kv_cache_append_multiple()); + runner.run_test("KV Cache Append Eviction", test_kv_cache_append_eviction()); + runner.run_test("KV Cache Append Full Window Eviction", test_kv_cache_append_full_window_eviction()); + runner.run_test("Padded Rollback No Eviction", test_padded_rollback_no_eviction()); + runner.run_test("Padded Rollback Eviction Overshoot", test_padded_rollback_eviction_overshoot()); + runner.run_test("Padded Rollback Only Pads Evict", test_padded_rollback_only_pads_evict()); + runner.run_test("Padded Rollback Empty Cache", test_padded_rollback_empty_cache()); + runner.run_test("Attention Cached Basic", test_attention_cached_basic()); + runner.run_test("KV Cache Slots Independent", test_kv_cache_slots_independent()); + runner.run_test("Batched Per-Slot Attention", test_batched_per_slot_attention()); + runner.run_test("Batched KV Append", test_batched_kv_append()); + runner.run_test("Attention Cached Multistep", test_attention_cached_multistep()); + runner.run_test("KV Cache Invalidate", test_kv_cache_invalidate()); + runner.run_test("Conv Cache State Init", test_conv_cache_state_init()); + runner.run_test("Conv Cache Append Basic", test_conv_cache_append_basic()); + runner.run_test("Conv Cache Circular", test_conv_cache_append_circular()); + runner.run_test("Conv Cache Persistent", test_conv_cache_persistent()); + runner.run_test("Conv Cache Initialize Populates Trailing", test_conv_cache_initialize_populates_trailing_rows()); + runner.run_test("Conv Cache Initialize Rejects Non-State Input", test_conv_cache_initialize_rejects_non_state_input()); + runner.run_test("Conv Cache Initialize Resets Dirty Cache", test_conv_cache_initialize_resets_dirty_cache()); + runner.run_test("Conv Cache Initialize Output Zero Byte", test_conv_cache_initialize_output_is_zero_byte()); + runner.run_test("Recurrent Cache State Initializes to Zero", test_recurrent_cache_state_initializes_to_zero()); + runner.run_test("Steal Cache Buffer Preserves Source Descriptor", test_steal_cache_buffer_preserves_source_descriptor()); + runner.run_test("Recurrent Cache Write Carries State", test_recurrent_cache_write_carries_state_across_executions()); + runner.run_test("Recurrent Cache Write Rejects Non-State Input", test_recurrent_cache_write_rejects_non_state_input()); + runner.run_test("Recurrent Cache Write Rejects Mismatch", test_recurrent_cache_write_rejects_shape_mismatch()); + runner.run_test("Allocate Buffers Preserves Recurrent Init", test_allocate_buffers_preserves_recurrent_state_init()); + runner.run_test("Conv Cache Initialize Rejects Hidden Dim Mismatch", test_conv_cache_initialize_rejects_hidden_dim_mismatch()); + runner.run_test("Conv Cache Initialize Single Row", test_conv_cache_initialize_single_row()); + runner.run_test("Allocate Buffers Idempotent", test_allocate_buffers_idempotent()); + runner.run_test("Recurrent Cache Write Idempotent", test_recurrent_cache_write_idempotent()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_dsp_ops.cpp b/cactus-graph/tests/test_dsp_ops.cpp new file mode 100644 index 000000000..3706c91f5 --- /dev/null +++ b/cactus-graph/tests/test_dsp_ops.cpp @@ -0,0 +1,210 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +using namespace TestUtils; + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +bool test_rfft_graph_op() { + CactusGraph g; + + const size_t n = 64; + size_t input = g.input({n}, Precision::FP32); + + std::vector data(n); + for (size_t i = 0; i < n; i++) { + data[i] = std::sin(2.0f * static_cast(M_PI) * 5.0f * i / n); + } + g.set_input(input, data.data(), Precision::FP32); + + size_t rfft_out = g.rfft(input); + g.execute(); + + float* output = static_cast(g.get_output(rfft_out)); + const auto& buf = g.get_output_buffer(rfft_out); + if (buf.total_size != (n / 2 + 1) * 2) return false; + + bool has_nonzero = false; + for (size_t i = 0; i < buf.total_size; i++) { + if (!std::isfinite(output[i])) return false; + if (std::abs(output[i]) > 1e-3f) has_nonzero = true; + } + return has_nonzero; +} + +bool test_irfft_graph_op() { + CactusGraph g; + + const size_t n = 64; + size_t input = g.input({n}, Precision::FP32); + + std::vector data(n); + for (size_t i = 0; i < n; i++) { + data[i] = std::sin(2.0f * static_cast(M_PI) * 3.0f * i / n); + } + g.set_input(input, data.data(), Precision::FP32); + + size_t freq = g.rfft(input); + size_t recovered = g.irfft(freq, n); + g.execute(); + + float* result = static_cast(g.get_output(recovered)); + for (size_t i = 0; i < n; i++) { + if (std::abs(result[i] - data[i]) > 1e-3f) return false; + } + return true; +} + +bool test_mel_filter_bank_graph_op() { + CactusGraph g; + + const size_t num_freq = 257; + const size_t num_mel = 80; + + size_t mel_node = g.mel_filter_bank(num_freq, num_mel, 0.0f, 8000.0f, 16000); + g.execute(); + + const auto& buf = g.get_output_buffer(mel_node); + if (buf.shape[0] != num_freq || buf.shape[1] != num_mel) return false; + + float* filters = static_cast(g.get_output(mel_node)); + bool has_nonzero = false; + for (size_t i = 0; i < buf.total_size; i++) { + if (filters[i] < 0.0f) return false; + if (filters[i] > 0.0f) has_nonzero = true; + } + return has_nonzero; +} + +bool test_spectrogram_graph_op() { + CactusGraph g; + + const size_t n_samples = 1600; + const size_t frame_length = 400; + const size_t hop = 160; + const size_t fft_len = 400; + const size_t num_freq = fft_len / 2 + 1; + const size_t num_mel = 80; + + size_t mel_node = g.mel_filter_bank(num_freq, num_mel, 0.0f, 8000.0f, 16000); + + size_t wav_input = g.input({n_samples}, Precision::FP32); + std::vector waveform(n_samples); + for (size_t i = 0; i < n_samples; i++) { + waveform[i] = std::sin(2.0f * static_cast(M_PI) * 440.0f * i / 16000.0f); + } + g.set_input(wav_input, waveform.data(), Precision::FP32); + + size_t spec = g.spectrogram(wav_input, mel_node, frame_length, hop, fft_len); + g.execute(); + + float* output = static_cast(g.get_output(spec)); + const auto& buf = g.get_output_buffer(spec); + + if (buf.shape[0] != num_mel) return false; + + bool has_nonzero = false; + for (size_t i = 0; i < buf.total_size; i++) { + if (!std::isfinite(output[i])) return false; + if (output[i] > 1e-8f) has_nonzero = true; + } + return has_nonzero; +} + +bool test_spectrogram_with_log() { + CactusGraph g; + + const size_t n_samples = 1600; + const size_t frame_length = 400; + const size_t hop = 160; + const size_t fft_len = 400; + const size_t num_freq = fft_len / 2 + 1; + const size_t num_mel = 80; + + size_t mel_node = g.mel_filter_bank(num_freq, num_mel, 0.0f, 8000.0f, 16000); + + size_t wav_input = g.input({n_samples}, Precision::FP32); + std::vector waveform(n_samples); + for (size_t i = 0; i < n_samples; i++) { + waveform[i] = std::sin(2.0f * static_cast(M_PI) * 440.0f * i / 16000.0f); + } + g.set_input(wav_input, waveform.data(), Precision::FP32); + + size_t spec = g.spectrogram(wav_input, mel_node, frame_length, hop, fft_len, + 2.0f, true, 0, 1e-10f, 1); + g.execute(); + + float* output = static_cast(g.get_output(spec)); + const auto& buf = g.get_output_buffer(spec); + + bool has_negative = false; + for (size_t i = 0; i < buf.total_size; i++) { + if (!std::isfinite(output[i])) return false; + if (output[i] < 0.0f) has_negative = true; + } + return has_negative; +} + +bool run_benchmarks() { + auto bench = [](const char* label, auto setup, auto run) { + setup(); + run(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) run(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + { + const size_t n = 512; + std::vector data(n, 1.0f); + + bench("rfft graph op 512", []{}, [&]{ + CactusGraph g; + size_t inp = g.input({n}, Precision::FP32); + g.set_input(inp, data.data(), Precision::FP32); + g.rfft(inp); + g.execute(); + }); + } + + { + const size_t n_samples = 16000; + const size_t num_freq = 201; + const size_t num_mel = 80; + std::vector waveform(n_samples, 0.1f); + + CactusGraph g; + size_t mel = g.mel_filter_bank(num_freq, num_mel, 0.0f, 8000.0f, 16000); + size_t wav = g.input({n_samples}, Precision::FP32); + g.set_input(wav, waveform.data(), Precision::FP32); + g.spectrogram(wav, mel, 400, 160, 400); + + bench("spectrogram 1s@16kHz", []{}, [&]{ + g.execute(); + }); + } + + return true; +} + +int main() { + TestUtils::TestRunner runner("DSP Graph Ops Tests"); + + runner.run_test("RFFT Graph Op", test_rfft_graph_op()); + runner.run_test("IRFFT Roundtrip Graph Op", test_irfft_graph_op()); + runner.run_test("Mel Filter Bank Graph Op", test_mel_filter_bank_graph_op()); + runner.run_test("Spectrogram Graph Op", test_spectrogram_graph_op()); + runner.run_test("Spectrogram with Log", test_spectrogram_with_log()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_dynamic_shapes.cpp b/cactus-graph/tests/test_dynamic_shapes.cpp new file mode 100644 index 000000000..fec55fdf0 --- /dev/null +++ b/cactus-graph/tests/test_dynamic_shapes.cpp @@ -0,0 +1,109 @@ +#include "test_utils.h" +#include +#include +#include +#include + +using namespace TestUtils; + +namespace { + +bool test_dynamic_batch_matches_static() { + const size_t K = 8, N = 8, MAXB = 4; + std::vector<__fp16> wd(K * N), bd(N), xd(MAXB * K); + fill_random_fp16(wd); + fill_random_fp16(bd); + fill_random_fp16(xd); + + CactusGraph dyn; + size_t x = dyn.input({1, K}, Precision::FP16); + size_t w = dyn.input({K, N}, Precision::FP16); + size_t mm = dyn.matmul(x, w, false); + size_t bias = dyn.input({1, N}, Precision::FP16); + size_t out = dyn.add(mm, bias); + + for (size_t B : {size_t(1), size_t(4), size_t(2)}) { + std::vector<__fp16> xb(xd.begin(), xd.begin() + B * K); + + dyn.set_runtime_input_shape(x, {B, K}); + dyn.set_input(x, xb.data(), Precision::FP16); + dyn.set_input(w, wd.data(), Precision::FP16); + dyn.set_input(bias, bd.data(), Precision::FP16); + dyn.execute(); + + if (dyn.get_output_buffer(out).shape != std::vector{B, N}) return false; + std::vector dvals(B * N); + const __fp16* dop = static_cast<__fp16*>(dyn.get_output(out)); + for (size_t i = 0; i < B * N; ++i) dvals[i] = static_cast(dop[i]); + + CactusGraph st; + size_t sx = st.input({B, K}, Precision::FP16); + size_t sw = st.input({K, N}, Precision::FP16); + size_t smm = st.matmul(sx, sw, false); + size_t sbias = st.input({1, N}, Precision::FP16); + size_t sout = st.add(smm, sbias); + st.set_input(sx, xb.data(), Precision::FP16); + st.set_input(sw, wd.data(), Precision::FP16); + st.set_input(sbias, bd.data(), Precision::FP16); + st.execute(); + + if (st.get_output_buffer(sout).shape != std::vector{B, N}) return false; + const __fp16* sop = static_cast<__fp16*>(st.get_output(sout)); + for (size_t i = 0; i < B * N; ++i) { + if (std::abs(dvals[i] - static_cast(sop[i])) > 1e-3f) return false; + } + st.hard_reset(); + } + dyn.hard_reset(); + return true; +} + +bool test_buffer_resize_tracks_bucket() { + BufferPool pool; + BufferDesc d({4}, Precision::FP16); + d.resize_from_pool(pool); + if (d.get_data() == nullptr || d.byte_size != 8) return false; + + d.set_shape({4096}); + d.resize_from_pool(pool); + if (d.byte_size != 8192 || d.get_data() == nullptr) return false; + + d.set_shape({4}); + d.resize_from_pool(pool); + if (d.byte_size != 8 || d.get_data() == nullptr) return false; + return true; +} + +bool test_dynamic_mask_roundtrip() { + CactusGraph g; + size_t x = g.input({1, 8}, Precision::FP16); + size_t w = g.input({8, 8}, Precision::FP16); + size_t mm = g.matmul(x, w, false); + size_t bias = g.input({1, 8}, Precision::FP16); + g.add(mm, bias); + g.set_runtime_input_shape(x, {1, 8}); + + std::string path = + (std::filesystem::temp_directory_path() / "cactus_dynamic_roundtrip.cactus").string(); + GraphFile::save_graph(g, path); + GraphFile::SerializedGraph sg = GraphFile::load_graph(path); + std::filesystem::remove(path); + + if (sg.header.version != 6) return false; + size_t dynamic_inputs = 0; + for (const auto& n : sg.nodes) { + if (n.op_type == OpType::INPUT && !n.dynamic_mask.empty()) dynamic_inputs++; + } + return dynamic_inputs == 1; +} + +} + +int main() { + TestUtils::TestRunner runner("Dynamic Shape Tests"); + runner.run_test("Dynamic batch matches static", test_dynamic_batch_matches_static()); + runner.run_test("Buffer resize tracks bucket", test_buffer_resize_tracks_bucket()); + runner.run_test("Dynamic mask roundtrip", test_dynamic_mask_roundtrip()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_graph.cpp b/cactus-graph/tests/test_graph.cpp new file mode 100644 index 000000000..77bdeae58 --- /dev/null +++ b/cactus-graph/tests/test_graph.cpp @@ -0,0 +1,138 @@ +#include "test_utils.h" +#include +#include +#include +#include + +using namespace TestUtils; + +bool test_complex_graph_structure() { + TestUtils::FP16TestFixture fixture("Complex Graph Structure"); + + size_t input_a = fixture.create_input({2, 2}); + size_t input_b = fixture.create_input({2, 2}); + size_t input_c = fixture.create_input({2, 2}); + + size_t add_ab = fixture.graph().add(input_a, input_b); + size_t mul_result = fixture.graph().multiply(add_ab, input_c); + size_t scalar_result = fixture.graph().scalar_add(mul_result, 1.0f); + + std::vector<__fp16> data_a = {1, 2, 3, 4}; + std::vector<__fp16> data_b = {2, 3, 4, 5}; + std::vector<__fp16> data_c = {2, 2, 2, 2}; + fixture.set_input_data(input_a, data_a); + fixture.set_input_data(input_b, data_b); + fixture.set_input_data(input_c, data_c); + + fixture.execute(); + + std::vector<__fp16> expected = {7, 11, 15, 19}; + return fixture.verify_output(scalar_result, expected); +} + +bool test_multiple_outputs() { + TestUtils::FP16TestFixture fixture("Multiple Outputs"); + + size_t input_a = fixture.create_input({3}); + size_t add_result = fixture.graph().scalar_add(input_a, 10.0f); + size_t mul_result = fixture.graph().scalar_multiply(input_a, 2.0f); + size_t combine_result = fixture.graph().add(add_result, mul_result); + + fixture.graph().retain_outputs({static_cast(add_result), static_cast(mul_result)}); + + std::vector<__fp16> data_a = {1, 2, 3}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected_add = {11, 12, 13}; + std::vector<__fp16> expected_mul = {2, 4, 6}; + std::vector<__fp16> expected_combine = {13, 16, 19}; + + return fixture.verify_output(add_result, expected_add) && + fixture.verify_output(mul_result, expected_mul) && + fixture.verify_output(combine_result, expected_combine); +} + +bool test_graph_reset() { + CactusGraph graph; + + size_t input_a = graph.input({2}, Precision::FP16); + size_t result_id = graph.scalar_add(input_a, 5.0f); + + std::vector<__fp16> data_a = {1, 2}; + graph.set_input(input_a, data_a.data(), Precision::FP16); + graph.execute(); + + __fp16* output1 = static_cast<__fp16*>(graph.get_output(result_id)); + if (std::abs(static_cast(output1[0]) - 6.0f) > 1e-2f || + std::abs(static_cast(output1[1]) - 7.0f) > 1e-2f) return false; + + graph.hard_reset(); + if (graph.get_node_count() != 0) return false; + + size_t new_input = graph.input({2}, Precision::FP16); + size_t new_result = graph.scalar_add(new_input, 5.0f); + + std::vector<__fp16> data_b = {10, 20}; + graph.set_input(new_input, data_b.data(), Precision::FP16); + graph.execute(); + + __fp16* output2 = static_cast<__fp16*>(graph.get_output(new_result)); + return (std::abs(static_cast(output2[0]) - 15.0f) < 1e-2f && + std::abs(static_cast(output2[1]) - 25.0f) < 1e-2f); +} + +bool run_benchmarks() { + std::vector<__fp16> data(4, static_cast<__fp16>(1.0f)); + + auto build_chain = [&](CactusGraph& g) { + size_t a = g.input({4}, Precision::FP16); + size_t b = g.input({4}, Precision::FP16); + g.set_input(a, data.data(), Precision::FP16); + g.set_input(b, data.data(), Precision::FP16); + size_t node = a; + for (int n = 0; n < 100; ++n) node = g.add(node, b); + return node; + }; + + auto bench = [](const char* label, auto fn) { + fn(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + bench("graph_construct 100 nodes", [&]{ + CactusGraph g; + build_chain(g); + }); + + bench("graph_execute 100 nodes", [&]{ + CactusGraph g; + build_chain(g); + g.execute(); + }); + + bench("graph_reset", [&]{ + CactusGraph g; + build_chain(g); + g.execute(); + g.hard_reset(); + }); + + return true; +} + +int main() { + TestUtils::TestRunner runner("Graph Structure Tests"); + + runner.run_test("Complex Graph Structure", test_complex_graph_structure()); + runner.run_test("Multiple Outputs", test_multiple_outputs()); + runner.run_test("Graph Reset", test_graph_reset()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_image.cpp b/cactus-graph/tests/test_image.cpp new file mode 100644 index 000000000..78231a675 --- /dev/null +++ b/cactus-graph/tests/test_image.cpp @@ -0,0 +1,129 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +using namespace TestUtils; + +bool test_image_preprocess_basic() { + CactusGraph g; + + const int w = 32, h = 32, ch = 3, ps = 16; + size_t pixel_count = static_cast(w) * h * ch; + + size_t input = g.input({pixel_count}, Precision::FP32); + std::vector pixels(pixel_count); + for (size_t i = 0; i < pixel_count; i++) + pixels[i] = static_cast(i % 256); + g.set_input(input, pixels.data(), Precision::FP32); + + size_t result = g.image_preprocess(input, w, h, w, h, ps, ch); + g.execute(); + + const auto& buf = g.get_output_buffer(result); + int ph = h / ps; + int pw = w / ps; + size_t expected_patches = static_cast(ph) * pw; + size_t patch_dim = static_cast(ps) * ps * ch; + + if (buf.shape[0] != expected_patches) return false; + if (buf.shape[1] != patch_dim) return false; + + float* out = static_cast(g.get_output(result)); + bool has_nonzero = false; + for (size_t i = 0; i < buf.total_size; i++) { + if (!std::isfinite(out[i])) return false; + if (std::abs(out[i]) > 1e-6f) has_nonzero = true; + } + return has_nonzero; +} + +bool test_image_preprocess_resize() { + CactusGraph g; + + const int src_w = 64, src_h = 64, ch = 3, ps = 16; + const int dst_w = 32, dst_h = 32; + size_t pixel_count = static_cast(src_w) * src_h * ch; + + size_t input = g.input({pixel_count}, Precision::FP32); + std::vector pixels(pixel_count, 128.0f); + g.set_input(input, pixels.data(), Precision::FP32); + + size_t result = g.image_preprocess(input, src_w, src_h, dst_w, dst_h, ps, ch); + g.execute(); + + const auto& buf = g.get_output_buffer(result); + int ph = dst_h / ps; + int pw = dst_w / ps; + if (buf.shape[0] != static_cast(ph * pw)) return false; + + float* out = static_cast(g.get_output(result)); + for (size_t i = 0; i < buf.total_size; i++) { + if (!std::isfinite(out[i])) return false; + } + return true; +} + +bool test_image_preprocess_normalize() { + CactusGraph g; + + const int w = 16, h = 16, ch = 3, ps = 16; + size_t pixel_count = static_cast(w) * h * ch; + + size_t input = g.input({pixel_count}, Precision::FP32); + std::vector pixels(pixel_count, 127.5f); + g.set_input(input, pixels.data(), Precision::FP32); + + float mean[3] = {0.5f, 0.5f, 0.5f}; + float std_dev[3] = {0.5f, 0.5f, 0.5f}; + size_t result = g.image_preprocess(input, w, h, w, h, ps, ch, 1.0f / 255.0f, mean, std_dev); + g.execute(); + + float* out = static_cast(g.get_output(result)); + float expected = (127.5f / 255.0f - 0.5f) / 0.5f; + if (std::abs(out[0] - expected) > 0.02f) return false; + + return true; +} + +bool run_benchmarks() { + auto bench = [](const char* label, auto setup, auto run) { + setup(); + run(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) run(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + { + const int w = 512, h = 512, ch = 3, ps = 16; + size_t pixel_count = static_cast(w) * h * ch; + std::vector pixels(pixel_count, 128.0f); + + bench("image_preprocess 512x512", []{}, [&]{ + CactusGraph g; + size_t inp = g.input({pixel_count}, Precision::FP32); + g.set_input(inp, pixels.data(), Precision::FP32); + g.image_preprocess(inp, w, h, w, h, ps, ch); + g.execute(); + }); + } + + return true; +} + +int main() { + TestUtils::TestRunner runner("Image Preprocessing Tests"); + + runner.run_test("Image Preprocess Basic", test_image_preprocess_basic()); + runner.run_test("Image Preprocess Resize", test_image_preprocess_resize()); + runner.run_test("Image Preprocess Normalize", test_image_preprocess_normalize()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_io.cpp b/cactus-graph/tests/test_io.cpp new file mode 100644 index 000000000..d792b6102 --- /dev/null +++ b/cactus-graph/tests/test_io.cpp @@ -0,0 +1,524 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include + +using namespace TestUtils; + +namespace { +template +void print_vector_inline(const std::vector& values) { + std::cout << "["; + for (size_t i = 0; i < values.size(); ++i) { + if (i > 0) std::cout << ", "; + std::cout << values[i]; + } + std::cout << "]"; +} +} + +bool test_node_save_load() { + try { + CactusGraph graph; + + size_t input_a = graph.input({2, 3}, Precision::FP16); + size_t input_b = graph.input({2, 3}, Precision::FP16); + size_t result_id = graph.add(input_a, input_b); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {10, 20, 30, 40, 50, 60}; + + graph.set_input(input_a, const_cast(static_cast(data_a.data())), Precision::FP16); + graph.set_input(input_b, const_cast(static_cast(data_b.data())), Precision::FP16); + graph.execute(); + + std::string filename = "test_graph_save_load.bin"; + GraphFile::save_node(graph, result_id, filename); + + CactusGraph new_graph; + size_t loaded_id = new_graph.mmap_weights(filename); + new_graph.execute(); + + __fp16* original_data = static_cast<__fp16*>(graph.get_output(result_id)); + __fp16* loaded_data = static_cast<__fp16*>(new_graph.get_output(loaded_id)); + + for (size_t i = 0; i < 6; ++i) { + if (std::abs(static_cast(original_data[i]) - static_cast(loaded_data[i])) > 1e-3f) { + graph.hard_reset(); + new_graph.hard_reset(); + std::remove(filename.c_str()); + return false; + } + } + + const auto& buf = new_graph.get_output_buffer(loaded_id); + bool result = (buf.shape == std::vector{2, 3}) && + (buf.precision == Precision::FP16) && + (buf.byte_size == 12); + + graph.hard_reset(); + new_graph.hard_reset(); + std::remove(filename.c_str()); + return result; + } catch (const std::exception& e) { + return false; + } +} + +bool test_graph_save_load() { + try { + const std::string filename = "test_graph_save_load.cg"; + + CactusGraph graph; + size_t input_a = graph.input({2, 3}, Precision::FP16); + size_t input_b = graph.input({2, 3}, Precision::FP16); + size_t sum_id = graph.add(input_a, input_b); + size_t pow_id = graph.pow(sum_id, 2.0f); + + graph.save(filename); + + GraphFile::SerializedGraph sg = GraphFile::load_graph(filename); + + if (sg.header.node_count != 4) { + std::cout << "[graph_save_load] unexpected node_count: " + << sg.header.node_count << " expected 4" << std::endl; + std::remove(filename.c_str()); + return false; + } + + if (sg.graph_inputs.size() != 2 || sg.graph_inputs[0] != input_a || + sg.graph_inputs[1] != input_b) { + std::cout << "[graph_save_load] unexpected graph_inputs:"; + for (uint32_t idx : sg.graph_inputs) { + std::cout << " " << idx; + } + std::cout << std::endl; + std::remove(filename.c_str()); + return false; + } + + if (sg.graph_outputs.size() != 1 || sg.graph_outputs[0] != pow_id) { + std::cout << "[graph_save_load] unexpected graph_outputs:"; + for (uint32_t idx : sg.graph_outputs) { + std::cout << " " << idx; + } + std::cout << std::endl; + std::remove(filename.c_str()); + return false; + } + + CactusGraph loaded = CactusGraph::load(filename); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {10, 20, 30, 40, 50, 60}; + + const size_t loaded_input_a = runtime_id_from_serialized_index(sg.graph_inputs[0]); + const size_t loaded_input_b = runtime_id_from_serialized_index(sg.graph_inputs[1]); + const size_t loaded_output_id = runtime_id_from_serialized_index(sg.graph_outputs[0]); + + loaded.set_input(loaded_input_a, data_a.data(), Precision::FP16); + loaded.set_input(loaded_input_b, data_b.data(), Precision::FP16); + loaded.execute(); + + __fp16* output = static_cast<__fp16*>(loaded.get_output(loaded_output_id)); + std::vector expected = { + 121.0f, 484.0f, 1089.0f, + 1936.0f, 3025.0f, 4356.0f + }; + + for (size_t i = 0; i < expected.size(); ++i) { + float got = static_cast(output[i]); + if (std::abs(got - expected[i]) > 1.5f) { + std::cout << "[graph_save_load] mismatch at index " << i + << ": got=" << got + << " expected=" << expected[i] << std::endl; + std::remove(filename.c_str()); + return false; + } + } + + std::remove(filename.c_str()); + return true; + } catch (const std::exception& e) { + std::cout << "[graph_save_load] exception: " << e.what() << std::endl; + return false; + } +} + +bool test_graph_save_load_roundtrip_execution() { + try { + const std::string filename = "test_graph_save_load_roundtrip.cg"; + + CactusGraph original; + size_t input_a = original.input({2, 3}, Precision::FP16); + size_t input_b = original.input({2, 3}, Precision::FP16); + size_t sum_id = original.add(input_a, input_b); + size_t pow_id = original.pow(sum_id, 2.0f); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {10, 20, 30, 40, 50, 60}; + + original.set_input(input_a, data_a.data(), Precision::FP16); + original.set_input(input_b, data_b.data(), Precision::FP16); + original.execute(); + + __fp16* original_output = static_cast<__fp16*>(original.get_output(pow_id)); + std::vector expected(6); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = static_cast(original_output[i]); + } + + original.save(filename); + + CactusGraph loaded = CactusGraph::load(filename); + GraphFile::SerializedGraph loaded_graph = GraphFile::load_graph(filename); + const size_t loaded_input_a = runtime_id_from_serialized_index(loaded_graph.graph_inputs[0]); + const size_t loaded_input_b = runtime_id_from_serialized_index(loaded_graph.graph_inputs[1]); + const size_t loaded_output_id = runtime_id_from_serialized_index(loaded_graph.graph_outputs[0]); + + loaded.set_input(loaded_input_a, data_a.data(), Precision::FP16); + loaded.set_input(loaded_input_b, data_b.data(), Precision::FP16); + loaded.execute(); + + __fp16* loaded_output = static_cast<__fp16*>(loaded.get_output(loaded_output_id)); + for (size_t i = 0; i < expected.size(); ++i) { + float got = static_cast(loaded_output[i]); + if (std::abs(got - expected[i]) > 1e-3f) { + std::cout << "[graph_save_load_roundtrip] mismatch at index " << i + << ": got=" << got + << " expected(original)=" << expected[i] << std::endl; + std::remove(filename.c_str()); + return false; + } + } + + std::remove(filename.c_str()); + return true; + } catch (const std::exception& e) { + std::cout << "[graph_save_load_roundtrip] exception: " << e.what() << std::endl; + return false; + } +} + +bool test_graph_save_load_supported_ops_roundtrip() { + try { + const std::string filename = "test_graph_supported_roundtrip.cg"; + + CactusGraph original; + + size_t a = original.input({2, 3}, Precision::FP16); + size_t b = original.input({2, 3}, Precision::FP16); + size_t m1 = original.input({2, 3}, Precision::FP16); + size_t m2 = original.input({3, 2}, Precision::FP16); + + size_t add_id = original.add(a, b); + size_t add_clipped_id = original.add_clipped(a, b); + size_t sub_id = original.subtract(a, b); + size_t mul_id = original.multiply(a, b); + size_t div_id = original.divide(a, b); + + size_t abs_id = original.abs(sub_id); + size_t relu_id = original.relu(sub_id); + size_t silu_id = original.silu(add_id); + size_t gelu_id = original.gelu(add_id); + size_t gelu_erf_id = original.gelu_erf(add_id); + size_t sigmoid_id = original.sigmoid(sub_id); + size_t tanh_id = original.tanh(sub_id); + + size_t pow_id = original.pow(abs_id, 2.0f); + size_t scalar_add_id = original.scalar_add(add_id, 1.5f); + size_t scalar_sub_id = original.scalar_subtract(add_id, 0.5f); + size_t scalar_mul_id = original.scalar_multiply(add_id, 2.0f); + size_t scalar_div_id = original.scalar_divide(add_id, 2.0f); + + size_t view_id = original.view(add_id, {3, 2}); + size_t reshape_id = original.reshape(add_id, {3, 2}); + size_t flatten_id = original.flatten(add_id); + size_t concat_id = original.concat(a, b, 1); + size_t cat_id = original.cat({a, b}, 1); + + size_t slice_id = original.slice(add_id, 1, 1, 2); + size_t index_id = original.index(add_id, 1, 0); + + size_t sum_id = original.sum(add_id, -1); + size_t mean_id = original.mean(add_id, 1); + size_t variance_id = original.variance(add_id, 1); + size_t min_id = original.min(add_id, 1); + size_t max_id = original.max(add_id, 1); + size_t softmax_id = original.softmax(add_id, 1); + + size_t matmul_id = original.matmul(m1, m2, false); + size_t persistent_id = original.persistent(add_id); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {6, 5, 4, 3, 2, 1}; + std::vector<__fp16> data_m1 = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_m2 = {1, 2, 3, 4, 5, 6}; + + original.set_input(a, data_a.data(), Precision::FP16); + original.set_input(b, data_b.data(), Precision::FP16); + original.set_input(m1, data_m1.data(), Precision::FP16); + original.set_input(m2, data_m2.data(), Precision::FP16); + original.execute(); + + std::vector check_nodes = { + add_id, add_clipped_id, sub_id, mul_id, div_id, + abs_id, relu_id, silu_id, gelu_id, gelu_erf_id, sigmoid_id, tanh_id, + pow_id, scalar_add_id, scalar_sub_id, scalar_mul_id, scalar_div_id, + view_id, reshape_id, flatten_id, concat_id, cat_id, + slice_id, index_id, + sum_id, mean_id, variance_id, min_id, max_id, softmax_id, + matmul_id, persistent_id + }; + + std::vector> expected_outputs; + std::vector> expected_shapes; + expected_outputs.reserve(check_nodes.size()); + expected_shapes.reserve(check_nodes.size()); + for (size_t node_id : check_nodes) { + const auto& buf = original.get_output_buffer(node_id); + __fp16* out = static_cast<__fp16*>(original.get_output(node_id)); + std::vector values(buf.total_size); + for (size_t i = 0; i < buf.total_size; ++i) { + values[i] = static_cast(out[i]); + } + expected_outputs.push_back(std::move(values)); + expected_shapes.push_back(buf.shape); + } + + original.save(filename); + + CactusGraph loaded = CactusGraph::load(filename); + GraphFile::SerializedGraph loaded_graph = GraphFile::load_graph(filename); + if (loaded_graph.graph_inputs.size() != 4 || loaded_graph.graph_outputs.empty()) { + std::remove(filename.c_str()); + return false; + } + + loaded.set_input(runtime_id_from_serialized_index(loaded_graph.graph_inputs[0]), data_a.data(), Precision::FP16); + loaded.set_input(runtime_id_from_serialized_index(loaded_graph.graph_inputs[1]), data_b.data(), Precision::FP16); + loaded.set_input(runtime_id_from_serialized_index(loaded_graph.graph_inputs[2]), data_m1.data(), Precision::FP16); + loaded.set_input(runtime_id_from_serialized_index(loaded_graph.graph_inputs[3]), data_m2.data(), Precision::FP16); + loaded.execute(); + + for (size_t node_idx = 0; node_idx < check_nodes.size(); ++node_idx) { + size_t node_id = check_nodes[node_idx]; + const auto& loaded_buf = loaded.get_output_buffer(node_id); + __fp16* loaded_out = static_cast<__fp16*>(loaded.get_output(node_id)); + + if (loaded_buf.shape != expected_shapes[node_idx]) { + std::cout << "[supported_roundtrip] shape mismatch for node " << node_id + << ": got="; + print_vector_inline(loaded_buf.shape); + std::cout << " expected="; + print_vector_inline(expected_shapes[node_idx]); + std::cout << std::endl; + std::remove(filename.c_str()); + return false; + } + + if (loaded_buf.total_size != expected_outputs[node_idx].size()) { + std::cout << "[supported_roundtrip] size mismatch for node " << node_id + << ": got=" << loaded_buf.total_size + << " expected=" << expected_outputs[node_idx].size() << std::endl; + std::remove(filename.c_str()); + return false; + } + + for (size_t i = 0; i < loaded_buf.total_size; ++i) { + float got = static_cast(loaded_out[i]); + float expected = expected_outputs[node_idx][i]; + if (std::abs(got - expected) > 1e-3f) { + std::cout << "[supported_roundtrip] mismatch for node " << node_id + << " at index " << i + << ": got=" << got + << " expected=" << expected << std::endl; + std::remove(filename.c_str()); + return false; + } + } + } + + std::remove(filename.c_str()); + return true; + } catch (const std::exception& e) { + std::cout << "[supported_roundtrip] exception: " << e.what() << std::endl; + return false; + } +} + +bool test_graph_save_for_inspection() { + try { + const std::string filename = "test_graph_inspect.cg"; + + CactusGraph graph; + size_t input_a = graph.input({2, 3}, Precision::FP16); + size_t input_b = graph.input({2, 3}, Precision::FP16); + size_t sum_id = graph.add(input_a, input_b); + size_t pow_id = graph.pow(sum_id, 2.0f); + + (void)input_a; + (void)input_b; + (void)sum_id; + (void)pow_id; + + graph.save(filename); + return true; + } catch (const std::exception& e) { + std::cout << "[graph_save_for_inspection] exception: " << e.what() << std::endl; + return false; + } +} + +bool test_save_load_preserves_recurrent_cache_persistence() { + try { + const std::string filename = "test_recurrent_save_load.cg"; + + CactusGraph graph; + const std::vector shape{2, 4}; + size_t cache = graph.recurrent_cache_state(shape, Precision::FP16); + (void)cache; + graph.save(filename); + + CactusGraph loaded = CactusGraph::load(filename); + loaded.execute(); + + size_t loaded_cache_id = 0; + bool found = false; + for (size_t i = 0; i < loaded.get_node_count(); ++i) { + const auto& node = loaded.nodes_[i]; + if (node->op_type == OpType::RECURRENT_CACHE_STATE) { + loaded_cache_id = node->id; + found = true; + break; + } + } + if (!found) { + std::remove(filename.c_str()); + return false; + } + + const __fp16* data = static_cast(loaded.get_output(loaded_cache_id)); + if (!data) { + std::remove(filename.c_str()); + return false; + } + for (size_t i = 0; i < 8; ++i) { + if (static_cast(data[i]) != 0.0f) { + std::remove(filename.c_str()); + return false; + } + } + + loaded.release_runtime_buffers(); + const __fp16* still_there = static_cast(loaded.get_output(loaded_cache_id)); + bool persisted = still_there != nullptr; + + std::remove(filename.c_str()); + return persisted; + } catch (const std::exception& e) { + std::cout << "[recurrent_save_load] exception: " << e.what() << std::endl; + return false; + } +} + +bool test_save_load_preserves_kv_cache_num_slots() { + try { + const std::string filename = "test_kv_slots_save_load.cg"; + + CactusGraph graph; + size_t cache = graph.kv_cache_state(128, 4, 64, 0, 4, 8); // num_slots = 8 + (void)cache; + graph.save(filename); + + CactusGraph loaded = CactusGraph::load(filename); + size_t loaded_cache_id = 0; + bool found = false; + for (size_t i = 0; i < loaded.get_node_count(); ++i) { + const auto& node = loaded.nodes_[i]; + if (node->op_type == OpType::KV_CACHE_STATE) { + loaded_cache_id = node->id; + found = true; + break; + } + } + std::remove(filename.c_str()); + if (!found) return false; + return loaded.get_node_cache_num_slots(loaded_cache_id) == 8; + } catch (const std::exception& e) { + std::cout << "[kv_slots_save_load] exception: " << e.what() << std::endl; + return false; + } +} + +bool run_benchmarks() { + const int ITERS = 100; + const std::string temp_file = "bench_io_50nodes.cg"; + + // graph_save 50 nodes + { + double total = 0.0; + for (int i = 0; i < ITERS; ++i) { + CactusGraph graph; + size_t node = graph.input({4}, Precision::FP16); + size_t other = graph.input({4}, Precision::FP16); + for (int n = 0; n < 50; ++n) { + node = graph.add(node, other); + } + TestUtils::Timer t; + graph.save(temp_file); + total += t.elapsed_ms(); + graph.hard_reset(); + } + std::remove(temp_file.c_str()); + double ms = total / ITERS; + std::cout << " \u26a1 " << std::left << std::setw(30) << "graph_save 50 nodes" + << std::fixed << std::setprecision(3) << ms << " ms\n"; + } + + // graph_load 50 nodes — build and save once, then load repeatedly + { + { + CactusGraph graph; + size_t node = graph.input({4}, Precision::FP16); + size_t other = graph.input({4}, Precision::FP16); + for (int n = 0; n < 50; ++n) { + node = graph.add(node, other); + } + graph.save(temp_file); + } + double total = 0.0; + for (int i = 0; i < ITERS; ++i) { + TestUtils::Timer t; + CactusGraph loaded = CactusGraph::load(temp_file); + total += t.elapsed_ms(); + loaded.hard_reset(); + } + std::remove(temp_file.c_str()); + double ms = total / ITERS; + std::cout << " \u26a1 " << std::left << std::setw(30) << "graph_load 50 nodes" + << std::fixed << std::setprecision(3) << ms << " ms\n"; + } + + return true; +} + +int main() { + TestUtils::TestRunner runner("IO / Serialization Tests"); + + runner.run_test("Node Save/Load", test_node_save_load()); + runner.run_test("Graph Save/Load", test_graph_save_load()); + runner.run_test("Graph Save/Load Roundtrip Execution", test_graph_save_load_roundtrip_execution()); + runner.run_test("Graph Save/Load Supported Ops Roundtrip", test_graph_save_load_supported_ops_roundtrip()); + runner.run_test("Graph Save For Inspection", test_graph_save_for_inspection()); + runner.run_test("Save/Load Preserves Recurrent Cache Persistence", + test_save_load_preserves_recurrent_cache_persistence()); + runner.run_test("Save/Load Preserves KV Cache Num Slots", + test_save_load_preserves_kv_cache_num_slots()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_metal_parity.cpp b/cactus-graph/tests/test_metal_parity.cpp new file mode 100644 index 000000000..ae8f2d548 --- /dev/null +++ b/cactus-graph/tests/test_metal_parity.cpp @@ -0,0 +1,572 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include +#include +#include + +using namespace TestUtils; + +namespace { + +std::vector<__fp16> random_fp16(size_t n, float lo = -1.0f, float hi = 1.0f, uint32_t seed = 42) { + std::mt19937 gen(seed); + std::uniform_real_distribution dis(lo, hi); + std::vector<__fp16> out(n); + for (auto& v : out) v = static_cast<__fp16>(dis(gen)); + return out; +} + +struct ParityCase { + std::vector> input_shapes; + std::vector> input_data; + std::vector> input_data_f32; + std::function&)> build; + float tolerance = 5e-2f; + + size_t add_input(const std::vector& shape, float lo = -1.0f, float hi = 1.0f) { + size_t n = 1; + for (size_t d : shape) n *= d; + input_shapes.push_back(shape); + input_data.push_back(random_fp16(n, lo, hi, 42 + (uint32_t)input_shapes.size())); + input_data_f32.push_back({}); + return input_shapes.size() - 1; + } + + size_t add_input_data(const std::vector& shape, std::vector<__fp16> data) { + input_shapes.push_back(shape); + input_data.push_back(std::move(data)); + input_data_f32.push_back({}); + return input_shapes.size() - 1; + } + + size_t add_input_f32(const std::vector& shape, const std::vector& data) { + input_shapes.push_back(shape); + input_data.push_back({}); + input_data_f32.push_back(data); + return input_shapes.size() - 1; + } + + std::vector run(const char* backend) { + if (cactus_backend_select(backend) != 0) return {}; + CactusGraph graph; + std::vector ids; + for (size_t i = 0; i < input_shapes.size(); ++i) { + bool f32 = !input_data_f32[i].empty(); + ids.push_back(graph.input(input_shapes[i], f32 ? Precision::FP32 : Precision::FP16)); + } + size_t out_id = build(graph, ids); + for (size_t i = 0; i < ids.size(); ++i) { + bool f32 = !input_data_f32[i].empty(); + graph.set_input(ids[i], + f32 ? static_cast(input_data_f32[i].data()) + : static_cast(input_data[i].data()), + f32 ? Precision::FP32 : Precision::FP16); + } + graph.execute(); + const auto& desc = graph.get_output_buffer(out_id); + std::vector result(desc.total_size); + if (desc.precision == Precision::FP16) { + const __fp16* p = static_cast(graph.get_output(out_id)); + for (size_t i = 0; i < result.size(); ++i) result[i] = static_cast(p[i]); + } else { + const float* p = static_cast(graph.get_output(out_id)); + for (size_t i = 0; i < result.size(); ++i) result[i] = p[i]; + } + graph.hard_reset(); + return result; + } + + bool check() { + std::vector cpu = run("cpu"); + std::vector metal = run("metal"); + if (cpu.size() != metal.size() || cpu.empty()) return false; + for (size_t i = 0; i < cpu.size(); ++i) { + float scale = std::max(1.0f, std::fabs(cpu[i])); + if (std::fabs(cpu[i] - metal[i]) > tolerance * scale) { + std::cout << " mismatch at " << i << ": cpu=" << cpu[i] + << " metal=" << metal[i] << "\n"; + return false; + } + } + return true; + } +}; + +bool metal_present() { + bool ok = cactus_backend_select("metal") == 0; + return ok; +} + +bool parity_per_op_pins() { + auto run_pinned = [&](const char* global, ComputeBackend matmul_backend) { + ParityCase c; + size_t a = c.add_input({8, 32}); + size_t b = c.add_input({32, 16}); + c.build = [matmul_backend](CactusGraph& g, const std::vector& in) { + size_t m = g.matmul(in[0], in[1], false, matmul_backend); + size_t s = g.silu(m); + return g.scalar_add(s, 1.0f); + }; + return c.run(global); + }; + std::vector ref = run_pinned("cpu", ComputeBackend::CPU); + std::vector cpu_pin_under_metal = run_pinned("metal", ComputeBackend::CPU); + std::vector metal_pin_under_cpu = run_pinned("cpu", ComputeBackend::METAL); + if (ref.empty() || cpu_pin_under_metal.size() != ref.size() || metal_pin_under_cpu.size() != ref.size()) + return false; + for (size_t i = 0; i < ref.size(); ++i) { + float scale = std::max(1.0f, std::fabs(ref[i])); + if (std::fabs(ref[i] - cpu_pin_under_metal[i]) > 5e-2f * scale) return false; + if (std::fabs(ref[i] - metal_pin_under_cpu[i]) > 5e-2f * scale) return false; + } + return true; +} + +bool parity_unary(size_t (CactusGraph::*fn)(size_t, ComputeBackend), float lo = -1.0f, float hi = 1.0f) { + ParityCase c; + c.add_input({4, 33}, lo, hi); + c.build = [fn](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], cactus_default_backend()); }; + return c.check(); +} + +bool parity_scalar(size_t (CactusGraph::*fn)(size_t, float, ComputeBackend), float p, float lo = -1.0f, float hi = 1.0f) { + ParityCase c; + c.add_input({4, 33}, lo, hi); + c.build = [fn, p](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], p, cactus_default_backend()); }; + return c.check(); +} + +bool parity_reduce(size_t (CactusGraph::*fn)(size_t, int, ComputeBackend)) { + for (int axis = 0; axis < 3; ++axis) { + ParityCase c; + c.add_input({3, 5, 17}); + c.build = [fn, axis](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], axis, cactus_default_backend()); }; + if (!c.check()) return false; + } + return true; +} + +bool parity_not_equal() { + ParityCase c; + c.add_input({2, 31}); + c.add_input({2, 31}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.not_equal(in[0], in[1]); }; + return c.check(); +} + +bool parity_broadcast_binary() { + { + ParityCase c; + c.add_input({4, 33}); + c.add_input({1, 33}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.multiply(in[0], in[1]); }; + if (!c.check()) return false; + } + { + ParityCase c; + c.add_input({4, 33}); + c.add_input({1, 33}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.not_equal(in[0], in[1]); }; + if (!c.check()) return false; + } + return true; +} + +bool parity_concat() { + for (int axis = 0; axis < 2; ++axis) { + ParityCase c; + c.add_input({3, 7}); + c.add_input({axis == 0 ? 5u : 3u, axis == 0 ? 7u : 11u}); + c.build = [axis](CactusGraph& g, const std::vector& in) { return g.concat(in[0], in[1], axis); }; + if (!c.check()) return false; + } + return true; +} + +bool parity_gather() { + ParityCase c; + c.add_input({16, 24}); + c.add_input_f32({6}, {3.0f, 0.0f, 15.0f, 7.0f, 7.0f, 2.0f}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.gather(in[0], in[1]); }; + return c.check(); +} + +bool parity_rope(bool gptj) { + ParityCase c; + c.add_input({1, 6, 4, 32}); + c.build = [gptj](CactusGraph& g, const std::vector& in) { + return gptj ? g.rope_gptj(in[0], 10000.0f, 3, 16) : g.rope(in[0], 10000.0f, 3); + }; + return c.check(); +} + +bool parity_maxpool1d() { + ParityCase c; + c.add_input({2, 5, 29}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.maxpool1d(in[0], 3, 2); }; + return c.check(); +} + +bool parity_bilinear() { + ParityCase c; + c.add_input({64, 12}); + c.build = [](CactusGraph& g, const std::vector& in) { + return g.bilinear_interpolation(in[0], 13, 11, false); + }; + return c.check(); +} + +bool parity_conv1d() { + ParityCase c; + c.add_input({2, 6, 31}); + c.add_input({4, 6, 5}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.conv1d(in[0], in[1], 2); }; + return c.check(); +} + +bool parity_conv1d_k7s3() { + ParityCase c; + c.add_input({1, 8, 40}); + c.add_input({8, 7, 16}); + c.add_input({16}); + c.build = [](CactusGraph& g, const std::vector& in) { return g.conv1d_k7s3(in[0], in[1], in[2]); }; + return c.check(); +} + +bool parity_conv1d_causal() { + for (size_t dil : {1u, 2u}) { + ParityCase c; + c.add_input({1, 21, 8}); + c.add_input({8, 1, 4}); + c.build = [dil](CactusGraph& g, const std::vector& in) { + return g.conv1d_causal(in[0], in[1], 4, dil); + }; + if (!c.check()) return false; + } + return true; +} + +bool parity_conv1d_dw_k9() { + ParityCase c; + c.add_input({1, 25, 6}); + c.add_input({6, 1, 9}); + c.add_input({6}); + c.build = [](CactusGraph& g, const std::vector& in) { + return g.conv1d_same_depthwise_k9(in[0], in[1], in[2]); + }; + return c.check(); +} + +bool parity_conv1d_pointwise() { + ParityCase c; + c.add_input({1, 13, 16}); + c.add_input({24, 16}); + c.add_input({24}); + c.build = [](CactusGraph& g, const std::vector& in) { + return g.conv1d_pointwise(in[0], in[1], in[2]); + }; + return c.check(); +} + +bool parity_conv2d(size_t (CactusGraph::*fn)(size_t, size_t, size_t, ComputeBackend), bool depthwise, bool pointwise) { + ParityCase c; + c.add_input({1, 4, 11, 13}); + if (depthwise) c.add_input({4, 1, 3, 3}); + else if (pointwise) c.add_input({6, 4, 1, 1}); + else c.add_input({6, 4, 3, 3}); + c.add_input({pointwise || depthwise ? (depthwise ? 4u : 6u) : 6u}); + c.build = [fn](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], in[1], in[2], cactus_default_backend()); }; + return c.check(); +} + +bool parity_batchnorm() { + ParityCase c; + c.add_input({2, 5, 9}); + c.add_input({5}); + c.add_input({5}); + c.add_input({5}); + c.add_input({5}, 0.5f, 1.5f); + c.build = [](CactusGraph& g, const std::vector& in) { + return g.batchnorm(in[0], in[1], in[2], in[3], in[4], 1); + }; + return c.check(); +} + +bool parity_groupnorm() { + ParityCase c; + c.add_input({2, 8, 6}); + c.add_input({8}); + c.add_input({8}); + c.build = [](CactusGraph& g, const std::vector& in) { + return g.groupnorm(in[0], in[1], in[2], 4); + }; + return c.check(); +} + +bool parity_cumsum() { + for (int axis = 0; axis < 2; ++axis) { + ParityCase c; + c.add_input({4, 19}); + c.build = [axis](CactusGraph& g, const std::vector& in) { return g.cumsum(in[0], axis); }; + if (!c.check()) return false; + } + return true; +} + +bool parity_softmax() { + ParityCase c; + c.add_input({5, 333}, -4.0f, 4.0f); + c.tolerance = 5e-3f; + c.build = [](CactusGraph& g, const std::vector& in) { return g.softmax(in[0], -1); }; + return c.check(); +} + +bool parity_rms_norm() { + for (size_t rows : {1u, 7u}) { + ParityCase c; + c.add_input({rows, 256}); + c.add_input({256}, 0.5f, 1.5f); + c.tolerance = 1e-2f; + c.build = [](CactusGraph& g, const std::vector& in) { + return g.rms_norm(in[0], in[1], 1e-5f); + }; + if (!c.check()) return false; + } + return true; +} + +bool parity_residual_rms_norm() { + for (size_t rows : {1u, 6u}) { + ParityCase c; + c.add_input({rows, 192}); + c.add_input({rows, 192}); + c.add_input({192}, 0.5f, 1.5f); + c.tolerance = 1e-2f; + c.build = [](CactusGraph& g, const std::vector& in) { + return g.rms_norm(g.add(in[0], in[1]), in[2], 1e-5f); + }; + if (!c.check()) return false; + } + return true; +} + +bool parity_elemwise_chain() { + ParityCase c; + c.add_input({4, 96}); + c.add_input({4, 96}); + c.add_input({4, 96}); + c.tolerance = 1e-2f; + c.build = [](CactusGraph& g, const std::vector& in) { + return g.add(g.multiply(g.silu(in[0]), in[1]), g.scalar_multiply(in[2], 0.5f)); + }; + return c.check(); +} + +bool parity_flash_attention() { + for (size_t T : {64u, 89u, 121u}) { + ParityCase c; + const size_t H = 2, D = 64, S = 729; + c.add_input({1, T, H, D}); + c.add_input({1, S, H, D}); + c.add_input({1, S, H, D}); + c.tolerance = 2e-2f; + c.build = [D](CactusGraph& g, const std::vector& in) { + return g.attention(in[0], in[1], in[2], 1.0f / std::sqrt((float)D), false); + }; + if (!c.check()) return false; + } + return true; +} + +bool parity_attention_causal() { + ParityCase c; + const size_t T = 37, H = 4, HKV = 2, D = 32; + c.add_input({1, T, H, D}); + c.add_input({1, T, HKV, D}); + c.add_input({1, T, HKV, D}); + c.tolerance = 2e-2f; + c.build = [D](CactusGraph& g, const std::vector& in) { + return g.attention(in[0], in[1], in[2], 1.0f / std::sqrt((float)D), true); + }; + return c.check(); +} + +bool parity_equality_exact() { + const size_t n = 2 * 31; + std::vector<__fp16> a = random_fp16(n, -1.0f, 1.0f, 7); + std::vector<__fp16> b = random_fp16(n, -1.0f, 1.0f, 8); + for (size_t i = 0; i < n; i += 2) b[i] = a[i]; + { + ParityCase c; + c.add_input_data({2, 31}, a); + c.add_input_data({2, 31}, b); + c.tolerance = 1e-3f; + c.build = [](CactusGraph& g, const std::vector& in) { return g.not_equal(in[0], in[1]); }; + if (!c.check()) return false; + } + { + std::vector<__fp16> z = a; + for (size_t i = 0; i < n; i += 2) z[i] = (__fp16)0.0f; + ParityCase c; + c.add_input_data({2, 31}, z); + c.tolerance = 1e-3f; + c.build = [](CactusGraph& g, const std::vector& in) { return g.scalar_not_equal(in[0], 0.0f); }; + if (!c.check()) return false; + } + return true; +} + +struct CacheStepData { + std::vector<__fp16> q, k, v; + size_t tokens; +}; + +std::vector run_cached_attention(const char* backend, size_t ceiling, size_t window, size_t sink, + size_t h, size_t kv, size_t d, + const std::vector& steps, size_t collect_from) { + if (cactus_backend_select(backend) != 0) return {}; + const size_t sentinel = std::numeric_limits::max(); + const float scale = 1.0f / std::sqrt((float)d); + std::vector collected; + CactusGraph g; + size_t kc = g.kv_cache_state(ceiling, kv, d, window, sink); + size_t vc = g.kv_cache_state(ceiling, kv, d, window, sink); + for (size_t si = 0; si < steps.size(); ++si) { + const auto& data = steps[si]; + size_t s = data.tokens; + size_t iq = g.input({1, s, h, d}, Precision::FP16); + size_t ik = g.input({1, s, kv, d}, Precision::FP16); + size_t iv = g.input({1, s, kv, d}, Precision::FP16); + g.set_input(iq, const_cast<__fp16*>(data.q.data()), Precision::FP16); + g.set_input(ik, const_cast<__fp16*>(data.k.data()), Precision::FP16); + g.set_input(iv, const_cast<__fp16*>(data.v.data()), Precision::FP16); + g.kv_cache_append(ik, kc, window, sink); + g.kv_cache_append(iv, vc, window, sink); + size_t attn = g.attention_cached(iq, ik, iv, kc, vc, scale, sentinel, window); + g.execute(); + if (si >= collect_from) { + const __fp16* p = static_cast(g.get_output(attn)); + for (size_t i = 0; i < s * h * d; ++i) collected.push_back((float)p[i]); + } + g.soft_reset(); + } + g.hard_reset(); + return collected; +} + +bool check_cached_attention(size_t ceiling, size_t window, size_t sink, + const std::vector& chunk_sizes, size_t decode_steps, + float tolerance, size_t collect_from = 0) { + const size_t h = 4, kv = 2, d = 32; + std::vector steps; + uint32_t seed = 100; + for (size_t s : chunk_sizes) { + steps.push_back({random_fp16(s * h * d, -1.0f, 1.0f, seed), + random_fp16(s * kv * d, -1.0f, 1.0f, seed + 1), + random_fp16(s * kv * d, -1.0f, 1.0f, seed + 2), s}); + seed += 3; + } + for (size_t i = 0; i < decode_steps; ++i) { + steps.push_back({random_fp16(h * d, -1.0f, 1.0f, seed), + random_fp16(kv * d, -1.0f, 1.0f, seed + 1), + random_fp16(kv * d, -1.0f, 1.0f, seed + 2), 1}); + seed += 3; + } + std::vector cpu = run_cached_attention("cpu", ceiling, window, sink, h, kv, d, steps, collect_from); + std::vector metal = run_cached_attention("metal", ceiling, window, sink, h, kv, d, steps, collect_from); + if (cpu.size() != metal.size() || cpu.empty()) return false; + for (size_t i = 0; i < cpu.size(); ++i) { + float scale = std::max(1.0f, std::fabs(cpu[i])); + if (std::fabs(cpu[i] - metal[i]) > tolerance * scale) { + std::cout << " mismatch at " << i << " (token " << i / (4 * 32) << "): cpu=" << cpu[i] + << " metal=" << metal[i] << "\n"; + return false; + } + } + return true; +} + +bool parity_sliding_window_cache() { + if (!check_cached_attention(/*ceiling=*/64, /*window=*/8, /*sink=*/2, + /*chunks=*/{5}, /*decode_steps=*/12, 5e-2f)) return false; + return check_cached_attention(/*ceiling=*/64, /*window=*/8, /*sink=*/2, + /*chunks=*/{4, 8}, /*decode_steps=*/6, 5e-2f, /*collect_from=*/2); +} + +bool parity_cache_growth() { + return check_cached_attention(/*ceiling=*/2048, /*window=*/0, /*sink=*/2, + /*chunks=*/{600}, /*decode_steps=*/3, 5e-2f); +} + +} + +int main() { + TestRunner runner("Metal Parity (CPU vs Metal per-op)"); + if (!metal_present()) { +#if defined(__APPLE__) + if (std::getenv("CACTUS_PARITY_ALLOW_SKIP")) { + std::cout << "Metal backend unavailable; skipping parity suite (CACTUS_PARITY_ALLOW_SKIP).\n"; + return 0; + } + std::cout << "FAIL: Metal backend unavailable on Apple hardware " + "(set CACTUS_PARITY_ALLOW_SKIP=1 to skip).\n"; + return 1; +#else + std::cout << "Metal backend unavailable; skipping parity suite.\n"; + return 0; +#endif + } + + runner.run_test("abs", parity_unary(&CactusGraph::abs)); + runner.run_test("scalar_exp", parity_unary(&CactusGraph::scalar_exp)); + runner.run_test("scalar_sqrt", parity_unary(&CactusGraph::scalar_sqrt, 0.1f, 2.0f)); + runner.run_test("scalar_cos", parity_unary(&CactusGraph::scalar_cos)); + runner.run_test("scalar_sin", parity_unary(&CactusGraph::scalar_sin)); + runner.run_test("scalar_log", parity_unary(&CactusGraph::scalar_log, 0.1f, 3.0f)); + runner.run_test("gelu_erf", parity_unary(&CactusGraph::gelu_erf)); + runner.run_test("sigmoid", parity_unary(&CactusGraph::sigmoid)); + runner.run_test("pow", parity_scalar(&CactusGraph::pow, 2.0f, 0.1f, 2.0f)); + runner.run_test("leaky_relu", parity_scalar(&CactusGraph::leaky_relu, 0.1f)); + runner.run_test("scalar_not_equal", parity_scalar(&CactusGraph::scalar_not_equal, 0.0f)); + runner.run_test("not_equal", parity_not_equal()); + runner.run_test("broadcast_binary", parity_broadcast_binary()); + runner.run_test("sum", parity_reduce(&CactusGraph::sum)); + runner.run_test("mean", parity_reduce(&CactusGraph::mean)); + runner.run_test("variance", parity_reduce(&CactusGraph::variance)); + runner.run_test("min", parity_reduce(&CactusGraph::min)); + runner.run_test("max", parity_reduce(&CactusGraph::max)); + runner.run_test("cumsum", parity_cumsum()); + runner.run_test("concat", parity_concat()); + runner.run_test("gather", parity_gather()); + runner.run_test("rope", parity_rope(false)); + runner.run_test("rope_gptj", parity_rope(true)); + runner.run_test("maxpool1d", parity_maxpool1d()); + runner.run_test("bilinear_interpolation", parity_bilinear()); + runner.run_test("conv1d", parity_conv1d()); + runner.run_test("conv1d_k7s3", parity_conv1d_k7s3()); + runner.run_test("conv1d_causal", parity_conv1d_causal()); + runner.run_test("conv1d_same_depthwise_k9", parity_conv1d_dw_k9()); + runner.run_test("conv1d_pointwise", parity_conv1d_pointwise()); + runner.run_test("conv2d_k3s2p1", parity_conv2d(&CactusGraph::conv2d_k3s2p1, false, false)); + runner.run_test("conv2d_k3s1p1", parity_conv2d(&CactusGraph::conv2d_k3s1p1, false, false)); + runner.run_test("conv2d_depthwise_k3s2p1", parity_conv2d(&CactusGraph::conv2d_depthwise_k3s2p1, true, false)); + runner.run_test("conv2d_pointwise_1x1", parity_conv2d(&CactusGraph::conv2d_pointwise_1x1, false, true)); + runner.run_test("batchnorm", parity_batchnorm()); + runner.run_test("groupnorm", parity_groupnorm()); + runner.run_test("softmax", parity_softmax()); + runner.run_test("rms_norm", parity_rms_norm()); + runner.run_test("residual_rms_norm (fused)", parity_residual_rms_norm()); + runner.run_test("elemwise_chain (fused)", parity_elemwise_chain()); + runner.run_test("flash_attention (T%64!=0)", parity_flash_attention()); + runner.run_test("attention_causal", parity_attention_causal()); + runner.run_test("equality_exact", parity_equality_exact()); + runner.run_test("sliding_window_cache (ring wrap)", parity_sliding_window_cache()); + runner.run_test("cache_growth", parity_cache_growth()); + runner.run_test("per_op_backend_pins", parity_per_op_pins()); + + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp new file mode 100644 index 000000000..ffb26372e --- /dev/null +++ b/cactus-graph/tests/test_nn.cpp @@ -0,0 +1,729 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include +#include + +using namespace TestUtils; + +bool test_matrix_multiplication() { + TestUtils::FP16TestFixture fixture("Matrix Multiplication"); + + size_t input_a = fixture.create_input({2, 3}); + size_t input_b = fixture.create_input({3, 2}); + size_t matmul_result = fixture.graph().matmul(input_a, input_b, false); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {1, 2, 3, 4, 5, 6}; + fixture.set_input_data(input_a, data_a); + fixture.set_input_data(input_b, data_b); + fixture.execute(); + + std::vector<__fp16> expected = {22, 28, 49, 64}; + return fixture.verify_output(matmul_result, expected); +} + +static size_t align_offset_test(size_t offset, size_t alignment) { + size_t remainder = offset % alignment; + return remainder == 0 ? offset : offset + (alignment - remainder); +} + +static Precision cq_precision_for_bits(uint32_t bits) { + switch (bits) { + case 1: return Precision::CQ1; + case 2: return Precision::CQ2; + case 3: return Precision::CQ3; + case 4: return Precision::CQ4; + default: throw std::runtime_error("unsupported CQ bits"); + } +} + +static void write_test_cq_weights( + const std::filesystem::path& path, + uint32_t bits, + size_t K, + size_t N, + size_t gs, + const std::vector& packed, + const std::vector<__fp16>& codebook, + const std::vector<__fp16>& input_scale, + const std::vector<__fp16>& input_scale_recip, + const std::vector<__fp16>& norms, + const std::vector& left_signs, + const std::vector& right_signs, + const std::vector& permutation) { + constexpr uint32_t CACTUS_MAGIC = 0x54434143; + constexpr uint32_t FLAG_HAS_SCALES = 1 << 0; + constexpr size_t HEADER_SIZE = 84; + constexpr uint32_t alignment = 32; + + const size_t ng = K / gs; + const size_t scales_bytes = + codebook.size() * sizeof(__fp16) + + input_scale.size() * sizeof(__fp16) + + input_scale_recip.size() * sizeof(__fp16) + + norms.size() * sizeof(__fp16) + + left_signs.size() + + right_signs.size() + + permutation.size() * sizeof(uint32_t); + + std::ofstream file(path, std::ios::binary); + if (!file) throw std::runtime_error("cannot open test CQ weights"); + + auto write_u32 = [&](uint32_t v) { file.write(reinterpret_cast(&v), sizeof(v)); }; + auto write_u64 = [&](uint64_t v) { file.write(reinterpret_cast(&v), sizeof(v)); }; + auto write_padding = [&](size_t bytes) { + char zero = 0; + for (size_t i = 0; i < bytes; ++i) file.write(&zero, 1); + }; + + write_u32(CACTUS_MAGIC); + write_u32(FLAG_HAS_SCALES); + write_u32(alignment); + write_u32(2); + write_u64(N); + write_u64(K); + write_u64(0); + write_u64(0); + write_u32(static_cast(cq_precision_for_bits(bits))); + write_u64(packed.size()); + write_u64(scales_bytes); + write_u32(static_cast(gs)); + write_u32(static_cast(ng)); + write_u64(N); + + size_t aligned_header = align_offset_test(HEADER_SIZE, alignment); + write_padding(aligned_header - HEADER_SIZE); + + file.write(reinterpret_cast(codebook.data()), codebook.size() * sizeof(__fp16)); + file.write(reinterpret_cast(input_scale.data()), input_scale.size() * sizeof(__fp16)); + file.write(reinterpret_cast(input_scale_recip.data()), input_scale_recip.size() * sizeof(__fp16)); + file.write(reinterpret_cast(norms.data()), norms.size() * sizeof(__fp16)); + file.write(reinterpret_cast(left_signs.data()), left_signs.size()); + file.write(reinterpret_cast(right_signs.data()), right_signs.size()); + file.write(reinterpret_cast(permutation.data()), permutation.size() * sizeof(uint32_t)); + + size_t data_start = align_offset_test(aligned_header + scales_bytes, alignment); + write_padding(data_start - (aligned_header + scales_bytes)); + file.write(reinterpret_cast(packed.data()), packed.size()); + if (!file) throw std::runtime_error("failed writing test CQ weights"); +} + +bool test_matmul_cq() { + // Test graph-level CQ matmul dispatch for every supported bit width. + const size_t M = 2, K = 128, N = 8; + const size_t gs = 128; + const size_t ng = K / gs; + + for (uint32_t bits : {1u, 2u, 3u, 4u}) { + const uint32_t cb_size = 1u << bits; + std::mt19937 gen(42 + bits); + std::uniform_real_distribution dist(-1.f, 1.f); + + std::vector<__fp16> A(M * K); + for (auto& v : A) v = static_cast<__fp16>(dist(gen)); + + uint32_t pgb = cactus_quant_packed_group_bytes(bits, gs); + std::vector packed(N * ng * pgb); + for (auto& v : packed) v = static_cast(gen() & 0xFF); + + std::vector<__fp16> codebook(cb_size), input_scale(K), input_scale_recip(K), norms(N * ng); + std::vector left_signs(gs), right_signs(gs); + std::vector permutation(gs); + + for (auto& v : codebook) v = static_cast<__fp16>(dist(gen)); + for (size_t i = 0; i < K; i++) { + float s = 0.5f + std::abs(dist(gen)); + input_scale[i] = static_cast<__fp16>(s); + input_scale_recip[i] = static_cast<__fp16>(1.f / s); + } + for (auto& v : norms) v = static_cast<__fp16>(dist(gen) * 0.1f); + for (auto& v : left_signs) v = (gen() & 1) ? 1 : -1; + for (auto& v : right_signs) v = (gen() & 1) ? 1 : -1; + for (uint32_t i = 0; i < gs; i++) permutation[i] = i; + + CactusQuantMatrix mat{ + .bits = bits, .K = static_cast(K), .N = static_cast(N), + .group_size = static_cast(gs), .num_groups = static_cast(ng), + .flags = 0, + .codebook = codebook.data(), + .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), + .norms = norms.data(), + .packed_indices = packed.data(), + .left_signs = left_signs.data(), + .right_signs = right_signs.data(), + .permutation = permutation.data(), + .rotation = nullptr, + .expanded = nullptr, + .norm_f32 = nullptr, + }; + + std::vector<__fp16> direct(M * N, static_cast<__fp16>(0)); + cactus_quant_matmul(&mat, A.data(), static_cast(M), direct.data()); + + auto path = std::filesystem::temp_directory_path() / + ("cactus_graph_cq" + std::to_string(bits) + "_matmul.weights"); + write_test_cq_weights(path, bits, K, N, gs, packed, codebook, input_scale, + input_scale_recip, norms, left_signs, right_signs, permutation); + + CactusGraph g; + size_t ia = g.input({M, K}, Precision::FP16); + size_t iw = g.mmap_weights(path.string()); + size_t out = g.matmul(ia, iw, true); + g.set_input(ia, A.data(), Precision::FP16); + g.execute(); + __fp16* graph_out = static_cast<__fp16*>(g.get_output(out)); + + bool ok = true; + for (size_t i = 0; i < M * N; i++) { + float actual = static_cast(graph_out[i]); + float expected = static_cast(direct[i]); + if (!std::isfinite(actual) || std::abs(actual - expected) > 1e-3f) { + ok = false; + break; + } + } + g.hard_reset(); + std::filesystem::remove(path); + if (!ok) return false; + } + return true; +} + +bool test_attention_int8_hybrid() { + const size_t b = 1, s = 1, h = 2, kv = 2, d = 16; + const size_t cache_len = 4; + const size_t num_groups = (d + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + + std::vector<__fp16> q(b * s * h * d), k_new(b * s * kv * d), v_new(b * s * kv * d); + std::vector k_cached(cache_len * kv * d, 10); + std::vector v_cached(cache_len * kv * d, 5); + std::vector k_scales(cache_len * kv * num_groups, 0.01f); + std::vector v_scales(cache_len * kv * num_groups, 0.01f); + + fill_random_fp16(q); + fill_random_fp16(k_new); + fill_random_fp16(v_new); + + float scale = 1.0f / std::sqrt(static_cast(d)); + + CactusGraph g; + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + size_t out = g.attention_int8_hybrid(iq, ik, iv, scale, 0, + k_cached.data(), v_cached.data(), + k_scales.data(), v_scales.data(), + cache_len, kv, d); + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k_new.data(), Precision::FP16); + g.set_input(iv, v_new.data(), Precision::FP16); + g.execute(); + + __fp16* result = static_cast<__fp16*>(g.get_output(out)); + size_t out_size = b * s * h * d; + for (size_t i = 0; i < out_size; i++) { + if (!std::isfinite(static_cast(result[i]))) return false; + } + + bool has_nonzero = false; + for (size_t i = 0; i < out_size; i++) { + if (std::abs(static_cast(result[i])) > 1e-6f) has_nonzero = true; + } + return has_nonzero; +} + +bool test_transpose() { + TestUtils::FP16TestFixture fixture("Transpose"); + + size_t input_a = fixture.create_input({2, 3}); + size_t transpose_result = fixture.graph().transpose(input_a); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected = {1, 4, 2, 5, 3, 6}; + return fixture.verify_output(transpose_result, expected); +} + +bool test_rms_norm() { + TestUtils::FP16TestFixture fixture("RMS Norm"); + + size_t input_a = fixture.create_input({1, 8}); + size_t weight = fixture.create_input({8}); + size_t norm_result = fixture.graph().rms_norm(input_a, weight); + + std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + std::vector<__fp16> weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + fixture.set_input_data(input_a, input_data); + fixture.set_input_data(weight, weight_data); + fixture.execute(); + + float sum_squares = 0.0f; + for (auto val : input_data) { + float v = static_cast(val); + sum_squares += v * v; + } + float rms = sqrtf(sum_squares / 8.0f + 1e-5f); + float inv_rms = 1.0f / rms; + + std::vector<__fp16> expected; + for (size_t i = 0; i < input_data.size(); i++) { + expected.push_back(static_cast<__fp16>(static_cast(input_data[i]) * inv_rms * static_cast(weight_data[i]))); + } + + return fixture.verify_output(norm_result, expected, 0.01f); +} + +bool test_softmax() { + TestUtils::FP16TestFixture fixture("Softmax"); + + size_t input_a = fixture.create_input({2, 3}); + size_t softmax_result = fixture.graph().softmax(input_a, -1); + + std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + fixture.set_input_data(input_a, input_data); + fixture.execute(); + + std::vector<__fp16> expected = {0.09003f, 0.24473f, 0.66524f, 0.09003f, 0.24473f, 0.66524f}; + return fixture.verify_output(softmax_result, expected, 0.01f); +} + +bool test_attention() { + TestUtils::FP16TestFixture fixture("Attention"); + + size_t query = fixture.create_input({1, 2, 1, 4}); + size_t key = fixture.create_input({1, 2, 1, 4}); + size_t value = fixture.create_input({1, 2, 1, 4}); + + size_t attention_result = fixture.graph().attention(query, key, value, 0.5f); + (void)attention_result; + + std::vector<__fp16> q_data = {1, 0, 0, 0, 0, 1, 0, 0}; + std::vector<__fp16> k_data = {1, 0, 0, 0, 0, 1, 0, 0}; + std::vector<__fp16> v_data = {1, 2, 3, 4, 5, 6, 7, 8}; + + fixture.set_input_data(query, q_data); + fixture.set_input_data(key, k_data); + fixture.set_input_data(value, v_data); + fixture.execute(); + + return true; +} + +bool test_reduction_operations() { + TestUtils::FP16TestFixture fixture("Reduction Operations"); + + size_t input_a = fixture.create_input({2, 3}); + size_t sum_all = fixture.graph().sum(input_a, -1); + size_t sum_axis0 = fixture.graph().sum(input_a, 0); + size_t sum_axis1 = fixture.graph().sum(input_a, 1); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected_all = {21}; + std::vector<__fp16> expected_axis0 = {5, 7, 9}; + std::vector<__fp16> expected_axis1 = {6, 15}; + + return fixture.verify_output(sum_all, expected_all) && + fixture.verify_output(sum_axis0, expected_axis0) && + fixture.verify_output(sum_axis1, expected_axis1); +} + +bool test_mean_operations() { + TestUtils::FP16TestFixture fixture("Mean Operations"); + + size_t input_a = fixture.create_input({2, 4}); + size_t mean_all = fixture.graph().mean(input_a, -1); + size_t mean_axis0 = fixture.graph().mean(input_a, 0); + size_t mean_axis1 = fixture.graph().mean(input_a, 1); + + std::vector<__fp16> data_a = {2, 4, 6, 8, 10, 12, 14, 16}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected_all = {9}; + std::vector<__fp16> expected_axis0 = {6, 8, 10, 12}; + std::vector<__fp16> expected_axis1 = {5, 13}; + + return fixture.verify_output(mean_all, expected_all) && + fixture.verify_output(mean_axis0, expected_axis0) && + fixture.verify_output(mean_axis1, expected_axis1); +} + +bool test_variance_operations() { + TestUtils::FP16TestFixture fixture("Variance Operations"); + + size_t input_a = fixture.create_input({1, 4}); + size_t var_axis1 = fixture.graph().variance(input_a, 1); + + std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f}; + fixture.set_input_data(input_a, input_data); + fixture.execute(); + + std::vector<__fp16> expected = {1.25f}; + return fixture.verify_output(var_axis1, expected, 0.01f); +} + +bool test_min_max_operations() { + TestUtils::FP16TestFixture fixture("Min/Max Operations"); + + size_t input_a = fixture.create_input({2, 3}); + size_t min_axis0 = fixture.graph().min(input_a, 0); + size_t max_axis0 = fixture.graph().max(input_a, 0); + size_t min_axis1 = fixture.graph().min(input_a, 1); + size_t max_axis1 = fixture.graph().max(input_a, 1); + + std::vector<__fp16> data_a = {6, 2, 8, 1, 5, 3}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected_min_axis0 = {1, 2, 3}; + std::vector<__fp16> expected_max_axis0 = {6, 5, 8}; + std::vector<__fp16> expected_min_axis1 = {2, 1}; + std::vector<__fp16> expected_max_axis1 = {8, 5}; + + return fixture.verify_output(min_axis0, expected_min_axis0) && + fixture.verify_output(max_axis0, expected_max_axis0) && + fixture.verify_output(min_axis1, expected_min_axis1) && + fixture.verify_output(max_axis1, expected_max_axis1); +} + +bool test_stft() { + const size_t N = 2, C_in = 1, L = 8, K = 4, stride = 2, num_fft_bins = 2; + const size_t C_out = 2 * num_fft_bins; + const size_t out_len = (L - K) / stride + 1; + + std::vector<__fp16> weight_data = { + (__fp16) 1, (__fp16) 1, (__fp16) 1, (__fp16) 1, + (__fp16) 1, (__fp16) 0, (__fp16)-1, (__fp16) 0, + (__fp16) 0, (__fp16) 0, (__fp16) 0, (__fp16) 0, + (__fp16) 0, (__fp16)-1, (__fp16) 0, (__fp16) 1, + }; + std::vector<__fp16> input_data = { + (__fp16)1, (__fp16)2, (__fp16)3, (__fp16)4, (__fp16)5, (__fp16)6, (__fp16)7, (__fp16)8, + (__fp16)0, (__fp16)1, (__fp16)0, (__fp16)-1, (__fp16)0, (__fp16)1, (__fp16)0, (__fp16)-1, + }; + + TestUtils::FP16TestFixture fx; + size_t inp = fx.create_input({N, C_in, L}); + size_t wt = fx.create_input({C_out, C_in, K}); + size_t out = fx.graph().stft(inp, wt, stride, num_fft_bins); + + if (fx.graph().get_output_buffer(out).shape != std::vector{N, C_out, out_len}) return false; + + fx.set_input_data(inp, input_data); + fx.set_input_data(wt, weight_data); + fx.execute(); + + const __fp16* cplx = fx.get_output(out); + const size_t out_bs = C_out * out_len; + const float tol = 0.1f; + + for (size_t t = 0; t < out_len; ++t) { + if (std::abs((float)cplx[1 * out_len + t] - (-2.0f)) > tol) return false; + if (std::abs((float)cplx[(1 + num_fft_bins) * out_len + t] - 2.0f) > tol) return false; + } + + const float batch1_bin1_imag[3] = {-2.0f, 2.0f, -2.0f}; + for (size_t t = 0; t < out_len; ++t) { + if (std::abs((float)cplx[out_bs + 1 * out_len + t] - 0.0f) > tol) return false; + if (std::abs((float)cplx[out_bs + (1 + num_fft_bins) * out_len + t] - batch1_bin1_imag[t]) > tol) return false; + } + + return true; +} + +template +static bool run_layernorm_case( + size_t batch, size_t feat, bool with_bias, float epsilon, + float weight_scale, float bias_val) +{ + const size_t total = batch * feat; + + std::vector input_f(total), weight_f(feat), bias_f(feat); + for (size_t b = 0; b < batch; ++b) + for (size_t j = 0; j < feat; ++j) + input_f[b * feat + j] = static_cast(j + 1); + for (size_t j = 0; j < feat; ++j) { + weight_f[j] = weight_scale; + bias_f[j] = bias_val; + } + + std::vector inp_data(total), w_data(feat), b_data(feat); + for (size_t i = 0; i < total; ++i) inp_data[i] = static_cast(input_f[i]); + for (size_t j = 0; j < feat; ++j) { + w_data[j] = static_cast(weight_f[j]); + b_data[j] = static_cast(bias_f[j]); + } + + TestUtils::TestFixture fx; + size_t inp_id = fx.create_input({batch, feat}); + size_t w_id = fx.create_input({feat}); + fx.set_input_data(inp_id, inp_data); + fx.set_input_data(w_id, w_data); + + size_t out_id; + if (with_bias) { + size_t b_id = fx.create_input({feat}); + fx.set_input_data(b_id, b_data); + out_id = fx.graph().layernorm(inp_id, w_id, b_id, epsilon); + } else { + out_id = fx.graph().layernorm(inp_id, w_id, epsilon); + } + + fx.execute(); + + std::vector expected(total); + for (size_t b = 0; b < batch; ++b) { + float mean = 0.0f; + for (size_t j = 0; j < feat; ++j) mean += input_f[b * feat + j]; + mean /= static_cast(feat); + float var = 0.0f; + for (size_t j = 0; j < feat; ++j) { + float d = input_f[b * feat + j] - mean; + var += d * d; + } + var /= static_cast(feat); + float inv_std = 1.0f / std::sqrt(var + epsilon); + for (size_t j = 0; j < feat; ++j) { + float val = (input_f[b * feat + j] - mean) * inv_std * weight_f[j]; + if (with_bias) val += bias_f[j]; + expected[b * feat + j] = static_cast(val); + } + } + + return fx.verify_output(out_id, expected, TestUtils::default_tolerance()); +} + +bool test_layernorm() { + struct Case { size_t batch, feat; bool fp32, with_bias; float epsilon, weight_scale, bias_val; }; + const std::vector cases = { + {1, 1, false, false, 1e-5f, 1.0f, 0.0f}, + {1, 7, false, false, 1e-5f, 1.0f, 0.0f}, + {1, 8, false, false, 1e-5f, 1.0f, 0.0f}, + {4, 8, false, false, 1e-5f, 1.0f, 0.0f}, + {4, 8, false, true, 1e-5f, 1.0f, 0.0f}, + {4, 8, true, false, 1e-5f, 1.0f, 0.0f}, + {4, 8, true, true, 1e-5f, 1.0f, 0.0f}, + {1, 8, false, false, 1.0f, 1.0f, 0.0f}, + {2, 16, false, true, 1e-5f, 0.5f, 0.3f}, + }; + + for (const auto& c : cases) { + bool ok = c.fp32 + ? run_layernorm_case(c.batch, c.feat, c.with_bias, c.epsilon, c.weight_scale, c.bias_val) + : run_layernorm_case<__fp16>(c.batch, c.feat, c.with_bias, c.epsilon, c.weight_scale, c.bias_val); + if (!ok) return false; + } + return true; +} + +static void apply_activation_reference(Activation act, const __fp16* in, __fp16* out, size_t n) { + switch (act) { + case Activation::GELU: cactus_gelu_f16(in, out, n); break; + case Activation::GELU_ERF: cactus_gelu_f16_erf(in, out, n); break; + case Activation::RELU: cactus_relu_f16(in, out, n); break; + case Activation::SIGMOID: cactus_sigmoid_f16(in, out, n); break; + case Activation::TANH: cactus_tanh_f16(in, out, n); break; + case Activation::SILU: cactus_silu_f16(in, out, n); break; + } +} + +// Verifies the MoE layer dispatches to the requested activation rather than +// silently falling back to SILU. With identity expert weights and a single +// token routed to a single expert at probability 1, the layer reduces to +// out = activation(hidden), so each activation must match its standalone kernel. +bool test_moe_activations() { + const size_t H = 4; + std::vector<__fp16> hidden = {(__fp16)-2.0f, (__fp16)-0.5f, (__fp16)0.5f, (__fp16)2.0f}; + + std::vector<__fp16> identity(H * H, (__fp16)0.0f); + for (size_t i = 0; i < H; ++i) identity[i * H + i] = (__fp16)1.0f; + + std::vector<__fp16> routing = {(__fp16)1.0f}; + std::vector topk = {0.0f}; + + const std::vector activations = { + Activation::SILU, Activation::GELU, Activation::GELU_ERF, + Activation::RELU, Activation::SIGMOID, Activation::TANH, + }; + + for (Activation act : activations) { + CactusGraph g; + size_t hidden_id = g.input({1, H}, Precision::FP16); + size_t routing_id = g.input({1, 1}, Precision::FP16); + size_t topk_id = g.input({1, 1}, Precision::FP32); + size_t w1_id = g.input({H, H}, Precision::FP16); + size_t w2_id = g.input({H, H}, Precision::FP16); + + size_t out = g.moe_layer(hidden_id, routing_id, topk_id, + {w1_id}, {w2_id}, + 1, 1, false, 1e-6f, 1.0f, act); + + g.set_input(hidden_id, hidden.data(), Precision::FP16); + g.set_input(routing_id, routing.data(), Precision::FP16); + g.set_input(topk_id, topk.data(), Precision::FP32); + g.set_input(w1_id, identity.data(), Precision::FP16); + g.set_input(w2_id, identity.data(), Precision::FP16); + g.execute(); + + std::vector<__fp16> expected(H); + apply_activation_reference(act, hidden.data(), expected.data(), H); + + __fp16* result = static_cast<__fp16*>(g.get_output(out)); + bool ok = TestUtils::compare_arrays(result, expected.data(), H, 0.02f); + g.hard_reset(); + if (!ok) return false; + } + return true; +} + +bool run_benchmarks() { + auto bench = [](const char* label, auto setup, auto run) { + setup(); + run(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) run(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + { + const size_t M = 1024, K = 1024, N = 1024; + std::vector<__fp16> a(M * K), b(N * K); + TestUtils::fill_random_fp16(a); + TestUtils::fill_random_fp16(b); + CactusGraph g; + size_t ia = g.input({M, K}, Precision::FP16); + size_t ib = g.input({N, K}, Precision::FP16); + g.matmul(ia, ib, true); + g.set_input(ia, a.data(), Precision::FP16); + g.set_input(ib, b.data(), Precision::FP16); + bench("matmul_f16 1024^3", []{}, [&]{ g.execute(); }); + } + { + // CQ4 matmul benchmark via cactus_quant_matmul (graph-level equivalent) + const size_t M = 1024, K = 1024, N = 1024, gs = 128; + const size_t ng = K / gs; + const uint32_t bits = 4, cb_size = 16; + std::mt19937 bgen(77); + std::uniform_real_distribution bdist(-1.f, 1.f); + + std::vector<__fp16> A(M * K), codebook(cb_size), input_sc(K), input_sc_r(K), norms_v(N * ng); + std::vector lsigns(gs), rsigns(gs); + std::vector perm(gs); + for (auto& v : A) v = static_cast<__fp16>(bdist(bgen)); + for (auto& v : codebook) v = static_cast<__fp16>(bdist(bgen)); + for (size_t i = 0; i < K; i++) { float s = 0.5f+std::abs(bdist(bgen)); input_sc[i]=(__fp16)s; input_sc_r[i]=(__fp16)(1.f/s); } + for (auto& v : norms_v) v = static_cast<__fp16>(bdist(bgen) * 0.1f); + for (auto& v : lsigns) v = (bgen()&1)?1:-1; + for (auto& v : rsigns) v = (bgen()&1)?1:-1; + for (uint32_t i = 0; i < gs; i++) perm[i] = i; + uint32_t pgb = cactus_quant_packed_group_bytes(bits, gs); + std::vector packed(N * ng * pgb); + for (auto& v : packed) v = static_cast(bgen() & 0xFF); + + CactusQuantMatrix mat{bits, (uint32_t)K, (uint32_t)N, (uint32_t)gs, (uint32_t)ng, + 0, + codebook.data(), input_sc.data(), input_sc_r.data(), norms_v.data(), + packed.data(), lsigns.data(), rsigns.data(), perm.data(), nullptr, nullptr, nullptr}; + + std::vector<__fp16> C(M * N); + bench("matmul_cq4 1024^3", []{}, [&]{ cactus_quant_matmul(&mat, A.data(), M, C.data()); }); + } + { + const size_t b = 1, s = 256, h = 16, kv = 8, d = 128; + std::vector<__fp16> q(b*s*h*d), k(b*s*kv*d), v(b*s*kv*d); + TestUtils::fill_random_fp16(q); + TestUtils::fill_random_fp16(k); + TestUtils::fill_random_fp16(v); + float scale = 1.0f / std::sqrt(static_cast(d)); + CactusGraph g; + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + g.attention(iq, ik, iv, scale); + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k.data(), Precision::FP16); + g.set_input(iv, v.data(), Precision::FP16); + bench("attention_f16 seq=256", []{}, [&]{ g.execute(); }); + } + { + const size_t b = 1, s = 1, h = 16, kv = 8, d = 128; + const size_t cache_len = 512; + std::vector<__fp16> q(b*s*h*d), k(b*s*kv*d), v(b*s*kv*d); + std::vector ck(cache_len*kv*d, 1), cv(cache_len*kv*d, 1); + size_t ng = (d + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + std::vector ks(cache_len*kv*ng, 0.01f), vs(cache_len*kv*ng, 0.01f); + TestUtils::fill_random_fp16(q); + TestUtils::fill_random_fp16(k); + TestUtils::fill_random_fp16(v); + float scale = 1.0f / std::sqrt(static_cast(d)); + CactusGraph g; + size_t iq = g.input({b, s, h, d}, Precision::FP16); + size_t ik = g.input({b, s, kv, d}, Precision::FP16); + size_t iv = g.input({b, s, kv, d}, Precision::FP16); + g.attention_int8_hybrid(iq, ik, iv, scale, 0, + ck.data(), cv.data(), ks.data(), vs.data(), cache_len, kv, d); + g.set_input(iq, q.data(), Precision::FP16); + g.set_input(ik, k.data(), Precision::FP16); + g.set_input(iv, v.data(), Precision::FP16); + bench("attention_int8 cache=512", []{}, [&]{ g.execute(); }); + } + { + const size_t batch = 1024, dim = 1024; + std::vector<__fp16> in(batch * dim), w(dim); + TestUtils::fill_random_fp16(in); + for (size_t i = 0; i < dim; i++) w[i] = static_cast<__fp16>(1.0f); + CactusGraph g; + size_t ii = g.input({batch, dim}, Precision::FP16); + size_t iw = g.input({dim}, Precision::FP16); + g.rms_norm(ii, iw, 1e-6f); + g.set_input(ii, in.data(), Precision::FP16); + g.set_input(iw, w.data(), Precision::FP16); + bench("rms_norm 1024x1024", []{}, [&]{ g.execute(); }); + } + { + const size_t rows = 1024, cols = 1024; + std::vector<__fp16> in(rows * cols); + TestUtils::fill_random_fp16(in); + CactusGraph g; + size_t ii = g.input({rows, cols}, Precision::FP16); + g.softmax(ii); + g.set_input(ii, in.data(), Precision::FP16); + bench("softmax 1024x1024", []{}, [&]{ g.execute(); }); + } + return true; +} + +int main() { + TestUtils::TestRunner runner("Neural Network Ops Tests"); + + runner.run_test("Matrix Multiplication", test_matrix_multiplication()); + runner.run_test("MatMul CQ", test_matmul_cq()); + runner.run_test("Transpose", test_transpose()); + runner.run_test("RMS Norm", test_rms_norm()); + runner.run_test("Softmax", test_softmax()); + runner.run_test("Attention", test_attention()); + runner.run_test("Attention INT8 Hybrid", test_attention_int8_hybrid()); + runner.run_test("Reduction Operations", test_reduction_operations()); + runner.run_test("Mean Operations", test_mean_operations()); + runner.run_test("Variance Operations", test_variance_operations()); + runner.run_test("Min/Max Operations", test_min_max_operations()); + runner.run_test("STFT Complex", test_stft()); + runner.run_test("LayerNorm", test_layernorm()); + runner.run_test("MoE Activations", test_moe_activations()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_ops.cpp b/cactus-graph/tests/test_ops.cpp new file mode 100644 index 000000000..009f7d1c4 --- /dev/null +++ b/cactus-graph/tests/test_ops.cpp @@ -0,0 +1,758 @@ +#include "test_utils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace TestUtils; + +bool test_abs() { + TestUtils::FP16TestFixture fixture("absval"); + size_t input = fixture.create_input({2, 3}); + + size_t abs_result = fixture.graph().abs(input); + std::vector<__fp16> data = {1, -2, 3, 4, -5, -6}; + + fixture.set_input_data(input, data); + fixture.execute(); + + std::vector<__fp16> expected = {1, 2, 3, 4, 5, 6}; + return fixture.verify_output(abs_result, expected); +} + +bool test_concat() { + TestUtils::FP16TestFixture fixture("Concat"); + + size_t input_a = fixture.create_input({2, 3}); + size_t input_b = fixture.create_input({2, 5}); + size_t concat_result = fixture.graph().concat(input_a, input_b, 1); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; + std::vector<__fp16> expected = {1, 2, 3, 1, 2, 3, 4, 5, 4, 5, 6, 6, 7, 8, 9, 10}; + fixture.set_input_data(input_a, data_a); + fixture.set_input_data(input_b, data_b); + fixture.execute(); + + return fixture.verify_output(concat_result, expected); +} + +bool test_cat() { + TestUtils::FP16TestFixture fixture("Cat (multiple input tensors)"); + + size_t input_a = fixture.create_input({2, 3}); + size_t input_b = fixture.create_input({2, 5}); + size_t input_c = fixture.create_input({2, 2}); + + size_t cat_result = fixture.graph().cat({input_a, input_b, input_c}, 1); + + std::vector<__fp16> data_a = {1, 2, 3, + 4, 5, 6}; + + std::vector<__fp16> data_b = {1, 2, 3, 4, 5, + 6, 7, 8, 9, 10}; + + std::vector<__fp16> data_c = {-1, -2, + -1, -2}; + + std::vector<__fp16> expected = { + 1, 2, 3, 1, 2, 3, 4, 5, -1, -2, + 4, 5, 6, 6, 7, 8, 9, 10, -1, -2 + }; + + fixture.set_input_data(input_a, data_a); + fixture.set_input_data(input_b, data_b); + fixture.set_input_data(input_c, data_c); + + fixture.execute(); + + return fixture.verify_output(cat_result, expected); +} + +bool test_view() { + TestUtils::FP16TestFixture fixture("View"); + + size_t input = fixture.create_input({2, 3}); + size_t view_result = fixture.graph().view(input, {3, 2}); + + std::vector<__fp16> data = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> expected = {1, 2, 3, 4, 5, 6}; + + fixture.set_input_data(input, data); + fixture.execute(); + + return fixture.verify_output(view_result, expected); +} + +bool test_flatten() { + TestUtils::FP16TestFixture fixture("Flatten"); + + size_t input = fixture.create_input({2, 3}); + size_t flatten_result = fixture.graph().flatten(input); + + std::vector<__fp16> data = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> expected = {1, 2, 3, 4, 5, 6}; + + fixture.set_input_data(input, data); + fixture.execute(); + + return fixture.verify_output(flatten_result, expected); +} + +bool test_basic_operations() { + TestUtils::FP16TestFixture fixture("Basic Operations"); + + size_t input_a = fixture.create_input({2, 3}); + size_t input_b = fixture.create_input({2, 3}); + size_t add_result = fixture.graph().add(input_a, input_b); + size_t mul_result = fixture.graph().multiply(add_result, input_a); + size_t scalar_result = fixture.graph().scalar_multiply(mul_result, 2.0f); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {2, 3, 4, 5, 6, 7}; + + fixture.set_input_data(input_a, data_a); + fixture.set_input_data(input_b, data_b); + fixture.execute(); + + std::vector<__fp16> expected(6); + for (int i = 0; i < 6; i++) { + float result = ((static_cast(data_a[i]) + static_cast(data_b[i])) * static_cast(data_a[i])) * 2.0f; + expected[i] = static_cast<__fp16>(result); + } + + return fixture.verify_output(scalar_result, expected); +} + +bool test_basic_addition() { + return TestUtils::test_basic_operation( + "Addition", + [](CactusGraph& graph, size_t a, size_t b) { return graph.add(a, b); }, + {1, 2, 3, 4}, + {5, 6, 7, 8}, + {6, 8, 10, 12} + ); +} + +bool test_basic_subtraction() { + return TestUtils::test_basic_operation( + "Subtraction", + [](CactusGraph& graph, size_t a, size_t b) { return graph.subtract(a, b); }, + {10, 8, 6, 4}, + {2, 3, 1, 2}, + {8, 5, 5, 2} + ); +} + +bool test_basic_multiplication() { + return TestUtils::test_basic_operation( + "Multiplication", + [](CactusGraph& graph, size_t a, size_t b) { return graph.multiply(a, b); }, + {2, 3, 4, 5}, + {3, 4, 2, 2}, + {6, 12, 8, 10} + ); +} + +bool test_basic_division() { + return TestUtils::test_basic_operation( + "Division", + [](CactusGraph& graph, size_t a, size_t b) { return graph.divide(a, b); }, + {12, 15, 8, 9}, + {3, 5, 2, 3}, + {4, 3, 4, 3} + ); +} + +bool test_scalar_operations() { + TestUtils::FP16TestFixture fixture("Scalar Operations"); + + size_t input_a = fixture.create_input({4}); + size_t add_result = fixture.graph().scalar_add(input_a, 5.0f); + size_t mul_result = fixture.graph().scalar_multiply(add_result, 2.0f); + + std::vector<__fp16> data_a = {1, 2, 3, 4}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected = {12, 14, 16, 18}; + return fixture.verify_output(mul_result, expected); +} + +bool test_scalar_subtract_divide() { + TestUtils::FP16TestFixture fixture("Scalar Subtract/Divide"); + + size_t input_a = fixture.create_input({4}); + size_t sub_result = fixture.graph().scalar_subtract(input_a, 2.0f); + size_t div_result = fixture.graph().scalar_divide(input_a, 2.0f); + + std::vector<__fp16> data_a = {10, 8, 6, 4}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected_sub = {8, 6, 4, 2}; + std::vector<__fp16> expected_div = {5, 4, 3, 2}; + return fixture.verify_output(sub_result, expected_sub) && + fixture.verify_output(div_result, expected_div); +} + +bool test_scalar_math_functions() { + TestUtils::FP16TestFixture fixture("Scalar Math Functions"); + + size_t input_a = fixture.create_input({3}); + size_t exp_result = fixture.graph().scalar_exp(input_a); + size_t sqrt_result = fixture.graph().scalar_sqrt(input_a); + size_t cos_result = fixture.graph().scalar_cos(input_a); + size_t sin_result = fixture.graph().scalar_sin(input_a); + size_t log_result = fixture.graph().scalar_log(input_a); + + std::vector<__fp16> input_data = {0.5f, 1.0f, 4.0f}; + fixture.set_input_data(input_a, input_data); + fixture.execute(); + + std::vector<__fp16> exp_expected = {1.64872f, 2.71828f, 54.5982f}; + std::vector<__fp16> sqrt_expected = {0.70711f, 1.0f, 2.0f}; + std::vector<__fp16> cos_expected = {0.87758f, 0.54030f, -0.65364f}; + std::vector<__fp16> sin_expected = {0.47943f, 0.84147f, -0.75680f}; + std::vector<__fp16> log_expected = {-0.69315f, 0.0f, 1.38629f}; + + return fixture.verify_output(exp_result, exp_expected, 0.01f) && + fixture.verify_output(sqrt_result, sqrt_expected, 0.01f) && + fixture.verify_output(cos_result, cos_expected, 0.01f) && + fixture.verify_output(sin_result, sin_expected, 0.01f) && + fixture.verify_output(log_result, log_expected, 0.01f); +} + +bool test_scalar_operations_with_pow() { + TestUtils::FP16TestFixture fixture("Scalar Operations with Pow"); + + size_t input_a = fixture.create_input({4}); + size_t add_result = fixture.graph().scalar_add(input_a, 5.0f); + size_t mul_result = fixture.graph().scalar_multiply(add_result, 2.0f); + size_t pow_result = fixture.graph().pow(mul_result, 2.0f); + + std::vector<__fp16> data_a = {1, 2, 3, 4}; + fixture.set_input_data(input_a, data_a); + fixture.execute(); + + std::vector<__fp16> expected = {144, 196, 256, 324}; + return fixture.verify_output(pow_result, expected); +} + +bool test_gather_operation() { + CactusGraph graph; + + size_t embeddings = graph.input({5, 3}, Precision::FP16); + size_t indices = graph.input({2, 2}, Precision::INT8); + size_t gathered = graph.gather(embeddings, indices); + + std::vector<__fp16> emb_data = { + 1, 2, 3, + 4, 5, 6, + 7, 8, 9, + 10, 11, 12, + 13, 14, 15 + }; + std::vector idx_data = {0, 2, 4, 1}; + + graph.set_input(embeddings, emb_data.data(), Precision::FP16); + graph.set_input(indices, idx_data.data(), Precision::INT8); + graph.execute(); + + __fp16* output = static_cast<__fp16*>(graph.get_output(gathered)); + std::vector<__fp16> expected = { + 1, 2, 3, + 7, 8, 9, + 13, 14, 15, + 4, 5, 6 + }; + + for (size_t i = 0; i < expected.size(); ++i) { + if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-3f) { + graph.hard_reset(); + return false; + } + } + + graph.hard_reset(); + return true; +} + +bool test_gather_1d_tensor() { + CactusGraph graph; + + size_t tensor = graph.input({8}, Precision::FP16); + size_t indices = graph.input({3}, Precision::INT8); + size_t gathered = graph.gather(tensor, indices); + + std::vector<__fp16> tensor_data = {10, 20, 30, 40, 50, 60, 70, 80}; + std::vector idx_data = {7, 2, 0}; + + graph.set_input(tensor, tensor_data.data(), Precision::FP16); + graph.set_input(indices, idx_data.data(), Precision::INT8); + graph.execute(); + + __fp16* output = static_cast<__fp16*>(graph.get_output(gathered)); + std::vector<__fp16> expected = {80, 30, 10}; + + for (size_t i = 0; i < expected.size(); ++i) { + if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-3f) { + graph.hard_reset(); + return false; + } + } + + graph.hard_reset(); + return true; +} + +bool test_gather_3d_tensor() { + TestUtils::FP16TestFixture fixture("Gather 3D Tensor"); + + size_t tensor = fixture.create_input({3, 2, 4}); + size_t indices = fixture.graph().input({2}, Precision::INT8); + size_t gathered = fixture.graph().gather(tensor, indices); + + std::vector<__fp16> tensor_data = { + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, + 9.0f, 10.0f, 11.0f, 12.0f, + 13.0f, 14.0f, 15.0f, 16.0f, + 17.0f, 18.0f, 19.0f, 20.0f, + 21.0f, 22.0f, 23.0f, 24.0f + }; + std::vector idx_data = {2, 0}; + + fixture.set_input_data(tensor, tensor_data); + fixture.graph().set_input(indices, idx_data.data(), Precision::INT8); + fixture.execute(); + + std::vector<__fp16> expected = { + 17.0f, 18.0f, 19.0f, 20.0f, + 21.0f, 22.0f, 23.0f, 24.0f, + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f + }; + + return fixture.verify_output(gathered, expected); +} + +bool test_mmap_gather() { + CactusGraph graph; + + std::vector<__fp16> embeddings_data = { + 1.0f, 2.0f, 3.0f, + 4.0f, 5.0f, 6.0f, + 7.0f, 8.0f, 9.0f, + 10.0f, 11.0f, 12.0f + }; + + size_t temp_embeddings = graph.input({4, 3}, Precision::FP16); + graph.set_input(temp_embeddings, embeddings_data.data(), Precision::FP16); + + const std::string temp_file = "test_embeddings.bin"; + GraphFile::save_node(graph, temp_embeddings, temp_file); + + graph.hard_reset(); + + size_t mmap_embeddings = graph.mmap_embeddings(temp_file); + size_t indices = graph.input({3}, Precision::INT8); + size_t gathered = graph.gather(mmap_embeddings, indices); + + std::vector idx_data = {2, 0, 3}; + graph.set_input(indices, idx_data.data(), Precision::INT8); + graph.execute(); + + std::vector<__fp16> expected = { + 7.0f, 8.0f, 9.0f, + 1.0f, 2.0f, 3.0f, + 10.0f, 11.0f, 12.0f + }; + + __fp16* output = static_cast<__fp16*>(graph.get_output(gathered)); + bool passed = true; + for (size_t i = 0; i < expected.size(); i++) { + if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 0.01f) { + passed = false; + break; + } + } + + std::remove(temp_file.c_str()); + + return passed; +} + +bool test_embedding_operation() { + CactusGraph graph; + + const size_t vocab_size = 4; + const size_t hidden_dim = 8; + + std::vector<__fp16> emb_data(vocab_size * hidden_dim); + for (size_t row = 0; row < vocab_size; ++row) { + for (size_t k = 0; k < hidden_dim; ++k) { + emb_data[row * hidden_dim + k] = static_cast<__fp16>((row + 1) * 10 + k); + } + } + + size_t embeddings = graph.input({vocab_size, hidden_dim}, Precision::FP16); + graph.set_input(embeddings, emb_data.data(), Precision::FP16); + + size_t indices = graph.input({4}, Precision::INT8); + size_t embedded = graph.embedding(embeddings, indices); + + std::vector idx_data = {0, 2, 3, 1}; + graph.set_input(indices, idx_data.data(), Precision::INT8); + graph.execute(); + + __fp16* output = static_cast<__fp16*>(graph.get_output(embedded)); + + std::vector expected = { + 10, 11, 12, 13, 14, 15, 16, 17, + 30, 31, 32, 33, 34, 35, 36, 37, + 40, 41, 42, 43, 44, 45, 46, 47, + 20, 21, 22, 23, 24, 25, 26, 27 + }; + + for (size_t i = 0; i < expected.size(); ++i) { + float out_val = static_cast(output[i]); + if (std::abs(out_val - expected[i]) > 0.5f) { + std::cerr << "Embedding mismatch at " << i << ": got " << out_val + << ", expected " << expected[i] << std::endl; + return false; + } + } + + return true; +} + +bool test_cq_embedding_operation() { + const uint32_t bits = 2; + const uint32_t vocab_size = 4; + const uint32_t hidden_dim = 16; + const uint32_t group_size = 16; + const uint32_t num_groups = hidden_dim / group_size; + const uint32_t cb_size = 1u << bits; + + std::vector<__fp16> codebook(cb_size); + codebook[0] = static_cast<__fp16>(-1.0f); + codebook[1] = static_cast<__fp16>(-0.5f); + codebook[2] = static_cast<__fp16>(0.5f); + codebook[3] = static_cast<__fp16>(1.0f); + + std::vector<__fp16> input_scale_recip(hidden_dim, static_cast<__fp16>(1.0f)); + + std::vector<__fp16> norms(vocab_size * num_groups); + for (uint32_t r = 0; r < vocab_size; r++) + for (uint32_t g = 0; g < num_groups; g++) + norms[r * num_groups + g] = static_cast<__fp16>(1.0f); + + std::vector left_signs(group_size, 1); + std::vector right_signs(group_size, 1); + + std::vector permutation(group_size); + for (uint32_t i = 0; i < group_size; i++) permutation[i] = i; + uint32_t pgb = cactus_quant_packed_group_bytes(bits, group_size); + std::vector packed(vocab_size * num_groups * pgb, 0); + for (uint32_t row = 0; row < vocab_size; row++) { + uint8_t idx = static_cast(row % cb_size); + for (uint32_t g = 0; g < num_groups; g++) { + uint8_t* p = packed.data() + (row * num_groups + g) * pgb; + // Pack: each byte holds 4 x 2-bit indices + for (uint32_t byte_i = 0; byte_i < pgb; byte_i++) { + p[byte_i] = static_cast(idx | (idx << 2) | (idx << 4) | (idx << 6)); + } + } + } + + std::vector<__fp16> expected(vocab_size * hidden_dim); + for (uint32_t row = 0; row < vocab_size; row++) { + cactus_quant_dequantize_hadamard_embedding_row( + bits, hidden_dim, group_size, num_groups, row, + packed.data(), codebook.data(), norms.data(), + input_scale_recip.data(), left_signs.data(), right_signs.data(), + permutation.data(), expected.data() + row * hidden_dim); + } + + CactusGraph graph; + + size_t emb_node = graph.input({vocab_size, hidden_dim}, Precision::CQ2); + graph.set_input(emb_node, packed.data(), Precision::CQ2); + + auto& buffer = graph.nodes_[graph.node_index_map_[emb_node]]->output_buffer; + buffer.group_size = group_size; + buffer.num_groups = num_groups; + buffer.cq_codebook = codebook.data(); + buffer.cq_input_scale = nullptr; + buffer.cq_input_scale_recip = input_scale_recip.data(); + buffer.cq_norms = norms.data(); + buffer.cq_left_signs = left_signs.data(); + buffer.cq_right_signs = right_signs.data(); + buffer.cq_permutation = permutation.data(); + buffer.cq_rotation = nullptr; + buffer.cq_flags = 0; + + size_t indices = graph.input({3}, Precision::INT8); + size_t embedded = graph.embedding(emb_node, indices); + + std::vector idx_data = {0, 2, 3}; + graph.set_input(indices, idx_data.data(), Precision::INT8); + graph.execute(); + + __fp16* output = static_cast<__fp16*>(graph.get_output(embedded)); + + size_t lookup[] = {0, 2, 3}; + for (size_t i = 0; i < 3; i++) { + const __fp16* exp_row = expected.data() + lookup[i] * hidden_dim; + for (size_t k = 0; k < hidden_dim; k++) { + float out_val = static_cast(output[i * hidden_dim + k]); + float exp_val = static_cast(exp_row[k]); + if (std::abs(out_val - exp_val) > 0.01f) { + std::cerr << "CQ Embedding mismatch at [" << i << "," << k << "]: got " + << out_val << ", expected " << exp_val << std::endl; + return false; + } + } + } + + return true; +} + +static void set_cq4_index(std::vector& packed, uint32_t K, uint32_t row, uint32_t col, uint8_t value) { + const size_t byte_index = static_cast(row) * (K / 2) + col / 2; + if ((col & 1u) == 0) { + packed[byte_index] = static_cast((packed[byte_index] & 0xF0) | (value & 0x0F)); + } else { + packed[byte_index] = static_cast((packed[byte_index] & 0x0F) | ((value & 0x0F) << 4)); + } +} + +static std::vector make_interleaved_cq4(const std::vector& row_major, uint32_t N, uint32_t K) { + const uint32_t pgb = cactus_quant_packed_group_bytes(4, K); + std::vector interleaved(static_cast(N / 4) * 4 * pgb, 0); + for (uint32_t row = 0; row < N; ++row) { + const uint8_t* src = row_major.data() + static_cast(row) * pgb; + uint8_t* panel = interleaved.data() + static_cast(row / 4) * 4 * pgb; + for (uint32_t k = 0; k < K; ++k) { + const uint8_t idx = static_cast((src[k / 2] >> ((k & 1u) * 4)) & 0x0F); + const uint32_t chunk = k / 8; + const uint32_t sub = k & 7u; + uint8_t& dst = panel[chunk * 16 + (row & 3u) * 4 + (sub & 3u)]; + if (sub & 4u) dst = static_cast((dst & 0x0F) | (idx << 4)); + else dst = static_cast((dst & 0xF0) | idx); + } + } + return interleaved; +} + +bool test_orthogonal_interleaved_embedding_row() { + const uint32_t N = 8; + const uint32_t K = 32; + std::vector<__fp16> codebook(16); + for (uint32_t i = 0; i < 16; ++i) codebook[i] = static_cast<__fp16>((static_cast(i) - 8) * 0.125f); + std::vector<__fp16> input_scale_recip(K, static_cast<__fp16>(1.0f)); + std::vector<__fp16> norms(N); + for (uint32_t row = 0; row < N; ++row) norms[row] = static_cast<__fp16>(0.5f + 0.03125f * row); + std::vector<__fp16> rotation(static_cast(K) * K, static_cast<__fp16>(0.0f)); + for (uint32_t i = 0; i < K; ++i) rotation[static_cast(i) * K + i] = static_cast<__fp16>(1.0f); + + std::vector row_major(static_cast(N) * K / 2, 0); + for (uint32_t row = 0; row < N; ++row) + for (uint32_t col = 0; col < K; ++col) + set_cq4_index(row_major, K, row, col, static_cast((row * 3 + col) & 0x0F)); + std::vector interleaved = make_interleaved_cq4(row_major, N, K); + + for (uint32_t row : {0u, 1u, 3u, 4u, 7u}) { + std::vector<__fp16> ref(K); + std::vector<__fp16> got(K); + cactus_quant_dequantize_orthogonal_embedding_row( + 4, K, row, row_major.data(), codebook.data(), norms.data(), input_scale_recip.data(), + rotation.data(), CACTUS_QUANT_FLAG_ORTHOGONAL, ref.data()); + cactus_quant_dequantize_orthogonal_embedding_row( + 4, K, row, interleaved.data(), codebook.data(), norms.data(), input_scale_recip.data(), + rotation.data(), CACTUS_QUANT_FLAG_ORTHOGONAL | CACTUS_QUANT_FLAG_INTERLEAVED_4ROW, got.data()); + for (uint32_t i = 0; i < K; ++i) { + if (std::abs(static_cast(ref[i]) - static_cast(got[i])) > 0.001f) return false; + } + } + return true; +} + +bool test_orthogonal_interleaved_lmhead_matmul() { + const uint32_t N = 8; + const uint32_t K = 32; + std::vector<__fp16> codebook(16); + for (uint32_t i = 0; i < 16; ++i) codebook[i] = static_cast<__fp16>((static_cast(i) - 8) * 0.125f); + std::vector<__fp16> input_scale_recip(K, static_cast<__fp16>(1.0f)); + std::vector<__fp16> norms(N); + for (uint32_t row = 0; row < N; ++row) norms[row] = static_cast<__fp16>(0.05f + 0.003f * row); + std::vector<__fp16> rotation(static_cast(K) * K, static_cast<__fp16>(0.0f)); + for (uint32_t i = 0; i < K; ++i) { + uint32_t j = (i * 5) % K; + rotation[static_cast(i) * K + j] = static_cast<__fp16>((i & 1u) ? -1.0f : 1.0f); + } + std::vector row_major(static_cast(N) * K / 2, 0); + for (uint32_t row = 0; row < N; ++row) + for (uint32_t col = 0; col < K; ++col) + set_cq4_index(row_major, K, row, col, static_cast((row + col * 5) & 0x0F)); + std::vector interleaved = make_interleaved_cq4(row_major, N, K); + std::vector<__fp16> activation(static_cast(2) * K); + for (uint32_t i = 0; i < K; ++i) { + activation[i] = static_cast<__fp16>((static_cast(i % 9) - 4) * 0.1f); + activation[K + i] = static_cast<__fp16>((static_cast((i * 3) % 11) - 5) * 0.07f); + } + + CactusQuantMatrix row_mat{ + .bits = 4, .K = K, .N = N, .group_size = K, .num_groups = 1, + .flags = CACTUS_QUANT_FLAG_ORTHOGONAL, .codebook = codebook.data(), + .input_scale = nullptr, .input_scale_recip = input_scale_recip.data(), .norms = norms.data(), + .packed_indices = row_major.data(), .left_signs = nullptr, .right_signs = nullptr, + .permutation = nullptr, .rotation = rotation.data(), .expanded = nullptr, .norm_f32 = nullptr, + }; + CactusQuantMatrix inter_mat = row_mat; + inter_mat.flags = CACTUS_QUANT_FLAG_ORTHOGONAL | CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + inter_mat.packed_indices = interleaved.data(); + std::vector<__fp16> ref(N); + std::vector<__fp16> got(N); + cactus_quant_orthogonal_matmul(&row_mat, activation.data(), 1, ref.data()); + cactus_quant_orthogonal_matmul(&inter_mat, activation.data(), 1, got.data()); + + float diff2 = 0.0f; + float ref2 = 0.0f; + for (uint32_t i = 0; i < N; ++i) { + const float d = static_cast(got[i]) - static_cast(ref[i]); + diff2 += d * d; + ref2 += static_cast(ref[i]) * static_cast(ref[i]); + } + if (std::sqrt(diff2 / std::max(ref2, 1.0e-12f)) >= 0.08f) return false; + + std::vector<__fp16> ref_fallback(static_cast(2) * N); + std::vector<__fp16> got_fallback(static_cast(2) * N); + cactus_quant_orthogonal_matmul(&row_mat, activation.data(), 2, ref_fallback.data()); + cactus_quant_orthogonal_matmul(&inter_mat, activation.data(), 2, got_fallback.data()); + for (uint32_t i = 0; i < 2 * N; ++i) { + if (std::abs(static_cast(got_fallback[i]) - static_cast(ref_fallback[i])) > 0.001f) { + return false; + } + } + return true; +} + +bool test_embedding_from_file() { + CactusGraph graph; + + std::vector<__fp16> embeddings_data = { + 1.0f, 5.0f, 9.0f, + 2.0f, 6.0f, 10.0f, + 3.0f, 7.0f, 11.0f, + 4.0f, 8.0f, 12.0f + }; + + size_t temp_embeddings = graph.input({4, 3}, Precision::FP16); + graph.set_input(temp_embeddings, embeddings_data.data(), Precision::FP16); + + const std::string temp_file = "test_embedding.bin"; + GraphFile::save_node(graph, temp_embeddings, temp_file); + + graph.hard_reset(); + + size_t indices = graph.input({3}, Precision::INT8); + size_t embedded = graph.embedding(temp_file, indices); + + std::vector idx_data = {2, 0, 3}; + graph.set_input(indices, idx_data.data(), Precision::INT8); + graph.execute(); + + std::vector<__fp16> expected = { + 3.0f, 7.0f, 11.0f, + 1.0f, 5.0f, 9.0f, + 4.0f, 8.0f, 12.0f + }; + + __fp16* output = static_cast<__fp16*>(graph.get_output(embedded)); + bool passed = true; + for (size_t i = 0; i < expected.size(); i++) { + if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 0.01f) { + passed = false; + break; + } + } + + std::remove(temp_file.c_str()); + + return passed; +} + +bool run_benchmarks() { + const size_t N = 1024 * 1024; + std::vector<__fp16> data_a(N), data_b(N); + TestUtils::fill_random_fp16(data_a); + TestUtils::fill_random_fp16(data_b); + + auto bench = [](const char* label, auto fn) { + fn(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + bench("add 1M (graph)", [&]{ + CactusGraph g; + size_t a = g.input({N}, Precision::FP16); + size_t b = g.input({N}, Precision::FP16); + g.add(a, b); + g.set_input(a, data_a.data(), Precision::FP16); + g.set_input(b, data_b.data(), Precision::FP16); + g.execute(); + }); + + bench("scalar_multiply 1M (graph)", [&]{ + CactusGraph g; + size_t a = g.input({N}, Precision::FP16); + g.scalar_multiply(a, 2.0f); + g.set_input(a, data_a.data(), Precision::FP16); + g.execute(); + }); + + return true; +} + +int main() { + TestUtils::TestRunner runner("Ops Tests"); + + runner.run_test("Abs Operation", test_abs()); + runner.run_test("Concat Operation", test_concat()); + runner.run_test("Cat Operation", test_cat()); + runner.run_test("View Operation", test_view()); + runner.run_test("Flatten Operation", test_flatten()); + runner.run_test("Basic Operations", test_basic_operations()); + runner.run_test("Basic Addition", test_basic_addition()); + runner.run_test("Basic Subtraction", test_basic_subtraction()); + runner.run_test("Basic Multiplication", test_basic_multiplication()); + runner.run_test("Basic Division", test_basic_division()); + runner.run_test("Scalar Operations", test_scalar_operations()); + runner.run_test("Scalar Subtract/Divide", test_scalar_subtract_divide()); + runner.run_test("Scalar Math Functions", test_scalar_math_functions()); + runner.run_test("Scalar Operations with Pow", test_scalar_operations_with_pow()); + runner.run_test("Gather Operation", test_gather_operation()); + runner.run_test("Gather 1D Tensor", test_gather_1d_tensor()); + runner.run_test("Gather 3D Tensor", test_gather_3d_tensor()); + runner.run_test("Memory-Mapped Gather", test_mmap_gather()); + runner.run_test("Embedding Operation", test_embedding_operation()); + runner.run_test("CQ Embedding Operation", test_cq_embedding_operation()); + runner.run_test("Orthogonal Interleaved Embedding Row", test_orthogonal_interleaved_embedding_row()); + runner.run_test("Orthogonal Interleaved LMHead MatMul", test_orthogonal_interleaved_lmhead_matmul()); + runner.run_test("Embedding from File", test_embedding_from_file()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_precision.cpp b/cactus-graph/tests/test_precision.cpp new file mode 100644 index 000000000..f962c6c12 --- /dev/null +++ b/cactus-graph/tests/test_precision.cpp @@ -0,0 +1,194 @@ +#include "test_utils.h" +#include +#include +#include +#include + +using namespace TestUtils; + +bool test_fp16_precision() { + TestUtils::FP16TestFixture fixture("FP16 Precision"); + + size_t input_a = fixture.create_input({3}); + size_t input_b = fixture.create_input({3}); + size_t result_id = fixture.graph().add(input_a, input_b); + + std::vector<__fp16> data_a = {1.5f, 2.5f, 3.5f}; + std::vector<__fp16> data_b = {0.5f, 1.5f, 2.5f}; + fixture.set_input_data(input_a, data_a); + fixture.set_input_data(input_b, data_b); + fixture.execute(); + + std::vector<__fp16> expected = {2.0f, 4.0f, 6.0f}; + return fixture.verify_output(result_id, expected); +} + +bool test_broadcast_shape_compatibility() { + TestUtils::FP16TestFixture fixture("Broadcast Shape Compatibility"); + + size_t a_id = fixture.create_input({2, 3}); + size_t b_id = fixture.create_input({2, 1}); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; + std::vector<__fp16> data_b = {10, 20}; + fixture.set_input_data(a_id, data_a); + fixture.set_input_data(b_id, data_b); + + size_t result_id = fixture.graph().add(a_id, b_id); + fixture.execute(); + + std::vector<__fp16> expected = {11, 12, 13, 24, 25, 26}; + return fixture.verify_output(result_id, expected); +} + +bool test_broadcast_scalar_tensor() { + TestUtils::FP16TestFixture fixture("Broadcast Scalar Tensor"); + + size_t a_id = fixture.create_input({2, 2}); + size_t b_id = fixture.create_input({1}); + + std::vector<__fp16> data_a = {1, 2, 3, 4}; + std::vector<__fp16> data_b = {5}; + fixture.set_input_data(a_id, data_a); + fixture.set_input_data(b_id, data_b); + + size_t result_id = fixture.graph().add(a_id, b_id); + fixture.execute(); + + std::vector<__fp16> expected = {6, 7, 8, 9}; + return fixture.verify_output(result_id, expected); +} + +bool test_broadcast_different_ranks() { + TestUtils::FP16TestFixture fixture("Broadcast Different Ranks"); + + size_t a_id = fixture.create_input({2, 2, 3}); + size_t b_id = fixture.create_input({2, 3}); + + std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + std::vector<__fp16> data_b = {1, 1, 1, 2, 2, 2}; + fixture.set_input_data(a_id, data_a); + fixture.set_input_data(b_id, data_b); + + size_t result_id = fixture.graph().add(a_id, b_id); + fixture.execute(); + + std::vector<__fp16> expected = {2, 3, 4, 6, 7, 8, 8, 9, 10, 12, 13, 14}; + return fixture.verify_output(result_id, expected); +} + +bool test_broadcast_fp16_precision() { + TestUtils::FP16TestFixture fixture("Broadcast FP16 Precision"); + + size_t a_id = fixture.create_input({2, 2}); + size_t b_id = fixture.create_input({1}); + + std::vector<__fp16> data_a = {1.5f, 2.5f, 3.5f, 4.5f}; + std::vector<__fp16> data_b = {0.5f}; + fixture.set_input_data(a_id, data_a); + fixture.set_input_data(b_id, data_b); + + size_t result_id = fixture.graph().add(a_id, b_id); + fixture.execute(); + + std::vector<__fp16> expected = {2.0f, 3.0f, 4.0f, 5.0f}; + return fixture.verify_output(result_id, expected); +} + +bool test_precision_traits() { + assert(PrecisionTraits::size_of(Precision::INT8) == 1); + assert(PrecisionTraits::size_of(Precision::FP32) == 4); + return true; +} + +bool test_graph_precision_construction() { + TestUtils::FP16TestFixture fixture("Graph Precision Construction"); + + size_t fp16_id = fixture.create_input({2, 3}, Precision::FP16); + size_t fp32_id = fixture.create_input({3, 4}, Precision::FP32); + + if (fixture.graph().get_output_buffer(fp16_id).precision != Precision::FP16) return false; + if (fixture.graph().get_output_buffer(fp16_id).shape[0] != 2) return false; + if (fixture.graph().get_output_buffer(fp16_id).shape[1] != 3) return false; + if (fixture.graph().get_output_buffer(fp16_id).byte_size != 12) return false; + + if (fixture.graph().get_output_buffer(fp32_id).precision != Precision::FP32) return false; + if (fixture.graph().get_output_buffer(fp32_id).shape[0] != 3) return false; + if (fixture.graph().get_output_buffer(fp32_id).shape[1] != 4) return false; + if (fixture.graph().get_output_buffer(fp32_id).byte_size != 48) return false; + + return true; +} + +bool test_precision_conversion() { + TestUtils::FP16TestFixture fixture("Precision Conversion"); + + size_t fp16_id = fixture.create_input({2, 2}, Precision::FP16); + std::vector<__fp16> data = {1, 2, 3, 4}; + fixture.set_input_data(fp16_id, data); + + size_t fp32_converted_id = fixture.graph().precision_cast(fp16_id, Precision::FP32); + fixture.execute(); + + float* fp32_data = static_cast(fixture.graph().get_output(fp32_converted_id)); + for (size_t i = 0; i < 4; ++i) { + if (std::abs(fp32_data[i] - static_cast(data[i])) >= 1e-3f) return false; + } + + return true; +} + +bool run_benchmarks() { + const size_t N = 1024 * 1024; + std::vector<__fp16> fp16_data(N); + std::vector<__fp16> bcast_a(1024 * 1024), bcast_b(1024); + TestUtils::fill_random_fp16(fp16_data); + TestUtils::fill_random_fp16(bcast_a); + TestUtils::fill_random_fp16(bcast_b); + + auto bench = [](const char* label, auto fn) { + fn(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + bench("precision_cast FP16->FP32 1M", [&]{ + CactusGraph g; + size_t a = g.input({N}, Precision::FP16); + g.precision_cast(a, Precision::FP32); + g.set_input(a, fp16_data.data(), Precision::FP16); + g.execute(); + }); + + bench("broadcast_add 1024x1024", [&]{ + CactusGraph g; + size_t a = g.input({1024, 1024}, Precision::FP16); + size_t b = g.input({1, 1024}, Precision::FP16); + g.add(a, b); + g.set_input(a, bcast_a.data(), Precision::FP16); + g.set_input(b, bcast_b.data(), Precision::FP16); + g.execute(); + }); + + return true; +} + +int main() { + TestUtils::TestRunner runner("Precision & Broadcast Tests"); + + runner.run_test("FP16 Precision", test_fp16_precision()); + runner.run_test("Broadcast Shape Compatibility", test_broadcast_shape_compatibility()); + runner.run_test("Broadcast Scalar Tensor", test_broadcast_scalar_tensor()); + runner.run_test("Broadcast Different Ranks", test_broadcast_different_ranks()); + runner.run_test("Broadcast FP16 Precision", test_broadcast_fp16_precision()); + runner.run_test("Precision Traits", test_precision_traits()); + runner.run_test("Graph Precision Construction", test_graph_precision_construction()); + runner.run_test("Precision Conversion", test_precision_conversion()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-graph/tests/test_utils.h b/cactus-graph/tests/test_utils.h new file mode 100644 index 000000000..511fd446e --- /dev/null +++ b/cactus-graph/tests/test_utils.h @@ -0,0 +1,189 @@ +#pragma once + +#include "../cactus_graph.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace TestUtils { + +class Timer { +public: + Timer() : start_(std::chrono::high_resolution_clock::now()) {} + double elapsed_ms() const { + auto now = std::chrono::high_resolution_clock::now(); + return std::chrono::duration(now - start_).count(); + } +private: + std::chrono::time_point start_; +}; + +class TestRunner { +public: + explicit TestRunner(const std::string& suite_name) : suite_(suite_name), passed_(0), failed_(0) { + std::cout << "\n╔══════════════════════════════════════════════════════════════════════════════════════╗\n" + << "║ Running " << std::left << std::setw(73) << suite_name << "║\n" + << "╚══════════════════════════════════════════════════════════════════════════════════════╝\n"; + } + + void run_test(const std::string& name, bool result) { + if (result) { + std::cout << "✓ PASS │ " << std::left << std::setw(25) << name << std::endl; + passed_++; + } else { + std::cout << "✗ FAIL │ " << std::left << std::setw(25) << name << std::endl; + failed_++; + } + } + + void run_bench(const std::string&, bool result) { + if (!result) failed_++; + } + + void print_benchmarks_header() const { + std::cout << "── benchmarks ──────────────────────────────────────────────────────────────────────────\n"; + } + + void print_summary() const { + std::cout << "────────────────────────────────────────────────────────────────────────────────────────\n"; + if (failed_ == 0) { + std::cout << "✓ All " << passed_ << " tests passed!\n"; + } else { + std::cout << "✗ " << failed_ << " of " << (passed_ + failed_) << " tests failed!\n"; + } + } + + bool all_passed() const { return failed_ == 0; } + +private: + std::string suite_; + int passed_; + int failed_; +}; + +template +constexpr Precision default_precision() { + if constexpr (std::is_same_v) return Precision::INT8; + else if constexpr (std::is_same_v) return Precision::FP16; + else return Precision::FP32; +} + +template +constexpr float default_tolerance() { + if constexpr (std::is_same_v) return 1e-2f; + else return 1e-6f; +} + +template +bool compare_arrays(const T* actual, const T* expected, size_t count, float tolerance = default_tolerance()) { + for (size_t i = 0; i < count; ++i) { + if constexpr (std::is_same_v) { + if (std::abs(static_cast(actual[i]) - static_cast(expected[i])) > tolerance) return false; + } else if constexpr (std::is_floating_point_v) { + if (std::abs(actual[i] - expected[i]) > tolerance) return false; + } else { + if (actual[i] != expected[i]) return false; + } + } + return true; +} + +template +class TestFixture { +public: + TestFixture(const std::string& = "") {} + ~TestFixture() { graph_.hard_reset(); } + + CactusGraph& graph() { return graph_; } + + size_t create_input(const std::vector& shape, Precision precision = default_precision()) { + return graph_.input(shape, precision); + } + + void set_input_data(size_t input_id, const std::vector& data, Precision precision = default_precision()) { + graph_.set_input(input_id, const_cast(static_cast(data.data())), precision); + } + + void execute() { graph_.execute(); } + + T* get_output(size_t node_id) { + return static_cast(graph_.get_output(node_id)); + } + + bool verify_output(size_t node_id, const std::vector& expected, float tolerance = default_tolerance()) { + return compare_arrays(get_output(node_id), expected.data(), expected.size(), tolerance); + } + +private: + CactusGraph graph_; +}; + +using Int8TestFixture = TestFixture; +using FP16TestFixture = TestFixture<__fp16>; + +inline void fill_random_fp16(std::vector<__fp16>& data) { + static std::mt19937 gen(42); + std::uniform_real_distribution dis(-1.0f, 1.0f); + for (auto& x : data) x = static_cast<__fp16>(dis(gen)); +} + +inline size_t runtime_id_from_serialized_index(uint32_t serialized_index) { + return static_cast(serialized_index); +} + +inline bool test_basic_operation(const std::string& op_name, + std::function op_func, + const std::vector<__fp16>& data_a, + const std::vector<__fp16>& data_b, + const std::vector<__fp16>& expected, + const std::vector& shape = {4}) { + (void)op_name; + CactusGraph graph; + size_t input_a = graph.input(shape, Precision::FP16); + size_t input_b = graph.input(shape, Precision::FP16); + size_t result_id = op_func(graph, input_a, input_b); + graph.set_input(input_a, const_cast(static_cast(data_a.data())), Precision::FP16); + graph.set_input(input_b, const_cast(static_cast(data_b.data())), Precision::FP16); + graph.execute(); + __fp16* output = static_cast<__fp16*>(graph.get_output(result_id)); + for (size_t i = 0; i < expected.size(); ++i) { + if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-2f) { + graph.hard_reset(); + return false; + } + } + graph.hard_reset(); + return true; +} + +inline bool test_scalar_operation(const std::string& op_name, + std::function op_func, + const std::vector<__fp16>& data, + float scalar, + const std::vector<__fp16>& expected, + const std::vector& shape = {4}) { + (void)op_name; + CactusGraph graph; + size_t input_a = graph.input(shape, Precision::FP16); + size_t result_id = op_func(graph, input_a, scalar); + graph.set_input(input_a, const_cast(static_cast(data.data())), Precision::FP16); + graph.execute(); + __fp16* output = static_cast<__fp16*>(graph.get_output(result_id)); + for (size_t i = 0; i < expected.size(); ++i) { + if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-2f) { + graph.hard_reset(); + return false; + } + } + graph.hard_reset(); + return true; +} + +} // namespace TestUtils diff --git a/cactus-kernels/CMakeLists.txt b/cactus-kernels/CMakeLists.txt new file mode 100644 index 000000000..20d76dc4f --- /dev/null +++ b/cactus-kernels/CMakeLists.txt @@ -0,0 +1,92 @@ +cmake_minimum_required(VERSION 3.10) +project(CactusKernels LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +add_library(cactus_kernels STATIC + src/attention.cpp + src/attention_hybrid.cpp + src/blas.cpp + src/conv.cpp + src/conv2d.cpp + src/dsp.cpp + src/fused.cpp + src/image.cpp + src/lstm.cpp + src/matmul.cpp + src/nn.cpp + src/norms_rope.cpp + src/quants.cpp + src/reduce.cpp + src/scalar.cpp +) + +set_target_properties(cactus_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON) + +target_include_directories(cactus_kernels + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/libs +) + +add_library(cactus_flags INTERFACE) + +target_compile_options(cactus_flags INTERFACE + -march=armv8.2-a+fp16+simd+dotprod+i8mm -pthread -O3 + -fvisibility=hidden -fvisibility-inlines-hidden + -ffunction-sections -fdata-sections -fno-rtti + $<$:-flto=thin> +) + +target_link_options(cactus_flags INTERFACE + $<$:-flto=thin> +) + +if(APPLE) + target_link_options(cactus_flags INTERFACE -Wl,-dead_strip) +else() + target_link_options(cactus_flags INTERFACE -Wl,--gc-sections) +endif() + +if(ANDROID) + target_link_options(cactus_flags INTERFACE -Wl,--exclude-libs,ALL) +endif() + +target_link_libraries(cactus_kernels PUBLIC cactus_flags) + +target_compile_options(cactus_kernels PRIVATE + -Wall -Wextra -pedantic -Wno-missing-field-initializers + $<$:-Wno-c99-extensions> + $<$:-Wno-pedantic> +) + +if(APPLE) + set(CMAKE_OSX_ARCHITECTURES "arm64") + target_compile_definitions(cactus_kernels PUBLIC ACCELERATE_NEW_LAPACK) + target_link_libraries(cactus_kernels PUBLIC "-framework Accelerate") + enable_language(OBJCXX) + set(_msl_in ${CMAKE_CURRENT_SOURCE_DIR}/src/cactus_kernels.metal) + set(_msl_out ${CMAKE_CURRENT_BINARY_DIR}/generated/cactus_kernels_msl.h) + add_custom_command( + OUTPUT ${_msl_out} + COMMAND ${CMAKE_COMMAND} -DIN=${_msl_in} -DOUT=${_msl_out} -DDELIM=CACTUS_MSL + -P ${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_msl.cmake + DEPENDS ${_msl_in} ${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_msl.cmake + VERBATIM) + target_sources(cactus_kernels PRIVATE src/metal_backend.mm ${_msl_out}) + target_include_directories(cactus_kernels PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/generated) + set_source_files_properties(src/metal_backend.mm PROPERTIES COMPILE_FLAGS "-fobjc-arc") + target_link_libraries(cactus_kernels PUBLIC "-framework Metal" "-framework Foundation" "-framework MetalPerformanceShaders") +else() + target_sources(cactus_kernels PRIVATE src/metal_backend_stub.cpp) +endif() + +if((BUILD_TESTING OR CACTUS_BUILD_TESTS) AND CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + file(GLOB TEST_SOURCES "tests/test_*.cpp") + foreach(TEST_SRC ${TEST_SOURCES}) + get_filename_component(TEST_NAME ${TEST_SRC} NAME_WE) + add_executable(${TEST_NAME} ${TEST_SRC}) + target_link_libraries(${TEST_NAME} PRIVATE cactus_kernels) + endforeach() +endif() diff --git a/cactus-kernels/build.sh b/cactus-kernels/build.sh new file mode 100755 index 000000000..7438faad9 --- /dev/null +++ b/cactus-kernels/build.sh @@ -0,0 +1,16 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")" + +echo "Building cactus-kernels..." + +rm -rf build +mkdir -p build +cd build + +cmake .. -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null 2>&1 +make -j$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) + +echo "cactus-kernels built successfully!" +echo "Library: $(pwd)/libcactus_kernels.a" diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h new file mode 100644 index 000000000..4260e9e2a --- /dev/null +++ b/cactus-kernels/cactus_kernels.h @@ -0,0 +1,832 @@ +#ifndef CACTUS_KERNELS_H +#define CACTUS_KERNELS_H + +#include +#include +#include + +#include "threading.h" + +enum class Precision { + INT8, + FP16, + FP32, + CQ1, + CQ2, + CQ3, + CQ4 +}; + +enum class ScalarOpType { + ADD, + SUBTRACT, + MULTIPLY, + DIVIDE, + ABS, + EXP, + POW, + SQRT, + COS, + SIN, + LOG +}; + +constexpr size_t KV_QUANT_GROUP_SIZE = 32; + +void cactus_add_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + size_t num_elements); + +void cactus_add_f16_clipped( + const __fp16* a, + const __fp16* b, + __fp16* output, + size_t num_elements); + +void cactus_subtract_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + size_t num_elements); + +void cactus_multiply_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + size_t num_elements); + +void cactus_divide_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + size_t num_elements); + +void cactus_add_scaled_f16( + const __fp16* base, + const __fp16* src, + __fp16* output, + size_t num_elements, + float scale); + +void cactus_scalar_op_f16( + const __fp16* input, + __fp16* output, + size_t num_elements, + float scalar_value, + ScalarOpType op_type); + +void cactus_add_broadcast_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + const size_t* a_strides, + const size_t* b_strides, + const size_t* output_shape, + size_t ndim); + +void cactus_subtract_broadcast_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + const size_t* a_strides, + const size_t* b_strides, + const size_t* output_shape, + size_t ndim); + +void cactus_multiply_broadcast_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + const size_t* a_strides, + const size_t* b_strides, + const size_t* output_shape, + size_t ndim); + +void cactus_divide_broadcast_f16( + const __fp16* a, + const __fp16* b, + __fp16* output, + const size_t* a_strides, + const size_t* b_strides, + const size_t* output_shape, + size_t ndim); + +double cactus_sum_all_f16(const __fp16* data, size_t num_elements); +double cactus_mean_all_f16(const __fp16* data, size_t num_elements); +double cactus_variance_all_f16(const __fp16* data, size_t num_elements); +__fp16 cactus_min_all_f16(const __fp16* data, size_t num_elements); +__fp16 cactus_max_all_f16(const __fp16* data, size_t num_elements); + +void cactus_sum_axis_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t axis_size, + size_t inner_size); + +void cactus_mean_axis_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t axis_size, + size_t inner_size); + +void cactus_variance_axis_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t axis_size, + size_t inner_size); + +void cactus_min_axis_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t axis_size, + size_t inner_size); + +void cactus_max_axis_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t axis_size, + size_t inner_size); + +void cactus_transpose_2d_f16( + const __fp16* source, + __fp16* destination, + size_t num_rows, + size_t num_cols, + size_t start_row, + size_t end_row); + +void cactus_transpose_f16( + const __fp16* source, + __fp16* destination, + const size_t* shape, + const size_t* permutation, + size_t ndim, + size_t start_idx, + size_t end_idx); + +void cactus_concat_f16( + const __fp16* input1, + const __fp16* input2, + __fp16* output, + const size_t* shape1, + const size_t* shape2, + const size_t* output_shape, + size_t ndims, + int axis); + +void cactus_cat_f16( + const __fp16** inputs, + __fp16* output, + const size_t** input_shapes, + const size_t* output_shape, + size_t num_inputs, + size_t rank, + int axis); + +void cactus_matmul_f16( + const __fp16* a, + const __fp16* b_transposed, + __fp16* c, + size_t M, + size_t K, + size_t N); + +enum CactusQuantFlags : uint32_t { + CACTUS_QUANT_FLAG_ORTHOGONAL = 1u << 2, + CACTUS_QUANT_FLAG_INTERLEAVED_4ROW = 1u << 3, +}; + +struct CactusQuantMatrix { + uint32_t bits; + uint32_t K; + uint32_t N; + uint32_t group_size; + uint32_t num_groups; + uint32_t flags; + const __fp16* codebook; + const __fp16* input_scale; + const __fp16* input_scale_recip; + const __fp16* norms; + const uint8_t* packed_indices; + const int8_t* left_signs; + const int8_t* right_signs; + const uint32_t* permutation; + const __fp16* rotation; + const int8_t* expanded; + const float* norm_f32; +}; + +uint32_t cactus_quant_packed_group_bytes(uint32_t bits, uint32_t group_size); + +void cactus_quant_4bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y); + +void cactus_quant_4bit_gemm( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C); + +void cactus_quant_2bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y); + +void cactus_quant_2bit_gemm( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C); + +void cactus_quant_1bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y); + +void cactus_quant_1bit_gemm( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C); + +void cactus_quant_3bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y); + +void cactus_quant_3bit_gemm( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C); + +void cactus_quant_matmul( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C); + +void cactus_quant_orthogonal_matmul( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C); + +void cactus_quant_4bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y); + +void cactus_quant_2bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y); + +void cactus_quant_3bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y); + +void cactus_quant_1bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y); + +void cactus_quant_dequantize_hadamard_embedding_row( + uint32_t bits, + uint32_t hidden_dim, + uint32_t group_size, + uint32_t num_groups, + size_t row, + const uint8_t* packed_base, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const int8_t* left_signs, + const int8_t* right_signs, + const uint32_t* permutation, + __fp16* out_row); + +void cactus_quant_dequantize_orthogonal_embedding_row( + uint32_t bits, + uint32_t K, + size_t row, + const uint8_t* packed_base, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const __fp16* rotation, + uint32_t flags, + __fp16* out_row); + +void cactus_rms_norm_f16( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t batch_size, + size_t dims, + float eps); + +void cactus_layer_norm_f16( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t batch_size, + size_t dims, + float eps); + +void cactus_batchnorm_f16( + const __fp16* input, + const float* weight, + const float* bias, + const float* running_mean, + const float* running_var, + __fp16* output, + size_t outer_size, + size_t channels, + size_t inner_size, + float epsilon); + +void cactus_batchnorm_f32( + const float* input, + const float* weight, + const float* bias, + const float* running_mean, + const float* running_var, + float* output, + size_t outer_size, + size_t channels, + size_t inner_size, + float epsilon); + +void cactus_softmax_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t vocab_size); + +void cactus_rope_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t num_heads, + size_t head_dim, + size_t start_pos, + float theta); + +void cactus_gpt_j_rope_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t num_heads, + size_t head_dim, + size_t rot_dim, + size_t start_pos, + float theta); + + +void cactus_relu_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_leaky_relu_f16(const __fp16* input, __fp16* output, size_t num_elements, float negative_slope); +void cactus_clamp_f16(const __fp16* input, __fp16* output, size_t num_elements, float lo, float hi); +void cactus_silu_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_gelu_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_gelu_f16_erf(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_sigmoid_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_tanh_f16(const __fp16* input, __fp16* output, size_t num_elements); + +void cactus_glu_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t split_size, + size_t inner_size); + +void cactus_glu_f32( + const float* input, + float* output, + size_t outer_size, + size_t split_size, + size_t inner_size); + + +void cactus_attention_f16( + const __fp16* queries, + const __fp16* keys, + const __fp16* values, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t kv_seq_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + const __fp16* mask, + size_t position_offset = 0, + size_t window_size = 0, + bool is_causal = true, + bool mask_is_additive = false, + bool mask_per_head = false, + size_t v_head_dim = 0, + float logit_cap = 0.0f); + +void cactus_attention_hybrid_int8_fp16( + const __fp16* queries, + const int8_t* keys_cached, + const int8_t* values_cached, + const float* k_scales, + const float* v_scales, + const __fp16* keys_new, + const __fp16* values_new, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t cache_len, + size_t new_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + size_t position_offset = 0, + bool is_causal = true, + size_t window_size = 0, + size_t group_size = KV_QUANT_GROUP_SIZE, + size_t v_head_dim = 0); + + +void cactus_conv1d_causal_depthwise_f16( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, + size_t L, + size_t C, + size_t K, + size_t dilation); + +void cactus_conv1d_f16_k3( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out, + size_t stride); + +void cactus_conv1d_f16( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out, + size_t K, + size_t stride); + +void cactus_conv1d_f16_k7s3_oc8( + const __fp16* input, + const __fp16* Wpack, + const __fp16* bias, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out); + +void cactus_conv1d_same_depthwise_f16_k9( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t L, + size_t C); + +void cactus_conv1d_pointwise_f16_gemm( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out); + +void cactus_conv2d_f16_k3s1p1_nchw( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C_in, + size_t H, + size_t W, + size_t C_out); + +void cactus_conv2d_f16_k3s2p1_nchw( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C_in, + size_t H, + size_t W, + size_t C_out); + +void cactus_conv2d_depthwise_f16_k3s2p1_nchw( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C, + size_t H, + size_t W); + +void cactus_conv2d_pointwise_f16_1x1_nchw_gemm( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C_in, + size_t H, + size_t W, + size_t C_out); + +void cactus_stft_f16( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out, + size_t K, + size_t stride, + size_t num_fft_bins); + +void cactus_bilinear_interpolation_f16( + const __fp16* input, + __fp16* output, + size_t src_height, + size_t src_width, + size_t embed_dim, + size_t dst_height, + size_t dst_width, + bool align_corners = true); + +void cactus_maxpool1d_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t channels, + size_t input_length, + size_t kernel_size, + size_t stride); + +void cactus_lstm_cell_f16( + const __fp16* x_input, + const __fp16* h_prev, + const __fp16* c_prev, + const __fp16* weight_ih, + const __fp16* weight_hh, + const __fp16* bias_ih, + const __fp16* bias_hh, + __fp16* h_new, + __fp16* c_new, + size_t batch_size, + size_t input_size, + size_t hidden_size); + +void cactus_bilstm_sequence_f16( + const __fp16* input, + const __fp16* weight_ih_fwd, + const __fp16* weight_hh_fwd, + const __fp16* bias_ih_fwd, + const __fp16* bias_hh_fwd, + const __fp16* weight_ih_bwd, + const __fp16* weight_hh_bwd, + const __fp16* bias_ih_bwd, + const __fp16* bias_hh_bwd, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t input_size, + size_t hidden_size); + +void cactus_gated_deltanet_decode_f16( + const __fp16* q_data, + const __fp16* k_data, + const __fp16* v_data, + const __fp16* g_data, + const __fp16* b_data, + const __fp16* s_data, + __fp16* out, + size_t B, + size_t Hq, + size_t Hv, + size_t K, + size_t V, + float scale); + +void cactus_gated_deltanet_prefill_f16( + const __fp16* q_data, + const __fp16* k_data, + const __fp16* v_data, + const __fp16* g_data, + const __fp16* b_data, + const __fp16* s_data, + __fp16* out, + size_t B, + size_t T, + size_t Hq, + size_t Hv, + size_t K, + size_t V, + size_t requested_chunk_size, + float scale); + + +void cactus_gaussian_topk_f16( + const __fp16* input, + __fp16* output, + size_t rows, + size_t cols, + float ppf); + +void cactus_altup_predict_f16( + const __fp16* coefs, + const __fp16* const* streams, + __fp16* output, + size_t n, + size_t seq_len, + size_t hidden_dim); + +void cactus_altup_correct_f16( + const __fp16* coefs, + const __fp16* innovation, + const __fp16* const* predictions, + __fp16* output, + size_t n, + size_t seq_len, + size_t hidden_dim); + + +void cactus_sample_f32( + const float* logits, + uint32_t* output, + size_t vocab_size, + float temperature, + float top_p, + size_t top_k, + size_t random_seed, + const float* bias_values = nullptr, + const uint32_t* bias_indices = nullptr, + size_t bias_count = 0); + +void cactus_sample_f16( + const __fp16* logits, + uint32_t* output, + size_t vocab_size, + float temperature, + float top_p, + size_t top_k, + size_t random_seed, + const float* bias_values = nullptr, + const uint32_t* bias_indices = nullptr, + size_t bias_count = 0); + +void cactus_sample_f32_ex( + const float* logits, + uint32_t* output, + size_t vocab_size, + float temperature, + float top_p, + float min_p, + float repetition_penalty, + size_t top_k, + size_t random_seed, + const float* bias_values = nullptr, + const uint32_t* bias_indices = nullptr, + size_t bias_count = 0); + +void cactus_sample_f16_ex( + const __fp16* logits, + uint32_t* output, + size_t vocab_size, + float temperature, + float top_p, + float min_p, + float repetition_penalty, + size_t top_k, + size_t random_seed, + const float* bias_values = nullptr, + const uint32_t* bias_indices = nullptr, + size_t bias_count = 0); + + +void cactus_int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale = 1.0f); +void cactus_fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale = 1.0f); +void cactus_fp16_to_fp32(const __fp16* src, float* dst, size_t count); +void cactus_fp32_to_fp16(const float* src, __fp16* dst, size_t count); +void cactus_int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale = 1.0f); +void cactus_fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale = 1.0f); +float cactus_fp16_max_abs(const __fp16* src, size_t count); + +void cactus_quantize_kv_fp16_to_int8( + const __fp16* src, + int8_t* dst, + float* scales, + size_t seq_len, + size_t kv_heads, + size_t head_dim, + size_t group_size = KV_QUANT_GROUP_SIZE); + +void cactus_rfft_f32_1d(const float* input, float* output, size_t n, const char* norm); +void cactus_irfft_f32_1d(const float* input, float* output, size_t n, const char* norm); +float cactus_hertz_to_mel(float freq, const char* mel_scale); +float cactus_mel_to_hertz(float mels, const char* mel_scale); + +void cactus_generate_mel_filter_bank( + float* mel_filters, int num_frequency_bins, int num_mel_filters, + float min_frequency, float max_frequency, int sampling_rate, + const char* norm, const char* mel_scale, bool triangularize_in_mel_space); + +void cactus_spectrogram_to_db( + float* spectrogram, size_t size, float reference, float min_value, + const float* db_range, float multiplier); + +void cactus_compute_spectrogram_f32( + const float* waveform, size_t waveform_length, + const float* window, size_t window_length, + size_t frame_length, size_t hop_length, const size_t* fft_length, + float* spectrogram, float power, + bool center, const char* pad_mode, bool onesided, + float dither, const float* preemphasis, + const float* mel_filters, size_t mel_filters_size, + float mel_floor, const char* log_mel, + float reference, float min_value, const float* db_range, + bool remove_dc_offset); + +unsigned char* cactus_image_load(const char* path, int* width, int* height, int* channels, int desired_channels); +int cactus_image_info(const char* path, int* width, int* height, int* channels); +void cactus_image_free(unsigned char* data); +const char* cactus_image_failure_reason(); + +void cactus_image_resize_uint8( + const unsigned char* input, int src_w, int src_h, + unsigned char* output, int dst_w, int dst_h, int channels); + +void cactus_image_resize_float( + const float* input, int src_w, int src_h, + float* output, int dst_w, int dst_h, int channels); + +void cactus_image_normalize( + const float* input, float* output, + int width, int height, int channels, + float rescale_factor, const float* mean, const float* std_dev); + +void cactus_image_to_patches( + const float* image, float* patches, + int width, int height, int channels, int patch_size); + +void cactus_image_convert_to_rgb( + const unsigned char* input, unsigned char* output, + int width, int height, int channels); + +inline size_t kv_scales_count( + size_t seq_len, + size_t kv_heads, + size_t head_dim, + size_t group_size = KV_QUANT_GROUP_SIZE) { + size_t num_groups = (head_dim + group_size - 1) / group_size; + return seq_len * kv_heads * num_groups; +} + +#endif diff --git a/cactus-kernels/cmake/embed_msl.cmake b/cactus-kernels/cmake/embed_msl.cmake new file mode 100644 index 000000000..500fde994 --- /dev/null +++ b/cactus-kernels/cmake/embed_msl.cmake @@ -0,0 +1,4 @@ +file(READ "${IN}" SRC_HEX HEX) +string(REGEX REPLACE "(..)" "0x\\1," SRC_BYTES "${SRC_HEX}") +file(WRITE "${OUT}" + "static const char kCactusMSL[] = {${SRC_BYTES}0x00};\n") diff --git a/libs/stb/stb_image.h b/cactus-kernels/libs/stb_image.h similarity index 100% rename from libs/stb/stb_image.h rename to cactus-kernels/libs/stb_image.h diff --git a/libs/stb/stb_image_resize2.h b/cactus-kernels/libs/stb_image_resize2.h similarity index 100% rename from libs/stb/stb_image_resize2.h rename to cactus-kernels/libs/stb_image_resize2.h diff --git a/cactus-kernels/metal_backend.h b/cactus-kernels/metal_backend.h new file mode 100644 index 000000000..9505aad48 --- /dev/null +++ b/cactus-kernels/metal_backend.h @@ -0,0 +1,227 @@ +#ifndef CACTUS_METAL_BACKEND_H +#define CACTUS_METAL_BACKEND_H + +#if defined(__APPLE__) +#define CACTUS_HAS_METAL 1 +#else +#define CACTUS_HAS_METAL 0 +#endif + +#include "cactus_kernels.h" +#include + +#ifdef __cplusplus +extern "C" { +#endif +bool cactus_metal_available(); +#ifdef __cplusplus +} +#endif + +void cactus_metal_set_active(bool active); +bool cactus_metal_active_mode(); + +void cactus_metal_quant_matmul(const CactusQuantMatrix* W, const __fp16* A, + uint32_t M, __fp16* C); + +void cactus_metal_session_begin(); +void cactus_metal_session_sync(); +void cactus_metal_session_flush(); +void cactus_metal_invalidate_host_wraps(); +void cactus_metal_trim_prefill_cache(); +void cactus_metal_session_end(); + +void* cactus_metal_alloc_shared(size_t bytes); +void* cactus_metal_alloc_pooled(size_t bytes); +void cactus_metal_free_shared(void* contents); + +bool cactus_metal_encode_copy(void* out, const void* in, size_t bytes); +bool cactus_metal_encode_binary(int op_type, void* out, const void* a, const void* b, size_t n); +bool cactus_metal_encode_scalar(int op_type, void* out, const void* in, size_t n, float param); +bool cactus_metal_encode_unary(int op_type, void* out, const void* in, size_t n); +bool cactus_metal_encode_swiglu(void* out, const void* gate, const void* up, size_t n, float scale); +bool cactus_metal_encode_rms_norm(void* out, const void* in, const void* weight, + size_t rows, size_t dim, float eps); +bool cactus_metal_encode_rms_norm_add(void* out, const void* in, const void* weight, const void* res, + size_t rows, size_t dim, float eps); +bool cactus_metal_encode_rms_norm_add_rms(void* h_out, void* xn_out, const void* in, const void* w1, + const void* res, const void* w2, + size_t rows, size_t dim, float eps, float out_scale); +bool cactus_metal_encode_rms_norm_add_scale(void* out, const void* in, const void* weight, const void* res, + size_t rows, size_t dim, float eps, float out_scale); +bool cactus_metal_encode_argmax(const void* logits, uint32_t vocab, void* out3, const void* bias); +bool cactus_metal_encode_softcap(void* out, const void* in, size_t n, float cap); +bool cactus_metal_encode_adjust_logits(void* logits, size_t vocab, + const uint32_t* recent, uint32_t n_recent, + int64_t suppressed, float penalty); +bool cactus_metal_encode_cast(void* out, int out_prec, const void* in, int in_prec, size_t n); +bool cactus_metal_encode_quant_matmul(void* out, const void* lhs, const CactusQuantMatrix* W); +bool cactus_metal_encode_quant_matmul_m(void* out, const void* lhs, const CactusQuantMatrix* W, uint32_t M); +bool cactus_metal_encode_transform_batch(const void* x, const CactusQuantMatrix* const* Ws, int B, void* const* codes); +bool cactus_metal_encode_gemv_precoded(void* out, const void* code, const CactusQuantMatrix* W); +bool cactus_metal_encode_transform_gemv(void* out, const void* x, const CactusQuantMatrix* W, const void* osw); +bool cactus_metal_transform_gemv_fits(uint32_t K); +bool cactus_metal_encode_gemv_cat(void* const* outs, const void* const* codes, + const CactusQuantMatrix* const* Ws, int B); +bool cactus_metal_encode_swiglu_transform(void* code, const void* gate, const void* up, + const CactusQuantMatrix* W, float scale); +bool cactus_metal_prewarm_quant(const CactusQuantMatrix* W); +bool cactus_metal_encode_rope_pair(void* out, const void* x, const void* c, const void* s, + uint32_t H, uint32_t D); +bool cactus_metal_encode_rope_pair_rms(void* out, const void* x, const void* w, + const void* c, const void* s, + uint32_t H, uint32_t D, float eps); +bool cactus_metal_encode_deltanet_decode(void* out, const void* q, const void* k, const void* v, + const void* g, const void* b, const void* s, + uint32_t B, uint32_t Hq, uint32_t Hv, + uint32_t K, uint32_t V, float scale); +bool cactus_metal_encode_deltanet_prefill(void* out, const void* q, const void* k, const void* v, + const void* g, const void* b, const void* s, + uint32_t B, uint32_t T, uint32_t Hq, uint32_t Hv, + uint32_t K, uint32_t V, float scale); +bool cactus_metal_encode_rms2_add_clip(void* out, const void* a, const void* wa, + const void* b, const void* wb, size_t dim, + float eps_a, float eps_b); +bool cactus_metal_encode_rms_norm_scale(void* out, const void* in, const void* weight, + size_t rows, size_t dim, float eps, float oscale); +bool cactus_metal_encode_softmax_topk(void* probs, void* topk, const void* in, + size_t rows, size_t cols, size_t k, float scale); +bool cactus_metal_encode_topk_rows(void* out, const void* in, size_t rows, size_t cols, size_t k); +bool cactus_metal_moe_cq4_ready(const CactusQuantMatrix* w1_0); +bool cactus_metal_moe_cq4_build(const CactusQuantMatrix* w1s, const CactusQuantMatrix* w3s, + const CactusQuantMatrix* w2s, uint32_t num_experts); +bool cactus_metal_encode_moe_gated_cq4(void* out, const void* hidden, const void* probs, + const void* topk, const CactusQuantMatrix* w1_0, + uint32_t num_experts, uint32_t top_k, uint32_t tokens, + uint32_t act, uint32_t normalize, float eps, float scaling); +bool cactus_metal_encode_quant_matmul_ortho(void* out, const void* act, void* code, + const CactusQuantMatrix* W); +bool cactus_metal_encode_embedding_ortho(void* out, uint32_t row, const CactusQuantMatrix* W, float scale); +bool cactus_metal_encode_embedding_hadamard(void* out, uint32_t row, const CactusQuantMatrix* W); +bool cactus_metal_encode_embedding_ortho_m(void* out, const CactusQuantMatrix* W, const uint32_t* rows, uint32_t M); +bool cactus_metal_encode_embedding_hadamard_m(void* out, const CactusQuantMatrix* W, const uint32_t* rows, uint32_t M); +bool cactus_metal_encode_gather_f16(void* out, const void* table, size_t table_bytes, const uint32_t* rows, uint32_t M, uint32_t D); + +bool cactus_metal_encode_attention_i8( + void* out, const void* q, const void* knew, const void* vnew, + const void* kc, const void* vc, const void* ks, const void* vs, + uint32_t num_q_heads, uint32_t num_kv_heads, uint32_t head_dim, uint32_t v_hdim, + uint32_t history_len, uint32_t total_keys, uint32_t kv_start, uint32_t kv_end, + float scale, size_t kc_bytes, size_t vc_bytes, size_t ks_bytes, size_t vs_bytes); + +bool cactus_metal_encode_attention_fused_i8( + void* out, const void* q, const void* kraw, const void* vraw, + void* kc, void* vc, void* ks, void* vs, + const void* qw, const void* kw, const void* vw, const void* cs, const void* sn, + uint32_t nqh, uint32_t hd, uint32_t vhd, + uint32_t kv_start, uint32_t kv_end, uint32_t slot, uint32_t has_new, + float eps, float scale, + size_t kc_bytes, size_t vc_bytes, size_t ks_bytes, size_t vs_bytes); + +bool cactus_metal_encode_attention_i8_prefill( + void* out, const void* q, const void* knew, const void* vnew, + const void* kc, const void* vc, const void* ks, const void* vs, + uint32_t num_q_heads, uint32_t num_kv_heads, uint32_t head_dim, uint32_t v_hdim, + uint32_t history_len, uint32_t new_len, uint32_t q_pos0, uint32_t window, uint32_t is_causal, uint32_t M, + float scale, size_t kc_bytes, size_t vc_bytes, size_t ks_bytes, size_t vs_bytes, + uint32_t sink, uint32_t ring); + +bool cactus_metal_encode_binary_f32(int op, void* out, const void* a, const void* b, size_t n); +bool cactus_metal_encode_scalar_f32(int op, void* out, const void* in, size_t n, float p); +bool cactus_metal_encode_unary_f32(int op, void* out, const void* in, size_t n); +bool cactus_metal_encode_clamp(void* out, const void* in, size_t n, float lo, float hi, int f32); +bool cactus_metal_encode_glu(void* out, const void* in, size_t split, size_t inner, size_t n_out); +bool cactus_metal_encode_layer_norm(void* out, const void* in, const void* w, const void* b, + size_t rows, size_t dim, float eps); +bool cactus_metal_encode_softmax_rows(void* out, const void* in, size_t rows, size_t cols); +bool cactus_metal_encode_conv1d_k3(void* out, const void* x, const void* w, int w_int8, + const void* w_scales, uint32_t w_gs, + uint32_t Cin, uint32_t L, uint32_t Cout, uint32_t Lout, uint32_t stride); +bool cactus_metal_encode_gemm_f16(void* out, const void* lhs, const void* rhs, + uint32_t M, uint32_t K, uint32_t N, int pretransposed); +bool cactus_metal_encode_attention_f16(void* out, const void* q, const void* k, const void* v, const void* mask, + uint32_t B, uint32_t T, uint32_t S, uint32_t HQ, uint32_t HKV, uint32_t D, uint32_t DV, + float scale, uint32_t causal, uint32_t pos_off, uint32_t window, float logit_cap, uint32_t mask_mode); +bool cactus_metal_encode_reduce_axis(int op, void* out, const void* in, uint32_t outer, + uint32_t axis_size, uint32_t inner, int f32); +bool cactus_metal_encode_cumsum(void* out, const void* in, uint32_t outer, + uint32_t axis_size, uint32_t inner, int f32); +bool cactus_metal_encode_concat2(void* out, const void* a, const void* b, + uint32_t a_outer, uint32_t b_outer, + uint32_t a_axis, uint32_t b_axis, uint32_t inner); +bool cactus_metal_encode_gather_f32idx(void* out, const void* table, const void* idx, + uint32_t rows, uint32_t D, size_t table_bytes); +bool cactus_metal_encode_rope_full(void* out, const void* in, uint32_t tokens, uint32_t S, + uint32_t H, uint32_t D, uint32_t rot, uint32_t pos0, + float theta, int gptj); +bool cactus_metal_encode_maxpool1d(void* out, const void* in, uint32_t NC, uint32_t L, + uint32_t Lout, uint32_t K, uint32_t stride); +bool cactus_metal_encode_bilinear(void* out, const void* in, uint32_t sh, uint32_t sw, + uint32_t dh, uint32_t dw, uint32_t E, int align); +bool cactus_metal_encode_conv1d_gen(void* out, const void* x, const void* w, const void* bias, + uint32_t N, uint32_t Cin, uint32_t L, uint32_t Cout, + uint32_t Lout, uint32_t K, uint32_t stride, int w_ck_co); +bool cactus_metal_encode_conv1d_nlc_dw(void* out, const void* x, const void* w, const void* bias, + uint32_t N, uint32_t L, uint32_t C, uint32_t K, + uint32_t dil, uint32_t pad); +bool cactus_metal_encode_conv2d(void* out, const void* x, const void* w, const void* bias, + uint32_t N, uint32_t Cin, uint32_t H, uint32_t W, uint32_t Cout, + uint32_t Ho, uint32_t Wo, uint32_t K, uint32_t stride, + uint32_t pad, int dw); +bool cactus_metal_encode_batchnorm(void* out, const void* x, const void* w, const void* b, + const void* rm, const void* rv, uint32_t C, uint32_t inner, + uint32_t total, float eps); +bool cactus_metal_encode_groupnorm(void* out, const void* x, const void* w, const void* b, + uint32_t N, uint32_t C, uint32_t S, uint32_t groups, float eps); +bool cactus_metal_encode_bias_add_rows(void* y, const void* bias, uint32_t C, uint32_t total); +bool cactus_metal_encode_elemwise_chain(void* out, const void* in, const float* steps, + uint32_t nsteps, const void* side0, const void* side1, + const void* side2, const size_t* side_elems, + size_t n, uint32_t flags, uint32_t inner); +bool cactus_metal_encode_rms_norm_add_rows(void* ysum, void* ynorm, const void* x, const void* res, + const void* w, uint32_t rows, uint32_t dim, float eps, + int clipped); +bool cactus_metal_encode_gemm_batch(void* out, const void* a, const void* b, + uint32_t M, uint32_t K, uint32_t N, uint32_t batch, int f32out, int f32a); +bool cactus_metal_encode_conv1d_dw(void* out, const void* x, const void* w, + uint32_t C, uint32_t L, uint32_t Lout, uint32_t K, uint32_t stride); +bool cactus_metal_encode_transpose2d(void* out, const void* in, uint32_t batch, uint32_t R, uint32_t C); +bool cactus_metal_encode_strided_copy(void* out, const void* in, const uint32_t* oshape, + const uint32_t* sstride, uint32_t ndim, uint32_t total, uint32_t base, size_t in_bytes, size_t out_bytes); + +bool cactus_metal_encode_bcast_binary(int op, void* out, const void* a, const void* b, + const uint32_t* oshape, const uint32_t* astride, const uint32_t* bstride, uint32_t ndim, uint32_t total, + size_t a_bytes, size_t b_bytes, size_t out_bytes); + +bool cactus_metal_encode_strided_scatter(void* out, const void* in, const uint32_t* ishape, + const uint32_t* ostride, uint32_t ndim, uint32_t total, uint32_t base, size_t in_bytes, size_t out_bytes); + +bool cactus_metal_encode_kv_append_i8(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, + size_t src_bytes, size_t int8_bytes, size_t scale_bytes); +bool cactus_metal_encode_kv_append_sliding_i8(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t keep_sink, uint32_t remaining, uint32_t shift_src, + uint32_t group_size, size_t src_bytes, size_t int8_bytes, size_t scale_bytes); +bool cactus_metal_encode_kv_append_sliding_i8_m(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t keep_sink, uint32_t remaining, uint32_t shift_src, + uint32_t group_size, uint32_t M, size_t src_bytes, size_t int8_bytes, size_t scale_bytes); + +bool cactus_metal_encode_kv_append_i8_m(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, uint32_t M, + size_t src_bytes, size_t int8_bytes, size_t scale_bytes); + +bool cactus_metal_encode_kv_append_ring_i8_m(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, uint32_t M, + uint32_t sink, uint32_t W, size_t src_bytes, size_t int8_bytes, size_t scale_bytes); + +bool cactus_metal_encode_conv_cache_append(void* out, const void* src, void* ring, + uint32_t hd, uint32_t ws, uint32_t nnew, uint32_t head0, uint32_t count_new, + uint32_t num_rows, int src_f32); + +bool cactus_metal_encode_rel_pos_bias(void* y, const void* q, const void* r, + uint32_t B, uint32_t T, uint32_t H, uint32_t D, uint32_t R, int r_batched, float scale); + +bool cactus_metal_encode_gemv_bias(void* out, const void* x, const void* w, const void* bias, + uint32_t K, uint32_t N, int pretransposed); + +#endif diff --git a/cactus-kernels/src/attention.cpp b/cactus-kernels/src/attention.cpp new file mode 100644 index 000000000..e431d16c5 --- /dev/null +++ b/cactus-kernels/src/attention.cpp @@ -0,0 +1,678 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +#endif + +#ifdef __APPLE__ +static void cactus_attention_f16_accelerate( + const __fp16* queries, + const __fp16* keys, + const __fp16* values, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t kv_seq_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + size_t v_head_dim, + float scale, + size_t position_offset, + bool is_causal +) { + constexpr size_t BLOCK_SIZE = 64; + + const size_t group_size = num_q_heads / num_kv_heads; + const size_t q_batch_stride = seq_len * num_q_heads * head_dim; + const size_t k_batch_stride = kv_seq_len * num_kv_heads * head_dim; + const size_t v_batch_stride = kv_seq_len * num_kv_heads * v_head_dim; + const size_t o_batch_stride = seq_len * num_q_heads * v_head_dim; + const size_t q_seq_stride = num_q_heads * head_dim; + const size_t k_seq_stride = num_kv_heads * head_dim; + const size_t v_seq_stride = num_kv_heads * v_head_dim; + const size_t o_seq_stride = num_q_heads * v_head_dim; + + static constexpr CactusThreading::ParallelConfig ATTENTION_BATCHED{1, 1}; + CactusThreading::parallel_for(batch_size * num_q_heads, ATTENTION_BATCHED, + [&](size_t start, size_t end) { + + std::vector Q_f32(seq_len * head_dim); + std::vector K_f32(BLOCK_SIZE * head_dim); + std::vector V_f32(BLOCK_SIZE * v_head_dim); + std::vector scores(seq_len * BLOCK_SIZE); + std::vector acc(seq_len * v_head_dim); + std::vector row_max(seq_len); + std::vector row_sum(seq_len); + + for (size_t work = start; work < end; ++work) { + const size_t batch = work / num_q_heads; + const size_t q_head = work % num_q_heads; + const size_t kv_head = q_head / group_size; + + for (size_t q = 0; q < seq_len; ++q) { + const __fp16* q_src = queries + batch*q_batch_stride + q*q_seq_stride + q_head*head_dim; + float* q_dst = Q_f32.data() + q * head_dim; + for (size_t d = 0; d < head_dim; d += 8) { + float16x8_t v = vld1q_f16(q_src + d); + vst1q_f32(q_dst + d, vcvt_f32_f16(vget_low_f16(v))); + vst1q_f32(q_dst + d + 4, vcvt_f32_f16(vget_high_f16(v))); + } + } + + std::fill(row_max.begin(), row_max.begin() + seq_len, -INFINITY); + std::fill(row_sum.begin(), row_sum.begin() + seq_len, 0.0f); + memset(acc.data(), 0, seq_len * v_head_dim * sizeof(float)); + + for (size_t kv0 = 0; kv0 < kv_seq_len; kv0 += BLOCK_SIZE) { + const size_t block_len = std::min(BLOCK_SIZE, kv_seq_len - kv0); + + size_t q_start = 0; + size_t active_rows = seq_len; + if (is_causal) { + if (kv0 > position_offset) { + q_start = kv0 - position_offset; + } + if (q_start >= seq_len) continue; + active_rows = seq_len - q_start; + } + + for (size_t i = 0; i < block_len; ++i) { + const __fp16* k_src = keys + batch*k_batch_stride + (kv0+i)*k_seq_stride + kv_head*head_dim; + float* k_dst = K_f32.data() + i * head_dim; + for (size_t d = 0; d < head_dim; d += 8) { + float16x8_t v = vld1q_f16(k_src + d); + vst1q_f32(k_dst + d, vcvt_f32_f16(vget_low_f16(v))); + vst1q_f32(k_dst + d + 4, vcvt_f32_f16(vget_high_f16(v))); + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + (int)active_rows, (int)block_len, (int)head_dim, + scale, + Q_f32.data() + q_start * head_dim, (int)head_dim, + K_f32.data(), (int)head_dim, + 0.0f, + scores.data(), (int)block_len); + + for (size_t r = 0; r < active_rows; ++r) { + const size_t q_pos = q_start + r; + const size_t abs_q = position_offset + q_pos; + float* s_row = scores.data() + r * block_len; + + size_t valid_len = block_len; + if (is_causal) { + if (abs_q < kv0) { + memset(s_row, 0, block_len * sizeof(float)); + continue; + } + valid_len = std::min(block_len, abs_q - kv0 + 1); + } + + float32x4_t vmax = vdupq_n_f32(-INFINITY); + size_t j = 0; + for (; j + 4 <= valid_len; j += 4) { + vmax = vmaxq_f32(vmax, vld1q_f32(s_row + j)); + } + float block_max = vmaxvq_f32(vmax); + for (; j < valid_len; ++j) { + block_max = std::max(block_max, s_row[j]); + } + + float prev_max = row_max[q_pos]; + float new_max = std::max(prev_max, block_max); + float scale_old = expf(prev_max - new_max); + + if (prev_max != -INFINITY) { + float* acc_row = acc.data() + q_pos * v_head_dim; + float32x4_t sv = vdupq_n_f32(scale_old); + for (size_t d = 0; d < v_head_dim; d += 4) { + float32x4_t a = vld1q_f32(acc_row + d); + vst1q_f32(acc_row + d, vmulq_f32(a, sv)); + } + } + row_sum[q_pos] = row_sum[q_pos] * scale_old; + row_max[q_pos] = new_max; + + float32x4_t vsum = vdupq_n_f32(0.0f); + float32x4_t vnmax = vdupq_n_f32(new_max); + float32x4_t log2e = vdupq_n_f32(1.442695f); + j = 0; + for (; j + 4 <= valid_len; j += 4) { + float32x4_t x = vmulq_f32(vsubq_f32(vld1q_f32(s_row + j), vnmax), log2e); + float32x4_t x_floor = vrndmq_f32(x); + int32x4_t xi = vcvtq_s32_f32(x_floor); + float32x4_t xf = vsubq_f32(x, x_floor); + float32x4_t t = vfmaq_n_f32(vdupq_n_f32(0.2246932f), xf, 0.0789673f); + t = vfmaq_f32(vdupq_n_f32(0.6963248f), t, xf); + float32x4_t y = vfmaq_f32(vdupq_n_f32(0.9999003f), t, xf); + xi = vshlq_n_s32(vaddq_s32(xi, vdupq_n_s32(127)), 23); + y = vmulq_f32(y, vreinterpretq_f32_s32(xi)); + uint32x4_t underflow = vcltq_f32(x, vdupq_n_f32(-126.0f)); + y = vbslq_f32(underflow, vdupq_n_f32(0.0f), y); + vst1q_f32(s_row + j, y); + vsum = vaddq_f32(vsum, y); + } + float block_sum = vaddvq_f32(vsum); + for (; j < valid_len; ++j) { + s_row[j] = expf(s_row[j] - new_max); + block_sum += s_row[j]; + } + if (valid_len < block_len) { + memset(s_row + valid_len, 0, (block_len - valid_len) * sizeof(float)); + } + row_sum[q_pos] += block_sum; + } + + for (size_t i = 0; i < block_len; ++i) { + const __fp16* v_src = values + batch*v_batch_stride + (kv0+i)*v_seq_stride + kv_head*v_head_dim; + float* v_dst = V_f32.data() + i * v_head_dim; + for (size_t d = 0; d < v_head_dim; d += 8) { + float16x8_t v = vld1q_f16(v_src + d); + vst1q_f32(v_dst + d, vcvt_f32_f16(vget_low_f16(v))); + vst1q_f32(v_dst + d + 4, vcvt_f32_f16(vget_high_f16(v))); + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + (int)active_rows, (int)v_head_dim, (int)block_len, + 1.0f, + scores.data(), (int)block_len, + V_f32.data(), (int)v_head_dim, + 1.0f, + acc.data() + q_start * v_head_dim, (int)v_head_dim); + } + + for (size_t q = 0; q < seq_len; ++q) { + __fp16* o = output + batch*o_batch_stride + q*o_seq_stride + q_head*v_head_dim; + float sum = row_sum[q]; + if (sum == 0.0f) { + memset(o, 0, v_head_dim * sizeof(__fp16)); + continue; + } + float inv = 1.0f / sum; + float32x4_t invv = vdupq_n_f32(inv); + float* acc_row = acc.data() + q * v_head_dim; + for (size_t d = 0; d < v_head_dim; d += 8) { + float32x4_t a0 = vmulq_f32(vld1q_f32(acc_row + d), invv); + float32x4_t a1 = vmulq_f32(vld1q_f32(acc_row + d + 4), invv); + vst1q_f16(o + d, vcombine_f16(vcvt_f16_f32(a0), vcvt_f16_f32(a1))); + } + } + } + }); +} +#endif + +static inline void cactus_attention_f16_fast( + const __fp16* queries, + const __fp16* keys, + const __fp16* values, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t kv_seq_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + size_t position_offset, + bool is_causal, + size_t window_size, + size_t v_head_dim +) { + constexpr size_t BLOCK_SIZE = 32; + const size_t qk_nblocks = head_dim / 8; + const size_t v_nblocks = v_head_dim / 8; + +#ifdef __APPLE__ + if (seq_len >= 64 && window_size == 0) { + cactus_attention_f16_accelerate( + queries, keys, values, output, + batch_size, seq_len, kv_seq_len, + num_q_heads, num_kv_heads, head_dim, v_head_dim, + scale, position_offset, is_causal + ); + return; + } +#endif + + const size_t group_size = num_q_heads / num_kv_heads; + const size_t q_batch_stride = seq_len * num_q_heads * head_dim; + const size_t kv_batch_stride = kv_seq_len * num_kv_heads * head_dim; + const size_t v_batch_stride = kv_seq_len * num_kv_heads * v_head_dim; + const size_t o_batch_stride = seq_len * num_q_heads * v_head_dim; + const size_t q_seq_stride = num_q_heads * head_dim; + const size_t kv_seq_stride = num_kv_heads * head_dim; + const size_t v_seq_stride = num_kv_heads * v_head_dim; + const size_t o_seq_stride = num_q_heads * v_head_dim; + + CactusThreading::parallel_for(batch_size * num_q_heads * seq_len, CactusThreading::Thresholds::ATTENTION, + [&](size_t start, size_t end) { + + float block_scores[BLOCK_SIZE]; + std::vector acc_lo(v_nblocks), acc_hi(v_nblocks); + + for (size_t work = start; work < end; ++work) { + const size_t batch = work / (num_q_heads * seq_len); + const size_t rem = work % (num_q_heads * seq_len); + const size_t q_head = rem / seq_len; + const size_t q_pos = rem % seq_len; + const size_t kv_head = q_head / group_size; + + const __fp16* q = queries + batch*q_batch_stride + q_pos*q_seq_stride + q_head*head_dim; + __fp16* o = output + batch*o_batch_stride + q_pos*o_seq_stride + q_head*v_head_dim; + + for (size_t i = 0; i < v_nblocks; i++) { + acc_lo[i] = vdupq_n_f32(0.f); + acc_hi[i] = vdupq_n_f32(0.f); + } + + float running_max = -INFINITY; + float running_sum = 0.f; + + const size_t abs_q = position_offset + q_pos; + size_t kv_end = is_causal ? std::min(kv_seq_len, abs_q + 1) : kv_seq_len; + size_t kv_start = (window_size > 0 && abs_q > window_size) ? abs_q - window_size : 0; + + for (size_t kv0 = kv_start; kv0 < kv_end; kv0 += BLOCK_SIZE) { + const size_t kv1 = std::min(kv0 + BLOCK_SIZE, kv_end); + float block_max = -INFINITY; + + for (size_t i = kv0; i < kv1; i++) { + float32x4_t s0 = vdupq_n_f32(0.f); + float32x4_t s1 = vdupq_n_f32(0.f); + + const __fp16* k = keys + batch*kv_batch_stride + i*kv_seq_stride + kv_head*head_dim; + + for (size_t d = 0; d < qk_nblocks; d++) { + float16x8_t qv = vld1q_f16(q + d*8); + float16x8_t kv = vld1q_f16(k + d*8); + + float32x4_t ql = vcvt_f32_f16(vget_low_f16(qv)); + float32x4_t qh = vcvt_f32_f16(vget_high_f16(qv)); + float32x4_t kl = vcvt_f32_f16(vget_low_f16(kv)); + float32x4_t kh = vcvt_f32_f16(vget_high_f16(kv)); + + s0 = vfmaq_f32(s0, ql, kl); + s1 = vfmaq_f32(s1, qh, kh); + } + + float score = vaddvq_f32(vaddq_f32(s0, s1)) * scale; + block_scores[i - kv0] = score; + block_max = std::max(block_max, score); + } + + float current_block_scale = 1.0f; + if (block_max > running_max) { + float scale_correction = expf(running_max - block_max); + running_sum *= scale_correction; + + for (size_t d = 0; d < v_nblocks; d++) { + acc_lo[d] = vmulq_n_f32(acc_lo[d], scale_correction); + acc_hi[d] = vmulq_n_f32(acc_hi[d], scale_correction); + } + running_max = block_max; + } else { + current_block_scale = expf(block_max - running_max); + } + + float block_sum = 0.f; + for (size_t i = 0; i < kv1 - kv0; i++) { + block_scores[i] = expf(block_scores[i] - block_max); + block_sum += block_scores[i]; + } + + for (size_t i = 0; i < kv1 - kv0; i++) { + const float attn_weight = block_scores[i] * current_block_scale; + if (attn_weight == 0.f) continue; + + const __fp16* v = values + batch*v_batch_stride + (kv0+i)*v_seq_stride + kv_head*v_head_dim; + float32x4_t wv = vdupq_n_f32(attn_weight); + + for (size_t d = 0; d < v_nblocks; d++) { + float16x8_t vv = vld1q_f16(v + d*8); + acc_lo[d] = vfmaq_f32(acc_lo[d], vcvt_f32_f16(vget_low_f16(vv)), wv); + acc_hi[d] = vfmaq_f32(acc_hi[d], vcvt_f32_f16(vget_high_f16(vv)), wv); + } + } + + running_sum += block_sum * current_block_scale; + } + + if (running_sum == 0.f) { + memset(o, 0, v_head_dim * sizeof(__fp16)); + continue; + } + + float inv = 1.f / running_sum; + float32x4_t invv = vdupq_n_f32(inv); + + for (size_t d = 0; d < v_nblocks; d++) { + float16x8_t out = vcombine_f16( + vcvt_f16_f32(vmulq_f32(acc_lo[d], invv)), + vcvt_f16_f32(vmulq_f32(acc_hi[d], invv)) + ); + vst1q_f16(o + d*8, out); + } + } + }); +} + +void cactus_attention_f16( + const __fp16* queries, + const __fp16* keys, + const __fp16* values, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t kv_seq_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + const __fp16* mask, + size_t position_offset, + size_t window_size, + bool is_causal, + bool mask_is_additive, + bool mask_per_head, + size_t v_head_dim, + float logit_cap +) { + if (v_head_dim == 0) v_head_dim = head_dim; + if (scale == 0.0f) { + scale = 1.0f / sqrtf(static_cast(head_dim)); + } + + if (mask == nullptr && head_dim % 8 == 0 && v_head_dim % 8 == 0 && logit_cap == 0.0f) { + cactus_attention_f16_fast( + queries, keys, values, output, + batch_size, seq_len, kv_seq_len, + num_q_heads, num_kv_heads, head_dim, + scale, position_offset, is_causal, window_size, v_head_dim + ); + return; + } + + constexpr size_t VECTOR_WIDTH = 8; + constexpr size_t BLOCK_SIZE = 32; + const size_t head_dim_aligned = (head_dim / VECTOR_WIDTH) * VECTOR_WIDTH; + const size_t v_head_dim_aligned = (v_head_dim / VECTOR_WIDTH) * VECTOR_WIDTH; + + const size_t group_size = num_q_heads / num_kv_heads; + + const size_t q_batch_stride = seq_len * num_q_heads * head_dim; + const size_t kv_batch_stride = kv_seq_len * num_kv_heads * head_dim; + const size_t v_batch_stride = kv_seq_len * num_kv_heads * v_head_dim; + const size_t o_batch_stride = seq_len * num_q_heads * v_head_dim; + const size_t q_seq_stride = num_q_heads * head_dim; + const size_t kv_seq_stride = num_kv_heads * head_dim; + const size_t v_seq_stride = num_kv_heads * v_head_dim; + const size_t o_seq_stride = num_q_heads * v_head_dim; + const size_t mask_batch_stride = mask + ? (mask_per_head ? (num_q_heads * seq_len * kv_seq_len) : (seq_len * kv_seq_len)) + : 0; + + CactusThreading::parallel_for(batch_size * num_q_heads * seq_len, CactusThreading::Thresholds::ATTENTION, + [=](size_t start_idx, size_t end_idx) { + std::vector block_scores(BLOCK_SIZE); + std::vector output_accum_low(v_head_dim_aligned / VECTOR_WIDTH * 2); + std::vector output_accum_high(v_head_dim_aligned / VECTOR_WIDTH * 2); + + const size_t v_tail_dims = v_head_dim - v_head_dim_aligned; + std::vector output_accum_tail(v_tail_dims, 0.0f); + + const float NEG_INF = -std::numeric_limits::infinity(); + const size_t used_vec_blocks = v_head_dim_aligned / VECTOR_WIDTH; + + for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { + const size_t batch_idx = work_idx / (num_q_heads * seq_len); + const size_t remainder = work_idx % (num_q_heads * seq_len); + const size_t q_head_idx = remainder / seq_len; + const size_t q_pos = remainder % seq_len; + + const size_t kv_head_idx = q_head_idx / group_size; + + const __fp16* Q_base = queries + batch_idx * q_batch_stride; + const __fp16* K_base = keys + batch_idx * kv_batch_stride; + const __fp16* V_base = values + batch_idx * v_batch_stride; + __fp16* O_base = output + batch_idx * o_batch_stride; + const __fp16* M = mask ? (mask + batch_idx * mask_batch_stride) : nullptr; + const __fp16* q_vec = Q_base + q_pos * q_seq_stride + q_head_idx * head_dim; + __fp16* o_vec = O_base + q_pos * o_seq_stride + q_head_idx * v_head_dim; + + float running_max = -std::numeric_limits::infinity(); + float running_sum = 0.0f; + + for (size_t i = 0; i < output_accum_low.size(); ++i) { + output_accum_low[i] = vdupq_n_f32(0.0f); + output_accum_high[i] = vdupq_n_f32(0.0f); + } + for (size_t i = 0; i < v_tail_dims; ++i) { + output_accum_tail[i] = 0.0f; + } + + const bool is_decode = (q_pos == seq_len - 1) && seq_len > 1; + const size_t absolute_q_pos = position_offset + q_pos; + + size_t kv_start = 0; + size_t kv_end = kv_seq_len; + + if (window_size > 0 && window_size < kv_seq_len) { + if (absolute_q_pos > window_size) { + kv_start = absolute_q_pos - window_size; + } + if (is_causal) { + kv_end = std::min(kv_end, absolute_q_pos + 1); + } + } else if (is_causal) { + kv_end = std::min(kv_end, absolute_q_pos + 1); + } + + for (size_t kv_block_start = kv_start; kv_block_start < kv_end; kv_block_start += BLOCK_SIZE) { + const size_t kv_block_end = std::min(kv_block_start + BLOCK_SIZE, kv_end); + const size_t block_size = kv_block_end - kv_block_start; + + float block_max = -std::numeric_limits::infinity(); + + if (!is_decode && is_causal && kv_block_start > absolute_q_pos) { + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + block_scores[kv_idx] = NEG_INF; + } + continue; + } + + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + const size_t kv_pos = kv_block_start + kv_idx; + + if (!is_decode && is_causal && kv_pos > absolute_q_pos) { + block_scores[kv_idx] = NEG_INF; + continue; + } + + const __fp16* k_vec = K_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; + + if (kv_idx + 1 < block_size) { + const __fp16* next_k_vec = K_base + (kv_pos + 1) * kv_seq_stride + kv_head_idx * head_dim; + __builtin_prefetch(next_k_vec, 0, 1); + } + + float32x4_t score_accum_low = vdupq_n_f32(0.0f); + float32x4_t score_accum_high = vdupq_n_f32(0.0f); + + for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { + float16x8_t q_vec_f16 = vld1q_f16(&q_vec[dim_block]); + float16x8_t k_vec_f16 = vld1q_f16(&k_vec[dim_block]); + + float32x4_t q_low = vcvt_f32_f16(vget_low_f16(q_vec_f16)); + float32x4_t q_high = vcvt_f32_f16(vget_high_f16(q_vec_f16)); + float32x4_t k_low = vcvt_f32_f16(vget_low_f16(k_vec_f16)); + float32x4_t k_high = vcvt_f32_f16(vget_high_f16(k_vec_f16)); + + score_accum_low = vfmaq_f32(score_accum_low, q_low, k_low); + score_accum_high = vfmaq_f32(score_accum_high, q_high, k_high); + } + + float score = vaddvq_f32(vaddq_f32(score_accum_low, score_accum_high)); + + for (size_t dim = head_dim_aligned; dim < head_dim; ++dim) { + score += static_cast(q_vec[dim]) * static_cast(k_vec[dim]); + } + + score *= scale; + + size_t absolute_q_pos = position_offset + q_pos; + + if (is_causal && kv_pos > absolute_q_pos) { + score = NEG_INF; + } + else if (window_size > 0 && kv_pos < absolute_q_pos && (absolute_q_pos - kv_pos) > window_size) { + score = NEG_INF; + } + else if (M) { + const size_t mask_index = mask_per_head + ? ((q_head_idx * seq_len + q_pos) * kv_seq_len + kv_pos) + : (q_pos * kv_seq_len + kv_pos); + const float mask_value = static_cast(M[mask_index]); + if (mask_is_additive) { + if (!std::isfinite(mask_value)) { + score = NEG_INF; + } else { + score += mask_value; + } + } else if (mask_value == 0.0f) { + score = NEG_INF; + } + } + + if (logit_cap > 0.0f && std::isfinite(score)) { + score = logit_cap * tanhf(score / logit_cap); + } + + block_scores[kv_idx] = score; + block_max = std::max(block_max, score); + } + + float current_block_scale = 1.0f; + + if (block_max > NEG_INF) { + if (block_max > running_max) { + float scale_correction = expf(running_max - block_max); + running_sum *= scale_correction; + + for (size_t i = 0; i < used_vec_blocks; ++i) { + output_accum_low[i] = vmulq_n_f32(output_accum_low[i], scale_correction); + output_accum_high[i] = vmulq_n_f32(output_accum_high[i], scale_correction); + } + for (size_t i = 0; i < v_tail_dims; ++i) { + output_accum_tail[i] *= scale_correction; + } + running_max = block_max; + } else { + current_block_scale = expf(block_max - running_max); + } + } + + float block_sum = 0.0f; + const size_t vec_size = (block_size / 4) * 4; + + for (size_t kv_idx = 0; kv_idx < vec_size; kv_idx += 4) { + float32x4_t scores = vld1q_f32(&block_scores[kv_idx]); + uint32x4_t inf_mask = vceqq_f32(scores, vdupq_n_f32(NEG_INF)); + + float32x4_t x = vsubq_f32(scores, vdupq_n_f32(block_max)); + x = vmulq_n_f32(x, 1.442695f); + float32x4_t x_floor = vrndmq_f32(x); + int32x4_t xi = vcvtq_s32_f32(x_floor); + float32x4_t xf = vsubq_f32(x, x_floor); + + float32x4_t t = vfmaq_n_f32(vdupq_n_f32(0.2246932f), xf, 0.0789673f); + t = vfmaq_f32(vdupq_n_f32(0.6963248f), t, xf); + float32x4_t y = vfmaq_f32(vdupq_n_f32(0.9999003f), t, xf); + + xi = vaddq_s32(xi, vdupq_n_s32(127)); + xi = vshlq_n_s32(xi, 23); + y = vmulq_f32(y, vreinterpretq_f32_s32(xi)); + + uint32x4_t underflow_mask = vcltq_f32(x, vdupq_n_f32(-126.0f)); + uint32x4_t zero_mask = vorrq_u32(inf_mask, underflow_mask); + y = vbslq_f32(zero_mask, vdupq_n_f32(0.0f), y); + + vst1q_f32(&block_scores[kv_idx], y); + block_sum += vaddvq_f32(y); + } + + for (size_t kv_idx = vec_size; kv_idx < block_size; ++kv_idx) { + if (block_scores[kv_idx] != NEG_INF) { + block_scores[kv_idx] = expf(block_scores[kv_idx] - block_max); + block_sum += block_scores[kv_idx]; + } else { + block_scores[kv_idx] = 0.0f; + } + } + + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + const float attn_weight = block_scores[kv_idx] * current_block_scale; + if (attn_weight == 0.0f) continue; + + const size_t kv_pos = kv_block_start + kv_idx; + const __fp16* v_vec = V_base + kv_pos * v_seq_stride + kv_head_idx * v_head_dim; + + const float32x4_t weight_vec = vdupq_n_f32(attn_weight); + + for (size_t dim_block = 0; dim_block < v_head_dim_aligned; dim_block += VECTOR_WIDTH) { + float16x8_t v_vec_f16 = vld1q_f16(&v_vec[dim_block]); + float32x4_t v_low = vcvt_f32_f16(vget_low_f16(v_vec_f16)); + float32x4_t v_high = vcvt_f32_f16(vget_high_f16(v_vec_f16)); + + size_t idx = dim_block / VECTOR_WIDTH; + output_accum_low[idx] = vfmaq_f32(output_accum_low[idx], v_low, weight_vec); + output_accum_high[idx] = vfmaq_f32(output_accum_high[idx], v_high, weight_vec); + } + + for (size_t dim = v_head_dim_aligned; dim < v_head_dim; ++dim) { + float val = attn_weight * static_cast(v_vec[dim]); + output_accum_tail[dim - v_head_dim_aligned] += val; + } + } + + running_sum += block_sum * current_block_scale; + } + + if (running_sum > 0.0f) { + const float inv_sum = 1.0f / running_sum; + const float32x4_t inv_sum_vec = vdupq_n_f32(inv_sum); + + for (size_t dim_block = 0; dim_block < v_head_dim_aligned; dim_block += VECTOR_WIDTH) { + size_t idx = dim_block / VECTOR_WIDTH; + float32x4_t final_low = vmulq_f32(output_accum_low[idx], inv_sum_vec); + float32x4_t final_high = vmulq_f32(output_accum_high[idx], inv_sum_vec); + + float16x4_t low_f16 = vcvt_f16_f32(final_low); + float16x4_t high_f16 = vcvt_f16_f32(final_high); + float16x8_t combined = vcombine_f16(low_f16, high_f16); + + vst1q_f16(&o_vec[dim_block], combined); + } + + for (size_t dim = v_head_dim_aligned; dim < v_head_dim; ++dim) { + o_vec[dim] = static_cast<__fp16>(output_accum_tail[dim - v_head_dim_aligned] * inv_sum); + } + } else { + for (size_t dim = 0; dim < v_head_dim; ++dim) { + o_vec[dim] = static_cast<__fp16>(0.0f); + } + } + } + }); +} + diff --git a/cactus-kernels/src/attention_hybrid.cpp b/cactus-kernels/src/attention_hybrid.cpp new file mode 100644 index 000000000..044ff1408 --- /dev/null +++ b/cactus-kernels/src/attention_hybrid.cpp @@ -0,0 +1,618 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include + +static void cactus_attention_hybrid_int8_fp16_decode_dot( + const __fp16* queries, + const int8_t* keys_cached, + const int8_t* values_cached, + const float* k_scales, + const float* v_scales, + const __fp16* keys_new, + const __fp16* values_new, + __fp16* output, + size_t batch_size, + size_t cache_len, + size_t new_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + size_t position_offset, + bool is_causal, + size_t window_size +) { + const size_t kv_seq_len = cache_len + new_len; + + constexpr size_t VECTOR_WIDTH = 8; + constexpr size_t BLOCK_SIZE = 64; + constexpr size_t QGROUP = 32; + constexpr size_t MAX_HEAD_DIM = 512; + constexpr size_t MAX_QUANT_GROUPS = MAX_HEAD_DIM / QGROUP; + constexpr size_t MAX_ACCUM_SLOTS = MAX_HEAD_DIM / VECTOR_WIDTH; + + const size_t num_quant_groups = head_dim / QGROUP; + const size_t num_accum_slots = head_dim / VECTOR_WIDTH; + const size_t gqa_group_size = num_q_heads / num_kv_heads; + + const size_t q_batch_stride = num_q_heads * head_dim; + const size_t kv_seq_stride = num_kv_heads * head_dim; + const size_t k_cached_batch_stride = cache_len * kv_seq_stride; + const size_t v_cached_batch_stride = cache_len * kv_seq_stride; + const size_t k_new_batch_stride = new_len * kv_seq_stride; + const size_t v_new_batch_stride = new_len * kv_seq_stride; + const size_t o_batch_stride = num_q_heads * head_dim; + + CactusThreading::parallel_for(batch_size * num_q_heads, CactusThreading::Thresholds::ATTENTION, + [=](size_t start_idx, size_t end_idx) { + alignas(16) int8_t q_int8[MAX_HEAD_DIM]; + float q_scales[MAX_QUANT_GROUPS]; + float block_scores[BLOCK_SIZE]; + float32x4_t output_accum_low[MAX_ACCUM_SLOTS]; + float32x4_t output_accum_high[MAX_ACCUM_SLOTS]; + float16x8_t block_accum[MAX_ACCUM_SLOTS]; + + for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { + const size_t batch_idx = work_idx / num_q_heads; + const size_t q_head_idx = work_idx % num_q_heads; + const size_t kv_head_idx = q_head_idx / gqa_group_size; + + const __fp16* q_vec = queries + batch_idx * q_batch_stride + q_head_idx * head_dim; + const int8_t* K_cached_base = keys_cached + batch_idx * k_cached_batch_stride; + const int8_t* V_cached_base = values_cached + batch_idx * v_cached_batch_stride; + const __fp16* K_new_base = keys_new + batch_idx * k_new_batch_stride; + const __fp16* V_new_base = values_new + batch_idx * v_new_batch_stride; + __fp16* o_vec = output + batch_idx * o_batch_stride + q_head_idx * head_dim; + + for (size_t qg = 0; qg < num_quant_groups; ++qg) { + const __fp16* q_grp = q_vec + qg * QGROUP; + float16x8_t amax_v = vabsq_f16(vld1q_f16(q_grp)); + for (size_t i = 1; i < QGROUP / VECTOR_WIDTH; ++i) + amax_v = vmaxq_f16(amax_v, vabsq_f16(vld1q_f16(q_grp + i * VECTOR_WIDTH))); + float amax = static_cast(vmaxvq_f16(amax_v)); + float q_scale = amax / 127.0f; + float inv = q_scale > 0.0f ? 127.0f / amax : 0.0f; + q_scales[qg] = q_scale; + + int8_t* qd = q_int8 + qg * QGROUP; + for (size_t i = 0; i < QGROUP / VECTOR_WIDTH; ++i) { + float16x8_t qf = vld1q_f16(q_grp + i * VECTOR_WIDTH); + float32x4_t lo = vmulq_n_f32(vcvt_f32_f16(vget_low_f16(qf)), inv); + float32x4_t hi = vmulq_n_f32(vcvt_f32_f16(vget_high_f16(qf)), inv); + int32x4_t lo_i = vcvtaq_s32_f32(lo); + int32x4_t hi_i = vcvtaq_s32_f32(hi); + int16x8_t pack16 = vcombine_s16(vqmovn_s32(lo_i), vqmovn_s32(hi_i)); + vst1_s8(qd + i * VECTOR_WIDTH, vqmovn_s16(pack16)); + } + } + + float running_max = -std::numeric_limits::infinity(); + float running_sum = 0.0f; + + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vdupq_n_f32(0.0f); + output_accum_high[i] = vdupq_n_f32(0.0f); + } + + const size_t absolute_q_pos = position_offset; + const size_t kv_end = is_causal ? std::min(kv_seq_len, absolute_q_pos + 1) : kv_seq_len; + const size_t kv_start_abs = (window_size > 0 && absolute_q_pos > window_size) + ? absolute_q_pos - window_size : 0; + const size_t kv_start = (position_offset > cache_len) ? 0 : kv_start_abs; + + for (size_t kv_block_start = kv_start; kv_block_start < kv_end; kv_block_start += BLOCK_SIZE) { + const size_t kv_block_end = std::min(kv_block_start + BLOCK_SIZE, kv_end); + const size_t block_size = kv_block_end - kv_block_start; + + float block_max = -std::numeric_limits::infinity(); + + const size_t cached_kv_end = std::min(kv_block_end, cache_len); + const size_t new_kv_start = std::max(kv_block_start, cache_len); + + size_t kv_pos = kv_block_start; + for (; kv_pos + 3 < cached_kv_end; kv_pos += 4) { + const int8_t* k1 = K_cached_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; + const int8_t* k2 = k1 + kv_seq_stride; + const int8_t* k3 = k2 + kv_seq_stride; + const int8_t* k4 = k3 + kv_seq_stride; + const float* ks1 = k_scales + (kv_pos * num_kv_heads + kv_head_idx) * num_quant_groups; + const float* ks2 = ks1 + num_kv_heads * num_quant_groups; + const float* ks3 = ks2 + num_kv_heads * num_quant_groups; + const float* ks4 = ks3 + num_kv_heads * num_quant_groups; + if (kv_pos + 8 < cached_kv_end) { + __builtin_prefetch(k1 + 4 * kv_seq_stride, 0, 0); + __builtin_prefetch(k1 + 5 * kv_seq_stride, 0, 0); + __builtin_prefetch(k1 + 6 * kv_seq_stride, 0, 0); + __builtin_prefetch(k1 + 7 * kv_seq_stride, 0, 0); + } + + float32x4_t sumv1 = vdupq_n_f32(0.0f); + float32x4_t sumv2 = vdupq_n_f32(0.0f); + float32x4_t sumv3 = vdupq_n_f32(0.0f); + float32x4_t sumv4 = vdupq_n_f32(0.0f); + + for (size_t qg = 0; qg < num_quant_groups; ++qg) { + int8x16_t q_lo = vld1q_s8(q_int8 + qg * QGROUP); + int8x16_t q_hi = vld1q_s8(q_int8 + qg * QGROUP + 16); + + int32x4_t d1 = vdupq_n_s32(0); + int32x4_t d2 = vdupq_n_s32(0); + int32x4_t d3 = vdupq_n_s32(0); + int32x4_t d4 = vdupq_n_s32(0); + + d1 = vdotq_s32(d1, q_lo, vld1q_s8(k1 + qg * QGROUP)); + d2 = vdotq_s32(d2, q_lo, vld1q_s8(k2 + qg * QGROUP)); + d3 = vdotq_s32(d3, q_lo, vld1q_s8(k3 + qg * QGROUP)); + d4 = vdotq_s32(d4, q_lo, vld1q_s8(k4 + qg * QGROUP)); + d1 = vdotq_s32(d1, q_hi, vld1q_s8(k1 + qg * QGROUP + 16)); + d2 = vdotq_s32(d2, q_hi, vld1q_s8(k2 + qg * QGROUP + 16)); + d3 = vdotq_s32(d3, q_hi, vld1q_s8(k3 + qg * QGROUP + 16)); + d4 = vdotq_s32(d4, q_hi, vld1q_s8(k4 + qg * QGROUP + 16)); + + float qg_q = q_scales[qg]; + sumv1 = vmlaq_n_f32(sumv1, vcvtq_f32_s32(d1), qg_q * ks1[qg]); + sumv2 = vmlaq_n_f32(sumv2, vcvtq_f32_s32(d2), qg_q * ks2[qg]); + sumv3 = vmlaq_n_f32(sumv3, vcvtq_f32_s32(d3), qg_q * ks3[qg]); + sumv4 = vmlaq_n_f32(sumv4, vcvtq_f32_s32(d4), qg_q * ks4[qg]); + } + float s1 = vaddvq_f32(sumv1) * scale; + float s2 = vaddvq_f32(sumv2) * scale; + float s3 = vaddvq_f32(sumv3) * scale; + float s4 = vaddvq_f32(sumv4) * scale; + block_scores[kv_pos - kv_block_start] = s1; + block_scores[kv_pos - kv_block_start + 1] = s2; + block_scores[kv_pos - kv_block_start + 2] = s3; + block_scores[kv_pos - kv_block_start + 3] = s4; + float local_max = std::max(std::max(s1, s2), std::max(s3, s4)); + if (local_max > block_max) block_max = local_max; + } + for (; kv_pos < cached_kv_end; ++kv_pos) { + const int8_t* k_vec = K_cached_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; + const float* k_scale_base = k_scales + (kv_pos * num_kv_heads + kv_head_idx) * num_quant_groups; + + float32x4_t sumv = vdupq_n_f32(0.0f); + for (size_t qg = 0; qg < num_quant_groups; ++qg) { + int8x16_t q_lo = vld1q_s8(q_int8 + qg * QGROUP); + int8x16_t q_hi = vld1q_s8(q_int8 + qg * QGROUP + 16); + int8x16_t k_lo = vld1q_s8(k_vec + qg * QGROUP); + int8x16_t k_hi = vld1q_s8(k_vec + qg * QGROUP + 16); + int32x4_t dot_acc = vdupq_n_s32(0); + dot_acc = vdotq_s32(dot_acc, q_lo, k_lo); + dot_acc = vdotq_s32(dot_acc, q_hi, k_hi); + sumv = vmlaq_n_f32(sumv, vcvtq_f32_s32(dot_acc), q_scales[qg] * k_scale_base[qg]); + } + float score = vaddvq_f32(sumv) * scale; + block_scores[kv_pos - kv_block_start] = score; + block_max = std::max(block_max, score); + } + + for (kv_pos = std::max(kv_pos, new_kv_start); kv_pos < kv_block_end; ++kv_pos) { + if (is_causal && kv_pos > absolute_q_pos) { + block_scores[kv_pos - kv_block_start] = -std::numeric_limits::infinity(); + continue; + } + const size_t new_pos = kv_pos - cache_len; + const __fp16* k_vec = K_new_base + new_pos * kv_seq_stride + kv_head_idx * head_dim; + float16x8_t s_acc = vdupq_n_f16((__fp16)0.0f); + for (size_t d = 0; d < head_dim; d += VECTOR_WIDTH) { + s_acc = vfmaq_f16(s_acc, vld1q_f16(q_vec + d), vld1q_f16(k_vec + d)); + } + float score = (vaddvq_f32(vcvt_f32_f16(vget_low_f16(s_acc))) + + vaddvq_f32(vcvt_f32_f16(vget_high_f16(s_acc)))) * scale; + block_scores[kv_pos - kv_block_start] = score; + block_max = std::max(block_max, score); + } + + if (block_max > -std::numeric_limits::infinity()) { + float scale_correction = expf(running_max - block_max); + running_sum *= scale_correction; + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vmulq_n_f32(output_accum_low[i], scale_correction); + output_accum_high[i] = vmulq_n_f32(output_accum_high[i], scale_correction); + } + running_max = block_max; + } + + float block_sum = 0.0f; + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + if (block_scores[kv_idx] != -std::numeric_limits::infinity()) { + block_scores[kv_idx] = expf(block_scores[kv_idx] - block_max); + block_sum += block_scores[kv_idx]; + } else { + block_scores[kv_idx] = 0.0f; + } + } + + for (size_t i = 0; i < num_accum_slots; ++i) + block_accum[i] = vdupq_n_f16((__fp16)0.0f); + + const size_t cached_block_end = std::min(kv_block_end, cache_len); + size_t v_kv = kv_block_start; + for (; v_kv + 3 < cached_block_end; v_kv += 4) { + const float w1 = block_scores[v_kv - kv_block_start]; + const float w2 = block_scores[v_kv + 1 - kv_block_start]; + const float w3 = block_scores[v_kv + 2 - kv_block_start]; + const float w4 = block_scores[v_kv + 3 - kv_block_start]; + if (w1 == 0.0f && w2 == 0.0f && w3 == 0.0f && w4 == 0.0f) continue; + + const int8_t* v1 = V_cached_base + v_kv * kv_seq_stride + kv_head_idx * head_dim; + const int8_t* v2 = v1 + kv_seq_stride; + const int8_t* v3 = v2 + kv_seq_stride; + const int8_t* v4 = v3 + kv_seq_stride; + const float* vs1 = v_scales + (v_kv * num_kv_heads + kv_head_idx) * num_quant_groups; + const float* vs2 = vs1 + num_kv_heads * num_quant_groups; + const float* vs3 = vs2 + num_kv_heads * num_quant_groups; + const float* vs4 = vs3 + num_kv_heads * num_quant_groups; + if (v_kv + 8 < cached_block_end) { + __builtin_prefetch(v1 + 4 * kv_seq_stride, 0, 0); + __builtin_prefetch(v1 + 5 * kv_seq_stride, 0, 0); + __builtin_prefetch(v1 + 6 * kv_seq_stride, 0, 0); + __builtin_prefetch(v1 + 7 * kv_seq_stride, 0, 0); + } + + for (size_t qg = 0; qg < num_quant_groups; ++qg) { + const float16x8_t ws1_vec = vdupq_n_f16(static_cast<__fp16>(w1 * vs1[qg])); + const float16x8_t ws2_vec = vdupq_n_f16(static_cast<__fp16>(w2 * vs2[qg])); + const float16x8_t ws3_vec = vdupq_n_f16(static_cast<__fp16>(w3 * vs3[qg])); + const float16x8_t ws4_vec = vdupq_n_f16(static_cast<__fp16>(w4 * vs4[qg])); + #pragma unroll + for (size_t i = 0; i < QGROUP / VECTOR_WIDTH; ++i) { + const size_t d = qg * QGROUP + i * VECTOR_WIDTH; + float16x8_t v1_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(v1 + d))); + float16x8_t v2_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(v2 + d))); + float16x8_t v3_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(v3 + d))); + float16x8_t v4_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(v4 + d))); + float16x8_t acc = block_accum[d / VECTOR_WIDTH]; + acc = vfmaq_f16(acc, v1_f16, ws1_vec); + acc = vfmaq_f16(acc, v2_f16, ws2_vec); + acc = vfmaq_f16(acc, v3_f16, ws3_vec); + acc = vfmaq_f16(acc, v4_f16, ws4_vec); + block_accum[d / VECTOR_WIDTH] = acc; + } + } + } + for (; v_kv < cached_block_end; ++v_kv) { + const float w = block_scores[v_kv - kv_block_start]; + if (w == 0.0f) continue; + const int8_t* v_vec = V_cached_base + v_kv * kv_seq_stride + kv_head_idx * head_dim; + const float* v_scale_base = v_scales + (v_kv * num_kv_heads + kv_head_idx) * num_quant_groups; + for (size_t qg = 0; qg < num_quant_groups; ++qg) { + const float16x8_t ws_vec = vdupq_n_f16(static_cast<__fp16>(w * v_scale_base[qg])); + #pragma unroll + for (size_t i = 0; i < QGROUP / VECTOR_WIDTH; ++i) { + const size_t d = qg * QGROUP + i * VECTOR_WIDTH; + float16x8_t v_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(v_vec + d))); + block_accum[d / VECTOR_WIDTH] = vfmaq_f16(block_accum[d / VECTOR_WIDTH], v_f16, ws_vec); + } + } + } + for (size_t kv_idx = std::max(v_kv, std::max(kv_block_start, cache_len)); kv_idx < kv_block_end; ++kv_idx) { + const float w = block_scores[kv_idx - kv_block_start]; + if (w == 0.0f) continue; + const size_t new_pos = kv_idx - cache_len; + const __fp16* v_vec = V_new_base + new_pos * kv_seq_stride + kv_head_idx * head_dim; + const float16x8_t w_vec = vdupq_n_f16(static_cast<__fp16>(w)); + for (size_t d = 0; d < head_dim; d += VECTOR_WIDTH) { + block_accum[d / VECTOR_WIDTH] = vfmaq_f16(block_accum[d / VECTOR_WIDTH], vld1q_f16(v_vec + d), w_vec); + } + } + + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vaddq_f32(output_accum_low[i], vcvt_f32_f16(vget_low_f16(block_accum[i]))); + output_accum_high[i] = vaddq_f32(output_accum_high[i], vcvt_f32_f16(vget_high_f16(block_accum[i]))); + } + + running_sum += block_sum; + } + + if (running_sum > 0.0f) { + const float inv_sum = 1.0f / running_sum; + const float32x4_t inv_sum_vec = vdupq_n_f32(inv_sum); + for (size_t d = 0; d < head_dim; d += VECTOR_WIDTH) { + size_t idx = d / VECTOR_WIDTH; + vst1q_f16(o_vec + d, vcombine_f16( + vcvt_f16_f32(vmulq_f32(output_accum_low[idx], inv_sum_vec)), + vcvt_f16_f32(vmulq_f32(output_accum_high[idx], inv_sum_vec)))); + } + } else { + memset(o_vec, 0, head_dim * sizeof(__fp16)); + } + } + }); +} + +void cactus_attention_hybrid_int8_fp16( + const __fp16* queries, + const int8_t* keys_cached, + const int8_t* values_cached, + const float* k_scales, + const float* v_scales, + const __fp16* keys_new, + const __fp16* values_new, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t cache_len, + size_t new_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + size_t position_offset, + bool is_causal, + size_t window_size, + size_t quant_group_size, + size_t v_head_dim +) { + if (v_head_dim == 0) v_head_dim = head_dim; + if (scale == 0.0f) { + scale = 1.0f / sqrtf(static_cast(head_dim)); + } + + if (seq_len == 1 && + head_dim == v_head_dim && + head_dim <= 512 && + head_dim % 32 == 0 && + quant_group_size == 32) { + cactus_attention_hybrid_int8_fp16_decode_dot( + queries, keys_cached, values_cached, k_scales, v_scales, + keys_new, values_new, output, + batch_size, cache_len, new_len, + num_q_heads, num_kv_heads, head_dim, + scale, position_offset, is_causal, window_size); + return; + } + + const size_t kv_seq_len = cache_len + new_len; + + constexpr size_t VECTOR_WIDTH = 8; + constexpr size_t BLOCK_SIZE = 32; + const size_t head_dim_aligned = (head_dim / VECTOR_WIDTH) * VECTOR_WIDTH; + const size_t v_head_dim_aligned = (v_head_dim / VECTOR_WIDTH) * VECTOR_WIDTH; + const size_t num_accum_slots = v_head_dim_aligned / VECTOR_WIDTH; + + const size_t gqa_group_size = num_q_heads / num_kv_heads; + const size_t num_quant_groups_k = (head_dim + quant_group_size - 1) / quant_group_size; + const size_t num_quant_groups_v = (v_head_dim + quant_group_size - 1) / quant_group_size; + + const size_t q_batch_stride = seq_len * num_q_heads * head_dim; + const size_t k_cached_batch_stride = cache_len * num_kv_heads * head_dim; + const size_t v_cached_batch_stride = cache_len * num_kv_heads * v_head_dim; + const size_t k_new_batch_stride = new_len * num_kv_heads * head_dim; + const size_t v_new_batch_stride = new_len * num_kv_heads * v_head_dim; + const size_t o_batch_stride = seq_len * num_q_heads * v_head_dim; + const size_t q_seq_stride = num_q_heads * head_dim; + const size_t k_seq_stride = num_kv_heads * head_dim; + const size_t v_seq_stride = num_kv_heads * v_head_dim; + const size_t o_seq_stride = num_q_heads * v_head_dim; + + CactusThreading::parallel_for(batch_size * num_q_heads * seq_len, CactusThreading::Thresholds::ATTENTION, + [=](size_t start_idx, size_t end_idx) { + float block_scores[BLOCK_SIZE]; + std::vector output_accum_low(num_accum_slots); + std::vector output_accum_high(num_accum_slots); + std::vector block_accum(num_accum_slots); + + for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { + const size_t batch_idx = work_idx / (num_q_heads * seq_len); + const size_t remainder = work_idx % (num_q_heads * seq_len); + const size_t q_head_idx = remainder / seq_len; + const size_t q_pos = remainder % seq_len; + + const size_t kv_head_idx = q_head_idx / gqa_group_size; + + const __fp16* Q_base = queries + batch_idx * q_batch_stride; + const int8_t* K_cached_base = keys_cached + batch_idx * k_cached_batch_stride; + const int8_t* V_cached_base = values_cached + batch_idx * v_cached_batch_stride; + const __fp16* K_new_base = keys_new + batch_idx * k_new_batch_stride; + const __fp16* V_new_base = values_new + batch_idx * v_new_batch_stride; + __fp16* O_base = output + batch_idx * o_batch_stride; + + const __fp16* q_vec = Q_base + q_pos * q_seq_stride + q_head_idx * head_dim; + __fp16* o_vec = O_base + q_pos * o_seq_stride + q_head_idx * v_head_dim; + + float running_max = -std::numeric_limits::infinity(); + float running_sum = 0.0f; + + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vdupq_n_f32(0.0f); + output_accum_high[i] = vdupq_n_f32(0.0f); + } + + const size_t absolute_q_pos = position_offset + q_pos; + size_t kv_end = is_causal ? std::min(kv_seq_len, cache_len + q_pos + 1) : kv_seq_len; + + size_t kv_start = 0; + if (window_size > 0 && absolute_q_pos > window_size) { + kv_start = absolute_q_pos - window_size; + } + + constexpr size_t SINK_SIZE = 4; + const size_t cache_abs_offset = (position_offset >= cache_len) ? (position_offset - cache_len) : 0; + + const size_t kv_block_start0 = (window_size > 0 && kv_start > 0) ? 0 + : (kv_start / BLOCK_SIZE) * BLOCK_SIZE; + + for (size_t kv_block_start = kv_block_start0; kv_block_start < kv_end; kv_block_start += BLOCK_SIZE) { + const size_t kv_block_end = std::min(kv_block_start + BLOCK_SIZE, kv_end); + const size_t block_size = kv_block_end - kv_block_start; + + float block_max = -std::numeric_limits::infinity(); + + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + const size_t kv_pos = kv_block_start + kv_idx; + + bool window_masked = false; + if (window_size > 0 && kv_start > 0) { + if (kv_pos < cache_len) { + if (cache_abs_offset == 0 || kv_pos >= SINK_SIZE) { + window_masked = (cache_abs_offset + kv_pos < kv_start); + } + } else { + window_masked = (kv_pos + cache_abs_offset < kv_start); + } + } + + if ((is_causal && kv_pos > absolute_q_pos) || window_masked) { + block_scores[kv_idx] = -std::numeric_limits::infinity(); + continue; + } + + float score = 0.0f; + + if (kv_pos < cache_len) { + if (k_scales != nullptr) { + const int8_t* k_vec = K_cached_base + kv_pos * k_seq_stride + kv_head_idx * head_dim; + const float* k_scale_base = k_scales + (kv_pos * num_kv_heads + kv_head_idx) * num_quant_groups_k; + + for (size_t quant_group = 0; quant_group < num_quant_groups_k; quant_group++) { + const size_t dim_base = quant_group * quant_group_size; + float16x8_t s_acc = vdupq_n_f16((__fp16)0.0f); + + #pragma unroll + for (size_t i = 0; i < 4; i++) { + const size_t dim_block = dim_base + i * VECTOR_WIDTH; + if (dim_block >= head_dim_aligned) break; + + float16x8_t q_f16 = vld1q_f16(&q_vec[dim_block]); + float16x8_t k_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(&k_vec[dim_block]))); + s_acc = vfmaq_f16(s_acc, q_f16, k_f16); + } + + float partial = vaddvq_f32(vcvt_f32_f16(vget_low_f16(s_acc))) + + vaddvq_f32(vcvt_f32_f16(vget_high_f16(s_acc))); + score += k_scale_base[quant_group] * partial; + } + } else { + const __fp16* k_vec = reinterpret_cast(K_cached_base) + + kv_pos * k_seq_stride + kv_head_idx * head_dim; + float16x8_t s_acc = vdupq_n_f16((__fp16)0.0f); + + for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { + float16x8_t q_f16 = vld1q_f16(&q_vec[dim_block]); + float16x8_t k_f16 = vld1q_f16(&k_vec[dim_block]); + s_acc = vfmaq_f16(s_acc, q_f16, k_f16); + } + + score = vaddvq_f32(vcvt_f32_f16(vget_low_f16(s_acc))) + + vaddvq_f32(vcvt_f32_f16(vget_high_f16(s_acc))); + } + } else { + const size_t new_pos = kv_pos - cache_len; + const __fp16* k_vec = K_new_base + new_pos * k_seq_stride + kv_head_idx * head_dim; + + float16x8_t s_acc = vdupq_n_f16((__fp16)0.0f); + + for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { + float16x8_t q_f16 = vld1q_f16(&q_vec[dim_block]); + float16x8_t k_f16 = vld1q_f16(&k_vec[dim_block]); + s_acc = vfmaq_f16(s_acc, q_f16, k_f16); + } + + score = vaddvq_f32(vcvt_f32_f16(vget_low_f16(s_acc))) + + vaddvq_f32(vcvt_f32_f16(vget_high_f16(s_acc))); + } + + score *= scale; + block_scores[kv_idx] = score; + block_max = std::max(block_max, score); + } + + if (block_max > -std::numeric_limits::infinity()) { + float scale_correction = expf(running_max - block_max); + running_sum *= scale_correction; + + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vmulq_n_f32(output_accum_low[i], scale_correction); + output_accum_high[i] = vmulq_n_f32(output_accum_high[i], scale_correction); + } + running_max = block_max; + } + + float block_sum = 0.0f; + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + if (block_scores[kv_idx] != -std::numeric_limits::infinity()) { + block_scores[kv_idx] = expf(block_scores[kv_idx] - block_max); + block_sum += block_scores[kv_idx]; + } else { + block_scores[kv_idx] = 0.0f; + } + } + + for (size_t i = 0; i < num_accum_slots; ++i) + block_accum[i] = vdupq_n_f16((__fp16)0.0f); + + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + const float attn_weight = block_scores[kv_idx]; + if (attn_weight == 0.0f) continue; + + const size_t kv_pos = kv_block_start + kv_idx; + + if (kv_pos < cache_len) { + if (v_scales != nullptr) { + const int8_t* v_vec = V_cached_base + kv_pos * v_seq_stride + kv_head_idx * v_head_dim; + const float* v_scale_base = v_scales + (kv_pos * num_kv_heads + kv_head_idx) * num_quant_groups_v; + + for (size_t quant_group = 0; quant_group < num_quant_groups_v; quant_group++) { + const size_t dim_base = quant_group * quant_group_size; + const float16x8_t ws_vec = vdupq_n_f16(static_cast<__fp16>(attn_weight * v_scale_base[quant_group])); + + #pragma unroll + for (size_t i = 0; i < 4; i++) { + const size_t dim_block = dim_base + i * VECTOR_WIDTH; + if (dim_block >= v_head_dim_aligned) break; + + float16x8_t v_f16 = vcvtq_f16_s16(vmovl_s8(vld1_s8(&v_vec[dim_block]))); + block_accum[dim_block / VECTOR_WIDTH] = vfmaq_f16(block_accum[dim_block / VECTOR_WIDTH], v_f16, ws_vec); + } + } + } else { + const __fp16* v_vec = reinterpret_cast(V_cached_base) + + kv_pos * v_seq_stride + kv_head_idx * v_head_dim; + const float16x8_t w_vec = vdupq_n_f16(static_cast<__fp16>(attn_weight)); + + for (size_t dim_block = 0; dim_block < v_head_dim_aligned; dim_block += VECTOR_WIDTH) { + block_accum[dim_block / VECTOR_WIDTH] = + vfmaq_f16(block_accum[dim_block / VECTOR_WIDTH], vld1q_f16(&v_vec[dim_block]), w_vec); + } + } + } else { + const size_t new_pos = kv_pos - cache_len; + const __fp16* v_vec = V_new_base + new_pos * v_seq_stride + kv_head_idx * v_head_dim; + const float16x8_t w_vec = vdupq_n_f16(static_cast<__fp16>(attn_weight)); + + for (size_t dim_block = 0; dim_block < v_head_dim_aligned; dim_block += VECTOR_WIDTH) { + block_accum[dim_block / VECTOR_WIDTH] = vfmaq_f16(block_accum[dim_block / VECTOR_WIDTH], vld1q_f16(&v_vec[dim_block]), w_vec); + } + } + } + + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vaddq_f32(output_accum_low[i], vcvt_f32_f16(vget_low_f16(block_accum[i]))); + output_accum_high[i] = vaddq_f32(output_accum_high[i], vcvt_f32_f16(vget_high_f16(block_accum[i]))); + } + + running_sum += block_sum; + } + + if (running_sum > 0.0f) { + const float inv_sum = 1.0f / running_sum; + const float32x4_t inv_sum_vec = vdupq_n_f32(inv_sum); + + for (size_t dim_block = 0; dim_block < v_head_dim_aligned; dim_block += VECTOR_WIDTH) { + size_t idx = dim_block / VECTOR_WIDTH; + vst1q_f16(&o_vec[dim_block], vcombine_f16( + vcvt_f16_f32(vmulq_f32(output_accum_low[idx], inv_sum_vec)), + vcvt_f16_f32(vmulq_f32(output_accum_high[idx], inv_sum_vec)))); + } + } else { + memset(o_vec, 0, v_head_dim * sizeof(__fp16)); + } + } + }); +} diff --git a/cactus/kernel/kernel_blas.cpp b/cactus-kernels/src/blas.cpp similarity index 77% rename from cactus/kernel/kernel_blas.cpp rename to cactus-kernels/src/blas.cpp index 92d289982..db58d4e88 100644 --- a/cactus/kernel/kernel_blas.cpp +++ b/cactus-kernels/src/blas.cpp @@ -1,5 +1,5 @@ -#include "kernel.h" -#include "kernel_utils.h" +#include "../cactus_kernels.h" +#include "threading.h" #include #include #include @@ -27,7 +27,7 @@ static inline void increment_coords(size_t* coords, const size_t* shape, size_t enum class BroadcastOp { ADD, SUB, MUL, DIV }; template -static inline float16x8_t broadcast_op_vec(float16x8_t a, float16x8_t b) { +static inline float16x8_t binop_vec(float16x8_t a, float16x8_t b) { if constexpr (Op == BroadcastOp::ADD) return vaddq_f16(a, b); else if constexpr (Op == BroadcastOp::SUB) return vsubq_f16(a, b); else if constexpr (Op == BroadcastOp::MUL) return vmulq_f16(a, b); @@ -35,13 +35,37 @@ static inline float16x8_t broadcast_op_vec(float16x8_t a, float16x8_t b) { } template -static inline __fp16 broadcast_op_scalar(__fp16 a, __fp16 b) { +static inline __fp16 binop_scalar(__fp16 a, __fp16 b) { if constexpr (Op == BroadcastOp::ADD) return a + b; else if constexpr (Op == BroadcastOp::SUB) return a - b; else if constexpr (Op == BroadcastOp::MUL) return a * b; else return a / b; } +template +static void elementwise_binop_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { + const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; + CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + constexpr size_t SIMD_WIDTH = 8; + const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; + + if (use_streaming) { + for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { + stream_store_f16x8(&output[i], binop_vec(vld1q_f16(&a[i]), vld1q_f16(&b[i]))); + } + } else { + for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { + vst1q_f16(&output[i], binop_vec(vld1q_f16(&a[i]), vld1q_f16(&b[i]))); + } + } + + for (size_t i = vectorized_end; i < end_idx; ++i) { + output[i] = binop_scalar(a[i], b[i]); + } + }); +} + template static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* output, const size_t* a_strides, const size_t* b_strides, @@ -85,7 +109,7 @@ static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* out const bool use_stream = total_elements >= STREAMING_STORE_THRESHOLD; if (a_inner_broadcast && b_inner_broadcast) { - __fp16 result = broadcast_op_scalar(a[a_base], b[b_base]); + __fp16 result = binop_scalar(a[a_base], b[b_base]); float16x8_t result_vec = vdupq_n_f16(result); if (use_stream) { for (size_t i = 0; i < vec_end; i += 8) { @@ -105,16 +129,16 @@ static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* out if (use_stream) { for (size_t i = 0; i < vec_end; i += 8) { float16x8_t b_vec = vld1q_f16(b_ptr + i); - stream_store_f16x8(out_ptr + i, broadcast_op_vec(a_vec, b_vec)); + stream_store_f16x8(out_ptr + i, binop_vec(a_vec, b_vec)); } } else { for (size_t i = 0; i < vec_end; i += 8) { float16x8_t b_vec = vld1q_f16(b_ptr + i); - vst1q_f16(out_ptr + i, broadcast_op_vec(a_vec, b_vec)); + vst1q_f16(out_ptr + i, binop_vec(a_vec, b_vec)); } } for (size_t i = vec_end; i < inner_size; ++i) { - out_ptr[i] = broadcast_op_scalar(a[a_base], b_ptr[i]); + out_ptr[i] = binop_scalar(a[a_base], b_ptr[i]); } } else if (b_inner_broadcast) { const __fp16* a_ptr = a + a_base; @@ -122,16 +146,16 @@ static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* out if (use_stream) { for (size_t i = 0; i < vec_end; i += 8) { float16x8_t a_vec = vld1q_f16(a_ptr + i); - stream_store_f16x8(out_ptr + i, broadcast_op_vec(a_vec, b_vec)); + stream_store_f16x8(out_ptr + i, binop_vec(a_vec, b_vec)); } } else { for (size_t i = 0; i < vec_end; i += 8) { float16x8_t a_vec = vld1q_f16(a_ptr + i); - vst1q_f16(out_ptr + i, broadcast_op_vec(a_vec, b_vec)); + vst1q_f16(out_ptr + i, binop_vec(a_vec, b_vec)); } } for (size_t i = vec_end; i < inner_size; ++i) { - out_ptr[i] = broadcast_op_scalar(a_ptr[i], b[b_base]); + out_ptr[i] = binop_scalar(a_ptr[i], b[b_base]); } } else { const __fp16* a_ptr = a + a_base; @@ -140,17 +164,17 @@ static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* out for (size_t i = 0; i < vec_end; i += 8) { float16x8_t a_vec = vld1q_f16(a_ptr + i); float16x8_t b_vec = vld1q_f16(b_ptr + i); - stream_store_f16x8(out_ptr + i, broadcast_op_vec(a_vec, b_vec)); + stream_store_f16x8(out_ptr + i, binop_vec(a_vec, b_vec)); } } else { for (size_t i = 0; i < vec_end; i += 8) { float16x8_t a_vec = vld1q_f16(a_ptr + i); float16x8_t b_vec = vld1q_f16(b_ptr + i); - vst1q_f16(out_ptr + i, broadcast_op_vec(a_vec, b_vec)); + vst1q_f16(out_ptr + i, binop_vec(a_vec, b_vec)); } } for (size_t i = vec_end; i < inner_size; ++i) { - out_ptr[i] = broadcast_op_scalar(a_ptr[i], b_ptr[i]); + out_ptr[i] = binop_scalar(a_ptr[i], b_ptr[i]); } } @@ -176,7 +200,7 @@ static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* out size_t a_idx = compute_linear_index(coords.data(), a_strides, ndim); size_t b_idx = compute_linear_index(coords.data(), b_strides, ndim); - output[linear_idx] = broadcast_op_scalar(a[a_idx], b[b_idx]); + output[linear_idx] = binop_scalar(a[a_idx], b[b_idx]); increment_coords(coords.data(), output_shape, ndim); } @@ -185,40 +209,7 @@ static void broadcast_op_optimized(const __fp16* a, const __fp16* b, __fp16* out } void cactus_add_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { - const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::ELEMENT_WISE, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 4; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&a[i + 256], 0, 0); - __builtin_prefetch(&b[i + 256], 0, 0); - float16x8_t v0 = vaddq_f16(vld1q_f16(&a[i]), vld1q_f16(&b[i])); - float16x8_t v1 = vaddq_f16(vld1q_f16(&a[i + 8]), vld1q_f16(&b[i + 8])); - float16x8_t v2 = vaddq_f16(vld1q_f16(&a[i + 16]), vld1q_f16(&b[i + 16])); - float16x8_t v3 = vaddq_f16(vld1q_f16(&a[i + 24]), vld1q_f16(&b[i + 24])); - stream_store_f16x8(&output[i], v0); - stream_store_f16x8(&output[i + 8], v1); - stream_store_f16x8(&output[i + 16], v2); - stream_store_f16x8(&output[i + 24], v3); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - stream_store_f16x8(&output[i], vaddq_f16(vld1q_f16(&a[i]), vld1q_f16(&b[i]))); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - vst1q_f16(&output[i], vaddq_f16(vld1q_f16(&a[i]), vld1q_f16(&b[i]))); - } - } - - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = a[i] + b[i]; - } - }); + elementwise_binop_f16(a, b, output, num_elements); } void cactus_add_f16_clipped(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { @@ -262,84 +253,40 @@ void cactus_add_f16_clipped(const __fp16* a, const __fp16* b, __fp16* output, si } void cactus_subtract_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { - const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::ELEMENT_WISE, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - if (use_streaming) { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t a_vec = vld1q_f16(&a[i]); - float16x8_t b_vec = vld1q_f16(&b[i]); - stream_store_f16x8(&output[i], vsubq_f16(a_vec, b_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t a_vec = vld1q_f16(&a[i]); - float16x8_t b_vec = vld1q_f16(&b[i]); - vst1q_f16(&output[i], vsubq_f16(a_vec, b_vec)); - } - } - - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = a[i] - b[i]; - } - }); + elementwise_binop_f16(a, b, output, num_elements); } void cactus_multiply_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { - const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::ELEMENT_WISE, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; + elementwise_binop_f16(a, b, output, num_elements); +} - if (use_streaming) { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t a_vec = vld1q_f16(&a[i]); - float16x8_t b_vec = vld1q_f16(&b[i]); - stream_store_f16x8(&output[i], vmulq_f16(a_vec, b_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t a_vec = vld1q_f16(&a[i]); - float16x8_t b_vec = vld1q_f16(&b[i]); - vst1q_f16(&output[i], vmulq_f16(a_vec, b_vec)); - } - } +void cactus_add_scaled_f16(const __fp16* base, const __fp16* src, __fp16* output, size_t num_elements, float scale) { + constexpr size_t SIMD_WIDTH = 8; + const float32x4_t vscale = vdupq_n_f32(scale); + const size_t vec_end = (num_elements / SIMD_WIDTH) * SIMD_WIDTH; - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = a[i] * b[i]; - } - }); -} + for (size_t i = 0; i < vec_end; i += SIMD_WIDTH) { + float16x8_t base_vec = vld1q_f16(base + i); + float16x8_t src_vec = vld1q_f16(src + i); -void cactus_divide_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { - const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::ELEMENT_WISE, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; + float32x4_t base_lo = vcvt_f32_f16(vget_low_f16(base_vec)); + float32x4_t base_hi = vcvt_f32_f16(vget_high_f16(base_vec)); + float32x4_t src_lo = vcvt_f32_f16(vget_low_f16(src_vec)); + float32x4_t src_hi = vcvt_f32_f16(vget_high_f16(src_vec)); - if (use_streaming) { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t a_vec = vld1q_f16(&a[i]); - float16x8_t b_vec = vld1q_f16(&b[i]); - stream_store_f16x8(&output[i], vdivq_f16(a_vec, b_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t a_vec = vld1q_f16(&a[i]); - float16x8_t b_vec = vld1q_f16(&b[i]); - vst1q_f16(&output[i], vdivq_f16(a_vec, b_vec)); - } - } + float32x4_t result_lo = vfmaq_f32(base_lo, src_lo, vscale); + float32x4_t result_hi = vfmaq_f32(base_hi, src_hi, vscale); - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = a[i] / b[i]; - } - }); + vst1q_f16(output + i, vcombine_f16(vcvt_f16_f32(result_lo), vcvt_f16_f32(result_hi))); + } + for (size_t i = vec_end; i < num_elements; ++i) { + output[i] = static_cast<__fp16>(static_cast(base[i]) + + static_cast(src[i]) * scale); + } +} + +void cactus_divide_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements) { + elementwise_binop_f16(a, b, output, num_elements); } void cactus_add_broadcast_f16(const __fp16* a, const __fp16* b, __fp16* output, @@ -384,6 +331,15 @@ void cactus_concat_f16(const __fp16* input1, const __fp16* input2, __fp16* outpu size_t axis_size1 = shape1[axis]; size_t axis_size2 = shape2[axis]; + size_t outer_size1 = 1; + size_t outer_size2 = 1; + for (size_t i = 0; i < static_cast(axis); ++i) { + outer_size1 *= shape1[i]; + outer_size2 *= shape2[i]; + } + if (outer_size1 == 0) outer_size1 = 1; + if (outer_size2 == 0) outer_size2 = 1; + size_t input1_stride = axis_size1 * inner_size; size_t input2_stride = axis_size2 * inner_size; size_t output_stride = (axis_size1 + axis_size2) * inner_size; @@ -391,8 +347,8 @@ void cactus_concat_f16(const __fp16* input1, const __fp16* input2, __fp16* outpu CactusThreading::parallel_for(outer_size, CactusThreading::Thresholds::ELEMENT_WISE, [&](size_t start, size_t end) { for (size_t outer = start; outer < end; ++outer) { - const __fp16* in1_ptr = input1 + outer * input1_stride; - const __fp16* in2_ptr = input2 + outer * input2_stride; + const __fp16* in1_ptr = input1 + (outer % outer_size1) * input1_stride; + const __fp16* in2_ptr = input2 + (outer % outer_size2) * input2_stride; __fp16* out_ptr = output + outer * output_stride; size_t copy_size1 = axis_size1 * inner_size; @@ -404,6 +360,45 @@ void cactus_concat_f16(const __fp16* input1, const __fp16* input2, __fp16* outpu }); } +void cactus_cat_f16(const __fp16** inputs, __fp16* output, const size_t** input_shapes, + const size_t* output_shape, size_t num_inputs, size_t rank, int axis) { + if (axis < 0) axis += rank; + + size_t outer_size = 1; + for (size_t i = 0; i < static_cast(axis); ++i) { + outer_size *= output_shape[i]; + } + + size_t inner_size = 1; + for (size_t i = axis + 1; i < rank; ++i) { + inner_size *= output_shape[i]; + } + + std::vector input_outer(num_inputs, 1); + for (size_t input_idx = 0; input_idx < num_inputs; ++input_idx) { + for (size_t i = 0; i < static_cast(axis); ++i) { + input_outer[input_idx] *= input_shapes[input_idx][i]; + } + if (input_outer[input_idx] == 0) input_outer[input_idx] = 1; + } + + CactusThreading::parallel_for(outer_size, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start, size_t end) { + for (size_t outer = start; outer < end; ++outer) { + __fp16* out_ptr = output + outer * inner_size * output_shape[axis]; + size_t offset = 0; + + for (size_t input_idx = 0; input_idx < num_inputs; ++input_idx) { + size_t src_outer = outer % input_outer[input_idx]; + const __fp16* in_ptr = inputs[input_idx] + src_outer * inner_size * input_shapes[input_idx][axis]; + size_t copy_size = input_shapes[input_idx][axis] * inner_size; + std::memcpy(out_ptr + offset, in_ptr, copy_size * sizeof(__fp16)); + offset += copy_size; + } + } + }); +} + void cactus_transpose_2d_f16(const __fp16* source, __fp16* destination, size_t num_rows, size_t num_cols, size_t start_row, size_t end_row) { constexpr size_t TILE_SIZE = 32; constexpr size_t VECTOR_WIDTH = 8; @@ -533,4 +528,4 @@ void cactus_transpose_f16(const __fp16* source, __fp16* destination, const size_ destination[idx] = source[src_idx]; } } -} \ No newline at end of file +} diff --git a/cactus-kernels/src/cactus_kernels.metal b/cactus-kernels/src/cactus_kernels.metal new file mode 100644 index 000000000..01c3b8b4d --- /dev/null +++ b/cactus-kernels/src/cactus_kernels.metal @@ -0,0 +1,3795 @@ +#include +using namespace metal; +#define ROWS 8u + +inline uint cq_idx(device const uchar* packed, uint k, uint bits) { + if (bits == 4u) return (packed[k >> 1] >> ((k & 1u) * 4u)) & 0xFu; + if (bits == 2u) return (packed[k >> 2] >> ((k & 3u) * 2u)) & 0x3u; + uint bit_pos = k * 3u; + uint byte_idx = bit_pos >> 3, bit_idx = bit_pos & 7u; + uint val = (uint)packed[byte_idx] >> bit_idx; + if (bit_idx > 5u) val |= (uint)packed[byte_idx + 1u] << (8u - bit_idx); + return val & 0x7u; +} + +inline float gelu_tanh(float x) { + float c = 0.7978845608028654f * (x + 0.044715f*x*x*x); + return 0.5f * x * (1.0f + precise::tanh(c)); +} + +static inline void cq4_hada128(thread float& x0, thread float& x1, + thread float& x2, thread float& x3, uint lane) { + float a0=x0+x1,a1=x0-x1,a2=x2+x3,a3=x2-x3; + x0=a0+a2; x1=a1+a3; x2=a0-a2; x3=a1-a3; + #pragma clang loop unroll(full) + for (uint d=1u; d<=16u; d<<=1){ + bool hi=(lane&d)!=0u; + float p0=simd_shuffle_xor(x0,d),p1=simd_shuffle_xor(x1,d),p2=simd_shuffle_xor(x2,d),p3=simd_shuffle_xor(x3,d); + x0=hi?p0-x0:x0+p0; x1=hi?p1-x1:x1+p1; x2=hi?p2-x2:x2+p2; x3=hi?p3-x3:x3+p3; + } + float s=rsqrt(128.0f); + x0*=s; x1*=s; x2*=s; x3*=s; +} + +static inline float cq4_dot64(half4 c0, half4 c1, half4 c2, half4 c3, + ushort4 w, threadgroup const float* cb) { + return (float)c0.x*cb[w[0]&0xF] + (float)c0.y*cb[(w[0]>>4)&0xF] + (float)c0.z*cb[(w[0]>>8)&0xF] + (float)c0.w*cb[(w[0]>>12)&0xF] + + (float)c1.x*cb[w[1]&0xF] + (float)c1.y*cb[(w[1]>>4)&0xF] + (float)c1.z*cb[(w[1]>>8)&0xF] + (float)c1.w*cb[(w[1]>>12)&0xF] + + (float)c2.x*cb[w[2]&0xF] + (float)c2.y*cb[(w[2]>>4)&0xF] + (float)c2.z*cb[(w[2]>>8)&0xF] + (float)c2.w*cb[(w[2]>>12)&0xF] + + (float)c3.x*cb[w[3]&0xF] + (float)c3.y*cb[(w[3]>>4)&0xF] + (float)c3.z*cb[(w[3]>>8)&0xF] + (float)c3.w*cb[(w[3]>>12)&0xF]; +} + +static inline device const uchar4* cq4_il_base(device const uchar* packed, uint n, uint ng, uint g, uint pgb) { + return (device const uchar4*)(packed + ((size_t)(n>>2)*ng+g)*4u*(size_t)pgb + (n&3u)*4u); +} + +static inline float cq4_dot64_il(half4 c0, half4 c1, half4 c2, half4 c3, + uchar4 A, uchar4 B, threadgroup const float* cb) { + return (float)c0.x*cb[A.x&0xFu] + (float)c0.y*cb[A.y&0xFu] + (float)c0.z*cb[A.z&0xFu] + (float)c0.w*cb[A.w&0xFu] + + (float)c1.x*cb[A.x>>4] + (float)c1.y*cb[A.y>>4] + (float)c1.z*cb[A.z>>4] + (float)c1.w*cb[A.w>>4] + + (float)c2.x*cb[B.x&0xFu] + (float)c2.y*cb[B.y&0xFu] + (float)c2.z*cb[B.z&0xFu] + (float)c2.w*cb[B.w&0xFu] + + (float)c3.x*cb[B.x>>4] + (float)c3.y*cb[B.y>>4] + (float)c3.z*cb[B.z>>4] + (float)c3.w*cb[B.w>>4]; +} + +static inline float cq4_norm_at(device const half* norms, uint n, uint ng, uint g, uint il) { + return (float)(il != 0u ? norms[(((size_t)(n>>2)*ng+g)<<2) + (n&3u)] + : norms[(size_t)n*ng+g]); +} + +static inline uint cq4_il_nib(device const uchar* packed, uint n, uint ng, uint g, uint pgb, uint e) { + uint v = e>>4, j = e&15u, sub = j>>2, b = j&3u; + device const uchar* panel = packed + ((size_t)(n>>2)*ng+g)*4u*(size_t)pgb; + uchar by = panel[(2u*v + (sub>>1))*16u + (n&3u)*4u + b]; + return (sub&1u) ? (uint)(by>>4) : (uint)(by&0xFu); +} + +static inline uint cq_il_idx(device const uchar* packed, uint n, uint ng, uint g, uint pgb, uint e, uint bits) { + if (bits == 4u) return cq4_il_nib(packed, n, ng, g, pgb, e); + device const uchar* panel = packed + ((size_t)(n>>2)*ng+g)*4u*(size_t)pgb; + uint r = n&3u; + if (bits == 2u) { + uchar by = panel[(e>>4)*16u + ((e&15u)>>2)*4u + r]; + return ((uint)by >> (2u*(e&3u))) & 3u; + } + device const uchar* ch = panel + (e>>2)*6u; + ulong word = (ulong)ch[0] | ((ulong)ch[1]<<8) | ((ulong)ch[2]<<16) + | ((ulong)ch[3]<<24) | ((ulong)ch[4]<<32) | ((ulong)ch[5]<<40); + return (uint)((word >> (r*12u + (e&3u)*3u)) & 7u); +} + +static inline float cq4_dot64_ilw(half4 c0, half4 c1, half4 c2, half4 c3, + uint A, uint B, threadgroup const float* cb) { + return (float)c0.x*cb[A&0xFu] + (float)c0.y*cb[(A>>8)&0xFu] + (float)c0.z*cb[(A>>16)&0xFu] + (float)c0.w*cb[(A>>24)&0xFu] + + (float)c1.x*cb[(A>>4)&0xFu] + (float)c1.y*cb[(A>>12)&0xFu] + (float)c1.z*cb[(A>>20)&0xFu] + (float)c1.w*cb[A>>28] + + (float)c2.x*cb[B&0xFu] + (float)c2.y*cb[(B>>8)&0xFu] + (float)c2.z*cb[(B>>16)&0xFu] + (float)c2.w*cb[(B>>24)&0xFu] + + (float)c3.x*cb[(B>>4)&0xFu] + (float)c3.y*cb[(B>>12)&0xFu] + (float)c3.z*cb[(B>>20)&0xFu] + (float)c3.w*cb[B>>28]; +} + +kernel void cq4_transform( + device const half* x [[buffer(0)]], + device const half* recip [[buffer(1)]], + device const char* lsign [[buffer(2)]], + device const char* rsign [[buffer(3)]], + device const uint* perm [[buffer(4)]], + device half* code [[buffer(5)]], + constant uint& gs [[buffer(6)]], + uint g [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* z [[threadgroup(0)]]) +{ + for (uint k=t; k=N) return; + uint K = num_groups*gs; + float acc=0; + for (uint base=lane*CQ4_VPL; base>4)*8u], B = pp[(off>>4)*8u+4u]; + float p = cq4_dot64_il(cbase[0], cbase[1], cbase[2], cbase[3], A, B, cb); + acc += (float)norms[(((size_t)(n>>2)*num_groups+g)<<2) + (n&3u)]*p; + } else { + device const ushort4* pr=(device const ushort4*)(packed + ((size_t)n*num_groups+g)*pgb + off/2u); + ushort4 w=pr[0]; + float p = 0; + #pragma clang loop unroll(full) + for (uint q=0;q<4;++q){ + half4 c=cbase[q]; + ushort ww = w[q]; + p += (float)c.x*cb[ww&0xF] + (float)c.y*cb[(ww>>4)&0xF] + + (float)c.z*cb[(ww>>8)&0xF] + (float)c.w*cb[(ww>>12)&0xF]; + } + acc += (float)norms[(size_t)n*num_groups+g]*p; + } + } + acc=simd_sum(acc); + if (lane==0) y[n]=(half)acc; +} + +#define CQ4_NR 2u +kernel void cq4_gemv_mr( + device const half* code [[buffer(0)]], + device const uchar* packed [[buffer(1)]], + device const half* codebook [[buffer(2)]], + device const half* norms [[buffer(3)]], + device half* y [[buffer(4)]], + constant uint& gs [[buffer(5)]], + constant uint& num_groups [[buffer(6)]], + constant uint& pgb [[buffer(7)]], + constant uint& N [[buffer(8)]], + constant uint& il [[buffer(9)]], + uint tg [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint tl [[thread_index_in_threadgroup]]) +{ + threadgroup float cb[16]; + if (tl<16) cb[tl]=(float)codebook[tl]; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint n0 = (tg*ROWS + sgid)*CQ4_NR; + if (n0>=N) return; + uint K = num_groups*gs; + float acc[CQ4_NR]; + #pragma clang loop unroll(full) + for (uint r=0;r>2)*num_groups+g)*4u*(size_t)pgb + + (n0&3u)*4u + (off>>4)*32u; + uint2 Aw = ((device const uint2*)pan)[0]; + uint2 Bw = ((device const uint2*)(pan+16u))[0]; + #pragma clang loop unroll(full) + for (uint r=0;r>2)*num_groups+g)<<2) + (n0&3u) + r]*p; + } + } else { + #pragma clang loop unroll(full) + for (uint r=0;r>2)*ng+g)*256u + (n0&3u)*4u + (off>>4)*32u; + uint2 Aw = ((device const uint2*)pan)[0]; + uint2 Bw = ((device const uint2*)(pan+16u))[0]; + #pragma clang loop unroll(full) + for (uint r=0u;r<2u;r++){ + float p = cq4_dot64_ilw(c0, c1, c2, c3, r?Aw.y:Aw.x, r?Bw.y:Bw.x, cb); + acc[r] += (float)norms[(((size_t)(n0>>2)*ng+g)<<2) + (n0&3u) + r]*p; + } + } + #pragma clang loop unroll(full) + for (uint r=0u;r<2u;r++){ + float a = simd_sum(acc[r]); + if (lane==0u){ + uint n = n0+r; + if (U.oswi != 0u) { + float gg = gelu_tanh((float)(half)a) * (float)osw[n]; + y[n] = (half)clamp(gg,-65504.0f,65504.0f); + } else { + y[n] = (half)a; + } + } + } +} + +kernel void cq4_swiglu_transform( + device const half* gate [[buffer(0)]], + device const half* up [[buffer(1)]], + device const half* recip [[buffer(2)]], + device const char* lsign [[buffer(3)]], + device const char* rsign [[buffer(4)]], + device const uint* perm [[buffer(5)]], + device half* code [[buffer(6)]], + constant float& scale [[buffer(7)]], + uint g [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + threadgroup float* zmem [[threadgroup(0)]]) +{ + uint b = g*128u + lane*4u; + uint k = lane*4u; + float x0,x1,x2,x3; + { + float g0 = gelu_tanh((float)gate[b+0]) * scale * (float)up[b+0]; + float g1 = gelu_tanh((float)gate[b+1]) * scale * (float)up[b+1]; + float g2 = gelu_tanh((float)gate[b+2]) * scale * (float)up[b+2]; + float g3 = gelu_tanh((float)gate[b+3]) * scale * (float)up[b+3]; + x0=(float)(half)clamp(g0,-65504.0f,65504.0f)*(float)recip[b+0]*(float)lsign[k+0]; + x1=(float)(half)clamp(g1,-65504.0f,65504.0f)*(float)recip[b+1]*(float)lsign[k+1]; + x2=(float)(half)clamp(g2,-65504.0f,65504.0f)*(float)recip[b+2]*(float)lsign[k+2]; + x3=(float)(half)clamp(g3,-65504.0f,65504.0f)*(float)recip[b+3]*(float)lsign[k+3]; + } + cq4_hada128(x0, x1, x2, x3, lane); + zmem[k+0]=x0*(float)rsign[k+0]; zmem[k+1]=x1*(float)rsign[k+1]; + zmem[k+2]=x2*(float)rsign[k+2]; zmem[k+3]=x3*(float)rsign[k+3]; + threadgroup_barrier(mem_flags::mem_threadgroup); + code[b+0]=(half)zmem[perm[k+0]]; code[b+1]=(half)zmem[perm[k+1]]; + code[b+2]=(half)zmem[perm[k+2]]; code[b+3]=(half)zmem[perm[k+3]]; +} + +struct CatU { uint ng, N0, N1, N2; }; +kernel void cq4_gemv_cat( + device const half* code0 [[buffer(0)]], device const half* code1 [[buffer(1)]], device const half* code2 [[buffer(2)]], + device const uchar* packed0[[buffer(3)]], device const uchar* packed1[[buffer(4)]], device const uchar* packed2[[buffer(5)]], + device const half* cbook0 [[buffer(6)]], device const half* cbook1 [[buffer(7)]], device const half* cbook2 [[buffer(8)]], + device const half* norms0 [[buffer(9)]], device const half* norms1 [[buffer(10)]], device const half* norms2 [[buffer(11)]], + device half* y0 [[buffer(12)]], device half* y1 [[buffer(13)]], device half* y2 [[buffer(14)]], + constant CatU& U [[buffer(15)]], + uint tg [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint tl [[thread_index_in_threadgroup]]) +{ + uint n0 = tg*16u; + uint b, nbase; + if (n0 < U.N0) { b = 0u; nbase = 0u; } + else if (n0 < U.N0 + U.N1) { b = 1u; nbase = U.N0; } + else { b = 2u; nbase = U.N0 + U.N1; } + device const half* code = (b==0u)?code0 :((b==1u)?code1 :code2 ); + device const uchar* packed= (b==0u)?packed0:((b==1u)?packed1:packed2); + device const half* cbook = (b==0u)?cbook0:((b==1u)?cbook1:cbook2); + device const half* norms = (b==0u)?norms0:((b==1u)?norms1:norms2); + device half* y = (b==0u)?y0 :((b==1u)?y1 :y2 ); + threadgroup float cb[16]; + if (tl<16) cb[tl]=(float)cbook[tl]; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint ng = U.ng; + uint K = ng*128u; + uint nl0 = (n0 - nbase) + sgid*2u; + float acc[2]; acc[0]=0; acc[1]=0; + for (uint base=lane*16u; base>2)*ng+g)*256u + (nl0&3u)*4u + (off>>4)*32u; + uint2 Aw = ((device const uint2*)pan)[0]; + uint2 Bw = ((device const uint2*)(pan+16u))[0]; + #pragma clang loop unroll(full) + for (uint r=0;r<2;r++){ + float p = cq4_dot64_ilw(c0, c1, c2, c3, r?Aw.y:Aw.x, r?Bw.y:Bw.x, cb); + acc[r] += (float)norms[(((size_t)(nl0>>2)*ng+g)<<2) + (nl0&3u) + r]*p; + } + } + #pragma clang loop unroll(full) + for (uint r=0;r<2;r++){ + float a = simd_sum(acc[r]); + if (lane==0) y[nl0+r]=(half)a; + } +} + +kernel void cq_gemv_mr_lowbit( + device const half* code [[buffer(0)]], + device const uchar* packed [[buffer(1)]], + device const half* codebook [[buffer(2)]], + device const half* norms [[buffer(3)]], + device half* y [[buffer(4)]], + constant uint& gs [[buffer(5)]], + constant uint& num_groups [[buffer(6)]], + constant uint& pgb [[buffer(7)]], + constant uint& N [[buffer(8)]], + constant uint& bits [[buffer(9)]], + constant uint& il [[buffer(10)]], + uint tg [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint tl [[thread_index_in_threadgroup]]) +{ + threadgroup float cb[8]; + uint cbn = 1u << bits; + if (tl < cbn) cb[tl] = (float)codebook[tl]; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint n0 = (tg*ROWS + sgid)*CQ4_NR; + if (n0 >= N) return; + uint K = num_groups*gs; + float acc[CQ4_NR]; + #pragma clang loop unroll(full) + for (uint r=0;r> 3; + #pragma clang loop unroll(full) + for (uint r=0;r>2)*num_groups+g)*4u*(size_t)pgb + byte_off*4u; + uint rr = n&3u; + w = (uint)pan[rr] | ((uint)pan[4u+rr]<<8) | ((uint)pan[8u+rr]<<16) | ((uint)pan[12u+rr]<<24); + } else { + w = *(device const uint*)(packed + ((size_t)n*num_groups+g)*pgb + byte_off); + } + w0=cb[w&3u]; w1=cb[(w>>2)&3u]; w2=cb[(w>>4)&3u]; w3=cb[(w>>6)&3u]; + w4=cb[(w>>8)&3u]; w5=cb[(w>>10)&3u]; w6=cb[(w>>12)&3u]; w7=cb[(w>>14)&3u]; + w8=cb[(w>>16)&3u]; w9=cb[(w>>18)&3u]; wa=cb[(w>>20)&3u]; wb=cb[(w>>22)&3u]; + wc=cb[(w>>24)&3u]; wd=cb[(w>>26)&3u]; we=cb[(w>>28)&3u]; wf=cb[(w>>30)&3u]; + } else { + ulong v; + if (il != 0u) { + device const uchar* pan = packed + ((size_t)(n>>2)*num_groups+g)*4u*(size_t)pgb + byte_off*4u; + uint rr = n&3u; + v = 0; + #pragma clang loop unroll(full) + for (uint q=0;q<4u;q++){ + device const uchar* ch = pan + q*6u; + ulong word = (ulong)ch[0] | ((ulong)ch[1]<<8) | ((ulong)ch[2]<<16) + | ((ulong)ch[3]<<24) | ((ulong)ch[4]<<32) | ((ulong)ch[5]<<40); + v |= ((word >> (rr*12u)) & 0xFFFul) << (12u*q); + } + } else { + device const ushort* ps = (device const ushort*)(packed + ((size_t)n*num_groups+g)*pgb + byte_off); + v = (ulong)ps[0] | ((ulong)ps[1] << 16) | ((ulong)ps[2] << 32); + } + w0=cb[v&7u]; w1=cb[(v>>3)&7u]; w2=cb[(v>>6)&7u]; w3=cb[(v>>9)&7u]; + w4=cb[(v>>12)&7u]; w5=cb[(v>>15)&7u]; w6=cb[(v>>18)&7u]; w7=cb[(v>>21)&7u]; + w8=cb[(v>>24)&7u]; w9=cb[(v>>27)&7u]; wa=cb[(v>>30)&7u]; wb=cb[(v>>33)&7u]; + wc=cb[(v>>36)&7u]; wd=cb[(v>>39)&7u]; we=cb[(v>>42)&7u]; wf=cb[(v>>45)&7u]; + } + float pacc = (float)c0.x*w0 + (float)c0.y*w1 + (float)c0.z*w2 + (float)c0.w*w3 + + (float)c1.x*w4 + (float)c1.y*w5 + (float)c1.z*w6 + (float)c1.w*w7 + + (float)c2.x*w8 + (float)c2.y*w9 + (float)c2.z*wa + (float)c2.w*wb + + (float)c3.x*wc + (float)c3.y*wd + (float)c3.z*we + (float)c3.w*wf; + acc[r] += cq4_norm_at(norms, n, num_groups, g, il)*pacc; + } + } + #pragma clang loop unroll(full) + for (uint r=0;r= N) return; + float acc = 0.f; + for (uint g = 0u; g < num_groups; ++g) { + uchar b = packed[((size_t)(n>>2)*num_groups + g)*128u + (lane>>2)*16u + (lane&3u)*4u + (n&3u)]; + device const char* a = act_i8 + (size_t)g*128u + lane*4u; + int partial = (int)a[0]*cb[b & 3u] + (int)a[1]*cb[(b >> 2) & 3u] + + (int)a[2]*cb[(b >> 4) & 3u] + (int)a[3]*cb[(b >> 6) & 3u]; + int idot = simd_sum(partial); + if (lane == 0u) { + float ns = (float)norms[(((size_t)(n>>2)*num_groups+g)<<2) + (n&3u)] * cb_scale; + float sc = ns * act_scales[g]; + acc = fma((float)idot, sc, acc); + } + } + if (lane == 0u) y[n] = (half)acc; +} + +kernel void cq4_transform_m( + device const half* x [[buffer(0)]], + device const half* recip [[buffer(1)]], + device const char* lsign [[buffer(2)]], + device const char* rsign [[buffer(3)]], + device const uint* perm [[buffer(4)]], + device half* code [[buffer(5)]], + constant uint& gs [[buffer(6)]], + constant uint& K [[buffer(7)]], + uint pos [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* z [[threadgroup(0)]]) +{ + uint ng = K/gs; + uint g = pos % ng, row = pos / ng; + size_t xb = (size_t)row*K + (size_t)g*gs; + for (uint k=t; k ma[4], mb[2]; + simdgroup_matrix mc[8]; + for (uint i=0;i<8u;i++) mc[i]=make_filled_simdgroup_matrix(0.f); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint loop_k=0; loop_k>2)*num_groups+g)<<2) + (nrow&3u)]:0.f; + float wv[16]; + if (nrow>4; + uchar4 A = pp[v4*8u], B = pp[v4*8u+4u]; + wv[0]=cb[A.x&0xFu]; wv[1]=cb[A.y&0xFu]; wv[2]=cb[A.z&0xFu]; wv[3]=cb[A.w&0xFu]; + wv[4]=cb[A.x>>4]; wv[5]=cb[A.y>>4]; wv[6]=cb[A.z>>4]; wv[7]=cb[A.w>>4]; + wv[8]=cb[B.x&0xFu]; wv[9]=cb[B.y&0xFu]; wv[10]=cb[B.z&0xFu]; wv[11]=cb[B.w&0xFu]; + wv[12]=cb[B.x>>4]; wv[13]=cb[B.y>>4]; wv[14]=cb[B.z>>4]; wv[15]=cb[B.w>>4]; + } + for (uint i=0;i<16u;i++){ + uint sx=2u*il0+i/8u, sy=(tl/NL0)/8u, lx=(tl/NL0)%8u, ly=i%8u, ib=8u*sx+sy; + float v=0.f; + if (loop_k+16u*il0+i>1u), n_sg = 32u*(sg&1u); + for (uint i=0;i<8u;i++) + simdgroup_store(mc[i], &Cs[(m_sg+8u*(i/4u))*64u + n_sg+8u*(i%4u)], 64u); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint idx=tl; idx<32u*64u; idx+=128u){ + uint m=idx/64u, n=idx%64u; + if (r1+m ma[4], mb[2]; + simdgroup_matrix mc[8]; + for (uint i=0;i<8u;i++) mc[i]=make_filled_simdgroup_matrix(0.f); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint loop_k=0; loop_k>1u), n_sg = 32u*(sg&1u); + for (uint i=0;i<8u;i++) + simdgroup_store(mc[i], &Cs[(m_sg+8u*(i/4u))*64u + n_sg+8u*(i%4u)], 64u); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint idx=tl; idx<32u*64u; idx+=128u){ + uint m=idx/64u, n=idx%64u; + if (r1+m> 5u; + uint j = tg*32u + c; + float acc = 0; + if (j < K) for (uint k=r; k= K) return; + device const half* rrow = rotation + (size_t)j*K; + float acc=0; + for (uint i=lane*4u; i=n) return; + uint row=i/D, d=i%D; + out[i] = table[(size_t)rows[row]*D + d]; +} + +kernel void copy_bytes(device const uchar* in [[buffer(0)]], device uchar* out [[buffer(1)]], + constant uint& n [[buffer(2)]], uint i [[thread_position_in_grid]]) { + if (i= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) { + float av=(float)a[i0+j], bv=(float)b[i0+j], r; + switch(op){ case 2: r=av-bv; break; case 3: r=av*bv; break; case 4: r=av/bv; break; + case 5: r=(av!=bv)?1.0f:0.0f; break; default: r=av+bv; } + if (op==1) r=clamp(r,-65500.0f,65500.0f); + y[i0+j]=(half)r; + } +} + +static inline float erf_approx(float x) { + float sgn = x < 0.0f ? -1.0f : 1.0f; + float ax = fabs(x); + float t = 1.0f / (1.0f + 0.3275911f * ax); + float poly = t * (0.254829592f + t * (-0.284496736f + t * (1.421413741f + + t * (-1.453152027f + t * 1.061405429f)))); + return sgn * (1.0f - poly * precise::exp(-ax * ax)); +} + +static inline float scalar_apply(float v, float p, int op) { + switch (op) { + case 0: return v+p; + case 1: return v-p; + case 3: return v/p; + case 4: return precise::exp(v); + case 5: return precise::sqrt(v); + case 6: return precise::cos(v); + case 7: return precise::sin(v); + case 8: return precise::log(v); + case 9: return fabs(v); + case 10: return precise::pow(v, p); + case 11: return (v != p) ? 1.0f : 0.0f; + case 12: return (v > 0.0f) ? v : v*p; + default: return v*p; + } +} + +kernel void scalar_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant int& op [[buffer(3)]], + constant float& p [[buffer(4)]], uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) y[i0+j]=(half)scalar_apply((float)in[i0+j], p, op); +} + +kernel void unary_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant int& op [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) { + float x=(float)in[i0+j], r; + if (op==0) r=gelu_tanh(x); + else if (op==1) r=precise::tanh(x); + else if (op==2) r=x/(1.0f+precise::exp(-x)); + else if (op==4) r=0.5f*x*(1.0f+erf_approx(x*0.70710678f)); + else if (op==5) r=1.0f/(1.0f+precise::exp(-x)); + else r=max(x,0.0f); + y[i0+j]=(half)r; + } +} + +kernel void swiglu_f16(device const half* gate [[buffer(0)]], device const half* up [[buffer(1)]], + device half* y [[buffer(2)]], constant uint& n [[buffer(3)]], + constant float& scale [[buffer(4)]], uint i [[thread_position_in_grid]]) { + if (i>=n) return; + float x=(float)gate[i]; + half g1=(half)gelu_tanh(x); + half g2=(half)((float)g1*scale); + y[i]=(half)((float)g2*(float)up[i]); +} + +kernel void rms_norm_add_simd_f16(device const half* x [[buffer(0)]], device const half* res [[buffer(1)]], + device const half* w [[buffer(2)]], device half* ysum [[buffer(3)]], + device half* ynorm [[buffer(4)]], constant uint& dim [[buffer(5)]], + constant float& eps [[buffer(6)]], constant uint& rows [[buffer(7)]], + constant uint& clipped [[buffer(8)]], + uint tg [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + uint row = tg * 4 + sgid; + if (row >= rows) return; + device const half* xr = x + (size_t)row * dim; + device const half* rr = res + (size_t)row * dim; + device half* os = ysum + (size_t)row * dim; + device half* on = ynorm + (size_t)row * dim; + float partial = 0; + for (uint i = lane; i < dim; i += 32) { + float v = (float)xr[i] + (float)rr[i]; + if (clipped) v = clamp(v, -65500.0f, 65500.0f); + half h = (half)v; + os[i] = h; + float f = (float)h; + partial += f * f; + } + partial += simd_shuffle_xor(partial, 16); + partial += simd_shuffle_xor(partial, 8); + partial += simd_shuffle_xor(partial, 4); + partial += simd_shuffle_xor(partial, 2); + partial += simd_shuffle_xor(partial, 1); + float inv = 1.0f / sqrt(partial / (float)dim + eps); + for (uint i = lane; i < dim; i += 32) on[i] = (half)((float)os[i] * inv * (float)w[i]); +} + +kernel void rms_norm_simd_f16(device const half* in [[buffer(0)]], device const half* w [[buffer(1)]], + device half* y [[buffer(2)]], constant uint& dim [[buffer(3)]], + constant float& eps [[buffer(4)]], constant uint& rows [[buffer(5)]], + uint tg [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + uint row = tg * 4 + sgid; + if (row >= rows) return; + device const half* x = in + (size_t)row * dim; + device half* o = y + (size_t)row * dim; + float partial = 0; + for (uint i = lane * 4; i + 3 < dim; i += 128) { + half4 v4 = *(device const half4*)(x + i); + float4 f = float4(v4); + partial += f.x*f.x + f.y*f.y + f.z*f.z + f.w*f.w; + } + for (uint i = (dim & ~3u) + lane; i < dim; i += 32) { float v = (float)x[i]; partial += v*v; } + partial += simd_shuffle_xor(partial, 16); + partial += simd_shuffle_xor(partial, 8); + partial += simd_shuffle_xor(partial, 4); + partial += simd_shuffle_xor(partial, 2); + partial += simd_shuffle_xor(partial, 1); + float inv = 1.0f / sqrt(partial / (float)dim + eps); + for (uint i = lane * 4; i + 3 < dim; i += 128) { + half4 v4 = *(device const half4*)(x + i); + half4 w4 = *(device const half4*)(w + i); + float4 f = float4(v4) * inv * float4(w4); + *(device half4*)(o + i) = half4(f); + } + for (uint i = (dim & ~3u) + lane; i < dim; i += 32) o[i] = (half)((float)x[i] * inv * (float)w[i]); +} + +kernel void rms_norm_f16(device const half* in [[buffer(0)]], device const half* w [[buffer(1)]], + device half* y [[buffer(2)]], constant uint& dim [[buffer(3)]], + constant float& eps [[buffer(4)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint nt [[threads_per_threadgroup]], threadgroup float* red [[threadgroup(0)]]) { + device const half* x = in + (size_t)row*dim; + device half* o = y + (size_t)row*dim; + float partial=0; for (uint i=t;i0; s>>=1){ if (t0; s>>=1){ if (t0; s>>=1){ if (t0; s>>=1){ if (t0; s>>=1){ if (t=total) return; + uint rem=i, src=base; + for (int d=int(ndim)-1; d>=0; --d){ uint c=rem%oshape[d]; rem/=oshape[d]; src+=c*sstride[d]; } + out[i]=in[src]; +} + +kernel void transpose2d_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& R [[buffer(2)]], constant uint& C [[buffer(3)]], + uint3 tgid [[threadgroup_position_in_grid]], + uint3 tp3 [[thread_position_in_threadgroup]]) { + uint2 tp = tp3.xy; + threadgroup half tile[32][33]; + uint b = tgid.z; + uint c0 = tgid.x * 32, r0 = tgid.y * 32; + device const half* src = in + (size_t)b * R * C; + device half* dst = y + (size_t)b * R * C; + for (uint dr = tp.y; dr < 32; dr += 8) { + uint r = r0 + dr, c = c0 + tp.x; + tile[dr][tp.x] = (r < R && c < C) ? src[(size_t)r * C + c] : (half)0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint dr = tp.y; dr < 32; dr += 8) { + uint c = c0 + dr, r = r0 + tp.x; + if (c < C && r < R) dst[(size_t)c * R + r] = tile[tp.x][dr]; + } +} + +kernel void strided_copy_rows_f16(device const half* in [[buffer(0)]], device half* out [[buffer(1)]], + constant uint* oshape [[buffer(2)]], constant uint* sstride [[buffer(3)]], + constant uint& ndim [[buffer(4)]], constant uint& rows [[buffer(5)]], constant uint& base [[buffer(6)]], + constant uint& inner4 [[buffer(7)]], + uint2 gid [[thread_position_in_grid]]) { + uint v = gid.x, row = gid.y; + if (v >= inner4 || row >= rows) return; + uint rem = row, src = base; + for (int d = int(ndim) - 1; d >= 0; --d) { uint c = rem % oshape[d]; rem /= oshape[d]; src += c * sstride[d]; } + device const half4* s4 = (device const half4*)(in + src); + device half4* d4 = (device half4*)(out + (size_t)row * inner4 * 4); + d4[v] = s4[v]; +} + +kernel void kv_append_i8(device const half* src [[buffer(0)]], device char* int8base [[buffer(1)]], + device float* scalebase [[buffer(2)]], constant uint& kv_heads [[buffer(3)]], + constant uint& hdim [[buffer(4)]], constant uint& current_len [[buffer(5)]], constant uint& group_size [[buffer(6)]], + uint gid [[thread_position_in_grid]]) { + uint num_groups = (hdim + group_size - 1)/group_size; + if (gid >= kv_heads*num_groups) return; + uint h = gid / num_groups, g = gid % num_groups; + uint gstart = g*group_size, gcount = min(group_size, hdim - gstart); + device const half* hs = src + (size_t)h*hdim + gstart; + float maxabs = 0; + for (uint k=0;k= remaining*per) return; + uint s=gid/per, hg=gid%per, h=hg/num_groups, g=hg%num_groups; + uint gstart=g*group_size, gcount=min(group_size, hdim-gstart); + uint i8s=kv_heads*hdim, scs=kv_heads*num_groups, srcseq=shift_src+s; + device const char* sp=int8base+(size_t)srcseq*i8s+(size_t)h*hdim+gstart; + device char* dp=scr_i8+(size_t)s*i8s+(size_t)h*hdim+gstart; + for (uint k=0;k= (remaining+1)*per) return; + uint s=gid/per, hg=gid%per, h=hg/num_groups, g=hg%num_groups; + uint gstart=g*group_size, gcount=min(group_size, hdim-gstart); + uint i8s=kv_heads*hdim, scs=kv_heads*num_groups, dstseq=keep_sink+s; + device char* dp=int8base+(size_t)dstseq*i8s+(size_t)h*hdim+gstart; + if (s < remaining) { + device const char* sp=scr_i8+(size_t)s*i8s+(size_t)h*hdim+gstart; + for (uint k=0;k= (remaining+M)*per) return; + uint s=gid/per, hg=gid%per, h=hg/num_groups, g=hg%num_groups; + uint gstart=g*group_size, gcount=min(group_size, hdim-gstart); + uint i8s=kv_heads*hdim, scs=kv_heads*num_groups, dstseq=keep_sink+s; + device char* dp=int8base+(size_t)dstseq*i8s+(size_t)h*hdim+gstart; + if (s < remaining) { + device const char* sp=scr_i8+(size_t)s*i8s+(size_t)h*hdim+gstart; + for (uint k=0;k=total) return; + uint rem=i, dst=base; + for (int d=int(ndim)-1; d>=0; --d){ uint c=rem%ishape[d]; rem/=ishape[d]; dst+=c*ostride[d]; } + out[dst]=in[i]; +} + +kernel void strided_scatter_rows_f16(device const half* in [[buffer(0)]], device half* out [[buffer(1)]], + constant uint* ishape [[buffer(2)]], constant uint* ostride [[buffer(3)]], + constant uint& ndim [[buffer(4)]], constant uint& rows [[buffer(5)]], constant uint& base [[buffer(6)]], + constant uint& inner4 [[buffer(7)]], + uint2 gid [[thread_position_in_grid]]) { + uint v = gid.x, row = gid.y; + if (v >= inner4 || row >= rows) return; + uint rem = row, dst = base; + for (int d = int(ndim) - 1; d >= 0; --d) { uint c = rem % ishape[d]; rem /= ishape[d]; dst += c * ostride[d]; } + device const half4* s4 = (device const half4*)(in + (size_t)row * inner4 * 4); + device half4* d4 = (device half4*)(out + dst); + d4[v] = s4[v]; +} + +kernel void binary_f32(device const float* a [[buffer(0)]], device const float* b [[buffer(1)]], + device float* y [[buffer(2)]], constant uint& n [[buffer(3)]], + constant int& op [[buffer(4)]], uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) { + float av=a[i0+j], bv=b[i0+j], r; + switch(op){ case 2: r=av-bv; break; case 3: r=av*bv; break; case 4: r=av/bv; break; + case 5: r=(av!=bv)?1.0f:0.0f; break; default: r=av+bv; } + if (op==1) r=clamp(r,-65500.0f,65500.0f); + y[i0+j]=r; + } +} +kernel void scalar_f32(device const float* in [[buffer(0)]], device float* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant int& op [[buffer(3)]], + constant float& p [[buffer(4)]], uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) y[i0+j]=scalar_apply(in[i0+j], p, op); +} +kernel void unary_f32(device const float* in [[buffer(0)]], device float* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant int& op [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) { + float x=in[i0+j], r; + if (op==0) r=gelu_tanh(x); + else if (op==1) r=precise::tanh(x); + else if (op==2) r=x/(1.0f+precise::exp(-x)); + else if (op==4) r=0.5f*x*(1.0f+erf_approx(x*0.70710678f)); + else if (op==5) r=1.0f/(1.0f+precise::exp(-x)); + else r=max(x,0.0f); + y[i0+j]=r; + } +} + +kernel void clamp_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant float& lo [[buffer(3)]], + constant float& hi [[buffer(4)]], uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) y[i0+j]=(half)clamp((float)in[i0+j], lo, hi); +} + +kernel void clamp_f32(device const float* in [[buffer(0)]], device float* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant float& lo [[buffer(3)]], + constant float& hi [[buffer(4)]], uint i [[thread_position_in_grid]]) { + uint i0 = i * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + for (uint j = 0; j < cnt; ++j) y[i0+j]=clamp(in[i0+j], lo, hi); +} +kernel void glu_f16(device const half* x [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& split [[buffer(2)]], constant uint& inner [[buffer(3)]], + constant uint& n [[buffer(4)]], uint i [[thread_position_in_grid]]) { + if (i>=n) return; + uint per_outer = split*inner; + uint o = i/per_outer, r = i%per_outer; + size_t base = (size_t)o*2u*per_outer; + float a = (float)x[base + r]; + float g = (float)x[base + per_outer + r]; + y[i] = (half)(a / (1.0f + precise::exp(-g))); +} + +kernel void layer_norm_f16(device const half* in [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* b [[buffer(2)]], device half* y [[buffer(3)]], + constant uint& dim [[buffer(4)]], constant float& eps [[buffer(5)]], + constant uint& has_bias [[buffer(6)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint nt [[threads_per_threadgroup]], threadgroup float* red [[threadgroup(0)]]) { + device const half* x = in + (size_t)row*dim; + device half* o = y + (size_t)row*dim; + float s=0, ss=0; + for (uint i=t;i0; st>>=1){ if (t0; s>>=1){ if (t0; s>>=1){ if (t0.0f ? 1.0f/red[0] : 0.0f; + for (uint i=t;i=Lout) return; + int center = (int)(gid.x*stride); + float acc=0; + device const half* wr = w + (size_t)gid.y*Cin*3u; + for (uint c=0;c0) ? (float)xc[center-1] : 0.0f; + float x1 = (float)xc[center]; + float x2 = (center+1<(int)L) ? (float)xc[center+1] : 0.0f; + acc = fma(x0,(float)wr[c*3u],acc); + acc = fma(x1,(float)wr[c*3u+1u],acc); + acc = fma(x2,(float)wr[c*3u+2u],acc); + } + y[(size_t)gid.y*Lout + gid.x] = (half)acc; +} + +struct AttnFU { uint T, S, HQ, HKV, D, DV; float scale; uint causal, pos_off, window; float logit_cap; uint mask_mode; }; + +static inline float attn_apply_mask(float score, device const half* mask, constant AttnFU& U, + uint b, uint h, uint t, uint i) { + if (U.mask_mode != 0u) { + size_t mi = (U.mask_mode>=3u) ? (((size_t)b*U.HQ + h)*U.T + t)*U.S + i + : ((size_t)b*U.T + t)*U.S + i; + float mv = (float)mask[mi]; + if (U.mask_mode==1u || U.mask_mode==3u) { score = isfinite(mv) ? score+mv : -INFINITY; } + else if (mv == 0.0f) score = -INFINITY; + } + if (U.logit_cap > 0.0f && score > -INFINITY) score = U.logit_cap * precise::tanh(score / U.logit_cap); + return score; +} + +kernel void attn_f16(device const half* q [[buffer(0)]], device const half* k [[buffer(1)]], + device const half* v [[buffer(2)]], device half* o [[buffer(3)]], + device const half* mask [[buffer(4)]], constant AttnFU& U [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], uint l [[thread_index_in_threadgroup]]) { + const uint QR = 4u; + uint t0 = tg.x * QR, h = tg.y, b = tg.z; + uint kvh = h / (U.HQ / U.HKV); + device const half* kb = k + (size_t)b*U.S*U.HKV*U.D + (size_t)kvh*U.D; + device const half* vb = v + (size_t)b*U.S*U.HKV*U.DV + (size_t)kvh*U.DV; + uint d4 = U.D/4u, dv4 = U.DV/4u; + threadgroup half4 qs[4u*32u]; + uint nq = min(QR, U.T - t0); + for (uint r=0;r0u) { uint aq0 = U.pos_off + t0; kstart = (aq0>U.window) ? aq0-U.window : 0u; } + size_t kstride = (size_t)U.HKV*U.D, vstride = (size_t)U.HKV*U.DV; + for (uint kv0=kstart; kv01u){ float4 q1 = float4(qs[32u+d]); + p.y = fma(kv4.x,q1.x,p.y); p.y = fma(kv4.y,q1.y,p.y); p.y = fma(kv4.z,q1.z,p.y); p.y = fma(kv4.w,q1.w,p.y); } + if (nq>2u){ float4 q2 = float4(qs[64u+d]); + p.z = fma(kv4.x,q2.x,p.z); p.z = fma(kv4.y,q2.y,p.z); p.z = fma(kv4.z,q2.z,p.z); p.z = fma(kv4.w,q2.w,p.z); } + if (nq>3u){ float4 q3 = float4(qs[96u+d]); + p.w = fma(kv4.x,q3.x,p.w); p.w = fma(kv4.y,q3.y,p.w); p.w = fma(kv4.z,q3.z,p.w); p.w = fma(kv4.w,q3.w,p.w); } + } + sc = p * U.scale; + for (uint r=0;r<4u;++r) { + if (r >= nq) { sc[r] = -INFINITY; continue; } + uint t = t0 + r, aq = U.pos_off + t; + if (U.causal && i > aq) { sc[r] = -INFINITY; continue; } + if (U.window>0u && i < aq && (aq - i) > U.window) { sc[r] = -INFINITY; continue; } + sc[r] = attn_apply_mask(sc[r], mask, U, b, h, t, i); + } + } + float4 wgt = float4(0.0f); + bool anyrow = false; + for (uint r=0;r<4u;++r) { + float bm = simd_max(sc[r]); + if (bm > -INFINITY) { + anyrow = true; + if (bm > m[r]) { + float corr = exp(m[r] - bm); + sum[r] *= corr; acc[r][0] *= corr; acc[r][1] *= corr; + m[r] = bm; + } + wgt[r] = (sc[r] > -INFINITY) ? exp(sc[r] - m[r]) : 0.0f; + sum[r] += simd_sum(wgt[r]); + } + } + if (!anyrow) continue; + uint nk = min(32u, kend_all - kv0); + uint half_id = l >> 4u, dl = l & 15u; + for (uint j2=0; j20.0f ? 1.0f/sum[r] : 0.0f; + device half4* orow = (device half4*)(o + ((size_t)b*U.T + t0 + r)*U.HQ*U.DV + (size_t)h*U.DV); + uint ci = 0u; + for (uint dd=dl; dd0u) { uint aq0 = U.pos_off + t0; kstart = (aq0>U.window) ? aq0-U.window : 0u; } + size_t kstride = (size_t)U.HKV*U.D, vstride = (size_t)U.HKV*U.DV; + for (uint kv0=kstart; kv0= nq) { sc[r] = -INFINITY; continue; } + sc[r] = p[r] * U.scale; + uint t = t0 + r, aq = U.pos_off + t; + if (U.causal && i > aq) { sc[r] = -INFINITY; continue; } + if (U.window>0u && i < aq && (aq - i) > U.window) { sc[r] = -INFINITY; continue; } + sc[r] = attn_apply_mask(sc[r], mask, U, b, h, t, i); + } + } + float wgt[8]; + bool anyrow = false; + for (uint r=0;r<8u;++r) { + float bm = simd_max(sc[r]); + if (bm > -INFINITY) { + anyrow = true; + if (bm > m[r]) { + float corr = exp(m[r] - bm); + sum[r] *= corr; acc[r] *= corr; + m[r] = bm; + } + wgt[r] = (sc[r] > -INFINITY) ? exp(sc[r] - m[r]) : 0.0f; + sum[r] += simd_sum(wgt[r]); + } else wgt[r] = 0.0f; + } + if (!anyrow) continue; + uint nk = min(32u, kend_all - kv0); + uint half_id = l >> 4u, dl = l & 15u; + for (uint j2=0; j20.0f ? 1.0f/sum[r] : 0.0f; + device half4* orow = (device half4*)(o + ((size_t)b*U.T + t0 + r)*U.HQ*U.DV + (size_t)h*U.DV); + float4 vv; + vv[0] = acc[r][0] + simd_shuffle_xor(acc[r][0], (ushort)16); + vv[1] = acc[r][1] + simd_shuffle_xor(acc[r][1], (ushort)16); + vv[2] = acc[r][2] + simd_shuffle_xor(acc[r][2], (ushort)16); + vv[3] = acc[r][3] + simd_shuffle_xor(acc[r][3], (ushort)16); + if (l < 16u && dl < dv4) orow[dl] = half4(vv*inv); + } +} + +kernel void attn_flash_f16(device const half* q [[buffer(0)]], device const half* k [[buffer(1)]], + device const half* v [[buffer(2)]], device half* y [[buffer(3)]], + device const half* mask [[buffer(4)]], + constant uint& T [[buffer(5)]], constant uint& S [[buffer(6)]], + constant uint& H [[buffer(7)]], constant float& scale [[buffer(8)]], + constant uint& mask_mode [[buffer(9)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + constexpr uint BQ = 64, BK = 32, D = 64, NT = 256; + threadgroup half KVs[2 * BK * D]; + threadgroup half* Ks = KVs; + threadgroup half* Vs = KVs + BK * D; + threadgroup float Ss[8][8 * BK]; + threadgroup half Ps[8][8 * BK]; + threadgroup float Dg[8][64]; + threadgroup float Ms[8][8]; + threadgroup float Ls[8][8]; + threadgroup uint Rf[8]; + + uint qb = tgid.x, h = tgid.y; + uint q0 = qb * BQ; + uint row0 = q0 + sgid * 8; + bool interior = q0 + BQ <= T; + + for (uint x = tid; x < 8 * 64; x += NT) Dg[x / 64][x % 64] = 0; + if (tid < 64) { Ms[tid / 8][tid % 8] = -INFINITY; Ls[tid / 8][tid % 8] = 0; } + + simdgroup_half8x8 qf[8]; + if (interior) { + for (uint d = 0; d < 8; ++d) + simdgroup_load(qf[d], q + ((size_t)row0 * H + h) * D + d * 8, (size_t)H * D); + } else { + threadgroup half* Qe = KVs; + for (uint x = tid; x < BQ * D; x += NT) { + uint r = x / D, d = x % D; + uint t = q0 + r; + Qe[x] = (t < T) ? q[((size_t)t * H + h) * D + d] : (half)0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint d = 0; d < 8; ++d) + simdgroup_load(qf[d], Qe + (size_t)sgid * 8 * D + d * 8, D); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 of[8]; + for (uint d = 0; d < 8; ++d) of[d] = make_filled_simdgroup_matrix(0.0f); + + for (uint kb = 0; kb < S; kb += BK) { + for (uint x = tid; x < BK * D / 4; x += NT) { + uint r = (x * 4) / D, d = (x * 4) % D; + uint sidx = kb + r; + half4 kv4 = half4(0.0h), vv4 = half4(0.0h); + if (sidx < S) { + kv4 = *(device const half4*)(k + ((size_t)sidx * H + h) * D + d); + vv4 = *(device const half4*)(v + ((size_t)sidx * H + h) * D + d); + } + *(threadgroup half4*)(Ks + r * D + d) = kv4; + *(threadgroup half4*)(Vs + r * D + d) = vv4; + } + if (lane == 0) Rf[sgid] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint j = 0; j < 4; ++j) { + simdgroup_float8x8 sf = make_filled_simdgroup_matrix(0.0f); + for (uint d = 0; d < 8; ++d) { + simdgroup_half8x8 kt; + simdgroup_load(kt, Ks + (size_t)j * 8 * D + d * 8, D, 0, true); + simdgroup_multiply_accumulate(sf, qf[d], kt, sf); + } + simdgroup_store(sf, Ss[sgid] + j * 8, BK); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + uint r = lane / 4, c0 = (lane % 4) * 8; + uint t = row0 + r; + float sv[8]; + float mloc = -INFINITY; + for (uint c = 0; c < 8; ++c) { + uint sidx = kb + c0 + c; + float val = Ss[sgid][r * BK + c0 + c] * scale; + bool dead = sidx >= S || t >= T; + if (!dead && mask_mode == 2u && mask[(size_t)t * S + sidx] == (half)0.0h) dead = true; + sv[c] = dead ? -INFINITY : val; + mloc = max(mloc, sv[c]); + } + mloc = max(mloc, simd_shuffle_xor(mloc, 1)); + mloc = max(mloc, simd_shuffle_xor(mloc, 2)); + float mold = Ms[sgid][r]; + float mnew = max(mold, mloc); + float sloc = 0; + for (uint c = 0; c < 8; ++c) { + float pv = (mnew == -INFINITY || sv[c] == -INFINITY) ? 0.0f : fast::exp(sv[c] - mnew); + Ps[sgid][r * BK + c0 + c] = (half)pv; + sloc += pv; + } + sloc += simd_shuffle_xor(sloc, 1); + sloc += simd_shuffle_xor(sloc, 2); + float alpha = (mold == -INFINITY) ? 0.0f : fast::exp(mold - mnew); + bool first = mold == -INFINITY; + if ((lane % 4) == 0) { + Ms[sgid][r] = mnew; + Ls[sgid][r] = Ls[sgid][r] * alpha + sloc; + if (mnew != mold && !first) Rf[sgid] = 1; + Dg[sgid][r * 8 + r] = alpha; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + bool rescale = Rf[sgid] != 0 || kb == 0; + simdgroup_half8x8 pf[4]; + for (uint j = 0; j < 4; ++j) + simdgroup_load(pf[j], Ps[sgid] + j * 8, BK); + if (rescale) { + simdgroup_float8x8 dgf; + simdgroup_load(dgf, Dg[sgid], 8); + for (uint d = 0; d < 8; ++d) { + simdgroup_float8x8 acc; + simdgroup_multiply(acc, dgf, of[d]); + for (uint j = 0; j < 4; ++j) { + simdgroup_half8x8 vf; + simdgroup_load(vf, Vs + (size_t)j * 8 * D + d * 8, D); + simdgroup_multiply_accumulate(acc, pf[j], vf, acc); + } + of[d] = acc; + } + } else { + for (uint d = 0; d < 8; ++d) { + for (uint j = 0; j < 4; ++j) { + simdgroup_half8x8 vf; + simdgroup_load(vf, Vs + (size_t)j * 8 * D + d * 8, D); + simdgroup_multiply_accumulate(of[d], pf[j], vf, of[d]); + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + { + uint r = lane / 4; + if ((lane % 4) == 0) { + float l = Ls[sgid][r]; + Dg[sgid][r * 8 + r] = l > 0.0f ? 1.0f / l : 0.0f; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_float8x8 dgf; + simdgroup_load(dgf, Dg[sgid], 8); + if (interior) { + for (uint d = 0; d < 8; ++d) { + simdgroup_multiply(of[d], dgf, of[d]); + simdgroup_store(of[d], Ss[sgid], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = lane; e < 64; e += 32) { + uint r2 = e / 8, cc = e % 8; + y[((size_t)(row0 + r2) * H + h) * D + d * 8 + cc] = (half)Ss[sgid][e]; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + } else { + threadgroup half* Oe = KVs; + for (uint d = 0; d < 8; ++d) { + simdgroup_multiply(of[d], dgf, of[d]); + simdgroup_store(of[d], Ss[sgid], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = lane; e < 64; e += 32) { + uint r2 = e / 8, cc = e % 8; + Oe[(size_t)sgid * 8 * D + r2 * D + d * 8 + cc] = (half)Ss[sgid][e]; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint x = tid; x < BQ * D; x += NT) { + uint r2 = x / D, d = x % D; + uint t = q0 + r2; + if (t < T) y[((size_t)t * H + h) * D + d] = Oe[x]; + } + } + } +} + +kernel void attn_maskfill_f16(device const half* mask [[buffer(0)]], device half* out [[buffer(1)]], + constant uint& ts [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) { + uint i = gid.x, h = gid.y; + if (i >= ts) return; + out[(size_t)h*ts + i] = (mask[i] == (half)0.0f) ? (half)(-30000.0f) : (half)0.0f; +} + +kernel void gemm_batch_f32a(device const float* a [[buffer(0)]], device const half* b [[buffer(1)]], + device uchar* y [[buffer(2)]], + constant uint& M [[buffer(3)]], constant uint& K [[buffer(4)]], + constant uint& N [[buffer(5)]], constant uint& f32out [[buffer(6)]], + uint2 gid [[thread_position_in_grid]]) { + uint mn = gid.x, bidx = gid.y; + if (mn >= M*N) return; + uint m = mn / N, nn = mn % N; + device const float* ar = a + (size_t)bidx*M*K + (size_t)m*K; + device const half* br = b + (size_t)bidx*K*N + nn; + float acc = 0; + for (uint k = 0; k < K; ++k) acc = fma(ar[k], (float)br[(size_t)k*N], acc); + size_t oi = (size_t)bidx*M*N + mn; + if (f32out != 0u) ((device float*)y)[oi] = acc; + else ((device half*)y)[oi] = (half)acc; +} + +kernel void gemm_batch_f16(device const half* a [[buffer(0)]], device const half* b [[buffer(1)]], + device uchar* y [[buffer(2)]], + constant uint& M [[buffer(3)]], constant uint& K [[buffer(4)]], + constant uint& N [[buffer(5)]], constant uint& f32out [[buffer(6)]], + uint2 gid [[thread_position_in_grid]]) { + uint mn = gid.x, bidx = gid.y; + if (mn >= M*N) return; + uint m = mn / N, nn = mn % N; + device const half* ar = a + (size_t)bidx*M*K + (size_t)m*K; + device const half* br = b + (size_t)bidx*K*N + nn; + float acc = 0; + for (uint k = 0; k < K; ++k) acc = fma((float)ar[k], (float)br[(size_t)k*N], acc); + size_t oi = (size_t)bidx*M*N + mn; + if (f32out != 0u) ((device float*)y)[oi] = acc; + else ((device half*)y)[oi] = (half)acc; +} + +kernel void conv1d_dw_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device half* y [[buffer(2)]], + constant uint& L [[buffer(3)]], constant uint& Lout [[buffer(4)]], + constant uint& Kk [[buffer(5)]], constant uint& stride [[buffer(6)]], + uint2 gid [[thread_position_in_grid]]) { + uint l = gid.x, c = gid.y; + if (l >= Lout) return; + device const half* xc = x + (size_t)c*L + (size_t)l*stride; + device const half* wc = w + (size_t)c*Kk; + float acc = 0; + for (uint k = 0; k < Kk; ++k) acc = fma((float)xc[k], (float)wc[k], acc); + y[(size_t)c*Lout + l] = (half)acc; +} + +kernel void reduce_axis_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& axis_size [[buffer(2)]], constant uint& inner [[buffer(3)]], + constant uint& n [[buffer(4)]], constant int& op [[buffer(5)]], + uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint outer = i / inner, in_i = i % inner; + size_t base = (size_t)outer * axis_size * inner + in_i; + float r; + if (op <= 2) { + float sum = 0; + for (uint a = 0; a < axis_size; ++a) sum += (float)in[base + (size_t)a*inner]; + float mean = sum / (float)max(axis_size, 1u); + if (op == 0) r = sum; + else if (op == 1) r = mean; + else { + float var = 0; + for (uint a = 0; a < axis_size; ++a) { float d = (float)in[base + (size_t)a*inner] - mean; var += d*d; } + r = var / (float)max(axis_size, 1u); + } + } else { + r = axis_size ? (float)in[base] : 0.0f; + for (uint a = 1; a < axis_size; ++a) { + float v = (float)in[base + (size_t)a*inner]; + r = (op == 3) ? min(r, v) : max(r, v); + } + } + y[i] = (half)r; +} + +kernel void reduce_axis_f32(device const float* in [[buffer(0)]], device float* y [[buffer(1)]], + constant uint& axis_size [[buffer(2)]], constant uint& inner [[buffer(3)]], + constant uint& n [[buffer(4)]], constant int& op [[buffer(5)]], + uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint outer = i / inner, in_i = i % inner; + size_t base = (size_t)outer * axis_size * inner + in_i; + float r; + if (op <= 2) { + float sum = 0; + for (uint a = 0; a < axis_size; ++a) sum += in[base + (size_t)a*inner]; + float mean = sum / (float)max(axis_size, 1u); + if (op == 0) r = sum; + else if (op == 1) r = mean; + else { + float var = 0; + for (uint a = 0; a < axis_size; ++a) { float d = in[base + (size_t)a*inner] - mean; var += d*d; } + r = var / (float)max(axis_size, 1u); + } + } else { + r = axis_size ? in[base] : 0.0f; + for (uint a = 1; a < axis_size; ++a) { + float v = in[base + (size_t)a*inner]; + r = (op == 3) ? min(r, v) : max(r, v); + } + } + y[i] = r; +} + +kernel void cumsum_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& axis_size [[buffer(2)]], constant uint& inner [[buffer(3)]], + constant uint& n [[buffer(4)]], uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint outer = i / inner, in_i = i % inner; + size_t base = (size_t)outer * axis_size * inner + in_i; + float run = 0; + for (uint a = 0; a < axis_size; ++a) { + run += (float)in[base + (size_t)a*inner]; + y[base + (size_t)a*inner] = (half)run; + } +} + +kernel void cumsum_f32(device const float* in [[buffer(0)]], device float* y [[buffer(1)]], + constant uint& axis_size [[buffer(2)]], constant uint& inner [[buffer(3)]], + constant uint& n [[buffer(4)]], uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint outer = i / inner, in_i = i % inner; + size_t base = (size_t)outer * axis_size * inner + in_i; + float run = 0; + for (uint a = 0; a < axis_size; ++a) { + run += in[base + (size_t)a*inner]; + y[base + (size_t)a*inner] = run; + } +} + +kernel void concat2_f16(device const half* a [[buffer(0)]], device const half* b [[buffer(1)]], + device half* y [[buffer(2)]], constant uint& a_axis [[buffer(3)]], + constant uint& b_axis [[buffer(4)]], constant uint& inner [[buffer(5)]], + constant uint& n [[buffer(6)]], constant uint& a_outer [[buffer(7)]], + constant uint& b_outer [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint in_i = i % inner, rest = i / inner; + uint ab = a_axis + b_axis; + uint ax = rest % ab, outer = rest / ab; + y[i] = (ax < a_axis) ? a[((size_t)(outer % a_outer)*a_axis + ax)*inner + in_i] + : b[((size_t)(outer % b_outer)*b_axis + (ax - a_axis))*inner + in_i]; +} + +kernel void gather_f32idx_f16(device const half* table [[buffer(0)]], device const float* rows [[buffer(1)]], + device half* out [[buffer(2)]], constant uint& D [[buffer(3)]], + constant uint& n [[buffer(4)]], uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint row = i / D, d = i % D; + out[i] = table[(size_t)(uint)rows[row]*D + d]; +} + +kernel void rope_full_f16(device const half* x [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& S [[buffer(2)]], constant uint& H [[buffer(3)]], + constant uint& D [[buffer(4)]], constant uint& rot [[buffer(5)]], + constant uint& pos0 [[buffer(6)]], constant float& theta [[buffer(7)]], + constant uint& gptj [[buffer(8)]], + uint2 gid [[thread_position_in_grid]]) { + uint idx = gid.x, tok = gid.y; + uint hr = rot / 2; + if (idx >= hr) { + uint d = rot + (idx - hr); + if (d >= D) return; + size_t off = (size_t)tok * D + d; + y[off] = x[off]; + return; + } + uint seq = (tok / H) % S; + float freq = 1.0f / precise::pow(theta, (2.0f * idx) / (float)rot); + float angle = (float)(pos0 + seq) * freq; + float c = (float)(half)precise::cos(angle); + float sn = (float)(half)precise::sin(angle); + uint d0 = gptj ? 2u * idx : idx; + uint d1 = gptj ? 2u * idx + 1u : idx + hr; + size_t off0 = (size_t)tok * D + d0, off1 = (size_t)tok * D + d1; + float v0 = (float)x[off0], v1 = (float)x[off1]; + y[off0] = (half)(v0*c - v1*sn); + y[off1] = (half)(v1*c + v0*sn); +} + +kernel void maxpool1d_f16(device const half* x [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& L [[buffer(2)]], constant uint& Lout [[buffer(3)]], + constant uint& K [[buffer(4)]], constant uint& stride [[buffer(5)]], + uint2 gid [[thread_position_in_grid]]) { + uint lo = gid.x, nc = gid.y; + if (lo >= Lout) return; + size_t base = (size_t)nc * L + (size_t)lo * stride; + float r = -INFINITY; + for (uint k = 0; k < K && lo * stride + k < L; ++k) r = max(r, (float)x[base + k]); + y[(size_t)nc * Lout + lo] = (half)r; +} + +kernel void bilinear_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& sh [[buffer(2)]], constant uint& sw [[buffer(3)]], + constant uint& dh [[buffer(4)]], constant uint& dw [[buffer(5)]], + constant uint& E [[buffer(6)]], constant uint& align [[buffer(7)]], + uint2 gid [[thread_position_in_grid]]) { + uint e = gid.x, pix = gid.y; + if (e >= E || pix >= dh * dw) return; + uint dy = pix / dw, dx = pix % dw; + float syf, sxf; + if (align) { + float sc_h = (sh > 1 && dh > 1) ? (float)(sh - 1) / (float)(dh - 1) : 0.0f; + float sc_w = (sw > 1 && dw > 1) ? (float)(sw - 1) / (float)(dw - 1) : 0.0f; + syf = dy * sc_h; sxf = dx * sc_w; + } else { + float sc_h = dh ? (float)sh / (float)dh : 0.0f; + float sc_w = dw ? (float)sw / (float)dw : 0.0f; + syf = clamp((dy + 0.5f) * sc_h - 0.5f, 0.0f, (float)sh - 1.0f); + sxf = clamp((dx + 0.5f) * sc_w - 0.5f, 0.0f, (float)sw - 1.0f); + } + int y0 = (int)floor(syf), x0 = (int)floor(sxf); + int y1 = min(y0 + 1, (int)sh - 1), x1 = min(x0 + 1, (int)sw - 1); + float fy = syf - y0, fx = sxf - x0; + float v00 = (float)in[((size_t)y0*sw + x0)*E + e], v01 = (float)in[((size_t)y0*sw + x1)*E + e]; + float v10 = (float)in[((size_t)y1*sw + x0)*E + e], v11 = (float)in[((size_t)y1*sw + x1)*E + e]; + y[(size_t)pix*E + e] = (half)(v00*(1-fx)*(1-fy) + v01*fx*(1-fy) + v10*(1-fx)*fy + v11*fx*fy); +} + +kernel void conv1d_gen_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], device half* y [[buffer(3)]], + constant uint& Cin [[buffer(4)]], constant uint& L [[buffer(5)]], + constant uint& Cout [[buffer(6)]], constant uint& Lout [[buffer(7)]], + constant uint& K [[buffer(8)]], constant uint& stride [[buffer(9)]], + constant uint& has_bias [[buffer(10)]], constant uint& w_ck_co [[buffer(11)]], + uint3 gid [[thread_position_in_grid]]) { + uint lo = gid.x, co = gid.y, n = gid.z; + if (lo >= Lout || co >= Cout) return; + float acc = has_bias ? (float)bias[co] : 0.0f; + for (uint ci = 0; ci < Cin; ++ci) { + device const half* xr = x + ((size_t)n*Cin + ci)*L + (size_t)lo*stride; + for (uint k = 0; k < K; ++k) { + float wv = w_ck_co ? (float)w[((size_t)ci*K + k)*Cout + co] + : (float)w[((size_t)co*Cin + ci)*K + k]; + acc = fma((float)xr[k], wv, acc); + } + } + y[((size_t)n*Cout + co)*Lout + lo] = (half)acc; +} + +kernel void conv1d_nlc_dw_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], device half* y [[buffer(3)]], + constant uint& L [[buffer(4)]], constant uint& C [[buffer(5)]], + constant uint& K [[buffer(6)]], constant uint& dil [[buffer(7)]], + constant uint& pad [[buffer(8)]], constant uint& has_bias [[buffer(9)]], + uint3 gid [[thread_position_in_grid]]) { + uint c = gid.x, l = gid.y, n = gid.z; + if (c >= C || l >= L) return; + float acc = has_bias ? (float)bias[c] : 0.0f; + for (uint k = 0; k < K; ++k) { + long p = (long)l - (long)pad + (long)k * dil; + if (p >= 0 && p < (long)L) + acc = fma((float)x[((size_t)n*L + p)*C + c], (float)w[(size_t)c*K + k], acc); + } + y[((size_t)n*L + l)*C + c] = (half)acc; +} + +kernel void conv2d_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* bias [[buffer(2)]], device half* y [[buffer(3)]], + constant uint& Cin [[buffer(4)]], constant uint& H [[buffer(5)]], + constant uint& W [[buffer(6)]], constant uint& Cout [[buffer(7)]], + constant uint& Ho [[buffer(8)]], constant uint& Wo [[buffer(9)]], + constant uint& K [[buffer(10)]], constant uint& stride [[buffer(11)]], + constant uint& pad [[buffer(12)]], constant uint& dw [[buffer(13)]], + constant uint& has_bias [[buffer(14)]], + uint3 gid [[thread_position_in_grid]]) { + uint wo0 = gid.x * 4, ho = gid.y, nco = gid.z; + if (wo0 >= Wo || ho >= Ho) return; + uint n = nco / Cout, co = nco % Cout; + float acc[4]; + float bv = has_bias ? (float)bias[co] : 0.0f; + for (uint q = 0; q < 4; ++q) acc[q] = bv; + uint ci0 = dw ? co : 0, ci1 = dw ? co + 1 : Cin; + for (uint ci = ci0; ci < ci1; ++ci) { + for (uint kh = 0; kh < K; ++kh) { + long h = (long)ho * stride - (long)pad + kh; + if (h < 0 || h >= (long)H) continue; + device const half* xr = x + (((size_t)n*Cin + ci)*H + h)*W; + for (uint kw = 0; kw < K; ++kw) { + float wv = dw ? (float)w[((size_t)co*K + kh)*K + kw] + : (float)w[(((size_t)co*Cin + ci)*K + kh)*K + kw]; + for (uint q = 0; q < 4; ++q) { + long ww = (long)(wo0 + q) * stride - (long)pad + kw; + if (ww >= 0 && ww < (long)W && wo0 + q < Wo) + acc[q] = fma((float)xr[ww], wv, acc[q]); + } + } + } + } + for (uint q = 0; q < 4 && wo0 + q < Wo; ++q) + y[(((size_t)n*Cout + co)*Ho + ho)*Wo + wo0 + q] = (half)acc[q]; +} + +kernel void batchnorm_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* b [[buffer(2)]], device const half* rm [[buffer(3)]], + device const half* rv [[buffer(4)]], device half* y [[buffer(5)]], + constant uint& C [[buffer(6)]], constant uint& inner [[buffer(7)]], + constant float& eps [[buffer(8)]], constant uint& n [[buffer(9)]], + uint i [[thread_position_in_grid]]) { + if (i >= n) return; + uint c = (i / inner) % C; + float inv = 1.0f / precise::sqrt((float)rv[c] + eps); + y[i] = (half)(((float)x[i] - (float)rm[c]) * inv * (float)w[c] + (float)b[c]); +} + +kernel void groupnorm_f16(device const half* x [[buffer(0)]], device const half* w [[buffer(1)]], + device const half* b [[buffer(2)]], device half* y [[buffer(3)]], + constant uint& cpg [[buffer(4)]], constant uint& S [[buffer(5)]], + constant uint& C [[buffer(6)]], constant float& eps [[buffer(7)]], + uint2 tgp [[threadgroup_position_in_grid]], + uint2 tp [[thread_position_in_threadgroup]], uint2 ntp [[threads_per_threadgroup]], + threadgroup float* red [[threadgroup(0)]]) { + uint g = tgp.x, n = tgp.y, t = tp.x, nt = ntp.x; + size_t base = ((size_t)n*C + (size_t)g*cpg) * S; + uint count = cpg * S; + float sum = 0, sq = 0; + for (uint i = t; i < count; i += nt) { float v = (float)x[base + i]; sum += v; sq += v*v; } + red[t] = sum; threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s2 = nt/2; s2 > 0; s2 >>= 1) { if (t < s2) red[t] += red[t+s2]; threadgroup_barrier(mem_flags::mem_threadgroup); } + sum = red[0]; threadgroup_barrier(mem_flags::mem_threadgroup); + red[t] = sq; threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s2 = nt/2; s2 > 0; s2 >>= 1) { if (t < s2) red[t] += red[t+s2]; threadgroup_barrier(mem_flags::mem_threadgroup); } + sq = red[0]; + float mean = sum / count; + float inv = 1.0f / precise::sqrt(sq / count - mean*mean + eps); + for (uint i = t; i < count; i += nt) { + uint c = g*cpg + i / S; + y[base + i] = (half)(((float)x[base + i] - mean) * inv * (float)w[c] + (float)b[c]); + } +} + +kernel void bias_add_rows_f16(device half* y [[buffer(0)]], device const half* bias [[buffer(1)]], + constant uint& C [[buffer(2)]], constant uint& n [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + if (i < n) y[i] = (half)((float)y[i] + (float)bias[i % C]); +} + +struct EwStep { int kind; int code; float p0; float p1; }; + +kernel void elemwise_chain_f16(device const uchar* in [[buffer(0)]], device uchar* y [[buffer(1)]], + constant EwStep* steps [[buffer(2)]], constant uint& nsteps [[buffer(3)]], + device const uchar* s0 [[buffer(4)]], device const uchar* s1 [[buffer(5)]], + device const uchar* s2 [[buffer(6)]], constant uint& n [[buffer(7)]], + constant uint& flags [[buffer(8)]], + constant uint& inner [[buffer(9)]], + uint gid [[thread_position_in_grid]]) { + uint i0 = gid * 4; + if (i0 >= n) return; + uint cnt = min(4u, n - i0); + float xv[4]; + bool in32 = (flags & 1u) != 0; + for (uint q = 0; q < cnt; ++q) + xv[q] = in32 ? ((device const float*)in)[i0+q] : (float)((device const half*)in)[i0+q]; + bool prec32 = in32; + for (uint s = 0; s < nsteps; ++s) { + EwStep st = steps[s]; + if (st.kind == 4) { prec32 = st.code != 0; continue; } + if (st.kind == 2) { + uint slot = (st.code >> 6) & 3u; + int op = st.code & 15; + bool rhs = (st.code & 16) != 0; + bool sf32 = (st.code & 32) != 0; + device const uchar* sp = slot == 0 ? s0 : slot == 1 ? s1 : s2; + uint bmode = (st.code >> 8) & 3u; + for (uint q = 0; q < cnt; ++q) { + uint si = i0 + q; + if (bmode == 1u) si = si / inner; + else if (bmode == 2u) si = si % inner; + float o = sf32 ? ((device const float*)sp)[si] : (float)((device const half*)sp)[si]; + float a = rhs ? o : xv[q], b = rhs ? xv[q] : o; + float r; + switch (op) { case 2: r=a-b; break; case 3: r=a*b; break; case 4: r=a/b; break; + case 5: r=(a!=b)?1.0f:0.0f; break; default: r=a+b; } + if (op == 1) r = clamp(r, -65500.0f, 65500.0f); + xv[q] = r; + } + } else if (st.kind == 0) { + for (uint q = 0; q < cnt; ++q) { + float x = xv[q]; + if (st.code == 0) x = gelu_tanh(x); + else if (st.code == 1) x = precise::tanh(x); + else if (st.code == 2) x = x/(1.0f+precise::exp(-x)); + else if (st.code == 4) x = 0.5f*x*(1.0f+erf_approx(x*0.70710678f)); + else if (st.code == 5) x = 1.0f/(1.0f+precise::exp(-x)); + else x = max(x, 0.0f); + xv[q] = x; + } + } else if (st.kind == 1) { + for (uint q = 0; q < cnt; ++q) xv[q] = scalar_apply(xv[q], st.p0, st.code); + } else { + for (uint q = 0; q < cnt; ++q) xv[q] = clamp(xv[q], st.p0, st.p1); + } + if (!prec32) for (uint q = 0; q < cnt; ++q) xv[q] = (float)(half)xv[q]; + } + if ((flags & 2u) != 0) { for (uint q = 0; q < cnt; ++q) ((device float*)y)[i0+q] = xv[q]; } + else { for (uint q = 0; q < cnt; ++q) ((device half*)y)[i0+q] = (half)xv[q]; } +} + +kernel void bcast_binary_rows_f16(device const half* a [[buffer(0)]], device const half* b [[buffer(1)]], + device half* out [[buffer(2)]], constant uint* oshape [[buffer(3)]], + constant uint* astride [[buffer(4)]], constant uint* bstride [[buffer(5)]], + constant uint& ndim [[buffer(6)]], constant uint& rows [[buffer(7)]], constant int& op [[buffer(8)]], + constant uint& inner [[buffer(9)]], constant uint& ainner [[buffer(10)]], constant uint& binner [[buffer(11)]], + uint2 gid [[thread_position_in_grid]]) { + uint v = gid.x, row = gid.y; + if (v >= inner || row >= rows) return; + uint rem = row, ai = 0, bi = 0; + for (int d = int(ndim) - 1; d >= 0; --d) { uint c = rem % oshape[d]; rem /= oshape[d]; ai += c * astride[d]; bi += c * bstride[d]; } + float av = (float)a[ai + v * ainner], bv = (float)b[bi + v * binner], r; + switch(op){ case 2:r=av-bv;break; case 3:r=av*bv;break; case 4:r=av/bv;break; + case 5:r=(av!=bv)?1.0f:0.0f;break; default:r=av+bv; } + if (op==1) r=clamp(r,-65500.0f,65500.0f); + out[(size_t)row * inner + v] = (half)r; +} + +kernel void bcast_binary_f16(device const half* a [[buffer(0)]], device const half* b [[buffer(1)]], + device half* out [[buffer(2)]], constant uint* oshape [[buffer(3)]], + constant uint* astride [[buffer(4)]], constant uint* bstride [[buffer(5)]], + constant uint& ndim [[buffer(6)]], constant uint& total [[buffer(7)]], constant int& op [[buffer(8)]], + uint i [[thread_position_in_grid]]) { + if (i>=total) return; + uint rem=i, ai=0, bi=0; + for (int d=int(ndim)-1; d>=0; --d){ uint c=rem%oshape[d]; rem/=oshape[d]; ai+=c*astride[d]; bi+=c*bstride[d]; } + float av=(float)a[ai], bv=(float)b[bi], r; + switch(op){ case 2:r=av-bv;break; case 3:r=av*bv;break; case 4:r=av/bv;break; + case 5:r=(av!=bv)?1.0f:0.0f;break; default:r=av+bv; } + if (op==1) r=clamp(r,-65500.0f,65500.0f); + out[i]=(half)r; +} + +kernel void attn_decode_i8( + device const half* q [[buffer(0)]], + device const half* knew [[buffer(1)]], + device const half* vnew [[buffer(2)]], + device const char* kc [[buffer(3)]], + device const char* vc [[buffer(4)]], + device const float* ks [[buffer(5)]], + device const float* vs [[buffer(6)]], + device half* out [[buffer(7)]], + constant uint& num_q_heads [[buffer(8)]], constant uint& num_kv_heads [[buffer(9)]], + constant uint& head_dim [[buffer(10)]], constant uint& v_hdim [[buffer(11)]], + constant uint& history_len [[buffer(12)]], constant float& scale [[buffer(13)]], + constant uint& kv_start [[buffer(14)]], constant uint& kv_end [[buffer(15)]], + device float* part_o [[buffer(16)]], device float* part_ml [[buffer(17)]], + constant uint& nwg [[buffer(18)]], + uint tg [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]], threadgroup float* smem [[threadgroup(0)]]) +{ + uint h = tg / nwg, w = tg % nwg; + uint kvh = h / (num_q_heads / num_kv_heads); + uint ngK = (head_dim + 31u)/32u, ngV = (v_hdim + 31u)/32u; + uint nsg = T / 32u; + device const half* qh = q + (size_t)h*head_dim; + + float qreg[16]; + { uint i=0; for (uint d=lane; d 0.0f ? 1.0f/gl : 0.0f; + for (uint d=t; d0; s2>>=1){ if (t0; s2>>=1){ if (t0; s2>>=1){ if (t 0.0f ? 1.0f/gl : 0.0f; + for (uint d=t; d 0.0f ? 1.0f/gl : 0.0f; + for (uint d=t; d 0u) ? sinkN : 0u; + uint Wsr = S + ringR; + uint rstart, nactive; + if (ringR > 0u && kv_end > Wsr) { rstart = kv_end - ringR; nactive = Wsr; } + else { rstart = S; nactive = kv_end; } + device const half* qh = q + ((size_t)m*num_q_heads + h)*head_dim; + threadgroup float* red = sc + maxsc; + if (nactive == 0u) { + for (uint d=t; d 0u) { + float lmax = -INFINITY; + for (uint j = t; j < nactive; j += T) { + uint k = (j < S) ? j : rstart + (j - S); + float dot = 0; + if (k < history_len) { + uint slot = (ringR > 0u && k >= Wsr) ? (S + ((k - S) % ringR)) : k; + device const char* kk = kc + ((size_t)slot*num_kv_heads + kvh)*head_dim; + device const float* kss = ks + ((size_t)slot*num_kv_heads + kvh)*ngK; + for (uint d=0; d0;s>>=1){ if(t0;s>>=1){ if(t 0 ? 1.0f/red[0] : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint d = t; d < v_hdim; d += T) { + float acc=0; + for (uint j=0; j= Wsr) ? (S + ((k - S) % ringR)) : k; + device const char* vvv = vc + ((size_t)slot*num_kv_heads+kvh)*v_hdim; + device const float* vss = vs + ((size_t)slot*num_kv_heads+kvh)*ngV; + vv = (float)vvv[d]*vss[d/32]; + } else { + device const half* vvv = vnew + ((size_t)(k-history_len)*num_kv_heads+kvh)*v_hdim; + vv = (float)vvv[d]; + } + acc += sc[j]*vv; + } + out[((size_t)m*num_q_heads + h)*v_hdim + d] = (half)(acc*inv); + } + } else { + uint TILE = maxsc; + float rmax = -INFINITY, rsum = 0.0f; + float acc[8]; + for (uint a=0;a<8;++a) acc[a]=0.0f; + for (uint tile=0; tile0;s>>=1){ if(t0;s>>=1){ if(t 0 ? 1.0f/rsum : 0.0f; + uint ai=0; + for (uint d=t; d> 5, lane = tl & 31u; + uint h = sg >> 1, hf = sg & 1u, kvh = 0u; + uint vd0 = hf*256u; + uint m0 = tgx*QB; + uint total_keys = history_len + new_len; + + uint pos_last = q_pos0 + m0 + (QB-1u); + uint kv_end_tile = min(total_keys, pos_last + 1u); + simdgroup_matrix O[NFH]; + + for (uint f=0; f(0.f); + if (hf==0u && lane < QB) { mrun[h*QB+lane] = -INFINITY; lrun[h*QB+lane]=0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k0=0; k0 C0=make_filled_simdgroup_matrix(0.f); + simdgroup_matrix C1=C0,C2=C0,C3=C0,C4=C0,C5=C0,C6=C0,C7=C0; + for (uint c=0; c>5, d=i&31u; + Qs[h*QB*LD + r*LD + d] = (m0+r A,B0,B1,B2,B3,B4,B5,B6,B7; + simdgroup_load(A, &Qs[h*QB*LD + kk], LD); + simdgroup_load(B0,&Ks[kk*LDK + 0u], LDK); + simdgroup_load(B1,&Ks[kk*LDK + 8u], LDK); + simdgroup_load(B2,&Ks[kk*LDK + 16u],LDK); + simdgroup_load(B3,&Ks[kk*LDK + 24u],LDK); + simdgroup_load(B4,&Ks[kk*LDK + 32u],LDK); + simdgroup_load(B5,&Ks[kk*LDK + 40u],LDK); + simdgroup_load(B6,&Ks[kk*LDK + 48u],LDK); + simdgroup_load(B7,&Ks[kk*LDK + 56u],LDK); + simdgroup_multiply_accumulate(C0,A,B0,C0); + simdgroup_multiply_accumulate(C1,A,B1,C1); + simdgroup_multiply_accumulate(C2,A,B2,C2); + simdgroup_multiply_accumulate(C3,A,B3,C3); + simdgroup_multiply_accumulate(C4,A,B4,C4); + simdgroup_multiply_accumulate(C5,A,B5,C5); + simdgroup_multiply_accumulate(C6,A,B6,C6); + simdgroup_multiply_accumulate(C7,A,B7,C7); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (hf==0u) { + simdgroup_store(C0,&Ss[h*QB*BK + 0u], BK); + simdgroup_store(C1,&Ss[h*QB*BK + 8u], BK); + simdgroup_store(C2,&Ss[h*QB*BK + 16u],BK); + simdgroup_store(C3,&Ss[h*QB*BK + 24u],BK); + simdgroup_store(C4,&Ss[h*QB*BK + 32u],BK); + simdgroup_store(C5,&Ss[h*QB*BK + 40u],BK); + simdgroup_store(C6,&Ss[h*QB*BK + 48u],BK); + simdgroup_store(C7,&Ss[h*QB*BK + 56u],BK); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (hf==0u && lane < QB) { + uint r=lane, qpos=q_pos0+m0+r; + threadgroup float* srow=&Ss[h*QB*BK + r*BK]; + float tmax=-INFINITY; + for (uint k=0;k-INFINITY)?exp(srow[k]-mnew):0.0f; + srow[k]=e; tsum+=e; + } + lrun[h*QB+r]=lrun[h*QB+r]*resc+tsum; + mrun[h*QB+r]=mnew; resc8[h*QB+r]=resc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (lane == 0u) { + int nr = 0; + for (uint r=0; r>3; Rsc[sg*QB*8u + e] *= resc8[h*QB+r]; } + simdgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_load(O[f], &Rsc[sg*QB*8u], 8u); + simdgroup_barrier(mem_flags::mem_threadgroup); + } + } + + if (hf==0u) for (uint i=lane; i>5, d=i&31u; + half val=(half)0; + if (k P,V0; + simdgroup_load(P, &Ps[h*QB*LDP + kk], LDP); + simdgroup_load(V0, &Vs[hf*BK*LD + kk*LD + fc], LD); + simdgroup_multiply_accumulate(O[f], P, V0, O[f]); + } + } + } + } + for (uint f=0; f>3, d=e&7u; + if (m0+r0?1.0f/lrun[h*QB+r]:0.0f; + out[((size_t)(m0+r)*num_q_heads + h)*VD + (vd0 + f*8u + d)] = (half)(Rsc[sg*QB*8u + e]*inv); + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } +} + +kernel void attn_prefill_mma_hd256( + device const half* q [[buffer(0)]], + device const half* knew [[buffer(1)]], + device const half* vnew [[buffer(2)]], + device const char* kc [[buffer(3)]], + device const char* vc [[buffer(4)]], + device const float* ks [[buffer(5)]], + device const float* vs [[buffer(6)]], + device half* out [[buffer(7)]], + constant uint& num_q_heads [[buffer(8)]], constant uint& num_kv_heads [[buffer(9)]], + constant uint& head_dim [[buffer(10)]], constant uint& v_hdim [[buffer(11)]], + constant uint& history_len [[buffer(12)]], constant float& scale [[buffer(13)]], + constant uint& q_pos0 [[buffer(14)]], constant uint& new_len [[buffer(15)]], + constant uint& Mtot [[buffer(16)]], + constant uint& sinkN [[buffer(17)]], constant uint& ringR [[buffer(18)]], + uint tgx [[threadgroup_position_in_grid]], uint tl [[thread_index_in_threadgroup]]) +{ + const uint NSG=8u, LD=40u, BK=64u, QB=8u, HD=256u, VD=256u, NFH=32u; + const uint LDK=72u, LDP=72u; + threadgroup float poolA[4096]; + threadgroup half poolB[4608]; + threadgroup float Rsc[NSG*QB*8u]; + threadgroup float mrun[8u*QB]; + threadgroup float lrun[8u*QB]; + threadgroup float resc8[8u*QB]; + threadgroup int need_rescale[NSG]; + + threadgroup half* Ks = (threadgroup half*)poolA; + threadgroup float* Ss = poolA; + threadgroup half* Vs = (threadgroup half*)poolA; + threadgroup half* Qs = poolB; + threadgroup half* Ps = poolB; + + uint sg = tl >> 5, lane = tl & 31u; + uint h = sg, kvh = 0u; + uint m0 = tgx*QB; + uint total_keys = history_len + new_len; + + uint S = (ringR > 0u) ? sinkN : 0u; + uint R = ringR, Wsr = S + R; + uint pos_first = q_pos0 + m0; + uint pos_last = q_pos0 + m0 + (QB-1u); + uint kv_end_tile = min(total_keys, pos_last + 1u); + bool windowed = (R > 0u && pos_first >= Wsr); + uint rstart = windowed ? (pos_first + 1u - R) : 0u; + uint nactive = windowed ? (S + (kv_end_tile - rstart)) : kv_end_tile; + simdgroup_matrix O[NFH]; + + for (uint f=0; f(0.f); + if (lane < QB) { mrun[h*QB+lane] = -INFINITY; lrun[h*QB+lane]=0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k0=0; k0 C0=make_filled_simdgroup_matrix(0.f); + simdgroup_matrix C1=C0,C2=C0,C3=C0,C4=C0,C5=C0,C6=C0,C7=C0; + for (uint c=0; c0u && kk>=Wsr)?(S+(kk-S)%R):kk; + device const char* kp=kc+((size_t)slot*num_kv_heads+kvh)*HD; + device const float* sp=ks+((size_t)slot*num_kv_heads+kvh)*(HD/32u); + val=(half)((float)kp[c+d]*sp[(c+d)/32u]); + } else { + device const half* kp=knew+((size_t)(kk-history_len)*num_kv_heads+kvh)*HD; + val=kp[c+d]; + } + } + Ks[d*LDK + k]=val; + } + for (uint i=lane; i>5, d=i&31u; + Qs[h*QB*LD + r*LD + d] = (m0+r A,B0,B1,B2,B3,B4,B5,B6,B7; + simdgroup_load(A, &Qs[h*QB*LD + kk], LD); + simdgroup_load(B0,&Ks[kk*LDK + 0u], LDK); + simdgroup_load(B1,&Ks[kk*LDK + 8u], LDK); + simdgroup_load(B2,&Ks[kk*LDK + 16u],LDK); + simdgroup_load(B3,&Ks[kk*LDK + 24u],LDK); + simdgroup_load(B4,&Ks[kk*LDK + 32u],LDK); + simdgroup_load(B5,&Ks[kk*LDK + 40u],LDK); + simdgroup_load(B6,&Ks[kk*LDK + 48u],LDK); + simdgroup_load(B7,&Ks[kk*LDK + 56u],LDK); + simdgroup_multiply_accumulate(C0,A,B0,C0); + simdgroup_multiply_accumulate(C1,A,B1,C1); + simdgroup_multiply_accumulate(C2,A,B2,C2); + simdgroup_multiply_accumulate(C3,A,B3,C3); + simdgroup_multiply_accumulate(C4,A,B4,C4); + simdgroup_multiply_accumulate(C5,A,B5,C5); + simdgroup_multiply_accumulate(C6,A,B6,C6); + simdgroup_multiply_accumulate(C7,A,B7,C7); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_store(C0,&Ss[h*QB*BK + 0u], BK); + simdgroup_store(C1,&Ss[h*QB*BK + 8u], BK); + simdgroup_store(C2,&Ss[h*QB*BK + 16u],BK); + simdgroup_store(C3,&Ss[h*QB*BK + 24u],BK); + simdgroup_store(C4,&Ss[h*QB*BK + 32u],BK); + simdgroup_store(C5,&Ss[h*QB*BK + 40u],BK); + simdgroup_store(C6,&Ss[h*QB*BK + 48u],BK); + simdgroup_store(C7,&Ss[h*QB*BK + 56u],BK); + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (lane < QB) { + uint r=lane, qpos=q_pos0+m0+r; + threadgroup float* srow=&Ss[h*QB*BK + r*BK]; + float tmax=-INFINITY; + for (uint k=0;kqpos); + float s = att ? srow[k]*scale : -INFINITY; + srow[k]=s; tmax=max(tmax,s); + } + float mo=mrun[h*QB+r], mnew=max(mo,tmax); + float resc=exp(mo-mnew), tsum=0.0f; + for (uint k=0;k-INFINITY)?exp(srow[k]-mnew):0.0f; + srow[k]=e; tsum+=e; + } + lrun[h*QB+r]=lrun[h*QB+r]*resc+tsum; + mrun[h*QB+r]=mnew; resc8[h*QB+r]=resc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (lane == 0u) { + int nr = 0; + for (uint r=0; r>3; Rsc[sg*QB*8u + e] *= resc8[h*QB+r]; } + simdgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_load(O[f], &Rsc[sg*QB*8u], 8u); + simdgroup_barrier(mem_flags::mem_threadgroup); + } + } + + for (uint i=lane; i>5, d=i&31u; + half val=(half)0; + if (k0u && kk>=Wsr)?(S+(kk-S)%R):kk; + device const char* vp=vc+((size_t)slot*num_kv_heads+kvh)*VD; + device const float* sp=vs+((size_t)slot*num_kv_heads+kvh)*(VD/32u); + val=(half)((float)vp[c+d]*sp[(c+d)/32u]); + } else { + device const half* vp=vnew+((size_t)(kk-history_len)*num_kv_heads+kvh)*VD; + val=vp[c+d]; + } + } + Vs[k*LD + d]=val; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint fc=0; fc<32u; fc+=8u){ + uint f = jc/8u + fc/8u; + for (uint kk=0; kk P,V0; + simdgroup_load(P, &Ps[h*QB*LDP + kk], LDP); + simdgroup_load(V0, &Vs[kk*LD + fc], LD); + simdgroup_multiply_accumulate(O[f], P, V0, O[f]); + } + } + } + } + for (uint f=0; f>3, d=e&7u; + if (m0+r0?1.0f/lrun[h*QB+r]:0.0f; + out[((size_t)(m0+r)*num_q_heads + h)*VD + (f*8u + d)] = (half)(Rsc[sg*QB*8u + e]*inv); + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } +} + +kernel void kv_append_i8_m(device const half* src [[buffer(0)]], device char* int8base [[buffer(1)]], + device float* scalebase [[buffer(2)]], constant uint& kv_heads [[buffer(3)]], + constant uint& hdim [[buffer(4)]], constant uint& current_len [[buffer(5)]], + constant uint& group_size [[buffer(6)]], constant uint& M [[buffer(7)]], + uint gid [[thread_position_in_grid]]) { + uint num_groups = (hdim + group_size - 1)/group_size, per = kv_heads*num_groups; + if (gid >= M*per) return; + uint i = gid / per, hg = gid % per, h = hg / num_groups, g = hg % num_groups; + uint gstart = g*group_size, gcount = min(group_size, hdim - gstart); + device const half* hs = src + (size_t)i*kv_heads*hdim + (size_t)h*hdim + gstart; + float maxabs = 0; + for (uint k=0;k= N) return; + float acc = 0; + if (tr != 0u) { + device const half* wr = w + (size_t)n * K; + for (uint k = lane; k < K; k += 32u) acc = fma((float)x[k], (float)wr[k], acc); + } else { + for (uint k = lane; k < K; k += 32u) acc = fma((float)x[k], (float)w[(size_t)k * N + n], acc); + } + acc = simd_sum(acc); + if (lane == 0) y[n] = (half)((float)(half)acc + (float)bias[n]); +} +kernel void rel_pos_bias_f16(device const half* q [[buffer(0)]], device const half* r [[buffer(1)]], + device half* y [[buffer(2)]], constant uint& T [[buffer(3)]], constant uint& H [[buffer(4)]], + constant uint& D [[buffer(5)]], constant uint& rbs [[buffer(6)]], + constant float& scale [[buffer(7)]], + uint3 gid [[thread_position_in_grid]]) { + uint tj = gid.x, h = gid.y, b = gid.z; + if (tj >= T * T || h >= H) return; + uint t = tj / T, j = tj % T; + device const half* qv = q + ((size_t)b * T + t) * H * D + h * D; + uint rel = (T - 1u) - t + j; + device const half* rv = r + (size_t)b * rbs + (size_t)rel * H * D + h * D; + float acc = 0; + for (uint d = 0; d < D; ++d) acc = fma((float)qv[d], (float)rv[d], acc); + y[((size_t)b * H + h) * T * T + (size_t)t * T + j] = (half)(acc * scale); +} +kernel void conv_cache_append_f16(device const uchar* src [[buffer(0)]], + device half* ring [[buffer(1)]], device half* out [[buffer(2)]], + constant uint& hd [[buffer(3)]], constant uint& ws [[buffer(4)]], + constant uint& nnew [[buffer(5)]], constant uint& head0 [[buffer(6)]], + constant uint& count_new [[buffer(7)]], constant uint& num_rows [[buffer(8)]], + constant uint& src_f32 [[buffer(9)]], + uint2 gid [[thread_position_in_grid]]) { + uint x = gid.x, y = gid.y; + if (x >= hd || y >= ws + nnew) return; + uint start_row = num_rows - nnew; + if (y >= ws) { + uint i = y - ws; + uint sidx = (start_row + i) * hd + x; + half v = src_f32 ? (half)((device const float*)src)[sidx] : ((device const half*)src)[sidx]; + ring[((head0 + i) % ws) * hd + x] = v; + return; + } + uint pad = ws - count_new; + if (y < pad) { out[y * hd + x] = 0.0h; return; } + uint a = count_new - 1u - (y - pad); + half v; + if (a < nnew) { + uint sidx = (num_rows - 1u - a) * hd + x; + v = src_f32 ? (half)((device const float*)src)[sidx] : ((device const half*)src)[sidx]; + } else { + uint b = a - nnew; + v = ring[((head0 + 2u * ws - 1u - b) % ws) * hd + x]; + } + out[y * hd + x] = v; +} +kernel void kv_append_ring_i8_m(device const half* src [[buffer(0)]], device char* int8base [[buffer(1)]], + device float* scalebase [[buffer(2)]], constant uint& kv_heads [[buffer(3)]], + constant uint& hdim [[buffer(4)]], constant uint& current_len [[buffer(5)]], + constant uint& group_size [[buffer(6)]], constant uint& M [[buffer(7)]], + constant uint& sink [[buffer(8)]], constant uint& W [[buffer(9)]], + uint gid [[thread_position_in_grid]]) { + uint num_groups = (hdim + group_size - 1)/group_size, per = kv_heads*num_groups; + if (gid >= M*per) return; + uint i = gid / per, hg = gid % per, h = hg / num_groups, g = hg % num_groups; + uint gstart = g*group_size, gcount = min(group_size, hdim - gstart); + device const half* hs = src + (size_t)i*kv_heads*hdim + (size_t)h*hdim + gstart; + float maxabs = 0; + for (uint k=0;k sink) ? (W - sink) : 1u; + uint slot = (pos < W) ? pos : sink + ((pos - sink) % R); + if (pos >= sink && pos + R < current_len + M) return; + device char* dst = int8base + (size_t)slot*i8s + (size_t)h*hdim + gstart; + for (uint k=0;k= U.vocab) continue; + float v = (float)logits[id]; + logits[id] = (half)(v > 0.0f ? v/U.penalty : v*U.penalty); + } + } + threadgroup_barrier(mem_flags::mem_device); + if (t == 0u && U.suppress_flag != 0u && U.suppress_id < U.vocab) logits[U.suppress_id] = half(-65504.0f); +} + +kernel void softcap_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& n [[buffer(2)]], constant float& cap [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + if (i>=n) return; + half v1 = (half)((float)in[i]/cap); + half v2 = (half)precise::tanh((float)v1); + y[i]=(half)((float)v2*cap); +} + +kernel void argmax_part(device const half* logits [[buffer(0)]], + device float* part [[buffer(1)]], + constant uint& V [[buffer(2)]], + constant uint& chunk [[buffer(3)]], + device const float* bias [[buffer(4)]], + constant uint& has_bias [[buffer(5)]], + uint p [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* sb [[threadgroup(0)]], + threadgroup uint* si [[threadgroup(1)]], + threadgroup float* ss [[threadgroup(2)]]) { + uint lo = p*chunk, hi = min(V, lo+chunk); + float b = -INFINITY, s = -INFINITY; uint bi = lo; + for (uint i = lo+t; i < hi; i += T) { + float v = (float)logits[i]; + if (has_bias != 0u) v += bias[i]; + if (v > b) { s = b; b = v; bi = i; } + else if (v > s) { s = v; } + } + sb[t] = b; si[t] = bi; ss[t] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = T >> 1; stride > 0u; stride >>= 1) { + if (t < stride) { + float ab = sb[t], asx = ss[t]; uint ai = si[t]; + float bb = sb[t+stride], bsx = ss[t+stride]; uint bidx = si[t+stride]; + if (ab > bb || (ab == bb && ai <= bidx)) { sb[t] = ab; si[t] = ai; ss[t] = max(asx, bb); } + else { sb[t] = bb; si[t] = bidx; ss[t] = max(bsx, ab); } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t == 0u) { part[p*3u] = sb[0]; part[p*3u+1u] = ss[0]; part[p*3u+2u] = (float)si[0]; } +} + +kernel void argmax_final(device const float* part [[buffer(0)]], + device float* out3 [[buffer(1)]], + constant uint& NP [[buffer(2)]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* sb [[threadgroup(0)]], + threadgroup uint* si [[threadgroup(1)]], + threadgroup float* ss [[threadgroup(2)]]) { + float b = -INFINITY, s = -INFINITY; uint bi = 0u; + for (uint i = t; i < NP; i += T) { + float pb = part[i*3u], ps = part[i*3u+1u]; uint pi = (uint)part[i*3u+2u]; + if (pb > b || (pb == b && pi <= bi)) { s = max(s, b); b = pb; bi = pi; s = max(s, ps); } + else { s = max(s, pb); } + } + sb[t] = b; si[t] = bi; ss[t] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = T >> 1; stride > 0u; stride >>= 1) { + if (t < stride) { + float ab = sb[t], asx = ss[t]; uint ai = si[t]; + float bb = sb[t+stride], bsx = ss[t+stride]; uint bidx = si[t+stride]; + if (ab > bb || (ab == bb && ai <= bidx)) { sb[t] = ab; si[t] = ai; ss[t] = max(asx, bb); } + else { sb[t] = bb; si[t] = bidx; ss[t] = max(bsx, ab); } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t == 0u) { out3[0] = sb[0]; out3[1] = ss[0]; out3[2] = (float)si[0]; } +} + +kernel void argmax_logits(device const half* logits [[buffer(0)]], + device float* out3 [[buffer(1)]], + constant uint& V [[buffer(2)]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* sb [[threadgroup(0)]], + threadgroup uint* si [[threadgroup(1)]], + threadgroup float* ss [[threadgroup(2)]]) { + float b = -INFINITY, s = -INFINITY; uint bi = 0u; + for (uint i = t; i < V; i += T) { + float v = (float)logits[i]; + if (v > b) { s = b; b = v; bi = i; } + else if (v > s) { s = v; } + } + sb[t] = b; si[t] = bi; ss[t] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = T >> 1; stride > 0u; stride >>= 1) { + if (t < stride) { + float ab = sb[t], asx = ss[t]; uint ai = si[t]; + float bb = sb[t+stride], bsx = ss[t+stride]; uint bidx = si[t+stride]; + if (ab > bb || (ab == bb && ai <= bidx)) { sb[t] = ab; si[t] = ai; ss[t] = max(asx, bb); } + else { sb[t] = bb; si[t] = bidx; ss[t] = max(bsx, ab); } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (t == 0u) { out3[0] = sb[0]; out3[1] = ss[0]; out3[2] = (float)si[0]; } +} + +kernel void topk_rows_f16(device const half* in [[buffer(0)]], + device float* out [[buffer(1)]], + constant uint& F [[buffer(2)]], + constant uint& k [[buffer(3)]], + constant uint& B [[buffer(4)]], + uint row [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + device const half* x = in + (size_t)row*F; + device float* idx_out = out + (size_t)row*k; + device float* val_out = out + (size_t)B*k + (size_t)row*k; + for (uint j = 0u; j < k; ++j) { + float best = -INFINITY; uint bi = 0xFFFFFFFFu; + for (uint i = lane; i < F; i += 32u) { + bool taken = false; + for (uint p = 0u; p < j; ++p) if ((uint)idx_out[p] == i) taken = true; + float v = (float)x[i]; + if (!taken && (v > best || (v == best && i < bi))) { best = v; bi = i; } + } + float gb = simd_max(best); + uint cand = (best == gb) ? bi : 0xFFFFFFFFu; + uint gi = simd_min(cand); + if (lane == 0u) { idx_out[j] = (float)gi; val_out[j] = gb; } + simdgroup_barrier(mem_flags::mem_device); + } +} + +kernel void cq4_moe_transform( + device const half* x [[buffer(0)]], + device const float* topk [[buffer(1)]], + device const half* recip_c [[buffer(2)]], + device const char* lsign_c [[buffer(3)]], + device const char* rsign_c [[buffer(4)]], + device const uint* perm_c [[buffer(5)]], + device half* code [[buffer(6)]], + constant uint& K [[buffer(7)]], + constant uint& k_valid [[buffer(8)]], + constant uint& xz_stride [[buffer(9)]], + constant uint& xy_stride [[buffer(10)]], + constant uint& topk_k [[buffer(11)]], + constant uint& estride [[buffer(12)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + threadgroup float* zmem [[threadgroup(0)]]) +{ + const uint g = tgp.x, slot = tgp.y, tok = tgp.z; + const uint e = (uint)(topk[(size_t)tok*topk_k + slot] + 0.5f); + device const half* recip = (device const half*)((device const uchar*)recip_c + (size_t)e*estride); + device const char* lsign = lsign_c + (size_t)e*estride; + device const char* rsign = rsign_c + (size_t)e*estride; + device const uint* perm = (device const uint*)((device const uchar*)perm_c + (size_t)e*estride); + device const half* xs = x + (size_t)tok*xz_stride + (size_t)slot*xy_stride; + uint b = g*128u + lane*4u; + uint k = lane*4u; + float xv[4]; + #pragma clang loop unroll(full) + for (uint j = 0u; j < 4u; ++j) { + float v = (b+j < k_valid) ? (float)xs[b+j] : 0.0f; + xv[j] = v*(float)recip[b+j]*(float)lsign[k+j]; + } + float x0 = xv[0], x1 = xv[1], x2 = xv[2], x3 = xv[3]; + cq4_hada128(x0, x1, x2, x3, lane); + zmem[k+0] = x0*(float)rsign[k+0]; zmem[k+1] = x1*(float)rsign[k+1]; + zmem[k+2] = x2*(float)rsign[k+2]; zmem[k+3] = x3*(float)rsign[k+3]; + threadgroup_barrier(mem_flags::mem_threadgroup); + device half* co = code + ((size_t)tok*topk_k + slot)*K + g*128u; + co[k+0] = (half)zmem[perm[k+0]]; co[k+1] = (half)zmem[perm[k+1]]; + co[k+2] = (half)zmem[perm[k+2]]; co[k+3] = (half)zmem[perm[k+3]]; +} + +kernel void cq4_moe_gemv_up( + device const half* code1 [[buffer(0)]], + device const half* code3 [[buffer(1)]], + device const float* topk [[buffer(2)]], + device const uchar* pk1 [[buffer(3)]], + device const half* nm1 [[buffer(4)]], + device const half* cb1 [[buffer(5)]], + device const uchar* pk3 [[buffer(6)]], + device const half* nm3 [[buffer(7)]], + device const half* cb3 [[buffer(8)]], + device half* y [[buffer(9)]], + constant uint& K [[buffer(10)]], + constant uint& N [[buffer(11)]], + constant uint& act [[buffer(12)]], + constant uint& topk_k [[buffer(13)]], + constant uint& estride1 [[buffer(14)]], + constant uint& estride3 [[buffer(15)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint tl [[thread_index_in_threadgroup]]) +{ + const uint slot = tgp.y, tok = tgp.z; + const uint slotg = tok*topk_k + slot; + const uint e = (uint)(topk[(size_t)tok*topk_k + slot] + 0.5f); + const uint ng = K/128u; + threadgroup float c1[16], c3[16]; + if (tl < 16u) { + c1[tl] = (float)((device const half*)((device const uchar*)cb1 + (size_t)e*estride1))[tl]; + c3[tl] = (float)((device const half*)((device const uchar*)cb3 + (size_t)e*estride3))[tl]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint n0 = (tgp.x*8u + sgid)*4u; + if (n0 >= N) return; + float gv[4] = {0,0,0,0}, uv[4] = {0,0,0,0}; + for (int m = 0; m < 2; ++m) { + device const uchar* pb = (m==0) ? pk1 + (size_t)e*estride1 : pk3 + (size_t)e*estride3; + device const half* nb = (m==0) ? (device const half*)((device const uchar*)nm1 + (size_t)e*estride1) + : (device const half*)((device const uchar*)nm3 + (size_t)e*estride3); + device const half* cd = (m==0) ? code1 + (size_t)slotg*K : code3 + (size_t)slotg*K; + threadgroup float* cb = (m==0) ? c1 : c3; + float a0=0, a1=0, a2=0, a3=0; + for (uint base = lane*16u; base < K; base += 512u) { + uint gg = base/128u, off = base - gg*128u; + device const half4* cp = (device const half4*)(cd + gg*128u + off); + half4 c0=cp[0], c1v=cp[1], c2v=cp[2], c3v=cp[3]; + device const uchar* pan = pb + ((size_t)(n0>>2)*ng+gg)*256u + (off>>4)*32u; + uint2 A01 = ((device const uint2*)pan)[0]; + uint2 A23 = ((device const uint2*)(pan+8u))[0]; + uint2 B01 = ((device const uint2*)(pan+16u))[0]; + uint2 B23 = ((device const uint2*)(pan+24u))[0]; + size_t nb4 = ((size_t)(n0>>2)*ng+gg)<<2; + a0 += (float)nb[nb4+0u]*cq4_dot64_ilw(c0, c1v, c2v, c3v, A01.x, B01.x, cb); + a1 += (float)nb[nb4+1u]*cq4_dot64_ilw(c0, c1v, c2v, c3v, A01.y, B01.y, cb); + a2 += (float)nb[nb4+2u]*cq4_dot64_ilw(c0, c1v, c2v, c3v, A23.x, B23.x, cb); + a3 += (float)nb[nb4+3u]*cq4_dot64_ilw(c0, c1v, c2v, c3v, A23.y, B23.y, cb); + } + a0 = simd_sum(a0); a1 = simd_sum(a1); a2 = simd_sum(a2); a3 = simd_sum(a3); + if (m == 0) { gv[0]=a0; gv[1]=a1; gv[2]=a2; gv[3]=a3; } + else { uv[0]=a0; uv[1]=a1; uv[2]=a2; uv[3]=a3; } + } + if (lane == 0u) { + device half* ys = y + (size_t)slotg*N; + for (uint r = 0u; r < 4u && n0+r < N; ++r) { + float gh = (float)(half)gv[r]; + float av = (act == 1u) ? gelu_tanh(gh) : gh/(1.0f + precise::exp(-gh)); + ys[n0+r] = (half)((float)(half)av*(float)(half)uv[r]); + } + } +} + +kernel void cq4_moe_transform2( + device const half* x [[buffer(0)]], + device const float* topk [[buffer(1)]], + device const half* recip_c1 [[buffer(2)]], + device const char* lsign_c1 [[buffer(3)]], + device const char* rsign_c1 [[buffer(4)]], + device const uint* perm_c1 [[buffer(5)]], + device const half* recip_c3 [[buffer(6)]], + device const char* lsign_c3 [[buffer(7)]], + device const char* rsign_c3 [[buffer(8)]], + device const uint* perm_c3 [[buffer(9)]], + device half* code1 [[buffer(10)]], + device half* code3 [[buffer(11)]], + constant uint& K [[buffer(12)]], + constant uint& topk_k [[buffer(13)]], + constant uint& estride1 [[buffer(14)]], + constant uint& estride3 [[buffer(15)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], + threadgroup float* zmem [[threadgroup(0)]]) +{ + const uint g = tgp.x, tok = tgp.z; + const bool w3 = tgp.y >= topk_k; + const uint slot = w3 ? tgp.y - topk_k : tgp.y; + const uint e = (uint)(topk[(size_t)tok*topk_k + slot] + 0.5f); + const uint es = w3 ? estride3 : estride1; + device const half* recip = (device const half*)((device const uchar*)(w3 ? recip_c3 : recip_c1) + (size_t)e*es); + device const char* lsign = (w3 ? lsign_c3 : lsign_c1) + (size_t)e*es; + device const char* rsign = (w3 ? rsign_c3 : rsign_c1) + (size_t)e*es; + device const uint* perm = (device const uint*)((device const uchar*)(w3 ? perm_c3 : perm_c1) + (size_t)e*es); + device const half* xs = x + (size_t)tok*K; + uint b = g*128u + lane*4u; + uint k = lane*4u; + float x0=(float)xs[b+0]*(float)recip[b+0]*(float)lsign[k+0]; + float x1=(float)xs[b+1]*(float)recip[b+1]*(float)lsign[k+1]; + float x2=(float)xs[b+2]*(float)recip[b+2]*(float)lsign[k+2]; + float x3=(float)xs[b+3]*(float)recip[b+3]*(float)lsign[k+3]; + cq4_hada128(x0, x1, x2, x3, lane); + zmem[k+0]=x0*(float)rsign[k+0]; zmem[k+1]=x1*(float)rsign[k+1]; + zmem[k+2]=x2*(float)rsign[k+2]; zmem[k+3]=x3*(float)rsign[k+3]; + threadgroup_barrier(mem_flags::mem_threadgroup); + device half* co = (w3 ? code3 : code1) + ((size_t)tok*topk_k + slot)*K + g*128u; + co[k+0]=(half)zmem[perm[k+0]]; co[k+1]=(half)zmem[perm[k+1]]; + co[k+2]=(half)zmem[perm[k+2]]; co[k+3]=(half)zmem[perm[k+3]]; +} + +kernel void cq4_moe_gemv_down_acc( + device const half* code2 [[buffer(0)]], + device const float* topk [[buffer(1)]], + device const half* probs [[buffer(2)]], + device const uchar* pk2 [[buffer(3)]], + device const half* nm2 [[buffer(4)]], + device const half* cb2 [[buffer(5)]], + device half* out [[buffer(6)]], + constant uint& K [[buffer(7)]], + constant uint& N [[buffer(8)]], + constant uint& topk_k [[buffer(9)]], + constant uint& normalize [[buffer(10)]], + constant float& eps [[buffer(11)]], + constant float& scaling [[buffer(12)]], + constant uint& E [[buffer(13)]], + constant uint& estride [[buffer(14)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint tl [[thread_index_in_threadgroup]], + threadgroup float* cbs [[threadgroup(0)]]) +{ + const uint tok = tgp.z; + device const float* tk = topk + (size_t)tok*topk_k; + device const half* pr = probs + (size_t)tok*E; + uint order[16]; + for (uint j = 0u; j < topk_k; ++j) order[j] = j; + for (uint a = 1u; a < topk_k; ++a) { + uint v = order[a]; + uint ev = (uint)(tk[v] + 0.5f); + uint b = a; + while (b > 0u && (uint)(tk[order[b-1u]] + 0.5f) > ev) { order[b] = order[b-1u]; --b; } + order[b] = v; + } + if (tl < topk_k*16u) { + uint s = tl >> 4, ci = tl & 15u; + uint e = (uint)(tk[s] + 0.5f); + cbs[tl] = (float)((device const half*)((device const uchar*)cb2 + (size_t)e*estride))[ci]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float denom = 1.0f; + if (normalize != 0u) { + float s = 0.0f; + for (uint j = 0u; j < topk_k; ++j) s += (float)pr[(uint)(tk[j] + 0.5f)]; + denom = s + eps; + } + const uint ng = K/128u; + uint n0 = (tgp.x*8u + sgid)*2u; + if (n0 >= N) return; + half res[2] = {0.0h, 0.0h}; + for (uint j = 0u; j < topk_k; ++j) { + uint slot = order[j]; + uint e = (uint)(tk[slot] + 0.5f); + float p = (float)pr[e]; + if (p <= 0.0f) continue; + float w = p/denom*scaling; + device const uchar* pbase = pk2 + (size_t)e*estride; + device const half* nbase = (device const half*)((device const uchar*)nm2 + (size_t)e*estride); + device const half* cd = code2 + ((size_t)tok*topk_k + slot)*K; + threadgroup const float* cb = cbs + (size_t)slot*16u; + float acc[2] = {0, 0}; + for (uint base = lane*16u; base < K; base += 512u) { + uint gg = base/128u, off = base - gg*128u; + device const half4* cp = (device const half4*)(cd + gg*128u + off); + half4 c0=cp[0], c1=cp[1], c2=cp[2], c3=cp[3]; + device const uchar* pan = pbase + ((size_t)(n0>>2)*ng+gg)*256u + (n0&3u)*4u + (off>>4)*32u; + uint2 Aw = ((device const uint2*)pan)[0]; + uint2 Bw = ((device const uint2*)(pan+16u))[0]; + #pragma clang loop unroll(full) + for (uint r = 0u; r < 2u; ++r) { + float pv = cq4_dot64_ilw(c0, c1, c2, c3, r?Aw.y:Aw.x, r?Bw.y:Bw.x, cb); + acc[r] += (float)nbase[(((size_t)(n0>>2)*ng+gg)<<2)+(n0&3u)+r]*pv; + } + } + #pragma clang loop unroll(full) + for (uint r = 0u; r < 2u; ++r) { + float a = simd_sum(acc[r]); + res[r] = (half)((float)res[r] + (float)(half)a*w); + } + } + if (lane == 0u) { + device half* ys = out + (size_t)tok*N; + for (uint r = 0u; r < 2u && n0+r < N; ++r) ys[n0+r] = res[r]; + } +} + +kernel void rope_pair_f16( + device const half* x [[buffer(0)]], + device const half* c [[buffer(1)]], + device const half* s [[buffer(2)]], + device half* y [[buffer(3)]], + constant uint& H [[buffer(4)]], + constant uint& D [[buffer(5)]], + uint2 g [[thread_position_in_grid]]) +{ + uint d = g.x, h = g.y; + if (d >= D || h >= H) return; + size_t idx = (size_t)h*D + d; + half xv = x[idx]; + half rh = (d < D/2u) ? (half)(-x[idx + D/2u]) : x[idx - D/2u]; + y[idx] = xv*c[d] + rh*s[d]; +} + +kernel void rms_norm_scale_f16( + device const half* in [[buffer(0)]], + device const half* w [[buffer(1)]], + device half* y [[buffer(2)]], + constant uint& D [[buffer(3)]], + constant float& eps [[buffer(4)]], + constant float& oscale [[buffer(5)]], + uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* red [[threadgroup(0)]]) +{ + device const half* x = in + (size_t)row*D; + device half* o = y + (size_t)row*D; + float ss = 0; + for (uint i = t; i < D; i += T) { float v = (float)x[i]; ss += v*v; } + red[t] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = T/2; s > 0; s >>= 1) { + if (t < s) red[t] += red[t+s]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float inv = rsqrt(red[0]/(float)D + eps); + for (uint i = t; i < D; i += T) { + half y0 = (half)((float)x[i]*inv*(float)w[i]); + o[i] = (half)((float)y0*oscale); + } +} + +kernel void softmax_topk_f16( + device const half* in [[buffer(0)]], + device half* probs [[buffer(1)]], + device float* tk [[buffer(2)]], + constant uint& E [[buffer(3)]], + constant uint& k [[buffer(4)]], + constant uint& B [[buffer(5)]], + constant float& scale [[buffer(6)]], + uint row [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) +{ + device const half* x = in + (size_t)row*E; + device half* pr = probs + (size_t)row*E; + device float* idx_out = tk + (size_t)row*k; + device float* val_out = tk + (size_t)B*k + (size_t)row*k; + float mx = -INFINITY; + for (uint i = lane; i < E; i += 32u) mx = max(mx, (float)(half)((float)x[i]*scale)); + mx = simd_max(mx); + float sum = 0; + for (uint i = lane; i < E; i += 32u) sum += exp((float)(half)((float)x[i]*scale) - mx); + sum = simd_sum(sum); + float inv = 1.0f/sum; + for (uint i = lane; i < E; i += 32u) pr[i] = (half)(exp((float)(half)((float)x[i]*scale) - mx)*inv); + for (uint j = 0u; j < k; ++j) { + float best = -INFINITY; uint bi = 0xFFFFFFFFu; + for (uint i = lane; i < E; i += 32u) { + bool taken = false; + for (uint p = 0u; p < j; ++p) if ((uint)idx_out[p] == i) taken = true; + float v = (float)(half)((float)x[i]*scale); + if (!taken && (v > best || (v == best && i < bi))) { best = v; bi = i; } + } + float gb = simd_max(best); + uint cand = (best == gb) ? bi : 0xFFFFFFFFu; + uint gi = simd_min(cand); + if (lane == 0u) { idx_out[j] = (float)gi; val_out[j] = gb; } + simdgroup_barrier(mem_flags::mem_device); + } +} + +kernel void rope_pair_rms_f16( + device const half* x [[buffer(0)]], + device const half* w [[buffer(1)]], + device const half* c [[buffer(2)]], + device const half* s [[buffer(3)]], + device half* y [[buffer(4)]], + constant uint& D [[buffer(5)]], + constant float& eps [[buffer(6)]], + uint h [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* red [[threadgroup(0)]], + threadgroup half* xn [[threadgroup(1)]]) +{ + device const half* xr = x + (size_t)h*D; + float ss = 0; + for (uint i = t; i < D; i += T) { float v = (float)xr[i]; ss += v*v; } + red[t] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint st = T/2; st > 0; st >>= 1) { + if (t < st) red[t] += red[t+st]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float inv = rsqrt(red[0]/(float)D + eps); + for (uint i = t; i < D; i += T) xn[i] = (half)((float)xr[i]*inv*(float)w[i]); + threadgroup_barrier(mem_flags::mem_threadgroup); + device half* yr = y + (size_t)h*D; + for (uint i = t; i < D; i += T) { + half xv = xn[i]; + half rh = (i < D/2u) ? (half)(-xn[i + D/2u]) : xn[i - D/2u]; + yr[i] = xv*c[i] + rh*s[i]; + } +} + +kernel void rms2_add_clip_f16( + device const half* a [[buffer(0)]], + device const half* wa [[buffer(1)]], + device const half* b [[buffer(2)]], + device const half* wb [[buffer(3)]], + device half* y [[buffer(4)]], + constant uint& D [[buffer(5)]], + constant float& eps_a [[buffer(6)]], + constant float& eps_b [[buffer(7)]], + uint t [[thread_position_in_threadgroup]], + uint T [[threads_per_threadgroup]], + threadgroup float* red [[threadgroup(0)]]) +{ + float sa = 0, sb = 0; + for (uint i = t; i < D; i += T) { + float va = (float)a[i], vb = (float)b[i]; + sa += va*va; sb += vb*vb; + } + red[t] = sa; red[T+t] = sb; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = T/2; s > 0; s >>= 1) { + if (t < s) { red[t] += red[t+s]; red[T+t] += red[T+t+s]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float inva = rsqrt(red[0]/(float)D + eps_a); + float invb = rsqrt(red[T]/(float)D + eps_b); + for (uint i = t; i < D; i += T) { + half ya = (half)((float)a[i]*inva*(float)wa[i]); + half yb = (half)((float)b[i]*invb*(float)wb[i]); + y[i] = (half)clamp((float)ya + (float)yb, -65500.0f, 65500.0f); + } +} + +kernel void gated_deltanet_decode_f16( + device const half* q [[buffer(0)]], + device const half* k [[buffer(1)]], + device const half* v [[buffer(2)]], + device const half* g [[buffer(3)]], + device const half* b [[buffer(4)]], + device const half* s [[buffer(5)]], + device half* y [[buffer(6)]], + constant uint& Hq [[buffer(7)]], + constant uint& Hv [[buffer(8)]], + constant uint& K [[buffer(9)]], + constant uint& V [[buffer(10)]], + constant float& scale [[buffer(11)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], + threadgroup half* kq [[threadgroup(0)]]) +{ + const uint h = tgp.x, batch = tgp.y; + const uint qk_head = h / (Hv / Hq); + const uint qk_base = (batch * Hq + qk_head) * K; + threadgroup half* kk = kq; + threadgroup half* qq = kq + K; + for (uint i = t; i < K; i += V) { + kk[i] = k[qk_base + i]; + qq[i] = q[qk_base + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t >= V) return; + + float gate_log = (float)g[batch * Hv + h]; + if (!isfinite(gate_log)) gate_log = -20.0f; + float beta = (float)b[batch * Hv + h]; + beta = isfinite(beta) ? clamp(beta, 0.0f, 1.0f) : 0.0f; + const float alpha = exp(clamp(gate_log, -20.0f, 6.0f)); + + const size_t s_stride = (size_t)Hv * V; + device const half* sp = s + ((size_t)batch * K) * s_stride + (size_t)h * V + t; + float proj = 0.0f; + for (uint kd = 0; kd < K; ++kd) proj += (float)sp[kd * s_stride] * (float)kk[kd]; + const float delta = ((float)v[((size_t)batch * Hv + h) * V + t] - alpha * proj) * beta; + + device half* yp = y + ((size_t)batch * (1 + K) + 1) * s_stride + (size_t)h * V + t; + float acc = 0.0f; + for (uint kd = 0; kd < K; ++kd) { + float s_new = (float)sp[kd * s_stride] * alpha + (float)kk[kd] * delta; + if (!isfinite(s_new)) s_new = 0.0f; + yp[kd * s_stride] = (half)s_new; + acc += s_new * (float)qq[kd]; + } + if (!isfinite(acc)) acc = 0.0f; + y[((size_t)batch * (1 + K)) * s_stride + (size_t)h * V + t] = (half)(acc * scale); +} + +kernel void gated_deltanet_prefill_f16( + device const half* q [[buffer(0)]], + device const half* k [[buffer(1)]], + device const half* v [[buffer(2)]], + device const half* g [[buffer(3)]], + device const half* b [[buffer(4)]], + device const half* s [[buffer(5)]], + device half* y [[buffer(6)]], + device float* st [[buffer(7)]], + constant uint& T [[buffer(8)]], + constant uint& Hq [[buffer(9)]], + constant uint& Hv [[buffer(10)]], + constant uint& K [[buffer(11)]], + constant uint& V [[buffer(12)]], + constant float& scale [[buffer(13)]], + uint3 tgp [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], + threadgroup half* kq [[threadgroup(0)]]) +{ + const uint h = tgp.x, batch = tgp.y; + const uint qk_head = h / (Hv / Hq); + threadgroup half* kk = kq; + threadgroup half* qq = kq + K; + const size_t hv_stride = (size_t)Hv * V; + device float* col = st + (((size_t)batch * Hv + h) * K) * V + t; + + if (t < V) { + device const half* sp = s + ((size_t)batch * K) * hv_stride + (size_t)h * V + t; + for (uint kd = 0; kd < K; ++kd) col[(size_t)kd * V] = (float)sp[kd * hv_stride]; + } + + for (uint step = 0; step < T; ++step) { + threadgroup_barrier(mem_flags::mem_threadgroup); + const size_t qk_base = ((size_t)(batch * T + step) * Hq + qk_head) * K; + for (uint i = t; i < K; i += V) { + kk[i] = k[qk_base + i]; + qq[i] = q[qk_base + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t >= V) continue; + + float gate_log = (float)g[(size_t)(batch * T + step) * Hv + h]; + if (!isfinite(gate_log)) gate_log = -20.0f; + float beta = (float)b[(size_t)(batch * T + step) * Hv + h]; + beta = isfinite(beta) ? clamp(beta, 0.0f, 1.0f) : 0.0f; + const float alpha = exp(clamp(gate_log, -20.0f, 6.0f)); + + float proj = 0.0f; + for (uint kd = 0; kd < K; ++kd) proj += col[(size_t)kd * V] * (float)kk[kd]; + const float delta = ((float)v[((size_t)(batch * T + step) * Hv + h) * V + t] - alpha * proj) * beta; + + float acc = 0.0f; + for (uint kd = 0; kd < K; ++kd) { + float s_new = col[(size_t)kd * V] * alpha + (float)kk[kd] * delta; + if (!isfinite(s_new)) s_new = 0.0f; + col[(size_t)kd * V] = s_new; + acc += s_new * (float)qq[kd]; + } + if (!isfinite(acc)) acc = 0.0f; + y[((size_t)(batch * (T + K) + step)) * hv_stride + (size_t)h * V + t] = (half)(acc * scale); + } + + if (t < V) { + for (uint kd = 0; kd < K; ++kd) + y[((size_t)(batch * (T + K) + T + kd)) * hv_stride + (size_t)h * V + t] = (half)col[(size_t)kd * V]; + } +} diff --git a/cactus-kernels/src/conv.cpp b/cactus-kernels/src/conv.cpp new file mode 100644 index 000000000..18990605b --- /dev/null +++ b/cactus-kernels/src/conv.cpp @@ -0,0 +1,964 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +#endif + +constexpr size_t T_TILE_F16 = 2; +#ifdef __APPLE__ +constexpr size_t ACCELERATE_K_THRESHOLD = 32; +constexpr size_t ACCELERATE_L_THRESHOLD = 128; +#endif + +#ifdef __APPLE__ +static void conv1d_causal_depthwise_f16_accelerate( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, size_t L, size_t C, size_t K, size_t dilation) +{ + const size_t in_bs = L * C; + const size_t out_bs = L * C; + CactusThreading::parallel_for_2d(N, C, CactusThreading::Thresholds::ATTENTION, [&](size_t n, size_t c) { + const __fp16* Xb = input + n * in_bs; + __fp16* Yb = output + n * out_bs; + const __fp16* Wc = weight + c * K; + + if (dilation == 1) { + + size_t i = 0; + for(; i + 7 < L; i += 8){ + float16x8_t acc = vdupq_n_f16(0.0f); + + for (size_t k = 0; k < K; ++k) { + float16x8_t input_vec = vld1q_f16(&Xb[(i + k) * C + c]); + float16x8_t weight_vec = vdupq_n_f16(Wc[k]); + acc = vfmaq_f16(acc, input_vec, weight_vec); + } + vst1q_f16(&Yb[i * C + c], acc); + } + + for (; i < L; ++i) { + float acc = 0.0f; + for (size_t k = 0; k < K; ++k) { + acc += float(Xb[(i + k) * C + c]) * float(Wc[k]); + } + Yb[i * C + c] = (__fp16)acc; + } + } else { + + for (size_t i = 0; i < L; ++i) { + float acc = 0.0f; + + for (size_t k = 0; k < K; ++k) { + size_t idx = i + k * dilation; + acc += float(Xb[idx * C + c]) * float(Wc[k]); + } + + Yb[i * C + c] = (__fp16)acc; + } + } + }); +} +#endif + +void cactus_conv1d_causal_depthwise_f16( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, size_t L, size_t C, size_t K, size_t dilation) +{ +#ifdef __APPLE__ + if (K >= ACCELERATE_K_THRESHOLD && L >= ACCELERATE_L_THRESHOLD) { + conv1d_causal_depthwise_f16_accelerate(input, weight, output, N, L, C, K, dilation); + return; + } +#endif + + const size_t in_bs = L * C; + const size_t out_bs = L * C; + + CactusThreading::parallel_for_2d(N, C, CactusThreading::Thresholds::ATTENTION, [&](size_t n, size_t c) { + const __fp16* Xb = input + n * in_bs; + __fp16* Yb = output + n * out_bs; + + std::vector wrev(K); + const __fp16* Wc = weight + c * K; + for (size_t k = 0; k < K; ++k) wrev[k] = (float)Wc[K - 1 - k]; + + for (size_t t0 = 0; t0 < L; t0 += T_TILE_F16) { + const size_t t1 = std::min(t0 + 1, L - 1); + + float32x4_t vacc0 = vdupq_n_f32(0.f); + float32x4_t vacc1 = vdupq_n_f32(0.f); + + size_t k = 0; + for (; k + 8 <= K; k += 8) { + + float x0_0=0, x1_0=0, x2_0=0, x3_0=0; + float x0_1=0, x1_1=0, x2_1=0, x3_1=0; + { + ptrdiff_t a0=(ptrdiff_t)t0-(ptrdiff_t)((k+0)*dilation); + ptrdiff_t a1=(ptrdiff_t)t0-(ptrdiff_t)((k+1)*dilation); + ptrdiff_t a2=(ptrdiff_t)t0-(ptrdiff_t)((k+2)*dilation); + ptrdiff_t a3=(ptrdiff_t)t0-(ptrdiff_t)((k+3)*dilation); + if (a0>=0) x0_0 = (float)Xb[(size_t)a0*C + c]; + if (a1>=0) x1_0 = (float)Xb[(size_t)a1*C + c]; + if (a2>=0) x2_0 = (float)Xb[(size_t)a2*C + c]; + if (a3>=0) x3_0 = (float)Xb[(size_t)a3*C + c]; + + ptrdiff_t b0=(ptrdiff_t)t1-(ptrdiff_t)((k+0)*dilation); + ptrdiff_t b1=(ptrdiff_t)t1-(ptrdiff_t)((k+1)*dilation); + ptrdiff_t b2=(ptrdiff_t)t1-(ptrdiff_t)((k+2)*dilation); + ptrdiff_t b3=(ptrdiff_t)t1-(ptrdiff_t)((k+3)*dilation); + if (b0>=0) x0_1 = (float)Xb[(size_t)b0*C + c]; + if (b1>=0) x1_1 = (float)Xb[(size_t)b1*C + c]; + if (b2>=0) x2_1 = (float)Xb[(size_t)b2*C + c]; + if (b3>=0) x3_1 = (float)Xb[(size_t)b3*C + c]; + } + float32x4_t xv0 = {x0_0,x1_0,x2_0,x3_0}; + float32x4_t yv0 = {x0_1,x1_1,x2_1,x3_1}; + float32x4_t wv0 = {wrev[k+0],wrev[k+1],wrev[k+2],wrev[k+3]}; + vacc0 = vfmaq_f32(vacc0, xv0, wv0); + vacc1 = vfmaq_f32(vacc1, yv0, wv0); + + float a0_0=0, a1_0=0, a2_0=0, a3_0=0; + float a0_1=0, a1_1=0, a2_1=0, a3_1=0; + { + ptrdiff_t a0i=(ptrdiff_t)t0-(ptrdiff_t)((k+4)*dilation); + ptrdiff_t a1i=(ptrdiff_t)t0-(ptrdiff_t)((k+5)*dilation); + ptrdiff_t a2i=(ptrdiff_t)t0-(ptrdiff_t)((k+6)*dilation); + ptrdiff_t a3i=(ptrdiff_t)t0-(ptrdiff_t)((k+7)*dilation); + if (a0i>=0) a0_0 = (float)Xb[(size_t)a0i*C + c]; + if (a1i>=0) a1_0 = (float)Xb[(size_t)a1i*C + c]; + if (a2i>=0) a2_0 = (float)Xb[(size_t)a2i*C + c]; + if (a3i>=0) a3_0 = (float)Xb[(size_t)a3i*C + c]; + + ptrdiff_t b0i=(ptrdiff_t)t1-(ptrdiff_t)((k+4)*dilation); + ptrdiff_t b1i=(ptrdiff_t)t1-(ptrdiff_t)((k+5)*dilation); + ptrdiff_t b2i=(ptrdiff_t)t1-(ptrdiff_t)((k+6)*dilation); + ptrdiff_t b3i=(ptrdiff_t)t1-(ptrdiff_t)((k+7)*dilation); + if (b0i>=0) a0_1 = (float)Xb[(size_t)b0i*C + c]; + if (b1i>=0) a1_1 = (float)Xb[(size_t)b1i*C + c]; + if (b2i>=0) a2_1 = (float)Xb[(size_t)b2i*C + c]; + if (b3i>=0) a3_1 = (float)Xb[(size_t)b3i*C + c]; + } + float32x4_t xv1 = {a0_0,a1_0,a2_0,a3_0}; + float32x4_t yv1 = {a0_1,a1_1,a2_1,a3_1}; + float32x4_t wv1 = {wrev[k+4],wrev[k+5],wrev[k+6],wrev[k+7]}; + vacc0 = vfmaq_f32(vacc0, xv1, wv1); + vacc1 = vfmaq_f32(vacc1, yv1, wv1); + } + + float acc0 = vaddvq_f32(vacc0); + float acc1 = vaddvq_f32(vacc1); + + for (; k < K; ++k) { + ptrdiff_t a=(ptrdiff_t)t0-(ptrdiff_t)(k*dilation); + if (a>=0) acc0 += wrev[k] * (float)Xb[(size_t)a*C + c]; + ptrdiff_t b=(ptrdiff_t)t1-(ptrdiff_t)(k*dilation); + if (b>=0) acc1 += wrev[k] * (float)Xb[(size_t)b*C + c]; + } + + Yb[t0*C + c] = (__fp16)acc0; + if (t0 + 1 < L) Yb[t1*C + c] = (__fp16)acc1; + } + }); +} + +void cactus_conv1d_f16_k3( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, size_t L, + size_t C_in, size_t C_out, + size_t stride +){ + const size_t out_len = ((L - 1) / stride) + 1; + const size_t in_bs = C_in * L; + const size_t out_bs = C_out * out_len; + + const size_t total_compute = N * C_out * out_len * C_in * 3; + CactusThreading::ParallelConfig config = (total_compute < 100000) + ? CactusThreading::ParallelConfig{SIZE_MAX, SIZE_MAX} + : CactusThreading::Thresholds::ATTENTION; + + CactusThreading::parallel_for_2d(N, C_out, config, [&](size_t n, size_t oc) { + const __fp16* Xb = input + n * in_bs; + __fp16* Yoc = output + n * out_bs + oc * out_len; + const __fp16* Woc = weight + oc * (C_in * 3); + + for (size_t out_idx = 0; out_idx < out_len; out_idx += 2) { + const size_t out_t0 = out_idx; + const bool have_t1 = (out_idx + 1) < out_len; + const size_t out_t1 = have_t1 ? (out_idx + 1) : out_idx; + + const size_t t0 = out_t0 * stride; + const size_t t1 = have_t1 ? (out_t1 * stride) : t0; + + float32x4_t acc0 = vdupq_n_f32(0.f); + float32x4_t acc1 = vdupq_n_f32(0.f); + + size_t ic = 0; + for (; ic + 16 <= C_in; ic += 16) { + for (size_t u = 0; u < 16; ++u) { + const __fp16* Xc = Xb + (ic + u) * L; + const __fp16* Wc = Woc + (ic + u) * 3; + + const float16x8_t wv = { + Wc[0], Wc[1], Wc[2], (__fp16)0, + Wc[0], Wc[1], Wc[2], (__fp16)0 + }; + + const ptrdiff_t tm0 = (ptrdiff_t)t0 - 1; + const ptrdiff_t tp0 = (ptrdiff_t)t0 + 1; + const ptrdiff_t tm1 = (ptrdiff_t)t1 - 1; + const ptrdiff_t tp1 = (ptrdiff_t)t1 + 1; + + const __fp16 x0m = (tm0 >= 0) ? Xc[tm0] : (__fp16)0; + const __fp16 x00 = Xc[t0]; + const __fp16 x0p = (tp0 < (ptrdiff_t)L) ? Xc[tp0] : (__fp16)0; + + __fp16 x1m = 0, x10 = 0, x1p = 0; + if (have_t1) { + x1m = (tm1 >= 0) ? Xc[tm1] : (__fp16)0; + x10 = Xc[t1]; + x1p = (tp1 < (ptrdiff_t)L) ? Xc[tp1] : (__fp16)0; + } + + const float16x8_t xv = { + x0m, x00, x0p, (__fp16)0, + x1m, x10, x1p, (__fp16)0 + }; + + const float16x4_t xv0_h = vget_low_f16(xv); + const float16x4_t wv0_h = vget_low_f16(wv); + acc0 = vfmaq_f32(acc0, vcvt_f32_f16(xv0_h), vcvt_f32_f16(wv0_h)); + + if (have_t1) { + const float16x4_t xv1_h = vget_high_f16(xv); + const float16x4_t wv1_h = vget_high_f16(wv); + acc1 = vfmaq_f32(acc1, vcvt_f32_f16(xv1_h), vcvt_f32_f16(wv1_h)); + } + } + } + + for (; ic < C_in; ++ic) { + const __fp16* Xc = Xb + ic * L; + const __fp16* Wc = Woc + ic * 3; + + const float16x8_t wv = { + Wc[0], Wc[1], Wc[2], (__fp16)0, + Wc[0], Wc[1], Wc[2], (__fp16)0 + }; + + const ptrdiff_t tm0 = (ptrdiff_t)t0 - 1; + const ptrdiff_t tp0 = (ptrdiff_t)t0 + 1; + const ptrdiff_t tm1 = (ptrdiff_t)t1 - 1; + const ptrdiff_t tp1 = (ptrdiff_t)t1 + 1; + + const __fp16 x0m = (tm0 >= 0) ? Xc[tm0] : (__fp16)0; + const __fp16 x00 = Xc[t0]; + const __fp16 x0p = (tp0 < (ptrdiff_t)L) ? Xc[tp0] : (__fp16)0; + + __fp16 x1m = 0, x10 = 0, x1p = 0; + if (have_t1) { + x1m = (tm1 >= 0) ? Xc[tm1] : (__fp16)0; + x10 = Xc[t1]; + x1p = (tp1 < (ptrdiff_t)L) ? Xc[tp1] : (__fp16)0; + } + + const float16x8_t xv = { + x0m, x00, x0p, (__fp16)0, + x1m, x10, x1p, (__fp16)0 + }; + + const float16x4_t xv0_h = vget_low_f16(xv); + const float16x4_t wv0_h = vget_low_f16(wv); + acc0 = vfmaq_f32(acc0, vcvt_f32_f16(xv0_h), vcvt_f32_f16(wv0_h)); + + if (have_t1) { + const float16x4_t xv1_h = vget_high_f16(xv); + const float16x4_t wv1_h = vget_high_f16(wv); + acc1 = vfmaq_f32(acc1, vcvt_f32_f16(xv1_h), vcvt_f32_f16(wv1_h)); + } + } + + float32x2_t s0 = vadd_f32(vget_low_f32(acc0), vget_high_f32(acc0)); + float sum0 = vget_lane_f32(s0, 0) + vget_lane_f32(s0, 1); + Yoc[out_t0] = (__fp16)sum0; + + if (have_t1) { + float32x2_t s1 = vadd_f32(vget_low_f32(acc1), vget_high_f32(acc1)); + float sum1 = vget_lane_f32(s1, 0) + vget_lane_f32(s1, 1); + Yoc[out_t1] = (__fp16)sum1; + } + } + }); +} + +#ifdef __APPLE__ +static void conv1d_f16_accelerate( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, size_t L, + size_t C_in, size_t C_out, + size_t K, + size_t stride +){ + const size_t out_len = ((L - K) / stride) + 1; + const size_t in_bs = C_in * L; + const size_t out_bs = C_out * out_len; + + const size_t total_compute = N * C_out * out_len * C_in * K; + CactusThreading::ParallelConfig config = (total_compute < 100000) + ? CactusThreading::ParallelConfig{SIZE_MAX, SIZE_MAX} + : CactusThreading::Thresholds::ATTENTION; + + CactusThreading::parallel_for_2d( + N, C_out, config, + [&](size_t n, size_t oc) { + + const __fp16* Xb = input + n * in_bs; + __fp16* Yoc = output + n * out_bs + oc * out_len; + const __fp16* Woc = weight + oc * (C_in * K); + const float b = bias ? (float)bias[oc] : 0.f; + + std::vector out_f32(out_len, b); + std::vector input_f32(L); + std::vector weight_f32(K); + + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Xc = Xb + ic * L; + const __fp16* Wc = Woc + ic * K; + + for (size_t i = 0; i < L; ++i) input_f32[i] = (float)Xc[i]; + for (size_t k = 0; k < K; ++k) weight_f32[k] = (float)Wc[k]; + + if (stride == 1) { + std::vector conv_out(out_len); + vDSP_conv(input_f32.data(), 1, weight_f32.data(), 1, + conv_out.data(), 1, out_len, K); + vDSP_vadd(out_f32.data(), 1, conv_out.data(), 1, + out_f32.data(), 1, out_len); + } else { + std::vector full_conv(L - K + 1); + vDSP_conv(input_f32.data(), 1, weight_f32.data(), 1, + full_conv.data(), 1, L - K + 1, K); + for (size_t out_t = 0; out_t < out_len; ++out_t) { + out_f32[out_t] += full_conv[out_t * stride]; + } + } + } + + for (size_t out_t = 0; out_t < out_len; ++out_t) { + Yoc[out_t] = (__fp16)out_f32[out_t]; + } + }); +} +#endif + +static void conv1d_f16_neon( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, size_t L, + size_t C_in, size_t C_out, + size_t K, + size_t stride +){ + const size_t out_len = ((L - K) / stride) + 1; + const size_t in_bs = C_in * L; + const size_t out_bs = C_out * out_len; + + const size_t total_compute = N * C_out * out_len * C_in * K; + CactusThreading::ParallelConfig config = (total_compute < 100000) + ? CactusThreading::ParallelConfig{SIZE_MAX, SIZE_MAX} + : CactusThreading::Thresholds::ATTENTION; + + CactusThreading::parallel_for_2d( + N, C_out, config, + [&](size_t n, size_t oc) { + + const __fp16* Xb = input + n * in_bs; + __fp16* Yoc = output + n * out_bs + oc * out_len; + const __fp16* Woc = weight + oc * (C_in * K); + const float b = bias ? (float)bias[oc] : 0.f; + + for (size_t out_t = 0; out_t < out_len; ++out_t) { + const size_t t = out_t * stride; + float sum = b; + + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Xc = Xb + ic * L + t; + const __fp16* Wc = Woc + ic * K; + + float32x4_t acc0 = vdupq_n_f32(0.f); + float32x4_t acc1 = vdupq_n_f32(0.f); + + size_t k = 0; + + for (; k + 8 <= K; k += 8) { + const float16x8_t xv = vld1q_f16(Xc + k); + const float16x8_t wv = vld1q_f16(Wc + k); + + acc0 = vfmaq_f32(acc0, + vcvt_f32_f16(vget_low_f16(xv)), + vcvt_f32_f16(vget_low_f16(wv))); + acc1 = vfmaq_f32(acc1, + vcvt_f32_f16(vget_high_f16(xv)), + vcvt_f32_f16(vget_high_f16(wv))); + } + + float32x2_t s1 = vadd_f32(vget_low_f32(acc0), vget_high_f32(acc0)); + sum += vget_lane_f32(s1, 0) + vget_lane_f32(s1, 1); + + float32x2_t s2 = vadd_f32(vget_low_f32(acc1), vget_high_f32(acc1)); + sum += vget_lane_f32(s2, 0) + vget_lane_f32(s2, 1); + + for (; k < K; ++k) { + sum += (float)Xc[k] * (float)Wc[k]; + } + } + + Yoc[out_t] = (__fp16)sum; + } + }); +} + +#ifdef __APPLE__ +static void conv1d_f16_gemm( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, size_t L, + size_t C_in, size_t C_out, + size_t K, + size_t stride +) { + const size_t out_len = (L - K) / stride + 1; + const size_t col_K = C_in * K; + + std::vector W_f32(C_out * col_K); + for (size_t i = 0; i < C_out * col_K; ++i) + W_f32[i] = static_cast(weight[i]); + + std::vector bias_f32; + if (bias) { + bias_f32.resize(C_out); + for (size_t i = 0; i < C_out; ++i) + bias_f32[i] = static_cast(bias[i]); + } + + std::vector col(col_K * out_len); + std::vector Y_f32(C_out * out_len); + + for (size_t n = 0; n < N; ++n) { + const __fp16* Xn = input + n * C_in * L; + __fp16* Yn = output + n * C_out * out_len; + + if (stride == 1) { + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Xc = Xn + ic * L; + for (size_t k = 0; k < K; ++k) { + float* dst = col.data() + (ic * K + k) * out_len; + const __fp16* src = Xc + k; + size_t t = 0; + for (; t + 8 <= out_len; t += 8) { + float16x8_t v = vld1q_f16(src + t); + vst1q_f32(dst + t, vcvt_f32_f16(vget_low_f16(v))); + vst1q_f32(dst + t + 4, vcvt_f32_f16(vget_high_f16(v))); + } + for (; t < out_len; ++t) + dst[t] = static_cast(src[t]); + } + } + } else { + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Xc = Xn + ic * L; + for (size_t k = 0; k < K; ++k) { + float* dst = col.data() + (ic * K + k) * out_len; + for (size_t t = 0; t < out_len; ++t) + dst[t] = static_cast(Xc[t * stride + k]); + } + } + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + static_cast(C_out), static_cast(out_len), static_cast(col_K), + 1.0f, W_f32.data(), static_cast(col_K), + col.data(), static_cast(out_len), + 0.0f, Y_f32.data(), static_cast(out_len)); + + if (bias) { + for (size_t oc = 0; oc < C_out; ++oc) { + float b = bias_f32[oc]; + const float* src = Y_f32.data() + oc * out_len; + __fp16* dst = Yn + oc * out_len; + size_t t = 0; + for (; t + 8 <= out_len; t += 8) { + float32x4_t v0 = vaddq_f32(vld1q_f32(src + t), vdupq_n_f32(b)); + float32x4_t v1 = vaddq_f32(vld1q_f32(src + t + 4), vdupq_n_f32(b)); + vst1q_f16(dst + t, vcombine_f16(vcvt_f16_f32(v0), vcvt_f16_f32(v1))); + } + for (; t < out_len; ++t) + dst[t] = static_cast<__fp16>(src[t] + b); + } + } else { + for (size_t oc = 0; oc < C_out; ++oc) { + const float* src = Y_f32.data() + oc * out_len; + __fp16* dst = Yn + oc * out_len; + size_t t = 0; + for (; t + 8 <= out_len; t += 8) { + float32x4_t v0 = vld1q_f32(src + t); + float32x4_t v1 = vld1q_f32(src + t + 4); + vst1q_f16(dst + t, vcombine_f16(vcvt_f16_f32(v0), vcvt_f16_f32(v1))); + } + for (; t < out_len; ++t) + dst[t] = static_cast<__fp16>(src[t]); + } + } + } +} +#endif + +void cactus_conv1d_f16( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, size_t L, + size_t C_in, size_t C_out, + size_t K, + size_t stride +){ +#ifdef __APPLE__ + if (stride > 1 || K < ACCELERATE_K_THRESHOLD) { + conv1d_f16_gemm(input, weight, bias, output, N, L, C_in, C_out, K, stride); + return; + } + if (K >= ACCELERATE_K_THRESHOLD && L >= ACCELERATE_L_THRESHOLD) { + conv1d_f16_accelerate(input, weight, bias, output, N, L, C_in, C_out, K, stride); + return; + } +#endif + conv1d_f16_neon(input, weight, bias, output, N, L, C_in, C_out, K, stride); +} + +inline void conv1d_k7s3_oc8_t4( + const __fp16* Xb, + const __fp16* Wpack, + const __fp16* bias, + __fp16* Yb, + size_t L, + size_t out_len, + size_t C_in, + size_t C_out, + size_t out_t, + size_t oc0 +){ + float32x4_t acc[4][2]; + float32x4_t b0 = vdupq_n_f32(0.f); + float32x4_t b1 = vdupq_n_f32(0.f); + if (bias) { + float16x8_t bv = vld1q_f16(bias + oc0); + b0 = vcvt_f32_f16(vget_low_f16(bv)); + b1 = vcvt_f32_f16(vget_high_f16(bv)); + } + + for (int j = 0; j < 4; ++j) { + acc[j][0] = b0; + acc[j][1] = b1; + } + + const size_t t_base = out_t * 3; + + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Wic = Wpack + (ic * 7) * C_out + oc0; + const __fp16* Xic = Xb + ic * L + t_base; + + for (int k = 0; k < 7; ++k) { + float16x8_t w_half = vld1q_f16(Wic + k * C_out); + float32x4_t w0 = vcvt_f32_f16(vget_low_f16(w_half)); + float32x4_t w1 = vcvt_f32_f16(vget_high_f16(w_half)); + for (int j = 0; j < 4; ++j) { + float x_val = (float)Xic[j * 3 + k]; + float32x4_t xv = vdupq_n_f32(x_val); + + acc[j][0] = vfmaq_f32(acc[j][0], w0, xv); + acc[j][1] = vfmaq_f32(acc[j][1], w1, xv); + } + } + } + + float tmp0[4], tmp1[4]; + for(int j=0; j<4; ++j) { + vst1q_f32(tmp0, acc[j][0]); + vst1q_f32(tmp1, acc[j][1]); + + for(int i=0; i<4; ++i) { + Yb[(oc0 + i) * out_len + out_t + j] = (__fp16)tmp0[i]; + } + for(int i=0; i<4; ++i) { + Yb[(oc0 + 4 + i) * out_len + out_t + j] = (__fp16)tmp1[i]; + } + } +} + +inline void conv1d_k7s3_oc8_scalar( + const __fp16* Xb, + const __fp16* Wpack, + const __fp16* bias, + __fp16* Yb, + size_t L, + size_t out_len, + size_t C_in, + size_t C_out, + size_t out_t, + size_t oc0 +){ + float32x4_t acc0 = vdupq_n_f32(0.f); + float32x4_t acc1 = vdupq_n_f32(0.f); + + if (bias) { + float16x8_t bv = vld1q_f16(bias + oc0); + acc0 = vcvt_f32_f16(vget_low_f16(bv)); + acc1 = vcvt_f32_f16(vget_high_f16(bv)); + } + + const size_t t_base = out_t * 3; + + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Wic = Wpack + (ic * 7) * C_out + oc0; + const __fp16* Xic = Xb + ic * L + t_base; + for (int k = 0; k < 7; ++k) { + float x_val = (float)Xic[k]; + float32x4_t xv = vdupq_n_f32(x_val); + + float16x8_t w_half = vld1q_f16(Wic + k * C_out); + acc0 = vfmaq_f32(acc0, vcvt_f32_f16(vget_low_f16(w_half)), xv); + acc1 = vfmaq_f32(acc1, vcvt_f32_f16(vget_high_f16(w_half)), xv); + } + } + + float tmp[8]; + vst1q_f32(tmp, acc0); + vst1q_f32(tmp+4, acc1); + + for(int i=0; i<8; ++i) { + if (oc0 + i < C_out) { + Yb[(oc0 + i) * out_len + out_t] = (__fp16)tmp[i]; + } + } +} + +void cactus_conv1d_f16_k7s3_oc8( + const __fp16* input, + const __fp16* Wpack, + const __fp16* bias, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out) +{ + if (L < 7) return; + size_t out_len = (L - 7) / 3 + 1; + size_t num_oc_blocks = (C_out + 7) / 8; + + CactusThreading::parallel_for_2d(N, num_oc_blocks, CactusThreading::Thresholds::SCALAR_EXPENSIVE, [=](size_t n, size_t ob) { + size_t oc0 = ob * 8; + const __fp16* Xb = input + n * (C_in * L); + __fp16* Yb = output + n * (C_out * out_len); + + size_t out_t = 0; + for (; out_t + 4 <= out_len; out_t += 4) { + conv1d_k7s3_oc8_t4(Xb, Wpack, bias, Yb, L, out_len, C_in, C_out, out_t, oc0); + } + + for (; out_t < out_len; ++out_t) { + conv1d_k7s3_oc8_scalar(Xb, Wpack, bias, Yb, L, out_len, C_in, C_out, out_t, oc0); + } + }); +} + + + +void cactus_bilinear_interpolation_f16(const __fp16* input, __fp16* output, size_t src_height, size_t src_width, size_t embed_dim, + size_t dst_height, size_t dst_width, bool align_corners) +{ + float scale_h, scale_w; + if (align_corners) { + scale_h = (src_height > 1 && dst_height > 1) + ? static_cast(src_height - 1) / static_cast(dst_height - 1) + : 0.0f; + scale_w = (src_width > 1 && dst_width > 1) + ? static_cast(src_width - 1) / static_cast(dst_width - 1) + : 0.0f; + } else { + scale_h = (dst_height > 0) + ? static_cast(src_height) / static_cast(dst_height) + : 0.0f; + scale_w = (dst_width > 0) + ? static_cast(src_width) / static_cast(dst_width) + : 0.0f; + } + + for (size_t dst_y = 0; dst_y < dst_height; ++dst_y) { + for (size_t dst_x = 0; dst_x < dst_width; ++dst_x) { + float src_y_float, src_x_float; + if (align_corners) { + src_y_float = static_cast(dst_y) * scale_h; + src_x_float = static_cast(dst_x) * scale_w; + } else { + src_y_float = (static_cast(dst_y) + 0.5f) * scale_h - 0.5f; + src_x_float = (static_cast(dst_x) + 0.5f) * scale_w - 0.5f; + if (src_y_float < 0.0f) src_y_float = 0.0f; + if (src_x_float < 0.0f) src_x_float = 0.0f; + const float max_y = static_cast(src_height) - 1.0f; + const float max_x = static_cast(src_width) - 1.0f; + if (src_y_float > max_y) src_y_float = max_y; + if (src_x_float > max_x) src_x_float = max_x; + } + + int y0 = static_cast(std::floor(src_y_float)); + int x0 = static_cast(std::floor(src_x_float)); + + int y1 = ((y0 + 1) < static_cast(src_height)) ? (y0 + 1) : (static_cast(src_height) - 1); + int x1 = ((x0 + 1) < static_cast(src_width)) ? (x0 + 1) : (static_cast(src_width) - 1); + + float dy = src_y_float - y0; + float dx = src_x_float - x0; + + float w00 = (1.0f - dx) * (1.0f - dy); + float w01 = dx * (1.0f - dy); + float w10 = (1.0f - dx) * dy; + float w11 = dx * dy; + + size_t idx00 = (y0 * static_cast(src_width) + x0) * static_cast(embed_dim); + size_t idx01 = (y0 * static_cast(src_width) + x1) * static_cast(embed_dim); + size_t idx10 = (y1 * static_cast(src_width) + x0) * static_cast(embed_dim); + size_t idx11 = (y1 * static_cast(src_width) + x1) * static_cast(embed_dim); + + size_t out_idx = (dst_y * dst_width + dst_x) * embed_dim; + + size_t d = 0; + for(; d + 8 <= embed_dim; d += 8){ + float16x8_t v00 = vld1q_f16(input + idx00 + d); + float16x8_t v01 = vld1q_f16(input + idx01 + d); + float16x8_t v10 = vld1q_f16(input + idx10 + d); + float16x8_t v11 = vld1q_f16(input + idx11 + d); + + float16x8_t result = vfmaq_f16( + vfmaq_f16( + vfmaq_f16(vdupq_n_f16(0.0f), v00, vdupq_n_f16(w00)), + v01, vdupq_n_f16(w01)), + v10, vdupq_n_f16(w10)); + result = vfmaq_f16(result, v11, vdupq_n_f16(w11)); + + vst1q_f16(output + out_idx + d, result); + } + + for (; d < embed_dim; ++d) { + output[out_idx + d] = + input[idx00 + d] * w00 + + input[idx01 + d] * w01 + + input[idx10 + d] * w10 + + input[idx11 + d] * w11; + } + } + } +} + +void cactus_stft_f16( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t N, size_t L, + size_t C_in, size_t /*C_out*/, + size_t K, size_t stride, + size_t num_fft_bins +) { + const size_t out_len = ((L - K) / stride) + 1; + const size_t in_bs = C_in * L; + const size_t out_bs = 2 * num_fft_bins * out_len; + + for (size_t n = 0; n < N; ++n) { + const __fp16* Xb = input + n * in_bs; + + for (size_t bin = 0; bin < num_fft_bins; ++bin) { + const __fp16* Wr = weight + bin * (C_in * K); + const __fp16* Wi = weight + (bin + num_fft_bins) * (C_in * K); + + for (size_t out_t = 0; out_t < out_len; ++out_t) { + const size_t t = out_t * stride; + float sum_real = 0.0f; + float sum_imag = 0.0f; + + for (size_t ic = 0; ic < C_in; ++ic) { + const __fp16* Xc = Xb + ic * L + t; + const __fp16* Wrc = Wr + ic * K; + const __fp16* Wic = Wi + ic * K; + + float32x4_t acc_r0 = vdupq_n_f32(0.f); + float32x4_t acc_r1 = vdupq_n_f32(0.f); + float32x4_t acc_i0 = vdupq_n_f32(0.f); + float32x4_t acc_i1 = vdupq_n_f32(0.f); + + size_t k = 0; + for (; k + 8 <= K; k += 8) { + const float16x8_t xv = vld1q_f16(Xc + k); + const float16x8_t wr = vld1q_f16(Wrc + k); + const float16x8_t wi = vld1q_f16(Wic + k); + + acc_r0 = vfmaq_f32(acc_r0, vcvt_f32_f16(vget_low_f16(xv)), vcvt_f32_f16(vget_low_f16(wr))); + acc_r1 = vfmaq_f32(acc_r1, vcvt_f32_f16(vget_high_f16(xv)), vcvt_f32_f16(vget_high_f16(wr))); + acc_i0 = vfmaq_f32(acc_i0, vcvt_f32_f16(vget_low_f16(xv)), vcvt_f32_f16(vget_low_f16(wi))); + acc_i1 = vfmaq_f32(acc_i1, vcvt_f32_f16(vget_high_f16(xv)), vcvt_f32_f16(vget_high_f16(wi))); + } + + sum_real += vaddvq_f32(acc_r0) + vaddvq_f32(acc_r1); + sum_imag += vaddvq_f32(acc_i0) + vaddvq_f32(acc_i1); + + for (; k < K; ++k) { + float x = (float)Xc[k]; + sum_real += x * (float)Wrc[k]; + sum_imag += x * (float)Wic[k]; + } + } + + output[n * out_bs + bin * out_len + out_t] = (__fp16)sum_real; + output[n * out_bs + (bin + num_fft_bins) * out_len + out_t] = (__fp16)sum_imag; + } + } + } +} + +void cactus_conv1d_same_depthwise_f16_k9( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, size_t L, size_t C) +{ + constexpr int K = 9; + + const size_t bs = L * C; + + CactusThreading::parallel_for_2d( + N, C, CactusThreading::Thresholds::ATTENTION, + [&](size_t n, size_t c) { + + const __fp16* Xb = input + n * bs; + __fp16* Yb = output + n * bs; + + const __fp16* Wc = weight + c * K; + const float b = bias ? (float)bias[c] : 0.0f; + + const float w0 = (float)Wc[0]; + const float w1 = (float)Wc[1]; + const float w2 = (float)Wc[2]; + const float w3 = (float)Wc[3]; + const float w4 = (float)Wc[4]; + const float w5 = (float)Wc[5]; + const float w6 = (float)Wc[6]; + const float w7 = (float)Wc[7]; + const float w8 = (float)Wc[8]; + + const float32x4_t wv0 = {w0, w1, w2, w3}; + const float32x4_t wv1 = {w4, w5, w6, w7}; + + for (size_t t0 = 0; t0 < L; t0 += T_TILE_F16) { + const bool have_t1 = (t0 + 1 < L); + const size_t t1 = have_t1 ? (t0 + 1) : t0; + + float x0_0=0, x1_0=0, x2_0=0, x3_0=0; + float x4_0=0, x5_0=0, x6_0=0, x7_0=0; + float x8_0=0; + + { + const ptrdiff_t i0 = (ptrdiff_t)t0 - 4; + const ptrdiff_t i1 = (ptrdiff_t)t0 - 3; + const ptrdiff_t i2 = (ptrdiff_t)t0 - 2; + const ptrdiff_t i3 = (ptrdiff_t)t0 - 1; + if (i0 >= 0 && i0 < (ptrdiff_t)L) x0_0 = (float)Xb[(size_t)i0 * C + c]; + if (i1 >= 0 && i1 < (ptrdiff_t)L) x1_0 = (float)Xb[(size_t)i1 * C + c]; + if (i2 >= 0 && i2 < (ptrdiff_t)L) x2_0 = (float)Xb[(size_t)i2 * C + c]; + if (i3 >= 0 && i3 < (ptrdiff_t)L) x3_0 = (float)Xb[(size_t)i3 * C + c]; + + const ptrdiff_t j4 = (ptrdiff_t)t0 + 0; + const ptrdiff_t j5 = (ptrdiff_t)t0 + 1; + const ptrdiff_t j6 = (ptrdiff_t)t0 + 2; + const ptrdiff_t j7 = (ptrdiff_t)t0 + 3; + const ptrdiff_t j8 = (ptrdiff_t)t0 + 4; + if (j4 >= 0 && j4 < (ptrdiff_t)L) x4_0 = (float)Xb[(size_t)j4 * C + c]; + if (j5 >= 0 && j5 < (ptrdiff_t)L) x5_0 = (float)Xb[(size_t)j5 * C + c]; + if (j6 >= 0 && j6 < (ptrdiff_t)L) x6_0 = (float)Xb[(size_t)j6 * C + c]; + if (j7 >= 0 && j7 < (ptrdiff_t)L) x7_0 = (float)Xb[(size_t)j7 * C + c]; + if (j8 >= 0 && j8 < (ptrdiff_t)L) x8_0 = (float)Xb[(size_t)j8 * C + c]; + } + + float32x4_t acc0v = vdupq_n_f32(0.f); + const float32x4_t xv0_0 = {x0_0, x1_0, x2_0, x3_0}; + const float32x4_t xv1_0 = {x4_0, x5_0, x6_0, x7_0}; + acc0v = vfmaq_f32(acc0v, xv0_0, wv0); + acc0v = vfmaq_f32(acc0v, xv1_0, wv1); + float acc0 = vaddvq_f32(acc0v) + (w8 * x8_0) + b; + + float acc1 = 0.f; + if (have_t1) { + float x0_1=0, x1_1=0, x2_1=0, x3_1=0; + float x4_1=0, x5_1=0, x6_1=0, x7_1=0; + float x8_1=0; + + { + const ptrdiff_t i0 = (ptrdiff_t)t1 - 4; + const ptrdiff_t i1 = (ptrdiff_t)t1 - 3; + const ptrdiff_t i2 = (ptrdiff_t)t1 - 2; + const ptrdiff_t i3 = (ptrdiff_t)t1 - 1; + if (i0 >= 0 && i0 < (ptrdiff_t)L) x0_1 = (float)Xb[(size_t)i0 * C + c]; + if (i1 >= 0 && i1 < (ptrdiff_t)L) x1_1 = (float)Xb[(size_t)i1 * C + c]; + if (i2 >= 0 && i2 < (ptrdiff_t)L) x2_1 = (float)Xb[(size_t)i2 * C + c]; + if (i3 >= 0 && i3 < (ptrdiff_t)L) x3_1 = (float)Xb[(size_t)i3 * C + c]; + + const ptrdiff_t j4 = (ptrdiff_t)t1 + 0; + const ptrdiff_t j5 = (ptrdiff_t)t1 + 1; + const ptrdiff_t j6 = (ptrdiff_t)t1 + 2; + const ptrdiff_t j7 = (ptrdiff_t)t1 + 3; + const ptrdiff_t j8 = (ptrdiff_t)t1 + 4; + if (j4 >= 0 && j4 < (ptrdiff_t)L) x4_1 = (float)Xb[(size_t)j4 * C + c]; + if (j5 >= 0 && j5 < (ptrdiff_t)L) x5_1 = (float)Xb[(size_t)j5 * C + c]; + if (j6 >= 0 && j6 < (ptrdiff_t)L) x6_1 = (float)Xb[(size_t)j6 * C + c]; + if (j7 >= 0 && j7 < (ptrdiff_t)L) x7_1 = (float)Xb[(size_t)j7 * C + c]; + if (j8 >= 0 && j8 < (ptrdiff_t)L) x8_1 = (float)Xb[(size_t)j8 * C + c]; + } + + float32x4_t acc1v = vdupq_n_f32(0.f); + const float32x4_t xv0_1 = {x0_1, x1_1, x2_1, x3_1}; + const float32x4_t xv1_1 = {x4_1, x5_1, x6_1, x7_1}; + acc1v = vfmaq_f32(acc1v, xv0_1, wv0); + acc1v = vfmaq_f32(acc1v, xv1_1, wv1); + acc1 = vaddvq_f32(acc1v) + (w8 * x8_1) + b; + } + + Yb[t0 * C + c] = (__fp16)acc0; + if (have_t1) Yb[t1 * C + c] = (__fp16)acc1; + } + }); +} + + diff --git a/cactus-kernels/src/conv2d.cpp b/cactus-kernels/src/conv2d.cpp new file mode 100644 index 000000000..768d85b1c --- /dev/null +++ b/cactus-kernels/src/conv2d.cpp @@ -0,0 +1,619 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include + +static void conv2d_k3_weights_to_f32( + const __fp16* weight, float* W_f32, + size_t C_out, size_t C_in +) { + const size_t col_K = C_in * 9; + for (size_t oc = 0; oc < C_out; ++oc) + for (size_t ic = 0; ic < C_in; ++ic) + for (size_t kh = 0; kh < 3; ++kh) + for (size_t kw = 0; kw < 3; ++kw) + W_f32[oc * col_K + ic * 9 + kh * 3 + kw] = + static_cast(weight[((oc * C_in + ic) * 3 + kh) * 3 + kw]); +} + +static void conv2d_gemm_bias_convert( + const float* W_f32, const float* col, const float* bias_f32, + __fp16* Yn, size_t C_out, size_t col_K, size_t spatial +) { + std::vector Y_f32(C_out * spatial); + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + static_cast(C_out), static_cast(spatial), static_cast(col_K), + 1.0f, W_f32, static_cast(col_K), + col, static_cast(spatial), + 0.0f, Y_f32.data(), static_cast(spatial)); + + for (size_t oc = 0; oc < C_out; ++oc) { + const float b = bias_f32[oc]; + const float* src = Y_f32.data() + oc * spatial; + __fp16* dst = Yn + oc * spatial; + size_t i = 0; + for (; i + 8 <= spatial; i += 8) { + float32x4_t v0 = vaddq_f32(vld1q_f32(src + i), vdupq_n_f32(b)); + float32x4_t v1 = vaddq_f32(vld1q_f32(src + i + 4), vdupq_n_f32(b)); + vst1q_f16(dst + i, vcombine_f16(vcvt_f16_f32(v0), vcvt_f16_f32(v1))); + } + for (; i < spatial; ++i) + dst[i] = static_cast<__fp16>(src[i] + b); + } +} +#endif + +void cactus_conv2d_f16_k3s2p1_nchw( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C_in, size_t H, size_t W, + size_t C_out +){ + if (H + 2 < 3 || W + 2 < 3) return; + const size_t H_out = (H - 1) / 2 + 1; + const size_t W_out = (W - 1) / 2 + 1; + +#ifdef __APPLE__ + { + const size_t spatial = H_out * W_out; + const size_t col_K = C_in * 9; + std::vector W_f32(C_out * col_K); + conv2d_k3_weights_to_f32(weight, W_f32.data(), C_out, C_in); + + std::vector bias_f32(C_out, 0.0f); + if (bias) + for (size_t i = 0; i < C_out; ++i) + bias_f32[i] = static_cast(bias[i]); + + std::vector col(col_K * spatial); + + for (size_t n = 0; n < N; ++n) { + const __fp16* Xn = input + n * C_in * H * W; + __fp16* Yn = output + n * C_out * spatial; + + for (size_t ic = 0; ic < C_in; ++ic) { + for (size_t kh = 0; kh < 3; ++kh) { + for (size_t kw = 0; kw < 3; ++kw) { + float* dst = col.data() + (ic * 9 + kh * 3 + kw) * spatial; + for (size_t oh = 0; oh < H_out; ++oh) { + ptrdiff_t ih = static_cast(oh) * 2 + static_cast(kh) - 1; + float* dst_row = dst + oh * W_out; + if (ih < 0 || ih >= static_cast(H)) { + memset(dst_row, 0, W_out * sizeof(float)); + continue; + } + const __fp16* src_row = Xn + ic * H * W + static_cast(ih) * W; + const ptrdiff_t iw_offset = static_cast(kw) - 1; + for (size_t ow = 0; ow < W_out; ++ow) { + ptrdiff_t iw = static_cast(ow) * 2 + iw_offset; + dst_row[ow] = (iw < 0 || iw >= static_cast(W)) + ? 0.0f : static_cast(src_row[iw]); + } + } + } + } + } + + conv2d_gemm_bias_convert(W_f32.data(), col.data(), bias_f32.data(), Yn, C_out, col_K, spatial); + } + return; + } +#endif + + auto in_idx = [&](size_t n, size_t ic, size_t ih, size_t iw) -> size_t { + return ((n * C_in + ic) * H + ih) * W + iw; + }; + auto w_idx = [&](size_t oc, size_t ic, size_t kh, size_t kw) -> size_t { + return (((oc * C_in + ic) * 3 + kh) * 3 + kw); + }; + auto out_idx = [&](size_t n, size_t oc, size_t oh, size_t ow) -> size_t { + return ((n * C_out + oc) * H_out + oh) * W_out + ow; + }; + + constexpr size_t OW_VEC = 4; + + const size_t total_compute = N * C_out * H_out * W_out * C_in * 9; + CactusThreading::ParallelConfig cfg = + (total_compute < 100000) ? CactusThreading::ParallelConfig{SIZE_MAX, SIZE_MAX} + : CactusThreading::Thresholds::ATTENTION; + + CactusThreading::parallel_for_2d(N, C_out, cfg, [&](size_t n, size_t oc) { + const float b0 = bias ? (float)bias[oc] : 0.0f; + + for (size_t oh = 0; oh < H_out; ++oh) { + const ptrdiff_t ih0 = (ptrdiff_t)oh * 2 - 1; + + const bool h_interior = (ih0 >= 0) && ((size_t)(ih0 + 2) < H); + + for (size_t ow = 0; ow < W_out; ) { + const size_t rem = W_out - ow; + + const bool do_vec = (rem >= OW_VEC); + const bool w_interior_vec = + do_vec && + (ow >= 1) && + (((ow + (OW_VEC - 1)) * 2 + 1) < W); + + float32x4_t vacc = vdupq_n_f32(b0); + + if (h_interior && w_interior_vec) { + for (size_t ic = 0; ic < C_in; ++ic) { + + const __fp16* row0 = input + in_idx(n, ic, (size_t)(ih0 + 0), 0); + const __fp16* row1 = input + in_idx(n, ic, (size_t)(ih0 + 1), 0); + const __fp16* row2 = input + in_idx(n, ic, (size_t)(ih0 + 2), 0); + const ptrdiff_t iw_base0 = (ptrdiff_t)ow * 2 - 1; + { + const float w00 = (float)weight[w_idx(oc, ic, 0, 0)]; + const float w01 = (float)weight[w_idx(oc, ic, 0, 1)]; + const float w02 = (float)weight[w_idx(oc, ic, 0, 2)]; + + float x0 = (float)row0[(size_t)(iw_base0 + 0)]; + float x1 = (float)row0[(size_t)(iw_base0 + 2)]; + float x2 = (float)row0[(size_t)(iw_base0 + 4)]; + float x3 = (float)row0[(size_t)(iw_base0 + 6)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w00)); + + x0 = (float)row0[(size_t)(iw_base0 + 1)]; + x1 = (float)row0[(size_t)(iw_base0 + 3)]; + x2 = (float)row0[(size_t)(iw_base0 + 5)]; + x3 = (float)row0[(size_t)(iw_base0 + 7)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w01)); + + x0 = (float)row0[(size_t)(iw_base0 + 2)]; + x1 = (float)row0[(size_t)(iw_base0 + 4)]; + x2 = (float)row0[(size_t)(iw_base0 + 6)]; + x3 = (float)row0[(size_t)(iw_base0 + 8)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w02)); + } + + { + const float w10 = (float)weight[w_idx(oc, ic, 1, 0)]; + const float w11 = (float)weight[w_idx(oc, ic, 1, 1)]; + const float w12 = (float)weight[w_idx(oc, ic, 1, 2)]; + const ptrdiff_t iw_base0 = (ptrdiff_t)ow * 2 - 1; + + float x0 = (float)row1[(size_t)(iw_base0 + 0)]; + float x1 = (float)row1[(size_t)(iw_base0 + 2)]; + float x2 = (float)row1[(size_t)(iw_base0 + 4)]; + float x3 = (float)row1[(size_t)(iw_base0 + 6)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w10)); + + x0 = (float)row1[(size_t)(iw_base0 + 1)]; + x1 = (float)row1[(size_t)(iw_base0 + 3)]; + x2 = (float)row1[(size_t)(iw_base0 + 5)]; + x3 = (float)row1[(size_t)(iw_base0 + 7)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w11)); + + x0 = (float)row1[(size_t)(iw_base0 + 2)]; + x1 = (float)row1[(size_t)(iw_base0 + 4)]; + x2 = (float)row1[(size_t)(iw_base0 + 6)]; + x3 = (float)row1[(size_t)(iw_base0 + 8)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w12)); + } + + { + const float w20 = (float)weight[w_idx(oc, ic, 2, 0)]; + const float w21 = (float)weight[w_idx(oc, ic, 2, 1)]; + const float w22 = (float)weight[w_idx(oc, ic, 2, 2)]; + const ptrdiff_t iw_base0 = (ptrdiff_t)ow * 2 - 1; + + float x0 = (float)row2[(size_t)(iw_base0 + 0)]; + float x1 = (float)row2[(size_t)(iw_base0 + 2)]; + float x2 = (float)row2[(size_t)(iw_base0 + 4)]; + float x3 = (float)row2[(size_t)(iw_base0 + 6)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w20)); + + x0 = (float)row2[(size_t)(iw_base0 + 1)]; + x1 = (float)row2[(size_t)(iw_base0 + 3)]; + x2 = (float)row2[(size_t)(iw_base0 + 5)]; + x3 = (float)row2[(size_t)(iw_base0 + 7)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w21)); + + x0 = (float)row2[(size_t)(iw_base0 + 2)]; + x1 = (float)row2[(size_t)(iw_base0 + 4)]; + x2 = (float)row2[(size_t)(iw_base0 + 6)]; + x3 = (float)row2[(size_t)(iw_base0 + 8)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0,x1,x2,x3}, vdupq_n_f32(w22)); + } + } + + __fp16* Y = output + out_idx(n, oc, oh, ow); + const float16x4_t yv = vcvt_f16_f32(vacc); + vst1_f16(Y, yv); + + ow += OW_VEC; + continue; + } + + const size_t tile = do_vec ? OW_VEC : 1; + float acc_s[OW_VEC] = {b0, b0, b0, b0}; + + for (size_t ic = 0; ic < C_in; ++ic) { + for (size_t kh = 0; kh < 3; ++kh) { + const ptrdiff_t ih = ih0 + (ptrdiff_t)kh; + if (ih < 0 || ih >= (ptrdiff_t)H) continue; + + const __fp16* row = input + in_idx(n, ic, (size_t)ih, 0); + + for (size_t kw = 0; kw < 3; ++kw) { + const float wv = (float)weight[w_idx(oc, ic, kh, kw)]; + for (size_t t = 0; t < tile; ++t) { + const size_t ow_t = ow + t; + const ptrdiff_t iw = (ptrdiff_t)ow_t * 2 - 1 + (ptrdiff_t)kw; + if (iw < 0 || iw >= (ptrdiff_t)W) continue; + acc_s[t] += (float)row[(size_t)iw] * wv; + } + } + } + } + + __fp16* Y = output + out_idx(n, oc, oh, ow); + for (size_t t = 0; t < tile; ++t) { + Y[t] = (__fp16)acc_s[t]; + } + ow += tile; + } + } + }); +} + +void cactus_conv2d_depthwise_f16_k3s2p1_nchw( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C, + size_t H, + size_t W +){ + if (H + 2 < 3 || W + 2 < 3) return; + const size_t H_out = (H - 1) / 2 + 1; + const size_t W_out = (W - 1) / 2 + 1; + + auto in_idx = [&](size_t n, size_t c, size_t ih, size_t iw) -> size_t { + return ((n * C + c) * H + ih) * W + iw; + }; + auto out_idx = [&](size_t n, size_t c, size_t oh, size_t ow) -> size_t { + return ((n * C + c) * H_out + oh) * W_out + ow; + }; + + constexpr size_t OW_VEC = 4; + const size_t total_compute = N * C * H_out * W_out * 9; + CactusThreading::ParallelConfig cfg = + (total_compute < 100000) ? CactusThreading::ParallelConfig{SIZE_MAX, SIZE_MAX} + : CactusThreading::Thresholds::ATTENTION; + + CactusThreading::parallel_for_2d(N, C, cfg, [&](size_t n, size_t c) { + const __fp16* Wc = weight + c * 9; + const float w00 = (float)Wc[0], w01 = (float)Wc[1], w02 = (float)Wc[2]; + const float w10 = (float)Wc[3], w11 = (float)Wc[4], w12 = (float)Wc[5]; + const float w20 = (float)Wc[6], w21 = (float)Wc[7], w22 = (float)Wc[8]; + const float b0 = bias ? (float)bias[c] : 0.0f; + + for (size_t oh = 0; oh < H_out; ++oh) { + const ptrdiff_t ih0 = (ptrdiff_t)oh * 2 - 1; + const bool h_interior = (ih0 >= 0) && ((size_t)(ih0 + 2) < H); + + for (size_t ow = 0; ow < W_out; ) { + const size_t rem = W_out - ow; + const bool do_vec = (rem >= OW_VEC); + const bool w_interior_vec = + do_vec && + (ow >= 1) && + (((ow + (OW_VEC - 1)) * 2 + 1) < W); + + if (h_interior && w_interior_vec) { + const __fp16* row0 = input + in_idx(n, c, (size_t)(ih0 + 0), 0); + const __fp16* row1 = input + in_idx(n, c, (size_t)(ih0 + 1), 0); + const __fp16* row2 = input + in_idx(n, c, (size_t)(ih0 + 2), 0); + const ptrdiff_t iw_base0 = (ptrdiff_t)ow * 2 - 1; + + float32x4_t vacc = vdupq_n_f32(b0); + float x0, x1, x2, x3; + + x0 = (float)row0[(size_t)(iw_base0 + 0)]; + x1 = (float)row0[(size_t)(iw_base0 + 2)]; + x2 = (float)row0[(size_t)(iw_base0 + 4)]; + x3 = (float)row0[(size_t)(iw_base0 + 6)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w00)); + + x0 = (float)row0[(size_t)(iw_base0 + 1)]; + x1 = (float)row0[(size_t)(iw_base0 + 3)]; + x2 = (float)row0[(size_t)(iw_base0 + 5)]; + x3 = (float)row0[(size_t)(iw_base0 + 7)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w01)); + + x0 = (float)row0[(size_t)(iw_base0 + 2)]; + x1 = (float)row0[(size_t)(iw_base0 + 4)]; + x2 = (float)row0[(size_t)(iw_base0 + 6)]; + x3 = (float)row0[(size_t)(iw_base0 + 8)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w02)); + + x0 = (float)row1[(size_t)(iw_base0 + 0)]; + x1 = (float)row1[(size_t)(iw_base0 + 2)]; + x2 = (float)row1[(size_t)(iw_base0 + 4)]; + x3 = (float)row1[(size_t)(iw_base0 + 6)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w10)); + + x0 = (float)row1[(size_t)(iw_base0 + 1)]; + x1 = (float)row1[(size_t)(iw_base0 + 3)]; + x2 = (float)row1[(size_t)(iw_base0 + 5)]; + x3 = (float)row1[(size_t)(iw_base0 + 7)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w11)); + + x0 = (float)row1[(size_t)(iw_base0 + 2)]; + x1 = (float)row1[(size_t)(iw_base0 + 4)]; + x2 = (float)row1[(size_t)(iw_base0 + 6)]; + x3 = (float)row1[(size_t)(iw_base0 + 8)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w12)); + + x0 = (float)row2[(size_t)(iw_base0 + 0)]; + x1 = (float)row2[(size_t)(iw_base0 + 2)]; + x2 = (float)row2[(size_t)(iw_base0 + 4)]; + x3 = (float)row2[(size_t)(iw_base0 + 6)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w20)); + + x0 = (float)row2[(size_t)(iw_base0 + 1)]; + x1 = (float)row2[(size_t)(iw_base0 + 3)]; + x2 = (float)row2[(size_t)(iw_base0 + 5)]; + x3 = (float)row2[(size_t)(iw_base0 + 7)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w21)); + + x0 = (float)row2[(size_t)(iw_base0 + 2)]; + x1 = (float)row2[(size_t)(iw_base0 + 4)]; + x2 = (float)row2[(size_t)(iw_base0 + 6)]; + x3 = (float)row2[(size_t)(iw_base0 + 8)]; + vacc = vfmaq_f32(vacc, (float32x4_t){x0, x1, x2, x3}, vdupq_n_f32(w22)); + + __fp16* Y = output + out_idx(n, c, oh, ow); + vst1_f16(Y, vcvt_f16_f32(vacc)); + ow += OW_VEC; + continue; + } + + const size_t tile = do_vec ? OW_VEC : 1; + float acc_s[OW_VEC] = {b0, b0, b0, b0}; + for (size_t t = 0; t < tile; ++t) { + const size_t ow_t = ow + t; + for (size_t kh = 0; kh < 3; ++kh) { + const ptrdiff_t ih = ih0 + (ptrdiff_t)kh; + if (ih < 0 || ih >= (ptrdiff_t)H) continue; + + const __fp16* row = input + in_idx(n, c, (size_t)ih, 0); + for (size_t kw = 0; kw < 3; ++kw) { + const ptrdiff_t iw = (ptrdiff_t)ow_t * 2 - 1 + (ptrdiff_t)kw; + if (iw < 0 || iw >= (ptrdiff_t)W) continue; + const float wv = (float)Wc[kh * 3 + kw]; + acc_s[t] += (float)row[(size_t)iw] * wv; + } + } + } + + __fp16* Y = output + out_idx(n, c, oh, ow); + for (size_t t = 0; t < tile; ++t) { + Y[t] = (__fp16)acc_s[t]; + } + ow += tile; + } + } + }); +} + +void cactus_conv2d_pointwise_f16_1x1_nchw_gemm( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C_in, + size_t H, + size_t W, + size_t C_out +){ + const size_t HW = H * W; + if (HW == 0 || N == 0) return; + + CactusThreading::parallel_for( + N, + CactusThreading::Thresholds::ATTENTION, + [&](size_t n_start, size_t n_end) { + std::vector<__fp16> packed_in(HW * C_in); + std::vector<__fp16> packed_out(HW * C_out); + + for (size_t n = n_start; n < n_end; ++n) { + const __fp16* Xn = input + n * C_in * HW; + __fp16* Yn = output + n * C_out * HW; + + for (size_t hw = 0; hw < HW; ++hw) { + const size_t h = hw / W; + const size_t w = hw - h * W; + __fp16* row = packed_in.data() + hw * C_in; + for (size_t ic = 0; ic < C_in; ++ic) { + row[ic] = Xn[(ic * H + h) * W + w]; + } + } + + cactus_matmul_f16( + packed_in.data(), + weight, + packed_out.data(), + HW, + C_in, + C_out + ); + + for (size_t hw = 0; hw < HW; ++hw) { + const size_t h = hw / W; + const size_t w = hw - h * W; + const __fp16* row = packed_out.data() + hw * C_out; + for (size_t oc = 0; oc < C_out; ++oc) { + float outv = (float)row[oc]; + if (bias) outv += (float)bias[oc]; + Yn[(oc * H + h) * W + w] = (__fp16)outv; + } + } + } + }); +} + +void cactus_conv1d_pointwise_f16_gemm( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t L, + size_t C_in, + size_t C_out +){ + const size_t M = N * L; + if (M == 0) return; + + cactus_matmul_f16(input, weight, output, M, C_in, C_out); + + if (bias) { + for (size_t m = 0; m < M; ++m) { + __fp16* row = output + m * C_out; + for (size_t oc = 0; oc < C_out; ++oc) { + row[oc] = (__fp16)((float)row[oc] + (float)bias[oc]); + } + } + } +} + +void cactus_conv2d_f16_k3s1p1_nchw( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t N, + size_t C_in, size_t H, size_t W, + size_t C_out +) { + const size_t H_out = H; + const size_t W_out = W; + +#ifdef __APPLE__ + const size_t col_K = C_in * 9; + std::vector W_f32(C_out * col_K); + conv2d_k3_weights_to_f32(weight, W_f32.data(), C_out, C_in); + + std::vector bias_f32(C_out, 0.0f); + if (bias) + for (size_t i = 0; i < C_out; ++i) + bias_f32[i] = static_cast(bias[i]); + + const size_t spatial = H_out * W_out; + std::vector col(col_K * spatial); + + for (size_t n = 0; n < N; ++n) { + const __fp16* Xn = input + n * C_in * H * W; + __fp16* Yn = output + n * C_out * spatial; + + for (size_t ic = 0; ic < C_in; ++ic) { + for (size_t kh = 0; kh < 3; ++kh) { + for (size_t kw = 0; kw < 3; ++kw) { + float* dst = col.data() + (ic * 9 + kh * 3 + kw) * spatial; + for (size_t oh = 0; oh < H_out; ++oh) { + ptrdiff_t ih = static_cast(oh) + static_cast(kh) - 1; + float* dst_row = dst + oh * W_out; + if (ih < 0 || ih >= static_cast(H)) { + memset(dst_row, 0, W_out * sizeof(float)); + continue; + } + const __fp16* src_row = Xn + ic * H * W + static_cast(ih) * W; + const ptrdiff_t iw_offset = static_cast(kw) - 1; + size_t ow_start = 0, ow_end = W_out; + if (iw_offset < 0) { dst_row[0] = 0.0f; ow_start = 1; } + if (iw_offset > 0) { dst_row[W_out - 1] = 0.0f; ow_end = W_out - 1; } + size_t ow = ow_start; + for (; ow + 8 <= ow_end; ow += 8) { + float16x8_t v = vld1q_f16(src_row + ow + iw_offset); + vst1q_f32(dst_row + ow, vcvt_f32_f16(vget_low_f16(v))); + vst1q_f32(dst_row + ow + 4, vcvt_f32_f16(vget_high_f16(v))); + } + for (; ow < ow_end; ++ow) + dst_row[ow] = static_cast(src_row[ow + iw_offset]); + } + } + } + } + + conv2d_gemm_bias_convert(W_f32.data(), col.data(), bias_f32.data(), Yn, C_out, col_K, spatial); + } + +#else + const size_t total_compute = N * C_out * H_out * W_out * C_in * 9; + CactusThreading::ParallelConfig cfg = + (total_compute < 100000) ? CactusThreading::ParallelConfig{SIZE_MAX, SIZE_MAX} + : CactusThreading::Thresholds::ATTENTION; + + CactusThreading::parallel_for_2d(N, C_out, cfg, [&](size_t n, size_t oc) { + const float b0 = bias ? static_cast(bias[oc]) : 0.0f; + for (size_t oh = 0; oh < H_out; ++oh) { + for (size_t ow = 0; ow < W_out; ++ow) { + float acc = b0; + for (size_t ic = 0; ic < C_in; ++ic) { + for (size_t kh = 0; kh < 3; ++kh) { + for (size_t kw = 0; kw < 3; ++kw) { + const ptrdiff_t ih = static_cast(oh) + kh - 1; + const ptrdiff_t iw = static_cast(ow) + kw - 1; + if (ih >= 0 && ih < static_cast(H) && + iw >= 0 && iw < static_cast(W)) { + acc += static_cast(input[((n * C_in + ic) * H + ih) * W + iw]) * + static_cast(weight[((oc * C_in + ic) * 3 + kh) * 3 + kw]); + } + } + } + } + output[((n * C_out + oc) * H_out + oh) * W_out + ow] = static_cast<__fp16>(acc); + } + } + }); +#endif +} + +void cactus_maxpool1d_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t channels, + size_t input_length, + size_t kernel_size, + size_t stride +) { + const size_t output_length = (input_length - kernel_size) / stride + 1; + + CactusThreading::parallel_for(batch_size * channels, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start, size_t end) { + for (size_t bc = start; bc < end; ++bc) { + const size_t b = bc / channels; + const size_t c = bc % channels; + + const __fp16* src = input + b * channels * input_length + c * input_length; + __fp16* dst = output + b * channels * output_length + c * output_length; + + for (size_t i = 0; i < output_length; ++i) { + const size_t in_start = i * stride; + float max_val = -std::numeric_limits::infinity(); + for (size_t k = 0; k < kernel_size; ++k) { + float val = static_cast(src[in_start + k]); + if (val > max_val) max_val = val; + } + dst[i] = static_cast<__fp16>(max_val); + } + } + }); +} diff --git a/cactus-kernels/src/dsp.cpp b/cactus-kernels/src/dsp.cpp new file mode 100644 index 000000000..bfe01eb76 --- /dev/null +++ b/cactus-kernels/src/dsp.cpp @@ -0,0 +1,565 @@ +#include "cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +#endif + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +static size_t bit_reverse(size_t x, size_t log2n) { + size_t result = 0; + for (size_t i = 0; i < log2n; i++) { + result = (result << 1) | (x & 1); + x >>= 1; + } + return result; +} + +static bool is_power_of_two(size_t n) { + return n != 0 && (n & (n - 1)) == 0; +} + +static size_t log2_power_of_two(size_t n) { + size_t log2n = 0; + for (size_t temp = n; temp > 1; temp >>= 1) { + log2n++; + } + return log2n; +} + +static size_t next_power_of_two(size_t n) { + size_t padded_n = 1; + while (padded_n < n) { + padded_n <<= 1; + } + return padded_n; +} + +enum class FFTNorm { BACKWARD, FORWARD, ORTHO }; + +static FFTNorm parse_fft_norm(const char* norm) { + if (!norm || std::strcmp(norm, "backward") == 0) return FFTNorm::BACKWARD; + if (std::strcmp(norm, "forward") == 0) return FFTNorm::FORWARD; + if (std::strcmp(norm, "ortho") == 0) return FFTNorm::ORTHO; + throw std::invalid_argument("norm must be one of {\"backward\",\"forward\",\"ortho\"}"); +} + +static float fft_norm_factor(size_t n, FFTNorm norm_mode, bool inverse) { + if (norm_mode == FFTNorm::ORTHO) return 1.0f / std::sqrt(static_cast(n)); + if (norm_mode == FFTNorm::FORWARD) return inverse ? 1.0f : (1.0f / static_cast(n)); + return inverse ? (1.0f / static_cast(n)) : 1.0f; +} + +static void ensure_irfft_scratch(std::vector& re, std::vector& im, size_t n) { + if (re.size() < n) re.resize(n); + if (im.size() < n) im.resize(n); +} + +static void fft_radix2(float* re, float* im, size_t n, bool inverse) { + if (!is_power_of_two(n)) return; + const size_t log2n = log2_power_of_two(n); + + for (size_t i = 0; i < n; i++) { + size_t j = bit_reverse(i, log2n); + if (i < j) { + std::swap(re[i], re[j]); + std::swap(im[i], im[j]); + } + } + + for (size_t s = 1; s <= log2n; s++) { + size_t m = size_t{1} << s; + size_t m2 = m >> 1; + float w_re = 1.0f; + float w_im = 0.0f; + const float angle = static_cast(M_PI) / static_cast(m2); + const float wm_re = std::cos(angle); + const float wm_im = inverse ? std::sin(angle) : -std::sin(angle); + + for (size_t j = 0; j < m2; j++) { + for (size_t k = j; k < n; k += m) { + size_t k_m2 = k + m2; + float t_re = w_re * re[k_m2] - w_im * im[k_m2]; + float t_im = w_re * im[k_m2] + w_im * re[k_m2]; + float u_re = re[k]; + float u_im = im[k]; + re[k] = u_re + t_re; + im[k] = u_im + t_im; + re[k_m2] = u_re - t_re; + im[k_m2] = u_im - t_im; + } + float new_w_re = w_re * wm_re - w_im * wm_im; + float new_w_im = w_re * wm_im + w_im * wm_re; + w_re = new_w_re; + w_im = new_w_im; + } + } +} + +static void idft_r_f32_1d_dft( + const float* input, float* output, const size_t n, const float norm_factor) { + const size_t in_len = n / 2 + 1; + const float two_pi_over_n = (2.0f * static_cast(M_PI)) / static_cast(n); + for (size_t t = 0; t < n; ++t) { + float sum = input[0]; + for (size_t k = 1; k < in_len; ++k) { + const float re = input[k * 2]; + const float im = input[k * 2 + 1]; + const float a = two_pi_over_n * static_cast(k * t); + const float c = std::cos(a); + const float s = std::sin(a); + const bool self_conjugate = (k * 2 == n); + if (self_conjugate) { + sum += re * c; + } else { + sum += 2.0f * (re * c - im * s); + } + } + output[t] = sum * norm_factor; + } +} + +void cactus_rfft_f32_1d(const float* input, float* output, size_t n, const char* norm) { + const size_t out_len = n / 2 + 1; + const FFTNorm norm_mode = parse_fft_norm(norm); + const float norm_factor = fft_norm_factor(n, norm_mode, false); + + if (n == 1) { + output[0] = input[0] * norm_factor; + output[1] = 0.0f; + return; + } + +#ifdef __APPLE__ + { + size_t log2n = 0; + size_t padded_n = 1; + while (padded_n < n) { + padded_n <<= 1; + log2n++; + } + + static std::mutex fft_mu; + static std::unordered_map fft_cache; + FFTSetup fft_setup; + { + std::lock_guard lk(fft_mu); + auto it = fft_cache.find(log2n); + if (it == fft_cache.end()) { + FFTSetup created = vDSP_create_fftsetup(log2n, FFT_RADIX2); + if (created) it = fft_cache.emplace(log2n, created).first; + else it = fft_cache.end(); + } + fft_setup = it != fft_cache.end() ? it->second : nullptr; + } + if (fft_setup) { + std::vector buffer(padded_n, 0.0f); + std::copy(input, input + n, buffer.begin()); + + DSPSplitComplex split; + std::vector real_part(padded_n / 2); + std::vector imag_part(padded_n / 2); + split.realp = real_part.data(); + split.imagp = imag_part.data(); + + vDSP_ctoz(reinterpret_cast(buffer.data()), 2, &split, 1, padded_n / 2); + vDSP_fft_zrip(fft_setup, &split, 1, log2n, FFT_FORWARD); + + float scale = 0.5f * norm_factor; + + output[0] = split.realp[0] * scale; + output[1] = 0.0f; + + if (out_len > 1) { + size_t nyquist_idx = padded_n / 2; + if (nyquist_idx < out_len) { + output[nyquist_idx * 2] = split.imagp[0] * scale; + output[nyquist_idx * 2 + 1] = 0.0f; + } + } + + for (size_t i = 1; i < out_len && i < padded_n / 2; i++) { + output[i * 2] = split.realp[i] * scale; + output[i * 2 + 1] = split.imagp[i] * scale; + } + + return; + } + } +#endif + + const size_t padded_n = next_power_of_two(n); + std::vector re(padded_n, 0.0f), im(padded_n, 0.0f); + std::copy(input, input + n, re.begin()); + fft_radix2(re.data(), im.data(), padded_n, false); + + for (size_t i = 0; i < out_len; i++) { + output[i * 2] = re[i] * norm_factor; + output[i * 2 + 1] = im[i] * norm_factor; + } +} + +void cactus_irfft_f32_1d(const float* input, float* output, size_t n, const char* norm) { + const size_t in_len = n / 2 + 1; + const FFTNorm norm_mode = parse_fft_norm(norm); + const float norm_factor = fft_norm_factor(n, norm_mode, true); + + if (n == 1) { + output[0] = input[0] * norm_factor; + return; + } + + if (!is_power_of_two(n)) { + idft_r_f32_1d_dft(input, output, n, norm_factor); + return; + } + + const size_t padded_n = next_power_of_two(n); + const size_t padded_half_plus_one = padded_n / 2 + 1; + thread_local std::vector re_scratch; + thread_local std::vector im_scratch; + ensure_irfft_scratch(re_scratch, im_scratch, padded_n); + std::fill(re_scratch.begin(), re_scratch.begin() + padded_n, 0.0f); + std::fill(im_scratch.begin(), im_scratch.begin() + padded_n, 0.0f); + float* re = re_scratch.data(); + float* im = im_scratch.data(); + for (size_t i = 0; i < in_len; ++i) { + re[i] = input[i * 2]; + im[i] = input[i * 2 + 1]; + } + + im[0] = 0.0f; + const bool has_padded_nyquist_bin = in_len == padded_half_plus_one; + if (has_padded_nyquist_bin) { + im[padded_n / 2] = 0.0f; + } + + const size_t mirror_limit = has_padded_nyquist_bin ? (in_len - 1) : in_len; + for (size_t i = 1; i < mirror_limit; ++i) { + const size_t mirrored = padded_n - i; + re[mirrored] = re[i]; + im[mirrored] = -im[i]; + } + + fft_radix2(re, im, padded_n, true); + + for (size_t i = 0; i < n; ++i) { + output[i] = re[i] * norm_factor; + } +} + +float cactus_hertz_to_mel(float freq, const char* mel_scale) { + if (std::strcmp(mel_scale, "htk") == 0) { + return 2595.0f * std::log10(1.0f + (freq / 700.0f)); + } else if (std::strcmp(mel_scale, "kaldi") == 0) { + return 1127.0f * std::log(1.0f + (freq / 700.0f)); + } + + const float min_log_hertz = 1000.0f; + const float min_log_mel = 15.0f; + const float logstep = 27.0f / std::log(6.4f); + float mels = 3.0f * freq / 200.0f; + + if (freq >= min_log_hertz) { + mels = min_log_mel + std::log(freq / min_log_hertz) * logstep; + } + + return mels; +} + +float cactus_mel_to_hertz(float mels, const char* mel_scale) { + if (std::strcmp(mel_scale, "htk") == 0) { + return 700.0f * (std::pow(10.0f, mels / 2595.0f) - 1.0f); + } else if (std::strcmp(mel_scale, "kaldi") == 0) { + return 700.0f * (std::exp(mels / 1127.0f) - 1.0f); + } + + const float min_log_hertz = 1000.0f; + const float min_log_mel = 15.0f; + const float logstep = std::log(6.4f) / 27.0f; + float freq = 200.0f * mels / 3.0f; + + if (mels >= min_log_mel) { + freq = min_log_hertz * std::exp(logstep * (mels - min_log_mel)); + } + + return freq; +} + +void cactus_generate_mel_filter_bank( + float* mel_filters, + int num_frequency_bins, + int num_mel_filters, + float min_frequency, + float max_frequency, + int sampling_rate, + const char* norm, + const char* mel_scale, + bool triangularize_in_mel_space) { + + if (norm != nullptr && std::strcmp(norm, "slaney") != 0) { + throw std::invalid_argument("norm must be one of None or \"slaney\""); + } + + if (std::strcmp(mel_scale, "htk") != 0 && std::strcmp(mel_scale, "kaldi") != 0 && std::strcmp(mel_scale, "slaney") != 0) { + throw std::invalid_argument("mel_scale should be one of \"htk\", \"slaney\" or \"kaldi\"."); + } + + if (num_frequency_bins < 2) { + throw std::invalid_argument("Require num_frequency_bins >= 2"); + } + + if (min_frequency > max_frequency) { + throw std::invalid_argument("Require min_frequency <= max_frequency"); + } + + const float mel_min = cactus_hertz_to_mel(min_frequency, mel_scale); + const float mel_max = cactus_hertz_to_mel(max_frequency, mel_scale); + + std::vector mel_freqs(num_mel_filters + 2); + for (int i = 0; i < num_mel_filters + 2; i++) { + mel_freqs[i] = mel_min + (mel_max - mel_min) * i / (num_mel_filters + 1); + } + + std::vector filter_freqs(num_mel_filters + 2); + for (int i = 0; i < num_mel_filters + 2; i++) { + filter_freqs[i] = cactus_mel_to_hertz(mel_freqs[i], mel_scale); + } + + std::vector fft_freqs(num_frequency_bins); + if (triangularize_in_mel_space) { + float fft_bin_width = static_cast(sampling_rate) / ((num_frequency_bins - 1) * 2); + for (int i = 0; i < num_frequency_bins; i++) { + fft_freqs[i] = cactus_hertz_to_mel(fft_bin_width * i, mel_scale); + } + filter_freqs = mel_freqs; + } else { + for (int i = 0; i < num_frequency_bins; i++) { + fft_freqs[i] = (static_cast(sampling_rate) / 2.0f) * i / (num_frequency_bins - 1); + } + } + + for (int i = 0; i < num_mel_filters; i++) { + float left_edge = filter_freqs[i]; + float center = filter_freqs[i + 1]; + float right_edge = filter_freqs[i + 2]; + + for (int j = 0; j < num_frequency_bins; j++) { + float freq = fft_freqs[j]; + float down_slope = (freq - left_edge) / (center - left_edge); + float up_slope = (right_edge - freq) / (right_edge - center); + mel_filters[i * num_frequency_bins + j] = std::max(0.0f, std::min(down_slope, up_slope)); + } + } + + if (norm != nullptr && std::strcmp(norm, "slaney") == 0) { + for (int i = 0; i < num_mel_filters; i++) { + float enorm = 2.0f / (filter_freqs[i + 2] - filter_freqs[i]); + for (int j = 0; j < num_frequency_bins; j++) { + mel_filters[i * num_frequency_bins + j] *= enorm; + } + } + } +} + +void cactus_spectrogram_to_db( + float* spectrogram, size_t size, + float reference, float min_value, + const float* db_range, float multiplier) { + + if (reference <= 0.0f) throw std::invalid_argument("reference must be greater than zero"); + if (min_value <= 0.0f) throw std::invalid_argument("min_value must be greater than zero"); + + reference = std::max(min_value, reference); + const float log_ref = std::log10(reference); + + CactusThreading::parallel_for(size, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { + for (size_t i = start; i < end; i++) { + float value = std::max(min_value, spectrogram[i]); + spectrogram[i] = multiplier * (std::log10(value) - log_ref); + } + }); + + if (db_range != nullptr) { + if (*db_range <= 0.0f) throw std::invalid_argument("db_range must be greater than zero"); + + float max_db = CactusThreading::parallel_reduce, float, std::function>( + size, CactusThreading::Thresholds::ALL_REDUCE, + [&](size_t start, size_t end) { + float local_max = -std::numeric_limits::infinity(); + for (size_t i = start; i < end; i++) local_max = std::max(local_max, spectrogram[i]); + return local_max; + }, + -std::numeric_limits::infinity(), + [](float a, float b) { return std::max(a, b); } + ); + + float min_db = max_db - *db_range; + CactusThreading::parallel_for(size, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { + for (size_t i = start; i < end; i++) spectrogram[i] = std::max(min_db, spectrogram[i]); + }); + } +} + +void cactus_compute_spectrogram_f32( + const float* waveform, size_t waveform_length, + const float* window, size_t window_length, + size_t frame_length, size_t hop_length, + const size_t* fft_length, + float* spectrogram, float power, + bool center, const char* pad_mode, bool onesided, + float dither, const float* preemphasis, + const float* mel_filters, size_t mel_filters_size, + float mel_floor, const char* log_mel, + float reference, float min_value, + const float* db_range, bool remove_dc_offset) { + + (void)onesided; + + size_t actual_fft_length = fft_length ? *fft_length : frame_length; + + if (frame_length > actual_fft_length) { + throw std::invalid_argument("frame_length may not be larger than fft_length"); + } + + std::vector hann_window; + const float* actual_window = window; + + if (window == nullptr) { + size_t length = frame_length + 1; + hann_window.resize(frame_length); + for (size_t i = 0; i < frame_length; i++) { + hann_window[i] = 0.5f * (1.0f - std::cos(2.0f * static_cast(M_PI) * i / (length - 1))); + } + actual_window = hann_window.data(); + } else if (window_length != frame_length) { + throw std::invalid_argument("window length must equal frame_length"); + } + + if (hop_length <= 0) throw std::invalid_argument("hop_length must be greater than zero"); + + if (power == 0.0f && mel_filters != nullptr) { + throw std::invalid_argument("Cannot compute mel spectrogram with power=0"); + } + + std::vector padded_waveform; + const float* input_waveform = waveform; + size_t input_length = waveform_length; + + if (center) { + size_t pad_length = frame_length / 2; + size_t padded_length = waveform_length + 2 * pad_length; + padded_waveform.assign(padded_length, 0.0f); + + if (std::strcmp(pad_mode, "reflect") == 0) { + for (size_t i = 0; i < pad_length; i++) padded_waveform[i] = waveform[pad_length - i]; + std::copy(waveform, waveform + waveform_length, padded_waveform.data() + pad_length); + for (size_t i = 0; i < pad_length; i++) padded_waveform[pad_length + waveform_length + i] = waveform[waveform_length - 2 - i]; + } else if (std::strcmp(pad_mode, "constant") == 0) { + std::copy(waveform, waveform + waveform_length, padded_waveform.data() + pad_length); + } else { + throw std::invalid_argument("Unsupported pad_mode"); + } + + input_waveform = padded_waveform.data(); + input_length = padded_length; + } + + const size_t num_frames = 1 + (input_length - frame_length) / hop_length; + const size_t num_frequency_bins = (actual_fft_length / 2) + 1; + + const size_t num_mel_bins = mel_filters != nullptr ? mel_filters_size / num_frequency_bins : 0; + const size_t spectrogram_bins = mel_filters != nullptr ? num_mel_bins : num_frequency_bins; + + std::vector temp_spectrogram(num_frames * num_frequency_bins); + + CactusThreading::parallel_for(num_frames, CactusThreading::Thresholds::SCALAR_EXPENSIVE, [&](size_t start_frame, size_t end_frame) { + std::vector local_buffer(actual_fft_length); + std::vector local_complex(num_frequency_bins * 2); + + for (size_t frame_idx = start_frame; frame_idx < end_frame; frame_idx++) { + size_t timestep = frame_idx * hop_length; + std::fill(local_buffer.begin(), local_buffer.end(), 0.0f); + + size_t available_length = std::min(frame_length, input_length - timestep); + std::copy(input_waveform + timestep, input_waveform + timestep + available_length, local_buffer.data()); + + if (dither != 0.0f) { + thread_local std::mt19937 rng(std::random_device{}()); + std::normal_distribution dist(0.0f, 1.0f); + for (size_t i = 0; i < frame_length; i++) local_buffer[i] += dither * dist(rng); + } + + if (remove_dc_offset) { + float mean = 0.0f; + for (size_t i = 0; i < frame_length; i++) mean += local_buffer[i]; + mean /= static_cast(frame_length); + for (size_t i = 0; i < frame_length; i++) local_buffer[i] -= mean; + } + + if (preemphasis != nullptr) { + float preemph_coef = *preemphasis; + for (size_t i = frame_length - 1; i > 0; i--) local_buffer[i] -= preemph_coef * local_buffer[i - 1]; + local_buffer[0] *= (1.0f - preemph_coef); + } + + for (size_t i = 0; i < frame_length; i++) local_buffer[i] *= actual_window[i]; + + cactus_rfft_f32_1d(local_buffer.data(), local_complex.data(), actual_fft_length, "backward"); + + for (size_t i = 0; i < num_frequency_bins; i++) { + float real = local_complex[i * 2]; + float imag = local_complex[i * 2 + 1]; + float magnitude = std::hypot(real, imag); + temp_spectrogram[frame_idx * num_frequency_bins + i] = std::pow(magnitude, power); + } + } + }); + + if (mel_filters != nullptr) { + CactusThreading::parallel_for_2d(num_mel_bins, num_frames, CactusThreading::Thresholds::AXIS_REDUCE, [&](size_t m, size_t t) { + float sum = 0.0f; + for (size_t f = 0; f < num_frequency_bins; f++) { + sum += mel_filters[m * num_frequency_bins + f] * temp_spectrogram[t * num_frequency_bins + f]; + } + spectrogram[m * num_frames + t] = std::max(mel_floor, sum); + }); + } else { + CactusThreading::parallel_for_2d(num_frames, num_frequency_bins, CactusThreading::Thresholds::AXIS_REDUCE, [&](size_t t, size_t f) { + spectrogram[f * num_frames + t] = temp_spectrogram[t * num_frequency_bins + f]; + }); + } + + if (power != 0.0f && log_mel != nullptr) { + const size_t total_elements = spectrogram_bins * num_frames; + + if (std::strcmp(log_mel, "log") == 0) { + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { + for (size_t i = start; i < end; i++) spectrogram[i] = std::log(spectrogram[i]); + }); + } else if (std::strcmp(log_mel, "log10") == 0) { + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { + for (size_t i = start; i < end; i++) spectrogram[i] = std::log10(spectrogram[i]); + }); + } else if (std::strcmp(log_mel, "dB") == 0) { + if (power == 1.0f) cactus_spectrogram_to_db(spectrogram, total_elements, reference, min_value, db_range, 20.0f); + else if (power == 2.0f) cactus_spectrogram_to_db(spectrogram, total_elements, reference, min_value, db_range, 10.0f); + else throw std::invalid_argument("Cannot use log_mel option 'dB' with this power value"); + } else { + throw std::invalid_argument("Unknown log_mel option"); + } + } +} diff --git a/cactus-kernels/src/fused.cpp b/cactus-kernels/src/fused.cpp new file mode 100644 index 000000000..fc187a2a0 --- /dev/null +++ b/cactus-kernels/src/fused.cpp @@ -0,0 +1,850 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include +#include +#include + +void cactus_gaussian_topk_f16( + const __fp16* input, + __fp16* output, + size_t rows, + size_t cols, + float ppf +) { + static constexpr float MAX_SAFE_ABS = 240.0f; + + for (size_t r = 0; r < rows; r++) { + const __fp16* row_in = input + r * cols; + __fp16* row_out = output + r * cols; + + float32x4_t sum_vec = vdupq_n_f32(0.0f); + size_t d = 0; + for (; d + 8 <= cols; d += 8) { + float16x8_t v = vld1q_f16(row_in + d); + sum_vec = vaddq_f32(sum_vec, vcvt_f32_f16(vget_low_f16(v))); + sum_vec = vaddq_f32(sum_vec, vcvt_f32_f16(vget_high_f16(v))); + } + float sum = vaddvq_f32(sum_vec); + for (; d < cols; d++) { + sum += static_cast(row_in[d]); + } + float mu = sum / static_cast(cols); + + float32x4_t mu_vec = vdupq_n_f32(mu); + float32x4_t max_abs_vec = vdupq_n_f32(0.0f); + d = 0; + for (; d + 8 <= cols; d += 8) { + float16x8_t v = vld1q_f16(row_in + d); + float32x4_t lo = vsubq_f32(vcvt_f32_f16(vget_low_f16(v)), mu_vec); + float32x4_t hi = vsubq_f32(vcvt_f32_f16(vget_high_f16(v)), mu_vec); + max_abs_vec = vmaxq_f32(max_abs_vec, vabsq_f32(lo)); + max_abs_vec = vmaxq_f32(max_abs_vec, vabsq_f32(hi)); + } + float max_abs = vmaxvq_f32(max_abs_vec); + for (; d < cols; d++) { + float diff = static_cast(row_in[d]) - mu; + float a = diff < 0.0f ? -diff : diff; + if (a > max_abs) max_abs = a; + } + + float scale = max_abs * (1.0f / MAX_SAFE_ABS); + if (scale < 1.0f) scale = 1.0f; + float inv_scale = 1.0f / scale; + + float32x4_t inv_scale_vec = vdupq_n_f32(inv_scale); + float32x4_t var_sum_vec = vdupq_n_f32(0.0f); + d = 0; + for (; d + 8 <= cols; d += 8) { + float16x8_t v = vld1q_f16(row_in + d); + float32x4_t lo = vmulq_f32(vsubq_f32(vcvt_f32_f16(vget_low_f16(v)), mu_vec), inv_scale_vec); + float32x4_t hi = vmulq_f32(vsubq_f32(vcvt_f32_f16(vget_high_f16(v)), mu_vec), inv_scale_vec); + var_sum_vec = vfmaq_f32(var_sum_vec, lo, lo); + var_sum_vec = vfmaq_f32(var_sum_vec, hi, hi); + } + float var_sum = vaddvq_f32(var_sum_vec); + for (; d < cols; d++) { + float ds = (static_cast(row_in[d]) - mu) * inv_scale; + var_sum += ds * ds; + } + float variance = var_sum / static_cast(cols); + float sigma = sqrtf(variance) * scale; + + float cutoff = mu + ppf * sigma; + float32x4_t cutoff_vec = vdupq_n_f32(cutoff); + float32x4_t zero_vec = vdupq_n_f32(0.0f); + d = 0; + for (; d + 8 <= cols; d += 8) { + float16x8_t v = vld1q_f16(row_in + d); + float32x4_t lo = vmaxq_f32(vsubq_f32(vcvt_f32_f16(vget_low_f16(v)), cutoff_vec), zero_vec); + float32x4_t hi = vmaxq_f32(vsubq_f32(vcvt_f32_f16(vget_high_f16(v)), cutoff_vec), zero_vec); + vst1q_f16(row_out + d, vcombine_f16(vcvt_f16_f32(lo), vcvt_f16_f32(hi))); + } + for (; d < cols; d++) { + float val = static_cast(row_in[d]) - cutoff; + row_out[d] = static_cast<__fp16>(val > 0.0f ? val : 0.0f); + } + } +} + +void cactus_altup_predict_f16( + const __fp16* coefs, + const __fp16* const* streams, + __fp16* output, + size_t n, + size_t seq_len, + size_t hidden_dim +) { + if (n > 8) throw std::runtime_error("cactus_altup_predict_f16 expects n <= 8"); + const size_t coef_stride = n * n; + + for (size_t i = 0; i < n; i++) { + for (size_t s = 0; s < seq_len; s++) { + const __fp16* coef_row = coefs + s * coef_stride + i * n; + __fp16* out_row = output + (i * seq_len + s) * hidden_dim; + const __fp16* src_i = streams[i] + s * hidden_dim; + + float c[8]; + for (size_t j = 0; j < n; j++) { + c[j] = static_cast(coef_row[j]); + } + + size_t d = 0; + for (; d + 8 <= hidden_dim; d += 8) { + float32x4_t acc_lo = vcvt_f32_f16(vget_low_f16(vld1q_f16(src_i + d))); + float32x4_t acc_hi = vcvt_f32_f16(vget_high_f16(vld1q_f16(src_i + d))); + + for (size_t j = 0; j < n; j++) { + float16x8_t sj = vld1q_f16(streams[j] + s * hidden_dim + d); + float32x4_t cj = vdupq_n_f32(c[j]); + acc_lo = vfmaq_f32(acc_lo, vcvt_f32_f16(vget_low_f16(sj)), cj); + acc_hi = vfmaq_f32(acc_hi, vcvt_f32_f16(vget_high_f16(sj)), cj); + } + + vst1q_f16(out_row + d, vcombine_f16(vcvt_f16_f32(acc_lo), vcvt_f16_f32(acc_hi))); + } + for (; d < hidden_dim; d++) { + float acc = static_cast(src_i[d]); + for (size_t j = 0; j < n; j++) { + acc += c[j] * static_cast(streams[j][s * hidden_dim + d]); + } + out_row[d] = static_cast<__fp16>(acc); + } + } + } +} + +void cactus_altup_correct_f16( + const __fp16* coefs, + const __fp16* innovation, + const __fp16* const* predictions, + __fp16* output, + size_t n, + size_t seq_len, + size_t hidden_dim +) { + for (size_t i = 0; i < n; i++) { + for (size_t s = 0; s < seq_len; s++) { + float ci = static_cast(coefs[s * n + i]); + float32x4_t ci_vec = vdupq_n_f32(ci); + + const __fp16* pred_row = predictions[i] + s * hidden_dim; + const __fp16* innov_row = innovation + s * hidden_dim; + __fp16* out_row = output + (i * seq_len + s) * hidden_dim; + + size_t d = 0; + for (; d + 8 <= hidden_dim; d += 8) { + float16x8_t pred = vld1q_f16(pred_row + d); + float16x8_t innov = vld1q_f16(innov_row + d); + + float32x4_t p_lo = vcvt_f32_f16(vget_low_f16(pred)); + float32x4_t p_hi = vcvt_f32_f16(vget_high_f16(pred)); + float32x4_t i_lo = vcvt_f32_f16(vget_low_f16(innov)); + float32x4_t i_hi = vcvt_f32_f16(vget_high_f16(innov)); + + float32x4_t out_lo = vfmaq_f32(p_lo, i_lo, ci_vec); + float32x4_t out_hi = vfmaq_f32(p_hi, i_hi, ci_vec); + + vst1q_f16(out_row + d, vcombine_f16(vcvt_f16_f32(out_lo), vcvt_f16_f32(out_hi))); + } + for (; d < hidden_dim; d++) { + out_row[d] = static_cast<__fp16>( + static_cast(pred_row[d]) + ci * static_cast(innov_row[d]) + ); + } + } + } +} + +namespace { + +inline float safe_exp_gate(float gate_log) { + const float clamped = std::max(-20.0f, std::min(6.0f, gate_log)); + return std::exp(clamped); +} + +inline void fold_state_scale(std::vector& state, float factor) { + const float32x4_t f4 = vdupq_n_f32(factor); + size_t i = 0; + for (; i + 8 <= state.size(); i += 8) { + vst1q_f32(&state[i], vmulq_f32(vld1q_f32(&state[i]), f4)); + vst1q_f32(&state[i + 4], vmulq_f32(vld1q_f32(&state[i + 4]), f4)); + } + for (; i < state.size(); ++i) state[i] *= factor; +} + +struct GatedDeltaChunkScratch { + std::vector state; + std::vector m_chunk; + std::vector q_chunk; + std::vector k_chunk; + std::vector v_chunk; + std::vector g_chunk; + std::vector beta_chunk; + std::vector n0_proj; + std::vector delta_chunk; + std::vector gram; + std::vector p_prev; + std::vector p_curr; + std::vector inv_p_curr; + std::vector coeff; + std::vector out_acc; + + void ensure(size_t k_dim, size_t v_dim, size_t c_max) { + const size_t kv = k_dim * v_dim; + if (state.size() < kv) state.resize(kv); + if (m_chunk.size() < kv) m_chunk.resize(kv); + if (q_chunk.size() < c_max * k_dim) q_chunk.resize(c_max * k_dim); + if (k_chunk.size() < c_max * k_dim) k_chunk.resize(c_max * k_dim); + if (v_chunk.size() < c_max * v_dim) v_chunk.resize(c_max * v_dim); + if (g_chunk.size() < c_max) g_chunk.resize(c_max); + if (beta_chunk.size() < c_max) beta_chunk.resize(c_max); + if (n0_proj.size() < c_max * v_dim) n0_proj.resize(c_max * v_dim); + if (delta_chunk.size() < c_max * v_dim) delta_chunk.resize(c_max * v_dim); + if (gram.size() < c_max * c_max) gram.resize(c_max * c_max); + if (p_prev.size() < c_max) p_prev.resize(c_max); + if (p_curr.size() < c_max) p_curr.resize(c_max); + if (inv_p_curr.size() < c_max) inv_p_curr.resize(c_max); + if (coeff.size() < c_max * c_max) coeff.resize(c_max * c_max); + if (out_acc.size() < v_dim) out_acc.resize(v_dim); + } +}; + +thread_local GatedDeltaChunkScratch g_gated_deltanet_chunk_scratch; + +inline size_t tuned_gated_deltanet_chunk_size(size_t requested_chunk, size_t k_dim, size_t v_dim) { + size_t chunk = requested_chunk == 0 ? 64 : requested_chunk; + + const char* env_chunk = std::getenv("CACTUS_GATED_DELTANET_CHUNK_SIZE"); + if (env_chunk != nullptr) { + long parsed = std::strtol(env_chunk, nullptr, 10); + if (parsed > 1) { + chunk = static_cast(parsed); + } + } else { + (void)k_dim; + (void)v_dim; + chunk = std::min(chunk, 16); + } + + return std::max(2, chunk); +} + +void gated_deltanet_step( + const __fp16* q_ptr, + const __fp16* k_ptr, + const __fp16* v_ptr, + float gate_log, + float beta, + float scale, + size_t k_dim, + size_t v_dim, + std::vector& state, + float& state_scale, + std::vector& proj, + std::vector& delta, + __fp16* out_ptr) { + const float gate_log_safe = std::isfinite(gate_log) ? gate_log : -20.0f; + const float beta_safe = std::isfinite(beta) ? std::min(1.0f, std::max(0.0f, beta)) : 0.0f; + + if (proj.size() != v_dim) { + proj.assign(v_dim, 0.0f); + } + if (delta.size() != v_dim) { + delta.assign(v_dim, 0.0f); + } + + const float prev_scale = state_scale; + state_scale = prev_scale * safe_exp_gate(gate_log_safe); + std::fill(proj.begin(), proj.end(), 0.0f); + + for (size_t kd = 0; kd < k_dim; ++kd) { + if (kd + 2 < k_dim) __builtin_prefetch(state.data() + (kd + 2) * v_dim, 0, 3); + const float k_val = static_cast(k_ptr[kd]); + const float* state_row = state.data() + kd * v_dim; + size_t vd = 0; + const float32x4_t k4 = vdupq_n_f32(k_val); + for (; vd + 4 <= v_dim; vd += 4) { + float32x4_t p = vld1q_f32(proj.data() + vd); + float32x4_t s = vld1q_f32(state_row + vd); + p = vfmaq_f32(p, s, k4); + vst1q_f32(proj.data() + vd, p); + } + for (; vd < v_dim; ++vd) { + proj[vd] += state_row[vd] * k_val; + } + } + { + size_t vd = 0; + const float32x4_t scale4 = vdupq_n_f32(state_scale); + for (; vd + 4 <= v_dim; vd += 4) { + float32x4_t p = vld1q_f32(proj.data() + vd); + p = vmulq_f32(p, scale4); + vst1q_f32(proj.data() + vd, p); + } + for (; vd < v_dim; ++vd) { + proj[vd] *= state_scale; + } + } + + for (size_t vd = 0; vd < v_dim; ++vd) { + const float v_val = static_cast(v_ptr[vd]); + delta[vd] = (v_val - proj[vd]) * beta_safe; + } + + if (!std::isfinite(state_scale) || + std::fabs(state_scale) < 1e-8f || + std::fabs(state_scale) > 1e8f) { + fold_state_scale(state, state_scale); + state_scale = 1.0f; + } + const float inv_scale = 1.0f / state_scale; + + for (size_t kd = 0; kd < k_dim; ++kd) { + if (kd + 2 < k_dim) __builtin_prefetch(state.data() + (kd + 2) * v_dim, 1, 3); + const float k_scaled = static_cast(k_ptr[kd]) * inv_scale; + float* state_row = state.data() + kd * v_dim; + size_t vd = 0; + const float32x4_t k4 = vdupq_n_f32(k_scaled); + for (; vd + 4 <= v_dim; vd += 4) { + float32x4_t s = vld1q_f32(state_row + vd); + float32x4_t d = vld1q_f32(delta.data() + vd); + s = vfmaq_f32(s, d, k4); + vst1q_f32(state_row + vd, s); + } + for (; vd < v_dim; ++vd) { + state_row[vd] += k_scaled * delta[vd]; + } + for (size_t vd = 0; vd < v_dim; ++vd) { + if (!std::isfinite(state_row[vd])) { + state_row[vd] = 0.0f; + } + } + } + + std::fill(proj.begin(), proj.end(), 0.0f); + for (size_t kd = 0; kd < k_dim; ++kd) { + if (kd + 2 < k_dim) __builtin_prefetch(state.data() + (kd + 2) * v_dim, 0, 3); + const float q_val = static_cast(q_ptr[kd]); + const float* state_row = state.data() + kd * v_dim; + size_t vd = 0; + const float32x4_t q4 = vdupq_n_f32(q_val); + for (; vd + 4 <= v_dim; vd += 4) { + float32x4_t o = vld1q_f32(proj.data() + vd); + float32x4_t s = vld1q_f32(state_row + vd); + o = vfmaq_f32(o, s, q4); + vst1q_f32(proj.data() + vd, o); + } + for (; vd < v_dim; ++vd) { + proj[vd] += state_row[vd] * q_val; + } + } + + const float out_scale = state_scale * scale; + for (size_t vd = 0; vd < v_dim; ++vd) { + float acc = proj[vd]; + if (!std::isfinite(acc)) { + acc = 0.0f; + } + out_ptr[vd] = static_cast<__fp16>(acc * out_scale); + } +} + +void gated_deltanet_prefill_old_f16( + const __fp16* q_data, + const __fp16* k_data, + const __fp16* v_data, + const __fp16* g_data, + const __fp16* b_data, + const __fp16* s_data, + __fp16* out, + size_t B, + size_t T, + size_t Hq, + size_t Hv, + size_t K, + size_t V, + float scale) { + const size_t out_seq = T + K; + const size_t qk_repeat = Hv / Hq; + + CactusThreading::parallel_for(B * Hv, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t bh_start, size_t bh_end) { + std::vector state(K * V); + std::vector proj(V); + std::vector delta(V); + + for (size_t bh = bh_start; bh < bh_end; ++bh) { + const size_t batch = bh / Hv; + const size_t v_head = bh % Hv; + const size_t qk_head = v_head / qk_repeat; + + for (size_t kd = 0; kd < K; ++kd) { + for (size_t vd = 0; vd < V; ++vd) { + const size_t s_idx = (((batch * K + kd) * Hv + v_head) * V + vd); + state[kd * V + vd] = static_cast(s_data[s_idx]); + } + } + + float state_scale = 1.0f; + for (size_t t = 0; t < T; ++t) { + const size_t qkv_k_base = (((batch * T + t) * Hq + qk_head) * K); + const size_t qkv_v_base = (((batch * T + t) * Hv + v_head) * V); + const size_t gb_idx = ((batch * T + t) * Hv + v_head); + + gated_deltanet_step( + q_data + qkv_k_base, + k_data + qkv_k_base, + v_data + qkv_v_base, + static_cast(g_data[gb_idx]), + static_cast(b_data[gb_idx]), + scale, + K, + V, + state, + state_scale, + proj, + delta, + out + (((batch * out_seq + t) * Hv + v_head) * V)); + } + + fold_state_scale(state, state_scale); + for (size_t kd = 0; kd < K; ++kd) { + for (size_t vd = 0; vd < V; ++vd) { + const size_t out_idx = (((batch * out_seq + (T + kd)) * Hv + v_head) * V + vd); + out[out_idx] = static_cast<__fp16>(state[kd * V + vd]); + } + } + } + }); +} + +void gated_deltanet_prefill_chunked_f16( + const __fp16* q_data, + const __fp16* k_data, + const __fp16* v_data, + const __fp16* g_data, + const __fp16* b_data, + const __fp16* s_data, + __fp16* out, + size_t B, + size_t T, + size_t Hq, + size_t Hv, + size_t K, + size_t V, + size_t chunk_size, + float scale) { + const size_t out_seq = T + K; + const size_t qk_repeat = Hv / Hq; + constexpr float kMinAbs = 1e-12f; + + CactusThreading::parallel_for(B * Hv, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t bh_start, size_t bh_end) { + auto& ws = g_gated_deltanet_chunk_scratch; + ws.ensure(K, V, chunk_size); + + float* state = ws.state.data(); + float* m_chunk = ws.m_chunk.data(); + float* q_chunk = ws.q_chunk.data(); + float* k_chunk = ws.k_chunk.data(); + float* v_chunk = ws.v_chunk.data(); + float* g_chunk = ws.g_chunk.data(); + float* beta_chunk = ws.beta_chunk.data(); + float* n0_proj = ws.n0_proj.data(); + float* delta_chunk = ws.delta_chunk.data(); + float* gram = ws.gram.data(); + float* p_prev = ws.p_prev.data(); + float* p_curr = ws.p_curr.data(); + float* inv_p_curr = ws.inv_p_curr.data(); + float* coeff = ws.coeff.data(); + float* out_acc = ws.out_acc.data(); + + for (size_t bh = bh_start; bh < bh_end; ++bh) { + const size_t batch = bh / Hv; + const size_t v_head = bh % Hv; + const size_t qk_head = v_head / qk_repeat; + + for (size_t kd = 0; kd < K; ++kd) { + const size_t s_idx = (((batch * K + kd) * Hv + v_head) * V); + cactus_fp16_to_fp32(s_data + s_idx, state + kd * V, V); + } + + for (size_t t0 = 0; t0 < T; t0 += chunk_size) { + const size_t C = std::min(chunk_size, T - t0); + const size_t CV = C * V; + const size_t CC = C * C; + + std::fill(n0_proj, n0_proj + CV, 0.0f); + std::fill(delta_chunk, delta_chunk + CV, 0.0f); + std::fill(gram, gram + CC, 0.0f); + std::fill(coeff, coeff + CC, 0.0f); + std::fill(m_chunk, m_chunk + (K * V), 0.0f); + + // Fused fp16→fp32 conversion + n0_proj: k_chunk stays in L1 + for (size_t ct = 0; ct < C; ++ct) { + const size_t t = t0 + ct; + const size_t qkv_k_base = (((batch * T + t) * Hq + qk_head) * K); + const size_t qkv_v_base = (((batch * T + t) * Hv + v_head) * V); + const size_t gb_idx = ((batch * T + t) * Hv + v_head); + + // Inline fp16→fp32 for q + { + const __fp16* src = q_data + qkv_k_base; + float* dst = q_chunk + ct * K; + size_t i = 0; + for (; i + 8 <= K; i += 8) { + float16x8_t in = vld1q_f16(src + i); + vst1q_f32(dst + i, vcvt_f32_f16(vget_low_f16(in))); + vst1q_f32(dst + i + 4, vcvt_f32_f16(vget_high_f16(in))); + } + for (; i < K; ++i) dst[i] = static_cast(src[i]); + } + + // Inline fp16→fp32 for k + { + const __fp16* src = k_data + qkv_k_base; + float* dst = k_chunk + ct * K; + size_t i = 0; + for (; i + 8 <= K; i += 8) { + float16x8_t in = vld1q_f16(src + i); + vst1q_f32(dst + i, vcvt_f32_f16(vget_low_f16(in))); + vst1q_f32(dst + i + 4, vcvt_f32_f16(vget_high_f16(in))); + } + for (; i < K; ++i) dst[i] = static_cast(src[i]); + } + + // Inline fp16→fp32 for v + { + const __fp16* src = v_data + qkv_v_base; + float* dst = v_chunk + ct * V; + size_t i = 0; + for (; i + 8 <= V; i += 8) { + float16x8_t in = vld1q_f16(src + i); + vst1q_f32(dst + i, vcvt_f32_f16(vget_low_f16(in))); + vst1q_f32(dst + i + 4, vcvt_f32_f16(vget_high_f16(in))); + } + for (; i < V; ++i) dst[i] = static_cast(src[i]); + } + + const float gate_log = static_cast(g_data[gb_idx]); + const float beta = static_cast(b_data[gb_idx]); + g_chunk[ct] = safe_exp_gate(std::isfinite(gate_log) ? gate_log : -20.0f); + beta_chunk[ct] = std::isfinite(beta) ? std::min(1.0f, std::max(0.0f, beta)) : 0.0f; + + // n0_proj fused here so k_chunk[ct] is still in L1 + const float* k_row = k_chunk + ct * K; + float* p_row = n0_proj + ct * V; + for (size_t kd = 0; kd < K; ++kd) { + const float kval = k_row[kd]; + const float* s_row = state + kd * V; + size_t vd = 0; + const float32x4_t kv4 = vdupq_n_f32(kval); + for (; vd + 8 <= V; vd += 8) { + float32x4_t p4 = vld1q_f32(p_row + vd); + const float32x4_t slo = vld1q_f32(s_row + vd); + const float32x4_t shi = vld1q_f32(s_row + vd + 4); + p4 = vfmaq_f32(p4, slo, kv4); + vst1q_f32(p_row + vd, p4); + p4 = vld1q_f32(p_row + vd + 4); + p4 = vfmaq_f32(p4, shi, kv4); + vst1q_f32(p_row + vd + 4, p4); + } + for (; vd < V; ++vd) { + p_row[vd] += kval * s_row[vd]; + } + } + } + + for (size_t t = 0; t < C; ++t) { + const float* kt = k_chunk + t * K; + for (size_t j = 0; j < t; ++j) { + const float* kj = k_chunk + j * K; + float dot = 0.0f; + size_t kd = 0; + float32x4_t acc4 = vdupq_n_f32(0.0f); + for (; kd + 8 <= K; kd += 8) { + const float32x4_t al = vld1q_f32(kt + kd); + const float32x4_t ah = vld1q_f32(kt + kd + 4); + const float32x4_t bl = vld1q_f32(kj + kd); + const float32x4_t bh = vld1q_f32(kj + kd + 4); + acc4 = vfmaq_f32(acc4, al, bl); + acc4 = vfmaq_f32(acc4, ah, bh); + } + dot += vaddvq_f32(acc4); + for (; kd < K; ++kd) { + dot += kt[kd] * kj[kd]; + } + gram[t * C + j] = dot; + } + } + + float running_p = 1.0f; + for (size_t t = 0; t < C; ++t) { + p_prev[t] = running_p; + running_p *= g_chunk[t]; + p_curr[t] = running_p; + inv_p_curr[t] = (std::fabs(running_p) > kMinAbs) ? (1.0f / running_p) : 0.0f; + } + + for (size_t t = 0; t < C; ++t) { + const float pp = p_curr[t]; + float* coeff_row = coeff + t * C; + for (size_t j = 0; j < t; ++j) { + const float inv_pc = inv_p_curr[j]; + coeff_row[j] = inv_pc == 0.0f ? 0.0f : (pp * inv_pc * gram[t * C + j]); + } + } + + for (size_t t = 0; t < C; ++t) { + const float pp = p_curr[t]; + const float bt = beta_chunk[t]; + const float* n0_row = n0_proj + t * V; + const float* v_row = v_chunk + t * V; + float* d_row = delta_chunk + t * V; + + size_t vd = 0; + const float32x4_t pp4 = vdupq_n_f32(pp); + const float32x4_t bt4 = vdupq_n_f32(bt); + for (; vd + 8 <= V; vd += 8) { + float32x4_t acc_lo = vmulq_f32(vld1q_f32(n0_row + vd), pp4); + float32x4_t acc_hi = vmulq_f32(vld1q_f32(n0_row + vd + 4), pp4); + for (size_t j = 0; j < t; ++j) { + const float c = coeff[t * C + j]; + const float* dj = delta_chunk + j * V + vd; + acc_lo = vmlaq_n_f32(acc_lo, vld1q_f32(dj), c); + acc_hi = vmlaq_n_f32(acc_hi, vld1q_f32(dj + 4), c); + } + vst1q_f32(d_row + vd, vmulq_f32(vsubq_f32(vld1q_f32(v_row + vd), acc_lo), bt4)); + vst1q_f32(d_row + vd + 4, vmulq_f32(vsubq_f32(vld1q_f32(v_row + vd + 4), acc_hi), bt4)); + } + for (; vd < V; ++vd) { + float acc = pp * n0_row[vd]; + for (size_t j = 0; j < t; ++j) { + acc += coeff[t * C + j] * delta_chunk[j * V + vd]; + } + d_row[vd] = bt * (v_row[vd] - acc); + } + } + + running_p = 1.0f; + for (size_t t = 0; t < C; ++t) { + const float gt = g_chunk[t]; + const float* k_row = k_chunk + t * K; + const float* q_row = q_chunk + t * K; + const float* d_row = delta_chunk + t * V; + + running_p *= gt; + const float p_run = running_p; + + for (size_t kd = 0; kd < K; ++kd) { + const float kval = k_row[kd]; + float* m_row = m_chunk + kd * V; + size_t vd = 0; + const float32x4_t gt4 = vdupq_n_f32(gt); + const float32x4_t kv4 = vdupq_n_f32(kval); + for (; vd + 8 <= V; vd += 8) { + float32x4_t ml = vld1q_f32(m_row + vd); + float32x4_t mh = vld1q_f32(m_row + vd + 4); + const float32x4_t dl = vld1q_f32(d_row + vd); + const float32x4_t dh = vld1q_f32(d_row + vd + 4); + ml = vmulq_f32(ml, gt4); + mh = vmulq_f32(mh, gt4); + ml = vfmaq_f32(ml, dl, kv4); + mh = vfmaq_f32(mh, dh, kv4); + vst1q_f32(m_row + vd, ml); + vst1q_f32(m_row + vd + 4, mh); + } + for (; vd < V; ++vd) { + m_row[vd] = gt * m_row[vd] + kval * d_row[vd]; + } + } + + const size_t out_base = (((batch * out_seq + (t0 + t)) * Hv + v_head) * V); + std::fill(out_acc, out_acc + V, 0.0f); + for (size_t kd = 0; kd < K; ++kd) { + const float qval = q_row[kd]; + const float* s_row = state + kd * V; + const float* m_row = m_chunk + kd * V; + size_t vd = 0; + const float32x4_t q4 = vdupq_n_f32(qval); + const float32x4_t p4 = vdupq_n_f32(p_run); + for (; vd + 8 <= V; vd += 8) { + float32x4_t acc4 = vld1q_f32(out_acc + vd); + float32x4_t acc4h = vld1q_f32(out_acc + vd + 4); + const float32x4_t sl = vld1q_f32(s_row + vd); + const float32x4_t sh = vld1q_f32(s_row + vd + 4); + const float32x4_t ml = vld1q_f32(m_row + vd); + const float32x4_t mh = vld1q_f32(m_row + vd + 4); + const float32x4_t nl = vfmaq_f32(ml, sl, p4); + const float32x4_t nh = vfmaq_f32(mh, sh, p4); + acc4 = vfmaq_f32(acc4, nl, q4); + acc4h = vfmaq_f32(acc4h, nh, q4); + vst1q_f32(out_acc + vd, acc4); + vst1q_f32(out_acc + vd + 4, acc4h); + } + for (; vd < V; ++vd) { + out_acc[vd] += qval * (p_run * s_row[vd] + m_row[vd]); + } + } + for (size_t vd = 0; vd < V; ++vd) { + out[out_base + vd] = static_cast<__fp16>(out_acc[vd] * scale); + } + } + + const float p_end = running_p; + for (size_t kd = 0; kd < K; ++kd) { + float* s_row = state + kd * V; + const float* m_row = m_chunk + kd * V; + size_t vd = 0; + const float32x4_t p4 = vdupq_n_f32(p_end); + for (; vd + 8 <= V; vd += 8) { + float32x4_t sl = vld1q_f32(s_row + vd); + float32x4_t sh = vld1q_f32(s_row + vd + 4); + const float32x4_t ml = vld1q_f32(m_row + vd); + const float32x4_t mh = vld1q_f32(m_row + vd + 4); + sl = vfmaq_f32(ml, sl, p4); + sh = vfmaq_f32(mh, sh, p4); + vst1q_f32(s_row + vd, sl); + vst1q_f32(s_row + vd + 4, sh); + } + for (; vd < V; ++vd) { + s_row[vd] = p_end * s_row[vd] + m_row[vd]; + } + } + } + + for (size_t kd = 0; kd < K; ++kd) { + const size_t out_idx = (((batch * out_seq + (T + kd)) * Hv + v_head) * V); + cactus_fp32_to_fp16(state + kd * V, out + out_idx, V); + } + } + }); +} + +} // namespace + +void cactus_gated_deltanet_decode_f16( + const __fp16* q_data, + const __fp16* k_data, + const __fp16* v_data, + const __fp16* g_data, + const __fp16* b_data, + const __fp16* s_data, + __fp16* out, + size_t B, + size_t Hq, + size_t Hv, + size_t K, + size_t V, + float scale) { + const size_t qk_repeat = Hv / Hq; + const size_t out_seq = 1 + K; + + CactusThreading::parallel_for(B * Hv, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t bh_start, size_t bh_end) { + std::vector state(K * V); + std::vector proj(V); + std::vector delta(V); + + for (size_t bh = bh_start; bh < bh_end; ++bh) { + const size_t batch = bh / Hv; + const size_t v_head = bh % Hv; + const size_t qk_head = v_head / qk_repeat; + + for (size_t kd = 0; kd < K; ++kd) { + for (size_t vd = 0; vd < V; ++vd) { + const size_t s_idx = (((batch * K + kd) * Hv + v_head) * V + vd); + state[kd * V + vd] = static_cast(s_data[s_idx]); + } + } + + float state_scale = 1.0f; + const size_t qk_base = ((batch * Hq + qk_head) * K); + const size_t v_base = ((batch * Hv + v_head) * V); + const float gate_log = static_cast(g_data[batch * Hv + v_head]); + const float beta = static_cast(b_data[batch * Hv + v_head]); + + gated_deltanet_step( + q_data + qk_base, + k_data + qk_base, + v_data + v_base, + gate_log, + beta, + scale, + K, + V, + state, + state_scale, + proj, + delta, + out + (((batch * out_seq) * Hv + v_head) * V)); + + fold_state_scale(state, state_scale); + for (size_t kd = 0; kd < K; ++kd) { + for (size_t vd = 0; vd < V; ++vd) { + const size_t out_idx = (((batch * out_seq + (1 + kd)) * Hv + v_head) * V + vd); + out[out_idx] = static_cast<__fp16>(state[kd * V + vd]); + } + } + } + }); +} + +void cactus_gated_deltanet_prefill_f16( + const __fp16* q_data, + const __fp16* k_data, + const __fp16* v_data, + const __fp16* g_data, + const __fp16* b_data, + const __fp16* s_data, + __fp16* out, + size_t B, + size_t T, + size_t Hq, + size_t Hv, + size_t K, + size_t V, + size_t requested_chunk_size, + float scale) { + const char* force_old = std::getenv("CACTUS_GATED_DELTANET_PREFILL_OLD"); + if (force_old != nullptr && std::atoi(force_old) != 0) { + gated_deltanet_prefill_old_f16(q_data, k_data, v_data, g_data, b_data, s_data, out, + B, T, Hq, Hv, K, V, scale); + return; + } + + if (requested_chunk_size <= 1) { + gated_deltanet_prefill_old_f16(q_data, k_data, v_data, g_data, b_data, s_data, out, + B, T, Hq, Hv, K, V, scale); + return; + } + + const size_t chunk_size = tuned_gated_deltanet_chunk_size(requested_chunk_size, K, V); + gated_deltanet_prefill_chunked_f16(q_data, k_data, v_data, g_data, b_data, s_data, out, + B, T, Hq, Hv, K, V, chunk_size, scale); +} diff --git a/cactus-kernels/src/image.cpp b/cactus-kernels/src/image.cpp new file mode 100644 index 000000000..2473af9b5 --- /dev/null +++ b/cactus-kernels/src/image.cpp @@ -0,0 +1,132 @@ +#define STB_IMAGE_IMPLEMENTATION +#define STB_IMAGE_RESIZE_IMPLEMENTATION + +#define STBI_NO_BMP +#define STBI_NO_PSD +#define STBI_NO_HDR +#define STBI_NO_PIC +#define STBI_NO_PNM +#define STBI_NO_TGA + +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wc99-extensions" +#pragma clang diagnostic ignored "-Wunused-parameter" +#elif defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wpedantic" +#pragma GCC diagnostic ignored "-Wunused-parameter" +#endif + +#include "stb_image.h" +#include "stb_image_resize2.h" + +#ifdef __clang__ +#pragma clang diagnostic pop +#elif defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + +#include "cactus_kernels.h" +#include +#include +#include +#include +#include + +unsigned char* cactus_image_load(const char* path, int* width, int* height, int* channels, int desired_channels) { + return stbi_load(path, width, height, channels, desired_channels); +} + +int cactus_image_info(const char* path, int* width, int* height, int* channels) { + return stbi_info(path, width, height, channels); +} + +void cactus_image_free(unsigned char* data) { + stbi_image_free(data); +} + +const char* cactus_image_failure_reason() { + return stbi_failure_reason(); +} + +void cactus_image_resize_uint8( + const unsigned char* input, int src_w, int src_h, + unsigned char* output, int dst_w, int dst_h, int channels) { + stbir_resize_uint8_linear(input, src_w, src_h, 0, + output, dst_w, dst_h, 0, + static_cast(channels)); +} + +void cactus_image_resize_float( + const float* input, int src_w, int src_h, + float* output, int dst_w, int dst_h, int channels) { + stbir_pixel_layout layout = (channels == 1) ? STBIR_1CHANNEL : + (channels == 3) ? STBIR_RGB : STBIR_RGBA; + stbir_resize_float_linear(input, src_w, src_h, 0, + output, dst_w, dst_h, 0, layout); +} + +void cactus_image_normalize( + const float* input, float* output, + int width, int height, int channels, + float rescale_factor, const float* mean, const float* std_dev) { + size_t total = static_cast(width) * height; + for (size_t i = 0; i < total; ++i) { + for (int c = 0; c < channels; ++c) { + size_t idx = i * channels + c; + float pixel = input[idx] * rescale_factor; + output[idx] = (pixel - mean[c]) / std_dev[c]; + } + } +} + +void cactus_image_to_patches( + const float* image, float* patches, + int width, int height, int channels, int patch_size) { + int ph = height / patch_size; + int pw = width / patch_size; + int patch_elements = patch_size * patch_size * channels; + + for (int py = 0; py < ph; ++py) { + for (int px = 0; px < pw; ++px) { + int patch_idx = py * pw + px; + for (int y = 0; y < patch_size; ++y) { + for (int x = 0; x < patch_size; ++x) { + int img_y = py * patch_size + y; + int img_x = px * patch_size + x; + int img_idx = (img_y * width + img_x) * channels; + int patch_offset = (y * patch_size + x) * channels; + for (int c = 0; c < channels; ++c) { + patches[patch_idx * patch_elements + patch_offset + c] = image[img_idx + c]; + } + } + } + } + } +} + +void cactus_image_convert_to_rgb( + const unsigned char* input, unsigned char* output, + int width, int height, int channels) { + int total = width * height; + if (channels == 1) { + for (int i = 0; i < total; ++i) { + output[i * 3 + 0] = input[i]; + output[i * 3 + 1] = input[i]; + output[i * 3 + 2] = input[i]; + } + } else if (channels == 4) { + for (int i = 0; i < total; ++i) { + output[i * 3 + 0] = input[i * 4 + 0]; + output[i * 3 + 1] = input[i * 4 + 1]; + output[i * 3 + 2] = input[i * 4 + 2]; + } + } else if (channels == 2) { + for (int i = 0; i < total; ++i) { + output[i * 3 + 0] = input[i * 2 + 0]; + output[i * 3 + 1] = input[i * 2 + 0]; + output[i * 3 + 2] = input[i * 2 + 0]; + } + } +} diff --git a/cactus-kernels/src/lstm.cpp b/cactus-kernels/src/lstm.cpp new file mode 100644 index 000000000..3889b3ff8 --- /dev/null +++ b/cactus-kernels/src/lstm.cpp @@ -0,0 +1,416 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include + +void cactus_lstm_cell_f16( + const __fp16* x_input, + const __fp16* h_prev, + const __fp16* c_prev, + const __fp16* weight_ih, + const __fp16* weight_hh, + const __fp16* bias_ih, + const __fp16* bias_hh, + __fp16* h_new, + __fp16* c_new, + size_t batch_size, + size_t input_size, + size_t hidden_size +) { + constexpr size_t SIMD_WIDTH = 8; + const size_t gate_size = 4 * hidden_size; + + std::vector<__fp16> gates_ih(batch_size * gate_size); + std::vector<__fp16> gates_hh(batch_size * gate_size); + + cactus_matmul_f16(x_input, weight_ih, gates_ih.data(), batch_size, input_size, gate_size); + cactus_matmul_f16(h_prev, weight_hh, gates_hh.data(), batch_size, hidden_size, gate_size); + + const size_t simd_end = (hidden_size / SIMD_WIDTH) * SIMD_WIDTH; + + CactusThreading::parallel_for(batch_size, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t batch_start, size_t batch_end) { + const float32x4_t one = vdupq_n_f32(1.0f); + + for (size_t b = batch_start; b < batch_end; ++b) { + const size_t gate_offset = b * gate_size; + const size_t hidden_offset = b * hidden_size; + + for (size_t h = 0; h < simd_end; h += SIMD_WIDTH) { + float16x8_t i_gate = vaddq_f16(vaddq_f16(vld1q_f16(&gates_ih[gate_offset + h]), + vld1q_f16(&gates_hh[gate_offset + h])), + vaddq_f16(vld1q_f16(&bias_ih[h]), + vld1q_f16(&bias_hh[h]))); + + float16x8_t f_gate = vaddq_f16(vaddq_f16(vld1q_f16(&gates_ih[gate_offset + hidden_size + h]), + vld1q_f16(&gates_hh[gate_offset + hidden_size + h])), + vaddq_f16(vld1q_f16(&bias_ih[hidden_size + h]), + vld1q_f16(&bias_hh[hidden_size + h]))); + + float16x8_t g_gate = vaddq_f16(vaddq_f16(vld1q_f16(&gates_ih[gate_offset + 2 * hidden_size + h]), + vld1q_f16(&gates_hh[gate_offset + 2 * hidden_size + h])), + vaddq_f16(vld1q_f16(&bias_ih[2 * hidden_size + h]), + vld1q_f16(&bias_hh[2 * hidden_size + h]))); + + float16x8_t o_gate = vaddq_f16(vaddq_f16(vld1q_f16(&gates_ih[gate_offset + 3 * hidden_size + h]), + vld1q_f16(&gates_hh[gate_offset + 3 * hidden_size + h])), + vaddq_f16(vld1q_f16(&bias_ih[3 * hidden_size + h]), + vld1q_f16(&bias_hh[3 * hidden_size + h]))); + + float32x4_t i_low = vcvt_f32_f16(vget_low_f16(i_gate)); + float32x4_t i_high = vcvt_f32_f16(vget_high_f16(i_gate)); + float16x8_t i_act = vcombine_f16(vcvt_f16_f32(vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(i_low))))), + vcvt_f16_f32(vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(i_high)))))); + + float32x4_t f_low = vcvt_f32_f16(vget_low_f16(f_gate)); + float32x4_t f_high = vcvt_f32_f16(vget_high_f16(f_gate)); + float16x8_t f_act = vcombine_f16(vcvt_f16_f32(vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(f_low))))), + vcvt_f16_f32(vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(f_high)))))); + + float32x4_t g_low = vcvt_f32_f16(vget_low_f16(g_gate)); + float32x4_t g_high = vcvt_f32_f16(vget_high_f16(g_gate)); + float16x8_t g_act = vcombine_f16(vcvt_f16_f32(fast_tanh_f32x4(g_low)), + vcvt_f16_f32(fast_tanh_f32x4(g_high))); + + float32x4_t o_low = vcvt_f32_f16(vget_low_f16(o_gate)); + float32x4_t o_high = vcvt_f32_f16(vget_high_f16(o_gate)); + float16x8_t o_act = vcombine_f16(vcvt_f16_f32(vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(o_low))))), + vcvt_f16_f32(vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(o_high)))))); + + float16x8_t c_prev_vec = vld1q_f16(&c_prev[hidden_offset + h]); + float16x8_t c_update = vfmaq_f16(vmulq_f16(f_act, c_prev_vec), i_act, g_act); + vst1q_f16(&c_new[hidden_offset + h], c_update); + + float32x4_t c_low = vcvt_f32_f16(vget_low_f16(c_update)); + float32x4_t c_high = vcvt_f32_f16(vget_high_f16(c_update)); + float16x8_t c_tanh = vcombine_f16(vcvt_f16_f32(fast_tanh_f32x4(c_low)), + vcvt_f16_f32(fast_tanh_f32x4(c_high))); + vst1q_f16(&h_new[hidden_offset + h], vmulq_f16(o_act, c_tanh)); + } + + for (size_t h = simd_end; h < hidden_size; ++h) { + float i_gate_val = static_cast(gates_ih[gate_offset + h] + + gates_hh[gate_offset + h] + + bias_ih[h] + bias_hh[h]); + float f_gate_val = static_cast(gates_ih[gate_offset + hidden_size + h] + + gates_hh[gate_offset + hidden_size + h] + + bias_ih[hidden_size + h] + bias_hh[hidden_size + h]); + float g_gate_val = static_cast(gates_ih[gate_offset + 2 * hidden_size + h] + + gates_hh[gate_offset + 2 * hidden_size + h] + + bias_ih[2 * hidden_size + h] + bias_hh[2 * hidden_size + h]); + float o_gate_val = static_cast(gates_ih[gate_offset + 3 * hidden_size + h] + + gates_hh[gate_offset + 3 * hidden_size + h] + + bias_ih[3 * hidden_size + h] + bias_hh[3 * hidden_size + h]); + + float i_act = 1.0f / (1.0f + expf(-i_gate_val)); + float f_act = 1.0f / (1.0f + expf(-f_gate_val)); + float g_act = tanhf(g_gate_val); + float o_act = 1.0f / (1.0f + expf(-o_gate_val)); + + float c_val = f_act * static_cast(c_prev[hidden_offset + h]) + i_act * g_act; + c_new[hidden_offset + h] = static_cast<__fp16>(c_val); + h_new[hidden_offset + h] = static_cast<__fp16>(o_act * tanhf(c_val)); + } + } + }); +} +#ifdef __APPLE__ +#include +#endif + +static void apply_lstm_gates_f32( + const float* __restrict gates, + float* __restrict c, + float* __restrict h, + size_t hidden_size +) { + const float32x4_t one = vdupq_n_f32(1.0f); + size_t i = 0; + + for (; i + 4 <= hidden_size; i += 4) { + float32x4_t i_raw = vld1q_f32(gates + i); + float32x4_t f_raw = vld1q_f32(gates + hidden_size + i); + float32x4_t g_raw = vld1q_f32(gates + 2 * hidden_size + i); + float32x4_t o_raw = vld1q_f32(gates + 3 * hidden_size + i); + + float32x4_t i_act = vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(i_raw)))); + float32x4_t f_act = vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(f_raw)))); + float32x4_t g_act = fast_tanh_f32x4(g_raw); + float32x4_t o_act = vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(o_raw)))); + + float32x4_t c_prev = vld1q_f32(c + i); + float32x4_t c_new = vfmaq_f32(vmulq_f32(f_act, c_prev), i_act, g_act); + vst1q_f32(c + i, c_new); + + vst1q_f32(h + i, vmulq_f32(o_act, fast_tanh_f32x4(c_new))); + } + + for (; i < hidden_size; ++i) { + float ig = 1.0f / (1.0f + expf(-gates[i])); + float fg = 1.0f / (1.0f + expf(-gates[hidden_size + i])); + float gg = tanhf(gates[2 * hidden_size + i]); + float og = 1.0f / (1.0f + expf(-gates[3 * hidden_size + i])); + float cv = fg * c[i] + ig * gg; + c[i] = cv; + h[i] = og * tanhf(cv); + } +} + +#ifndef __APPLE__ +static void apply_lstm_gates_f16( + const __fp16* __restrict gates, + __fp16* __restrict c, + __fp16* __restrict h, + size_t hidden_size +) { + const size_t gate_size = 4 * hidden_size; + std::vector gates_f32(gate_size); + std::vector c_f32(hidden_size); + std::vector h_f32(hidden_size); + + cactus_fp16_to_fp32(gates, gates_f32.data(), gate_size); + cactus_fp16_to_fp32(c, c_f32.data(), hidden_size); + + apply_lstm_gates_f32(gates_f32.data(), c_f32.data(), h_f32.data(), hidden_size); + + cactus_fp32_to_fp16(c_f32.data(), c, hidden_size); + cactus_fp32_to_fp16(h_f32.data(), h, hidden_size); +} + +static inline float hsum_f16x8_f32(float16x8_t v) { + float16x4_t lo = vget_low_f16(v); + float16x4_t hi = vget_high_f16(v); + float16x4_t s4 = vadd_f16(lo, hi); + float16x4_t s2 = vadd_f16(s4, vext_f16(s4, s4, 2)); + float16x4_t s1 = vadd_f16(s2, vext_f16(s2, s2, 1)); + return static_cast(vget_lane_f16(s1, 0)); +} + +static void lstm_gemv_f16_neon( + const __fp16* __restrict x, + const __fp16* __restrict W, + const __fp16* __restrict bias, + __fp16* __restrict out, + size_t K, size_t N +) { + const size_t K16 = (K / 16) * 16; + const size_t K8 = (K / 8) * 8; + size_t n = 0; + for (; n + 4 <= N; n += 4) { + const __fp16* w0 = W + n * K; + const __fp16* w1 = W + (n + 1) * K; + const __fp16* w2 = W + (n + 2) * K; + const __fp16* w3 = W + (n + 3) * K; + float16x8_t a0 = vdupq_n_f16(0); + float16x8_t a1 = vdupq_n_f16(0); + float16x8_t a2 = vdupq_n_f16(0); + float16x8_t a3 = vdupq_n_f16(0); + for (size_t k = 0; k < K16; k += 16) { + float16x8_t xlo = vld1q_f16(x + k); + float16x8_t xhi = vld1q_f16(x + k + 8); + a0 = vfmaq_f16(a0, xlo, vld1q_f16(w0 + k)); + a0 = vfmaq_f16(a0, xhi, vld1q_f16(w0 + k + 8)); + a1 = vfmaq_f16(a1, xlo, vld1q_f16(w1 + k)); + a1 = vfmaq_f16(a1, xhi, vld1q_f16(w1 + k + 8)); + a2 = vfmaq_f16(a2, xlo, vld1q_f16(w2 + k)); + a2 = vfmaq_f16(a2, xhi, vld1q_f16(w2 + k + 8)); + a3 = vfmaq_f16(a3, xlo, vld1q_f16(w3 + k)); + a3 = vfmaq_f16(a3, xhi, vld1q_f16(w3 + k + 8)); + } + for (size_t k = K16; k < K8; k += 8) { + float16x8_t xv = vld1q_f16(x + k); + a0 = vfmaq_f16(a0, xv, vld1q_f16(w0 + k)); + a1 = vfmaq_f16(a1, xv, vld1q_f16(w1 + k)); + a2 = vfmaq_f16(a2, xv, vld1q_f16(w2 + k)); + a3 = vfmaq_f16(a3, xv, vld1q_f16(w3 + k)); + } + float s0 = hsum_f16x8_f32(a0); + float s1 = hsum_f16x8_f32(a1); + float s2 = hsum_f16x8_f32(a2); + float s3 = hsum_f16x8_f32(a3); + for (size_t k = K8; k < K; ++k) { + float xv = static_cast(x[k]); + s0 += xv * static_cast(w0[k]); + s1 += xv * static_cast(w1[k]); + s2 += xv * static_cast(w2[k]); + s3 += xv * static_cast(w3[k]); + } + out[n] = static_cast<__fp16>(s0 + static_cast(bias[n])); + out[n + 1] = static_cast<__fp16>(s1 + static_cast(bias[n + 1])); + out[n + 2] = static_cast<__fp16>(s2 + static_cast(bias[n + 2])); + out[n + 3] = static_cast<__fp16>(s3 + static_cast(bias[n + 3])); + } + for (; n < N; ++n) { + const __fp16* w = W + n * K; + float16x8_t acc = vdupq_n_f16(0); + for (size_t k = 0; k < K8; k += 8) + acc = vfmaq_f16(acc, vld1q_f16(x + k), vld1q_f16(w + k)); + float s = hsum_f16x8_f32(acc); + for (size_t k = K8; k < K; ++k) + s += static_cast(x[k]) * static_cast(w[k]); + out[n] = static_cast<__fp16>(s + static_cast(bias[n])); + } +} +#endif + +void cactus_bilstm_sequence_f16( + const __fp16* input, + const __fp16* weight_ih_fwd, + const __fp16* weight_hh_fwd, + const __fp16* bias_ih_fwd, + const __fp16* bias_hh_fwd, + const __fp16* weight_ih_bwd, + const __fp16* weight_hh_bwd, + const __fp16* bias_ih_bwd, + const __fp16* bias_hh_bwd, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t input_size, + size_t hidden_size +) { + const size_t gate_size = 4 * hidden_size; + const size_t combined_K = input_size + hidden_size; + const size_t output_size = 2 * hidden_size; + + std::vector bias_fwd_f32(gate_size); + std::vector bias_bwd_f32(gate_size); + for (size_t g = 0; g < gate_size; ++g) { + bias_fwd_f32[g] = static_cast(bias_ih_fwd[g]) + static_cast(bias_hh_fwd[g]); + bias_bwd_f32[g] = static_cast(bias_ih_bwd[g]) + static_cast(bias_hh_bwd[g]); + } + +#ifdef __APPLE__ + std::vector W_fwd_f32(gate_size * combined_K); + std::vector W_bwd_f32(gate_size * combined_K); + for (size_t g = 0; g < gate_size; ++g) { + cactus_fp16_to_fp32(weight_ih_fwd + g * input_size, + W_fwd_f32.data() + g * combined_K, input_size); + cactus_fp16_to_fp32(weight_hh_fwd + g * hidden_size, + W_fwd_f32.data() + g * combined_K + input_size, hidden_size); + cactus_fp16_to_fp32(weight_ih_bwd + g * input_size, + W_bwd_f32.data() + g * combined_K, input_size); + cactus_fp16_to_fp32(weight_hh_bwd + g * hidden_size, + W_bwd_f32.data() + g * combined_K + input_size, hidden_size); + } + + for (size_t b = 0; b < batch_size; ++b) { + const __fp16* batch_in = input + b * seq_len * input_size; + __fp16* batch_out = output + b * seq_len * output_size; + + auto& pool = CactusThreading::get_thread_pool(); + auto bwd_future = pool.enqueue([&, batch_in, batch_out]() { + std::vector xh(combined_K); + std::vector gates(gate_size); + std::vector h(hidden_size, 0.0f); + std::vector c(hidden_size, 0.0f); + + for (size_t idx = 0; idx < seq_len; ++idx) { + const size_t t = seq_len - 1 - idx; + cactus_fp16_to_fp32(batch_in + t * input_size, xh.data(), input_size); + memcpy(xh.data() + input_size, h.data(), hidden_size * sizeof(float)); + + memcpy(gates.data(), bias_bwd_f32.data(), gate_size * sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + static_cast(gate_size), static_cast(combined_K), + 1.0f, W_bwd_f32.data(), static_cast(combined_K), + xh.data(), 1, 1.0f, gates.data(), 1); + + apply_lstm_gates_f32(gates.data(), c.data(), h.data(), hidden_size); + + cactus_fp32_to_fp16(h.data(), batch_out + t * output_size + hidden_size, hidden_size); + } + }); + + { + std::vector xh(combined_K); + std::vector gates(gate_size); + std::vector h(hidden_size, 0.0f); + std::vector c(hidden_size, 0.0f); + + for (size_t t = 0; t < seq_len; ++t) { + cactus_fp16_to_fp32(batch_in + t * input_size, xh.data(), input_size); + memcpy(xh.data() + input_size, h.data(), hidden_size * sizeof(float)); + + memcpy(gates.data(), bias_fwd_f32.data(), gate_size * sizeof(float)); + cblas_sgemv(CblasRowMajor, CblasNoTrans, + static_cast(gate_size), static_cast(combined_K), + 1.0f, W_fwd_f32.data(), static_cast(combined_K), + xh.data(), 1, 1.0f, gates.data(), 1); + + apply_lstm_gates_f32(gates.data(), c.data(), h.data(), hidden_size); + + cactus_fp32_to_fp16(h.data(), batch_out + t * output_size, hidden_size); + } + } + + bwd_future.get(); + } + +#else + std::vector<__fp16> W_fwd_f16(gate_size * combined_K); + std::vector<__fp16> W_bwd_f16(gate_size * combined_K); + for (size_t g = 0; g < gate_size; ++g) { + memcpy(W_fwd_f16.data() + g * combined_K, + weight_ih_fwd + g * input_size, input_size * sizeof(__fp16)); + memcpy(W_fwd_f16.data() + g * combined_K + input_size, + weight_hh_fwd + g * hidden_size, hidden_size * sizeof(__fp16)); + memcpy(W_bwd_f16.data() + g * combined_K, + weight_ih_bwd + g * input_size, input_size * sizeof(__fp16)); + memcpy(W_bwd_f16.data() + g * combined_K + input_size, + weight_hh_bwd + g * hidden_size, hidden_size * sizeof(__fp16)); + } + + std::vector<__fp16> bias_fwd_f16(gate_size); + std::vector<__fp16> bias_bwd_f16(gate_size); + for (size_t g = 0; g < gate_size; ++g) { + bias_fwd_f16[g] = static_cast<__fp16>(bias_fwd_f32[g]); + bias_bwd_f16[g] = static_cast<__fp16>(bias_bwd_f32[g]); + } + + for (size_t b = 0; b < batch_size; ++b) { + const __fp16* batch_in = input + b * seq_len * input_size; + __fp16* batch_out = output + b * seq_len * output_size; + + auto& pool = CactusThreading::get_thread_pool(); + auto bwd_future = pool.enqueue([&, batch_in, batch_out]() { + std::vector<__fp16> xh(combined_K); + std::vector<__fp16> gates(gate_size); + std::vector<__fp16> h(hidden_size, static_cast<__fp16>(0)); + std::vector<__fp16> c(hidden_size, static_cast<__fp16>(0)); + + for (size_t idx = 0; idx < seq_len; ++idx) { + const size_t t = seq_len - 1 - idx; + memcpy(xh.data(), batch_in + t * input_size, input_size * sizeof(__fp16)); + memcpy(xh.data() + input_size, h.data(), hidden_size * sizeof(__fp16)); + lstm_gemv_f16_neon(xh.data(), W_bwd_f16.data(), bias_bwd_f16.data(), + gates.data(), combined_K, gate_size); + apply_lstm_gates_f16(gates.data(), c.data(), h.data(), hidden_size); + memcpy(batch_out + t * output_size + hidden_size, + h.data(), hidden_size * sizeof(__fp16)); + } + }); + + { + std::vector<__fp16> xh(combined_K); + std::vector<__fp16> gates(gate_size); + std::vector<__fp16> h(hidden_size, static_cast<__fp16>(0)); + std::vector<__fp16> c(hidden_size, static_cast<__fp16>(0)); + + for (size_t t = 0; t < seq_len; ++t) { + memcpy(xh.data(), batch_in + t * input_size, input_size * sizeof(__fp16)); + memcpy(xh.data() + input_size, h.data(), hidden_size * sizeof(__fp16)); + lstm_gemv_f16_neon(xh.data(), W_fwd_f16.data(), bias_fwd_f16.data(), + gates.data(), combined_K, gate_size); + apply_lstm_gates_f16(gates.data(), c.data(), h.data(), hidden_size); + memcpy(batch_out + t * output_size, + h.data(), hidden_size * sizeof(__fp16)); + } + } + + bwd_future.get(); + } +#endif +} diff --git a/cactus-kernels/src/matmul.cpp b/cactus-kernels/src/matmul.cpp new file mode 100644 index 000000000..2873ac911 --- /dev/null +++ b/cactus-kernels/src/matmul.cpp @@ -0,0 +1,2768 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +constexpr size_t ACCELERATE_M_THRESHOLD = 4; +constexpr size_t ACCELERATE_K_THRESHOLD = 256; +#endif + +// Do NOT Remove: Uncomment for testing on various paths +// ----- +// TEMPORARY: Force fallback path for testing on DOTPROD devices +// #undef __ARM_FEATURE_DOTPROD + +#if defined(__ARM_FEATURE_DOTPROD) + #define CACTUS_DOTQ_LANE(acc, b, a, lane) vdotq_laneq_s32(acc, b, a, lane) +#else + static inline int32x4_t cactus_dotq_with_pattern(int32x4_t acc, int8x16_t b, int8x8_t a_pattern) { + int8x8_t b_lo = vget_low_s8(b); + int8x8_t b_hi = vget_high_s8(b); + + int16x8_t prod_lo = vmull_s8(b_lo, a_pattern); + int16x8_t prod_hi = vmull_s8(b_hi, a_pattern); + + int32x4_t sum_lo = vpaddlq_s16(prod_lo); + int32x4_t sum_hi = vpaddlq_s16(prod_hi); + + int32x2_t final_lo = vpadd_s32(vget_low_s32(sum_lo), vget_high_s32(sum_lo)); + int32x2_t final_hi = vpadd_s32(vget_low_s32(sum_hi), vget_high_s32(sum_hi)); + + return vaddq_s32(acc, vcombine_s32(final_lo, final_hi)); + } + + static inline int32x4_t cactus_dotq_lane0(int32x4_t acc, int8x16_t b, int8x16_t a) { + int8x8_t a_lo = vget_low_s8(a); + int8x8_t a_pattern = vreinterpret_s8_s32(vdup_lane_s32(vreinterpret_s32_s8(a_lo), 0)); + return cactus_dotq_with_pattern(acc, b, a_pattern); + } + + static inline int32x4_t cactus_dotq_lane1(int32x4_t acc, int8x16_t b, int8x16_t a) { + int8x8_t a_lo = vget_low_s8(a); + int8x8_t a_pattern = vreinterpret_s8_s32(vdup_lane_s32(vreinterpret_s32_s8(a_lo), 1)); + return cactus_dotq_with_pattern(acc, b, a_pattern); + } + + static inline int32x4_t cactus_dotq_lane2(int32x4_t acc, int8x16_t b, int8x16_t a) { + int8x8_t a_hi = vget_high_s8(a); + int8x8_t a_pattern = vreinterpret_s8_s32(vdup_lane_s32(vreinterpret_s32_s8(a_hi), 0)); + return cactus_dotq_with_pattern(acc, b, a_pattern); + } + + static inline int32x4_t cactus_dotq_lane3(int32x4_t acc, int8x16_t b, int8x16_t a) { + int8x8_t a_hi = vget_high_s8(a); + int8x8_t a_pattern = vreinterpret_s8_s32(vdup_lane_s32(vreinterpret_s32_s8(a_hi), 1)); + return cactus_dotq_with_pattern(acc, b, a_pattern); + } + + #define CACTUS_DOTQ_LANE(acc, b, a, lane) cactus_dotq_lane##lane(acc, b, a) +#endif + +static inline __fp16 hsum_f16x8(float16x8_t v) { + float16x4_t lo = vget_low_f16(v); + float16x4_t hi = vget_high_f16(v); + float16x4_t sum4 = vadd_f16(lo, hi); + float16x4_t sum2 = vadd_f16(sum4, vext_f16(sum4, sum4, 2)); + float16x4_t sum1 = vadd_f16(sum2, vext_f16(sum2, sum2, 1)); + return vget_lane_f16(sum1, 0); +} + +namespace { + +static inline float16x8_t cactus_quant_signs_to_f16(const int8_t* signs, uint32_t offset) { + if (signs == nullptr) return vdupq_n_f16(1); + return vcvtq_f16_s16(vmovl_s8(vld1_s8(signs + offset))); +} + +static inline float16x8_t cactus_quant_input_scale_recip8(const CactusQuantMatrix& W, uint32_t offset) { + if (W.input_scale_recip != nullptr) { + return vld1q_f16(W.input_scale_recip + offset); + } + if (W.input_scale != nullptr) { + return vdivq_f16(vdupq_n_f16(1), vld1q_f16(W.input_scale + offset)); + } + return vdupq_n_f16(1); +} + +static inline __fp16 cactus_quant_input_scale_recip1(const CactusQuantMatrix& W, uint32_t offset) { + if (W.input_scale_recip != nullptr) return W.input_scale_recip[offset]; + if (W.input_scale != nullptr) return static_cast<__fp16>(1.0f / static_cast(W.input_scale[offset])); + return static_cast<__fp16>(1); +} + +static inline const __fp16* cactus_quant_scale_ptr(const CactusQuantMatrix& W, uint32_t row, uint32_t group) { + return W.norms + static_cast(row) * W.num_groups + group; +} + +static inline const uint8_t* cactus_quant_packed_chunk_ptr( + const CactusQuantMatrix& W, + uint32_t row, + uint32_t group, + uint32_t k) { + return W.packed_indices + + (((static_cast(row) * W.num_groups + group) + * cactus_quant_packed_group_bytes(W.bits, W.group_size)) + + (static_cast(k) * W.bits) / 8); +} + +static inline void tq_interleave_4x_s8(const int8x16_t row0, const int8x16_t row1, + const int8x16_t row2, const int8x16_t row3, + int8_t* dst) { + int32x4_t r0 = vreinterpretq_s32_s8(row0), r1 = vreinterpretq_s32_s8(row1); + int32x4_t r2 = vreinterpretq_s32_s8(row2), r3 = vreinterpretq_s32_s8(row3); + int32x4_t t01l = vzip1q_s32(r0, r1), t01h = vzip2q_s32(r0, r1); + int32x4_t t23l = vzip1q_s32(r2, r3), t23h = vzip2q_s32(r2, r3); + vst1q_s8(dst, vreinterpretq_s8_s64(vzip1q_s64(vreinterpretq_s64_s32(t01l), vreinterpretq_s64_s32(t23l)))); + vst1q_s8(dst + 16, vreinterpretq_s8_s64(vzip2q_s64(vreinterpretq_s64_s32(t01l), vreinterpretq_s64_s32(t23l)))); + vst1q_s8(dst + 32, vreinterpretq_s8_s64(vzip1q_s64(vreinterpretq_s64_s32(t01h), vreinterpretq_s64_s32(t23h)))); + vst1q_s8(dst + 48, vreinterpretq_s8_s64(vzip2q_s64(vreinterpretq_s64_s32(t01h), vreinterpretq_s64_s32(t23h)))); +} + +static inline float16x8_t cactus_tq4_lookup_codebook8(uint8x8_t nibbles, uint8x16x2_t cb_bytes) { + uint8x8_t byte_offsets = vshl_n_u8(nibbles, 1); + uint8x8_t byte_offsets_hi = vadd_u8(byte_offsets, vdup_n_u8(1)); + uint8x8x2_t zipped = vzip_u8(byte_offsets, byte_offsets_hi); + uint8x16_t byte_idx = vcombine_u8(zipped.val[0], zipped.val[1]); + return vreinterpretq_f16_u8(vqtbl2q_u8(cb_bytes, byte_idx)); +} + +static inline uint8x8_t cactus_tq2_unpack_8x2bit_le(uint8_t b0, uint8_t b1) { + uint64_t idx_word = + ((uint64_t)((b0 ) & 0x3) ) | + ((uint64_t)((b0 >> 2) & 0x3) << 8) | + ((uint64_t)((b0 >> 4) & 0x3) << 16) | + ((uint64_t)((b0 >> 6) & 0x3) << 24) | + ((uint64_t)((b1 ) & 0x3) << 32) | + ((uint64_t)((b1 >> 2) & 0x3) << 40) | + ((uint64_t)((b1 >> 4) & 0x3) << 48) | + ((uint64_t)((b1 >> 6) & 0x3) << 56); + return vcreate_u8(idx_word); +} + +static inline float16x8_t cactus_tq2_lookup_codebook8(uint8x8_t indices, uint8x8_t cb_bytes) { + uint8x8_t off_lo = vshl_n_u8(indices, 1); + uint8x8_t off_hi = vadd_u8(off_lo, vdup_n_u8(1)); + uint8x8x2_t zipped = vzip_u8(off_lo, off_hi); + uint8x16_t byte_idx = vcombine_u8(zipped.val[0], zipped.val[1]); + uint8x16_t lut = vcombine_u8(cb_bytes, cb_bytes); + return vreinterpretq_f16_u8(vqtbl1q_u8(lut, byte_idx)); +} + +static void cactus_quant_fwht128_f16(__fp16* x) { + float16x8_t v[16]; + for (int i = 0; i < 16; ++i) v[i] = vld1q_f16(x + i * 8); + for (int i = 0; i < 16; ++i) { + float16x8_t r = vreinterpretq_f16_u16(vrev32q_u16(vreinterpretq_u16_f16(v[i]))); + float16x8_t s = vaddq_f16(v[i], r); + float16x8_t d = vsubq_f16(v[i], r); + v[i] = vreinterpretq_f16_u16(vtrn1q_u16(vreinterpretq_u16_f16(s), vreinterpretq_u16_f16(d))); + } + for (int i = 0; i < 16; ++i) { + float32x4_t f32 = vreinterpretq_f32_f16(v[i]); + float16x8_t a = vreinterpretq_f16_f32(vtrn1q_f32(f32, f32)); + float16x8_t b = vreinterpretq_f16_f32(vtrn2q_f32(f32, f32)); + float16x8_t s = vaddq_f16(a, b); + float16x8_t d = vsubq_f16(a, b); + v[i] = vreinterpretq_f16_f32(vtrn1q_f32(vreinterpretq_f32_f16(s), vreinterpretq_f32_f16(d))); + } + for (int i = 0; i < 16; ++i) { + float16x4_t lo = vget_low_f16(v[i]); + float16x4_t hi = vget_high_f16(v[i]); + v[i] = vcombine_f16(vadd_f16(lo, hi), vsub_f16(lo, hi)); + } + + auto pass = [&](int s) { + for (int base = 0; base < 16; base += (s << 1)) { + for (int j = 0; j < s; ++j) { + float16x8_t a = v[base + j]; + float16x8_t b = v[base + j + s]; + v[base + j] = vaddq_f16(a, b); + v[base + j + s] = vsubq_f16(a, b); + } + } + }; + pass(1); + pass(2); + pass(4); + pass(8); + + float16x8_t inv = vdupq_n_f16(static_cast<__fp16>(1.0f / std::sqrt(128.0f))); + for (int i = 0; i < 16; ++i) { + vst1q_f16(x + i * 8, vmulq_f16(v[i], inv)); + } +} + +static void cactus_quant_fwht_f16(__fp16* x, uint32_t n) { + for (uint32_t h = 1; h < n; h <<= 1) { + for (uint32_t i = 0; i < n; i += (h << 1)) { + for (uint32_t j = i; j < i + h; j += 8) { + if (j + 8 <= i + h) { + float16x8_t a = vld1q_f16(x + j); + float16x8_t b = vld1q_f16(x + j + h); + vst1q_f16(x + j, vaddq_f16(a, b)); + vst1q_f16(x + j + h, vsubq_f16(a, b)); + } else { + for (uint32_t k = j; k < i + h; ++k) { + __fp16 a = x[k]; + __fp16 b = x[k + h]; + x[k] = static_cast<__fp16>(a + b); + x[k + h] = static_cast<__fp16>(a - b); + } + } + } + } + } + + const float16x8_t inv_v = vdupq_n_f16(static_cast<__fp16>(1.0f / std::sqrt(static_cast(n)))); + uint32_t k = 0; + for (; k + 8 <= n; k += 8) { + vst1q_f16(x + k, vmulq_f16(vld1q_f16(x + k), inv_v)); + } + const __fp16 inv = static_cast<__fp16>(1.0f / std::sqrt(static_cast(n))); + for (; k < n; ++k) { + x[k] = static_cast<__fp16>(x[k] * inv); + } +} + +static void cactus_quant_transform_hadamard_group( + const CactusQuantMatrix& W, + const __fp16* x_group, + uint32_t group, + __fp16* code_basis) { + const uint32_t gs = W.group_size; + __fp16 tmp[256]; + __fp16* work = (W.permutation == nullptr) ? code_basis : tmp; + + uint32_t k = 0; + for (; k + 8 <= gs; k += 8) { + const uint32_t offset = group * gs + k; + float16x8_t x_v = vld1q_f16(x_group + k); + x_v = vmulq_f16(x_v, cactus_quant_input_scale_recip8(W, offset)); + float16x8_t s_v = cactus_quant_signs_to_f16(W.left_signs, k); + vst1q_f16(work + k, vmulq_f16(x_v, s_v)); + } + for (; k < gs; ++k) { + const uint32_t offset = group * gs + k; + const float sign = W.left_signs ? static_cast(W.left_signs[k]) : 1.0f; + const float scale = static_cast(cactus_quant_input_scale_recip1(W, offset)); + work[k] = static_cast<__fp16>(static_cast(x_group[k]) * scale * sign); + } + + if (gs == 128) { + cactus_quant_fwht128_f16(work); + } else { + cactus_quant_fwht_f16(work, gs); + } + + k = 0; + for (; k + 8 <= gs; k += 8) { + float16x8_t w_v = vld1q_f16(work + k); + float16x8_t s_v = cactus_quant_signs_to_f16(W.right_signs, k); + vst1q_f16(work + k, vmulq_f16(w_v, s_v)); + } + for (; k < gs; ++k) { + const float sign = W.right_signs ? static_cast(W.right_signs[k]) : 1.0f; + work[k] = static_cast<__fp16>(static_cast(work[k]) * sign); + } + + if (work != code_basis) { + for (uint32_t j = 0; j < gs; ++j) { + code_basis[j] = work[W.permutation[j]]; + } + } +} + +static void cactus_quant_transform_hadamard_activations( + const CactusQuantMatrix& W, + const __fp16* A, + uint32_t M, + __fp16* code_basis) { + const size_t work_items = static_cast(M) * W.num_groups; + CactusThreading::parallel_for( + work_items, + CactusThreading::ParallelConfig{16, 1}, + [&](size_t start, size_t end) { + for (size_t idx = start; idx < end; ++idx) { + const size_t m = idx / W.num_groups; + const size_t g = idx - m * W.num_groups; + cactus_quant_transform_hadamard_group( + W, + A + m * W.K + g * W.group_size, + static_cast(g), + code_basis + m * W.K + g * W.group_size); + } + }); +} + +template +static void cactus_quant_parallel_ranges(size_t total_work, size_t work_per_thread, WorkFunc work_func) { + if (total_work == 0) return; + if (work_per_thread == 0) work_per_thread = 1; + + auto& pool = CactusThreading::get_thread_pool(); + size_t num_threads = std::min(pool.num_workers(), (total_work + work_per_thread - 1) / work_per_thread); + num_threads = std::min(num_threads, total_work); + if (num_threads <= 1) { + work_func(0, total_work); + return; + } + + pool.enqueue_n_threads(total_work, num_threads, work_func); + pool.wait_all(); +} + +static size_t cactus_quant_gemv_sb_per_thread() { + static const size_t v = [] { + const char* e = getenv("CACTUS_GEMV_SB_PER_THREAD"); + const int i = e ? atoi(e) : 8; + return static_cast(i > 0 ? i : 8); + }(); + return v; +} + +template +static void cactus_quant_two_phase_run(size_t nt, uint32_t num_groups, uint32_t n_items, + PhaseA&& phase_a, PhaseB&& phase_b) { + std::atomic next_group{0}; + std::atomic groups_done{0}; + std::atomic next_item{0}; + std::atomic workers_done{0}; + auto worker = [&](size_t wid) { + for (uint32_t g; (g = next_group.fetch_add(1, std::memory_order_relaxed)) < num_groups; ) { + phase_a(g); + groups_done.fetch_add(1, std::memory_order_release); + } + while (groups_done.load(std::memory_order_acquire) < num_groups) {} + for (;;) { + const uint32_t seen = next_item.load(std::memory_order_relaxed); + if (seen >= n_items) break; + const uint32_t want = (n_items - seen > 4u * nt) ? 4u : 1u; + const uint32_t item = next_item.fetch_add(want, std::memory_order_relaxed); + if (item >= n_items) break; + phase_b(wid, item, std::min(want, n_items - item)); + } + }; + auto& pool = CactusThreading::get_thread_pool(); + // Main joins as worker 0 and spins: a cv sleep costs ~5-10us per call, material at decode rates. + pool.enqueue_n_threads(nt - 1, nt - 1, [&](size_t wid, size_t) { + worker(wid + 1); + workers_done.fetch_add(1, std::memory_order_release); + }); + worker(0); + while (workers_done.load(std::memory_order_acquire) < nt - 1) {} +} + +static void cactus_quant_matmul_f32_segment_accum( + const __fp16* __restrict__ A, + size_t a_stride, + const __fp16* __restrict__ B_tile, + float* __restrict__ C_f32, + size_t M, + size_t Kseg, + size_t actual_n) { + constexpr size_t TILE_M = 4; + constexpr size_t TILE_N_MAX = 16; + + for (size_t m_start = 0; m_start < M; m_start += TILE_M) { + const size_t actual_m = std::min(TILE_M, M - m_start); + + float16x8_t acc[TILE_M][TILE_N_MAX]; + for (size_t mi = 0; mi < actual_m; ++mi) + for (size_t ni = 0; ni < actual_n; ++ni) + acc[mi][ni] = vdupq_n_f16(0); + + for (size_t k = 0; k < Kseg; k += 16) { + float16x8_t a_lo[TILE_M], a_hi[TILE_M]; + for (size_t mi = 0; mi < actual_m; ++mi) { + const __fp16* ap = A + (m_start + mi) * a_stride + k; + a_lo[mi] = vld1q_f16(ap); + a_hi[mi] = vld1q_f16(ap + 8); + } + for (size_t ni = 0; ni < actual_n; ++ni) { + const __fp16* bp = B_tile + ni * Kseg + k; + float16x8_t b_lo = vld1q_f16(bp); + float16x8_t b_hi = vld1q_f16(bp + 8); + for (size_t mi = 0; mi < actual_m; ++mi) { + acc[mi][ni] = vfmaq_f16(acc[mi][ni], a_lo[mi], b_lo); + acc[mi][ni] = vfmaq_f16(acc[mi][ni], a_hi[mi], b_hi); + } + } + } + + for (size_t mi = 0; mi < actual_m; ++mi) + for (size_t ni = 0; ni < actual_n; ++ni) + C_f32[(m_start + mi) * TILE_N_MAX + ni] += static_cast(hsum_f16x8(acc[mi][ni])); + } +} + +struct CactusTQ4ScaledDecoder { + uint8x16x2_t cb_bytes; + + explicit CactusTQ4ScaledDecoder(const CactusQuantMatrix& W) { + cb_bytes.val[0] = vld1q_u8(reinterpret_cast(W.codebook)); + cb_bytes.val[1] = vld1q_u8(reinterpret_cast(W.codebook) + 16); + } + + void operator()(const CactusQuantMatrix& W, uint32_t row, uint32_t group, __fp16* dst) const { + const __fp16 rn = *cactus_quant_scale_ptr(W, row, group); + const float16x8_t rn_v = vdupq_n_f16(rn); + for (uint32_t k = 0; k < W.group_size; k += 16) { + const uint8_t* packed = cactus_quant_packed_chunk_ptr(W, row, group, k); + uint8x8_t bytes = vld1_u8(packed); + uint8x8_t lo = vand_u8(bytes, vdup_n_u8(0x0F)); + uint8x8_t hi = vshr_n_u8(bytes, 4); + vst1q_f16(dst + k, + vmulq_f16(cactus_tq4_lookup_codebook8(vzip1_u8(lo, hi), cb_bytes), rn_v)); + vst1q_f16(dst + k + 8, + vmulq_f16(cactus_tq4_lookup_codebook8(vzip2_u8(lo, hi), cb_bytes), rn_v)); + } + } +}; + +struct CactusTQ2ScaledDecoder { + uint8x8_t cb_bytes; + + explicit CactusTQ2ScaledDecoder(const CactusQuantMatrix& W) + : cb_bytes(vld1_u8(reinterpret_cast(W.codebook))) {} + + void operator()(const CactusQuantMatrix& W, uint32_t row, uint32_t group, __fp16* dst) const { + const __fp16 rn = *cactus_quant_scale_ptr(W, row, group); + const float16x8_t rn_v = vdupq_n_f16(rn); + for (uint32_t k = 0; k < W.group_size; k += 8) { + const uint8_t* packed = cactus_quant_packed_chunk_ptr(W, row, group, k); + uint8x8_t indices = cactus_tq2_unpack_8x2bit_le(packed[0], packed[1]); + vst1q_f16(dst + k, + vmulq_f16(cactus_tq2_lookup_codebook8(indices, cb_bytes), rn_v)); + } + } +}; + +template +static void cactus_quant_group_gemm( + const CactusQuantMatrix& W, + const __fp16* A, + uint32_t M, + __fp16* C, + DecodeGroup decode_group) { + static_assert(Bits >= 1 && Bits <= 4); + if (W.bits != Bits) return; + + constexpr size_t TILE_N = 16; + const size_t n_blocks = (W.N + TILE_N - 1) / TILE_N; + + + CactusThreading::parallel_gemm_tiles(M, n_blocks, + [&, decode_group](size_t block_start, size_t block_end) { + thread_local std::vector<__fp16> b_tile; + thread_local std::vector c_accum; + if (b_tile.size() < TILE_N * W.group_size) + b_tile.resize(TILE_N * W.group_size); + if (c_accum.size() < M * TILE_N) + c_accum.resize(M * TILE_N); + + for (size_t block = block_start; block < block_end; ++block) { + const size_t n_start = block * TILE_N; + const size_t actual_n = std::min(TILE_N, static_cast(W.N) - n_start); + + std::fill(c_accum.begin(), c_accum.begin() + M * TILE_N, 0.0f); + + for (uint32_t g = 0; g < W.num_groups; ++g) { + for (size_t ni = 0; ni < actual_n; ++ni) { + decode_group(W, static_cast(n_start + ni), g, + b_tile.data() + ni * W.group_size); + } + cactus_quant_matmul_f32_segment_accum( + A + static_cast(g) * W.group_size, + W.K, + b_tile.data(), + c_accum.data(), + M, + W.group_size, + actual_n); + } + + for (size_t m = 0; m < M; ++m) + for (size_t ni = 0; ni < actual_n; ++ni) + C[m * W.N + n_start + ni] = static_cast<__fp16>(c_accum[m * TILE_N + ni]); + } + }); +} + +static bool cactus_quant_valid_common(const CactusQuantMatrix* W, const void* A, void* C) { + if (W == nullptr || A == nullptr || C == nullptr) return false; + if (W->K == 0 || W->N == 0 || W->group_size == 0 || W->num_groups == 0) return false; + if (W->group_size > 256) return false; + if ((W->group_size & (W->group_size - 1)) != 0) return false; + if (W->K != W->group_size * W->num_groups) return false; + if (W->codebook == nullptr || W->norms == nullptr || W->packed_indices == nullptr) return false; + return true; +} + +} // namespace + +static void cactus_quant_dispatch_group_gemm( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C) { + switch (W->bits) { + case 1: cactus_quant_1bit_gemm(W, A, M, C); return; + case 2: cactus_quant_2bit_gemm(W, A, M, C); return; + case 3: cactus_quant_3bit_gemm(W, A, M, C); return; + case 4: cactus_quant_4bit_gemm(W, A, M, C); return; + default: return; + } +} + +static void cactus_matmul_f16_worker( + const __fp16* a, + const __fp16* b_transposed, + __fp16* c, + size_t /*M*/, + size_t K, + size_t N, + size_t start_row, + size_t end_row +) { + constexpr size_t TILE_M = 4; + constexpr size_t TILE_N = 4; + const size_t K16 = (K / 16) * 16; + const size_t K8 = (K / 8) * 8; + + for (size_t row_block = start_row; row_block < end_row; row_block += TILE_M) { + const size_t m_end = std::min(row_block + TILE_M, end_row); + + for (size_t col_block = 0; col_block < N; col_block += TILE_N) { + const size_t n_end = std::min(col_block + TILE_N, N); + + float16x8_t acc[TILE_M][TILE_N]; + for (size_t m = 0; m < TILE_M; ++m) + for (size_t n = 0; n < TILE_N; ++n) + acc[m][n] = vdupq_n_f16(0); + + for (size_t k = 0; k < K16; k += 16) { + float16x8_t a0_lo = (row_block < m_end) ? vld1q_f16(a + row_block * K + k) : vdupq_n_f16(0); + float16x8_t a0_hi = (row_block < m_end) ? vld1q_f16(a + row_block * K + k + 8) : vdupq_n_f16(0); + float16x8_t a1_lo = (row_block + 1 < m_end) ? vld1q_f16(a + (row_block + 1) * K + k) : vdupq_n_f16(0); + float16x8_t a1_hi = (row_block + 1 < m_end) ? vld1q_f16(a + (row_block + 1) * K + k + 8) : vdupq_n_f16(0); + float16x8_t a2_lo = (row_block + 2 < m_end) ? vld1q_f16(a + (row_block + 2) * K + k) : vdupq_n_f16(0); + float16x8_t a2_hi = (row_block + 2 < m_end) ? vld1q_f16(a + (row_block + 2) * K + k + 8) : vdupq_n_f16(0); + float16x8_t a3_lo = (row_block + 3 < m_end) ? vld1q_f16(a + (row_block + 3) * K + k) : vdupq_n_f16(0); + float16x8_t a3_hi = (row_block + 3 < m_end) ? vld1q_f16(a + (row_block + 3) * K + k + 8) : vdupq_n_f16(0); + + for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { + float16x8_t b_lo = vld1q_f16(b_transposed + (col_block + ni) * K + k); + float16x8_t b_hi = vld1q_f16(b_transposed + (col_block + ni) * K + k + 8); + + acc[0][ni] = vfmaq_f16(acc[0][ni], a0_lo, b_lo); + acc[0][ni] = vfmaq_f16(acc[0][ni], a0_hi, b_hi); + acc[1][ni] = vfmaq_f16(acc[1][ni], a1_lo, b_lo); + acc[1][ni] = vfmaq_f16(acc[1][ni], a1_hi, b_hi); + acc[2][ni] = vfmaq_f16(acc[2][ni], a2_lo, b_lo); + acc[2][ni] = vfmaq_f16(acc[2][ni], a2_hi, b_hi); + acc[3][ni] = vfmaq_f16(acc[3][ni], a3_lo, b_lo); + acc[3][ni] = vfmaq_f16(acc[3][ni], a3_hi, b_hi); + } + } + + for (size_t k = K16; k < K8; k += 8) { + float16x8_t a0_v = (row_block < m_end) ? vld1q_f16(a + row_block * K + k) : vdupq_n_f16(0); + float16x8_t a1_v = (row_block + 1 < m_end) ? vld1q_f16(a + (row_block + 1) * K + k) : vdupq_n_f16(0); + float16x8_t a2_v = (row_block + 2 < m_end) ? vld1q_f16(a + (row_block + 2) * K + k) : vdupq_n_f16(0); + float16x8_t a3_v = (row_block + 3 < m_end) ? vld1q_f16(a + (row_block + 3) * K + k) : vdupq_n_f16(0); + + for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { + float16x8_t b_v = vld1q_f16(b_transposed + (col_block + ni) * K + k); + acc[0][ni] = vfmaq_f16(acc[0][ni], a0_v, b_v); + acc[1][ni] = vfmaq_f16(acc[1][ni], a1_v, b_v); + acc[2][ni] = vfmaq_f16(acc[2][ni], a2_v, b_v); + acc[3][ni] = vfmaq_f16(acc[3][ni], a3_v, b_v); + } + } + + for (size_t k = K8; k < K; ++k) { + for (size_t mi = 0; mi < TILE_M && row_block + mi < m_end; ++mi) { + __fp16 av = a[(row_block + mi) * K + k]; + for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { + __fp16 bv = b_transposed[(col_block + ni) * K + k]; + acc[mi][ni] = vsetq_lane_f16(vgetq_lane_f16(acc[mi][ni], 0) + av * bv, acc[mi][ni], 0); + } + } + } + + for (size_t mi = 0; mi < TILE_M && row_block + mi < m_end; ++mi) { + for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { + c[(row_block + mi) * N + col_block + ni] = hsum_f16x8(acc[mi][ni]); + } + } + } + } +} + +void cactus_matmul_f16( + const __fp16* a, + const __fp16* b_transposed, + __fp16* c, + size_t M, + size_t K, + size_t N +) { + +#ifdef __APPLE__ + if (K >= ACCELERATE_K_THRESHOLD && M >= ACCELERATE_M_THRESHOLD) { + const size_t a_len = M * K; + const size_t b_len = N * K; + const size_t c_len = M * N; + + std::vector A_f32(a_len); + std::vector BT_f32(b_len); + std::vector C_f32(c_len); + + for (size_t i = 0; i < a_len; i++) A_f32[i] = (float)a[i]; + for (size_t i = 0; i < b_len; i++) BT_f32[i] = (float)b_transposed[i]; + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, + (int)M, (int)N, (int)K, + 1.0f, A_f32.data(), (int)K, + BT_f32.data(), (int)K, + 0.0f, C_f32.data(), (int)N); + + for (size_t i = 0; i < c_len; i++) { + float v = C_f32[i]; + if (v > 65504.f) v = 65504.f; + else if (v < -65504.f) v = -65504.f; + c[i] = (__fp16)v; + } + return; + } +#endif + + constexpr size_t TILE_M = 4; + const size_t num_row_blocks = (M + TILE_M - 1) / TILE_M; + + CactusThreading::parallel_for(num_row_blocks, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [=](size_t start_block, size_t end_block) { + for (size_t block_idx = start_block; block_idx < end_block; ++block_idx) { + size_t start_row = block_idx * TILE_M; + size_t end_row = std::min(start_row + TILE_M, M); + + cactus_matmul_f16_worker( + a, b_transposed, c, + M, K, N, + start_row, end_row + ); + + } + }); +} + +uint32_t cactus_quant_packed_group_bytes(uint32_t bits, uint32_t group_size) { + if (bits == 0 || bits > 4) return 0; + return (group_size * bits + 7) / 8; +} + +static inline float tq_quantize_group_i8(const __fp16* src, int8_t* dst, uint32_t gs); +static inline float tq_quantize_codebook_i8(const __fp16* codebook, int8_t* cb_i8, uint32_t cb_size); +static inline int8x16_t tq_expand_i8_16(const uint8_t* packed, uint32_t bits, int8x16_t cb_lut); + +template +__attribute__((always_inline)) static inline void tq_expand_i8_32( + const uint8_t* packed, int8x16_t cb_lut, + int8x16_t& out0, int8x16_t& out1) { + if constexpr (Bits == 4) { + uint8x16_t bytes = vld1q_u8(packed); + uint8x16_t lo = vandq_u8(bytes, vdupq_n_u8(0x0F)); + uint8x16_t hi = vshrq_n_u8(bytes, 4); + out0 = vqtbl1q_s8(cb_lut, vzip1q_u8(lo, hi)); + out1 = vqtbl1q_s8(cb_lut, vzip2q_u8(lo, hi)); + } else { + out0 = tq_expand_i8_16(packed, Bits, cb_lut); + out1 = tq_expand_i8_16(packed + (16 * Bits) / 8, Bits, cb_lut); + } +} + +template +__attribute__((always_inline)) static inline void cactus_quant_sdot_gemv_int8( + const CactusQuantMatrix* W, + const __fp16* code_basis, + __fp16* y) { + static_assert(Bits >= 1 && Bits <= 4); + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(Bits, gs); + constexpr uint32_t cb_size = 1u << Bits; + + constexpr size_t INT8_TILE_N = 16; + const size_t int8_n_blocks = (W->N + INT8_TILE_N - 1) / INT8_TILE_N; + + thread_local std::vector act_i8_buf; + thread_local std::vector act_scales_buf; + if (act_i8_buf.size() < W->K) act_i8_buf.resize(W->K); + if (act_scales_buf.size() < W->num_groups) act_scales_buf.resize(W->num_groups); + for (uint32_t g = 0; g < W->num_groups; ++g) { + act_scales_buf[g] = tq_quantize_group_i8( + code_basis + static_cast(g) * gs, + act_i8_buf.data() + static_cast(g) * gs, gs); + } + const int8_t* act_i8 = act_i8_buf.data(); + const float* act_scales = act_scales_buf.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, cb_size); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + auto expand_group4 = [&](size_t cache_block, uint32_t g, int8_t* dst, float* norm_scale) { + const uint32_t n_vecs = gs / 16; + const size_t n_start4 = cache_block * 4; + + const uint8_t* row_ptrs[4]; + bool row_valid[4]; + for (size_t ni = 0; ni < 4; ++ni) { + const size_t n_abs = n_start4 + ni; + row_valid[ni] = n_abs < W->N; + if (row_valid[ni]) { + row_ptrs[ni] = W->packed_indices + (n_abs * W->num_groups + g) * pgb; + norm_scale[ni] = static_cast(W->norms[n_abs * W->num_groups + g]) * cb_scale; + } else { + row_ptrs[ni] = nullptr; + norm_scale[ni] = 0.0f; + } + } + + uint32_t v = 0; + for (; v + 1 < n_vecs; v += 2) { + int8x16_t r0a, r0b, r1a, r1b, r2a, r2b, r3a, r3b; + const size_t off = (v * 16 * Bits) / 8; + if (row_valid[0]) tq_expand_i8_32(row_ptrs[0] + off, cb_lut, r0a, r0b); + else { r0a = vdupq_n_s8(0); r0b = vdupq_n_s8(0); } + if (row_valid[1]) tq_expand_i8_32(row_ptrs[1] + off, cb_lut, r1a, r1b); + else { r1a = vdupq_n_s8(0); r1b = vdupq_n_s8(0); } + if (row_valid[2]) tq_expand_i8_32(row_ptrs[2] + off, cb_lut, r2a, r2b); + else { r2a = vdupq_n_s8(0); r2b = vdupq_n_s8(0); } + if (row_valid[3]) tq_expand_i8_32(row_ptrs[3] + off, cb_lut, r3a, r3b); + else { r3a = vdupq_n_s8(0); r3b = vdupq_n_s8(0); } + tq_interleave_4x_s8(r0a, r1a, r2a, r3a, dst + v * 64); + tq_interleave_4x_s8(r0b, r1b, r2b, r3b, dst + (v + 1) * 64); + } + for (; v < n_vecs; ++v) { + int8x16_t r0, r1, r2, r3; + const size_t off = (v * 16 * Bits) / 8; + r0 = row_valid[0] ? tq_expand_i8_16(row_ptrs[0] + off, Bits, cb_lut) : vdupq_n_s8(0); + r1 = row_valid[1] ? tq_expand_i8_16(row_ptrs[1] + off, Bits, cb_lut) : vdupq_n_s8(0); + r2 = row_valid[2] ? tq_expand_i8_16(row_ptrs[2] + off, Bits, cb_lut) : vdupq_n_s8(0); + r3 = row_valid[3] ? tq_expand_i8_16(row_ptrs[3] + off, Bits, cb_lut) : vdupq_n_s8(0); + tq_interleave_4x_s8(r0, r1, r2, r3, dst + v * 64); + } + }; + + cactus_quant_parallel_ranges(int8_n_blocks, 16, [&](size_t block_start, size_t block_end) { + for (size_t block = block_start; block < block_end; ++block) { + const size_t n_start = block * INT8_TILE_N; + const size_t actual_n = std::min(INT8_TILE_N, static_cast(W->N) - n_start); + float acc[INT8_TILE_N] = {}; + + for (uint32_t g = 0; g < W->num_groups; ++g) { + const int8_t* a_group = act_i8 + static_cast(g) * gs; + const float act_scale = act_scales[g]; + + if (actual_n == INT8_TILE_N) { + int32x4_t dot0 = vdupq_n_s32(0); + int32x4_t dot1 = vdupq_n_s32(0); + int32x4_t dot2 = vdupq_n_s32(0); + int32x4_t dot3 = vdupq_n_s32(0); + + const size_t cache_block0 = n_start / 4; + alignas(16) int8_t b_stack0[256 * 4]; + alignas(16) int8_t b_stack1[256 * 4]; + alignas(16) int8_t b_stack2[256 * 4]; + alignas(16) int8_t b_stack3[256 * 4]; + float norm_stack0[4], norm_stack1[4], norm_stack2[4], norm_stack3[4]; + + expand_group4(cache_block0 + 0, g, b_stack0, norm_stack0); + expand_group4(cache_block0 + 1, g, b_stack1, norm_stack1); + expand_group4(cache_block0 + 2, g, b_stack2, norm_stack2); + expand_group4(cache_block0 + 3, g, b_stack3, norm_stack3); + + // Dual-accumulator panels break the 8-deep SDOT dependency chain + // into two 4-deep chains running in parallel. Critical for SDOT + // throughput on Apple Silicon (4 sdot/cycle, ~3 cycle latency). + // Tested quad-accum: tied with dual (OoO already extracts ILP). + int32x4_t dot0b = vdupq_n_s32(0); + int32x4_t dot1b = vdupq_n_s32(0); + int32x4_t dot2b = vdupq_n_s32(0); + int32x4_t dot3b = vdupq_n_s32(0); + + for (uint32_t k = 0; k < gs; k += 32) { + int8x16_t a_lo = vld1q_s8(a_group + k); + int8x16_t a_hi = vld1q_s8(a_group + k + 16); + + #define TQ_SDOT_PANEL_T(DOT_A, DOT_B, BASE) do { \ + const int8_t* bk = (BASE) + k * 4; \ + DOT_A = CACTUS_DOTQ_LANE(DOT_A, vld1q_s8(bk), a_lo, 0); \ + DOT_B = CACTUS_DOTQ_LANE(DOT_B, vld1q_s8(bk + 16), a_lo, 1); \ + DOT_A = CACTUS_DOTQ_LANE(DOT_A, vld1q_s8(bk + 32), a_lo, 2); \ + DOT_B = CACTUS_DOTQ_LANE(DOT_B, vld1q_s8(bk + 48), a_lo, 3); \ + DOT_A = CACTUS_DOTQ_LANE(DOT_A, vld1q_s8(bk + 64), a_hi, 0); \ + DOT_B = CACTUS_DOTQ_LANE(DOT_B, vld1q_s8(bk + 80), a_hi, 1); \ + DOT_A = CACTUS_DOTQ_LANE(DOT_A, vld1q_s8(bk + 96), a_hi, 2); \ + DOT_B = CACTUS_DOTQ_LANE(DOT_B, vld1q_s8(bk + 112), a_hi, 3); \ + } while (0) + + TQ_SDOT_PANEL_T(dot0, dot0b, b_stack0); + TQ_SDOT_PANEL_T(dot1, dot1b, b_stack1); + TQ_SDOT_PANEL_T(dot2, dot2b, b_stack2); + TQ_SDOT_PANEL_T(dot3, dot3b, b_stack3); + #undef TQ_SDOT_PANEL_T + } + + dot0 = vaddq_s32(dot0, dot0b); + dot1 = vaddq_s32(dot1, dot1b); + dot2 = vaddq_s32(dot2, dot2b); + dot3 = vaddq_s32(dot3, dot3b); + + const float32x4_t scale0 = vmulq_n_f32(vld1q_f32(norm_stack0), act_scale); + const float32x4_t scale1 = vmulq_n_f32(vld1q_f32(norm_stack1), act_scale); + const float32x4_t scale2 = vmulq_n_f32(vld1q_f32(norm_stack2), act_scale); + const float32x4_t scale3 = vmulq_n_f32(vld1q_f32(norm_stack3), act_scale); + + vst1q_f32(acc + 0, vfmaq_f32(vld1q_f32(acc + 0), vcvtq_f32_s32(dot0), scale0)); + vst1q_f32(acc + 4, vfmaq_f32(vld1q_f32(acc + 4), vcvtq_f32_s32(dot1), scale1)); + vst1q_f32(acc + 8, vfmaq_f32(vld1q_f32(acc + 8), vcvtq_f32_s32(dot2), scale2)); + vst1q_f32(acc + 12, vfmaq_f32(vld1q_f32(acc + 12), vcvtq_f32_s32(dot3), scale3)); + continue; + } + + for (size_t ni4 = 0; ni4 < actual_n; ni4 += 4) { + const size_t n_abs = n_start + ni4; + const size_t cache_block = n_abs / 4; + alignas(16) int8_t b_stack[256 * 4]; + float norm_stack[4]; + expand_group4(cache_block, g, b_stack, norm_stack); + const float32x4_t scale_v = vmulq_n_f32(vld1q_f32(norm_stack), act_scale); + + int32x4_t dot = vdupq_n_s32(0); + for (uint32_t k = 0; k < gs; k += 32) { + const int8_t* bk = b_stack + k * 4; + int8x16_t a_lo = vld1q_s8(a_group + k); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk), a_lo, 0); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 16), a_lo, 1); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 32), a_lo, 2); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 48), a_lo, 3); + int8x16_t a_hi = vld1q_s8(a_group + k + 16); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 64), a_hi, 0); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 80), a_hi, 1); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 96), a_hi, 2); + dot = CACTUS_DOTQ_LANE(dot, vld1q_s8(bk + 112), a_hi, 3); + } + + float32x4_t contrib = vmulq_f32(vcvtq_f32_s32(dot), scale_v); + float tmp[4]; + vst1q_f32(tmp, contrib); + const size_t lane_count = std::min(4, actual_n - ni4); + for (size_t lane = 0; lane < lane_count; ++lane) { + acc[ni4 + lane] += tmp[lane]; + } + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + y[n_start + ni] = static_cast<__fp16>(acc[ni]); + } + } + }); +} + +void cactus_quant_4bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 4 || (W->group_size % 16) != 0) return; + + constexpr size_t TILE_N = 12; + const size_t n_blocks = (W->N + TILE_N - 1) / TILE_N; + const uint32_t gs = W->group_size; + + uint8x16x2_t cb_bytes; + cb_bytes.val[0] = vld1q_u8(reinterpret_cast(W->codebook)); + cb_bytes.val[1] = vld1q_u8(reinterpret_cast(W->codebook) + 16); + + const uint32_t pgb = cactus_quant_packed_group_bytes(4, gs); + thread_local std::vector<__fp16> code_basis_buf; + if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); + cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); + const __fp16* code_basis = code_basis_buf.data(); + + if ((gs % 32) == 0 && gs <= 256) { + cactus_quant_sdot_gemv_int8<4>(W, code_basis, y); + return; + } + + cactus_quant_parallel_ranges(n_blocks, 16, [&](size_t block_start, size_t block_end) { + for (size_t block = block_start; block < block_end; ++block) { + const size_t n_start = block * TILE_N; + const size_t actual_n = std::min(TILE_N, static_cast(W->N) - n_start); + float acc[TILE_N] = {}; + + for (uint32_t g = 0; g < W->num_groups; ++g) { + const __fp16* z = code_basis + static_cast(g) * gs; + + float16x8_t acc0[TILE_N]; + float16x8_t acc1[TILE_N]; + for (size_t ni = 0; ni < TILE_N; ++ni) { + acc0[ni] = vdupq_n_f16(0); + acc1[ni] = vdupq_n_f16(0); + } + + for (uint32_t k = 0; k < gs; k += 16) { + float16x8_t z0 = vld1q_f16(z + k); + float16x8_t z1 = vld1q_f16(z + k + 8); + for (size_t ni = 0; ni < actual_n; ++ni) { + const uint8_t* p = W->packed_indices + + (static_cast(n_start + ni) * W->num_groups + g) * pgb + + k / 2; + uint8x8_t bytes = vld1_u8(p); + uint8x8_t lo = vand_u8(bytes, vdup_n_u8(0x0F)); + uint8x8_t hi = vshr_n_u8(bytes, 4); + + float16x8_t cv0 = cactus_tq4_lookup_codebook8(vzip1_u8(lo, hi), cb_bytes); + float16x8_t cv1 = cactus_tq4_lookup_codebook8(vzip2_u8(lo, hi), cb_bytes); + acc0[ni] = vfmaq_f16(acc0[ni], z0, cv0); + acc1[ni] = vfmaq_f16(acc1[ni], z1, cv1); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + float rn = static_cast(W->norms[(n_start + ni) * W->num_groups + g]); + acc[ni] += rn * + (static_cast(hsum_f16x8(acc0[ni])) + + static_cast(hsum_f16x8(acc1[ni]))); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + y[n_start + ni] = static_cast<__fp16>(acc[ni]); + } + } + }); +} + +void cactus_quant_2bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 2 || (W->group_size % 8) != 0) return; + + constexpr size_t TILE_N = 16; + const size_t n_blocks = (W->N + TILE_N - 1) / TILE_N; + const uint32_t gs = W->group_size; + + uint8x8_t cb_bytes = vld1_u8(reinterpret_cast(W->codebook)); + + const uint32_t pgb = cactus_quant_packed_group_bytes(2, gs); + thread_local std::vector<__fp16> code_basis_buf; + if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); + cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); + const __fp16* code_basis = code_basis_buf.data(); + + if ((gs % 32) == 0 && gs <= 256) { + cactus_quant_sdot_gemv_int8<2>(W, code_basis, y); + return; + } + + cactus_quant_parallel_ranges(n_blocks, 16, [&](size_t block_start, size_t block_end) { + for (size_t block = block_start; block < block_end; ++block) { + const size_t n_start = block * TILE_N; + const size_t actual_n = std::min(TILE_N, static_cast(W->N) - n_start); + float acc[TILE_N] = {}; + + for (uint32_t g = 0; g < W->num_groups; ++g) { + const __fp16* z = code_basis + static_cast(g) * gs; + + float16x8_t accv[TILE_N]; + for (size_t ni = 0; ni < TILE_N; ++ni) { + accv[ni] = vdupq_n_f16(0); + } + + for (uint32_t k = 0; k < gs; k += 8) { + float16x8_t z_v = vld1q_f16(z + k); + for (size_t ni = 0; ni < actual_n; ++ni) { + const uint8_t* p = W->packed_indices + + (static_cast(n_start + ni) * W->num_groups + g) * pgb + + k / 4; + uint8x8_t indices = cactus_tq2_unpack_8x2bit_le(p[0], p[1]); + float16x8_t cv = cactus_tq2_lookup_codebook8(indices, cb_bytes); + accv[ni] = vfmaq_f16(accv[ni], z_v, cv); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + float rn = static_cast(W->norms[(n_start + ni) * W->num_groups + g]); + acc[ni] += rn * static_cast(hsum_f16x8(accv[ni])); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + y[n_start + ni] = static_cast<__fp16>(acc[ni]); + } + } + }); +} + +template +static void cactus_quant_Nbit_gemm_impl( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C, + uint32_t gs_align) { + if (!cactus_quant_valid_common(W, A, C) || M == 0) return; + if (W->bits != Bits || (W->group_size % gs_align) != 0) return; + thread_local std::vector<__fp16> code_basis_buf; + if (code_basis_buf.size() < static_cast(M) * W->K) { + code_basis_buf.resize(static_cast(M) * W->K); + } + cactus_quant_transform_hadamard_activations(*W, A, M, code_basis_buf.data()); + cactus_quant_group_gemm(*W, code_basis_buf.data(), M, C, Decoder(*W)); +} + +void cactus_quant_4bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C) { + cactus_quant_Nbit_gemm_impl<4, CactusTQ4ScaledDecoder>(W, A, M, C, 16); +} + +void cactus_quant_2bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C) { + cactus_quant_Nbit_gemm_impl<2, CactusTQ2ScaledDecoder>(W, A, M, C, 8); +} + + + +struct CactusTQ1ScaledDecoder { + uint8x8_t cb_bytes; + + explicit CactusTQ1ScaledDecoder(const CactusQuantMatrix& W) + : cb_bytes(vld1_u8(reinterpret_cast(W.codebook))) {} + + void operator()(const CactusQuantMatrix& W, uint32_t row, uint32_t group, __fp16* dst) const { + const __fp16 rn = *cactus_quant_scale_ptr(W, row, group); + const float16x8_t rn_v = vdupq_n_f16(rn); + for (uint32_t k = 0; k < W.group_size; k += 8) { + const uint8_t* packed = cactus_quant_packed_chunk_ptr(W, row, group, k); + + uint8_t byte = packed[0]; + uint64_t idx_word = + ((uint64_t)((byte >> 0) & 1) ) | + ((uint64_t)((byte >> 1) & 1) << 8) | + ((uint64_t)((byte >> 2) & 1) << 16) | + ((uint64_t)((byte >> 3) & 1) << 24) | + ((uint64_t)((byte >> 4) & 1) << 32) | + ((uint64_t)((byte >> 5) & 1) << 40) | + ((uint64_t)((byte >> 6) & 1) << 48) | + ((uint64_t)((byte >> 7) & 1) << 56); + uint8x8_t indices = vcreate_u8(idx_word); + + uint8x8_t off_lo = vshl_n_u8(indices, 1); + uint8x8_t off_hi = vadd_u8(off_lo, vdup_n_u8(1)); + uint8x8x2_t zipped = vzip_u8(off_lo, off_hi); + uint8x16_t byte_idx = vcombine_u8(zipped.val[0], zipped.val[1]); + uint8x16_t lut = vcombine_u8(cb_bytes, cb_bytes); + float16x8_t cv = vreinterpretq_f16_u8(vqtbl1q_u8(lut, byte_idx)); + vst1q_f16(dst + k, vmulq_f16(cv, rn_v)); + } + } +}; + +void cactus_quant_1bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 1 || (W->group_size % 8) != 0) return; + + constexpr size_t TILE_N = 12; + const size_t n_blocks = (W->N + TILE_N - 1) / TILE_N; + const uint32_t gs = W->group_size; + + if ((gs % 32) == 0 && gs <= 256) { + thread_local std::vector<__fp16> code_basis_buf; + if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); + cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); + cactus_quant_sdot_gemv_int8<1>(W, code_basis_buf.data(), y); + return; + } + + uint8x8_t cb_bytes = vld1_u8(reinterpret_cast(W->codebook)); + uint8x16_t cb_lut = vcombine_u8(cb_bytes, cb_bytes); + + const uint32_t pgb = cactus_quant_packed_group_bytes(1, gs); + + cactus_quant_parallel_ranges(n_blocks, 16, [&](size_t block_start, size_t block_end) { + __fp16 z_buf[256]; + + for (size_t block = block_start; block < block_end; ++block) { + const size_t n_start = block * TILE_N; + const size_t actual_n = std::min(TILE_N, static_cast(W->N) - n_start); + float acc[TILE_N] = {}; + + for (uint32_t g = 0; g < W->num_groups; ++g) { + cactus_quant_transform_hadamard_group(*W, x + g * gs, g, z_buf); + const __fp16* z = z_buf; + + float16x8_t accv[TILE_N]; + for (size_t ni = 0; ni < TILE_N; ++ni) accv[ni] = vdupq_n_f16(0); + + for (uint32_t k = 0; k < gs; k += 8) { + float16x8_t z_v = vld1q_f16(z + k); + for (size_t ni = 0; ni < actual_n; ++ni) { + const uint8_t* p = W->packed_indices + + (static_cast(n_start + ni) * W->num_groups + g) * pgb + + k / 8; + uint8_t byte = p[0]; + uint64_t idx_word = + ((uint64_t)((byte >> 0) & 1) ) | + ((uint64_t)((byte >> 1) & 1) << 8) | + ((uint64_t)((byte >> 2) & 1) << 16) | + ((uint64_t)((byte >> 3) & 1) << 24) | + ((uint64_t)((byte >> 4) & 1) << 32) | + ((uint64_t)((byte >> 5) & 1) << 40) | + ((uint64_t)((byte >> 6) & 1) << 48) | + ((uint64_t)((byte >> 7) & 1) << 56); + uint8x8_t indices = vcreate_u8(idx_word); + uint8x8_t off_lo = vshl_n_u8(indices, 1); + uint8x8_t off_hi = vadd_u8(off_lo, vdup_n_u8(1)); + uint8x8x2_t zipped = vzip_u8(off_lo, off_hi); + uint8x16_t byte_idx = vcombine_u8(zipped.val[0], zipped.val[1]); + float16x8_t cv = vreinterpretq_f16_u8(vqtbl1q_u8(cb_lut, byte_idx)); + accv[ni] = vfmaq_f16(accv[ni], z_v, cv); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + float rn = static_cast(W->norms[(n_start + ni) * W->num_groups + g]); + acc[ni] += rn * static_cast(hsum_f16x8(accv[ni])); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + y[n_start + ni] = static_cast<__fp16>(acc[ni]); + } + } + }); +} + +void cactus_quant_1bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C) { + cactus_quant_Nbit_gemm_impl<1, CactusTQ1ScaledDecoder>(W, A, M, C, 8); +} + + + +static inline uint8x8_t cactus_tq3_unpack_8x3bit(const uint8_t* packed) { + + uint32_t b0 = packed[0], b1 = packed[1], b2 = packed[2]; + uint32_t word = b0 | (b1 << 8) | (b2 << 16); + uint64_t idx_word = + ((uint64_t)((word >> 0) & 0x7) ) | + ((uint64_t)((word >> 3) & 0x7) << 8) | + ((uint64_t)((word >> 6) & 0x7) << 16) | + ((uint64_t)((word >> 9) & 0x7) << 24) | + ((uint64_t)((word >> 12) & 0x7) << 32) | + ((uint64_t)((word >> 15) & 0x7) << 40) | + ((uint64_t)((word >> 18) & 0x7) << 48) | + ((uint64_t)((word >> 21) & 0x7) << 56); + return vcreate_u8(idx_word); +} + +static inline float16x8_t cactus_tq3_lookup_codebook8(uint8x8_t indices, uint8x16_t cb_bytes) { + uint8x8_t off_lo = vshl_n_u8(indices, 1); + uint8x8_t off_hi = vadd_u8(off_lo, vdup_n_u8(1)); + uint8x8x2_t zipped = vzip_u8(off_lo, off_hi); + uint8x16_t byte_idx = vcombine_u8(zipped.val[0], zipped.val[1]); + return vreinterpretq_f16_u8(vqtbl1q_u8(cb_bytes, byte_idx)); +} + +struct CactusTQ3ScaledDecoder { + uint8x16_t cb_bytes; + + explicit CactusTQ3ScaledDecoder(const CactusQuantMatrix& W) + : cb_bytes(vld1q_u8(reinterpret_cast(W.codebook))) {} + + void operator()(const CactusQuantMatrix& W, uint32_t row, uint32_t group, __fp16* dst) const { + const __fp16 rn = *cactus_quant_scale_ptr(W, row, group); + const float16x8_t rn_v = vdupq_n_f16(rn); + for (uint32_t k = 0; k < W.group_size; k += 8) { + const uint8_t* packed = cactus_quant_packed_chunk_ptr(W, row, group, k); + uint8x8_t indices = cactus_tq3_unpack_8x3bit(packed); + float16x8_t cv = cactus_tq3_lookup_codebook8(indices, cb_bytes); + vst1q_f16(dst + k, vmulq_f16(cv, rn_v)); + } + } +}; + +void cactus_quant_3bit_gemv( + const CactusQuantMatrix* W, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 3 || (W->group_size % 8) != 0) return; + + constexpr size_t TILE_N = 12; + const size_t n_blocks = (W->N + TILE_N - 1) / TILE_N; + const uint32_t gs = W->group_size; + + if ((gs % 32) == 0 && gs <= 256) { + thread_local std::vector<__fp16> code_basis_buf; + if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); + cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); + cactus_quant_sdot_gemv_int8<3>(W, code_basis_buf.data(), y); + return; + } + + uint8x16_t cb_bytes = vld1q_u8(reinterpret_cast(W->codebook)); + + const uint32_t pgb = cactus_quant_packed_group_bytes(3, gs); + + cactus_quant_parallel_ranges(n_blocks, 16, [&](size_t block_start, size_t block_end) { + __fp16 z_buf[256]; + + for (size_t block = block_start; block < block_end; ++block) { + const size_t n_start = block * TILE_N; + const size_t actual_n = std::min(TILE_N, static_cast(W->N) - n_start); + float acc[TILE_N] = {}; + + for (uint32_t g = 0; g < W->num_groups; ++g) { + cactus_quant_transform_hadamard_group(*W, x + g * gs, g, z_buf); + const __fp16* z = z_buf; + + float16x8_t accv[TILE_N]; + for (size_t ni = 0; ni < TILE_N; ++ni) accv[ni] = vdupq_n_f16(0); + + for (uint32_t k = 0; k < gs; k += 8) { + float16x8_t z_v = vld1q_f16(z + k); + for (size_t ni = 0; ni < actual_n; ++ni) { + const uint8_t* p = W->packed_indices + + (static_cast(n_start + ni) * W->num_groups + g) * pgb + + (k * 3) / 8; + uint8x8_t indices = cactus_tq3_unpack_8x3bit(p); + float16x8_t cv = cactus_tq3_lookup_codebook8(indices, cb_bytes); + accv[ni] = vfmaq_f16(accv[ni], z_v, cv); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + float rn = static_cast(W->norms[(n_start + ni) * W->num_groups + g]); + acc[ni] += rn * static_cast(hsum_f16x8(accv[ni])); + } + } + + for (size_t ni = 0; ni < actual_n; ++ni) { + y[n_start + ni] = static_cast<__fp16>(acc[ni]); + } + } + }); +} + +void cactus_quant_3bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C) { + cactus_quant_Nbit_gemm_impl<3, CactusTQ3ScaledDecoder>(W, A, M, C, 8); +} + +static inline float tq_quantize_codebook_i8(const __fp16* codebook, int8_t* cb_i8, uint32_t cb_size) { + float max_abs = 0.f; + for (uint32_t i = 0; i < cb_size; i++) { + float v = std::abs(static_cast(codebook[i])); + if (v > max_abs) max_abs = v; + } + float scale = max_abs / 127.f; + if (scale < 1e-10f) scale = 1e-10f; + float inv = 1.f / scale; + for (uint32_t i = 0; i < cb_size; i++) + cb_i8[i] = static_cast(std::round(static_cast(codebook[i]) * inv)); + return scale; +} + +static inline float tq_quantize_group_i8(const __fp16* src, int8_t* dst, uint32_t gs) { + float32x4_t max_vec = vdupq_n_f32(0.f); + uint32_t k = 0; + for (; k + 8 <= gs; k += 8) { + float16x8_t v = vld1q_f16(src + k); + max_vec = vmaxq_f32(max_vec, vabsq_f32(vcvt_f32_f16(vget_low_f16(v)))); + max_vec = vmaxq_f32(max_vec, vabsq_f32(vcvt_f32_f16(vget_high_f16(v)))); + } + float max_abs = vmaxvq_f32(max_vec); + for (; k < gs; k++) { + float v = std::abs(static_cast(src[k])); + if (v > max_abs) max_abs = v; + } + float scale = max_abs / 127.f; + if (scale < 1e-10f) scale = 1e-10f; + float32x4_t inv_vec = vdupq_n_f32(1.f / scale); + k = 0; + for (; k + 8 <= gs; k += 8) { + float16x8_t v = vld1q_f16(src + k); + int32x4_t i0 = vcvtnq_s32_f32(vmulq_f32(vcvt_f32_f16(vget_low_f16(v)), inv_vec)); + int32x4_t i1 = vcvtnq_s32_f32(vmulq_f32(vcvt_f32_f16(vget_high_f16(v)), inv_vec)); + vst1_s8(dst + k, vqmovn_s16(vcombine_s16(vqmovn_s32(i0), vqmovn_s32(i1)))); + } + for (; k < gs; k++) + dst[k] = static_cast(std::round(static_cast(src[k]) / scale)); + return scale; +} + + +static inline uint32_t tq_extract_idx(const uint8_t* packed, uint32_t k, uint32_t bits) { + if (bits == 4) return (packed[k / 2] >> ((k & 1) * 4)) & 0xF; + if (bits == 2) return (packed[k / 4] >> ((k & 3) * 2)) & 0x3; + if (bits == 1) return (packed[k / 8] >> (k & 7)) & 0x1; + const uint32_t bit_pos = k * 3; + const uint32_t byte_idx = bit_pos / 8; + const uint32_t bit_idx = bit_pos & 7; + uint32_t val = static_cast(packed[byte_idx]) >> bit_idx; + if (bit_idx > 5) { + val |= static_cast(packed[byte_idx + 1]) << (8 - bit_idx); + } + return val & 0x7; +} + +static inline int8x16_t tq_expand_i8_16(const uint8_t* packed, uint32_t bits, int8x16_t cb_lut) { + if (bits == 4) { + uint8x8_t bytes = vld1_u8(packed); + uint8x8_t lo = vand_u8(bytes, vdup_n_u8(0x0F)); + uint8x8_t hi = vshr_n_u8(bytes, 4); + uint8x16_t idx = vcombine_u8(vzip1_u8(lo, hi), vzip2_u8(lo, hi)); + return vqtbl1q_s8(cb_lut, idx); + } else if (bits == 2) { + static const uint8_t shuf[16] = {0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3}; + static const int8_t shifts[16] = {0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6}; + uint32_t r; + std::memcpy(&r, packed, 4); + uint8x16_t s = vcombine_u8(vcreate_u8(static_cast(r)), vdup_n_u8(0)); + uint8x16_t bytes = vqtbl1q_u8(s, vld1q_u8(shuf)); + uint8x16_t idx = vandq_u8(vshlq_u8(bytes, vld1q_s8(shifts)), vdupq_n_u8(3)); + return vqtbl1q_s8(cb_lut, idx); + } else if (bits == 1) { + static const uint8_t shuf[16] = {0,0,0,0,0,0,0,0, 1,1,1,1,1,1,1,1}; + static const int8_t shifts[16] = {0,-1,-2,-3,-4,-5,-6,-7, 0,-1,-2,-3,-4,-5,-6,-7}; + uint8x16_t s = vcombine_u8(vld1_u8(packed), vdup_n_u8(0)); + uint8x16_t bytes = vqtbl1q_u8(s, vld1q_u8(shuf)); + uint8x16_t idx = vandq_u8(vshlq_u8(bytes, vld1q_s8(shifts)), vdupq_n_u8(1)); + return vqtbl1q_s8(cb_lut, idx); + } else { // bits == 3 + static const uint8_t shuf_lo[16] = {0,0,0,1,1,1,2,2,3,3,3,4,4,4,5,5}; + static const int8_t sh_lo[16] = {0,-3,-6,-1,-4,-7,-2,-5,0,-3,-6,-1,-4,-7,-2,-5}; + static const uint8_t shuf_hi[16] = {0xFF,0xFF,1,0xFF,0xFF,2,0xFF,0xFF,0xFF,0xFF,4,0xFF,0xFF,5,0xFF,0xFF}; + static const int8_t sh_hi[16] = {0,0,2,0,0,1,0,0,0,0,2,0,0,1,0,0}; + uint8x8_t src = vld1_u8(packed); + uint8x16_t s = vcombine_u8(src, src); + uint8x16_t lo = vshlq_u8(vqtbl1q_u8(s, vld1q_u8(shuf_lo)), vld1q_s8(sh_lo)); + uint8x16_t hi = vshlq_u8(vqtbl1q_u8(s, vld1q_u8(shuf_hi)), vld1q_s8(sh_hi)); + uint8x16_t idx = vandq_u8(vorrq_u8(lo, hi), vdupq_n_u8(7)); + return vqtbl1q_s8(cb_lut, idx); + } +} + +static void tq_preexpand_weights_interleaved( + const CactusQuantMatrix* W, + uint32_t bits, uint32_t num_groups, uint32_t gs, uint32_t pgb, + const int8x16_t& cb_lut, float cb_scale, + size_t N_blocks, + int8_t* w_il, float* n_f32) { + const uint32_t n_vecs = gs / 16; + const size_t panel_bytes = static_cast(4) * pgb; + + alignas(16) static const uint8_t reorder4_tbl[16] = {0,1,2,3, 8,9,10,11, 4,5,6,7, 12,13,14,15}; + const uint8x16_t reorder4 = vld1q_u8(reorder4_tbl); + alignas(16) static const uint8_t spread2_tbl[16] = {0,0,0,0, 4,4,4,4, 8,8,8,8, 12,12,12,12}; + alignas(16) static const int8_t shifts2_tbl[16] = {0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6}; + const uint8x16_t spread2 = vld1q_u8(spread2_tbl); + const int8x16_t shifts2 = vld1q_s8(shifts2_tbl); + const uint8x16_t mask2 = vdupq_n_u8(0x03); + + for (size_t nb = 0; nb < N_blocks; ++nb) { + size_t n_start = nb * 4; + size_t valid_n = std::min(size_t(4), static_cast(W->N) - n_start); + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* panel = W->packed_indices + (nb * num_groups + g) * panel_bytes; + int8x16_t exp4[4][16]; + + if (bits == 4) { + for (uint32_t v = 0; v < n_vecs; ++v) { + const uint8_t* c0 = panel + (2 * v + 0) * 16; + const uint8_t* c1 = panel + (2 * v + 1) * 16; + uint32_t b0[4], b1[4]; + std::memcpy(b0, c0, 16); + std::memcpy(b1, c1, 16); + for (size_t r = 0; r < 4; ++r) { + uint8x8_t bytes_r = vcreate_u8( + static_cast(b0[r]) | + (static_cast(b1[r]) << 32)); + uint8x8_t lo = vand_u8(bytes_r, vdup_n_u8(0x0F)); + uint8x8_t hi = vshr_n_u8(bytes_r, 4); + uint8x16_t combined = vcombine_u8(lo, hi); + uint8x16_t reordered = vqtbl1q_u8(combined, reorder4); + exp4[r][v] = vqtbl1q_s8(cb_lut, reordered); + } + } + } else if (bits == 2) { + const uint32_t chunks = gs / 16; + for (uint32_t c = 0; c < chunks; ++c) { + uint8x16_t bytes = vld1q_u8(panel + c * 16); + for (size_t r = 0; r < 4; ++r) { + uint8x16_t pick_r = vaddq_u8(spread2, vdupq_n_u8(static_cast(r))); + uint8x16_t row_bytes = vqtbl1q_u8(bytes, pick_r); + uint8x16_t shifted = vshlq_u8(row_bytes, shifts2); + uint8x16_t idx = vandq_u8(shifted, mask2); + exp4[r][c] = vqtbl1q_s8(cb_lut, idx); + } + } + } else if (bits == 1) { + const uint32_t chunks = gs / 32; + alignas(16) uint8_t row_idx[4][256]; + for (uint32_t c = 0; c < chunks; ++c) { + const uint8_t* ch = panel + c * 16; + for (size_t r = 0; r < 4; ++r) { + for (uint32_t p = 0; p < 4; ++p) { + uint8_t byte = ch[p * 4 + r]; + uint8_t* dst_p = row_idx[r] + c * 32 + p * 8; + dst_p[0] = (byte >> 0) & 0x1; + dst_p[1] = (byte >> 1) & 0x1; + dst_p[2] = (byte >> 2) & 0x1; + dst_p[3] = (byte >> 3) & 0x1; + dst_p[4] = (byte >> 4) & 0x1; + dst_p[5] = (byte >> 5) & 0x1; + dst_p[6] = (byte >> 6) & 0x1; + dst_p[7] = (byte >> 7) & 0x1; + } + } + } + for (size_t r = 0; r < 4; ++r) + for (uint32_t v = 0; v < n_vecs; ++v) + exp4[r][v] = vqtbl1q_s8(cb_lut, vld1q_u8(row_idx[r] + v * 16)); + } else if (bits == 3) { + const uint32_t chunks = gs / 4; + alignas(16) uint8_t row_idx[4][256]; + for (uint32_t c = 0; c < chunks; ++c) { + const uint8_t* ch = panel + c * 6; + uint64_t word = 0; + std::memcpy(&word, ch, 6); + for (size_t r = 0; r < 4; ++r) { + uint32_t bit_pos = static_cast(r) * 12; + row_idx[r][c * 4 + 0] = static_cast((word >> (bit_pos + 0)) & 0x7); + row_idx[r][c * 4 + 1] = static_cast((word >> (bit_pos + 3)) & 0x7); + row_idx[r][c * 4 + 2] = static_cast((word >> (bit_pos + 6)) & 0x7); + row_idx[r][c * 4 + 3] = static_cast((word >> (bit_pos + 9)) & 0x7); + } + } + for (size_t r = 0; r < 4; ++r) + for (uint32_t v = 0; v < n_vecs; ++v) + exp4[r][v] = vqtbl1q_s8(cb_lut, vld1q_u8(row_idx[r] + v * 16)); + } + + for (size_t ni = valid_n; ni < 4; ++ni) + for (uint32_t v = 0; v < n_vecs; ++v) exp4[ni][v] = vdupq_n_s8(0); + + int8_t* dst = w_il + (nb * num_groups + g) * gs * 4; + for (uint32_t v = 0; v < n_vecs; ++v) + tq_interleave_4x_s8(exp4[0][v], exp4[1][v], exp4[2][v], exp4[3][v], dst + v * 64); + + float* nd = n_f32 + (nb * num_groups + g) * 4; + for (size_t ni = 0; ni < 4; ++ni) + nd[ni] = (n_start + ni < W->N) + ? static_cast(W->norms[(nb * num_groups + g) * 4 + ni]) * cb_scale + : 0.f; + } + } +} + +static void tq_preexpand_weights( + const CactusQuantMatrix* W, + uint32_t bits, uint32_t num_groups, uint32_t gs, uint32_t pgb, + const int8x16_t& cb_lut, float cb_scale, + size_t N_blocks, + int8_t* w_il, float* n_f32) { + if (W != nullptr + && (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0 + && bits >= 1 && bits <= 4 + && (W->N % 4) == 0 + && (gs % 32) == 0 + && gs <= 256) { + tq_preexpand_weights_interleaved(W, bits, num_groups, gs, pgb, cb_lut, cb_scale, N_blocks, w_il, n_f32); + return; + } + const uint32_t n_vecs = gs / 16; + for (size_t nb = 0; nb < N_blocks; ++nb) { + size_t n_start = nb * 4; + size_t valid_n = std::min(size_t(4), static_cast(W->N) - n_start); + for (uint32_t g = 0; g < num_groups; ++g) { + int8x16_t exp4[4][16]; + for (size_t ni = 0; ni < valid_n; ++ni) { + const uint8_t* p = W->packed_indices + (static_cast(n_start+ni)*num_groups+g)*pgb; + for (uint32_t v = 0; v < n_vecs; ++v) + exp4[ni][v] = tq_expand_i8_16(p + (v*16*bits)/8, bits, cb_lut); + } + for (size_t ni = valid_n; ni < 4; ++ni) + for (uint32_t v = 0; v < n_vecs; ++v) exp4[ni][v] = vdupq_n_s8(0); + + int8_t* dst = w_il + (nb*num_groups+g)*gs*4; + for (uint32_t v = 0; v < n_vecs; ++v) + tq_interleave_4x_s8(exp4[0][v], exp4[1][v], exp4[2][v], exp4[3][v], dst + v*64); + float* nd = n_f32 + (nb*num_groups+g)*4; + for (size_t ni = 0; ni < 4; ++ni) + nd[ni] = (n_start+ni < W->N) ? static_cast(W->norms[(n_start+ni)*num_groups+g]) * cb_scale : 0.f; + } + } +} + +static void cactus_quant_4bit_gemm_interleaved( + const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); + +void cactus_quant_matmul( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C) { + if (W != nullptr && (W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL) != 0) { + cactus_quant_orthogonal_matmul(W, A, M, C); + return; + } + if (W != nullptr && (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0) { + if (M == 1) { + if (W->bits == 4) { + cactus_quant_4bit_gemv_interleaved(W, W->packed_indices, W->norms, A, C); + return; + } + if (W->bits == 3) { + cactus_quant_3bit_gemv_interleaved(W, W->packed_indices, W->norms, A, C); + return; + } + if (W->bits == 2) { + cactus_quant_2bit_gemv_interleaved(W, W->packed_indices, W->norms, A, C); + return; + } + if (W->bits == 1) { + cactus_quant_1bit_gemv_interleaved(W, W->packed_indices, W->norms, A, C); + return; + } + } else if (W->bits == 4 && (W->group_size % 32) == 0 && W->group_size <= 256 + && (W->N % 4) == 0 && cactus_quant_valid_common(W, A, C)) { + cactus_quant_4bit_gemm_interleaved(W, A, M, C); + return; + } + } + if (!cactus_quant_valid_common(W, A, C) || M == 0) return; + + const uint32_t gs = W->group_size; + const uint32_t bits = W->bits; + const uint32_t num_groups = W->num_groups; + const uint32_t pgb = cactus_quant_packed_group_bytes(bits, gs); + const bool has_expanded_layout = W->expanded != nullptr && W->norm_f32 != nullptr; + const bool has_sdot_group_layout = (gs % 32) == 0; + const size_t N_blocks = (W->N + 3) / 4; + const uint32_t codebook_size = 1u << bits; + const size_t expanded_size = N_blocks * num_groups * gs * 4; + const size_t expanded_norms_size = N_blocks * num_groups * 4; + + if (M == 1) { + if (!has_expanded_layout || !has_sdot_group_layout) { + if (bits == 4 && (gs % 16) == 0) { cactus_quant_4bit_gemv(W, A, C); return; } + if (bits == 2 && (gs % 8) == 0) { cactus_quant_2bit_gemv(W, A, C); return; } + if (bits == 1 && (gs % 8) == 0) { cactus_quant_1bit_gemv(W, A, C); return; } + if (bits == 3 && (gs % 8) == 0) { cactus_quant_3bit_gemv(W, A, C); return; } + } + + if (!has_sdot_group_layout) { + cactus_quant_dispatch_group_gemm(W, A, M, C); + return; + } + + const size_t act_size = W->K; + static thread_local std::vector<__fp16> tl_code_basis; + if (tl_code_basis.size() < act_size) tl_code_basis.resize(act_size); + __fp16* code_basis_ptr = tl_code_basis.data(); + cactus_quant_transform_hadamard_activations(*W, A, M, code_basis_ptr); + + // INT8 SDOT GEMV + thread_local std::vector tl_act_i8; + thread_local std::vector tl_act_scales; + if (tl_act_i8.size() < act_size) tl_act_i8.resize(act_size); + if (tl_act_scales.size() < num_groups) tl_act_scales.resize(num_groups); + + for (uint32_t g = 0; g < num_groups; g++) { + tl_act_scales[g] = tq_quantize_group_i8( + code_basis_ptr + g * gs, + tl_act_i8.data() + g * gs, gs); + } + + const int8_t* act_i8 = tl_act_i8.data(); + const float* act_scales = tl_act_scales.data(); + + // Use pre-cached expanded weights if available, otherwise build inline + const int8_t* w_il = has_expanded_layout ? W->expanded : nullptr; + const float* n_f32 = has_expanded_layout ? W->norm_f32 : nullptr; + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, codebook_size); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + std::vector w_il_buf; + std::vector n_f32_buf; + if (!w_il) { + w_il_buf.resize(expanded_size); + n_f32_buf.resize(expanded_norms_size); + tq_preexpand_weights(W, bits, num_groups, gs, pgb, cb_lut, cb_scale, + N_blocks, w_il_buf.data(), n_f32_buf.data()); + w_il = w_il_buf.data(); + n_f32 = n_f32_buf.data(); + } + + auto& pool = CactusThreading::get_thread_pool(); + size_t num_threads = CactusThreading::GemmThreading::get_gemv_threads(N_blocks, pool.num_workers()); + num_threads = std::min(num_threads, N_blocks); + + auto process_blocks = [&](size_t block_start, size_t block_end) { + for (size_t n_block = block_start; n_block < block_end; ++n_block) { + float32x4_t running_sum = vdupq_n_f32(0.f); + const size_t nb_off = n_block * num_groups; + + uint32_t g = 0; + for (; g + 1 < num_groups; g += 2) { + const int8_t* a0 = act_i8 + g * gs; + const int8_t* a1 = act_i8 + (g + 1) * gs; + const int8_t* b0 = w_il + (nb_off + g) * gs * 4; + const int8_t* b1 = w_il + (nb_off + g + 1) * gs * 4; + + __builtin_prefetch(b1 + gs * 4, 0, 3); + + int32x4_t acc0 = vdupq_n_s32(0); + int32x4_t acc1 = vdupq_n_s32(0); + + for (uint32_t k = 0; k < gs; k += 32) { + int8x16_t av = vld1q_s8(a0 + k); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4), av, 0); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 16), av, 1); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 32), av, 2); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 48), av, 3); + av = vld1q_s8(a0 + k + 16); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 64), av, 0); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 80), av, 1); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 96), av, 2); + acc0 = CACTUS_DOTQ_LANE(acc0, vld1q_s8(b0 + k*4 + 112),av, 3); + } + + for (uint32_t k = 0; k < gs; k += 32) { + int8x16_t av = vld1q_s8(a1 + k); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4), av, 0); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 16), av, 1); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 32), av, 2); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 48), av, 3); + av = vld1q_s8(a1 + k + 16); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 64), av, 0); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 80), av, 1); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 96), av, 2); + acc1 = CACTUS_DOTQ_LANE(acc1, vld1q_s8(b1 + k*4 + 112),av, 3); + } + + float32x4_t s0 = vmulq_n_f32(vld1q_f32(n_f32 + (nb_off + g) * 4), act_scales[g]); + float32x4_t s1 = vmulq_n_f32(vld1q_f32(n_f32 + (nb_off + g + 1) * 4), act_scales[g + 1]); + running_sum = vmlaq_f32(running_sum, vcvtq_f32_s32(acc0), s0); + running_sum = vmlaq_f32(running_sum, vcvtq_f32_s32(acc1), s1); + } + + for (; g < num_groups; ++g) { + const int8_t* a_group = act_i8 + g * gs; + const int8_t* b_base = w_il + (nb_off + g) * gs * 4; + + int32x4_t acc = vdupq_n_s32(0); + for (uint32_t k = 0; k < gs; k += 32) { + int8x16_t av = vld1q_s8(a_group + k); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4), av, 0); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 16), av, 1); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 32), av, 2); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 48), av, 3); + av = vld1q_s8(a_group + k + 16); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 64), av, 0); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 80), av, 1); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 96), av, 2); + acc = CACTUS_DOTQ_LANE(acc, vld1q_s8(b_base + k*4 + 112),av, 3); + } + + float32x4_t scale = vmulq_n_f32(vld1q_f32(n_f32 + (nb_off + g) * 4), act_scales[g]); + running_sum = vmlaq_f32(running_sum, vcvtq_f32_s32(acc), scale); + } + + const size_t n_start = n_block * 4; + const size_t actual_n = std::min(size_t(4), static_cast(W->N) - n_start); + float16x4_t result = vcvt_f16_f32(running_sum); + if (actual_n == 4) { + vst1_f16(C + n_start, result); + } else { + for (size_t ni = 0; ni < actual_n; ni++) { + C[n_start + ni] = vget_lane_f16(result, 0); + result = vext_f16(result, result, 1); + } + } + } + }; + + if (num_threads <= 1) { + process_blocks(0, N_blocks); + } else { + pool.enqueue_n_threads(N_blocks, num_threads, process_blocks); + pool.wait_all(); + } + + return; + } + + if (!has_sdot_group_layout) { + cactus_quant_dispatch_group_gemm(W, A, M, C); + return; + } + + const size_t act_size = static_cast(M) * W->K; + std::vector<__fp16> heap_code_basis(act_size); + __fp16* code_basis_ptr = heap_code_basis.data(); + cactus_quant_transform_hadamard_activations(*W, A, M, code_basis_ptr); + + // M > 1: full pre-expansion amortized over M rows + const size_t scales_size = static_cast(M) * num_groups; + std::vector heap_act_i8(act_size); + std::vector heap_act_scales(scales_size); + int8_t* act_i8_ptr = heap_act_i8.data(); + float* act_scales_ptr = heap_act_scales.data(); + + for (uint32_t m = 0; m < M; m++) { + for (uint32_t g = 0; g < num_groups; g++) { + act_scales_ptr[m * num_groups + g] = tq_quantize_group_i8( + code_basis_ptr + m * W->K + g * gs, + act_i8_ptr + m * W->K + g * gs, gs); + } + } + + constexpr size_t TILE_M = 8; + + struct ExpandEntry { + uint64_t fingerprint = 0; + std::vector weights; + std::vector norms; + }; + static std::mutex s_expand_mutex; + static std::unordered_map s_expand_cache; + const size_t packed_bytes = static_cast(W->N) * num_groups * pgb; + uint64_t fp = 1469598103934665603ull; + auto fp_mix = [&fp](uint64_t v) { fp ^= v; fp *= 1099511628211ull; }; + fp_mix(bits); fp_mix(W->K); fp_mix(W->N); fp_mix(gs); + const uint8_t* pb = static_cast(W->packed_indices); + for (size_t i = 0; i < 64 && i < packed_bytes; ++i) fp_mix(pb[i]); + for (size_t i = packed_bytes > 64 ? packed_bytes - 64 : 0; i < packed_bytes; ++i) fp_mix(pb[i]); + const int8_t* w_il; + const float* n_f32; + { + std::lock_guard lock(s_expand_mutex); + auto& entry = s_expand_cache[W->packed_indices]; + if (entry.weights.empty() || entry.fingerprint != fp) { + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, codebook_size); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + entry.weights.assign(expanded_size, 0); + entry.norms.assign(expanded_norms_size, 0.0f); + tq_preexpand_weights(W, bits, num_groups, gs, pgb, cb_lut, cb_scale, + N_blocks, entry.weights.data(), entry.norms.data()); + entry.fingerprint = fp; + } + w_il = entry.weights.data(); + n_f32 = entry.norms.data(); + } + + const size_t M_blocks = (M + TILE_M - 1) / TILE_M; + const size_t total_tiles = M_blocks * N_blocks; + + CactusThreading::parallel_gemm_tiles(M, total_tiles, + [&](size_t tile_start, size_t tile_end) { + for (size_t tile_idx = tile_start; tile_idx < tile_end; ++tile_idx) { + const size_t m_block = tile_idx / N_blocks; + const size_t n_block = tile_idx % N_blocks; + const size_t m_start = m_block * TILE_M; + const size_t m_end = std::min(m_start + TILE_M, static_cast(M)); + const size_t actual_m = m_end - m_start; + const size_t n_start = n_block * 4; + const size_t actual_n = std::min(size_t(4), static_cast(W->N) - n_start); + + const int8_t* a_rows[TILE_M]; + for (size_t mi = 0; mi < TILE_M; mi++) { + size_t row = m_start + (mi < actual_m ? mi : actual_m - 1); + a_rows[mi] = act_i8_ptr + row * W->K; + } + + float32x4_t running_sum[TILE_M]; + for (size_t mi = 0; mi < TILE_M; mi++) + running_sum[mi] = vdupq_n_f32(0.f); + + for (uint32_t g = 0; g < num_groups; ++g) { + const int8_t* b_base = w_il + (n_block * num_groups + g) * gs * 4; + float32x4_t norms_v = vld1q_f32(n_f32 + (n_block * num_groups + g) * 4); + + __builtin_prefetch(b_base + gs * 4, 0, 3); + + int32x4_t row_acc[TILE_M]; + for (size_t mi = 0; mi < TILE_M; mi++) row_acc[mi] = vdupq_n_s32(0); + + for (uint32_t k = 0; k < gs; k += 32) { + const int8_t* bk = b_base + k * 4; + int8x16_t b00 = vld1q_s8(bk); + int8x16_t b01 = vld1q_s8(bk + 16); + int8x16_t b02 = vld1q_s8(bk + 32); + int8x16_t b03 = vld1q_s8(bk + 48); + int8x16_t b10 = vld1q_s8(bk + 64); + int8x16_t b11 = vld1q_s8(bk + 80); + int8x16_t b12 = vld1q_s8(bk + 96); + int8x16_t b13 = vld1q_s8(bk + 112); + +#define TQ_GEMM_ROW(ROW) do { \ + const int8_t* ap = a_rows[ROW] + g * gs + k; \ + int8x16_t a_lo = vld1q_s8(ap); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b00, a_lo, 0); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b01, a_lo, 1); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b02, a_lo, 2); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b03, a_lo, 3); \ + int8x16_t a_hi = vld1q_s8(ap + 16); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b10, a_hi, 0); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b11, a_hi, 1); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b12, a_hi, 2); \ + row_acc[ROW] = CACTUS_DOTQ_LANE(row_acc[ROW], b13, a_hi, 3); \ + } while(0) + + TQ_GEMM_ROW(0); + TQ_GEMM_ROW(1); + TQ_GEMM_ROW(2); + TQ_GEMM_ROW(3); + TQ_GEMM_ROW(4); + TQ_GEMM_ROW(5); + TQ_GEMM_ROW(6); + TQ_GEMM_ROW(7); + #undef TQ_GEMM_ROW + } + + for (size_t mi = 0; mi < actual_m; ++mi) { + float a_sc = act_scales_ptr[(m_start + mi) * num_groups + g]; + running_sum[mi] = vmlaq_f32(running_sum[mi], + vcvtq_f32_s32(row_acc[mi]), vmulq_n_f32(norms_v, a_sc)); + } + } + + for (size_t mi = 0; mi < actual_m; ++mi) { + float16x4_t r = vcvt_f16_f32(running_sum[mi]); + if (actual_n == 4) { + vst1_f16(C + (m_start + mi) * W->N + n_start, r); + } else { + for (size_t ni = 0; ni < actual_n; ni++) { + C[(m_start + mi) * W->N + n_start + ni] = vget_lane_f16(r, 0); + r = vext_f16(r, r, 1); + } + } + } + } + }); +} + +static void tq_fwht_normalized_f32(std::vector& x) { + const uint32_t n = static_cast(x.size()); + for (uint32_t h = 1; h < n; h <<= 1) { + for (uint32_t i = 0; i < n; i += (h << 1)) { + for (uint32_t j = i; j < i + h; ++j) { + float a = x[j]; + float b = x[j + h]; + x[j] = a + b; + x[j + h] = a - b; + } + } + } + const float inv = 1.0f / std::sqrt(static_cast(n)); + for (float& v : x) v *= inv; +} + +void cactus_quant_dequantize_hadamard_embedding_row( + uint32_t bits, + uint32_t hidden_dim, + uint32_t group_size, + uint32_t num_groups, + size_t row, + const uint8_t* packed_base, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const int8_t* left_signs, + const int8_t* right_signs, + const uint32_t* permutation, + __fp16* out_row) { + if (!packed_base || !codebook || !norms || !out_row || bits == 0 || bits > 4) return; + if (hidden_dim == 0 || group_size == 0 || num_groups == 0) return; + if ((group_size & (group_size - 1)) != 0) return; + if (group_size > 256) return; + + const uint32_t packed_group_bytes = cactus_quant_packed_group_bytes(bits, group_size); + std::vector rotated(group_size); + for (uint32_t g = 0; g < num_groups; ++g) { + std::fill(rotated.begin(), rotated.end(), 0.0f); + const uint8_t* packed = packed_base + (static_cast(row) * num_groups + g) * packed_group_bytes; + + for (uint32_t k = 0; k < group_size; ++k) { + uint8_t idx = static_cast(tq_extract_idx(packed, k, bits)); + uint32_t dst = permutation ? permutation[k] : k; + float rs = right_signs ? static_cast(right_signs[dst]) : 1.0f; + rotated[dst] = static_cast(codebook[idx]) * rs; + } + + tq_fwht_normalized_f32(rotated); + + const float norm = static_cast(norms[static_cast(row) * num_groups + g]); + for (uint32_t k = 0; k < group_size; ++k) { + uint32_t col = g * group_size + k; + float ls = left_signs ? static_cast(left_signs[k]) : 1.0f; + float scale = input_scale_recip ? static_cast(input_scale_recip[col]) : 1.0f; + out_row[col] = static_cast<__fp16>(rotated[k] * ls * norm * scale); + } + } + + for (uint32_t col = num_groups * group_size; col < hidden_dim; ++col) { + out_row[col] = static_cast<__fp16>(0); + } +} + +void cactus_quant_dequantize_orthogonal_embedding_row( + uint32_t bits, + uint32_t K, + size_t row, + const uint8_t* packed_base, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const __fp16* rotation, + uint32_t flags, + __fp16* out_row) { + if (!packed_base || !codebook || !norms || !rotation || !out_row || bits == 0 || bits > 4) return; + if (K == 0) return; + + const uint32_t packed_group_bytes = cactus_quant_packed_group_bytes(bits, K); + const bool interleaved = (flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0; + if (interleaved && (bits != 4 || (K % 8) != 0)) return; + const uint8_t* packed = interleaved ? nullptr : packed_base + static_cast(row) * packed_group_bytes; + const float norm = interleaved + ? static_cast(norms[(row / 4) * 4 + (row & 3u)]) + : static_cast(norms[row]); + + std::vector dq(K); + for (uint32_t i = 0; i < K; ++i) { + uint32_t idx = 0; + if (interleaved) { + const size_t panel_bytes = static_cast(4) * packed_group_bytes; + const uint8_t* panel = packed_base + (row / 4) * panel_bytes; + const uint32_t chunk = i / 8; + const uint32_t sub = i & 7u; + const uint8_t packed_byte = panel[chunk * 16 + (row & 3u) * 4 + (sub & 3u)]; + idx = (sub & 4u) ? static_cast(packed_byte >> 4) + : static_cast(packed_byte & 0x0F); + } else { + idx = tq_extract_idx(packed, i, bits); + } + dq[i] = static_cast(codebook[idx]); + } + + for (uint32_t j = 0; j < K; ++j) { + float acc = 0.0f; + for (uint32_t i = 0; i < K; ++i) { + acc += dq[i] * static_cast(rotation[static_cast(j) * K + i]); + } + float scale = input_scale_recip ? static_cast(input_scale_recip[j]) : 1.0f; + out_row[j] = static_cast<__fp16>(acc * norm * scale); + } +} + +static inline uint32_t tq_extract_interleaved_4row_4bit( + const uint8_t* packed_base, + uint32_t K, + size_t row, + uint32_t k) { + const uint32_t packed_group_bytes = cactus_quant_packed_group_bytes(4, K); + const size_t panel_bytes = static_cast(4) * packed_group_bytes; + const uint8_t* panel = packed_base + (row / 4) * panel_bytes; + const uint32_t chunk = k / 8; + const uint32_t sub = k & 7u; + const uint8_t packed_byte = panel[chunk * 16 + (row & 3u) * 4 + (sub & 3u)]; + return (sub & 4u) ? static_cast(packed_byte >> 4) + : static_cast(packed_byte & 0x0F); +} + +static bool cactus_quant_orthogonal_interleaved_lmhead_matmul( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C) { + if (!W || !A || !C || M == 0) return false; + if (W->bits != 4 || W->num_groups != 1 || W->group_size != W->K) return false; + if ((W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) == 0) return false; + if ((W->K % 32) != 0 || (W->N % 4) != 0) return false; + + const uint32_t K = W->K; + const uint32_t N = W->N; + const uint32_t packed_group_bytes = cactus_quant_packed_group_bytes(4, K); + const size_t panel_bytes = static_cast(4) * packed_group_bytes; + const size_t N_blocks = N / 4; + + thread_local std::vector ar32; + thread_local std::vector act_i8; + thread_local std::vector act_scales; + if (ar32.size() < static_cast(M) * K) ar32.resize(static_cast(M) * K); + if (act_i8.size() < static_cast(M) * K) act_i8.resize(static_cast(M) * K); + if (act_scales.size() < M) act_scales.resize(M); + + for (uint32_t m = 0; m < M; ++m) { + float* ar = ar32.data() + static_cast(m) * K; + std::fill(ar, ar + K, 0.0f); + const __fp16* a_row = A + static_cast(m) * K; + for (uint32_t k = 0; k < K; ++k) { + float a_val = static_cast(a_row[k]); + if (W->input_scale_recip) a_val *= static_cast(W->input_scale_recip[k]); + if (a_val == 0.0f) continue; + const __fp16* r = W->rotation + static_cast(k) * K; + const float32x4_t av = vdupq_n_f32(a_val); + uint32_t i = 0; + for (; i + 8 <= K; i += 8) { + const float16x8_t rv = vld1q_f16(r + i); + vst1q_f32(ar + i, vfmaq_f32(vld1q_f32(ar + i), av, vcvt_f32_f16(vget_low_f16(rv)))); + vst1q_f32(ar + i + 4, vfmaq_f32(vld1q_f32(ar + i + 4), av, vcvt_f32_f16(vget_high_f16(rv)))); + } + for (; i < K; ++i) ar[i] += a_val * static_cast(r[i]); + } + float max_abs = 0.0f; + for (uint32_t i = 0; i < K; ++i) max_abs = std::max(max_abs, std::abs(ar[i])); + const float as = max_abs > 0.0f ? max_abs / 127.0f : 1.0f; + act_scales[m] = as; + const float inv_as = 1.0f / as; + int8_t* qi = act_i8.data() + static_cast(m) * K; + for (uint32_t i = 0; i < K; ++i) { + int q = static_cast(std::lrintf(ar[i] * inv_as)); + q = std::max(-127, std::min(127, q)); + qi[i] = static_cast(q); + } + } + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 16); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + const int8_t* act_base = act_i8.data(); + const float* scales_base = act_scales.data(); + constexpr uint32_t TILE_M = 8; + + cactus_quant_parallel_ranges(N_blocks, 64, [&](size_t block_start, size_t block_end) { + for (size_t nb = block_start; nb < block_end; ++nb) { + const uint8_t* panel = W->packed_indices + nb * panel_bytes; + float32x4_t norm_v = vmulq_n_f32(vcvt_f32_f16(vld1_f16(W->norms + nb * 4)), cb_scale); + + for (uint32_t m0 = 0; m0 < M; m0 += TILE_M) { + const uint32_t tile = std::min(TILE_M, M - m0); + int32x4_t acc_a[TILE_M]; + int32x4_t acc_b[TILE_M]; + for (uint32_t mi = 0; mi < TILE_M; ++mi) { + acc_a[mi] = vdupq_n_s32(0); + acc_b[mi] = vdupq_n_s32(0); + } + for (uint32_t kb = 0; kb < K; kb += 16) { + const uint8x16_t p0 = vld1q_u8(panel + (kb / 8 + 0) * 16); + const int8x16_t w0 = vqtbl1q_s8(cb_lut, vandq_u8(p0, lo_mask)); + const int8x16_t w1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(p0, 4)); + const uint8x16_t p1 = vld1q_u8(panel + (kb / 8 + 1) * 16); + const int8x16_t w2 = vqtbl1q_s8(cb_lut, vandq_u8(p1, lo_mask)); + const int8x16_t w3 = vqtbl1q_s8(cb_lut, vshrq_n_u8(p1, 4)); + for (uint32_t mi = 0; mi < tile; ++mi) { + const int8x16_t a_v = vld1q_s8(act_base + static_cast(m0 + mi) * K + kb); + acc_a[mi] = CACTUS_DOTQ_LANE(acc_a[mi], w0, a_v, 0); + acc_b[mi] = CACTUS_DOTQ_LANE(acc_b[mi], w1, a_v, 1); + acc_a[mi] = CACTUS_DOTQ_LANE(acc_a[mi], w2, a_v, 2); + acc_b[mi] = CACTUS_DOTQ_LANE(acc_b[mi], w3, a_v, 3); + } + } + for (uint32_t mi = 0; mi < tile; ++mi) { + const int32x4_t dot = vaddq_s32(acc_a[mi], acc_b[mi]); + const float32x4_t scale = vmulq_n_f32(norm_v, scales_base[m0 + mi]); + vst1_f16(C + static_cast(m0 + mi) * N + nb * 4, + vcvt_f16_f32(vmulq_f32(vcvtq_f32_s32(dot), scale))); + } + } + } + }); + return true; +} + +void cactus_quant_orthogonal_matmul( + const CactusQuantMatrix* W, + const __fp16* A, + uint32_t M, + __fp16* C) { + if (!W || !A || !C || M == 0) return; + if (!W->rotation || !W->packed_indices || !W->codebook || !W->norms) return; + + const uint32_t K = W->K; + const uint32_t N = W->N; + const uint32_t bits = W->bits; + const uint32_t pgb = cactus_quant_packed_group_bytes(bits, K); + const bool interleaved = (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0; + + if (cactus_quant_orthogonal_interleaved_lmhead_matmul(W, A, M, C)) { + return; + } + const bool invalid_interleaved = + interleaved && (bits != 4 || W->num_groups != 1 || W->group_size != K || + (K % 8) != 0 || (N % 4) != 0); + if (invalid_interleaved) { + throw std::runtime_error("Invalid orthogonal INTERLEAVED_4ROW matmul layout"); + } + + const uint32_t n_cb = 1u << bits; + + std::vector cb_f32(n_cb); + for (uint32_t c = 0; c < n_cb; ++c) + cb_f32[c] = static_cast(W->codebook[c]); + + std::vector A_rot(static_cast(M) * K, 0.0f); + for (uint32_t m = 0; m < M; ++m) { + const __fp16* a_row = A + static_cast(m) * K; + float* ar_row = A_rot.data() + static_cast(m) * K; + for (uint32_t k = 0; k < K; ++k) { + float a_val = static_cast(a_row[k]); + if (W->input_scale_recip) a_val *= static_cast(W->input_scale_recip[k]); + if (a_val == 0.0f) continue; + const __fp16* R_row_k = W->rotation + static_cast(k) * K; + float32x4_t av = vdupq_n_f32(a_val); + uint32_t i = 0; + for (; i + 8 <= K; i += 8) { + float16x8_t rv = vld1q_f16(R_row_k + i); + float32x4_t r_lo = vcvt_f32_f16(vget_low_f16(rv)); + float32x4_t r_hi = vcvt_f32_f16(vget_high_f16(rv)); + vst1q_f32(ar_row + i, vfmaq_f32(vld1q_f32(ar_row + i), av, r_lo)); + vst1q_f32(ar_row + i + 4, vfmaq_f32(vld1q_f32(ar_row + i + 4), av, r_hi)); + } + for (; i < K; ++i) + ar_row[i] += a_val * static_cast(R_row_k[i]); + } + } + + if (!interleaved && bits == 4 && (K % 16) == 0) { + __fp16 cb_f16[16]; + uint8_t cb_lo_arr[16], cb_hi_arr[16]; + for (uint32_t c = 0; c < 16; ++c) { + cb_f16[c] = static_cast<__fp16>(cb_f32[c]); + cb_lo_arr[c] = reinterpret_cast(&cb_f16[c])[0]; + cb_hi_arr[c] = reinterpret_cast(&cb_f16[c])[1]; + } + uint8x16_t cb_lo_tbl = vld1q_u8(cb_lo_arr); + uint8x16_t cb_hi_tbl = vld1q_u8(cb_hi_arr); + + std::vector<__fp16> A_rot_f16; + A_rot_f16.resize(static_cast(M) * K); + for (uint32_t m = 0; m < M; ++m) { + const float* ar = A_rot.data() + static_cast(m) * K; + __fp16* arf = A_rot_f16.data() + static_cast(m) * K; + uint32_t i = 0; + for (; i + 8 <= K; i += 8) { + float32x4_t lo = vld1q_f32(ar + i); + float32x4_t hi = vld1q_f32(ar + i + 4); + vst1q_f16(arf + i, vcombine_f16(vcvt_f16_f32(lo), vcvt_f16_f32(hi))); + } + for (; i < K; ++i) arf[i] = static_cast<__fp16>(ar[i]); + } + + CactusThreading::parallel_for( + static_cast(N), + CactusThreading::ParallelConfig{64, 1}, + [&](size_t n_start, size_t n_end) { + for (size_t n = n_start; n < n_end; ++n) { + const uint8_t* packed = W->packed_indices + n * pgb; + float norm_n = static_cast(W->norms[n]); + for (uint32_t m = 0; m < M; ++m) { + const __fp16* arf = A_rot_f16.data() + static_cast(m) * K; + float32x4_t acc0 = vdupq_n_f32(0.f); + float32x4_t acc1 = vdupq_n_f32(0.f); + float32x4_t acc2 = vdupq_n_f32(0.f); + float32x4_t acc3 = vdupq_n_f32(0.f); + + for (uint32_t i = 0; i < K; i += 16) { + uint8x8_t raw8 = vld1_u8(packed + i / 2); + uint8x8_t lo_nibs = vand_u8(raw8, vdup_n_u8(0x0F)); + uint8x8_t hi_nibs = vshr_n_u8(raw8, 4); + uint8x8x2_t zipped = vzip_u8(lo_nibs, hi_nibs); + uint8x16_t indices16 = vcombine_u8(zipped.val[0], zipped.val[1]); + + uint8x16_t lo_bytes = vqtbl1q_u8(cb_lo_tbl, indices16); + uint8x16_t hi_bytes = vqtbl1q_u8(cb_hi_tbl, indices16); + + uint8x16x2_t fp16_bytes = vzipq_u8(lo_bytes, hi_bytes); + float16x8_t cb_lo = vreinterpretq_f16_u8(fp16_bytes.val[0]); + float16x8_t cb_hi = vreinterpretq_f16_u8(fp16_bytes.val[1]); + + float16x8_t ar_lo = vld1q_f16(arf + i); + float16x8_t ar_hi = vld1q_f16(arf + i + 8); + + acc0 = vfmaq_f32(acc0, vcvt_f32_f16(vget_low_f16(cb_lo)), + vcvt_f32_f16(vget_low_f16(ar_lo))); + acc1 = vfmaq_f32(acc1, vcvt_f32_f16(vget_high_f16(cb_lo)), + vcvt_f32_f16(vget_high_f16(ar_lo))); + acc2 = vfmaq_f32(acc2, vcvt_f32_f16(vget_low_f16(cb_hi)), + vcvt_f32_f16(vget_low_f16(ar_hi))); + acc3 = vfmaq_f32(acc3, vcvt_f32_f16(vget_high_f16(cb_hi)), + vcvt_f32_f16(vget_high_f16(ar_hi))); + } + + float acc = vaddvq_f32(vaddq_f32(vaddq_f32(acc0, acc1), + vaddq_f32(acc2, acc3))); + C[static_cast(m) * N + n] = static_cast<__fp16>(acc * norm_n); + } + } + }); + return; + } + + CactusThreading::parallel_for( + static_cast(N), + CactusThreading::ParallelConfig{64, 1}, + [&](size_t n_start, size_t n_end) { + for (size_t n = n_start; n < n_end; ++n) { + const uint8_t* packed = W->packed_indices + n * pgb; + float norm_n = static_cast(W->norms[n]); + for (uint32_t m = 0; m < M; ++m) { + const float* ar = A_rot.data() + static_cast(m) * K; + float acc = 0.0f; + for (uint32_t i = 0; i < K; ++i) { + uint8_t idx = interleaved + ? static_cast(tq_extract_interleaved_4row_4bit(W->packed_indices, K, n, i)) + : static_cast(tq_extract_idx(packed, i, bits)); + acc += cb_f32[idx] * ar[i]; + } + C[static_cast(m) * N + n] = static_cast<__fp16>(acc * norm_n); + } + } + }); +} + +static void cactus_quant_interleaved4_gemv_blocks( + const CactusQuantMatrix* W, const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, const int8_t* act_i8, const float* act_scales, + const int8x16_t cb_lut, const float cb_scale, + size_t block_start, size_t block_end, __fp16* y) { + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(4, gs); + const uint32_t num_groups = W->num_groups; + const size_t panel_bytes = 4 * pgb; + const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + size_t nb = block_start; + size_t aligned_end = (block_end / 2) * 2; + + for (; nb + 2 <= aligned_end; nb += 2) { + float32x4_t acc0 = vdupq_n_f32(0.f); + float32x4_t acc1 = vdupq_n_f32(0.f); + + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p0 = packed_interleaved + ((nb + 0) * num_groups + g) * panel_bytes; + const uint8_t* p1 = packed_interleaved + ((nb + 1) * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + + int32x4_t d0a = vdupq_n_s32(0), d0b = vdupq_n_s32(0); + int32x4_t d1a = vdupq_n_s32(0), d1b = vdupq_n_s32(0); + + for (uint32_t kb = 0; kb < gs; kb += 16) { + int8x16_t a_v = vld1q_s8(a_grp + kb); + + #define PROC_PANEL(P, DA, DB) do { \ + const uint8_t* c0 = (P) + (kb / 8 + 0) * 16; \ + const uint8_t* c1 = (P) + (kb / 8 + 1) * 16; \ + uint8x16_t b0 = vld1q_u8(c0); \ + int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); \ + int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); \ + DA = vdotq_laneq_s32(DA, wl0, a_v, 0); \ + DB = vdotq_laneq_s32(DB, wh0, a_v, 1); \ + uint8x16_t b1 = vld1q_u8(c1); \ + int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); \ + int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); \ + DA = vdotq_laneq_s32(DA, wl1, a_v, 2); \ + DB = vdotq_laneq_s32(DB, wh1, a_v, 3); \ + } while (0) + + PROC_PANEL(p0, d0a, d0b); + PROC_PANEL(p1, d1a, d1b); + #undef PROC_PANEL + } + + int32x4_t dot0 = vaddq_s32(d0a, d0b); + int32x4_t dot1 = vaddq_s32(d1a, d1b); + + float32x4_t n0 = vcvt_f32_f16(vld1_f16(norms_interleaved + ((nb + 0) * num_groups + g) * 4)); + float32x4_t n1 = vcvt_f32_f16(vld1_f16(norms_interleaved + ((nb + 1) * num_groups + g) * 4)); + float scale_grp = cb_scale * a_sc; + acc0 = vfmaq_f32(acc0, vcvtq_f32_s32(dot0), vmulq_n_f32(n0, scale_grp)); + acc1 = vfmaq_f32(acc1, vcvtq_f32_s32(dot1), vmulq_n_f32(n1, scale_grp)); + } + + vst1_f16(y + (nb + 0) * 4, vcvt_f16_f32(acc0)); + vst1_f16(y + (nb + 1) * 4, vcvt_f16_f32(acc1)); + } + + for (; nb < block_end; ++nb) { + float32x4_t acc = vdupq_n_f32(0.f); + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + + int32x4_t dot_a = vdupq_n_s32(0); + int32x4_t dot_b = vdupq_n_s32(0); + + for (uint32_t kb = 0; kb < gs; kb += 16) { + int8x16_t a_v = vld1q_s8(a_grp + kb); + uint8x16_t b0 = vld1q_u8(p_base + (kb / 8 + 0) * 16); + int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); + int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); + dot_a = vdotq_laneq_s32(dot_a, wl0, a_v, 0); + dot_b = vdotq_laneq_s32(dot_b, wh0, a_v, 1); + uint8x16_t b1 = vld1q_u8(p_base + (kb / 8 + 1) * 16); + int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); + int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); + dot_a = vdotq_laneq_s32(dot_a, wl1, a_v, 2); + dot_b = vdotq_laneq_s32(dot_b, wh1, a_v, 3); + } + + int32x4_t dot = vaddq_s32(dot_a, dot_b); + float32x4_t norm = vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)); + norm = vmulq_n_f32(norm, cb_scale * a_sc); + acc = vfmaq_f32(acc, vcvtq_f32_s32(dot), norm); + } + vst1_f16(y + nb * 4, vcvt_f16_f32(acc)); + } +} + +static void cactus_quant_interleaved4_gemm_blocks( + const CactusQuantMatrix* W, const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, const int8_t* act_i8, const float* act_scales, + uint32_t M, const int8x16_t cb_lut, const float cb_scale, + size_t block_start, size_t block_end, __fp16* C) { + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(4, gs); + const uint32_t num_groups = W->num_groups; + const uint32_t N = W->N; + const uint32_t K = W->K; + const size_t panel_bytes = 4 * pgb; + const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + constexpr uint32_t TILE_M = 8; + + for (size_t nb = block_start; nb < block_end; ++nb) { + for (uint32_t m0 = 0; m0 < M; m0 += TILE_M) { + const uint32_t tile = std::min(TILE_M, M - m0); + float32x4_t acc[TILE_M]; + for (uint32_t mi = 0; mi < TILE_M; ++mi) acc[mi] = vdupq_n_f32(0.f); + + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; + float32x4_t norm_g = vmulq_n_f32( + vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)), cb_scale); + + int32x4_t dot_a[TILE_M]; + int32x4_t dot_b[TILE_M]; + for (uint32_t mi = 0; mi < TILE_M; ++mi) { dot_a[mi] = vdupq_n_s32(0); dot_b[mi] = vdupq_n_s32(0); } + + for (uint32_t kb = 0; kb < gs; kb += 16) { + const uint8x16_t b0 = vld1q_u8(p_base + (kb / 8 + 0) * 16); + const int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); + const int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); + const uint8x16_t b1 = vld1q_u8(p_base + (kb / 8 + 1) * 16); + const int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); + const int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); + for (uint32_t mi = 0; mi < tile; ++mi) { + const int8x16_t a_v = vld1q_s8(act_i8 + static_cast(m0 + mi) * K + g * gs + kb); + dot_a[mi] = vdotq_laneq_s32(dot_a[mi], wl0, a_v, 0); + dot_b[mi] = vdotq_laneq_s32(dot_b[mi], wh0, a_v, 1); + dot_a[mi] = vdotq_laneq_s32(dot_a[mi], wl1, a_v, 2); + dot_b[mi] = vdotq_laneq_s32(dot_b[mi], wh1, a_v, 3); + } + } + for (uint32_t mi = 0; mi < tile; ++mi) { + const int32x4_t dot = vaddq_s32(dot_a[mi], dot_b[mi]); + const float a_sc = act_scales[static_cast(m0 + mi) * num_groups + g]; + acc[mi] = vfmaq_f32(acc[mi], vcvtq_f32_s32(dot), vmulq_n_f32(norm_g, a_sc)); + } + } + for (uint32_t mi = 0; mi < tile; ++mi) + vst1_f16(C + static_cast(m0 + mi) * N + nb * 4, vcvt_f16_f32(acc[mi])); + } + } +} + +static void cactus_quant_4bit_gemm_interleaved( + const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C) { + const uint32_t gs = W->group_size; + const uint32_t num_groups = W->num_groups; + const size_t N_blocks = W->N / 4; + + thread_local std::vector act_i8; + thread_local std::vector act_scales; + if (act_i8.size() < static_cast(M) * W->K) act_i8.resize(static_cast(M) * W->K); + if (act_scales.size() < static_cast(M) * num_groups) act_scales.resize(static_cast(M) * num_groups); + + for (uint32_t m = 0; m < M; ++m) { + __fp16 basis[256]; + const __fp16* a_row = A + static_cast(m) * W->K; + for (uint32_t g = 0; g < num_groups; ++g) { + cactus_quant_transform_hadamard_group(*W, a_row + static_cast(g) * gs, g, basis); + act_scales[static_cast(m) * num_groups + g] = + tq_quantize_group_i8(basis, act_i8.data() + static_cast(m) * W->K + g * gs, gs); + } + } + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 16); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + const int8_t* act_base = act_i8.data(); + const float* scales_base = act_scales.data(); + + cactus_quant_parallel_ranges(N_blocks, 64, [&](size_t block_start, size_t block_end) { + cactus_quant_interleaved4_gemm_blocks(W, W->packed_indices, W->norms, + act_base, scales_base, M, cb_lut, cb_scale, + block_start, block_end, C); + }); +} + +void cactus_quant_4bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 4) return; + if (W->N % 4 != 0) return; + if ((W->group_size % 32) != 0) return; + if (W->group_size > 256) return; + + const uint32_t gs = W->group_size; + const uint32_t num_groups = W->num_groups; + const size_t N_blocks = W->N / 4; + const size_t n_chunks = (N_blocks + 15) / 16; + auto& pool = CactusThreading::get_thread_pool(); + const size_t sb_per_thread = cactus_quant_gemv_sb_per_thread(); + const size_t nt_budget = std::max(1, (n_chunks + sb_per_thread - 1) / sb_per_thread); + const size_t nt = std::min(pool.num_workers(), std::min(nt_budget, n_chunks)); + + static thread_local std::vector tl_il_act_i8; + static thread_local std::vector tl_il_act_scales; + if (tl_il_act_i8.size() < W->K) tl_il_act_i8.resize(W->K); + if (tl_il_act_scales.size() < num_groups) tl_il_act_scales.resize(num_groups); + int8_t* act_i8 = tl_il_act_i8.data(); + float* act_scales = tl_il_act_scales.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 16); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + auto phase_a_group = [&](uint32_t g) { + __fp16 basis[256]; + cactus_quant_transform_hadamard_group(*W, x + static_cast(g) * gs, g, basis); + act_scales[g] = tq_quantize_group_i8(basis, act_i8 + static_cast(g) * gs, gs); + }; + + if (nt <= 1) { + for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); + cactus_quant_interleaved4_gemv_blocks(W, packed_interleaved, norms_interleaved, + act_i8, act_scales, cb_lut, cb_scale, + 0, N_blocks, y); + return; + } + + cactus_quant_two_phase_run(nt, num_groups, static_cast(n_chunks), phase_a_group, + [&](size_t, uint32_t ck, uint32_t cnt) { + const size_t b0 = static_cast(ck) * 16; + const size_t b1 = std::min(N_blocks, b0 + static_cast(cnt) * 16); + cactus_quant_interleaved4_gemv_blocks(W, packed_interleaved, norms_interleaved, + act_i8, act_scales, cb_lut, cb_scale, + b0, b1, y); + }); +} + +void cactus_quant_3bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 3 || W->N % 4 != 0 || (W->group_size % 32) != 0 || W->group_size > 256) return; + const uint32_t gs = W->group_size; + const uint32_t num_groups = W->num_groups; + const size_t panel_bytes = static_cast(gs) * 3 / 2; + const size_t N_blocks = W->N / 4; + const size_t n_chunks = (N_blocks + 15) / 16; + auto& pool = CactusThreading::get_thread_pool(); + const size_t sb_per_thread = cactus_quant_gemv_sb_per_thread(); + const size_t nt_budget = std::max(1, (n_chunks + sb_per_thread - 1) / sb_per_thread); + const size_t nt = std::min(pool.num_workers(), std::min(nt_budget, n_chunks)); + + static thread_local std::vector tl_act_i8; + static thread_local std::vector tl_act_scales; + if (tl_act_i8.size() < W->K) tl_act_i8.resize(W->K); + if (tl_act_scales.size() < num_groups) tl_act_scales.resize(num_groups); + int8_t* act_i8 = tl_act_i8.data(); + float* act_scales = tl_act_scales.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 8); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + auto phase_a_group = [&](uint32_t g) { + __fp16 basis[256]; + cactus_quant_transform_hadamard_group(*W, x + static_cast(g) * gs, g, basis); + act_scales[g] = tq_quantize_group_i8(basis, act_i8 + static_cast(g) * gs, gs); + }; + + auto phase_b = [&](size_t b0, size_t b1) { + for (size_t nb = b0; nb < b1; ++nb) { + float32x4_t acc = vdupq_n_f32(0.f); + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + int32x4_t dot_a = vdupq_n_s32(0); + int32x4_t dot_b = vdupq_n_s32(0); + for (uint32_t kb = 0; kb < gs; kb += 16) { + int8x16_t a_v = vld1q_s8(a_grp + kb); + const uint8_t* band_base = p_base + (kb / 16) * (4 * 6); + int8x16_t w0 = tq_expand_i8_16(band_base + 0, 3, cb_lut); + int8x16_t w1 = tq_expand_i8_16(band_base + 6, 3, cb_lut); + int8x16_t w2 = tq_expand_i8_16(band_base + 12, 3, cb_lut); + int8x16_t w3 = tq_expand_i8_16(band_base + 18, 3, cb_lut); + dot_a = vdotq_laneq_s32(dot_a, w0, a_v, 0); + dot_b = vdotq_laneq_s32(dot_b, w1, a_v, 1); + dot_a = vdotq_laneq_s32(dot_a, w2, a_v, 2); + dot_b = vdotq_laneq_s32(dot_b, w3, a_v, 3); + } + int32x4_t dot = vaddq_s32(dot_a, dot_b); + float32x4_t norm = vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)); + norm = vmulq_n_f32(norm, cb_scale * a_sc); + acc = vfmaq_f32(acc, vcvtq_f32_s32(dot), norm); + } + vst1_f16(y + nb * 4, vcvt_f16_f32(acc)); + } + }; + + if (nt <= 1) { + for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); + phase_b(0, N_blocks); + return; + } + cactus_quant_two_phase_run(nt, num_groups, static_cast(n_chunks), phase_a_group, + [&](size_t, uint32_t ck, uint32_t cnt) { + const size_t b0 = static_cast(ck) * 16; + const size_t b1 = std::min(N_blocks, b0 + static_cast(cnt) * 16); + phase_b(b0, b1); + }); +} + +void cactus_quant_2bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 2 || W->N % 4 != 0 || (W->group_size % 32) != 0 || W->group_size > 256) return; + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(2, gs); + const uint32_t num_groups = W->num_groups; + const size_t panel_bytes = 4 * pgb; + const size_t N_blocks = W->N / 4; + const size_t n_chunks = (N_blocks + 15) / 16; + auto& pool = CactusThreading::get_thread_pool(); + const size_t sb_per_thread = cactus_quant_gemv_sb_per_thread(); + const size_t nt_budget = std::max(1, (n_chunks + sb_per_thread - 1) / sb_per_thread); + const size_t nt = std::min(pool.num_workers(), std::min(nt_budget, n_chunks)); + + static thread_local std::vector tl_act_i8; + static thread_local std::vector tl_act_scales; + if (tl_act_i8.size() < W->K) tl_act_i8.resize(W->K); + if (tl_act_scales.size() < num_groups) tl_act_scales.resize(num_groups); + int8_t* act_i8 = tl_act_i8.data(); + float* act_scales = tl_act_scales.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 4); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + const int8x16_t shifts = vcombine_s8( + (int8x8_t){0,-2,-4,-6, 0,-2,-4,-6}, + (int8x8_t){0,-2,-4,-6, 0,-2,-4,-6}); + const uint8x16_t lookup_s0 = (uint8x16_t){0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3}; + const uint8x16_t lookup_s1 = (uint8x16_t){4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7}; + const uint8x16_t lookup_s2 = (uint8x16_t){8,8,8,8, 9,9,9,9, 10,10,10,10, 11,11,11,11}; + const uint8x16_t lookup_s3 = (uint8x16_t){12,12,12,12, 13,13,13,13, 14,14,14,14, 15,15,15,15}; + const uint8x16_t idx_mask = vdupq_n_u8(0x03); + + auto phase_a_group = [&](uint32_t g) { + __fp16 basis[256]; + cactus_quant_transform_hadamard_group(*W, x + static_cast(g) * gs, g, basis); + act_scales[g] = tq_quantize_group_i8(basis, act_i8 + static_cast(g) * gs, gs); + }; + + auto phase_b = [&](size_t b0, size_t b1) { + for (size_t nb = b0; nb < b1; ++nb) { + float32x4_t acc = vdupq_n_f32(0.f); + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + int32x4_t dot_a = vdupq_n_s32(0); + int32x4_t dot_b = vdupq_n_s32(0); + for (uint32_t chunk = 0; chunk < gs / 16; ++chunk) { + int8x16_t a_v = vld1q_s8(a_grp + chunk * 16); + uint8x16_t bytes = vld1q_u8(p_base + chunk * 16); + auto unpack_set = [&](uint8x16_t lookup) -> int8x16_t { + uint8x16_t spread = vqtbl1q_u8(bytes, lookup); + uint8x16_t shifted = vshlq_u8(spread, shifts); + uint8x16_t idx = vandq_u8(shifted, idx_mask); + return vqtbl1q_s8(cb_lut, idx); + }; + int8x16_t w0 = unpack_set(lookup_s0); + int8x16_t w1 = unpack_set(lookup_s1); + int8x16_t w2 = unpack_set(lookup_s2); + int8x16_t w3 = unpack_set(lookup_s3); + dot_a = vdotq_laneq_s32(dot_a, w0, a_v, 0); + dot_b = vdotq_laneq_s32(dot_b, w1, a_v, 1); + dot_a = vdotq_laneq_s32(dot_a, w2, a_v, 2); + dot_b = vdotq_laneq_s32(dot_b, w3, a_v, 3); + } + int32x4_t dot = vaddq_s32(dot_a, dot_b); + float32x4_t norm = vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)); + norm = vmulq_n_f32(norm, cb_scale * a_sc); + acc = vfmaq_f32(acc, vcvtq_f32_s32(dot), norm); + } + vst1_f16(y + nb * 4, vcvt_f16_f32(acc)); + } + }; + + if (nt <= 1) { + for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); + phase_b(0, N_blocks); + return; + } + cactus_quant_two_phase_run(nt, num_groups, static_cast(n_chunks), phase_a_group, + [&](size_t, uint32_t ck, uint32_t cnt) { + const size_t b0 = static_cast(ck) * 16; + const size_t b1 = std::min(N_blocks, b0 + static_cast(cnt) * 16); + phase_b(b0, b1); + }); +} + +void cactus_quant_1bit_gemv_interleaved( + const CactusQuantMatrix* W, + const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, + const __fp16* x, + __fp16* y) { + if (!cactus_quant_valid_common(W, x, y)) return; + if (W->bits != 1 || W->N % 4 != 0 || (W->group_size % 32) != 0 || W->group_size > 256) return; + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(1, gs); + const uint32_t num_groups = W->num_groups; + + thread_local std::vector<__fp16> code_basis_buf; + if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); + cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); + const __fp16* code_basis = code_basis_buf.data(); + + thread_local std::vector act_i8_buf; + thread_local std::vector act_scales_buf; + if (act_i8_buf.size() < W->K) act_i8_buf.resize(W->K); + if (act_scales_buf.size() < num_groups) act_scales_buf.resize(num_groups); + for (uint32_t g = 0; g < num_groups; ++g) { + act_scales_buf[g] = tq_quantize_group_i8( + code_basis + static_cast(g) * gs, + act_i8_buf.data() + static_cast(g) * gs, gs); + } + const int8_t* act_i8 = act_i8_buf.data(); + const float* act_scales = act_scales_buf.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 2); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + const int8x16_t shifts_lo = vcombine_s8( + (int8x8_t){0,-1,-2,-3, 0,-1,-2,-3}, + (int8x8_t){0,-1,-2,-3, 0,-1,-2,-3}); + const int8x16_t shifts_hi = vcombine_s8( + (int8x8_t){-4,-5,-6,-7, -4,-5,-6,-7}, + (int8x8_t){-4,-5,-6,-7, -4,-5,-6,-7}); + const uint8x16_t idx_mask = vdupq_n_u8(0x01); + const uint8x16_t lookup_s0 = (uint8x16_t){0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3}; + const uint8x16_t lookup_s1 = (uint8x16_t){4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7}; + const uint8x16_t lookup_s2 = (uint8x16_t){8,8,8,8, 9,9,9,9, 10,10,10,10, 11,11,11,11}; + const uint8x16_t lookup_s3 = (uint8x16_t){12,12,12,12, 13,13,13,13, 14,14,14,14, 15,15,15,15}; + + const size_t panel_bytes = 4 * pgb; + const size_t N_blocks = W->N / 4; + + cactus_quant_parallel_ranges(N_blocks, 64, [&](size_t block_start, size_t block_end) { + for (size_t nb = block_start; nb < block_end; ++nb) { + float32x4_t acc = vdupq_n_f32(0.f); + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + + int32x4_t dot_a = vdupq_n_s32(0); + int32x4_t dot_b = vdupq_n_s32(0); + + for (uint32_t kb = 0; kb < gs; kb += 32) { + int8x16_t a_lo = vld1q_s8(a_grp + kb); + int8x16_t a_hi = vld1q_s8(a_grp + kb + 16); + uint8x16_t bytes = vld1q_u8(p_base + (kb / 32) * 16); + + auto unpack_kq = [&](uint8x16_t lookup, int8x16_t sh) -> int8x16_t { + uint8x16_t spread = vqtbl1q_u8(bytes, lookup); + uint8x16_t shifted = vshlq_u8(spread, sh); + uint8x16_t idx = vandq_u8(shifted, idx_mask); + return vqtbl1q_s8(cb_lut, idx); + }; + + int8x16_t w0_lo = unpack_kq(lookup_s0, shifts_lo); + int8x16_t w0_hi = unpack_kq(lookup_s0, shifts_hi); + int8x16_t w1_lo = unpack_kq(lookup_s1, shifts_lo); + int8x16_t w1_hi = unpack_kq(lookup_s1, shifts_hi); + int8x16_t w2_lo = unpack_kq(lookup_s2, shifts_lo); + int8x16_t w2_hi = unpack_kq(lookup_s2, shifts_hi); + int8x16_t w3_lo = unpack_kq(lookup_s3, shifts_lo); + int8x16_t w3_hi = unpack_kq(lookup_s3, shifts_hi); + + dot_a = vdotq_laneq_s32(dot_a, w0_lo, a_lo, 0); + dot_b = vdotq_laneq_s32(dot_b, w0_hi, a_lo, 1); + dot_a = vdotq_laneq_s32(dot_a, w1_lo, a_lo, 2); + dot_b = vdotq_laneq_s32(dot_b, w1_hi, a_lo, 3); + dot_a = vdotq_laneq_s32(dot_a, w2_lo, a_hi, 0); + dot_b = vdotq_laneq_s32(dot_b, w2_hi, a_hi, 1); + dot_a = vdotq_laneq_s32(dot_a, w3_lo, a_hi, 2); + dot_b = vdotq_laneq_s32(dot_b, w3_hi, a_hi, 3); + } + + int32x4_t dot = vaddq_s32(dot_a, dot_b); + float32x4_t norm = vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)); + norm = vmulq_n_f32(norm, cb_scale * a_sc); + acc = vfmaq_f32(acc, vcvtq_f32_s32(dot), norm); + } + vst1_f16(y + nb * 4, vcvt_f16_f32(acc)); + } + }); +} diff --git a/cactus-kernels/src/metal_backend.mm b/cactus-kernels/src/metal_backend.mm new file mode 100644 index 000000000..0c4a5f236 --- /dev/null +++ b/cactus-kernels/src/metal_backend.mm @@ -0,0 +1,2489 @@ + + +#include "metal_backend.h" + +#if CACTUS_HAS_METAL + +#import +#import +#include +#include +#if TARGET_OS_OSX +#include +#endif +#import +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cactus_kernels_msl.h" +#include "threading.h" +#include + +namespace { + +struct MetalCtx { + id dev = nil; + id queue = nil; + id psoT = nil; + id psoTsimd = nil; + id psoTbatch = nil; + id psoMega = nil, psoSwiT = nil, psoCat = nil; + id psoG = nil, psoGmr = nil, psoGmrLow = nil; + id psoActQ = nil, psoG2i8 = nil; + id psoTm = nil; + id psoGmma = nil, psoGdense = nil; + id psoRotW = nil, psoEmbOW = nil, psoRmsAddRms = nil; + id psoEmbH = nil; + id psoEmbOm = nil, psoEmbHm = nil; + id psoCopy=nil, psoBinary=nil, psoScalar=nil, psoUnary=nil, psoRms=nil, psoSwiglu=nil, psoRmsAdd=nil, psoRmsAddScale=nil; + id psoCF16F32=nil, psoCF32F16=nil, psoCI8F16=nil, psoCF16I8=nil; + id psoAttn=nil, psoAttnC=nil, psoAttnF=nil, psoStrided=nil, psoScatter=nil, psoBcast=nil, psoKvAppend=nil; + id psoAttnPre=nil, psoAttnPreMma2=nil, psoAttnPreHd256=nil; + id psoKvAppendM=nil, psoKvAppendRingM=nil; + id psoSlideS=nil, psoSlideR=nil, psoSlideRM=nil; + id psoArgmax=nil; + id psoArgmaxP=nil, psoArgmaxF=nil, psoSoftcap=nil, psoAdjust=nil; + id psoGather=nil; + id psoBinF32=nil, psoScaF32=nil, psoUnaF32=nil; + id psoClampF16=nil, psoClampF32=nil, psoGlu=nil; + id psoLayerNorm=nil, psoSoftmaxR=nil, psoConv1dK3=nil, psoAttnDense=nil, psoAttnD64=nil; + id psoGemmBatch=nil, psoGemmBatchF32A=nil, psoConvDw=nil, psoStridedRows=nil, psoMaskFill=nil; + id psoReduceF16=nil, psoReduceF32=nil, psoCumsumF16=nil, psoCumsumF32=nil; + id psoConcat2=nil, psoGatherF32Idx=nil, psoRopeFull=nil, psoMaxpool1d=nil; + id psoBilinear=nil, psoConv1dGen=nil, psoConv1dNlcDw=nil, psoConv2d=nil; + id psoBatchnorm=nil, psoGroupnorm=nil, psoBiasAddRows=nil, psoEwChain=nil; + id psoAttnFlash=nil, psoRmsSimd=nil, psoScatterRows=nil; + id psoTranspose2d=nil, psoBcastRows=nil, psoRmsAddSimd=nil; + id psoConvCacheAppend=nil, psoRelPosBias=nil, psoGemvBias=nil; + id psoTopkRows=nil, psoMoeT=nil, psoMoeUp=nil; + id psoMoeT2=nil, psoMoeDownAcc=nil, psoRopePair=nil; + id psoRmsScale=nil, psoSoftmaxTopk=nil, psoRopePairRms=nil; + id psoRms2AddClip=nil, psoDeltanet=nil, psoDeltanetPre=nil; + id dummy=nil; + bool ok = false; + + MetalCtx() { @autoreleasepool { + dev = MTLCreateSystemDefaultDevice(); + if (!dev) return; + queue = [dev newCommandQueue]; + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:kCactusMSL] + options:nil error:&err]; + if (!lib) { if (err) fprintf(stderr,"[cactus-metal] MSL compile failed: %s\n",[[err localizedDescription] UTF8String]); return; } + auto pso = [&](const char* name) -> id { + NSError* e=nil; + id f=[lib newFunctionWithName:[NSString stringWithUTF8String:name]]; + id p = f ? [dev newComputePipelineStateWithFunction:f error:&e] : nil; + if (!p) fprintf(stderr,"[cactus-metal] pipeline '%s' failed: %s\n", name, e?[[e localizedDescription] UTF8String]:"function not found"); + return p; + }; + psoT=pso("cq4_transform"); psoTsimd=pso("cq4_transform_simd"); psoG=pso("cq4_gemv"); psoTbatch=pso("cq4_transform_batch"); psoGmr=pso("cq4_gemv_mr"); + psoMega=pso("cq4_transform_gemv"); psoSwiT=pso("cq4_swiglu_transform"); psoCat=pso("cq4_gemv_cat"); + psoTm=pso("cq4_transform_m"); psoGmma=pso("cq4_gemm_mma"); + psoGdense=pso("cq4_gemm_dense_f16"); + psoGmrLow=pso("cq_gemv_mr_lowbit"); + psoActQ=pso("cq_act_quant_i8"); psoG2i8=pso("cq2_gemv_i8"); + psoRotW=pso("lmhead_rotate_wide"); psoEmbOW=pso("emb_ortho_wide"); psoRmsAddRms=pso("rms_norm_add_rms_f16"); + psoEmbH=pso("emb_hadamard"); + psoEmbOm=pso("emb_ortho_m"); psoEmbHm=pso("emb_hadamard_m"); + psoGather=pso("gather_f16"); + psoCopy=pso("copy_bytes"); psoBinary=pso("binary_f16"); psoScalar=pso("scalar_f16"); + psoUnary=pso("unary_f16"); psoRms=pso("rms_norm_f16"); psoSwiglu=pso("swiglu_f16"); psoRmsAdd=pso("rms_norm_add_f16"); psoRmsAddScale=pso("rms_norm_add_scale_f16"); + psoCF16F32=pso("cast_f16_f32"); psoCF32F16=pso("cast_f32_f16"); + psoCI8F16=pso("cast_i8_f16"); psoCF16I8=pso("cast_f16_i8"); + psoAttn=pso("attn_decode_i8"); psoAttnC=pso("attn_decode_combine"); psoAttnF=pso("attn_decode_fused_i8"); + psoStrided=pso("strided_copy_f16"); psoBcast=pso("bcast_binary_f16"); + psoAttnPre=pso("attn_prefill_i8"); psoAttnPreMma2=pso("attn_prefill_mma2"); psoAttnPreHd256=pso("attn_prefill_mma_hd256"); psoKvAppendM=pso("kv_append_i8_m"); + psoKvAppendRingM=pso("kv_append_ring_i8_m"); + psoSlideS=pso("kv_slide_save"); psoSlideR=pso("kv_slide_restore"); psoSlideRM=pso("kv_slide_restore_m"); + psoScatter=pso("strided_scatter_f16"); psoKvAppend=pso("kv_append_i8"); + psoArgmax=pso("argmax_logits"); + psoArgmaxP=pso("argmax_part"); psoArgmaxF=pso("argmax_final"); psoSoftcap=pso("softcap_f16"); + psoAdjust=pso("adjust_logits"); + psoBinF32=pso("binary_f32"); psoScaF32=pso("scalar_f32"); psoUnaF32=pso("unary_f32"); + psoClampF16=pso("clamp_f16"); psoClampF32=pso("clamp_f32"); psoGlu=pso("glu_f16"); + psoLayerNorm=pso("layer_norm_f16"); psoSoftmaxR=pso("softmax_rows_f16"); + psoConv1dK3=pso("conv1d_k3_f16"); psoAttnDense=pso("attn_f16"); psoAttnD64=pso("attn_f16_d64"); + psoGemmBatch=pso("gemm_batch_f16"); psoGemmBatchF32A=pso("gemm_batch_f32a"); psoConvDw=pso("conv1d_dw_f16"); + psoStridedRows=pso("strided_copy_rows_f16"); psoMaskFill=pso("attn_maskfill_f16"); + psoReduceF16=pso("reduce_axis_f16"); psoReduceF32=pso("reduce_axis_f32"); + psoCumsumF16=pso("cumsum_f16"); psoCumsumF32=pso("cumsum_f32"); + psoConcat2=pso("concat2_f16"); psoGatherF32Idx=pso("gather_f32idx_f16"); + psoRopeFull=pso("rope_full_f16"); psoMaxpool1d=pso("maxpool1d_f16"); + psoBilinear=pso("bilinear_f16"); psoConv1dGen=pso("conv1d_gen_f16"); + psoConv1dNlcDw=pso("conv1d_nlc_dw_f16"); psoConv2d=pso("conv2d_f16"); + psoBatchnorm=pso("batchnorm_f16"); psoGroupnorm=pso("groupnorm_f16"); + psoBiasAddRows=pso("bias_add_rows_f16"); psoEwChain=pso("elemwise_chain_f16"); + psoAttnFlash=pso("attn_flash_f16"); psoRmsSimd=pso("rms_norm_simd_f16"); + psoScatterRows=pso("strided_scatter_rows_f16"); psoTranspose2d=pso("transpose2d_f16"); + psoBcastRows=pso("bcast_binary_rows_f16"); psoRmsAddSimd=pso("rms_norm_add_simd_f16"); + psoConvCacheAppend=pso("conv_cache_append_f16"); psoRelPosBias=pso("rel_pos_bias_f16"); + psoGemvBias=pso("gemv_bias_f16"); + psoTopkRows=pso("topk_rows_f16"); psoMoeT=pso("cq4_moe_transform"); + psoMoeUp=pso("cq4_moe_gemv_up"); psoMoeT2=pso("cq4_moe_transform2"); + psoMoeDownAcc=pso("cq4_moe_gemv_down_acc"); psoRopePair=pso("rope_pair_f16"); + psoRmsScale=pso("rms_norm_scale_f16"); psoSoftmaxTopk=pso("softmax_topk_f16"); + psoRopePairRms=pso("rope_pair_rms_f16"); psoRms2AddClip=pso("rms2_add_clip_f16"); + psoDeltanet=pso("gated_deltanet_decode_f16"); psoDeltanetPre=pso("gated_deltanet_prefill_f16"); + dummy=[dev newBufferWithLength:16 options:MTLResourceStorageModeShared]; + ok = psoT&&psoG&&psoTm&&psoGmma&&psoRotW&&psoEmbH&&psoEmbOm&&psoEmbHm&&psoCopy&&psoBinary&&psoScalar&&psoUnary&&psoRms&&psoSwiglu&&psoRmsAdd&&psoCF16F32&&psoCF32F16&&psoCI8F16&&psoCF16I8 + &&psoAttn&&psoAttnC&&psoAttnPre&&psoAttnPreMma2&&psoKvAppendM&&psoKvAppendRingM&&psoSlideS&&psoSlideR&&psoSlideRM&&psoStrided&&psoBcast&&psoScatter&&psoKvAppend&&psoArgmax&&psoGather; + }} +}; + +MetalCtx& ctx() { static MetalCtx c; return c; } + +id owned_shared(size_t len) { + size_t alen = ((len ? len : 1) + 16383) & ~(size_t)16383; + void* p = mmap(nullptr, alen, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); + if (p == MAP_FAILED) + return [ctx().dev newBufferWithLength:alen options:MTLResourceStorageModeShared]; + id b = [ctx().dev newBufferWithBytesNoCopy:p length:alen + options:MTLResourceStorageModeShared + deallocator:^(void* ptr, NSUInteger l){ munmap(ptr, l); }]; + if (!b) { + munmap(p, alen); + return [ctx().dev newBufferWithLength:alen options:MTLResourceStorageModeShared]; + } + return b; +} + +inline id buf(const void* p, size_t bytes) { + if (!p || bytes == 0) return nil; + id b = owned_shared(bytes); + if (b) std::memcpy([b contents], p, bytes); + return b; +} + +struct ResW { + id packed=nil, norms=nil, codebook=nil, lsign=nil, rsign=nil, perm=nil, recip=nil; + id rotation=nil; + id cb_i8=nil; + float cb_scale=1.f; + size_t packed_off=0, norms_off=0; + size_t cb_off=0, ls_off=0, rs_off=0, pm_off=0, rc_off=0, rot_off=0; + uint32_t il=0; + bool ok=false; +}; +std::unordered_map g_resident; + +static uint64_t resident_key(const CactusQuantMatrix* W) { + uint64_t h = 1469598103934665603ull; + auto mix = [&](uint64_t v) { h ^= v; h *= 1099511628211ull; }; + mix(W->bits); mix(W->K); mix(W->N); mix(W->group_size); mix(W->num_groups); mix(W->flags); + const uint8_t* p = (const uint8_t*)W->packed_indices; + size_t pkb = (size_t)W->N * W->num_groups * ((W->group_size * W->bits + 7u) / 8u); + if (!p || pkb == 0) return h; + size_t take = pkb < 64 ? pkb : 64; + for (size_t i = 0; i < take; ++i) mix(p[i]); + for (size_t i = pkb - take; i < pkb; ++i) mix(p[i]); + return h; +} +id g_attn_scores = nil; +std::unordered_map g_resident_emb; +std::mutex g_resident_mu; + +struct MoESet4 { + id pk=nil, nm=nil, rc=nil, ls=nil, rs=nil, pm=nil, cb=nil; + size_t pk_off=0, nm_off=0, rc_off=0, ls_off=0, rs_off=0, pm_off=0, cb_off=0; + uint32_t K=0, N=0; + uint32_t stride=0; +}; +struct MoECat4 { MoESet4 w1, w3, w2; bool ok=false; }; +std::unordered_map g_moe_cat; + +std::pair, size_t> wrapHostPtr(const void* p, size_t bytes); + +bool bind_component(const void* p, size_t bytes, size_t align, + id __strong& b, size_t& off) { + if (!p || bytes == 0) return true; + if (((uintptr_t)p & (align-1)) != 0) return false; + auto w = wrapHostPtr(p, bytes); + if (!w.first) return false; + b = w.first; + off = w.second; + return true; +} + +ResW& resident(const CactusQuantMatrix* W) { + std::lock_guard lk(g_resident_mu); + uint64_t key = resident_key(W); + auto it = g_resident.find(key); + if (it != g_resident.end()) return it->second; + const uint32_t K=W->K, N=W->N, ng=W->num_groups, gs=W->group_size, bits=W->bits; + const uint32_t pgb=(gs*bits+7u)/8u; + ResW r; + r.ok = bind_component(W->codebook, (size_t)(1u<left_signs, gs, 1, r.lsign, r.ls_off) + && bind_component(W->right_signs, gs, 1, r.rsign, r.rs_off) + && bind_component(W->permutation, (size_t)gs*sizeof(uint32_t), 4, r.perm, r.pm_off) + && bind_component(W->input_scale_recip, (size_t)K*sizeof(__fp16), 2, r.recip, r.rc_off) + && (!W->rotation || bind_component(W->rotation, (size_t)K*K*sizeof(__fp16), 8, r.rotation, r.rot_off)); + if (bits == 2u) { + int8_t cbq[4]; + float max_abs = 0.f; + for (uint32_t i = 0; i < 4u; i++) { + float v = std::fabs((float)W->codebook[i]); + if (v > max_abs) max_abs = v; + } + float scale = max_abs / 127.f; + if (scale < 1e-10f) scale = 1e-10f; + float inv = 1.f / scale; + for (uint32_t i = 0; i < 4u; i++) + cbq[i] = (int8_t)std::round((float)W->codebook[i] * inv); + r.cb_i8 = buf(cbq, 4); + r.cb_scale = scale; + } + size_t pkb=(size_t)N*ng*pgb, nmb=(size_t)N*ng*sizeof(__fp16); + const bool interleaved = (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0; + if (((uintptr_t)W->packed_indices & 7u) != 0 || ((uintptr_t)W->norms & 1u) != 0) { + r.ok = false; + } else if (r.ok) { + auto wp = wrapHostPtr(W->packed_indices, pkb); + auto wn = wrapHostPtr(W->norms, nmb); + if (wp.first && wn.first) { + r.packed = wp.first; r.packed_off = wp.second; + r.norms = wn.first; r.norms_off = wn.second; + r.il = interleaved ? 1 : 0; + } else { + r.ok = false; + } + } + return g_resident.emplace(key, r).first->second; +} + +id g_cmd = nil; +id g_enc = nil; +std::unordered_map>> g_free; +std::vector> g_pending; +std::map> g_shared; +bool g_active = false; +id g_code_buf = nil; +id g_code_buf_m = nil; + +inline size_t bucket(size_t b) { return (b + 16383) & ~size_t(16383); } + +static const uint32_t g_mr_rows = 8u * 2u; + +inline bool tg_mem_ok(size_t bytes) { + return bytes <= (size_t)[ctx().dev maxThreadgroupMemoryLength]; +} + +id ensureEncoder() { + if (!g_enc) { + if (!g_cmd) g_cmd = [ctx().queue commandBuffer]; + g_enc = [g_cmd computeCommandEncoder]; + } + return g_enc; +} + +id recycled(size_t bytes) { + size_t bk = bucket(bytes ? bytes : 1); + auto it = g_free.find(bk); + id b; + if (it != g_free.end() && !it->second.empty()) { b = it->second.back(); it->second.pop_back(); } + else b = owned_shared(bk); + g_pending.push_back(b); + return b; +} + +std::map> g_wrapped; +std::unordered_map> g_readonly; +std::map, MPSMatrix*> g_mps_mats; +std::pair, size_t> wrapHostPtr(const void* p, size_t bytes) { + uintptr_t a = (uintptr_t)p, base = a & ~(uintptr_t)16383u; + size_t need = (a - base) + (bytes ? bytes : 1); + auto it = g_wrapped.upper_bound(a); + if (it != g_wrapped.begin()) { + --it; + if (a >= it->first && a + (bytes?bytes:1) <= it->first + (size_t)it->second.length) + return { it->second, (size_t)(a - it->first) }; + } + size_t wraplen = (need + 16383u) & ~(size_t)16383u; + uintptr_t lo = base, hi = base + wraplen; +#if TARGET_OS_OSX + // Expand the wrap to the whole VM region so neighboring host pointers hit + // the same MTLBuffer. mach_vm.h is macOS-only; elsewhere the page-aligned + // span above is used as-is (correct, just more wrap entries). + { + mach_vm_address_t raddr = a; + mach_vm_size_t rsize = 0; + vm_region_basic_info_data_64_t info; + mach_msg_type_number_t cnt = VM_REGION_BASIC_INFO_COUNT_64; + mach_port_t obj = MACH_PORT_NULL; + if (mach_vm_region(mach_task_self(), &raddr, &rsize, VM_REGION_BASIC_INFO_64, + (vm_region_info_t)&info, &cnt, &obj) == KERN_SUCCESS + && raddr <= base && raddr + rsize >= hi + && (info.protection & VM_PROT_READ) && (info.protection & VM_PROT_WRITE) + && rsize <= (512u << 20)) { + lo = (uintptr_t)raddr; + hi = (uintptr_t)(raddr + rsize); + } + if (obj != MACH_PORT_NULL) mach_port_deallocate(mach_task_self(), obj); + } +#endif + auto ov = g_wrapped.lower_bound(lo); + if (ov != g_wrapped.begin()) { + auto prev = std::prev(ov); + if (prev->first + (size_t)prev->second.length > lo) ov = prev; + } + while (ov != g_wrapped.end() && ov->first < hi) { + if (ov->first + (size_t)ov->second.length > lo) { + if (ov->first < lo) lo = ov->first; + if (ov->first + (size_t)ov->second.length > hi) hi = ov->first + (size_t)ov->second.length; + ov = g_wrapped.erase(ov); + } else ++ov; + } + id b = [ctx().dev newBufferWithBytesNoCopy:(void*)lo length:(size_t)(hi - lo) + options:MTLResourceStorageModeShared deallocator:nil]; + if (!b) return { nil, 0 }; + g_wrapped[lo] = b; + return { b, (size_t)(a - lo) }; +} + +std::pair, size_t> bufForPtrOff(const void* p, size_t bytes) { + uintptr_t a = reinterpret_cast(p); + auto it = g_shared.upper_bound(a); + if (it != g_shared.begin()) { + --it; + uintptr_t base = it->first; + if (a < base + static_cast(it->second.length)) + return { it->second, static_cast(a - base) }; + } + + auto wr = wrapHostPtr(p, bytes); + if (wr.first) return wr; + + size_t nb = bytes ? bytes : 1; + id b = recycled(bytes); + std::memcpy([b contents], p, nb); + return { b, 0 }; +} + +inline void setBufAt(const void* p, size_t bytes, int idx) { + auto pr = bufForPtrOff(p, bytes); + [g_enc setBuffer:pr.first offset:pr.second atIndex:idx]; +} + +} + +bool cactus_metal_available() { return ctx().ok; } + +void cactus_metal_set_active(bool a) { g_active = a; } +bool cactus_metal_active_mode() { return ctx().ok && g_active; } + +extern "C" void* objc_autoreleasePoolPush(void); +extern "C" void objc_autoreleasePoolPop(void*); +static void* g_arpool = nullptr; + +void cactus_metal_session_begin() { + if (!g_arpool) g_arpool = objc_autoreleasePoolPush(); +} +static id g_last_cmd = nil; +static std::atomic g_cmd_failed{false}; +static void watch_cmd_errors(id cmd) { + [cmd addCompletedHandler:^(id cb) { + if (cb.status == MTLCommandBufferStatusError) { + g_cmd_failed.store(true, std::memory_order_relaxed); + fprintf(stderr, "[cactus-metal] command buffer failed: %s\n", + cb.error ? [[cb.error localizedDescription] UTF8String] : "unknown"); + } + }]; +} +static void reclaim_dead_slabs(); +void cactus_metal_session_sync() { + if (g_enc) { [g_enc endEncoding]; g_enc = nil; } + if (g_cmd) { + watch_cmd_errors(g_cmd); + [g_cmd commit]; + g_last_cmd = g_cmd; + g_cmd = nil; + } + if (g_last_cmd) { + [g_last_cmd waitUntilCompleted]; + g_last_cmd = nil; + } + for (id b : g_pending) g_free[(size_t)b.length].push_back(b); + g_pending.clear(); + reclaim_dead_slabs(); +} +void cactus_metal_session_end() { + cactus_metal_session_sync(); + if (g_arpool) { objc_autoreleasePoolPop(g_arpool); g_arpool = nullptr; } +} +void cactus_metal_invalidate_host_wraps() { + std::lock_guard lk(g_resident_mu); + g_attn_scores = nil; + g_wrapped.clear(); + g_readonly.clear(); + g_resident.clear(); + g_resident_emb.clear(); + g_mps_mats.clear(); + g_moe_cat.clear(); +} + +void cactus_metal_session_flush() { + if (g_enc) { [g_enc endEncoding]; g_enc = nil; } + if (g_cmd) { watch_cmd_errors(g_cmd); [g_cmd commit]; g_last_cmd = g_cmd; g_cmd = nil; } +} + + +void* cactus_metal_alloc_shared(size_t bytes) { + if (!ctx().ok) return nullptr; + size_t bk = bucket(bytes ? bytes : 1); + id b = nil; + auto it = g_free.find(bk); + if (it != g_free.end() && !it->second.empty()) { b = it->second.back(); it->second.pop_back(); } + else b = owned_shared(bk); + void* c = [b contents]; + g_shared[reinterpret_cast(c)] = b; + return c; +} + +struct SlabInfo { size_t cap = 0, used = 0, live = 0; }; +static char* g_slab = nullptr; +static std::map g_slabs; + +static void reclaim_dead_slabs() { + for (auto it = g_slabs.begin(); it != g_slabs.end();) { + SlabInfo& s = it->second; + if (s.live != 0 || s.used == 0) { ++it; continue; } + if ((char*)it->first == g_slab) { s.used = 0; ++it; continue; } + g_shared.erase(it->first); + it = g_slabs.erase(it); + } +} + +void* cactus_metal_alloc_pooled(size_t bytes) { + if (!ctx().ok) return nullptr; + if (bytes > (4u << 20)) return cactus_metal_alloc_shared(bytes); + size_t need = (bytes + 255) & ~size_t(255); + SlabInfo* cur = nullptr; + if (g_slab) { + auto it = g_slabs.find(reinterpret_cast(g_slab)); + if (it != g_slabs.end()) cur = &it->second; + } + if (!cur || cur->used + need > cur->cap) { + size_t cap = 32u << 20; + id b = owned_shared(cap); + if (!b) return cactus_metal_alloc_shared(bytes); + g_slab = (char*)[b contents]; + g_shared[reinterpret_cast(g_slab)] = b; + cur = &g_slabs[reinterpret_cast(g_slab)]; + *cur = SlabInfo{cap, 0, 0}; + } + void* p = g_slab + cur->used; + cur->used += need; + cur->live += 1; + return p; +} +void cactus_metal_free_shared(void* contents) { + uintptr_t a = reinterpret_cast(contents); + auto sit = g_slabs.upper_bound(a); + if (sit != g_slabs.begin()) { + --sit; + if (a < sit->first + sit->second.cap) { + if (sit->second.live > 0) --sit->second.live; + return; + } + } + auto it = g_shared.find(a); + if (it != g_shared.end()) { g_pending.push_back(it->second); g_shared.erase(it); } +} + +bool cactus_metal_encode_copy(void* out, const void* in, size_t bytes) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t n=(uint32_t)bytes; + [g_enc setComputePipelineState:ctx().psoCopy]; + setBufAt(in, bytes, 0); + setBufAt(out, bytes, 1); + [g_enc setBytes:&n length:4 atIndex:2]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_binary(int op, void* out, const void* a, const void* b, size_t n) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; int o=op; + [g_enc setComputePipelineState:ctx().psoBinary]; + setBufAt(a, n*2, 0); setBufAt(b, n*2, 1); setBufAt(out, n*2, 2); + [g_enc setBytes:&nn length:4 atIndex:3]; [g_enc setBytes:&o length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake((nn+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_scalar(int op, void* out, const void* in, size_t n, float param) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; int o=op; float p=param; + [g_enc setComputePipelineState:ctx().psoScalar]; + setBufAt(in, n*2, 0); setBufAt(out, n*2, 1); + [g_enc setBytes:&nn length:4 atIndex:2]; [g_enc setBytes:&o length:4 atIndex:3]; [g_enc setBytes:&p length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake((nn+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_unary(int op, void* out, const void* in, size_t n) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; int o=op; + [g_enc setComputePipelineState:ctx().psoUnary]; + setBufAt(in, n*2, 0); setBufAt(out, n*2, 1); + [g_enc setBytes:&nn length:4 atIndex:2]; [g_enc setBytes:&o length:4 atIndex:3]; + [g_enc dispatchThreads:MTLSizeMake((nn+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_swiglu(void* out, const void* gate, const void* up, size_t n, float scale) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; float s=scale; + [g_enc setComputePipelineState:ctx().psoSwiglu]; + setBufAt(gate, n*2, 0); setBufAt(up, n*2, 1); setBufAt(out, n*2, 2); + [g_enc setBytes:&nn length:4 atIndex:3]; [g_enc setBytes:&s length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake(nn,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_rms_norm(void* out, const void* in, const void* weight, + size_t rows, size_t dim, float eps) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t d=(uint32_t)dim; float e=eps; + if (rows > 1 && ctx().psoRmsSimd) { + uint32_t r=(uint32_t)rows; + [g_enc setComputePipelineState:ctx().psoRmsSimd]; + setBufAt(in, rows*dim*2, 0); setBufAt(weight, dim*2, 1); setBufAt(out, rows*dim*2, 2); + [g_enc setBytes:&d length:4 atIndex:3]; [g_enc setBytes:&e length:4 atIndex:4]; + [g_enc setBytes:&r length:4 atIndex:5]; + [g_enc dispatchThreadgroups:MTLSizeMake((rows+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; + } + [g_enc setComputePipelineState:ctx().psoRms]; + setBufAt(in, rows*dim*2, 0); setBufAt(weight, dim*2, 1); setBufAt(out, rows*dim*2, 2); + [g_enc setBytes:&d length:4 atIndex:3]; [g_enc setBytes:&e length:4 atIndex:4]; + [g_enc setThreadgroupMemoryLength:256*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_rms_norm_add(void* out, const void* in, const void* weight, const void* res, + size_t rows, size_t dim, float eps) { + if (!ctx().ok) return false; + ensureEncoder(); + uint32_t d=(uint32_t)dim; float e=eps; + [g_enc setComputePipelineState:ctx().psoRmsAdd]; + setBufAt(in, rows*dim*2, 0); setBufAt(weight, dim*2, 1); setBufAt(res, rows*dim*2, 2); setBufAt(out, rows*dim*2, 3); + [g_enc setBytes:&d length:4 atIndex:4]; [g_enc setBytes:&e length:4 atIndex:5]; + [g_enc setThreadgroupMemoryLength:256*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_rms_norm_add_scale(void* out, const void* in, const void* weight, const void* res, + size_t rows, size_t dim, float eps, float out_scale) { + if (!ctx().ok || !ctx().psoRmsAddScale) return false; + ensureEncoder(); + uint32_t d=(uint32_t)dim; float e=eps, os=out_scale; + [g_enc setComputePipelineState:ctx().psoRmsAddScale]; + setBufAt(in, rows*dim*2, 0); setBufAt(weight, dim*2, 1); setBufAt(res, rows*dim*2, 2); setBufAt(out, rows*dim*2, 3); + [g_enc setBytes:&d length:4 atIndex:4]; [g_enc setBytes:&e length:4 atIndex:5]; [g_enc setBytes:&os length:4 atIndex:6]; + [g_enc setThreadgroupMemoryLength:256*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_rms_norm_add_rms(void* h_out, void* xn_out, const void* in, const void* w1, + const void* res, const void* w2, + size_t rows, size_t dim, float eps, float out_scale) { + if (!ctx().ok || !ctx().psoRmsAddRms) return false; + ensureEncoder(); + uint32_t d=(uint32_t)dim; float e=eps, os=out_scale; + [g_enc setComputePipelineState:ctx().psoRmsAddRms]; + setBufAt(in, rows*dim*2, 0); setBufAt(w1, dim*2, 1); setBufAt(res, rows*dim*2, 2); + setBufAt(h_out, rows*dim*2, 3); setBufAt(w2, dim*2, 4); setBufAt(xn_out, rows*dim*2, 5); + [g_enc setBytes:&d length:4 atIndex:6]; [g_enc setBytes:&e length:4 atIndex:7]; [g_enc setBytes:&os length:4 atIndex:8]; + [g_enc setThreadgroupMemoryLength:256*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_argmax(const void* logits, uint32_t vocab, void* out3, const void* bias) { + if (!ctx().ok) return false; + if (bias && !(ctx().psoArgmaxP && ctx().psoArgmaxF && vocab >= 32768u)) return false; + ensureEncoder(); + uint32_t V = vocab; + if (ctx().psoArgmaxP && ctx().psoArgmaxF && V >= 32768u) { + const uint32_t NP = 128u, T = 256u; + uint32_t chunk = (V + NP - 1u)/NP; + uint32_t has_bias = bias ? 1u : 0u; + id part = recycled((size_t)NP*3*sizeof(float)); + [g_enc setComputePipelineState:ctx().psoArgmaxP]; + setBufAt(logits, (size_t)vocab*2, 0); + [g_enc setBuffer:part offset:0 atIndex:1]; + [g_enc setBytes:&V length:4 atIndex:2]; [g_enc setBytes:&chunk length:4 atIndex:3]; + if (bias) setBufAt(bias, (size_t)vocab*sizeof(float), 4); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:4]; + [g_enc setBytes:&has_bias length:4 atIndex:5]; + [g_enc setThreadgroupMemoryLength:T*sizeof(float) atIndex:0]; + [g_enc setThreadgroupMemoryLength:T*sizeof(uint) atIndex:1]; + [g_enc setThreadgroupMemoryLength:T*sizeof(float) atIndex:2]; + [g_enc dispatchThreadgroups:MTLSizeMake(NP,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + [g_enc setComputePipelineState:ctx().psoArgmaxF]; + [g_enc setBuffer:part offset:0 atIndex:0]; + setBufAt(out3, 3*sizeof(float), 1); + [g_enc setBytes:&NP length:4 atIndex:2]; + [g_enc setThreadgroupMemoryLength:NP*sizeof(float) atIndex:0]; + [g_enc setThreadgroupMemoryLength:NP*sizeof(uint) atIndex:1]; + [g_enc setThreadgroupMemoryLength:NP*sizeof(float) atIndex:2]; + [g_enc dispatchThreadgroups:MTLSizeMake(1,1,1) threadsPerThreadgroup:MTLSizeMake(NP,1,1)]; + return true; + } + uint32_t T = 1024; + [g_enc setComputePipelineState:ctx().psoArgmax]; + setBufAt(logits, (size_t)vocab*2, 0); + setBufAt(out3, 3*sizeof(float), 1); + [g_enc setBytes:&V length:4 atIndex:2]; + [g_enc setThreadgroupMemoryLength:T*sizeof(float) atIndex:0]; + [g_enc setThreadgroupMemoryLength:T*sizeof(uint) atIndex:1]; + [g_enc setThreadgroupMemoryLength:T*sizeof(float) atIndex:2]; + [g_enc dispatchThreadgroups:MTLSizeMake(1,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + return true; +} +bool cactus_metal_encode_adjust_logits(void* logits, size_t vocab, + const uint32_t* recent, uint32_t n_recent, + int64_t suppressed, float penalty) { + if (!ctx().ok || !ctx().psoAdjust) return false; + if (n_recent > 4096u) return false; + ensureEncoder(); + id rb = ctx().dummy; + if (n_recent) { rb = recycled((size_t)n_recent*4); std::memcpy([rb contents], recent, (size_t)n_recent*4); } + struct { uint32_t n_recent, suppress_flag, suppress_id, vocab_n; float penalty; } U = + { n_recent, + (uint32_t)(suppressed >= 0 ? 1 : 0), (uint32_t)(suppressed >= 0 ? suppressed : 0), + (uint32_t)vocab, penalty }; + [g_enc setComputePipelineState:ctx().psoAdjust]; + setBufAt(logits, vocab*2, 0); + [g_enc setBuffer:rb offset:0 atIndex:1]; + [g_enc setBytes:&U length:sizeof(U) atIndex:2]; + [g_enc dispatchThreadgroups:MTLSizeMake(1,1,1) threadsPerThreadgroup:MTLSizeMake(1024,1,1)]; + return true; +} + +bool cactus_metal_encode_softcap(void* out, const void* in, size_t n, float cap) { + if (!ctx().ok || !ctx().psoSoftcap) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; float c=cap; + [g_enc setComputePipelineState:ctx().psoSoftcap]; + setBufAt(in, n*2, 0); setBufAt(out, n*2, 1); + [g_enc setBytes:&nn length:4 atIndex:2]; [g_enc setBytes:&c length:4 atIndex:3]; + [g_enc dispatchThreads:MTLSizeMake(nn,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_cast(void* out, int out_prec, const void* in, int in_prec, size_t n) { + if (!ctx().ok) return false; + + auto esz=[](int p){ return p==0?1: p==1?2: p==2?4: 0; }; + if (in_prec==out_prec) return cactus_metal_encode_copy(out, in, n*esz(in_prec)); + id pso=nil; + if (in_prec==1&&out_prec==2) pso=ctx().psoCF16F32; + else if (in_prec==2&&out_prec==1) pso=ctx().psoCF32F16; + else if (in_prec==0&&out_prec==1) pso=ctx().psoCI8F16; + else if (in_prec==1&&out_prec==0) pso=ctx().psoCF16I8; + else return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; + [g_enc setComputePipelineState:pso]; + setBufAt(in, n*esz(in_prec), 0); setBufAt(out, n*esz(out_prec), 1); + [g_enc setBytes:&nn length:4 atIndex:2]; + [g_enc dispatchThreads:MTLSizeMake(nn,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +static inline bool quant_fast_eligible(const CactusQuantMatrix* W) { + return W->bits == 4 && W->group_size >= 128 && (W->group_size % 128) == 0 && (W->N % 4) == 0 && + (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) && + !(W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL) && W->input_scale_recip && W->left_signs && + W->right_signs && W->permutation && W->codebook && W->norms && W->packed_indices; +} + +static inline bool quant_lowbit_eligible(const CactusQuantMatrix* W) { + return (W->bits == 2 || W->bits == 3) && W->group_size == 128 && (W->N % 16) == 0 && + (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) && + !(W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL) && + W->input_scale_recip && W->left_signs && W->right_signs && W->permutation && + W->codebook && W->norms && W->packed_indices; +} + +void cactus_metal_trim_prefill_cache() { + std::lock_guard lk(g_resident_mu); + g_mps_mats.clear(); + g_free.clear(); + g_attn_scores = nil; + g_code_buf = nil; + g_code_buf_m = nil; +} + +bool cactus_metal_encode_quant_matmul(void* out, const void* lhs, const CactusQuantMatrix* W) { + if (!ctx().ok || !W) return false; + const uint32_t gs=W->group_size, K=W->K, N=W->N, ng=W->num_groups; + if ((W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL) && W->rotation && ng==1 && gs==K && (N%4)==0) { + static void* ortho_code = nullptr; static uint32_t ortho_code_k = 0; + if (ortho_code_k < K) { + if (ortho_code) cactus_metal_free_shared(ortho_code); + ortho_code = cactus_metal_alloc_shared((size_t)K*sizeof(__fp16)); + ortho_code_k = K; + } + if (ortho_code) return cactus_metal_encode_quant_matmul_ortho(out, lhs, ortho_code, W); + } + const bool fast = quant_fast_eligible(W); + const bool lowbit = !fast && ctx().psoGmrLow && quant_lowbit_eligible(W); + if (!fast && !lowbit) return false; + const uint32_t bits = W->bits, pgbw = (gs*bits+7u)/8u; + ResW& rw = resident(W); + if (!rw.ok) return false; + + size_t code_bytes = (size_t)ng*gs*sizeof(__fp16); + if (!g_code_buf || (size_t)g_code_buf.length < code_bytes) + g_code_buf = owned_shared(code_bytes); + if (!g_code_buf) return false; + ensureEncoder(); + bool simdT = (gs==128u && ctx().psoTsimd); + [g_enc setComputePipelineState:(simdT?ctx().psoTsimd:ctx().psoT)]; + setBufAt(lhs, (size_t)K*2, 0); [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:1]; + [g_enc setBuffer:rw.lsign offset:rw.ls_off atIndex:2]; [g_enc setBuffer:rw.rsign offset:rw.rs_off atIndex:3]; + [g_enc setBuffer:rw.perm offset:rw.pm_off atIndex:4]; [g_enc setBuffer:g_code_buf offset:0 atIndex:5]; + [g_enc setBytes:&gs length:4 atIndex:6]; [g_enc setThreadgroupMemoryLength:gs*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(ng,1,1) threadsPerThreadgroup:MTLSizeMake(simdT?32u:(gs>1024u?1024u:gs),1,1)]; + if (lowbit && bits == 2u && gs == 128u && ctx().psoActQ && ctx().psoG2i8 && rw.cb_i8) { + id actq = recycled((size_t)ng*gs); + id ascl = recycled((size_t)ng*sizeof(float)); + [g_enc setComputePipelineState:ctx().psoActQ]; + [g_enc setBuffer:g_code_buf offset:0 atIndex:0]; + [g_enc setBuffer:actq offset:0 atIndex:1]; [g_enc setBuffer:ascl offset:0 atIndex:2]; + [g_enc setBytes:&gs length:4 atIndex:3]; + [g_enc dispatchThreadgroups:MTLSizeMake(ng,1,1) threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + [g_enc setComputePipelineState:ctx().psoG2i8]; + [g_enc setBuffer:actq offset:0 atIndex:0]; [g_enc setBuffer:ascl offset:0 atIndex:1]; + [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:2]; [g_enc setBuffer:rw.cb_i8 offset:0 atIndex:3]; + [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:4]; + setBufAt(out, (size_t)N*2, 5); + [g_enc setBytes:&ng length:4 atIndex:6]; [g_enc setBytes:&N length:4 atIndex:7]; + [g_enc setBytes:&rw.cb_scale length:4 atIndex:8]; + [g_enc dispatchThreadgroups:MTLSizeMake((N+7u)/8u,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; + } + bool umr = ctx().psoGmr && (N % g_mr_rows == 0u); + [g_enc setComputePipelineState:(lowbit?ctx().psoGmrLow:(umr?ctx().psoGmr:ctx().psoG))]; + [g_enc setBuffer:g_code_buf offset:0 atIndex:0]; [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:1]; + [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:2]; [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:3]; + setBufAt(out, (size_t)N*2, 4); + [g_enc setBytes:&gs length:4 atIndex:5]; [g_enc setBytes:&ng length:4 atIndex:6]; + [g_enc setBytes:&pgbw length:4 atIndex:7]; [g_enc setBytes:&N length:4 atIndex:8]; + if (lowbit) { [g_enc setBytes:&bits length:4 atIndex:9]; [g_enc setBytes:&rw.il length:4 atIndex:10]; } + else [g_enc setBytes:&rw.il length:4 atIndex:9]; + uint32_t ROWS=8; + uint32_t grid = (lowbit || umr) ? (N+g_mr_rows-1u)/g_mr_rows : (N+ROWS-1u)/ROWS; + [g_enc dispatchThreadgroups:MTLSizeMake(grid,1,1) threadsPerThreadgroup:MTLSizeMake(ROWS*32,1,1)]; + return true; +} + +bool cactus_metal_encode_transform_batch(const void* x, const CactusQuantMatrix* const* Ws, int B, void* const* codes) { + if (!ctx().ok || B < 1 || B > 3 || !ctx().psoTbatch) return false; + const uint32_t gs=Ws[0]->group_size, K=Ws[0]->K, ng=Ws[0]->num_groups; + if (gs != 128u || !quant_fast_eligible(Ws[0])) return false; + ResW& r0 = resident(Ws[0]); + ResW* rw[3] = { &r0, &r0, &r0 }; + for (int b=1;bok) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoTbatch]; + setBufAt(x, (size_t)K*2, 0); + [g_enc setBuffer:r0.recip offset:r0.rc_off atIndex:1]; + for (int b=0;b<3;b++) { + ResW* R = rw[blsign offset:R->ls_off atIndex:2+b]; + [g_enc setBuffer:R->rsign offset:R->rs_off atIndex:5+b]; + [g_enc setBuffer:R->perm offset:R->pm_off atIndex:8+b]; + setBufAt(codes[bgroup_size, K=W->K, N=W->N, ng=W->num_groups, pgb=(gs*4u+7u)/8u; + if (!quant_fast_eligible(W)) return false; + ResW& rw = resident(W); + if (!rw.ok) return false; + ensureEncoder(); + bool umr = ctx().psoGmr && (N % g_mr_rows == 0u); + [g_enc setComputePipelineState:(umr?ctx().psoGmr:ctx().psoG)]; + setBufAt(code, (size_t)K*2, 0); [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:1]; + [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:2]; [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:3]; + setBufAt(out, (size_t)N*2, 4); + [g_enc setBytes:&gs length:4 atIndex:5]; [g_enc setBytes:&ng length:4 atIndex:6]; + [g_enc setBytes:&pgb length:4 atIndex:7]; [g_enc setBytes:&N length:4 atIndex:8]; + [g_enc setBytes:&rw.il length:4 atIndex:9]; + uint32_t ROWS=8; + uint32_t grid = umr ? (N+g_mr_rows-1u)/g_mr_rows : (N+ROWS-1u)/ROWS; + [g_enc dispatchThreadgroups:MTLSizeMake(grid,1,1) threadsPerThreadgroup:MTLSizeMake(ROWS*32,1,1)]; + return true; +} + +bool cactus_metal_transform_gemv_fits(uint32_t K) { + return ctx().ok && tg_mem_ok((size_t)K*2 + 8*128*sizeof(float) + 64); +} + +bool cactus_metal_encode_transform_gemv(void* out, const void* x, const CactusQuantMatrix* W, const void* osw) { + if (!ctx().ok || !ctx().psoMega) return false; + const uint32_t gs=W->group_size, K=W->K, ng=W->num_groups, N=W->N; + if (gs != 128u || !quant_fast_eligible(W) || (N % 16u) != 0u) return false; + if (!cactus_metal_transform_gemv_fits(K)) return false; + ResW& rw = resident(W); + if (!rw.ok) return false; + ensureEncoder(); + struct { uint32_t K, ng, oswi; } U = { K, ng, (uint32_t)(osw ? 1 : 0) }; + [g_enc setComputePipelineState:ctx().psoMega]; + setBufAt(x, (size_t)K*2, 0); + [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:1]; + [g_enc setBuffer:rw.lsign offset:rw.ls_off atIndex:2]; [g_enc setBuffer:rw.rsign offset:rw.rs_off atIndex:3]; + [g_enc setBuffer:rw.perm offset:rw.pm_off atIndex:4]; [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:5]; + [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:6]; [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:7]; + setBufAt(out, (size_t)N*2, 8); + if (osw) setBufAt(osw, (size_t)N*2, 9); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:9]; + [g_enc setBytes:&U length:sizeof(U) atIndex:10]; + [g_enc setThreadgroupMemoryLength:(size_t)K*2 atIndex:0]; + [g_enc setThreadgroupMemoryLength:8*128*sizeof(float) atIndex:1]; + [g_enc dispatchThreadgroups:MTLSizeMake(N/16u,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_swiglu_transform(void* code, const void* gate, const void* up, + const CactusQuantMatrix* W, float scale) { + if (!ctx().ok || !ctx().psoSwiT || W->group_size != 128u || !quant_fast_eligible(W)) return false; + const uint32_t K=W->K, ng=W->num_groups; + ResW& rw = resident(W); + if (!rw.ok) return false; + ensureEncoder(); + float s = scale; + [g_enc setComputePipelineState:ctx().psoSwiT]; + setBufAt(gate, (size_t)K*2, 0); setBufAt(up, (size_t)K*2, 1); + [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:2]; + [g_enc setBuffer:rw.lsign offset:rw.ls_off atIndex:3]; [g_enc setBuffer:rw.rsign offset:rw.rs_off atIndex:4]; + [g_enc setBuffer:rw.perm offset:rw.pm_off atIndex:5]; + setBufAt(code, (size_t)K*2, 6); + [g_enc setBytes:&s length:4 atIndex:7]; + [g_enc setThreadgroupMemoryLength:128*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(ng,1,1) threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + return true; +} + +bool cactus_metal_encode_gemv_cat(void* const* outs, const void* const* codes, + const CactusQuantMatrix* const* Ws, int B) { + if (!ctx().ok || !ctx().psoCat || B < 1 || B > 3) return false; + const uint32_t gs=Ws[0]->group_size, K=Ws[0]->K, ng=Ws[0]->num_groups; + if (gs != 128u) return false; + uint32_t Ns[3]={0,0,0}; + for (int i=0;iK != K || (Ws[i]->N % 16u) != 0u) return false; + Ns[i]=Ws[i]->N; + } + ResW* rw[3]; rw[0]=&resident(Ws[0]); rw[1]=rw[0]; rw[2]=rw[0]; + for (int i=1;iok) return false; + ensureEncoder(); + struct { uint32_t ng, N0, N1, N2; } U = { ng, Ns[0], Ns[1], Ns[2] }; + [g_enc setComputePipelineState:ctx().psoCat]; + for (int i=0;i<3;i++){ + int j = ipacked offset:rw[j]->packed_off atIndex:3+i]; + [g_enc setBuffer:rw[j]->codebook offset:rw[j]->cb_off atIndex:6+i]; + [g_enc setBuffer:rw[j]->norms offset:rw[j]->norms_off atIndex:9+i]; + setBufAt(outs[j], (size_t)Ns[j]*2, 12+i); + } + [g_enc setBytes:&U length:sizeof(U) atIndex:15]; + uint32_t total = Ns[0]+Ns[1]+Ns[2]; + [g_enc dispatchThreadgroups:MTLSizeMake(total/16u,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_prewarm_quant(const CactusQuantMatrix* W) { + if (!ctx().ok) return false; + if (!quant_fast_eligible(W) && !quant_lowbit_eligible(W)) return false; + return resident(W).ok; +} + +namespace { +bool moe_component_offsets(const CactusQuantMatrix& W, uintptr_t& lo, uintptr_t& hi, size_t offs[7]) { + const uintptr_t pg = 16384; + struct Span { const void* p; size_t len; }; + Span spans[7] = { + { W.packed_indices, (size_t)W.N*W.num_groups*64u }, + { W.norms, (size_t)W.N*W.num_groups*2 }, + { W.input_scale_recip, (size_t)W.K*2 }, + { W.left_signs, 128 }, + { W.right_signs, 128 }, + { W.permutation, 512 }, + { W.codebook, 32 }, + }; + uintptr_t mn = UINTPTR_MAX, mx = 0; + for (const auto& s : spans) { + if (!s.p) return false; + uintptr_t a = (uintptr_t)s.p; + if (a < mn) mn = a; + if (a + s.len > mx) mx = a + s.len; + } + lo = mn & ~(pg - 1); + hi = (mx + pg - 1) & ~(pg - 1); + for (int i = 0; i < 7; i++) offs[i] = (uintptr_t)spans[i].p - lo; + return true; +} + +bool moe_build_set4_zerocopy(MoESet4& S, const CactusQuantMatrix* Ws, uint32_t count) { + if (count == 0) return false; + uintptr_t lo0, hi0; + size_t offs[7]; + if (!moe_component_offsets(Ws[0], lo0, hi0, offs)) return false; + const size_t slot = hi0 - lo0; + if (slot > 0xFFFFFFFFull) return false; + if ((offs[0] & 7u) || (offs[1] & 1u) || (offs[2] & 1u) || (offs[6] & 1u) || (offs[5] & 3u)) return false; + for (uint32_t e = 1; e < count; ++e) { + uintptr_t lo, hi; + size_t o[7]; + if (!moe_component_offsets(Ws[e], lo, hi, o)) return false; + if (hi - lo != slot) return false; + for (int i = 0; i < 7; i++) if (o[i] != offs[i]) return false; + } + const size_t total = slot * count; + void* arena = mmap(nullptr, total, PROT_NONE, MAP_ANON | MAP_PRIVATE, -1, 0); + if (arena == MAP_FAILED) return false; + for (uint32_t e = 0; e < count; ++e) { + uintptr_t lo, hi; + size_t o[7]; + moe_component_offsets(Ws[e], lo, hi, o); + vm_address_t dst = (vm_address_t)((uintptr_t)arena + (size_t)e*slot); + vm_prot_t curp = 0, maxp = 0; + kern_return_t kr = vm_remap(mach_task_self(), &dst, slot, 0, + VM_FLAGS_FIXED | VM_FLAGS_OVERWRITE, mach_task_self(), + (vm_address_t)lo, FALSE, &curp, &maxp, VM_INHERIT_NONE); + if (kr != KERN_SUCCESS || dst != (vm_address_t)((uintptr_t)arena + (size_t)e*slot) + || !(curp & VM_PROT_READ)) { + munmap(arena, total); + return false; + } + } + id ab = [ctx().dev newBufferWithBytesNoCopy:arena length:total + options:MTLResourceStorageModeShared deallocator:nil]; + if (!ab) { munmap(arena, total); return false; } + CactusThreading::parallel_for(count, CactusThreading::ParallelConfig(1, 1), + [&](size_t start, size_t end) { + volatile uint64_t acc = 0; + for (size_t e = start; e < end; ++e) + for (size_t o = 0; o < slot; o += 16384) + acc += *((const volatile uint8_t*)arena + e*slot + o); + (void)acc; + }); + S.K=Ws[0].K; S.N=Ws[0].N; + S.pk=S.nm=S.rc=S.ls=S.rs=S.pm=S.cb=ab; + S.pk_off=offs[0]; S.nm_off=offs[1]; S.rc_off=offs[2]; S.ls_off=offs[3]; + S.rs_off=offs[4]; S.pm_off=offs[5]; S.cb_off=offs[6]; + S.stride = (uint32_t)slot; + return true; +} + +bool moe_build_set4(MoESet4& S, const CactusQuantMatrix* Ws, uint32_t count) { + const uint32_t K=Ws[0].K, N=Ws[0].N, gs=Ws[0].group_size, ng=Ws[0].num_groups; + if (gs != 128u || Ws[0].bits != 4u || (N % 4u) != 0u || (K % 128u) != 0u) return false; + for (uint32_t e = 0; e < count; ++e) { + const CactusQuantMatrix* W = &Ws[e]; + if (!quant_fast_eligible(W) || W->K != K || W->N != N || W->num_groups != ng) return false; + } + return moe_build_set4_zerocopy(S, Ws, count); +} +} + +bool cactus_metal_moe_cq4_ready(const CactusQuantMatrix* w1_0) { + std::lock_guard lk(g_resident_mu); + auto it = g_moe_cat.find(resident_key(w1_0)); + return it != g_moe_cat.end() && it->second.ok; +} + +bool cactus_metal_moe_cq4_build(const CactusQuantMatrix* w1s, const CactusQuantMatrix* w3s, + const CactusQuantMatrix* w2s, uint32_t num_experts) { + if (!ctx().ok || !ctx().psoMoeT || !ctx().psoMoeT2 || !ctx().psoMoeUp || !ctx().psoMoeDownAcc) + return false; + const uint64_t key = resident_key(&w1s[0]); + { + std::lock_guard lk(g_resident_mu); + auto it = g_moe_cat.find(key); + if (it != g_moe_cat.end()) return it->second.ok; + } + MoECat4 cat; + cat.ok = moe_build_set4(cat.w1, w1s, num_experts) + && moe_build_set4(cat.w3, w3s, num_experts) + && moe_build_set4(cat.w2, w2s, num_experts) + && cat.w1.K == cat.w3.K && cat.w1.N == cat.w3.N && cat.w2.K >= cat.w1.N; + if (!cat.ok) cat = MoECat4{}; + std::lock_guard lk(g_resident_mu); + auto ins = g_moe_cat.emplace(key, cat); + return ins.first->second.ok; +} + +bool cactus_metal_encode_moe_gated_cq4(void* out, const void* hidden, const void* probs, + const void* topk, const CactusQuantMatrix* w1_0, + uint32_t num_experts, uint32_t top_k, uint32_t tokens, + uint32_t act, uint32_t normalize, float eps, float scaling) { + if (!ctx().ok || !ctx().psoMoeT || top_k == 0 || top_k > 16 || tokens == 0) return false; + g_resident_mu.lock(); + auto it = g_moe_cat.find(resident_key(w1_0)); + bool hit = it != g_moe_cat.end() && it->second.ok; + g_resident_mu.unlock(); + if (!hit) return false; + MoECat4& C = it->second; + const uint32_t K1=C.w1.K, N1=C.w1.N, K2=C.w2.K, N2=C.w2.N; + const size_t slots = (size_t)tokens*top_k; + ensureEncoder(); + + id code1 = recycled(slots*K1*2); + id code3 = recycled(slots*K1*2); + id gu = recycled(slots*N1*2); + id code2 = recycled(slots*K2*2); + + { + [g_enc setComputePipelineState:ctx().psoMoeT2]; + setBufAt(hidden, (size_t)tokens*K1*2, 0); + setBufAt(topk, slots*sizeof(float), 1); + [g_enc setBuffer:C.w1.rc offset:C.w1.rc_off atIndex:2]; + [g_enc setBuffer:C.w1.ls offset:C.w1.ls_off atIndex:3]; + [g_enc setBuffer:C.w1.rs offset:C.w1.rs_off atIndex:4]; + [g_enc setBuffer:C.w1.pm offset:C.w1.pm_off atIndex:5]; + [g_enc setBuffer:C.w3.rc offset:C.w3.rc_off atIndex:6]; + [g_enc setBuffer:C.w3.ls offset:C.w3.ls_off atIndex:7]; + [g_enc setBuffer:C.w3.rs offset:C.w3.rs_off atIndex:8]; + [g_enc setBuffer:C.w3.pm offset:C.w3.pm_off atIndex:9]; + [g_enc setBuffer:code1 offset:0 atIndex:10]; + [g_enc setBuffer:code3 offset:0 atIndex:11]; + uint32_t Kc=K1, tk=top_k; + [g_enc setBytes:&Kc length:4 atIndex:12]; [g_enc setBytes:&tk length:4 atIndex:13]; + [g_enc setBytes:&C.w1.stride length:4 atIndex:14]; + [g_enc setBytes:&C.w3.stride length:4 atIndex:15]; + [g_enc setThreadgroupMemoryLength:128*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(K1/128u, 2u*top_k, tokens) + threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + } + { + [g_enc setComputePipelineState:ctx().psoMoeUp]; + [g_enc setBuffer:code1 offset:0 atIndex:0]; [g_enc setBuffer:code3 offset:0 atIndex:1]; + setBufAt(topk, slots*sizeof(float), 2); + [g_enc setBuffer:C.w1.pk offset:C.w1.pk_off atIndex:3]; [g_enc setBuffer:C.w1.nm offset:C.w1.nm_off atIndex:4]; [g_enc setBuffer:C.w1.cb offset:C.w1.cb_off atIndex:5]; + [g_enc setBuffer:C.w3.pk offset:C.w3.pk_off atIndex:6]; [g_enc setBuffer:C.w3.nm offset:C.w3.nm_off atIndex:7]; [g_enc setBuffer:C.w3.cb offset:C.w3.cb_off atIndex:8]; + [g_enc setBuffer:gu offset:0 atIndex:9]; + uint32_t Kc=K1, Nc=N1, a=act, tk=top_k; + [g_enc setBytes:&Kc length:4 atIndex:10]; [g_enc setBytes:&Nc length:4 atIndex:11]; + [g_enc setBytes:&a length:4 atIndex:12]; [g_enc setBytes:&tk length:4 atIndex:13]; + [g_enc setBytes:&C.w1.stride length:4 atIndex:14]; + [g_enc setBytes:&C.w3.stride length:4 atIndex:15]; + [g_enc dispatchThreadgroups:MTLSizeMake((N1+31u)/32u, top_k, tokens) + threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + } + { + [g_enc setComputePipelineState:ctx().psoMoeT]; + [g_enc setBuffer:gu offset:0 atIndex:0]; + setBufAt(topk, slots*sizeof(float), 1); + [g_enc setBuffer:C.w2.rc offset:C.w2.rc_off atIndex:2]; + [g_enc setBuffer:C.w2.ls offset:C.w2.ls_off atIndex:3]; + [g_enc setBuffer:C.w2.rs offset:C.w2.rs_off atIndex:4]; + [g_enc setBuffer:C.w2.pm offset:C.w2.pm_off atIndex:5]; + [g_enc setBuffer:code2 offset:0 atIndex:6]; + uint32_t Kc=K2, kv=N1, xz=top_k*N1, xy=N1, tk=top_k; + [g_enc setBytes:&Kc length:4 atIndex:7]; [g_enc setBytes:&kv length:4 atIndex:8]; + [g_enc setBytes:&xz length:4 atIndex:9]; [g_enc setBytes:&xy length:4 atIndex:10]; + [g_enc setBytes:&tk length:4 atIndex:11]; + [g_enc setBytes:&C.w2.stride length:4 atIndex:12]; + [g_enc setThreadgroupMemoryLength:128*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(K2/128u, top_k, tokens) + threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + } + { + [g_enc setComputePipelineState:ctx().psoMoeDownAcc]; + [g_enc setBuffer:code2 offset:0 atIndex:0]; + setBufAt(topk, slots*sizeof(float), 1); + setBufAt(probs, (size_t)tokens*num_experts*2, 2); + [g_enc setBuffer:C.w2.pk offset:C.w2.pk_off atIndex:3]; [g_enc setBuffer:C.w2.nm offset:C.w2.nm_off atIndex:4]; [g_enc setBuffer:C.w2.cb offset:C.w2.cb_off atIndex:5]; + setBufAt(out, (size_t)tokens*N2*2, 6); + uint32_t Kc=K2, Nc=N2, tk=top_k, nz=normalize, Ec=num_experts; + [g_enc setBytes:&Kc length:4 atIndex:7]; [g_enc setBytes:&Nc length:4 atIndex:8]; + [g_enc setBytes:&tk length:4 atIndex:9]; [g_enc setBytes:&nz length:4 atIndex:10]; + [g_enc setBytes:&eps length:4 atIndex:11]; [g_enc setBytes:&scaling length:4 atIndex:12]; + [g_enc setBytes:&Ec length:4 atIndex:13]; + [g_enc setBytes:&C.w2.stride length:4 atIndex:14]; + [g_enc setThreadgroupMemoryLength:top_k*16*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake((N2+15u)/16u, 1, tokens) + threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + } + return true; +} + +bool cactus_metal_encode_deltanet_decode(void* out, const void* q, const void* k, const void* v, + const void* g, const void* b, const void* s, + uint32_t B, uint32_t Hq, uint32_t Hv, + uint32_t K, uint32_t V, float scale) { + if (!ctx().ok || !ctx().psoDeltanet) return false; + if (B == 0 || Hq == 0 || Hv == 0 || K == 0 || V == 0 || V > 1024 || (Hv % Hq) != 0) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoDeltanet]; + setBufAt(q, (size_t)B*Hq*K*2, 0); + setBufAt(k, (size_t)B*Hq*K*2, 1); + setBufAt(v, (size_t)B*Hv*V*2, 2); + setBufAt(g, (size_t)B*Hv*2, 3); + setBufAt(b, (size_t)B*Hv*2, 4); + setBufAt(s, (size_t)B*K*Hv*V*2, 5); + setBufAt(out, (size_t)B*(1+K)*Hv*V*2, 6); + [g_enc setBytes:&Hq length:4 atIndex:7]; [g_enc setBytes:&Hv length:4 atIndex:8]; + [g_enc setBytes:&K length:4 atIndex:9]; [g_enc setBytes:&V length:4 atIndex:10]; + [g_enc setBytes:&scale length:4 atIndex:11]; + [g_enc setThreadgroupMemoryLength:(size_t)K*4 atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(Hv, B, 1) threadsPerThreadgroup:MTLSizeMake(V, 1, 1)]; + return true; +} + +bool cactus_metal_encode_deltanet_prefill(void* out, const void* q, const void* k, const void* v, + const void* g, const void* b, const void* s, + uint32_t B, uint32_t T, uint32_t Hq, uint32_t Hv, + uint32_t K, uint32_t V, float scale) { + if (!ctx().ok || !ctx().psoDeltanetPre) return false; + if (B == 0 || T == 0 || Hq == 0 || Hv == 0 || K == 0 || V == 0 || V > 1024 || (Hv % Hq) != 0) return false; + ensureEncoder(); + id scratch = recycled((size_t)B*Hv*K*V*sizeof(float)); + [g_enc setComputePipelineState:ctx().psoDeltanetPre]; + setBufAt(q, (size_t)B*T*Hq*K*2, 0); + setBufAt(k, (size_t)B*T*Hq*K*2, 1); + setBufAt(v, (size_t)B*T*Hv*V*2, 2); + setBufAt(g, (size_t)B*T*Hv*2, 3); + setBufAt(b, (size_t)B*T*Hv*2, 4); + setBufAt(s, (size_t)B*K*Hv*V*2, 5); + setBufAt(out, (size_t)B*(T+K)*Hv*V*2, 6); + [g_enc setBuffer:scratch offset:0 atIndex:7]; + [g_enc setBytes:&T length:4 atIndex:8]; [g_enc setBytes:&Hq length:4 atIndex:9]; + [g_enc setBytes:&Hv length:4 atIndex:10]; [g_enc setBytes:&K length:4 atIndex:11]; + [g_enc setBytes:&V length:4 atIndex:12]; [g_enc setBytes:&scale length:4 atIndex:13]; + [g_enc setThreadgroupMemoryLength:(size_t)K*4 atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(Hv, B, 1) threadsPerThreadgroup:MTLSizeMake(V, 1, 1)]; + return true; +} + +bool cactus_metal_encode_rms2_add_clip(void* out, const void* a, const void* wa, + const void* b, const void* wb, size_t dim, + float eps_a, float eps_b) { + if (!ctx().ok || !ctx().psoRms2AddClip || dim == 0) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoRms2AddClip]; + setBufAt(a, dim*2, 0); setBufAt(wa, dim*2, 1); + setBufAt(b, dim*2, 2); setBufAt(wb, dim*2, 3); + setBufAt(out, dim*2, 4); + uint32_t d=(uint32_t)dim; float ea=eps_a, eb=eps_b; + [g_enc setBytes:&d length:4 atIndex:5]; [g_enc setBytes:&ea length:4 atIndex:6]; + [g_enc setBytes:&eb length:4 atIndex:7]; + [g_enc setThreadgroupMemoryLength:512*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(1,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_rms_norm_scale(void* out, const void* in, const void* weight, + size_t rows, size_t dim, float eps, float oscale) { + if (!ctx().ok || !ctx().psoRmsScale || rows == 0 || dim == 0) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoRmsScale]; + setBufAt(in, rows*dim*2, 0); setBufAt(weight, dim*2, 1); setBufAt(out, rows*dim*2, 2); + uint32_t d=(uint32_t)dim; float e=eps, sc=oscale; + [g_enc setBytes:&d length:4 atIndex:3]; [g_enc setBytes:&e length:4 atIndex:4]; + [g_enc setBytes:&sc length:4 atIndex:5]; + [g_enc setThreadgroupMemoryLength:256*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_softmax_topk(void* probs, void* topk, const void* in, + size_t rows, size_t cols, size_t k, float scale) { + if (!ctx().ok || !ctx().psoSoftmaxTopk || k == 0 || k > 16 || rows == 0 || cols == 0) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoSoftmaxTopk]; + setBufAt(in, rows*cols*2, 0); + setBufAt(probs, rows*cols*2, 1); + setBufAt(topk, rows*k*2*sizeof(float), 2); + uint32_t E=(uint32_t)cols, kc=(uint32_t)k, B=(uint32_t)rows; float sc=scale; + [g_enc setBytes:&E length:4 atIndex:3]; + [g_enc setBytes:&kc length:4 atIndex:4]; + [g_enc setBytes:&B length:4 atIndex:5]; + [g_enc setBytes:&sc length:4 atIndex:6]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + return true; +} + +bool cactus_metal_encode_rope_pair(void* out, const void* x, const void* c, const void* s, + uint32_t H, uint32_t D) { + if (!ctx().ok || !ctx().psoRopePair || H == 0 || D == 0 || (D % 2) != 0) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoRopePair]; + setBufAt(x, (size_t)H*D*2, 0); + setBufAt(c, (size_t)D*2, 1); + setBufAt(s, (size_t)D*2, 2); + setBufAt(out, (size_t)H*D*2, 3); + [g_enc setBytes:&H length:4 atIndex:4]; + [g_enc setBytes:&D length:4 atIndex:5]; + [g_enc dispatchThreads:MTLSizeMake(D, H, 1) threadsPerThreadgroup:MTLSizeMake(64, 4, 1)]; + return true; +} + +bool cactus_metal_encode_rope_pair_rms(void* out, const void* x, const void* w, + const void* c, const void* s, + uint32_t H, uint32_t D, float eps) { + if (!ctx().ok || !ctx().psoRopePairRms || H == 0 || D == 0 || (D % 2) != 0 || D > 1024) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoRopePairRms]; + setBufAt(x, (size_t)H*D*2, 0); + setBufAt(w, (size_t)D*2, 1); + setBufAt(c, (size_t)D*2, 2); + setBufAt(s, (size_t)D*2, 3); + setBufAt(out, (size_t)H*D*2, 4); + float e = eps; + [g_enc setBytes:&D length:4 atIndex:5]; + [g_enc setBytes:&e length:4 atIndex:6]; + uint32_t T = D < 256 ? D : 256; + [g_enc setThreadgroupMemoryLength:T*sizeof(float) atIndex:0]; + [g_enc setThreadgroupMemoryLength:(size_t)D*2 atIndex:1]; + [g_enc dispatchThreadgroups:MTLSizeMake(H,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + return true; +} + +bool cactus_metal_encode_topk_rows(void* out, const void* in, size_t rows, size_t cols, size_t k) { + if (!ctx().ok || !ctx().psoTopkRows || k == 0 || k > 16 || rows == 0 || cols == 0) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoTopkRows]; + setBufAt(in, rows*cols*2, 0); + setBufAt(out, rows*k*2*sizeof(float), 1); + uint32_t F=(uint32_t)cols, kc=(uint32_t)k, B=(uint32_t)rows; + [g_enc setBytes:&F length:4 atIndex:2]; + [g_enc setBytes:&kc length:4 atIndex:3]; + [g_enc setBytes:&B length:4 atIndex:4]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + return true; +} + +bool cactus_metal_encode_quant_matmul_m(void* out, const void* lhs, const CactusQuantMatrix* W, uint32_t M) { + if (!ctx().ok || !W) return false; + if (M == 1 && (W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL)) + return cactus_metal_encode_quant_matmul(out, lhs, W); + + const uint32_t gs=W->group_size, K=W->K, N=W->N, ng=W->num_groups; + const uint32_t bits=W->bits, pgb=(gs*bits+7u)/8u; + const bool fast = ctx().ok && quant_fast_eligible(W); + if (!fast || !ctx().psoGmma) return false; + ResW& rw = resident(W); + if (!rw.ok) return false; + size_t code_bytes = (size_t)M*ng*gs*sizeof(__fp16); + if (!g_code_buf_m || (size_t)g_code_buf_m.length < code_bytes) + g_code_buf_m = owned_shared(code_bytes); + if (!g_code_buf_m) return false; + ensureEncoder(); + + [g_enc setComputePipelineState:ctx().psoTm]; + setBufAt(lhs, (size_t)M*K*2, 0); [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:1]; + [g_enc setBuffer:rw.lsign offset:rw.ls_off atIndex:2]; [g_enc setBuffer:rw.rsign offset:rw.rs_off atIndex:3]; + [g_enc setBuffer:rw.perm offset:rw.pm_off atIndex:4]; [g_enc setBuffer:g_code_buf_m offset:0 atIndex:5]; + [g_enc setBytes:&gs length:4 atIndex:6]; [g_enc setBytes:&K length:4 atIndex:7]; + [g_enc setThreadgroupMemoryLength:gs*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake((size_t)ng*M,1,1) threadsPerThreadgroup:MTLSizeMake(gs>1024u?1024u:gs,1,1)]; + + [g_enc setComputePipelineState:ctx().psoGmma]; + [g_enc setBuffer:g_code_buf_m offset:0 atIndex:0]; [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:1]; + [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:2]; [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:3]; + setBufAt(out, (size_t)M*N*2, 4); + [g_enc setBytes:&gs length:4 atIndex:5]; [g_enc setBytes:&ng length:4 atIndex:6]; + [g_enc setBytes:&pgb length:4 atIndex:7]; [g_enc setBytes:&N length:4 atIndex:8]; [g_enc setBytes:&M length:4 atIndex:9]; + [g_enc dispatchThreadgroups:MTLSizeMake((M+31)/32,(N+63)/64,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} + +bool cactus_metal_encode_quant_matmul_ortho(void* out, const void* act, void* code, + const CactusQuantMatrix* W) { + if (!ctx().ok || !W->rotation) return false; + const uint32_t K=W->K, N=W->N, ng=W->num_groups, gs=W->group_size, bits=W->bits; + const uint32_t pgb=(gs*bits+7u)/8u; + const bool lowbit = (bits == 2 || bits == 3); + if (bits != 4 && !lowbit) return false; + if (lowbit && (!ctx().psoGmrLow || (N % 16) != 0)) return false; + if (ng != 1 || gs != K || (N % 4) != 0) return false; + ResW& rw = resident(W); + if (!rw.ok || !rw.rotation || !rw.packed) return false; + ensureEncoder(); + + [g_enc setComputePipelineState:ctx().psoRotW]; + setBufAt(act, (size_t)K*2, 0); [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:1]; + [g_enc setBuffer:rw.rotation offset:rw.rot_off atIndex:2]; setBufAt(code, (size_t)K*2, 3); + [g_enc setBytes:&K length:4 atIndex:4]; + [g_enc setThreadgroupMemoryLength:256*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake((K+31)/32,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + + bool umr = ctx().psoGmr && (N % g_mr_rows == 0u); + [g_enc setComputePipelineState:(lowbit?ctx().psoGmrLow:(umr?ctx().psoGmr:ctx().psoG))]; + setBufAt(code, (size_t)K*2, 0); [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:1]; + [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:2]; [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:3]; + setBufAt(out, (size_t)N*2, 4); + [g_enc setBytes:&gs length:4 atIndex:5]; [g_enc setBytes:&ng length:4 atIndex:6]; + [g_enc setBytes:&pgb length:4 atIndex:7]; [g_enc setBytes:&N length:4 atIndex:8]; + if (lowbit) { [g_enc setBytes:&bits length:4 atIndex:9]; [g_enc setBytes:&rw.il length:4 atIndex:10]; } + else [g_enc setBytes:&rw.il length:4 atIndex:9]; + uint32_t ROWS=8; + uint32_t grid = (lowbit || umr) ? (N+g_mr_rows-1u)/g_mr_rows : (N+ROWS-1u)/ROWS; + [g_enc dispatchThreadgroups:MTLSizeMake(grid,1,1) threadsPerThreadgroup:MTLSizeMake(ROWS*32,1,1)]; + return true; +} + +bool cactus_metal_encode_embedding_ortho(void* out, uint32_t row, const CactusQuantMatrix* W, float scale) { + if (!ctx().ok || !W->rotation || (W->bits != 4 && W->bits != 2 && W->bits != 3)) return false; + const uint32_t K=W->K, ng=W->num_groups, gs=W->group_size, bits=W->bits; + if (ng != 1 || gs != K) return false; + if (!tg_mem_ok((size_t)K*sizeof(float))) return false; + if (bits != 4 && !(ctx().psoEmbOW && (K % 8u) == 0u)) return false; + if (!ctx().psoEmbOW || (K % 8u) != 0u) return false; + ResW& rw = resident(W); + if (!rw.ok || !rw.packed || !rw.rotation || !rw.codebook || !rw.norms || !rw.recip) return false; + ensureEncoder(); + { + [g_enc setComputePipelineState:ctx().psoEmbOW]; + [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:0]; [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:1]; + [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:2]; [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:3]; + [g_enc setBuffer:rw.rotation offset:rw.rot_off atIndex:4]; setBufAt(out, (size_t)K*2, 5); + [g_enc setBytes:&K length:4 atIndex:6]; [g_enc setBytes:&row length:4 atIndex:7]; + [g_enc setBytes:&scale length:4 atIndex:8]; [g_enc setBytes:&bits length:4 atIndex:9]; + [g_enc setBytes:&rw.il length:4 atIndex:10]; + [g_enc setThreadgroupMemoryLength:(size_t)K*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(K/8u,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; + } +} + +ResW& resident_emb_meta(const CactusQuantMatrix* W) { + std::lock_guard lk(g_resident_mu); + auto it = g_resident_emb.find(W->packed_indices); + if (it != g_resident_emb.end()) return it->second; + const uint32_t K=W->K, gs=W->group_size, ng=W->num_groups, bits=W->bits; + const uint32_t pgb=(gs*bits+7u)/8u; + ResW r; + bind_component(W->codebook, (size_t)(1u<bits)*sizeof(__fp16), 2, r.codebook, r.cb_off); + bind_component(W->left_signs, gs, 1, r.lsign, r.ls_off); + bind_component(W->right_signs, gs, 1, r.rsign, r.rs_off); + bind_component(W->permutation, (size_t)gs*sizeof(uint32_t), 4, r.perm, r.pm_off); + bind_component(W->input_scale_recip, (size_t)K*sizeof(__fp16), 2, r.recip, r.rc_off); + if (W->N && ng) { + auto wp = wrapHostPtr(W->packed_indices, (size_t)W->N*ng*pgb); + if (wp.first && ((uintptr_t)W->norms & 1u) == 0) { + auto wn = wrapHostPtr(W->norms, (size_t)W->N*ng*sizeof(__fp16)); + if (wn.first) { + r.packed = wp.first; r.packed_off = wp.second; + r.norms = wn.first; r.norms_off = wn.second; + } + } + } + return g_resident_emb.emplace(W->packed_indices, r).first->second; +} + +bool cactus_metal_encode_embedding_hadamard(void* out, uint32_t row, const CactusQuantMatrix* W) { + if (!ctx().ok || (W->bits != 4 && W->bits != 2 && W->bits != 3) || + (W->flags & (CACTUS_QUANT_FLAG_ORTHOGONAL | CACTUS_QUANT_FLAG_INTERLEAVED_4ROW))) return false; + const uint32_t K=W->K, gs=W->group_size, ng=W->num_groups, bits=W->bits, pgb=(gs*bits+7u)/8u; + if (gs > 256 || (gs & (gs-1)) != 0 || !W->packed_indices || !W->norms || !W->codebook + || !W->left_signs || !W->right_signs || !W->permutation || !W->input_scale_recip) return false; + ResW& rm = resident_emb_meta(W); + if (!rm.codebook || !rm.recip || !rm.lsign || !rm.rsign || !rm.perm) return false; + ensureEncoder(); + size_t row_bytes = (size_t)ng*pgb; + if (!rm.packed || !rm.norms) return false; + [g_enc setComputePipelineState:ctx().psoEmbH]; + [g_enc setBuffer:rm.packed offset:rm.packed_off + (size_t)row*row_bytes atIndex:0]; + [g_enc setBuffer:rm.norms offset:rm.norms_off + (size_t)row*ng*sizeof(__fp16) atIndex:2]; + [g_enc setBuffer:rm.codebook offset:rm.cb_off atIndex:1]; + [g_enc setBuffer:rm.recip offset:rm.rc_off atIndex:3]; + [g_enc setBuffer:rm.lsign offset:rm.ls_off atIndex:4]; [g_enc setBuffer:rm.rsign offset:rm.rs_off atIndex:5]; + [g_enc setBuffer:rm.perm offset:rm.pm_off atIndex:6]; setBufAt(out, (size_t)K*2, 7); + [g_enc setBytes:&gs length:4 atIndex:8]; [g_enc setBytes:&bits length:4 atIndex:9]; + [g_enc setThreadgroupMemoryLength:(size_t)gs*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(ng,1,1) threadsPerThreadgroup:MTLSizeMake(gs,1,1)]; + return true; +} + +static id registerReadonly(const void* p, size_t bytes) { + if (!p) return nil; + auto it = g_readonly.find(p); + if (it != g_readonly.end()) return it->second; + uintptr_t a = (uintptr_t)p, base = a & ~(uintptr_t)16383u; + size_t wraplen = (((a - base) + bytes + 16383u) & ~(size_t)16383u); + id b = [ctx().dev newBufferWithBytesNoCopy:(void*)base length:wraplen + options:MTLResourceStorageModeShared deallocator:nil]; + if (b) g_readonly[p] = b; + return b; +} + +bool cactus_metal_encode_embedding_ortho_m(void* out, const CactusQuantMatrix* W, const uint32_t* rows, uint32_t M) { + if (!ctx().ok || !ctx().psoEmbOm || !W->rotation || (W->bits != 4 && W->bits != 2 && W->bits != 3)) return false; + const uint32_t K=W->K, ng=W->num_groups, gs=W->group_size, bits=W->bits; + if (ng != 1 || gs != K) return false; + ResW& rw = resident(W); + if (!rw.packed || !rw.rotation || !rw.codebook || !rw.norms || !rw.recip) return false; + ensureEncoder(); + id rb = recycled((size_t)M*sizeof(uint32_t)); + std::memcpy([rb contents], rows, (size_t)M*sizeof(uint32_t)); + [g_enc setComputePipelineState:ctx().psoEmbOm]; + [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:0]; [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:1]; + [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:2]; [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:3]; + [g_enc setBuffer:rw.rotation offset:rw.rot_off atIndex:4]; [g_enc setBuffer:rb offset:0 atIndex:5]; + setBufAt(out, (size_t)M*K*2, 6); [g_enc setBytes:&K length:4 atIndex:7]; [g_enc setBytes:&M length:4 atIndex:8]; + [g_enc setBytes:&bits length:4 atIndex:9]; + [g_enc setBytes:&rw.il length:4 atIndex:10]; + [g_enc dispatchThreadgroups:MTLSizeMake(((K+15)/16)*((M+15)/16),1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_embedding_hadamard_m(void* out, const CactusQuantMatrix* W, const uint32_t* rows, uint32_t M) { + if (!ctx().ok || !ctx().psoEmbHm || (W->bits != 4 && W->bits != 2 && W->bits != 3) || + (W->flags & (CACTUS_QUANT_FLAG_ORTHOGONAL | CACTUS_QUANT_FLAG_INTERLEAVED_4ROW))) return false; + const uint32_t K=W->K, gs=W->group_size, ng=W->num_groups, bits=W->bits, pgb=(gs*bits+7u)/8u; + if (gs > 256 || (gs & (gs-1)) != 0 || !W->packed_indices || !W->norms || !W->codebook + || !W->left_signs || !W->right_signs || !W->permutation || !W->input_scale_recip) return false; + ResW& rm = resident_emb_meta(W); + if (!rm.codebook || !rm.recip || !rm.lsign || !rm.rsign || !rm.perm) return false; + ensureEncoder(); + size_t rowb = (size_t)ng*pgb; + id pk = recycled((size_t)M*rowb); + id nm = recycled((size_t)M*ng*sizeof(__fp16)); + for (uint32_t m=0; mpacked_indices+(size_t)rows[m]*rowb, rowb); + std::memcpy((__fp16*)[nm contents]+(size_t)m*ng, (const __fp16*)W->norms+(size_t)rows[m]*ng, (size_t)ng*sizeof(__fp16)); + } + [g_enc setComputePipelineState:ctx().psoEmbHm]; + [g_enc setBuffer:pk offset:0 atIndex:0]; [g_enc setBuffer:rm.codebook offset:rm.cb_off atIndex:1]; + [g_enc setBuffer:nm offset:0 atIndex:2]; [g_enc setBuffer:rm.recip offset:rm.rc_off atIndex:3]; + [g_enc setBuffer:rm.lsign offset:rm.ls_off atIndex:4]; [g_enc setBuffer:rm.rsign offset:rm.rs_off atIndex:5]; + [g_enc setBuffer:rm.perm offset:rm.pm_off atIndex:6]; setBufAt(out, (size_t)M*K*2, 7); + [g_enc setBytes:&gs length:4 atIndex:8]; [g_enc setBytes:&ng length:4 atIndex:9]; [g_enc setBytes:&K length:4 atIndex:10]; + [g_enc setBytes:&bits length:4 atIndex:11]; + [g_enc setThreadgroupMemoryLength:(size_t)gs*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake((size_t)ng*M,1,1) threadsPerThreadgroup:MTLSizeMake(gs,1,1)]; + return true; +} + + +bool cactus_metal_encode_gather_f16(void* out, const void* table, size_t table_bytes, + const uint32_t* rows, uint32_t M, uint32_t D) { + if (!ctx().ok || !ctx().psoGather || M == 0 || D == 0) return false; + id tb = registerReadonly(table, table_bytes); + if (!tb) return false; + size_t toff = (uintptr_t)table & (uintptr_t)16383u; + ensureEncoder(); + id rb = recycled((size_t)M*sizeof(uint32_t)); + std::memcpy([rb contents], rows, (size_t)M*sizeof(uint32_t)); + uint32_t n = M*D; + [g_enc setComputePipelineState:ctx().psoGather]; + [g_enc setBuffer:tb offset:toff atIndex:0]; [g_enc setBuffer:rb offset:0 atIndex:1]; + setBufAt(out, (size_t)n*2, 2); + [g_enc setBytes:&D length:4 atIndex:3]; [g_enc setBytes:&n length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +static void setCacheAt(const void* p, size_t bytes, int idx) { + if (p && bytes) setBufAt(p, bytes, idx); + else [g_enc setBuffer:ctx().dummy offset:0 atIndex:idx]; +} + +bool cactus_metal_encode_attention_i8( + void* out, const void* q, const void* knew, const void* vnew, + const void* kc, const void* vc, const void* ks, const void* vs, + uint32_t num_q_heads, uint32_t num_kv_heads, uint32_t head_dim, uint32_t v_hdim, + uint32_t history_len, uint32_t total_keys, uint32_t kv_start, uint32_t kv_end, + float scale, size_t kc_bytes, size_t vc_bytes, size_t ks_bytes, size_t vs_bytes) { + if (!ctx().ok) return false; + if (kv_end <= kv_start || (num_q_heads % num_kv_heads) != 0) return false; + if (head_dim > 512u || v_hdim > 512u) return false; + ensureEncoder(); + auto setInputs = [&]() { + setBufAt(q, (size_t)num_q_heads*head_dim*2, 0); + if (total_keys > history_len && knew && vnew) { + setBufAt(knew, (size_t)num_kv_heads*head_dim*2, 1); + setBufAt(vnew, (size_t)num_kv_heads*v_hdim*2, 2); + } else { + [g_enc setBuffer:ctx().dummy offset:0 atIndex:1]; [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + } + setCacheAt(kc, kc_bytes, 3); setCacheAt(vc, vc_bytes, 4); + setCacheAt(ks, ks_bytes, 5); setCacheAt(vs, vs_bytes, 6); + }; + + const uint32_t T = 256, nsg = T / 32u; + const uint32_t R = kv_end - kv_start; + uint32_t nwg = R / 24u; if (nwg < 1u) nwg = 1u; if (nwg > 32u) nwg = 32u; + id partO = ctx().dummy, partML = ctx().dummy; + if (nwg > 1u) { + partO = recycled((size_t)num_q_heads*nwg*v_hdim*sizeof(float)); + partML = recycled((size_t)num_q_heads*nwg*2*sizeof(float)); + } + [g_enc setComputePipelineState:ctx().psoAttn]; + setInputs(); + setBufAt(out, (size_t)num_q_heads*v_hdim*2, 7); + [g_enc setBytes:&num_q_heads length:4 atIndex:8]; [g_enc setBytes:&num_kv_heads length:4 atIndex:9]; + [g_enc setBytes:&head_dim length:4 atIndex:10]; [g_enc setBytes:&v_hdim length:4 atIndex:11]; + [g_enc setBytes:&history_len length:4 atIndex:12];[g_enc setBytes:&scale length:4 atIndex:13]; + [g_enc setBytes:&kv_start length:4 atIndex:14]; [g_enc setBytes:&kv_end length:4 atIndex:15]; + [g_enc setBuffer:partO offset:0 atIndex:16]; [g_enc setBuffer:partML offset:0 atIndex:17]; + [g_enc setBytes:&nwg length:4 atIndex:18]; + [g_enc setThreadgroupMemoryLength:((size_t)nsg*v_hdim + 2*nsg)*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(num_q_heads*nwg,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + if (nwg > 1u) { + [g_enc setComputePipelineState:ctx().psoAttnC]; + [g_enc setBuffer:partO offset:0 atIndex:0]; [g_enc setBuffer:partML offset:0 atIndex:1]; + setBufAt(out, (size_t)num_q_heads*v_hdim*2, 2); + [g_enc setBytes:&v_hdim length:4 atIndex:3]; [g_enc setBytes:&nwg length:4 atIndex:4]; + [g_enc dispatchThreadgroups:MTLSizeMake(num_q_heads,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + } + return true; +} + +bool cactus_metal_encode_attention_fused_i8( + void* out, const void* q, const void* kraw, const void* vraw, + void* kc, void* vc, void* ks, void* vs, + const void* qw, const void* kw, const void* vw, const void* cs, const void* sn, + uint32_t nqh, uint32_t hd, uint32_t vhd, + uint32_t kv_start, uint32_t kv_end, uint32_t slot, uint32_t has_new, + float eps, float scale, + size_t kc_bytes, size_t vc_bytes, size_t ks_bytes, size_t vs_bytes) { + if (!ctx().ok || !ctx().psoAttnF) return false; + if (kv_end <= kv_start) return false; + ensureEncoder(); + const uint32_t T = 256, nsg = T / 32u; + const uint32_t R = kv_end - kv_start; + uint32_t nwg = R / 24u; if (nwg < 1u) nwg = 1u; if (nwg > 32u) nwg = 32u; + id partO = ctx().dummy, partML = ctx().dummy; + if (nwg > 1u) { + partO = recycled((size_t)nqh*nwg*vhd*sizeof(float)); + partML = recycled((size_t)nqh*nwg*2*sizeof(float)); + } + struct { uint32_t nqh, nkvh, hd, vhd, hist, kv_start, kv_end, nwg, slot, has_new; float eps, scale; } U = + { nqh, 1u, hd, vhd, 0u, kv_start, kv_end, nwg, slot, has_new, eps, scale }; + [g_enc setComputePipelineState:ctx().psoAttnF]; + setBufAt(q, (size_t)nqh*hd*2, 0); + if (has_new && kraw && vraw) { setBufAt(kraw, (size_t)hd*2, 1); setBufAt(vraw, (size_t)vhd*2, 2); } + else { [g_enc setBuffer:ctx().dummy offset:0 atIndex:1]; [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; } + setBufAt(kc, kc_bytes, 3); setBufAt(vc, vc_bytes, 4); + setBufAt(ks, ks_bytes, 5); setBufAt(vs, vs_bytes, 6); + setBufAt(out, (size_t)nqh*vhd*2, 7); + setBufAt(qw, (size_t)hd*2, 8); + if (has_new && kw) setBufAt(kw, (size_t)hd*2, 9); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:9]; + if (has_new && vw) setBufAt(vw, (size_t)vhd*2, 10); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:10]; + setBufAt(cs, (size_t)hd*2, 11); setBufAt(sn, (size_t)hd*2, 12); + [g_enc setBuffer:partO offset:0 atIndex:13]; [g_enc setBuffer:partML offset:0 atIndex:14]; + [g_enc setBytes:&U length:sizeof(U) atIndex:15]; + [g_enc setThreadgroupMemoryLength:((size_t)nsg*vhd + 2*nsg + 256)*sizeof(float) atIndex:0]; + [g_enc setThreadgroupMemoryLength:(size_t)(2*hd + vhd)*2 atIndex:1]; + [g_enc dispatchThreadgroups:MTLSizeMake(nqh*nwg,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + if (nwg > 1u) { + [g_enc setComputePipelineState:ctx().psoAttnC]; + [g_enc setBuffer:partO offset:0 atIndex:0]; [g_enc setBuffer:partML offset:0 atIndex:1]; + setBufAt(out, (size_t)nqh*vhd*2, 2); + [g_enc setBytes:&vhd length:4 atIndex:3]; [g_enc setBytes:&nwg length:4 atIndex:4]; + [g_enc dispatchThreadgroups:MTLSizeMake(nqh,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + } + return true; +} + +bool cactus_metal_encode_attention_i8_prefill( + void* out, const void* q, const void* knew, const void* vnew, + const void* kc, const void* vc, const void* ks, const void* vs, + uint32_t num_q_heads, uint32_t num_kv_heads, uint32_t head_dim, uint32_t v_hdim, + uint32_t history_len, uint32_t new_len, uint32_t q_pos0, uint32_t window, uint32_t is_causal, uint32_t M, + float scale, size_t kc_bytes, size_t vc_bytes, size_t ks_bytes, size_t vs_bytes, + uint32_t sink, uint32_t ring) { + if (!ctx().ok) return false; + uint32_t total_keys = history_len + new_len; + const uint32_t T = 128; + if (total_keys == 0 || M == 0 || (num_q_heads % num_kv_heads) != 0) return false; + uint32_t maxsc = (ring > 0u) ? ((total_keys > sink + ring) ? (sink + ring) : total_keys) : 256u; + if (maxsc > 7936u) return false; + bool mma2 = (ring == 0u && window == 0u && is_causal && head_dim == 512u && v_hdim == 512u + && num_q_heads == 8u && num_kv_heads == 1u); + bool hd256 = (ctx().psoAttnPreHd256 != nil && is_causal && head_dim == 256u && v_hdim == 256u + && num_q_heads == 8u && num_kv_heads == 1u); + if (mma2 || hd256) { + ensureEncoder(); + [g_enc setComputePipelineState:(hd256 ? ctx().psoAttnPreHd256 : ctx().psoAttnPreMma2)]; + setBufAt(q, (size_t)M*num_q_heads*head_dim*2, 0); + if (new_len > 0 && knew && vnew) { + setBufAt(knew, (size_t)new_len*num_kv_heads*head_dim*2, 1); + setBufAt(vnew, (size_t)new_len*num_kv_heads*v_hdim*2, 2); + } else { + [g_enc setBuffer:ctx().dummy offset:0 atIndex:1]; [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + } + setCacheAt(kc, kc_bytes, 3); setCacheAt(vc, vc_bytes, 4); setCacheAt(ks, ks_bytes, 5); setCacheAt(vs, vs_bytes, 6); + setBufAt(out, (size_t)M*num_q_heads*v_hdim*2, 7); + [g_enc setBytes:&num_q_heads length:4 atIndex:8]; [g_enc setBytes:&num_kv_heads length:4 atIndex:9]; + [g_enc setBytes:&head_dim length:4 atIndex:10]; [g_enc setBytes:&v_hdim length:4 atIndex:11]; + [g_enc setBytes:&history_len length:4 atIndex:12];[g_enc setBytes:&scale length:4 atIndex:13]; + [g_enc setBytes:&q_pos0 length:4 atIndex:14]; [g_enc setBytes:&new_len length:4 atIndex:15]; + [g_enc setBytes:&M length:4 atIndex:16]; + if (hd256) { [g_enc setBytes:&sink length:4 atIndex:17]; [g_enc setBytes:&ring length:4 atIndex:18]; } + uint32_t mtiles = (M + 7u)/8u; + [g_enc dispatchThreadgroups:MTLSizeMake(mtiles, 1, 1) threadsPerThreadgroup:MTLSizeMake(hd256?256u:512u,1,1)]; + return true; + } + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoAttnPre]; + setBufAt(q, (size_t)M*num_q_heads*head_dim*2, 0); + if (new_len > 0 && knew && vnew) { + setBufAt(knew, (size_t)new_len*num_kv_heads*head_dim*2, 1); + setBufAt(vnew, (size_t)new_len*num_kv_heads*v_hdim*2, 2); + } else { + [g_enc setBuffer:ctx().dummy offset:0 atIndex:1]; [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + } + setCacheAt(kc, kc_bytes, 3); setCacheAt(vc, vc_bytes, 4); setCacheAt(ks, ks_bytes, 5); setCacheAt(vs, vs_bytes, 6); + setBufAt(out, (size_t)M*num_q_heads*v_hdim*2, 7); + [g_enc setBytes:&num_q_heads length:4 atIndex:8]; [g_enc setBytes:&num_kv_heads length:4 atIndex:9]; + [g_enc setBytes:&head_dim length:4 atIndex:10]; [g_enc setBytes:&v_hdim length:4 atIndex:11]; + [g_enc setBytes:&history_len length:4 atIndex:12];[g_enc setBytes:&scale length:4 atIndex:13]; + [g_enc setBytes:&q_pos0 length:4 atIndex:14]; [g_enc setBytes:&new_len length:4 atIndex:15]; + [g_enc setBytes:&window length:4 atIndex:16]; [g_enc setBytes:&is_causal length:4 atIndex:17]; + [g_enc setBytes:&maxsc length:4 atIndex:18]; + [g_enc setBytes:&sink length:4 atIndex:19]; [g_enc setBytes:&ring length:4 atIndex:20]; + [g_enc setThreadgroupMemoryLength:((size_t)maxsc + T)*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(M*num_q_heads,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + return true; +} + +bool cactus_metal_encode_binary_f32(int op, void* out, const void* a, const void* b, size_t n) { + if (!ctx().ok || !ctx().psoBinF32) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; int o=op; + [g_enc setComputePipelineState:ctx().psoBinF32]; + setBufAt(a, n*4, 0); setBufAt(b, n*4, 1); setBufAt(out, n*4, 2); + [g_enc setBytes:&nn length:4 atIndex:3]; [g_enc setBytes:&o length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake((n+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_scalar_f32(int op, void* out, const void* in, size_t n, float p) { + if (!ctx().ok || !ctx().psoScaF32) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; int o=op; + [g_enc setComputePipelineState:ctx().psoScaF32]; + setBufAt(in, n*4, 0); setBufAt(out, n*4, 1); + [g_enc setBytes:&nn length:4 atIndex:2]; [g_enc setBytes:&o length:4 atIndex:3]; [g_enc setBytes:&p length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake((n+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_unary_f32(int op, void* out, const void* in, size_t n) { + if (!ctx().ok || !ctx().psoUnaF32) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; int o=op; + [g_enc setComputePipelineState:ctx().psoUnaF32]; + setBufAt(in, n*4, 0); setBufAt(out, n*4, 1); + [g_enc setBytes:&nn length:4 atIndex:2]; [g_enc setBytes:&o length:4 atIndex:3]; + [g_enc dispatchThreads:MTLSizeMake((n+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_clamp(void* out, const void* in, size_t n, float lo, float hi, int f32) { + id pso = f32 ? ctx().psoClampF32 : ctx().psoClampF16; + if (!ctx().ok || !pso) return false; + ensureEncoder(); + uint32_t nn=(uint32_t)n; size_t es = f32?4:2; + [g_enc setComputePipelineState:pso]; + setBufAt(in, n*es, 0); setBufAt(out, n*es, 1); + [g_enc setBytes:&nn length:4 atIndex:2]; [g_enc setBytes:&lo length:4 atIndex:3]; [g_enc setBytes:&hi length:4 atIndex:4]; + size_t grid = (n + 3) / 4; + [g_enc dispatchThreads:MTLSizeMake(grid,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_glu(void* out, const void* in, size_t split, size_t inner, size_t n_out) { + if (!ctx().ok || !ctx().psoGlu) return false; + ensureEncoder(); + uint32_t sp=(uint32_t)split, in_=(uint32_t)inner, nn=(uint32_t)n_out; + [g_enc setComputePipelineState:ctx().psoGlu]; + setBufAt(in, n_out*4, 0); setBufAt(out, n_out*2, 1); + [g_enc setBytes:&sp length:4 atIndex:2]; [g_enc setBytes:&in_ length:4 atIndex:3]; [g_enc setBytes:&nn length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake(n_out,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_layer_norm(void* out, const void* in, const void* w, const void* b, + size_t rows, size_t dim, float eps) { + if (!ctx().ok || !ctx().psoLayerNorm) return false; + ensureEncoder(); + uint32_t d=(uint32_t)dim, hb = b ? 1u : 0u; + uint32_t T = dim >= 1024 ? 1024u : 256u; + [g_enc setComputePipelineState:ctx().psoLayerNorm]; + setBufAt(in, rows*dim*2, 0); setBufAt(w, dim*2, 1); + if (b) setBufAt(b, dim*2, 2); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + setBufAt(out, rows*dim*2, 3); + [g_enc setBytes:&d length:4 atIndex:4]; [g_enc setBytes:&eps length:4 atIndex:5]; [g_enc setBytes:&hb length:4 atIndex:6]; + [g_enc setThreadgroupMemoryLength:2*T*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + return true; +} +bool cactus_metal_encode_softmax_rows(void* out, const void* in, size_t rows, size_t cols) { + if (!ctx().ok || !ctx().psoSoftmaxR) return false; + ensureEncoder(); + uint32_t c=(uint32_t)cols; + uint32_t T = cols >= 1024 ? 1024u : 256u; + [g_enc setComputePipelineState:ctx().psoSoftmaxR]; + setBufAt(in, rows*cols*2, 0); setBufAt(out, rows*cols*2, 1); + [g_enc setBytes:&c length:4 atIndex:2]; + [g_enc setThreadgroupMemoryLength:T*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(T,1,1)]; + return true; +} +bool cactus_metal_encode_conv1d_k3(void* out, const void* x, const void* w, int w_int8, + const void* w_scales, uint32_t w_gs, + uint32_t Cin, uint32_t L, uint32_t Cout, uint32_t Lout, uint32_t stride) { + if (!ctx().ok || !ctx().psoConv1dK3) return false; + ensureEncoder(); + id dq = nil; + if (w_int8) { + size_t rows = Cout, cols = (size_t)Cin*3; + dq = recycled(rows*cols*2); + if (!dq) return false; + __fp16* dst = (__fp16*)[dq contents]; + const int8_t* src = (const int8_t*)w; + if (w_gs == 0 || !w_scales) { + for (size_t i = 0; i < rows*cols; ++i) dst[i] = (__fp16)src[i]; + } else { + const __fp16* scales = (const __fp16*)w_scales; + size_t ngr = cols / w_gs; + for (size_t r = 0; r < rows; ++r) + for (size_t c = 0; c < cols; ++c) + dst[r*cols+c] = (__fp16)((float)src[r*cols+c] * (float)scales[r*ngr + c/w_gs]); + } + } + [g_enc setComputePipelineState:ctx().psoConv1dK3]; + setBufAt(x, (size_t)Cin*L*2, 0); + if (dq) [g_enc setBuffer:dq offset:0 atIndex:1]; else setBufAt(w, (size_t)Cout*Cin*3*2, 1); + setBufAt(out, (size_t)Cout*Lout*2, 2); + [g_enc setBytes:&Cin length:4 atIndex:3]; [g_enc setBytes:&L length:4 atIndex:4]; + [g_enc setBytes:&Lout length:4 atIndex:5]; [g_enc setBytes:&stride length:4 atIndex:6]; + [g_enc dispatchThreads:MTLSizeMake(Lout,Cout,1) threadsPerThreadgroup:MTLSizeMake(64,4,1)]; + return true; +} +static MPSMatrixMultiplication* mps_gemm_ab(uint32_t M, uint32_t N, uint32_t K, bool tr, float alpha, float beta) { + static std::map, MPSMatrixMultiplication*> cache; + uint32_t ab, bb; + std::memcpy(&ab, &alpha, 4); + std::memcpy(&bb, &beta, 4); + auto key = std::make_tuple(M,N,K,tr,((uint64_t)ab<<32)|bb); + auto it = cache.find(key); + if (it != cache.end()) return it->second; + MPSMatrixMultiplication* mm = [[MPSMatrixMultiplication alloc] + initWithDevice:ctx().dev transposeLeft:false transposeRight:tr + resultRows:M resultColumns:N interiorColumns:K alpha:alpha beta:beta]; + cache[key] = mm; + return mm; +} +static MPSMatrixMultiplication* mps_gemm_a(uint32_t M, uint32_t N, uint32_t K, bool tr, float alpha) { + return mps_gemm_ab(M, N, K, tr, alpha, 0.0f); +} +static MPSMatrixMultiplication* mps_gemm(uint32_t M, uint32_t N, uint32_t K, bool tr) { + return mps_gemm_a(M, N, K, tr, 1.0f); +} + +bool cactus_metal_encode_gemm_f16(void* out, const void* lhs, const void* rhs, + uint32_t M, uint32_t K, uint32_t N, int pretransposed) { + if (!ctx().ok) return false; + static const bool mps_ok = MPSSupportsMTLDevice(ctx().dev); + auto a = bufForPtrOff(lhs, (size_t)M*K*2); + auto b = bufForPtrOff(rhs, (size_t)N*K*2); + auto c = bufForPtrOff(out, (size_t)M*N*2); + int use_tr = pretransposed; + if (mps_ok && (a.second % 16)==0 && (b.second % 16)==0 && (c.second % 16)==0 + && (K % 8)==0 && (N % 8)==0) { + if (g_enc) { [g_enc endEncoding]; g_enc = nil; } + if (!g_cmd) g_cmd = [ctx().queue commandBuffer]; + MPSMatrixDescriptor* da = [MPSMatrixDescriptor matrixDescriptorWithRows:M columns:K + rowBytes:(size_t)K*2 dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor* db = use_tr + ? [MPSMatrixDescriptor matrixDescriptorWithRows:N columns:K rowBytes:(size_t)K*2 dataType:MPSDataTypeFloat16] + : [MPSMatrixDescriptor matrixDescriptorWithRows:K columns:N rowBytes:(size_t)N*2 dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor* dc = [MPSMatrixDescriptor matrixDescriptorWithRows:M columns:N + rowBytes:(size_t)N*2 dataType:MPSDataTypeFloat16]; + auto mat = [](id buf, size_t off, MPSMatrixDescriptor* d) -> MPSMatrix* { + auto key = std::make_tuple((__bridge void*)buf, off, (size_t)d.rows, (size_t)d.columns); + auto it = g_mps_mats.find(key); + if (it != g_mps_mats.end()) return it->second; + if (g_mps_mats.size() > 4096) g_mps_mats.clear(); + MPSMatrix* m = [[MPSMatrix alloc] initWithBuffer:buf offset:off descriptor:d]; + g_mps_mats[key] = m; + return m; + }; + MPSMatrix* ma = mat(a.first, a.second, da); + MPSMatrix* mb = mat(b.first, b.second, db); + MPSMatrix* mc = mat(c.first, c.second, dc); + [mps_gemm(M, N, K, use_tr != 0) encodeToCommandBuffer:g_cmd + leftMatrix:ma rightMatrix:mb resultMatrix:mc]; + return true; + } + if (!ctx().psoGdense || !ctx().psoStrided) return false; + ensureEncoder(); + id tbuf = nil; + if (!pretransposed) { + tbuf = recycled((size_t)K*N*2); + if (!tbuf) return false; + uint32_t oshape[2] = { N, K }, sstride[2] = { 1u, N }; + uint32_t total = N*K; + [g_enc setComputePipelineState:ctx().psoStrided]; + setBufAt(rhs, (size_t)K*N*2, 0); [g_enc setBuffer:tbuf offset:0 atIndex:1]; + [g_enc setBytes:oshape length:8 atIndex:2]; [g_enc setBytes:sstride length:8 atIndex:3]; + uint32_t nd=2, base=0; + [g_enc setBytes:&nd length:4 atIndex:4]; [g_enc setBytes:&total length:4 atIndex:5]; [g_enc setBytes:&base length:4 atIndex:6]; + [g_enc dispatchThreads:MTLSizeMake(total,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + } + [g_enc setComputePipelineState:ctx().psoGdense]; + setBufAt(lhs, (size_t)M*K*2, 0); + if (tbuf) [g_enc setBuffer:tbuf offset:0 atIndex:1]; else setBufAt(rhs, (size_t)N*K*2, 1); + setBufAt(out, (size_t)M*N*2, 2); + [g_enc setBytes:&K length:4 atIndex:3]; [g_enc setBytes:&N length:4 atIndex:4]; [g_enc setBytes:&M length:4 atIndex:5]; + [g_enc dispatchThreadgroups:MTLSizeMake((M+31)/32,(N+63)/64,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} +bool cactus_metal_encode_attention_f16(void* out, const void* q, const void* k, const void* v, const void* mask, + uint32_t B, uint32_t T, uint32_t S, uint32_t HQ, uint32_t HKV, uint32_t D, uint32_t DV, + float scale, uint32_t causal, uint32_t pos_off, uint32_t window, float logit_cap, uint32_t mask_mode) { + if (!ctx().ok || !ctx().psoAttnDense) return false; + if (D > 128 || DV > 128 || (D % 4) != 0 || (DV % 4) != 0 || HKV == 0 || HQ % HKV != 0) return false; + static const bool mps_attn_ok = MPSSupportsMTLDevice(ctx().dev); + bool mask_ok = mask_mode == 0 || (mask_mode == 2 && mask && ctx().psoMaskFill); + if (ctx().psoAttnFlash && B == 1 && HQ == HKV && !causal && window == 0 + && logit_cap == 0.0f && D == 64 && DV == 64 && S >= 256 + && (mask_mode == 0 || (mask_mode == 2 && mask))) { + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoAttnFlash]; + setBufAt(q, (size_t)T*HQ*D*2, 0); setBufAt(k, (size_t)S*HKV*D*2, 1); + setBufAt(v, (size_t)S*HKV*DV*2, 2); setBufAt(out, (size_t)T*HQ*DV*2, 3); + if (mask_mode == 2 && mask) setBufAt(mask, (size_t)T*S*2, 4); + else [g_enc setBuffer:ctx().dummy offset:0 atIndex:4]; + [g_enc setBytes:&T length:4 atIndex:5]; [g_enc setBytes:&S length:4 atIndex:6]; + [g_enc setBytes:&HQ length:4 atIndex:7]; [g_enc setBytes:&scale length:4 atIndex:8]; + [g_enc setBytes:&mask_mode length:4 atIndex:9]; + [g_enc dispatchThreadgroups:MTLSizeMake((T + 63) / 64, HQ, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + return true; + } + if (mps_attn_ok && B == 1 && HQ == HKV && !causal && mask_ok && window == 0 + && logit_cap == 0.0f && S >= 512 && (S % 8) == 0 && (D % 8) == 0 && (DV % 8) == 0 + && ctx().psoSoftmaxR) { + auto qb = bufForPtrOff(q, (size_t)T*HQ*D*2); + auto kb = bufForPtrOff(k, (size_t)S*HKV*D*2); + auto vb = bufForPtrOff(v, (size_t)S*HKV*DV*2); + auto ob = bufForPtrOff(out, (size_t)T*HQ*DV*2); + bool aligned = (qb.second % 16)==0 && (kb.second % 16)==0 && (vb.second % 16)==0 && (ob.second % 16)==0 + && ((D*2) % 16)==0 && ((DV*2) % 16)==0; + if (aligned) { + size_t need = (size_t)HQ*T*S*2; + if (!g_attn_scores || g_attn_scores.length < need) + g_attn_scores = owned_shared(need); + id sbuf = g_attn_scores; + if (sbuf) { + float beta = 0.0f; + if (mask_mode == 2) { + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoMaskFill]; + setBufAt(mask, (size_t)T*S*2, 0); + [g_enc setBuffer:sbuf offset:0 atIndex:1]; + uint32_t ts = T*S; + [g_enc setBytes:&ts length:4 atIndex:2]; + [g_enc dispatchThreads:MTLSizeMake(ts, HQ, 1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + beta = 1.0f; + } + if (g_enc) { [g_enc endEncoding]; g_enc = nil; } + if (!g_cmd) g_cmd = [ctx().queue commandBuffer]; + MPSMatrixDescriptor* dq = [MPSMatrixDescriptor matrixDescriptorWithRows:T columns:D + rowBytes:(size_t)HQ*D*2 dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor* dk = [MPSMatrixDescriptor matrixDescriptorWithRows:S columns:D + rowBytes:(size_t)HKV*D*2 dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor* ds = [MPSMatrixDescriptor matrixDescriptorWithRows:T columns:S + rowBytes:(size_t)S*2 dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor* dv = [MPSMatrixDescriptor matrixDescriptorWithRows:S columns:DV + rowBytes:(size_t)HKV*DV*2 dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor* do2 = [MPSMatrixDescriptor matrixDescriptorWithRows:T columns:DV + rowBytes:(size_t)HQ*DV*2 dataType:MPSDataTypeFloat16]; + MPSMatrixMultiplication* mqk = mps_gemm_ab(T, S, D, true, scale, beta); + for (uint32_t h = 0; h < HQ; ++h) { + MPSMatrix* mq = [[MPSMatrix alloc] initWithBuffer:qb.first offset:qb.second + (size_t)h*D*2 descriptor:dq]; + MPSMatrix* mk = [[MPSMatrix alloc] initWithBuffer:kb.first offset:kb.second + (size_t)h*D*2 descriptor:dk]; + MPSMatrix* msc = [[MPSMatrix alloc] initWithBuffer:sbuf offset:(size_t)h*T*S*2 descriptor:ds]; + [mqk encodeToCommandBuffer:g_cmd leftMatrix:mq rightMatrix:mk resultMatrix:msc]; + } + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoSoftmaxR]; + [g_enc setBuffer:sbuf offset:0 atIndex:0]; + [g_enc setBuffer:sbuf offset:0 atIndex:1]; + uint32_t cols = S; + [g_enc setBytes:&cols length:4 atIndex:2]; + [g_enc setThreadgroupMemoryLength:256*4 atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake((size_t)HQ*T,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + [g_enc endEncoding]; g_enc = nil; + MPSMatrixMultiplication* mpv = mps_gemm(T, DV, S, false); + for (uint32_t h = 0; h < HQ; ++h) { + MPSMatrix* msc = [[MPSMatrix alloc] initWithBuffer:sbuf offset:(size_t)h*T*S*2 descriptor:ds]; + MPSMatrix* mv = [[MPSMatrix alloc] initWithBuffer:vb.first offset:vb.second + (size_t)h*DV*2 descriptor:dv]; + MPSMatrix* mo = [[MPSMatrix alloc] initWithBuffer:ob.first offset:ob.second + (size_t)h*DV*2 descriptor:do2]; + [mpv encodeToCommandBuffer:g_cmd leftMatrix:msc rightMatrix:mv resultMatrix:mo]; + } + return true; + } + } + } + ensureEncoder(); + struct { uint32_t T,S,HQ,HKV,D,DV; float scale; uint32_t causal,pos_off,window; float logit_cap; uint32_t mask_mode; } U = + { T,S,HQ,HKV,D,DV, scale, causal,pos_off,window, logit_cap, mask_mode }; + bool d64 = (D <= 64 && DV <= 64 && ctx().psoAttnD64); + [g_enc setComputePipelineState:(d64 ? ctx().psoAttnD64 : ctx().psoAttnDense)]; + setBufAt(q, (size_t)B*T*HQ*D*2, 0); setBufAt(k, (size_t)B*S*HKV*D*2, 1); + setBufAt(v, (size_t)B*S*HKV*DV*2, 2); setBufAt(out, (size_t)B*T*HQ*DV*2, 3); + if (mask_mode && mask) { + size_t mb = (mask_mode>=3u) ? (size_t)B*HQ*T*S*2 : (size_t)B*T*S*2; + setBufAt(mask, mb, 4); + } else [g_enc setBuffer:ctx().dummy offset:0 atIndex:4]; + [g_enc setBytes:&U length:sizeof(U) atIndex:5]; + [g_enc dispatchThreadgroups:MTLSizeMake(d64 ? (T+7)/8 : (T+3)/4,HQ,B) threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + return true; +} + +bool cactus_metal_encode_reduce_axis(int op, void* out, const void* in, uint32_t outer, + uint32_t axis_size, uint32_t inner, int f32) { + if (!ctx().ok || !ctx().psoReduceF16 || !ctx().psoReduceF32) return false; + ensureEncoder(); + size_t es = f32 ? 4 : 2; + uint32_t n = outer * inner; + [g_enc setComputePipelineState:(f32 ? ctx().psoReduceF32 : ctx().psoReduceF16)]; + setBufAt(in, (size_t)outer*axis_size*inner*es, 0); setBufAt(out, (size_t)n*es, 1); + [g_enc setBytes:&axis_size length:4 atIndex:2]; [g_enc setBytes:&inner length:4 atIndex:3]; + [g_enc setBytes:&n length:4 atIndex:4]; [g_enc setBytes:&op length:4 atIndex:5]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_cumsum(void* out, const void* in, uint32_t outer, + uint32_t axis_size, uint32_t inner, int f32) { + if (!ctx().ok || !ctx().psoCumsumF16 || !ctx().psoCumsumF32) return false; + ensureEncoder(); + size_t es = f32 ? 4 : 2; + uint32_t n = outer * inner; + size_t bytes = (size_t)outer*axis_size*inner*es; + [g_enc setComputePipelineState:(f32 ? ctx().psoCumsumF32 : ctx().psoCumsumF16)]; + setBufAt(in, bytes, 0); setBufAt(out, bytes, 1); + [g_enc setBytes:&axis_size length:4 atIndex:2]; [g_enc setBytes:&inner length:4 atIndex:3]; + [g_enc setBytes:&n length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_concat2(void* out, const void* a, const void* b, + uint32_t a_outer, uint32_t b_outer, + uint32_t a_axis, uint32_t b_axis, uint32_t inner) { + if (!ctx().ok || !ctx().psoConcat2 || a_outer == 0 || b_outer == 0) return false; + ensureEncoder(); + uint32_t outer = a_outer > b_outer ? a_outer : b_outer; + uint32_t n = outer * (a_axis + b_axis) * inner; + [g_enc setComputePipelineState:ctx().psoConcat2]; + setBufAt(a, (size_t)a_outer*a_axis*inner*2, 0); setBufAt(b, (size_t)b_outer*b_axis*inner*2, 1); + setBufAt(out, (size_t)n*2, 2); + [g_enc setBytes:&a_axis length:4 atIndex:3]; [g_enc setBytes:&b_axis length:4 atIndex:4]; + [g_enc setBytes:&inner length:4 atIndex:5]; [g_enc setBytes:&n length:4 atIndex:6]; + [g_enc setBytes:&a_outer length:4 atIndex:7]; [g_enc setBytes:&b_outer length:4 atIndex:8]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_gather_f32idx(void* out, const void* table, const void* idx, + uint32_t rows, uint32_t D, size_t table_bytes) { + if (!ctx().ok || !ctx().psoGatherF32Idx) return false; + ensureEncoder(); + uint32_t n = rows * D; + [g_enc setComputePipelineState:ctx().psoGatherF32Idx]; + setBufAt(table, table_bytes, 0); setBufAt(idx, (size_t)rows*4, 1); setBufAt(out, (size_t)n*2, 2); + [g_enc setBytes:&D length:4 atIndex:3]; [g_enc setBytes:&n length:4 atIndex:4]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_rope_full(void* out, const void* in, uint32_t tokens, uint32_t S, + uint32_t H, uint32_t D, uint32_t rot, uint32_t pos0, + float theta, int gptj) { + if (!ctx().ok || !ctx().psoRopeFull) return false; + ensureEncoder(); + uint32_t gj = gptj ? 1u : 0u; + [g_enc setComputePipelineState:ctx().psoRopeFull]; + setBufAt(in, (size_t)tokens*D*2, 0); setBufAt(out, (size_t)tokens*D*2, 1); + [g_enc setBytes:&S length:4 atIndex:2]; [g_enc setBytes:&H length:4 atIndex:3]; + [g_enc setBytes:&D length:4 atIndex:4]; [g_enc setBytes:&rot length:4 atIndex:5]; + [g_enc setBytes:&pos0 length:4 atIndex:6]; [g_enc setBytes:&theta length:4 atIndex:7]; + [g_enc setBytes:&gj length:4 atIndex:8]; + uint32_t gx = rot/2 + (D > rot ? D - rot : 0); + if (gx == 0) gx = 1; + [g_enc dispatchThreads:MTLSizeMake(gx, tokens, 1) threadsPerThreadgroup:MTLSizeMake(std::min(gx, 64u), 4, 1)]; + return true; +} +bool cactus_metal_encode_maxpool1d(void* out, const void* in, uint32_t NC, uint32_t L, + uint32_t Lout, uint32_t K, uint32_t stride) { + if (!ctx().ok || !ctx().psoMaxpool1d) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoMaxpool1d]; + setBufAt(in, (size_t)NC*L*2, 0); setBufAt(out, (size_t)NC*Lout*2, 1); + [g_enc setBytes:&L length:4 atIndex:2]; [g_enc setBytes:&Lout length:4 atIndex:3]; + [g_enc setBytes:&K length:4 atIndex:4]; [g_enc setBytes:&stride length:4 atIndex:5]; + [g_enc dispatchThreads:MTLSizeMake(Lout, NC, 1) threadsPerThreadgroup:MTLSizeMake(std::min(Lout, 64u), 4, 1)]; + return true; +} +bool cactus_metal_encode_bilinear(void* out, const void* in, uint32_t sh, uint32_t sw, + uint32_t dh, uint32_t dw, uint32_t E, int align) { + if (!ctx().ok || !ctx().psoBilinear) return false; + ensureEncoder(); + uint32_t al = align ? 1u : 0u; + [g_enc setComputePipelineState:ctx().psoBilinear]; + setBufAt(in, (size_t)sh*sw*E*2, 0); setBufAt(out, (size_t)dh*dw*E*2, 1); + [g_enc setBytes:&sh length:4 atIndex:2]; [g_enc setBytes:&sw length:4 atIndex:3]; + [g_enc setBytes:&dh length:4 atIndex:4]; [g_enc setBytes:&dw length:4 atIndex:5]; + [g_enc setBytes:&E length:4 atIndex:6]; [g_enc setBytes:&al length:4 atIndex:7]; + [g_enc dispatchThreads:MTLSizeMake(E, dh*dw, 1) threadsPerThreadgroup:MTLSizeMake(std::min(E, 64u), 4, 1)]; + return true; +} +bool cactus_metal_encode_conv1d_gen(void* out, const void* x, const void* w, const void* bias, + uint32_t N, uint32_t Cin, uint32_t L, uint32_t Cout, + uint32_t Lout, uint32_t K, uint32_t stride, int w_ck_co) { + if (!ctx().ok || !ctx().psoConv1dGen) return false; + ensureEncoder(); + uint32_t hb = bias ? 1u : 0u, wl = w_ck_co ? 1u : 0u; + [g_enc setComputePipelineState:ctx().psoConv1dGen]; + setBufAt(x, (size_t)N*Cin*L*2, 0); setBufAt(w, (size_t)Cout*Cin*K*2, 1); + if (bias) setBufAt(bias, (size_t)Cout*2, 2); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + setBufAt(out, (size_t)N*Cout*Lout*2, 3); + [g_enc setBytes:&Cin length:4 atIndex:4]; [g_enc setBytes:&L length:4 atIndex:5]; + [g_enc setBytes:&Cout length:4 atIndex:6]; [g_enc setBytes:&Lout length:4 atIndex:7]; + [g_enc setBytes:&K length:4 atIndex:8]; [g_enc setBytes:&stride length:4 atIndex:9]; + [g_enc setBytes:&hb length:4 atIndex:10]; [g_enc setBytes:&wl length:4 atIndex:11]; + [g_enc dispatchThreads:MTLSizeMake(Lout, Cout, N) threadsPerThreadgroup:MTLSizeMake(std::min(Lout, 64u), 4, 1)]; + return true; +} +bool cactus_metal_encode_conv1d_nlc_dw(void* out, const void* x, const void* w, const void* bias, + uint32_t N, uint32_t L, uint32_t C, uint32_t K, + uint32_t dil, uint32_t pad) { + if (!ctx().ok || !ctx().psoConv1dNlcDw) return false; + ensureEncoder(); + uint32_t hb = bias ? 1u : 0u; + [g_enc setComputePipelineState:ctx().psoConv1dNlcDw]; + setBufAt(x, (size_t)N*L*C*2, 0); setBufAt(w, (size_t)C*K*2, 1); + if (bias) setBufAt(bias, (size_t)C*2, 2); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + setBufAt(out, (size_t)N*L*C*2, 3); + [g_enc setBytes:&L length:4 atIndex:4]; [g_enc setBytes:&C length:4 atIndex:5]; + [g_enc setBytes:&K length:4 atIndex:6]; [g_enc setBytes:&dil length:4 atIndex:7]; + [g_enc setBytes:&pad length:4 atIndex:8]; [g_enc setBytes:&hb length:4 atIndex:9]; + [g_enc dispatchThreads:MTLSizeMake(C, L, N) threadsPerThreadgroup:MTLSizeMake(std::min(C, 64u), 4, 1)]; + return true; +} +bool cactus_metal_encode_conv2d(void* out, const void* x, const void* w, const void* bias, + uint32_t N, uint32_t Cin, uint32_t H, uint32_t W, uint32_t Cout, + uint32_t Ho, uint32_t Wo, uint32_t K, uint32_t stride, + uint32_t pad, int dw) { + if (!ctx().ok || !ctx().psoConv2d) return false; + ensureEncoder(); + uint32_t hb = bias ? 1u : 0u, dwf = dw ? 1u : 0u; + size_t wbytes = dw ? (size_t)Cout*K*K*2 : (size_t)Cout*Cin*K*K*2; + [g_enc setComputePipelineState:ctx().psoConv2d]; + setBufAt(x, (size_t)N*Cin*H*W*2, 0); setBufAt(w, wbytes, 1); + if (bias) setBufAt(bias, (size_t)Cout*2, 2); else [g_enc setBuffer:ctx().dummy offset:0 atIndex:2]; + setBufAt(out, (size_t)N*Cout*Ho*Wo*2, 3); + [g_enc setBytes:&Cin length:4 atIndex:4]; [g_enc setBytes:&H length:4 atIndex:5]; + [g_enc setBytes:&W length:4 atIndex:6]; [g_enc setBytes:&Cout length:4 atIndex:7]; + [g_enc setBytes:&Ho length:4 atIndex:8]; [g_enc setBytes:&Wo length:4 atIndex:9]; + [g_enc setBytes:&K length:4 atIndex:10]; [g_enc setBytes:&stride length:4 atIndex:11]; + [g_enc setBytes:&pad length:4 atIndex:12]; [g_enc setBytes:&dwf length:4 atIndex:13]; + [g_enc setBytes:&hb length:4 atIndex:14]; + uint32_t wq = (Wo + 3) / 4; + [g_enc dispatchThreads:MTLSizeMake(wq, Ho, (size_t)N*Cout) threadsPerThreadgroup:MTLSizeMake(std::min(wq, 32u), std::min(Ho, 8u), 1)]; + return true; +} +bool cactus_metal_encode_batchnorm(void* out, const void* x, const void* w, const void* b, + const void* rm, const void* rv, uint32_t C, uint32_t inner, + uint32_t total, float eps) { + if (!ctx().ok || !ctx().psoBatchnorm) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoBatchnorm]; + setBufAt(x, (size_t)total*2, 0); setBufAt(w, (size_t)C*2, 1); setBufAt(b, (size_t)C*2, 2); + setBufAt(rm, (size_t)C*2, 3); setBufAt(rv, (size_t)C*2, 4); setBufAt(out, (size_t)total*2, 5); + [g_enc setBytes:&C length:4 atIndex:6]; [g_enc setBytes:&inner length:4 atIndex:7]; + [g_enc setBytes:&eps length:4 atIndex:8]; [g_enc setBytes:&total length:4 atIndex:9]; + [g_enc dispatchThreads:MTLSizeMake(total,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_groupnorm(void* out, const void* x, const void* w, const void* b, + uint32_t N, uint32_t C, uint32_t S, uint32_t groups, float eps) { + if (!ctx().ok || !ctx().psoGroupnorm || groups == 0 || C % groups != 0) return false; + ensureEncoder(); + uint32_t cpg = C / groups; + [g_enc setComputePipelineState:ctx().psoGroupnorm]; + setBufAt(x, (size_t)N*C*S*2, 0); setBufAt(w, (size_t)C*2, 1); setBufAt(b, (size_t)C*2, 2); + setBufAt(out, (size_t)N*C*S*2, 3); + [g_enc setBytes:&cpg length:4 atIndex:4]; [g_enc setBytes:&S length:4 atIndex:5]; + [g_enc setBytes:&C length:4 atIndex:6]; [g_enc setBytes:&eps length:4 atIndex:7]; + [g_enc setThreadgroupMemoryLength:256*4 atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(groups, N, 1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_bias_add_rows(void* y, const void* bias, uint32_t C, uint32_t total) { + if (!ctx().ok || !ctx().psoBiasAddRows) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoBiasAddRows]; + setBufAt(y, (size_t)total*2, 0); setBufAt(bias, (size_t)C*2, 1); + [g_enc setBytes:&C length:4 atIndex:2]; [g_enc setBytes:&total length:4 atIndex:3]; + [g_enc dispatchThreads:MTLSizeMake(total,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_elemwise_chain(void* out, const void* in, const float* steps, + uint32_t nsteps, const void* side0, const void* side1, + const void* side2, const size_t* side_elems, + size_t n, uint32_t flags, uint32_t inner) { + if (!ctx().ok || !ctx().psoEwChain || nsteps == 0 || nsteps > 12) return false; + ensureEncoder(); + uint32_t ns = nsteps, nn = (uint32_t)n; + size_t ein = (flags & 1u) ? 4 : 2, eout = (flags & 2u) ? 4 : 2; + [g_enc setComputePipelineState:ctx().psoEwChain]; + setBufAt(in, n*ein, 0); setBufAt(out, n*eout, 1); + [g_enc setBytes:steps length:nsteps*16 atIndex:2]; + [g_enc setBytes:&ns length:4 atIndex:3]; + setBufAt(side0 ? side0 : in, (side0 ? side_elems[0] : n)*((flags & 4u) ? 4 : 2), 4); + setBufAt(side1 ? side1 : in, (side1 ? side_elems[1] : n)*((flags & 8u) ? 4 : 2), 5); + setBufAt(side2 ? side2 : in, (side2 ? side_elems[2] : n)*((flags & 16u) ? 4 : 2), 6); + [g_enc setBytes:&nn length:4 atIndex:7]; + [g_enc setBytes:&flags length:4 atIndex:8]; + uint32_t innr = inner; + [g_enc setBytes:&innr length:4 atIndex:9]; + [g_enc dispatchThreads:MTLSizeMake((nn+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_rms_norm_add_rows(void* ysum, void* ynorm, const void* x, const void* res, + const void* w, uint32_t rows, uint32_t dim, float eps, + int clipped) { + if (!ctx().ok || !ctx().psoRmsAddSimd) return false; + ensureEncoder(); + uint32_t cl = clipped ? 1u : 0u; + [g_enc setComputePipelineState:ctx().psoRmsAddSimd]; + setBufAt(x, (size_t)rows*dim*2, 0); setBufAt(res, (size_t)rows*dim*2, 1); + setBufAt(w, (size_t)dim*2, 2); setBufAt(ysum, (size_t)rows*dim*2, 3); + setBufAt(ynorm, (size_t)rows*dim*2, 4); + [g_enc setBytes:&dim length:4 atIndex:5]; [g_enc setBytes:&eps length:4 atIndex:6]; + [g_enc setBytes:&rows length:4 atIndex:7]; [g_enc setBytes:&cl length:4 atIndex:8]; + [g_enc dispatchThreadgroups:MTLSizeMake((rows+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} +bool cactus_metal_encode_gemm_batch(void* out, const void* a, const void* b, + uint32_t M, uint32_t K, uint32_t N, uint32_t batch, int f32out, int f32a) { + if (!ctx().ok || !ctx().psoGemmBatch || !ctx().psoGemmBatchF32A) return false; + ensureEncoder(); + uint32_t fo = f32out ? 1u : 0u; + [g_enc setComputePipelineState:(f32a ? ctx().psoGemmBatchF32A : ctx().psoGemmBatch)]; + setBufAt(a, (size_t)batch*M*K*(f32a?4:2), 0); setBufAt(b, (size_t)batch*K*N*2, 1); + setBufAt(out, (size_t)batch*M*N*(f32out?4:2), 2); + [g_enc setBytes:&M length:4 atIndex:3]; [g_enc setBytes:&K length:4 atIndex:4]; + [g_enc setBytes:&N length:4 atIndex:5]; [g_enc setBytes:&fo length:4 atIndex:6]; + [g_enc dispatchThreads:MTLSizeMake((size_t)M*N, batch, 1) threadsPerThreadgroup:MTLSizeMake(64,4,1)]; + return true; +} +bool cactus_metal_encode_conv1d_dw(void* out, const void* x, const void* w, + uint32_t C, uint32_t L, uint32_t Lout, uint32_t K, uint32_t stride) { + if (!ctx().ok || !ctx().psoConvDw) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoConvDw]; + setBufAt(x, (size_t)C*L*2, 0); setBufAt(w, (size_t)C*K*2, 1); setBufAt(out, (size_t)C*Lout*2, 2); + [g_enc setBytes:&L length:4 atIndex:3]; [g_enc setBytes:&Lout length:4 atIndex:4]; + [g_enc setBytes:&K length:4 atIndex:5]; [g_enc setBytes:&stride length:4 atIndex:6]; + [g_enc dispatchThreads:MTLSizeMake(Lout, C, 1) threadsPerThreadgroup:MTLSizeMake(64,4,1)]; + return true; +} +bool cactus_metal_encode_transpose2d(void* out, const void* in, uint32_t batch, uint32_t R, uint32_t C) { + if (!ctx().ok || !ctx().psoTranspose2d) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoTranspose2d]; + setBufAt(in, (size_t)batch*R*C*2, 0); setBufAt(out, (size_t)batch*R*C*2, 1); + [g_enc setBytes:&R length:4 atIndex:2]; [g_enc setBytes:&C length:4 atIndex:3]; + [g_enc dispatchThreadgroups:MTLSizeMake((C+31)/32, (R+31)/32, batch) threadsPerThreadgroup:MTLSizeMake(32,8,1)]; + return true; +} +bool cactus_metal_encode_strided_copy(void* out, const void* in, const uint32_t* oshape, + const uint32_t* sstride, uint32_t ndim, uint32_t total, uint32_t base, size_t in_bytes, size_t out_bytes) { + if (!ctx().ok || ndim == 0 || ndim > 8) return false; + ensureEncoder(); + if (ndim >= 2 && sstride[ndim-1] == 1 && (oshape[ndim-1] % 4) == 0 && ctx().psoStridedRows) { + uint32_t inner = oshape[ndim-1]; + uint32_t inner4 = inner / 4; + uint32_t nd2 = ndim - 1; + uint32_t rows = total / inner; + [g_enc setComputePipelineState:ctx().psoStridedRows]; + setBufAt(in, in_bytes, 0); setBufAt(out, out_bytes, 1); + [g_enc setBytes:oshape length:nd2*4 atIndex:2]; [g_enc setBytes:sstride length:nd2*4 atIndex:3]; + [g_enc setBytes:&nd2 length:4 atIndex:4]; [g_enc setBytes:&rows length:4 atIndex:5]; + [g_enc setBytes:&base length:4 atIndex:6]; [g_enc setBytes:&inner4 length:4 atIndex:7]; + [g_enc dispatchThreads:MTLSizeMake(inner4, rows, 1) threadsPerThreadgroup:MTLSizeMake(std::min(inner4, 64u), 4, 1)]; + return true; + } + [g_enc setComputePipelineState:ctx().psoStrided]; + setBufAt(in, in_bytes, 0); setBufAt(out, out_bytes, 1); + [g_enc setBytes:oshape length:ndim*4 atIndex:2]; [g_enc setBytes:sstride length:ndim*4 atIndex:3]; + [g_enc setBytes:&ndim length:4 atIndex:4]; [g_enc setBytes:&total length:4 atIndex:5]; [g_enc setBytes:&base length:4 atIndex:6]; + [g_enc dispatchThreads:MTLSizeMake(total,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_bcast_binary(int op, void* out, const void* a, const void* b, + const uint32_t* oshape, const uint32_t* astride, const uint32_t* bstride, uint32_t ndim, uint32_t total, + size_t a_bytes, size_t b_bytes, size_t out_bytes) { + if (!ctx().ok || ndim == 0 || ndim > 8) return false; + ensureEncoder(); + int o = op; + if (ndim >= 2 && oshape[ndim-1] > 1 && ctx().psoBcastRows) { + uint32_t inner = oshape[ndim-1]; + uint32_t ainner = astride[ndim-1], binner = bstride[ndim-1]; + uint32_t nd2 = ndim - 1; + uint32_t rows = total / inner; + [g_enc setComputePipelineState:ctx().psoBcastRows]; + setBufAt(a, a_bytes, 0); setBufAt(b, b_bytes, 1); setBufAt(out, out_bytes, 2); + [g_enc setBytes:oshape length:nd2*4 atIndex:3]; [g_enc setBytes:astride length:nd2*4 atIndex:4]; + [g_enc setBytes:bstride length:nd2*4 atIndex:5]; + [g_enc setBytes:&nd2 length:4 atIndex:6]; [g_enc setBytes:&rows length:4 atIndex:7]; + [g_enc setBytes:&o length:4 atIndex:8]; [g_enc setBytes:&inner length:4 atIndex:9]; + [g_enc setBytes:&ainner length:4 atIndex:10]; [g_enc setBytes:&binner length:4 atIndex:11]; + [g_enc dispatchThreads:MTLSizeMake(inner, rows, 1) threadsPerThreadgroup:MTLSizeMake(std::min(inner, 64u), 4, 1)]; + return true; + } + [g_enc setComputePipelineState:ctx().psoBcast]; + setBufAt(a, a_bytes, 0); setBufAt(b, b_bytes, 1); setBufAt(out, out_bytes, 2); + [g_enc setBytes:oshape length:ndim*4 atIndex:3]; [g_enc setBytes:astride length:ndim*4 atIndex:4]; [g_enc setBytes:bstride length:ndim*4 atIndex:5]; + [g_enc setBytes:&ndim length:4 atIndex:6]; [g_enc setBytes:&total length:4 atIndex:7]; [g_enc setBytes:&o length:4 atIndex:8]; + [g_enc dispatchThreads:MTLSizeMake(total,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +bool cactus_metal_encode_strided_scatter(void* out, const void* in, const uint32_t* ishape, + const uint32_t* ostride, uint32_t ndim, uint32_t total, uint32_t base, size_t in_bytes, size_t out_bytes) { + if (!ctx().ok || ndim == 0 || ndim > 8) return false; + ensureEncoder(); + if (ndim >= 2 && ostride[ndim-1] == 1 && (ishape[ndim-1] % 4) == 0 && ctx().psoScatterRows) { + uint32_t inner = ishape[ndim-1]; + uint32_t inner4 = inner / 4; + uint32_t nd2 = ndim - 1; + uint32_t rows = total / inner; + [g_enc setComputePipelineState:ctx().psoScatterRows]; + setBufAt(in, in_bytes, 0); setBufAt(out, out_bytes, 1); + [g_enc setBytes:ishape length:nd2*4 atIndex:2]; [g_enc setBytes:ostride length:nd2*4 atIndex:3]; + [g_enc setBytes:&nd2 length:4 atIndex:4]; [g_enc setBytes:&rows length:4 atIndex:5]; + [g_enc setBytes:&base length:4 atIndex:6]; [g_enc setBytes:&inner4 length:4 atIndex:7]; + [g_enc dispatchThreads:MTLSizeMake(inner4, rows, 1) threadsPerThreadgroup:MTLSizeMake(std::min(inner4, 64u), 4, 1)]; + return true; + } + [g_enc setComputePipelineState:ctx().psoScatter]; + setBufAt(in, in_bytes, 0); setBufAt(out, out_bytes, 1); + [g_enc setBytes:ishape length:ndim*4 atIndex:2]; [g_enc setBytes:ostride length:ndim*4 atIndex:3]; + [g_enc setBytes:&ndim length:4 atIndex:4]; [g_enc setBytes:&total length:4 atIndex:5]; [g_enc setBytes:&base length:4 atIndex:6]; + [g_enc dispatchThreads:MTLSizeMake(total,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} +static bool encode_kv_append(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, uint32_t M, + size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + if (!ctx().ok) return false; + uint32_t num_groups = (hdim + group_size - 1)/group_size, n = M*kv_heads*num_groups; + ensureEncoder(); + [g_enc setComputePipelineState:(M == 1 ? ctx().psoKvAppend : ctx().psoKvAppendM)]; + setBufAt(src, src_bytes, 0); setBufAt(int8base, int8_bytes, 1); setBufAt(scalebase, scale_bytes, 2); + [g_enc setBytes:&kv_heads length:4 atIndex:3]; [g_enc setBytes:&hdim length:4 atIndex:4]; + [g_enc setBytes:¤t_len length:4 atIndex:5]; [g_enc setBytes:&group_size length:4 atIndex:6]; + if (M != 1) [g_enc setBytes:&M length:4 atIndex:7]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(n<256?n:256,1,1)]; + return true; +} + +bool cactus_metal_encode_kv_append_i8(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, + size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + return encode_kv_append(src, int8base, scalebase, kv_heads, hdim, current_len, group_size, 1, + src_bytes, int8_bytes, scale_bytes); +} + +bool cactus_metal_encode_kv_append_i8_m(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, uint32_t M, + size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + return encode_kv_append(src, int8base, scalebase, kv_heads, hdim, current_len, group_size, M, + src_bytes, int8_bytes, scale_bytes); +} + +bool cactus_metal_encode_gemv_bias(void* out, const void* x, const void* w, const void* bias, + uint32_t K, uint32_t N, int pretransposed) { + if (!ctx().ok || !ctx().psoGemvBias || !bias) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoGemvBias]; + setBufAt(x, (size_t)K*2, 0); + setBufAt(w, (size_t)K*N*2, 1); + setBufAt(bias, (size_t)N*2, 2); + setBufAt(out, (size_t)N*2, 3); + uint32_t trv = pretransposed ? 1u : 0u; + [g_enc setBytes:&K length:4 atIndex:4]; [g_enc setBytes:&N length:4 atIndex:5]; + [g_enc setBytes:&trv length:4 atIndex:6]; + [g_enc dispatchThreadgroups:MTLSizeMake((N+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} + +bool cactus_metal_encode_rel_pos_bias(void* y, const void* q, const void* r, + uint32_t B, uint32_t T, uint32_t H, uint32_t D, uint32_t R, int r_batched, float scale) { + if (!ctx().ok || !ctx().psoRelPosBias) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoRelPosBias]; + setBufAt(q, (size_t)B*T*H*D*2, 0); + setBufAt(r, (size_t)(r_batched?B:1)*R*H*D*2, 1); + setBufAt(y, (size_t)B*H*T*T*2, 2); + uint32_t rbs = r_batched ? R*H*D : 0u; + [g_enc setBytes:&T length:4 atIndex:3]; [g_enc setBytes:&H length:4 atIndex:4]; + [g_enc setBytes:&D length:4 atIndex:5]; [g_enc setBytes:&rbs length:4 atIndex:6]; + [g_enc setBytes:&scale length:4 atIndex:7]; + [g_enc dispatchThreads:MTLSizeMake((size_t)T*T, H, B) + threadsPerThreadgroup:MTLSizeMake(64, H < 4 ? H : 4, 1)]; + return true; +} + +bool cactus_metal_encode_conv_cache_append(void* out, const void* src, void* ring, + uint32_t hd, uint32_t ws, uint32_t nnew, uint32_t head0, uint32_t count_new, + uint32_t num_rows, int src_f32) { + if (!ctx().ok || !ctx().psoConvCacheAppend) return false; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoConvCacheAppend]; + setBufAt(src, (size_t)num_rows*hd*(src_f32?4:2), 0); + setBufAt(ring, (size_t)ws*hd*2, 1); + setBufAt(out, (size_t)ws*hd*2, 2); + uint32_t f32 = src_f32 ? 1u : 0u; + [g_enc setBytes:&hd length:4 atIndex:3]; [g_enc setBytes:&ws length:4 atIndex:4]; + [g_enc setBytes:&nnew length:4 atIndex:5]; [g_enc setBytes:&head0 length:4 atIndex:6]; + [g_enc setBytes:&count_new length:4 atIndex:7]; [g_enc setBytes:&num_rows length:4 atIndex:8]; + [g_enc setBytes:&f32 length:4 atIndex:9]; + [g_enc dispatchThreads:MTLSizeMake(hd, ws + nnew, 1) + threadsPerThreadgroup:MTLSizeMake(hd < 64 ? hd : 64, 4, 1)]; + return true; +} + +bool cactus_metal_encode_kv_append_ring_i8_m(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t current_len, uint32_t group_size, uint32_t M, + uint32_t sink, uint32_t W, size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + if (!ctx().ok) return false; + uint32_t num_groups = (hdim + group_size - 1)/group_size, n = M*kv_heads*num_groups; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoKvAppendRingM]; + setBufAt(src, src_bytes, 0); setBufAt(int8base, int8_bytes, 1); setBufAt(scalebase, scale_bytes, 2); + [g_enc setBytes:&kv_heads length:4 atIndex:3]; [g_enc setBytes:&hdim length:4 atIndex:4]; + [g_enc setBytes:¤t_len length:4 atIndex:5]; [g_enc setBytes:&group_size length:4 atIndex:6]; + [g_enc setBytes:&M length:4 atIndex:7]; + [g_enc setBytes:&sink length:4 atIndex:8]; [g_enc setBytes:&W length:4 atIndex:9]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(n<256?n:256,1,1)]; + return true; +} + +static bool encode_kv_append_sliding(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t keep_sink, uint32_t remaining, uint32_t shift_src, + uint32_t group_size, uint32_t M, size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + if (!ctx().ok) return false; + uint32_t num_groups = (hdim + group_size - 1)/group_size, per = kv_heads*num_groups; + id scrI8 = recycled((size_t)(remaining?remaining:1)*kv_heads*hdim); + id scrSc = recycled((size_t)(remaining?remaining:1)*kv_heads*num_groups*sizeof(float)); + ensureEncoder(); + if (remaining > 0) { + uint32_t n = remaining*per; + [g_enc setComputePipelineState:ctx().psoSlideS]; + setBufAt(int8base, int8_bytes, 0); setBufAt(scalebase, scale_bytes, 1); + [g_enc setBuffer:scrI8 offset:0 atIndex:2]; [g_enc setBuffer:scrSc offset:0 atIndex:3]; + [g_enc setBytes:&kv_heads length:4 atIndex:4]; [g_enc setBytes:&hdim length:4 atIndex:5]; + [g_enc setBytes:&group_size length:4 atIndex:6]; [g_enc setBytes:&shift_src length:4 atIndex:7]; + [g_enc setBytes:&remaining length:4 atIndex:8]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(n<256?n:256,1,1)]; + } + uint32_t n2 = (remaining+M)*per; + [g_enc setComputePipelineState:(M == 1 ? ctx().psoSlideR : ctx().psoSlideRM)]; + setBufAt(src, src_bytes, 0); setBufAt(int8base, int8_bytes, 1); setBufAt(scalebase, scale_bytes, 2); + [g_enc setBuffer:scrI8 offset:0 atIndex:3]; [g_enc setBuffer:scrSc offset:0 atIndex:4]; + [g_enc setBytes:&kv_heads length:4 atIndex:5]; [g_enc setBytes:&hdim length:4 atIndex:6]; + [g_enc setBytes:&group_size length:4 atIndex:7]; [g_enc setBytes:&keep_sink length:4 atIndex:8]; + [g_enc setBytes:&remaining length:4 atIndex:9]; + if (M != 1) [g_enc setBytes:&M length:4 atIndex:10]; + [g_enc dispatchThreads:MTLSizeMake(n2,1,1) threadsPerThreadgroup:MTLSizeMake(n2<256?n2:256,1,1)]; + return true; +} + +bool cactus_metal_encode_kv_append_sliding_i8(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t keep_sink, uint32_t remaining, uint32_t shift_src, + uint32_t group_size, size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + return encode_kv_append_sliding(src, int8base, scalebase, kv_heads, hdim, keep_sink, remaining, + shift_src, group_size, 1, src_bytes, int8_bytes, scale_bytes); +} + +bool cactus_metal_encode_kv_append_sliding_i8_m(const void* src, void* int8base, void* scalebase, + uint32_t kv_heads, uint32_t hdim, uint32_t keep_sink, uint32_t remaining, uint32_t shift_src, + uint32_t group_size, uint32_t M, size_t src_bytes, size_t int8_bytes, size_t scale_bytes) { + return encode_kv_append_sliding(src, int8base, scalebase, kv_heads, hdim, keep_sink, remaining, + shift_src, group_size, M, src_bytes, int8_bytes, scale_bytes); +} + +void cactus_metal_quant_matmul(const CactusQuantMatrix* W, const __fp16* A, + uint32_t M, __fp16* C) { + const uint32_t gs = W->group_size; + const bool fast = ctx().ok && M == 1 && quant_fast_eligible(W); + + if (!fast) { cactus_quant_matmul(W, A, M, C); return; } + + @autoreleasepool { + const uint32_t K = W->K, N = W->N, ng = W->num_groups; + const uint32_t pgb = (gs * 4u + 7u) / 8u; + + ResW& rw = resident(W); + if (!rw.ok) { cactus_quant_matmul(W, A, M, C); return; } + id bx = buf(A, (size_t)K * sizeof(__fp16)); + id brecip = rw.recip; + id bls = rw.lsign; + id brs = rw.rsign; + id bperm = rw.perm; + id bpk = rw.packed; + id bcb = rw.codebook; + id bnorm = rw.norms; + id bcode = owned_shared((size_t)ng * gs * sizeof(__fp16)); + id by = owned_shared((size_t)N * sizeof(__fp16)); + + id cmd = [ctx().queue commandBuffer]; + id e = [cmd computeCommandEncoder]; + + bool simdT = (gs==128u && ctx().psoTsimd); + [e setComputePipelineState:(simdT?ctx().psoTsimd:ctx().psoT)]; + [e setBuffer:bx offset:0 atIndex:0]; [e setBuffer:brecip offset:rw.rc_off atIndex:1]; + [e setBuffer:bls offset:rw.ls_off atIndex:2]; [e setBuffer:brs offset:rw.rs_off atIndex:3]; + [e setBuffer:bperm offset:rw.pm_off atIndex:4]; [e setBuffer:bcode offset:0 atIndex:5]; + [e setBytes:&gs length:sizeof(uint32_t) atIndex:6]; + [e setThreadgroupMemoryLength:gs * sizeof(float) atIndex:0]; + [e dispatchThreadgroups:MTLSizeMake(ng,1,1) threadsPerThreadgroup:MTLSizeMake(simdT?32u:(gs>1024u?1024u:gs),1,1)]; + + [e setComputePipelineState:ctx().psoG]; + [e setBuffer:bcode offset:0 atIndex:0]; [e setBuffer:bpk offset:rw.packed_off atIndex:1]; + [e setBuffer:bcb offset:rw.cb_off atIndex:2]; [e setBuffer:bnorm offset:rw.norms_off atIndex:3]; + [e setBuffer:by offset:0 atIndex:4]; + [e setBytes:&gs length:sizeof(uint32_t) atIndex:5]; + [e setBytes:&ng length:sizeof(uint32_t) atIndex:6]; + [e setBytes:&pgb length:sizeof(uint32_t) atIndex:7]; + [e setBytes:&N length:sizeof(uint32_t) atIndex:8]; + [e setBytes:&rw.il length:sizeof(uint32_t) atIndex:9]; + const uint32_t ROWS = 8; + [e dispatchThreadgroups:MTLSizeMake((N+ROWS-1)/ROWS,1,1) + threadsPerThreadgroup:MTLSizeMake(ROWS*32,1,1)]; + + [e endEncoding]; + [cmd commit]; + [cmd waitUntilCompleted]; + + std::memcpy(C, [by contents], (size_t)N * sizeof(__fp16)); + } +} + +#endif diff --git a/cactus-kernels/src/metal_backend_stub.cpp b/cactus-kernels/src/metal_backend_stub.cpp new file mode 100644 index 000000000..9ecafc05b --- /dev/null +++ b/cactus-kernels/src/metal_backend_stub.cpp @@ -0,0 +1,136 @@ + +#include "metal_backend.h" + +#if !CACTUS_HAS_METAL +bool cactus_metal_available() { return false; } +void cactus_metal_set_active(bool) {} +bool cactus_metal_active_mode() { return false; } +bool cactus_metal_encode_softcap(void*, const void*, size_t, float) { return false; } +bool cactus_metal_encode_adjust_logits(void*, size_t, const uint32_t*, uint32_t, + int64_t, float) { return false; } +bool cactus_metal_encode_transform_gemv(void*, const void*, const CactusQuantMatrix*, const void*) { return false; } +bool cactus_metal_transform_gemv_fits(uint32_t) { return false; } +bool cactus_metal_encode_gemv_cat(void* const*, const void* const*, + const CactusQuantMatrix* const*, int) { return false; } +bool cactus_metal_encode_swiglu_transform(void*, const void*, const void*, + const CactusQuantMatrix*, float) { return false; } +void cactus_metal_quant_matmul(const CactusQuantMatrix* W, const __fp16* A, + uint32_t M, __fp16* C) { + cactus_quant_matmul(W, A, M, C); +} +void cactus_metal_session_begin() {} +void cactus_metal_session_sync() {} +void cactus_metal_session_flush() {} +void cactus_metal_invalidate_host_wraps() {} +void cactus_metal_trim_prefill_cache() {} +void cactus_metal_session_end() {} +void* cactus_metal_alloc_shared(size_t) { return nullptr; } +void* cactus_metal_alloc_pooled(size_t) { return nullptr; } +void cactus_metal_free_shared(void*) {} +bool cactus_metal_encode_copy(void*, const void*, size_t) { return false; } +bool cactus_metal_encode_binary(int, void*, const void*, const void*, size_t) { return false; } +bool cactus_metal_encode_scalar(int, void*, const void*, size_t, float) { return false; } +bool cactus_metal_encode_unary(int, void*, const void*, size_t) { return false; } +bool cactus_metal_encode_swiglu(void*, const void*, const void*, size_t, float) { return false; } +bool cactus_metal_encode_rms_norm(void*, const void*, const void*, size_t, size_t, float) { return false; } +bool cactus_metal_encode_rms_norm_add(void*, const void*, const void*, const void*, size_t, size_t, float) { return false; } +bool cactus_metal_encode_rms_norm_add_rms(void*, void*, const void*, const void*, const void*, const void*, + size_t, size_t, float, float) { return false; } +bool cactus_metal_encode_rms_norm_add_scale(void*, const void*, const void*, const void*, size_t, size_t, float, float) { return false; } +bool cactus_metal_encode_argmax(const void*, uint32_t, void*, const void*) { return false; } +bool cactus_metal_encode_cast(void*, int, const void*, int, size_t) { return false; } +bool cactus_metal_encode_quant_matmul(void*, const void*, const CactusQuantMatrix*) { return false; } +bool cactus_metal_encode_quant_matmul_m(void*, const void*, const CactusQuantMatrix*, uint32_t) { return false; } +bool cactus_metal_encode_transform_batch(const void*, const CactusQuantMatrix* const*, int, void* const*) { return false; } +bool cactus_metal_encode_gemv_precoded(void*, const void*, const CactusQuantMatrix*) { return false; } +bool cactus_metal_prewarm_quant(const CactusQuantMatrix*) { return false; } +bool cactus_metal_encode_quant_matmul_ortho(void*, const void*, void*, const CactusQuantMatrix*) { return false; } +bool cactus_metal_encode_embedding_ortho(void*, uint32_t, const CactusQuantMatrix*, float) { return false; } +bool cactus_metal_encode_embedding_hadamard(void*, uint32_t, const CactusQuantMatrix*) { return false; } +bool cactus_metal_encode_embedding_ortho_m(void*, const CactusQuantMatrix*, const uint32_t*, uint32_t) { return false; } +bool cactus_metal_encode_embedding_hadamard_m(void*, const CactusQuantMatrix*, const uint32_t*, uint32_t) { return false; } +bool cactus_metal_encode_gather_f16(void*, const void*, size_t, const uint32_t*, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_attention_i8(void*, const void*, const void*, const void*, const void*, const void*, + const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, + float, size_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_attention_fused_i8(void*, const void*, const void*, const void*, + void*, void*, void*, void*, const void*, const void*, const void*, const void*, const void*, + uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, + float, float, size_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_attention_i8_prefill(void*, const void*, const void*, const void*, const void*, const void*, + const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, float, size_t, size_t, size_t, size_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_binary_f32(int, void*, const void*, const void*, size_t) { return false; } +bool cactus_metal_encode_scalar_f32(int, void*, const void*, size_t, float) { return false; } +bool cactus_metal_encode_unary_f32(int, void*, const void*, size_t) { return false; } +bool cactus_metal_encode_clamp(void*, const void*, size_t, float, float, int) { return false; } +bool cactus_metal_encode_glu(void*, const void*, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_layer_norm(void*, const void*, const void*, const void*, size_t, size_t, float) { return false; } +bool cactus_metal_encode_softmax_rows(void*, const void*, size_t, size_t) { return false; } +bool cactus_metal_encode_conv1d_k3(void*, const void*, const void*, int, const void*, uint32_t, + uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_gemm_f16(void*, const void*, const void*, uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_attention_f16(void*, const void*, const void*, const void*, const void*, + uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, + float, uint32_t, uint32_t, uint32_t, float, uint32_t) { return false; } +bool cactus_metal_encode_reduce_axis(int, void*, const void*, uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_cumsum(void*, const void*, uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_concat2(void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_gather_f32idx(void*, const void*, const void*, uint32_t, uint32_t, size_t) { return false; } +bool cactus_metal_encode_rope_full(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float, int) { return false; } +bool cactus_metal_encode_maxpool1d(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_bilinear(void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_conv1d_gen(void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_conv1d_nlc_dw(void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_conv2d(void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_batchnorm(void*, const void*, const void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, float) { return false; } +bool cactus_metal_encode_groupnorm(void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, float) { return false; } +bool cactus_metal_encode_bias_add_rows(void*, const void*, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_elemwise_chain(void*, const void*, const float*, uint32_t, const void*, const void*, const void*, const size_t*, size_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_rms_norm_add_rows(void*, void*, const void*, const void*, const void*, uint32_t, uint32_t, float, int) { return false; } +bool cactus_metal_encode_gemm_batch(void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, int, int) { return false; } +bool cactus_metal_encode_conv1d_dw(void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_transpose2d(void*, const void*, uint32_t, uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_strided_copy(void*, const void*, const uint32_t*, const uint32_t*, + uint32_t, uint32_t, uint32_t, size_t, size_t) { return false; } +bool cactus_metal_encode_bcast_binary(int, void*, const void*, const void*, const uint32_t*, + const uint32_t*, const uint32_t*, uint32_t, uint32_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_strided_scatter(void*, const void*, const uint32_t*, const uint32_t*, + uint32_t, uint32_t, uint32_t, size_t, size_t) { return false; } +bool cactus_metal_encode_kv_append_i8(const void*, void*, void*, uint32_t, uint32_t, uint32_t, + uint32_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_kv_append_sliding_i8(const void*, void*, void*, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, uint32_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_kv_append_sliding_i8_m(const void*, void*, void*, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, uint32_t, uint32_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_kv_append_i8_m(const void*, void*, void*, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_kv_append_ring_i8_m(const void*, void*, void*, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, uint32_t, uint32_t, size_t, size_t, size_t) { return false; } +bool cactus_metal_encode_conv_cache_append(void*, const void*, void*, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_rel_pos_bias(void*, const void*, const void*, uint32_t, uint32_t, uint32_t, + uint32_t, uint32_t, int, float) { return false; } +bool cactus_metal_encode_gemv_bias(void*, const void*, const void*, const void*, + uint32_t, uint32_t, int) { return false; } +bool cactus_metal_encode_rope_pair(void*, const void*, const void*, const void*, + uint32_t, uint32_t) { return false; } +bool cactus_metal_encode_rope_pair_rms(void*, const void*, const void*, const void*, + const void*, uint32_t, uint32_t, float) { return false; } +bool cactus_metal_encode_deltanet_decode(void*, const void*, const void*, const void*, + const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float) { return false; } +bool cactus_metal_encode_deltanet_prefill(void*, const void*, const void*, const void*, + const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float) { return false; } +bool cactus_metal_encode_rms2_add_clip(void*, const void*, const void*, + const void*, const void*, size_t, float, float) { return false; } +bool cactus_metal_encode_rms_norm_scale(void*, const void*, const void*, + size_t, size_t, float, float) { return false; } +bool cactus_metal_encode_softmax_topk(void*, void*, const void*, + size_t, size_t, size_t, float) { return false; } +bool cactus_metal_encode_topk_rows(void*, const void*, size_t, size_t, size_t) { return false; } +bool cactus_metal_moe_cq4_ready(const CactusQuantMatrix*) { return false; } +bool cactus_metal_moe_cq4_build(const CactusQuantMatrix*, const CactusQuantMatrix*, + const CactusQuantMatrix*, uint32_t) { return false; } +bool cactus_metal_encode_moe_gated_cq4(void*, const void*, const void*, const void*, const CactusQuantMatrix*, + uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float, float) { return false; } +#endif diff --git a/cactus/kernel/kernel_nn.cpp b/cactus-kernels/src/nn.cpp similarity index 54% rename from cactus/kernel/kernel_nn.cpp rename to cactus-kernels/src/nn.cpp index 65c791ff5..c56de09db 100644 --- a/cactus/kernel/kernel_nn.cpp +++ b/cactus-kernels/src/nn.cpp @@ -1,5 +1,5 @@ -#include "kernel.h" -#include "kernel_utils.h" +#include "../cactus_kernels.h" +#include "threading.h" #include #include #include @@ -7,38 +7,61 @@ #include #include #include -#include + +void cactus_relu_f16(const __fp16* input, __fp16* output, size_t num_elements) { + for (size_t i = 0; i < num_elements; ++i) { + __fp16 x = input[i]; + output[i] = x > static_cast<__fp16>(0) ? x : static_cast<__fp16>(0); + } +} + +void cactus_leaky_relu_f16(const __fp16* input, __fp16* output, size_t num_elements, float negative_slope) { + for (size_t i = 0; i < num_elements; ++i) { + __fp16 x = input[i]; + output[i] = x > static_cast<__fp16>(0) ? x : static_cast<__fp16>(static_cast(x) * negative_slope); + } +} + +void cactus_clamp_f16(const __fp16* input, __fp16* output, size_t num_elements, float lo, float hi) { + const float16x8_t lo_vec = vdupq_n_f16(static_cast<__fp16>(lo)); + const float16x8_t hi_vec = vdupq_n_f16(static_cast<__fp16>(hi)); + constexpr size_t SIMD_WIDTH = 8; + const size_t vectorized_end = (num_elements / SIMD_WIDTH) * SIMD_WIDTH; + + for (size_t i = 0; i < vectorized_end; i += SIMD_WIDTH) { + float16x8_t x = vld1q_f16(&input[i]); + x = vmaxq_f16(x, lo_vec); + x = vminq_f16(x, hi_vec); + vst1q_f16(&output[i], x); + } + const __fp16 lo_h = static_cast<__fp16>(lo); + const __fp16 hi_h = static_cast<__fp16>(hi); + for (size_t i = vectorized_end; i < num_elements; ++i) { + __fp16 x = input[i]; + if (x < lo_h) x = lo_h; + if (x > hi_h) x = hi_h; + output[i] = x; + } +} void cactus_silu_f16(const __fp16* input, __fp16* output, size_t num_elements) { CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_EXPENSIVE, [&](size_t start_idx, size_t end_idx) { constexpr size_t SIMD_WIDTH = 8; const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - const float32x4_t one_f32 = vdupq_n_f32(1.0f); for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { float16x8_t x = vld1q_f16(&input[i]); - float32x4_t x_low = vcvt_f32_f16(vget_low_f16(x)); float32x4_t x_high = vcvt_f32_f16(vget_high_f16(x)); - - float32x4_t exp_low = fast_exp_f32x4(vnegq_f32(x_low)); - float32x4_t exp_high = fast_exp_f32x4(vnegq_f32(x_high)); - - float32x4_t sigmoid_low = vdivq_f32(one_f32, vaddq_f32(one_f32, exp_low)); - float32x4_t sigmoid_high = vdivq_f32(one_f32, vaddq_f32(one_f32, exp_high)); - - float32x4_t silu_low = vmulq_f32(x_low, sigmoid_low); - float32x4_t silu_high = vmulq_f32(x_high, sigmoid_high); - - float16x8_t silu = vcombine_f16(vcvt_f16_f32(silu_low), vcvt_f16_f32(silu_high)); - vst1q_f16(&output[i], silu); + float32x4_t silu_low = vmulq_f32(x_low, fast_sigmoid_f32x4(x_low)); + float32x4_t silu_high = vmulq_f32(x_high, fast_sigmoid_f32x4(x_high)); + vst1q_f16(&output[i], vcombine_f16(vcvt_f16_f32(silu_low), vcvt_f16_f32(silu_high))); } for (size_t i = vectorized_end; i < end_idx; ++i) { float x_f32 = static_cast(input[i]); - float sigmoid = 1.0f / (1.0f + expf(-x_f32)); - output[i] = static_cast<__fp16>(x_f32 * sigmoid); + output[i] = static_cast<__fp16>(x_f32 / (1.0f + expf(-x_f32))); } }); } @@ -149,6 +172,241 @@ void cactus_gelu_f16_erf(const __fp16* input, __fp16* output, size_t num_element ); } +void cactus_sigmoid_f16(const __fp16* input, __fp16* output, size_t num_elements) { + CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start_idx, size_t end_idx) { + constexpr size_t SIMD_WIDTH = 8; + const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; + + for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { + float16x8_t x = vld1q_f16(&input[i]); + float32x4_t s_low = fast_sigmoid_f32x4(vcvt_f32_f16(vget_low_f16(x))); + float32x4_t s_high = fast_sigmoid_f32x4(vcvt_f32_f16(vget_high_f16(x))); + vst1q_f16(&output[i], vcombine_f16(vcvt_f16_f32(s_low), vcvt_f16_f32(s_high))); + } + + for (size_t i = vectorized_end; i < end_idx; ++i) { + float x_f32 = static_cast(input[i]); + output[i] = static_cast<__fp16>(1.0f / (1.0f + expf(-x_f32))); + } + }); +} + +void cactus_tanh_f16(const __fp16* input, __fp16* output, size_t num_elements) { + CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start_idx, size_t end_idx) { + constexpr size_t SIMD_WIDTH = 8; + const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; + + for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { + float16x8_t x_f16 = vld1q_f16(&input[i]); + + float32x4_t x_low = vcvt_f32_f16(vget_low_f16(x_f16)); + float32x4_t x_high = vcvt_f32_f16(vget_high_f16(x_f16)); + + float32x4_t tanh_low = fast_tanh_f32x4(x_low); + float32x4_t tanh_high = fast_tanh_f32x4(x_high); + + float16x8_t tanh_f16 = vcombine_f16(vcvt_f16_f32(tanh_low), vcvt_f16_f32(tanh_high)); + vst1q_f16(&output[i], tanh_f16); + } + + for (size_t i = vectorized_end; i < end_idx; ++i) { + float x = static_cast(input[i]); + output[i] = static_cast<__fp16>(std::tanh(x)); + } + }); +} + +void cactus_glu_f16( + const __fp16* input, + __fp16* output, + size_t outer_size, + size_t split_size, + size_t inner_size +) { + if (outer_size == 0 || split_size == 0 || inner_size == 0) return; + + const size_t axis_size = split_size * 2; + const size_t rows = outer_size * split_size; + + CactusThreading::parallel_for(rows, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t row_start, size_t row_end) { + for (size_t row = row_start; row < row_end; ++row) { + const size_t outer_idx = row / split_size; + const size_t split_idx = row - outer_idx * split_size; + const size_t in_base = outer_idx * axis_size * inner_size; + const size_t out_base = outer_idx * split_size * inner_size; + + const __fp16* a = input + in_base + split_idx * inner_size; + const __fp16* b = input + in_base + (split_idx + split_size) * inner_size; + __fp16* y = output + out_base + split_idx * inner_size; + + size_t i = 0; + for (; i + 8 <= inner_size; i += 8) { + const float16x8_t av = vld1q_f16(a + i); + const float16x8_t bv = vld1q_f16(b + i); + + const float32x4_t a_lo = vcvt_f32_f16(vget_low_f16(av)); + const float32x4_t a_hi = vcvt_f32_f16(vget_high_f16(av)); + + const float32x4_t y_lo = vmulq_f32(a_lo, fast_sigmoid_f32x4(vcvt_f32_f16(vget_low_f16(bv)))); + const float32x4_t y_hi = vmulq_f32(a_hi, fast_sigmoid_f32x4(vcvt_f32_f16(vget_high_f16(bv)))); + vst1q_f16(y + i, vcombine_f16(vcvt_f16_f32(y_lo), vcvt_f16_f32(y_hi))); + } + + for (; i < inner_size; ++i) { + const float gate = 1.0f / (1.0f + expf(-static_cast(b[i]))); + y[i] = static_cast<__fp16>(static_cast(a[i]) * gate); + } + } + }); +} + +void cactus_glu_f32( + const float* input, + float* output, + size_t outer_size, + size_t split_size, + size_t inner_size +) { + if (outer_size == 0 || split_size == 0 || inner_size == 0) return; + + const size_t axis_size = split_size * 2; + const size_t rows = outer_size * split_size; + + CactusThreading::parallel_for(rows, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t row_start, size_t row_end) { + for (size_t row = row_start; row < row_end; ++row) { + const size_t outer_idx = row / split_size; + const size_t split_idx = row - outer_idx * split_size; + const size_t in_base = outer_idx * axis_size * inner_size; + const size_t out_base = outer_idx * split_size * inner_size; + + const float* a = input + in_base + split_idx * inner_size; + const float* b = input + in_base + (split_idx + split_size) * inner_size; + float* y = output + out_base + split_idx * inner_size; + + size_t i = 0; + for (; i + 4 <= inner_size; i += 4) { + const float32x4_t av = vld1q_f32(a + i); + vst1q_f32(y + i, vmulq_f32(av, fast_sigmoid_f32x4(vld1q_f32(b + i)))); + } + + for (; i < inner_size; ++i) { + const float gate = 1.0f / (1.0f + expf(-b[i])); + y[i] = a[i] * gate; + } + } + }); +} + +static void batchnorm_apply_f16( + const __fp16* input, __fp16* output, + const float* scale, const float* shift, + size_t channels, size_t inner_size, size_t rows +) { + CactusThreading::parallel_for(rows, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t row_start, size_t row_end) { + for (size_t row = row_start; row < row_end; ++row) { + const size_t c = row % channels; + const float32x4_t scv = vdupq_n_f32(scale[c]); + const float32x4_t shv = vdupq_n_f32(shift[c]); + + const __fp16* x = input + row * inner_size; + __fp16* y = output + row * inner_size; + + size_t i = 0; + for (; i + 8 <= inner_size; i += 8) { + const float16x8_t xv = vld1q_f16(x + i); + float32x4_t lo = vfmaq_f32(shv, vcvt_f32_f16(vget_low_f16(xv)), scv); + float32x4_t hi = vfmaq_f32(shv, vcvt_f32_f16(vget_high_f16(xv)), scv); + vst1q_f16(y + i, vcombine_f16(vcvt_f16_f32(lo), vcvt_f16_f32(hi))); + } + + for (; i < inner_size; ++i) { + y[i] = static_cast<__fp16>(static_cast(x[i]) * scale[c] + shift[c]); + } + } + }); +} + +static void batchnorm_apply_f32( + const float* input, float* output, + const float* scale, const float* shift, + size_t channels, size_t inner_size, size_t rows +) { + CactusThreading::parallel_for(rows, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t row_start, size_t row_end) { + for (size_t row = row_start; row < row_end; ++row) { + const size_t c = row % channels; + const float32x4_t scv = vdupq_n_f32(scale[c]); + const float32x4_t shv = vdupq_n_f32(shift[c]); + + const float* x = input + row * inner_size; + float* y = output + row * inner_size; + + size_t i = 0; + for (; i + 4 <= inner_size; i += 4) { + vst1q_f32(y + i, vfmaq_f32(shv, vld1q_f32(x + i), scv)); + } + + for (; i < inner_size; ++i) { + y[i] = x[i] * scale[c] + shift[c]; + } + } + }); +} + +static void batchnorm_precompute( + const float* weight, const float* bias, + const float* running_mean, const float* running_var, + float* scale, float* shift, + size_t channels, float epsilon +) { + for (size_t c = 0; c < channels; ++c) { + const float inv_std = 1.0f / std::sqrt(running_var[c] + epsilon); + scale[c] = weight[c] * inv_std; + shift[c] = bias[c] - running_mean[c] * scale[c]; + } +} + +void cactus_batchnorm_f16( + const __fp16* input, + const float* weight, + const float* bias, + const float* running_mean, + const float* running_var, + __fp16* output, + size_t outer_size, + size_t channels, + size_t inner_size, + float epsilon +) { + if (outer_size == 0 || channels == 0 || inner_size == 0) return; + std::vector scale(channels), shift(channels); + batchnorm_precompute(weight, bias, running_mean, running_var, scale.data(), shift.data(), channels, epsilon); + batchnorm_apply_f16(input, output, scale.data(), shift.data(), channels, inner_size, outer_size * channels); +} + +void cactus_batchnorm_f32( + const float* input, + const float* weight, + const float* bias, + const float* running_mean, + const float* running_var, + float* output, + size_t outer_size, + size_t channels, + size_t inner_size, + float epsilon +) { + if (outer_size == 0 || channels == 0 || inner_size == 0) return; + std::vector scale(channels), shift(channels); + batchnorm_precompute(weight, bias, running_mean, running_var, scale.data(), shift.data(), channels, epsilon); + batchnorm_apply_f32(input, output, scale.data(), shift.data(), channels, inner_size, outer_size * channels); +} + void kernel_softmax_f16_single(const __fp16* input, __fp16* output, size_t vocab_size) { constexpr size_t SIMD_WIDTH = 8; @@ -254,32 +512,28 @@ void cactus_sample_f32(const float* logits, uint32_t* output, size_t vocab_size, float temperature, float top_p, size_t top_k, size_t random_seed, const float* bias_values, const uint32_t* bias_indices, size_t bias_count) { + cactus_sample_f32_ex(logits, output, vocab_size, + temperature, top_p, 0.15f, 1.1f, + top_k, random_seed, + bias_values, bias_indices, bias_count); +} - if (temperature == 0.0f && top_p <= 0.0f && top_k == 0) { - if (vocab_size == 0) { - output[0] = 0; - return; - } - size_t best_idx = 0; - float best_val = logits[0]; - for (size_t i = 1; i < vocab_size; ++i) { - float val = logits[i]; - if (val > best_val) { - best_val = val; - best_idx = i; - } - } - output[0] = static_cast(best_idx); +void cactus_sample_f32_ex(const float* logits, uint32_t* output, size_t vocab_size, + float temperature, float top_p, float min_p, float repetition_penalty, + size_t top_k, size_t random_seed, + const float* bias_values, const uint32_t* bias_indices, + size_t bias_count) { + if (vocab_size == 0) { + output[0] = 0; return; } - std::vector filtered_logits(vocab_size); + const bool has_bias = bias_values && bias_indices && bias_count > 0; - for (size_t i = 0; i < vocab_size; ++i) { - filtered_logits[i] = logits[i]; - } + std::vector filtered_logits(vocab_size); + std::memcpy(filtered_logits.data(), logits, vocab_size * sizeof(float)); - if (bias_values && bias_indices && bias_count > 0) { + if (has_bias) { for (size_t i = 0; i < bias_count; ++i) { uint32_t idx = bias_indices[i]; if (idx < vocab_size) { @@ -288,11 +542,19 @@ void cactus_sample_f32(const float* logits, uint32_t* output, size_t vocab_size, } } + if (temperature == 0.0f) { + auto it = std::max_element(filtered_logits.begin(), filtered_logits.end()); + output[0] = static_cast(std::distance(filtered_logits.begin(), it)); + return; + } + if (temperature > 0) { for (size_t i = 0; i < vocab_size; ++i) { filtered_logits[i] /= temperature; } } + + (void)repetition_penalty; if (top_k > 0) { std::vector> logit_pairs; @@ -313,7 +575,6 @@ void cactus_sample_f32(const float* logits, uint32_t* output, size_t vocab_size, } } - constexpr float min_p = 0.15f; if (min_p > 0.0f) { float max_logit = *std::max_element(filtered_logits.begin(), filtered_logits.end()); if (!std::isinf(max_logit)) { @@ -447,205 +708,27 @@ void cactus_sample_f16(const __fp16* logits, uint32_t* output, size_t vocab_size float temperature, float top_p, size_t top_k, size_t random_seed, const float* bias_values, const uint32_t* bias_indices, size_t bias_count) { + cactus_sample_f16_ex(logits, output, vocab_size, + temperature, top_p, 0.15f, 1.1f, + top_k, random_seed, + bias_values, bias_indices, bias_count); +} - if (temperature == 0.0f && top_p <= 0.0f && top_k == 0) { - if (vocab_size == 0) { - output[0] = 0; - return; - } - size_t best_idx = 0; - float best_val = static_cast(logits[0]); - for (size_t i = 1; i < vocab_size; ++i) { - float val = static_cast(logits[i]); - if (val > best_val) { - best_val = val; - best_idx = i; - } - } - output[0] = static_cast(best_idx); - return; - } - - std::vector<__fp16> filtered_logits(vocab_size); - - std::memcpy(filtered_logits.data(), logits, vocab_size * sizeof(__fp16)); - - if (bias_values && bias_indices && bias_count > 0) { - for (size_t i = 0; i < bias_count; ++i) { - uint32_t idx = bias_indices[i]; - if (idx < vocab_size) { - filtered_logits[idx] = static_cast<__fp16>(static_cast(filtered_logits[idx]) + bias_values[i]); - } - } - } - - if (temperature > 0) { - __fp16 inv_temp = static_cast<__fp16>(1.0f / temperature); - float16x8_t inv_temp_vec = vdupq_n_f16(inv_temp); - size_t i = 0; - for (; i + 8 <= vocab_size; i += 8) { - float16x8_t logits_vec = vld1q_f16(&logits[i]); - float16x8_t scaled = vmulq_f16(logits_vec, inv_temp_vec); - vst1q_f16(&filtered_logits[i], scaled); - } - for (; i < vocab_size; ++i) { - filtered_logits[i] = logits[i] * inv_temp; - } - } else { - std::memcpy(filtered_logits.data(), logits, vocab_size * sizeof(__fp16)); - } - - static std::vector token_history; - static const size_t MAX_HISTORY = 128; - static const float REPETITION_PENALTY = 1.1f; - - if (!token_history.empty() && REPETITION_PENALTY != 1.0f) { - const __fp16 penalty_inv = static_cast<__fp16>(1.0f / REPETITION_PENALTY); - const __fp16 penalty = static_cast<__fp16>(REPETITION_PENALTY); - - for (uint32_t prev_token : token_history) { - if (prev_token < vocab_size) { - filtered_logits[prev_token] = (filtered_logits[prev_token] > static_cast<__fp16>(0)) - ? static_cast<__fp16>(filtered_logits[prev_token] * penalty_inv) - : static_cast<__fp16>(filtered_logits[prev_token] * penalty); - } - } - } - - if (top_k > 0) { - std::vector> logit_pairs; - logit_pairs.reserve(vocab_size); - for (size_t i = 0; i < vocab_size; ++i) { - logit_pairs.emplace_back(filtered_logits[i], i); - } - std::partial_sort(logit_pairs.begin(), - logit_pairs.begin() + std::min(top_k, vocab_size), - logit_pairs.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - - if (top_k < vocab_size) { - __fp16 kth_value = logit_pairs[top_k - 1].first; - __fp16 neg_inf = static_cast<__fp16>(-std::numeric_limits::infinity()); - for (size_t i = 0; i < vocab_size; ++i) { - if (filtered_logits[i] < kth_value) { - filtered_logits[i] = neg_inf; - } - } - } - } - - if (top_p > 0.0f && top_p < 1.0f) { - std::vector> sorted_logits; - sorted_logits.reserve(vocab_size); - __fp16 neg_inf = static_cast<__fp16>(-std::numeric_limits::infinity()); - for (size_t i = 0; i < vocab_size; ++i) { - if (filtered_logits[i] != neg_inf) { - sorted_logits.emplace_back(filtered_logits[i], i); - } - } - std::sort(sorted_logits.begin(), sorted_logits.end(), - [](const auto& a, const auto& b) { return a.first > b.first; }); - - __fp16 max_logit = sorted_logits.empty() ? static_cast<__fp16>(0.0f) : sorted_logits[0].first; - std::vector temp_probs; - temp_probs.reserve(sorted_logits.size()); - float sum = 0.0f; - for (const auto& pair : sorted_logits) { - float prob = std::exp(static_cast(pair.first - max_logit)); - temp_probs.push_back(prob); - sum += prob; - } - - for (float& prob : temp_probs) { - prob /= sum; - } - - float cumulative_prob = 0.0f; - std::vector indices_to_remove(sorted_logits.size(), false); - bool threshold_reached = false; - for (size_t i = 0; i < sorted_logits.size(); ++i) { - cumulative_prob += temp_probs[i]; - if (cumulative_prob > top_p && i > 0) { - threshold_reached = true; - } - if (threshold_reached) { - indices_to_remove[i] = true; - } - } - - if (!indices_to_remove.empty()) { - for (size_t i = 1; i < indices_to_remove.size(); ++i) { - indices_to_remove[i] = indices_to_remove[i-1] || indices_to_remove[i]; - } - indices_to_remove[0] = false; - } - - for (size_t i = 0; i < sorted_logits.size(); ++i) { - if (indices_to_remove[i]) { - filtered_logits[sorted_logits[i].second] = neg_inf; - } - } - } - - __fp16 max_logit = *std::max_element(filtered_logits.begin(), filtered_logits.end()); - __fp16 neg_inf = static_cast<__fp16>(-std::numeric_limits::infinity()); - if (max_logit == neg_inf) { +void cactus_sample_f16_ex(const __fp16* logits, uint32_t* output, size_t vocab_size, + float temperature, float top_p, float min_p, float repetition_penalty, + size_t top_k, size_t random_seed, + const float* bias_values, const uint32_t* bias_indices, + size_t bias_count) { + if (vocab_size == 0) { output[0] = 0; return; } - - std::vector probs(vocab_size); - float sum = 0.0f; - for (size_t i = 0; i < vocab_size; ++i) { - if (filtered_logits[i] == neg_inf) { - probs[i] = 0.0f; - } else { - probs[i] = std::exp(static_cast(filtered_logits[i] - max_logit)); - sum += probs[i]; - } - } - if (sum == 0.0f) { - output[0] = 0; - return; - } + std::vector logits_f32(vocab_size); + cactus_fp16_to_fp32(logits, logits_f32.data(), vocab_size); - for (size_t i = 0; i < vocab_size; ++i) { - probs[i] /= sum; - } - - uint32_t actual_seed = (random_seed == 0) ? std::random_device{}() : random_seed; - std::mt19937 gen(actual_seed); - std::uniform_real_distribution dist(0.0f, 1.0f); - float sample = dist(gen); - - float cumulative = 0.0f; - for (size_t i = 0; i < vocab_size; ++i) { - cumulative += probs[i]; - if (cumulative >= sample) { - output[0] = static_cast(i); - token_history.push_back(output[0]); - if (token_history.size() > MAX_HISTORY) { - token_history.erase(token_history.begin()); - } - return; - } - } - - for (size_t i = vocab_size; i > 0; --i) { - if (probs[i-1] > 0.0f) { - output[0] = static_cast(i-1); - token_history.push_back(output[0]); - if (token_history.size() > MAX_HISTORY) { - token_history.erase(token_history.begin()); - } - return; - } - } - - output[0] = 0; - token_history.push_back(output[0]); - if (token_history.size() > MAX_HISTORY) { - token_history.erase(token_history.begin()); - } + cactus_sample_f32_ex(logits_f32.data(), output, vocab_size, + temperature, top_p, min_p, repetition_penalty, + top_k, random_seed, + bias_values, bias_indices, bias_count); } diff --git a/cactus-kernels/src/norms_rope.cpp b/cactus-kernels/src/norms_rope.cpp new file mode 100644 index 000000000..b2f466871 --- /dev/null +++ b/cactus-kernels/src/norms_rope.cpp @@ -0,0 +1,424 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include +#include +#include + +void cactus_rms_norm_f16( + const __fp16* input, + const __fp16* weight, + __fp16* output, + size_t batch_size, + size_t dims, + float eps +) { + constexpr size_t SIMD_WIDTH = 8; + constexpr size_t UNROLL_FACTOR = 2; + constexpr size_t TILE_SIZE = SIMD_WIDTH * UNROLL_FACTOR; + + for (size_t b = 0; b < batch_size; ++b) { + const __fp16* input_row = input + b * dims; + __fp16* output_row = output + b * dims; + + float32x4_t sum_squares_vec[UNROLL_FACTOR * 2]; + for (size_t u = 0; u < UNROLL_FACTOR * 2; u++) { + sum_squares_vec[u] = vdupq_n_f32(0.0f); + } + + size_t i = 0; + const size_t tile_end = (dims >= TILE_SIZE) ? dims - TILE_SIZE + 1 : 0; + + for (; i < tile_end; i += TILE_SIZE) { + for (size_t u = 0; u < UNROLL_FACTOR; u++) { + float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); + float32x4_t input_low = vcvt_f32_f16(vget_low_f16(input_vec)); + float32x4_t input_high = vcvt_f32_f16(vget_high_f16(input_vec)); + sum_squares_vec[u * 2] = vfmaq_f32(sum_squares_vec[u * 2], input_low, input_low); + sum_squares_vec[u * 2 + 1] = vfmaq_f32(sum_squares_vec[u * 2 + 1], input_high, input_high); + } + } + + const size_t simd_end = (dims >= SIMD_WIDTH) ? dims - SIMD_WIDTH + 1 : 0; + for (; i < simd_end; i += SIMD_WIDTH) { + float16x8_t input_vec = vld1q_f16(&input_row[i]); + float32x4_t input_low = vcvt_f32_f16(vget_low_f16(input_vec)); + float32x4_t input_high = vcvt_f32_f16(vget_high_f16(input_vec)); + sum_squares_vec[0] = vfmaq_f32(sum_squares_vec[0], input_low, input_low); + sum_squares_vec[1] = vfmaq_f32(sum_squares_vec[1], input_high, input_high); + } + + float32x4_t total_sum = sum_squares_vec[0]; + for (size_t u = 1; u < UNROLL_FACTOR * 2; u++) { + total_sum = vaddq_f32(total_sum, sum_squares_vec[u]); + } + float sum_squares = vaddvq_f32(total_sum); + + for (; i < dims; ++i) { + float val = static_cast(input_row[i]); + sum_squares += val * val; + } + + float rms = sqrtf(sum_squares / static_cast(dims) + eps); + float inv_rms = 1.0f / rms; + float16x8_t inv_rms_vec = vdupq_n_f16(static_cast<__fp16>(inv_rms)); + + i = 0; + for (; i < tile_end; i += TILE_SIZE) { + for (size_t u = 0; u < UNROLL_FACTOR; u++) { + float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); + float16x8_t weight_vec = vld1q_f16(&weight[i + u * SIMD_WIDTH]); + float16x8_t norm_vec = vmulq_f16(vmulq_f16(input_vec, inv_rms_vec), weight_vec); + vst1q_f16(&output_row[i + u * SIMD_WIDTH], norm_vec); + } + } + + for (; i < simd_end; i += SIMD_WIDTH) { + float16x8_t input_vec = vld1q_f16(&input_row[i]); + float16x8_t weight_vec = vld1q_f16(&weight[i]); + float16x8_t norm_vec = vmulq_f16(vmulq_f16(input_vec, inv_rms_vec), weight_vec); + vst1q_f16(&output_row[i], norm_vec); + } + + for (; i < dims; ++i) { + output_row[i] = static_cast<__fp16>(static_cast(input_row[i]) * inv_rms * static_cast(weight[i])); + } + } +} + +void cactus_layer_norm_f16( + const __fp16* input, + const __fp16* weight, + const __fp16* bias, + __fp16* output, + size_t batch_size, + size_t dims, + float eps +) { + constexpr size_t SIMD_WIDTH = 8; + constexpr size_t UNROLL_FACTOR = 3; + constexpr size_t TILE_SIZE = SIMD_WIDTH * UNROLL_FACTOR; + + const size_t tile_end = (dims >= TILE_SIZE) ? dims - TILE_SIZE + 1 : 0; + const size_t simd_end = (dims >= SIMD_WIDTH) ? dims - SIMD_WIDTH + 1 : 0; + + for (size_t b = 0; b < batch_size; ++b) { + const __fp16* input_row = input + b * dims; + __fp16* output_row = output + b * dims; + + float32x4_t sum_input_vec[UNROLL_FACTOR * 2]; + float32x4_t sum_squares_vec[UNROLL_FACTOR * 2]; + for (size_t u = 0; u < UNROLL_FACTOR * 2; u++) { + sum_input_vec[u] = vdupq_n_f32(0.0f); + sum_squares_vec[u] = vdupq_n_f32(0.0f); + } + + size_t i = 0; + + for (; i < tile_end; i += TILE_SIZE) { + for (size_t u = 0; u < UNROLL_FACTOR; u++) { + float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); + float32x4_t input_low = vcvt_f32_f16(vget_low_f16(input_vec)); + float32x4_t input_high = vcvt_f32_f16(vget_high_f16(input_vec)); + + sum_input_vec[u * 2] = vaddq_f32(sum_input_vec[u * 2], input_low); + sum_input_vec[u * 2 + 1] = vaddq_f32(sum_input_vec[u * 2 + 1], input_high); + + sum_squares_vec[u * 2] = vfmaq_f32(sum_squares_vec[u * 2], input_low, input_low); + sum_squares_vec[u * 2 + 1] = vfmaq_f32(sum_squares_vec[u * 2 + 1], input_high, input_high); + } + } + + for (; i < simd_end; i += SIMD_WIDTH) { + float16x8_t input_vec = vld1q_f16(&input_row[i]); + float32x4_t input_low = vcvt_f32_f16(vget_low_f16(input_vec)); + float32x4_t input_high = vcvt_f32_f16(vget_high_f16(input_vec)); + sum_input_vec[0] = vaddq_f32(sum_input_vec[0], input_low); + sum_input_vec[1] = vaddq_f32(sum_input_vec[1], input_high); + sum_squares_vec[0] = vfmaq_f32(sum_squares_vec[0], input_low, input_low); + sum_squares_vec[1] = vfmaq_f32(sum_squares_vec[1], input_high, input_high); + } + + float32x4_t total_sum_inputs = sum_input_vec[0]; + float32x4_t total_sum_squares = sum_squares_vec[0]; + for (size_t u = 1; u < UNROLL_FACTOR * 2; u++) { + total_sum_inputs = vaddq_f32(total_sum_inputs, sum_input_vec[u]); + total_sum_squares = vaddq_f32(total_sum_squares, sum_squares_vec[u]); + } + + float sum_inputs = vaddvq_f32(total_sum_inputs); + float sum_squares = vaddvq_f32(total_sum_squares); + for (; i < dims; ++i) { + float val = static_cast(input_row[i]); + sum_inputs += val; + sum_squares += val * val; + } + + float mean = sum_inputs / static_cast(dims); + float mean_squares = sum_squares / static_cast(dims); + float variance = mean_squares - mean * mean; + if (variance < 0.0f) variance = 0.0f; + float inv_std = 1.0f / sqrtf(variance + eps); + + float16x8_t mean_vec = vdupq_n_f16(static_cast<__fp16>(mean)); + float16x8_t inv_std_vec = vdupq_n_f16(static_cast<__fp16>(inv_std)); + + i = 0; + if (bias) { + for (; i < tile_end; i += TILE_SIZE) { + for (size_t u = 0; u < UNROLL_FACTOR; u++) { + float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); + float16x8_t weight_vec = vld1q_f16(&weight[i + u * SIMD_WIDTH]); + float16x8_t bias_vec = vld1q_f16(&bias[i + u * SIMD_WIDTH]); + float16x8_t out_vec = vmulq_f16(vmulq_f16(vsubq_f16(input_vec, mean_vec), inv_std_vec), weight_vec); + out_vec = vaddq_f16(out_vec, bias_vec); + vst1q_f16(&output_row[i + u * SIMD_WIDTH], out_vec); + } + } + + for (; i < simd_end; i += SIMD_WIDTH) { + float16x8_t input_vec = vld1q_f16(&input_row[i]); + float16x8_t weight_vec = vld1q_f16(&weight[i]); + float16x8_t bias_vec = vld1q_f16(&bias[i]); + float16x8_t out_vec = vmulq_f16(vmulq_f16(vsubq_f16(input_vec, mean_vec), inv_std_vec), weight_vec); + out_vec = vaddq_f16(out_vec, bias_vec); + vst1q_f16(&output_row[i], out_vec); + } + + for (; i < dims; ++i) { + output_row[i] = static_cast<__fp16>((static_cast(input_row[i]) - mean) * inv_std * static_cast(weight[i]) + static_cast(bias[i])); + } + } else { + for (; i < tile_end; i += TILE_SIZE) { + for (size_t u = 0; u < UNROLL_FACTOR; u++) { + float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); + float16x8_t weight_vec = vld1q_f16(&weight[i + u * SIMD_WIDTH]); + float16x8_t out_vec = vmulq_f16(vmulq_f16(vsubq_f16(input_vec, mean_vec), inv_std_vec), weight_vec); + vst1q_f16(&output_row[i + u * SIMD_WIDTH], out_vec); + } + } + + for (; i < simd_end; i += SIMD_WIDTH) { + float16x8_t input_vec = vld1q_f16(&input_row[i]); + float16x8_t weight_vec = vld1q_f16(&weight[i]); + float16x8_t out_vec = vmulq_f16(vmulq_f16(vsubq_f16(input_vec, mean_vec), inv_std_vec), weight_vec); + vst1q_f16(&output_row[i], out_vec); + } + + for (; i < dims; ++i) { + output_row[i] = static_cast<__fp16>((static_cast(input_row[i]) - mean) * inv_std * static_cast(weight[i])); + } + } + } +} + +namespace CactusRoPEF16 { + +struct RoPECacheF16 { + std::vector<__fp16> cos_table; + std::vector<__fp16> sin_table; + size_t max_seq_len; + size_t head_dim; + float theta; + bool initialized; + + RoPECacheF16() : max_seq_len(0), head_dim(0), theta(0.0f), initialized(false) {} +}; + +static thread_local std::vector rope_caches_f16; +static thread_local RoPECacheF16* active_rope_cache_f16 = nullptr; + +void precompute_rope_tables_f16(size_t seq_len, size_t head_dim, float theta) { + RoPECacheF16* cache = nullptr; + for (auto& candidate : rope_caches_f16) { + if (candidate.initialized && candidate.head_dim == head_dim && candidate.theta == theta) { + cache = &candidate; + break; + } + } + if (!cache) { + rope_caches_f16.emplace_back(); + cache = &rope_caches_f16.back(); + cache->head_dim = head_dim; + cache->theta = theta; + } + + active_rope_cache_f16 = cache; + if (cache->initialized && cache->max_seq_len >= seq_len) { + return; + } + + const size_t half_dim = head_dim / 2; + const size_t table_size = seq_len * half_dim; + + size_t start_pos = 0; + if (cache->initialized) { + start_pos = cache->max_seq_len; + } + + cache->cos_table.resize(table_size); + cache->sin_table.resize(table_size); + + for (size_t pos = start_pos; pos < seq_len; ++pos) { + const float pos_float = static_cast(pos); + for (size_t i = 0; i < half_dim; ++i) { + const float freq = 1.0f / powf(theta, (2.0f * i) / head_dim); + const float angle = pos_float * freq; + + const size_t idx = pos * half_dim + i; + cache->cos_table[idx] = static_cast<__fp16>(cosf(angle)); + cache->sin_table[idx] = static_cast<__fp16>(sinf(angle)); + } + } + + cache->max_seq_len = seq_len; + cache->initialized = true; +} + +} + +void cactus_rope_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t num_heads, + size_t head_dim, + size_t start_pos, + float theta +) { + const size_t half_dim = head_dim / 2; + + CactusRoPEF16::precompute_rope_tables_f16(seq_len + start_pos, head_dim, theta); + + const auto& cache = *CactusRoPEF16::active_rope_cache_f16; + const __fp16* cos_cache = cache.cos_table.data() + start_pos * half_dim; + const __fp16* sin_cache = cache.sin_table.data() + start_pos * half_dim; + + CactusThreading::parallel_for(batch_size * seq_len, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start_idx, size_t end_idx) { + for (size_t idx = start_idx; idx < end_idx; ++idx) { + const size_t batch_idx = idx / seq_len; + const size_t seq_idx = idx % seq_len; + + for (size_t head_idx = 0; head_idx < num_heads; ++head_idx) { + const size_t offset = ((batch_idx * seq_len + seq_idx) * num_heads + head_idx) * head_dim; + const __fp16* input_ptr = input + offset; + __fp16* output_ptr = output + offset; + + const __fp16* cos_ptr = cos_cache + seq_idx * half_dim; + const __fp16* sin_ptr = sin_cache + seq_idx * half_dim; + + constexpr size_t SIMD_WIDTH = 8; + const size_t vectorized_half_dim = (half_dim / SIMD_WIDTH) * SIMD_WIDTH; + + for (size_t i = 0; i < vectorized_half_dim; i += SIMD_WIDTH) { + float16x8_t cos_vec = vld1q_f16(&cos_ptr[i]); + float16x8_t sin_vec = vld1q_f16(&sin_ptr[i]); + + float16x8_t x_first_half = vld1q_f16(&input_ptr[i]); + float16x8_t x_second_half = vld1q_f16(&input_ptr[i + half_dim]); + + float16x8_t first_result = vfmsq_f16(vmulq_f16(x_first_half, cos_vec), x_second_half, sin_vec); + float16x8_t second_result = vfmaq_f16(vmulq_f16(x_second_half, cos_vec), x_first_half, sin_vec); + + vst1q_f16(&output_ptr[i], first_result); + vst1q_f16(&output_ptr[i + half_dim], second_result); + } + + for (size_t i = vectorized_half_dim; i < half_dim; ++i) { + const __fp16 cos_val = cos_ptr[i]; + const __fp16 sin_val = sin_ptr[i]; + + const __fp16 x_first_half = input_ptr[i]; + const __fp16 x_second_half = input_ptr[i + half_dim]; + + output_ptr[i] = x_first_half * cos_val - x_second_half * sin_val; + + output_ptr[i + half_dim] = x_second_half * cos_val + x_first_half * sin_val; + } + } + } + }); +} + +void cactus_gpt_j_rope_f16( + const __fp16* input, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t num_heads, + size_t head_dim, + size_t rot_dim, + size_t start_pos, + float theta +) { + const size_t half_rot_dim = rot_dim / 2; + + CactusRoPEF16::precompute_rope_tables_f16(seq_len + start_pos, rot_dim, theta); + + const auto& cache = *CactusRoPEF16::active_rope_cache_f16; + const __fp16* cos_cache = cache.cos_table.data() + start_pos * half_rot_dim; + const __fp16* sin_cache = cache.sin_table.data() + start_pos * half_rot_dim; + + CactusThreading::parallel_for(batch_size * seq_len, CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start_idx, size_t end_idx) { + for (size_t idx = start_idx; idx < end_idx; ++idx) { + const size_t batch_idx = idx / seq_len; + const size_t seq_idx = idx % seq_len; + + for (size_t head_idx = 0; head_idx < num_heads; ++head_idx) { + const size_t offset = ((batch_idx * seq_len + seq_idx) * num_heads + head_idx) * head_dim; + const __fp16* input_ptr = input + offset; + __fp16* output_ptr = output + offset; + + const __fp16* cos_ptr = cos_cache + seq_idx * half_rot_dim; + const __fp16* sin_ptr = sin_cache + seq_idx * half_rot_dim; + + constexpr size_t SIMD_WIDTH = 8; + const size_t vectorized_half_rot_dim = (half_rot_dim / SIMD_WIDTH) * SIMD_WIDTH; + + for (size_t i = 0; i < vectorized_half_rot_dim; i += SIMD_WIDTH) { + float16x8_t cos_vec = vld1q_f16(&cos_ptr[i]); + float16x8_t sin_vec = vld1q_f16(&sin_ptr[i]); + + float16x8x2_t x_vec = vld2q_f16(&input_ptr[2*i]); + float16x8_t x_first_half = x_vec.val[0]; + float16x8_t x_second_half = x_vec.val[1]; + + float16x8_t first_result = vfmsq_f16(vmulq_f16(x_first_half, cos_vec), x_second_half, sin_vec); + float16x8_t second_result = vfmaq_f16(vmulq_f16(x_second_half, cos_vec), x_first_half, sin_vec); + + float16x8x2_t t; + t.val[0] = first_result; + t.val[1] = second_result; + vst2q_f16(&output_ptr[2*i], t); + } + + for (size_t i = vectorized_half_rot_dim; i < half_rot_dim; ++i) { + const __fp16 cos_val = cos_ptr[i]; + const __fp16 sin_val = sin_ptr[i]; + + const __fp16 x_first_half = input_ptr[2*i]; + const __fp16 x_second_half = input_ptr[2*i + 1]; + + output_ptr[2*i] = x_first_half * cos_val - x_second_half * sin_val; + + output_ptr[2*i + 1] = x_second_half * cos_val + x_first_half * sin_val; + } + + constexpr size_t TAIL_SIMD_WIDTH = 8; + size_t copy_idx = rot_dim; + const size_t copy_end_vec = (head_dim / TAIL_SIMD_WIDTH) * TAIL_SIMD_WIDTH; + + for (; copy_idx + TAIL_SIMD_WIDTH <= copy_end_vec; copy_idx += TAIL_SIMD_WIDTH) { + float16x8_t v = vld1q_f16(&input_ptr[copy_idx]); + vst1q_f16(&output_ptr[copy_idx], v); + } + for (; copy_idx < head_dim; ++copy_idx) { + output_ptr[copy_idx] = input_ptr[copy_idx]; + } + } + } + }); +} diff --git a/cactus/kernel/kernel_quants.cpp b/cactus-kernels/src/quants.cpp similarity index 83% rename from cactus/kernel/kernel_quants.cpp rename to cactus-kernels/src/quants.cpp index 40798b842..58332a503 100644 --- a/cactus/kernel/kernel_quants.cpp +++ b/cactus-kernels/src/quants.cpp @@ -1,7 +1,8 @@ -#include "kernel.h" -#include "kernel_utils.h" +#include "../cactus_kernels.h" +#include "threading.h" #include #include +#include #include void cactus_int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale) { @@ -292,60 +293,3 @@ void cactus_quantize_kv_fp16_to_int8( }); } -void cactus_unpack_int4_to_int8(const uint8_t* packed, int8_t* unpacked, size_t unpacked_count) { - const size_t packed_count = (unpacked_count + 1) / 2; - - CactusThreading::parallel_for(packed_count, CactusThreading::Thresholds::ELEMENT_WISE, - [packed, unpacked, unpacked_count](size_t start, size_t end) { - const size_t simd_end = start + ((end - start) / 16) * 16; - - for (size_t i = start; i < simd_end; i += 16) { - - uint8x16_t input = vld1q_u8(&packed[i]); - uint8x16_t low_mask = vdupq_n_u8(0x0F); - uint8x16_t low_nibbles = vandq_u8(input, low_mask); - - uint8x16_t high_nibbles = vshrq_n_u8(input, 4); - uint8x16_t sign_bit = vdupq_n_u8(0x08); - uint8x16_t sign_extend = vdupq_n_u8(0x10); - - uint8x16_t low_sign = vandq_u8(low_nibbles, sign_bit); - uint8x16_t low_has_sign = vcgtq_u8(low_sign, vdupq_n_u8(0)); - uint8x16_t low_correction = vandq_u8(low_has_sign, sign_extend); - int8x16_t low_signed = vreinterpretq_s8_u8(vsubq_u8(low_nibbles, low_correction)); - - uint8x16_t high_sign = vandq_u8(high_nibbles, sign_bit); - uint8x16_t high_has_sign = vcgtq_u8(high_sign, vdupq_n_u8(0)); - uint8x16_t high_correction = vandq_u8(high_has_sign, sign_extend); - int8x16_t high_signed = vreinterpretq_s8_u8(vsubq_u8(high_nibbles, high_correction)); - - int8x16x2_t interleaved = vzipq_s8(low_signed, high_signed); - - size_t out_idx = i * 2; - if (out_idx + 16 <= unpacked_count) { - vst1q_s8(&unpacked[out_idx], interleaved.val[0]); - } - if (out_idx + 32 <= unpacked_count) { - vst1q_s8(&unpacked[out_idx + 16], interleaved.val[1]); - } - } - - for (size_t i = simd_end; i < end; ++i) { - uint8_t byte = packed[i]; - - int8_t low = byte & 0x0F; - if (low & 0x08) low |= 0xF0; - - int8_t high = (byte >> 4) & 0x0F; - if (high & 0x08) high |= 0xF0; - - size_t out_idx = i * 2; - if (out_idx < unpacked_count) { - unpacked[out_idx] = low; - } - if (out_idx + 1 < unpacked_count) { - unpacked[out_idx + 1] = high; - } - } - }); -} \ No newline at end of file diff --git a/cactus-kernels/src/reduce.cpp b/cactus-kernels/src/reduce.cpp new file mode 100644 index 000000000..716308ae2 --- /dev/null +++ b/cactus-kernels/src/reduce.cpp @@ -0,0 +1,327 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include + +template +static void axis_reduce_f32_impl(const __fp16* input, __fp16* output, + size_t outer_size, size_t axis_size, size_t inner_size, + FinalizeFn finalize) { + if (inner_size == 1) { + CactusThreading::parallel_for(outer_size, CactusThreading::Thresholds::AXIS_REDUCE, + [&](size_t start_outer, size_t end_outer) { + constexpr size_t W = SIMD_F16_WIDTH; + const size_t vec_axis = simd_align(axis_size); + + for (size_t outer = start_outer; outer < end_outer; ++outer) { + const __fp16* row = input + outer * axis_size; + float32x4_t sum_lo = vdupq_n_f32(0.0f); + float32x4_t sum_hi = vdupq_n_f32(0.0f); + + for (size_t a = 0; a < vec_axis; a += W) { + float32x4_t lo, hi; + f16x8_split_f32(vld1q_f16(row + a), lo, hi); + sum_lo = vaddq_f32(sum_lo, lo); + sum_hi = vaddq_f32(sum_hi, hi); + } + + float total = vaddvq_f32(vaddq_f32(sum_lo, sum_hi)); + for (size_t a = vec_axis; a < axis_size; ++a) { + total += static_cast(row[a]); + } + + output[outer] = static_cast<__fp16>(finalize(total, axis_size)); + } + }); + return; + } + + CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, + [&](size_t outer, size_t inner) { + constexpr size_t W = SIMD_F16_WIDTH; + const size_t vec_axis = simd_align(axis_size); + + float32x4_t sum_lo = vdupq_n_f32(0.0f); + float32x4_t sum_hi = vdupq_n_f32(0.0f); + + for (size_t a = 0; a < vec_axis; a += W) { + __fp16 values[W]; + for (size_t j = 0; j < W; j++) { + values[j] = input[outer * axis_size * inner_size + (a + j) * inner_size + inner]; + } + float32x4_t lo, hi; + f16x8_split_f32(vld1q_f16(values), lo, hi); + sum_lo = vaddq_f32(sum_lo, lo); + sum_hi = vaddq_f32(sum_hi, hi); + } + + float total = vaddvq_f32(vaddq_f32(sum_lo, sum_hi)); + + for (size_t a = vec_axis; a < axis_size; a++) { + total += static_cast(input[outer * axis_size * inner_size + a * inner_size + inner]); + } + + output[outer * inner_size + inner] = static_cast<__fp16>(finalize(total, axis_size)); + }); +} + +template +static void axis_reduce_f16_impl(const __fp16* input, __fp16* output, + size_t outer_size, size_t axis_size, size_t inner_size, + float16x8_t init_vec, __fp16 init_scalar, + SimdReduceOp simd_reduce, ScalarReduceOp scalar_reduce) { + if (inner_size == 1) { + CactusThreading::parallel_for(outer_size, CactusThreading::Thresholds::AXIS_REDUCE, + [&](size_t start_outer, size_t end_outer) { + constexpr size_t W = SIMD_F16_WIDTH; + const size_t vec_axis = simd_align(axis_size); + + for (size_t outer = start_outer; outer < end_outer; ++outer) { + const __fp16* row = input + outer * axis_size; + float16x8_t acc = init_vec; + + for (size_t a = 0; a < vec_axis; a += W) { + acc = simd_reduce(acc, vld1q_f16(row + a)); + } + + __fp16 result = init_scalar; + __fp16 arr[W]; + vst1q_f16(arr, acc); + for (size_t j = 0; j < W; ++j) { + result = scalar_reduce(result, arr[j]); + } + + for (size_t a = vec_axis; a < axis_size; ++a) { + result = scalar_reduce(result, row[a]); + } + + output[outer] = result; + } + }); + return; + } + + CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, + [&](size_t outer, size_t inner) { + constexpr size_t W = SIMD_F16_WIDTH; + const size_t vec_axis = simd_align(axis_size); + float16x8_t acc = init_vec; + + for (size_t a = 0; a < vec_axis; a += W) { + __fp16 values[W]; + for (size_t j = 0; j < W; j++) { + values[j] = input[outer * axis_size * inner_size + (a + j) * inner_size + inner]; + } + acc = simd_reduce(acc, vld1q_f16(values)); + } + + __fp16 result = init_scalar; + __fp16 arr[8]; + vst1q_f16(arr, acc); + for (int j = 0; j < 8; j++) { + result = scalar_reduce(result, arr[j]); + } + + for (size_t a = vec_axis; a < axis_size; a++) { + result = scalar_reduce(result, input[outer * axis_size * inner_size + a * inner_size + inner]); + } + + output[outer * inner_size + inner] = result; + }); +} + +double cactus_sum_all_f16(const __fp16* data, size_t num_elements) { + return CactusThreading::parallel_reduce( + num_elements, CactusThreading::Thresholds::ALL_REDUCE, + [&](size_t start, size_t end) -> double { + const size_t vec_end = start + simd_align(end - start); + float32x4_t sum_lo = vdupq_n_f32(0.0f); + float32x4_t sum_hi = vdupq_n_f32(0.0f); + + for (size_t i = start; i < vec_end; i += SIMD_F16_WIDTH) { + float32x4_t lo, hi; + f16x8_split_f32(vld1q_f16(&data[i]), lo, hi); + sum_lo = vaddq_f32(sum_lo, lo); + sum_hi = vaddq_f32(sum_hi, hi); + } + + double s = static_cast(vaddvq_f32(vaddq_f32(sum_lo, sum_hi))); + for (size_t i = vec_end; i < end; ++i) s += static_cast(data[i]); + return s; + }, + 0.0, [](double a, double b) { return a + b; } + ); +} + +void cactus_sum_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { + axis_reduce_f32_impl(input, output, outer_size, axis_size, inner_size, + [](float total, size_t) { return total; }); +} + +double cactus_mean_all_f16(const __fp16* data, size_t num_elements) { + return cactus_sum_all_f16(data, num_elements) / static_cast(num_elements); +} + +void cactus_mean_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { + axis_reduce_f32_impl(input, output, outer_size, axis_size, inner_size, + [](float total, size_t axis_size) { return total / static_cast(axis_size); }); +} + +struct VarianceState { + double sum; + double sum_sq; + VarianceState() : sum(0.0), sum_sq(0.0) {} + VarianceState(double s, double sq) : sum(s), sum_sq(sq) {} +}; + +double cactus_variance_all_f16(const __fp16* data, size_t num_elements) { + VarianceState result = CactusThreading::parallel_reduce( + num_elements, CactusThreading::Thresholds::ALL_REDUCE, + [&](size_t start, size_t end) -> VarianceState { + const size_t vec_end = start + simd_align(end - start); + + float32x4_t sum_lo = vdupq_n_f32(0.0f), sum_hi = vdupq_n_f32(0.0f); + float32x4_t sq_lo = vdupq_n_f32(0.0f), sq_hi = vdupq_n_f32(0.0f); + + for (size_t i = start; i < vec_end; i += SIMD_F16_WIDTH) { + float32x4_t lo, hi; + f16x8_split_f32(vld1q_f16(&data[i]), lo, hi); + sum_lo = vaddq_f32(sum_lo, lo); + sum_hi = vaddq_f32(sum_hi, hi); + sq_lo = vfmaq_f32(sq_lo, lo, lo); + sq_hi = vfmaq_f32(sq_hi, hi, hi); + } + + double sum = static_cast(vaddvq_f32(vaddq_f32(sum_lo, sum_hi))); + double sum_sq = static_cast(vaddvq_f32(vaddq_f32(sq_lo, sq_hi))); + + for (size_t i = vec_end; i < end; ++i) { + double x = static_cast(data[i]); + sum += x; + sum_sq += x * x; + } + return VarianceState(sum, sum_sq); + }, + VarianceState(), + [](const VarianceState& a, const VarianceState& b) { + return VarianceState(a.sum + b.sum, a.sum_sq + b.sum_sq); + } + ); + + double mean = result.sum / static_cast(num_elements); + double mean_sq = result.sum_sq / static_cast(num_elements); + return mean_sq - mean * mean; +} + +void cactus_variance_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { + if (inner_size == 1) { + CactusThreading::parallel_for(outer_size, CactusThreading::Thresholds::AXIS_REDUCE, + [&](size_t start_outer, size_t end_outer) { + constexpr size_t W = SIMD_F16_WIDTH; + const size_t vec_axis = simd_align(axis_size); + + for (size_t outer = start_outer; outer < end_outer; ++outer) { + const __fp16* row = input + outer * axis_size; + float32x4_t sum_lo = vdupq_n_f32(0.0f); + float32x4_t sum_hi = vdupq_n_f32(0.0f); + float32x4_t sq_lo = vdupq_n_f32(0.0f); + float32x4_t sq_hi = vdupq_n_f32(0.0f); + + for (size_t a = 0; a < vec_axis; a += W) { + float32x4_t lo, hi; + f16x8_split_f32(vld1q_f16(row + a), lo, hi); + sum_lo = vaddq_f32(sum_lo, lo); + sum_hi = vaddq_f32(sum_hi, hi); + sq_lo = vfmaq_f32(sq_lo, lo, lo); + sq_hi = vfmaq_f32(sq_hi, hi, hi); + } + + float sum = vaddvq_f32(vaddq_f32(sum_lo, sum_hi)); + float sum_sq = vaddvq_f32(vaddq_f32(sq_lo, sq_hi)); + for (size_t a = vec_axis; a < axis_size; ++a) { + const float x = static_cast(row[a]); + sum += x; + sum_sq += x * x; + } + + const float mean = sum / static_cast(axis_size); + const float mean_sq = sum_sq / static_cast(axis_size); + output[outer] = static_cast<__fp16>(mean_sq - mean * mean); + } + }); + return; + } + + CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, + [&](size_t outer, size_t inner) { + float sum = 0.0f, sum_sq = 0.0f; + for (size_t a = 0; a < axis_size; a++) { + float x = static_cast(input[outer * axis_size * inner_size + a * inner_size + inner]); + sum += x; + sum_sq += x * x; + } + float mean = sum / static_cast(axis_size); + float mean_sq = sum_sq / static_cast(axis_size); + output[outer * inner_size + inner] = static_cast<__fp16>(mean_sq - mean * mean); + }); +} + +__fp16 cactus_min_all_f16(const __fp16* data, size_t num_elements) { + return CactusThreading::parallel_reduce( + num_elements, CactusThreading::Thresholds::ALL_REDUCE, + [&](size_t start, size_t end) -> __fp16 { + const size_t vec_end = start + simd_align(end - start); + float16x8_t acc = vdupq_n_f16(static_cast<__fp16>(65504.0f)); + + for (size_t i = start; i < vec_end; i += SIMD_F16_WIDTH) { + acc = vminq_f16(acc, vld1q_f16(&data[i])); + } + + __fp16 result = static_cast<__fp16>(65504.0f); + __fp16 arr[8]; + vst1q_f16(arr, acc); + for (int j = 0; j < 8; j++) result = std::min(result, arr[j]); + for (size_t i = vec_end; i < end; ++i) result = std::min(result, data[i]); + return result; + }, + static_cast<__fp16>(65504.0f), + [](__fp16 a, __fp16 b) { return std::min(a, b); } + ); +} + +void cactus_min_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { + axis_reduce_f16_impl(input, output, outer_size, axis_size, inner_size, + vdupq_n_f16(static_cast<__fp16>(65504.0f)), static_cast<__fp16>(65504.0f), + [](float16x8_t a, float16x8_t b) { return vminq_f16(a, b); }, + [](__fp16 a, __fp16 b) { return std::min(a, b); }); +} + +__fp16 cactus_max_all_f16(const __fp16* data, size_t num_elements) { + return CactusThreading::parallel_reduce( + num_elements, CactusThreading::Thresholds::ALL_REDUCE, + [&](size_t start, size_t end) -> __fp16 { + const size_t vec_end = start + simd_align(end - start); + float16x8_t acc = vdupq_n_f16(static_cast<__fp16>(-65504.0f)); + + for (size_t i = start; i < vec_end; i += SIMD_F16_WIDTH) { + acc = vmaxq_f16(acc, vld1q_f16(&data[i])); + } + + __fp16 result = static_cast<__fp16>(-65504.0f); + __fp16 arr[8]; + vst1q_f16(arr, acc); + for (int j = 0; j < 8; j++) result = std::max(result, arr[j]); + for (size_t i = vec_end; i < end; ++i) result = std::max(result, data[i]); + return result; + }, + static_cast<__fp16>(-65504.0f), + [](__fp16 a, __fp16 b) { return std::max(a, b); } + ); +} + +void cactus_max_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { + axis_reduce_f16_impl(input, output, outer_size, axis_size, inner_size, + vdupq_n_f16(static_cast<__fp16>(-65504.0f)), static_cast<__fp16>(-65504.0f), + [](float16x8_t a, float16x8_t b) { return vmaxq_f16(a, b); }, + [](__fp16 a, __fp16 b) { return std::max(a, b); }); +} diff --git a/cactus-kernels/src/scalar.cpp b/cactus-kernels/src/scalar.cpp new file mode 100644 index 000000000..44ce2e22e --- /dev/null +++ b/cactus-kernels/src/scalar.cpp @@ -0,0 +1,98 @@ +#include "../cactus_kernels.h" +#include "threading.h" +#include +#include + +void cactus_scalar_op_f16(const __fp16* input, __fp16* output, size_t num_elements, float scalar_value, ScalarOpType op_type) { + const __fp16 scalar_f16 = static_cast<__fp16>(scalar_value); + const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; + + switch (op_type) { + case ScalarOpType::ADD: { + const float16x8_t sv = vdupq_n_f16(scalar_f16); + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_BASIC, + [sv](float16x8_t v) { return vaddq_f16(v, sv); }, + [scalar_f16](__fp16 v) { return static_cast<__fp16>(v + scalar_f16); }); + break; + } + case ScalarOpType::SUBTRACT: { + const float16x8_t sv = vdupq_n_f16(scalar_f16); + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_BASIC, + [sv](float16x8_t v) { return vsubq_f16(v, sv); }, + [scalar_f16](__fp16 v) { return static_cast<__fp16>(v - scalar_f16); }); + break; + } + case ScalarOpType::MULTIPLY: { + const float16x8_t sv = vdupq_n_f16(scalar_f16); + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_BASIC, + [sv](float16x8_t v) { return vmulq_f16(v, sv); }, + [scalar_f16](__fp16 v) { return static_cast<__fp16>(v * scalar_f16); }); + break; + } + case ScalarOpType::DIVIDE: { + const float16x8_t sv = vdupq_n_f16(scalar_f16); + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_BASIC, + [sv](float16x8_t v) { return vdivq_f16(v, sv); }, + [scalar_f16](__fp16 v) { return static_cast<__fp16>(v / scalar_f16); }); + break; + } + case ScalarOpType::ABS: + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_BASIC, + [](float16x8_t v) { return vabsq_f16(v); }, + [](__fp16 v) { return static_cast<__fp16>(std::abs(static_cast(v))); }); + break; + + case ScalarOpType::EXP: + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [](float16x8_t v) { return apply_f32_op_on_f16x8(v, fast_exp_f32x4); }, + [](__fp16 v) { return static_cast<__fp16>(std::exp(static_cast(v))); }, 2); + break; + + case ScalarOpType::SQRT: + elementwise_op_f16(input, output, num_elements, use_streaming, + CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [](float16x8_t v) { return apply_f32_op_on_f16x8(v, vsqrtq_f32); }, + [](__fp16 v) { return static_cast<__fp16>(std::sqrt(static_cast(v))); }, 2); + break; + + case ScalarOpType::POW: + CactusThreading::parallel_for(num_elements, + CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start, size_t end) { + for (size_t i = start; i < end; ++i) { + output[i] = static_cast<__fp16>( + std::pow(static_cast(input[i]), static_cast(scalar_f16))); + } + }); + break; + + case ScalarOpType::COS: + case ScalarOpType::SIN: + CactusThreading::parallel_for(num_elements, + CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start, size_t end) { + for (size_t i = start; i < end; ++i) { + float val = static_cast(input[i]); + float result = (op_type == ScalarOpType::COS) ? std::cos(val) : std::sin(val); + output[i] = static_cast<__fp16>(result); + } + }); + break; + + case ScalarOpType::LOG: + CactusThreading::parallel_for(num_elements, + CactusThreading::Thresholds::SCALAR_EXPENSIVE, + [&](size_t start, size_t end) { + for (size_t i = start; i < end; ++i) { + output[i] = static_cast<__fp16>(std::log(static_cast(input[i]))); + } + }); + break; + } +} diff --git a/cactus-kernels/src/threading.h b/cactus-kernels/src/threading.h new file mode 100644 index 000000000..e5795329f --- /dev/null +++ b/cactus-kernels/src/threading.h @@ -0,0 +1,825 @@ +#ifndef KERNEL_UTILS_H +#define KERNEL_UTILS_H + +#include +#if defined(__APPLE__) +#include +#include +#endif +#if defined(__ANDROID__) +#include +#include +#include +#include +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +constexpr size_t NEON_VECTOR_SIZE = 16; +constexpr size_t STREAMING_STORE_THRESHOLD = 32768; + +inline void stream_store_f16x8(__fp16* dst, float16x8_t val) { +#if defined(__aarch64__) + float16x4_t lo = vget_low_f16(val); + float16x4_t hi = vget_high_f16(val); + __asm__ __volatile__( + "stnp %d0, %d1, [%2]" + : + : "w"(lo), "w"(hi), "r"(dst) + : "memory" + ); +#else + vst1q_f16(dst, val); +#endif +} + +inline bool cpu_has_i8mm() { +#if defined(__aarch64__) + static std::once_flag once; + static bool has = false; + + std::call_once(once, []() { +#if defined(__APPLE__) + int ret = 0; + size_t size = sizeof(ret); + if (sysctlbyname("hw.optional.arm.FEAT_I8MM", &ret, &size, nullptr, 0) == 0) { + has = (ret == 1); + } +#elif defined(__ANDROID__) + unsigned long hwcap2 = getauxval(AT_HWCAP2); + #ifndef HWCAP2_I8MM + #define HWCAP2_I8MM (1 << 13) + #endif + has = (hwcap2 & HWCAP2_I8MM) != 0; +#endif + }); + + return has; +#else + return false; +#endif +} + +inline float32x4_t fast_exp_f32x4(float32x4_t x) { + const float32x4_t log2e = vdupq_n_f32(1.44269504088896341f); + const float32x4_t ln2_hi = vdupq_n_f32(6.93145751953125e-1f); + const float32x4_t ln2_lo = vdupq_n_f32(1.42860682030941723212e-6f); + + const float32x4_t p0 = vdupq_n_f32(1.9875691500e-4f); + const float32x4_t p1 = vdupq_n_f32(1.3981999507e-3f); + const float32x4_t p2 = vdupq_n_f32(8.3334519073e-3f); + const float32x4_t p3 = vdupq_n_f32(4.1665795894e-2f); + const float32x4_t p4 = vdupq_n_f32(1.6666665459e-1f); + const float32x4_t p5 = vdupq_n_f32(5.0000001201e-1f); + const float32x4_t one = vdupq_n_f32(1.0f); + + x = vmaxq_f32(x, vdupq_n_f32(-87.0f)); + x = vminq_f32(x, vdupq_n_f32(87.0f)); + + float32x4_t z = vmulq_f32(x, log2e); + float32x4_t n = vrndnq_f32(z); + + float32x4_t r = vfmsq_f32(x, n, ln2_hi); + r = vfmsq_f32(r, n, ln2_lo); + + float32x4_t r2 = vmulq_f32(r, r); + float32x4_t p = p0; + p = vfmaq_f32(p1, p, r); + p = vfmaq_f32(p2, p, r); + p = vfmaq_f32(p3, p, r); + p = vfmaq_f32(p4, p, r); + p = vfmaq_f32(p5, p, r); + p = vfmaq_f32(r, p, r2); + p = vaddq_f32(p, one); + + int32x4_t ni = vcvtq_s32_f32(n); + int32x4_t exp_bits = vshlq_n_s32(vaddq_s32(ni, vdupq_n_s32(127)), 23); + float32x4_t scale = vreinterpretq_f32_s32(exp_bits); + + return vmulq_f32(p, scale); +} + +alignas(16) inline constexpr float kFastTanhAlpha[7][4] = { + { 4.89352455891786e-03f, 4.89352455891786e-03f, 4.89352455891786e-03f, 4.89352455891786e-03f }, + { 6.37261928875436e-04f, 6.37261928875436e-04f, 6.37261928875436e-04f, 6.37261928875436e-04f }, + { 1.48572235717979e-05f, 1.48572235717979e-05f, 1.48572235717979e-05f, 1.48572235717979e-05f }, + { 5.12229709037114e-08f, 5.12229709037114e-08f, 5.12229709037114e-08f, 5.12229709037114e-08f }, + {-8.60467152213735e-11f,-8.60467152213735e-11f,-8.60467152213735e-11f,-8.60467152213735e-11f }, + { 2.00018790482477e-13f, 2.00018790482477e-13f, 2.00018790482477e-13f, 2.00018790482477e-13f }, + {-2.76076847742355e-16f,-2.76076847742355e-16f,-2.76076847742355e-16f,-2.76076847742355e-16f }, +}; +alignas(16) inline constexpr float kFastTanhBeta[4][4] = { + { 4.89352518554385e-03f, 4.89352518554385e-03f, 4.89352518554385e-03f, 4.89352518554385e-03f }, + { 2.26843463243900e-03f, 2.26843463243900e-03f, 2.26843463243900e-03f, 2.26843463243900e-03f }, + { 1.18534705686654e-04f, 1.18534705686654e-04f, 1.18534705686654e-04f, 1.18534705686654e-04f }, + { 1.19825839466702e-06f, 1.19825839466702e-06f, 1.19825839466702e-06f, 1.19825839466702e-06f }, +}; +alignas(16) inline constexpr float kFastTanhClampHi[4] = { 9.0f, 9.0f, 9.0f, 9.0f }; +alignas(16) inline constexpr float kFastTanhClampLo[4] = {-9.0f,-9.0f,-9.0f,-9.0f }; + +inline float32x4_t fast_tanh_f32x4(float32x4_t x) { + x = vmaxq_f32(vld1q_f32(kFastTanhClampLo), vminq_f32(vld1q_f32(kFastTanhClampHi), x)); + float32x4_t x2 = vmulq_f32(x, x); + float32x4_t p = vfmaq_f32(vld1q_f32(kFastTanhAlpha[5]), vld1q_f32(kFastTanhAlpha[6]), x2); + p = vfmaq_f32(vld1q_f32(kFastTanhAlpha[4]), p, x2); + p = vfmaq_f32(vld1q_f32(kFastTanhAlpha[3]), p, x2); + p = vfmaq_f32(vld1q_f32(kFastTanhAlpha[2]), p, x2); + p = vfmaq_f32(vld1q_f32(kFastTanhAlpha[1]), p, x2); + p = vfmaq_f32(vld1q_f32(kFastTanhAlpha[0]), p, x2); + p = vmulq_f32(p, x); + float32x4_t q = vfmaq_f32(vld1q_f32(kFastTanhBeta[2]), vld1q_f32(kFastTanhBeta[3]), x2); + q = vfmaq_f32(vld1q_f32(kFastTanhBeta[1]), q, x2); + q = vfmaq_f32(vld1q_f32(kFastTanhBeta[0]), q, x2); + return vdivq_f32(p, q); +} + +constexpr size_t SIMD_F16_WIDTH = 8; + +inline size_t simd_align(size_t count, size_t width = SIMD_F16_WIDTH) { + return (count / width) * width; +} + +inline void f16x8_split_f32(float16x8_t v, float32x4_t& lo, float32x4_t& hi) { + lo = vcvt_f32_f16(vget_low_f16(v)); + hi = vcvt_f32_f16(vget_high_f16(v)); +} + +inline float16x8_t f32_merge_f16(float32x4_t lo, float32x4_t hi) { + return vcombine_f16(vcvt_f16_f32(lo), vcvt_f16_f32(hi)); +} + +inline float32x4_t fast_sigmoid_f32x4(float32x4_t x) { + const float32x4_t one = vdupq_n_f32(1.0f); + return vdivq_f32(one, vaddq_f32(one, fast_exp_f32x4(vnegq_f32(x)))); +} + +template +inline float16x8_t apply_f32_op_on_f16x8(float16x8_t v, F32x4Op op) { + float32x4_t lo, hi; + f16x8_split_f32(v, lo, hi); + return f32_merge_f16(op(lo), op(hi)); +} + +namespace CactusThreading { + +#if defined(__ANDROID__) + static constexpr size_t ANDROID_DYNAMIC_CHUNK_MULTIPLIER = 16; + class ThreadPool; + inline ThreadPool& get_thread_pool(); + + struct CoreTopology { + std::vector performance_cores; + std::vector performance_core_capacities; + std::vector all_cores; + + static CoreTopology& get() { + static CoreTopology topo = detect(); + return topo; + } + + private: + static int read_sysfs_int(const char* path) { + std::ifstream f(path); + if (!f.is_open()) return -1; + int val = -1; + f >> val; + return val; + } + + static CoreTopology detect() { + CoreTopology topo; + constexpr int MAX_CPUS = 16; + std::vector> core_caps; + + for (int i = 0; i < MAX_CPUS; ++i) { + char path[128]; + + snprintf(path, sizeof(path), + "/sys/devices/system/cpu/cpu%d/cpu_capacity", i); + int cap = read_sysfs_int(path); + if (cap > 0) { + core_caps.push_back({i, cap}); + topo.all_cores.push_back(i); + continue; + } + + snprintf(path, sizeof(path), + "/sys/devices/system/cpu/cpu%d/cpufreq/cpuinfo_max_freq", i); + int freq = read_sysfs_int(path); + if (freq > 0) { + core_caps.push_back({i, freq}); + topo.all_cores.push_back(i); + } + } + + if (core_caps.empty()) return topo; + + int max_cap = 0; + for (auto& [id, cap] : core_caps) { + max_cap = std::max(max_cap, cap); + } + + int threshold = static_cast(max_cap * 0.70); + for (auto& [id, cap] : core_caps) { + if (cap >= threshold) { + topo.performance_cores.push_back(id); + topo.performance_core_capacities.push_back(cap); + } + } + + return topo; + } + }; + + inline bool pin_current_thread_to_cores(const std::vector& cores) { + if (cores.empty()) return false; + cpu_set_t current_mask; + const bool has_current_mask = sched_getaffinity(0, sizeof(current_mask), ¤t_mask) == 0; + cpu_set_t mask; + CPU_ZERO(&mask); + bool selected = false; + for (int core : cores) { + if (!has_current_mask || CPU_ISSET(core, ¤t_mask)) { + CPU_SET(core, &mask); + selected = true; + } + } + if (!selected && has_current_mask) return true; + return sched_setaffinity(0, sizeof(mask), &mask) == 0; + } + + struct ThreadAffinityState { + cpu_set_t original_mask{}; + bool has_original_mask{false}; + + static ThreadAffinityState& current() { + static thread_local ThreadAffinityState state; + return state; + } + + void capture_once() { + if (has_original_mask) return; + has_original_mask = sched_getaffinity(0, sizeof(original_mask), &original_mask) == 0; + } + + void restore() const { + if (has_original_mask) { + sched_setaffinity(0, sizeof(original_mask), &original_mask); + } + } + }; + + struct CpuTimeSample { + uint64_t idle{0}; + uint64_t total{0}; + bool valid{false}; + }; + + inline std::vector read_cpu_time_samples() { + std::ifstream f("/proc/stat"); + std::vector samples; + std::string label; + while (f >> label) { + if (label.size() <= 3 || label[0] != 'c' || label[1] != 'p' || label[2] != 'u') { + std::string rest; + std::getline(f, rest); + continue; + } + + int cpu = 0; + for (size_t i = 3; i < label.size(); ++i) { + if (label[i] < '0' || label[i] > '9') { + cpu = -1; + break; + } + cpu = cpu * 10 + (label[i] - '0'); + } + if (cpu < 0) continue; + + uint64_t user = 0, nice = 0, system = 0, idle = 0, iowait = 0; + uint64_t irq = 0, softirq = 0, steal = 0, guest = 0, guest_nice = 0; + f >> user >> nice >> system >> idle >> iowait >> irq >> softirq >> steal >> guest >> guest_nice; + if (samples.size() <= static_cast(cpu)) samples.resize(static_cast(cpu) + 1); + auto& sample = samples[static_cast(cpu)]; + sample.idle = idle + iowait; + sample.total = user + nice + system + idle + iowait + irq + softirq + steal + guest + guest_nice; + sample.valid = sample.total > 0; + } + return samples; + } + + inline int select_load_aware_performance_core( + const std::vector& before, + const std::vector& after + ) { + auto& topo = CoreTopology::get(); + if (topo.performance_cores.empty()) return -1; + + cpu_set_t current_mask; + const bool has_current_mask = sched_getaffinity(0, sizeof(current_mask), ¤t_mask) == 0; + int best_core = -1; + double best_score = -1.0; + int max_cap = 1; + for (int cap : topo.performance_core_capacities) max_cap = std::max(max_cap, cap); + + for (size_t i = 0; i < topo.performance_cores.size(); ++i) { + int core = topo.performance_cores[i]; + if (has_current_mask && !CPU_ISSET(core, ¤t_mask)) continue; + + double busy = 0.0; + if (static_cast(core) < before.size() && static_cast(core) < after.size()) { + const auto& prev = before[static_cast(core)]; + const auto& curr = after[static_cast(core)]; + if (prev.valid && curr.valid && curr.total >= prev.total && curr.idle >= prev.idle) { + uint64_t total_delta = curr.total - prev.total; + uint64_t idle_delta = curr.idle - prev.idle; + if (total_delta > 0 && idle_delta <= total_delta) { + busy = 1.0 - (static_cast(idle_delta) / static_cast(total_delta)); + } + } + } + + int cap = i < topo.performance_core_capacities.size() ? topo.performance_core_capacities[i] : max_cap; + double score = (static_cast(cap) / static_cast(max_cap)) / (1.0 + busy); + if (score > best_score || (score == best_score && core > best_core)) { + best_score = score; + best_core = core; + } + } + return best_core; + } + + inline void prepare_current_thread_for_cactus_work() { + auto& affinity = ThreadAffinityState::current(); + affinity.capture_once(); + (void)get_thread_pool(); + affinity.restore(); + + auto before = read_cpu_time_samples(); + if (!before.empty()) { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + int core = select_load_aware_performance_core(before, read_cpu_time_samples()); + if (core >= 0) { + pin_current_thread_to_cores({core}); + } + } +#else + inline void prepare_current_thread_for_cactus_work() {} +#endif + + class ThreadPool { + private: + static constexpr size_t MAX_WORKERS = 16; + + std::vector workers; + std::deque> tasks; + + std::mutex mutex; + std::condition_variable work_available; + std::condition_variable work_done; + + bool stop{false}; + std::atomic pending_tasks{0}; + size_t num_workers_; + + void worker_thread() { + while (true) { + std::function task; + { + std::unique_lock lock(mutex); + work_available.wait(lock, [this] { + return stop || !tasks.empty(); + }); + + if (stop && tasks.empty()) { + return; + } + + task = std::move(tasks.front()); + tasks.pop_front(); + } + + task(); + + if (pending_tasks.fetch_sub(1, std::memory_order_acq_rel) == 1) { + std::lock_guard lock(mutex); + work_done.notify_one(); + } + } + } + + public: + explicit ThreadPool(size_t num_threads = std::thread::hardware_concurrency()) + : stop(false), pending_tasks(0) { + num_workers_ = std::min(num_threads, MAX_WORKERS); + if (num_workers_ == 0) num_workers_ = 1; + +#if defined(__ANDROID__) + auto& topo = CoreTopology::get(); + if (!topo.performance_cores.empty()) { + num_workers_ = std::min(num_workers_, topo.performance_cores.size()); + } +#endif + + workers.reserve(num_workers_); + for (size_t i = 0; i < num_workers_; ++i) { + workers.emplace_back([this, i]() { +#if defined(__ANDROID__) + auto& perf = CoreTopology::get().performance_cores; + if (!perf.empty()) { + pin_current_thread_to_cores({perf[i % perf.size()]}); + } +#else + (void)i; +#endif + worker_thread(); + }); + } + } + + ~ThreadPool() { + { + std::lock_guard lock(mutex); + stop = true; + } + work_available.notify_all(); + for (auto& worker : workers) { + if (worker.joinable()) { + worker.join(); + } + } + } + + template + auto enqueue(F&& f) -> std::future { + using return_type = decltype(f()); + + auto task = std::make_shared>( + std::forward(f) + ); + + std::future res = task->get_future(); + + { + std::lock_guard lock(mutex); + pending_tasks.fetch_add(1, std::memory_order_relaxed); + tasks.emplace_back([task](){ (*task)(); }); + } + work_available.notify_one(); + + return res; + } + + template + void enqueue_batch(size_t total_work, F task_func) { + if (total_work == 0) return; + + const size_t num_tasks = std::min(num_workers_, total_work); + const size_t per_worker = total_work / num_tasks; + const size_t remainder = total_work % num_tasks; + + { + std::lock_guard lock(mutex); + pending_tasks.fetch_add(num_tasks, std::memory_order_relaxed); + + for (size_t w = 0; w < num_tasks; ++w) { + size_t start = w * per_worker + std::min(w, remainder); + size_t end = start + per_worker + (w < remainder ? 1 : 0); + tasks.emplace_back([=]() { task_func(start, end); }); + } + } + work_available.notify_all(); + } + + void wait_all() { + std::unique_lock lock(mutex); + work_done.wait(lock, [this] { + return pending_tasks.load(std::memory_order_acquire) == 0; + }); + } + + template + void enqueue_n_threads(size_t total_work, size_t num_threads, F task_func) { + if (total_work == 0 || num_threads == 0) return; + + num_threads = std::min(num_threads, std::min(num_workers_, total_work)); + size_t num_tasks = num_threads; +#if defined(__ANDROID__) + if (num_threads > 1) { + num_tasks = std::min(total_work, std::max(num_threads, num_threads * ANDROID_DYNAMIC_CHUNK_MULTIPLIER)); + } +#endif + const size_t per_task = total_work / num_tasks; + const size_t remainder = total_work % num_tasks; + + { + std::lock_guard lock(mutex); + pending_tasks.fetch_add(num_tasks, std::memory_order_relaxed); + + for (size_t t = 0; t < num_tasks; ++t) { + size_t start = t * per_task + std::min(t, remainder); + size_t end = start + per_task + (t < remainder ? 1 : 0); + tasks.emplace_back([=]() { task_func(start, end); }); + } + } + work_available.notify_all(); + } + + size_t num_workers() const { return num_workers_; } + }; + + inline ThreadPool& get_thread_pool() { + static ThreadPool pool; + return pool; + } + + struct ParallelConfig { + size_t min_work_gate; + size_t work_per_thread; + + constexpr ParallelConfig(size_t gate, size_t per_thread) + : min_work_gate(gate), work_per_thread(per_thread) {} + }; + + inline size_t get_optimal_thread_count(size_t total_work, ParallelConfig config) { + if (total_work < config.min_work_gate) return 1; + + size_t pool_size = get_thread_pool().num_workers(); + size_t num_threads = (total_work + config.work_per_thread - 1) / config.work_per_thread; + return std::min(pool_size, std::max(static_cast(1), num_threads)); + } + + struct Thresholds { + #if defined(__ANDROID__) + static constexpr ParallelConfig ATTENTION{64, 32}; + static constexpr ParallelConfig ELEMENT_WISE{5000, 2500}; + static constexpr ParallelConfig AXIS_REDUCE{1000, 500}; + static constexpr ParallelConfig ALL_REDUCE{10000, 5000}; + static constexpr ParallelConfig SCALAR_BASIC{30000, 15000}; + static constexpr ParallelConfig SCALAR_EXPENSIVE{10000, 5000}; + #else // Apple + static constexpr ParallelConfig ATTENTION{32, 16}; + static constexpr ParallelConfig ELEMENT_WISE{5000, 2500}; + static constexpr ParallelConfig AXIS_REDUCE{1000, 500}; + static constexpr ParallelConfig ALL_REDUCE{10000, 5000}; + static constexpr ParallelConfig SCALAR_BASIC{5000, 2500}; + static constexpr ParallelConfig SCALAR_EXPENSIVE{2500, 1250}; + #endif + }; + + struct GemmThreading { + #if defined(__ANDROID__) + static size_t get_num_threads(size_t M, size_t pool_size) { + if (M <= 1) return 1; + return pool_size; + } + static size_t get_gemv_threads(size_t /*N_blocks*/, size_t /*pool_size*/) { + return 1; + } + #elif defined(__APPLE__) && TARGET_OS_IPHONE + static constexpr size_t GEMV_MIN_N_BLOCKS = 512; + static size_t get_num_threads(size_t M, size_t pool_size) { + if (M <= 1) return std::min(pool_size, static_cast(2)); + return pool_size; + } + static size_t get_gemv_threads(size_t N_blocks, size_t pool_size) { + if (N_blocks < GEMV_MIN_N_BLOCKS) return 1; + return std::min(pool_size, static_cast(3)); + } + #else + static constexpr size_t GEMV_MIN_N_BLOCKS = 256; + static size_t get_num_threads(size_t M, size_t pool_size) { + if (M <= 1) return std::min(pool_size, static_cast(4)); + return pool_size; + } + static size_t get_gemv_threads(size_t N_blocks, size_t pool_size) { + if (N_blocks < GEMV_MIN_N_BLOCKS) return 1; + if (N_blocks < 512) return std::min(pool_size, static_cast(2)); + return std::min(pool_size, static_cast(5)); + } + #endif + }; + + inline size_t& get_gemm_thread_override() { + static size_t override_threads = 0; + return override_threads; + } + + inline void set_gemm_threads(size_t num_threads) { + get_gemm_thread_override() = num_threads; + } + + inline void reset_gemm_threads() { + get_gemm_thread_override() = 0; + } + + class TaskHandle { + private: + std::vector> futures_; + bool auto_wait_; + + public: + TaskHandle(bool auto_wait = true) : auto_wait_(auto_wait) {} + + ~TaskHandle() { + if (auto_wait_) { + wait(); + } + } + + TaskHandle(TaskHandle&&) = default; + TaskHandle& operator=(TaskHandle&&) = default; + TaskHandle(const TaskHandle&) = delete; + TaskHandle& operator=(const TaskHandle&) = delete; + + void add_future(std::future&& f) { + futures_.push_back(std::move(f)); + } + + void wait() { + for (auto& f : futures_) { + if (f.valid()) { + f.wait(); + } + } + futures_.clear(); + } + + bool is_ready() const { + for (const auto& f : futures_) { + if (f.valid() && f.wait_for(std::chrono::seconds(0)) != std::future_status::ready) { + return false; + } + } + return true; + } + + size_t task_count() const { return futures_.size(); } + }; + + template + TaskHandle parallel_for(size_t total_work, ParallelConfig config, WorkFunc work_func, bool wait = true) { + const size_t num_threads = get_optimal_thread_count(total_work, config); + TaskHandle handle(!wait); + + if (num_threads == 1) { + if (wait) { + work_func(0, total_work); + return handle; + } + auto& pool = get_thread_pool(); + handle.add_future(pool.enqueue([work_func, total_work]() { + work_func(0, total_work); + })); + return handle; + } + + auto& pool = get_thread_pool(); +#if defined(__ANDROID__) + if (wait) { + pool.enqueue_n_threads(total_work, num_threads, work_func); + pool.wait_all(); + return handle; + } +#endif + + const size_t work_per_thread = total_work / num_threads; + + for (size_t t = 0; t < num_threads; ++t) { + handle.add_future(pool.enqueue([work_func, t, num_threads, work_per_thread, total_work]() { + const size_t start_idx = t * work_per_thread; + const size_t end_idx = (t == num_threads - 1) ? total_work : (t + 1) * work_per_thread; + work_func(start_idx, end_idx); + })); + } + + if (wait) { + handle.wait(); + } + return handle; + } + + template + void parallel_for_2d(size_t outer_size, size_t inner_size, ParallelConfig config, WorkFunc work_func) { + const size_t total_work = outer_size * inner_size; + parallel_for(total_work, config, [&](size_t start_idx, size_t end_idx) { + for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { + const size_t outer = work_idx / inner_size; + const size_t inner = work_idx % inner_size; + work_func(outer, inner); + } + }); + } + + template + ResultType parallel_reduce(size_t total_work, ParallelConfig config, + WorkFunc work_func, ResultType init_value, CombineFunc combine_func) { + const size_t num_threads = get_optimal_thread_count(total_work, config); + + if (num_threads == 1) { + return work_func(0, total_work); + } + + auto& pool = get_thread_pool(); + std::vector> futures; + std::vector partial_results(num_threads, init_value); + const size_t work_per_thread = total_work / num_threads; + + for (size_t t = 0; t < num_threads; ++t) { + futures.push_back(pool.enqueue([work_func, t, num_threads, work_per_thread, total_work]() -> ResultType { + const size_t start_idx = t * work_per_thread; + const size_t end_idx = (t == num_threads - 1) ? total_work : (t + 1) * work_per_thread; + return work_func(start_idx, end_idx); + })); + } + + ResultType result = init_value; + for (auto& future : futures) { + result = combine_func(result, future.get()); + } + return result; + } + + template + void parallel_gemm_tiles(size_t M, size_t total_tiles, WorkFunc work_func) { + auto& pool = get_thread_pool(); + + size_t override = get_gemm_thread_override(); + size_t num_threads = (override > 0) ? override : GemmThreading::get_num_threads(M, pool.num_workers()); + num_threads = std::min(num_threads, total_tiles); + + if (num_threads <= 1) { + work_func(0, total_tiles); + return; + } + + pool.enqueue_n_threads(total_tiles, num_threads, work_func); + pool.wait_all(); + } + +} + +template +void elementwise_op_f16(const __fp16* input, __fp16* output, size_t num_elements, + bool use_streaming, CactusThreading::ParallelConfig config, + SimdOp simd_op, ScalarOp scalar_op, size_t unroll = 4) { + CactusThreading::parallel_for(num_elements, config, + [&](size_t start, size_t end) { + const size_t n = end - start; + const size_t vec_end = start + simd_align(n); + + if (use_streaming && unroll >= 4) { + const size_t unrolled_end = start + simd_align(n, SIMD_F16_WIDTH * 4); + for (size_t i = start; i < unrolled_end; i += SIMD_F16_WIDTH * 4) { + __builtin_prefetch(&input[i + 256], 0, 0); + float16x8_t v0 = simd_op(vld1q_f16(&input[i])); + float16x8_t v1 = simd_op(vld1q_f16(&input[i + 8])); + float16x8_t v2 = simd_op(vld1q_f16(&input[i + 16])); + float16x8_t v3 = simd_op(vld1q_f16(&input[i + 24])); + stream_store_f16x8(&output[i], v0); + stream_store_f16x8(&output[i + 8], v1); + stream_store_f16x8(&output[i + 16], v2); + stream_store_f16x8(&output[i + 24], v3); + } + for (size_t i = unrolled_end; i < vec_end; i += SIMD_F16_WIDTH) { + stream_store_f16x8(&output[i], simd_op(vld1q_f16(&input[i]))); + } + } else if (use_streaming && unroll >= 2) { + const size_t unrolled_end = start + simd_align(n, SIMD_F16_WIDTH * 2); + for (size_t i = start; i < unrolled_end; i += SIMD_F16_WIDTH * 2) { + __builtin_prefetch(&input[i + 128], 0, 0); + float16x8_t v0 = simd_op(vld1q_f16(&input[i])); + float16x8_t v1 = simd_op(vld1q_f16(&input[i + 8])); + stream_store_f16x8(&output[i], v0); + stream_store_f16x8(&output[i + 8], v1); + } + for (size_t i = unrolled_end; i < vec_end; i += SIMD_F16_WIDTH) { + stream_store_f16x8(&output[i], simd_op(vld1q_f16(&input[i]))); + } + } else { + for (size_t i = start; i < vec_end; i += SIMD_F16_WIDTH) { + vst1q_f16(&output[i], simd_op(vld1q_f16(&input[i]))); + } + } + for (size_t i = vec_end; i < end; ++i) { + output[i] = scalar_op(input[i]); + } + }); +} + +#endif // KERNEL_UTILS_H diff --git a/libs/audio/wav.h b/cactus-kernels/src/wav.h similarity index 61% rename from libs/audio/wav.h rename to cactus-kernels/src/wav.h index d67311fbe..576ef43c1 100644 --- a/libs/audio/wav.h +++ b/cactus-kernels/src/wav.h @@ -70,12 +70,13 @@ inline AudioFP32 load_wav_fp32(const std::string& path) { } size_t num_samples = data_size / 2; - std::vector tmp(num_samples); + std::vector raw(num_samples); + file.read(reinterpret_cast(raw.data()), data_size); + std::vector tmp(num_samples); + constexpr float scale = 1.0f / 32768.0f; for (size_t i = 0; i < num_samples; i++) { - int16_t s; - file.read(reinterpret_cast(&s), 2); - tmp[i] = float(s) / 32768.0f; + tmp[i] = float(raw[i]) * scale; } std::vector mono; @@ -103,19 +104,58 @@ inline std::vector resample_to_16k_fp32( const int sr_out = 16000; if (sr_in == sr_out) return in; - double ratio = double(sr_out) / double(sr_in); - size_t out_len = size_t(in.size() * ratio); + const double ratio = double(sr_out) / double(sr_in); + const size_t out_len = size_t(in.size() * ratio); + const size_t in_len = in.size(); + + if (sr_in <= sr_out) { + std::vector out(out_len); + for (size_t i = 0; i < out_len; i++) { + double pos = i / ratio; + size_t i0 = (size_t)pos; + double frac = pos - i0; + out[i] = (i0 + 1 < in_len) + ? float((1.0 - frac) * in[i0] + frac * in[i0 + 1]) + : in.back(); + } + return out; + } + + const double cutoff = ratio; + constexpr int SINC_HALF_WIDTH = 16; + const double M_PI_VAL = 3.14159265358979323846; std::vector out(out_len); for (size_t i = 0; i < out_len; i++) { - double pos = i / ratio; - size_t i0 = (size_t)pos; - double frac = pos - i0; - - out[i] = (i0 + 1 < in.size()) - ? float((1.0 - frac) * in[i0] + frac * in[i0 + 1]) - : in.back(); + double center = i / ratio; + int left = (int)std::ceil(center) - SINC_HALF_WIDTH; + int right = (int)std::floor(center) + SINC_HALF_WIDTH; + + double sum = 0.0; + double weight_sum = 0.0; + + for (int j = left; j <= right; j++) { + if (j < 0 || j >= (int)in_len) continue; + + double x = center - j; + double sinc_val; + if (std::fabs(x) < 1e-9) { + sinc_val = cutoff; + } else { + double pix = M_PI_VAL * x * cutoff; + sinc_val = cutoff * std::sin(pix) / pix; + } + + double window_pos = (x + SINC_HALF_WIDTH) / (2.0 * SINC_HALF_WIDTH); + double window = 0.5 * (1.0 - std::cos(2.0 * M_PI_VAL * window_pos)); + + double w = sinc_val * window; + sum += in[j] * w; + weight_sum += w; + } + + out[i] = (weight_sum > 1e-9) ? float(sum / weight_sum) : 0.0f; } return out; } diff --git a/cactus-kernels/test.sh b/cactus-kernels/test.sh new file mode 100755 index 000000000..9c6d1ca01 --- /dev/null +++ b/cactus-kernels/test.sh @@ -0,0 +1,40 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")" + +SUITE="" +while [[ $# -gt 0 ]]; do + case $1 in + --suite) SUITE="${2:?--suite needs an argument}"; shift 2 ;; + *) echo "Unknown arg: $1" >&2; exit 2 ;; + esac +done + +echo "Building and testing cactus-kernels..." + +rm -rf build +mkdir -p build +cd build + +cmake .. -DCACTUS_BUILD_TESTS=ON -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null +make -j$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) + +echo "" +FAILED=0 +if [ -n "$SUITE" ]; then + target="./test_${SUITE}" + if [ -x "$target" ]; then + "$target" || FAILED=1 + else + echo "Test not found: $target" >&2 + FAILED=1 + fi +else + for t in ./test_*; do + [ -x "$t" ] || continue + "$t" || FAILED=1 + done +fi + +exit $FAILED diff --git a/cactus-kernels/tests/test_attention.cpp b/cactus-kernels/tests/test_attention.cpp new file mode 100644 index 000000000..e6089ba6a --- /dev/null +++ b/cactus-kernels/tests/test_attention.cpp @@ -0,0 +1,138 @@ +#include "test_utils.h" +#include +#include + +using namespace TestUtils; + +bool test_rms_norm() { + const size_t batch = 4, dim = 128; + std::vector<__fp16> input(batch * dim), weight(dim), output(batch * dim); + fill_random_fp16(input, -1.0f, 1.0f); + for (size_t i = 0; i < dim; i++) weight[i] = static_cast<__fp16>(1.0f); + cactus_rms_norm_f16(input.data(), weight.data(), output.data(), batch, dim, 1e-6f); + for (size_t b = 0; b < batch; b++) { + float sum_sq = 0.0f; + for (size_t d = 0; d < dim; d++) { float v = static_cast(output[b * dim + d]); sum_sq += v * v; } + if (std::abs(std::sqrt(sum_sq / dim) - 1.0f) > 0.05f) return false; + } + return true; +} + +bool test_layer_norm() { + const size_t batch = 4, dim = 128; + std::vector<__fp16> input(batch * dim), weight(dim), bias(dim), output(batch * dim); + fill_random_fp16(input, -1.0f, 1.0f); + for (size_t i = 0; i < dim; i++) { weight[i] = static_cast<__fp16>(1.0f); bias[i] = static_cast<__fp16>(0.0f); } + cactus_layer_norm_f16(input.data(), weight.data(), bias.data(), output.data(), batch, dim, 1e-5f); + for (size_t b = 0; b < batch; b++) { + float mean = 0.0f; + for (size_t d = 0; d < dim; d++) mean += static_cast(output[b * dim + d]); + mean /= dim; + if (std::abs(mean) > 0.05f) return false; + } + return true; +} + +bool test_softmax() { + const size_t rows = 8, vocab = 128; + std::vector<__fp16> input(rows * vocab), output(rows * vocab); + fill_random_fp16(input, -5.0f, 5.0f); + cactus_softmax_f16(input.data(), output.data(), rows, 1, vocab); + for (size_t r = 0; r < rows; r++) { + float sum = 0.0f; + for (size_t j = 0; j < vocab; j++) { + float v = static_cast(output[r * vocab + j]); + if (v < 0.0f || v > 1.0f) return false; + sum += v; + } + if (std::abs(sum - 1.0f) > 0.01f) return false; + } + return true; +} + +bool test_rope() { + const size_t batch = 1, seq = 4, heads = 2, dim = 16; + std::vector<__fp16> input(batch * seq * heads * dim), output(batch * seq * heads * dim); + for (auto& v : input) v = static_cast<__fp16>(1.0f); + cactus_rope_f16(input.data(), output.data(), batch, seq, heads, dim, 0, 10000.0f); + for (size_t i = 0; i < heads * dim; i++) + if (std::abs(static_cast(output[i]) - 1.0f) > 0.01f) return false; + bool changed = false; + for (size_t i = heads * dim; i < output.size(); i++) + if (std::abs(static_cast(output[i]) - 1.0f) > 0.001f) { changed = true; break; } + return changed; +} + +bool test_attention_f16() { + const size_t batch = 1, seq = 8, heads = 2, kv_heads = 2, dim = 16; + std::vector<__fp16> q(batch * seq * heads * dim), k(batch * seq * kv_heads * dim); + std::vector<__fp16> v(batch * seq * kv_heads * dim), out(batch * seq * heads * dim); + fill_random_fp16(q, -0.5f, 0.5f); fill_random_fp16(k, -0.5f, 0.5f); fill_random_fp16(v, -0.5f, 0.5f); + float scale = 1.0f / std::sqrt(static_cast(dim)); + cactus_attention_f16(q.data(), k.data(), v.data(), out.data(), batch, seq, seq, heads, kv_heads, dim, scale, + nullptr, 0, 0, true, false, false, 0, 0.0f); + for (size_t i = 0; i < out.size(); i++) + if (!std::isfinite(static_cast(out[i]))) return false; + return true; +} + +bool run_benchmarks() { + auto bench = [](const char* label, auto fn) { + fn(); + Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + { + const size_t b = 1024, d = 1024; + std::vector<__fp16> in(b * d), w(d), out(b * d); + fill_random_fp16(in); for (size_t i = 0; i < d; i++) w[i] = static_cast<__fp16>(1.0f); + bench("rms_norm 1024x1024", [&]{ cactus_rms_norm_f16(in.data(), w.data(), out.data(), b, d, 1e-6f); }); + } + { + const size_t b = 1024, d = 1024; + std::vector<__fp16> in(b * d), w(d), bias(d), out(b * d); + fill_random_fp16(in); + for (size_t i = 0; i < d; i++) { w[i] = static_cast<__fp16>(1.0f); bias[i] = static_cast<__fp16>(0.0f); } + bench("layer_norm 1024x1024", [&]{ cactus_layer_norm_f16(in.data(), w.data(), bias.data(), out.data(), b, d, 1e-5f); }); + } + { + const size_t rows = 1024, cols = 1024; + std::vector<__fp16> in(rows * cols), out(rows * cols); + fill_random_fp16(in, -5.0f, 5.0f); + bench("softmax 1024x1024", [&]{ cactus_softmax_f16(in.data(), out.data(), rows, 1, cols); }); + } + { + const size_t b = 1, s = 256, h = 16, d = 128; + std::vector<__fp16> in(b * s * h * d), out(b * s * h * d); + fill_random_fp16(in, -0.5f, 0.5f); + bench("rope 256x16x128", [&]{ cactus_rope_f16(in.data(), out.data(), b, s, h, d, 0, 10000.0f); }); + } + { + const size_t b = 1, s = 256, h = 16, kv = 8, d = 128; + std::vector<__fp16> q(b * s * h * d), k(b * s * kv * d), v(b * s * kv * d), out(b * s * h * d); + fill_random_fp16(q, -0.3f, 0.3f); fill_random_fp16(k, -0.3f, 0.3f); fill_random_fp16(v, -0.3f, 0.3f); + float sc = 1.0f / std::sqrt(static_cast(d)); + bench("attention seq=256 h=16 d=128", [&]{ + cactus_attention_f16(q.data(), k.data(), v.data(), out.data(), b, s, s, h, kv, d, sc, + nullptr, 0, 0, true, false, false, 0, 0.0f); + }); + } + return true; +} + +int main() { + TestRunner runner("Attention, RoPE & Normalization"); + runner.run_test("rms_norm", test_rms_norm()); + runner.run_test("layer_norm", test_layer_norm()); + runner.run_test("softmax", test_softmax()); + runner.run_test("rope", test_rope()); + runner.run_test("attention_f16", test_attention_f16()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_conv.cpp b/cactus-kernels/tests/test_conv.cpp new file mode 100644 index 000000000..42b916c1e --- /dev/null +++ b/cactus-kernels/tests/test_conv.cpp @@ -0,0 +1,152 @@ +#include "test_utils.h" +#include +#include + +using namespace TestUtils; + +bool test_conv1d_k3() { + const size_t N = 1, L = 32, C_in = 8, C_out = 16, stride = 1; + const size_t L_out = ((L - 1) / stride) + 1; + std::vector<__fp16> input(N * L * C_in), weight(C_out * 3 * C_in), output(N * L_out * C_out); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + + cactus_conv1d_f16_k3(input.data(), weight.data(), output.data(), N, L, C_in, C_out, stride); + + for (size_t i = 0; i < N * L_out * C_out; i++) { + if (!std::isfinite(static_cast(output[i]))) { + std::cerr << " conv1d_k3: non-finite at " << i << "\n"; + return false; + } + } + return true; +} + +bool test_conv1d_causal_depthwise() { + const size_t N = 1, L = 16, C = 32, K = 3; + std::vector<__fp16> input(N * L * C), weight(C * K), output(N * L * C); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + + cactus_conv1d_causal_depthwise_f16(input.data(), weight.data(), output.data(), N, L, C, K, 1); + + for (size_t i = 0; i < N * L * C; i++) { + if (!std::isfinite(static_cast(output[i]))) { + std::cerr << " conv1d_causal: non-finite at " << i << "\n"; + return false; + } + } + return true; +} + +bool test_stft_complex() { + const size_t N = 1, L = 128, C_in = 1, K = 64, stride = 32; + const size_t num_fft_bins = K / 2; + const size_t C_out = K; + const size_t num_frames = (L - K) / stride + 1; + std::vector<__fp16> input(N * L * C_in), weight(C_out * K * C_in); + std::vector<__fp16> output(N * num_frames * num_fft_bins * 2); + fill_random_fp16(input, -1.0f, 1.0f); + fill_random_fp16(weight, -0.5f, 0.5f); + + cactus_stft_f16(input.data(), weight.data(), output.data(), N, L, C_in, C_out, K, stride, num_fft_bins); + + for (size_t i = 0; i < output.size(); i++) { + if (!std::isfinite(static_cast(output[i]))) { + std::cerr << " stft: non-finite at " << i << "\n"; + return false; + } + } + return true; +} + +bool test_maxpool1d() { + const size_t batch = 1, channels = 4, length = 16, kernel = 3, stride = 2; + const size_t out_len = (length - kernel) / stride + 1; + std::vector<__fp16> input(batch * channels * length), output(batch * channels * out_len); + fill_random_fp16(input, -5.0f, 5.0f); + + cactus_maxpool1d_f16(input.data(), output.data(), batch, channels, length, kernel, stride); + + for (size_t c = 0; c < channels; c++) { + for (size_t o = 0; o < out_len; o++) { + float expected_max = -1e9f; + for (size_t k = 0; k < kernel; k++) { + float v = static_cast(input[c * length + o * stride + k]); + if (v > expected_max) expected_max = v; + } + float actual = static_cast(output[c * out_len + o]); + if (std::abs(actual - expected_max) > 0.01f) { + std::cerr << " maxpool1d mismatch [c=" << c << ",o=" << o << "]: " + << actual << " vs " << expected_max << "\n"; + return false; + } + } + } + return true; +} + +bool run_benchmarks() { + { + const size_t N = 1, L = 3000, C_in = 80, C_out = 512, stride = 1; + const size_t L_out = ((L - 1) / stride) + 1; + std::vector<__fp16> input(N * L * C_in), weight(C_out * 3 * C_in), output(N * L_out * C_out); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + cactus_conv1d_f16_k3(input.data(), weight.data(), output.data(), N, L, C_in, C_out, stride); + Timer t; + for (int i = 0; i < 100; i++) + cactus_conv1d_f16_k3(input.data(), weight.data(), output.data(), N, L, C_in, C_out, stride); + double ms = t.elapsed_ms() / 100.0; + double gflops = (2.0 * N * L_out * C_out * 3 * C_in) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "conv1d_k3 1x3000x80->512" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + } + { + const size_t N = 1, L = 256, C = 128, K = 3, stride = 1; + std::vector<__fp16> input(N * L * C), weight(C * K), output(N * L * C); + fill_random_fp16(input, -0.5f, 0.5f); + fill_random_fp16(weight, -0.5f, 0.5f); + cactus_conv1d_causal_depthwise_f16(input.data(), weight.data(), output.data(), N, L, C, K, stride); + Timer t; + for (int i = 0; i < 100; i++) + cactus_conv1d_causal_depthwise_f16(input.data(), weight.data(), output.data(), N, L, C, K, stride); + double ms = t.elapsed_ms() / 100.0; + double gflops = (2.0 * N * L * C * K) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "conv1d_causal 1x256x128 k3" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + } + { + const size_t N = 1, L = 1024, C_in = 1, K = 64, stride = 32; + const size_t num_fft_bins = K / 2; + const size_t C_out = K; + const size_t num_frames = (L - K) / stride + 1; + std::vector<__fp16> input(N * L * C_in), weight(C_out * K * C_in); + std::vector<__fp16> output(N * num_frames * num_fft_bins * 2); + fill_random_fp16(input, -1.0f, 1.0f); + fill_random_fp16(weight, -0.5f, 0.5f); + cactus_stft_f16(input.data(), weight.data(), output.data(), N, L, C_in, C_out, K, stride, num_fft_bins); + Timer t; + for (int i = 0; i < 100; i++) + cactus_stft_f16(input.data(), weight.data(), output.data(), N, L, C_in, C_out, K, stride, num_fft_bins); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(28) << "stft 1x1024x1 k64" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << ms << " ms\n"; + } + return true; +} + +int main() { + TestRunner runner("Convolution & Pooling"); + runner.run_test("conv1d_k3", test_conv1d_k3()); + runner.run_test("conv1d_causal_depthwise", test_conv1d_causal_depthwise()); + runner.run_test("stft_complex", test_stft_complex()); + runner.run_test("maxpool1d", test_maxpool1d()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_dsp.cpp b/cactus-kernels/tests/test_dsp.cpp new file mode 100644 index 000000000..f55bd0fdd --- /dev/null +++ b/cactus-kernels/tests/test_dsp.cpp @@ -0,0 +1,210 @@ +#include "test_utils.h" +#include +#include +#include + +using namespace TestUtils; + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +bool test_rfft_irfft_roundtrip() { + const size_t n = 256; + std::vector input(n); + for (size_t i = 0; i < n; i++) { + input[i] = std::sin(2.0f * static_cast(M_PI) * 10.0f * i / n) + + 0.5f * std::cos(2.0f * static_cast(M_PI) * 30.0f * i / n); + } + + std::vector freq((n / 2 + 1) * 2); + cactus_rfft_f32_1d(input.data(), freq.data(), n, "backward"); + + std::vector recovered(n); + cactus_irfft_f32_1d(freq.data(), recovered.data(), n, "backward"); + + for (size_t i = 0; i < n; i++) { + if (std::abs(recovered[i] - input[i]) > 1e-3f) { + std::cerr << " roundtrip mismatch at " << i << ": " << recovered[i] << " vs " << input[i] << "\n"; + return false; + } + } + return true; +} + +bool test_rfft_dc_signal() { + const size_t n = 64; + std::vector input(n, 3.0f); + std::vector freq((n / 2 + 1) * 2, 0.0f); + + cactus_rfft_f32_1d(input.data(), freq.data(), n, "backward"); + + if (std::abs(freq[0] - 3.0f * n) > 1.0f) { + std::cerr << " DC bin: " << freq[0] << " expected " << 3.0f * n << "\n"; + return false; + } + + for (size_t i = 1; i < n / 2 + 1; i++) { + float mag = std::hypot(freq[i * 2], freq[i * 2 + 1]); + if (mag > 1.0f) { + std::cerr << " non-DC bin " << i << " has magnitude " << mag << "\n"; + return false; + } + } + return true; +} + +bool test_mel_filter_bank() { + const int num_freq_bins = 257; + const int num_mel_filters = 80; + const float min_freq = 0.0f; + const float max_freq = 8000.0f; + const int sampling_rate = 16000; + + std::vector filters(num_mel_filters * num_freq_bins, 0.0f); + cactus_generate_mel_filter_bank( + filters.data(), num_freq_bins, num_mel_filters, + min_freq, max_freq, sampling_rate, "slaney", "slaney", false); + + for (int m = 0; m < num_mel_filters; m++) { + bool has_nonzero = false; + for (int f = 0; f < num_freq_bins; f++) { + float val = filters[m * num_freq_bins + f]; + if (val < 0.0f) { + std::cerr << " negative filter value at mel=" << m << " freq=" << f << "\n"; + return false; + } + if (val > 0.0f) has_nonzero = true; + } + if (!has_nonzero) { + std::cerr << " mel filter " << m << " is all zeros\n"; + return false; + } + } + return true; +} + +bool test_hertz_mel_roundtrip() { + const char* scales[] = {"htk", "kaldi", "slaney"}; + + for (const char* scale : scales) { + for (float freq : {100.0f, 440.0f, 1000.0f, 4000.0f, 8000.0f}) { + float mel = cactus_hertz_to_mel(freq, scale); + float back = cactus_mel_to_hertz(mel, scale); + if (std::abs(back - freq) > 0.5f) { + std::cerr << " roundtrip fail for scale=" << scale << " freq=" << freq + << " got " << back << "\n"; + return false; + } + } + } + return true; +} + +bool test_spectrogram_basic() { + const size_t n_samples = 1600; + const size_t n_fft = 400; + const size_t hop = 160; + const size_t num_freq_bins = n_fft / 2 + 1; + const int num_mel_bins = 80; + + std::vector waveform(n_samples); + for (size_t i = 0; i < n_samples; i++) { + waveform[i] = std::sin(2.0f * static_cast(M_PI) * 440.0f * i / 16000.0f); + } + + std::vector mel_filters(num_mel_bins * num_freq_bins); + cactus_generate_mel_filter_bank( + mel_filters.data(), num_freq_bins, num_mel_bins, + 0.0f, 8000.0f, 16000, "slaney", "slaney", false); + + size_t pad_length = n_fft / 2; + size_t padded_length = n_samples + 2 * pad_length; + size_t num_frames = 1 + (padded_length - n_fft) / hop; + + std::vector spectrogram(num_mel_bins * num_frames); + size_t fft_size = n_fft; + + cactus_compute_spectrogram_f32( + waveform.data(), n_samples, + nullptr, 0, + n_fft, hop, &fft_size, + spectrogram.data(), 2.0f, + true, "reflect", true, + 0.0f, nullptr, + mel_filters.data(), mel_filters.size(), + 1e-10f, nullptr, + 1.0f, 1e-10f, nullptr, false); + + bool has_nonzero = false; + for (size_t i = 0; i < spectrogram.size(); i++) { + if (!std::isfinite(spectrogram[i])) { + std::cerr << " non-finite spectrogram at " << i << "\n"; + return false; + } + if (spectrogram[i] > 1e-8f) has_nonzero = true; + } + + if (!has_nonzero) { + std::cerr << " spectrogram is all zeros\n"; + return false; + } + return true; +} + +bool test_spectrogram_to_db() { + std::vector data = {1.0f, 10.0f, 100.0f, 1000.0f}; + cactus_spectrogram_to_db(data.data(), data.size(), 1.0f, 1e-10f, nullptr, 10.0f); + + float expected[] = {0.0f, 10.0f, 20.0f, 30.0f}; + for (size_t i = 0; i < 4; i++) { + if (std::abs(data[i] - expected[i]) > 0.01f) { + std::cerr << " dB mismatch at " << i << ": " << data[i] << " vs " << expected[i] << "\n"; + return false; + } + } + return true; +} + +bool run_benchmarks() { + auto bench = [](const char* label, auto fn) { + fn(); + TestUtils::Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + std::cout << " ⚡ " << std::left << std::setw(30) << label + << std::fixed << std::setprecision(3) << ms << " ms\n"; + }; + + { + const size_t n = 512; + std::vector input(n, 1.0f); + std::vector output((n / 2 + 1) * 2); + bench("rfft 512", [&]{ cactus_rfft_f32_1d(input.data(), output.data(), n, "backward"); }); + } + + { + const size_t n = 512; + std::vector input((n / 2 + 1) * 2, 0.0f); + input[0] = 1.0f; + std::vector output(n); + bench("irfft 512", [&]{ cactus_irfft_f32_1d(input.data(), output.data(), n, "backward"); }); + } + + return true; +} + +int main() { + TestUtils::TestRunner runner("DSP Kernel Tests"); + + runner.run_test("RFFT/IRFFT Roundtrip", test_rfft_irfft_roundtrip()); + runner.run_test("RFFT DC Signal", test_rfft_dc_signal()); + runner.run_test("Mel Filter Bank", test_mel_filter_bank()); + runner.run_test("Hertz/Mel Roundtrip", test_hertz_mel_roundtrip()); + runner.run_test("Spectrogram Basic", test_spectrogram_basic()); + runner.run_test("Spectrogram to dB", test_spectrogram_to_db()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_elementwise.cpp b/cactus-kernels/tests/test_elementwise.cpp new file mode 100644 index 000000000..bae21e4e8 --- /dev/null +++ b/cactus-kernels/tests/test_elementwise.cpp @@ -0,0 +1,191 @@ +#include "test_utils.h" +#include +#include + +using namespace TestUtils; + +bool test_add_f16() { + const size_t n = 1024; + std::vector<__fp16> a(n), b(n), out(n), expected(n); + fill_random_fp16(a); fill_random_fp16(b); + for (size_t i = 0; i < n; i++) expected[i] = a[i] + b[i]; + cactus_add_f16(a.data(), b.data(), out.data(), n); + return compare_arrays(out.data(), expected.data(), n); +} + +bool test_subtract_f16() { + const size_t n = 1024; + std::vector<__fp16> a(n), b(n), out(n), expected(n); + fill_random_fp16(a); fill_random_fp16(b); + for (size_t i = 0; i < n; i++) expected[i] = a[i] - b[i]; + cactus_subtract_f16(a.data(), b.data(), out.data(), n); + return compare_arrays(out.data(), expected.data(), n); +} + +bool test_multiply_f16() { + const size_t n = 1024; + std::vector<__fp16> a(n), b(n), out(n), expected(n); + fill_random_fp16(a); fill_random_fp16(b); + for (size_t i = 0; i < n; i++) expected[i] = a[i] * b[i]; + cactus_multiply_f16(a.data(), b.data(), out.data(), n); + return compare_arrays(out.data(), expected.data(), n); +} + +bool test_divide_f16() { + const size_t n = 1024; + std::vector<__fp16> a(n), b(n), out(n), expected(n); + fill_random_fp16(a); + fill_random_fp16(b, 0.5f, 2.0f); + for (size_t i = 0; i < n; i++) expected[i] = a[i] / b[i]; + cactus_divide_f16(a.data(), b.data(), out.data(), n); + return compare_arrays(out.data(), expected.data(), n, 5e-2f); +} + +bool test_add_clipped_f16() { + const size_t n = 256; + std::vector<__fp16> a(n), b(n), out(n); + fill_random_fp16(a, -100.0f, 100.0f); + fill_random_fp16(b, -100.0f, 100.0f); + cactus_add_f16_clipped(a.data(), b.data(), out.data(), n); + for (size_t i = 0; i < n; i++) { + if (!std::isfinite(static_cast(out[i]))) return false; + } + return true; +} + +bool test_scalar_ops_f16() { + const size_t n = 256; + std::vector<__fp16> in(n), out(n); + + fill_random_fp16(in, 0.1f, 2.0f); + cactus_scalar_op_f16(in.data(), out.data(), n, 3.0f, ScalarOpType::ADD); + for (size_t i = 0; i < n; i++) { + float diff = std::abs(static_cast(out[i]) - (static_cast(in[i]) + 3.0f)); + if (diff > 0.05f) return false; + } + + fill_random_fp16(in, 0.1f, 2.0f); + cactus_scalar_op_f16(in.data(), out.data(), n, 2.5f, ScalarOpType::MULTIPLY); + for (size_t i = 0; i < n; i++) { + float diff = std::abs(static_cast(out[i]) - (static_cast(in[i]) * 2.5f)); + if (diff > 0.05f) return false; + } + + fill_random_fp16(in, -2.0f, 2.0f); + cactus_scalar_op_f16(in.data(), out.data(), n, 0.0f, ScalarOpType::EXP); + for (size_t i = 0; i < n; i++) { + float actual = static_cast(out[i]); + float ref = std::exp(static_cast(in[i])); + float rel_err = std::abs(actual - ref) / std::max(std::abs(ref), 1e-6f); + if (rel_err > 0.01f) return false; + } + return true; +} + +bool test_transpose_2d_f16() { + const size_t rows = 64, cols = 128; + std::vector<__fp16> src(rows * cols), dst(rows * cols); + fill_random_fp16(src); + cactus_transpose_2d_f16(src.data(), dst.data(), rows, cols, 0, rows); + for (size_t r = 0; r < rows; r++) + for (size_t c = 0; c < cols; c++) + if (std::abs(static_cast(src[r * cols + c]) - static_cast(dst[c * rows + r])) > 1e-4f) + return false; + return true; +} + +bool test_fast_tanh_f32x4() { + constexpr float TOL = 1e-5f; + for (int i = -10000; i <= 10000; ++i) { + float x = 0.001f * static_cast(i); + float32x4_t r = fast_tanh_f32x4(vdupq_n_f32(x)); + float lanes[4]; + vst1q_f32(lanes, r); + float want = std::tanh(x); + for (int k = 0; k < 4; ++k) { + if (std::fabs(lanes[k] - want) > TOL) return false; + } + } + return true; +} + +bool run_benchmarks() { + auto bench_binary = [](const char* label, void(*fn)(const __fp16*, const __fp16*, __fp16*, size_t)) { + const size_t n = 1024 * 1024; + std::vector<__fp16> a(n), b(n), out(n); + fill_random_fp16(a); fill_random_fp16(b); + fn(a.data(), b.data(), out.data(), n); + Timer t; + for (int i = 0; i < 100; i++) fn(a.data(), b.data(), out.data(), n); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (3.0 * n * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << label + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + }; + + bench_binary("add 1M", cactus_add_f16); + bench_binary("subtract 1M", cactus_subtract_f16); + bench_binary("multiply 1M", cactus_multiply_f16); + bench_binary("divide 1M", cactus_divide_f16); + + { + const size_t n = 1024 * 1024; + std::vector<__fp16> in(n), out(n); + fill_random_fp16(in, 0.1f, 2.0f); + cactus_scalar_op_f16(in.data(), out.data(), n, 2.0f, ScalarOpType::MULTIPLY); + Timer t; + for (int i = 0; i < 100; i++) cactus_scalar_op_f16(in.data(), out.data(), n, 2.0f, ScalarOpType::MULTIPLY); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (2.0 * n * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "scalar_multiply 1M" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + + { + const size_t n = 1024 * 1024; + std::vector<__fp16> in(n), out(n); + fill_random_fp16(in, -2.0f, 2.0f); + cactus_scalar_op_f16(in.data(), out.data(), n, 0.0f, ScalarOpType::EXP); + Timer t; + for (int i = 0; i < 100; i++) cactus_scalar_op_f16(in.data(), out.data(), n, 0.0f, ScalarOpType::EXP); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (2.0 * n * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "scalar_exp 1M" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + + { + const size_t rows = 1024, cols = 1024; + std::vector<__fp16> src(rows * cols), dst(rows * cols); + fill_random_fp16(src); + cactus_transpose_2d_f16(src.data(), dst.data(), rows, cols, 0, rows); + Timer t; + for (int i = 0; i < 100; i++) cactus_transpose_2d_f16(src.data(), dst.data(), rows, cols, 0, rows); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (2.0 * rows * cols * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "transpose 1024x1024" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + + return true; +} + +int main() { + TestRunner runner("Elementwise & Scalar Ops"); + runner.run_test("add_f16", test_add_f16()); + runner.run_test("subtract_f16", test_subtract_f16()); + runner.run_test("multiply_f16", test_multiply_f16()); + runner.run_test("divide_f16", test_divide_f16()); + runner.run_test("add_clipped_f16", test_add_clipped_f16()); + runner.run_test("scalar_ops_f16", test_scalar_ops_f16()); + runner.run_test("transpose_2d_f16", test_transpose_2d_f16()); + runner.run_test("fast_tanh_f32x4", test_fast_tanh_f32x4()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_matmul.cpp b/cactus-kernels/tests/test_matmul.cpp new file mode 100644 index 000000000..74f917974 --- /dev/null +++ b/cactus-kernels/tests/test_matmul.cpp @@ -0,0 +1,596 @@ +#include "test_utils.h" +#include +#include +#include +#include + +using namespace TestUtils; + +static uint8_t unpack_index(const uint8_t* base, uint32_t bits, uint32_t k); + +struct SyntheticCQ { + uint32_t bits, K, N, group_size, num_groups; + std::vector<__fp16> codebook; + std::vector<__fp16> input_scale; + std::vector<__fp16> input_scale_recip; + std::vector<__fp16> norms; + std::vector left_signs; + std::vector right_signs; + std::vector permutation; + std::vector packed; + + SyntheticCQ(uint32_t b, uint32_t k, uint32_t n, uint32_t gs, uint32_t seed = 42) + : bits(b), K(k), N(n), group_size(gs), num_groups(k / gs) { + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.f, 1.f); + + uint32_t cb_size = 1u << bits; + codebook.resize(cb_size); + for (auto& v : codebook) v = static_cast<__fp16>(dist(gen)); + + input_scale.resize(K); + input_scale_recip.resize(K); + for (uint32_t i = 0; i < K; i++) { + float s = 0.5f + std::abs(dist(gen)); + input_scale[i] = static_cast<__fp16>(s); + input_scale_recip[i] = static_cast<__fp16>(1.f / s); + } + + norms.resize(size_t(N) * num_groups); + for (auto& v : norms) v = static_cast<__fp16>(dist(gen) * 0.1f); + + left_signs.resize(group_size); + right_signs.resize(group_size); + for (auto& v : left_signs) v = (gen() & 1) ? 1 : -1; + for (auto& v : right_signs) v = (gen() & 1) ? 1 : -1; + + permutation.resize(group_size); + for (uint32_t i = 0; i < group_size; i++) permutation[i] = i; + + size_t packed_bytes = size_t(N) * num_groups * cactus_quant_packed_group_bytes(bits, group_size); + packed.resize(packed_bytes); + for (auto& v : packed) v = static_cast(gen() & 0xFF); + } + + std::vector expanded_buf; + std::vector norm_f32_buf; + + void preexpand() { + int8_t cb_i8[16] = {}; + float cb_max = 0.f; + uint32_t cb_size = 1u << bits; + for (uint32_t i = 0; i < cb_size; i++) { + float v = std::abs(static_cast(codebook[i])); + if (v > cb_max) cb_max = v; + } + float cb_sc = cb_max / 127.f; + if (cb_sc < 1e-10f) cb_sc = 1e-10f; + for (uint32_t i = 0; i < cb_size; i++) + cb_i8[i] = static_cast(std::round(static_cast(codebook[i]) / cb_sc)); + int8x16_t cb_lut = vld1q_s8(cb_i8); + + size_t N_blocks = (N + 3) / 4; + uint32_t pgb = cactus_quant_packed_group_bytes(bits, group_size); + expanded_buf.resize(N_blocks * num_groups * group_size * 4); + norm_f32_buf.resize(N_blocks * num_groups * 4); + + auto expand16 = [&](const uint8_t* p) -> int8x16_t { + if (bits == 4) { + uint8x8_t bytes = vld1_u8(p); + return vqtbl1q_s8(cb_lut, vcombine_u8(vzip1_u8(vand_u8(bytes,vdup_n_u8(0x0F)),vshr_n_u8(bytes,4)), + vzip2_u8(vand_u8(bytes,vdup_n_u8(0x0F)),vshr_n_u8(bytes,4)))); + } else if (bits == 2) { + uint8_t b0=p[0],b1=p[1],b2=p[2],b3=p[3]; + uint64_t lo=((uint64_t)(b0&3))|((uint64_t)((b0>>2)&3)<<8)|((uint64_t)((b0>>4)&3)<<16)|((uint64_t)((b0>>6)&3)<<24)| + ((uint64_t)(b1&3)<<32)|((uint64_t)((b1>>2)&3)<<40)|((uint64_t)((b1>>4)&3)<<48)|((uint64_t)((b1>>6)&3)<<56); + uint64_t hi=((uint64_t)(b2&3))|((uint64_t)((b2>>2)&3)<<8)|((uint64_t)((b2>>4)&3)<<16)|((uint64_t)((b2>>6)&3)<<24)| + ((uint64_t)(b3&3)<<32)|((uint64_t)((b3>>2)&3)<<40)|((uint64_t)((b3>>4)&3)<<48)|((uint64_t)((b3>>6)&3)<<56); + return vqtbl1q_s8(cb_lut, vcombine_u8(vcreate_u8(lo),vcreate_u8(hi))); + } else if (bits == 1) { + uint8_t b0=p[0],b1=p[1]; + uint64_t lo=((uint64_t)((b0>>0)&1))|((uint64_t)((b0>>1)&1)<<8)|((uint64_t)((b0>>2)&1)<<16)|((uint64_t)((b0>>3)&1)<<24)| + ((uint64_t)((b0>>4)&1)<<32)|((uint64_t)((b0>>5)&1)<<40)|((uint64_t)((b0>>6)&1)<<48)|((uint64_t)((b0>>7)&1)<<56); + uint64_t hi=((uint64_t)((b1>>0)&1))|((uint64_t)((b1>>1)&1)<<8)|((uint64_t)((b1>>2)&1)<<16)|((uint64_t)((b1>>3)&1)<<24)| + ((uint64_t)((b1>>4)&1)<<32)|((uint64_t)((b1>>5)&1)<<40)|((uint64_t)((b1>>6)&1)<<48)|((uint64_t)((b1>>7)&1)<<56); + return vqtbl1q_s8(cb_lut, vcombine_u8(vcreate_u8(lo),vcreate_u8(hi))); + } else { + uint64_t raw=0; std::memcpy(&raw,p,6); + uint64_t lo=0,hi=0; + for(int i=0;i<8;i++) lo|=((raw>>(i*3))&7ULL)<<(i*8); + for(int i=0;i<8;i++) hi|=((raw>>((i+8)*3))&7ULL)<<(i*8); + return vqtbl1q_s8(cb_lut, vcombine_u8(vcreate_u8(lo),vcreate_u8(hi))); + } + }; + + for (size_t nb = 0; nb < N_blocks; ++nb) { + size_t n_start = nb * 4; + size_t valid_n = std::min(size_t(4), static_cast(N) - n_start); + for (uint32_t g = 0; g < num_groups; ++g) { + int8x16_t exp4[4][16]; + uint32_t n_vecs = group_size / 16; + for (size_t ni = 0; ni < valid_n; ++ni) { + const uint8_t* pk = packed.data() + (static_cast(n_start+ni)*num_groups+g)*pgb; + for (uint32_t v = 0; v < n_vecs; ++v) + exp4[ni][v] = expand16(pk + (v*16*bits)/8); + } + for (size_t ni = valid_n; ni < 4; ++ni) + for (uint32_t v = 0; v < n_vecs; ++v) exp4[ni][v] = vdupq_n_s8(0); + + int8_t* dst = expanded_buf.data() + (nb*num_groups+g)*group_size*4; + for (uint32_t v = 0; v < n_vecs; ++v) { + int32x4_t r0=vreinterpretq_s32_s8(exp4[0][v]),r1=vreinterpretq_s32_s8(exp4[1][v]); + int32x4_t r2=vreinterpretq_s32_s8(exp4[2][v]),r3=vreinterpretq_s32_s8(exp4[3][v]); + int32x4_t t01l=vzip1q_s32(r0,r1),t01h=vzip2q_s32(r0,r1); + int32x4_t t23l=vzip1q_s32(r2,r3),t23h=vzip2q_s32(r2,r3); + vst1q_s8(dst+v*64, vreinterpretq_s8_s64(vzip1q_s64(vreinterpretq_s64_s32(t01l),vreinterpretq_s64_s32(t23l)))); + vst1q_s8(dst+v*64+16, vreinterpretq_s8_s64(vzip2q_s64(vreinterpretq_s64_s32(t01l),vreinterpretq_s64_s32(t23l)))); + vst1q_s8(dst+v*64+32, vreinterpretq_s8_s64(vzip1q_s64(vreinterpretq_s64_s32(t01h),vreinterpretq_s64_s32(t23h)))); + vst1q_s8(dst+v*64+48, vreinterpretq_s8_s64(vzip2q_s64(vreinterpretq_s64_s32(t01h),vreinterpretq_s64_s32(t23h)))); + } + float* nd = norm_f32_buf.data() + (nb*num_groups+g)*4; + for (size_t ni = 0; ni < 4; ++ni) + nd[ni] = (n_start+ni < N) ? static_cast(norms[(n_start+ni)*num_groups+g]) * cb_sc : 0.f; + } + } + } + + CactusQuantMatrix matrix() const { + return CactusQuantMatrix{ + .bits = bits, .K = K, .N = N, + .group_size = group_size, .num_groups = num_groups, + .flags = 0, + .codebook = codebook.data(), + .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), + .norms = norms.data(), + .packed_indices = packed.data(), + .left_signs = left_signs.data(), + .right_signs = right_signs.data(), + .permutation = permutation.data(), + .rotation = nullptr, + .expanded = expanded_buf.empty() ? nullptr : expanded_buf.data(), + .norm_f32 = norm_f32_buf.empty() ? nullptr : norm_f32_buf.data(), + }; + } + + // INTERLEAVED_4ROW encoding: exact inverse of the shipped decoder. Per 8-byte half, + // low nibbles = k 0-3 and high nibbles = k 4-7 of the half's K-range. + std::vector packed_il; + std::vector<__fp16> norms_il; + + void make_interleaved() { + if (bits != 4 || (N % 4) != 0) return; + const uint32_t pgb = cactus_quant_packed_group_bytes(4, group_size); + const size_t NB = N / 4; + packed_il.assign(NB * num_groups * 4 * (size_t)pgb, 0); + norms_il.resize(NB * num_groups * 4); + for (size_t nb = 0; nb < NB; ++nb) + for (uint32_t g = 0; g < num_groups; ++g) { + uint8_t* panel = packed_il.data() + (nb * num_groups + g) * 4 * (size_t)pgb; + for (uint32_t r = 0; r < 4; ++r) { + const size_t n = nb * 4 + r; + const uint8_t* row = packed.data() + (n * num_groups + g) * pgb; + for (uint32_t v = 0; v < group_size / 16; ++v) + for (uint32_t b = 0; b < 4; ++b) { + auto idx = [&](uint32_t k) { return unpack_index(row, 4, k); }; + panel[(2 * v) * 16 + r * 4 + b] = + (uint8_t)(idx(16 * v + b) | (idx(16 * v + 4 + b) << 4)); + panel[(2 * v + 1) * 16 + r * 4 + b] = + (uint8_t)(idx(16 * v + 8 + b) | (idx(16 * v + 12 + b) << 4)); + } + norms_il[(nb * num_groups + g) * 4 + r] = norms[n * num_groups + g]; + } + } + } + + CactusQuantMatrix matrix_interleaved() { + if (packed_il.empty()) make_interleaved(); + return CactusQuantMatrix{ + .bits = bits, .K = K, .N = N, + .group_size = group_size, .num_groups = num_groups, + .flags = CACTUS_QUANT_FLAG_INTERLEAVED_4ROW, + .codebook = codebook.data(), + .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), + .norms = norms_il.data(), + .packed_indices = packed_il.data(), + .left_signs = left_signs.data(), + .right_signs = right_signs.data(), + .permutation = permutation.data(), + .rotation = nullptr, + .expanded = nullptr, + .norm_f32 = nullptr, + }; + } +}; + +static void fwht_f32(float* x, uint32_t n) { + for (uint32_t h = 1; h < n; h <<= 1) + for (uint32_t i = 0; i < n; i += h << 1) + for (uint32_t j = i; j < i + h; ++j) { + float a = x[j], b = x[j + h]; + x[j] = a + b; x[j + h] = a - b; + } + float s = 1.f / std::sqrt(static_cast(n)); + for (uint32_t i = 0; i < n; ++i) x[i] *= s; +} + +static uint8_t unpack_index(const uint8_t* base, uint32_t bits, uint32_t k) { + switch (bits) { + case 1: return (base[k / 8] >> (k % 8)) & 0x1u; + case 2: return (base[k / 4] >> ((k & 3u) * 2u)) & 0x3u; + case 3: { + uint32_t bit_offset = k * 3; + uint32_t byte_idx = bit_offset / 8; + uint32_t bit_idx = bit_offset % 8; + uint32_t word = static_cast(base[byte_idx]) >> bit_idx; + if (bit_idx > 5) { + word |= static_cast(base[byte_idx + 1]) << (8 - bit_idx); + } + return word & 0x7u; + } + case 4: return (k & 1u) ? (base[k / 2] >> 4) : (base[k / 2] & 0x0Fu); + default: return 0; + } +} + +static void cq_reference_gemv_f32(const SyntheticCQ& w, const float* x, float* y) { + uint32_t pgb = cactus_quant_packed_group_bytes(w.bits, w.group_size); + for (uint32_t n = 0; n < w.N; ++n) { + for (uint32_t g = 0; g < w.num_groups; ++g) { + uint32_t base_k = g * w.group_size; + std::vector z(w.group_size); + for (uint32_t k = 0; k < w.group_size; ++k) + z[k] = x[base_k + k] / static_cast(w.input_scale[base_k + k]) + * static_cast(w.left_signs[k]); + fwht_f32(z.data(), w.group_size); + for (uint32_t k = 0; k < w.group_size; ++k) + z[k] *= static_cast(w.right_signs[k]); + + const uint8_t* packed_row = w.packed.data() + (size_t(n) * w.num_groups + g) * pgb; + float gsum = 0.f; + for (uint32_t k = 0; k < w.group_size; ++k) { + uint8_t idx = unpack_index(packed_row, w.bits, k); + gsum += z[k] * static_cast(w.codebook[idx]); + } + y[n] += static_cast(w.norms[size_t(n) * w.num_groups + g]) * gsum; + } + } +} + +static double compute_mse(const float* ref, const __fp16* actual, size_t n) { + double sum = 0.0; + for (size_t i = 0; i < n; i++) { + double diff = static_cast(ref[i]) - static_cast(actual[i]); + sum += diff * diff; + } + return sum / static_cast(n); +} + +// ══════════════════════════════════════════════════════════════════════════════ +// Correctness tests +// ══════════════════════════════════════════════════════════════════════════════ + +bool test_matmul_f16() { + const size_t M = 4, K = 1024, N = 64; + std::vector<__fp16> a(M * K), b(N * K), c(M * N); + fill_random_fp16(a, -0.5f, 0.5f); + fill_random_fp16(b, -0.5f, 0.5f); + cactus_matmul_f16(a.data(), b.data(), c.data(), M, K, N); + for (size_t i = 0; i < M; i++) + for (size_t j = 0; j < N; j++) { + float ref = 0.0f; + for (size_t k = 0; k < K; k++) + ref += static_cast(a[i * K + k]) * static_cast(b[j * K + k]); + if (std::abs(static_cast(c[i * N + j]) - ref) > 1.0f) return false; + } + return true; +} + +bool test_cq_correctness(uint32_t bits) { + const uint32_t K = 1024, N = 64, gs = 128; + SyntheticCQ cq(bits, K, N, gs, 123); + CactusQuantMatrix mat = cq.matrix(); + + std::mt19937 gen(77); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + + // FP32 reference + std::vector ref(N, 0.f); + cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); + + // FP16 kernel + std::vector<__fp16> x_f16(K), y_f16(N, static_cast<__fp16>(0)); + for (size_t i = 0; i < K; i++) x_f16[i] = static_cast<__fp16>(x_f32[i]); + cactus_quant_matmul(&mat, x_f16.data(), 1, y_f16.data()); + + double mse = compute_mse(ref.data(), y_f16.data(), N); + + double threshold = 0.1; + if (mse > threshold) { + std::cerr << " cq" << bits << " MSE=" << mse << " > " << threshold << "\n"; + return false; + } + return true; +} + +bool run_benchmarks() { + auto bench = [](const char* label, size_t M, size_t K, size_t N, auto fn) { + fn(); + Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + double gflops = (2.0 * M * K * N) / (ms * 1e6); + std::cout << " \u26A1 " << std::left << std::setw(28) << label + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + }; + + const size_t K = 1024, N = 1024; + const size_t M_batch = 1024; + const uint32_t gs = 128; + + // FP16 + { + std::vector<__fp16> a(K), b(N * K), c(N); + fill_random_fp16(a, -0.5f, 0.5f); fill_random_fp16(b, -0.5f, 0.5f); + bench("matmul_f16 1x1024x1024", 1, K, N, [&]{ cactus_matmul_f16(a.data(), b.data(), c.data(), 1, K, N); }); + } + { + std::vector<__fp16> a(M_batch * K), b(N * K), c(M_batch * N); + fill_random_fp16(a, -0.5f, 0.5f); fill_random_fp16(b, -0.5f, 0.5f); + bench("matmul_f16 1024^3", M_batch, K, N, [&]{ cactus_matmul_f16(a.data(), b.data(), c.data(), M_batch, K, N); }); + } + + // TQ1 + { + SyntheticCQ cq(1, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(K), y(N); + fill_random_fp16(x, -1.f, 1.f); + bench("matmul_cq1 1x1024x1024", 1, K, N, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(1, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> A(M_batch * K), C(M_batch * N); + fill_random_fp16(A, -1.f, 1.f); + bench("matmul_cq1 1024^3", M_batch, K, N, [&]{ cactus_quant_matmul(&mat, A.data(), M_batch, C.data()); }); + } + + { + SyntheticCQ cq(2, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(K), y(N); + fill_random_fp16(x, -1.f, 1.f); + bench("matmul_cq2 1x1024x1024", 1, K, N, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(2, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> A(M_batch * K), C(M_batch * N); + fill_random_fp16(A, -1.f, 1.f); + bench("matmul_cq2 1024^3", M_batch, K, N, [&]{ cactus_quant_matmul(&mat, A.data(), M_batch, C.data()); }); + } + + { + SyntheticCQ cq(3, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(K), y(N); + fill_random_fp16(x, -1.f, 1.f); + bench("matmul_cq3 1x1024x1024", 1, K, N, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(3, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> A(M_batch * K), C(M_batch * N); + fill_random_fp16(A, -1.f, 1.f); + bench("matmul_cq3 1024^3", M_batch, K, N, [&]{ cactus_quant_matmul(&mat, A.data(), M_batch, C.data()); }); + } + + { + SyntheticCQ cq(4, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(K), y(N); + fill_random_fp16(x, -1.f, 1.f); + bench("matmul_cq4 1x1024x1024", 1, K, N, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(4, K, N, gs); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> A(M_batch * K), C(M_batch * N); + fill_random_fp16(A, -1.f, 1.f); + bench("matmul_cq4 1024^3", M_batch, K, N, [&]{ cactus_quant_matmul(&mat, A.data(), M_batch, C.data()); }); + } + + auto bench2k = [](const char* label, size_t M, size_t K, size_t N, auto fn) { + fn(); + Timer t; + for (int i = 0; i < 10; i++) fn(); + double ms = t.elapsed_ms() / 10.0; + double gflops = (2.0 * M * K * N) / (ms * 1e6); + std::cout << " \u26A1 " << std::left << std::setw(28) << label + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + }; + + const size_t K2 = 2048, N2 = 2048; + const size_t M2 = 2048; + const uint32_t gs2 = 128; + + { + std::vector<__fp16> a(K2), b(N2 * K2), c(N2); + fill_random_fp16(a, -0.5f, 0.5f); fill_random_fp16(b, -0.5f, 0.5f); + bench2k("matmul_f16 1x2048x2048", 1, K2, N2, [&]{ cactus_matmul_f16(a.data(), b.data(), c.data(), 1, K2, N2); }); + } + { + std::vector<__fp16> a(M2 * K2), b(N2 * K2), c(M2 * N2); + fill_random_fp16(a, -0.5f, 0.5f); fill_random_fp16(b, -0.5f, 0.5f); + bench2k("matmul_f16 2048^3", M2, K2, N2, [&]{ cactus_matmul_f16(a.data(), b.data(), c.data(), M2, K2, N2); }); + } + { + SyntheticCQ cq(2, K2, N2, gs2); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(K2), y(N2); + fill_random_fp16(x, -1.f, 1.f); + bench2k("matmul_cq2 1x2048x2048", 1, K2, N2, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(2, K2, N2, gs2); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> A(M2 * K2), C(M2 * N2); + fill_random_fp16(A, -1.f, 1.f); + bench2k("matmul_cq2 2048^3", M2, K2, N2, [&]{ cactus_quant_matmul(&mat, A.data(), M2, C.data()); }); + } + { + SyntheticCQ cq(4, K2, N2, gs2); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(K2), y(N2); + fill_random_fp16(x, -1.f, 1.f); + bench2k("matmul_cq4 1x2048x2048", 1, K2, N2, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(4, K2, N2, gs2); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> A(M2 * K2), C(M2 * N2); + fill_random_fp16(A, -1.f, 1.f); + bench2k("matmul_cq4 2048^3", M2, K2, N2, [&]{ cactus_quant_matmul(&mat, A.data(), M2, C.data()); }); + } + + auto bench_model = [](const char* label, size_t M, size_t K, size_t N, auto fn) { + fn(); + Timer t; + for (int i = 0; i < 5; i++) fn(); + double ms = t.elapsed_ms() / 5.0; + double gflops = (2.0 * M * K * N) / (ms * 1e6); + std::cout << " \u26A1 " << std::left << std::setw(28) << label + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gflops << " GFLOPS\n"; + }; + + const size_t Km = 2304, Nm = 9216; + const uint32_t gsm = 128; + + { + std::vector<__fp16> a(Km), b(Nm * Km), c(Nm); + fill_random_fp16(a, -0.5f, 0.5f); fill_random_fp16(b, -0.5f, 0.5f); + bench_model("f16 1x2304x9216", 1, Km, Nm, [&]{ cactus_matmul_f16(a.data(), b.data(), c.data(), 1, Km, Nm); }); + } + { + SyntheticCQ cq(1, Km, Nm, gsm); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(Km), y(Nm); + fill_random_fp16(x, -1.f, 1.f); + bench_model("cq1 1x2304x9216", 1, Km, Nm, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(2, Km, Nm, gsm); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(Km), y(Nm); + fill_random_fp16(x, -1.f, 1.f); + bench_model("cq2 1x2304x9216", 1, Km, Nm, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + { + SyntheticCQ cq(4, Km, Nm, gsm); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> x(Km), y(Nm); + fill_random_fp16(x, -1.f, 1.f); + bench_model("cq4 1x2304x9216", 1, Km, Nm, [&]{ cactus_quant_matmul(&mat, x.data(), 1, y.data()); }); + } + + return true; +} + + +void print_mse_report() { + const uint32_t K = 1024, N = 256, gs = 128; + + std::mt19937 gen(99); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + std::vector<__fp16> x_f16(K); + for (size_t i = 0; i < K; i++) x_f16[i] = static_cast<__fp16>(x_f32[i]); + + std::cout << "── MSE vs FP32 reference ──────────────────────────────────────────────────────────\n"; + + for (uint32_t bits : {1u, 2u, 3u, 4u}) { + SyntheticCQ cq(bits, K, N, gs, 55 + bits); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + + std::vector ref(N, 0.f); + cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); + + std::vector<__fp16> y(N, static_cast<__fp16>(0)); + cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + + double mse = compute_mse(ref.data(), y.data(), N); + double max_err = 0.0; + for (size_t i = 0; i < N; i++) { + double err = std::abs(static_cast(ref[i]) - static_cast(y[i])); + max_err = std::max(max_err, err); + } + + std::cout << " TQ" << bits << " │ MSE=" << std::scientific << std::setprecision(4) << mse + << " max_err=" << std::fixed << std::setprecision(5) << max_err << "\n"; + } +} + +// Legacy IL CQ4 vs the FP32 oracle through the real dispatch. +static bool test_cq4_interleaved(double& mse_out, + uint32_t K = 1024, uint32_t N = 192, uint32_t gs = 128) { + SyntheticCQ cq(4, K, N, gs, 777); + CactusQuantMatrix mat = cq.matrix_interleaved(); + + std::mt19937 gen(31); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + std::vector ref(N, 0.f); + cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); + + std::vector<__fp16> x_f16(K), y(N, (__fp16)0); + for (size_t i = 0; i < K; i++) x_f16[i] = (__fp16)x_f32[i]; + cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + mse_out = compute_mse(ref.data(), y.data(), N); + return mse_out <= 0.1; +} + +int main() { + TestRunner runner("Matrix Multiplication"); + runner.run_test("matmul_f16", test_matmul_f16()); + runner.run_test("matmul_cq1", test_cq_correctness(1)); + runner.run_test("matmul_cq2", test_cq_correctness(2)); + runner.run_test("matmul_cq3", test_cq_correctness(3)); + runner.run_test("matmul_cq4", test_cq_correctness(4)); + { + double m1 = 0; + runner.run_test("matmul_cq4_il", test_cq4_interleaved(m1)); + // N=4164 -> 66 chunks incl. a 1-block tail: exercises the multi-thread fused driver. + double m_mt = 0; + runner.run_test("matmul_cq4_il_mt", test_cq4_interleaved(m_mt, 1024, 4164, 128)); + } + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + print_mse_report(); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_quant.cpp b/cactus-kernels/tests/test_quant.cpp new file mode 100644 index 000000000..ee0bc3625 --- /dev/null +++ b/cactus-kernels/tests/test_quant.cpp @@ -0,0 +1,199 @@ +#include "test_utils.h" +#include +#include + +using namespace TestUtils; + +bool test_fp16_fp32_roundtrip() { + const size_t n = 256; + std::vector<__fp16> src(n); + std::vector fp32(n); + std::vector<__fp16> dst(n); + fill_random_fp16(src, -10.0f, 10.0f); + + cactus_fp16_to_fp32(src.data(), fp32.data(), n); + cactus_fp32_to_fp16(fp32.data(), dst.data(), n); + + return compare_arrays(src.data(), dst.data(), n, 1e-3f); +} + +bool test_int8_fp32_roundtrip() { + const size_t n = 256; + std::vector src(n), dst(n); + std::vector quantized(n); + for (size_t i = 0; i < n; i++) src[i] = static_cast(i % 127) - 63.0f; + + float scale = 1.0f; + cactus_fp32_to_int8(src.data(), quantized.data(), n, scale); + cactus_int8_to_fp32(quantized.data(), dst.data(), n, scale); + + for (size_t i = 0; i < n; i++) { + float expected = std::max(-127.0f, std::min(127.0f, std::round(src[i]))); + if (std::abs(dst[i] - expected) > 1.0f) { + std::cerr << " int8 roundtrip mismatch at " << i << ": " << dst[i] << " vs " << expected << "\n"; + return false; + } + } + return true; +} + +bool test_int8_fp16_conversion() { + const size_t n = 256; + std::vector src(n); + std::vector<__fp16> dst(n); + for (size_t i = 0; i < n; i++) src[i] = static_cast(i % 127 - 63); + + cactus_int8_to_fp16(src.data(), dst.data(), n, 0.1f); + + for (size_t i = 0; i < n; i++) { + float expected = static_cast(src[i]) * 0.1f; + float actual = static_cast(dst[i]); + if (std::abs(actual - expected) > 0.05f) { + std::cerr << " int8_to_fp16 mismatch at " << i << ": " << actual << " vs " << expected << "\n"; + return false; + } + } + return true; +} + +bool test_fp16_max_abs() { + const size_t n = 256; + std::vector<__fp16> data(n); + fill_random_fp16(data, -5.0f, 5.0f); + data[42] = static_cast<__fp16>(-99.0f); + + float result = cactus_fp16_max_abs(data.data(), n); + if (std::abs(result - 99.0f) > 0.5f) { + std::cerr << " fp16_max_abs: " << result << " (expected ~99.0)\n"; + return false; + } + return true; +} + +bool test_kv_quantize_int8() { + const size_t seq = 4, kv_heads = 2, head_dim = 64; + const size_t group_size = 32; + const size_t num_groups = head_dim / group_size; + std::vector<__fp16> src(seq * kv_heads * head_dim); + std::vector dst(seq * kv_heads * head_dim); + std::vector scales(seq * kv_heads * num_groups); + fill_random_fp16(src, -2.0f, 2.0f); + + cactus_quantize_kv_fp16_to_int8(src.data(), dst.data(), scales.data(), + seq, kv_heads, head_dim, group_size); + + for (size_t i = 0; i < scales.size(); i++) { + if (scales[i] <= 0.0f || !std::isfinite(scales[i])) { + std::cerr << " kv_quantize: invalid scale at " << i << ": " << scales[i] << "\n"; + return false; + } + } + + for (size_t s = 0; s < seq; s++) { + for (size_t h = 0; h < kv_heads; h++) { + for (size_t g = 0; g < num_groups; g++) { + float scale = scales[(s * kv_heads + h) * num_groups + g]; + for (size_t k = 0; k < group_size; k++) { + size_t idx = (s * kv_heads + h) * head_dim + g * group_size + k; + float original = static_cast(src[idx]); + float dequant = static_cast(dst[idx]) * scale; + if (std::abs(original - dequant) > 0.1f) { + std::cerr << " kv_quantize: dequant error at " << idx << ": " + << dequant << " vs " << original << "\n"; + return false; + } + } + } + } + } + return true; +} + +bool run_benchmarks() { + const size_t n = 1024 * 1024; + + { + std::vector<__fp16> src(n); + std::vector dst(n); + fill_random_fp16(src); + cactus_fp16_to_fp32(src.data(), dst.data(), n); + Timer t; + for (int i = 0; i < 100; i++) cactus_fp16_to_fp32(src.data(), dst.data(), n); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (static_cast(n) * (sizeof(__fp16) + sizeof(float))) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "fp16_to_fp32 1M" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + { + std::vector src(n); + std::vector<__fp16> dst(n); + for (size_t i = 0; i < n; i++) src[i] = static_cast(i % 1000) * 0.001f; + cactus_fp32_to_fp16(src.data(), dst.data(), n); + Timer t; + for (int i = 0; i < 100; i++) cactus_fp32_to_fp16(src.data(), dst.data(), n); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (static_cast(n) * (sizeof(float) + sizeof(__fp16))) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "fp32_to_fp16 1M" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + { + std::vector src(n); + std::vector<__fp16> dst(n); + for (size_t i = 0; i < n; i++) src[i] = static_cast(i % 127 - 63); + cactus_int8_to_fp16(src.data(), dst.data(), n, 0.01f); + Timer t; + for (int i = 0; i < 100; i++) cactus_int8_to_fp16(src.data(), dst.data(), n, 0.01f); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (static_cast(n) * (sizeof(int8_t) + sizeof(__fp16))) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "int8_to_fp16 1M" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + { + std::vector<__fp16> src(n); + std::vector dst(n); + fill_random_fp16(src, -1.0f, 1.0f); + cactus_fp16_to_int8(src.data(), dst.data(), n, 0.01f); + Timer t; + for (int i = 0; i < 100; i++) cactus_fp16_to_int8(src.data(), dst.data(), n, 0.01f); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (static_cast(n) * (sizeof(__fp16) + sizeof(int8_t))) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "fp16_to_int8 1M" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + { + const size_t seq = 4, kv_heads = 8, head_dim = 64, group_size = 32; + const size_t num_groups = head_dim / group_size; + std::vector<__fp16> src(seq * kv_heads * head_dim); + std::vector dst(seq * kv_heads * head_dim); + std::vector scales(seq * kv_heads * num_groups); + fill_random_fp16(src, -2.0f, 2.0f); + cactus_quantize_kv_fp16_to_int8(src.data(), dst.data(), scales.data(), seq, kv_heads, head_dim, group_size); + Timer t; + for (int i = 0; i < 100; i++) + cactus_quantize_kv_fp16_to_int8(src.data(), dst.data(), scales.data(), seq, kv_heads, head_dim, group_size); + double ms = t.elapsed_ms() / 100.0; + size_t nbytes = seq * kv_heads * head_dim; + double gb_s = (static_cast(nbytes) * (sizeof(__fp16) + sizeof(int8_t))) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << "kv_quantize 4x8x64" + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + } + return true; +} + +int main() { + TestRunner runner("Quantization & Conversion"); + runner.run_test("fp16_fp32_roundtrip", test_fp16_fp32_roundtrip()); + runner.run_test("int8_fp32_roundtrip", test_int8_fp32_roundtrip()); + runner.run_test("int8_fp16_conversion", test_int8_fp16_conversion()); + runner.run_test("fp16_max_abs", test_fp16_max_abs()); + runner.run_test("kv_quantize_int8", test_kv_quantize_int8()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_reduce.cpp b/cactus-kernels/tests/test_reduce.cpp new file mode 100644 index 000000000..a6c3753c2 --- /dev/null +++ b/cactus-kernels/tests/test_reduce.cpp @@ -0,0 +1,270 @@ +#include "test_utils.h" +#include +#include + +using namespace TestUtils; + +bool test_sum_all() { + const size_t n = 256; + std::vector<__fp16> data(n); + fill_random_fp16(data, -1.0f, 1.0f); + + double result = cactus_sum_all_f16(data.data(), n); + double expected = 0.0; + for (size_t i = 0; i < n; i++) expected += static_cast(static_cast(data[i])); + + if (std::abs(result - expected) > 0.5) { + std::cerr << " sum_all: " << result << " vs " << expected << "\n"; + return false; + } + return true; +} + +bool test_mean_all() { + const size_t n = 256; + std::vector<__fp16> data(n); + fill_random_fp16(data, -1.0f, 1.0f); + + double result = cactus_mean_all_f16(data.data(), n); + double expected = 0.0; + for (size_t i = 0; i < n; i++) expected += static_cast(static_cast(data[i])); + expected /= n; + + if (std::abs(result - expected) > 0.01) { + std::cerr << " mean_all: " << result << " vs " << expected << "\n"; + return false; + } + return true; +} + +bool test_variance_all() { + const size_t n = 256; + std::vector<__fp16> data(n); + fill_random_fp16(data, -1.0f, 1.0f); + + double result = cactus_variance_all_f16(data.data(), n); + + double mean = 0.0; + for (size_t i = 0; i < n; i++) mean += static_cast(data[i]); + mean /= n; + double var = 0.0; + for (size_t i = 0; i < n; i++) { + double d = static_cast(data[i]) - mean; + var += d * d; + } + var /= n; + + if (std::abs(result - var) > 0.02) { + std::cerr << " variance_all: " << result << " vs " << var << "\n"; + return false; + } + return true; +} + +bool test_min_max_all() { + const size_t n = 256; + std::vector<__fp16> data(n); + fill_random_fp16(data, -10.0f, 10.0f); + + __fp16 result_min = cactus_min_all_f16(data.data(), n); + __fp16 result_max = cactus_max_all_f16(data.data(), n); + + __fp16 expected_min = data[0], expected_max = data[0]; + for (size_t i = 1; i < n; i++) { + if (data[i] < expected_min) expected_min = data[i]; + if (data[i] > expected_max) expected_max = data[i]; + } + + if (result_min != expected_min || result_max != expected_max) { + std::cerr << " min/max: got [" << static_cast(result_min) << ", " + << static_cast(result_max) << "] expected [" + << static_cast(expected_min) << ", " + << static_cast(expected_max) << "]\n"; + return false; + } + return true; +} + +bool test_sum_axis() { + const size_t outer = 4, axis = 8, inner = 16; + std::vector<__fp16> input(outer * axis * inner), output(outer * inner); + fill_random_fp16(input, -1.0f, 1.0f); + + cactus_sum_axis_f16(input.data(), output.data(), outer, axis, inner); + + for (size_t o = 0; o < outer; o++) { + for (size_t i = 0; i < inner; i++) { + float ref = 0.0f; + for (size_t a = 0; a < axis; a++) { + ref += static_cast(input[(o * axis + a) * inner + i]); + } + float actual = static_cast(output[o * inner + i]); + if (std::abs(actual - ref) > 0.1f) { + std::cerr << " sum_axis mismatch [" << o << "," << i << "]: " + << actual << " vs " << ref << "\n"; + return false; + } + } + } + return true; +} + +bool test_neon_axis_inner1_correctness() { + const size_t outer_size = 2; + const size_t axis_size = 9; + const size_t inner_size = 1; + + std::vector<__fp16> input = { + 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, + 2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f, 14.0f, 16.0f, 18.0f + }; + std::vector<__fp16> output(outer_size * inner_size); + + cactus_sum_axis_f16(input.data(), output.data(), outer_size, axis_size, inner_size); + std::vector<__fp16> exp_sum = {45.0f, 90.0f}; + if (!TestUtils::compare_arrays(output.data(), exp_sum.data(), exp_sum.size(), 1e-2f)) return false; + + cactus_mean_axis_f16(input.data(), output.data(), outer_size, axis_size, inner_size); + std::vector<__fp16> exp_mean = {5.0f, 10.0f}; + if (!TestUtils::compare_arrays(output.data(), exp_mean.data(), exp_mean.size(), 1e-2f)) return false; + + cactus_min_axis_f16(input.data(), output.data(), outer_size, axis_size, inner_size); + std::vector<__fp16> exp_min = {1.0f, 2.0f}; + if (!TestUtils::compare_arrays(output.data(), exp_min.data(), exp_min.size(), 1e-2f)) return false; + + cactus_max_axis_f16(input.data(), output.data(), outer_size, axis_size, inner_size); + std::vector<__fp16> exp_max = {9.0f, 18.0f}; + if (!TestUtils::compare_arrays(output.data(), exp_max.data(), exp_max.size(), 1e-2f)) return false; + + return true; +} + +bool test_neon_variance_axis_inner1_correctness() { + const size_t outer_size = 2; + const size_t axis_size = 9; + const size_t inner_size = 1; + + std::vector<__fp16> input = { + 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, + 2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f, 14.0f, 16.0f, 18.0f + }; + std::vector<__fp16> output(outer_size * inner_size); + std::vector<__fp16> expected = {6.6667f, 26.6667f}; + + cactus_variance_axis_f16(input.data(), output.data(), outer_size, axis_size, inner_size); + return TestUtils::compare_arrays(output.data(), expected.data(), expected.size(), 0.05f); +} + +bool test_neon_variance_axis_non_inner1_correctness() { + const size_t outer_size = 2; + const size_t axis_size = 9; + const size_t inner_size = 3; + + std::vector<__fp16> input(outer_size * axis_size * inner_size); + auto idx = [&](size_t outer, size_t axis, size_t inner) { + return outer * axis_size * inner_size + axis * inner_size + inner; + }; + + for (size_t a = 0; a < axis_size; ++a) { + input[idx(0, a, 0)] = static_cast<__fp16>(1.0f + static_cast(a)); + input[idx(0, a, 1)] = static_cast<__fp16>(2.0f + static_cast(a)); + input[idx(0, a, 2)] = static_cast<__fp16>(-4.0f + static_cast(a)); + + input[idx(1, a, 0)] = static_cast<__fp16>(2.0f + 2.0f * static_cast(a)); + input[idx(1, a, 1)] = static_cast<__fp16>(3.0f + 2.0f * static_cast(a)); + input[idx(1, a, 2)] = static_cast<__fp16>(-8.0f + 2.0f * static_cast(a)); + } + + std::vector<__fp16> output(outer_size * inner_size); + std::vector<__fp16> expected = { + 6.6667f, 6.6667f, 6.6667f, + 26.6667f, 26.6667f, 26.6667f + }; + + cactus_variance_axis_f16(input.data(), output.data(), outer_size, axis_size, inner_size); + return TestUtils::compare_arrays(output.data(), expected.data(), expected.size(), 0.05f); +} + +bool run_benchmarks(TestRunner& runner) { + (void)runner; + auto benchmark = [&](const std::string& label, size_t n, auto fn) { + fn(); + Timer t; + for (int i = 0; i < 100; i++) fn(); + double ms = t.elapsed_ms() / 100.0; + double gb_s = (static_cast(n) * sizeof(__fp16)) / (ms * 1e6); + std::cout << " ⚡ " << std::left << std::setw(28) << label + << std::fixed << std::setprecision(3) << ms << "ms " + << std::setprecision(1) << gb_s << " GB/s\n"; + }; + + const size_t n = 1024 * 1024; + std::vector<__fp16> data(n); + fill_random_fp16(data, -1.0f, 1.0f); + + volatile double sink_d = 0.0; + volatile __fp16 sink_h = static_cast<__fp16>(0.0f); + + benchmark("sum_all 1M", n, [&]{ sink_d = cactus_sum_all_f16(data.data(), n); }); + benchmark("mean_all 1M", n, [&]{ sink_d = cactus_mean_all_f16(data.data(), n); }); + benchmark("variance_all 1M", n, [&]{ sink_d = cactus_variance_all_f16(data.data(), n); }); + benchmark("min_all 1M", n, [&]{ sink_h = cactus_min_all_f16(data.data(), n); }); + benchmark("max_all 1M", n, [&]{ sink_h = cactus_max_all_f16(data.data(), n); }); + + const size_t outer = 1024; + const size_t axis = 1024; + const size_t inner1 = 1; + std::vector<__fp16> axis_input(outer * axis * inner1); + std::vector<__fp16> axis_output(outer * inner1); + fill_random_fp16(axis_input, -1.0f, 1.0f); + + benchmark("sum_axis 1024x1024", axis_input.size(), [&]{ + cactus_sum_axis_f16(axis_input.data(), axis_output.data(), outer, axis, inner1); + }); + benchmark("mean_axis 1024x1024", axis_input.size(), [&]{ + cactus_mean_axis_f16(axis_input.data(), axis_output.data(), outer, axis, inner1); + }); + benchmark("variance_axis 1024x1024", axis_input.size(), [&]{ + cactus_variance_axis_f16(axis_input.data(), axis_output.data(), outer, axis, inner1); + }); + benchmark("min_axis 1024x1024", axis_input.size(), [&]{ + cactus_min_axis_f16(axis_input.data(), axis_output.data(), outer, axis, inner1); + }); + benchmark("max_axis 1024x1024", axis_input.size(), [&]{ + cactus_max_axis_f16(axis_input.data(), axis_output.data(), outer, axis, inner1); + }); + + const size_t channels = 4; + std::vector<__fp16> variance_input_non_inner1(outer * axis * channels); + std::vector<__fp16> variance_output_non_inner1(outer * channels); + fill_random_fp16(variance_input_non_inner1, -1.0f, 1.0f); + + benchmark("variance_axis 1024x1024x4", variance_input_non_inner1.size(), [&]{ + cactus_variance_axis_f16( + variance_input_non_inner1.data(), + variance_output_non_inner1.data(), + outer, + axis, + channels + ); + }); + + (void)sink_d; (void)sink_h; + return true; +} + +int main() { + TestRunner runner("Reduction Operations"); + runner.run_test("sum_all", test_sum_all()); + runner.run_test("mean_all", test_mean_all()); + runner.run_test("variance_all", test_variance_all()); + runner.run_test("min_max_all", test_min_max_all()); + runner.run_test("sum_axis", test_sum_axis()); + runner.run_test("Kernel Sum/Mean/Min/Max Axis Inner1 FP16 Correctness", test_neon_axis_inner1_correctness()); + runner.run_test("Kernel Variance Axis Inner1 FP16 Correctness", test_neon_variance_axis_inner1_correctness()); + runner.run_test("Kernel Variance Axis Non-Inner1 FP16 Correctness", test_neon_variance_axis_non_inner1_correctness()); + runner.print_benchmarks_header(); + runner.run_bench("benchmarks", run_benchmarks(runner)); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-kernels/tests/test_utils.h b/cactus-kernels/tests/test_utils.h new file mode 100644 index 000000000..5d2a14ea5 --- /dev/null +++ b/cactus-kernels/tests/test_utils.h @@ -0,0 +1,86 @@ +#pragma once + +#include "../cactus_kernels.h" +#include "../src/threading.h" +#include +#include +#include +#include +#include +#include +#include + +namespace TestUtils { + +class Timer { +public: + Timer() : start_(std::chrono::high_resolution_clock::now()) {} + double elapsed_ms() const { + auto now = std::chrono::high_resolution_clock::now(); + return std::chrono::duration(now - start_).count(); + } +private: + std::chrono::time_point start_; +}; + +class TestRunner { +public: + explicit TestRunner(const std::string& suite_name) : suite_(suite_name), passed_(0), failed_(0) { + std::cout << "\n╔══════════════════════════════════════════════════════════════════════════════════════╗\n" + << "║ Running " << std::left << std::setw(73) << suite_name << "║\n" + << "╚══════════════════════════════════════════════════════════════════════════════════════╝\n"; + } + + void run_test(const std::string& name, bool result) { + if (result) { + std::cout << "✓ PASS │ " << std::left << std::setw(25) << name << std::endl; + passed_++; + } else { + std::cout << "✗ FAIL │ " << std::left << std::setw(25) << name << std::endl; + failed_++; + } + } + + void run_bench(const std::string& /*name*/, bool result) { + if (!result) failed_++; + } + + void print_benchmarks_header() const { + std::cout << "── benchmarks ──────────────────────────────────────────────────────────────────────────\n"; + } + + void print_summary() const { + std::cout << "────────────────────────────────────────────────────────────────────────────────────────\n"; + if (failed_ == 0) { + std::cout << "✓ All " << passed_ << " tests passed!\n"; + } else { + std::cout << "✗ " << failed_ << " of " << (passed_ + failed_) << " tests failed!\n"; + } + } + + bool all_passed() const { return failed_ == 0; } + +private: + std::string suite_; + int passed_; + int failed_; +}; + +inline bool compare_arrays(const __fp16* a, const __fp16* b, size_t n, float tol = 1e-2f) { + for (size_t i = 0; i < n; i++) { + if (std::abs(static_cast(a[i]) - static_cast(b[i])) > tol) { + std::cerr << " mismatch at [" << i << "]: " << static_cast(a[i]) + << " vs " << static_cast(b[i]) << std::endl; + return false; + } + } + return true; +} + +inline void fill_random_fp16(std::vector<__fp16>& v, float lo = -1.0f, float hi = 1.0f) { + static std::mt19937 gen(42); + std::uniform_real_distribution dis(lo, hi); + for (auto& x : v) x = static_cast<__fp16>(dis(gen)); +} + +} // namespace TestUtils diff --git a/cactus/CMakeLists.txt b/cactus/CMakeLists.txt deleted file mode 100644 index 7b056ceb2..000000000 --- a/cactus/CMakeLists.txt +++ /dev/null @@ -1,124 +0,0 @@ -cmake_minimum_required(VERSION 3.10) -project(Cactus LANGUAGES CXX) - -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED True) - -if(NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE Release) - message(STATUS "CMAKE_BUILD_TYPE was not set, defaulting to Release.") -endif() - -set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/) -set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/) -set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) - -file(READ "${CMAKE_CURRENT_SOURCE_DIR}/../CACTUS_VERSION" CACTUS_VERSION) -string(STRIP "${CACTUS_VERSION}" CACTUS_VERSION) - -if(APPLE) - set(CMAKE_OSX_ARCHITECTURES "arm64") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -arch arm64 -march=armv8.2-a+fp16+simd+dotprod+i8mm -pthread -Wall -Wextra -pedantic -O3 -Wno-missing-field-initializers") - add_compile_definitions( - __ARM_NEON=1 - __ARM_FEATURE_FP16_VECTOR_ARITHMETIC=1 - __ARM_FEATURE_DOTPROD=1 - __ARM_FEATURE_MATMUL_INT8=1 - ) - - enable_language(OBJCXX) - find_package(CURL REQUIRED) - find_library(COREML_FRAMEWORK CoreML REQUIRED) - find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED) -else() - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=armv8.2-a+fp16+simd+dotprod+i8mm -pthread -Wall -Wextra -pedantic -O3 -Wno-missing-field-initializers") - add_compile_definitions( - __ARM_NEON=1 - __ARM_FEATURE_FP16_VECTOR_ARITHMETIC=1 - __ARM_FEATURE_DOTPROD=1 - __ARM_FEATURE_MATMUL_INT8=1 - ) -endif() - -file(GLOB ENGINE_SOURCES "engine/*.cpp") -file(GLOB GRAPH_SOURCES "graph/*.cpp") -file(GLOB KERNEL_SOURCES "kernel/*.cpp") -file(GLOB FFI_SOURCES "ffi/*.cpp") -file(GLOB MODEL_SOURCES "models/*.cpp") -file(GLOB INDEX_SOURCES "index/*.cpp") - -file(GLOB NPU_SOURCES "npu/*.cpp") - -set(CACTUS_PRO_LIB "${CMAKE_CURRENT_SOURCE_DIR}/../libs/libcactus_pro.a") -if(EXISTS ${CACTUS_PRO_LIB}) - message(STATUS "Found libcactus_pro.a - NPU acceleration enabled") - set(CACTUS_HAS_PRO TRUE) -else() - message(STATUS "libcactus_pro.a not found - NPU acceleration disabled (using stubs)") - set(CACTUS_HAS_PRO FALSE) -endif() - -set(COMMON_SOURCES - ${ENGINE_SOURCES} - ${GRAPH_SOURCES} - ${KERNEL_SOURCES} - ${FFI_SOURCES} - ${MODEL_SOURCES} - ${INDEX_SOURCES} - ${NPU_SOURCES} -) - -set(COMMON_INCLUDES - ${CMAKE_CURRENT_SOURCE_DIR} - engine - graph - kernel - ffi - models - index - npu -) - -set(COMMON_DEFINITIONS - PLATFORM_CPU_ONLY=1 - CACTUS_VERSION="${CACTUS_VERSION}" -) - -function(configure_cactus_target target_name) - target_compile_definitions(${target_name} PUBLIC ${COMMON_DEFINITIONS}) - target_include_directories(${target_name} PUBLIC ${COMMON_INCLUDES}) - - if(CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang" OR CMAKE_CXX_COMPILER_ID STREQUAL "Clang") - target_compile_options(${target_name} PRIVATE -Wno-c99-extensions) - elseif(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") - target_compile_options(${target_name} PRIVATE -Wno-pedantic) - endif() - - if(CACTUS_HAS_PRO) - if(APPLE) - target_link_libraries(${target_name} PUBLIC "-force_load ${CACTUS_PRO_LIB}") - else() - target_link_libraries(${target_name} PUBLIC - "-Wl,--whole-archive" ${CACTUS_PRO_LIB} "-Wl,--no-whole-archive") - endif() - endif() - - if(APPLE) - target_link_libraries(${target_name} PUBLIC - ${COREML_FRAMEWORK} - ${FOUNDATION_FRAMEWORK} - ${CURL_LIBRARIES} - ) - target_include_directories(${target_name} PUBLIC ${CURL_INCLUDE_DIRS}) - endif() -endfunction() - -add_library(cactus STATIC ${COMMON_SOURCES}) -configure_cactus_target(cactus) - -add_library(cactus_ffi SHARED ${COMMON_SOURCES}) -configure_cactus_target(cactus_ffi) -set_target_properties(cactus_ffi PROPERTIES - POSITION_INDEPENDENT_CODE ON - OUTPUT_NAME "cactus" -) diff --git a/cactus/cactus.h b/cactus/cactus.h deleted file mode 100644 index 727801389..000000000 --- a/cactus/cactus.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef CACTUS_H -#define CACTUS_H - -#include "graph/graph.h" -#include "kernel/kernel.h" -#include "kernel/kernel_utils.h" -#include "engine/engine.h" -#include "models/model.h" -#include "ffi/cactus_ffi.h" -#include "ffi/cactus_telemetry.h" -#include "npu/npu.h" - -#endif // CACTUS_H \ No newline at end of file diff --git a/cactus/engine/engine.h b/cactus/engine/engine.h deleted file mode 100644 index a2cceb2d5..000000000 --- a/cactus/engine/engine.h +++ /dev/null @@ -1,798 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include - -#include "../graph/graph.h" - -#ifdef __clang__ -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wc99-extensions" -#pragma clang diagnostic ignored "-Wunused-parameter" -#elif defined(__GNUC__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wpedantic" -#pragma GCC diagnostic ignored "-Wunused-parameter" -#endif - -extern "C" { - #include "../../libs/stb/stb_image.h" - #include "../../libs/stb/stb_image_resize2.h" -} - -#ifdef __clang__ -#pragma clang diagnostic pop -#elif defined(__GNUC__) -#pragma GCC diagnostic pop -#endif - -class CactusGraph; - -namespace cactus { -namespace npu { - class NPUPrefill; -} -namespace engine { - -class Siglip2Preprocessor; - -struct Config { - uint32_t vocab_size = 151936; - uint32_t bos_token_id = 151643; - uint32_t eos_token_id = 151645; - uint32_t num_layers = 28; - uint32_t hidden_dim = 1024; - uint32_t ffn_intermediate_dim = 3072; - uint32_t attention_heads = 16; - uint32_t attention_kv_heads = 8; - uint32_t attention_head_dim = 128; - float layer_norm_eps = 1e-6f; - float rope_theta = 1000000.0f; - uint32_t num_experts = 0; - uint32_t num_shared_experts = 0; - uint32_t num_top_experts = 0; - uint32_t moe_every_n_layers = 0; - bool tie_word_embeddings = true; - - uint32_t vision_hidden_dim = 0; - uint32_t vision_num_layers = 0; - uint32_t vision_attention_heads = 0; - uint32_t vision_image_size = 0; - uint32_t vision_patch_size = 0; - uint32_t vision_num_channels = 3; - uint32_t vision_embed_dim = 0; - uint32_t visual_tokens_per_img = 0; - bool use_pixel_shuffle = false; - uint32_t pixel_shuffle_factor = 1; - bool use_image_tokens = false; - bool use_layout_tags = false; - uint32_t image_seq_len = 64; - - uint32_t global_image_size = 2048; - uint32_t max_tile_size = 512; - float rescale_factor = 0.00392156862745098f; - float image_mean = 0.5f; - float image_std = 0.5f; - - uint32_t downsample_factor = 2; - uint32_t min_tiles = 2; - uint32_t max_tiles = 10; - bool use_thumbnail = true; - uint32_t min_image_tokens = 64; - uint32_t max_image_tokens = 256; - uint32_t max_num_patches = 1024; - uint32_t tile_size = 512; - float max_pixels_tolerance = 2.0f; - bool do_image_splitting = true; - - enum class ModelType {QWEN = 0, GEMMA = 1, SMOL = 2, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7}; - ModelType model_type = ModelType::QWEN; - - enum class ModelVariant {DEFAULT = 0, VLM = 1, EXTRACT = 2, RAG = 3}; - ModelVariant model_variant = ModelVariant::DEFAULT; - - enum class Activation {GELU = 0, SILU = 1}; - Activation activation = Activation::SILU; - - enum class Backend {CPU = 0, NPU = 1}; - Backend default_backend = Backend::CPU; - - enum class Precision {INT8 = 0, FP16 = 1, FP32 = 2}; - Precision precision = Precision::FP32; - - float default_temperature = 0.6f; - float default_top_p = 0.95f; - size_t default_top_k = 20; - - std::vector layer_types; - size_t conv_L_cache = 0; - - bool from_json(const std::string& json_path); - std::string to_json() const; -}; - - - -struct MergeRule { - std::string first; - std::string second; - std::string merged; - uint32_t priority; - - MergeRule(const std::string& f, const std::string& s, const std::string& m, uint32_t p) - : first(f), second(s), merged(m), priority(p) {} -}; - - -struct ChatMessage { - std::string role; - std::string content; - std::string name; - std::vector images; -}; - - - -class Tokenizer { -public: - virtual ~Tokenizer() = default; - - virtual std::vector encode(const std::string& text) const = 0; - virtual std::string decode(const std::vector& tokens) const = 0; - - virtual std::vector apply_chat_template(const std::vector& messages, bool add_generation_prompt = true) const; - virtual std::string format_chat_prompt(const std::vector& messages, bool add_generation_prompt = true, const std::string& tools_json = "") const; - - virtual uint32_t get_vocab_size() const = 0; - virtual uint32_t get_unk_token() const = 0; - virtual uint32_t get_bos_token() const = 0; - virtual uint32_t get_eos_token() const = 0; - virtual bool has_chat_template() const { return has_chat_template_; } - - virtual bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) = 0; - - uint32_t get_image_token_id() const { return image_token_id_; } - uint32_t get_fake_token_id() const { return fake_token_id_; } - uint32_t get_global_img_token_id() const { return global_img_token_id_; } - -protected: - enum class ModelType { UNKNOWN, QWEN, GEMMA, LFM2, SMOL, BERT, WHISPER}; - ModelType model_type_ = ModelType::UNKNOWN; - enum class ModelVariant { DEFAULT, VLM, EXTRACT, RAG}; - ModelVariant model_variant_ = ModelVariant::DEFAULT; - bool has_chat_template_ = false; - std::string chat_template_; - - uint32_t image_token_id_ = 396; - uint32_t fake_token_id_ = 49189; - uint32_t global_img_token_id_ = 49152; - - void detect_model_type(const std::string& config_path); - std::string format_qwen_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; - std::string format_gemma_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; - std::string format_lfm2_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; - std::string format_lfm2_vl_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; - std::string format_smol_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; -}; - -class BPETokenizer : public Tokenizer { -public: - BPETokenizer(); - ~BPETokenizer(); - - bool load_vocabulary_mmap(const std::string& vocab_file, const std::string& merges_file); - bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) override; - - std::vector encode(const std::string& text) const override; - std::string decode(const std::vector& tokens) const override; - - uint32_t get_vocab_size() const override { return vocab_size_; } - uint32_t get_unk_token() const override { return unk_token_id_; } - uint32_t get_bos_token() const override { return bos_token_id_; } - uint32_t get_eos_token() const override { return eos_token_id_; } - -private: - std::unordered_map token_to_id_; - std::vector id_to_token_; - std::vector merge_rules_; - std::unordered_map merge_map_; - - uint32_t vocab_size_; - uint32_t unk_token_id_; - uint32_t bos_token_id_; - uint32_t eos_token_id_; - - void* vocab_mmap_ptr_; - size_t vocab_mmap_size_; - - void* merges_mmap_ptr_; - size_t merges_mmap_size_; - - std::vector apply_bpe(const std::vector& tokens) const; - std::pair find_best_merge_fast(const std::vector& tokens) const; - - std::string bytes_to_unicode(const std::string& text) const; - std::string unicode_to_bytes(const std::string& text) const; - std::vector byte_level_split(const std::string& text) const; - - void cleanup_mmap(); - -private: - mutable std::unordered_map byte_to_unicode_; - mutable std::unordered_map unicode_to_byte_; - void init_byte_mappings() const; - - std::unordered_map special_tokens_; - std::vector split_with_special_tokens(const std::string& text) const; - void load_special_tokens(const std::string& config_file); - - void load_chat_template(const std::string& template_file); - - std::unordered_map tool_tokens_; - bool has_tool_support_; - void load_tokenizer_config(const std::string& config_file); -}; - -class SPTokenizer : public Tokenizer { -public: - SPTokenizer(); - ~SPTokenizer(); - - bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) override; - - std::vector encode(const std::string& text) const override; - std::string decode(const std::vector& tokens) const override; - - uint32_t get_vocab_size() const override { return vocab_size_; } - uint32_t get_unk_token() const override { return unk_token_id_; } - uint32_t get_bos_token() const override { return bos_token_id_; } - uint32_t get_eos_token() const override { return eos_token_id_; } - -private: - struct TrieNode { - std::unordered_map> children; - int32_t token_id = -1; - float score = 0.0f; - }; - - std::unique_ptr trie_root_; - std::unordered_map token_to_id_; - std::vector id_to_token_; - std::vector token_scores_; - - uint32_t vocab_size_; - uint32_t unk_token_id_; - uint32_t bos_token_id_; - uint32_t eos_token_id_; - uint32_t pad_token_id_; - - void* vocab_mmap_ptr_; - size_t vocab_mmap_size_; - - void build_trie(); - std::vector> tokenize_with_trie(const std::string& text) const; - std::string preprocess_text(const std::string& text) const; - std::string postprocess_text(const std::string& text) const; - std::vector split_by_unicode_spaces(const std::string& text) const; - - void cleanup_mmap(); - - std::unordered_map special_tokens_; - std::vector split_with_special_tokens(const std::string& text) const; - void load_special_tokens(const std::string& config_file); - - void load_chat_template(const std::string& template_file); -}; - -class ConvCache { -public: - struct CircularView { - const void* ptr1; - size_t len1; - const void* ptr2; - size_t len2; - size_t total_len; - }; - - void init(size_t layers, size_t hidden_dim, size_t window_len, Precision model_precision); - CircularView get_window(size_t layer) const; - void update(CactusGraph* gb, size_t layer, const size_t latest_token); - void reset(); - - bool is_empty() const { return num_layers == 0; } - - size_t num_layers = 0; - size_t hidden_size = 0; - size_t window_size = 0; - Precision precision = Precision::FP32; - size_t element_size = 4; - -private: - struct LayerState { - std::vector data; - size_t head = 0; - size_t count = 0; - }; - - std::vector layer_states; -}; - -struct KVCache { - static constexpr size_t DEFAULT_WINDOW_SIZE = 1024; - static constexpr size_t DEFAULT_SINK_SIZE = 4; - - struct LayerCache { - std::vector keys; - std::vector values; - std::vector key_scales; - std::vector value_scales; - }; - - std::vector layer_caches; - - size_t window_size = DEFAULT_WINDOW_SIZE; - size_t sink_size = DEFAULT_SINK_SIZE; - size_t current_seq_len = 0; - size_t total_seq_len = 0; - size_t max_seq_len = 2048; - size_t num_kv_heads = 0; - size_t head_dim = 0; - size_t num_layers = 0; - Precision precision; - size_t element_size = 4; - - void set_window_size(size_t window, size_t sink = DEFAULT_SINK_SIZE); - size_t get_effective_seq_len() const { return current_seq_len; } - size_t get_total_seq_len() const { return total_seq_len; } - - void init(size_t num_layers, size_t max_seq, size_t num_kv_heads, size_t head_dim, Precision model_precision); - void reset(); - void update_from_graph(CactusGraph* gb, const std::vector& k_nodes, - const std::vector& v_nodes, size_t seq_len, - size_t num_layers, size_t kv_heads, size_t head_dim); - - void update_from_npu(size_t layer_idx, const __fp16* k_data, const __fp16* v_data, - size_t num_tokens, size_t kv_heads, size_t head_dim); - - bool is_empty() const { return current_seq_len == 0; } - bool is_int8() const { return precision == Precision::INT8; } - void* get_key_ptr(size_t layer); - void* get_value_ptr(size_t layer); - - struct CircularView { - const void* ptr1; - const void* ptr2; - size_t len1; - size_t len2; - size_t total_len; - }; - - CircularView get_key_view(size_t layer); - CircularView get_value_view(size_t layer); - - const int8_t* get_keys_int8(size_t layer) const; - const int8_t* get_values_int8(size_t layer) const; - const float* get_key_scales(size_t layer) const; - const float* get_value_scales(size_t layer) const; -}; - -class ToolCallConstrainer { -public: - enum class State { - DONE, - - QWEN_START, - QWEN_EXPECT_OPEN_BRACE, - QWEN_EXPECT_NAME_KEY, - QWEN_EXPECT_NAME_COLON, - QWEN_EXPECT_NAME_VALUE, - QWEN_EXPECT_COMMA, - QWEN_EXPECT_ARGS_KEY, - QWEN_EXPECT_ARGS_COLON, - QWEN_IN_ARGUMENTS, - QWEN_EXPECT_CLOSE_BRACE, - QWEN_EXPECT_END, - - LFM_START, - LFM_EXPECT_BRACKET, - LFM_IN_FUNC_NAME, - LFM_EXPECT_PAREN, - LFM_IN_ARGUMENTS, - LFM_EXPECT_BRACKET_CLOSE, - LFM_EXPECT_END, - - GEMMA_START, - GEMMA_EXPECT_CALL, - GEMMA_IN_FUNC_NAME, - GEMMA_EXPECT_BRACE, - GEMMA_IN_ARGUMENTS, - GEMMA_EXPECT_END - }; - - void init(Config::ModelType model_type, - const std::vector& function_names, - Tokenizer* tokenizer); - - const std::unordered_map& get_bias() const { return current_bias_; } - - void update(uint32_t token_id, const std::string& decoded_text); - - void reset(); - - bool is_active() const { return active_; } - -private: - bool active_ = false; - State state_ = State::QWEN_START; - Config::ModelType model_type_ = Config::ModelType::QWEN; - Tokenizer* tokenizer_ = nullptr; - - std::vector function_names_; - std::string generated_text_; - int brace_depth_ = 0; - - std::unordered_set qwen_tool_call_start_tokens_; - std::unordered_set qwen_tool_call_end_tokens_; - std::unordered_set open_brace_tokens_; - std::unordered_set close_brace_tokens_; - std::unordered_set colon_tokens_; - std::unordered_set comma_tokens_; - std::unordered_set name_key_tokens_; - std::unordered_set args_key_tokens_; - std::unordered_set quote_tokens_; - std::unordered_set backtick_tokens_; - std::unordered_set all_func_name_tokens_; - std::unordered_map> func_name_sequences_; - - std::unordered_set tool_start_tokens_; - std::unordered_set tool_end_tokens_; - std::unordered_set bracket_open_tokens_; - std::unordered_set bracket_close_tokens_; - std::unordered_set paren_open_tokens_; - std::unordered_set paren_close_tokens_; - std::unordered_set equals_tokens_; - - std::unordered_set gemma_call_start_tokens_; - std::unordered_set gemma_call_end_tokens_; - std::unordered_set gemma_response_start_tokens_; - std::unordered_set gemma_call_prefix_tokens_; - std::unordered_set escape_tokens_; - - std::unordered_map current_bias_; - - void compute_bias(); - void tokenize_grammar_elements(); - void add_tokens_for_string(const std::string& str, std::unordered_set& token_set); -}; - -class Model { -public: - struct DebugNode { - uint32_t layer_idx; - std::string name; - size_t node_id; - }; - - Model(); - explicit Model(const Config& config); - virtual ~Model(); - - const Config& get_config() const { return config_; } - Tokenizer* get_tokenizer() const { return tokenizer_.get(); } - const std::vector& get_debug_nodes() const; - - virtual bool init(const std::string& model_folder, size_t context_size, const std::string& system_prompt = "", bool do_warmup = true); - - virtual bool init(CactusGraph* external_graph, const std::string& model_folder, size_t context_size, - const std::string& system_prompt = "", bool do_warmup = true); - - virtual uint32_t decode(const std::vector& tokens, float temperature = -1.0f, float top_p = -1.0f, - size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr); - - virtual void prefill(const std::vector& tokens, size_t chunk_size = 256, const std::string& profile_file = ""); - - virtual uint32_t decode_with_images(const std::vector& tokens, const std::vector& image_paths, - float temperature = -1.0f, float top_p = -1.0f, - size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr); - - virtual uint32_t decode_with_audio(const std::vector& tokens, const std::vector& mel_bins, float temperature = 0.0f, float top_p = 0.0f, - size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr); - - std::vector get_embeddings(const std::vector& tokens, bool pooled = true, bool normalize = false, const std::string& profile_file = ""); - - virtual std::vector get_image_embeddings(const std::string& image_path); - - virtual std::vector get_audio_embeddings(const std::vector& mel_bins); - - virtual void reset_cache() { kv_cache_.reset(); } - - double score_tokens_window_logprob(const std::vector& tokens, size_t start, size_t end, size_t context, size_t* tokens_scored); - - - - void set_cache_window(size_t window_size, size_t sink_size = 4) { kv_cache_.set_window_size(window_size, sink_size); } - - bool load_npu_prefill(const std::string& model_path); - bool has_npu_prefill() const; - size_t get_prefill_chunk_size() const; - - void set_tool_constraints(const std::vector& function_names); - void clear_tool_constraints(); - void update_tool_constraints(uint32_t token_id); - - void* graph_handle_; - -protected: - virtual size_t forward(const std::vector& tokens, bool use_cache = false) = 0; - - virtual size_t forward(const std::vector& mel_bins, const std::vector& tokens, bool use_cache = false); - - virtual void load_weights_to_graph(CactusGraph* gb) = 0; - - virtual size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) = 0; - - virtual size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const = 0; - virtual size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) = 0; - void update_kv_cache(CactusGraph* gb, size_t seq_len); - virtual void post_init() {} - virtual void post_execute_updates(CactusGraph*, size_t) {} - Config config_; - std::unique_ptr tokenizer_; - - bool initialized_; - float attention_scale_; - -protected: - KVCache kv_cache_; - std::vector cache_k_output_nodes_; - std::vector cache_v_output_nodes_; - - std::string embedding_file_path_; - size_t embedding_node_id_; - std::string model_folder_path_; - size_t output_weight_node_id_; - - mutable std::vector debug_nodes_; - - void capture_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id) const; - void clear_debug_nodes(); - - bool init_internal(CactusGraph* gb, const std::string& model_folder, size_t context_size, - const std::string& system_prompt, bool do_warmup); - bool owns_graph_; - - std::unique_ptr npu_prefill_; - void prefill_npu(const std::vector& tokens); - virtual std::vector<__fp16> get_token_embeddings(const std::vector& tokens); - - ToolCallConstrainer tool_constrainer_; -}; - -std::unique_ptr create_model(const std::string& model_folder); - -class Siglip2Preprocessor { -public: - struct Config { - int patch_size = 16; - int downsample_factor = 2; - int min_tiles = 2; - int max_tiles = 10; - bool use_thumbnail = true; - int min_image_tokens = 64; - int max_image_tokens = 256; - int max_num_patches = 1024; - int tile_size = 512; - float max_pixels_tolerance = 2.0f; - bool do_resize = true; - bool do_rescale = true; - bool do_normalize = true; - bool do_convert_rgb = true; - bool do_image_splitting = true; - float rescale_factor = 1.0f / 255.0f; - float image_mean[3] = {0.5f, 0.5f, 0.5f}; - float image_std[3] = {0.5f, 0.5f, 0.5f}; - }; - - struct PreprocessedImage { - std::vector pixel_values; - std::vector pixel_attention_mask; - std::vector> spatial_shapes; - std::vector pixel_values_shape; - std::vector pixel_attention_mask_shape; - std::vector spatial_shapes_shape; - int num_patches_height; - int num_patches_width; - int actual_num_patches; - int num_tiles; - int patch_dim; - int max_patches_per_tile; - - int image_rows; - int image_cols; - int image_height; - int image_width; - int tokens_per_tile; - int thumbnail_tokens; - - ~PreprocessedImage(); - }; - - struct SpatialShapeResult { - std::vector> shapes; - int grid_rows; - int grid_cols; - }; - - explicit Siglip2Preprocessor(const Config& config); - Siglip2Preprocessor(); - ~Siglip2Preprocessor(); - - PreprocessedImage preprocess_from_file(const std::string& image_path); - PreprocessedImage preprocess_from_memory(const unsigned char* img_data, int width, int height, int channels); - SpatialShapeResult compute_spatial_shapes(int height, int width); - -private: - Config config_; - - std::vector convert_to_rgb(const unsigned char* img_data, int width, int height, int channels); - std::pair smart_resize(int height, int width); - bool is_image_too_large(int height, int width); - std::pair get_grid_layout(int height, int width); - std::pair find_closest_aspect_ratio(float aspect_ratio, int width, int height); - std::vector resize_image(const unsigned char* img_data, int src_width, int src_height, - int dst_width, int dst_height, int channels); - std::vector normalize_image(const float* img_data, int width, int height, int channels); - std::vector> convert_image_to_patches( - const std::vector& image, int width, int height, int channels, int patch_size); - PreprocessedImage pad_patches(const std::vector>& tile_patches, - const std::vector>& spatial_shapes, - int patch_dim, - int max_patches_per_tile); - int round_by_factor(int number, int factor); -}; - -class AudioProcessor { -public: - struct SpectrogramConfig { - size_t n_fft = 400; - size_t hop_length = 160; - size_t frame_length = 400; - float power = 2.0f; - bool center = true; - const char* pad_mode = "reflect"; - bool onesided = true; - float dither = 0.0f; - float mel_floor = 1e-10f; - const char* log_mel = nullptr; - float reference = 1.0f; - float min_value = 1e-10f; - bool remove_dc_offset = false; - }; - - AudioProcessor(); - ~AudioProcessor(); - - void init_mel_filters(size_t num_frequency_bins, size_t num_mel_filters, - float min_freq, float max_freq, size_t sampling_rate); - - std::vector compute_spectrogram( - const std::vector& waveform, - const SpectrogramConfig& config); - - const std::vector& get_mel_filters() const { return mel_filters_; } - - size_t get_num_mel_filters() const { return num_mel_filters_; } - size_t get_num_frequency_bins() const { return num_frequency_bins_; } - -private: - std::vector mel_filters_; - size_t num_frequency_bins_; - size_t num_mel_filters_; -}; - -namespace index { - constexpr uint32_t MAGIC = 0x43414354; - constexpr uint32_t VERSION = 1; - - struct Document { - int id; - std::vector embedding; - std::string content; - std::string metadata; - }; - - struct QueryResult { - int doc_id; - float score; - }; - - struct QueryOptions { - size_t top_k = 10; - float score_threshold = -1.0f; - }; - - class Index { - public: - Index(const std::string& index_path, const std::string& data_path, size_t embedding_dim); - ~Index(); - - Index(const Index&) = delete; - Index& operator=(const Index&) = delete; - Index(Index&&) = delete; - Index& operator=(Index&&) = delete; - - void add_documents(const std::vector& documents); - void delete_documents(const std::vector& doc_ids); - std::vector get_documents(const std::vector& doc_ids); - std::vector> query(const std::vector>& embeddings, const QueryOptions& options); - void compact(); - - private: - struct IndexHeader { - uint32_t magic; - uint32_t version; - uint32_t embedding_dim; - uint32_t num_documents; - }; - - struct IndexEntry { - int32_t doc_id; - uint64_t data_offset; - uint8_t flags; // bit 0: tombstone - - const __fp16* embedding() const { - return reinterpret_cast(this + 1); - } - - static size_t size(size_t embedding_dim) { - return sizeof(IndexEntry) + embedding_dim * sizeof(__fp16); - } - }; - - struct DataHeader { - uint32_t magic; - uint32_t version; - }; - - struct DataEntry { - uint16_t content_len; - uint16_t metadata_len; - - const char* content() const { - return reinterpret_cast(this + 1); - } - - const char* metadata() const { - return content() + content_len; - } - }; - - void parse_index_header(); - void parse_data_header(); - void build_doc_id_map(); - void validate_documents(const std::vector& documents); - void validate_doc_ids(const std::vector& doc_ids); - ssize_t write_full(int fd, const void* buf, size_t count); - - std::unordered_map doc_id_map_; - - std::string index_path_, data_path_; - size_t embedding_dim_; - size_t index_entry_size_; - uint32_t num_documents_; - - int index_fd_, data_fd_; - void *mapped_index_, *mapped_data_; - size_t index_file_size_, data_file_size_; - }; -} // namespace index - -} -} \ No newline at end of file diff --git a/cactus/engine/engine_cache.cpp b/cactus/engine/engine_cache.cpp deleted file mode 100644 index 6e39b20ab..000000000 --- a/cactus/engine/engine_cache.cpp +++ /dev/null @@ -1,608 +0,0 @@ -#include "engine.h" -#include "../graph/graph.h" -#include "../kernel/kernel.h" -#include -#include -#include - -namespace cactus { -namespace engine { - -void KVCache::init(size_t layers, size_t max_seq, size_t kv_heads, size_t dim, Precision model_precision) { - num_layers = layers; - max_seq_len = max_seq; - num_kv_heads = kv_heads; - head_dim = dim; - precision = model_precision; - element_size = PrecisionTraits::size_of(precision); - - layer_caches.resize(num_layers); - - current_seq_len = 0; - total_seq_len = 0; -} - -void KVCache::set_window_size(size_t window, size_t sink) { - window_size = window; - sink_size = sink; - - if (num_kv_heads > 0 && head_dim > 0 && window_size > 0) { - size_t cache_bytes = window_size * num_kv_heads * head_dim * element_size; - size_t num_groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; - size_t num_scales = window_size * num_kv_heads * num_groups; - - for (auto& cache : layer_caches) { - cache.keys.resize(cache_bytes); - cache.values.resize(cache_bytes); - std::memset(cache.keys.data(), 0, cache_bytes); - std::memset(cache.values.data(), 0, cache_bytes); - - cache.key_scales.resize(num_scales); - cache.value_scales.resize(num_scales); - std::fill(cache.key_scales.begin(), cache.key_scales.end(), 1.0f); - std::fill(cache.value_scales.begin(), cache.value_scales.end(), 1.0f); - } - } -} - -void KVCache::reset() { - current_seq_len = 0; - total_seq_len = 0; -} - -void* KVCache::get_key_ptr(size_t layer) { - if (current_seq_len == 0 || layer >= num_layers) return nullptr; - return layer_caches[layer].keys.data(); -} - -void* KVCache::get_value_ptr(size_t layer) { - if (current_seq_len == 0 || layer >= num_layers) return nullptr; - return layer_caches[layer].values.data(); -} - -KVCache::CircularView KVCache::get_key_view(size_t layer) { - CircularView view; - if (layer >= num_layers || current_seq_len == 0) { - view.ptr1 = nullptr; - view.ptr2 = nullptr; - view.len1 = 0; - view.len2 = 0; - view.total_len = 0; - return view; - } - - view.ptr1 = layer_caches[layer].keys.data(); - view.ptr2 = nullptr; - view.len1 = current_seq_len; - view.len2 = 0; - view.total_len = current_seq_len; - return view; -} - -KVCache::CircularView KVCache::get_value_view(size_t layer) { - CircularView view; - if (layer >= num_layers || current_seq_len == 0) { - view.ptr1 = nullptr; - view.ptr2 = nullptr; - view.len1 = 0; - view.len2 = 0; - view.total_len = 0; - return view; - } - - view.ptr1 = layer_caches[layer].values.data(); - view.ptr2 = nullptr; - view.len1 = current_seq_len; - view.len2 = 0; - view.total_len = current_seq_len; - return view; -} - -void KVCache::update_from_graph(CactusGraph* gb, const std::vector& k_nodes, - const std::vector& v_nodes, size_t seq_len, - size_t layers, size_t kv_heads, size_t dim) { - size_t old_seq_len = current_seq_len; - size_t new_total_len = old_seq_len + seq_len; - size_t elements_per_token = kv_heads * dim; - size_t num_groups = (dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; - size_t scales_per_token = kv_heads * num_groups; - size_t bytes_per_token = elements_per_token * element_size; - - total_seq_len += seq_len; - - size_t effective_seq_len; - bool use_sliding_window = (window_size > 0 && new_total_len > window_size); - - if (use_sliding_window) { - effective_seq_len = window_size; - } else { - effective_seq_len = new_total_len; - } - - bool any_layer_updated = false; - - for (size_t layer_idx = 0; layer_idx < layers; layer_idx++) { - auto& cache = layer_caches[layer_idx]; - - void* k_output = gb->get_output(k_nodes[layer_idx]); - void* v_output = gb->get_output(v_nodes[layer_idx]); - - if (k_output && v_output) { - const auto& k_buffer = gb->get_output_buffer(k_nodes[layer_idx]); - const auto& v_buffer = gb->get_output_buffer(v_nodes[layer_idx]); - - size_t expected_elements = new_total_len * elements_per_token; - - if (k_buffer.total_size == expected_elements && v_buffer.total_size == expected_elements) { - any_layer_updated = true; - - if (!use_sliding_window) { - size_t total_bytes = new_total_len * bytes_per_token; - cache.keys.resize(total_bytes); - cache.values.resize(total_bytes); - - if (precision == Precision::INT8) { - size_t num_scales = new_total_len * scales_per_token; - cache.key_scales.resize(num_scales); - cache.value_scales.resize(num_scales); - - cactus_quantize_kv_fp16_to_int8( - static_cast(k_output), - reinterpret_cast(cache.keys.data()), - cache.key_scales.data(), - new_total_len, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - static_cast(v_output), - reinterpret_cast(cache.values.data()), - cache.value_scales.data(), - new_total_len, kv_heads, dim); - } else { - std::memcpy(cache.keys.data(), k_output, total_bytes); - std::memcpy(cache.values.data(), v_output, total_bytes); - } - } else { - size_t cache_bytes = window_size * bytes_per_token; - size_t remaining_window = window_size - sink_size; - size_t skip_tokens = new_total_len - window_size; - - bool first_slide = (cache.keys.size() != cache_bytes); - - if (first_slide) { - cache.keys.resize(cache_bytes); - cache.values.resize(cache_bytes); - } - - if (precision == Precision::INT8) { - size_t num_scales = window_size * scales_per_token; - if (first_slide) { - cache.key_scales.resize(num_scales); - cache.value_scales.resize(num_scales); - } - - const __fp16* k_fp16 = static_cast(k_output); - const __fp16* v_fp16 = static_cast(v_output); - - if (first_slide) { - cactus_quantize_kv_fp16_to_int8( - k_fp16, - reinterpret_cast(cache.keys.data()), - cache.key_scales.data(), - sink_size, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - v_fp16, - reinterpret_cast(cache.values.data()), - cache.value_scales.data(), - sink_size, kv_heads, dim); - } - - size_t src_offset = (sink_size + skip_tokens) * elements_per_token; - size_t dst_offset = sink_size * elements_per_token; - size_t scale_dst_offset = sink_size * scales_per_token; - - cactus_quantize_kv_fp16_to_int8( - k_fp16 + src_offset, - reinterpret_cast(cache.keys.data()) + dst_offset, - cache.key_scales.data() + scale_dst_offset, - remaining_window, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - v_fp16 + src_offset, - reinterpret_cast(cache.values.data()) + dst_offset, - cache.value_scales.data() + scale_dst_offset, - remaining_window, kv_heads, dim); - } else { - if (first_slide) { - size_t sink_bytes = sink_size * bytes_per_token; - std::memcpy(cache.keys.data(), k_output, sink_bytes); - std::memcpy(cache.values.data(), v_output, sink_bytes); - } - - const uint8_t* k_src = static_cast(k_output) + - (sink_size + skip_tokens) * bytes_per_token; - const uint8_t* v_src = static_cast(v_output) + - (sink_size + skip_tokens) * bytes_per_token; - size_t sink_bytes = sink_size * bytes_per_token; - size_t recent_bytes = remaining_window * bytes_per_token; - - std::memcpy(cache.keys.data() + sink_bytes, k_src, recent_bytes); - std::memcpy(cache.values.data() + sink_bytes, v_src, recent_bytes); - } - } - } - else if (seq_len * elements_per_token == k_buffer.total_size && - seq_len * elements_per_token == v_buffer.total_size) { - any_layer_updated = true; - - if (!use_sliding_window) { - size_t old_bytes = old_seq_len * bytes_per_token; - size_t new_bytes = seq_len * bytes_per_token; - size_t total_bytes = old_bytes + new_bytes; - cache.keys.resize(total_bytes); - cache.values.resize(total_bytes); - - if (precision == Precision::INT8) { - size_t num_scales = (old_seq_len + seq_len) * scales_per_token; - cache.key_scales.resize(num_scales); - cache.value_scales.resize(num_scales); - - size_t dst_offset = old_seq_len * elements_per_token; - size_t scale_offset = old_seq_len * scales_per_token; - - cactus_quantize_kv_fp16_to_int8( - static_cast(k_output), - reinterpret_cast(cache.keys.data()) + dst_offset, - cache.key_scales.data() + scale_offset, - seq_len, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - static_cast(v_output), - reinterpret_cast(cache.values.data()) + dst_offset, - cache.value_scales.data() + scale_offset, - seq_len, kv_heads, dim); - } else { - std::memcpy(cache.keys.data() + old_bytes, k_output, new_bytes); - std::memcpy(cache.values.data() + old_bytes, v_output, new_bytes); - } - } else { - size_t cache_bytes = window_size * bytes_per_token; - - if (cache.keys.size() != cache_bytes) { - cache.keys.resize(cache_bytes); - cache.values.resize(cache_bytes); - } - - size_t tokens_to_shift = window_size - sink_size - seq_len; - size_t sink_bytes = sink_size * bytes_per_token; - - if (tokens_to_shift > 0 && old_seq_len > sink_size) { - size_t shift_src = old_seq_len - tokens_to_shift; - if (shift_src > sink_size) { - size_t shift_bytes = tokens_to_shift * bytes_per_token; - - std::memmove(cache.keys.data() + sink_bytes, - cache.keys.data() + shift_src * bytes_per_token, - shift_bytes); - std::memmove(cache.values.data() + sink_bytes, - cache.values.data() + shift_src * bytes_per_token, - shift_bytes); - - if (precision == Precision::INT8) { - size_t sink_scale_offset = sink_size * scales_per_token; - size_t shift_src_scale_offset = shift_src * scales_per_token; - size_t shift_scale_count = tokens_to_shift * scales_per_token; - - std::memmove(cache.key_scales.data() + sink_scale_offset, - cache.key_scales.data() + shift_src_scale_offset, - shift_scale_count * sizeof(float)); - std::memmove(cache.value_scales.data() + sink_scale_offset, - cache.value_scales.data() + shift_src_scale_offset, - shift_scale_count * sizeof(float)); - } - } - } - - size_t append_token_offset = window_size - seq_len; - size_t append_bytes_offset = append_token_offset * bytes_per_token; - size_t new_bytes = seq_len * bytes_per_token; - - if (precision == Precision::INT8) { - size_t append_offset = append_token_offset * elements_per_token; - size_t scale_offset = append_token_offset * scales_per_token; - - cactus_quantize_kv_fp16_to_int8( - static_cast(k_output), - reinterpret_cast(cache.keys.data()) + append_offset, - cache.key_scales.data() + scale_offset, - seq_len, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - static_cast(v_output), - reinterpret_cast(cache.values.data()) + append_offset, - cache.value_scales.data() + scale_offset, - seq_len, kv_heads, dim); - } else { - std::memcpy(cache.keys.data() + append_bytes_offset, k_output, new_bytes); - std::memcpy(cache.values.data() + append_bytes_offset, v_output, new_bytes); - } - } - } - } - } - - if (any_layer_updated) { - current_seq_len = effective_seq_len; - } -} - -void KVCache::update_from_npu(size_t layer_idx, const __fp16* k_data, const __fp16* v_data, - size_t num_tokens, size_t kv_heads, size_t dim) { - if (layer_idx >= num_layers || !k_data || !v_data || num_tokens == 0) { - return; - } - - auto& cache = layer_caches[layer_idx]; - size_t old_seq_len = current_seq_len; - size_t new_total_len = old_seq_len + num_tokens; - size_t elements_per_token = kv_heads * dim; - size_t bytes_per_token = elements_per_token * element_size; - size_t num_groups = (dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; - size_t scales_per_token = kv_heads * num_groups; - - if (layer_idx == 0) { - total_seq_len += num_tokens; - } - - bool use_sliding_window = (window_size > 0 && new_total_len > window_size); - size_t effective_seq_len = use_sliding_window ? window_size : new_total_len; - - if (!use_sliding_window) { - size_t total_bytes = new_total_len * bytes_per_token; - cache.keys.resize(total_bytes); - cache.values.resize(total_bytes); - - size_t num_scales = new_total_len * scales_per_token; - cache.key_scales.resize(num_scales); - cache.value_scales.resize(num_scales); - - size_t dst_offset = old_seq_len * elements_per_token; - size_t scale_offset = old_seq_len * scales_per_token; - - cactus_quantize_kv_fp16_to_int8( - k_data, - reinterpret_cast(cache.keys.data()) + dst_offset, - cache.key_scales.data() + scale_offset, - num_tokens, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - v_data, - reinterpret_cast(cache.values.data()) + dst_offset, - cache.value_scales.data() + scale_offset, - num_tokens, kv_heads, dim); - } else { - size_t cache_bytes = window_size * bytes_per_token; - - if (cache.keys.size() != cache_bytes) { - cache.keys.resize(cache_bytes); - cache.values.resize(cache_bytes); - - size_t num_scales = window_size * scales_per_token; - cache.key_scales.resize(num_scales); - cache.value_scales.resize(num_scales); - } - - size_t remaining_window = window_size - sink_size; - - if (num_tokens >= remaining_window) { - size_t skip_tokens = num_tokens - remaining_window; - size_t dst_offset = sink_size * elements_per_token; - size_t scale_offset = sink_size * scales_per_token; - - cactus_quantize_kv_fp16_to_int8( - k_data + skip_tokens * elements_per_token, - reinterpret_cast(cache.keys.data()) + dst_offset, - cache.key_scales.data() + scale_offset, - remaining_window, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - v_data + skip_tokens * elements_per_token, - reinterpret_cast(cache.values.data()) + dst_offset, - cache.value_scales.data() + scale_offset, - remaining_window, kv_heads, dim); - } else { - size_t tokens_to_shift = remaining_window - num_tokens; - - if (tokens_to_shift > 0 && old_seq_len > sink_size) { - size_t shift_src = old_seq_len - tokens_to_shift; - if (shift_src > sink_size) { - size_t sink_offset = sink_size * elements_per_token; - size_t shift_src_offset = shift_src * elements_per_token; - size_t shift_bytes = tokens_to_shift * bytes_per_token; - - std::memmove(cache.keys.data() + sink_offset, - cache.keys.data() + shift_src_offset, - shift_bytes); - std::memmove(cache.values.data() + sink_offset, - cache.values.data() + shift_src_offset, - shift_bytes); - - size_t sink_scale_offset = sink_size * scales_per_token; - size_t shift_src_scale_offset = shift_src * scales_per_token; - size_t shift_scale_count = tokens_to_shift * scales_per_token; - - std::memmove(cache.key_scales.data() + sink_scale_offset, - cache.key_scales.data() + shift_src_scale_offset, - shift_scale_count * sizeof(float)); - std::memmove(cache.value_scales.data() + sink_scale_offset, - cache.value_scales.data() + shift_src_scale_offset, - shift_scale_count * sizeof(float)); - } - } - - size_t append_token_offset = window_size - num_tokens; - size_t append_offset = append_token_offset * elements_per_token; - size_t scale_offset = append_token_offset * scales_per_token; - - cactus_quantize_kv_fp16_to_int8( - k_data, - reinterpret_cast(cache.keys.data()) + append_offset, - cache.key_scales.data() + scale_offset, - num_tokens, kv_heads, dim); - - cactus_quantize_kv_fp16_to_int8( - v_data, - reinterpret_cast(cache.values.data()) + append_offset, - cache.value_scales.data() + scale_offset, - num_tokens, kv_heads, dim); - } - } - - if (layer_idx == num_layers - 1) { - current_seq_len = effective_seq_len; - } -} - -const int8_t* KVCache::get_keys_int8(size_t layer) const { - if (layer >= num_layers || current_seq_len == 0) { - return nullptr; - } - return reinterpret_cast(layer_caches[layer].keys.data()); -} - -const int8_t* KVCache::get_values_int8(size_t layer) const { - if (layer >= num_layers || current_seq_len == 0) { - return nullptr; - } - return reinterpret_cast(layer_caches[layer].values.data()); -} - -const float* KVCache::get_key_scales(size_t layer) const { - if (layer >= num_layers || current_seq_len == 0) { - return nullptr; - } - return layer_caches[layer].key_scales.data(); -} - -const float* KVCache::get_value_scales(size_t layer) const { - if (layer >= num_layers || current_seq_len == 0) { - return nullptr; - } - return layer_caches[layer].value_scales.data(); -} - -void ConvCache::init(size_t layers, size_t hidden_dim, size_t window_len, Precision model_precision) { - num_layers = layers; - hidden_size = hidden_dim; - window_size = window_len; - precision = model_precision; - element_size = PrecisionTraits::size_of(precision); - - size_t state_bytes = window_size * hidden_size * element_size; - layer_states.resize(num_layers); - for (auto& state : layer_states) { - state.data.resize(state_bytes); - std::memset(state.data.data(), 0, state_bytes); - state.head = 0; - state.count = 0; - } -} - -ConvCache::CircularView ConvCache::get_window(size_t layer) const { - CircularView view; - if (layer >= num_layers) { - view.ptr1 = nullptr; - view.len1 = 0; - view.ptr2 = nullptr; - view.len2 = 0; - view.total_len = 0; - return view; - } - - const auto& state = layer_states[layer]; - if (state.count == 0) { - view.ptr1 = nullptr; - view.len1 = 0; - view.ptr2 = nullptr; - view.len2 = 0; - view.total_len = 0; - return view; - } - - size_t stride = hidden_size * element_size; - - if (state.count < window_size) { - view.ptr1 = state.data.data(); - view.len1 = state.count; - view.ptr2 = nullptr; - view.len2 = 0; - view.total_len = state.count; - return view; - } - - view.ptr1 = state.data.data(); - view.len1 = state.head; - view.ptr2 = state.data.data() + state.head * stride; - view.len2 = window_size - state.head; - view.total_len = window_size; - return view; -} - -void ConvCache::update(CactusGraph* gb, size_t layer, const size_t bx_node) { - if (layer >= num_layers || !bx_node || window_size == 0 || hidden_size == 0) { - return; - } - - auto& state = layer_states[layer]; - const void* output_ptr = gb->get_output(bx_node); - if (!output_ptr) { - return; - } - - const auto& buffer = gb->get_output_buffer(bx_node); - const size_t stride_bytes = hidden_size * element_size; - - size_t rows = 1; - if (!buffer.shape.empty()) { - if (buffer.shape.size() == 1) { - rows = 1; - } else { - rows = buffer.shape[0]; - } - } - - if (buffer.total_size > 0 && hidden_size > 0) { - size_t inferred = buffer.total_size / hidden_size; - if (inferred > 0) { - rows = inferred; - } - } - - if (rows == 0) { - return; - } - - size_t copy_rows = std::min(rows, window_size); - size_t start_row = rows > window_size ? rows - window_size : 0; - - const uint8_t* src = static_cast(output_ptr) + start_row * stride_bytes; - - for (size_t i = 0; i < copy_rows; ++i) { - std::memcpy(state.data.data() + state.head * stride_bytes, src + i * stride_bytes, stride_bytes); - state.head = (state.head + 1) % window_size; - if (state.count < window_size) { - ++state.count; - } - } -} - -void ConvCache::reset() { - for (auto& state : layer_states) { - std::fill(state.data.begin(), state.data.end(), 0); - state.head = 0; - state.count = 0; - } -} - -} -} \ No newline at end of file diff --git a/cactus/engine/engine_constraints.cpp b/cactus/engine/engine_constraints.cpp deleted file mode 100644 index 323484e09..000000000 --- a/cactus/engine/engine_constraints.cpp +++ /dev/null @@ -1,726 +0,0 @@ -#include "engine.h" - -namespace cactus { -namespace engine { - -constexpr float FORCE_BIAS = 500.0f; -constexpr float BLOCK_BIAS = -500.0f; - -void ToolCallConstrainer::add_tokens_for_string(const std::string& str, std::unordered_set& token_set) { - if (!tokenizer_) return; - auto tokens = tokenizer_->encode(str); - for (uint32_t t : tokens) { - token_set.insert(t); - } -} - -void ToolCallConstrainer::tokenize_grammar_elements() { - if (!tokenizer_) return; - - open_brace_tokens_.clear(); - close_brace_tokens_.clear(); - colon_tokens_.clear(); - comma_tokens_.clear(); - name_key_tokens_.clear(); - args_key_tokens_.clear(); - quote_tokens_.clear(); - backtick_tokens_.clear(); - all_func_name_tokens_.clear(); - func_name_sequences_.clear(); - tool_start_tokens_.clear(); - tool_end_tokens_.clear(); - bracket_open_tokens_.clear(); - bracket_close_tokens_.clear(); - paren_open_tokens_.clear(); - paren_close_tokens_.clear(); - equals_tokens_.clear(); - - add_tokens_for_string("`", backtick_tokens_); - add_tokens_for_string("``", backtick_tokens_); - add_tokens_for_string("```", backtick_tokens_); - add_tokens_for_string("````", backtick_tokens_); - add_tokens_for_string("```json", backtick_tokens_); - add_tokens_for_string("```JSON", backtick_tokens_); - add_tokens_for_string("``` json", backtick_tokens_); - add_tokens_for_string("```\n", backtick_tokens_); - add_tokens_for_string("` ", backtick_tokens_); - - if (model_type_ == Config::ModelType::LFM2) { - add_tokens_for_string("<|tool_call_start|>", tool_start_tokens_); - add_tokens_for_string("<|tool_call_end|>", tool_end_tokens_); - add_tokens_for_string("[", bracket_open_tokens_); - add_tokens_for_string("]", bracket_close_tokens_); - add_tokens_for_string("(", paren_open_tokens_); - add_tokens_for_string(")", paren_close_tokens_); - add_tokens_for_string("=", equals_tokens_); - add_tokens_for_string(",", comma_tokens_); - add_tokens_for_string("\"", quote_tokens_); - - for (const auto& name : function_names_) { - auto tokens = tokenizer_->encode(name); - func_name_sequences_[name] = tokens; - for (uint32_t t : tokens) { - all_func_name_tokens_.insert(t); - } - } - } else if (model_type_ == Config::ModelType::GEMMA) { - gemma_call_start_tokens_.clear(); - gemma_call_end_tokens_.clear(); - gemma_call_prefix_tokens_.clear(); - escape_tokens_.clear(); - gemma_response_start_tokens_.clear(); - - add_tokens_for_string("", gemma_call_start_tokens_); - add_tokens_for_string("", gemma_call_end_tokens_); - add_tokens_for_string("", gemma_response_start_tokens_); - add_tokens_for_string("call:", gemma_call_prefix_tokens_); - add_tokens_for_string("", escape_tokens_); - - add_tokens_for_string("{", open_brace_tokens_); - add_tokens_for_string("}", close_brace_tokens_); - add_tokens_for_string(":", colon_tokens_); - add_tokens_for_string(",", comma_tokens_); - - for (const auto& name : function_names_) { - auto tokens = tokenizer_->encode(name); - func_name_sequences_[name] = tokens; - for (uint32_t t : tokens) { - all_func_name_tokens_.insert(t); - } - } - } else { - qwen_tool_call_start_tokens_.clear(); - qwen_tool_call_end_tokens_.clear(); - - add_tokens_for_string("", qwen_tool_call_start_tokens_); - add_tokens_for_string("", qwen_tool_call_end_tokens_); - - add_tokens_for_string("{", open_brace_tokens_); - add_tokens_for_string("}", close_brace_tokens_); - add_tokens_for_string(":", colon_tokens_); - add_tokens_for_string(",", comma_tokens_); - add_tokens_for_string("\"", quote_tokens_); - - add_tokens_for_string("\"name\"", name_key_tokens_); - add_tokens_for_string("name", name_key_tokens_); - - add_tokens_for_string("\"arguments\"", args_key_tokens_); - add_tokens_for_string("arguments", args_key_tokens_); - - for (const auto& name : function_names_) { - std::string quoted_name = "\"" + name + "\""; - auto tokens = tokenizer_->encode(quoted_name); - func_name_sequences_[name] = tokens; - for (uint32_t t : tokens) { - all_func_name_tokens_.insert(t); - } - auto unquoted_tokens = tokenizer_->encode(name); - for (uint32_t t : unquoted_tokens) { - all_func_name_tokens_.insert(t); - } - } - } -} - -void ToolCallConstrainer::init(Config::ModelType model_type, - const std::vector& function_names, - Tokenizer* tokenizer) { - model_type_ = model_type; - function_names_ = function_names; - tokenizer_ = tokenizer; - generated_text_.clear(); - brace_depth_ = 0; - active_ = !function_names.empty() && tokenizer != nullptr; - - if (model_type_ == Config::ModelType::LFM2) { - state_ = State::LFM_START; - } else if (model_type_ == Config::ModelType::GEMMA) { - state_ = State::GEMMA_START; - } else { - state_ = State::QWEN_START; - } - - if (!active_) { - return; - } - - tokenize_grammar_elements(); - compute_bias(); -} - -void ToolCallConstrainer::update(uint32_t /*token_id*/, const std::string& decoded_text) { - if (!active_) return; - - generated_text_ += decoded_text; - - if (model_type_ == Config::ModelType::LFM2) { - switch (state_) { - case State::LFM_START: - if (generated_text_.find("<|tool_call_start|>") != std::string::npos) { - state_ = State::LFM_EXPECT_BRACKET; - generated_text_.clear(); - } - break; - - case State::LFM_EXPECT_BRACKET: - if (generated_text_.find("[") != std::string::npos) { - state_ = State::LFM_IN_FUNC_NAME; - generated_text_.clear(); - } - break; - - case State::LFM_IN_FUNC_NAME: - for (const auto& name : function_names_) { - if (generated_text_.find(name) != std::string::npos) { - state_ = State::LFM_EXPECT_PAREN; - generated_text_.clear(); - break; - } - } - break; - - case State::LFM_EXPECT_PAREN: - if (generated_text_.find("(") != std::string::npos) { - state_ = State::LFM_IN_ARGUMENTS; - generated_text_.clear(); - } - break; - - case State::LFM_IN_ARGUMENTS: - if (generated_text_.find(")") != std::string::npos) { - state_ = State::LFM_EXPECT_BRACKET_CLOSE; - } - break; - - case State::LFM_EXPECT_BRACKET_CLOSE: - if (generated_text_.find("]") != std::string::npos) { - state_ = State::LFM_EXPECT_END; - generated_text_.clear(); - } - break; - - case State::LFM_EXPECT_END: - if (generated_text_.find("<|tool_call_end|>") != std::string::npos) { - state_ = State::DONE; - generated_text_.clear(); - } - break; - - default: - break; - } - } else if (model_type_ == Config::ModelType::GEMMA) { - switch (state_) { - case State::GEMMA_START: - if (generated_text_.find("") != std::string::npos) { - state_ = State::GEMMA_EXPECT_CALL; - generated_text_.clear(); - } - break; - - case State::GEMMA_EXPECT_CALL: - if (generated_text_.find("call:") != std::string::npos) { - state_ = State::GEMMA_IN_FUNC_NAME; - generated_text_.clear(); - } - break; - - case State::GEMMA_IN_FUNC_NAME: - for (const auto& name : function_names_) { - if (generated_text_.find(name) != std::string::npos) { - state_ = State::GEMMA_EXPECT_BRACE; - generated_text_.clear(); - break; - } - } - break; - - case State::GEMMA_EXPECT_BRACE: - if (generated_text_.find("{") != std::string::npos) { - state_ = State::GEMMA_IN_ARGUMENTS; - brace_depth_ = 1; - generated_text_.clear(); - } - break; - - case State::GEMMA_IN_ARGUMENTS: - for (char c : decoded_text) { - if (c == '{') brace_depth_++; - else if (c == '}') { - brace_depth_--; - if (brace_depth_ == 0) { - state_ = State::GEMMA_EXPECT_END; - generated_text_.clear(); - break; - } - } - } - break; - - case State::GEMMA_EXPECT_END: - if (generated_text_.find("") != std::string::npos) { - state_ = State::DONE; - generated_text_.clear(); - } - break; - - default: - break; - } - } else { - switch (state_) { - case State::QWEN_START: - if (generated_text_.find("") != std::string::npos) { - state_ = State::QWEN_EXPECT_OPEN_BRACE; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_OPEN_BRACE: - if (generated_text_.find("{") != std::string::npos) { - state_ = State::QWEN_EXPECT_NAME_KEY; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_NAME_KEY: - if (generated_text_.find("name") != std::string::npos) { - state_ = State::QWEN_EXPECT_NAME_COLON; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_NAME_COLON: - if (generated_text_.find(":") != std::string::npos) { - state_ = State::QWEN_EXPECT_NAME_VALUE; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_NAME_VALUE: - for (const auto& name : function_names_) { - if (generated_text_.find(name) != std::string::npos) { - state_ = State::QWEN_EXPECT_COMMA; - generated_text_.clear(); - break; - } - } - break; - - case State::QWEN_EXPECT_COMMA: - if (generated_text_.find(",") != std::string::npos) { - state_ = State::QWEN_EXPECT_ARGS_KEY; - generated_text_.clear(); - } else if (generated_text_.find("}") != std::string::npos) { - state_ = State::QWEN_EXPECT_END; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_ARGS_KEY: - if (generated_text_.find("arguments") != std::string::npos) { - state_ = State::QWEN_EXPECT_ARGS_COLON; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_ARGS_COLON: - if (generated_text_.find(":") != std::string::npos) { - state_ = State::QWEN_IN_ARGUMENTS; - brace_depth_ = 0; - generated_text_.clear(); - } - break; - - case State::QWEN_IN_ARGUMENTS: - for (char c : decoded_text) { - if (c == '{') brace_depth_++; - else if (c == '}') { - if (brace_depth_ > 0) { - brace_depth_--; - } else { - state_ = State::QWEN_EXPECT_CLOSE_BRACE; - break; - } - } - } - break; - - case State::QWEN_EXPECT_CLOSE_BRACE: - if (generated_text_.find("}") != std::string::npos) { - state_ = State::QWEN_EXPECT_END; - generated_text_.clear(); - } - break; - - case State::QWEN_EXPECT_END: - if (generated_text_.find("") != std::string::npos) { - state_ = State::DONE; - generated_text_.clear(); - } - break; - - default: - break; - } - } - - compute_bias(); -} - -void ToolCallConstrainer::compute_bias() { - current_bias_.clear(); - - if (!active_) return; - - for (uint32_t t : backtick_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - - if (model_type_ == Config::ModelType::LFM2) { - switch (state_) { - case State::LFM_START: - for (uint32_t t : tool_start_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : bracket_open_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : paren_open_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::LFM_EXPECT_BRACKET: - for (uint32_t t : bracket_open_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : paren_open_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : bracket_close_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::LFM_IN_FUNC_NAME: - for (uint32_t t : all_func_name_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : bracket_close_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : paren_close_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : equals_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::LFM_EXPECT_PAREN: - for (uint32_t t : paren_open_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : bracket_close_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : equals_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::LFM_IN_ARGUMENTS: - for (uint32_t t : paren_close_tokens_) { - current_bias_[t] = 15.0f; - } - for (uint32_t t : equals_tokens_) { - current_bias_[t] = 10.0f; - } - for (uint32_t t : comma_tokens_) { - current_bias_[t] = 8.0f; - } - for (uint32_t t : quote_tokens_) { - current_bias_[t] = 5.0f; - } - for (uint32_t t : bracket_close_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : tool_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::LFM_EXPECT_BRACKET_CLOSE: - for (uint32_t t : bracket_close_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : paren_open_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : equals_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::LFM_EXPECT_END: - for (uint32_t t : tool_end_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : bracket_open_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : paren_open_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - default: - break; - } - } else if (model_type_ == Config::ModelType::GEMMA) { - for (uint32_t t : gemma_response_start_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - - switch (state_) { - case State::GEMMA_START: - for (uint32_t t : gemma_call_start_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_EXPECT_CALL: - for (uint32_t t : gemma_call_prefix_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_IN_FUNC_NAME: - for (uint32_t t : all_func_name_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_EXPECT_BRACE: - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_IN_ARGUMENTS: - for (uint32_t t : colon_tokens_) { - current_bias_[t] = 10.0f; - } - for (uint32_t t : comma_tokens_) { - current_bias_[t] = 8.0f; - } - for (uint32_t t : escape_tokens_) { - current_bias_[t] = 5.0f; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = 3.0f; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = 3.0f; - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_EXPECT_END: - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : gemma_call_start_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - default: - break; - } - } else { - switch (state_) { - case State::QWEN_START: - for (uint32_t t : qwen_tool_call_start_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::QWEN_EXPECT_OPEN_BRACE: - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : qwen_tool_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::QWEN_EXPECT_NAME_KEY: - for (uint32_t t : name_key_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : quote_tokens_) { - current_bias_[t] = 5.0f; - } - for (uint32_t t : all_func_name_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : args_key_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::QWEN_EXPECT_NAME_COLON: - for (uint32_t t : colon_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - break; - - case State::QWEN_EXPECT_NAME_VALUE: - for (uint32_t t : all_func_name_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : quote_tokens_) { - current_bias_[t] = 5.0f; - } - break; - - case State::QWEN_EXPECT_COMMA: - for (uint32_t t : comma_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = 5.0f; - } - break; - - case State::QWEN_EXPECT_ARGS_KEY: - for (uint32_t t : args_key_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : quote_tokens_) { - current_bias_[t] = 5.0f; - } - break; - - case State::QWEN_EXPECT_ARGS_COLON: - for (uint32_t t : colon_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - break; - - case State::QWEN_IN_ARGUMENTS: - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = 3.0f; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = 3.0f; - } - for (uint32_t t : colon_tokens_) { - current_bias_[t] = 2.0f; - } - for (uint32_t t : comma_tokens_) { - current_bias_[t] = 2.0f; - } - for (uint32_t t : quote_tokens_) { - current_bias_[t] = 2.0f; - } - for (uint32_t t : qwen_tool_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::QWEN_EXPECT_CLOSE_BRACE: - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - break; - - case State::QWEN_EXPECT_END: - for (uint32_t t : qwen_tool_call_end_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : qwen_tool_call_start_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - default: - break; - } - } -} - -void ToolCallConstrainer::reset() { - generated_text_.clear(); - current_bias_.clear(); - brace_depth_ = 0; - - if (model_type_ == Config::ModelType::LFM2) { - state_ = State::LFM_START; - } else if (model_type_ == Config::ModelType::GEMMA) { - state_ = State::GEMMA_START; - } else { - state_ = State::QWEN_START; - } - - if (active_) { - compute_bias(); - } -} - - -void Model::set_tool_constraints(const std::vector& function_names) { - tool_constrainer_.init(config_.model_type, function_names, tokenizer_.get()); -} - -void Model::clear_tool_constraints() { - tool_constrainer_.reset(); - tool_constrainer_.init(config_.model_type, {}, tokenizer_.get()); -} - -void Model::update_tool_constraints(uint32_t token_id) { - if (tool_constrainer_.is_active() && tokenizer_) { - std::string decoded = tokenizer_->decode({token_id}); - tool_constrainer_.update(token_id, decoded); - } -} - -} // namespace engine -} // namespace cactus \ No newline at end of file diff --git a/cactus/engine/engine_image.cpp b/cactus/engine/engine_image.cpp deleted file mode 100644 index 9d38cd371..000000000 --- a/cactus/engine/engine_image.cpp +++ /dev/null @@ -1,600 +0,0 @@ -#define STB_IMAGE_IMPLEMENTATION -#define STB_IMAGE_RESIZE_IMPLEMENTATION - -#define STBI_NO_BMP -#define STBI_NO_PSD -#define STBI_NO_HDR -#define STBI_NO_PIC -#define STBI_NO_PNM -#define STBI_NO_TGA - -#include "engine.h" -#include -#include -#include -#include -#include -#include -#include - -namespace cactus { -namespace engine { - -Siglip2Preprocessor::PreprocessedImage::~PreprocessedImage() { - pixel_values.clear(); - pixel_attention_mask.clear(); -} - -Siglip2Preprocessor::Siglip2Preprocessor(const Config& config) - : config_(config) {} - -Siglip2Preprocessor::Siglip2Preprocessor() : config_() {} - -Siglip2Preprocessor::~Siglip2Preprocessor() = default; - -int Siglip2Preprocessor::round_by_factor(int number, int factor) { - if (factor == 0) { - return number; - } - double scaled = static_cast(number) / static_cast(factor); - long long rounded = std::llround(scaled); - - return static_cast(rounded * factor); -} - -std::pair Siglip2Preprocessor::smart_resize(int height, int width) { - const int total_factor = config_.patch_size * config_.downsample_factor; - const int64_t min_pixels = static_cast(config_.min_image_tokens) * - config_.patch_size * config_.patch_size * - config_.downsample_factor * config_.downsample_factor; - const int64_t max_pixels = static_cast(config_.max_image_tokens) * - config_.patch_size * config_.patch_size * - config_.downsample_factor * config_.downsample_factor; - - int h_bar = std::max(total_factor, round_by_factor(height, total_factor)); - int w_bar = std::max(total_factor, round_by_factor(width, total_factor)); - - - if (static_cast(h_bar) * static_cast(w_bar) > max_pixels) { - double beta = std::sqrt((static_cast(height) * static_cast(width)) / - static_cast(max_pixels)); - h_bar = std::max(total_factor, - static_cast(std::floor(height / beta / total_factor)) * total_factor); - w_bar = std::max(total_factor, - static_cast(std::floor(width / beta / total_factor)) * total_factor); - } else if (static_cast(h_bar) * static_cast(w_bar) < min_pixels) { - double beta = std::sqrt(static_cast(min_pixels) / - (static_cast(height) * static_cast(width))); - h_bar = std::max(total_factor, - static_cast(std::ceil(height * beta / total_factor)) * total_factor); - w_bar = std::max(total_factor, - static_cast(std::ceil(width * beta / total_factor)) * total_factor); - } - - return {w_bar, h_bar}; -} - - -bool Siglip2Preprocessor::is_image_too_large(int height, int width) { - const int total_factor = config_.patch_size * config_.downsample_factor; - const int h_bar = std::max(config_.patch_size, round_by_factor(height, total_factor)); - const int w_bar = std::max(config_.patch_size, round_by_factor(width, total_factor)); - - const int64_t pixels = static_cast(h_bar) * static_cast(w_bar); - const double max_pixels = static_cast(config_.max_image_tokens) * - config_.patch_size * config_.patch_size * - config_.downsample_factor * config_.downsample_factor * - config_.max_pixels_tolerance; - - bool result = static_cast(pixels) > max_pixels; - - return result; -} - - -std::pair Siglip2Preprocessor::find_closest_aspect_ratio(float aspect_ratio, int width, int height) { - float best_ratio_diff = std::numeric_limits::infinity(); - std::pair best_ratio = {1, 1}; - int area = width * height; - - std::vector> target_ratios; - for (int n = config_.min_tiles; n <= config_.max_tiles; ++n) { - for (int w = 1; w <= n; ++w) { - for (int h = 1; h <= n; ++h) { - int total_tiles = w * h; - if (total_tiles >= config_.min_tiles && total_tiles <= config_.max_tiles) { - target_ratios.push_back({w, h}); - } - } - } - } - - std::sort(target_ratios.begin(), target_ratios.end()); - target_ratios.erase(std::unique(target_ratios.begin(), target_ratios.end()), target_ratios.end()); - std::sort(target_ratios.begin(), target_ratios.end(), [](const auto& a, const auto& b) { - return (a.first * a.second) < (b.first * b.second); - }); - - for (const auto& ratio : target_ratios) { - float target_aspect_ratio = static_cast(ratio.first) / ratio.second; - float ratio_diff = std::abs(aspect_ratio - target_aspect_ratio); - - if (ratio_diff < best_ratio_diff) { - best_ratio_diff = ratio_diff; - best_ratio = ratio; - } else if (ratio_diff == best_ratio_diff) { - int target_area = config_.tile_size * config_.tile_size * ratio.first * ratio.second; - if (area > 0.5f * target_area) { - best_ratio = ratio; - } - } - } - - - return best_ratio; -} - -std::pair Siglip2Preprocessor::get_grid_layout(int height, int width) { - float aspect_ratio = (float)width / (float)height; - auto [grid_width, grid_height] = find_closest_aspect_ratio(aspect_ratio, width, height); - - int target_width = config_.tile_size * grid_width; - int target_height = config_.tile_size * grid_height; - - return {target_width, target_height}; -} - -Siglip2Preprocessor::SpatialShapeResult Siglip2Preprocessor::compute_spatial_shapes(int height, int width) { - if (height <= 0 || width <= 0) { - throw std::runtime_error("Image dimensions must be positive"); - } - - if (config_.patch_size <= 0) { - throw std::runtime_error("Patch size must be positive"); - } - - const int patch = config_.patch_size; - auto [resized_width, resized_height] = smart_resize(height, width); - const bool should_split = config_.do_image_splitting && is_image_too_large(height, width); - - - SpatialShapeResult result; - result.grid_rows = 1; - result.grid_cols = 1; - - if (should_split) { - if (config_.tile_size % patch != 0) { - throw std::runtime_error("Tile size must be divisible by patch size"); - } - - auto [grid_target_width, grid_target_height] = get_grid_layout(height, width); - result.grid_cols = grid_target_width / config_.tile_size; - result.grid_rows = grid_target_height / config_.tile_size; - - - const int patches_per_tile_side = config_.tile_size / patch; - const auto tile_shape = std::make_pair(patches_per_tile_side, patches_per_tile_side); - - result.shapes.reserve(static_cast(result.grid_rows) * result.grid_cols + 1); - for (int row = 0; row < result.grid_rows; ++row) { - for (int col = 0; col < result.grid_cols; ++col) { - result.shapes.push_back(tile_shape); - - } - } - - if (config_.use_thumbnail && result.grid_rows * result.grid_cols != 1) { - if (resized_height % patch != 0 || resized_width % patch != 0) { - throw std::runtime_error("Resized thumbnail dimensions must be divisible by patch size"); - } - result.shapes.emplace_back(resized_height / patch, resized_width / patch); - - } - } else { - int target_width = resized_width; - int target_height = resized_height; - if (!config_.do_resize) { - target_width = width; - target_height = height; - } - - if (target_height % patch != 0 || target_width % patch != 0) { - throw std::runtime_error("Target dimensions must be divisible by patch size"); - } - - result.shapes.emplace_back(target_height / patch, target_width / patch); - - } - - return result; -} - - -Siglip2Preprocessor::PreprocessedImage Siglip2Preprocessor::preprocess_from_file(const std::string& image_path) { - int width, height, channels; - unsigned char* img_data = stbi_load(image_path.c_str(), &width, &height, &channels, 0); - - if (!img_data) { - throw std::runtime_error("Failed to load image: " + image_path + " - " + std::string(stbi_failure_reason())); - } - - PreprocessedImage result; - result = preprocess_from_memory(img_data, width, height, channels); - - stbi_image_free(img_data); - return result; -} - -Siglip2Preprocessor::PreprocessedImage Siglip2Preprocessor::preprocess_from_memory( - const unsigned char* img_data, int width, int height, int channels) { - if (!img_data) { - throw std::runtime_error("Invalid image data pointer"); - } - - const int expected_channels = 3; - std::vector rgb_data; - const unsigned char* source_data = img_data; - int source_channels = channels; - - if (config_.do_convert_rgb && channels != expected_channels) { - rgb_data = convert_to_rgb(img_data, width, height, channels); - source_data = rgb_data.data(); - source_channels = expected_channels; - - } - - if (source_channels != expected_channels) { - throw std::runtime_error("Image must have 3 channels (RGB)"); - } - - const int patch = config_.patch_size; - const int downsample = config_.downsample_factor; - const int patch_dim = patch * patch * expected_channels; - - if (patch <= 0) { - throw std::runtime_error("Patch size must be positive"); - } - if (config_.tile_size % patch != 0) { - throw std::runtime_error("Tile size must be divisible by patch size"); - } - - auto [resized_width, resized_height] = smart_resize(height, width); - - - const int patches_per_tile_side = config_.tile_size / patch; - const int tile_patch_count = patches_per_tile_side * patches_per_tile_side; - const int thumbnail_patch_cap = config_.max_image_tokens * downsample * downsample; - int max_patches_per_tile = config_.max_num_patches; - max_patches_per_tile = std::max(max_patches_per_tile, tile_patch_count); - max_patches_per_tile = std::max(max_patches_per_tile, thumbnail_patch_cap); - - const bool allow_splitting = config_.do_image_splitting; - const bool should_split = allow_splitting && is_image_too_large(height, width); - - - size_t expected_tiles = should_split ? static_cast(config_.max_tiles) : 1; - if (should_split && config_.use_thumbnail) { - expected_tiles += 1; - } - - std::vector> tile_patches; - tile_patches.reserve(expected_tiles); - std::vector> spatial_shapes; - spatial_shapes.reserve(expected_tiles); - - auto normalize_and_patchify = [&](const float* data_ptr, int img_width, int img_height) { - if (img_height % patch != 0 || img_width % patch != 0) { - throw std::runtime_error("Image dimensions must be divisible by patch size"); - } - std::vector normalized = normalize_image(data_ptr, img_width, img_height, expected_channels); - auto patches = convert_image_to_patches(normalized, img_width, img_height, expected_channels, patch); - std::vector flattened(patches.size() * patch_dim); - for (size_t idx = 0; idx < patches.size(); ++idx) { - std::copy(patches[idx].begin(), patches[idx].end(), flattened.begin() + idx * patch_dim); - } - tile_patches.push_back(std::move(flattened)); - spatial_shapes.emplace_back(img_height / patch, img_width / patch); - }; - - int grid_rows = 1; - int grid_cols = 1; - bool thumbnail_added = false; - - if (should_split) { - auto [grid_target_width, grid_target_height] = get_grid_layout(height, width); - grid_cols = grid_target_width / config_.tile_size; - grid_rows = grid_target_height / config_.tile_size; - printf("[debug] grid_target_width=%d grid_target_height=%d grid_cols=%d grid_rows=%d\n", - grid_target_width, grid_target_height, grid_cols, grid_rows); - - - std::vector resized_grid = resize_image( - source_data, width, height, grid_target_width, grid_target_height, expected_channels); - - std::vector tile_buffer( - static_cast(config_.tile_size) * config_.tile_size * expected_channels); - - for (int row = 0; row < grid_rows; ++row) { - for (int col = 0; col < grid_cols; ++col) { - for (int y = 0; y < config_.tile_size; ++y) { - const float* src_row = resized_grid.data() + - ((row * config_.tile_size + y) * grid_target_width + col * config_.tile_size) * - expected_channels; - float* dst_row = tile_buffer.data() + (y * config_.tile_size) * expected_channels; - std::copy_n(src_row, - static_cast(config_.tile_size) * expected_channels, - dst_row); - } - normalize_and_patchify(tile_buffer.data(), config_.tile_size, config_.tile_size); - - } - } - - if (config_.use_thumbnail && grid_rows * grid_cols != 1) { - std::vector thumbnail_bytes = resize_image( - source_data, width, height, resized_width, resized_height, expected_channels); - normalize_and_patchify(thumbnail_bytes.data(), resized_width, resized_height); - thumbnail_added = true; - - } - } else { - int target_width = resized_width; - int target_height = resized_height; - - const bool needs_resize = config_.do_resize && (width != target_width || height != target_height); - - std::vector resized_image; - if (needs_resize) { - resized_image = resize_image(source_data, width, height, target_width, target_height, expected_channels); - normalize_and_patchify(resized_image.data(), target_width, target_height); - - } else { - resized_image.resize(static_cast(width) * height * expected_channels); - for (size_t idx = 0; idx < resized_image.size(); ++idx) { - resized_image[idx] = static_cast(source_data[idx]); - } - normalize_and_patchify(resized_image.data(), width, height); - target_width = width; - target_height = height; - resized_width = target_width; - resized_height = target_height; - - } - - grid_rows = 1; - grid_cols = 1; - } - - PreprocessedImage result = pad_patches(tile_patches, spatial_shapes, patch_dim, max_patches_per_tile); - - - result.image_rows = grid_rows; - result.image_cols = grid_cols; - result.image_width = resized_width; - result.image_height = resized_height; - - auto compute_tokens = [&](int patches_h, int patches_w) -> int { - int tokens_h = (patches_h + downsample - 1) / downsample; - int tokens_w = (patches_w + downsample - 1) / downsample; - return tokens_h * tokens_w; - }; - - result.tokens_per_tile = spatial_shapes.empty() - ? 0 - : compute_tokens(spatial_shapes.front().first, spatial_shapes.front().second); - - - if (thumbnail_added && !spatial_shapes.empty()) { - const auto& thumb_shape = spatial_shapes.back(); - result.thumbnail_tokens = compute_tokens(thumb_shape.first, thumb_shape.second); - - } else { - result.thumbnail_tokens = 0; - } - - return result; -} - -std::vector Siglip2Preprocessor::convert_to_rgb( - const unsigned char* img_data, int width, int height, int channels) { - - std::vector rgb_data(width * height * 3); - - if (channels == 1) { - for (int i = 0; i < width * height; ++i) { - rgb_data[i * 3 + 0] = img_data[i]; - rgb_data[i * 3 + 1] = img_data[i]; - rgb_data[i * 3 + 2] = img_data[i]; - } - } else if (channels == 4) { - for (int i = 0; i < width * height; ++i) { - rgb_data[i * 3 + 0] = img_data[i * 4 + 0]; - rgb_data[i * 3 + 1] = img_data[i * 4 + 1]; - rgb_data[i * 3 + 2] = img_data[i * 4 + 2]; - } - } else if (channels == 2) { - for (int i = 0; i < width * height; ++i) { - rgb_data[i * 3 + 0] = img_data[i * 2 + 0]; - rgb_data[i * 3 + 1] = img_data[i * 2 + 0]; - rgb_data[i * 3 + 2] = img_data[i * 2 + 0]; - } - } else { - throw std::runtime_error("Unsupported number of channels: " + std::to_string(channels)); - } - - return rgb_data; -} - -std::vector Siglip2Preprocessor::resize_image( - const unsigned char* img_data, int src_width, int src_height, - int dst_width, int dst_height, int channels) { - - const size_t src_elements = static_cast(src_width) * src_height * channels; - std::vector src_float(src_elements); - for (size_t idx = 0; idx < src_elements; ++idx) { - src_float[idx] = static_cast(img_data[idx]); - } - - - std::vector resized_data(static_cast(dst_width) * dst_height * channels); - - stbir_pixel_layout layout = (channels == 1) ? STBIR_1CHANNEL : - (channels == 3) ? STBIR_RGB : STBIR_RGBA; - - float* result = stbir_resize_float_linear( - src_float.data(), src_width, src_height, 0, - resized_data.data(), dst_width, dst_height, 0, - layout - ); - - if (!result) { - throw std::runtime_error("Failed to resize image"); - } - - - return resized_data; -} - -std::vector Siglip2Preprocessor::normalize_image( - const float* img_data, int width, int height, int channels) { - - size_t total_pixels = width * height * channels; - std::vector normalized(total_pixels); - - for (size_t i = 0; i < static_cast(width * height); ++i) { - for (int c = 0; c < channels; ++c) { - size_t idx = i * channels + c; - float pixel = img_data[idx]; - - if (config_.do_rescale) { - pixel *= config_.rescale_factor; - } - - if (config_.do_normalize) { - pixel = (pixel - config_.image_mean[c]) / config_.image_std[c]; - } - - normalized[idx] = pixel; - } - } - - - return normalized; -} - -std::vector> Siglip2Preprocessor::convert_image_to_patches( - const std::vector& image, int width, int height, int channels, int patch_size) { - - int num_patches_height = height / patch_size; - int num_patches_width = width / patch_size; - int num_patches = num_patches_height * num_patches_width; - int patch_elements = patch_size * patch_size * channels; - - std::vector> patches(num_patches, std::vector(patch_elements)); - - for (int ph = 0; ph < num_patches_height; ++ph) { - for (int pw = 0; pw < num_patches_width; ++pw) { - int patch_idx = ph * num_patches_width + pw; - - for (int y = 0; y < patch_size; ++y) { - for (int x = 0; x < patch_size; ++x) { - int img_y = ph * patch_size + y; - int img_x = pw * patch_size + x; - int img_idx = (img_y * width + img_x) * channels; - int patch_offset = (y * patch_size + x) * channels; - - for (int c = 0; c < channels; ++c) { - patches[patch_idx][patch_offset + c] = image[img_idx + c]; - } - } - } - } - } - - return patches; -} - -Siglip2Preprocessor::PreprocessedImage Siglip2Preprocessor::pad_patches( - const std::vector>& tile_patches, - const std::vector>& spatial_shapes, - int patch_dim, - int max_patches_per_tile) { - - if (tile_patches.size() != spatial_shapes.size()) { - throw std::runtime_error("Mismatch between tile data and spatial shapes"); - } - - PreprocessedImage result; - - const int num_tiles = static_cast(tile_patches.size()); - result.num_tiles = num_tiles; - result.patch_dim = patch_dim; - result.max_patches_per_tile = max_patches_per_tile; - result.spatial_shapes = spatial_shapes; - result.actual_num_patches = 0; - - const size_t total_values = static_cast(num_tiles) * max_patches_per_tile * patch_dim; - result.pixel_values.assign(total_values, 0.0f); - result.pixel_attention_mask.assign(static_cast(num_tiles) * max_patches_per_tile, 0); - - - for (int tile_idx = 0; tile_idx < num_tiles; ++tile_idx) { - const auto& [patches_h, patches_w] = spatial_shapes[tile_idx]; - const int actual_patches = patches_h * patches_w; - - if (actual_patches > max_patches_per_tile) { - throw std::runtime_error("Actual patches exceed max_patches_per_tile"); - } - - const auto& flattened = tile_patches[tile_idx]; - const size_t expected_size = static_cast(actual_patches) * patch_dim; - if (flattened.size() != expected_size) { - throw std::runtime_error("Tile patch data has unexpected size"); - } - - float* destination = result.pixel_values.data() + - static_cast(tile_idx) * max_patches_per_tile * patch_dim; - std::memcpy(destination, flattened.data(), expected_size * sizeof(float)); - - - int mask_offset = tile_idx * max_patches_per_tile; - for (int p = 0; p < actual_patches; ++p) { - result.pixel_attention_mask[mask_offset + p] = 1; - } - - result.actual_num_patches += actual_patches; - } - - if (!spatial_shapes.empty()) { - result.num_patches_height = spatial_shapes.front().first; - result.num_patches_width = spatial_shapes.front().second; - } else { - result.num_patches_height = 0; - result.num_patches_width = 0; - } - - result.pixel_values_shape = { - static_cast(num_tiles), - static_cast(max_patches_per_tile), - static_cast(patch_dim) - }; - - - result.pixel_attention_mask_shape = { - static_cast(num_tiles), - static_cast(max_patches_per_tile) - }; - - result.spatial_shapes_shape = { - static_cast(num_tiles), - static_cast(2) - }; - - - return result; -} - -} // namespace engine -} // namespace cactus - - diff --git a/cactus/engine/engine_model.cpp b/cactus/engine/engine_model.cpp deleted file mode 100644 index be7033bfe..000000000 --- a/cactus/engine/engine_model.cpp +++ /dev/null @@ -1,844 +0,0 @@ -#include "engine.h" -#include "../models/model.h" -#include "../graph/graph.h" -#include "../npu/npu.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace cactus { -namespace engine { - - -Model::Model() - : graph_handle_(nullptr), - config_(), - tokenizer_(nullptr), - initialized_(false), - attention_scale_(0.0f), - output_weight_node_id_(0), - owns_graph_(false) { -} - -Model::Model(const Config& config) - : graph_handle_(nullptr), - config_(config), - tokenizer_(nullptr), - initialized_(false), - attention_scale_(0.0f), - output_weight_node_id_(0), - owns_graph_(false) { -} - -Model::~Model() { - if (graph_handle_ && owns_graph_) { - delete static_cast(graph_handle_); - } -} - -bool Model::init(const std::string& model_folder, size_t context_size, const std::string& system_prompt, bool do_warmup) { - if (initialized_) { - return true; - } - auto* gb = new CactusGraph(); - graph_handle_ = gb; - owns_graph_ = true; - embedding_file_path_ = model_folder + "/token_embeddings.weights"; - return init_internal(gb, model_folder, context_size, system_prompt, do_warmup); -} - -bool Model::init(CactusGraph* external_graph, const std::string& model_folder, size_t context_size, - const std::string& system_prompt, bool do_warmup) { - if (!external_graph) { - throw std::invalid_argument("External graph pointer must not be null"); - } - if (initialized_) { - graph_handle_ = external_graph; - owns_graph_ = false; - return true; - } - - owns_graph_ = false; - graph_handle_ = external_graph; - return init_internal(external_graph, model_folder, context_size, system_prompt, do_warmup); -} - -bool Model::init_internal(CactusGraph* gb, const std::string& model_folder, size_t context_size, - const std::string& system_prompt, bool do_warmup) { - - CACTUS_LOG_DEBUG("model", "Initializing model from: " << model_folder); - model_folder_path_ = model_folder; - std::string config_path = model_folder + "/config.txt"; - - if (!config_.from_json(config_path)) { - CACTUS_LOG_ERROR("model", "Model initialization failed - config not loaded from: " << model_folder); - return false; - } - - std::string vocab_file = model_folder + "/vocab.txt"; - std::string merges_file = model_folder + "/merges.txt"; - std::string tokenizer_config_file = model_folder + "/tokenizer_config.txt"; - - std::ifstream merges_check(merges_file); - bool has_merges = false; - if (merges_check.is_open()) { - std::string line; - int line_count = 0; - while (std::getline(merges_check, line) && line_count < 10) { - if (!line.empty() && line[0] != '#') { - has_merges = true; - break; - } - line_count++; - } - merges_check.close(); - } - - if (has_merges) { - tokenizer_ = std::make_unique(); - } else { - tokenizer_ = std::make_unique(); - } - - if (!tokenizer_->load_vocabulary_with_config(vocab_file, merges_file, tokenizer_config_file)) { - return false; - } - - graph_handle_ = gb; - - if(config_.model_type == Config::ModelType::WHISPER){ - embedding_file_path_ = model_folder+"/decoder_token_embeddings.weights"; - } - else{ - embedding_file_path_ = model_folder + "/token_embeddings.weights"; - } - - load_weights_to_graph(gb); - - if (config_.model_type == Config::ModelType::GEMMA) { - attention_scale_ = 1.0f / std::sqrt(256.0f); - } else { - attention_scale_ = 1.0f / std::sqrt(static_cast(config_.attention_head_dim)); - } - - Precision cache_precision = (config_.model_type == Config::ModelType::WHISPER) - ? Precision::FP16 - : Precision::INT8; - kv_cache_.init(config_.num_layers, context_size, config_.attention_kv_heads, config_.attention_head_dim, cache_precision); - - size_t window_size = std::min(context_size, size_t(512)); - size_t sink_size = 4; - const char* env_window = std::getenv("CACTUS_KV_WINDOW_SIZE"); - const char* env_sink = std::getenv("CACTUS_KV_SINK_SIZE"); - if (env_window) { - window_size = std::stoul(env_window); - } - if (env_sink) { - sink_size = std::stoul(env_sink); - } - kv_cache_.set_window_size(window_size, sink_size); - cache_k_output_nodes_.resize(config_.num_layers); - cache_v_output_nodes_.resize(config_.num_layers); - - post_init(); - - initialized_ = true; - - if (do_warmup && config_.model_type != Config::ModelType::WHISPER) { - std::string warmup_text = system_prompt.empty() ? "Hello" : system_prompt; - auto warmup_tokens = tokenizer_->encode(warmup_text); - forward(warmup_tokens); - auto* gb = static_cast(graph_handle_); - gb->execute(); - } - - reset_cache(); - return true; -} - -size_t Model::forward(const std::vector& /*mel_bins*/, const std::vector& tokens, bool use_cache){ - return forward(tokens, use_cache); -} - -void Model::prefill(const std::vector& tokens, size_t chunk_size, const std::string& profile_file) { - if (tokens.empty()) { - return; - } - - if (has_npu_prefill()) { - size_t npu_chunk_size = static_cast(npu_prefill_->get_chunk_size()); - if (tokens.size() > npu_chunk_size) { - prefill_npu(tokens); - return; - } - } - - auto* gb = static_cast(graph_handle_); - - auto process_chunk = [&](const std::vector& chunk) { - forward(chunk, true); - - if (!profile_file.empty()) { - gb->execute(profile_file); - } else { - gb->execute(profile_file); - } - - post_execute_updates(gb, chunk.size()); - update_kv_cache(gb, chunk.size()); - }; - - if (tokens.size() <= chunk_size) { - process_chunk(tokens); - return; - } - - size_t num_full_chunks = (tokens.size() - 1) / chunk_size; - - for (size_t chunk_idx = 0; chunk_idx < num_full_chunks; ++chunk_idx) { - size_t start = chunk_idx * chunk_size; - size_t end = start + chunk_size; - std::vector chunk(tokens.begin() + start, tokens.begin() + end); - if (chunk_idx == 1) { - gb->set_prefill_mode(true); - } - process_chunk(chunk); - } - - gb->set_prefill_mode(false); - size_t final_start = num_full_chunks * chunk_size; - std::vector final_chunk(tokens.begin() + final_start, tokens.end()); - process_chunk(final_chunk); -} - -uint32_t Model::decode(const std::vector& tokens, float temperature, float top_p, - size_t top_k, const std::string& profile_file, float* out_entropy) { - - if (temperature < 0) { - temperature = config_.default_temperature; - } - if (top_p < 0) { - top_p = config_.default_top_p; - } - if (top_k == 0) { - top_k = config_.default_top_k; - } - - auto final_hidden = forward(tokens, true); - - auto* gb = static_cast(graph_handle_); - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - auto last_hidden = gb->index(final_hidden, tokens.size() - 1, 0); - const auto& last_hidden_buf = gb->get_output_buffer(last_hidden); - size_t hidden_dim = last_hidden_buf.shape[0]; - last_hidden = gb->reshape(last_hidden, {1, hidden_dim}); - - auto logits_node_id = gb->matmul(last_hidden, output_weight_node_id_, true, backend); - auto sampled_token_id = gb->sample(logits_node_id, temperature, top_p, top_k, tool_constrainer_.get_bias()); - - if (!profile_file.empty()) { - gb->execute(profile_file); - } else { - gb->execute(); - } - - if (out_entropy) { - const auto& logits_buf = gb->get_output_buffer(logits_node_id); - void* logits_ptr = gb->get_output(logits_node_id); - size_t vocab_size = logits_buf.shape.back(); - - std::vector logits(vocab_size); - if (logits_buf.precision == Precision::FP32) { - float* src = static_cast(logits_ptr); - std::copy(src, src + vocab_size, logits.begin()); - } else if (logits_buf.precision == Precision::FP16) { - __fp16* src = static_cast<__fp16*>(logits_ptr); - Quantization::fp16_to_fp32(src, logits.data(), vocab_size); - } else { - int8_t* src = static_cast(logits_ptr); - Quantization::int8_to_fp32(src, logits.data(), vocab_size, 1.0f); - } - - float max_logit = *std::max_element(logits.begin(), logits.end()); - double sum_exp = 0.0; - for (size_t i = 0; i < vocab_size; ++i) { - sum_exp += std::exp(static_cast(logits[i] - max_logit)); - } - double log_sum_exp = static_cast(max_logit) + std::log(sum_exp); - - double entropy = 0.0; - for (size_t i = 0; i < vocab_size; ++i) { - double log_prob = static_cast(logits[i]) - log_sum_exp; - double prob = std::exp(log_prob); - if (prob > 1e-10) { - entropy -= prob * log_prob; - } - } - - double max_entropy = std::log(static_cast(vocab_size)); - *out_entropy = static_cast(entropy / max_entropy); - } - - post_execute_updates(gb, tokens.size()); - update_kv_cache(gb, tokens.size()); - - auto* output_ptr = gb->get_output(sampled_token_id); - return *static_cast(output_ptr); -} - -uint32_t Model::decode_with_audio(const std::vector& tokens, const std::vector& /*mel_bins*/, float temperature, float top_p, size_t top_k, const std::string& profile_file, float* out_entropy){ - return decode(tokens, temperature, top_p, top_k, profile_file, out_entropy); -} - -uint32_t Model::decode_with_images(const std::vector& tokens, const std::vector& image_paths, - float temperature, float top_p, size_t top_k, const std::string& profile_file, float* out_entropy) { - (void)image_paths; - return decode(tokens, temperature, top_p, top_k, profile_file, out_entropy); -} - -std::vector Model::get_image_embeddings(const std::string& /*image_path*/) { - throw std::runtime_error("Image embeddings not supported for this model type"); -} - -std::vector Model::get_audio_embeddings(const std::vector& /*mel_bins*/) { - throw std::runtime_error("Audio embeddings not supported for this model type"); -} - -void Model::update_kv_cache(CactusGraph* gb, size_t seq_len) { - kv_cache_.update_from_graph(gb, cache_k_output_nodes_, cache_v_output_nodes_, - seq_len, config_.num_layers, config_.attention_kv_heads, - config_.attention_head_dim); -} - - -std::vector Model::get_embeddings(const std::vector& tokens, bool pooled, bool normalize, const std::string& profile_file) { - std::vector embeddings; - auto final_hidden = forward(tokens); - - auto* gb = static_cast(graph_handle_); - auto* output_ptr = gb->get_output(final_hidden); - const auto& output_buffer = gb->get_output_buffer(final_hidden); - - if (pooled) { - auto pooled_hidden = gb->mean(final_hidden, 0); - - if (!profile_file.empty()) { - gb->execute(profile_file); - } else { - gb->execute(); - } - post_execute_updates(gb, tokens.size()); - auto* pooled_ptr = gb->get_output(pooled_hidden); - const auto& pooled_buffer = gb->get_output_buffer(pooled_hidden); - - size_t hidden_dim = pooled_buffer.total_size; - embeddings.resize(hidden_dim); - - if (pooled_buffer.precision == Precision::FP32) { - float* pooled_data = static_cast(pooled_ptr); - std::copy(pooled_data, pooled_data + hidden_dim, embeddings.begin()); - } else if (pooled_buffer.precision == Precision::FP16) { - __fp16* pooled_data = static_cast<__fp16*>(pooled_ptr); - Quantization::fp16_to_fp32(pooled_data, embeddings.data(), hidden_dim); - } else if (pooled_buffer.precision == Precision::INT8) { - int8_t* pooled_data = static_cast(pooled_ptr); - Quantization::int8_to_fp32(pooled_data, embeddings.data(), hidden_dim, 1.0f); - } - } else { - if (!profile_file.empty()) { - gb->execute(profile_file); - } else { - gb->execute(); - } - post_execute_updates(gb, tokens.size()); - - size_t total_size = output_buffer.total_size; - embeddings.resize(total_size); - - if (output_buffer.precision == Precision::FP32) { - float* hidden_states = static_cast(output_ptr); - std::copy(hidden_states, hidden_states + total_size, embeddings.begin()); - } else if (output_buffer.precision == Precision::FP16) { - __fp16* hidden_states = static_cast<__fp16*>(output_ptr); - for (size_t i = 0; i < total_size; i++) { - embeddings[i] = static_cast(hidden_states[i]); - } - } else if (output_buffer.precision == Precision::INT8) { - int8_t* hidden_states = static_cast(output_ptr); - for (size_t i = 0; i < total_size; i++) { - embeddings[i] = static_cast(hidden_states[i]); - } - } - } - - if (normalize && !embeddings.empty()) { - float norm_sq = 0.0f; - for (float v : embeddings) { - norm_sq += v * v; - } - float norm = std::sqrt(norm_sq); - if (norm > 1e-12f) { - float inv_norm = 1.0f / norm; - for (float& v : embeddings) { - v *= inv_norm; - } - } - } - - kv_cache_.reset(); - - return embeddings; -} - -bool Config::from_json(const std::string& config_path) { - std::ifstream file(config_path); - if (!file) { - CACTUS_LOG_ERROR("config", "Failed to open config file: " << config_path); - return false; - } - - std::string line; - while (std::getline(file, line)) { - if (line.empty() || line[0] == '#') continue; - - size_t eq_pos = line.find('='); - if (eq_pos == std::string::npos) continue; - - std::string key = line.substr(0, eq_pos); - std::string value = line.substr(eq_pos + 1); - - key.erase(0, key.find_first_not_of(" \t")); - key.erase(key.find_last_not_of(" \t") + 1); - value.erase(0, value.find_first_not_of(" \t")); - value.erase(value.find_last_not_of(" \t") + 1); - - if (key == "vocab_size") vocab_size = static_cast(std::stoul(value)); - else if (key == "bos_token_id") bos_token_id = static_cast(std::stoul(value)); - else if (key == "eos_token_id") eos_token_id = static_cast(std::stoul(value)); - else if (key == "num_layers") num_layers = static_cast(std::stoul(value)); - else if (key == "hidden_dim") hidden_dim = static_cast(std::stoul(value)); - else if (key == "ffn_intermediate_dim") ffn_intermediate_dim = static_cast(std::stoul(value)); - else if (key == "attention_heads") attention_heads = static_cast(std::stoul(value)); - else if (key == "attention_kv_heads") attention_kv_heads = static_cast(std::stoul(value)); - else if (key == "attention_head_dim") attention_head_dim = static_cast(std::stoul(value)); - else if (key == "layer_norm_eps") layer_norm_eps = std::stof(value); - else if (key == "rope_theta") rope_theta = std::stof(value); - else if (key == "num_experts") num_experts = static_cast(std::stoul(value)); - else if (key == "num_shared_experts") num_shared_experts = static_cast(std::stoul(value)); - else if (key == "num_top_experts") num_top_experts = static_cast(std::stoul(value)); - else if (key == "moe_every_n_layers") moe_every_n_layers = static_cast(std::stoul(value)); - else if (key == "tie_word_embeddings") tie_word_embeddings = (value == "true" || value == "1"); - else if (key == "vision_hidden_dim") vision_hidden_dim = static_cast(std::stoul(value)); - else if (key == "vision_num_layers") vision_num_layers = static_cast(std::stoul(value)); - else if (key == "vision_attention_heads") vision_attention_heads = static_cast(std::stoul(value)); - else if (key == "vision_image_size") vision_image_size = static_cast(std::stoul(value)); - else if (key == "vision_patch_size") vision_patch_size = static_cast(std::stoul(value)); - else if (key == "vision_num_channels") vision_num_channels = static_cast(std::stoul(value)); - else if (key == "vision_embed_dim") vision_embed_dim = static_cast(std::stoul(value)); - else if (key == "visual_tokens_per_img") visual_tokens_per_img = static_cast(std::stoul(value)); - else if (key == "use_pixel_shuffle") use_pixel_shuffle = (value == "true" || value == "1"); - else if (key == "pixel_shuffle_factor") pixel_shuffle_factor = static_cast(std::stoul(value)); - else if (key == "use_image_tokens") use_image_tokens = (value == "true" || value == "1"); - else if (key == "use_layout_tags") use_layout_tags = (value == "true" || value == "1"); - else if (key == "image_seq_len") image_seq_len = static_cast(std::stoul(value)); - else if (key == "global_image_size") global_image_size = static_cast(std::stoul(value)); - else if (key == "max_tile_size") max_tile_size = static_cast(std::stoul(value)); - else if (key == "rescale_factor") rescale_factor = std::stof(value); - else if (key == "image_mean") image_mean = std::stof(value); - else if (key == "image_std") image_std = std::stof(value); - else if (key == "downsample_factor") downsample_factor = static_cast(std::stoul(value)); - else if (key == "min_tiles") min_tiles = static_cast(std::stoul(value)); - else if (key == "max_tiles") max_tiles = static_cast(std::stoul(value)); - else if (key == "use_thumbnail") use_thumbnail = (value == "true" || value == "1"); - else if (key == "min_image_tokens") min_image_tokens = static_cast(std::stoul(value)); - else if (key == "max_image_tokens") max_image_tokens = static_cast(std::stoul(value)); - else if (key == "tile_size") tile_size = static_cast(std::stoul(value)); - else if (key == "max_pixels_tolerance") max_pixels_tolerance = std::stof(value); - else if (key == "do_image_splitting") do_image_splitting = (value == "true" || value == "1"); - else if (key == "precision") { - if (value == "INT8") precision = Precision::INT8; - else if (value == "FP16") precision = Precision::FP16; - else precision = Precision::FP32; - } - else if (key == "model_type") { - if (value == "gemma" || value == "GEMMA") model_type = ModelType::GEMMA; - else if (value == "lfm2" || value == "LFM2") model_type = ModelType::LFM2; - else if (value == "smol" || value == "SMOL" || value == "Smol") model_type = ModelType::SMOL; - else if (value == "bert" || value == "BERT") model_type = ModelType::NOMIC; - else if (value == "whisper" || value == "WHISPER") model_type = ModelType::WHISPER; - else model_type = ModelType::QWEN; - } - else if (key == "model_variant") { - std::string v = value; - std::transform(v.begin(), v.end(), v.begin(), ::tolower); - if (v == "vlm") model_variant = ModelVariant::VLM; - else if (v == "extract") model_variant = ModelVariant::EXTRACT; - else if (v == "rag") model_variant = ModelVariant::RAG; - else model_variant = ModelVariant::DEFAULT; - } - else if (key == "conv_L_cache") conv_L_cache = static_cast(std::stoul(value)); - else if (key == "layer_types") { - layer_types.clear(); - std::string sanitized; - sanitized.reserve(value.size()); - for (char c : value) { - if (c == '[' || c == ']' || c == '\'' || c == '"') { - continue; - } - sanitized.push_back(c); - } - std::stringstream ss(sanitized); - std::string item; - while (std::getline(ss, item, ',')) { - if (!item.empty()) { - item.erase(0, item.find_first_not_of(" \t")); - item.erase(item.find_last_not_of(" \t") + 1); - if (!item.empty()) layer_types.push_back(item); - } - } - } - } - - if (model_type == ModelType::GEMMA) { - default_temperature = 1.0f; - default_top_p = 0.95f; - default_top_k = 64; - } else if (model_type == ModelType::SMOL) { - default_temperature = 0.2f; - default_top_p = 0.95f; - default_top_k = 20; - } else if (model_type == ModelType::LFM2) { - default_temperature = 0.3f; - default_top_p = 0.95f; - default_top_k = 20; - } else if (model_type == ModelType::QWEN) { - default_temperature = 0.6f; - default_top_p = 0.95f; - default_top_k = 20; - } else if (model_type == ModelType::QWEN) { - default_temperature = 0.7f; - default_top_p = 0.8f; - default_top_k = 20; - } else if (model_type == ModelType::WHISPER) { - default_temperature = 0.0f; - default_top_p = 0.0f; - default_top_k = 0; - } - - return true; -} - -std::string Config::to_json() const { - return "{}"; -} - -std::unique_ptr create_model(const std::string& model_folder) { - CACTUS_LOG_DEBUG("model", "Creating model from: " << model_folder); - Config config; - std::string config_path = model_folder + "/config.txt"; - - if (!config.from_json(config_path)) { - CACTUS_LOG_ERROR("model", "Failed to create model - cannot load config from: " << model_folder); - return nullptr; - } - - const bool has_vision_support = - config.use_image_tokens || - config.vision_num_layers > 0 || - config.vision_embed_dim > 0 || - config.vision_hidden_dim > 0 || - config.visual_tokens_per_img > 0; - - if (config.model_type == Config::ModelType::LFM2 && has_vision_support) { - return std::make_unique(config); - } - - switch (config.model_type) { - case Config::ModelType::QWEN: - return std::make_unique(config); - case Config::ModelType::GEMMA: - return std::make_unique(config); - case Config::ModelType::LFM2: - return std::make_unique(config); - case Config::ModelType::SMOL: - return std::make_unique(config); - case Config::ModelType::NOMIC: - return std::make_unique(config); - case Config::ModelType::WHISPER: - return std::make_unique(config); - default: - return std::make_unique(config); - } -} - -void Model::capture_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id) const { - auto* graph = static_cast(graph_handle_); - if (!graph) { - return; - } - graph->capture_debug_node(layer_idx, name, node_id); -} - -void Model::clear_debug_nodes() { - auto* graph = static_cast(graph_handle_); - if (!graph) { - return; - } - graph->clear_debug_nodes(); -} - -const std::vector& Model::get_debug_nodes() const { - auto* graph = static_cast(graph_handle_); - debug_nodes_.clear(); - if (!graph) { - return debug_nodes_; - } - - const auto& entries = graph->get_debug_nodes(); - debug_nodes_.reserve(entries.size()); - for (const auto& entry : entries) { - debug_nodes_.push_back({entry.layer_idx, entry.name, entry.node_id}); - } - return debug_nodes_; -} - -bool Model::load_npu_prefill(const std::string& model_path) { - CACTUS_LOG_DEBUG("npu", "Attempting to load NPU prefill from: " << model_path); - - npu_prefill_ = npu::create_prefill(); - if (!npu_prefill_) { - CACTUS_LOG_DEBUG("npu", "NPU prefill creation failed (not supported on this device)"); - return false; - } - - bool loaded = npu_prefill_->load(model_path); - if (loaded) { - CACTUS_LOG_INFO("npu", "NPU prefill loaded successfully from: " << model_path); - } else { - CACTUS_LOG_DEBUG("npu", "NPU prefill model not found at: " << model_path); - } - return loaded; -} - -bool Model::has_npu_prefill() const { - return npu_prefill_ && npu_prefill_->is_available(); -} - -size_t Model::get_prefill_chunk_size() const { - if (has_npu_prefill()) { - return static_cast(npu_prefill_->get_chunk_size()); - } - return 256; // default chunk size -} - -std::vector<__fp16> Model::get_token_embeddings(const std::vector& tokens) { - auto* gb = static_cast(graph_handle_); - if (!gb || tokens.empty()) { - return {}; - } - - gb->soft_reset(); - - size_t tok_input = gb->input({tokens.size()}, Precision::FP32); - std::vector tok_f(tokens.size()); - for (size_t i = 0; i < tokens.size(); i++) { - tok_f[i] = static_cast(tokens[i]); - } - gb->set_input(tok_input, tok_f.data(), Precision::FP32); - - size_t embedding_node = gb->embedding(embedding_node_id_, tok_input); - - gb->execute(); - - const auto& emb_buf = gb->get_output_buffer(embedding_node); - void* emb_ptr = gb->get_output(embedding_node); - - size_t num_tokens = tokens.size(); - size_t hidden_dim = config_.hidden_dim; - std::vector<__fp16> embeddings(num_tokens * hidden_dim); - - if (emb_buf.precision == Precision::FP16) { - __fp16* src = static_cast<__fp16*>(emb_ptr); - std::copy(src, src + num_tokens * hidden_dim, embeddings.begin()); - } else if (emb_buf.precision == Precision::FP32) { - float* src = static_cast(emb_ptr); - for (size_t i = 0; i < num_tokens * hidden_dim; i++) { - embeddings[i] = static_cast<__fp16>(src[i]); - } - } else if (emb_buf.precision == Precision::INT8) { - int8_t* src = static_cast(emb_ptr); - for (size_t i = 0; i < num_tokens * hidden_dim; i++) { - embeddings[i] = static_cast<__fp16>(src[i]); - } - } - - return embeddings; -} - -void Model::prefill_npu(const std::vector& tokens) { - if (!npu_prefill_ || !npu_prefill_->is_available()) { - throw std::runtime_error("NPU prefill not available"); - } - - const int chunk_size = npu_prefill_->get_chunk_size(); - const int hidden_dim = npu_prefill_->get_hidden_dim(); - const int num_layers = npu_prefill_->get_num_layers(); - const int num_kv_heads = npu_prefill_->get_num_kv_heads(); - const int head_dim = npu_prefill_->get_head_dim(); - - std::vector<__fp16> all_embeddings = get_token_embeddings(tokens); - if (all_embeddings.empty()) { - throw std::runtime_error("Failed to get token embeddings for NPU prefill"); - } - - if (config_.model_type == Config::ModelType::GEMMA) { - float scale = std::sqrt(static_cast(hidden_dim)); - for (size_t i = 0; i < all_embeddings.size(); i++) { - all_embeddings[i] = __fp16(static_cast(all_embeddings[i]) * scale); - } - } - - size_t num_tokens = tokens.size(); - size_t num_chunks = (num_tokens + chunk_size - 1) / chunk_size; - - for (size_t c = 0; c < num_chunks; c++) { - size_t start = c * chunk_size; - size_t actual_tokens = std::min(static_cast(chunk_size), num_tokens - start); - - std::vector<__fp16> chunk_embeddings(chunk_size * hidden_dim, __fp16(0)); - std::copy(all_embeddings.begin() + start * hidden_dim, - all_embeddings.begin() + (start + actual_tokens) * hidden_dim, - chunk_embeddings.begin()); - - int position_offset = static_cast(start); - - npu::NPUPrefillDirectResult direct_result = npu_prefill_->prefill_chunk_direct(chunk_embeddings, position_offset); - - if (direct_result.valid) { - for (int layer_idx = 0; layer_idx < num_layers; layer_idx++) { - const auto& k_ref = direct_result.k_caches[layer_idx]; - const auto& v_ref = direct_result.v_caches[layer_idx]; - - if (k_ref.data && v_ref.data) { - kv_cache_.update_from_npu(layer_idx, k_ref.data, v_ref.data, - actual_tokens, num_kv_heads, head_dim); - } - } - } - } -} - -double Model::score_tokens_window_logprob( - const std::vector& tokens, - size_t start, - size_t end, - size_t context, - size_t* tokens_scored -) { - if (tokens_scored) - *tokens_scored = 0; - - if (tokens.empty()) - return 0.0; - - if (end > tokens.size()) - end = tokens.size(); - - if (start >= end) - return 0.0; - - if (start == 0) - start = 1; - - if (start >= end) - return 0.0; - - const size_t target_len = end - start; - const size_t ctx_begin = (start > context) ? (start - context) : 0; - - if (end < 2) return 0.0; - const size_t input_end = end - 1; - - if (input_end <= ctx_begin) - return 0.0; - - std::vector input_tokens(tokens.begin() + ctx_begin,tokens.begin() + input_end); - - if (tokens_scored) - *tokens_scored = target_len; - - reset_cache(); - - auto* gb = static_cast(graph_handle_); - const auto backend = (config_.default_backend == Config::Backend::CPU) ? ComputeBackend::CPU : ComputeBackend::NPU; - - const size_t hidden_node = forward(input_tokens, /*use_cache=*/false); - const auto& hidden_buf = gb->get_output_buffer(hidden_node); - - if (hidden_buf.shape.size() != 2) { - throw std::runtime_error("Expected hidden to be rank-2 [L, hidden_dim]"); - } - - - const size_t first_pos = start - ctx_begin - 1; - const size_t hidden_slice = gb->slice(hidden_node, /*axis=*/0, first_pos, target_len); - bool transpose_w = true; - const size_t logits_node = gb->matmul(hidden_slice, output_weight_node_id_, transpose_w, backend); - gb->execute(); - - const auto& logits_buf = gb->get_output_buffer(logits_node); - if (logits_buf.shape.size() != 2) - throw std::runtime_error("Expected logits to be rank-2 [T, vocab]"); - - const size_t T = logits_buf.shape[0]; - const size_t vocab_size = logits_buf.shape[1]; - - if (T != target_len) - throw std::runtime_error("Logits T dimension does not match target_len"); - - void* logits_ptr = gb->get_output(logits_node); - std::vector row(vocab_size); - double total_logprob = 0.0; - - for (size_t i = 0; i < target_len; ++i) { - const uint32_t y = tokens[start + i]; - if (y >= vocab_size) - throw std::runtime_error("Target token out of vocab range"); - - if (logits_buf.precision == Precision::FP32) { - const float* src = static_cast(logits_ptr) + i * vocab_size; - std::memcpy(row.data(), src, vocab_size * sizeof(float)); - } - else if (logits_buf.precision == Precision::FP16) { - const __fp16* src = static_cast(logits_ptr) + i * vocab_size; - Quantization::fp16_to_fp32(const_cast<__fp16*>(src), row.data(), vocab_size); - } - else { - const int8_t* src = static_cast(logits_ptr) + i * vocab_size; - Quantization::int8_to_fp32(const_cast(src), row.data(), vocab_size, 1.0f); - } - - float max_logit = *std::max_element(row.begin(), row.end()); - double sum = 0.0; - - for (size_t j = 0; j < vocab_size; ++j) - sum += std::exp(double(row[j] - max_logit)); - - const double lse = double(max_logit) + std::log(sum); - total_logprob += double(row[y]) - lse; - } - - return total_logprob; -} -} -} \ No newline at end of file diff --git a/cactus/engine/engine_signals.cpp b/cactus/engine/engine_signals.cpp deleted file mode 100644 index b334dbe7a..000000000 --- a/cactus/engine/engine_signals.cpp +++ /dev/null @@ -1,555 +0,0 @@ -#include "engine.h" -#include "../kernel/kernel_utils.h" -#include -#include -#include -#include -#include -#include - -#ifndef M_PI -#define M_PI 3.14159265358979323846 -#endif - -namespace cactus { -namespace engine { - -static void to_db( - float* spectrogram, - size_t size, - float reference, - float min_value, - const float* db_range, - float multiplier) -{ - if (reference <= 0.0f) { - throw std::invalid_argument("reference must be greater than zero"); - } - if (min_value <= 0.0f) { - throw std::invalid_argument("min_value must be greater than zero"); - } - - reference = std::max(min_value, reference); - const float log_ref = std::log10(reference); - - CactusThreading::parallel_for(size, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - float value = std::max(min_value, spectrogram[i]); - spectrogram[i] = multiplier * (std::log10(value) - log_ref); - } - }); - - if (db_range != nullptr) { - if (*db_range <= 0.0f) { - throw std::invalid_argument("db_range must be greater than zero"); - } - - float max_db = CactusThreading::parallel_reduce, float, std::function>( - size, CactusThreading::Thresholds::ALL_REDUCE, - [&](size_t start, size_t end) { - float local_max = -std::numeric_limits::infinity(); - for (size_t i = start; i < end; i++) { - local_max = std::max(local_max, spectrogram[i]); - } - return local_max; - }, - -std::numeric_limits::infinity(), - [](float a, float b) { return std::max(a, b); } - ); - - float min_db = max_db - *db_range; - CactusThreading::parallel_for(size, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - spectrogram[i] = std::max(min_db, spectrogram[i]); - } - }); - } -} - -static size_t bit_reverse(size_t x, size_t log2n) { - size_t result = 0; - for (size_t i = 0; i < log2n; i++) { - result = (result << 1) | (x & 1); - x >>= 1; - } - return result; -} - -static void fft_radix2(float* re, float* im, size_t n) { - if (n == 0 || (n & (n - 1)) != 0) return; - - size_t log2n = 0; - for (size_t temp = n; temp > 1; temp >>= 1) log2n++; - - for (size_t i = 0; i < n; i++) { - size_t j = bit_reverse(i, log2n); - if (i < j) { - std::swap(re[i], re[j]); - std::swap(im[i], im[j]); - } - } - - for (size_t s = 1; s <= log2n; s++) { - size_t m = 1 << s; - size_t m2 = m >> 1; - float w_re = 1.0f; - float w_im = 0.0f; - float wm_re = std::cos(static_cast(M_PI) / static_cast(m2)); - float wm_im = -std::sin(static_cast(M_PI) / static_cast(m2)); - - for (size_t j = 0; j < m2; j++) { - for (size_t k = j; k < n; k += m) { - size_t k_m2 = k + m2; - float t_re = w_re * re[k_m2] - w_im * im[k_m2]; - float t_im = w_re * im[k_m2] + w_im * re[k_m2]; - float u_re = re[k]; - float u_im = im[k]; - re[k] = u_re + t_re; - im[k] = u_im + t_im; - re[k_m2] = u_re - t_re; - im[k_m2] = u_im - t_im; - } - float new_w_re = w_re * wm_re - w_im * wm_im; - float new_w_im = w_re * wm_im + w_im * wm_re; - w_re = new_w_re; - w_im = new_w_im; - } - } -} - -static void rfft_f32_1d(const float* input, float* output, const size_t n, const char* norm) { - const size_t out_len = n / 2 + 1; - - float norm_factor = 1.0f; - if (norm) { - if (std::strcmp(norm, "backward") == 0) { - norm_factor = 1.0f; - } else if (std::strcmp(norm, "forward") == 0) { - norm_factor = 1.0f / static_cast(n); - } else if (std::strcmp(norm, "ortho") == 0) { - norm_factor = 1.0f / std::sqrt(static_cast(n)); - } else { - throw std::invalid_argument("norm must be one of {\"backward\",\"forward\",\"ortho\"}"); - } - } - - if ((n & (n - 1)) == 0 && n >= 4) { - std::vector re(n), im(n, 0.0f); - std::copy(input, input + n, re.begin()); - - fft_radix2(re.data(), im.data(), n); - - for (size_t i = 0; i < out_len; i++) { - output[i * 2] = re[i] * norm_factor; - output[i * 2 + 1] = im[i] * norm_factor; - } - } else { - const float two_pi_over_n = 2.0f * static_cast(M_PI) / static_cast(n); - for (size_t i = 0; i < out_len; i++) { - float re = 0.0f; - float im = 0.0f; - const float base = -two_pi_over_n * static_cast(i); - for (size_t j = 0; j < n; j++) { - const float angle = base * static_cast(j); - const float input_val = input[j]; - re += input_val * std::cos(angle); - im += input_val * std::sin(angle); - } - output[i * 2] = re * norm_factor; - output[i * 2 + 1] = im * norm_factor; - } - } -} - -static float hertz_to_mel(float freq, const char* mel_scale) { - if (std::strcmp(mel_scale, "htk") == 0) { - return 2595.0f * std::log10(1.0f + (freq / 700.0f)); - } else if (std::strcmp(mel_scale, "kaldi") == 0) { - return 1127.0f * std::log(1.0f + (freq / 700.0f)); - } - - const float min_log_hertz = 1000.0f; - const float min_log_mel = 15.0f; - const float logstep = 27.0f / std::log(6.4f); - float mels = 3.0f * freq / 200.0f; - - if (freq >= min_log_hertz) { - mels = min_log_mel + std::log(freq / min_log_hertz) * logstep; - } - - return mels; -} - -static float mel_to_hertz(float mels, const char* mel_scale) { - if (std::strcmp(mel_scale, "htk") == 0) { - return 700.0f * (std::pow(10.0f, mels / 2595.0f) - 1.0f); - } else if (std::strcmp(mel_scale, "kaldi") == 0) { - return 700.0f * (std::exp(mels / 1127.0f) - 1.0f); - } - - const float min_log_hertz = 1000.0f; - const float min_log_mel = 15.0f; - const float logstep = std::log(6.4f) / 27.0f; - float freq = 200.0f * mels / 3.0f; - - if (mels >= min_log_mel) { - freq = min_log_hertz * std::exp(logstep * (mels - min_log_mel)); - } - - return freq; -} - -static void generate_mel_filter_bank( - float* mel_filters, - const int num_frequency_bins, - const int num_mel_filters, - const float min_frequency, - const float max_frequency, - const int sampling_rate, - const char* norm, - const char* mel_scale, - const bool triangularize_in_mel_space) -{ - if (norm != nullptr && std::strcmp(norm, "slaney") != 0) { - throw std::invalid_argument("norm must be one of None or \"slaney\""); - } - - if (std::strcmp(mel_scale, "htk") != 0 && std::strcmp(mel_scale, "kaldi") != 0 && std::strcmp(mel_scale, "slaney") != 0) { - throw std::invalid_argument("mel_scale should be one of \"htk\", \"slaney\" or \"kaldi\"."); - } - - if (num_frequency_bins < 2) { - throw std::invalid_argument( - "Require num_frequency_bins: " + std::to_string(num_frequency_bins) + " >= 2"); - } - - if (min_frequency > max_frequency) { - throw std::invalid_argument( - "Require min_frequency: " + std::to_string(min_frequency) + - " <= max_frequency: " + std::to_string(max_frequency)); - } - - const float mel_min = hertz_to_mel(min_frequency, mel_scale); - const float mel_max = hertz_to_mel(max_frequency, mel_scale); - - std::vector mel_freqs(num_mel_filters + 2); - for (int i = 0; i < num_mel_filters + 2; i++) { - mel_freqs[i] = mel_min + (mel_max - mel_min) * i / (num_mel_filters + 1); - } - - std::vector filter_freqs(num_mel_filters + 2); - for (int i = 0; i < num_mel_filters + 2; i++) { - filter_freqs[i] = mel_to_hertz(mel_freqs[i], mel_scale); - } - - std::vector fft_freqs(num_frequency_bins); - if (triangularize_in_mel_space) { - float fft_bin_width = static_cast(sampling_rate) / ((num_frequency_bins - 1) * 2); - for (int i = 0; i < num_frequency_bins; i++) { - fft_freqs[i] = hertz_to_mel(fft_bin_width * i, mel_scale); - } - filter_freqs = mel_freqs; - } else { - for (int i = 0; i < num_frequency_bins; i++) { - fft_freqs[i] = (static_cast(sampling_rate) / 2.0f) * i / (num_frequency_bins - 1); - } - } - - for (int i = 0; i < num_mel_filters; i++) { - float left_edge = filter_freqs[i]; - float center = filter_freqs[i + 1]; - float right_edge = filter_freqs[i + 2]; - - for (int j = 0; j < num_frequency_bins; j++) { - float freq = fft_freqs[j]; - float down_slope = (freq - left_edge) / (center - left_edge); - float up_slope = (right_edge - freq) / (right_edge - center); - - mel_filters[i * num_frequency_bins + j] = std::max(0.0f, std::min(down_slope, up_slope)); - } - } - - if (norm != nullptr && std::strcmp(norm, "slaney") == 0) { - for (int i = 0; i < num_mel_filters; i++) { - float enorm = 2.0f / (filter_freqs[i + 2] - filter_freqs[i]); - for (int j = 0; j < num_frequency_bins; j++) { - mel_filters[i * num_frequency_bins + j] *= enorm; - } - } - } -} - -static void compute_spectrogram_f32( - const float* waveform, - size_t waveform_length, - const float* window, - size_t window_length, - size_t frame_length, - size_t hop_length, - const size_t* fft_length, - float* spectrogram, - float power, - bool center, - const char* pad_mode, - bool onesided [[maybe_unused]], - float dither, - const float* preemphasis, - const float* mel_filters, - size_t mel_filters_size, - float mel_floor, - const char* log_mel, - float reference, - float min_value, - const float* db_range, - bool remove_dc_offset) -{ - size_t actual_fft_length; - if (fft_length == nullptr) { - actual_fft_length = frame_length; - } else { - actual_fft_length = *fft_length; - } - - if (frame_length > actual_fft_length) { - throw std::invalid_argument( - "frame_length (" + std::to_string(frame_length) + - ") may not be larger than fft_length (" + - std::to_string(actual_fft_length) + ")"); - } - - std::vector hann_window; - const float* actual_window = window; - - if (window == nullptr) { - size_t length = frame_length + 1; - hann_window.resize(frame_length); - for (size_t i = 0; i < frame_length; i++) { - hann_window[i] = 0.5f * (1.0f - std::cos(2.0f * static_cast(M_PI) * i / (length - 1))); - } - actual_window = hann_window.data(); - } else if (window_length != frame_length) { - throw std::invalid_argument( - "Length of the window (" + std::to_string(window_length) + - ") must equal frame_length (" + std::to_string(frame_length) + ")"); - } - - if (hop_length <= 0) { - throw std::invalid_argument("hop_length must be greater than zero"); - } - - if (power == 0.0f && mel_filters != nullptr) { - throw std::invalid_argument( - "You have provided `mel_filters` but `power` is `None`. " - "Mel spectrogram computation is not yet supported for complex-valued spectrogram. " - "Specify `power` to fix this issue."); - } - - std::vector padded_waveform; - const float* input_waveform = waveform; - size_t input_length = waveform_length; - - if (center) { - size_t pad_length = frame_length / 2; - size_t padded_length = waveform_length + 2 * pad_length; - padded_waveform.resize(padded_length); - - if (std::strcmp(pad_mode, "reflect") == 0) { - for (size_t i = 0; i < pad_length; i++) { - padded_waveform[i] = waveform[pad_length - i]; - } - - std::copy(waveform, waveform + waveform_length, padded_waveform.data() + pad_length); - - for (size_t i = 0; i < pad_length; i++) { - padded_waveform[pad_length + waveform_length + i] = waveform[waveform_length - 2 - i]; - } - } else { - throw std::invalid_argument("Unsupported pad_mode: " + std::string(pad_mode)); - } - - input_waveform = padded_waveform.data(); - input_length = padded_length; - } - - const size_t num_frames = 1 + (input_length - frame_length) / hop_length; - const size_t num_frequency_bins = (actual_fft_length / 2) + 1; - - std::vector buffer(actual_fft_length); - std::vector raw_complex_frequencies(num_frequency_bins * 2); - - const size_t num_mel_bins = mel_filters != nullptr ? mel_filters_size / num_frequency_bins : 0; - const size_t spectrogram_bins = mel_filters != nullptr ? num_mel_bins : num_frequency_bins; - - std::vector temp_spectrogram(num_frames * num_frequency_bins); - - CactusThreading::parallel_for(num_frames, CactusThreading::Thresholds::SCALAR_EXPENSIVE, [&](size_t start_frame, size_t end_frame) { - std::vector local_buffer(actual_fft_length); - std::vector local_complex_frequencies(num_frequency_bins * 2); - - for (size_t frame_idx = start_frame; frame_idx < end_frame; frame_idx++) { - size_t timestep = frame_idx * hop_length; - std::fill(local_buffer.begin(), local_buffer.end(), 0.0f); - - size_t available_length = std::min(frame_length, input_length - timestep); - std::copy(input_waveform + timestep, input_waveform + timestep + available_length, local_buffer.data()); - - if (dither != 0.0f) { - for (size_t i = 0; i < frame_length; i++) { - float u1 = static_cast(rand()) / static_cast(RAND_MAX); - float u2 = static_cast(rand()) / static_cast(RAND_MAX); - float randn = std::sqrt(-2.0f * std::log(u1)) * std::cos(2.0f * static_cast(M_PI) * u2); - local_buffer[i] += dither * randn; - } - } - - if (remove_dc_offset) { - float mean = 0.0f; - for (size_t i = 0; i < frame_length; i++) { - mean += local_buffer[i]; - } - mean /= static_cast(frame_length); - - for (size_t i = 0; i < frame_length; i++) { - local_buffer[i] -= mean; - } - } - - if (preemphasis != nullptr) { - float preemph_coef = *preemphasis; - for (size_t i = frame_length - 1; i > 0; i--) { - local_buffer[i] -= preemph_coef * local_buffer[i - 1]; - } - local_buffer[0] *= (1.0f - preemph_coef); - } - - for (size_t i = 0; i < frame_length; i++) { - local_buffer[i] *= actual_window[i]; - } - - rfft_f32_1d(local_buffer.data(), local_complex_frequencies.data(), actual_fft_length, "backward"); - - for (size_t i = 0; i < num_frequency_bins; i++) { - float real = local_complex_frequencies[i * 2]; - float imag = local_complex_frequencies[i * 2 + 1]; - float magnitude = std::hypot(real, imag); - temp_spectrogram[frame_idx * num_frequency_bins + i] = std::pow(magnitude, power); - } - } - }); - - if (mel_filters != nullptr) { - CactusThreading::parallel_for_2d(num_mel_bins, num_frames, CactusThreading::Thresholds::AXIS_REDUCE, [&](size_t m, size_t t) { - float sum = 0.0f; - for (size_t f = 0; f < num_frequency_bins; f++) { - sum += mel_filters[m * num_frequency_bins + f] * temp_spectrogram[t * num_frequency_bins + f]; - } - spectrogram[m * num_frames + t] = std::max(mel_floor, sum); - }); - } else { - CactusThreading::parallel_for_2d(num_frames, num_frequency_bins, CactusThreading::Thresholds::AXIS_REDUCE, [&](size_t t, size_t f) { - spectrogram[f * num_frames + t] = temp_spectrogram[t * num_frequency_bins + f]; - }); - } - - if (power != 0.0f && log_mel != nullptr) { - const size_t total_elements = spectrogram_bins * num_frames; - - if (std::strcmp(log_mel, "log") == 0) { - CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - spectrogram[i] = std::log(spectrogram[i]); - } - }); - } else if (std::strcmp(log_mel, "log10") == 0) { - CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ALL_REDUCE, [&](size_t start, size_t end) { - for (size_t i = start; i < end; i++) { - spectrogram[i] = std::log10(spectrogram[i]); - } - }); - } else if (std::strcmp(log_mel, "dB") == 0) { - if (power == 1.0f) { - to_db(spectrogram, total_elements, reference, min_value, db_range, 20.0f); - } else if (power == 2.0f) { - to_db(spectrogram, total_elements, reference, min_value, db_range, 10.0f); - } else { - throw std::invalid_argument( - "Cannot use log_mel option 'dB' with power " + std::to_string(power)); - } - } else { - throw std::invalid_argument("Unknown log_mel option: " + std::string(log_mel)); - } - } -} - -AudioProcessor::AudioProcessor() - : num_frequency_bins_(0), num_mel_filters_(0) {} - -AudioProcessor::~AudioProcessor() = default; - -void AudioProcessor::init_mel_filters(size_t num_frequency_bins, - size_t num_mel_filters, - float min_freq, - float max_freq, - size_t sampling_rate) { - num_frequency_bins_ = num_frequency_bins; - num_mel_filters_ = num_mel_filters; - mel_filters_.resize(num_mel_filters * num_frequency_bins); - - generate_mel_filter_bank( - mel_filters_.data(), - num_frequency_bins, - num_mel_filters, - min_freq, - max_freq, - sampling_rate, - "slaney", - "slaney", - false - ); -} - -std::vector AudioProcessor::compute_spectrogram( - const std::vector& waveform, - const SpectrogramConfig& config) { - - if (mel_filters_.empty()) { - throw std::runtime_error("Mel filters not initialized. Call init_mel_filters() first."); - } - - const size_t n_samples = waveform.size(); - const size_t pad_length = config.center ? config.frame_length / 2 : 0; - const size_t padded_length = n_samples + 2 * pad_length; - const size_t num_frames = 1 + (padded_length - config.frame_length) / config.hop_length; - - std::vector output(num_mel_filters_ * num_frames); - - compute_spectrogram_f32( - waveform.data(), - waveform.size(), - nullptr, - 0, - config.frame_length, - config.hop_length, - &config.n_fft, - output.data(), - config.power, - config.center, - config.pad_mode, - config.onesided, - config.dither, - nullptr, - mel_filters_.data(), - mel_filters_.size(), - config.mel_floor, - config.log_mel, - config.reference, - config.min_value, - nullptr, - config.remove_dc_offset - ); - - return output; -} - -} -} diff --git a/cactus/engine/engine_tokenizer.cpp b/cactus/engine/engine_tokenizer.cpp deleted file mode 100644 index 44e4d5cb7..000000000 --- a/cactus/engine/engine_tokenizer.cpp +++ /dev/null @@ -1,393 +0,0 @@ -#include "engine.h" -#include -#include -#include - -namespace cactus { -namespace engine { - -void Tokenizer::detect_model_type(const std::string& config_path) { - std::ifstream file(config_path); - if (!file.is_open()) { - model_type_ = ModelType::UNKNOWN; - return; - } - - std::string line; - while (std::getline(file, line)) { - size_t pos = line.find("model_type"); - if (pos != std::string::npos) { - std::transform(line.begin(), line.end(), line.begin(), ::tolower); - - if (line.find("qwen") != std::string::npos) { - model_type_ = ModelType::QWEN; - break; - } else if (line.find("gemma") != std::string::npos) { - model_type_ = ModelType::GEMMA; - break; - } else if(line.find("lfm2") != std::string::npos) { - model_type_ = ModelType::LFM2; - } else if (line.find("smol") != std::string::npos) { - model_type_ = ModelType::SMOL; - break; - } else if (line.find("bert") != std::string::npos) { - model_type_ = ModelType::BERT; - break; - } else if (line.find("whisper") != std::string::npos) { - model_type_ = ModelType::WHISPER; - break; - } else { - model_type_ = ModelType::UNKNOWN; - } - } - } - file.clear(); - file.seekg(0); - - while (std::getline(file, line)) { - size_t pos2 = line.find("model_variant"); - if (pos2 != std::string::npos) { - std::transform(line.begin(), line.end(), line.begin(), ::tolower); - - if (line.find("vlm") != std::string::npos) { - model_variant_ = ModelVariant::VLM; - break; - } else if (line.find("extract") != std::string::npos) { - model_variant_ = ModelVariant::EXTRACT; - break; - } else if (line.find("rag") != std::string::npos) { - model_variant_ = ModelVariant::RAG; - break; - } else { - model_variant_ = ModelVariant::DEFAULT; - } - } - } - file.close(); -} - -std::vector Tokenizer::apply_chat_template(const std::vector& messages, bool add_generation_prompt) const { - std::string formatted_prompt = format_chat_prompt(messages, add_generation_prompt); - return encode(formatted_prompt); -} - -std::string Tokenizer::format_chat_prompt(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const { - bool has_images = false; - for (const auto& msg : messages) { - if (!msg.images.empty()) { - has_images = true; - break; - } - } - if (model_type_ == ModelType::LFM2 && has_images) { - return format_lfm2_vl_style(messages, add_generation_prompt, tools_json); - } - - switch (model_type_) { - case ModelType::QWEN: - return format_qwen_style(messages, add_generation_prompt, tools_json); - case ModelType::GEMMA: - return format_gemma_style(messages, add_generation_prompt, tools_json); - case ModelType::LFM2: - return format_lfm2_style(messages, add_generation_prompt, tools_json); - case ModelType::SMOL: - return format_smol_style(messages, add_generation_prompt, tools_json); - default: - return format_qwen_style(messages, add_generation_prompt, tools_json); - } -} - -std::string Tokenizer::format_qwen_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const { - std::string result; - - if (!tools_json.empty()) { - result += "<|im_start|>system\n"; - - bool has_system_msg = false; - for (const auto& msg : messages) { - if (msg.role == "system") { - result += msg.content; - result += "\n\n"; - has_system_msg = true; - break; - } - } - - result += "# Tools\n\n"; - result += "You may call one or more functions to assist with the user query.\n\n"; - result += "You are provided with function signatures within XML tags:\n"; - result += "\n"; - result += tools_json; - result += "\n\n\n"; - result += "For each function call, return a json object with function name and arguments within XML tags:\n"; - result += "\n"; - result += "{\"name\": , \"arguments\": }\n"; - result += ""; - result += "<|im_end|>\n"; - - for (const auto& msg : messages) { - if (msg.role == "system" && has_system_msg) { - continue; - } else if (msg.role == "user") { - result += "<|im_start|>user\n" + msg.content + "<|im_end|>\n"; - } else if (msg.role == "assistant") { - result += "<|im_start|>assistant\n" + msg.content + "<|im_end|>\n"; - } - } - } else { - for (const auto& msg : messages) { - if (msg.role == "system") { - result += "<|im_start|>system\n" + msg.content + "<|im_end|>\n"; - } else if (msg.role == "user") { - result += "<|im_start|>user\n" + msg.content + "<|im_end|>\n"; - } else if (msg.role == "assistant") { - result += "<|im_start|>assistant\n" + msg.content + "<|im_end|>\n"; - } - } - } - - if (add_generation_prompt) { - if (!tools_json.empty()) { - result += "<|im_start|>assistant\n\n\n\n"; - } else { - result += "<|im_start|>assistant\n"; - } - } - - return result; -} - -std::string Tokenizer::format_lfm2_style(const std::vector& messages, - bool add_generation_prompt, - const std::string& tools_json) const -{ - std::string result = "<|startoftext|>"; - - std::string sys_content; - bool has_system_msg = false; - for (const auto& msg : messages) { - if (msg.role == "system") { - sys_content = msg.content; - has_system_msg = true; - break; - } - } - - if (!tools_json.empty()) { - if (!sys_content.empty()) { - sys_content += "\n"; - } - sys_content += "List of tools: <|tool_list_start|>["; - if (!tools_json.empty()) { - sys_content += "\n"; - sys_content += tools_json; - sys_content += "\n"; - } - sys_content += "]<|tool_list_end|>"; - sys_content += "\n\nWhen you need to call a tool, use this exact format:\n"; - sys_content += "<|tool_call_start|>[function_name(arg1=\"value1\", arg2=\"value2\")]<|tool_call_end|>\n"; - sys_content += "You can call multiple tools by using multiple tool call blocks."; - } - - if (!sys_content.empty()) { - result += "<|im_start|>system\n"; - result += sys_content; - result += "<|im_end|>\n"; - } - - for (const auto& msg : messages) { - if (msg.role == "system" && has_system_msg) { - has_system_msg = false; - continue; - } - result += "<|im_start|>" + msg.role + "\n"; - if (msg.role == "tool") { - result += "<|tool_response_start|>"; - result += msg.content; - result += "<|tool_response_end|>"; - } else { - result += msg.content; - } - result += "<|im_end|>\n"; - } - - if (add_generation_prompt) { - result += "<|im_start|>assistant\n"; - } - - return result; -} - -std::string Tokenizer::format_lfm2_vl_style( - const std::vector& messages, - bool add_generation_prompt, - const std::string& tools_json) const -{ - if (!tools_json.empty()) { - return "ERROR: Tool calls are not supported for LFM2-VL models"; - } - - std::string result = "<|startoftext|>"; - - for (const auto& msg : messages) { - result += "<|im_start|>" + msg.role + "\n"; - result += msg.content; - for (const auto& image_path : msg.images) { - int width = 0, height = 0, channels = 0; - unsigned char* img_data = stbi_load(image_path.c_str(), &width, &height, &channels, 0); - - if (img_data) { - Siglip2Preprocessor preprocessor; - auto shape_result = preprocessor.compute_spatial_shapes(height, width); - int downsample_factor = 2; - bool use_thumbnail = true; - int grid_rows = shape_result.grid_rows; - int grid_cols = shape_result.grid_cols; - int num_tiles = grid_rows * grid_cols; - result += "<|image_start|>"; - - if (num_tiles > 1) { - for (int tile_idx = 0; tile_idx < num_tiles; ++tile_idx) { - int row = tile_idx / grid_cols; - int col = tile_idx % grid_cols; - - result += "<|img_row_" + std::to_string(row + 1) + "_col_" + std::to_string(col + 1) + "|>"; - auto [tile_height, tile_width] = shape_result.shapes[tile_idx]; - int tile_tokens = (tile_height * tile_width) / (downsample_factor * downsample_factor); - - for (int t = 0; t < tile_tokens; ++t) { - result += ""; - } - } - if (use_thumbnail && static_cast(num_tiles) < shape_result.shapes.size()) { - result += "<|img_thumbnail|>"; - - auto [thumb_height, thumb_width] = shape_result.shapes[num_tiles]; - int thumbnail_tokens = (thumb_height * thumb_width) / (downsample_factor * downsample_factor); - - for (int t = 0; t < thumbnail_tokens; ++t) { - result += ""; - } - } - } else if (num_tiles == 1) { - auto [thumb_height, thumb_width] = shape_result.shapes[0]; - int thumbnail_tokens = (thumb_height * thumb_width) / (downsample_factor * downsample_factor); - - for (int t = 0; t < thumbnail_tokens; ++t) { - result += ""; - } - } - - result += "<|image_end|>"; - - stbi_image_free(img_data); - } else { - result += ""; - } - } - - result += "<|im_end|>\n"; - } - - if (add_generation_prompt) { - result += "<|im_start|>assistant\n"; - } - - return result; -} - - -std::string Tokenizer::format_gemma_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const { - std::string result = ""; - - std::string system_content; - size_t start_idx = 0; - - if (!messages.empty() && (messages[0].role == "system" || messages[0].role == "developer")) { - system_content = messages[0].content; - start_idx = 1; - } - - if (!tools_json.empty() || !system_content.empty()) { - result += "developer\n"; - if (!system_content.empty()) { - result += system_content; - if (!tools_json.empty()) { - result += "\n"; - } - } - if (!tools_json.empty()) { - result += "You are a model that can do function calling with the following functions."; - result += tools_json; - result += "\n\nWhen you decide to call a function, output it in this exact format:\n"; - result += "call:function_name{arg1:value1,arg2:value2}"; - } - result += "\n"; - } - - std::string prev_message_type; - - for (size_t i = start_idx; i < messages.size(); i++) { - const auto& msg = messages[i]; - - if (msg.role == "tool") { - std::string func_name = msg.name.empty() ? "tool" : msg.name; - result += "response:" + func_name + "{value:" + msg.content + "}"; - prev_message_type = "tool_response"; - } else if (msg.role == "user") { - if (prev_message_type != "tool_response") { - result += "user\n"; - } - result += msg.content; - result += "\n"; - prev_message_type = "content"; - } else if (msg.role == "assistant" || msg.role == "model") { - if (prev_message_type != "tool_response") { - result += "model\n"; - } - result += msg.content; - result += "\n"; - prev_message_type = "content"; - } - } - - if (add_generation_prompt) { - if (prev_message_type != "tool_response") { - result += "model\n"; - } - } - - return result; -} - -std::string Tokenizer::format_smol_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const { - if (!tools_json.empty()) { - return "ERROR: Tool calls are currently not supported for Smol models"; - } - - std::string result; - - if (!messages.empty() && messages.front().role != "system") { - result += "<|im_start|>system\n"; - result += "You are a helpful AI assistant named SmolLM, trained by Hugging Face"; - result += "<|im_end|>\n"; - } - - for (const auto& msg : messages) { - result += "<|im_start|>"; - result += msg.role; - result += "\n"; - result += msg.content; - result += "<|im_end|>\n"; - } - - if (add_generation_prompt) { - result += "<|im_start|>assistant\n"; - } - - return result; -} - - -} // namespace engine -} // namespace cactus \ No newline at end of file diff --git a/cactus/ffi/cactus_complete.cpp b/cactus/ffi/cactus_complete.cpp deleted file mode 100644 index b8ad0bf95..000000000 --- a/cactus/ffi/cactus_complete.cpp +++ /dev/null @@ -1,424 +0,0 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" -#include "cactus_telemetry.h" -#include -#include - -using namespace cactus::engine; -using namespace cactus::ffi; - -static constexpr size_t ROLLING_ENTROPY_WINDOW = 10; - -extern "C" { - -int cactus_complete( - cactus_model_t model, - const char* messages_json, - char* response_buffer, - size_t buffer_size, - const char* options_json, - const char* tools_json, - cactus_token_callback callback, - void* user_data -) { - if (!model) { - std::string error_msg = last_error_message.empty() ? - "Model not initialized. Check model path and files." : last_error_message; - CACTUS_LOG_ERROR("complete", error_msg); - handle_error_response(error_msg, response_buffer, buffer_size); - return -1; - } - - if (!messages_json || !response_buffer || buffer_size == 0) { - CACTUS_LOG_ERROR("complete", "Invalid parameters: messages_json, response_buffer, or buffer_size"); - handle_error_response("Invalid parameters", response_buffer, buffer_size); - return -1; - } - - try { - auto start_time = std::chrono::high_resolution_clock::now(); - - auto* handle = static_cast(model); - auto* tokenizer = handle->model->get_tokenizer(); - handle->should_stop = false; - - std::vector image_paths; - auto messages = parse_messages_json(messages_json, image_paths); - - if (messages.empty()) { - CACTUS_LOG_ERROR("complete", "No messages provided in request"); - handle_error_response("No messages provided", response_buffer, buffer_size); - return -1; - } - - if (handle->corpus_index) { - std::string query; - for (auto it = messages.rbegin(); it != messages.rend(); ++it) { - if (it->role == "user") { - query = it->content; - break; - } - } - - if (!query.empty()) { - std::string rag_context = retrieve_rag_context(handle, query); - if (!rag_context.empty()) { - if (!messages.empty() && messages[0].role == "system") { - messages[0].content = rag_context + messages[0].content; - } else { - ChatMessage system_msg; - system_msg.role = "system"; - system_msg.content = rag_context + "Answer the user's question using ONLY the context above. Do not use any prior knowledge. If the answer cannot be found in the context, respond with \"I don't have enough information to answer that.\""; - messages.insert(messages.begin(), system_msg); - } - } - } - } - - float temperature, top_p, confidence_threshold; - size_t top_k, max_tokens, tool_rag_top_k; - std::vector stop_sequences; - bool force_tools; - parse_options_json(options_json ? options_json : "", - temperature, top_p, top_k, max_tokens, stop_sequences, force_tools, tool_rag_top_k, confidence_threshold); - - std::vector tools; - if (tools_json && strlen(tools_json) > 0) - tools = parse_tools_json(tools_json); - - if (tool_rag_top_k > 0 && tools.size() > tool_rag_top_k) { - std::string query; - for (auto it = messages.rbegin(); it != messages.rend(); ++it) { - if (it->role == "user") { - query = it->content; - break; - } - } - if (!query.empty()) { - tools = select_relevant_tools(handle, query, tools, tool_rag_top_k); - } - } - - if (force_tools && !tools.empty()) { - std::vector function_names; - function_names.reserve(tools.size()); - for (const auto& tool : tools) { - function_names.push_back(tool.name); - } - handle->model->set_tool_constraints(function_names); - - if (temperature == 0.0f) { - temperature = 0.01f; - } - } - - Config::ModelType model_type = handle->model->get_config().model_type; - std::string formatted_tools; - if (model_type == Config::ModelType::GEMMA) { - formatted_tools = gemma::format_tools(tools); - } else { - formatted_tools = format_tools_for_prompt(tools); - } - std::string full_prompt = tokenizer->format_chat_prompt(messages, true, formatted_tools); - - if (full_prompt.find("ERROR:") == 0) { - CACTUS_LOG_ERROR("complete", "Prompt formatting failed: " << full_prompt.substr(6)); - handle_error_response(full_prompt.substr(6), response_buffer, buffer_size); - return -1; - } - - std::vector current_prompt_tokens = tokenizer->encode(full_prompt); - - CACTUS_LOG_DEBUG("complete", "Prompt tokens: " << current_prompt_tokens.size() << ", max_tokens: " << max_tokens); - - std::vector tokens_to_process; - - bool has_images = !image_paths.empty(); - bool is_prefix = !has_images && - (current_prompt_tokens.size() >= handle->processed_tokens.size()) && - std::equal(handle->processed_tokens.begin(), handle->processed_tokens.end(), current_prompt_tokens.begin()); - - if (handle->processed_tokens.empty() || !is_prefix) { - if (!has_images) { - handle->model->reset_cache(); - } - tokens_to_process = current_prompt_tokens; - } else { - tokens_to_process.assign(current_prompt_tokens.begin() + handle->processed_tokens.size(), current_prompt_tokens.end()); - } - - size_t prompt_tokens = tokens_to_process.size(); - - std::vector> stop_token_sequences; - stop_token_sequences.push_back({tokenizer->get_eos_token()}); - for (const auto& stop_seq : stop_sequences) - stop_token_sequences.push_back(tokenizer->encode(stop_seq)); - - if (model_type == Config::ModelType::GEMMA && !tools.empty()) { - stop_token_sequences.push_back(tokenizer->encode("")); - stop_token_sequences.push_back(tokenizer->encode("")); - } - - std::vector generated_tokens; - double time_to_first_token = 0.0; - uint32_t next_token; - float first_token_entropy = 0.0f; - - if (tokens_to_process.empty()) { - if (handle->processed_tokens.empty()) { - handle_error_response("Cannot generate from empty prompt", response_buffer, buffer_size); - return -1; - } - std::vector last_token_vec = { handle->processed_tokens.back() }; - next_token = handle->model->decode(last_token_vec, temperature, top_p, top_k, "", &first_token_entropy); - } else { - if (!image_paths.empty()) { - next_token = handle->model->decode_with_images(tokens_to_process, image_paths, temperature, top_p, top_k, "", &first_token_entropy); - } else { - size_t prefill_chunk_size = handle->model->get_prefill_chunk_size(); - - if (tokens_to_process.size() > 1) { - std::vector prefill_tokens(tokens_to_process.begin(), - tokens_to_process.end() - 1); - handle->model->prefill(prefill_tokens, prefill_chunk_size); - - std::vector last_token = {tokens_to_process.back()}; - next_token = handle->model->decode(last_token, temperature, top_p, top_k, "", &first_token_entropy); - } else { - next_token = handle->model->decode(tokens_to_process, temperature, top_p, top_k, "", &first_token_entropy); - } - } - } - - handle->processed_tokens = current_prompt_tokens; - - auto token_end = std::chrono::high_resolution_clock::now(); - time_to_first_token = std::chrono::duration_cast(token_end - start_time).count() / 1000.0; - - float confidence = 1.0f - first_token_entropy; - - if (confidence < confidence_threshold) { - double prefill_tps = time_to_first_token > 0 ? (prompt_tokens * 1000.0) / time_to_first_token : 0.0; - std::string result = construct_cloud_handoff_json(confidence, time_to_first_token, prefill_tps, prompt_tokens); - if (result.length() >= buffer_size) { - handle_error_response("Response buffer too small", response_buffer, buffer_size); - return -1; - } - std::strcpy(response_buffer, result.c_str()); - - CactusTelemetry::getInstance().recordCompletion( - handle->model_name, - true, - time_to_first_token, - 0.0, - time_to_first_token, - static_cast(prompt_tokens), - "" - ); - - return static_cast(result.length()); - } - - generated_tokens.push_back(next_token); - handle->processed_tokens.push_back(next_token); - - if (force_tools && !tools.empty()) { - handle->model->update_tool_constraints(next_token); - } - - std::vector entropy_window; - entropy_window.push_back(first_token_entropy); - float entropy_sum = first_token_entropy; - float total_entropy_sum = first_token_entropy; - size_t total_entropy_count = 1; - bool entropy_spike_handoff = false; - - if (!matches_stop_sequence(generated_tokens, stop_token_sequences)) { - if (callback) { - std::string new_text = tokenizer->decode({next_token}); - callback(new_text.c_str(), next_token, user_data); - } - - for (size_t i = 1; i < max_tokens; i++) { - if (handle->should_stop) break; - - float token_entropy = 0.0f; - next_token = handle->model->decode({next_token}, temperature, top_p, top_k, "", &token_entropy); - generated_tokens.push_back(next_token); - handle->processed_tokens.push_back(next_token); - - total_entropy_sum += token_entropy; - total_entropy_count++; - - entropy_window.push_back(token_entropy); - entropy_sum += token_entropy; - if (entropy_window.size() > ROLLING_ENTROPY_WINDOW) { - entropy_sum -= entropy_window.front(); - entropy_window.erase(entropy_window.begin()); - } - - float rolling_mean_entropy = entropy_sum / entropy_window.size(); - float rolling_confidence = 1.0f - rolling_mean_entropy; - if (rolling_confidence < confidence_threshold) { - entropy_spike_handoff = true; - break; - } - - if (force_tools && !tools.empty()) { - handle->model->update_tool_constraints(next_token); - } - - if (matches_stop_sequence(generated_tokens, stop_token_sequences)) break; - - if (callback) { - std::string new_text = tokenizer->decode({next_token}); - callback(new_text.c_str(), next_token, user_data); - } - } - } - - float mean_entropy = total_entropy_sum / static_cast(total_entropy_count); - confidence = 1.0f - mean_entropy; - - if (force_tools && !tools.empty()) { - handle->model->clear_tool_constraints(); - } - - auto end_time = std::chrono::high_resolution_clock::now(); - double total_time_ms = std::chrono::duration_cast(end_time - start_time).count() / 1000.0; - - size_t completion_tokens = generated_tokens.size(); - double decode_time_ms = total_time_ms - time_to_first_token; - double prefill_tps = time_to_first_token > 0 ? (prompt_tokens * 1000.0) / time_to_first_token : 0.0; - double decode_tps = (completion_tokens > 1 && decode_time_ms > 0) ? ((completion_tokens - 1) * 1000.0) / decode_time_ms : 0.0; - - std::string response_text = tokenizer->decode(generated_tokens); - - std::string regular_response; - std::vector function_calls; - parse_function_calls_from_response(response_text, regular_response, function_calls); - - std::string result = construct_response_json(regular_response, function_calls, time_to_first_token, - total_time_ms, prefill_tps, decode_tps, prompt_tokens, - completion_tokens, confidence, entropy_spike_handoff); - - if (result.length() >= buffer_size) { - handle_error_response("Response buffer too small", response_buffer, buffer_size); - return -1; - } - - std::strcpy(response_buffer, result.c_str()); - - CactusTelemetry::getInstance().recordCompletion( - handle->model_name, - true, - time_to_first_token, - decode_tps, - total_time_ms, - static_cast(prompt_tokens + completion_tokens), - "" - ); - - return static_cast(result.length()); - - } catch (const std::exception& e) { - CACTUS_LOG_ERROR("complete", "Exception: " << e.what()); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordCompletion( - handle ? handle->model_name : "unknown", - false, - 0.0, 0.0, 0.0, 0, - std::string(e.what()) - ); - - handle_error_response(e.what(), response_buffer, buffer_size); - return -1; - } catch (...) { - CACTUS_LOG_ERROR("complete", "Unknown exception during completion"); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordCompletion( - handle ? handle->model_name : "unknown", - false, - 0.0, 0.0, 0.0, 0, - "Unknown exception" - ); - - handle_error_response("Unknown error during completion", response_buffer, buffer_size); - return -1; - } -} - -int cactus_tokenize( - cactus_model_t model, - const char* text, - uint32_t* token_buffer, - size_t token_buffer_len, - size_t* out_token_len -) { - if (!model || !text || !out_token_len) return -1; - - try { - auto* handle = static_cast(model); - auto* tokenizer = handle->model->get_tokenizer(); - - std::vector toks = tokenizer->encode(std::string(text)); - *out_token_len = toks.size(); - - if (!token_buffer || token_buffer_len == 0) return 0; - if (token_buffer_len < toks.size()) return -2; - - std::memcpy(token_buffer, toks.data(), toks.size() * sizeof(uint32_t)); - return 0; - } catch (...) { - return -1; - } -} - -int cactus_score_window( - cactus_model_t model, - const uint32_t* tokens, - size_t token_len, - size_t start, - size_t end, - size_t context, - char* response_buffer, - size_t buffer_size -) { - if (!model || !tokens || token_len == 0 || !response_buffer || buffer_size == 0) { - handle_error_response("Invalid parameters", response_buffer, buffer_size); - return -1; - } - - try { - auto* handle = static_cast(model); - - std::vector vec(tokens, tokens + token_len); - - size_t scored = 0; - double logprob = handle->model->score_tokens_window_logprob(vec, start, end, context, &scored); - - std::ostringstream oss; - oss << "{" - << "\"success\":true," - << "\"logprob\":" << std::setprecision(10) << logprob << "," - << "\"tokens\":" << scored - << "}"; - - std::string result = oss.str(); - if (result.size() >= buffer_size) { - handle_error_response("Response buffer too small", response_buffer, buffer_size); - return -1; - } - - std::strcpy(response_buffer, result.c_str()); - return (int)result.size(); - - } catch (const std::exception& e) { - handle_error_response(e.what(), response_buffer, buffer_size); - return -1; - } -} - -} diff --git a/cactus/ffi/cactus_stream.cpp b/cactus/ffi/cactus_stream.cpp deleted file mode 100644 index 2cf5aa0b5..000000000 --- a/cactus/ffi/cactus_stream.cpp +++ /dev/null @@ -1,359 +0,0 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" -#include -#include - -using namespace cactus::ffi; - -static std::string extract_json_string_value(const std::string& json, const std::string& key) { - std::string pattern = "\"" + key + "\":"; - size_t pos = json.find(pattern); - if (pos == std::string::npos) return ""; - - size_t colon_pos = json.find(':', pos); - if (colon_pos == std::string::npos) return ""; - - size_t start_quote = json.find('"', colon_pos); - if (start_quote == std::string::npos) return ""; - - std::string result; - size_t i = start_quote + 1; - while (i < json.length()) { - if (json[i] == '"' && (i == start_quote + 1 || json[i - 1] != '\\')) { - break; - } - if (json[i] == '\\' && i + 1 < json.length()) { - char next = json[i + 1]; - if (next == '"') { - result += '"'; - i += 2; - } else if (next == '\\') { - result += '\\'; - i += 2; - } else if (next == 'n') { - result += '\n'; - i += 2; - } else if (next == 'r') { - result += '\r'; - i += 2; - } else if (next == 't') { - result += '\t'; - i += 2; - } else { - result += json[i]; - i++; - } - } else { - result += json[i]; - i++; - } - } - return result; -} - -static bool fuzzy_match(const std::string& a, const std::string& b, size_t n, double threshold) { - if (!n) return false; - if (a.size() < n || b.size() < n) return false; - - std::vector dp(n + 1); - size_t dp_im1_jm1; - - for (size_t j = 0; j <= n; ++j) dp[j] = j; - - for (size_t i = 1; i <= n; ++i) { - dp_im1_jm1 = dp[0]; - dp[0] = i; - - for (size_t j = 1; j <= n; ++j) { - size_t dp_im1_j = dp[j]; - - if (a[i - 1] == b[j - 1]) { - dp[j] = dp_im1_jm1; - } else { - dp[j] = std::min({ - dp[j] + 1, - dp[j - 1] + 1, - dp_im1_jm1 + 1 - }); - } - - dp_im1_jm1 = dp_im1_j; - } - } - - return 1.0 - static_cast(dp[n]) / static_cast(n) >= threshold; -} - -static std::string get_last_n_words(const std::string& text, size_t n) { - if (text.empty() || n == 0) return ""; - - size_t word_count = 0; - bool in_word = false; - - for (size_t i = text.length(); i-- > 0; ) { - bool is_space = std::isspace(text[i]); - if (is_space && in_word) { - ++word_count; - in_word = false; - if (word_count == n) { - return text.substr(i + 1); - } - } else if (!is_space) { - in_word = true; - } - } - - return text; -} - -static void parse_stream_transcribe_process_options(const std::string& json, double& confirmation_threshold) { - confirmation_threshold = 0.95; - - if (json.empty()) { - return; - } - - size_t pos = json.find("\"confirmation_threshold\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - confirmation_threshold = std::stod(json.substr(pos)); - } -} - -struct CactusStreamTranscribeHandle { - CactusModelHandle* model_handle; - - std::string confirmed; - std::string pending; - - std::vector audio_buffer; - - std::string last_n_words; - std::string previous_transcription; - size_t previous_audio_buffer_size; - - char transcribe_response_buffer[8192]; -}; - -extern "C" { - -cactus_stream_transcribe_t cactus_stream_transcribe_init(cactus_model_t model) { - if (!model) { - last_error_message = "Model not initialized. Check model path and files."; - CACTUS_LOG_ERROR("stream_transcribe_init", last_error_message); - return nullptr; - } - - try { - auto* model_handle = static_cast(model); - if (!model_handle->model) { - last_error_message = "Invalid model handle."; - CACTUS_LOG_ERROR("stream_transcribe_init", last_error_message); - return nullptr; - } - - auto* stream_handle = new CactusStreamTranscribeHandle(); - stream_handle->model_handle = model_handle; - stream_handle->previous_audio_buffer_size = 0; - stream_handle->transcribe_response_buffer[0] = '\0'; - - CACTUS_LOG_INFO("stream_transcribe_init", - "Stream transcription initialized for model: " << model_handle->model_name); - - return stream_handle; - } catch (const std::exception& e) { - last_error_message = "Exception during stream_transcribe_init: " + std::string(e.what()); - CACTUS_LOG_ERROR("stream_transcribe_init", last_error_message); - return nullptr; - } catch (...) { - last_error_message = "Unknown exception during stream transcription initialization"; - CACTUS_LOG_ERROR("stream_transcribe_init", last_error_message); - return nullptr; - } -} - -int cactus_stream_transcribe_insert( - cactus_stream_transcribe_t stream, - const uint8_t* pcm_buffer, - size_t pcm_buffer_size -) { - if (!stream) { - last_error_message = "Stream not initialized."; - CACTUS_LOG_ERROR("stream_transcribe_insert", last_error_message); - return -1; - } - - if (!pcm_buffer || pcm_buffer_size == 0) { - last_error_message = "Invalid parameters: pcm_buffer or pcm_buffer_size"; - CACTUS_LOG_ERROR("stream_transcribe_insert", last_error_message); - return -1; - } - - try { - auto* handle = static_cast(stream); - handle->audio_buffer.insert(handle->audio_buffer.end(), pcm_buffer, - pcm_buffer + pcm_buffer_size); - - CACTUS_LOG_DEBUG("stream_transcribe_insert", - "Inserted " << pcm_buffer_size << " bytes, buffer size: " << handle->audio_buffer.size()); - - return 0; - } catch (const std::exception& e) { - last_error_message = "Exception during stream_transcribe_insert: " + std::string(e.what()); - CACTUS_LOG_ERROR("stream_transcribe_insert", last_error_message); - return -1; - } catch (...) { - last_error_message = "Unknown exception during audio buffer insertion"; - CACTUS_LOG_ERROR("stream_transcribe_insert", last_error_message); - return -1; - } -} - -int cactus_stream_transcribe_process( - cactus_stream_transcribe_t stream, - char* response_buffer, - size_t buffer_size, - const char* options_json -) { - if (!stream) { - last_error_message = "Stream not initialized."; - CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); - return -1; - } - - if (!response_buffer || buffer_size == 0) { - last_error_message = "Invalid parameters: response_buffer or buffer_size"; - CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); - return -1; - } - - try { - auto* handle = static_cast(stream); - - double confirmation_threshold; - parse_stream_transcribe_process_options( - options_json ? options_json : "", - confirmation_threshold - ); - - std::string prompt = "<|startofprev|>" - + handle->last_n_words - + "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"; - - const int result = cactus_transcribe( - handle->model_handle, - nullptr, - prompt.c_str(), - handle->transcribe_response_buffer, - sizeof(handle->transcribe_response_buffer), - nullptr, - nullptr, - nullptr, - handle->audio_buffer.data(), - handle->audio_buffer.size()); - - cactus_reset(handle->model_handle); - - if (result < 0) { - last_error_message = "Transcription failed in stream process."; - CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); - handle_error_response(last_error_message, response_buffer, buffer_size); - return -1; - } - - std::string json_str(handle->transcribe_response_buffer); - std::string response = extract_json_string_value(json_str, "response"); - std::string json_response = "{\"success\":true,\"confirmed\":\"" + - escape_json_string(handle->confirmed) + "\",\"pending\":\"" + - escape_json_string(response) + "\"}"; - - if (json_response.length() >= buffer_size) { - last_error_message = "Response buffer too small"; - CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); - handle_error_response(last_error_message, response_buffer, buffer_size); - return -1; - } - - std::strcpy(response_buffer, json_response.c_str()); - - const size_t n = std::min(handle->previous_transcription.size(), response.size()); - if (fuzzy_match(handle->previous_transcription, response, n, confirmation_threshold)) { - handle->audio_buffer.erase( - handle->audio_buffer.begin(), - handle->audio_buffer.begin() + handle->previous_audio_buffer_size - ); - handle->last_n_words = get_last_n_words(handle->last_n_words + handle->previous_transcription, 200); - handle->confirmed = std::move(handle->previous_transcription); - handle->previous_transcription.clear(); - handle->previous_audio_buffer_size = 0; - } else { - handle->confirmed.clear(); - handle->previous_transcription = std::move(response); - handle->previous_audio_buffer_size = handle->audio_buffer.size(); - } - - return static_cast(json_response.length()); - } catch (const std::exception& e) { - last_error_message = "Exception during stream_transcribe_process: " + std::string(e.what()); - CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); - handle_error_response(e.what(), response_buffer, buffer_size); - return -1; - } catch (...) { - last_error_message = "Unknown exception during stream transcription processing"; - CACTUS_LOG_ERROR("stream_transcribe_process", last_error_message); - handle_error_response("Unknown error during stream processing", response_buffer, buffer_size); - return -1; - } -} - -int cactus_stream_transcribe_finalize( - cactus_stream_transcribe_t stream, - char* response_buffer, - size_t buffer_size -) { - if (!stream) { - last_error_message = "Stream not initialized."; - CACTUS_LOG_ERROR("stream_transcribe_finalize", last_error_message); - return -1; - } - - if (!response_buffer || buffer_size == 0) { - last_error_message = "Invalid parameters: response_buffer or buffer_size"; - CACTUS_LOG_ERROR("stream_transcribe_finalize", last_error_message); - return -1; - } - - try { - auto* handle = static_cast(stream); - - std::string json_response = "{\"success\":true,\"confirmed\":\"" + - escape_json_string(handle->confirmed + handle->previous_transcription) + "\"}"; - - if (json_response.length() >= buffer_size) { - last_error_message = "Response buffer too small"; - CACTUS_LOG_ERROR("stream_transcribe_finalize", last_error_message); - handle_error_response(last_error_message, response_buffer, buffer_size); - return -1; - } - - std::strcpy(response_buffer, json_response.c_str()); - - return static_cast(json_response.length()); - } catch (const std::exception& e) { - last_error_message = "Exception during stream_transcribe_finalize: " + std::string(e.what()); - CACTUS_LOG_ERROR("stream_transcribe_finalize", last_error_message); - handle_error_response(e.what(), response_buffer, buffer_size); - return -1; - } catch (...) { - last_error_message = "Unknown exception during stream transcription finalization"; - CACTUS_LOG_ERROR("stream_transcribe_finalize", last_error_message); - handle_error_response("Unknown error during stream finalization", response_buffer, buffer_size); - return -1; - } -} - -void cactus_stream_transcribe_destroy(cactus_stream_transcribe_t stream) { - if (stream) delete static_cast(stream); -} - -} diff --git a/cactus/ffi/cactus_telemetry.cpp b/cactus/ffi/cactus_telemetry.cpp deleted file mode 100644 index 7f9cccdb9..000000000 --- a/cactus/ffi/cactus_telemetry.cpp +++ /dev/null @@ -1,21 +0,0 @@ -#include "cactus_telemetry.h" - -extern "C" { - -void cactus_set_telemetry_token(const char* token) { - auto& telemetry = cactus::ffi::CactusTelemetry::getInstance(); - if (token && token[0] != '\0') { - telemetry.setTelemetryToken(token); - telemetry.setEnabled(true); - } else { - telemetry.setEnabled(false); - } -} - -void cactus_set_pro_key(const char* pro_key) { - if (pro_key) { - cactus::ffi::DeviceManager::setProKey(pro_key); - } -} - -} \ No newline at end of file diff --git a/cactus/ffi/cactus_telemetry.h b/cactus/ffi/cactus_telemetry.h deleted file mode 100644 index a48b6fc87..000000000 --- a/cactus/ffi/cactus_telemetry.h +++ /dev/null @@ -1,656 +0,0 @@ -#ifndef CACTUS_TELEMETRY_H -#define CACTUS_TELEMETRY_H - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "cactus_utils.h" - -#if defined(__APPLE__) -#include -#endif - -#if defined(__APPLE__) && (!TARGET_OS_IPHONE) && !defined(__ANDROID__) -#define CACTUS_TELEMETRY_ENABLED -#include -#include -#include -#endif - -namespace cactus { -namespace ffi { - -enum class TelemetryEventType { - Init, - Completion, - Embedding, - Transcription -}; - -struct TelemetryMetrics { - TelemetryEventType event_type; - std::string model; - - double ttft_ms = 0.0; - double tps = 0.0; - double response_time_ms = 0.0; - int tokens = 0; - - bool success = false; - std::string message; - - std::chrono::system_clock::time_point timestamp; -}; - - -class HttpClient { -public: - struct Response { - bool success; - int status_code; - std::string body; - }; - - static Response postJson( - const std::string& url, - const std::map& headers, - const std::string& json_body - ); - -private: - static size_t writeCallback(void* contents, size_t size, size_t nmemb, void* userp); -}; - -inline size_t HttpClient::writeCallback(void* contents, size_t size, size_t nmemb, void* userp) { - size_t totalSize = size * nmemb; - std::string* response = static_cast(userp); - response->append(static_cast(contents), totalSize); - return totalSize; -} - -inline HttpClient::Response HttpClient::postJson( - [[maybe_unused]] const std::string& url, - [[maybe_unused]] const std::map& headers, - [[maybe_unused]] const std::string& json_body -) { -#ifdef CACTUS_TELEMETRY_ENABLED - Response response; - response.success = false; - response.status_code = 0; - - CURL* curl = curl_easy_init(); - if (!curl) { - return response; - } - - curl_easy_setopt(curl, CURLOPT_URL, url.c_str()); - - curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 1L); - curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 2L); - - curl_easy_setopt(curl, CURLOPT_POST, 1L); - - curl_easy_setopt(curl, CURLOPT_POSTFIELDS, json_body.c_str()); - curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, json_body.length()); - - struct curl_slist* header_list = nullptr; - for (const auto& header : headers) { - std::string header_str = header.first + ": " + header.second; - header_list = curl_slist_append(header_list, header_str.c_str()); - } - curl_easy_setopt(curl, CURLOPT_HTTPHEADER, header_list); - - curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, writeCallback); - curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response.body); - - curl_easy_setopt(curl, CURLOPT_TIMEOUT, 5L); - - CURLcode res = curl_easy_perform(curl); - - if (res == CURLE_OK) { - long response_code = 0; - curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &response_code); - response.status_code = static_cast(response_code); - response.success = (response_code >= 200 && response_code < 300); - - if (!response.success && !response.body.empty()) { - std::cerr << "[Telemetry] Response body: " << response.body << std::endl; - } - } else { - std::cerr << "[Telemetry] HTTP POST failed: " << curl_easy_strerror(res) << std::endl; - } - - if (header_list) { - curl_slist_free_all(header_list); - } - curl_easy_cleanup(curl); - - return response; -#else - (void)url; - (void)headers; - (void)json_body; - Response response; - response.success = false; - response.status_code = 0; - return response; -#endif -} - -class DeviceManager { -public: - static std::string getDeviceId(); - static std::string getProjectId(); - static std::map getDeviceMetadata(); - static std::string registerDevice(const std::string& device_id = "", const std::string& pro_key = ""); - - static void setProKey(const std::string& key); - static std::string getProKey(); - -private: - static std::string getConfigPath(); - static std::map readConfig(); - static void writeConfig(const std::map& config); - - static std::string pro_key_; -}; - -inline std::string DeviceManager::pro_key_ = ""; - -inline void DeviceManager::setProKey(const std::string& key) { - pro_key_ = key; -} - -inline std::string DeviceManager::getProKey() { - return pro_key_; -} - -inline std::string DeviceManager::getConfigPath() { - const char* home = getenv("HOME"); - if (!home) { - home = "/tmp"; - } - - std::string cactus_dir = std::string(home) + "/.cactus"; - - mkdir(cactus_dir.c_str(), 0755); - - return cactus_dir + "/telemetry_config.json"; -} - -inline std::map DeviceManager::readConfig() { - std::map config; - std::string path = getConfigPath(); - std::ifstream file(path); - - if (file.is_open()) { - std::stringstream buffer; - buffer << file.rdbuf(); - file.close(); - - std::string content = buffer.str(); - - const std::string project_id_key = "\"project_id\":\""; - size_t project_pos = content.find(project_id_key); - if (project_pos != std::string::npos) { - size_t start = project_pos + project_id_key.length(); - size_t end = content.find("\"", start); - if (end != std::string::npos) { - config["project_id"] = content.substr(start, end - start); - } - } - } - - return config; -} - -inline void DeviceManager::writeConfig(const std::map& config) { - std::string path = getConfigPath(); - std::ofstream file(path); - - if (file.is_open()) { - file << "{\n"; - - auto device_it = config.find("device_id"); - if (device_it != config.end()) { - file << " \"device_id\":\"" << device_it->second << "\""; - } - - auto project_it = config.find("project_id"); - if (project_it != config.end()) { - if (device_it != config.end()) { - file << ",\n"; - } - file << " \"project_id\":\"" << project_it->second << "\""; - } - - file << "\n}\n"; - file.close(); - } -} - -inline std::string DeviceManager::getDeviceId() { - auto config = readConfig(); - std::string pro_key = getProKey(); - - std::string project_id = config["project_id"]; - if (project_id.empty()) { - project_id = generateUUID(); - } - - config["project_id"] = project_id; - writeConfig(config); - - const char* device_id_cstr = get_device_id(pro_key.c_str()); - if (device_id_cstr != nullptr) { - std::string device_id = std::string(device_id_cstr); - size_t pipe_pos = device_id.find('|'); - if (pipe_pos != std::string::npos) { - std::string device_part = device_id.substr(0, pipe_pos); - std::string pro_key_part = device_id.substr(pipe_pos + 1); - setProKey(pro_key_part); - return registerDevice(device_part, pro_key_part); - } - return device_id; - } - return registerDevice("", pro_key); -} - -inline std::string DeviceManager::getProjectId() { - auto config = readConfig(); - std::string project_id = config["project_id"]; - - if (!project_id.empty()) { - return project_id; - } - - project_id = generateUUID(); - std::cerr << "[Device Manager] Generated new project ID: " << project_id << std::endl; - - config["project_id"] = project_id; - writeConfig(config); - - return project_id; -} - -inline std::string DeviceManager::registerDevice(const std::string& device_id, const std::string& pro_key) { -#ifdef CACTUS_TELEMETRY_ENABLED - static const std::string SUPABASE_URL = "https://vlqqczxwyaodtcdmdmlw.supabase.co"; - static const std::string SUPABASE_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InZscXFjenh3eWFvZHRjZG1kbWx3Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3NTE1MTg2MzIsImV4cCI6MjA2NzA5NDYzMn0.nBzqGuK9j6RZ6mOPWU2boAC_5H9XDs-fPpo5P3WZYbI"; - - auto metadata = getDeviceMetadata(); - - std::ostringstream json; - json << "{"; - - if (!device_id.empty()) { - json << "\"device_id\":\"" << generateUUID() << "\""; - } else { - json << "\"device_data\":{" - << "\"model\":\"" << metadata["model"] << "\","; - json << "\"os\":\"" << metadata["os"] << "\","; - json << "\"os_version\":\"" << metadata["os_version"] << "\","; - json << "\"brand\":\"" << metadata["brand"] << "\""; - json << "}"; - } - - json << ",\"cactus_pro_key\":\"" << pro_key << "\""; - json << "}"; - - std::string payload = json.str(); - - std::map headers; - headers["Content-Type"] = "application/json"; - - std::string url = SUPABASE_URL + "/functions/v1/device-registration"; - - auto response = HttpClient::postJson(url, headers, payload); - - if (response.success && !response.body.empty()) { - const char* registered_id_cstr = register_app(response.body.c_str()); - std::string registered_id = (registered_id_cstr && registered_id_cstr[0] != '\0') - ? std::string(registered_id_cstr) - : std::string(); - - if (!registered_id.empty()) { - std::cerr << "[Device Registration] SUCCESS - Device registered!" << std::endl; - return registered_id; - } - - std::cerr << "[Device Registration] FAILED - Could not parse ID from response" << std::endl; - return ""; - } else { - std::cerr << "[Device Registration] FAILED - Direct table insertion unsuccessful" << std::endl; - return ""; - } -#else - return ""; -#endif -} - -inline std::map DeviceManager::getDeviceMetadata() { - std::map metadata; - -#ifdef CACTUS_TELEMETRY_ENABLED - struct utsname system_info; - if (uname(&system_info) == 0) { - metadata["os"] = "macOS"; - metadata["os_version"] = system_info.release; - metadata["architecture"] = system_info.machine; - metadata["model"] = system_info.machine; - metadata["brand"] = "apple"; - } -#else - metadata["os"] = "unknown"; - metadata["os_version"] = "unknown"; - metadata["architecture"] = "unknown"; - metadata["model"] = "unknown"; - metadata["brand"] = "unknown"; -#endif - - return metadata; -} - -class LogRecord { -public: - static std::string buildJson( - const TelemetryMetrics& metrics, - const std::string& project_id, - const std::string& device_id, - const std::string& telemetry_token - ); - -private: - static std::string escapeJson(const std::string& input); - static std::string formatTimestamp(const std::chrono::system_clock::time_point& timestamp); - static std::string eventTypeToString(TelemetryEventType type); -}; - -inline std::string LogRecord::escapeJson(const std::string& input) { - std::ostringstream output; - for (char c : input) { - switch (c) { - case '"': output << "\\\""; break; - case '\\': output << "\\\\"; break; - case '\b': output << "\\b"; break; - case '\f': output << "\\f"; break; - case '\n': output << "\\n"; break; - case '\r': output << "\\r"; break; - case '\t': output << "\\t"; break; - default: - if (static_cast(c) < 0x20) { - output << "\\u" << std::hex << std::setw(4) << std::setfill('0') - << static_cast(c); - } else { - output << c; - } - } - } - return output.str(); -} - -inline std::string LogRecord::formatTimestamp(const std::chrono::system_clock::time_point& timestamp) { - auto time_t = std::chrono::system_clock::to_time_t(timestamp); - auto ms = std::chrono::duration_cast( - timestamp.time_since_epoch()) % 1000; - - std::tm tm; -#ifdef _WIN32 - gmtime_s(&tm, &time_t); -#else - gmtime_r(&time_t, &tm); -#endif - - std::ostringstream oss; - oss << std::put_time(&tm, "%Y-%m-%dT%H:%M:%S") - << '.' << std::setfill('0') << std::setw(3) << ms.count() << 'Z'; - return oss.str(); -} - -inline std::string LogRecord::eventTypeToString(TelemetryEventType type) { - switch (type) { - case TelemetryEventType::Init: - return "init"; - case TelemetryEventType::Completion: - return "completion"; - case TelemetryEventType::Embedding: - return "embedding"; - case TelemetryEventType::Transcription: - return "transcription"; - default: - return "unknown"; - } -} - -inline std::string LogRecord::buildJson( - const TelemetryMetrics& metrics, - const std::string& project_id, - const std::string& device_id, - const std::string& telemetry_token -) { - std::ostringstream json; - json << std::fixed << std::setprecision(2); - - json << "{"; - json << "\"event_type\":\"" << eventTypeToString(metrics.event_type) << "\","; - json << "\"model\":\"" << escapeJson(metrics.model) << "\","; - json << "\"success\":" << (metrics.success ? "true" : "false") << ","; - json << "\"project_id\":\"" << project_id << "\","; - json << "\"device_id\":\"" << device_id << "\","; - json << "\"telemetry_token\":\"" << telemetry_token << "\","; - json << "\"framework\":\"cpp\","; - json << "\"framework_version\":\"" << getVersion() << "\""; - - json << ",\"ttft\":" << metrics.ttft_ms; - json << ",\"tps\":" << metrics.tps; - json << ",\"response_time\":" << metrics.response_time_ms; - json << ",\"tokens\":" << metrics.tokens; - - if (!metrics.message.empty()) { - json << ",\"message\":\"" << escapeJson(metrics.message) << "\""; - } - - json << "}"; - return json.str(); -} - -class CactusTelemetry { -public: - static CactusTelemetry& getInstance(); - - void setEnabled(bool enabled); - void setTelemetryToken(const std::string& token); - void setProjectId(const std::string& project_id); - void ensureInitialized(); - - void recordEvent(const TelemetryMetrics& metrics); - - void recordInit(const std::string& model, bool success, const std::string& message = ""); - - void recordCompletion(const std::string& model, bool success, - double ttft_ms, double tps, double response_time_ms, - int tokens, const std::string& message = ""); - - void recordEmbedding(const std::string& model, bool success, - const std::string& message = ""); - - void recordTranscription(const std::string& model, bool success, - double ttft_ms, double tps, double response_time_ms, - int tokens, const std::string& message = ""); - - bool isEnabled() const; - -private: - CactusTelemetry(); - ~CactusTelemetry() = default; - - CactusTelemetry(const CactusTelemetry&) = delete; - CactusTelemetry& operator=(const CactusTelemetry&) = delete; - - void sendToSupabase(const TelemetryMetrics& metrics); - - bool enabled_ = false; - bool initialized_ = false; - std::string telemetry_token_; - std::string project_id_; - std::string device_id_; - - mutable std::mutex mutex_; -}; - -static const std::string SUPABASE_URL = "https://vlqqczxwyaodtcdmdmlw.supabase.co"; -static const std::string SUPABASE_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InZscXFjenh3eWFvZHRjZG1kbWx3Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3NTE1MTg2MzIsImV4cCI6MjA2NzA5NDYzMn0.nBzqGuK9j6RZ6mOPWU2boAC_5H9XDs-fPpo5P3WZYbI"; - -inline CactusTelemetry& CactusTelemetry::getInstance() { - static CactusTelemetry instance; - return instance; -} - -inline CactusTelemetry::CactusTelemetry() { - // Device ID and project ID are now initialized lazily -} - -inline void CactusTelemetry::setEnabled(bool enabled) { - std::lock_guard lock(mutex_); - enabled_ = enabled; -} - -inline void CactusTelemetry::setTelemetryToken(const std::string& token) { - std::lock_guard lock(mutex_); - telemetry_token_ = token; -} - -inline void CactusTelemetry::setProjectId(const std::string& project_id) { - std::lock_guard lock(mutex_); - project_id_ = project_id; -} - -inline void CactusTelemetry::ensureInitialized() { - std::lock_guard lock(mutex_); - if (!initialized_) { - std::string pro_key = DeviceManager::getProKey(); - if (!pro_key.empty()) { - get_device_id(pro_key.c_str()); - } -#ifdef CACTUS_TELEMETRY_ENABLED - device_id_ = DeviceManager::getDeviceId(); - project_id_ = DeviceManager::getProjectId(); -#endif - initialized_ = true; - } -} - -inline bool CactusTelemetry::isEnabled() const { - std::lock_guard lock(mutex_); - return enabled_ && !telemetry_token_.empty(); -} - -inline void CactusTelemetry::sendToSupabase([[maybe_unused]] const TelemetryMetrics& metrics) { -#ifdef CACTUS_TELEMETRY_ENABLED - std::string telemetry_token; - std::string project_id; - std::string device_id; - { - std::lock_guard lock(mutex_); - telemetry_token = telemetry_token_; - project_id = project_id_; - device_id = device_id_; - } - - std::string log_json = LogRecord::buildJson(metrics, project_id, device_id, telemetry_token); - - std::string payload = "[" + log_json + "]"; - - std::map headers; - headers["apikey"] = SUPABASE_KEY; - headers["Authorization"] = "Bearer " + SUPABASE_KEY; - headers["Content-Type"] = "application/json"; - headers["Prefer"] = "return=minimal"; - headers["Content-Profile"] = "cactus"; - - std::string url = SUPABASE_URL + "/rest/v1/logs"; - HttpClient::postJson(url, headers, payload); -#else - (void)metrics; -#endif -} - -inline void CactusTelemetry::recordEvent(const TelemetryMetrics& metrics) { - if (!isEnabled()) { - return; - } - std::thread([this, metrics]() { - sendToSupabase(metrics); - }).detach(); -} - -inline void CactusTelemetry::recordInit(const std::string& model, bool success, - const std::string& message) { - TelemetryMetrics metrics; - metrics.event_type = TelemetryEventType::Init; - metrics.model = model; - metrics.success = success; - metrics.message = message; - metrics.timestamp = std::chrono::system_clock::now(); - - recordEvent(metrics); -} - -inline void CactusTelemetry::recordCompletion(const std::string& model, bool success, - double ttft_ms, double tps, double response_time_ms, - int tokens, const std::string& message) { - TelemetryMetrics metrics; - metrics.event_type = TelemetryEventType::Completion; - metrics.model = model; - metrics.success = success; - metrics.ttft_ms = ttft_ms; - metrics.tps = tps; - metrics.response_time_ms = response_time_ms; - metrics.tokens = tokens; - metrics.message = message; - metrics.timestamp = std::chrono::system_clock::now(); - - recordEvent(metrics); -} - -inline void CactusTelemetry::recordEmbedding(const std::string& model, bool success, - const std::string& message) { - TelemetryMetrics metrics; - metrics.event_type = TelemetryEventType::Embedding; - metrics.model = model; - metrics.success = success; - metrics.message = message; - metrics.timestamp = std::chrono::system_clock::now(); - - recordEvent(metrics); -} - -inline void CactusTelemetry::recordTranscription(const std::string& model, bool success, - double ttft_ms, double tps, double response_time_ms, - int tokens, const std::string& message) { - TelemetryMetrics metrics; - metrics.event_type = TelemetryEventType::Transcription; - metrics.model = model; - metrics.success = success; - metrics.response_time_ms = response_time_ms; - metrics.ttft_ms = ttft_ms; - metrics.tps = tps; - metrics.tokens = tokens; - metrics.message = message; - metrics.timestamp = std::chrono::system_clock::now(); - - recordEvent(metrics); -} - -} // namespace ffi -} // namespace cactus - -#endif // CACTUS_TELEMETRY_H diff --git a/cactus/ffi/cactus_transcribe.cpp b/cactus/ffi/cactus_transcribe.cpp deleted file mode 100644 index fb935f364..000000000 --- a/cactus/ffi/cactus_transcribe.cpp +++ /dev/null @@ -1,307 +0,0 @@ -#include "cactus_ffi.h" -#include "cactus_utils.h" -#include "cactus_telemetry.h" -#include "../../libs/audio/wav.h" -#include -#include -#include -#include - -using namespace cactus::engine; -using namespace cactus::ffi; - -static constexpr size_t WHISPER_TARGET_FRAMES = 3000; -static constexpr int WHISPER_SAMPLE_RATE = 16000; -static constexpr size_t WHISPER_MAX_DECODER_POSITIONS = 448; - -static AudioProcessor::SpectrogramConfig get_whisper_spectrogram_config() { - AudioProcessor::SpectrogramConfig cfg{}; - cfg.n_fft = 400; - cfg.frame_length = 400; - cfg.hop_length = 160; - cfg.power = 2.0f; - cfg.center = true; - cfg.pad_mode = "reflect"; - cfg.onesided = true; - cfg.dither = 0.0f; - cfg.mel_floor = 1e-10f; - cfg.log_mel = "log10"; - cfg.reference = 1.0f; - cfg.min_value = 1e-10f; - cfg.remove_dc_offset = true; - return cfg; -} - -static std::vector normalize_mel(std::vector& mel, size_t n_mels) { - size_t n_frames = mel.size() / n_mels; - - float max_val = -std::numeric_limits::infinity(); - for (float v : mel) - if (v > max_val) max_val = v; - - float min_allowed = max_val - 8.0f; - for (float& v : mel) { - if (v < min_allowed) v = min_allowed; - v = (v + 4.0f) / 4.0f; - } - - if (n_frames != WHISPER_TARGET_FRAMES) { - std::vector fixed(n_mels * WHISPER_TARGET_FRAMES, 0.0f); - size_t copy_frames = std::min(n_frames, WHISPER_TARGET_FRAMES); - for (size_t m = 0; m < n_mels; ++m) { - const float* src = &mel[m * n_frames]; - float* dst = &fixed[m * WHISPER_TARGET_FRAMES]; - std::copy(src, src + copy_frames, dst); - } - return fixed; - } - return mel; -} - -static std::vector compute_whisper_mel_from_pcm(const int16_t* pcm_samples, size_t num_samples, int sample_rate_in) { - if (!pcm_samples || num_samples == 0) return {}; - - std::vector waveform_fp32(num_samples); - for (size_t i = 0; i < num_samples; i++) - waveform_fp32[i] = static_cast(pcm_samples[i]) / 32768.0f; - - std::vector waveform_16k = resample_to_16k_fp32(waveform_fp32, sample_rate_in); - if (waveform_16k.empty()) return {}; - - auto cfg = get_whisper_spectrogram_config(); - const size_t num_mel_filters = 80; - const size_t num_frequency_bins = cfg.n_fft / 2 + 1; - - AudioProcessor ap; - ap.init_mel_filters(num_frequency_bins, num_mel_filters, 0.0f, 8000.0f, WHISPER_SAMPLE_RATE); - std::vector mel = ap.compute_spectrogram(waveform_16k, cfg); - - if (mel.empty()) return mel; - return normalize_mel(mel, num_mel_filters); -} - -static std::vector compute_whisper_mel_from_wav(const std::string& wav_path) { - AudioFP32 audio = load_wav(wav_path); - std::vector waveform_16k = resample_to_16k_fp32(audio.samples, audio.sample_rate); - - auto cfg = get_whisper_spectrogram_config(); - const size_t num_mel_filters = 80; - const size_t num_frequency_bins = cfg.n_fft / 2 + 1; - - AudioProcessor ap; - ap.init_mel_filters(num_frequency_bins, num_mel_filters, 0.0f, 8000.0f, WHISPER_SAMPLE_RATE); - std::vector mel = ap.compute_spectrogram(waveform_16k, cfg); - - if (mel.empty()) return mel; - return normalize_mel(mel, num_mel_filters); -} - -extern "C" { - -int cactus_transcribe( - cactus_model_t model, - const char* audio_file_path, - const char* prompt, - char* response_buffer, - size_t buffer_size, - const char* options_json, - cactus_token_callback callback, - void* user_data, - const uint8_t* pcm_buffer, - size_t pcm_buffer_size -) { - if (!model) { - std::string error_msg = last_error_message.empty() ? "Model not initialized." : last_error_message; - CACTUS_LOG_ERROR("transcribe", error_msg); - handle_error_response(error_msg, response_buffer, buffer_size); - return -1; - } - - if (!prompt || !response_buffer || buffer_size == 0) { - CACTUS_LOG_ERROR("transcribe", "Invalid parameters: prompt, response_buffer, or buffer_size"); - handle_error_response("Invalid parameters", response_buffer, buffer_size); - return -1; - } - - if (!audio_file_path && (!pcm_buffer || pcm_buffer_size == 0)) { - CACTUS_LOG_ERROR("transcribe", "No audio input provided"); - handle_error_response("Either audio_file_path or pcm_buffer must be provided", response_buffer, buffer_size); - return -1; - } - - if (audio_file_path && pcm_buffer && pcm_buffer_size > 0) { - CACTUS_LOG_ERROR("transcribe", "Both audio_file_path and pcm_buffer provided"); - handle_error_response("Cannot provide both audio_file_path and pcm_buffer", response_buffer, buffer_size); - return -1; - } - - if (pcm_buffer && pcm_buffer_size > 0 && (pcm_buffer_size < 2 || pcm_buffer_size % 2 != 0)) { - CACTUS_LOG_ERROR("transcribe", "Invalid pcm_buffer_size: " << pcm_buffer_size); - handle_error_response("pcm_buffer_size must be even and at least 2 bytes", response_buffer, buffer_size); - return -1; - } - - try { - auto start_time = std::chrono::high_resolution_clock::now(); - auto* handle = static_cast(model); - std::lock_guard lock(handle->model_mutex); - handle->should_stop = false; - - float temperature, top_p, confidence_threshold; - size_t top_k, max_tokens, tool_rag_top_k; - std::vector stop_sequences; - bool force_tools; - parse_options_json(options_json ? options_json : "", temperature, top_p, top_k, max_tokens, stop_sequences, force_tools, tool_rag_top_k, confidence_threshold); - - std::vector mel_bins; - if (audio_file_path == nullptr) { - const int16_t* pcm_samples = reinterpret_cast(pcm_buffer); - size_t num_samples = pcm_buffer_size / 2; - mel_bins = compute_whisper_mel_from_pcm(pcm_samples, num_samples, WHISPER_SAMPLE_RATE); - } else { - mel_bins = compute_whisper_mel_from_wav(audio_file_path); - } - - if (mel_bins.empty()) { - CACTUS_LOG_ERROR("transcribe", "Computed mel spectrogram is empty"); - handle_error_response("Computed mel spectrogram is empty", response_buffer, buffer_size); - return -1; - } - - CACTUS_LOG_DEBUG("transcribe", "Mel spectrogram computed, size: " << mel_bins.size()); - - auto* tokenizer = handle->model->get_tokenizer(); - if (!tokenizer) { - CACTUS_LOG_ERROR("transcribe", "Tokenizer unavailable"); - handle_error_response("Tokenizer unavailable", response_buffer, buffer_size); - return -1; - } - - std::vector tokens = tokenizer->encode(std::string(prompt)); - if (tokens.empty()) { - CACTUS_LOG_ERROR("transcribe", "Decoder input tokens empty after encoding prompt"); - handle_error_response("Decoder input tokens empty", response_buffer, buffer_size); - return -1; - } - - size_t max_allowed_tokens = WHISPER_MAX_DECODER_POSITIONS - tokens.size(); - if (max_tokens > max_allowed_tokens) { - CACTUS_LOG_WARN("transcribe", "max_tokens exceeds limit, reducing to " << max_allowed_tokens); - max_tokens = max_allowed_tokens; - } - - std::vector> stop_token_sequences; - stop_token_sequences.push_back({ tokenizer->get_eos_token() }); - - double time_to_first_token = 0.0; - size_t completion_tokens = 0; - std::vector generated_tokens; - std::string final_text; - - uint32_t next_token = handle->model->decode_with_audio(tokens, mel_bins, temperature, top_p, top_k); - { - auto t_first = std::chrono::high_resolution_clock::now(); - time_to_first_token = std::chrono::duration_cast(t_first - start_time).count() / 1000.0; - } - - generated_tokens.push_back(next_token); - tokens.push_back(next_token); - completion_tokens++; - - std::string piece = tokenizer->decode({ next_token }); - final_text += piece; - if (callback) callback(piece.c_str(), next_token, user_data); - - if (!matches_stop_sequence(generated_tokens, stop_token_sequences)) { - for (size_t i = 1; i < max_tokens; ++i) { - if (handle->should_stop) break; - - next_token = handle->model->decode_with_audio(tokens, mel_bins, temperature, top_p, top_k); - generated_tokens.push_back(next_token); - tokens.push_back(next_token); - completion_tokens++; - - piece = tokenizer->decode({ next_token }); - final_text += piece; - if (callback) callback(piece.c_str(), next_token, user_data); - - if (matches_stop_sequence(generated_tokens, stop_token_sequences)) break; - } - } - - auto end_time = std::chrono::high_resolution_clock::now(); - double total_time_ms = std::chrono::duration_cast(end_time - start_time).count() / 1000.0; - double decode_time_ms = std::max(0.0, total_time_ms - time_to_first_token); - - size_t prompt_tokens = 0; - if (!tokens.empty() && completion_tokens <= tokens.size()) - prompt_tokens = tokens.size() - completion_tokens; - - double prefill_tps = time_to_first_token > 0 ? (prompt_tokens * 1000.0) / time_to_first_token : 0.0; - double decode_tps = (completion_tokens > 1 && decode_time_ms > 0.0) ? ((completion_tokens - 1) * 1000.0) / decode_time_ms : 0.0; - - std::string cleaned_text = final_text; - const std::string token_to_remove = "<|startoftranscript|>"; - size_t pos = 0; - while ((pos = cleaned_text.find(token_to_remove, pos)) != std::string::npos) { - cleaned_text.erase(pos, token_to_remove.length()); - } - - std::string json = construct_response_json(cleaned_text, {}, time_to_first_token, total_time_ms, prefill_tps, decode_tps, prompt_tokens, completion_tokens); - - if (json.size() >= buffer_size) { - handle_error_response("Response buffer too small", response_buffer, buffer_size); - return -1; - } - - std::strcpy(response_buffer, json.c_str()); - - CactusTelemetry::getInstance().recordTranscription( - handle->model_name, - true, - time_to_first_token, - decode_tps, - total_time_ms, - completion_tokens, - "" - ); - - return static_cast(json.size()); - } - catch (const std::exception& e) { - CACTUS_LOG_ERROR("transcribe", "Exception: " << e.what()); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordTranscription( - handle->model_name, - false, - 0.0, - 0, - 0.0, - 0, - std::string(e.what()) - ); - - handle_error_response(e.what(), response_buffer, buffer_size); - return -1; - } - catch (...) { - CACTUS_LOG_ERROR("transcribe", "Unknown exception during transcription"); - - auto* handle = static_cast(model); - CactusTelemetry::getInstance().recordTranscription( - handle->model_name, - false, - 0.0, - 0, - 0.0, - 0, - "Unknown error in transcribe" - ); - - handle_error_response("Unknown error in transcribe", response_buffer, buffer_size); - return -1; - } -} - -} diff --git a/cactus/ffi/cactus_utils.h b/cactus/ffi/cactus_utils.h deleted file mode 100644 index 0c481c5c3..000000000 --- a/cactus/ffi/cactus_utils.h +++ /dev/null @@ -1,634 +0,0 @@ -#ifndef CACTUS_UTILS_H -#define CACTUS_UTILS_H - -#include "../engine/engine.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __APPLE__ -#include -#include -#elif defined(_WIN32) -#include -#include -#elif defined(__linux__) || defined(__ANDROID__) -#include -#endif - -inline size_t get_memory_footprint_bytes() { -#ifdef __APPLE__ - task_vm_info_data_t vm_info; - mach_msg_type_number_t count = TASK_VM_INFO_COUNT; - if (task_info(mach_task_self(), TASK_VM_INFO, (task_info_t)&vm_info, &count) == KERN_SUCCESS) - return vm_info.phys_footprint; -#elif defined(_WIN32) - PROCESS_MEMORY_COUNTERS_EX pmc; - if (GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS*)&pmc, sizeof(pmc))) - return pmc.PrivateUsage; -#elif defined(__linux__) || defined(__ANDROID__) - std::ifstream statm("/proc/self/statm"); - if (statm.is_open()) { - size_t size, resident; - statm >> size >> resident; - return resident * sysconf(_SC_PAGESIZE); - } -#endif - return 0; -} - -inline double get_ram_usage_mb() { - return get_memory_footprint_bytes() / (1024.0 * 1024.0); -} - -struct CactusModelHandle { - std::unique_ptr model; - std::atomic should_stop; - std::vector processed_tokens; - std::mutex model_mutex; - std::string model_name; - std::unique_ptr corpus_index; - std::string corpus_dir; - size_t corpus_embedding_dim = 0; - std::vector> tool_embeddings; - std::vector tool_texts; - - CactusModelHandle() : should_stop(false) {} -}; - -extern std::string last_error_message; - -bool matches_stop_sequence(const std::vector& generated_tokens, - const std::vector>& stop_sequences); - -std::string retrieve_rag_context(CactusModelHandle* handle, const std::string& query); - -namespace cactus { -namespace ffi { - -#ifndef CACTUS_VERSION -#define CACTUS_VERSION "unknown" -#endif - -inline const char* getVersion() { - return CACTUS_VERSION; -} - -inline std::string generateUUID() { -#ifdef __APPLE__ - uuid_t uuid; - uuid_generate_random(uuid); - char uuid_str[37]; - uuid_unparse_lower(uuid, uuid_str); - return std::string(uuid_str); -#endif -} - -struct ToolFunction { - std::string name; - std::string description; - std::unordered_map parameters; -}; - -} // namespace ffi -} // namespace cactus - -std::vector select_relevant_tools( - CactusModelHandle* handle, - const std::string& query, - const std::vector& all_tools, - size_t top_k); - -#include "gemma_tools.h" - -namespace cactus { -namespace ffi { - -inline std::string escape_json_string(const std::string& s) { - std::ostringstream o; - for (char c : s) { - if (c == '"') o << "\\\""; - else if (c == '\n') o << "\\n"; - else if (c == '\r') o << "\\r"; - else if (c == '\t') o << "\\t"; - else if (c == '\\') o << "\\\\"; - else o << c; - } - return o.str(); -} - -inline void handle_error_response(const std::string& error_message, char* response_buffer, size_t buffer_size) { - std::ostringstream json; - json << "{"; - json << "\"success\":false,"; - json << "\"error\":\"" << escape_json_string(error_message) << "\","; - json << "\"cloud_handoff\":false,"; - json << "\"response\":null,"; - json << "\"function_calls\":[],"; - json << "\"confidence\":0.0,"; - json << "\"time_to_first_token_ms\":0.0,"; - json << "\"total_time_ms\":0.0,"; - json << "\"prefill_tps\":0.0,"; - json << "\"decode_tps\":0.0,"; - json << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() << ","; - json << "\"prefill_tokens\":0,"; - json << "\"decode_tokens\":0,"; - json << "\"total_tokens\":0"; - json << "}"; - std::string error_json = json.str(); - if (response_buffer && error_json.length() < buffer_size) { - std::strcpy(response_buffer, error_json.c_str()); - } -} - -inline std::vector parse_messages_json(const std::string& json, - std::vector& out_image_paths) { - std::vector messages; - out_image_paths.clear(); - - size_t pos = json.find('['); - if (pos == std::string::npos) { - throw std::runtime_error("Invalid JSON: expected array"); - } - - pos = json.find('{', pos); - while (pos != std::string::npos) { - cactus::engine::ChatMessage msg; - - size_t obj_start = pos; - int brace_count = 1; - size_t obj_end = obj_start + 1; - while (obj_end < json.length() && brace_count > 0) { - if (json[obj_end] == '{') brace_count++; - else if (json[obj_end] == '}') brace_count--; - obj_end++; - } - - size_t role_pos = json.find("\"role\"", pos); - if (role_pos == std::string::npos || role_pos >= obj_end) break; - - size_t role_start = json.find('"', role_pos + 6) + 1; - size_t role_end = json.find('"', role_start); - msg.role = json.substr(role_start, role_end - role_start); - - size_t content_pos = json.find("\"content\"", role_end); - if (content_pos != std::string::npos && content_pos < obj_end) { - size_t content_start = json.find('"', content_pos + 9) + 1; - size_t content_end = content_start; - - while (content_end < json.length()) { - content_end = json.find('"', content_end); - if (content_end == std::string::npos) break; - if (json[content_end - 1] != '\\') break; - content_end++; - } - - msg.content = json.substr(content_start, content_end - content_start); - - size_t escape_pos = 0; - while ((escape_pos = msg.content.find("\\n", escape_pos)) != std::string::npos) { - msg.content.replace(escape_pos, 2, "\n"); - escape_pos += 1; - } - escape_pos = 0; - while ((escape_pos = msg.content.find("\\\"", escape_pos)) != std::string::npos) { - msg.content.replace(escape_pos, 2, "\""); - escape_pos += 1; - } - } - - size_t images_pos = json.find("\"images\"", pos); - if (images_pos != std::string::npos && images_pos < obj_end) { - size_t array_start = json.find('[', images_pos); - if (array_start != std::string::npos && array_start < obj_end) { - size_t array_end = json.find(']', array_start); - if (array_end != std::string::npos && array_end < obj_end) { - size_t img_pos = array_start; - while (true) { - img_pos = json.find('"', img_pos + 1); - if (img_pos == std::string::npos || img_pos >= array_end) break; - - size_t img_start = img_pos + 1; - size_t img_end = json.find('"', img_start); - if (img_end == std::string::npos || img_end > array_end) break; - - std::string img_path = json.substr(img_start, img_end - img_start); - - std::filesystem::path p(img_path); - img_path = std::filesystem::absolute(p).string(); - - msg.images.push_back(img_path); - out_image_paths.push_back(img_path); - img_pos = img_end; - } - } - } - } - - messages.push_back(msg); - - pos = json.find('{', obj_end); - } - - return messages; -} - -inline std::vector parse_tools_json(const std::string& json) { - std::vector tools; - - if (json.empty()) return tools; - - size_t pos = json.find('['); - if (pos == std::string::npos) return tools; - - pos = json.find("\"function\"", pos); - while (pos != std::string::npos) { - ToolFunction tool; - - size_t name_pos = json.find("\"name\"", pos); - if (name_pos != std::string::npos) { - size_t name_start = json.find('"', name_pos + 6) + 1; - size_t name_end = json.find('"', name_start); - tool.name = json.substr(name_start, name_end - name_start); - } - - size_t desc_pos = json.find("\"description\"", pos); - if (desc_pos != std::string::npos) { - size_t desc_start = json.find('"', desc_pos + 13) + 1; - size_t desc_end = json.find('"', desc_start); - tool.description = json.substr(desc_start, desc_end - desc_start); - } - - size_t params_pos = json.find("\"parameters\"", pos); - if (params_pos != std::string::npos) { - size_t params_start = json.find('{', params_pos); - if (params_start != std::string::npos) { - int brace_count = 1; - size_t params_end = params_start + 1; - while (params_end < json.length() && brace_count > 0) { - if (json[params_end] == '{') brace_count++; - else if (json[params_end] == '}') brace_count--; - params_end++; - } - tool.parameters["schema"] = json.substr(params_start, params_end - params_start); - } - } - - tools.push_back(tool); - - pos = json.find("\"function\"", name_pos); - } - - return tools; -} - -inline void parse_options_json(const std::string& json, - float& temperature, float& top_p, - size_t& top_k, size_t& max_tokens, - std::vector& stop_sequences, - bool& force_tools, - size_t& tool_rag_top_k, - float& confidence_threshold) { - temperature = 0.0f; - top_p = 0.0f; - top_k = 0; - max_tokens = 100; - force_tools = false; - tool_rag_top_k = 2; // 0 = disabled, N = select top N relevant tools - confidence_threshold = 0.7f; // trigger cloud handoff when confidence < this value - stop_sequences.clear(); - - if (json.empty()) return; - - size_t pos = json.find("\"temperature\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - temperature = std::stof(json.substr(pos)); - } - - pos = json.find("\"top_p\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - top_p = std::stof(json.substr(pos)); - } - - pos = json.find("\"top_k\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - top_k = std::stoul(json.substr(pos)); - } - - pos = json.find("\"max_tokens\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - max_tokens = std::stoul(json.substr(pos)); - } - - pos = json.find("\"force_tools\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - while (pos < json.length() && std::isspace(json[pos])) pos++; - force_tools = (json.substr(pos, 4) == "true"); - } - - pos = json.find("\"tool_rag_top_k\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - tool_rag_top_k = std::stoul(json.substr(pos)); - } - - pos = json.find("\"confidence_threshold\""); - if (pos != std::string::npos) { - pos = json.find(':', pos) + 1; - confidence_threshold = std::stof(json.substr(pos)); - } - - pos = json.find("\"stop_sequences\""); - if (pos != std::string::npos) { - pos = json.find('[', pos); - if (pos != std::string::npos) { - size_t end_pos = json.find(']', pos); - size_t seq_pos = json.find('"', pos); - - while (seq_pos != std::string::npos && seq_pos < end_pos) { - size_t seq_start = seq_pos + 1; - size_t seq_end = json.find('"', seq_start); - if (seq_end != std::string::npos) { - stop_sequences.push_back(json.substr(seq_start, seq_end - seq_start)); - } - seq_pos = json.find('"', seq_end + 1); - } - } - } -} - -inline std::string format_tools_for_prompt(const std::vector& tools) { - if (tools.empty()) return ""; - std::string formatted_tools_json; - for (size_t i = 0; i < tools.size(); i++) { - if (i > 0) formatted_tools_json += "\n"; - formatted_tools_json += "{\"type\":\"function\",\"function\":{\"name\":\"" - + tools[i].name - + "\",\"description\":\"" - + tools[i].description + "\""; - if (tools[i].parameters.find("schema") != tools[i].parameters.end()) { - formatted_tools_json += ",\"parameters\":" + tools[i].parameters.at("schema"); - } - formatted_tools_json += "}}"; - } - return formatted_tools_json; -} - -inline void parse_function_calls_from_response(const std::string& response_text, - std::string& regular_response, - std::vector& function_calls) { - regular_response = response_text; - function_calls.clear(); - - gemma::parse_function_calls(regular_response, function_calls); - - // Parse Qwen-style function calls: {"name": "...", "arguments": {...}} - const std::string QWEN_TOOL_START = ""; - const std::string QWEN_TOOL_END = ""; - size_t qwen_start_pos = 0; - - while ((qwen_start_pos = regular_response.find(QWEN_TOOL_START, qwen_start_pos)) != std::string::npos) { - size_t content_start = qwen_start_pos + QWEN_TOOL_START.length(); - size_t qwen_end_pos = regular_response.find(QWEN_TOOL_END, content_start); - - if (qwen_end_pos != std::string::npos) { - std::string json_content = regular_response.substr(content_start, qwen_end_pos - content_start); - - size_t first = json_content.find_first_not_of(" \t\n\r"); - size_t last = json_content.find_last_not_of(" \t\n\r"); - if (first != std::string::npos && last != std::string::npos) { - json_content = json_content.substr(first, last - first + 1); - } - - if (json_content.size() > 2 && json_content[0] == '{' && - json_content.find("\"name\"") != std::string::npos) { - function_calls.push_back(json_content); - } - - regular_response.erase(qwen_start_pos, qwen_end_pos + QWEN_TOOL_END.length() - qwen_start_pos); - } else { - break; - } - } - - // Parse LFM2-style function calls: <|tool_call_start|>[name(args)]<|tool_call_end|> - const std::string TOOL_CALL_START = "<|tool_call_start|>"; - const std::string TOOL_CALL_END = "<|tool_call_end|>"; - size_t tool_start_pos = 0; - - while ((tool_start_pos = regular_response.find(TOOL_CALL_START, tool_start_pos)) != std::string::npos) { - size_t content_start = tool_start_pos + TOOL_CALL_START.length(); - size_t tool_end_pos = response_text.find(TOOL_CALL_END, content_start); - - if (tool_end_pos != std::string::npos) { - std::string tool_content = response_text.substr(content_start, tool_end_pos - content_start); - - if (tool_content.size() > 2 && tool_content[0] == '[' && tool_content[tool_content.size()-1] == ']') { - tool_content = tool_content.substr(1, tool_content.size() - 2); - - size_t paren_pos = tool_content.find('('); - if (paren_pos != std::string::npos) { - std::string func_name = tool_content.substr(0, paren_pos); - std::string args_str = tool_content.substr(paren_pos + 1); - - if (!args_str.empty() && args_str.back() == ')') { - args_str.pop_back(); - } - - std::string json_call = "{\"name\":\"" + func_name + "\",\"arguments\":{"; - - size_t arg_pos = 0; - bool first_arg = true; - while (arg_pos < args_str.length()) { - while (arg_pos < args_str.length() && std::isspace(args_str[arg_pos])) arg_pos++; - - size_t eq_pos = args_str.find('=', arg_pos); - if (eq_pos == std::string::npos) break; - - std::string arg_name = args_str.substr(arg_pos, eq_pos - arg_pos); - - size_t val_start = eq_pos + 1; - size_t val_end = val_start; - - if (val_start < args_str.length() && args_str[val_start] == '"') { - val_start++; - val_end = args_str.find('"', val_start); - if (val_end == std::string::npos) break; - } else { - val_end = args_str.find(',', val_start); - if (val_end == std::string::npos) val_end = args_str.length(); - } - - std::string arg_value = args_str.substr(val_start, val_end - val_start); - - if (!first_arg) json_call += ","; - json_call += "\"" + arg_name + "\":\"" + arg_value + "\""; - first_arg = false; - - arg_pos = args_str.find(',', val_end); - if (arg_pos != std::string::npos) { - arg_pos++; - } else { - break; - } - } - - json_call += "}}"; - function_calls.push_back(json_call); - } - } - - regular_response.erase(tool_start_pos, tool_end_pos + TOOL_CALL_END.length() - tool_start_pos); - tool_start_pos = tool_end_pos + TOOL_CALL_END.length(); - } else { - break; - } - } - - const char* FUNCTION_CALL_MARKER = "\"function_call\""; - size_t search_pos = 0; - const size_t text_len = regular_response.length(); - - while (search_pos < text_len) { - size_t marker_pos = regular_response.find(FUNCTION_CALL_MARKER, search_pos); - if (marker_pos == std::string::npos) break; - - size_t json_start = regular_response.find('{', marker_pos); - if (json_start == std::string::npos) break; - - int brace_count = 1; - size_t json_end = json_start + 1; - while (json_end < text_len && brace_count > 0) { - char c = regular_response[json_end]; - brace_count += (c == '{') - (c == '}'); - json_end++; - } - - if (brace_count == 0) { - function_calls.push_back(regular_response.substr(json_start, json_end - json_start)); - regular_response = regular_response.substr(0, marker_pos); - size_t last_bracket = regular_response.rfind('{'); - if(last_bracket != std::string::npos) { - regular_response = regular_response.substr(0, last_bracket); - } - } - search_pos = json_end; - } -} - -inline std::string construct_response_json(const std::string& regular_response, - const std::vector& function_calls, - double time_to_first_token, - double total_time_ms, - double prefill_tps, - double decode_tps, - size_t prompt_tokens, - size_t completion_tokens, - float confidence = 0.0f, - bool cloud_handoff = false) { - std::ostringstream json; - json << "{"; - json << "\"success\":" << (cloud_handoff ? "false" : "true") << ","; - json << "\"error\":null,"; - json << "\"cloud_handoff\":" << (cloud_handoff ? "true" : "false") << ","; - json << "\"response\":\"" << escape_json_string(regular_response) << "\","; - json << "\"function_calls\":["; - for (size_t i = 0; i < function_calls.size(); ++i) { - if (i > 0) json << ","; - json << function_calls[i]; - } - json << "],"; - json << "\"confidence\":" << std::fixed << std::setprecision(4) << confidence << ","; - json << "\"time_to_first_token_ms\":" << std::fixed << std::setprecision(2) << time_to_first_token << ","; - json << "\"total_time_ms\":" << std::fixed << std::setprecision(2) << total_time_ms << ","; - json << "\"prefill_tps\":" << std::fixed << std::setprecision(2) << prefill_tps << ","; - json << "\"decode_tps\":" << std::fixed << std::setprecision(2) << decode_tps << ","; - json << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() << ","; - json << "\"prefill_tokens\":" << prompt_tokens << ","; - json << "\"decode_tokens\":" << completion_tokens << ","; - json << "\"total_tokens\":" << (prompt_tokens + completion_tokens); - json << "}"; - return json.str(); -} - -inline std::string construct_cloud_handoff_json(float confidence, - double time_to_first_token, - double prefill_tps, - size_t prompt_tokens) { - std::ostringstream json; - json << "{"; - json << "\"success\":false,"; - json << "\"error\":null,"; - json << "\"cloud_handoff\":true,"; - json << "\"response\":null,"; - json << "\"function_calls\":[],"; - json << "\"confidence\":" << std::fixed << std::setprecision(4) << confidence << ","; - json << "\"time_to_first_token_ms\":" << std::fixed << std::setprecision(2) << time_to_first_token << ","; - json << "\"total_time_ms\":" << std::fixed << std::setprecision(2) << time_to_first_token << ","; - json << "\"prefill_tps\":" << std::fixed << std::setprecision(2) << prefill_tps << ","; - json << "\"decode_tps\":0.0,"; - json << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() << ","; - json << "\"prefill_tokens\":" << prompt_tokens << ","; - json << "\"decode_tokens\":0,"; - json << "\"total_tokens\":" << prompt_tokens; - json << "}"; - return json.str(); -} - -} // namespace ffi -} // namespace cactus - -#ifdef __cplusplus -extern "C" { -#endif - -const char* cactus_get_last_error(); - -__attribute__((weak)) -const char* register_app(const char* encrypted_data); - -__attribute__((weak)) -const char* get_device_id(const char* current_token); - -#ifdef __cplusplus -} -#endif - -#ifdef __cplusplus -extern "C" { - -__attribute__((weak)) -inline const char* register_app(const char* encrypted_data) { - (void)encrypted_data; - static thread_local std::string uuid_storage; - uuid_storage = cactus::ffi::generateUUID(); - return uuid_storage.c_str(); -} - -__attribute__((weak)) -inline const char* get_device_id(const char* current_token) { - (void)current_token; - static thread_local std::string uuid_storage; - uuid_storage = cactus::ffi::generateUUID(); - return uuid_storage.c_str(); -} -} -#endif - -#endif // CACTUS_UTILS_H diff --git a/cactus/ffi/gemma_tools.h b/cactus/ffi/gemma_tools.h deleted file mode 100644 index 912de57b8..000000000 --- a/cactus/ffi/gemma_tools.h +++ /dev/null @@ -1,549 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include - -namespace gemma { - -inline std::string to_upper(const std::string& s) { - std::string result = s; - for (auto& c : result) c = std::toupper(c); - return result; -} - -inline std::string escape(const std::string& s) { - return "" + s + ""; -} - -inline void skip_whitespace(const std::string& json, size_t& pos) { - while (pos < json.length() && std::isspace(json[pos])) pos++; -} - -inline std::string extract_json_string(const std::string& json, size_t& pos) { - std::string value; - while (pos < json.length() && json[pos] != '"') { - if (json[pos] == '\\' && pos + 1 < json.length()) { - pos++; - if (json[pos] == 'n') value += '\n'; - else if (json[pos] == 't') value += '\t'; - else if (json[pos] == 'r') value += '\r'; - else if (json[pos] == '"') value += '"'; - else if (json[pos] == '\\') value += '\\'; - else value += json[pos]; - } else { - value += json[pos]; - } - pos++; - } - if (pos < json.length()) pos++; - return value; -} - -std::string format_argument(const std::string& json, size_t& pos, bool escape_keys); -std::string format_parameters(const std::string& properties_json, const std::string& /*required_json*/); - -inline std::string format_argument(const std::string& json, size_t& pos, bool escape_keys = true) { - skip_whitespace(json, pos); - if (pos >= json.length()) return ""; - - char c = json[pos]; - - if (c == '"') { - std::string value = extract_json_string(json, pos); - return escape(value); - } else if (c == '{') { - std::string result = "{"; - pos++; - bool first = true; - - while (pos < json.length()) { - skip_whitespace(json, pos); - if (pos >= json.length() || json[pos] == '}') { pos++; break; } - if (json[pos] == ',') { pos++; continue; } - - if (json[pos] != '"') break; - pos++; - std::string key = extract_json_string(json, pos); - - skip_whitespace(json, pos); - if (pos < json.length() && json[pos] == ':') pos++; - - std::string value = format_argument(json, pos, escape_keys); - - if (!first) result += ","; - first = false; - if (escape_keys) { - result += escape(key) + ":" + value; - } else { - result += key + ":" + value; - } - } - result += "}"; - return result; - } else if (c == '[') { - std::string result = "["; - pos++; - bool first = true; - - while (pos < json.length()) { - skip_whitespace(json, pos); - if (pos >= json.length() || json[pos] == ']') { pos++; break; } - if (json[pos] == ',') { pos++; continue; } - - std::string value = format_argument(json, pos, escape_keys); - - if (!first) result += ","; - first = false; - result += value; - } - result += "]"; - return result; - } else if (json.compare(pos, 4, "true") == 0) { - pos += 4; - return "true"; - } else if (json.compare(pos, 5, "false") == 0) { - pos += 5; - return "false"; - } else if (json.compare(pos, 4, "null") == 0) { - pos += 4; - return "null"; - } else { - size_t start = pos; - while (pos < json.length() && (std::isdigit(json[pos]) || json[pos] == '.' || - json[pos] == '-' || json[pos] == '+' || json[pos] == 'e' || json[pos] == 'E')) { - pos++; - } - return json.substr(start, pos - start); - } -} - -inline std::map parse_json_object_raw(const std::string& json, size_t& pos) { - std::map result; - skip_whitespace(json, pos); - if (pos >= json.length() || json[pos] != '{') return result; - pos++; - - while (pos < json.length()) { - skip_whitespace(json, pos); - if (pos >= json.length() || json[pos] == '}') { pos++; break; } - if (json[pos] == ',') { pos++; continue; } - - if (json[pos] != '"') break; - pos++; - std::string key = extract_json_string(json, pos); - - skip_whitespace(json, pos); - if (pos < json.length() && json[pos] == ':') pos++; - skip_whitespace(json, pos); - - size_t value_start = pos; - if (json[pos] == '"') { - pos++; - while (pos < json.length() && json[pos] != '"') { - if (json[pos] == '\\') pos++; - pos++; - } - pos++; - } else if (json[pos] == '{') { - int depth = 1; - pos++; - while (pos < json.length() && depth > 0) { - if (json[pos] == '{') depth++; - else if (json[pos] == '}') depth--; - else if (json[pos] == '"') { - pos++; - while (pos < json.length() && json[pos] != '"') { - if (json[pos] == '\\') pos++; - pos++; - } - } - pos++; - } - } else if (json[pos] == '[') { - int depth = 1; - pos++; - while (pos < json.length() && depth > 0) { - if (json[pos] == '[') depth++; - else if (json[pos] == ']') depth--; - else if (json[pos] == '"') { - pos++; - while (pos < json.length() && json[pos] != '"') { - if (json[pos] == '\\') pos++; - pos++; - } - } - pos++; - } - } else { - while (pos < json.length() && json[pos] != ',' && json[pos] != '}') pos++; - } - result[key] = json.substr(value_start, pos - value_start); - } - return result; -} - -inline std::string get_json_string_value(const std::string& json, size_t pos) { - skip_whitespace(json, pos); - if (pos < json.length() && json[pos] == '"') { - pos++; - return extract_json_string(json, pos); - } - return ""; -} - -inline std::string format_parameters(const std::string& properties_json, const std::string& /*required_json*/) { - static const std::set standard_keys = {"description", "type", "properties", "required", "nullable"}; - - size_t pos = 0; - auto properties = parse_json_object_raw(properties_json, pos); - - std::string result; - bool first = true; - - for (const auto& [key, value_json] : properties) { - if (standard_keys.count(key)) continue; - - if (!first) result += ","; - first = false; - - size_t prop_pos = 0; - auto prop_obj = parse_json_object_raw(value_json, prop_pos); - - result += key + ":{"; - - if (prop_obj.count("description")) { - std::string desc = get_json_string_value(prop_obj["description"], 0); - result += "description:" + escape(desc); - } - - std::string type_val; - if (prop_obj.count("type")) { - type_val = get_json_string_value(prop_obj["type"], 0); - } - - if (to_upper(type_val) == "STRING") { - if (prop_obj.count("enum")) { - size_t enum_pos = 0; - std::string enum_formatted = format_argument(prop_obj["enum"], enum_pos, true); - result += ",enum:" + enum_formatted; - } - } else if (to_upper(type_val) == "OBJECT") { - if (prop_obj.count("properties")) { - std::string nested_required; - if (prop_obj.count("required")) { - nested_required = prop_obj["required"]; - } - result += ",properties:{" + format_parameters(prop_obj["properties"], nested_required) + "}"; - } - if (prop_obj.count("required")) { - result += ",required:["; - size_t req_pos = 0; - skip_whitespace(prop_obj["required"], req_pos); - if (req_pos < prop_obj["required"].length() && prop_obj["required"][req_pos] == '[') { - req_pos++; - bool req_first = true; - while (req_pos < prop_obj["required"].length()) { - skip_whitespace(prop_obj["required"], req_pos); - if (prop_obj["required"][req_pos] == ']') break; - if (prop_obj["required"][req_pos] == ',') { req_pos++; continue; } - if (prop_obj["required"][req_pos] == '"') { - req_pos++; - std::string req_item = extract_json_string(prop_obj["required"], req_pos); - if (!req_first) result += ","; - req_first = false; - result += escape(req_item); - } - } - } - result += "]"; - } - } else if (to_upper(type_val) == "ARRAY") { - if (prop_obj.count("items")) { - result += ",items:{"; - size_t items_pos = 0; - auto items_obj = parse_json_object_raw(prop_obj["items"], items_pos); - bool items_first = true; - - for (const auto& [item_key, item_value] : items_obj) { - if (!items_first) result += ","; - items_first = false; - - if (item_key == "properties") { - std::string items_required; - if (items_obj.count("required")) { - items_required = items_obj["required"]; - } - result += "properties:{" + format_parameters(item_value, items_required) + "}"; - } else if (item_key == "required") { - result += "required:["; - size_t req_pos = 0; - skip_whitespace(item_value, req_pos); - if (req_pos < item_value.length() && item_value[req_pos] == '[') { - req_pos++; - bool req_first = true; - while (req_pos < item_value.length()) { - skip_whitespace(item_value, req_pos); - if (item_value[req_pos] == ']') break; - if (item_value[req_pos] == ',') { req_pos++; continue; } - if (item_value[req_pos] == '"') { - req_pos++; - std::string req_item = extract_json_string(item_value, req_pos); - if (!req_first) result += ","; - req_first = false; - result += escape(req_item); - } - } - } - result += "]"; - } else if (item_key == "type") { - std::string item_type = get_json_string_value(item_value, 0); - result += "type:" + escape(to_upper(item_type)); - } else { - size_t val_pos = 0; - result += item_key + ":" + format_argument(item_value, val_pos, true); - } - } - result += "}"; - } - } - - if (!type_val.empty()) { - result += ",type:" + escape(to_upper(type_val)); - } - - result += "}"; - } - - return result; -} - -inline std::string format_function_declaration(const std::string& name, - const std::string& description, - const std::string& params_json) { - std::string result = "declaration:" + name + "{"; - result += "description:" + escape(description); - - if (!params_json.empty()) { - result += ",parameters:{"; - - size_t pos = 0; - auto params = parse_json_object_raw(params_json, pos); - - if (params.count("properties")) { - std::string required_json; - if (params.count("required")) { - required_json = params["required"]; - } - result += "properties:{" + format_parameters(params["properties"], required_json) + "}"; - } - - if (params.count("required")) { - result += ",required:["; - size_t req_pos = 0; - skip_whitespace(params["required"], req_pos); - if (req_pos < params["required"].length() && params["required"][req_pos] == '[') { - req_pos++; - bool first = true; - while (req_pos < params["required"].length()) { - skip_whitespace(params["required"], req_pos); - if (params["required"][req_pos] == ']') break; - if (params["required"][req_pos] == ',') { req_pos++; continue; } - if (params["required"][req_pos] == '"') { - req_pos++; - std::string item = extract_json_string(params["required"], req_pos); - if (!first) result += ","; - first = false; - result += escape(item); - } - } - } - result += "]"; - } - - if (params.count("type")) { - std::string type_val = get_json_string_value(params["type"], 0); - result += ",type:" + escape(to_upper(type_val)); - } - - result += "}"; - } - - result += "}"; - return result; -} - -template -inline std::string format_tools(const std::vector& tools) { - if (tools.empty()) return ""; - - std::string result; - for (const auto& tool : tools) { - result += ""; - std::string params_json; - auto it = tool.parameters.find("schema"); - if (it != tool.parameters.end()) { - params_json = it->second; - } - - result += format_function_declaration(tool.name, tool.description, params_json); - result += ""; - } - return result; -} - - -inline std::string unescape(const std::string& s) { - const std::string ESCAPE_TAG = ""; - std::string result = s; - size_t pos = 0; - while ((pos = result.find(ESCAPE_TAG, pos)) != std::string::npos) { - result.erase(pos, ESCAPE_TAG.length()); - } - return result; -} - -inline std::string args_to_json(const std::string& args_content) { - std::string result = "{"; - size_t pos = 0; - bool first = true; - - if (!args_content.empty() && args_content[0] == '{') pos = 1; - - while (pos < args_content.length()) { - while (pos < args_content.length() && std::isspace(args_content[pos])) pos++; - if (pos >= args_content.length() || args_content[pos] == '}') break; - if (args_content[pos] == ',') { pos++; continue; } - - size_t key_start = pos; - while (pos < args_content.length() && args_content[pos] != ':') pos++; - std::string key = args_content.substr(key_start, pos - key_start); - if (pos < args_content.length()) pos++; - - std::string value; - while (pos < args_content.length() && std::isspace(args_content[pos])) pos++; - - if (pos < args_content.length()) { - if (args_content.compare(pos, 8, "") == 0) { - pos += 8; - size_t val_end = args_content.find("", pos); - if (val_end != std::string::npos) { - value = "\"" + args_content.substr(pos, val_end - pos) + "\""; - pos = val_end + 8; - } - } else if (args_content[pos] == '{') { - int depth = 1; - size_t start = pos; - pos++; - while (pos < args_content.length() && depth > 0) { - if (args_content[pos] == '{') depth++; - else if (args_content[pos] == '}') depth--; - pos++; - } - value = args_to_json(args_content.substr(start, pos - start)); - } else if (args_content[pos] == '[') { - int depth = 1; - size_t start = pos; - pos++; - while (pos < args_content.length() && depth > 0) { - if (args_content[pos] == '[') depth++; - else if (args_content[pos] == ']') depth--; - pos++; - } - std::string arr_content = args_content.substr(start + 1, pos - start - 2); - value = "["; - size_t arr_pos = 0; - bool first_item = true; - while (arr_pos < arr_content.length()) { - while (arr_pos < arr_content.length() && (std::isspace(arr_content[arr_pos]) || arr_content[arr_pos] == ',')) arr_pos++; - if (arr_pos >= arr_content.length()) break; - - if (!first_item) value += ","; - first_item = false; - - if (arr_content.compare(arr_pos, 8, "") == 0) { - arr_pos += 8; - size_t end = arr_content.find("", arr_pos); - if (end != std::string::npos) { - value += "\"" + arr_content.substr(arr_pos, end - arr_pos) + "\""; - arr_pos = end + 8; - } - } else { - size_t end = arr_content.find_first_of(",]", arr_pos); - if (end == std::string::npos) end = arr_content.length(); - value += arr_content.substr(arr_pos, end - arr_pos); - arr_pos = end; - } - } - value += "]"; - } else { - size_t val_start = pos; - while (pos < args_content.length() && args_content[pos] != ',' && args_content[pos] != '}') { - pos++; - } - value = args_content.substr(val_start, pos - val_start); - while (!value.empty() && std::isspace(value.back())) value.pop_back(); - } - } - - if (!first) result += ","; - first = false; - result += "\"" + key + "\":" + value; - } - - result += "}"; - return result; -} - -inline void parse_function_calls(std::string& response, std::vector& function_calls) { - const std::string CALL_START = ""; - const std::string CALL_END = ""; - size_t pos = 0; - - while ((pos = response.find(CALL_START, pos)) != std::string::npos) { - size_t content_start = pos + CALL_START.length(); - size_t call_end_pos = response.find(CALL_END, content_start); - - size_t content_end = (call_end_pos != std::string::npos) ? call_end_pos : response.length(); - std::string call_content = response.substr(content_start, content_end - content_start); - - if (call_content.compare(0, 5, "call:") == 0) { - size_t brace_pos = call_content.find('{'); - - if (brace_pos == std::string::npos) { - size_t sep_pos = call_content.find_first_of(", ", 5); - if (sep_pos != std::string::npos) { - std::string func_name = call_content.substr(5, sep_pos - 5); - size_t args_start = sep_pos + 1; - while (args_start < call_content.length() && - (call_content[args_start] == ' ' || call_content[args_start] == ',')) { - args_start++; - } - std::string args_content = "{" + call_content.substr(args_start); - if (args_content.back() != '}') args_content += "}"; - - std::string args_json = args_to_json(args_content); - std::string json_call = "{\"name\":\"" + func_name + "\",\"arguments\":" + args_json + "}"; - function_calls.push_back(json_call); - } - } else { - std::string func_name = call_content.substr(5, brace_pos - 5); - std::string args_content = call_content.substr(brace_pos); - if (args_content.back() != '}') args_content += "}"; - - std::string args_json = args_to_json(args_content); - std::string json_call = "{\"name\":\"" + func_name + "\",\"arguments\":" + args_json + "}"; - function_calls.push_back(json_call); - } - } - - size_t erase_end = (call_end_pos != std::string::npos) ? - call_end_pos + CALL_END.length() : response.length(); - response.erase(pos, erase_end - pos); - } -} - -} // namespace gemma \ No newline at end of file diff --git a/cactus/graph/graph.h b/cactus/graph/graph.h deleted file mode 100644 index 1caf3f375..000000000 --- a/cactus/graph/graph.h +++ /dev/null @@ -1,571 +0,0 @@ -#ifndef GRAPH_H -#define GRAPH_H - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace cactus { - -enum class LogLevel { - DEBUG = 0, - INFO = 1, - WARN = 2, - ERROR = 3, - NONE = 4 -}; - -class Logger { -public: - static Logger& instance() { - static Logger logger; - return logger; - } - - void set_level(LogLevel level) { min_level_ = level; } - LogLevel get_level() const { return min_level_; } - - void set_callback(std::function cb) { - std::lock_guard lock(mutex_); - callback_ = cb; - } - - void log(LogLevel level, const std::string& component, const std::string& message) { - if (level < min_level_) return; - - std::lock_guard lock(mutex_); - - if (callback_) { - callback_(level, component, message); - } else { - std::cerr << "[" << level_string(level) << "] [" << component << "] " << message << std::endl; - } - - if (level == LogLevel::ERROR) { - last_error_ = "[" + component + "] " + message; - } - } - - const std::string& last_error() const { return last_error_; } - void clear_error() { last_error_.clear(); } - -private: - Logger() : min_level_(LogLevel::WARN) {} - - static const char* level_string(LogLevel level) { - switch (level) { - case LogLevel::DEBUG: return "DEBUG"; - case LogLevel::INFO: return "INFO"; - case LogLevel::WARN: return "WARN"; - case LogLevel::ERROR: return "ERROR"; - default: return "?"; - } - } - - LogLevel min_level_; - std::mutex mutex_; - std::string last_error_; - std::function callback_; -}; - -} // namespace cactus - -#define CACTUS_LOG(level, component, msg) \ - do { \ - if (static_cast(level) >= static_cast(cactus::Logger::instance().get_level())) { \ - std::ostringstream _cactus_log_ss; \ - _cactus_log_ss << msg; \ - cactus::Logger::instance().log(level, component, _cactus_log_ss.str()); \ - } \ - } while(0) - -#define CACTUS_LOG_DEBUG(component, msg) CACTUS_LOG(cactus::LogLevel::DEBUG, component, msg) -#define CACTUS_LOG_INFO(component, msg) CACTUS_LOG(cactus::LogLevel::INFO, component, msg) -#define CACTUS_LOG_WARN(component, msg) CACTUS_LOG(cactus::LogLevel::WARN, component, msg) -#define CACTUS_LOG_ERROR(component, msg) CACTUS_LOG(cactus::LogLevel::ERROR, component, msg) - -namespace GraphFile { - class MappedFile; -} - -enum class Precision { - INT8, - FP16, - FP32, - INT4 -}; - -enum class ComputeBackend { - CPU, - NPU -}; - -enum class OpType { - INPUT, PRECISION_CAST, - ADD, ADD_CLIPPED, SUBTRACT, MULTIPLY, DIVIDE, - MATMUL, TRANSPOSE, RESHAPE, SLICE, GATHER, EMBEDDING, - BILINEAR_INTERPOLATION, - SUM, MEAN, VARIANCE, MIN, MAX, - RMS_NORM, ROPE, SOFTMAX, ATTENTION, ATTENTION_INT8_HYBRID, CONV1D_CAUSAL, CONV1D_K3, - SCALAR_ADD, SCALAR_SUBTRACT, SCALAR_MULTIPLY, SCALAR_DIVIDE, SCALAR_EXP, SCALAR_SQRT, SCALAR_COS, SCALAR_SIN, - SILU, GELU, GELU_ERF, - SAMPLE, CONCAT, - SCATTER_TOPK, - TOPK, LAYERNORM, - INDEX, - QUANTIZE_ACTIVATIONS, -}; - -struct PrecisionTraits { - static constexpr size_t size_of(Precision prec) { - switch (prec) { - case Precision::INT8: return 1; - case Precision::FP16: return 2; - case Precision::FP32: return 4; - case Precision::INT4: return 1; - } - return 1; - } - - static constexpr size_t packed_size_of(Precision prec, size_t count) { - switch (prec) { - case Precision::INT4: return (count + 1) / 2; - default: return count * size_of(prec); - } - } - - static constexpr bool is_integer(Precision prec) { - switch (prec) { - case Precision::INT8: return true; - case Precision::INT4: return true; - case Precision::FP16: return false; - case Precision::FP32: return false; - } - return true; - } - - static constexpr bool is_floating_point(Precision prec) { - switch (prec) { - case Precision::INT8: return false; - case Precision::INT4: return false; - case Precision::FP16: return true; - case Precision::FP32: return true; - } - return false; - } -}; - -namespace Quantization { - void int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale = 1.0f); - void fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale = 1.0f); - void fp16_to_fp32(const __fp16* src, float* dst, size_t count); - void fp32_to_fp16(const float* src, __fp16* dst, size_t count); - void int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale = 1.0f); - void fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale = 1.0f); -} - -struct TensorConfig { - Precision default_precision = Precision::INT8; - Precision compute_precision = Precision::INT8; - Precision output_precision = Precision::INT8; - bool auto_mixed_precision = false; - bool enable_int4_packing = true; - - static TensorConfig& global(); -}; - -struct BroadcastInfo { - std::vector output_shape; - bool needs_broadcasting; - - static BroadcastInfo compute(const std::vector& lhs, const std::vector& rhs); -}; - -class BufferPool; - -struct BufferDesc { - std::vector shape; - size_t total_size; - size_t byte_size; - std::unique_ptr data; - void* external_data; - char* pooled_data; - Precision precision; - - size_t group_size = 0; - size_t num_groups = 0; - void* scales_data = nullptr; - std::unique_ptr owned_scales; - - const void* packed_int4_data = nullptr; - size_t packed_int4_size = 0; - - void* activation_scales_data = nullptr; - std::unique_ptr owned_activation_scales; - size_t num_rows_for_activation_scales = 0; - - BufferDesc(); - BufferDesc(const std::vector& s, Precision prec = Precision::INT8); - ~BufferDesc(); - - BufferDesc(BufferDesc&& other) noexcept; - BufferDesc& operator=(BufferDesc&& other) noexcept; - - BufferDesc(const BufferDesc&) = delete; - BufferDesc& operator=(const BufferDesc&) = delete; - - void* get_data(); - const void* get_data() const; - - template - T* data_as() { return static_cast(get_data()); } - - template - const T* data_as() const { return static_cast(get_data()); } - - const __fp16* scales_as_fp16() const { - return reinterpret_cast(scales_data); - } - bool is_grouped_int8() const { - return precision == Precision::INT8 && group_size > 0; - } - bool is_packed_int4() const { - return packed_int4_data != nullptr && packed_int4_size > 0; - } - const uint8_t* packed_int4_as_uint8() const { - return reinterpret_cast(packed_int4_data); - } - void set_grouped_scales(size_t gs, size_t ng, void* scales_ptr) { - group_size = gs; - num_groups = ng; - scales_data = scales_ptr; - } - void set_packed_int4(const void* packed_data, size_t packed_size) { - packed_int4_data = packed_data; - packed_int4_size = packed_size; - } - - bool has_activation_scales() const { - return activation_scales_data != nullptr && num_rows_for_activation_scales > 0; - } - const float* activation_scales_as_float() const { - return reinterpret_cast(activation_scales_data); - } - float* activation_scales_as_float() { - return reinterpret_cast(activation_scales_data); - } - void allocate_activation_scales(size_t num_rows) { - num_rows_for_activation_scales = num_rows; - owned_activation_scales = std::make_unique(num_rows * sizeof(float)); - activation_scales_data = owned_activation_scales.get(); - } - void set_activation_scales(void* scales_ptr, size_t num_rows) { - activation_scales_data = scales_ptr; - num_rows_for_activation_scales = num_rows; - } - - void allocate(); - void allocate_from_pool(BufferPool& pool); - void release_to_pool(BufferPool& pool); - void set_external(void* ptr); -}; - -struct OpParams { - float scalar = 0.0f; - float scale = 1.0f; - float theta = 10000.0f; - float epsilon = 1e-6f; - int axis = -1; - bool pretransposed_rhs = false; - size_t position_offset = 0; - size_t slice_start = 0; - size_t slice_length = 0; - size_t window_size = 0; - bool is_causal = true; - std::vector new_shape; - std::vector permutation; - Precision output_precision = Precision::INT8; - BroadcastInfo broadcast_info; - ComputeBackend backend = ComputeBackend::CPU; - - size_t dilation = 1; - size_t stride = 1; - float temperature = 1.0f; - float top_p = 1.0f; - size_t top_k = 0; - size_t random_seed = 0; - - size_t index_value = 0; - size_t num_classes = 0; - size_t dst_height = 0; - size_t dst_width = 0; - - std::vector bias_values; - std::vector bias_indices; - - const int8_t* cached_keys_int8 = nullptr; - const int8_t* cached_values_int8 = nullptr; - const float* cached_k_scales = nullptr; - const float* cached_v_scales = nullptr; - size_t cache_seq_len = 0; - size_t num_kv_heads = 0; - size_t head_dim = 0; -}; - -struct GraphNode { - size_t id; - OpType op_type; - std::vector input_ids; - BufferDesc output_buffer; - OpParams params; - - GraphNode(size_t node_id, OpType type); -}; - -template -void dispatch_binary_op(OpType op, const T* lhs, const T* rhs, T* output, size_t count); - -template -void dispatch_unary_op(OpType op, const T* input, T* output, size_t count, float param = 0.0f); - -void compute_node_optimized(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_matmul_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_transpose_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_reduce_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_fused_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_reshape_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_precision_cast_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_sample_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_scatter_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_layernorm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -void compute_index_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); - -void shrink_thread_local_buffers(); - -class BufferPool { -public: - BufferPool() = default; - ~BufferPool() = default; - - BufferPool(const BufferPool&) = delete; - BufferPool& operator=(const BufferPool&) = delete; - - char* acquire(size_t byte_size); - void release(char* ptr, size_t byte_size); - void clear(); - - size_t active_bytes() const { return active_bytes_; } - size_t pool_bytes() const { return pool_bytes_; } - size_t peak_bytes() const { return peak_bytes_; } - -private: - std::unordered_map>> free_buffers_; - size_t active_bytes_ = 0; - size_t pool_bytes_ = 0; - size_t peak_bytes_ = 0; - - size_t round_up_size(size_t size) const; -}; - -namespace ValidationUtils { - void validate_tensor_dims(const std::vector& shape, size_t required_dims, const std::string& op_name); - void validate_precision(Precision actual, Precision required, const std::string& op_name); - void validate_input_count(size_t actual, size_t required, const std::string& op_name); -} - - -class CactusGraph { -public: - CactusGraph(); - - struct DebugNodeEntry { - uint32_t layer_idx; - std::string name; - size_t node_id; - }; - - size_t input(const std::vector& shape, Precision precision = Precision::INT8); - size_t precision_cast(size_t input, Precision target_precision); - size_t quantize_activations(size_t input); - - size_t add(size_t input1, size_t input2); - size_t add_clipped(size_t input1, size_t input2); - size_t subtract(size_t input1, size_t input2); - size_t multiply(size_t input1, size_t input2); - size_t divide(size_t input1, size_t input2); - - - size_t scalar_add(size_t input, float value); - size_t scalar_subtract(size_t input, float value); - size_t scalar_multiply(size_t input, float value); - size_t scalar_divide(size_t input, float value); - size_t scalar_exp(size_t input); - size_t scalar_sqrt(size_t input); - size_t scalar_cos(size_t input); - size_t scalar_sin(size_t input); - - size_t silu(size_t input); - size_t gelu(size_t input); - size_t gelu_erf(size_t input); - - size_t matmul(size_t input1, size_t input2, bool pretransposed_rhs = false, ComputeBackend backend = ComputeBackend::CPU); - size_t transpose(size_t input, ComputeBackend backend = ComputeBackend::CPU); - size_t transposeN(size_t input, const std::vector& permutation, ComputeBackend backend = ComputeBackend::CPU); - size_t reshape(size_t input, const std::vector& new_shape); - size_t slice(size_t input, int axis, size_t start, size_t length); - size_t index(size_t input, size_t index_value, int dim); - - size_t sum(size_t input, int axis); - size_t mean(size_t input, int axis); - size_t variance(size_t input, int axis); - size_t min(size_t input, int axis); - size_t max(size_t input, int axis); - - size_t gather(size_t embeddings, size_t indices); - size_t mmap_embeddings(const std::string& filename); - size_t mmap_weights(const std::string& filename); - size_t load_weights(const std::string& filename); - void set_grouped_scales(size_t node_id, size_t group_size, size_t num_groups, void* scales_ptr); - - void release_weight_pages(size_t node_id); - void prefetch_weight_pages(size_t node_id); - void release_all_weight_pages(); - size_t embedding(const std::string& filename, size_t indices); - size_t embedding(size_t embedding_tensor, size_t indices); - size_t bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width); - - size_t layernorm(size_t input, size_t weight, size_t bias, float epsilon = 1e-5f); - size_t topk(size_t input, size_t k); - size_t rms_norm(size_t input, size_t weight, float epsilon = 1e-5f); - size_t rope(size_t input, float theta, size_t position_offset = 0, ComputeBackend backend = ComputeBackend::CPU); - size_t softmax(size_t input, int axis = -1); - size_t attention(size_t query, size_t key, size_t value, float scale, bool is_causal = true, ComputeBackend backend = ComputeBackend::CPU); - size_t attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, ComputeBackend backend = ComputeBackend::CPU); - size_t attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, size_t window_size, ComputeBackend backend = ComputeBackend::CPU); - - size_t attention_int8_hybrid(size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset, - const int8_t* cached_keys, const int8_t* cached_values, - const float* k_scales, const float* v_scales, - size_t cache_len, size_t num_kv_heads, size_t head_dim); - - size_t conv1d_causal(size_t input, size_t weight, size_t kernel_size, size_t dilation = 1); - size_t conv1d_k3(size_t input, size_t weight, size_t stride); - - size_t sample(size_t logits, float temperature = 0.6f, float top_p = 0.95f, size_t top_k = 20, - const std::unordered_map& logit_bias = {}); - - size_t concat(size_t input1, size_t input2, int axis = 0); - size_t scatter_topk(size_t indices, size_t values, size_t num_classes); - - void set_input(size_t node_id, const void* data, Precision precision); - void set_external_input(size_t node_id, void* data, Precision precision); - void* get_output(size_t node_id); - - void execute(const std::string& profile_file = ""); - void hard_reset(); - void soft_reset(); - void soft_reset_keep_pool(); - void set_prefill_mode(bool enabled) { prefill_mode_ = enabled; } - - void register_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id); - void capture_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id); - const std::vector& get_debug_nodes() const; - void clear_debug_nodes(); - - size_t add_node(OpType op_type, const std::vector& inputs, const std::vector& output_shape, const OpParams& params = {}); - const BufferDesc& get_output_buffer(size_t node_id) const; - void allocate_buffers(); - size_t get_node_count() const; - - std::vector> nodes_; - std::unordered_map node_index_map_; - -private: - size_t next_node_id_; - std::vector> mapped_files_; - std::unordered_map weight_cache_; - std::unordered_map node_to_mapped_file_; - std::vector debug_nodes_; - BufferPool buffer_pool_; - bool prefill_mode_ = false; -}; - - -namespace GraphFile { - struct LoadedNode { - size_t node_id; - std::vector shape; - Precision precision; - size_t byte_size; - }; - - void save_node(CactusGraph& graph, size_t node_id, const std::string& filename); - LoadedNode load_into_graph(CactusGraph& graph, const std::string& filename); - - class MappedFile { - public: - MappedFile(const std::string& filename); - ~MappedFile(); - - MappedFile(const MappedFile&) = delete; - MappedFile& operator=(const MappedFile&) = delete; - MappedFile(MappedFile&& other) noexcept; - MappedFile& operator=(MappedFile&& other) noexcept; - - const std::vector& shape() const; - Precision precision() const; - Precision effective_precision() const { - return is_int4_ ? Precision::INT8 : precision_; - } - size_t byte_size() const; - - size_t group_size() const { return group_size_; } - size_t num_groups() const { return num_groups_; } - const void* scales_data() const; - const void* raw_packed_data() const; // Get raw mmap'd data without unpacking (for INT4) - bool is_int4() const { return is_int4_; } - - void* data(); - const void* data() const; - - template - const T* typed_data() const; - - LoadedNode load_into_graph(CactusGraph& graph) const; - - void release_pages(); - void prefetch_pages(); - - private: - int fd_; - void* mapped_data_; - size_t file_size_, data_offset_; - std::vector shape_; - Precision precision_; - size_t byte_size_; - size_t group_size_ = 0; - size_t num_groups_ = 0; - size_t scales_offset_ = 0; - size_t scales_bytes_ = 0; - uint32_t version_ = 1; - uint32_t alignment_ = 32; - bool is_int4_ = false; - mutable std::unique_ptr unpacked_int4_data_; - void parse_header(); - void apply_madvise_hints(); - void unpack_int4_if_needed() const; - }; - - MappedFile mmap_load(const std::string& filename); -} - -#endif diff --git a/cactus/graph/graph_builder.cpp b/cactus/graph/graph_builder.cpp deleted file mode 100644 index 9bcf77bc5..000000000 --- a/cactus/graph/graph_builder.cpp +++ /dev/null @@ -1,600 +0,0 @@ -#include "graph.h" -#include -#include -#include -#include - -size_t CactusGraph::input(const std::vector& shape, Precision precision) { - return add_node(OpType::INPUT, {}, shape, {.output_precision = precision}); -} - -size_t CactusGraph::add(size_t input1, size_t input2) { - const auto& lhs_buffer = get_output_buffer(input1); - const auto& rhs_buffer = get_output_buffer(input2); - - BroadcastInfo broadcast_info = BroadcastInfo::compute(lhs_buffer.shape, rhs_buffer.shape); - OpParams params{.broadcast_info = broadcast_info}; - - return add_node(OpType::ADD, {input1, input2}, broadcast_info.output_shape, params); -} - -size_t CactusGraph::add_clipped(size_t input1, size_t input2) { - const auto& lhs_buffer = get_output_buffer(input1); - const auto& rhs_buffer = get_output_buffer(input2); - - BroadcastInfo broadcast_info = BroadcastInfo::compute(lhs_buffer.shape, rhs_buffer.shape); - OpParams params{.broadcast_info = broadcast_info}; - - return add_node(OpType::ADD_CLIPPED, {input1, input2}, broadcast_info.output_shape, params); -} - -size_t CactusGraph::subtract(size_t input1, size_t input2) { - const auto& lhs_buffer = get_output_buffer(input1); - const auto& rhs_buffer = get_output_buffer(input2); - - BroadcastInfo broadcast_info = BroadcastInfo::compute(lhs_buffer.shape, rhs_buffer.shape); - OpParams params{.broadcast_info = broadcast_info}; - - return add_node(OpType::SUBTRACT, {input1, input2}, broadcast_info.output_shape, params); -} - -size_t CactusGraph::multiply(size_t input1, size_t input2) { - const auto& lhs_buffer = get_output_buffer(input1); - const auto& rhs_buffer = get_output_buffer(input2); - - BroadcastInfo broadcast_info = BroadcastInfo::compute(lhs_buffer.shape, rhs_buffer.shape); - OpParams params{.broadcast_info = broadcast_info}; - - return add_node(OpType::MULTIPLY, {input1, input2}, broadcast_info.output_shape, params); -} - -size_t CactusGraph::divide(size_t input1, size_t input2) { - const auto& lhs_buffer = get_output_buffer(input1); - const auto& rhs_buffer = get_output_buffer(input2); - - BroadcastInfo broadcast_info = BroadcastInfo::compute(lhs_buffer.shape, rhs_buffer.shape); - OpParams params{.broadcast_info = broadcast_info}; - - return add_node(OpType::DIVIDE, {input1, input2}, broadcast_info.output_shape, params); -} - -size_t CactusGraph::matmul(size_t input1, size_t input2, bool pretransposed_rhs, ComputeBackend backend) { - const auto& lhs_buffer = get_output_buffer(input1); - const auto& rhs_buffer = get_output_buffer(input2); - - if (lhs_buffer.shape.size() != 2 || rhs_buffer.shape.size() != 2) { - throw std::invalid_argument("Matrix multiplication requires 2D tensors"); - } - - size_t M = lhs_buffer.shape[0]; - size_t K = lhs_buffer.shape[1]; - size_t rhs_K = pretransposed_rhs ? rhs_buffer.shape[1] : rhs_buffer.shape[0]; - size_t N = pretransposed_rhs ? rhs_buffer.shape[0] : rhs_buffer.shape[1]; - - if (K != rhs_K) { - throw std::invalid_argument("Matrix dimensions incompatible for multiplication"); - } - - std::vector output_shape = {M, N}; - OpParams params{.pretransposed_rhs = pretransposed_rhs, .backend = backend}; - return add_node(OpType::MATMUL, {input1, input2}, output_shape, params); -} - -size_t CactusGraph::transpose(size_t input, ComputeBackend backend) { - const auto& input_buffer = get_output_buffer(input); - std::vector output_shape = input_buffer.shape; - - if (output_shape.size() >= 2) { - std::swap(output_shape[output_shape.size()-2], output_shape[output_shape.size()-1]); - } - - std::vector permutation; - for (size_t i = 0; i < input_buffer.shape.size(); ++i) { - permutation.push_back(i); - } - if (permutation.size() >= 2) { - std::swap(permutation[permutation.size()-2], permutation[permutation.size()-1]); - } - - OpParams params{.permutation = permutation, .backend = backend}; - return add_node(OpType::TRANSPOSE, {input}, output_shape, params); -} - -size_t CactusGraph::transposeN(size_t input, const std::vector& permutation, ComputeBackend backend) { - const auto& input_buffer = get_output_buffer(input); - if (permutation.size() != input_buffer.shape.size()) { - throw std::runtime_error("transposeN permutation size must match tensor rank"); - } - std::vector output_shape(permutation.size()); - for (size_t i = 0; i < permutation.size(); ++i) { - size_t p = permutation[i]; - if (p >= input_buffer.shape.size()) { - throw std::runtime_error("transposeN permutation index out of range"); - } - output_shape[i] = input_buffer.shape[p]; - } - OpParams params{.permutation = permutation, .backend = backend}; - return add_node(OpType::TRANSPOSE, {input}, output_shape, params); -} - -size_t CactusGraph::reshape(size_t input, const std::vector& new_shape) { - OpParams params{.new_shape = new_shape}; - return add_node(OpType::RESHAPE, {input}, new_shape, params); -} - -size_t CactusGraph::index(size_t input, size_t index_value, int dim) { - const auto& input_buffer = get_output_buffer(input); - const auto& shape = input_buffer.shape; - - if (shape.empty()) { - throw std::invalid_argument("Cannot index a scalar tensor"); - } - - int actual_dim = dim; - if (actual_dim < 0) { - actual_dim += static_cast(shape.size()); - } - - if (actual_dim < 0 || static_cast(actual_dim) >= shape.size()) { - throw std::invalid_argument("Index dimension out of bounds"); - } - - if (index_value >= shape[actual_dim]) { - throw std::invalid_argument("Index value " + std::to_string(index_value) + - " out of bounds for dimension " + std::to_string(actual_dim) + - " with size " + std::to_string(shape[actual_dim])); - } - - std::vector output_shape; - for (size_t i = 0; i < shape.size(); ++i) { - if (static_cast(i) != actual_dim) { - output_shape.push_back(shape[i]); - } - } - - if (output_shape.empty()) { - output_shape = {1}; - } - - OpParams params{.axis = actual_dim, .output_precision = input_buffer.precision, .index_value = index_value}; - return add_node(OpType::INDEX, {input}, output_shape, params); -} - -size_t CactusGraph::sum(size_t input, int axis) { - const auto& input_buffer = get_output_buffer(input); - std::vector output_shape; - - if (axis == -1) { - output_shape = {1}; - } else { - output_shape = input_buffer.shape; - output_shape.erase(output_shape.begin() + axis); - if (output_shape.empty()) { - output_shape = {1}; - } - } - - OpParams params{.axis = axis, .output_precision = input_buffer.precision}; - return add_node(OpType::SUM, {input}, output_shape, params); -} - -size_t CactusGraph::mean(size_t input, int axis) { - const auto& input_buffer = get_output_buffer(input); - std::vector output_shape; - - if (axis == -1) { - output_shape = {1}; - } else { - output_shape = input_buffer.shape; - output_shape.erase(output_shape.begin() + axis); - if (output_shape.empty()) { - output_shape = {1}; - } - } - - OpParams params{.axis = axis, .output_precision = input_buffer.precision}; - return add_node(OpType::MEAN, {input}, output_shape, params); -} - -size_t CactusGraph::variance(size_t input, int axis) { - const auto& input_buffer = get_output_buffer(input); - std::vector output_shape; - - if (axis == -1) { - output_shape = {1}; - } else { - output_shape = input_buffer.shape; - output_shape.erase(output_shape.begin() + axis); - if (output_shape.empty()) { - output_shape = {1}; - } - } - - OpParams params{.axis = axis, .output_precision = input_buffer.precision}; - return add_node(OpType::VARIANCE, {input}, output_shape, params); -} - -size_t CactusGraph::min(size_t input, int axis) { - const auto& input_buffer = get_output_buffer(input); - std::vector output_shape; - - if (axis == -1) { - output_shape = {1}; - } else { - output_shape = input_buffer.shape; - output_shape.erase(output_shape.begin() + axis); - if (output_shape.empty()) { - output_shape = {1}; - } - } - - OpParams params{.axis = axis, .output_precision = input_buffer.precision}; - return add_node(OpType::MIN, {input}, output_shape, params); -} - -size_t CactusGraph::max(size_t input, int axis) { - const auto& input_buffer = get_output_buffer(input); - std::vector output_shape; - - if (axis == -1) { - output_shape = {1}; - } else { - output_shape = input_buffer.shape; - output_shape.erase(output_shape.begin() + axis); - if (output_shape.empty()) { - output_shape = {1}; - } - } - - OpParams params{.axis = axis, .output_precision = input_buffer.precision}; - return add_node(OpType::MAX, {input}, output_shape, params); -} - -size_t CactusGraph::rms_norm(size_t input, size_t weight, float epsilon) { - OpParams params{.epsilon = epsilon}; - return add_node(OpType::RMS_NORM, {input, weight}, {}, params); -} - -size_t CactusGraph::rope(size_t input, float theta, size_t position_offset, ComputeBackend backend) { - OpParams params{.theta = theta, .position_offset = position_offset, .backend = backend}; - return add_node(OpType::ROPE, {input}, {}, params); -} - -size_t CactusGraph::softmax(size_t input, int axis) { - OpParams params{.axis = axis}; - return add_node(OpType::SOFTMAX, {input}, {}, params); -} - -size_t CactusGraph::topk(size_t input, size_t k) { - const auto& input_buffer = get_output_buffer(input); - - if (input_buffer.shape.empty()) { - throw std::runtime_error("TopK requires non-empty input tensor"); - } - - std::vector output_shape = {2, input_buffer.shape[0], k}; - OpParams params{.output_precision = Precision::FP32, .top_k = k}; - - return add_node(OpType::TOPK, {input}, output_shape, params); -} - -size_t CactusGraph::layernorm(size_t input, size_t weight, size_t bias, float epsilon) { - OpParams params{.epsilon = epsilon}; - return add_node(OpType::LAYERNORM, {input, weight, bias}, {}, params); -} - -size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, bool is_causal, ComputeBackend backend) { - OpParams params{.scale = scale, .is_causal = is_causal, .backend = backend}; - return add_node(OpType::ATTENTION, {query, key, value}, {}, params); -} - -size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, ComputeBackend backend) { - OpParams params{.scale = scale, .position_offset = position_offset, .backend = backend}; - return add_node(OpType::ATTENTION, {query, key, value}, {}, params); -} - -size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, size_t window_size, ComputeBackend backend) { - OpParams params{.scale = scale, .position_offset = position_offset, .window_size = window_size, .backend = backend}; - return add_node(OpType::ATTENTION, {query, key, value}, {}, params); -} - -size_t CactusGraph::attention_int8_hybrid(size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset, - const int8_t* cached_keys, const int8_t* cached_values, - const float* k_scales, const float* v_scales, - size_t cache_len, size_t num_kv_heads, size_t head_dim) { - OpParams params; - params.scale = scale; - params.position_offset = position_offset; - params.cached_keys_int8 = cached_keys; - params.cached_values_int8 = cached_values; - params.cached_k_scales = k_scales; - params.cached_v_scales = v_scales; - params.cache_seq_len = cache_len; - params.num_kv_heads = num_kv_heads; - params.head_dim = head_dim; - return add_node(OpType::ATTENTION_INT8_HYBRID, {query, key_new, value_new}, {}, params); -} - -size_t CactusGraph::conv1d_causal(size_t input, size_t weight, size_t, size_t dilation) { - OpParams params{.dilation = dilation}; - return add_node(OpType::CONV1D_CAUSAL, {input, weight}, {}, params); -} - -size_t CactusGraph::conv1d_k3(size_t input, size_t weight, size_t stride) { - const auto& xin = get_output_buffer(input); - const auto& w = get_output_buffer(weight); - - if (xin.shape.size() != 3) throw std::runtime_error("conv1d_k3 expects N,C,L"); - if (w.shape.size() != 3) throw std::runtime_error("weight must be [C_out,C_in,3]"); - if (w.shape[1] != xin.shape[1]) throw std::runtime_error("C_in mismatch in conv1d_k3"); - if (w.shape[2] != 3) throw std::runtime_error("K=3 expected in conv1d_k3"); - - const size_t N = xin.shape[0]; - const size_t L = xin.shape[2]; - const size_t C_out= w.shape[0]; - const size_t K = w.shape[2]; - - const size_t pad = 1; - const size_t L_out = (L + 2 * pad - K) / stride + 1; - - OpParams params{}; - params.stride = stride; - params.output_precision = xin.precision; - - std::vector out_shape{N, C_out, L_out}; - return add_node(OpType::CONV1D_K3, {input, weight}, out_shape, params); -} - -size_t CactusGraph::concat(size_t input1, size_t input2, int axis) { - const auto& buffer1 = get_output_buffer(input1); - const auto& buffer2 = get_output_buffer(input2); - - if (buffer1.shape.size() != buffer2.shape.size()) { - throw std::runtime_error("Concat requires inputs with same number of dimensions"); - } - - std::vector output_shape = buffer1.shape; - size_t ndims = output_shape.size(); - - if (axis < 0) axis += ndims; - if (axis < 0 || static_cast(axis) >= ndims) { - throw std::runtime_error("Invalid axis for concat operation"); - } - - for (size_t i = 0; i < ndims; ++i) { - if (i != static_cast(axis) && buffer1.shape[i] != buffer2.shape[i]) { - throw std::runtime_error("Concat inputs must have same shape except on concat axis"); - } - } - - output_shape[axis] = buffer1.shape[axis] + buffer2.shape[axis]; - - OpParams params; - params.axis = axis; - return add_node(OpType::CONCAT, {input1, input2}, output_shape, params); -} - -size_t CactusGraph::scatter_topk(size_t indices, size_t values, size_t num_classes) { - const auto& indices_buffer = get_output_buffer(indices); - const auto& values_buffer = get_output_buffer(values); - - if (indices_buffer.shape != values_buffer.shape) { - throw std::runtime_error("ScatterTopK requires indices and values with identical shapes"); - } - if (indices_buffer.shape.size() != 2) { - throw std::runtime_error("ScatterTopK currently supports 2D tensors [batch, top_k]"); - } - if (indices_buffer.precision != Precision::FP32 || values_buffer.precision != Precision::FP32) { - throw std::runtime_error("ScatterTopK expects FP32 indices and values"); - } - - std::vector output_shape = {num_classes, indices_buffer.shape[0]}; - OpParams params{.output_precision = Precision::FP32, .num_classes = num_classes}; - return add_node(OpType::SCATTER_TOPK, {indices, values}, output_shape, params); -} - -size_t CactusGraph::sample(size_t logits, float temperature, float top_p, size_t top_k, - const std::unordered_map& logit_bias) { - const auto& logits_buffer = get_output_buffer(logits); - - if (logits_buffer.shape.empty()) { - throw std::runtime_error("Sample requires non-empty logits tensor"); - } - - OpParams params; - params.temperature = temperature; - params.top_p = top_p; - params.top_k = top_k; - params.random_seed = std::chrono::high_resolution_clock::now().time_since_epoch().count(); - params.output_precision = Precision::FP32; - - if (!logit_bias.empty()) { - params.bias_indices.reserve(logit_bias.size()); - params.bias_values.reserve(logit_bias.size()); - for (const auto& [idx, val] : logit_bias) { - params.bias_indices.push_back(idx); - params.bias_values.push_back(val); - } - } - - std::vector output_shape = {1}; - return add_node(OpType::SAMPLE, {logits}, output_shape, params); -} - -size_t CactusGraph::scalar_add(size_t input, float value) { - OpParams params{}; - params.scalar = value; - params.output_precision = get_output_buffer(input).precision; - return add_node(OpType::SCALAR_ADD, {input}, {}, params); -} - -size_t CactusGraph::scalar_subtract(size_t input, float value) { - OpParams params{}; - params.scalar = value; - params.output_precision = get_output_buffer(input).precision; - return add_node(OpType::SCALAR_SUBTRACT, {input}, {}, params); -} - -size_t CactusGraph::scalar_multiply(size_t input, float value) { - OpParams params{}; - params.scalar = value; - params.output_precision = get_output_buffer(input).precision; - return add_node(OpType::SCALAR_MULTIPLY, {input}, {}, params); -} - -size_t CactusGraph::scalar_divide(size_t input, float value) { - OpParams params{}; - params.scalar = value; - params.output_precision = get_output_buffer(input).precision; - return add_node(OpType::SCALAR_DIVIDE, {input}, {}, params); -} - -size_t CactusGraph::scalar_exp(size_t input) { - return add_node(OpType::SCALAR_EXP, {input}, {}); -} - -size_t CactusGraph::scalar_sqrt(size_t input) { - return add_node(OpType::SCALAR_SQRT, {input}, {}); -} - -size_t CactusGraph::scalar_cos(size_t input) { - return add_node(OpType::SCALAR_COS, {input}, {}); -} - -size_t CactusGraph::scalar_sin(size_t input) { - return add_node(OpType::SCALAR_SIN, {input}, {}); -} - -size_t CactusGraph::silu(size_t input) { - return add_node(OpType::SILU, {input}, {}); -} - -size_t CactusGraph::gelu(size_t input) { - return add_node(OpType::GELU, {input}, {}); -} - -size_t CactusGraph::gelu_erf(size_t input) { - return add_node(OpType::GELU_ERF, {input}, {}); -} - -size_t CactusGraph::gather(size_t tensor, size_t indices) { - const auto& tensor_buffer = get_output_buffer(tensor); - const auto& idx_shape = get_output_buffer(indices).shape; - - if (tensor_buffer.shape.empty()) { - throw std::runtime_error("Cannot gather from scalar tensor"); - } - - std::vector output_shape = idx_shape; - for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { - output_shape.push_back(tensor_buffer.shape[i]); - } - - OpParams params; - params.output_precision = tensor_buffer.precision; - - return add_node(OpType::GATHER, {tensor, indices}, output_shape, params); -} - -size_t CactusGraph::slice(size_t input, int axis, size_t start, size_t length) { - const auto& input_buffer = get_output_buffer(input); - if (input_buffer.shape.empty()) { - throw std::runtime_error("Cannot slice a scalar tensor"); - } - - size_t axis_index = static_cast(axis); - size_t axis_size = input_buffer.shape[axis_index]; - - if (start + length > axis_size) { - throw std::runtime_error("Slice range extends beyond axis size"); - } - - std::vector output_shape = input_buffer.shape; - output_shape[axis_index] = length; - - OpParams params; - params.axis = axis_index; - params.slice_start = start; - params.slice_length = length; - params.output_precision = input_buffer.precision; - - return add_node(OpType::SLICE, {input}, output_shape, params); -} - -size_t CactusGraph::embedding(size_t embedding_tensor, size_t indices) { - const auto& emb_buffer = get_output_buffer(embedding_tensor); - const auto& idx_shape = get_output_buffer(indices).shape; - - if (emb_buffer.shape.size() != 2) { - throw std::runtime_error("Embedding tensor must be 2D [vocab_size, hidden_dim]"); - } - - std::vector output_shape = idx_shape; - output_shape.push_back(emb_buffer.shape[1]); - - OpParams params; - params.output_precision = (emb_buffer.precision == Precision::INT8) ? Precision::FP16 : emb_buffer.precision; - - return add_node(OpType::EMBEDDING, {embedding_tensor, indices}, output_shape, params); -} - -size_t CactusGraph::bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width) { - const auto& pos_embeds_buffer = get_output_buffer(pos_embeds); - size_t embed_dim = pos_embeds_buffer.shape[1]; - std::vector output_shape = {dst_height * dst_width, embed_dim}; - - OpParams params; - params.dst_height = dst_height; - params.dst_width = dst_width; - params.output_precision = Precision::FP16; - - return add_node(OpType::BILINEAR_INTERPOLATION, {pos_embeds}, output_shape, params); -} - -size_t CactusGraph::precision_cast(size_t input, Precision target_precision) { - OpParams params{}; - params.output_precision = target_precision; - return add_node(OpType::PRECISION_CAST, {input}, {}, params); -} - -size_t CactusGraph::quantize_activations(size_t input) { - const auto& input_buffer = get_output_buffer(input); - - if (input_buffer.precision != Precision::FP16) { - throw std::invalid_argument("quantize_activations requires FP16 input"); - } - - OpParams params{}; - params.output_precision = Precision::INT8; - return add_node(OpType::QUANTIZE_ACTIVATIONS, {input}, input_buffer.shape, params); -} - -size_t CactusGraph::add_node(OpType op_type, const std::vector& inputs, const std::vector& output_shape, const OpParams& params) { - auto node = std::make_unique(next_node_id_, op_type); - node->input_ids = inputs; - node->params = params; - - std::vector result_shape = output_shape; - if (result_shape.empty() && !inputs.empty()) { - result_shape = nodes_[node_index_map_[inputs[0]]]->output_buffer.shape; - } - - Precision result_precision = params.output_precision; - if (op_type == OpType::PRECISION_CAST) { - result_precision = params.output_precision; - } else if (result_precision == Precision::INT8 && !inputs.empty()) { - result_precision = nodes_[node_index_map_[inputs[0]]]->output_buffer.precision; - } - - node->output_buffer = BufferDesc(result_shape, result_precision); - - size_t node_id = next_node_id_++; - node_index_map_[node_id] = nodes_.size(); - nodes_.push_back(std::move(node)); - - return node_id; -} - -const BufferDesc& CactusGraph::get_output_buffer(size_t node_id) const { - return nodes_[node_index_map_.at(node_id)]->output_buffer; -} diff --git a/cactus/graph/graph_execute.cpp b/cactus/graph/graph_execute.cpp deleted file mode 100644 index e0233c1be..000000000 --- a/cactus/graph/graph_execute.cpp +++ /dev/null @@ -1,707 +0,0 @@ -#include "graph.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -extern void compute_binary_op_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_unary_op_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_activation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_reduce_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_reshape_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_precision_cast_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); - -extern void compute_matmul_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_rms_norm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_rope_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_softmax_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_attention_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_attention_int8_hybrid_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_layernorm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_conv1d_causal_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_conv1d_k3_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void shrink_thread_local_buffers(); - -extern void compute_transpose_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_gather_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_slice_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_embedding_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_concat_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_index_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_bilinear_interpolation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); - -extern void compute_sample_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_scatter_topk_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); -extern void compute_quantize_activations_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map); - -static const char* op_type_names[] = { - "INPUT", "PRECISION_CAST", - "ADD", "ADD_CLIPPED", "SUBTRACT", "MULTIPLY", "DIVIDE", - "MATMUL", "TRANSPOSE", "RESHAPE", "SLICE", "GATHER", "EMBEDDING", - "BILINEAR_INTERPOLATION", - "SUM", "MEAN", "VARIANCE", "MIN", "MAX", - "RMS_NORM", "ROPE", "SOFTMAX", "ATTENTION", "ATTENTION_INT8_HYBRID", "CONV1D_CAUSAL", "CONV1D_K3", - "SCALAR_ADD", "SCALAR_SUBTRACT", "SCALAR_MULTIPLY", "SCALAR_DIVIDE", - "SCALAR_EXP", "SCALAR_SQRT", "SCALAR_COS", "SCALAR_SIN", - "SILU", "GELU", "GELU_ERF", "SAMPLE", "CONCAT", - "SCATTER_TOPK", - "TOPK", "LAYERNORM", - "INDEX", - "QUANTIZE_ACTIVATIONS" -}; - -static const char* get_op_name(OpType op) { - return op_type_names[static_cast(op)]; -} - -void compute_node_optimized(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - switch (node.op_type) { - case OpType::INPUT: - break; - - case OpType::ADD: - case OpType::ADD_CLIPPED: - case OpType::SUBTRACT: - case OpType::MULTIPLY: - case OpType::DIVIDE: - compute_binary_op_node(node, nodes, node_index_map); - break; - - case OpType::SCALAR_ADD: - case OpType::SCALAR_SUBTRACT: - case OpType::SCALAR_MULTIPLY: - case OpType::SCALAR_DIVIDE: - case OpType::SCALAR_EXP: - case OpType::SCALAR_SQRT: - case OpType::SCALAR_COS: - case OpType::SCALAR_SIN: - compute_unary_op_node(node, nodes, node_index_map); - break; - - case OpType::SILU: - case OpType::GELU: - case OpType::GELU_ERF: - compute_activation_node(node, nodes, node_index_map); - break; - - case OpType::SUM: - case OpType::MEAN: - case OpType::VARIANCE: - case OpType::MIN: - case OpType::MAX: - compute_reduce_node(node, nodes, node_index_map); - break; - - case OpType::RESHAPE: - compute_reshape_node(node, nodes, node_index_map); - break; - - case OpType::PRECISION_CAST: - compute_precision_cast_node(node, nodes, node_index_map); - break; - - case OpType::MATMUL: - compute_matmul_node(node, nodes, node_index_map); - break; - - case OpType::RMS_NORM: - compute_rms_norm_node(node, nodes, node_index_map); - break; - - case OpType::ROPE: - compute_rope_node(node, nodes, node_index_map); - break; - - case OpType::SOFTMAX: - compute_softmax_node(node, nodes, node_index_map); - break; - - case OpType::ATTENTION: - compute_attention_node(node, nodes, node_index_map); - break; - - case OpType::ATTENTION_INT8_HYBRID: - compute_attention_int8_hybrid_node(node, nodes, node_index_map); - break; - - case OpType::LAYERNORM: - compute_layernorm_node(node, nodes, node_index_map); - break; - - case OpType::CONV1D_CAUSAL: - compute_conv1d_causal_node(node, nodes, node_index_map); - break; - - case OpType::CONV1D_K3: - compute_conv1d_k3_node(node, nodes, node_index_map); - break; - - case OpType::TRANSPOSE: - compute_transpose_node(node, nodes, node_index_map); - break; - - case OpType::GATHER: - compute_gather_node(node, nodes, node_index_map); - break; - - case OpType::SLICE: - compute_slice_node(node, nodes, node_index_map); - break; - - case OpType::EMBEDDING: - compute_embedding_node(node, nodes, node_index_map); - break; - - case OpType::CONCAT: - compute_concat_node(node, nodes, node_index_map); - break; - - case OpType::INDEX: - compute_index_node(node, nodes, node_index_map); - break; - - case OpType::BILINEAR_INTERPOLATION: - compute_bilinear_interpolation_node(node, nodes, node_index_map); - break; - - case OpType::SAMPLE: - compute_sample_node(node, nodes, node_index_map); - break; - - case OpType::TOPK: - compute_topk_node(node, nodes, node_index_map); - break; - - case OpType::SCATTER_TOPK: - compute_scatter_topk_node(node, nodes, node_index_map); - break; - - case OpType::QUANTIZE_ACTIVATIONS: - compute_quantize_activations_node(node, nodes, node_index_map); - break; - - default: - throw std::runtime_error("Unknown operation type: " + std::to_string(static_cast(node.op_type))); - } -} - -void CactusGraph::set_input(size_t node_id, const void* data, Precision) { - auto& node = *nodes_[node_index_map_[node_id]]; - if (node.op_type != OpType::INPUT) { - throw std::invalid_argument("Can only set data on input nodes"); - } - - if (node.output_buffer.is_packed_int4()) { - throw std::invalid_argument("Cannot use set_input on packed INT4 buffer - use set_external_input instead"); - } - - if (!node.output_buffer.data && !node.output_buffer.external_data) { - node.output_buffer.allocate(); - } - - std::memcpy(node.output_buffer.get_data(), data, node.output_buffer.byte_size); -} - -void CactusGraph::set_external_input(size_t node_id, void* data, Precision) { - auto& node = *nodes_[node_index_map_[node_id]]; - if (node.op_type != OpType::INPUT) { - throw std::invalid_argument("Can only set data on input nodes"); - } - - node.output_buffer.set_external(data); -} - -void* CactusGraph::get_output(size_t node_id) { - auto& buffer = nodes_[node_index_map_[node_id]]->output_buffer; - if (!buffer.get_data()) { - buffer.allocate(); - } - return buffer.get_data(); -} - -void CactusGraph::execute(const std::string& profile_file) { - std::vector last_use(nodes_.size(), 0); - for (size_t i = 0; i < nodes_.size(); ++i) { - for (size_t input_id : nodes_[i]->input_ids) { - auto it = node_index_map_.find(input_id); - if (it != node_index_map_.end()) { - last_use[it->second] = std::max(last_use[it->second], i); - } - } - } - - BufferPool& pool = buffer_pool_; - - auto get_env_int = [](const char* name, int fallback) -> int { - const char* val = std::getenv(name); - return val ? std::atoi(val) : fallback; - }; - - auto get_env_str = [](const char* name) -> std::string { - const char* val = std::getenv(name); - return val ? std::string(val) : std::string(); - }; - - bool capture_to_stdout = get_env_int("CACTUS_CAPTURE_STDOUT", 0) != 0; - std::string capture_file_path = get_env_str("CACTUS_CAPTURE_FILE"); - bool capture_requested = get_env_int("CACTUS_CAPTURE_ENABLE", 0) != 0; - - if (!capture_requested) { - capture_requested = capture_to_stdout || !capture_file_path.empty(); - } else if (capture_file_path.empty() && !capture_to_stdout) { - capture_to_stdout = true; - } - - size_t capture_preview_count = static_cast(get_env_int("CACTUS_CAPTURE_PREVIEW_COUNT", 8)); - size_t capture_max_elements = static_cast(get_env_int("CACTUS_CAPTURE_MAX_ELEMENTS", 65536)); - - bool enable_profiling = !profile_file.empty(); - std::ofstream profile_out; - std::ostream* out = &std::cout; - - if (enable_profiling) { - profile_out.open(profile_file); - if (profile_out.is_open()) { - out = &profile_out; - } - } - - auto total_start = std::chrono::high_resolution_clock::now(); - - if (enable_profiling) { - *out << "=== Graph Execution Profile ===" << std::endl; - *out << std::left << std::setw(15) << "Operation" - << std::setw(12) << "Time (ms)" - << std::setw(20) << "Output Shape" - << "Backend" << std::endl; - *out << std::string(60, '-') << std::endl; - } - - for (size_t node_idx = 0; node_idx < nodes_.size(); ++node_idx) { - auto& node = nodes_[node_idx]; - - if (node->op_type != OpType::INPUT) { - node->output_buffer.allocate_from_pool(pool); - } - - if (enable_profiling && node->op_type != OpType::INPUT) { - auto start = std::chrono::high_resolution_clock::now(); - - compute_node_optimized(*node, nodes_, node_index_map_); - - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - double ms = duration.count() / 1000.0; - - std::string shape_str = "["; - for (size_t i = 0; i < node->output_buffer.shape.size(); ++i) { - if (i > 0) shape_str += ","; - shape_str += std::to_string(node->output_buffer.shape[i]); - } - shape_str += "]"; - - std::string values_str = ""; - if (node->output_buffer.get_data()) { - size_t num_values = std::min(size_t(5), node->output_buffer.total_size); - values_str = " values=["; - - if (node->output_buffer.precision == Precision::FP32) { - if (node->op_type == OpType::SAMPLE) { - uint32_t* uint32_data = reinterpret_cast(node->output_buffer.get_data()); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) values_str += ","; - values_str += std::to_string(uint32_data[i]); - } - } else { - float* float_data = reinterpret_cast(node->output_buffer.get_data()); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) values_str += ","; - values_str += std::to_string(float_data[i]).substr(0, 6); - } - } - } else if (node->output_buffer.precision == Precision::FP16) { - __fp16* fp16_data = reinterpret_cast<__fp16*>(node->output_buffer.get_data()); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) values_str += ","; - values_str += std::to_string(static_cast(fp16_data[i])).substr(0, 6); - } - } else if (node->output_buffer.precision == Precision::INT8) { - int8_t* int8_data = reinterpret_cast(node->output_buffer.get_data()); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) values_str += ","; - values_str += std::to_string(static_cast(int8_data[i])); - } - } - - if (node->output_buffer.total_size > 5) { - values_str += ",..."; - } - values_str += "]"; - } - - std::string weights_str = ""; - if ((node->op_type == OpType::RMS_NORM || node->op_type == OpType::MATMUL || - node->op_type == OpType::GATHER || node->op_type == OpType::EMBEDDING || - node->op_type == OpType::ATTENTION || node->op_type == OpType::CONCAT) && - node->input_ids.size() >= 2) { - const auto& weight_node = nodes_[node_index_map_.at(node->input_ids[1])]; - if (weight_node->output_buffer.get_data()) { - size_t num_values = std::min(size_t(5), weight_node->output_buffer.total_size); - weights_str = " weights=["; - - if (weight_node->output_buffer.precision == Precision::FP32) { - const float* float_data = weight_node->output_buffer.data_as(); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) weights_str += ","; - weights_str += std::to_string(float_data[i]).substr(0, 6); - } - } else if (weight_node->output_buffer.precision == Precision::FP16) { - const __fp16* fp16_data = weight_node->output_buffer.data_as<__fp16>(); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) weights_str += ","; - weights_str += std::to_string(static_cast(fp16_data[i])).substr(0, 6); - } - } else if (weight_node->output_buffer.precision == Precision::INT8) { - const int8_t* int8_data = weight_node->output_buffer.data_as(); - for (size_t i = 0; i < num_values; ++i) { - if (i > 0) weights_str += ","; - weights_str += std::to_string(static_cast(int8_data[i])); - } - } - - if (weight_node->output_buffer.total_size > 5) { - weights_str += ",..."; - } - weights_str += "]"; - } - } - - *out << std::left << std::setw(15) << get_op_name(node->op_type) - << std::setw(12) << std::fixed << std::setprecision(3) << ms - << std::setw(20) << shape_str - << values_str << weights_str << std::endl; - } else { - compute_node_optimized(*node, nodes_, node_index_map_); - } - } - - std::unique_ptr capture_file_stream; - std::vector capture_outputs; - - if (capture_requested) { - if (capture_to_stdout) { - capture_outputs.push_back(&std::cout); - } - - if (!capture_file_path.empty()) { - std::filesystem::path capture_path(capture_file_path); - if (capture_path.has_parent_path()) { - std::error_code ec; - std::filesystem::create_directories(capture_path.parent_path(), ec); - } - - auto stream_ptr = std::make_unique(capture_path, std::ios::out | std::ios::app); - if (stream_ptr->is_open()) { - capture_outputs.push_back(stream_ptr.get()); - capture_file_stream = std::move(stream_ptr); - } else { - std::cerr << "Failed to open capture file: " << capture_path << std::endl; - } - } - - if (capture_outputs.empty()) { - capture_requested = false; - } - } - - if (capture_requested) { - auto precision_to_string = [](Precision p) -> const char* { - switch (p) { - case Precision::FP32: return "FP32"; - case Precision::FP16: return "FP16"; - case Precision::INT8: return "INT8"; - default: return "UNKNOWN"; - } - }; - - auto format_double = [](double value) { - std::ostringstream oss; - oss << std::fixed << std::setprecision(6) << value; - return oss.str(); - }; - - auto now = std::chrono::system_clock::now(); - std::time_t now_time = std::chrono::system_clock::to_time_t(now); - std::tm time_info{}; -#if defined(_WIN32) - localtime_s(&time_info, &now_time); -#else - localtime_r(&now_time, &time_info); -#endif - - auto write_header = [&](std::ostream& stream) { - stream << "=== Graph Debug Capture ===" << std::endl; - stream << "Timestamp: " << std::put_time(&time_info, "%Y-%m-%d %H:%M:%S") << std::endl; - stream << "Captured nodes: " << debug_nodes_.size() << std::endl; - stream << std::string(60, '-') << std::endl; - }; - - auto write_separator = [](std::ostream& stream) { - stream << std::string(60, '-') << std::endl; - }; - - if (debug_nodes_.empty()) { - for (auto* stream : capture_outputs) { - write_header(*stream); - *stream << "No debug nodes registered on this graph." << std::endl; - write_separator(*stream); - stream->flush(); - } - } else { - for (auto* stream : capture_outputs) { - write_header(*stream); - } - - for (const auto& entry : debug_nodes_) { - auto node_it = node_index_map_.find(entry.node_id); - const GraphNode* node_ptr = nullptr; - if (node_it != node_index_map_.end()) { - node_ptr = nodes_[node_it->second].get(); - } - - if (!node_ptr) { - for (auto* stream : capture_outputs) { - *stream << "Layer " << entry.layer_idx << " - " << entry.name - << " (node " << entry.node_id << ")" << std::endl; - *stream << " Data: " << std::endl; - write_separator(*stream); - } - continue; - } - - const BufferDesc& buffer = node_ptr->output_buffer; - const void* data_ptr = buffer.get_data(); - size_t total_size = buffer.total_size; - - std::ostringstream shape_ss; - shape_ss << "["; - for (size_t i = 0; i < buffer.shape.size(); ++i) { - if (i > 0) { - shape_ss << ","; - } - shape_ss << buffer.shape[i]; - } - shape_ss << "]"; - std::string shape_str = shape_ss.str(); - - bool has_data = data_ptr != nullptr && total_size > 0; - size_t elements_to_process = total_size; - bool truncated = false; - if (has_data && elements_to_process > capture_max_elements && capture_max_elements > 0) { - elements_to_process = capture_max_elements; - truncated = true; - } - - std::vector preview_values; - if (capture_preview_count > 0) { - preview_values.reserve(std::min(capture_preview_count, elements_to_process)); - } - - double min_val = std::numeric_limits::infinity(); - double max_val = -std::numeric_limits::infinity(); - long double sum = 0.0L; - long double sum_sq = 0.0L; - - if (has_data && elements_to_process > 0) { - auto accumulate = [&](float value, size_t index) { - double v = static_cast(value); - min_val = std::min(min_val, v); - max_val = std::max(max_val, v); - sum += static_cast(value); - sum_sq += static_cast(value) * static_cast(value); - if (capture_preview_count > 0 && index < capture_preview_count) { - preview_values.push_back(value); - } - }; - - if (buffer.precision == Precision::FP32) { - const float* typed = static_cast(data_ptr); - for (size_t i = 0; i < elements_to_process; ++i) { - accumulate(typed[i], i); - } - } else if (buffer.precision == Precision::FP16) { - const __fp16* typed = reinterpret_cast(data_ptr); - for (size_t i = 0; i < elements_to_process; ++i) { - accumulate(static_cast(typed[i]), i); - } - } else if (buffer.precision == Precision::INT8) { - const int8_t* typed = reinterpret_cast(data_ptr); - for (size_t i = 0; i < elements_to_process; ++i) { - accumulate(static_cast(typed[i]), i); - } - } else { - has_data = false; - } - } else { - has_data = false; - } - - size_t processed_count = has_data ? elements_to_process : 0; - long double mean_ld = processed_count > 0 ? sum / processed_count : 0.0L; - long double variance_ld = processed_count > 0 ? (sum_sq / processed_count) - (mean_ld * mean_ld) : 0.0L; - if (variance_ld < 0.0L) { - variance_ld = 0.0L; - } - double mean_val = static_cast(mean_ld); - double stddev_val = processed_count > 0 ? std::sqrt(static_cast(variance_ld)) : 0.0; - - std::ostringstream preview_ss; - if (capture_preview_count > 0 && !preview_values.empty()) { - preview_ss << "["; - for (size_t i = 0; i < preview_values.size(); ++i) { - if (i > 0) { - preview_ss << ", "; - } - preview_ss << format_double(static_cast(preview_values[i])); - } - if (processed_count > preview_values.size()) { - if (!preview_values.empty()) { - preview_ss << ", ..."; - } else { - preview_ss << "..."; - } - } - preview_ss << "]"; - } - - for (auto* stream : capture_outputs) { - *stream << "Layer " << entry.layer_idx << " - " << entry.name - << " (node " << entry.node_id << ")" << std::endl; - *stream << " Shape: " << shape_str << " Precision: " << precision_to_string(buffer.precision) << std::endl; - if (!has_data) { - *stream << " Data: " << std::endl; - } else { - *stream << " Stats: min=" << format_double(min_val) - << " max=" << format_double(max_val) - << " mean=" << format_double(mean_val) - << " std=" << format_double(stddev_val) << std::endl; - if (truncated || processed_count < total_size) { - *stream << " Note: stats computed on first " << processed_count - << " of " << total_size << " values" << std::endl; - } - if (capture_preview_count > 0 && !preview_values.empty()) { - *stream << " Preview: " << preview_ss.str() << std::endl; - } - } - write_separator(*stream); - } - } - - for (auto* stream : capture_outputs) { - stream->flush(); - } - } - } - - if (enable_profiling) { - auto total_end = std::chrono::high_resolution_clock::now(); - auto total_duration = std::chrono::duration_cast(total_end - total_start); - double total_ms = total_duration.count() / 1000.0; - - *out << std::string(60, '-') << std::endl; - *out << "Total execution time: " << std::fixed << std::setprecision(3) << total_ms << " ms" << std::endl; - *out << "================================" << std::endl; - - if (profile_out.is_open()) { - profile_out.close(); - } - } -} - -void CactusGraph::hard_reset() { - nodes_.clear(); - node_index_map_.clear(); - mapped_files_.clear(); - weight_cache_.clear(); - next_node_id_ = 0; - debug_nodes_.clear(); - buffer_pool_.clear(); -} - -void CactusGraph::soft_reset() { - std::set cached_node_ids; - for (const auto& cache_entry : weight_cache_) { - cached_node_ids.insert(cache_entry.second); - } - - size_t max_preserved_id = 0; - for (const auto& node : nodes_) { - if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || - cached_node_ids.count(node->id)) { - max_preserved_id = std::max(max_preserved_id, node->id); - } - } - - auto preserved_nodes = std::move(nodes_); - auto preserved_index_map = std::move(node_index_map_); - - nodes_.clear(); - node_index_map_.clear(); - - for (auto& node : preserved_nodes) { - if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || - cached_node_ids.count(node->id)) { - size_t index = nodes_.size(); - node_index_map_[node->id] = index; - nodes_.push_back(std::move(node)); - } - } - - next_node_id_ = max_preserved_id + 1; - debug_nodes_.clear(); - if (!prefill_mode_) { - buffer_pool_.clear(); - shrink_thread_local_buffers(); - } -} - -void CactusGraph::soft_reset_keep_pool() { - std::set cached_node_ids; - for (const auto& cache_entry : weight_cache_) { - cached_node_ids.insert(cache_entry.second); - } - - size_t max_preserved_id = 0; - for (const auto& node : nodes_) { - if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || - cached_node_ids.count(node->id)) { - max_preserved_id = std::max(max_preserved_id, node->id); - } - } - - auto preserved_nodes = std::move(nodes_); - - nodes_.clear(); - node_index_map_.clear(); - - for (auto& node : preserved_nodes) { - if ((node->op_type == OpType::INPUT && node->output_buffer.external_data) || - cached_node_ids.count(node->id)) { - size_t index = nodes_.size(); - node_index_map_[node->id] = index; - nodes_.push_back(std::move(node)); - } - } - - next_node_id_ = max_preserved_id + 1; - debug_nodes_.clear(); -} diff --git a/cactus/graph/graph_io.cpp b/cactus/graph/graph_io.cpp deleted file mode 100644 index ec6c92cfb..000000000 --- a/cactus/graph/graph_io.cpp +++ /dev/null @@ -1,607 +0,0 @@ -#include "graph.h" -#include "../kernel/kernel.h" -#include -#include -#include -#include -#include -#include -#include - -namespace { - constexpr uint32_t CACTUS_MAGIC = 0x54434143; // "CACT" in little-endian - constexpr uint32_t TENSOR_FORMAT_VERSION = 1; - constexpr uint32_t FLAG_HAS_SCALES = 1 << 0; - constexpr size_t HEADER_SIZE = 80; - - inline size_t align_offset(size_t offset, size_t alignment) { - size_t remainder = offset % alignment; - if (remainder == 0) return offset; - return offset + (alignment - remainder); - } -} - - -size_t CactusGraph::mmap_embeddings(const std::string& filename) { - auto mapped_file = std::make_unique(filename); - - const auto& shape = mapped_file->shape(); - if (shape.size() != 2) { - throw std::runtime_error("Memory-mapped embeddings must be 2D [vocab_size, embedding_dim]"); - } - - Precision precision = mapped_file->effective_precision(); - - size_t node_id = input(shape, precision); - set_external_input(node_id, const_cast(mapped_file->data()), precision); - - if (precision == Precision::INT8 && mapped_file->group_size() > 0) { - set_grouped_scales(node_id, mapped_file->group_size(), mapped_file->num_groups(), - const_cast(mapped_file->scales_data())); - } - - size_t file_idx = mapped_files_.size(); - mapped_files_.push_back(std::move(mapped_file)); - node_to_mapped_file_[node_id] = file_idx; - return node_id; -} - -size_t CactusGraph::mmap_weights(const std::string& filename) { - auto it = weight_cache_.find(filename); - if (it != weight_cache_.end()) { - return it->second; - } - - auto mapped_file = std::make_unique(filename); - - const auto& shape = mapped_file->shape(); - Precision precision = mapped_file->effective_precision(); - - size_t node_id = input(shape, precision); - set_external_input(node_id, const_cast(mapped_file->data()), precision); - - if (precision == Precision::INT8 && mapped_file->group_size() > 0) { - set_grouped_scales(node_id, mapped_file->group_size(), mapped_file->num_groups(), - const_cast(mapped_file->scales_data())); - } - - size_t file_idx = mapped_files_.size(); - mapped_files_.push_back(std::move(mapped_file)); - node_to_mapped_file_[node_id] = file_idx; - weight_cache_[filename] = node_id; - return node_id; -} - -void CactusGraph::release_weight_pages(size_t node_id) { - auto it = node_to_mapped_file_.find(node_id); - if (it != node_to_mapped_file_.end() && it->second < mapped_files_.size()) { - mapped_files_[it->second]->release_pages(); - } -} - -void CactusGraph::prefetch_weight_pages(size_t node_id) { - auto it = node_to_mapped_file_.find(node_id); - if (it != node_to_mapped_file_.end() && it->second < mapped_files_.size()) { - mapped_files_[it->second]->prefetch_pages(); - } -} - -void CactusGraph::release_all_weight_pages() { - for (auto& mf : mapped_files_) { - if (mf) mf->release_pages(); - } -} - -size_t CactusGraph::load_weights(const std::string& filename) { - auto it = weight_cache_.find(filename); - if (it != weight_cache_.end()) { - return it->second; - } - - auto loaded = GraphFile::load_into_graph(*this, filename); - weight_cache_[filename] = loaded.node_id; - return loaded.node_id; -} - -void CactusGraph::set_grouped_scales(size_t node_id, size_t group_size, size_t num_groups, void* scales_ptr) { - auto it = node_index_map_.find(node_id); - if (it != node_index_map_.end()) { - nodes_[it->second]->output_buffer.set_grouped_scales(group_size, num_groups, scales_ptr); - } -} - -size_t CactusGraph::embedding(const std::string& filename, size_t indices) { - auto mapped_file = std::make_unique(filename); - - const auto& shape = mapped_file->shape(); - if (shape.size() != 2) { - throw std::runtime_error("Embedding file must contain 2D tensor [vocab_size, hidden_dim]"); - } - - Precision precision = mapped_file->effective_precision(); - size_t embeddings_node = input(shape, precision); - set_external_input(embeddings_node, const_cast(mapped_file->data()), precision); - - if (precision == Precision::INT8 && mapped_file->group_size() > 0) { - set_grouped_scales(embeddings_node, mapped_file->group_size(), mapped_file->num_groups(), - const_cast(mapped_file->scales_data())); - } - - mapped_files_.push_back(std::move(mapped_file)); - - const auto& idx_shape = get_output_buffer(indices).shape; - std::vector output_shape = idx_shape; - output_shape.push_back(shape[1]); - - OpParams params; - params.output_precision = (precision == Precision::INT8) ? Precision::FP16 : precision; - - return add_node(OpType::EMBEDDING, {embeddings_node, indices}, output_shape, params); -} - - -namespace GraphFile { - -void save_node(CactusGraph& graph, size_t node_id, const std::string& filename) { - graph.execute(); - void* data = graph.get_output(node_id); - - const auto& buffer = graph.get_output_buffer(node_id); - const auto& shape = buffer.shape; - Precision precision = buffer.precision; - - std::ofstream file(filename, std::ios::binary); - if (!file) { - throw std::runtime_error("Cannot open file for writing: " + filename); - } - - size_t total_elements = 1; - for (size_t dim : shape) { - total_elements *= dim; - } - - size_t element_size = PrecisionTraits::size_of(precision); - size_t byte_size = total_elements * element_size; - - bool has_scales = (precision == Precision::INT8 && buffer.is_grouped_int8() && buffer.scales_data); - size_t N = shape.size() >= 1 ? shape[0] : 1; - size_t scales_bytes = has_scales ? (N * buffer.num_groups * sizeof(__fp16)) : 0; - - uint32_t ndim = static_cast(shape.size()); - uint32_t flags = has_scales ? FLAG_HAS_SCALES : 0; - uint32_t alignment = 32; - - uint32_t magic = CACTUS_MAGIC; - uint32_t version = TENSOR_FORMAT_VERSION; - file.write(reinterpret_cast(&magic), sizeof(magic)); - file.write(reinterpret_cast(&version), sizeof(version)); - file.write(reinterpret_cast(&flags), sizeof(flags)); - file.write(reinterpret_cast(&alignment), sizeof(alignment)); - file.write(reinterpret_cast(&ndim), sizeof(ndim)); - - for (uint32_t i = 0; i < 4; i++) { - uint64_t dim_val = (i < shape.size()) ? static_cast(shape[i]) : 0; - file.write(reinterpret_cast(&dim_val), sizeof(dim_val)); - } - - uint32_t prec_val = static_cast(precision); - file.write(reinterpret_cast(&prec_val), sizeof(prec_val)); - - uint64_t data_bytes = static_cast(byte_size); - uint64_t scales_bytes_val = static_cast(scales_bytes); - file.write(reinterpret_cast(&data_bytes), sizeof(data_bytes)); - file.write(reinterpret_cast(&scales_bytes_val), sizeof(scales_bytes_val)); - - uint32_t group_size = has_scales ? static_cast(buffer.group_size) : 0; - uint32_t num_groups = has_scales ? static_cast(buffer.num_groups) : 0; - file.write(reinterpret_cast(&group_size), sizeof(group_size)); - file.write(reinterpret_cast(&num_groups), sizeof(num_groups)); - - size_t header_end = 80; - size_t aligned_header = align_offset(header_end, alignment); - size_t header_padding = aligned_header - header_end; - for (size_t i = 0; i < header_padding; i++) { - char zero = 0; - file.write(&zero, 1); - } - - if (has_scales) { - file.write(static_cast(buffer.scales_data), scales_bytes); - - size_t scales_end = aligned_header + scales_bytes; - size_t data_start = align_offset(scales_end, alignment); - size_t scales_padding = data_start - scales_end; - for (size_t i = 0; i < scales_padding; i++) { - char zero = 0; - file.write(&zero, 1); - } - } - - file.write(static_cast(data), byte_size); - - if (!file) { - throw std::runtime_error("Error writing node data to file: " + filename); - } -} - -LoadedNode load_into_graph(CactusGraph& graph, const std::string& filename) { - std::ifstream file(filename, std::ios::binary); - if (!file) { - throw std::runtime_error("Cannot open file for reading: " + filename); - } - - char header[HEADER_SIZE]; - file.read(header, HEADER_SIZE); - if (!file) { - throw std::runtime_error("Error reading file header: " + filename); - } - - const char* ptr = header; - size_t offset = 0; - - uint32_t magic = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - if (magic != CACTUS_MAGIC) { - throw std::runtime_error("Invalid tensor file: missing CACT magic number"); - } - - offset += sizeof(uint32_t); - offset += sizeof(uint32_t); - - uint32_t alignment = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - if (alignment == 0) alignment = 1; - - uint32_t ndim = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - std::vector shape; - for (uint32_t i = 0; i < 4; i++) { - uint64_t dim_val = *reinterpret_cast(ptr + offset); - offset += sizeof(uint64_t); - if (i < ndim && dim_val > 0) { - shape.push_back(static_cast(dim_val)); - } - } - - uint32_t prec_val = *reinterpret_cast(ptr + offset); - Precision precision = static_cast(prec_val); - offset += sizeof(uint32_t); - - size_t byte_size = *reinterpret_cast(ptr + offset); - offset += sizeof(uint64_t); - - size_t scales_bytes = *reinterpret_cast(ptr + offset); - offset += sizeof(uint64_t); - - size_t group_size = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - size_t num_groups = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - size_t aligned_header = align_offset(HEADER_SIZE, alignment); - size_t padding = aligned_header - HEADER_SIZE; - if (padding > 0) { - file.seekg(padding, std::ios::cur); - } - - std::vector scales_buffer; - if (scales_bytes > 0) { - scales_buffer.resize(scales_bytes); - file.read(scales_buffer.data(), scales_bytes); - - size_t scales_end = aligned_header + scales_bytes; - size_t data_start = align_offset(scales_end, alignment); - size_t scales_padding = data_start - scales_end; - if (scales_padding > 0) { - file.seekg(scales_padding, std::ios::cur); - } - } - - std::vector buffer(byte_size); - file.read(buffer.data(), byte_size); - - if (!file || file.gcount() != static_cast(byte_size)) { - throw std::runtime_error("Error reading node data: " + filename); - } - - size_t node_id = graph.input(shape, precision); - graph.set_input(node_id, buffer.data(), precision); - - if (scales_bytes > 0 && group_size > 0 && num_groups > 0) { - auto& node_buffer = graph.nodes_[graph.node_index_map_.at(node_id)]->output_buffer; - node_buffer.owned_scales = std::make_unique(scales_buffer.size()); - std::memcpy(node_buffer.owned_scales.get(), scales_buffer.data(), scales_buffer.size()); - node_buffer.set_grouped_scales(group_size, num_groups, node_buffer.owned_scales.get()); - } - - return {node_id, shape, precision, byte_size}; -} - -MappedFile mmap_load(const std::string& filename) { - return MappedFile(filename); -} - -// MappedFile implementation - -MappedFile::MappedFile(const std::string& filename) - : fd_(-1), mapped_data_(nullptr), file_size_(0), data_offset_(0) { - fd_ = open(filename.c_str(), O_RDONLY); - if (fd_ == -1) { - throw std::runtime_error("Cannot open file for mapping: " + filename); - } - - struct stat st; - if (fstat(fd_, &st) == -1) { - close(fd_); - throw std::runtime_error("Cannot get file size: " + filename); - } - file_size_ = st.st_size; - - mapped_data_ = mmap(nullptr, file_size_, PROT_READ, MAP_SHARED, fd_, 0); - if (mapped_data_ == MAP_FAILED) { - close(fd_); - throw std::runtime_error("Cannot map file: " + filename); - } - - close(fd_); - fd_ = -1; - - parse_header(); - apply_madvise_hints(); -} - -MappedFile::~MappedFile() { - if (mapped_data_ != nullptr && mapped_data_ != MAP_FAILED) { - madvise(mapped_data_, file_size_, MADV_DONTNEED); - munmap(mapped_data_, file_size_); - mapped_data_ = nullptr; - } - if (fd_ != -1) { - close(fd_); - fd_ = -1; - } -} - -MappedFile::MappedFile(MappedFile&& other) noexcept - : fd_(other.fd_), mapped_data_(other.mapped_data_), file_size_(other.file_size_), - data_offset_(other.data_offset_), shape_(std::move(other.shape_)), - precision_(other.precision_), byte_size_(other.byte_size_), - group_size_(other.group_size_), num_groups_(other.num_groups_), - scales_offset_(other.scales_offset_), scales_bytes_(other.scales_bytes_), - version_(other.version_), alignment_(other.alignment_), - is_int4_(other.is_int4_), - unpacked_int4_data_(std::move(other.unpacked_int4_data_)) { - other.fd_ = -1; - other.mapped_data_ = nullptr; - other.file_size_ = 0; - other.is_int4_ = false; -} - -MappedFile& MappedFile::operator=(MappedFile&& other) noexcept { - if (this != &other) { - if (mapped_data_ != nullptr && mapped_data_ != MAP_FAILED) { - munmap(mapped_data_, file_size_); - } - if (fd_ != -1) { - close(fd_); - } - - fd_ = other.fd_; - mapped_data_ = other.mapped_data_; - file_size_ = other.file_size_; - data_offset_ = other.data_offset_; - shape_ = std::move(other.shape_); - precision_ = other.precision_; - byte_size_ = other.byte_size_; - group_size_ = other.group_size_; - num_groups_ = other.num_groups_; - scales_offset_ = other.scales_offset_; - scales_bytes_ = other.scales_bytes_; - version_ = other.version_; - alignment_ = other.alignment_; - is_int4_ = other.is_int4_; - unpacked_int4_data_ = std::move(other.unpacked_int4_data_); - - other.fd_ = -1; - other.mapped_data_ = nullptr; - other.file_size_ = 0; - other.is_int4_ = false; - } - return *this; -} - -const std::vector& MappedFile::shape() const { - return shape_; -} - -Precision MappedFile::precision() const { - return precision_; -} - -size_t MappedFile::byte_size() const { - return byte_size_; -} - -const void* MappedFile::scales_data() const { - return static_cast(mapped_data_) + scales_offset_; -} - -const void* MappedFile::raw_packed_data() const { - return static_cast(mapped_data_) + data_offset_; -} - -void MappedFile::unpack_int4_if_needed() const { - if (!is_int4_ || unpacked_int4_data_) { - return; - } - - size_t unpacked_count = byte_size_ * 2; - unpacked_int4_data_ = std::make_unique(unpacked_count); - - const uint8_t* packed = reinterpret_cast( - static_cast(mapped_data_) + data_offset_); - - cactus_unpack_int4_to_int8(packed, unpacked_int4_data_.get(), unpacked_count); -} - -void* MappedFile::data() { - if (is_int4_) { - unpack_int4_if_needed(); - return unpacked_int4_data_.get(); - } - return static_cast(mapped_data_) + data_offset_; -} - -const void* MappedFile::data() const { - if (is_int4_) { - unpack_int4_if_needed(); - return unpacked_int4_data_.get(); - } - return static_cast(mapped_data_) + data_offset_; -} - -template -const T* MappedFile::typed_data() const { - return static_cast(data()); -} - -LoadedNode MappedFile::load_into_graph(CactusGraph& graph) const { - Precision eff_prec = effective_precision(); - - size_t node_id = graph.input(shape_, eff_prec); - - if (is_int4_) { - auto& node_buffer = graph.nodes_[graph.node_index_map_.at(node_id)]->output_buffer; - node_buffer.set_external(const_cast(raw_packed_data())); - node_buffer.set_packed_int4(raw_packed_data(), byte_size_); - node_buffer.byte_size = byte_size_; - - if (group_size_ > 0) { - graph.set_grouped_scales(node_id, group_size_, num_groups_, - const_cast(scales_data())); - } - } else { - graph.set_external_input(node_id, const_cast(data()), eff_prec); - - if ((eff_prec == Precision::INT8) && group_size_ > 0) { - graph.set_grouped_scales(node_id, group_size_, num_groups_, - const_cast(scales_data())); - } - } - - return {node_id, shape_, eff_prec, byte_size_}; -} - -void MappedFile::parse_header() { - if (file_size_ < HEADER_SIZE) { - throw std::runtime_error("File too small: insufficient data for header"); - } - - const char* ptr = static_cast(mapped_data_); - size_t offset = 0; - - uint32_t magic = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - if (magic != CACTUS_MAGIC) { - throw std::runtime_error("Invalid tensor file: missing CACT magic number"); - } - - version_ = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - uint32_t flags = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - (void)flags; - - alignment_ = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - if (alignment_ == 0) alignment_ = 1; - - uint32_t ndim = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - shape_.clear(); - for (uint32_t i = 0; i < 4; i++) { - uint64_t dim_val = *reinterpret_cast(ptr + offset); - offset += sizeof(uint64_t); - if (i < ndim && dim_val > 0) { - shape_.push_back(static_cast(dim_val)); - } - } - - uint32_t prec_val = *reinterpret_cast(ptr + offset); - precision_ = static_cast(prec_val); - offset += sizeof(uint32_t); - - is_int4_ = (precision_ == Precision::INT4); - - byte_size_ = *reinterpret_cast(ptr + offset); - offset += sizeof(uint64_t); - - scales_bytes_ = *reinterpret_cast(ptr + offset); - offset += sizeof(uint64_t); - - group_size_ = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - num_groups_ = *reinterpret_cast(ptr + offset); - offset += sizeof(uint32_t); - - size_t aligned_header = align_offset(HEADER_SIZE, alignment_); - - if (scales_bytes_ > 0) { - scales_offset_ = aligned_header; - size_t scales_end = scales_offset_ + scales_bytes_; - data_offset_ = align_offset(scales_end, alignment_); - } else { - scales_offset_ = 0; - data_offset_ = aligned_header; - } - - if (data_offset_ + byte_size_ > file_size_) { - throw std::runtime_error("File corrupted: data extends beyond file size (v2)"); - } -} - -void MappedFile::apply_madvise_hints() { - if (scales_bytes_ > 0 && scales_offset_ > 0) { - madvise(static_cast(mapped_data_) + scales_offset_, scales_bytes_, MADV_WILLNEED); - } - - madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_SEQUENTIAL); - - if (byte_size_ > 1024 * 1024) { - madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_WILLNEED); - } -} - -void MappedFile::release_pages() { - if (mapped_data_ == nullptr || mapped_data_ == MAP_FAILED) return; - - if (scales_bytes_ > 0 && scales_offset_ > 0) { - madvise(static_cast(mapped_data_) + scales_offset_, scales_bytes_, MADV_DONTNEED); - } - madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_DONTNEED); - - if (is_int4_ && unpacked_int4_data_) { - unpacked_int4_data_.reset(); - } -} - -void MappedFile::prefetch_pages() { - if (mapped_data_ == nullptr || mapped_data_ == MAP_FAILED) return; - - if (scales_bytes_ > 0 && scales_offset_ > 0) { - madvise(static_cast(mapped_data_) + scales_offset_, scales_bytes_, MADV_WILLNEED); - } - madvise(static_cast(mapped_data_) + data_offset_, byte_size_, MADV_WILLNEED); -} - -template const int8_t* MappedFile::typed_data() const; -template const float* MappedFile::typed_data() const; -template const uint16_t* MappedFile::typed_data() const; -template const uint8_t* MappedFile::typed_data() const; - -} // namespace GraphFile diff --git a/cactus/graph/graph_ops_math.cpp b/cactus/graph/graph_ops_math.cpp deleted file mode 100644 index a59359b8f..000000000 --- a/cactus/graph/graph_ops_math.cpp +++ /dev/null @@ -1,376 +0,0 @@ -#include "graph.h" -#include "../kernel/kernel.h" -#include -#include - -namespace Quantization { - void int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale) { - cactus_int8_to_fp32(src, dst, count, scale); - } - - void fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale) { - cactus_fp32_to_int8(src, dst, count, scale); - } - - void fp16_to_fp32(const __fp16* src, float* dst, size_t count) { - cactus_fp16_to_fp32(src, dst, count); - } - - void fp32_to_fp16(const float* src, __fp16* dst, size_t count) { - cactus_fp32_to_fp16(src, dst, count); - } - - void int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale) { - cactus_int8_to_fp16(src, dst, count, scale); - } - - void fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale) { - cactus_fp16_to_int8(src, dst, count, scale); - } -} - -static std::vector compute_strides(const std::vector& shape, const std::vector& target_shape) { - std::vector strides(target_shape.size()); - - size_t shape_offset = target_shape.size() - shape.size(); - - for (size_t i = 0; i < target_shape.size(); ++i) { - if (i < shape_offset) { - strides[i] = 0; - } else { - size_t dim_idx = i - shape_offset; - if (shape[dim_idx] == 1) { - strides[i] = 0; - } else { - strides[i] = 1; - for (size_t j = dim_idx + 1; j < shape.size(); ++j) { - strides[i] *= shape[j]; - } - } - } - } - - return strides; -} - -void dispatch_binary_op_f16(OpType op, const __fp16* lhs, const __fp16* rhs, __fp16* output, size_t count) { - switch (op) { - case OpType::ADD: - cactus_add_f16(lhs, rhs, output, count); - break; - case OpType::ADD_CLIPPED: - cactus_add_f16_clipped(lhs, rhs, output, count); - break; - case OpType::SUBTRACT: - cactus_subtract_f16(lhs, rhs, output, count); - break; - case OpType::MULTIPLY: - cactus_multiply_f16(lhs, rhs, output, count); - break; - case OpType::DIVIDE: - cactus_divide_f16(lhs, rhs, output, count); - break; - default: - break; - } -} - -void dispatch_unary_op_f16(OpType op, const __fp16* input, __fp16* output, size_t count, float param) { - ScalarOpType scalar_op; - switch (op) { - case OpType::SCALAR_ADD: scalar_op = ScalarOpType::ADD; break; - case OpType::SCALAR_SUBTRACT: scalar_op = ScalarOpType::SUBTRACT; break; - case OpType::SCALAR_MULTIPLY: scalar_op = ScalarOpType::MULTIPLY; break; - case OpType::SCALAR_DIVIDE: scalar_op = ScalarOpType::DIVIDE; break; - case OpType::SCALAR_EXP: scalar_op = ScalarOpType::EXP; break; - case OpType::SCALAR_SQRT: scalar_op = ScalarOpType::SQRT; break; - case OpType::SCALAR_COS: scalar_op = ScalarOpType::COS; break; - case OpType::SCALAR_SIN: scalar_op = ScalarOpType::SIN; break; - default: return; - } - - cactus_scalar_op_f16(input, output, count, param, scalar_op); -} - -void compute_binary_op_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& lhs = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& rhs = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - - if (lhs.precision != Precision::FP16) { - throw std::runtime_error("Binary operations only support FP16 precision"); - } - - if (node.params.broadcast_info.needs_broadcasting) { - std::vector lhs_strides = compute_strides(lhs.shape, node.params.broadcast_info.output_shape); - std::vector rhs_strides = compute_strides(rhs.shape, node.params.broadcast_info.output_shape); - - switch (node.op_type) { - case OpType::ADD: - case OpType::ADD_CLIPPED: - cactus_add_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - lhs_strides.data(), rhs_strides.data(), - node.params.broadcast_info.output_shape.data(), - node.params.broadcast_info.output_shape.size()); - break; - case OpType::SUBTRACT: - cactus_subtract_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - lhs_strides.data(), rhs_strides.data(), - node.params.broadcast_info.output_shape.data(), - node.params.broadcast_info.output_shape.size()); - break; - case OpType::MULTIPLY: - cactus_multiply_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - lhs_strides.data(), rhs_strides.data(), - node.params.broadcast_info.output_shape.data(), - node.params.broadcast_info.output_shape.size()); - break; - case OpType::DIVIDE: - cactus_divide_broadcast_f16(lhs.data_as<__fp16>(), rhs.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - lhs_strides.data(), rhs_strides.data(), - node.params.broadcast_info.output_shape.data(), - node.params.broadcast_info.output_shape.size()); - break; - default: break; - } - } else { - dispatch_binary_op_f16(node.op_type, lhs.data_as<__fp16>(), - rhs.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size); - } -} - -void compute_unary_op_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - - if (input.precision != Precision::FP16) { - throw std::runtime_error("Scalar operations only support FP16 precision"); - } - - dispatch_unary_op_f16(node.op_type, input.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size, node.params.scalar); -} - -void compute_activation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - - if (input.precision != Precision::FP16) { - throw std::runtime_error("Activation operations only support FP16 precision"); - } - - switch (node.op_type) { - case OpType::SILU: - cactus_silu_f16(input.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size); - break; - case OpType::GELU: - cactus_gelu_f16(input.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size); - break; - case OpType::GELU_ERF: - cactus_gelu_f16_erf(input.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size); - break; - default: - break; - } -} - -void compute_reduce_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - int axis = node.params.axis; - - if (input_buffer.precision != Precision::FP16) { - throw std::runtime_error("Reduction operations only support FP16 precision"); - } - - if (axis == -1) { - switch (node.op_type) { - case OpType::SUM: { - double result = cactus_sum_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); - node.output_buffer.data_as<__fp16>()[0] = static_cast<__fp16>(result); - break; - } - case OpType::MEAN: { - double result = cactus_mean_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); - node.output_buffer.data_as<__fp16>()[0] = static_cast<__fp16>(result); - break; - } - case OpType::VARIANCE: { - double result = cactus_variance_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); - node.output_buffer.data_as<__fp16>()[0] = static_cast<__fp16>(result); - break; - } - case OpType::MIN: { - __fp16 result = cactus_min_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); - node.output_buffer.data_as<__fp16>()[0] = result; - break; - } - case OpType::MAX: { - __fp16 result = cactus_max_all_f16(input_buffer.data_as<__fp16>(), input_buffer.total_size); - node.output_buffer.data_as<__fp16>()[0] = result; - break; - } - default: break; - } - } else { - const auto& shape = input_buffer.shape; - size_t axis_idx = static_cast(axis); - - size_t outer_size = 1; - for (size_t i = 0; i < axis_idx; i++) { - outer_size *= shape[i]; - } - - size_t axis_size = shape[axis_idx]; - - size_t inner_size = 1; - for (size_t i = axis_idx + 1; i < shape.size(); i++) { - inner_size *= shape[i]; - } - - switch (node.op_type) { - case OpType::SUM: - cactus_sum_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - outer_size, axis_size, inner_size); - break; - case OpType::MEAN: - cactus_mean_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - outer_size, axis_size, inner_size); - break; - case OpType::VARIANCE: - cactus_variance_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - outer_size, axis_size, inner_size); - break; - case OpType::MIN: - cactus_min_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - outer_size, axis_size, inner_size); - break; - case OpType::MAX: - cactus_max_axis_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - outer_size, axis_size, inner_size); - break; - default: break; - } - } -} - -void compute_reshape_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - - if (input_buffer.is_packed_int4()) { - throw std::runtime_error("Reshape operation not supported on packed INT4 data"); - } - - size_t input_total_elements = input_buffer.total_size; - size_t output_total_elements = node.output_buffer.total_size; - - if (input_total_elements != output_total_elements) { - throw std::runtime_error("Reshape operation: input elements (" + std::to_string(input_total_elements) + - ") must match output elements (" + std::to_string(output_total_elements) + ")"); - } - - std::memcpy(node.output_buffer.get_data(), input_buffer.get_data(), input_buffer.byte_size); -} - -void compute_precision_cast_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_node = *nodes[node_index_map.at(node.input_ids[0])]; - - if (input_node.output_buffer.is_packed_int4()) { - throw std::runtime_error("Precision cast not supported on packed INT4 data"); - } - - if (input_node.output_buffer.precision == node.output_buffer.precision) { - std::memcpy(node.output_buffer.get_data(), input_node.output_buffer.get_data(), input_node.output_buffer.byte_size); - return; - } - - size_t count = input_node.output_buffer.total_size; - - if (input_node.output_buffer.precision == Precision::INT8 && node.output_buffer.precision == Precision::FP32) { - if (input_node.output_buffer.is_grouped_int8()) { - const int8_t* src = input_node.output_buffer.data_as(); - float* dst = node.output_buffer.data_as(); - const __fp16* scales = input_node.output_buffer.scales_as_fp16(); - size_t group_size = input_node.output_buffer.group_size; - - const auto& shape = input_node.output_buffer.shape; - if (shape.size() == 2) { - size_t N = shape[0]; - size_t K = shape[1]; - size_t num_groups = K / group_size; - for (size_t row = 0; row < N; ++row) { - for (size_t col = 0; col < K; ++col) { - size_t idx = row * K + col; - size_t group_idx = col / group_size; - float scale = static_cast(scales[row * num_groups + group_idx]); - dst[idx] = static_cast(src[idx]) * scale; - } - } - } else if (shape.size() == 1) { - size_t K = shape[0]; - for (size_t col = 0; col < K; ++col) { - size_t group_idx = col / group_size; - float scale = static_cast(scales[group_idx]); - dst[col] = static_cast(src[col]) * scale; - } - } else { - Quantization::int8_to_fp32(src, dst, count, 1.0f); - } - } else { - Quantization::int8_to_fp32(input_node.output_buffer.data_as(), node.output_buffer.data_as(), count, 1.0f); - } - } else if (input_node.output_buffer.precision == Precision::FP32 && node.output_buffer.precision == Precision::INT8) { - Quantization::fp32_to_int8(input_node.output_buffer.data_as(), node.output_buffer.data_as(), count, 1.0f); - } else if (input_node.output_buffer.precision == Precision::FP16 && node.output_buffer.precision == Precision::FP32) { - Quantization::fp16_to_fp32(input_node.output_buffer.data_as<__fp16>(), node.output_buffer.data_as(), count); - } else if (input_node.output_buffer.precision == Precision::FP32 && node.output_buffer.precision == Precision::FP16) { - Quantization::fp32_to_fp16(input_node.output_buffer.data_as(), node.output_buffer.data_as<__fp16>(), count); - } else if (input_node.output_buffer.precision == Precision::INT8 && node.output_buffer.precision == Precision::FP16) { - if (input_node.output_buffer.is_grouped_int8()) { - const int8_t* src = input_node.output_buffer.data_as(); - __fp16* dst = node.output_buffer.data_as<__fp16>(); - const __fp16* scales = input_node.output_buffer.scales_as_fp16(); - size_t group_size = input_node.output_buffer.group_size; - - const auto& shape = input_node.output_buffer.shape; - if (shape.size() == 2) { - size_t N = shape[0]; - size_t K = shape[1]; - size_t num_groups = K / group_size; - for (size_t row = 0; row < N; ++row) { - for (size_t col = 0; col < K; ++col) { - size_t idx = row * K + col; - size_t group_idx = col / group_size; - float scale = static_cast(scales[row * num_groups + group_idx]); - dst[idx] = static_cast<__fp16>(src[idx] * scale); - } - } - } else if (shape.size() == 1) { - size_t K = shape[0]; - for (size_t col = 0; col < K; ++col) { - size_t group_idx = col / group_size; - float scale = static_cast(scales[group_idx]); - dst[col] = static_cast<__fp16>(src[col] * scale); - } - } else { - Quantization::int8_to_fp16(src, dst, count, 1.0f); - } - } else { - Quantization::int8_to_fp16(input_node.output_buffer.data_as(), node.output_buffer.data_as<__fp16>(), count, 1.0f); - } - } else if (input_node.output_buffer.precision == Precision::FP16 && node.output_buffer.precision == Precision::INT8) { - Quantization::fp16_to_int8(input_node.output_buffer.data_as<__fp16>(), node.output_buffer.data_as(), count, 1.0f); - } else { - throw std::runtime_error("Unsupported precision conversion from " + - std::to_string(static_cast(input_node.output_buffer.precision)) + - " to " + std::to_string(static_cast(node.output_buffer.precision))); - } -} diff --git a/cactus/graph/graph_ops_nn.cpp b/cactus/graph/graph_ops_nn.cpp deleted file mode 100644 index d753edcda..000000000 --- a/cactus/graph/graph_ops_nn.cpp +++ /dev/null @@ -1,608 +0,0 @@ -#include "graph.h" -#include "../kernel/kernel.h" -#include "../kernel/kernel_utils.h" -#include -#include -#include -#include -#include - -namespace { - thread_local std::vector<__fp16> transpose_buffer_fp16; - thread_local std::vector quant_activation_buffer; - thread_local std::vector quant_scales_buffer; - - thread_local const __fp16* cached_quant_src = nullptr; - thread_local size_t cached_quant_M = 0; - thread_local size_t cached_quant_K = 0; - - void ensure_transpose_buffer_fp16(size_t required_size) { - if (transpose_buffer_fp16.size() < required_size) { - transpose_buffer_fp16.resize(required_size); - } - } - - void ensure_quant_buffers(size_t M, size_t K) { - size_t required_data = M * K; - if (quant_activation_buffer.size() < required_data) { - quant_activation_buffer.resize(required_data); - } - if (quant_scales_buffer.size() < M) { - quant_scales_buffer.resize(M); - } - } - - void quantize_activations_fp16_to_int8(const __fp16* src, int8_t* dst, float* scales, size_t M, size_t K) { - if (src == cached_quant_src && M == cached_quant_M && K == cached_quant_K) { - return; - } - - constexpr size_t PARALLEL_THRESHOLD = 16; - - if (M >= PARALLEL_THRESHOLD) { - CactusThreading::parallel_for(M, CactusThreading::Thresholds::ELEMENT_WISE, - [src, dst, scales, K](size_t m_start, size_t m_end) { - for (size_t m = m_start; m < m_end; m++) { - float max_abs = cactus_fp16_max_abs(src + m * K, K); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - scales[m] = scale; - cactus_fp16_to_int8(src + m * K, dst + m * K, K, scale); - } - }); - } else { - for (size_t m = 0; m < M; m++) { - float max_abs = cactus_fp16_max_abs(src + m * K, K); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - scales[m] = scale; - cactus_fp16_to_int8(src + m * K, dst + m * K, K, scale); - } - } - - cached_quant_src = src; - cached_quant_M = M; - cached_quant_K = K; - } -} - -void shrink_thread_local_buffers() { - std::vector<__fp16>().swap(transpose_buffer_fp16); - std::vector().swap(quant_activation_buffer); - std::vector().swap(quant_scales_buffer); - cached_quant_src = nullptr; - cached_quant_M = 0; - cached_quant_K = 0; -} - -void compute_quantize_activations_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& shape = input_buffer.shape; - - if (input_buffer.precision != Precision::FP16) { - throw std::runtime_error("QUANTIZE_ACTIVATIONS requires FP16 input"); - } - - if (shape.size() < 2) { - throw std::runtime_error("QUANTIZE_ACTIVATIONS requires at least 2D tensor"); - } - - size_t K = shape.back(); - size_t M = 1; - for (size_t i = 0; i < shape.size() - 1; ++i) { - M *= shape[i]; - } - - if (!node.output_buffer.has_activation_scales() || - node.output_buffer.num_rows_for_activation_scales != M) { - node.output_buffer.allocate_activation_scales(M); - } - - const __fp16* src = input_buffer.data_as<__fp16>(); - int8_t* dst = node.output_buffer.data_as(); - float* scales = node.output_buffer.activation_scales_as_float(); - - constexpr size_t PARALLEL_THRESHOLD = 16; - - if (M >= PARALLEL_THRESHOLD) { - CactusThreading::parallel_for(M, CactusThreading::Thresholds::ELEMENT_WISE, - [src, dst, scales, K](size_t m_start, size_t m_end) { - for (size_t m = m_start; m < m_end; m++) { - float max_abs = cactus_fp16_max_abs(src + m * K, K); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - scales[m] = scale; - cactus_fp16_to_int8(src + m * K, dst + m * K, K, scale); - } - }); - } else { - for (size_t m = 0; m < M; m++) { - float max_abs = cactus_fp16_max_abs(src + m * K, K); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - scales[m] = scale; - cactus_fp16_to_int8(src + m * K, dst + m * K, K, scale); - } - } -} - -void compute_matmul_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& lhs_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& rhs_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - const auto& lhs_shape = lhs_buffer.shape; - const auto& rhs_shape = rhs_buffer.shape; - - size_t M = lhs_shape[lhs_shape.size() - 2]; - size_t K = lhs_shape[lhs_shape.size() - 1]; - size_t N = node.params.pretransposed_rhs ? - rhs_shape[rhs_shape.size() - 2] : rhs_shape[rhs_shape.size() - 1]; - - bool pretransposed_rhs = node.params.pretransposed_rhs; - - ComputeBackend backend = node.params.backend; - - if (backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU matrix multiplication not yet implemented"); - } - - const bool lhs_is_prequantized_int8 = (lhs_buffer.precision == Precision::INT8 && - lhs_buffer.has_activation_scales()); - - if (rhs_buffer.is_packed_int4()) { - const uint8_t* rhs_packed = rhs_buffer.packed_int4_as_uint8(); - const __fp16* rhs_scales = rhs_buffer.scales_as_fp16(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (!pretransposed_rhs) { - throw std::runtime_error("INT4 matmul requires pretransposed weights"); - } - - const int8_t* lhs_int8; - const float* lhs_scales; - - if (lhs_is_prequantized_int8) { - lhs_int8 = lhs_buffer.data_as(); - lhs_scales = lhs_buffer.activation_scales_as_float(); - } else if (lhs_buffer.precision == Precision::FP16) { - const __fp16* lhs = lhs_buffer.data_as<__fp16>(); - ensure_quant_buffers(M, K); - quantize_activations_fp16_to_int8(lhs, quant_activation_buffer.data(), - quant_scales_buffer.data(), M, K); - lhs_int8 = quant_activation_buffer.data(); - lhs_scales = quant_scales_buffer.data(); - } else { - throw std::runtime_error("INT4 matmul requires INT8 (pre-quantized) or FP16 activations"); - } - - cactus_matmul_int4(lhs_int8, lhs_scales, - rhs_packed, rhs_scales, output, - M, K, N, rhs_buffer.group_size); - - } else if (rhs_buffer.is_grouped_int8()) { - const int8_t* rhs = rhs_buffer.data_as(); - const __fp16* rhs_scales = rhs_buffer.scales_as_fp16(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (!pretransposed_rhs) { - throw std::runtime_error("Group-wise INT8 matmul requires pretransposed weights"); - } - - const int8_t* lhs_int8; - const float* lhs_scales; - - if (lhs_is_prequantized_int8) { - lhs_int8 = lhs_buffer.data_as(); - lhs_scales = lhs_buffer.activation_scales_as_float(); - } else if (lhs_buffer.precision == Precision::FP16) { - const __fp16* lhs = lhs_buffer.data_as<__fp16>(); - ensure_quant_buffers(M, K); - quantize_activations_fp16_to_int8(lhs, quant_activation_buffer.data(), - quant_scales_buffer.data(), M, K); - lhs_int8 = quant_activation_buffer.data(); - lhs_scales = quant_scales_buffer.data(); - } else { - throw std::runtime_error("INT8 matmul requires INT8 (pre-quantized) or FP16 activations"); - } - - cactus_matmul_int8(lhs_int8, lhs_scales, - rhs, rhs_scales, output, - M, K, N, rhs_buffer.group_size); - - } else { - if (lhs_buffer.precision != Precision::FP16) { - throw std::runtime_error("FP16 matmul requires FP16 activations"); - } - - const __fp16* lhs = lhs_buffer.data_as<__fp16>(); - const __fp16* rhs = rhs_buffer.data_as<__fp16>(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (pretransposed_rhs) { - cactus_matmul_f16(lhs, rhs, output, M, K, N); - } else { - size_t transpose_size = rhs_shape[0] * rhs_shape[1]; - ensure_transpose_buffer_fp16(transpose_size); - - cactus_transpose_2d_f16(rhs, transpose_buffer_fp16.data(), - rhs_shape[0], rhs_shape[1], 0, rhs_shape[0]); - cactus_matmul_f16(lhs, transpose_buffer_fp16.data(), output, M, K, N); - } - } -} - -void compute_rms_norm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& weight_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - - if (input_buffer.shape.size() != 2) { - throw std::runtime_error("RMS normalization requires 2D input tensor [batch_size, dims], got " + - std::to_string(input_buffer.shape.size()) + "D tensor"); - } - - size_t batch_size = input_buffer.shape[0]; - size_t dims = input_buffer.shape[1]; - - if (input_buffer.precision != Precision::FP16) { - throw std::runtime_error("RMS normalization only supports FP16 precision"); - } - - cactus_rms_norm_f16(input_buffer.data_as<__fp16>(), weight_buffer.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), batch_size, dims, node.params.epsilon); -} - -void compute_rope_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU RoPE operation not yet implemented"); - } - - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& shape = input_buffer.shape; - - if (shape.size() < 4) { - throw std::runtime_error("RoPE operation requires 4D tensor with shape [batch, seq_len, num_heads, head_dim], got " + - std::to_string(shape.size()) + "D tensor"); - } - - if (input_buffer.precision != Precision::FP16 || node.output_buffer.precision != Precision::FP16) { - throw std::runtime_error("RoPE operation only supports FP16 precision"); - } - - size_t batch_size = shape[0]; - size_t seq_len = shape[1]; - size_t num_heads = shape[2]; - size_t head_dim = shape[3]; - - cactus_rope_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - batch_size, seq_len, num_heads, head_dim, node.params.position_offset, node.params.theta); -} - -void compute_softmax_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& shape = input_buffer.shape; - - if (shape.size() < 2) { - throw std::runtime_error("Softmax operation requires at least 2D tensor, got " + - std::to_string(shape.size()) + "D tensor"); - } - - if (input_buffer.precision != Precision::FP16) { - throw std::runtime_error("Softmax operation only supports FP16 precision"); - } - - size_t batch_size = 1; - for (size_t i = 0; i < shape.size() - 1; i++) { - batch_size *= shape[i]; - } - size_t vocab_size = shape[shape.size() - 1]; - - cactus_softmax_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - batch_size, 1, vocab_size); -} - -void compute_attention_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU attention operation not yet implemented"); - } - - if (node.input_ids.size() < 3) { - throw std::runtime_error("Attention operation requires 3 inputs (query, key, value), got " + - std::to_string(node.input_ids.size()) + " inputs"); - } - - const auto& query_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& key_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - const auto& value_buffer = nodes[node_index_map.at(node.input_ids[2])]->output_buffer; - const auto& q_shape = query_buffer.shape; - const auto& k_shape = key_buffer.shape; - - if (q_shape.size() < 4) { - throw std::runtime_error("Attention operation requires 4D tensors [batch, seq_len, num_heads, head_dim], got " + - std::to_string(q_shape.size()) + "D tensor"); - } - - if (query_buffer.precision != Precision::FP16) { - throw std::runtime_error("Attention operation only supports FP16 precision"); - } - - size_t batch_size = q_shape[0]; - size_t seq_len = q_shape[1]; - size_t num_q_heads = q_shape[2]; - size_t head_dim = q_shape[3]; - size_t num_kv_heads = k_shape[2]; - size_t kv_seq_len = key_buffer.shape[1]; - - cactus_attention_f16(query_buffer.data_as<__fp16>(), key_buffer.data_as<__fp16>(), - value_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(), - batch_size, seq_len, kv_seq_len, num_q_heads, num_kv_heads, head_dim, node.params.scale, nullptr, - node.params.position_offset, node.params.window_size, node.params.is_causal); -} - -void compute_attention_int8_hybrid_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& query_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& key_new_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - const auto& value_new_buffer = nodes[node_index_map.at(node.input_ids[2])]->output_buffer; - const auto& q_shape = query_buffer.shape; - - if (q_shape.size() < 4) { - throw std::runtime_error("ATTENTION_INT8_HYBRID requires 4D query tensor"); - } - - size_t batch_size = q_shape[0]; - size_t seq_len = q_shape[1]; - size_t num_q_heads = q_shape[2]; - size_t head_dim = node.params.head_dim; - size_t num_kv_heads = node.params.num_kv_heads; - size_t cache_len = node.params.cache_seq_len; - size_t new_len = key_new_buffer.shape[1]; - - cactus_attention_hybrid_int8_fp16( - query_buffer.data_as<__fp16>(), - node.params.cached_keys_int8, - node.params.cached_values_int8, - node.params.cached_k_scales, - node.params.cached_v_scales, - key_new_buffer.data_as<__fp16>(), - value_new_buffer.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - batch_size, seq_len, cache_len, new_len, - num_q_heads, num_kv_heads, head_dim, - node.params.scale, node.params.position_offset, true - ); -} - -void compute_layernorm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& weight_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - const auto& bias_buffer = nodes[node_index_map.at(node.input_ids[2])]->output_buffer; - float epsilon = node.params.epsilon; - - if (input_buffer.shape.empty()) { - throw std::runtime_error("LayerNorm requires non-empty input tensor"); - } - - size_t feature_size = input_buffer.shape.back(); - size_t batch_size = input_buffer.total_size / feature_size; - - std::vector input_float(input_buffer.total_size); - std::vector weight_float(feature_size); - std::vector bias_float(feature_size); - - if (input_buffer.precision == Precision::INT8) { - throw std::runtime_error("LayerNorm currently does not support INT8 input"); - } else if (input_buffer.precision == Precision::FP16) { - const __fp16* input_fp16 = input_buffer.data_as<__fp16>(); - for (size_t i = 0; i < input_buffer.total_size; ++i) { - input_float[i] = static_cast(input_fp16[i]); - } - } else { - std::memcpy(input_float.data(), input_buffer.data_as(), input_buffer.total_size * sizeof(float)); - } - - if (weight_buffer.precision == Precision::INT8) { - throw std::runtime_error("LayerNorm currently does not support INT8 weight"); - } else if (weight_buffer.precision == Precision::FP16) { - const __fp16* weight_fp16 = weight_buffer.data_as<__fp16>(); - for (size_t i = 0; i < feature_size; ++i) { - weight_float[i] = static_cast(weight_fp16[i]); - } - } else { - std::memcpy(weight_float.data(), weight_buffer.data_as(), feature_size * sizeof(float)); - } - - if (bias_buffer.precision == Precision::INT8) { - throw std::runtime_error("LayerNorm currently does not support INT8 bias"); - } else if (bias_buffer.precision == Precision::FP16) { - const __fp16* bias_fp16 = bias_buffer.data_as<__fp16>(); - for (size_t i = 0; i < feature_size; ++i) { - bias_float[i] = static_cast(bias_fp16[i]); - } - } else { - std::memcpy(bias_float.data(), bias_buffer.data_as(), feature_size * sizeof(float)); - } - - std::vector output_float(input_buffer.total_size); - for (size_t b = 0; b < batch_size; ++b) { - const float* input_row = input_float.data() + b * feature_size; - float* output_row = output_float.data() + b * feature_size; - - float mean = 0.0f; - for (size_t i = 0; i < feature_size; ++i) { - mean += input_row[i]; - } - mean /= feature_size; - - float variance = 0.0f; - for (size_t i = 0; i < feature_size; ++i) { - float diff = input_row[i] - mean; - variance += diff * diff; - } - variance /= feature_size; - - float std_inv = 1.0f / std::sqrt(variance + epsilon); - for (size_t i = 0; i < feature_size; ++i) { - output_row[i] = (input_row[i] - mean) * std_inv * weight_float[i] + bias_float[i]; - } - } - - if (node.output_buffer.precision == Precision::INT8) { - throw std::runtime_error("LayerNorm currently does not support INT8 output"); - } else if (node.output_buffer.precision == Precision::FP16) { - __fp16* output_fp16 = node.output_buffer.data_as<__fp16>(); - for (size_t i = 0; i < input_buffer.total_size; ++i) { - output_fp16[i] = static_cast<__fp16>(output_float[i]); - } - } else { - std::memcpy(node.output_buffer.data_as(), output_float.data(), input_buffer.total_size * sizeof(float)); - } -} - -void compute_conv1d_causal_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU causal convolution operation not yet implemented"); - } - - const auto& X = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& W = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - auto& Y = node.output_buffer; - - if (X.shape.size() != 3) { - throw std::runtime_error("Causal conv requires 3D input [batch, seq_len, in_channels]"); - } - if (W.shape.size() != 3) { - throw std::runtime_error("Weight must be 3D"); - } - - const size_t N = X.shape[0]; - const size_t L = X.shape[1]; - const size_t C_in = X.shape[2]; - const size_t W0 = W.shape[0]; - const size_t W1 = W.shape[1]; - const size_t K = W.shape[2]; - const size_t dil = node.params.dilation; - if (dil < 1) throw std::runtime_error("dilation must be >= 1"); - - size_t M = 1; - size_t C_out = 0; - - assert((W1 == 1) && (W0 % C_in == 0) && "Only depthwise causal convolution is supported currently"); - M = W0 / C_in; - C_out = C_in * M; - - Y.shape = { N, L, C_out }; - Y.precision = X.precision; - - if (W.precision == Precision::INT8) { - const size_t W_size = W0 * W1 * K; - const int8_t* W_int8 = W.data_as(); - - std::vector<__fp16> W_fp16(W_size); - - if (W.is_grouped_int8()) { - const __fp16* scales = W.scales_as_fp16(); - const size_t K_total = W1 * K; - const size_t group_size = W.group_size; - const size_t num_groups = K_total / group_size; - - for (size_t row = 0; row < W0; ++row) { - for (size_t col = 0; col < K_total; ++col) { - size_t idx = row * K_total + col; - size_t group_idx = col / group_size; - float scale = static_cast(scales[row * num_groups + group_idx]); - W_fp16[idx] = static_cast<__fp16>(W_int8[idx] * scale); - } - } - } else { - for (size_t i = 0; i < W_size; ++i) { - W_fp16[i] = static_cast<__fp16>(W_int8[i]); - } - } - - cactus_conv1d_causal_depthwise_f16( - X.data_as<__fp16>(), W_fp16.data(), Y.data_as<__fp16>(), - N, L, C_in, K, dil); - } else if (W.precision == Precision::FP16) { - cactus_conv1d_causal_depthwise_f16( - X.data_as<__fp16>(), W.data_as<__fp16>(), Y.data_as<__fp16>(), - N, L, C_in, K, dil); - } else { - throw std::runtime_error("Depthwise causal conv supports INT8/FP16 weights"); - } -} - -void compute_conv1d_k3_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU causal convolution operation not yet implemented"); - } - - const auto& X = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& W = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - auto& Y = node.output_buffer; - - if (X.shape.size() != 3) - throw std::runtime_error("Conv requires 3D input [N, C_in, L]!"); - - if (W.shape.size() != 3) - throw std::runtime_error("Weight must be [C_out, C_in, 3]!"); - - const size_t N = X.shape[0]; - const size_t C_in = X.shape[1]; - const size_t L = X.shape[2]; - - const size_t C_out = W.shape[0]; - const size_t K = W.shape[2]; - const size_t stride = node.params.stride; - - if (K != 3) - throw std::runtime_error("Conv1d_k3 only supports K=3!"); - - size_t L_out = ((L - 1) / stride) + 1; - Y.shape = { N, C_out, L_out }; - Y.precision = X.precision; - - if (X.precision != Precision::FP16) { - throw std::runtime_error("Conv1d_k3 only supports FP16 activations"); - } - - if (W.precision == Precision::INT8) { - const size_t W_size = C_out * C_in * K; - const int8_t* W_int8 = W.data_as(); - - std::vector<__fp16> W_fp16(W_size); - - if (W.is_grouped_int8()) { - const __fp16* scales = W.scales_as_fp16(); - const size_t K_total = C_in * K; - const size_t group_size = W.group_size; - const size_t num_groups = K_total / group_size; - - for (size_t row = 0; row < C_out; ++row) { - for (size_t col = 0; col < K_total; ++col) { - size_t idx = row * K_total + col; - size_t group_idx = col / group_size; - float scale = static_cast(scales[row * num_groups + group_idx]); - W_fp16[idx] = static_cast<__fp16>(W_int8[idx] * scale); - } - } - } else { - for (size_t i = 0; i < W_size; ++i) { - W_fp16[i] = static_cast<__fp16>(W_int8[i]); - } - } - - cactus_conv1d_f16_k3( - X.data_as<__fp16>(), - W_fp16.data(), - Y.data_as<__fp16>(), - N, L, C_in, C_out, stride - ); - } else if (W.precision == Precision::FP16) { - cactus_conv1d_f16_k3( - X.data_as<__fp16>(), - W.data_as<__fp16>(), - Y.data_as<__fp16>(), - N, L, C_in, C_out, stride - ); - } else { - throw std::runtime_error("Conv1d_k3 only supports FP16 and INT8 weights"); - } -} diff --git a/cactus/graph/graph_ops_tensor.cpp b/cactus/graph/graph_ops_tensor.cpp deleted file mode 100644 index e119da09a..000000000 --- a/cactus/graph/graph_ops_tensor.cpp +++ /dev/null @@ -1,420 +0,0 @@ -#include "graph.h" -#include "../kernel/kernel.h" -#include -#include - -void compute_transpose_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU transpose operation not yet implemented"); - } - - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - - if (input_buffer.precision != Precision::FP16) { - throw std::runtime_error("Transpose only supports FP16 precision"); - } - - const auto& permutation = node.params.permutation; - - const __fp16* input = input_buffer.data_as<__fp16>(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - cactus_transpose_f16(input, output, input_buffer.shape.data(), permutation.data(), permutation.size(), 0, input_buffer.total_size); -} - -void compute_gather_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& tensor_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& indices_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - - size_t first_dim = tensor_buffer.shape[0]; - size_t element_size = 1; - for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { - element_size *= tensor_buffer.shape[i]; - } - - size_t num_indices = indices_buffer.total_size; - size_t bytes_per_element = element_size * PrecisionTraits::size_of(tensor_buffer.precision); - - if (tensor_buffer.precision == Precision::INT8) { - const int8_t* tensor_data = tensor_buffer.data_as(); - int8_t* output = node.output_buffer.data_as(); - - const bool is_grouped = tensor_buffer.is_grouped_int8(); - __fp16* gathered_scales = nullptr; - const __fp16* src_scales = nullptr; - size_t num_groups = 0; - - if (is_grouped) { - num_groups = tensor_buffer.num_groups; - src_scales = tensor_buffer.scales_as_fp16(); - size_t scales_bytes = num_indices * num_groups * sizeof(__fp16); - node.output_buffer.owned_scales = std::make_unique(scales_bytes); - gathered_scales = reinterpret_cast<__fp16*>(node.output_buffer.owned_scales.get()); - } - - if (indices_buffer.precision == Precision::INT8) { - const int8_t* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= first_dim) { - throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); - } - std::memcpy(output + i * element_size, tensor_data + idx * element_size, bytes_per_element); - if (is_grouped) { - for (size_t g = 0; g < num_groups; g++) { - gathered_scales[i * num_groups + g] = src_scales[idx * num_groups + g]; - } - } - } - } else { - const float* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= first_dim) { - throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); - } - std::memcpy(output + i * element_size, tensor_data + idx * element_size, bytes_per_element); - if (is_grouped) { - for (size_t g = 0; g < num_groups; g++) { - gathered_scales[i * num_groups + g] = src_scales[idx * num_groups + g]; - } - } - } - } - - if (is_grouped) { - node.output_buffer.group_size = tensor_buffer.group_size; - node.output_buffer.num_groups = num_groups; - node.output_buffer.scales_data = gathered_scales; - } - } else if (tensor_buffer.precision == Precision::FP16) { - const __fp16* tensor_data = tensor_buffer.data_as<__fp16>(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (indices_buffer.precision == Precision::INT8) { - const int8_t* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= first_dim) { - throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); - } - std::memcpy(output + i * element_size, tensor_data + idx * element_size, bytes_per_element); - } - } else { - const float* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= first_dim) { - throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); - } - std::memcpy(output + i * element_size, tensor_data + idx * element_size, bytes_per_element); - } - } - } else { - const float* tensor_data = tensor_buffer.data_as(); - float* output = node.output_buffer.data_as(); - - if (indices_buffer.precision == Precision::INT8) { - const int8_t* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= first_dim) { - throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); - } - std::memcpy(output + i * element_size, tensor_data + idx * element_size, bytes_per_element); - } - } else { - const float* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= first_dim) { - throw std::runtime_error("Gather index " + std::to_string(idx) + " out of bounds for dimension " + std::to_string(first_dim)); - } - std::memcpy(output + i * element_size, tensor_data + idx * element_size, bytes_per_element); - } - } - } -} - -void compute_slice_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - auto* input_node = nodes[node_index_map.at(node.input_ids[0])].get(); - auto& input_buffer = input_node->output_buffer; - - const size_t axis_index = static_cast(node.params.axis); - - const size_t axis_size = input_buffer.shape[axis_index]; - const size_t slice_start = node.params.slice_start; - size_t slice_length = node.params.slice_length; - - if (slice_length == 0) { - slice_length = axis_size - slice_start; - } - - const size_t element_size = PrecisionTraits::size_of(input_buffer.precision); - - if (axis_index == 0) { - size_t inner_elements = 1; - for (size_t i = 1; i < input_buffer.shape.size(); ++i) { - inner_elements *= input_buffer.shape[i]; - } - - auto* base_ptr = static_cast(input_buffer.get_data()); - if (!base_ptr) { - throw std::runtime_error("Slice input buffer is not available"); - } - - const size_t byte_offset = slice_start * inner_elements * element_size; - - node.output_buffer.set_external(base_ptr + byte_offset); - node.output_buffer.precision = input_buffer.precision; - - if (input_buffer.is_grouped_int8()) { - size_t num_groups = input_buffer.num_groups; - size_t scales_bytes = slice_length * num_groups * sizeof(__fp16); - node.output_buffer.owned_scales = std::make_unique(scales_bytes); - __fp16* sliced_scales = reinterpret_cast<__fp16*>(node.output_buffer.owned_scales.get()); - const __fp16* input_scales = input_buffer.scales_as_fp16(); - - for (size_t i = 0; i < slice_length; i++) { - for (size_t g = 0; g < num_groups; g++) { - sliced_scales[i * num_groups + g] = input_scales[(slice_start + i) * num_groups + g]; - } - } - - node.output_buffer.group_size = input_buffer.group_size; - node.output_buffer.num_groups = num_groups; - node.output_buffer.scales_data = sliced_scales; - } - return; - } - - const char* input_ptr = static_cast(input_buffer.get_data()); - if (!input_ptr) { - throw std::runtime_error("Slice input buffer is not available"); - } - - size_t inner_elements = 1; - for (size_t i = axis_index + 1; i < input_buffer.shape.size(); ++i) { - inner_elements *= input_buffer.shape[i]; - } - - size_t outer_elements = 1; - for (size_t i = 0; i < axis_index; ++i) { - outer_elements *= input_buffer.shape[i]; - } - - const size_t copy_block_elements = slice_length * inner_elements; - const size_t axis_stride_elements = axis_size * inner_elements; - const size_t copy_block_bytes = copy_block_elements * element_size; - const size_t axis_stride_bytes = axis_stride_elements * element_size; - - node.output_buffer.external_data = nullptr; - node.output_buffer.allocate(); - node.output_buffer.precision = input_buffer.precision; - - auto* output_ptr = static_cast(node.output_buffer.get_data()); - if (!output_ptr) { - throw std::runtime_error("Slice output buffer could not be allocated"); - } - - for (size_t outer = 0; outer < outer_elements; ++outer) { - const char* src = input_ptr + outer * axis_stride_bytes + slice_start * inner_elements * element_size; - char* dst = output_ptr + outer * copy_block_bytes; - std::memcpy(dst, src, copy_block_bytes); - } -} - -void compute_embedding_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& embeddings_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& indices_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - - size_t vocab_size = embeddings_buffer.shape[0]; - size_t hidden_dim = embeddings_buffer.shape[1]; - size_t num_indices = indices_buffer.total_size; - - if (embeddings_buffer.precision == Precision::INT8) { - const int8_t* embeddings = embeddings_buffer.data_as(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (embeddings_buffer.is_grouped_int8()) { - const __fp16* scales = embeddings_buffer.scales_as_fp16(); - size_t group_size = embeddings_buffer.group_size; - size_t num_groups = embeddings_buffer.num_groups; - - auto dequant_row = [&](size_t i, size_t idx) { - const int8_t* emb_row = embeddings + idx * hidden_dim; - __fp16* out_row = output + i * hidden_dim; - - for (size_t g = 0; g < num_groups; g++) { - float scale = (float)scales[idx * num_groups + g]; - size_t k_start = g * group_size; - size_t k_end = std::min(k_start + group_size, hidden_dim); - for (size_t k = k_start; k < k_end; k++) { - out_row[k] = static_cast<__fp16>(emb_row[k] * scale); - } - } - }; - - if (indices_buffer.precision == Precision::FP32) { - const float* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= vocab_size) { - throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); - } - dequant_row(i, idx); - } - } else { - const int8_t* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= vocab_size) { - throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); - } - dequant_row(i, idx); - } - } - } else { - if (indices_buffer.precision == Precision::FP32) { - const float* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= vocab_size) { - throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); - } - for (size_t j = 0; j < hidden_dim; j++) { - output[i * hidden_dim + j] = static_cast<__fp16>(embeddings[idx * hidden_dim + j]); - } - } - } else { - const int8_t* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= vocab_size) { - throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); - } - for (size_t j = 0; j < hidden_dim; j++) { - output[i * hidden_dim + j] = static_cast<__fp16>(embeddings[idx * hidden_dim + j]); - } - } - } - } - } else if (embeddings_buffer.precision == Precision::FP16) { - const __fp16* embeddings = embeddings_buffer.data_as<__fp16>(); - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (indices_buffer.precision == Precision::FP32) { - const float* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= vocab_size) { - throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); - } - for (size_t j = 0; j < hidden_dim; j++) { - output[i * hidden_dim + j] = embeddings[idx * hidden_dim + j]; - } - } - } else { - const int8_t* indices = indices_buffer.data_as(); - for (size_t i = 0; i < num_indices; i++) { - size_t idx = static_cast(indices[i]); - if (idx >= vocab_size) { - throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); - } - for (size_t j = 0; j < hidden_dim; j++) { - output[i * hidden_dim + j] = embeddings[idx * hidden_dim + j]; - } - } - } - } else { - throw std::runtime_error("Embedding only supports INT8 and FP16 precision"); - } -} - -void compute_concat_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input1_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& input2_buffer = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; - - std::vector shape1 = input1_buffer.shape; - std::vector shape2 = input2_buffer.shape; - std::vector output_shape = node.output_buffer.shape; - - if (input1_buffer.precision != Precision::FP16) { - throw std::runtime_error("Concat operation only supports FP16 precision"); - } - cactus_concat_f16(input1_buffer.data_as<__fp16>(), input2_buffer.data_as<__fp16>(), - node.output_buffer.data_as<__fp16>(), - shape1.data(), shape2.data(), output_shape.data(), - shape1.size(), node.params.axis); -} - -void compute_index_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& input_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - const auto& input_shape = input_buffer.shape; - - int dim = node.params.axis; - size_t index_value = node.params.index_value; - - size_t element_size = PrecisionTraits::size_of(input_buffer.precision); - const char* input_data = static_cast(input_buffer.get_data()); - char* output_data = static_cast(node.output_buffer.get_data()); - - if (dim == 0) { - size_t slice_size = input_buffer.total_size / input_shape[0]; - size_t offset_bytes = index_value * slice_size * element_size; - node.output_buffer.set_external(const_cast(input_data) + offset_bytes); - return; - } - - std::vector input_strides(input_shape.size()); - input_strides[input_shape.size() - 1] = 1; - for (int i = static_cast(input_shape.size()) - 2; i >= 0; --i) { - input_strides[i] = input_strides[i + 1] * input_shape[i + 1]; - } - - size_t slice_size = input_strides[dim]; - size_t outer_size = input_buffer.total_size / input_strides[dim - 1]; - size_t dim_stride = input_strides[dim]; - size_t block_size = dim_stride * input_shape[dim]; - - size_t output_idx = 0; - for (size_t outer_idx = 0; outer_idx < outer_size; ++outer_idx) { - size_t input_base = outer_idx * block_size + index_value * dim_stride; - - std::memcpy(output_data + output_idx * element_size, - input_data + input_base * element_size, - slice_size * element_size); - - output_idx += slice_size; - } -} - -void compute_bilinear_interpolation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - const auto& pos_embeds_buffer = nodes[node_index_map.at(node.input_ids[0])]->output_buffer; - - size_t total_pos_embeds = pos_embeds_buffer.shape[0]; - size_t embed_dim = pos_embeds_buffer.shape[1]; - - size_t src_height = static_cast(std::sqrt(total_pos_embeds)); - size_t src_width = src_height; - - size_t dst_height = node.params.dst_height; - size_t dst_width = node.params.dst_width; - - __fp16* output = node.output_buffer.data_as<__fp16>(); - - if (pos_embeds_buffer.precision == Precision::FP16) { - const __fp16* input = pos_embeds_buffer.data_as<__fp16>(); - cactus_bilinear_interpolation_f16(input, output, src_height, src_width, embed_dim, - dst_height, dst_width); - } - else if (pos_embeds_buffer.precision == Precision::INT8) { - std::vector<__fp16> input_fp16(total_pos_embeds * embed_dim); - cactus_int8_to_fp16(pos_embeds_buffer.data_as(), input_fp16.data(), - total_pos_embeds * embed_dim); - cactus_bilinear_interpolation_f16(input_fp16.data(), output, src_height, src_width, embed_dim, - dst_height, dst_width); - } - else { - throw std::runtime_error("BILINEAR_INTERPOLATION only supports INT8 and FP16 input precision"); - } -} diff --git a/cactus/kernel/kernel.h b/cactus/kernel/kernel.h deleted file mode 100644 index dea1344ba..000000000 --- a/cactus/kernel/kernel.h +++ /dev/null @@ -1,165 +0,0 @@ -#ifndef KERNEL_H -#define KERNEL_H - -#include -#include - -enum class ScalarOpType { - ADD, - SUBTRACT, - MULTIPLY, - DIVIDE, - EXP, - SQRT, - COS, - SIN -}; - -constexpr size_t KV_QUANT_GROUP_SIZE = 128; - -void cactus_add_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); -void cactus_add_f16_clipped(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); -void cactus_subtract_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); -void cactus_multiply_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); -void cactus_divide_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); - -void cactus_add_broadcast_f16(const __fp16* a, const __fp16* b, __fp16* output, - const size_t* a_strides, const size_t* b_strides, - const size_t* output_shape, size_t ndim); -void cactus_subtract_broadcast_f16(const __fp16* a, const __fp16* b, __fp16* output, - const size_t* a_strides, const size_t* b_strides, - const size_t* output_shape, size_t ndim); -void cactus_multiply_broadcast_f16(const __fp16* a, const __fp16* b, __fp16* output, - const size_t* a_strides, const size_t* b_strides, - const size_t* output_shape, size_t ndim); -void cactus_divide_broadcast_f16(const __fp16* a, const __fp16* b, __fp16* output, - const size_t* a_strides, const size_t* b_strides, - const size_t* output_shape, size_t ndim); - -void cactus_scalar_op_f16(const __fp16* input, __fp16* output, size_t num_elements, float scalar_value, ScalarOpType op_type); - -void cactus_matmul_int8(const int8_t* A, const float* A_scales, - const int8_t* B, const __fp16* B_scales, - __fp16* C, size_t M, size_t K, size_t N, size_t group_size); - -void cactus_matmul_int4(const int8_t* A, const float* A_scales, - const uint8_t* B_packed, const __fp16* B_scales, - __fp16* C, size_t M, size_t K, size_t N, size_t group_size); - -void cactus_matmul_f16(const __fp16* a, const __fp16* b_transposed, __fp16* c, - size_t M, size_t K, size_t N); - -void cactus_transpose_2d_f16(const __fp16* source, __fp16* destination, - size_t num_rows, size_t num_cols, size_t start_row, size_t end_row); -void cactus_transpose_f16(const __fp16* source, __fp16* destination, const size_t* shape, - const size_t* permutation, size_t ndim, size_t start_idx, size_t end_idx); - -double cactus_sum_all_f16(const __fp16* data, size_t num_elements); -void cactus_sum_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size); - -double cactus_mean_all_f16(const __fp16* data, size_t num_elements); -void cactus_mean_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size); - -double cactus_variance_all_f16(const __fp16* data, size_t num_elements); -void cactus_variance_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size); - -__fp16 cactus_min_all_f16(const __fp16* data, size_t num_elements); -void cactus_min_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size); - -__fp16 cactus_max_all_f16(const __fp16* data, size_t num_elements); -void cactus_max_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size); - -void cactus_rms_norm_f16(const __fp16* input, const __fp16* weight, __fp16* output, - size_t batch_size, size_t dims, float eps); - -void cactus_rope_f16(const __fp16* input, __fp16* output, size_t batch_size, size_t seq_len, - size_t num_heads, size_t head_dim, size_t start_pos, float theta); - -void cactus_softmax_f16(const __fp16* input, __fp16* output, size_t batch_size, - size_t seq_len, size_t vocab_size); - -void cactus_silu_f16(const __fp16* input, __fp16* output, size_t num_elements); - -void cactus_gelu_f16(const __fp16* input, __fp16* output, size_t num_elements); - -void cactus_gelu_f16_erf(const __fp16* input, __fp16* output, size_t num_elements); - -void cactus_attention_f16(const __fp16* queries, const __fp16* keys, const __fp16* values, __fp16* output, - size_t batch_size, size_t seq_len, size_t kv_seq_len, size_t num_q_heads, size_t num_kv_heads, - size_t head_dim, float scale, const __fp16* mask, size_t position_offset = 0, size_t window_size = 0, - bool is_causal = true); - -void cactus_attention_hybrid_int8_fp16( - const __fp16* queries, - const int8_t* keys_cached, - const int8_t* values_cached, - const float* k_scales, - const float* v_scales, - const __fp16* keys_new, - const __fp16* values_new, - __fp16* output, - size_t batch_size, size_t seq_len, size_t cache_len, size_t new_len, - size_t num_q_heads, size_t num_kv_heads, size_t head_dim, - float scale, size_t position_offset = 0, bool is_causal = true, - size_t group_size = KV_QUANT_GROUP_SIZE); - -void cactus_conv1d_causal_depthwise_f16( - const __fp16* input, - const __fp16* weight, - __fp16* output, - size_t N, - size_t L, - size_t C, - size_t K, - size_t dilation); - -void cactus_conv1d_f16_k3( - const __fp16* input, - const __fp16* weight, - __fp16* output, - size_t N, - size_t L, - size_t C_in, - size_t C_out, - size_t stride -); - -void cactus_bilinear_interpolation_f16(const __fp16* input, __fp16* output, size_t src_height, size_t src_width, size_t embed_dim, - size_t dst_height, size_t dst_width); - -void cactus_sample_f32(const float* logits, uint32_t* output, size_t vocab_size, - float temperature, float top_p, size_t top_k, size_t random_seed, - const float* bias_values = nullptr, const uint32_t* bias_indices = nullptr, - size_t bias_count = 0); -void cactus_sample_f16(const __fp16* logits, uint32_t* output, size_t vocab_size, - float temperature, float top_p, size_t top_k, size_t random_seed, - const float* bias_values = nullptr, const uint32_t* bias_indices = nullptr, - size_t bias_count = 0); - -void cactus_concat_f16(const __fp16* input1, const __fp16* input2, __fp16* output, - const size_t* shape1, const size_t* shape2, const size_t* output_shape, - size_t ndims, int axis); - -void cactus_int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale = 1.0f); -void cactus_fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale = 1.0f); -void cactus_fp16_to_fp32(const __fp16* src, float* dst, size_t count); -void cactus_fp32_to_fp16(const float* src, __fp16* dst, size_t count); -void cactus_int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale = 1.0f); -void cactus_fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale = 1.0f); -float cactus_fp16_max_abs(const __fp16* src, size_t count); - -void cactus_quantize_kv_fp16_to_int8( - const __fp16* src, - int8_t* dst, - float* scales, - size_t seq_len, size_t kv_heads, size_t head_dim, - size_t group_size = KV_QUANT_GROUP_SIZE); - -inline size_t kv_scales_count(size_t seq_len, size_t kv_heads, size_t head_dim, size_t group_size = KV_QUANT_GROUP_SIZE) { - size_t num_groups = (head_dim + group_size - 1) / group_size; - return seq_len * kv_heads * num_groups; -} - -void cactus_unpack_int4_to_int8(const uint8_t* packed, int8_t* unpacked, size_t unpacked_count); - -#endif diff --git a/cactus/kernel/kernel_attention.cpp b/cactus/kernel/kernel_attention.cpp deleted file mode 100644 index 076fc93bd..000000000 --- a/cactus/kernel/kernel_attention.cpp +++ /dev/null @@ -1,676 +0,0 @@ -#include "kernel.h" -#include "kernel_utils.h" -#include -#include -#include -#include -#include -#include - -void cactus_attention_f16( - const __fp16* queries, - const __fp16* keys, - const __fp16* values, - __fp16* output, - size_t batch_size, - size_t seq_len, - size_t kv_seq_len, - size_t num_q_heads, - size_t num_kv_heads, - size_t head_dim, - float scale, - const __fp16* mask, - size_t position_offset, - size_t window_size, - bool is_causal -) { - if (scale == 0.0f) { - scale = 1.0f / sqrtf(static_cast(head_dim)); - } - - constexpr size_t VECTOR_WIDTH = 8; - constexpr size_t BLOCK_SIZE = 32; - const size_t head_dim_aligned = (head_dim / VECTOR_WIDTH) * VECTOR_WIDTH; - - const size_t group_size = num_q_heads / num_kv_heads; - - const size_t q_batch_stride = seq_len * num_q_heads * head_dim; - const size_t kv_batch_stride = kv_seq_len * num_kv_heads * head_dim; - const size_t o_batch_stride = seq_len * num_q_heads * head_dim; - const size_t q_seq_stride = num_q_heads * head_dim; - const size_t kv_seq_stride = num_kv_heads * head_dim; - const size_t o_seq_stride = num_q_heads * head_dim; - const size_t mask_batch_stride = mask ? seq_len * kv_seq_len : 0; - - CactusThreading::parallel_for(batch_size * num_q_heads * seq_len, CactusThreading::Thresholds::ATTENTION, - [=](size_t start_idx, size_t end_idx) { - std::vector block_scores(BLOCK_SIZE); - std::vector output_accum_low(head_dim_aligned / VECTOR_WIDTH * 2); - std::vector output_accum_high(head_dim_aligned / VECTOR_WIDTH * 2); - - for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { - const size_t batch_idx = work_idx / (num_q_heads * seq_len); - const size_t remainder = work_idx % (num_q_heads * seq_len); - const size_t q_head_idx = remainder / seq_len; - const size_t q_pos = remainder % seq_len; - - const size_t kv_head_idx = q_head_idx / group_size; - - const __fp16* Q_base = queries + batch_idx * q_batch_stride; - const __fp16* K_base = keys + batch_idx * kv_batch_stride; - const __fp16* V_base = values + batch_idx * kv_batch_stride; - __fp16* O_base = output + batch_idx * o_batch_stride; - const __fp16* M = mask ? (mask + batch_idx * mask_batch_stride) : nullptr; - const __fp16* q_vec = Q_base + q_pos * q_seq_stride + q_head_idx * head_dim; - __fp16* o_vec = O_base + q_pos * o_seq_stride + q_head_idx * head_dim; - - float running_max = -std::numeric_limits::infinity(); - float running_sum = 0.0f; - - for (size_t i = 0; i < output_accum_low.size(); ++i) { - output_accum_low[i] = vdupq_n_f32(0.0f); - output_accum_high[i] = vdupq_n_f32(0.0f); - } - - const bool is_decode = (q_pos == seq_len - 1) && seq_len > 1; - const size_t absolute_q_pos = position_offset + q_pos; - - size_t kv_start = 0; - size_t kv_end = kv_seq_len; - - if (window_size > 0 && window_size < kv_seq_len) { - if (absolute_q_pos > window_size) { - kv_start = absolute_q_pos - window_size; - } - if (is_causal) { - kv_end = std::min(kv_end, absolute_q_pos + 1); - } - } else if (is_causal) { - kv_end = std::min(kv_end, absolute_q_pos + 1); - } - - for (size_t kv_block_start = kv_start; kv_block_start < kv_end; kv_block_start += BLOCK_SIZE) { - const size_t kv_block_end = std::min(kv_block_start + BLOCK_SIZE, kv_seq_len); - const size_t block_size = kv_block_end - kv_block_start; - - float block_max = -std::numeric_limits::infinity(); - - if (!is_decode && is_causal && kv_block_start > absolute_q_pos) { - for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { - block_scores[kv_idx] = -std::numeric_limits::infinity(); - } - continue; - } - - for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { - const size_t kv_pos = kv_block_start + kv_idx; - - if (!is_decode && is_causal && kv_pos > absolute_q_pos) { - block_scores[kv_idx] = -std::numeric_limits::infinity(); - continue; - } - - const __fp16* k_vec = K_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; - - if (kv_idx + 1 < block_size) { - const __fp16* next_k_vec = K_base + (kv_pos + 1) * kv_seq_stride + kv_head_idx * head_dim; - __builtin_prefetch(next_k_vec, 0, 1); - } - - float32x4_t score_accum_low = vdupq_n_f32(0.0f); - float32x4_t score_accum_high = vdupq_n_f32(0.0f); - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - float16x8_t q_vec_f16 = vld1q_f16(&q_vec[dim_block]); - float16x8_t k_vec_f16 = vld1q_f16(&k_vec[dim_block]); - - float32x4_t q_low = vcvt_f32_f16(vget_low_f16(q_vec_f16)); - float32x4_t q_high = vcvt_f32_f16(vget_high_f16(q_vec_f16)); - float32x4_t k_low = vcvt_f32_f16(vget_low_f16(k_vec_f16)); - float32x4_t k_high = vcvt_f32_f16(vget_high_f16(k_vec_f16)); - - score_accum_low = vfmaq_f32(score_accum_low, q_low, k_low); - score_accum_high = vfmaq_f32(score_accum_high, q_high, k_high); - } - - float score = vaddvq_f32(vaddq_f32(score_accum_low, score_accum_high)); - - for (size_t dim = head_dim_aligned; dim < head_dim; ++dim) { - score += static_cast(q_vec[dim]) * static_cast(k_vec[dim]); - } - - score *= scale; - - size_t absolute_q_pos = position_offset + q_pos; - - if (is_causal && kv_pos > absolute_q_pos) { - score = -std::numeric_limits::infinity(); - } - else if (window_size > 0 && kv_pos < absolute_q_pos && (absolute_q_pos - kv_pos) > window_size) { - score = -std::numeric_limits::infinity(); - } - else if (M && static_cast(M[q_pos * kv_seq_len + kv_pos]) == 0.0f) { - score = -std::numeric_limits::infinity(); - } - - block_scores[kv_idx] = score; - block_max = std::max(block_max, score); - } - - if (block_max > -std::numeric_limits::infinity()) { - float scale_correction = expf(running_max - block_max); - running_sum *= scale_correction; - - for (size_t i = 0; i < output_accum_low.size() / 2; ++i) { - output_accum_low[i] = vmulq_n_f32(output_accum_low[i], scale_correction); - output_accum_high[i] = vmulq_n_f32(output_accum_high[i], scale_correction); - } - running_max = block_max; - } - - float block_sum = 0.0f; - const size_t vec_size = (block_size / 4) * 4; - - for (size_t kv_idx = 0; kv_idx < vec_size; kv_idx += 4) { - float32x4_t scores = vld1q_f32(&block_scores[kv_idx]); - uint32x4_t inf_mask = vceqq_f32(scores, vdupq_n_f32(-std::numeric_limits::infinity())); - - float32x4_t x = vsubq_f32(scores, vdupq_n_f32(block_max)); - x = vmulq_n_f32(x, 1.442695f); - - int32x4_t xi = vcvtq_s32_f32(x); - float32x4_t xf = vsubq_f32(x, vcvtq_f32_s32(xi)); - - float32x4_t y = vfmaq_n_f32(vdupq_n_f32(1.0f), xf, 0.6931472f); - y = vfmaq_f32(y, vmulq_f32(xf, xf), vdupq_n_f32(0.2402265f)); - - xi = vaddq_s32(xi, vdupq_n_s32(127)); - xi = vshlq_n_s32(xi, 23); - y = vmulq_f32(y, vreinterpretq_f32_s32(xi)); - - y = vbslq_f32(inf_mask, vdupq_n_f32(0.0f), y); - - vst1q_f32(&block_scores[kv_idx], y); - block_sum += vaddvq_f32(y); - } - - for (size_t kv_idx = vec_size; kv_idx < block_size; ++kv_idx) { - if (block_scores[kv_idx] != -std::numeric_limits::infinity()) { - block_scores[kv_idx] = expf(block_scores[kv_idx] - block_max); - block_sum += block_scores[kv_idx]; - } else { - block_scores[kv_idx] = 0.0f; - } - } - - for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { - const float attn_weight = block_scores[kv_idx]; - if (attn_weight == 0.0f) continue; - - const size_t kv_pos = kv_block_start + kv_idx; - const __fp16* v_vec = V_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; - - const float32x4_t weight_vec = vdupq_n_f32(attn_weight); - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - float16x8_t v_vec_f16 = vld1q_f16(&v_vec[dim_block]); - float32x4_t v_low = vcvt_f32_f16(vget_low_f16(v_vec_f16)); - float32x4_t v_high = vcvt_f32_f16(vget_high_f16(v_vec_f16)); - - size_t idx = dim_block / VECTOR_WIDTH; - output_accum_low[idx] = vfmaq_f32(output_accum_low[idx], v_low, weight_vec); - output_accum_high[idx] = vfmaq_f32(output_accum_high[idx], v_high, weight_vec); - } - } - - running_sum += block_sum; - } - - if (running_sum > 0.0f) { - const float inv_sum = 1.0f / running_sum; - const float32x4_t inv_sum_vec = vdupq_n_f32(inv_sum); - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - size_t idx = dim_block / VECTOR_WIDTH; - float32x4_t final_low = vmulq_f32(output_accum_low[idx], inv_sum_vec); - float32x4_t final_high = vmulq_f32(output_accum_high[idx], inv_sum_vec); - - float16x4_t low_f16 = vcvt_f16_f32(final_low); - float16x4_t high_f16 = vcvt_f16_f32(final_high); - float16x8_t combined = vcombine_f16(low_f16, high_f16); - - vst1q_f16(&o_vec[dim_block], combined); - } - - for (size_t dim = head_dim_aligned; dim < head_dim; ++dim) { - o_vec[dim] = static_cast<__fp16>(0.0f); - } - } else { - for (size_t dim = 0; dim < head_dim; ++dim) { - o_vec[dim] = static_cast<__fp16>(0.0f); - } - } - } - }); -} - -void cactus_attention_hybrid_int8_fp16( - const __fp16* queries, - const int8_t* keys_cached, - const int8_t* values_cached, - const float* k_scales, - const float* v_scales, - const __fp16* keys_new, - const __fp16* values_new, - __fp16* output, - size_t batch_size, - size_t seq_len, - size_t cache_len, - size_t new_len, - size_t num_q_heads, - size_t num_kv_heads, - size_t head_dim, - float scale, - size_t position_offset, - bool is_causal, - size_t quant_group_size -) { - if (scale == 0.0f) { - scale = 1.0f / sqrtf(static_cast(head_dim)); - } - - const size_t kv_seq_len = cache_len + new_len; - - constexpr size_t VECTOR_WIDTH = 8; - constexpr size_t BLOCK_SIZE = 32; - const size_t head_dim_aligned = (head_dim / VECTOR_WIDTH) * VECTOR_WIDTH; - - const size_t gqa_group_size = num_q_heads / num_kv_heads; // GQA group size - const size_t num_quant_groups = (head_dim + quant_group_size - 1) / quant_group_size; - - const size_t q_batch_stride = seq_len * num_q_heads * head_dim; - const size_t kv_cached_batch_stride = cache_len * num_kv_heads * head_dim; - const size_t kv_new_batch_stride = new_len * num_kv_heads * head_dim; - const size_t o_batch_stride = seq_len * num_q_heads * head_dim; - const size_t q_seq_stride = num_q_heads * head_dim; - const size_t kv_seq_stride = num_kv_heads * head_dim; - const size_t o_seq_stride = num_q_heads * head_dim; - - CactusThreading::parallel_for(batch_size * num_q_heads * seq_len, CactusThreading::Thresholds::ATTENTION, - [=](size_t start_idx, size_t end_idx) { - std::vector block_scores(BLOCK_SIZE); - std::vector output_accum_low(head_dim_aligned / VECTOR_WIDTH * 2); - std::vector output_accum_high(head_dim_aligned / VECTOR_WIDTH * 2); - - for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { - const size_t batch_idx = work_idx / (num_q_heads * seq_len); - const size_t remainder = work_idx % (num_q_heads * seq_len); - const size_t q_head_idx = remainder / seq_len; - const size_t q_pos = remainder % seq_len; - - const size_t kv_head_idx = q_head_idx / gqa_group_size; - - const __fp16* Q_base = queries + batch_idx * q_batch_stride; - const int8_t* K_cached_base = keys_cached + batch_idx * kv_cached_batch_stride; - const int8_t* V_cached_base = values_cached + batch_idx * kv_cached_batch_stride; - const __fp16* K_new_base = keys_new + batch_idx * kv_new_batch_stride; - const __fp16* V_new_base = values_new + batch_idx * kv_new_batch_stride; - __fp16* O_base = output + batch_idx * o_batch_stride; - - const __fp16* q_vec = Q_base + q_pos * q_seq_stride + q_head_idx * head_dim; - __fp16* o_vec = O_base + q_pos * o_seq_stride + q_head_idx * head_dim; - - float running_max = -std::numeric_limits::infinity(); - float running_sum = 0.0f; - - for (size_t i = 0; i < output_accum_low.size(); ++i) { - output_accum_low[i] = vdupq_n_f32(0.0f); - output_accum_high[i] = vdupq_n_f32(0.0f); - } - - const size_t absolute_q_pos = position_offset + q_pos; - size_t kv_end = is_causal ? std::min(kv_seq_len, absolute_q_pos + 1) : kv_seq_len; - - for (size_t kv_block_start = 0; kv_block_start < kv_end; kv_block_start += BLOCK_SIZE) { - const size_t kv_block_end = std::min(kv_block_start + BLOCK_SIZE, kv_end); - const size_t block_size = kv_block_end - kv_block_start; - - float block_max = -std::numeric_limits::infinity(); - - for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { - const size_t kv_pos = kv_block_start + kv_idx; - - if (is_causal && kv_pos > absolute_q_pos) { - block_scores[kv_idx] = -std::numeric_limits::infinity(); - continue; - } - - float32x4_t score_accum_low = vdupq_n_f32(0.0f); - float32x4_t score_accum_high = vdupq_n_f32(0.0f); - - if (kv_pos < cache_len) { - const int8_t* k_vec = K_cached_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; - const float* k_scale_base = k_scales + (kv_pos * num_kv_heads + kv_head_idx) * num_quant_groups; - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - size_t quant_group = dim_block / quant_group_size; - float k_scale = k_scale_base[quant_group]; - float32x4_t k_scale_vec = vdupq_n_f32(k_scale); - - float16x8_t q_vec_f16 = vld1q_f16(&q_vec[dim_block]); - float32x4_t q_low = vcvt_f32_f16(vget_low_f16(q_vec_f16)); - float32x4_t q_high = vcvt_f32_f16(vget_high_f16(q_vec_f16)); - - int8x8_t k_vec_i8 = vld1_s8(&k_vec[dim_block]); - int16x8_t k_vec_i16 = vmovl_s8(k_vec_i8); - float32x4_t k_low = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(k_vec_i16))), k_scale_vec); - float32x4_t k_high = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(k_vec_i16))), k_scale_vec); - - score_accum_low = vfmaq_f32(score_accum_low, q_low, k_low); - score_accum_high = vfmaq_f32(score_accum_high, q_high, k_high); - } - } else { - const size_t new_pos = kv_pos - cache_len; - const __fp16* k_vec = K_new_base + new_pos * kv_seq_stride + kv_head_idx * head_dim; - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - float16x8_t q_vec_f16 = vld1q_f16(&q_vec[dim_block]); - float16x8_t k_vec_f16 = vld1q_f16(&k_vec[dim_block]); - - float32x4_t q_low = vcvt_f32_f16(vget_low_f16(q_vec_f16)); - float32x4_t q_high = vcvt_f32_f16(vget_high_f16(q_vec_f16)); - float32x4_t k_low = vcvt_f32_f16(vget_low_f16(k_vec_f16)); - float32x4_t k_high = vcvt_f32_f16(vget_high_f16(k_vec_f16)); - - score_accum_low = vfmaq_f32(score_accum_low, q_low, k_low); - score_accum_high = vfmaq_f32(score_accum_high, q_high, k_high); - } - } - - float score = vaddvq_f32(vaddq_f32(score_accum_low, score_accum_high)) * scale; - block_scores[kv_idx] = score; - block_max = std::max(block_max, score); - } - - if (block_max > -std::numeric_limits::infinity()) { - float scale_correction = expf(running_max - block_max); - running_sum *= scale_correction; - - for (size_t i = 0; i < output_accum_low.size() / 2; ++i) { - output_accum_low[i] = vmulq_n_f32(output_accum_low[i], scale_correction); - output_accum_high[i] = vmulq_n_f32(output_accum_high[i], scale_correction); - } - running_max = block_max; - } - - float block_sum = 0.0f; - for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { - if (block_scores[kv_idx] != -std::numeric_limits::infinity()) { - block_scores[kv_idx] = expf(block_scores[kv_idx] - block_max); - block_sum += block_scores[kv_idx]; - } else { - block_scores[kv_idx] = 0.0f; - } - } - - for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { - const float attn_weight = block_scores[kv_idx]; - if (attn_weight == 0.0f) continue; - - const size_t kv_pos = kv_block_start + kv_idx; - const float32x4_t weight_vec = vdupq_n_f32(attn_weight); - - if (kv_pos < cache_len) { - const int8_t* v_vec = V_cached_base + kv_pos * kv_seq_stride + kv_head_idx * head_dim; - const float* v_scale_base = v_scales + (kv_pos * num_kv_heads + kv_head_idx) * num_quant_groups; - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - size_t quant_group = dim_block / quant_group_size; - float v_scale = v_scale_base[quant_group]; - float32x4_t v_scale_vec = vdupq_n_f32(v_scale); - - int8x8_t v_vec_i8 = vld1_s8(&v_vec[dim_block]); - int16x8_t v_vec_i16 = vmovl_s8(v_vec_i8); - float32x4_t v_low = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(v_vec_i16))), v_scale_vec); - float32x4_t v_high = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(v_vec_i16))), v_scale_vec); - - size_t idx = dim_block / VECTOR_WIDTH; - output_accum_low[idx] = vfmaq_f32(output_accum_low[idx], v_low, weight_vec); - output_accum_high[idx] = vfmaq_f32(output_accum_high[idx], v_high, weight_vec); - } - } else { - const size_t new_pos = kv_pos - cache_len; - const __fp16* v_vec = V_new_base + new_pos * kv_seq_stride + kv_head_idx * head_dim; - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - float16x8_t v_vec_f16 = vld1q_f16(&v_vec[dim_block]); - float32x4_t v_low = vcvt_f32_f16(vget_low_f16(v_vec_f16)); - float32x4_t v_high = vcvt_f32_f16(vget_high_f16(v_vec_f16)); - - size_t idx = dim_block / VECTOR_WIDTH; - output_accum_low[idx] = vfmaq_f32(output_accum_low[idx], v_low, weight_vec); - output_accum_high[idx] = vfmaq_f32(output_accum_high[idx], v_high, weight_vec); - } - } - } - - running_sum += block_sum; - } - - if (running_sum > 0.0f) { - const float inv_sum = 1.0f / running_sum; - const float32x4_t inv_sum_vec = vdupq_n_f32(inv_sum); - - for (size_t dim_block = 0; dim_block < head_dim_aligned; dim_block += VECTOR_WIDTH) { - size_t idx = dim_block / VECTOR_WIDTH; - float32x4_t final_low = vmulq_f32(output_accum_low[idx], inv_sum_vec); - float32x4_t final_high = vmulq_f32(output_accum_high[idx], inv_sum_vec); - - float16x4_t low_f16 = vcvt_f16_f32(final_low); - float16x4_t high_f16 = vcvt_f16_f32(final_high); - float16x8_t combined = vcombine_f16(low_f16, high_f16); - - vst1q_f16(&o_vec[dim_block], combined); - } - } else { - for (size_t dim = 0; dim < head_dim; ++dim) { - o_vec[dim] = static_cast<__fp16>(0.0f); - } - } - } - }); -} - -void cactus_rms_norm_f16( - const __fp16* input, - const __fp16* weight, - __fp16* output, - size_t batch_size, - size_t dims, - float eps -) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL_FACTOR = 2; - constexpr size_t TILE_SIZE = SIMD_WIDTH * UNROLL_FACTOR; - - for (size_t b = 0; b < batch_size; ++b) { - const __fp16* input_row = input + b * dims; - __fp16* output_row = output + b * dims; - - float32x4_t sum_squares_vec[UNROLL_FACTOR * 2]; - for (size_t u = 0; u < UNROLL_FACTOR * 2; u++) { - sum_squares_vec[u] = vdupq_n_f32(0.0f); - } - - size_t i = 0; - const size_t tile_end = (dims >= TILE_SIZE) ? dims - TILE_SIZE + 1 : 0; - - for (; i < tile_end; i += TILE_SIZE) { - for (size_t u = 0; u < UNROLL_FACTOR; u++) { - float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); - float32x4_t input_low = vcvt_f32_f16(vget_low_f16(input_vec)); - float32x4_t input_high = vcvt_f32_f16(vget_high_f16(input_vec)); - sum_squares_vec[u * 2] = vfmaq_f32(sum_squares_vec[u * 2], input_low, input_low); - sum_squares_vec[u * 2 + 1] = vfmaq_f32(sum_squares_vec[u * 2 + 1], input_high, input_high); - } - } - - const size_t simd_end = (dims >= SIMD_WIDTH) ? dims - SIMD_WIDTH + 1 : 0; - for (; i < simd_end; i += SIMD_WIDTH) { - float16x8_t input_vec = vld1q_f16(&input_row[i]); - float32x4_t input_low = vcvt_f32_f16(vget_low_f16(input_vec)); - float32x4_t input_high = vcvt_f32_f16(vget_high_f16(input_vec)); - sum_squares_vec[0] = vfmaq_f32(sum_squares_vec[0], input_low, input_low); - sum_squares_vec[1] = vfmaq_f32(sum_squares_vec[1], input_high, input_high); - } - - float32x4_t total_sum = sum_squares_vec[0]; - for (size_t u = 1; u < UNROLL_FACTOR * 2; u++) { - total_sum = vaddq_f32(total_sum, sum_squares_vec[u]); - } - float sum_squares = vaddvq_f32(total_sum); - - for (; i < dims; ++i) { - float val = static_cast(input_row[i]); - sum_squares += val * val; - } - - float rms = sqrtf(sum_squares / static_cast(dims) + eps); - float inv_rms = 1.0f / rms; - float16x8_t inv_rms_vec = vdupq_n_f16(static_cast<__fp16>(inv_rms)); - - i = 0; - for (; i < tile_end; i += TILE_SIZE) { - for (size_t u = 0; u < UNROLL_FACTOR; u++) { - float16x8_t input_vec = vld1q_f16(&input_row[i + u * SIMD_WIDTH]); - float16x8_t weight_vec = vld1q_f16(&weight[i + u * SIMD_WIDTH]); - float16x8_t norm_vec = vmulq_f16(vmulq_f16(input_vec, inv_rms_vec), weight_vec); - vst1q_f16(&output_row[i + u * SIMD_WIDTH], norm_vec); - } - } - - for (; i < simd_end; i += SIMD_WIDTH) { - float16x8_t input_vec = vld1q_f16(&input_row[i]); - float16x8_t weight_vec = vld1q_f16(&weight[i]); - float16x8_t norm_vec = vmulq_f16(vmulq_f16(input_vec, inv_rms_vec), weight_vec); - vst1q_f16(&output_row[i], norm_vec); - } - - for (; i < dims; ++i) { - output_row[i] = static_cast<__fp16>(static_cast(input_row[i]) * inv_rms * static_cast(weight[i])); - } - } -} - -namespace CactusRoPEF16 { - -struct RoPECacheF16 { - std::vector<__fp16> cos_table; - std::vector<__fp16> sin_table; - size_t max_seq_len; - size_t head_dim; - float theta; - bool initialized; - - RoPECacheF16() : max_seq_len(0), head_dim(0), theta(0.0f), initialized(false) {} -}; - -static thread_local RoPECacheF16 rope_cache_f16; - -void precompute_rope_tables_f16(size_t seq_len, size_t head_dim, float theta) { - if (rope_cache_f16.initialized && - rope_cache_f16.max_seq_len >= seq_len && - rope_cache_f16.head_dim == head_dim && - rope_cache_f16.theta == theta) { - return; - } - - const size_t half_dim = head_dim / 2; - const size_t table_size = seq_len * half_dim; - - rope_cache_f16.cos_table.resize(table_size); - rope_cache_f16.sin_table.resize(table_size); - - for (size_t pos = 0; pos < seq_len; ++pos) { - const float pos_float = static_cast(pos); - for (size_t i = 0; i < half_dim; ++i) { - const float freq = 1.0f / powf(theta, (2.0f * i) / head_dim); - const float angle = pos_float * freq; - - const size_t idx = pos * half_dim + i; - rope_cache_f16.cos_table[idx] = static_cast<__fp16>(cosf(angle)); - rope_cache_f16.sin_table[idx] = static_cast<__fp16>(sinf(angle)); - } - } - - rope_cache_f16.max_seq_len = seq_len; - rope_cache_f16.head_dim = head_dim; - rope_cache_f16.theta = theta; - rope_cache_f16.initialized = true; -} - -} - -void cactus_rope_f16( - const __fp16* input, - __fp16* output, - size_t batch_size, - size_t seq_len, - size_t num_heads, - size_t head_dim, - size_t start_pos, - float theta -) { - const size_t half_dim = head_dim / 2; - - CactusRoPEF16::precompute_rope_tables_f16(seq_len + start_pos, head_dim, theta); - - const __fp16* cos_cache = CactusRoPEF16::rope_cache_f16.cos_table.data() + start_pos * half_dim; - const __fp16* sin_cache = CactusRoPEF16::rope_cache_f16.sin_table.data() + start_pos * half_dim; - - CactusThreading::parallel_for(batch_size * seq_len, CactusThreading::Thresholds::SCALAR_EXPENSIVE, - [&](size_t start_idx, size_t end_idx) { - for (size_t idx = start_idx; idx < end_idx; ++idx) { - const size_t batch_idx = idx / seq_len; - const size_t seq_idx = idx % seq_len; - - for (size_t head_idx = 0; head_idx < num_heads; ++head_idx) { - const size_t offset = ((batch_idx * seq_len + seq_idx) * num_heads + head_idx) * head_dim; - const __fp16* input_ptr = input + offset; - __fp16* output_ptr = output + offset; - - const __fp16* cos_ptr = cos_cache + seq_idx * half_dim; - const __fp16* sin_ptr = sin_cache + seq_idx * half_dim; - - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_half_dim = (half_dim / SIMD_WIDTH) * SIMD_WIDTH; - - for (size_t i = 0; i < vectorized_half_dim; i += SIMD_WIDTH) { - float16x8_t cos_vec = vld1q_f16(&cos_ptr[i]); - float16x8_t sin_vec = vld1q_f16(&sin_ptr[i]); - - float16x8_t x_first_half = vld1q_f16(&input_ptr[i]); - float16x8_t x_second_half = vld1q_f16(&input_ptr[i + half_dim]); - - float16x8_t first_result = vfmsq_f16(vmulq_f16(x_first_half, cos_vec), x_second_half, sin_vec); - float16x8_t second_result = vfmaq_f16(vmulq_f16(x_second_half, cos_vec), x_first_half, sin_vec); - - vst1q_f16(&output_ptr[i], first_result); - vst1q_f16(&output_ptr[i + half_dim], second_result); - } - - for (size_t i = vectorized_half_dim; i < half_dim; ++i) { - const __fp16 cos_val = cos_ptr[i]; - const __fp16 sin_val = sin_ptr[i]; - - const __fp16 x_first_half = input_ptr[i]; - const __fp16 x_second_half = input_ptr[i + half_dim]; - - output_ptr[i] = x_first_half * cos_val - x_second_half * sin_val; - - output_ptr[i + half_dim] = x_second_half * cos_val + x_first_half * sin_val; - } - } - } - }); -} \ No newline at end of file diff --git a/cactus/kernel/kernel_conv.cpp b/cactus/kernel/kernel_conv.cpp deleted file mode 100644 index 71bf8e9d6..000000000 --- a/cactus/kernel/kernel_conv.cpp +++ /dev/null @@ -1,284 +0,0 @@ -#include "kernel.h" -#include "kernel_utils.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -constexpr size_t T_TILE_F16 = 2; - -void cactus_conv1d_causal_depthwise_f16( - const __fp16* input, - const __fp16* weight, - __fp16* output, - size_t N, size_t L, size_t C, size_t K, size_t dilation) -{ - const size_t in_bs = L * C; - const size_t out_bs = L * C; - - CactusThreading::parallel_for_2d(N, C, CactusThreading::Thresholds::ATTENTION, [&](size_t n, size_t c) { - const __fp16* Xb = input + n * in_bs; - __fp16* Yb = output + n * out_bs; - - std::vector wrev(K); - const __fp16* Wc = weight + c * K; - for (size_t k = 0; k < K; ++k) wrev[k] = (float)Wc[K - 1 - k]; - - for (size_t t0 = 0; t0 < L; t0 += T_TILE_F16) { - const size_t t1 = std::min(t0 + 1, L - 1); - - float32x4_t vacc0 = vdupq_n_f32(0.f); - float32x4_t vacc1 = vdupq_n_f32(0.f); - - size_t k = 0; - for (; k + 8 <= K; k += 8) { - - float x0_0=0, x1_0=0, x2_0=0, x3_0=0; - float x0_1=0, x1_1=0, x2_1=0, x3_1=0; - { - ptrdiff_t a0=(ptrdiff_t)t0-(ptrdiff_t)((k+0)*dilation); - ptrdiff_t a1=(ptrdiff_t)t0-(ptrdiff_t)((k+1)*dilation); - ptrdiff_t a2=(ptrdiff_t)t0-(ptrdiff_t)((k+2)*dilation); - ptrdiff_t a3=(ptrdiff_t)t0-(ptrdiff_t)((k+3)*dilation); - if (a0>=0) x0_0 = (float)Xb[(size_t)a0*C + c]; - if (a1>=0) x1_0 = (float)Xb[(size_t)a1*C + c]; - if (a2>=0) x2_0 = (float)Xb[(size_t)a2*C + c]; - if (a3>=0) x3_0 = (float)Xb[(size_t)a3*C + c]; - - ptrdiff_t b0=(ptrdiff_t)t1-(ptrdiff_t)((k+0)*dilation); - ptrdiff_t b1=(ptrdiff_t)t1-(ptrdiff_t)((k+1)*dilation); - ptrdiff_t b2=(ptrdiff_t)t1-(ptrdiff_t)((k+2)*dilation); - ptrdiff_t b3=(ptrdiff_t)t1-(ptrdiff_t)((k+3)*dilation); - if (b0>=0) x0_1 = (float)Xb[(size_t)b0*C + c]; - if (b1>=0) x1_1 = (float)Xb[(size_t)b1*C + c]; - if (b2>=0) x2_1 = (float)Xb[(size_t)b2*C + c]; - if (b3>=0) x3_1 = (float)Xb[(size_t)b3*C + c]; - } - float32x4_t xv0 = {x0_0,x1_0,x2_0,x3_0}; - float32x4_t yv0 = {x0_1,x1_1,x2_1,x3_1}; - float32x4_t wv0 = {wrev[k+0],wrev[k+1],wrev[k+2],wrev[k+3]}; - vacc0 = vfmaq_f32(vacc0, xv0, wv0); - vacc1 = vfmaq_f32(vacc1, yv0, wv0); - - float a0_0=0, a1_0=0, a2_0=0, a3_0=0; - float a0_1=0, a1_1=0, a2_1=0, a3_1=0; - { - ptrdiff_t a0i=(ptrdiff_t)t0-(ptrdiff_t)((k+4)*dilation); - ptrdiff_t a1i=(ptrdiff_t)t0-(ptrdiff_t)((k+5)*dilation); - ptrdiff_t a2i=(ptrdiff_t)t0-(ptrdiff_t)((k+6)*dilation); - ptrdiff_t a3i=(ptrdiff_t)t0-(ptrdiff_t)((k+7)*dilation); - if (a0i>=0) a0_0 = (float)Xb[(size_t)a0i*C + c]; - if (a1i>=0) a1_0 = (float)Xb[(size_t)a1i*C + c]; - if (a2i>=0) a2_0 = (float)Xb[(size_t)a2i*C + c]; - if (a3i>=0) a3_0 = (float)Xb[(size_t)a3i*C + c]; - - ptrdiff_t b0i=(ptrdiff_t)t1-(ptrdiff_t)((k+4)*dilation); - ptrdiff_t b1i=(ptrdiff_t)t1-(ptrdiff_t)((k+5)*dilation); - ptrdiff_t b2i=(ptrdiff_t)t1-(ptrdiff_t)((k+6)*dilation); - ptrdiff_t b3i=(ptrdiff_t)t1-(ptrdiff_t)((k+7)*dilation); - if (b0i>=0) a0_1 = (float)Xb[(size_t)b0i*C + c]; - if (b1i>=0) a1_1 = (float)Xb[(size_t)b1i*C + c]; - if (b2i>=0) a2_1 = (float)Xb[(size_t)b2i*C + c]; - if (b3i>=0) a3_1 = (float)Xb[(size_t)b3i*C + c]; - } - float32x4_t xv1 = {a0_0,a1_0,a2_0,a3_0}; - float32x4_t yv1 = {a0_1,a1_1,a2_1,a3_1}; - float32x4_t wv1 = {wrev[k+4],wrev[k+5],wrev[k+6],wrev[k+7]}; - vacc0 = vfmaq_f32(vacc0, xv1, wv1); - vacc1 = vfmaq_f32(vacc1, yv1, wv1); - } - - float acc0 = vaddvq_f32(vacc0); - float acc1 = vaddvq_f32(vacc1); - - for (; k < K; ++k) { - ptrdiff_t a=(ptrdiff_t)t0-(ptrdiff_t)(k*dilation); - if (a>=0) acc0 += wrev[k] * (float)Xb[(size_t)a*C + c]; - ptrdiff_t b=(ptrdiff_t)t1-(ptrdiff_t)(k*dilation); - if (b>=0) acc1 += wrev[k] * (float)Xb[(size_t)b*C + c]; - } - - Yb[t0*C + c] = (__fp16)acc0; - if (t0 + 1 < L) Yb[t1*C + c] = (__fp16)acc1; - } - }); -} - -void cactus_conv1d_f16_k3( - const __fp16* input, - const __fp16* weight, - __fp16* output, - size_t N, size_t L, - size_t C_in, size_t C_out, - size_t stride -){ - const size_t out_len = ((L - 1) / stride) + 1; - const size_t in_bs = C_in * L; - const size_t out_bs = C_out * out_len; - - CactusThreading::parallel_for_2d(N, C_out, CactusThreading::Thresholds::ATTENTION, [&](size_t n, size_t oc) { - const __fp16* Xb = input + n * in_bs; - __fp16* Yoc = output + n * out_bs + oc * out_len; - const __fp16* Woc = weight + oc * (C_in * 3); - - for (size_t out_idx = 0; out_idx < out_len; out_idx += 2) { - const size_t out_t0 = out_idx; - const bool have_t1 = (out_idx + 1) < out_len; - const size_t out_t1 = have_t1 ? (out_idx + 1) : out_idx; - - const size_t t0 = out_t0 * stride; - const size_t t1 = have_t1 ? (out_t1 * stride) : t0; - - float32x4_t acc0 = vdupq_n_f32(0.f); - float32x4_t acc1 = vdupq_n_f32(0.f); - - size_t ic = 0; - for (; ic + 16 <= C_in; ic += 16) { - for (size_t u = 0; u < 16; ++u) { - const __fp16* Xc = Xb + (ic + u) * L; - const __fp16* Wc = Woc + (ic + u) * 3; - - const float16x8_t wv = { - Wc[0], Wc[1], Wc[2], (__fp16)0, - Wc[0], Wc[1], Wc[2], (__fp16)0 - }; - - const ptrdiff_t tm0 = (ptrdiff_t)t0 - 1; - const ptrdiff_t tp0 = (ptrdiff_t)t0 + 1; - const ptrdiff_t tm1 = (ptrdiff_t)t1 - 1; - const ptrdiff_t tp1 = (ptrdiff_t)t1 + 1; - - const __fp16 x0m = (tm0 >= 0) ? Xc[tm0] : (__fp16)0; - const __fp16 x00 = Xc[t0]; - const __fp16 x0p = (tp0 < (ptrdiff_t)L) ? Xc[tp0] : (__fp16)0; - - __fp16 x1m = 0, x10 = 0, x1p = 0; - if (have_t1) { - x1m = (tm1 >= 0) ? Xc[tm1] : (__fp16)0; - x10 = Xc[t1]; - x1p = (tp1 < (ptrdiff_t)L) ? Xc[tp1] : (__fp16)0; - } - - const float16x8_t xv = { - x0m, x00, x0p, (__fp16)0, - x1m, x10, x1p, (__fp16)0 - }; - - const float16x4_t xv0_h = vget_low_f16(xv); - const float16x4_t wv0_h = vget_low_f16(wv); - acc0 = vfmaq_f32(acc0, vcvt_f32_f16(xv0_h), vcvt_f32_f16(wv0_h)); - - if (have_t1) { - const float16x4_t xv1_h = vget_high_f16(xv); - const float16x4_t wv1_h = vget_high_f16(wv); - acc1 = vfmaq_f32(acc1, vcvt_f32_f16(xv1_h), vcvt_f32_f16(wv1_h)); - } - } - } - - for (; ic < C_in; ++ic) { - const __fp16* Xc = Xb + ic * L; - const __fp16* Wc = Woc + ic * 3; - - const float16x8_t wv = { - Wc[0], Wc[1], Wc[2], (__fp16)0, - Wc[0], Wc[1], Wc[2], (__fp16)0 - }; - - const ptrdiff_t tm0 = (ptrdiff_t)t0 - 1; - const ptrdiff_t tp0 = (ptrdiff_t)t0 + 1; - const ptrdiff_t tm1 = (ptrdiff_t)t1 - 1; - const ptrdiff_t tp1 = (ptrdiff_t)t1 + 1; - - const __fp16 x0m = (tm0 >= 0) ? Xc[tm0] : (__fp16)0; - const __fp16 x00 = Xc[t0]; - const __fp16 x0p = (tp0 < (ptrdiff_t)L) ? Xc[tp0] : (__fp16)0; - - __fp16 x1m = 0, x10 = 0, x1p = 0; - if (have_t1) { - x1m = (tm1 >= 0) ? Xc[tm1] : (__fp16)0; - x10 = Xc[t1]; - x1p = (tp1 < (ptrdiff_t)L) ? Xc[tp1] : (__fp16)0; - } - - const float16x8_t xv = { - x0m, x00, x0p, (__fp16)0, - x1m, x10, x1p, (__fp16)0 - }; - - const float16x4_t xv0_h = vget_low_f16(xv); - const float16x4_t wv0_h = vget_low_f16(wv); - acc0 = vfmaq_f32(acc0, vcvt_f32_f16(xv0_h), vcvt_f32_f16(wv0_h)); - - if (have_t1) { - const float16x4_t xv1_h = vget_high_f16(xv); - const float16x4_t wv1_h = vget_high_f16(wv); - acc1 = vfmaq_f32(acc1, vcvt_f32_f16(xv1_h), vcvt_f32_f16(wv1_h)); - } - } - - float32x2_t s0 = vadd_f32(vget_low_f32(acc0), vget_high_f32(acc0)); - float sum0 = vget_lane_f32(s0, 0) + vget_lane_f32(s0, 1); - Yoc[out_t0] = (__fp16)sum0; - - if (have_t1) { - float32x2_t s1 = vadd_f32(vget_low_f32(acc1), vget_high_f32(acc1)); - float sum1 = vget_lane_f32(s1, 0) + vget_lane_f32(s1, 1); - Yoc[out_t1] = (__fp16)sum1; - } - } - }); -} - -void cactus_bilinear_interpolation_f16(const __fp16* input, __fp16* output, size_t src_height, size_t src_width, size_t embed_dim, - size_t dst_height, size_t dst_width) -{ - float scale_h = (src_height > 1 && dst_height > 1) - ? static_cast(src_height - 1) / static_cast(dst_height - 1) - : 0.0f; - float scale_w = (src_width > 1 && dst_width > 1) - ? static_cast(src_width - 1) / static_cast(dst_width - 1) - : 0.0f; - - for (size_t dst_y = 0; dst_y < dst_height; ++dst_y) { - for (size_t dst_x = 0; dst_x < dst_width; ++dst_x) { - float src_y_float = dst_y * scale_h; - float src_x_float = dst_x * scale_w; - - int y0 = static_cast(std::floor(src_y_float)); - int x0 = static_cast(std::floor(src_x_float)); - - int y1 = ((y0 + 1) < static_cast(src_height)) ? (y0 + 1) : (static_cast(src_height) - 1); - int x1 = ((x0 + 1) < static_cast(src_width)) ? (x0 + 1) : (static_cast(src_width) - 1); - - float dy = src_y_float - y0; - float dx = src_x_float - x0; - - float w00 = (1.0f - dx) * (1.0f - dy); - float w01 = dx * (1.0f - dy); - float w10 = (1.0f - dx) * dy; - float w11 = dx * dy; - - size_t idx00 = (y0 * static_cast(src_width) + x0) * static_cast(embed_dim); - size_t idx01 = (y0 * static_cast(src_width) + x1) * static_cast(embed_dim); - size_t idx10 = (y1 * static_cast(src_width) + x0) * static_cast(embed_dim); - size_t idx11 = (y1 * static_cast(src_width) + x1) * static_cast(embed_dim); - - size_t out_idx = (dst_y * dst_width + dst_x) * embed_dim; - - for (size_t d = 0; d < embed_dim; ++d) { - float result = - static_cast(input[idx00 + d]) * w00 + - static_cast(input[idx01 + d]) * w01 + - static_cast(input[idx10 + d]) * w10 + - static_cast(input[idx11 + d]) * w11; - output[out_idx + d] = static_cast<__fp16>(result); - } - } - } -} \ No newline at end of file diff --git a/cactus/kernel/kernel_gemm.cpp b/cactus/kernel/kernel_gemm.cpp deleted file mode 100644 index 77496c2df..000000000 --- a/cactus/kernel/kernel_gemm.cpp +++ /dev/null @@ -1,529 +0,0 @@ -#include "kernel.h" -#include "kernel_utils.h" -#include -#include -#include -#include - -static inline __fp16 hsum_f16x8(float16x8_t v) { - float16x4_t lo = vget_low_f16(v); - float16x4_t hi = vget_high_f16(v); - float16x4_t sum4 = vadd_f16(lo, hi); - float16x4_t sum2 = vadd_f16(sum4, vext_f16(sum4, sum4, 2)); - float16x4_t sum1 = vadd_f16(sum2, vext_f16(sum2, sum2, 1)); - return vget_lane_f16(sum1, 0); -} - -static void cactus_matmul_f16_worker( - const __fp16* a, - const __fp16* b_transposed, - __fp16* c, - size_t /*M*/, - size_t K, - size_t N, - size_t start_row, - size_t end_row -) { - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 4; - const size_t K16 = (K / 16) * 16; - const size_t K8 = (K / 8) * 8; - - for (size_t row_block = start_row; row_block < end_row; row_block += TILE_M) { - const size_t m_end = std::min(row_block + TILE_M, end_row); - - for (size_t col_block = 0; col_block < N; col_block += TILE_N) { - const size_t n_end = std::min(col_block + TILE_N, N); - - float16x8_t acc[TILE_M][TILE_N]; - for (size_t m = 0; m < TILE_M; ++m) - for (size_t n = 0; n < TILE_N; ++n) - acc[m][n] = vdupq_n_f16(0); - - for (size_t k = 0; k < K16; k += 16) { - float16x8_t a0_lo = (row_block < m_end) ? vld1q_f16(a + row_block * K + k) : vdupq_n_f16(0); - float16x8_t a0_hi = (row_block < m_end) ? vld1q_f16(a + row_block * K + k + 8) : vdupq_n_f16(0); - float16x8_t a1_lo = (row_block + 1 < m_end) ? vld1q_f16(a + (row_block + 1) * K + k) : vdupq_n_f16(0); - float16x8_t a1_hi = (row_block + 1 < m_end) ? vld1q_f16(a + (row_block + 1) * K + k + 8) : vdupq_n_f16(0); - float16x8_t a2_lo = (row_block + 2 < m_end) ? vld1q_f16(a + (row_block + 2) * K + k) : vdupq_n_f16(0); - float16x8_t a2_hi = (row_block + 2 < m_end) ? vld1q_f16(a + (row_block + 2) * K + k + 8) : vdupq_n_f16(0); - float16x8_t a3_lo = (row_block + 3 < m_end) ? vld1q_f16(a + (row_block + 3) * K + k) : vdupq_n_f16(0); - float16x8_t a3_hi = (row_block + 3 < m_end) ? vld1q_f16(a + (row_block + 3) * K + k + 8) : vdupq_n_f16(0); - - for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { - float16x8_t b_lo = vld1q_f16(b_transposed + (col_block + ni) * K + k); - float16x8_t b_hi = vld1q_f16(b_transposed + (col_block + ni) * K + k + 8); - - acc[0][ni] = vfmaq_f16(acc[0][ni], a0_lo, b_lo); - acc[0][ni] = vfmaq_f16(acc[0][ni], a0_hi, b_hi); - acc[1][ni] = vfmaq_f16(acc[1][ni], a1_lo, b_lo); - acc[1][ni] = vfmaq_f16(acc[1][ni], a1_hi, b_hi); - acc[2][ni] = vfmaq_f16(acc[2][ni], a2_lo, b_lo); - acc[2][ni] = vfmaq_f16(acc[2][ni], a2_hi, b_hi); - acc[3][ni] = vfmaq_f16(acc[3][ni], a3_lo, b_lo); - acc[3][ni] = vfmaq_f16(acc[3][ni], a3_hi, b_hi); - } - } - - for (size_t k = K16; k < K8; k += 8) { - float16x8_t a0_v = (row_block < m_end) ? vld1q_f16(a + row_block * K + k) : vdupq_n_f16(0); - float16x8_t a1_v = (row_block + 1 < m_end) ? vld1q_f16(a + (row_block + 1) * K + k) : vdupq_n_f16(0); - float16x8_t a2_v = (row_block + 2 < m_end) ? vld1q_f16(a + (row_block + 2) * K + k) : vdupq_n_f16(0); - float16x8_t a3_v = (row_block + 3 < m_end) ? vld1q_f16(a + (row_block + 3) * K + k) : vdupq_n_f16(0); - - for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { - float16x8_t b_v = vld1q_f16(b_transposed + (col_block + ni) * K + k); - acc[0][ni] = vfmaq_f16(acc[0][ni], a0_v, b_v); - acc[1][ni] = vfmaq_f16(acc[1][ni], a1_v, b_v); - acc[2][ni] = vfmaq_f16(acc[2][ni], a2_v, b_v); - acc[3][ni] = vfmaq_f16(acc[3][ni], a3_v, b_v); - } - } - - for (size_t k = K8; k < K; ++k) { - for (size_t mi = 0; mi < TILE_M && row_block + mi < m_end; ++mi) { - __fp16 av = a[(row_block + mi) * K + k]; - for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { - __fp16 bv = b_transposed[(col_block + ni) * K + k]; - acc[mi][ni] = vsetq_lane_f16(vgetq_lane_f16(acc[mi][ni], 0) + av * bv, acc[mi][ni], 0); - } - } - } - - for (size_t mi = 0; mi < TILE_M && row_block + mi < m_end; ++mi) { - for (size_t ni = 0; ni < TILE_N && col_block + ni < n_end; ++ni) { - c[(row_block + mi) * N + col_block + ni] = hsum_f16x8(acc[mi][ni]); - } - } - } - } -} - -void cactus_matmul_f16( - const __fp16* a, - const __fp16* b_transposed, - __fp16* c, - size_t M, - size_t K, - size_t N -) { - constexpr size_t TILE_M = 4; - const size_t num_row_blocks = (M + TILE_M - 1) / TILE_M; - - CactusThreading::parallel_for(num_row_blocks, CactusThreading::Thresholds::SCALAR_EXPENSIVE, - [=](size_t start_block, size_t end_block) { - for (size_t block_idx = start_block; block_idx < end_block; ++block_idx) { - size_t start_row = block_idx * TILE_M; - size_t end_row = std::min(start_row + TILE_M, M); - - cactus_matmul_f16_worker( - a, b_transposed, c, - M, K, N, - start_row, end_row - ); - } - }); -} - -void cactus_matmul_int8( - const int8_t* A, - const float* A_scales, - const int8_t* B, - const __fp16* B_scales, - __fp16* C, - size_t M, size_t K, size_t N, - size_t group_size -) { - if (M == 0 || K == 0 || N == 0) return; - - #if defined(__APPLE__) && defined(__arm64__) - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 8; - #else - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 4; - #endif - - const size_t num_groups = K / group_size; - const size_t num_row_tiles = (M + TILE_M - 1) / TILE_M; - const size_t num_col_tiles = (N + TILE_N - 1) / TILE_N; - const size_t total_tiles = num_row_tiles * num_col_tiles; - - CactusThreading::parallel_gemm_tiles(M, total_tiles, - [=](size_t tile_start, size_t tile_end) { - for (size_t tile_idx = tile_start; tile_idx < tile_end; ++tile_idx) { - const size_t tile_row = tile_idx / num_col_tiles; - const size_t tile_col = tile_idx % num_col_tiles; - const size_t m_start = tile_row * TILE_M; - const size_t m_end = std::min(m_start + TILE_M, M); - const size_t n_start = tile_col * TILE_N; - const size_t n_end = std::min(n_start + TILE_N, N); - const size_t actual_m = m_end - m_start; - const size_t actual_n = n_end - n_start; - - float running_sum[TILE_M][TILE_N] = {{0.0f}}; - - for (size_t ni = 0; ni < actual_n; ni++) { - __builtin_prefetch(&B_scales[(n_start + ni) * num_groups], 0, 3); - } - - for (size_t g = 0; g < num_groups; g++) { - const size_t k_base = g * group_size; - - int32_t group_acc[TILE_M][TILE_N] = {{0}}; - - if (cactus_has_i8mm()) { - size_t mi = 0; - for (; mi + 1 < actual_m; mi += 2) { - size_t ni = 0; - for (; ni + 3 < actual_n; ni += 4) { - int32x4_t acc01 = vdupq_n_s32(0); - int32x4_t acc23 = vdupq_n_s32(0); - - const int8_t* a_base0 = A + (m_start + mi) * K + k_base; - const int8_t* a_base1 = A + (m_start + mi + 1) * K + k_base; - const int8_t* b_base0 = B + (n_start + ni) * K + k_base; - const int8_t* b_base1 = B + (n_start + ni + 1) * K + k_base; - const int8_t* b_base2 = B + (n_start + ni + 2) * K + k_base; - const int8_t* b_base3 = B + (n_start + ni + 3) * K + k_base; - - for (size_t k_offset = 0; k_offset < group_size; k_offset += 64) { - #pragma unroll - for (int kk = 0; kk < 64; kk += 16) { - int8x16_t a_lo = vcombine_s8(vld1_s8(a_base0 + k_offset + kk), - vld1_s8(a_base1 + k_offset + kk)); - int8x16_t a_hi = vcombine_s8(vld1_s8(a_base0 + k_offset + kk + 8), - vld1_s8(a_base1 + k_offset + kk + 8)); - - int8x16_t b01_lo = vcombine_s8(vld1_s8(b_base0 + k_offset + kk), - vld1_s8(b_base1 + k_offset + kk)); - int8x16_t b01_hi = vcombine_s8(vld1_s8(b_base0 + k_offset + kk + 8), - vld1_s8(b_base1 + k_offset + kk + 8)); - acc01 = accum_matmul(acc01, a_lo, b01_lo); - acc01 = accum_matmul(acc01, a_hi, b01_hi); - - int8x16_t b23_lo = vcombine_s8(vld1_s8(b_base2 + k_offset + kk), - vld1_s8(b_base3 + k_offset + kk)); - int8x16_t b23_hi = vcombine_s8(vld1_s8(b_base2 + k_offset + kk + 8), - vld1_s8(b_base3 + k_offset + kk + 8)); - acc23 = accum_matmul(acc23, a_lo, b23_lo); - acc23 = accum_matmul(acc23, a_hi, b23_hi); - } - } - - group_acc[mi][ni] += vgetq_lane_s32(acc01, 0); - group_acc[mi][ni + 1] += vgetq_lane_s32(acc01, 1); - group_acc[mi + 1][ni] += vgetq_lane_s32(acc01, 2); - group_acc[mi + 1][ni + 1] += vgetq_lane_s32(acc01, 3); - group_acc[mi][ni + 2] += vgetq_lane_s32(acc23, 0); - group_acc[mi][ni + 3] += vgetq_lane_s32(acc23, 1); - group_acc[mi + 1][ni + 2] += vgetq_lane_s32(acc23, 2); - group_acc[mi + 1][ni + 3] += vgetq_lane_s32(acc23, 3); - } - - for (; ni + 1 < actual_n; ni += 2) { - int32x4_t acc0 = vdupq_n_s32(0); - - for (size_t k_offset = 0; k_offset < group_size; k_offset += 64) { - const int8_t* a_ptr0 = A + (m_start + mi) * K + k_base + k_offset; - const int8_t* a_ptr1 = A + (m_start + mi + 1) * K + k_base + k_offset; - const int8_t* b_ptr0 = B + (n_start + ni) * K + k_base + k_offset; - const int8_t* b_ptr1 = B + (n_start + ni + 1) * K + k_base + k_offset; - - for (int kk = 0; kk < 64; kk += 8) { - int8x16_t a_vec = vcombine_s8(vld1_s8(a_ptr0 + kk), vld1_s8(a_ptr1 + kk)); - int8x16_t b_vec = vcombine_s8(vld1_s8(b_ptr0 + kk), vld1_s8(b_ptr1 + kk)); - acc0 = accum_matmul(acc0, a_vec, b_vec); - } - } - - group_acc[mi][ni] += vgetq_lane_s32(acc0, 0); - group_acc[mi][ni + 1] += vgetq_lane_s32(acc0, 1); - group_acc[mi + 1][ni] += vgetq_lane_s32(acc0, 2); - group_acc[mi + 1][ni + 1] += vgetq_lane_s32(acc0, 3); - } - - for (; ni < actual_n; ni++) { - for (size_t mii = mi; mii < mi + 2; mii++) { - for (size_t k_offset = 0; k_offset < group_size; k_offset += 64) { - const int8_t* a_ptr = A + (m_start + mii) * K + k_base + k_offset; - const int8_t* b_ptr = B + (n_start + ni) * K + k_base + k_offset; - int32x4_t sum = vdupq_n_s32(0); - sum = accum_dot(sum, vld1q_s8(a_ptr), vld1q_s8(b_ptr)); - sum = accum_dot(sum, vld1q_s8(a_ptr + 16), vld1q_s8(b_ptr + 16)); - sum = accum_dot(sum, vld1q_s8(a_ptr + 32), vld1q_s8(b_ptr + 32)); - sum = accum_dot(sum, vld1q_s8(a_ptr + 48), vld1q_s8(b_ptr + 48)); - group_acc[mii][ni] += vaddvq_s32(sum); - } - } - } - } - - for (; mi < actual_m; mi++) { - for (size_t k_offset = 0; k_offset < group_size; k_offset += 64) { - const int8_t* a_ptr = A + (m_start + mi) * K + k_base + k_offset; - int8x16_t a_vec0 = vld1q_s8(a_ptr); - int8x16_t a_vec1 = vld1q_s8(a_ptr + 16); - int8x16_t a_vec2 = vld1q_s8(a_ptr + 32); - int8x16_t a_vec3 = vld1q_s8(a_ptr + 48); - - for (size_t ni = 0; ni < actual_n; ni++) { - const int8_t* b_ptr = B + (n_start + ni) * K + k_base + k_offset; - int32x4_t sum = vdupq_n_s32(0); - sum = accum_dot(sum, a_vec0, vld1q_s8(b_ptr)); - sum = accum_dot(sum, a_vec1, vld1q_s8(b_ptr + 16)); - sum = accum_dot(sum, a_vec2, vld1q_s8(b_ptr + 32)); - sum = accum_dot(sum, a_vec3, vld1q_s8(b_ptr + 48)); - group_acc[mi][ni] += vaddvq_s32(sum); - } - } - } - } else { - for (size_t k_offset = 0; k_offset < group_size; k_offset += 64) { - int8x16_t b_vec0[TILE_N], b_vec1[TILE_N], b_vec2[TILE_N], b_vec3[TILE_N]; - for (size_t ni = 0; ni < actual_n; ni++) { - const int8_t* b_ptr = B + (n_start + ni) * K + k_base + k_offset; - b_vec0[ni] = vld1q_s8(b_ptr); - b_vec1[ni] = vld1q_s8(b_ptr + 16); - b_vec2[ni] = vld1q_s8(b_ptr + 32); - b_vec3[ni] = vld1q_s8(b_ptr + 48); - } - - for (size_t mi = 0; mi < actual_m; mi++) { - const int8_t* a_ptr = A + (m_start + mi) * K + k_base + k_offset; - int8x16_t a_vec0 = vld1q_s8(a_ptr); - int8x16_t a_vec1 = vld1q_s8(a_ptr + 16); - int8x16_t a_vec2 = vld1q_s8(a_ptr + 32); - int8x16_t a_vec3 = vld1q_s8(a_ptr + 48); - - for (size_t ni = 0; ni < actual_n; ni++) { - int32x4_t sum = vdupq_n_s32(0); - sum = accum_dot(sum, a_vec0, b_vec0[ni]); - sum = accum_dot(sum, a_vec1, b_vec1[ni]); - sum = accum_dot(sum, a_vec2, b_vec2[ni]); - sum = accum_dot(sum, a_vec3, b_vec3[ni]); - group_acc[mi][ni] += vaddvq_s32(sum); - } - } - } - } - - for (size_t ni = 0; ni < actual_n; ni++) { - const float b_scale = (float)B_scales[(n_start + ni) * num_groups + g]; - for (size_t mi = 0; mi < actual_m; mi++) { - running_sum[mi][ni] += (float)group_acc[mi][ni] * b_scale; - } - } - } - - for (size_t mi = 0; mi < actual_m; mi++) { - const float a_scale = A_scales[m_start + mi]; - for (size_t ni = 0; ni < actual_n; ni++) { - C[(m_start + mi) * N + (n_start + ni)] = (__fp16)(running_sum[mi][ni] * a_scale); - } - } - } // tile_idx - }); -} - -void cactus_matmul_int4( - const int8_t* A, - const float* A_scales, - const uint8_t* B_packed, - const __fp16* B_scales, - __fp16* C, - size_t M, size_t K, size_t N, - size_t group_size -) { - if (M == 0 || K == 0 || N == 0) return; - - #if defined(__APPLE__) && defined(__arm64__) - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 8; - #else - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 4; - #endif - - const size_t num_groups = K / group_size; - const size_t K_packed = K / 2; - - const size_t num_row_tiles = (M + TILE_M - 1) / TILE_M; - const size_t num_col_tiles = (N + TILE_N - 1) / TILE_N; - const size_t total_tiles = num_row_tiles * num_col_tiles; - - if (M == 1) { - const int8_t* a_row = A; - const float a_scale = A_scales[0]; - - CactusThreading::parallel_for(N, CactusThreading::Thresholds::ELEMENT_WISE, - [=](size_t n_start, size_t n_end) { - for (size_t n = n_start; n < n_end; n++) { - const uint8_t* b_col = B_packed + n * K_packed; - const __fp16* col_scales = &B_scales[n * num_groups]; - - float sum = 0.0f; - for (size_t g = 0; g < num_groups; g++) { - int32_t group_sum = int4_dot_m1_asm( - a_row + g * group_size, - b_col + (g * group_size) / 2, - group_size - ); - sum += (float)group_sum * (float)col_scales[g]; - } - C[n] = (__fp16)(sum * a_scale); - } - }); - return; - } - - CactusThreading::parallel_gemm_tiles(M, total_tiles, - [=](size_t tile_start, size_t tile_end) { - alignas(64) int8_t B_unpacked[TILE_N][128]; - - for (size_t tile_idx = tile_start; tile_idx < tile_end; ++tile_idx) { - const size_t tile_row = tile_idx / num_col_tiles; - const size_t tile_col = tile_idx % num_col_tiles; - const size_t m_start = tile_row * TILE_M; - const size_t m_end = std::min(m_start + TILE_M, M); - const size_t n_start = tile_col * TILE_N; - const size_t n_end = std::min(n_start + TILE_N, N); - const size_t actual_m = m_end - m_start; - const size_t actual_n = n_end - n_start; - - float running_sum[TILE_M][TILE_N] = {{0.0f}}; - - for (size_t ni = 0; ni < actual_n; ni++) { - __builtin_prefetch(&B_scales[(n_start + ni) * num_groups], 0, 3); - } - - for (size_t g = 0; g < num_groups; g++) { - const size_t k_base = g * group_size; - const size_t k_base_packed = k_base / 2; - - int32_t group_acc[TILE_M][TILE_N] = {{0}}; - - for (size_t ni = 0; ni < actual_n; ni++) { - const uint8_t* b_ptr = B_packed + (n_start + ni) * K_packed + k_base_packed; - int8_t* dst = B_unpacked[ni]; - - for (size_t k = 0; k < group_size; k += 32) { - uint8x16_t packed = vld1q_u8(b_ptr + k / 2); - int8x16_t lo, hi; - unpack_int4_to_int8x32(packed, lo, hi); - vst1q_s8(dst + k, lo); - vst1q_s8(dst + k + 16, hi); - } - } - - if (cactus_has_i8mm()) { - size_t mi = 0; - for (; mi + 1 < actual_m; mi += 2) { - const int8_t* a_base0 = A + (m_start + mi) * K + k_base; - const int8_t* a_base1 = A + (m_start + mi + 1) * K + k_base; - - size_t ni = 0; - for (; ni + 3 < actual_n; ni += 4) { - int32x4_t acc01 = vdupq_n_s32(0); - int32x4_t acc23 = vdupq_n_s32(0); - - const int8_t* b0 = B_unpacked[ni]; - const int8_t* b1 = B_unpacked[ni + 1]; - const int8_t* b2 = B_unpacked[ni + 2]; - const int8_t* b3 = B_unpacked[ni + 3]; - - for (size_t k_offset = 0; k_offset < group_size; k_offset += 8) { - int8x8_t a0_8 = vld1_s8(a_base0 + k_offset); - int8x8_t a1_8 = vld1_s8(a_base1 + k_offset); - int8x16_t a_combined = vcombine_s8(a0_8, a1_8); - - int8x8_t b0_8 = vld1_s8(b0 + k_offset); - int8x8_t b1_8 = vld1_s8(b1 + k_offset); - int8x8_t b2_8 = vld1_s8(b2 + k_offset); - int8x8_t b3_8 = vld1_s8(b3 + k_offset); - - acc01 = accum_matmul(acc01, a_combined, vcombine_s8(b0_8, b1_8)); - acc23 = accum_matmul(acc23, a_combined, vcombine_s8(b2_8, b3_8)); - } - - group_acc[mi][ni] += vgetq_lane_s32(acc01, 0); - group_acc[mi][ni + 1] += vgetq_lane_s32(acc01, 1); - group_acc[mi + 1][ni] += vgetq_lane_s32(acc01, 2); - group_acc[mi + 1][ni + 1] += vgetq_lane_s32(acc01, 3); - group_acc[mi][ni + 2] += vgetq_lane_s32(acc23, 0); - group_acc[mi][ni + 3] += vgetq_lane_s32(acc23, 1); - group_acc[mi + 1][ni + 2] += vgetq_lane_s32(acc23, 2); - group_acc[mi + 1][ni + 3] += vgetq_lane_s32(acc23, 3); - } - - // Handle remaining columns - for (; ni < actual_n; ni++) { - const int8_t* b_col = B_unpacked[ni]; - for (size_t mii = mi; mii < mi + 2 && mii < actual_m; mii++) { - const int8_t* a_ptr = A + (m_start + mii) * K + k_base; - int32x4_t sum = vdupq_n_s32(0); - for (size_t k_offset = 0; k_offset < group_size; k_offset += 16) { - sum = accum_dot(sum, vld1q_s8(a_ptr + k_offset), vld1q_s8(b_col + k_offset)); - } - group_acc[mii][ni] += vaddvq_s32(sum); - } - } - } - - for (; mi < actual_m; mi++) { - const int8_t* a_base = A + (m_start + mi) * K + k_base; - for (size_t ni = 0; ni < actual_n; ni++) { - const int8_t* b_col = B_unpacked[ni]; - int32x4_t sum = vdupq_n_s32(0); - for (size_t k_offset = 0; k_offset < group_size; k_offset += 16) { - sum = accum_dot(sum, vld1q_s8(a_base + k_offset), vld1q_s8(b_col + k_offset)); - } - group_acc[mi][ni] += vaddvq_s32(sum); - } - } - } else { - for (size_t mi = 0; mi < actual_m; mi++) { - const int8_t* a_ptr = A + (m_start + mi) * K + k_base; - - for (size_t ni = 0; ni < actual_n; ni++) { - const int8_t* b_col = B_unpacked[ni]; - int32x4_t sum = vdupq_n_s32(0); - - for (size_t k_offset = 0; k_offset < group_size; k_offset += 64) { - int8x16_t a0 = vld1q_s8(a_ptr + k_offset); - int8x16_t a1 = vld1q_s8(a_ptr + k_offset + 16); - int8x16_t a2 = vld1q_s8(a_ptr + k_offset + 32); - int8x16_t a3 = vld1q_s8(a_ptr + k_offset + 48); - - int8x16_t b0 = vld1q_s8(b_col + k_offset); - int8x16_t b1 = vld1q_s8(b_col + k_offset + 16); - int8x16_t b2 = vld1q_s8(b_col + k_offset + 32); - int8x16_t b3 = vld1q_s8(b_col + k_offset + 48); - - sum = accum_dot(sum, a0, b0); - sum = accum_dot(sum, a1, b1); - sum = accum_dot(sum, a2, b2); - sum = accum_dot(sum, a3, b3); - } - group_acc[mi][ni] += vaddvq_s32(sum); - } - } - } - - for (size_t ni = 0; ni < actual_n; ni++) { - const float b_scale = (float)B_scales[(n_start + ni) * num_groups + g]; - for (size_t mi = 0; mi < actual_m; mi++) { - running_sum[mi][ni] += (float)group_acc[mi][ni] * b_scale; - } - } - } - - for (size_t mi = 0; mi < actual_m; mi++) { - const float a_scale = A_scales[m_start + mi]; - for (size_t ni = 0; ni < actual_n; ni++) { - C[(m_start + mi) * N + (n_start + ni)] = (__fp16)(running_sum[mi][ni] * a_scale); - } - } - } // tile_idx - }); -} \ No newline at end of file diff --git a/cactus/kernel/kernel_reduce.cpp b/cactus/kernel/kernel_reduce.cpp deleted file mode 100644 index de0cafa69..000000000 --- a/cactus/kernel/kernel_reduce.cpp +++ /dev/null @@ -1,323 +0,0 @@ -#include "kernel.h" -#include "kernel_utils.h" -#include -#include -#include -#include -#include - -double cactus_sum_all_f16(const __fp16* data, size_t num_elements) { - return CactusThreading::parallel_reduce( - num_elements, CactusThreading::Thresholds::ALL_REDUCE, - [&](size_t start_idx, size_t end_idx) -> double { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - float16x8_t sum_vec = vdupq_n_f16(0.0f); - - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t input_vec = vld1q_f16(&data[i]); - sum_vec = vaddq_f16(sum_vec, input_vec); - } - - double thread_sum = 0.0; - __fp16 sum_array[8]; - vst1q_f16(sum_array, sum_vec); - for (int j = 0; j < 8; j++) { - thread_sum += static_cast(sum_array[j]); - } - - for (size_t i = vectorized_end; i < end_idx; ++i) { - thread_sum += static_cast(data[i]); - } - - return thread_sum; - }, - 0.0, - [](double a, double b) { return a + b; } - ); -} - -void cactus_sum_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { - CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, - [&](size_t outer, size_t inner) { - float16x8_t sum_vec = vdupq_n_f16(0.0f); - - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_axis = (axis_size / SIMD_WIDTH) * SIMD_WIDTH; - - for (size_t a = 0; a < vectorized_axis; a += SIMD_WIDTH) { - __fp16 values[SIMD_WIDTH]; - for (size_t j = 0; j < SIMD_WIDTH; j++) { - size_t idx = outer * axis_size * inner_size + (a + j) * inner_size + inner; - values[j] = input[idx]; - } - float16x8_t input_vec = vld1q_f16(values); - sum_vec = vaddq_f16(sum_vec, input_vec); - } - - __fp16 total_sum = 0.0f; - __fp16 sum_array[8]; - vst1q_f16(sum_array, sum_vec); - for (int j = 0; j < 8; j++) { - total_sum += sum_array[j]; - } - - for (size_t a = vectorized_axis; a < axis_size; a++) { - size_t idx = outer * axis_size * inner_size + a * inner_size + inner; - total_sum += input[idx]; - } - - size_t output_idx = outer * inner_size + inner; - output[output_idx] = total_sum; - }); -} - -double cactus_mean_all_f16(const __fp16* data, size_t num_elements) { - double sum = cactus_sum_all_f16(data, num_elements); - return sum / static_cast(num_elements); -} - -void cactus_mean_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { - CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, - [&](size_t outer, size_t inner) { - float16x8_t sum_vec = vdupq_n_f16(0.0f); - - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_axis = (axis_size / SIMD_WIDTH) * SIMD_WIDTH; - - for (size_t a = 0; a < vectorized_axis; a += SIMD_WIDTH) { - __fp16 values[SIMD_WIDTH]; - for (size_t j = 0; j < SIMD_WIDTH; j++) { - size_t idx = outer * axis_size * inner_size + (a + j) * inner_size + inner; - values[j] = input[idx]; - } - float16x8_t input_vec = vld1q_f16(values); - sum_vec = vaddq_f16(sum_vec, input_vec); - } - - __fp16 total_sum = 0.0f; - __fp16 sum_array[8]; - vst1q_f16(sum_array, sum_vec); - for (int j = 0; j < 8; j++) { - total_sum += sum_array[j]; - } - - for (size_t a = vectorized_axis; a < axis_size; a++) { - size_t idx = outer * axis_size * inner_size + a * inner_size + inner; - total_sum += input[idx]; - } - - size_t output_idx = outer * inner_size + inner; - output[output_idx] = total_sum / static_cast<__fp16>(axis_size); - }); -} - -struct VarianceState { - double sum; - double sum_sq; - - VarianceState() : sum(0.0), sum_sq(0.0) {} - VarianceState(double s, double sq) : sum(s), sum_sq(sq) {} -}; - -double cactus_variance_all_f16(const __fp16* data, size_t num_elements) { - VarianceState result = CactusThreading::parallel_reduce( - num_elements, CactusThreading::Thresholds::ALL_REDUCE, - [&](size_t start_idx, size_t end_idx) -> VarianceState { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - float32x4_t sum_vec_lo = vdupq_n_f32(0.0f); - float32x4_t sum_vec_hi = vdupq_n_f32(0.0f); - float32x4_t sum_sq_vec_lo = vdupq_n_f32(0.0f); - float32x4_t sum_sq_vec_hi = vdupq_n_f32(0.0f); - - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t x = vld1q_f16(&data[i]); - float32x4_t x_lo = vcvt_f32_f16(vget_low_f16(x)); - float32x4_t x_hi = vcvt_f32_f16(vget_high_f16(x)); - - sum_vec_lo = vaddq_f32(sum_vec_lo, x_lo); - sum_vec_hi = vaddq_f32(sum_vec_hi, x_hi); - sum_sq_vec_lo = vfmaq_f32(sum_sq_vec_lo, x_lo, x_lo); - sum_sq_vec_hi = vfmaq_f32(sum_sq_vec_hi, x_hi, x_hi); - } - - // Reduce vectors to scalars - double sum = static_cast(vaddvq_f32(vaddq_f32(sum_vec_lo, sum_vec_hi))); - double sum_sq = static_cast(vaddvq_f32(vaddq_f32(sum_sq_vec_lo, sum_sq_vec_hi))); - - for (size_t i = vectorized_end; i < end_idx; ++i) { - double x = static_cast(data[i]); - sum += x; - sum_sq += x * x; - } - - return VarianceState(sum, sum_sq); - }, - VarianceState(), - [](const VarianceState& a, const VarianceState& b) { - return VarianceState(a.sum + b.sum, a.sum_sq + b.sum_sq); - } - ); - - double mean = result.sum / static_cast(num_elements); - double mean_sq = result.sum_sq / static_cast(num_elements); - return mean_sq - mean * mean; -} - -void cactus_variance_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { - - CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, - [&](size_t outer, size_t inner) { - float sum = 0.0f; - float sum_sq = 0.0f; - - for (size_t a = 0; a < axis_size; a++) { - size_t idx = outer * axis_size * inner_size + a * inner_size + inner; - float x = static_cast(input[idx]); - sum += x; - sum_sq += x * x; - } - - float mean = sum / static_cast(axis_size); - float mean_sq = sum_sq / static_cast(axis_size); - size_t output_idx = outer * inner_size + inner; - output[output_idx] = static_cast<__fp16>(mean_sq - mean * mean); - }); -} - -__fp16 cactus_min_all_f16(const __fp16* data, size_t num_elements) { - return CactusThreading::parallel_reduce( - num_elements, CactusThreading::Thresholds::ALL_REDUCE, - [&](size_t start_idx, size_t end_idx) -> __fp16 { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - float16x8_t min_vec = vdupq_n_f16(static_cast<__fp16>(65504.0f)); - - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t input_vec = vld1q_f16(&data[i]); - min_vec = vminq_f16(min_vec, input_vec); - } - - __fp16 thread_min = static_cast<__fp16>(65504.0f); - __fp16 min_array[8]; - vst1q_f16(min_array, min_vec); - for (int j = 0; j < 8; j++) { - thread_min = std::min(thread_min, min_array[j]); - } - - for (size_t i = vectorized_end; i < end_idx; ++i) { - thread_min = std::min(thread_min, data[i]); - } - - return thread_min; - }, - static_cast<__fp16>(65504.0f), - [](__fp16 a, __fp16 b) { return std::min(a, b); } - ); -} - -void cactus_min_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { - CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, - [&](size_t outer, size_t inner) { - float16x8_t min_vec = vdupq_n_f16(static_cast<__fp16>(65504.0f)); - - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_axis = (axis_size / SIMD_WIDTH) * SIMD_WIDTH; - - for (size_t a = 0; a < vectorized_axis; a += SIMD_WIDTH) { - __fp16 values[SIMD_WIDTH]; - for (size_t j = 0; j < SIMD_WIDTH; j++) { - size_t idx = outer * axis_size * inner_size + (a + j) * inner_size + inner; - values[j] = input[idx]; - } - float16x8_t input_vec = vld1q_f16(values); - min_vec = vminq_f16(min_vec, input_vec); - } - - __fp16 min_val = static_cast<__fp16>(65504.0f); - __fp16 min_array[8]; - vst1q_f16(min_array, min_vec); - for (int j = 0; j < 8; j++) { - min_val = std::min(min_val, min_array[j]); - } - - for (size_t a = vectorized_axis; a < axis_size; a++) { - size_t idx = outer * axis_size * inner_size + a * inner_size + inner; - min_val = std::min(min_val, input[idx]); - } - - size_t output_idx = outer * inner_size + inner; - output[output_idx] = min_val; - }); -} - -__fp16 cactus_max_all_f16(const __fp16* data, size_t num_elements) { - return CactusThreading::parallel_reduce( - num_elements, CactusThreading::Thresholds::ALL_REDUCE, - [&](size_t start_idx, size_t end_idx) -> __fp16 { - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - float16x8_t max_vec = vdupq_n_f16(static_cast<__fp16>(-65504.0f)); - - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t input_vec = vld1q_f16(&data[i]); - max_vec = vmaxq_f16(max_vec, input_vec); - } - - __fp16 thread_max = static_cast<__fp16>(-65504.0f); - __fp16 max_array[8]; - vst1q_f16(max_array, max_vec); - for (int j = 0; j < 8; j++) { - thread_max = std::max(thread_max, max_array[j]); - } - - for (size_t i = vectorized_end; i < end_idx; ++i) { - thread_max = std::max(thread_max, data[i]); - } - - return thread_max; - }, - static_cast<__fp16>(-65504.0f), - [](__fp16 a, __fp16 b) { return std::max(a, b); } - ); -} - -void cactus_max_axis_f16(const __fp16* input, __fp16* output, size_t outer_size, size_t axis_size, size_t inner_size) { - CactusThreading::parallel_for_2d(outer_size, inner_size, CactusThreading::Thresholds::AXIS_REDUCE, - [&](size_t outer, size_t inner) { - float16x8_t max_vec = vdupq_n_f16(static_cast<__fp16>(-65504.0f)); - - constexpr size_t SIMD_WIDTH = 8; - const size_t vectorized_axis = (axis_size / SIMD_WIDTH) * SIMD_WIDTH; - - for (size_t a = 0; a < vectorized_axis; a += SIMD_WIDTH) { - __fp16 values[SIMD_WIDTH]; - for (size_t j = 0; j < SIMD_WIDTH; j++) { - size_t idx = outer * axis_size * inner_size + (a + j) * inner_size + inner; - values[j] = input[idx]; - } - float16x8_t input_vec = vld1q_f16(values); - max_vec = vmaxq_f16(max_vec, input_vec); - } - - __fp16 max_val = static_cast<__fp16>(-65504.0f); - __fp16 max_array[8]; - vst1q_f16(max_array, max_vec); - for (int j = 0; j < 8; j++) { - max_val = std::max(max_val, max_array[j]); - } - - for (size_t a = vectorized_axis; a < axis_size; a++) { - size_t idx = outer * axis_size * inner_size + a * inner_size + inner; - max_val = std::max(max_val, input[idx]); - } - - size_t output_idx = outer * inner_size + inner; - output[output_idx] = max_val; - }); -} diff --git a/cactus/kernel/kernel_scalar.cpp b/cactus/kernel/kernel_scalar.cpp deleted file mode 100644 index 2ad9d4d9e..000000000 --- a/cactus/kernel/kernel_scalar.cpp +++ /dev/null @@ -1,264 +0,0 @@ -#include "kernel.h" -#include "kernel_utils.h" -#include -#include -#include - -void cactus_scalar_op_f16(const __fp16* input, __fp16* output, size_t num_elements, float scalar_value, ScalarOpType op_type) { - const __fp16 scalar_f16 = static_cast<__fp16>(scalar_value); - const bool use_streaming = num_elements >= STREAMING_STORE_THRESHOLD; - - switch (op_type) { - case ScalarOpType::ADD: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_BASIC, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 4; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - const float16x8_t scalar_vec = vdupq_n_f16(scalar_f16); - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&input[i + 256], 0, 0); - float16x8_t v0 = vaddq_f16(vld1q_f16(&input[i]), scalar_vec); - float16x8_t v1 = vaddq_f16(vld1q_f16(&input[i + 8]), scalar_vec); - float16x8_t v2 = vaddq_f16(vld1q_f16(&input[i + 16]), scalar_vec); - float16x8_t v3 = vaddq_f16(vld1q_f16(&input[i + 24]), scalar_vec); - stream_store_f16x8(&output[i], v0); - stream_store_f16x8(&output[i + 8], v1); - stream_store_f16x8(&output[i + 16], v2); - stream_store_f16x8(&output[i + 24], v3); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - stream_store_f16x8(&output[i], vaddq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - vst1q_f16(&output[i], vaddq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = input[i] + scalar_f16; - } - }); - break; - } - - case ScalarOpType::SUBTRACT: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_BASIC, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 4; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - const float16x8_t scalar_vec = vdupq_n_f16(scalar_f16); - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&input[i + 256], 0, 0); - float16x8_t v0 = vsubq_f16(vld1q_f16(&input[i]), scalar_vec); - float16x8_t v1 = vsubq_f16(vld1q_f16(&input[i + 8]), scalar_vec); - float16x8_t v2 = vsubq_f16(vld1q_f16(&input[i + 16]), scalar_vec); - float16x8_t v3 = vsubq_f16(vld1q_f16(&input[i + 24]), scalar_vec); - stream_store_f16x8(&output[i], v0); - stream_store_f16x8(&output[i + 8], v1); - stream_store_f16x8(&output[i + 16], v2); - stream_store_f16x8(&output[i + 24], v3); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - stream_store_f16x8(&output[i], vsubq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - vst1q_f16(&output[i], vsubq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = input[i] - scalar_f16; - } - }); - break; - } - - case ScalarOpType::MULTIPLY: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_BASIC, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 4; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - const float16x8_t scalar_vec = vdupq_n_f16(scalar_f16); - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&input[i + 256], 0, 0); - float16x8_t v0 = vmulq_f16(vld1q_f16(&input[i]), scalar_vec); - float16x8_t v1 = vmulq_f16(vld1q_f16(&input[i + 8]), scalar_vec); - float16x8_t v2 = vmulq_f16(vld1q_f16(&input[i + 16]), scalar_vec); - float16x8_t v3 = vmulq_f16(vld1q_f16(&input[i + 24]), scalar_vec); - stream_store_f16x8(&output[i], v0); - stream_store_f16x8(&output[i + 8], v1); - stream_store_f16x8(&output[i + 16], v2); - stream_store_f16x8(&output[i + 24], v3); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - stream_store_f16x8(&output[i], vmulq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - vst1q_f16(&output[i], vmulq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = input[i] * scalar_f16; - } - }); - break; - } - - case ScalarOpType::DIVIDE: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_BASIC, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 4; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - const float16x8_t scalar_vec = vdupq_n_f16(scalar_f16); - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&input[i + 256], 0, 0); - float16x8_t v0 = vdivq_f16(vld1q_f16(&input[i]), scalar_vec); - float16x8_t v1 = vdivq_f16(vld1q_f16(&input[i + 8]), scalar_vec); - float16x8_t v2 = vdivq_f16(vld1q_f16(&input[i + 16]), scalar_vec); - float16x8_t v3 = vdivq_f16(vld1q_f16(&input[i + 24]), scalar_vec); - stream_store_f16x8(&output[i], v0); - stream_store_f16x8(&output[i + 8], v1); - stream_store_f16x8(&output[i + 16], v2); - stream_store_f16x8(&output[i + 24], v3); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - stream_store_f16x8(&output[i], vdivq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - vst1q_f16(&output[i], vdivq_f16(vld1q_f16(&input[i]), scalar_vec)); - } - } - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = input[i] / scalar_f16; - } - }); - break; - } - - case ScalarOpType::EXP: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_EXPENSIVE, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 2; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&input[i + 128], 0, 0); - float16x8_t in0 = vld1q_f16(&input[i]); - float16x8_t in1 = vld1q_f16(&input[i + 8]); - float16x8_t r0 = vcombine_f16( - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_low_f16(in0)))), - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_high_f16(in0))))); - float16x8_t r1 = vcombine_f16( - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_low_f16(in1)))), - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_high_f16(in1))))); - stream_store_f16x8(&output[i], r0); - stream_store_f16x8(&output[i + 8], r1); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t in_vec = vld1q_f16(&input[i]); - float16x8_t result = vcombine_f16( - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_low_f16(in_vec)))), - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_high_f16(in_vec))))); - stream_store_f16x8(&output[i], result); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t in_vec = vld1q_f16(&input[i]); - float16x8_t result = vcombine_f16( - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_low_f16(in_vec)))), - vcvt_f16_f32(fast_exp_f32x4(vcvt_f32_f16(vget_high_f16(in_vec))))); - vst1q_f16(&output[i], result); - } - } - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = static_cast<__fp16>(std::exp(static_cast(input[i]))); - } - }); - break; - } - - case ScalarOpType::SQRT: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_EXPENSIVE, - [&](size_t start_idx, size_t end_idx) { - constexpr size_t SIMD_WIDTH = 8; - constexpr size_t UNROLL = 2; - const size_t unrolled_end = start_idx + ((end_idx - start_idx) / (SIMD_WIDTH * UNROLL)) * (SIMD_WIDTH * UNROLL); - const size_t vectorized_end = start_idx + ((end_idx - start_idx) / SIMD_WIDTH) * SIMD_WIDTH; - - if (use_streaming) { - for (size_t i = start_idx; i < unrolled_end; i += SIMD_WIDTH * UNROLL) { - __builtin_prefetch(&input[i + 128], 0, 0); - float16x8_t in0 = vld1q_f16(&input[i]); - float16x8_t in1 = vld1q_f16(&input[i + 8]); - float16x8_t r0 = vcombine_f16( - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_low_f16(in0)))), - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_high_f16(in0))))); - float16x8_t r1 = vcombine_f16( - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_low_f16(in1)))), - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_high_f16(in1))))); - stream_store_f16x8(&output[i], r0); - stream_store_f16x8(&output[i + 8], r1); - } - for (size_t i = unrolled_end; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t in_vec = vld1q_f16(&input[i]); - float16x8_t result = vcombine_f16( - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_low_f16(in_vec)))), - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_high_f16(in_vec))))); - stream_store_f16x8(&output[i], result); - } - } else { - for (size_t i = start_idx; i < vectorized_end; i += SIMD_WIDTH) { - float16x8_t in_vec = vld1q_f16(&input[i]); - float16x8_t result = vcombine_f16( - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_low_f16(in_vec)))), - vcvt_f16_f32(vsqrtq_f32(vcvt_f32_f16(vget_high_f16(in_vec))))); - vst1q_f16(&output[i], result); - } - } - for (size_t i = vectorized_end; i < end_idx; ++i) { - output[i] = static_cast<__fp16>(std::sqrt(static_cast(input[i]))); - } - }); - break; - } - - case ScalarOpType::COS: - case ScalarOpType::SIN: { - CactusThreading::parallel_for(num_elements, CactusThreading::Thresholds::SCALAR_EXPENSIVE, - [&](size_t start_idx, size_t end_idx) { - for (size_t i = start_idx; i < end_idx; ++i) { - float val = static_cast(input[i]); - float result; - switch (op_type) { - case ScalarOpType::COS: result = std::cos(val); break; - case ScalarOpType::SIN: result = std::sin(val); break; - default: result = val; break; - } - output[i] = static_cast<__fp16>(result); - } - }); - break; - } - } -} \ No newline at end of file diff --git a/cactus/kernel/kernel_utils.h b/cactus/kernel/kernel_utils.h deleted file mode 100644 index 5628006cb..000000000 --- a/cactus/kernel/kernel_utils.h +++ /dev/null @@ -1,610 +0,0 @@ -#ifndef KERNEL_UTILS_H -#define KERNEL_UTILS_H - -#include -#if defined(__APPLE__) -#include -#include -#include -#endif -#if defined(__ANDROID__) -#include -#include -#endif -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -constexpr size_t NEON_VECTOR_SIZE = 16; -constexpr size_t STREAMING_STORE_THRESHOLD = 32768; - -inline void stream_store_f16x8(__fp16* dst, float16x8_t val) { -#if defined(__aarch64__) - float16x4_t lo = vget_low_f16(val); - float16x4_t hi = vget_high_f16(val); - __asm__ __volatile__( - "stnp %d0, %d1, [%2]" - : - : "w"(lo), "w"(hi), "r"(dst) - : "memory" - ); -#else - vst1q_f16(dst, val); -#endif -} - -#if defined(__ARM_FEATURE_DOTPROD) -inline int32x4_t accum_dot(int32x4_t acc, int8x16_t a, int8x16_t b) { - return vdotq_s32(acc, a, b); -} -#else -inline int32x4_t accum_dot(int32x4_t acc, int8x16_t a, int8x16_t b) { - int16x8_t prod_low = vmull_s8(vget_low_s8(a), vget_low_s8(b)); - int32x4_t acc_high = vpaddlq_s16(vmull_s8(vget_high_s8(a), vget_high_s8(b))); - return vaddq_s32(vaddq_s32(acc, vpaddlq_s16(prod_low)), acc_high); -} -#endif - -// I8MM support: runtime detection on Android, compile-time on Apple -#if defined(__ANDROID__) && defined(__aarch64__) - -inline bool cactus_has_i8mm() { - static int8_t supported = -1; - if (supported == -1) { - unsigned long hwcaps = getauxval(AT_HWCAP2); - supported = (hwcaps & HWCAP2_I8MM) ? 1 : 0; - } - return supported; -} - -__attribute__((target("arch=armv8.2-a+i8mm"))) -inline int32x4_t accum_matmul(int32x4_t acc, int8x16_t a, int8x16_t b) { - return vmmlaq_s32(acc, a, b); -} - -#elif defined(__APPLE__) && defined(__aarch64__) - -inline bool cactus_has_i8mm() { - static int8_t supported = -1; - if (supported == -1) { - int v = 0; - size_t sz = sizeof(v); - supported = sysctlbyname("hw.optional.arm.FEAT_I8MM", &v, &sz, nullptr, 0) == 0 && v; - } - return supported; -} - -__attribute__((target("i8mm"))) -inline int32x4_t accum_matmul(int32x4_t acc, int8x16_t a, int8x16_t b) { - return vmmlaq_s32(acc, a, b); -} - -#else - -inline bool cactus_has_i8mm() { - return false; -} - -#endif - -inline float16x8_t accum_f16_dot(float16x8_t acc, float16x8_t a_low, float16x8_t a_high, - float16x8_t b_low, float16x8_t b_high) { - acc = vfmaq_f16(acc, a_low, b_low); - return vfmaq_f16(acc, a_high, b_high); -} - -inline float32x4_t fast_exp_f32x4(float32x4_t x) { - const float32x4_t log2e = vdupq_n_f32(1.4426950408889634f); - const float32x4_t ln2 = vdupq_n_f32(0.6931471805599453f); - - const float32x4_t c0 = vdupq_n_f32(1.0f); - const float32x4_t c1 = vdupq_n_f32(0.6931471805599453f); - const float32x4_t c2 = vdupq_n_f32(0.2402265069591007f); - const float32x4_t c3 = vdupq_n_f32(0.05550410866482158f); - const float32x4_t c4 = vdupq_n_f32(0.009618129842071803f); - - x = vmaxq_f32(x, vdupq_n_f32(-87.0f)); - x = vminq_f32(x, vdupq_n_f32(87.0f)); - - float32x4_t z = vmulq_f32(x, log2e); - - int32x4_t zi = vcvtq_s32_f32(z); - float32x4_t zf = vsubq_f32(z, vcvtq_f32_s32(zi)); - - uint32x4_t neg_mask = vcltq_f32(zf, vdupq_n_f32(0.0f)); - zi = vsubq_s32(zi, vandq_s32(vreinterpretq_s32_u32(neg_mask), vdupq_n_s32(1))); - zf = vaddq_f32(zf, vreinterpretq_f32_u32(vandq_u32(neg_mask, vreinterpretq_u32_f32(vdupq_n_f32(1.0f))))); - - float32x4_t zf_ln2 = vmulq_f32(zf, ln2); - float32x4_t p = c4; - p = vfmaq_f32(c3, p, zf_ln2); - p = vfmaq_f32(c2, p, zf_ln2); - p = vfmaq_f32(c1, p, zf_ln2); - p = vfmaq_f32(c0, p, zf_ln2); - - int32x4_t exp_bits = vshlq_n_s32(vaddq_s32(zi, vdupq_n_s32(127)), 23); - float32x4_t scale = vreinterpretq_f32_s32(exp_bits); - - return vmulq_f32(p, scale); -} - -inline float32x4_t fast_tanh_f32x4(float32x4_t x) { - const float32x4_t one = vdupq_n_f32(1.0f); - const float32x4_t neg_one = vdupq_n_f32(-1.0f); - - uint32x4_t pos_sat = vcgtq_f32(x, vdupq_n_f32(4.5f)); - uint32x4_t neg_sat = vcltq_f32(x, vdupq_n_f32(-4.5f)); - - const float32x4_t c27 = vdupq_n_f32(27.0f); - const float32x4_t c9 = vdupq_n_f32(9.0f); - - float32x4_t x2 = vmulq_f32(x, x); - float32x4_t num = vaddq_f32(c27, x2); - float32x4_t den = vfmaq_f32(c27, c9, x2); - - float32x4_t result = vmulq_f32(x, vdivq_f32(num, den)); - - result = vbslq_f32(pos_sat, one, result); - result = vbslq_f32(neg_sat, neg_one, result); - - return result; -} - -inline int8x16_t unpack_int4_lo(uint8x16_t packed) { - uint8x16_t lo = vandq_u8(packed, vdupq_n_u8(0x0F)); - uint8x16_t sign_mask = vcgtq_u8(lo, vdupq_n_u8(7)); - uint8x16_t correction = vandq_u8(sign_mask, vdupq_n_u8(16)); - return vreinterpretq_s8_u8(vsubq_u8(lo, correction)); -} - -inline int8x16_t unpack_int4_hi(uint8x16_t packed) { - uint8x16_t hi = vshrq_n_u8(packed, 4); - uint8x16_t sign_mask = vcgtq_u8(hi, vdupq_n_u8(7)); - uint8x16_t correction = vandq_u8(sign_mask, vdupq_n_u8(16)); - return vreinterpretq_s8_u8(vsubq_u8(hi, correction)); -} - -inline void unpack_int4_to_int8x32(uint8x16_t packed, int8x16_t& out_lo, int8x16_t& out_hi) { - int8x16_t lo_nibbles = unpack_int4_lo(packed); - int8x16_t hi_nibbles = unpack_int4_hi(packed); - int8x16x2_t interleaved = vzipq_s8(lo_nibbles, hi_nibbles); - out_lo = interleaved.val[0]; - out_hi = interleaved.val[1]; -} - -inline int32x4_t int4_dot_asm(int32x4_t acc, uint8x16_t packed, int8x16_t a_lo, int8x16_t a_hi) { -#if defined(__aarch64__) - int8x16_t b_lo, b_hi; - - __asm__ __volatile__ ( - "movi v16.16b, #0x0F \n" // low nibble mask - "movi v17.16b, #7 \n" // sign threshold - "movi v18.16b, #16 \n" // sign correction - - "and %[b_lo].16b, %[packed].16b, v16.16b \n" - - "ushr %[b_hi].16b, %[packed].16b, #4 \n" - - "cmgt v19.16b, %[b_lo].16b, v17.16b \n" - "and v19.16b, v19.16b, v18.16b \n" - "sub %[b_lo].16b, %[b_lo].16b, v19.16b \n" - - "cmgt v20.16b, %[b_hi].16b, v17.16b \n" - "and v20.16b, v20.16b, v18.16b \n" - "sub %[b_hi].16b, %[b_hi].16b, v20.16b \n" - - "zip1 v21.16b, %[b_lo].16b, %[b_hi].16b \n" - "zip2 v22.16b, %[b_lo].16b, %[b_hi].16b \n" - - ".arch armv8.2-a+dotprod \n" - "sdot %[acc].4s, %[a_lo].16b, v21.16b \n" - "sdot %[acc].4s, %[a_hi].16b, v22.16b \n" - - : [acc] "+w"(acc), [b_lo] "=w"(b_lo), [b_hi] "=w"(b_hi) - : [packed] "w"(packed), [a_lo] "w"(a_lo), [a_hi] "w"(a_hi) - : "v16", "v17", "v18", "v19", "v20", "v21", "v22" - ); - - return acc; -#else - int8x16_t b_lo, b_hi; - unpack_int4_to_int8x32(packed, b_lo, b_hi); - acc = accum_dot(acc, a_lo, b_lo); - acc = accum_dot(acc, a_hi, b_hi); - return acc; -#endif -} - -inline int32_t int4_dot_m1_asm(const int8_t* a_ptr, const uint8_t* b_packed, size_t group_size) { -#if defined(__aarch64__) - int32x4_t acc = vdupq_n_s32(0); - - for (size_t k = 0; k < group_size; k += 64) { - uint8x16_t p0 = vld1q_u8(b_packed + k/2); - uint8x16_t p1 = vld1q_u8(b_packed + k/2 + 16); - - int8x16_t a0 = vld1q_s8(a_ptr + k); - int8x16_t a1 = vld1q_s8(a_ptr + k + 16); - int8x16_t a2 = vld1q_s8(a_ptr + k + 32); - int8x16_t a3 = vld1q_s8(a_ptr + k + 48); - - acc = int4_dot_asm(acc, p0, a0, a1); - acc = int4_dot_asm(acc, p1, a2, a3); - } - - return vaddvq_s32(acc); -#else - int32x4_t acc = vdupq_n_s32(0); - for (size_t k = 0; k < group_size; k += 32) { - uint8x16_t packed = vld1q_u8(b_packed + k/2); - int8x16_t b_lo, b_hi; - unpack_int4_to_int8x32(packed, b_lo, b_hi); - acc = accum_dot(acc, vld1q_s8(a_ptr + k), b_lo); - acc = accum_dot(acc, vld1q_s8(a_ptr + k + 16), b_hi); - } - return vaddvq_s32(acc); -#endif -} - -namespace CactusThreading { - - class ThreadPool { - private: - static constexpr size_t MAX_WORKERS = 16; - - std::vector workers; - std::deque> tasks; - - std::mutex mutex; - std::condition_variable work_available; - std::condition_variable work_done; - - bool stop{false}; - std::atomic pending_tasks{0}; - size_t num_workers_; - - void worker_thread() { - while (true) { - std::function task; - { - std::unique_lock lock(mutex); - work_available.wait(lock, [this] { - return stop || !tasks.empty(); - }); - - if (stop && tasks.empty()) { - return; - } - - task = std::move(tasks.front()); - tasks.pop_front(); - } - - task(); - - if (pending_tasks.fetch_sub(1, std::memory_order_acq_rel) == 1) { - std::lock_guard lock(mutex); - work_done.notify_one(); - } - } - } - - public: - explicit ThreadPool(size_t num_threads = std::thread::hardware_concurrency()) - : stop(false), pending_tasks(0) { - num_workers_ = std::min(num_threads, MAX_WORKERS); - if (num_workers_ == 0) num_workers_ = 1; - workers.reserve(num_workers_); - for (size_t i = 0; i < num_workers_; ++i) { - workers.emplace_back(&ThreadPool::worker_thread, this); - } - } - - ~ThreadPool() { - { - std::lock_guard lock(mutex); - stop = true; - } - work_available.notify_all(); - for (auto& worker : workers) { - if (worker.joinable()) { - worker.join(); - } - } - } - - template - auto enqueue(F&& f) -> std::future { - using return_type = decltype(f()); - - auto task = std::make_shared>( - std::forward(f) - ); - - std::future res = task->get_future(); - - { - std::lock_guard lock(mutex); - pending_tasks.fetch_add(1, std::memory_order_relaxed); - tasks.emplace_back([task](){ (*task)(); }); - } - work_available.notify_one(); - - return res; - } - - template - void enqueue_batch(size_t total_work, F task_func) { - if (total_work == 0) return; - - const size_t num_tasks = std::min(num_workers_, total_work); - const size_t per_worker = total_work / num_tasks; - const size_t remainder = total_work % num_tasks; - - { - std::lock_guard lock(mutex); - pending_tasks.fetch_add(num_tasks, std::memory_order_relaxed); - - for (size_t w = 0; w < num_tasks; ++w) { - size_t start = w * per_worker + std::min(w, remainder); - size_t end = start + per_worker + (w < remainder ? 1 : 0); - tasks.emplace_back([=]() { task_func(start, end); }); - } - } - work_available.notify_all(); - } - - void wait_all() { - std::unique_lock lock(mutex); - work_done.wait(lock, [this] { - return pending_tasks.load(std::memory_order_acquire) == 0; - }); - } - - template - void enqueue_n_threads(size_t total_work, size_t num_threads, F task_func) { - if (total_work == 0 || num_threads == 0) return; - - num_threads = std::min(num_threads, std::min(num_workers_, total_work)); - const size_t per_thread = total_work / num_threads; - const size_t remainder = total_work % num_threads; - - { - std::lock_guard lock(mutex); - pending_tasks.fetch_add(num_threads, std::memory_order_relaxed); - - for (size_t t = 0; t < num_threads; ++t) { - size_t start = t * per_thread + std::min(t, remainder); - size_t end = start + per_thread + (t < remainder ? 1 : 0); - tasks.emplace_back([=]() { task_func(start, end); }); - } - } - work_available.notify_all(); - } - - size_t num_workers() const { return num_workers_; } - }; - - inline ThreadPool& get_thread_pool() { - static ThreadPool pool; - return pool; - } - - struct ParallelConfig { - size_t min_work_gate; - size_t work_per_thread; - - constexpr ParallelConfig(size_t gate, size_t per_thread) - : min_work_gate(gate), work_per_thread(per_thread) {} - }; - - inline size_t get_optimal_thread_count(size_t total_work, ParallelConfig config) { - if (total_work < config.min_work_gate) return 1; - - size_t pool_size = get_thread_pool().num_workers(); - size_t num_threads = (total_work + config.work_per_thread - 1) / config.work_per_thread; - return std::min(pool_size, std::max(static_cast(1), num_threads)); - } - - struct Thresholds { - #if defined(__ANDROID__) - static constexpr ParallelConfig ATTENTION{64, 32}; - static constexpr ParallelConfig ELEMENT_WISE{5000, 2500}; - static constexpr ParallelConfig AXIS_REDUCE{1000, 500}; - static constexpr ParallelConfig ALL_REDUCE{10000, 5000}; - static constexpr ParallelConfig SCALAR_BASIC{30000, 15000}; - static constexpr ParallelConfig SCALAR_EXPENSIVE{10000, 5000}; - #else // Apple - static constexpr ParallelConfig ATTENTION{32, 16}; - static constexpr ParallelConfig ELEMENT_WISE{5000, 2500}; - static constexpr ParallelConfig AXIS_REDUCE{1000, 500}; - static constexpr ParallelConfig ALL_REDUCE{10000, 5000}; - static constexpr ParallelConfig SCALAR_BASIC{5000, 2500}; - static constexpr ParallelConfig SCALAR_EXPENSIVE{2500, 1250}; - #endif - }; - - struct GemmThreading { - #if defined(__ANDROID__) - static size_t get_num_threads(size_t M, size_t pool_size) { - if (M <= 1) return 1; - return pool_size; - } - #elif defined(__APPLE__) && TARGET_OS_IPHONE - static size_t get_num_threads(size_t M, size_t pool_size) { - if (M <= 1) return std::min(pool_size, static_cast(2)); - return pool_size; - } - #else // Mac - static size_t get_num_threads(size_t M, size_t pool_size) { - if (M <= 1) return std::min(pool_size, static_cast(4)); - return pool_size; - } - #endif - }; - - inline size_t& get_gemm_thread_override() { - static size_t override_threads = 0; - return override_threads; - } - - inline void set_gemm_threads(size_t num_threads) { - get_gemm_thread_override() = num_threads; - } - - inline void reset_gemm_threads() { - get_gemm_thread_override() = 0; - } - - class TaskHandle { - private: - std::vector> futures_; - bool auto_wait_; - - public: - TaskHandle(bool auto_wait = true) : auto_wait_(auto_wait) {} - - ~TaskHandle() { - if (auto_wait_) { - wait(); - } - } - - TaskHandle(TaskHandle&&) = default; - TaskHandle& operator=(TaskHandle&&) = default; - TaskHandle(const TaskHandle&) = delete; - TaskHandle& operator=(const TaskHandle&) = delete; - - void add_future(std::future&& f) { - futures_.push_back(std::move(f)); - } - - void wait() { - for (auto& f : futures_) { - if (f.valid()) { - f.wait(); - } - } - futures_.clear(); - } - - bool is_ready() const { - for (const auto& f : futures_) { - if (f.valid() && f.wait_for(std::chrono::seconds(0)) != std::future_status::ready) { - return false; - } - } - return true; - } - - size_t task_count() const { return futures_.size(); } - }; - - template - TaskHandle parallel_for(size_t total_work, ParallelConfig config, WorkFunc work_func, bool wait = true) { - const size_t num_threads = get_optimal_thread_count(total_work, config); - TaskHandle handle(!wait); - - if (num_threads == 1) { - if (wait) { - work_func(0, total_work); - return handle; - } - auto& pool = get_thread_pool(); - handle.add_future(pool.enqueue([work_func, total_work]() { - work_func(0, total_work); - })); - return handle; - } - - auto& pool = get_thread_pool(); - const size_t work_per_thread = total_work / num_threads; - - for (size_t t = 0; t < num_threads; ++t) { - handle.add_future(pool.enqueue([work_func, t, num_threads, work_per_thread, total_work]() { - const size_t start_idx = t * work_per_thread; - const size_t end_idx = (t == num_threads - 1) ? total_work : (t + 1) * work_per_thread; - work_func(start_idx, end_idx); - })); - } - - if (wait) { - handle.wait(); - } - return handle; - } - - template - void parallel_for_2d(size_t outer_size, size_t inner_size, ParallelConfig config, WorkFunc work_func) { - const size_t total_work = outer_size * inner_size; - parallel_for(total_work, config, [&](size_t start_idx, size_t end_idx) { - for (size_t work_idx = start_idx; work_idx < end_idx; ++work_idx) { - const size_t outer = work_idx / inner_size; - const size_t inner = work_idx % inner_size; - work_func(outer, inner); - } - }); - } - - template - ResultType parallel_reduce(size_t total_work, ParallelConfig config, - WorkFunc work_func, ResultType init_value, CombineFunc combine_func) { - const size_t num_threads = get_optimal_thread_count(total_work, config); - - if (num_threads == 1) { - return work_func(0, total_work); - } - - auto& pool = get_thread_pool(); - std::vector> futures; - std::vector partial_results(num_threads, init_value); - const size_t work_per_thread = total_work / num_threads; - - for (size_t t = 0; t < num_threads; ++t) { - futures.push_back(pool.enqueue([work_func, t, num_threads, work_per_thread, total_work]() -> ResultType { - const size_t start_idx = t * work_per_thread; - const size_t end_idx = (t == num_threads - 1) ? total_work : (t + 1) * work_per_thread; - return work_func(start_idx, end_idx); - })); - } - - ResultType result = init_value; - for (auto& future : futures) { - result = combine_func(result, future.get()); - } - return result; - } - - template - void parallel_gemm_tiles(size_t M, size_t total_tiles, WorkFunc work_func) { - auto& pool = get_thread_pool(); - - size_t override = get_gemm_thread_override(); - size_t num_threads = (override > 0) ? override : GemmThreading::get_num_threads(M, pool.num_workers()); - num_threads = std::min(num_threads, total_tiles); - - if (num_threads <= 1) { - work_func(0, total_tiles); - return; - } - - pool.enqueue_n_threads(total_tiles, num_threads, work_func); - pool.wait_all(); - } - -} - - -#endif // KERNEL_UTILS_H \ No newline at end of file diff --git a/cactus/models/model.h b/cactus/models/model.h deleted file mode 100644 index 8ce34d148..000000000 --- a/cactus/models/model.h +++ /dev/null @@ -1,732 +0,0 @@ -#pragma once - -#include "../engine/engine.h" -#include "../npu/npu.h" - -namespace cactus { -namespace engine { - - - -class QwenModel : public Model { -public: - QwenModel(); - explicit QwenModel(const Config& config); - ~QwenModel() override = default; - -protected: - size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const override; - - size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t forward(const std::vector& tokens, bool use_cache = false) override; - void load_weights_to_graph(CactusGraph* gb) override; - -private: - struct WeightNodeIDs { - size_t output_weight; - size_t output_norm_weight; - - struct LayerWeights { - size_t attn_q_weight; - size_t attn_k_weight; - size_t attn_v_weight; - size_t attn_output_weight; - size_t input_layernorm_weight; - size_t attn_q_norm_weight; - size_t attn_k_norm_weight; - size_t pre_feedforward_layernorm_weight; - size_t post_feedforward_layernorm_weight; - size_t ffn_gate_weight; - size_t ffn_up_weight; - size_t ffn_down_weight; - size_t post_attention_layernorm_weight; - }; - - std::vector layers; - } weight_nodes_; -}; - - - -class GemmaModel : public Model { -public: - GemmaModel(); - explicit GemmaModel(const Config& config); - ~GemmaModel() override = default; - -protected: - size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const override; - - size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t forward(const std::vector& tokens, bool use_cache = false) override; - void load_weights_to_graph(CactusGraph* gb) override; - void post_init() override; - -private: - struct WeightNodeIDs { - size_t output_weight; - size_t output_norm_weight; - - struct LayerWeights { - size_t attn_q_weight; - size_t attn_k_weight; - size_t attn_v_weight; - size_t attn_output_weight; - size_t input_layernorm_weight; - size_t attn_q_norm_weight; - size_t attn_k_norm_weight; - size_t pre_feedforward_layernorm_weight; - size_t post_feedforward_layernorm_weight; - size_t ffn_gate_weight; - size_t ffn_up_weight; - size_t ffn_down_weight; - size_t post_attention_layernorm_weight; - }; - - std::vector layers; - } weight_nodes_; -}; -class SmolModel : public Model{ -public: - SmolModel(); - explicit SmolModel(const Config& config); - ~SmolModel() override = default; - -protected: - size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const override; - - size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t forward(const std::vector& tokens, bool use_cache = false) override; - void load_weights_to_graph(CactusGraph* gb) override; - -protected: - struct WeightNodeIDs { - size_t output_weight; - size_t output_norm_weight; - - struct LayerWeights { - size_t attn_q_weight; - size_t attn_k_weight; - size_t attn_v_weight; - size_t attn_output_weight; - size_t input_layernorm_weight; - size_t ffn_gate_weight; - size_t ffn_up_weight; - size_t ffn_down_weight; - size_t post_attention_layernorm_weight; - }; - - std::vector layers; - } weight_nodes_; -}; - - -class Siglip2VisionModel : public Model { - friend class Lfm2VlModel; - -public: - struct VisionEmbeddingResult { - size_t combined_embeddings; - std::vector tile_embeddings; - }; - - Siglip2VisionModel(); - explicit Siglip2VisionModel(const Config& cfg); - ~Siglip2VisionModel() override = default; - virtual size_t forward_vision(const Siglip2Preprocessor::PreprocessedImage& preprocessed_image); - virtual size_t forward_vision(CactusGraph* gb, - const Siglip2Preprocessor::PreprocessedImage& preprocessed_image, - ComputeBackend backend); - std::vector get_image_features(const std::string& image_path); - std::vector get_image_features(const Siglip2Preprocessor::PreprocessedImage& preprocessed_image); - std::vector get_image_embedding(const std::string& image_path); - size_t get_image_features_node(const Siglip2Preprocessor::PreprocessedImage& preprocessed_image); - Siglip2Preprocessor& get_preprocessor() { return preprocessor_; } - const Siglip2Preprocessor& get_preprocessor() const { return preprocessor_; } - -protected: - VisionEmbeddingResult build_vision_embeddings(CactusGraph* gb, - const Siglip2Preprocessor::PreprocessedImage& preprocessed_image, - ComputeBackend backend); - - size_t build_vision_transformer_layer(CactusGraph* gb, size_t hidden_states, uint32_t layer_idx, - ComputeBackend backend); - - size_t build_vision_attention(CactusGraph* gb, size_t hidden_states, uint32_t layer_idx, - ComputeBackend backend); - - size_t build_vision_mlp(CactusGraph* gb, size_t hidden_states, uint32_t layer_idx, - ComputeBackend backend); - - void load_weights_to_graph(CactusGraph* gb) override; - size_t forward(const std::vector& tokens, bool use_cache = false) override; - size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const override; - size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - -protected: - struct VisionWeightNodeIDs { - size_t patch_embedding_weight; - size_t patch_embedding_bias; - size_t position_embedding; - size_t post_layernorm_weight; - size_t post_layernorm_bias; - - struct VisionLayerWeights { - size_t attn_q_weight; - size_t attn_k_weight; - size_t attn_v_weight; - size_t attn_output_weight; - size_t attn_q_bias; - size_t attn_k_bias; - size_t attn_v_bias; - size_t attn_output_bias; - size_t layer_norm1_weight; - size_t layer_norm1_bias; - size_t layer_norm2_weight; - size_t layer_norm2_bias; - size_t mlp_fc1_weight; - size_t mlp_fc1_bias; - size_t mlp_fc2_weight; - size_t mlp_fc2_bias; - }; - - std::vector vision_layers; - } vision_weight_nodes_; - - Siglip2Preprocessor preprocessor_; - - std::unique_ptr npu_encoder_; - bool use_npu_encoder_ = false; -}; - - -class LFM2Model : public Model { - friend class Lfm2VlModel; - -public: - LFM2Model(); - explicit LFM2Model(const Config& config); - ~LFM2Model() override = default; - - bool is_cache_empty() const; - - bool init(const std::string& model_folder, size_t context_size, const std::string& system_prompt = "", bool do_warmup = true) override; - bool init(CactusGraph* external_graph, const std::string& model_folder, size_t context_size, - const std::string& system_prompt = "", bool do_warmup = true) override; - -protected: - using Model::forward; - size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t build_conv1d(CactusGraph* gb, size_t input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache); - - size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const override; - - size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t forward(const std::vector& tokens, bool use_cache = false) override; - size_t forward(CactusGraph* gb, const std::vector& tokens, ComputeBackend backend, bool use_cache = false); - size_t forward(CactusGraph* gb, size_t input_embeddings, size_t seq_len, ComputeBackend backend, bool use_cache = false); - void post_init() override; - void post_execute_updates(CactusGraph* gb, size_t seq_len) override; - void reset_cache() override; - void load_weights_to_graph(CactusGraph* gb) override; - -private: - - struct WeightNodeIDs { - size_t output_weight; - size_t output_norm_weight; - - struct LayerWeights { - size_t attn_q_weight; - size_t attn_k_weight; - size_t attn_v_weight; - size_t attn_output_weight; - size_t attn_q_norm_weight; - size_t attn_k_norm_weight; - - size_t conv_depthwise_weight; - size_t conv_in_proj_weight; - size_t conv_out_proj_weight; - - size_t input_layernorm_weight; - size_t post_attention_layernorm_weight; - size_t ffn_gate_weight; - size_t ffn_up_weight; - size_t ffn_down_weight; - }; - - enum class LayerType : uint8_t { ATTENTION, CONV }; - - struct LayerEntry { - LayerType type; - LayerWeights weights; - }; - - std::vector layers; - } weight_nodes_; - - ConvCache conv_cache_; - std::vector conv_cache_bx_nodes_; - bool last_forward_used_cache_ = false; -}; - - -class NomicModel : public Model { -public: - NomicModel(); - explicit NomicModel(const Config& config); - ~NomicModel() override = default; - -protected: - size_t build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const override; - - size_t build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0) override; - - size_t forward(const std::vector& tokens, bool use_cache = false) override; - - void load_weights_to_graph(CactusGraph* gb) override; - -private: - size_t build_standard_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const; - - size_t build_moe_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const; - - struct WeightNodeIDs { - size_t embedding_layernorm_weight; - size_t embedding_layernorm_bias; - - struct LayerWeights { - size_t attn_q_weight; - size_t attn_k_weight; - size_t attn_v_weight; - size_t attn_q_bias; - size_t attn_k_bias; - size_t attn_v_bias; - size_t attn_output_weight; - size_t attn_output_bias; - size_t ffn_up_weight; - size_t ffn_up_bias; - size_t ffn_norm_1_weight; - size_t ffn_norm_1_bias; - size_t ffn_down_weight; - size_t ffn_down_bias; - size_t ffn_norm_2_weight; - size_t ffn_norm_2_bias; - size_t mlp_router_layer_weight; - size_t mlp_experts_bias; - std::vector mlp_experts_mlp1_weight; - std::vector mlp_experts_mlp2_weight; - }; - - std::vector layers; - } weight_nodes_; -}; - -class WhisperModel : public Model { -public: - WhisperModel(); - explicit WhisperModel(const Config& config); - ~WhisperModel() override = default; - -protected: - size_t build_attention(CactusGraph*, size_t, uint32_t,ComputeBackend, bool, size_t) override { - throw std::runtime_error("Whisper: build_attention unused"); - } - - size_t build_mlp(CactusGraph*, size_t, uint32_t, ComputeBackend) const override { - throw std::runtime_error("Whisper: build_mlp unused"); - } - - size_t build_transformer_block(CactusGraph*, size_t, uint32_t, ComputeBackend, bool, size_t) override { - throw std::runtime_error("Whisper: build_transformer_block unused"); - } - - size_t forward(const std::vector& /*tokens*/, bool /*use_cache*/ = false) override { - throw std::runtime_error("Whisper requires mel+token forward()."); - } - - size_t forward(const std::vector& mel_bins, const std::vector& tokens, bool use_cache = false) override; - - void run_encoder(const std::vector& mel_bins); - void reset_graph_side_cache_nodes(); - - size_t run_decoder_step(const std::vector& tokens, bool use_cache, bool last_token_only); - - void load_weights_to_graph(CactusGraph* gb) override; - - size_t build_encoder_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0); - - size_t build_decoder_self_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0); - - size_t build_encoder_self_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0); - - size_t build_encoder_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend); - - size_t build_decoder_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const; - - size_t build_encoder_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0); - - size_t build_decoder_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache = false, size_t position_offset = 0); - - size_t build_conv1d(CactusGraph* gb, size_t input); - - uint32_t decode_with_audio(const std::vector& tokens, const std::vector& mel_bins, - float temperature = 0.0f, float top_p = 0.0f, size_t top_k = 0, const std::string& profile_file = "", float* out_entropy = nullptr) override; - - std::vector get_audio_embeddings(const std::vector& mel_bins) override; - - void reset_cache() override; - -private: - struct WeightNodeIDs { - size_t output_weight; - size_t output_norm_weight; - - size_t decoder_norm_weight; - size_t decoder_norm_bias; - size_t decoder_position_embeddings_weight; - - size_t encoder_position_embeddings; - size_t encoder_conv1_weight; - size_t encoder_conv1_bias; - size_t encoder_conv2_weight; - size_t encoder_conv2_bias; - size_t encoder_norm_weight; - size_t encoder_norm_bias; - - size_t encoder_output; - - struct LayerWeights { - //Decoder layers - size_t decoder_output_norm_bias; - size_t decoder_output_norm_weight; - size_t decoder_position_embeddings_weight; - size_t decoder_token_embeddings_weight; - - size_t decoder_encoder_attn_q_weight; - size_t decoder_encoder_attn_k_weight; - size_t decoder_encoder_attn_v_weight; - size_t decoder_encoder_attn_q_bias; - size_t decoder_encoder_attn_v_bias; - size_t decoder_encoder_attn_output_weight; - size_t decoder_encoder_attn_output_bias; - - size_t decoder_post_encoder_layernorm_weight; - size_t decoder_post_encoder_layernorm_bias; - - size_t decoder_ffn1_weight; - size_t decoder_ffn1_bias; - size_t decoder_ffn2_weight; - size_t decoder_ffn2_bias; - - size_t decoder_post_ffn_layernorm_weight; - size_t decoder_post_ffn_layernorm_bias; - - size_t decoder_self_attn_q_weight; - size_t decoder_self_attn_k_weight; - size_t decoder_self_attn_v_weight; - size_t decoder_self_attn_q_bias; - size_t decoder_self_attn_v_bias; - size_t decoder_self_attn_output_weight; - size_t decoder_self_attn_output_bias; - - size_t decoder_post_attn_layernorm_weight; - size_t decoder_post_attn_layernorm_bias; - - //Encoder layers - size_t encoder_ffn1_weight; - size_t encoder_ffn1_bias; - size_t encoder_ffn2_weight; - size_t encoder_ffn2_bias; - - size_t encoder_post_ffn_layernorm_weight; - size_t encoder_post_ffn_layernorm_bias; - - size_t encoder_self_attn_q_weight; - size_t encoder_self_attn_k_weight; - size_t encoder_self_attn_v_weight; - size_t encoder_self_attn_q_bias; - size_t encoder_self_attn_v_bias; - size_t encoder_self_attn_output_weight; - size_t encoder_self_attn_output_bias; - - size_t encoder_post_attn_layernorm_weight; - size_t encoder_post_attn_layernorm_bias; - }; - - std::vector layers; - } weight_nodes_; - - bool encoder_ready_ = false; - size_t last_new_tokens_; - std::vector encoder_output_host_; - std::vector encoder_output_shape_; - size_t last_conv1_node_ = 0; - size_t last_conv2_node_ = 0; - size_t last_encoder_post_norm_node_ = 0; - size_t last_enc_plus_pos_node_ = 0; - size_t encoder_transformer_block_0 = 0; - size_t encoder_pre_gelu = 0; - size_t encoder_post_gelu = 0; - - std::vector encoder_block_out_nodes_; - std::vector encoder_output_bytes_; - Precision encoder_output_precision_ = Precision::FP32; - - std::vector suppress_tokens_ = { - 1, - 2, - 7, - 8, - 9, - 10, - 14, - 25, - 26, - 27, - 28, - 29, - 31, - 58, - 59, - 60, - 61, - 62, - 63, - 90, - 91, - 92, - 93, - 359, - 503, - 522, - 542, - 873, - 893, - 902, - 918, - 922, - 931, - 1350, - 1853, - 1982, - 2460, - 2627, - 3246, - 3253, - 3268, - 3536, - 3846, - 3961, - 4183, - 4667, - 6585, - 6647, - 7273, - 9061, - 9383, - 10428, - 10929, - 11938, - 12033, - 12331, - 12562, - 13793, - 14157, - 14635, - 15265, - 15618, - 16553, - 16604, - 18362, - 18956, - 20075, - 21675, - 22520, - 26130, - 26161, - 26435, - 28279, - 29464, - 31650, - 32302, - 32470, - 36865, - 42863, - 47425, - 49870, - 50254, - 50258, - 50358, - 50359, - 50360, - 50361, - 50362 - }; - - std::vector begin_suppress_tokens_ = { - 220, - 50257 - }; - - bool first_decode_step_ = true; - - std::vector encoder_k_nodes_; - std::vector encoder_v_nodes_; - - std::vector> encoder_k_host_; - std::vector> encoder_v_host_; - std::vector> encoder_k_shape_; - std::vector> encoder_v_shape_; - Precision encoder_kv_precision_ = Precision::FP32; - bool encoder_kv_ready_ = false; - - std::unique_ptr npu_encoder_; - bool use_npu_encoder_ = false; - -}; - - -class Lfm2VlModel : public Model { -public: - Lfm2VlModel(); - explicit Lfm2VlModel(const Config& config); - ~Lfm2VlModel() override = default; - - bool init(const std::string& model_folder, size_t context_size, const std::string& system_prompt = "", bool do_warmup = true) override; - size_t forward(const std::vector& tokens, bool use_cache = false) override; - - uint32_t decode(const std::vector& tokens, - float temperature = -1.0f, - float top_p = -1.0f, - size_t top_k = 0, - const std::string& profile_file = "", - float* out_entropy = nullptr) override; - - void prefill(const std::vector& tokens, size_t chunk_size = 256, const std::string& profile_file = "") override; - - uint32_t decode_with_images( - const std::vector& tokens, - const std::vector& image_paths, - float temperature = -1.0f, - float top_p = -1.0f, - size_t top_k = 0, - const std::string& profile_file = "", - float* out_entropy = nullptr) override; - - void reset_cache() override; - std::vector get_image_embeddings(const std::string& image_path) override; - -protected: - size_t build_attention(CactusGraph*, size_t, uint32_t, ComputeBackend, bool, size_t) override; - size_t build_mlp(CactusGraph*, size_t, uint32_t, ComputeBackend) const override; - size_t build_transformer_block(CactusGraph*, size_t, uint32_t, ComputeBackend, bool, size_t) override; - - void load_weights_to_graph(CactusGraph* gb) override; - -private: - struct ProjectedTileFeature { - size_t node_id; - size_t token_count; - }; - - struct TextEmbeddingInput { - size_t input_node; - std::vector tokens; - }; - - struct MergedEmbeddingResult { - size_t node_id; - size_t seq_len; - }; - - struct ForwardImageResult { - size_t final_hidden_node; - size_t seq_len; - }; - std::vector get_image_features( - CactusGraph* gb, - const Siglip2Preprocessor::PreprocessedImage& preprocessed_image, - ComputeBackend backend); - - ForwardImageResult forward_images( - CactusGraph* gb, - const std::vector& tokens, - const std::vector& image_paths, - ComputeBackend backend, - bool use_cache); - size_t build_multimodal_projector( - CactusGraph* gb, - size_t image_features, - size_t tile_h, - size_t tile_w, - ComputeBackend backend); - size_t pixel_unshuffle(CactusGraph* gb, size_t hidden_states, size_t height, size_t width, size_t channels); - MergedEmbeddingResult merge_image_text_embeddings( - CactusGraph* gb, - const std::vector& tokens, - const std::vector>& image_embedding_nodes, - std::vector& text_embedding_inputs); - Siglip2VisionModel vision_tower_; - LFM2Model language_model_; - Siglip2Preprocessor preprocessor_; - struct ProjectorWeights { - size_t layer_norm_weight; - size_t layer_norm_bias; - size_t linear_1_weight; - size_t linear_1_bias; - size_t linear_2_weight; - size_t linear_2_bias; - } projector_weights_; - - bool vision_weights_loaded_ = false; - bool language_weights_loaded_ = false; - - bool image_prefill_completed_ = false; - size_t last_token_count_ = 0; -}; - -} -} diff --git a/cactus/models/model_gemma.cpp b/cactus/models/model_gemma.cpp deleted file mode 100644 index 4e71fef60..000000000 --- a/cactus/models/model_gemma.cpp +++ /dev/null @@ -1,196 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include "../npu/npu.h" -#include -#include -#include - -namespace cactus { -namespace engine { - -GemmaModel::GemmaModel() : Model() {} - -GemmaModel::GemmaModel(const Config& config) : Model(config) { - weight_nodes_.layers.resize(config.num_layers); -} - -void GemmaModel::post_init() { - kv_cache_.set_window_size(0, 0); -} - -void GemmaModel::load_weights_to_graph(CactusGraph* gb) { - embedding_node_id_ = gb->mmap_embeddings(embedding_file_path_); - weight_nodes_.output_norm_weight = gb->mmap_weights(model_folder_path_ + "/output_norm.weights"); - - if (config_.tie_word_embeddings) { - weight_nodes_.output_weight = embedding_node_id_; - output_weight_node_id_ = embedding_node_id_; - } else { - weight_nodes_.output_weight = gb->mmap_weights(model_folder_path_ + "/output_weight.weights"); - output_weight_node_id_ = weight_nodes_.output_weight; - } - - for (uint32_t i = 0; i < config_.num_layers; i++) { - auto& layer = weight_nodes_.layers[i]; - std::string layer_prefix = model_folder_path_ + "/layer_" + std::to_string(i) + "_"; - layer.attn_q_weight = gb->mmap_weights(layer_prefix + "attn_q.weights"); - layer.attn_k_weight = gb->mmap_weights(layer_prefix + "attn_k.weights"); - layer.attn_v_weight = gb->mmap_weights(layer_prefix + "attn_v.weights"); - layer.attn_output_weight = gb->mmap_weights(layer_prefix + "attn_output.weights"); - layer.input_layernorm_weight = gb->mmap_weights(layer_prefix + "input_norm.weights"); - layer.attn_q_norm_weight = gb->mmap_weights(layer_prefix + "attn_q_norm.weights"); - layer.attn_k_norm_weight = gb->mmap_weights(layer_prefix + "attn_k_norm.weights"); - layer.ffn_gate_weight = gb->mmap_weights(layer_prefix + "ffn_gate.weights"); - layer.ffn_up_weight = gb->mmap_weights(layer_prefix + "ffn_up.weights"); - layer.ffn_down_weight = gb->mmap_weights(layer_prefix + "ffn_down.weights"); - layer.post_attention_layernorm_weight = gb->mmap_weights(layer_prefix + "post_attn_norm.weights"); - layer.pre_feedforward_layernorm_weight = gb->mmap_weights(layer_prefix + "pre_ffn_norm.weights"); - layer.post_feedforward_layernorm_weight = gb->mmap_weights(layer_prefix + "post_ffn_norm.weights"); - } - - if (npu::is_npu_available()) { - std::string npu_prefill_path = model_folder_path_ + "/model.mlpackage"; - load_npu_prefill(npu_prefill_path); - } -} - -size_t GemmaModel::build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto q_proj = gb->matmul(normalized_input, layer.attn_q_weight, true, backend); - auto k_proj = gb->matmul(normalized_input, layer.attn_k_weight, true, backend); - auto v_proj = gb->matmul(normalized_input, layer.attn_v_weight, true, backend); - - const auto& q_shape = gb->get_output_buffer(q_proj).shape; - size_t batch_seq = q_shape[0]; - size_t num_heads = config_.attention_heads; - size_t head_dim = config_.attention_head_dim; - q_proj = gb->reshape(q_proj, {batch_seq * num_heads, head_dim}); - q_proj = gb->rms_norm(q_proj, layer.attn_q_norm_weight, config_.layer_norm_eps); - q_proj = gb->reshape(q_proj, {batch_seq, num_heads * head_dim}); - - size_t num_kv_heads = config_.attention_kv_heads; - k_proj = gb->reshape(k_proj, {batch_seq * num_kv_heads, head_dim}); - k_proj = gb->rms_norm(k_proj, layer.attn_k_norm_weight, config_.layer_norm_eps); - k_proj = gb->reshape(k_proj, {batch_seq, num_kv_heads * head_dim}); - - size_t seq_len = batch_seq; - - auto q_proj_4d = gb->reshape(q_proj, {1, seq_len, config_.attention_heads, config_.attention_head_dim}); - auto k_proj_4d = gb->reshape(k_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - auto v_proj_4d = gb->reshape(v_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - - bool is_global_attention = ((layer_idx + 1) % 6) == 0; - - if (config_.rope_theta > 0) { - float rope_freq = is_global_attention ? 1000000.0f : 10000.0f; - - q_proj_4d = gb->rope(q_proj_4d, rope_freq, position_offset); - k_proj_4d = gb->rope(k_proj_4d, rope_freq, position_offset); - } - - if (use_cache) { - cache_k_output_nodes_[layer_idx] = k_proj_4d; - cache_v_output_nodes_[layer_idx] = v_proj_4d; - } - - size_t attn_output_4d; - size_t window_size = is_global_attention ? 0 : 512; - - if (use_cache && !kv_cache_.is_empty()) { - attn_output_4d = gb->attention_int8_hybrid( - q_proj_4d, k_proj_4d, v_proj_4d, - attention_scale_, position_offset, - kv_cache_.get_keys_int8(layer_idx), - kv_cache_.get_values_int8(layer_idx), - kv_cache_.get_key_scales(layer_idx), - kv_cache_.get_value_scales(layer_idx), - kv_cache_.current_seq_len, num_kv_heads, head_dim - ); - } else { - attn_output_4d = gb->attention(q_proj_4d, k_proj_4d, v_proj_4d, attention_scale_, position_offset, window_size); - } - auto attn_output = gb->reshape(attn_output_4d, {seq_len, config_.attention_head_dim * config_.attention_heads}); - return gb->matmul(attn_output, layer.attn_output_weight, true, backend); -} - - -size_t GemmaModel::build_mlp(CactusGraph* gb, size_t input, uint32_t layer_idx, - ComputeBackend backend) const { - const auto& layer = weight_nodes_.layers[layer_idx]; - - size_t gate_output = gb->matmul(input, layer.ffn_gate_weight, true, backend); - size_t up_output = gb->matmul(input, layer.ffn_up_weight, true, backend); - size_t gate_gelu = gb->gelu(gate_output); - size_t gated = gb->multiply(gate_gelu, up_output); - return gb->matmul(gated, layer.ffn_down_weight, true, backend); -} - -size_t GemmaModel::build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto normalized_input = gb->rms_norm(hidden, layer.input_layernorm_weight, config_.layer_norm_eps); - auto attn_output = build_attention(gb, normalized_input, layer_idx, backend, use_cache, position_offset); - - auto normalized_attn = gb->rms_norm(attn_output, layer.post_attention_layernorm_weight, config_.layer_norm_eps); - - auto after_attention = gb->add_clipped(hidden, normalized_attn); - - auto pre_mlp_norm = gb->rms_norm(after_attention, layer.pre_feedforward_layernorm_weight, config_.layer_norm_eps); - auto mlp_output = build_mlp(gb, pre_mlp_norm, layer_idx, backend); - - auto normalized_mlp = gb->rms_norm(mlp_output, layer.post_feedforward_layernorm_weight, config_.layer_norm_eps); - - return gb->add_clipped(after_attention, normalized_mlp); -} - - -size_t GemmaModel::forward(const std::vector& tokens, bool use_cache) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - if (tokens.empty()) { - throw std::runtime_error("Token sequence cannot be empty"); - } - - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - - auto seq_len = static_cast(tokens.size()); - - size_t position_offset = use_cache ? kv_cache_.get_total_seq_len() : 0; - - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - auto input_node_id = gb->input({seq_len}, Precision::FP32); - auto hidden = gb->embedding(embedding_node_id_, input_node_id); - - float embed_scale = std::sqrt(static_cast(config_.hidden_dim)); - hidden = gb->scalar_multiply(hidden, embed_scale); - - static std::set skip_layers = {}; - for (uint32_t layer_idx = 0; layer_idx < config_.num_layers; layer_idx++) { - if (skip_layers.count(layer_idx)) { - continue; - } - hidden = build_transformer_block(gb, hidden, layer_idx, backend, use_cache, position_offset); - } - - auto final_hidden = gb->rms_norm(hidden, weight_nodes_.output_norm_weight, config_.layer_norm_eps); - - std::vector input_data(seq_len); - for (size_t i = 0; i < seq_len; i++) { - input_data[i] = static_cast(tokens[i]); - } - gb->set_input(input_node_id, input_data.data(), Precision::FP32); - - return final_hidden; -} - -} -} \ No newline at end of file diff --git a/cactus/models/model_lfm2.cpp b/cactus/models/model_lfm2.cpp deleted file mode 100644 index 978492d00..000000000 --- a/cactus/models/model_lfm2.cpp +++ /dev/null @@ -1,398 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include -#include -#include -#include -#include -#include -#include -namespace cactus { -namespace engine { -LFM2Model::LFM2Model() : Model() { - weight_nodes_.layers.resize(config_.num_layers); - conv_cache_bx_nodes_.assign(config_.num_layers, 0); -} -LFM2Model::LFM2Model(const Config& config) : Model(config) { - weight_nodes_.layers.resize(config_.num_layers); - conv_cache_bx_nodes_.assign(config_.num_layers, 0); -} -void LFM2Model::post_init() { - if (config_.conv_L_cache > 0) { - Precision cache_precision = Precision::FP16; - size_t conv_window = config_.conv_L_cache > 0 ? config_.conv_L_cache - 1 : 0; - conv_cache_.init(config_.num_layers, config_.hidden_dim, conv_window, cache_precision); - - } - last_forward_used_cache_ = false; - -} -void LFM2Model::reset_cache() { - Model::reset_cache(); - - if (conv_cache_.window_size > 0) { - conv_cache_.reset(); - - } -} -bool LFM2Model::is_cache_empty() const { - return kv_cache_.is_empty(); -} -bool LFM2Model::init(const std::string& model_folder, size_t context_size, const std::string& system_prompt, bool do_warmup) { - if (!Model::init(model_folder, context_size, system_prompt, do_warmup)) { - return false; - } - - if (weight_nodes_.layers.size() != config_.num_layers) { - weight_nodes_.layers.resize(config_.num_layers); - - } - conv_cache_bx_nodes_.assign(config_.num_layers, 0); - - return true; -} -bool LFM2Model::init(CactusGraph* external_graph, const std::string& model_folder, size_t context_size, - const std::string& system_prompt, bool do_warmup) { - if (!Model::init(external_graph, model_folder, context_size, system_prompt, do_warmup)) { - return false; - } - - if (weight_nodes_.layers.size() != config_.num_layers) { - weight_nodes_.layers.resize(config_.num_layers); - - } - conv_cache_bx_nodes_.assign(config_.num_layers, 0); - - return true; -} -void LFM2Model::load_weights_to_graph(CactusGraph* gb) { - if (weight_nodes_.layers.size() != config_.num_layers) { - weight_nodes_.layers.resize(config_.num_layers); - } - if (conv_cache_bx_nodes_.size() != config_.num_layers) { - conv_cache_bx_nodes_.assign(config_.num_layers, 0); - } - embedding_node_id_ = gb->mmap_embeddings(embedding_file_path_); - weight_nodes_.output_norm_weight = gb->mmap_weights(model_folder_path_ + "/output_norm.weights"); - - if (config_.tie_word_embeddings) { - weight_nodes_.output_weight = embedding_node_id_; - output_weight_node_id_ = embedding_node_id_; - - } else { - weight_nodes_.output_weight = gb->mmap_weights(model_folder_path_ + "/output_weight.weights"); - output_weight_node_id_ = weight_nodes_.output_weight; - - } - for (uint32_t i = 0; i < config_.num_layers; i++) { - auto& layer_entry = weight_nodes_.layers[i]; - auto& layer = layer_entry.weights; - std::string layer_prefix = model_folder_path_ + "/layer_" + std::to_string(i) + "_"; - - bool is_conv_layer = false; - if (i < config_.layer_types.size()) { - std::string layer_type = config_.layer_types[i]; - is_conv_layer = (layer_type == "conv" || layer_type == "CONV"); - - } - if (is_conv_layer) { - layer_entry.type = WeightNodeIDs::LayerType::CONV; - layer.conv_in_proj_weight = gb->mmap_weights(layer_prefix + "conv_in_proj.weights"); - layer.conv_out_proj_weight = gb->mmap_weights(layer_prefix + "conv_out_proj.weights"); - layer.conv_depthwise_weight = gb->mmap_weights(layer_prefix + "conv_depthwise.weights"); - - } else { - layer_entry.type = WeightNodeIDs::LayerType::ATTENTION; - layer.attn_q_weight = gb->mmap_weights(layer_prefix + "attn_q.weights"); - layer.attn_k_weight = gb->mmap_weights(layer_prefix + "attn_k.weights"); - layer.attn_v_weight = gb->mmap_weights(layer_prefix + "attn_v.weights"); - layer.attn_output_weight = gb->mmap_weights(layer_prefix + "attn_output.weights"); - layer.attn_q_norm_weight = gb->mmap_weights(layer_prefix + "attn_q_norm.weights"); - layer.attn_k_norm_weight = gb->mmap_weights(layer_prefix + "attn_k_norm.weights"); - - } - layer.input_layernorm_weight = gb->mmap_weights(layer_prefix + "input_norm.weights"); - layer.post_attention_layernorm_weight = gb->mmap_weights(layer_prefix + "post_attn_norm.weights"); - layer.ffn_gate_weight = gb->mmap_weights(layer_prefix + "ffn_gate.weights"); - layer.ffn_up_weight = gb->mmap_weights(layer_prefix + "ffn_up.weights"); - layer.ffn_down_weight = gb->mmap_weights(layer_prefix + "ffn_down.weights"); - - } - -} -size_t LFM2Model::build_conv1d(CactusGraph* gb, size_t input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache) { - const auto& layer_entry = weight_nodes_.layers[layer_idx]; - const auto& layer = layer_entry.weights; - size_t in_proj = gb->matmul(input, layer.conv_in_proj_weight, true, backend); - const auto& in_proj_buf = gb->get_output_buffer(in_proj); - if (in_proj_buf.shape.size() != 2 || (in_proj_buf.shape[1] % 3) != 0) { - throw std::runtime_error("Conv in_proj output must be [L, 3*C]"); - } - - const size_t L = in_proj_buf.shape[0]; - const size_t C = in_proj_buf.shape[1] / 3; - - size_t triplet = gb->reshape(in_proj, {L, static_cast(3), C}); - size_t B = gb->slice(triplet, 1, 0, 1); - size_t Cg = gb->slice(triplet, 1, 1, 1); - size_t X = gb->slice(triplet, 1, 2, 1); - - B = gb->reshape(B, {L, C}); - - Cg = gb->reshape(Cg, {L, C}); - - X = gb->reshape(X, {L, C}); - size_t Bx = gb->multiply(B, X); - if (use_cache) { - conv_cache_bx_nodes_[layer_idx] = Bx; - - } else { - conv_cache_bx_nodes_[layer_idx] = 0; - - } - const auto& wbuf = gb->get_output_buffer(layer.conv_depthwise_weight); - size_t K = wbuf.shape.back(); - - size_t conv_w = layer.conv_depthwise_weight; - if (wbuf.shape.size() == 2) { - K = wbuf.shape[1]; - conv_w = gb->reshape(conv_w, {wbuf.shape[0], static_cast(1), K}); - } else if (wbuf.shape.size() != 3){ - throw std::runtime_error("Unexpected depthwise weight rank"); - } - K = wbuf.shape.back(); - - size_t conv_input_lc = Bx; - if (use_cache && conv_cache_.window_size > 0) { - auto view = conv_cache_.get_window(layer_idx); - std::vector segments; - if (view.len2 > 0) { - size_t left_node = gb->input({view.len2, C}, conv_cache_.precision); - gb->set_input(left_node, view.ptr2, conv_cache_.precision); - segments.push_back(left_node); - - } - if (view.len1 > 0) { - size_t right_node = gb->input({view.len1, C}, conv_cache_.precision); - gb->set_input(right_node, view.ptr1, conv_cache_.precision); - segments.push_back(right_node); - - } - if (!segments.empty()) { - conv_input_lc = segments[0]; - for (size_t idx = 1; idx < segments.size(); ++idx) { - conv_input_lc = gb->concat(conv_input_lc, segments[idx], 0); - - } - conv_input_lc = gb->concat(conv_input_lc, Bx, 0); - - } - } - - const auto& conv_input_buf = gb->get_output_buffer(conv_input_lc); - size_t total_len = conv_input_buf.shape[0]; - - size_t x_nlc = gb->reshape(conv_input_lc, {static_cast(1), total_len, C}); - const size_t dilation = 1; - - size_t y_nlc = gb->conv1d_causal(x_nlc, conv_w, K, dilation); - size_t start = total_len > L ? total_len - L : 0; - size_t y_slice = gb->slice(y_nlc, 1, start, L); - - size_t y_lc = gb->reshape(y_slice, {L, C}); - size_t gated = gb->multiply(Cg, y_lc); - size_t projected = gb->matmul(gated, layer.conv_out_proj_weight, true, backend); - return projected; -} - -size_t LFM2Model::build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer_entry = weight_nodes_.layers[layer_idx]; - const auto& layer = layer_entry.weights; - auto q_proj_linear = gb->matmul(normalized_input, layer.attn_q_weight, true, backend); - auto k_proj_linear = gb->matmul(normalized_input, layer.attn_k_weight, true, backend); - auto v_proj_linear = gb->matmul(normalized_input, layer.attn_v_weight, true, backend); - const auto& q_shape = gb->get_output_buffer(q_proj_linear).shape; - size_t batch_seq = q_shape[0]; - size_t num_heads = config_.attention_heads; - size_t head_dim = config_.attention_head_dim; - - auto q_proj_reshaped = gb->reshape(q_proj_linear, {batch_seq * num_heads, head_dim}); - - auto q_proj_norm = gb->rms_norm(q_proj_reshaped, layer.attn_q_norm_weight, config_.layer_norm_eps); - - auto q_proj = gb->reshape(q_proj_norm, {batch_seq, num_heads * head_dim}); - size_t num_kv_heads = config_.attention_kv_heads; - auto k_proj_reshaped = gb->reshape(k_proj_linear, {batch_seq * num_kv_heads, head_dim}); - - auto k_proj_norm = gb->rms_norm(k_proj_reshaped, layer.attn_k_norm_weight, config_.layer_norm_eps); - - auto k_proj = gb->reshape(k_proj_norm, {batch_seq, num_kv_heads * head_dim}); - size_t seq_len = batch_seq; - - auto q_proj_4d = gb->reshape(q_proj, {1, seq_len, config_.attention_heads, config_.attention_head_dim}); - auto k_proj_4d = gb->reshape(k_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - auto v_proj_4d = gb->reshape(v_proj_linear, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - - if (config_.rope_theta > 0) { - q_proj_4d = gb->rope(q_proj_4d, config_.rope_theta, position_offset); - k_proj_4d = gb->rope(k_proj_4d, config_.rope_theta, position_offset); - } - if (use_cache) { - cache_k_output_nodes_[layer_idx] = k_proj_4d; - cache_v_output_nodes_[layer_idx] = v_proj_4d; - } - - size_t attn_output_4d; - - if (use_cache && !kv_cache_.is_empty()) { - attn_output_4d = gb->attention_int8_hybrid( - q_proj_4d, k_proj_4d, v_proj_4d, - attention_scale_, position_offset, - kv_cache_.get_keys_int8(layer_idx), - kv_cache_.get_values_int8(layer_idx), - kv_cache_.get_key_scales(layer_idx), - kv_cache_.get_value_scales(layer_idx), - kv_cache_.current_seq_len, num_kv_heads, head_dim - ); - } else { - attn_output_4d = gb->attention(q_proj_4d, k_proj_4d, v_proj_4d, attention_scale_, position_offset); - } - auto attn_output = gb->reshape(attn_output_4d, {seq_len, config_.attention_head_dim * config_.attention_heads}); - auto projected = gb->matmul(attn_output, layer.attn_output_weight, true, backend); - return projected; -} - -size_t LFM2Model::build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, ComputeBackend backend) const { - const auto& layer_entry = weight_nodes_.layers[layer_idx]; - const auto& layer = layer_entry.weights; - auto gate = gb->matmul(normalized_h, layer.ffn_gate_weight, true, backend); - auto up = gb->matmul(normalized_h, layer.ffn_up_weight, true, backend); - auto activated = gb->multiply(gb->silu(gate), up); - - auto down = gb->matmul(activated, layer.ffn_down_weight, true, backend); - - return down; -} - -size_t LFM2Model::build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer_entry = weight_nodes_.layers[layer_idx]; - const auto& layer = layer_entry.weights; - auto normalized_input = gb->rms_norm(hidden, layer.input_layernorm_weight, config_.layer_norm_eps); - size_t block_output; - if (layer_entry.type == WeightNodeIDs::LayerType::CONV) { - block_output = build_conv1d(gb, normalized_input, layer_idx, backend, use_cache); - } else { - if (layer_idx < conv_cache_bx_nodes_.size()) { - conv_cache_bx_nodes_[layer_idx] = 0; - } - block_output = build_attention(gb, normalized_input, layer_idx, backend, use_cache, position_offset); - } - - auto after_block = gb->add(hidden, block_output); - auto normalized_after_block = gb->rms_norm(after_block, layer.post_attention_layernorm_weight, config_.layer_norm_eps); - auto mlp_output = build_mlp(gb, normalized_after_block, layer_idx, backend); - auto block_result = gb->add(after_block, mlp_output); - return block_result; -} - -size_t LFM2Model::forward(CactusGraph* gb, size_t input_embeddings, size_t seq_len, - ComputeBackend backend, bool use_cache) { - if (seq_len == 0) { - throw std::runtime_error("Sequence length must be greater than zero"); - } - - if (conv_cache_bx_nodes_.size() != config_.num_layers) { - conv_cache_bx_nodes_.assign(config_.num_layers, 0); - - } - std::fill(conv_cache_bx_nodes_.begin(), conv_cache_bx_nodes_.end(), 0); - - last_forward_used_cache_ = use_cache; - - if (!use_cache && conv_cache_.window_size > 0) { - conv_cache_.reset(); - - } - size_t position_offset = use_cache ? kv_cache_.get_total_seq_len() : 0; - - size_t hidden = input_embeddings; - - for (uint32_t layer_idx = 0; layer_idx < config_.num_layers; layer_idx++) { - hidden = build_transformer_block(gb, hidden, layer_idx, backend, use_cache, position_offset); - - } - - auto final_hidden = gb->rms_norm(hidden, weight_nodes_.output_norm_weight, config_.layer_norm_eps); - - return final_hidden; -} -size_t LFM2Model::forward(CactusGraph* gb, const std::vector& tokens, - ComputeBackend backend, bool use_cache) { - if (tokens.empty()) { - throw std::runtime_error("Token sequence cannot be empty"); - } - auto seq_len = static_cast(tokens.size()); - auto input_node_id = gb->input({seq_len}, Precision::FP32); - const auto& embedding_buffer = gb->get_output_buffer(embedding_node_id_); - - for (size_t i = 0; i < embedding_buffer.shape.size(); ++i) { - - if (i + 1 < embedding_buffer.shape.size()) { - - } - } - - auto hidden = gb->embedding(embedding_node_id_, input_node_id); - auto final_hidden = forward(gb, hidden, seq_len, backend, use_cache); - - std::vector input_data(seq_len); - for (size_t i = 0; i < seq_len; i++) { - input_data[i] = static_cast(tokens[i]); - } - gb->set_input(input_node_id, input_data.data(), Precision::FP32); - - return final_hidden; -} -size_t LFM2Model::forward(const std::vector& tokens, bool use_cache) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - return forward(gb, tokens, backend, use_cache); -} -void LFM2Model::post_execute_updates(CactusGraph* gb, size_t /*seq_len*/) { - if (conv_cache_bx_nodes_.empty()) { - return; - } - - if (!last_forward_used_cache_ || conv_cache_.window_size == 0) { - std::fill(conv_cache_bx_nodes_.begin(), conv_cache_bx_nodes_.end(), 0); - - return; - } - size_t layer_count = std::min(conv_cache_bx_nodes_.size(), weight_nodes_.layers.size()); - for (size_t layer_idx = 0; layer_idx < layer_count; ++layer_idx) { - if (weight_nodes_.layers[layer_idx].type != WeightNodeIDs::LayerType::CONV) { - conv_cache_bx_nodes_[layer_idx] = 0; - continue; - } - size_t bx_node = conv_cache_bx_nodes_[layer_idx]; - if (bx_node != 0) { - conv_cache_.update(gb, layer_idx, bx_node); - } - conv_cache_bx_nodes_[layer_idx] = 0; - } - last_forward_used_cache_ = false; - -} -} -} \ No newline at end of file diff --git a/cactus/models/model_lfm2vl.cpp b/cactus/models/model_lfm2vl.cpp deleted file mode 100644 index 2da651126..000000000 --- a/cactus/models/model_lfm2vl.cpp +++ /dev/null @@ -1,599 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include -#include -#include -#include -#include - -namespace cactus { -namespace engine { - -Lfm2VlModel::Lfm2VlModel() : Model() { - config_.model_type = Config::ModelType::LFM2; -} - -Lfm2VlModel::Lfm2VlModel(const Config& config) - : Model(config), - vision_tower_(config), - language_model_(config) { - Siglip2Preprocessor::Config preprocessor_config; - preprocessor_config.patch_size = static_cast(config.vision_patch_size); - preprocessor_config.downsample_factor = static_cast(config.downsample_factor); - preprocessor_config.min_tiles = static_cast(config.min_tiles); - preprocessor_config.max_tiles = static_cast(config.max_tiles); - preprocessor_config.use_thumbnail = config.use_thumbnail; - preprocessor_config.min_image_tokens = static_cast(config.min_image_tokens); - preprocessor_config.max_image_tokens = static_cast(config.max_image_tokens); - preprocessor_config.max_num_patches = static_cast(config.max_num_patches); - preprocessor_config.tile_size = static_cast(config.tile_size); - preprocessor_config.max_pixels_tolerance = config.max_pixels_tolerance; - preprocessor_config.do_resize = true; - preprocessor_config.do_rescale = true; - preprocessor_config.do_normalize = true; - preprocessor_config.do_convert_rgb = true; - preprocessor_config.do_image_splitting = config.do_image_splitting; - preprocessor_config.rescale_factor = config.rescale_factor; - preprocessor_config.image_mean[0] = config.image_mean; - preprocessor_config.image_mean[1] = config.image_mean; - preprocessor_config.image_mean[2] = config.image_mean; - preprocessor_config.image_std[0] = config.image_std; - preprocessor_config.image_std[1] = config.image_std; - preprocessor_config.image_std[2] = config.image_std; - - preprocessor_ = Siglip2Preprocessor(preprocessor_config); -} - -bool Lfm2VlModel::init(const std::string& model_folder, size_t context_size, const std::string& system_prompt, bool do_warmup) { - - if (!Model::init(model_folder, context_size, system_prompt, false)) { - return false; - } - auto* shared_graph = static_cast(graph_handle_); - if (!shared_graph) { - throw std::runtime_error("Shared graph was not initialized for Lfm2VlModel"); - } - std::string vision_folder = model_folder; - if (!vision_tower_.init(shared_graph, vision_folder, context_size, "", false)) { - throw std::runtime_error("Failed to initialize vision tower"); - } - - vision_weights_loaded_ = true; - if (!language_model_.init(shared_graph, model_folder, context_size, system_prompt, false)) { - throw std::runtime_error("Failed to initialize language model"); - } - - language_model_.output_weight_node_id_ = output_weight_node_id_; - - language_weights_loaded_ = true; - - if (do_warmup) { - std::string warmup_text = system_prompt.empty() ? "Hello" : system_prompt; - auto warmup_tokens = tokenizer_->encode(warmup_text); - language_model_.decode(warmup_tokens, config_.default_temperature, config_.default_top_p, config_.default_top_k, ""); - language_model_.reset_cache(); - } - - return true; -} - -void Lfm2VlModel::reset_cache() { - Model::reset_cache(); - language_model_.reset_cache(); - image_prefill_completed_ = false; - last_token_count_ = 0; -} - -void Lfm2VlModel::load_weights_to_graph(CactusGraph* gb) { - namespace fs = std::filesystem; - fs::path base(model_folder_path_); - - auto resolve_weight = [&](const std::string& primary, const std::string& fallback = "") -> std::string { - fs::path primary_path = base / primary; - if (fs::exists(primary_path)) { - return primary_path.string(); - } - if (!fallback.empty()) { - fs::path fallback_path = base / fallback; - if (fs::exists(fallback_path)) { - return fallback_path.string(); - } - } - return primary_path.string(); - }; - - projector_weights_.layer_norm_weight = gb->mmap_weights(resolve_weight("projector_layer_norm.weights")); - - projector_weights_.layer_norm_bias = gb->mmap_weights(resolve_weight("projector_layer_norm.bias.weights")); - - projector_weights_.linear_1_weight = gb->mmap_weights(resolve_weight("projector_linear_1.weights", "projector_linear1.weights")); - projector_weights_.linear_1_bias = gb->mmap_weights(resolve_weight("projector_linear_1.bias.weights", "projector_linear1.bias.weights")); - projector_weights_.linear_2_weight = gb->mmap_weights(resolve_weight("projector_linear_2.weights", "projector_linear2.weights")); - projector_weights_.linear_2_bias = gb->mmap_weights(resolve_weight("projector_linear_2.bias.weights", "projector_linear2.bias.weights")); - output_weight_node_id_ = gb->mmap_weights(resolve_weight("output_weight.weights")); -} - -size_t Lfm2VlModel::pixel_unshuffle(CactusGraph* gb, size_t hidden_states, - size_t height, size_t width, size_t channels) { - - const size_t factor = config_.downsample_factor; - const size_t new_height = height / factor; - const size_t new_width = width / factor; - size_t step1 = gb->reshape(hidden_states, {1, height, new_width, channels * factor}); - step1 = gb->transposeN(step1, {0, 2, 1, 3}); - size_t step2 = gb->reshape(step1, {1, new_width, new_height, channels * factor * factor}); - size_t result = gb->transposeN(step2, {0, 2, 1, 3}); - return result; -} - -size_t Lfm2VlModel::build_multimodal_projector(CactusGraph* gb, size_t image_features, - size_t tile_h, size_t tile_w, ComputeBackend backend) { - const size_t vision_hidden = config_.vision_embed_dim; - - const auto& input_buf = gb->get_output_buffer(image_features); - size_t image_features_fp16 = (input_buf.precision == Precision::FP16) - ? image_features - : gb->precision_cast(image_features, Precision::FP16); - - size_t unshuffled = pixel_unshuffle(gb, image_features_fp16, tile_h, tile_w, vision_hidden); - const size_t factor = config_.downsample_factor; - const size_t new_h = tile_h / factor; - const size_t new_w = tile_w / factor; - const size_t in_channels = vision_hidden * factor * factor; - const size_t seq_len = new_h * new_w; - size_t flattened = gb->reshape(unshuffled, {seq_len, in_channels}); - size_t normalized = gb->layernorm(flattened, projector_weights_.layer_norm_weight, - projector_weights_.layer_norm_bias, config_.layer_norm_eps); - size_t hidden = gb->matmul(normalized, projector_weights_.linear_1_weight, true, backend); - hidden = gb->add(hidden, projector_weights_.linear_1_bias); - hidden = gb->gelu(hidden); - size_t output = gb->matmul(hidden, projector_weights_.linear_2_weight, true, backend); - output = gb->add(output, projector_weights_.linear_2_bias); - return output; -} - -std::vector Lfm2VlModel::get_image_embeddings(const std::string& image_path) { - return vision_tower_.get_image_embedding(image_path); -} - -std::vector Lfm2VlModel::get_image_features( - CactusGraph* gb, - const Siglip2Preprocessor::PreprocessedImage& preprocessed_image, - ComputeBackend backend) { - - size_t vision_output = vision_tower_.forward_vision(gb, preprocessed_image, backend); - std::vector projected_features; - projected_features.reserve(preprocessed_image.spatial_shapes.size()); - - size_t offset = 0; - for (size_t tile_idx = 0; tile_idx < preprocessed_image.spatial_shapes.size(); ++tile_idx) { - const auto& shape = preprocessed_image.spatial_shapes[tile_idx]; - const size_t tile_h = shape.first; - const size_t tile_w = shape.second; - const size_t tile_tokens = tile_h * tile_w; - const size_t factor = config_.downsample_factor; - if (factor == 0) { - throw std::runtime_error("Downsample factor must be greater than zero"); - } - if (tile_h % factor != 0 || tile_w % factor != 0) { - throw std::runtime_error("Tile dimensions must be divisible by downsample factor"); - } - const size_t new_h = tile_h / factor; - const size_t new_w = tile_w / factor; - const size_t projected_tokens = new_h * new_w; - - size_t tile_features = gb->slice(vision_output, 0, offset, tile_tokens); - offset += tile_tokens; - - size_t reshaped = gb->reshape(tile_features, {1, tile_h, tile_w, config_.vision_embed_dim}); - size_t projected = build_multimodal_projector(gb, reshaped, tile_h, tile_w, backend); - ProjectedTileFeature feature{}; - feature.node_id = projected; - feature.token_count = projected_tokens; - projected_features.push_back(feature); - - } - - return projected_features; -} - -Lfm2VlModel::MergedEmbeddingResult Lfm2VlModel::merge_image_text_embeddings( - CactusGraph* gb, - const std::vector& tokens, - const std::vector>& image_embedding_nodes, - std::vector& text_embedding_inputs) { - - text_embedding_inputs.clear(); - - Tokenizer* tokenizer = language_model_.get_tokenizer(); - if (!tokenizer) { - throw std::runtime_error("Tokenizer must be initialized before merging embeddings"); - } - - const uint32_t image_token_id = tokenizer->get_image_token_id(); - - auto get_token_id = [tokenizer](const std::string& token_text) -> uint32_t { - auto encoded = tokenizer->encode(token_text); - if (encoded.size() != 1) { - - throw std::runtime_error("Expected single token encoding for " + token_text); - } - return encoded[0]; - }; - - const uint32_t image_start_id = get_token_id("<|image_start|>"); - const uint32_t image_end_id = get_token_id("<|image_end|>"); - - std::vector sequence_nodes; - sequence_nodes.reserve(tokens.size() + image_embedding_nodes.size()); - - std::vector current_segment; - current_segment.reserve(tokens.size()); - - size_t total_seq_len = 0; - - auto flush_segment = [&](void) { - if (current_segment.empty()) { - - return; - } - - const size_t segment_len = current_segment.size(); - TextEmbeddingInput segment; - segment.tokens.swap(current_segment); - segment.input_node = gb->input({segment.tokens.size()}, Precision::FP32); - - const auto& embedding_buffer = gb->get_output_buffer(language_model_.embedding_node_id_); - - for (size_t i = 0; i < embedding_buffer.shape.size(); ++i) { - (void)i; - } - - size_t embedding_node = gb->embedding(language_model_.embedding_node_id_, segment.input_node); - - text_embedding_inputs.push_back(std::move(segment)); - - sequence_nodes.push_back(embedding_node); - - total_seq_len += segment_len; - - current_segment.clear(); - - }; - - size_t token_index = 0; - size_t image_index = 0; - - while (token_index < tokens.size()) { - uint32_t token_id = tokens[token_index]; - - if (token_id == image_start_id) { - flush_segment(); - - if (image_index >= image_embedding_nodes.size()) { - - throw std::runtime_error("Encountered <|image_start|> without corresponding image features"); - } - - current_segment.push_back(token_id); - ++token_index; - - const auto& tiles = image_embedding_nodes[image_index]; - size_t tile_index = 0; - while (token_index < tokens.size()) { - uint32_t inner_token = tokens[token_index]; - if (inner_token == image_token_id) { - flush_segment(); - - if (tile_index >= tiles.size()) { - - throw std::runtime_error("More placeholders than projected tile features"); - } - - const auto& tile = tiles[tile_index++]; - sequence_nodes.push_back(tile.node_id); - - total_seq_len += tile.token_count; - - for (size_t count = 0; count < tile.token_count; ++count) { - if (token_index >= tokens.size()) { - throw std::runtime_error("Insufficient tokens for projected features"); - } - if (tokens[token_index] != image_token_id) { - throw std::runtime_error("Unexpected token encountered within image feature span"); - } - ++token_index; - - } - continue; - } - - current_segment.push_back(inner_token); - - ++token_index; - - if (inner_token == image_end_id) { - flush_segment(); - break; - } - } - - if (tile_index != tiles.size()) { - if (tile_index < tiles.size()) { - for (size_t remaining = tile_index; remaining < tiles.size(); ++remaining) { - (void)tiles[remaining]; - } - } - throw std::runtime_error("Unused projected tile features remain after processing image block"); - } - - ++image_index; - - } else { - current_segment.push_back(token_id); - - ++token_index; - - } - } - - flush_segment(); - if (image_index != image_embedding_nodes.size()) { - throw std::runtime_error("Not all image features were consumed while merging embeddings"); - } - - if (sequence_nodes.empty()) { - throw std::runtime_error("Failed to build embedding sequence from provided tokens"); - } - - size_t merged = sequence_nodes[0]; - for (size_t idx = 1; idx < sequence_nodes.size(); ++idx) { - merged = gb->concat(merged, sequence_nodes[idx], 0); - - } - - return MergedEmbeddingResult{merged, total_seq_len}; -} - -size_t Lfm2VlModel::build_attention(CactusGraph*, size_t, uint32_t, ComputeBackend, bool, size_t) { - throw std::runtime_error("build_attention should not be called directly on Lfm2VlModel"); -} - -size_t Lfm2VlModel::build_mlp(CactusGraph*, size_t, uint32_t, ComputeBackend) const { - throw std::runtime_error("build_mlp should not be called directly on Lfm2VlModel"); -} - -size_t Lfm2VlModel::build_transformer_block(CactusGraph*, size_t, uint32_t, ComputeBackend, bool, size_t) { - throw std::runtime_error("build_transformer_block should not be called directly on Lfm2VlModel"); -} - -size_t Lfm2VlModel::forward(const std::vector& tokens, bool use_cache) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - size_t final_hidden = language_model_.forward(gb, tokens, backend, use_cache); - - return final_hidden; -} - -uint32_t Lfm2VlModel::decode(const std::vector& tokens, - float temperature, - float top_p, - size_t top_k, - const std::string& profile_file, - float* out_entropy) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - if (temperature < 0) { - temperature = config_.default_temperature; - } - if (top_p < 0) { - top_p = config_.default_top_p; - } - if (top_k == 0) { - top_k = config_.default_top_k; - } - - image_prefill_completed_ = false; - last_token_count_ = tokens.size(); - - return language_model_.decode(tokens, temperature, top_p, top_k, profile_file, out_entropy); -} - -void Lfm2VlModel::prefill(const std::vector& tokens, size_t chunk_size, const std::string& profile_file) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - image_prefill_completed_ = false; - last_token_count_ = tokens.size(); - - language_model_.prefill(tokens, chunk_size, profile_file); -} - -Lfm2VlModel::ForwardImageResult Lfm2VlModel::forward_images( - CactusGraph* gb, - const std::vector& tokens, - const std::vector& image_paths, - ComputeBackend backend, - bool use_cache) { - if (!gb) { - throw std::runtime_error("Graph must be initialized before forwarding"); - } - if (tokens.empty()) { - throw std::runtime_error("Token sequence cannot be empty"); - } - - std::vector> all_image_embeddings; - all_image_embeddings.reserve(image_paths.size()); - for (const auto& image_path : image_paths) { - auto preprocessed = preprocessor_.preprocess_from_file(image_path); - - auto image_features = get_image_features(gb, preprocessed, backend); - - all_image_embeddings.push_back(std::move(image_features)); - } - - std::vector text_embedding_inputs; - text_embedding_inputs.reserve(tokens.size() / 4 + 1); - - auto merged_embeddings = merge_image_text_embeddings(gb, tokens, all_image_embeddings, text_embedding_inputs); - if (merged_embeddings.seq_len == 0) { - throw std::runtime_error("Merged embedding sequence length cannot be zero"); - } - - for (const auto& embedding_input : text_embedding_inputs) { - if (embedding_input.tokens.empty()) { - continue; - } - - std::vector segment_data(embedding_input.tokens.size()); - for (size_t i = 0; i < embedding_input.tokens.size(); ++i) { - segment_data[i] = static_cast(embedding_input.tokens[i]); - } - gb->set_input(embedding_input.input_node, segment_data.data(), Precision::FP32); - } - size_t final_hidden = language_model_.forward(gb, merged_embeddings.node_id, merged_embeddings.seq_len, backend, use_cache); - return ForwardImageResult{final_hidden, merged_embeddings.seq_len}; -} - -uint32_t Lfm2VlModel::decode_with_images( - const std::vector& tokens, - const std::vector& image_paths, - float temperature, - float top_p, - size_t top_k, - const std::string& profile_file, - float* out_entropy) { - - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - if (image_paths.empty()) { - - image_prefill_completed_ = false; - last_token_count_ = tokens.size(); - return language_model_.decode(tokens, temperature, top_p, top_k, profile_file, out_entropy); - } - - if (temperature < 0) { - temperature = config_.default_temperature; - } - if (top_p < 0) { - top_p = config_.default_top_p; - } - if (top_k == 0) { - top_k = config_.default_top_k; - } - - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - bool cache_empty = language_model_.is_cache_empty(); - bool need_prefill = cache_empty || !image_prefill_completed_; - - if (!need_prefill && tokens.size() <= last_token_count_) { - - reset_cache(); - need_prefill = true; - } - - size_t seq_len_for_updates = 0; - size_t final_hidden_node = 0; - - if (need_prefill) { - auto forward_result = forward_images(gb, tokens, image_paths, backend, true); - - final_hidden_node = forward_result.final_hidden_node; - seq_len_for_updates = forward_result.seq_len; - image_prefill_completed_ = true; - last_token_count_ = tokens.size(); - } else { - size_t delta = tokens.size() - last_token_count_; - if (delta > tokens.size()) { - delta = tokens.size(); - } - if (delta == 0) { - if (tokens.empty()) { - throw std::runtime_error("Token sequence cannot be empty for cached decode step"); - } - delta = 1; - - } - std::vector incremental_tokens(tokens.end() - delta, tokens.end()); - - final_hidden_node = language_model_.forward(gb, incremental_tokens, backend, true); - - seq_len_for_updates = incremental_tokens.size(); - last_token_count_ = tokens.size(); - } - - auto logits_node_id = gb->matmul(final_hidden_node, language_model_.output_weight_node_id_, true, backend); - auto sampled_token_id = gb->sample(logits_node_id, temperature, top_p, top_k); - if (!profile_file.empty()) { - gb->execute(profile_file); - - } else { - gb->execute(); - - } - - if (out_entropy) { - const auto& logits_buf = gb->get_output_buffer(logits_node_id); - void* logits_ptr = gb->get_output(logits_node_id); - size_t vocab_size = logits_buf.shape.back(); - - std::vector logits(vocab_size); - if (logits_buf.precision == Precision::FP32) { - float* src = static_cast(logits_ptr); - std::copy(src, src + vocab_size, logits.begin()); - } else if (logits_buf.precision == Precision::FP16) { - __fp16* src = static_cast<__fp16*>(logits_ptr); - Quantization::fp16_to_fp32(src, logits.data(), vocab_size); - } else { - int8_t* src = static_cast(logits_ptr); - Quantization::int8_to_fp32(src, logits.data(), vocab_size, 1.0f); - } - - float max_logit = *std::max_element(logits.begin(), logits.end()); - double sum_exp = 0.0; - for (size_t i = 0; i < vocab_size; ++i) { - sum_exp += std::exp(static_cast(logits[i] - max_logit)); - } - double log_sum_exp = static_cast(max_logit) + std::log(sum_exp); - - double entropy = 0.0; - for (size_t i = 0; i < vocab_size; ++i) { - double log_prob = static_cast(logits[i]) - log_sum_exp; - double prob = std::exp(log_prob); - if (prob > 1e-10) { - entropy -= prob * log_prob; - } - } - - double max_entropy = std::log(static_cast(vocab_size)); - *out_entropy = static_cast(entropy / max_entropy); - } - - language_model_.post_execute_updates(gb, seq_len_for_updates); - language_model_.update_kv_cache(gb, seq_len_for_updates); - - auto* output_ptr = gb->get_output(sampled_token_id); - return *static_cast(output_ptr); -} - -} -} \ No newline at end of file diff --git a/cactus/models/model_nomic.cpp b/cactus/models/model_nomic.cpp deleted file mode 100644 index 69d753fce..000000000 --- a/cactus/models/model_nomic.cpp +++ /dev/null @@ -1,217 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include -#include - -namespace cactus { -namespace engine { - -NomicModel::NomicModel() : Model() {} - -NomicModel::NomicModel(const Config& config) : Model(config) { - weight_nodes_.layers.resize(config.num_layers); -} - -void NomicModel::load_weights_to_graph(CactusGraph* gb) { - embedding_node_id_ = gb->mmap_embeddings(embedding_file_path_); - weight_nodes_.embedding_layernorm_weight = gb->mmap_weights(model_folder_path_ + "/embedding_layernorm.weight"); - weight_nodes_.embedding_layernorm_bias = gb->mmap_weights(model_folder_path_ + "/embedding_layernorm.bias"); - - for (uint32_t i = 0; i < config_.num_layers; i++) { - auto& layer = weight_nodes_.layers[i]; - std::string layer_prefix = model_folder_path_ + "/layer_" + std::to_string(i) + "_"; - layer.attn_q_weight = gb->mmap_weights(layer_prefix + "attn_q.weights"); - layer.attn_k_weight = gb->mmap_weights(layer_prefix + "attn_k.weights"); - layer.attn_v_weight = gb->mmap_weights(layer_prefix + "attn_v.weights"); - layer.attn_q_bias = gb->mmap_weights(layer_prefix + "attn_q.bias"); - layer.attn_k_bias = gb->mmap_weights(layer_prefix + "attn_k.bias"); - layer.attn_v_bias = gb->mmap_weights(layer_prefix + "attn_v.bias"); - - layer.attn_output_weight = gb->mmap_weights(layer_prefix + "attn_output.weights"); - layer.attn_output_bias = gb->mmap_weights(layer_prefix + "attn_output.bias"); - layer.ffn_norm_1_weight = gb->mmap_weights(layer_prefix + "norm1.weights"); - layer.ffn_norm_1_bias = gb->mmap_weights(layer_prefix + "norm1.bias"); - layer.ffn_norm_2_weight = gb->mmap_weights(layer_prefix + "norm2.weights"); - layer.ffn_norm_2_bias = gb->mmap_weights(layer_prefix + "norm2.bias"); - - if ((i + 1) % config_.moe_every_n_layers != 0) { - layer.ffn_up_weight = gb->mmap_weights(layer_prefix + "mlp_fc1.weights"); - layer.ffn_up_bias = gb->mmap_weights(layer_prefix + "mlp_fc1.bias"); - layer.ffn_down_weight = gb->mmap_weights(layer_prefix + "mlp_fc2.weights"); - layer.ffn_down_bias = gb->mmap_weights(layer_prefix + "mlp_fc2.bias"); - - } else { - layer.mlp_router_layer_weight = gb->mmap_weights(layer_prefix + "mlp_router.layer.weights"); - layer.mlp_experts_bias = gb->mmap_weights(layer_prefix + "mlp_experts.bias"); - - for (uint32_t j = 0; j < config_.num_experts; j++) { - layer.mlp_experts_mlp1_weight.push_back(gb->mmap_weights(layer_prefix + "mlp_expert_" + std::to_string(j) + ".mlp1.weights")); - layer.mlp_experts_mlp2_weight.push_back(gb->mmap_weights(layer_prefix + "mlp_expert_" + std::to_string(j) + ".mlp2.weights")); - } - } - } -} - -size_t NomicModel::build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - if (use_cache) { - throw std::runtime_error("NomicModel does not support generation, it's an encoder model"); - } - (void)position_offset; - - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto q_proj = gb->matmul(normalized_input, layer.attn_q_weight, true, backend); - q_proj = gb->add(q_proj, layer.attn_q_bias); - - auto k_proj = gb->matmul(normalized_input, layer.attn_k_weight, true, backend); - k_proj = gb->add(k_proj, layer.attn_k_bias); - - auto v_proj = gb->matmul(normalized_input, layer.attn_v_weight, true, backend); - v_proj = gb->add(v_proj, layer.attn_v_bias); - - const auto& q_shape = gb->get_output_buffer(q_proj).shape; - const size_t seq_len = q_shape[0]; - const size_t num_heads = config_.attention_heads; - const size_t head_dim = config_.attention_head_dim; - - if (num_heads == 0 || head_dim == 0) { - throw std::runtime_error("Invalid attention head configuration for Nomic model"); - } - - auto reshape_to_heads = [&](size_t tensor) { - return gb->reshape(tensor, {1, seq_len, num_heads, head_dim}); - }; - - auto q_proj_4d = reshape_to_heads(q_proj); - auto k_proj_4d = reshape_to_heads(k_proj); - auto v_proj_4d = reshape_to_heads(v_proj); - - if (config_.rope_theta > 0) { - q_proj_4d = gb->rope(q_proj_4d, config_.rope_theta, 0); - k_proj_4d = gb->rope(k_proj_4d, config_.rope_theta, 0); - } - - auto attn_output_4d = gb->attention(q_proj_4d, k_proj_4d, v_proj_4d, attention_scale_, false); - auto attn_output = gb->reshape(attn_output_4d, {seq_len, num_heads * head_dim}); - - auto output = gb->matmul(attn_output, layer.attn_output_weight, true, backend); - output = gb->add(output, layer.attn_output_bias); - return output; -} - -size_t NomicModel::build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const { - if ((layer_idx + 1) % config_.moe_every_n_layers != 0) { - return build_standard_mlp(gb, normalized_h, layer_idx, backend); - } else { - return build_moe_mlp(gb, normalized_h, layer_idx, backend); - } -} - -size_t NomicModel::build_standard_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const { - const auto& layer = weight_nodes_.layers[layer_idx]; - auto hidden = gb->matmul(normalized_h, layer.ffn_up_weight, true, backend); - hidden = gb->add(hidden, layer.ffn_up_bias); - hidden = gb->gelu(hidden); - hidden = gb->matmul(hidden, layer.ffn_down_weight, true, backend); - hidden = gb->add(hidden, layer.ffn_down_bias); - return hidden; -} - -size_t NomicModel::build_moe_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const { - const auto& layer = weight_nodes_.layers[layer_idx]; - const auto& router_shape = gb->get_output_buffer(layer.mlp_router_layer_weight).shape; - if (router_shape.size() != 2) { - throw std::runtime_error("MoE router weight must be 2D"); - } - - const size_t num_experts = config_.num_experts != 0 ? config_.num_experts : router_shape[0]; - const size_t seq_len = gb->get_output_buffer(normalized_h).shape[0]; - - auto gate_weights = gb->matmul(normalized_h, layer.mlp_router_layer_weight, true, backend); - auto gate_probs = gb->softmax(gate_weights); - - auto topk_indices_and_values = gb->topk(gate_probs, config_.num_top_experts); - auto topk_idx = gb->index(topk_indices_and_values, 0, 0); - auto topk_w = gb->index(topk_indices_and_values, 1, 0); - auto expert_token_weights_matrix = gb->scatter_topk(topk_idx, topk_w, num_experts); - - size_t expert_outputs = 0; - for (size_t e = 0; e < num_experts; e++) { - auto new_expert_output = gb->matmul(normalized_h, layer.mlp_experts_mlp1_weight[e], true, backend); - new_expert_output = gb->gelu(new_expert_output); - new_expert_output = gb->matmul(new_expert_output, layer.mlp_experts_mlp2_weight[e], true, backend); - - auto expert_token_weights = gb->index(expert_token_weights_matrix, e, 0); - const auto& expert_weights_prec = gb->get_output_buffer(new_expert_output).precision; - if (gb->get_output_buffer(expert_token_weights).precision != expert_weights_prec) { - expert_token_weights = gb->precision_cast(expert_token_weights, expert_weights_prec); - } - expert_token_weights = gb->reshape(expert_token_weights, {seq_len, 1}); - - new_expert_output = gb->multiply(new_expert_output, expert_token_weights); - if (e == 0) { - expert_outputs = new_expert_output; - } else { - expert_outputs = gb->add(expert_outputs, new_expert_output); - } - } - - return gb->add(expert_outputs, layer.mlp_experts_bias); -} - -size_t NomicModel::build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - if (use_cache) { - throw std::runtime_error("NomicModel does not support generation, it's an encoder model"); - } - (void)position_offset; - const auto& layer = weight_nodes_.layers[layer_idx]; - auto attn_output = build_attention(gb, hidden, layer_idx, backend); - auto residual = gb->add(hidden, attn_output); - auto normalized_residual = gb->layernorm(residual, layer.ffn_norm_1_weight, layer.ffn_norm_1_bias, config_.layer_norm_eps); - auto mlp_output = build_mlp(gb, normalized_residual, layer_idx, backend); - auto final_residual = gb->add(normalized_residual, mlp_output); - auto normalized_final_residual = gb->layernorm(final_residual, layer.ffn_norm_2_weight, layer.ffn_norm_2_bias, config_.layer_norm_eps); - return normalized_final_residual; -} - -size_t NomicModel::forward(const std::vector& tokens, bool use_cache) { - if (use_cache) { - throw std::runtime_error("NomicModel does not support generation, it's an encoder model"); - } - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - - size_t seq_len = static_cast(tokens.size()); - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - size_t input_node_id = gb->input({seq_len}, Precision::FP32); - std::vector input_data(seq_len); - for (size_t i = 0; i < seq_len; i++) { - input_data[i] = static_cast(tokens[i]); - } - gb->set_input(input_node_id, input_data.data(), Precision::FP32); - - size_t hidden = gb->embedding(embedding_node_id_, input_node_id); - - hidden = gb->layernorm(hidden, weight_nodes_.embedding_layernorm_weight, weight_nodes_.embedding_layernorm_bias, config_.layer_norm_eps); - - static std::set skip_layers = {}; - for (uint32_t layer_idx = 0; layer_idx < config_.num_layers; layer_idx++) { - if (skip_layers.count(layer_idx)) { - continue; - } - hidden = build_transformer_block(gb, hidden, layer_idx, backend); - } - - return hidden; -} - -} -} diff --git a/cactus/models/model_qwen.cpp b/cactus/models/model_qwen.cpp deleted file mode 100644 index b08e57d87..000000000 --- a/cactus/models/model_qwen.cpp +++ /dev/null @@ -1,175 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include "../npu/npu.h" -#include -#include -#include - -namespace cactus { -namespace engine { - -QwenModel::QwenModel() : Model() {} - -QwenModel::QwenModel(const Config& config) : Model(config) { - weight_nodes_.layers.resize(config.num_layers); -} - -void QwenModel::load_weights_to_graph(CactusGraph* gb) { - embedding_node_id_ = gb->mmap_embeddings(embedding_file_path_); - weight_nodes_.output_norm_weight = gb->mmap_weights(model_folder_path_ + "/output_norm.weights"); - - if (config_.tie_word_embeddings) { - weight_nodes_.output_weight = embedding_node_id_; - output_weight_node_id_ = embedding_node_id_; - } else { - weight_nodes_.output_weight = gb->mmap_weights(model_folder_path_ + "/output_weight.weights"); - output_weight_node_id_ = weight_nodes_.output_weight; - } - - for (uint32_t i = 0; i < config_.num_layers; i++) { - auto& layer = weight_nodes_.layers[i]; - std::string layer_prefix = model_folder_path_ + "/layer_" + std::to_string(i) + "_"; - layer.attn_q_weight = gb->mmap_weights(layer_prefix + "attn_q.weights"); - layer.attn_k_weight = gb->mmap_weights(layer_prefix + "attn_k.weights"); - layer.attn_v_weight = gb->mmap_weights(layer_prefix + "attn_v.weights"); - layer.attn_output_weight = gb->mmap_weights(layer_prefix + "attn_output.weights"); - layer.input_layernorm_weight = gb->mmap_weights(layer_prefix + "input_norm.weights"); - layer.attn_q_norm_weight = gb->mmap_weights(layer_prefix + "attn_q_norm.weights"); - layer.attn_k_norm_weight = gb->mmap_weights(layer_prefix + "attn_k_norm.weights"); - layer.ffn_gate_weight = gb->mmap_weights(layer_prefix + "ffn_gate.weights"); - layer.ffn_up_weight = gb->mmap_weights(layer_prefix + "ffn_up.weights"); - layer.ffn_down_weight = gb->mmap_weights(layer_prefix + "ffn_down.weights"); - layer.post_attention_layernorm_weight = gb->mmap_weights(layer_prefix + "post_attn_norm.weights"); - } - - if (npu::is_npu_available()) { - std::string npu_prefill_path = model_folder_path_ + "/model.mlpackage"; - load_npu_prefill(npu_prefill_path); - } -} - -size_t QwenModel::build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto q_proj = gb->matmul(normalized_input, layer.attn_q_weight, true, backend); - auto k_proj = gb->matmul(normalized_input, layer.attn_k_weight, true, backend); - auto v_proj = gb->matmul(normalized_input, layer.attn_v_weight, true, backend); - - const auto& q_shape = gb->get_output_buffer(q_proj).shape; - size_t batch_seq = q_shape[0]; - size_t num_heads = config_.attention_heads; - size_t head_dim = config_.attention_head_dim; - q_proj = gb->reshape(q_proj, {batch_seq * num_heads, head_dim}); - q_proj = gb->rms_norm(q_proj, layer.attn_q_norm_weight, config_.layer_norm_eps); - q_proj = gb->reshape(q_proj, {batch_seq, num_heads * head_dim}); - - size_t num_kv_heads = config_.attention_kv_heads; - k_proj = gb->reshape(k_proj, {batch_seq * num_kv_heads, head_dim}); - k_proj = gb->rms_norm(k_proj, layer.attn_k_norm_weight, config_.layer_norm_eps); - k_proj = gb->reshape(k_proj, {batch_seq, num_kv_heads * head_dim}); - - size_t seq_len = batch_seq; - - auto q_proj_4d = gb->reshape(q_proj, {1, seq_len, config_.attention_heads, config_.attention_head_dim}); - auto k_proj_4d = gb->reshape(k_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - auto v_proj_4d = gb->reshape(v_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - - if (config_.rope_theta > 0) { - q_proj_4d = gb->rope(q_proj_4d, config_.rope_theta, position_offset); - k_proj_4d = gb->rope(k_proj_4d, config_.rope_theta, position_offset); - } - - size_t attn_output_4d; - - if (use_cache) { - cache_k_output_nodes_[layer_idx] = k_proj_4d; - cache_v_output_nodes_[layer_idx] = v_proj_4d; - } - - if (use_cache && !kv_cache_.is_empty()) { - attn_output_4d = gb->attention_int8_hybrid( - q_proj_4d, k_proj_4d, v_proj_4d, - attention_scale_, position_offset, - kv_cache_.get_keys_int8(layer_idx), - kv_cache_.get_values_int8(layer_idx), - kv_cache_.get_key_scales(layer_idx), - kv_cache_.get_value_scales(layer_idx), - kv_cache_.current_seq_len, num_kv_heads, head_dim - ); - } else { - attn_output_4d = gb->attention(q_proj_4d, k_proj_4d, v_proj_4d, attention_scale_, position_offset); - } - - auto attn_output = gb->reshape(attn_output_4d, {seq_len, config_.attention_head_dim * config_.attention_heads}); - return gb->matmul(attn_output, layer.attn_output_weight, true, backend); -} - - -size_t QwenModel::build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const { - const auto& layer = weight_nodes_.layers[layer_idx]; - size_t gate_output = gb->matmul(normalized_h, layer.ffn_gate_weight, true, backend); - size_t up_output = gb->matmul(normalized_h, layer.ffn_up_weight, true, backend); - size_t gate_silu = gb->silu(gate_output); - size_t gated = gb->multiply(gate_silu, up_output); - return gb->matmul(gated, layer.ffn_down_weight, true, backend); -} - - -size_t QwenModel::build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer = weight_nodes_.layers[layer_idx]; - auto normalized_input = gb->rms_norm(hidden, layer.input_layernorm_weight, config_.layer_norm_eps); - auto attn_output = build_attention(gb, normalized_input, layer_idx, backend, use_cache, position_offset); - auto after_attention = gb->add(hidden, attn_output); - auto normalized_after_attention = gb->rms_norm(after_attention, layer.post_attention_layernorm_weight, config_.layer_norm_eps); - auto mlp_output = build_mlp(gb, normalized_after_attention, layer_idx, backend); - return gb->add(after_attention, mlp_output); -} - - -size_t QwenModel::forward(const std::vector& tokens, bool use_cache) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - if (tokens.empty()) { - throw std::runtime_error("Token sequence cannot be empty"); - } - - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - - auto seq_len = static_cast(tokens.size()); - - size_t position_offset = use_cache ? kv_cache_.get_total_seq_len() : 0; - - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - auto input_node_id = gb->input({seq_len}, Precision::FP32); - auto hidden = gb->embedding(embedding_node_id_, input_node_id); - - std::vector input_data(seq_len); - for (size_t i = 0; i < seq_len; i++) { - input_data[i] = static_cast(tokens[i]); - } - gb->set_input(input_node_id, input_data.data(), Precision::FP32); - - static std::set skip_layers = {}; - for (uint32_t layer_idx = 0; layer_idx < config_.num_layers; layer_idx++) { - if (skip_layers.count(layer_idx)) { - continue; - } - hidden = build_transformer_block(gb, hidden, layer_idx, backend, use_cache, position_offset); - } - - auto final_hidden = gb->rms_norm(hidden, weight_nodes_.output_norm_weight, config_.layer_norm_eps); - - return final_hidden; -} - -} -} \ No newline at end of file diff --git a/cactus/models/model_siglip2.cpp b/cactus/models/model_siglip2.cpp deleted file mode 100644 index aea1eefcd..000000000 --- a/cactus/models/model_siglip2.cpp +++ /dev/null @@ -1,414 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include "../kernel/kernel.h" -#include -#include -#include -#include - -namespace cactus { -namespace engine { - -Siglip2VisionModel::Siglip2VisionModel() : Model() { - config_.model_type = Config::ModelType::SIGLIP2; -} - -Siglip2VisionModel::Siglip2VisionModel(const Config& cfg) : Model(cfg) { - Siglip2Preprocessor::Config preprocessor_config; - preprocessor_config.patch_size = static_cast(config_.vision_patch_size); - preprocessor_config.downsample_factor = static_cast(config_.downsample_factor); - preprocessor_config.min_tiles = static_cast(config_.min_tiles); - preprocessor_config.max_tiles = static_cast(config_.max_tiles); - preprocessor_config.use_thumbnail = config_.use_thumbnail; - preprocessor_config.min_image_tokens = static_cast(config_.min_image_tokens); - preprocessor_config.max_image_tokens = static_cast(config_.max_image_tokens); - preprocessor_config.max_num_patches = static_cast(config_.max_num_patches); - preprocessor_config.tile_size = static_cast(config_.tile_size); - preprocessor_config.max_pixels_tolerance = config_.max_pixels_tolerance; - preprocessor_config.do_resize = true; - preprocessor_config.do_rescale = true; - preprocessor_config.do_normalize = true; - preprocessor_config.do_convert_rgb = true; - preprocessor_config.do_image_splitting = config_.do_image_splitting; - preprocessor_config.rescale_factor = config_.rescale_factor; - preprocessor_config.image_mean[0] = config_.image_mean; - preprocessor_config.image_mean[1] = config_.image_mean; - preprocessor_config.image_mean[2] = config_.image_mean; - preprocessor_config.image_std[0] = config_.image_std; - preprocessor_config.image_std[1] = config_.image_std; - preprocessor_config.image_std[2] = config_.image_std; - - preprocessor_ = Siglip2Preprocessor(preprocessor_config); -} - -void Siglip2VisionModel::load_weights_to_graph(CactusGraph* gb) { - std::string base = model_folder_path_ + "/"; - - if (npu::is_npu_available()) { - std::string npu_encoder_path = model_folder_path_ + "/model.mlpackage"; - npu_encoder_ = npu::create_encoder(); - if (npu_encoder_ && npu_encoder_->load(npu_encoder_path)) { - use_npu_encoder_ = true; - - std::vector typical_input_shape = { - static_cast(config_.max_num_patches), - static_cast(config_.vision_embed_dim) - }; - npu_encoder_->preallocate(typical_input_shape, "x", ""); - } else { - use_npu_encoder_ = false; - npu_encoder_.reset(); - } - } - - // Always load patch embedding and position embedding weights (needed for both CPU and NPU paths) - vision_weight_nodes_.patch_embedding_weight = gb->mmap_weights(base + "vision_patch_embedding.weights"); - vision_weight_nodes_.patch_embedding_bias = gb->mmap_weights(base + "vision_patch_embedding.bias.weights"); - vision_weight_nodes_.position_embedding = gb->mmap_weights(base + "vision_position_embedding.weights"); - - if (!use_npu_encoder_) { - vision_weight_nodes_.vision_layers.resize(config_.vision_num_layers); - - vision_weight_nodes_.post_layernorm_weight = gb->mmap_weights(base + "vision_post_layernorm.weights"); - vision_weight_nodes_.post_layernorm_bias = gb->mmap_weights(base + "vision_post_layernorm.bias.weights"); - - for (uint32_t i = 0; i < vision_weight_nodes_.vision_layers.size(); ++i) { - auto& layer = vision_weight_nodes_.vision_layers[i]; - std::string prefix = base + "vision_layer_" + std::to_string(i) + "_"; - - layer.attn_q_weight = gb->mmap_weights(prefix + "self_attn_q.weights"); - layer.attn_q_bias = gb->mmap_weights(prefix + "self_attn_q.bias.weights"); - layer.attn_k_weight = gb->mmap_weights(prefix + "self_attn_k.weights"); - layer.attn_k_bias = gb->mmap_weights(prefix + "self_attn_k.bias.weights"); - layer.attn_v_weight = gb->mmap_weights(prefix + "self_attn_v.weights"); - layer.attn_v_bias = gb->mmap_weights(prefix + "self_attn_v.bias.weights"); - layer.attn_output_weight = gb->mmap_weights(prefix + "self_attn_out.weights"); - layer.attn_output_bias = gb->mmap_weights(prefix + "self_attn_out.bias.weights"); - - layer.layer_norm1_weight = gb->mmap_weights(prefix + "layer_norm1.weights"); - layer.layer_norm1_bias = gb->mmap_weights(prefix + "layer_norm1.bias.weights"); - layer.layer_norm2_weight = gb->mmap_weights(prefix + "layer_norm2.weights"); - layer.layer_norm2_bias = gb->mmap_weights(prefix + "layer_norm2.bias.weights"); - - layer.mlp_fc1_weight = gb->mmap_weights(prefix + "ffn_fc1.weights"); - layer.mlp_fc1_bias = gb->mmap_weights(prefix + "ffn_fc1.bias.weights"); - layer.mlp_fc2_weight = gb->mmap_weights(prefix + "ffn_fc2.weights"); - layer.mlp_fc2_bias = gb->mmap_weights(prefix + "ffn_fc2.bias.weights"); - } - } -} - -Siglip2VisionModel::VisionEmbeddingResult Siglip2VisionModel::build_vision_embeddings( - CactusGraph* gb, - const Siglip2Preprocessor::PreprocessedImage& preprocessed_image, - ComputeBackend backend) { - const int num_tiles = preprocessed_image.num_tiles; - const int max_patches = preprocessed_image.max_patches_per_tile; - const int patch_dim = preprocessed_image.patch_dim; - - const size_t expected_size = static_cast(num_tiles) * static_cast(max_patches) * - static_cast(patch_dim); - if (preprocessed_image.pixel_values.size() != expected_size) { - throw std::runtime_error( - "Pixel values size mismatch: expected " + std::to_string(expected_size) + - " (tiles=" + std::to_string(num_tiles) + " * max_patches=" + std::to_string(max_patches) + - " * patch_dim=" + std::to_string(patch_dim) + ") but got " + - std::to_string(preprocessed_image.pixel_values.size())); - } - for (size_t i = 0; i < std::min(100, preprocessed_image.pixel_values.size()); ++i) { - float val = preprocessed_image.pixel_values[i]; - if (std::isnan(val) || std::isinf(val)) { - throw std::runtime_error( - "Invalid value in pixel_values at index " + std::to_string(i) + ": " + std::to_string(val)); - } - } - - size_t reshaped_weight = gb->reshape( - vision_weight_nodes_.patch_embedding_weight, - {static_cast(config_.vision_embed_dim), static_cast(patch_dim)}); - - size_t patch_bias = vision_weight_nodes_.patch_embedding_bias; - std::vector tile_embeddings; - tile_embeddings.reserve(static_cast(num_tiles)); - - for (int tile_idx = 0; tile_idx < num_tiles; ++tile_idx) { - const auto& shape = preprocessed_image.spatial_shapes[tile_idx]; - const int tile_h = shape.first; - const int tile_w = shape.second; - const int actual_patches = tile_h * tile_w; - if (actual_patches <= 0) { - continue; - } - - const float* tile_data = preprocessed_image.pixel_values.data() + - static_cast(tile_idx) * static_cast(max_patches) * - static_cast(patch_dim); - - size_t tile_input_fp32 = gb->input( - {static_cast(actual_patches), static_cast(patch_dim)}, Precision::FP32); - gb->set_input(tile_input_fp32, tile_data, Precision::FP32); - size_t tile_input = gb->precision_cast(tile_input_fp32, Precision::FP16); - size_t tile_patch = gb->matmul(tile_input, reshaped_weight, true, backend); - size_t tile_bias = gb->add(tile_patch, patch_bias); - size_t tile_pos = gb->bilinear_interpolation( - vision_weight_nodes_.position_embedding, - static_cast(tile_h), - static_cast(tile_w)); - size_t tile_pos_cast = gb->precision_cast(tile_pos, Precision::FP16); - size_t tile_embed = gb->add(tile_bias, tile_pos_cast); - tile_embeddings.push_back(tile_embed); - } - - if (tile_embeddings.empty()) { - throw std::runtime_error("No valid tiles produced embeddings in build_vision_embeddings"); - } - auto concat_nodes = [&](const std::vector& nodes) { - if (nodes.empty()) { - throw std::runtime_error("Attempted to concatenate an empty node list"); - } - size_t combined = nodes.front(); - for (size_t i = 1; i < nodes.size(); ++i) { - combined = gb->concat(combined, nodes[i], /*axis=*/0); - } - return combined; - }; - - size_t embeddings = concat_nodes(tile_embeddings); - return VisionEmbeddingResult{embeddings, std::move(tile_embeddings)}; -} - -size_t Siglip2VisionModel::build_vision_attention(CactusGraph* gb, size_t hidden_states, - uint32_t layer_idx, ComputeBackend backend) { - const auto& layer = vision_weight_nodes_.vision_layers[layer_idx]; - - size_t q = gb->matmul(hidden_states, layer.attn_q_weight, true, backend); - q = gb->add(q, layer.attn_q_bias); - size_t k = gb->matmul(hidden_states, layer.attn_k_weight, true, backend); - k = gb->add(k, layer.attn_k_bias); - size_t v = gb->matmul(hidden_states, layer.attn_v_weight, true, backend); - v = gb->add(v, layer.attn_v_bias); - const size_t num_heads = static_cast(config_.vision_attention_heads); - const size_t head_dim = static_cast(config_.vision_embed_dim / config_.vision_attention_heads); - const auto& q_buf = gb->get_output_buffer(q); - size_t seq_len = q_buf.shape[0]; - - size_t q_4d = gb->reshape(q, {1, seq_len, num_heads, head_dim}); - size_t k_4d = gb->reshape(k, {1, seq_len, num_heads, head_dim}); - size_t v_4d = gb->reshape(v, {1, seq_len, num_heads, head_dim}); - - float scale = 1.0f / std::sqrt(static_cast(head_dim)); - size_t attn_output = gb->attention(q_4d, k_4d, v_4d, scale, false, backend); - size_t attn_2d = gb->reshape(attn_output, {seq_len, num_heads * head_dim}); - size_t output = gb->matmul(attn_2d, layer.attn_output_weight, true, backend); - output = gb->add(output, layer.attn_output_bias); - return output; -} - -size_t Siglip2VisionModel::build_vision_mlp(CactusGraph* gb, size_t hidden_states, - uint32_t layer_idx, ComputeBackend backend) { - const auto& layer = vision_weight_nodes_.vision_layers[layer_idx]; - - size_t fc1_output = gb->matmul(hidden_states, layer.mlp_fc1_weight, true, backend); - fc1_output = gb->add(fc1_output, layer.mlp_fc1_bias); - size_t activated = gb->gelu(fc1_output); - size_t fc2_output = gb->matmul(activated, layer.mlp_fc2_weight, true, backend); - fc2_output = gb->add(fc2_output, layer.mlp_fc2_bias); - return fc2_output; -} - -size_t Siglip2VisionModel::build_vision_transformer_layer(CactusGraph* gb, size_t hidden_states, - uint32_t layer_idx, ComputeBackend backend) { - const auto& layer = vision_weight_nodes_.vision_layers[layer_idx]; - - size_t residual = hidden_states; - size_t normalized = gb->layernorm(hidden_states, layer.layer_norm1_weight, - layer.layer_norm1_bias, config_.layer_norm_eps); - size_t attn_output = build_vision_attention(gb, normalized, layer_idx, backend); - hidden_states = gb->add(residual, attn_output); - residual = hidden_states; - normalized = gb->layernorm(hidden_states, layer.layer_norm2_weight, - layer.layer_norm2_bias, config_.layer_norm_eps); - size_t mlp_output = build_vision_mlp(gb, normalized, layer_idx, backend); - hidden_states = gb->add(residual, mlp_output); - return hidden_states; -} - -size_t Siglip2VisionModel::forward_vision( - CactusGraph* gb, - const Siglip2Preprocessor::PreprocessedImage& preprocessed_image, - ComputeBackend backend) { - - if (use_npu_encoder_ && npu_encoder_ && npu_encoder_->is_available()) { - // NPU path: build patch embeddings + position embeddings on CPU, then run transformer on NPU - auto embedding_result = build_vision_embeddings(gb, preprocessed_image, backend); - - size_t total_patches = 0; - for (const auto& shape : preprocessed_image.spatial_shapes) { - total_patches += shape.first * shape.second; - } - - gb->execute(); - - const auto& embed_buffer = gb->get_output_buffer(embedding_result.combined_embeddings); - void* embed_ptr = gb->get_output(embedding_result.combined_embeddings); - - std::vector<__fp16> embed_f16; - const __fp16* input_ptr; - if (embed_buffer.precision == Precision::FP16) { - input_ptr = static_cast(embed_ptr); - } else { - embed_f16.resize(total_patches * config_.vision_embed_dim); - cactus_fp32_to_fp16(static_cast(embed_ptr), - embed_f16.data(), embed_f16.size()); - input_ptr = embed_f16.data(); - } - - std::vector input_shape = { - static_cast(total_patches), - static_cast(config_.vision_embed_dim) - }; - - __fp16* output_buffer = npu_encoder_->get_output_buffer(); - if (output_buffer) { - size_t elements = npu_encoder_->encode( - input_ptr, output_buffer, input_shape, "x", ""); - - if (elements > 0) { - gb->soft_reset(); - size_t vision_output = gb->input({total_patches, config_.vision_embed_dim}, - Precision::FP16); - gb->set_input(vision_output, output_buffer, Precision::FP16); - return vision_output; - } - } else { - std::vector<__fp16> npu_output(total_patches * config_.vision_embed_dim); - size_t elements = npu_encoder_->encode( - input_ptr, npu_output.data(), input_shape, "x", ""); - - if (elements > 0) { - gb->soft_reset(); - size_t vision_output = gb->input({total_patches, config_.vision_embed_dim}, - Precision::FP16); - gb->set_input(vision_output, npu_output.data(), Precision::FP16); - return vision_output; - } - } - - throw std::runtime_error("NPU encoder failed"); - } - - // CPU path: full forward pass through transformer layers - auto embedding_result = build_vision_embeddings(gb, preprocessed_image, backend); - - auto concat_nodes = [&](const std::vector& nodes) { - if (nodes.empty()) { - throw std::runtime_error("Attempted to concatenate an empty node list in forward_vision"); - } - size_t combined = nodes.front(); - for (size_t i = 1; i < nodes.size(); ++i) { - combined = gb->concat(combined, nodes[i], /*axis=*/0); - } - return combined; - }; - - std::vector tile_outputs; - tile_outputs.reserve(embedding_result.tile_embeddings.size()); - - for (size_t tile_idx = 0; tile_idx < embedding_result.tile_embeddings.size(); ++tile_idx) { - size_t hidden_states = embedding_result.tile_embeddings[tile_idx]; - - for (uint32_t layer_idx = 0; layer_idx < config_.vision_num_layers; ++layer_idx) { - hidden_states = build_vision_transformer_layer(gb, hidden_states, layer_idx, backend); - } - - hidden_states = gb->layernorm(hidden_states, - vision_weight_nodes_.post_layernorm_weight, - vision_weight_nodes_.post_layernorm_bias, - config_.layer_norm_eps); - tile_outputs.push_back(hidden_states); - } - - if (tile_outputs.empty()) { - throw std::runtime_error("No tile outputs generated in forward_vision"); - } - - size_t combined_output = concat_nodes(tile_outputs); - - return combined_output; -} - -size_t Siglip2VisionModel::forward_vision(const Siglip2Preprocessor::PreprocessedImage& preprocessed_image) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - auto backend = config_.default_backend == Config::Backend::CPU ? ComputeBackend::CPU : ComputeBackend::NPU; - return forward_vision(gb, preprocessed_image, backend); -} - -std::vector Siglip2VisionModel::get_image_features(const std::string& image_path) { - auto preprocessed = preprocessor_.preprocess_from_file(image_path); - return get_image_features(preprocessed); -} - -size_t Siglip2VisionModel::get_image_features_node(const Siglip2Preprocessor::PreprocessedImage& preprocessed_image) { - return forward_vision(preprocessed_image); -} - -std::vector Siglip2VisionModel::get_image_features(const Siglip2Preprocessor::PreprocessedImage& preprocessed_image) { - size_t last_hidden_state = forward_vision(preprocessed_image); - - auto* gb = static_cast(graph_handle_); - gb->execute(); - - const auto& output_buf = gb->get_output_buffer(last_hidden_state); - size_t total_elements = 1; - for (auto dim : output_buf.shape) { - total_elements *= dim; - } - - std::vector features(total_elements); - void* output_data = gb->get_output(last_hidden_state); - const float* output_ptr = static_cast(output_data); - std::copy(output_ptr, output_ptr + total_elements, features.begin()); - return features; -} - -std::vector Siglip2VisionModel::get_image_embedding(const std::string& image_path) { - auto preprocessed = preprocessor_.preprocess_from_file(image_path); - size_t last_hidden_state = forward_vision(preprocessed); - - auto* gb = static_cast(graph_handle_); - size_t pooled = gb->mean(last_hidden_state, 0); - gb->execute(); - - const auto& output_buf = gb->get_output_buffer(pooled); - size_t hidden_dim = output_buf.total_size; - - std::vector embedding(hidden_dim); - void* output_data = gb->get_output(pooled); - const float* output_ptr = static_cast(output_data); - std::copy(output_ptr, output_ptr + hidden_dim, embedding.begin()); - return embedding; -} - -size_t Siglip2VisionModel::forward(const std::vector&, bool) {return 0;} - -size_t Siglip2VisionModel::build_attention(CactusGraph*, size_t, uint32_t, - ComputeBackend, bool, size_t) { - return 0; -} - -size_t Siglip2VisionModel::build_mlp(CactusGraph*, size_t, uint32_t, - ComputeBackend) const { - return 0; -} - -size_t Siglip2VisionModel::build_transformer_block(CactusGraph*, size_t, uint32_t, - ComputeBackend, bool, size_t) { - return 0; -} - -} -} diff --git a/cactus/models/model_smol.cpp b/cactus/models/model_smol.cpp deleted file mode 100644 index e1f877b86..000000000 --- a/cactus/models/model_smol.cpp +++ /dev/null @@ -1,151 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include -#include -#include - -namespace cactus { -namespace engine { - -SmolModel::SmolModel() : Model() {} - -SmolModel::SmolModel(const Config& config) : Model(config) { - weight_nodes_.layers.resize(config.num_layers); -} - -void SmolModel::load_weights_to_graph(CactusGraph* gb) { - embedding_node_id_ = gb->mmap_embeddings(embedding_file_path_); - weight_nodes_.output_norm_weight = gb->mmap_weights(model_folder_path_ + "/output_norm.weights"); - - if (config_.tie_word_embeddings) { - weight_nodes_.output_weight = embedding_node_id_; - output_weight_node_id_ = embedding_node_id_; - } else { - weight_nodes_.output_weight = gb->mmap_weights(model_folder_path_ + "/output_weight.weights"); - output_weight_node_id_ = weight_nodes_.output_weight; - } - - for (uint32_t i = 0; i < config_.num_layers; i++) { - auto& layer = weight_nodes_.layers[i]; - std::string layer_prefix = model_folder_path_ + "/layer_" + std::to_string(i) + "_"; - layer.attn_q_weight = gb->mmap_weights(layer_prefix + "attn_q.weights"); - layer.attn_k_weight = gb->mmap_weights(layer_prefix + "attn_k.weights"); - layer.attn_v_weight = gb->mmap_weights(layer_prefix + "attn_v.weights"); - layer.attn_output_weight = gb->mmap_weights(layer_prefix + "attn_output.weights"); - layer.input_layernorm_weight = gb->mmap_weights(layer_prefix + "input_norm.weights"); - layer.ffn_gate_weight = gb->mmap_weights(layer_prefix + "ffn_gate.weights"); - layer.ffn_up_weight = gb->mmap_weights(layer_prefix + "ffn_up.weights"); - layer.ffn_down_weight = gb->mmap_weights(layer_prefix + "ffn_down.weights"); - layer.post_attention_layernorm_weight = gb->mmap_weights(layer_prefix + "post_attn_norm.weights"); - } -} - -size_t SmolModel::SmolModel::build_attention(CactusGraph* gb, size_t normalized_input, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto q_proj = gb->matmul(normalized_input, layer.attn_q_weight, true, backend); - auto k_proj = gb->matmul(normalized_input, layer.attn_k_weight, true, backend); - auto v_proj = gb->matmul(normalized_input, layer.attn_v_weight, true, backend); - - const auto& q_shape = gb->get_output_buffer(q_proj).shape; - size_t seq_len = q_shape[0]; - - auto q_proj_4d = gb->reshape(q_proj, {1, seq_len, config_.attention_heads, config_.attention_head_dim}); - auto k_proj_4d = gb->reshape(k_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - auto v_proj_4d = gb->reshape(v_proj, {1, seq_len, config_.attention_kv_heads, config_.attention_head_dim}); - - if (config_.rope_theta > 0) { - q_proj_4d = gb->rope(q_proj_4d, config_.rope_theta, position_offset); - k_proj_4d = gb->rope(k_proj_4d, config_.rope_theta, position_offset); - } - - if (use_cache) { - cache_k_output_nodes_[layer_idx] = k_proj_4d; - cache_v_output_nodes_[layer_idx] = v_proj_4d; - } - - size_t attn_output_4d; - - if (use_cache && !kv_cache_.is_empty()) { - attn_output_4d = gb->attention_int8_hybrid( - q_proj_4d, k_proj_4d, v_proj_4d, - attention_scale_, position_offset, - kv_cache_.get_keys_int8(layer_idx), - kv_cache_.get_values_int8(layer_idx), - kv_cache_.get_key_scales(layer_idx), - kv_cache_.get_value_scales(layer_idx), - kv_cache_.current_seq_len, config_.attention_kv_heads, config_.attention_head_dim - ); - } else { - attn_output_4d = gb->attention(q_proj_4d, k_proj_4d, v_proj_4d, attention_scale_, position_offset); - } - auto attn_output = gb->reshape(attn_output_4d, {seq_len, config_.attention_head_dim * config_.attention_heads}); - return gb->matmul(attn_output, layer.attn_output_weight, true, backend); -} - -size_t SmolModel::build_mlp(CactusGraph* gb, size_t normalized_h, uint32_t layer_idx, - ComputeBackend backend) const { - const auto& layer = weight_nodes_.layers[layer_idx]; - size_t gate_output = gb->matmul(normalized_h, layer.ffn_gate_weight, true, backend); - size_t up_output = gb->matmul(normalized_h, layer.ffn_up_weight, true, backend); - size_t gate_silu = gb->silu(gate_output); - size_t gated = gb->multiply(gate_silu, up_output); - return gb->matmul(gated, layer.ffn_down_weight, true, backend); -} - -size_t SmolModel::build_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, - ComputeBackend backend, bool use_cache, size_t position_offset) { - const auto& layer = weight_nodes_.layers[layer_idx]; - auto normalized_input = gb->rms_norm(hidden, layer.input_layernorm_weight, config_.layer_norm_eps); - auto attn_output = build_attention(gb, normalized_input, layer_idx, backend, use_cache, position_offset); - auto after_attention = gb->add(hidden, attn_output); - auto normalized_after_attention = gb->rms_norm(after_attention, layer.post_attention_layernorm_weight, config_.layer_norm_eps); - auto mlp_output = build_mlp(gb, normalized_after_attention, layer_idx, backend); - return gb->add(after_attention, mlp_output); -} - -size_t SmolModel::forward(const std::vector& tokens, bool use_cache) { - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized - call init() first"); - } - - if (tokens.empty()) { - throw std::runtime_error("Token sequence cannot be empty"); - } - - auto* gb = static_cast(graph_handle_); - gb->soft_reset(); - - auto seq_len = static_cast(tokens.size()); - - size_t position_offset = use_cache ? kv_cache_.get_total_seq_len() : 0; - - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - auto input_node_id = gb->input({seq_len}, Precision::FP32); - auto hidden = gb->embedding(embedding_node_id_, input_node_id); - - static std::set skip_layers = {}; - for (uint32_t layer_idx = 0; layer_idx < config_.num_layers; layer_idx++) { - if (skip_layers.count(layer_idx)) { - continue; - } - hidden = build_transformer_block(gb, hidden, layer_idx, backend, use_cache, position_offset); - } - - auto final_hidden = gb->rms_norm(hidden, weight_nodes_.output_norm_weight, config_.layer_norm_eps); - - std::vector input_data(seq_len); - for (size_t i = 0; i < seq_len; i++) { - input_data[i] = static_cast(tokens[i]); - } - gb->set_input(input_node_id, input_data.data(), Precision::FP32); - - return final_hidden; -} - -} -} \ No newline at end of file diff --git a/cactus/models/model_whisper.cpp b/cactus/models/model_whisper.cpp deleted file mode 100644 index 4dd303d51..000000000 --- a/cactus/models/model_whisper.cpp +++ /dev/null @@ -1,895 +0,0 @@ -#include "model.h" -#include "../graph/graph.h" -#include "../npu/npu.h" -#include "../kernel/kernel.h" -#include -#include -#include -#include - -namespace cactus { -namespace engine { - - -struct ConvDebugNodes { - size_t conv1; - size_t conv2; - size_t conv2_transposed; - size_t output; -}; - - -WhisperModel::WhisperModel() : Model() {} - -WhisperModel::WhisperModel(const Config& config) : Model(config) { - weight_nodes_.layers.resize(config.num_layers); - - float hd = static_cast(config.attention_head_dim); - if (hd <= 0.0f) { - hd = 64.0f; - } - - attention_scale_ = 1.0f / std::sqrt(hd); - - encoder_block_out_nodes_.resize(config.num_layers, 0); - encoder_k_nodes_.assign(config.num_layers, 0); - encoder_v_nodes_.assign(config.num_layers, 0); - -} - -void WhisperModel::load_weights_to_graph(CactusGraph* gb) { - - embedding_node_id_ = gb->mmap_embeddings(embedding_file_path_); - - weight_nodes_.decoder_norm_weight = gb->mmap_weights(model_folder_path_ + "/decoder_norm.weights"); - weight_nodes_.decoder_norm_bias = gb->mmap_weights(model_folder_path_ + "/decoder_norm.bias"); - weight_nodes_.decoder_position_embeddings_weight = gb->mmap_weights(model_folder_path_ + "/decoder_position_embeddings.weights"); - - if (config_.tie_word_embeddings) { - weight_nodes_.output_weight = embedding_node_id_; - output_weight_node_id_ = embedding_node_id_; - } else { - weight_nodes_.output_weight = gb->mmap_weights(model_folder_path_ + "/output_weight.weights"); - output_weight_node_id_ = weight_nodes_.output_weight; - } - - if (npu::is_npu_available()) { - std::string npu_encoder_path = model_folder_path_ + "/model.mlpackage"; - npu_encoder_ = npu::create_encoder(); - if (npu_encoder_ && npu_encoder_->load(npu_encoder_path)) { - use_npu_encoder_ = true; - - std::vector typical_input_shape = {1, 80, 3000}; - npu_encoder_->preallocate(typical_input_shape, "x", ""); - } else { - use_npu_encoder_ = false; - npu_encoder_.reset(); - } - } - - if (!use_npu_encoder_) { - weight_nodes_.encoder_position_embeddings = gb->mmap_weights(model_folder_path_ + "/encoder_position_embeddings.weights"); - weight_nodes_.encoder_conv1_weight = gb->mmap_weights(model_folder_path_ + "/encoder_conv1_weight.weights"); - weight_nodes_.encoder_conv1_bias = gb->mmap_weights(model_folder_path_ + "/encoder_conv1_bias.bias"); - weight_nodes_.encoder_conv2_weight = gb->mmap_weights(model_folder_path_ + "/encoder_conv2_weight.weights"); - weight_nodes_.encoder_conv2_bias = gb->mmap_weights(model_folder_path_ + "/encoder_conv2_bias.bias"); - weight_nodes_.encoder_norm_weight = gb->mmap_weights(model_folder_path_ + "/encoder_norm_weight.weights"); - weight_nodes_.encoder_norm_bias = gb->mmap_weights(model_folder_path_ + "/encoder_norm_bias.bias"); - } - - for (uint32_t i = 0; i < config_.num_layers; i++) { - auto& layer = weight_nodes_.layers[i]; - - // Decoder Layers (always needed) - std::string layer_prefix = model_folder_path_ + "/decoder.layer_" + std::to_string(i) + "_"; - - layer.decoder_encoder_attn_k_weight = gb->mmap_weights(layer_prefix + "encoder_attn_k.weights"); - layer.decoder_encoder_attn_q_weight = gb->mmap_weights(layer_prefix + "encoder_attn_q.weights"); - layer.decoder_encoder_attn_v_weight = gb->mmap_weights(layer_prefix + "encoder_attn_v.weights"); - layer.decoder_encoder_attn_output_weight = gb->mmap_weights(layer_prefix + "encoder_attn_output.weights"); - layer.decoder_encoder_attn_q_bias = gb->mmap_weights(layer_prefix + "encoder_attn_q.bias"); - layer.decoder_encoder_attn_v_bias = gb->mmap_weights(layer_prefix + "encoder_attn_v.bias"); - layer.decoder_encoder_attn_output_bias = gb->mmap_weights(layer_prefix + "encoder_attn_output.bias"); - - layer.decoder_post_encoder_layernorm_weight = gb->mmap_weights(layer_prefix + "encoder_attn_norm.weights"); - layer.decoder_post_encoder_layernorm_bias = gb->mmap_weights(layer_prefix + "encoder_attn_norm.bias"); - - layer.decoder_ffn1_weight = gb->mmap_weights(layer_prefix + "mlp_fc1.weights"); - layer.decoder_ffn1_bias = gb->mmap_weights(layer_prefix + "mlp_fc1.bias"); - layer.decoder_ffn2_weight = gb->mmap_weights(layer_prefix + "mlp_fc2.weights"); - layer.decoder_ffn2_bias = gb->mmap_weights(layer_prefix + "mlp_fc2.bias"); - - layer.decoder_post_ffn_layernorm_weight = gb->mmap_weights(layer_prefix + "final_norm.weights"); - layer.decoder_post_ffn_layernorm_bias = gb->mmap_weights(layer_prefix + "final_norm.bias"); - - layer.decoder_self_attn_k_weight = gb->mmap_weights(layer_prefix + "self_attn_k.weights"); - layer.decoder_self_attn_q_weight = gb->mmap_weights(layer_prefix + "self_attn_q.weights"); - layer.decoder_self_attn_v_weight = gb->mmap_weights(layer_prefix + "self_attn_v.weights"); - layer.decoder_self_attn_output_weight = gb->mmap_weights(layer_prefix + "self_attn_output.weights"); - layer.decoder_self_attn_q_bias = gb->mmap_weights(layer_prefix + "self_attn_q.bias"); - layer.decoder_self_attn_v_bias = gb->mmap_weights(layer_prefix + "self_attn_v.bias"); - layer.decoder_self_attn_output_bias = gb->mmap_weights(layer_prefix + "self_attn_output.bias"); - - layer.decoder_post_attn_layernorm_weight = gb->mmap_weights(layer_prefix + "self_attn_norm.weights"); - layer.decoder_post_attn_layernorm_bias = gb->mmap_weights(layer_prefix + "self_attn_norm.bias"); - - if (!use_npu_encoder_) { - layer_prefix = model_folder_path_ + "/encoder.layer_" + std::to_string(i) + "_"; - - layer.encoder_ffn1_weight = gb->mmap_weights(layer_prefix + "mlp_fc1.weights"); - layer.encoder_ffn1_bias = gb->mmap_weights(layer_prefix + "mlp_fc1.bias"); - layer.encoder_ffn2_weight = gb->mmap_weights(layer_prefix + "mlp_fc2.weights"); - layer.encoder_ffn2_bias = gb->mmap_weights(layer_prefix + "mlp_fc2.bias"); - - layer.encoder_post_ffn_layernorm_weight = gb->mmap_weights(layer_prefix + "final_norm.weights"); - layer.encoder_post_ffn_layernorm_bias = gb->mmap_weights(layer_prefix + "final_norm.bias"); - - layer.encoder_self_attn_k_weight = gb->mmap_weights(layer_prefix + "self_attn_k.weights"); - layer.encoder_self_attn_q_weight = gb->mmap_weights(layer_prefix + "self_attn_q.weights"); - layer.encoder_self_attn_v_weight = gb->mmap_weights(layer_prefix + "self_attn_v.weights"); - layer.encoder_self_attn_output_weight = gb->mmap_weights(layer_prefix + "self_attn_output.weights"); - layer.encoder_self_attn_q_bias = gb->mmap_weights(layer_prefix + "self_attn_q.bias"); - layer.encoder_self_attn_v_bias = gb->mmap_weights(layer_prefix + "self_attn_v.bias"); - layer.encoder_self_attn_output_bias = gb->mmap_weights(layer_prefix + "self_attn_output.bias"); - - layer.encoder_post_attn_layernorm_weight = gb->mmap_weights(layer_prefix + "self_attn_norm.weights"); - layer.encoder_post_attn_layernorm_bias = gb->mmap_weights(layer_prefix + "self_attn_norm.bias"); - } - } -} - -size_t WhisperModel::build_encoder_mlp(CactusGraph* gb, size_t input, uint32_t layer_idx, ComputeBackend backend) { - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto ffn1_weight = gb->matmul(input, layer.encoder_ffn1_weight, true, backend); - auto ffn1_bias = gb->add(ffn1_weight, layer.encoder_ffn1_bias); - - encoder_pre_gelu = ffn1_bias; - - auto ffn1_act = gb->gelu_erf(ffn1_bias); - - encoder_post_gelu = ffn1_act; - - auto ffn2_weight = gb->matmul(ffn1_act, layer.encoder_ffn2_weight, true, backend); - auto ffn2_bias = gb->add(ffn2_weight, layer.encoder_ffn2_bias); - return ffn2_bias; -} - -size_t WhisperModel::build_decoder_mlp(CactusGraph* gb, size_t input, uint32_t layer_idx, ComputeBackend backend) const { - const auto& layer = weight_nodes_.layers[layer_idx]; - auto ffn1_weight = gb->matmul(input, layer.decoder_ffn1_weight, true, backend); - auto ffn1_bias = gb->add(ffn1_weight, layer.decoder_ffn1_bias); - auto ffn1_act = gb->gelu_erf(ffn1_bias); - auto ffn2_weight = gb->matmul(ffn1_act, layer.decoder_ffn2_weight, true, backend); - auto ffn2_bias = gb->add(ffn2_weight, layer.decoder_ffn2_bias); - return ffn2_bias; -} - - - -size_t WhisperModel::build_encoder_attention(CactusGraph* gb, size_t input, uint32_t layer_idx, ComputeBackend backend, bool use_cache, size_t /*position_offset*/){ - - const auto& layer = weight_nodes_.layers[layer_idx]; - - size_t q = gb->matmul(input, layer.decoder_encoder_attn_q_weight, true, backend); - q = gb->add(q, layer.decoder_encoder_attn_q_bias); - - const auto& q_buf = gb->get_output_buffer(q); - if (q_buf.shape.size() != 2) { - throw std::runtime_error("encoder cross-attn: q must be [T_dec, D]"); - } - - size_t T_dec = q_buf.shape[0]; - size_t q_heads = config_.attention_heads; - size_t kv_heads = config_.attention_kv_heads; - size_t head_dim = config_.attention_head_dim; - - q = gb->reshape(q, {1, T_dec, q_heads, head_dim}); - - size_t k_4d = 0; - size_t v_4d = 0; - - if (use_cache && encoder_kv_ready_) { - const auto& k_shape = encoder_k_shape_[layer_idx]; - const auto& v_shape = encoder_v_shape_[layer_idx]; - - size_t cache_k_node = gb->input(k_shape, encoder_kv_precision_); - size_t cache_v_node = gb->input(v_shape, encoder_kv_precision_); - - gb->set_input(cache_k_node, encoder_k_host_[layer_idx].data(), encoder_kv_precision_); - gb->set_input(cache_v_node, encoder_v_host_[layer_idx].data(), encoder_kv_precision_); - - k_4d = cache_k_node; - v_4d = cache_v_node; - } else { - size_t enc_norm = weight_nodes_.encoder_output; - - size_t k = gb->matmul(enc_norm, layer.decoder_encoder_attn_k_weight, true, backend); - size_t v = gb->matmul(enc_norm, layer.decoder_encoder_attn_v_weight, true, backend); - v = gb->add(v, layer.decoder_encoder_attn_v_bias); - - const auto& k_buf = gb->get_output_buffer(k); - if (k_buf.shape.size() != 2) { - throw std::runtime_error("encoder cross-attn: k must be [T_enc, D]"); - } - size_t T_enc = k_buf.shape[0]; - - k_4d = gb->reshape(k, {1, T_enc, kv_heads, head_dim}); - v_4d = gb->reshape(v, {1, T_enc, kv_heads, head_dim}); - - if (!encoder_kv_ready_) { - encoder_k_nodes_[layer_idx] = k_4d; - encoder_v_nodes_[layer_idx] = v_4d; - } - } - - size_t attn = gb->attention(q, k_4d, v_4d, attention_scale_, false); - - attn = gb->reshape(attn, {T_dec, q_heads * head_dim}); - size_t out = gb->matmul(attn, layer.decoder_encoder_attn_output_weight, true, backend); - out = gb->add(out, layer.decoder_encoder_attn_output_bias); - - return out; -} - -void WhisperModel::reset_graph_side_cache_nodes() { - cache_k_output_nodes_.assign(config_.num_layers, 0); - cache_v_output_nodes_.assign(config_.num_layers, 0); -} - -void WhisperModel::reset_cache() { - Model::reset_cache(); - encoder_ready_ = false; - encoder_kv_ready_ = false; - first_decode_step_ = true; - encoder_output_host_.clear(); - encoder_k_host_.clear(); - encoder_v_host_.clear(); - encoder_k_shape_.clear(); - encoder_v_shape_.clear(); -} - -size_t WhisperModel::build_decoder_self_attention(CactusGraph* gb, size_t input, uint32_t layer_idx, ComputeBackend backend, bool use_cache, size_t position_offset){ - const auto& layer = weight_nodes_.layers[layer_idx]; - - auto q = gb->matmul(input, layer.decoder_self_attn_q_weight, true, backend); - q = gb->add(q, layer.decoder_self_attn_q_bias); - - auto k = gb->matmul(input, layer.decoder_self_attn_k_weight, true, backend); - auto v = gb->matmul(input, layer.decoder_self_attn_v_weight, true, backend); - v = gb->add(v, layer.decoder_self_attn_v_bias); - - const auto& q_shape = gb->get_output_buffer(q).shape; - if (q_shape.size() != 2) { - throw std::runtime_error("decoder self-attn: q must be [T_new, D]"); - } - - size_t seq_new = q_shape[0]; - size_t num_heads = config_.attention_heads; - size_t head_dim = config_.attention_head_dim; - size_t num_kv_heads = config_.attention_kv_heads; - - auto q_4d = gb->reshape(q, {1, seq_new, num_heads, head_dim}); - auto k_4d = gb->reshape(k, {1, seq_new, num_kv_heads, head_dim}); - auto v_4d = gb->reshape(v, {1, seq_new, num_kv_heads, head_dim}); - - size_t final_k = k_4d; - size_t final_v = v_4d; - - if (use_cache && !kv_cache_.is_empty()) { - auto k_view = kv_cache_.get_key_view(layer_idx); - auto v_view = kv_cache_.get_value_view(layer_idx); - - if (!k_view.ptr1 || !v_view.ptr1) { - throw std::runtime_error("KV cache view is empty but kv_cache_.is_empty()==false"); - } - - size_t cache_len = kv_cache_.current_seq_len; - - size_t cache_k_node = gb->input( - {1, cache_len, num_kv_heads, head_dim}, - kv_cache_.precision - ); - size_t cache_v_node = gb->input( - {1, cache_len, num_kv_heads, head_dim}, - kv_cache_.precision - ); - - if (k_view.ptr2 == nullptr && v_view.ptr2 == nullptr) { - gb->set_input(cache_k_node, k_view.ptr1, kv_cache_.precision); - gb->set_input(cache_v_node, v_view.ptr1, kv_cache_.precision); - } else { - gb->set_input(cache_k_node, kv_cache_.get_key_ptr(layer_idx), kv_cache_.precision); - gb->set_input(cache_v_node, kv_cache_.get_value_ptr(layer_idx), kv_cache_.precision); - } - - final_k = gb->concat(cache_k_node, k_4d, 1); - final_v = gb->concat(cache_v_node, v_4d, 1); - } - - if (use_cache) { - cache_k_output_nodes_[layer_idx] = final_k; - cache_v_output_nodes_[layer_idx] = final_v; - } else { - cache_k_output_nodes_[layer_idx] = k_4d; - cache_v_output_nodes_[layer_idx] = v_4d; - } - - auto attn_out_4d = gb->attention(q_4d, final_k, final_v, attention_scale_, position_offset); - auto attn_out = gb->reshape(attn_out_4d, {seq_new, num_heads * head_dim}); - - auto output = gb->matmul(attn_out, layer.decoder_self_attn_output_weight, true, backend); - output = gb->add(output, layer.decoder_self_attn_output_bias); - return output; -} - -size_t WhisperModel::build_encoder_self_attention(CactusGraph* gb, size_t input, uint32_t layer_idx, ComputeBackend backend, bool use_cache, size_t /*position_offset*/){ - const auto& layer = weight_nodes_.layers[layer_idx]; - - if(use_cache) - throw std::runtime_error("The encoder attention layers are not auto-regressive, and thus don't use KV caching!"); - - auto q = gb->matmul(input, layer.encoder_self_attn_q_weight, true, backend); - q = gb->add(q, layer.encoder_self_attn_q_bias); - auto v = gb->matmul(input, layer.encoder_self_attn_v_weight, true, backend); - v = gb->add(v, layer.encoder_self_attn_v_bias); - auto k = gb->matmul(input, layer.encoder_self_attn_k_weight, true, backend); - - size_t seq_len = gb->get_output_buffer(q).shape[0]; - size_t num_heads = config_.attention_heads; - size_t head_dim = config_.attention_head_dim; - - q = gb->reshape(q, {1, seq_len, num_heads, head_dim}); - k = gb->reshape(k, {1, seq_len, num_heads, head_dim}); - v = gb->reshape(v, {1, seq_len, num_heads, head_dim}); - - auto attn = gb->attention(q, k, v, attention_scale_, false); - - attn = gb->reshape(attn, {seq_len, num_heads * head_dim}); - - auto output = gb->matmul(attn, layer.encoder_self_attn_output_weight, true, backend); - output = gb->add(output, layer.encoder_self_attn_output_bias); - - return output; -} - -size_t WhisperModel::build_conv1d(CactusGraph* gb, size_t input) -{ - size_t conv_input = input; - const auto& xbuf = gb->get_output_buffer(input); - - if (xbuf.precision == Precision::INT8) { - conv_input = gb->precision_cast(input, Precision::FP16); - } - - size_t conv1 = gb->conv1d_k3(conv_input, weight_nodes_.encoder_conv1_weight, 1); - - auto bias1_shape = gb->get_output_buffer(weight_nodes_.encoder_conv1_bias).shape; - size_t C1 = bias1_shape[0]; - size_t bias1 = gb->reshape(weight_nodes_.encoder_conv1_bias, {1, C1, 1}); - conv1 = gb->add(conv1, bias1); - - last_conv1_node_ = conv1; - - conv1 = gb->gelu_erf(conv1); - - size_t conv2 = gb->conv1d_k3(conv1, weight_nodes_.encoder_conv2_weight, 2); - auto bias2_shape = gb->get_output_buffer(weight_nodes_.encoder_conv2_bias).shape; - size_t C2 = bias2_shape[0]; - size_t bias2 = gb->reshape(weight_nodes_.encoder_conv2_bias, {1, C2, 1}); - conv2 = gb->add(conv2, bias2); - - last_conv2_node_ = conv2; - - conv2 = gb->gelu_erf(conv2); - - const auto& buf = gb->get_output_buffer(conv2); - - size_t conv2_transposed; - if (buf.precision == Precision::FP16) { - conv2_transposed = gb->transpose(conv2, ComputeBackend::CPU); - } else { - size_t conv2_f16 = gb->precision_cast(conv2, Precision::FP16); - conv2_transposed = gb->transpose(conv2_f16, ComputeBackend::CPU); - } - - return conv2_transposed; -} - -size_t WhisperModel::build_encoder_transformer_block( - CactusGraph* gb, - size_t hidden, - uint32_t layer_idx, - ComputeBackend backend, - bool use_cache, - size_t position_offset) -{ - const auto& layer = weight_nodes_.layers[layer_idx]; - - size_t ln1 = gb->layernorm( - hidden, - layer.encoder_post_attn_layernorm_weight, - layer.encoder_post_attn_layernorm_bias - ); - - size_t sa = build_encoder_self_attention( - gb, ln1, layer_idx, backend, use_cache, position_offset - ); - - size_t x_post_sa = gb->add(hidden, sa); - - size_t ln2 = gb->layernorm( - x_post_sa, - layer.encoder_post_ffn_layernorm_weight, - layer.encoder_post_ffn_layernorm_bias - ); - - size_t ffn_out = build_encoder_mlp( - gb, ln2, layer_idx, backend - ); - - size_t out = gb->add(x_post_sa, ffn_out); - - if (layer_idx < encoder_block_out_nodes_.size()) { - encoder_block_out_nodes_[layer_idx] = out; - } - - return out; -} - -size_t WhisperModel::build_decoder_transformer_block(CactusGraph* gb, size_t hidden, uint32_t layer_idx, ComputeBackend backend, bool use_cache, size_t position_offset){ - const auto& layer = weight_nodes_.layers[layer_idx]; - - size_t ln1 = gb->layernorm(hidden, layer.decoder_post_attn_layernorm_weight, layer.decoder_post_attn_layernorm_bias); - size_t sa = build_decoder_self_attention(gb, ln1, layer_idx, backend, use_cache, position_offset); - size_t x_post_sa = gb->add(hidden, sa); - - size_t ln2 = gb->layernorm(x_post_sa, layer.decoder_post_encoder_layernorm_weight, layer.decoder_post_encoder_layernorm_bias); - size_t ca = build_encoder_attention(gb, ln2, layer_idx, backend, use_cache, position_offset); - size_t x_post_ca = gb->add(x_post_sa, ca); - - size_t ln3 = gb->layernorm(x_post_ca,layer.decoder_post_ffn_layernorm_weight,layer.decoder_post_ffn_layernorm_bias); - size_t ffn_out = build_decoder_mlp(gb, ln3, layer_idx, backend); - size_t x_post_ffn = gb->add(x_post_ca, ffn_out); - - return x_post_ffn; - -} - -void WhisperModel::run_encoder(const std::vector& mel_bins) -{ - if (mel_bins.size() % 80 != 0) - throw std::runtime_error("Mel bins length must be divisible by 80."); - - size_t T_mel = mel_bins.size() / 80; - if (T_mel == 0) - throw std::runtime_error("Mel bins has zero frames."); - - auto* gb = static_cast(graph_handle_); - if (!gb) - throw std::runtime_error("Graph handle is null in run_encoder."); - - if (use_npu_encoder_ && npu_encoder_ && npu_encoder_->is_available()) { - std::vector out_shape = npu_encoder_->get_output_shape(); - size_t T_enc, D_enc; - if (out_shape.size() == 3) { - T_enc = static_cast(out_shape[1]); - D_enc = static_cast(out_shape[2]); - } else if (out_shape.size() == 2) { - T_enc = static_cast(out_shape[0]); - D_enc = static_cast(out_shape[1]); - } else { - throw std::runtime_error("NPU encoder output has unexpected shape"); - } - - std::vector<__fp16> mel_bins_f16(mel_bins.size()); - cactus_fp32_to_fp16(mel_bins.data(), mel_bins_f16.data(), mel_bins.size()); - - std::vector input_shape = {1, 80, static_cast(T_mel)}; - - __fp16* output_buffer = npu_encoder_->get_output_buffer(); - if (output_buffer) { - size_t elements_written = npu_encoder_->encode( - mel_bins_f16.data(), - output_buffer, - input_shape, - "x", - "" - ); - - if (elements_written > 0) { - size_t enc_output_node = gb->input({T_enc, D_enc}, Precision::FP16); - gb->set_input(enc_output_node, output_buffer, Precision::FP16); - - weight_nodes_.encoder_output = enc_output_node; - return; - } - } else { - std::vector<__fp16> npu_output(T_enc * D_enc); - size_t elements_written = npu_encoder_->encode( - mel_bins_f16.data(), - npu_output.data(), - input_shape, - "x", - "" - ); - - if (elements_written > 0) { - size_t enc_output_node = gb->input({T_enc, D_enc}, Precision::FP16); - gb->set_input(enc_output_node, npu_output.data(), Precision::FP16); - - weight_nodes_.encoder_output = enc_output_node; - return; - } - } - } - - auto backend = - (config_.default_backend == Config::Backend::CPU) - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - size_t mel_input = 0; - std::vector<__fp16> mel_bins_f16(mel_bins.size()); - cactus_fp32_to_fp16(mel_bins.data(), mel_bins_f16.data(), mel_bins.size()); - - mel_input = gb->input({1, 80, T_mel}, Precision::FP16); - gb->set_input(mel_input, mel_bins_f16.data(), Precision::FP16); - - size_t conv2_transposed = build_conv1d(gb, mel_input); - - const auto& conv_shape = gb->get_output_buffer(conv2_transposed).shape; - if (conv_shape.size() != 3 || conv_shape[0] != 1) - throw std::runtime_error("Conv2 transpose should be [1, T_enc, D]."); - - size_t T_enc = conv_shape[1]; - size_t D_enc = conv_shape[2]; - - size_t pos_slice = gb->slice(weight_nodes_.encoder_position_embeddings, 0, 0, T_enc); - - size_t h2d = gb->reshape(conv2_transposed, {T_enc, D_enc}); - - auto& h2d_buf = gb->get_output_buffer(h2d); - auto& pos_buf = gb->get_output_buffer(pos_slice); - - if (pos_buf.precision != h2d_buf.precision) { - pos_slice = gb->precision_cast(pos_slice, h2d_buf.precision); - } - - size_t h_pos = gb->add(h2d, pos_slice); - last_enc_plus_pos_node_ = h_pos; - - size_t h = h_pos; - for (uint32_t i = 0; i < config_.num_layers; ++i){ - h = build_encoder_transformer_block(gb, h, i, backend, false, 0); - if (i == 0) { - encoder_transformer_block_0 = h; - } - } - - size_t h_norm = gb->layernorm( - h, - weight_nodes_.encoder_norm_weight, - weight_nodes_.encoder_norm_bias - ); - last_encoder_post_norm_node_ = h_norm; - - - weight_nodes_.encoder_output = h_norm; -} - - - -size_t WhisperModel::run_decoder_step(const std::vector& tokens, bool use_cache, bool last_token_only) { - auto* gb = static_cast(graph_handle_); - - const size_t full_len = tokens.size(); - if (full_len == 0) { - throw std::runtime_error("Decoder token list cannot be empty."); - } - - auto backend = config_.default_backend == Config::Backend::CPU - ? ComputeBackend::CPU - : ComputeBackend::NPU; - - size_t start_idx = (use_cache && kv_cache_.current_seq_len > 0) ? full_len - 1 : 0; - size_t new_tokens = full_len - start_idx; - - size_t tok_input = gb->input({new_tokens}, Precision::FP32); - std::vector tok_f(new_tokens); - for (size_t i = 0; i < new_tokens; i++) { - tok_f[i] = static_cast(tokens[start_idx + i]); - } - gb->set_input(tok_input, tok_f.data(), Precision::FP32); - - size_t dec_hidden = gb->embedding(embedding_node_id_, tok_input); - - size_t position_offset = kv_cache_.current_seq_len; - size_t dec_pos = gb->slice(weight_nodes_.decoder_position_embeddings_weight,0,position_offset,new_tokens); - - { - const auto& h_buf = gb->get_output_buffer(dec_hidden); - const auto& pos_buf = gb->get_output_buffer(dec_pos); - - size_t pos_node_for_add = dec_pos; - if (pos_buf.precision != h_buf.precision) { - pos_node_for_add = gb->precision_cast(dec_pos, h_buf.precision); - } - - dec_hidden = gb->add(dec_hidden, pos_node_for_add); - } - - for (uint32_t layer_idx = 0; layer_idx < config_.num_layers; ++layer_idx) { - dec_hidden = build_decoder_transformer_block( - gb, - dec_hidden, - layer_idx, - backend, - use_cache, - position_offset - ); - } - - size_t dec_norm = gb->layernorm( - dec_hidden, - weight_nodes_.decoder_norm_weight, - weight_nodes_.decoder_norm_bias - ); - - size_t logits_input = dec_norm; - if (last_token_only) { - size_t row_index = new_tokens - 1; - logits_input = gb->slice(logits_input, 0, row_index, 1); - } - - auto w_shape = gb->get_output_buffer(output_weight_node_id_).shape; - - size_t logits = gb->matmul(logits_input, output_weight_node_id_, true, backend); - - last_new_tokens_ = new_tokens; - return logits; -} - - -size_t WhisperModel::forward(const std::vector& mel_bins, const std::vector& tokens, bool use_cache) -{ - - if (!initialized_ || !graph_handle_) { - throw std::runtime_error("Model not initialized"); - } - - if (!use_cache) { - kv_cache_.reset(); - kv_cache_.current_seq_len = 0; - reset_graph_side_cache_nodes(); - run_encoder(mel_bins); - } - - return run_decoder_step(tokens, use_cache, false); -} - -std::vector WhisperModel::get_audio_embeddings(const std::vector& mel_bins) { - run_encoder(mel_bins); - - auto* gb = static_cast(graph_handle_); - - size_t pooled = gb->mean(weight_nodes_.encoder_output, 0); - gb->execute(); - - const auto& output_buf = gb->get_output_buffer(pooled); - size_t hidden_dim = output_buf.total_size; - - std::vector embedding(hidden_dim); - void* output_data = gb->get_output(pooled); - const float* output_ptr = static_cast(output_data); - std::copy(output_ptr, output_ptr + hidden_dim, embedding.begin()); - - reset_cache(); - return embedding; -} - -uint32_t WhisperModel::decode_with_audio( - const std::vector& tokens, - const std::vector& mel_bins, - float temperature, - float top_p, - size_t top_k, - const std::string& profile_file, - float* out_entropy) -{ - if (!initialized_ || !graph_handle_) - throw std::runtime_error("Model not initialized - call init() first"); - if (tokens.empty()) - throw std::runtime_error("Token sequence cannot be empty"); - if (mel_bins.empty()) - throw std::runtime_error("Mel bins cannot be empty in Whisper decode_with_audio"); - - auto* gb = static_cast(graph_handle_); - - bool cold_start = !encoder_ready_; - size_t logits_node = 0; - - uint32_t bos = static_cast(get_tokenizer()->get_bos_token()); - - std::vector full_tokens; - full_tokens.reserve(tokens.size() + 1); - full_tokens.push_back(bos); - full_tokens.insert(full_tokens.end(), tokens.begin(), tokens.end()); - - if (cold_start) - { - gb->soft_reset(); - kv_cache_.reset(); - kv_cache_.current_seq_len = 0; - reset_graph_side_cache_nodes(); - - encoder_kv_ready_ = false; - encoder_k_nodes_.assign(config_.num_layers, 0); - encoder_v_nodes_.assign(config_.num_layers, 0); - encoder_k_host_.clear(); - encoder_v_host_.clear(); - encoder_k_shape_.clear(); - encoder_v_shape_.clear(); - - first_decode_step_ = true; - run_encoder(mel_bins); - logits_node = run_decoder_step(full_tokens, false, false); - } - - else - { - gb->soft_reset(); - reset_graph_side_cache_nodes(); - - if (encoder_output_host_.empty()) - throw std::runtime_error("Missing encoder_output_host_ in warm step!"); - - size_t enc_node = gb->input(encoder_output_shape_, encoder_output_precision_); - gb->set_input(enc_node, encoder_output_host_.data(), encoder_output_precision_); - weight_nodes_.encoder_output = enc_node; - - std::vector last_token_vec = { tokens.back() }; - logits_node = run_decoder_step(last_token_vec, true, true); - } - - size_t sampled_token_id = gb->sample(logits_node, temperature, top_p, top_k); - if (!profile_file.empty()) gb->execute(profile_file); - else gb->execute(); - - if (cold_start) - { - auto& out_buf = gb->get_output_buffer(weight_nodes_.encoder_output); - - encoder_output_shape_ = out_buf.shape; - encoder_output_precision_ = out_buf.precision; - - size_t total_elems = 1; - for (auto s : out_buf.shape) - total_elems *= s; - - size_t elem_size = 0; - switch (out_buf.precision) { - case Precision::FP32: elem_size = sizeof(float); break; - case Precision::FP16: elem_size = sizeof(uint16_t); break; - case Precision::INT8: elem_size = sizeof(int8_t); break; - default: - throw std::runtime_error("Unsupported encoder_output precision in WhisperModel"); - } - - const size_t total_bytes = total_elems * elem_size; - - encoder_output_host_.resize(total_bytes); - std::memcpy( - encoder_output_host_.data(), - gb->get_output(weight_nodes_.encoder_output), - total_bytes - ); - - { - if (config_.num_layers == 0) { - throw std::runtime_error("WhisperModel: num_layers is zero?"); - } - - auto& k0_buf = gb->get_output_buffer(encoder_k_nodes_[0]); - encoder_kv_precision_ = k0_buf.precision; - - encoder_k_host_.resize(config_.num_layers); - encoder_v_host_.resize(config_.num_layers); - encoder_k_shape_.resize(config_.num_layers); - encoder_v_shape_.resize(config_.num_layers); - - size_t kv_elem_size = 0; - switch (encoder_kv_precision_) { - case Precision::FP32: kv_elem_size = sizeof(float); break; - case Precision::FP16: kv_elem_size = sizeof(uint16_t); break; - case Precision::INT8: kv_elem_size = sizeof(int8_t); break; - default: - throw std::runtime_error("Unsupported encoder K/V precision in WhisperModel"); - } - - for (uint32_t i = 0; i < config_.num_layers; ++i) { - size_t k_node = encoder_k_nodes_[i]; - size_t v_node = encoder_v_nodes_[i]; - - auto& k_buf = gb->get_output_buffer(k_node); - auto& v_buf = gb->get_output_buffer(v_node); - - encoder_k_shape_[i] = k_buf.shape; - encoder_v_shape_[i] = v_buf.shape; - - size_t k_elems = 1; - for (auto s : k_buf.shape) k_elems *= s; - size_t v_elems = 1; - for (auto s : v_buf.shape) v_elems *= s; - - encoder_k_host_[i].resize(k_elems * kv_elem_size); - encoder_v_host_[i].resize(v_elems * kv_elem_size); - - std::memcpy( - encoder_k_host_[i].data(), - gb->get_output(k_node), - k_elems * kv_elem_size - ); - std::memcpy( - encoder_v_host_[i].data(), - gb->get_output(v_node), - v_elems * kv_elem_size - ); - } - - encoder_kv_ready_ = true; - } - - encoder_ready_ = true; - } - - - if (out_entropy) { - const auto& logits_buf = gb->get_output_buffer(logits_node); - void* logits_ptr = gb->get_output(logits_node); - size_t vocab_size = logits_buf.shape.back(); - - std::vector logits(vocab_size); - if (logits_buf.precision == Precision::FP32) { - float* src = static_cast(logits_ptr); - std::copy(src, src + vocab_size, logits.begin()); - } else if (logits_buf.precision == Precision::FP16) { - __fp16* src = static_cast<__fp16*>(logits_ptr); - Quantization::fp16_to_fp32(src, logits.data(), vocab_size); - } else { - int8_t* src = static_cast(logits_ptr); - Quantization::int8_to_fp32(src, logits.data(), vocab_size, 1.0f); - } - - float max_logit = *std::max_element(logits.begin(), logits.end()); - double sum_exp = 0.0; - for (size_t i = 0; i < vocab_size; ++i) { - sum_exp += std::exp(static_cast(logits[i] - max_logit)); - } - double log_sum_exp = static_cast(max_logit) + std::log(sum_exp); - - double entropy = 0.0; - for (size_t i = 0; i < vocab_size; ++i) { - double log_prob = static_cast(logits[i]) - log_sum_exp; - double prob = std::exp(log_prob); - if (prob > 1e-10) { - entropy -= prob * log_prob; - } - } - - double max_entropy = std::log(static_cast(vocab_size)); - *out_entropy = static_cast(entropy / max_entropy); - } - - post_execute_updates(gb, full_tokens.size()); - update_kv_cache(gb, last_new_tokens_); - - auto* out_ptr = gb->get_output(sampled_token_id); - uint32_t sampled = *reinterpret_cast(out_ptr); - - return sampled; -} - - -} -} diff --git a/cactus/npu/npu.cpp b/cactus/npu/npu.cpp deleted file mode 100644 index 12b65f4a3..000000000 --- a/cactus/npu/npu.cpp +++ /dev/null @@ -1,69 +0,0 @@ - -#include "npu.h" -#include -#include -#if defined(__APPLE__) -#include -#endif -#if defined(__APPLE__) && (TARGET_OS_IPHONE) - const char* lib_name = "@rpath/cactus_util.framework/cactus_util"; -#else - const char* lib_name = nullptr; -#endif - -namespace cactus { -namespace npu { - -typedef std::unique_ptr (*create_encoder_fn)(); -typedef bool (*is_npu_available_fn)(); -typedef std::unique_ptr (*create_prefill_fn)(); - -static void* g_cactus_util_handle = nullptr; - -template -T get_runtime_symbol(const char* symbol_name) { - if (!lib_name) { - return nullptr; - } - void* handle = dlopen(lib_name, RTLD_NOW | RTLD_GLOBAL); - if (handle) { - g_cactus_util_handle = handle; - } - if (g_cactus_util_handle) { - void* sym = dlsym(g_cactus_util_handle, symbol_name); - if (sym) { - return reinterpret_cast(sym); - } - } - return nullptr; -} - -__attribute__((weak, visibility("default"))) -std::unique_ptr create_encoder() { - auto strong_fn = get_runtime_symbol("_ZN6cactus3npu14create_encoderEv"); - if (strong_fn) { - return strong_fn(); - } - return nullptr; -} - -__attribute__((weak, visibility("default"))) -std::unique_ptr create_prefill() { - auto strong_fn = get_runtime_symbol("_ZN6cactus3npu14create_prefillEv"); - if (strong_fn) { - return strong_fn(); - } - return nullptr; -} - -__attribute__((weak, visibility("default"))) -bool is_npu_available() { - auto strong_fn = get_runtime_symbol("_ZN6cactus3npu16is_npu_availableEv"); - if (strong_fn) { - return strong_fn(); - } - return false; -} - -} // namespace npu -} // namespace cactus \ No newline at end of file diff --git a/cactus/npu/npu.h b/cactus/npu/npu.h deleted file mode 100644 index a150685ce..000000000 --- a/cactus/npu/npu.h +++ /dev/null @@ -1,76 +0,0 @@ -#ifndef CACTUS_NPU_H -#define CACTUS_NPU_H - -#include -#include -#include -#include - -namespace cactus { -namespace npu { - -class NPUEncoder { -public: - virtual ~NPUEncoder() = default; - - virtual bool load(const std::string& model_path) = 0; - - virtual bool preallocate(const std::vector& input_shape, - const std::string& input_name = "x", - const std::string& output_name = "") = 0; - - virtual size_t encode(const __fp16* input, - __fp16* output, - const std::vector& shape, - const std::string& input_name = "x", - const std::string& output_name = "") = 0; - - virtual bool is_available() const = 0; - - virtual std::vector get_input_shape() const = 0; - - virtual std::vector get_output_shape() const = 0; - - virtual __fp16* get_output_buffer() = 0; - - virtual size_t get_output_buffer_size() const = 0; -}; - -struct NPUBufferRef { - const __fp16* data; - size_t count; -}; - -struct NPUPrefillDirectResult { - NPUBufferRef hidden; - std::vector k_caches; - std::vector v_caches; - bool valid; -}; - -class NPUPrefill { -public: - virtual ~NPUPrefill() = default; - virtual bool load(const std::string& model_path) = 0; - virtual bool is_available() const = 0; - virtual int get_chunk_size() const = 0; - virtual int get_hidden_dim() const = 0; - virtual int get_num_layers() const = 0; - virtual int get_num_kv_heads() const = 0; - virtual int get_head_dim() const = 0; - - virtual NPUPrefillDirectResult prefill_chunk_direct( - const std::vector<__fp16>& embeddings, - int position_offset = 0, - const std::string& input_name = "x") = 0; -}; - - -std::unique_ptr create_encoder(); -std::unique_ptr create_prefill(); -bool is_npu_available(); - -} // namespace npu -} // namespace cactus - -#endif // CACTUS_NPU_H \ No newline at end of file diff --git a/docs/cactus_engine.md b/docs/cactus_engine.md index 45062bc9b..3e16d8ec0 100644 --- a/docs/cactus_engine.md +++ b/docs/cactus_engine.md @@ -1,3 +1,9 @@ +--- +title: "Cactus Engine FFI API Reference" +description: "C API documentation for Cactus on-device AI inference engine. Supports text completion, vision, transcription, embeddings, RAG, tool calling, and cloud handoff." +keywords: ["on-device AI", "mobile inference", "LLM API", "C FFI", "edge AI", "transcription", "embeddings", "RAG", "tool calling"] +--- + # Cactus Engine FFI Documentation The Cactus Engine provides a clean C FFI (Foreign Function Interface) for integrating the LLM inference engine into various applications. This documentation covers all available functions, their parameters, and usage examples. @@ -11,9 +17,19 @@ Before using the Cactus Engine, you need to download model weights: cactus download LiquidAI/LFM2-1.2B cactus download LiquidAI/LFM2-VL-450M cactus download openai/whisper-small + +# Optional: set your Cactus Cloud API key for automatic cloud fallback +cactus auth ``` -Weights are saved to the `weights/` directory and can be loaded using `cactus_init()`. +`cactus download` fetches a **pre-built runtime bundle** (CQ weights + serialized +graph + manifest) from +[huggingface.co/Cactus-Compute](https://huggingface.co/Cactus-Compute) into +`weights/-cq/`. The result can be loaded directly via +`cactus_init()`. + +For models not on Cactus-Compute, build a bundle from source with +`cactus convert ` (quantizes the weights and builds the runtime graph). ## Types @@ -24,6 +40,13 @@ An opaque pointer type representing a loaded model instance. This handle is used typedef void* cactus_model_t; ``` +### `cactus_index_t` +An opaque pointer type representing a vector index instance. + +```c +typedef void* cactus_index_t; +``` + ### `cactus_token_callback` Callback function type for streaming token generation. Called for each generated token during completion. @@ -35,6 +58,13 @@ typedef void (*cactus_token_callback)( ); ``` +### `cactus_log_callback_t` +Callback function type for log messages. Installed via `cactus_log_set_callback`. + +```c +typedef void (*cactus_log_callback_t)(int level, const char* component, const char* message, void* user_data); +``` + ## Core Functions ### `cactus_init` @@ -43,7 +73,8 @@ Initializes a model from disk and prepares it for inference. ```c cactus_model_t cactus_init( const char* model_path, // Path to the model directory - const char* corpus_dir // Optional path to corpus directory for RAG (can be NULL) + const char* corpus_dir, // Optional path to corpus directory for RAG (can be NULL) + bool cache_index // false = always rebuild index, true = load cached if available ); ``` @@ -51,14 +82,14 @@ cactus_model_t cactus_init( **Example:** ```c -cactus_model_t model = cactus_init("../../weights/qwen3-600m", NULL); +cactus_model_t model = cactus_init("../../weights/qwen3-600m", NULL, false); if (!model) { fprintf(stderr, "Failed to initialize model\n"); return -1; } // with RAG corpus -cactus_model_t rag_model = cactus_init("../../weights/lfm2-rag", "./documents"); +cactus_model_t rag_model = cactus_init("../../weights/lfm2-rag", "./documents", true); ``` ### `cactus_complete` @@ -73,7 +104,9 @@ int cactus_complete( const char* options_json, // Optional generation options (can be NULL) const char* tools_json, // Optional tools definition (can be NULL) cactus_token_callback callback, // Optional streaming callback (can be NULL) - void* user_data // User data for callback (can be NULL) + void* user_data, // User data for callback (can be NULL) + const uint8_t* pcm_buffer, // Optional raw PCM audio buffer (can be NULL) + size_t pcm_buffer_size // Size of PCM buffer in bytes (0 when not used) ); ``` @@ -94,17 +127,38 @@ int cactus_complete( ] ``` +**Messages with Audio (for multimodal models like Gemma4):** +```json +[ + {"role": "user", "content": "Transcribe the audio.", "audio": ["/path/to/audio.wav"]} +] +``` + +**Messages with Images and Audio:** +```json +[ + {"role": "user", "content": "Describe the image and transcribe the audio.", "images": ["/path/to/image.jpg"], "audio": ["/path/to/audio.wav"]} +] +``` + **Options Format:** ```json { "max_tokens": 256, "temperature": 0.7, "top_p": 0.95, + "min_p": 0.15, + "repetition_penalty": 1.1, "top_k": 40, "stop_sequences": ["<|im_end|>", ""], + "include_stop_sequences": false, "force_tools": false, "tool_rag_top_k": 2, - "confidence_threshold": 0.7 + "confidence_threshold": 0.7, + "auto_handoff": true, + "cloud_timeout_ms": 15000, + "handoff_with_images": true, + "enable_thinking_if_supported": false } ``` @@ -114,12 +168,19 @@ int cactus_complete( | `temperature` | float | 0.0 | Sampling temperature | | `top_p` | float | 0.0 | Top-p (nucleus) sampling | | `top_k` | int | 0 | Top-k sampling | +| `min_p` | float | 0.15 | Minimum probability threshold relative to max probability | +| `repetition_penalty` | float | 1.1 | Penalize previously generated tokens (1.0 disables) | | `stop_sequences` | array | [] | Stop generation on these strings | +| `include_stop_sequences` | bool | false | Include stop sequence tokens in the response | | `force_tools` | bool | false | Constrain output to tool call format | | `tool_rag_top_k` | int | 2 | Select top-k relevant tools via Tool RAG (0 = disabled, use all tools) | -| `confidence_threshold` | float | 0.7 | Minimum confidence for local generation; triggers cloud_handoff when below | +| `confidence_threshold` | float | model-dependent | Minimum confidence for local generation; triggers cloud_handoff when below. Resolved in this order: `0.5` if the bundle ships a `handoff_probe.bin`; else the model's `default_cloud_handoff_threshold` (Gemma 4 = `0.81`); else `0.7`. | +| `auto_handoff` | bool | true | Automatically attempt cloud handoff when confidence is low | +| `cloud_timeout_ms` | int | 15000 | Timeout in milliseconds for cloud handoff requests | +| `handoff_with_images` | bool | true | Allow cloud handoff for requests that include images | +| `enable_thinking_if_supported` | bool | false | Enable chain-of-thought thinking blocks for models that support it | -**Response Format** (all fields always present): +**Response Format:** ```json { "success": true, @@ -127,7 +188,9 @@ int cactus_complete( "cloud_handoff": false, "response": "I am an AI assistant.", "function_calls": [], + "segments": [], "confidence": 0.85, + "confidence_threshold": 0.7, "time_to_first_token_ms": 150.5, "total_time_ms": 1250.3, "prefill_tps": 166.1, @@ -139,15 +202,42 @@ int cactus_complete( } ``` -**Cloud Handoff Response** (when model detects low confidence): +`confidence_threshold` is the resolved value `confidence` is compared against — model-dependent (see the options table above), or whatever you pass; the `0.7` here is just the fallback default. `cloud_handoff` becomes `true` when `confidence` drops below it. + +The `thinking` field is only present in the JSON when the model produced a chain-of-thought block: ```json { - "success": false, + "success": true, + "error": null, + "cloud_handoff": false, + "response": "The answer is 4.", + "thinking": "Let me consider this... 2+2 equals 4.", + "function_calls": [], + "segments": [], + "confidence": 0.91, + "confidence_threshold": 0.7, + "time_to_first_token_ms": 150.5, + "total_time_ms": 1250.3, + "prefill_tps": 166.1, + "decode_tps": 45.2, + "ram_usage_mb": 245.67, + "prefill_tokens": 25, + "decode_tokens": 8, + "total_tokens": 33 +} +``` + +**Cloud Handoff Response** (when model detects low confidence and cloud handoff succeeds): +```json +{ + "success": true, "error": null, "cloud_handoff": true, - "response": null, + "response": "Cloud-provided answer.", "function_calls": [], + "segments": [], "confidence": 0.18, + "confidence_threshold": 0.7, "time_to_first_token_ms": 45.2, "total_time_ms": 45.2, "prefill_tps": 619.5, @@ -159,7 +249,7 @@ int cactus_complete( } ``` -When `cloud_handoff` is true, the model's confidence dropped below `confidence_threshold` (default: 0.7). The application should defer to a cloud-based model for better results. +When `cloud_handoff` is true, the model's confidence dropped below the resolved `confidence_threshold` (see the request options above) and the response was fulfilled by a cloud-based model. The `response` field contains the cloud-provided answer. **Error Response:** ```json @@ -169,18 +259,20 @@ When `cloud_handoff` is true, the model's confidence dropped below `confidence_t "cloud_handoff": false, "response": null, "function_calls": [], - "confidence": 0.0, + "confidence": null, "time_to_first_token_ms": 0.0, "total_time_ms": 0.0, "prefill_tps": 0.0, "decode_tps": 0.0, - "ram_usage_mb": 0.0, + "ram_usage_mb": 245.67, "prefill_tokens": 0, "decode_tokens": 0, "total_tokens": 0 } ``` +Note: `ram_usage_mb` reflects actual current RAM usage even in error responses. + **Response with Function Call:** ```json { @@ -191,10 +283,12 @@ When `cloud_handoff` is true, the model's confidence dropped below `confidence_t "function_calls": [ { "name": "get_weather", - "arguments": "{\"location\": \"San Francisco, CA, USA\"}" + "arguments": {"location": "San Francisco, CA, USA"} } ], + "segments": [], "confidence": 0.92, + "confidence_threshold": 0.7, "time_to_first_token_ms": 120.0, "total_time_ms": 450.5, "prefill_tps": 375.0, @@ -217,7 +311,93 @@ const char* messages = "[{\"role\": \"user\", \"content\": \"Tell me a story\"}] char response[8192]; int result = cactus_complete(model, messages, response, sizeof(response), - NULL, NULL, streaming_callback, NULL); + NULL, NULL, streaming_callback, NULL, NULL, 0); +``` + +### `cactus_prefill` +Pre-processes input text and populates the KV cache without generating output tokens. This reduces latency for future calls to `cactus_complete`. + +```c +int cactus_prefill( + cactus_model_t model, // Model handle + const char* messages_json, // JSON array of messages + char* response_buffer, // Buffer for response JSON + size_t buffer_size, // Size of response buffer + const char* options_json, // Optional generation options (can be NULL) + const char* tools_json, // Optional tools definition (can be NULL) + const uint8_t* pcm_buffer, // Optional raw PCM audio buffer (can be NULL) + size_t pcm_buffer_size // Size of PCM buffer in bytes (0 when not used) +); +``` + +**Returns:** Number of bytes written to response_buffer on success, negative value on error. + +**Message Format:** Same as `cactus_complete` (see above) + +**Options Format:** Same as `cactus_complete` (see above) + +**Response Format:** +```json +{ + "success": true, + "error": null, + "prefill_tokens": 25, + "prefill_tps": 166.1, + "total_time_ms": 150.5, + "ram_usage_mb": 245.67 +} +``` + +**Error Response:** +```json +{ + "success": false, + "error": "Error message here", + "prefill_tokens": 0, + "prefill_tps": 0.0, + "total_time_ms": 0.0, + "ram_usage_mb": 245.67 +} +``` + +**Example:** +```c +const char* tools = R"([{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City, State, Country"} + }, + "required": ["location"] + } + } +}])"; + +const char* base_messages = R"([ + { "role": "system", "content": "You are a helpful assistant." }, + { "role": "user", "content": "What is the weather in Paris?" }, + { "role": "assistant", "content": "<|tool_call_start|>get_weather(location=\"Paris\")<|tool_call_end|>" }, + { "role": "tool", "content": "{\"name\": \"get_weather\", \"content\": \"Sunny, 72°F\"}" }, + { "role": "assistant", "content": "It's sunny and 72°F in Paris!" } +])"; + +char prefill_response[1024]; +cactus_prefill(model, base_messages, prefill_response, sizeof(prefill_response), NULL, tools, NULL, 0); + +const char* completion_messages = R"([ + { "role": "system", "content": "You are a helpful assistant." }, + { "role": "user", "content": "What is the weather in Paris?" }, + { "role": "assistant", "content": "<|tool_call_start|>get_weather(location=\"Paris\")<|tool_call_end|>" }, + { "role": "tool", "content": "{\"name\": \"get_weather\", \"content\": \"Sunny, 72°F\"}" }, + { "role": "assistant", "content": "It's sunny and 72°F in Paris!" }, + { "role": "user", "content": "What about SF?" } +])"; +char response[4096]; +cactus_complete(model, completion_messages, response, sizeof(response), NULL, tools, NULL, NULL, NULL, 0); ``` ### `cactus_tokenize` @@ -233,7 +413,7 @@ int cactus_tokenize( ); ``` -**Returns:** 0 on success, negative value on error +**Returns:** 0 on success; -1 on invalid parameters or tokenization error; -2 if `token_buffer_len` is smaller than the number of tokens produced (but `*out_token_len` is still set to the required count). Pass `NULL` for `token_buffer` and `0` for `token_buffer_len` to query the token count without copying. **Example:** ```c @@ -251,6 +431,34 @@ if (result == 0) { } ``` +### `cactus_render_prompt` +Renders the chat-templated prompt string for the given messages without running inference. Useful for debugging prompt formatting, estimating token budgets, or feeding the rendered text to another tool. + +```c +int cactus_render_prompt( + cactus_model_t model, // Model handle + const char* messages_json, // JSON array of messages (same format as cactus_complete) + const char* options_json, // Optional generation options (can be NULL) + const char* tools_json, // Optional tools definition (can be NULL) + char* prompt_buffer, // Buffer for the rendered prompt text + size_t buffer_size // Size of prompt_buffer +); +``` + +**Returns:** The rendered prompt length in bytes (excluding the null terminator) on success; -1 on invalid parameters or rendering error; -2 if the buffer is too small to hold the rendered prompt and its null terminator. + +**Note:** The output is plain prompt text, not JSON. + +**Example:** +```c +const char* messages = "[{\"role\": \"user\", \"content\": \"Hello!\"}]"; +char prompt[4096]; +int len = cactus_render_prompt(model, messages, NULL, NULL, prompt, sizeof(prompt)); +if (len >= 0) { + printf("Rendered prompt:\n%s\n", prompt); +} +``` + ### `cactus_score_window` Scores a window of tokens for perplexity calculation or token probability analysis. @@ -269,6 +477,18 @@ int cactus_score_window( **Returns:** Number of bytes written to response_buffer on success, negative value on error +**Response Format:** +```json +{ + "success": true, + "logprob": -12.3456789012, + "tokens": 4 +} +``` + +- `logprob`: Total log-probability of the scored token window +- `tokens`: Number of tokens scored in the window + **Example:** ```c uint32_t tokens[256]; @@ -278,18 +498,36 @@ cactus_tokenize(model, "The quick brown fox", tokens, 256, &num_tokens); char response[4096]; int result = cactus_score_window(model, tokens, num_tokens, 0, num_tokens, 512, response, sizeof(response)); -if (result > 0) { +if (result >= 0) { printf("Scores: %s\n", response); } ``` +### `cactus_benchmark_tokens` +Runs a prefill + decode benchmark on the given prompt tokens and returns timing JSON +(prefill ms, decode ms, tokens/sec). Useful for measuring inference performance on +a specific model + device without running the full chat loop. + +```c +int cactus_benchmark_tokens( + cactus_model_t model, + const uint32_t* prompt_tokens, + size_t prompt_token_len, + size_t decode_token_len, + char* response_buffer, + size_t buffer_size +); +``` + +Returns the number of bytes written to `response_buffer` on success, `-1` on error. + ### `cactus_transcribe` -Transcribes audio to text using a Whisper model. Supports both file-based and buffer-based audio input. +Transcribes audio to text using Whisper or Parakeet TDT models. Supports both file-based and buffer-based audio input. ```c int cactus_transcribe( - cactus_model_t model, // Model handle (must be Whisper model) - const char* audio_file_path, // Path to audio file (WAV, MP3, etc.) - can be NULL if using pcm_buffer + cactus_model_t model, // Model handle (Whisper or Parakeet TDT model) + const char* audio_file_path, // Path to WAV file (16-bit PCM) - can be NULL if using pcm_buffer const char* prompt, // Optional prompt to guide transcription (can be NULL) char* response_buffer, // Buffer for response JSON size_t buffer_size, // Size of response buffer @@ -297,21 +535,65 @@ int cactus_transcribe( cactus_token_callback callback, // Optional streaming callback (can be NULL) void* user_data, // User data for callback (can be NULL) const uint8_t* pcm_buffer, // Optional raw PCM audio buffer (can be NULL if using file) - size_t pcm_buffer_size // Size of PCM buffer in bytes + size_t pcm_buffer_size // Size of PCM buffer in bytes (must be even and >= 2) ); ``` **Returns:** Number of bytes written to response_buffer on success, negative value on error +**Note:** Exactly one of `audio_file_path` or `pcm_buffer` must be provided; passing both or neither returns -1. The file path must point to a 16-bit PCM WAV file. The `pcm_buffer` must contain 16-bit signed PCM samples at 16 kHz and `pcm_buffer_size` must be even and at least 2. + +**Options Format:** +```json +{ + "max_tokens": 448, + "language": "en", + "timestamps": true +} +``` + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `max_tokens` | int | auto | Maximum tokens to generate. When unset, it defaults to the larger of 100 and an audio-length estimate (`audio_sec × 20` for Whisper, `audio_sec × 30` for Parakeet). For Whisper the result is then capped so the prompt tokens plus generated tokens fit the decoder's 448-position limit. | +| `language` | string | model default | Whisper only. Two-letter language code (e.g. `en`, `es`, `de`) substituted into the decoder prompt's language token. Ignored by Parakeet and when an explicit `prompt` is supplied. | +| `timestamps` | bool | false | Whisper only. Decodes timestamp tokens and populates `segments` with `{start, end, text}` entries (seconds). Empty otherwise, including all Parakeet transcription. | + +**Response Format:** +```json +{ + "success": true, + "error": null, + "cloud_handoff": false, + "response": "Transcribed text here.", + "function_calls": [], + "segments": [], + "confidence": 1.0, + "confidence_threshold": -1.0, + "time_to_first_token_ms": 120.0, + "total_time_ms": 450.0, + "prefill_tps": 50.0, + "decode_tps": 30.0, + "ram_usage_mb": 512.34, + "prefill_tokens": 10, + "decode_tokens": 15, + "total_tokens": 25 +} +``` + +- `response`: Full transcription text +- `segments`: `{start, end, text}` entries (seconds), populated only for Whisper when the `timestamps` option is set; empty otherwise, including all Parakeet transcription +- `cloud_handoff`: Always false for transcription +- `confidence_threshold`: `-1.0` (unset) — transcription does not resolve a cloud-handoff threshold + **Example (file-based):** ```c -cactus_model_t whisper = cactus_init("../../weights/whisper-small", NULL); +cactus_model_t whisper = cactus_init("../../weights/whisper-small", NULL, false); char response[16384]; int result = cactus_transcribe(whisper, "audio.wav", NULL, response, sizeof(response), NULL, NULL, NULL, NULL, 0); -if (result > 0) { +if (result >= 0) { printf("Transcription: %s\n", response); } ``` @@ -326,145 +608,81 @@ int result = cactus_transcribe(whisper, NULL, NULL, pcm_data, pcm_size); ``` -### `cactus_stream_transcribe_t` -An opaque pointer type representing a streaming transcription session. Used for real-time audio transcription with incremental confirmation. - -```c -typedef void* cactus_stream_transcribe_t; -``` - -### `cactus_stream_transcribe_init` -Initializes a new streaming transcription session for a transcription model. - -```c -cactus_stream_transcribe_t cactus_stream_transcribe_init( - cactus_model_t model // Model handle (must be Whisper model) -); -``` - -**Returns:** Stream handle on success, NULL on failure +### Streaming transcription +Transcribe continuously while audio is still being captured, instead of waiting for a complete recording. You push PCM as it arrives and read back text as soon as it stabilizes. Supported for the dedicated speech models (Whisper and Parakeet TDT). -**Example:** -```c -cactus_model_t whisper = cactus_init("../../weights/whisper-small", NULL); +The session emits text in two parts on every call: -cactus_stream_transcribe_t stream = cactus_stream_transcribe_init(whisper); -if (!stream) { - fprintf(stderr, "Failed to initialize stream: %s\n", cactus_get_last_error()); - return -1; -} -``` +- **`confirmed`** — newly finalized words. Append them to your running transcript; they never change. +- **`pending`** — the current best guess for the still-changing tail. Replace it on every call (do not append it); it is for live display only. -### `cactus_stream_transcribe_insert` -Inserts audio samples into the streaming transcription buffer. Audio should be 16-bit PCM, 16kHz, mono. +Text is confirmed once two successive re-transcriptions of the audio agree (LocalAgreement, compared per **segment** for Whisper) or once a following token starts a new word (Parakeet TDT). Confirmed audio is dropped from the front of the buffer as the window advances, so memory stays bounded and the active window never approaches the model's fixed input length. +#### `cactus_stream_transcribe_start` +Opens a streaming session bound to an already-initialized speech model. ```c -int cactus_stream_transcribe_insert( - cactus_stream_transcribe_t stream, // Stream handle - const uint8_t* pcm_buffer, // Raw PCM audio data (16-bit, 16kHz, mono) - size_t pcm_buffer_size // Size of PCM buffer in bytes +cactus_stream_transcribe_t cactus_stream_transcribe_start( + cactus_model_t model, // Whisper or Parakeet TDT model handle + const char* options_json // Optional (can be NULL); forwarded to cactus_transcribe ); ``` +**Returns:** an opaque session handle, or `NULL` on error (see `cactus_get_last_error`). Free it with `cactus_stream_transcribe_stop`. -**Returns:** 0 on success, negative value on error - -**Example:** -```c -uint8_t audio_chunk[32000]; // 1 second at 16kHz * 2 bytes - -int result = cactus_stream_transcribe_insert(stream, audio_chunk, sizeof(audio_chunk)); -if (result < 0) { - fprintf(stderr, "Insert failed: %s\n", cactus_get_last_error()); -} -``` - -### `cactus_stream_transcribe_process` -Processes the accumulated audio buffer and returns confirmed and pending transcription results. +`options_json` (optional) is forwarded to the underlying `cactus_transcribe` call **for Whisper only** (e.g. `language`); the Parakeet TDT path ignores it. Chunking and segmentation are handled internally with no user-facing tunables. +#### `cactus_stream_transcribe_process` +Feeds the next slice of audio. Input is 16-bit signed PCM, **16 kHz, mono** — the same format as `cactus_transcribe`'s `pcm_buffer`. Feed reasonably small chunks (≈ 0.1–2 s) for low latency. ```c int cactus_stream_transcribe_process( - cactus_stream_transcribe_t stream, // Stream handle - char* response_buffer, // Buffer for response JSON - size_t buffer_size, // Size of response buffer - const char* options_json // Optional processing options (can be NULL) + cactus_stream_transcribe_t stream, + const uint8_t* pcm_buffer, // 16-bit PCM, 16 kHz, mono + size_t pcm_buffer_size, // size in bytes + char* response_buffer, + size_t buffer_size ); ``` +**Returns:** bytes written to `response_buffer`, or -1 on error. -**Returns:** Number of bytes written to response_buffer on success, negative value on error - -**Options Format:** +**Response Format:** (also includes timing stats: `decode_tps`, `total_time_ms`, `time_to_first_token_ms`, `decode_tokens`) ```json -{ - "confirmation_threshold": 0.95 -} -``` - -- `confirmation_threshold`: Threshold (0.0-1.0) for confirming transcription segments. Higher values require more stability before confirmation. Default: 0.95 - -**Response Format:** -```json -{ - "success": true, - "confirmed": "text confirmed from previous call", - "pending": "current transcription result" -} +{ "success": true, "confirmed": "the quick brown", "pending": "fox jumps", "decode_tps": 0, "total_time_ms": 0, "time_to_first_token_ms": 0, "decode_tokens": 0 } ``` -- `confirmed`: Text that was confirmed from the previous call (append to final transcription) -- `pending`: Current transcription result (may be confirmed in next call if stable) - -**Example:** +#### `cactus_stream_transcribe_stop` +Flushes any buffered audio, returns the final `confirmed` text, and destroys the session. ```c -char response[32768]; -int result = cactus_stream_transcribe_process(stream, response, sizeof(response), "{\"confirmation_threshold\": 0.90}"); - -if (result > 0) { - printf("Response: %s\n", response); -} -``` - -### `cactus_stream_transcribe_finalize` -Finalizes the streaming session and confirms any remaining transcription. - -```c -int cactus_stream_transcribe_finalize( - cactus_stream_transcribe_t stream, // Stream handle - char* response_buffer, // Buffer for response JSON - size_t buffer_size // Size of response buffer +int cactus_stream_transcribe_stop( + cactus_stream_transcribe_t stream, + char* response_buffer, // may be NULL to discard the final text + size_t buffer_size ); ``` - -**Returns:** Number of bytes written to response_buffer on success, negative value on error - -**Response Format:** -```json -{ - "success": true, - "confirmed": "Final confirmed transcription segment" -} -``` +**Returns:** bytes written (0 when discarded), or -1 on error. **Example:** ```c -char final_response[32768]; -int result = cactus_stream_transcribe_finalize(stream, final_response, sizeof(final_response)); -if (result > 0) { - printf("Final: %s\n", final_response); -} -``` +cactus_model_t whisper = cactus_init("../../weights/whisper-base", NULL, false); +cactus_stream_transcribe_t stream = cactus_stream_transcribe_start(whisper, NULL); -### `cactus_stream_transcribe_destroy` -Releases all resources associated with the streaming transcription session. +char response[16384]; -```c -void cactus_stream_transcribe_destroy( - cactus_stream_transcribe_t stream // Stream handle -); +// Push audio as it is captured (here, 0.5s chunks of 16kHz mono PCM16). +for (each chunk) { + int rc = cactus_stream_transcribe_process( + stream, chunk_pcm, chunk_bytes, response, sizeof(response)); + if (rc < 0) break; + // Parse "confirmed" from response and append it to your transcript; + // show "confirmed-so-far + pending" live. +} + +// Flush the tail; its "confirmed" holds the final words. +cactus_stream_transcribe_stop(stream, response, sizeof(response)); ``` -**Example:** -```c -cactus_stream_transcribe_destroy(stream); +The `transcribe` CLI uses this streaming path for **live microphone** transcription (no file; a colored UI with running captions, press Enter to stop). With a **file** it does a one-shot transcription (long files are windowed internally, so there is no 30s limit): +```bash +cactus transcribe openai/whisper-base # live microphone (press Enter to stop) +cactus transcribe openai/whisper-base --file audio.wav # one-shot file transcription (any length) ``` ### `cactus_embed` @@ -475,13 +693,13 @@ int cactus_embed( cactus_model_t model, // Model handle const char* text, // Text to embed float* embeddings_buffer, // Buffer for embedding vector - size_t buffer_size, // Buffer size in bytes + size_t buffer_size, // Size of embeddings_buffer in bytes size_t* embedding_dim, // Output: actual embedding dimensions bool normalize // Whether to L2-normalize the output vector ); ``` -**Returns:** 0 on success, negative value on error +**Returns:** Number of float elements written to embeddings_buffer on success; -1 on invalid parameters, tokenization error, or other failure; -2 if `buffer_size` (in bytes) is smaller than `embedding_dim * sizeof(float)` **Example:** ```c @@ -490,7 +708,7 @@ float embeddings[2048]; size_t actual_dim = 0; int result = cactus_embed(model, text, embeddings, sizeof(embeddings), &actual_dim, true); -if (result == 0) { +if (result >= 0) { printf("Generated %zu-dimensional embedding\n", actual_dim); } ``` @@ -505,21 +723,20 @@ int cactus_image_embed( cactus_model_t model, // Model handle (must support vision) const char* image_path, // Path to image file float* embeddings_buffer, // Buffer for embedding vector - size_t buffer_size, // Buffer size in bytes + size_t buffer_size, // Size of embeddings_buffer in bytes size_t* embedding_dim // Output: actual embedding dimensions ); ``` -**Returns:** 0 on success, negative value on error +**Returns:** Number of float elements written to embeddings_buffer on success; -1 on invalid parameters or embedding failure; -2 if `buffer_size` (in bytes) is smaller than `embedding_dim * sizeof(float)` **Example:** ```c float image_embeddings[1024]; size_t dim = 0; -int result = cactus_image_embed(model, "photo.jpg", image_embeddings, - sizeof(image_embeddings), &dim); -if (result == 0) { +int result = cactus_image_embed(model, "photo.jpg", image_embeddings, sizeof(image_embeddings), &dim); +if (result >= 0) { printf("Image embedding dimension: %zu\n", dim); } ``` @@ -532,20 +749,19 @@ int cactus_audio_embed( cactus_model_t model, // Model handle (must support audio) const char* audio_path, // Path to audio file float* embeddings_buffer, // Buffer for embedding vector - size_t buffer_size, // Buffer size in bytes + size_t buffer_size, // Size of embeddings_buffer in bytes size_t* embedding_dim // Output: actual embedding dimensions ); ``` -**Returns:** 0 on success, negative value on error +**Returns:** Number of float elements written to embeddings_buffer on success; -1 on invalid parameters or embedding failure; -2 if `buffer_size` (in bytes) is smaller than `embedding_dim * sizeof(float)` **Example:** ```c float audio_embeddings[768]; size_t dim = 0; -int result = cactus_audio_embed(model, "speech.wav", audio_embeddings, - sizeof(audio_embeddings), &dim); +int result = cactus_audio_embed(model, "speech.wav", audio_embeddings, sizeof(audio_embeddings), &dim); ``` ### `cactus_stop` @@ -576,7 +792,7 @@ void control_callback(const char* token, uint32_t token_id, void* user_data) { struct ControlData control = {model, 0, 50}; cactus_complete(model, messages, response, sizeof(response), - NULL, NULL, control_callback, &control); + NULL, NULL, control_callback, &control, NULL, 0); ``` ### `cactus_reset` @@ -605,26 +821,36 @@ int cactus_rag_query( ); ``` -**Returns:** 0 on success, negative value on error +**Returns:** Number of bytes written to response_buffer on success; 0 when the query cannot be executed (no corpus index, no tokenizer, empty query, or dimension mismatch) — response_buffer contains `{"chunks":[],"error":"..."}` in those cases; also 0 when the query executes but returns no results — response_buffer contains `{"chunks":[]}` with no `error` field; -1 on error (invalid params, buffer too small, or exception) **Response Format:** ```json -[ - {"text": "Relevant chunk 1...", "score": 0.85}, - {"text": "Relevant chunk 2...", "score": 0.72} -] +{ + "chunks": [ + {"score": 0.85, "source": "document.txt", "content": "Relevant chunk 1..."}, + {"score": 0.72, "source": "document.txt", "content": "Relevant chunk 2..."} + ] +} +``` + +When the query cannot be executed (no corpus index, no tokenizer, empty query, or dimension mismatch), `chunks` is empty and an `error` field is present: +```json +{ + "chunks": [], + "error": "No corpus index loaded" +} ``` **Example:** ```c // Initialize model with corpus -cactus_model_t model = cactus_init("path/to/model", "./documents"); +cactus_model_t model = cactus_init("path/to/model", "./documents", true); // Query for relevant chunks char response[65536]; int result = cactus_rag_query(model, "What is machine learning?", response, sizeof(response), 5); -if (result == 0) { +if (result >= 0) { printf("Retrieved chunks: %s\n", response); } ``` @@ -647,57 +873,21 @@ Returns the last error message from the Cactus engine. const char* cactus_get_last_error(void); ``` -**Returns:** Error message string, or NULL if no error +**Returns:** Error message string (never NULL; empty string if no error) **Example:** ```c -cactus_model_t model = cactus_init("invalid/path", NULL); +cactus_model_t model = cactus_init("invalid/path", NULL, false); if (!model) { const char* error = cactus_get_last_error(); fprintf(stderr, "Error: %s\n", error); } ``` -### `cactus_set_telemetry_token` -Sets the telemetry token for usage tracking. Pass NULL or empty string to disable telemetry. - -```c -void cactus_set_telemetry_token(const char* token); -``` - -**Example:** -```c -cactus_set_telemetry_token("your-telemetry-token"); -cactus_set_telemetry_token(NULL); // disable -``` - -### `cactus_set_pro_key` -Sets the pro key to enable NPU acceleration on supported devices (Apple Neural Engine). - -```c -void cactus_set_pro_key(const char* pro_key); -``` - -**Example:** -```c -cactus_set_pro_key("your-pro-key"); - -cactus_model_t model = cactus_init("path/to/model", NULL); -``` - -**Note:** The pro key should be set before initializing any models to ensure NPU acceleration is enabled. - ## Vector Index APIs The vector index APIs provide persistent storage and retrieval of embeddings for RAG (Retrieval-Augmented Generation) applications. -### `cactus_index_t` -An opaque pointer type representing a vector index instance. - -```c -typedef void* cactus_index_t; -``` - ### `cactus_index_init` Initializes or opens a vector index from disk. @@ -776,11 +966,11 @@ int cactus_index_get( const int* ids, // Array of document IDs to retrieve size_t ids_count, // Number of IDs char** document_buffers, // Output: document text buffers - size_t* document_buffer_sizes, // Sizes of document buffers + size_t* document_buffer_sizes, // Sizes of document buffers (in bytes) char** metadata_buffers, // Output: metadata JSON buffers - size_t* metadata_buffer_sizes, // Sizes of metadata buffers + size_t* metadata_buffer_sizes, // Sizes of metadata buffers (in bytes) float** embedding_buffers, // Output: embedding buffers - size_t* embedding_buffer_sizes // Sizes of embedding buffers (in bytes) + size_t* embedding_buffer_sizes // Sizes of embedding buffers (in float elements, not bytes) ); ``` @@ -795,19 +985,27 @@ int cactus_index_query( const float** embeddings, // Array of query embeddings size_t embeddings_count, // Number of query embeddings size_t embedding_dim, // Dimension of each embedding - const char* options_json, // Query options (e.g., {"k": 10}) + const char* options_json, // Query options (e.g., {"top_k": 10, "score_threshold": 0.5}) int** id_buffers, // Output: arrays of result IDs - size_t* id_buffer_sizes, // Sizes of ID buffers + size_t* id_buffer_sizes, // In: capacity of each id_buffer; Out: actual result count written float** score_buffers, // Output: arrays of similarity scores - size_t* score_buffer_sizes // Sizes of score buffers + size_t* score_buffer_sizes // In: capacity of each score_buffer; Out: actual result count written ); ``` **Returns:** 0 on success, negative value on error +**Options JSON fields:** + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `top_k` | int | 10 | Maximum number of results to return per query | +| `score_threshold` | float | -1.0 | Minimum similarity score threshold; results below this are excluded (-1.0 disables filtering) | + **Example:** ```c float query_emb[768]; +size_t dim; cactus_embed(model, "search query", query_emb, sizeof(query_emb), &dim, true); const float* queries[] = {query_emb}; @@ -818,10 +1016,11 @@ float* score_bufs[] = {result_scores}; size_t id_sizes[] = {10}; size_t score_sizes[] = {10}; -cactus_index_query(index, queries, 1, 768, "{\"k\": 10}", +cactus_index_query(index, queries, 1, 768, "{\"top_k\": 10}", id_bufs, id_sizes, score_bufs, score_sizes); -for (int i = 0; i < 10; i++) { +// id_sizes[0] is updated to the actual number of results returned +for (size_t i = 0; i < id_sizes[0]; i++) { printf("ID: %d, Score: %.4f\n", result_ids[i], result_scores[i]); } ``` @@ -852,10 +1051,10 @@ void cactus_index_destroy(cactus_index_t index); ### Complete RAG Example ```c -#include "cactus_ffi.h" +#include "cactus_engine.h" int main() { - cactus_model_t embed_model = cactus_init("path/to/embed-model", NULL); + cactus_model_t embed_model = cactus_init("path/to/embed-model", NULL, false); cactus_index_t index = cactus_index_init("./rag_index", 768); const char* docs[] = { @@ -885,7 +1084,7 @@ int main() { size_t id_sizes[] = {3}; size_t score_sizes[] = {3}; - cactus_index_query(index, queries, 1, 768, "{\"k\": 3}", + cactus_index_query(index, queries, 1, 768, "{\"top_k\": 3}", id_bufs, id_sizes, score_bufs, score_sizes); printf("Top result ID: %d (score: %.4f)\n", result_ids[0], result_scores[0]); @@ -900,11 +1099,11 @@ int main() { ### Basic Conversation ```c -#include "cactus_ffi.h" +#include "cactus_engine.h" #include int main() { - cactus_model_t model = cactus_init("path/to/model", NULL); + cactus_model_t model = cactus_init("path/to/model", NULL, false); if (!model) return -1; const char* messages = @@ -915,8 +1114,8 @@ int main() { char response[4096]; int result = cactus_complete(model, messages, response, - sizeof(response), NULL, NULL, NULL, NULL); - if (result > 0) { + sizeof(response), NULL, NULL, NULL, NULL, NULL, 0); + if (result >= 0) { printf("Response: %s\n", response); } @@ -927,10 +1126,10 @@ int main() { ### Vision-Language Model (VLM) ```c -#include "cactus_ffi.h" +#include "cactus_engine.h" int main() { - cactus_model_t vlm = cactus_init("path/to/lfm2-vlm", NULL); + cactus_model_t vlm = cactus_init("path/to/lfm2-vlm", NULL, false); if (!vlm) return -1; const char* messages = @@ -940,8 +1139,8 @@ int main() { char response[8192]; int result = cactus_complete(vlm, messages, response, sizeof(response), - NULL, NULL, NULL, NULL); - if (result > 0) { + NULL, NULL, NULL, NULL, NULL, 0); + if (result >= 0) { printf("%s\n", response); } @@ -953,7 +1152,7 @@ int main() { ### Tool Calling ```c const char* tools = - "[{\"function\": {" + "[{\"type\": \"function\", \"function\": {" " \"name\": \"get_weather\"," " \"description\": \"Get weather for a location\"," " \"parameters\": {" @@ -969,7 +1168,7 @@ const char* messages = "[{\"role\": \"user\", \"content\": \"What's the weather char response[4096]; int result = cactus_complete(model, messages, response, sizeof(response), - NULL, tools, NULL, NULL); + NULL, tools, NULL, NULL, NULL, 0); printf("Response: %s\n", response); ``` @@ -997,7 +1196,7 @@ printf("Similarity: %.4f\n", similarity); ### Audio Transcription with Whisper ```c -#include "cactus_ffi.h" +#include "cactus_engine.h" #include void transcription_callback(const char* token, uint32_t token_id, void* user_data) { @@ -1006,7 +1205,7 @@ void transcription_callback(const char* token, uint32_t token_id, void* user_dat } int main() { - cactus_model_t whisper = cactus_init("path/to/whisper-small", NULL); + cactus_model_t whisper = cactus_init("path/to/whisper-small", NULL, false); if (!whisper) return -1; char response[32768]; @@ -1022,7 +1221,7 @@ int main() { ### Multimodal Retrieval ```c -#include "cactus_ffi.h" +#include "cactus_engine.h" #include int find_similar_image(cactus_model_t model, const char* query, @@ -1056,32 +1255,6 @@ int find_similar_image(cactus_model_t model, const char* query, } ``` -## Supported Model Types - -| Model Type | Text | Vision | Audio | Embeddings | Description | -|------------|------|--------|-------|------------|-------------| -| Qwen | ✓ | - | - | ✓ | Qwen/Qwen2/Qwen3 language models | -| Gemma | ✓ | - | - | ✓ | Google Gemma models | -| LFM2 | ✓ | ✓ | - | ✓ | Liquid Foundation Models | -| Smol | ✓ | - | - | ✓ | SmolLM compact models | -| Nomic | - | - | - | ✓ | Nomic embedding models | -| Whisper | - | - | ✓ | ✓ | OpenAI Whisper transcription | -| Siglip2 | - | ✓ | - | ✓ | Vision encoder for embeddings | - -## Environment Variables - -| Variable | Default | Description | -|----------|---------|-------------| -| `CACTUS_KV_WINDOW_SIZE` | 512 | Sliding window size for KV cache | -| `CACTUS_KV_SINK_SIZE` | 4 | Number of attention sink tokens to preserve | - -**Example:** -```bash -export CACTUS_KV_WINDOW_SIZE=1024 -export CACTUS_KV_SINK_SIZE=8 -./my_app -``` - ## Best Practices 1. **Always Check Return Values**: Functions return negative values on error @@ -1111,4 +1284,76 @@ Common error scenarios: 2. **Streaming for UX**: Use callbacks for responsive user interfaces 3. **Early Stopping**: Use `cactus_stop()` to avoid unnecessary generation 4. **Batch Embeddings**: When possible, process multiple texts in sequence without resetting -5. **KV Cache Tuning**: Adjust `CACTUS_KV_WINDOW_SIZE` based on your context needs \ No newline at end of file + +## Logging + +### `cactus_log_set_level` +Sets the minimum log level. Messages below this level are suppressed. + +```c +void cactus_log_set_level(int level); +// level: 0=DEBUG, 1=INFO, 2=WARN (default), 3=ERROR, 4=NONE +``` + +### `cactus_log_set_callback` +Installs a callback to receive log messages. Pass NULL to remove the callback. + +```c +typedef void (*cactus_log_callback_t)(int level, const char* component, const char* message, void* user_data); + +void cactus_log_set_callback(cactus_log_callback_t callback, void* user_data); +``` + +**Example:** +```c +void my_log(int level, const char* component, const char* message, void* user_data) { + printf("[%d] %s: %s\n", level, component, message); +} + +cactus_log_set_level(1); // INFO and above +cactus_log_set_callback(my_log, NULL); +``` + +## Telemetry + +These functions configure anonymous usage telemetry sent to Cactus Compute. Telemetry is opt-out and contains no user data. + +### `cactus_set_telemetry_environment` +Identifies the calling framework and cache directory. + +```c +void cactus_set_telemetry_environment(const char* framework, const char* cache_location, const char* version); +``` + +### `cactus_set_app_id` +Associates telemetry events with an application identifier. + +```c +void cactus_set_app_id(const char* app_id); +``` + +### `cactus_telemetry_flush` +Flushes pending telemetry events. + +```c +void cactus_telemetry_flush(void); +``` + +### `cactus_telemetry_shutdown` +Flushes and shuts down the telemetry subsystem. Call before process exit. + +```c +void cactus_telemetry_shutdown(void); +``` + +## See Also + +- [Cactus Graph API](/docs/cactus_graph.md) — Low-level computational graph for custom tensor operations +- [Cactus Index API](/docs/cactus_index.md) — On-device vector database for RAG applications +- [Fine-tuning Guide](/docs/finetuning.md) — Deploy Unsloth LoRA fine-tunes to mobile +- [Runtime Compatibility](/docs/compatibility.md) — Weight versioning across releases +- [Python Binding](/python/) — Python bindings for the Engine API +- [Swift Binding](/bindings/swift/) — Swift bindings for iOS and macOS +- [Kotlin Binding](/bindings/kotlin/) — Kotlin Multiplatform bindings +- [Flutter Binding](/bindings/flutter/) — Dart FFI bindings for mobile apps +- [Rust Binding](/bindings/rust/) — Raw `extern "C"` FFI declarations diff --git a/docs/cactus_graph.md b/docs/cactus_graph.md index a1d8300e4..a07a48d51 100644 --- a/docs/cactus_graph.md +++ b/docs/cactus_graph.md @@ -1,3 +1,9 @@ +--- +title: "Cactus Graph API Documentation" +description: "Computational graph framework for building and executing tensor operations on mobile devices. Supports matrix multiplication, attention, normalization, and INT8/FP16/CQ1-CQ4 precision." +keywords: ["computation graph", "tensor operations", "mobile AI", "matrix multiplication", "attention mechanism", "INT8", "FP16", "CQ1", "CQ4"] +--- + # Cactus Graph API Documentation The Cactus Graph API provides a computational graph framework for building and executing tensor operations. It supports multiple precision types, broadcasting, and optimized execution for neural network inference. @@ -17,7 +23,7 @@ Before using the Cactus Graph API, set up your development environment: ```bash # Setup the environment and install dependencies -./setup +source ./setup # Build the Cactus library cactus build @@ -29,17 +35,21 @@ cactus test ## Core Concepts ### Precision Types -The framework supports three precision types for tensors: +The framework supports the following precision types for tensors: ```cpp enum class Precision { - INT4, INT8, - FP16 + FP16, + FP32, + CQ1, + CQ2, + CQ3, + CQ4 }; ``` -**Note:** INT4 tensors use packed storage (2 values per byte) and automatically unpack to INT8 for computation. +**Note:** CQ1–CQ4 are Cactus custom quantization formats. ### Graph Construction The `CactusGraph` class manages the computational graph: @@ -57,14 +67,14 @@ For testing, use the provided fixtures that handle memory management: ```cpp TestUtils::Int8TestFixture fixture("My Test"); -TestUtils::FloatTestFixture fixture("Float Test"); +TestUtils::FP16TestFixture fixture("Float Test"); ``` ## Getting Started ### Basic Example ```cpp -#include "cactus/graph/graph.h" +#include "cactus_graph.h" CactusGraph graph; size_t a = graph.input({4}, Precision::INT8); @@ -106,6 +116,7 @@ size_t exp_result = graph.scalar_exp(input); // e^input size_t sqrt_result = graph.scalar_sqrt(input); // √input size_t cos_result = graph.scalar_cos(input); // cos(input) size_t sin_result = graph.scalar_sin(input); // sin(input) +size_t log_result = graph.scalar_log(input); // ln(input) ``` ### Matrix Operations @@ -178,6 +189,61 @@ size_t rope_output = graph.rope(input, theta, position_offset); ```cpp size_t silu_out = graph.silu(input); size_t gelu_out = graph.gelu(input); +size_t gelu_erf_out = graph.gelu_erf(input); // GeLU with erf approximation +size_t sigmoid_out = graph.sigmoid(input); +size_t tanh_out = graph.tanh(input); +size_t relu_out = graph.relu(input); +size_t glu_out = graph.glu(input, axis); // Gated Linear Unit +``` + +#### Convolution Operations + +Each conv has two overloads: with and without bias. + +```cpp +// 1D convolutions (no bias / with bias overloads) +size_t conv1d_out = graph.conv1d(input, weight, stride); +size_t conv1d_with_bias = graph.conv1d(input, weight, bias, stride); +size_t conv1d_k3_out = graph.conv1d_k3(input, weight, stride); +size_t conv1d_causal = graph.conv1d_causal(input, weight, kernel_size, dilation); +size_t conv1d_pointwise = graph.conv1d_pointwise(input, weight); // or (input, weight, bias) +size_t conv1d_depthwise = graph.conv1d_same_depthwise_k9(input, weight); // or (input, weight, bias) + +// 2D convolutions +size_t conv2d_out = graph.conv2d_k3s2p1(input, weight); // or (input, weight, bias) +size_t conv2d_dw = graph.conv2d_depthwise_k3s2p1(input, weight); // or (input, weight, bias) +size_t conv2d_pw = graph.conv2d_pointwise_1x1(input, weight); // or (input, weight, bias) +``` + +#### Normalization +```cpp +size_t groupnorm_out = graph.groupnorm(input, weight, bias, num_groups, epsilon); +size_t batchnorm_out = graph.batchnorm(input, weight, bias, running_mean, running_var, axis, epsilon); +``` + +#### Recurrent & Sequence Operations +```cpp +size_t lstm_out = graph.lstm_cell(input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh); +size_t deltanet_out = graph.gated_deltanet_decode(query, key, value, gate_log, beta, initial_state, scale); +size_t deltanet_prefill = graph.gated_deltanet_prefill(query, key, value, gate_log, beta, initial_state, chunk_size, scale); +``` + +#### Mixture of Experts (MoE) +```cpp +// Gated SwiGLU MoE: pass three weight sets (w1=gate, w3=up, w2=down) +size_t moe_gated = graph.moe_layer(hidden, routing_probs, topk_indices, + w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor); + +// Ungated MoE: pass two weight sets (w1=up, w2=down) and the activation explicitly +size_t moe_ungated = graph.moe_layer(hidden, routing_probs, topk_indices, + w1_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor, activation); +``` + +#### Signal Processing +```cpp +size_t stft_out = graph.stft(input, weight, stride, num_fft_bins); ``` ### Indexing and Gathering @@ -206,6 +272,7 @@ size_t weights = graph.mmap_weights("model_weights.bin"); #### Concatenation ```cpp size_t concatenated = graph.concat(tensor1, tensor2, axis); +size_t multi_cat = graph.cat({tensor1, tensor2, tensor3}, axis); // cat multiple tensors ``` #### Slicing @@ -223,9 +290,30 @@ size_t indexed = graph.index(input, index_value, dimension); size_t topk_values = graph.topk(input, k); ``` +#### Persistent Nodes +```cpp +size_t persistent = graph.persistent(source_node); // cache result across executions +``` + +#### AltUp Operations +```cpp +size_t prediction = graph.altup_predict(coefs, streams, num_streams); +size_t correction = graph.altup_correct(coefs, innovation, predictions, num_predictions); +``` + +#### Bilinear Interpolation +```cpp +size_t interpolated = graph.bilinear_interpolation(pos_embeds, dst_height, dst_width); +``` + #### Sampling ```cpp +// Defaults: temperature=0.6, top_p=0.95, top_k=20 size_t sampled = graph.sample(logits, temperature, top_p, top_k); + +// With per-token bias map and repetition penalty / min_p: +size_t sampled_ext = graph.sample_with_options( + logits, temperature, top_p, min_p, repetition_penalty, top_k, /*logit_bias=*/{}); ``` ## Advanced Features @@ -251,25 +339,48 @@ size_t result = graph.add(a, b); // {2,3} -> {2,2,3} ```cpp size_t int8_tensor = graph.input({4}, Precision::INT8); size_t fp16_tensor = graph.precision_cast(int8_tensor, Precision::FP16); -graph.set_quantization_scale(node_id, scale); ``` ### Graph Persistence +#### Saving graphs +```cpp +const std::string filename = "test_graph_save_load.cg"; + +CactusGraph graph; +size_t input_a = graph.input({2, 3}, Precision::FP16); +size_t input_b = graph.input({2, 3}, Precision::FP16); +size_t sum_id = graph.add(input_a, input_b); +graph.save(filename); + +``` +#### Loading Graphs +```cpp +CactusGraph loaded = CactusGraph::load(filename); +std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; +std::vector<__fp16> data_b = {10, 20, 30, 40, 50, 60}; +loaded.set_input(0, data_a.data(), Precision::FP16); +loaded.set_input(1, data_b.data(), Precision::FP16); +loaded.execute(); +``` + #### Saving Nodes ```cpp GraphFile::save_node(graph, node_id, "output.bin"); ``` -#### Loading Nodes +#### GraphFile Namespace ```cpp -CactusGraph new_graph; -auto loaded = GraphFile::load_into_graph(new_graph, "output.bin"); -size_t node_id = loaded.node_id; -std::vector shape = loaded.shape; -Precision precision = loaded.precision; +GraphFile::save_graph(graph, "graph.cactus"); +GraphFile::SerializedGraph serialized = GraphFile::load_graph("graph.cactus"); +CactusGraph loaded = CactusGraph::from_serialized(serialized); ``` +The `GraphFile` namespace currently exposes `save_graph`, `load_graph`, and +`save_node`. There is no per-node loader; reload a saved subgraph by +serializing the parent graph with `save_graph` and reading it back with +`load_graph`. + ### Graph Management #### Execution @@ -350,7 +461,7 @@ size_t final_embed = graph.add(embeddings, pos_embed); ### Similarity Computation ```cpp -TestUtils::FloatTestFixture fixture("Similarity"); +TestUtils::FP16TestFixture fixture("Similarity"); size_t text1 = fixture.create_input({1, 768}, Precision::FP16); size_t text2 = fixture.create_input({1, 768}, Precision::FP16); @@ -369,7 +480,7 @@ size_t similarity = fixture.graph().divide(dot_product, fixture.graph().multiply ## Best Practices ### Memory Management -1. **Use appropriate precision**: INT4/INT8 for memory efficiency, FP16 for accuracy +1. **Use appropriate precision**: CQ1-CQ4/INT8 for memory efficiency, FP16 for accuracy 2. **Memory-map large tensors**: Use `mmap_embeddings()` for vocabulary tables 3. **Reset graphs**: Call `hard_reset()` when switching between different models 4. **External buffers**: Use `set_external_input()` to avoid copying large inputs @@ -378,10 +489,11 @@ size_t similarity = fixture.graph().divide(dot_product, fixture.graph().multiply 1. **Batch operations**: Process multiple samples together 2. **Pre-transpose weights**: Use `pretransposed_rhs=true` for matmul when possible 3. **Fused operations**: The framework automatically fuses compatible operations -4. **Backend selection**: Use NPU backend for supported operations: +4. **Backend selection**: Every op accepts a backend (`ComputeBackend::CPU` or `ComputeBackend::METAL`). When omitted, it defaults to `cactus_default_backend()` — auto (best available), or the one forced with `cactus_set_backend("cpu"|"metal")`. Pin a single op: ```cpp - size_t result = graph.matmul(a, b, false, ComputeBackend::NPU); + size_t result = graph.matmul(a, b, false, ComputeBackend::METAL); ``` + Python mirrors this: `g.matmul(a, b, backend=Graph.METAL)`, default `backend=None` follows the global backend. ### Graph Construction 1. **Build once, execute many**: Construct the graph once, run with different inputs @@ -395,6 +507,65 @@ size_t similarity = fixture.graph().divide(dot_product, fixture.graph().multiply 3. **Test edge cases**: Include tests for broadcasting, empty tensors, and large inputs 4. **Check precision**: Test operations with different precision types +### Contributing Graph Operations +1. **Define the op in core graph types** +Add the new OpType in `cactus-graph/cactus_graph.h` +If the op needs additional parameters, add the fields to OpParams in the same file + +2. **Add a graph builder API** +Add a builder method in `cactus-graph/src/builder.cpp` and its declaration in +`cactus-graph/cactus_graph.h` +Follow the pattern of existing builder methods, e.g. for a new "relu" op: +```cpp +size_t CactusGraph::relu(size_t input) { + return add_node(OpType::RELU, {input}, {}); +} +``` +3. **Implement the op in the execution engine** +Implement the kernel or graph op code in the relevant file, usually in `cactus-kernels/src/` +Register the new op in the dispatch table in `cactus-graph/src/execute.cpp` for the supported backends (CPU, Metal) + +4. **Export op in FFI bindings** +- header: `cactus-graph/cactus_graph.h` (in the `extern "C"` block) +- implementation: `cactus-graph/src/graph_ffi.cpp` + +5. **Add python ctypes declaration** +Add `_lib.cactus_graph_my_new_op.argtypes/restype` in `python/cactus/bindings/cactus.py` + +6. **Add python graph wrapper** +Add `Graph.my_new_op(...)` in `python/cactus/bindings/cactus.py`, and optionally a Tensor + convenience method. + +7. **Add serialization schema entry if needed** +If your op has extra parameters that need to be saved/loaded that are not in the +default node, add new ParamField enum values. + +If the op has any graph-persistent params: + +- add any new ParamField enum values in cactus-graph/src/param_io.cpp +- add read/write logic for those fields +- add the op’s schema entry in `op_schema(...)` + +If the op has no params, you may not need to touch schema beyond maybe adding an +empty schema entry. + +The syntax pattern there is: +``` +{OpType::MY_NEW_OP, { + {ParamField::Alpha, FieldPersistence::Persistent}, + {ParamField::Mode, FieldPersistence::Persistent}, +}}, +``` +8. **Add test coverage** +Add unit tests to `cactus-graph/tests/test_graph.cpp` covering the native graph function, and +add python tests in `python/tests/test_graph.py` covering the Python API and end-to-end execution. + +and then support those fields in `write_field(...)` / `read_field(...)`. + +If a field is runtime-only, mark it RuntimeOnly instead of Persistent. + + + ### Error Handling ```cpp try { @@ -411,10 +582,10 @@ try { ```cpp CactusGraph graph; size_t x = graph.input({batch, dim}, Precision::FP16); -x = graph.linear(x, weight1, bias1); +x = graph.add(graph.matmul(x, weight1), bias1); x = graph.gelu(x); x = graph.layernorm(x, ln_weight1, ln_bias1); -x = graph.linear(x, weight2, bias2); +x = graph.add(graph.matmul(x, weight2), bias2); ``` ### Residual Connections @@ -432,4 +603,10 @@ size_t path1 = graph.matmul(input, weight1); path1 = graph.silu(path1); size_t path2 = graph.matmul(input, weight2); size_t output = graph.multiply(path1, path2); -``` \ No newline at end of file +``` + +## See Also + +- [Cactus Engine API](/docs/cactus_engine.md) — High-level inference API built on top of Cactus Graph +- [Cactus Index API](/docs/cactus_index.md) — On-device vector database for embedding storage and search +- [Runtime Compatibility](/docs/compatibility.md) — Weight versioning across releases diff --git a/docs/cactus_hybrid.md b/docs/cactus_hybrid.md new file mode 100644 index 000000000..9e4f5fbfd --- /dev/null +++ b/docs/cactus_hybrid.md @@ -0,0 +1,197 @@ +--- +title: "Cactus Hybrid Inference" +description: "Serve most queries locally, route the hard ones to the cloud. A lightweight probe on the local model decides when to ask for help." +keywords: ["hybrid inference", "cloud handoff", "routing", "on-device AI", "confidence", "edge-cloud"] +--- + +# Cactus Hybrid Inference + +Even a well-quantized local model trails a frontier cloud model on the hardest queries. +Cactus Hybrid Inference closes the gap: serve most queries locally, and route only the +ones the local model is likely to get wrong to the cloud. + +A lightweight probe (~65k parameters) attached to the local model's hidden states emits +a single confidence score per generation. Queries above a threshold are sent to the +cloud. The result: frontier-class accuracy at a fraction of the cloud cost, with most +data never leaving the device. + +## Why Hybrid? + +A 2B-parameter model running locally handles the majority of user queries well. But +there's a long tail of hard questions — complex reasoning, domain-specific knowledge, +ambiguous multimodal inputs — where a frontier cloud model (Gemini Pro, Claude, GPT-4) +would get it right and the local model won't. + +The naive solutions don't work: + +- **Always local**: misses the hard queries. Users notice. +- **Always cloud**: expensive, slow, and sends all data off-device. +- **Token entropy as confidence**: anti-predictive on transcription and vision MCQ. + High entropy often means the model is considering multiple valid candidates, not + that it's wrong. + +The probe solves this by reading the model's internal state directly — a much richer +signal than output token probabilities. + +## What the Probe Does + +The probe reads the local model's hidden states after generation and outputs a single +score: the probability the model got this query wrong. Above a threshold → route to +cloud. Below → serve locally. + +The probe: + +- **Adds negligible latency** — 65k parameters, runs after generation is already complete +- **Handles all modalities** — text, vision, and audio through the same architecture, since + it reads post-fusion hidden states where modalities are already merged +- **Transfers to audio with zero audio training** — trained on text + vision only, yet + achieves 0.79–0.88 AUROC on audio benchmarks +- **Works across quantization levels** — same probe checkpoint at FP16, 4-bit, 3-bit; + no per-quantization retraining needed + +## Results + +Evaluated on 12 hold-out benchmarks spanning text, vision, and audio — none seen during +training. Cloud target: Gemini 3.1 Pro. All results averaged over 3 seeds. + +### Routing Quality (AUROC) + +AUROC measures how well the probe separates wrong answers from right ones (higher = better). +0.5 is random, 1.0 is perfect: + +| Hold-out | Modality | Probe | Token Entropy | +|---|---|---|---| +| MMLU | text MCQ | **0.770** | 0.697 | +| MMLU-Pro | text MCQ | **0.771** | 0.692 | +| ARC-Easy | text MCQ | **0.888** | 0.655 | +| ARC-Challenge | text MCQ | **0.834** | 0.646 | +| GSM8K (3-shot) | text gen | **0.782** | 0.731 | +| MMBench-EN-Dev | vision MCQ | **0.840** | 0.435 | +| ChartQA | vision QA | **0.779** | 0.615 | +| DocVQA | vision QA | **0.781** | 0.512 | +| MMAU | audio MCQ | **0.789** | 0.517 | +| GigaSpeech | audio | **0.876** | 0.343 | +| Earnings-22 | audio | **0.839** | 0.323 | +| LibriSpeech | audio | **0.822** | 0.427 | +| **Mean** | | **0.814** | **0.549** | + +The probe outperforms token entropy on every benchmark. On audio and vision, token +entropy is actively harmful (AUROC < 0.5 means it routes the *right* answers to the +cloud and keeps the wrong ones local). + +### Cross-Modality Transfer + +The strongest result: the probe was trained on **zero audio data**, yet achieves +0.79–0.88 AUROC on four audio benchmarks (two transcription, one audio MCQ, one +out-of-domain transcription). This rules out surface-level explanations — the probe +is reading a modality-independent correctness signal from the hidden state, not +memorizing patterns from training data. + +This means hybrid inference requires no per-task probe tuning and no re-training when +the model is re-quantized. + +### How Much Handoff Is Needed? + +At FP16, the hybrid system matches Gemini 3.1 Flash-Lite accuracy on most benchmarks +by routing only 15–35% of queries to the cloud: + +| Benchmark | Handoff to match Flash-Lite (FP16) | At 4-bit | At 3-bit | +|---|---|---|---| +| ChartQA | 15–20% | 25–30% | 40–50% | +| MMBench | 30–35% | 40–45% | 50–55% | +| LibriSpeech | 25–30% | 35–40% | 55–65% | +| GigaSpeech | 30–35% | 40–45% | 50–55% | +| MMAU | 30–35% | 35–40% | 50–55% | +| MMLU-Pro | 45–55% | ~90% | n/a | + +At 4-bit quantization, the deployment cost is modest: 5–10 extra percentage points of +handoff versus FP16 on most benchmarks. At 3-bit it grows to ~20–35 points. MMLU-Pro +is the hardest case — the local model's accuracy drops sharply under quantization, +leaving less signal for the probe to work with. + +### Probe Architecture Ablations + +How much does the attention pool matter vs. simpler alternatives? + +| Head | Mean AUROC | +|---|---| +| **Learned-query attention pool + MLP** (ours) | **0.814** | +| Mean pool + MLP | 0.781 | +| Linear (logistic regression on mean pool) | 0.773 | +| Token entropy | 0.549 | + +Mean pooling captures most of the value. The attention pool opens a gap on +long-generation modalities (audio: +6 points on MMAU) and loosely-structured tasks +(MMLU-Pro: +9 points), where selectively weighting tokens within a generation matters +more than pooling-of-any-kind. + +## Where It's Weakest + +1. **When the local model is broken.** On MMLU-Pro at 4-bit, local accuracy drops to + 0.28 and probe AUROC falls toward 0.5. Probe utility scales with local-model + competence — when the model is broken, there's no correctness signal to read. + +2. **When the cloud target is worse.** On DocVQA and GSM8K, Gemini Pro is worse than + Flash-Lite, so routing to Pro can never help. A production system would select the + cloud target per task. + +3. **High handoff regime.** Above 50% handoff, the probe's marginal value over random + routing narrows. The probe's value is concentrated in the low-handoff regime — which + is the regime that matters for deployment cost. + +## Using Hybrid Inference + +Hybrid inference is built into the Cactus Engine. Set a confidence threshold in the +completion options: + +```c +const char* options = R"({ + "confidence_threshold": 0.7, + "auto_handoff": true, + "cloud_timeout_ms": 15000 +})"; + +char response[4096]; +cactus_complete(model, messages, response, sizeof(response), options, NULL, NULL, NULL, NULL, 0); +``` + +When the model's confidence drops below the threshold, the request is automatically +routed to the cloud. The response JSON indicates whether handoff occurred: + +```json +{ + "success": true, + "cloud_handoff": true, + "response": "Cloud-provided answer.", + "confidence": 0.18, + "confidence_threshold": 0.7 +} +``` + +In Python: + +```python +import json +from cactus import get_bundle_dir, cactus_init, cactus_complete + +model = cactus_init(str(get_bundle_dir("google/gemma-4-E2B-it")), None, False) +options = json.dumps({ + "confidence_threshold": 0.7, + "auto_handoff": True +}) +messages = json.dumps([{"role": "user", "content": "Explain quantum entanglement"}]) +result = cactus_complete(model, messages, options, None, None) + +if result["cloud_handoff"]: + print("Answered by cloud") +else: + print(f"Answered locally (confidence: {result['confidence']:.2f})") +``` + +Set your cloud API key with `cactus auth` to enable automatic handoff. + +## See Also + +- [Cactus Engine API](/docs/cactus_engine.md) — `confidence_threshold` and `auto_handoff` options +- [Cactus Quants](/docs/cactus_quants.md) — Quantization that the probe works across +- [Gemma 4 on Cactus](/blog/gemma4.md) — The base model used for hybrid inference diff --git a/docs/cactus_index.md b/docs/cactus_index.md index da9fca261..a0e669117 100644 --- a/docs/cactus_index.md +++ b/docs/cactus_index.md @@ -1,3 +1,9 @@ +--- +title: "Cactus Index FFI API Reference" +description: "On-device vector database C API for storing and querying embeddings. Supports cosine similarity search, memory-mapped files, and RAG applications." +keywords: ["vector database", "embeddings", "cosine similarity", "RAG", "on-device search", "C FFI"] +--- + # Cactus Index FFI Documentation The Cactus Index provides a clean C FFI (Foreign Function Interface) for integrating a vector database into various applications. This documentation covers all available functions, their parameters, and usage examples. @@ -5,7 +11,7 @@ The Cactus Index provides a clean C FFI (Foreign Function Interface) for integra ## Getting Started The index uses memory-mapped files: -- `index.bin`: Embeddings (FP16) and metadata pointers +- `index.bin`: Embeddings (FP16) plus per-doc offsets into `data.bin` - `data.bin`: Document content and metadata (UTF-8) All embeddings are automatically normalized to unit length @@ -274,7 +280,7 @@ int cactus_index_query( **Returns:** 0 on success, -1 on error (if buffers too small, no data copied) -**Example** +**Example:** ```c float query1[768] = {/* ... */}; float query2[768] = {/* ... */}; @@ -332,6 +338,15 @@ Releases all resources associated with the index. void cactus_index_destroy(cactus_index_t index); ``` +**Parameters:** +- `index`: Index handle from `cactus_index_init` (required) + +**Example:** +```c +cactus_index_destroy(index); +index = NULL; +``` + ## Examples ### Create and Populate @@ -547,3 +562,12 @@ cactus_index_destroy(new_index); 4. **Thread Safety**: One index instance per thread 5. **Batching**: Add 100-1000 documents per call for best performance 6. **Errors**: Use `cactus_get_last_error()` for error details + +## See Also + +- [Cactus Engine API](/docs/cactus_engine.md) — LLM inference, embeddings (`cactus_embed`), and RAG query APIs +- [Cactus Graph API](/docs/cactus_graph.md) — Low-level computational graph for custom tensor operations +- [Python Binding](/python/) — Python bindings with vector index support +- [Swift Binding](/bindings/swift/) — Swift vector index functions +- [Kotlin Binding](/bindings/kotlin/) — Kotlin vector index functions +- [Flutter Binding](/bindings/flutter/) — Dart vector index functions diff --git a/docs/cactus_kernels.md b/docs/cactus_kernels.md new file mode 100644 index 000000000..33cbb6dbe --- /dev/null +++ b/docs/cactus_kernels.md @@ -0,0 +1,473 @@ +--- +title: "Cactus Kernels API Reference" +description: "ARM NEON SIMD kernel library for matrix multiplication, attention, convolution, quantization, DSP, and image processing. The computational foundation of the Cactus inference engine." +keywords: ["ARM NEON", "SIMD", "kernels", "matmul", "attention", "quantization", "CQ", "DSP", "image processing"] +--- + +# Cactus Kernels API Documentation + +The Cactus Kernels layer (`cactus-kernels`) provides hand-tuned ARM NEON SIMD implementations of the core operations used by the graph and engine layers. All kernels operate on FP16 (`__fp16`) data by default and are designed for mobile ARM chips (Apple Silicon, Snapdragon, Exynos, Tensor, etc.). + +Header: `cactus-kernels/cactus_kernels.h` + +## Precision Types + +```cpp +enum class Precision { + INT8, // 8-bit integer (KV cache quantization) + FP16, // 16-bit float (default compute precision) + FP32, // 32-bit float + CQ1, // Cactus Quant 1-bit + CQ2, // Cactus Quant 2-bit + CQ3, // Cactus Quant 3-bit + CQ4 // Cactus Quant 4-bit +}; +``` + +## Element-wise Arithmetic + +```cpp +void cactus_add_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); +void cactus_subtract_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); +void cactus_multiply_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); +void cactus_divide_f16(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); +void cactus_add_f16_clipped(const __fp16* a, const __fp16* b, __fp16* output, size_t num_elements); +void cactus_add_scaled_f16(const __fp16* base, const __fp16* src, __fp16* output, size_t num_elements, float scale); +``` + +### Broadcast Variants + +For tensors with different shapes, broadcast versions handle stride-based indexing: + +```cpp +void cactus_add_broadcast_f16( + const __fp16* a, const __fp16* b, __fp16* output, + const size_t* a_strides, const size_t* b_strides, + const size_t* output_shape, size_t ndim); +// Also: subtract, multiply, divide broadcast variants +``` + +### Scalar Operations + +```cpp +enum class ScalarOpType { ADD, SUBTRACT, MULTIPLY, DIVIDE, ABS, EXP, POW, SQRT, COS, SIN, LOG }; + +void cactus_scalar_op_f16( + const __fp16* input, __fp16* output, size_t num_elements, + float scalar_value, ScalarOpType op_type); +``` + +## Reductions + +```cpp +// Full reductions (all elements) +double cactus_sum_all_f16(const __fp16* data, size_t num_elements); +double cactus_mean_all_f16(const __fp16* data, size_t num_elements); +double cactus_variance_all_f16(const __fp16* data, size_t num_elements); +__fp16 cactus_min_all_f16(const __fp16* data, size_t num_elements); +__fp16 cactus_max_all_f16(const __fp16* data, size_t num_elements); + +// Axis reductions +void cactus_sum_axis_f16(const __fp16* input, __fp16* output, + size_t outer_size, size_t axis_size, size_t inner_size); +// Also: mean, variance, min, max axis variants +``` + +## Matrix Multiplication + +### FP16 Matmul + +```cpp +void cactus_matmul_f16( + const __fp16* a, // (M, K) + const __fp16* b_transposed, // (N, K) — pre-transposed RHS + __fp16* c, // (M, N) output + size_t M, size_t K, size_t N); +``` + +### Cactus Quant (CQ) Matmul + +CQ quantization uses Hadamard rotation + per-group codebooks for 1-4 bit weight compression: + +```cpp +struct CactusQuantMatrix { + uint32_t bits; // 1, 2, 3, or 4 + uint32_t K, N; // matrix dimensions + uint32_t group_size; // elements per codebook group + uint32_t num_groups; + uint32_t flags; // ORTHOGONAL, INTERLEAVED_4ROW + const __fp16* codebook; + const __fp16* input_scale; + const __fp16* input_scale_recip; + const __fp16* norms; + const uint8_t* packed_indices; + const int8_t* left_signs; + const int8_t* right_signs; + const uint32_t* permutation; + const __fp16* rotation; + const int8_t* expanded; + const float* norm_f32; +}; + +// Unified dispatch (picks gemv or gemm based on M) +void cactus_quant_matmul(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); + +// Bit-width specific +void cactus_quant_4bit_gemv(const CactusQuantMatrix* W, const __fp16* x, __fp16* y); +void cactus_quant_4bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); +void cactus_quant_2bit_gemv(const CactusQuantMatrix* W, const __fp16* x, __fp16* y); +void cactus_quant_2bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); +void cactus_quant_1bit_gemv(const CactusQuantMatrix* W, const __fp16* x, __fp16* y); +void cactus_quant_1bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); +void cactus_quant_3bit_gemv(const CactusQuantMatrix* W, const __fp16* x, __fp16* y); +void cactus_quant_3bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); + +// Orthogonal rotation variant +void cactus_quant_orthogonal_matmul(const CactusQuantMatrix* W, const __fp16* A, uint32_t M, __fp16* C); +``` + +### Embedding Dequantization + +```cpp +// Hadamard-rotated embedding row dequantization +void cactus_quant_dequantize_hadamard_embedding_row( + uint32_t bits, uint32_t hidden_dim, uint32_t group_size, uint32_t num_groups, + size_t row, const uint8_t* packed_base, const __fp16* codebook, + const __fp16* norms, const __fp16* input_scale_recip, + const int8_t* left_signs, const int8_t* right_signs, + const uint32_t* permutation, __fp16* out_row); + +// Orthogonal rotation variant +void cactus_quant_dequantize_orthogonal_embedding_row( + uint32_t bits, uint32_t K, size_t row, + const uint8_t* packed_base, const __fp16* codebook, + const __fp16* norms, const __fp16* input_scale_recip, + const __fp16* rotation, uint32_t flags, __fp16* out_row); +``` + +## Attention + +### Standard FP16 Attention + +```cpp +void cactus_attention_f16( + const __fp16* queries, // (batch, seq, num_q_heads, head_dim) + const __fp16* keys, // (batch, kv_seq, num_kv_heads, head_dim) + const __fp16* values, // (batch, kv_seq, num_kv_heads, head_dim) + __fp16* output, + size_t batch_size, size_t seq_len, size_t kv_seq_len, + size_t num_q_heads, size_t num_kv_heads, size_t head_dim, + float scale, + const __fp16* mask, + size_t position_offset = 0, + size_t window_size = 0, // 0 = full context, >0 = sliding window + bool is_causal = true, + bool mask_is_additive = false, + bool mask_per_head = false, + size_t v_head_dim = 0, // 0 = same as head_dim + float logit_cap = 0.0f); +``` + +Supports grouped-query attention (GQA) when `num_q_heads > num_kv_heads`, sliding window attention, additive masks, and logit soft-capping. + +### Hybrid INT8/FP16 Attention (KV Cache) + +For models with INT8-quantized KV cache and FP16 new tokens: + +```cpp +void cactus_attention_hybrid_int8_fp16( + const __fp16* queries, + const int8_t* keys_cached, // INT8 quantized cache + const int8_t* values_cached, + const float* k_scales, // per-group dequant scales + const float* v_scales, + const __fp16* keys_new, // FP16 new tokens + const __fp16* values_new, + __fp16* output, + size_t batch_size, size_t seq_len, + size_t cache_len, size_t new_len, + size_t num_q_heads, size_t num_kv_heads, size_t head_dim, + float scale, + size_t position_offset = 0, + bool is_causal = true, + size_t window_size = 0, + size_t group_size = 32, // KV_QUANT_GROUP_SIZE + size_t v_head_dim = 0); +``` + +## Normalization + +```cpp +void cactus_rms_norm_f16( + const __fp16* input, const __fp16* weight, __fp16* output, + size_t batch_size, size_t dims, float eps); + +void cactus_layer_norm_f16( + const __fp16* input, const __fp16* weight, const __fp16* bias, __fp16* output, + size_t batch_size, size_t dims, float eps); + +void cactus_batchnorm_f16( + const __fp16* input, const float* weight, const float* bias, + const float* running_mean, const float* running_var, __fp16* output, + size_t outer_size, size_t channels, size_t inner_size, float epsilon); + +void cactus_softmax_f16( + const __fp16* input, __fp16* output, + size_t batch_size, size_t seq_len, size_t vocab_size); +``` + +## Positional Encoding + +```cpp +void cactus_rope_f16( + const __fp16* input, __fp16* output, + size_t batch_size, size_t seq_len, size_t num_heads, size_t head_dim, + size_t start_pos, float theta); + +void cactus_gpt_j_rope_f16( + const __fp16* input, __fp16* output, + size_t batch_size, size_t seq_len, size_t num_heads, size_t head_dim, + size_t rot_dim, size_t start_pos, float theta); +``` + +## Activation Functions + +```cpp +void cactus_relu_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_leaky_relu_f16(const __fp16* input, __fp16* output, size_t num_elements, float negative_slope); +void cactus_clamp_f16(const __fp16* input, __fp16* output, size_t num_elements, float lo, float hi); +void cactus_silu_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_gelu_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_gelu_f16_erf(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_sigmoid_f16(const __fp16* input, __fp16* output, size_t num_elements); +void cactus_tanh_f16(const __fp16* input, __fp16* output, size_t num_elements); + +void cactus_glu_f16(const __fp16* input, __fp16* output, + size_t outer_size, size_t split_size, size_t inner_size); +``` + +## Convolution + +### 1D Convolution + +```cpp +void cactus_conv1d_f16(const __fp16* input, const __fp16* weight, const __fp16* bias, + __fp16* output, size_t N, size_t L, size_t C_in, size_t C_out, size_t K, size_t stride); + +void cactus_conv1d_f16_k3(const __fp16* input, const __fp16* weight, __fp16* output, + size_t N, size_t L, size_t C_in, size_t C_out, size_t stride); + +void cactus_conv1d_causal_depthwise_f16(const __fp16* input, const __fp16* weight, __fp16* output, + size_t N, size_t L, size_t C, size_t K, size_t dilation); + +void cactus_conv1d_same_depthwise_f16_k9(const __fp16* input, const __fp16* weight, + const __fp16* bias, __fp16* output, size_t N, size_t L, size_t C); + +void cactus_conv1d_pointwise_f16_gemm(const __fp16* input, const __fp16* weight, + const __fp16* bias, __fp16* output, size_t N, size_t L, size_t C_in, size_t C_out); +``` + +### 2D Convolution + +```cpp +void cactus_conv2d_f16_k3s2p1_nchw(...); // 3x3 stride-2 pad-1 +void cactus_conv2d_depthwise_f16_k3s2p1_nchw(...); // depthwise 3x3 stride-2 +void cactus_conv2d_pointwise_f16_1x1_nchw_gemm(...); // 1x1 pointwise via GEMM +void cactus_conv2d_f16_k3s1p1_nchw(...); // 3x3 stride-1 pad-1 +``` + +## Recurrent Layers + +```cpp +void cactus_lstm_cell_f16( + const __fp16* x_input, const __fp16* h_prev, const __fp16* c_prev, + const __fp16* weight_ih, const __fp16* weight_hh, + const __fp16* bias_ih, const __fp16* bias_hh, + __fp16* h_new, __fp16* c_new, + size_t batch_size, size_t input_size, size_t hidden_size); + +void cactus_bilstm_sequence_f16( + const __fp16* input, + const __fp16* weight_ih_fwd, const __fp16* weight_hh_fwd, const __fp16* bias_ih_fwd, const __fp16* bias_hh_fwd, + const __fp16* weight_ih_bwd, const __fp16* weight_hh_bwd, const __fp16* bias_ih_bwd, const __fp16* bias_hh_bwd, + __fp16* output, size_t batch_size, size_t seq_len, size_t input_size, size_t hidden_size); + +void cactus_gated_deltanet_decode_f16( + const __fp16* q_data, const __fp16* k_data, const __fp16* v_data, + const __fp16* g_data, const __fp16* b_data, const __fp16* s_data, + __fp16* out, size_t B, size_t Hq, size_t Hv, size_t K, size_t V, float scale); + +void cactus_gated_deltanet_prefill_f16( + const __fp16* q_data, const __fp16* k_data, const __fp16* v_data, + const __fp16* g_data, const __fp16* b_data, const __fp16* s_data, + __fp16* out, size_t B, size_t T, size_t Hq, size_t Hv, + size_t K, size_t V, size_t requested_chunk_size, float scale); +``` + +## Sampling + +```cpp +void cactus_sample_f16( + const __fp16* logits, uint32_t* output, size_t vocab_size, + float temperature, float top_p, size_t top_k, size_t random_seed, + const float* bias_values = nullptr, + const uint32_t* bias_indices = nullptr, + size_t bias_count = 0); + +// Extended version with min_p and repetition penalty +void cactus_sample_f16_ex( + const __fp16* logits, uint32_t* output, size_t vocab_size, + float temperature, float top_p, float min_p, float repetition_penalty, + size_t top_k, size_t random_seed, + const float* bias_values = nullptr, + const uint32_t* bias_indices = nullptr, + size_t bias_count = 0); +``` + +## DSP (Digital Signal Processing) + +```cpp +void cactus_rfft_f32_1d(const float* input, float* output, size_t n, const char* norm); +void cactus_irfft_f32_1d(const float* input, float* output, size_t n, const char* norm); +float cactus_hertz_to_mel(float freq, const char* mel_scale); +float cactus_mel_to_hertz(float mels, const char* mel_scale); + +void cactus_generate_mel_filter_bank( + float* mel_filters, int num_frequency_bins, int num_mel_filters, + float min_frequency, float max_frequency, int sampling_rate, + const char* norm, const char* mel_scale, bool triangularize_in_mel_space); + +void cactus_compute_spectrogram_f32( + const float* waveform, size_t waveform_length, + const float* window, size_t window_length, + size_t frame_length, size_t hop_length, const size_t* fft_length, + float* spectrogram, float power, + bool center, const char* pad_mode, bool onesided, + float dither, const float* preemphasis, + const float* mel_filters, size_t mel_filters_size, + float mel_floor, const char* log_mel, + float reference, float min_value, const float* db_range, + bool remove_dc_offset); +``` + +## Image Processing + +```cpp +// Load image from file (JPEG, PNG, BMP, etc.) +unsigned char* cactus_image_load(const char* path, int* width, int* height, int* channels, int desired_channels); +void cactus_image_free(unsigned char* data); + +// Resize +void cactus_image_resize_uint8(const unsigned char* input, int src_w, int src_h, + unsigned char* output, int dst_w, int dst_h, int channels); +void cactus_image_resize_float(const float* input, int src_w, int src_h, + float* output, int dst_w, int dst_h, int channels); + +// Normalize with mean/std +void cactus_image_normalize(const float* input, float* output, + int width, int height, int channels, + float rescale_factor, const float* mean, const float* std_dev); + +// Convert to vision transformer patches +void cactus_image_to_patches(const float* image, float* patches, + int width, int height, int channels, int patch_size); + +// Channel conversion +void cactus_image_convert_to_rgb(const unsigned char* input, unsigned char* output, + int width, int height, int channels); +``` + +## Precision Conversion + +```cpp +void cactus_fp16_to_fp32(const __fp16* src, float* dst, size_t count); +void cactus_fp32_to_fp16(const float* src, __fp16* dst, size_t count); +void cactus_int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale = 1.0f); +void cactus_fp16_to_int8(const __fp16* src, int8_t* dst, size_t count, float scale = 1.0f); +void cactus_int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale = 1.0f); +void cactus_fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale = 1.0f); +float cactus_fp16_max_abs(const __fp16* src, size_t count); +``` + +## KV Cache Quantization + +```cpp +constexpr size_t KV_QUANT_GROUP_SIZE = 32; + +void cactus_quantize_kv_fp16_to_int8( + const __fp16* src, int8_t* dst, float* scales, + size_t seq_len, size_t kv_heads, size_t head_dim, + size_t group_size = KV_QUANT_GROUP_SIZE); + +inline size_t kv_scales_count( + size_t seq_len, size_t kv_heads, size_t head_dim, + size_t group_size = KV_QUANT_GROUP_SIZE); +``` + +## Miscellaneous + +```cpp +void cactus_stft_f16(const __fp16* input, const __fp16* weight, __fp16* output, + size_t N, size_t L, size_t C_in, size_t C_out, size_t K, size_t stride, size_t num_fft_bins); + +void cactus_bilinear_interpolation_f16(const __fp16* input, __fp16* output, + size_t src_height, size_t src_width, size_t embed_dim, + size_t dst_height, size_t dst_width, bool align_corners = true); + +void cactus_maxpool1d_f16(const __fp16* input, __fp16* output, + size_t batch_size, size_t channels, size_t input_length, + size_t kernel_size, size_t stride); + +void cactus_altup_predict_f16(const __fp16* coefs, const __fp16* const* streams, + __fp16* output, size_t n, size_t seq_len, size_t hidden_dim); + +void cactus_altup_correct_f16(const __fp16* coefs, const __fp16* innovation, + const __fp16* const* predictions, __fp16* output, + size_t n, size_t seq_len, size_t hidden_dim); + +void cactus_gaussian_topk_f16(const __fp16* input, __fp16* output, + size_t rows, size_t cols, float ppf); +``` + +## Directory Structure + +``` +cactus-kernels/ + cactus_kernels.h # public API (this file) + libs/ + stb_image.h # vendored image loading + stb_image_resize2.h # vendored image resizing + src/ + threading.h # thread pool utilities + wav.h # WAV file loading + 16 kHz resampling + attention.cpp # attention kernels (FP16) + attention_hybrid.cpp # hybrid INT8/FP16 attention + blas.cpp # BLAS-backed paths + conv.cpp # conv1d variants, STFT + conv2d.cpp # conv2d variants + dsp.cpp # rfft, irfft, mel filter bank, spectrogram + fused.cpp # fused op kernels + image.cpp # image load/resize/normalize/patches + lstm.cpp # LSTM cell, BiLSTM sequence + matmul.cpp # FP16/INT8 GEMM + nn.cpp # activations (relu/silu/gelu/...), glu, clamp + norms_rope.cpp # rms_norm, layer_norm, RoPE + quants.cpp # CQ 1-4 bit GEMV/GEMM, dequantization + reduce.cpp # reductions (sum/mean/var/min/max + axis variants) + scalar.cpp # scalar elementwise ops + tests/ + test_utils.h # test runner, fp16 comparison helpers + test_attention.cpp + test_conv.cpp + test_dsp.cpp + test_elementwise.cpp + test_matmul.cpp + test_quant.cpp + test_reduce.cpp +``` + +## See Also + +- [Cactus Graph API](/docs/cactus_graph.md) — Computation graph built on top of these kernels +- [Cactus Engine API](/docs/cactus_engine.md) — High-level inference API +- [TurboQuant-H](/blog/turboquant-h.md) — 2-bit embedding quantization using the CQ kernel infrastructure diff --git a/docs/cactus_quants.md b/docs/cactus_quants.md new file mode 100644 index 000000000..a656803bd --- /dev/null +++ b/docs/cactus_quants.md @@ -0,0 +1,210 @@ +--- +title: "Cactus Quants (CQ)" +description: "Rotation-and-codebook post-training quantization for on-device multimodal models. One recipe for every weight tensor — linears, embeddings, vision towers, audio towers — from 4-bit down to 1-bit." +keywords: ["quantization", "PTQ", "CQ", "Cactus Quants", "on-device AI", "mobile inference", "Hadamard", "codebook", "GPTQ", "AWQ", "HQQ"] +--- + +# Cactus Quants + +Cactus Quants (CQ) is the post-training quantization system used by Cactus to compress +models for on-device deployment. It applies a single rotation-and-codebook recipe to +every weight tensor in a multimodal model — transformer linears, vision encoders, audio +encoders, cross-modal bridges, per-layer embedding tables, and the shared token embedding +— at bit widths from 4-bit down to 1-bit. + +## Why Another Quantization Method? + +A 2B-parameter model in FP16 is roughly 4 GB on disk. At 4-bit it's 1 GB. At 2-bit +it's 500 MB. Edge deployment is a storage and memory problem. + +Existing PTQ methods (GPTQ, AWQ, HQQ) work well at 4 bits on transformer linear layers. +But they were not designed for the full multimodal setting: + +- **They only quantize linears.** Per-layer embedding tables, modality towers, and + vision-language bridges are left at FP16. In models like Gemma 4 E2B, per-layer + embeddings alone are 60% of the model — more than half the storage is untouched. + +- **They break below 4 bits.** At 3-bit, accuracy drops 15-30 points on hard benchmarks. + At 2-bit, every baseline collapses to near-random on generation tasks. + +- **They don't handle embeddings.** Embedding tables are gathered by token ID, not + multiplied against activations. Methods that depend on an activation Hessian + (GPTQ, AWQ) have no signal to work with. + +CQ solves all three problems with one unified recipe. + +## Results + +We evaluated CQ against GPTQ, AWQ, and HQQ on three sub-2B multimodal model families: +**Gemma 4 E2B**, **Qwen3-VL-2B**, and **LFM-2.5-VL-1.6B**. All baselines use group +size 128 where applicable. Results averaged over 3 seeds. + +### At 4 Bits: Competitive with the Best Baseline + +At 4-bit, the field is closely contested. CQ leads or ties on most metrics across +Gemma 4 and Qwen3-VL-2B: + +| | CQ-4bit | GPTQ-4bit | AWQ-4bit | HQQ-4bit | FP16 | +|---|---|---|---|---|---| +| MMLU (Gemma 4) | 59.45 | 59.48 | **59.66** | 57.23 | 62.33 | +| GSM8K (Gemma 4) | **71.20** | 68.33 | 70.53 | 58.87 | 73.67 | +| HumanEval (Gemma 4) | **57.11** | 50.61 | 52.64 | 43.29 | 54.88 | +| ARC-E (Gemma 4) | **73.73** | 72.00 | 73.00 | 72.00 | 73.80 | + +On LFM-2.5-VL-1.6B, HQQ leads at 4-bit on several metrics — CQ's advantage doesn't +appear until bits get tighter. + +### At 3 Bits: CQ Dominates + +This is where the gap opens. At 3-bit, CQ sweeps all 8 text metrics on Gemma 4, with +margins of up to **+32 points** on the hardest benchmarks: + +| | CQ-3bit | GPTQ-3bit | AWQ-3bit | HQQ-3bit | +|---|---|---|---|---| +| GSM8K (Gemma 4) | **52.67** | 16.40 | 21.87 | 1.13 | +| HumanEval (Gemma 4) | **45.12** | 8.54 | 13.01 | 4.88 | +| MMLU (Gemma 4) | **54.56** | 49.26 | 49.22 | 36.02 | +| BFCL Parallel (Gemma 4) | **81.17** | 41.00 | 58.00 | 0.50 | + +The same pattern holds on Qwen3-VL-2B (CQ leads 5 of 7 text metrics, up to +30 points +on GSM8K) and LFM-2.5-VL-1.6B (CQ leads 5 of 8 metrics). + +### At 2 Bits: CQ Is the Only Survivor + +At 2-bit, every baseline collapses. CQ is the only method retaining usable accuracy: + +| | CQ-2bit | GPTQ-2bit | AWQ-2bit | HQQ-2bit | +|---|---|---|---|---| +| ARC-E (Gemma 4) | **50.80** | 24.20 | 27.40 | 23.00 | +| MMLU (Gemma 4) | **33.18** | 23.40 | 24.28 | 25.06 | +| BFCL Simple (Gemma 4) | **18.75** | 0.00 | 0.00 | 0.00 | +| BFCL Multi (Gemma 4) | **13.67** | 0.00 | 0.00 | 0.00 | + +On Qwen3-VL-2B at 2-bit, CQ retains 50.5 ARC-E and 29.4 MMLU while baselines are +at near-random. On vision, CQ holds 27.9 MMMU and 31.2 AI2D where the next-best +method scores below 7. + +### Function Calling + +Tool calling requires longer-horizon coherence than multiple-choice reasoning, making +it the hardest task to preserve under quantization: + +| | CQ | GPTQ | AWQ | HQQ | +|---|---|---|---|---| +| 4-bit BFCL Parallel-Multi (Gemma 4) | **83.33** | 80.33 | 80.50 | 81.00 | +| 3-bit BFCL Parallel-Multi (Gemma 4) | **79.67** | 38.83 | 31.33 | 0.00 | +| 2-bit BFCL Parallel-Multi (Gemma 4) | **1.33** | 0.00 | 0.00 | 0.00 | + +At 3-bit, CQ retains near-FP16 function calling while baselines lose 40-80 points. + +### Multimodal: Vision and Audio + +CQ quantizes modality towers with the same recipe. Under a joint (M, L) schedule where +M is the modality-tower bit width and L is the LLM bit width: + +| Setting | CQ ChartQA | AWQ ChartQA | HQQ ChartQA | +|---|---|---|---| +| M4 + L4 | **51.00** | 49.00 | 49.33 | +| M2 + L4 | **48.44** | 14.22 | 14.33 | + +At 2-bit modality tower with 4-bit LLM, AWQ and HQQ collapse on ChartQA from ~49 to +~14 while CQ retains 48.44. + +### Transcription (Whisper and Parakeet) + +On speech models, CQ is the best method on Parakeet at every bit width. The standout: +2-bit Parakeet-1.1B where CQ holds 0.147 WER while every baseline collapses to ≥ 1.0. + +| | CQ | GPTQ | AWQ | HQQ | +|---|---|---|---|---| +| Parakeet-1.1B 4-bit WER | **0.115** | 0.125 | 0.116 | 0.125 | +| Parakeet-1.1B 3-bit WER | **0.115** | 0.298 | 0.157 | 0.298 | +| Parakeet-1.1B 2-bit WER | **0.147** | 1.043 | 1.084 | 1.043 | + +On Whisper, CQ is competitive but doesn't consistently lead — it's within 1% of the +best method at most settings. + +### Mixed-Precision Allocation + +CQ supports sensitivity-driven mixed precision: rank tensors by their impact on +downstream accuracy, keep high-sensitivity tensors at higher bits, push low-sensitivity +ones lower. + +The sweet spot on Gemma 4: allocate 4-bit to 68 high-sensitivity LLM units and 3-bit +to the remaining 207 (3.26-bit average). For CQ, this costs ~2 points vs uniform 4-bit. +For GPTQ/AWQ/HQQ, the same allocation costs 7-15 points — mixed precision is +approximately free for CQ but expensive for baselines. + +A systematic sweep shows that promoting just 25% of LLM tensors from 2-bit to 3-bit +(2.50 average bits) recovers ~75% of the CQ2-to-CQ3 accuracy gap. The marginal value +of a bit is sharply non-uniform across tensors. + +### Production Bundles + +Real deployment uses different bit widths for different tensor types. The production +bundle for Gemma 4 E2B: + +| Component | Bits | Rationale | +|---|---|---| +| LLM linears | 4 | Accuracy-sensitive | +| Per-layer embeddings (PLI) | 2 | 60% of model, compresses 7.5x | +| Shared token embeddings | 4 | Small, accuracy-sensitive | +| Vision tower | 2 | Compresses gracefully | +| Audio tower | 2 | Holds understanding, some WER cost | + +**Result: 4.79 GB (FP16) → ~0.9 GB on disk. 5.3x reduction.** + +At ~3.0 bits-per-weight, this production bundle: + +- Matches or exceeds Unsloth UD-Q2_K_XL (~3.4 bpw) on 10 of 12 text/tool metrics + despite a tighter bit budget +- Retains MMLU 59.14 (FP16: 62.33), GSM8K 74.00 (FP16: 73.67) +- Retains BFCL Parallel-Multi 78.00 (FP16: 78.00) +- Retains MMAU 58.67 (FP16: 56.00) +- Largest concession: LibriSpeech WER degrades from ~5.9 to 10.42 at 2-bit audio tower + +The mixed-precision bundle at ~2.7 bpw is the deployment sweet spot: half the bit budget +of Unsloth UD-Q4_K_XL while matching it on most metrics. + +## Where CQ Is Weakest + +Three patterns are visible: + +1. **Audio transcription at 2-bit modality tower.** LibriSpeech WER degrades from ~5.9 + to ~10.4. Audio understanding (MMAU) is preserved — the degradation is concentrated + in token-level acoustic decoding, not higher-level reasoning. + +2. **Tool calling at 2-bit.** BFCL Parallel-Multi is near zero at 2-bit for every method + including CQ. Function calling requires longer-horizon coherence that 2-bit noise breaks. + +3. **LFM at 4-bit.** On LFM-2.5-VL-1.6B, HQQ wins several metrics at 4-bit. CQ's + advantage only appears at 3-bit and below on this family, suggesting that at sufficient + bit budget the choice of quantizer matters less than per-model factors. + +## Using CQ Weights + +`cactus convert` quantizes the weights to CQ and builds the runtime bundle +(graph + manifest) in one step. Run via `cactus run`: +`cactus convert ` builds a runnable bundle locally; `cactus run` runs a +bundle path or a model id. + +```bash +cactus convert google/gemma-4-E2B-it ./gemma4-bundle +cactus convert google/gemma-4-E2B-it ./gemma4-bundle --bits 2 +cactus run ./gemma4-bundle +``` + +For models on huggingface.co/Cactus-Compute, `cactus run ` (or `cactus +download `) fetches the pre-built bundle, building locally if none is +published. + +The CQ matmul kernels live in `cactus-kernels/src/matmul.cpp`. At inference, +the kernel applies a Walsh-Hadamard transform to the FP16 activation (not the weight), +then performs a fused codebook-lookup + norm-scale + matmul in a single NEON pass. + +## See Also + +- [Cactus Kernels](/docs/cactus_kernels.md) — CQ GEMV/GEMM kernel implementations +- [TurboQuant-H](/blog/turboquant-h.md) — The 2-bit embedding quantization technique for AltUp models +- [Cactus Engine API](/docs/cactus_engine.md) — Run quantized models via the C API +- [Fine-tuning Guide](/docs/finetuning.md) — Convert LoRA fine-tunes to CQ format diff --git a/docs/cactus_transpiler.md b/docs/cactus_transpiler.md new file mode 100644 index 000000000..de0a7e601 --- /dev/null +++ b/docs/cactus_transpiler.md @@ -0,0 +1,518 @@ +--- +title: "Cactus Transpiler" +description: "Convert any PyTorch model to run on-device with Cactus. Step-by-step guide from a simple PyTorch model to on-device inference." +keywords: ["transpiler", "PyTorch", "torch.export", "on-device AI", "model conversion", "custom models"] +--- + +# Cactus Transpiler + +The Cactus Transpiler takes a PyTorch model and turns it into a Cactus runtime graph +that runs on ARM devices. You write normal PyTorch, and the transpiler handles the rest: +capturing the computation graph, fusing ops, quantizing weights, and producing a bundle +that loads with zero-copy memory mapping. + +## From PyTorch to On-Device in 2 Steps + +### Step 1: Write a PyTorch Model + +Any `torch.nn.Module` works. Here's a small feedforward classifier: + +```python +import torch +import torch.nn as nn + +class TextClassifier(nn.Module): + def __init__(self, vocab_size=32000, embed_dim=256, num_classes=4): + super().__init__() + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.fc1 = nn.Linear(embed_dim, 128) + self.fc2 = nn.Linear(128, num_classes) + + def forward(self, input_ids): + x = self.embedding(input_ids) # (batch, seq, embed_dim) + x = x.mean(dim=1) # (batch, embed_dim) + x = torch.nn.functional.relu(self.fc1(x)) + x = self.fc2(x) + return x + +model = TextClassifier() +torch.save(model.state_dict(), "classifier.pt") +``` + +### Step 2: Convert and Run + +`cactus convert` quantizes the weights to Cactus format (writing a +`weights_manifest.json`) and then captures the graph, optimizes it, and saves a +self-contained bundle — all in one step: + +```bash +cactus convert ./classifier.pt ./classifier-cactus --bits 4 + +# Run the converted bundle +cactus run ./classifier-cactus +``` + +That's it. The output is a self-contained bundle with memory-mapped weights that +loads in milliseconds on any ARM device. + +--- + +## How It Works + +The transpiler is a six-stage compiler pipeline: + +``` +PyTorch nn.Module + | + v ++------------------+ +| 1. Canonicalize | Wrap the model in a task-specific adapter +| model | (causal LM, encoder, multimodal, etc.) ++--------+---------+ + | + v ++------------------+ +| 2. torch.export | Capture the forward pass as an FX graph +| | (no Python control flow, tensor-only) ++--------+---------+ + | + v ++------------------+ +| 3. Import to IR | FX nodes -> canonical IRGraph +| | Bind converted weight files from manifest ++--------+---------+ + | + v ++------------------+ +| 4. Canonicalize | Normalize ops, constant-fold, remove no-ops, +| IR | insert precision casts, dead code elimination ++--------+---------+ + | + v ++------------------+ +| 5. Fuse patterns | Collapse subgraphs into single ops: +| | RMSNorm, RoPE, attention, MLP, LSTM, etc. ++--------+---------+ + | + v ++------------------+ +| 6. Lower to | IRNodes -> CactusGraph ops +| CactusGraph | Constants -> mmap_weights() or inline tensors ++------------------+ + | + v + Saved bundle: graph.cactus + manifest.json + weights +``` + +### Stage 1: Canonicalize the Model + +The transpiler doesn't try to handle every model's forward signature as-is. It wraps +the model in a thin adapter that exposes a stable, task-specific interface: + +| Task | Input | Output | +|------|-------|--------| +| `causal_lm_logits` | `input_ids` | next-token logits | +| `multimodal_causal_lm_logits` | `input_ids`, image/audio features | next-token logits | +| `encoder_hidden_states` | `input_features` | hidden states | +| `ctc_logits` | `input_features` | CTC logits | +| `seq2seq_transcription` | `input_features` | transcription token IDs | +| `tdt_transcription` | `input_features` | TDT token IDs | + +This is where model-family-specific knowledge lives (Gemma4, Qwen, LFM2, Parakeet, etc.). +The adapter also injects `use_cache=False` and `return_dict=False` so the export +boundary stays tensor-only. + +### Stage 2: Capture with `torch.export` + +The adapted model is captured using `torch.export`, producing an FX graph where every +operation is a canonical ATen function call (e.g. `aten.linear.default`, `aten.silu.default`). + +This is the generality boundary. After capture, the transpiler works with ATen op names, +not layer names. It doesn't care whether your model calls its attention layer `self_attn` +or `mha` — it recognizes the computation pattern. + +### Stage 3: Import to IR + +The FX graph is converted into Cactus's intermediate representation (`IRGraph`), a simple +dataflow graph of `IRNode`s connected by `IRValue`s. + +During import, constants that correspond to converted weight files are resolved through +`weights_manifest.json`. This metadata follows the value through the pipeline so that +lowering can use `mmap_weights()` instead of embedding the tensor into the graph. + +### Stage 4: Canonicalize the IR + +Multi-pass cleanup: + +- Rename ops to canonical spellings +- Rewrite `transpose`/`movedim` into `permute` +- Constant-fold trivial expressions +- Materialize `ones`, `zeros`, `arange` constants +- Remove no-op casts, views, slices +- Insert FP16 precision casts where the runtime requires them +- Dead code elimination + +### Stage 5: Fuse Patterns + +The optimizer recognizes computational subgraphs and collapses them into single +high-level ops: + +| Pattern | Fused Op | +|---------|----------| +| `variance -> rsqrt -> mul -> mul` | `rms_norm` | +| `cos/sin interleave on position freqs` | `rope` | +| `Q @ K^T -> scale -> mask -> softmax -> @ V` | `attention` | +| `gate_proj * up_proj -> activation -> down_proj` | `dense_mlp_tq_fused` | +| `conv -> batchnorm -> activation` | `conv_module` | +| `input/forget/cell/output gates` | `lstm_cell` | + +These fused ops map directly to optimized Cactus kernels. + +### Stage 6: Lower to CactusGraph + +Each `IRNode` becomes one or more `CactusGraph` operations. Constants become: + +- `graph.mmap_weights(path)` — zero-copy memory-mapped from disk (most weights) +- Inline scalar values +- Materialized tensor nodes (small constants baked into the graph) + +The output is a `TranspiledGraph` containing the runtime graph, input/output handles, +and weight binding metadata. + +--- + +## Converting HuggingFace Models + +The most common use case is converting a model from HuggingFace. `cactus convert` +quantizes the weights and builds the runtime graph in one step: + +```bash +# 1. Convert (CQ weights + runtime bundle) +cactus convert google/gemma-4-E2B-it ./gemma4-weights --bits 4 + +# 2. Run +cactus run ./gemma4-weights --prompt "Hello!" +``` + +### Text Models + +```bash +cactus convert Qwen/Qwen3-0.6B ./qwen3-weights --bits 4 \ + --prompt "The capital of France is" +``` + +### Multimodal Models (Gemma4) + +Multimodal models are automatically split into components (`vision_encoder`, +`audio_encoder`, `lm_encoder`, `decoder`) that are each independently captured +and optimized: + +```bash +cactus convert google/gemma-4-E2B-it ./gemma4-weights --bits 4 \ + --task multimodal_causal_lm_logits \ + --image-file photo.jpg \ + --audio-file speech.wav \ + --prompt "Describe what you see and hear." +``` + +### Speech Models (Parakeet TDT) + +```bash +cactus convert nvidia/parakeet-tdt-0.6b-v3 ./parakeet-weights --bits 4 \ + --audio-file recording.wav \ + --task tdt_transcription +``` + +--- + +## Running Saved Bundles + +Transpiled bundles are saved to `./weights/-cq/` by default (alongside +their CQ weights). Run them later without re-transpiling: + +```bash +# Text +cactus run ./weights/qwen3-0.6b-cq4 --prompt "Write a haiku" + +# Multimodal +cactus run ./weights/gemma-4-e2b-it-cq4 \ + --image photo.jpg \ + --audio speech.wav \ + --prompt "What do you see?" + +# Audio +cactus run ./weights/parakeet-tdt-0.6b-v3-cq4 \ + --audio meeting.wav +``` + +`cactus run` accepts either a HuggingFace model id (downloads the bundle) or a +local bundle directory path (detected by the presence of +`components/manifest.json`). The standalone `run-transpiled` subcommand was +removed — `cactus run` handles both cases. + +--- + +## Using the Python API Directly + +You can also drive the transpiler from Python: + +```python +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.lower import transpile_ir +from cactus.bindings.cactus import Graph +import torch +import numpy as np + +# Define your model +class MyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(64, 32) + + def forward(self, x): + return torch.relu(self.linear(x)) + +model = MyModel().eval() +example_input = torch.randn(1, 64) + +# Capture and transpile +captured = capture_model(model, (example_input,)) +transpiled = transpile_ir(captured.ir_graph) + +# Run +transpiled.set_input(0, np.random.randn(1, 64).astype(np.float16)) +outputs = transpiled.execute() +result = outputs[0].numpy() +print("Output shape:", result.shape) # (1, 32) +``` + +### Generic JAX User Graphs + +JAX and Flax models can be bundled by providing params and one or more graph +entrypoints. The caller owns tokenization, masks, sampling, and any prefill/decode +loop; Cactus captures the supplied functions, writes FP16 mmap weights, and saves +component graphs. + +```python +import jax.numpy as jnp +import numpy as np + +from cactus.transpile.capture_jax import JaxGraphSpec +from cactus.transpile.jax_user_graph_bundle import build_jax_user_graph_bundle +from cactus.transpile.jax_user_graph_bundle import load_jax_user_graph_bundle + + +def model(params, x): + return jnp.maximum(x @ params["w"] + params["b"], 0) + + +params = {"w": jnp.ones((4, 3), jnp.float16), "b": jnp.zeros((3,), jnp.float16)} +example_x = jnp.ones((2, 4), jnp.float16) + +result = build_jax_user_graph_bundle( + params=params, + output_dir="/tmp/my-jax-bundle", + model_id="my-jax-model", + specs=(JaxGraphSpec(name="forward", fn=model, example_args=(example_x,)),), +) + +loaded = load_jax_user_graph_bundle(result.output_dir) +y = loaded.execute("forward", np.ones((2, 4), np.float16))[0].numpy() +``` + +For encoder-decoder models, provide each graph boundary explicitly: + +```python +params, model, tokenizer, config = load_needle_model() +src = tokenizer.encode("What time is it in Tokyo?") +tgt = tokenizer.encode("") +src_mask = make_padding_mask(src, config.pad_token_id) +tgt_mask = make_causal_mask(tgt.shape[1]) +encoder_out, enc_mask = model.apply({"params": params}, src, src_mask, method=model.encode_text) + +result = build_jax_user_graph_bundle( + params=params, + output_dir="weights/needle", + model_id="needle", + task="encoder-decoder", + specs=( + JaxGraphSpec( + name="encoder", + fn=lambda params, src, src_mask: model.apply( + {"params": params}, src, src_mask, method=model.encode_text + ), + example_args=(src, src_mask), + output_names=("encoder_out", "encoder_mask"), + ), + JaxGraphSpec( + name="decoder_prefill", + fn=lambda params, tgt, encoder_out, tgt_mask, cross_mask: model.apply( + {"params": params}, tgt, encoder_out, tgt_mask, cross_mask, method=model.decode + ), + example_args=(tgt, encoder_out, tgt_mask, enc_mask), + output_names=("logits",), + ), + ), +) + +loaded = load_jax_user_graph_bundle(result.output_dir) +encoder_out, enc_mask = loaded.execute("encoder", src, src_mask) +logits = loaded.execute("decoder_prefill", tgt, encoder_out, tgt_mask, enc_mask)[0].numpy() +next_token = int(np.argmax(logits[0, -1])) +``` + +Each component is saved under `components//` with `graph.cactus`, +`raw_ir.json`, and `optimized_ir.json`. For decode, expose a separate +`decoder_step` graph. The generic JAX path does not create or manage internal +KV cache yet; component output `Tensor`s can be passed directly into another +component. Use `loaded.reset()` between independent generations. + +--- + +## Artifact Layout + +A transpiled bundle looks like this: + +``` +weights/-cq/ + raw_ir.json # IR before optimization (debugging) + optimized_ir.json # IR after fusion passes (debugging) + graph.cactus # serialized runtime graph + graph_bindings.json # weight binding metadata + result.json # execution results (if --execute-after-transpile) + components/ # for multimodal models + manifest.json # component order, input/output names + vision_encoder/ + graph.cactus + bound_constants/ + audio_encoder/ + graph.cactus + lm_encoder/ + graph.cactus + decoder/ + graph.cactus +``` + +The `manifest.json` tells the runtime loader how to chain the components and where +to find the memory-mapped weight files. The raw/optimized IR JSONs are also valid +reload targets — the runtime can re-lower from saved IR as a fallback. + +--- + +## CLI Reference + +### `cactus convert` + +```bash +cactus convert [output-dir] [options] +``` + +Quantizes the weights to CQ and builds the runtime graph. Pass `--weights-only` to +stop after the CQ weights. The graph-build options below are the same ones the +transpiler accepts. + +The decoder graph is emitted with a dynamic batch axis and a single baked KV-cache +slot, so single-stream decode stays lean while fixed-batch decode (`decode_batch` / +`generate_batch`) can size the KV-cache slot pool at runtime via +`Model::set_decode_slots(N)`. No separate conversion flag is needed. + +| Option | Description | +|--------|-------------| +| `--bits 1\|2\|3\|4` | CQ quantization bits (default: 4) | +| `--weights-only` | Stop after CQ quantization; skip the runtime graph | +| `--weights-dir ` | Path to converted CQ weights (default: `weights/`) | +| `--task ` | Force task type (default: `auto` — inferred from model config). Choices: `causal_lm_logits`, `multimodal_causal_lm_logits`, `ctc_logits`, `encoder_hidden_states`, `seq2seq_transcription`, `tdt_transcription` | +| `--prompt ` | Representative prompt for shape capture | +| `--image-file ` | Representative image (repeatable) | +| `--audio-file ` | Representative audio file | +| `--max-new-tokens ` | Generation room to preallocate for causal decode graphs | +| `--artifact-dir ` | Output directory (default: `weights/`) | +| `--execute-after-transpile` | Run the graph after saving | +| `--skip-reference-compare` | Skip PyTorch vs transpiled comparison | +| `--component-pipeline auto\|on\|off` | Use split component graph transpilation when supported | +| `--components ` | Comma-separated component subset for component-pipeline models | +| `--trust-remote-code` | Allow HF remote code during the build | +| `--local-files-only` | Require HF files to already be local | +| `--allow-unconverted-weights` | Build against an unconverted source checkpoint | +| `--no-fuse-rms-norm` | Disable RMSNorm fusion | +| `--no-fuse-rope` | Disable RoPE fusion | +| `--no-fuse-attention` | Disable attention fusion | + +### `cactus run` + +```bash +cactus run [options] +``` + +`model_id` may be a HuggingFace model id (downloads the matching bundle from +huggingface.co/Cactus-Compute) or a local path to a bundle directory. + +| Option | Description | +|--------|-------------| +| `--bits 1\|2\|3\|4\|2.54\|3.26` | CQ quantization bits when downloading (default: 4) | +| `--token ` | HuggingFace token (gated models) | +| `--prompt ` | Input prompt | +| `--input-ids ` | Comma-separated token IDs | +| `--image ` | Image file for multimodal bundles | +| `--audio ` | Audio file for speech / audio bundles | +| `--system ` | System prompt | +| `--thinking` | Enable thinking / reasoning mode | +| `--max-new-tokens ` | Max tokens to generate | +| `--result-json ` | Save bundle result payload to JSON | +| `--no-cloud-handoff` | Disable automatic cloud handoff | +| `--confidence-threshold ` | Confidence threshold below which to hand off to cloud | +| `--cloud-timeout-ms ` | Max wait time for cloud handoff | +| `--no-cloud-tele` | Disable cloud telemetry (write to cache only) | +| `--reconvert` | Force local rebuild from source | + +--- + +## Supported Model Families + +The transpiler works with any model that `torch.export` can capture, but has +tested adapters for: + +| Family | Tasks | Component Split | +|--------|-------|-----------------| +| Gemma 3/4 | text, multimodal | yes (vision + audio + LM + decoder) | +| Qwen 3/3.5 | text, vision | no | +| LFM2/2.5 | text, vision | no | +| Whisper | encoder | no | +| Parakeet CTC | encoder | no | +| Parakeet TDT | transcription | yes (encoder + decoder) | + +Adding a new model family means writing a small adapter in `model_adapters.py` that +wraps the HF forward into the canonical task interface. The rest of the pipeline +(capture, optimize, lower) is model-agnostic. + +--- + +## Design Decisions + +**ATen ops, not layer names.** After `torch.export`, the transpiler dispatches on +canonical ATen operation names (`aten.linear.default` -> `linear`). It never reads +HuggingFace module paths like `model.layers.0.self_attn`. This means the fusion +passes work on any model that produces the same computational pattern, regardless +of how the code is organized. + +**Weight binding through manifests.** The transpiler resolves weights exclusively +through `weights_manifest.json` written by `cactus convert`. It does not guess +filenames from layer names. This is intentional — it keeps the contract between +conversion and transpilation explicit and auditable. + +**Component splitting for multimodal.** Complex models like Gemma4 are split into +independent components (vision encoder, audio encoder, LM encoder, decoder) before +capture. Each component goes through the full optimize/lower pipeline independently. +This keeps each graph small enough for `torch.export` to handle and lets the runtime +load/unload components independently. + +**Saved IR as fallback.** Bundles include both the serialized `.cactus` graph and +the raw/optimized IR JSON. If the binary graph format changes between versions, the +runtime can re-lower from saved IR instead of requiring a full re-transpile. + +## See Also + +- [Cactus Engine API](/docs/cactus_engine.md) — C API for inference +- [Cactus Graph API](/docs/cactus_graph.md) — Computation graph reference +- [Fine-tuning Guide](/docs/finetuning.md) — Train LoRA fine-tunes and deploy to mobile +- [Transpiler source](https://github.com/cactus-compute/cactus/tree/main/python/cactus/transpile) — Full implementation diff --git a/docs/choose-bindings.md b/docs/choose-bindings.md new file mode 100644 index 000000000..df3a90eb7 --- /dev/null +++ b/docs/choose-bindings.md @@ -0,0 +1,48 @@ +# Choose Your Binding + +Not sure which binding to use? + +Pick the right one for your platform and use case: + +| | React Native | Flutter | Kotlin | Swift | Python | Rust | CLI | C++ | +|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| +| **Platforms** | iOS, Android | iOS, Android, macOS | Android, iOS (KMP) | iOS, macOS | Arm Linux, macOS | Arm Linux, macOS | macOS, Arm Linux | Arm Linux, macOS, iOS, Android | +| **Install** | build / source | build | build | build | pip | build / source | brew / source | header | +| **LLM** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **Streaming** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **Vision** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **Audio** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **Transcription** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **Function Calling** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **RAG / Embeddings** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| **Cloud Fallback** | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | + +## Quick Recommendations + +!!! tip "Building a mobile app?" + - **React Native** -- already have a React Native project? Just `npm install` and go + - **Flutter** -- cross-platform mobile + Mac with full native bindings + - **Kotlin** -- native Android apps or Kotlin Multiplatform + - **Swift** -- native iOS/macOS apps with Metal acceleration + +!!! tip "Server-side or scripting?" + - **Python** -- server-side inference, batch processing, or rapid prototyping + - **CLI** -- quick model testing and interactive sessions without writing code + +!!! tip "Embedding in a native app?" + - **C++** -- game engines, native desktop apps, or any C/C++ project + - **Rust** -- systems-level integration with safe FFI bindings + +## Binding Documentation + +- **[React Native](/bindings/react-native/)** -- Native bridge modules over the C API for iOS and Android +- **[Python](/python/)** -- Module-level FFI bindings, mirrors the C API +- **[Swift](/bindings/swift/)** -- XCFramework for iOS/macOS with Metal support +- **[Kotlin](/bindings/kotlin/)** -- JNI bindings + Kotlin Multiplatform support +- **[Flutter](/bindings/flutter/)** -- Dart FFI bindings for Android, iOS, and macOS +- **[Rust](/bindings/rust/)** -- Raw `extern "C"` FFI declarations +- **[C++ / Engine API](/docs/cactus_engine.md)** -- Direct C FFI for maximum control + +## Getting Started + +Once you've picked your binding, head to the **[Quickstart](quickstart.md)** to install and run your first completion. diff --git a/docs/compatibility.md b/docs/compatibility.md new file mode 100644 index 000000000..9c800202a --- /dev/null +++ b/docs/compatibility.md @@ -0,0 +1,38 @@ +--- +title: "Runtime & Weights Compatibility" +description: "How Cactus runtime versions map to model weight versions on HuggingFace. Explains versioning, compatibility checks, and when to re-download weights." +keywords: ["versioning", "compatibility", "model weights", "HuggingFace", "Cactus runtime"] +--- + +# Runtime & Weights Compatibility + +Some Cactus releases change the internal weight format. When this happens, cached weights from an older version will not load with a newer runtime and must be re-downloaded. + +Breaking weight changes are called out in the [release notes](https://github.com/cactus-compute/cactus/releases). + +## How Versioning Works + +Weights are published to [Hugging Face](https://huggingface.co/Cactus-Compute) and **only re-tagged when they actually change**. If a release does not affect the weight format, the previous tag remains — no new upload. + +``` +Runtime v1.7 -> weights tagged v1.7 on HF +Runtime v1.8 -> no new tag (unchanged) - still use v1.7 +... +Runtime v1.14 -> no new tag - still use v1.7 +Runtime v1.15 -> new tag v1.15 (changed!) - must update +``` + +**The rule:** use the latest HF weight tag that is ≤ your runtime version. + +## Checking Compatibility + +1. Open your model on [huggingface.co/Cactus-Compute](https://huggingface.co/Cactus-Compute) +2. Click **Files and versions → open branch dropdown from Main** +3. Find the latest tag that is ≤ your runtime version +4. If your local weights use an older tag, re-download them + +## See Also + +- [Cactus Engine API](/docs/cactus_engine.md) — Full inference API reference +- [Fine-tuning Guide](/docs/finetuning.md) — Convert and deploy custom fine-tunes +- [HuggingFace Weights](https://huggingface.co/Cactus-Compute) — Official Cactus model weights diff --git a/docs/finetuning.md b/docs/finetuning.md index ccfa42008..539aea5a1 100644 --- a/docs/finetuning.md +++ b/docs/finetuning.md @@ -1,9 +1,15 @@ +--- +title: "Deploying Unsloth Fine-Tunes to Cactus for Phones" +description: "Step-by-step guide to training LoRA fine-tunes with Unsloth and deploying them to iOS and Android devices using Cactus inference engine." +keywords: ["fine-tuning", "LoRA", "Unsloth", "mobile deployment", "on-device AI", "iOS", "Android", "Qwen", "Gemma"] +--- + # Deploying Unsloth Fine-Tunes to Cactus for Phones -- Cactus is an inference engine for mobile devices, macs and ARM chips like Raspberry Pi. +- Cactus is an inference engine for mobile devices, Macs, and ARM chips like Raspberry Pi. - At INT8, Cactus runs `Qwen3-0.6B` and `LFM2-1.2B` at `60-70 toks/sec` on iPhone 17 Pro, `13-18 toks/sec` on budget Pixel 6a. - INT4 quantization provides ~50% memory reduction with minimal quality loss. -- Task-Specific INT8 tunes of `Gemma3-270m` hit `150 toks/sec` on iPhone 17 Pro and `23 toks/sec` on Raspberry Pi. +- Task-Specific INT8 tunes of `Gemma3-270m` hit `150 toks/sec` on iPhone 17 Pro and `23 toks/sec` on Raspberry Pi 5. ## Quick Start @@ -54,18 +60,23 @@ git clone https://github.com/cactus-compute/cactus && cd cactus && source ./setu cactus convert Qwen/Qwen3-0.6B ./my-qwen3-0.6b --lora ./my-lora-adapter # From HuggingFace Hub: Use the correct base model! -cactus convert Qwen/Qwen3-0.6B ./my-qwen3-0.6b --lora username/my-lora-adapter +cactus convert Qwen/Qwen3-0.6B ./my-qwen3-0.6b --lora username/my-lora-adapter ``` Logo ### 4. Run -Test your model on Mac: +Build the native library, then run (the convert above already produced the runtime +bundle alongside the CQ weights): ```bash +cactus build cactus run ./my-qwen3-0.6b ``` + +`cactus run` accepts the bundle path directly, or auto-builds from a HF model id +when no bundle exists locally. Logo ### 5. Use in iOS/macOS App @@ -75,50 +86,55 @@ Build the native library: ```bash cactus build --apple ``` -```bash +``` Build complete! Total time: 58 seconds Static libraries: - Device: /Users/henry/Desktop/cactus/apple/libcactus-device.a - Simulator: /Users/henry/Desktop/cactus/apple/libcactus-simulator.a + Device: /apple/libcactus_engine-device.a + Simulator: /apple/libcactus_engine-simulator.a XCFrameworks: - iOS: /Users/henry/Desktop/cactus/apple/cactus-ios.xcframework - macOS: /Users/henry/Desktop/cactus/apple/cactus-macos.xcframework -Apple build complete! -(venv) henry@Henrys-MacBook-Air cactus % + iOS: /apple/cactus-ios.xcframework + macOS: /apple/cactus-macos.xcframework +Building Cactus for Apple platforms complete! ``` Link `cactus-ios.xcframework` to your Xcode project, then: ```swift import Foundation +import cactus // module map → cactus_engine.h -// Load model from app bundle let modelPath = Bundle.main.path(forResource: "my-model", ofType: nil)! -let model = cactus_init(modelPath, nil) +let model = cactus_init(modelPath, nil, false) -// Run completion let messages = "[{\"role\":\"user\",\"content\":\"Hello!\"}]" -var response = [CChar](repeating: 0, count: 4096) -cactus_complete(model, messages, &response, response.count, nil, nil, nil, nil) -print(String(cString: response)) +var buf = [Int8](repeating: 0, count: 65536) +buf.withUnsafeMutableBufferPointer { ptr in + _ = cactus_complete(model, messages, ptr.baseAddress, ptr.count, + nil, nil, nil, nil, nil, 0) +} +let resultJson = String(cString: buf) +if let data = resultJson.data(using: .utf8), + let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let response = obj["response"] as? String { + print(response) +} -// Cleanup cactus_destroy(model) ``` Logo You can now build iOS apps using the following code, but to see performance on any device while testing, -run cactus tests by plugging any iphone to your Mac then running: +run cactus tests by plugging any iPhone to your Mac then running: ```bash -cactus test -- --ios +cactus test --model --transcription-model --ios ``` Cactus demo apps will eventually expand to using your custom fine-tunes. -Also, `cactus run` will allow plugging in a phone, -such that the interactive session use the phone chips, +Also, `cactus run` will allow plugging in a phone, +such that the interactive session uses the phone chips, this way you can test before fully building out your apps. ### 6. Use in Android App @@ -128,47 +144,57 @@ Build the native library: ```bash cactus build --android ``` -```bash +``` Build complete! -Shared library location: /Users/henry/Desktop/cactus/android/libcactus.so -Static library location: /Users/henry/Desktop/cactus/android/libcactus.a -Android build complete! -(venv) henry@Henrys-MacBook-Air cactus % +Shared library location: /android/libcactus_engine.so +Static library location: /android/libcactus_engine.a +Building Cactus for Android complete! ``` -Copy `libcactus.so` to `app/src/main/jniLibs/arm64-v8a/`, then: +Copy `libcactus_engine.so` to `app/src/main/jniLibs/arm64-v8a/`, then: ```kotlin -class CactusWrapper { - init { System.loadLibrary("cactus") } +import com.cactus.* +import org.json.JSONObject - external fun init(modelPath: String, contextSize: Long, corpusDir: String?): Long - external fun complete(model: Long, messagesJson: String, bufferSize: Int): String - external fun destroy(model: Long) -} +val model = cactusInit("/data/local/tmp/my-model", null, false) -// Usage -val cactus = CactusWrapper() -val model = cactus.init("/data/local/tmp/my-model", 2048, null) -val response = cactus.complete(model, """[{"role":"user","content":"Hello!"}]""", 4096) -cactus.destroy(model) +val resultJson = cactusComplete( + model, + """[{"role":"user","content":"Hello!"}]""", + null, // optionsJson + null, // toolsJson + null, // callback + null, // pcmData +) +val response = JSONObject(resultJson).getString("response") +println(response) +cactusDestroy(model) ``` -You can now build ANdroid apps using the following code, +You can now build Android apps using the following code, but to see performance on any device while testing, -run cactus tests by plugging any android phone to your Mac then running: +run cactus tests by plugging any Android phone to your Mac then running: ```bash -cactus test -- --android +cactus test --model --transcription-model --android ``` Cactus demo apps will eventually expand to using your custom fine-tunes. -Also, `cactus run` will allow plugging in a phone, -such that the interactive session use the phone chips, +Also, `cactus run` will allow plugging in a phone, +such that the interactive session uses the phone chips, this way you can test before fully building out your apps. ## Resources -- Supported Base Models: `Qwen3, Gemma3, LFM2, SmolLM2` +- Supported Base Models: `Qwen3, Qwen3.5, Gemma3, LFM2, LFM2.5` - Full API reference: [Cactus Engine](https://github.com/cactus-compute/cactus/blob/main/docs/cactus_engine.md) - Learn more and report bugs: [Cactus](https://github.com/cactus-compute/cactus/tree/main) + +## See Also + +- [Cactus Engine API](/docs/cactus_engine.md) — Full C API reference for inference, streaming, and tool calling +- [Runtime Compatibility](/docs/compatibility.md) — Ensure your weights match your Cactus runtime version +- [Python Binding](/python/) — Use fine-tuned models from Python +- [Swift Binding](/bindings/swift/) — Deploy fine-tuned models in iOS/macOS apps +- [Kotlin Binding](/bindings/kotlin/) — Deploy fine-tuned models in Android apps diff --git a/docs/quickstart.md b/docs/quickstart.md new file mode 100644 index 000000000..19e3ac0cb --- /dev/null +++ b/docs/quickstart.md @@ -0,0 +1,173 @@ +# Quickstart + +Install Cactus and run your first on-device AI completion. + +## Installation + +=== "React Native" + + --8<-- "react-native/README.md:install" + + ### Platform Integration + + --8<-- "react-native/README.md:integration" + +=== "Flutter" + + --8<-- "flutter/README.md:install" + + ### Platform Integration + + --8<-- "flutter/README.md:integration" + +=== "Kotlin" + + --8<-- "kotlin/README.md:install" + + ### Platform Integration + + --8<-- "kotlin/README.md:integration" + +=== "Swift" + + --8<-- "swift/README.md:install" + + ### Platform Integration + + --8<-- "swift/README.md:integration" + +=== "Python" + + --8<-- "python/README.md:install" + +=== "Rust" + + --8<-- "rust/README.md:install" + +=== "CLI" + + **Homebrew (macOS):** + + ```bash + brew install cactus-compute/cactus/cactus + ``` + + **From Source (macOS):** + + ```bash + brew install cmake + git clone https://github.com/cactus-compute/cactus && cd cactus && source ./setup && cactus build --python + ``` + + **From Source (Linux):** + + ```bash + sudo apt-get install python3.12 python3.12-venv python3-pip cmake build-essential libcurl4-openssl-dev + git clone https://github.com/cactus-compute/cactus && cd cactus && source ./setup && cactus build --python + ``` + +=== "C++" + + Include the Cactus header in your project: + + ```cpp + #include + ``` + + See the [Cactus repository](https://github.com/cactus-compute/cactus) for CMake build instructions. + +--- + +## Your First Completion + +=== "React Native" + + --8<-- "react-native/README.md:example" + +=== "Flutter" + + --8<-- "flutter/README.md:example" + +=== "Kotlin" + + --8<-- "kotlin/README.md:example" + +=== "Swift" + + --8<-- "swift/README.md:example" + +=== "Python" + + --8<-- "python/README.md:example" + +=== "Rust" + + ```rust + use std::ffi::CString; + use std::os::raw::c_char; + + mod cactus; + + fn main() { + unsafe { + let model_path = CString::new("path/to/weight/folder").unwrap(); + let model = cactus::cactus_init(model_path.as_ptr(), std::ptr::null(), false); + + let messages = CString::new( + r#"[{"role": "user", "content": "What is the capital of France?"}]"# + ).unwrap(); + + let mut response = vec![0u8; 4096]; + cactus::cactus_complete( + model, messages.as_ptr(), + response.as_mut_ptr() as *mut c_char, response.len(), + std::ptr::null(), std::ptr::null(), + None, std::ptr::null_mut(), + std::ptr::null(), 0, + ); + + println!("{}", String::from_utf8_lossy(&response)); + cactus::cactus_destroy(model); + } + } + ``` + +=== "CLI" + + ```bash + cactus run + ``` + +=== "C++" + + ```cpp + #include + + cactus_model_t model = cactus_init( + "path/to/weight/folder", + "path/to/rag/documents", + false + ); + + const char* messages = R"([ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"} + ])"; + + char response[4096]; + int result = cactus_complete( + model, messages, response, sizeof(response), + nullptr, nullptr, nullptr, nullptr, + nullptr, 0 + ); + ``` + +--- + +## Next Steps + +- **[Engine API](cactus_engine.md)** -- Full inference API reference +- **[Speech Transcription](cactus_engine.md#cactus_transcribe)** -- Live & file transcription (`cactus transcribe`; one-shot + streaming) +- **[Graph API](cactus_graph.md)** -- Zero-copy computation graph for custom models +- **[Fine-tuning & Deployment](finetuning.md)** -- Convert and deploy custom fine-tunes +- **[Choose Your Binding](choose-bindings.md)** -- Help picking the right binding for your project diff --git a/flutter/README.md b/flutter/README.md deleted file mode 100644 index bdb5d0e59..000000000 --- a/flutter/README.md +++ /dev/null @@ -1,314 +0,0 @@ -# Cactus for Flutter - -Run AI models on-device with dart:ffi direct bindings for iOS, macOS, and Android. - -## Building - -```bash -cactus build --flutter -``` - -Build output: - -| File | Platform | -|------|----------| -| `libcactus.so` | Android (arm64-v8a) | -| `cactus-ios.xcframework` | iOS | -| `cactus-macos.xcframework` | macOS | - -see the main [README.md](../README.md) for how to use CLI & download weight - -## Integration - -### Android - -1. Copy `libcactus.so` to `android/app/src/main/jniLibs/arm64-v8a/` -2. Copy `cactus.dart` to your `lib/` folder - -### iOS - -1. Copy `cactus-ios.xcframework` to your `ios/` folder -2. Open `ios/Runner.xcworkspace` in Xcode -3. Drag the xcframework into the project -4. In Runner target > General > "Frameworks, Libraries, and Embedded Content", set to "Embed & Sign" -5. Copy `cactus.dart` to your `lib/` folder - -### macOS - -1. Copy `cactus-macos.xcframework` to your `macos/` folder -2. Open `macos/Runner.xcworkspace` in Xcode -3. Drag the xcframework into the project -4. In Runner target > General > "Frameworks, Libraries, and Embedded Content", set to "Embed & Sign" -5. Copy `cactus.dart` to your `lib/` folder - -## Usage - -### Basic Completion - -```dart -import 'cactus.dart'; - -final model = Cactus.create('/path/to/model.gguf'); -final result = model.complete('What is the capital of France?'); -print(result.text); -model.dispose(); -``` - -### Chat Messages - -```dart -final model = Cactus.create(modelPath); -final result = model.completeMessages([ - Message.system('You are a helpful assistant.'), - Message.user('What is 2 + 2?'), -]); -print(result.text); -model.dispose(); -``` - -### Completion Options - -```dart -final options = CompletionOptions( - temperature: 0.7, - topP: 0.9, - topK: 40, - maxTokens: 256, - stopSequences: ['\n\n'], -); - -final result = model.complete('Write a haiku:', options: options); -``` - -### Audio Transcription - -```dart -// From file -final result = model.transcribe('/path/to/audio.wav'); - -// From PCM data (16kHz mono) -final pcmData = Uint8List.fromList([...]); // Your PCM bytes -final result = model.transcribePcm(pcmData); -``` - -### Streaming Transcription - -```dart -final stream = model.createStreamTranscriber(); -stream.insert(audioChunk1); -stream.insert(audioChunk2); - -final partial = stream.process(); -print('Partial: ${partial.text}'); - -final finalResult = stream.finalize(); -print('Final: ${finalResult.text}'); - -stream.dispose(); -``` - -### Embeddings - -```dart -final embedding = model.embed('Hello, world!'); -final imageEmbedding = model.imageEmbed('/path/to/image.jpg'); -final audioEmbedding = model.audioEmbed('/path/to/audio.wav'); -``` - -### Tokenization - -```dart -final tokens = model.tokenize('Hello, world!'); -final scores = model.scoreWindow(tokens, 0, tokens.length, 512); -``` - -### RAG (Retrieval-Augmented Generation) - -```dart -final model = Cactus.create( - '/path/to/model.gguf', - corpusDir: '/path/to/documents', -); - -final result = model.complete('What does the documentation say about X?'); -``` - -### Vector Index - -```dart -final index = CactusIndex.create('/path/to/index', embeddingDim: 384); - -index.add( - ids: [1, 2], - documents: ['Document 1', 'Document 2'], - embeddings: [ - model.embed('Document 1'), - model.embed('Document 2'), - ], -); - -final results = index.query(model.embed('search query'), topK: 5); -for (final r in results) { - print('ID: ${r.id}, Score: ${r.score}'); -} - -index.dispose(); -``` - -## API Reference - -### Cactus - -```dart -class Cactus { - static Cactus create(String modelPath, {String? corpusDir}); - - CompletionResult complete(String prompt, {CompletionOptions options, void Function(String, int)? onToken}); - CompletionResult completeMessages(List messages, {CompletionOptions options, List>? tools, void Function(String, int)? onToken}); - - TranscriptionResult transcribe(String audioPath, {String? prompt, TranscriptionOptions options}); - TranscriptionResult transcribePcm(Uint8List pcmData, {String? prompt, TranscriptionOptions options}); - - List embed(String text, {bool normalize = true}); - List imageEmbed(String imagePath); - List audioEmbed(String audioPath); - String ragQuery(String query, {int topK = 5}); - - List tokenize(String text); - String scoreWindow(List tokens, int start, int end, int context); - StreamTranscriber createStreamTranscriber(); - - void reset(); - void stop(); - void dispose(); - - static String getLastError(); - static void setTelemetryToken(String token); - static void setProKey(String key); -} -``` - -### Message - -```dart -class Message { - static Message system(String content); - static Message user(String content); - static Message assistant(String content); -} -``` - -### CompletionOptions - -```dart -class CompletionOptions { - final double temperature; - final double topP; - final int topK; - final int maxTokens; - final List stopSequences; - final double confidenceThreshold; - - static const defaultOptions; -} -``` - -### CompletionResult - -```dart -class CompletionResult { - final String text; - final List>? functionCalls; - final int promptTokens; - final int completionTokens; - final double timeToFirstToken; - final double totalTime; - final double prefillTokensPerSecond; - final double decodeTokensPerSecond; - final double confidence; - final bool needsCloudHandoff; -} -``` - -### TranscriptionResult - -```dart -class TranscriptionResult { - final String text; - final List>? segments; - final double totalTime; -} -``` - -### StreamTranscriber - -```dart -class StreamTranscriber { - void insert(Uint8List pcmData); - TranscriptionResult process({String? language}); - TranscriptionResult finalize(); - void dispose(); -} -``` - -### CactusIndex - -```dart -class CactusIndex { - static CactusIndex create(String indexDir, {required int embeddingDim}); - - void add({required List ids, required List documents, required List> embeddings, List? metadatas}); - void delete(List ids); - List query(List embedding, {int topK = 5}); - void compact(); - void dispose(); -} - -class IndexResult { - final int id; - final double score; -} -``` - -## Bundling Model Weights - -Models must be accessible via file path at runtime. - -### Android - -Copy from assets to internal storage on first launch: - -```dart -import 'package:flutter/services.dart'; -import 'package:path_provider/path_provider.dart'; -import 'dart:io'; - -Future getModelPath() async { - final dir = await getApplicationDocumentsDirectory(); - final modelFile = File('${dir.path}/model.gguf'); - - if (!await modelFile.exists()) { - final data = await rootBundle.load('assets/model.gguf'); - await modelFile.writeAsBytes(data.buffer.asUint8List()); - } - - return modelFile.path; -} -``` - -### iOS/macOS - -Add model to bundle and access via path: - -```dart -import 'package:path_provider/path_provider.dart'; - -final path = '${Directory.current.path}/model.gguf'; -``` - -## Requirements - -- Flutter 3.0+ -- Dart 2.17+ -- iOS 14.0+ / macOS 13.0+ -- Android API 24+ / arm64-v8a diff --git a/flutter/build.sh b/flutter/build.sh deleted file mode 100755 index bbcb3fab1..000000000 --- a/flutter/build.sh +++ /dev/null @@ -1,44 +0,0 @@ -#!/bin/bash -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" -FLUTTER_DIR="$SCRIPT_DIR" - -echo "Building Cactus for Flutter..." -echo "" - -# Build Android -echo "=== Building Android ===" -if [ -f "$PROJECT_ROOT/android/build.sh" ]; then - bash "$PROJECT_ROOT/android/build.sh" - cp "$PROJECT_ROOT/android/libcactus.so" "$FLUTTER_DIR/" - echo "Copied libcactus.so" -else - echo "Warning: android/build.sh not found, skipping Android build" -fi - -echo "" - -# Build Apple (iOS/macOS) -echo "=== Building Apple (iOS/macOS) ===" -if [ -f "$PROJECT_ROOT/apple/build.sh" ]; then - bash "$PROJECT_ROOT/apple/build.sh" - - rm -rf "$FLUTTER_DIR/cactus-ios.xcframework" - cp -R "$PROJECT_ROOT/apple/cactus-ios.xcframework" "$FLUTTER_DIR/" - echo "Copied cactus-ios.xcframework" - - rm -rf "$FLUTTER_DIR/cactus-macos.xcframework" - cp -R "$PROJECT_ROOT/apple/cactus-macos.xcframework" "$FLUTTER_DIR/" - echo "Copied cactus-macos.xcframework" -else - echo "Warning: apple/build.sh not found, skipping Apple build" -fi - -echo "" -echo "=== Build Complete ===" -echo "" -echo "Output:" -echo " flutter/libcactus.so" -echo " flutter/cactus-ios.xcframework" -echo " flutter/cactus-macos.xcframework" diff --git a/flutter/cactus.dart b/flutter/cactus.dart deleted file mode 100644 index aee2d5985..000000000 --- a/flutter/cactus.dart +++ /dev/null @@ -1,1209 +0,0 @@ -import 'dart:ffi'; -import 'dart:io'; -import 'dart:convert'; -import 'dart:typed_data'; - -typedef CactusModelT = Pointer; -typedef CactusIndexT = Pointer; -typedef CactusStreamTranscribeT = Pointer; - -typedef TokenCallbackNative = Void Function( - Pointer token, Uint32 tokenId, Pointer userData); -typedef TokenCallbackDart = void Function( - Pointer token, int tokenId, Pointer userData); - -typedef CactusInitNative = CactusModelT Function( - Pointer modelPath, Pointer corpusDir); -typedef CactusDestroyNative = Void Function(CactusModelT model); -typedef CactusResetNative = Void Function(CactusModelT model); -typedef CactusStopNative = Void Function(CactusModelT model); - -typedef CactusCompleteNative = Int32 Function( - CactusModelT model, - Pointer messagesJson, - Pointer responseBuffer, - IntPtr bufferSize, - Pointer optionsJson, - Pointer toolsJson, - Pointer> callback, - Pointer userData); - -typedef CactusTokenizeNative = Int32 Function( - CactusModelT model, - Pointer text, - Pointer tokenBuffer, - IntPtr tokenBufferLen, - Pointer outTokenLen); - -typedef CactusScoreWindowNative = Int32 Function( - CactusModelT model, - Pointer tokens, - IntPtr tokenLen, - IntPtr start, - IntPtr end, - IntPtr context, - Pointer responseBuffer, - IntPtr bufferSize); - -typedef CactusTranscribeNative = Int32 Function( - CactusModelT model, - Pointer audioFilePath, - Pointer prompt, - Pointer responseBuffer, - IntPtr bufferSize, - Pointer optionsJson, - Pointer> callback, - Pointer userData, - Pointer pcmBuffer, - IntPtr pcmBufferSize); - -typedef CactusStreamTranscribeInitNative = CactusStreamTranscribeT Function( - CactusModelT model); -typedef CactusStreamTranscribeInsertNative = Int32 Function( - CactusStreamTranscribeT stream, Pointer pcmBuffer, IntPtr pcmBufferSize); -typedef CactusStreamTranscribeProcessNative = Int32 Function( - CactusStreamTranscribeT stream, - Pointer responseBuffer, - IntPtr bufferSize, - Pointer optionsJson); -typedef CactusStreamTranscribeFinalizeNative = Int32 Function( - CactusStreamTranscribeT stream, Pointer responseBuffer, IntPtr bufferSize); -typedef CactusStreamTranscribeDestroyNative = Void Function( - CactusStreamTranscribeT stream); - -typedef CactusEmbedNative = Int32 Function( - CactusModelT model, - Pointer text, - Pointer embeddingsBuffer, - IntPtr bufferSize, - Pointer embeddingDim, - Bool normalize); - -typedef CactusImageEmbedNative = Int32 Function( - CactusModelT model, - Pointer imagePath, - Pointer embeddingsBuffer, - IntPtr bufferSize, - Pointer embeddingDim); - -typedef CactusAudioEmbedNative = Int32 Function( - CactusModelT model, - Pointer audioPath, - Pointer embeddingsBuffer, - IntPtr bufferSize, - Pointer embeddingDim); - -typedef CactusRagQueryNative = Int32 Function( - CactusModelT model, - Pointer query, - Pointer responseBuffer, - IntPtr bufferSize, - IntPtr topK); - -typedef CactusIndexInitNative = CactusIndexT Function( - Pointer indexDir, IntPtr embeddingDim); -typedef CactusIndexAddNative = Int32 Function( - CactusIndexT index, - Pointer ids, - Pointer> documents, - Pointer> metadatas, - Pointer> embeddings, - IntPtr count, - IntPtr embeddingDim); -typedef CactusIndexDeleteNative = Int32 Function( - CactusIndexT index, Pointer ids, IntPtr idsCount); -typedef CactusIndexQueryNative = Int32 Function( - CactusIndexT index, - Pointer> embeddings, - IntPtr embeddingsCount, - IntPtr embeddingDim, - Pointer optionsJson, - Pointer> idBuffers, - Pointer idBufferSizes, - Pointer> scoreBuffers, - Pointer scoreBufferSizes); -typedef CactusIndexCompactNative = Int32 Function(CactusIndexT index); -typedef CactusIndexDestroyNative = Void Function(CactusIndexT index); - -typedef CactusGetLastErrorNative = Pointer Function(); -typedef CactusSetTelemetryTokenNative = Void Function(Pointer token); -typedef CactusSetProKeyNative = Void Function(Pointer proKey); - - -typedef CactusInitDart = CactusModelT Function( - Pointer modelPath, Pointer corpusDir); -typedef CactusDestroyDart = void Function(CactusModelT model); -typedef CactusResetDart = void Function(CactusModelT model); -typedef CactusStopDart = void Function(CactusModelT model); - -typedef CactusCompleteDart = int Function( - CactusModelT model, - Pointer messagesJson, - Pointer responseBuffer, - int bufferSize, - Pointer optionsJson, - Pointer toolsJson, - Pointer> callback, - Pointer userData); - -typedef CactusTokenizeDart = int Function( - CactusModelT model, - Pointer text, - Pointer tokenBuffer, - int tokenBufferLen, - Pointer outTokenLen); - -typedef CactusScoreWindowDart = int Function( - CactusModelT model, - Pointer tokens, - int tokenLen, - int start, - int end, - int context, - Pointer responseBuffer, - int bufferSize); - -typedef CactusTranscribeDart = int Function( - CactusModelT model, - Pointer audioFilePath, - Pointer prompt, - Pointer responseBuffer, - int bufferSize, - Pointer optionsJson, - Pointer> callback, - Pointer userData, - Pointer pcmBuffer, - int pcmBufferSize); - -typedef CactusStreamTranscribeInitDart = CactusStreamTranscribeT Function( - CactusModelT model); -typedef CactusStreamTranscribeInsertDart = int Function( - CactusStreamTranscribeT stream, Pointer pcmBuffer, int pcmBufferSize); -typedef CactusStreamTranscribeProcessDart = int Function( - CactusStreamTranscribeT stream, - Pointer responseBuffer, - int bufferSize, - Pointer optionsJson); -typedef CactusStreamTranscribeFinalizeDart = int Function( - CactusStreamTranscribeT stream, Pointer responseBuffer, int bufferSize); -typedef CactusStreamTranscribeDestroyDart = void Function( - CactusStreamTranscribeT stream); - -typedef CactusEmbedDart = int Function( - CactusModelT model, - Pointer text, - Pointer embeddingsBuffer, - int bufferSize, - Pointer embeddingDim, - bool normalize); - -typedef CactusImageEmbedDart = int Function( - CactusModelT model, - Pointer imagePath, - Pointer embeddingsBuffer, - int bufferSize, - Pointer embeddingDim); - -typedef CactusAudioEmbedDart = int Function( - CactusModelT model, - Pointer audioPath, - Pointer embeddingsBuffer, - int bufferSize, - Pointer embeddingDim); - -typedef CactusRagQueryDart = int Function( - CactusModelT model, - Pointer query, - Pointer responseBuffer, - int bufferSize, - int topK); - -typedef CactusIndexInitDart = CactusIndexT Function( - Pointer indexDir, int embeddingDim); -typedef CactusIndexAddDart = int Function( - CactusIndexT index, - Pointer ids, - Pointer> documents, - Pointer> metadatas, - Pointer> embeddings, - int count, - int embeddingDim); -typedef CactusIndexDeleteDart = int Function( - CactusIndexT index, Pointer ids, int idsCount); -typedef CactusIndexQueryDart = int Function( - CactusIndexT index, - Pointer> embeddings, - int embeddingsCount, - int embeddingDim, - Pointer optionsJson, - Pointer> idBuffers, - Pointer idBufferSizes, - Pointer> scoreBuffers, - Pointer scoreBufferSizes); -typedef CactusIndexCompactDart = int Function(CactusIndexT index); -typedef CactusIndexDestroyDart = void Function(CactusIndexT index); - -typedef CactusGetLastErrorDart = Pointer Function(); -typedef CactusSetTelemetryTokenDart = void Function(Pointer token); -typedef CactusSetProKeyDart = void Function(Pointer proKey); - - -DynamicLibrary _loadLibrary() { - if (Platform.isAndroid) { - return DynamicLibrary.open('libcactus.so'); - } else if (Platform.isIOS) { - return DynamicLibrary.process(); - } else if (Platform.isMacOS) { - return DynamicLibrary.process(); - } else { - throw UnsupportedError('Platform not supported: ${Platform.operatingSystem}'); - } -} - -final _lib = _loadLibrary(); - - -final _cactusInit = - _lib.lookupFunction('cactus_init'); -final _cactusDestroy = - _lib.lookupFunction('cactus_destroy'); -final _cactusReset = - _lib.lookupFunction('cactus_reset'); -final _cactusStop = - _lib.lookupFunction('cactus_stop'); -final _cactusComplete = - _lib.lookupFunction('cactus_complete'); -final _cactusTokenize = - _lib.lookupFunction('cactus_tokenize'); -final _cactusScoreWindow = _lib - .lookupFunction('cactus_score_window'); -final _cactusTranscribe = - _lib.lookupFunction('cactus_transcribe'); -final _cactusStreamTranscribeInit = _lib.lookupFunction< - CactusStreamTranscribeInitNative, - CactusStreamTranscribeInitDart>('cactus_stream_transcribe_init'); -final _cactusStreamTranscribeInsert = _lib.lookupFunction< - CactusStreamTranscribeInsertNative, - CactusStreamTranscribeInsertDart>('cactus_stream_transcribe_insert'); -final _cactusStreamTranscribeProcess = _lib.lookupFunction< - CactusStreamTranscribeProcessNative, - CactusStreamTranscribeProcessDart>('cactus_stream_transcribe_process'); -final _cactusStreamTranscribeFinalize = _lib.lookupFunction< - CactusStreamTranscribeFinalizeNative, - CactusStreamTranscribeFinalizeDart>('cactus_stream_transcribe_finalize'); -final _cactusStreamTranscribeDestroy = _lib.lookupFunction< - CactusStreamTranscribeDestroyNative, - CactusStreamTranscribeDestroyDart>('cactus_stream_transcribe_destroy'); -final _cactusEmbed = - _lib.lookupFunction('cactus_embed'); -final _cactusImageEmbed = - _lib.lookupFunction('cactus_image_embed'); -final _cactusAudioEmbed = - _lib.lookupFunction('cactus_audio_embed'); -final _cactusRagQuery = - _lib.lookupFunction('cactus_rag_query'); -final _cactusIndexInit = - _lib.lookupFunction('cactus_index_init'); -final _cactusIndexAdd = - _lib.lookupFunction('cactus_index_add'); -final _cactusIndexDelete = - _lib.lookupFunction('cactus_index_delete'); -final _cactusIndexQuery = - _lib.lookupFunction('cactus_index_query'); -final _cactusIndexCompact = - _lib.lookupFunction('cactus_index_compact'); -final _cactusIndexDestroy = - _lib.lookupFunction('cactus_index_destroy'); -final _cactusGetLastError = _lib - .lookupFunction('cactus_get_last_error'); -final _cactusSetTelemetryToken = _lib.lookupFunction('cactus_set_telemetry_token'); -final _cactusSetProKey = - _lib.lookupFunction('cactus_set_pro_key'); - -// ---------------------------------------------------------------------------- -// Helper Extensions -// ---------------------------------------------------------------------------- - -extension Utf8Pointer on String { - Pointer toNativeUtf8({Allocator allocator = malloc}) { - final units = utf8.encode(this); - final ptr = allocator(units.length + 1); - final list = ptr.asTypedList(units.length + 1); - list.setAll(0, units); - list[units.length] = 0; - return ptr.cast(); - } -} - -extension Utf8PointerExtension on Pointer { - String toDartString() { - if (this == nullptr) return ''; - final codeUnits = []; - var i = 0; - while (true) { - final byte = cast().elementAt(i).value; - if (byte == 0) break; - codeUnits.add(byte); - i++; - } - return utf8.decode(codeUnits); - } -} - -final malloc = _MallocAllocator(); - -class _MallocAllocator implements Allocator { - @override - Pointer allocate(int byteCount, {int? alignment}) { - return calloc.allocate(byteCount, alignment: alignment); - } - - @override - void free(Pointer pointer) { - calloc.free(pointer); - } -} - -final calloc = _Calloc(); - -class _Calloc implements Allocator { - static final _callocPtr = _lib.lookup Function(IntPtr, IntPtr)>>('calloc'); - static final _freePtr = _lib.lookup)>>('free'); - static final _calloc = _callocPtr.asFunction Function(int, int)>(); - static final _free = _freePtr.asFunction)>(); - - @override - Pointer allocate(int byteCount, {int? alignment}) { - return _calloc(byteCount, 1).cast(); - } - - @override - void free(Pointer pointer) { - _free(pointer.cast()); - } -} - -class Message { - final String role; - final String content; - - Message._(this.role, this.content); - - factory Message.system(String content) => Message._('system', content); - factory Message.user(String content) => Message._('user', content); - factory Message.assistant(String content) => Message._('assistant', content); - - Map toJson() => {'role': role, 'content': content}; -} - -class CompletionOptions { - final double temperature; - final double topP; - final int topK; - final int maxTokens; - final List stopSequences; - final double confidenceThreshold; - - const CompletionOptions({ - this.temperature = 0.7, - this.topP = 0.9, - this.topK = 40, - this.maxTokens = 512, - this.stopSequences = const [], - this.confidenceThreshold = 0.0, - }); - - static const defaultOptions = CompletionOptions(); - - Map toJson() => { - 'temperature': temperature, - 'top_p': topP, - 'top_k': topK, - 'max_tokens': maxTokens, - 'stop_sequences': stopSequences, - 'confidence_threshold': confidenceThreshold, - }; -} - -class CompletionResult { - final String text; - final List>? functionCalls; - final int promptTokens; - final int completionTokens; - final double timeToFirstToken; - final double totalTime; - final double prefillTokensPerSecond; - final double decodeTokensPerSecond; - final double confidence; - final bool needsCloudHandoff; - - CompletionResult({ - required this.text, - this.functionCalls, - this.promptTokens = 0, - this.completionTokens = 0, - this.timeToFirstToken = 0.0, - this.totalTime = 0.0, - this.prefillTokensPerSecond = 0.0, - this.decodeTokensPerSecond = 0.0, - this.confidence = 1.0, - this.needsCloudHandoff = false, - }); - - factory CompletionResult.fromJson(Map json) { - return CompletionResult( - text: json['text'] ?? '', - functionCalls: json['function_calls'] != null - ? List>.from(json['function_calls']) - : null, - promptTokens: json['prompt_tokens'] ?? 0, - completionTokens: json['completion_tokens'] ?? 0, - timeToFirstToken: (json['time_to_first_token'] ?? 0.0).toDouble(), - totalTime: (json['total_time'] ?? 0.0).toDouble(), - prefillTokensPerSecond: (json['prefill_tokens_per_second'] ?? 0.0).toDouble(), - decodeTokensPerSecond: (json['decode_tokens_per_second'] ?? 0.0).toDouble(), - confidence: (json['confidence'] ?? 1.0).toDouble(), - needsCloudHandoff: json['needs_cloud_handoff'] ?? false, - ); - } -} - -class TranscriptionOptions { - final String? language; - final bool translate; - - const TranscriptionOptions({ - this.language, - this.translate = false, - }); - - static const defaultOptions = TranscriptionOptions(); - - Map toJson() => { - if (language != null) 'language': language, - 'translate': translate, - }; -} - -class TranscriptionResult { - final String text; - final List>? segments; - final double totalTime; - - TranscriptionResult({ - required this.text, - this.segments, - this.totalTime = 0.0, - }); - - factory TranscriptionResult.fromJson(Map json) { - return TranscriptionResult( - text: json['text'] ?? '', - segments: json['segments'] != null - ? List>.from(json['segments']) - : null, - totalTime: (json['total_time'] ?? 0.0).toDouble(), - ); - } -} - -class IndexResult { - final int id; - final double score; - - IndexResult({required this.id, required this.score}); -} - -class Cactus { - final CactusModelT _handle; - bool _disposed = false; - - Cactus._(this._handle); - - static Cactus create(String modelPath, {String? corpusDir}) { - final modelPathPtr = modelPath.toNativeUtf8(); - final corpusDirPtr = corpusDir?.toNativeUtf8() ?? nullptr; - - try { - final handle = _cactusInit(modelPathPtr, corpusDirPtr); - if (handle == nullptr) { - throw CactusException('Failed to initialize model: ${getLastError()}'); - } - return Cactus._(handle); - } finally { - calloc.free(modelPathPtr); - if (corpusDirPtr != nullptr) calloc.free(corpusDirPtr); - } - } - - void _checkNotDisposed() { - if (_disposed) { - throw StateError('Cactus instance has been disposed'); - } - } - - CompletionResult complete( - String prompt, { - CompletionOptions options = CompletionOptions.defaultOptions, - void Function(String token, int tokenId)? onToken, - }) { - return completeMessages( - [Message.user(prompt)], - options: options, - onToken: onToken, - ); - } - - CompletionResult completeMessages( - List messages, { - CompletionOptions options = CompletionOptions.defaultOptions, - List>? tools, - void Function(String token, int tokenId)? onToken, - }) { - _checkNotDisposed(); - - final messagesJson = jsonEncode(messages.map((m) => m.toJson()).toList()); - final optionsJson = jsonEncode(options.toJson()); - final toolsJson = tools != null ? jsonEncode(tools) : null; - - const bufferSize = 1024 * 1024; // 1MB buffer - final responseBuffer = calloc(bufferSize); - final messagesPtr = messagesJson.toNativeUtf8(); - final optionsPtr = optionsJson.toNativeUtf8(); - final toolsPtr = toolsJson?.toNativeUtf8() ?? nullptr; - - Pointer> callbackPtr = nullptr; - if (onToken != null) { - } - - try { - final result = _cactusComplete( - _handle, - messagesPtr, - responseBuffer.cast(), - bufferSize, - optionsPtr, - toolsPtr, - callbackPtr, - nullptr, - ); - - if (result != 0) { - throw CactusException('Completion failed: ${getLastError()}'); - } - - final responseStr = responseBuffer.cast().toDartString(); - final responseJson = jsonDecode(responseStr) as Map; - return CompletionResult.fromJson(responseJson); - } finally { - calloc.free(responseBuffer); - calloc.free(messagesPtr); - calloc.free(optionsPtr); - if (toolsPtr != nullptr) calloc.free(toolsPtr); - } - } - - List tokenize(String text) { - _checkNotDisposed(); - - const maxTokens = 8192; - final tokenBuffer = calloc(maxTokens); - final outTokenLen = calloc(1); - final textPtr = text.toNativeUtf8(); - - try { - final result = _cactusTokenize( - _handle, - textPtr, - tokenBuffer, - maxTokens, - outTokenLen, - ); - - if (result != 0) { - throw CactusException('Tokenization failed: ${getLastError()}'); - } - - final tokenCount = outTokenLen.value; - return List.generate(tokenCount, (i) => tokenBuffer[i]); - } finally { - calloc.free(tokenBuffer); - calloc.free(outTokenLen); - calloc.free(textPtr); - } - } - - String scoreWindow(List tokens, int start, int end, int context) { - _checkNotDisposed(); - - final tokenBuffer = calloc(tokens.length); - for (var i = 0; i < tokens.length; i++) { - tokenBuffer[i] = tokens[i]; - } - - const bufferSize = 65536; - final responseBuffer = calloc(bufferSize); - - try { - final result = _cactusScoreWindow( - _handle, - tokenBuffer, - tokens.length, - start, - end, - context, - responseBuffer.cast(), - bufferSize, - ); - - if (result != 0) { - throw CactusException('Score window failed: ${getLastError()}'); - } - - return responseBuffer.cast().toDartString(); - } finally { - calloc.free(tokenBuffer); - calloc.free(responseBuffer); - } - } - - TranscriptionResult transcribe( - String audioPath, { - String? prompt, - TranscriptionOptions options = TranscriptionOptions.defaultOptions, - }) { - _checkNotDisposed(); - - const bufferSize = 1024 * 1024; - final responseBuffer = calloc(bufferSize); - final audioPathPtr = audioPath.toNativeUtf8(); - final promptPtr = prompt?.toNativeUtf8() ?? nullptr; - final optionsJson = jsonEncode(options.toJson()); - final optionsPtr = optionsJson.toNativeUtf8(); - - try { - final result = _cactusTranscribe( - _handle, - audioPathPtr, - promptPtr, - responseBuffer.cast(), - bufferSize, - optionsPtr, - nullptr, - nullptr, - nullptr, - 0, - ); - - if (result != 0) { - throw CactusException('Transcription failed: ${getLastError()}'); - } - - final responseStr = responseBuffer.cast().toDartString(); - final responseJson = jsonDecode(responseStr) as Map; - return TranscriptionResult.fromJson(responseJson); - } finally { - calloc.free(responseBuffer); - calloc.free(audioPathPtr); - if (promptPtr != nullptr) calloc.free(promptPtr); - calloc.free(optionsPtr); - } - } - - TranscriptionResult transcribePcm( - Uint8List pcmData, { - String? prompt, - TranscriptionOptions options = TranscriptionOptions.defaultOptions, - }) { - _checkNotDisposed(); - - const bufferSize = 1024 * 1024; - final responseBuffer = calloc(bufferSize); - final promptPtr = prompt?.toNativeUtf8() ?? nullptr; - final optionsJson = jsonEncode(options.toJson()); - final optionsPtr = optionsJson.toNativeUtf8(); - final pcmBuffer = calloc(pcmData.length); - pcmBuffer.asTypedList(pcmData.length).setAll(0, pcmData); - - try { - final result = _cactusTranscribe( - _handle, - nullptr, - promptPtr, - responseBuffer.cast(), - bufferSize, - optionsPtr, - nullptr, - nullptr, - pcmBuffer, - pcmData.length, - ); - - if (result != 0) { - throw CactusException('Transcription failed: ${getLastError()}'); - } - - final responseStr = responseBuffer.cast().toDartString(); - final responseJson = jsonDecode(responseStr) as Map; - return TranscriptionResult.fromJson(responseJson); - } finally { - calloc.free(responseBuffer); - if (promptPtr != nullptr) calloc.free(promptPtr); - calloc.free(optionsPtr); - calloc.free(pcmBuffer); - } - } - - StreamTranscriber createStreamTranscriber() { - _checkNotDisposed(); - final handle = _cactusStreamTranscribeInit(_handle); - if (handle == nullptr) { - throw CactusException('Failed to create stream transcriber: ${getLastError()}'); - } - return StreamTranscriber._(handle); - } - - List embed(String text, {bool normalize = true}) { - _checkNotDisposed(); - - const maxDim = 8192; - final embeddingsBuffer = calloc(maxDim); - final embeddingDim = calloc(1); - final textPtr = text.toNativeUtf8(); - - try { - final result = _cactusEmbed( - _handle, - textPtr, - embeddingsBuffer, - maxDim, - embeddingDim, - normalize, - ); - - if (result != 0) { - throw CactusException('Embedding failed: ${getLastError()}'); - } - - final dim = embeddingDim.value; - return List.generate(dim, (i) => embeddingsBuffer[i]); - } finally { - calloc.free(embeddingsBuffer); - calloc.free(embeddingDim); - calloc.free(textPtr); - } - } - - List imageEmbed(String imagePath) { - _checkNotDisposed(); - - const maxDim = 8192; - final embeddingsBuffer = calloc(maxDim); - final embeddingDim = calloc(1); - final imagePathPtr = imagePath.toNativeUtf8(); - - try { - final result = _cactusImageEmbed( - _handle, - imagePathPtr, - embeddingsBuffer, - maxDim, - embeddingDim, - ); - - if (result != 0) { - throw CactusException('Image embedding failed: ${getLastError()}'); - } - - final dim = embeddingDim.value; - return List.generate(dim, (i) => embeddingsBuffer[i]); - } finally { - calloc.free(embeddingsBuffer); - calloc.free(embeddingDim); - calloc.free(imagePathPtr); - } - } - - List audioEmbed(String audioPath) { - _checkNotDisposed(); - - const maxDim = 8192; - final embeddingsBuffer = calloc(maxDim); - final embeddingDim = calloc(1); - final audioPathPtr = audioPath.toNativeUtf8(); - - try { - final result = _cactusAudioEmbed( - _handle, - audioPathPtr, - embeddingsBuffer, - maxDim, - embeddingDim, - ); - - if (result != 0) { - throw CactusException('Audio embedding failed: ${getLastError()}'); - } - - final dim = embeddingDim.value; - return List.generate(dim, (i) => embeddingsBuffer[i]); - } finally { - calloc.free(embeddingsBuffer); - calloc.free(embeddingDim); - calloc.free(audioPathPtr); - } - } - - String ragQuery(String query, {int topK = 5}) { - _checkNotDisposed(); - - const bufferSize = 1024 * 1024; - final responseBuffer = calloc(bufferSize); - final queryPtr = query.toNativeUtf8(); - - try { - final result = _cactusRagQuery( - _handle, - queryPtr, - responseBuffer.cast(), - bufferSize, - topK, - ); - - if (result != 0) { - throw CactusException('RAG query failed: ${getLastError()}'); - } - - return responseBuffer.cast().toDartString(); - } finally { - calloc.free(responseBuffer); - calloc.free(queryPtr); - } - } - - void reset() { - _checkNotDisposed(); - _cactusReset(_handle); - } - - void stop() { - _checkNotDisposed(); - _cactusStop(_handle); - } - - void dispose() { - if (!_disposed) { - _cactusDestroy(_handle); - _disposed = true; - } - } - - static String getLastError() { - return _cactusGetLastError().toDartString(); - } - - static void setTelemetryToken(String token) { - final tokenPtr = token.toNativeUtf8(); - try { - _cactusSetTelemetryToken(tokenPtr); - } finally { - calloc.free(tokenPtr); - } - } - - static void setProKey(String key) { - final keyPtr = key.toNativeUtf8(); - try { - _cactusSetProKey(keyPtr); - } finally { - calloc.free(keyPtr); - } - } -} - -class StreamTranscriber { - final CactusStreamTranscribeT _handle; - bool _disposed = false; - - StreamTranscriber._(this._handle); - - void _checkNotDisposed() { - if (_disposed) { - throw StateError('StreamTranscriber has been disposed'); - } - } - - void insert(Uint8List pcmData) { - _checkNotDisposed(); - - final pcmBuffer = calloc(pcmData.length); - pcmBuffer.asTypedList(pcmData.length).setAll(0, pcmData); - - try { - final result = _cactusStreamTranscribeInsert(_handle, pcmBuffer, pcmData.length); - if (result != 0) { - throw CactusException('Stream insert failed: ${Cactus.getLastError()}'); - } - } finally { - calloc.free(pcmBuffer); - } - } - - TranscriptionResult process({String? language}) { - _checkNotDisposed(); - - const bufferSize = 1024 * 1024; - final responseBuffer = calloc(bufferSize); - final optionsJson = language != null ? jsonEncode({'language': language}) : null; - final optionsPtr = optionsJson?.toNativeUtf8() ?? nullptr; - - try { - final result = _cactusStreamTranscribeProcess( - _handle, - responseBuffer.cast(), - bufferSize, - optionsPtr, - ); - - if (result != 0) { - throw CactusException('Stream process failed: ${Cactus.getLastError()}'); - } - - final responseStr = responseBuffer.cast().toDartString(); - final responseJson = jsonDecode(responseStr) as Map; - return TranscriptionResult.fromJson(responseJson); - } finally { - calloc.free(responseBuffer); - if (optionsPtr != nullptr) calloc.free(optionsPtr); - } - } - - TranscriptionResult finalize() { - _checkNotDisposed(); - - const bufferSize = 1024 * 1024; - final responseBuffer = calloc(bufferSize); - - try { - final result = _cactusStreamTranscribeFinalize( - _handle, - responseBuffer.cast(), - bufferSize, - ); - - if (result != 0) { - throw CactusException('Stream finalize failed: ${Cactus.getLastError()}'); - } - - final responseStr = responseBuffer.cast().toDartString(); - final responseJson = jsonDecode(responseStr) as Map; - return TranscriptionResult.fromJson(responseJson); - } finally { - calloc.free(responseBuffer); - } - } - - void dispose() { - if (!_disposed) { - _cactusStreamTranscribeDestroy(_handle); - _disposed = true; - } - } -} - -class CactusIndex { - final CactusIndexT _handle; - final int embeddingDim; - bool _disposed = false; - - CactusIndex._(this._handle, this.embeddingDim); - - static CactusIndex create(String indexDir, {required int embeddingDim}) { - final indexDirPtr = indexDir.toNativeUtf8(); - - try { - final handle = _cactusIndexInit(indexDirPtr, embeddingDim); - if (handle == nullptr) { - throw CactusException('Failed to create index: ${Cactus.getLastError()}'); - } - return CactusIndex._(handle, embeddingDim); - } finally { - calloc.free(indexDirPtr); - } - } - - void _checkNotDisposed() { - if (_disposed) { - throw StateError('CactusIndex has been disposed'); - } - } - - void add({ - required List ids, - required List documents, - required List> embeddings, - List? metadatas, - }) { - _checkNotDisposed(); - - if (ids.length != documents.length || ids.length != embeddings.length) { - throw ArgumentError('ids, documents, and embeddings must have the same length'); - } - - final count = ids.length; - final idsPtr = calloc(count); - for (var i = 0; i < count; i++) { - idsPtr[i] = ids[i]; - } - - final documentsPtr = calloc>(count); - for (var i = 0; i < count; i++) { - documentsPtr[i] = documents[i].toNativeUtf8(); - } - - final metadatasPtr = metadatas != null ? calloc>(count) : nullptr; - if (metadatas != null) { - for (var i = 0; i < count; i++) { - metadatasPtr[i] = metadatas[i].toNativeUtf8(); - } - } - - final embeddingsPtr = calloc>(count); - for (var i = 0; i < count; i++) { - final embedding = embeddings[i]; - final embPtr = calloc(embedding.length); - for (var j = 0; j < embedding.length; j++) { - embPtr[j] = embedding[j]; - } - embeddingsPtr[i] = embPtr; - } - - try { - final result = _cactusIndexAdd( - _handle, - idsPtr, - documentsPtr, - metadatasPtr, - embeddingsPtr, - count, - embeddingDim, - ); - - if (result != 0) { - throw CactusException('Index add failed: ${Cactus.getLastError()}'); - } - } finally { - calloc.free(idsPtr); - for (var i = 0; i < count; i++) { - calloc.free(documentsPtr[i]); - calloc.free(embeddingsPtr[i]); - if (metadatasPtr != nullptr) calloc.free(metadatasPtr[i]); - } - calloc.free(documentsPtr); - calloc.free(embeddingsPtr); - if (metadatasPtr != nullptr) calloc.free(metadatasPtr); - } - } - - void delete(List ids) { - _checkNotDisposed(); - - final idsPtr = calloc(ids.length); - for (var i = 0; i < ids.length; i++) { - idsPtr[i] = ids[i]; - } - - try { - final result = _cactusIndexDelete(_handle, idsPtr, ids.length); - if (result != 0) { - throw CactusException('Index delete failed: ${Cactus.getLastError()}'); - } - } finally { - calloc.free(idsPtr); - } - } - - List query(List embedding, {int topK = 5}) { - _checkNotDisposed(); - - final embPtr = calloc(embedding.length); - for (var i = 0; i < embedding.length; i++) { - embPtr[i] = embedding[i]; - } - - final embPtrPtr = calloc>(1); - embPtrPtr[0] = embPtr; - - final optionsJson = jsonEncode({'top_k': topK}); - final optionsPtr = optionsJson.toNativeUtf8(); - - final idBuffers = calloc>(1); - final idBufferSizes = calloc(1); - final scoreBuffers = calloc>(1); - final scoreBufferSizes = calloc(1); - - try { - final result = _cactusIndexQuery( - _handle, - embPtrPtr, - 1, - embedding.length, - optionsPtr, - idBuffers, - idBufferSizes, - scoreBuffers, - scoreBufferSizes, - ); - - if (result != 0) { - throw CactusException('Index query failed: ${Cactus.getLastError()}'); - } - - final resultCount = idBufferSizes[0]; - final results = []; - for (var i = 0; i < resultCount; i++) { - results.add(IndexResult( - id: idBuffers[0][i], - score: scoreBuffers[0][i], - )); - } - - if (idBuffers[0] != nullptr) calloc.free(idBuffers[0]); - if (scoreBuffers[0] != nullptr) calloc.free(scoreBuffers[0]); - - return results; - } finally { - calloc.free(embPtr); - calloc.free(embPtrPtr); - calloc.free(optionsPtr); - calloc.free(idBuffers); - calloc.free(idBufferSizes); - calloc.free(scoreBuffers); - calloc.free(scoreBufferSizes); - } - } - - void compact() { - _checkNotDisposed(); - final result = _cactusIndexCompact(_handle); - if (result != 0) { - throw CactusException('Index compact failed: ${Cactus.getLastError()}'); - } - } - - void dispose() { - if (!_disposed) { - _cactusIndexDestroy(_handle); - _disposed = true; - } - } -} - -class CactusException implements Exception { - final String message; - CactusException(this.message); - - @override - String toString() => 'CactusException: $message'; -} diff --git a/libs/libcactus_pro.a b/libs/libcactus_pro.a deleted file mode 100644 index 0f969ba2e..000000000 Binary files a/libs/libcactus_pro.a and /dev/null differ diff --git a/llms.txt b/llms.txt new file mode 100644 index 000000000..420d17e95 --- /dev/null +++ b/llms.txt @@ -0,0 +1,55 @@ +# Cactus + +> Energy-efficient AI inference engine for running LLMs, vision models, and speech models on phones, wearables, Macs, and ARM devices like Raspberry Pi. Built by Cactus Compute, Inc. (YC S25). + +Cactus has three layers: Cactus Kernels (ARM SIMD kernels for Apple, Snapdragon, Google Tensor, Exynos, MediaTek, Raspberry Pi), Cactus Graph (zero-copy computation graph, like PyTorch for mobile), and Cactus Engine (high-level inference API with OpenAI-compatible chat completion, tool calling, auto-RAG, vision, audio, transcription, embeddings, and cloud handoff). + +Supported models include Gemma 4 (E2B, E4B with vision + audio), Gemma 3/3n, Qwen 3/3.5, LiquidAI LFM2/LFM2.5 (incl. VL and MoE), Whisper, Moonshine, Parakeet (CTC/TDT), and Nomic Embed. Models run at INT4/INT8 precision with Metal GPU acceleration on supported hardware. Pre-converted weights at https://huggingface.co/Cactus-Compute. + +## Docs + +- [Engine API](docs/cactus_engine.md): C FFI for LLM inference, chat completion, streaming, tool calling, transcription, embeddings, RAG, vision, and audio +- [Graph API](docs/cactus_graph.md): Computational graph framework for tensor operations, matrix multiplication, attention, normalization, convolutions, MoE layers, and activation functions +- [Vector Index](docs/cactus_index.md): On-device vector database with cosine similarity search for RAG applications +- [Fine-tuning Guide](docs/finetuning.md): Deploying Unsloth LoRA fine-tunes to mobile devices via Cactus +- [Compatibility](docs/compatibility.md): Runtime and weight versioning across releases + +## Bindings + +- [Python](python/README.md): Python bindings via ctypes FFI +- [Swift](bindings/swift/README.md): Swift bindings via C module map (iOS, macOS) +- [Kotlin](bindings/kotlin/README.md): Kotlin bindings via JNI (Android) and cinterop (iOS via KMP) +- [Flutter](bindings/flutter/README.md): Dart FFI bindings (iOS, Android) +- [Rust](bindings/rust/README.md): Raw `extern "C"` declarations +- [React Native](bindings/react-native/README.md): React Native bridge modules (iOS, Android) + +## Blog + +- [TurboQuant-H](blog/turboquant-h.md): Hadamard rotation for 2-bit per-layer embedding quantization — 4x PLI compression, 40% total model size reduction +- [Gemma 4 on Cactus](blog/gemma4.md): Native multimodal voice, vision, and audio on-device with hybrid cloud handoff +- [Hybrid Transcription](blog/hybrid_transcription.md): Sub-150ms transcription with cloud-level accuracy using on-device/cloud hybrid inference +- [LFM2 24B Review](blog/lfm2_24b_a2b.md): Running LFM2-24B MoE (A2B) locally on Mac for coding use cases +- [Parakeet CTC 1.1B Review](blog/parakeet.md): Sub-200ms on-device transcription with NVIDIA Parakeet CTC 1.1B on Apple Silicon +- [LFM-2.5-350m](blog/lfm2.5_350m.md): 140 tok/sec single-core INT8 inference across seven devices + +## Key Concepts + +- **Cloud Handoff**: When the on-device model's confidence drops below a threshold, Cactus signals `cloud_handoff: true` so the app can route to a cloud API +- **Auto-RAG**: Pass a corpus directory at init and Cactus automatically retrieves relevant context for each query +- **Tool Calling**: Define tools as JSON and Cactus returns structured `function_calls` in the response +- **Tool RAG**: When many tools are defined, Cactus selects the top-k most relevant tools per query +- **Streaming**: All completion and transcription APIs support token-level streaming callbacks +- **Audio in Messages**: Gemma 4 models accept `"audio": ["path.wav"]` in message JSON for native audio understanding +- **Vision in Messages**: VLM models (Gemma 4, LFM2-VL, Qwen3.5) accept `"images": ["path.png"]` in message JSON +- **Thinking**: Models that support chain-of-thought (Gemma 4) return a `thinking` field when `enable_thinking_if_supported` is set +- **TurboQuant-H**: Hadamard rotation + Lloyd-Max codebook quantization for 2-bit per-layer embeddings, reducing embedding-dominated models by 40% +- **INT4 Quantization**: Lossy weight quantization for ~50% memory reduction with minimal quality loss +- **Metal Acceleration**: Apple GPU support for faster inference on Apple silicon + +## Links + +- GitHub: https://github.com/cactus-compute/cactus +- Docs: https://docs.cactuscompute.com/ +- Website: https://cactuscompute.com/ +- HuggingFace Weights: https://huggingface.co/Cactus-Compute +- Reddit: https://www.reddit.com/r/cactuscompute/ diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 000000000..b7bd930cd --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,113 @@ +site_name: Cactus Docs +site_url: https://docs.cactuscompute.com/ +site_description: "Low-latency AI engine for mobile devices and wearables" +site_author: Cactus Compute +repo_url: https://github.com/cactus-compute/cactus +repo_name: cactus-compute/cactus + +docs_dir: site_docs + +theme: + name: material + custom_dir: .github/docs-overrides + logo: assets/logo_white.png + favicon: assets/logo_white.png + font: false + features: + - navigation.tabs + - navigation.sections + - navigation.expand + - navigation.top + - navigation.footer + - search.suggest + - search.highlight + - content.code.copy + - content.tabs.link + - toc.follow + palette: + - scheme: slate + primary: black + accent: white + toggle: + icon: material/brightness-4 + name: Switch to light mode + - scheme: default + primary: black + accent: black + toggle: + icon: material/brightness-7 + name: Switch to dark mode + +markdown_extensions: + - pymdownx.highlight: + anchor_linenums: true + - pymdownx.superfences + - pymdownx.tabbed: + alternate_style: true + - admonition + - pymdownx.details + - attr_list + - pymdownx.emoji: + emoji_index: !!python/name:material.extensions.emoji.twemoji + emoji_generator: !!python/name:material.extensions.emoji.to_svg + - md_in_html + - toc: + permalink: true + - tables + - pymdownx.inlinehilite + - pymdownx.snippets: + base_path: [site_docs] + - pymdownx.arithmatex: + generic: true + +plugins: + - search + +extra: + version: + provider: mike + default: latest + social: + - icon: fontawesome/brands/github + link: https://github.com/cactus-compute/cactus + - icon: fontawesome/brands/reddit + link: https://www.reddit.com/r/cactuscompute/ + +extra_javascript: + - javascripts/mathjax.js + - https://unpkg.com/mathjax@3/es5/tex-mml-chtml.js + +extra_css: + - stylesheets/custom.css + +nav: + - Home: index.md + - Quickstart: docs/quickstart.md + - Choose Your Binding: docs/choose-bindings.md + - Bindings: + - React Native: react-native/README.md + - Python: python/README.md + - Swift: swift/README.md + - Kotlin: kotlin/README.md + - Flutter: flutter/README.md + - Rust: rust/README.md + - Core APIs: + - Engine API: docs/cactus_engine.md + - Graph API: docs/cactus_graph.md + - Kernels: docs/cactus_kernels.md + - Index API: docs/cactus_index.md + - Cactus Quants: docs/cactus_quants.md + - Hybrid Inference: docs/cactus_hybrid.md + - Transpiler: docs/cactus_transpiler.md + - Guides: + - Fine-tuning & Deployment: docs/finetuning.md + - Runtime Compatibility: docs/compatibility.md + - Contributing: CONTRIBUTING.md + - Blog: + - All Posts: blog/README.md + - "TurboQuant-H": blog/turboquant-h.md + - "Gemma 4 on Cactus": blog/gemma4.md + - Hybrid Transcription: blog/hybrid_transcription.md + - "LFM2-24B on Mac": blog/lfm2_24b_a2b.md + - "Parakeet CTC 1.1B": blog/parakeet.md + - "LFM-2.5-350m": blog/lfm2.5_350m.md diff --git a/python/README.md b/python/README.md index ba0d6ea0d..91d1a452b 100644 --- a/python/README.md +++ b/python/README.md @@ -1,288 +1,393 @@ +--- +title: "Cactus Python Package" +description: "Python package and ctypes bindings for the Cactus on-device AI inference engine." +keywords: ["Python package", "Python bindings", "on-device AI", "Python FFI", "embeddings", "transcription", "RAG"] +--- + # Cactus Python Package Python bindings for Cactus Engine via FFI. Auto-installed when you run `source ./setup`. +> **Model bundles:** Pre-built runtime bundles for all supported models at [huggingface.co/Cactus-Compute](https://huggingface.co/Cactus-Compute). + ## Getting Started + ```bash -# Setup environment -source ./setup - -# Build shared library for Python +git clone https://github.com/cactus-compute/cactus && cd cactus && source ./setup cactus build --python +``` + -# Download models +```bash +# Download pre-built bundles cactus download LiquidAI/LFM2-VL-450M cactus download openai/whisper-small + +# Optional: set your Cactus Cloud API key for automatic cloud fallback +cactus auth ``` ## Quick Example + ```python -from cactus import cactus_init, cactus_complete, cactus_destroy +from cactus import ensure_model, cactus_init, cactus_complete, cactus_destroy import json -model = cactus_init("weights/lfm2-vl-450m") - -messages = [{"role": "user", "content": "What is 2+2?"}] -response = json.loads(cactus_complete(model, messages)) -print(response["response"]) +# Downloads the pre-built bundle from HuggingFace if not already present +bundle = ensure_model("LiquidAI/LFM2-VL-450M") +model = cactus_init(str(bundle), None, False) +messages = json.dumps([{"role": "user", "content": "What is 2+2?"}]) +result = cactus_complete(model, messages, None, None, None) +print(result["response"]) cactus_destroy(model) ``` + ## API Reference -### `cactus_init(model_path, corpus_dir=None)` +All functions are module-level and mirror the C FFI directly. Handles are plain `int` values (C pointers). -Initialize a model and return its handle. +### Model Downloads -| Parameter | Type | Description | -|-----------|------|-------------| -| `model_path` | `str` | Path to model weights directory | -| `corpus_dir` | `str` | Optional path to RAG corpus directory for document Q&A | +Download pre-built bundles programmatically (no CLI needed): ```python -model = cactus_init("weights/lfm2-vl-450m") -rag_model = cactus_init("weights/lfm2-rag", corpus_dir="./documents") +from cactus import ensure_model, get_bundle_dir + +# ensure_model downloads the pre-built bundle if missing, returns its Path +bundle = ensure_model("openai/whisper-tiny") + +# Or resolve the expected on-disk location explicitly +bundle_dir = get_bundle_dir("openai/whisper-tiny", bits=4) +# -> Path("weights/whisper-tiny-cq4") ``` -### `cactus_complete(model, messages, **options)` - -Run chat completion. Returns JSON string with response and metrics. - -| Parameter | Type | Description | -|-----------|------|-------------| -| `model` | handle | Model handle from `cactus_init` | -| `messages` | `list\|str` | List of message dicts or JSON string | -| `tools` | `list` | Optional tool definitions for function calling | -| `temperature` | `float` | Sampling temperature | -| `top_p` | `float` | Top-p sampling | -| `top_k` | `int` | Top-k sampling | -| `max_tokens` | `int` | Maximum tokens to generate | -| `stop_sequences` | `list` | Stop sequences | -| `force_tools` | `bool` | Constrain output to tool call format | -| `tool_rag_top_k` | `int` | Select top-k relevant tools via Tool RAG (default: 2, 0 = use all tools) | -| `confidence_threshold` | `float` | Minimum confidence for local generation (default: 0.7, triggers cloud_handoff when below) | -| `callback` | `fn` | Streaming callback `fn(token, token_id, user_data)` | +### Init / Lifecycle ```python -# Basic completion -messages = [{"role": "user", "content": "Hello!"}] -response = cactus_complete(model, messages, max_tokens=100) -print(json.loads(response)["response"]) - -# With tools -tools = [{ - "name": "get_weather", - "description": "Get weather for a location", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"] - } -}] -response = cactus_complete(model, messages, tools=tools) +model = cactus_init(model_path: str, corpus_dir: str | None, cache_index: bool) -> int +cactus_destroy(model: int) +cactus_reset(model: int) # clear KV cache +cactus_stop(model: int) # abort ongoing generation +cactus_get_last_error() -> str | None +``` + +### Completion + +Returns a `dict` with `success`, `error`, `cloud_handoff`, `response`, optional `thinking` (only present when the model emits chain-of-thought content, placed before `function_calls`), `function_calls`, `segments` (always `[]` for completion — populated only for Whisper transcription with the `timestamps` option), `confidence`, timing stats (`time_to_first_token_ms`, `total_time_ms`, `prefill_tps`, `decode_tps`, `ram_usage_mb`), and token counts (`prefill_tokens`, `decode_tokens`, `total_tokens`). -# Streaming -def on_token(token, token_id, user_data): - print(token, end="", flush=True) +```python +result = cactus_complete( + model: int, + messages: list | str, # list of {role, content} dicts or JSON string + options: dict | str | None, # optional inference options + tools: list | str | None, # optional tool definitions + callback: Callable[[str, int], None] | None, # streaming token callback + pcm_data: list[int] | None = None # optional raw audio bytes +) -> dict +``` + +```python +# With options and streaming +options = json.dumps({"max_tokens": 256, "temperature": 0.7}) +def on_token(token, token_id): print(token, end="", flush=True) -cactus_complete(model, messages, callback=on_token) +result = cactus_complete(model, messages, options, None, on_token) +if result["cloud_handoff"]: + # response already contains cloud result + pass ``` -**Response format** (all fields always present): +**Response format:** ```json { "success": true, "error": null, "cloud_handoff": false, - "response": "Hello! How can I help?", + "response": "4", "function_calls": [], - "confidence": 0.85, + "segments": [], + "confidence": 0.92, + "confidence_threshold": 0.7, "time_to_first_token_ms": 45.2, "total_time_ms": 163.7, "prefill_tps": 619.5, "decode_tps": 168.4, - "ram_usage_mb": 245.67, + "ram_usage_mb": 512.3, "prefill_tokens": 28, - "decode_tokens": 50, - "total_tokens": 78 + "decode_tokens": 12, + "total_tokens": 40 } ``` -**Cloud handoff response** (when model detects low confidence): +### Prefill + +Pre-processes input text and populates the KV cache without generating output tokens. This reduces latency for subsequent calls to `cactus_complete`. + +```python +cactus_prefill( + model: int, + messages: list | str, # list of {role, content} dicts or JSON string + options: dict | str | None, # optional inference options + tools: list | str | None, # optional tool definitions + pcm_data: list[int] | None = None # optional raw audio bytes +) -> dict +``` + +```python +tools = json.dumps([{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City, State, Country"} + }, + "required": ["location"] + } + } +}]) + +messages = json.dumps([ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "assistant", "content": "<|tool_call_start|>get_weather(location=\"Paris\")<|tool_call_end|>"}, + {"role": "tool", "content": "{\"name\": \"get_weather\", \"content\": \"Sunny, 72°F\"}"}, + {"role": "assistant", "content": "It's sunny and 72°F in Paris!"} +]) +cactus_prefill(model, messages, None, tools) + +completion_messages = json.dumps([ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "assistant", "content": "<|tool_call_start|>get_weather(location=\"Paris\")<|tool_call_end|>"}, + {"role": "tool", "content": "{\"name\": \"get_weather\", \"content\": \"Sunny, 72°F\"}"}, + {"role": "assistant", "content": "It's sunny and 72°F in Paris!"}, + {"role": "user", "content": "What about SF?"} +]) +result = cactus_complete(model, completion_messages, None, tools, None) +``` + +**Response format:** ```json { - "success": false, + "success": true, "error": null, - "cloud_handoff": true, - "response": null, - "function_calls": [], - "confidence": 0.18, - "time_to_first_token_ms": 45.2, - "total_time_ms": 45.2, - "prefill_tps": 619.5, - "decode_tps": 0.0, - "ram_usage_mb": 245.67, - "prefill_tokens": 28, - "decode_tokens": 0, - "total_tokens": 28 + "prefill_tokens": 25, + "prefill_tps": 166.1, + "total_time_ms": 150.5, + "ram_usage_mb": 245.67 } ``` -When `cloud_handoff` is `True`, the model's confidence dropped below `confidence_threshold` (default: 0.7) and recommends deferring to a cloud-based model for better results. Handle this in your application: +### Transcription + +Returns a `dict` with the `response` field (transcribed text) and a `segments` array of `{start, end, text}` objects. `segments` is populated only for Whisper models when the `timestamps` option is set (`{"timestamps": True}`); it is empty otherwise, including for all Parakeet transcription. ```python -result = json.loads(cactus_complete(model, messages)) -if result["cloud_handoff"]: - # Defer to cloud API (e.g., OpenAI, Anthropic) - response = call_cloud_api(messages) -else: - response = result["response"] +result = cactus_transcribe( + model: int, + audio_path: str | None, + prompt: str | None, + options: dict | str | None, + callback: Callable[[str, int], None] | None, + pcm_data: list[int] | bytes | None +) -> dict ``` -### `cactus_transcribe(model, audio_path, prompt="")` - -Transcribe audio using a Whisper model. Returns JSON string. - -| Parameter | Type | Description | -|-----------|------|-------------| -| `model` | handle | Whisper model handle | -| `audio_path` | `str` | Path to audio file (WAV) | -| `prompt` | `str` | Whisper prompt for language/task | +**Custom vocabulary** biases the decoder toward domain-specific words (supported for Whisper and Moonshine models). Pass `custom_vocabulary` and `vocabulary_boost` in `options`: ```python -whisper = cactus_init("weights/whisper-small") -prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" -response = cactus_transcribe(whisper, "audio.wav", prompt=prompt) -print(json.loads(response)["response"]) -cactus_destroy(whisper) +options = json.dumps({ + "custom_vocabulary": ["Omeprazole", "HIPAA", "Cactus"], + "vocabulary_boost": 3.0 +}) +result = cactus_transcribe(model, "medical_notes.wav", None, options, None, None) ``` -### `cactus_embed(model, text, normalize=False)` +```python +result = cactus_transcribe(model, "/path/to/audio.wav", None, {"timestamps": True}, None, None) +print(result["response"]) +for seg in result["segments"]: + print(f"[{seg['start']:.3f}s - {seg['end']:.3f}s] {seg['text']}") +``` -Get text embeddings. Returns list of floats. +### Streaming transcription -| Parameter | Type | Description | -|-----------|------|-------------| -| `model` | handle | Model handle | -| `text` | `str` | Text to embed | -| `normalize` | `bool` | L2-normalize embeddings (default: False) | +Transcribe continuously while audio is still being captured (Whisper and Parakeet TDT). Open a session, push 16 kHz mono 16-bit PCM chunks, and read text back as it stabilizes: `confirmed` words are final (append them to your transcript), `pending` is the volatile tail (replace it each call, for live display only). ```python -embedding = cactus_embed(model, "Hello world") -print(f"Dimension: {len(embedding)}") +stream = cactus_stream_transcribe_start(model: int, options: dict | str | None) -> int +result = cactus_stream_transcribe_process(stream: int, pcm_data: bytes) -> dict # {"confirmed": str, "pending": str, plus per-call timing stats} +result = cactus_stream_transcribe_stop(stream: int) -> dict # {"confirmed": str, "pending": ""}; destroys the session ``` -### `cactus_image_embed(model, image_path)` - -Get image embeddings from a VLM. Returns list of floats. +`options` is forwarded to `cactus_transcribe` for **Whisper only** (e.g. `language`, `max_tokens`); the Parakeet TDT path ignores it. Chunking is handled internally. ```python -embedding = cactus_image_embed(model, "image.png") +stream = cactus_stream_transcribe_start(model, {"language": "en"}) +transcript = "" +for chunk in pcm_chunks: # each chunk: 16 kHz mono 16-bit PCM bytes + out = cactus_stream_transcribe_process(stream, chunk) + transcript += out["confirmed"] # show out["pending"] separately as a live preview +transcript += cactus_stream_transcribe_stop(stream)["confirmed"] ``` -### `cactus_audio_embed(model, audio_path)` - -Get audio embeddings from a Whisper model. Returns list of floats. +### Embeddings ```python -embedding = cactus_audio_embed(whisper, "audio.wav") +embedding = cactus_embed(model: int, text: str, normalize: bool) -> list[float] +embedding = cactus_image_embed(model: int, image_path: str) -> list[float] +embedding = cactus_audio_embed(model: int, audio_path: str) -> list[float] ``` -### `cactus_reset(model)` +### Tokenization + +```python +tokens = cactus_tokenize(model: int, text: str) -> list[int] +result = cactus_score_window(model: int, tokens: list[int], start: int, end: int, context: int) -> dict +``` -Reset model state (clear KV cache). Call between unrelated conversations. +### RAG ```python -cactus_reset(model) +result = cactus_rag_query(model: int, query: str, top_k: int) -> dict ``` -### `cactus_stop(model)` +Returns a `dict` with a `chunks` array. Each chunk has `score` (float), `source` (str, from document metadata), and `content` (str): -Stop an ongoing generation (useful with streaming callbacks). +```json +{ + "chunks": [ + {"score": 0.0142, "source": "doc.txt", "content": "relevant passage..."} + ] +} +``` + +### Vector Index ```python -cactus_stop(model) +index = cactus_index_init(index_dir: str, embedding_dim: int) -> int +cactus_index_add(index: int, ids: list[int], documents: list[str], + metadatas: list[str] | None, embeddings: list[list[float]]) +cactus_index_delete(index: int, ids: list[int]) +result = cactus_index_get(index: int, ids: list[int]) -> dict +result = cactus_index_query(index: int, embedding: list[float], options: dict | str | None) -> dict +cactus_index_compact(index: int) +cactus_index_destroy(index: int) ``` -### `cactus_destroy(model)` +`cactus_index_query` returns `{"results":[{"id":,"score":}, ...]}`. `cactus_index_get` returns `{"results":[{"document":"...","metadata":,"embedding":[...]}, ...]}`. -Free model memory. Always call when done. +### Logging ```python -cactus_destroy(model) +cactus_log_set_level(level: int) # 0=DEBUG 1=INFO 2=WARN (default) 3=ERROR 4=NONE +cactus_log_set_callback(callback: Callable[[int, str, str], None] | None) ``` -### `cactus_get_last_error()` - -Get the last error message, or `None` if no error. +### Telemetry ```python -error = cactus_get_last_error() -if error: - print(f"Error: {error}") +cactus_set_telemetry_environment(framework: str, cache_location: str | None, version: str | None) +cactus_set_app_id(app_id: str) +cactus_telemetry_flush() +cactus_telemetry_shutdown() ``` -### `cactus_set_pro_key(key)` +Functions that return a value raise `RuntimeError` on failure. `cactus_index_add`, `cactus_index_delete`, and `cactus_index_compact` also raise `RuntimeError` on failure despite not returning a value. Truly void functions that never raise: `cactus_destroy`, `cactus_reset`, `cactus_stop`, `cactus_index_destroy`, logging and telemetry functions. + +## Vision (VLM) -Set Cactus Pro key for NPU acceleration (Apple devices). +Pass images in the messages content for vision-language models (LFM2-VL, LFM2.5-VL, Gemma4, Qwen3.5): ```python -cactus_set_pro_key("your-key") # email founders@cactuscompute.com +messages = json.dumps([{ + "role": "user", + "content": "Describe this image", + "images": ["path/to/image.png"] +}]) +result = cactus_complete(model, messages, None, None, None) +print(result["response"]) ``` -### `cactus_tokenize(model, text)` +## Audio (Multimodal) -Tokenize text. Returns list of token IDs. +Pass audio files in messages for models with native audio understanding (Gemma4): ```python -tokens = cactus_tokenize(model, "Hello world") -print(tokens) # [1234, 5678, ...] +messages = json.dumps([{ + "role": "user", + "content": "Transcribe the audio.", + "audio": ["path/to/audio.wav"] +}]) +result = cactus_complete(model, messages, None, None, None) +print(result["response"]) + +# Combined vision + audio +messages = json.dumps([{ + "role": "user", + "content": "Describe the image and transcribe the audio.", + "images": ["path/to/image.png"], + "audio": ["path/to/audio.wav"] +}]) +result = cactus_complete(model, messages, None, None, None) ``` -### `cactus_rag_query(model, query, top_k=5)` - -Query RAG corpus for relevant text chunks. Requires model initialized with `corpus_dir`. +## Compute Graph -| Parameter | Type | Description | -|-----------|------|-------------| -| `model` | handle | Model handle (must have corpus_dir set) | -| `query` | `str` | Query text | -| `top_k` | `int` | Number of chunks to retrieve (default: 5) | +The `Graph` API provides a tensor computation graph for building and executing dataflow pipelines on the Cactus kernel layer: ```python -model = cactus_init("weights/lfm2-rag", corpus_dir="./documents") -chunks = cactus_rag_query(model, "What is machine learning?", top_k=3) -for chunk in chunks: - print(f"Score: {chunk['score']:.2f} - {chunk['text'][:100]}...") +from cactus.bindings.cactus import Graph +import numpy as np + +g = Graph() +a = g.input((2, 2)) +b = g.input((2, 2)) +y = ((a - b) * (a + b)).abs().pow(2.0).view((4,)) + +g.set_input(a, np.array([[2, 4], [6, 8]], dtype=np.float16)) +g.set_input(b, np.array([[1, 2], [3, 4]], dtype=np.float16)) +g.execute() + +print(y.numpy()) # [9. 144. 729. 2304.] ``` -## Vision (VLM) +Supported ops: `+`, `-`, `*`, `/`, `abs`, `pow`, `view`, `flatten`, `concat`, `cat`, `relu`, `sigmoid`, `tanh`, `gelu`, `softmax`. -Pass images in the messages for vision-language models: +Every op accepts a `backend=` kwarg (`Graph.CPU` or `Graph.METAL`) to pin that op; by default it follows the global backend — auto (best available), or the one forced with `cactus_set_backend("cpu"|"metal")`: ```python -vlm = cactus_init("weights/lfm2-vl-450m") - -messages = [{ - "role": "user", - "content": "Describe this image", - "images": ["path/to/image.png"] -}] -response = cactus_complete(vlm, messages) -print(json.loads(response)["response"]) +m = g.matmul(a, b, backend=Graph.METAL) ``` -## Full Example +## Testing -See `python/example.py` for a complete example covering: -- Text completion -- Text/image/audio embeddings -- Vision (VLM) -- Speech transcription +Run the full test suite: ```bash -python python/example.py -``` \ No newline at end of file +python python/test.py # compact output +python python/test.py -v # verbose +``` + +Tests are in `python/tests/` — bindings, CLI, server, graph, model, transpile, +and component-partition coverage. Add a new `test_*.py` to extend. + +## See Also + +- `Cactus Engine API` — Full C API reference that the Python bindings wrap +- `Cactus Index API` — Vector database API for RAG applications +- `Fine-tuning Guide` — Train and deploy custom LoRA fine-tunes +- `Runtime Compatibility` — Weight versioning across releases +- [Apple Build Step](/apple/) — Builds Apple native artifacts used by bindings +- [Android Build Step](/android/) — Builds Android native artifacts used by bindings +- [Swift Bindings](/bindings/swift/) — Swift C-module bindings +- [Kotlin Bindings](/bindings/kotlin/) — Kotlin/JNI bindings +- [Flutter Bindings](/bindings/flutter/) — Dart FFI bindings +- [Rust Bindings](/bindings/rust/) — Raw Rust FFI declarations diff --git a/python/cactus/__init__.py b/python/cactus/__init__.py new file mode 100644 index 000000000..1af766aee --- /dev/null +++ b/python/cactus/__init__.py @@ -0,0 +1,59 @@ +"""Cactus — on-device AI inference.""" + +from ._version import __version__ + +from .cli.download import ensure_model, get_weights_dir, get_bundle_dir, get_model_dir_name +from .cli import main + +__all__ = [ + "__version__", + "ensure_model", + "get_weights_dir", + "get_bundle_dir", + "get_model_dir_name", + "main", + "Graph", + "Tensor", + "cactus_init", + "cactus_set_backend", + "cactus_destroy", + "cactus_reset", + "cactus_stop", + "cactus_complete", + "cactus_prefill", + "cactus_embed", + "cactus_image_embed", + "cactus_audio_embed", + "cactus_transcribe", + "cactus_stream_transcribe_start", + "cactus_stream_transcribe_process", + "cactus_stream_transcribe_stop", + "cactus_tokenize", + "cactus_render_prompt", + "cactus_score_window", + "cactus_rag_query", + "cactus_index_init", + "cactus_index_add", + "cactus_index_delete", + "cactus_index_query", + "cactus_index_get", + "cactus_index_compact", + "cactus_index_destroy", + "cactus_set_app_id", + "cactus_set_telemetry_environment", + "cactus_telemetry_flush", + "cactus_telemetry_shutdown", + "cactus_log_set_level", + "cactus_log_set_callback", + "cactus_get_last_error", + "cactus_preprocess_audio_features", +] + +_FFI_NAMES = frozenset(n for n in __all__ if n.startswith("cactus_")) + + +def __getattr__(name): + if name in ("Graph", "Tensor") or name in _FFI_NAMES: + from .bindings import cactus + return getattr(cactus, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/src/__main__.py b/python/cactus/__main__.py similarity index 100% rename from python/src/__main__.py rename to python/cactus/__main__.py diff --git a/python/cactus/_version.py b/python/cactus/_version.py new file mode 100644 index 000000000..f627338ef --- /dev/null +++ b/python/cactus/_version.py @@ -0,0 +1,13 @@ +"""Stamped by CI before building; dev fallback reads CACTUS_VERSION.""" +from pathlib import Path as _Path + + +def _dev_version(): + vfile = _Path(__file__).resolve().parent.parent.parent / "CACTUS_VERSION" + if vfile.exists(): + raw = vfile.read_text().strip().lstrip("v") + return raw + ".0" if raw.count(".") == 1 else raw + return "0.0.0+dev" + + +__version__ = _dev_version() diff --git a/python/cactus/assets/test.wav b/python/cactus/assets/test.wav new file mode 100644 index 000000000..99638e372 Binary files /dev/null and b/python/cactus/assets/test.wav differ diff --git a/python/cactus/assets/test_monkey.png b/python/cactus/assets/test_monkey.png new file mode 100644 index 000000000..cab0ac1f4 Binary files /dev/null and b/python/cactus/assets/test_monkey.png differ diff --git a/python/cactus/bindings/__init__.py b/python/cactus/bindings/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py new file mode 100644 index 000000000..5901d6a4c --- /dev/null +++ b/python/cactus/bindings/cactus.py @@ -0,0 +1,2988 @@ +"""Cactus Python FFI bindings.""" +import ctypes +import json +import platform +from pathlib import Path + +import numpy as np + +TokenCallback = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_uint32, ctypes.c_void_p) + +_LIB_NAME = "libcactus_engine.dylib" if platform.system() == "Darwin" else "libcactus_engine.so" + + +def _find_library(): + bundled = Path(__file__).parent / "lib" / _LIB_NAME + if bundled.exists(): + return bundled + + dev_build = Path(__file__).parent.parent.parent.parent / "cactus-engine" / "build" / _LIB_NAME + if dev_build.exists(): + return dev_build + + raise RuntimeError( + f"Cactus library ({_LIB_NAME}) not found.\n" + f"Install with: pip install cactus-compute\n" + f"Or build from source: cactus build --python" + ) + + +_LIB_PATH = _find_library() +try: + _lib = ctypes.CDLL(str(_LIB_PATH)) +except OSError as exc: + raise RuntimeError( + f"Cactus library found at {_LIB_PATH} but failed to load: {exc}.\n" + "Common causes: wrong arch (the library is built for arm64 only), " + "missing libcurl on Linux, or a stale build. Rebuild with " + "`cactus build --python`." + ) from exc + + +def _bind_optional(name, argtypes, restype): + try: + fn = getattr(_lib, name) + except AttributeError: + return None + fn.argtypes = argtypes + fn.restype = restype + return fn + +cactus_graph_t = ctypes.c_void_p +cactus_node_t = ctypes.c_uint64 + +class cactus_tensor_info_t(ctypes.Structure): + _fields_ = [ + ("precision", ctypes.c_int32), + ("rank", ctypes.c_size_t), + ("shape", ctypes.c_size_t * 8), + ("num_elements", ctypes.c_size_t), + ("byte_size", ctypes.c_size_t), + ] + +_lib.cactus_set_telemetry_environment.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p] +_lib.cactus_set_telemetry_environment.restype = None +_lib.cactus_set_telemetry_environment(b"python", None, None) + +_lib.cactus_init.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_bool] +_lib.cactus_init.restype = ctypes.c_void_p + +_lib.cactus_set_backend.argtypes = [ctypes.c_char_p] +_lib.cactus_set_backend.restype = ctypes.c_int + +# cactus graph API +_lib.cactus_graph_create.restype = cactus_graph_t +_lib.cactus_graph_destroy.argtypes = [cactus_graph_t] +_lib.cactus_graph_destroy.restype = None +_lib.cactus_graph_hard_reset.argtypes = [cactus_graph_t] +_lib.cactus_graph_hard_reset.restype = ctypes.c_int +_lib.cactus_graph_set_node_backend.argtypes = [cactus_graph_t, cactus_node_t, ctypes.c_int32] +_lib.cactus_graph_set_node_backend.restype = ctypes.c_int + +_lib.cactus_graph_save.argtypes = [cactus_graph_t, ctypes.c_char_p] +_lib.cactus_graph_save.restype = ctypes.c_int + +_lib.cactus_graph_load.argtypes = [ctypes.c_char_p] +_lib.cactus_graph_load.restype = cactus_graph_t + +_lib.cactus_graph_input.argtypes = [ + cactus_graph_t, + ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, + ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_input.restype = ctypes.c_int + +_lib.cactus_graph_set_input.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_void_p, ctypes.c_int32 +] +_lib.cactus_graph_set_input.restype = ctypes.c_int +_lib.cactus_graph_set_external_input.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_void_p, ctypes.c_int32 +] +_lib.cactus_graph_set_external_input.restype = ctypes.c_int +_bind_optional( + "cactus_graph_mark_embedded_input", + [cactus_graph_t, cactus_node_t], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_set_runtime_input_shape", + [cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_set_input_dynamic_dims", + [cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t], + ctypes.c_int, +) + +_lib.cactus_graph_add.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_add.restype = ctypes.c_int +_lib.cactus_graph_add_clipped.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_add_clipped.restype = ctypes.c_int + +_lib.cactus_graph_subtract.argtypes = [cactus_graph_t, cactus_node_t, + cactus_node_t, ctypes.POINTER(cactus_node_t)] +_lib.cactus_graph_subtract.restype = ctypes.c_int + +_lib.cactus_graph_multiply.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_multiply.restype = ctypes.c_int + +_lib.cactus_graph_divide.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_divide.restype = ctypes.c_int +_bind_optional( + "cactus_graph_not_equal", + [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) + +_lib.cactus_graph_precision_cast.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_precision_cast.restype = ctypes.c_int +_bind_optional( + "cactus_graph_quantize_activations", + [cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) + +_lib.cactus_graph_scalar_add.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_add.restype = ctypes.c_int +_lib.cactus_graph_scalar_subtract.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_subtract.restype = ctypes.c_int +_lib.cactus_graph_scalar_multiply.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_multiply.restype = ctypes.c_int +_lib.cactus_graph_scalar_divide.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_divide.restype = ctypes.c_int +_bind_optional( + "cactus_graph_scalar_floor_divide", + [cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_scalar_not_equal", + [cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_lib.cactus_graph_scalar_exp.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_exp.restype = ctypes.c_int +_lib.cactus_graph_scalar_sqrt.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_sqrt.restype = ctypes.c_int +_lib.cactus_graph_scalar_cos.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_cos.restype = ctypes.c_int +_lib.cactus_graph_scalar_sin.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_sin.restype = ctypes.c_int +_lib.cactus_graph_scalar_log.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scalar_log.restype = ctypes.c_int +_bind_optional( + "cactus_graph_masked_select_prefix", + [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_masked_scatter", + [cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) + +_lib.cactus_graph_abs.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_abs.restype = ctypes.c_int + +_lib.cactus_graph_pow.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_pow.restype = ctypes.c_int + +_lib.cactus_graph_view.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, + ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_view.restype = ctypes.c_int + +_lib.cactus_graph_flatten.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_flatten.restype = ctypes.c_int +_lib.cactus_graph_reshape.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_reshape.restype = ctypes.c_int +_bind_optional( + "cactus_graph_expand", + [cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_lib.cactus_graph_transpose.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_transpose.restype = ctypes.c_int +_lib.cactus_graph_transpose_n.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_transpose_n.restype = ctypes.c_int +_lib.cactus_graph_slice.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_slice.restype = ctypes.c_int +_lib.cactus_graph_index.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_index.restype = ctypes.c_int + +_lib.cactus_graph_concat.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_concat.restype = ctypes.c_int + +_lib.cactus_graph_cat.argtypes = [ + cactus_graph_t, ctypes.POINTER(cactus_node_t), ctypes.c_size_t, ctypes.c_int32, + ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_cat.restype = ctypes.c_int +_lib.cactus_graph_matmul.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_matmul.restype = ctypes.c_int +_lib.cactus_graph_gather.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gather.restype = ctypes.c_int +_lib.cactus_graph_embedding_from_tensor.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_embedding_from_tensor.restype = ctypes.c_int +_lib.cactus_graph_embedding_from_file.argtypes = [ + cactus_graph_t, ctypes.c_char_p, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_embedding_from_file.restype = ctypes.c_int +_lib.cactus_graph_mmap_embeddings.argtypes = [ + cactus_graph_t, ctypes.c_char_p, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_mmap_embeddings.restype = ctypes.c_int +_lib.cactus_graph_mmap_weights.argtypes = [ + cactus_graph_t, ctypes.c_char_p, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_mmap_weights.restype = ctypes.c_int +_bind_optional( + "cactus_graph_bind_mmap_weights", + [cactus_graph_t, cactus_node_t, ctypes.c_char_p], + ctypes.c_int, +) +_lib.cactus_graph_bilinear_interpolation.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_bilinear_interpolation.restype = ctypes.c_int +_bind_optional( + "cactus_graph_set_grouped_scales", + [cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.c_void_p], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_set_interleaved", + [cactus_graph_t, cactus_node_t, ctypes.c_bool, ctypes.c_size_t], + ctypes.c_int, +) +_lib.cactus_graph_release_weight_pages.argtypes = [cactus_graph_t, cactus_node_t] +_lib.cactus_graph_release_weight_pages.restype = ctypes.c_int +_lib.cactus_graph_prefetch_weight_pages.argtypes = [cactus_graph_t, cactus_node_t] +_lib.cactus_graph_prefetch_weight_pages.restype = ctypes.c_int +_lib.cactus_graph_release_all_weight_pages.argtypes = [cactus_graph_t] +_lib.cactus_graph_release_all_weight_pages.restype = ctypes.c_int + +_lib.cactus_graph_sum.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_sum.restype = ctypes.c_int +_lib.cactus_graph_mean.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_mean.restype = ctypes.c_int +_lib.cactus_graph_variance.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_variance.restype = ctypes.c_int +_lib.cactus_graph_min.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_min.restype = ctypes.c_int +_lib.cactus_graph_max.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_max.restype = ctypes.c_int +_lib.cactus_graph_cumsum.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_cumsum.restype = ctypes.c_int + +_lib.cactus_graph_relu.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_relu.restype = ctypes.c_int +_lib.cactus_graph_silu.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_silu.restype = ctypes.c_int +_lib.cactus_graph_gelu.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gelu.restype = ctypes.c_int +_lib.cactus_graph_gelu_erf.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gelu_erf.restype = ctypes.c_int +_lib.cactus_graph_sigmoid.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_sigmoid.restype = ctypes.c_int +_lib.cactus_graph_tanh.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_tanh.restype = ctypes.c_int +_lib.cactus_graph_glu.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_glu.restype = ctypes.c_int + +_lib.cactus_graph_layernorm.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.c_bool, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_layernorm.restype = ctypes.c_int +_lib.cactus_graph_groupnorm.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_size_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_groupnorm.restype = ctypes.c_int +_lib.cactus_graph_batchnorm.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_int32, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_batchnorm.restype = ctypes.c_int +_lib.cactus_graph_rms_norm.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_rms_norm.restype = ctypes.c_int +_lib.cactus_graph_topk.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_topk.restype = ctypes.c_int +_lib.cactus_graph_rope.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_rope.restype = ctypes.c_int +_lib.cactus_graph_rope_gptj.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_rope_gptj.restype = ctypes.c_int +_lib.cactus_graph_softmax.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_softmax.restype = ctypes.c_int +_lib.cactus_graph_attention.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.c_bool, + ctypes.c_size_t, ctypes.c_size_t, ctypes.c_bool, cactus_node_t, ctypes.c_bool, + ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_attention.restype = ctypes.c_int +_lib.cactus_graph_kv_cache_state.argtypes = [ + cactus_graph_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_kv_cache_state.restype = ctypes.c_int +_lib.cactus_graph_kv_cache_append.argtypes = [ + cactus_graph_t, + cactus_node_t, + cactus_node_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_kv_cache_append.restype = ctypes.c_int +_lib.cactus_graph_attention_cached.argtypes = [ + cactus_graph_t, + cactus_node_t, + cactus_node_t, + cactus_node_t, + cactus_node_t, + cactus_node_t, + ctypes.c_float, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_attention_cached.restype = ctypes.c_int +_lib.cactus_graph_conv_cache_state.argtypes = [ + cactus_graph_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_conv_cache_state.restype = ctypes.c_int +_lib.cactus_graph_conv_cache_append.argtypes = [ + cactus_graph_t, + cactus_node_t, + cactus_node_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_conv_cache_append.restype = ctypes.c_int +_lib.cactus_graph_conv_cache_initialize.argtypes = [ + cactus_graph_t, + cactus_node_t, + cactus_node_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_conv_cache_initialize.restype = ctypes.c_int +_lib.cactus_graph_recurrent_cache_state.argtypes = [ + cactus_graph_t, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.c_int, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_recurrent_cache_state.restype = ctypes.c_int +_lib.cactus_graph_recurrent_cache_write.argtypes = [ + cactus_graph_t, + cactus_node_t, + cactus_node_t, + ctypes.POINTER(cactus_node_t), +] +_lib.cactus_graph_recurrent_cache_write.restype = ctypes.c_int +_lib.cactus_graph_rel_pos_bias.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_rel_pos_bias.restype = ctypes.c_int +_lib.cactus_graph_attention_int8_hybrid.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, + ctypes.POINTER(ctypes.c_int8), ctypes.POINTER(ctypes.c_int8), + ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float), + ctypes.c_size_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_attention_int8_hybrid.restype = ctypes.c_int +_lib.cactus_graph_rfft.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_rfft.restype = ctypes.c_int +_lib.cactus_graph_irfft.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_irfft.restype = ctypes.c_int +_lib.cactus_graph_mel_filter_bank.argtypes = [ + cactus_graph_t, ctypes.c_size_t, ctypes.c_size_t, + ctypes.c_float, ctypes.c_float, ctypes.c_size_t, + ctypes.c_int, ctypes.c_int, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_mel_filter_bank.restype = ctypes.c_int +_lib.cactus_graph_spectrogram.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, + ctypes.c_size_t, ctypes.c_size_t, ctypes.c_size_t, + ctypes.c_float, ctypes.c_bool, ctypes.c_int, + ctypes.c_float, ctypes.c_int, + ctypes.c_float, ctypes.c_float, ctypes.c_bool, + ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_spectrogram.restype = ctypes.c_int +_lib.cactus_graph_image_preprocess.argtypes = [ + cactus_graph_t, cactus_node_t, + ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_int, ctypes.c_int, ctypes.c_float, + ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float), ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_image_preprocess.restype = ctypes.c_int +_lib.cactus_graph_conv1d_causal.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv1d_causal.restype = ctypes.c_int +_lib.cactus_graph_conv1d_k3.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv1d_k3.restype = ctypes.c_int +_lib.cactus_graph_conv1d_k7s3.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv1d_k7s3.restype = ctypes.c_int +_lib.cactus_graph_conv1d.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv1d.restype = ctypes.c_int +_lib.cactus_graph_conv1d_same_depthwise_k9.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv1d_same_depthwise_k9.restype = ctypes.c_int +_lib.cactus_graph_conv1d_pointwise.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv1d_pointwise.restype = ctypes.c_int +_bind_optional( + "cactus_graph_clamp", + [cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_float, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_lib.cactus_graph_conv2d_k3s2p1.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv2d_k3s2p1.restype = ctypes.c_int +_lib.cactus_graph_conv2d_depthwise_k3s2p1.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv2d_depthwise_k3s2p1.restype = ctypes.c_int +_lib.cactus_graph_conv2d_pointwise_1x1.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_conv2d_pointwise_1x1.restype = ctypes.c_int +_bind_optional( + "cactus_graph_conv2d_k3s1p1", + [cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, cactus_node_t, ctypes.POINTER(cactus_node_t)], + ctypes.c_int, +) +_bind_optional( + "cactus_graph_conv2d", + [ + cactus_graph_t, + cactus_node_t, + cactus_node_t, + ctypes.c_bool, + cactus_node_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.POINTER(cactus_node_t), + ], + ctypes.c_int, +) +_lib.cactus_graph_lstm_cell.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_lstm_cell.restype = ctypes.c_int +_lib.cactus_graph_gated_deltanet_decode.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gated_deltanet_decode.restype = ctypes.c_int +_lib.cactus_graph_gated_deltanet_prefill.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_size_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gated_deltanet_prefill.restype = ctypes.c_int +_lib.cactus_graph_stft.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_stft.restype = ctypes.c_int +_lib.cactus_graph_altup_predict.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t), ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_altup_predict.restype = ctypes.c_int +_lib.cactus_graph_altup_correct.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t), ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_altup_correct.restype = ctypes.c_int +_lib.cactus_graph_gaussian_topk.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gaussian_topk.restype = ctypes.c_int +_lib.cactus_graph_moe_layer_gated.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, + ctypes.POINTER(cactus_node_t), ctypes.POINTER(cactus_node_t), ctypes.POINTER(cactus_node_t), + ctypes.c_size_t, ctypes.c_size_t, ctypes.c_bool, ctypes.c_float, ctypes.c_float, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_moe_layer_gated.restype = ctypes.c_int +_bind_optional( + "cactus_graph_dense_mlp_tq_fused", + [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, + ctypes.c_float, ctypes.POINTER(cactus_node_t) + ], + ctypes.c_int, +) +_lib.cactus_graph_moe_layer_ungated.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, + ctypes.POINTER(cactus_node_t), ctypes.POINTER(cactus_node_t), + ctypes.c_size_t, ctypes.c_size_t, ctypes.c_bool, ctypes.c_float, ctypes.c_float, ctypes.c_int32, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_moe_layer_ungated.restype = ctypes.c_int +_lib.cactus_graph_sample.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_float, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_sample.restype = ctypes.c_int +_lib.cactus_graph_scatter_topk.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_scatter_topk.restype = ctypes.c_int +_lib.cactus_graph_persistent.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_persistent.restype = ctypes.c_int +_lib.cactus_graph_is_populated.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_int32) +] +_lib.cactus_graph_is_populated.restype = ctypes.c_int +_lib.cactus_graph_invalidate_persistent.argtypes = [cactus_graph_t, cactus_node_t] +_lib.cactus_graph_invalidate_persistent.restype = ctypes.c_int + +_lib.cactus_graph_execute.argtypes = [cactus_graph_t] +_lib.cactus_graph_execute.restype = ctypes.c_int + +_lib.cactus_graph_get_output_ptr.argtypes = [cactus_graph_t, cactus_node_t, + ctypes.POINTER(ctypes.c_void_p)] +_lib.cactus_graph_get_output_ptr.restype = ctypes.c_int + +_lib.cactus_graph_get_output_info.argtypes = [ + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_tensor_info_t) +] +_lib.cactus_graph_get_output_info.restype = ctypes.c_int + +_lib.cactus_complete.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_size_t, + ctypes.c_char_p, ctypes.c_char_p, TokenCallback, ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t +] +_lib.cactus_complete.restype = ctypes.c_int + +_lib.cactus_prefill.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_size_t, + ctypes.c_char_p, ctypes.c_char_p, + ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t +] +_lib.cactus_prefill.restype = ctypes.c_int + +_lib.cactus_transcribe.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p, + ctypes.c_size_t, ctypes.c_char_p, TokenCallback, ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t +] +_lib.cactus_transcribe.restype = ctypes.c_int + +_lib.cactus_stream_transcribe_start.argtypes = [ctypes.c_void_p, ctypes.c_char_p] +_lib.cactus_stream_transcribe_start.restype = ctypes.c_void_p + +_lib.cactus_stream_transcribe_process.argtypes = [ + ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t, + ctypes.c_char_p, ctypes.c_size_t +] +_lib.cactus_stream_transcribe_process.restype = ctypes.c_int + +_lib.cactus_stream_transcribe_stop.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t] +_lib.cactus_stream_transcribe_stop.restype = ctypes.c_int + +_lib.cactus_embed.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), + ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_bool +] +_lib.cactus_embed.restype = ctypes.c_int + +_lib.cactus_image_embed.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), + ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t) +] +_lib.cactus_image_embed.restype = ctypes.c_int + +_lib.cactus_audio_embed.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), + ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t) +] +_lib.cactus_audio_embed.restype = ctypes.c_int + +try: + _lib.cactus_preprocess_audio_features.argtypes = [ + ctypes.c_char_p, ctypes.c_char_p, ctypes.c_size_t, + ctypes.POINTER(ctypes.c_float), ctypes.c_size_t, + ctypes.POINTER(ctypes.c_size_t), ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t) + ] + _lib.cactus_preprocess_audio_features.restype = ctypes.c_int +except AttributeError: + pass + + +_lib.cactus_reset.argtypes = [ctypes.c_void_p] +_lib.cactus_reset.restype = None + +_lib.cactus_stop.argtypes = [ctypes.c_void_p] +_lib.cactus_stop.restype = None + +_lib.cactus_destroy.argtypes = [ctypes.c_void_p] +_lib.cactus_destroy.restype = None + +_lib.cactus_get_last_error.argtypes = [] +_lib.cactus_get_last_error.restype = ctypes.c_char_p + +_lib.cactus_tokenize.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.POINTER(ctypes.c_uint32), + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_size_t), +] +_lib.cactus_tokenize.restype = ctypes.c_int + +_bind_optional( + "cactus_render_prompt", + [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int, +) + +_lib.cactus_score_window.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint32), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.c_size_t, +] +_lib.cactus_score_window.restype = ctypes.c_int + +_lib.cactus_rag_query.argtypes = [ + ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, + ctypes.c_size_t, ctypes.c_size_t +] +_lib.cactus_rag_query.restype = ctypes.c_int + +_lib.cactus_index_init.argtypes = [ctypes.c_char_p, ctypes.c_size_t] +_lib.cactus_index_init.restype = ctypes.c_void_p + +_lib.cactus_index_add.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_char_p), + ctypes.POINTER(ctypes.c_char_p), + ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), + ctypes.c_size_t, + ctypes.c_size_t +] +_lib.cactus_index_add.restype = ctypes.c_int + +_lib.cactus_index_delete.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int), + ctypes.c_size_t +] +_lib.cactus_index_delete.restype = ctypes.c_int + +_lib.cactus_index_query.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_char_p, + ctypes.POINTER(ctypes.POINTER(ctypes.c_int)), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), + ctypes.POINTER(ctypes.c_size_t) +] +_lib.cactus_index_query.restype = ctypes.c_int + +_lib.cactus_index_compact.argtypes = [ctypes.c_void_p] +_lib.cactus_index_compact.restype = ctypes.c_int + +_lib.cactus_index_destroy.argtypes = [ctypes.c_void_p] +_lib.cactus_index_destroy.restype = None + +_lib.cactus_index_get.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int), + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_char_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_char_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), + ctypes.POINTER(ctypes.c_size_t) +] +_lib.cactus_index_get.restype = ctypes.c_int + +_lib.cactus_set_app_id.argtypes = [ctypes.c_char_p] +_lib.cactus_set_app_id.restype = None + +_lib.cactus_telemetry_flush.argtypes = [] +_lib.cactus_telemetry_flush.restype = None + +_lib.cactus_telemetry_shutdown.argtypes = [] +_lib.cactus_telemetry_shutdown.restype = None + +LogCallback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_void_p) + +_lib.cactus_log_set_level.argtypes = [ctypes.c_int] +_lib.cactus_log_set_level.restype = None + +_lib.cactus_log_set_callback.argtypes = [LogCallback, ctypes.c_void_p] +_lib.cactus_log_set_callback.restype = None + + +def _enc(s): + """Encode a string to bytes for C, or pass through None/bytes.""" + if s is None: + return None + return s.encode() if isinstance(s, str) else s + + +def _to_json(obj): + """Accept str, bytes, dict, list, or None — return encoded bytes for C.""" + if obj is None: + return None + if isinstance(obj, (dict, list)): + return json.dumps(obj).encode() + if isinstance(obj, str): + return obj.encode() + return obj + + +def _from_json(buf): + """Decode a ctypes string buffer to a dict. Returns {} on empty.""" + text = buf.value.decode("utf-8", errors="ignore") + if not text: + return {} + try: + return json.loads(text) + except json.JSONDecodeError as exc: + if len(buf.value) >= len(buf) - 1: + raise RuntimeError( + f"C engine filled the {len(buf)} byte response buffer; " + "response was truncated before valid JSON could be parsed." + ) from exc + raise + + +def _prepare_pcm(pcm_data): + """Marshal pcm_data bytes to (ctypes_ptr, size) for C. Returns (None, 0) if None.""" + if pcm_data is None: + return None, 0 + pcm_arr = (ctypes.c_uint8 * len(pcm_data))(*pcm_data) + return ctypes.cast(pcm_arr, ctypes.POINTER(ctypes.c_uint8)), len(pcm_data) + + +def _make_token_callback(callback): + """Wrap a Python callback(text, token_id) into a C-compatible TokenCallback.""" + if not callback: + return TokenCallback() + def _bridge(token_bytes, token_id, _): + callback(token_bytes.decode("utf-8", errors="ignore") if token_bytes else "", token_id) + return TokenCallback(_bridge) + + +def cactus_get_last_error(): + """Returns the last error message from the C runtime, or None.""" + result = _lib.cactus_get_last_error() + return result.decode() if result else None + + +def _err(default): + return cactus_get_last_error() or default + + +# ── Telemetry ──────────────────────────────────────────────────────── + + +def cactus_set_telemetry_environment(framework, cache_location, version): + """Sets telemetry environment (framework name, cache path, version).""" + _lib.cactus_set_telemetry_environment(_enc(framework), _enc(cache_location), _enc(version)) + + +def cactus_set_app_id(app_id): + """Sets the application identifier for telemetry.""" + _lib.cactus_set_app_id(_enc(app_id)) + + +def cactus_telemetry_flush(): + """Flushes pending telemetry events.""" + _lib.cactus_telemetry_flush() + + +def cactus_telemetry_shutdown(): + """Shuts down the telemetry subsystem.""" + _lib.cactus_telemetry_shutdown() + + +# ── Model lifecycle ────────────────────────────────────────────────── + + +def cactus_set_backend(backend): + """Select the inference backend ("cpu" or "metal"); unset defaults to auto (best available). Returns 0 on success, -1 on failure.""" + return _lib.cactus_set_backend(_enc(backend)) + + +def cactus_init(model_path, corpus_dir=None, cache_index=False): + """Load a model from disk. + + Args: + model_path: Path to the converted model weights directory. + corpus_dir: Optional path to a RAG corpus directory. + cache_index: Whether to cache the RAG index to disk. + + Returns: + An opaque model handle. Pass to other cactus_* functions. + Call cactus_destroy() when done. + """ + handle = _lib.cactus_init(_enc(model_path), _enc(corpus_dir), cache_index) + if not handle: + raise RuntimeError(_err("Failed to initialize model")) + return handle + + +def cactus_destroy(model): + """Free all resources associated with a model handle. Safe to call on + a null/already-destroyed handle (no-op).""" + if not model: + return + _lib.cactus_destroy(model) + + +def cactus_reset(model): + """Clear the KV cache, resetting the model to a fresh conversation state. + No-op when ``model`` is null.""" + if not model: + return + _lib.cactus_reset(model) + + +def cactus_stop(model): + """Signal the model to stop the current generation early. No-op when + ``model`` is null.""" + if not model: + return + _lib.cactus_stop(model) + + +# ── LLM completion ─────────────────────────────────────────────────── + + +def cactus_complete(model, messages, options=None, tools=None, callback=None, pcm_data=None): + """Run chat completion. + + Args: + model: Model handle from cactus_init(). + messages: List of message dicts, e.g. [{"role": "user", "content": "Hi"}]. + Also accepts a pre-serialized JSON string. + options: Optional dict of generation options (temperature, max_tokens, etc.). + tools: Optional list of tool definitions for function calling. + callback: Optional function(text, token_id) called on each generated token. + pcm_data: Optional raw PCM audio bytes for audio-capable models. + + Returns: + A dict with the completion response and metrics. + """ + buf = ctypes.create_string_buffer(1 << 20) + cb = _make_token_callback(callback) + pcm_ptr, pcm_size = _prepare_pcm(pcm_data) + rc = _lib.cactus_complete( + model, _to_json(messages), buf, len(buf), + _to_json(options), _to_json(tools), cb, None, pcm_ptr, pcm_size, + ) + if rc < 0: + raise RuntimeError(_err("Completion failed")) + return _from_json(buf) + + +def cactus_prefill(model, messages, options=None, tools=None, pcm_data=None): + """Pre-fill the KV cache with messages without generating a response. + + Returns: + A dict with prefill stats (tokens processed, latency, etc.). + """ + buf = ctypes.create_string_buffer(1 << 20) + pcm_ptr, pcm_size = _prepare_pcm(pcm_data) + rc = _lib.cactus_prefill( + model, _to_json(messages), buf, len(buf), + _to_json(options), _to_json(tools), pcm_ptr, pcm_size, + ) + if rc < 0: + raise RuntimeError(_err("Prefill failed")) + return _from_json(buf) + + +# ── Audio / speech ─────────────────────────────────────────────────── + + +def cactus_transcribe(model, audio_path, prompt=None, options=None, callback=None, pcm_data=None): + """Transcribe audio to text. + + Args: + model: Model handle. + audio_path: Path to a WAV audio file. + prompt: Optional prompt to guide transcription. + options: Optional dict of transcription options. Set ``timestamps: True`` (Whisper + only) to populate the ``segments`` array. + callback: Optional function(text, token_id) for streaming tokens. + pcm_data: Optional raw PCM audio bytes (alternative to audio_path). + + Returns: + A dict with the transcribed text under ``response`` and a ``segments`` array of + ``{start, end, text}`` objects, populated only for Whisper when ``timestamps`` is set + (empty otherwise, including all Parakeet transcription). + """ + buf = ctypes.create_string_buffer(1 << 20) + cb = _make_token_callback(callback) + pcm_ptr, pcm_size = _prepare_pcm(pcm_data) + rc = _lib.cactus_transcribe( + model, _enc(audio_path), _enc(prompt), buf, len(buf), + _to_json(options), cb, None, pcm_ptr, pcm_size, + ) + if rc < 0: + raise RuntimeError(_err("Transcription failed")) + return _from_json(buf) + + +def cactus_stream_transcribe_start(model, options=None): + """Open a streaming transcription session on a Whisper or Parakeet TDT model. + + Args: + model: Model handle. + options: Optional dict forwarded to the underlying transcribe call + (e.g. {"language": "en", "max_tokens": 256}); chunking is + handled internally. + + Returns: + An opaque stream handle. Feed audio with cactus_stream_transcribe_process + and finish with cactus_stream_transcribe_stop. + """ + stream = _lib.cactus_stream_transcribe_start(model, _to_json(options)) + if not stream: + raise RuntimeError(_err("Failed to start streaming transcription")) + return stream + + +def cactus_stream_transcribe_process(stream, pcm_data): + """Feed the next slice of audio to a streaming session. + + Args: + stream: Stream handle from cactus_stream_transcribe_start. + pcm_data: 16 kHz mono 16-bit PCM audio bytes for this chunk. + + Returns: + A dict with "confirmed" (newly finalized text, append it), "pending" (volatile tail, + replace it each call), and per-call stats ("decode_tps", + "total_time_ms", "time_to_first_token_ms", "decode_tokens"). + """ + buf = ctypes.create_string_buffer(1 << 16) + pcm_ptr, pcm_size = _prepare_pcm(pcm_data) + rc = _lib.cactus_stream_transcribe_process(stream, pcm_ptr, pcm_size, buf, len(buf)) + if rc < 0: + raise RuntimeError(_err("Stream transcription failed")) + return _from_json(buf) + + +def cactus_stream_transcribe_stop(stream): + """Flush buffered audio, return the final text, and destroy the session. + + Args: + stream: Stream handle from cactus_stream_transcribe_start. + + Returns: + A dict {"success": True, "confirmed": , "pending": ""}. The handle is invalid + after this call. + """ + buf = ctypes.create_string_buffer(1 << 16) + rc = _lib.cactus_stream_transcribe_stop(stream, buf, len(buf)) + if rc < 0: + raise RuntimeError(_err("Stream transcription stop failed")) + return _from_json(buf) + + +def cactus_preprocess_audio_features(audio_path, model_type, mel_bins, capacity): + """Compute mel spectrogram features from an audio file. + + Returns: + A tuple (values, mel_bins, frames) where values is a list of floats. + """ + if not hasattr(_lib, "cactus_preprocess_audio_features"): + raise RuntimeError("cactus_preprocess_audio_features is unavailable; rebuild with cactus build --python") + buf = (ctypes.c_float * int(capacity))() + feature_count = ctypes.c_size_t() + out_mels = ctypes.c_size_t() + out_frames = ctypes.c_size_t() + rc = _lib.cactus_preprocess_audio_features( + _enc(audio_path), _enc(model_type), + ctypes.c_size_t(int(mel_bins)), + buf, ctypes.sizeof(buf), + ctypes.byref(feature_count), ctypes.byref(out_mels), ctypes.byref(out_frames), + ) + if rc < 0: + raise RuntimeError(_err("Audio feature preprocessing failed")) + return list(buf[:feature_count.value]), int(out_mels.value), int(out_frames.value) + + +# ── Embeddings ─────────────────────────────────────────────────────── + + +def cactus_embed(model, text, normalize=True): + """Compute a text embedding. + + Args: + model: Model handle. + text: The text to embed. + normalize: Whether to L2-normalize the embedding (default True). + + Returns: + A list of floats (the embedding vector). + """ + buf = (ctypes.c_float * 4096)() + dim = ctypes.c_size_t() + rc = _lib.cactus_embed(model, _enc(text), buf, ctypes.sizeof(buf), ctypes.byref(dim), normalize) + if rc < 0: + raise RuntimeError(_err("Embedding failed")) + return list(buf[:dim.value]) + + +def cactus_image_embed(model, image_path): + """Compute an image embedding. Returns a list of floats.""" + buf = (ctypes.c_float * 4096)() + dim = ctypes.c_size_t() + rc = _lib.cactus_image_embed(model, _enc(image_path), buf, ctypes.sizeof(buf), ctypes.byref(dim)) + if rc < 0: + raise RuntimeError(_err("Image embedding failed")) + return list(buf[:dim.value]) + + +def cactus_audio_embed(model, audio_path): + """Compute an audio embedding. Returns a list of floats.""" + buf = (ctypes.c_float * 4096)() + dim = ctypes.c_size_t() + rc = _lib.cactus_audio_embed(model, _enc(audio_path), buf, ctypes.sizeof(buf), ctypes.byref(dim)) + if rc < 0: + raise RuntimeError(_err("Audio embedding failed")) + return list(buf[:dim.value]) + + +# ── Tokenization ───────────────────────────────────────────────────── + + +def cactus_tokenize(model, text): + """Tokenize text into token IDs. Returns a list of ints.""" + max_tokens = 8192 + arr = (ctypes.c_uint32 * max_tokens)() + n = ctypes.c_size_t(0) + rc = _lib.cactus_tokenize(model, _enc(text), arr, max_tokens, ctypes.byref(n)) + if rc < 0: + raise RuntimeError(_err("Tokenization failed")) + return list(arr[:n.value]) + + +def cactus_render_prompt(model, messages, options=None, tools=None): + """Render the chat-template prompt string for messages without generating.""" + if not hasattr(_lib, "cactus_render_prompt"): + raise RuntimeError("cactus_render_prompt is unavailable; rebuild with cactus build --python") + buf = ctypes.create_string_buffer(1 << 20) + rc = _lib.cactus_render_prompt( + model, _to_json(messages), _to_json(options), _to_json(tools), buf, len(buf), + ) + if rc < 0: + raise RuntimeError(_err("Prompt rendering failed")) + return buf.value.decode("utf-8", errors="ignore") + + +def cactus_score_window(model, tokens, start, end, context): + """Score a window of tokens for log-probabilities. + + Returns: + A dict with token-level log-probability scores. + """ + buf = ctypes.create_string_buffer(65536) + arr = (ctypes.c_uint32 * len(tokens))(*tokens) + rc = _lib.cactus_score_window(model, arr, len(tokens), start, end, context, buf, len(buf)) + if rc < 0: + raise RuntimeError(_err("Score window failed")) + return _from_json(buf) + + +# ── RAG ────────────────────────────────────────────────────────────── + + +def cactus_rag_query(model, query, top_k=5): + """Query the RAG corpus for relevant documents. + + Returns: + A dict with ranked results. + """ + buf = ctypes.create_string_buffer(65536) + rc = _lib.cactus_rag_query(model, _enc(query), buf, len(buf), top_k) + if rc < 0: + raise RuntimeError(_err("RAG query failed")) + return _from_json(buf) + + +# ── Vector index ───────────────────────────────────────────────────── + + +def cactus_index_init(index_dir, embedding_dim): + """Create a vector index for semantic search. + + Args: + index_dir: Directory to persist the index. + embedding_dim: Dimensionality of the embedding vectors. + + Returns: + An opaque index handle. Call cactus_index_destroy() when done. + """ + handle = _lib.cactus_index_init(_enc(index_dir), embedding_dim) + if not handle: + raise RuntimeError(_err("Failed to initialize index")) + return handle + + +def cactus_index_add(index, ids, documents, metadatas=None, embeddings=None): + """Add documents with embeddings to the index. + + Args: + index: Index handle from cactus_index_init(). + ids: List of integer document IDs. + documents: List of document strings. + metadatas: Optional list of metadata strings (one per document). + embeddings: List of embedding vectors (list of floats each). + """ + count = len(ids) + if len(documents) != count: + raise ValueError(f"documents length ({len(documents)}) must match ids length ({count})") + if embeddings and len(embeddings) != count: + raise ValueError(f"embeddings length ({len(embeddings)}) must match ids length ({count})") + if metadatas and len(metadatas) != count: + raise ValueError(f"metadatas length ({len(metadatas)}) must match ids length ({count})") + embedding_dim = len(embeddings[0]) if embeddings else 0 + if embeddings and any(len(e) != embedding_dim for e in embeddings): + raise ValueError("all embedding vectors must share the same dimension") + + ids_arr = (ctypes.c_int * count)(*ids) + docs_arr = (ctypes.c_char_p * count)() + for i, doc in enumerate(documents): + docs_arr[i] = _enc(doc) + + meta_arr = None + if metadatas: + meta_arr = (ctypes.c_char_p * count)() + for i, meta in enumerate(metadatas): + meta_arr[i] = _enc(meta) + + emb_ptrs = (ctypes.POINTER(ctypes.c_float) * count)() + emb_arrays = [] + for i, emb in enumerate(embeddings or []): + arr = (ctypes.c_float * len(emb))(*emb) + emb_arrays.append(arr) + emb_ptrs[i] = ctypes.cast(arr, ctypes.POINTER(ctypes.c_float)) + + rc = _lib.cactus_index_add(index, ids_arr, docs_arr, meta_arr, emb_ptrs, count, embedding_dim) + if rc < 0: + raise RuntimeError(_err("Failed to add to index")) + + +def cactus_index_delete(index, ids): + """Remove documents from the index by ID.""" + ids_arr = (ctypes.c_int * len(ids))(*ids) + rc = _lib.cactus_index_delete(index, ids_arr, len(ids)) + if rc < 0: + raise RuntimeError(_err("Failed to delete from index")) + + +def cactus_index_query(index, embedding, options=None): + """Query the index by embedding vector. + + Returns: + A dict with "results" — a list of {"id": int, "score": float}. + """ + result_capacity = 1000 + embedding_dim = len(embedding) + emb_arr = (ctypes.c_float * embedding_dim)(*embedding) + emb_ptr = ctypes.cast(emb_arr, ctypes.POINTER(ctypes.c_float)) + id_buffer = (ctypes.c_int * result_capacity)() + score_buffer = (ctypes.c_float * result_capacity)() + id_ptr = ctypes.cast(id_buffer, ctypes.POINTER(ctypes.c_int)) + score_ptr = ctypes.cast(score_buffer, ctypes.POINTER(ctypes.c_float)) + id_size = ctypes.c_size_t(result_capacity) + score_size = ctypes.c_size_t(result_capacity) + rc = _lib.cactus_index_query( + index, ctypes.pointer(emb_ptr), 1, embedding_dim, _to_json(options), + ctypes.pointer(id_ptr), ctypes.byref(id_size), + ctypes.pointer(score_ptr), ctypes.byref(score_size), + ) + if rc < 0: + raise RuntimeError(_err("Index query failed")) + n = id_size.value + return {"results": [{"id": int(id_buffer[i]), "score": float(score_buffer[i])} for i in range(n)]} + + +_INDEX_DOC_BUF_SIZE = 4096 +_INDEX_EMB_BUF_SIZE = 4096 + + +def cactus_index_get(index, ids): + """Retrieve documents by ID from the index. + + Returns: + A dict with "results" — a list of {"document", "metadata", "embedding"}. + """ + count = len(ids) + ids_arr = (ctypes.c_int * count)(*ids) + doc_raw = [ctypes.create_string_buffer(_INDEX_DOC_BUF_SIZE) for _ in range(count)] + doc_ptrs = (ctypes.c_char_p * count)() + doc_sizes = (ctypes.c_size_t * count)() + meta_raw = [ctypes.create_string_buffer(_INDEX_DOC_BUF_SIZE) for _ in range(count)] + meta_ptrs = (ctypes.c_char_p * count)() + meta_sizes = (ctypes.c_size_t * count)() + emb_raw = [(ctypes.c_float * _INDEX_EMB_BUF_SIZE)() for _ in range(count)] + emb_ptrs = (ctypes.POINTER(ctypes.c_float) * count)() + emb_sizes = (ctypes.c_size_t * count)() + for i in range(count): + doc_ptrs[i] = ctypes.cast(doc_raw[i], ctypes.c_char_p) + doc_sizes[i] = _INDEX_DOC_BUF_SIZE + meta_ptrs[i] = ctypes.cast(meta_raw[i], ctypes.c_char_p) + meta_sizes[i] = _INDEX_DOC_BUF_SIZE + emb_ptrs[i] = ctypes.cast(emb_raw[i], ctypes.POINTER(ctypes.c_float)) + emb_sizes[i] = _INDEX_EMB_BUF_SIZE + rc = _lib.cactus_index_get( + index, ids_arr, count, + doc_ptrs, doc_sizes, meta_ptrs, meta_sizes, emb_ptrs, emb_sizes, + ) + if rc < 0: + raise RuntimeError(_err("Failed to get from index")) + results = [] + for i in range(count): + doc = doc_raw[i].value.decode("utf-8", errors="ignore") + meta = meta_raw[i].value.decode("utf-8", errors="ignore") + emb = list(emb_raw[i][:emb_sizes[i]]) + results.append({"document": doc, "metadata": meta or None, "embedding": emb}) + return {"results": results} + + +def cactus_index_compact(index): + """Compact the index storage on disk to reclaim space from deletions.""" + rc = _lib.cactus_index_compact(index) + if rc < 0: + raise RuntimeError(_err("Failed to compact index")) + + +def cactus_index_destroy(index): + """Free all resources associated with an index handle. Safe to call on a + null/already-destroyed handle (no-op).""" + if not index: + return + _lib.cactus_index_destroy(index) + + +# ── Logging ────────────────────────────────────────────────────────── + + +def cactus_log_set_level(level): + """Set the log level: 0=DEBUG, 1=INFO, 2=WARN, 3=ERROR, 4=NONE.""" + _lib.cactus_log_set_level(level) + + +_log_callback_ref = None + + +def cactus_log_set_callback(callback): + """Set a log callback: callback(level, component, message). Pass None to clear.""" + global _log_callback_ref + if callback is None: + _log_callback_ref = None + _lib.cactus_log_set_callback(LogCallback(), None) + return + + def _bridge(level, component_bytes, message_bytes, _): + callback( + level, + component_bytes.decode("utf-8", errors="ignore") if component_bytes else "", + message_bytes.decode("utf-8", errors="ignore") if message_bytes else "", + ) + + _log_callback_ref = LogCallback(_bridge) + _lib.cactus_log_set_callback(_log_callback_ref, None) + + +class Graph: + INT8 = 0 + FP16 = 1 + FP32 = 2 + CQ1 = 3 + CQ2 = 4 + CQ3 = 5 + CQ4 = 6 + CPU = 0 + METAL = 1 + ACT_SILU = 0 + ACT_GELU = 1 + ACT_GELU_ERF = 2 + ACT_RELU = 3 + ACT_SIGMOID = 4 + ACT_TANH = 5 + + def __init__(self): + self.h = _lib.cactus_graph_create() + if not self.h: + raise RuntimeError(_err("cactus_graph_create failed")) + + def save(self, filename): + rc = _lib.cactus_graph_save(self.h, str(filename).encode()) + if rc != 0: + raise RuntimeError(_err("graph_save failed")) + + @classmethod + def load(cls, filename): + h = _lib.cactus_graph_load(str(filename).encode()) + if not h: + raise RuntimeError(_err("cactus_graph_load failed")) + obj = cls.__new__(cls) + obj.h = h + return obj + + def __del__(self): + h = getattr(self, "h", None) + if h: + _lib.cactus_graph_destroy(h) + self.h = None + + def input(self, shape, dtype=FP16, dynamic_dims=None): + shape = tuple(int(x) for x in shape) + arr = (ctypes.c_size_t * len(shape))(*shape) + out = cactus_node_t() + rc = _lib.cactus_graph_input(self.h, arr, len(shape), int(dtype), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_input failed")) + tensor = self._tensor_from_node(out.value) + if dynamic_dims is not None: + mask = (ctypes.c_uint8 * len(dynamic_dims))(*[1 if int(d) else 0 for d in dynamic_dims]) + rc = _lib.cactus_graph_set_input_dynamic_dims(self.h, cactus_node_t(tensor.id), mask, len(dynamic_dims)) + if rc != 0: + raise RuntimeError(_err("graph_set_input_dynamic_dims failed")) + return tensor + + def set_input(self, tensor, data, dtype=None): + if not isinstance(tensor, Tensor): + raise TypeError("tensor must be a Tensor") + if tensor.g is not self: + raise ValueError("tensor belongs to a different graph") + target_dtype = int(tensor.dtype if dtype is None else dtype) + arr = self._coerce_input_array(data, target_dtype) + info = self._get_output_info(tensor.id) + expected_shape = tuple(int(x) for x in info["shape"]) + expected_num_elements = int(info["num_elements"]) + expected_byte_size = int(info["byte_size"]) + if tuple(int(x) for x in arr.shape) != expected_shape: + raise ValueError( + "graph input shape mismatch for node " + f"{tensor.id}: expected {expected_shape}, got {tuple(int(x) for x in arr.shape)}" + ) + if int(arr.size) != expected_num_elements: + raise ValueError( + "graph input element-count mismatch for node " + f"{tensor.id}: expected {expected_num_elements}, got {int(arr.size)}" + ) + if int(arr.nbytes) != expected_byte_size: + raise ValueError( + "graph input byte-size mismatch for node " + f"{tensor.id}: expected {expected_byte_size}, got {int(arr.nbytes)}" + ) + rc = _lib.cactus_graph_set_input( + self.h, + cactus_node_t(tensor.id), + arr.ctypes.data_as(ctypes.c_void_p), + target_dtype, + ) + if rc != 0: + raise RuntimeError(_err("graph_set_input failed")) + + def set_external_input(self, tensor, data_ptr, dtype=None): + if not isinstance(tensor, Tensor): + raise TypeError("tensor must be a Tensor") + if tensor.g is not self: + raise ValueError("tensor belongs to a different graph") + target_dtype = int(tensor.dtype if dtype is None else dtype) + ptr = ctypes.c_void_p(data_ptr if isinstance(data_ptr, int) else int(data_ptr)) + rc = _lib.cactus_graph_set_external_input( + self.h, + cactus_node_t(tensor.id), + ptr, + target_dtype, + ) + if rc != 0: + raise RuntimeError(_err("graph_set_external_input failed")) + + def set_runtime_input_shape(self, tensor, shape): + if not isinstance(tensor, Tensor): + raise TypeError("tensor must be a Tensor") + if tensor.g is not self: + raise ValueError("tensor belongs to a different graph") + shape = tuple(int(x) for x in shape) + arr = (ctypes.c_size_t * len(shape))(*shape) + rc = _lib.cactus_graph_set_runtime_input_shape(self.h, cactus_node_t(tensor.id), arr, len(shape)) + if rc != 0: + raise RuntimeError(_err("graph_set_runtime_input_shape failed")) + + def mark_embedded_input(self, tensor): + if not isinstance(tensor, Tensor): + raise TypeError("tensor must be a Tensor") + if tensor.g is not self: + raise ValueError("tensor belongs to a different graph") + rc = _lib.cactus_graph_mark_embedded_input(self.h, cactus_node_t(tensor.id)) + if rc != 0: + raise RuntimeError(_err("graph_mark_embedded_input failed")) + + def hard_reset(self): + rc = _lib.cactus_graph_hard_reset(self.h) + if rc != 0: + raise RuntimeError(_err("graph_hard_reset failed")) + + def execute(self): + rc = _lib.cactus_graph_execute(self.h) + if rc != 0: + raise RuntimeError(_err("graph_execute failed")) + + def add(self, a, b, backend=None): + return self._binary("cactus_graph_add", a, b, backend=backend) + + def add_clipped(self, a, b, backend=None): + return self._binary("cactus_graph_add_clipped", a, b, backend=backend) + + def subtract(self, a, b, backend=None): + return self._binary("cactus_graph_subtract", a, b, backend=backend) + + def multiply(self, a, b, backend=None): + return self._binary("cactus_graph_multiply", a, b, backend=backend) + + def divide(self, a, b, backend=None): + return self._binary("cactus_graph_divide", a, b, backend=backend) + + def not_equal(self, a, b, backend=None): + return self._binary("cactus_graph_not_equal", a, b, backend=backend) + + def abs(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_abs(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_abs failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def pow(self, x, exponent, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_pow(self.h, cactus_node_t(x.id), ctypes.c_float(float(exponent)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_pow failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def precision_cast(self, x, dtype, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_precision_cast(self.h, cactus_node_t(x.id), int(dtype), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_precision_cast failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def quantize_activations(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_quantize_activations(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_quantize_activations failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def _scalar(self, fn_name, x, value=None, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + fn = getattr(_lib, fn_name) + if value is None: + rc = fn(self.h, cactus_node_t(x.id), ctypes.byref(out)) + else: + rc = fn(self.h, cactus_node_t(x.id), ctypes.c_float(float(value)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(f"{fn_name} failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def scalar_add(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_add", x, value, backend=backend) + + def scalar_subtract(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_subtract", x, value, backend=backend) + + def scalar_multiply(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_multiply", x, value, backend=backend) + + def scalar_divide(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_divide", x, value, backend=backend) + + def scalar_floor_divide(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_floor_divide", x, value, backend=backend) + + def scalar_not_equal(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_not_equal", x, value, backend=backend) + + def scalar_exp(self, x, backend=None): + return self._scalar("cactus_graph_scalar_exp", x, backend=backend) + + def scalar_sqrt(self, x, backend=None): + return self._scalar("cactus_graph_scalar_sqrt", x, backend=backend) + + def scalar_cos(self, x, backend=None): + return self._scalar("cactus_graph_scalar_cos", x, backend=backend) + + def scalar_sin(self, x, backend=None): + return self._scalar("cactus_graph_scalar_sin", x, backend=backend) + + def scalar_log(self, x, backend=None): + return self._scalar("cactus_graph_scalar_log", x, backend=backend) + + def clamp(self, x, lo, hi, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_clamp( + self.h, + cactus_node_t(x.id), + ctypes.c_float(float(lo)), + ctypes.c_float(float(hi)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_clamp failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def masked_select_prefix(self, x, mask, backend=None): + return self._binary("cactus_graph_masked_select_prefix", x, mask, backend=backend) + + def masked_scatter(self, x, mask, source, backend=None): + x = self._ensure_tensor(x) + mask = self._ensure_tensor(mask) + source = self._ensure_tensor(source) + out = cactus_node_t() + rc = _lib.cactus_graph_masked_scatter( + self.h, + cactus_node_t(x.id), + cactus_node_t(mask.id), + cactus_node_t(source.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_masked_scatter failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def view(self, x, shape, backend=None): + x = self._ensure_tensor(x) + shape = tuple(int(v) for v in shape) + arr = (ctypes.c_size_t * len(shape))(*shape) + out = cactus_node_t() + rc = _lib.cactus_graph_view(self.h, cactus_node_t(x.id), arr, len(shape), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_view failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def reshape(self, x, shape, backend=None): + x = self._ensure_tensor(x) + shape = tuple(int(v) for v in shape) + arr = (ctypes.c_size_t * len(shape))(*shape) + out = cactus_node_t() + rc = _lib.cactus_graph_reshape(self.h, cactus_node_t(x.id), arr, len(shape), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_reshape failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def expand(self, x, shape, backend=None): + x = self._ensure_tensor(x) + shape = tuple(int(v) for v in shape) + arr = (ctypes.c_size_t * len(shape))(*shape) + out = cactus_node_t() + rc = _lib.cactus_graph_expand(self.h, cactus_node_t(x.id), arr, len(shape), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_expand failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def flatten(self, x, start_dim=0, end_dim=-1, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_flatten( + self.h, + cactus_node_t(x.id), + ctypes.c_int32(int(start_dim)), + ctypes.c_int32(int(end_dim)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_flatten failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def slice(self, x, axis, start, length=0, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_slice( + self.h, + cactus_node_t(x.id), + ctypes.c_int32(int(axis)), + ctypes.c_size_t(int(start)), + ctypes.c_size_t(int(length)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_slice failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def index(self, x, index_value, axis=0, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_index( + self.h, + cactus_node_t(x.id), + ctypes.c_size_t(int(index_value)), + ctypes.c_int32(int(axis)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_index failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def transpose(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_transpose( + self.h, + cactus_node_t(x.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_transpose failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def permute(self, x, permutation, backend=None): + x = self._ensure_tensor(x) + permutation = tuple(int(v) for v in permutation) + arr = (ctypes.c_size_t * len(permutation))(*permutation) + out = cactus_node_t() + rc = _lib.cactus_graph_transpose_n( + self.h, + cactus_node_t(x.id), + arr, + len(permutation), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_transpose_n failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None): + a = self._ensure_tensor(a) + b = self._ensure_tensor(b) + out = cactus_node_t() + rc = _lib.cactus_graph_matmul( + self.h, + cactus_node_t(a.id), + cactus_node_t(b.id), + ctypes.c_bool(bool(pretransposed_rhs)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_matmul failed")) + result = self._apply_backend(self._tensor_from_node(out.value), backend) + if output_dtype is not None and int(output_dtype) != int(result.dtype): + result = self.precision_cast(result, int(output_dtype)) + return result + + def gather(self, tensor, indices, backend=None): + tensor = self._ensure_tensor(tensor) + indices = self._ensure_tensor(indices) + out = cactus_node_t() + rc = _lib.cactus_graph_gather( + self.h, + cactus_node_t(tensor.id), + cactus_node_t(indices.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_gather failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def embedding_from_tensor(self, embedding_tensor, indices, backend=None): + embedding_tensor = self._ensure_tensor(embedding_tensor) + indices = self._ensure_tensor(indices) + out = cactus_node_t() + rc = _lib.cactus_graph_embedding_from_tensor( + self.h, cactus_node_t(embedding_tensor.id), cactus_node_t(indices.id), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_embedding_from_tensor failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def embedding_from_file(self, filename, indices, backend=None): + indices = self._ensure_tensor(indices) + out = cactus_node_t() + rc = _lib.cactus_graph_embedding_from_file(self.h, str(filename).encode(), cactus_node_t(indices.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_embedding_from_file failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def mmap_embeddings(self, filename): + out = cactus_node_t() + rc = _lib.cactus_graph_mmap_embeddings(self.h, str(filename).encode(), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_mmap_embeddings failed")) + return self._tensor_from_node(out.value) + + def mmap_weights(self, filename): + out = cactus_node_t() + rc = _lib.cactus_graph_mmap_weights(self.h, str(filename).encode(), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_mmap_weights failed")) + return self._tensor_from_node(out.value) + + def bind_mmap_weights(self, tensor, filename): + tensor = self._ensure_tensor(tensor) + rc = _lib.cactus_graph_bind_mmap_weights(self.h, cactus_node_t(tensor.id), str(filename).encode()) + if rc != 0: + raise RuntimeError(_err("graph_bind_mmap_weights failed")) + + def bilinear_interpolation(self, pos_embeds, dst_height, dst_width, backend=None): + pos_embeds = self._ensure_tensor(pos_embeds) + out = cactus_node_t() + rc = _lib.cactus_graph_bilinear_interpolation( + self.h, cactus_node_t(pos_embeds.id), ctypes.c_size_t(int(dst_height)), ctypes.c_size_t(int(dst_width)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_bilinear_interpolation failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def set_grouped_scales(self, tensor, group_size, num_groups, scales): + tensor = self._ensure_tensor(tensor) + arr = np.ascontiguousarray(scales, dtype=np.float16) + rc = _lib.cactus_graph_set_grouped_scales( + self.h, + cactus_node_t(tensor.id), + ctypes.c_size_t(int(group_size)), + ctypes.c_size_t(int(num_groups)), + arr.ctypes.data_as(ctypes.c_void_p), + ) + if rc != 0: + raise RuntimeError(_err("graph_set_grouped_scales failed")) + + def set_interleaved(self, tensor, interleaved=True, original_n=0): + tensor = self._ensure_tensor(tensor) + rc = _lib.cactus_graph_set_interleaved( + self.h, cactus_node_t(tensor.id), ctypes.c_bool(bool(interleaved)), ctypes.c_size_t(int(original_n)) + ) + if rc != 0: + raise RuntimeError(_err("graph_set_interleaved failed")) + + def release_weight_pages(self, tensor): + tensor = self._ensure_tensor(tensor) + rc = _lib.cactus_graph_release_weight_pages(self.h, cactus_node_t(tensor.id)) + if rc != 0: + raise RuntimeError(_err("graph_release_weight_pages failed")) + + def prefetch_weight_pages(self, tensor): + tensor = self._ensure_tensor(tensor) + rc = _lib.cactus_graph_prefetch_weight_pages(self.h, cactus_node_t(tensor.id)) + if rc != 0: + raise RuntimeError(_err("graph_prefetch_weight_pages failed")) + + def release_all_weight_pages(self): + rc = _lib.cactus_graph_release_all_weight_pages(self.h) + if rc != 0: + raise RuntimeError(_err("graph_release_all_weight_pages failed")) + + def concat(self, a, b, axis=0, backend=None): + a = self._ensure_tensor(a) + b = self._ensure_tensor(b) + out = cactus_node_t() + rc = _lib.cactus_graph_concat( + self.h, + cactus_node_t(a.id), + cactus_node_t(b.id), + ctypes.c_int32(int(axis)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_concat failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def cat(self, tensors, axis=0, backend=None): + tensors = [self._ensure_tensor(t) for t in tensors] + if not tensors: + raise ValueError("cat requires at least one tensor") + if len(tensors) == 1: + return tensors[0] + ids = (cactus_node_t * len(tensors))(*(cactus_node_t(t.id) for t in tensors)) + out = cactus_node_t() + rc = _lib.cactus_graph_cat( + self.h, ids, ctypes.c_size_t(len(tensors)), + ctypes.c_int32(int(axis)), ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_cat failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def groupnorm(self, x, weight, bias, num_groups, eps=1e-5, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + bias = self._ensure_tensor(bias) + out = cactus_node_t() + rc = _lib.cactus_graph_groupnorm( + self.h, + cactus_node_t(x.id), + cactus_node_t(weight.id), + cactus_node_t(bias.id), + ctypes.c_size_t(int(num_groups)), + ctypes.c_float(float(eps)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_group_norm failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def group_norm(self, x, weight, bias, num_groups, eps=1e-5, backend=None): + return self.groupnorm(x, weight, bias, num_groups, eps=eps, backend=backend) + + def layernorm(self, x, weight, bias=None, eps=1e-5, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + has_bias = bias is not None + bias_node = cactus_node_t(0) + if has_bias: + bias = self._ensure_tensor(bias) + bias_node = cactus_node_t(bias.id) + out = cactus_node_t() + rc = _lib.cactus_graph_layernorm( + self.h, + cactus_node_t(x.id), + cactus_node_t(weight.id), + bias_node, + ctypes.c_float(float(eps)), + ctypes.c_bool(has_bias), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_layer_norm failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def layer_norm(self, x, weight, bias=None, eps=1e-5, backend=None): + return self.layernorm(x, weight, bias=bias, eps=eps, backend=backend) + + def batchnorm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + bias = self._ensure_tensor(bias) + running_mean = self._ensure_tensor(running_mean) + running_var = self._ensure_tensor(running_var) + out = cactus_node_t() + rc = _lib.cactus_graph_batchnorm( + self.h, + cactus_node_t(x.id), + cactus_node_t(weight.id), + cactus_node_t(bias.id), + cactus_node_t(running_mean.id), + cactus_node_t(running_var.id), + ctypes.c_int32(int(axis)), + ctypes.c_float(float(eps)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_batchnorm failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def batch_norm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5, backend=None): + return self.batchnorm(x, weight, bias, running_mean, running_var, axis=axis, eps=eps, backend=backend) + + def rms_norm(self, x, weight, eps=1e-5, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + out = cactus_node_t() + rc = _lib.cactus_graph_rms_norm( + self.h, + cactus_node_t(x.id), + cactus_node_t(weight.id), + ctypes.c_float(float(eps)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_rms_norm failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def topk(self, x, k, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_topk(self.h, cactus_node_t(x.id), ctypes.c_size_t(int(k)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_topk failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def rope(self, x, theta, position_offset=0, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_rope( + self.h, cactus_node_t(x.id), ctypes.c_float(float(theta)), ctypes.c_size_t(int(position_offset)), + ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_rope failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def rope_gptj(self, x, theta, position_offset=0, rot_dim=0, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_rope_gptj( + self.h, cactus_node_t(x.id), ctypes.c_float(float(theta)), ctypes.c_size_t(int(position_offset)), + ctypes.c_size_t(int(rot_dim)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_rope_gptj failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def _reduce(self, fn_name, x, axis, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = getattr(_lib, fn_name)(self.h, cactus_node_t(x.id), ctypes.c_int32(int(axis)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(f"{fn_name} failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def sum(self, x, axis, backend=None): + return self._reduce("cactus_graph_sum", x, axis, backend=backend) + + def mean(self, x, axis, backend=None): + return self._reduce("cactus_graph_mean", x, axis, backend=backend) + + def variance(self, x, axis, backend=None): + return self._reduce("cactus_graph_variance", x, axis, backend=backend) + + def min(self, x, axis, backend=None): + return self._reduce("cactus_graph_min", x, axis, backend=backend) + + def max(self, x, axis, backend=None): + return self._reduce("cactus_graph_max", x, axis, backend=backend) + + def cumsum(self, x, axis, backend=None): + return self._reduce("cactus_graph_cumsum", x, axis, backend=backend) + + def softmax(self, x, axis=-1, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_softmax( + self.h, + cactus_node_t(x.id), + ctypes.c_int32(int(axis)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_softmax failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def attention(self, query, key, value, scale, is_causal=True, position_offset=0, window_size=0, + backend=None, mask=None, additive_mask=False): + query = self._ensure_tensor(query) + key = self._ensure_tensor(key) + value = self._ensure_tensor(value) + mask_node = cactus_node_t(0) + use_mask = mask is not None + if use_mask: + mask_node = cactus_node_t(self._ensure_tensor(mask).id) + out = cactus_node_t() + rc = _lib.cactus_graph_attention( + self.h, + cactus_node_t(query.id), + cactus_node_t(key.id), + cactus_node_t(value.id), + ctypes.c_float(float(scale)), + ctypes.c_bool(bool(is_causal)), + ctypes.c_size_t(int(position_offset)), + ctypes.c_size_t(int(window_size)), + ctypes.c_bool(use_mask), + mask_node, + ctypes.c_bool(bool(additive_mask)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_attention failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def kv_cache_state(self, max_seq_len, num_kv_heads, head_dim, window_size=0, sink_size=0, num_slots=1, backend=None): + out = cactus_node_t() + rc = _lib.cactus_graph_kv_cache_state( + self.h, + ctypes.c_size_t(int(max_seq_len)), + ctypes.c_size_t(int(num_kv_heads)), + ctypes.c_size_t(int(head_dim)), + ctypes.c_size_t(int(window_size)), + ctypes.c_size_t(int(sink_size)), + ctypes.c_size_t(int(num_slots)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_kv_cache_state failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def kv_cache_append(self, new_kv, cache_state, window_size=0, sink_size=0, backend=None): + new_kv = self._ensure_tensor(new_kv) + cache_state = self._ensure_tensor(cache_state) + out = cactus_node_t() + rc = _lib.cactus_graph_kv_cache_append( + self.h, + cactus_node_t(new_kv.id), + cactus_node_t(cache_state.id), + ctypes.c_size_t(int(window_size)), + ctypes.c_size_t(int(sink_size)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_kv_cache_append failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def attention_cached( + self, + query, + key_new, + value_new, + k_cache_state, + v_cache_state, + scale, + position_offset=0, + window_size=0, + v_head_dim=0, + backend=None, + ): + query = self._ensure_tensor(query) + key_new = self._ensure_tensor(key_new) + value_new = self._ensure_tensor(value_new) + k_cache_state = self._ensure_tensor(k_cache_state) + v_cache_state = self._ensure_tensor(v_cache_state) + out = cactus_node_t() + rc = _lib.cactus_graph_attention_cached( + self.h, + cactus_node_t(query.id), + cactus_node_t(key_new.id), + cactus_node_t(value_new.id), + cactus_node_t(k_cache_state.id), + cactus_node_t(v_cache_state.id), + ctypes.c_float(float(scale)), + ctypes.c_size_t(int(position_offset)), + ctypes.c_size_t(int(window_size)), + ctypes.c_size_t(int(v_head_dim)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_attention_cached failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv_cache_state(self, window_size, hidden_dim, backend=None): + out = cactus_node_t() + rc = _lib.cactus_graph_conv_cache_state( + self.h, + ctypes.c_size_t(int(window_size)), + ctypes.c_size_t(int(hidden_dim)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_conv_cache_state failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv_cache_append(self, new_data, cache_state, backend=None): + new_data = self._ensure_tensor(new_data) + cache_state = self._ensure_tensor(cache_state) + out = cactus_node_t() + rc = _lib.cactus_graph_conv_cache_append( + self.h, + cactus_node_t(new_data.id), + cactus_node_t(cache_state.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_conv_cache_append failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv_cache_initialize(self, rows, cache_state, backend=None): + rows = self._ensure_tensor(rows) + cache_state = self._ensure_tensor(cache_state) + out = cactus_node_t() + rc = _lib.cactus_graph_conv_cache_initialize( + self.h, + cactus_node_t(rows.id), + cactus_node_t(cache_state.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_conv_cache_initialize failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def recurrent_cache_state(self, shape, dtype=FP16, backend=None): + shape = tuple(int(x) for x in shape) + arr = (ctypes.c_size_t * len(shape))(*shape) + out = cactus_node_t() + rc = _lib.cactus_graph_recurrent_cache_state( + self.h, + arr, + ctypes.c_size_t(len(shape)), + ctypes.c_int(int(dtype)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_recurrent_cache_state failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def recurrent_cache_write(self, new_value, cache_input, backend=None): + new_value = self._ensure_tensor(new_value) + cache_input = self._ensure_tensor(cache_input) + out = cactus_node_t() + rc = _lib.cactus_graph_recurrent_cache_write( + self.h, + cactus_node_t(new_value.id), + cactus_node_t(cache_input.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_recurrent_cache_write failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def rel_pos_bias(self, query, relative_key, scale, backend=None): + query = self._ensure_tensor(query) + relative_key = self._ensure_tensor(relative_key) + out = cactus_node_t() + rc = _lib.cactus_graph_rel_pos_bias( + self.h, cactus_node_t(query.id), cactus_node_t(relative_key.id), ctypes.c_float(float(scale)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_rel_pos_bias failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def attention_int8_hybrid(self, query, key_new, value_new, scale, position_offset, + cached_keys, cached_values, k_scales, v_scales, + cache_len, num_kv_heads, head_dim, window_size=0, backend=None): + query = self._ensure_tensor(query) + key_new = self._ensure_tensor(key_new) + value_new = self._ensure_tensor(value_new) + ck = np.ascontiguousarray(cached_keys, dtype=np.int8) + cv = np.ascontiguousarray(cached_values, dtype=np.int8) + ks = np.ascontiguousarray(k_scales, dtype=np.float32) + vs = np.ascontiguousarray(v_scales, dtype=np.float32) + out = cactus_node_t() + rc = _lib.cactus_graph_attention_int8_hybrid( + self.h, + cactus_node_t(query.id), + cactus_node_t(key_new.id), + cactus_node_t(value_new.id), + ctypes.c_float(float(scale)), + ctypes.c_size_t(int(position_offset)), + ck.ctypes.data_as(ctypes.POINTER(ctypes.c_int8)), + cv.ctypes.data_as(ctypes.POINTER(ctypes.c_int8)), + ks.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + vs.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + ctypes.c_size_t(int(cache_len)), + ctypes.c_size_t(int(num_kv_heads)), + ctypes.c_size_t(int(head_dim)), + ctypes.c_size_t(int(window_size)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError("graph_attention_int8_hybrid failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def relu(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_relu(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_relu failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def silu(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_silu(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_silu failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def gelu(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_gelu(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_gelu failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def gelu_erf(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_gelu_erf(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_gelu_erf failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def sigmoid(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_sigmoid(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_sigmoid failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def tanh(self, x, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_tanh(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_tanh failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def glu(self, x, axis=-1, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_glu(self.h, cactus_node_t(x.id), ctypes.c_int32(int(axis)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_glu failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv1d_causal(self, x, weight, kernel_size, dilation, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + out = cactus_node_t() + rc = _lib.cactus_graph_conv1d_causal( + self.h, cactus_node_t(x.id), cactus_node_t(weight.id), + ctypes.c_size_t(int(kernel_size)), ctypes.c_size_t(int(dilation)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError("graph_conv1d_causal failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv1d_k3(self, x, weight, stride=1, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + out = cactus_node_t() + rc = _lib.cactus_graph_conv1d_k3( + self.h, cactus_node_t(x.id), cactus_node_t(weight.id), ctypes.c_size_t(int(stride)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError("graph_conv1d_k3 failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv1d_k7s3(self, x, weight, bias, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + bias = self._ensure_tensor(bias) + out = cactus_node_t() + rc = _lib.cactus_graph_conv1d_k7s3( + self.h, cactus_node_t(x.id), cactus_node_t(weight.id), cactus_node_t(bias.id), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError("graph_conv1d_k7s3 failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv1d(self, x, weight, bias=None, stride=1, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv1d", x, weight, bias, ctypes.c_size_t(int(stride)), backend=backend) + + def conv1d_same_depthwise_k9(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv1d_same_depthwise_k9", x, weight, bias, backend=backend) + + def conv1d_pointwise(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv1d_pointwise", x, weight, bias, backend=backend) + + def conv2d_k3s2p1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_k3s2p1", x, weight, bias, backend=backend) + + def conv2d_depthwise_k3s2p1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_depthwise_k3s2p1", x, weight, bias, backend=backend) + + def conv2d_pointwise_1x1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_pointwise_1x1", x, weight, bias, backend=backend) + + def conv2d_k3s1p1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_k3s1p1", x, weight, bias, backend=backend) + + def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, backend=None): + return self._conv_with_optional_bias( + "cactus_graph_conv2d", + x, + weight, + bias, + ctypes.c_size_t(int(stride)), + ctypes.c_size_t(int(padding)), + ctypes.c_size_t(int(dilation)), + ctypes.c_size_t(int(groups)), + backend=backend, + ) + + def _conv_with_optional_bias(self, fn_name, x, weight, bias=None, *extra, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + has_bias = bias is not None + bias_node = cactus_node_t(0 if bias is None else self._ensure_tensor(bias).id) + out = cactus_node_t() + rc = getattr(_lib, fn_name)( + self.h, cactus_node_t(x.id), cactus_node_t(weight.id), ctypes.c_bool(has_bias), bias_node, *extra, ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(f"{fn_name} failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def lstm_cell(self, input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh, backend=None): + tensors = [self._ensure_tensor(t) for t in (input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh)] + out = cactus_node_t() + rc = _lib.cactus_graph_lstm_cell(self.h, *(cactus_node_t(t.id) for t in tensors), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_lstm_cell failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def gated_deltanet_decode(self, query, key, value, gate_log, beta, initial_state, scale, backend=None): + tensors = [self._ensure_tensor(t) for t in (query, key, value, gate_log, beta, initial_state)] + out = cactus_node_t() + rc = _lib.cactus_graph_gated_deltanet_decode( + self.h, *(cactus_node_t(t.id) for t in tensors), ctypes.c_float(float(scale)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_gated_deltanet_decode failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def gated_deltanet_prefill(self, query, key, value, gate_log, beta, initial_state, chunk_size, scale, backend=None): + tensors = [self._ensure_tensor(t) for t in (query, key, value, gate_log, beta, initial_state)] + out = cactus_node_t() + rc = _lib.cactus_graph_gated_deltanet_prefill( + self.h, *(cactus_node_t(t.id) for t in tensors), ctypes.c_size_t(int(chunk_size)), ctypes.c_float(float(scale)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_gated_deltanet_prefill failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def stft(self, x, weight, stride, num_fft_bins, backend=None): + x = self._ensure_tensor(x) + weight = self._ensure_tensor(weight) + out = cactus_node_t() + rc = _lib.cactus_graph_stft( + self.h, cactus_node_t(x.id), cactus_node_t(weight.id), + ctypes.c_size_t(int(stride)), ctypes.c_size_t(int(num_fft_bins)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_stft failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def altup_predict(self, coefs, streams, backend=None): + coefs = self._ensure_tensor(coefs) + streams = [self._ensure_tensor(t) for t in streams] + ids = (cactus_node_t * len(streams))(*(cactus_node_t(t.id) for t in streams)) + out = cactus_node_t() + rc = _lib.cactus_graph_altup_predict(self.h, cactus_node_t(coefs.id), ids, ctypes.c_size_t(len(streams)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_altup_predict failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def altup_correct(self, coefs, innovation, predictions, backend=None): + coefs = self._ensure_tensor(coefs) + innovation = self._ensure_tensor(innovation) + predictions = [self._ensure_tensor(t) for t in predictions] + ids = (cactus_node_t * len(predictions))(*(cactus_node_t(t.id) for t in predictions)) + out = cactus_node_t() + rc = _lib.cactus_graph_altup_correct( + self.h, cactus_node_t(coefs.id), cactus_node_t(innovation.id), ids, ctypes.c_size_t(len(predictions)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_altup_correct failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def gaussian_topk(self, x, ppf, backend=None): + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_gaussian_topk(self.h, cactus_node_t(x.id), ctypes.c_float(float(ppf)), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_gaussian_topk failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def moe_layer_gated(self, hidden, routing_probs, topk_indices, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, routed_scaling_factor=1.0, activation=0, backend=None): + hidden = self._ensure_tensor(hidden) + routing_probs = self._ensure_tensor(routing_probs) + topk_indices = self._ensure_tensor(topk_indices) + w1 = (cactus_node_t * len(w1_weights))(*(cactus_node_t(self._ensure_tensor(t).id) for t in w1_weights)) + w3 = (cactus_node_t * len(w3_weights))(*(cactus_node_t(self._ensure_tensor(t).id) for t in w3_weights)) + w2 = (cactus_node_t * len(w2_weights))(*(cactus_node_t(self._ensure_tensor(t).id) for t in w2_weights)) + out = cactus_node_t() + rc = _lib.cactus_graph_moe_layer_gated( + self.h, cactus_node_t(hidden.id), cactus_node_t(routing_probs.id), cactus_node_t(topk_indices.id), + w1, w3, w2, ctypes.c_size_t(int(num_experts)), ctypes.c_size_t(int(num_experts_per_tok)), + ctypes.c_bool(bool(normalize_routing)), ctypes.c_float(float(epsilon)), + ctypes.c_float(float(routed_scaling_factor)), ctypes.c_int32(int(activation)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_moe_layer_gated failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def dense_mlp_tq_fused(self, hidden, gate_weight, up_weight, down_weight, product_scale=1.0, backend=None): + hidden = self._ensure_tensor(hidden) + gate_weight = self._ensure_tensor(gate_weight) + up_weight = self._ensure_tensor(up_weight) + down_weight = self._ensure_tensor(down_weight) + out = cactus_node_t() + rc = _lib.cactus_graph_dense_mlp_tq_fused( + self.h, + cactus_node_t(hidden.id), + cactus_node_t(gate_weight.id), + cactus_node_t(up_weight.id), + cactus_node_t(down_weight.id), + ctypes.c_float(float(product_scale)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_dense_mlp_tq_fused failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def moe_layer_ungated(self, hidden, routing_probs, topk_indices, w1_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, + routed_scaling_factor=1.0, activation=ACT_GELU, backend=None): + hidden = self._ensure_tensor(hidden) + routing_probs = self._ensure_tensor(routing_probs) + topk_indices = self._ensure_tensor(topk_indices) + w1 = (cactus_node_t * len(w1_weights))(*(cactus_node_t(self._ensure_tensor(t).id) for t in w1_weights)) + w2 = (cactus_node_t * len(w2_weights))(*(cactus_node_t(self._ensure_tensor(t).id) for t in w2_weights)) + out = cactus_node_t() + rc = _lib.cactus_graph_moe_layer_ungated( + self.h, cactus_node_t(hidden.id), cactus_node_t(routing_probs.id), cactus_node_t(topk_indices.id), + w1, w2, ctypes.c_size_t(int(num_experts)), ctypes.c_size_t(int(num_experts_per_tok)), + ctypes.c_bool(bool(normalize_routing)), ctypes.c_float(float(epsilon)), + ctypes.c_float(float(routed_scaling_factor)), ctypes.c_int32(int(activation)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_moe_layer_ungated failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def sample(self, logits, temperature=0.6, top_p=0.95, top_k=20, backend=None): + logits = self._ensure_tensor(logits) + out = cactus_node_t() + rc = _lib.cactus_graph_sample( + self.h, cactus_node_t(logits.id), ctypes.c_float(float(temperature)), + ctypes.c_float(float(top_p)), ctypes.c_size_t(int(top_k)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_sample failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def scatter_topk(self, indices, values, num_classes, backend=None): + indices = self._ensure_tensor(indices) + values = self._ensure_tensor(values) + out = cactus_node_t() + rc = _lib.cactus_graph_scatter_topk( + self.h, cactus_node_t(indices.id), cactus_node_t(values.id), ctypes.c_size_t(int(num_classes)), ctypes.byref(out) + ) + if rc != 0: + raise RuntimeError(_err("graph_scatter_topk failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def persistent(self, source_node, backend=None): + source_node = self._ensure_tensor(source_node) + out = cactus_node_t() + rc = _lib.cactus_graph_persistent(self.h, cactus_node_t(source_node.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_persistent failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def is_populated(self, persistent_node): + persistent_node = self._ensure_tensor(persistent_node) + out_is_populated = ctypes.c_int32() + rc = _lib.cactus_graph_is_populated(self.h, cactus_node_t(persistent_node.id), ctypes.byref(out_is_populated)) + if rc != 0: + raise RuntimeError(_err("graph_is_populated failed")) + return bool(out_is_populated.value) + + def invalidate_persistent(self, persistent_node): + persistent_node = self._ensure_tensor(persistent_node) + rc = _lib.cactus_graph_invalidate_persistent(self.h, cactus_node_t(persistent_node.id)) + if rc != 0: + raise RuntimeError(_err("graph_invalidate_persistent failed")) + + # ── Audio / signal processing ─────────────────────────────────── + + def rfft(self, x, backend=None): + """Real-to-complex FFT.""" + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_rfft(self.h, cactus_node_t(x.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_rfft failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def irfft(self, x, output_length, backend=None): + """Complex-to-real inverse FFT.""" + x = self._ensure_tensor(x) + out = cactus_node_t() + rc = _lib.cactus_graph_irfft( + self.h, cactus_node_t(x.id), + ctypes.c_size_t(int(output_length)), ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_irfft failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def mel_filter_bank(self, num_frequency_bins, num_mel_filters, + min_frequency, max_frequency, sampling_rate, + norm_type=0, scale_type=0, backend=None): + """Generate a mel-scale filter bank tensor.""" + out = cactus_node_t() + rc = _lib.cactus_graph_mel_filter_bank( + self.h, + ctypes.c_size_t(int(num_frequency_bins)), + ctypes.c_size_t(int(num_mel_filters)), + ctypes.c_float(float(min_frequency)), + ctypes.c_float(float(max_frequency)), + ctypes.c_size_t(int(sampling_rate)), + ctypes.c_int(int(norm_type)), + ctypes.c_int(int(scale_type)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_mel_filter_bank failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def spectrogram(self, waveform, mel_filters, frame_length, hop_length, + fft_length, power=2.0, center=True, pad_mode=0, + mel_floor=1e-10, log_mel_mode=0, dither=0.0, + preemphasis=0.0, remove_dc_offset=False, backend=None): + """Compute a spectrogram from a waveform tensor.""" + waveform = self._ensure_tensor(waveform) + mel_filters = self._ensure_tensor(mel_filters) + out = cactus_node_t() + rc = _lib.cactus_graph_spectrogram( + self.h, + cactus_node_t(waveform.id), + cactus_node_t(mel_filters.id), + ctypes.c_size_t(int(frame_length)), + ctypes.c_size_t(int(hop_length)), + ctypes.c_size_t(int(fft_length)), + ctypes.c_float(float(power)), + ctypes.c_bool(bool(center)), + ctypes.c_int(int(pad_mode)), + ctypes.c_float(float(mel_floor)), + ctypes.c_int(int(log_mel_mode)), + ctypes.c_float(float(dither)), + ctypes.c_float(float(preemphasis)), + ctypes.c_bool(bool(remove_dc_offset)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_spectrogram failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + # ── Image processing ───────────────────────────────────────────── + + def image_preprocess(self, pixel_input, src_width, src_height, + target_width, target_height, patch_size, channels, + rescale_factor, mean, std_dev, backend=None): + """Resize, normalize, and patch an image tensor. + + *mean* and *std_dev* are per-channel float sequences (length = channels). + """ + pixel_input = self._ensure_tensor(pixel_input) + mean_arr = (ctypes.c_float * len(mean))(*[float(v) for v in mean]) + std_arr = (ctypes.c_float * len(std_dev))(*[float(v) for v in std_dev]) + out = cactus_node_t() + rc = _lib.cactus_graph_image_preprocess( + self.h, cactus_node_t(pixel_input.id), + ctypes.c_int(int(src_width)), + ctypes.c_int(int(src_height)), + ctypes.c_int(int(target_width)), + ctypes.c_int(int(target_height)), + ctypes.c_int(int(patch_size)), + ctypes.c_int(int(channels)), + ctypes.c_float(float(rescale_factor)), + mean_arr, std_arr, ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_image_preprocess failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + + # ── Introspection ──────────────────────────────────────────────── + + def output_info(self, x): + x = self._ensure_tensor(x) + return self._get_output_info(x.id) + + def _apply_backend(self, tensor, backend): + if backend is None: + return tensor + rc = _lib.cactus_graph_set_node_backend(self.h, cactus_node_t(tensor.id), ctypes.c_int32(int(backend))) + if rc != 0: + raise RuntimeError(_err("set_node_backend failed")) + return tensor + + def _binary(self, fn_name, a, b, backend=None): + a = self._ensure_tensor(a) + b = self._ensure_tensor(b) + out = cactus_node_t() + rc = getattr(_lib, fn_name)(self.h, cactus_node_t(a.id), cactus_node_t(b.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(f"{fn_name} failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def _ensure_tensor(self, x): + if not isinstance(x, Tensor): + raise TypeError("expected Tensor") + if x.g is not self: + raise ValueError("tensor belongs to a different graph") + return x + + def _get_output_info(self, node_id): + info = cactus_tensor_info_t() + rc = _lib.cactus_graph_get_output_info(self.h, cactus_node_t(node_id), ctypes.byref(info)) + if rc != 0: + raise RuntimeError(_err("graph_get_output_info failed")) + shape = tuple(int(info.shape[i]) for i in range(int(info.rank))) + return { + "precision": int(info.precision), + "rank": int(info.rank), + "shape": shape, + "num_elements": int(info.num_elements), + "byte_size": int(info.byte_size), + } + + def _tensor_from_node(self, node_id): + meta = self._get_output_info(node_id) + return Tensor(self, int(node_id), meta["shape"], meta["precision"]) + + def _coerce_input_array(self, data, precision): + if isinstance(data, Tensor): + arr = data.numpy() + else: + arr = np.asarray(data) + if precision == self.INT8: + arr = np.ascontiguousarray(arr, dtype=np.int8) + elif precision == self.FP16: + arr = np.ascontiguousarray(arr, dtype=np.float16) + elif precision == self.FP32: + arr = np.ascontiguousarray(arr, dtype=np.float32) + elif precision in (self.CQ1, self.CQ2, self.CQ3, self.CQ4): + arr = np.ascontiguousarray(arr, dtype=np.uint8) + else: + raise RuntimeError("unsupported precision") + return arr + + +class Tensor: + def __init__(self, g, node_id, shape, dtype): + self.g = g + self.id = int(node_id) + self.shape = tuple(shape) + self.dtype = int(dtype) + + def __add__(self, other): + return self.g.add(self, other) + + def __sub__(self, other): + return self.g.subtract(self, other) + + def __mul__(self, other): + return self.g.multiply(self, other) + + def __truediv__(self, other): + return self.g.divide(self, other) + + def __ne__(self, other): + if isinstance(other, Tensor): + return self.g.not_equal(self, other) + return self.g.scalar_not_equal(self, other) + + def abs(self): + return self.g.abs(self) + + def pow(self, exponent): + return self.g.pow(self, exponent) + + def precision_cast(self, dtype): + return self.g.precision_cast(self, dtype) + + def quantize_activations(self): + return self.g.quantize_activations(self) + + def scalar_add(self, value): + return self.g.scalar_add(self, value) + + def scalar_subtract(self, value): + return self.g.scalar_subtract(self, value) + + def scalar_multiply(self, value): + return self.g.scalar_multiply(self, value) + + def scalar_divide(self, value): + return self.g.scalar_divide(self, value) + + def scalar_floor_divide(self, value): + return self.g.scalar_floor_divide(self, value) + + def scalar_not_equal(self, value): + return self.g.scalar_not_equal(self, value) + + def scalar_exp(self): + return self.g.scalar_exp(self) + + def scalar_sqrt(self): + return self.g.scalar_sqrt(self) + + def scalar_cos(self): + return self.g.scalar_cos(self) + + def scalar_sin(self): + return self.g.scalar_sin(self) + + def scalar_log(self): + return self.g.scalar_log(self) + + def clamp(self, lo, hi): + return self.g.clamp(self, lo, hi) + + def relu(self): + return self.g.relu(self) + + def sigmoid(self): + return self.g.sigmoid(self) + + def tanh(self): + return self.g.tanh(self) + + def gelu(self): + return self.g.gelu(self) + + def gelu_erf(self): + return self.g.gelu_erf(self) + + def silu(self): + return self.g.silu(self) + + def view(self, shape): + return self.g.view(self, shape) + + def reshape(self, shape): + return self.g.reshape(self, shape) + + def expand(self, shape): + return self.g.expand(self, shape) + + def flatten(self, start_dim=0, end_dim=-1): + return self.g.flatten(self, start_dim=start_dim, end_dim=end_dim) + + def slice(self, axis, start, length=0): + return self.g.slice(self, axis, start, length=length) + + def index(self, index_value, axis=0): + return self.g.index(self, index_value, axis=axis) + + def transpose(self, backend=None): + return self.g.transpose(self, backend=backend) + + def permute(self, permutation, backend=None): + return self.g.permute(self, permutation, backend=backend) + + def concat(self, other, axis=0): + return self.g.concat(self, other, axis=axis) + + def cat(self, tensors, axis=0): + return self.g.cat([self] + tensors, axis=axis) + + def groupnorm(self, weight, bias, num_groups, eps=1e-5): + return self.g.groupnorm(self, weight, bias, num_groups, eps=eps) + + def layernorm(self, weight, bias=None, eps=1e-5): + return self.g.layernorm(self, weight, bias=bias, eps=eps) + + def batchnorm(self, weight, bias, running_mean, running_var, axis=1, eps=1e-5): + return self.g.batchnorm(self, weight, bias, running_mean, running_var, axis=axis, eps=eps) + + def group_norm(self, weight, bias, num_groups, eps=1e-5): + return self.groupnorm(weight, bias, num_groups, eps=eps) + + def layer_norm(self, weight, bias=None, eps=1e-5): + return self.layernorm(weight, bias=bias, eps=eps) + + def batch_norm(self, weight, bias, running_mean, running_var, axis=1, eps=1e-5): + return self.batchnorm(weight, bias, running_mean, running_var, axis=axis, eps=eps) + + def rms_norm(self, weight, eps=1e-5): + return self.g.rms_norm(self, weight, eps=eps) + + def softmax(self, axis=-1): + return self.g.softmax(self, axis) + + def glu(self, axis=-1): + return self.g.glu(self, axis=axis) + + def matmul(self, other, pretransposed_rhs=False, backend=None, output_dtype=None): + return self.g.matmul( + self, + other, + pretransposed_rhs=pretransposed_rhs, + backend=backend, + output_dtype=output_dtype, + ) + + def sum(self, axis): + return self.g.sum(self, axis) + + def mean(self, axis): + return self.g.mean(self, axis) + + def variance(self, axis): + return self.g.variance(self, axis) + + def min(self, axis): + return self.g.min(self, axis) + + def max(self, axis): + return self.g.max(self, axis) + + def cumsum(self, axis): + return self.g.cumsum(self, axis) + + def numpy(self): + info = cactus_tensor_info_t() + rc = _lib.cactus_graph_get_output_info(self.g.h, cactus_node_t(self.id), ctypes.byref(info)) + if rc != 0: + raise RuntimeError(_err("graph_get_output_info failed")) + + out_ptr = ctypes.c_void_p() + rc = _lib.cactus_graph_get_output_ptr(self.g.h, cactus_node_t(self.id), ctypes.byref(out_ptr)) + if rc != 0 or not out_ptr.value: + raise RuntimeError(_err("graph_get_output_ptr failed")) + + rank = int(info.rank) + shape = tuple(int(info.shape[i]) for i in range(rank)) + num_elements = int(info.num_elements) + precision = int(info.precision) + + if precision == Graph.FP16: + arr = np.ctypeslib.as_array((ctypes.c_uint16 * num_elements).from_address(out_ptr.value)).view(np.float16) + elif precision == Graph.FP32: + arr = np.ctypeslib.as_array((ctypes.c_float * num_elements).from_address(out_ptr.value)) + elif precision == Graph.INT8: + arr = np.ctypeslib.as_array((ctypes.c_int8 * num_elements).from_address(out_ptr.value)) + elif precision in (Graph.CQ1, Graph.CQ2, Graph.CQ3, Graph.CQ4): + arr = np.ctypeslib.as_array((ctypes.c_uint8 * int(info.byte_size)).from_address(out_ptr.value)) + return arr.copy() + else: + raise RuntimeError("unsupported precision") + + return arr.reshape(shape).copy() + + def __repr__(self): + return f"Tensor(id={self.id}, shape={self.shape}, dtype={self.dtype})" diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py new file mode 100644 index 000000000..a308d13c3 --- /dev/null +++ b/python/cactus/cli/__init__.py @@ -0,0 +1,504 @@ +import sys +import argparse + +from .. import __version__ +from .common import ( + DEFAULT_MODEL_ID, + DEFAULT_TRANSCRIPTION_MODEL_ID, + DEFAULT_TEST_MODEL_ID, + DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, +) +from .download import cmd_download +from .compile import cmd_build +from .serve import cmd_serve +from .transcribe import cmd_transcribe +from .test import cmd_test, cmd_benchmark, COMPONENTS +from .convert import cmd_convert +from .upload import cmd_upload +from .run import cmd_run +from .list import cmd_list + +from .auth import cmd_auth +from .clean import cmd_clean +from .code import cmd_code +from .utils import bits_arg, ALLOWED_BITS + +_BITS_METAVAR = "{" + ",".join(str(b) for b in ALLOWED_BITS) + "}" + + +def _telemetry_parent(): + """Args shared by commands that support telemetry toggle.""" + p = argparse.ArgumentParser(add_help=False) + p.add_argument("--no-cloud-tele", action="store_true", + help="Disable cloud telemetry (write to cache only)") + return p + + +def _build_parent(mixed: bool = True): + """Bundle-build flags shared by model-preparing commands + (mixed=False limits --bits to uniform 1-4 for build-only commands).""" + p = argparse.ArgumentParser(add_help=False) + if mixed: + p.add_argument("--bits", type=bits_arg, default=4, metavar=_BITS_METAVAR, + help="CQ quantization: uniform 1-4 or gemma-4 mixed 3.26/2.54 (default: 4)") + else: + p.add_argument("--bits", type=int, choices=[1, 2, 3, 4], default=4, + help="CQ quantization (default: 4)") + p.add_argument("--token", help="HuggingFace token") + p.add_argument("--reconvert", action="store_true", + help="Force local rebuild from source") + return p + + +def _engine_test_parent(): + """Args shared by `test` and `benchmark`.""" + p = argparse.ArgumentParser(add_help=False) + p.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") + p.add_argument("--model", dest="model_id", default=None, + type=_hf_id_or_path, + help=f"HF model ID under test (default: {DEFAULT_TEST_MODEL_ID})") + p.add_argument("--transcription-model", dest="transcription_model_id", default=None, + type=_hf_id_or_path, + help=f"HF transcription model ID under test (default: {DEFAULT_TEST_TRANSCRIPTION_MODEL_ID})") + p.add_argument("--android", action="store_true", help="Run tests on Android") + p.add_argument("--ios", action="store_true", help="Run tests on iOS") + p.add_argument("--enable-telemetry", action="store_true", + help="Enable cloud telemetry (disabled by default in tests)") + return p + + +def _positive_int(value): + n = int(value) + if n <= 0: + raise argparse.ArgumentTypeError(f"must be > 0, got {n}") + return n + + +def _non_negative_int(value): + n = int(value) + if n < 0: + raise argparse.ArgumentTypeError(f"must be >= 0, got {n}") + return n + + +def _port_int(value): + n = int(value) + if n < 1 or n > 65535: + raise argparse.ArgumentTypeError(f"port must be in 1..65535, got {n}") + return n + + +def _unit_float(value): + f = float(value) + if not (0.0 <= f <= 1.0): + raise argparse.ArgumentTypeError(f"must be in [0.0, 1.0], got {f}") + return f + + +def _hf_id_or_path(value): + v = (value or "").strip() + if "/" not in v: + raise argparse.ArgumentTypeError( + f"invalid model {value!r}. Use a HuggingFace id like 'openai/whisper-base' " + f"or a path like '/abs/path' or './bundle'." + ) + return v + + + + +def create_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, + usage=argparse.SUPPRESS, + description=f""" + + ----------------------------------------------------------------- + + Cactus CLI: + + ----------------------------------------------------------------- + + cactus auth manage cloud API key + --status show key status + --clear remove saved key + + cactus run [model|path] run a model (default: {DEFAULT_MODEL_ID}) + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --image image file for VLM inference + --audio audio file for audio chat + --system system prompt + --prompt send prompt immediately + --thinking enable thinking/reasoning mode + --token HuggingFace token (gated models) + --reconvert force local convert fallback + + cactus transcribe [model] live microphone transcription with a model + --file audio file to transcribe (WAV) + --language language code (default: en) + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --token HuggingFace token (gated models) + --reconvert force local rebuild from source + + cactus download [model] fetch a prebuilt bundle, else build locally (default: {DEFAULT_MODEL_ID}) + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --token HuggingFace token (gated models) + --reconvert force local rebuild from source + + cactus convert [dir] build a runnable bundle locally (skips prebuilt fetch) + --bits 1|2|3|4 CQ quantization (default: 4) + --token HuggingFace token (gated models) + --reconvert force local rebuild from source + --lora merge a LoRA adapter before converting + + cactus serve [model] OpenAI-compatible local HTTP server + --host bind address (default: 127.0.0.1) + --port port (default: 8080) + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --token HuggingFace token (gated models) + --reconvert force local rebuild from source + --no-cloud-handoff disable automatic cloud handoff + --confidence-threshold <0..1> handoff to cloud below this confidence + --cloud-timeout-ms max wait for cloud handoff before local fallback + + cactus list list downloaded models + + cactus build build cactus libraries + --apple Apple (iOS/macOS) + --android Android + --python shared lib for Python FFI + + cactus test run the test suite + --component kernels | graph | engine | all + (default: all) + --model default: {DEFAULT_TEST_MODEL_ID} + --transcription-model default: {DEFAULT_TEST_TRANSCRIPTION_MODEL_ID} + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --token HuggingFace token (gated models) + --reconvert force local rebuild of test models + --suite run a single test suite from any + component (kernels, graph, or engine) + --list list components and suites + --ios run on connected iPhone + --android run on connected Android + --enable-telemetry send cloud telemetry (off by default) + + cactus benchmark run the engine benchmark suite + --model default: {DEFAULT_TEST_MODEL_ID} + --transcription-model default: {DEFAULT_TEST_TRANSCRIPTION_MODEL_ID} + --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) + --backend cpu|metal inference backend (default: auto) + --ios run on connected iPhone + --android run on connected Android + + cactus clean delete build artifacts, weights, venv + + cactus --help show this help + + ----------------------------------------------------------------- + + Advanced / build pipeline — convert runs automatically inside run, + serve, transcribe and download; reach for it only to control the + build (custom flags, LoRA, debugging): + + cactus convert [dir] HuggingFace -> runnable cactus bundle + (quantizes weights to CQ, then builds + the runtime graph) + --bits 1|2|3|4 CQ quantization (default: 4) + --token HuggingFace token (defaults to $HF_TOKEN) + --reconvert force weight conversion from source + --lora merge a LoRA adapter before converting + --weights-only stop after CQ weights (skip the graph) + --weights-dir path to CQ weights (default: weights/) + --task task (default: auto, inferred from config) + --artifact-dir write bundle here (default: weights/) + --prompt representative prompt for shape capture + --system-prompt system prompt for multimodal chat + --enable-thinking enable thinking markers when supported + --input-ids token ids for causal-LM shape capture + --image-file representative image (repeatable) + --audio-file representative audio file (WAV) + --max-new-tokens preallocate decode context for causal LM + --component-pipeline auto|on|off force component-pipeline graph build + --components subset of components to build + --torch-dtype float16 | float32 | bfloat16 + --trust-remote-code allow HF remote code during the build + --local-files-only require model/processor to be local + --low-memory-load graph capture with meta tensors + --allow-unconverted-weights debug-only: skip the CQ-weights check + --execute-after-transpile run a reference execution after building + --graph-filename override saved graph filename + --skip-reference-compare skip PyTorch comparison (with --execute-…) + --no-fuse-rms-norm disable RMSNorm fusion + --no-fuse-rope disable RoPE fusion + --no-fuse-attention disable attention fusion + --no-fuse-attention-block disable attention-block fusion + --no-fuse-add-clipped disable add-clipped fusion + --no-fuse-gated-deltanet disable gated DeltaNet fusion + + ----------------------------------------------------------------- +""" + ) + + parser.add_argument("--version", action="version", version=f"cactus {__version__}") + + subparsers = parser.add_subparsers(dest='command') + subparsers.required = False + + for action in parser._actions: + if isinstance(action, argparse._SubParsersAction): + action.help = argparse.SUPPRESS + + parser._action_groups = [] + + download_parser = subparsers.add_parser("download", + help="Download a pre-built bundle from huggingface.co/Cactus-Compute", + parents=[_build_parent()]) + download_parser.add_argument("model_id", nargs="?", default=DEFAULT_MODEL_ID, + type=_hf_id_or_path, + help=f"HuggingFace model id (default: {DEFAULT_MODEL_ID})") + + build_parser = subparsers.add_parser("build", help="Build cactus libraries") + build_group = build_parser.add_mutually_exclusive_group() + build_group.add_argument("--apple", action="store_true", + help="Build for Apple (iOS/macOS)") + build_group.add_argument("--android", action="store_true", + help="Build for Android") + build_group.add_argument("--python", action="store_true", + help="Build shared library for Python FFI") + + run_parser = subparsers.add_parser("run", help="Run a model (downloads bundle if needed)", + parents=[_telemetry_parent(), _build_parent()]) + run_parser.add_argument("model_id", nargs="?", default=DEFAULT_MODEL_ID, + type=_hf_id_or_path, + help=f"HuggingFace model id or local bundle path (default: {DEFAULT_MODEL_ID})") + run_parser.add_argument("--image", + help="Path to image file for VLM inference (attached to first message)") + run_parser.add_argument("--audio", + help="Path to audio file (WAV) for audio chat (attached to first message)") + run_parser.add_argument("--system", + help="System prompt to prepend to all messages") + run_parser.add_argument("--prompt", + help="Initial prompt to send immediately") + run_parser.add_argument("--input-ids", default=None, + help="Comma-separated token ids for causal-LM bundles") + run_parser.add_argument("--input-ids-file", default=None, + help="File containing token ids for causal-LM bundles") + run_parser.add_argument("--max-new-tokens", type=_positive_int, default=None, + help="Maximum tokens to generate for causal-LM bundles") + run_parser.add_argument("--result-json", default=None, + help="Optional path to save bundle results as JSON") + run_parser.add_argument("--thinking", action="store_true", + help="Enable thinking/reasoning for models that support it") + run_parser.add_argument("--no-cloud-handoff", action="store_true", + help="Disable automatic cloud handoff for this run") + run_parser.add_argument("--confidence-threshold", type=_unit_float, default=None, + help="Confidence threshold below which local completions may hand off to cloud") + run_parser.add_argument("--cloud-timeout-ms", type=_non_negative_int, default=None, + help="Maximum time to wait for cloud handoff before falling back locally") + run_parser.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") + + transcribe_parser = subparsers.add_parser("transcribe", help="Transcribe audio with a model", + parents=[_telemetry_parent(), _build_parent()]) + transcribe_parser.add_argument("model_id", nargs="?", default=DEFAULT_TRANSCRIPTION_MODEL_ID, + type=_hf_id_or_path, + help=f"HuggingFace model id (default: {DEFAULT_TRANSCRIPTION_MODEL_ID})") + transcribe_parser.add_argument("--file", dest="audio_file", default=None, + help="Audio file to transcribe (WAV)") + transcribe_parser.add_argument("--language", default="en", + help="Language code (default: en)") + + serve_parser = subparsers.add_parser("serve", help="OpenAI-compatible local HTTP server", + parents=[_telemetry_parent(), _build_parent()]) + serve_parser.add_argument("model_id", nargs="?", default=None, + type=_hf_id_or_path, + help="HuggingFace model id (e.g. openai/whisper-base) or bundle path") + serve_parser.add_argument("--host", default="127.0.0.1", + help="Bind address (default: 127.0.0.1)") + serve_parser.add_argument("--port", type=_port_int, default=8080, + help="Port (default: 8080)") + serve_parser.add_argument("--no-cloud-handoff", action="store_true", + help="Disable automatic cloud handoff for all requests") + serve_parser.add_argument("--confidence-threshold", type=_unit_float, default=None, + help="Confidence threshold below which completions hand off to cloud (1.0 forces cloud handoff)") + serve_parser.add_argument("--cloud-timeout-ms", type=_non_negative_int, default=None, + help="Maximum time to wait for cloud handoff before falling back locally") + serve_parser.add_argument("--no-access-log", action="store_true", + help="Disable per-request HTTP access logging") + serve_parser.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") + + code_parser = subparsers.add_parser("code", help="Run the Cactus coding agent (TUI / print mode)", + parents=[_build_parent()]) + code_parser.add_argument("--serve-model", default=None, + help="If no server is running, start `cactus serve` with this model") + code_parser.add_argument("--host", default="127.0.0.1", + help="Server bind address to use/start (default: 127.0.0.1)") + code_parser.add_argument("--port", type=_port_int, default=8080, + help="Server port to use/start (default: 8080)") + code_parser.add_argument("--no-serve", action="store_true", + help="Do not auto-start a server; require one to already be running") + code_parser.add_argument("--no-cloud-handoff", action="store_true", + help="Disable automatic cloud handoff on the auto-started server") + code_parser.add_argument("--confidence-threshold", type=_unit_float, default=None, + help="Confidence threshold below which completions hand off to cloud") + code_parser.add_argument("--cloud-timeout-ms", type=_non_negative_int, default=None, + help="Maximum time to wait for cloud handoff before falling back locally") + code_parser.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") + code_parser.add_argument("agent_args", nargs=argparse.REMAINDER, + help="Arguments passed through to the coding agent (prefix with -- )") + + test_parser = subparsers.add_parser("test", help="Run the test suite", + parents=[_build_parent(), _engine_test_parent()]) + test_parser.add_argument("--component", choices=COMPONENTS, default="all", + help="Component to test (default: all)") + test_parser.add_argument("--suite", default=None, + help="Run a single test suite by name; resolved across all components (e.g. llm → engine)") + test_parser.add_argument("--list", action="store_true", + help="List available components and engine tests, then exit") + + subparsers.add_parser("benchmark", help="Run the engine benchmark suite", + parents=[_build_parent(), _engine_test_parent()]) + + auth_parser = subparsers.add_parser("auth", help="Manage cloud API key") + auth_parser.add_argument("--clear", action="store_true", + help="Remove the saved API key") + auth_parser.add_argument("--status", action="store_true", + help="Show current key status") + + clean_parser = subparsers.add_parser("clean", help="Delete build artifacts, downloaded weights, and venv") + clean_parser.add_argument("--yes", "-y", action="store_true", help="Skip confirmation prompt") + + subparsers.add_parser("list", help="List downloaded models") + + convert_parser = subparsers.add_parser("convert", + help="Convert a HuggingFace model into a runnable cactus bundle") + convert_parser.add_argument("model_id", type=_hf_id_or_path, + help="HuggingFace model id (e.g. openai/whisper-base)") + convert_parser.add_argument("output_dir", nargs="?", default=None, + help="Output directory (default: weights/)") + convert_parser.add_argument("--bits", type=int, choices=[1, 2, 3, 4], default=4, + help="CQ quantization bits (default: 4)") + convert_parser.add_argument("--token", help="HuggingFace token") + convert_parser.add_argument("--lora", + help="Path or HF id of a LoRA adapter to merge before converting (requires `peft`)") + convert_parser.add_argument("--reconvert", action="store_true", + help="Force conversion from source") + convert_parser.add_argument("--skip-model-load", action="store_true", + help="Convert directly from checkpoint tensors without loading the full HF model object") + convert_parser.add_argument("--low-memory-load", action="store_true", + help="Avoid loading checkpoint tensors during graph capture") + convert_parser.add_argument("--weights-only", action="store_true", + help="Only quantize weights; skip building the runtime graph bundle") + convert_parser.add_argument("--weights-dir", + help="CQ weights directory (default: weights/)") + convert_parser.add_argument("--task", default="auto", + choices=["auto", "causal_lm_logits", "multimodal_causal_lm_logits", + "ctc_logits", "encoder_hidden_states", + "seq2seq_transcription", "tdt_transcription"], + help="Graph task (default: auto, from model config)") + convert_parser.add_argument("--prompt", + help="Prompt for causal/multimodal shape capture") + convert_parser.add_argument("--system-prompt", default=None, + help="System prompt for multimodal chat formats") + convert_parser.add_argument("--enable-thinking", action="store_true", + help="Enable thinking markers when the prompt supports them") + convert_parser.add_argument("--input-ids", default=None, + help="Comma-separated token ids for causal-LM shape capture") + convert_parser.add_argument("--image-file", action="append", default=[], + help="Image for multimodal shape capture (repeatable)") + convert_parser.add_argument("--audio-file", + help="Audio file (WAV) for audio/multimodal shape capture") + convert_parser.add_argument("--max-new-tokens", type=_positive_int, default=None, + help="Decode context to preallocate for causal LM (default: 32)") + convert_parser.add_argument("--component-pipeline", default="auto", choices=["auto", "on", "off"], + help="Split-component graph when supported (default: auto)") + convert_parser.add_argument("--components", + help="Comma-separated component subset (e.g. vision_encoder,decoder)") + convert_parser.add_argument("--torch-dtype", default=None, + choices=["float16", "float32", "bfloat16"], + help="Torch dtype for HF loading (default: float16)") + convert_parser.add_argument("--trust-remote-code", action="store_true", + help="Pass trust_remote_code=True to HF loaders") + convert_parser.add_argument("--local-files-only", action="store_true", + help="Require model/processor to already be local") + convert_parser.add_argument("--allow-unconverted-weights", action="store_true", + help="Debug: build the graph without CQ weights") + convert_parser.add_argument("--execute-after-transpile", action="store_true", + help="Run a reference execution after building the graph") + convert_parser.add_argument("--artifact-dir", + help="Graph bundle output directory (default: weights/)") + convert_parser.add_argument("--graph-filename", default=None, + help="Saved graph filename (default: graph.cactus)") + convert_parser.add_argument("--skip-reference-compare", action="store_true", + help="Skip PyTorch comparison (requires --execute-after-transpile)") + convert_parser.add_argument("--no-fuse-rms-norm", action="store_true", + help="Disable RMSNorm fusion") + convert_parser.add_argument("--no-fuse-rope", action="store_true", + help="Disable RoPE fusion") + convert_parser.add_argument("--no-fuse-attention", action="store_true", + help="Disable attention fusion") + convert_parser.add_argument("--no-fuse-attention-block", action="store_true", + help="Disable attention-block fusion") + convert_parser.add_argument("--no-fuse-add-clipped", action="store_true", + help="Disable add-clipped fusion") + convert_parser.add_argument("--no-fuse-gated-deltanet", action="store_true", + help="Disable gated DeltaNet fusion") + convert_parser.add_argument("--cache-context-length", default=None, + help="KV cache context length for cached decode graphs (default: model config)") + + upload_parser = subparsers.add_parser("upload", + parents=[_build_parent(mixed=False)], + help="Build a runnable bundle locally and upload it to Cactus-Compute on HuggingFace") + upload_parser.add_argument("model_id", type=_hf_id_or_path, + help="HuggingFace model id (e.g. openai/whisper-base)") + + return parser + + + +_COMMANDS = { + "download": cmd_download, + "build": cmd_build, + "run": cmd_run, + "serve": cmd_serve, + "transcribe": cmd_transcribe, + "test": cmd_test, + "benchmark": cmd_benchmark, + "list": cmd_list, + + "auth": cmd_auth, + "clean": cmd_clean, + "code": cmd_code, + "convert": cmd_convert, + "upload": cmd_upload, +} + + +_REPO_ONLY = {"build", "test", "benchmark", "clean"} + + +def main(): + from .common import is_repo_checkout + + parser = create_parser() + args = parser.parse_args() + + if args.command in _REPO_ONLY and not is_repo_checkout(): + print(f"Error: `cactus {args.command}` requires a git clone of the cactus repository.") + print("See: https://github.com/cactus-compute/cactus") + sys.exit(1) + + handler = _COMMANDS.get(args.command) + if handler: + sys.exit(handler(args)) + else: + parser.print_help() + sys.exit(1) + + +if __name__ == '__main__': + main() diff --git a/python/cactus/cli/auth.py b/python/cactus/cli/auth.py new file mode 100644 index 000000000..38989fba7 --- /dev/null +++ b/python/cactus/cli/auth.py @@ -0,0 +1,41 @@ +import getpass +import sys + +from .common import print_color, mask_key, CYAN, GREEN, NC, YELLOW + + +def cmd_auth(args): + from .config_utils import CactusConfig + + config = CactusConfig() + + if args.clear: + config.clear_api_key() + print_color(GREEN, "API key cleared.") + return 0 + + try: + api_key = config.get_api_key() + except (OSError, ValueError): + print_color(YELLOW, "Existing config is unreadable; set or clear a key to reset it.") + api_key = None + + if api_key: + print(f"Current API key: {mask_key(api_key)}") + else: + print("No API key set.") + + if args.status: + return 0 + + if not sys.stdin.isatty(): + print_color(YELLOW, "stdin is not a TTY; refusing interactive key entry. Set CACTUS_CLOUD_KEY env var instead.") + return 0 + + print() + print(f"Get your cloud key at {CYAN}https://www.cactuscompute.com/dashboard/api-keys{NC}") + new_key = getpass.getpass("Enter new API key (press Enter to skip): ").strip() + if new_key: + config.set_api_key(new_key) + print_color(GREEN, f"API key saved: {mask_key(new_key)}") + return 0 diff --git a/python/cactus/cli/clean.py b/python/cactus/cli/clean.py new file mode 100644 index 000000000..0901997e8 --- /dev/null +++ b/python/cactus/cli/clean.py @@ -0,0 +1,121 @@ +import os +import shutil +import subprocess +import sys +from pathlib import Path + +from .common import ( + PROJECT_ROOT, + print_color, + RED, GREEN, YELLOW, BLUE, +) + + +def cmd_clean(args): + """Remove all build artifacts, caches, and downloaded weights.""" + print_color(BLUE, "Cleaning all build artifacts from Cactus project...") + print(f"Project root: {PROJECT_ROOT}") + print() + + destructive_targets = [ + PROJECT_ROOT / "weights", + PROJECT_ROOT / "transpiled", + PROJECT_ROOT / "venv", + ] + present = [t for t in destructive_targets if t.exists()] + if present and not args.yes: + print_color(YELLOW, "This also deletes the following (not just build artifacts):") + for t in present: + print(f" - {t}") + print_color(YELLOW, "Downloaded weights are NOT re-fetched by setup and will be lost.") + if not sys.stdin.isatty(): + print_color(RED, "Refusing to proceed without confirmation. Re-run with --yes to skip this prompt.") + return 0 + if input("Continue? [y/N]: ").strip().lower() not in ("y", "yes"): + print_color(YELLOW, "Aborted.") + return 0 + + print("Stopping any running Cactus server...") + try: + stopped = subprocess.run(["pkill", "-f", "cactus serve"], capture_output=True) + print("Stopped running Cactus server(s)." if stopped.returncode == 0 else "No running Cactus server found.") + except FileNotFoundError: + print_color(YELLOW, "Could not stop the server automatically; stop it manually if one is running.") + print() + + preserve_roots = [ + (PROJECT_ROOT / "cactus-engine" / "libs" / "curl").resolve(), + (PROJECT_ROOT / "android" / "mbedtls").resolve(), + (PROJECT_ROOT / "libs" / "mbedtls").resolve(), + ] + + def should_preserve(path): + resolved = path.resolve() + return any(resolved.is_relative_to(root) for root in preserve_roots) + + def remove_if_exists(path): + if path.is_dir(): + print(f"Removing: {path}") + shutil.rmtree(path, ignore_errors=True) + + for tree in ("venv", "weights", "transpiled"): + remove_if_exists(PROJECT_ROOT / tree) + remove_if_exists(PROJECT_ROOT / "python" / "cactus" / "bin") + remove_if_exists(PROJECT_ROOT / "python" / "cactus" / "code") + + print() + print("Sweeping build, dependency, and cache artifacts across the tree...") + + def is_artifact_dir(name): + return ( + name in {"node_modules", "dist", "__pycache__", ".pytest_cache"} + or name == "build" + or name.startswith(("build-", "build_")) + or name.endswith((".egg-info", ".xcframework")) + ) + + artifact_file_suffixes = {".so", ".a", ".bin", ".dylib"} + + removed_dirs = 0 + removed_files = 0 + for root, dirs, files in os.walk(PROJECT_ROOT, topdown=True): + root_path = Path(root) + kept = [] + for d in dirs: + if d == ".git": + continue + full = root_path / d + if is_artifact_dir(d) and not should_preserve(full): + print(f"Removing: {full}") + shutil.rmtree(full, ignore_errors=True) + removed_dirs += 1 + else: + kept.append(d) + dirs[:] = kept + for f in files: + full = root_path / f + if full.suffix in artifact_file_suffixes and not should_preserve(full): + try: + full.unlink() + removed_files += 1 + except OSError: + pass + print(f"Removed {removed_dirs} build/dependency directories and {removed_files} compiled artifacts") + + print() + print_color(GREEN, "Clean complete!") + print("All build artifacts, weights, and venv have been removed.") + print() + + print_color(BLUE, "Re-running setup...") + setup_script = PROJECT_ROOT / "setup" + result = subprocess.run( + ["bash", "-c", f"source {setup_script}"], + cwd=PROJECT_ROOT + ) + if result.returncode == 0: + print_color(GREEN, "Setup complete!") + else: + print_color(YELLOW, "Setup had issues. Please run manually:") + print(" source ./setup") + return 0 diff --git a/python/cactus/cli/code.py b/python/cactus/cli/code.py new file mode 100644 index 000000000..8c35449c5 --- /dev/null +++ b/python/cactus/cli/code.py @@ -0,0 +1,190 @@ +import os +import shutil +import subprocess +import sys +import tempfile +import time +import urllib.request +from pathlib import Path + +from .common import BLUE, DEFAULT_MODEL_ID, GREEN, RED, YELLOW, is_valid_bundle, print_color, weights_root + +_PACKAGE_DIR = Path(__file__).resolve().parent.parent +_PROJECT_ROOT = _PACKAGE_DIR.parent.parent + + +def _find_agent_cli() -> Path | None: + candidates = ( + _PACKAGE_DIR / "code" / "node_modules" / "@earendil-works" / "pi-coding-agent" / "dist" / "cli.js", + _PACKAGE_DIR / "code" / "packages" / "coding-agent" / "dist" / "cli.js", + _PROJECT_ROOT / "cactus-code" / "packages" / "coding-agent" / "dist" / "cli.js", + ) + for cli in candidates: + if cli.exists(): + return cli + return None + + +def _ensure_agent_built() -> Path | None: + cli = _find_agent_cli() + if cli is not None: + return cli + source_dir = _PROJECT_ROOT / "cactus-code" + if not (source_dir / "package.json").exists(): + return None + npm = shutil.which("npm") + if not npm: + print_color(RED, "Error: npm is required to build the Cactus coding agent on first run.") + print_color(YELLOW, "Install Node.js >= 22 (includes npm) from https://nodejs.org and retry.") + return None + print_color(BLUE, "Building the Cactus coding agent (first run, this can take a minute) ...") + try: + subprocess.run([npm, "install", "--no-audit", "--no-fund"], cwd=source_dir, check=True) + subprocess.run([npm, "run", "build"], cwd=source_dir, check=True) + except subprocess.CalledProcessError: + print_color(RED, "Error: failed to build the Cactus coding agent.") + return None + return _find_agent_cli() + + +def _discover_local_model() -> str | None: + try: + root = weights_root() + if not root.is_dir(): + return None + bundles = sorted(p for p in root.iterdir() if p.is_dir() and is_valid_bundle(p)) + return str(bundles[0]) if bundles else None + except Exception: + return None + + +def _clear_console() -> None: + if sys.stdout.isatty(): + sys.stdout.write("\033[2J\033[3J\033[H") + sys.stdout.flush() + + +def _base_url(host: str, port: int) -> str: + env = os.environ.get("CACTUS_BASE_URL") + if env: + return env.rstrip("/") + return f"http://{host}:{port}/v1" + + +def _server_alive(base_url: str, timeout: float = 1.5) -> bool: + try: + with urllib.request.urlopen(f"{base_url}/models", timeout=timeout) as resp: + return resp.status == 200 + except Exception: + return False + + +def _echo_log_tail(log_path: Path, pos: int) -> int: + try: + with open(log_path, "r", errors="replace") as f: + f.seek(pos) + chunk = f.read() + new_pos = f.tell() + if chunk: + sys.stdout.write(chunk) + sys.stdout.flush() + return new_pos + except Exception: + return pos + + +def _wait_for_server(base_url: str, proc: "subprocess.Popen | None" = None, timeout: float = 1800.0, + log_path: "Path | None" = None) -> bool: + deadline = time.time() + timeout + pos = 0 + while time.time() < deadline: + if log_path is not None: + pos = _echo_log_tail(log_path, pos) + if _server_alive(base_url): + return True + if proc is not None and proc.poll() is not None: + return False + time.sleep(0.5) + return False + + +def cmd_code(args) -> int: + node = shutil.which("node") + if not node: + print_color(RED, "Error: Node.js (>=22) is required to run `cactus code`.") + print_color(YELLOW, "Install it from https://nodejs.org and try again.") + return 1 + + agent_was_built = _find_agent_cli() is None + agent_cli = _ensure_agent_built() + if agent_cli is None: + print_color(RED, "Error: the Cactus coding agent is not available in this installation.") + return 1 + + base_url = _base_url(args.host, args.port) + started_server: subprocess.Popen | None = None + + if not _server_alive(base_url): + if args.no_serve or os.environ.get("CACTUS_BASE_URL"): + print_color(RED, f"Error: no Cactus server reachable at {base_url}.") + print_color(YELLOW, "Start one with `cactus serve ` first.") + return 1 + model_flags_given = ( + getattr(args, "reconvert", False) + or getattr(args, "bits", 4) != 4 + or bool(getattr(args, "token", None)) + ) + if args.serve_model: + serve_model = args.serve_model + elif model_flags_given: + serve_model = DEFAULT_MODEL_ID + else: + serve_model = _discover_local_model() or DEFAULT_MODEL_ID + if serve_model == DEFAULT_MODEL_ID and not args.serve_model: + print_color(BLUE, f"No model specified; using the default ({DEFAULT_MODEL_ID}), downloading on first run ...") + print_color(BLUE, f"Starting Cactus server with {serve_model} ...") + serve_cmd = [ + sys.executable, "-m", "cactus", "serve", serve_model, + "--host", args.host, "--port", str(args.port), "--no-access-log", + "--bits", str(getattr(args, "bits", 4)), + ] + if getattr(args, "backend", None): + serve_cmd += ["--backend", args.backend] + if getattr(args, "token", None): + serve_cmd += ["--token", args.token] + if getattr(args, "reconvert", False): + serve_cmd.append("--reconvert") + if getattr(args, "no_cloud_handoff", False): + serve_cmd.append("--no-cloud-handoff") + if getattr(args, "confidence_threshold", None) is not None: + serve_cmd += ["--confidence-threshold", str(args.confidence_threshold)] + if getattr(args, "cloud_timeout_ms", None) is not None: + serve_cmd += ["--cloud-timeout-ms", str(args.cloud_timeout_ms)] + serve_log_path = Path(tempfile.gettempdir()) / "cactus-code-serve.log" + serve_log = open(serve_log_path, "w") + started_server = subprocess.Popen(serve_cmd, stdout=serve_log, stderr=subprocess.STDOUT) + if not _wait_for_server(base_url, started_server, log_path=serve_log_path): + print_color(RED, "Error: the Cactus server did not become ready in time.") + print_color(YELLOW, f"See server logs: {serve_log_path}") + started_server.terminate() + return 1 + print_color(GREEN, f"Cactus server ready (logs: {serve_log_path})") + + env = dict(os.environ) + env["CACTUS_BASE_URL"] = base_url + env["PI_PACKAGE_DIR"] = str(agent_cli.parent.parent) + agent_args = [a for a in (args.agent_args or []) if a != "--"] + interactive = not any(a in ("-p", "--print") for a in agent_args) + if interactive and (started_server is not None or agent_was_built): + _clear_console() + try: + return subprocess.run([node, str(agent_cli), *agent_args], env=env).returncode + except KeyboardInterrupt: + return 130 + finally: + if started_server is not None: + started_server.terminate() + try: + started_server.wait(timeout=10) + except Exception: + started_server.kill() diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py new file mode 100644 index 000000000..87ae3daae --- /dev/null +++ b/python/cactus/cli/common.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +import os +import subprocess +import sys +from pathlib import Path + +SCRIPT_DIR = Path(__file__).resolve().parent + + +def _looks_like_project_root(path): + return ( + (path / "python" / "cactus" / "cli").is_dir() + and (path / "cactus-kernels").is_dir() + ) + + +PROJECT_ROOT = SCRIPT_DIR.parent.parent.parent + + +def is_repo_checkout(): + return _looks_like_project_root(PROJECT_ROOT) + + +def weights_root() -> Path: + if is_repo_checkout(): + return PROJECT_ROOT / "weights" + return Path.home() / ".cache" / "cactus" / "weights" + + +def transpiled_root() -> Path: + if is_repo_checkout(): + return PROJECT_ROOT / "transpiled" + return Path.home() / ".cache" / "cactus" / "transpiled" + + +DEFAULT_MODEL_ID = "google/gemma-4-E2B-it" +DEFAULT_TRANSCRIPTION_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3" +DEFAULT_TEST_MODEL_ID = DEFAULT_MODEL_ID +DEFAULT_TEST_TRANSCRIPTION_MODEL_ID = DEFAULT_TRANSCRIPTION_MODEL_ID + + +RED = '\033[0;31m' +GREEN = '\033[0;32m' +YELLOW = '\033[1;33m' +BLUE = '\033[0;34m' +CYAN = '\033[1;36m' +NC = '\033[0m' + + +def _color_enabled(): + if os.environ.get("NO_COLOR"): + return False + return sys.stdout.isatty() + + +def print_color(color, message): + if _color_enabled(): + print(f"{color}{message}{NC}") + else: + print(message) + + +def mask_key(key): + return key[:4] + "..." + key[-4:] if len(key) >= 8 else "***" + + +def convert_toolchain_error(): + """Message if the model-conversion toolchain (torch/transformers) is missing, else None.""" + import importlib.util + missing = [m for m in ("torch", "transformers") + if importlib.util.find_spec(m) is None] + if not missing: + return None + return ( + "the model-conversion toolchain is not installed (missing: " + + ", ".join(missing) + ").\n" + " Converting a model from source needs it. Either:\n" + " - use a prebuilt model: cactus run or cactus download \n" + " - install the toolchain: pip install \"cactus-compute[convert]\"" + ) + + +BIN_DIR = SCRIPT_DIR.parent / "bin" + + +def apply_cloud_api_key_env() -> None: + from .config_utils import CactusConfig + try: + api_key = CactusConfig().get_api_key() + except (OSError, ValueError): + return + if api_key: + os.environ["CACTUS_CLOUD_KEY"] = api_key + + +def apply_runtime_env(args) -> None: + """Prepare the process env for an inference command: honour --no-cloud-tele + and load the stored cloud API key.""" + if getattr(args, "no_cloud_tele", False): + os.environ["CACTUS_NO_CLOUD_TELE"] = "1" + apply_cloud_api_key_env() + + +def is_valid_bundle(path) -> bool: + """A runnable v2 bundle has both config.txt and the components manifest.""" + path = Path(path) + return (path / "config.txt").exists() and (path / "components" / "manifest.json").exists() + + +def launch_binary(name, *args) -> int: + """Resolve a bundled binary and exec it with str-coerced args. Returns its + exit code, or 1 if the binary is unavailable.""" + binary = resolve_binary(name) + if binary is None: + return 1 + return subprocess.run([str(binary), *(str(a) for a in args)]).returncode + + +def _auto_build_binaries() -> bool: + """Build the native runtime and CLI binaries, equivalent to `cactus build`.""" + from argparse import Namespace + from .compile import cmd_build + + return cmd_build(Namespace(apple=False, android=False, python=False)) == 0 + + +def resolve_binary(name): + path = BIN_DIR / name + if path.exists(): + return path + + # First run in a source checkout: build automatically instead of erroring out. + if is_repo_checkout(): + print_color(YELLOW, f"{name} binary not found; building Cactus (first run, this may take a minute)...") + if _auto_build_binaries() and path.exists(): + return path + + print_color(RED, f"{name} binary not found at {path}.") + return None diff --git a/python/cactus/cli/compile.py b/python/cactus/cli/compile.py new file mode 100644 index 000000000..acada64f7 --- /dev/null +++ b/python/cactus/cli/compile.py @@ -0,0 +1,219 @@ +import shutil +import subprocess +import platform +from pathlib import Path + +from .common import ( + PROJECT_ROOT, + print_color, + RED, GREEN, YELLOW, BLUE, +) + + +def check_command(cmd): + return shutil.which(cmd) is not None + + +def run_command(cmd, cwd=None): + if isinstance(cmd, str): + cmd = [cmd] + return subprocess.run(cmd, cwd=cwd) + + +def check_libcurl(): + if platform.system() == 'Darwin': + return True + + if check_command('pkg-config'): + result = subprocess.run(['pkg-config', '--exists', 'libcurl'], capture_output=True) + if result.returncode == 0: + return True + + curl_paths = [ + '/usr/include/curl/curl.h', + '/usr/include/x86_64-linux-gnu/curl/curl.h', + '/usr/include/aarch64-linux-gnu/curl/curl.h', + '/usr/local/include/curl/curl.h', + ] + for path in curl_paths: + if Path(path).exists(): + return True + + return False + + +def _detect_sdl2() -> tuple[list[str], list[str]]: + if not check_command("pkg-config"): + return [], [] + + if subprocess.run(["pkg-config", "--exists", "sdl2"], capture_output=True).returncode != 0: + return [], [] + + cflags = subprocess.run(["pkg-config", "--cflags", "sdl2"], capture_output=True, text=True) + libs = subprocess.run(["pkg-config", "--libs", "sdl2"], capture_output=True, text=True) + if cflags.returncode != 0 or libs.returncode != 0: + return [], [] + + return ( + ["-DHAVE_SDL2"] + cflags.stdout.strip().split(), + libs.stdout.strip().split(), + ) + + +def build_binary( + name: str, + lib_path: Path, + *, + sdl2: tuple[list[str], list[str]] | None = None, +) -> int: + tests_dir = PROJECT_ROOT / "cactus-engine" / "tests" + source = tests_dir / f"{name}.cpp" + build_dir = tests_dir / "build" + build_dir.mkdir(parents=True, exist_ok=True) + + is_darwin = platform.system() == "Darwin" + sdl2_flags, sdl2_link = sdl2 if sdl2 is not None else _detect_sdl2() + + include_dirs = [ + PROJECT_ROOT, + PROJECT_ROOT / "cactus-engine", + PROJECT_ROOT / "cactus-graph", + PROJECT_ROOT / "cactus-kernels", + PROJECT_ROOT / "cactus-kernels" / "src", + ] + + print(f"Compiling {name}.cpp...") + + if is_darwin: + vendored_curl = PROJECT_ROOT / "cactus-engine" / "libs" / "curl" / "lib" / "libcurl.a" + curl_link = [str(vendored_curl)] if vendored_curl.exists() else ["-lcurl"] + compiler = "clang++" + cmd = [ + compiler, "-std=c++20", "-O3", + "-DACCELERATE_NEW_LAPACK", + *[f"-I{d}" for d in include_dirs], + *sdl2_flags, + str(source), str(lib_path), + "-o", name, + *curl_link, + "-framework", "Accelerate", + "-framework", "Foundation", + "-framework", "Metal", + "-framework", "MetalPerformanceShaders", + "-framework", "Security", + "-framework", "SystemConfiguration", + "-framework", "CFNetwork", + *sdl2_link, + ] + else: + compiler = "g++" + cmd = [ + compiler, "-std=c++20", "-O3", + *[f"-I{d}" for d in include_dirs], + *sdl2_flags, + str(source), str(lib_path), + "-o", name, + "-lcurl", "-pthread", + *sdl2_link, + ] + + if not check_command(compiler): + print_color(RED, f"Error: {compiler} is not installed") + return 1 + + result = subprocess.run(cmd, cwd=build_dir) + if result.returncode != 0: + print_color(RED, f"{name} build failed") + return 1 + + bin_dir = Path(__file__).resolve().parent.parent / "bin" + bin_dir.mkdir(parents=True, exist_ok=True) + shutil.copy2(build_dir / name, bin_dir / name) + print_color(GREEN, f"Build complete: {bin_dir / name}") + return 0 + + +def cmd_build(args): + if args.apple: + return _build_with_script("apple", "Building Cactus for Apple platforms") + if args.android: + return _build_with_script("android", "Building Cactus for Android") + if args.python: + return cmd_build_python() + + print_color(BLUE, "Building Cactus library...") + print("=" * 24) + + if not check_command('cmake'): + print_color(RED, "Error: CMake is not installed") + print(" macOS: brew install cmake") + print(" Ubuntu: sudo apt-get install cmake build-essential") + return 1 + + if not check_libcurl(): + print_color(RED, "Error: libcurl development libraries not found") + print(" macOS: brew install curl") + print(" Ubuntu: sudo apt-get install libcurl4-openssl-dev") + return 1 + + cactus_dir = PROJECT_ROOT / "cactus-engine" + lib_path = cactus_dir / "build" / "libcactus_engine.a" + + if run_command(str(cactus_dir / "build.sh"), cwd=cactus_dir).returncode != 0: + print_color(RED, "Failed to build cactus library") + return 1 + + sdl2 = _detect_sdl2() + if sdl2[0]: + print_color(GREEN, "SDL2 found - voice input enabled") + else: + print_color(YELLOW, "SDL2 not found - voice input disabled") + print_color(YELLOW, "Install SDL2 to enable voice input: brew install sdl2 (macOS)") + + rc = build_binary("run", lib_path, sdl2=([], [])) + if rc != 0: + return rc + rc = build_binary("transcribe", lib_path, sdl2=sdl2) + if rc != 0: + return rc + + print_color(GREEN, "Cactus library built successfully!") + print(f"Library location: {lib_path}") + + return 0 + + +def _build_with_script(subdir, title): + print_color(BLUE, f"{title}...") + + if subdir == "apple" and platform.system() != "Darwin": + print_color(RED, "Error: Apple builds require macOS") + return 1 + + if run_command(str(PROJECT_ROOT / subdir / "build.sh"), cwd=PROJECT_ROOT / subdir).returncode != 0: + print_color(RED, f"{title} failed") + return 1 + + print_color(GREEN, f"{title} complete!") + return 0 + + +def cmd_build_python(): + print_color(BLUE, "Building Cactus for Python...") + + if not check_command('cmake'): + print_color(RED, "Error: CMake is not installed") + print(" macOS: brew install cmake") + print(" Ubuntu: sudo apt-get install cmake") + return 1 + + cactus_dir = PROJECT_ROOT / "cactus-engine" + if run_command(str(cactus_dir / "build.sh"), cwd=cactus_dir).returncode != 0: + print_color(RED, "Build failed") + return 1 + + lib_name = "libcactus_engine.dylib" if platform.system() == "Darwin" else "libcactus_engine.so" + lib_path = cactus_dir / "build" / lib_name + print_color(GREEN, "Python build complete!") + print(f"Library: {lib_path}") + return 0 diff --git a/python/cactus/cli/config_utils.py b/python/cactus/cli/config_utils.py new file mode 100644 index 000000000..6b6ecd154 --- /dev/null +++ b/python/cactus/cli/config_utils.py @@ -0,0 +1,57 @@ +import json +import os +import tempfile +from pathlib import Path + + +class CactusConfig: + def __init__(self): + self.config_dir = Path.home() / ".cactus" + self.config_file = self.config_dir / "config.json" + self.config_dir.mkdir(exist_ok=True) + try: + os.chmod(self.config_dir, 0o700) + except OSError: + pass + + def load_config(self): + if self.config_file.exists(): + return json.loads(self.config_file.read_text()) + return {} + + def save_config(self, config): + payload = json.dumps(config, indent=2) + fd, tmp_path = tempfile.mkstemp(prefix="config.", suffix=".tmp", dir=str(self.config_dir)) + try: + with os.fdopen(fd, "w") as f: + f.write(payload) + os.chmod(tmp_path, 0o600) + os.replace(tmp_path, self.config_file) + except Exception: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + + def _load_or_empty(self): + try: + return self.load_config() + except (OSError, ValueError): + return {} + + def get_api_key(self): + env_key = os.getenv("CACTUS_CLOUD_KEY") or os.getenv("CACTUS_CLOUD_API_KEY") + if env_key: + return env_key + return self.load_config().get("api_key") or None + + def set_api_key(self, key): + config = self._load_or_empty() + config["api_key"] = key + self.save_config(config) + + def clear_api_key(self): + config = self._load_or_empty() + config.pop("api_key", None) + self.save_config(config) diff --git a/python/cactus/cli/convert.py b/python/cactus/cli/convert.py new file mode 100644 index 000000000..f6da1b1db --- /dev/null +++ b/python/cactus/cli/convert.py @@ -0,0 +1,170 @@ +import shutil +import tempfile +from pathlib import Path + +from .common import GREEN, RED, YELLOW, print_color +from .download import get_bundle_dir, get_weights_dir + + +def _merge_lora_adapter(base_model_id, lora_path, token=None): + """Merge a LoRA/PEFT adapter into the base model, save to a temp dir, return path.""" + try: + from peft import PeftModel + except ImportError: + print_color(RED, "Error: `peft` is required for LoRA merging.") + print("Install with: pip install peft") + return None + + from transformers import AutoModelForCausalLM, AutoTokenizer + + print_color(YELLOW, f"Loading base model: {base_model_id}") + base = AutoModelForCausalLM.from_pretrained( + base_model_id, token=token, trust_remote_code=True, + ) + tokenizer = AutoTokenizer.from_pretrained( + base_model_id, token=token, trust_remote_code=True, + ) + + print_color(YELLOW, f"Loading LoRA adapter: {lora_path}") + merged = PeftModel.from_pretrained(base, lora_path, token=token).merge_and_unload() + + out_dir = Path(tempfile.mkdtemp(prefix="cactus_lora_merged_")) + print_color(YELLOW, f"Saving merged model to: {out_dir}") + merged.save_pretrained(out_dir) + tokenizer.save_pretrained(out_dir) + + lora_tok = Path(lora_path) / "tokenizer_config.json" + if lora_tok.is_file(): + shutil.copy2(lora_tok, out_dir / "tokenizer_config.json") + + print_color(GREEN, "LoRA merge complete") + return str(out_dir) + + +def cmd_convert(args): + """Convert a HuggingFace model into a runnable cactus bundle: quantize weights to CQ, + then build the runtime graph. Pass --weights-only to stop after the weight conversion. + """ + from .model import ensure_weights + + source_model_id = args.model_id + merged_dir = None + + if args.lora: + merged_dir = _merge_lora_adapter(args.model_id, args.lora, token=args.token) + if merged_dir is None: + return 1 + source_model_id = merged_dir + + output_dir = args.output_dir or str( + get_bundle_dir(args.model_id, bits=args.bits) + ) + + try: + ensure_weights( + source_model_id, + bits=args.bits, + token=args.token, + reconvert=args.reconvert, + output_dir=output_dir, + skip_model_load=bool(getattr(args, "skip_model_load", False)), + ) + if getattr(args, "weights_only", False): + return 0 + args.weights_dir = args.weights_dir or output_dir + args.artifact_dir = args.artifact_dir or output_dir + return cmd_transpile(args) + except RuntimeError as e: + print_color(RED, f"Conversion error: {e}") + return 1 + finally: + if merged_dir: + shutil.rmtree(merged_dir, ignore_errors=True) + + +def cmd_transpile(args): + """Build the runtime graph bundle from already-converted CQ weights.""" + from .transpile import run_transpile + from .model import _default_multimodal_assets + from cactus.transpile.component_plan import infer_component_plan_from_output + + extra_args = [] + if args.weights_dir: + extra_args.extend(["--weights-dir", args.weights_dir]) + if args.task and args.task != "auto": + extra_args.extend(["--task", args.task]) + if args.prompt is not None: + extra_args.extend(["--prompt", args.prompt]) + if args.system_prompt is not None: + extra_args.extend(["--system-prompt", args.system_prompt]) + if args.enable_thinking: + extra_args.append("--enable-thinking") + if args.input_ids is not None: + extra_args.extend(["--input-ids", args.input_ids]) + + image_files = list(args.image_file or []) + audio_file = args.audio_file + weights_dir = Path(args.weights_dir).expanduser() if args.weights_dir else get_weights_dir(args.model_id) + from .model import _AUDIO_TASKS + plan = infer_component_plan_from_output(str(weights_dir), model_id=args.model_id) + default_needs_image = bool( + plan.needs_image if plan is not None else args.task == "multimodal_causal_lm_logits" + ) + default_needs_audio = bool( + plan.needs_audio if plan is not None + else (args.task in _AUDIO_TASKS or args.task == "multimodal_causal_lm_logits") + ) + if (default_needs_image and not image_files) or (default_needs_audio and not audio_file): + default_images, default_audio = _default_multimodal_assets() + if default_needs_image and not image_files: + image_files = default_images + if default_needs_audio and not audio_file and default_audio: + audio_file = default_audio + + for img in image_files: + extra_args.extend(["--image-file", img]) + if audio_file: + extra_args.extend(["--audio-file", audio_file]) + if args.max_new_tokens is not None: + extra_args.extend(["--max-new-tokens", str(args.max_new_tokens)]) + if args.component_pipeline and args.component_pipeline != "auto": + extra_args.extend(["--component-pipeline", args.component_pipeline]) + if args.components: + extra_args.extend(["--components", args.components]) + if args.torch_dtype: + extra_args.extend(["--torch-dtype", args.torch_dtype]) + if args.token: + extra_args.extend(["--token", args.token]) + if args.trust_remote_code: + extra_args.append("--trust-remote-code") + if args.local_files_only: + extra_args.append("--local-files-only") + if getattr(args, "low_memory_load", False): + extra_args.append("--low-memory-load") + if args.artifact_dir: + extra_args.extend(["--artifact-dir", args.artifact_dir]) + if args.graph_filename: + extra_args.extend(["--graph-filename", args.graph_filename]) + if args.skip_reference_compare: + extra_args.append("--skip-reference-compare") + if args.no_fuse_rms_norm: + extra_args.append("--no-fuse-rms-norm") + if args.no_fuse_rope: + extra_args.append("--no-fuse-rope") + if args.no_fuse_attention: + extra_args.append("--no-fuse-attention") + if args.no_fuse_attention_block: + extra_args.append("--no-fuse-attention-block") + if args.no_fuse_add_clipped: + extra_args.append("--no-fuse-add-clipped") + if args.no_fuse_gated_deltanet: + extra_args.append("--no-fuse-gated-deltanet") + if args.cache_context_length is not None: + extra_args.extend(["--cache-context-length", str(args.cache_context_length)]) + + return run_transpile( + args.model_id, + extra_args=extra_args, + execute_after_transpile=args.execute_after_transpile, + allow_unconverted_weights=args.allow_unconverted_weights, + ) diff --git a/python/cactus/cli/download.py b/python/cactus/cli/download.py new file mode 100644 index 000000000..110eded47 --- /dev/null +++ b/python/cactus/cli/download.py @@ -0,0 +1,61 @@ +"""Download prebuilt bundles from huggingface.co/Cactus-Compute.""" +from pathlib import Path + +from .common import ( + BLUE, GREEN, + print_color, weights_root, +) + + +def get_model_dir_name(model_id: str) -> str: + return model_id.split("/")[-1].lower() + + +def get_weights_dir(model_id: str) -> Path: + return weights_root() / get_model_dir_name(model_id) + + +def get_bundle_dir(model_id: str, *, bits: int | float = 4) -> Path: + from .utils import variant_suffix + return weights_root() / f"{get_model_dir_name(model_id)}-{variant_suffix(bits)}" + + +def ensure_model(model_id: str) -> Path: + return download_bundle(model_id) + + +def download_bundle(model_id: str, *, bits: int | float = 4, + token: str | None = None, output_dir: Path | None = None) -> Path: + from .utils import ( + download_cq_archive, + list_hf_cq_archives, + resolve_archive, + resolve_weight_revision, + suggested_cq_repo, + variant_suffix, + ) + + repo_id = suggested_cq_repo(model_id) + local_name = get_model_dir_name(model_id) + bundle_dir = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits) + + revision = resolve_weight_revision(repo_id, token=token) + label = variant_suffix(bits) + print() + print_color(BLUE, f"Fetching {repo_id} [{label}] @ {revision or 'main'}") + + archives = list_hf_cq_archives(repo_id, token=token, revision=revision) + if not archives: + raise RuntimeError(f"no bundles published at {repo_id}") + + resolution = resolve_archive(repo_id, local_name, archives, bits) + download_cq_archive(resolution, bundle_dir, token=token, revision=revision) + print_color(GREEN, f"Ready at {bundle_dir}") + return bundle_dir + + +def cmd_download(args) -> int: + from .model import prepare_bundle + + bundle = prepare_bundle(args, fail_prefix=f"Failed to prepare {args.model_id}") + return 1 if bundle is None else 0 diff --git a/python/cactus/cli/list.py b/python/cactus/cli/list.py new file mode 100644 index 000000000..973a9a767 --- /dev/null +++ b/python/cactus/cli/list.py @@ -0,0 +1,103 @@ +import itertools +import stat as _stat +import struct + +from .common import CYAN, print_color, transpiled_root, weights_root + +_PREC_TO_BITS = {3: 1, 4: 2, 5: 3, 6: 4} + + +def _read_config(path): + pairs = ( + line.split("=", 1) + for line in path.read_text(errors="ignore").splitlines() + if "=" in line and not line.lstrip().startswith("#") + ) + return {k.strip(): v.strip() for k, v in pairs} + + +def _human_size(n): + for unit in ("B", "KB", "MB", "GB"): + if n < 1024 or unit == "GB": + return f"{n:.1f} {unit}" + n /= 1024 + + +def _cq_precision(weights_file): + """Read the precision code at byte offset 48 of a Cactus tensor header.""" + try: + with weights_file.open("rb") as f: + head = f.read(52) + except OSError: + return None + if len(head) < 52 or head[:4] != b"CACT": + return None + return struct.unpack_from("= sample_cap or scanned >= scan_cap: + break + if not counts: + return "—" + return f"CQ{max(counts, key=counts.get)}" + + +def _dir_size(path): + total = 0 + for f in path.rglob("*"): + try: + st = f.lstat() + except OSError: + continue + if _stat.S_ISREG(st.st_mode): + total += st.st_size + return total + + +def _collect(roots): + models = [] + for root in roots: + if not root.is_dir(): + continue + for p in sorted(root.iterdir()): + if not p.is_dir() or not (p / "config.txt").is_file(): + continue + cfg = _read_config(p / "config.txt") + models.append((p, cfg.get("model_type", "?"), _quant_label(p), _dir_size(p))) + return models + + +def cmd_list(_args): + models = _collect((weights_root(), transpiled_root())) + print_color(CYAN, "Available models") + if not models: + print(" (none)") + return 0 + name_w = max(len(p.name) for p, _, _, _ in models) + type_w = max(len("type"), max(len(t) for _, t, _, _ in models)) + quant_w = max(len("quant"), max(len(q) for _, _, q, _ in models)) + print(f" {'name':<{name_w}} {'type':<{type_w}} {'quant':<{quant_w}} {'size':>10} location") + for p, model_type, quant, size in models: + print(f" {p.name:<{name_w}} {model_type:<{type_w}} {quant:<{quant_w}} {_human_size(size):>10} {p.parent}") + return 0 diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py new file mode 100644 index 000000000..1a8a4a508 --- /dev/null +++ b/python/cactus/cli/model.py @@ -0,0 +1,391 @@ +"""Model resolution, weight management, and bundle preparation.""" +from __future__ import annotations + +import os +import shutil +from dataclasses import dataclass +from pathlib import Path + +from .common import GREEN, PROJECT_ROOT, RED, YELLOW, print_color + + + + +def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load=False): + """Download from HuggingFace and run CQ conversion.""" + if bits not in (1, 2, 3, 4): + raise SystemExit( + f"CQ{bits} is a mixed-precision variant, available only as a prebuilt " + f"download; local conversion supports uniform bits 1-4" + ) + from .common import convert_toolchain_error + err = convert_toolchain_error() + if err: + raise RuntimeError(err) + print_color(YELLOW, f"Converting {model_id} from HuggingFace source...") + from ..convert.cli import main as cq_main + + cq_args = [ + "convert", "--model", model_id, + "--out", str(weights_dir), + "--bits", str(bits), + ] + if skip_model_load: + cq_args.append("--skip-model-load") + if token: + os.environ["HF_TOKEN"] = token + os.environ["HUGGING_FACE_HUB_TOKEN"] = token + cq_main(cq_args) + + print_color(GREEN, f"Model converted and ready at {weights_dir}") + return weights_dir + + +def ensure_weights(model_id, *, bits=4, token=None, reconvert=False, output_dir=None, skip_model_load=False): + from .download import get_bundle_dir + + weights_dir = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits) + + if reconvert and weights_dir.exists(): + print_color(YELLOW, "Removing cached weights for reconversion...") + shutil.rmtree(weights_dir) + + if weights_dir.exists() and (weights_dir / "config.txt").exists(): + print_color(GREEN, f"Model weights found at {weights_dir}") + return weights_dir + + if weights_dir.exists(): + print_color(YELLOW, "Removing incomplete weights from a previous run...") + shutil.rmtree(weights_dir) + + return _convert_from_source( + model_id, + bits=bits, + token=token, + weights_dir=weights_dir, + skip_model_load=skip_model_load, + ) + + + +_DEFAULT_MULTIMODAL_PROMPT = ( + "Respond with 2 lines. The first should be a description of the image, " + "and the second should be a transcription of the audio" +) +_DEFAULT_TEXT_PROMPT = "Hello" + + +@dataclass(frozen=True) +class _TranspileSpec: + task: str + components: tuple[str, ...] = () + default_max_new_tokens: int | None = None + needs_image: bool = False + needs_audio: bool = False + force_component_pipeline: bool = False + + +def _spec_from_plan(plan): + return _TranspileSpec( + task=plan.task, + components=tuple(plan.components or ()), + default_max_new_tokens=plan.default_max_new_tokens, + needs_image=bool(plan.needs_image), + needs_audio=bool(plan.needs_audio), + force_component_pipeline=bool(plan.force_component_pipeline), + ) + + +def _infer_transpile_spec(*, task, plan): + if task != "auto": + if plan is not None and task == plan.task: + return _spec_from_plan(plan) + return _TranspileSpec( + task=task, + needs_image=task == "multimodal_causal_lm_logits", + needs_audio=task in { + "tdt_transcription", "seq2seq_transcription", + "ctc_logits", "encoder_hidden_states", + "multimodal_causal_lm_logits", + }, + force_component_pipeline=task in { + "tdt_transcription", "seq2seq_transcription", + "multimodal_causal_lm_logits", + }, + ) + + if plan is None: + return _TranspileSpec(task="causal_lm_logits") + + return _spec_from_plan(plan) + + +def _default_max_new_tokens(spec): + if spec.default_max_new_tokens is not None: + return int(spec.default_max_new_tokens) + return { + "seq2seq_transcription": 128, + "multimodal_causal_lm_logits": 512, + "causal_lm_logits": 128, + }.get(spec.task, 32) + + +def _default_multimodal_assets(): + """Return bundled test image/audio paths for multimodal shape capture.""" + candidates = ( + Path(__file__).resolve().parent.parent / "assets", + PROJECT_ROOT / "cactus-engine" / "tests" / "assets", + ) + def _find(name): + return next((d / name for d in candidates if (d / name).exists()), None) + image = _find("test_monkey.png") + audio = _find("test.wav") + return ([str(image)] if image else []), (str(audio) if audio else None) + + +def _default_audio_asset(): + _, audio = _default_multimodal_assets() + return audio + + +def _remove_stale_transpile_artifacts(output_dir): + for relative in ( + "components", + "transpile_entrypoints.json", + "raw_ir.json", + "optimized_ir.json", + "graph.cactus", + "graph_bindings.json", + "result.json", + ): + path = output_dir / relative + if path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + for pattern in ("raw_ir_*.json", "optimized_ir_*.json"): + for path in output_dir.glob(pattern): + if path.is_file(): + path.unlink() + + +def _has_transpiled_bundle(path): + return (path / "components" / "manifest.json").exists() + + +_AUDIO_TASKS = frozenset({ + "tdt_transcription", "seq2seq_transcription", + "ctc_logits", "encoder_hidden_states", +}) + + + + +def resolve_bundle_dir(model_id): + path = Path(model_id).expanduser() + if path.is_dir() and _has_transpiled_bundle(path): + return path + return None + + +@dataclass(frozen=True) +class TranspileOptions: + task: str = "auto" + prompt: str | None = None + image_files: list[str] | None = None + audio_file: str | None = None + max_new_tokens: int | None = None + component_pipeline: str = "auto" + components: str | None = None + system_prompt: str | None = None + trust_remote_code: bool = False + local_files_only: bool = False + cache_context_length: str | int | None = None + + +def ensure_runnable_bundle(model_id, *, bits=4, token=None, + reconvert=False, prebuilt=True, output_dir=None, + transpile=None): + """Resolve a runnable bundle via the full fallback ladder, building if needed. + + Rungs, in order: (1) a local bundle path, (2) a cached prior build, + (3) a prebuilt bundle on HuggingFace, (4) local convert + transpile. + With prebuilt=False, rung (3) is skipped and the bundle is always built + locally when not already present. + Raises RuntimeError if every rung fails. + """ + from .download import download_bundle, get_bundle_dir + + local = resolve_bundle_dir(model_id) + if local is not None: + return local + + if str(model_id).startswith(("/", "./", "../", "~")) and not Path(model_id).expanduser().exists(): + raise RuntimeError(f"path not found: {model_id}") + + cached = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits) + if _has_transpiled_bundle(cached) and not reconvert: + return cached + + if prebuilt and not reconvert: + try: + return download_bundle(model_id, bits=bits, + token=token, output_dir=cached) + except (RuntimeError, OSError) as exc: + print_color(YELLOW, f"No prebuilt bundle ({exc}); building locally") + + opts = transpile or TranspileOptions() + return ensure_bundle(model_id, bits=bits, token=token, + reconvert=reconvert, output_dir=cached, transpile=opts) + + +def prepare_bundle(args, *, model_id=None, transpile=None, prebuilt=True, + output_dir=None, fail_prefix="Model setup failed"): + """Return a runnable bundle, with uniform error handling shared by every + model command. Returns the bundle Path, or None (after printing the error) + on failure.""" + try: + return ensure_runnable_bundle( + args.model_id if model_id is None else model_id, + bits=getattr(args, "bits", 4), + token=getattr(args, "token", None), + reconvert=getattr(args, "reconvert", False), + prebuilt=prebuilt, + output_dir=output_dir, + transpile=transpile, + ) + except (RuntimeError, OSError, ValueError) as exc: + print_color(RED, f"{fail_prefix}: {exc}") + return None + + +def ensure_bundle(model_id, *, bits=4, token=None, + reconvert=False, output_dir=None, transpile=None, + skip_model_load=False): + from .common import convert_toolchain_error + err = convert_toolchain_error() + if err: + raise RuntimeError(err) + from .download import get_bundle_dir + from .transpile import run_transpile + from cactus.transpile.component_plan import infer_component_plan_from_output + + opts = transpile or TranspileOptions() + + if output_dir is not None: + output_dir = Path(output_dir).expanduser().resolve() + else: + output_dir = get_bundle_dir(model_id, bits=bits) + + ensure_weights( + model_id, bits=bits, token=token, + reconvert=reconvert, output_dir=output_dir, + skip_model_load=skip_model_load, + ) + + if _has_transpiled_bundle(output_dir): + return output_dir + + plan = infer_component_plan_from_output(str(output_dir), model_id=model_id) + spec = _infer_transpile_spec(task=opts.task, plan=plan) + _remove_stale_transpile_artifacts(output_dir) + + spec_prompt = opts.prompt + spec_image_files = list(opts.image_files or []) + spec_audio_file = opts.audio_file + + if spec_prompt is None and spec.task == "multimodal_causal_lm_logits": + spec_prompt = _DEFAULT_MULTIMODAL_PROMPT + elif spec_prompt is None and spec.task == "causal_lm_logits": + spec_prompt = _DEFAULT_TEXT_PROMPT + + effective_component_pipeline = opts.component_pipeline + effective_components = opts.components + + if spec.task == "multimodal_causal_lm_logits": + needs_image = spec.needs_image + needs_audio = spec.needs_audio + if not needs_image and not needs_audio: + needs_image = bool(spec_image_files) + needs_audio = bool(spec_audio_file) + if (needs_image and not spec_image_files) or (needs_audio and not spec_audio_file): + default_images, default_audio = _default_multimodal_assets() + if needs_image and not spec_image_files: + spec_image_files = default_images + if needs_audio and not spec_audio_file: + spec_audio_file = default_audio + print_color( + YELLOW, + "Multimodal transpile needs representative media shapes; " + "using bundled tiny test assets.", + ) + if needs_image and not spec_image_files: + raise RuntimeError("Building this multimodal model requires --image-file.") + if needs_audio and not spec_audio_file: + raise RuntimeError("Building this multimodal model requires --audio-file.") + + if effective_component_pipeline == "auto" and spec.force_component_pipeline: + effective_component_pipeline = "on" + if effective_components is None and spec.components: + effective_components = ",".join(spec.components) + + used_default_audio = False + if spec.task in _AUDIO_TASKS and not spec_audio_file: + spec_audio_file = _default_audio_asset() + used_default_audio = spec_audio_file is not None + if spec.task in _AUDIO_TASKS and used_default_audio: + print_color( + YELLOW, + f"{spec.task} transpile needs a representative audio shape; " + "using bundled tiny test audio asset.", + ) + elif spec.task in _AUDIO_TASKS and not spec_audio_file: + raise RuntimeError(f"Building a {spec.task} model requires --audio-file.") + + effective_max_new_tokens = opts.max_new_tokens or _default_max_new_tokens(spec) + + extra_args = [ + "--weights-dir", str(output_dir), + "--artifact-dir", str(output_dir), + "--task", spec.task, + "--max-new-tokens", str(effective_max_new_tokens), + "--component-pipeline", effective_component_pipeline, + ] + if spec_prompt is not None: + extra_args.extend(["--prompt", spec_prompt]) + if effective_components: + extra_args.extend(["--components", str(effective_components)]) + for img in spec_image_files: + extra_args.extend(["--image-file", img]) + if spec_audio_file: + extra_args.extend(["--audio-file", str(spec_audio_file)]) + if opts.system_prompt: + extra_args.extend(["--system-prompt", str(opts.system_prompt)]) + if token: + extra_args.extend(["--token", token]) + if opts.trust_remote_code or spec.task == "multimodal_causal_lm_logits": + extra_args.append("--trust-remote-code") + if opts.local_files_only: + extra_args.append("--local-files-only") + if opts.cache_context_length is not None: + extra_args.extend(["--cache-context-length", str(opts.cache_context_length)]) + + rc = run_transpile(model_id, extra_args=extra_args) + if rc != 0: + raise RuntimeError(f"Build failed for {model_id}") + + try: + from cactus.convert.handoff_probe import ( + export_gemma4_handoff_probe, + export_parakeet_handoff_probe, + ) + + if export_gemma4_handoff_probe(output_dir, model_id=model_id): + print_color(GREEN, f"Gemma4 cloud handoff probe packaged into {output_dir}") + if export_parakeet_handoff_probe(output_dir, model_id=model_id): + print_color(GREEN, f"Parakeet cloud handoff probe packaged into {output_dir}") + except Exception as e: + print_color(YELLOW, f"Warning: failed to package cloud handoff probe: {e}") + + print_color(GREEN, f"Model built at {output_dir}") + return output_dir diff --git a/python/cactus/cli/run.py b/python/cactus/cli/run.py new file mode 100644 index 000000000..38078a3a8 --- /dev/null +++ b/python/cactus/cli/run.py @@ -0,0 +1,54 @@ +from pathlib import Path + +from .common import GREEN, RED, apply_runtime_env, launch_binary, print_color + + +def cmd_run(args) -> int: + apply_runtime_env(args) + + if args.image: + args.image = str(Path(args.image).expanduser()) + if not Path(args.image).is_file(): + print_color(RED, f"Image not found: {args.image}") + return 1 + if args.audio: + args.audio = str(Path(args.audio).expanduser()) + if not Path(args.audio).is_file(): + print_color(RED, f"Audio not found: {args.audio}") + return 1 + if args.result_json: + args.result_json = str(Path(args.result_json).expanduser()) + + from .model import TranspileOptions, prepare_bundle + + bundle_dir = prepare_bundle(args, transpile=TranspileOptions( + image_files=[args.image] if args.image else None, + audio_file=args.audio, + )) + if bundle_dir is None: + return 1 + + cmd = [str(bundle_dir)] + for flag, value in ( + ("--system", args.system), + ("--prompt", args.prompt), + ("--image", args.image), + ("--audio", args.audio), + ("--input-ids", args.input_ids), + ("--input-ids-file", args.input_ids_file), + ("--max-new-tokens", args.max_new_tokens), + ("--result-json", args.result_json), + ("--confidence-threshold", args.confidence_threshold), + ("--cloud-timeout-ms", args.cloud_timeout_ms), + ("--backend", args.backend), + ): + if value is not None: + cmd.extend([flag, str(value)]) + if args.thinking: + cmd.append("--thinking") + if args.no_cloud_handoff: + cmd.append("--no-cloud-handoff") + + print_color(GREEN, f"Running: {bundle_dir}") + print() + return launch_binary("run", *cmd) diff --git a/python/cactus/cli/runtime.py b/python/cactus/cli/runtime.py new file mode 100644 index 000000000..95cd3f2c3 --- /dev/null +++ b/python/cactus/cli/runtime.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import platform +import shutil +import subprocess +from pathlib import Path + +from .common import PROJECT_ROOT, YELLOW, print_color + + +def _static_library_path(): + return PROJECT_ROOT / "cactus-engine" / "build" / "libcactus_engine.a" + + +def ensure_library(): + lib = _static_library_path() + if lib.exists(): + return lib + + build_script = PROJECT_ROOT / "cactus-engine" / "build.sh" + if subprocess.run([str(build_script)], cwd=build_script.parent).returncode != 0: + raise RuntimeError("Failed to build the Cactus static runtime") + return lib + + +def _python_runtime_library_path(): + suffix = ".dylib" if platform.system() == "Darwin" else ".so" + bundled = Path(__file__).resolve().parent.parent / "bindings" / "lib" / f"libcactus_engine{suffix}" + if bundled.exists(): + return bundled + return PROJECT_ROOT / "cactus-engine" / "build" / f"libcactus_engine{suffix}" + + +def _public_cactus_api_symbols(static_lib): + cmd = ( + ["nm", "-gU", str(static_lib)] if platform.system() == "Darwin" + else ["nm", "-g", "--defined-only", str(static_lib)] + ) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"`{' '.join(cmd)}` failed: {result.stderr.strip()}") + + def is_cactus_api(symbol): + name = symbol[1:] if symbol.startswith("_") else symbol + return name.startswith("cactus_") + + symbols = sorted({ + parts[-1] for line in result.stdout.splitlines() + if (parts := line.split()) and is_cactus_api(parts[-1]) + }) + if not symbols: + raise RuntimeError(f"No public cactus_* symbols in {static_lib}") + return symbols + + +def _link_python_runtime_library(*, static_lib, library_path): + build_dir = library_path.parent + build_dir.mkdir(parents=True, exist_ok=True) + if library_path.exists(): + library_path.unlink() + + compiler_name = "clang++" if platform.system() == "Darwin" else "g++" + compiler = shutil.which(compiler_name) + if compiler is None: + raise RuntimeError(f"{compiler_name} is not installed; cannot link libcactus_engine") + + exported_symbols = _public_cactus_api_symbols(static_lib) + if platform.system() == "Darwin": + command = [ + compiler, + "-dynamiclib", + "-o", str(library_path), + *[f"-Wl,-u,{s}" for s in exported_symbols], + str(static_lib), + "-Wl,-install_name,@rpath/libcactus_engine.dylib", + "-lcurl", + "-framework", "Accelerate", + "-framework", "Foundation", + "-framework", "Metal", + "-framework", "MetalPerformanceShaders", + "-framework", "Security", + "-framework", "SystemConfiguration", + "-framework", "CFNetwork", + ] + else: + command = [ + compiler, + "-shared", + "-o", str(library_path), + *[f"-Wl,--undefined={s}" for s in exported_symbols], + str(static_lib), + "-lcurl", "-pthread", "-ldl", "-lm", + ] + + if subprocess.run(command, cwd=build_dir).returncode != 0: + raise RuntimeError(f"Failed to link the Cactus shared runtime: {library_path}") + + +def ensure_python_runtime_library(): + library_path = _python_runtime_library_path() + static_lib = _static_library_path() + + if ( + library_path.exists() + and (not static_lib.exists() + or library_path.stat().st_mtime >= static_lib.stat().st_mtime) + ): + return library_path + + print_color(YELLOW, "Preparing Cactus shared runtime...") + if not static_lib.exists(): + static_lib = ensure_library() + _link_python_runtime_library(static_lib=static_lib, library_path=library_path) + return library_path diff --git a/python/cactus/cli/serve.py b/python/cactus/cli/serve.py new file mode 100644 index 000000000..f244743de --- /dev/null +++ b/python/cactus/cli/serve.py @@ -0,0 +1,90 @@ +from pathlib import Path + +from .common import BLUE, GREEN, RED, YELLOW, apply_runtime_env, is_valid_bundle, print_color, weights_root +from .download import get_weights_dir + + +def _resolve_model_arg(model: str | None) -> tuple[Path | None, str | None]: + if not model: + return None, None + path = Path(model).expanduser() + if path.is_dir(): + return path, path.name + candidate = weights_root() / model + if candidate.is_dir(): + return candidate, candidate.name + hf_candidate = get_weights_dir(model) + if hf_candidate.is_dir(): + return hf_candidate, hf_candidate.name + return None, model + + +def _ensure_engine_library() -> bool: + from .runtime import ensure_python_runtime_library + + try: + ensure_python_runtime_library() + return True + except Exception as exc: + print_color(RED, f"Error: could not build the Cactus engine: {exc}") + print_color(YELLOW, "Build it manually with `cactus build --python` (needs cmake and a C++ toolchain).") + return False + + +def cmd_serve(args): + """Start the OpenAI-compatible HTTP server.""" + apply_runtime_env(args) + if not _ensure_engine_library(): + return 1 + model_path, model_name = _resolve_model_arg(args.model_id) + if args.model_id and model_path is None: + from .model import prepare_bundle + built = prepare_bundle(args, fail_prefix=f"Error: could not prepare {args.model_id}") + if built is None: + return 1 + model_path, model_name = built, built.name + if model_path is not None and not is_valid_bundle(model_path): + print_color(RED, f"Error: not a valid v2 Cactus bundle: {model_path}") + print("Expected config.txt and components/manifest.json.") + return 1 + + try: + import uvicorn + except ImportError: + print_color(RED, "Error: uvicorn not installed. Install the serve extra or run `pip install fastapi uvicorn python-multipart`.") + return 1 + + backend = getattr(args, "backend", None) + if backend: + from ..bindings.cactus import cactus_set_backend + if cactus_set_backend(backend) != 0: + print_color(YELLOW, f"Backend '{backend}' not available; using auto") + + try: + from ..server import create_app + except ImportError: + print_color(RED, "Error: server dependencies not installed. Install the serve extra or run `pip install fastapi uvicorn python-multipart`.") + return 1 + + try: + application = create_app( + weights_root=weights_root(), + model_path=model_path, + default_model=model_name, + auto_handoff=not args.no_cloud_handoff, + confidence_threshold=args.confidence_threshold, + cloud_timeout_ms=args.cloud_timeout_ms, + ) + except RuntimeError as exc: + print_color(RED, f"Error: {exc}") + print("Prepare a bundle first with `cactus run ` (or `cactus convert `).") + return 1 + + models = sorted(application.state.registry.models) + print_color(GREEN, f"Available models: {', '.join(models)}") + if args.host not in {"127.0.0.1", "localhost", "::1"}: + print_color(YELLOW, f"Warning: binding to {args.host} exposes the server beyond loopback; no auth is enforced.") + print_color(BLUE, f"Starting server on {args.host}:{args.port}") + uvicorn.run(application, host=args.host, port=args.port, log_level="info", + access_log=not getattr(args, "no_access_log", False)) + return 0 diff --git a/python/cactus/cli/test.py b/python/cactus/cli/test.py new file mode 100644 index 000000000..720f9268b --- /dev/null +++ b/python/cactus/cli/test.py @@ -0,0 +1,112 @@ +import os +import subprocess + +from .common import ( + BLUE, DEFAULT_TEST_MODEL_ID, DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, + PROJECT_ROOT, RED, YELLOW, apply_cloud_api_key_env, print_color, +) + +COMPONENTS = ("kernels", "graph", "engine", "all") + + +def _list_component_suites(component): + tests_dir = PROJECT_ROOT / f"cactus-{component}" / "tests" + if not tests_dir.exists(): + return [] + return sorted( + f.stem.removeprefix("test_") + for f in tests_dir.glob("test_*.cpp") + if f.stem != "test_utils" + ) + + +def _components_with_suite(suite): + return tuple(c for c in ("kernels", "graph", "engine") if suite in _list_component_suites(c)) + + +def _component_args(component, args): + cmd = [str(PROJECT_ROOT / f"cactus-{component}" / "test.sh")] + if args.suite: + cmd.extend(["--suite", args.suite]) + if component == "engine": + cmd.extend(["--model", args.model_id]) + cmd.extend(["--transcription-model", args.transcription_model_id]) + if getattr(args, "backend", None): + cmd.extend(["--backend", args.backend]) + if args.android: + cmd.append("--android") + if args.ios: + cmd.append("--ios") + return cmd + + +def cmd_benchmark(args): + args.component = "engine" + args.suite = "benchmark" + args.list = False + return cmd_test(args) + + +def cmd_test(args): + if getattr(args, "list", False): + print_color(BLUE, "Components:") + for c in COMPONENTS: + print(f" {c}") + print_color(BLUE, "\nSuites by component:") + for c in ("kernels", "graph", "engine"): + suites = _list_component_suites(c) + if suites: + print(f" {c}: {', '.join(suites)}") + return 0 + + if args.suite: + matches = _components_with_suite(args.suite) + if not matches: + print_color(RED, f"unknown suite '{args.suite}'.") + print_color(RED, "Run `cactus test --list` to see available suites.") + return 2 + if args.component != "all" and args.component not in matches: + print_color(RED, + f"suite '{args.suite}' does not exist in component '{args.component}'. " + f"Found in: {', '.join(matches)}.") + return 2 + targets = matches if args.component == "all" else (args.component,) + else: + targets = ("kernels", "graph", "engine") if args.component == "all" else (args.component,) + + if "engine" in targets: + from .model import prepare_bundle + model_dir = prepare_bundle(args, model_id=args.model_id or DEFAULT_TEST_MODEL_ID, + fail_prefix="Model setup failed") + if model_dir is None: + return 1 + args.model_id = str(model_dir) + stt_dir = prepare_bundle(args, model_id=args.transcription_model_id or DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, + fail_prefix="Model setup failed") + if stt_dir is None: + return 1 + args.transcription_model_id = str(stt_dir) + + apply_cloud_api_key_env() + env = os.environ.copy() + if args.enable_telemetry: + env.pop("CACTUS_NO_CLOUD_TELE", None) + else: + env["CACTUS_NO_CLOUD_TELE"] = "1" + + for c in targets: + cmd = _component_args(c, args) + rc = subprocess.run(cmd, cwd=PROJECT_ROOT, env=env).returncode + if rc != 0: + if c == "engine" and args.ios: + print_color(YELLOW, + "If the app could not build or launch on your device:\n" + "Turn on Developer Mode:\n" + " Go to Settings → Privacy & Security → Developer Mode.\n" + " Toggle it on, restart, then confirm with your passcode.\n" + "Trust the developer:\n" + " Go to Settings → General → VPN & Device Management.\n" + " Under Enterprise App, tap the developer name.\n" + " Tap Trust “[developer name]”.") + return rc + return 0 diff --git a/python/cactus/cli/transcribe.py b/python/cactus/cli/transcribe.py new file mode 100644 index 000000000..a4d00acea --- /dev/null +++ b/python/cactus/cli/transcribe.py @@ -0,0 +1,31 @@ +from pathlib import Path + +from .common import apply_runtime_env, launch_binary, print_color, RED, GREEN + + +def cmd_transcribe(args): + from .model import prepare_bundle, TranspileOptions + + if args.audio_file: + audio_path = Path(args.audio_file).expanduser() + if not audio_path.is_file(): + print_color(RED, f"Audio file not found: {audio_path}") + return 1 + args.audio_file = str(audio_path) + + apply_runtime_env(args) + + bundle_dir = prepare_bundle(args, transpile=TranspileOptions(audio_file=args.audio_file)) + if bundle_dir is None: + return 1 + + cmd = [str(bundle_dir)] + if args.audio_file: + cmd.append(args.audio_file) + if args.language: + cmd.extend(["--language", args.language]) + + print_color(GREEN, f"Starting Cactus transcription with model: {args.model_id}") + print() + + return launch_binary("transcribe", *cmd) diff --git a/python/cactus/cli/transpile.py b/python/cactus/cli/transpile.py new file mode 100644 index 000000000..05f77711d --- /dev/null +++ b/python/cactus/cli/transpile.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +from .common import PROJECT_ROOT, RED, YELLOW, print_color +from .download import get_weights_dir + + +def _weights_dir_looks_transpile_ready(weights_dir): + root = Path(weights_dir).expanduser().resolve() + if not root.is_dir(): + return False + if (root / "weights_manifest.json").exists(): + return True + if any(root.glob("*.cq[1-4].weights")): + return True + return (root / "config.txt").exists() and any(root.glob("*.weights")) + + +def _extra_args_has_option(extra_args, option): + prefix = f"{option}=" + return any(arg == option or arg.startswith(prefix) for arg in extra_args) + + +def _prepend_python_path(env): + python_root = str(PROJECT_ROOT / "python") + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = python_root if not existing else f"{python_root}{os.pathsep}{existing}" + + +def run_transpile(model_id, *, extra_args, execute_after_transpile=False, + allow_unconverted_weights=False): + from .common import convert_toolchain_error + from .runtime import ensure_python_runtime_library + + err = convert_toolchain_error() + if err: + print_color(RED, f"Error: {err}") + return 1 + + extra_args = list(extra_args or []) + command = [sys.executable, "-m", "cactus.transpile.hf_model", "--model-id", model_id] + default_weights_dir = get_weights_dir(model_id) + if not default_weights_dir.name: + print_color(RED, f"Error: cannot derive a unique output dir from model_id {model_id!r}.") + print_color(YELLOW, "Pass --weights-dir and --artifact-dir explicitly.") + return 1 + if not _extra_args_has_option(extra_args, "--weights-dir"): + if _weights_dir_looks_transpile_ready(default_weights_dir): + command.extend(["--weights-dir", str(default_weights_dir)]) + if not _extra_args_has_option(extra_args, "--artifact-dir"): + command.extend(["--artifact-dir", str(default_weights_dir)]) + + if allow_unconverted_weights: + command.append("--allow-unconverted-weights") + if not execute_after_transpile and "--skip-execute" not in extra_args: + command.append("--skip-execute") + command.extend(extra_args) + + try: + transpile_lib = ensure_python_runtime_library() + except RuntimeError as exc: + print_color(RED, f"Error: {exc}") + return 1 + + env = os.environ.copy() + env["CACTUS_LIB_PATH"] = str(transpile_lib) + _prepend_python_path(env) + return subprocess.run(command, cwd=PROJECT_ROOT, env=env).returncode diff --git a/python/cactus/cli/upload.py b/python/cactus/cli/upload.py new file mode 100644 index 000000000..f38339659 --- /dev/null +++ b/python/cactus/cli/upload.py @@ -0,0 +1,52 @@ +import os +import tempfile +import zipfile +from pathlib import Path + +from .common import BLUE, GREEN, PROJECT_ROOT, print_color +from .download import get_model_dir_name +from .model import prepare_bundle +from .utils import suggested_cq_repo, variant_suffix + + +def cmd_upload(args): + """Build a runnable bundle locally and upload it to Cactus-Compute on HuggingFace.""" + token = args.token or os.getenv("HF_TOKEN") + if not token: + raise SystemExit("missing HuggingFace token: pass --token or set HF_TOKEN") + + repo_id = suggested_cq_repo(args.model_id) + stem = f"{get_model_dir_name(args.model_id)}-{variant_suffix(args.bits)}" + archive_name = f"{stem}.zip" + tag = (PROJECT_ROOT / "CACTUS_VERSION").read_text(encoding="utf-8").strip() + + bundle_dir = prepare_bundle(args, prebuilt=False, fail_prefix=f"Build failed for {args.model_id}") + if bundle_dir is None: + return 1 + + from huggingface_hub import HfApi + + api = HfApi(token=token) + api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True) + + with tempfile.TemporaryDirectory(prefix="cactus-upload-") as tmp: + archive_path = Path(tmp) / archive_name + with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: + for path in sorted(p for p in bundle_dir.rglob("*") if p.is_file()): + zf.write(path, f"{stem}/{path.relative_to(bundle_dir).as_posix()}") + print_color(BLUE, f"Uploading {archive_name} ({archive_path.stat().st_size / 1e6:.1f} MB) -> {repo_id}") + api.upload_file( + path_or_fileobj=str(archive_path), + path_in_repo=archive_name, + repo_id=repo_id, + repo_type="model", + commit_message=f"Upload {archive_name} ({tag})", + ) + + try: + api.delete_tag(repo_id=repo_id, tag=tag, repo_type="model") + except Exception: + pass + api.create_tag(repo_id=repo_id, tag=tag, repo_type="model", revision="main") + print_color(GREEN, f"Done: {repo_id} @ {tag}") + return 0 diff --git a/python/cactus/cli/utils.py b/python/cactus/cli/utils.py new file mode 100644 index 000000000..82d374a92 --- /dev/null +++ b/python/cactus/cli/utils.py @@ -0,0 +1,424 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import tarfile +import tempfile +import urllib.parse +import urllib.request +import zipfile +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + + +CQ_VARIANT_RE = re.compile(r"-cq(\d+(?:\.\d+)?)$", re.IGNORECASE) +ARCHIVE_SUFFIXES = (".zip", ".tar", ".tar.gz", ".tgz") + +ALLOWED_BITS = (1, 2, 3, 4, 2.54, 3.26) + + +def normalize_bits(bits: str | int | float) -> int | float: + value = float(bits) + return int(value) if value == int(value) else value + + +def variant_suffix(bits: int | float) -> str: + return f"cq{normalize_bits(bits)}" + + +def bits_arg(value: str) -> int | float: + try: + bits = normalize_bits(value) + except (TypeError, ValueError): + raise argparse.ArgumentTypeError(f"invalid --bits value {value!r}") + if bits not in ALLOWED_BITS: + allowed = ", ".join(str(b) for b in ALLOWED_BITS) + raise argparse.ArgumentTypeError(f"unsupported --bits {value!r}; choose from: {allowed}") + return bits + + +@dataclass(frozen=True) +class CqArchive: + filename: str + bits: int | float + size: int | None = None + sha256: str | None = None + + @property + def suffix(self) -> str: + return variant_suffix(self.bits) + + +@dataclass(frozen=True) +class CqResolution: + repo_id: str + local_name: str + archive: CqArchive + available: tuple[CqArchive, ...] + + +def suggested_cq_repo(model_id: str) -> str: + name = str(model_id).strip().replace("_", "-").split("/")[-1] + name = re.sub(r"-cq\d*(?:\.\d+)?$", "", name, flags=re.IGNORECASE) + return f"Cactus-Compute/{name}" + + +def archive_stem(filename: str) -> str: + for suffix in ARCHIVE_SUFFIXES: + if filename.lower().endswith(suffix): + return filename[: -len(suffix)] + return filename + + +def parse_cq_variant(filename: str) -> int | float | None: + stem = archive_stem(Path(filename).name) + m = CQ_VARIANT_RE.search(stem) + if not m: + return None + return normalize_bits(m.group(1)) + + +def is_supported_archive(filename: str) -> bool: + lower = filename.lower() + return any(lower.endswith(suffix) for suffix in ARCHIVE_SUFFIXES) + + +def archives_from_repo_files(repo_files: Iterable[str], sizes: dict[str, int] | None = None, + sha256s: dict[str, str] | None = None) -> tuple[CqArchive, ...]: + sizes = sizes or {} + sha256s = sha256s or {} + archives = [] + for filename in repo_files: + if not is_supported_archive(filename): + continue + bits = parse_cq_variant(filename) + if bits is None: + continue + archives.append(CqArchive(filename=filename, bits=bits, + size=sizes.get(filename), sha256=sha256s.get(filename))) + return tuple(sorted(archives, key=lambda a: (isinstance(a.bits, float), a.bits))) + + +VERSION_TAG_RE = re.compile(r"^v?(\d+(?:\.\d+){0,2})$", re.IGNORECASE) + + +def parse_version(text: str | None) -> tuple[int, int, int] | None: + if not text: + return None + m = VERSION_TAG_RE.match(str(text).strip()) + if not m: + return None + parts = [int(p) for p in m.group(1).split(".")] + parts += [0] * (3 - len(parts)) + return tuple(parts[:3]) + + +def list_version_tags(repo_id: str, *, token=None) -> tuple[tuple[str, tuple[int, int, int]], ...]: + try: + from huggingface_hub import HfApi + refs = HfApi().list_repo_refs(repo_id, token=token) + names = [t.name for t in refs.tags] + except ImportError: + quoted = urllib.parse.quote(repo_id, safe="/") + url = f"https://huggingface.co/api/models/{quoted}/refs" + headers = {"User-Agent": "cactus-cq-downloader"} + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=30) as response: + data = json.load(response) + names = [t.get("name", "") for t in data.get("tags", [])] + out = [] + for name in names: + version = parse_version(name) + if version is not None: + out.append((name, version)) + return tuple(out) + + +def resolve_weight_revision(repo_id: str, runtime_version: str | None = None, + *, token=None) -> str | None: + if runtime_version is None: + from .. import __version__ as runtime_version + runtime = parse_version(runtime_version) + if runtime is None: + return None + try: + tags = list_version_tags(repo_id, token=token) + except Exception: + return None + eligible = [(name, version) for name, version in tags if version <= runtime] + if not eligible: + return None + return max(eligible, key=lambda item: item[1])[0] + + +def resolve_archive(repo_id: str, local_name: str, archives: Iterable[CqArchive], + bits: int | float) -> CqResolution: + available = tuple(archives) + if not available: + raise RuntimeError(f"No CQ archives found in {repo_id}") + + wanted = variant_suffix(bits) + match = next((a for a in available if a.bits == bits), None) + if not match: + choices = ", ".join(a.suffix for a in available) + raise RuntimeError(f"{wanted} not found in {repo_id}. Available: {choices}") + + return CqResolution( + repo_id=repo_id, + local_name=local_name, + archive=match, + available=available, + ) + + +def _safe_zip_extract(zip_path: Path, out_dir: Path) -> None: + out_dir = out_dir.resolve() + with zipfile.ZipFile(zip_path, "r") as zf: + for info in zf.infolist(): + target = (out_dir / info.filename).resolve() + if target != out_dir and out_dir not in target.parents: + raise RuntimeError(f"Unsafe path in zip archive: {info.filename}") + mode = (info.external_attr >> 16) & 0o170000 + if mode in (0o120000, 0o10000): + raise RuntimeError(f"Refusing unsafe archive member: {info.filename}") + zf.extractall(out_dir) + + +def _safe_tar_extract(tar_path: Path, out_dir: Path) -> None: + out_dir = out_dir.resolve() + with tarfile.open(tar_path, "r:*") as tf: + for member in tf.getmembers(): + target = (out_dir / member.name).resolve() + if target != out_dir and out_dir not in target.parents: + raise RuntimeError(f"Unsafe path in tar archive: {member.name}") + if member.issym() or member.islnk(): + raise RuntimeError(f"Refusing link in tar archive: {member.name}") + tf.extractall(out_dir) + + +def safe_extract_archive(archive_path: Path, out_dir: Path) -> None: + lower = archive_path.name.lower() + if lower.endswith(".zip"): + _safe_zip_extract(archive_path, out_dir) + elif lower.endswith((".tar", ".tar.gz", ".tgz")): + _safe_tar_extract(archive_path, out_dir) + else: + raise RuntimeError(f"Unsupported CQ archive format: {archive_path.name}") + + +def promote_single_root(output_dir: Path) -> None: + if (output_dir / "config.txt").exists(): + return + + children = [path for path in output_dir.iterdir() if path.name != "__MACOSX"] + if len(children) != 1 or not children[0].is_dir(): + raise RuntimeError("CQ archive must contain config.txt at the archive root or under one top-level directory") + + nested = children[0] + if not (nested / "config.txt").exists(): + raise RuntimeError("CQ archive has one top-level directory, but it does not contain config.txt") + + for child in nested.iterdir(): + target = output_dir / child.name + if target.exists(): + raise RuntimeError(f"Cannot promote archive member over existing path: {target}") + child.rename(target) + nested.rmdir() + + +def validate_extracted_bundle(output_dir: Path) -> None: + required = ( + "config.txt", + "vocab.txt", + "components/manifest.json", + ) + missing = [name for name in required if not (output_dir / name).exists()] + if missing: + raise RuntimeError(f"Downloaded bundle is missing required file(s): {', '.join(missing)}") + + if not any(output_dir.rglob("*.weights")): + raise RuntimeError("Downloaded bundle contains no weight files") + + tokenizer_config_txt = output_dir / "tokenizer_config.txt" + if not tokenizer_config_txt.exists(): + return + + tokenizer_config = {} + for line in tokenizer_config_txt.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + tokenizer_config[key.strip()] = value.strip() + + tokenizer_type = tokenizer_config.get("tokenizer_type", "").lower() + optional_required = ["special_tokens.json"] + if tokenizer_type == "sentencepiece": + optional_required.append("merges.txt") + else: + optional_required.append("tokenizer.json") + if tokenizer_type == "bpe": + optional_required.append("merges.txt") + + missing_sidecars = [name for name in optional_required if not (output_dir / name).exists()] + if missing_sidecars: + raise RuntimeError(f"Downloaded bundle is missing tokenizer sidecar file(s): {', '.join(missing_sidecars)}") + + +def sha256_file(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def verify_archive_sha256(archive_path: Path, expected_sha256: str | None) -> None: + if not expected_sha256: + return + actual = sha256_file(archive_path) + if actual.lower() != expected_sha256.lower(): + raise RuntimeError( + f"Downloaded archive checksum mismatch for {archive_path.name}: " + f"expected {expected_sha256}, got {actual}" + ) + + +def _download_with_urllib(resolution: CqResolution, *, token=None, revision=None) -> Path: + quoted_repo = urllib.parse.quote(resolution.repo_id, safe="") + quoted_file = urllib.parse.quote(resolution.archive.filename) + revision = revision or "main" + url = f"https://huggingface.co/{resolution.repo_id}/resolve/{urllib.parse.quote(str(revision), safe='')}/{quoted_file}" + + cache_root = Path(tempfile.gettempdir()) / "cactus-cq-cache" + archive_dir = cache_root / quoted_repo / str(revision) + archive_dir.mkdir(parents=True, exist_ok=True) + archive_path = archive_dir / Path(resolution.archive.filename).name + if archive_path.exists() and ( + resolution.archive.size is None or archive_path.stat().st_size == resolution.archive.size + ): + return archive_path + + headers = {"User-Agent": "cactus-cq-downloader"} + if token: + headers["Authorization"] = f"Bearer {token}" + request = urllib.request.Request(url, headers=headers) + tmp_path = archive_path.with_suffix(archive_path.suffix + ".part") + with urllib.request.urlopen(request, timeout=60) as response, tmp_path.open("wb") as out: + shutil.copyfileobj(response, out, length=1024 * 1024) + tmp_path.rename(archive_path) + return archive_path + + +def write_download_metadata(output_dir: Path, resolution: CqResolution, archive_path: Path) -> None: + metadata = { + "repo_id": resolution.repo_id, + "local_name": resolution.local_name, + "archive": resolution.archive.filename, + "bits": resolution.archive.bits, + "archive_size": resolution.archive.size if resolution.archive.size is not None else archive_path.stat().st_size, + "archive_sha256": resolution.archive.sha256, + } + if metadata["archive_sha256"] is None and os.getenv("CACTUS_CQ_HASH_DOWNLOAD", "") == "1": + metadata["archive_sha256"] = sha256_file(archive_path) + (output_dir / ".cactus_cq_download.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") + + +def download_cq_archive(resolution: CqResolution, output_dir: Path, *, token=None, + revision=None) -> Path: + + if output_dir.exists() and (output_dir / "components" / "manifest.json").exists(): + return output_dir + if output_dir.exists() and any(output_dir.iterdir()): + raise RuntimeError(f"Output directory already exists and is not a complete bundle: {output_dir}") + + try: + from huggingface_hub import hf_hub_download + archive_path = Path(hf_hub_download( + repo_id=resolution.repo_id, + filename=resolution.archive.filename, + repo_type="model", + token=token, + revision=revision, + )) + except ImportError: + archive_path = _download_with_urllib( + resolution, + token=token, + revision=revision, + ) + verify_archive_sha256(archive_path, resolution.archive.sha256) + + output_dir.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=str(output_dir.parent)) as tmp: + tmp_dir = Path(tmp) + safe_extract_archive(archive_path, tmp_dir) + promote_single_root(tmp_dir) + validate_extracted_bundle(tmp_dir) + write_download_metadata(tmp_dir, resolution, archive_path) + if output_dir.exists(): + shutil.rmtree(output_dir) + tmp_dir.rename(output_dir) + + return output_dir + + +def list_hf_cq_archives(repo_id: str, *, token=None, revision=None) -> tuple[CqArchive, ...]: + try: + from huggingface_hub import HfApi + except ImportError: + quoted = urllib.parse.quote(repo_id, safe="/") + url = f"https://huggingface.co/api/models/{quoted}?blobs=true" + if revision: + url += f"&revision={urllib.parse.quote(str(revision), safe='')}" + headers = {"User-Agent": "cactus-cq-downloader"} + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=30) as response: + data = json.load(response) + files = [] + sizes = {} + sha256s = {} + for sibling in data.get("siblings", []): + filename = sibling.get("rfilename") + if not filename: + continue + files.append(filename) + lfs = sibling.get("lfs") or {} + if "size" in lfs: + sizes[filename] = int(lfs["size"]) + elif "size" in sibling: + sizes[filename] = int(sibling["size"]) + if "sha256" in lfs: + sha256s[filename] = lfs["sha256"] + return archives_from_repo_files(files, sizes=sizes, sha256s=sha256s) + + try: + info = HfApi().model_info(repo_id, revision=revision, token=token, files_metadata=True) + except Exception as exc: + raise RuntimeError(f"could not query {repo_id} on huggingface.co: {exc}") from exc + files = [] + sizes = {} + sha256s = {} + for sibling in info.siblings: + filename = sibling.rfilename + files.append(filename) + size = getattr(sibling, "size", None) + if size is not None: + sizes[filename] = int(size) + lfs = getattr(sibling, "lfs", None) + if isinstance(lfs, dict): + if "sha256" in lfs: + sha256s[filename] = lfs["sha256"] + if "size" in lfs: + sizes[filename] = int(lfs["size"]) + return archives_from_repo_files(files, sizes=sizes, sha256s=sha256s) + diff --git a/python/cactus/convert/__init__.py b/python/cactus/convert/__init__.py new file mode 100644 index 000000000..dd47a10d0 --- /dev/null +++ b/python/cactus/convert/__init__.py @@ -0,0 +1 @@ +"""CQ conversion pipeline for cactus weights.""" diff --git a/python/cactus/convert/assets/gemma-4-e2b-it/probe.pt b/python/cactus/convert/assets/gemma-4-e2b-it/probe.pt new file mode 100644 index 000000000..434ec3750 Binary files /dev/null and b/python/cactus/convert/assets/gemma-4-e2b-it/probe.pt differ diff --git a/python/cactus/convert/cactus_adapters/__init__.py b/python/cactus/convert/cactus_adapters/__init__.py new file mode 100644 index 000000000..bec954e32 --- /dev/null +++ b/python/cactus/convert/cactus_adapters/__init__.py @@ -0,0 +1,2 @@ +"""Vendored cactus compatibility helpers.""" + diff --git a/python/cactus/convert/cactus_adapters/config_utils.py b/python/cactus/convert/cactus_adapters/config_utils.py new file mode 100644 index 000000000..012f658d7 --- /dev/null +++ b/python/cactus/convert/cactus_adapters/config_utils.py @@ -0,0 +1,626 @@ +import math + + +def cfg_get(c, key, default=None): + """Get a config value from a dict or object, with fallback default.""" + if c is None: + return default + try: + if isinstance(c, dict): + return c.get(key, default) + except Exception: + pass + try: + return getattr(c, key, default) + except Exception: + return default + + +def detect_model_type(cfg, config, output_dir=None): + """Detect the model architecture type from config.""" + model_type_str = str(cfg_get(cfg, 'model_type', cfg_get(config, 'model_type', '')) or '').lower().strip() + decoding_cfg = cfg_get(cfg, 'decoding', cfg_get(config, 'decoding', None)) + decoding_model_type = str(cfg_get(decoding_cfg, 'model_type', '')).lower() + loss_cfg = cfg_get(cfg, 'loss', cfg_get(config, 'loss', None)) + loss_name = str(cfg_get(loss_cfg, 'loss_name', '')).lower() + + # NeMo Parakeet-TDT configs often do not expose HF-style model_type names. + if decoding_model_type == 'tdt' or loss_name == 'tdt': + return 'parakeet_tdt' + + if model_type_str in {'nomic_bert', 'nomic-bert'} or 'nomic' in model_type_str: + return 'nomic' + if 'gemma4' in model_type_str or 'tinyllama' in model_type_str: + return 'gemma4' + elif 'gemma3n' in model_type_str: + return 'gemma3n' + elif 'gemma' in model_type_str: + return 'gemma' + elif 'lfm2' in model_type_str: + return 'lfm2' + elif 'qwen' in model_type_str: + return 'qwen' + elif 'moonshine' in model_type_str: + return 'moonshine' + elif 'needle' in model_type_str: + return 'needle' + elif 'llama' in model_type_str: + if output_dir and 'smol' in str(output_dir): + return 'smol' + else: + return 'llama' + elif 'youtu' in model_type_str: + return 'youtu' + elif 'bert' in model_type_str: + return 'bert' + elif 'whisper' in model_type_str: + return 'whisper' + elif 'parakeet_tdt' in model_type_str or 'parakeet-tdt' in model_type_str: + return 'parakeet_tdt' + elif 'parakeet' in model_type_str: + return 'parakeet' + else: + if model_type_str: + print(f" Warning: Unknown model type '{model_type_str}', defaulting to 'qwen'") + return 'qwen' + + +def resolve_audio_fft_length(audio_cfg): + sampling_rate = int(cfg_get(audio_cfg, 'sampling_rate', 16000)) + frame_length_ms = float(cfg_get(audio_cfg, 'frame_length_ms', 20.0)) + fft_overdrive = bool(cfg_get(audio_cfg, 'fft_overdrive', False)) + frame_length = int(round(sampling_rate * frame_length_ms / 1000.0)) + fft_length = 2 ** math.ceil(math.log2(max(1, frame_length))) + if fft_overdrive: + fft_length *= 2 + return fft_length + + +def extract_base_config(cfg, config): + """Extract base model configuration parameters.""" + rope_parameters = cfg_get(cfg, 'rope_parameters', {}) + rope_theta = cfg_get(cfg, 'rope_theta', None) + partial_rotary_factor = cfg_get(cfg, 'partial_rotary_factor', None) + if rope_theta is None and isinstance(rope_parameters, dict): + rope_theta = rope_parameters.get('rope_theta', None) + if partial_rotary_factor is None and isinstance(rope_parameters, dict): + partial_rotary_factor = rope_parameters.get('partial_rotary_factor', None) + if rope_theta is None: + rope_theta = cfg_get(config, 'rope_theta', 10000.0) + + num_experts_per_tok = cfg_get(cfg, 'num_experts_per_tok', cfg_get(cfg, 'moe_top_k', cfg_get(cfg, 'num_top_experts', cfg_get(cfg, 'top_k_experts', 0)))) or 0 + decoder_start_token_id = cfg_get( + cfg, + 'decoder_start_token_id', + cfg_get( + config, + 'decoder_start_token_id', + cfg_get(cfg, 'bos_token_id', cfg_get(config, 'bos_token_id', 0)), + ), + ) + + base = { + 'vocab_size': cfg_get(cfg, 'vocab_size', cfg_get(config, 'vocab_size', 0)), + 'hidden_dim': cfg_get(cfg, 'hidden_size', cfg_get(cfg, 'hidden_dim', 0)), + 'num_layers': int(cfg_get(cfg, 'num_hidden_layers', cfg_get(cfg, 'num_layers', 0) or 0)), + 'attention_heads': cfg_get(cfg, 'num_attention_heads', 0), + 'attention_kv_heads': cfg_get(cfg, 'num_key_value_heads', cfg_get(cfg, 'num_attention_heads', 0)), + 'ffn_intermediate_dim': cfg_get(cfg, 'intermediate_size', cfg_get(cfg, 'n_inner', 0)), + 'context_length': cfg_get(cfg, 'max_position_embeddings', cfg_get(cfg, 'max_sequence_length', 0)), + 'rope_theta': rope_theta, + 'attention_head_dim': int(cfg_get(cfg, 'head_dim', int(cfg_get(cfg, 'hidden_size', cfg_get(cfg, 'hidden_dim', 0)) // max(1, cfg_get(cfg, 'num_attention_heads', 1))))), + 'layer_norm_eps': cfg_get(cfg, 'layer_norm_eps', cfg_get(cfg, 'layer_norm_epsilon', cfg_get(cfg, 'rms_norm_eps', cfg_get(cfg, 'norm_eps', 1e-6)))), + 'decoder_start_token_id': int(decoder_start_token_id or 0), + 'num_experts': cfg_get(cfg, 'num_experts', 0) or 0, + 'num_shared_experts': cfg_get(cfg, 'num_shared_experts', 0), + 'num_top_experts': num_experts_per_tok, + 'num_experts_per_tok': num_experts_per_tok, + 'moe_every_n_layers': cfg_get(cfg, 'moe_every_n_layers', 0), + } + if partial_rotary_factor is not None: + base['partial_rotary_factor'] = float(partial_rotary_factor) + + layer_types = cfg_get(cfg, 'layer_types', None) + if isinstance(layer_types, (list, tuple)) and layer_types: + base['layer_types'] = list(layer_types) + + conv_l_cache = cfg_get(cfg, 'conv_L_cache', None) + if conv_l_cache is not None: + base['conv_L_cache'] = int(conv_l_cache) + + linear_num_key_heads = cfg_get(cfg, 'linear_num_key_heads', None) + linear_key_head_dim = cfg_get(cfg, 'linear_key_head_dim', None) + linear_num_value_heads = cfg_get(cfg, 'linear_num_value_heads', None) + linear_value_head_dim = cfg_get(cfg, 'linear_value_head_dim', None) + if linear_num_key_heads is not None: + base['linear_num_key_heads'] = int(linear_num_key_heads) + if linear_key_head_dim is not None: + base['linear_key_head_dim'] = int(linear_key_head_dim) + if linear_num_value_heads is not None: + base['linear_num_value_heads'] = int(linear_num_value_heads) + if linear_value_head_dim is not None: + base['linear_value_head_dim'] = int(linear_value_head_dim) + if linear_num_key_heads is not None and linear_key_head_dim is not None: + base['linear_q_proj_dim'] = int(linear_num_key_heads) * int(linear_key_head_dim) + base['linear_k_proj_dim'] = int(linear_num_key_heads) * int(linear_key_head_dim) + if linear_num_value_heads is not None and linear_value_head_dim is not None: + base['linear_v_proj_dim'] = int(linear_num_value_heads) * int(linear_value_head_dim) + + return base + + +def extract_vision_config(config, vision_cfg): + """Extract vision encoder configuration for VLM models.""" + vision_hidden = int(cfg_get(vision_cfg, 'hidden_size', 0)) + vision_image_size = cfg_get(vision_cfg, 'image_size', cfg_get(vision_cfg, 'size', {}).get('longest_edge', 0) if isinstance(cfg_get(vision_cfg, 'size', {}), dict) else cfg_get(vision_cfg, 'image_size', 0)) + vision_patch = int(cfg_get(vision_cfg, 'patch_size', 0)) + vision_heads = int(cfg_get(vision_cfg, 'num_attention_heads', 0)) + vision_num_layers = int(cfg_get(vision_cfg, 'num_hidden_layers', cfg_get(vision_cfg, 'num_layers', 0) or 0)) + num_channels = int(cfg_get(vision_cfg, 'num_channels', 3)) + visual_tokens_per_img = 0 + try: + if vision_patch > 0 and vision_image_size > 0: + per_side = vision_image_size // vision_patch + visual_tokens_per_img = per_side * per_side + except Exception: + visual_tokens_per_img = 0 + + pixel_shuffle_factor = int(cfg_get(config, 'scale_factor', cfg_get(vision_cfg, 'scale_factor', 1) or 1)) + downsample_factor = int(cfg_get(config, 'downsample_factor', 2)) + + vision_head_dim = int(cfg_get(vision_cfg, 'head_dim', 64)) + vision_kv_heads = int(cfg_get(vision_cfg, 'num_key_value_heads', vision_heads)) + vision_intermediate_size = int(cfg_get(vision_cfg, 'intermediate_size', 4 * vision_hidden)) + vision_position_embedding_size = int(cfg_get(vision_cfg, 'position_embedding_size', 10240)) + vision_pooling_kernel_size = int(cfg_get(vision_cfg, 'pooling_kernel_size', 3)) + vision_default_output_length = cfg_get(vision_cfg, 'default_output_length', 280) + if isinstance(vision_default_output_length, (list, tuple)): + vision_default_output_length = vision_default_output_length[0] + vision_default_output_length = int(vision_default_output_length) + + vision_rope_params = cfg_get(vision_cfg, 'rope_parameters', {}) + vision_rope_theta = 100.0 + if isinstance(vision_rope_params, dict): + for v in vision_rope_params.values(): + if isinstance(v, dict) and 'rope_theta' in v: + vision_rope_theta = float(v['rope_theta']) + break + + return { + 'vision_hidden_size': int(vision_hidden), + 'vision_num_layers': int(vision_num_layers), + 'vision_image_size': int(vision_image_size), + 'vision_patch_size': int(vision_patch), + 'vision_attention_heads': int(vision_heads), + 'vision_embed_dim': int(vision_hidden), + 'num_channels': int(num_channels), + 'visual_tokens_per_img': int(visual_tokens_per_img), + 'use_pixel_shuffle': bool(pixel_shuffle_factor > 1), + 'pixel_shuffle_factor': int(pixel_shuffle_factor), + 'use_image_tokens': bool(cfg_get(config, 'image_token_id', None) is not None), + 'image_token_id': int(cfg_get(config, 'image_token_id', 0)), + 'use_layout_tags': False, + 'downsample_factor': int(downsample_factor), + 'vision_head_dim': vision_head_dim, + 'vision_kv_heads': vision_kv_heads, + 'vision_intermediate_size': vision_intermediate_size, + 'vision_position_embedding_size': vision_position_embedding_size, + 'vision_pooling_kernel_size': vision_pooling_kernel_size, + 'vision_default_output_length': vision_default_output_length, + 'vision_rope_theta': vision_rope_theta, + } + + +def extract_lfm2_config(cfg): + """Extract LFM2-specific configuration parameters.""" + layer_types = cfg_get(cfg, 'layer_types', []) + conv_L_cache = cfg_get(cfg, 'conv_L_cache', 0) + moe_intermediate_size = cfg_get(cfg, 'moe_intermediate_size', 0) + num_dense_layers = cfg_get(cfg, 'num_dense_layers', 0) + num_experts_per_tok = cfg_get(cfg, 'num_experts_per_tok', 0) + norm_topk_prob = bool(cfg_get(cfg, 'norm_topk_prob', False)) + use_expert_bias = bool(cfg_get(cfg, 'use_expert_bias', False)) + routed_scaling_factor = float(cfg_get(cfg, 'routed_scaling_factor', 1.0)) + return { + 'layer_types': layer_types, + 'conv_L_cache': conv_L_cache, + 'moe_intermediate_size': int(moe_intermediate_size), + 'num_dense_layers': int(num_dense_layers), + 'num_experts_per_tok': int(num_experts_per_tok), + 'norm_topk_prob': norm_topk_prob, + 'use_expert_bias': use_expert_bias, + 'routed_scaling_factor': routed_scaling_factor, + } + + +def extract_youtu_config(cfg): + rope_scaling = cfg_get(cfg, 'rope_scaling', {}) or {} + return { + 'kv_lora_rank': int(cfg_get(cfg, 'kv_lora_rank', 0)), + 'q_lora_rank': int(cfg_get(cfg, 'q_lora_rank', 0)), + 'qk_head_dim': int(cfg_get(cfg, 'qk_head_dim', 0)), + 'qk_nope_head_dim': int(cfg_get(cfg, 'qk_nope_head_dim', 0)), + 'qk_rope_head_dim': int(cfg_get(cfg, 'qk_rope_head_dim', 0)), + 'v_head_dim': int(cfg_get(cfg, 'v_head_dim', 0)), + 'rope_interleave': bool(cfg_get(cfg, 'rope_interleave', False)), + 'attention_bias': bool(cfg_get(cfg, 'attention_bias', False)), + 'rope_scaling_factor': float(cfg_get(rope_scaling, 'factor', 1.0)), + 'rope_mscale_all_dim': float(cfg_get(rope_scaling, 'mscale_all_dim', 0.0)), + } + + +def extract_moonshine_config(cfg): + """Extract Moonshine-specific configuration parameters.""" + rot_factor = getattr(cfg, "partial_rotary_factor", 0.9) + return { + 'partial_rotary_factor': rot_factor, + } + + +def extract_whisper_config(config): + """Extract Cactus runtime config for HF Whisper checkpoints.""" + hidden_dim = int(cfg_get(config, 'd_model', cfg_get(config, 'hidden_size', 0))) + attention_heads = int(cfg_get(config, 'decoder_attention_heads', cfg_get(config, 'encoder_attention_heads', 0))) + encoder_layers = int(cfg_get(config, 'encoder_layers', cfg_get(config, 'num_encoder_layers', 0))) + decoder_layers = int(cfg_get(config, 'decoder_layers', cfg_get(config, 'num_decoder_layers', 0))) + decoder_start_token_id = int(cfg_get(config, 'decoder_start_token_id', cfg_get(config, 'bos_token_id', 0)) or 0) + decoder_prompt_token_ids = [decoder_start_token_id] + forced_decoder_ids = cfg_get(config, 'forced_decoder_ids', None) or [] + for item in sorted(forced_decoder_ids, key=lambda pair: int(pair[0]) if isinstance(pair, (list, tuple)) and pair else 0): + if isinstance(item, (list, tuple)) and len(item) == 2: + decoder_prompt_token_ids.append(int(item[1])) + return { + 'vocab_size': int(cfg_get(config, 'vocab_size', 0)), + 'hidden_dim': hidden_dim, + 'num_layers': max(encoder_layers, decoder_layers), + 'num_encoder_layers': encoder_layers, + 'num_decoder_layers': decoder_layers, + 'attention_heads': attention_heads, + 'attention_kv_heads': attention_heads, + 'attention_head_dim': int(hidden_dim // max(1, attention_heads)), + 'ffn_intermediate_dim': int(cfg_get(config, 'decoder_ffn_dim', cfg_get(config, 'encoder_ffn_dim', 0))), + 'context_length': int(cfg_get(config, 'max_target_positions', cfg_get(config, 'max_position_embeddings', 0))), + 'layer_norm_eps': float(cfg_get(config, 'layer_norm_eps', 1e-5)), + 'rope_theta': 0.0, + 'num_mel_bins': int(cfg_get(config, 'num_mel_bins', 80)), + 'pad_token_id': int(cfg_get(config, 'pad_token_id', 0)), + 'bos_token_id': int(cfg_get(config, 'bos_token_id', 0)), + 'eos_token_id': int(cfg_get(config, 'eos_token_id', 0)), + 'decoder_start_token_id': decoder_start_token_id, + 'decoder_prompt_token_ids': decoder_prompt_token_ids, + 'tie_word_embeddings': bool(cfg_get(config, 'tie_word_embeddings', True)), + } + + +def extract_parakeet_config(config): + """Extract Cactus runtime config for Parakeet CTC checkpoints.""" + encoder_cfg = cfg_get(config, 'encoder_config', None) + if encoder_cfg is None: + raise ValueError("Parakeet conversion requires encoder_config in model config") + + hidden_dim = int(cfg_get(encoder_cfg, 'hidden_size', 0)) + attention_heads = int(cfg_get(encoder_cfg, 'num_attention_heads', 0)) + layer_norm_eps = cfg_get(encoder_cfg, 'layer_norm_eps', cfg_get(encoder_cfg, 'norm_eps', 1e-5)) + if layer_norm_eps is None: + layer_norm_eps = 1e-5 + rope_theta = cfg_get(encoder_cfg, 'rope_theta', 0.0) + if rope_theta is None: + rope_theta = 0.0 + + return { + 'vocab_size': int(cfg_get(config, 'vocab_size', 0)), + 'hidden_dim': hidden_dim, + 'num_layers': int(cfg_get(encoder_cfg, 'num_hidden_layers', 0)), + 'attention_heads': attention_heads, + 'attention_kv_heads': int(cfg_get(encoder_cfg, 'num_key_value_heads', attention_heads)), + 'attention_head_dim': int(hidden_dim // max(1, attention_heads)), + 'ffn_intermediate_dim': int(cfg_get(encoder_cfg, 'intermediate_size', 0)), + 'context_length': int(cfg_get(encoder_cfg, 'max_position_embeddings', 0)), + 'layer_norm_eps': float(layer_norm_eps), + 'rope_theta': float(rope_theta), + 'conv_kernel_size': int(cfg_get(encoder_cfg, 'conv_kernel_size', 9)), + 'subsampling_conv_kernel_size': int(cfg_get(encoder_cfg, 'subsampling_conv_kernel_size', 3)), + 'subsampling_conv_stride': int(cfg_get(encoder_cfg, 'subsampling_conv_stride', 2)), + 'subsampling_conv_channels': int(cfg_get(encoder_cfg, 'subsampling_conv_channels', 256)), + 'subsampling_factor': int(cfg_get(encoder_cfg, 'subsampling_factor', 8)), + 'num_mel_bins': int(cfg_get(encoder_cfg, 'num_mel_bins', 80)), + 'pad_token_id': int(cfg_get(config, 'pad_token_id', 0)), + 'encoder_hidden_act': cfg_get(encoder_cfg, 'hidden_act', 'silu'), + } + + +def extract_parakeet_tdt_config(config): + """Extract Cactus runtime config for Parakeet TDT checkpoints.""" + encoder_cfg = cfg_get(config, 'encoder', cfg_get(config, 'encoder_config', None)) + if encoder_cfg is None: + raise ValueError("Parakeet TDT conversion requires encoder config") + + preprocessor_cfg = cfg_get(config, 'preprocessor', {}) + decoder_cfg = cfg_get(config, 'decoder', {}) + prediction_cfg = cfg_get(decoder_cfg, 'prediction', {}) + joint_cfg = cfg_get(config, 'joint', {}) + jointnet_cfg = cfg_get(joint_cfg, 'jointnet', {}) + model_defaults_cfg = cfg_get(config, 'model_defaults', {}) + + hidden_dim = int(cfg_get(encoder_cfg, 'd_model', cfg_get(encoder_cfg, 'hidden_size', 0))) + attention_heads = int(cfg_get(encoder_cfg, 'n_heads', cfg_get(encoder_cfg, 'num_attention_heads', 0))) + ff_intermediate = int(cfg_get(encoder_cfg, 'ffn_hidden_size', cfg_get(encoder_cfg, 'intermediate_size', 0))) + if ff_intermediate == 0: + ff_intermediate = int(round(hidden_dim * float(cfg_get(encoder_cfg, 'ff_expansion_factor', 4.0)))) + + labels = cfg_get(config, 'labels', []) + if not isinstance(labels, (list, tuple)): + labels = [] + tdt_durations = cfg_get(model_defaults_cfg, 'tdt_durations', cfg_get(config, 'durations', [])) + if not isinstance(tdt_durations, (list, tuple)): + tdt_durations = [] + vocab_size = int(cfg_get(decoder_cfg, 'vocab_size', cfg_get(config, 'vocab_size', len(labels)))) + blank_id_cfg = cfg_get( + cfg_get(config, 'decoding', {}), + 'blank_id', + cfg_get(decoder_cfg, 'blank_id', cfg_get(config, 'blank_token_id', None)), + ) + try: + tdt_blank_id = int(blank_id_cfg) if blank_id_cfg is not None else vocab_size + except (TypeError, ValueError): + tdt_blank_id = vocab_size + if tdt_blank_id < 0: + tdt_blank_id = vocab_size + + return { + 'vocab_size': vocab_size, + 'hidden_dim': hidden_dim, + 'num_layers': int(cfg_get(encoder_cfg, 'n_layers', cfg_get(encoder_cfg, 'num_hidden_layers', 0))), + 'attention_heads': attention_heads, + 'attention_kv_heads': int(cfg_get(encoder_cfg, 'n_kv_heads', attention_heads)), + 'attention_head_dim': int(hidden_dim // max(1, attention_heads)), + 'ffn_intermediate_dim': ff_intermediate, + 'context_length': int(cfg_get(encoder_cfg, 'max_position_embeddings', 0)), + 'layer_norm_eps': float(cfg_get(encoder_cfg, 'layer_norm_eps', cfg_get(encoder_cfg, 'norm_eps', 1e-5)) or 1e-5), + 'rope_theta': float(cfg_get(encoder_cfg, 'rope_theta', 0.0) or 0.0), + 'conv_kernel_size': int(cfg_get(encoder_cfg, 'conv_kernel_size', 9)), + 'subsampling_conv_kernel_size': int(cfg_get(encoder_cfg, 'subsampling_conv_kernel_size', 3)), + 'subsampling_conv_stride': int(cfg_get(encoder_cfg, 'subsampling_conv_stride', 2)), + 'subsampling_conv_channels': int(cfg_get(encoder_cfg, 'subsampling_conv_channels', 256)), + 'subsampling_factor': int(cfg_get(encoder_cfg, 'subsampling_factor', 8)), + 'num_mel_bins': int(cfg_get(preprocessor_cfg, 'features', cfg_get(encoder_cfg, 'feat_in', cfg_get(encoder_cfg, 'num_mel_bins', 128)))), + 'pad_token_id': int(cfg_get(config, 'pad_token_id', 0)), + 'encoder_hidden_act': cfg_get(encoder_cfg, 'activation', cfg_get(encoder_cfg, 'hidden_act', 'silu')), + 'predictor_hidden_dim': int(cfg_get(prediction_cfg, 'pred_hidden', cfg_get(config, 'decoder_hidden_size', 0))), + 'predictor_num_layers': int(cfg_get(prediction_cfg, 'pred_rnn_layers', cfg_get(config, 'num_decoder_layers', 0))), + 'tdt_joint_dim': int(cfg_get(jointnet_cfg, 'joint_hidden', cfg_get(config, 'decoder_hidden_size', 0))), + 'tdt_num_durations': len(tdt_durations), + 'tdt_durations': [int(v) for v in tdt_durations], + 'tdt_blank_id': tdt_blank_id, + } + + +def extract_complex_gemma_config(cfg, root_config): + """Extract configuration parameters for Gemma3n and Gemma4 models.""" + altup_num_inputs = int(cfg_get(cfg, 'altup_num_inputs', cfg_get(root_config, 'altup_num_inputs', 4))) + laurel_rank = int(cfg_get(cfg, 'laurel_rank', cfg_get(root_config, 'laurel_rank', 64))) + hidden_size_per_layer_input_raw = cfg_get(cfg, 'hidden_size_per_layer_input', + cfg_get(root_config, 'hidden_size_per_layer_input', None)) + hidden_size_per_layer_input = int(hidden_size_per_layer_input_raw) if hidden_size_per_layer_input_raw is not None else 0 + rope_local_base_freq = float(cfg_get(cfg, 'rope_local_base_freq', + cfg_get(root_config, 'rope_local_base_freq', 10000.0))) + + num_kv_shared_layers = int(cfg_get(cfg, 'num_kv_shared_layers', + cfg_get(root_config, 'num_kv_shared_layers', 0))) + sliding_window = int(cfg_get(cfg, 'sliding_window', + cfg_get(root_config, 'sliding_window', 512))) + + rope_theta = cfg_get(root_config, 'rope_theta', None) + if rope_theta is None: + rope_theta = cfg_get(cfg, 'rope_theta', 1000000.0) + + attention_types = cfg_get(cfg, 'attention_type_pattern', cfg_get(root_config, 'attention_type_pattern', None)) + if attention_types is None: + attention_types = cfg_get(cfg, 'attention_types', cfg_get(root_config, 'attention_types', None)) + if attention_types is None: + attention_types = cfg_get(cfg, 'layer_types', cfg_get(root_config, 'layer_types', None)) + layer_types = [] + if attention_types: + num_layers = int(cfg_get(cfg, 'num_hidden_layers', cfg_get(cfg, 'num_layers', 30))) + if isinstance(attention_types, (list, tuple)): + pattern = list(attention_types) + else: + pattern = [str(attention_types)] + pattern_len = len(pattern) + for i in range(num_layers): + at = str(pattern[i % pattern_len]).lower() + if 'global' in at or 'full' in at: + layer_types.append('global') + else: + layer_types.append('sliding') + + final_logit_softcapping = float(cfg_get(cfg, 'final_logit_softcapping', + cfg_get(root_config, 'final_logit_softcapping', 30.0))) + + query_pre_attn_scalar = int(cfg_get(cfg, 'query_pre_attn_scalar', + cfg_get(root_config, 'query_pre_attn_scalar', 0))) + + rope_params = cfg_get(cfg, 'rope_parameters', cfg_get(root_config, 'rope_parameters', {})) + global_rope = rope_params.get('full_attention', {}) if isinstance(rope_params, dict) else {} + global_partial_rotary_factor = float(cfg_get(cfg, 'global_partial_rotary_factor', + global_rope.get('partial_rotary_factor', 1.0) if isinstance(global_rope, dict) else 1.0)) + + if rope_theta is None or rope_theta == 1000000.0: + global_rope_theta = global_rope.get('rope_theta', None) if isinstance(global_rope, dict) else None + if global_rope_theta is not None: + rope_theta = float(global_rope_theta) + + sliding_rope = rope_params.get('sliding_attention', {}) if isinstance(rope_params, dict) else {} + if isinstance(sliding_rope, dict) and 'rope_theta' in sliding_rope: + rope_local_base_freq = float(sliding_rope['rope_theta']) + + activation_sparsity = cfg_get(cfg, 'activation_sparsity_pattern', + cfg_get(root_config, 'activation_sparsity_pattern', None)) + + activation_sparsity_ppf = None + if activation_sparsity: + import torch, math + activation_sparsity_ppf = [] + for s in activation_sparsity: + if s > 0: + activation_sparsity_ppf.append( + round(math.sqrt(2) * torch.erfinv(torch.tensor(2.0 * s - 1.0)).item(), 7)) + else: + activation_sparsity_ppf.append(0.0) + + query_pre_attn_scalar = int(cfg_get(cfg, 'query_pre_attn_scalar', + cfg_get(root_config, 'query_pre_attn_scalar', 0))) + + rope_params = cfg_get(cfg, 'rope_parameters', cfg_get(root_config, 'rope_parameters', {})) + global_rope = rope_params.get('full_attention', {}) if isinstance(rope_params, dict) else {} + global_partial_rotary_factor = float(cfg_get(cfg, 'global_partial_rotary_factor', + global_rope.get('partial_rotary_factor', 1.0) if isinstance(global_rope, dict) else 1.0)) + + if rope_theta is None or rope_theta == 1000000.0: + global_rope_theta = global_rope.get('rope_theta', None) if isinstance(global_rope, dict) else None + if global_rope_theta is not None: + rope_theta = float(global_rope_theta) + + sliding_rope = rope_params.get('sliding_attention', {}) if isinstance(rope_params, dict) else {} + if isinstance(sliding_rope, dict) and 'rope_theta' in sliding_rope: + rope_local_base_freq = float(sliding_rope['rope_theta']) + + global_head_dim = cfg_get(cfg, 'global_head_dim', cfg_get(root_config, 'global_head_dim', None)) + num_global_kv_heads = cfg_get(cfg, 'num_global_key_value_heads', + cfg_get(root_config, 'num_global_key_value_heads', None)) + expert_intermediate_size = cfg_get(cfg, 'expert_intermediate_size', + cfg_get(root_config, 'expert_intermediate_size', None)) + if expert_intermediate_size is None: + expert_intermediate_size = cfg_get(cfg, 'moe_intermediate_size', + cfg_get(root_config, 'moe_intermediate_size', None)) + top_k_experts = cfg_get(cfg, 'top_k_experts', cfg_get(root_config, 'top_k_experts', None)) + if top_k_experts is None: + top_k_experts = cfg_get(cfg, 'num_experts_per_tok', cfg_get(root_config, 'num_experts_per_tok', None)) + attention_k_eq_v = bool(cfg_get(cfg, 'attention_k_eq_v', + cfg_get(root_config, 'attention_k_eq_v', False))) + vocab_size_per_layer_input = cfg_get(cfg, 'vocab_size_per_layer_input', + cfg_get(root_config, 'vocab_size_per_layer_input', None)) + sliding_window_pattern = cfg_get(cfg, '_sliding_window_pattern', + cfg_get(root_config, '_sliding_window_pattern', None)) + enable_moe_block = bool(cfg_get(cfg, 'enable_moe_block', + cfg_get(root_config, 'enable_moe_block', False))) + + result = { + 'altup_num_inputs': altup_num_inputs, + 'laurel_rank': laurel_rank, + 'hidden_size_per_layer_input': hidden_size_per_layer_input, + 'num_kv_shared_layers': num_kv_shared_layers, + 'sliding_window': sliding_window, + 'rope_local_base_freq': rope_local_base_freq, + 'rope_theta': float(rope_theta), + 'final_logit_softcapping': final_logit_softcapping, + 'query_pre_attn_scalar': query_pre_attn_scalar, + 'global_partial_rotary_factor': global_partial_rotary_factor, + 'attention_k_eq_v': attention_k_eq_v, + 'enable_moe_block': enable_moe_block, + } + if top_k_experts is not None: + result['num_top_experts'] = int(top_k_experts) + result['num_experts_per_tok'] = int(top_k_experts) + if global_head_dim is not None: + result['global_head_dim'] = int(global_head_dim) + if num_global_kv_heads is not None: + result['num_global_key_value_heads'] = int(num_global_kv_heads) + if expert_intermediate_size is not None: + result['expert_intermediate_size'] = int(expert_intermediate_size) + if vocab_size_per_layer_input is not None: + result['vocab_size_per_layer_input'] = int(vocab_size_per_layer_input) + if sliding_window_pattern is not None: + result['sliding_window_pattern'] = int(sliding_window_pattern) + if layer_types: + result['layer_types'] = layer_types + if activation_sparsity_ppf: + result['activation_sparsity_ppf'] = activation_sparsity_ppf + return result + +def extract_audio_config(config, audio_cfg): + """Extract audio encoder configuration for multimodal models.""" + hidden = int(cfg_get(audio_cfg, 'hidden_size', 1024)) + num_heads = int(cfg_get(audio_cfg, 'conf_num_attention_heads', 8)) + sscp_channels = cfg_get(audio_cfg, 'sscp_conv_channel_size', [128, 32]) + if not isinstance(sscp_channels, (list, tuple)): + sscp_channels = [128, 32] + output_proj = cfg_get(audio_cfg, 'output_proj_dims', None) + if output_proj is None: + output_proj = 0 + return { + 'audio_hidden_dim': hidden, + 'audio_num_layers': int(cfg_get(audio_cfg, 'conf_num_hidden_layers', 12)), + 'audio_num_heads': num_heads, + 'audio_head_dim': hidden // max(1, num_heads), + 'audio_input_feat_size': int(cfg_get(audio_cfg, 'input_feat_size', 128)), + 'audio_conf_conv_kernel_size': int(cfg_get(audio_cfg, 'conf_conv_kernel_size', 5)), + 'audio_chunk_size': int(cfg_get(audio_cfg, 'conf_attention_chunk_size', 12)), + 'audio_context_left': int(cfg_get(audio_cfg, 'conf_attention_context_left', 13)), + 'audio_context_right': int(cfg_get(audio_cfg, 'conf_attention_context_right', 0)), + 'audio_logit_cap': float(cfg_get(audio_cfg, 'conf_attention_logit_cap', 50.0)), + 'audio_residual_weight': float(cfg_get(audio_cfg, 'conf_residual_weight', 0.5)), + 'audio_output_proj_dims': int(output_proj), + 'audio_vocab_size': int(cfg_get(audio_cfg, 'vocab_size', 128)), + 'audio_vocab_offset': int(cfg_get(audio_cfg, 'vocab_offset', 0)), + 'audio_soft_tokens': int(cfg_get(audio_cfg, 'audio_soft_tokens_per_image', 188)), + 'audio_sscp_conv0_channels': int(sscp_channels[0]), + 'audio_sscp_conv1_channels': int(sscp_channels[1]), + 'audio_sscp_conv_eps': float(cfg_get(audio_cfg, 'sscp_conv_eps', 1e-3)), + 'audio_rms_norm_eps': float(cfg_get(audio_cfg, 'rms_norm_eps', 1e-6)), + 'audio_fft_length': resolve_audio_fft_length(audio_cfg), + 'audio_token_id': int(cfg_get(config, 'audio_token_id', 0)), + } + + +def is_vlm_model(config): + """Check if a model config indicates a vision-language model.""" + text_cfg = cfg_get(config, 'text_config', None) + vision_cfg = cfg_get(config, 'vision_config', None) + return text_cfg is not None or vision_cfg is not None + + +def is_lfm2_vl(model_name, cfg): + """Check if the model is an LFM2 vision-language model.""" + model_type = str(cfg_get(cfg, 'model_type', '') or '').lower().strip() + if model_type.replace('_', '-') == "lfm2-vl": + return True + architectures = cfg_get(cfg, 'architectures', []) + if isinstance(architectures, (list, tuple)): + for architecture in architectures: + normalized_arch = str(architecture or '').lower().replace('_', '').replace('-', '') + if normalized_arch == 'lfm2vlforconditionalgeneration': + return True + name = (model_name or "").lower() + return ( + "lfm2-vl" in name + or "lfm2_vl" in name + or "lfm2.5-vl" in name + or "lfm2.5_vl" in name + ) + + +def pick_dtype(): + """Select torch dtype for model loading — bf16 preferred, float16 fallback.""" + import torch + try: + torch.zeros(1, dtype=torch.bfloat16) + return torch.bfloat16 + except Exception: + return torch.float16 + + +def vision_weight_sanity_check(model): + """Verify vision tower weights are properly initialized.""" + ok = True + vt = getattr(model, "vision_tower", None) + try: + emb = vt.vision_model.embeddings + w_mean = emb.patch_embedding.weight.detach().abs().mean().item() + p_mean = emb.position_embedding.weight.detach().abs().mean().item() + print(f"[sanity] |patch W| mean={w_mean:.5f} |pos W| mean={p_mean:.5f}") + if w_mean < 1e-3 or p_mean < 1e-3: + ok = False + except Exception: + pass + return ok diff --git a/python/cactus/convert/cactus_adapters/tensor_io.py b/python/cactus/convert/cactus_adapters/tensor_io.py new file mode 100644 index 000000000..ba978c6aa --- /dev/null +++ b/python/cactus/convert/cactus_adapters/tensor_io.py @@ -0,0 +1,647 @@ +import numpy as np +import struct +from .weight_patterns import EMBED_NAMES +from ..model_adapters.naming import GEMMA4_WEIGHT_SCALE, gemma4_scale_factor + +try: + import torch +except ImportError: + torch = None + + +GROUP_SIZE = 32 +INTERLEAVE_BLOCK = 4 + +CACTUS_MAGIC = b'CACT' +CACTUS_ALIGNMENT = 32 + +FLAG_HAS_SCALES = 1 << 0 +FLAG_PAGE_ALIGNED = 1 << 1 +FLAG_TRANSPOSED = 1 << 2 +FLAG_INTERLEAVED = 1 << 3 + + +def align_offset(offset: int, alignment: int) -> int: + """Round up offset to next alignment boundary.""" + remainder = offset % alignment + if remainder == 0: + return offset + return offset + (alignment - remainder) + + +def compute_padding(current_offset: int, alignment: int) -> bytes: + """Compute padding bytes needed to reach alignment boundary.""" + aligned = align_offset(current_offset, alignment) + padding_size = aligned - current_offset + return b'\x00' * padding_size + + +def interleave_weights(data: np.ndarray, block_size: int = INTERLEAVE_BLOCK) -> tuple[np.ndarray, int]: + """Interleave rows for SIMD-friendly GEMM access using vdotq_laneq_s32. + + Input: data[N, K] - row-major weights + Output: data_interleaved[N_padded/block_size, K/4, block_size, 4] flattened + original_N - the original N before padding + + Memory layout after interleaving (GGML-style): + For each block of 4 rows and each group of 4 K positions: + [row0_k0..k3, row1_k0..k3, row2_k0..k3, row3_k0..k3, row0_k4..k7, ...] + + This enables vdotq_laneq_s32 to broadcast activation lanes efficiently: + - Load 16 bytes of activation: a[k:k+16] (4 groups of 4) + - Load 16 bytes of weights: 4 columns x 4 K values + - Use lane broadcast to compute 4 output columns simultaneously + """ + N, K = data.shape + original_N = N + + if N % block_size != 0: + pad_n = block_size - (N % block_size) + data = np.pad(data, ((0, pad_n), (0, 0)), mode='constant', constant_values=0) + N = data.shape[0] + + if K % 4 != 0: + pad_k = 4 - (K % 4) + data = np.pad(data, ((0, 0), (0, pad_k)), mode='constant', constant_values=0) + K = data.shape[1] + + data = data.reshape(N // block_size, block_size, K // 4, 4) + data = data.transpose(0, 2, 1, 3) + return data.reshape(-1), original_N + + +def interleave_scales(scales: np.ndarray, block_size: int = INTERLEAVE_BLOCK) -> tuple[np.ndarray, int]: + """Interleave scales to match interleaved weight layout. + + Input: scales[N, num_groups] + Output: scales_interleaved[N_padded/block_size, num_groups, block_size] flattened + original_N + """ + N, num_groups = scales.shape + original_N = N + + if N % block_size != 0: + pad_n = block_size - (N % block_size) + scales = np.pad(scales, ((0, pad_n), (0, 0)), mode='constant', constant_values=1e-10) + N = scales.shape[0] + + scales = scales.reshape(N // block_size, block_size, num_groups) + scales = scales.transpose(0, 2, 1) + return scales.reshape(-1), original_N + + +def pack_int4_pairs(data: np.ndarray) -> np.ndarray: + """Pack INT4 values (stored as int8) into bytes using planar layout. + + Input: array of int8 values in range [-8, 7], length must be a multiple of 32 + Output: array of uint8 with half the length + + Packing format (planar, groups of 32): + For each group of 32 values, the first 16 are stored in the low nibbles + and the next 16 in the high nibbles of 16 consecutive bytes. + Nibbles are stored in two's complement form (value & 0x0F). + """ + assert len(data) % 32 == 0, "Data length must be a multiple of 32 for INT4 planar packing" + + groups = data.reshape(-1, 32) + low = (groups[:, :16].astype(np.int8).view(np.uint8) & 0x0F).astype(np.uint8) + high = ((groups[:, 16:].astype(np.int8).view(np.uint8) & 0x0F).astype(np.uint8)) << 4 + + return (low | high).astype(np.uint8).reshape(-1) + + +def save_depthwise_conv_int8_with_header(tensor, output_path): + """Save [C, 1, K] depthwise conv weights as grouped INT8 without padding K.""" + if torch is not None and isinstance(tensor, torch.Tensor): + t = tensor.detach().cpu() + if t.dtype == torch.bfloat16: + t = t.float() + data = t.numpy() + else: + data = np.array(tensor) + + shape = list(data.shape) + if len(shape) != 3 or shape[1] != 1: + raise ValueError(f"depthwise conv INT8 expects [C, 1, K], got {shape}") + + C, _, K = shape + if K <= 0: + raise ValueError("depthwise conv kernel width must be positive") + + data_2d = data.reshape(C, K).astype(np.float32) + group_abs_max = np.max(np.abs(data_2d), axis=1, keepdims=True) + scales = np.maximum(group_abs_max / 127.0, 1e-10).astype(np.float32) + quantized = np.clip(np.round(data_2d / scales), -128, 127).astype(np.int8).reshape(-1) + scales_fp16 = scales.reshape(-1).astype(np.float16) + + with open(output_path, 'wb') as f: + ndim = 3 + data_bytes = quantized.size + scales_bytes = scales_fp16.size * 2 + flags = FLAG_HAS_SCALES + + f.write(CACTUS_MAGIC) + f.write(struct.pack(' 0 else float('inf') + data_flat = data[:original_N, :].flatten() + dequant_flat = dequantized[:original_N, :].flatten() + cos_sim = np.dot(data_flat, dequant_flat) / (np.linalg.norm(data_flat) * np.linalg.norm(dequant_flat) + 1e-10) + del data_flat, dequantized + + quantized_interleaved, _ = interleave_weights(quantized_2d, INTERLEAVE_BLOCK) + scales_interleaved, _ = interleave_scales(scales, INTERLEAVE_BLOCK) + scales_fp16 = scales_interleaved.astype(np.float16) + + N_padded = ((N + INTERLEAVE_BLOCK - 1) // INTERLEAVE_BLOCK) * INTERLEAVE_BLOCK + + if stats_tracker: + stats_tracker['int8_tensors'] += 1 + stats_tracker['quantized_parameters'] += original_N * K + stats_tracker['mse_values'].append(mse_error) + stats_tracker['snr_values'].append(snr_db) + stats_tracker['cos_sim_values'].append(cos_sim) + stats_tracker['total_tensors'] += 1 + stats_tracker['total_parameters'] += original_N * K + + del data + + with open(output_path, 'wb') as f: + ndim = 2 + data_bytes = quantized_interleaved.size + scales_bytes = scales_fp16.size * 2 + flags = FLAG_HAS_SCALES | FLAG_INTERLEAVED + if transpose: + flags |= FLAG_TRANSPOSED + + f.write(CACTUS_MAGIC) # 4 bytes + f.write(struct.pack(' 0 else float('inf') + data_flat = data[:original_N, :].flatten() + dequant_flat = dequantized[:original_N, :].flatten() + cos_sim = np.dot(data_flat, dequant_flat) / (np.linalg.norm(data_flat) * np.linalg.norm(dequant_flat) + 1e-10) + del data_flat, dequantized + + quantized_interleaved, _ = interleave_weights(quantized_2d, INTERLEAVE_BLOCK) + scales_interleaved, _ = interleave_scales(scales, INTERLEAVE_BLOCK) + scales_fp16 = scales_interleaved.astype(np.float16) + + packed_data = pack_int4_pairs(quantized_interleaved) + + N_padded = ((N + INTERLEAVE_BLOCK - 1) // INTERLEAVE_BLOCK) * INTERLEAVE_BLOCK + + if stats_tracker: + stats_tracker['int4_tensors'] += 1 + stats_tracker['quantized_parameters'] += original_N * K + stats_tracker['mse_values'].append(mse_error) + stats_tracker['snr_values'].append(snr_db) + stats_tracker['cos_sim_values'].append(cos_sim) + stats_tracker['total_tensors'] += 1 + stats_tracker['total_parameters'] += original_N * K + + del data + + with open(output_path, 'wb') as f: + ndim = 2 + data_bytes = packed_data.size + scales_bytes = scales_fp16.size * 2 + flags = FLAG_HAS_SCALES | FLAG_INTERLEAVED + if transpose: + flags |= FLAG_TRANSPOSED + + f.write(CACTUS_MAGIC) # 4 bytes + f.write(struct.pack(' 0 else float('inf') + cos_sim = np.dot(data, dequantized) / (np.linalg.norm(data) * np.linalg.norm(dequantized) + 1e-10) + del dequantized + + scales_fp16 = scales.flatten().astype(np.float16) + + if stats_tracker: + stats_tracker['int8_tensors'] += 1 + stats_tracker['quantized_parameters'] += data.size + stats_tracker['mse_values'].append(mse_error) + stats_tracker['snr_values'].append(snr_db) + stats_tracker['cos_sim_values'].append(cos_sim) + stats_tracker['total_tensors'] += 1 + stats_tracker['total_parameters'] += data.size + + with open(output_path, 'wb') as f: + ndim = len(shape) + data_bytes = quantized_flat.size + scales_bytes = scales_fp16.size * 2 + flags = FLAG_HAS_SCALES + if transpose: + flags |= FLAG_TRANSPOSED + + f.write(CACTUS_MAGIC) + f.write(struct.pack('= 1 else 0)) + + header_size = 84 + f.write(compute_padding(header_size, CACTUS_ALIGNMENT)) + + f.write(data_flat.tobytes()) + + +def format_config_value(value): + """Format a config value for writing to config.txt.""" + if isinstance(value, bool): + return 'true' if value else 'false' + if isinstance(value, (list, tuple)): + return ','.join(str(v) for v in value) + return str(value) + + +def create_quantization_stats(): + """Create an empty stats tracker dictionary for quantization metrics.""" + return { + 'total_tensors': 0, + 'int8_tensors': 0, + 'int4_tensors': 0, + 'fp16_tensors': 0, + 'total_parameters': 0, + 'quantized_parameters': 0, + 'mse_values': [], + 'snr_values': [], + 'cos_sim_values': [], + 'saturation_warnings': 0 + } + + +def print_quantization_summary(quantization_stats, args=None): + """Print a summary of quantization statistics.""" + int8_count = quantization_stats.get('int8_tensors', 0) + int4_count = quantization_stats.get('int4_tensors', 0) + fp16_count = quantization_stats.get('fp16_tensors', 0) + quantized_count = int8_count + int4_count + + if quantized_count > 0: + mse_values = np.array(quantization_stats['mse_values']) + snr_values = np.array(quantization_stats['snr_values']) + cos_sim_values = np.array(quantization_stats['cos_sim_values']) + + print("\nQuantization Summary:") + print(f"MSE - Mean: {np.mean(mse_values):.2e}, Max: {np.max(mse_values):.2e}, Median: {np.median(mse_values):.2e}, Min: {np.min(mse_values):.2e}") + print(f"SNR - Mean: {np.mean(snr_values):.1f}dB, Max: {np.max(snr_values):.1f}dB, Median: {np.median(snr_values):.1f}dB, Min: {np.min(snr_values):.1f}dB") + print(f"CosSim - Mean: {np.mean(cos_sim_values):.6f}, Max: {np.max(cos_sim_values):.6f}, Median: {np.median(cos_sim_values):.6f}, Min: {np.min(cos_sim_values):.6f}") + print(f"Processed {int8_count} INT8 tensors, {int4_count} INT4 tensors, {fp16_count} FP16 tensors") diff --git a/python/cactus/convert/cactus_adapters/tokenizer.py b/python/cactus/convert/cactus_adapters/tokenizer.py new file mode 100644 index 000000000..46ee4c86b --- /dev/null +++ b/python/cactus/convert/cactus_adapters/tokenizer.py @@ -0,0 +1,597 @@ +import json +from pathlib import Path + +try: + from huggingface_hub import hf_hub_download +except ImportError: + hf_hub_download = None + +SENTENCEPIECE_MODEL_TYPES = { + 'gemma', 'gemma3n', 'gemma4', 'llama', 'smol', 'bert', 't5', +} + +BPE_MODEL_TYPES = { + 'qwen', 'qwen3_5', 'lfm2', + 'whisper', 'moonshine', + 'parakeet', 'parakeet_tdt', +} + + +def _is_metaspace_normalizer(normalizer): + return ( + isinstance(normalizer, dict) + and normalizer.get("type") == "Replace" + and normalizer.get("pattern", {}).get("String") == " " + and normalizer.get("content") == "▁" + ) + + +def _decoder_has_type(decoder, decoder_type): + if not isinstance(decoder, dict): + return False + if decoder.get("type") == decoder_type: + return True + if decoder.get("type") == "Sequence": + return any(_decoder_has_type(item, decoder_type) for item in decoder.get("decoders", [])) + return False + + +def _is_replace_metaspace_decoder(decoder): + if not isinstance(decoder, dict): + return False + if decoder.get("type") == "Replace": + return ( + decoder.get("pattern", {}).get("String") == "▁" + and decoder.get("content") == " " + ) + if decoder.get("type") == "Sequence": + return any(_is_replace_metaspace_decoder(item) for item in decoder.get("decoders", [])) + return False + + +def convert_hf_tokenizer(tokenizer, output_dir, token=None, model_id=None, labels=None, model_type=None): + """Convert a HuggingFace tokenizer to Cactus format.""" + model_name_l = (model_id or getattr(tokenizer, 'name_or_path', '') or '').lower() + + # Parakeet-TDT exports labels directly in config.json (8192 token classes). + # HF tokenizer object for this repo exposes only a byte-level 256 vocab, + # which breaks runtime decode. Prefer labels when provided. + if labels and isinstance(labels, (list, tuple)) and 'parakeet-tdt' in model_name_l: + id_to_token = [str(tok) for tok in labels] + vocab_size = len(id_to_token) + + vocab_output = output_dir / "vocab.txt" + with open(vocab_output, 'w', encoding='utf-8') as f: + for token_id, token_str in enumerate(id_to_token): + f.write(f"{token_id}\t{token_str}\n") + print(f" Saved Parakeet-TDT vocabulary from labels (ID\\ttoken format, {vocab_size} tokens)") + + merges_output = output_dir / "merges.txt" + with open(merges_output, 'w', encoding='utf-8', newline='') as f: + f.write("#version: 0.2\n") + + token_to_id = {tok: i for i, tok in enumerate(id_to_token)} + unk_id = token_to_id.get('', 0) + pad_id = token_to_id.get('', getattr(tokenizer, 'pad_token_id', 0) or 0) + eos_id = token_to_id.get('<|endoftext|>', getattr(tokenizer, 'eos_token_id', None)) + if eos_id is None: + eos_id = pad_id + bos_id = token_to_id.get('<|startoftranscript|>', None) + + special_token_ids = { + 'pad_token_id': int(pad_id), + 'unk_token_id': int(unk_id), + 'eos_token_id': int(eos_id), + } + if bos_id is not None: + special_token_ids['bos_token_id'] = int(bos_id) + + # Keep special token map dense for known control tags in labels. + special_tokens = {} + for token_id, tok in enumerate(id_to_token): + if tok.startswith('<') and tok.endswith('>'): + special_tokens[token_id] = tok + # Ensure core IDs are present even if token strings differ. + special_tokens.setdefault(int(pad_id), id_to_token[int(pad_id)] if int(pad_id) < vocab_size else "") + special_tokens.setdefault(int(unk_id), id_to_token[int(unk_id)] if int(unk_id) < vocab_size else "") + special_tokens.setdefault(int(eos_id), id_to_token[int(eos_id)] if int(eos_id) < vocab_size else "<|endoftext|>") + + special_tokens_output = output_dir / "special_tokens.json" + with open(special_tokens_output, 'w', encoding='utf-8') as f: + json.dump({ + **special_token_ids, + "vocab_size": vocab_size, + "model_max_length": getattr(tokenizer, 'model_max_length', 131072), + "special_tokens": special_tokens, + "additional_special_tokens": [], + }, f, indent=2, ensure_ascii=False) + + tokenizer_config_output = output_dir / "tokenizer_config.txt" + with open(tokenizer_config_output, 'w') as f: + f.write(f"vocab_size={vocab_size}\n") + for key, value in special_token_ids.items(): + f.write(f"{key}={value}\n") + f.write(f"model_max_length={getattr(tokenizer, 'model_max_length', 131072)}\n") + f.write("tokenizer_type=sentencepiece\n") + f.write("has_chat_template=false\n") + return + + tokenizer_json_data = {} + tokenizer_json_path = output_dir / "tokenizer.json" + try: + tokenizer.save_pretrained(output_dir) + if tokenizer_json_path.exists(): + with open(tokenizer_json_path, 'r', encoding='utf-8') as f: + tokenizer_json_data = json.load(f) + + unused_files = [ + "tokenizer_config.json", + "special_tokens_map.json", + "added_tokens.json", + "chat_template.jinja", + ] + for filename in unused_files: + filepath = output_dir / filename + if filepath.exists(): + filepath.unlink() + except Exception as e: + print(f" Warning: Could not save tokenizer JSON: {e}") + + source_dir = Path(getattr(tokenizer, "name_or_path", "") or "") + source_tokenizer_model = source_dir / "tokenizer.model" + output_tokenizer_model = output_dir / "tokenizer.model" + if (model_type == "needle" or "needle" in model_name_l) and ( + output_tokenizer_model.exists() or source_tokenizer_model.exists() + ): + tokenizer_model_path = output_tokenizer_model if output_tokenizer_model.exists() else source_tokenizer_model + convert_sentencepiece_tokenizer( + tokenizer_model_path, + output_dir, + getattr(tokenizer, "model_max_length", 1024), + ) + print(" Saved Needle SentencePiece tokenizer") + return + + tokenizer_model = tokenizer_json_data.get("model", {}) if tokenizer_json_data else {} + tokenizer_model_type = str(tokenizer_model.get("type", "")).upper() + is_unigram = tokenizer_model_type == "UNIGRAM" + is_sentencepiece = is_unigram or (tokenizer_model_type != "BPE" and model_type in SENTENCEPIECE_MODEL_TYPES) + + # Unigram (e.g. XLM-RoBERTa, used by nomic-embed) carries per-token scores that + # the SentencePiece runtime needs for Viterbi segmentation. + unigram_scores: dict[int, float] = {} + if is_unigram and isinstance(tokenizer_model.get("vocab"), list): + for token_id, entry in enumerate(tokenizer_model["vocab"]): + if isinstance(entry, (list, tuple)) and len(entry) == 2: + unigram_scores[token_id] = float(entry[1]) + + if tokenizer_model_type == "BPE" and tokenizer_model.get("vocab"): + vocab = tokenizer_model["vocab"] + vocab_size = max(vocab.values()) + 1 + id_to_token = [""] * vocab_size + for token_str, token_id in vocab.items(): + if token_id < vocab_size: + id_to_token[token_id] = token_str + else: + vocab = tokenizer.get_vocab() + vocab_size = (max(vocab.values()) + 1) if vocab else 0 + id_to_token = [""] * vocab_size + for token_str, token_id in vocab.items(): + if token_id < vocab_size: + id_to_token[token_id] = token_str + + # vocab.txt is written later, after special tokens are collected + + + merges_output = output_dir / "merges.txt" + + def write_merges_file(merges_list): + with open(merges_output, 'w', encoding='utf-8', newline='') as f: + f.write("#version: 0.2\n") + for merge in merges_list: + f.write(f"{' '.join(merge)}\n") + + merges_written = False + + if not is_sentencepiece and tokenizer_json_data: + merges_from_json = tokenizer_json_data.get("model", {}).get("merges", []) or [] + write_merges_file(merges_from_json) + merges_written = True + + if not merges_written and hf_hub_download: + try: + import shutil + merges_file = hf_hub_download(repo_id=tokenizer.name_or_path, filename="merges.txt", token=token) + shutil.copy2(merges_file, merges_output) + merges_written = True + except Exception: + pass + + if not merges_written and hasattr(tokenizer, 'backend_tokenizer') and tokenizer.backend_tokenizer: + backend = tokenizer.backend_tokenizer + merges = [] + + if hasattr(backend, 'model'): + model = backend.model + if hasattr(model, 'merges'): + merges = model.merges + + write_merges_file(merges) + merges_written = True + + if not merges_written: + write_merges_file([]) + + + special_tokens = {} + special_token_ids = {} + + if hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id is not None: + special_token_ids['eos_token_id'] = tokenizer.eos_token_id + special_tokens[tokenizer.eos_token_id] = tokenizer.eos_token or "<|endoftext|>" + + if hasattr(tokenizer, 'pad_token_id') and tokenizer.pad_token_id is not None: + special_token_ids['pad_token_id'] = tokenizer.pad_token_id + special_tokens[tokenizer.pad_token_id] = tokenizer.pad_token or "<|endoftext|>" + + if hasattr(tokenizer, 'bos_token_id') and tokenizer.bos_token_id is not None: + special_token_ids['bos_token_id'] = tokenizer.bos_token_id + special_tokens[tokenizer.bos_token_id] = tokenizer.bos_token or "<|startoftext|>" + + if hasattr(tokenizer, 'unk_token_id') and tokenizer.unk_token_id is not None: + special_token_ids['unk_token_id'] = tokenizer.unk_token_id + special_tokens[tokenizer.unk_token_id] = tokenizer.unk_token or "<|unknown|>" + + core_token_fallbacks = { + 'pad_token_id': '', + 'eos_token_id': '', + 'bos_token_id': '', + 'unk_token_id': '', + } + vocab_lookup = {token: token_id for token_id, token in enumerate(id_to_token) if token} + for key, token_str in core_token_fallbacks.items(): + if key not in special_token_ids and token_str in vocab_lookup: + token_id = vocab_lookup[token_str] + special_token_ids[key] = token_id + special_tokens[token_id] = token_str + + if 'eos_token_id' not in special_token_ids and 'parakeet' in model_name_l: + pad_id = special_token_ids.get('pad_token_id') + if pad_id is not None: + special_token_ids['eos_token_id'] = pad_id + + additional_special_tokens = [] + if hasattr(tokenizer, 'additional_special_tokens'): + for token_str in tokenizer.additional_special_tokens or []: + token_id = tokenizer.convert_tokens_to_ids(token_str) + if token_id != tokenizer.unk_token_id: + special_tokens[token_id] = token_str + additional_special_tokens.append({"token": token_str, "id": token_id}) + + for token_info in tokenizer_json_data.get("added_tokens", []) or []: + token_str = token_info.get("content") + token_id = token_info.get("id") + if token_str is None or token_id is None: + continue + special_tokens[int(token_id)] = token_str + if not any(item["token"] == token_str and item["id"] == int(token_id) for item in additional_special_tokens): + additional_special_tokens.append({"token": token_str, "id": int(token_id)}) + + name_or_path_l = model_name_l or getattr(tokenizer, 'name_or_path', '').lower() + if 'gemma' in (model_type or '') or 'gemma' in name_or_path_l: + gemma_special_tokens = { + '': None, + '': None, + '': None, + '': None, + # Gemma 3 function calling tokens + '': None, + '': None, + '': None, + '': None, + '': None, + '': None, + '': None + } + + vocab = tokenizer.get_vocab() + for token_str in gemma_special_tokens.keys(): + if token_str in vocab: + token_id = vocab[token_str] + gemma_special_tokens[token_str] = token_id + special_tokens[token_id] = token_str + print(f" Found Gemma special token: {token_str} (ID: {token_id})") + + missing_tokens = [k for k, v in gemma_special_tokens.items() if v is None] + if missing_tokens: + unk_id = getattr(tokenizer, 'unk_token_id', None) + for token_str in missing_tokens: + token_id = tokenizer.convert_tokens_to_ids(token_str) + if token_id != unk_id and token_id is not None: + gemma_special_tokens[token_str] = token_id + special_tokens[token_id] = token_str + print(f" Found Gemma special token: {token_str} (ID: {token_id})") + + if gemma_special_tokens[''] is None: + hardcoded_ids = { + '': 105, + '': 106 + } + for token_str, token_id in hardcoded_ids.items(): + if token_str in gemma_special_tokens and gemma_special_tokens[token_str] is None: + if token_id not in special_tokens: + gemma_special_tokens[token_str] = token_id + special_tokens[token_id] = token_str + print(f" Using hardcoded Gemma special token: {token_str} (ID: {token_id})") + + chat_template_data = {} + if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template: + chat_template_output = output_dir / "chat_template.jinja2" + with open(chat_template_output, 'w', encoding='utf-8') as f: + f.write(tokenizer.chat_template) + chat_template_data["chat_template"] = tokenizer.chat_template + elif (output_dir / "chat_template.jinja2").exists(): + chat_template_data["chat_template"] = (output_dir / "chat_template.jinja2").read_text(encoding='utf-8') + + tokenizer_full_config = {} + added_tokens_decoder = {} + tool_tokens = {} + + def add_additional_special_token(token_str): + if not isinstance(token_str, str) or not token_str: + return + token_id = None + try: + converted = tokenizer.convert_tokens_to_ids(token_str) + if converted is not None and converted != getattr(tokenizer, 'unk_token_id', None): + token_id = int(converted) + except Exception: + token_id = None + if token_id is None: + token_id = vocab_lookup.get(token_str) + if token_id is None: + return + special_tokens[token_id] = token_str + if not any(item["token"] == token_str and item["id"] == int(token_id) for item in additional_special_tokens): + additional_special_tokens.append({"token": token_str, "id": int(token_id)}) + + try: + config_path = None + if hasattr(tokenizer, 'name_or_path') and hf_hub_download: + try: + local_candidate = Path(tokenizer.name_or_path) / "tokenizer_config.json" + if Path(tokenizer.name_or_path).is_dir() and local_candidate.exists(): + config_path = str(local_candidate) + else: + config_path = hf_hub_download(repo_id=tokenizer.name_or_path, filename="tokenizer_config.json", token=token) + with open(config_path, 'r') as f: + tokenizer_full_config = json.load(f) + + if 'chat_template' in tokenizer_full_config and not chat_template_data: + chat_template_output = output_dir / "chat_template.jinja2" + with open(chat_template_output, 'w', encoding='utf-8') as f: + f.write(tokenizer_full_config['chat_template']) + chat_template_data["chat_template"] = tokenizer_full_config['chat_template'] + + if 'added_tokens_decoder' in tokenizer_full_config: + added_tokens_decoder = tokenizer_full_config['added_tokens_decoder'] + + print(" Extracting special tokens from tokenizer_config.json...") + for token_id_str, token_info in added_tokens_decoder.items(): + content = token_info.get('content', '') + token_id = int(token_id_str) + + tool_related = ['', '', + '', '', + '', '', + '', '', + # Gemma 3 function calling tokens + '', '', + '', '', + '', '', + ''] + + if any(x == content for x in tool_related): + tool_tokens[token_id] = token_info + print(f" Found tool token: {content} (ID: {token_id})") + special_tokens[token_id] = content + + for token_str in tokenizer_full_config.get("additional_special_tokens", []) or []: + add_additional_special_token(token_str) + + except Exception as e: + print(f" Note: Could not load full tokenizer config: {e}") + pass + except Exception: + pass + + try: + if hasattr(tokenizer, 'name_or_path') and Path(tokenizer.name_or_path).is_dir(): + special_map_path = Path(tokenizer.name_or_path) / "special_tokens_map.json" + if special_map_path.exists(): + special_map = json.loads(special_map_path.read_text(encoding="utf-8")) + for token_str in special_map.get("additional_special_tokens", []) or []: + add_additional_special_token(token_str) + except Exception: + pass + + + # Extend id_to_token to include special/added tokens so vocab.txt is self-contained + if special_tokens: + max_special_id = max(special_tokens.keys()) + if max_special_id >= len(id_to_token): + id_to_token.extend([""] * (max_special_id - len(id_to_token) + 1)) + for token_id, token_str in special_tokens.items(): + id_to_token[token_id] = token_str + + vocab_output = output_dir / "vocab.txt" + with open(vocab_output, 'w', encoding='utf-8') as f: + for token_id, token_str in enumerate(id_to_token): + if token_str: + if unigram_scores: + f.write(f"{token_id}\t{token_str}\t{unigram_scores.get(token_id, 0.0)}\n") + else: + f.write(f"{token_id}\t{token_str}\n") + vocab_fmt = "ID\\ttoken\\tscore" if unigram_scores else "ID\\ttoken" + print(f" Saved tokenizer vocabulary ({vocab_fmt} format)") + + special_tokens_output = output_dir / "special_tokens.json" + with open(special_tokens_output, 'w', encoding='utf-8') as f: + json.dump({ + **special_token_ids, + "vocab_size": len(id_to_token), + "model_max_length": getattr(tokenizer, 'model_max_length', 131072), + "special_tokens": special_tokens, + "additional_special_tokens": additional_special_tokens, + **chat_template_data + }, f, indent=2, ensure_ascii=False) + + normalizer = "none" + decoder = "none" + byte_fallback = False + vocab_format = "id_tab_token" + tokenizer_type = "sentencepiece" if is_sentencepiece else "bpe" + + if tokenizer_model_type == "BPE": + tokenizer_type = "bpe" + byte_fallback = bool(tokenizer_model.get("byte_fallback", False)) + if _is_metaspace_normalizer(tokenizer_json_data.get("normalizer")): + normalizer = "metaspace" + elif _decoder_has_type(tokenizer_json_data.get("decoder"), "ByteFallback"): + normalizer = "byte_level" + + if _is_replace_metaspace_decoder(tokenizer_json_data.get("decoder")): + decoder = "replace_metaspace" + elif _decoder_has_type(tokenizer_json_data.get("decoder"), "ByteFallback"): + decoder = "byte_level" + elif is_sentencepiece: + tokenizer_type = "sentencepiece" + if is_unigram: + normalizer = "metaspace" + decoder = "replace_metaspace" + + tokenizer_config_output = output_dir / "tokenizer_config.txt" + with open(tokenizer_config_output, 'w') as f: + f.write(f"vocab_size={len(id_to_token)}\n") + for key, value in special_token_ids.items(): + f.write(f"{key}={value}\n") + f.write(f"model_max_length={getattr(tokenizer, 'model_max_length', 131072)}\n") + f.write(f"tokenizer_type={tokenizer_type}\n") + f.write(f"vocab_format={vocab_format}\n") + f.write(f"normalizer={normalizer}\n") + f.write(f"decoder={decoder}\n") + f.write(f"byte_fallback={'true' if byte_fallback else 'false'}\n") + if is_unigram: + f.write("sp_model_type=unigram\n") + f.write("sp_add_dummy_prefix=true\n") + f.write("sp_byte_fallback=false\n") + + if chat_template_data: + f.write("has_chat_template=true\n") + else: + f.write("has_chat_template=false\n") + if len(tool_tokens) > 0: + f.write("has_tool_support=true\n") + f.write(f"tool_token_count={len(tool_tokens)}\n") + + +def _read_varint(data, pos): + shift, value = 0, 0 + while True: + byte = data[pos]; pos += 1 + value |= (byte & 0x7F) << shift + if not (byte & 0x80): return value, pos + shift += 7 + + +def _skip_proto(data, pos, wire_type): + if wire_type == 0: _, pos = _read_varint(data, pos); return pos + if wire_type == 1: return pos + 8 + if wire_type == 2: length, pos = _read_varint(data, pos); return pos + length + if wire_type == 5: return pos + 4 + raise ValueError(f"Unsupported wire type: {wire_type}") + + +def parse_sentencepiece_pieces(path): + import struct + data = Path(path).read_bytes() + pos, pieces = 0, [] + while pos < len(data): + tag, pos = _read_varint(data, pos) + if (tag >> 3) != 1 or (tag & 7) != 2: + pos = _skip_proto(data, pos, tag & 7); continue + msg_len, pos = _read_varint(data, pos) + end = pos + msg_len; msg = data[pos:end]; pos = end + inner_pos, piece, score = 0, None, 0.0 + while inner_pos < len(msg): + itag, inner_pos = _read_varint(msg, inner_pos) + ifield, iwire = itag >> 3, itag & 7 + if ifield == 1 and iwire == 2: + vlen, inner_pos = _read_varint(msg, inner_pos) + piece = msg[inner_pos:inner_pos + vlen].decode("utf-8", errors="replace") + inner_pos += vlen + elif ifield == 2 and iwire == 5: + score = struct.unpack("", 0) + eos_id = piece_to_id.get("", 1) + bos_id = piece_to_id.get("", 2) + unk_id = piece_to_id.get("", 3) + + special = {pad_id: "", eos_id: "", bos_id: "", unk_id: ""} + tool_tokens = [] + for tok in ("", ""): + tid = piece_to_id.get(tok) + if tid is not None: + special[tid] = tok + tool_tokens.append({"token": tok, "id": tid}) + + return { + "vocab_size": len(pieces), "pad_token_id": pad_id, "eos_token_id": eos_id, + "bos_token_id": bos_id, "unk_token_id": unk_id, "model_max_length": model_max_length, + "sp_model_type": "bpe", "sp_add_dummy_prefix": True, + "sp_remove_extra_whitespaces": True, "sp_escape_whitespaces": True, + "sp_byte_fallback": True, "special_tokens": {str(k): v for k, v in special.items()}, + "additional_special_tokens": tool_tokens, + } + + +def _write_sentencepiece_files(output_dir, pieces, meta): + with open(output_dir / "vocab.txt", "w", encoding="utf-8") as f: + for i, p in enumerate(pieces): + f.write(f"{i}\t{p['piece']}\t{p['score']}\n") + + with open(output_dir / "merges.txt", "w", encoding="utf-8") as f: + f.write("#version: 0.2\n") + + with open(output_dir / "special_tokens.json", "w", encoding="utf-8") as f: + json.dump(meta, f, indent=2, ensure_ascii=False) + + tool_tokens = meta.get("additional_special_tokens", []) + with open(output_dir / "tokenizer_config.txt", "w", encoding="utf-8") as f: + f.write(f"vocab_size={meta['vocab_size']}\n") + f.write(f"eos_token_id={meta['eos_token_id']}\npad_token_id={meta['pad_token_id']}\n") + f.write(f"bos_token_id={meta['bos_token_id']}\nunk_token_id={meta['unk_token_id']}\n") + f.write(f"model_max_length={meta['model_max_length']}\n") + f.write("tokenizer_type=sentencepiece\nsp_model_type=bpe\n") + f.write("sp_add_dummy_prefix=true\nsp_remove_extra_whitespaces=true\n") + f.write("sp_escape_whitespaces=true\nsp_byte_fallback=true\n") + f.write("has_chat_template=false\n") + if tool_tokens: + f.write(f"has_tool_support=true\ntool_token_count={len(tool_tokens)}\n") + + +def convert_sentencepiece_tokenizer(tokenizer_path, output_dir, model_max_length=131072): + output_dir = Path(output_dir) + pieces = parse_sentencepiece_pieces(tokenizer_path) + meta = _build_sentencepiece_metadata(pieces, model_max_length) + _write_sentencepiece_files(output_dir, pieces, meta) diff --git a/python/cactus/convert/cactus_adapters/weight_patterns.py b/python/cactus/convert/cactus_adapters/weight_patterns.py new file mode 100644 index 000000000..67add2b28 --- /dev/null +++ b/python/cactus/convert/cactus_adapters/weight_patterns.py @@ -0,0 +1,417 @@ +EMBED_NAMES = [ + 'model.language_model.embed_tokens.weight', + 'model.text_model.embed_tokens.weight', + 'model.embed_tokens.weight', + 'embed_tokens.weight', + 'embeddings.weight', + 'transformer.wte.weight', + 'model.decoder.embed_tokens.weight', + 'decoder.embed_tokens.weight', + 'token_embeddings' +] + +MOONSHINE_GLOBAL_WEIGHTS = [ + ('model.encoder.conv1.weight', 'encoder_conv1_weight.weights'), + ('model.encoder.conv2.weight', 'encoder_conv2_weight.weights'), + ('model.encoder.conv2.bias', 'encoder_conv2_bias.bias'), + ('model.encoder.conv3.weight', 'encoder_conv3_weight.weights'), + ('model.encoder.conv3.bias', 'encoder_conv3_bias.bias'), + ('model.encoder.groupnorm.weight', 'encoder_norm_weight.weights'), + ('model.encoder.groupnorm.bias', 'encoder_norm_bias.bias'), + ('model.encoder.layer_norm.weight', 'encoder_layer_norm_weight.weights'), + ('model.encoder.layer_norm.bias', 'encoder_layer_norm_bias.bias'), + ('model.decoder.norm.weight', 'output_norm.weights'), + ('model.decoder.norm.bias', 'output_norm.bias'), + ('decoder.norm.weight', 'output_norm.weights'), + ('proj_out.weight', 'output_weight.weights'), + ('model.proj_out.weight', 'output_weight.weights'), + ('model.decoder.embed_tokens.weight', 'token_embeddings.weights'), + ('model.decoder.embed_positions.weight', 'decoder_position_embeddings.weights'), +] + +OUTPUT_NAMES = [ + 'lm_head.weight', + 'output.weight', + 'transformer.lm_head.weight', + 'model.text_model.lm_head.weight' +] + +OUTPUT_NORM_NAMES = [ + 'model.norm.weight', + 'norm.weight', + 'final_layernorm.weight', + 'transformer.ln_f.weight', + 'model.embedding_norm.weight', + 'model.language_model.embedding_norm.weight', + 'model.language_model.norm.weight', + 'model.text_model.norm.weight', + 'decoder.norm.weight' +] + +LAYER_PREFIXES = [ + 'model.language_model.layers.{i}.', + 'model.text_model.layers.{i}.', + 'model.layers.{i}.', + 'layers.{i}.', + 'transformer.h.{i}.', + 'encoder.layers.{i}.', + 'decoder.layers.{i}.', + 'model.encoder.layers.{i}.', + 'model.decoder.layers.{i}.' +] + +VISION_ITEMS = [ + ('model.vision_tower.vision_model.embeddings.patch_embedding.weight', 'vision_patch_embedding.weights'), + ('model.vision_model.embeddings.patch_embedding.weight', 'vision_patch_embedding.weights'), + ('model.vision_tower.vision_model.embeddings.patch_embedding.bias', 'vision_patch_embedding.bias.weights'), + ('model.vision_model.embeddings.patch_embedding.bias', 'vision_patch_embedding.bias.weights'), + ('model.vision_tower.vision_model.embeddings.position_embedding.weight', 'vision_position_embedding.weights'), + ('model.vision_model.embeddings.position_embedding.weight', 'vision_position_embedding.weights'), + ('model.vision_tower.vision_model.post_layernorm.weight', 'vision_post_layernorm.weights'), + ('model.vision_model.post_layernorm.weight', 'vision_post_layernorm.weights'), + ('model.vision_tower.vision_model.post_layernorm.bias', 'vision_post_layernorm.bias.weights'), + ('model.vision_model.post_layernorm.bias', 'vision_post_layernorm.bias.weights') +] + +PROJECTOR_WEIGHTS = [ + ('model.multi_modal_projector.linear_1.weight', 'projector_linear1.weights'), + ('model.multi_modal_projector.linear_1.bias', 'projector_linear1.bias.weights'), + ('model.multi_modal_projector.linear_2.weight', 'projector_linear2.weights'), + ('model.multi_modal_projector.linear_2.bias', 'projector_linear2.bias.weights'), + ('model.multi_modal_projector.layer_norm.weight', 'projector_layer_norm.weights'), + ('model.multi_modal_projector.layer_norm.bias', 'projector_layer_norm.bias.weights'), +] + +CONNECTOR_KEYS = [ + 'model.connector.modality_projection.proj.weight', + 'connector.modality_projection.proj.weight', + 'model.connector.proj.weight', + 'connector.proj.weight' +] + +GEMMA3N_GLOBAL_WEIGHTS = [ + ('model.language_model.altup_projections.0.weight', 'altup_proj_0.weights'), + ('model.language_model.altup_projections.1.weight', 'altup_proj_1.weights'), + ('model.language_model.altup_projections.2.weight', 'altup_proj_2.weights'), + ('model.language_model.altup_unembed_projections.0.weight', 'altup_unembed_proj_0.weights'), + ('model.language_model.altup_unembed_projections.1.weight', 'altup_unembed_proj_1.weights'), + ('model.language_model.altup_unembed_projections.2.weight', 'altup_unembed_proj_2.weights'), + ('model.language_model.embed_tokens_per_layer.weight', 'embed_tokens_per_layer.weights'), + ('model.language_model.per_layer_model_projection.weight', 'per_layer_model_proj.weights'), + ('model.language_model.per_layer_projection_norm.weight', 'per_layer_proj_norm.weights'), + ('model.embed_vision.embedding.weight', 'embed_vision_embedding.weights'), + ('model.embed_vision.embedding_projection.weight', 'embed_vision_proj.weights'), + ('model.embed_vision.soft_embedding_norm.weight', 'embed_vision_soft_norm.weights'), + ('model.embed_vision.hard_embedding_norm.weight', 'embed_vision_hard_norm.weights'), + ('model.embed_audio.embedding.weight', 'embed_audio_embedding.weights'), + ('model.embed_audio.embedding_projection.weight', 'embed_audio_proj.weights'), + ('model.embed_audio.soft_embedding_norm.weight', 'embed_audio_soft_norm.weights'), + ('model.embed_audio.hard_embedding_norm.weight', 'embed_audio_hard_norm.weights'), +] + +GEMMA3N_VISION_TOWER_PREFIX = 'model.vision_tower.timm_model.' +GEMMA3N_AUDIO_TOWER_PREFIX = 'model.audio_tower.' + +GEMMA4_GLOBAL_WEIGHTS = [ + ('model.language_model.embed_tokens_per_layer.weight', 'embed_tokens_per_layer.weights'), + ('model.language_model.per_layer_model_projection.weight', 'per_layer_model_proj.weights'), + ('model.language_model.per_layer_projection_norm.weight', 'per_layer_proj_norm.weights'), + ('model.embed_vision.embedding.weight', 'embed_vision_embedding.weights'), + ('model.embed_vision.embedding_projection.weight', 'embed_vision_proj.weights'), + ('model.embed_audio.embedding.weight', 'embed_audio_embedding.weights'), + ('model.embed_audio.embedding_projection.weight', 'embed_audio_proj.weights'), +] + +GEMMA4_VISION_TOWER_PREFIX = 'model.vision_tower.' +GEMMA4_AUDIO_TOWER_PREFIX = 'model.audio_tower.' + +WHISPER_GLOBAL_WEIGHTS = [ + ('decoder.embed_tokens.weight', 'decoder_token_embeddings.weights'), + ('decoder.embed_positions.weight', 'decoder_position_embeddings.weights'), + ('decoder.layer_norm.weight', 'decoder_norm.weights'), + ('decoder.layer_norm.bias', 'decoder_norm.bias'), + ('proj_out.weight', 'output_weight.weights'), + ('encoder.embed_positions.weight', 'encoder_position_embeddings.weights'), + ('encoder.conv1.bias', 'encoder_conv1_bias.bias'), + ('encoder.conv1.weight', 'encoder_conv1_weight.weights'), + ('encoder.conv2.bias', 'encoder_conv2_bias.bias'), + ('encoder.conv2.weight', 'encoder_conv2_weight.weights'), + ('encoder.layer_norm.bias', 'encoder_norm_bias.bias'), + ('encoder.layer_norm.weight', 'encoder_norm_weight.weights') +] + +PARAKEET_GLOBAL_WEIGHTS = [ + (('encoder.subsampling.layers.0.weight',), 'subsampling_conv0_weight.weights'), + (('encoder.subsampling.layers.0.bias',), 'subsampling_conv0_bias.bias'), + (('encoder.subsampling.layers.2.weight',), 'subsampling_depthwise1_weight.weights'), + (('encoder.subsampling.layers.2.bias',), 'subsampling_depthwise1_bias.bias'), + (('encoder.subsampling.layers.3.weight',), 'subsampling_pointwise1_weight.weights'), + (('encoder.subsampling.layers.3.bias',), 'subsampling_pointwise1_bias.bias'), + (('encoder.subsampling.layers.5.weight',), 'subsampling_depthwise2_weight.weights'), + (('encoder.subsampling.layers.5.bias',), 'subsampling_depthwise2_bias.bias'), + (('encoder.subsampling.layers.6.weight',), 'subsampling_pointwise2_weight.weights'), + (('encoder.subsampling.layers.6.bias',), 'subsampling_pointwise2_bias.bias'), + (('encoder.subsampling.linear.weight',), 'subsampling_linear_weight.weights'), + (('encoder.subsampling.linear.bias',), 'subsampling_linear_bias.bias'), + (('ctc_head.weight',), 'ctc_head_weight.weights'), + (('ctc_head.bias',), 'ctc_head_bias.bias'), +] + +PARAKEET_TDT_GLOBAL_WEIGHTS = [ + (('encoder.pre_encode.conv.0.weight', 'encoder.subsampling.layers.0.weight'), 'subsampling_conv0_weight.weights'), + (('encoder.pre_encode.conv.0.bias', 'encoder.subsampling.layers.0.bias'), 'subsampling_conv0_bias.bias'), + (('encoder.pre_encode.conv.2.weight', 'encoder.subsampling.layers.2.weight'), 'subsampling_depthwise1_weight.weights'), + (('encoder.pre_encode.conv.2.bias', 'encoder.subsampling.layers.2.bias'), 'subsampling_depthwise1_bias.bias'), + (('encoder.pre_encode.conv.3.weight', 'encoder.subsampling.layers.3.weight'), 'subsampling_pointwise1_weight.weights'), + (('encoder.pre_encode.conv.3.bias', 'encoder.subsampling.layers.3.bias'), 'subsampling_pointwise1_bias.bias'), + (('encoder.pre_encode.conv.5.weight', 'encoder.subsampling.layers.5.weight'), 'subsampling_depthwise2_weight.weights'), + (('encoder.pre_encode.conv.5.bias', 'encoder.subsampling.layers.5.bias'), 'subsampling_depthwise2_bias.bias'), + (('encoder.pre_encode.conv.6.weight', 'encoder.subsampling.layers.6.weight'), 'subsampling_pointwise2_weight.weights'), + (('encoder.pre_encode.conv.6.bias', 'encoder.subsampling.layers.6.bias'), 'subsampling_pointwise2_bias.bias'), + (('encoder.pre_encode.out.weight', 'encoder.subsampling.linear.weight'), 'subsampling_linear_weight.weights'), + (('encoder.pre_encode.out.bias', 'encoder.subsampling.linear.bias'), 'subsampling_linear_bias.bias'), + (('decoder.prediction.embed.weight', 'decoder.embedding.weight'), 'tdt_predictor_embed.weights'), + (('joint.enc.weight', 'encoder_projector.weight'), 'tdt_joint_enc.weights'), + (('joint.enc.bias', 'encoder_projector.bias'), 'tdt_joint_enc.bias'), + (('joint.pred.weight', 'decoder.decoder_projector.weight'), 'tdt_joint_pred.weights'), + (('joint.pred.bias', 'decoder.decoder_projector.bias'), 'tdt_joint_pred.bias'), + (('joint.joint_net.2.weight', 'joint.joint_net.0.weight', 'joint.head.weight'), 'tdt_joint_out.weights'), + (('joint.joint_net.2.bias', 'joint.joint_net.0.bias', 'joint.head.bias'), 'tdt_joint_out.bias'), +] + +PARAKEET_LAYER_WEIGHTS = [ + (('feed_forward1.linear1.weight',), 'ff1_linear1.weights'), + (('feed_forward1.linear1.bias',), 'ff1_linear1.bias'), + (('feed_forward1.linear2.weight',), 'ff1_linear2.weights'), + (('feed_forward1.linear2.bias',), 'ff1_linear2.bias'), + (('feed_forward2.linear1.weight',), 'ff2_linear1.weights'), + (('feed_forward2.linear1.bias',), 'ff2_linear1.bias'), + (('feed_forward2.linear2.weight',), 'ff2_linear2.weights'), + (('feed_forward2.linear2.bias',), 'ff2_linear2.bias'), + (('self_attn.q_proj.weight', 'self_attn.linear_q.weight', 'self_attention.q_proj.weight', 'self_attention.linear_q.weight'), 'self_attn_q.weights'), + (('self_attn.q_proj.bias', 'self_attn.linear_q.bias', 'self_attention.q_proj.bias', 'self_attention.linear_q.bias'), 'self_attn_q.bias'), + (('self_attn.k_proj.weight', 'self_attn.linear_k.weight', 'self_attention.k_proj.weight', 'self_attention.linear_k.weight'), 'self_attn_k.weights'), + (('self_attn.k_proj.bias', 'self_attn.linear_k.bias', 'self_attention.k_proj.bias', 'self_attention.linear_k.bias'), 'self_attn_k.bias'), + (('self_attn.v_proj.weight', 'self_attn.linear_v.weight', 'self_attention.v_proj.weight', 'self_attention.linear_v.weight'), 'self_attn_v.weights'), + (('self_attn.v_proj.bias', 'self_attn.linear_v.bias', 'self_attention.v_proj.bias', 'self_attention.linear_v.bias'), 'self_attn_v.bias'), + (('self_attn.o_proj.weight', 'self_attn.linear_out.weight', 'self_attention.o_proj.weight', 'self_attention.linear_out.weight'), 'self_attn_output.weights'), + (('self_attn.o_proj.bias', 'self_attn.linear_out.bias', 'self_attention.o_proj.bias', 'self_attention.linear_out.bias'), 'self_attn_output.bias'), + (('self_attn.relative_k_proj.weight', 'self_attn.linear_pos.weight', 'self_attention.relative_k_proj.weight', 'self_attention.linear_pos.weight'), 'self_attn_relative_k.weights'), + (('self_attn.bias_u', 'self_attn.pos_bias_u', 'self_attention.bias_u', 'self_attention.pos_bias_u'), 'self_attn_bias_u.weights'), + (('self_attn.bias_v', 'self_attn.pos_bias_v', 'self_attention.bias_v', 'self_attention.pos_bias_v'), 'self_attn_bias_v.weights'), + (('conv.pointwise_conv1.weight',), 'conv_pointwise1.weights'), + (('conv.pointwise_conv1.bias',), 'conv_pointwise1.bias'), + (('conv.depthwise_conv.weight',), 'conv_depthwise.weights'), + (('conv.depthwise_conv.bias',), 'conv_depthwise.bias'), + (('conv.pointwise_conv2.weight',), 'conv_pointwise2.weights'), + (('conv.pointwise_conv2.bias',), 'conv_pointwise2.bias'), + (('conv.norm.weight', 'conv.batch_norm.weight'), 'conv_batchnorm_weight.weights'), + (('conv.norm.bias', 'conv.batch_norm.bias'), 'conv_batchnorm_bias.bias'), + (('conv.norm.running_mean', 'conv.batch_norm.running_mean'), 'conv_batchnorm_running_mean.weights'), + (('conv.norm.running_var', 'conv.batch_norm.running_var'), 'conv_batchnorm_running_var.weights'), + (('norm_feed_forward1.weight',), 'norm_ff1.weights'), + (('norm_feed_forward1.bias',), 'norm_ff1.bias'), + (('norm_self_att.weight',), 'norm_self_attn.weights'), + (('norm_self_att.bias',), 'norm_self_attn.bias'), + (('norm_conv.weight',), 'norm_conv.weights'), + (('norm_conv.bias',), 'norm_conv.bias'), + (('norm_feed_forward2.weight',), 'norm_ff2.weights'), + (('norm_feed_forward2.bias',), 'norm_ff2.bias'), + (('norm_out.weight',), 'norm_out.weights'), + (('norm_out.bias',), 'norm_out.bias'), +] + + +NEEDLE_GLOBAL_WEIGHTS = [ + ('model.encoder.final_norm.weight', 'encoder_layer_norm_weight.weights'), + ('model.decoder.norm.weight', 'output_norm.weights'), +] + +NEEDLE_ENCODER_LAYER_WEIGHTS = [ + ('ZCRMSNorm_0.scale', 'input_norm.weights', False), + ('attn_gate', 'attn_gate.weights', False), + ('ZCRMSNorm_1.scale', 'post_attn_norm.weights', False), + ('self_attn.q_proj.kernel', 'attn_q.weights', True), + ('self_attn.k_proj.kernel', 'attn_k.weights', True), + ('self_attn.v_proj.kernel', 'attn_v.weights', True), + ('self_attn.out_proj.kernel', 'attn_output.weights', True), + ('self_attn.q_norm.scale', 'attn_q_norm.weights', False), + ('self_attn.k_norm.scale', 'attn_k_norm.weights', False), + ('FeedForward_0.gate_proj.kernel', 'ffn_gate.weights', True), + ('FeedForward_0.up_proj.kernel', 'ffn_up.weights', True), + ('FeedForward_0.down_proj.kernel', 'mlp_fc2.weights', True), +] + +NEEDLE_DECODER_LAYER_WEIGHTS = [ + ('ZCRMSNorm_0.scale', 'input_norm.weights', False), + ('ZCRMSNorm_1.scale', 'post_attn_norm.weights', False), + ('self_attn_gate', 'self_attn_gate.weights', False), + ('cross_attn_gate', 'cross_attn_gate.weights', False), + ('ZCRMSNorm_2.scale', 'final_norm.weights', False), + ('self_attn.q_proj.kernel', 'attn_q.weights', True), + ('self_attn.k_proj.kernel', 'attn_k.weights', True), + ('self_attn.v_proj.kernel', 'attn_v.weights', True), + ('self_attn.out_proj.kernel', 'attn_output.weights', True), + ('self_attn.q_norm.scale', 'attn_q_norm.weights', False), + ('self_attn.k_norm.scale', 'attn_k_norm.weights', False), + ('cross_attn.q_proj.kernel', 'encoder_attn_q.weights', True), + ('cross_attn.k_proj.kernel', 'encoder_attn_k.weights', True), + ('cross_attn.v_proj.kernel', 'encoder_attn_v.weights', True), + ('cross_attn.out_proj.kernel', 'encoder_attn_output.weights', True), + ('cross_attn.q_norm.scale', 'encoder_attn_q_norm.weights', False), + ('cross_attn.k_norm.scale', 'encoder_attn_k_norm.weights', False), + ('FeedForward_0.gate_proj.kernel', 'ffn_gate.weights', True), + ('FeedForward_0.up_proj.kernel', 'ffn_up.weights', True), + ('FeedForward_0.down_proj.kernel', 'mlp_fc2.weights', True), +] + + +def get_layer_weight_patterns(i, precision, model_type=None, skip_kv=False): + is_whisper = model_type == 'whisper' + is_qwen_family = isinstance(model_type, str) and ('qwen' in model_type) + is_youtu = model_type == 'youtu' + + patterns = [ + # Youtu MLA attention weights + (['self_attn.q_a_proj.weight'], precision, f'layer_{i}_attn_q_a.weights', False) if is_youtu else None, + (['self_attn.q_a_layernorm.weight'], precision, f'layer_{i}_attn_q_a_norm.weights', False) if is_youtu else None, + (['self_attn.q_b_proj.weight'], precision, f'layer_{i}_attn_q_b.weights', False) if is_youtu else None, + (['self_attn.kv_a_proj_with_mqa.weight'], precision, f'layer_{i}_attn_kv_a.weights', False) if is_youtu else None, + (['self_attn.kv_a_layernorm.weight'], precision, f'layer_{i}_attn_kv_a_norm.weights', False) if is_youtu else None, + (['self_attn.kv_b_proj.weight'], precision, f'layer_{i}_attn_kv_b.weights', False) if is_youtu else None, + (['self_attn.q_proj.weight', 'attn.q_proj.weight', 'attn.c_attn.weight'], precision, f'layer_{i}_attn_q.weights', False) if not is_whisper and not is_youtu else None, + (['self_attn.k_proj.weight', 'attn.k_proj.weight'], precision, f'layer_{i}_attn_k.weights', False) if not is_whisper and not skip_kv and not is_youtu else None, + (['self_attn.v_proj.weight', 'attn.v_proj.weight'], precision, f'layer_{i}_attn_v.weights', False) if not is_whisper and not skip_kv and not is_youtu else None, + (['self_attn.o_proj.weight', 'attn.o_proj.weight', 'attn.c_proj.weight', 'self_attn.out_proj.weight'], precision, f'layer_{i}_attn_output.weights', False) if not is_whisper else None, + # Qwen3.5 linear-attention path + (['linear_attn.in_proj_qkv.weight'], precision, f'layer_{i}_linear_attn_qkv.weights', False) if is_qwen_family else None, + (['linear_attn.in_proj_a.weight'], precision, f'layer_{i}_linear_attn_a.weights', False) if is_qwen_family else None, + (['linear_attn.in_proj_b.weight'], precision, f'layer_{i}_linear_attn_b.weights', False) if is_qwen_family else None, + (['linear_attn.in_proj_z.weight'], precision, f'layer_{i}_linear_attn_z.weights', False) if is_qwen_family else None, + (['linear_attn.out_proj.weight'], precision, f'layer_{i}_linear_attn_output.weights', False) if is_qwen_family else None, + (['linear_attn.norm.weight'], precision, f'layer_{i}_linear_attn_norm.weights', False) if is_qwen_family else None, + (['linear_attn.conv1d.weight'], precision, f'layer_{i}_linear_attn_conv1d.weights', False) if is_qwen_family else None, + (['linear_attn.A_log'], precision, f'layer_{i}_linear_attn_A_log.weights', False) if is_qwen_family else None, + (['linear_attn.dt_bias'], precision, f'layer_{i}_linear_attn_dt_bias.weights', False) if is_qwen_family else None, + ([ + 'self_attn.deltanet_gate_proj.weight', + 'self_attn.gated_deltanet_gate_proj.weight', + 'self_attn.attn_gate_proj.weight', + 'self_attn.f_gate_proj.weight', + 'self_attn.attn_f_gate_proj.weight', + 'self_attn.attn_gate.weight', + 'self_attn.attn_f_gate.weight', + ], precision, f'layer_{i}_deltanet_gate.weights', False) if is_qwen_family else None, + ([ + 'self_attn.deltanet_beta_proj.weight', + 'self_attn.gated_deltanet_beta_proj.weight', + 'self_attn.attn_beta_proj.weight', + 'self_attn.f_beta_proj.weight', + 'self_attn.attn_f_beta_proj.weight', + 'self_attn.attn_beta.weight', + 'self_attn.attn_f_beta.weight', + ], precision, f'layer_{i}_deltanet_beta.weights', False) if is_qwen_family else None, + (['input_layernorm.weight', 'ln_1.weight', 'operator_norm.weight'], precision, f'layer_{i}_input_norm.weights', False), + (['self_attn.q_norm.weight', 'self_attn.q_layernorm.weight'], precision, f'layer_{i}_attn_q_norm.weights', False), + (['self_attn.k_norm.weight', 'self_attn.k_layernorm.weight'], precision, f'layer_{i}_attn_k_norm.weights', False) if not skip_kv else None, + (['mlp.gate_proj.weight', 'mlp.c_fc.weight', 'feed_forward.w1.weight', 'ff.ff_proj.weight'], precision, f'layer_{i}_ffn_gate.weights', False), + (['mlp.up_proj.weight', 'feed_forward.w3.weight', 'ff.ff_noact.weight'], precision, f'layer_{i}_ffn_up.weights', False), + (['mlp.down_proj.weight', 'mlp.c_proj.weight', 'feed_forward.w2.weight', 'ff.ff_out.weight'], precision, f'layer_{i}_ffn_down.weights', False), + (['feed_forward.gate.weight'], precision, f'layer_{i}_moe_router.weights', False), + (['feed_forward.expert_bias'], precision, f'layer_{i}_moe_expert_bias.weights', False), + (['feed_forward.experts.{channel}.w1.weight'], precision, f'layer_{i}_moe_expert_{{channel}}_w1.weights', False), + (['feed_forward.experts.{channel}.w3.weight'], precision, f'layer_{i}_moe_expert_{{channel}}_w3.weights', False), + (['feed_forward.experts.{channel}.w2.weight'], precision, f'layer_{i}_moe_expert_{{channel}}_w2.weights', False), + (['moe.gate_proj'], precision, f'layer_{i}_moe_gate_proj.weights', False), + (['moe.up_proj'], precision, f'layer_{i}_moe_up_proj.weights', False), + (['moe.down_proj'], precision, f'layer_{i}_moe_down_proj.weights', False), + (['moe.per_expert_scale'], 'FP16', f'layer_{i}_moe_per_expert_scale.weights', False), + (['router.proj.weight'], precision, f'layer_{i}_router_proj.weights', False), + (['router.scale'], 'FP16', f'layer_{i}_router_scale.weights', False), + (['post_attention_layernorm.weight', 'ln_2.weight', 'ffn_norm.weight', 'norm2.weight'], precision, f'layer_{i}_post_attn_norm.weights', False), + (['pre_feedforward_layernorm.weight'], precision, f'layer_{i}_pre_ffn_norm.weights', False), + (['post_feedforward_layernorm.weight'], precision, f'layer_{i}_post_ffn_norm.weights', False), + (['post_feedforward_layernorm_1.weight'], precision, f'layer_{i}_post_ffn_norm_1.weights', False), + (['post_feedforward_layernorm_2.weight'], precision, f'layer_{i}_post_ffn_norm_2.weights', False), + (['pre_feedforward_layernorm_2.weight'], precision, f'layer_{i}_pre_ffn_norm_2.weights', False), + (['conv.in_proj.weight'], precision, f'layer_{i}_conv_in_proj.weights', False), + (['conv.out_proj.weight'], precision, f'layer_{i}_conv_out_proj.weights', False), + (['conv.conv.weight'], precision, f'layer_{i}_conv_depthwise.weights', False), + (['attn.Wqkv.bias'], precision, f'layer_{i}_attn_{{channel}}.bias', False), + (['attn.Wqkv.weight'], precision, f'layer_{i}_attn_{{channel}}.weights', False), + (['attn.out_proj.bias', 'attention.to_out.bias'], precision, f'layer_{i}_attn_output.bias', False), + (['attn.out_proj.weight', 'attention.to_out.weight'], precision, f'layer_{i}_attn_output.weights', False), + (['mlp.fc1.bias', 'ff.ff.0.bias'], precision, f'layer_{i}_mlp_fc1.bias', False), + (['mlp.fc1.weight', 'ff.ff.0.weight'], precision, f'layer_{i}_mlp_fc1.weights', False), + (['mlp.fc2.bias', 'ff.ff.2.bias'], precision, f'layer_{i}_mlp_fc2.bias', False), + (['mlp.fc2.weight', 'ff.ff.2.weight'], precision, f'layer_{i}_mlp_fc2.weights', False), + (['norm1.bias'], precision, f'layer_{i}_norm1.bias', False), + (['norm1.weight'], precision, f'layer_{i}_norm1.weights', False), + (['norm2.bias'], precision, f'layer_{i}_norm2.bias', False), + (['norm2.weight'], precision, f'layer_{i}_norm2.weights', False), + (['mlp.experts.bias'], precision, f'layer_{i}_mlp_experts.bias', False), + (['mlp.experts.mlp.w1'], precision, f'layer_{i}_mlp_expert_{{channel}}.mlp1.weights', False), + (['mlp.experts.mlp.w2'], precision, f'layer_{i}_mlp_expert_{{channel}}.mlp2.weights', True), + (['mlp.router.layer.weight'], precision, f'layer_{i}_mlp_router.layer.weights', False), + (['encoder_attn.q_proj.weight', 'attention.to_q.weight'], precision, f'layer_{i}_encoder_attn_q.weights', False), + (['encoder_attn.k_proj.weight', 'attention.to_k.weight'], precision, f'layer_{i}_encoder_attn_k.weights', False), + (['encoder_attn.v_proj.weight', 'attention.to_v.weight'], precision, f'layer_{i}_encoder_attn_v.weights', False), + (['encoder_attn.out_proj.weight', 'encoder_attn.o_proj.weight'], precision, f'layer_{i}_encoder_attn_output.weights', False), + (['encoder_attn.q_proj.bias'], precision, f'layer_{i}_encoder_attn_q.bias', False), + (['encoder_attn.v_proj.bias'], precision, f'layer_{i}_encoder_attn_v.bias', False), + (['encoder_attn.out_proj.bias', 'encoder_attn.o_proj.bias'], precision, f'layer_{i}_encoder_attn_output.bias', False), + (['encoder_attn_layer_norm.weight'], precision, f'layer_{i}_encoder_attn_norm.weights', False), + (['encoder_attn_layer_norm.bias'], precision, f'layer_{i}_encoder_attn_norm.bias', False), + (['fc1.weight'], precision, f'layer_{i}_mlp_fc1.weights', False), + (['fc1.bias'], precision, f'layer_{i}_mlp_fc1.bias', False), + (['fc2.weight'], precision, f'layer_{i}_mlp_fc2.weights', False), + (['fc2.bias'], precision, f'layer_{i}_mlp_fc2.bias', False), + (['final_layer_norm.weight', 'final_layernorm.weight'], precision, f'layer_{i}_final_norm.weights', False), + (['final_layer_norm.bias', 'final_layernorm.bias'], precision, f'layer_{i}_final_norm.bias', False), + + # Whisper-only: separate self_attn_* outputs (non-Whisper uses attn_* above) + (['self_attn.q_proj.weight'], precision, f'layer_{i}_self_attn_q.weights', False) if is_whisper else None, + (['self_attn.k_proj.weight'], precision, f'layer_{i}_self_attn_k.weights', False) if is_whisper else None, + (['self_attn.v_proj.weight'], precision, f'layer_{i}_self_attn_v.weights', False) if is_whisper else None, + (['self_attn.q_proj.bias'], precision, f'layer_{i}_self_attn_q.bias', False) if is_whisper else None, + (['self_attn.v_proj.bias'], precision, f'layer_{i}_self_attn_v.bias', False) if is_whisper else None, + (['self_attn.out_proj.weight'], precision, f'layer_{i}_self_attn_output.weights', False) if is_whisper else None, + (['self_attn.out_proj.bias'], precision, f'layer_{i}_self_attn_output.bias', False) if is_whisper else None, + (['self_attn_layer_norm.weight'], precision, f'layer_{i}_self_attn_norm.weights', False), + (['self_attn_layer_norm.bias'], precision, f'layer_{i}_self_attn_norm.bias', False), + (['altup.router_norm.weight'], precision, f'layer_{i}_altup_router_norm.weights', False), + (['altup.prediction_coefs.weight'], 'FP16', f'layer_{i}_altup_prediction_coefs.weights', False), + (['altup.correction_coefs.weight'], 'FP16', f'layer_{i}_altup_correction_coefs.weights', False), + (['altup.correct_output_scale'], 'FP16', f'layer_{i}_altup_correct_output_scale.weights', False), + (['altup.modality_router.weight'], precision, f'layer_{i}_altup_modality_router.weights', False), + (['laurel.linear_left.weight'], precision, f'layer_{i}_laurel_left.weights', False), + (['laurel.linear_right.weight'], precision, f'layer_{i}_laurel_right.weights', False), + (['laurel.post_laurel_norm.weight'], precision, f'layer_{i}_laurel_norm.weights', False), + (['per_layer_projection.weight'], precision, f'layer_{i}_per_layer_proj.weights', False), + (['per_layer_input_gate.weight'], precision, f'layer_{i}_per_layer_gate.weights', False), + (['post_per_layer_input_norm.weight'], precision, f'layer_{i}_post_per_layer_norm.weights', False), + (['layer_scalar'], 'FP16', f'layer_{i}_layer_scalar.weights', False), + ] + + return [p for p in patterns if p is not None] + + +def get_vision_layer_weights(i_v, vpref): + return [ + (vpref + 'layer_norm1.weight', f'vision_layer_{i_v}_layer_norm1.weights'), + (vpref + 'layer_norm1.bias', f'vision_layer_{i_v}_layer_norm1.bias.weights'), + (vpref + 'layer_norm2.weight', f'vision_layer_{i_v}_layer_norm2.weights'), + (vpref + 'layer_norm2.bias', f'vision_layer_{i_v}_layer_norm2.bias.weights'), + (vpref + 'mlp.fc1.weight', f'vision_layer_{i_v}_ffn_fc1.weights'), + (vpref + 'mlp.fc1.bias', f'vision_layer_{i_v}_ffn_fc1.bias.weights'), + (vpref + 'mlp.fc2.weight', f'vision_layer_{i_v}_ffn_fc2.weights'), + (vpref + 'mlp.fc2.bias', f'vision_layer_{i_v}_ffn_fc2.bias.weights'), + (vpref + 'self_attn.q_proj.weight', f'vision_layer_{i_v}_self_attn_q.weights'), + (vpref + 'self_attn.k_proj.weight', f'vision_layer_{i_v}_self_attn_k.weights'), + (vpref + 'self_attn.v_proj.weight', f'vision_layer_{i_v}_self_attn_v.weights'), + (vpref + 'self_attn.out_proj.weight', f'vision_layer_{i_v}_self_attn_out.weights'), + (vpref + 'self_attn.q_proj.bias', f'vision_layer_{i_v}_self_attn_q.bias.weights'), + (vpref + 'self_attn.k_proj.bias', f'vision_layer_{i_v}_self_attn_k.bias.weights'), + (vpref + 'self_attn.v_proj.bias', f'vision_layer_{i_v}_self_attn_v.bias.weights'), + (vpref + 'self_attn.out_proj.bias', f'vision_layer_{i_v}_self_attn_out.bias.weights'), + ] diff --git a/python/cactus/convert/calibration/__init__.py b/python/cactus/convert/calibration/__init__.py new file mode 100644 index 000000000..0fa427ead --- /dev/null +++ b/python/cactus/convert/calibration/__init__.py @@ -0,0 +1,2 @@ +"""Calibration example loading and Hessian collection.""" + diff --git a/python/cactus/convert/calibration/hessian.py b/python/cactus/convert/calibration/hessian.py new file mode 100644 index 000000000..62e109a8d --- /dev/null +++ b/python/cactus/convert/calibration/hessian.py @@ -0,0 +1,440 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +import hashlib +import json +from pathlib import Path +import time + +try: + import torch +except Exception: # pragma: no cover + torch = None + + +@dataclass +class HessianStats: + hessians: dict[str, object] = field(default_factory=dict) + diag: dict[str, object] = field(default_factory=dict) + samples: dict[str, int] = field(default_factory=dict) + unresolved_targets: list[str] = field(default_factory=list) + errors: dict[str, int] = field(default_factory=dict) + error_samples: list[dict[str, str]] = field(default_factory=list) + rows: dict[str, int] = field(default_factory=dict) + timings: dict[str, float] = field(default_factory=dict) + + +def _flatten_inputs(x): + if isinstance(x, (tuple, list)): + x = x[0] + if x.ndim == 1: + x = x.unsqueeze(0) + return x.reshape(-1, x.shape[-1]).detach().float() + + +def collect_text_hessians(model, tokenizer, texts: list[str], target_names: set[str], device: str, max_length: int = 512) -> HessianStats: + if torch is None: + return HessianStats() + stats = HessianStats() + handles = [] + modules = dict(model.named_modules()) + for name in target_names: + module = modules.get(name) + if module is None or not hasattr(module, "weight"): + continue + + def hook(_module, inputs, _name=name): + x = _flatten_inputs(inputs) + if x.numel() == 0: + return + h = (x.T @ x).detach().to("cpu") + if _name not in stats.hessians: + stats.hessians[_name] = h + stats.diag[_name] = torch.diag(h).clone() + stats.samples[_name] = int(x.shape[0]) + else: + stats.hessians[_name] += h + stats.diag[_name] += torch.diag(h) + stats.samples[_name] += int(x.shape[0]) + + handles.append(module.register_forward_pre_hook(hook)) + if not handles or not texts: + return stats + try: + model.eval() + model.to(device) + with torch.no_grad(): + for text in texts: + batch = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length) + batch = {k: v.to(device) for k, v in batch.items()} + try: + model(**batch) + except Exception: + continue + finally: + for h in handles: + h.remove() + for name, count in list(stats.samples.items()): + if count > 0: + stats.hessians[name] = stats.hessians[name] / float(count) + stats.diag[name] = stats.diag[name] / float(count) + return stats + + +def _load_audio_16k(path: Path): + import numpy as np + import soundfile as sf + + audio, sr = sf.read(path) + audio = np.asarray(audio, dtype=np.float32) + if audio.ndim > 1: + audio = audio.mean(axis=-1) + if sr != 16000: + old_x = np.linspace(0.0, 1.0, num=audio.shape[0], endpoint=False) + new_len = max(1, int(round(audio.shape[0] * 16000 / sr))) + new_x = np.linspace(0.0, 1.0, num=new_len, endpoint=False) + audio = np.interp(new_x, old_x, audio).astype(np.float32) + return audio + + +def _read_jsonl(path: str | Path, limit: int | None): + import json + + if limit is not None and limit <= 0: + return [] + rows = [] + with Path(path).open("r", encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + rows.append(json.loads(line)) + if limit is not None and len(rows) >= limit: + break + return rows + + +def _row_language_text(row: dict) -> str: + if "messages" in row: + return "\n".join(f"{m.get('role', 'user')}: {m.get('content', '')}" for m in row["messages"]) + if "prompt_text" in row or "completion_text" in row: + return (row.get("prompt_text", "") + row.get("completion_text", "")).strip() + return (row.get("prompt", "") + row.get("completion", "")).strip() + + +def _register_hooks(model, target_names: set[str]) -> tuple[HessianStats, list, dict[str, object]]: + stats = HessianStats() + handles = [] + gpu_accum: dict[str, object] = {} + modules = dict(model.named_modules()) + + def add_cpu_hessian(name: str, h): + h = h.detach().to("cpu") + if name not in stats.hessians: + stats.hessians[name] = h + stats.diag[name] = torch.diag(h).clone() + else: + stats.hessians[name] += h + stats.diag[name] += torch.diag(h) + + for name in target_names: + module = modules.get(name) + if module is None or not hasattr(module, "weight"): + stats.unresolved_targets.append(name) + continue + + def hook(_module, inputs, _name=name): + x = _flatten_inputs(inputs) + if x.numel() == 0: + return + h = (x.T @ x).detach() + if h.is_cuda: + if _name not in gpu_accum: + gpu_accum[_name] = h + else: + gpu_accum[_name] += h + else: + add_cpu_hessian(_name, h) + stats.samples[_name] = stats.samples.get(_name, 0) + int(x.shape[0]) + + handles.append(module.register_forward_pre_hook(hook)) + return stats, handles, {"gpu_accum": gpu_accum, "add_cpu_hessian": add_cpu_hessian} + + +def _record_error(stats: HessianStats, exc: Exception, context: str) -> None: + key = type(exc).__name__ + stats.errors[key] = stats.errors.get(key, 0) + 1 + if len(stats.error_samples) < 5: + stats.error_samples.append({"context": context, "type": key, "message": str(exc)[:500]}) + + +def collect_manifest_hessians( + model, + processor, + manifest: dict, + limits: dict[str, int | None], + target_names: set[str], + device: str, + max_length: int = 1024, + adapter=None, + batch_size: int = 1, + gpu_flush_interval: int = 16, + progress_path: str | Path | None = None, + preprocessed_cache_dir: str | Path | None = None, +) -> HessianStats: + if torch is None: + return HessianStats() + started = time.perf_counter() + stats, handles, hook_state = _register_hooks(model, target_names) + if not handles: + return stats + model.eval() + model.to(device) + + tokenizer = getattr(processor, "tokenizer", processor) + batch_size = max(1, int(batch_size or 1)) + gpu_flush_interval = max(1, int(gpu_flush_interval or 1)) + progress_file = Path(progress_path) if progress_path else None + cache_dir = Path(preprocessed_cache_dir) if preprocessed_cache_dir else None + if cache_dir is not None: + cache_dir.mkdir(parents=True, exist_ok=True) + forwards = 0 + + def add_time(name: str, delta: float): + stats.timings[name] = stats.timings.get(name, 0.0) + float(delta) + + def flush_gpu(): + gpu_accum = hook_state["gpu_accum"] + add_cpu_hessian = hook_state["add_cpu_hessian"] + if not gpu_accum: + return + t0 = time.perf_counter() + for name, h in list(gpu_accum.items()): + add_cpu_hessian(name, h) + gpu_accum.clear() + if device != "cpu" and torch.cuda.is_available(): + torch.cuda.synchronize(device) + add_time("hessian_gpu_flush_seconds", time.perf_counter() - t0) + + def write_progress(): + if progress_file is None: + return + payload = { + "elapsed_seconds": time.perf_counter() - started, + "forwards": forwards, + "rows": dict(stats.rows), + "target_modules": len(target_names), + "resolved_modules": len(handles), + "modules_with_samples": len(stats.samples), + "sample_counts_min": min(stats.samples.values()) if stats.samples else 0, + "sample_counts_max": max(stats.samples.values()) if stats.samples else 0, + "unresolved_targets": len(stats.unresolved_targets), + "errors": dict(stats.errors), + "timings": dict(stats.timings), + } + progress_file.parent.mkdir(parents=True, exist_ok=True) + progress_file.write_text(json.dumps(payload, indent=2), encoding="utf-8") + + def maybe_flush(): + nonlocal forwards + forwards += 1 + if forwards % gpu_flush_interval == 0: + flush_gpu() + write_progress() + + def cache_key(kind: str, source: str, payload: object) -> Path | None: + if cache_dir is None: + return None + raw = json.dumps({"kind": kind, "source": source, "payload": payload}, sort_keys=True, default=str).encode("utf-8") + return cache_dir / f"{hashlib.sha1(raw).hexdigest()}.pt" + + def tensor_cache_load(path: Path | None): + if path is None or not path.exists(): + return None + t0 = time.perf_counter() + try: + obj = torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + obj = torch.load(path, map_location="cpu") + add_time("preprocess_cache_read_seconds", time.perf_counter() - t0) + return obj + + def tensor_cache_save(path: Path | None, obj) -> None: + if path is None or path.exists(): + return + t0 = time.perf_counter() + torch.save(obj, path) + add_time("preprocess_cache_write_seconds", time.perf_counter() - t0) + + def run_inputs(inputs, context: str): + try: + model_dtype = next(model.parameters()).dtype + except Exception: + model_dtype = None + t0 = time.perf_counter() + inputs = { + k: ( + v.to(device=device, dtype=model_dtype) + if model_dtype is not None and torch.is_tensor(v) and v.is_floating_point() + else (v.to(device) if hasattr(v, "to") else v) + ) + for k, v in inputs.items() + } + add_time("input_to_device_seconds", time.perf_counter() - t0) + with torch.no_grad(): + try: + t0 = time.perf_counter() + model(**inputs) + if device != "cpu" and torch.cuda.is_available(): + torch.cuda.synchronize(device) + add_time(f"{context}_forward_seconds", time.perf_counter() - t0) + maybe_flush() + except Exception as exc: + _record_error(stats, exc, context) + return + + def run_text_batch(texts: list[str]): + if not texts: + return + stats.rows["language"] = stats.rows.get("language", 0) + len(texts) + key = cache_key("language", "tokenizer", texts) + inputs = tensor_cache_load(key) + if inputs is None: + t0 = time.perf_counter() + try: + inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=max_length) + except Exception: + if len(texts) == 1: + inputs = tokenizer(texts[0], return_tensors="pt", truncation=True, max_length=max_length) + add_time("language_preprocess_seconds", time.perf_counter() - t0) + tensor_cache_save(key, inputs) + run_inputs(inputs, "language") + return + for text in texts: + run_text_batch([text]) + return + add_time("language_preprocess_seconds", time.perf_counter() - t0) + tensor_cache_save(key, inputs) + run_inputs(inputs, "language") + + try: + lang = manifest.get("language") + if lang: + path = lang["path"] if isinstance(lang, dict) else lang + text_batch: list[str] = [] + for row in _read_jsonl(path, limits.get("language")): + text = _row_language_text(row) + if text: + text_batch.append(text) + if len(text_batch) >= batch_size: + run_text_batch(text_batch) + text_batch = [] + continue + if text_batch: + run_text_batch(text_batch) + text_batch = [] + ids = row.get("full_token_ids") or ((row.get("prompt_token_ids") or []) + (row.get("completion_token_ids") or [])) + if ids: + stats.rows["language_ids"] = stats.rows.get("language_ids", 0) + 1 + run_inputs({"input_ids": torch.tensor([ids[:max_length]], dtype=torch.long)}, "language_ids") + if text_batch: + run_text_batch(text_batch) + + vision = manifest.get("vision") + if vision and hasattr(processor, "apply_chat_template"): + from PIL import Image + + base = Path(vision.get("base_dir", ".")) + vision_batch = [] + vision_keys = [] + + def run_vision_batch(batch, keys): + if not batch: + return + stats.rows["vision"] = stats.rows.get("vision", 0) + len(batch) + key = cache_key("vision", "processor", keys) + inputs = tensor_cache_load(key) + if inputs is None: + t0 = time.perf_counter() + try: + inputs = processor.apply_chat_template(batch, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt") + except Exception: + if len(batch) == 1: + raise + stats.rows["vision"] -= len(batch) + for item, item_key in zip(batch, keys): + run_vision_batch([item], [item_key]) + return + add_time("vision_preprocess_seconds", time.perf_counter() - t0) + tensor_cache_save(key, inputs) + run_inputs(inputs, "vision") + + for row in _read_jsonl(vision["path"], limits.get("vision")): + image_rel = row.get("image_path") or (row.get("image_paths") or [None])[0] + if not image_rel: + continue + image_path = base / image_rel + if not image_path.exists(): + continue + image = Image.open(image_path).convert("RGB") + prompt = row.get("prompt_text", row.get("prompt", "")) + vision_batch.append([{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}]) + vision_keys.append({"image_path": str(image_path), "prompt": prompt}) + if len(vision_batch) >= batch_size: + run_vision_batch(vision_batch, vision_keys) + vision_batch = [] + vision_keys = [] + if vision_batch: + run_vision_batch(vision_batch, vision_keys) + + audio = manifest.get("audio") + if audio and hasattr(processor, "apply_chat_template"): + base = Path(audio.get("base_dir", ".")) + for row in _read_jsonl(audio["path"], limits.get("audio")): + audio_rel = row.get("audio_path") + if not audio_rel: + continue + audio_path = base / audio_rel + if not audio_path.exists(): + continue + audio_arr = _load_audio_16k(audio_path) + messages = [[{"role": "user", "content": [{"type": "audio", "audio": audio_arr}, {"type": "text", "text": row.get("prompt_text", row.get("prompt", ""))}]}]] + stats.rows["audio"] = stats.rows.get("audio", 0) + 1 + run_inputs(processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"), "audio") + + transcription = manifest.get("transcription") + if transcription: + base = Path(transcription.get("base_dir", ".")) + for row in _read_jsonl(transcription["path"], limits.get("transcription")): + inputs = None + if adapter is not None: + try: + inputs = adapter.build_calibration_inputs(row, processor, "transcription", base) + except Exception as exc: + _record_error(stats, exc, "transcription_builder") + inputs = None + if inputs is None and hasattr(processor, "apply_chat_template"): + audio_rel = row.get("audio_path") + if not audio_rel: + continue + audio_path = base / audio_rel + if not audio_path.exists(): + continue + audio_arr = _load_audio_16k(audio_path) + messages = [[{"role": "user", "content": [{"type": "audio", "audio": audio_arr}, {"type": "text", "text": row.get("prompt_text", row.get("prompt", ""))}]}]] + inputs = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt") + if inputs is not None: + stats.rows["transcription"] = stats.rows.get("transcription", 0) + 1 + run_inputs(inputs, "transcription") + finally: + flush_gpu() + write_progress() + for h in handles: + h.remove() + + for name, count in list(stats.samples.items()): + if count > 0: + stats.hessians[name] = stats.hessians[name] / float(count) + stats.diag[name] = stats.diag[name] / float(count) + stats.timings["total_seconds"] = time.perf_counter() - started + write_progress() + return stats diff --git a/python/cactus/convert/calibration/loaders.py b/python/cactus/convert/calibration/loaders.py new file mode 100644 index 000000000..e8844347f --- /dev/null +++ b/python/cactus/convert/calibration/loaders.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + + +def read_jsonl(path: str | Path, limit: int | None = None) -> list[dict[str, Any]]: + rows = [] + with Path(path).open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + rows.append(json.loads(line)) + if limit is not None and len(rows) >= limit: + break + return rows + + +def load_manifest(path: str | Path | None) -> dict[str, Any]: + if path is None: + return {} + with Path(path).open("r", encoding="utf-8") as f: + return json.load(f) + + +def language_text(row: dict[str, Any]) -> str: + if "messages" in row: + return "\n".join(f"{m.get('role', 'user')}: {m.get('content', '')}" for m in row["messages"]) + if "prompt_text" in row or "completion_text" in row: + return (row.get("prompt_text", "") + row.get("completion_text", "")).strip() + return (row.get("prompt", "") + row.get("completion", "")).strip() + + +def embedding_texts(row: dict[str, Any]) -> list[str]: + if "texts" in row: + return [str(x) for x in row["texts"]] + if "text" in row: + return [str(row["text"])] + return [] + + +def load_text_calibration(manifest: dict[str, Any], limits: dict[str, int | None]) -> list[str]: + texts: list[str] = [] + language = manifest.get("language") + if language: + path = language["path"] if isinstance(language, dict) else language + texts.extend(language_text(r) for r in read_jsonl(path, limits.get("language"))) + embedding = manifest.get("embedding") + if embedding: + path = embedding["path"] if isinstance(embedding, dict) else embedding + for row in read_jsonl(path, limits.get("embedding")): + texts.extend(embedding_texts(row)) + return [t for t in texts if t] diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py new file mode 100644 index 000000000..3a494f961 --- /dev/null +++ b/python/cactus/convert/cli.py @@ -0,0 +1,645 @@ +from __future__ import annotations + +import argparse +import json +import shutil +from collections.abc import Mapping +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np + +try: + import torch +except Exception: # pragma: no cover + torch = None + +from .calibration.hessian import collect_manifest_hessians +from .calibration.loaders import language_text, load_manifest, read_jsonl +from .cactus_adapters.config_utils import ( + cfg_get, +) +from .cactus_adapters.tensor_io import ( + save_depthwise_conv_int8_with_header, + save_pointwise_conv1d_int8_with_header, + save_tensor_with_header, +) +from .export.files import copy_runtime_files, write_config_txt +from .export.qdq import convert_qdq +from .export.reports import print_summary, write_reports +from .export.validate import validate_qdq +from .model_adapters.detection import SUPPORTED_FAMILIES, detect_family +from .model_adapters.adapters import adapter_for_family +from .model_adapters.nemo import ensure_parakeet_tdt_nemo_source +from .quantization.cq import quantize_hadamard, quantize_orthogonal, write_cq_tensor +from .compat import patch_transformers_import_compat + +ALPHA = 0.25 + + +def _hf_cache_dir() -> str | None: + import os + + return os.environ.get("HF_HUB_CACHE") or os.environ.get("HF_HOME") or None + + +def _load_hf(model_id_or_path: str, device: str, *, skip_model_load: bool = False): + import logging, warnings + logging.getLogger("transformers").setLevel(logging.ERROR) + warnings.filterwarnings("ignore", message=".*You are using a model of type.*") + for note in patch_transformers_import_compat(): + print(f"note={note}") + from ..models.needle import register_with_transformers + + register_with_transformers() + nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) + if nemo_export is not None: + model_id_or_path = nemo_export + try: + from transformers import AutoConfig, AutoModel + except Exception as exc: # pragma: no cover + raise RuntimeError("transformers is required for conversion") from exc + try: + cfg = AutoConfig.from_pretrained(model_id_or_path, trust_remote_code=True, local_files_only=Path(model_id_or_path).exists()) + except Exception: + try: + if Path(model_id_or_path).exists(): + cfg_path = Path(model_id_or_path) / "config.json" + else: + from huggingface_hub import hf_hub_download + cfg_path = Path(hf_hub_download(model_id_or_path, "config.json", cache_dir=_hf_cache_dir())) + cfg = json.loads(cfg_path.read_text(encoding="utf-8")) + except Exception as exc: + raise RuntimeError(f"Could not load model {model_id_or_path!r}: {exc}") from exc + family = detect_family(cfg, "auto") + adapter = adapter_for_family(family) + processor = adapter.load_processor(model_id_or_path) + if skip_model_load: + return cfg, processor, None + model_cls = adapter.model_class(cfg) + model_type = str(cfg_get(cfg, "model_type", "") or "").lower() + if isinstance(cfg, dict) and model_type == "parakeet_tdt": + return cfg, processor, None + try: + model = model_cls.from_pretrained( + model_id_or_path, + torch_dtype=torch.float16 if torch is not None else None, + trust_remote_code=True, + low_cpu_mem_usage=True, + local_files_only=Path(model_id_or_path).exists(), + ) + except Exception as primary_exc: + try: + model = AutoModel.from_pretrained( + model_id_or_path, + torch_dtype=torch.float16 if torch is not None else None, + trust_remote_code=True, + low_cpu_mem_usage=True, + local_files_only=Path(model_id_or_path).exists(), + ) + except Exception as auto_exc: + print(f"warning: could not load model object ({model_cls.__name__}: {primary_exc}; AutoModel: {auto_exc})") + model = None + if model is None: + return cfg, processor, model + model.eval() + if torch is not None and device != "cpu": + try: + model.to(device) + except Exception: + model.to("cpu") + return cfg, processor, model + + +def _module_name(param_name: str) -> str: + return param_name[:-7] if param_name.endswith(".weight") else param_name + + +def _tensor_shape(tensor) -> tuple[int, ...]: + return tuple(int(x) for x in tensor.shape) + + +def _tensor_bytes(path: Path | None) -> int: + return path.stat().st_size if path and path.exists() else 0 + + +def _bits_for_component(component: str, args: argparse.Namespace) -> int: + if component == "language" and args.language_bits is not None: + return int(args.language_bits) + if component == "vision" and args.vision_bits is not None: + return int(args.vision_bits) + if component in {"audio", "transcription"} and args.audio_bits is not None: + return int(args.audio_bits) + if component == "embedding" and args.embedding_bits is not None: + return int(args.embedding_bits) + return int(args.bits) + + +def _load_checkpoint_state_dict(model_id_or_path: str) -> dict[str, Any] | None: + nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) + if nemo_export is not None: + model_id_or_path = nemo_export + root = Path(model_id_or_path) + if not root.exists() or not root.is_dir(): + try: + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + try: + model_file = hf_hub_download(model_id_or_path, "model.safetensors", cache_dir=_hf_cache_dir()) + return load_file(model_file) + except Exception: + index_file = hf_hub_download(model_id_or_path, "model.safetensors.index.json", cache_dir=_hf_cache_dir()) + index = json.loads(Path(index_file).read_text(encoding="utf-8")) + shard_names = sorted(set(index.get("weight_map", {}).values())) + if not shard_names: + return None + state: dict[str, Any] = {} + for shard in shard_names: + shard_path = hf_hub_download(model_id_or_path, shard, cache_dir=_hf_cache_dir()) + for key, tensor in load_file(shard_path).items(): + if key in state: + raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") + state[key] = tensor + return state + except RuntimeError: + raise + except Exception: + return None + safetensor_files = sorted(root.glob("*.safetensors")) + if safetensor_files: + try: + from safetensors import safe_open + except Exception: + return None + state: dict[str, Any] = {} + for file in safetensor_files: + with safe_open(file, framework="pt", device="cpu") as sf: + for key in sf.keys(): + if key in state: + raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") + state[key] = sf.get_tensor(key) + return state + if torch is None: + return None + bin_files = sorted(root.glob("pytorch_model*.bin")) + if not bin_files: + return None + state = {} + for file in bin_files: + shard = torch.load(file, map_location="cpu") + shard_state = shard.get("state_dict", shard) if isinstance(shard, dict) else shard + for key, tensor in shard_state.items(): + if key in state: + raise RuntimeError(f"duplicate tensor key {key!r} across checkpoint shards") + state[key] = tensor + return state + + +def _augment_state_dict_for_family(state_dict: dict[str, Any], family: str) -> dict[str, Any]: + return adapter_for_family(family).normalize_state_dict(state_dict).state_dict + + +def _save_fallback_tensor(tensor, out_path: Path, precision: str, family: str) -> None: + if len(_tensor_shape(tensor)) > 4: + if torch is not None and isinstance(tensor, torch.Tensor): + tensor = tensor.reshape(tensor.shape[0], -1) + else: + arr = np.asarray(tensor) + tensor = arr.reshape(arr.shape[0], -1) + if precision == "INT8": + shape = _tensor_shape(tensor) + if "conv_pointwise" in out_path.name and len(shape) == 3 and shape[2] == 1: + save_pointwise_conv1d_int8_with_header(tensor, out_path) + return + if "conv_depthwise.weights" in out_path.name and len(shape) == 3 and shape[1] == 1: + save_depthwise_conv_int8_with_header(tensor, out_path) + return + save_tensor_with_header(tensor, out_path, precision="INT8", model_type=family, allow_int8_bias=True) + return + save_tensor_with_header(tensor, out_path, precision="FP16", model_type=family) + + +def _scale_cq_norms(cq, factor: float): + if factor == 1.0: + return cq + return replace(cq, norms=(cq.norms.astype(np.float32) * float(factor)).astype(np.float16)) + + +def _validate_cq_layout(policy, shape: tuple[int, ...], source_name: str, output_name: str) -> None: + if getattr(policy, "layout", "row_major") != "interleaved_4row": + return + if policy.rotation != "orthogonal" or int(policy.bits or 0) != 4: + raise RuntimeError(f"{source_name}: INTERLEAVED_4ROW output {output_name} requires orthogonal CQ4") + if len(shape) != 2: + raise RuntimeError(f"{source_name}: INTERLEAVED_4ROW output {output_name} requires rank-2 tensor, got shape={shape}") + n, k = int(shape[0]), int(shape[1]) + if n % 4 != 0 or k % 32 != 0: + raise RuntimeError( + f"{source_name}: INTERLEAVED_4ROW output {output_name} requires N % 4 == 0 and K % 32 == 0, got shape={shape}" + ) + + +def _adapt_tensor_for_cactus(tensor, output_name: str | None, family: str): + match = type("_CompatMatch", (), {"output_name": output_name})() + transformed, _name = adapter_for_family(family).transform_tensor(match, tensor) + return transformed + + +def _input_scale_from_diag(diag: np.ndarray | None, tensor) -> np.ndarray | None: + if diag is None: + return None + w = tensor.detach().float().cpu().numpy() if torch is not None and isinstance(tensor, torch.Tensor) else np.asarray(tensor, dtype=np.float32) + d = np.asarray(diag, dtype=np.float32).reshape(-1) + if d.shape[0] != w.shape[1]: + return None + x_abs = np.sqrt(np.clip(d, 1e-12, None)).clip(min=1e-6) + w_abs = np.mean(np.abs(w), axis=0).clip(min=1e-6) + raw = np.power(x_abs, ALPHA) / np.power(w_abs, 1.0 - ALPHA) + raw = raw / np.exp(np.mean(np.log(np.clip(raw, 1e-6, None)))) + return np.clip(raw, 1.0 / 8.0, 8.0).astype(np.float32) + + +def _row_table_scale(tensor, token_ids: list[int]) -> np.ndarray | None: + if not token_ids: + return None + table = tensor.detach().float().cpu().numpy() if torch is not None and isinstance(tensor, torch.Tensor) else np.asarray(tensor, dtype=np.float32) + if table.ndim != 2: + return None + ids = np.asarray(token_ids, dtype=np.int64) + ids = np.clip(ids, 0, table.shape[0] - 1) + x_abs = np.mean(np.abs(table[ids]), axis=0).clip(min=1e-6) + w_abs = np.mean(np.abs(table), axis=0).clip(min=1e-6) + raw = np.power(x_abs, ALPHA) / np.power(w_abs, 1.0 - ALPHA) + raw = raw / np.exp(np.mean(np.log(np.clip(raw, 1e-6, None)))) + return np.clip(raw, 1.0 / 8.0, 8.0).astype(np.float32) + + +def _collect_manifest_token_ids( + processor, + manifest: dict[str, Any], + limits: dict[str, int | None], + max_tokens: int = 32768, +) -> list[int]: + section = manifest.get("language") or manifest.get("transcription") + if not section: + return [] + section_name = "language" if manifest.get("language") else "transcription" + path = section["path"] if isinstance(section, dict) else section + out: list[int] = [] + for row in read_jsonl(path, limits.get(section_name)): + if processor is not None and section_name == "language": + text = language_text(row) + if text: + try: + encoded = processor(text, return_tensors=None) + ids = encoded.get("input_ids", encoded) if isinstance(encoded, Mapping) else encoded + if ids and isinstance(ids[0], list): + ids = ids[0] + out.extend(int(x) for x in ids) + if len(out) >= max_tokens: + return out[:max_tokens] + continue + except Exception: + pass + for key in ("full_token_ids", "input_ids", "prompt_token_ids", "completion_token_ids"): + ids = row.get(key) + if isinstance(ids, list): + out.extend(int(x) for x in ids if isinstance(x, int)) + if len(out) >= max_tokens: + return out[:max_tokens] + if len(out) >= max_tokens: + return out[:max_tokens] + return out[:max_tokens] + + +def _save_hessian_artifacts(out_dir: Path, stats) -> dict[str, int]: + out_dir.mkdir(parents=True, exist_ok=True) + metadata = {"sample_counts": dict(stats.samples)} + for attr in ("unresolved_targets", "errors", "error_samples", "rows", "timings"): + if hasattr(stats, attr): + metadata[attr] = getattr(stats, attr) + if not stats.hessians: + (out_dir / "hessian_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + return dict(stats.samples) + try: + from safetensors.torch import save_file + tensors = {} + if torch is not None: + for name, h in stats.hessians.items(): + tensors[name.replace("/", "__")] = h.detach().cpu() + for name, d in stats.diag.items(): + tensors[name.replace("/", "__") + ".diag"] = d.detach().cpu() + save_file(tensors, str(out_dir / "hessians.safetensors")) + except Exception as exc: + metadata["save_warning"] = str(exc) + (out_dir / "hessian_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + return dict(stats.samples) + + +def _load_hessian_artifacts(cache_dir: Path): + from .calibration.hessian import HessianStats + + metadata_path = cache_dir / "hessian_metadata.json" + tensor_path = cache_dir / "hessians.safetensors" + if not metadata_path.exists(): + raise FileNotFoundError(f"{cache_dir} does not contain hessian_metadata.json") + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + stats = HessianStats(samples={str(k): int(v) for k, v in metadata.get("sample_counts", {}).items()}) + stats.unresolved_targets = list(metadata.get("unresolved_targets", [])) + stats.errors = {str(k): int(v) for k, v in metadata.get("errors", {}).items()} + stats.error_samples = list(metadata.get("error_samples", [])) + stats.rows = {str(k): int(v) for k, v in metadata.get("rows", {}).items()} + stats.timings = {str(k): float(v) for k, v in metadata.get("timings", {}).items()} + if tensor_path.exists(): + if torch is None: + raise RuntimeError("torch is required to load Hessian cache") + from safetensors.torch import load_file + + tensors = load_file(str(tensor_path), device="cpu") + for key, tensor in tensors.items(): + name = key.replace("__", "/") + if name.endswith(".diag"): + stats.diag[name[:-5]] = tensor + else: + stats.hessians[name] = tensor + return stats + + +def _config_dict(cfg) -> dict[str, Any]: + if hasattr(cfg, "to_dict"): + return cfg.to_dict() + if isinstance(cfg, dict): + return cfg + return dict(getattr(cfg, "__dict__", {})) + + +def _resolve_device(device: str) -> str: + if device != "auto": + return device + if torch is not None: + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + if torch.cuda.is_available(): + return "cuda" + return "cpu" + + +def convert(args: argparse.Namespace) -> None: + if args.bits not in {1, 2, 3, 4}: + raise SystemExit("--bits must be one of 1,2,3,4") + args.device = _resolve_device(args.device) + out_dir = Path(args.out) + if args.force and args.hessian_cache_in: + try: + if Path(args.hessian_cache_in).resolve() == out_dir.resolve(): + raise SystemExit("--hessian-cache-in cannot point at --out when --force is set; --force deletes --out before cache load") + except FileNotFoundError: + pass + if out_dir.exists(): + if not args.force: + raise SystemExit(f"{out_dir} exists; pass --force to replace it") + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + cfg, processor, model = _load_hf(args.model, args.device, skip_model_load=bool(args.skip_model_load)) + runtime_source = ensure_parakeet_tdt_nemo_source(args.model, cache_dir=_hf_cache_dir()) or args.model + family = detect_family(cfg, args.model_family) + adapter = adapter_for_family(family) + checkpoint_state = _load_checkpoint_state_dict(args.model) + if checkpoint_state is None: + if model is None: + raise RuntimeError(f"could not load model object or checkpoint tensors from {args.model}") + checkpoint_state = model.state_dict() + normalized_state = adapter.normalize_state_dict(checkpoint_state) + state_dict = normalized_state.state_dict + model_config = adapter.runtime_config(cfg) + model_config["model_type"] = adapter.runtime_model_type() + write_config_txt({**_config_dict(cfg), **model_config}, out_dir) + copy_runtime_files(runtime_source, out_dir, token=getattr(args, "token", None), cache_dir=_hf_cache_dir()) + try: + from .cactus_adapters.tokenizer import convert_hf_tokenizer + + tokenizer = getattr(processor, "tokenizer", processor) + if tokenizer is not None: + convert_hf_tokenizer(tokenizer, out_dir, model_id=args.model, model_type=family) + except Exception as exc: + if args.strict: + raise RuntimeError(f"failed to write Cactus tokenizer files: {exc}") from exc + + manifest = load_manifest(args.calibration_manifest) + limits = { + "language": args.max_language_examples, + "vision": args.max_vision_examples, + "audio": args.max_audio_examples, + "transcription": args.max_transcription_examples, + "embedding": args.max_embedding_examples, + } + rows: list[dict[str, Any]] = [] + pending: list[tuple[str, Any, Any, Any]] = [] + num_layers = max( + int(model_config.get("num_layers", 0) or 0), + int(cfg_get(cfg, "num_encoder_layers", 0) or 0), + int(cfg_get(cfg, "num_decoder_layers", 0) or 0), + ) or None + target_modules: set[str] = set() + for name, tensor in state_dict.items(): + match = adapter.name_tensor(name, tensor, num_layers) + requested_bits = _bits_for_component(match.component, args) + policy = adapter.policy(match, _tensor_shape(tensor), requested_bits) + if policy.action == "convert" and policy.use_gptq: + target = adapter.module_target_name(name, model) + if target: + target_modules.add(target) + pending.append((name, tensor, match, policy)) + + if target_modules and model is None and not args.hessian_cache_in: + msg = "model object failed to load; GPTQ Hessian calibration is disabled and use_gptq tensors will fall back to RTN" + if args.strict: + raise RuntimeError(msg) + print(f"warning: {msg}") + + if args.hessian_cache_in: + hessian_stats = _load_hessian_artifacts(Path(args.hessian_cache_in)) + else: + progress_path = out_dir / "hessian_progress.json" if args.hessian_progress else None + hessian_stats = ( + collect_manifest_hessians( + model, + processor, + manifest, + limits, + target_modules, + args.device, + adapter=adapter, + batch_size=args.calibration_batch_size, + gpu_flush_interval=args.hessian_gpu_flush_interval, + progress_path=progress_path, + preprocessed_cache_dir=args.preprocessed_cache_dir, + ) + if processor is not None and model is not None and target_modules + else None + ) + hessian_samples = _save_hessian_artifacts(out_dir, hessian_stats) if hessian_stats is not None else {} + if args.hessian_cache_out and hessian_stats is not None: + _save_hessian_artifacts(Path(args.hessian_cache_out), hessian_stats) + hessians_np = {} + diag_np = {} + if hessian_stats is not None and torch is not None: + for name, h in hessian_stats.hessians.items(): + hessians_np[name] = h.detach().cpu().numpy().astype(np.float32) + for name, d in hessian_stats.diag.items(): + diag_np[name] = d.detach().cpu().numpy().astype(np.float32) + + scale_token_ids = _collect_manifest_token_ids(processor, manifest, limits) + + for name, tensor, match, policy in pending: + emissions = adapter.expand_tensor(match, tensor) + if not emissions: + emissions = [type("_IgnoredEmission", (), {"output_name": None, "tensor": tensor, "transform": "none", "source_names": None})()] + provenance = normalized_state.provenance.get(name) + source_names = provenance.source_names if provenance else [name] + for emission in emissions: + qdq_restore = getattr(emission, "qdq_restore", None) or (provenance.qdq_restore if provenance else "hf_key") + out_path = out_dir / emission.output_name if emission.output_name else None + emit_tensor = emission.tensor + emit_match = replace(match, output_name=emission.output_name) + requested_bits = _bits_for_component(emit_match.component, args) + emit_policy = adapter.policy(emit_match, _tensor_shape(emit_tensor), requested_bits) + manifest_transform = provenance.transform if provenance and provenance.transform != "none" else emission.transform + emit_source_names = emission.source_names or source_names + status = emit_policy.action + precision = emit_policy.precision + gptq_used = False + hessian_missing_reason = None + module_name = adapter.module_target_name(name, model) or _module_name(name) + try: + if not match.recognized and args.strict: + raise RuntimeError(f"unrecognized tensor in strict mode: {name}") + if emit_policy.action == "convert" and out_path is not None: + _validate_cq_layout(emit_policy, _tensor_shape(emit_tensor), name, out_path.name) + if emit_policy.use_gptq and int(hessian_samples.get(module_name, 0)) <= 0: + hessian_missing_reason = "expected GPTQ target had zero samples" + if args.strict: + raise RuntimeError(f"{name}: {hessian_missing_reason} ({module_name})") + hessian = hessians_np.get(module_name) + input_scale = None + if emit_policy.rotation == "orthogonal" or name.endswith("embed_tokens_per_layer.weight"): + input_scale = _row_table_scale(emit_tensor, scale_token_ids) + else: + input_scale = _input_scale_from_diag(diag_np.get(module_name), emit_tensor) + if emit_policy.rotation == "orthogonal": + cq = quantize_orthogonal(emit_tensor, bits=int(emit_policy.bits or 4), input_scale=input_scale) + else: + cq = quantize_hadamard( + emit_tensor, + bits=int(emit_policy.bits or requested_bits), + hessian=hessian, + use_gptq=emit_policy.use_gptq, + input_scale=input_scale, + ) + cq = _scale_cq_norms(cq, adapter.scale_factor(out_path.name)) + if getattr(emit_policy, "layout", "row_major") == "interleaved_4row": + cq = replace(cq, interleaved_4row=True) + write_cq_tensor(out_path, cq) + gptq_used = cq.gptq_used + status = "converted" if match.recognized else "unrecognized" + elif emit_policy.action == "fallback" and out_path is not None: + _save_fallback_tensor(emit_tensor, out_path, precision, family) + status = "fallback" if match.recognized else "unrecognized" + elif emit_policy.action == "ignored": + status = "ignored" + except Exception as exc: + if args.strict or getattr(emit_policy, "layout", "row_major") == "interleaved_4row": + raise + if out_path is not None: + _save_fallback_tensor(emit_tensor, out_path, "FP16", family) + status = "fallback" if match.recognized else "unrecognized" + precision = "FP16" + emit_policy = replace(emit_policy, action="fallback", precision="FP16", fallback_reason=str(exc)) + rows.append({ + "source_name": name, + "hf_name": match.hf_name or name, + "adapter_name": match.adapter_name or name, + "output_file": str(out_path.name) if out_path else None, + "shape": list(_tensor_shape(emit_tensor)), + "dtype": str(getattr(emit_tensor, "dtype", "")), + "component": emit_policy.component, + "policy": emit_policy.action, + "precision": precision, + "bits": emit_policy.bits, + "status": status, + "required": bool(emit_policy.action != "ignored"), + "fallback_reason": emit_policy.fallback_reason, + "hessian_samples": int(hessian_samples.get(module_name, 0)), + "gptq_used": bool(gptq_used), + "bytes": _tensor_bytes(out_path), + "scale_factor": float(adapter.scale_factor(out_path.name)) if out_path else 1.0, + "recognized": bool(match.recognized), + "adapter_family": family, + "source_names": emit_source_names, + "transform": manifest_transform, + "qdq_restore": qdq_restore, + "hessian_missing_reason": hessian_missing_reason, + }) + + summary = write_reports(out_dir, rows) + print_summary(summary) + unrecognized = [r["source_name"] for r in rows if r["status"] == "unrecognized"] + if unrecognized and not args.strict: + print(f"\nWarning: {len(unrecognized)} unrecognized tensors were exported with generic names or FP16 fallback.") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="python -m cactus.convert") + sub = parser.add_subparsers(dest="command", required=True) + p = sub.add_parser("convert") + p.add_argument("--model", required=True) + p.add_argument("--out", required=True) + p.add_argument("--bits", type=int, choices=[1, 2, 3, 4], required=True) + p.add_argument("--language-bits", type=int, choices=[1, 2, 3, 4]) + p.add_argument("--vision-bits", type=int, choices=[1, 2, 3, 4]) + p.add_argument("--audio-bits", type=int, choices=[1, 2, 3, 4]) + p.add_argument("--embedding-bits", type=int, choices=[1, 2, 3, 4]) + p.add_argument("--calibration-manifest") + p.add_argument("--device", default="auto") + p.add_argument("--model-family", default="auto", choices=sorted(SUPPORTED_FAMILIES)) + p.add_argument("--strict", action="store_true") + p.add_argument("--force", action="store_true") + p.add_argument("--skip-model-load", action="store_true") + p.add_argument("--max-language-examples", type=int) + p.add_argument("--max-vision-examples", type=int) + p.add_argument("--max-audio-examples", type=int) + p.add_argument("--max-transcription-examples", type=int) + p.add_argument("--max-embedding-examples", type=int) + p.add_argument("--hessian-cache-in") + p.add_argument("--hessian-cache-out") + p.add_argument("--calibration-batch-size", type=int, default=1) + p.add_argument("--hessian-gpu-flush-interval", type=int, default=16) + p.add_argument("--preprocessed-cache-dir") + p.add_argument("--hessian-progress", action="store_true") + p.set_defaults(func=convert) + q = sub.add_parser("qdq") + q.add_argument("input", type=Path, help="Cactus `.weights` directory or tar archive") + q.add_argument("--out", type=Path, required=True, help="Output HF-style QDQ checkpoint directory") + q.add_argument("--dtype", choices=["float16", "bfloat16"], default="float16") + q.add_argument("--model-family", default="auto", choices=sorted(SUPPORTED_FAMILIES)) + q.add_argument("--shard-size-gb", type=float, default=4.0) + q.add_argument("--row-batch-size", type=int, default=2048) + q.add_argument("--tmp-dir", type=Path) + q.add_argument("--force", action="store_true") + q.set_defaults(func=lambda args: print(f"done: wrote {convert_qdq(args)['written_count']} tensors")) + v = sub.add_parser("validate") + v.add_argument("--source-model", required=True) + v.add_argument("--qdq", required=True) + v.add_argument("--out") + v.add_argument("--strict", action="store_true") + v.add_argument("--model-family", default="auto", choices=sorted(SUPPORTED_FAMILIES)) + v.set_defaults(func=validate_qdq) + return parser + + +def main(argv: list[str] | None = None) -> None: + parser = build_parser() + args = parser.parse_args(argv) + args.func(args) diff --git a/python/cactus/convert/compat.py b/python/cactus/convert/compat.py new file mode 100644 index 000000000..99d4fb5ec --- /dev/null +++ b/python/cactus/convert/compat.py @@ -0,0 +1,34 @@ +from __future__ import annotations + + +def patch_transformers_import_compat() -> list[str]: + """Keep Transformers processor/model imports working in local dev envs. + + Some pyenv Python builds are missing the stdlib `_lzma` extension. Importing + `AutoProcessor` can still reach torchvision datasets, which import `lzma` + even when conversion only needs processor metadata. The transpiler already + has the battle-tested patch, so reuse it before converter-side HF imports. + """ + + notes: list[str] = [] + try: + from cactus.transpile.runtime_support import ( + ensure_transformers_supports_gemma4, + patch_torch_flex_attention_compat, + patch_transformers_torchvision_probe, + ) + except Exception: + return notes + + for patch in ( + patch_transformers_torchvision_probe, + patch_torch_flex_attention_compat, + ensure_transformers_supports_gemma4, + ): + try: + note = patch() + except Exception: + note = None + if note: + notes.append(str(note)) + return notes diff --git a/python/cactus/convert/export/__init__.py b/python/cactus/convert/export/__init__.py new file mode 100644 index 000000000..1602684f5 --- /dev/null +++ b/python/cactus/convert/export/__init__.py @@ -0,0 +1,2 @@ +"""Cactus export helpers.""" + diff --git a/python/cactus/convert/export/files.py b/python/cactus/convert/export/files.py new file mode 100644 index 000000000..4529a4db8 --- /dev/null +++ b/python/cactus/convert/export/files.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import json +import shutil +from pathlib import Path + + +TOKENIZER_FILES = [ + "config.json", + "vocab.json", + "merges.txt", + "tokenizer.json", + "tokenizer.model", + "vocab.txt", + "tokenizer_config.json", + "special_tokens_map.json", + "added_tokens.json", + "chat_template.json", + "chat_template.jinja", + "generation_config.json", + "preprocessor_config.json", + "image_processor_config.json", + "video_preprocessor_config.json", + "processor_config.json", +] + + +def copy_runtime_files( + model_path: str | Path, + out_dir: Path, + *, + token: str | None = None, + cache_dir: str | None = None, +) -> None: + src_dir = Path(model_path) + if not src_dir.exists() or not src_dir.is_dir(): + try: + from huggingface_hub import snapshot_download + + src_dir = Path( + snapshot_download( + repo_id=str(model_path), + allow_patterns=TOKENIZER_FILES, + token=token, + cache_dir=cache_dir, + ) + ) + except Exception: + return + for name in TOKENIZER_FILES: + src = src_dir / name + if src.exists(): + shutil.copy2(src, out_dir / name) + + +def write_config_txt(config: dict, out_dir: Path) -> None: + def fmt(v): + if isinstance(v, bool): + return "true" if v else "false" + if isinstance(v, (list, tuple)): + return ",".join(str(x) for x in v) + return str(v) + with (out_dir / "config.txt").open("w", encoding="utf-8") as f: + for key in sorted(config): + val = config[key] + if isinstance(val, (str, int, float, bool, list, tuple)): + f.write(f"{key}={fmt(val)}\n") + (out_dir / "hf_config.json").write_text(json.dumps(config, indent=2, sort_keys=True), encoding="utf-8") diff --git a/python/cactus/convert/export/qdq.py b/python/cactus/convert/export/qdq.py new file mode 100644 index 000000000..5c73081d6 --- /dev/null +++ b/python/cactus/convert/export/qdq.py @@ -0,0 +1,564 @@ +from __future__ import annotations + +import json +import math +import os +import shutil +import struct +import tarfile +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from safetensors.torch import save_file +from scipy.linalg import hadamard + +from ..model_adapters.detection import detect_family +from ..model_adapters.adapters import adapter_for_family + +CACTUS_MAGIC = b"CACT" +HEADER_SIZE = 84 +ALIGNMENT_DEFAULT = 32 +FLAG_ORTHOGONAL_ROTATION = 1 << 1 +FLAG_INTERLEAVED_4ROW = 1 << 2 +FLAG_INTERLEAVED = 1 << 3 + +PRECISION_INT8 = 0 +PRECISION_FP16 = 1 +PRECISION_FP32 = 2 +PRECISION_CQ = {3: 1, 4: 2, 5: 3, 6: 4} + +CONFIG_FILES = { + "config.json", + "generation_config.json", + "vocab.json", + "merges.txt", + "tokenizer.json", + "tokenizer_config.json", + "tokenizer_config.txt", + "special_tokens_map.json", + "processor_config.json", + "preprocessor_config.json", + "image_processor_config.json", + "video_preprocessor_config.json", + "chat_template.json", + "chat_template.jinja", + "tokenizer.model", + "config.txt", +} + + +@dataclass(frozen=True) +class CactusHeader: + path: Path + flags: int + alignment: int + ndim: int + dims: tuple[int, int, int, int] + precision: int + data_bytes: int + scales_bytes: int + group_size: int + num_groups: int + original_n: int + + @property + def shape(self) -> tuple[int, ...]: + return tuple(int(x) for x in self.dims[: self.ndim]) + + @property + def bits(self) -> int: + if self.precision not in PRECISION_CQ: + raise ValueError(f"{self.path}: precision {self.precision} is not CQ") + return PRECISION_CQ[self.precision] + + +class ShardedSafetensorsWriter: + def __init__(self, out_dir: Path, shard_size_bytes: int) -> None: + self.out_dir = out_dir + self.shard_size_bytes = int(shard_size_bytes) + self.current: dict[str, torch.Tensor] = {} + self.current_bytes = 0 + self.shards: list[tuple[Path, list[str]]] = [] + self.weight_map: dict[str, str] = {} + self.total_size = 0 + + def add(self, key: str, tensor: torch.Tensor) -> None: + tensor = tensor.detach().cpu().contiguous() + nbytes = tensor.numel() * tensor.element_size() + if self.current and self.current_bytes + nbytes > self.shard_size_bytes: + self.flush() + self.current[key] = tensor + self.current_bytes += nbytes + self.total_size += nbytes + if self.current_bytes >= self.shard_size_bytes: + self.flush() + + def flush(self) -> None: + if not self.current: + return + tmp_path = self.out_dir / f"model-{len(self.shards) + 1:05d}.safetensors" + save_file(self.current, tmp_path) + self.shards.append((tmp_path, sorted(self.current))) + self.current = {} + self.current_bytes = 0 + + def close(self) -> None: + self.flush() + if len(self.shards) == 1: + old, keys = self.shards[0] + final = self.out_dir / "model.safetensors" + old.rename(final) + for key in keys: + self.weight_map[key] = final.name + return + total = len(self.shards) + for idx, (old, keys) in enumerate(self.shards, start=1): + final = self.out_dir / f"model-{idx:05d}-of-{total:05d}.safetensors" + old.rename(final) + for key in keys: + self.weight_map[key] = final.name + index = { + "metadata": {"total_size": str(self.total_size)}, + "weight_map": dict(sorted(self.weight_map.items())), + } + (self.out_dir / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n", encoding="utf-8") + + +def align_offset(offset: int, alignment: int) -> int: + rem = offset % alignment + return offset if rem == 0 else offset + alignment - rem + + +def safe_extract_tar(tar_path: Path, out_dir: Path) -> None: + out_resolved = out_dir.resolve() + with tarfile.open(tar_path) as tf: + for member in tf.getmembers(): + target = (out_dir / member.name).resolve() + if os.path.commonpath([str(out_resolved), str(target)]) != str(out_resolved): + raise RuntimeError(f"refusing unsafe tar member path: {member.name}") + tf.extractall(out_dir) + + +def materialize_input(input_path: Path, tmp_dir: Path | None) -> tuple[Path, tempfile.TemporaryDirectory[str] | None]: + if input_path.is_dir(): + return input_path, None + if not tarfile.is_tarfile(input_path): + raise ValueError(f"{input_path} is neither a directory nor a tar archive") + tmp = tempfile.TemporaryDirectory(dir=str(tmp_dir) if tmp_dir else None) + root = Path(tmp.name) + safe_extract_tar(input_path, root) + children = [p for p in root.iterdir() if not p.name.startswith(".")] + return (children[0] if len(children) == 1 and children[0].is_dir() else root), tmp + + +def find_cactus_root(root: Path) -> Path: + if (root / "conversion_manifest.json").exists(): + return root + weights = [p for p in root.rglob("*.weights") if p.is_file()] + if not weights: + raise ValueError(f"no .weights files found under {root}") + counts: dict[Path, int] = {} + for path in weights: + counts[path.parent] = counts.get(path.parent, 0) + 1 + return max(counts, key=counts.get) + + +def copy_config_files(src_root: Path, out_dir: Path) -> list[str]: + copied: list[str] = [] + for path in src_root.rglob("*"): + if path.is_file() and path.name in CONFIG_FILES: + rel = path.relative_to(src_root) + target = out_dir / rel + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, target) + copied.append(str(rel)) + return sorted(set(copied)) + + +def read_header(path: Path) -> CactusHeader: + raw = path.read_bytes()[:HEADER_SIZE] + if len(raw) != HEADER_SIZE: + raise ValueError(f"{path}: too small for cactus header") + if raw[:4] != CACTUS_MAGIC: + raise ValueError(f"{path}: bad magic {raw[:4]!r}") + fields = struct.unpack(" 4: + raise ValueError(f"{path}: invalid ndim={ndim}") + if alignment <= 0: + alignment = ALIGNMENT_DEFAULT + scales_offset = align_offset(HEADER_SIZE, alignment) + data_offset = align_offset(scales_offset + scales_bytes, alignment) + expected = data_offset + data_bytes + actual = path.stat().st_size + if actual != expected: + raise ValueError(f"{path}: size mismatch actual={actual} expected={expected}") + return CactusHeader(path, flags, alignment, ndim, (d0, d1, d2, d3), precision, data_bytes, scales_bytes, group_size, num_groups, original_n) + + +def unpack_lsb_values(packed: np.ndarray, count: int, bits: int) -> np.ndarray: + raw_bits = np.unpackbits(packed.astype(np.uint8, copy=False), bitorder="little")[: count * bits] + raw_bits = raw_bits.reshape(count, bits) + out = np.zeros(count, dtype=np.uint8) + for bit in range(bits): + out |= raw_bits[:, bit].astype(np.uint8) << bit + return out + + +def unpack_interleaved_4row_4bit(packed: np.ndarray, rows: int, cols: int) -> np.ndarray: + if rows % 4 != 0 or cols % 8 != 0: + raise ValueError(f"INTERLEAVED_4ROW CQ4 requires rows % 4 == 0 and cols % 8 == 0, got {(rows, cols)}") + chunks = cols // 8 + view = packed.reshape(rows // 4, chunks, 4, 4) + lo = (view & 0x0F).astype(np.uint8) + hi = ((view >> 4) & 0x0F).astype(np.uint8) + out = np.concatenate([lo, hi], axis=3).transpose(0, 2, 1, 3).copy() + return out.reshape(rows, cols) + + +def dequantize_fp_file(path: Path, header: CactusHeader, out_dtype: torch.dtype) -> torch.Tensor: + offset = align_offset(HEADER_SIZE, header.alignment) + dtype = np.float16 if header.precision == PRECISION_FP16 else np.float32 + arr = np.fromfile(path, dtype=dtype, count=math.prod(header.shape), offset=offset) + return torch.from_numpy(arr.reshape(header.shape).copy()).to(out_dtype) + + +def dequantize_int8_file(path: Path, header: CactusHeader, out_dtype: torch.dtype) -> torch.Tensor: + scales_offset = align_offset(HEADER_SIZE, header.alignment) + data_offset = align_offset(scales_offset + header.scales_bytes, header.alignment) + with path.open("rb") as f: + f.seek(scales_offset) + scales_blob = f.read(header.scales_bytes) + f.seek(data_offset) + q = np.frombuffer(f.read(header.data_bytes), dtype=np.int8).copy() + count = math.prod(header.shape) + if q.size < count: + raise ValueError(f"{path}: INT8 payload too small for shape {header.shape}") + q = q[:count].reshape(header.shape).astype(np.float32) + if header.scales_bytes > 0 and header.group_size > 0: + scales = np.frombuffer(scales_blob, dtype=np.float16).astype(np.float32) + if header.flags & FLAG_INTERLEAVED: + raise ValueError(f"{path}: interleaved grouped INT8 QDQ is not implemented") + if len(header.shape) == 1: + group_ids = np.arange(count) // header.group_size + arr = (q.reshape(-1) * scales[group_ids]).reshape(header.shape) + if header.original_n and header.original_n < arr.shape[0]: + arr = arr[: header.original_n] + elif len(header.shape) == 2: + n, k = header.shape + num_groups = k // header.group_size + if scales.size < n * num_groups: + raise ValueError(f"{path}: grouped INT8 scale metadata too small") + scales = scales[: n * num_groups].reshape(n, num_groups) + arr = np.empty((n, k), dtype=np.float32) + for col in range(k): + arr[:, col] = q[:, col] * scales[:, col // header.group_size] + elif len(header.shape) >= 3: + n = header.shape[0] + k_total = math.prod(header.shape[1:]) + if k_total % header.group_size != 0: + raise ValueError(f"{path}: grouped INT8 inner size must be divisible by group size") + num_groups = k_total // header.group_size + if scales.size < n * num_groups: + raise ValueError(f"{path}: grouped INT8 scale metadata too small") + q2 = q.reshape(n, k_total) + scales = scales[: n * num_groups].reshape(n, num_groups) + arr2 = np.empty((n, k_total), dtype=np.float32) + for col in range(k_total): + arr2[:, col] = q2[:, col] * scales[:, col // header.group_size] + arr = arr2.reshape(header.shape) + else: + raise ValueError(f"{path}: grouped INT8 QDQ supports rank >= 1 tensors only") + else: + arr = q + return torch.from_numpy(arr.copy()).to(out_dtype) + + +def read_cq_payload(path: Path, header: CactusHeader) -> tuple[bytes, np.ndarray]: + scales_offset = align_offset(HEADER_SIZE, header.alignment) + data_offset = align_offset(scales_offset + header.scales_bytes, header.alignment) + with path.open("rb") as f: + f.seek(scales_offset) + scales_blob = f.read(header.scales_bytes) + f.seek(data_offset) + packed = np.frombuffer(f.read(header.data_bytes), dtype=np.uint8).copy() + return scales_blob, packed + + +def parse_normal_metadata(blob: bytes, header: CactusHeader): + n, k = header.shape + pos = 0 + codebook = np.frombuffer(blob, dtype=np.float16, count=1 << header.bits, offset=pos).astype(np.float32) + pos += (1 << header.bits) * 2 + input_scale = np.frombuffer(blob, dtype=np.float16, count=k, offset=pos).astype(np.float32) + pos += k * 4 + norms = np.frombuffer(blob, dtype=np.float16, count=n * header.num_groups, offset=pos).astype(np.float32).reshape(n, header.num_groups) + pos += n * header.num_groups * 2 + left = np.frombuffer(blob, dtype=np.int8, count=header.group_size, offset=pos).astype(np.float32) + pos += header.group_size + right = np.frombuffer(blob, dtype=np.int8, count=header.group_size, offset=pos).astype(np.float32) + pos += header.group_size + perm = np.frombuffer(blob, dtype=" torch.Tensor: + if header.ndim != 2: + raise ValueError(f"{path}: CQ tensors must be 2D, got shape={header.shape}") + n, k = header.shape + bits = header.bits + scales_blob, packed = read_cq_payload(path, header) + if header.flags & FLAG_ORTHOGONAL_ROTATION: + packed_row_bytes = math.ceil(k * bits / 8) + codebook, input_scale, norms, rotation = parse_orthogonal_metadata(scales_blob, header) + out = torch.empty(n, k, dtype=out_dtype) + interleaved = bool(header.flags & FLAG_INTERLEAVED_4ROW) + if interleaved: + if bits != 4 or header.num_groups != 1 or header.group_size != k: + raise ValueError(f"{path}: orthogonal INTERLEAVED_4ROW requires full-width CQ4") + packed_indices = unpack_interleaved_4row_4bit(packed, n, k) + norms = norms.reshape(n // 4, 4).reshape(n) + else: + packed_rows = packed.reshape(n, packed_row_bytes) + rt = rotation.float().T.contiguous() + scale = input_scale.float().unsqueeze(0) + codebook = codebook.float() + for start in range(0, n, row_batch_size): + end = min(start + row_batch_size, n) + if interleaved: + idx_np = packed_indices[start:end] + else: + idx_np = np.stack([unpack_lsb_values(row, k, bits) for row in packed_rows[start:end]]) + idx = torch.from_numpy(idx_np.astype(np.int64, copy=False)) + recon = (codebook[idx] @ rt) * norms[start:end].float().unsqueeze(1) + out[start:end] = (recon / scale).to(out_dtype) + return out + packed_group_bytes = math.ceil(header.group_size * bits / 8) + codebook, input_scale, norms, left, right, perm = parse_normal_metadata(scales_blob, header) + signs = set(int(x) for x in left.tolist()) | set(int(x) for x in right.tolist()) + if not signs.issubset({-1, 1}): + raise ValueError(f"{path}: signs must be +/-1, got {sorted(signs)}") + if sorted(int(x) for x in perm.tolist()) != list(range(header.group_size)): + raise ValueError(f"{path}: permutation is not bijective") + base_h = torch.from_numpy((hadamard(header.group_size, dtype=float) / math.sqrt(header.group_size)).astype(np.float32)) + rotation = (left.float().unsqueeze(1) * base_h * right.float().unsqueeze(0))[:, perm.long()].contiguous() + rt = rotation.T.contiguous() + scale = input_scale.float().unsqueeze(0) + codebook = codebook.float() + + interleaved = bool(header.flags & FLAG_INTERLEAVED_4ROW) + if interleaved: + if bits not in (1, 2, 3, 4): + raise ValueError(f"{path}: INTERLEAVED_4ROW only valid for 1..4-bit CQ, got {bits}") + if n % 4 != 0: + raise ValueError(f"{path}: INTERLEAVED_4ROW requires N % 4 == 0") + if header.group_size % 32 != 0: + raise ValueError(f"{path}: INTERLEAVED_4ROW requires group_size % 32 == 0") + norms = norms.reshape(n // 4, header.num_groups, 4).permute(0, 2, 1).reshape(n, header.num_groups).contiguous() + nb = n // 4 + view_packed = np.asarray(packed) + + if bits == 4: + chunks = header.group_size // 8 + view = view_packed.reshape(nb, header.num_groups, chunks, 4, 4) + lo = (view & 0x0F).astype(np.uint8) + hi = ((view >> 4) & 0x0F).astype(np.uint8) + per_chunk = np.concatenate([lo, hi], axis=4) + elif bits == 2: + chunks = header.group_size // 16 + view = view_packed.reshape(nb, header.num_groups, chunks, 4, 4) + view = view.transpose(0, 1, 2, 4, 3) + extracts = [] + for b in range(4): + extracts.append(((view >> (b * 2)) & 0x3).astype(np.uint8)) + stacked = np.stack(extracts, axis=-1) + per_chunk = stacked.reshape(nb, header.num_groups, chunks, 4, 16) + elif bits == 3: + chunks_per_group = header.group_size // 4 + panel_bytes = header.group_size * 3 // 2 + assert view_packed.size == nb * header.num_groups * panel_bytes, \ + f"expected {nb*header.num_groups*panel_bytes} bytes, got {view_packed.size}" + chunk_view = view_packed.reshape(nb, header.num_groups, chunks_per_group, 6) + word = np.zeros((nb, header.num_groups, chunks_per_group), dtype=np.uint64) + for b in range(6): + word |= chunk_view[..., b].astype(np.uint64) << (b * 8) + chunk_idx = np.empty((nb, header.num_groups, chunks_per_group, 16), dtype=np.uint8) + for i in range(16): + chunk_idx[..., i] = ((word >> (i * 3)) & 0x7).astype(np.uint8) + chunk_idx = chunk_idx.reshape(nb, header.num_groups, chunks_per_group, 4, 4) + per_chunk = chunk_idx + else: + chunks = header.group_size // 32 + view = view_packed.reshape(nb, header.num_groups, chunks, 4, 4) + view = view.transpose(0, 1, 2, 4, 3) + extracts = [] + for b in range(8): + extracts.append(((view >> b) & 0x1).astype(np.uint8)) + stacked = np.stack(extracts, axis=-1) + per_chunk = stacked.reshape(nb, header.num_groups, chunks, 4, 32) + + per_chunk = per_chunk.transpose(0, 3, 1, 2, 4) + idx_full = per_chunk.reshape(n, header.num_groups, header.group_size) + out = torch.empty(n, k, dtype=out_dtype) + for start in range(0, n, row_batch_size): + end = min(start + row_batch_size, n) + idx = torch.from_numpy(idx_full[start:end].astype(np.int64, copy=False)) + recon = (codebook[idx] @ rt) * norms[start:end].float().unsqueeze(-1) + out[start:end] = (recon.reshape(end - start, k) / scale).to(out_dtype) + return out + + out = torch.empty(n, k, dtype=out_dtype) + packed_groups = packed.reshape(n, header.num_groups, packed_group_bytes) + for start in range(0, n, row_batch_size): + end = min(start + row_batch_size, n) + idx_np = np.empty((end - start, header.num_groups, header.group_size), dtype=np.uint8) + for row_i, row_groups in enumerate(packed_groups[start:end]): + for group_i, group in enumerate(row_groups): + idx_np[row_i, group_i] = unpack_lsb_values(group, header.group_size, bits) + idx = torch.from_numpy(idx_np.astype(np.int64, copy=False)) + recon = (codebook[idx] @ rt) * norms[start:end].float().unsqueeze(-1) + out[start:end] = (recon.reshape(end - start, k) / scale).to(out_dtype) + return out + + +def load_conversion_manifest(cactus_root: Path) -> list[dict[str, Any]]: + path = cactus_root / "conversion_manifest.json" + if not path.exists(): + raise ValueError(f"{cactus_root} does not contain conversion_manifest.json; manifest-driven QDQ is required") + rows = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(rows, list): + raise ValueError(f"{path}: expected list") + return rows + + +def output_keys_for_row(row: dict[str, Any], family: str) -> list[str]: + return adapter_for_family(family).qdq_output_keys(row) + + +def dequantize_path(path: Path, out_dtype: torch.dtype, row_batch_size: int) -> tuple[torch.Tensor, CactusHeader]: + header = read_header(path) + if header.precision in (PRECISION_FP16, PRECISION_FP32): + tensor = dequantize_fp_file(path, header, out_dtype) + elif header.precision == PRECISION_INT8: + tensor = dequantize_int8_file(path, header, out_dtype) + elif header.precision in PRECISION_CQ: + tensor = dequantize_cq_file(path, header, out_dtype, row_batch_size) + else: + raise ValueError(f"{path}: unsupported precision={header.precision}") + return tensor, header + + +def trim_to_manifest_shape(tensor: torch.Tensor, shape: list[int] | tuple[int, ...] | None) -> torch.Tensor: + if not shape: + return tensor + expected = tuple(int(x) for x in shape) + if tuple(tensor.shape) == expected: + return tensor + if len(expected) == 1 and tensor.ndim == 1 and tensor.shape[0] >= expected[0]: + return tensor[: expected[0]] + if len(expected) == 2 and tensor.ndim == 2 and tensor.shape[0] == expected[0] and tensor.shape[1] >= expected[1]: + return tensor[:, : expected[1]] + if tensor.numel() == math.prod(expected): + return tensor.reshape(expected) + return tensor + + +def infer_family(cactus_root: Path, requested: str) -> str: + if requested != "auto": + return requested + config_path = cactus_root / "config.json" + if not config_path.exists(): + return "generic" + try: + from transformers import AutoConfig + + cfg = AutoConfig.from_pretrained(str(cactus_root), trust_remote_code=True, local_files_only=True) + return detect_family(cfg, "auto") + except Exception: + data = json.loads(config_path.read_text(encoding="utf-8")) + model_type = str(data.get("model_type", "")).lower() + arch = " ".join(data.get("architectures") or []).lower() + if "gemma4" in model_type or "gemma4" in arch: + return "gemma4" + return "generic" + + +def convert_qdq(args) -> dict[str, Any]: + if args.out.exists(): + if not args.force: + raise SystemExit(f"{args.out} exists; pass --force to replace it") + shutil.rmtree(args.out) + args.out.mkdir(parents=True, exist_ok=True) + out_dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + root, tmp = materialize_input(args.input, args.tmp_dir) + try: + cactus_root = find_cactus_root(root) + family = infer_family(cactus_root, args.model_family) + adapter = adapter_for_family(family) + rows = load_conversion_manifest(cactus_root) + writer = ShardedSafetensorsWriter(args.out, int(args.shard_size_gb * (1024**3))) + report = { + "input": str(args.input), + "cactus_root": str(cactus_root), + "family": family, + "dtype": args.dtype, + "copied_config_files": copy_config_files(cactus_root, args.out), + "written": [], + "skipped": [], + } + seen: set[str] = set() + for row in rows: + output_file = row.get("output_file") + if not output_file: + report["skipped"].append({"source_name": row.get("source_name"), "reason": "no output file"}) + continue + path = cactus_root / str(output_file) + if not path.exists(): + if row.get("required", row.get("status") != "ignored"): + raise FileNotFoundError(path) + report["skipped"].append({"source_name": row.get("source_name"), "reason": "missing ignored output"}) + continue + keys = output_keys_for_row(row, family) + for key in keys: + if key in seen: + raise RuntimeError(f"duplicate QDQ output key: {key}") + tensor, header = dequantize_path(path, out_dtype, args.row_batch_size) + scale_factor = float(row.get("scale_factor") or adapter.scale_factor(path.name)) + if scale_factor != 1.0 and torch.is_floating_point(tensor): + tensor = tensor / scale_factor + tensor = trim_to_manifest_shape(tensor, row.get("shape")) + for key in keys: + writer.add(key, tensor) + seen.add(key) + report["written"].append({ + "file": path.name, + "key": key, + "shape": list(tensor.shape), + "precision": header.precision, + "scale_factor": scale_factor, + }) + del tensor + writer.close() + report["written_count"] = len(report["written"]) + (args.out / "qdq_conversion_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + return report + finally: + if tmp is not None: + tmp.cleanup() diff --git a/python/cactus/convert/export/reports.py b/python/cactus/convert/export/reports.py new file mode 100644 index 000000000..47e35f73a --- /dev/null +++ b/python/cactus/convert/export/reports.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import json +from collections import Counter +from pathlib import Path +from typing import Any + + +def write_reports(out_dir: Path, rows: list[dict[str, Any]]) -> dict[str, Any]: + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "conversion_manifest.json").write_text(json.dumps(rows, indent=2, sort_keys=True), encoding="utf-8") + weights = [] + for row in rows: + output_file = str(row.get("output_file") or "") + if output_file.endswith((".weights", ".bias")): + weights.append({ + "source_name": row.get("source_name"), + "hf_name": row.get("hf_name") or row.get("source_name"), + "adapter_name": row.get("adapter_name") or row.get("source_name"), + "output_name": row.get("output_file"), + "shape": row.get("shape"), + "precision": row.get("precision"), + "status": row.get("status"), + "component": row.get("component"), + "scale_factor": row.get("scale_factor", 1.0), + "adapter_family": row.get("adapter_family"), + "source_names": row.get("source_names") or [row.get("source_name")], + "transform": row.get("transform", "none"), + "qdq_restore": row.get("qdq_restore", "hf_key"), + }) + (out_dir / "weights_manifest.json").write_text(json.dumps({"weights": weights}, indent=2, sort_keys=True), encoding="utf-8") + by_status = Counter(r["status"] for r in rows) + by_component = Counter(r["component"] for r in rows) + by_precision = Counter(r["precision"] for r in rows) + fallback = Counter(r.get("fallback_reason") or "" for r in rows if r["status"] in {"fallback", "unrecognized"}) + total_bytes = sum(int(r.get("bytes", 0) or 0) for r in rows) + summary = { + "total_tensors": len(rows), + "counts_by_status": dict(by_status), + "counts_by_component": dict(by_component), + "counts_by_precision": dict(by_precision), + "fallback_reasons": dict(fallback), + "total_bytes": total_bytes, + } + (out_dir / "conversion_summary.json").write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8") + return summary + + +def print_summary(summary: dict[str, Any]) -> None: + print("\nConversion summary") + print("------------------") + for key in ["converted", "fallback", "ignored", "unrecognized"]: + print(f"{key:13s} {summary.get('counts_by_status', {}).get(key, 0)}") + print(f"{'total bytes':13s} {summary.get('total_bytes', 0)}") diff --git a/python/cactus/convert/export/validate.py b/python/cactus/convert/export/validate.py new file mode 100644 index 000000000..f1adc4466 --- /dev/null +++ b/python/cactus/convert/export/validate.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import load_file + +from ..model_adapters.adapters import adapter_for_family +from ..model_adapters.detection import detect_family + + +def _dtype_name(dtype) -> str: + return str(dtype).replace("torch.", "") + + +def _is_integer_dtype(dtype_name: str) -> bool: + return any(x in dtype_name for x in ("int", "uint", "bool")) + + +def load_safetensor_index(root: Path) -> dict[str, tuple[tuple[int, ...], str]]: + out: dict[str, tuple[tuple[int, ...], str]] = {} + for shard in sorted(root.glob("*.safetensors")): + tensors = load_file(shard) + for key, tensor in tensors.items(): + out[key] = (tuple(int(x) for x in tensor.shape), _dtype_name(tensor.dtype)) + del tensors + if out: + return out + raise ValueError(f"no safetensors found under {root}") + + +def _state_index(state: dict) -> dict[str, tuple[tuple[int, ...], str]]: + return { + key: (tuple(int(x) for x in tensor.shape), _dtype_name(tensor.dtype)) + for key, tensor in state.items() + } + + +def load_source_state_index(model_path: str | Path, family: str = "auto") -> dict[str, tuple[tuple[int, ...], str]]: + path = Path(model_path) + if path.exists() and path.is_dir() and list(path.glob("*.safetensors")): + source = load_file(sorted(path.glob("*.safetensors"))[0]) + for shard in sorted(path.glob("*.safetensors"))[1:]: + source.update(load_file(shard)) + if family != "auto": + source = adapter_for_family(family).normalize_state_dict(source).state_dict + return _state_index(source) + try: + from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + + resolved_family = family + if resolved_family == "auto": + try: + cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True, local_files_only=path.exists()) + resolved_family = detect_family(cfg, "auto") + except Exception: + resolved_family = "generic" + + try: + model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True) + except Exception: + model = AutoModel.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True) + state = adapter_for_family(resolved_family).normalize_state_dict(model.state_dict()).state_dict + return _state_index(state) + except Exception as exc: + raise RuntimeError(f"could not load source model state from {model_path}") from exc + + +def validate_qdq(args) -> dict[str, Any]: + family = getattr(args, "model_family", "auto") + source = load_source_state_index(args.source_model, family) + qdq = load_safetensor_index(Path(args.qdq)) + source_keys = set(source) + qdq_keys = set(qdq) + missing = sorted(source_keys - qdq_keys) + extra = sorted(qdq_keys - source_keys) + shape_mismatches = [] + dtype_mismatches = [] + for key in sorted(source_keys & qdq_keys): + source_shape, source_dtype = source[key] + qdq_shape, qdq_dtype = qdq[key] + if source_shape != qdq_shape: + shape_mismatches.append({"key": key, "source": list(source_shape), "qdq": list(qdq_shape)}) + if _is_integer_dtype(source_dtype) and source_dtype != qdq_dtype: + dtype_mismatches.append({"key": key, "source": source_dtype, "qdq": qdq_dtype}) + report = { + "source_model": str(args.source_model), + "qdq": str(args.qdq), + "source_tensors": len(source), + "qdq_tensors": len(qdq), + "missing": missing, + "extra": extra, + "shape_mismatches": shape_mismatches, + "dtype_mismatches": dtype_mismatches, + "ok": not missing and not extra and not shape_mismatches and not dtype_mismatches, + } + if args.out: + Path(args.out).write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2)) + if args.strict and not report["ok"]: + raise SystemExit(1) + return report diff --git a/python/cactus/convert/handoff_probe.py b/python/cactus/convert/handoff_probe.py new file mode 100644 index 000000000..9dd8f9c24 --- /dev/null +++ b/python/cactus/convert/handoff_probe.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +import io +import hashlib +import json +import struct +import zipfile +from pathlib import Path +from typing import Any + + +_PROBE_MAGIC = b"CHP10P6\0" +_ORDERED_KEYS = ( + "norm.weight", + "norm.bias", + "proj.weight", + "proj.bias", + "attn_query", + "head.0.weight", + "head.0.bias", + "head.2.weight", + "head.2.bias", + "head.4.weight", + "head.4.bias", +) + + +def _packaged_asset_dir(model_id: str | None) -> Path | None: + if not model_id: + return None + from ..cli.download import get_model_dir_name + + return Path(__file__).resolve().parent / "assets" / get_model_dir_name(model_id) + + +def _candidate_probe_files(output_dir: Path, model_id: str | None = None) -> list[Path]: + cwd = Path.cwd() + candidates: list[Path] = [] + asset_dir = _packaged_asset_dir(model_id) + if asset_dir is not None: + candidates.append(asset_dir / "probe.pt") + candidates += [ + output_dir / "probe.pt", + output_dir / "global_attn_probe_v10p6.pt", + cwd / "probe.pt", + cwd / "v10p6_probe_release" / "global_attn_probe_v10p6.pt", + Path.home() / "Downloads" / "probe.pt", + Path.home() / "Downloads" / "v10p6_probe_release" / "global_attn_probe_v10p6.pt", + ] + return candidates + + +def _candidate_probe_zips(output_dir: Path) -> list[Path]: + cwd = Path.cwd() + return [ + output_dir / "v10p6_probe_release.zip", + cwd / "v10p6_probe_release.zip", + Path.home() / "Downloads" / "v10p6_probe_release.zip", + ] + + +def _load_checkpoint_from_zip(zip_path: Path) -> Any: + import torch + + with zipfile.ZipFile(zip_path) as zf: + names = set(zf.namelist()) + for name in ( + "v10p6_probe_release/global_attn_probe_v10p6.pt", + "global_attn_probe_v10p6.pt", + "probe.pt", + ): + if name in names: + with zf.open(name) as f: + return torch.load(io.BytesIO(f.read()), map_location="cpu") + raise FileNotFoundError(f"no probe checkpoint found in {zip_path}") + + +def _state_dict_from_checkpoint(checkpoint: Any) -> dict[str, Any]: + if isinstance(checkpoint, dict): + for key in ("state_dict", "model_state"): + value = checkpoint.get(key) + if isinstance(value, dict): + return value + return checkpoint + + +def _load_checkpoint(output_dir: Path, model_id: str | None = None) -> tuple[Any, str] | tuple[None, None]: + import torch + + for path in _candidate_probe_files(output_dir, model_id): + if path.exists(): + return torch.load(path, map_location="cpu"), str(path) + for path in _candidate_probe_zips(output_dir): + if path.exists(): + return _load_checkpoint_from_zip(path), str(path) + return None, None + + +def _tensor_hash(state: dict[str, Any]) -> str: + digest = hashlib.sha256() + for key in sorted(_ORDERED_KEYS): + tensor = state[key].detach().cpu().contiguous().float() + digest.update(key.encode("utf-8")) + digest.update(str(tuple(tensor.shape)).encode("utf-8")) + digest.update(tensor.numpy().tobytes(order="C")) + return digest.hexdigest() + + +def _probe_state_dict(checkpoint: Any) -> dict[str, Any]: + state = _state_dict_from_checkpoint(checkpoint) + missing = [key for key in _ORDERED_KEYS if key not in state] + if missing: + raise RuntimeError(f"handoff probe checkpoint missing tensors: {', '.join(missing)}") + return state + + +def _write_probe_bundle(out_dir: Path, state: dict[str, Any], *, feat_dim: int, t_h: int, + h1: int, h2: int, fmt: str, layer: int, source: str | None) -> bool: + out_dir.mkdir(parents=True, exist_ok=True) + probe_path = out_dir / "handoff_probe.bin" + with probe_path.open("wb") as f: + f.write(_PROBE_MAGIC) + f.write(struct.pack(" bool: + """Package the Gemma4 v10p6 cloud-handoff probe into a C++-readable bundle file.""" + model_key = (model_id or "").lower() + if model_key and "gemma-4" not in model_key and "gemma4" not in model_key: + return False + + out_dir = Path(output_dir) + checkpoint, source = _load_checkpoint(out_dir, model_id) + if checkpoint is None: + return False + + state = _probe_state_dict(checkpoint) + return _write_probe_bundle( + out_dir, state, feat_dim=1536, t_h=32, h1=128, h2=64, + fmt="cactus_handoff_probe_v10p6", layer=28, source=source, + ) + + +_PARAKEET_PROBE_FILE = "parakeet_v3_probe.pt" +_PARAKEET_PROBE_LAYER = 23 + + +def _candidate_parakeet_probe_files(output_dir: Path, model_id: str | None = None) -> list[Path]: + candidates: list[Path] = [] + asset_dir = _packaged_asset_dir(model_id) + if asset_dir is not None: + candidates.append(asset_dir / _PARAKEET_PROBE_FILE) + candidates += [ + output_dir / _PARAKEET_PROBE_FILE, + Path.cwd() / _PARAKEET_PROBE_FILE, + Path.home() / "Downloads" / _PARAKEET_PROBE_FILE, + ] + return candidates + + +def export_parakeet_handoff_probe(output_dir: str | Path, *, model_id: str | None = None) -> bool: + """Package the Parakeet-TDT-v3 cloud-handoff probe into a C++-readable bundle file.""" + model_key = (model_id or "").lower() + if model_key and "parakeet" not in model_key: + return False + + out_dir = Path(output_dir) + checkpoint = source = None + for path in _candidate_parakeet_probe_files(out_dir, model_id): + if path.exists(): + import torch + + checkpoint, source = torch.load(path, map_location="cpu"), str(path) + break + if checkpoint is None: + return False + + state = _probe_state_dict(checkpoint) + return _write_probe_bundle( + out_dir, state, + feat_dim=int(state["norm.weight"].shape[0]), + t_h=int(state["proj.weight"].shape[0]), + h1=int(state["head.0.weight"].shape[0]), + h2=int(state["head.2.weight"].shape[0]), + fmt="cactus_handoff_probe_parakeet", layer=_PARAKEET_PROBE_LAYER, source=source, + ) diff --git a/python/cactus/convert/interleave_orthogonal_cq4.py b/python/cactus/convert/interleave_orthogonal_cq4.py new file mode 100644 index 000000000..a7d999766 --- /dev/null +++ b/python/cactus/convert/interleave_orthogonal_cq4.py @@ -0,0 +1,97 @@ +# Temporary artifact helper: rewrites full-width orthogonal CQ4 weights into +# INTERLEAVED_4ROW storage without full model reconversion. +from __future__ import annotations + +import argparse +import struct +from pathlib import Path + +import numpy as np + +from .export.qdq import ( + ALIGNMENT_DEFAULT, + CACTUS_MAGIC, + FLAG_INTERLEAVED_4ROW, + FLAG_ORTHOGONAL_ROTATION, + HEADER_SIZE, + align_offset, + read_cq_payload, + read_header, + unpack_lsb_values, +) +from .quantization.cq import PRECISION_CQ, pack_indices_interleaved_4row + + +def _padding(offset: int, alignment: int) -> bytes: + return b"\0" * (align_offset(offset, alignment) - offset) + + +def interleave_orthogonal_cq4_file(input_path: Path, output_path: Path, *, force: bool = False) -> None: + if output_path.exists() and not force: + raise FileExistsError(f"{output_path} exists; pass --force to overwrite") + + header = read_header(input_path) + if header.ndim != 2: + raise ValueError(f"{input_path}: expected 2D CQ tensor") + n, k = header.shape + if header.precision != PRECISION_CQ[4] or header.bits != 4: + raise ValueError(f"{input_path}: expected CQ4") + if (header.flags & FLAG_ORTHOGONAL_ROTATION) == 0: + raise ValueError(f"{input_path}: expected orthogonal CQ rotation") + if header.flags & FLAG_INTERLEAVED_4ROW: + raise ValueError(f"{input_path}: already INTERLEAVED_4ROW") + if header.num_groups != 1 or header.group_size != k: + raise ValueError( + f"{input_path}: expected one full-width group, got group_size={header.group_size} num_groups={header.num_groups}" + ) + if n % 4 != 0 or k % 32 != 0: + raise ValueError(f"{input_path}: expected N % 4 == 0 and K % 32 == 0, got shape={header.shape}") + + scales_blob, packed = read_cq_payload(input_path, header) + rows = packed.reshape(n, k // 2) + indices = np.stack([unpack_lsb_values(row, k, 4) for row in rows]) + interleaved = pack_indices_interleaved_4row(indices, k, 4) + + output_path.parent.mkdir(parents=True, exist_ok=True) + alignment = header.alignment or ALIGNMENT_DEFAULT + scales_offset = align_offset(HEADER_SIZE, alignment) + scales_end = scales_offset + len(scales_blob) + data_offset = align_offset(scales_end, alignment) + flags = header.flags | FLAG_INTERLEAVED_4ROW + + with output_path.open("wb") as out: + out.write(CACTUS_MAGIC) + out.write(struct.pack(" None: + parser = argparse.ArgumentParser(description="Rewrite a full-width orthogonal CQ4 file to 4-row interleaved storage.") + parser.add_argument("input", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--force", action="store_true") + args = parser.parse_args() + interleave_orthogonal_cq4_file(args.input, args.output, force=args.force) + print(f"wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/python/cactus/convert/model_adapters/__init__.py b/python/cactus/convert/model_adapters/__init__.py new file mode 100644 index 000000000..9b3530067 --- /dev/null +++ b/python/cactus/convert/model_adapters/__init__.py @@ -0,0 +1,2 @@ +"""Model-family detection, naming, and tensor policy.""" + diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py new file mode 100644 index 000000000..41eaa86dc --- /dev/null +++ b/python/cactus/convert/model_adapters/adapters.py @@ -0,0 +1,958 @@ +from __future__ import annotations + +from dataclasses import dataclass, field, replace +import json +from pathlib import Path +import re +from typing import Any + +import numpy as np + +try: + import torch +except Exception: # pragma: no cover + torch = None + +from ..cactus_adapters.config_utils import ( + extract_audio_config, + extract_base_config, + extract_complex_gemma_config, + extract_lfm2_config, + extract_parakeet_config, + extract_parakeet_tdt_config, + extract_vision_config, + extract_whisper_config, +) +from .naming import NameMatch, cactus_name_for_tensor, gemma3_scale_factor, gemma4_scale_factor, restore_hf_key_for_family +from .policy import TensorPolicy, policy_for_tensor +from ..compat import patch_transformers_import_compat + + +@dataclass(frozen=True) +class TensorProvenance: + source_names: list[str] + transform: str = "none" + qdq_restore: str = "hf_key" + + +@dataclass +class NormalizedState: + state_dict: dict[str, Any] + provenance: dict[str, TensorProvenance] = field(default_factory=dict) + + +@dataclass(frozen=True) +class TensorEmission: + output_name: str + tensor: Any + transform: str = "none" + source_names: list[str] | None = None + qdq_restore: str | None = None + + +class FamilyAdapter: + family = "generic" + + def __init__(self, family: str | None = None) -> None: + if family is not None: + self.family = family + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + text_cfg = _cfg_get(cfg, "text_config", None) + base_cfg = text_cfg if text_cfg is not None else cfg + return extract_base_config(base_cfg, cfg) + + def runtime_model_type(self) -> str: + return self.family + + def model_class(self, cfg: Any): + patch_transformers_import_compat() + from transformers import AutoModel, AutoModelForCausalLM + + arch = " ".join(_cfg_get(cfg, "architectures", []) or []).lower() + if not any(x in arch for x in ["causallm", "conditionalgeneration", "forctc"]): + return AutoModel + return AutoModelForCausalLM + + def load_processor(self, model_id_or_path: str): + patch_transformers_import_compat() + from transformers import AutoProcessor, AutoTokenizer + + try: + return AutoProcessor.from_pretrained(model_id_or_path, trust_remote_code=True, local_files_only=Path(model_id_or_path).exists()) + except Exception: + try: + return AutoTokenizer.from_pretrained(model_id_or_path, trust_remote_code=True, local_files_only=Path(model_id_or_path).exists()) + except Exception: + return None + + def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: + return NormalizedState(state_dict=dict(state_dict)) + + def name_tensor(self, source_name: str, _tensor: Any, num_layers: int | None) -> NameMatch: + return cactus_name_for_tensor(source_name, self.family, num_layers) + + def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + return policy_for_tensor(match, shape, requested_bits, self.family) + + def transform_tensor(self, match: NameMatch, tensor: Any) -> tuple[Any, str]: + return tensor, "none" + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + tensor, transform = self.transform_tensor(match, tensor) + if match.transpose: + tensor = _transpose_last_two(tensor) + transform = f"{transform}+transpose" if transform != "none" else "transpose" + return [TensorEmission(match.output_name, tensor, transform)] + + def module_target_name(self, source_name: str, _model: Any) -> str | None: + return source_name[:-7] if source_name.endswith(".weight") else source_name + + def build_calibration_inputs(self, row: dict[str, Any], processor: Any, modality: str, base_dir) -> dict[str, Any] | None: + if modality == "transcription": + return None + return None + + def qdq_output_keys(self, row: dict[str, Any]) -> list[str]: + if row.get("qdq_restore") == "runtime_key": + output_file = row.get("output_file") + if not output_file: + raise ValueError(f"manifest row has no runtime output file: {row}") + key = str(output_file) + return [key[:-8] if key.endswith(".weights") else key] + key = row.get("hf_name") or row.get("source_name") + if not key: + raise ValueError(f"manifest row has no hf/source key: {row}") + return [restore_hf_key_for_family(str(key), self.family)] + + def scale_factor(self, output_name: str) -> float: + return 1.0 + + +class Gemma4Adapter(FamilyAdapter): + family = "gemma4" + + def __init__(self) -> None: + super().__init__() + self.kv_shared_from_layer = 15 + + def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: + out = dict(state_dict) + provenance: dict[str, TensorProvenance] = {} + down_pattern = re.compile( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.(moe|experts)\.down_proj" + ) + for source_name, tensor in list(state_dict.items()): + match = down_pattern.fullmatch(source_name) + if match is None: + continue + shape = _tensor_shape(tensor) + if len(shape) != 3: + continue + block_name = match.group(2) + scale_name = source_name[: -len(".down_proj")] + ".per_expert_scale" + if block_name == "experts": + prefix = source_name[: -len(".experts.down_proj")] + scale_name = f"{prefix}.router.per_expert_scale" + scale = state_dict.get(scale_name) + if scale is None: + raise ValueError( + f"gemma4 MoE down_proj {source_name!r} is missing per-expert scale {scale_name!r}" + ) + scale_shape = _tensor_shape(scale) + if len(scale_shape) != 1 or int(shape[0]) != int(scale_shape[0]): + raise ValueError( + f"gemma4 MoE down_proj {source_name!r} scale {scale_name!r} shape {scale_shape} " + f"does not match {int(shape[0])} experts" + ) + if torch is not None and isinstance(tensor, torch.Tensor): + scale_tensor = scale if isinstance(scale, torch.Tensor) else torch.as_tensor(scale) + scale_value = scale_tensor.to(device=tensor.device, dtype=tensor.dtype).reshape(-1, 1, 1) + out[source_name] = (tensor * scale_value).contiguous() + else: + out[source_name] = (np.asarray(tensor) * np.asarray(scale).reshape(-1, 1, 1)).copy() + provenance[source_name] = TensorProvenance( + [source_name, scale_name], + "gemma4_moe_down_proj_per_expert_scale", + "hf_key", + ) + return NormalizedState(out, provenance) + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + text_cfg = _cfg_get(cfg, "text_config", None) + base_cfg = text_cfg if text_cfg is not None else cfg + config = extract_base_config(base_cfg, cfg) + config.update(extract_complex_gemma_config(base_cfg, cfg)) + num_layers = int(config.get("num_layers", 0) or 0) + num_kv_shared_layers = int(config.get("num_kv_shared_layers", 0) or 0) + if num_layers > 0 and num_kv_shared_layers > 0: + self.kv_shared_from_layer = max(0, num_layers - num_kv_shared_layers) + vision_cfg = _cfg_get(cfg, "vision_config", None) + if vision_cfg: + config.update(extract_vision_config(cfg, vision_cfg)) + audio_cfg = _cfg_get(cfg, "audio_config", None) + if audio_cfg: + config.update(extract_audio_config(cfg, audio_cfg)) + return config + + def model_class(self, cfg: Any): + from transformers import AutoModel + + try: + from transformers import Gemma4ForConditionalGeneration + + return Gemma4ForConditionalGeneration + except Exception: + return AutoModel + + def scale_factor(self, output_name: str) -> float: + return gemma4_scale_factor(output_name) + + def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + if match.output_name and ( + "moe_router" in match.output_name + or "moe_expert_bias" in match.output_name + or "moe_per_expert_scale" in match.output_name + or match.source_name.endswith(".router.proj.weight") + or match.source_name.endswith(".router.scale") + ): + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe routing tensor") + policy = super().policy(match, shape, requested_bits) + name = match.source_name + if name == "model.embed_vision.embedding_projection.weight": + return policy + if match.output_name and "moe_expert_" in match.output_name: + return replace(policy, use_gptq=False) + if ".self_attn." in name and (name.endswith(".k_proj.weight") or name.endswith(".v_proj.weight")): + parts = name.split(".") + try: + layer = int(parts[3]) if parts[:3] == ["model", "language_model", "layers"] else None + except Exception: + layer = None + if layer is not None and layer >= self.kv_shared_from_layer: + return replace(policy, use_gptq=False, fallback_reason=policy.fallback_reason or "shared KV tensor has no per-layer hook module") + return policy + + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.(?:moe|experts)\.(gate_up_proj|down_proj)", + source_name, + ) + if match is not None: + layer_index = int(match.group(1)) + kind = match.group(2) + if kind == "gate_up_proj": + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_gate_up.weights" + else: + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_w2.weights" + return NameMatch(source_name, output_name, "language", True, hf_name=source_name, adapter_name=source_name) + return super().name_tensor(source_name, tensor, num_layers) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + source_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.(?:moe|experts)\.(gate_up_proj|down_proj)", + match.source_name, + ) + if source_match is None: + return super().expand_tensor(match, tensor) + + layer_index = int(source_match.group(1)) + kind = source_match.group(2) + shape = _tensor_shape(tensor) + if len(shape) != 3: + raise ValueError(f"Gemma4-MoE expert tensor expects rank 3, got {shape}") + num_experts = int(shape[0]) + emissions: list[TensorEmission] = [] + + def _piece(value, expert_idx: int, start: int | None = None, end: int | None = None): + if torch is not None and isinstance(value, torch.Tensor): + selected = value[expert_idx] if start is None else value[expert_idx, start:end, :] + return selected.contiguous() + array = np.asarray(value) + selected = array[expert_idx] if start is None else array[expert_idx, start:end, :] + return selected.copy() + + def _source_names(weight_name: str, expert_idx: int) -> list[str]: + return [ + f"model.language_model.layers.{layer_index}.moe.{weight_name}.{expert_idx}", + f"model.layers.{layer_index}.moe.{weight_name}.{expert_idx}", + f"layers.{layer_index}.moe.{weight_name}.{expert_idx}", + f"backbone.layers.{layer_index}.moe.{weight_name}.{expert_idx}", + ] + + if kind == "gate_up_proj": + if int(shape[1]) % 2 != 0: + raise ValueError(f"Gemma4-MoE gate_up_proj second dim must be even, got {shape}") + intermediate = int(shape[1]) // 2 + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w1.weights", + _piece(tensor, expert_idx, 0, intermediate), + "gemma4_moe_gate_up_split_gate", + source_names=_source_names("w1_weights", expert_idx), + qdq_restore="runtime_key", + )) + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w3.weights", + _piece(tensor, expert_idx, intermediate, int(shape[1])), + "gemma4_moe_gate_up_split_up", + source_names=_source_names("w3_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w2.weights", + _piece(tensor, expert_idx), + "gemma4_moe_down_proj_split_scaled", + source_names=_source_names("w2_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ( + ".moe.gate_up_proj" in source_name + or ".moe.down_proj" in source_name + or ".experts.gate_up_proj" in source_name + or ".experts.down_proj" in source_name + ): + return None + target = super().module_target_name(source_name, _model) + if target and source_name.startswith("model.vision_tower."): + modules = dict(_model.named_modules()) if _model is not None else {} + if target in modules and hasattr(modules[target], "weight"): + return target + linear_target = target if target.endswith(".linear") else f"{target}.linear" + if linear_target in modules: + return linear_target + return target + + +class QwenAdapter(FamilyAdapter): + family = "qwen" + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + config = super().runtime_config(cfg) + model_type = str(_cfg_get(cfg, "model_type", "") or "").lower() + if model_type == "qwen2_moe" or _cfg_get(cfg, "moe_intermediate_size", None) is not None: + config.update({ + "moe_intermediate_size": int(_cfg_get(cfg, "moe_intermediate_size", 0) or 0), + "shared_expert_intermediate_size": int(_cfg_get(cfg, "shared_expert_intermediate_size", 0) or 0), + "decoder_sparse_step": int(_cfg_get(cfg, "decoder_sparse_step", 0) or 0), + "num_experts": int(_cfg_get(cfg, "num_experts", config.get("num_experts", 0)) or 0), + "num_experts_per_tok": int(_cfg_get(cfg, "num_experts_per_tok", config.get("num_experts_per_tok", 0)) or 0), + "norm_topk_prob": bool(_cfg_get(cfg, "norm_topk_prob", False)), + }) + return config + + def model_class(self, cfg: Any): + from transformers import AutoModel, AutoModelForCausalLM + + arch = " ".join(_cfg_get(cfg, "architectures", []) or []).lower() + model_type = str(_cfg_get(cfg, "model_type", "") or "").lower().replace("_", "-") + if "qwen3vlforconditionalgeneration" in arch or model_type == "qwen3-vl": + try: + from transformers import Qwen3VLForConditionalGeneration + + return Qwen3VLForConditionalGeneration + except Exception: + return AutoModel + return AutoModelForCausalLM + + def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + if match.output_name and ( + "moe_router" in match.output_name + or "moe_expert_" in match.output_name + or "shared_expert_gate" in match.output_name + ): + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe routing/expert tensor") + return super().policy(match, shape, requested_bits) + + def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ".mlp.experts.gate_up_proj" in source_name or ".mlp.experts.down_proj" in source_name: + return None + return super().module_target_name(source_name, _model) + + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)", + source_name, + ) + if match is not None: + layer_index = int(match.group(1)) + kind = match.group(2) + if kind == "gate_up_proj": + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_gate_up.weights" + else: + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_w2.weights" + return NameMatch(source_name, output_name, "language", True, hf_name=source_name, adapter_name=source_name) + + router_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.mlp\.gate\.weight", + source_name, + ) + if router_match is not None: + layer_index = int(router_match.group(1)) + return NameMatch( + source_name, + f"layer_{layer_index}_moe_router.weights", + "language", + True, + hf_name=source_name, + adapter_name=source_name, + ) + + return super().name_tensor(source_name, tensor, num_layers) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + source_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)", + match.source_name, + ) + if source_match is None: + return super().expand_tensor(match, tensor) + + layer_index = int(source_match.group(1)) + kind = source_match.group(2) + shape = _tensor_shape(tensor) + if len(shape) != 3: + raise ValueError(f"Qwen2-MoE expert tensor expects rank 3, got {shape}") + num_experts = int(shape[0]) + emissions: list[TensorEmission] = [] + + def _piece(value, expert_idx: int, start: int | None = None, end: int | None = None): + if torch is not None and isinstance(value, torch.Tensor): + selected = value[expert_idx] if start is None else value[expert_idx, start:end, :] + return selected.contiguous() + array = np.asarray(value) + selected = array[expert_idx] if start is None else array[expert_idx, start:end, :] + return selected.copy() + + def _source_names(weight_name: str, expert_idx: int) -> list[str]: + return [ + f"model.layers.{layer_index}.mlp.{weight_name}.{expert_idx}", + f"layers.{layer_index}.mlp.{weight_name}.{expert_idx}", + f"backbone.layers.{layer_index}.mlp.{weight_name}.{expert_idx}", + ] + + if kind == "gate_up_proj": + if int(shape[1]) % 2 != 0: + raise ValueError(f"Qwen2-MoE gate_up_proj second dim must be even, got {shape}") + intermediate = int(shape[1]) // 2 + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w1.weights", + _piece(tensor, expert_idx, 0, intermediate), + "qwen2_moe_gate_up_split_gate", + source_names=_source_names("w1_weights", expert_idx), + qdq_restore="runtime_key", + )) + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w3.weights", + _piece(tensor, expert_idx, intermediate, int(shape[1])), + "qwen2_moe_gate_up_split_up", + source_names=_source_names("w3_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w2.weights", + _piece(tensor, expert_idx), + "qwen2_moe_down_proj_split", + source_names=_source_names("w2_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + +class WhisperAdapter(FamilyAdapter): + family = "whisper" + + def model_class(self, cfg: Any): + from transformers import AutoModelForSpeechSeq2Seq + + return AutoModelForSpeechSeq2Seq + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return extract_whisper_config(cfg) + + def build_calibration_inputs(self, row: dict[str, Any], processor: Any, modality: str, base_dir) -> dict[str, Any] | None: + if modality != "transcription" or processor is None: + return None + audio_rel = row.get("audio_path") + if not audio_rel: + return None + audio_path = base_dir / audio_rel + if not audio_path.exists(): + return None + audio = _load_audio_16k(audio_path) + try: + return processor(audio, sampling_rate=16000, return_tensors="pt") + except Exception: + return None + + +class ParakeetAdapter(FamilyAdapter): + family = "parakeet" + + def model_class(self, cfg: Any): + from transformers import AutoModel, AutoModelForCTC + + arch = " ".join(_cfg_get(cfg, "architectures", []) or []).lower() + model_type = str(_cfg_get(cfg, "model_type", "") or "").lower() + if "parakeetforctc" in arch or model_type == "parakeet_ctc": + try: + return AutoModelForCTC + except Exception: + try: + from transformers import ParakeetForCTC + + return ParakeetForCTC + except Exception: + return AutoModel + return AutoModel + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return extract_parakeet_config(cfg) + + def transform_tensor(self, match: NameMatch, tensor: Any) -> tuple[Any, str]: + out = match.output_name + if out not in { + "subsampling_conv0_weight.weights", + "subsampling_depthwise1_weight.weights", + "subsampling_pointwise1_weight.weights", + "subsampling_depthwise2_weight.weights", + "subsampling_pointwise2_weight.weights", + }: + return tensor, "none" + if len(_tensor_shape(tensor)) != 4: + return tensor, "none" + shape = _tensor_shape(tensor) + is_hwio_k3 = shape[1] == 3 and shape[2] == 3 and shape[3] == 1 + is_hwio_pw = shape[1] == 1 and shape[2] == 1 and shape[3] > 1 + if not (is_hwio_k3 or is_hwio_pw): + return tensor, "none" + if torch is not None and isinstance(tensor, torch.Tensor): + return tensor.permute(0, 3, 1, 2).contiguous(), "parakeet_conv_hwio_to_oihw" + return np.transpose(np.asarray(tensor), (0, 3, 1, 2)).copy(), "parakeet_conv_hwio_to_oihw" + + +class NomicAdapter(FamilyAdapter): + family = "nomic" + + def __init__(self) -> None: + super().__init__() + self.num_experts = 0 + + def runtime_model_type(self) -> str: + return "bert" + + def model_class(self, cfg: Any): + from transformers import AutoModel + + return AutoModel + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + hidden_dim = int(_cfg_get(cfg, "n_embd", _cfg_get(cfg, "hidden_size", 0))) + heads = int(_cfg_get(cfg, "n_head", _cfg_get(cfg, "num_attention_heads", 0))) + self.num_experts = int(_cfg_get(cfg, "num_experts", 0) or 0) + return { + "vocab_size": int(_cfg_get(cfg, "vocab_size", 0)), + "hidden_dim": hidden_dim, + "num_layers": int(_cfg_get(cfg, "n_layer", _cfg_get(cfg, "num_hidden_layers", 0))), + "attention_heads": heads, + "attention_kv_heads": heads, + "attention_head_dim": int(hidden_dim // max(1, heads)), + "ffn_intermediate_dim": int(_cfg_get(cfg, "n_inner", _cfg_get(cfg, "intermediate_size", 0))), + "context_length": int(_cfg_get(cfg, "n_positions", _cfg_get(cfg, "max_position_embeddings", 0))), + "rope_theta": float(_cfg_get(cfg, "rotary_emb_base", _cfg_get(cfg, "rope_theta", 10000.0)) or 10000.0), + "layer_norm_eps": float(_cfg_get(cfg, "layer_norm_epsilon", _cfg_get(cfg, "layer_norm_eps", 1e-5)) or 1e-5), + "num_experts": self.num_experts, + "num_shared_experts": int(_cfg_get(cfg, "num_shared_experts", 0) or 0), + "num_top_experts": int(_cfg_get(cfg, "moe_top_k", _cfg_get(cfg, "num_top_experts", 0)) or 0), + "num_experts_per_tok": int(_cfg_get(cfg, "moe_top_k", _cfg_get(cfg, "num_experts_per_tok", 0)) or 0), + "moe_every_n_layers": int(_cfg_get(cfg, "moe_every_n_layers", 0) or 0), + "tie_word_embeddings": True, + } + + def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: + out = dict(state_dict) + provenance: dict[str, TensorProvenance] = {} + word = state_dict.get("embeddings.word_embeddings.weight") + if word is not None: + token_type = state_dict.get("embeddings.token_type_embeddings.weight") + if token_type is not None: + out["token_embeddings"] = word + token_type + sources = ["embeddings.word_embeddings.weight", "embeddings.token_type_embeddings.weight"] + transform = "nomic_word_token_type_embedding_sum" + else: + out["token_embeddings"] = word + sources = ["embeddings.word_embeddings.weight"] + transform = "nomic_word_embedding_alias" + out.pop("embeddings.word_embeddings.weight", None) + out.pop("embeddings.token_type_embeddings.weight", None) + provenance["token_embeddings"] = TensorProvenance(sources, transform, "adapter_key") + if "emb_ln.weight" in out: + out["embedding_layernorm.weight"] = out.pop("emb_ln.weight") + provenance["embedding_layernorm.weight"] = TensorProvenance(["emb_ln.weight"], "nomic_embedding_layernorm_rename", "adapter_key") + if "emb_ln.bias" in out: + out["embedding_layernorm.bias"] = out.pop("emb_ln.bias") + provenance["embedding_layernorm.bias"] = TensorProvenance(["emb_ln.bias"], "nomic_embedding_layernorm_rename", "adapter_key") + return NormalizedState(out, provenance) + + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + globals_map = { + "token_embeddings": "token_embeddings.weights", + "embedding_layernorm.weight": "embedding_layernorm.weight", + "embedding_layernorm.bias": "embedding_layernorm.bias", + } + if source_name in globals_map: + return NameMatch(source_name, globals_map[source_name], "embedding", True, hf_name=source_name, adapter_name=source_name) + norm2 = _nomic_layer_suffix(source_name, ".norm2.weight") + if norm2 is not None: + return NameMatch(source_name, f"layer_{norm2}_norm2.weights", "language", True, hf_name=source_name, adapter_name=source_name) + for suffix, template, transpose in ( + (".attn.Wqkv.weight", "layer_{i}_attn_qkv.weights", False), + (".attn.Wqkv.bias", "layer_{i}_attn_qkv.bias", False), + (".mlp.experts.mlp.w1", "layer_{i}_mlp_experts_w1.weights", False), + (".mlp.experts.mlp.w2", "layer_{i}_mlp_experts_w2.weights", True), + ): + layer = _nomic_layer_suffix(source_name, suffix) + if layer is not None: + return NameMatch(source_name, template.format(i=layer), "language", True, transpose, hf_name=source_name, adapter_name=source_name) + return cactus_name_for_tensor(source_name, self.family, num_layers) + + def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + # The MoE router is tiny ([num_experts, hidden]) but decides expert selection; + # 4-bit quantizing it corrupts routing, so keep it in FP16. + if ".mlp.router.layer.weight" in match.source_name: + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe router precision-sensitive") + policy = super().policy(match, shape, requested_bits) + if policy.use_gptq and ".mlp.experts.mlp." in match.source_name: + return replace(policy, use_gptq=False) + return policy + + def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ".mlp.experts.mlp." in source_name: + return None + return super().module_target_name(source_name, _model) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + out = match.output_name + if "{channel}" in out: + if "attn_{channel}" in out: + return split_channel_tensor(tensor, out, ["q", "k", "v"], "nomic_qkv_split", match.transpose) + if "mlp_expert_{channel}" in out: + num_experts = self.num_experts or 1 + return split_channel_tensor(tensor, out, [str(i) for i in range(num_experts)], "nomic_moe_expert_split", match.transpose) + return super().expand_tensor(match, tensor) + + +class ParakeetTDTAdapter(ParakeetAdapter): + family = "parakeet_tdt" + + def model_class(self, cfg: Any): + from transformers import AutoModel + + return AutoModel + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + return extract_parakeet_tdt_config(cfg) + + def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: + augmented = dict(state_dict) + provenance: dict[str, TensorProvenance] = {} + for prefix in ("decoder.lstm", "decoder.prediction.dec_rnn.lstm"): + i = 0 + while True: + ih_key = f"{prefix}.bias_ih_l{i}" + hh_key = f"{prefix}.bias_hh_l{i}" + out_key = f"{prefix}.bias_l{i}" + if ih_key not in state_dict and hh_key not in state_dict: + break + sources = [k for k in (ih_key, hh_key) if k in state_dict] + if ih_key in state_dict and hh_key in state_dict: + augmented[out_key] = state_dict[ih_key] + state_dict[hh_key] + elif ih_key in state_dict: + augmented[out_key] = state_dict[ih_key] + else: + augmented[out_key] = state_dict[hh_key] + augmented.pop(ih_key, None) + augmented.pop(hh_key, None) + provenance[out_key] = TensorProvenance(sources, "parakeet_tdt_lstm_bias_sum", "hf_key") + i += 1 + return NormalizedState(augmented, provenance) + + +class Lfm2Adapter(FamilyAdapter): + family = "lfm2" + + def __init__(self) -> None: + super().__init__() + self.num_experts = 0 + + def model_class(self, cfg: Any): + from transformers import AutoModelForCausalLM + + arch = " ".join(_cfg_get(cfg, "architectures", []) or []).lower() + model_type = str(_cfg_get(cfg, "model_type", "") or "").lower().replace("_", "-") + if "lfm2vlforconditionalgeneration" in arch or model_type == "lfm2-vl": + try: + from transformers import Lfm2VlForConditionalGeneration + + return Lfm2VlForConditionalGeneration + except Exception: + pass + return AutoModelForCausalLM + + def runtime_config(self, cfg: Any) -> dict[str, Any]: + config = super().runtime_config(cfg) + config.update(extract_lfm2_config(cfg)) + self.num_experts = int(_cfg_get(cfg, "num_experts", config.get("num_experts", 0)) or 0) + return config + + def load_processor(self, model_id_or_path: str): + processor = super().load_processor(model_id_or_path) + if processor is not None and hasattr(processor, "tokenizer"): + return processor + root = Path(model_id_or_path) + has_vl_processor_config = ( + root.exists() + and (root / "tokenizer.json").exists() + and (root / "preprocessor_config.json").exists() + ) + if not has_vl_processor_config: + return processor + from transformers import Lfm2VlImageProcessorFast, Lfm2VlProcessor, PreTrainedTokenizerFast + + cfg_path = root / "tokenizer_config.json" + cfg = json.loads(cfg_path.read_text(encoding="utf-8")) if cfg_path.exists() else {} + tokenizer = PreTrainedTokenizerFast( + tokenizer_file=str(root / "tokenizer.json"), + bos_token=cfg.get("bos_token", "<|startoftext|>"), + eos_token=cfg.get("eos_token", "<|im_end|>"), + pad_token=cfg.get("pad_token", "<|pad|>"), + clean_up_tokenization_spaces=cfg.get("clean_up_tokenization_spaces", True), + model_max_length=cfg.get("model_max_length", int(1e30)), + ) + for key in ("image_token", "image_start_token", "image_end_token", "image_thumbnail"): + value = cfg.get(key) + if value is None: + continue + setattr(tokenizer, key, value) + setattr(tokenizer, f"{key}_id", tokenizer.convert_tokens_to_ids(value)) + tokenizer.model_specific_special_tokens = cfg.get("model_specific_special_tokens", {}) + image_processor = Lfm2VlImageProcessorFast.from_pretrained(str(root), local_files_only=True) + return Lfm2VlProcessor(image_processor=image_processor, tokenizer=tokenizer) + + def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + if match.output_name and ( + "moe_router" in match.output_name + or "moe_expert_bias" in match.output_name + ): + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe routing tensor") + policy = super().policy(match, shape, requested_bits) + if match.source_name.endswith("vision_model.embeddings.patch_embedding.weight"): + return replace(policy, use_gptq=False, fallback_reason=policy.fallback_reason or "vision patch embedding has no linear Hessian target") + if match.output_name and "moe_expert_" in match.output_name: + return replace(policy, use_gptq=False) + return policy + + def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ".feed_forward.experts.gate_up_proj" in source_name or ".feed_forward.experts.down_proj" in source_name: + return None + if source_name.startswith("model.vision_tower.vision_model."): + source_name = source_name.replace("model.vision_tower.vision_model.", "model.vision_tower.", 1) + return super().module_target_name(source_name, _model) + + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.feed_forward\.experts\.(gate_up_proj|down_proj)", + source_name, + ) + if match is not None: + layer_index = int(match.group(1)) + kind = match.group(2) + if kind == "gate_up_proj": + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_gate_up.weights" + else: + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_w2.weights" + return NameMatch(source_name, output_name, "language", True, hf_name=source_name, adapter_name=source_name) + return super().name_tensor(source_name, tensor, num_layers) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + source_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.feed_forward\.experts\.(gate_up_proj|down_proj)", + match.source_name, + ) + if source_match is None: + return super().expand_tensor(match, tensor) + + layer_index = int(source_match.group(1)) + kind = source_match.group(2) + shape = _tensor_shape(tensor) + if len(shape) != 3: + raise ValueError(f"LFM2-MoE expert tensor expects rank 3, got {shape}") + num_experts = int(shape[0]) + self.num_experts = self.num_experts or num_experts + emissions: list[TensorEmission] = [] + + def _piece(value, expert_idx: int, start: int | None = None, end: int | None = None): + if torch is not None and isinstance(value, torch.Tensor): + selected = value[expert_idx] if start is None else value[expert_idx, start:end, :] + return selected.contiguous() + array = np.asarray(value) + selected = array[expert_idx] if start is None else array[expert_idx, start:end, :] + return selected.copy() + + def _source_names(weight_name: str, expert_idx: int) -> list[str]: + return [ + f"model.layers.{layer_index}.feed_forward.{weight_name}.{expert_idx}", + f"layers.{layer_index}.feed_forward.{weight_name}.{expert_idx}", + f"backbone.layers.{layer_index}.feed_forward.{weight_name}.{expert_idx}", + ] + + if kind == "gate_up_proj": + if int(shape[1]) % 2 != 0: + raise ValueError(f"LFM2-MoE gate_up_proj second dim must be even, got {shape}") + intermediate = int(shape[1]) // 2 + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w1.weights", + _piece(tensor, expert_idx, 0, intermediate), + "lfm2_moe_gate_up_split_gate", + source_names=_source_names("w1_weights", expert_idx), + qdq_restore="runtime_key", + )) + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w3.weights", + _piece(tensor, expert_idx, intermediate, int(shape[1])), + "lfm2_moe_gate_up_split_up", + source_names=_source_names("w3_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w2.weights", + _piece(tensor, expert_idx), + "lfm2_moe_down_proj_split", + source_names=_source_names("w2_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + +class Gemma3Adapter(FamilyAdapter): + family = "gemma3" + + def scale_factor(self, output_name: str) -> float: + return gemma3_scale_factor(output_name) + + +ADAPTERS: dict[str, FamilyAdapter] = { + "generic": FamilyAdapter(), + "gemma3": Gemma3Adapter(), + "gemma4": Gemma4Adapter(), + "qwen": QwenAdapter(), + "lfm2": Lfm2Adapter(), + "moonshine": FamilyAdapter("moonshine"), + "nomic": NomicAdapter(), + "needle": FamilyAdapter("needle"), + "whisper": WhisperAdapter(), + "parakeet": ParakeetAdapter(), + "parakeet_tdt": ParakeetTDTAdapter(), +} + + +def adapter_for_family(family: str) -> FamilyAdapter: + return ADAPTERS.get(family, ADAPTERS["generic"]) + + +def _cfg_get(c, key, default=None): + if c is None: + return default + if isinstance(c, dict): + return c.get(key, default) + return getattr(c, key, default) + + +def _tensor_shape(tensor) -> tuple[int, ...]: + return tuple(int(x) for x in tensor.shape) + + +def _nomic_layer_suffix(name: str, suffix: str) -> int | None: + prefix = "encoder.layers." + if not name.startswith(prefix) or not name.endswith(suffix): + return None + middle = name[len(prefix) : -len(suffix)] + if not middle.isdigit(): + return None + return int(middle) + + +def _transpose_last_two(tensor): + if len(_tensor_shape(tensor)) < 2: + return tensor + if torch is not None and isinstance(tensor, torch.Tensor): + return tensor.transpose(-1, -2).contiguous() + return np.swapaxes(np.asarray(tensor), -1, -2).copy() + + +def split_channel_tensor( + tensor, + output_template: str, + channel_names: list[str] | None, + transform: str, + transpose: bool = False, +) -> list[TensorEmission]: + shape = _tensor_shape(tensor) + if len(shape) not in {1, 2}: + raise ValueError(f"channel split expects rank 1 or 2 tensor, got {shape}") + channels = len(channel_names) if channel_names is not None else None + if channels is None: + raise ValueError("channel_names are required unless an adapter overrides split behavior") + if shape[0] % channels != 0: + raise ValueError(f"first tensor dimension {shape[0]} is not divisible by {channels}") + if torch is not None and isinstance(tensor, torch.Tensor): + pieces = tensor.reshape(channels, -1, *shape[1:]) + else: + pieces = np.asarray(tensor).reshape(channels, -1, *shape[1:]) + emissions = [] + for idx, channel in enumerate(channel_names): + piece = pieces[idx] + if transpose: + piece = _transpose_last_two(piece) + out = output_template.replace("{channel}", str(channel)) + emissions.append(TensorEmission(out, piece, transform if not transpose else f"{transform}+transpose", qdq_restore="runtime_key")) + return emissions + + +def _load_audio_16k(path): + import soundfile as sf + + audio, sr = sf.read(path) + audio = np.asarray(audio, dtype=np.float32) + if audio.ndim > 1: + audio = audio.mean(axis=-1) + if sr != 16000: + old_x = np.linspace(0.0, 1.0, num=audio.shape[0], endpoint=False) + new_len = max(1, int(round(audio.shape[0] * 16000 / sr))) + new_x = np.linspace(0.0, 1.0, num=new_len, endpoint=False) + audio = np.interp(new_x, old_x, audio).astype(np.float32) + return audio diff --git a/python/cactus/convert/model_adapters/detection.py b/python/cactus/convert/model_adapters/detection.py new file mode 100644 index 000000000..69cb2aa09 --- /dev/null +++ b/python/cactus/convert/model_adapters/detection.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from ..cactus_adapters.config_utils import cfg_get, detect_model_type + + +SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "generic"} + + +def detect_family(config, requested: str = "auto") -> str: + if requested != "auto": + return requested + text_config = cfg_get(config, "text_config", None) + base = text_config if text_config is not None else config + detected = detect_model_type(base, config) + if detected == "gemma": + raw = str(cfg_get(base, "model_type", cfg_get(config, "model_type", "")) or "").lower() + if "gemma3" in raw: + return "gemma3" + return "generic" + if detected in {"gemma4", "qwen", "lfm2", "whisper", "moonshine", "parakeet", "parakeet_tdt", "nomic", "needle"}: + return detected + return "generic" diff --git a/python/cactus/convert/model_adapters/naming.py b/python/cactus/convert/model_adapters/naming.py new file mode 100644 index 000000000..3b66871af --- /dev/null +++ b/python/cactus/convert/model_adapters/naming.py @@ -0,0 +1,428 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Iterable + +from ..cactus_adapters import weight_patterns as wp + + +@dataclass(frozen=True) +class NameMatch: + source_name: str + output_name: str | None + component: str + recognized: bool + transpose: bool = False + hf_name: str | None = None + adapter_name: str | None = None + + +GEMMA4_WEIGHT_SCALE = 16.0 +GEMMA4_MULT_BASENAMES = {"ffn_gate", "ffn_up", "per_layer_gate", "moe_gate_proj", "moe_up_proj"} +GEMMA4_DIV_BASENAMES = { + "token_embeddings", + "output_weight", + "embed_vision_proj", + "embed_audio_proj", +} + + +def remap_gemma4_audio_key(key: str) -> str: + if "audio_tower" not in key: + return key + new_key = re.sub(r"subsample_conv_projection\.layer(\d+)\.", r"subsample_conv_projection.conv_\1.", key) + new_key = re.sub(r"audio_tower\.layers\.", "audio_tower.conformer.", new_key) + new_key = new_key.replace(".feed_forward1.", ".ffw_layer_start.") + new_key = new_key.replace(".feed_forward2.", ".ffw_layer_end.") + new_key = re.sub(r"\.self_attn\.(q_proj|k_proj|v_proj)\.", r".attention.attn.\1.", new_key) + new_key = new_key.replace(".self_attn.per_dim_scale", ".attention.attn.per_dim_scale") + new_key = new_key.replace(".self_attn.relative_k_proj.", ".attention.attn.relative_position_embedding.pos_proj.") + new_key = new_key.replace(".self_attn.post.", ".attention.post.") + new_key = new_key.replace(".norm_pre_attn.", ".attention.pre_attn_norm.") + new_key = new_key.replace(".norm_post_attn.", ".attention.post_norm.") + new_key = re.sub(r"\.norm_out\.", ".norm.", new_key) + return new_key + + +def restore_gemma4_audio_key(key: str) -> str: + if not key.startswith("model.audio_tower."): + return key + new_key = re.sub(r"subsample_conv_projection\.conv_(\d+)\.", r"subsample_conv_projection.layer\1.", key) + new_key = new_key.replace("model.audio_tower.conformer.", "model.audio_tower.layers.") + new_key = new_key.replace(".ffw_layer_start.", ".feed_forward1.") + new_key = new_key.replace(".ffw_layer_end.", ".feed_forward2.") + new_key = re.sub(r"\.attention\.attn\.([qkv])_proj\.", r".self_attn.\1_proj.", new_key) + new_key = new_key.replace(".attention.attn.per_dim_scale", ".self_attn.per_dim_scale") + new_key = new_key.replace(".attention.attn.relative_position_embedding.pos_proj.", ".self_attn.relative_k_proj.") + new_key = new_key.replace(".attention.post.", ".self_attn.post.") + new_key = new_key.replace(".attention.pre_attn_norm.", ".norm_pre_attn.") + new_key = new_key.replace(".attention.post_norm.", ".norm_post_attn.") + new_key = re.sub(r"^model\.audio_tower\.layers\.(\d+)\.norm\.weight$", r"model.audio_tower.layers.\1.norm_out.weight", new_key) + return new_key + + +def remap_whisper_key(key: str) -> str: + if key.startswith("model.encoder.") or key.startswith("model.decoder."): + return key[len("model."):] + return key + + +def adapter_key_for_family(key: str, family: str) -> str: + if family == "gemma4": + return remap_gemma4_audio_key(key) + if family == "whisper": + return remap_whisper_key(key) + return key + + +def restore_hf_key_for_family(key: str, family: str) -> str: + if family == "gemma4": + return restore_gemma4_audio_key(key) + return key + + +def _scale_basename(out_name: str) -> str: + base = out_name.removesuffix(".weights").removesuffix(".bias") + parts = base.split("_", 2) + if len(parts) == 3 and parts[0] == "layer" and parts[1].isdigit(): + base = parts[2] + return base + + +def gemma3_scale_factor(out_name: str) -> float: + base = _scale_basename(out_name) + if base in {"ffn_gate", "ffn_up"}: + return GEMMA4_WEIGHT_SCALE + if base == "token_embeddings": + return 1.0 / GEMMA4_WEIGHT_SCALE + return 1.0 + + +def gemma4_scale_factor(out_name: str) -> float: + base = _scale_basename(out_name) + if base in GEMMA4_MULT_BASENAMES: + return GEMMA4_WEIGHT_SCALE + if base in GEMMA4_DIV_BASENAMES: + return 1.0 / GEMMA4_WEIGHT_SCALE + if any(x in out_name for x in [ + "input_norm", + "post_attn_norm", + "pre_ffn_norm", + "post_ffn_norm", + "post_per_layer_norm", + "post_proj_norm", + ]): + return 1.0 / GEMMA4_WEIGHT_SCALE + if "router_scale" in out_name: + return 1.0 / GEMMA4_WEIGHT_SCALE + if out_name == "output_norm.weights": + return GEMMA4_WEIGHT_SCALE + return 1.0 + + +def remap_state_dict_for_family(state_dict: dict, family: str) -> dict: + if family != "gemma4": + return state_dict + return {remap_gemma4_audio_key(k): v for k, v in state_dict.items()} + + +def _tower_output_name(hf_key: str, strip_prefix: str, add_prefix: str) -> str: + name = hf_key[len(strip_prefix):] + if name.endswith(".weight"): + name = name[:-7] + ext = ".weights" + elif name.endswith(".bias"): + name = name[:-5] + ext = ".bias" + else: + ext = ".weights" + if name.endswith(".linear"): + name = name[:-7] + elif name.endswith("_linear"): + name = name[:-7] + return add_prefix + name.replace(".", "_") + ext + + +def _vision_layer_match(name: str) -> str | None: + prefixes = ( + "model.vision_tower.vision_model.encoder.layers.", + "model.vision_model.encoder.layers.", + ) + for prefix in prefixes: + m = re.match(rf"^{re.escape(prefix)}(\d+)\.", name) + if not m: + continue + idx = int(m.group(1)) + layer_prefix = f"{prefix}{idx}." + for source, out in wp.get_vision_layer_weights(idx, layer_prefix): + if name == source: + return out + return None + + +def _qwen_visual_match(name: str) -> str | None: + global_map = { + "model.visual.patch_embed.proj.weight": "vision_patch_embedding.weights", + "model.visual.patch_embed.proj.bias": "vision_patch_embedding.bias.weights", + "model.visual.pos_embed.weight": "vision_position_embedding.weights", + "model.visual.merger.norm.weight": "vision_merger_norm.weights", + "model.visual.merger.norm.bias": "vision_merger_norm.bias.weights", + "model.visual.merger.linear_fc1.weight": "vision_merger_linear_fc1.weights", + "model.visual.merger.linear_fc1.bias": "vision_merger_linear_fc1.bias.weights", + "model.visual.merger.linear_fc2.weight": "vision_merger_linear_fc2.weights", + "model.visual.merger.linear_fc2.bias": "vision_merger_linear_fc2.bias.weights", + } + if name in global_map: + return global_map[name] + + m = re.match(r"^model\.visual\.blocks\.(\d+)\.(.+)$", name) + if m: + idx = int(m.group(1)) + suffix = m.group(2) + suffix_map = { + "norm1.weight": "layer_norm1.weights", + "norm1.bias": "layer_norm1.bias.weights", + "norm2.weight": "layer_norm2.weights", + "norm2.bias": "layer_norm2.bias.weights", + "mlp.linear_fc1.weight": "ffn_fc1.weights", + "mlp.linear_fc1.bias": "ffn_fc1.bias.weights", + "mlp.linear_fc2.weight": "ffn_fc2.weights", + "mlp.linear_fc2.bias": "ffn_fc2.bias.weights", + "attn.qkv.weight": "self_attn_qkv.weights", + "attn.qkv.bias": "self_attn_qkv.bias.weights", + "attn.proj.weight": "self_attn_out.weights", + "attn.proj.bias": "self_attn_out.bias.weights", + } + out_suffix = suffix_map.get(suffix) + if out_suffix: + return f"vision_layer_{idx}_{out_suffix}" + + m = re.match(r"^model\.visual\.deepstack_merger_list\.(\d+)\.(.+)$", name) + if m: + idx = int(m.group(1)) + suffix = m.group(2) + suffix_map = { + "norm.weight": "norm.weights", + "norm.bias": "norm.bias.weights", + "linear_fc1.weight": "linear_fc1.weights", + "linear_fc1.bias": "linear_fc1.bias.weights", + "linear_fc2.weight": "linear_fc2.weights", + "linear_fc2.bias": "linear_fc2.bias.weights", + } + out_suffix = suffix_map.get(suffix) + if out_suffix: + return f"vision_deepstack_merger_{idx}_{out_suffix}" + return None + + +def _candidate_table_match(name: str, table: Iterable) -> str | None: + for item in table: + candidates, out_name = item[0], item[1] + if isinstance(candidates, str): + candidates = (candidates,) + if name in candidates: + return out_name + return None + + +def _whisper_layer_match(name: str, i: int) -> tuple[str, bool] | None: + for block in ("encoder", "decoder"): + prefix = f"{block}.layers.{i}." + if not name.startswith(prefix): + continue + suffix = name[len(prefix):] + for patterns, _precision, out_name, transpose in wp.get_layer_weight_patterns(i, "CQ", model_type="whisper"): + for pat in patterns: + if suffix == pat: + return f"{block}.{out_name}", transpose + return None + + +def _parakeet_layer_match(name: str, i: int) -> str | None: + prefix = f"encoder.layers.{i}." + if not name.startswith(prefix): + return None + suffix = name[len(prefix):] + out_suffix = _candidate_table_match(suffix, wp.PARAKEET_LAYER_WEIGHTS) + return f"layer_{i}_{out_suffix}" if out_suffix else None + + +def _parakeet_tdt_predictor_match(name: str) -> str | None: + old = re.fullmatch(r"decoder\.prediction\.dec_rnn\.lstm\.(\d+)\.(Wx|Wh|bias)", name) + if old: + suffix = {"Wx": "weight_ih.weights", "Wh": "weight_hh.weights", "bias": "bias.weights"}[old.group(2)] + return f"tdt_predictor_lstm_{int(old.group(1))}_{suffix}" + new = re.fullmatch(r"decoder\.lstm\.(weight_ih|weight_hh|bias)_l(\d+)", name) + if new: + suffix = {"weight_ih": "weight_ih.weights", "weight_hh": "weight_hh.weights", "bias": "bias.weights"}[new.group(1)] + return f"tdt_predictor_lstm_{int(new.group(2))}_{suffix}" + nemo = re.fullmatch(r"decoder\.prediction\.dec_rnn\.lstm\.(weight_ih|weight_hh|bias)_l(\d+)", name) + if nemo: + suffix = {"weight_ih": "weight_ih.weights", "weight_hh": "weight_hh.weights", "bias": "bias.weights"}[nemo.group(1)] + return f"tdt_predictor_lstm_{int(nemo.group(2))}_{suffix}" + return None + + +def _needle_gate_match(name: str) -> str | None: + m = re.fullmatch(r"model\.encoder\.layers\.(\d+)\.attn_gate", name) + if m: + return f"encoder_layer_{int(m.group(1))}_attn_gate.weights" + m = re.fullmatch(r"model\.decoder\.layers\.(\d+)\.(self_attn_gate|cross_attn_gate)", name) + if m: + return f"layer_{int(m.group(1))}_{m.group(2)}.weights" + return None + + +def _suffix_layer_match(name: str, i: int, family: str) -> tuple[str, bool] | None: + if family == "whisper": + return _whisper_layer_match(name, i) + if family in {"parakeet", "parakeet_tdt"}: + out = _parakeet_layer_match(name, i) + if out: + return out, False + if family == "needle": + out = _needle_gate_match(name) + if out: + return out, False + prefixes = [p.format(i=i) for p in wp.LAYER_PREFIXES] + [f"h.{i}."] + for p in prefixes: + if not name.startswith(p): + continue + suffix = name[len(p):] + for patterns, _precision, out_name, transpose in wp.get_layer_weight_patterns(i, "CQ", model_type=family): + for pat in patterns: + if "{channel}" in pat: + rx = re.escape(pat).replace(r"\{channel\}", r"(\d+)") + m = re.fullmatch(rx, suffix) + if m: + matched_out = out_name.format(channel=m.group(1)) + if family == "needle" and p == f"model.encoder.layers.{i}.": + matched_out = matched_out.replace(f"layer_{i}_", f"encoder_layer_{i}_", 1) + return matched_out, transpose + elif suffix == pat: + matched_out = out_name + if family == "needle" and p == f"model.encoder.layers.{i}.": + matched_out = matched_out.replace(f"layer_{i}_", f"encoder_layer_{i}_", 1) + return matched_out, transpose + if suffix.endswith(".weight"): + prefix = f"encoder_layer_{i}_" if family == "needle" and p == f"model.encoder.layers.{i}." else f"layer_{i}_" + return f"{prefix}{suffix[:-7].replace('.', '_')}.weights", False + if suffix.endswith(".bias"): + prefix = f"encoder_layer_{i}_" if family == "needle" and p == f"model.encoder.layers.{i}." else f"layer_{i}_" + return f"{prefix}{suffix[:-5].replace('.', '_')}.bias", False + return None + + +def _global_match(name: str, family: str) -> str | None: + if family == "whisper": + found = _candidate_table_match(name, wp.WHISPER_GLOBAL_WEIGHTS) + if found: + return found + elif family == "parakeet": + found = _candidate_table_match(name, wp.PARAKEET_GLOBAL_WEIGHTS) + if found: + return found + elif family == "parakeet_tdt": + for table in (wp.PARAKEET_GLOBAL_WEIGHTS, wp.PARAKEET_TDT_GLOBAL_WEIGHTS): + found = _candidate_table_match(name, table) + if found: + return found + elif family == "needle": + found = _candidate_table_match(name, wp.NEEDLE_GLOBAL_WEIGHTS) + if found: + return found + if name in {"wte.weight", "word_embeddings.weight"}: + return "token_embeddings.weights" + if name in {"wpe.weight", "position_embeddings.weight"}: + return "position_embeddings.weights" + if name == "ln_f.weight": + return "output_norm.weights" + if name == "ln_f.bias": + return "output_norm.bias" + if name in wp.EMBED_NAMES: + return "token_embeddings.weights" + if name in wp.OUTPUT_NAMES: + return "output_weight.weights" + if name in wp.OUTPUT_NORM_NAMES: + return "output_norm.weights" + tables: list[Iterable] = [wp.VISION_ITEMS, wp.PROJECTOR_WEIGHTS] + if family == "gemma4": + tables.append(wp.GEMMA4_GLOBAL_WEIGHTS) + elif family == "gemma3n": + tables.append(wp.GEMMA3N_GLOBAL_WEIGHTS) + elif family == "moonshine": + tables.append(wp.MOONSHINE_GLOBAL_WEIGHTS) + for table in tables: + found = _candidate_table_match(name, table) + if found: + return found + if family == "parakeet_tdt": + predictor_name = _parakeet_tdt_predictor_match(name) + if predictor_name: + return predictor_name + return None + + +def component_for_name(name: str, output_name: str | None = None) -> str: + joined = f"{name} {output_name or ''}".lower() + if "parakeet" in joined or "tdt_" in joined or "ctc_" in joined: + return "transcription" + if "audio" in joined or "encoder.conv" in joined or "subsample" in joined: + return "audio" + if "vision" in joined or "visual" in joined or "image" in joined or "projector" in joined: + return "vision" + if "embed" in joined or "token_embedding" in joined or "lm_head" in joined or "output_weight" in joined: + return "embedding" + if "decoder" in joined and ("whisper" in joined or "moonshine" in joined): + return "transcription" + return "language" + + +def _component_for_family(name: str, output_name: str | None, family: str) -> str: + if family in {"parakeet", "parakeet_tdt"}: + return "transcription" + if family == "whisper": + out = output_name or "" + if out in {"decoder_token_embeddings.weights", "output_weight.weights"}: + return "embedding" + return "transcription" + return component_for_name(name, output_name) + + +def cactus_name_for_tensor(name: str, family: str, num_layers: int | None = None) -> NameMatch: + hf_name = name + adapter_name = adapter_key_for_family(name, family) + if family in {"parakeet", "parakeet_tdt"} and adapter_name.endswith(".conv.norm.num_batches_tracked"): + return NameMatch(name, None, "transcription", True, hf_name=hf_name, adapter_name=adapter_name) + global_name = _global_match(adapter_name, family) + if global_name: + return NameMatch(name, global_name, _component_for_family(hf_name, global_name, family), True, hf_name=hf_name, adapter_name=adapter_name) + vision_name = _vision_layer_match(adapter_name) + if vision_name: + return NameMatch(name, vision_name, "vision", True, hf_name=hf_name, adapter_name=adapter_name) + if family == "qwen": + qwen_vision_name = _qwen_visual_match(adapter_name) + if qwen_vision_name: + return NameMatch(name, qwen_vision_name, "vision", True, hf_name=hf_name, adapter_name=adapter_name) + if family == "gemma4": + if adapter_name.startswith(wp.GEMMA4_VISION_TOWER_PREFIX): + out = _tower_output_name(adapter_name, wp.GEMMA4_VISION_TOWER_PREFIX, "vision_") + return NameMatch(name, out, "vision", True, hf_name=hf_name, adapter_name=adapter_name) + if adapter_name.startswith(wp.GEMMA4_AUDIO_TOWER_PREFIX): + out = _tower_output_name(adapter_name, wp.GEMMA4_AUDIO_TOWER_PREFIX, "audio_") + return NameMatch(name, out, "audio", True, hf_name=hf_name, adapter_name=adapter_name) + max_layers = int(num_layers or 160) + for i in range(max_layers): + layer = _suffix_layer_match(adapter_name, i, family) + if layer: + out, transpose = layer + return NameMatch(name, out, _component_for_family(hf_name, out, family), True, transpose, hf_name=hf_name, adapter_name=adapter_name) + if adapter_name.endswith(".weight") or adapter_name.endswith(".bias"): + generic = adapter_name.replace(".", "_") + if generic.endswith("_weight"): + generic = generic[:-7] + ".weights" + elif generic.endswith("_bias"): + generic = generic[:-5] + ".bias" + return NameMatch(name, generic, component_for_name(hf_name, generic), False, hf_name=hf_name, adapter_name=adapter_name) + return NameMatch(name, None, component_for_name(hf_name), False, hf_name=hf_name, adapter_name=adapter_name) diff --git a/python/cactus/convert/model_adapters/nemo.py b/python/cactus/convert/model_adapters/nemo.py new file mode 100644 index 000000000..af828a65c --- /dev/null +++ b/python/cactus/convert/model_adapters/nemo.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import json +import shutil +import tarfile +from pathlib import Path +from typing import Any + + +_NEMO_EXPORT_FILES = { + "model_config.yaml", + "model_weights.ckpt", + "vocab.txt", +} + + +def ensure_parakeet_tdt_nemo_source( + model_id_or_path: str, + *, + token: str | None = None, + cache_dir: str | None = None, +) -> str | None: + nemo_path = _find_single_nemo(model_id_or_path, token=token, cache_dir=cache_dir) + if nemo_path is None: + return None + + out = nemo_path.parent / f"{nemo_path.stem}-cactus-hf" + if (out / "config.json").exists() and (out / "pytorch_model.bin").exists() and (out / "vocab.txt").exists(): + return str(out) + + tmp = out.with_name(f"{out.name}.tmp") + if tmp.exists(): + shutil.rmtree(tmp) + tmp.mkdir(parents=True) + + with tarfile.open(nemo_path, "r:*") as archive: + for member in archive.getmembers(): + name = Path(member.name).name + if not _keep_nemo_member(name): + continue + extracted = archive.extractfile(member) + if extracted is None: + continue + with extracted, (tmp / name).open("wb") as dst: + shutil.copyfileobj(extracted, dst, length=1024 * 1024) + + weights = tmp / "model_weights.ckpt" + config = _read_yaml(tmp / "model_config.yaml") + if not weights.exists() or not config: + raise RuntimeError(f"{nemo_path} is missing model_config.yaml or model_weights.ckpt") + + hf_config = _parakeet_tdt_config_from_nemo(config) + (tmp / "config.json").write_text(json.dumps(hf_config, indent=2, sort_keys=True), encoding="utf-8") + weights.rename(tmp / "pytorch_model.bin") + _write_parakeet_tdt_tokenizer_files(tmp, hf_config) + + if out.exists(): + shutil.rmtree(out) + tmp.rename(out) + return str(out) + + +def _find_single_nemo(model_id_or_path: str, *, token: str | None, cache_dir: str | None) -> Path | None: + path = Path(model_id_or_path) + if path.suffix == ".nemo" and path.is_file(): + return path + if path.is_dir(): + files = sorted(path.glob("*.nemo")) + return files[0] if len(files) == 1 else None + + try: + from huggingface_hub import hf_hub_download, list_repo_files + + files = [name for name in list_repo_files(model_id_or_path, token=token) if name.endswith(".nemo")] + if len(files) != 1: + return None + return Path(hf_hub_download(model_id_or_path, files[0], token=token, cache_dir=cache_dir)) + except Exception: + return None + + +def _keep_nemo_member(name: str) -> bool: + return ( + name in _NEMO_EXPORT_FILES + or name.endswith("_vocab.txt") + or name.endswith("_tokenizer.model") + or name.endswith("_tokenizer.vocab") + ) + + +def _read_yaml(path: Path) -> dict[str, Any]: + try: + import yaml + except Exception as exc: + raise RuntimeError("PyYAML is required to import Parakeet TDT .nemo checkpoints") from exc + if not path.exists(): + return {} + loaded = yaml.safe_load(path.read_text(encoding="utf-8")) + return loaded if isinstance(loaded, dict) else {} + + +def _get(mapping: dict[str, Any], key: str, default: Any = None) -> Any: + value = mapping.get(key, default) + return default if value is None else value + + +def _parakeet_tdt_config_from_nemo(root: dict[str, Any]) -> dict[str, Any]: + decoder = _get(root, "decoder", {}) or {} + prednet = _get(decoder, "prednet", _get(decoder, "prediction", {})) or {} + joint = _get(root, "joint", {}) or {} + labels = [str(token) for token in (_get(joint, "vocabulary", []) or [])] + + decoder_vocab = int(_get(decoder, "vocab_size", len(labels))) + + config = dict(root) + config.update(architectures=["ParakeetForTDT"], model_type="parakeet_tdt", labels=labels) + config["decoder"] = {**decoder, "prediction": dict(prednet), "vocab_size": decoder_vocab} + config["joint"] = {**joint, "vocabulary": labels} + return config + + +def _write_parakeet_tdt_tokenizer_files(root: Path, config: dict[str, Any]) -> None: + tokenizer_model = next(iter(sorted(root.glob("*_tokenizer.model"))), None) + if tokenizer_model is not None and not (root / "tokenizer.model").exists(): + tokenizer_model.rename(root / "tokenizer.model") + + labels = config.get("labels") + if isinstance(labels, list) and labels: + with (root / "vocab.txt").open("w", encoding="utf-8") as f: + for idx, token in enumerate(labels): + f.write(f"{idx}\t{token}\n") + else: + nemo_vocab = next(iter(sorted(root.glob("*_vocab.txt"))), None) + if nemo_vocab is not None and not (root / "vocab.txt").exists(): + nemo_vocab.rename(root / "vocab.txt") + + (root / "tokenizer_config.json").write_text( + json.dumps({"model_type": "parakeet_tdt", "tokenizer_class": "SentencePieceProcessor"}, indent=2), + encoding="utf-8", + ) diff --git a/python/cactus/convert/model_adapters/policy.py b/python/cactus/convert/model_adapters/policy.py new file mode 100644 index 000000000..578218aef --- /dev/null +++ b/python/cactus/convert/model_adapters/policy.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .naming import NameMatch + + +@dataclass(frozen=True) +class TensorPolicy: + action: str + precision: str + bits: int | None + component: str + use_gptq: bool + rotation: str + fallback_reason: str | None = None + layout: str = "row_major" + + +def policy_for_tensor(match: NameMatch, shape: tuple[int, ...], user_bits: int, family: str) -> TensorPolicy: + component = match.component + name = match.source_name + out = match.output_name or "" + if match.output_name is None: + return TensorPolicy("ignored", "none", None, component, False, "none", "no output filename") + if "norm" in out.lower(): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "norm tensor") + if family == "gemma4" and name == "model.embed_vision.embedding_projection.weight": + return TensorPolicy("fallback", "FP16", None, component, False, "none", "vision embedding projection scale-sensitive") + if family in {"parakeet", "parakeet_tdt"} and out.endswith(".bias") and ( + "conv_" in out or out.startswith("subsampling_") or out.startswith("ctc_head_") + ): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "conv bias tensor") + if family in {"parakeet", "parakeet_tdt"} and "lstm" in out.lower(): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "lstm recurrent tensor") + if family == "parakeet_tdt" and out.startswith("tdt_"): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "tdt decoder tensor") + if family in {"parakeet", "parakeet_tdt"} and "self_attn_bias_" in out: + return TensorPolicy("fallback", "FP16", None, component, False, "none", "relative attention bias tensor") + if family in {"parakeet", "parakeet_tdt"} and "conv_pointwise" in out and len(shape) == 3 and shape[2] == 1: + return TensorPolicy("fallback", "INT8", 8, component, False, "none", "pointwise conv tensor") + if "conv_depthwise.weights" in out and len(shape) == 3 and shape[1] == 1: + return TensorPolicy("fallback", "INT8", 8, component, False, "none", "depthwise conv tensor") + if family in {"parakeet", "parakeet_tdt"} and out.startswith("layer_") and ( + "conv_pointwise" in out or "conv_depthwise" in out + ): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "conformer conv tensor") + if family == "whisper" and out.startswith("encoder.layer_"): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "whisper encoder tensor") + if family == "gemma4" and component == "audio" and (name.endswith(".bias") or out.endswith(".bias") or ".bias." in out): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "audio bias tensor") + if name.endswith(".bias") or out.endswith(".bias") or ".bias." in out: + return TensorPolicy("fallback", "FP16", None, component, False, "none", "bias tensor") + if len(shape) != 2: + return TensorPolicy("fallback", "FP16", None, component, False, "none", "non-2d tensor") + if "position_embedding" in out.lower() or "pos_embed" in out.lower() or "embed_positions" in out.lower(): + return TensorPolicy("fallback", "FP16", None, component, False, "none", "position embedding tensor") + if family == "gemma4" and name == "model.language_model.embed_tokens_per_layer.weight": + return TensorPolicy("convert", f"CQ{user_bits}", user_bits, component, False, "hadamard") + if family == "gemma4" and component in {"audio", "vision"}: + return TensorPolicy("fallback", "FP16", None, component, False, "none", "gemma4 media tower accuracy") + output_head_or_tied_embedding = out in {"token_embeddings.weights", "decoder_token_embeddings.weights", "output_weight.weights"} + if component == "embedding" or output_head_or_tied_embedding: + use_interleaved = ( + output_head_or_tied_embedding + and len(shape) == 2 + and int(shape[0]) % 4 == 0 + and int(shape[1]) % 32 == 0 + ) + layout = "interleaved_4row" if use_interleaved else "row_major" + return TensorPolicy("convert", "CQ4", 4, component, False, "orthogonal", layout=layout) + if component == "audio" or component == "transcription": + return TensorPolicy("convert", f"CQ{user_bits}", user_bits, component, False, "hadamard") + return TensorPolicy("convert", f"CQ{user_bits}", user_bits, component, True, "hadamard") diff --git a/python/cactus/convert/quantization/__init__.py b/python/cactus/convert/quantization/__init__.py new file mode 100644 index 000000000..6f30b7f5b --- /dev/null +++ b/python/cactus/convert/quantization/__init__.py @@ -0,0 +1,2 @@ +"""CQ quantization and writer implementations.""" + diff --git a/python/cactus/convert/quantization/cq.py b/python/cactus/convert/quantization/cq.py new file mode 100644 index 000000000..2d3942389 --- /dev/null +++ b/python/cactus/convert/quantization/cq.py @@ -0,0 +1,410 @@ +from __future__ import annotations + +import math +import struct +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +try: + import torch +except Exception: # pragma: no cover + torch = None + +from ..cactus_adapters.tensor_io import ( + CACTUS_ALIGNMENT, + CACTUS_MAGIC, + align_offset, + compute_padding, +) + +HEADER_SIZE = 84 +GROUP_SIZE = 128 +PRECISION_CQ = {1: 3, 2: 4, 3: 5, 4: 6} +FLAG_ORTHOGONAL_ROTATION = 1 << 1 +FLAG_INTERLEAVED_4ROW = 1 << 2 + + +@dataclass(frozen=True) +class CQTensor: + indices: np.ndarray + norms: np.ndarray + input_scale: np.ndarray + bits: int + group_size: int = GROUP_SIZE + rotation_family: str = "hadamard" + seed: int = 1234 + gptq_used: bool = False + interleaved_4row: bool = False + + +_CODEBOOK_CACHE: dict[tuple[int, int], np.ndarray] = {} +_HADAMARD_CACHE: dict[tuple[int, int], tuple[np.ndarray, np.ndarray, np.ndarray]] = {} +_ORTHO_CACHE: dict[tuple[int, int], np.ndarray] = {} + + +def make_codebook(group_dim: int, bits: int, grid_size: int = 200_001) -> np.ndarray: + key = (int(group_dim), int(bits)) + if key in _CODEBOOK_CACHE: + return _CODEBOOK_CACHE[key] + try: + from scipy.special import gammaln + except Exception as exc: # pragma: no cover + raise RuntimeError("scipy is required for CQ codebook generation") from exc + n_centroids = 1 << bits + grid = np.linspace(-1.0, 1.0, grid_size, dtype=np.float64) + log_c = gammaln(group_dim / 2.0) - 0.5 * math.log(math.pi) - gammaln((group_dim - 1.0) / 2.0) + weights = math.exp(log_c) * np.power(np.clip(1.0 - grid * grid, 0.0, None), (group_dim - 3.0) / 2.0) + quantiles = np.linspace(0.0, 1.0, n_centroids + 2, dtype=np.float64)[1:-1] + cum = np.cumsum(weights) + cum /= cum[-1] + centroids = np.sort(np.clip(np.interp(quantiles, cum, grid), -1.0, 1.0)) + for _ in range(200): + bounds = np.concatenate(([-1.0], (centroids[:-1] + centroids[1:]) / 2.0, [1.0])) + updated = centroids.copy() + for i in range(n_centroids): + mask = (grid >= bounds[i]) & (grid <= bounds[i + 1] if i == n_centroids - 1 else grid < bounds[i + 1]) + ws = weights[mask] + if ws.sum() > 0: + updated[i] = float((grid[mask] * ws).sum() / ws.sum()) + if np.max(np.abs(updated - centroids)) < 1e-8: + centroids = updated + break + centroids = np.sort(np.clip(updated, -1.0, 1.0)) + _CODEBOOK_CACHE[key] = centroids.astype(np.float32) + return _CODEBOOK_CACHE[key] + + +def make_hadamard_components(group_dim: int, seed: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + key = (int(group_dim), int(seed)) + if key in _HADAMARD_CACHE: + return _HADAMARD_CACHE[key] + if torch is None: # pragma: no cover + raise RuntimeError("torch is required for deterministic CQ rotations") + g = torch.Generator(device="cpu").manual_seed(seed + 17 * group_dim) + left = (2 * torch.randint(0, 2, (group_dim,), generator=g, dtype=torch.int64) - 1).to(torch.int8).numpy() + right = (2 * torch.randint(0, 2, (group_dim,), generator=g, dtype=torch.int64) - 1).to(torch.int8).numpy() + perm = torch.randperm(group_dim, generator=g).to(torch.int64).numpy().astype(np.uint32) + _HADAMARD_CACHE[key] = (left, right, perm) + return _HADAMARD_CACHE[key] + + +def make_hadamard_matrix(group_dim: int, seed: int) -> np.ndarray: + try: + from scipy.linalg import hadamard + except Exception as exc: # pragma: no cover + raise RuntimeError("scipy is required for Hadamard CQ rotation") from exc + left, right, perm = make_hadamard_components(group_dim, seed) + base = hadamard(group_dim, dtype=np.float32) / math.sqrt(group_dim) + return (left.astype(np.float32)[:, None] * base * right.astype(np.float32)[None, :])[:, perm] + + +def make_orthogonal_rotation(group_dim: int, seed: int) -> np.ndarray: + key = (int(group_dim), int(seed)) + if key in _ORTHO_CACHE: + return _ORTHO_CACHE[key] + if torch is None: # pragma: no cover + raise RuntimeError("torch is required for deterministic orthogonal CQ rotations") + g = torch.Generator(device="cpu").manual_seed(seed + 17 * group_dim) + a = torch.randn(group_dim, group_dim, generator=g) + q, r = torch.linalg.qr(a, mode="reduced") + d = torch.sign(torch.diagonal(r)) + d[d == 0] = 1 + _ORTHO_CACHE[key] = (q * d).numpy().astype(np.float32) + return _ORTHO_CACHE[key] + + +def pack_indices_lsb(indices: np.ndarray, group_size: int, bits: int) -> np.ndarray: + n, k = indices.shape + if bits == 8: + return indices.astype(np.uint8, copy=False) + if not 1 <= bits <= 4: + raise ValueError(f"bits must be 1..4, got {bits}") + if k % group_size != 0: + raise ValueError(f"K={k} is not divisible by group_size={group_size}") + chunk = 8 + bytes_per_chunk = chunk * bits // 8 + chunks_per_group = group_size // chunk + grouped = indices.reshape(n, k // group_size, chunks_per_group, chunk).astype(np.uint64) + word = np.zeros(grouped.shape[:-1], dtype=np.uint64) + for i in range(chunk): + word |= grouped[..., i] << (i * bits) + out = np.zeros(word.shape + (bytes_per_chunk,), dtype=np.uint8) + for byte in range(bytes_per_chunk): + out[..., byte] = ((word >> (8 * byte)) & 0xFF).astype(np.uint8) + return out.reshape(n, k * bits // 8) + + +def pack_indices_interleaved_4row_4bit(indices: np.ndarray, group_size: int) -> np.ndarray: + n, k = indices.shape + if n % 4 != 0 or group_size % 8 != 0 or k % group_size != 0: + raise ValueError(f"interleaved-4row-4bit constraints not met: n={n}, gs={group_size}, k={k}") + nb = n // 4 + ng = k // group_size + chunks = group_size // 8 + g = indices.reshape(nb, 4, ng, chunks, 8).astype(np.uint8) + lo = g[..., 0:4] + hi = g[..., 4:8] + bytes_arr = ((lo & 0x0F) | ((hi & 0x0F) << 4)).astype(np.uint8) + bytes_arr = np.transpose(bytes_arr, (0, 2, 3, 1, 4)).copy() + return bytes_arr.reshape(-1) + + +def pack_indices_interleaved_4row_2bit(indices: np.ndarray, group_size: int) -> np.ndarray: + n, k = indices.shape + if n % 4 != 0 or group_size % 16 != 0 or k % group_size != 0: + raise ValueError(f"interleaved-4row-2bit constraints not met: n={n}, gs={group_size}, k={k}") + nb = n // 4 + ng = k // group_size + chunks = group_size // 16 + g = indices.reshape(nb, 4, ng, chunks, 4, 4).astype(np.uint8) + packed = ((g[..., 0] & 0x3) + | ((g[..., 1] & 0x3) << 2) + | ((g[..., 2] & 0x3) << 4) + | ((g[..., 3] & 0x3) << 6)).astype(np.uint8) + packed = np.transpose(packed, (0, 2, 3, 4, 1)).copy() + return packed.reshape(-1) + + +def pack_indices_interleaved_4row_1bit(indices: np.ndarray, group_size: int) -> np.ndarray: + n, k = indices.shape + if n % 4 != 0 or group_size % 32 != 0 or k % group_size != 0: + raise ValueError(f"interleaved-4row-1bit constraints not met: n={n}, gs={group_size}, k={k}") + nb = n // 4 + ng = k // group_size + chunks = group_size // 32 + g = indices.reshape(nb, 4, ng, chunks, 4, 8).astype(np.uint8) + packed = np.zeros((nb, 4, ng, chunks, 4), dtype=np.uint8) + for bit in range(8): + packed |= ((g[..., bit] & 0x1) << bit).astype(np.uint8) + packed = np.transpose(packed, (0, 2, 3, 4, 1)).copy() + return packed.reshape(-1) + + +def pack_indices_interleaved_4row_3bit(indices: np.ndarray, group_size: int) -> np.ndarray: + n, k = indices.shape + if n % 4 != 0 or group_size % 4 != 0 or k % group_size != 0: + raise ValueError(f"interleaved-4row-3bit constraints not met: n={n}, gs={group_size}, k={k}") + nb = n // 4 + ng = k // group_size + chunks_per_group = group_size // 4 + g = indices.reshape(nb, 4, ng, chunks_per_group, 4).astype(np.uint64) & 0x7 + g = np.transpose(g, (0, 2, 3, 1, 4)) + g = g.reshape(nb, ng, chunks_per_group, 16) + word = np.zeros((nb, ng, chunks_per_group), dtype=np.uint64) + for i in range(16): + word |= (g[..., i] & 0x7) << (i * 3) + out = np.zeros((nb, ng, chunks_per_group, 6), dtype=np.uint8) + for b in range(6): + out[..., b] = ((word >> (b * 8)) & 0xFF).astype(np.uint8) + return out.reshape(-1) + + +def pack_indices_interleaved_4row(indices: np.ndarray, group_size: int, bits: int) -> np.ndarray: + if bits == 4: + return pack_indices_interleaved_4row_4bit(indices, group_size) + if bits == 3: + return pack_indices_interleaved_4row_3bit(indices, group_size) + if bits == 2: + return pack_indices_interleaved_4row_2bit(indices, group_size) + if bits == 1: + return pack_indices_interleaved_4row_1bit(indices, group_size) + raise ValueError(f"interleaved-4row not supported for {bits}-bit quantization") + + +def _as_numpy_fp32(weight) -> np.ndarray: + if torch is not None and isinstance(weight, torch.Tensor): + return weight.detach().float().cpu().numpy() + return np.asarray(weight, dtype=np.float32) + + +def _prepare_input_scale(input_scale: np.ndarray | None, k: int) -> np.ndarray: + if input_scale is None: + return np.ones(k, dtype=np.float32) + scale = np.asarray(input_scale, dtype=np.float32).reshape(-1) + if scale.shape[0] > k: + scale = scale[:k] + elif scale.shape[0] < k: + scale = np.pad(scale, (0, k - scale.shape[0]), constant_values=1.0) + return np.clip(scale, 1e-8, 65504.0).astype(np.float32, copy=False) + + +def _gptq_correct_group(work: np.ndarray, recon: np.ndarray, h_inv: np.ndarray | None, start: int, stop: int) -> None: + if h_inv is None or stop >= work.shape[1]: + return + try: + m_bb = h_inv[start:stop, start:stop].astype(np.float32, copy=False) + m_bs = h_inv[start:stop, stop:].astype(np.float32, copy=False) + update = np.linalg.solve(m_bb + np.eye(m_bb.shape[0], dtype=np.float32) * 1e-6, m_bs) + work[:, stop:] -= (work[:, start:stop] - recon) @ update + except Exception: + return + + +def quantize_hadamard( + weight, + bits: int, + hessian: np.ndarray | None = None, + use_gptq: bool = False, + seed: int = 1234, + input_scale: np.ndarray | None = None, +) -> CQTensor: + w = _as_numpy_fp32(weight) + if w.ndim != 2: + raise ValueError(f"hadamard CQ requires a 2D tensor, got {w.shape}") + n, k = w.shape + original_k = k + if k % GROUP_SIZE != 0: + pad = GROUP_SIZE - (k % GROUP_SIZE) + w = np.pad(w, ((0, 0), (0, pad)), mode="constant") + k = w.shape[1] + input_scale = _prepare_input_scale(input_scale, k) + work = w * input_scale[None, :] + rot = make_hadamard_matrix(GROUP_SIZE, seed) + codebook = make_codebook(GROUP_SIZE, bits) + groups = k // GROUP_SIZE + indices = np.zeros((n, k), dtype=np.uint8) + norms = np.zeros((n, groups), dtype=np.float16) + h_inv = None + if use_gptq and hessian is not None: + try: + h = np.asarray(hessian, dtype=np.float32) + if h.shape[0] == original_k and h.shape[0] == h.shape[1]: + if original_k != k: + h = np.pad(h, ((0, k - original_k), (0, k - original_k)), mode="constant") + np.fill_diagonal(h[original_k:, original_k:], 1.0) + s = np.clip(input_scale.astype(np.float32), 1e-6, None) + h = h / (s[:, None] * s[None, :]) + h = h + np.eye(h.shape[0], dtype=np.float32) * (0.01 * np.mean(np.diag(h)) + 1e-6) + h_inv = np.linalg.inv(h) + except Exception: + h_inv = None + for g in range(groups): + start, stop = g * GROUP_SIZE, (g + 1) * GROUP_SIZE + group = work[:, start:stop] + row_norms = np.linalg.norm(group, axis=1).clip(min=1e-8).astype(np.float32) + rotated = (group / row_norms[:, None]) @ rot + idx = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) + recon = (codebook[idx] @ rot.T) * row_norms[:, None] + indices[:, start:stop] = idx + norms[:, g] = row_norms.astype(np.float16) + _gptq_correct_group(work, recon, h_inv, start, stop) + return CQTensor(indices=indices, norms=norms, input_scale=input_scale.astype(np.float16), bits=bits, gptq_used=bool(use_gptq and h_inv is not None)) + + +def quantize_orthogonal(weight, bits: int = 4, seed: int = 1234, input_scale: np.ndarray | None = None) -> CQTensor: + w = _as_numpy_fp32(weight) + if w.ndim != 2: + raise ValueError(f"orthogonal CQ requires a 2D tensor, got {w.shape}") + n, k = w.shape + input_scale = _prepare_input_scale(input_scale, k) + work = w * input_scale[None, :] + rot = make_orthogonal_rotation(k, seed) + codebook = make_codebook(k, bits) + norms = np.linalg.norm(work, axis=1).clip(min=1e-8).astype(np.float32) + indices = np.zeros((n, k), dtype=np.uint8) + for start in range(0, n, 1024): + stop = min(start + 1024, n) + rotated = (work[start:stop] / norms[start:stop, None]) @ rot + indices[start:stop] = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) + return CQTensor(indices=indices, norms=norms.astype(np.float16).reshape(n, 1), input_scale=input_scale.astype(np.float16), bits=bits, group_size=k, rotation_family="orthogonal") + + +_EMBEDDING_TENSOR_STEMS = frozenset({ + "token_embeddings", + "embed_tokens", + "embed_tokens_per_layer", +}) + + +def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: + out_path.parent.mkdir(parents=True, exist_ok=True) + n, k = cq.indices.shape + group_size = int(cq.group_size) + groups = k // group_size + + is_embedding = out_path.stem in _EMBEDDING_TENSOR_STEMS + + if cq.interleaved_4row: + if cq.bits not in (1, 2, 3, 4): + raise ValueError(f"INTERLEAVED_4ROW only supports CQ1..CQ4, got CQ{cq.bits}") + if n % 4 != 0: + raise ValueError(f"INTERLEAVED_4ROW requires N % 4 == 0, got N={n}") + if cq.rotation_family == "orthogonal": + if cq.bits != 4 or group_size != k: + raise ValueError("orthogonal INTERLEAVED_4ROW requires full-width CQ4") + if k % 32 != 0: + raise ValueError(f"orthogonal INTERLEAVED_4ROW requires K % 32 == 0, got K={k}") + elif group_size % 32 != 0 or group_size > 256: + raise ValueError(f"INTERLEAVED_4ROW requires group_size % 32 == 0 and group_size <= 256, got {group_size}") + + interleaved = ( + cq.interleaved_4row + or ( + cq.bits in (1, 2, 3, 4) + and cq.rotation_family != "orthogonal" + and n % 4 == 0 + and group_size % 32 == 0 + and group_size <= 256 + and not is_embedding + ) + ) + + codebook_f32 = make_codebook(group_size, cq.bits).astype(np.float32) + norms_f32 = cq.norms.astype(np.float32, copy=True) + if interleaved and cq.rotation_family != "orthogonal": + cb_max = float(np.max(np.abs(codebook_f32))) + cb_factor = cb_max / 127.0 if cb_max > 1e-12 else 1.0 / 127.0 + codebook_f32 = codebook_f32 / cb_factor + norms_f32 = norms_f32 * cb_factor + + codebook = codebook_f32.astype(np.float16) + input_scale = cq.input_scale.astype(np.float16, copy=False).reshape(k) + recip = np.minimum(1.0 / np.maximum(input_scale.astype(np.float32), 1e-8), 65504.0).astype(np.float16) + + if interleaved: + norms_blob = norms_f32.reshape(n // 4, 4, groups).transpose(0, 2, 1).astype(np.float16).copy().tobytes() + packed = pack_indices_interleaved_4row(cq.indices, group_size, cq.bits) + else: + norms_blob = norms_f32.astype(np.float16).tobytes() + packed = pack_indices_lsb(cq.indices, group_size, cq.bits) + + parts = [codebook.tobytes(), input_scale.tobytes(), recip.tobytes(), norms_blob] + flags = 0 + if interleaved: + flags |= FLAG_INTERLEAVED_4ROW + if cq.rotation_family == "orthogonal": + flags |= FLAG_ORTHOGONAL_ROTATION + parts.append(make_orthogonal_rotation(group_size, cq.seed).astype(np.float16).tobytes()) + else: + left, right, perm = make_hadamard_components(group_size, cq.seed) + parts.extend([left.tobytes(), right.tobytes(), perm.astype(" 0 + assert scales_bytes > 0 + + +def test_orthogonal_embedding_is_cq4(tmp_path): + w = np.random.default_rng(1).standard_normal((4, 16), dtype=np.float32) + cq = quantize_orthogonal(w, bits=4) + out = tmp_path / "embed.weights" + write_cq_tensor(out, cq) + precision = struct.unpack_from("": 0, "<|startoftext|>": 1, "<|im_end|>": 2, "": 3, "hello": 4}, unk_token="<|pad|>")) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "TokenizersBackend", + "bos_token": "<|startoftext|>", + "eos_token": "<|im_end|>", + "pad_token": "<|pad|>", + "image_token": "", + "image_start_token": "", + "image_end_token": "", + "image_thumbnail": "", + } + ), + encoding="utf-8", + ) + (tmp_path / "preprocessor_config.json").write_text( + json.dumps( + { + "image_processor_type": "Lfm2VlImageProcessorFast", + "do_resize": True, + "size": {"height": 512, "width": 512}, + "do_rescale": True, + "rescale_factor": 1 / 255, + "do_normalize": True, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + "do_pad": True, + "data_format": "channels_first", + } + ), + encoding="utf-8", + ) + + processor = adapter_for_family("lfm2").load_processor(str(tmp_path)) + assert isinstance(processor, Lfm2VlProcessor) + assert processor.image_token == "" + assert processor.image_token_id == 3 + + +def test_lfm2_moe_packed_gate_up_splits_into_runtime_experts(): + adapter = adapter_for_family("lfm2") + tensor = torch.arange(2 * 6 * 4, dtype=torch.float16).reshape(2, 6, 4) + + match = adapter.name_tensor( + "model.layers.5.feed_forward.experts.gate_up_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w1.weights", + "layer_5_moe_expert_0_w3.weights", + "layer_5_moe_expert_1_w1.weights", + "layer_5_moe_expert_1_w3.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0, :3, :]) + assert torch.equal(emissions[1].tensor, tensor[0, 3:, :]) + assert "model.layers.5.feed_forward.w1_weights.0" in emissions[0].source_names + assert "model.layers.5.feed_forward.w3_weights.1" in emissions[3].source_names + + +def test_lfm2_moe_packed_down_splits_into_runtime_experts(): + adapter = adapter_for_family("lfm2") + tensor = torch.arange(2 * 4 * 3, dtype=torch.float16).reshape(2, 4, 3) + + match = adapter.name_tensor( + "model.layers.5.feed_forward.experts.down_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w2.weights", + "layer_5_moe_expert_1_w2.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0]) + assert torch.equal(emissions[1].tensor, tensor[1]) + assert "model.layers.5.feed_forward.w2_weights.0" in emissions[0].source_names diff --git a/python/cactus/convert/tests/test_naming_qdq.py b/python/cactus/convert/tests/test_naming_qdq.py new file mode 100644 index 000000000..f5d2976e8 --- /dev/null +++ b/python/cactus/convert/tests/test_naming_qdq.py @@ -0,0 +1,419 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from safetensors.torch import load_file + +from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header +from cactus.convert.export.qdq import convert_qdq +from cactus.convert.model_adapters.adapters import adapter_for_family +from cactus.convert.model_adapters.naming import cactus_name_for_tensor, restore_hf_key_for_family + + +def test_gemma4_audio_aux_roundtrip_name(): + hf = "model.audio_tower.layers.0.self_attn.q_proj.input_min" + match = cactus_name_for_tensor(hf, "gemma4", 35) + assert match.output_name == "audio_conformer_0_attention_attn_q_proj_input_min.weights" + assert match.hf_name == hf + assert match.adapter_name == "model.audio_tower.conformer.0.attention.attn.q_proj.input_min" + assert restore_hf_key_for_family(match.adapter_name, "gemma4") == hf + assert "q_proj_proj" not in restore_hf_key_for_family(match.adapter_name, "gemma4") + + +def test_needle_encoder_layers_beyond_decoder_depth_are_recognized(): + match = cactus_name_for_tensor("model.encoder.layers.11.self_attn.k_proj.weight", "needle", 12) + assert match.recognized + assert match.output_name == "encoder_layer_11_attn_k.weights" + + +def test_needle_global_norms_map_to_bundle_names(): + encoder_norm = cactus_name_for_tensor("model.encoder.final_norm.weight", "needle") + assert encoder_norm.recognized + assert encoder_norm.output_name == "encoder_layer_norm_weight.weights" + decoder_norm = cactus_name_for_tensor("model.decoder.norm.weight", "needle") + assert decoder_norm.recognized + assert decoder_norm.output_name == "output_norm.weights" + + +def test_manifest_qdq_restores_aux_stats_without_suffix(tmp_path): + cactus = tmp_path / "cactus" + out = tmp_path / "qdq" + cactus.mkdir() + tensor = torch.tensor([0.25], dtype=torch.float16) + output_file = "audio_conformer_0_attention_attn_q_proj_input_min.weights" + save_tensor_with_header(tensor, cactus / output_file, precision="FP16") + (cactus / "conversion_manifest.json").write_text( + """[ + { + "source_name": "model.audio_tower.conformer.0.attention.attn.q_proj.input_min", + "hf_name": "model.audio_tower.layers.0.self_attn.q_proj.input_min", + "adapter_name": "model.audio_tower.conformer.0.attention.attn.q_proj.input_min", + "output_file": "audio_conformer_0_attention_attn_q_proj_input_min.weights", + "shape": [1], + "dtype": "torch.float16", + "component": "audio", + "policy": "fallback", + "precision": "FP16", + "status": "fallback", + "required": true, + "scale_factor": 1.0 + } +]""", + encoding="utf-8", + ) + args = SimpleNamespace( + input=cactus, + out=out, + dtype="float16", + model_family="gemma4", + shard_size_gb=1.0, + row_batch_size=64, + tmp_dir=None, + force=True, + ) + report = convert_qdq(args) + assert report["written_count"] == 1 + tensors = load_file(out / "model.safetensors") + assert set(tensors) == {"model.audio_tower.layers.0.self_attn.q_proj.input_min"} + assert tensors["model.audio_tower.layers.0.self_attn.q_proj.input_min"].shape == (1,) + + +def test_manifest_qdq_trims_int8_bias_padding(tmp_path): + cactus = tmp_path / "cactus" + out = tmp_path / "qdq" + cactus.mkdir() + output_file = "projector_linear1.bias.weights" + save_tensor_with_header(torch.tensor([1.0, -2.0, 3.0]), cactus / output_file, precision="INT8", allow_int8_bias=True) + (cactus / "conversion_manifest.json").write_text( + """[ + { + "source_name": "model.multi_modal_projector.linear_1.bias", + "hf_name": "model.multi_modal_projector.linear_1.bias", + "adapter_name": "model.multi_modal_projector.linear_1.bias", + "output_file": "projector_linear1.bias.weights", + "shape": [3], + "dtype": "torch.float32", + "component": "vision", + "policy": "fallback", + "precision": "INT8", + "status": "fallback", + "required": true, + "scale_factor": 1.0 + } +]""", + encoding="utf-8", + ) + args = SimpleNamespace( + input=cactus, + out=out, + dtype="float16", + model_family="generic", + shard_size_gb=1.0, + row_batch_size=64, + tmp_dir=None, + force=True, + ) + convert_qdq(args) + tensors = load_file(out / "model.safetensors") + assert tensors["model.multi_modal_projector.linear_1.bias"].shape == (3,) + + +def test_manifest_qdq_uses_adapter_output_keys(tmp_path): + cactus = tmp_path / "cactus" + out = tmp_path / "qdq" + cactus.mkdir() + output_file = "tdt_predictor_lstm_0_bias.weights" + save_tensor_with_header(torch.tensor([1.0, 2.0]), cactus / output_file, precision="FP16") + (cactus / "conversion_manifest.json").write_text( + """[ + { + "source_name": "decoder.lstm.bias_l0", + "hf_name": "decoder.lstm.bias_l0", + "adapter_name": "decoder.lstm.bias_l0", + "source_names": ["decoder.lstm.bias_ih_l0", "decoder.lstm.bias_hh_l0"], + "transform": "parakeet_tdt_lstm_bias_sum", + "qdq_restore": "hf_key", + "output_file": "tdt_predictor_lstm_0_bias.weights", + "shape": [2], + "dtype": "torch.float32", + "component": "transcription", + "policy": "fallback", + "precision": "FP16", + "status": "fallback", + "required": true, + "scale_factor": 1.0 + } +]""", + encoding="utf-8", + ) + args = SimpleNamespace( + input=cactus, + out=out, + dtype="float16", + model_family="parakeet_tdt", + shard_size_gb=1.0, + row_batch_size=64, + tmp_dir=None, + force=True, + ) + convert_qdq(args) + tensors = load_file(out / "model.safetensors") + assert set(tensors) == {"decoder.lstm.bias_l0"} + + +def test_manifest_qdq_restores_int8_depthwise_conv_shape(tmp_path): + from cactus.convert.cactus_adapters.tensor_io import save_depthwise_conv_int8_with_header + + cactus = tmp_path / "cactus" + out = tmp_path / "qdq" + cactus.mkdir() + output_file = "layer_0_conv_depthwise.weights" + save_depthwise_conv_int8_with_header(torch.randn(4, 1, 3), cactus / output_file) + (cactus / "conversion_manifest.json").write_text( + """[ + { + "source_name": "model.layers.0.conv.conv.weight", + "hf_name": "model.layers.0.conv.conv.weight", + "adapter_name": "model.layers.0.conv.conv.weight", + "output_file": "layer_0_conv_depthwise.weights", + "shape": [4, 1, 3], + "dtype": "torch.float32", + "component": "language", + "policy": "fallback", + "precision": "INT8", + "status": "fallback", + "required": true, + "scale_factor": 1.0 + } +]""", + encoding="utf-8", + ) + args = SimpleNamespace( + input=cactus, + out=out, + dtype="float16", + model_family="lfm2", + shard_size_gb=1.0, + row_batch_size=64, + tmp_dir=None, + force=True, + ) + convert_qdq(args) + tensors = load_file(out / "model.safetensors") + assert tensors["model.layers.0.conv.conv.weight"].shape == (4, 1, 3) + + +def test_lfm_vl_siglip2_layer_names_are_cactus_runtime_names(): + cases = { + "model.vision_tower.vision_model.encoder.layers.3.layer_norm1.bias": "vision_layer_3_layer_norm1.bias.weights", + "model.vision_tower.vision_model.encoder.layers.3.layer_norm2.weight": "vision_layer_3_layer_norm2.weights", + "model.vision_tower.vision_model.encoder.layers.3.mlp.fc2.weight": "vision_layer_3_ffn_fc2.weights", + "model.vision_tower.vision_model.encoder.layers.3.self_attn.out_proj.bias": "vision_layer_3_self_attn_out.bias.weights", + } + for source, expected in cases.items(): + match = cactus_name_for_tensor(source, "lfm2", 16) + assert match.recognized + assert match.component == "vision" + assert match.output_name == expected + + +def test_qwen_vl_visual_names_are_cactus_style(): + cases = { + "model.visual.patch_embed.proj.weight": "vision_patch_embedding.weights", + "model.visual.blocks.12.attn.qkv.bias": "vision_layer_12_self_attn_qkv.bias.weights", + "model.visual.blocks.12.mlp.linear_fc2.weight": "vision_layer_12_ffn_fc2.weights", + "model.visual.blocks.12.norm1.weight": "vision_layer_12_layer_norm1.weights", + "model.visual.merger.norm.bias": "vision_merger_norm.bias.weights", + "model.visual.deepstack_merger_list.2.linear_fc1.bias": "vision_deepstack_merger_2_linear_fc1.bias.weights", + } + for source, expected in cases.items(): + match = cactus_name_for_tensor(source, "qwen", 28) + assert match.recognized + assert match.component == "vision" + assert match.output_name == expected + + +def test_whisper_layer_names_include_runtime_block_prefix(): + cases = { + "encoder.layers.2.self_attn.q_proj.weight": "encoder.layer_2_self_attn_q.weights", + "model.encoder.layers.2.self_attn.q_proj.weight": "encoder.layer_2_self_attn_q.weights", + "model.encoder.conv1.weight": "encoder_conv1_weight.weights", + "encoder.layers.2.fc1.bias": "encoder.layer_2_mlp_fc1.bias", + "decoder.layers.4.encoder_attn.out_proj.bias": "decoder.layer_4_encoder_attn_output.bias", + "model.decoder.layers.4.encoder_attn.out_proj.bias": "decoder.layer_4_encoder_attn_output.bias", + "decoder.layers.4.self_attn_layer_norm.weight": "decoder.layer_4_self_attn_norm.weights", + "model.decoder.embed_tokens.weight": "decoder_token_embeddings.weights", + "proj_out.weight": "output_weight.weights", + } + for source, expected in cases.items(): + match = cactus_name_for_tensor(source, "whisper", 8) + assert match.recognized + assert match.output_name == expected + if source.startswith(("encoder.layers.", "decoder.layers.")): + assert match.component == "transcription" + + +def test_whisper_unknown_layer_name_is_not_recognized(): + match = cactus_name_for_tensor("encoder.layers.2.some_new_projection.weight", "whisper", 8) + assert not match.recognized + assert match.output_name == "encoder_layers_2_some_new_projection.weights" + + +def test_parakeet_ctc_names_are_cactus_runtime_names(): + cases = { + "encoder.subsampling.layers.0.weight": "subsampling_conv0_weight.weights", + "encoder.subsampling.linear.bias": "subsampling_linear_bias.bias", + "ctc_head.weight": "ctc_head_weight.weights", + "encoder.layers.3.feed_forward1.linear1.weight": "layer_3_ff1_linear1.weights", + "encoder.layers.3.self_attn.relative_k_proj.weight": "layer_3_self_attn_relative_k.weights", + "encoder.layers.3.conv.norm.running_var": "layer_3_conv_batchnorm_running_var.weights", + "encoder.layers.3.norm_self_att.bias": "layer_3_norm_self_attn.bias", + } + for source, expected in cases.items(): + match = cactus_name_for_tensor(source, "parakeet", 24) + assert match.recognized + assert match.component == "transcription" + assert match.output_name == expected + + +def test_parakeet_tdt_names_are_cactus_runtime_names(): + cases = { + "encoder.pre_encode.conv.0.weight": "subsampling_conv0_weight.weights", + "decoder.embedding.weight": "tdt_predictor_embed.weights", + "decoder.prediction.embed.weight": "tdt_predictor_embed.weights", + "decoder.prediction.dec_rnn.lstm.1.Wx": "tdt_predictor_lstm_1_weight_ih.weights", + "decoder.lstm.weight_ih_l1": "tdt_predictor_lstm_1_weight_ih.weights", + "decoder.prediction.dec_rnn.lstm.1.Wh": "tdt_predictor_lstm_1_weight_hh.weights", + "decoder.lstm.weight_hh_l1": "tdt_predictor_lstm_1_weight_hh.weights", + "decoder.prediction.dec_rnn.lstm.1.bias": "tdt_predictor_lstm_1_bias.weights", + "decoder.lstm.bias_l1": "tdt_predictor_lstm_1_bias.weights", + "encoder_projector.weight": "tdt_joint_enc.weights", + "decoder.decoder_projector.bias": "tdt_joint_pred.bias", + "joint.joint_net.0.weight": "tdt_joint_out.weights", + "joint.head.bias": "tdt_joint_out.bias", + "encoder.layers.5.self_attention.linear_out.bias": "layer_5_self_attn_output.bias", + } + for source, expected in cases.items(): + match = cactus_name_for_tensor(source, "parakeet_tdt", 24) + assert match.recognized + assert match.component == "transcription" + assert match.output_name == expected + + +def test_parakeet_batchnorm_tracking_tensors_are_ignored(): + match = cactus_name_for_tensor("encoder.layers.0.conv.norm.num_batches_tracked", "parakeet", 24) + assert match.recognized + assert match.output_name is None + + +def test_nomic_normalizes_global_tensors(): + adapter = adapter_for_family("nomic") + state = { + "embeddings.word_embeddings.weight": torch.ones(4, 3), + "embeddings.token_type_embeddings.weight": torch.full((1, 3), 2.0), + "emb_ln.weight": torch.arange(3.0), + "emb_ln.bias": torch.arange(3.0) + 10, + } + normalized = adapter.normalize_state_dict(state) + assert set(normalized.state_dict) == {"token_embeddings", "embedding_layernorm.weight", "embedding_layernorm.bias"} + assert torch.equal(normalized.state_dict["token_embeddings"], torch.full((4, 3), 3.0)) + assert normalized.provenance["token_embeddings"].source_names == [ + "embeddings.word_embeddings.weight", + "embeddings.token_type_embeddings.weight", + ] + assert normalized.provenance["token_embeddings"].qdq_restore == "adapter_key" + assert adapter.name_tensor("token_embeddings", normalized.state_dict["token_embeddings"], 12).output_name == "token_embeddings.weights" + assert adapter.name_tensor("embedding_layernorm.weight", normalized.state_dict["embedding_layernorm.weight"], 12).output_name == "embedding_layernorm.weight" + + +def test_nomic_norm2_weight_uses_runtime_name(): + adapter = adapter_for_family("nomic") + match = adapter.name_tensor("encoder.layers.3.norm2.weight", torch.ones(768), 12) + assert match.recognized + assert match.output_name == "layer_3_norm2.weights" + + +def test_nomic_keeps_qkv_and_moe_experts_fused(): + # The v2 transpile path binds graph weights by their HF parameter name, so the + # converter emits one fused tensor per HF parameter (no q/k/v or per-expert split). + # w2 is stored transposed so the second expert matmul can consume it as a direct + # linear weight in the transpiled graph. + adapter = adapter_for_family("nomic") + adapter.num_experts = 8 + + qkv = torch.arange(2304 * 2, dtype=torch.float32).reshape(2304, 2) + match = adapter.name_tensor("encoder.layers.0.attn.Wqkv.weight", qkv, 12) + emissions = adapter.expand_tensor(match, qkv) + assert [e.output_name for e in emissions] == ["layer_0_attn_qkv.weights"] + assert tuple(emissions[0].tensor.shape) == (2304, 2) + + w1 = torch.empty(24576, 2) + match = adapter.name_tensor("encoder.layers.1.mlp.experts.mlp.w1", w1, 12) + emissions = adapter.expand_tensor(match, w1) + assert [e.output_name for e in emissions] == ["layer_1_mlp_experts_w1.weights"] + assert tuple(emissions[0].tensor.shape) == (24576, 2) + + w2 = torch.empty(24576, 2) + match = adapter.name_tensor("encoder.layers.1.mlp.experts.mlp.w2", w2, 12) + emissions = adapter.expand_tensor(match, w2) + assert [e.output_name for e in emissions] == ["layer_1_mlp_experts_w2.weights"] + assert tuple(emissions[0].tensor.shape) == (2, 24576) + + +def test_nomic_qdq_runtime_keys_are_unique(tmp_path): + cactus = tmp_path / "cactus" + out = tmp_path / "qdq" + cactus.mkdir() + save_tensor_with_header(torch.ones(2, 3), cactus / "layer_0_attn_q.weights", precision="FP16") + save_tensor_with_header(torch.ones(2, 3) * 2, cactus / "layer_0_attn_k.weights", precision="FP16") + (cactus / "conversion_manifest.json").write_text( + """[ + { + "source_name": "encoder.layers.0.attn.Wqkv.weight", + "hf_name": "encoder.layers.0.attn.Wqkv.weight", + "adapter_name": "encoder.layers.0.attn.Wqkv.weight", + "output_file": "layer_0_attn_q.weights", + "shape": [2, 3], + "dtype": "torch.float32", + "component": "language", + "policy": "fallback", + "precision": "FP16", + "status": "fallback", + "required": true, + "qdq_restore": "runtime_key", + "scale_factor": 1.0 + }, + { + "source_name": "encoder.layers.0.attn.Wqkv.weight", + "hf_name": "encoder.layers.0.attn.Wqkv.weight", + "adapter_name": "encoder.layers.0.attn.Wqkv.weight", + "output_file": "layer_0_attn_k.weights", + "shape": [2, 3], + "dtype": "torch.float32", + "component": "language", + "policy": "fallback", + "precision": "FP16", + "status": "fallback", + "required": true, + "qdq_restore": "runtime_key", + "scale_factor": 1.0 + } +]""", + encoding="utf-8", + ) + report = convert_qdq( + SimpleNamespace( + input=cactus, + out=out, + dtype="float16", + model_family="nomic", + shard_size_gb=1.0, + row_batch_size=64, + tmp_dir=None, + force=True, + ) + ) + tensors = load_file(out / "model.safetensors") + assert report["written_count"] == 2 + assert set(tensors) == {"layer_0_attn_q", "layer_0_attn_k"} diff --git a/python/cactus/convert/tests/test_policy.py b/python/cactus/convert/tests/test_policy.py new file mode 100644 index 000000000..384811311 --- /dev/null +++ b/python/cactus/convert/tests/test_policy.py @@ -0,0 +1,293 @@ +from cactus.convert.model_adapters.naming import cactus_name_for_tensor +from cactus.convert.model_adapters.policy import policy_for_tensor +from cactus.convert.cli import _bits_for_component, _validate_cq_layout +from cactus.convert.model_adapters.detection import detect_family +from cactus.convert.cactus_adapters.config_utils import extract_parakeet_tdt_config, extract_whisper_config +from cactus.convert.cli import _augment_state_dict_for_family +from cactus.convert.model_adapters.adapters import adapter_for_family +import pytest +import torch + + +@pytest.mark.parametrize("source_name", ["model.embed_tokens.weight", "lm_head.weight"]) +def test_policy_embedding_lm_head_interleaved_cq4(source_name): + match = cactus_name_for_tensor(source_name, "generic", 1) + p = policy_for_tensor(match, (100, 64), 2, "generic") + assert p.precision == "CQ4" + assert p.rotation == "orthogonal" + assert p.layout == "interleaved_4row" + + +def test_policy_lfm_qwen_tied_embeddings_interleaved_cq4(): + for family in ("lfm2", "qwen"): + embed = cactus_name_for_tensor("model.language_model.embed_tokens.weight", family, 1) + embed_policy = policy_for_tensor(embed, (65536, 2048), 2, family) + assert embed.output_name == "token_embeddings.weights" + assert embed_policy.precision == "CQ4" + assert embed_policy.rotation == "orthogonal" + assert embed_policy.layout == "interleaved_4row" + _validate_cq_layout(embed_policy, (65536, 2048), embed.source_name, embed.output_name) + + lm_head = cactus_name_for_tensor("lm_head.weight", family, 1) + lm_head_policy = policy_for_tensor(lm_head, (65536, 2048), 2, family) + assert lm_head.output_name == "output_weight.weights" + assert lm_head_policy.precision == "CQ4" + assert lm_head_policy.rotation == "orthogonal" + assert lm_head_policy.layout == "interleaved_4row" + _validate_cq_layout(lm_head_policy, (65536, 2048), lm_head.source_name, lm_head.output_name) + + +def test_interleaved_lm_head_policy_falls_back_to_row_major_for_unsupported_shape(): + match = cactus_name_for_tensor("lm_head.weight", "generic", 1) + p = policy_for_tensor(match, (101, 64), 2, "generic") + assert p.layout == "row_major" + _validate_cq_layout(p, (101, 64), "lm_head.weight", "output_weight.weights") + + +def test_interleaved_lm_head_policy_uses_interleaved_for_supported_shape(): + match = cactus_name_for_tensor("lm_head.weight", "generic", 1) + p = policy_for_tensor(match, (128, 256), 2, "generic") + assert p.layout == "interleaved_4row" + + +def test_adapter_registry_preserves_family_names(): + assert adapter_for_family("qwen").family == "qwen" + assert adapter_for_family("lfm2").family == "lfm2" + + +def test_policy_gemma4_pli_honors_requested_bits(): + name = "model.language_model.embed_tokens_per_layer.weight" + match = cactus_name_for_tensor(name, "gemma4", 1) + p = policy_for_tensor(match, (10, 128), 4, "gemma4") + assert p.precision == "CQ4" + assert p.bits == 4 + assert p.rotation == "hadamard" + + +def test_gemma4_adapter_disables_gptq_for_unhookable_tensors(): + adapter = adapter_for_family("gemma4") + shared_kv = cactus_name_for_tensor("model.language_model.layers.15.self_attn.k_proj.weight", "gemma4", 35) + shared_policy = adapter.policy(shared_kv, (256, 1536), 4) + assert shared_policy.precision == "CQ4" + assert not shared_policy.use_gptq + + vision_proj = cactus_name_for_tensor("model.embed_vision.embedding_projection.weight", "gemma4", 35) + vision_policy = adapter.policy(vision_proj, (1536, 1152), 4) + assert vision_policy.precision == "FP16" + assert not vision_policy.use_gptq + + vision_tower = cactus_name_for_tensor("model.vision_tower.encoder.layers.0.mlp.down_proj.linear.weight", "gemma4", 35) + vision_tower_policy = adapter.policy(vision_tower, (1152, 4304), 4) + assert vision_tower_policy.precision == "FP16" + assert not vision_tower_policy.use_gptq + assert adapter.module_target_name("model.vision_tower.encoder.layers.0.mlp.down_proj.linear.weight", None) == "model.vision_tower.encoder.layers.0.mlp.down_proj.linear" + + +def test_policy_audio_no_gptq(): + match = cactus_name_for_tensor("model.audio_tower.output_proj.weight", "gemma4", 1) + p = policy_for_tensor(match, (128, 128), 4, "gemma4") + assert p.component == "audio" + assert p.precision == "FP16" + assert not p.use_gptq + + +def test_policy_gemma4_audio_bias_fp16(): + match = cactus_name_for_tensor("model.audio_tower.output_proj.bias", "gemma4", 1) + p = policy_for_tensor(match, (1536,), 4, "gemma4") + assert p.precision == "FP16" + assert p.bits is None + + +def test_policy_position_embedding_fp16(): + match = cactus_name_for_tensor("model.vision_tower.vision_model.embeddings.position_embedding.weight", "lfm2", 1) + p = policy_for_tensor(match, (256, 1152), 4, "lfm2") + assert p.precision == "FP16" + assert p.fallback_reason == "position embedding tensor" + + +def test_policy_lfm_depthwise_conv_int8(): + match = cactus_name_for_tensor("model.layers.0.conv.conv.weight", "lfm2", 16) + p = policy_for_tensor(match, (1024, 1, 3), 4, "lfm2") + assert p.precision == "INT8" + assert p.fallback_reason == "depthwise conv tensor" + + +def test_cli_component_bit_overrides(): + class Args: + bits = 4 + language_bits = 2 + vision_bits = 4 + audio_bits = 3 + embedding_bits = None + + assert _bits_for_component("language", Args) == 2 + assert _bits_for_component("vision", Args) == 4 + assert _bits_for_component("audio", Args) == 3 + assert _bits_for_component("transcription", Args) == 3 + assert _bits_for_component("embedding", Args) == 4 + + +def test_policy_parakeet_transcription_no_gptq(): + match = cactus_name_for_tensor("encoder.layers.0.self_attn.q_proj.weight", "parakeet", 24) + p = policy_for_tensor(match, (1024, 1024), 4, "parakeet") + assert p.component == "transcription" + assert p.precision == "CQ4" + assert not p.use_gptq + + +def test_parakeet_tdt_hf_config_detection_and_extraction(): + cfg = { + "model_type": "parakeet_tdt", + "vocab_size": 8193, + "pad_token_id": 2, + "blank_token_id": 8192, + "decoder_hidden_size": 640, + "num_decoder_layers": 2, + "durations": [0, 1, 2, 3, 4], + "encoder_config": { + "hidden_size": 1024, + "num_hidden_layers": 24, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "intermediate_size": 4096, + "max_position_embeddings": 5000, + "conv_kernel_size": 9, + "subsampling_conv_kernel_size": 3, + "subsampling_conv_stride": 2, + "subsampling_conv_channels": 256, + "subsampling_factor": 8, + "num_mel_bins": 128, + "hidden_act": "silu", + }, + } + assert detect_family(cfg) == "parakeet_tdt" + extracted = extract_parakeet_tdt_config(cfg) + assert extracted["num_layers"] == 24 + assert extracted["num_mel_bins"] == 128 + assert extracted["predictor_hidden_dim"] == 640 + assert extracted["predictor_num_layers"] == 2 + assert extracted["tdt_durations"] == [0, 1, 2, 3, 4] + assert extracted["tdt_blank_id"] == 8192 + + +def test_whisper_hf_config_extraction(): + cfg = { + "model_type": "whisper", + "vocab_size": 51865, + "d_model": 384, + "encoder_layers": 4, + "decoder_layers": 4, + "encoder_attention_heads": 6, + "decoder_attention_heads": 6, + "encoder_ffn_dim": 1536, + "decoder_ffn_dim": 1536, + "max_target_positions": 448, + "num_mel_bins": 80, + "pad_token_id": 50257, + "bos_token_id": 50257, + "eos_token_id": 50257, + "tie_word_embeddings": False, + } + extracted = extract_whisper_config(cfg) + assert detect_family(cfg) == "whisper" + assert extracted["hidden_dim"] == 384 + assert extracted["num_encoder_layers"] == 4 + assert extracted["num_decoder_layers"] == 4 + assert extracted["attention_head_dim"] == 64 + assert extracted["context_length"] == 448 + assert not extracted["tie_word_embeddings"] + + +def test_parakeet_tdt_lstm_biases_are_combined_for_runtime(): + state = { + "decoder.lstm.bias_ih_l0": torch.tensor([1.0, 2.0]), + "decoder.lstm.bias_hh_l0": torch.tensor([0.5, -1.0]), + } + out = _augment_state_dict_for_family(state, "parakeet_tdt") + assert "decoder.lstm.bias_ih_l0" not in out + assert "decoder.lstm.bias_hh_l0" not in out + assert torch.equal(out["decoder.lstm.bias_l0"], torch.tensor([1.5, 1.0])) + + +def test_parakeet_tdt_lstm_bias_normalization_records_provenance(): + state = { + "decoder.lstm.bias_ih_l0": torch.tensor([1.0, 2.0]), + "decoder.lstm.bias_hh_l0": torch.tensor([0.5, -1.0]), + } + normalized = adapter_for_family("parakeet_tdt").normalize_state_dict(state) + provenance = normalized.provenance["decoder.lstm.bias_l0"] + assert provenance.source_names == ["decoder.lstm.bias_ih_l0", "decoder.lstm.bias_hh_l0"] + assert provenance.transform == "parakeet_tdt_lstm_bias_sum" + + +def test_policy_parakeet_pointwise_conv_int8_and_conv_bias_fp16(): + weight = cactus_name_for_tensor("encoder.layers.0.conv.pointwise_conv1.weight", "parakeet_tdt", 24) + weight_policy = policy_for_tensor(weight, (2048, 1024, 1), 4, "parakeet_tdt") + assert weight_policy.precision == "INT8" + assert weight_policy.fallback_reason == "pointwise conv tensor" + + bias = cactus_name_for_tensor("encoder.layers.0.conv.pointwise_conv1.bias", "parakeet_tdt", 24) + bias_policy = policy_for_tensor(bias, (2048,), 4, "parakeet_tdt") + assert bias_policy.precision == "FP16" + assert bias_policy.fallback_reason == "conv bias tensor" + + +def test_parakeet_adapter_conv_transform_is_scoped(): + adapter = adapter_for_family("parakeet_tdt") + match = cactus_name_for_tensor("encoder.pre_encode.conv.0.weight", "parakeet_tdt", 24) + tensor = torch.randn(256, 3, 3, 1) + transformed, transform = adapter.transform_tensor(match, tensor) + assert transform == "parakeet_conv_hwio_to_oihw" + assert transformed.shape == (256, 1, 3, 3) + + other = cactus_name_for_tensor("encoder.layers.0.feed_forward1.linear1.weight", "parakeet_tdt", 24) + untouched, transform = adapter.transform_tensor(other, torch.randn(4, 4)) + assert transform == "none" + assert untouched.shape == (4, 4) + + +def _nomic_cfg() -> dict[str, object]: + return { + "model_type": "nomic_bert", + "vocab_size": 250048, + "n_embd": 768, + "n_head": 12, + "n_layer": 12, + "n_inner": 3072, + "n_positions": 2048, + "num_experts": 8, + "moe_top_k": 2, + "moe_every_n_layers": 2, + } + + +def test_nomic_detection_and_runtime_config(): + adapter = adapter_for_family("nomic") + assert detect_family(_nomic_cfg(), "auto") == "nomic" + assert adapter.runtime_model_type() == "bert" + runtime = adapter.runtime_config(_nomic_cfg()) + assert runtime["num_layers"] == 12 + assert runtime["hidden_dim"] == 768 + assert runtime["attention_heads"] == 12 + assert runtime["num_experts"] == 8 + assert runtime["num_experts_per_tok"] == 2 + assert runtime["moe_every_n_layers"] == 2 + + +def test_nomic_uses_gptq_only_for_hookable_linear_modules(): + adapter = adapter_for_family("nomic") + qkv_match = adapter.name_tensor("encoder.layers.0.attn.Wqkv.weight", torch.empty(2304, 768), 12) + qkv_policy = adapter.policy(qkv_match, (768, 768), 4) + assert qkv_policy.use_gptq + assert adapter.module_target_name("encoder.layers.0.attn.Wqkv.weight", None) == "encoder.layers.0.attn.Wqkv" + + expert_match = adapter.name_tensor("encoder.layers.1.mlp.experts.mlp.w1", torch.empty(24576, 768), 12) + expert_policy = adapter.policy(expert_match, (3072, 768), 4) + assert not expert_policy.use_gptq + assert adapter.module_target_name("encoder.layers.1.mlp.experts.mlp.w1", None) is None + + +def test_nomic_router_stays_fp16(): + adapter = adapter_for_family("nomic") + match = adapter.name_tensor("encoder.layers.1.mlp.router.layer.weight", torch.empty(8, 768), 12) + assert adapter.policy(match, (8, 768), 4).precision == "FP16" diff --git a/python/cactus/convert/tests/test_qwen_adapter.py b/python/cactus/convert/tests/test_qwen_adapter.py new file mode 100644 index 000000000..e3e3e48f9 --- /dev/null +++ b/python/cactus/convert/tests/test_qwen_adapter.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import torch + +from cactus.convert.model_adapters.adapters import adapter_for_family + + +def test_qwen2_moe_packed_gate_up_splits_into_runtime_experts(): + adapter = adapter_for_family("qwen") + tensor = torch.arange(2 * 6 * 4, dtype=torch.float16).reshape(2, 6, 4) + + match = adapter.name_tensor( + "model.layers.5.mlp.experts.gate_up_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w1.weights", + "layer_5_moe_expert_0_w3.weights", + "layer_5_moe_expert_1_w1.weights", + "layer_5_moe_expert_1_w3.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0, :3, :]) + assert torch.equal(emissions[1].tensor, tensor[0, 3:, :]) + assert "model.layers.5.mlp.w1_weights.0" in emissions[0].source_names + assert "model.layers.5.mlp.w3_weights.1" in emissions[3].source_names + + +def test_qwen2_moe_packed_down_splits_into_runtime_experts(): + adapter = adapter_for_family("qwen") + tensor = torch.arange(2 * 4 * 3, dtype=torch.float16).reshape(2, 4, 3) + + match = adapter.name_tensor( + "model.layers.5.mlp.experts.down_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w2.weights", + "layer_5_moe_expert_1_w2.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0]) + assert torch.equal(emissions[1].tensor, tensor[1]) + assert "model.layers.5.mlp.w2_weights.0" in emissions[0].source_names + + +def test_qwen2_moe_router_uses_runtime_router_name_and_fp16_policy(): + adapter = adapter_for_family("qwen") + tensor = torch.randn(60, 2048, dtype=torch.float16) + + match = adapter.name_tensor( + "model.layers.5.mlp.gate.weight", + tensor, + num_layers=24, + ) + policy = adapter.policy(match, tuple(tensor.shape), requested_bits=4) + + assert match.output_name == "layer_5_moe_router.weights" + assert policy.action == "fallback" + assert policy.precision == "FP16" diff --git a/python/cactus/convert/tests/test_reports.py b/python/cactus/convert/tests/test_reports.py new file mode 100644 index 000000000..625170e4d --- /dev/null +++ b/python/cactus/convert/tests/test_reports.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from cactus.convert.export.reports import write_reports + + +def test_weights_manifest_includes_bias_tensors(tmp_path: Path) -> None: + write_reports( + tmp_path, + [ + { + "source_name": "decoder.layers.0.fc1.bias", + "output_file": "decoder.layer_0_mlp_fc1.bias", + "shape": [4], + "precision": "FP16", + "status": "fallback", + "component": "transcription", + }, + { + "source_name": "decoder.layers.0.fc1.weight", + "output_file": "decoder.layer_0_mlp_fc1.weights", + "shape": [4, 4], + "precision": "CQ4", + "status": "converted", + "component": "transcription", + }, + ], + ) + + manifest = json.loads((tmp_path / "weights_manifest.json").read_text()) + output_names = {row["output_name"] for row in manifest["weights"]} + + assert "decoder.layer_0_mlp_fc1.bias" in output_names + assert "decoder.layer_0_mlp_fc1.weights" in output_names diff --git a/python/cactus/models/__init__.py b/python/cactus/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/models/needle/__init__.py b/python/cactus/models/needle/__init__.py new file mode 100644 index 000000000..8436b054a --- /dev/null +++ b/python/cactus/models/needle/__init__.py @@ -0,0 +1,25 @@ +"""Needle encoder-decoder model (mirrors the files published on Cactus-Compute/needle).""" + + +def register_with_transformers() -> None: + """Register the needle architecture with transformers' Auto* factories. + + Idempotent; imports torch/sentencepiece lazily so importing cactus stays light. + """ + from transformers import ( + AutoConfig, + AutoModel, + AutoModelForCausalLM, + AutoModelForSeq2SeqLM, + AutoTokenizer, + ) + + from .configuration_needle import NeedleConfig + from .modeling_needle import NeedleForCausalLM, NeedleModel + from .tokenization_needle import NeedleTokenizer + + AutoConfig.register("needle", NeedleConfig, exist_ok=True) + AutoModel.register(NeedleConfig, NeedleModel, exist_ok=True) + AutoModelForCausalLM.register(NeedleConfig, NeedleForCausalLM, exist_ok=True) + AutoModelForSeq2SeqLM.register(NeedleConfig, NeedleForCausalLM, exist_ok=True) + AutoTokenizer.register(NeedleConfig, slow_tokenizer_class=NeedleTokenizer, exist_ok=True) diff --git a/python/cactus/models/needle/configuration_needle.py b/python/cactus/models/needle/configuration_needle.py new file mode 100644 index 000000000..4fd2d3c98 --- /dev/null +++ b/python/cactus/models/needle/configuration_needle.py @@ -0,0 +1,63 @@ +"""Hugging Face configuration for Needle.""" + +from __future__ import annotations + +from transformers import PretrainedConfig + + +class NeedleConfig(PretrainedConfig): + model_type = "needle" + + def __init__( + self, + vocab_size: int = 8192, + hidden_size: int | None = None, + d_model: int = 512, + num_attention_heads: int | None = None, + num_heads: int = 8, + num_key_value_heads: int | None = None, + num_kv_heads: int = 4, + num_encoder_layers: int = 12, + num_decoder_layers: int = 8, + rope_theta: float = 10000.0, + rms_norm_eps: float = 1e-6, + pad_token_id: int = 0, + eos_token_id: int = 1, + bos_token_id: int = 2, + unk_token_id: int = 3, + decoder_start_token_id: int | None = None, + tie_word_embeddings: bool = True, + torch_dtype: str = "bfloat16", + **kwargs, + ) -> None: + kwargs.pop("is_encoder_decoder", None) + hidden_size = int(hidden_size if hidden_size is not None else d_model) + num_attention_heads = int(num_attention_heads if num_attention_heads is not None else num_heads) + num_key_value_heads = int(num_key_value_heads if num_key_value_heads is not None else num_kv_heads) + decoder_start_token_id = eos_token_id if decoder_start_token_id is None else decoder_start_token_id + + super().__init__( + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + bos_token_id=bos_token_id, + unk_token_id=unk_token_id, + decoder_start_token_id=decoder_start_token_id, + tie_word_embeddings=tie_word_embeddings, + is_encoder_decoder=True, + torch_dtype=torch_dtype, + **kwargs, + ) + + self.vocab_size = int(vocab_size) + self.hidden_size = hidden_size + self.d_model = hidden_size + self.num_attention_heads = num_attention_heads + self.num_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.num_kv_heads = num_key_value_heads + self.num_encoder_layers = int(num_encoder_layers) + self.num_decoder_layers = int(num_decoder_layers) + self.num_hidden_layers = int(num_decoder_layers) + self.rope_theta = float(rope_theta) + self.rms_norm_eps = float(rms_norm_eps) + self.attention_head_dim = hidden_size // max(1, num_attention_heads) diff --git a/python/cactus/models/needle/modeling_needle.py b/python/cactus/models/needle/modeling_needle.py new file mode 100644 index 000000000..b8e3cbb0c --- /dev/null +++ b/python/cactus/models/needle/modeling_needle.py @@ -0,0 +1,429 @@ +"""Minimal PyTorch Needle model for Cactus conversion.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +from torch import nn +import torch.nn.functional as F +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput + +from .configuration_needle import NeedleConfig + + +class NeedleRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.eps = float(eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + variance = x.float().pow(2).mean(dim=-1, keepdim=True) + x = x.float() * torch.rsqrt(variance + self.eps) + return x.to(dtype=dtype) * (1.0 + self.weight.to(dtype=dtype)) + + +def _padding_mask(input_ids: torch.Tensor, pad_token_id: int) -> torch.Tensor: + return (input_ids != int(pad_token_id))[:, None, None, :] + + +def _causal_mask(seq_len: int, device: torch.device) -> torch.Tensor: + return torch.ones((seq_len, seq_len), dtype=torch.bool, device=device).tril()[None, None, :, :] + + +def _build_inv_freq(head_dim: int, theta: float) -> torch.Tensor: + return 1.0 / (float(theta) ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / float(head_dim))) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + cos = cos.unsqueeze(2) + sin = sin.unsqueeze(2) + return (x * cos) + (_rotate_half(x) * sin) + + +def _rotary_tables( + inv_freq: torch.Tensor, + batch_size: int, + seq_len: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + position_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) + inv_freq = inv_freq[None, :, None].float().expand(batch_size, -1, 1).to(device) + freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype) + + +def _rotary_tables_for_position_ids( + inv_freq: torch.Tensor, + position_ids: torch.Tensor, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(position_ids.device) + freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype) + + +def _add_clipped(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return torch.clamp(a + b, min=-65500.0, max=65500.0) + + +class NeedleAttention(nn.Module): + def __init__(self, config: NeedleConfig) -> None: + super().__init__() + self.hidden_size = int(config.hidden_size) + self.num_heads = int(config.num_attention_heads) + self.num_key_value_heads = int(config.num_key_value_heads) + self.head_dim = self.hidden_size // self.num_heads + kv_size = self.num_key_value_heads * self.head_dim + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.k_proj = nn.Linear(self.hidden_size, kv_size, bias=False) + self.v_proj = nn.Linear(self.hidden_size, kv_size, bias=False) + self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.q_norm = NeedleRMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = NeedleRMSNorm(self.head_dim, config.rms_norm_eps) + self.scale = 1.0 / math.sqrt(float(self.head_dim)) + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: torch.Tensor, + attention_mask: torch.Tensor | None, + rope: tuple[torch.Tensor, torch.Tensor] | None, + ) -> torch.Tensor: + batch, q_len, _ = hidden_states.shape + kv_len = key_value_states.shape[1] + q = self.q_proj(hidden_states).view(batch, q_len, self.num_heads, self.head_dim) + k = self.k_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim) + v = self.v_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim) + q = self.q_norm(q) + k = self.k_norm(k) + if rope is not None: + cos, sin = rope + q = _apply_rope(q, cos, sin) + k = _apply_rope(k, cos, sin) + out = F.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + scale=self.scale, + enable_gqa=self.num_heads != self.num_key_value_heads, + ) + return self.out_proj(out.transpose(1, 2).contiguous().view(batch, q_len, self.hidden_size)) + + def project_kv(self, key_value_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch, kv_len, _ = key_value_states.shape + k = self.k_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim) + v = self.v_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim) + k = self.k_norm(k) + return k.contiguous(), v.contiguous() + + def forward_with_kv( + self, + hidden_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor | None, + rope: tuple[torch.Tensor, torch.Tensor] | None, + ) -> torch.Tensor: + batch, q_len, _ = hidden_states.shape + q = self.q_proj(hidden_states).view(batch, q_len, self.num_heads, self.head_dim) + q = self.q_norm(q) + if rope is not None: + cos, sin = rope + q = _apply_rope(q, cos, sin) + out = F.scaled_dot_product_attention( + q.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + scale=self.scale, + enable_gqa=self.num_heads != self.num_key_value_heads, + ) + return self.out_proj(out.transpose(1, 2).contiguous().view(batch, q_len, self.hidden_size)) + + +class NeedleEncoderLayer(nn.Module): + def __init__(self, config: NeedleConfig) -> None: + super().__init__() + self.input_layernorm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = NeedleAttention(config) + self.attn_gate = nn.Parameter(torch.zeros(1)) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + rope: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + normed = self.input_layernorm(hidden_states) + attn = self.self_attn(normed, normed, attention_mask, rope) + return _add_clipped(hidden_states, torch.sigmoid(self.attn_gate).to(dtype=attn.dtype) * attn) + + +class NeedleDecoderLayer(nn.Module): + def __init__(self, config: NeedleConfig) -> None: + super().__init__() + self.input_layernorm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = NeedleAttention(config) + self.self_attn_gate = nn.Parameter(torch.zeros(1)) + self.encoder_attn_layer_norm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps) + self.encoder_attn = NeedleAttention(config) + self.cross_attn_gate = nn.Parameter(torch.zeros(1)) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + self_mask: torch.Tensor, + encoder_mask: torch.Tensor, + rope: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + normed = self.input_layernorm(hidden_states) + attn = self.self_attn(normed, normed, self_mask, rope) + hidden_states = _add_clipped(hidden_states, torch.sigmoid(self.self_attn_gate).to(dtype=attn.dtype) * attn) + attn = self.encoder_attn(self.encoder_attn_layer_norm(hidden_states), encoder_hidden_states, encoder_mask, None) + return _add_clipped(hidden_states, torch.sigmoid(self.cross_attn_gate).to(dtype=attn.dtype) * attn) + + +class NeedleEncoder(nn.Module): + def __init__(self, config: NeedleConfig) -> None: + super().__init__() + self.layers = nn.ModuleList([NeedleEncoderLayer(config) for _ in range(config.num_encoder_layers)]) + self.final_norm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps) + self.head_dim = config.hidden_size // config.num_attention_heads + self.rope_theta = float(config.rope_theta) + self.register_buffer("inv_freq", _build_inv_freq(self.head_dim, self.rope_theta), persistent=False) + + def reset_rope(self) -> None: + self.inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=self.inv_freq.device) + + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=hidden_states.device) + rope = _rotary_tables( + inv_freq, + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.device, + hidden_states.dtype, + ) + for layer in self.layers: + hidden_states = layer(hidden_states, attention_mask, rope) + return self.final_norm(hidden_states) + + +class NeedleDecoder(nn.Module): + def __init__(self, config: NeedleConfig) -> None: + super().__init__() + self.layers = nn.ModuleList([NeedleDecoderLayer(config) for _ in range(config.num_decoder_layers)]) + self.norm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps) + self.head_dim = config.hidden_size // config.num_attention_heads + self.rope_theta = float(config.rope_theta) + self.register_buffer("inv_freq", _build_inv_freq(self.head_dim, self.rope_theta), persistent=False) + + def reset_rope(self) -> None: + self.inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=self.inv_freq.device) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + self_mask: torch.Tensor, + encoder_mask: torch.Tensor, + ) -> torch.Tensor: + inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=hidden_states.device) + rope = _rotary_tables( + inv_freq, + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.device, + hidden_states.dtype, + ) + for layer in self.layers: + hidden_states = layer(hidden_states, encoder_hidden_states, self_mask, encoder_mask, rope) + return self.norm(hidden_states) + + +class NeedleModel(PreTrainedModel): + config_class = NeedleConfig + base_model_prefix = "model" + main_input_name = "input_ids" + + def __init__(self, config: NeedleConfig) -> None: + super().__init__(config) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.embed_scale = math.sqrt(float(config.hidden_size)) + self.encoder = NeedleEncoder(config) + self.decoder = NeedleDecoder(config) + self.post_init() + self.reset_rope() + + def reset_rope(self) -> None: + self.encoder.reset_rope() + self.decoder.reset_rope() + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed_tokens + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.embed_tokens = value + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + **_: Any, + ) -> BaseModelOutput: + decoder_input_ids = input_ids if decoder_input_ids is None else decoder_input_ids + encoder_mask = _padding_mask(input_ids, self.config.pad_token_id) + if attention_mask is not None: + encoder_mask = encoder_mask & attention_mask[:, None, None, :].to(dtype=torch.bool) + encoder_hidden = self.embed_tokens(input_ids) * self.embed_scale + encoder_hidden = self.encoder(encoder_hidden, encoder_mask) + + self_mask = _causal_mask(decoder_input_ids.shape[1], decoder_input_ids.device) + decoder_hidden = self.embed_tokens(decoder_input_ids) * self.embed_scale + decoder_hidden = self.decoder(decoder_hidden, encoder_hidden, self_mask, encoder_mask) + return BaseModelOutput(last_hidden_state=decoder_hidden) + + +class NeedleForCausalLM(PreTrainedModel): + config_class = NeedleConfig + base_model_prefix = "model" + main_input_name = "input_ids" + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config: NeedleConfig) -> None: + super().__init__(config) + self.model = NeedleModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + self.model.reset_rope() + self.tie_weights() + + def get_encoder(self) -> NeedleEncoder: + return self.model.encoder + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.model.embed_tokens = value + + def get_output_embeddings(self) -> nn.Linear: + return self.lm_head + + def set_output_embeddings(self, value: nn.Linear) -> None: + self.lm_head = value + + def tie_weights(self, *args: Any, **kwargs: Any) -> None: + del args, kwargs + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> Seq2SeqLMOutput: + hidden_states = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + **kwargs, + ).last_hidden_state + return Seq2SeqLMOutput(logits=self.lm_head(hidden_states)) + + def cactus_source_encode( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + base_mask = _padding_mask(input_ids, self.config.pad_token_id) + base_mask = base_mask & attention_mask[:, None, None, :].to(dtype=torch.bool) + encoder_mask = base_mask.expand( + -1, + int(self.config.num_attention_heads), + input_ids.shape[1], + -1, + ).contiguous() + encoder_hidden = self.model.embed_tokens(input_ids) * self.model.embed_scale + encoder_hidden = self.model.encoder(encoder_hidden, encoder_mask) + decoder_mask = base_mask.expand( + -1, + int(self.config.num_attention_heads), + 1, + -1, + ).contiguous() + return encoder_hidden, decoder_mask.to(dtype=encoder_hidden.dtype) + + def cactus_decoder_cross_kv( + self, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + del encoder_attention_mask + outputs: list[torch.Tensor] = [] + for layer in self.model.decoder.layers: + k, v = layer.encoder_attn.project_kv(encoder_hidden_states) + outputs.extend((k, v)) + return tuple(outputs) + + def cactus_decoder_step( + self, + decoder_input_ids: torch.Tensor, + position_ids: torch.Tensor, + encoder_attention_mask: torch.Tensor, + *cross_kv: torch.Tensor, + ) -> torch.Tensor: + encoder_attention_mask = encoder_attention_mask != 0 + hidden_states = self.model.embed_tokens(decoder_input_ids) * self.model.embed_scale + inv_freq = _build_inv_freq( + self.model.decoder.head_dim, + self.model.decoder.rope_theta, + ).to(device=hidden_states.device) + rope = _rotary_tables_for_position_ids( + inv_freq, + position_ids.to(dtype=torch.long), + hidden_states.dtype, + ) + for layer_index, layer in enumerate(self.model.decoder.layers): + normed = layer.input_layernorm(hidden_states) + attn = layer.self_attn(normed, normed, None, rope) + hidden_states = _add_clipped(hidden_states, torch.sigmoid(layer.self_attn_gate).to(dtype=attn.dtype) * attn) + + cross_attn = layer.encoder_attn.forward_with_kv( + layer.encoder_attn_layer_norm(hidden_states), + cross_kv[layer_index * 2], + cross_kv[layer_index * 2 + 1], + encoder_attention_mask, + None, + ) + hidden_states = _add_clipped(hidden_states, torch.sigmoid(layer.cross_attn_gate).to(dtype=cross_attn.dtype) * cross_attn) + hidden_states = self.model.decoder.norm(hidden_states) + return self.lm_head(hidden_states) + + def _init_weights(self, module: nn.Module) -> None: + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + elif isinstance(module, NeedleRMSNorm): + nn.init.zeros_(module.weight) diff --git a/python/cactus/models/needle/tokenization_needle.py b/python/cactus/models/needle/tokenization_needle.py new file mode 100644 index 000000000..8fe3450d2 --- /dev/null +++ b/python/cactus/models/needle/tokenization_needle.py @@ -0,0 +1,121 @@ +"""Slow SentencePiece tokenizer for Needle.""" + +from __future__ import annotations + +import os +import shutil +from typing import Any + +import sentencepiece as spm +from transformers import PreTrainedTokenizer + + +VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"} + + +class NeedleTokenizer(PreTrainedTokenizer): + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file: str, + unk_token: str = "", + bos_token: str = "", + eos_token: str = "", + pad_token: str = "", + tool_call_token: str = "", + tools_token: str = "", + **kwargs: Any, + ) -> None: + self.vocab_file = vocab_file + self.sp_model = spm.SentencePieceProcessor() + self.sp_model.Load(vocab_file) + self.sp = self.sp_model + self.tool_call_token = tool_call_token + self.tools_token = tools_token + additional = list(kwargs.pop("additional_special_tokens", []) or []) + for token in (tool_call_token, tools_token): + if token not in additional: + additional.append(token) + super().__init__( + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + additional_special_tokens=additional, + **kwargs, + ) + + @property + def vocab_size(self) -> int: + return int(self.sp_model.GetPieceSize()) + + @property + def tool_call_token_id(self) -> int: + return int(self.sp_model.PieceToId(self.tool_call_token)) + + @property + def tools_token_id(self) -> int: + return int(self.sp_model.PieceToId(self.tools_token)) + + def get_vocab(self) -> dict[str, int]: + vocab = {self.sp_model.IdToPiece(i): i for i in range(self.vocab_size)} + vocab.update(self.added_tokens_encoder) + return vocab + + def _tokenize(self, text: str) -> list[str]: + return list(self.sp_model.EncodeAsPieces(text)) + + def _convert_token_to_id(self, token: str) -> int: + return int(self.sp_model.PieceToId(token)) + + def _convert_id_to_token(self, index: int) -> str: + return str(self.sp_model.IdToPiece(int(index))) + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + return self.sp_model.DecodePieces(tokens) + + def build_inputs_with_special_tokens( + self, + token_ids_0: list[int], + token_ids_1: list[int] | None = None, + ) -> list[int]: + if token_ids_1 is None: + return list(token_ids_0) + return list(token_ids_0) + list(token_ids_1) + + def get_special_tokens_mask( + self, + token_ids_0: list[int], + token_ids_1: list[int] | None = None, + already_has_special_tokens: bool = False, + ) -> list[int]: + if already_has_special_tokens: + all_ids = list(token_ids_0) + else: + all_ids = self.build_inputs_with_special_tokens(token_ids_0, token_ids_1) + special = { + self.pad_token_id, + self.eos_token_id, + self.bos_token_id, + self.unk_token_id, + self.tool_call_token_id, + self.tools_token_id, + } + return [1 if token_id in special else 0 for token_id in all_ids] + + def create_token_type_ids_from_sequences( + self, + token_ids_0: list[int], + token_ids_1: list[int] | None = None, + ) -> list[int]: + return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1)) + + def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]: + os.makedirs(save_directory, exist_ok=True) + out_name = "tokenizer.model" if filename_prefix is None else f"{filename_prefix}-tokenizer.model" + out_path = os.path.join(save_directory, out_name) + if os.path.abspath(self.vocab_file) != os.path.abspath(out_path): + shutil.copyfile(self.vocab_file, out_path) + return (out_path,) diff --git a/python/cactus/py.typed b/python/cactus/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/server.py b/python/cactus/server.py new file mode 100644 index 000000000..7b38344b1 --- /dev/null +++ b/python/cactus/server.py @@ -0,0 +1,922 @@ +"""OpenAI-compatible local HTTP server for Cactus v2 bundles.""" +from __future__ import annotations + +import asyncio +import base64 +import json +import logging +import math +import os +import re +import tempfile +import time +import uuid +from contextlib import asynccontextmanager +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from typing import Any + +from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile +from fastapi.responses import PlainTextResponse, StreamingResponse +from PIL import Image +from pydantic import BaseModel, Field + +from .bindings.cactus import ( + cactus_complete, + cactus_destroy, + cactus_embed, + cactus_get_last_error, + cactus_init, + cactus_reset, + cactus_transcribe, +) +from .cli.download import get_model_dir_name, get_weights_dir +from .cli.common import DEFAULT_MODEL_ID, is_valid_bundle, weights_root as default_weights_root + +LOGGER = logging.getLogger(__name__) + +LLM_MODEL_TYPES = {"gemma", "gemma3n", "gemma4", "lfm2", "qwen", "qwen3p5", "needle", "youtu"} +STT_MODEL_TYPES = {"whisper", "parakeet_tdt", "parakeet-tdt"} + + +@dataclass(frozen=True) +class ModelInfo: + id: str + path: Path + model_type: str + context_length: int + created: int + supports_embedding: bool = False + + +class ModelRegistry: + def __init__(self, weights_root: Path, extra_model: Path | None = None): + self.weights_root = weights_root + self.models: dict[str, ModelInfo] = {} + self._discover(weights_root) + if extra_model is not None: + info = self._info_for_dir(extra_model) + if info is None: + raise RuntimeError(f"Not a valid v2 Cactus bundle: {extra_model}") + self.models[info.id] = info + + @staticmethod + def _read_config_field(model_dir: Path, field: str) -> str: + config = model_dir / "config.txt" + if not config.exists(): + return "" + prefix = f"{field}=" + for line in config.read_text(encoding="utf-8").splitlines(): + if line.startswith(prefix): + return line.split("=", 1)[1].strip() + return "" + + @staticmethod + def _supports_embedding(model_dir: Path) -> bool: + """A bundle can produce text embeddings if it carries a text_embedding + component (Nomic/BERT encoder) or a decoder_embed_chunk component + (causal-LM hidden-state embeddings).""" + manifest = model_dir / "components" / "manifest.json" + try: + data = json.loads(manifest.read_text(encoding="utf-8")) + except Exception: + return False + names = {str(c.get("component", "")) for c in data.get("components", []) if isinstance(c, dict)} + return "text_embedding" in names or "decoder_embed_chunk" in names + + @classmethod + def _info_for_dir(cls, path: Path) -> ModelInfo | None: + display_id = path.expanduser().name + resolved = path.expanduser().resolve() + if not is_valid_bundle(resolved): + return None + context_raw = cls._read_config_field(resolved, "context_length") + try: + context_length = int(context_raw or 0) + except ValueError: + context_length = 0 + stat = resolved.stat() + return ModelInfo( + id=display_id, + path=resolved, + model_type=cls._read_config_field(resolved, "model_type"), + context_length=context_length, + created=int(stat.st_mtime), + supports_embedding=cls._supports_embedding(resolved), + ) + + def _discover(self, root: Path) -> None: + if not root.exists(): + return + for entry in sorted(root.iterdir()): + if not entry.is_dir(): + continue + info = self._info_for_dir(entry) + if info is not None: + self.models[info.id] = info + + def _resolve(self, model_id: str) -> ModelInfo | None: + """Match an exact bundle id, or fall back to the HuggingFace id / bare + stem (e.g. 'google/gemma-4-E2B-it' -> 'gemma-4-e2b-it-cq4') when it + is unambiguous.""" + info = self.models.get(model_id) + if info is not None: + return info + stem = get_model_dir_name(model_id) + if stem in self.models: + return self.models[stem] + variants = [m for m in self.models.values() if re.match(rf"{re.escape(stem)}-cq\d", m.id)] + return variants[0] if len(variants) == 1 else None + + def require(self, model_id: str) -> ModelInfo: + info = self._resolve(model_id) + if info is None: + available = ", ".join(sorted(self.models)) or "none" + raise HTTPException( + status_code=404, + detail=f"Model '{model_id}' is not available (available: {available})", + ) + return info + + def default_llm(self, preferred: str | None = None) -> ModelInfo: + if preferred: + info = self._resolve(preferred) + if info is None: + raise RuntimeError(f"Requested model '{preferred}' is not a valid v2 Cactus bundle") + if info.model_type not in LLM_MODEL_TYPES: + raise RuntimeError(f"Requested model '{preferred}' is not an LLM bundle") + return info + preferred_id = get_weights_dir(DEFAULT_MODEL_ID).name + if preferred_id in self.models and self.models[preferred_id].model_type in LLM_MODEL_TYPES: + return self.models[preferred_id] + for info in sorted(self.models.values(), key=lambda x: x.id): + if info.model_type in LLM_MODEL_TYPES: + return info + raise RuntimeError("No valid LLM bundles found. Prepare a transpiled bundle before running `cactus serve`.") + + def default_stt(self) -> ModelInfo | None: + for info in sorted(self.models.values(), key=lambda x: x.id): + if info.model_type in STT_MODEL_TYPES: + return info + return None + + def list_openai_models(self) -> list[dict[str, Any]]: + out = [] + for info in sorted(self.models.values(), key=lambda x: x.id): + out.append({ + "id": info.id, + "object": "model", + "created": info.created, + "owned_by": "cactus", + "context_window": info.context_length, + "model_type": info.model_type, + }) + return out + + +class _ModelSlot: + def __init__(self, info: ModelInfo, handle): + self.info = info + self.handle = handle + self.lock = asyncio.Lock() + self.last_used = time.monotonic() + self.active_requests = 0 + + def touch(self) -> None: + self.last_used = time.monotonic() + + +class ModelManager: + def __init__(self, registry: ModelRegistry, *, max_warm: int = 2): + self.registry = registry + self.max_warm = max_warm + self.slots: dict[str, _ModelSlot] = {} + self.swap_lock = asyncio.Lock() + + def _load(self, info: ModelInfo): + try: + return cactus_init(str(info.path)) + except RuntimeError as exc: + err = cactus_get_last_error() or str(exc) + raise HTTPException(status_code=500, detail=f"Failed to load model '{info.id}': {err}") from exc + + async def _get_slot(self, model_id: str) -> _ModelSlot: + info = self.registry.require(model_id) + async with self.swap_lock: + slot = self.slots.get(info.id) + if slot is not None: + slot.touch() + return slot + + if len(self.slots) >= self.max_warm: + idle = [s for s in self.slots.values() if s.active_requests == 0 and not s.lock.locked()] + if not idle: + raise HTTPException(status_code=503, detail="All warm model slots are busy") + victim = min(idle, key=lambda s: s.last_used) + self.slots.pop(victim.info.id, None) + cactus_destroy(victim.handle) + + handle = await asyncio.get_running_loop().run_in_executor(None, self._load, info) + slot = _ModelSlot(info, handle) + self.slots[info.id] = slot + return slot + + @asynccontextmanager + async def acquire(self, model_id: str): + slot = await self._get_slot(model_id) + async with self.swap_lock: + if slot.info.id not in self.slots: + raise HTTPException(status_code=503, detail="Model slot was evicted before use") + slot.active_requests += 1 + try: + yield slot + finally: + async with self.swap_lock: + slot.active_requests = max(0, slot.active_requests - 1) + slot.touch() + + async def preload(self, model_id: str) -> None: + async with self.acquire(model_id): + return + + def shutdown(self) -> None: + for slot in self.slots.values(): + cactus_destroy(slot.handle) + self.slots.clear() + + +class Permissive(BaseModel): + model_config = {"extra": "allow"} + + +class ChatMessage(Permissive): + role: str + content: str | list[Any] | None = None + tool_call_id: str | None = None + tool_calls: list[Any] | None = None + + +class ToolFunction(Permissive): + name: str + description: str | None = None + parameters: dict[str, Any] | None = None + + +class Tool(Permissive): + type: str = "function" + function: ToolFunction + + +class ToolChoiceFunction(Permissive): + name: str + + +class ToolChoiceObject(Permissive): + type: str = "function" + function: ToolChoiceFunction + + +class ChatRequest(Permissive): + model: str + messages: list[ChatMessage] = Field(min_length=1) + temperature: float | None = None + top_p: float | None = None + top_k: int | None = None + max_tokens: int | None = None + max_completion_tokens: int | None = None + stop: str | list[str] | None = None + stream: bool = False + reasoning_effort: str | None = None + tools: list[Tool] | None = None + tool_choice: str | ToolChoiceObject | None = None + + +class EmbeddingRequest(Permissive): + model: str + input: str | list[str] + + +_DATA_URI_RE = re.compile(r"^data:(?P[^;,]+)?(?P;base64)?,(?P.*)$", re.DOTALL) + +_ENGINE_IMAGE_FORMATS = {"PNG", "JPEG", "GIF"} + + +def _decodable_image(source: Any) -> bool: + try: + with Image.open(source) as im: + if im.format not in _ENGINE_IMAGE_FORMATS or im.width <= 0 or im.height <= 0: + return False + im.load() + return True + except Exception: + return False + + +def _materialize_image(url: str) -> str | None: + if not isinstance(url, str) or not url: + return None + if url.startswith("data:"): + match = _DATA_URI_RE.match(url) + if not match: + return None + mime = (match.group("mime") or "image/png").lower() + ext = "." + mime.split("/")[-1].split("+")[0] if "/" in mime else ".png" + payload = match.group("data") + try: + raw = base64.b64decode(payload) if match.group("b64") else payload.encode("utf-8") + except ValueError: + return None + if not _decodable_image(BytesIO(raw)): + return None + fd, path = tempfile.mkstemp(suffix=ext, prefix="cactus_img_") + with os.fdopen(fd, "wb") as fh: + fh.write(raw) + return path + if url.startswith("file://"): + url = url[len("file://"):] + if Path(url).exists(): + return url if _decodable_image(url) else None + return None + + +def _flatten_message(msg: ChatMessage) -> dict[str, Any]: + out: dict[str, Any] = {"role": msg.role} + if isinstance(msg.content, list): + parts = [] + images: list[str] = [] + for part in msg.content: + if isinstance(part, dict) and part.get("type") == "text": + parts.append(str(part.get("text", ""))) + elif isinstance(part, dict) and part.get("type") == "image_url": + raw = part.get("image_url") + url = raw.get("url") if isinstance(raw, dict) else (raw if isinstance(raw, str) else part.get("url")) + resolved = _materialize_image(url) if isinstance(url, str) else None + if resolved: + images.append(resolved) + elif isinstance(part, str): + parts.append(part) + out["content"] = "\n".join(parts) + if images: + out["images"] = images + elif msg.content is not None: + out["content"] = msg.content + if msg.tool_call_id is not None: + out["tool_call_id"] = msg.tool_call_id + if msg.tool_calls is not None: + out["tool_calls"] = msg.tool_calls + return out + + +def _translate_tools(tools: list[Tool] | None, tool_choice) -> tuple[list[dict[str, Any]] | None, bool]: + if not tools or tool_choice == "none": + return None, False + translated = [ + {"type": "function", "function": {"name": t.function.name, "description": t.function.description or "", "parameters": t.function.parameters or {}}} + for t in tools + ] + if isinstance(tool_choice, ToolChoiceObject): + selected = [t for t in translated if t["function"]["name"] == tool_choice.function.name] + return selected or None, True + return translated, tool_choice == "required" + + +def _make_tool_calls(function_calls: list[Any], *, with_index: bool = False) -> list[dict[str, Any]]: + out = [] + for i, call in enumerate(function_calls): + if not isinstance(call, dict): + continue + args = call.get("arguments", {}) + entry: dict[str, Any] = { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": call.get("name", ""), + "arguments": json.dumps(args) if isinstance(args, dict) else str(args), + }, + } + if with_index: + entry["index"] = len(out) + out.append(entry) + return out + + +def _clamp_finite(value: Any, lo: float, hi: float) -> float | None: + """Clamp a sampling param to a finite, sane range so extreme/non-finite + values (e.g. 1e308, inf, nan) can't overflow the native sampler.""" + if value is None: + return None + try: + v = float(value) + except (TypeError, ValueError): + return None + if not math.isfinite(v): + return None + return max(lo, min(hi, v)) + + +def _chat_options(req: ChatRequest) -> dict[str, Any]: + options: dict[str, Any] = {} + temperature = _clamp_finite(req.temperature, 0.0, 2.0) + options["temperature"] = temperature if temperature is not None else 0.7 + top_p = _clamp_finite(req.top_p, 0.0, 1.0) + if top_p is not None: + options["top_p"] = top_p + if req.top_k is not None: + options["top_k"] = max(0, min(int(req.top_k), 1 << 20)) + max_tokens = req.max_tokens if req.max_tokens is not None else req.max_completion_tokens + options["max_tokens"] = max_tokens if max_tokens is not None else 4096 + if req.stop: + options["stop_sequences"] = [req.stop] if isinstance(req.stop, str) else req.stop + if req.reasoning_effort and req.reasoning_effort.lower() not in ("none", "off"): + options["enable_thinking_if_supported"] = True + return options + + +_HARMONY_MARKER = "<|channel" +_HARMONY_THOUGHT_CHANNELS = {"thought", "thinking", "analysis", "reflection"} +_HARMONY_CHANNEL_RE = re.compile(r"<\|channel\|?>\s*([a-zA-Z_]+)\s*\n?") +_HARMONY_CONTROL_RE = re.compile(r"<\|[^>]*>") + + +def _split_harmony_channels(text: str) -> tuple[str, str | None]: + if not text or _HARMONY_MARKER not in text: + return text, None + segments = _HARMONY_CHANNEL_RE.split(text) + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + if segments[0].strip(): + content_parts.append(segments[0]) + for i in range(1, len(segments) - 1, 2): + body = segments[i + 1] + target = reasoning_parts if segments[i].lower() in _HARMONY_THOUGHT_CHANNELS else content_parts + target.append(body) + content = _HARMONY_CONTROL_RE.sub("", "".join(content_parts)).strip() + reasoning = _HARMONY_CONTROL_RE.sub("", "\n".join(reasoning_parts)).strip() or None + if not content and reasoning: + return reasoning, None + return content, reasoning + + +class _HarmonyStreamSplitter: + def __init__(self) -> None: + self._buf = "" + self._channel = "final" + + def _kind(self) -> str: + return "reasoning" if self._channel in _HARMONY_THOUGHT_CHANNELS else "content" + + def _hold_partial(self, s: str) -> str: + for k in range(min(len(s), len(_HARMONY_MARKER) - 1), 0, -1): + if _HARMONY_MARKER.startswith(s[-k:]): + return s[:-k] + return s + + def feed(self, text: str) -> list[tuple[str, str]]: + self._buf += text + out: list[tuple[str, str]] = [] + while True: + idx = self._buf.find(_HARMONY_MARKER) + if idx == -1: + safe = self._hold_partial(self._buf) + if safe: + out.append((self._kind(), _HARMONY_CONTROL_RE.sub("", safe))) + self._buf = self._buf[len(safe):] + break + before = self._buf[:idx] + if before: + out.append((self._kind(), _HARMONY_CONTROL_RE.sub("", before))) + rest = self._buf[idx + len(_HARMONY_MARKER):] + m = re.match(r"\|?>\s*([a-zA-Z_]+)(?:[ \t]*\n|(?=<))", rest) + if not m: + self._buf = self._buf[idx:] + break + self._channel = m.group(1).lower() + self._buf = rest[m.end():] + return out + + def flush(self) -> list[tuple[str, str]]: + if not self._buf: + return [] + out = [(self._kind(), _HARMONY_CONTROL_RE.sub("", self._buf))] + self._buf = "" + return out + +_TOOLCALL_PAIRS = ( + ("<|tool_call>", ""), + ("", ""), + ("", ""), + ("<|tool_call_start|>", "<|tool_call_end|>"), +) +_TOOLCALL_OPENERS = tuple(o for o, _ in _TOOLCALL_PAIRS) + + +def _strip_tool_call_markers(text: str) -> str: + if not text or ("call:" not in text and "tool_call" not in text + and "function_call" not in text and '<|"|>' not in text): + return text + s = text + for open_m, close_m in _TOOLCALL_PAIRS: + while True: + ci = s.find(close_m) + if ci == -1: + break + oi = s.rfind(open_m, 0, ci) + if oi == -1: + oi = s.rfind("call:", 0, ci) + if oi == -1: + s = s[:ci] + s[ci + len(close_m):] + else: + s = s[:oi] + s[ci + len(close_m):] + cut = None + for open_m in _TOOLCALL_OPENERS: + oi = s.find(open_m) + if oi != -1: + cut = oi if cut is None else min(cut, oi) + if cut is not None: + s = s[:cut] + ci = s.find("call:") + while ci != -1: + j = ci + 5 + k = j + while k < len(s) and (s[k].isalnum() or s[k] in "_-."): + k += 1 + m = k + while m < len(s) and s[m] == " ": + m += 1 + if k > j and m < len(s) and s[m] == "{": + s = s[:ci] + break + ci = s.find("call:", ci + 5) + s = s.replace('<|"|>', '"') + return s.strip() if s != text else text + + +def _build_chat_response(result: dict[str, Any], model_id: str, request_id: str) -> dict[str, Any]: + function_calls = result.get("function_calls") or [] + tool_calls = _make_tool_calls(function_calls) + has_tool_calls = bool(tool_calls) + prefill = int(result.get("prefill_tokens") or 0) + decode = int(result.get("decode_tokens") or 0) + content, reasoning = _split_harmony_channels(result.get("response", "") or "") + content = _strip_tool_call_markers(content) + message: dict[str, Any] = {"role": "assistant", "content": None if has_tool_calls else content} + if reasoning: + message["reasoning_content"] = reasoning + if has_tool_calls: + message["tool_calls"] = tool_calls + return { + "id": request_id, + "object": "chat.completion", + "created": int(time.time()), + "model": model_id, + "choices": [{ + "index": 0, + "message": message, + "logprobs": None, + "finish_reason": "tool_calls" if has_tool_calls else "stop", + }], + "usage": { + "prompt_tokens": prefill, + "completion_tokens": decode, + "total_tokens": prefill + decode, + }, + "cloud_handoff": bool(result.get("cloud_handoff", False)), + } + + +def _event(data: dict[str, Any]) -> str: + return f"data: {json.dumps(data)}\n\n" + + +async def _stream_completion(manager: ModelManager, req: ChatRequest, request_id: str, messages, options, tools): + queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue() + loop = asyncio.get_running_loop() + + def on_token(token: str, token_id: int) -> None: + loop.call_soon_threadsafe(queue.put_nowait, ("token", token)) + + async def run_inference(): + error = None + result = None + try: + async with manager.acquire(req.model) as slot: + async with slot.lock: + cactus_reset(slot.handle) + result = await loop.run_in_executor( + None, + lambda: cactus_complete(slot.handle, messages, options, tools, on_token), + ) + except Exception as exc: + LOGGER.exception("Streaming completion failed") + error = exc + finally: + loop.call_soon_threadsafe(queue.put_nowait, ("done", (result, error))) + + task = asyncio.create_task(run_inference()) + created = int(time.time()) + yield _event({ + "id": request_id, + "object": "chat.completion.chunk", + "created": created, + "model": req.model, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "logprobs": None, "finish_reason": None}], + }) + + result = None + error = None + splitter = _HarmonyStreamSplitter() + + def _delta_chunk(piece_kind: str, text: str) -> str: + field = "reasoning_content" if piece_kind == "reasoning" else "content" + return _event({ + "id": request_id, + "object": "chat.completion.chunk", + "created": created, + "model": req.model, + "choices": [{"index": 0, "delta": {field: text}, "logprobs": None, "finish_reason": None}], + }) + + tools_present = bool(tools) + + def _consume(pieces: list[tuple[str, str]]) -> list[str]: + out: list[str] = [] + for piece_kind, piece_text in pieces: + if not piece_text: + continue + if tools_present and piece_kind != "reasoning": + continue + out.append(_delta_chunk(piece_kind, piece_text)) + return out + + while True: + kind, value = await queue.get() + if kind == "token": + for event in _consume(splitter.feed(value)): + yield event + elif kind == "done": + result, error = value + break + + for event in _consume(splitter.flush()): + yield event + + await task + if error is not None: + # OpenAI streaming spec embeds errors in a `data:` line with an `error` + # object; clients (openai-python, vercel ai) parse `event:` lines as + # comments. We also emit a final chunk so the client gets a + # finish_reason before [DONE] and never hangs. + yield f"data: {json.dumps({'error': {'message': str(error), 'type': 'server_error'}})}\n\n" + yield _event({ + "id": request_id, + "object": "chat.completion.chunk", + "created": created, + "model": req.model, + "choices": [{"index": 0, "delta": {}, "logprobs": None, "finish_reason": "stop"}], + }) + yield "data: [DONE]\n\n" + return + + function_calls = (result or {}).get("function_calls") or [] + tool_calls = _make_tool_calls(function_calls, with_index=True) + finish_reason = "tool_calls" if tool_calls else "stop" + if tools_present and not tool_calls: + clean_content, _ = _split_harmony_channels((result or {}).get("response", "") or "") + clean_content = _strip_tool_call_markers(clean_content) + if clean_content: + yield _delta_chunk("content", clean_content) + if tool_calls: + yield _event({ + "id": request_id, + "object": "chat.completion.chunk", + "created": created, + "model": req.model, + "choices": [{"index": 0, "delta": {"tool_calls": tool_calls}, "logprobs": None, "finish_reason": None}], + }) + final = { + "id": request_id, + "object": "chat.completion.chunk", + "created": created, + "model": req.model, + "choices": [{"index": 0, "delta": {}, "logprobs": None, "finish_reason": finish_reason}], + "usage": { + "prompt_tokens": int((result or {}).get("prefill_tokens") or 0), + "completion_tokens": int((result or {}).get("decode_tokens") or 0), + "total_tokens": int((result or {}).get("total_tokens") or 0), + }, + "cloud_handoff": bool((result or {}).get("cloud_handoff", False)), + } + yield _event(final) + yield "data: [DONE]\n\n" + + +def _requested_granularities(primary: list[str] | None, bracketed: list[str] | None) -> list[str]: + values = [] + for item in (primary or []) + (bracketed or []): + if item: + values.append(item) + return values + + +def create_app( + *, + weights_root: Path | None = None, + model_path: Path | None = None, + default_model: str | None = None, + max_warm: int = 2, + preload: bool = True, + auto_handoff: bool = True, + confidence_threshold: float | None = None, + cloud_timeout_ms: int | None = None, +) -> FastAPI: + root = Path(weights_root) if weights_root is not None else default_weights_root() + registry = ModelRegistry(root, extra_model=model_path) + if default_model is not None: + selected = registry.models.get(default_model) + if selected is None: + raise RuntimeError(f"Requested model '{default_model}' is not a valid v2 Cactus bundle") + else: + try: + selected = registry.default_llm() + except RuntimeError: + # No LLM available — allow serving a non-LLM bundle (e.g. an embedding model). + available = sorted(registry.models.values(), key=lambda info: info.id) + if not available: + raise + selected = available[0] + manager = ModelManager(registry, max_warm=max_warm) + + @asynccontextmanager + async def lifespan(app: FastAPI): + if preload: + await manager.preload(selected.id) + try: + yield + finally: + manager.shutdown() + + app = FastAPI(title="Cactus", version="0.1.0", lifespan=lifespan) + app.state.registry = registry + app.state.manager = manager + app.state.default_model = selected.id + app.state.default_stt_model = registry.default_stt().id if registry.default_stt() else None + app.state.auto_handoff = auto_handoff + app.state.confidence_threshold = confidence_threshold + app.state.cloud_timeout_ms = cloud_timeout_ms + + @app.get("/v1/models") + async def list_models(request: Request): + reg: ModelRegistry = request.app.state.registry + return {"object": "list", "data": reg.list_openai_models()} + + @app.post("/v1/chat/completions") + async def chat_completions(request: Request, req: ChatRequest): + reg: ModelRegistry = request.app.state.registry + info = reg.require(req.model) + if info.model_type not in LLM_MODEL_TYPES: + raise HTTPException(status_code=400, detail=f"Model '{req.model}' is not an LLM model") + messages = [_flatten_message(m) for m in req.messages] + tools, force_tools = _translate_tools(req.tools, req.tool_choice) + options = _chat_options(req) + if "max_tokens" not in options: + options["max_tokens"] = info.context_length if info.context_length > 0 else 8192 + state = request.app.state + if not state.auto_handoff: + options["auto_handoff"] = False + if state.confidence_threshold is not None: + options["confidence_threshold"] = state.confidence_threshold + if state.cloud_timeout_ms is not None: + options["cloud_timeout_ms"] = state.cloud_timeout_ms + if force_tools: + options["force_tools"] = True + request_id = f"chatcmpl-{uuid.uuid4().hex[:29]}" + mgr: ModelManager = request.app.state.manager + if req.stream: + return StreamingResponse( + _stream_completion(mgr, req, request_id, messages, options or None, tools), + media_type="text/event-stream", + ) + async with mgr.acquire(req.model) as slot: + async with slot.lock: + cactus_reset(slot.handle) + result = await asyncio.get_running_loop().run_in_executor( + None, + lambda: cactus_complete(slot.handle, messages, options or None, tools, None), + ) + if not result.get("success", False): + raise HTTPException(status_code=500, detail=result.get("error") or "completion failed") + return _build_chat_response(result, req.model, request_id) + + @app.post("/v1/audio/transcriptions") + async def create_transcription( + request: Request, + file: UploadFile = File(...), + model: str = Form(""), + language: str | None = Form(None), + prompt: str | None = Form(None), + response_format: str = Form("json"), + temperature: float | None = Form(None), + timestamp_granularities: list[str] | None = Form(None), + timestamp_granularities_array: list[str] | None = Form(None, alias="timestamp_granularities[]"), + ): + if response_format not in {"json", "text", "verbose_json"}: + raise HTTPException(status_code=400, detail=f"Unsupported transcription response_format: {response_format}") + granularities = _requested_granularities(timestamp_granularities, timestamp_granularities_array) + if "word" in granularities: + raise HTTPException(status_code=400, detail="Word timestamp granularity is not supported") + if any(g != "segment" for g in granularities): + raise HTTPException(status_code=400, detail=f"Unsupported timestamp granularity: {', '.join(granularities)}") + suffix = Path(file.filename or "").suffix.lower() + if suffix != ".wav": + raise HTTPException(status_code=400, detail="Only .wav transcription uploads are supported for now") + model_id = model or request.app.state.default_stt_model + if not model_id: + raise HTTPException(status_code=400, detail="No STT model available") + reg: ModelRegistry = request.app.state.registry + info = reg.require(model_id) + if info.model_type not in STT_MODEL_TYPES: + raise HTTPException(status_code=400, detail=f"Model '{model_id}' is not a speech-to-text model") + import tempfile + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp.write(await file.read()) + tmp_path = Path(tmp.name) + try: + options: dict[str, Any] = {} + if temperature is not None: + options["temperature"] = temperature + if language: + options["language"] = language + if response_format == "verbose_json": + options["timestamps"] = True + mgr: ModelManager = request.app.state.manager + async with mgr.acquire(model_id) as slot: + async with slot.lock: + cactus_reset(slot.handle) + result = await asyncio.get_running_loop().run_in_executor( + None, + lambda: cactus_transcribe(slot.handle, str(tmp_path), prompt, options or None, None), + ) + finally: + tmp_path.unlink(missing_ok=True) + if not result.get("success", False): + raise HTTPException(status_code=500, detail=result.get("error") or "transcription failed") + text = result.get("response", "") + if response_format == "text": + return PlainTextResponse(text) + if response_format == "verbose_json": + segments = result.get("segments") or [] + return { + "task": "transcribe", + "language": language or "", + "duration": float(segments[-1].get("end", 0.0)) if segments else 0.0, + "text": text, + "segments": [ + {"id": i, + "start": float(seg.get("start", 0.0)), + "end": float(seg.get("end", 0.0)), + "text": seg.get("text", "")} + for i, seg in enumerate(segments) + ], + } + return {"text": text} + + @app.post("/v1/embeddings") + async def create_embeddings(request: Request, req: EmbeddingRequest): + reg: ModelRegistry = request.app.state.registry + info = reg.require(req.model) + if not info.supports_embedding: + raise HTTPException(status_code=400, detail=f"Model '{req.model}' is not an embedding model") + inputs = [req.input] if isinstance(req.input, str) else list(req.input) + if not inputs or any(not text for text in inputs): + raise HTTPException(status_code=400, detail="'input' must not be empty") + mgr: ModelManager = request.app.state.manager + + def _embed_all(handle) -> list[list[float]]: + vectors = [] + for text in inputs: + cactus_reset(handle) + vectors.append(cactus_embed(handle, text, True)) + return vectors + + async with mgr.acquire(req.model) as slot: + async with slot.lock: + try: + vectors = await asyncio.get_running_loop().run_in_executor( + None, lambda: _embed_all(slot.handle) + ) + except Exception as exc: + raise HTTPException(status_code=500, detail=str(exc)) + data = [ + {"object": "embedding", "index": i, "embedding": vector} + for i, vector in enumerate(vectors) + ] + return { + "object": "list", + "data": data, + "model": req.model, + "usage": {"prompt_tokens": 0, "total_tokens": 0}, + } + + return app diff --git a/python/cactus/transpile/README.md b/python/cactus/transpile/README.md new file mode 100644 index 000000000..978a57e4b --- /dev/null +++ b/python/cactus/transpile/README.md @@ -0,0 +1,1022 @@ +# Cactus Transpiler + +This folder contains the Python side of the Cactus transpiler: the code that takes +a Hugging Face or PyTorch model, captures it with `torch.export`, converts that +export into Cactus' intermediate representation (`IRGraph`), canonicalizes and +optimizes the IR, lowers it into a runtime `Graph`, and optionally saves or runs +the result. + +The transpiler has two main output shapes: + +- A single lowered graph for simple tasks such as text logits or encoder-only audio. +- A component bundle for models that are easier to split into staged graphs, most + notably Gemma4 multimodal and Parakeet TDT. + +This README is a code-first walkthrough of the actual flow in the repo today. + +## Related Entry Points + +These files are outside `python/cactus/transpile`, but they are the top-level +entry points for the transpiler (driven by `cactus convert` for end users; the +`cactus transpile` command is hidden from the CLI menu and used internally): + +- `python/cactus/cli/convert.py` (`cmd_convert` → `cmd_transpile`) +- `python/cactus/cli/transpile.py` (`run_transpile`) +- `python/cactus/transpile/hf_model.py` + +There is no separate `cactus transpile` subcommand. `cactus convert ` +quantizes the weights to CQ and then builds the runtime graph in one step; the graph +build is `run_transpile`, a thin wrapper around `python/cactus/transpile/hf_model.py`. +The packaged module is the real top-level transpile program. + +## CLI Entry Flow + +### `cactus convert` + +`cmd_convert` (`python/cactus/cli/convert.py`) quantizes the weights via `ensure_weights`, +then (unless `--weights-only`) calls `cmd_transpile`, which forwards options to +`run_transpile`: + +```python +def run_transpile(model_id, *, extra_args, execute_after_transpile=False, ...): + transpile_lib = _ensure_python_runtime_library() + model_id = resolve_model_id_alias(model_id) + command = [sys.executable, "-m", "cactus.transpile.hf_model", "--model-id", model_id] + if "--weights-dir" not in extra_args: + default_weights_dir = get_weights_dir(model_id) + if _weights_dir_looks_transpile_ready(default_weights_dir): + command.extend(["--weights-dir", str(default_weights_dir)]) + else: + return 1 + if not execute_after_transpile and "--skip-execute" not in extra_args: + command.append("--skip-execute") + command.extend(extra_args) +``` + +Two important details: + +- The graph build saves artifacts and skips execution by default + (pass `--execute-after-transpile` to run a reference execution). +- Building the runtime graph requires converted Cactus CQ weights. `cactus convert` + produces them before this stage; or pass `--weights-dir` to an existing converted folder. + +### `cactus run` + +`cactus run ` is the single user-facing entry. Resolution +order (`python/cactus/cli/model.py:ensure_runnable_bundle`): + +1. `` is an existing path with `components/manifest.json` → run as-is. +2. Cached download at `weights/-cq[-]/` → run. +3. Fresh download from `huggingface.co/Cactus-Compute` via `download_bundle`. +4. Fallback: local `ensure_bundle` (convert + transpile for custom models not on HF). + +Bundle handoff is to the native `run` binary (`python/cactus/bin/run`), which +loads the bundle via `cactus_init()`. + +## End-To-End Timeline + +## 1. Parse Args, Infer Task, Load Model Bundle + +The real transpile program lives in `python/cactus/transpile/hf_model.py`. +Its `main()` function decides what kind of model/task it is dealing with, validates +optional converted weights, and loads the HF model plus tokenizer/processor: + +```python +if args.task == "auto": + inferred_task = _infer_task_from_config(args.model_id) + ... + if image_files or (args.audio_file and has_multimodal_config and inferred_task == "causal_lm_logits"): + task = "multimodal_causal_lm_logits" + else: + task = inferred_task +else: + task = args.task + +validated_weights_dir = _validate_weights_dir(args.weights_dir.strip() or None, model_id=args.model_id) +weights_dir = str(validated_weights_dir) if validated_weights_dir is not None else None + +model_source, processor_or_tokenizer, model, model_config = _load_transformers_bundle( + model_id=args.model_id, + task=task, + torch_dtype=torch_dtype, + token=args.token, + trust_remote_code=args.trust_remote_code, + local_files_only=args.local_files_only, +) +``` + +### What decides the task? + +`_infer_task_from_config()` looks at: + +- `config.json` architectures +- model type +- known special cases like Parakeet TDT and Whisper +- model id substrings as a last fallback + +Today the main task shapes are: + +- `causal_lm_logits` +- `multimodal_causal_lm_logits` +- `ctc_logits` +- `encoder_hidden_states` +- `tdt_transcription` + +## 2. Prepare Example Inputs + +Before capture, the transpiler builds one concrete example input batch. That +example batch drives both `torch.export` capture and later optional execution. + +The high-level branches are: + +- `_prepare_text_inputs()` for text-only causal LM capture +- `_prepare_audio_inputs()` for speech/audio tasks +- `_prepare_gemma4_multimodal_inputs()` for Gemma4 multimodal + +Example: + +```python +if task == "causal_lm_logits": + prepared = _prepare_text_inputs(...) +elif task == "multimodal_causal_lm_logits": + prepared = _prepare_gemma4_multimodal_inputs(...) +elif task == "tdt_transcription": + prepared = _prepare_audio_inputs(...) +else: + prepared = _prepare_audio_inputs(...) +``` + +The `PreparedInputs` object carries: + +- `names`: logical input names like `input_ids` or `input_features` +- `tensors`: the actual example tensors +- `metadata`: prompt/audio/image/sample-rate info that later gets written into + manifests + +## 3. Canonicalize The Model Interface + +The transpiler does not try to capture every model's native HF forward signature +as-is. It first wraps the model in an adapter that exposes a smaller, more stable +task-specific interface. + +That happens in `model_adapters.py`: + +```python +def canonicalize_model_interface(model, task="causal_lm_logits", *, input_names=None, weights_dir=None): + family = _family_key(model) + ... + if task == "causal_lm_logits": + ... + elif task == "multimodal_causal_lm_logits": + adapter_factory = lambda inner_model: Gemma4MultimodalCausalLMLogitsAdapter( + inner_model, + input_names=resolved_input_names, + weights_dir=weights_dir, + ) + elif task == "ctc_logits": + adapter_factory = lambda inner_model: CTCLogitsAdapter(...) + elif task == "encoder_hidden_states": + adapter_factory = lambda inner_model: EncoderHiddenStatesAdapter(...) +``` + +This step is where family-specific knowledge lives: + +- Gemma/Gemma3/Gemma4/Qwen causal LM wrappers +- Gemma4 multimodal wrappers +- generic CTC and encoder wrappers + +The adapter is also where graph metadata and import hints get attached. + +## 3b. Import PyTorch Ops, Not Model Layer Names + +The compiler boundary is the `torch.export` graph. After capture, Cactus imports +`call_function` nodes by their canonical ATen operation name, not by Hugging Face +module names. For example, `torch.ops.aten.linear.default`, `aten.linear.default`, +and equivalent exported spellings normalize through `aten_ops.py` to the same +canonical IR op: + +```python +aten.linear.default -> linear +aten.silu.default -> silu +aten.conv1d.default -> conv1d +aten.rms_norm.default -> rms_norm +``` + +This is the main generality boundary: model wrappers may still prepare inputs or +split components, but actual compute lowering is ATen op / pattern based. Layer +names are allowed as structural hints for component boundaries and exact +converted-weight manifest aliases; they should not be required to decide how a +PyTorch op maps to a Cactus op. + +## 4. Decide Between Monolithic Capture And Component Capture + +After input prep, the script asks `build_component_module_specs()` whether the +current model supports a split component pipeline: + +```python +component_specs = build_component_module_specs( + model, + task=task, + named_tensors=_named_tensor_store(prepared), + weights_dir=weights_dir, +) +``` + +Today that returns component specs for: + +- Gemma4 multimodal: `vision_encoder`, `audio_encoder`, `lm_encoder`, `decoder` +- Parakeet TDT: `audio_encoder`, `decoder` + +If component specs exist and `--component-pipeline` is `auto` or `on`, the script +takes the component path. Otherwise it captures a single wrapped model graph. + +## 5A. Component Pipeline Path + +The component path is driven by `_run_component_pipeline_transpile()` in +`python/cactus/transpile/hf_model.py`: + +```python +for spec in component_specs: + captured_components[spec.component] = capture_component_spec(spec, fusion_config=fusion_config) +``` + +Each component spec is captured by `component_pipeline.capture_component_spec()`: + +```python +captured = capture_model(wrapped, spec.example_inputs) +raw_ir_graph = copy.deepcopy(captured.ir_graph) +optimized_ir_graph = copy.deepcopy(captured.ir_graph) +canonicalize_exported_graph(optimized_ir_graph) +optimize_graph(optimized_ir_graph, config=fusion_config) +transpiled_graph = transpile_preoptimized_ir(copy.deepcopy(optimized_ir_graph)) +``` + +So each component goes through the same capture/import/canonicalize/optimize/lower +pipeline independently. + +For Gemma4 multimodal, `model_adapters.py` builds the specs by first computing +example image/audio features, then generating example inputs for the downstream +`lm_encoder` and `decoder` stages. + +## 5B. Monolithic Capture Path + +The non-component path wraps the canonical adapter in `TranspileWrapper` and runs +the full graph through capture: + +```python +wrapper = TranspileWrapper(canonical.module, weights_dir=weights_dir).eval() + +captured = capture_model(wrapper, prepared.tensors) +raw_ir_graph = copy.deepcopy(captured.ir_graph) + +canonicalize_exported_graph(captured.ir_graph) +optimize_graph(captured.ir_graph, config=fusion_config) +tg = _lower_preoptimized_ir(captured.ir_graph) +``` + +`TranspileWrapper` mainly exists to pass `weights_dir` into the graph metadata: + +```python +def get_transpile_metadata(self) -> dict[str, object]: + ... + if self.weights_dir: + graph_meta["weights_dir"] = self.weights_dir +``` + +## 6. Capture With `torch.export` + +`capture_pytorch.py` owns the `torch.export` boundary: + +```python +normalized_kwargs = _inject_export_safe_kwargs(model, _normalize_kwargs(kwargs)) +... +transpile_metadata = _collect_transpile_metadata(model, example_args, example_kwargs) + +ep = export(model, args=example_args, kwargs=example_kwargs, strict=strict) +... +ir_graph = import_captured_to_ir(raw_captured, strict=strict) +``` + +Important behavior here: + +- HF-style `use_cache=False` and `return_dict=False` are injected automatically + when the model signature supports them. +- Metadata providers are collected from any module that defines + `get_transpile_metadata()` or `transpile_metadata()`. +- The captured export is immediately imported into the repo's own IR. + +The capture output is a `CapturedModel` dataclass holding: + +- the exported program +- the FX graph +- the lifted state dict +- the imported `IRGraph` +- the example inputs +- the collected metadata + +## 7. Import Exported FX Into `IRGraph` + +`import_ir.py` converts the exported FX/Export graph into the repo's intermediate +representation: + +```python +ir = IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[], constants={}, meta=dict(graph_meta)) +ctx = ImportContext(strict=strict, transpile_metadata=transpile_metadata) +weights_dir = resolve_transpile_weights_dir(ir.meta) + +for node in captured.graph.nodes: + if node.op == "placeholder": + ... + if node.op == "get_attr": + ... + if node.op == "call_function": + import_call_function(...) + if node.op == "output": + import_output(ir, node, ctx) + +apply_import_semantics(ir) +annotate_ir_components(ir) +verify_ir(ir) +``` + +### What gets imported? + +- Real graph inputs become `IRGraph.inputs` +- lifted parameters/buffers/constants become `IRGraph.constants` +- FX `call_function` nodes become `IRNode`s via `importers.py` + +### How weight binding metadata enters the IR + +When a captured constant corresponds to a known converted weight file, +`import_captured_to_ir()` resolves that mapping early: + +```python +binding = resolve_weight_binding(weights_dir=weights_dir, source_name=target) +... +ir.values[value_id_str].meta.update( + { + "path": binding.path, + "kind": binding.kind, + "source_name": binding.source_name, + } +) +``` + +That metadata is what later lets lowering use `g.mmap_weights()` instead of +embedding the constant payload into the graph. Embedding lookups still lower to +`embedding_from_tensor()`, but the backing table is represented as a normal mmap +weight tensor so tied output heads can reuse the same constant safely. +By default this resolver only trusts explicit `weights_manifest.json` entries +written by `cactus convert` plus wrapper-local aliases derived from those entries. +It does not guess old model-specific filenames from layer names. + +## 8. Importers: FX Target To Canonical IR Node + +`importers.py` is the large target-dispatch layer. It normalizes FX/PyTorch targets +to canonical op names and creates `IRNode`s with stable attributes. + +The dispatch point is: + +```python +def import_call_function(ir, node, ctx, *, shape, dtype, torch_op): + op = normalize_target(torch_op) + importer = OP_IMPORTERS.get(op) + if importer is None: + import_opaque_call_function(...) + else: + importer(...) +``` + +Import is intentionally ATen/op based. It does not apply module-path hints such +as `model.layers.N.self_attn`; model-specific behavior should live in explicit +preprocessing/component boundaries or in topology-based fusion passes. + +Some notable importer behaviors: + +- `normalize.py` reduces many `aten.*` spellings into one canonical name. +- Unknown ops are imported as `kind="opaque"` IR nodes rather than being silently + dropped. +- Import hints can inject attrs/meta into matching nodes after import. +- Some nested exported subgraphs such as `wrap_with_set_grad_enabled` are inlined + explicitly in `import_ir.py`. + +## 9. Import-Time Semantic Rewrites + +Immediately after raw import, `import_semantics.py` performs a first semantic pass: + +```python +def apply_import_semantics(graph: IRGraph) -> IRGraph: + _rewrite_explicit_attention(graph) + _tag_explicit_linear(graph) + _rewrite_early_rope(graph) + rebuild_graph(graph) + verify_ir(graph) + return graph +``` + +This step is intentionally early. It rewrites export patterns into higher-level +meaning before the main optimizer sees them. + +Examples: + +- `scaled_dot_product_attention` becomes semantic `attention` +- explicit linears get tagged +- RoPE patterns can become semantic `rope` before later fusion passes + +## 10. Canonical Cleanup + +`canonicalize/cleanup.py` turns imported IR into a stricter canonical form: + +```python +def canonicalize_exported_graph(graph: IRGraph, *, max_passes: int = 8) -> IRGraph: + for _ in range(max_passes): + for node_id in list(graph.order): + _canonicalize_node(graph, node) + rebuild_graph(graph) + + for node_id in list(graph.order): + _simplify_node(graph, node) + rebuild_graph(graph) + + if _legalize_precisions(graph): + rebuild_graph(graph) + + dce(graph) +``` + +What this pass is responsible for: + +- renaming ops into canonical spellings +- turning shape-only ops into a stable `view` +- rewriting `transpose` and `movedim` into `permute` +- constant-folding trivial expressions +- materializing `ones`, `zeros`, and `arange` constants +- removing no-op casts/views/slices +- inserting FP16 legalizing casts where the runtime requires them +- dead-code elimination + +It also records unsupported canonical ops in `graph.meta["canonical_unsupported_op_counts"]`. + +## 11. Fusion And Graph Optimization + +`optimize_graph.py` runs the semantic fusion passes: + +```python +def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None) -> IRGraph: + canonicalize_exported_graph(graph) + ... + if config.enable_rms_norm: + fuse_rms_norm(graph) + fuse_rms_norm_scale_multiply(graph) + if config.enable_rope: + fuse_rope(graph) + if config.enable_attention: + fuse_attention(graph) + ... + normalize_gemma4_decoder_attention_semantics(graph) + fuse_dense_mlp_tq(graph) + ... + annotate_gold_patterns(graph) + _prune_unused_inputs(graph) +``` + +This is where recognizable subgraphs collapse into semantic or fused nodes such as: + +- `rms_norm` +- `rope` +- `attention` +- `attention_block` +- `self_attention_block` +- `conv_module` +- `lstm_cell` +- `gated_deltanet_prefill` +- `gated_deltanet_decode` + +### Gemma4-specific optimization logic + +Gemma4 gets extra handling in this pass: + +- import hints and graph metadata carry layer types and sliding window info +- exported masks may be elided and rewritten into `window_size` +- some decoder attention behavior is normalized specifically for the saved runtime + +## 12. Lower IR Into A Cactus `Graph` + +`lower.py` is the last compiler stage. It maps `IRValue`s and `IRNode`s into the +runtime `Graph` object from `cactus.bindings.cactus`. + +The high-level structure is: + +```python +def transpile_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: + g = Graph() + env = {} + + for value_id in ir.inputs: + tensor = _lower_input_value(g, ir.values[value_id]) + env[value_id] = tensor + + for value_id, const in ir.constants.items(): + binding = _lookup_weight_binding(value) + if binding is not None: + binding = ensure_binding_compatible(binding, source_tensor=const) + lowered_const = _lower_constant_value(g, value, const, binding=binding) + env[value_id] = lowered_const + + for node_id in ir.order: + outputs = _lower_ir_node(g, ir.nodes[node_id], env, ir) +``` + +### Constant lowering + +This is the critical constant binding behavior: + +```python +def _lower_constant_value(g, value, const, *, binding=None): + if binding is not None: + return g.mmap_weights(binding.path) + + if tensor_value.numel() == 1: + return tensor_value.item() + + return _materialize_constant_tensor(g, tensor_value) +``` + +So constants can become: + +- `g.mmap_weights(path)` +- an inline scalar +- a materialized tensor node baked into the graph + +### Where weight compatibility is repaired + +`weight_compat.py` sits between IR metadata and actual lowering. It can rewrite a +binding to a more executable companion file, for example: + +- convert INT8 embeddings to a cached CQ embedding format +- materialize CQ4 companion files for legacy packed INT4 tensors + +### What lowering implements + +`_lower_ir_node()` is the actual op dispatcher. It covers: + +- elementwise math and comparisons +- views, reshape, expand, permute, transpose +- matmul and linear +- embedding and gather/index paths +- reductions and softmax +- conv1d/conv2d +- norms +- fused attention-family nodes +- fused LSTM and Conformer-style blocks + +If an IR op survives all prior passes and still has no lowering, `lower.py` raises: + +```python +raise NotImplementedError(f"unsupported IR op in lowering: {op}") +``` + +### Gemma4 decoder attention fallback + +One especially important special case is that Gemma4 decoder attention can bypass +the runtime attention kernel and be lowered manually as matmul + mask + softmax: + +```python +def _should_lower_gemma4_decoder_attention_without_kernel(ir, node): + family = str(ir.meta.get("adapter_family") or ir.meta.get("family") or "").strip().lower() + component = str(ir.meta.get("component", "") or "").strip().lower() + if family != "gemma4" or component != "decoder": + return False + return node.op in {"attention", "scaled_dot_product_attention"} +``` + +That path exists because the transpiled Gemma4 decoder has some runtime/kernel +compatibility constraints that are easier to control in lowered graph code. + +## 13. Save Artifacts + +After lowering, `hf_model.py` writes the IR and bundle artifacts. + +For the monolithic path it saves: + +- `raw_ir.json` +- `optimized_ir.json` +- `graph.cactus` +- `graph_bindings.json` +- component subgraphs under `components/` + +For the component path it saves: + +- `raw_ir.json` +- `optimized_ir.json` +- per-component IRs +- `components/manifest.json` +- one lowered graph per component + +The component bundle writer is `_write_component_bundle()`: + +```python +manifest_path = bundle_dir / "manifest.json" +_write_json( + manifest_path, + { + "model_id": model_id, + "model_source": model_source, + "task": task, + "family": family, + "component_order": component_order, + "inputs": _serialize_json_compatible(inputs_metadata), + "components": manifest_components, + }, +) +``` + +Why both a `.cactus` graph and a binding manifest? + +Because `Graph.save()` stores the graph structure, but the file-backed weight +bindings need to be reattached later. The manifests carry: + +- logical input/output names +- runtime node ids +- paths to bound weights or saved Cactus tensor constants + +## 14. Optional Execution And Reference Compare + +If execution is enabled, the transpiler module will run the lowered graph immediately. + +By task: + +- text LM: run graph, report next token +- Gemma4 multimodal: run pipeline, report next token +- Parakeet TDT: run component encoder + custom decode loop, report transcript +- generic audio: run encoder graph and report output shapes + +Many paths also compare the transpiled output against a PyTorch reference forward +unless `--skip-reference-compare` is set. + +## 15. Run A Saved Bundle Later + +Saved bundles are executed by the native C++ `run` binary, not by Python. The +Python CLI just resolves the bundle path and exec's the binary: + +``` +cactus run + → python/cactus/cli/run.py:cmd_run + → subprocess.run([run_bin, bundle_dir, ...]) + → run invokes cactus_init(bundle_dir) from libcactus_engine +``` + +The runtime reads `components/manifest.json` from the bundle directory, mmaps +the `.weights` files, and loads the serialized `.cactus` graphs. + +## Major Situations And How They Are Handled + +### Hugging Face models that normally return caches or `ModelOutput` + +`capture_pytorch.py` forces export-safe kwargs: + +- `use_cache=False` +- `return_dict=False` + +This keeps the capture boundary tensor-only. + +### Weights directory exists but has no manifest + +`hf_model.py` fails fast: + +```python +if not manifest_path.exists(): + raise RuntimeError("weights_dir is missing weights_manifest.json") +``` + +This is intentional. Converted CQ/Cactus weights are the source of truth for +transpiled bundles; raw filename guessing is not part of the runtime contract. + +### Weights directory resolves, but no bindings match + +The transpile script fails hard and suggests re-running conversion: + +- component path: if component graphs resolve `0` bindings +- monolithic path: if full optimized IR resolves `0` bindings + +### Local snapshot exists but is incomplete + +`_load_model_source()` will reject incomplete local snapshots when +`--local-files-only` is set. + +### Gemma4 transformers support is missing locally + +`runtime_support.py` and `hf_model.py` search alternate site-packages and +patch import compatibility for: + +- missing `transformers.models.gemma4` +- missing `_lzma` +- torchvision import probes +- missing `torchvision::nms` +- missing `flex_attention.AuxRequest` + +### Gemma4 local checkpoint key mismatches + +`_repair_gemma4_checkpoint_weights()` remaps local checkpoint keys before loading +the model. + +### Gemma4 multimodal prompt drift + +The transpiler normalizes prompt and tokenization details to match native Cactus: + +- inserts missing image/audio placeholder tokens +- builds a Cactus-style chat prompt when requested +- expands HF newline-merge tokens `108` and `109` back into repeated newline token + `107` +- can replace processor-generated image/audio tensors with native-like versions + +### Parakeet TDT does not run as a single generic forward + +`tdt_runtime.py` provides: + +- a local PyTorch reconstruction of the model +- audio preprocessing +- a greedy TDT decode loop +- component specs for an encoder graph and a recurrent decoder-step graph + +### Some runtime symbols are not present in the current shared library + +`runtime_compat.py` monkey-patches the Python `Graph` wrapper so the transpiler can +still build graphs against the current runtime. It provides compatibility behavior +for items like: + +- generic `matmul` +- `gather(axis=0)` +- limited `conv2d` +- `not_equal` +- scalar math legalization + +### Opaque import is allowed, but execution may still fail later + +If `importers.py` does not know a target, it creates an opaque IR node instead of +dropping it. That preserves information, but lowering still needs a real lowering +case. Opaque nodes are therefore a sign that later lowering may stop with +`unsupported IR op`. + +## Artifact Layout + +A typical transpile artifact directory looks like this: + +```text +transpiled// + raw_ir.json + optimized_ir.json + graph.cactus + graph_bindings.json + result.json + components/ + manifest.json + vision_encoder/ + graph.cactus + bound_constants/ + ... + audio_encoder/ + graph.cactus + lm_encoder/ + graph.cactus + decoder/ + graph.cactus +``` + +Notes: + +- `graph.cactus` is the monolithic lowered graph, when one exists. +- `components/manifest.json` is the portable bundle manifest. +- `graph_bindings.json` and per-component `bound_constant_bindings` tell the loader + how to reattach mmap-backed weights and saved Cactus tensor constants. + +## File-By-File Map + +### Core transpiler files + +| File | Responsibility | +| --- | --- | +| `audio_preprocess.py` | WAV loading, resampling, generic log-mel extraction, Cactus audio frontend wrappers, Parakeet-native features, Gemma4-native audio features. | +| `capture_pytorch.py` | `torch.export` capture wrapper, export-safe kwargs injection, metadata collection, shape propagation, and `CapturedModel` creation. | +| `component_partition.py` | Heuristic component labeling of IR nodes/values and subgraph extraction into `audio_encoder` / `vision_encoder` / `lm_encoder` / `decoder`. | +| `component_pipeline.py` | Capture/optimize/lower loop for one component spec, plus runtime execution of a staged pipeline of components. | +| `graph_ir.py` | Definitions for `IRValue`, `IRNode`, `IRGraph`, deep-copy behavior, and IR consistency verification. | +| `import_ir.py` | Top-level importer from exported FX graph into `IRGraph`; resolves lifted constants, weight binding metadata, and inlined subgraphs. | +| `import_semantics.py` | Early semantic rewrites done immediately after import, especially attention and RoPE recognition. | +| `importers.py` | Large target-to-IR importer registry; converts normalized PyTorch ops into canonical `IRNode`s. | +| `lower.py` | Final lowering from optimized IR into runtime `Graph`, including op lowering, constant binding, fused op handling, and execution wrapper. | +| `model_adapters.py` | Family- and task-specific forward adapters, Gemma4 multimodal pipeline adapters, task canonicalization, and component spec builders. | +| `model_patterns.py` | Topology/op catalog used to annotate recognizable transformer structures after optimization. | +| `model_profiles.py` | Compact model-family profile constants, including compatibility modules, graph context defaults, and weight-name aliases used by fallback runtimes. | +| `multimodal_runtime.py` | Shared multimodal runtime helpers used by bundle execution: prompt normalization, image/audio limiting, and native-like multimodal preprocessing. | +| `normalize.py` | Target-name normalization (`aten.*` to canonical names) and dtype-name normalization into IR spellings. | +| `ops.py` | Canonical op registry and alias table for IR ops, including backend op names and attr schemas. | +| `optimize_graph.py` | Main fusion/optimization pass driver, pattern fusion, Gemma4 attention normalization, and gold-pattern annotation. | +| `runtime_support.py` | Shared import/runtime compatibility helpers for optional Transformers modules, torchvision probes, and flex-attention shims. | +| `tdt_runtime.py` | TDT fallback runtime: local recurrent decoder loop, config loader, and encoder/decoder component-spec generation. | +| `runtime_compat.py` | Runtime wrapper patches that make the current Python `Graph` API usable for transpiled graphs even when some C symbols are absent. | +| `weight_binding.py` | Resolves IR constants to converted `.weights` files through `weights_manifest.json`; also resolves default weights dirs. | +| `weight_compat.py` | Makes bound weights executable with the current runtime by materializing compatibility companion files when needed. | + +### Canonicalization subdirectory + +| File | Responsibility | +| --- | --- | +| `canonicalize/cleanup.py` | Multi-pass canonical cleanup, constant folding, no-op removal, attr normalization, and precision legalization. | +| `canonicalize/utils.py` | Shared graph rebuild, DCE, node removal, dtype helpers, and constant materialization utilities. | + +### Fusion subdirectory + +| File | Responsibility | +| --- | --- | +| `fusion/__init__.py` | Re-exports the fusion matchers used by `optimize_graph.py`. | +| `fusion/common.py` | Shared fusion helper types and graph traversal helpers. | +| `fusion/linear.py` | Matcher for linear projection patterns. | +| `fusion/rms_norm.py` | Matcher for RMSNorm computational subgraphs. | +| `fusion/rope.py` | Matcher for rotary embedding patterns. | +| `fusion/attention.py` | Matchers for attention, attention blocks, self-attention blocks, and GQA-style projection patterns. | +| `fusion/conv.py` | Matcher for Conformer-like convolution modules. | +| `fusion/lstm.py` | Matcher for LSTM cell subgraphs. | +| `fusion/mlp.py` | Matcher for gated MLP branches such as GELU/SiLU gate-up-down patterns. | +| `fusion/rel_pos_bias.py` | Matcher for relative-position-bias subgraphs. | +| `fusion/deltanet.py` | Matcher for fused gated DeltaNet patterns. | + +## CLI Usage + +## 1. Basic Convert + +The simplest invocation (internal command, not in the CLI menu) is: + +```bash +cactus convert +``` + +This quantizes the weights to CQ and then builds the runtime graph. The graph stage: + +- requires the converted Cactus CQ weights (produced by the same command, or passed + with `--weights-dir`; default lookup is `weights/`) +- resolves the runtime library +- launches `python/cactus/transpile/hf_model.py` +- saves artifacts under `weights//` (alongside the converted CQ weights; override with `--artifact-dir`) +- adds `--skip-execute`, so it does not run the graph unless you ask + +To stop after the CQ weights (skip the graph), pass `--weights-only`. To also execute +a reference run right after building the graph, pass `--execute-after-transpile`. + +## 2. Graph Options + +The graph-build options are parsed by `python/cactus/transpile/hf_model.py` and exposed +on `cactus convert`. + +Common ones: + +- `--task` +- `--prompt` +- `--input-ids` +- `--audio-file` +- `--image-file` +- `--weights-dir` +- `--artifact-dir` +- `--component-pipeline` +- `--skip-reference-compare` +- `--graph-filename` +- `--no-fuse-rms-norm` +- `--no-fuse-rope` +- `--no-fuse-attention` +- `--allow-unconverted-weights` for compiler-only debugging + +Examples: + +### Text causal LM + +```bash +cactus convert /path/to/converted_weights --bits 4 \ + --prompt "The capital of France is" \ + --artifact-dir ./transpiled/text_demo +``` + +### Gemma4 multimodal + +```bash +cactus convert \ + --task multimodal_causal_lm_logits \ + --image-file /path/to/image.jpg \ + --audio-file /path/to/audio.wav \ + --prompt "Describe what is happening in the image and audio." +``` + +### Whisper-style encoder capture + +```bash +cactus convert /path/to/converted_weights --bits 4 \ + --audio-file /path/to/sample.wav \ + --task encoder_hidden_states +``` + +### Parakeet TDT + +```bash +cactus convert \ + --audio-file /path/to/sample.wav \ + --task tdt_transcription +``` + +If you want the full direct script help: + +```bash +python -m cactus.transpile.hf_model --help +``` + +## 3. Run A Saved Transpiled Bundle + +Use the bundle root directory: + +```bash +cactus run /path/to/transpiled/model +``` + +### Text causal LM bundle + +```bash +cactus run /path/to/transpiled/model \ + --prompt "The capital of France is" +``` + +Or pass token ids directly: + +```bash +cactus run /path/to/transpiled/model \ + --input-ids 2,13,42 +``` + +### Multimodal Gemma4 bundle + +```bash +cactus run /path/to/transpiled/model \ + --image /path/to/image.jpg \ + --audio /path/to/audio.wav \ + --prompt "What do you observe?" \ + --system "Answer concisely." +``` + +### Audio encoder / transcription bundle + +```bash +cactus run /path/to/transpiled/model \ + --audio /path/to/sample.wav +``` + +### Save result payload + +```bash +cactus run /path/to/transpiled/model \ + --prompt "Hello" \ + --result-json ./result.json +``` + +Internally, `cli/model.py:ensure_runnable_bundle` checks whether the positional +argument is a local path with `components/manifest.json` and dispatches the +native `run` binary against it. + +## Current Runtime Limitations + +As of the current code: + +- Saved bundle execution is implemented for `causal_lm_logits`, + `multimodal_causal_lm_logits`, `encoder_hidden_states`, and + `tdt_transcription`. +- `ctc_logits` can be transpiled, but the runtime does not currently implement + a saved-bundle executor for that task. +- The component pipeline is model-family-specific; it is not a universal split-IR + path. + +## Useful Environment Variables + +- `CACTUS_TRANSPILER_WEIGHTS_DIR`: + fallback weights directory used by `weight_binding.py` +- `CACTUS_TRANSPILER_WEIGHTS_DIR_`: + family-specific weights directory override +- `CACTUS_TRANSPILER_DEBUG_MMAP=1`: + log constant binding behavior during lowering +- `CACTUS_GEMMA4_CAPTURE_FP32=1`: + temporarily promote some Gemma4 text modules to FP32 during capture on CPU + +## Practical Summary + +If you only want the high-level lifecycle, it is: + +1. `cactus convert` quantizes the weights, then launches `python/cactus/transpile/hf_model.py`. +2. The script infers a task, loads the HF model, and creates one example input batch. +3. `model_adapters.py` wraps the model into a smaller canonical task interface. +4. `capture_pytorch.py` runs `torch.export`. +5. `import_ir.py` + `importers.py` convert the export to `IRGraph`. +6. `import_semantics.py` adds early semantic rewrites. +7. `canonicalize/cleanup.py` normalizes and simplifies the IR. +8. `optimize_graph.py` fuses patterns into higher-level semantic ops. +9. `lower.py` builds a runtime `Graph`, mmaps bound weights when possible, and + materializes everything else. +10. The script writes `raw_ir.json`, `optimized_ir.json`, `graph.cactus`, and + `components/manifest.json`. +11. `cactus run` later resolves the bundle path and invokes the native `run` + binary, which mmaps the weights, loads the serialized graph, and runs it + via `cactus_init()` + `cactus_complete()` (see `cactus-engine`). diff --git a/python/cactus/transpile/__init__.py b/python/cactus/transpile/__init__.py new file mode 100644 index 000000000..47c8ee364 --- /dev/null +++ b/python/cactus/transpile/__init__.py @@ -0,0 +1 @@ +"""Cactus graph transpiler package.""" diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py new file mode 100644 index 000000000..9096b8515 --- /dev/null +++ b/python/cactus/transpile/aten_ops.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class AtenOpRule: + torch_prefix: str + ir_op: str + + +def canonical_torch_op(target: Any) -> str: + """Return a stable torch/export op spelling independent of Python repr noise.""" + + if isinstance(target, str): + value = target + else: + rendered = str(target).strip() + if rendered.startswith("aten.") or rendered.startswith("": "getitem", + "aten.view.default": "reshape", + "aten.reshape.default": "reshape", + "aten._unsafe_view.default": "reshape", + "aten.new_empty.default": "new_empty", + "aten.flatten.using_ints": "flatten", + "aten.t.default": "transpose", + "aten.gelu.erf": "gelu_erf", + "cactus_transpile.lfm2_moe_layer_gated.default": "lfm2_moe_layer_gated", + "lfm2_moe_layer_gated.default": "lfm2_moe_layer_gated", + "cactus_transpile.qwen2_moe_layer_gated.default": "qwen2_moe_layer_gated", + "qwen2_moe_layer_gated.default": "qwen2_moe_layer_gated", + "cactus_transpile.gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated", + "gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated", + # Keep the legacy importer key until the importer itself is renamed. + "aten.diff.default": "aten.diff.default", +} + + +_PREFIX_RULES: tuple[AtenOpRule, ...] = ( + AtenOpRule("aten.mul_", "multiply_inplace"), + AtenOpRule("aten._assert_tensor_metadata", "identity"), + AtenOpRule("aten.lift_fresh_copy", "identity"), + AtenOpRule("aten.clone", "identity"), + AtenOpRule("aten.alias", "identity"), + AtenOpRule("aten.detach", "identity"), + AtenOpRule("aten.dropout", "identity"), + AtenOpRule("aten.native_dropout", "identity"), + AtenOpRule("aten.arange.start", "arange"), + AtenOpRule("aten.arange", "arange"), + AtenOpRule("aten.add", "add"), + AtenOpRule("aten.sub", "subtract"), + AtenOpRule("aten.mul", "multiply"), + AtenOpRule("aten.floor_divide", "floor_divide"), + AtenOpRule("aten.div", "divide"), + AtenOpRule("aten.to.dtype", "precision_cast"), + AtenOpRule("aten.to.device", "identity"), + AtenOpRule("aten.type_as", "type_as"), + AtenOpRule("aten.neg", "negate"), + AtenOpRule("aten.ne.", "not_equal"), + AtenOpRule("aten.eq.", "equal"), + AtenOpRule("aten.gt.", "greater"), + AtenOpRule("aten.ge.", "greater_equal"), + AtenOpRule("aten.lt.", "less"), + AtenOpRule("aten.le.", "less_equal"), + AtenOpRule("aten.logical_or", "logical_or"), + AtenOpRule("aten.__or__", "logical_or"), + AtenOpRule("aten.logical_and", "logical_and"), + AtenOpRule("aten.__and__", "logical_and"), + AtenOpRule("aten.logical_not", "logical_not"), + AtenOpRule("aten.bitwise_not", "logical_not"), + AtenOpRule("aten.abs", "abs"), + AtenOpRule("aten.clamp", "clamp"), + AtenOpRule("aten.where", "where"), + AtenOpRule("aten.masked_scatter", "masked_scatter"), + AtenOpRule("aten.masked_fill", "masked_fill"), + AtenOpRule("aten.expand", "expand"), + AtenOpRule("aten.repeat", "repeat"), + AtenOpRule("aten.exp", "scalar_exp"), + AtenOpRule("aten.sqrt", "scalar_sqrt"), + AtenOpRule("aten.reciprocal", "reciprocal"), + AtenOpRule("aten.log", "scalar_log"), + AtenOpRule("aten.rsqrt", "rsqrt"), + AtenOpRule("aten.cos", "cos"), + AtenOpRule("aten.sin", "sin"), + AtenOpRule("aten.squeeze", "squeeze"), + AtenOpRule("aten.unsqueeze", "unsqueeze"), + AtenOpRule("aten.transpose", "transpose"), + AtenOpRule("aten.permute", "permute"), + AtenOpRule("aten.numpy_T", "numpy_T"), + AtenOpRule("aten.movedim", "movedim"), + AtenOpRule("aten.contiguous", "contiguous"), + AtenOpRule("aten.mm", "matmul"), + AtenOpRule("aten.matmul", "matmul"), + AtenOpRule("aten.bmm", "matmul"), + AtenOpRule("aten.linear", "linear"), + AtenOpRule("aten.addmm", "addmm"), + AtenOpRule("aten.relu", "relu"), + AtenOpRule("aten.silu", "silu"), + AtenOpRule("aten.gelu", "gelu"), + AtenOpRule("aten.sigmoid", "sigmoid"), + AtenOpRule("aten.glu", "glu"), + AtenOpRule("aten.softplus", "softplus"), + AtenOpRule("aten.tanh", "tanh"), + AtenOpRule("aten.softmax", "softmax"), + AtenOpRule("aten.scaled_dot_product_attention", "scaled_dot_product_attention"), + AtenOpRule("aten.sum", "sum"), + AtenOpRule("aten.mean", "mean"), + AtenOpRule("aten.all", "min"), + AtenOpRule("aten.var", "variance"), + AtenOpRule("aten.minimum", "minimum"), + AtenOpRule("aten.maximum", "maximum"), + AtenOpRule("aten.min", "min"), + AtenOpRule("aten.amin", "min"), + AtenOpRule("aten.max", "max"), + AtenOpRule("aten.amax", "max"), + AtenOpRule("aten.cat", "cat"), + AtenOpRule("aten.split_with_sizes", "split_with_sizes"), + AtenOpRule("aten.chunk", "chunk"), + AtenOpRule("aten.ones", "ones"), + AtenOpRule("aten.slice_scatter", "slice_scatter"), + AtenOpRule("aten.select_scatter", "select_scatter"), + AtenOpRule("aten.slice", "slice"), + AtenOpRule("aten.select", "index"), + AtenOpRule("aten.gather", "gather"), + AtenOpRule("aten.pad", "pad"), + AtenOpRule("aten.one_hot", "one_hot"), + AtenOpRule("aten.tril", "tril"), + AtenOpRule("aten.unfold", "unfold"), + AtenOpRule("aten.embedding", "embedding"), + AtenOpRule("aten.conv1d", "conv1d"), + AtenOpRule("aten.conv2d", "conv2d"), + AtenOpRule("aten.pow", "pow"), + AtenOpRule("aten.layer_norm", "layer_norm"), + AtenOpRule("aten.native_layer_norm", "layer_norm"), + AtenOpRule("aten.group_norm", "group_norm"), + AtenOpRule("aten.batch_norm", "batch_norm"), + AtenOpRule("aten._native_batch_norm_legit_no_training", "batch_norm"), + AtenOpRule("aten.rms_norm", "rms_norm"), +) + + +_PREFIX_RULES_BY_LENGTH: tuple[AtenOpRule, ...] = tuple( + sorted(_PREFIX_RULES, key=lambda rule: len(rule.torch_prefix), reverse=True) +) + + +def canonical_ir_op(target: Any) -> str: + torch_op = canonical_torch_op(target) + exact = _EXACT_RULES.get(torch_op) + if exact is not None: + return exact + if torch_op in {"aten.ne", "aten.eq", "aten.gt", "aten.ge", "aten.lt", "aten.le"}: + return { + "aten.ne": "not_equal", + "aten.eq": "equal", + "aten.gt": "greater", + "aten.ge": "greater_equal", + "aten.lt": "less", + "aten.le": "less_equal", + }[torch_op] + for rule in _PREFIX_RULES_BY_LENGTH: + if torch_op.startswith(rule.torch_prefix): + return rule.ir_op + return torch_op + + +def is_aten_op(target: Any) -> bool: + return canonical_torch_op(target).startswith("aten.") diff --git a/python/cactus/transpile/audio_preprocess.py b/python/cactus/transpile/audio_preprocess.py new file mode 100644 index 000000000..5e6d0aeea --- /dev/null +++ b/python/cactus/transpile/audio_preprocess.py @@ -0,0 +1,400 @@ +from __future__ import annotations + +import os +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from functools import lru_cache +from pathlib import Path + +import numpy as np +import torch +from scipy.io import wavfile +from scipy.signal import resample_poly + + +_PARAKEET_SAMPLE_RATE = 16000 +_PARAKEET_NUM_MELS = 128 +_PARAKEET_N_FFT = 512 +_PARAKEET_FRAME_LENGTH = 400 +_PARAKEET_HOP_LENGTH = 160 +_PARAKEET_PREEMPHASIS = 0.97 +_PARAKEET_LOG_FLOOR = np.float32(2**-24) +_DEFAULT_MAX_AUDIO_SECONDS = 30.0 + + +def audio_duration_limit_seconds() -> float: + raw = os.environ.get("CACTUS_TRANSPILER_MAX_AUDIO_SECONDS", str(_DEFAULT_MAX_AUDIO_SECONDS)) + try: + value = float(raw) + except (TypeError, ValueError): + return _DEFAULT_MAX_AUDIO_SECONDS + return max(0.0, value) + + +def normalize_audio_samples(samples: np.ndarray) -> np.ndarray: + if samples.dtype == np.uint8: + return ((samples.astype(np.float32) - 128.0) / 128.0).astype(np.float32) + if np.issubdtype(samples.dtype, np.integer): + info = np.iinfo(samples.dtype) + scale = float(max(abs(info.min), info.max)) + if scale <= 0: + scale = 1.0 + return (samples.astype(np.float32) / scale).astype(np.float32) + return samples.astype(np.float32) + + +def _limit_waveform_duration( + waveform: np.ndarray, + *, + sample_rate: int, + max_seconds: float | None, +) -> np.ndarray: + if max_seconds is None or max_seconds <= 0.0: + return waveform + max_samples = int(round(float(sample_rate) * float(max_seconds))) + if max_samples <= 0 or int(waveform.shape[0]) <= max_samples: + return waveform + return np.ascontiguousarray(waveform[:max_samples]) + + +def load_audio_waveform( + audio_file: str | Path, + *, + target_sample_rate: int, + max_seconds: float | None = None, +) -> np.ndarray: + if max_seconds is None: + max_seconds = audio_duration_limit_seconds() + sample_rate, waveform = wavfile.read(str(audio_file)) + if waveform.ndim > 1: + waveform = waveform.mean(axis=1) + waveform = normalize_audio_samples(np.asarray(waveform)) + if int(sample_rate) != int(target_sample_rate): + gcd = int(np.gcd(int(sample_rate), int(target_sample_rate))) + waveform = resample_poly( + waveform, + int(target_sample_rate) // gcd, + int(sample_rate) // gcd, + ).astype(np.float32) + sample_rate = int(target_sample_rate) + waveform = _limit_waveform_duration( + waveform.astype(np.float32), + sample_rate=int(target_sample_rate), + max_seconds=max_seconds, + ) + return waveform.astype(np.float32) + + +@contextmanager +def limited_audio_file( + audio_file: str | Path, + *, + target_sample_rate: int = _PARAKEET_SAMPLE_RATE, + max_seconds: float | None = None, +) -> Iterator[Path]: + if max_seconds is None: + max_seconds = audio_duration_limit_seconds() + if max_seconds <= 0.0: + yield Path(audio_file).expanduser().resolve() + return + + waveform = load_audio_waveform( + audio_file, + target_sample_rate=target_sample_rate, + max_seconds=max_seconds, + ) + temp_audio_path: Path | None = None + try: + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file: + temp_audio_path = Path(temp_file.name) + wavfile.write( + str(temp_audio_path), + int(target_sample_rate), + (np.clip(waveform, -1.0, 1.0) * np.float32(32767.0)).astype(np.int16), + ) + yield temp_audio_path + finally: + if temp_audio_path is not None: + try: + temp_audio_path.unlink() + except OSError: + pass + + +def _hz_to_mel(hz: np.ndarray) -> np.ndarray: + return 2595.0 * np.log10(1.0 + hz / 700.0) + + +def _mel_to_hz(mel: np.ndarray) -> np.ndarray: + return 700.0 * (np.power(10.0, mel / 2595.0) - 1.0) + + +@lru_cache(maxsize=16) +def _mel_filter_bank(num_mels: int, sample_rate: int, n_fft: int) -> np.ndarray: + try: + from transformers.audio_utils import mel_filter_bank as transformers_mel_filter_bank # type: ignore + + return transformers_mel_filter_bank( + num_frequency_bins=n_fft // 2 + 1, + num_mel_filters=num_mels, + min_frequency=0.0, + max_frequency=sample_rate / 2.0, + sampling_rate=sample_rate, + norm="slaney", + mel_scale="slaney", + ).astype(np.float32) + except Exception: + pass + + num_freq_bins = n_fft // 2 + 1 + mel_points = np.linspace( + _hz_to_mel(np.array([0.0], dtype=np.float32))[0], + _hz_to_mel(np.array([sample_rate / 2.0], dtype=np.float32))[0], + num_mels + 2, + ) + hz_points = _mel_to_hz(mel_points) + bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(np.int64) + + filters = np.zeros((num_mels, num_freq_bins), dtype=np.float32) + for index in range(1, num_mels + 1): + left = max(0, min(num_freq_bins - 1, int(bins[index - 1]))) + center = max(left + 1, min(num_freq_bins - 1, int(bins[index]))) + right = max(center + 1, min(num_freq_bins, int(bins[index + 1]))) + for freq in range(left, center): + filters[index - 1, freq] = float(freq - left) / float(max(center - left, 1)) + for freq in range(center, right): + filters[index - 1, freq] = float(right - freq) / float(max(right - center, 1)) + return filters + + +def generic_log_mel_features( + waveform: np.ndarray, + *, + sample_rate: int, + num_mels: int, + n_fft: int, + hop_length: int, + frame_length: int, + preemphasis: float | None = None, + mel_floor: np.float32 = _PARAKEET_LOG_FLOOR, + normalize_active_frames_only: bool = True, +) -> tuple[np.ndarray, int]: + waveform_tensor = torch.from_numpy(waveform).to(torch.float32).unsqueeze(0) + if preemphasis is not None and abs(preemphasis) > 0.0: + waveform_tensor = torch.cat( + [ + waveform_tensor[:, :1], + waveform_tensor[:, 1:] - float(preemphasis) * waveform_tensor[:, :-1], + ], + dim=1, + ) + + window = torch.hann_window(frame_length, periodic=False, dtype=torch.float32) + stft = torch.stft( + waveform_tensor, + n_fft=n_fft, + hop_length=hop_length, + win_length=frame_length, + window=window, + return_complex=True, + center=True, + pad_mode="constant", + ) + magnitudes = torch.view_as_real(stft) + magnitudes = torch.sqrt(magnitudes.pow(2).sum(-1)) + power = magnitudes.pow(2) + + feature_length = max(1, int(waveform.shape[0] // hop_length)) + mel_filters = torch.from_numpy(_mel_filter_bank(num_mels, sample_rate, n_fft)).to(torch.float32) + mel_spec = power.transpose(1, 2) @ mel_filters + mel_spec = torch.log(torch.clamp(mel_spec, min=float(mel_floor))) + + if normalize_active_frames_only: + attention_mask = ( + torch.arange(mel_spec.shape[1], dtype=torch.long, device=mel_spec.device).unsqueeze(0) + < feature_length + ) + mask = attention_mask.unsqueeze(-1).to(torch.float32) + masked = mel_spec * mask + mean = masked.sum(dim=1) / float(feature_length) + centered = masked - mean.unsqueeze(1) + denom = float(max(1, feature_length - 1)) + variance = (centered.pow(2) * mask).sum(dim=1) / denom + mel_spec = (mel_spec - mean.unsqueeze(1)) / torch.sqrt(variance.unsqueeze(1) + np.float32(1e-5)) + mel_spec = mel_spec * mask + else: + # Native Parakeet normalizes across the full centered-STFT frame count, + # then trims to waveform_samples // hop_length. Keep that order for parity. + frame_count = int(mel_spec.shape[1]) + mean = mel_spec.mean(dim=1) + centered = mel_spec - mean.unsqueeze(1) + denom = float(max(1, frame_count - 1)) + variance = centered.pow(2).sum(dim=1) / denom + mel_spec = centered / torch.sqrt(variance.unsqueeze(1) + np.float32(1e-5)) + return mel_spec[0].cpu().numpy().astype(np.float32), feature_length + + +def prepare_cactus_audio_features( + audio_file: str | Path, + *, + model_type: str, + expected_frames: int | None, + expected_mels: int, + torch_dtype: torch.dtype, + layout: str = "frames_mels", +) -> tuple[torch.Tensor, int]: + from cactus.bindings.cactus import cactus_preprocess_audio_features + + capacity_frames = int(expected_frames) if isinstance(expected_frames, int) and expected_frames > 0 else 0 + if capacity_frames <= 0: + lowered_model_type = model_type.lower() + if "whisper" in lowered_model_type: + capacity_frames = 3000 + else: + sample_rate, waveform = wavfile.read(str(audio_file)) + samples = int(waveform.shape[0]) if hasattr(waveform, "shape") else len(waveform) + if sample_rate <= 0: + sample_rate = _PARAKEET_SAMPLE_RATE + capacity_frames = max( + 1, + int(np.ceil(samples * (_PARAKEET_SAMPLE_RATE / float(sample_rate)) / _PARAKEET_HOP_LENGTH)) + 8, + ) + if "whisper" in model_type.lower(): + capacity_frames = max(capacity_frames, 3000) + capacity = max(1, capacity_frames * int(expected_mels)) + with limited_audio_file(audio_file, target_sample_rate=_PARAKEET_SAMPLE_RATE) as frontend_audio_file: + values, mel_bins, frames = cactus_preprocess_audio_features( + str(frontend_audio_file), + model_type, + int(expected_mels), + capacity, + ) + if int(mel_bins) != int(expected_mels): + raise ValueError(f"Cactus audio frontend returned {mel_bins} mel bins, expected {expected_mels}") + feature_array = np.asarray(values, dtype=np.float32).reshape(int(mel_bins), int(frames)) + active_frames = int(frames) + if isinstance(expected_frames, int) and expected_frames > 0: + active_frames = min(active_frames, expected_frames) + if layout == "frames_mels": + feature_array = feature_array[:, :active_frames].T + if isinstance(expected_frames, int) and expected_frames > active_frames: + feature_array = np.pad(feature_array, ((0, expected_frames - active_frames), (0, 0)), mode="constant") + elif layout == "mels_frames": + feature_array = feature_array[:, :active_frames] + if isinstance(expected_frames, int) and expected_frames > active_frames: + feature_array = np.pad(feature_array, ((0, 0), (0, expected_frames - active_frames)), mode="constant") + else: + raise ValueError(f"unsupported Cactus audio feature layout: {layout!r}") + tensor = torch.from_numpy(np.ascontiguousarray(feature_array)).unsqueeze(0).to(dtype=torch_dtype) + return tensor, active_frames + + +def prepare_native_parakeet_audio_features( + audio_file: str | Path, + *, + expected_frames: int | None, + expected_mels: int, + torch_dtype: torch.dtype, +) -> tuple[torch.Tensor, int]: + try: + return prepare_cactus_audio_features( + audio_file, + model_type="parakeet", + expected_frames=expected_frames, + expected_mels=expected_mels, + torch_dtype=torch_dtype, + ) + except Exception: + pass + + waveform = load_audio_waveform(audio_file, target_sample_rate=_PARAKEET_SAMPLE_RATE) + features, feature_length = generic_log_mel_features( + waveform, + sample_rate=_PARAKEET_SAMPLE_RATE, + num_mels=expected_mels or _PARAKEET_NUM_MELS, + n_fft=_PARAKEET_N_FFT, + hop_length=_PARAKEET_HOP_LENGTH, + frame_length=_PARAKEET_FRAME_LENGTH, + preemphasis=_PARAKEET_PREEMPHASIS, + mel_floor=_PARAKEET_LOG_FLOOR, + normalize_active_frames_only=False, + ) + active_frames = feature_length + if isinstance(expected_frames, int) and expected_frames > 0: + active_frames = min(active_frames, expected_frames) + features = features[:active_frames, :] + if features.shape[1] != expected_mels: + raise ValueError( + f"feature mel dimension mismatch: expected {expected_mels}, got {features.shape[1]}" + ) + if isinstance(expected_frames, int) and expected_frames > active_frames: + features = np.pad(features, ((0, expected_frames - active_frames), (0, 0)), mode="constant") + tensor = torch.from_numpy(np.ascontiguousarray(features)).unsqueeze(0).to(dtype=torch_dtype) + return tensor, active_frames + + +def prepare_native_gemma4_audio_features( + audio_file: str | Path, + *, + expected_mels: int, + torch_dtype: torch.dtype, + max_seconds: float | None = None, + pad_to_max_seconds: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, int]: + if max_seconds is None: + max_seconds = audio_duration_limit_seconds() + waveform = load_audio_waveform( + audio_file, + target_sample_rate=_PARAKEET_SAMPLE_RATE, + max_seconds=max_seconds, + ) + # Native Gemma4 pads to the nearest 320 samples, adds 160 samples of + # semicausal left padding, and then runs a 321-point analysis frame. + estimated_frames = int(np.ceil((len(waveform) + 640) / 160.0)) + 8 + capacity = max(1, estimated_frames * int(expected_mels)) + values, mel_bins, frames = prepare_cactus_audio_features_raw( + audio_file, + model_type="gemma4", + mel_bins=int(expected_mels), + capacity=capacity, + max_seconds=max_seconds, + ) + if int(mel_bins) != int(expected_mels): + raise ValueError(f"Cactus Gemma4 audio frontend returned {mel_bins} mel bins, expected {expected_mels}") + max_frames = estimated_frames + if max_seconds is not None and max_seconds > 0.0: + max_samples = int(round(float(_PARAKEET_SAMPLE_RATE) * float(max_seconds))) + max_frames = int(np.ceil((max_samples + 640) / 160.0)) + 8 + active_frames = min(int(frames), int(max_frames)) + features = np.asarray(values, dtype=np.float32).reshape(int(frames), int(mel_bins))[:active_frames, :] + output_frames = max(active_frames, int(max_frames)) if pad_to_max_seconds else active_frames + if output_frames > active_frames: + features = np.pad(features, ((0, output_frames - active_frames), (0, 0)), mode="constant") + feature_tensor = torch.from_numpy(np.ascontiguousarray(features)).unsqueeze(0).to(dtype=torch_dtype) + mask = torch.zeros((1, int(output_frames)), dtype=torch.bool) + mask[:, :active_frames] = True + return feature_tensor, mask, int(active_frames) + + +def prepare_cactus_audio_features_raw( + audio_file: str | Path, + *, + model_type: str, + mel_bins: int, + capacity: int, + max_seconds: float | None = None, +) -> tuple[list[float], int, int]: + from cactus.bindings.cactus import cactus_preprocess_audio_features + + with limited_audio_file( + audio_file, + target_sample_rate=_PARAKEET_SAMPLE_RATE, + max_seconds=max_seconds, + ) as frontend_audio_file: + return cactus_preprocess_audio_features( + str(frontend_audio_file), + model_type, + int(mel_bins), + int(capacity), + ) diff --git a/python/cactus/transpile/canonicalize/cleanup.py b/python/cactus/transpile/canonicalize/cleanup.py new file mode 100644 index 000000000..c68e00447 --- /dev/null +++ b/python/cactus/transpile/canonicalize/cleanup.py @@ -0,0 +1,707 @@ +from __future__ import annotations + +from collections import Counter + +import torch + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.ops import canonicalize_op +from cactus.transpile.ops import get_op +from cactus.transpile.ops import has_op +from cactus.transpile.canonicalize.utils import bypass_unary_node +from cactus.transpile.canonicalize.utils import dce +from cactus.transpile.canonicalize.utils import ir_dtype_to_torch +from cactus.transpile.canonicalize.utils import materialize_constant_output +from cactus.transpile.canonicalize.utils import normalize_dim +from cactus.transpile.canonicalize.utils import normalize_dtype_name +from cactus.transpile.canonicalize.utils import numel +from cactus.transpile.canonicalize.utils import rebuild_graph + + +COMPILER_SUPPORTED_OPS = { + "arange", + "ones", +} + +CANONICAL_UNSUPPORTED_RENAMES = { + "aten.__and__.Tensor": "logical_and", + "aten.__or__.Tensor": "logical_or", + "aten.cumsum.default": "cumsum", + "aten.eq.Scalar": "equal", + "aten.eq.Tensor": "equal", + "aten.ge.Scalar": "greater_equal", + "aten.ge.Tensor": "greater_equal", + "aten.gt.Scalar": "greater", + "aten.gt.Tensor": "greater", + "aten.index.Tensor": "advanced_index", + "aten.le.Scalar": "less_equal", + "aten.le.Tensor": "less_equal", + "aten.lt.Scalar": "less", + "aten.lt.Tensor": "less", +} + +FP32_SUPPORTED_ALL_INPUT_OPS = { + "add", + "subtract", + "multiply", + "divide", + "precision_cast", + "view", + "permute", + "transpose", + "slice", + "index", + "advanced_index", + "gather", + "matmul", + "linear", + "addmm", + "layer_norm", + "batch_norm", + "glu", + "sum", + "mean", + "variance", + "min", + "max", + "scalar_add", + "scalar_subtract", + "scalar_subtract_reverse", + "scalar_multiply", + "scalar_divide", + "scalar_floor_divide", + "scalar_not_equal", + "scalar_equal", + "scalar_greater", + "scalar_greater_equal", + "scalar_less", + "scalar_less_equal", + "abs", + "negate", + "pow", + "scalar_exp", + "scalar_sqrt", + "scalar_log", + "scalar_cos", + "scalar_sin", +} + +FP32_SUPPORTED_INPUT_INDICES = { + # The embedding op dequantizes its weight (input 0) directly (CQ or FP16 when + # bound), and indices (input 1) stay integral — neither needs fp16 legalization. + "embedding": {0, 1}, +} + +FP16_ONLY_OUTPUT_OPS = { + "scalar_divide_reverse", + "not_equal", + "equal", + "greater", + "greater_equal", + "less", + "less_equal", + "logical_and", + "logical_or", + "logical_not", + "one_hot", + "scalar_not_equal", + "scalar_equal", + "scalar_greater", + "scalar_greater_equal", + "scalar_less", + "scalar_less_equal", + "cat", + "softmax", + "rms_norm", + "embedding", + "group_norm", + "lstm_cell", + "relu", + "silu", + "gelu", + "gelu_erf", + "sigmoid", + "tanh", + "attention", + "attention_block", + "conv_module", + "dense_mlp_tq_fused", + "lfm2_moe_layer_gated", + "qwen2_moe_layer_gated", + "gemma4_moe_layer_gated", + "rope", +} + + +def canonicalize_exported_graph(graph: IRGraph, *, max_passes: int = 8) -> IRGraph: + verify_ir(graph) + + for _ in range(max_passes): + changed = False + + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + if _canonicalize_node(graph, node): + changed = True + + rebuild_graph(graph) + + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + if _simplify_node(graph, node): + changed = True + + rebuild_graph(graph) + + if _legalize_precisions(graph): + changed = True + rebuild_graph(graph) + + dce(graph) + + if not changed: + break + + unsupported_counts = summarize_unsupported_ops(graph) + graph.meta = dict(graph.meta) + graph.meta["canonical_unsupported_ops"] = tuple(sorted(unsupported_counts)) + graph.meta["canonical_unsupported_op_counts"] = unsupported_counts + verify_ir(graph) + return graph + + +def summarize_unsupported_ops(graph: IRGraph) -> dict[str, int]: + counts: Counter[str] = Counter() + for node_id in graph.order: + node = graph.nodes[node_id] + if not _is_supported_without_backend_kernel(node.op): + counts[node.op] += 1 + return dict(sorted(counts.items())) + + +def _canonicalize_node(graph: IRGraph, node: IRNode) -> bool: + original_op = node.op + + if original_op in {"aten.new_ones.default", "aten.ones.default"}: + return _materialize_new_ones_constant(graph, node) + + if original_op in {"aten.new_zeros.default", "aten.zeros.default", "new_zeros"}: + return _materialize_new_zeros_constant(graph, node) + + if original_op == "arange": + return _materialize_arange_constant(graph, node) + + if original_op in {"identity", "contiguous"}: + return bypass_unary_node(graph, node) + + if original_op in CANONICAL_UNSUPPORTED_RENAMES: + node.op = CANONICAL_UNSUPPORTED_RENAMES[original_op] + return True + + if original_op == "multiply_inplace": + node.op = "multiply" + return True + + if original_op == "view" and "shape" not in node.attrs: + return _rewrite_view_from_output_shape(graph, node) + + if original_op in {"reshape", "flatten", "unsqueeze", "squeeze"}: + return _rewrite_view_from_output_shape(graph, node) + + if original_op == "expand": + return _rewrite_expand_if_view_equivalent(graph, node) + + if original_op == "transpose": + return _rewrite_transpose_as_permute(graph, node) + + if original_op == "movedim": + return _rewrite_movedim_as_permute(graph, node) + + if original_op == "type_as": + return _rewrite_type_as(graph, node) + + if original_op == "concat": + node.op = "cat" + if "dim" in node.attrs and "axis" not in node.attrs: + node.attrs = {**node.attrs, "axis": int(node.attrs["dim"])} + return True + + if original_op == "scaled_dot_product_attention": + node.op = "attention" + node.attrs = { + "scale": float(node.attrs.get("scale", 0.0)), + "is_causal": bool(node.attrs.get("is_causal", False)), + "window_size": int(node.attrs.get("window_size", 0)), + **({"mask": node.attrs["mask"]} if "mask" in node.attrs else {}), + **({"additive_mask": bool(node.attrs.get("additive_mask", False))} if "additive_mask" in node.attrs else {}), + } + return True + + canonical_op = canonicalize_op(original_op) + if canonical_op != original_op: + node.op = canonical_op + _normalize_attr_spellings(node) + if has_op(node.op): + _filter_attrs_to_schema(node) + return True + + if has_op(node.op): + _normalize_attr_spellings(node) + _filter_attrs_to_schema(node) + return True + + return False + + +def _simplify_node(graph: IRGraph, node: IRNode) -> bool: + if _fold_constant_node(graph, node): + return True + + if node.op == "precision_cast": + return _remove_noop_precision_cast(graph, node) + + if node.op == "view": + if _remove_noop_view(graph, node): + return True + return _collapse_view_chain(graph, node) + + if node.op == "slice": + return _remove_noop_slice(graph, node) + + if node.op == "scalar_add": + return _remove_noop_scalar(graph, node, neutral=0.0) + + if node.op == "scalar_subtract": + return _remove_noop_scalar(graph, node, neutral=0.0) + + if node.op == "scalar_multiply": + return _remove_noop_scalar(graph, node, neutral=1.0) + + if node.op == "scalar_divide": + return _remove_noop_scalar(graph, node, neutral=1.0) + + return False + + +def _fold_constant_node(graph: IRGraph, node: IRNode) -> bool: + if len(node.outputs) != 1: + return False + + output_value = graph.values.get(node.outputs[0]) + if output_value is None: + return False + + if node.op in FP16_ONLY_OUTPUT_OPS: + output_value.dtype = "fp16" + node.meta["dtype"] = "fp16" + + output_torch_dtype = ir_dtype_to_torch(output_value.dtype) if output_value.dtype is not None else None + + constant_inputs = [graph.constants.get(value_id) for value_id in node.inputs] + if any(_is_external_weight_constant(graph, value_id) for value_id in node.inputs): + return False + if node.op == "view" and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): + if output_value.shape is None: + return False + return materialize_constant_output( + graph, + node, + constant_inputs[0].reshape(tuple(int(dim) for dim in output_value.shape)).to(dtype=output_torch_dtype), + ) + + if node.op == "expand" and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): + if output_value.shape is None: + return False + return materialize_constant_output( + graph, + node, + constant_inputs[0].expand(tuple(int(dim) for dim in output_value.shape)).clone().to(dtype=output_torch_dtype), + ) + + if node.op == "precision_cast" and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): + dtype = normalize_dtype_name(node.attrs.get("dtype")) + if dtype is None: + return False + return materialize_constant_output(graph, node, constant_inputs[0].to(dtype=ir_dtype_to_torch(dtype))) + + if node.op in {"cos", "scalar_cos"} and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): + return materialize_constant_output(graph, node, torch.cos(constant_inputs[0])) + + if node.op in {"sin", "scalar_sin"} and len(constant_inputs) == 1 and isinstance(constant_inputs[0], torch.Tensor): + return materialize_constant_output(graph, node, torch.sin(constant_inputs[0])) + + if node.op in {"scalar_add", "scalar_subtract", "scalar_multiply", "scalar_divide"}: + if len(constant_inputs) != 1 or not isinstance(constant_inputs[0], torch.Tensor): + return False + scalar = float(node.attrs["value"]) + if node.op == "scalar_add": + result = constant_inputs[0] + scalar + elif node.op == "scalar_subtract": + result = constant_inputs[0] - scalar + elif node.op == "scalar_multiply": + result = constant_inputs[0] * scalar + else: + result = constant_inputs[0] / scalar + return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) + + if node.op in {"scalar_not_equal", "scalar_equal", "scalar_greater", "scalar_greater_equal", "scalar_less", "scalar_less_equal"}: + if len(constant_inputs) != 1 or not isinstance(constant_inputs[0], torch.Tensor): + return False + scalar = float(node.attrs["value"]) + if node.op == "scalar_not_equal": + result = constant_inputs[0] != scalar + elif node.op == "scalar_equal": + result = constant_inputs[0] == scalar + elif node.op == "scalar_greater": + result = constant_inputs[0] > scalar + elif node.op == "scalar_greater_equal": + result = constant_inputs[0] >= scalar + elif node.op == "scalar_less": + result = constant_inputs[0] < scalar + else: + result = constant_inputs[0] <= scalar + return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) + + if node.op in {"not_equal", "equal", "greater", "greater_equal", "less", "less_equal", "logical_and", "logical_or"}: + if len(constant_inputs) != 2 or not all(isinstance(value, torch.Tensor) for value in constant_inputs): + return False + lhs, rhs = constant_inputs + if node.op == "not_equal": + result = lhs != rhs + elif node.op == "equal": + result = lhs == rhs + elif node.op == "greater": + result = lhs > rhs + elif node.op == "greater_equal": + result = lhs >= rhs + elif node.op == "less": + result = lhs < rhs + elif node.op == "less_equal": + result = lhs <= rhs + elif node.op == "logical_and": + result = torch.logical_and(lhs != 0, rhs != 0) + else: + result = torch.logical_or(lhs != 0, rhs != 0) + return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) + + if node.op == "logical_not": + if len(constant_inputs) != 1 or not isinstance(constant_inputs[0], torch.Tensor): + return False + result = torch.logical_not(constant_inputs[0] != 0) + return materialize_constant_output(graph, node, result.to(dtype=output_torch_dtype)) + + return False + + +def _is_external_weight_constant(graph: IRGraph, value_id: str) -> bool: + if value_id not in graph.constants: + return False + value = graph.values.get(value_id) + if value is None or not isinstance(value.meta, dict): + return False + meta = value.meta + return ( + str(meta.get("kind", "") or "") == "weight" + or bool(str(meta.get("source_name", "") or "")) + or bool(str(meta.get("path", "") or "")) + ) + + +def _legalize_precisions(graph: IRGraph) -> bool: + changed = False + counter = 0 + + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + if node.op in FP32_SUPPORTED_ALL_INPUT_OPS: + continue + + inserted: list[IRNode] = [] + allowed_indices = FP32_SUPPORTED_INPUT_INDICES.get(node.op, set()) + for input_index, value_id in enumerate(list(node.inputs)): + if input_index in allowed_indices: + continue + value = graph.values.get(value_id) + if value is None or normalize_dtype_name(value.dtype) != "fp32": + continue + + counter += 1 + cast_node_id = _unique_ir_id( + f"{node.id}_legalize_fp16_{counter}", + graph.nodes.keys(), + graph.order, + ) + cast_value_id = _unique_ir_id( + f"{value_id}_legalize_fp16_{counter}", + graph.values.keys(), + graph.constants.keys(), + ) + cast_node = IRNode( + id=cast_node_id, + op="precision_cast", + inputs=[value_id], + outputs=[cast_value_id], + attrs={"dtype": "fp16"}, + meta={"inserted_by": "canonicalize_exported_graph"}, + ) + graph.values[cast_value_id] = IRValue( + id=cast_value_id, + shape=value.shape, + dtype="fp16", + producer=cast_node_id, + users=[], + ) + node.inputs[input_index] = cast_value_id + inserted.append(cast_node) + changed = True + + if inserted: + insert_at = graph.order.index(node_id) + for offset, cast_node in enumerate(inserted): + graph.nodes[cast_node.id] = cast_node + graph.order.insert(insert_at + offset, cast_node.id) + + if node.op in FP16_ONLY_OUTPUT_OPS: + for output_id in node.outputs: + output_value = graph.values.get(output_id) + if output_value is not None: + output_value.dtype = "fp16" + node.meta["dtype"] = "fp16" + + return changed + + +def _unique_ir_id(base: str, *existing_groups) -> str: + existing: set[str] = set() + for group in existing_groups: + existing.update(str(item) for item in group) + if base not in existing: + return base + suffix = 1 + while True: + candidate = f"{base}_{suffix}" + if candidate not in existing: + return candidate + suffix += 1 + + +def _rewrite_view_from_output_shape(graph: IRGraph, node: IRNode) -> bool: + if not node.outputs: + return False + output_value = graph.values.get(node.outputs[0]) + if output_value is None or output_value.shape is None: + return False + node.op = "view" + node.attrs = {"shape": tuple(int(dim) for dim in output_value.shape)} + return True + + +def _rewrite_expand_if_view_equivalent(graph: IRGraph, node: IRNode) -> bool: + if not node.inputs or not node.outputs: + return False + input_value = graph.values.get(node.inputs[0]) + output_value = graph.values.get(node.outputs[0]) + if input_value is None or output_value is None: + return False + if input_value.shape is None or output_value.shape is None: + return False + if numel(input_value.shape) != numel(output_value.shape): + node.meta["canonicalization_blocked"] = "expand_requires_broadcast" + return False + node.op = "view" + node.attrs = {"shape": tuple(int(dim) for dim in output_value.shape)} + return True + + +def _rewrite_transpose_as_permute(graph: IRGraph, node: IRNode) -> bool: + if not node.inputs: + return False + input_value = graph.values.get(node.inputs[0]) + if input_value is None or input_value.shape is None: + return False + rank = len(input_value.shape) + dim0 = normalize_dim(int(node.attrs["dim0"]), rank) + dim1 = normalize_dim(int(node.attrs["dim1"]), rank) + permutation = list(range(rank)) + permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] + node.op = "permute" + node.attrs = {"permutation": tuple(permutation)} + return True + + +def _rewrite_movedim_as_permute(graph: IRGraph, node: IRNode) -> bool: + if not node.inputs: + return False + input_value = graph.values.get(node.inputs[0]) + if input_value is None or input_value.shape is None: + return False + rank = len(input_value.shape) + source = [normalize_dim(int(dim), rank) for dim in node.attrs["source"]] + destination = [normalize_dim(int(dim), rank) for dim in node.attrs["destination"]] + if len(source) != len(destination): + return False + + remaining = [dim for dim in range(rank) if dim not in source] + permutation = list(remaining) + for dst, src in sorted(zip(destination, source)): + permutation.insert(dst, src) + + node.op = "permute" + node.attrs = {"permutation": tuple(permutation)} + return True + + +def _rewrite_type_as(graph: IRGraph, node: IRNode) -> bool: + if not node.inputs: + return False + target_dtype = None + if node.outputs: + output_value = graph.values.get(node.outputs[0]) + if output_value is not None: + target_dtype = output_value.dtype + if target_dtype is None and len(node.inputs) > 1: + target_value = graph.values.get(node.inputs[1]) + if target_value is not None: + target_dtype = target_value.dtype + if target_dtype is None: + return False + node.op = "precision_cast" + node.inputs = [node.inputs[0]] + node.attrs = {"dtype": normalize_dtype_name(target_dtype)} + return True + + +def _normalize_attr_spellings(node: IRNode) -> None: + attrs = dict(node.attrs) + if "dim" in attrs and "axis" not in attrs and node.op in {"softmax", "sum", "mean", "variance", "min", "max", "cat", "slice", "index"}: + attrs["axis"] = attrs.pop("dim") + if node.op == "permute" and "dims" in attrs and "permutation" not in attrs: + attrs["permutation"] = tuple(int(dim) for dim in attrs.pop("dims")) + if node.op == "slice": + attrs.setdefault("step", 1) + if node.op == "precision_cast" and "dtype" in attrs: + attrs["dtype"] = normalize_dtype_name(attrs["dtype"]) + node.attrs = attrs + + +def _filter_attrs_to_schema(node: IRNode) -> None: + schema = get_op(node.op) + if not schema.attrs: + node.attrs = {} + return + node.attrs = {name: node.attrs[name] for name in schema.attrs if name in node.attrs} + + +def _remove_noop_precision_cast(graph: IRGraph, node: IRNode) -> bool: + if len(node.inputs) != 1 or len(node.outputs) != 1: + return False + input_value = graph.values.get(node.inputs[0]) + output_value = graph.values.get(node.outputs[0]) + target_dtype = normalize_dtype_name(node.attrs.get("dtype")) + input_dtype = None if input_value is None else normalize_dtype_name(input_value.dtype) + output_dtype = None if output_value is None else normalize_dtype_name(output_value.dtype) + if target_dtype is None or target_dtype == input_dtype or (input_dtype is not None and output_dtype == input_dtype): + return bypass_unary_node(graph, node) + return False + + +def _remove_noop_view(graph: IRGraph, node: IRNode) -> bool: + if len(node.inputs) != 1 or len(node.outputs) != 1: + return False + input_value = graph.values.get(node.inputs[0]) + output_value = graph.values.get(node.outputs[0]) + if input_value is None or output_value is None: + return False + if input_value.shape is None or output_value.shape is None: + return False + if tuple(input_value.shape) == tuple(output_value.shape): + return bypass_unary_node(graph, node) + return False + + +def _collapse_view_chain(graph: IRGraph, node: IRNode) -> bool: + if len(node.inputs) != 1: + return False + input_value = graph.values.get(node.inputs[0]) + if input_value is None or input_value.producer is None: + return False + producer = graph.nodes.get(input_value.producer) + if producer is None or producer.op != "view" or len(producer.inputs) != 1: + return False + node.inputs = [producer.inputs[0]] + return True + + +def _remove_noop_slice(graph: IRGraph, node: IRNode) -> bool: + if len(node.inputs) != 1 or len(node.outputs) != 1: + return False + input_value = graph.values.get(node.inputs[0]) + output_value = graph.values.get(node.outputs[0]) + if input_value is None or output_value is None: + return False + if input_value.shape is None or output_value.shape is None: + return False + if tuple(input_value.shape) == tuple(output_value.shape): + return bypass_unary_node(graph, node) + return False + + +def _remove_noop_scalar(graph: IRGraph, node: IRNode, *, neutral: float) -> bool: + value = node.attrs.get("value") + if value is None: + return False + if float(value) != float(neutral): + return False + return bypass_unary_node(graph, node) + + +def _materialize_new_ones_constant(graph: IRGraph, node: IRNode) -> bool: + if len(node.outputs) != 1: + return False + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + if output_value is None or output_value.shape is None or output_value.dtype is None: + return False + torch_dtype = ir_dtype_to_torch(output_value.dtype) + const = torch.ones(tuple(int(dim) for dim in output_value.shape), dtype=torch_dtype) + return materialize_constant_output(graph, node, const) + + +def _materialize_new_zeros_constant(graph: IRGraph, node: IRNode) -> bool: + if len(node.outputs) != 1: + return False + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + if output_value is None or output_value.shape is None or output_value.dtype is None: + return False + torch_dtype = ir_dtype_to_torch(output_value.dtype) + const = torch.zeros(tuple(int(dim) for dim in output_value.shape), dtype=torch_dtype) + return materialize_constant_output(graph, node, const) + + +def _materialize_arange_constant(graph: IRGraph, node: IRNode) -> bool: + if len(node.outputs) != 1: + return False + start = int(node.attrs.get("start", 0)) + end = node.attrs.get("end") + if end is None: + return False + step = int(node.attrs.get("step", 1)) + dtype = normalize_dtype_name(node.attrs.get("dtype", "fp32")) + torch_dtype = ir_dtype_to_torch(dtype) + const = torch.arange(start, int(end), step=step, dtype=torch_dtype) + return materialize_constant_output(graph, node, const) + + +def _is_supported_without_backend_kernel(op_name: str) -> bool: + return has_op(op_name) or op_name in COMPILER_SUPPORTED_OPS diff --git a/python/cactus/transpile/canonicalize/utils.py b/python/cactus/transpile/canonicalize/utils.py new file mode 100644 index 000000000..7ee113b9a --- /dev/null +++ b/python/cactus/transpile/canonicalize/utils.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +from typing import Any + +import torch + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.model_profiles import profile_for_family + + +def normalize_dtype_name(dtype: Any) -> str | None: + if dtype is None: + return None + if not isinstance(dtype, str): + return str(dtype) + if dtype.startswith("torch."): + torch_name = dtype.split(".", 1)[1] + mapping = { + "bool": "bool", + "int8": "int8", + "int16": "int16", + "int32": "int32", + "int64": "int64", + "float16": "fp16", + "bfloat16": "bf16", + "float32": "fp32", + "float64": "fp64", + } + return mapping.get(torch_name, dtype) + return dtype + + +def ir_dtype_to_torch(dtype: str) -> torch.dtype: + normalized = normalize_dtype_name(dtype) + if normalized == "bool": + return torch.bool + if normalized == "int8": + return torch.int8 + if normalized == "int16": + return torch.int16 + if normalized == "int32": + return torch.int32 + if normalized == "int64": + return torch.int64 + if normalized == "fp16": + return torch.float16 + if normalized == "bf16": + return torch.bfloat16 + if normalized == "fp32": + return torch.float32 + if normalized == "fp64": + return torch.float64 + raise NotImplementedError(f"unsupported IR dtype for constant materialization: {dtype}") + + +def normalize_dim(dim: int, rank: int) -> int: + if dim < 0: + dim += rank + return dim + + +def numel(shape: tuple[int, ...]) -> int: + total = 1 + for dim in shape: + total *= int(dim) + return total + + +def refresh_value_users(graph: IRGraph) -> None: + for value in graph.values.values(): + value.users = [] + for node_id in graph.order: + node = graph.nodes[node_id] + for input_id in node.inputs: + value = graph.values.get(input_id) + if value is not None: + value.users.append(node_id) + + +def rebuild_graph(graph: IRGraph) -> None: + new_values: dict[str, IRValue] = {} + + for value_id in graph.inputs: + old = graph.values.get(value_id, IRValue(id=value_id)) + new_values[value_id] = IRValue( + id=value_id, + shape=old.shape, + dtype=old.dtype, + producer=None, + users=[], + meta=dict(old.meta), + ) + + for value_id in graph.constants.keys(): + old = graph.values.get(value_id, IRValue(id=value_id)) + new_values[value_id] = IRValue( + id=value_id, + shape=old.shape, + dtype=old.dtype, + producer=None, + users=[], + meta=dict(old.meta), + ) + + for node_id in graph.order: + node = graph.nodes[node_id] + for output_id in node.outputs: + old = graph.values.get(output_id, IRValue(id=output_id)) + new_values[output_id] = IRValue( + id=output_id, + shape=old.shape, + dtype=old.dtype, + producer=node_id, + users=[], + meta=dict(old.meta), + ) + + for node_id in graph.order: + node = graph.nodes[node_id] + for input_id in node.inputs: + if input_id not in new_values: + old = graph.values.get(input_id, IRValue(id=input_id)) + new_values[input_id] = IRValue( + id=input_id, + shape=old.shape, + dtype=old.dtype, + producer=old.producer, + users=[], + meta=dict(old.meta), + ) + new_values[input_id].users.append(node_id) + + for output_id in graph.outputs: + if output_id not in new_values: + old = graph.values.get(output_id, IRValue(id=output_id)) + new_values[output_id] = IRValue( + id=output_id, + shape=old.shape, + dtype=old.dtype, + producer=old.producer, + users=[], + meta=dict(old.meta), + ) + + graph.values = new_values + verify_ir(graph) + + +def dce(graph: IRGraph) -> None: + live_values = set(graph.outputs) + live_nodes: set[str] = set() + + changed = True + while changed: + changed = False + for node_id in reversed(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + if not any(output in live_values for output in node.outputs): + continue + if node_id not in live_nodes: + live_nodes.add(node_id) + changed = True + for input_id in node.inputs: + if input_id not in live_values: + live_values.add(input_id) + changed = True + + graph.order = [node_id for node_id in graph.order if node_id in live_nodes] + graph.nodes = {node_id: node for node_id, node in graph.nodes.items() if node_id in live_nodes} + + live_value_ids = set(graph.inputs) | set(graph.outputs) + for node in graph.nodes.values(): + live_value_ids.update(node.inputs) + live_value_ids.update(node.outputs) + + graph.values = {value_id: value for value_id, value in graph.values.items() if value_id in live_value_ids} + graph.constants = {value_id: value for value_id, value in graph.constants.items() if value_id in live_value_ids} + rebuild_graph(graph) + + +def remove_node(graph: IRGraph, node_id: str) -> None: + if node_id in graph.nodes: + del graph.nodes[node_id] + graph.order = [existing_id for existing_id in graph.order if existing_id != node_id] + + +def drop_value_if_unused(graph: IRGraph, value_id: str) -> None: + if value_id in graph.inputs or value_id in graph.outputs or value_id in graph.constants: + return + value = graph.values.get(value_id) + if value is None: + return + if value.producer is None and not value.users: + del graph.values[value_id] + + +def bypass_unary_node(graph: IRGraph, node: IRNode) -> bool: + if len(node.inputs) != 1 or len(node.outputs) != 1: + return False + input_id = node.inputs[0] + output_id = node.outputs[0] + + for other in graph.nodes.values(): + other.inputs = [input_id if value_id == output_id else value_id for value_id in other.inputs] + graph.outputs = [input_id if value_id == output_id else value_id for value_id in graph.outputs] + + remove_node(graph, node.id) + if output_id in graph.values: + del graph.values[output_id] + drop_value_if_unused(graph, input_id) + return True + + +def _is_gemma4_graph(graph: IRGraph) -> bool: + family = str(graph.meta.get("adapter_family") or graph.meta.get("family") or "").lower() + profile = profile_for_family(family) + return profile is not None and profile.family == "gemma4" + + +def _weight_binding_meta(value: IRValue | None) -> dict[str, object] | None: + if value is None: + return None + path = value.meta.get("path") + kind = value.meta.get("kind") + source_name = value.meta.get("source_name") + if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): + return dict(value.meta) + return None + + +def _gemma4_materialized_binding_meta(graph: IRGraph, node: IRNode) -> dict[str, object] | None: + if not _is_gemma4_graph(graph) or len(node.inputs) != 1: + return None + if node.op == "precision_cast": + pass + elif node.op == "scalar_add" and float(node.attrs.get("value", 0.0)) in (0.0, 1.0): + pass + else: + return None + + value_id = node.inputs[0] + for _ in range(4): + value = graph.values.get(value_id) + binding_meta = _weight_binding_meta(value) + if binding_meta is not None: + binding_meta["materialized_from_value_id"] = value_id + binding_meta["materialized_by_op"] = node.op + return binding_meta + producer_id = value.producer if value is not None else None + producer_node = graph.nodes.get(producer_id) if producer_id is not None else None + if producer_node is None or producer_node.op not in {"precision_cast", "type_as"}: + break + if len(producer_node.inputs) != 1: + break + value_id = producer_node.inputs[0] + return None + + +def materialize_constant_output(graph: IRGraph, node: IRNode, value: Any) -> bool: + if len(node.outputs) != 1: + return False + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + if output_value is None: + return False + binding_meta = _gemma4_materialized_binding_meta(graph, node) + if binding_meta is not None: + output_value.meta.update(binding_meta) + graph.constants[output_id] = value + output_value.producer = None + remove_node(graph, node.id) + return True diff --git a/python/cactus/transpile/capture_jax.py b/python/cactus/transpile/capture_jax.py new file mode 100644 index 000000000..bef4b50e3 --- /dev/null +++ b/python/cactus/transpile/capture_jax.py @@ -0,0 +1,1663 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +import copy +from dataclasses import dataclass +import math +from typing import Any + +import numpy as np +import torch + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.jax_semantic_rewrites import apply_jax_semantic_rewrites +from cactus.transpile.lower import TranspiledGraph +from cactus.transpile.lower import transpile_ir +from cactus.transpile.weight_binding import resolve_weight_binding + + +_BINARY_OPS = { + "add": "add", + "sub": "subtract", + "mul": "multiply", + "div": "divide", +} + +_BINARY_EVALUATORS = { + "add": np.add, + "sub": np.subtract, + "mul": np.multiply, + "div": np.divide, +} + +_UNARY_OPS = { + "logistic": "sigmoid", + "tanh": "tanh", + "sqrt": "scalar_sqrt", + "exp": "scalar_exp", + "log": "scalar_log", + "cos": "scalar_cos", + "sin": "scalar_sin", + "neg": "negate", + "abs": "abs", +} + +_UNARY_EVALUATORS = { + "sqrt": np.sqrt, + "exp": np.exp, + "log": np.log, + "cos": np.cos, + "sin": np.sin, + "neg": np.negative, + "logistic": lambda value: 1.0 / (1.0 + np.exp(-value)), + "tanh": np.tanh, +} + +_COMPARE_OPS = { + "eq": "equal", + "ne": "not_equal", + "lt": "less", + "le": "less_equal", + "gt": "greater", + "ge": "greater_equal", +} + +_RHS_SCALAR_COMPARE_OPS = { + "eq": "scalar_equal", + "ne": "scalar_not_equal", + "lt": "scalar_less", + "le": "scalar_less_equal", + "gt": "scalar_greater", + "ge": "scalar_greater_equal", +} + +_LHS_SCALAR_COMPARE_OPS = { + "eq": "scalar_equal", + "ne": "scalar_not_equal", + "lt": "scalar_greater", + "le": "scalar_greater_equal", + "gt": "scalar_less", + "ge": "scalar_less_equal", +} + + +def _dtype_to_ir(dtype: Any) -> str: + dtype = np.dtype(dtype) + if dtype.name == "bfloat16": + return "fp32" + if dtype == np.dtype(np.float16): + return "fp16" + if dtype == np.dtype(np.float32): + return "fp32" + if dtype == np.dtype(np.float64): + return "fp32" + if dtype == np.dtype(np.int8): + return "int8" + if dtype in { + np.dtype(np.int16), + np.dtype(np.int32), + np.dtype(np.int64), + np.dtype(np.uint8), + np.dtype(np.uint16), + np.dtype(np.uint32), + np.dtype(np.uint64), + np.dtype(np.bool_), + }: + return "fp32" + raise NotImplementedError(f"unsupported JAX dtype: {dtype}") + + +def _shape_dtype_from_aval(aval: Any) -> tuple[tuple[int, ...] | None, str | None]: + shape = getattr(aval, "shape", None) + dtype = getattr(aval, "dtype", None) + ir_shape = None if shape is None else tuple(int(dim) for dim in shape) + if ir_shape == (): + ir_shape = (1,) + ir_dtype = None if dtype is None else _dtype_to_ir(dtype) + return ir_shape, ir_dtype + + +@dataclass(frozen=True) +class _SyntheticAval: + shape: tuple[int, ...] + dtype: Any + + +def _constant_to_torch(value: Any) -> torch.Tensor: + array = np.asarray(value) + if array.dtype.name == "bfloat16": + array = array.astype(np.float32) + if array.dtype == np.dtype(np.float64): + array = array.astype(np.float32) + if np.issubdtype(array.dtype, np.floating): + compare_array = array.astype(np.float32, copy=False) + array = np.where(compare_array < -1.0e30, -65504.0, array) + array = np.where(compare_array > 1.0e30, 65504.0, array) + elif array.dtype in { + np.dtype(np.int16), + np.dtype(np.int32), + np.dtype(np.int64), + np.dtype(np.uint8), + np.dtype(np.uint16), + np.dtype(np.uint32), + np.dtype(np.uint64), + np.dtype(np.bool_), + }: + array = array.astype(np.float32) + return torch.from_numpy(np.array(array)) + + +def _tree_path_name(path: Sequence[Any]) -> str: + parts: list[str] = [] + for entry in path: + if hasattr(entry, "key"): + part = str(entry.key) + elif hasattr(entry, "name"): + part = str(entry.name) + elif hasattr(entry, "idx"): + part = str(entry.idx) + else: + part = str(entry) + parts.append(part) + return ".".join(parts) or "param" + + +def _flatten_named_leaves(tree_util: Any, params: Any) -> list[tuple[str, np.ndarray]]: + leaves: list[tuple[str, np.ndarray]] = [] + for path, value in tree_util.tree_flatten_with_path(params)[0]: + leaves.append((_tree_path_name(path), np.asarray(value))) + return leaves + + +def _weight_binding_fields(meta: dict[str, object]) -> dict[str, object] | None: + path = meta.get("path") + kind = meta.get("kind") + source_name = meta.get("source_name") + if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): + return {"path": path, "kind": kind, "source_name": source_name} + return None + + +def _propagate_weight_binding_meta(graph: IRGraph) -> None: + changed = True + while changed: + changed = False + for value_id in list(graph.constants): + value = graph.values[value_id] + if _weight_binding_fields(value.meta) is not None: + continue + if value.meta.get("derived_by_op") != "convert_element_type": + continue + derived_from = value.meta.get("derived_from_value_ids") + if not isinstance(derived_from, (tuple, list)): + continue + for source_id in derived_from: + source_value = graph.values.get(str(source_id)) + if source_value is None: + continue + binding_meta = _weight_binding_fields(source_value.meta) + if binding_meta is None: + continue + value.meta.update(binding_meta) + value.meta["materialized_from_value_id"] = str(source_id) + graph.meta.setdefault("weight_bindings", {})[value_id] = dict(value.meta) + changed = True + break + + +def _binding_meta( + *, + name: str, + weights_dir: str | None, + explicit: dict[str, dict[str, str]], +) -> dict[str, str]: + if name in explicit: + return dict(explicit[name]) + binding = resolve_weight_binding(weights_dir=weights_dir, source_name=name) + if binding is None: + return {} + return { + "path": binding.path, + "kind": binding.kind, + "source_name": binding.source_name, + } + + +def _freeze_leading_param_inputs( + graph: IRGraph, + named_leaves: Sequence[tuple[str, np.ndarray]], + *, + weights_dir: str | None, + explicit: dict[str, dict[str, str]], +) -> None: + if not named_leaves: + return + if len(graph.inputs) < len(named_leaves): + raise ValueError( + f"JAX graph has {len(graph.inputs)} inputs, cannot freeze {len(named_leaves)} parameter leaves" + ) + param_input_ids = list(graph.inputs[: len(named_leaves)]) + graph.inputs = graph.inputs[len(named_leaves) :] + graph.meta["weight_bindings"] = {} + for value_id, (name, leaf_array) in zip(param_input_ids, named_leaves, strict=True): + value = graph.values[value_id] + tensor = _constant_to_torch(leaf_array) + value.shape = tuple(int(dim) for dim in tensor.shape) + value.dtype = _dtype_to_ir(tensor.numpy().dtype) + value.meta.update( + { + "source_name": name, + "jax_param_input": True, + **_binding_meta(name=name, weights_dir=weights_dir, explicit=explicit), + } + ) + graph.constants[value_id] = tensor + if "path" in value.meta: + graph.meta["weight_bindings"][value_id] = dict(value.meta) + + +def _literal_value(literal: Any) -> Any: + value = getattr(literal, "val", literal) + if isinstance(value, np.ndarray): + return value.item() if value.ndim == 0 else value + if hasattr(value, "item"): + try: + return value.item() + except Exception: + return value + return value + + +class _JaxImportContext: + def __init__(self) -> None: + self.var_ids: dict[int, str] = {} + self.gather_index_aliases: dict[str, str] = {} + self.next_value_id = 0 + self.next_node_id = 0 + + def value_id(self, var: Any) -> str: + key = id(var) + existing = self.var_ids.get(key) + if existing is not None: + return existing + name = f"v{self.next_value_id}" + self.next_value_id += 1 + self.var_ids[key] = name + return name + + def bind_value(self, var: Any, value_id: str) -> None: + self.var_ids[id(var)] = value_id + + def node_id(self, op: str) -> str: + name = f"n{self.next_node_id}_{op}" + self.next_node_id += 1 + return name + + +@dataclass(frozen=True) +class JaxGraphSpec: + name: str + fn: Callable[..., Any] + example_args: Sequence[Any] + role: str = "generic" + input_names: Sequence[str] | None = None + output_names: Sequence[str] | None = None + graph_meta: dict[str, object] | None = None + + +@dataclass +class CapturedJaxGraph: + spec: JaxGraphSpec + raw_ir_graph: IRGraph + ir_graph: IRGraph + graph: TranspiledGraph + + def execute(self, *args: Any) -> list[Any]: + if len(args) != len(self.spec.example_args): + raise ValueError( + f"graph {self.spec.name!r} expected {len(self.spec.example_args)} inputs, got {len(args)}" + ) + self.graph.set_inputs( + [ + _coerce_runtime_input(arg, example) + for arg, example in zip(args, self.spec.example_args, strict=True) + ] + ) + return self.graph.execute() + + +@dataclass +class CapturedJaxGraphBundle: + graphs: dict[str, CapturedJaxGraph] + params: Any + weights_dir: str | None = None + + def execute(self, graph_name: str, *args: Any) -> list[Any]: + if graph_name not in self.graphs: + available = ", ".join(sorted(self.graphs)) or "" + raise ValueError(f"unknown JAX graph {graph_name!r}; available graphs: {available}") + return self.graphs[graph_name].execute(*args) + + +def _add_value(graph: IRGraph, value_id: str, aval: Any, *, producer: str | None = None) -> None: + shape, dtype = _shape_dtype_from_aval(aval) + graph.add_value(IRValue(id=value_id, shape=shape, dtype=dtype, producer=producer)) + + +def _coerce_runtime_input(value: Any, example: Any) -> np.ndarray: + array = np.asarray(value) + example_array = np.asarray(example) + target_dtype = example_array.dtype + if target_dtype.name == "bfloat16": + target_dtype = np.dtype(np.float32) + if array.dtype != target_dtype: + array = array.astype(target_dtype, copy=False) + if array.shape == (): + array = array.reshape((1,)) + return array + + +def _add_constant( + graph: IRGraph, + ctx: _JaxImportContext, + var: Any, + value: Any, + *, + source_name: str, + meta: dict[str, object] | None = None, +) -> str: + value_id = ctx.value_id(var) + if value_id in graph.values: + return value_id + tensor = _constant_to_torch(value) + graph.add_value( + IRValue( + id=value_id, + shape=tuple(int(dim) for dim in tensor.shape), + dtype=_dtype_to_ir(tensor.numpy().dtype), + producer=None, + meta={"source_name": source_name, **dict(meta or {})}, + ) + ) + graph.constants[value_id] = tensor + return value_id + + +def _register_node( + graph: IRGraph, + node: IRNode, + *, + out_avals: Sequence[Any], +) -> None: + graph.add_node(node) + graph.order.append(node.id) + for output_id, aval in zip(node.outputs, out_avals, strict=True): + shape, dtype = _shape_dtype_from_aval(aval) + graph.values[output_id].shape = shape + graph.values[output_id].dtype = dtype + + +def _primitive_name(eqn: Any) -> str: + primitive = getattr(eqn, "primitive", None) + return str(getattr(primitive, "name", primitive)) + + +def _is_literal(var: Any) -> bool: + return type(var).__name__ == "Literal" + + +def _ensure_literal_constant(graph: IRGraph, ctx: _JaxImportContext, literal: Any) -> str: + value_id = ctx.value_id(literal) + if value_id in graph.values: + return value_id + value = _literal_value(literal) + return _add_constant(graph, ctx, literal, value, source_name=f"literal:{value!r}") + + +def _derived_meta(input_ids: Sequence[str], *, op: str) -> dict[str, object]: + if not input_ids: + return {"derived_by_op": op} + return {"derived_from_value_ids": tuple(input_ids), "derived_by_op": op} + + +def _input_ids(graph: IRGraph, ctx: _JaxImportContext, invars: Sequence[Any]) -> list[str]: + result: list[str] = [] + for var in invars: + if _is_literal(var): + result.append(_ensure_literal_constant(graph, ctx, var)) + else: + result.append(ctx.value_id(var)) + return result + + +def _out_avals(eqn: Any) -> tuple[Any, ...]: + return tuple(getattr(var, "aval", None) for var in eqn.outvars) + + +def _literal_number(var: Any) -> float | None: + if not _is_literal(var): + return None + value = _literal_value(var) + if isinstance(value, (bool, np.bool_)): + return float(value) + if isinstance(value, (int, float, np.number)): + return float(value) + return None + + +def _constant_scalar(graph: IRGraph, value_id: str) -> float | None: + value = graph.constants.get(value_id) + if value is None: + return None + if graph.values[value_id].meta.get("jax_closed_constant"): + return None + array = np.asarray(value) + if array.shape != (): + return None + return float(array.item()) + + +def _constant_singleton_scalar(graph: IRGraph, value_id: str) -> float | None: + value = graph.constants.get(value_id) + if value is None: + return None + array = np.asarray(value) + if array.size != 1: + return None + return float(array.reshape(()).item()) + + +def _constant_scalar_or_singleton(graph: IRGraph, value_id: str) -> float | None: + scalar = _constant_scalar(graph, value_id) + if scalar is not None: + return scalar + singleton = _constant_singleton_scalar(graph, value_id) + if singleton is not None: + return singleton + value = graph.constants.get(value_id) + if value is None: + return None + array = np.asarray(value) + if array.size == 0: + return None + first = float(array.reshape(-1)[0].item()) + if np.allclose(array, first, rtol=0.0, atol=0.0): + return first + return None + + +def _replace_scalar_constant_with_broadcast( + graph: IRGraph, + value_id: str, + output_var: Any, + *, + source_op: str, +) -> str: + scalar = _constant_singleton_scalar(graph, value_id) + if scalar is None: + return value_id + shape = tuple(int(dim) for dim in getattr(output_var.aval, "shape", ())) + if not shape or _product(shape) <= 1: + return value_id + dtype = getattr(output_var.aval, "dtype", np.float32) + replacement_id = f"{value_id}__{source_op}_broadcast" + if replacement_id in graph.values: + return replacement_id + tensor = _constant_to_torch(np.full(shape, scalar, dtype=np.dtype(dtype))) + graph.add_value( + IRValue( + id=replacement_id, + shape=shape, + dtype=_dtype_to_ir(tensor.numpy().dtype), + producer=None, + meta={ + "source_name": f"broadcasted_scalar:{source_op}", + **_derived_meta((value_id,), op=f"{source_op}:scalar_broadcast"), + }, + ) + ) + graph.constants[replacement_id] = tensor + return replacement_id + + +def _where_scalar_value(value: float) -> float: + if value < -1.0e30: + return -65504.0 + if value > 1.0e30: + return 65504.0 + return value + + +def _constant_array(graph: IRGraph, value_id: str) -> np.ndarray | None: + value = graph.constants.get(value_id) + if value is None: + return None + if graph.values[value_id].meta.get("jax_closed_constant"): + return None + return np.asarray(value) + + +def _try_fold_constant_unary( + graph: IRGraph, + ctx: _JaxImportContext, + *, + prim: str, + input_id: str, + outvar: Any, +) -> bool: + evaluator = _UNARY_EVALUATORS.get(prim) + array = _constant_array(graph, input_id) + if evaluator is None or array is None: + return False + _add_constant( + graph, + ctx, + outvar, + evaluator(array), + source_name=f"folded:{prim}", + meta=_derived_meta((input_id,), op=prim), + ) + return True + + +def _try_fold_constant_binary( + graph: IRGraph, + ctx: _JaxImportContext, + *, + prim: str, + input_ids: Sequence[str], + outvar: Any, +) -> bool: + evaluator = _BINARY_EVALUATORS.get(prim) + if evaluator is None: + return False + lhs_array = _constant_array(graph, input_ids[0]) + rhs_array = _constant_array(graph, input_ids[1]) + if lhs_array is None or rhs_array is None: + return False + _add_constant( + graph, + ctx, + outvar, + evaluator(lhs_array, rhs_array), + source_name=f"folded:{prim}", + meta=_derived_meta(input_ids, op=prim), + ) + return True + + +def _alias_output(ctx: _JaxImportContext, outvar: Any, source_id: str) -> None: + ctx.bind_value(outvar, source_id) + + +def _product(values: Sequence[int]) -> int: + result = 1 + for value in values: + result *= int(value) + return result + + +def _conv_static_int(values: Sequence[int], *, name: str) -> int: + ints = tuple(int(value) for value in values) + if not ints: + return 1 + if any(value != ints[0] for value in ints): + raise NotImplementedError(f"JAX conv_general_dilated requires uniform {name}, got {ints}") + return ints[0] + + +def _conv_padding_int(padding: Sequence[Sequence[int]]) -> int: + pairs = tuple((int(left), int(right)) for left, right in padding) + if not pairs: + return 0 + if any(left != right for left, right in pairs): + raise NotImplementedError(f"JAX conv_general_dilated requires symmetric padding, got {pairs}") + return _conv_static_int(tuple(left for left, _ in pairs), name="padding") + + +def _conv_output_permutation(out_spec: Sequence[int]) -> tuple[int, ...]: + spec = tuple(int(axis) for axis in out_spec) + return tuple(spec.index(final_axis) for final_axis in range(len(spec))) + + +def _jaxpr_invars(jaxpr_like: Any) -> Sequence[Any]: + return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "invars", ())) + + +def _jaxpr_outvars(jaxpr_like: Any) -> Sequence[Any]: + return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "outvars", ())) + + +def _jaxpr_eqns(jaxpr_like: Any) -> Sequence[Any]: + return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "eqns", ())) + + +def _jaxpr_constvars(jaxpr_like: Any) -> Sequence[Any]: + return tuple(getattr(getattr(jaxpr_like, "jaxpr", jaxpr_like), "constvars", ())) + + +def _jaxpr_consts(jaxpr_like: Any) -> Sequence[Any]: + return tuple(getattr(jaxpr_like, "consts", ()) or ()) + + +def _inline_jaxpr( + graph: IRGraph, + ctx: _JaxImportContext, + jaxpr_like: Any, + input_ids: Sequence[str], + outvars: Sequence[Any], +) -> None: + inner_output_ids = _inline_jaxpr_outputs(graph, ctx, jaxpr_like, input_ids) + if len(outvars) != len(inner_output_ids): + raise NotImplementedError("nested JAXPR output arity mismatch") + for outer_var, inner_output_id in zip(outvars, inner_output_ids, strict=True): + _alias_output(ctx, outer_var, inner_output_id) + + +def _inline_jaxpr_outputs( + graph: IRGraph, + ctx: _JaxImportContext, + jaxpr_like: Any, + input_ids: Sequence[str], +) -> list[str]: + invars = _jaxpr_invars(jaxpr_like) + if len(input_ids) != len(invars): + raise NotImplementedError("nested JAXPR input arity mismatch") + invar_keys = {id(var) for var in invars} + for inner_eqn in _jaxpr_eqns(jaxpr_like): + for outvar in inner_eqn.outvars: + if id(outvar) not in invar_keys: + ctx.var_ids.pop(id(outvar), None) + for inner_var, input_id in zip(invars, input_ids, strict=True): + ctx.bind_value(inner_var, input_id) + for index, (constvar, const) in enumerate(zip(_jaxpr_constvars(jaxpr_like), _jaxpr_consts(jaxpr_like), strict=True)): + _add_constant(graph, ctx, constvar, const, source_name=f"nested_const_{index}") + for inner_eqn in _jaxpr_eqns(jaxpr_like): + _import_eqn(graph, ctx, inner_eqn) + return [ctx.value_id(inner_var) for inner_var in _jaxpr_outvars(jaxpr_like)] + + +def _generated_value( + graph: IRGraph, + ctx: _JaxImportContext, + *, + stem: str, + shape: tuple[int, ...] | None, + dtype: str | None, + producer: str, + meta: dict[str, object] | None = None, +) -> str: + value_id = f"v{ctx.next_value_id}_{stem}" + ctx.next_value_id += 1 + if value_id in graph.values: + raise ValueError(f"duplicate generated IR value id: {value_id}") + graph.values[value_id] = IRValue( + id=value_id, + shape=shape, + dtype=dtype, + producer=producer, + meta=dict(meta or {}), + ) + return value_id + + +def _register_generated_node( + graph: IRGraph, + node: IRNode, + *, + output_shapes: Sequence[tuple[int, ...] | None], + output_dtypes: Sequence[str | None], +) -> None: + if node.id in graph.nodes: + raise ValueError(f"duplicate IR node id: {node.id}") + graph.nodes[node.id] = node + graph.order.append(node.id) + for output_id, shape, dtype in zip(node.outputs, output_shapes, output_dtypes, strict=True): + graph.values[output_id].shape = shape + graph.values[output_id].dtype = dtype + + +def _scan_slice_input( + graph: IRGraph, + ctx: _JaxImportContext, + value_id: str, + *, + iteration: int, + input_index: int, +) -> str: + source_value = graph.values[value_id] + if source_value.shape is None or not source_value.shape: + return value_id + slice_node_id = ctx.node_id("scan_slice") + sliced_shape = (1, *tuple(int(dim) for dim in source_value.shape[1:])) + sliced_id = _generated_value( + graph, + ctx, + stem=f"scan{iteration}_{input_index}", + shape=sliced_shape, + dtype=source_value.dtype, + producer=slice_node_id, + meta={ + "derived_from_value_ids": (value_id,), + "derived_by_op": "scan_slice", + }, + ) + slice_node = IRNode( + slice_node_id, + "slice", + [value_id], + [sliced_id], + attrs={"axis": 0, "start": iteration, "end": iteration + 1, "step": 1}, + meta={"jax_generated": "scan_unroll"}, + ) + _register_generated_node(graph, slice_node, output_shapes=(sliced_shape,), output_dtypes=(source_value.dtype,)) + view_node_id = ctx.node_id("scan_slice_view") + output_shape = tuple(int(dim) for dim in source_value.shape[1:]) + if output_shape == () and sliced_shape == (1,): + return sliced_id + output_id = _generated_value( + graph, + ctx, + stem=f"scan{iteration}_{input_index}_view", + shape=output_shape, + dtype=source_value.dtype, + producer=view_node_id, + meta={ + "derived_from_value_ids": (sliced_id,), + "derived_by_op": "scan_slice_view", + }, + ) + view_node = IRNode( + view_node_id, + "view", + [sliced_id], + [output_id], + attrs={"shape": output_shape}, + meta={"jax_generated": "scan_unroll"}, + ) + _register_generated_node(graph, view_node, output_shapes=(output_shape,), output_dtypes=(source_value.dtype,)) + return output_id + + +def _add_iota_constant(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> None: + params = dict(getattr(eqn, "params", {}) or {}) + shape = tuple(int(dim) for dim in params["shape"]) + dimension = int(params["dimension"]) + dtype = params.get("dtype", getattr(eqn.outvars[0].aval, "dtype", np.float32)) + values = np.arange(shape[dimension], dtype=np.dtype(dtype)) + reshape = [1] * len(shape) + reshape[dimension] = shape[dimension] + values = np.broadcast_to(values.reshape(reshape), shape) + _add_constant(graph, ctx, eqn.outvars[0], values, source_name=f"iota:{shape}:{dimension}") + + +def _import_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> None: + prim = _primitive_name(eqn) + params = dict(getattr(eqn, "params", {}) or {}) + + if prim in {"jit", "remat2"}: + _inline_jaxpr(graph, ctx, params["jaxpr"], _input_ids(graph, ctx, eqn.invars), eqn.outvars) + return + + if prim in {"custom_jvp_call", "custom_jvp_call_jaxpr"}: + call_jaxpr = params.get("call_jaxpr") or params.get("fun_jaxpr") or params.get("jaxpr") + if call_jaxpr is None: + raise NotImplementedError("JAX custom_jvp_call import requires a primal call_jaxpr") + _inline_jaxpr(graph, ctx, call_jaxpr, _input_ids(graph, ctx, eqn.invars), eqn.outvars) + return + + if prim == "scan": + length = int(params.get("length", 1)) + num_consts = int(params.get("num_consts", 0)) + num_carry = int(params.get("num_carry", len(eqn.outvars))) + input_ids = _input_ids(graph, ctx, eqn.invars) + const_ids = input_ids[:num_consts] + carry_ids = input_ids[num_consts : num_consts + num_carry] + xs_ids = input_ids[num_consts + num_carry :] + if len(eqn.outvars) != num_carry: + raise NotImplementedError("JAX scan import currently supports carry-only scan outputs") + for iteration in range(length): + sliced_xs = [ + _scan_slice_input(graph, ctx, value_id, iteration=iteration, input_index=index) + for index, value_id in enumerate(xs_ids) + ] + body_outputs = _inline_jaxpr_outputs(graph, ctx, params["jaxpr"], [*const_ids, *carry_ids, *sliced_xs]) + if len(body_outputs) != num_carry: + raise NotImplementedError("JAX scan body output arity mismatch") + carry_ids = body_outputs + for outvar, carry_id in zip(eqn.outvars, carry_ids, strict=True): + _alias_output(ctx, outvar, carry_id) + return + + if prim == "stop_gradient": + _alias_output(ctx, eqn.outvars[0], _input_ids(graph, ctx, eqn.invars)[0]) + return + + if prim == "iota": + _add_iota_constant(graph, ctx, eqn) + return + + node = _node_for_eqn(graph, ctx, eqn) + if node is None: + return + _register_node(graph, node, out_avals=_out_avals(eqn)) + + +def _node_for_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> IRNode | None: + prim = _primitive_name(eqn) + inputs = _input_ids(graph, ctx, eqn.invars) + outputs = [ctx.value_id(var) for var in eqn.outvars] + node_id = ctx.node_id(prim) + params = dict(getattr(eqn, "params", {}) or {}) + meta = {"jax_primitive": prim} + + if prim in _BINARY_OPS: + if _try_fold_constant_binary(graph, ctx, prim=prim, input_ids=inputs, outvar=eqn.outvars[0]): + return None + if prim == "div": + lhs_literal = _literal_number(eqn.invars[0]) + if lhs_literal is not None: + reciprocal_id = f"{outputs[0]}__reciprocal" + reciprocal_node = IRNode( + node_id, + "pow", + [inputs[1]], + [reciprocal_id], + attrs={"exponent": -1.0}, + meta=meta, + ) + _register_node(graph, reciprocal_node, out_avals=(eqn.outvars[0].aval,)) + if lhs_literal == 1.0: + _alias_output(ctx, eqn.outvars[0], reciprocal_id) + return None + return IRNode( + ctx.node_id("scalar_div_mul"), + "scalar_multiply", + [reciprocal_id], + outputs, + attrs={"value": lhs_literal}, + meta=meta, + ) + if prim in _BINARY_OPS: + if len(inputs) == 2 and len(outputs) == 1: + inputs = [ + _replace_scalar_constant_with_broadcast(graph, input_id, eqn.outvars[0], source_op=prim) + for input_id in inputs + ] + return IRNode(node_id, _BINARY_OPS[prim], inputs, outputs, meta=meta) + if prim in _UNARY_OPS: + if _try_fold_constant_unary(graph, ctx, prim=prim, input_id=inputs[0], outvar=eqn.outvars[0]): + return None + return IRNode(node_id, _UNARY_OPS[prim], inputs, outputs, meta=meta) + if prim == "log1p": + add_id = f"{outputs[0]}__log1p_add" + _register_node( + graph, + IRNode( + node_id, + "scalar_add", + inputs, + [add_id], + attrs={"value": 1.0}, + meta=meta, + ), + out_avals=(eqn.outvars[0].aval,), + ) + return IRNode(ctx.node_id("log1p_log"), "scalar_log", [add_id], outputs, meta=meta) + if prim == "expm1": + exp_id = f"{outputs[0]}__expm1_exp" + _register_node( + graph, + IRNode(node_id, "scalar_exp", inputs, [exp_id], meta=meta), + out_avals=(eqn.outvars[0].aval,), + ) + return IRNode(ctx.node_id("expm1_sub"), "scalar_add", [exp_id], outputs, attrs={"value": -1.0}, meta=meta) + if prim in _COMPARE_OPS: + lhs_scalar = _constant_scalar(graph, inputs[0]) + rhs_scalar = _constant_scalar(graph, inputs[1]) + if rhs_scalar is not None: + return IRNode(node_id, _RHS_SCALAR_COMPARE_OPS[prim], [inputs[0]], outputs, attrs={"value": rhs_scalar}, meta=meta) + if lhs_scalar is not None: + return IRNode(node_id, _LHS_SCALAR_COMPARE_OPS[prim], [inputs[1]], outputs, attrs={"value": lhs_scalar}, meta=meta) + return IRNode(node_id, _COMPARE_OPS[prim], inputs, outputs, meta=meta) + if prim == "and": + return IRNode(node_id, "logical_and", inputs, outputs, meta=meta) + if prim == "or": + return IRNode(node_id, "logical_or", inputs, outputs, meta=meta) + if prim == "square": + return IRNode(node_id, "multiply", [inputs[0], inputs[0]], outputs, meta=meta) + if prim == "rsqrt": + return IRNode(node_id, "pow", inputs, outputs, attrs={"exponent": -0.5}, meta=meta) + if prim == "integer_pow": + exponent = int(params["y"]) + if exponent == 2: + return IRNode(node_id, "multiply", [inputs[0], inputs[0]], outputs, meta=meta) + return IRNode(node_id, "pow", inputs, outputs, attrs={"exponent": float(exponent)}, meta=meta) + if prim == "pow": + rhs_literal = _literal_number(eqn.invars[1]) + if rhs_literal is not None: + return IRNode(node_id, "pow", [inputs[0]], outputs, attrs={"exponent": float(rhs_literal)}, meta=meta) + lhs_literal = _literal_number(eqn.invars[0]) + if lhs_literal is not None and lhs_literal > 0.0: + intermediate = f"{outputs[0]}__logmul" + scale_node = IRNode( + node_id, + "scalar_multiply", + [inputs[1]], + [intermediate], + attrs={"value": math.log(lhs_literal)}, + meta=meta, + ) + _register_node(graph, scale_node, out_avals=(eqn.outvars[0].aval,)) + return IRNode(ctx.node_id("pow_exp"), "scalar_exp", [intermediate], outputs, meta=meta) + raise NotImplementedError("JAX pow import only supports positive scalar base") + if prim == "dot_general": + dimension_numbers = params.get("dimension_numbers") + if dimension_numbers is not None: + ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) = dimension_numbers + lhs_shape = tuple(getattr(eqn.invars[0].aval, "shape", ())) + rhs_shape = tuple(getattr(eqn.invars[1].aval, "shape", ())) + lhs_contract = tuple(int(dim) for dim in lhs_contract) + rhs_contract = tuple(int(dim) for dim in rhs_contract) + lhs_batch = tuple(int(dim) for dim in lhs_batch) + rhs_batch = tuple(int(dim) for dim in rhs_batch) + if lhs_batch or rhs_batch: + if tuple(lhs_shape[dim] for dim in lhs_batch) != tuple(rhs_shape[dim] for dim in rhs_batch): + raise NotImplementedError("JAX dot_general batch dimensions must have matching sizes") + expected_rhs_contract = (len(rhs_batch),) + if ( + tuple(lhs_contract) != (len(lhs_shape) - 1,) + or tuple(rhs_contract) != expected_rhs_contract + or len(rhs_shape) > 2 + ): + batch_shape = tuple(lhs_shape[dim] for dim in lhs_batch) + lhs_contract_shape = tuple(lhs_shape[dim] for dim in lhs_contract) + rhs_contract_shape = tuple(rhs_shape[dim] for dim in rhs_contract) + if lhs_contract_shape != rhs_contract_shape: + raise NotImplementedError("JAX dot_general contraction dimensions must have matching sizes") + lhs_non_contract = tuple( + dim for dim in range(len(lhs_shape)) if dim not in set(lhs_batch) | set(lhs_contract) + ) + rhs_non_contract = tuple( + dim for dim in range(len(rhs_shape)) if dim not in set(rhs_batch) | set(rhs_contract) + ) + lhs_order = lhs_batch + lhs_non_contract + lhs_contract + rhs_order = rhs_batch + rhs_contract + rhs_non_contract + lhs_non_shape = tuple(lhs_shape[dim] for dim in lhs_non_contract) + rhs_non_shape = tuple(rhs_shape[dim] for dim in rhs_non_contract) + lhs_matrix_shape = batch_shape + (_product(lhs_non_shape), _product(lhs_contract_shape)) + rhs_matrix_shape = batch_shape + (_product(rhs_contract_shape), _product(rhs_non_shape)) + output_shape = batch_shape + lhs_non_shape + rhs_non_shape + dtype = getattr(eqn.outvars[0].aval, "dtype", getattr(eqn.invars[0].aval, "dtype", np.float32)) + + lhs_id = inputs[0] + if lhs_order != tuple(range(len(lhs_shape))): + lhs_id = f"{outputs[0]}__lhs_transpose" + _register_node( + graph, + IRNode( + ctx.node_id("dot_general_lhs_transpose"), + "permute", + [inputs[0]], + [lhs_id], + attrs={"permutation": lhs_order}, + meta=meta, + ), + out_avals=(_SyntheticAval(tuple(lhs_shape[dim] for dim in lhs_order), dtype),), + ) + lhs_2d = f"{outputs[0]}__lhs_reshape" + _register_node( + graph, + IRNode( + ctx.node_id("dot_general_lhs_reshape"), + "reshape", + [lhs_id], + [lhs_2d], + attrs={"shape": lhs_matrix_shape}, + meta=meta, + ), + out_avals=(_SyntheticAval(lhs_matrix_shape, dtype),), + ) + + rhs_id = inputs[1] + if rhs_order != tuple(range(len(rhs_shape))): + rhs_id = f"{outputs[0]}__rhs_transpose" + _register_node( + graph, + IRNode( + ctx.node_id("dot_general_rhs_transpose"), + "permute", + [inputs[1]], + [rhs_id], + attrs={"permutation": rhs_order}, + meta=meta, + ), + out_avals=(_SyntheticAval(tuple(rhs_shape[dim] for dim in rhs_order), dtype),), + ) + rhs_2d = f"{outputs[0]}__rhs_reshape" + _register_node( + graph, + IRNode( + ctx.node_id("dot_general_rhs_reshape"), + "reshape", + [rhs_id], + [rhs_2d], + attrs={"shape": rhs_matrix_shape}, + meta=meta, + ), + out_avals=(_SyntheticAval(rhs_matrix_shape, dtype),), + ) + + matmul_id = f"{outputs[0]}__matmul" + _register_node( + graph, + IRNode(ctx.node_id("dot_general_matmul"), "matmul", [lhs_2d, rhs_2d], [matmul_id], meta=meta), + out_avals=(_SyntheticAval(batch_shape + (lhs_matrix_shape[-2], rhs_matrix_shape[-1]), dtype),), + ) + return IRNode( + ctx.node_id("dot_general_output_reshape"), + "reshape", + [matmul_id], + outputs, + attrs={"shape": output_shape}, + meta=meta, + ) + return IRNode(node_id, "matmul", inputs, outputs, meta=meta) + if prim == "conv_general_dilated": + if len(inputs) != 2 or len(outputs) != 1: + raise NotImplementedError("JAX conv_general_dilated import supports input and kernel only") + if int(params.get("batch_group_count", 1)) != 1: + raise NotImplementedError("JAX conv_general_dilated batch_group_count is unsupported") + lhs_dilation = tuple(int(value) for value in params.get("lhs_dilation", ())) + if any(value != 1 for value in lhs_dilation): + raise NotImplementedError(f"JAX conv_general_dilated lhs_dilation is unsupported: {lhs_dilation}") + dimension_numbers = params.get("dimension_numbers") + if dimension_numbers is None: + raise NotImplementedError("JAX conv_general_dilated requires static dimension_numbers") + + lhs_spec = tuple(int(axis) for axis in dimension_numbers.lhs_spec) + rhs_spec = tuple(int(axis) for axis in dimension_numbers.rhs_spec) + out_spec = tuple(int(axis) for axis in dimension_numbers.out_spec) + lhs_shape = tuple(int(dim) for dim in getattr(eqn.invars[0].aval, "shape", ())) + rhs_shape = tuple(int(dim) for dim in getattr(eqn.invars[1].aval, "shape", ())) + out_shape = tuple(int(dim) for dim in getattr(eqn.outvars[0].aval, "shape", ())) + rank = len(lhs_shape) + if rank not in {3, 4}: + raise NotImplementedError(f"JAX conv_general_dilated supports 1D/2D convs, got rank {rank}") + if tuple(sorted(lhs_spec)) != tuple(range(rank)) or tuple(sorted(rhs_spec)) != tuple(range(rank)): + raise NotImplementedError("JAX conv_general_dilated requires complete dimension specs") + if tuple(sorted(out_spec)) != tuple(range(rank)): + raise NotImplementedError("JAX conv_general_dilated requires complete output dimension spec") + + stride = _conv_static_int(params.get("window_strides", (1,) * (rank - 2)), name="stride") + dilation = _conv_static_int(params.get("rhs_dilation", (1,) * (rank - 2)), name="rhs_dilation") + padding = _conv_padding_int(params.get("padding", ((0, 0),) * (rank - 2))) + groups = int(params.get("feature_group_count", 1)) + dtype = getattr(eqn.outvars[0].aval, "dtype", getattr(eqn.invars[0].aval, "dtype", np.float32)) + + x_id = inputs[0] + x_shape = tuple(lhs_shape[axis] for axis in lhs_spec) + if lhs_spec != tuple(range(rank)): + x_id = f"{outputs[0]}__conv_lhs" + _register_node( + graph, + IRNode( + ctx.node_id("conv_lhs_transpose"), + "permute", + [inputs[0]], + [x_id], + attrs={"permutation": lhs_spec}, + meta=meta, + ), + out_avals=(_SyntheticAval(x_shape, dtype),), + ) + + weight_id = inputs[1] + weight_shape = tuple(rhs_shape[axis] for axis in rhs_spec) + if rhs_spec != tuple(range(rank)): + weight_id = f"{outputs[0]}__conv_rhs" + _register_node( + graph, + IRNode( + ctx.node_id("conv_rhs_transpose"), + "permute", + [inputs[1]], + [weight_id], + attrs={"permutation": rhs_spec}, + meta=meta, + ), + out_avals=(_SyntheticAval(weight_shape, dtype),), + ) + + conv_id = outputs[0] + output_permutation = _conv_output_permutation(out_spec) + conv_shape = tuple(out_shape[axis] for axis in out_spec) + if output_permutation != tuple(range(rank)): + conv_id = f"{outputs[0]}__conv" + _register_node( + graph, + IRNode( + node_id, + "conv1d" if rank == 3 else "conv2d", + [x_id, weight_id], + [conv_id], + attrs={"stride": stride, "padding": padding, "dilation": dilation, "groups": groups}, + meta=meta, + ), + out_avals=(_SyntheticAval(conv_shape, dtype),), + ) + if conv_id == outputs[0]: + return None + return IRNode( + ctx.node_id("conv_output_transpose"), + "permute", + [conv_id], + outputs, + attrs={"permutation": output_permutation}, + meta=meta, + ) + if prim == "reshape": + if inputs[0] in ctx.gather_index_aliases: + ctx.gather_index_aliases[outputs[0]] = ctx.gather_index_aliases[inputs[0]] + new_shape = tuple(int(dim) for dim in params["new_sizes"]) + input_shape = tuple(int(dim) for dim in getattr(eqn.invars[0].aval, "shape", ())) + source_value = graph.values.get(inputs[0]) + graph_input_shape = tuple(int(dim) for dim in (source_value.shape or ())) if source_value is not None else input_shape + if new_shape == () and graph_input_shape == (1,): + _alias_output(ctx, eqn.outvars[0], inputs[0]) + return None + return IRNode( + node_id, + "reshape", + inputs, + outputs, + attrs={"shape": new_shape}, + meta=meta, + ) + if prim == "squeeze": + shape = tuple(int(dim) for dim in getattr(eqn.outvars[0].aval, "shape", ())) + return IRNode(node_id, "reshape", inputs, outputs, attrs={"shape": shape}, meta=meta) + if prim == "transpose": + permutation = params.get("permutation") + if permutation is None: + permutation = params.get("permutation_or_none") + if permutation is None: + raise NotImplementedError("JAX transpose missing permutation") + return IRNode( + node_id, + "permute", + inputs, + outputs, + attrs={"permutation": tuple(int(dim) for dim in permutation)}, + meta=meta, + ) + if prim == "convert_element_type": + array = _constant_array(graph, inputs[0]) + if array is not None: + _add_constant( + graph, + ctx, + eqn.outvars[0], + array.astype(np.dtype(params["new_dtype"])), + source_name="folded:convert_element_type", + meta=_derived_meta(inputs, op=prim), + ) + return None + return IRNode( + node_id, + "precision_cast", + inputs, + outputs, + attrs={"dtype": _dtype_to_ir(params["new_dtype"])}, + meta=meta, + ) + if prim == "broadcast_in_dim": + shape = tuple(int(dim) for dim in params["shape"]) + scalar = _constant_singleton_scalar(graph, inputs[0]) + if scalar is not None: + _add_constant( + graph, + ctx, + eqn.outvars[0], + np.full(shape, scalar, dtype=np.float32), + source_name="folded:broadcast_in_dim", + meta=_derived_meta(inputs, op=prim), + ) + return None + if inputs[0] in ctx.gather_index_aliases: + ctx.gather_index_aliases[outputs[0]] = ctx.gather_index_aliases[inputs[0]] + broadcast_dimensions = tuple(int(dim) for dim in params.get("broadcast_dimensions", ())) + input_shape = tuple(graph.values[inputs[0]].shape or ()) + if broadcast_dimensions and broadcast_dimensions != tuple(range(len(input_shape))): + reshape_shape = [1] * len(shape) + for input_axis, output_axis in enumerate(broadcast_dimensions): + reshape_shape[output_axis] = input_shape[input_axis] + reshaped_shape = tuple(reshape_shape) + if reshaped_shape == shape: + return IRNode(node_id, "view", inputs, outputs, attrs={"shape": reshaped_shape}, meta=meta) + reshaped_id = _generated_value( + graph, + ctx, + stem="broadcast_base", + shape=reshaped_shape, + dtype=graph.values[inputs[0]].dtype, + producer=node_id, + meta=_derived_meta(inputs, op=f"{prim}:reshape"), + ) + reshape_node = IRNode(node_id, "view", inputs, [reshaped_id], attrs={"shape": reshaped_shape}, meta=meta) + _register_generated_node( + graph, + reshape_node, + output_shapes=(reshaped_shape,), + output_dtypes=(graph.values[inputs[0]].dtype,), + ) + node_id = ctx.node_id(prim) + return IRNode(node_id, "expand", [reshaped_id], outputs, attrs={"shape": shape}, meta=meta) + return IRNode(node_id, "expand", inputs, outputs, attrs={"shape": shape}, meta=meta) + if prim == "concatenate": + return IRNode(node_id, "cat", inputs, outputs, attrs={"axis": int(params["dimension"])}, meta=meta) + if prim == "stack": + axis = int(params.get("axis", 0)) + input_shape = tuple(getattr(eqn.invars[0].aval, "shape", ())) + rank = len(input_shape) + 1 + if axis < 0: + axis += rank + if axis < 0 or axis > len(input_shape): + raise NotImplementedError(f"JAX stack axis out of range: {axis}") + dtype = getattr(eqn.outvars[0].aval, "dtype", getattr(eqn.invars[0].aval, "dtype", np.float32)) + expanded_inputs: list[str] = [] + for input_index, input_id in enumerate(inputs): + expanded_id = f"{outputs[0]}__stack_{input_index}" + expanded_shape = input_shape[:axis] + (1,) + input_shape[axis:] + _register_node( + graph, + IRNode( + ctx.node_id("stack_view"), + "view", + [input_id], + [expanded_id], + attrs={"shape": expanded_shape}, + meta=meta, + ), + out_avals=(_SyntheticAval(expanded_shape, dtype),), + ) + expanded_inputs.append(expanded_id) + return IRNode(node_id, "cat", expanded_inputs, outputs, attrs={"axis": axis}, meta=meta) + if prim == "tile": + reps = tuple(int(value) for value in params.get("reps", ())) + if any(rep <= 0 for rep in reps): + raise NotImplementedError(f"JAX tile import requires positive reps, got {reps}") + input_shape = tuple(int(dim) for dim in getattr(eqn.invars[0].aval, "shape", ())) + if len(reps) < len(input_shape): + reps = (1,) * (len(input_shape) - len(reps)) + reps + current_id = inputs[0] + current_shape = input_shape + dtype = getattr(eqn.invars[0].aval, "dtype", getattr(eqn.outvars[0].aval, "dtype", np.float32)) + if len(reps) > len(current_shape): + current_shape = (1,) * (len(reps) - len(current_shape)) + current_shape + viewed_id = f"{outputs[0]}__tile_view" + _register_node( + graph, + IRNode( + ctx.node_id("tile_view"), + "view", + [current_id], + [viewed_id], + attrs={"shape": current_shape}, + meta=meta, + ), + out_avals=(_SyntheticAval(current_shape, dtype),), + ) + current_id = viewed_id + for axis, rep in enumerate(reps): + if rep == 1: + continue + next_id = outputs[0] if axis == len(reps) - 1 and all(factor == 1 for factor in reps[axis + 1 :]) else f"{outputs[0]}__tile_axis_{axis}" + next_shape = current_shape[:axis] + (current_shape[axis] * rep,) + current_shape[axis + 1 :] + _register_node( + graph, + IRNode( + ctx.node_id("tile_cat"), + "cat", + [current_id] * rep, + [next_id], + attrs={"axis": axis}, + meta=meta, + ), + out_avals=(_SyntheticAval(next_shape, dtype),), + ) + current_id = next_id + current_shape = next_shape + if current_id != outputs[0]: + _alias_output(ctx, eqn.outvars[0], current_id) + return None + if prim == "split": + axis = int(params.get("axis", 0)) + start = 0 + for output_id, output_var, size in zip(outputs, eqn.outvars, params.get("sizes", ()), strict=True): + end = start + int(size) + _register_node( + graph, + IRNode( + ctx.node_id("split_slice"), + "slice", + inputs, + [output_id], + attrs={"axis": axis, "start": start, "end": end, "step": 1}, + meta=meta, + ), + out_avals=(output_var.aval,), + ) + start = end + return None + if prim == "slice": + starts = tuple(int(value) for value in params["start_indices"]) + limits = tuple(int(value) for value in params["limit_indices"]) + raw_strides = params.get("strides") + strides = (1,) * len(starts) if raw_strides is None else tuple(int(value) for value in raw_strides) + input_shape = getattr(eqn.invars[0].aval, "shape", ()) + changed_axes = [axis for axis, (start, limit, stride) in enumerate(zip(starts, limits, strides, strict=True)) if start != 0 or limit != input_shape[axis] or stride != 1] + if not changed_axes: + _alias_output(ctx, eqn.outvars[0], inputs[0]) + return None + if len(changed_axes) != 1: + raise NotImplementedError("JAX slice import supports one non-trivial axis") + axis = changed_axes[0] + return IRNode( + node_id, + "slice", + inputs, + outputs, + attrs={"axis": axis, "start": starts[axis], "end": limits[axis], "step": strides[axis]}, + meta=meta, + ) + if prim == "select_n": + if len(inputs) != 3: + raise NotImplementedError("JAX select_n import supports ternary select") + ctx.gather_index_aliases[outputs[0]] = inputs[1] + false_scalar = _constant_scalar(graph, inputs[1]) + true_scalar = _constant_scalar(graph, inputs[2]) + where_inputs = [inputs[0]] + attrs: dict[str, object] = {} + if true_scalar is None: + where_inputs.append(inputs[2]) + else: + attrs["true_is_scalar"] = True + attrs["true_value"] = _where_scalar_value(true_scalar) + if false_scalar is None: + where_inputs.append(inputs[1]) + else: + attrs["false_is_scalar"] = True + attrs["false_value"] = _where_scalar_value(false_scalar) + return IRNode(node_id, "where", where_inputs, outputs, attrs=attrs, meta=meta) + if prim == "gather": + if len(inputs) != 2: + raise NotImplementedError("JAX gather import supports embedding-style gather") + indices_id = ctx.gather_index_aliases.get(inputs[1], inputs[1]) + index_shape = tuple(getattr(eqn.invars[1].aval, "shape", ())) + if index_shape and index_shape[-1] == 1: + squeezed_id = f"{indices_id}__squeezed" + squeeze_shape = tuple(int(dim) for dim in index_shape[:-1]) + if squeezed_id not in graph.values: + squeeze_node = IRNode( + node_id, + "reshape", + [indices_id], + [squeezed_id], + attrs={"shape": squeeze_shape}, + meta=meta, + ) + _register_node(graph, squeeze_node, out_avals=(eqn.invars[1].aval,)) + graph.values[squeezed_id].shape = squeeze_shape + if indices_id in ctx.gather_index_aliases: + ctx.gather_index_aliases[squeezed_id] = ctx.gather_index_aliases[indices_id] + indices_id = squeezed_id + node_id = ctx.node_id("gather_embedding") + indices_id = ctx.gather_index_aliases.get(indices_id, indices_id) + return IRNode(node_id, "embedding", [inputs[0], indices_id], outputs, meta=meta) + if prim in {"max", "min"}: + lhs_literal = _literal_number(eqn.invars[0]) + rhs_literal = _literal_number(eqn.invars[1]) + identity = -math.inf if prim == "max" else math.inf + if lhs_literal == identity: + _alias_output(ctx, eqn.outvars[0], inputs[1]) + return None + if rhs_literal == identity: + _alias_output(ctx, eqn.outvars[0], inputs[0]) + return None + + lhs_scalar = _constant_scalar(graph, inputs[0]) + rhs_scalar = _constant_scalar(graph, inputs[1]) + condition_id = f"{outputs[0]}__{prim}_condition" + condition_shape = tuple(getattr(eqn.outvars[0].aval, "shape", ())) + condition_dtype = np.bool_ + if rhs_scalar is not None: + compare_op = "scalar_greater" if prim == "max" else "scalar_less" + _register_node( + graph, + IRNode( + ctx.node_id(f"{prim}_condition"), + compare_op, + [inputs[0]], + [condition_id], + attrs={"value": rhs_scalar}, + meta=meta, + ), + out_avals=(_SyntheticAval(condition_shape, condition_dtype),), + ) + attrs = {"false_is_scalar": True, "false_value": _where_scalar_value(rhs_scalar)} + return IRNode(ctx.node_id(f"{prim}_where"), "where", [condition_id, inputs[0]], outputs, attrs=attrs, meta=meta) + if lhs_scalar is not None: + compare_op = "scalar_less" if prim == "max" else "scalar_greater" + _register_node( + graph, + IRNode( + ctx.node_id(f"{prim}_condition"), + compare_op, + [inputs[1]], + [condition_id], + attrs={"value": lhs_scalar}, + meta=meta, + ), + out_avals=(_SyntheticAval(condition_shape, condition_dtype),), + ) + attrs = {"true_is_scalar": True, "true_value": _where_scalar_value(lhs_scalar)} + return IRNode(ctx.node_id(f"{prim}_where"), "where", [condition_id, inputs[1]], outputs, attrs=attrs, meta=meta) + + compare_op = "greater" if prim == "max" else "less" + _register_node( + graph, + IRNode(ctx.node_id(f"{prim}_condition"), compare_op, inputs, [condition_id], meta=meta), + out_avals=(_SyntheticAval(condition_shape, condition_dtype),), + ) + return IRNode(ctx.node_id(f"{prim}_where"), "where", [condition_id, inputs[0], inputs[1]], outputs, meta=meta) + if prim in {"reduce_sum", "reduce_max", "reduce_min"}: + reduce_ops = { + "reduce_sum": "sum", + "reduce_max": "max", + "reduce_min": "min", + } + return IRNode( + node_id, + reduce_ops[prim], + inputs, + outputs, + attrs={"axis": tuple(int(axis) for axis in params.get("axes", ()))}, + meta=meta, + ) + + raise NotImplementedError(f"unsupported JAX primitive: {prim}") + + +def capture_jax_function( + fn: Callable[..., Any], + example_args: Sequence[Any], + *, + constant_names: Sequence[str] | None = None, + weight_bindings: dict[str, dict[str, str]] | None = None, + weights_dir: str | None = None, + graph_meta: dict[str, object] | None = None, +) -> IRGraph: + try: + import jax + except Exception as exc: # pragma: no cover - depends on optional dependency + raise RuntimeError("capture_jax_function requires the optional jax package") from exc + + closed = jax.make_jaxpr(fn)(*example_args) + jaxpr = closed.jaxpr + constants = tuple(getattr(closed, "consts", ()) or ()) + names = tuple(constant_names or ()) + bindings = weight_bindings or {} + ctx = _JaxImportContext() + graph = IRGraph( + values={}, + nodes={}, + order=[], + inputs=[], + outputs=[], + constants={}, + meta={ + "frontend": "jax", + "adapter_family": "generic", + **({"weights_dir": weights_dir} if weights_dir else {}), + **dict(graph_meta or {}), + }, + ) + + for var in jaxpr.invars: + value_id = ctx.value_id(var) + _add_value(graph, value_id, var.aval) + graph.inputs.append(value_id) + + for index, (var, value) in enumerate(zip(jaxpr.constvars, constants, strict=True)): + value_id = ctx.value_id(var) + name = names[index] if index < len(names) else f"const_{index}" + tensor = _constant_to_torch(value) + graph.add_value( + IRValue( + id=value_id, + shape=tuple(int(dim) for dim in tensor.shape), + dtype=_dtype_to_ir(tensor.numpy().dtype), + meta={ + "source_name": name, + "jax_closed_constant": True, + **_binding_meta(name=name, weights_dir=weights_dir, explicit=bindings), + }, + ) + ) + graph.constants[value_id] = tensor + if "path" in graph.values[value_id].meta: + graph.meta.setdefault("weight_bindings", {})[value_id] = dict(graph.values[value_id].meta) + + for eqn in jaxpr.eqns: + _import_eqn(graph, ctx, eqn) + + graph.outputs = [ctx.value_id(var) for var in jaxpr.outvars] + apply_jax_semantic_rewrites(graph) + verify_ir(graph) + return graph + + +def capture_jax_function_with_params( + fn: Callable[..., Any], + params: Any, + example_args: Sequence[Any], + *, + weights_dir: str | None = None, + weight_bindings: dict[str, dict[str, str]] | None = None, + graph_meta: dict[str, object] | None = None, +) -> IRGraph: + try: + import jax + except Exception as exc: # pragma: no cover - depends on optional dependency + raise RuntimeError("capture_jax_function_with_params requires the optional jax package") from exc + + named_leaves = _flatten_named_leaves(jax.tree_util, params) + flat_leaves, treedef = jax.tree_util.tree_flatten(params) + if len(flat_leaves) != len(named_leaves): + raise ValueError("JAX param flattening produced mismatched leaf counts") + leaf_count = len(flat_leaves) + + def bound_fn(*flat_params_and_args: Any) -> Any: + flat_param_values = flat_params_and_args[:leaf_count] + runtime_args = flat_params_and_args[leaf_count:] + rebuilt_params = jax.tree_util.tree_unflatten(treedef, flat_param_values) + return fn(rebuilt_params, *runtime_args) + + graph = capture_jax_function( + bound_fn, + (*flat_leaves, *example_args), + weight_bindings=weight_bindings, + weights_dir=weights_dir, + graph_meta=graph_meta, + ) + explicit = weight_bindings or {} + _freeze_leading_param_inputs( + graph, + named_leaves, + weights_dir=weights_dir, + explicit=explicit, + ) + _propagate_weight_binding_meta(graph) + verify_ir(graph) + return graph + + +def capture_jax_graphs( + params: Any, + specs: Sequence[JaxGraphSpec], + *, + weights_dir: str | None = None, + weight_bindings: dict[str, dict[str, str]] | None = None, + graph_meta: dict[str, object] | None = None, +) -> CapturedJaxGraphBundle: + seen_names: set[str] = set() + for spec in specs: + if spec.name in seen_names: + raise ValueError(f"duplicate JAX graph spec name: {spec.name!r}") + seen_names.add(spec.name) + if spec.input_names is not None and len(spec.input_names) != len(spec.example_args): + raise ValueError( + f"graph {spec.name!r} has {len(spec.input_names)} input names for " + f"{len(spec.example_args)} example inputs" + ) + + captured: dict[str, CapturedJaxGraph] = {} + for spec in specs: + ir = capture_jax_function_with_params( + spec.fn, + params, + tuple(spec.example_args), + weights_dir=weights_dir, + weight_bindings=weight_bindings, + graph_meta={ + **dict(graph_meta or {}), + **dict(spec.graph_meta or {}), + "jax_graph_name": spec.name, + "jax_graph_role": spec.role, + **({"input_names": tuple(spec.input_names)} if spec.input_names is not None else {}), + **({"output_names": tuple(spec.output_names)} if spec.output_names is not None else {}), + }, + ) + raw_ir = copy.deepcopy(ir) + transpiled_graph = transpile_ir(ir) + captured[spec.name] = CapturedJaxGraph( + spec=spec, + raw_ir_graph=raw_ir, + ir_graph=ir, + graph=transpiled_graph, + ) + return CapturedJaxGraphBundle(graphs=captured, params=params, weights_dir=weights_dir) + + +__all__ = [ + "CapturedJaxGraph", + "CapturedJaxGraphBundle", + "JaxGraphSpec", + "capture_jax_function", + "capture_jax_function_with_params", + "capture_jax_graphs", +] diff --git a/python/cactus/transpile/capture_pytorch.py b/python/cactus/transpile/capture_pytorch.py new file mode 100644 index 000000000..22c63833e --- /dev/null +++ b/python/cactus/transpile/capture_pytorch.py @@ -0,0 +1,375 @@ +from dataclasses import dataclass +import inspect +from contextlib import contextmanager +from typing import Any +from typing import Iterable + +import torch +from torch.export import export +from torch.fx.passes.shape_prop import ShapeProp + +from cactus.transpile.aten_ops import canonical_torch_op +from cactus.transpile.graph_ir import IRGraph + + +class CapturePhaseError(RuntimeError): + def __init__(self, phase: str, message: str, *, cause: Exception | None = None): + super().__init__(f"[capture:{phase}] {message}") + self.phase = phase + self.cause = cause + + +@dataclass +class TensorMetadata: + shape: tuple[Any, ...] + dtype: Any + stride: tuple[Any, ...] | None = None + device: Any | None = None + + +@dataclass +class CapturedModel: + exported_program: Any + graph_module: Any + graph: Any + state_dict: dict[str, Any] + ir_graph: IRGraph + source_module: Any + example_args: tuple[Any, ...] + example_kwargs: dict[str, Any] + strict: bool + transpile_metadata: dict[str, Any] + + def named_parameters(self): + return self.exported_program.named_parameters() + + def named_buffers(self): + return self.exported_program.named_buffers() + + def parameters_dict(self) -> dict[str, Any]: + return dict(self.named_parameters()) + + def buffers_dict(self) -> dict[str, Any]: + return dict(self.named_buffers()) + + def placeholders(self) -> list[Any]: + return [node for node in self.graph.nodes if node.op == "placeholder"] + + def outputs(self) -> list[Any]: + return [node for node in self.graph.nodes if node.op == "output"] + + +def _normalize_args(args: Any) -> tuple[Any, ...]: + if args is None: + return () + if isinstance(args, tuple): + return args + if isinstance(args, list): + return tuple(args) + return (args,) + + +def _normalize_kwargs(kwargs: dict[str, Any] | None) -> dict[str, Any]: + return {} if kwargs is None else dict(kwargs) + + +def _clone_example_value(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().clone() + if isinstance(value, tuple): + return tuple(_clone_example_value(v) for v in value) + if isinstance(value, list): + return [_clone_example_value(v) for v in value] + if isinstance(value, dict): + return {k: _clone_example_value(v) for k, v in value.items()} + return value + + +def _clone_examples(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[tuple[Any, ...], dict[str, Any]]: + return ( + tuple(_clone_example_value(arg) for arg in args), + {k: _clone_example_value(v) for k, v in kwargs.items()}, + ) + + +def _inject_export_safe_kwargs(model: torch.nn.Module, kwargs: dict[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(model.forward) + except (TypeError, ValueError): + return kwargs + + updated = dict(kwargs) + parameters = signature.parameters + + # HF causal LM forwards often default to returning cache objects and ModelOutput + # containers, which torch.export cannot serialize cleanly. Keep the capture + # boundary tensor-only unless the caller explicitly asked otherwise. + if "use_cache" in parameters and "use_cache" not in updated: + updated["use_cache"] = False + if "return_dict" in parameters and "return_dict" not in updated: + updated["return_dict"] = False + return updated + + +def _call_transpile_metadata_provider(module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: + for attr_name in ("get_transpile_metadata", "transpile_metadata"): + provider = getattr(module, attr_name, None) + if not callable(provider): + continue + try: + metadata = provider(args=args, kwargs=kwargs) + except TypeError: + metadata = provider() + if metadata is None: + return {} + if not isinstance(metadata, dict): + raise TypeError(f"{type(module).__name__}.{attr_name} must return a dict or None") + return metadata + return {} + + +def _collect_transpile_metadata(model: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: + graph_meta: dict[str, Any] = {} + provider_graph_meta: dict[str, dict[str, Any]] = {} + seen_modules: set[int] = set() + + for module_path, module in model.named_modules(): + module_id = id(module) + if module_id in seen_modules: + continue + seen_modules.add(module_id) + + metadata = _call_transpile_metadata_provider(module, args, kwargs) + if not metadata: + continue + + provider_key = module_path or type(module).__name__ + provider_meta = metadata.get("graph", {}) + if isinstance(provider_meta, dict) and provider_meta: + provider_graph_meta[provider_key] = dict(provider_meta) + if not module_path: + graph_meta.update(provider_meta) + + if provider_graph_meta: + graph_meta.setdefault("transpile_metadata_providers", provider_graph_meta) + + if not graph_meta: + return {} + return {"graph": graph_meta} + + +def _prepare_import_graph_module(ep: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + graph_module = ep.graph_module + if kwargs: + return graph_module + + try: + graph_module = ep.module() + ShapeProp(graph_module).propagate(*args) + except Exception: + return ep.graph_module + return graph_module + + +@contextmanager +def _suppress_transformers_model_output_registration(): + """Avoid tracing a Transformers global-set mutation inside ModelOutput init. + + Recent Transformers versions lazily register ModelOutput dataclasses as + PyTrees from their __post_init__. That side effect is irrelevant for our + capture wrapper because the exported boundary returns tensors, but Dynamo + sees the intermediate set membership check and rejects it. + """ + + try: + import transformers.utils.generic as hf_generic + except Exception: + yield + return + + original = getattr(hf_generic, "_register_model_output_pytree_node", None) + if original is None: + yield + return + + def _noop_register_model_output_pytree_node(*_args: Any, **_kwargs: Any) -> None: + return None + + hf_generic._register_model_output_pytree_node = _noop_register_model_output_pytree_node + try: + yield + finally: + hf_generic._register_model_output_pytree_node = original + + +def capture_model(model, args, kwargs=None, *, strict=True) -> CapturedModel: + if not isinstance(model, torch.nn.Module): + raise TypeError("model must be a torch.nn.Module") + + normalized_args = _normalize_args(args) + normalized_kwargs = _inject_export_safe_kwargs(model, _normalize_kwargs(kwargs)) + + if not normalized_args and not normalized_kwargs: + raise ValueError("capture_model requires example args or kwargs") + + # nn.Module.eval() is an in-place mutation on the caller's model. Remember + # the prior training flag so capture failures don't leave the caller's + # module silently in eval mode. + was_training = bool(getattr(model, "training", False)) + model.eval() + try: + example_args, example_kwargs = _clone_examples(normalized_args, normalized_kwargs) + transpile_metadata = _collect_transpile_metadata(model, example_args, example_kwargs) + + try: + with _suppress_transformers_model_output_registration(): + ep = export(model, args=example_args, kwargs=example_kwargs, strict=strict) + except Exception as exc: + raise CapturePhaseError( + "export", + f"torch.export failed for model={type(model).__name__} strict={strict}: {exc}", + cause=exc, + ) from exc + from cactus.transpile.import_ir import import_captured_to_ir + + import_graph_module = _prepare_import_graph_module(ep, example_args, example_kwargs) + + raw_captured = CapturedModel( + exported_program=ep, + graph_module=import_graph_module, + graph=import_graph_module.graph, + state_dict=dict(ep.state_dict), + ir_graph=IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[]), + source_module=model, + example_args=example_args, + example_kwargs=example_kwargs, + strict=strict, + transpile_metadata=transpile_metadata, + ) + try: + ir_graph = import_captured_to_ir(raw_captured, strict=strict) + except Exception as exc: + raise CapturePhaseError( + "import", + f"failed to import exported graph to IR for model={type(model).__name__} strict={strict}: {exc}", + cause=exc, + ) from exc + except Exception: + if was_training: + try: + model.train() + except Exception: + pass + raise + + return CapturedModel( + exported_program=ep, + graph_module=ep.graph_module, + graph=ep.graph, + state_dict=dict(ep.state_dict), + ir_graph=ir_graph, + source_module=model, + example_args=example_args, + example_kwargs=example_kwargs, + strict=strict, + transpile_metadata=transpile_metadata, + ) + + +def capture_model_with_fallback(model, args, kwargs=None) -> CapturedModel: + try: + return capture_model(model, args, kwargs=kwargs, strict=True) + except Exception: + return capture_model(model, args, kwargs=kwargs, strict=False) + + +def resolve_attr(root: Any, target: str) -> Any: + obj = root + for atom in target.split("."): + obj = getattr(obj, atom) + return obj + + +def format_target(node: Any) -> str: + return canonical_torch_op(node.target) + + +def _try_materialize_int_tuple(values: Iterable[Any]) -> tuple[int, ...] | None: + materialized: list[int] = [] + try: + for value in values: + materialized.append(int(value)) + except Exception: + return None + return tuple(materialized) + + +def get_tensor_metadata(node: Any) -> TensorMetadata | None: + meta = getattr(node, "meta", {}) or {} + tensor_meta = meta.get("tensor_meta") + if tensor_meta is not None: + shape = tuple(getattr(tensor_meta, "shape", ())) + dtype = getattr(tensor_meta, "dtype", None) + stride_value = getattr(tensor_meta, "stride", None) + stride = _try_materialize_int_tuple(stride_value) if isinstance(stride_value, Iterable) else None + device = getattr(tensor_meta, "device", None) + if shape or dtype is not None or stride is not None or device is not None: + return TensorMetadata( + shape=shape, + dtype=dtype, + stride=stride, + device=device, + ) + + value = meta.get("val") + if value is None: + return None + + if isinstance(value, torch.Tensor): + stride = _try_materialize_int_tuple(value.stride()) if value.layout == torch.strided else None + return TensorMetadata( + shape=tuple(value.shape), + dtype=value.dtype, + stride=stride, + device=value.device, + ) + + shape = tuple(getattr(value, "shape", ())) + dtype = getattr(value, "dtype", None) + stride_value = getattr(value, "stride", None) + stride = tuple(stride_value) if isinstance(stride_value, Iterable) else None + device = getattr(value, "device", None) + if not shape and dtype is None and stride is None and device is None: + return None + return TensorMetadata(shape=shape, dtype=dtype, stride=stride, device=device) + + +def get_shape(node: Any) -> tuple[Any, ...] | None: + metadata = get_tensor_metadata(node) + return None if metadata is None else metadata.shape + + +def get_dtype(node: Any) -> Any | None: + metadata = get_tensor_metadata(node) + return None if metadata is None else metadata.dtype + + +def dump_graph(captured: CapturedModel, *, include_meta: bool = True) -> str: + lines: list[str] = [] + for index, node in enumerate(captured.graph.nodes): + lines.append(f"[{index}] {node.op} {node.name} target={node.target}") + lines.append(f" args={node.args}") + lines.append(f" kwargs={node.kwargs}") + if include_meta: + tensor_meta = get_tensor_metadata(node) + if tensor_meta is not None: + lines.append( + " tensor_meta=" + f"shape={tensor_meta.shape} dtype={tensor_meta.dtype} " + f"stride={tensor_meta.stride} device={tensor_meta.device}" + ) + else: + lines.append(f" meta={getattr(node, 'meta', {})}") + return "\n".join(lines) + + diff --git a/python/cactus/transpile/component_partition.py b/python/cactus/transpile/component_partition.py new file mode 100644 index 000000000..a0a76dd32 --- /dev/null +++ b/python/cactus/transpile/component_partition.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +import copy +from collections import Counter + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.model_profiles import profile_for_family + + +COMPONENT_AUDIO_ENCODER = "audio_encoder" +COMPONENT_VISION_ENCODER = "vision_encoder" +COMPONENT_LM_ENCODER = "lm_encoder" +COMPONENT_DECODER = "decoder" +COMPONENT_UNSPECIFIED = "unspecified" + +COMPONENT_ORDER = ( + COMPONENT_AUDIO_ENCODER, + COMPONENT_VISION_ENCODER, + COMPONENT_LM_ENCODER, + COMPONENT_DECODER, +) + + +def _normalized_name_candidates( + *, + value_id: str | None = None, + meta: dict[str, object] | None = None, +) -> tuple[str, ...]: + names: list[str] = [] + + def _add(value: object) -> None: + if not isinstance(value, str): + return + normalized = value.strip().lower() + if normalized and normalized not in names: + names.append(normalized) + + if value_id is not None: + _add(value_id) + if isinstance(meta, dict): + for key in ("logical_name", "input_name", "output_name", "component_role"): + _add(meta.get(key)) + return tuple(names) + + +def _contains_any(name: str, needles: tuple[str, ...]) -> bool: + return any(needle in name for needle in needles) + + +def classify_component_semantic_candidates( + candidates: tuple[str, ...], + *, + task: str, +) -> str | None: + if not candidates: + return None + + for name in candidates: + if _contains_any( + name, + ( + "pixel_values", + "image_features", + ), + ): + return COMPONENT_VISION_ENCODER + + for name in candidates: + if _contains_any( + name, + ( + "input_features", + "audio_features", + "speech_features", + "mel_features", + ), + ): + return COMPONENT_AUDIO_ENCODER + + for name in candidates: + if _contains_any( + name, + ( + "logits", + "next_token", + "ctc_head", + "ctc_logits", + "decoder_output", + ), + ): + return COMPONENT_DECODER + + for name in candidates: + if task == "multimodal_causal_lm_logits" and _contains_any( + name, + ( + "input_ids", + "inputs_embeds", + "text_embeds", + "position_ids", + "placeholder", + "merge", + "masked_scatter", + "token_type", + "attention_mask", + "pixel_position", + ), + ): + return COMPONENT_LM_ENCODER + + return None + + +def classify_value_component(value_id: str, value: IRValue, *, family: str, task: str) -> str | None: + return classify_component_semantic_candidates( + _normalized_name_candidates( + value_id=value_id, + meta=value.meta, + ), + task=task, + ) + + +def classify_node_component(node: IRNode, *, family: str, task: str) -> str | None: + if task == "multimodal_causal_lm_logits" and node.op in {"embedding", "masked_scatter", "advanced_index"}: + return COMPONENT_LM_ENCODER + return classify_component_semantic_candidates( + _normalized_name_candidates(meta=node.meta), + task=task, + ) + + +def _task_default_component(*, family: str, task: str) -> str: + if task == "multimodal_causal_lm_logits": + return COMPONENT_DECODER + if task == "causal_lm_logits": + return COMPONENT_DECODER + if task in {"ctc_logits", "encoder_hidden_states"}: + return COMPONENT_AUDIO_ENCODER + profile = profile_for_family(family) + if profile is not None and profile.default_task == "tdt_transcription": + return COMPONENT_AUDIO_ENCODER + return COMPONENT_UNSPECIFIED + + +def _rebuild_users(graph: IRGraph) -> None: + for value in graph.values.values(): + value.users = [] + for value_id in graph.inputs: + graph.values[value_id].producer = None + for value_id in graph.constants: + graph.values[value_id].producer = None + for node_id in graph.order: + node = graph.nodes[node_id] + for output_id in node.outputs: + graph.values[output_id].producer = node_id + for input_id in node.inputs: + graph.values[input_id].users.append(node_id) + + +def annotate_ir_components(graph: IRGraph) -> None: + family = str(graph.meta.get("adapter_family", "") or "").lower() + task = str(graph.meta.get("task", "") or "").lower() + graph_component = str(graph.meta.get("component", "") or "").strip() + if graph_component in COMPONENT_ORDER: + for value in graph.values.values(): + value.meta["component"] = graph_component + for node_id in graph.order: + node = graph.nodes[node_id] + node.meta["component"] = graph_component + for output_id in node.outputs: + graph.values[output_id].meta["component"] = graph_component + graph.meta["component_counts"] = {graph_component: len(graph.order)} + return + + default_component = _task_default_component(family=family, task=task) + + for value_id, value in graph.values.items(): + component = classify_value_component(value_id, value, family=family, task=task) + if component is not None: + value.meta["component"] = component + + for node_id in graph.order: + node = graph.nodes[node_id] + component = classify_node_component(node, family=family, task=task) + if component is not None: + node.meta["component"] = component + for output_id in node.outputs: + graph.values[output_id].meta["component"] = component + + for _ in range(8): + changed = False + for node_id in graph.order: + node = graph.nodes[node_id] + component = node.meta.get("component") + if component not in COMPONENT_ORDER: + input_components = { + graph.values[input_id].meta.get("component") + for input_id in node.inputs + if graph.values[input_id].meta.get("component") in COMPONENT_ORDER + } + if len(input_components) == 1: + component = next(iter(input_components)) + node.meta["component"] = component + changed = True + component = node.meta.get("component") + if component in COMPONENT_ORDER: + for output_id in node.outputs: + if graph.values[output_id].meta.get("component") != component: + graph.values[output_id].meta["component"] = component + changed = True + + for node_id in reversed(graph.order): + node = graph.nodes[node_id] + component = node.meta.get("component") + if component in COMPONENT_ORDER: + continue + output_components: set[str] = set() + for output_id in node.outputs: + value = graph.values[output_id] + for user_id in value.users: + user_component = graph.nodes[user_id].meta.get("component") + if user_component in COMPONENT_ORDER: + output_components.add(user_component) + if len(output_components) == 1: + component = next(iter(output_components)) + node.meta["component"] = component + for output_id in node.outputs: + graph.values[output_id].meta["component"] = component + changed = True + + if not changed: + break + + for value in graph.values.values(): + component = value.meta.get("component") + if component not in COMPONENT_ORDER and default_component != COMPONENT_UNSPECIFIED: + value.meta["component"] = default_component + + for node_id in graph.order: + node = graph.nodes[node_id] + component = node.meta.get("component") + if component not in COMPONENT_ORDER and default_component != COMPONENT_UNSPECIFIED: + node.meta["component"] = default_component + for output_id in node.outputs: + graph.values[output_id].meta["component"] = default_component + + counts = Counter( + str(graph.nodes[node_id].meta.get("component", COMPONENT_UNSPECIFIED)) + for node_id in graph.order + ) + graph.meta["component_counts"] = dict(sorted(counts.items())) + + +def summarize_ir_components(graph: IRGraph) -> dict[str, int]: + annotate_ir_components(graph) + counts = graph.meta.get("component_counts", {}) + if isinstance(counts, dict): + return {str(key): int(value) for key, value in counts.items()} + return {} + + +def _clone_value_for_subgraph(value: IRValue, *, as_input: bool) -> IRValue: + cloned = copy.deepcopy(value) + if as_input: + cloned.producer = None + cloned.users = [] + return cloned + + +def extract_component_subgraphs(graph: IRGraph) -> dict[str, IRGraph]: + annotate_ir_components(graph) + component_graphs: dict[str, IRGraph] = {} + + for component in COMPONENT_ORDER: + node_ids = [ + node_id + for node_id in graph.order + if graph.nodes[node_id].meta.get("component") == component + ] + if not node_ids: + continue + + node_id_set = set(node_ids) + inputs: list[str] = [] + outputs: list[str] = [] + constants: dict[str, object] = {} + values: dict[str, IRValue] = {} + nodes: dict[str, IRNode] = {} + + for node_id in node_ids: + node = copy.deepcopy(graph.nodes[node_id]) + nodes[node_id] = node + for input_id in node.inputs: + if input_id in graph.constants: + constants[input_id] = graph.constants[input_id] + if input_id not in values: + values[input_id] = _clone_value_for_subgraph(graph.values[input_id], as_input=True) + continue + producer = graph.values[input_id].producer + if producer not in node_id_set and input_id not in inputs: + inputs.append(input_id) + values[input_id] = _clone_value_for_subgraph(graph.values[input_id], as_input=True) + for output_id in node.outputs: + if output_id not in values: + values[output_id] = _clone_value_for_subgraph(graph.values[output_id], as_input=False) + + for node_id in node_ids: + node = graph.nodes[node_id] + for output_id in node.outputs: + value = graph.values[output_id] + if output_id in graph.outputs or any(user_id not in node_id_set for user_id in value.users): + if output_id not in outputs: + outputs.append(output_id) + + component_graph = IRGraph( + values=values, + nodes=nodes, + order=list(node_ids), + inputs=inputs, + outputs=outputs, + constants=constants, + meta={ + **copy.deepcopy(graph.meta), + "component": component, + }, + ) + _rebuild_users(component_graph) + verify_ir(component_graph) + component_graphs[component] = component_graph + + return component_graphs diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py new file mode 100644 index 000000000..f161319cc --- /dev/null +++ b/python/cactus/transpile/component_pipeline.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +import copy +from dataclasses import dataclass +from dataclasses import field +import re +from typing import Any + +import numpy as np +import torch + +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.lower import TranspiledGraph +from cactus.transpile.lower import transpile_preoptimized_ir +from cactus.transpile.optimize_graph import FusionConfig +from cactus.transpile.optimize_graph import optimize_graph +from cactus.transpile.optimize_graph import precompute_rope_tables + + +@dataclass +class ComponentModuleSpec: + component: str + module: torch.nn.Module + example_inputs: tuple[torch.Tensor, ...] + input_keys: tuple[str, ...] + output_keys: tuple[str, ...] + graph_meta: dict[str, object] = field(default_factory=dict) + metadata: dict[str, object] = field(default_factory=dict) + dynamic_batch_axis: int | None = None + + +@dataclass +class CapturedComponent: + component: str + input_keys: tuple[str, ...] + output_keys: tuple[str, ...] + raw_ir_graph: IRGraph + optimized_ir_graph: IRGraph + transpiled_graph: TranspiledGraph + metadata: dict[str, object] = field(default_factory=dict) + + +_GRAPH_ARG_INPUT_RE = re.compile(r"^v_args_(\d+)$") + + +def _align_input_keys_to_graph_inputs( + declared_input_keys: tuple[str, ...], + graph_inputs: tuple[str, ...] | list[str], +) -> tuple[str, ...]: + aligned: list[str] = [] + for value_id in graph_inputs: + match = _GRAPH_ARG_INPUT_RE.match(str(value_id)) + if match is None: + return declared_input_keys + input_index = int(match.group(1)) + if input_index < 0 or input_index >= len(declared_input_keys): + return declared_input_keys + aligned.append(declared_input_keys[input_index]) + return tuple(aligned) + + +def _call_optional_capture_hook(spec: ComponentModuleSpec, attr_name: str) -> None: + hook = getattr(spec.module, attr_name, None) + if not callable(hook): + return + try: + hook( + component=spec.component, + example_inputs=spec.example_inputs, + input_keys=spec.input_keys, + output_keys=spec.output_keys, + ) + except TypeError: + hook() + + +def capture_component_spec( + spec: ComponentModuleSpec, + *, + fusion_config: FusionConfig | None = None, +) -> CapturedComponent: + wrapped = _ComponentCaptureWrapper(spec.module, graph_meta=spec.graph_meta).eval() + capture_error: Exception | None = None + _call_optional_capture_hook(spec, "prepare_for_capture") + try: + captured = capture_model(wrapped, spec.example_inputs) + except Exception as exc: + capture_error = exc + raise + finally: + try: + _call_optional_capture_hook(spec, "restore_after_capture") + except Exception: + if capture_error is None: + raise + raw_ir_graph = copy.deepcopy(captured.ir_graph) + optimized_ir_graph = copy.deepcopy(captured.ir_graph) + canonicalize_exported_graph(optimized_ir_graph) + optimize_graph(optimized_ir_graph, config=fusion_config, precompute_rope=False) + # Bake rope tables only into the lowered graph.cactus; the saved optimized IR stays un-baked. + lowered_ir = copy.deepcopy(optimized_ir_graph) + if precompute_rope_tables(lowered_ir): + canonicalize_exported_graph(lowered_ir) + if spec.dynamic_batch_axis is not None: + axis = int(spec.dynamic_batch_axis) + for vid in lowered_ir.inputs: + v = lowered_ir.values.get(vid) + if v is None or v.shape is None or axis >= len(v.shape): + continue + mask = [False] * len(v.shape) + mask[axis] = True + v.meta["dynamic_dims"] = tuple(mask) + transpiled_graph = transpile_preoptimized_ir(lowered_ir) + if len(spec.output_keys) != len(transpiled_graph.outputs): + raise ValueError( + f"component {spec.component} declared {len(spec.output_keys)} output keys but lowered " + f"{len(transpiled_graph.outputs)} graph outputs" + ) + aligned_input_keys = _align_input_keys_to_graph_inputs(spec.input_keys, optimized_ir_graph.inputs) + return CapturedComponent( + component=spec.component, + input_keys=aligned_input_keys, + output_keys=spec.output_keys, + raw_ir_graph=raw_ir_graph, + optimized_ir_graph=optimized_ir_graph, + transpiled_graph=transpiled_graph, + metadata=dict(spec.metadata), + ) + + +def execute_component_pipeline( + components: list[CapturedComponent], + *, + initial_store: dict[str, Any], +) -> tuple[dict[str, np.ndarray], dict[str, list[np.ndarray]]]: + store: dict[str, np.ndarray] = {} + for key, value in initial_store.items(): + store[key] = _to_numpy(value) + + outputs_by_component: dict[str, list[np.ndarray]] = {} + for component in components: + runtime_inputs = [] + for input_key in component.input_keys: + if input_key not in store: + raise KeyError( + f"component {component.component} is missing pipeline input {input_key!r}" + ) + runtime_inputs.append(store[input_key]) + component.transpiled_graph.set_inputs(runtime_inputs) + raw_outputs = component.transpiled_graph.execute() + numpy_outputs = [output.numpy().copy() for output in raw_outputs] + if len(numpy_outputs) != len(component.output_keys): + raise ValueError( + f"component {component.component} produced {len(numpy_outputs)} outputs, " + f"expected {len(component.output_keys)}" + ) + for output_key, value in zip(component.output_keys, numpy_outputs, strict=True): + store[output_key] = value + outputs_by_component[component.component] = numpy_outputs + + return store, outputs_by_component + + +def initial_store_from_named_tensors( + keys: tuple[str, ...] | list[str], + tensors: tuple[torch.Tensor, ...] | list[torch.Tensor], +) -> dict[str, np.ndarray]: + if len(keys) != len(tensors): + raise ValueError(f"expected {len(keys)} tensors, got {len(tensors)}") + return { + str(key): _to_numpy(tensor) + for key, tensor in zip(keys, tensors, strict=True) + } + + +def _to_numpy(value: Any) -> np.ndarray: + if isinstance(value, np.ndarray): + return np.ascontiguousarray(value) + if isinstance(value, torch.Tensor): + return np.ascontiguousarray(value.detach().cpu().numpy()) + raise TypeError(f"unsupported pipeline value type: {type(value).__name__}") + + +class _ComponentCaptureWrapper(torch.nn.Module): + def __init__(self, module: torch.nn.Module, *, graph_meta: dict[str, object]): + super().__init__() + self.module = module + self.graph_meta = dict(graph_meta) + + def forward(self, *args: torch.Tensor) -> Any: + return self.module(*args) + + def get_transpile_metadata(self) -> dict[str, object]: + metadata: dict[str, object] = {} + provider = getattr(self.module, "get_transpile_metadata", None) + if callable(provider): + provided = provider() + if isinstance(provided, dict): + metadata.update(provided) + graph_meta = {} + existing = metadata.get("graph", {}) + if isinstance(existing, dict): + graph_meta.update(existing) + graph_meta.update(self.graph_meta) + metadata["graph"] = graph_meta + return metadata diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py new file mode 100644 index 000000000..5c8fdd751 --- /dev/null +++ b/python/cactus/transpile/component_plan.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Mapping + +from cactus.transpile.model_profiles import ModelProfile +from cactus.transpile.model_profiles import profile_for_model_id +from cactus.transpile.model_profiles import profile_for_model_type + + +@dataclass(frozen=True) +class ComponentPlan: + task: str + components: tuple[str, ...] = () + default_max_new_tokens: int | None = None + needs_image: bool = False + needs_audio: bool = False + force_component_pipeline: bool = False + + +def _plan_from_profile( + profile: ModelProfile, + config: Mapping[str, object] | None = None, +) -> ComponentPlan | None: + if profile.default_task is None or not profile.default_components: + return None + has_vision: bool | None = None + has_audio: bool | None = None + if config is not None and profile.needs_image: + has_vision = _has_dict_config(config, "vision_config", "visual_config", "image_config") + if not has_vision: + architectures = [str(a).lower() for a in (config.get("architectures") or ())] + for arch_marker, components in profile.text_component_plans: + if any(arch_marker in a for a in architectures): + return ComponentPlan( + task="causal_lm_logits", + components=components, + force_component_pipeline=True, + ) + return ComponentPlan( + task="causal_lm_logits", + components=("decoder", "decoder_step"), + force_component_pipeline=True, + ) + if config is not None and profile.needs_audio: + has_audio = _has_dict_config(config, "audio_config", "speech_config", "acoustic_config") + components = profile.default_components + needs_audio = profile.needs_audio + if profile.needs_image and profile.needs_audio and has_vision and has_audio is False: + components = tuple(component for component in components if component != "audio_encoder") + needs_audio = False + return ComponentPlan( + task=profile.default_task, + components=components, + default_max_new_tokens=profile.default_max_new_tokens, + needs_image=profile.needs_image, + needs_audio=needs_audio, + force_component_pipeline=profile.force_component_pipeline, + ) + + +def _load_json(path: Path) -> dict[str, object]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + return value if isinstance(value, dict) else {} + + +def _load_config_txt(path: Path) -> dict[str, str]: + config: dict[str, str] = {} + try: + lines = path.read_text(encoding="utf-8").splitlines() + except Exception: + return config + for line in lines: + if "=" not in line: + continue + key, value = line.split("=", 1) + config[key.strip()] = value.strip() + return config + + +def _architecture_names(config: Mapping[str, object], config_txt: Mapping[str, str] | None = None) -> tuple[str, ...]: + raw_architectures = config.get("architectures", []) + names = tuple(str(value) for value in raw_architectures if isinstance(value, str)) + if names or config_txt is None: + return names + raw_txt = config_txt.get("architectures") + if not raw_txt: + return () + return tuple(value.strip() for value in raw_txt.split(",") if value.strip()) + + +def _has_dict_config(config: Mapping[str, object], *keys: str) -> bool: + return any(isinstance(config.get(key), dict) for key in keys) + + +def _has_positive_txt_int(config_txt: Mapping[str, str] | None, *keys: str) -> bool: + if config_txt is None: + return False + for key in keys: + try: + if int(config_txt.get(key, "0") or 0) > 0: + return True + except ValueError: + continue + return False + + +def _model_type(config: Mapping[str, object], config_txt: Mapping[str, str] | None = None) -> str: + value = config.get("model_type") + if value is None and config_txt is not None: + value = config_txt.get("model_type") + return str(value or "").strip().lower() + + +def _lowered_architectures(config: Mapping[str, object], config_txt: Mapping[str, str] | None = None) -> tuple[str, ...]: + return tuple(value.lower() for value in _architecture_names(config, config_txt)) + + +def _is_tdt_config(config: Mapping[str, object], model_type: str, lowered_id: str) -> bool: + decoding_cfg = config.get("decoding") + if isinstance(decoding_cfg, Mapping) and str(decoding_cfg.get("model_type", "") or "").lower() == "tdt": + return True + loss_cfg = config.get("loss") + if isinstance(loss_cfg, Mapping) and str(loss_cfg.get("loss_name", "") or "").lower() == "tdt": + return True + return model_type == "parakeet_tdt" or "parakeet-tdt" in lowered_id + + +def _looks_like_vision_language_model( + *, + model_type: str, + architectures: tuple[str, ...], + lowered_id: str, + has_vision: bool, +) -> bool: + if not has_vision: + return False + if any(token in model_type for token in ("vl", "vision", "image", "multimodal")): + return True + if any( + token in architecture + for architecture in architectures + for token in ("vision", "image", "vl", "imagetexttotext", "conditionalgeneration") + ): + return True + return any(token in lowered_id for token in ("-vl", "_vl", "vision", "image")) + + +def infer_component_plan_from_config( + config: Mapping[str, object], + *, + model_id: str = "", + config_txt: Mapping[str, str] | None = None, +) -> ComponentPlan | None: + model_type = _model_type(config, config_txt) + lowered_id = model_id.lower() + architectures = _lowered_architectures(config, config_txt) + profile = profile_for_model_type(model_type) or profile_for_model_id(model_id) + profile_plan = _plan_from_profile(profile, config) if profile is not None else None + if profile_plan is not None: + return profile_plan + + if "nomic" in model_type or "nomic-embed" in lowered_id or any("nomicbert" in value for value in architectures): + return ComponentPlan( + task="text_embedding", + components=("text_embedding",), + force_component_pipeline=True, + ) + + if _is_tdt_config(config, model_type, lowered_id): + return ComponentPlan( + task="tdt_transcription", + components=("audio_encoder", "decoder"), + needs_audio=True, + force_component_pipeline=True, + ) + + if model_type == "whisper" or "whisper" in lowered_id or any("whisper" in value for value in architectures): + return ComponentPlan( + task="seq2seq_transcription", + components=("audio_encoder", "decoder"), + needs_audio=True, + force_component_pipeline=True, + ) + + if any("ctc" in value for value in architectures) or "ctc" in lowered_id: + return ComponentPlan( + task="ctc_logits", + components=("audio_encoder",), + needs_audio=True, + ) + + has_vision = ( + _has_dict_config(config, "vision_config", "visual_config", "image_config") + or _has_positive_txt_int(config_txt, "vision_num_layers") + ) + has_audio = ( + _has_dict_config(config, "audio_config", "speech_config", "acoustic_config") + or ( + _has_dict_config(config, "encoder_config") + and ( + "audio" in model_type + or "speech" in model_type + or "audio_token_index" in config + or any(token in lowered_id for token in ("speech", "audio")) + or any(token in value for value in architectures for token in ("speech", "audio")) + ) + ) + or _has_positive_txt_int(config_txt, "audio_num_layers") + ) + if has_vision or has_audio: + if has_audio or _looks_like_vision_language_model( + model_type=model_type, + architectures=architectures, + lowered_id=lowered_id, + has_vision=has_vision, + ): + components: list[str] = [] + if has_vision: + components.append("vision_encoder") + if has_audio: + components.append("audio_encoder") + components.extend(("lm_encoder", "decoder")) + return ComponentPlan( + task="multimodal_causal_lm_logits", + components=tuple(components), + needs_image=has_vision, + needs_audio=has_audio, + force_component_pipeline=True, + ) + + if any("causallm" in value for value in architectures): + return ComponentPlan(task="causal_lm_logits", components=("decoder",)) + if any(token in lowered_id or token in model_type for token in ("qwen", "gemma", "llama", "mistral", "lfm")): + return ComponentPlan(task="causal_lm_logits", components=("decoder",)) + return None + + +def infer_component_plan_from_output(output_dir: str | Path, *, model_id: str = "") -> ComponentPlan | None: + root = Path(output_dir) + hf_config = _load_json(root / "hf_config.json") + config_txt = _load_config_txt(root / "config.txt") + return infer_component_plan_from_config(hf_config, model_id=model_id, config_txt=config_txt) diff --git a/python/cactus/transpile/fusion/__init__.py b/python/cactus/transpile/fusion/__init__.py new file mode 100644 index 000000000..463478d1c --- /dev/null +++ b/python/cactus/transpile/fusion/__init__.py @@ -0,0 +1,47 @@ +from cactus.transpile.fusion.conv import ConvModuleMatch +from cactus.transpile.fusion.conv import match_conv_module +from cactus.transpile.fusion.attention import AttentionBlockMatch +from cactus.transpile.fusion.attention import AttentionMatch +from cactus.transpile.fusion.attention import SelfAttentionBlockMatch +from cactus.transpile.fusion.attention import match_attention +from cactus.transpile.fusion.attention import match_attention_block +from cactus.transpile.fusion.attention import match_self_attention_block +from cactus.transpile.fusion.deltanet import GatedDeltaNetMatch +from cactus.transpile.fusion.deltanet import match_gated_deltanet +from cactus.transpile.fusion.linear import LinearMatch +from cactus.transpile.fusion.linear import match_linear +from cactus.transpile.fusion.lstm import LSTMCellMatch +from cactus.transpile.fusion.lstm import match_lstm_cell +from cactus.transpile.fusion.mlp import GatedMLPMatch +from cactus.transpile.fusion.mlp import match_gated_mlp +from cactus.transpile.fusion.rel_pos_bias import RelPosBiasMatch +from cactus.transpile.fusion.rel_pos_bias import match_rel_pos_bias +from cactus.transpile.fusion.rms_norm import RMSNormMatch +from cactus.transpile.fusion.rms_norm import match_rms_norm +from cactus.transpile.fusion.rope import RoPEMatch +from cactus.transpile.fusion.rope import match_rope + +__all__ = [ + "ConvModuleMatch", + "AttentionBlockMatch", + "AttentionMatch", + "SelfAttentionBlockMatch", + "GatedDeltaNetMatch", + "GatedMLPMatch", + "LinearMatch", + "LSTMCellMatch", + "RelPosBiasMatch", + "RMSNormMatch", + "RoPEMatch", + "match_attention", + "match_attention_block", + "match_self_attention_block", + "match_conv_module", + "match_gated_deltanet", + "match_gated_mlp", + "match_linear", + "match_lstm_cell", + "match_rel_pos_bias", + "match_rms_norm", + "match_rope", +] diff --git a/python/cactus/transpile/fusion/attention.py b/python/cactus/transpile/fusion/attention.py new file mode 100644 index 000000000..d6a8e14e0 --- /dev/null +++ b/python/cactus/transpile/fusion/attention.py @@ -0,0 +1,550 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.fusion.common import collect_node_ids +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.fusion.linear import match_linear +from cactus.transpile.fusion.rms_norm import match_rms_norm +from cactus.transpile.fusion.rope import match_rope + + +@dataclass(frozen=True) +class AttentionMatch: + query_value_id: str + key_value_id: str + value_value_id: str + source_input_value_ids: tuple[str, str, str] + weight_value_ids: tuple[str | None, str | None, str | None] + has_rope: bool + has_qk_norm: bool + has_gqa_repeat: bool + is_causal: bool + scale: float + window_size: int + node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class AttentionBlockMatch: + attention_node_id: str + query_value_id: str + key_value_id: str + value_value_id: str + mask_value_id: str | None + gate_value_id: str | None + output_projection_weight_value_id: str + output_projection_bias_value_id: str | None + attention_output_shape: tuple[int, ...] + is_causal: bool + additive_mask: bool + scale: float + window_size: int + node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class SelfAttentionBlockMatch: + attention_node_id: str + hidden_value_id: str + query_weight_value_id: str + query_projection_bias_value_id: str | None + query_add_value_id: str | None + rel_query_add_value_id: str | None + key_weight_value_id: str + key_projection_bias_value_id: str | None + value_weight_value_id: str + value_projection_bias_value_id: str | None + mask_value_id: str | None + relative_key_input_value_id: str | None + relative_key_weight_value_id: str | None + relative_key_projection_bias_value_id: str | None + gate_value_id: str | None + output_projection_weight_value_id: str + output_projection_bias_value_id: str | None + query_shape: tuple[int, ...] + key_shape: tuple[int, ...] + value_shape: tuple[int, ...] + relative_key_shape: tuple[int, ...] | None + attention_output_shape: tuple[int, ...] + is_causal: bool + additive_mask: bool + scale: float + rel_pos_scale: float | None + window_size: int + node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class _ProjectedAttentionInputMatch: + hidden_value_id: str + weight_value_id: str + projection_bias_value_id: str | None + projected_shape: tuple[int, ...] + add_value_id: str | None + node_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class _RelativeKeyProjectionMatch: + input_value_id: str + weight_value_id: str + projection_bias_value_id: str | None + projected_shape: tuple[int, ...] + node_ids: tuple[str, ...] + + +def match_attention(graph: IRGraph, node: IRNode) -> AttentionMatch | None: + if node.op not in {"attention", "scaled_dot_product_attention"} or len(node.inputs) < 3: + return None + + q_info = _extract_attention_input(graph, node.inputs[0], role="q") + k_info = _extract_attention_input(graph, node.inputs[1], role="k") + v_info = _extract_attention_input(graph, node.inputs[2], role="v") + if q_info is None or k_info is None or v_info is None: + return None + + has_rope = bool(q_info["has_rope"] and k_info["has_rope"]) + has_qk_norm = bool(q_info["has_rms_norm"] and k_info["has_rms_norm"]) + has_gqa = bool(k_info["has_gqa_repeat"] or v_info["has_gqa_repeat"] or node.attrs.get("enable_gqa", False)) + node_ids = {node.id, *q_info["node_ids"], *k_info["node_ids"], *v_info["node_ids"]} + + return AttentionMatch( + query_value_id=node.inputs[0], + key_value_id=node.inputs[1], + value_value_id=node.inputs[2], + source_input_value_ids=(q_info["source_input"], k_info["source_input"], v_info["source_input"]), + weight_value_ids=(q_info.get("weight_value_id"), k_info.get("weight_value_id"), v_info.get("weight_value_id")), + has_rope=has_rope, + has_qk_norm=has_qk_norm, + has_gqa_repeat=has_gqa, + is_causal=bool(node.attrs.get("is_causal", False)), + scale=float(node.attrs.get("scale", 0.0)), + window_size=int(node.attrs.get("window_size", 0)), + node_ids=tuple(sorted(node_ids)), + ) + + +def match_self_attention_block(graph: IRGraph, node: IRNode) -> SelfAttentionBlockMatch | None: + projection = match_linear(graph, node) + if projection is None: + return None + + output_path = _extract_attention_output_path(graph, projection.input_value_id) + if output_path is None: + return None + attention_node = graph.nodes.get(output_path["attention_node_id"]) + if attention_node is None or attention_node.op not in {"attention", "scaled_dot_product_attention"} or len(attention_node.inputs) < 3: + return None + + query_match = _extract_projected_attention_input(graph, attention_node.inputs[0], role="q") + key_match = _extract_projected_attention_input(graph, attention_node.inputs[1], role="k") + value_match = _extract_projected_attention_input(graph, attention_node.inputs[2], role="v") + if query_match is None or key_match is None or value_match is None: + return None + + if not ( + query_match.hidden_value_id == key_match.hidden_value_id == value_match.hidden_value_id + ): + return None + + mask_value_id = attention_node.inputs[3] if len(attention_node.inputs) > 3 else None + rel_query_add_value_id: str | None = None + relative_key_input_value_id: str | None = None + relative_key_weight_value_id: str | None = None + relative_key_projection_bias_value_id: str | None = None + relative_key_shape: tuple[int, ...] | None = None + rel_pos_scale: float | None = None + rel_pos_node_ids: tuple[str, ...] = () + + if mask_value_id is not None: + rel_pos_node = producer(graph, strip_passthrough(graph, mask_value_id)) + if rel_pos_node is not None and rel_pos_node.op == "rel_pos_bias" and len(rel_pos_node.inputs) == 2: + rel_query_match = _extract_projected_attention_input(graph, rel_pos_node.inputs[0], role="q") + relative_key_match = _extract_relative_key_projection(graph, rel_pos_node.inputs[1]) + if rel_query_match is None or relative_key_match is None: + return None + if not ( + rel_query_match.hidden_value_id == query_match.hidden_value_id + and rel_query_match.weight_value_id == query_match.weight_value_id + and rel_query_match.projection_bias_value_id == query_match.projection_bias_value_id + and rel_query_match.projected_shape == query_match.projected_shape + ): + return None + rel_query_add_value_id = rel_query_match.add_value_id + relative_key_input_value_id = relative_key_match.input_value_id + relative_key_weight_value_id = relative_key_match.weight_value_id + relative_key_projection_bias_value_id = relative_key_match.projection_bias_value_id + relative_key_shape = relative_key_match.projected_shape + rel_pos_scale = float(rel_pos_node.attrs.get("scale", 1.0)) + rel_pos_node_ids = collect_node_ids( + rel_pos_node.id, + rel_pos_node.meta.get("rel_pos_bias_nodes", ()), + rel_query_match.node_ids, + relative_key_match.node_ids, + ) + mask_value_id = None + + return SelfAttentionBlockMatch( + attention_node_id=attention_node.id, + hidden_value_id=query_match.hidden_value_id, + query_weight_value_id=query_match.weight_value_id, + query_projection_bias_value_id=query_match.projection_bias_value_id, + query_add_value_id=query_match.add_value_id, + rel_query_add_value_id=rel_query_add_value_id, + key_weight_value_id=key_match.weight_value_id, + key_projection_bias_value_id=key_match.projection_bias_value_id, + value_weight_value_id=value_match.weight_value_id, + value_projection_bias_value_id=value_match.projection_bias_value_id, + mask_value_id=mask_value_id, + relative_key_input_value_id=relative_key_input_value_id, + relative_key_weight_value_id=relative_key_weight_value_id, + relative_key_projection_bias_value_id=relative_key_projection_bias_value_id, + gate_value_id=output_path["gate_value_id"], + output_projection_weight_value_id=projection.weight_value_id, + output_projection_bias_value_id=projection.bias_value_id, + query_shape=query_match.projected_shape, + key_shape=key_match.projected_shape, + value_shape=value_match.projected_shape, + relative_key_shape=relative_key_shape, + attention_output_shape=output_path["attention_output_shape"], + is_causal=bool(attention_node.attrs.get("is_causal", True)), + additive_mask=bool(attention_node.attrs.get("additive_mask", False)), + scale=float(attention_node.attrs.get("scale", 0.0)), + rel_pos_scale=rel_pos_scale, + window_size=int(attention_node.attrs.get("window_size", 0)), + node_ids=collect_node_ids( + attention_node.id, + projection.node_ids, + output_path["node_ids"], + query_match.node_ids, + key_match.node_ids, + value_match.node_ids, + rel_pos_node_ids, + ), + ) + + +def match_attention_block(graph: IRGraph, node: IRNode) -> AttentionBlockMatch | None: + projection = match_linear(graph, node) + if projection is None: + return None + + output_path = _extract_attention_output_path(graph, projection.input_value_id) + if output_path is None: + return None + attention_node = graph.nodes.get(output_path["attention_node_id"]) + if attention_node is None or attention_node.op not in {"attention", "scaled_dot_product_attention"} or len(attention_node.inputs) < 3: + return None + + return AttentionBlockMatch( + attention_node_id=attention_node.id, + query_value_id=attention_node.inputs[0], + key_value_id=attention_node.inputs[1], + value_value_id=attention_node.inputs[2], + mask_value_id=attention_node.inputs[3] if len(attention_node.inputs) > 3 else None, + gate_value_id=output_path["gate_value_id"], + output_projection_weight_value_id=projection.weight_value_id, + output_projection_bias_value_id=projection.bias_value_id, + attention_output_shape=output_path["attention_output_shape"], + is_causal=bool(attention_node.attrs.get("is_causal", True)), + additive_mask=bool(attention_node.attrs.get("additive_mask", False)), + scale=float(attention_node.attrs.get("scale", 0.0)), + window_size=int(attention_node.attrs.get("window_size", 0)), + node_ids=tuple(sorted({attention_node.id, *projection.node_ids, *output_path["node_ids"]})), + ) + + +def _extract_attention_output_path(graph: IRGraph, value_id: str) -> dict[str, object] | None: + current = value_id + gate_value_id: str | None = None + + while True: + node = producer(graph, current) + if node is None: + return None + if node.op in {"precision_cast", "contiguous"} and len(node.inputs) == 1: + current = node.inputs[0] + continue + if node.op == "type_as" and len(node.inputs) >= 1: + current = node.inputs[0] + continue + break + + node = producer(graph, current) + if node is not None and node.op == "multiply" and len(node.inputs) == 2: + for attn_candidate, gate_candidate in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + attn_match = _extract_attention_output_path(graph, attn_candidate) + if attn_match is None: + continue + attn_match["gate_value_id"] = gate_candidate + attn_match["node_ids"] = tuple(sorted({node.id, *attn_match["node_ids"]})) + return attn_match + + reshape_node = producer(graph, current) + if reshape_node is None or reshape_node.op not in {"reshape", "view"} or len(reshape_node.inputs) != 1: + return None + + attention_output_shape = tuple(int(v) for v in reshape_node.attrs.get("shape", ())) + current = reshape_node.inputs[0] + transpose_node = producer(graph, current) + if transpose_node is None: + return None + + if transpose_node.op == "transpose": + if int(transpose_node.attrs.get("dim0", -1)) != 1 or int(transpose_node.attrs.get("dim1", -1)) != 2: + return None + current = transpose_node.inputs[0] + elif transpose_node.op == "permute": + permutation = tuple(int(v) for v in transpose_node.attrs.get("permutation", ())) + if permutation != (0, 2, 1, 3): + return None + current = transpose_node.inputs[0] + else: + return None + + attention_node = producer(graph, current) + if attention_node is None or attention_node.op not in {"attention", "scaled_dot_product_attention"}: + return None + + return { + "attention_node_id": attention_node.id, + "gate_value_id": gate_value_id, + "attention_output_shape": attention_output_shape, + "node_ids": tuple(sorted({reshape_node.id, transpose_node.id})), + } + + +def _extract_projected_attention_input(graph: IRGraph, value_id: str, *, role: str) -> _ProjectedAttentionInputMatch | None: + current = strip_passthrough(graph, value_id) + add_node = producer(graph, current) + if add_node is not None and add_node.op == "add" and len(add_node.inputs) == 2: + for projected_candidate, add_candidate in ((add_node.inputs[0], add_node.inputs[1]), (add_node.inputs[1], add_node.inputs[0])): + projected = _extract_transposed_projection(graph, projected_candidate) + if projected is None: + continue + if not _looks_like_attention_addend(graph, add_candidate, projected.projected_shape): + continue + return _ProjectedAttentionInputMatch( + hidden_value_id=projected.hidden_value_id, + weight_value_id=projected.weight_value_id, + projection_bias_value_id=projected.projection_bias_value_id, + projected_shape=projected.projected_shape, + add_value_id=strip_passthrough(graph, add_candidate), + node_ids=collect_node_ids(add_node.id, projected.node_ids), + ) + + if role == "q": + return None + + return _extract_transposed_projection(graph, current) + + +def _extract_transposed_projection(graph: IRGraph, value_id: str) -> _ProjectedAttentionInputMatch | None: + current = strip_passthrough(graph, value_id) + layout_node = producer(graph, current) + if layout_node is None or len(layout_node.inputs) != 1: + return None + + if layout_node.op == "permute": + permutation = tuple(int(v) for v in layout_node.attrs.get("permutation", ())) + if permutation != (0, 2, 1, 3): + return None + elif layout_node.op == "transpose": + if int(layout_node.attrs.get("dim0", -1)) != 1 or int(layout_node.attrs.get("dim1", -1)) != 2: + return None + else: + return None + + projected_value_id = strip_passthrough(graph, layout_node.inputs[0]) + projected_value = graph.values.get(projected_value_id) + if projected_value is None or projected_value.shape is None or len(projected_value.shape) != 4: + return None + + reshape_node = producer(graph, projected_value_id) + if reshape_node is None or reshape_node.op not in {"reshape", "view"} or len(reshape_node.inputs) != 1: + return None + + linear_node = producer(graph, reshape_node.inputs[0]) + if linear_node is None: + return None + linear = match_linear(graph, linear_node) + if linear is None: + return None + + return _ProjectedAttentionInputMatch( + hidden_value_id=linear.input_value_id, + weight_value_id=linear.weight_value_id, + projection_bias_value_id=linear.bias_value_id, + projected_shape=tuple(int(v) for v in projected_value.shape), + add_value_id=None, + node_ids=collect_node_ids(layout_node.id, reshape_node.id, linear.node_ids), + ) + + +def _extract_relative_key_projection(graph: IRGraph, value_id: str) -> _RelativeKeyProjectionMatch | None: + current = strip_passthrough(graph, value_id) + projected_value = graph.values.get(current) + if projected_value is None or projected_value.shape is None or len(projected_value.shape) != 4: + return None + + reshape_node = producer(graph, current) + if reshape_node is None or reshape_node.op not in {"reshape", "view"} or len(reshape_node.inputs) != 1: + return None + + linear_node = producer(graph, reshape_node.inputs[0]) + if linear_node is None: + return None + linear = match_linear(graph, linear_node) + if linear is None: + return None + + return _RelativeKeyProjectionMatch( + input_value_id=linear.input_value_id, + weight_value_id=linear.weight_value_id, + projection_bias_value_id=linear.bias_value_id, + projected_shape=tuple(int(v) for v in projected_value.shape), + node_ids=collect_node_ids(reshape_node.id, linear.node_ids), + ) + + +def _looks_like_attention_addend(graph: IRGraph, value_id: str, target_shape: tuple[int, ...]) -> bool: + value = graph.values.get(strip_passthrough(graph, value_id)) + if value is None or value.shape is None: + return False + + shape = tuple(int(v) for v in value.shape) + if len(shape) != 4 or len(target_shape) != 4: + return False + + batch, seq_len, heads, head_dim = target_shape + if shape[0] not in (1, batch) or shape[3] != head_dim: + return False + + # Allow either BHSD-style tensors like [B, H, 1, D] or BSHD-style tensors like [B, 1, H, D]. + if shape[1] == heads and shape[2] in (1, seq_len): + return True + if shape[1] in (1, seq_len) and shape[2] == heads: + return True + return False + + +def _extract_attention_input(graph: IRGraph, value_id: str, *, role: str) -> dict[str, object] | None: + node_ids: set[str] = set() + current = strip_passthrough(graph, value_id) + has_gqa_repeat = False + + while True: + node = producer(graph, current) + if node is None: + return { + "source_input": current, + "weight_value_id": None, + "has_rope": False, + "has_rms_norm": False, + "has_gqa_repeat": has_gqa_repeat, + "node_ids": tuple(sorted(node_ids)), + } + node_ids.add(node.id) + + if node.op in {"reshape", "view", "transpose", "permute"}: + current = node.inputs[0] + continue + + rope = match_rope(graph, current) + if rope is not None: + rope_info = _extract_attention_input(graph, rope.input_value_id, role=role) + if rope_info is None: + return None + rope_info["has_rope"] = True + rope_info["node_ids"] = tuple(sorted(set(rope_info["node_ids"]) | node_ids | set(rope.node_ids))) + rope_info["has_gqa_repeat"] = bool(rope_info["has_gqa_repeat"] or has_gqa_repeat) + return rope_info + + rms = match_rms_norm(graph, node) + if rms is not None: + if role == "v": + return None + rms_info = _extract_attention_input(graph, rms.input_value_id, role=role) + if rms_info is None: + return None + rms_info["has_rms_norm"] = True + rms_info["node_ids"] = tuple(sorted(set(rms_info["node_ids"]) | node_ids | set(rms.node_ids))) + rms_info["has_gqa_repeat"] = bool(rms_info["has_gqa_repeat"] or has_gqa_repeat) + return rms_info + + if _looks_like_gqa_repeat(graph, current): + has_gqa_repeat = True + current = _unwrap_gqa_repeat(graph, current) + if current is None: + return None + continue + + linear_node = producer(graph, current) + if linear_node is None: + return { + "source_input": strip_passthrough(graph, current), + "weight_value_id": None, + "has_rope": False, + "has_rms_norm": False, + "has_gqa_repeat": has_gqa_repeat, + "node_ids": tuple(sorted(node_ids)), + } + linear = match_linear(graph, linear_node) + if linear is None: + return { + "source_input": strip_passthrough(graph, current), + "weight_value_id": None, + "has_rope": False, + "has_rms_norm": False, + "has_gqa_repeat": has_gqa_repeat, + "node_ids": tuple(sorted(node_ids)), + } + return { + "source_input": strip_passthrough(graph, linear.input_value_id), + "weight_value_id": linear.weight_value_id, + "has_rope": False, + "has_rms_norm": False, + "has_gqa_repeat": has_gqa_repeat, + "node_ids": tuple(sorted(set(node_ids) | set(linear.node_ids))), + } + + +def _looks_like_gqa_repeat(graph: IRGraph, value_id: str) -> bool: + return _unwrap_gqa_repeat(graph, value_id) is not None + + +def _unwrap_gqa_repeat(graph: IRGraph, value_id: str) -> str | None: + current = strip_passthrough(graph, value_id) + reshape_node = producer(graph, current) + if reshape_node is None or reshape_node.op not in {"reshape", "view"}: + return None + current = reshape_node.inputs[0] + expand_node = producer(graph, current) + if expand_node is None or expand_node.op != "expand": + return None + current = expand_node.inputs[0] + while True: + node = producer(graph, current) + if node is None: + return None + if node.op == "slice": + current = node.inputs[0] + continue + if node.op != "unsqueeze": + return None + if int(node.attrs.get("dim", -999)) != 2: + return None + break + base_value_id = strip_passthrough(graph, node.inputs[0]) + base_shape = graph.values.get(base_value_id).shape if graph.values.get(base_value_id) is not None else None + target_shape = tuple(int(v) for v in expand_node.attrs.get("shape", ())) + if base_shape is None or len(base_shape) != 4 or len(target_shape) != 5: + return None + expected = (int(base_shape[0]), int(base_shape[1]), int(target_shape[2]), int(base_shape[2]), int(base_shape[3])) + if expected != target_shape: + return None + return base_value_id diff --git a/python/cactus/transpile/fusion/common.py b/python/cactus/transpile/fusion/common.py new file mode 100644 index 000000000..d463cb804 --- /dev/null +++ b/python/cactus/transpile/fusion/common.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode + + +@dataclass(frozen=True) +class FusionMatch: + kind: str + anchor_node_id: str + node_ids: tuple[str, ...] + value_ids: tuple[str, ...] + attrs: dict[str, object] + + +def producer(graph: IRGraph, value_id: str) -> IRNode | None: + value = graph.values.get(value_id) + if value is None or value.producer is None: + return None + return graph.nodes.get(value.producer) + + +def strip_passthrough(graph: IRGraph, value_id: str) -> str: + current = value_id + while True: + node = producer(graph, current) + if node is None: + return current + if node.op in {"precision_cast", "contiguous"} and len(node.inputs) == 1: + current = node.inputs[0] + continue + if node.op == "type_as" and len(node.inputs) >= 1: + current = node.inputs[0] + continue + return current + + +def strip_layout_passthrough(graph: IRGraph, value_id: str) -> str: + current = value_id + while True: + node = producer(graph, current) + if node is None: + return current + if node.op in {"precision_cast", "contiguous", "reshape", "view", "unsqueeze", "expand"} and len(node.inputs) == 1: + current = node.inputs[0] + continue + if node.op == "type_as" and len(node.inputs) >= 1: + current = node.inputs[0] + continue + return current + + +def collect_node_ids(*parts: Any) -> tuple[str, ...]: + node_ids: set[str] = set() + for part in parts: + if part is None: + continue + if isinstance(part, str): + node_ids.add(part) + continue + if isinstance(part, IRNode): + node_ids.add(part.id) + continue + if isinstance(part, (list, tuple, set)): + for item in part: + if isinstance(item, str): + node_ids.add(item) + elif isinstance(item, IRNode): + node_ids.add(item.id) + return tuple(sorted(node_ids)) diff --git a/python/cactus/transpile/fusion/conv.py b/python/cactus/transpile/fusion/conv.py new file mode 100644 index 000000000..265501f9a --- /dev/null +++ b/python/cactus/transpile/fusion/conv.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.fusion.common import collect_node_ids +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode + + +@dataclass(frozen=True) +class ConvModuleMatch: + input_value_id: str + pointwise1_weight_value_id: str + pointwise1_bias_value_id: str | None + depthwise_weight_value_id: str + depthwise_bias_value_id: str | None + batch_norm_weight_value_id: str + batch_norm_bias_value_id: str + batch_norm_running_mean_value_id: str + batch_norm_running_var_value_id: str + pointwise2_weight_value_id: str + pointwise2_bias_value_id: str | None + eps: float + depthwise_kernel_size: int + depthwise_padding: int + node_ids: tuple[str, ...] + + +def match_conv_module(graph: IRGraph, node: IRNode) -> ConvModuleMatch | None: + if node.op != "permute" or len(node.inputs) != 1 or len(node.outputs) != 1: + return None + if tuple(node.attrs.get("permutation", ())) != (0, 2, 1): + return None + + pointwise2 = producer(graph, strip_passthrough(graph, node.inputs[0])) + if not _is_pointwise_conv1d(pointwise2): + return None + + silu = producer(graph, strip_passthrough(graph, pointwise2.inputs[0])) + if silu is None or silu.op != "silu" or len(silu.inputs) != 1: + return None + + batch_norm = producer(graph, strip_passthrough(graph, silu.inputs[0])) + if batch_norm is None or batch_norm.op != "batch_norm" or len(batch_norm.inputs) != 5: + return None + + depthwise = producer(graph, strip_passthrough(graph, batch_norm.inputs[0])) + if not _is_same_depthwise_conv1d(depthwise): + return None + + glu = producer(graph, strip_passthrough(graph, depthwise.inputs[0])) + if glu is None or glu.op != "glu" or len(glu.inputs) != 1: + return None + if int(glu.attrs.get("axis", -1)) not in {1, -2}: + return None + + pointwise1 = producer(graph, strip_passthrough(graph, glu.inputs[0])) + if not _is_pointwise_conv1d(pointwise1): + return None + + input_permute = producer(graph, strip_passthrough(graph, pointwise1.inputs[0])) + if input_permute is None or input_permute.op != "permute" or len(input_permute.inputs) != 1: + return None + if tuple(input_permute.attrs.get("permutation", ())) != (0, 2, 1): + return None + + depthwise_weight_shape = graph.values.get(depthwise.inputs[1]).shape if depthwise.inputs[1] in graph.values else None + if depthwise_weight_shape is None or len(depthwise_weight_shape) != 3: + return None + depthwise_kernel_size = int(depthwise_weight_shape[2]) + + return ConvModuleMatch( + input_value_id=strip_passthrough(graph, input_permute.inputs[0]), + pointwise1_weight_value_id=pointwise1.inputs[1], + pointwise1_bias_value_id=pointwise1.inputs[2] if len(pointwise1.inputs) > 2 else None, + depthwise_weight_value_id=depthwise.inputs[1], + depthwise_bias_value_id=depthwise.inputs[2] if len(depthwise.inputs) > 2 else None, + batch_norm_weight_value_id=batch_norm.inputs[1], + batch_norm_bias_value_id=batch_norm.inputs[2], + batch_norm_running_mean_value_id=batch_norm.inputs[3], + batch_norm_running_var_value_id=batch_norm.inputs[4], + pointwise2_weight_value_id=pointwise2.inputs[1], + pointwise2_bias_value_id=pointwise2.inputs[2] if len(pointwise2.inputs) > 2 else None, + eps=float(batch_norm.attrs.get("eps", 1e-5)), + depthwise_kernel_size=depthwise_kernel_size, + depthwise_padding=int(depthwise.attrs.get("padding", 0)), + node_ids=collect_node_ids( + node, + pointwise2, + silu, + batch_norm, + depthwise, + glu, + pointwise1, + input_permute, + ), + ) + + +def _is_pointwise_conv1d(node: IRNode | None) -> bool: + return ( + node is not None + and node.op == "conv1d" + and len(node.inputs) >= 2 + and int(node.attrs.get("stride", 1)) == 1 + and int(node.attrs.get("padding", 0)) == 0 + and int(node.attrs.get("dilation", 1)) == 1 + and int(node.attrs.get("groups", 1)) == 1 + ) + + +def _is_same_depthwise_conv1d(node: IRNode | None) -> bool: + if node is None or node.op != "conv1d" or len(node.inputs) < 2: + return False + groups = int(node.attrs.get("groups", 1)) + padding = int(node.attrs.get("padding", 0)) + stride = int(node.attrs.get("stride", 1)) + dilation = int(node.attrs.get("dilation", 1)) + if stride != 1 or dilation != 1: + return False + return groups > 1 and padding >= 0 diff --git a/python/cactus/transpile/fusion/deltanet.py b/python/cactus/transpile/fusion/deltanet.py new file mode 100644 index 000000000..94f4e9299 --- /dev/null +++ b/python/cactus/transpile/fusion/deltanet.py @@ -0,0 +1,357 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.fusion.linear import match_linear +from cactus.transpile.fusion.rms_norm import match_rms_norm +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode + + +@dataclass(frozen=True) +class GatedDeltaNetMatch: + output_value_id: str + normalized_input_value_id: str + qkv_weight_value_id: str + a_weight_value_id: str + b_weight_value_id: str + norm_weight_value_id: str + z_weight_value_id: str | None + dt_bias_value_id: str | None + a_log_value_id: str | None + conv_weight_value_id: str | None + num_k_heads: int + num_v_heads: int + key_dim: int + value_dim: int + eps: float + chunk_size: int + mode: str + node_ids: tuple[str, ...] + + +def match_gated_deltanet(graph: IRGraph, node: IRNode) -> GatedDeltaNetMatch | None: + if node.op not in {"view", "reshape"} or len(node.outputs) != 1: + return None + output_value = graph.values.get(node.outputs[0]) + if output_value is None or output_value.shape is None or len(output_value.shape) != 3: + return None + + body_value_id = strip_passthrough(graph, node.inputs[0]) + body_node = producer(graph, body_value_id) + if body_node is None: + return None + + z_weight_value_id: str | None = None + rms_match = None + normalized_input_value_id: str | None = None + node_ids: set[str] = {node.id} + + if body_node.op == "multiply" and len(body_node.inputs) == 2: + for rms_value_id, z_value_id in ((body_node.inputs[0], body_node.inputs[1]), (body_node.inputs[1], body_node.inputs[0])): + rms_match = _match_rms_from_value(graph, rms_value_id) + if rms_match is None: + continue + z_weight_value_id, normalized_input_value_id, z_node_ids = _extract_z_branch(graph, z_value_id) + if normalized_input_value_id is None: + continue + node_ids.update(z_node_ids) + node_ids.add(body_node.id) + break + if rms_match is None or normalized_input_value_id is None: + return None + else: + rms_match = _match_rms_from_value(graph, body_value_id) + if rms_match is None: + return None + + y2d_value_id = rms_match.input_value_id + if rms_match.weight_value_id is None: + return None + norm_weight_value_id = rms_match.weight_value_id + eps = float(rms_match.eps) + node_ids.update(rms_match.node_ids) + + if normalized_input_value_id is None: + normalized_input_value_id = _infer_normalized_input_from_y_branch(graph, y2d_value_id) + if normalized_input_value_id is None: + return None + + qkv_info = _find_qkv_linear(graph, normalized_input_value_id) + if qkv_info is None: + return None + qkv_weight_value_id, conv_weight_value_id, num_k_heads, key_dim, num_v_heads, value_dim, qkv_node_ids = qkv_info + node_ids.update(qkv_node_ids) + + a_info = _find_a_linear(graph, normalized_input_value_id) + b_info = _find_b_linear(graph, normalized_input_value_id) + if a_info is None or b_info is None: + return None + a_weight_value_id, dt_bias_value_id, a_log_value_id, a_node_ids = a_info + b_weight_value_id, b_node_ids = b_info + node_ids.update(a_node_ids) + node_ids.update(b_node_ids) + + mode, chunk_size = _classify_mode(graph, y2d_value_id, match_node_ids=node_ids) + + return GatedDeltaNetMatch( + output_value_id=node.outputs[0], + normalized_input_value_id=normalized_input_value_id, + qkv_weight_value_id=qkv_weight_value_id, + a_weight_value_id=a_weight_value_id, + b_weight_value_id=b_weight_value_id, + norm_weight_value_id=norm_weight_value_id, + z_weight_value_id=z_weight_value_id, + dt_bias_value_id=dt_bias_value_id, + a_log_value_id=a_log_value_id, + conv_weight_value_id=conv_weight_value_id, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + key_dim=key_dim, + value_dim=value_dim, + eps=eps, + chunk_size=chunk_size, + mode=mode, + node_ids=tuple(sorted(node_ids)), + ) + + +def _match_rms_from_value(graph: IRGraph, value_id: str): + node = producer(graph, strip_passthrough(graph, value_id)) + if node is None: + return None + return match_rms_norm(graph, node) + + +def _extract_z_branch(graph: IRGraph, value_id: str) -> tuple[str | None, str | None, tuple[str, ...]]: + current = strip_passthrough(graph, value_id) + silu_node = producer(graph, current) + if silu_node is None or silu_node.op != "silu" or len(silu_node.inputs) != 1: + return None, None, () + linear_value_id = strip_passthrough(graph, silu_node.inputs[0]) + linear_node = producer(graph, linear_value_id) + while linear_node is not None and linear_node.op in {"view", "reshape"} and len(linear_node.inputs) == 1: + linear_value_id = strip_passthrough(graph, linear_node.inputs[0]) + linear_node = producer(graph, linear_value_id) + if linear_node is None: + return None, None, () + linear_match = match_linear(graph, linear_node) + if linear_match is None: + return None, None, () + return ( + linear_match.weight_value_id, + linear_match.input_value_id, + tuple(sorted({silu_node.id, *linear_match.node_ids})), + ) + + +def _infer_normalized_input_from_y_branch(graph: IRGraph, value_id: str) -> str | None: + current = strip_passthrough(graph, value_id) + node = producer(graph, current) + visited: set[str] = set() + while node is not None and node.id not in visited: + visited.add(node.id) + for input_value_id in node.inputs: + candidate = producer(graph, strip_passthrough(graph, input_value_id)) + if candidate is None: + continue + linear_match = match_linear(graph, candidate) + if linear_match is not None: + return linear_match.input_value_id + if len(node.inputs) != 1: + return None + node = producer(graph, strip_passthrough(graph, node.inputs[0])) + return None + + +def _find_qkv_linear( + graph: IRGraph, + normalized_input_value_id: str, +) -> tuple[str, str | None, int, int, int, int, tuple[str, ...]] | None: + for node_id in graph.order: + node = graph.nodes[node_id] + linear_match = match_linear(graph, node) + if linear_match is None or linear_match.input_value_id != normalized_input_value_id: + continue + split_node = _find_descendant_op(graph, linear_match.output_value_id, {"split_with_sizes"}, max_depth=8) + if split_node is None: + continue + split_sizes = tuple(int(v) for v in split_node.attrs.get("sizes", ())) + if len(split_sizes) != 3: + continue + + conv_node = _find_descendant_op(graph, linear_match.output_value_id, {"conv1d"}, max_depth=6) + conv_weight_value_id = conv_node.inputs[1] if conv_node is not None and len(conv_node.inputs) >= 2 else None + + q_shape = None + v_shape = None + node_ids = {node.id, split_node.id} + if conv_node is not None: + node_ids.add(conv_node.id) + + split_value = graph.values.get(split_node.outputs[0]) + if split_value is None: + continue + for getitem_node_id in split_value.users: + getitem_node = graph.nodes.get(getitem_node_id) + if getitem_node is None or getitem_node.op != "getitem" or not getitem_node.outputs: + continue + shape = _find_rank4_shape(graph, getitem_node.outputs[0], max_depth=5) + if shape is None: + continue + branch_has_sum = _find_descendant_op(graph, getitem_node.outputs[0], {"sum"}, max_depth=8) is not None + node_ids.add(getitem_node.id) + if branch_has_sum and q_shape is None: + q_shape = shape + elif not branch_has_sum and v_shape is None: + v_shape = shape + + if q_shape is None or v_shape is None: + continue + + num_k_heads = int(q_shape[2]) + key_dim = int(q_shape[3]) + num_v_heads = int(v_shape[2]) + value_dim = int(v_shape[3]) + if num_k_heads <= 0 or key_dim <= 0 or num_v_heads <= 0 or value_dim <= 0: + continue + + return ( + linear_match.weight_value_id, + conv_weight_value_id, + num_k_heads, + key_dim, + num_v_heads, + value_dim, + tuple(sorted(node_ids | set(linear_match.node_ids))), + ) + return None + + +def _find_a_linear( + graph: IRGraph, + normalized_input_value_id: str, +) -> tuple[str, str | None, str | None, tuple[str, ...]] | None: + for node_id in graph.order: + node = graph.nodes[node_id] + linear_match = match_linear(graph, node) + if linear_match is None or linear_match.input_value_id != normalized_input_value_id: + continue + softplus_node = _find_descendant_op(graph, linear_match.output_value_id, {"softplus"}, max_depth=6) + if softplus_node is None: + continue + dt_bias_value_id = None + softplus_input_node = producer(graph, strip_passthrough(graph, softplus_node.inputs[0])) + if softplus_input_node is not None and softplus_input_node.op == "add": + for input_value_id in softplus_input_node.inputs: + value = graph.values.get(input_value_id) + if value is not None and value.producer is None and input_value_id in graph.constants: + dt_bias_value_id = input_value_id + break + a_log_value_id = None + multiply_node = _find_descendant_op(graph, softplus_node.outputs[0], {"multiply"}, max_depth=6) + if multiply_node is not None: + for input_value_id in multiply_node.inputs: + a_log_value_id = _extract_exp_constant_input(graph, input_value_id) + if a_log_value_id is not None: + break + node_ids = {softplus_node.id, *linear_match.node_ids} + if softplus_input_node is not None: + node_ids.add(softplus_input_node.id) + if multiply_node is not None: + node_ids.add(multiply_node.id) + return linear_match.weight_value_id, dt_bias_value_id, a_log_value_id, tuple(sorted(node_ids)) + return None + + +def _find_b_linear(graph: IRGraph, normalized_input_value_id: str) -> tuple[str, tuple[str, ...]] | None: + for node_id in graph.order: + node = graph.nodes[node_id] + linear_match = match_linear(graph, node) + if linear_match is None or linear_match.input_value_id != normalized_input_value_id: + continue + sigmoid_node = _find_descendant_op(graph, linear_match.output_value_id, {"sigmoid"}, max_depth=4) + if sigmoid_node is None: + continue + return linear_match.weight_value_id, tuple(sorted({sigmoid_node.id, *linear_match.node_ids})) + return None + + +def _extract_exp_constant_input(graph: IRGraph, value_id: str) -> str | None: + node = producer(graph, strip_passthrough(graph, value_id)) + if node is None: + value = graph.values.get(value_id) + if value is not None and value.producer is None and value_id in graph.constants: + return value_id + return None + if node.op == "scalar_multiply": + return _extract_exp_constant_input(graph, node.inputs[0]) + if node.op == "scalar_exp": + source = strip_passthrough(graph, node.inputs[0]) + value = graph.values.get(source) + if value is not None and value.producer is None and source in graph.constants: + return source + return None + + +def _find_rank4_shape(graph: IRGraph, value_id: str, *, max_depth: int) -> tuple[int, int, int, int] | None: + frontier = [(value_id, 0)] + seen: set[str] = set() + while frontier: + current, depth = frontier.pop(0) + if current in seen or depth > max_depth: + continue + seen.add(current) + value = graph.values.get(current) + if value is not None and value.shape is not None and len(value.shape) == 4: + shape = tuple(int(dim) for dim in value.shape) + return shape # type: ignore[return-value] + for user_id in graph.values.get(current, None).users if graph.values.get(current, None) is not None else (): + user = graph.nodes.get(user_id) + if user is None or not user.outputs: + continue + frontier.extend((output_id, depth + 1) for output_id in user.outputs) + return None + + +def _find_descendant_op(graph: IRGraph, value_id: str, targets: set[str], *, max_depth: int) -> IRNode | None: + frontier = [(value_id, 0)] + seen: set[str] = set() + while frontier: + current, depth = frontier.pop(0) + if current in seen or depth > max_depth: + continue + seen.add(current) + value = graph.values.get(current) + if value is None: + continue + for user_id in value.users: + user = graph.nodes.get(user_id) + if user is None or not user.outputs: + continue + if user.op in targets: + return user + frontier.extend((output_id, depth + 1) for output_id in user.outputs) + return None + + +def _classify_mode(graph: IRGraph, y2d_value_id: str, *, match_node_ids: set[str]) -> tuple[str, int]: + del y2d_value_id + chunk_size = _infer_prefill_chunk_size(graph, match_node_ids) + # Current exported graphs are uncached/full-sequence captures, so fuse the + # linear-attention block to the prefill kernel by default. Decode should be + # emitted only once the capture path models cache/state explicitly. + return "prefill", chunk_size + + +def _infer_prefill_chunk_size(graph: IRGraph, match_node_ids: set[str]) -> int: + for node_id in match_node_ids: + node = graph.nodes.get(node_id) + if node is None or node.op != "aten.triu.default" or not node.inputs: + continue + ones_value = graph.values.get(node.inputs[0]) + if ones_value is not None and ones_value.shape and len(ones_value.shape) >= 1: + return int(ones_value.shape[-1]) + return 64 diff --git a/python/cactus/transpile/fusion/linear.py b/python/cactus/transpile/fusion/linear.py new file mode 100644 index 000000000..4841019f3 --- /dev/null +++ b/python/cactus/transpile/fusion/linear.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough + + +@dataclass(frozen=True) +class LinearMatch: + input_value_id: str + weight_value_id: str + bias_value_id: str | None + output_value_id: str + node_ids: tuple[str, ...] + kind: str + + +def match_linear(graph: IRGraph, node: IRNode) -> LinearMatch | None: + if node.op == "linear" and len(node.inputs) >= 2 and len(node.outputs) == 1: + return LinearMatch( + input_value_id=strip_passthrough(graph, node.inputs[0]), + weight_value_id=node.inputs[1], + bias_value_id=node.inputs[2] if bool(node.attrs.get("has_bias", False)) and len(node.inputs) > 2 else None, + output_value_id=node.outputs[0], + node_ids=(node.id,), + kind="linear", + ) + + if node.op == "addmm" and len(node.inputs) == 3 and len(node.outputs) == 1: + return LinearMatch( + input_value_id=strip_passthrough(graph, node.inputs[1]), + weight_value_id=node.inputs[2], + bias_value_id=node.inputs[0], + output_value_id=node.outputs[0], + node_ids=(node.id,), + kind="addmm", + ) + + if node.op != "add" or len(node.inputs) != 2 or len(node.outputs) != 1: + return None + + for matmul_value_id, bias_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + matmul_node = producer(graph, strip_passthrough(graph, matmul_value_id)) + if matmul_node is None or matmul_node.op != "matmul" or len(matmul_node.inputs) != 2: + continue + return LinearMatch( + input_value_id=strip_passthrough(graph, matmul_node.inputs[0]), + weight_value_id=matmul_node.inputs[1], + bias_value_id=bias_value_id, + output_value_id=node.outputs[0], + node_ids=tuple(sorted({node.id, matmul_node.id})), + kind="matmul_add", + ) + + return None diff --git a/python/cactus/transpile/fusion/lstm.py b/python/cactus/transpile/fusion/lstm.py new file mode 100644 index 000000000..64831240e --- /dev/null +++ b/python/cactus/transpile/fusion/lstm.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.fusion.common import collect_node_ids +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode + + +@dataclass(frozen=True) +class LSTMCellMatch: + anchor_node_id: str + node_ids: tuple[str, ...] + x_value_id: str + h_prev_value_id: str + c_prev_value_id: str + weight_ih_value_id: str + weight_hh_value_id: str + bias_ih_value_id: str + bias_hh_value_id: str | None + h_output_value_id: str + c_output_value_id: str + + +def match_lstm_cell(graph: IRGraph, node: IRNode) -> LSTMCellMatch | None: + if node.op != "multiply" or len(node.inputs) != 2 or len(node.outputs) != 1: + return None + + for output_gate_value_id, c_tanh_value_id in ( + (node.inputs[0], node.inputs[1]), + (node.inputs[1], node.inputs[0]), + ): + output_gate = _extract_gate_activation(graph, output_gate_value_id, expected_activation="sigmoid") + if output_gate is None or output_gate["index"] != 3: + continue + + c_tanh_node = producer(graph, strip_passthrough(graph, c_tanh_value_id)) + if c_tanh_node is None or c_tanh_node.op != "tanh" or len(c_tanh_node.inputs) != 1: + continue + + c_next_value_id = strip_passthrough(graph, c_tanh_node.inputs[0]) + c_next_node = producer(graph, c_next_value_id) + if c_next_node is None or c_next_node.op != "add" or len(c_next_node.inputs) != 2 or len(c_next_node.outputs) != 1: + continue + + for forget_term_value_id, input_term_value_id in ( + (c_next_node.inputs[0], c_next_node.inputs[1]), + (c_next_node.inputs[1], c_next_node.inputs[0]), + ): + forget_term = _extract_forget_term(graph, forget_term_value_id) + input_term = _extract_input_term(graph, input_term_value_id) + if forget_term is None or input_term is None: + continue + + input_gate = input_term["input_gate"] + cell_gate = input_term["cell_gate"] + forget_gate = forget_term["forget_gate"] + chunk_node_id = output_gate["chunk_node_id"] + if ( + input_gate["chunk_node_id"] != chunk_node_id + or forget_gate["chunk_node_id"] != chunk_node_id + or cell_gate["chunk_node_id"] != chunk_node_id + ): + continue + if input_gate["index"] != 0 or forget_gate["index"] != 1 or cell_gate["index"] != 2: + continue + + chunk_node = graph.nodes.get(chunk_node_id) + if chunk_node is None or len(chunk_node.inputs) != 1: + continue + + gates = _extract_lstm_gate_sum(graph, chunk_node.inputs[0]) + if gates is None: + continue + + return LSTMCellMatch( + anchor_node_id=node.id, + node_ids=collect_node_ids( + node, + c_tanh_node, + c_next_node, + output_gate["node_ids"], + forget_gate["node_ids"], + input_gate["node_ids"], + cell_gate["node_ids"], + forget_term["node_ids"], + input_term["node_ids"], + chunk_node, + gates["node_ids"], + ), + x_value_id=gates["x_value_id"], + h_prev_value_id=gates["h_prev_value_id"], + c_prev_value_id=forget_term["c_prev_value_id"], + weight_ih_value_id=gates["weight_ih_value_id"], + weight_hh_value_id=gates["weight_hh_value_id"], + bias_ih_value_id=gates["bias_value_ids"][0], + bias_hh_value_id=gates["bias_value_ids"][1] if len(gates["bias_value_ids"]) > 1 else None, + h_output_value_id=node.outputs[0], + c_output_value_id=c_next_node.outputs[0], + ) + + return None + + +def _extract_gate_activation( + graph: IRGraph, + value_id: str, + *, + expected_activation: str, +) -> dict[str, object] | None: + activation_node = producer(graph, strip_passthrough(graph, value_id)) + if activation_node is None or activation_node.op != expected_activation or len(activation_node.inputs) != 1: + return None + + getitem_node = producer(graph, strip_passthrough(graph, activation_node.inputs[0])) + if getitem_node is None or getitem_node.op != "getitem" or len(getitem_node.inputs) != 1: + return None + + chunk_node = producer(graph, strip_passthrough(graph, getitem_node.inputs[0])) + if chunk_node is None or chunk_node.op != "chunk": + return None + if int(chunk_node.attrs.get("chunks", 0)) != 4: + return None + + return { + "chunk_node_id": chunk_node.id, + "index": int(getitem_node.attrs.get("index", -1)), + "preactivation_value_id": getitem_node.inputs[0], + "node_ids": (activation_node.id, getitem_node.id), + } + + +def _extract_forget_term(graph: IRGraph, value_id: str) -> dict[str, object] | None: + mul_node = producer(graph, strip_passthrough(graph, value_id)) + if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: + return None + + for gate_value_id, c_prev_value_id in ( + (mul_node.inputs[0], mul_node.inputs[1]), + (mul_node.inputs[1], mul_node.inputs[0]), + ): + forget_gate = _extract_gate_activation(graph, gate_value_id, expected_activation="sigmoid") + if forget_gate is None: + continue + return { + "forget_gate": forget_gate, + "c_prev_value_id": strip_passthrough(graph, c_prev_value_id), + "node_ids": (mul_node.id,), + } + return None + + +def _extract_input_term(graph: IRGraph, value_id: str) -> dict[str, object] | None: + mul_node = producer(graph, strip_passthrough(graph, value_id)) + if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: + return None + + for input_gate_value_id, cell_gate_value_id in ( + (mul_node.inputs[0], mul_node.inputs[1]), + (mul_node.inputs[1], mul_node.inputs[0]), + ): + input_gate = _extract_gate_activation(graph, input_gate_value_id, expected_activation="sigmoid") + if input_gate is None: + continue + cell_tanh_node = producer(graph, strip_passthrough(graph, cell_gate_value_id)) + if cell_tanh_node is None or cell_tanh_node.op != "tanh" or len(cell_tanh_node.inputs) != 1: + continue + cell_gate = _extract_gate_activation(graph, cell_gate_value_id, expected_activation="tanh") + if cell_gate is None: + continue + return { + "input_gate": input_gate, + "cell_gate": cell_gate, + "node_ids": (mul_node.id, cell_tanh_node.id), + } + return None + + +def _extract_lstm_gate_sum(graph: IRGraph, value_id: str) -> dict[str, object] | None: + terms, node_ids = _flatten_add_terms(graph, strip_passthrough(graph, value_id)) + linear_nodes: list[IRNode] = [] + bias_value_ids: list[str] = [] + for term_value_id in terms: + term_value_id = strip_passthrough(graph, term_value_id) + term_node = producer(graph, term_value_id) + if term_node is not None and term_node.op == "linear" and len(term_node.inputs) >= 2: + linear_nodes.append(term_node) + continue + if term_value_id in graph.constants: + bias_value_ids.append(term_value_id) + continue + return None + + if len(linear_nodes) != 2 or len(bias_value_ids) not in {1, 2}: + return None + + first_linear, second_linear = linear_nodes + return { + "x_value_id": first_linear.inputs[0], + "weight_ih_value_id": first_linear.inputs[1], + "h_prev_value_id": second_linear.inputs[0], + "weight_hh_value_id": second_linear.inputs[1], + "bias_value_ids": tuple(bias_value_ids), + "node_ids": (*node_ids, first_linear.id, second_linear.id), + } + + +def _flatten_add_terms(graph: IRGraph, value_id: str) -> tuple[list[str], tuple[str, ...]]: + node = producer(graph, strip_passthrough(graph, value_id)) + if node is None or node.op != "add" or len(node.inputs) != 2: + return [value_id], () + left_terms, left_nodes = _flatten_add_terms(graph, node.inputs[0]) + right_terms, right_nodes = _flatten_add_terms(graph, node.inputs[1]) + return [*left_terms, *right_terms], (*left_nodes, *right_nodes, node.id) diff --git a/python/cactus/transpile/fusion/mlp.py b/python/cactus/transpile/fusion/mlp.py new file mode 100644 index 000000000..56a219514 --- /dev/null +++ b/python/cactus/transpile/fusion/mlp.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.fusion.linear import match_linear + + +@dataclass(frozen=True) +class GatedMLPMatch: + input_value_id: str + gate_weight_value_id: str + up_weight_value_id: str + down_weight_value_id: str + activation: str + bias_value_ids: tuple[str | None, str | None, str | None] + node_ids: tuple[str, ...] + + +def match_gated_mlp(graph: IRGraph, node: IRNode) -> GatedMLPMatch | None: + down = match_linear(graph, node) + if down is None: + return None + + mul_node = graph.nodes.get(graph.values[down.input_value_id].producer) if graph.values.get(down.input_value_id) is not None else None + if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: + return None + + for activated_value_id, up_value_id in ((mul_node.inputs[0], mul_node.inputs[1]), (mul_node.inputs[1], mul_node.inputs[0])): + activated_node = graph.nodes.get(graph.values[strip_passthrough(graph, activated_value_id)].producer) if graph.values.get(strip_passthrough(graph, activated_value_id)) is not None else None + if activated_node is None or activated_node.op not in {"gelu", "silu"}: + continue + + gate = match_linear(graph, graph.nodes[graph.values[strip_passthrough(graph, activated_node.inputs[0])].producer]) if graph.values.get(strip_passthrough(graph, activated_node.inputs[0])) is not None and graph.values[strip_passthrough(graph, activated_node.inputs[0])].producer in graph.nodes else None + up = match_linear(graph, graph.nodes[graph.values[strip_passthrough(graph, up_value_id)].producer]) if graph.values.get(strip_passthrough(graph, up_value_id)) is not None and graph.values[strip_passthrough(graph, up_value_id)].producer in graph.nodes else None + if gate is None or up is None: + continue + if strip_passthrough(graph, gate.input_value_id) != strip_passthrough(graph, up.input_value_id): + continue + + return GatedMLPMatch( + input_value_id=strip_passthrough(graph, gate.input_value_id), + gate_weight_value_id=gate.weight_value_id, + up_weight_value_id=up.weight_value_id, + down_weight_value_id=down.weight_value_id, + activation=activated_node.op, + bias_value_ids=(gate.bias_value_id, up.bias_value_id, down.bias_value_id), + node_ids=tuple(sorted({*gate.node_ids, activated_node.id, *up.node_ids, mul_node.id, *down.node_ids})), + ) + + return None diff --git a/python/cactus/transpile/fusion/rel_pos_bias.py b/python/cactus/transpile/fusion/rel_pos_bias.py new file mode 100644 index 000000000..a0e709c9c --- /dev/null +++ b/python/cactus/transpile/fusion/rel_pos_bias.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode + + +@dataclass(frozen=True) +class RelPosBiasMatch: + query_value_id: str + relative_key_value_id: str + scale: float + node_ids: tuple[str, ...] + + +def match_rel_pos_bias(graph: IRGraph, node: IRNode) -> RelPosBiasMatch | None: + if node.op != "scalar_multiply" or len(node.inputs) != 1: + return None + + gather_match = _match_gather_rel_pos_bias(graph, node) + if gather_match is not None: + return gather_match + + final_slice = producer(graph, strip_passthrough(graph, node.inputs[0])) + if final_slice is None or final_slice.op != "slice": + return None + if int(final_slice.attrs.get("axis", -1)) != 3 or int(final_slice.attrs.get("start", -1)) != 0: + return None + + reshape_after_shift = producer(graph, strip_passthrough(graph, final_slice.inputs[0])) + if reshape_after_shift is None or reshape_after_shift.op not in {"reshape", "view"} or len(reshape_after_shift.inputs) != 1: + return None + + shift_slice = producer(graph, strip_passthrough(graph, reshape_after_shift.inputs[0])) + if shift_slice is None or shift_slice.op != "slice": + return None + if int(shift_slice.attrs.get("axis", -1)) != 2 or int(shift_slice.attrs.get("start", -1)) != 1: + return None + + reshape_for_shift = producer(graph, strip_passthrough(graph, shift_slice.inputs[0])) + if reshape_for_shift is None or reshape_for_shift.op not in {"reshape", "view"} or len(reshape_for_shift.inputs) != 1: + return None + + pad = producer(graph, strip_passthrough(graph, reshape_for_shift.inputs[0])) + if pad is None or pad.op != "pad" or len(pad.inputs) != 1: + return None + if tuple(int(v) for v in pad.attrs.get("pads", ())) != (1, 0): + return None + if str(pad.attrs.get("mode", "constant")) != "constant": + return None + if float(pad.attrs.get("value", 0.0)) != 0.0: + return None + + matmul = producer(graph, strip_passthrough(graph, pad.inputs[0])) + if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: + return None + + rhs_permute = producer(graph, strip_passthrough(graph, matmul.inputs[1])) + if rhs_permute is None or rhs_permute.op != "permute" or len(rhs_permute.inputs) != 1: + return None + if tuple(int(v) for v in rhs_permute.attrs.get("permutation", ())) != (0, 2, 3, 1): + return None + + query_value_id = strip_passthrough(graph, matmul.inputs[0]) + relative_key_value_id = strip_passthrough(graph, rhs_permute.inputs[0]) + + query_shape = graph.values.get(query_value_id).shape if query_value_id in graph.values else None + relative_key_shape = graph.values.get(relative_key_value_id).shape if relative_key_value_id in graph.values else None + output_shape = graph.values.get(node.outputs[0]).shape if node.outputs and node.outputs[0] in graph.values else None + if query_shape is None or relative_key_shape is None or output_shape is None: + return None + if len(query_shape) != 4 or len(relative_key_shape) != 4 or len(output_shape) != 4: + return None + + batch, heads, seq_len, head_dim = (int(v) for v in query_shape) + rel_batch, rel_len, rel_heads, rel_dim = (int(v) for v in relative_key_shape) + out_batch, out_heads, out_t, out_s = (int(v) for v in output_shape) + if rel_batch not in {1, batch}: + return None + if rel_heads != heads or rel_dim != head_dim: + return None + if rel_len < 2 * seq_len - 1: + return None + if (out_batch, out_heads, out_t, out_s) != (batch, heads, seq_len, seq_len): + return None + + return RelPosBiasMatch( + query_value_id=query_value_id, + relative_key_value_id=relative_key_value_id, + scale=float(node.attrs.get("value", 1.0)), + node_ids=tuple( + sorted( + { + node.id, + final_slice.id, + reshape_after_shift.id, + shift_slice.id, + reshape_for_shift.id, + pad.id, + matmul.id, + rhs_permute.id, + } + ) + ), + ) + + +def _match_gather_rel_pos_bias(graph: IRGraph, node: IRNode) -> RelPosBiasMatch | None: + gather = producer(graph, strip_passthrough(graph, node.inputs[0])) + if gather is None or gather.op != "gather" or len(gather.inputs) < 2: + return None + if int(gather.attrs.get("axis", -1)) not in {-1, 3}: + return None + + matmul = producer(graph, strip_passthrough(graph, gather.inputs[0])) + if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: + return None + + query_heads = producer(graph, strip_passthrough(graph, matmul.inputs[0])) + rhs_transpose = producer(graph, strip_passthrough(graph, matmul.inputs[1])) + if query_heads is None or query_heads.op != "permute" or len(query_heads.inputs) != 1: + return None + if rhs_transpose is None or rhs_transpose.op != "permute" or len(rhs_transpose.inputs) != 1: + return None + if tuple(int(v) for v in query_heads.attrs.get("permutation", ())) != (0, 2, 1, 3): + return None + if tuple(int(v) for v in rhs_transpose.attrs.get("permutation", ())) != (0, 1, 3, 2): + return None + + relative_key_heads = producer(graph, strip_passthrough(graph, rhs_transpose.inputs[0])) + if relative_key_heads is None or relative_key_heads.op != "permute" or len(relative_key_heads.inputs) != 1: + return None + if tuple(int(v) for v in relative_key_heads.attrs.get("permutation", ())) != (0, 2, 1, 3): + return None + + query_value_id = strip_passthrough(graph, query_heads.inputs[0]) + relative_key_value_id = strip_passthrough(graph, relative_key_heads.inputs[0]) + output_value_id = node.outputs[0] if node.outputs else "" + + query_shape = graph.values.get(query_value_id).shape if query_value_id in graph.values else None + relative_key_shape = graph.values.get(relative_key_value_id).shape if relative_key_value_id in graph.values else None + output_shape = graph.values.get(output_value_id).shape if output_value_id in graph.values else None + if query_shape is None or relative_key_shape is None or output_shape is None: + return None + if len(query_shape) != 4 or len(relative_key_shape) != 4 or len(output_shape) != 4: + return None + + batch, seq_len, heads, head_dim = (int(v) for v in query_shape) + rel_batch, rel_len, rel_heads, rel_dim = (int(v) for v in relative_key_shape) + out_batch, out_heads, out_t, out_s = (int(v) for v in output_shape) + if rel_batch not in {1, batch}: + return None + if rel_heads != heads or rel_dim != head_dim: + return None + if rel_len < 2 * seq_len - 1: + return None + if (out_batch, out_heads, out_t, out_s) != (batch, heads, seq_len, seq_len): + return None + + return RelPosBiasMatch( + query_value_id=query_value_id, + relative_key_value_id=relative_key_value_id, + scale=float(node.attrs.get("value", 1.0)), + node_ids=tuple( + sorted( + { + node.id, + gather.id, + matmul.id, + query_heads.id, + rhs_transpose.id, + relative_key_heads.id, + } + ) + ), + ) diff --git a/python/cactus/transpile/fusion/rms_norm.py b/python/cactus/transpile/fusion/rms_norm.py new file mode 100644 index 000000000..322833744 --- /dev/null +++ b/python/cactus/transpile/fusion/rms_norm.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough + + +@dataclass(frozen=True) +class RMSNormMatch: + input_value_id: str + weight_value_id: str | None + weight_offset: float + eps: float + node_ids: tuple[str, ...] + + +def match_rms_norm(graph: IRGraph, node: IRNode) -> RMSNormMatch | None: + anchor_value_id = node.inputs[0] if node.op == "type_as" and node.inputs else node.outputs[0] + anchor_value_id = strip_passthrough(graph, anchor_value_id) + anchor_node = producer(graph, anchor_value_id) + if anchor_node is None or anchor_node.op != "multiply" or len(anchor_node.inputs) != 2: + return None + + for normed_value_id, scale_value_id in ( + (anchor_node.inputs[0], anchor_node.inputs[1]), + (anchor_node.inputs[1], anchor_node.inputs[0]), + ): + scale = _extract_scale(graph, scale_value_id) + norm = _extract_normed_value(graph, normed_value_id) + if scale is None or norm is None: + continue + if strip_passthrough(graph, norm["input_value_id"]) != strip_passthrough(graph, norm["pow_input_value_id"]): + continue + return RMSNormMatch( + input_value_id=norm["input_value_id"], + weight_value_id=scale["weight_value_id"], + weight_offset=float(scale["offset"]), + eps=float(norm["eps"]), + node_ids=tuple(sorted({anchor_node.id, *norm["node_ids"], *scale["node_ids"]})), + ) + + norm = _extract_normed_value(graph, anchor_node.outputs[0]) + if norm is None: + return None + if strip_passthrough(graph, norm["input_value_id"]) != strip_passthrough(graph, norm["pow_input_value_id"]): + return None + return RMSNormMatch( + input_value_id=norm["input_value_id"], + weight_value_id=None, + weight_offset=0.0, + eps=float(norm["eps"]), + node_ids=tuple(sorted({anchor_node.id, *norm["node_ids"]})), + ) + + +def _extract_scale(graph: IRGraph, value_id: str) -> dict[str, object] | None: + current = value_id + offset = 0.0 + node_ids: set[str] = set() + producer_node = producer(graph, current) + if producer_node is not None and producer_node.op == "scalar_add": + add_value = float(producer_node.attrs.get("value", 0.0)) + if add_value in (0.0, 1.0): + offset = add_value + node_ids.add(producer_node.id) + current = producer_node.inputs[0] + current = strip_passthrough(graph, current) + if current not in graph.constants or not isinstance(graph.constants[current], torch.Tensor): + return None + return {"weight_value_id": current, "offset": offset, "node_ids": tuple(sorted(node_ids))} + + +def _extract_normed_value(graph: IRGraph, value_id: str) -> dict[str, object] | None: + mul = producer(graph, strip_passthrough(graph, value_id)) + if mul is None or mul.op != "multiply" or len(mul.inputs) != 2: + return None + for x_candidate, rsqrt_candidate in ((mul.inputs[0], mul.inputs[1]), (mul.inputs[1], mul.inputs[0])): + rsqrt = _extract_rsqrt_chain(graph, rsqrt_candidate) + if rsqrt is None: + continue + if strip_passthrough(graph, x_candidate) != strip_passthrough(graph, rsqrt["pow_input_value_id"]): + continue + return { + "input_value_id": strip_passthrough(graph, x_candidate), + "pow_input_value_id": rsqrt["pow_input_value_id"], + "eps": rsqrt["eps"], + "node_ids": tuple(sorted({mul.id, *rsqrt["node_ids"]})), + } + return None + + +def _extract_rsqrt_chain(graph: IRGraph, value_id: str) -> dict[str, object] | None: + pow_node = producer(graph, strip_passthrough(graph, value_id)) + if pow_node is None or pow_node.op != "pow" or float(pow_node.attrs.get("exponent", 0.0)) != -0.5: + return None + add_node = producer(graph, pow_node.inputs[0]) + if add_node is None or add_node.op != "scalar_add": + return None + mean_node = producer(graph, add_node.inputs[0]) + if mean_node is None or mean_node.op != "mean" or not bool(mean_node.attrs.get("keepdim", False)): + return None + pow2_node = producer(graph, mean_node.inputs[0]) + if pow2_node is None or pow2_node.op != "pow" or float(pow2_node.attrs.get("exponent", 0.0)) != 2.0: + return None + if not _reduces_last_dim(graph, mean_node, pow2_node.inputs[0]): + return None + return { + "pow_input_value_id": pow2_node.inputs[0], + "eps": float(add_node.attrs.get("value", 0.0)), + "node_ids": (pow_node.id, add_node.id, mean_node.id, pow2_node.id), + } + + +def _reduces_last_dim(graph: IRGraph, mean_node: IRNode, input_value_id: str) -> bool: + value = graph.values.get(input_value_id) + if value is None or value.shape is None: + return False + rank = len(value.shape) + if rank == 0: + return False + axis = mean_node.attrs.get("axis") + if isinstance(axis, (list, tuple)): + if len(axis) != 1: + return False + axis = axis[0] + if not isinstance(axis, int): + return False + return axis % rank == rank - 1 diff --git a/python/cactus/transpile/fusion/rope.py b/python/cactus/transpile/fusion/rope.py new file mode 100644 index 000000000..25118b868 --- /dev/null +++ b/python/cactus/transpile/fusion/rope.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_layout_passthrough +from cactus.transpile.fusion.common import strip_passthrough + + +@dataclass(frozen=True) +class RoPEMatch: + input_value_id: str + theta: float + position_offset: int + partial: bool + node_ids: tuple[str, ...] + + +def match_rope(graph: IRGraph, value_id: str) -> RoPEMatch | None: + desc = _extract_rope_descriptor(graph, value_id) + if desc is None: + return None + return RoPEMatch( + input_value_id=strip_passthrough(graph, desc["input_value_id"]), + theta=float(desc["theta"]), + position_offset=int(desc["position_offset"]), + partial=bool(desc.get("partial_rope", False)), + node_ids=tuple(sorted(set(desc.get("node_ids", ())))), + ) + + +def _extract_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: + current = strip_layout_passthrough(graph, value_id) + node = producer(graph, current) + if node is None: + return None + if node.op == "rope": + return { + "input_value_id": node.inputs[0], + "theta": float(node.attrs.get("theta", 0.0)), + "position_offset": int(node.attrs.get("position_offset", 0)), + "partial_rope": False, + "node_ids": (node.id,), + } + classic = _extract_classic_rope_descriptor(graph, current) + if classic is not None: + classic["partial_rope"] = False + return classic + if node.op == "cat" and len(node.inputs) == 2: + for rope_branch_value_id, passthrough_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + rope_branch = _extract_rope_descriptor(graph, rope_branch_value_id) + if rope_branch is None: + continue + passthrough_node = producer(graph, strip_layout_passthrough(graph, passthrough_value_id)) + if passthrough_node is None or passthrough_node.op != "slice": + continue + if strip_passthrough(graph, passthrough_node.inputs[0]) != strip_passthrough(graph, rope_branch["input_value_id"]): + continue + rope_branch["partial_rope"] = True + rope_branch["node_ids"] = tuple(sorted({node.id, passthrough_node.id, *rope_branch["node_ids"]})) + return rope_branch + return None + + +def _extract_classic_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: + node = producer(graph, strip_layout_passthrough(graph, value_id)) + if node is None or node.op != "add" or len(node.inputs) != 2: + return None + for direct_mult_id, rotated_mult_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + direct_mult = producer(graph, direct_mult_id) + rotated_mult = producer(graph, rotated_mult_id) + if direct_mult is None or direct_mult.op != "multiply" or rotated_mult is None or rotated_mult.op != "multiply": + continue + direct_match = _extract_direct_rope_branch(graph, direct_mult) + rotated_match = _extract_rotated_rope_branch(graph, rotated_mult) + if direct_match is None or rotated_match is None: + continue + input_value_id = strip_passthrough(graph, direct_match["input_value_id"]) + if input_value_id != strip_passthrough(graph, rotated_match["input_value_id"]): + continue + cos_info = _extract_rope_trig(graph, direct_match["trig_value_id"], expected="scalar_cos") + sin_info = _extract_rope_trig(graph, rotated_match["trig_value_id"], expected="scalar_sin") + if cos_info is None or sin_info is None: + continue + if abs(float(cos_info["theta"]) - float(sin_info["theta"])) > 1e-2: + continue + if int(cos_info["position_offset"]) != int(sin_info["position_offset"]): + continue + return { + "input_value_id": input_value_id, + "theta": cos_info["theta"], + "position_offset": cos_info["position_offset"], + "node_ids": tuple(sorted({node.id, *direct_match["node_ids"], *rotated_match["node_ids"], *cos_info["node_ids"], *sin_info["node_ids"]})), + } + return None + + +def _extract_direct_rope_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: + for input_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + if _unwrap_trig_value(graph, trig_value_id, expected="scalar_cos") is not None: + return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} + return None + + +def _extract_rotated_rope_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: + for rotated_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + input_value_id = _extract_rotate_half_source(graph, rotated_value_id) + if input_value_id is None: + continue + if _unwrap_trig_value(graph, trig_value_id, expected="scalar_sin") is None: + continue + return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} + return None + + +def _extract_rotate_half_source(graph: IRGraph, value_id: str) -> str | None: + cat_node = producer(graph, strip_passthrough(graph, value_id)) + if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: + return None + neg_node = producer(graph, cat_node.inputs[0]) + right_slice = producer(graph, cat_node.inputs[1]) + if neg_node is None or neg_node.op != "scalar_multiply" or float(neg_node.attrs.get("value", 0.0)) != -1.0: + return None + left_slice = producer(graph, neg_node.inputs[0]) + if left_slice is None or right_slice is None or left_slice.op != "slice" or right_slice.op != "slice": + return None + if strip_passthrough(graph, left_slice.inputs[0]) != strip_passthrough(graph, right_slice.inputs[0]): + return None + return left_slice.inputs[0] + + +def _extract_rope_trig(graph: IRGraph, value_id: str, *, expected: str) -> dict[str, object] | None: + trig_info = _unwrap_trig_value(graph, value_id, expected=expected) + if trig_info is None: + return None + trig_node = trig_info["trig_node"] + cat_node = producer(graph, trig_node.inputs[0]) + if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2 or cat_node.inputs[0] != cat_node.inputs[1]: + return None + angle_info = _extract_rope_angle_source(graph, cat_node.inputs[0]) + if angle_info is None: + return None + matmul_node = angle_info["matmul_node"] + inv_freq_const_id = None + arange_node = None + for input_id in matmul_node.inputs: + if inv_freq_const_id is None: + inv_freq_const_id = _find_constant_ancestor(graph, input_id) + if arange_node is None: + arange_node = _find_arange_ancestor(graph, input_id) + if inv_freq_const_id is None or arange_node is None: + return None + theta = _infer_rope_theta(graph.constants[inv_freq_const_id]) + if theta is None: + return None + return { + "theta": theta, + "position_offset": int(arange_node.attrs.get("start", 0)), + "node_ids": tuple(sorted({trig_node.id, cat_node.id, matmul_node.id})), + } + + +def _unwrap_trig_value(graph: IRGraph, value_id: str, *, expected: str) -> dict[str, object] | None: + current = value_id + while True: + node = producer(graph, current) + if node is None: + return None + if node.op in {"unsqueeze", "precision_cast", "reshape", "view", "expand"}: + current = node.inputs[0] + continue + if node.op == "type_as": + current = node.inputs[0] + continue + if node.op == "scalar_multiply": + current = node.inputs[0] + continue + if node.op != expected: + return None + return {"trig_node": node} + + +def _extract_rope_angle_source(graph: IRGraph, value_id: str) -> dict[str, object] | None: + current = value_id + while True: + current = strip_layout_passthrough(graph, current) + node = producer(graph, current) + if node is None: + return None + if node.op == "permute": + matmul_node = producer(graph, node.inputs[0]) + if matmul_node is None or matmul_node.op != "matmul": + return None + return {"matmul_node": matmul_node} + if node.op == "matmul": + return {"matmul_node": node} + if node.op in {"index", "slice"} and len(node.inputs) == 1: + current = node.inputs[0] + continue + return None + + +def _find_constant_ancestor(graph: IRGraph, value_id: str) -> str | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + current = strip_layout_passthrough(graph, current) + if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): + return current + node = producer(graph, current) + if node is None or len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _find_arange_ancestor(graph: IRGraph, value_id: str) -> IRNode | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + current = strip_layout_passthrough(graph, current) + node = producer(graph, current) + if node is None: + return None + if node.op == "arange": + return node + if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): + return None + if len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _infer_rope_theta(value: Any) -> float | None: + if not isinstance(value, torch.Tensor): + return None + flat = value.detach().cpu().float().reshape(-1) + if flat.numel() < 2: + return None + second = float(flat[1].item()) + if second <= 0.0: + return None + inv_count = flat.numel() + return float((1.0 / second) ** inv_count) diff --git a/python/cactus/transpile/graph_ir.py b/python/cactus/transpile/graph_ir.py new file mode 100644 index 000000000..632585c57 --- /dev/null +++ b/python/cactus/transpile/graph_ir.py @@ -0,0 +1,201 @@ +import copy +from dataclasses import dataclass +from dataclasses import field + +import torch + + +def _clone_ir_payload(value, memo): + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, torch.Tensor): + # IR constants may be real tensors or export-time tensor wrappers. + # They are treated as immutable graph payloads, so preserve them by + # reference instead of triggering Tensor.__deepcopy__. + return value + if isinstance(value, tuple): + return tuple(_clone_ir_payload(item, memo) for item in value) + if isinstance(value, list): + return [_clone_ir_payload(item, memo) for item in value] + if isinstance(value, dict): + return { + copy.deepcopy(key, memo): _clone_ir_payload(inner, memo) + for key, inner in value.items() + } + try: + return copy.deepcopy(value, memo) + except Exception: + return value + + +@dataclass +class IRValue: + id: str + shape: tuple[int, ...] | None = None + dtype: str | None = None + producer: str | None = None + users: list[str] = field(default_factory=list) + meta: dict[str, object] = field(default_factory=dict) + + def __deepcopy__(self, memo): + existing = memo.get(id(self)) + if existing is not None: + return existing + cloned = IRValue( + id=self.id, + shape=_clone_ir_payload(self.shape, memo), + dtype=self.dtype, + producer=self.producer, + users=list(self.users), + meta=_clone_ir_payload(self.meta, memo), + ) + memo[id(self)] = cloned + return cloned + +@dataclass +class IRNode: + id: str + op: str + inputs: list[str] + outputs: list[str] + attrs: dict[str, object] = field(default_factory=dict) + meta: dict[str, object] = field(default_factory=dict) + kind: str = "generic" + + def __deepcopy__(self, memo): + existing = memo.get(id(self)) + if existing is not None: + return existing + cloned = IRNode( + id=self.id, + op=self.op, + inputs=list(self.inputs), + outputs=list(self.outputs), + attrs=_clone_ir_payload(self.attrs, memo), + meta=_clone_ir_payload(self.meta, memo), + kind=self.kind, + ) + memo[id(self)] = cloned + return cloned + +@dataclass +class IRGraph: + values: dict[str, IRValue] + nodes: dict[str, IRNode] + order: list[str] + inputs: list[str] + outputs: list[str] + constants: dict[str, object] = field(default_factory=dict) + meta: dict[str, object] = field(default_factory=dict) + + def __deepcopy__(self, memo): + existing = memo.get(id(self)) + if existing is not None: + return existing + cloned = IRGraph( + values={value_id: copy.deepcopy(value, memo) for value_id, value in self.values.items()}, + nodes={node_id: copy.deepcopy(node, memo) for node_id, node in self.nodes.items()}, + order=list(self.order), + inputs=list(self.inputs), + outputs=list(self.outputs), + constants={ + value_id: _clone_ir_payload(constant, memo) + for value_id, constant in self.constants.items() + }, + meta=_clone_ir_payload(self.meta, memo), + ) + memo[id(self)] = cloned + return cloned + + def add_node(self, node: IRNode) -> None: + if node.id in self.nodes: + raise ValueError(f"duplicate IR node id: {node.id}") + self.nodes[node.id] = node + for output in node.outputs: + if output in self.values: + raise ValueError(f"duplicate IR value id: {output}") + self.values[output] = IRValue(id=output, producer=node.id) + + def add_value(self, value: IRValue) -> None: + if value.id in self.values: + raise ValueError(f"duplicate IR value id: {value.id}") + self.values[value.id] = value + + +def verify_ir(graph: IRGraph) -> None: + order_set = set(graph.order) + node_keys = set(graph.nodes.keys()) + if len(order_set) != len(graph.order): + raise ValueError("IR order contains duplicate node ids") + if order_set != node_keys: + missing_in_order = node_keys - order_set + missing_in_nodes = order_set - node_keys + raise ValueError( + "IR order/nodes mismatch: " + f"missing_in_order={sorted(missing_in_order)} " + f"missing_in_nodes={sorted(missing_in_nodes)}" + ) + + seen_nodes: set[str] = set() + seen_outputs: set[str] = set() + expected_users: dict[str, list[str]] = {} + + for value_id in graph.inputs: + if value_id not in graph.values: + raise ValueError(f"IR input missing value entry: {value_id}") + expected_users.setdefault(value_id, []) + + for value_id in graph.constants: + if value_id not in graph.values: + raise ValueError(f"IR constant missing value entry: {value_id}") + expected_users.setdefault(value_id, []) + + for node_id in graph.order: + if node_id in seen_nodes: + raise ValueError(f"duplicate node in IR order: {node_id}") + seen_nodes.add(node_id) + + node = graph.nodes[node_id] + if not node.outputs: + raise ValueError(f"IR node has no outputs: {node_id}") + + for input_id in node.inputs: + if input_id not in graph.values: + raise ValueError(f"IR node {node_id} references missing input value {input_id}") + expected_users.setdefault(input_id, []).append(node_id) + producer = graph.values[input_id].producer + if producer is not None and producer not in graph.nodes: + raise ValueError( + f"IR value {input_id} for node {node_id} has missing producer {producer}" + ) + + for output_id in node.outputs: + if output_id in seen_outputs: + raise ValueError(f"IR value produced more than once: {output_id}") + seen_outputs.add(output_id) + if output_id not in graph.values: + raise ValueError(f"IR node {node_id} missing output value entry {output_id}") + value = graph.values[output_id] + if value.producer != node_id: + raise ValueError( + f"IR value {output_id} producer mismatch: expected {node_id}, got {value.producer}" + ) + expected_users.setdefault(output_id, []) + + for output_id in graph.outputs: + if output_id not in graph.values: + raise ValueError(f"IR graph output missing value entry: {output_id}") + + for value_id, value in graph.values.items(): + expected = expected_users.get(value_id, []) + if value.users != expected: + raise ValueError( + f"IR value {value_id} user list mismatch: expected {expected}, got {value.users}" + ) + if value.producer is None and value_id not in graph.inputs and value_id not in graph.constants: + if value_id in seen_outputs: + raise ValueError(f"IR value {value_id} is produced but has no producer") + if value.producer is not None and value.producer not in graph.nodes: + raise ValueError( + f"IR value {value_id} references unknown producer {value.producer}" + ) diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py new file mode 100644 index 000000000..7bf5d2ca3 --- /dev/null +++ b/python/cactus/transpile/hf_model.py @@ -0,0 +1,3526 @@ +from __future__ import annotations + +import argparse +import copy +import gc +import hashlib +import itertools +import json +import os +import re +import shutil +import sys +from collections.abc import Mapping +from collections import Counter +from dataclasses import dataclass +from dataclasses import fields +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +PYTHON_ROOT = Path(__file__).resolve().parents[2] +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.transpile.runtime_compat import Graph +from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header +from cactus.convert.model_adapters.nemo import ensure_parakeet_tdt_nemo_source +from cactus.transpile.audio_preprocess import generic_log_mel_features as _generic_log_mel_features +from cactus.transpile.audio_preprocess import load_audio_waveform as _load_audio_waveform +from cactus.transpile.audio_preprocess import prepare_cactus_audio_features +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.component_partition import extract_component_subgraphs +from cactus.transpile.component_partition import summarize_ir_components +from cactus.transpile.component_plan import infer_component_plan_from_config +from cactus.transpile.component_pipeline import capture_component_spec +from cactus.transpile.component_pipeline import execute_component_pipeline +from cactus.transpile.multimodal_runtime import prepare_gemma4_multimodal_inputs as _shared_prepare_gemma4_multimodal_inputs +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.lower import TranspiledGraph +from cactus.transpile.lower import _lower_constant_value +from cactus.transpile.lower import _lower_input_value +from cactus.transpile.lower import _lower_ir_node +from cactus.transpile.lower import _lookup_weight_binding +from cactus.transpile.media_limits import resize_static_image +from cactus.transpile.model_adapters import UnsupportedComponentsError +from cactus.transpile.model_adapters import build_component_module_specs +from cactus.transpile.model_adapters import canonicalize_model_interface +from cactus.transpile.model_profiles import multimodal_context_tokens_for_model_type +from cactus.transpile.optimize_graph import FusionConfig +from cactus.transpile.optimize_graph import optimize_graph +from cactus.transpile.runtime_support import ensure_transformers_supports_model_type as _ensure_transformers_supports_profiled_model_type +from cactus.transpile.runtime_support import patch_torch_flex_attention_compat as _patch_torch_flex_attention_compat +from cactus.transpile.runtime_support import patch_transformers_torchvision_probe as _patch_transformers_torchvision_probe +from cactus.transpile.tdt_runtime import greedy_decode_parakeet_tdt_token_ids +from cactus.transpile.tdt_runtime import load_tdt_local_model +from cactus.transpile.tdt_runtime import prepare_parakeet_tdt_audio_features +from cactus.transpile.weight_compat import ensure_binding_compatible + +_DEFAULT_CAUSAL_PROMPT = "The capital of France is" +_DEFAULT_MULTIMODAL_CONTEXT_TOKENS = 2048 + + +@dataclass +class PreparedInputs: + names: tuple[str, ...] + tensors: tuple[torch.Tensor, ...] + metadata: dict[str, object] + + +def _ensure_transformers_supports_model_type(model_type: str) -> str | None: + return _ensure_transformers_supports_profiled_model_type(model_type) + + +def _resolve_local_snapshot(model_id_or_path: str) -> str | None: + explicit = Path(model_id_or_path) + if explicit.exists(): + nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path) + if nemo_export is not None: + return nemo_export + return str(explicit) + + snapshots_dir = ( + Path.home() + / ".cache" + / "huggingface" + / "hub" + / ("models--" + model_id_or_path.replace("/", "--")) + / "snapshots" + ) + if not snapshots_dir.exists(): + return None + snapshots = sorted(path for path in snapshots_dir.iterdir() if path.is_dir()) + if not snapshots: + return None + snapshot = str(snapshots[-1]) + nemo_export = ensure_parakeet_tdt_nemo_source(snapshot) + if nemo_export is not None: + return nemo_export + return snapshot + + +def _snapshot_has_model_weights(path: str | Path) -> bool: + root = Path(path) + if not root.exists() or not root.is_dir(): + return False + candidates = ( + "model.safetensors", + "model.safetensors.index.json", + "pytorch_model.bin", + "pytorch_model.bin.index.json", + ) + return any((root / name).exists() for name in candidates) + + +def _validate_weights_dir(weights_dir: str | None, *, model_id: str) -> Path | None: + if not weights_dir: + return None + + root = Path(weights_dir).resolve() + if not root.exists(): + raise RuntimeError( + f"weights_dir does not exist: {root}\n" + "\n" + f"Create the folder first with:\n" + f" cactus convert {model_id} {root}\n" + ) + + manifest_path = root / "weights_manifest.json" + if not manifest_path.exists(): + raise RuntimeError( + f"weights_dir is missing weights_manifest.json: {manifest_path}\n" + "\n" + "The transpiler binds weights only through the converted CQ/Cactus " + "manifest; it no longer guesses filenames from model-specific layer names.\n" + "\n" + f"Re-convert with the current converter:\n" + f" cactus convert {model_id} {root}\n" + ) + + return root + + +def _serialize_json_compatible(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): _serialize_json_compatible(inner) for key, inner in value.items()} + if isinstance(value, (list, tuple)): + return [_serialize_json_compatible(inner) for inner in value] + if isinstance(value, Path): + return str(value) + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, torch.Tensor): + payload = { + "type": "torch.Tensor", + "dtype": str(value.dtype), + "shape": list(value.shape), + } + if value.numel() == 1 and not value.is_meta: + payload["data"] = value.detach().cpu().reshape(-1)[0].item() + return payload + if isinstance(value, np.ndarray): + payload = { + "type": "numpy.ndarray", + "dtype": str(value.dtype), + "shape": list(value.shape), + } + if value.size == 1: + payload["data"] = np.asarray(value).reshape(-1)[0].item() + return payload + if hasattr(value, "__dataclass_fields__"): + return { + field.name: _serialize_json_compatible(getattr(value, field.name)) + for field in fields(value) + } + try: + return repr(value) + except Exception: + return f"<{type(value).__module__}.{type(value).__name__}>" + + +def _slugify_model_artifact_name(value: str) -> str: + return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._-") or "model" + + +def _default_artifact_dir_for_model(model_id: str) -> Path: + return (Path.cwd() / "transpiled" / _slugify_model_artifact_name(model_id)).resolve() + + +def _graph_to_dict(graph) -> dict[str, object]: + return { + "meta": _serialize_json_compatible(graph.meta), + "inputs": list(graph.inputs), + "outputs": list(graph.outputs), + "constants": { + value_id: _serialize_json_compatible(constant) + for value_id, constant in graph.constants.items() + }, + "values": { + value_id: _serialize_json_compatible(value) + for value_id, value in graph.values.items() + }, + "nodes": [ + _serialize_json_compatible(graph.nodes[node_id]) + for node_id in graph.order + ], + } + + +def _write_json(path: Path, payload: dict[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def _component_artifact_name(prefix: str, component: str) -> str: + return f"{prefix}_{component}.json" + + +def _component_graphs_to_payload(component_graphs: dict[str, IRGraph]) -> dict[str, object]: + return { + component: _graph_to_dict(graph) + for component, graph in component_graphs.items() + } + + +def _binding_entries_by_node_id( + bindings: list[dict[str, object]], +) -> dict[int, dict[str, object]]: + result: dict[int, dict[str, object]] = {} + for binding in bindings: + try: + result[int(binding["node_id"])] = binding + except Exception: + continue + return result + + +def _externalize_shared_rope_tables( + transpiled_graph: TranspiledGraph, + *, + weights_root: Path, + registry: dict[str, str], +) -> int: + value_ids = getattr(transpiled_graph, "bound_constant_value_ids", {}) or {} + already_bound = { + int(binding["node_id"]) + for binding in transpiled_graph.bound_constant_bindings + if isinstance(binding, dict) and "node_id" in binding + } + externalized = 0 + for constant in transpiled_graph.bound_constants: + node_id = int(constant.id) + if node_id in already_bound: + continue + value_id = str(value_ids.get(node_id, "")) + if not value_id.startswith("c_rope_table_"): + continue + data = np.ascontiguousarray(constant.numpy()) + digest = hashlib.sha1(data.tobytes()).hexdigest()[:16] + filename = registry.get(digest) + if filename is None: + filename = f"rope_table_{digest}.weights" + save_tensor_with_header(data, weights_root / filename, precision="FP16") + registry[digest] = filename + transpiled_graph.bound_constant_bindings.append( + { + "node_id": node_id, + "value_id": value_id, + "path": filename, + "kind": "weight", + "source_name": value_id, + } + ) + externalized += 1 + return externalized + + +def _embed_materialized_bound_constants(transpiled_graph: TranspiledGraph) -> int: + """Serialize small graph-local constants into the .cactus graph itself. + + CQ model weights remain external mmap files. Only materialized constants + without an existing weight binding are embedded. + """ + existing_bindings = list(_serialize_json_compatible(transpiled_graph.bound_constant_bindings)) + external_node_ids = set(_binding_entries_by_node_id(existing_bindings)) + mark_embedded_input = getattr(transpiled_graph.graph, "mark_embedded_input", None) + if not callable(mark_embedded_input): + raise RuntimeError( + "Cactus runtime does not support embedded graph constants; rebuild with cactus build --python" + ) + + embedded_count = 0 + for constant_tensor in transpiled_graph.bound_constants: + node_id = int(constant_tensor.id) + if node_id in external_node_ids: + continue + mark_embedded_input(constant_tensor) + embedded_count += 1 + return embedded_count + + +def _write_graph_binding_manifest( + *, + artifact_dir: Path, + filename: str, + model_id: str, + model_source: str, + task: str, + family: str, + inputs_metadata: dict[str, object], + transpiled_graph: TranspiledGraph, +) -> Path: + manifest_path = artifact_dir / filename + _write_json( + manifest_path, + { + "model_id": model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(inputs_metadata), + "runtime_input_node_ids": [int(tensor.id) for tensor in transpiled_graph.runtime_inputs], + "output_node_ids": [int(tensor.id) for tensor in transpiled_graph.outputs], + "bound_constant_bindings": _serialize_json_compatible(transpiled_graph.bound_constant_bindings), + }, + ) + return manifest_path + + +def _export_lfm2_vl_position_embedding_grid( + *, + model, + artifact_dir: Path, + component_metadata: dict[str, dict[str, object]], +) -> None: + root = getattr(model, "model", None) + vision_tower = getattr(root, "vision_tower", None) + vision_model = getattr(vision_tower, "vision_model", None) + embeddings = getattr(vision_model, "embeddings", None) + if embeddings is None: + return + pos_size = int(embeddings.position_embedding_size) + embed_dim = int(embeddings.embed_dim) + grid = ( + embeddings.position_embedding.weight.detach().cpu().to(torch.float32) + .reshape(pos_size, pos_size, embed_dim) + .contiguous() + .numpy() + ) + rel_path = Path("components") / "vision_encoder" / "position_embedding_grid.f32" + out_path = artifact_dir / rel_path + out_path.parent.mkdir(parents=True, exist_ok=True) + grid.tofile(out_path) + meta = component_metadata.setdefault("vision_encoder", {}) + meta["position_embedding_grid_path"] = str(rel_path) + meta["position_embedding_grid_shape"] = f"{pos_size},{pos_size},{embed_dim}" + print(f"saved_vision_position_embedding_grid={out_path}") + + +def _write_component_bundle( + *, + artifact_dir: Path, + model_id: str, + model_source: str, + task: str, + family: str, + inputs_metadata: dict[str, object], + raw_component_graphs: dict[str, IRGraph], + optimized_component_graphs: dict[str, IRGraph], + transpiled_component_graphs: dict[str, TranspiledGraph] | None = None, + component_io_signatures: dict[str, dict[str, tuple[str, ...]]] | None = None, + component_metadata: dict[str, dict[str, object]] | None = None, + graph_filename: str = "graph.cactus", +) -> Path: + bundle_dir = artifact_dir / "components" + component_order = [ + component + for component in ( + "source_encoder", + "audio_encoder", + "vision_encoder", + "vision_projector", + "text_embedding", + "lm_encoder", + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + "decoder_embed_chunk", + "decoder", + "decoder_cross_kv", + "lm_encoder_step", + "lm_encoder_media_step", + "decoder_media_step", + "decoder_step", + "unspecified", + ) + if component in raw_component_graphs or component in optimized_component_graphs + ] + extra_components = sorted( + component + for component in set(raw_component_graphs) | set(optimized_component_graphs) + if component not in component_order + ) + component_order.extend(extra_components) + + manifest_components: list[dict[str, object]] = [] + + rope_table_registry: dict[str, str] = {} + for component in component_order: + raw_graph = raw_component_graphs.get(component) + optimized_graph = optimized_component_graphs.get(component) + transpiled_graph = None if transpiled_component_graphs is None else transpiled_component_graphs.get(component) + component_dir = bundle_dir / component + raw_relpath = None + optimized_relpath = None + graph_relpath = None + if raw_graph is not None: + raw_relpath = Path(component) / "raw_ir.json" + _write_json( + bundle_dir / raw_relpath, + { + "model_id": model_id, + "model_source": model_source, + "task": task, + "family": family, + "component": component, + "inputs": _serialize_json_compatible(inputs_metadata), + "graph": _graph_to_dict(raw_graph), + }, + ) + if optimized_graph is not None: + optimized_relpath = Path(component) / "optimized_ir.json" + _write_json( + bundle_dir / optimized_relpath, + { + "model_id": model_id, + "model_source": model_source, + "task": task, + "family": family, + "component": component, + "inputs": _serialize_json_compatible(inputs_metadata), + "graph": _graph_to_dict(optimized_graph), + }, + ) + if transpiled_graph is not None: + graph_relpath = Path(component) / graph_filename + component_dir.mkdir(parents=True, exist_ok=True) + shutil.rmtree(component_dir / "bound_constants", ignore_errors=True) + _externalize_shared_rope_tables( + transpiled_graph, + weights_root=artifact_dir, + registry=rope_table_registry, + ) + _embed_materialized_bound_constants(transpiled_graph) + transpiled_graph.graph.save(bundle_dir / graph_relpath) + + graph_for_signature = optimized_graph or raw_graph + if graph_for_signature is None: + continue + manifest_components.append( + { + "component": component, + "directory": str(component_dir.relative_to(artifact_dir)), + "raw_ir": None if raw_relpath is None else str((bundle_dir / raw_relpath).relative_to(artifact_dir)), + "optimized_ir": None if optimized_relpath is None else str((bundle_dir / optimized_relpath).relative_to(artifact_dir)), + "graph": None if graph_relpath is None else str((bundle_dir / graph_relpath).relative_to(artifact_dir)), + "inputs": list(graph_for_signature.inputs), + "outputs": list(graph_for_signature.outputs), + "logical_inputs": list((component_io_signatures or {}).get(component, {}).get("input_keys", ())), + "logical_outputs": list((component_io_signatures or {}).get(component, {}).get("output_keys", ())), + "metadata": _serialize_json_compatible((component_metadata or {}).get(component, {})), + "node_count": len(graph_for_signature.order), + "weight_binding_count": _count_weight_bindings(graph_for_signature), + "runtime_input_node_ids": [] if transpiled_graph is None else [int(tensor.id) for tensor in transpiled_graph.runtime_inputs], + "output_node_ids": [] if transpiled_graph is None else [int(tensor.id) for tensor in transpiled_graph.outputs], + "cache_state_node_ids": [] if transpiled_graph is None else [ + { + "layer_key": str(layer_key), + "key": int(key_tensor.id), + "value": int(value_tensor.id), + } + for layer_key, key_tensor, value_tensor in getattr(transpiled_graph, "cache_state_tensors", []) + ], + "bound_constant_bindings": [] if transpiled_graph is None else ( + list(_serialize_json_compatible(transpiled_graph.bound_constant_bindings)) + ), + } + ) + + manifest_path = bundle_dir / "manifest.json" + manifest_payload: dict[str, object] = { + "model_id": model_id, + "model_source": model_source, + "task": task, + "family": family, + "component_order": component_order, + "inputs": _serialize_json_compatible(inputs_metadata), + "components": manifest_components, + } + _write_json(manifest_path, manifest_payload) + return manifest_path + + +def _named_tensor_store(prepared: PreparedInputs) -> dict[str, torch.Tensor]: + return { + name: tensor + for name, tensor in zip(prepared.names, prepared.tensors, strict=True) + } + + +def _aggregate_component_counts(component_graphs: dict[str, IRGraph]) -> dict[str, int]: + return { + component: len(graph.order) + for component, graph in sorted(component_graphs.items()) + } + + +def _aggregate_component_op_counts(component_graphs: dict[str, IRGraph]) -> Counter: + counter: Counter[str] = Counter() + for graph in component_graphs.values(): + counter.update(graph.nodes[node_id].op for node_id in graph.order) + return counter + + +def _count_component_weight_bindings(component_graphs: dict[str, IRGraph]) -> int: + return sum(_count_weight_bindings(graph) for graph in component_graphs.values()) + + +def _execute_gemma4_component_pipeline( + *, + component_graphs, + prepared: PreparedInputs, +) -> np.ndarray: + store, _ = execute_component_pipeline( + [component_graphs[name] for name in ("vision_encoder", "audio_encoder", "lm_encoder", "decoder")], + initial_store=_named_tensor_store(prepared), + ) + logits = store["logits"] + return np.asarray(logits, dtype=np.float32) + + +def _run_parakeet_tdt_component_decode( + *, + component_graphs, + model, + prepared: PreparedInputs, +) -> dict[str, object]: + store, _ = execute_component_pipeline( + [component_graphs["audio_encoder"]], + initial_store=_named_tensor_store(prepared), + ) + encoder_hidden_states = np.asarray(store["encoder_hidden_states"]) + batch_size = int(encoder_hidden_states.shape[0]) + if batch_size != 1: + raise ValueError("Parakeet TDT component decode currently expects batch size 1") + + hidden_dtype = prepared.tensors[0].dtype + initial_states = model.initial_decoder_state( + batch_size=batch_size, + device=torch.device("cpu"), + dtype=hidden_dtype, + ) + initial_state_arrays = tuple(state.detach().cpu().numpy() for state in initial_states) + decoder_component = component_graphs["decoder"] + + def _step( + frame: np.ndarray, + token_id: int, + state_values: tuple[np.ndarray, ...], + ) -> tuple[np.ndarray, tuple[np.ndarray, ...]]: + input_store: dict[str, object] = { + "encoder_frame": np.ascontiguousarray(frame), + "token_ids": np.full((batch_size,), token_id, dtype=np.int64), + } + for index in range(model.config.predictor_num_layers): + input_store[f"state_h_{index}"] = np.ascontiguousarray(state_values[index * 2]) + input_store[f"state_c_{index}"] = np.ascontiguousarray(state_values[index * 2 + 1]) + + runtime_inputs = [input_store[key] for key in decoder_component.input_keys] + decoder_component.transpiled_graph.set_inputs(runtime_inputs) + outputs = decoder_component.transpiled_graph.execute() + logits = outputs[0].numpy().astype(np.float32, copy=False) + next_states = tuple(output.numpy() for output in outputs[1:]) + return logits, next_states + + emitted = greedy_decode_parakeet_tdt_token_ids( + config=model.config, + encoder_hidden_states=encoder_hidden_states, + initial_states=initial_state_arrays, + step=_step, + ) + + return { + "token_ids": emitted, + "transcript": model.decode_token_ids(emitted), + "encoder_hidden_shape": list(encoder_hidden_states.shape), + } + + +def _run_component_pipeline_transpile( + *, + args, + task: str, + family: str, + model_source: str, + model, + prepared: PreparedInputs, + component_specs, + fusion_config: FusionConfig, + weights_dir: str | None, + artifact_dir: Path | None, + processor_or_tokenizer, + canonical, +) -> int: + print(f"model_id={args.model_id}") + print(f"model_source={model_source}") + print(f"task={task}") + print(f"adapter_family={family}") + print("adapter_module=component_pipeline") + print(f"input_names={','.join(prepared.names)}") + for name, tensor in zip(prepared.names, prepared.tensors, strict=True): + print(f"input_{name}_shape={list(tensor.shape)}") + if weights_dir: + print(f"weights_dir={weights_dir}") + + print("capture_begin=true", flush=True) + captured_components = {} + for spec in component_specs: + print(f"capture_component_begin={spec.component}", flush=True) + captured_components[spec.component] = capture_component_spec(spec, fusion_config=fusion_config) + print(f"capture_component_done={spec.component}", flush=True) + print("capture_done=true", flush=True) + + raw_component_graphs = { + name: captured.raw_ir_graph + for name, captured in captured_components.items() + } + optimized_component_graphs = { + name: captured.optimized_ir_graph + for name, captured in captured_components.items() + } + transpiled_component_graphs = { + name: captured.transpiled_graph + for name, captured in captured_components.items() + } + component_io_signatures = { + name: { + "input_keys": tuple(captured.input_keys), + "output_keys": tuple(captured.output_keys), + } + for name, captured in captured_components.items() + } + component_metadata = { + name: dict(captured.metadata) + for name, captured in captured_components.items() + } + + raw_ir_nodes = sum(len(graph.order) for graph in raw_component_graphs.values()) + optimized_ir_nodes = sum(len(graph.order) for graph in optimized_component_graphs.values()) + binding_count = _count_component_weight_bindings(optimized_component_graphs) + raw_component_counts = _aggregate_component_counts(raw_component_graphs) + optimized_component_counts = _aggregate_component_counts(optimized_component_graphs) + op_counts = _aggregate_component_op_counts(optimized_component_graphs) + + print(f"raw_ir_nodes={raw_ir_nodes}") + print(f"optimized_ir_nodes={optimized_ir_nodes}") + print(f"weight_bindings={binding_count}") + if raw_component_counts: + print( + "raw_components=" + + ",".join(f"{name}:{count}" for name, count in sorted(raw_component_counts.items())) + ) + if optimized_component_counts: + print( + "optimized_components=" + + ",".join(f"{name}:{count}" for name, count in sorted(optimized_component_counts.items())) + ) + print( + "ops=" + f"attention:{op_counts.get('attention', 0)} " + f"conv1d:{op_counts.get('conv1d', 0)} " + f"conv2d:{op_counts.get('conv2d', 0)} " + f"batch_norm:{op_counts.get('batch_norm', 0)} " + f"layer_norm:{op_counts.get('layer_norm', 0)} " + f"rms_norm:{op_counts.get('rms_norm', 0)} " + f"rope:{op_counts.get('rope', 0)} " + f"linear:{op_counts.get('linear', 0)}" + ) + + if weights_dir and binding_count == 0: + raise RuntimeError( + f"No weight bindings were resolved from {weights_dir}\n" + "\n" + "The weights folder exists, but none of the component graphs matched entries in weights_manifest.json.\n" + "\n" + f"Recommended fix:\n" + f" cactus convert {args.model_id} {weights_dir}\n" + ) + + if artifact_dir is not None: + artifact_dir.mkdir(parents=True, exist_ok=True) + _write_json( + artifact_dir / "raw_ir.json", + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "component_order": [spec.component for spec in component_specs], + "components": _component_graphs_to_payload(raw_component_graphs), + }, + ) + _write_json( + artifact_dir / "optimized_ir.json", + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "component_order": [spec.component for spec in component_specs], + "components": _component_graphs_to_payload(optimized_component_graphs), + }, + ) + for component, component_graph in raw_component_graphs.items(): + _write_json( + artifact_dir / _component_artifact_name("raw_ir", component), + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "component": component, + "inputs": _serialize_json_compatible(prepared.metadata), + "graph": _graph_to_dict(component_graph), + }, + ) + for component, component_graph in optimized_component_graphs.items(): + _write_json( + artifact_dir / _component_artifact_name("optimized_ir", component), + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "component": component, + "inputs": _serialize_json_compatible(prepared.metadata), + "graph": _graph_to_dict(component_graph), + }, + ) + print(f"saved_raw_ir={artifact_dir / 'raw_ir.json'}") + print(f"saved_optimized_ir={artifact_dir / 'optimized_ir.json'}") + for component in raw_component_graphs: + print(f"saved_raw_component_ir_{component}={artifact_dir / _component_artifact_name('raw_ir', component)}") + for component in optimized_component_graphs: + print( + f"saved_optimized_component_ir_{component}=" + f"{artifact_dir / _component_artifact_name('optimized_ir', component)}" + ) + + if family == "lfm2_vl" and "vision_encoder" in component_metadata: + _export_lfm2_vl_position_embedding_grid( + model=model, + artifact_dir=artifact_dir, + component_metadata=component_metadata, + ) + + component_manifest_path = _write_component_bundle( + artifact_dir=artifact_dir, + model_id=args.model_id, + model_source=model_source, + task=task, + family=family, + inputs_metadata=prepared.metadata, + raw_component_graphs=raw_component_graphs, + optimized_component_graphs=optimized_component_graphs, + transpiled_component_graphs=transpiled_component_graphs, + component_io_signatures=component_io_signatures, + component_metadata=component_metadata, + graph_filename=args.graph_filename, + ) + print(f"saved_component_bundle_manifest={component_manifest_path}") + for component in transpiled_component_graphs: + print( + f"saved_component_graph_{component}=" + f"{artifact_dir / 'components' / component / args.graph_filename}" + ) + print("note=saved split component graphs; use the component bundle manifest to rebind mmap weights/embeddings when loading") + + if args.skip_execute: + return 0 + + if task == "multimodal_causal_lm_logits": + print("execute_begin=true") + transpiled_output = _execute_gemma4_component_pipeline( + component_graphs=captured_components, + prepared=prepared, + ) + print("execute_done=true") + tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) + if args.skip_reference_compare: + result_payload = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "output_shape": list(transpiled_output.shape), + "raw_ir_nodes": raw_ir_nodes, + "optimized_ir_nodes": optimized_ir_nodes, + "weight_bindings": binding_count, + "reference_compare_skipped": True, + } + print(f"output_shape={list(transpiled_output.shape)}") + transpiled_next = int(np.argmax(transpiled_output[0, -1])) + print(f"transpiled_next_token_id={transpiled_next}") + result_payload["transpiled_next_token_id"] = transpiled_next + if hasattr(tokenizer_like, "decode"): + transpiled_token = tokenizer_like.decode([transpiled_next]) + print(f"transpiled_next_token={transpiled_token!r}") + result_payload["transpiled_next_token"] = transpiled_token + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + + print("reference_begin=true") + with torch.no_grad(): + reference_output = canonical.module(*prepared.tensors).detach().float().cpu().numpy() + print("reference_done=true") + max_abs_diff = float(np.max(np.abs(reference_output - transpiled_output))) + mean_abs_diff = float(np.mean(np.abs(reference_output - transpiled_output))) + hf_next = int(np.argmax(reference_output[0, -1])) + transpiled_next = int(np.argmax(transpiled_output[0, -1])) + print(f"hf_next_token_id={hf_next}") + print(f"transpiled_next_token_id={transpiled_next}") + print(f"logits_max_abs_diff={max_abs_diff:.6f}") + print(f"logits_mean_abs_diff={mean_abs_diff:.6f}") + if hasattr(tokenizer_like, "decode"): + print(f"hf_next_token={tokenizer_like.decode([hf_next])!r}") + print(f"transpiled_next_token={tokenizer_like.decode([transpiled_next])!r}") + result_payload = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "output_shape": list(reference_output.shape), + "raw_ir_nodes": raw_ir_nodes, + "optimized_ir_nodes": optimized_ir_nodes, + "weight_bindings": binding_count, + "max_abs_diff": max_abs_diff, + "mean_abs_diff": mean_abs_diff, + "hf_next_token_id": hf_next, + "transpiled_next_token_id": transpiled_next, + } + if hasattr(tokenizer_like, "decode"): + result_payload["hf_next_token"] = tokenizer_like.decode([hf_next]) + result_payload["transpiled_next_token"] = tokenizer_like.decode([transpiled_next]) + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + + if task == "causal_lm_logits": + tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) + result_payload = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "raw_ir_nodes": raw_ir_nodes, + "optimized_ir_nodes": optimized_ir_nodes, + "weight_bindings": binding_count, + } + if "decoder" in captured_components: + print("execute_begin=true") + initial_store = { + name: tensor + for name, tensor in zip(prepared.names, prepared.tensors, strict=True) + } + store, _ = execute_component_pipeline( + [captured_components["decoder"]], + initial_store=initial_store, + ) + print("execute_done=true") + transpiled_output = np.asarray(store["logits"]) + transpiled_next = int(np.argmax(transpiled_output[0, -1])) + print(f"output_shape={list(transpiled_output.shape)}") + print(f"transpiled_next_token_id={transpiled_next}") + result_payload["output_shape"] = list(transpiled_output.shape) + result_payload["transpiled_next_token_id"] = transpiled_next + if hasattr(tokenizer_like, "decode"): + transpiled_token = tokenizer_like.decode([transpiled_next]) + print(f"transpiled_next_token={transpiled_token!r}") + result_payload["transpiled_next_token"] = transpiled_token + if not args.skip_reference_compare and canonical is not None: + print("reference_begin=true") + with torch.no_grad(): + reference_output = canonical.module(*prepared.tensors).detach().float().cpu().numpy() + print("reference_done=true") + max_abs_diff = float(np.max(np.abs(reference_output - transpiled_output))) + mean_abs_diff = float(np.mean(np.abs(reference_output - transpiled_output))) + hf_next = int(np.argmax(reference_output[0, -1])) + print(f"hf_next_token_id={hf_next}") + print(f"logits_max_abs_diff={max_abs_diff:.6f}") + print(f"logits_mean_abs_diff={mean_abs_diff:.6f}") + result_payload["hf_next_token_id"] = hf_next + result_payload["max_abs_diff"] = max_abs_diff + result_payload["mean_abs_diff"] = mean_abs_diff + if hasattr(tokenizer_like, "decode"): + result_payload["hf_next_token"] = tokenizer_like.decode([hf_next]) + else: + result_payload["reference_compare_skipped"] = True + result_payload["note"] = "causal LM component bundle contains only cached decode step components" + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + + if task == "tdt_transcription": + print("execute_begin=true") + transpiled_decode = _run_parakeet_tdt_component_decode( + component_graphs=captured_components, + model=model, + prepared=prepared, + ) + print("execute_done=true") + print(f"transpiled_transcript={transpiled_decode['transcript']!r}") + result_payload = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": family, + "inputs": _serialize_json_compatible(prepared.metadata), + "raw_ir_nodes": raw_ir_nodes, + "optimized_ir_nodes": optimized_ir_nodes, + "weight_bindings": binding_count, + "transpiled_token_ids": transpiled_decode["token_ids"], + "transpiled_transcript": transpiled_decode["transcript"], + "encoder_hidden_shape": transpiled_decode["encoder_hidden_shape"], + } + if args.skip_reference_compare: + result_payload["reference_compare_skipped"] = True + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + + print("reference_begin=true") + reference_token_ids = model.greedy_decode_token_ids(prepared.tensors[0]) + reference_transcript = model.decode_token_ids(reference_token_ids) + print("reference_done=true") + print(f"reference_transcript={reference_transcript!r}") + print(f"transcript_match={reference_transcript == transpiled_decode['transcript']}") + result_payload["reference_token_ids"] = reference_token_ids + result_payload["reference_transcript"] = reference_transcript + result_payload["transcript_match"] = bool(reference_transcript == transpiled_decode["transcript"]) + result_payload["token_id_match"] = bool(reference_token_ids == transpiled_decode["token_ids"]) + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + return 0 + + raise NotImplementedError(f"component pipeline execution is not implemented for task={task}") + + +def _parse_dtype(name: str) -> torch.dtype: + normalized = name.strip().lower() + mapping = { + "float16": torch.float16, + "fp16": torch.float16, + "float32": torch.float32, + "fp32": torch.float32, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + } + if normalized not in mapping: + raise ValueError(f"unsupported torch dtype: {name}") + return mapping[normalized] + + +def _to_meta_tensor(value: torch.Tensor) -> torch.Tensor: + if value.is_meta: + return value + return torch.empty_strided( + tuple(value.shape), + tuple(value.stride()), + dtype=value.dtype, + device="meta", + requires_grad=value.requires_grad, + ) + + +def _prepared_inputs_for_meta_capture(prepared: PreparedInputs) -> PreparedInputs: + return PreparedInputs( + names=prepared.names, + tensors=tuple(_to_meta_tensor(tensor) for tensor in prepared.tensors), + metadata=prepared.metadata, + ) + + +def _materialize_meta_rope_buffers(model: torch.nn.Module, *, dtype: torch.dtype) -> None: + def _compute_rope( + compute_default, + config: object, + *, + layer_type: str | None = None, + rope_init_fn=None, + rope_type: str | None = None, + ) -> torch.Tensor: + fn = rope_init_fn or compute_default + kwargs: dict[str, object] = {"device": torch.device("cpu")} + if layer_type is not None: + kwargs["layer_type"] = layer_type + if layer_type == "full_attention" and rope_type == "proportional": + kwargs["head_dim_key"] = "global_head_dim" + try: + materialized, _ = fn(config, **kwargs) + except TypeError: + materialized, _ = fn(config=config, **kwargs) + return materialized.to(dtype=dtype) + + for module in model.modules(): + compute_default = getattr(module, "compute_default_rope_parameters", None) + config = getattr(module, "config", None) + if not callable(compute_default) or config is None: + continue + buffers = getattr(module, "_buffers", None) + if not isinstance(buffers, dict): + continue + inv_freq = buffers.get("inv_freq") + if isinstance(inv_freq, torch.Tensor) and inv_freq.is_meta: + materialized = _compute_rope(compute_default, config) + buffers["inv_freq"] = materialized + original = buffers.get("original_inv_freq") + if isinstance(original, torch.Tensor) and original.is_meta: + buffers["original_inv_freq"] = materialized.clone() + + rope_init_fns = getattr(module, "rope_init_fns", None) + rope_types = getattr(module, "rope_type", None) + layer_types = getattr(module, "layer_types", ()) + if not isinstance(layer_types, (set, tuple, list)): + continue + for layer_type in layer_types: + if not isinstance(layer_type, str): + continue + key = f"{layer_type}_inv_freq" + inv_freq = buffers.get(key) + if not isinstance(inv_freq, torch.Tensor) or not inv_freq.is_meta: + continue + rope_init_fn = rope_init_fns.get(layer_type) if isinstance(rope_init_fns, dict) else None + rope_type = rope_types.get(layer_type) if isinstance(rope_types, dict) else None + materialized = _compute_rope( + compute_default, + config, + layer_type=layer_type, + rope_init_fn=rope_init_fn, + rope_type=rope_type if isinstance(rope_type, str) else None, + ) + buffers[key] = materialized + original = buffers.get(f"{layer_type}_original_inv_freq") + if isinstance(original, torch.Tensor) and original.is_meta: + buffers[f"{layer_type}_original_inv_freq"] = materialized.clone() + + +def _materialize_meta_scalar_buffers(model: torch.nn.Module, *, dtype: torch.dtype) -> None: + for module in model.modules(): + buffers = getattr(module, "_buffers", None) + if not isinstance(buffers, dict): + continue + embed_scale = buffers.get("embed_scale") + scalar_embed_scale = getattr(module, "scalar_embed_scale", None) + if isinstance(embed_scale, torch.Tensor) and embed_scale.is_meta and scalar_embed_scale is not None: + buffers["embed_scale"] = torch.tensor( + float(scalar_embed_scale), + dtype=embed_scale.dtype if embed_scale.dtype.is_floating_point else dtype, + device="cpu", + ) + + +class TranspileWrapper(torch.nn.Module): + def __init__(self, adapter_module: torch.nn.Module, *, weights_dir: str | None = None): + super().__init__() + self.adapter = adapter_module + self.weights_dir = weights_dir + + def forward(self, *bound_inputs: torch.Tensor) -> torch.Tensor: + return self.adapter(*bound_inputs) + + def get_transpile_metadata(self) -> dict[str, object]: + metadata: dict[str, object] = {} + provider = getattr(self.adapter, "get_transpile_metadata", None) + if callable(provider): + provided = provider() + if isinstance(provided, dict): + metadata.update(provided) + graph_meta: dict[str, object] = {} + base_graph = metadata.get("graph", {}) + if isinstance(base_graph, dict): + graph_meta.update(base_graph) + if self.weights_dir: + graph_meta["weights_dir"] = self.weights_dir + metadata["graph"] = graph_meta + return metadata + + +def _tie_lfm2_vl_lm_head_if_needed(model: torch.nn.Module) -> str | None: + if str(getattr(getattr(model, "config", None), "model_type", "") or "").lower() != "lfm2_vl": + return None + lm_head = getattr(model, "lm_head", None) + model_root = getattr(model, "model", None) + language_model = getattr(model_root, "language_model", None) + embed_tokens = getattr(language_model, "embed_tokens", None) + if not isinstance(lm_head, torch.nn.Linear) or not isinstance(embed_tokens, torch.nn.Embedding): + return None + if tuple(lm_head.weight.shape) != tuple(embed_tokens.weight.shape): + return None + if lm_head.weight.is_meta or embed_tokens.weight.is_meta: + return None + if lm_head.weight.data_ptr() == embed_tokens.weight.data_ptr(): + return None + lm_head.weight = embed_tokens.weight + return "tied LFM2-VL lm_head.weight to language_model.embed_tokens.weight" + + +def _load_local_torch_state_dict(model_source: str) -> dict[str, torch.Tensor] | None: + root = Path(model_source) + if not root.exists() or not root.is_dir(): + return None + + safetensors_path = root / "model.safetensors" + if safetensors_path.exists(): + try: + from safetensors.torch import load_file # type: ignore + + return load_file(str(safetensors_path)) + except Exception: + return None + + pytorch_path = root / "pytorch_model.bin" + if pytorch_path.exists(): + try: + loaded = torch.load(str(pytorch_path), map_location="cpu", weights_only=True) + except TypeError: + loaded = torch.load(str(pytorch_path), map_location="cpu") + return loaded if isinstance(loaded, dict) else None + + return None + + +def _remap_gemma4_checkpoint_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + remapped: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + new_key = key + if "audio_tower" in new_key: + new_key = re.sub(r"subsample_conv_projection\.layer(\d+)\.", r"subsample_conv_projection.conv_\1.", new_key) + new_key = new_key.replace("audio_tower.layers.", "audio_tower.conformer.") + new_key = new_key.replace(".feed_forward1.", ".ffw_layer_start.") + new_key = new_key.replace(".feed_forward2.", ".ffw_layer_end.") + new_key = re.sub(r"\.self_attn\.(q_proj|k_proj|v_proj)\.", r".attention.attn.\1.", new_key) + new_key = new_key.replace(".self_attn.per_dim_scale", ".attention.attn.per_dim_scale") + new_key = new_key.replace(".self_attn.relative_k_proj.", ".attention.attn.relative_position_embedding.pos_proj.") + new_key = new_key.replace(".self_attn.post.", ".attention.post.") + new_key = new_key.replace(".norm_pre_attn.", ".attention.pre_attn_norm.") + new_key = new_key.replace(".norm_post_attn.", ".attention.post_norm.") + new_key = new_key.replace(".norm_out.", ".norm.") + new_key = new_key.replace(".linear.weight", ".weight") + remapped[new_key] = value + if new_key.endswith(".attention.attn.per_dim_scale"): + remapped[new_key.replace(".per_dim_scale", ".per_dim_key_scale")] = value + if "lm_head.weight" not in remapped: + tied_embedding = remapped.get("model.language_model.embed_tokens.weight") + if tied_embedding is not None: + remapped["lm_head.weight"] = tied_embedding + return remapped + + +def _repair_gemma4_checkpoint_weights(model: torch.nn.Module, model_source: str) -> dict[str, object]: + multimodal_backbone = getattr(model, "model", model) + audio_tower = getattr(multimodal_backbone, "audio_tower", None) + # Older/local Gemma4 variants used Cactus-style module names such as + # audio_tower.conformer and bare clippable-linear weights. Current HF + # Gemma4 classes use audio_tower.layers and *.linear.weight, and + # from_pretrained already loads them correctly. Do not reload a legacy + # remap into the current layout, because that leaves real weights missing. + if hasattr(audio_tower, "layers"): + return {"applied": False, "reason": "current HF Gemma4 layout uses audio_tower.layers"} + if not hasattr(audio_tower, "conformer"): + return {"applied": False, "reason": "current HF Gemma4 layout does not need legacy key remap"} + + raw_state_dict = _load_local_torch_state_dict(model_source) + if raw_state_dict is None: + return {"applied": False, "reason": "no local checkpoint state_dict"} + + remapped_state_dict = _remap_gemma4_checkpoint_state_dict(raw_state_dict) + load_result = model.load_state_dict(remapped_state_dict, strict=False) + return { + "applied": True, + "missing_keys": list(load_result.missing_keys), + "unexpected_keys": list(load_result.unexpected_keys), + } + + +def _load_config_json(model_id_or_path: str) -> dict[str, object]: + local_snapshot = _resolve_local_snapshot(model_id_or_path) + config_source = Path(local_snapshot) / "config.json" if local_snapshot else None + if config_source is not None and config_source.exists(): + return json.loads(config_source.read_text()) + explicit = Path(model_id_or_path) / "config.json" + if explicit.exists(): + return json.loads(explicit.read_text()) + return {} + + +def _load_optional_json(model_id_or_path: str, filename: str) -> dict[str, object]: + local_snapshot = _resolve_local_snapshot(model_id_or_path) + candidate = Path(local_snapshot) / filename if local_snapshot else None + if candidate is not None and candidate.exists(): + return json.loads(candidate.read_text()) + explicit = Path(model_id_or_path) / filename + if explicit.exists(): + return json.loads(explicit.read_text()) + return {} + + +def _load_gemma4_tokenizer_fallback(source_candidates: list[str]) -> object | None: + try: + from transformers import PreTrainedTokenizerFast # type: ignore + except Exception: + return None + + for source in source_candidates: + root = Path(source) + if not root.exists() or not root.is_dir(): + continue + tokenizer_json = root / "tokenizer.json" + tokenizer_config_path = root / "tokenizer_config.json" + if not tokenizer_json.exists() or not tokenizer_config_path.exists(): + continue + + tokenizer_config = json.loads(tokenizer_config_path.read_text()) + tokenizer = PreTrainedTokenizerFast( + tokenizer_file=str(tokenizer_json), + bos_token=tokenizer_config.get("bos_token"), + eos_token=tokenizer_config.get("eos_token"), + unk_token=tokenizer_config.get("unk_token"), + pad_token=tokenizer_config.get("pad_token"), + mask_token=tokenizer_config.get("mask_token"), + padding_side=str(tokenizer_config.get("padding_side", "right")), + additional_special_tokens=list(tokenizer_config.get("extra_special_tokens", []) or []), + ) + model_max_length = tokenizer_config.get("model_max_length") + if isinstance(model_max_length, int): + tokenizer.model_max_length = model_max_length + + chat_template_path = root / "chat_template.jinja" + if chat_template_path.exists(): + tokenizer.chat_template = chat_template_path.read_text() + + for token_attr in ( + "image_token", + "audio_token", + "boi_token", + "eoi_token", + "boa_token", + "eoa_token", + ): + token_value = tokenizer_config.get(token_attr) + if isinstance(token_value, str): + setattr(tokenizer, token_attr, token_value) + setattr(tokenizer, f"{token_attr}_id", tokenizer.convert_tokens_to_ids(token_value)) + return tokenizer + + return None + + +def _load_gemma4_processor_fallback( + *, + source_candidates: list[str], + common_kwargs: dict[str, object], +) -> object | None: + try: + from transformers.models.gemma4.feature_extraction_gemma4 import Gemma4AudioFeatureExtractor # type: ignore + from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor # type: ignore + from transformers.models.gemma4.processing_gemma4 import Gemma4Processor # type: ignore + except Exception: + return None + + tokenizer = None + processor_config: dict[str, object] = {} + for source in source_candidates: + try: + try: + from transformers import AutoTokenizer # type: ignore + + tokenizer = AutoTokenizer.from_pretrained(source, **common_kwargs) + except Exception: + tokenizer = _load_gemma4_tokenizer_fallback([source]) + processor_config = _load_optional_json(source, "processor_config.json") + if tokenizer is not None and processor_config: + break + except Exception: + continue + if tokenizer is None or not processor_config: + return None + + feature_config = dict(processor_config.get("feature_extractor", {}) or {}) + image_config = dict(processor_config.get("image_processor", {}) or {}) + feature_config.pop("feature_extractor_type", None) + image_config.pop("image_processor_type", None) + + feature_extractor = Gemma4AudioFeatureExtractor(**feature_config) + image_processor = Gemma4ImageProcessor(**image_config) + + processor_kwargs = {} + for key in ("image_seq_length", "audio_seq_length", "audio_ms_per_token"): + if key in processor_config: + processor_kwargs[key] = processor_config[key] + + return Gemma4Processor( + feature_extractor=feature_extractor, + image_processor=image_processor, + tokenizer=tokenizer, + **processor_kwargs, + ) + + +def _infer_task_from_config(model_id_or_path: str) -> str: + config = _load_config_json(model_id_or_path) + plan = infer_component_plan_from_config(config, model_id=model_id_or_path) + if plan is not None: + return plan.task + architectures = [str(value) for value in config.get("architectures", []) if isinstance(value, str)] + model_type = str(config.get("model_type", "") or "").lower() + decoding_cfg = config.get("decoding") + if isinstance(decoding_cfg, dict) and str(decoding_cfg.get("model_type", "") or "").lower() == "tdt": + return "tdt_transcription" + loss_cfg = config.get("loss") + if isinstance(loss_cfg, dict) and str(loss_cfg.get("loss_name", "") or "").lower() == "tdt": + return "tdt_transcription" + + if "nomic" in model_type or any("NomicBert" in value for value in architectures): + return "text_embedding" + if any("CausalLM" in value for value in architectures): + return "causal_lm_logits" + if any("CTC" in value for value in architectures): + return "ctc_logits" + if model_type == "whisper": + return "seq2seq_transcription" + if any("ConditionalGeneration" in value and "Whisper" in value for value in architectures): + return "seq2seq_transcription" + + lowered_id = model_id_or_path.lower() + if "nomic-embed" in lowered_id: + return "text_embedding" + if "parakeet-tdt" in lowered_id: + return "tdt_transcription" + if "whisper" in lowered_id: + return "seq2seq_transcription" + if "ctc" in lowered_id: + return "ctc_logits" + if any(token in lowered_id for token in ("qwen", "gemma", "llama", "mistral", "lfm")): + return "causal_lm_logits" + + raise RuntimeError( + f"Could not infer transpile task for {model_id_or_path}.\n" + "\n" + "Pass one explicitly with --task, for example:\n" + " --task causal_lm_logits\n" + " --task tdt_transcription\n" + " --task ctc_logits\n" + " --task encoder_hidden_states\n" + " --task seq2seq_transcription\n" + ) + + +def _resolve_audio_sample_rate(processor: object) -> int: + for attr_name in ("feature_extractor", "tokenizer"): + child = getattr(processor, attr_name, None) + sample_rate = getattr(child, "sampling_rate", None) + if isinstance(sample_rate, int) and sample_rate > 0: + return sample_rate + sample_rate = getattr(processor, "sampling_rate", None) + if isinstance(sample_rate, int) and sample_rate > 0: + return sample_rate + return 16000 + + +def _infer_fallback_audio_input_names(config: dict[str, object], task: str) -> tuple[str, ...]: + model_type = str(config.get("model_type", "") or "").lower() + if model_type == "whisper": + return ("input_features",) + if task == "seq2seq_transcription": + return ("input_features",) + if task == "encoder_hidden_states": + return ("input_features",) + if task == "tdt_transcription": + return ("input_features",) + audio_cfg = config.get("audio_config") + if isinstance(audio_cfg, dict) and any(key in audio_cfg for key in ("features", "input_feat_size", "num_mel_bins")): + return ("input_features", "attention_mask") if task == "ctc_logits" else ("input_features",) + encoder_cfg = config.get("encoder") + if isinstance(encoder_cfg, dict) and any(key in encoder_cfg for key in ("feat_in", "num_mel_bins")): + return ("input_features", "attention_mask") if task == "ctc_logits" else ("input_features",) + return ("input_values", "attention_mask") if task == "ctc_logits" else ("input_values",) + + +def _resolve_encoder_module(model: torch.nn.Module) -> torch.nn.Module | None: + get_encoder = getattr(model, "get_encoder", None) + if callable(get_encoder): + encoder = get_encoder() + if isinstance(encoder, torch.nn.Module): + return encoder + encoder = getattr(model, "encoder", None) + if isinstance(encoder, torch.nn.Module): + return encoder + model_attr = getattr(model, "model", None) + if model_attr is not None: + encoder = getattr(model_attr, "encoder", None) + if isinstance(encoder, torch.nn.Module): + return encoder + return None + + +def _infer_expected_input_feature_frames(model: torch.nn.Module) -> int | None: + config = getattr(model, "config", None) + max_source_positions = getattr(config, "max_source_positions", None) + if not isinstance(max_source_positions, int) or max_source_positions <= 0: + return None + + encoder = _resolve_encoder_module(model) + if encoder is None: + return None + + stride_product = 1 + found_conv = False + for child in encoder.children(): + if isinstance(child, torch.nn.Conv1d): + stride = child.stride[0] if isinstance(child.stride, tuple) else child.stride + stride_product *= int(stride) + found_conv = True + + if not found_conv: + return None + return int(max_source_positions) * stride_product + + +def _prepare_fallback_audio_inputs( + *, + input_names: tuple[str, ...], + config: dict[str, object], + preprocessor_config: dict[str, object], + model: torch.nn.Module, + task: str, + audio_file: str, + torch_dtype: torch.dtype, +) -> PreparedInputs: + if not input_names: + input_names = _infer_fallback_audio_input_names(config, task) + preprocessor_root = config.get("preprocessor") + preprocessor_root = preprocessor_root if isinstance(preprocessor_root, dict) else {} + encoder_root = config.get("encoder") + encoder_root = encoder_root if isinstance(encoder_root, dict) else {} + + sample_rate = int( + preprocessor_config.get( + "sampling_rate", + preprocessor_root.get("sample_rate", config.get("sampling_rate", 16000)), + ) + or 16000 + ) + num_mels = int( + preprocessor_config.get( + "feature_size", + preprocessor_root.get( + "features", + encoder_root.get("feat_in", config.get("num_mel_bins", config.get("feature_size", 80))), + ), + ) + or 80 + ) + encoder_cfg = config.get("encoder_config") + if isinstance(encoder_cfg, dict): + num_mels = int(encoder_cfg.get("num_mel_bins", encoder_cfg.get("feat_in", num_mels)) or num_mels) + raw_hop_length = preprocessor_config.get( + "hop_length", + preprocessor_root.get("window_stride", config.get("hop_length", 160)), + ) + if raw_hop_length is None: + raw_hop_length = 160 + if isinstance(raw_hop_length, float) and raw_hop_length > 0.0 and raw_hop_length < 1.0: + hop_length = max(1, int(round(float(sample_rate) * raw_hop_length))) + else: + hop_length = int(raw_hop_length) + + n_fft = int(preprocessor_config.get("n_fft", config.get("n_fft", preprocessor_root.get("n_fft", 400))) or 400) + raw_frame_length = preprocessor_config.get( + "win_length", + preprocessor_config.get( + "frame_length", + preprocessor_root.get("window_size", preprocessor_root.get("n_window_size", config.get("frame_length", n_fft))), + ), + ) + if raw_frame_length is None: + raw_frame_length = n_fft + if isinstance(raw_frame_length, float) and raw_frame_length > 0.0 and raw_frame_length < 1.0: + frame_length = max(1, int(round(float(sample_rate) * raw_frame_length))) + else: + frame_length = int(raw_frame_length) + preemphasis = preprocessor_config.get("preemphasis", preprocessor_root.get("preemph")) + if preemphasis is not None: + preemphasis = float(preemphasis) + waveform = _load_audio_waveform(audio_file, target_sample_rate=sample_rate) + + tensors: list[torch.Tensor] = [] + if input_names and input_names[0] == "input_features": + model_type = str(config.get("model_type", "") or "").lower() + is_parakeet_like = ( + task == "tdt_transcription" + or "parakeet" in model_type + or "parakeet" in type(model).__module__.lower() + ) + if is_parakeet_like: + expected_frames = _infer_expected_input_feature_frames(model) + parakeet_features, _ = prepare_parakeet_tdt_audio_features( + audio_file, + expected_frames=expected_frames, + expected_mels=num_mels, + torch_dtype=torch_dtype, + ) + tensors.append(parakeet_features) + return PreparedInputs( + names=input_names[: len(tensors)], + tensors=tuple(tensors), + metadata={ + "audio_file": str(Path(audio_file).resolve()), + "sample_rate": sample_rate, + "fallback_audio_preprocessor": True, + "native_parakeet_audio_preprocessor": True, + "input_shapes": { + name: list(tensor.shape) + for name, tensor in zip(input_names, tensors) + }, + }, + ) + is_whisper_like = ( + str(config.get("model_type", "") or "").lower() == "whisper" + or "transformers.models.whisper." in type(model).__module__ + ) + if is_whisper_like: + expected_frames = _infer_expected_input_feature_frames(model) or 3000 + try: + whisper_features, active_frames = prepare_cactus_audio_features( + audio_file, + model_type="whisper", + expected_frames=expected_frames, + expected_mels=num_mels, + torch_dtype=torch_dtype, + layout="mels_frames", + ) + tensors.append(whisper_features) + return PreparedInputs( + names=input_names[: len(tensors)], + tensors=tuple(tensors), + metadata={ + "audio_file": str(Path(audio_file).resolve()), + "sample_rate": sample_rate, + "fallback_audio_preprocessor": True, + "native_cactus_audio_preprocessor": True, + "active_feature_frames": active_frames, + "input_shapes": { + name: list(tensor.shape) + for name, tensor in zip(input_names, tensors) + }, + }, + ) + except Exception: + pass + features, feature_length = _generic_log_mel_features( + waveform, + sample_rate=sample_rate, + num_mels=num_mels, + n_fft=n_fft, + hop_length=hop_length, + frame_length=frame_length, + preemphasis=preemphasis, + ) + expected_frames = _infer_expected_input_feature_frames(model) + attention_mask: np.ndarray | None = None + target_frames = feature_length + if isinstance(expected_frames, int) and expected_frames > 0: + target_frames = expected_frames + + active_frames = min(feature_length, target_frames) + features = features[:active_frames, :] + if target_frames > active_frames: + pad_width = target_frames - active_frames + features = np.pad(features, ((0, pad_width), (0, 0)), mode="constant") + attention_mask = np.zeros((target_frames,), dtype=np.bool_) + attention_mask[:active_frames] = True + + if is_whisper_like: + features = np.ascontiguousarray(features.T) + + tensors.append(torch.from_numpy(features).unsqueeze(0).to(dtype=torch_dtype)) + if attention_mask is not None and len(input_names) > 1 and input_names[1] == "attention_mask": + tensors.append(torch.from_numpy(attention_mask).unsqueeze(0)) + else: + input_values = torch.from_numpy(waveform).unsqueeze(0).to(dtype=torch_dtype) + tensors.append(input_values) + if len(input_names) > 1 and input_names[1] == "attention_mask": + tensors.append(torch.ones_like(input_values, dtype=torch.float32)) + + return PreparedInputs( + names=input_names[: len(tensors)], + tensors=tuple(tensors), + metadata={ + "audio_file": str(Path(audio_file).resolve()), + "sample_rate": sample_rate, + "fallback_audio_preprocessor": True, + "input_shapes": { + name: list(tensor.shape) + for name, tensor in zip(input_names, tensors) + }, + }, + ) + + +def _prepare_audio_inputs( + processor: object | None, + *, + input_names: tuple[str, ...], + config: dict[str, object], + preprocessor_config: dict[str, object], + model: torch.nn.Module, + task: str, + audio_file: str, + torch_dtype: torch.dtype, +) -> PreparedInputs: + if processor is None: + return _prepare_fallback_audio_inputs( + input_names=input_names, + config=config, + preprocessor_config=preprocessor_config, + model=model, + task=task, + audio_file=audio_file, + torch_dtype=torch_dtype, + ) + + model_type = str(config.get("model_type", "") or "").lower() + if model_type == "whisper" or "transformers.models.whisper." in type(model).__module__: + return _prepare_fallback_audio_inputs( + input_names=input_names, + config=config, + preprocessor_config=preprocessor_config, + model=model, + task=task, + audio_file=audio_file, + torch_dtype=torch_dtype, + ) + + sample_rate = _resolve_audio_sample_rate(processor) + waveform = _load_audio_waveform(audio_file, target_sample_rate=sample_rate) + batch = processor( + waveform, + sampling_rate=sample_rate, + return_tensors="pt", + ) + + preferred_keys = tuple(input_names) + tuple( + key for key in ("input_features", "input_values", "attention_mask") if key not in input_names + ) + tensor_keys = [key for key, value in batch.items() if isinstance(value, torch.Tensor)] + ordered_keys = [key for key in preferred_keys if key in tensor_keys] + ordered_keys.extend(key for key in tensor_keys if key not in ordered_keys) + if not ordered_keys: + raise RuntimeError(f"processor did not return tensor inputs for audio file: {audio_file}") + + tensors: list[torch.Tensor] = [] + for key in ordered_keys: + value = batch[key] + if not isinstance(value, torch.Tensor): + continue + if torch.is_floating_point(value): + value = value.to(dtype=torch_dtype) + tensors.append(value) + + return PreparedInputs( + names=tuple(ordered_keys[: len(tensors)]), + tensors=tuple(tensors), + metadata={ + "audio_file": str(Path(audio_file).resolve()), + "sample_rate": sample_rate, + "fallback_audio_preprocessor": False, + "input_shapes": { + name: list(tensor.shape) + for name, tensor in zip(ordered_keys, tensors) + }, + }, + ) + + +def _contains_token_subsequence(token_ids: list[int], subsequence: list[int]) -> bool: + if not subsequence: + return False + width = len(subsequence) + return any(token_ids[index : index + width] == subsequence for index in range(len(token_ids) - width + 1)) + + +def _resolve_whisper_decoder_prompt_token_ids( + tokenizer_or_processor: object, + *, + prompt: str | None, + decoder_start_token_id: int | None, + forced_decoder_ids: list[list[int]] | tuple[tuple[int, int], ...] | None = None, +) -> list[int]: + tokenizer = getattr(tokenizer_or_processor, "tokenizer", tokenizer_or_processor) + encode = getattr(tokenizer, "encode", None) + if not callable(encode): + encode = None + + normalized_prompt = (prompt or "").strip() + if normalized_prompt == _DEFAULT_CAUSAL_PROMPT: + normalized_prompt = "" + + prompt_token_ids: list[int] = [] + if decoder_start_token_id is not None: + prompt_token_ids.append(int(decoder_start_token_id)) + + normalized_forced_ids: list[tuple[int, int]] = [] + if forced_decoder_ids is not None: + for item in forced_decoder_ids: + if ( + isinstance(item, (list, tuple)) + and len(item) == 2 + and isinstance(item[0], int) + and isinstance(item[1], int) + ): + normalized_forced_ids.append((int(item[0]), int(item[1]))) + normalized_forced_ids.sort(key=lambda pair: pair[0]) + + next_position = 1 + for position, token_id in normalized_forced_ids: + if position != next_position: + break + prompt_token_ids.append(int(token_id)) + next_position += 1 + + if normalized_prompt: + if encode is None: + raise RuntimeError("Whisper transpile requires a tokenizer to encode a non-empty prompt") + try: + encoded_prompt = encode(normalized_prompt, add_special_tokens=False) + except TypeError: + encoded_prompt = encode(normalized_prompt) + prompt_token_ids.extend(int(value) for value in encoded_prompt) + + if not prompt_token_ids: + if encode is None: + raise RuntimeError("Whisper transpile requires a tokenizer or decoder_start_token_id") + try: + prompt_token_ids = [int(value) for value in encode("<|startoftranscript|>", add_special_tokens=False)] + except TypeError: + prompt_token_ids = [int(value) for value in encode("<|startoftranscript|>")] + if not prompt_token_ids and decoder_start_token_id is not None: + prompt_token_ids = [int(decoder_start_token_id)] + + return prompt_token_ids + + +def _augment_whisper_seq2seq_metadata( + prepared: PreparedInputs, + *, + tokenizer_or_processor: object, + model: torch.nn.Module, + prompt: str | None, + max_new_tokens: int, +) -> PreparedInputs: + config = getattr(model, "config", None) + decoder_start_token_id = getattr(config, "decoder_start_token_id", None) + eos_token_id = getattr(config, "eos_token_id", None) + pad_token_id = getattr(config, "pad_token_id", None) + max_target_positions = int(getattr(config, "max_target_positions", 0) or 0) + suppress_tokens = [int(value) for value in (getattr(config, "suppress_tokens", None) or [])] + begin_suppress_tokens = [int(value) for value in (getattr(config, "begin_suppress_tokens", None) or [])] + + decoder_input_ids = _resolve_whisper_decoder_prompt_token_ids( + tokenizer_or_processor, + prompt=prompt, + decoder_start_token_id=int(decoder_start_token_id) if isinstance(decoder_start_token_id, int) else None, + forced_decoder_ids=getattr(config, "forced_decoder_ids", None), + ) + if not decoder_input_ids: + raise RuntimeError("Whisper transpile could not resolve decoder prompt token ids") + + target_token_count = len(decoder_input_ids) + max(1, int(max_new_tokens)) + if max_target_positions > 0: + target_token_count = min(target_token_count, max_target_positions) + target_token_count = max(target_token_count, len(decoder_input_ids)) + + metadata = dict(prepared.metadata) + metadata.update( + { + "decoder_input_ids": [int(value) for value in decoder_input_ids], + "decoder_start_token_id": None if decoder_start_token_id is None else int(decoder_start_token_id), + "eos_token_id": None if eos_token_id is None else int(eos_token_id), + "pad_token_id": None if pad_token_id is None else int(pad_token_id), + "target_token_count": int(target_token_count), + "max_target_positions": int(max_target_positions), + "suppress_tokens": suppress_tokens, + "begin_suppress_tokens": begin_suppress_tokens, + } + ) + return PreparedInputs( + names=prepared.names, + tensors=prepared.tensors, + metadata=metadata, + ) + + +def _tokenize_text_prompt( + tokenizer: object, + prompt: str, + *, + enable_thinking_if_supported: bool = False, +) -> torch.Tensor: + apply_chat_template = getattr(tokenizer, "apply_chat_template", None) + if callable(apply_chat_template): + try: + encoded = apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + enable_thinking=bool(enable_thinking_if_supported), + ) + if isinstance(encoded, torch.Tensor) and encoded.ndim == 2: + return encoded.to(dtype=torch.long) + if isinstance(encoded, Mapping): + input_ids = encoded.get("input_ids") + if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2: + return input_ids.to(dtype=torch.long) + except Exception: + pass + encoded = tokenizer(prompt, return_tensors="pt") + return encoded["input_ids"].to(dtype=torch.long) + + +def _resolve_text_padding_token_id(tokenizer: object | None) -> int: + for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): + token_id = getattr(tokenizer, attr_name, None) if tokenizer is not None else None + if isinstance(token_id, int) and token_id >= 0: + return int(token_id) + return 0 + + +def _resolve_graph_safe_text_padding_token_id( + tokenizer: object | None, + prompt_input_ids: torch.Tensor, +) -> int: + # Keep the tokenizer pad ID in the captured graph metadata. The lowerer keeps + # large token-id comparisons in FP32, so Qwen-style high pad IDs are safe and + # runtime padding stays consistent with the traced mask. + return _resolve_text_padding_token_id(tokenizer) + + +def _prepare_text_inputs( + tokenizer: object, + *, + prompt: str, + input_ids_text: str | None, + max_new_tokens: int, + enable_thinking_if_supported: bool = False, +) -> PreparedInputs: + if input_ids_text: + token_ids = [int(part.strip()) for part in input_ids_text.split(",") if part.strip()] + if not token_ids: + raise ValueError("--input-ids was provided but no ids were parsed") + prompt_input_ids = torch.tensor([token_ids], dtype=torch.long) + else: + prompt_input_ids = _tokenize_text_prompt( + tokenizer, + prompt, + enable_thinking_if_supported=enable_thinking_if_supported, + ) + if prompt_input_ids.ndim != 2 or int(prompt_input_ids.shape[0]) != 1: + raise ValueError( + "causal_lm_logits transpile currently expects prompt input ids with shape [1, T], " + f"got {tuple(int(dim) for dim in prompt_input_ids.shape)}" + ) + if int(max_new_tokens) < 0: + raise ValueError("--max-new-tokens must be non-negative") + + prompt_token_count = int(prompt_input_ids.shape[1]) + target_token_count = prompt_token_count + int(max_new_tokens) + padding_token_id = _resolve_graph_safe_text_padding_token_id(tokenizer, prompt_input_ids) + if target_token_count > prompt_token_count: + input_ids = torch.full((1, target_token_count), padding_token_id, dtype=torch.long) + input_ids[:, :prompt_token_count] = prompt_input_ids + else: + input_ids = prompt_input_ids + return PreparedInputs( + names=("input_ids",), + tensors=(input_ids,), + metadata={ + "prompt": prompt, + "prompt_input_ids": prompt_input_ids.tolist(), + "input_ids": input_ids.tolist(), + "prompt_token_count": prompt_token_count, + "target_token_count": target_token_count, + "max_new_tokens": int(max_new_tokens), + "padding_token_id": int(padding_token_id), + "enable_thinking": bool(enable_thinking_if_supported), + }, + ) + + +_NOMIC_EMBED_SEQ_LEN = 256 + + +def _prepare_text_embedding_inputs( + tokenizer: object, + *, + prompt: str, + input_ids_text: str | None, + seq_len: int = _NOMIC_EMBED_SEQ_LEN, +) -> PreparedInputs: + base = _prepare_text_inputs( + tokenizer, + prompt=prompt, + input_ids_text=input_ids_text, + max_new_tokens=0, + ) + prompt_input_ids = base.tensors[0] + prompt_token_count = int(prompt_input_ids.shape[1]) + padding_token_id = int(base.metadata.get("padding_token_id", 0)) + target = max(1, int(seq_len)) + input_ids = torch.full((1, target), padding_token_id, dtype=torch.long) + attention_mask = torch.zeros((1, target), dtype=torch.long) + keep = min(prompt_token_count, target) + input_ids[:, :keep] = prompt_input_ids[:, :keep] + attention_mask[:, :keep] = 1 + return PreparedInputs( + names=("input_ids", "attention_mask"), + tensors=(input_ids, attention_mask), + metadata={ + "prompt": prompt, + "prompt_token_count": prompt_token_count, + "target_token_count": target, + "padding_token_id": padding_token_id, + }, + ) + + +def _add_multimodal_generation_headroom( + prepared: PreparedInputs, + *, + tokenizer: object | None, + max_new_tokens: int, + min_context_tokens: int | None = None, +) -> PreparedInputs: + requested = max(0, int(max_new_tokens)) + if "input_ids" not in prepared.names: + return prepared + + tensor_by_name = dict(zip(prepared.names, prepared.tensors, strict=True)) + input_ids = tensor_by_name.get("input_ids") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: + return prepared + + prompt_token_count = int(input_ids.shape[1]) + min_context = max(0, int(min_context_tokens or 0)) + target_token_count = max(prompt_token_count + requested, min_context) + if target_token_count <= prompt_token_count: + return prepared + padding_token_id = _resolve_graph_safe_text_padding_token_id(tokenizer, input_ids) + + padded_tensors: list[torch.Tensor] = [] + for name, tensor in zip(prepared.names, prepared.tensors, strict=True): + if ( + name == "position_ids" + and tensor.ndim == 3 + and int(tensor.shape[1]) == 1 + and int(tensor.shape[2]) == prompt_token_count + ): + padded = torch.zeros( + (int(tensor.shape[0]), 1, target_token_count), + dtype=tensor.dtype, + device=tensor.device, + ) + padded[:, :, :prompt_token_count] = tensor + padded_tensors.append(padded) + continue + if ( + name not in {"input_ids", "attention_mask", "token_type_ids", "mm_token_type_ids"} + or tensor.ndim != 2 + or int(tensor.shape[0]) != 1 + or int(tensor.shape[1]) != prompt_token_count + ): + padded_tensors.append(tensor) + continue + pad_value = int(padding_token_id) if name == "input_ids" else 0 + padded = torch.full( + (1, target_token_count), + pad_value, + dtype=tensor.dtype, + device=tensor.device, + ) + padded[:, :prompt_token_count] = tensor + padded_tensors.append(padded) + + metadata = dict(prepared.metadata) + metadata["prompt_token_count"] = prompt_token_count + metadata["target_token_count"] = target_token_count + metadata["max_new_tokens"] = max(requested, target_token_count - prompt_token_count) + metadata["padding_token_id"] = int(padding_token_id) + input_shapes = dict(metadata.get("input_shapes") or {}) + for name, tensor in zip(prepared.names, padded_tensors, strict=True): + input_shapes[name] = [int(dim) for dim in tensor.shape] + metadata["input_shapes"] = input_shapes + return PreparedInputs( + names=prepared.names, + tensors=tuple(padded_tensors), + metadata=metadata, + ) + + +def _multimodal_context_token_floor(model_type: str = "") -> int: + default_context = multimodal_context_tokens_for_model_type( + model_type, + _DEFAULT_MULTIMODAL_CONTEXT_TOKENS, + ) + raw = os.environ.get( + "CACTUS_TRANSPILER_MULTIMODAL_CONTEXT_TOKENS", + str(default_context), + ) + try: + return max(0, int(raw)) + except (TypeError, ValueError): + return default_context + + +_LFM2_VL_MULTIMODAL_INPUT_ORDER = ( + "input_ids", + "attention_mask", + "pixel_values", + "spatial_shapes", + "pixel_attention_mask", +) + +_QWEN3_5_MULTIMODAL_INPUT_ORDER = ( + "input_ids", + "attention_mask", + "position_ids", + "pixel_values", + "image_grid_thw", +) + + +def _load_image_inputs(image_files: tuple[str, ...]) -> list[object]: + if not image_files: + return [] + try: + from PIL import Image # type: ignore + except Exception as exc: # pragma: no cover + raise RuntimeError(f"Pillow is required for --image-file: {exc}") from exc + + images: list[object] = [] + for image_file in image_files: + path = Path(image_file).resolve() + if not path.exists(): + raise RuntimeError(f"image_file does not exist: {path}") + with Image.open(path) as image: + images.append(resize_static_image(image.convert("RGB")).copy()) + return images + + +def _prepare_gemma4_multimodal_inputs( + processor: object | None, + *, + prompt: str, + image_files: tuple[str, ...], + audio_file: str | None, + torch_dtype: torch.dtype, + system_prompt: str = "", + enable_thinking_if_supported: bool = False, + use_gemma4_chat_template: bool = False, +) -> PreparedInputs: + if processor is None: + raise RuntimeError("multimodal Gemma4 transpile requires an AutoProcessor with image and audio support") + shared = _shared_prepare_gemma4_multimodal_inputs( + processor, + prompt=prompt, + image_files=image_files, + audio_file=audio_file, + torch_dtype=torch_dtype, + system_prompt=system_prompt, + enable_thinking_if_supported=enable_thinking_if_supported, + use_gemma4_chat_template=use_gemma4_chat_template, + ) + return PreparedInputs( + names=shared.names, + tensors=shared.tensors, + metadata=shared.metadata, + ) + + +def _prepare_lfm2_vl_multimodal_inputs( + processor: object | None, + *, + prompt: str, + image_files: tuple[str, ...], + torch_dtype: torch.dtype, + system_prompt: str = "", + enable_thinking_if_supported: bool = False, +) -> PreparedInputs: + if processor is None: + raise RuntimeError("LFM2-VL multimodal transpile requires an AutoProcessor with image support") + images = _load_image_inputs(image_files) + if not images: + raise RuntimeError("LFM2-VL multimodal transpile requires at least one --image-file") + + user_content: list[dict[str, object]] = [{"type": "image", "image": image} for image in images] + user_content.append({"type": "text", "text": prompt.strip()}) + + messages: list[dict[str, object]] = [] + normalized_system = system_prompt.strip() + if normalized_system: + messages.append({"role": "system", "content": normalized_system}) + messages.append({"role": "user", "content": user_content}) + + batch: Mapping[str, object] + apply_chat_template = getattr(processor, "apply_chat_template", None) + thinking_prefix = "\n" if enable_thinking_if_supported else "\n\n\n\n" + image_placeholders = "".join("<|vision_start|><|image_pad|><|vision_end|>" for _ in images) + fallback_text = ( + f"<|im_start|>user\n{image_placeholders}{prompt.strip()}<|im_end|>\n" + f"<|im_start|>assistant\n{thinking_prefix}" + ) + if callable(apply_chat_template): + try: + batch = apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + except ValueError as exc: + if "chat template" not in str(exc).lower(): + raise + batch = processor(text=fallback_text, images=images, return_tensors="pt") + else: + batch = processor(text=fallback_text, images=images, return_tensors="pt") + + tensors: list[torch.Tensor] = [] + names: list[str] = [] + input_shapes: dict[str, list[int]] = {} + for key in _LFM2_VL_MULTIMODAL_INPUT_ORDER: + value = batch.get(key) if hasattr(batch, "get") else None + if not isinstance(value, torch.Tensor): + raise RuntimeError(f"LFM2-VL processor did not return required tensor input: {key}") + if torch.is_floating_point(value): + value = value.to(dtype=torch_dtype) + elif key == "pixel_attention_mask": + value = value.to(dtype=torch.int64) + else: + value = value.to(dtype=torch.long) + names.append(key) + tensors.append(value) + input_shapes[key] = [int(dim) for dim in value.shape] + + return PreparedInputs( + names=tuple(names), + tensors=tuple(tensors), + metadata={ + "prompt": prompt, + "system_prompt": system_prompt, + "image_files": [str(Path(path).resolve()) for path in image_files], + "input_shapes": input_shapes, + "enable_thinking": bool(enable_thinking_if_supported), + }, + ) + + +def _compute_qwen3_5_position_ids( + input_ids: torch.Tensor, + mm_token_type_ids: torch.Tensor, + image_grid_thw: torch.Tensor, + attention_mask: torch.Tensor | None, + *, + spatial_merge_size: int = 2, +) -> torch.Tensor: + position_ids = torch.zeros( + 3, + int(input_ids.shape[0]), + int(input_ids.shape[1]), + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_grid_iter = iter(image_grid_thw) + for batch_idx in range(int(input_ids.shape[0])): + token_types = mm_token_type_ids[batch_idx] + valid_mask = attention_mask[batch_idx].bool() if attention_mask is not None else None + if valid_mask is not None: + token_types = token_types[valid_mask] + + current_pos = 0 + position_parts: list[torch.Tensor] = [] + for modality_type, group in itertools.groupby(enumerate(token_types.tolist()), lambda item: int(item[1])): + group_items = list(group) + token_count = group_items[-1][0] - group_items[0][0] + 1 + if int(modality_type) == 0: + part = torch.arange(token_count, dtype=input_ids.dtype, device=input_ids.device) + position_parts.append(part.view(1, -1).expand(3, -1) + current_pos) + current_pos += token_count + continue + + grid_thw = next(image_grid_iter) + grid_t = int(grid_thw[0].item()) + grid_h = int(grid_thw[1].item()) + grid_w = int(grid_thw[2].item()) + llm_grid_t = grid_t + llm_grid_h = grid_h // int(spatial_merge_size) + llm_grid_w = grid_w // int(spatial_merge_size) + image_seq_length = llm_grid_t * llm_grid_h * llm_grid_w + position_width = torch.arange( + current_pos, + current_pos + llm_grid_w, + dtype=input_ids.dtype, + device=input_ids.device, + ).repeat(llm_grid_h * llm_grid_t) + position_height = torch.arange( + current_pos, + current_pos + llm_grid_h, + dtype=input_ids.dtype, + device=input_ids.device, + ).repeat_interleave(llm_grid_w * llm_grid_t) + position_temporal = torch.full( + (image_seq_length,), + current_pos, + dtype=input_ids.dtype, + device=input_ids.device, + ) + position_parts.append(torch.stack([position_temporal, position_height, position_width], dim=0)) + current_pos += max(grid_h, grid_w) // int(spatial_merge_size) + + if not position_parts: + continue + positions = torch.cat(position_parts, dim=1) + if valid_mask is not None: + position_ids[:, batch_idx, valid_mask] = positions.to(position_ids.device) + else: + position_ids[:, batch_idx, : positions.shape[1]] = positions.to(position_ids.device) + return position_ids + + +def _prepare_qwen3_5_multimodal_inputs( + processor: object | None, + *, + prompt: str, + image_files: tuple[str, ...], + torch_dtype: torch.dtype, + system_prompt: str = "", + enable_thinking_if_supported: bool = False, +) -> PreparedInputs: + if processor is None: + raise RuntimeError("Qwen3.5 multimodal transpile requires an AutoProcessor with image support") + images = _load_image_inputs(image_files) + if not images: + raise RuntimeError("Qwen3.5 multimodal transpile requires at least one --image-file") + + user_content: list[dict[str, object]] = [{"type": "image", "image": image} for image in images] + user_content.append({"type": "text", "text": prompt.strip()}) + + messages: list[dict[str, object]] = [] + normalized_system = system_prompt.strip() + if normalized_system: + messages.append({"role": "system", "content": normalized_system}) + messages.append({"role": "user", "content": user_content}) + + apply_chat_template = getattr(processor, "apply_chat_template", None) + if callable(apply_chat_template): + batch = apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + else: + batch = processor(text=prompt, images=images, return_tensors="pt") + + input_ids = batch.get("input_ids") if hasattr(batch, "get") else None + attention_mask = batch.get("attention_mask") if hasattr(batch, "get") else None + mm_token_type_ids = batch.get("mm_token_type_ids") if hasattr(batch, "get") else None + image_grid_thw = batch.get("image_grid_thw") if hasattr(batch, "get") else None + if not all(isinstance(value, torch.Tensor) for value in (input_ids, attention_mask, mm_token_type_ids, image_grid_thw)): + raise RuntimeError("Qwen3.5 processor did not return required multimodal text/grid tensors") + spatial_merge_size = int(getattr(getattr(processor, "image_processor", None), "merge_size", 2) or 2) + position_ids = _compute_qwen3_5_position_ids( + input_ids.to(dtype=torch.long), + mm_token_type_ids.to(dtype=torch.long), + image_grid_thw.to(dtype=torch.long), + attention_mask.to(dtype=torch.long), + spatial_merge_size=spatial_merge_size, + ) + batch["position_ids"] = position_ids + + tensors: list[torch.Tensor] = [] + names: list[str] = [] + input_shapes: dict[str, list[int]] = {} + for key in _QWEN3_5_MULTIMODAL_INPUT_ORDER: + value = batch.get(key) if hasattr(batch, "get") else None + if not isinstance(value, torch.Tensor): + raise RuntimeError(f"Qwen3.5 processor did not return required tensor input: {key}") + if torch.is_floating_point(value): + value = value.to(dtype=torch_dtype) + else: + value = value.to(dtype=torch.long) + names.append(key) + tensors.append(value) + input_shapes[key] = [int(dim) for dim in value.shape] + + return PreparedInputs( + names=tuple(names), + tensors=tuple(tensors), + metadata={ + "prompt": prompt, + "system_prompt": system_prompt, + "image_files": [str(Path(path).resolve()) for path in image_files], + "input_shapes": input_shapes, + "enable_thinking": bool(enable_thinking_if_supported), + "spatial_merge_size": spatial_merge_size, + }, + ) + + +def _load_model_source(model_id: str, *, local_files_only: bool, require_weights: bool = True) -> str: + local_snapshot = _resolve_local_snapshot(model_id) + if local_snapshot and _snapshot_has_model_weights(local_snapshot): + return local_snapshot + if local_snapshot and not require_weights: + return local_snapshot + if local_snapshot and local_files_only: + raise RuntimeError( + f"Found local snapshot for {model_id}, but it is incomplete and has no model weights:\n" + f" {local_snapshot}\n" + "\n" + "Re-run without --local-files-only to let transformers download the missing weights." + ) + return model_id + + +def _load_transformers_bundle( + *, + model_id: str, + task: str, + torch_dtype: torch.dtype, + token: str | None, + trust_remote_code: bool, + local_files_only: bool, + low_memory_load: bool, +): + config = _load_config_json(model_id) + config_model_type = str(config.get("model_type", "") or "").lower() + if not config_model_type: + normalized_model_id = model_id.lower().replace("_", "-") + if "qwen3-vl" in normalized_model_id: + config_model_type = "qwen3_vl" + elif "lfm2.5-vl" in normalized_model_id or "lfm2-vl" in normalized_model_id: + config_model_type = "lfm2_vl" + if config_model_type: + config = {"model_type": config_model_type} + external_transformers_site_packages = _ensure_transformers_supports_model_type(config_model_type) + patch_note = _patch_transformers_torchvision_probe() + if patch_note: + print(f"note={patch_note}") + flex_patch_note = _patch_torch_flex_attention_compat() + if flex_patch_note: + print(f"note={flex_patch_note}") + if external_transformers_site_packages: + print(f"note=using external transformers install for {config_model_type}: {external_transformers_site_packages}") + + try: + from transformers import AutoFeatureExtractor # type: ignore + from transformers import AutoConfig # type: ignore + from transformers import AutoModel # type: ignore + from transformers import AutoModelForCTC # type: ignore + from transformers import AutoModelForCausalLM # type: ignore + from transformers import AutoModelForImageTextToText # type: ignore + from transformers import AutoModelForSeq2SeqLM # type: ignore + from transformers import AutoModelForSpeechSeq2Seq # type: ignore + from transformers import AutoProcessor # type: ignore + from transformers import AutoTokenizer # type: ignore + except Exception as exc: # pragma: no cover + raise RuntimeError(f"transformers is not available: {exc}") from exc + + model_source = _load_model_source( + model_id, + local_files_only=local_files_only, + require_weights=not low_memory_load, + ) + source_candidates = [] + for candidate in (model_source, model_id): + if candidate not in source_candidates: + source_candidates.append(candidate) + common_kwargs: dict[str, object] = { + "trust_remote_code": trust_remote_code, + "local_files_only": local_files_only, + } + if token: + common_kwargs["token"] = token + + hf_config = None + if low_memory_load: + if task == "tdt_transcription": + raise RuntimeError("--low-memory-load is not supported for TDT transcription models") + try: + hf_config = AutoConfig.from_pretrained(model_source, **common_kwargs) + except Exception as exc: + raise RuntimeError( + f"--low-memory-load requires loading the HF config without checkpoint weights: {exc}" + ) from exc + + def _from_config(loader, *, trust_remote_code_override: bool | None = None): + if hf_config is None: + raise RuntimeError("internal error: low-memory config was not loaded") + with torch.device("meta"): + try: + model = loader.from_config( + hf_config, + trust_remote_code=( + trust_remote_code + if trust_remote_code_override is None + else bool(trust_remote_code_override) + ), + ) + except TypeError: + model = loader.from_config(hf_config) + model.to(dtype=torch_dtype) + _materialize_meta_rope_buffers(model, dtype=torch_dtype) + _materialize_meta_scalar_buffers(model, dtype=torch_dtype) + return model.eval() + + if task == "tdt_transcription": + model = load_tdt_local_model(model_source, torch_dtype=torch_dtype).eval() + return model_source, None, model, config + + if task == "text_embedding": + # nomic-bert weights are produced from the model's remote modeling code + # (fused Wqkv / experts.mlp.w1,w2); the transformers-native nomic_bert has a + # different architecture, so the remote code is required to match the bundle. + remote_kwargs = {**common_kwargs, "trust_remote_code": True} + tokenizer = None + tokenizer_errors: list[str] = [] + for source in source_candidates: + try: + tokenizer = AutoTokenizer.from_pretrained(source, **remote_kwargs) + break + except Exception as exc: + tokenizer_errors.append(f"{source}: {exc}") + if tokenizer is None: + raise RuntimeError( + f"Could not load tokenizer for {model_id}:\n" + + "\n".join(tokenizer_errors) + ) + if low_memory_load: + model = _from_config(AutoModel, trust_remote_code_override=True) + else: + model = AutoModel.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **remote_kwargs, + ).eval() + return model_source, tokenizer, model, config + + if task == "causal_lm_logits": + tokenizer = None + tokenizer_errors: list[str] = [] + for source in source_candidates: + try: + tokenizer = AutoTokenizer.from_pretrained(source, **common_kwargs) + break + except Exception as exc: + tokenizer_errors.append(f"{source}: {exc}") + if tokenizer is None: + raise RuntimeError( + f"Could not load tokenizer for {model_id}:\n" + + "\n".join(tokenizer_errors) + ) + if config_model_type in {"lfm2_vl", "qwen3_5", "qwen3_vl"}: + if low_memory_load: + model = _from_config(AutoModelForImageTextToText) + else: + model = AutoModelForImageTextToText.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + tie_note = _tie_lfm2_vl_lm_head_if_needed(model) + if tie_note: + print(f"note={tie_note}") + elif config_model_type in {"qwen3_5", "qwen3_vl"} and isinstance(config.get("text_config"), dict): + if low_memory_load: + model = _from_config(AutoModelForImageTextToText) + else: + model = AutoModelForImageTextToText.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + else: + if low_memory_load: + model = _from_config(AutoModelForCausalLM) + else: + model = AutoModelForCausalLM.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + return model_source, tokenizer, model, config + if task == "multimodal_causal_lm_logits": + processor = None + processor_errors: list[str] = [] + for source in source_candidates: + try: + processor = AutoProcessor.from_pretrained(source, **common_kwargs) + break + except Exception as exc: + processor_errors.append(f"{source}: {exc}") + if processor is None and config_model_type == "gemma4": + processor = _load_gemma4_processor_fallback( + source_candidates=source_candidates, + common_kwargs=common_kwargs, + ) + if processor is not None: + print("note=using manual gemma4 processor fallback") + if processor is None: + processor_config_hint = "" + if config_model_type == "gemma4": + processor_config_hint = ( + "\n" + "Gemma4 multimodal transpile needs a processor bundle, not just tokenizer/model weights.\n" + "Your local snapshot may be missing files such as `processor_config.json` or modality-specific\n" + "preprocessor configs. Use an official Gemma4 snapshot that includes the processor, or let\n" + "transformers download one by re-running without `--local-files-only`.\n" + ) + raise RuntimeError( + f"Could not load processor for {model_id}:\n" + + "\n".join(processor_errors) + + processor_config_hint + ) + + if config_model_type in {"lfm2_vl", "qwen3_5", "qwen3_vl"}: + if low_memory_load: + model = _from_config(AutoModelForImageTextToText) + else: + model = AutoModelForImageTextToText.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + tie_note = _tie_lfm2_vl_lm_head_if_needed(model) + if tie_note: + print(f"note={tie_note}") + else: + if low_memory_load: + model = _from_config(AutoModelForCausalLM) + else: + model = AutoModelForCausalLM.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + if config_model_type == "gemma4" and not low_memory_load: + repair_result = _repair_gemma4_checkpoint_weights(model, model_source) + if repair_result.get("applied"): + missing = repair_result.get("missing_keys", []) + unexpected = repair_result.get("unexpected_keys", []) + print( + "note=applied gemma4 checkpoint key remap" + f" missing_after={len(missing)} unexpected_after={len(unexpected)}" + ) + return model_source, processor, model, config + + processor = None + processor_errors: list[str] = [] + missing_optional_audio_dep: str | None = None + for source in source_candidates: + for loader in (AutoProcessor, AutoFeatureExtractor): + try: + processor = loader.from_pretrained(source, **common_kwargs) + break + except Exception as exc: + processor_errors.append(f"{loader.__name__}@{source}: {exc}") + if isinstance(exc, ImportError) and "requires the librosa library" in str(exc): + missing_optional_audio_dep = "librosa" + break + if processor is not None: + break + if missing_optional_audio_dep is not None: + break + if processor is None: + if missing_optional_audio_dep == "librosa": + print("note=falling back to built-in audio preprocessing because the HF feature extractor requires librosa") + else: + print("note=falling back to built-in audio preprocessing because no HF processor/feature extractor was available") + + if task == "ctc_logits": + model_loaders = (AutoModelForCTC, AutoModel) + elif task in {"encoder_hidden_states", "seq2seq_transcription"}: + model_loaders = (AutoModelForSpeechSeq2Seq, AutoModelForSeq2SeqLM, AutoModel) + else: + raise NotImplementedError(f"unsupported generic HF task: {task}") + + load_errors: list[str] = [] + for loader in model_loaders: + try: + if low_memory_load: + model = _from_config(loader) + else: + model = loader.from_pretrained( + model_source, + dtype=torch_dtype, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + return model_source, processor, model, config + except Exception as exc: + load_errors.append(f"{loader.__name__}: {exc}") + + raise RuntimeError( + f"Could not load model for task={task} from {model_source}.\n" + "\n".join(load_errors) + ) + + +def _load_optional_tokenizer( + *, + model_id: str, + model_source: str, + token: str | None, + trust_remote_code: bool, + local_files_only: bool, +): + try: + from transformers import AutoTokenizer # type: ignore + except Exception: + return None + + source_candidates = [] + for candidate in (model_source, model_id): + if candidate not in source_candidates: + source_candidates.append(candidate) + + common_kwargs: dict[str, object] = { + "trust_remote_code": trust_remote_code, + "local_files_only": local_files_only, + } + if token: + common_kwargs["token"] = token + + for source in source_candidates: + try: + return AutoTokenizer.from_pretrained(source, **common_kwargs) + except Exception: + continue + return None + + +def _ctc_greedy_decode_token_ids(logits: np.ndarray, *, blank_token_id: int | None) -> list[int]: + if logits.ndim != 3 or logits.shape[0] < 1: + raise ValueError(f"expected CTC logits with shape [batch, time, vocab], got {list(logits.shape)}") + + raw_ids = np.argmax(logits[0], axis=-1).tolist() + collapsed: list[int] = [] + previous: int | None = None + for token_id in raw_ids: + if token_id != previous: + collapsed.append(int(token_id)) + previous = int(token_id) + + if blank_token_id is None: + return collapsed + return [token_id for token_id in collapsed if int(token_id) != int(blank_token_id)] + + +def _decode_token_ids(tokenizer: object, token_ids: list[int]) -> str: + decode = getattr(tokenizer, "decode", None) + if not callable(decode): + raise TypeError(f"tokenizer does not expose decode(): {type(tokenizer).__name__}") + try: + return str(decode(token_ids, skip_special_tokens=True)) + except TypeError: + return str(decode(token_ids)) + + +def _count_weight_bindings(ir_graph: IRGraph) -> int: + count = 0 + for value in ir_graph.values.values(): + if isinstance(value.meta, dict) and isinstance(value.meta.get("path"), str): + count += 1 + return count + + +def _lower_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: + verify_ir(ir) + graph = Graph() + graph._transpile_materialized_constants = [] # type: ignore[attr-defined] + env: dict[str, Any] = {} + runtime_inputs = [] + bound_constants = [] + bound_constant_bindings = [] + bound_constant_value_ids: dict[int, str] = {} + + for value_id in ir.inputs: + value = ir.values[value_id] + tensor = _lower_input_value(graph, value) + env[value_id] = tensor + runtime_inputs.append(tensor) + + for value_id, const in ir.constants.items(): + value = ir.values[value_id] + binding = _lookup_weight_binding(value) + if binding is not None: + binding = ensure_binding_compatible(binding, source_tensor=const) + lowered_const = _lower_constant_value(graph, value, const, binding=binding) + env[value_id] = lowered_const + if hasattr(lowered_const, "g") and hasattr(lowered_const, "id"): + bound_constants.append(lowered_const) + bound_constant_value_ids[int(lowered_const.id)] = str(value_id) + if binding is not None: + bound_constant_bindings.append( + { + "node_id": int(lowered_const.id), + "value_id": str(value_id), + "path": binding.path, + "kind": binding.kind, + "source_name": binding.source_name, + } + ) + + for node_id in ir.order: + node = ir.nodes[node_id] + outputs = _lower_ir_node(graph, node, env, ir) + if len(outputs) != len(node.outputs): + raise ValueError( + f"node {node.id} produced {len(outputs)} outputs, expected {len(node.outputs)}" + ) + for output_id, tensor in zip(node.outputs, outputs): + env[output_id] = tensor + + outputs = [env[value_id] for value_id in ir.outputs] + seen_bound_constant_ids = {int(tensor.id) for tensor in bound_constants} + for tensor in getattr(graph, "_transpile_materialized_constants", []): + tensor_id = int(tensor.id) + if tensor_id in seen_bound_constant_ids: + continue + bound_constants.append(tensor) + seen_bound_constant_ids.add(tensor_id) + return TranspiledGraph( + graph=graph, + runtime_inputs=runtime_inputs, + bound_constants=bound_constants, + bound_constant_bindings=bound_constant_bindings, + bound_constant_value_ids=bound_constant_value_ids, + outputs=outputs, + ) + + +def main() -> int: + from cactus.models.needle import register_with_transformers + + register_with_transformers() + parser = argparse.ArgumentParser( + description=( + "Load a Hugging Face model, canonicalize it into a generic transpile task, " + "capture it with the Cactus transpiler, lower it to a Cactus Graph, and " + "optionally save artifacts or run the lowered graph." + ) + ) + parser.add_argument( + "--model-id", + required=True, + help="Hugging Face model id or local snapshot path.", + ) + parser.add_argument( + "--task", + default="auto", + choices=( + "auto", + "causal_lm_logits", + "multimodal_causal_lm_logits", + "ctc_logits", + "encoder_hidden_states", + "seq2seq_transcription", + "tdt_transcription", + "text_embedding", + ), + help="Transpile task. Use auto to infer from config/model id.", + ) + parser.add_argument( + "--prompt", + default=_DEFAULT_CAUSAL_PROMPT, + help="Prompt used for causal_lm_logits or multimodal_causal_lm_logits when --input-ids is not set.", + ) + parser.add_argument( + "--system-prompt", + default="", + help="Optional system prompt used for Gemma4 multimodal chat-style prompt construction.", + ) + parser.add_argument( + "--enable-thinking", + action="store_true", + help="Enable model-specific thinking markers when the multimodal prompt format supports them.", + ) + parser.add_argument( + "--input-ids", + default="", + help="Optional comma-separated token ids for causal_lm_logits.", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=32, + help=( + "For causal_lm_logits bundles, preallocate prompt-plus-generation context " + "so the saved bundle can do greedy autoregressive decoding." + ), + ) + parser.add_argument( + "--cache-context-length", + default="auto", + help="KV cache context length for cached decode graphs. Use auto to read the model config.", + ) + parser.add_argument( + "--audio-file", + default="", + help="Path to a WAV file for audio or multimodal tasks.", + ) + parser.add_argument( + "--image-file", + action="append", + default=[], + help="Path to an image file for multimodal tasks. Repeat to pass multiple images.", + ) + parser.add_argument( + "--torch-dtype", + default="float16", + help="Torch dtype for model loading: float16, float32, or bfloat16.", + ) + parser.add_argument( + "--token", + default=os.environ.get("HF_TOKEN"), + help="Optional Hugging Face token. Defaults to HF_TOKEN.", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Pass trust_remote_code=True to transformers loaders.", + ) + parser.add_argument( + "--local-files-only", + action="store_true", + help="Require the model/processor to already exist locally.", + ) + parser.add_argument( + "--low-memory-load", + action="store_true", + help=( + "Instantiate the HF model on the meta device for graph capture instead of " + "loading checkpoint tensors into RAM. Requires converted CQ weights and " + "only supports bundle generation, not execution." + ), + ) + parser.add_argument( + "--weights-dir", + default="", + help="Converted Cactus CQ weights directory for mmap weight binding.", + ) + parser.add_argument( + "--allow-unconverted-weights", + action="store_true", + help="Debug only: allow transpiling without converted Cactus CQ weights.", + ) + parser.add_argument( + "--artifact-dir", + default="", + help="Directory where raw_ir.json, optimized_ir.json, graph.cactus, component graphs, and result.json are saved. Defaults to ./transpiled//.", + ) + parser.add_argument( + "--graph-filename", + default="graph.cactus", + help="Filename to use for Graph.save() inside --artifact-dir.", + ) + parser.add_argument( + "--component-pipeline", + default="auto", + choices=("auto", "on", "off"), + help="Use a split component pipeline when the model adapter supports it.", + ) + parser.add_argument( + "--components", + default=None, + help=( + "Optional comma-separated component subset for component-pipeline models " + "(for example: vision_encoder,audio_encoder,lm_encoder,decoder)." + ), + ) + parser.add_argument("--no-fuse-gated-deltanet", action="store_true") + parser.add_argument("--no-fuse-rms-norm", action="store_true") + parser.add_argument("--no-fuse-rope", action="store_true") + parser.add_argument("--no-fuse-attention", action="store_true") + parser.add_argument("--no-fuse-attention-block", action="store_true") + parser.add_argument("--no-fuse-add-clipped", action="store_true") + parser.add_argument( + "--skip-execute", + action="store_true", + help="Stop after lowering instead of running the transpiled graph.", + ) + parser.add_argument( + "--skip-reference-compare", + action="store_true", + help="Run the transpiled graph but skip the follow-up PyTorch reference pass.", + ) + args = parser.parse_args() + + validated_weights_dir = _validate_weights_dir(args.weights_dir.strip() or None, model_id=args.model_id) + if validated_weights_dir is None and not args.allow_unconverted_weights: + raise RuntimeError( + "Building the runtime graph requires converted Cactus CQ weights.\n" + "\n" + "Run:\n" + f" cactus convert {args.model_id} --bits 4\n" + "\n" + "For compiler-only debugging, pass --allow-unconverted-weights." + ) + if args.low_memory_load: + if validated_weights_dir is None: + raise RuntimeError("--low-memory-load requires --weights-dir with converted Cactus CQ weights") + if not args.skip_execute: + print("note=low_memory_load_forces_skip_execute=true") + args.skip_execute = True + args.skip_reference_compare = True + + image_files = tuple(str(path) for path in args.image_file if str(path).strip()) + if args.task == "auto": + inferred_task = _infer_task_from_config(args.model_id) + config_for_auto = _load_config_json(args.model_id) + plan_for_auto = infer_component_plan_from_config(config_for_auto, model_id=args.model_id) + has_multimodal_config = ( + plan_for_auto is not None + and plan_for_auto.task == "multimodal_causal_lm_logits" + ) + if image_files or (args.audio_file and has_multimodal_config and inferred_task == "causal_lm_logits"): + task = plan_for_auto.task if plan_for_auto is not None else "multimodal_causal_lm_logits" + else: + task = inferred_task + else: + task = args.task + torch_dtype = _parse_dtype(args.torch_dtype) + weights_dir = str(validated_weights_dir) if validated_weights_dir is not None else None + artifact_dir = Path(args.artifact_dir).resolve() if args.artifact_dir else _default_artifact_dir_for_model(args.model_id) + + model_source, processor_or_tokenizer, model, model_config = _load_transformers_bundle( + model_id=args.model_id, + task=task, + torch_dtype=torch_dtype, + token=args.token, + trust_remote_code=args.trust_remote_code, + local_files_only=args.local_files_only, + low_memory_load=bool(args.low_memory_load), + ) + preprocessor_config = _load_optional_json(model_source, "preprocessor_config.json") + if not preprocessor_config: + preprocessor_config = _load_optional_json(args.model_id, "preprocessor_config.json") + auxiliary_tokenizer = None + if task == "ctc_logits": + auxiliary_tokenizer = _load_optional_tokenizer( + model_id=args.model_id, + model_source=model_source, + token=args.token, + trust_remote_code=args.trust_remote_code, + local_files_only=args.local_files_only, + ) + + canonical = None + if task == "causal_lm_logits": + prepared = _prepare_text_inputs( + processor_or_tokenizer, + prompt=args.prompt, + input_ids_text=args.input_ids.strip() or None, + max_new_tokens=int(args.max_new_tokens), + enable_thinking_if_supported=args.enable_thinking, + ) + canonical = canonicalize_model_interface( + model, + task=task, + input_names=prepared.names, + weights_dir=weights_dir, + inputs_metadata=prepared.metadata, + ) + elif task == "text_embedding": + prepared = _prepare_text_embedding_inputs( + processor_or_tokenizer, + prompt=args.prompt, + input_ids_text=args.input_ids.strip() or None, + ) + canonical = canonicalize_model_interface( + model, + task=task, + input_names=prepared.names, + weights_dir=weights_dir, + inputs_metadata=prepared.metadata, + ) + elif task == "multimodal_causal_lm_logits": + config_model_type = str(model_config.get("model_type", "") or "").lower() + if config_model_type == "lfm2_vl": + prepared = _prepare_lfm2_vl_multimodal_inputs( + processor_or_tokenizer, + prompt=args.prompt, + image_files=image_files, + torch_dtype=torch_dtype, + system_prompt=args.system_prompt, + enable_thinking_if_supported=args.enable_thinking, + ) + elif config_model_type in {"qwen3_5", "qwen3_vl"}: + prepared = _prepare_qwen3_5_multimodal_inputs( + processor_or_tokenizer, + prompt=args.prompt, + image_files=image_files, + torch_dtype=torch_dtype, + system_prompt=args.system_prompt, + enable_thinking_if_supported=args.enable_thinking, + ) + else: + prepared = _prepare_gemma4_multimodal_inputs( + processor_or_tokenizer, + prompt=args.prompt, + image_files=image_files, + audio_file=args.audio_file.strip() or None, + torch_dtype=torch_dtype, + system_prompt=args.system_prompt, + enable_thinking_if_supported=args.enable_thinking, + use_gemma4_chat_template=True, + ) + prepared = _add_multimodal_generation_headroom( + prepared, + tokenizer=getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer), + max_new_tokens=int(args.max_new_tokens), + min_context_tokens=_multimodal_context_token_floor(config_model_type), + ) + canonical = canonicalize_model_interface( + model, + task=task, + input_names=prepared.names, + weights_dir=weights_dir, + ) + prime_static_features = getattr(canonical.module, "prime_static_multimodal_features", None) + if callable(prime_static_features) and not args.low_memory_load: + prime_static_features(*prepared.tensors) + elif task == "tdt_transcription": + if not args.audio_file: + raise RuntimeError(f"--audio-file is required for task={task}") + prepared = _prepare_audio_inputs( + processor_or_tokenizer, + input_names=("input_features",), + config=model_config, + preprocessor_config=preprocessor_config, + model=model, + task=task, + audio_file=args.audio_file, + torch_dtype=torch_dtype, + ) + elif task == "seq2seq_transcription": + if not args.audio_file: + raise RuntimeError(f"--audio-file is required for task={task}") + prepared = _prepare_audio_inputs( + processor_or_tokenizer, + input_names=("input_features",), + config=model_config, + preprocessor_config=preprocessor_config, + model=model, + task=task, + audio_file=args.audio_file, + torch_dtype=torch_dtype, + ) + if getattr(model_config, "get", None) is not None or getattr(model, "config", None) is not None: + prepared = _augment_whisper_seq2seq_metadata( + prepared, + tokenizer_or_processor=processor_or_tokenizer, + model=model, + prompt=args.prompt, + max_new_tokens=int(args.max_new_tokens), + ) + else: + if not args.audio_file: + raise RuntimeError(f"--audio-file is required for task={task}") + canonical = canonicalize_model_interface(model, task=task) + prepared = _prepare_audio_inputs( + processor_or_tokenizer, + input_names=canonical.input_names, + config=model_config, + preprocessor_config=preprocessor_config, + model=model, + task=task, + audio_file=args.audio_file, + torch_dtype=torch_dtype, + ) + canonical = canonicalize_model_interface( + model, + task=task, + input_names=prepared.names, + ) + defer_low_memory_meta_inputs = ( + bool(args.low_memory_load) + and task == "multimodal_causal_lm_logits" + and config_model_type == "gemma4" + ) + if args.low_memory_load and not defer_low_memory_meta_inputs: + prepared = _prepared_inputs_for_meta_capture(prepared) + requested_components = None + if args.components: + requested_components = tuple( + component.strip() + for component in str(args.components).split(",") + if component.strip() + ) + plan_derived_components = False + if requested_components is None and task == "causal_lm_logits": + inferred_plan = infer_component_plan_from_config( + _load_config_json(args.model_id), model_id=args.model_id + ) + if ( + inferred_plan is not None + and inferred_plan.task == task + and inferred_plan.force_component_pipeline + and inferred_plan.components + ): + requested_components = tuple(inferred_plan.components) + plan_derived_components = True + spec_kwargs = dict( + task=task, + named_tensors=_named_tensor_store(prepared), + weights_dir=weights_dir, + inputs_metadata=prepared.metadata, + cache_context_length=args.cache_context_length, + ) + try: + component_specs = build_component_module_specs( + model, components=requested_components, **spec_kwargs + ) + except UnsupportedComponentsError as exc: + if not plan_derived_components: + raise + print(f"warning: dropping inferred component plan ({exc}); using builder defaults", file=sys.stderr) + component_specs = build_component_module_specs(model, components=None, **spec_kwargs) + use_component_pipeline = False + if args.component_pipeline == "on": + if component_specs is None: + raise RuntimeError( + f"--component-pipeline=on was requested, but {type(model).__name__} does not expose component specs for task={task}" + ) + use_component_pipeline = True + elif args.component_pipeline == "auto" and component_specs is not None: + use_component_pipeline = True + + if use_component_pipeline: + family = canonical.family if canonical is not None else getattr(model, "family", type(model).__name__) + return _run_component_pipeline_transpile( + args=args, + task=task, + family=str(family), + model_source=model_source, + model=model, + prepared=prepared, + component_specs=component_specs, + fusion_config=FusionConfig( + enable_gated_deltanet=not args.no_fuse_gated_deltanet, + enable_rms_norm=not args.no_fuse_rms_norm, + enable_rope=not args.no_fuse_rope, + enable_attention=not args.no_fuse_attention, + enable_attention_block=not args.no_fuse_attention_block, + enable_add_clipped=not args.no_fuse_add_clipped, + ), + weights_dir=weights_dir, + artifact_dir=artifact_dir, + processor_or_tokenizer=processor_or_tokenizer, + canonical=canonical, + ) + + if canonical is None: + raise RuntimeError(f"task={task} requires a component pipeline but none was selected") + if args.low_memory_load and defer_low_memory_meta_inputs: + prepared = _prepared_inputs_for_meta_capture(prepared) + wrapper = TranspileWrapper(canonical.module, weights_dir=weights_dir).eval() + + print(f"model_id={args.model_id}") + print(f"model_source={model_source}") + print(f"task={task}") + print(f"adapter_family={canonical.family}") + print(f"adapter_module={type(canonical.module).__name__}") + print(f"input_names={','.join(prepared.names)}") + for name, tensor in zip(prepared.names, prepared.tensors): + print(f"input_{name}_shape={list(tensor.shape)}") + if weights_dir: + print(f"weights_dir={weights_dir}") + + print("capture_begin=true", flush=True) + prepare_cpu_float32_capture = getattr(canonical.module, "prepare_cpu_float32_capture", None) + restore_cpu_float32_capture = getattr(canonical.module, "restore_cpu_float32_capture", None) + if callable(prepare_cpu_float32_capture): + prepare_cpu_float32_capture() + try: + captured = capture_model(wrapper, prepared.tensors) + finally: + if callable(restore_cpu_float32_capture): + restore_cpu_float32_capture() + print("capture_done=true", flush=True) + captured.ir_graph.meta.setdefault("task", task) + captured.ir_graph.meta.setdefault("adapter_family", canonical.family) + raw_ir_graph = copy.deepcopy(captured.ir_graph) + + fusion_config = FusionConfig( + enable_gated_deltanet=not args.no_fuse_gated_deltanet, + enable_rms_norm=not args.no_fuse_rms_norm, + enable_rope=not args.no_fuse_rope, + enable_attention=not args.no_fuse_attention, + enable_attention_block=not args.no_fuse_attention_block, + enable_add_clipped=not args.no_fuse_add_clipped, + ) + + print("canonicalize_begin=true", flush=True) + canonicalize_exported_graph(captured.ir_graph) + print("canonicalize_done=true", flush=True) + print("optimize_begin=true", flush=True) + optimize_graph(captured.ir_graph, config=fusion_config) + print("optimize_done=true", flush=True) + print("lower_begin=true", flush=True) + tg = _lower_preoptimized_ir(captured.ir_graph) + print("lower_done=true", flush=True) + + optimized_ir_graph = copy.deepcopy(captured.ir_graph) + binding_count = _count_weight_bindings(optimized_ir_graph) + op_counts = Counter(optimized_ir_graph.nodes[node_id].op for node_id in optimized_ir_graph.order) + raw_component_counts = summarize_ir_components(raw_ir_graph) + optimized_component_counts = summarize_ir_components(optimized_ir_graph) + raw_component_graphs = extract_component_subgraphs(raw_ir_graph) + optimized_component_graphs = extract_component_subgraphs(optimized_ir_graph) + transpiled_component_graphs: dict[str, TranspiledGraph] | None = None + + print(f"raw_ir_nodes={len(raw_ir_graph.order)}") + print(f"optimized_ir_nodes={len(optimized_ir_graph.order)}") + print(f"weight_bindings={binding_count}") + if raw_component_counts: + print( + "raw_components=" + + ",".join(f"{name}:{count}" for name, count in sorted(raw_component_counts.items())) + ) + if optimized_component_counts: + print( + "optimized_components=" + + ",".join(f"{name}:{count}" for name, count in sorted(optimized_component_counts.items())) + ) + print( + "ops=" + f"attention:{op_counts.get('attention', 0)} " + f"conv1d:{op_counts.get('conv1d', 0)} " + f"conv2d:{op_counts.get('conv2d', 0)} " + f"batch_norm:{op_counts.get('batch_norm', 0)} " + f"layer_norm:{op_counts.get('layer_norm', 0)} " + f"rms_norm:{op_counts.get('rms_norm', 0)} " + f"rope:{op_counts.get('rope', 0)} " + f"linear:{op_counts.get('linear', 0)}" + ) + + if weights_dir and binding_count == 0: + raise RuntimeError( + f"No weight bindings were resolved from {weights_dir}\n" + "\n" + "The weights folder exists, but none of the captured constants matched entries in weights_manifest.json.\n" + "\n" + f"Recommended fix:\n" + f" cactus convert {args.model_id} {weights_dir}\n" + ) + + if artifact_dir is not None: + artifact_dir.mkdir(parents=True, exist_ok=True) + transpiled_component_graphs = { + component: _lower_preoptimized_ir(copy.deepcopy(component_graph)) + for component, component_graph in optimized_component_graphs.items() + } + _write_json( + artifact_dir / "raw_ir.json", + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": canonical.family, + "inputs": _serialize_json_compatible(prepared.metadata), + "graph": _graph_to_dict(raw_ir_graph), + }, + ) + _write_json( + artifact_dir / "optimized_ir.json", + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": canonical.family, + "inputs": _serialize_json_compatible(prepared.metadata), + "graph": _graph_to_dict(optimized_ir_graph), + }, + ) + for component, component_graph in raw_component_graphs.items(): + _write_json( + artifact_dir / _component_artifact_name("raw_ir", component), + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": canonical.family, + "component": component, + "inputs": _serialize_json_compatible(prepared.metadata), + "graph": _graph_to_dict(component_graph), + }, + ) + for component, component_graph in optimized_component_graphs.items(): + _write_json( + artifact_dir / _component_artifact_name("optimized_ir", component), + { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": canonical.family, + "component": component, + "inputs": _serialize_json_compatible(prepared.metadata), + "graph": _graph_to_dict(component_graph), + }, + ) + print(f"saved_raw_ir={artifact_dir / 'raw_ir.json'}") + print(f"saved_optimized_ir={artifact_dir / 'optimized_ir.json'}") + for component in raw_component_graphs: + print(f"saved_raw_component_ir_{component}={artifact_dir / _component_artifact_name('raw_ir', component)}") + for component in optimized_component_graphs: + print( + f"saved_optimized_component_ir_{component}=" + f"{artifact_dir / _component_artifact_name('optimized_ir', component)}" + ) + + component_manifest_path = _write_component_bundle( + artifact_dir=artifact_dir, + model_id=args.model_id, + model_source=model_source, + task=task, + family=canonical.family, + inputs_metadata=prepared.metadata, + raw_component_graphs=raw_component_graphs, + optimized_component_graphs=optimized_component_graphs, + transpiled_component_graphs=transpiled_component_graphs, + graph_filename=args.graph_filename, + ) + print(f"saved_component_bundle_manifest={component_manifest_path}") + for component in optimized_component_graphs: + print( + f"saved_component_graph_{component}=" + f"{artifact_dir / 'components' / component / args.graph_filename}" + ) + graph_path = artifact_dir / args.graph_filename + _embed_materialized_bound_constants(tg) + tg.graph.save(graph_path) + print(f"saved_graph={graph_path}") + graph_binding_manifest_path = _write_graph_binding_manifest( + artifact_dir=artifact_dir, + filename="graph_bindings.json", + model_id=args.model_id, + model_source=model_source, + task=task, + family=canonical.family, + inputs_metadata=prepared.metadata, + transpiled_graph=tg, + ) + print(f"saved_graph_bindings={graph_binding_manifest_path}") + if binding_count > 0: + print( + "note=saved graph structure without embedded mmap bindings; use the saved component/full-graph manifests " + "to rebind weights/embeddings from the weights folder when loading" + ) + + if args.skip_execute: + return 0 + + tg.set_inputs([tensor.cpu().numpy() for tensor in prepared.tensors]) + print("execute_begin=true") + transpiled_output = tg.execute()[0].numpy().astype(np.float32) + print("execute_done=true") + + if args.skip_reference_compare: + result_payload = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": canonical.family, + "inputs": _serialize_json_compatible(prepared.metadata), + "output_shape": list(transpiled_output.shape), + "raw_ir_nodes": len(raw_ir_graph.order), + "optimized_ir_nodes": len(optimized_ir_graph.order), + "weight_bindings": binding_count, + "reference_compare_skipped": True, + } + print(f"output_shape={list(transpiled_output.shape)}") + if task in {"causal_lm_logits", "multimodal_causal_lm_logits"}: + tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) + transpiled_next = int(np.argmax(transpiled_output[0, -1])) + print(f"transpiled_next_token_id={transpiled_next}") + result_payload["transpiled_next_token_id"] = transpiled_next + if hasattr(tokenizer_like, "decode"): + transpiled_token = tokenizer_like.decode([transpiled_next]) + print(f"transpiled_next_token={transpiled_token!r}") + result_payload["transpiled_next_token"] = transpiled_token + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + del model + del wrapper + gc.collect() + return 0 + + print("reference_begin=true") + with torch.no_grad(): + reference_output = wrapper(*prepared.tensors).detach().float().cpu().numpy() + print("reference_done=true") + + max_abs_diff = float(np.max(np.abs(reference_output - transpiled_output))) + mean_abs_diff = float(np.mean(np.abs(reference_output - transpiled_output))) + result_payload: dict[str, object] = { + "model_id": args.model_id, + "model_source": model_source, + "task": task, + "family": canonical.family, + "inputs": _serialize_json_compatible(prepared.metadata), + "output_shape": list(reference_output.shape), + "raw_ir_nodes": len(raw_ir_graph.order), + "optimized_ir_nodes": len(optimized_ir_graph.order), + "weight_bindings": binding_count, + "max_abs_diff": max_abs_diff, + "mean_abs_diff": mean_abs_diff, + } + + if task == "causal_lm_logits": + tokenizer = processor_or_tokenizer + hf_next = int(np.argmax(reference_output[0, -1])) + transpiled_next = int(np.argmax(transpiled_output[0, -1])) + print(f"hf_next_token_id={hf_next}") + print(f"transpiled_next_token_id={transpiled_next}") + print(f"logits_max_abs_diff={max_abs_diff:.6f}") + print(f"logits_mean_abs_diff={mean_abs_diff:.6f}") + print(f"hf_next_token={tokenizer.decode([hf_next])!r}") + print(f"transpiled_next_token={tokenizer.decode([transpiled_next])!r}") + result_payload.update( + { + "hf_next_token_id": hf_next, + "transpiled_next_token_id": transpiled_next, + "hf_next_token": tokenizer.decode([hf_next]), + "transpiled_next_token": tokenizer.decode([transpiled_next]), + } + ) + elif task == "ctc_logits": + blank_token_id = getattr(model.config, "pad_token_id", None) + if blank_token_id is None and auxiliary_tokenizer is not None: + blank_token_id = getattr(auxiliary_tokenizer, "pad_token_id", None) + print(f"output_shape={list(reference_output.shape)}") + print(f"output_max_abs_diff={max_abs_diff:.6f}") + print(f"output_mean_abs_diff={mean_abs_diff:.6f}") + if auxiliary_tokenizer is not None: + hf_token_ids = _ctc_greedy_decode_token_ids(reference_output, blank_token_id=blank_token_id) + transpiled_token_ids = _ctc_greedy_decode_token_ids(transpiled_output, blank_token_id=blank_token_id) + hf_transcript = _decode_token_ids(auxiliary_tokenizer, hf_token_ids) + transpiled_transcript = _decode_token_ids(auxiliary_tokenizer, transpiled_token_ids) + print(f"hf_transcript={hf_transcript!r}") + print(f"transpiled_transcript={transpiled_transcript!r}") + result_payload.update( + { + "blank_token_id": None if blank_token_id is None else int(blank_token_id), + "hf_transcript_token_ids": hf_token_ids, + "transpiled_transcript_token_ids": transpiled_token_ids, + "hf_transcript": hf_transcript, + "transpiled_transcript": transpiled_transcript, + } + ) + else: + print(f"output_shape={list(reference_output.shape)}") + print(f"output_max_abs_diff={max_abs_diff:.6f}") + print(f"output_mean_abs_diff={mean_abs_diff:.6f}") + + if artifact_dir is not None: + _write_json(artifact_dir / "result.json", result_payload) + print(f"saved_result={artifact_dir / 'result.json'}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/cactus/transpile/import_ir.py b/python/cactus/transpile/import_ir.py new file mode 100644 index 000000000..19cd05112 --- /dev/null +++ b/python/cactus/transpile/import_ir.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch +from torch.export.graph_signature import ConstantArgument +from torch.export.graph_signature import InputKind + +from cactus.transpile.capture_pytorch import format_target +from cactus.transpile.capture_pytorch import get_dtype +from cactus.transpile.capture_pytorch import get_shape +from cactus.transpile.capture_pytorch import resolve_attr +from cactus.transpile.component_partition import annotate_ir_components +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.importers import extract_literals +from cactus.transpile.importers import ImportContext +from cactus.transpile.importers import import_call_function +from cactus.transpile.importers import import_get_attr +from cactus.transpile.importers import import_output +from cactus.transpile.importers import import_placeholder +from cactus.transpile.importers import value_id +from cactus.transpile.normalize import dtype_to_ir +from cactus.transpile.import_semantics import apply_import_semantics +from cactus.transpile.weight_binding import resolve_transpile_weights_dir +from cactus.transpile.weight_binding import resolve_weight_binding + + +def import_captured_to_ir(captured: Any, *, strict: bool = True) -> IRGraph: + transpile_metadata = getattr(captured, "transpile_metadata", {}) or {} + graph_meta = transpile_metadata.get("graph", {}) + ir = IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[], constants={}, meta=dict(graph_meta)) + ctx = ImportContext(strict=strict, transpile_metadata=transpile_metadata) + weights_dir = resolve_transpile_weights_dir(ir.meta) + if weights_dir: + ir.meta["weights_dir"] = weights_dir + placeholder_specs: dict[str, Any] = {} + inline_counter = 0 + + def _try_register_weight_binding(value_id_str: str, target: str, value: Any) -> None: + if not isinstance(value, torch.Tensor): + return + binding = resolve_weight_binding( + weights_dir=weights_dir, + source_name=target, + ) + if binding is None: + return + bindings = ir.meta.setdefault("weight_bindings", {}) + if isinstance(bindings, dict): + bindings[value_id_str] = { + "path": binding.path, + "kind": binding.kind, + "source_name": binding.source_name, + } + if value_id_str in ir.values: + ir.values[value_id_str].meta.update( + { + "path": binding.path, + "kind": binding.kind, + "source_name": binding.source_name, + } + ) + + def _make_inlined_node(node: Any, *, inline_index: int) -> Any: + return SimpleNamespace( + name=f"inl_{inline_index}_{node.name}", + op=node.op, + target=node.target, + args=node.args, + kwargs=getattr(node, "kwargs", {}), + meta=getattr(node, "meta", {}), + ) + + def _resolve_target(target: str) -> Any: + candidates: list[str] = [] + parts = target.split(".") + for i in range(len(parts)): + candidate = ".".join(parts[i:]) + if candidate and candidate not in candidates: + candidates.append(candidate) + + exported_program = getattr(captured, "exported_program", None) + if exported_program is not None: + tensor_constants = getattr(exported_program, "tensor_constants", None) + if tensor_constants is not None: + for candidate in candidates: + if candidate in tensor_constants: + return tensor_constants[candidate] + constants = getattr(exported_program, "constants", None) + if constants is not None: + for candidate in candidates: + if candidate in constants: + return constants[candidate] + + for candidate in candidates: + try: + return resolve_attr(captured.graph_module, candidate) + except AttributeError: + pass + + for candidate in candidates: + if candidate in captured.state_dict: + return captured.state_dict[candidate] + + source_module = getattr(captured, "source_module", None) + if source_module is not None: + for candidate in candidates: + try: + return resolve_attr(source_module, candidate) + except AttributeError: + pass + + raise NotImplementedError(f"could not resolve get_attr target: {target}") from None + + def _collect_output_ids(arg: Any, local_alias: dict[str, str]) -> list[str]: + if hasattr(arg, "name") and hasattr(arg, "op"): + name = arg.name + if name in local_alias: + return [local_alias[name]] + return [ctx.resolve_value_id(arg)] + if isinstance(arg, (tuple, list)): + out: list[str] = [] + for item in arg: + out.extend(_collect_output_ids(item, local_alias)) + return out + return [] + + def _inline_wrap_with_set_grad_enabled(node: Any, shape: Any, dtype: Any) -> None: + nonlocal inline_counter + subgraph_ref = node.args[1] + if not hasattr(subgraph_ref, "op") or subgraph_ref.op != "get_attr": + raise NotImplementedError("wrap_with_set_grad_enabled without get_attr submodule is not supported") + + subgraph_module = _resolve_target(subgraph_ref.target) + if not hasattr(subgraph_module, "graph"): + raise NotImplementedError("wrap_with_set_grad_enabled target is not a GraphModule") + + inline_counter += 1 + inline_index = inline_counter + local_alias: dict[str, str] = {} + subgraph_inputs = list(node.args[2:]) + input_index = 0 + + ctx.push_local_aliases(local_alias) + try: + for sub_node in subgraph_module.graph.nodes: + sub_shape = get_shape(sub_node) + sub_dtype = dtype_to_ir(get_dtype(sub_node)) + sub_desc = ( + f"inlined:{node.name}:{sub_node.op}:{getattr(sub_node, 'name', '')}:" + f"{format_target(sub_node) if sub_node.op == 'call_function' else getattr(sub_node, 'target', '')}" + ) + ctx.push_node(sub_desc) + try: + if sub_node.op == "placeholder": + if input_index >= len(subgraph_inputs): + raise NotImplementedError("subgraph placeholder/input arity mismatch during inlining") + local_alias[sub_node.name] = ctx.resolve_value_id(subgraph_inputs[input_index]) + input_index += 1 + continue + + if sub_node.op == "get_attr": + prefixed = _make_inlined_node(sub_node, inline_index=inline_index) + value = _resolve_target(sub_node.target) + import_get_attr(ir, prefixed, ctx, value, shape=sub_shape, dtype=sub_dtype, source_name=str(sub_node.target)) + _try_register_weight_binding(f"v_{prefixed.name}", str(sub_node.target), value) + local_alias[sub_node.name] = f"v_{prefixed.name}" + continue + + if sub_node.op == "call_function": + prefixed = _make_inlined_node(sub_node, inline_index=inline_index) + import_call_function( + ir, + prefixed, + ctx, + shape=sub_shape, + dtype=sub_dtype, + torch_op=format_target(sub_node), + ) + local_alias[sub_node.name] = f"v_{prefixed.name}" + continue + + if sub_node.op == "output": + ctx.tuple_aliases[node.name] = _collect_output_ids(sub_node.args[0], local_alias) + return + + raise NotImplementedError(f"unsupported inlined subgraph node op: {sub_node.op}") + finally: + ctx.pop_node() + finally: + ctx.pop_local_aliases() + + raise NotImplementedError("inlined subgraph had no output node") + + for spec in captured.exported_program.graph_signature.input_specs: + arg = getattr(spec, "arg", None) + name = getattr(arg, "name", None) + if name is not None: + placeholder_specs[name] = spec + + for node in captured.graph.nodes: + shape = get_shape(node) + dtype = dtype_to_ir(get_dtype(node)) + node_desc = f"top:{node.op}:{getattr(node, 'name', '')}:{getattr(node, 'target', '')}" + ctx.push_node(node_desc) + try: + if node.op == "placeholder": + spec = placeholder_specs.get(node.name) + if spec is not None and isinstance(getattr(spec, "arg", None), ConstantArgument): + import_get_attr( + ir, + node, + ctx, + spec.arg.value, + shape=shape, + dtype=dtype, + source_name=node.name, + ) + continue + if spec is not None and spec.kind in {InputKind.PARAMETER, InputKind.BUFFER, InputKind.CONSTANT_TENSOR}: + target = getattr(spec, "target", None) + if target is None: + raise NotImplementedError(f"placeholder {node.name} has no target for lifted constant input") + value = _resolve_target(target) + import_get_attr(ir, node, ctx, value, shape=shape, dtype=dtype, source_name=str(target)) + _try_register_weight_binding(value_id(node, ctx), str(target), value) + continue + import_placeholder(ir, node, ctx, shape=shape, dtype=dtype) + continue + + if node.op == "get_attr": + if str(node.target).startswith("submod_"): + continue + value = _resolve_target(node.target) + import_get_attr(ir, node, ctx, value, shape=shape, dtype=dtype, source_name=str(node.target)) + _try_register_weight_binding(value_id(node, ctx), str(node.target), value) + continue + + if node.op == "call_function": + torch_op = format_target(node) + if torch_op.endswith(".wrap_with_set_grad_enabled") or torch_op == "wrap_with_set_grad_enabled": + _inline_wrap_with_set_grad_enabled(node, shape, dtype) + continue + if torch_op == "": + source = node.args[0] + if hasattr(source, "name") and source.name in ctx.tuple_aliases: + index_value = extract_literals(node.args[1]) + if not isinstance(index_value, int): + raise NotImplementedError(f"unsupported tuple getitem index: {node.args!r}") + alias_seq = ctx.tuple_aliases[source.name] + if index_value < 0 or index_value >= len(alias_seq): + raise NotImplementedError( + f"tuple getitem index {index_value} out of range for " + f"inlined subgraph {source.name!r} (size {len(alias_seq)})" + ) + ctx.alias_values[node.name] = alias_seq[index_value] + continue + import_call_function(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op) + continue + + if node.op == "call_module" and str(getattr(node, "target", "")) == "_guards_fn" and not node.users: + continue + + if node.op == "output": + import_output(ir, node, ctx) + continue + + raise NotImplementedError(f"unsupported FX node op: {node.op}") + finally: + ctx.pop_node() + + apply_import_semantics(ir) + annotate_ir_components(ir) + verify_ir(ir) + return ir diff --git a/python/cactus/transpile/import_semantics.py b/python/cactus/transpile/import_semantics.py new file mode 100644 index 000000000..0bf84d3e7 --- /dev/null +++ b/python/cactus/transpile/import_semantics.py @@ -0,0 +1,337 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.canonicalize.utils import rebuild_graph +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import verify_ir + + +@dataclass(frozen=True) +class EarlyRoPEMatch: + anchor_node_id: str + input_value_id: str + theta: float + position_offset: int + partial: bool + node_ids: tuple[str, ...] + + +def apply_import_semantics(graph: IRGraph) -> IRGraph: + _rewrite_explicit_attention(graph) + _tag_explicit_linear(graph) + _rewrite_early_rope(graph) + rebuild_graph(graph) + verify_ir(graph) + return graph + + +def _rewrite_explicit_attention(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op != "scaled_dot_product_attention": + continue + additive_mask = bool(node.attrs.get("additive_mask", False)) + if not additive_mask and len(node.inputs) > 3: + mask_value = graph.values.get(node.inputs[3]) + if mask_value is not None and mask_value.dtype not in (None, "bool"): + additive_mask = True + node.op = "attention" + node.kind = "semantic" + node.meta["semantic_source"] = "import" + node.attrs = { + "scale": float(node.attrs.get("scale", 0.0)), + "is_causal": bool(node.attrs.get("is_causal", False)), + "window_size": int(node.attrs.get("window_size", 0)), + **({"enable_gqa": bool(node.attrs.get("enable_gqa", False))} if "enable_gqa" in node.attrs else {}), + **({"additive_mask": True} if additive_mask else {}), + } + + +def _tag_explicit_linear(graph: IRGraph) -> None: + for node in graph.nodes.values(): + if node.op not in {"linear", "addmm"}: + continue + node.meta.setdefault("semantic_source", "import") + + +def _rewrite_early_rope(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or not node.outputs: + continue + match = match_early_rope(graph, node.outputs[0]) + if match is None or match.anchor_node_id != node.id: + continue + if match.partial: + node.meta["semantic_hint"] = "partial_rope" + node.meta["rope_theta_hint"] = float(match.theta) + node.meta["rope_position_offset_hint"] = int(match.position_offset) + continue + node.op = "rope" + node.inputs = [match.input_value_id] + node.attrs = { + "theta": float(match.theta), + "position_offset": int(match.position_offset), + } + node.kind = "semantic" + node.meta["semantic_source"] = "import" + + +def match_early_rope(graph: IRGraph, value_id: str) -> EarlyRoPEMatch | None: + desc = _extract_rope_descriptor(graph, value_id) + if desc is None: + return None + return EarlyRoPEMatch( + anchor_node_id=str(desc["anchor_node_id"]), + input_value_id=str(desc["input_value_id"]), + theta=float(desc["theta"]), + position_offset=int(desc["position_offset"]), + partial=bool(desc.get("partial_rope", False)), + node_ids=tuple(sorted(set(desc.get("node_ids", ())))), + ) + + +def _extract_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: + current = _strip_layout_passthrough(graph, value_id) + node = producer(graph, current) + if node is None: + return None + if node.op == "rope": + return { + "anchor_node_id": node.id, + "input_value_id": node.inputs[0], + "theta": float(node.attrs.get("theta", 0.0)), + "position_offset": int(node.attrs.get("position_offset", 0)), + "partial_rope": False, + "node_ids": (node.id,), + } + classic = _extract_classic_rope_descriptor(graph, current) + if classic is not None: + classic["partial_rope"] = False + return classic + if node.op == "cat" and len(node.inputs) == 2: + for rope_branch_value_id, passthrough_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + rope_branch = _extract_rope_descriptor(graph, rope_branch_value_id) + if rope_branch is None: + continue + passthrough = producer(graph, _strip_layout_passthrough(graph, passthrough_value_id)) + if passthrough is None or passthrough.op != "slice": + continue + if strip_passthrough(graph, passthrough.inputs[0]) != strip_passthrough(graph, rope_branch["input_value_id"]): + continue + rope_branch["anchor_node_id"] = node.id + rope_branch["partial_rope"] = True + rope_branch["node_ids"] = tuple(sorted({node.id, passthrough.id, *rope_branch["node_ids"]})) + return rope_branch + return None + + +def _extract_classic_rope_descriptor(graph: IRGraph, value_id: str) -> dict[str, object] | None: + node = producer(graph, _strip_layout_passthrough(graph, value_id)) + if node is None or node.op != "add" or len(node.inputs) != 2: + return None + for direct_mult_id, rotated_mult_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + direct_mult = producer(graph, direct_mult_id) + rotated_mult = producer(graph, rotated_mult_id) + if direct_mult is None or rotated_mult is None: + continue + if direct_mult.op != "multiply" or rotated_mult.op != "multiply": + continue + direct = _extract_direct_branch(graph, direct_mult) + rotated = _extract_rotated_branch(graph, rotated_mult) + if direct is None or rotated is None: + continue + input_value_id = strip_passthrough(graph, direct["input_value_id"]) + if input_value_id != strip_passthrough(graph, rotated["input_value_id"]): + continue + cos_info = _extract_rope_trig(graph, direct["trig_value_id"], expected=("cos", "scalar_cos")) + sin_info = _extract_rope_trig(graph, rotated["trig_value_id"], expected=("sin", "scalar_sin")) + if cos_info is None or sin_info is None: + continue + if abs(float(cos_info["theta"]) - float(sin_info["theta"])) > 1e-2: + continue + if int(cos_info["position_offset"]) != int(sin_info["position_offset"]): + continue + return { + "anchor_node_id": node.id, + "input_value_id": input_value_id, + "theta": float(cos_info["theta"]), + "position_offset": int(cos_info["position_offset"]), + "node_ids": tuple(sorted({node.id, *direct["node_ids"], *rotated["node_ids"], *cos_info["node_ids"], *sin_info["node_ids"]})), + } + return None + + +def _extract_direct_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: + for input_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + if _unwrap_trig_value(graph, trig_value_id, expected=("cos", "scalar_cos")) is not None: + return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} + return None + + +def _extract_rotated_branch(graph: IRGraph, node: IRNode) -> dict[str, object] | None: + for rotated_value_id, trig_value_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + input_value_id = _extract_rotate_half_source(graph, rotated_value_id) + if input_value_id is None: + continue + if _unwrap_trig_value(graph, trig_value_id, expected=("sin", "scalar_sin")) is None: + continue + return {"input_value_id": input_value_id, "trig_value_id": trig_value_id, "node_ids": (node.id,)} + return None + + +def _extract_rotate_half_source(graph: IRGraph, value_id: str) -> str | None: + cat_node = producer(graph, strip_passthrough(graph, value_id)) + if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: + return None + neg_node = producer(graph, cat_node.inputs[0]) + right_slice = producer(graph, cat_node.inputs[1]) + if neg_node is None or right_slice is None: + return None + if neg_node.op != "scalar_multiply" or float(neg_node.attrs.get("value", 0.0)) != -1.0: + return None + left_slice = producer(graph, neg_node.inputs[0]) + if left_slice is None or left_slice.op != "slice" or right_slice.op != "slice": + return None + if strip_passthrough(graph, left_slice.inputs[0]) != strip_passthrough(graph, right_slice.inputs[0]): + return None + return left_slice.inputs[0] + + +def _extract_rope_trig(graph: IRGraph, value_id: str, *, expected: tuple[str, ...]) -> dict[str, object] | None: + trig = _unwrap_trig_value(graph, value_id, expected=expected) + if trig is None: + return None + trig_node = trig["trig_node"] + cat_node = producer(graph, trig_node.inputs[0]) + if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2 or cat_node.inputs[0] != cat_node.inputs[1]: + return None + angle_info = _extract_angle_source(graph, cat_node.inputs[0]) + if angle_info is None: + return None + matmul_node = angle_info["matmul_node"] + inv_freq_const_id = None + arange_node = None + for input_id in matmul_node.inputs: + if inv_freq_const_id is None: + inv_freq_const_id = _find_constant_ancestor(graph, input_id) + if arange_node is None: + arange_node = _find_arange_ancestor(graph, input_id) + if inv_freq_const_id is None or arange_node is None: + return None + theta = _infer_rope_theta(graph.constants[inv_freq_const_id]) + if theta is None: + return None + return { + "theta": theta, + "position_offset": int(arange_node.attrs.get("start", 0)), + "node_ids": tuple(sorted({trig_node.id, cat_node.id, matmul_node.id})), + } + + +def _unwrap_trig_value(graph: IRGraph, value_id: str, *, expected: tuple[str, ...]) -> dict[str, object] | None: + current = value_id + while True: + node = producer(graph, current) + if node is None: + return None + if node.op in {"unsqueeze", "reshape", "view", "expand", "precision_cast"}: + current = node.inputs[0] + continue + if node.op == "type_as": + current = node.inputs[0] + continue + if node.op == "scalar_multiply": + current = node.inputs[0] + continue + if node.op not in expected: + return None + return {"trig_node": node} + + +def _extract_angle_source(graph: IRGraph, value_id: str) -> dict[str, object] | None: + current = value_id + while True: + current = _strip_layout_passthrough(graph, current) + node = producer(graph, current) + if node is None: + return None + if node.op in {"transpose", "permute"}: + matmul_node = producer(graph, node.inputs[0]) + if matmul_node is None or matmul_node.op != "matmul": + return None + return {"matmul_node": matmul_node} + if node.op == "matmul": + return {"matmul_node": node} + if node.op in {"index", "slice"} and len(node.inputs) == 1: + current = node.inputs[0] + continue + return None + + +def _find_constant_ancestor(graph: IRGraph, value_id: str) -> str | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + current = _strip_layout_passthrough(graph, current) + if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): + return current + node = producer(graph, current) + if node is None or len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _find_arange_ancestor(graph: IRGraph, value_id: str) -> IRNode | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + current = _strip_layout_passthrough(graph, current) + node = producer(graph, current) + if node is None: + return None + if node.op == "arange": + return node + if current in graph.constants and isinstance(graph.constants[current], torch.Tensor): + return None + if len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _strip_layout_passthrough(graph: IRGraph, value_id: str) -> str: + current = value_id + while True: + node = producer(graph, current) + if node is None: + return current + if node.op in {"precision_cast", "contiguous", "reshape", "view", "unsqueeze", "expand"} and len(node.inputs) == 1: + current = node.inputs[0] + continue + if node.op == "type_as" and len(node.inputs) >= 1: + current = node.inputs[0] + continue + return current + + +def _infer_rope_theta(value: Any) -> float | None: + if not isinstance(value, torch.Tensor): + return None + flat = value.detach().cpu().float().reshape(-1) + if flat.numel() < 2: + return None + second = float(flat[1].item()) + if second <= 0.0: + return None + return float((1.0 / second) ** flat.numel()) diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py new file mode 100644 index 000000000..9666f7c03 --- /dev/null +++ b/python/cactus/transpile/importers.py @@ -0,0 +1,1960 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from dataclasses import field +from types import SimpleNamespace +from typing import Any + +import torch + +from cactus.transpile.aten_ops import canonical_torch_op +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.normalize import dtype_to_ir +from cactus.transpile.normalize import normalize_target + + +class UnsupportedImportError(NotImplementedError): + pass + + +@dataclass +class ImportContext: + strict: bool = True + alias_values: dict[str, str] = field(default_factory=dict) + tuple_aliases: dict[str, list[str]] = field(default_factory=dict) + local_alias_stack: list[dict[str, str]] = field(default_factory=list) + node_stack: list[str] = field(default_factory=list) + transpile_metadata: dict[str, object] = field(default_factory=dict) + + def push_local_aliases(self, aliases: dict[str, str]) -> None: + self.local_alias_stack.append(aliases) + + def pop_local_aliases(self) -> None: + if not self.local_alias_stack: + raise RuntimeError("local alias stack underflow") + self.local_alias_stack.pop() + + def resolve_value_id(self, node: Any) -> str: + if is_fx_node(node): + name = getattr(node, "name", "") + for aliases in reversed(self.local_alias_stack): + if name in aliases: + return aliases[name] + resolved = self.alias_values.get(name) + if resolved is not None: + return resolved + return f"v_{name}" + raise TypeError(f"cannot resolve value id for non-FX node: {type(node).__name__}") + + def fail(self, message: str) -> None: + if self.node_stack: + message = f"{message} [context: {' > '.join(self.node_stack)}]" + if self.strict: + raise UnsupportedImportError(message) + raise UnsupportedImportError(f"{message} [non-strict import has no generic fallback]") + + def push_node(self, description: str) -> None: + self.node_stack.append(description) + + def pop_node(self) -> None: + if self.node_stack: + self.node_stack.pop() + + +def value_id(node: Any, ctx: ImportContext) -> str: + return ctx.resolve_value_id(node) + + +def node_id(node: Any) -> str: + return f"n_{node.name}" + + +def is_fx_node(value: Any) -> bool: + return hasattr(value, "name") and hasattr(value, "op") + + +def extract_input_ids(arg: Any, ctx: ImportContext) -> list[str]: + if is_fx_node(arg): + return [value_id(arg, ctx)] + if isinstance(arg, (tuple, list)): + ids: list[str] = [] + for item in arg: + ids.extend(extract_input_ids(item, ctx)) + return ids + if isinstance(arg, dict): + ids: list[str] = [] + for item in arg.values(): + ids.extend(extract_input_ids(item, ctx)) + return ids + return [] + + +def extract_literals(value: Any) -> Any: + if is_fx_node(value): + meta = getattr(value, "meta", {}) or {} + meta_value = meta.get("val") + if meta_value is not None and not isinstance(meta_value, torch.Tensor): + return extract_literals(meta_value) + return None + if isinstance(value, torch.Size): + return tuple(extract_literals(v) for v in value) + if isinstance(value, (tuple, list)): + return type(value)(extract_literals(v) for v in value) + if isinstance(value, dict): + return {k: extract_literals(v) for k, v in value.items()} + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, (int, float, bool, str)) or value is None: + return value + return None + + +def require_literal_bool(arg: Any, ctx: "ImportContext", name: str) -> bool: + literal = extract_literals(arg) + if not isinstance(literal, (bool, int)): + ctx.fail(f"{name} must be a boolean literal, got {literal!r}") + return bool(literal) + + +def make_prefixed_node(node: Any, prefix: str) -> Any: + return SimpleNamespace( + name=f"{prefix}{node.name}", + op=node.op, + target=node.target, + args=node.args, + kwargs=getattr(node, "kwargs", {}), + meta=getattr(node, "meta", {}), + ) + + +def add_value_if_missing(ir: IRGraph, value: IRValue) -> None: + if value.id not in ir.values: + ir.add_value(value) + + +def register_node( + ir: IRGraph, + ir_node: IRNode, + *, + shape: tuple[int, ...] | None, + dtype: str | None, +) -> None: + ir.add_node(ir_node) + ir.order.append(ir_node.id) + for output_id in ir_node.outputs: + ir.values[output_id].shape = shape + ir.values[output_id].dtype = dtype + for input_id in ir_node.inputs: + if input_id in ir.values: + ir.values[input_id].users.append(ir_node.id) + + +def _extract_numeric_literal(value: Any) -> float | int | None: + literal = extract_literals(value) + if isinstance(literal, bool): + return int(literal) + if isinstance(literal, (int, float)): + return literal + return None + + +def _extract_static_int(value: Any) -> int | None: + if is_fx_node(value): + meta = getattr(value, "meta", {}) or {} + meta_value = meta.get("val") + if meta_value is not None and not isinstance(meta_value, torch.Tensor): + return _extract_static_int(meta_value) + return None + + literal = extract_literals(value) + if isinstance(literal, bool): + return int(literal) + if isinstance(literal, int): + return int(literal) + if isinstance(literal, float) and float(literal).is_integer(): + return int(literal) + + try: + return int(value) + except Exception: + return None + + +def _extract_int_sequence_literal(value: Any) -> tuple[int, ...] | None: + if is_fx_node(value): + meta = getattr(value, "meta", {}) or {} + meta_value = meta.get("val") + if meta_value is not None and not isinstance(meta_value, torch.Tensor): + return _extract_int_sequence_literal(meta_value) + return None + + if isinstance(value, torch.Size): + value = tuple(value) + + if isinstance(value, (tuple, list)): + result: list[int] = [] + for item in value: + item_value = _extract_static_int(item) + if item_value is None: + return None + result.append(item_value) + return tuple(result) + + scalar_value = _extract_static_int(value) + if scalar_value is not None: + return (scalar_value,) + return None + + +def _try_materialize_shape(shape: tuple[Any, ...] | None) -> tuple[int, ...] | None: + if shape is None: + return None + result: list[int] = [] + for dim in shape: + dim_value = _extract_static_int(dim) + if dim_value is None: + return None + result.append(dim_value) + return tuple(result) + + +def _extract_fx_tensor_shape(value: Any) -> tuple[Any, ...] | None: + if not is_fx_node(value): + return None + meta = getattr(value, "meta", {}) or {} + meta_value = meta.get("val") + if isinstance(meta_value, torch.Tensor): + return tuple(meta_value.shape) + tensor_meta = meta.get("tensor_meta") + tensor_shape = getattr(tensor_meta, "shape", None) + if tensor_shape is not None: + return tuple(tensor_shape) + return None + + +def _base_meta(shape: tuple[int, ...] | None, dtype: str | None, torch_op: str, node: Any) -> dict[str, object]: + aten_op = canonical_torch_op(torch_op) + return { + "shape": shape, + "dtype": dtype, + "torch_op": aten_op, + "aten_op": aten_op, + "torch_name": node.name, + } + + +def _extract_fx_torch_dtype(node: Any) -> torch.dtype | None: + if not is_fx_node(node): + return None + meta = getattr(node, "meta", {}) or {} + value = meta.get("val") + dtype = getattr(value, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + tensor_meta = meta.get("tensor_meta") + dtype = getattr(tensor_meta, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + return None + + +def _import_scalar_binary( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, + op_name: str, +) -> None: + lhs, rhs = node.args[0], node.args[1] + lhs_is_node = is_fx_node(lhs) + rhs_is_node = is_fx_node(rhs) + lhs_literal = _extract_numeric_literal(lhs) + rhs_literal = _extract_numeric_literal(rhs) + + if lhs_is_node and rhs_is_node: + ir_node = IRNode( + id=node_id(node), + op=op_name, + inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + return + + if op_name == "add": + if lhs_is_node and rhs_literal is not None: + scalar_op = "scalar_add" + inputs = [value_id(lhs, ctx)] + attrs = {"value": rhs_literal} + elif rhs_is_node and lhs_literal is not None: + scalar_op = "scalar_add" + inputs = [value_id(rhs, ctx)] + attrs = {"value": lhs_literal} + else: + scalar_op = None + elif op_name == "multiply": + if lhs_is_node and rhs_literal is not None: + scalar_op = "scalar_multiply" + inputs = [value_id(lhs, ctx)] + attrs = {"value": rhs_literal} + elif rhs_is_node and lhs_literal is not None: + scalar_op = "scalar_multiply" + inputs = [value_id(rhs, ctx)] + attrs = {"value": lhs_literal} + else: + scalar_op = None + elif op_name == "subtract": + if lhs_is_node and rhs_literal is not None: + scalar_op = "scalar_subtract" + inputs = [value_id(lhs, ctx)] + attrs = {"value": rhs_literal} + elif rhs_is_node and lhs_literal is not None: + scalar_op = "scalar_subtract_reverse" + inputs = [value_id(rhs, ctx)] + attrs = {"value": lhs_literal} + else: + scalar_op = None + elif op_name == "divide": + if lhs_is_node and rhs_literal is not None: + scalar_op = "scalar_divide" + inputs = [value_id(lhs, ctx)] + attrs = {"value": rhs_literal} + elif rhs_is_node and lhs_literal is not None: + scalar_op = "scalar_divide_reverse" + inputs = [value_id(rhs, ctx)] + attrs = {"value": lhs_literal} + else: + scalar_op = None + else: + scalar_op = None + + if scalar_op is None: + ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") + + ir_node = IRNode( + id=node_id(node), + op=scalar_op, + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_placeholder( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, +) -> None: + add_value_if_missing(ir, IRValue(id=value_id(node, ctx), shape=shape, dtype=dtype, producer=None)) + ir.inputs.append(value_id(node, ctx)) + + +def import_get_attr( + ir: IRGraph, + node: Any, + ctx: ImportContext, + value: Any, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + source_name: str | None = None, +) -> None: + add_value_if_missing(ir, IRValue(id=value_id(node, ctx), shape=shape, dtype=dtype, producer=None)) + if isinstance(value, torch.Tensor): + ir.constants[value_id(node, ctx)] = value if value.is_meta else value.detach().cpu() + else: + ir.constants[value_id(node, ctx)] = value + ir.values[value_id(node, ctx)].meta["source_name"] = source_name + if source_name is not None: + bindings = ir.meta.setdefault("weight_bindings", {}) + if isinstance(bindings, dict): + bindings.setdefault(value_id(node, ctx), {})["source_name"] = source_name + + +def import_output(ir: IRGraph, node: Any, ctx: ImportContext) -> None: + ir.outputs.extend(extract_input_ids(node.args[0], ctx)) + + +def import_call_function( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + op = normalize_target(torch_op) + importer = OP_IMPORTERS.get(op) + if importer is None: + import_opaque_call_function(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name=op) + else: + importer(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op) + + +def import_opaque_call_function( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, + op_name: str, +) -> None: + ir_node = IRNode( + id=node_id(node), + op=op_name, + inputs=extract_input_ids(getattr(node, "args", ()), ctx) + extract_input_ids(getattr(node, "kwargs", {}), ctx), + outputs=[value_id(node, ctx)], + attrs={ + "opaque": True, + "torch_op": torch_op, + "args": extract_literals(getattr(node, "args", ())), + "kwargs": extract_literals(getattr(node, "kwargs", {})), + }, + meta=_base_meta(shape, dtype, torch_op, node), + kind="opaque", + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_add(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="add") + + +def import_subtract(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="subtract") + + +def import_multiply(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="multiply") + + +def import_multiply_inplace(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="multiply") + + +def import_divide(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + _import_scalar_binary(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="divide") + + +def import_not_equal(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + lhs, rhs = node.args[0], node.args[1] + lhs_is_node = is_fx_node(lhs) + rhs_is_node = is_fx_node(rhs) + lhs_literal = _extract_numeric_literal(lhs) + rhs_literal = _extract_numeric_literal(rhs) + + if lhs_is_node and rhs_is_node: + ir_node = IRNode( + id=node_id(node), + op="not_equal", + inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + return + + if lhs_is_node and rhs_literal is not None: + inputs = [value_id(lhs, ctx)] + attrs = {"value": rhs_literal} + elif rhs_is_node and lhs_literal is not None: + inputs = [value_id(rhs, ctx)] + attrs = {"value": lhs_literal} + else: + ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") + + ir_node = IRNode( + id=node_id(node), + op="scalar_not_equal", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_equal(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + lhs, rhs = node.args[0], node.args[1] + lhs_is_node = is_fx_node(lhs) + rhs_is_node = is_fx_node(rhs) + lhs_literal = _extract_numeric_literal(lhs) + rhs_literal = _extract_numeric_literal(rhs) + + if lhs_is_node and rhs_is_node: + ir_node = IRNode( + id=node_id(node), + op="equal", + inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + return + + if lhs_is_node and rhs_literal is not None: + inputs = [value_id(lhs, ctx)] + attrs = {"value": rhs_literal} + elif rhs_is_node and lhs_literal is not None: + inputs = [value_id(rhs, ctx)] + attrs = {"value": lhs_literal} + else: + ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") + + ir_node = IRNode( + id=node_id(node), + op="scalar_equal", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_compare(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str, op_name: str) -> None: + lhs, rhs = node.args[0], node.args[1] + lhs_is_node = is_fx_node(lhs) + rhs_is_node = is_fx_node(rhs) + lhs_literal = _extract_numeric_literal(lhs) + rhs_literal = _extract_numeric_literal(rhs) + + if lhs_is_node and rhs_is_node: + ir_node = IRNode( + id=node_id(node), + op=op_name, + inputs=[value_id(lhs, ctx), value_id(rhs, ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + return + + if lhs_is_node and rhs_literal is not None: + inputs = [value_id(lhs, ctx)] + scalar_op = f"scalar_{op_name}" + scalar_value = rhs_literal + elif rhs_is_node and lhs_literal is not None: + reverse_ops = { + "greater": "less", + "greater_equal": "less_equal", + "less": "greater", + "less_equal": "greater_equal", + } + inputs = [value_id(rhs, ctx)] + scalar_op = f"scalar_{reverse_ops[op_name]}" + scalar_value = lhs_literal + else: + ctx.fail(f"unsupported scalar form for {torch_op}: {node.args!r}") + + ir_node = IRNode( + id=node_id(node), + op=scalar_op, + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs={"value": scalar_value}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_binary(op_name: str): + def _importer( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, + ) -> None: + ir_node = IRNode( + id=node_id(node), + op=op_name, + inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + return _importer + + +def import_where( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + condition = node.args[0] + true_value = node.args[1] + false_value = node.args[2] + + inputs = [value_id(condition, ctx)] + attrs: dict[str, object] = {} + + if is_fx_node(true_value): + inputs.append(value_id(true_value, ctx)) + attrs["true_is_scalar"] = False + else: + literal = _extract_numeric_literal(true_value) + if literal is None: + ctx.fail(f"unsupported true branch for {torch_op}: {node.args!r}") + attrs["true_is_scalar"] = True + attrs["true_value"] = float(literal) + + if is_fx_node(false_value): + inputs.append(value_id(false_value, ctx)) + attrs["false_is_scalar"] = False + else: + literal = _extract_numeric_literal(false_value) + if literal is None: + ctx.fail(f"unsupported false branch for {torch_op}: {node.args!r}") + attrs["false_is_scalar"] = True + attrs["false_value"] = float(literal) + + ir_node = IRNode( + id=node_id(node), + op="where", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_negate(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="scalar_multiply", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"value": -1.0}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_precision_cast(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + target_dtype = extract_literals(node.args[1]) if len(node.args) > 1 else None + ir_node = IRNode( + id=node_id(node), + op="precision_cast", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"dtype": target_dtype}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_arange(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + attrs: dict[str, object] = {} + if len(node.args) == 1: + attrs["start"] = 0 + attrs["end"] = int(node.args[0]) + elif len(node.args) >= 2: + attrs["start"] = int(node.args[0]) + attrs["end"] = int(node.args[1]) + else: + ctx.fail(f"unsupported arange signature for {torch_op}: {node.args!r}") + + if len(node.args) > 2 and extract_literals(node.args[2]) is not None: + attrs["step"] = int(node.args[2]) + + target_dtype = None + if "dtype" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dtype"]) is not None: + target_dtype = extract_literals(node.kwargs["dtype"]) + elif dtype is not None: + target_dtype = dtype + if target_dtype is not None: + attrs["dtype"] = target_dtype + + ir_node = IRNode( + id=node_id(node), + op="arange", + inputs=[], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_type_as(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="type_as", + inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_identity(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="identity", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_reshape(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + target_shape = _extract_int_sequence_literal(node.args[1]) + if target_shape is None: + ctx.fail(f"unsupported reshape shape for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="reshape", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"shape": target_shape}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_flatten(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="flatten", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"start_dim": int(node.args[1]), "end_dim": int(node.args[2])}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_unsqueeze(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + dim = int(node.args[1]) + attrs: dict[str, object] = {"dim": dim} + materialized_shape = _try_materialize_shape(shape) + if materialized_shape is not None: + attrs["shape"] = materialized_shape + ir_node = IRNode( + id=node_id(node), + op="unsqueeze", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_squeeze(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + dim = int(node.args[1]) if len(node.args) > 1 and node.args[1] is not None else 0 + attrs: dict[str, object] = {"dim": dim} + materialized_shape = _try_materialize_shape(shape) + if materialized_shape is not None: + attrs["shape"] = materialized_shape + ir_node = IRNode( + id=node_id(node), + op="squeeze", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_masked_scatter( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + ir_node = IRNode( + id=node_id(node), + op="masked_scatter", + inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx), value_id(node.args[2], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_masked_fill( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + value = _extract_numeric_literal(node.args[2]) + if value is None: + ctx.fail(f"unsupported masked_fill value for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="masked_fill", + inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], + outputs=[value_id(node, ctx)], + attrs={"value": float(value)}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_expand(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + target_shape = _extract_int_sequence_literal(node.args[1]) + if target_shape is None and shape is not None: + target_shape = _extract_int_sequence_literal(shape) + if target_shape is None: + ctx.fail(f"unsupported expand shape for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="expand", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"shape": target_shape}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_repeat(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + repeats = _extract_int_sequence_literal(node.args[1]) if len(node.args) > 1 else None + if repeats is None: + ctx.fail(f"unsupported repeat factors for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="repeat", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"repeats": repeats}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_one_hot(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + num_classes = _extract_static_int(node.args[1]) if len(node.args) > 1 else None + if (num_classes is None or num_classes <= 0) and shape is not None and shape: + num_classes = int(shape[-1]) + if num_classes is None or num_classes <= 0: + ctx.fail(f"unsupported one_hot num_classes for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="one_hot", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"num_classes": int(num_classes)}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_tril(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + diagonal = int(extract_literals(node.args[1])) if len(node.args) > 1 and extract_literals(node.args[1]) is not None else 0 + ir_node = IRNode( + id=node_id(node), + op="tril", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"diagonal": diagonal}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_unfold(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 4: + ctx.fail(f"unsupported unfold args for {torch_op}: {node.args!r}") + dimension = extract_literals(node.args[1]) + size = extract_literals(node.args[2]) + step = extract_literals(node.args[3]) + if not isinstance(dimension, int) or not isinstance(size, int) or not isinstance(step, int): + ctx.fail(f"unsupported unfold args for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="unfold", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"dimension": dimension, "size": size, "step": step}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_transpose(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + attrs = {"dim0": 0, "dim1": 1} if torch_op == "aten.t.default" else {"dim0": int(node.args[1]), "dim1": int(node.args[2])} + ir_node = IRNode( + id=node_id(node), + op="transpose", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_permute(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + permutation = _extract_int_sequence_literal(node.args[1]) + if permutation is None: + ctx.fail(f"unsupported permute dims for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="permute", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"permutation": permutation}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_numpy_t(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + input_shape = _extract_fx_tensor_shape(node.args[0]) + rank_source = input_shape if input_shape is not None else shape + if rank_source is None: + ctx.fail(f"unsupported numpy_T rank inference for {torch_op}: {node.args!r}") + permutation = tuple(range(len(rank_source) - 1, -1, -1)) + ir_node = IRNode( + id=node_id(node), + op="permute", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"permutation": permutation}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_movedim(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + source = _extract_int_sequence_literal(node.args[1]) + destination = _extract_int_sequence_literal(node.args[2]) + if source is None or destination is None: + ctx.fail(f"unsupported movedim dims for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="movedim", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"source": source, "destination": destination}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_matmul(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="matmul", + inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_linear(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] + attrs = {"has_bias": False} + if len(node.args) > 2 and node.args[2] is not None: + inputs.append(value_id(node.args[2], ctx)) + attrs["has_bias"] = True + ir_node = IRNode( + id=node_id(node), + op="linear", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_addmm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + kwargs = getattr(node, "kwargs", {}) + beta = _extract_numeric_literal(node.args[3]) if len(node.args) > 3 else _extract_numeric_literal(kwargs.get("beta")) + alpha = _extract_numeric_literal(node.args[4]) if len(node.args) > 4 else _extract_numeric_literal(kwargs.get("alpha")) + if (beta is not None and float(beta) != 1.0) or (alpha is not None and float(alpha) != 1.0): + ctx.fail(f"unsupported scaled addmm (beta/alpha != 1) for {torch_op}: {node.args!r} {kwargs!r}") + ir_node = IRNode( + id=node_id(node), + op="addmm", + inputs=[value_id(node.args[0], ctx), value_id(node.args[1], ctx), value_id(node.args[2], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_unary(op_name: str): + def _importer(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op=op_name, + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + return _importer + + +def import_clamp(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + min_value = extract_literals(node.args[1]) if len(node.args) > 1 else None + max_value = extract_literals(node.args[2]) if len(node.args) > 2 else None + attrs: dict[str, object] = {} + inputs = [value_id(node.args[0], ctx)] + if isinstance(min_value, (int, float)): + attrs["min"] = float(min_value) + elif len(node.args) > 1 and is_fx_node(node.args[1]): + attrs["has_min_tensor"] = True + inputs.append(value_id(node.args[1], ctx)) + if isinstance(max_value, (int, float)): + attrs["max"] = float(max_value) + elif len(node.args) > 2 and is_fx_node(node.args[2]): + attrs["has_max_tensor"] = True + inputs.append(value_id(node.args[2], ctx)) + ir_node = IRNode( + id=node_id(node), + op="clamp", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_pow_rsqrt(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="pow", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"exponent": -0.5}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_reciprocal(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="pow", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"exponent": -1.0}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_floor_divide( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + if len(node.args) < 2: + ctx.fail(f"unsupported floor_divide signature for {torch_op}: {node.args!r}") + divisor = _extract_numeric_literal(node.args[1]) + if divisor is None: + ctx.fail(f"unsupported floor_divide divisor for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="scalar_floor_divide", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"value": float(divisor)}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_softmax(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="softmax", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"axis": int(node.args[1])}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_glu(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + axis = -1 + if len(node.args) > 1 and extract_literals(node.args[1]) is not None: + axis = int(node.args[1]) + ir_node = IRNode( + id=node_id(node), + op="glu", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"axis": axis}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_scaled_dot_product_attention( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + attrs: dict[str, object] = {} + inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx), value_id(node.args[2], ctx)] + if len(node.args) > 3: + mask_literal = extract_literals(node.args[3]) + if mask_literal is not None: + attrs["mask"] = mask_literal + elif is_fx_node(node.args[3]): + inputs.append(value_id(node.args[3], ctx)) + mask_dtype = _extract_fx_torch_dtype(node.args[3]) + if mask_dtype is not None and mask_dtype != torch.bool: + attrs["additive_mask"] = True + if len(node.args) > 4 and extract_literals(node.args[4]) is not None: + attrs["dropout_p"] = float(node.args[4]) + if len(node.args) > 5 and extract_literals(node.args[5]) is not None: + attrs["is_causal"] = bool(node.args[5]) + if "scale" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["scale"]) is not None: + attrs["scale"] = float(node.kwargs["scale"]) + if "enable_gqa" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["enable_gqa"]) is not None: + attrs["enable_gqa"] = bool(node.kwargs["enable_gqa"]) + ir_node = IRNode( + id=node_id(node), + op="scaled_dot_product_attention", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_reduce(op_name: str): + def _importer(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + attrs: dict[str, object] = {} + if len(node.args) > 1 and extract_literals(node.args[1]) is not None: + attrs["axis"] = extract_literals(node.args[1]) + keepdim = None + if len(node.args) > 2 and extract_literals(node.args[2]) is not None: + keepdim = bool(extract_literals(node.args[2])) + elif "keepdim" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["keepdim"]) is not None: + keepdim = bool(extract_literals(node.kwargs["keepdim"])) + if keepdim is not None: + attrs["keepdim"] = keepdim + ir_node = IRNode( + id=node_id(node), + op=op_name, + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + return _importer + + +def import_cat(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + tensors_arg = node.args[0] + axis = int(node.args[1]) if len(node.args) > 1 else 0 + ir_node = IRNode( + id=node_id(node), + op="cat", + inputs=extract_input_ids(tensors_arg, ctx), + outputs=[value_id(node, ctx)], + attrs={"axis": axis}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_split_with_sizes(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 2: + ctx.fail(f"unsupported split_with_sizes signature for {torch_op}: {node.args!r}") + sizes = _extract_int_sequence_literal(node.args[1]) + if sizes is None: + ctx.fail(f"unsupported split sizes for {torch_op}: {node.args!r}") + axis = -1 + if len(node.args) > 2 and extract_literals(node.args[2]) is not None: + axis = int(extract_literals(node.args[2])) + ir_node = IRNode( + id=node_id(node), + op="split_with_sizes", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"sizes": sizes, "axis": axis}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_chunk(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 2: + ctx.fail(f"unsupported chunk signature for {torch_op}: {node.args!r}") + chunks = extract_literals(node.args[1]) + if not isinstance(chunks, int): + ctx.fail(f"unsupported chunk count for {torch_op}: {node.args!r}") + axis = 0 + if len(node.args) > 2 and extract_literals(node.args[2]) is not None: + axis = int(extract_literals(node.args[2])) + ir_node = IRNode( + id=node_id(node), + op="chunk", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"chunks": int(chunks), "axis": axis}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_unbind(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 1: + ctx.fail(f"unsupported unbind signature for {torch_op}: {node.args!r}") + axis = 0 + if len(node.args) > 1 and extract_literals(node.args[1]) is not None: + axis = int(extract_literals(node.args[1])) + ir_node = IRNode( + id=node_id(node), + op="unbind", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"axis": axis}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_pow(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + base_literal = _extract_numeric_literal(node.args[0]) if len(node.args) > 0 else None + attrs: dict[str, object] = {} + exponent_literal = extract_literals(node.args[1]) if len(node.args) > 1 else None + + if base_literal is not None and isinstance(exponent_literal, (int, float)): + add_value_if_missing(ir, IRValue(id=value_id(node, ctx), shape=shape, dtype=dtype, producer=None)) + ir.constants[value_id(node, ctx)] = base_literal**exponent_literal + return + + if base_literal is not None: + if not (len(node.args) > 1 and is_fx_node(node.args[1])): + ctx.fail(f"unsupported pow signature for {torch_op}: scalar base with non-literal exponent {node.args!r}") + if float(base_literal) <= 0.0: + ctx.fail(f"unsupported pow signature for {torch_op}: non-positive scalar base {base_literal!r}") + + exponent_value_id = value_id(node.args[1], ctx) + mul_output_id = f"{value_id(node, ctx)}__pow_scalar_base_log_mul" + mul_node = IRNode( + id=f"{node_id(node)}__pow_scalar_base_log_mul", + op="scalar_multiply", + inputs=[exponent_value_id], + outputs=[mul_output_id], + attrs={"value": float(math.log(float(base_literal)))}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, mul_node, shape=shape, dtype=dtype) + + exp_node = IRNode( + id=node_id(node), + op="scalar_exp", + inputs=[mul_output_id], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, exp_node, shape=shape, dtype=dtype) + return + + inputs = [value_id(node.args[0], ctx)] + if exponent_literal is None and len(node.args) > 1 and is_fx_node(node.args[1]): + inputs.append(value_id(node.args[1], ctx)) + else: + attrs["exponent"] = exponent_literal + ir_node = IRNode( + id=node_id(node), + op="pow", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_slice(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + attrs = {"axis": int(node.args[1]), "start": int(node.args[2]), "end": int(node.args[3])} + if len(node.args) > 4 and extract_literals(node.args[4]) is not None: + attrs["step"] = int(node.args[4]) + ir_node = IRNode( + id=node_id(node), + op="slice", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_diff(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + inputs = list(getattr(node, "args", ())) + if not inputs: + ctx.fail(f"unsupported diff signature for {torch_op}: missing input tensor") + + source = inputs[0] + if not is_fx_node(source): + ctx.fail(f"unsupported diff signature for {torch_op}: source is not FX node") + + n_value = 1 + if len(inputs) > 1 and extract_literals(inputs[1]) is not None: + n_value = int(extract_literals(inputs[1])) + elif "n" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["n"]) is not None: + n_value = int(extract_literals(node.kwargs["n"])) + if n_value != 1: + ctx.fail(f"unsupported diff order for {torch_op}: n={n_value}") + + dim_value = None + if len(inputs) > 2 and extract_literals(inputs[2]) is not None: + dim_value = int(extract_literals(inputs[2])) + elif "dim" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dim"]) is not None: + dim_value = int(extract_literals(node.kwargs["dim"])) + if dim_value is None: + ctx.fail(f"unsupported diff signature for {torch_op}: missing integer dim") + + prepend_value = None + append_value = None + prepend_node = None + append_node = None + if len(inputs) > 3: + if is_fx_node(inputs[3]): + prepend_node = inputs[3] + else: + prepend_value = extract_literals(inputs[3]) + if len(inputs) > 4: + if is_fx_node(inputs[4]): + append_node = inputs[4] + else: + append_value = extract_literals(inputs[4]) + if "prepend" in getattr(node, "kwargs", {}): + prepend_arg = node.kwargs["prepend"] + if is_fx_node(prepend_arg): + prepend_node = prepend_arg + prepend_value = None + else: + prepend_value = extract_literals(prepend_arg) + if "append" in getattr(node, "kwargs", {}): + append_arg = node.kwargs["append"] + if is_fx_node(append_arg): + append_node = append_arg + append_value = None + else: + append_value = extract_literals(append_arg) + if prepend_value is not None or append_value is not None: + ctx.fail(f"unsupported diff signature for {torch_op}: literal prepend/append are not supported") + + source_value_id = value_id(source, ctx) + source_shape = None + if source_value_id in ir.values: + source_shape = ir.values[source_value_id].shape + if source_shape is None: + source_shape = shape + if source_shape is None: + ctx.fail(f"unsupported diff import for {torch_op}: missing source shape") + + rank = len(source_shape) + normalized_dim = dim_value if dim_value >= 0 else dim_value + rank + if normalized_dim < 0 or normalized_dim >= rank: + ctx.fail(f"unsupported diff dimension for {torch_op}: dim={dim_value} rank={rank}") + + dim_extent = int(source_shape[normalized_dim]) + if dim_extent < 1: + ctx.fail(f"unsupported diff input extent for {torch_op}: dim size {dim_extent}") + + concat_inputs = [source_value_id] + concat_shapes = [tuple(int(v) for v in source_shape)] + + def _append_concat_input(arg: Any, label: str) -> None: + if arg is None: + return + arg_value_id = value_id(arg, ctx) + arg_value = ir.values.get(arg_value_id) + arg_shape = arg_value.shape if arg_value is not None else None + if arg_shape is None: + ctx.fail(f"unsupported diff import for {torch_op}: missing {label} shape") + arg_shape = tuple(int(v) for v in arg_shape) + if len(arg_shape) != rank: + ctx.fail( + f"unsupported diff import for {torch_op}: {label} rank {len(arg_shape)} does not match source rank {rank}" + ) + for axis, (arg_dim, src_dim) in enumerate(zip(arg_shape, source_shape, strict=True)): + if axis == normalized_dim: + continue + if int(arg_dim) != int(src_dim): + ctx.fail( + f"unsupported diff import for {torch_op}: {label} shape {arg_shape} " + f"is incompatible with source shape {source_shape}" + ) + concat_inputs.append(arg_value_id) + concat_shapes.append(arg_shape) + + if prepend_node is not None: + concat_inputs = [] + concat_shapes = [] + _append_concat_input(prepend_node, "prepend") + concat_inputs.append(source_value_id) + concat_shapes.append(tuple(int(v) for v in source_shape)) + if append_node is not None: + _append_concat_input(append_node, "append") + elif append_node is not None: + _append_concat_input(append_node, "append") + + augmented_shape = list(source_shape) + augmented_shape[normalized_dim] = sum(int(shape_item[normalized_dim]) for shape_item in concat_shapes) + augmented_shape_tuple = tuple(int(v) for v in augmented_shape) + diff_shape = list(augmented_shape) + diff_shape[normalized_dim] = max(0, int(augmented_shape[normalized_dim]) - 1) + diff_shape_tuple = tuple(int(v) for v in diff_shape) + + diff_source_value_id = source_value_id + if len(concat_inputs) > 1: + concat_output_id = f"{value_id(node, ctx)}__diff_source" + concat_node = IRNode( + id=f"{node_id(node)}__diff_source", + op="cat", + inputs=concat_inputs, + outputs=[concat_output_id], + attrs={"axis": normalized_dim}, + meta=_base_meta(augmented_shape_tuple, dtype, torch_op, node), + ) + register_node(ir, concat_node, shape=augmented_shape_tuple, dtype=dtype) + diff_source_value_id = concat_output_id + + left_shape = list(diff_shape_tuple) + left_shape_tuple = tuple(int(v) for v in left_shape) + + left_node = IRNode( + id=f"{node_id(node)}__diff_left", + op="slice", + inputs=[diff_source_value_id], + outputs=[f"{value_id(node, ctx)}__diff_left"], + attrs={"axis": normalized_dim, "start": 0, "end": int(augmented_shape[normalized_dim]) - 1, "step": 1}, + meta=_base_meta(left_shape_tuple, dtype, torch_op, node), + ) + register_node(ir, left_node, shape=left_shape_tuple, dtype=dtype) + + right_node = IRNode( + id=f"{node_id(node)}__diff_right", + op="slice", + inputs=[diff_source_value_id], + outputs=[f"{value_id(node, ctx)}__diff_right"], + attrs={"axis": normalized_dim, "start": 1, "end": int(augmented_shape[normalized_dim]), "step": 1}, + meta=_base_meta(left_shape_tuple, dtype, torch_op, node), + ) + register_node(ir, right_node, shape=left_shape_tuple, dtype=dtype) + + diff_node = IRNode( + id=node_id(node), + op="subtract", + inputs=[right_node.outputs[0], left_node.outputs[0]], + outputs=[value_id(node, ctx)], + meta=_base_meta(diff_shape_tuple, dtype, torch_op, node), + ) + register_node(ir, diff_node, shape=diff_shape_tuple, dtype=dtype) + + +def import_index(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if torch_op.startswith("aten.select"): + axis = int(node.args[1]) + index_value = int(node.args[2]) + else: + index_value = int(node.args[1]) + axis = int(node.args[2]) if len(node.args) > 2 else 0 + ir_node = IRNode( + id=node_id(node), + op="index", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"index_value": index_value, "axis": axis}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_gather(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="gather", + inputs=[value_id(node.args[0], ctx), value_id(node.args[2], ctx)], + outputs=[value_id(node, ctx)], + attrs={"axis": int(node.args[1])}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_embedding(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] + attrs: dict[str, object] = {} + if len(node.args) > 2 and extract_literals(node.args[2]) is not None: + attrs["padding_idx"] = int(node.args[2]) + if len(node.args) > 3 and extract_literals(node.args[3]) is not None: + attrs["scale_grad_by_freq"] = bool(node.args[3]) + if len(node.args) > 4 and extract_literals(node.args[4]) is not None: + attrs["sparse"] = bool(node.args[4]) + ir_node = IRNode( + id=node_id(node), + op="embedding", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def _extract_int_list_or_scalar(value: Any, *, default: int) -> int: + literal = extract_literals(value) + if literal is None: + return default + if isinstance(literal, int): + return int(literal) + if isinstance(literal, (tuple, list)) and literal: + first = literal[0] + if isinstance(first, int): + return int(first) + raise UnsupportedImportError(f"unsupported integer/list literal: {value!r}") + + +def import_conv1d(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 2: + ctx.fail(f"unsupported conv1d signature for {torch_op}: {node.args!r}") + + inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] + if len(node.args) > 2 and is_fx_node(node.args[2]): + inputs.append(value_id(node.args[2], ctx)) + + stride = _extract_int_list_or_scalar(node.args[3], default=1) if len(node.args) > 3 else 1 + padding = _extract_int_list_or_scalar(node.args[4], default=0) if len(node.args) > 4 else 0 + dilation = _extract_int_list_or_scalar(node.args[5], default=1) if len(node.args) > 5 else 1 + groups = _extract_int_list_or_scalar(node.args[6], default=1) if len(node.args) > 6 else 1 + + ir_node = IRNode( + id=node_id(node), + op="conv1d", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs={"stride": stride, "padding": padding, "dilation": dilation, "groups": groups}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_conv2d(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 2: + ctx.fail(f"unsupported conv2d signature for {torch_op}: {node.args!r}") + + inputs = [value_id(node.args[0], ctx), value_id(node.args[1], ctx)] + if len(node.args) > 2 and is_fx_node(node.args[2]): + inputs.append(value_id(node.args[2], ctx)) + + stride = _extract_int_list_or_scalar(node.args[3], default=1) if len(node.args) > 3 else 1 + padding = _extract_int_list_or_scalar(node.args[4], default=0) if len(node.args) > 4 else 0 + dilation = _extract_int_list_or_scalar(node.args[5], default=1) if len(node.args) > 5 else 1 + groups = _extract_int_list_or_scalar(node.args[6], default=1) if len(node.args) > 6 else 1 + + ir_node = IRNode( + id=node_id(node), + op="conv2d", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs={"stride": stride, "padding": padding, "dilation": dilation, "groups": groups}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_pad(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if len(node.args) < 2: + ctx.fail(f"unsupported pad signature for {torch_op}: {node.args!r}") + + pads = extract_literals(node.args[1]) + if not isinstance(pads, (tuple, list)) or not all(isinstance(v, int) for v in pads): + ctx.fail(f"unsupported pads for {torch_op}: {node.args!r}") + + mode = "constant" + if len(node.args) > 2 and extract_literals(node.args[2]) is not None: + mode = str(extract_literals(node.args[2])) + if "mode" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["mode"]) is not None: + mode = str(extract_literals(node.kwargs["mode"])) + + value = 0.0 + if len(node.args) > 3 and extract_literals(node.args[3]) is not None: + value = float(extract_literals(node.args[3])) + if "value" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["value"]) is not None: + value = float(extract_literals(node.kwargs["value"])) + + ir_node = IRNode( + id=node_id(node), + op="pad", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"pads": tuple(int(v) for v in pads), "mode": mode, "value": value}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_ones(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if not node.args: + ctx.fail(f"unsupported ones signature for {torch_op}: {node.args!r}") + + output_shape = _extract_int_sequence_literal(node.args[0]) + if output_shape is None: + ctx.fail(f"unsupported ones shape for {torch_op}: {node.args!r}") + + node_dtype = dtype + if "dtype" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dtype"]) is not None: + node_dtype = dtype_to_ir(extract_literals(node.kwargs["dtype"])) + elif len(node.args) > 1 and extract_literals(node.args[1]) is not None: + node_dtype = dtype_to_ir(extract_literals(node.args[1])) + if node_dtype is None: + node_dtype = "fp32" + + ir_node = IRNode( + id=node_id(node), + op="ones", + inputs=[], + outputs=[value_id(node, ctx)], + attrs={"shape": output_shape, "dtype": node_dtype}, + meta=_base_meta(shape or output_shape, node_dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape or output_shape, dtype=node_dtype) + + +def import_new_empty(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + output_shape = _try_materialize_shape(shape) + if output_shape is None and len(node.args) > 1: + output_shape = _extract_int_sequence_literal(node.args[1]) + if output_shape is None: + ctx.fail(f"unsupported new_empty shape for {torch_op}: {node.args!r}") + elif math.prod(output_shape) != 0: + ctx.fail(f"new_empty requires an empty (zero-numel) shape, got {output_shape} for {torch_op}") + + node_dtype = dtype + if "dtype" in getattr(node, "kwargs", {}) and extract_literals(node.kwargs["dtype"]) is not None: + node_dtype = dtype_to_ir(extract_literals(node.kwargs["dtype"])) + if node_dtype is None: + node_dtype = "fp32" + + ir_node = IRNode( + id=node_id(node), + op="new_empty", + inputs=[], + outputs=[value_id(node, ctx)], + attrs={"shape": output_shape, "dtype": node_dtype}, + meta=_base_meta(shape or output_shape, node_dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape or output_shape, dtype=node_dtype) + + +def import_layer_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + if not node.args: + ctx.fail(f"unsupported layer_norm signature for {torch_op}: {node.args!r}") + + normalized_shape_value = None + if len(node.args) > 1: + normalized_shape_value = _extract_int_sequence_literal(node.args[1]) + if normalized_shape_value is None: + normalized_shape_value = _extract_int_sequence_literal(getattr(node, "kwargs", {}).get("normalized_shape")) + if normalized_shape_value is None: + ctx.fail(f"unsupported normalized_shape for {torch_op}: {node.args!r} {getattr(node, 'kwargs', {})!r}") + + eps_value = None + if len(node.args) > 4: + eps_value = _extract_numeric_literal(node.args[4]) + if eps_value is None: + eps_value = _extract_numeric_literal(getattr(node, "kwargs", {}).get("eps")) + if eps_value is None: + eps_value = 1e-5 + + inputs = [value_id(node.args[0], ctx)] + weight = node.args[2] if len(node.args) > 2 else getattr(node, "kwargs", {}).get("weight") + bias = node.args[3] if len(node.args) > 3 else getattr(node, "kwargs", {}).get("bias") + if not is_fx_node(weight): + ctx.fail(f"unsupported layer_norm weight form for {torch_op}: {node.args!r} {getattr(node, 'kwargs', {})!r}") + inputs.append(value_id(weight, ctx)) + if bias is None: + pass + elif is_fx_node(bias): + inputs.append(value_id(bias, ctx)) + else: + ctx.fail(f"unsupported layer_norm bias form for {torch_op}: {node.args!r} {getattr(node, 'kwargs', {})!r}") + + attrs = {"normalized_shape": tuple(int(v) for v in normalized_shape_value), "eps": float(eps_value)} + ir_node = IRNode( + id=node_id(node), + op="layer_norm", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_rms_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + inputs = [value_id(node.args[0], ctx), value_id(node.args[2], ctx)] + attrs = {"normalized_shape": tuple(int(v) for v in node.args[1]), "eps": float(node.args[3])} + ir_node = IRNode( + id=node_id(node), + op="rms_norm", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_group_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + inputs = [value_id(node.args[0], ctx), value_id(node.args[2], ctx), value_id(node.args[3], ctx)] + attrs = {"num_groups": int(node.args[1]), "eps": float(node.args[4])} + ir_node = IRNode( + id=node_id(node), + op="group_norm", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_batch_norm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + inputs = [ + value_id(node.args[0], ctx), + value_id(node.args[1], ctx), + value_id(node.args[2], ctx), + value_id(node.args[3], ctx), + value_id(node.args[4], ctx), + ] + attrs = {"training": bool(node.args[5]), "momentum": float(node.args[6]), "eps": float(node.args[7])} + ir_node = IRNode( + id=node_id(node), + op="batch_norm", + inputs=inputs, + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_contiguous(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + ir_node = IRNode( + id=node_id(node), + op="contiguous", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_getitem(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + index_value = extract_literals(node.args[1]) + if not isinstance(index_value, int): + ctx.fail(f"unsupported getitem index for {torch_op}: {node.args!r}") + ir_node = IRNode( + id=node_id(node), + op="getitem", + inputs=[value_id(node.args[0], ctx)], + outputs=[value_id(node, ctx)], + attrs={"index": index_value}, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def _import_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op, op, has_expert_bias, activation): + num_leading = 3 if has_expert_bias else 2 + num_args = num_leading + 8 + (1 if has_expert_bias else 0) + if len(node.args) != num_args: + ctx.fail(f"{op} expected {num_args} args, got {len(node.args)}") + w1_inputs = extract_input_ids(node.args[num_leading], ctx) + w3_inputs = extract_input_ids(node.args[num_leading + 1], ctx) + w2_inputs = extract_input_ids(node.args[num_leading + 2], ctx) + num_experts = int(extract_literals(node.args[num_leading + 3])) + num_experts_per_tok = int(extract_literals(node.args[num_leading + 4])) + if len(w1_inputs) != num_experts or len(w3_inputs) != num_experts or len(w2_inputs) != num_experts: + ctx.fail( + f"{op} expert input count mismatch: " + f"num_experts={num_experts} w1={len(w1_inputs)} w3={len(w3_inputs)} w2={len(w2_inputs)}" + ) + flags = num_leading + 5 + attrs = {"num_experts": num_experts, "num_experts_per_tok": num_experts_per_tok} + if has_expert_bias: + attrs["use_expert_bias"] = require_literal_bool(node.args[flags], ctx, "use_expert_bias") + flags += 1 + attrs["normalize_routing"] = require_literal_bool(node.args[flags], ctx, "normalize_routing") + attrs["epsilon"] = float(extract_literals(node.args[flags + 1])) + attrs["routed_scaling_factor"] = float(extract_literals(node.args[flags + 2])) + attrs["activation"] = activation + leading_inputs = [value_id(node.args[i], ctx) for i in range(num_leading)] + ir_node = IRNode( + id=node_id(node), + op=op, + inputs=[*leading_inputs, *w1_inputs, *w3_inputs, *w2_inputs], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_lfm2_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op): + _import_moe_layer_gated(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, + op="lfm2_moe_layer_gated", has_expert_bias=True, activation="silu") + + +def import_qwen2_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op): + _import_moe_layer_gated(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, + op="qwen2_moe_layer_gated", has_expert_bias=False, activation="silu") + + +def import_gemma4_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op): + _import_moe_layer_gated(ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, + op="gemma4_moe_layer_gated", has_expert_bias=False, activation="gelu") + + +OP_IMPORTERS = { + "arange": import_arange, + "identity": import_identity, + "add": import_add, + "subtract": import_subtract, + "multiply": import_multiply, + "multiply_inplace": import_multiply_inplace, + "divide": import_divide, + "precision_cast": import_precision_cast, + "type_as": import_type_as, + "negate": import_negate, + "abs": import_unary("abs"), + "clamp": import_clamp, + "not_equal": import_not_equal, + "equal": import_equal, + "greater": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( + ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="greater" + ), + "greater_equal": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( + ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="greater_equal" + ), + "less": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( + ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="less" + ), + "less_equal": lambda ir, node, ctx, *, shape, dtype, torch_op: import_compare( + ir, node, ctx, shape=shape, dtype=dtype, torch_op=torch_op, op_name="less_equal" + ), + "logical_and": import_binary("logical_and"), + "logical_or": import_binary("logical_or"), + "logical_not": import_unary("logical_not"), + "where": import_where, + "masked_scatter": import_masked_scatter, + "masked_fill": import_masked_fill, + "cos": import_unary("cos"), + "sin": import_unary("sin"), + "floor_divide": import_floor_divide, + "scalar_exp": import_unary("scalar_exp"), + "scalar_sqrt": import_unary("scalar_sqrt"), + "reciprocal": import_reciprocal, + "scalar_log": import_unary("scalar_log"), + "rsqrt": import_pow_rsqrt, + "reshape": import_reshape, + "flatten": import_flatten, + "unsqueeze": import_unsqueeze, + "squeeze": import_squeeze, + "expand": import_expand, + "repeat": import_repeat, + "one_hot": import_one_hot, + "tril": import_tril, + "unfold": import_unfold, + "numpy_T": import_numpy_t, + "transpose": import_transpose, + "permute": import_permute, + "movedim": import_movedim, + "matmul": import_matmul, + "linear": import_linear, + "addmm": import_addmm, + "relu": import_unary("relu"), + "silu": import_unary("silu"), + "gelu": import_unary("gelu"), + "gelu_erf": import_unary("gelu_erf"), + "sigmoid": import_unary("sigmoid"), + "glu": import_glu, + "softplus": import_unary("softplus"), + "tanh": import_unary("tanh"), + "softmax": import_softmax, + "scaled_dot_product_attention": import_scaled_dot_product_attention, + "sum": import_reduce("sum"), + "mean": import_reduce("mean"), + "variance": import_reduce("variance"), + "min": import_reduce("min"), + "max": import_reduce("max"), + "cat": import_cat, + "split_with_sizes": import_split_with_sizes, + "chunk": import_chunk, + "unbind": import_unbind, + "aten.unbind.int": import_unbind, + "ones": import_ones, + "new_empty": import_new_empty, + "pad": import_pad, + "pow": import_pow, + "aten.diff.default": import_diff, + "slice": import_slice, + "index": import_index, + "gather": import_gather, + "embedding": import_embedding, + "conv1d": import_conv1d, + "conv2d": import_conv2d, + "layer_norm": import_layer_norm, + "rms_norm": import_rms_norm, + "group_norm": import_group_norm, + "batch_norm": import_batch_norm, + "contiguous": import_contiguous, + "getitem": import_getitem, + "lfm2_moe_layer_gated": import_lfm2_moe_layer_gated, + "qwen2_moe_layer_gated": import_qwen2_moe_layer_gated, + "gemma4_moe_layer_gated": import_gemma4_moe_layer_gated, +} diff --git a/python/cactus/transpile/jax_semantic_rewrites.py b/python/cactus/transpile/jax_semantic_rewrites.py new file mode 100644 index 000000000..d0dd63f15 --- /dev/null +++ b/python/cactus/transpile/jax_semantic_rewrites.py @@ -0,0 +1,875 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.import_semantics import apply_import_semantics + + +@dataclass(frozen=True) +class JaxPattern: + name: str + apply: Callable[[IRGraph], None] + + +def _constant_scalar_or_singleton(graph: IRGraph, value_id: str) -> float | None: + value = graph.constants.get(value_id) + if value is None: + return None + array = np.asarray(value) + if graph.values[value_id].meta.get("jax_closed_constant") and array.size != 1: + return None + if array.size == 0: + return None + first = float(array.reshape(-1)[0].item()) + if np.allclose(array, first, rtol=0.0, atol=0.0): + return first + return None + + +def _rebuild_users(graph: IRGraph) -> None: + for value in graph.values.values(): + value.users.clear() + for node_id in graph.order: + node = graph.nodes[node_id] + for input_id in node.inputs: + graph.values[input_id].users.append(node_id) + + +def _producer(graph: IRGraph, value_id: str) -> IRNode | None: + value = graph.values.get(value_id) + if value is None or value.producer is None: + return None + return graph.nodes.get(value.producer) + + +def _strip_simple_wrappers(graph: IRGraph, value_id: str) -> str: + current = value_id + for _ in range(8): + producer = _producer(graph, current) + if producer is None or producer.op not in {"precision_cast", "reshape", "view", "expand"} or not producer.inputs: + return current + current = producer.inputs[0] + return current + + +def _trace_reduction_source(graph: IRGraph, value_id: str) -> str | None: + current = _strip_simple_wrappers(graph, value_id) + node = _producer(graph, current) + if node is not None and node.op == "divide" and len(node.inputs) == 2: + divisor = _constant_scalar_or_singleton(graph, node.inputs[1]) + if divisor is None: + return None + current = _strip_simple_wrappers(graph, node.inputs[0]) + node = _producer(graph, current) + if node is None or node.op not in {"mean", "sum"} or not node.inputs: + return None + return _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, node.inputs[0])) + + +def _trace_square_reduction_source(graph: IRGraph, value_id: str, term_trace) -> str | None: + current = _strip_simple_wrappers(graph, value_id) + node = _producer(graph, current) + if node is not None and node.op == "divide" and len(node.inputs) == 2: + divisor = _constant_scalar_or_singleton(graph, node.inputs[1]) + if divisor is None: + return None + current = _strip_simple_wrappers(graph, node.inputs[0]) + node = _producer(graph, current) + if node is None or node.op not in {"mean", "sum"} or not node.inputs: + return None + square_node = _producer(graph, _strip_simple_wrappers(graph, node.inputs[0])) + if square_node is None or square_node.op != "multiply" or len(square_node.inputs) != 2: + return None + lhs = term_trace(graph, square_node.inputs[0]) + rhs = term_trace(graph, square_node.inputs[1]) + if lhs is None or rhs is None: + return None + if _strip_simple_wrappers(graph, lhs) != _strip_simple_wrappers(graph, rhs): + return None + return _strip_simple_wrappers(graph, lhs) + + +def _trace_raw_square_source(graph: IRGraph, value_id: str) -> str | None: + return _strip_simple_wrappers(graph, _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, value_id))) + + +def _trace_mean_square_source(graph: IRGraph, value_id: str) -> str | None: + return _trace_square_reduction_source(graph, value_id, _trace_raw_square_source) + + +def _trace_layer_norm_variance_source(graph: IRGraph, value_id: str) -> str | None: + current = _strip_simple_wrappers(graph, value_id) + node = _producer(graph, current) + if node is not None and node.op == "where" and node.inputs: + current = _strip_simple_wrappers(graph, node.inputs[-1]) + node = _producer(graph, current) + if node is None or node.op != "subtract" or len(node.inputs) != 2: + return None + mean_square_source = _trace_mean_square_source(graph, node.inputs[0]) + square_mean_node = _producer(graph, _strip_simple_wrappers(graph, node.inputs[1])) + if mean_square_source is None or square_mean_node is None or square_mean_node.op != "multiply": + return None + if len(square_mean_node.inputs) != 2: + return None + lhs_mean_source = _trace_reduction_source(graph, square_mean_node.inputs[0]) + rhs_mean_source = _trace_reduction_source(graph, square_mean_node.inputs[1]) + if lhs_mean_source is None or rhs_mean_source is None: + return None + if _strip_simple_wrappers(graph, lhs_mean_source) != _strip_simple_wrappers(graph, rhs_mean_source): + return None + if _strip_simple_wrappers(graph, mean_square_source) != _strip_simple_wrappers(graph, lhs_mean_source): + return None + return mean_square_source + + +def _trace_layer_norm_centered_source(graph: IRGraph, value_id: str) -> str | None: + centered = _trace_layer_norm_centered(graph, value_id) + return None if centered is None else centered[0] + + +def _trace_layer_norm_centered(graph: IRGraph, value_id: str) -> tuple[str, str] | None: + centered_id = _strip_simple_wrappers(graph, value_id) + node = _producer(graph, centered_id) + if node is None or node.op != "subtract" or len(node.inputs) != 2: + return None + source_id = _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, node.inputs[0])) + mean_source_id = _trace_reduction_source(graph, node.inputs[1]) + if mean_source_id is None: + return None + if _strip_simple_wrappers(graph, source_id) != _strip_simple_wrappers(graph, mean_source_id): + return None + return _strip_simple_wrappers(graph, source_id), centered_id + + +def _trace_inv_std_addend(graph: IRGraph, value_id: str) -> str | None: + current = _strip_simple_wrappers(graph, value_id) + node = _producer(graph, current) + if node is None: + return None + if node.op == "pow" and node.inputs: + exponent = float(node.attrs.get("exponent", 0.0)) + if exponent == -0.5: + return node.inputs[0] + if exponent == -1.0: + return _trace_inv_std_addend(graph, node.inputs[0]) + if node.op == "scalar_sqrt" and node.inputs: + return node.inputs[0] + if node.op == "divide" and len(node.inputs) == 2: + numerator = _constant_scalar_or_singleton(graph, node.inputs[0]) + if numerator == 1.0: + return _trace_inv_std_addend(graph, node.inputs[1]) or node.inputs[1] + return _trace_inv_std_addend(graph, node.inputs[1]) + return None + + +def _trace_norm_inv_std( + graph: IRGraph, + value_id: str, + default_eps: float, + variance_trace, +) -> tuple[str, float] | None: + add_id = _trace_inv_std_addend(graph, value_id) + add_node = _producer(graph, _strip_simple_wrappers(graph, add_id)) if add_id is not None else None + if add_node is None or add_node.op != "add" or len(add_node.inputs) != 2: + return None + + eps = default_eps + source_id: str | None = None + for add_input in add_node.inputs: + scalar = _constant_scalar_or_singleton(graph, add_input) + if scalar is not None: + eps = float(scalar) + continue + variance_source = variance_trace(graph, add_input) + if variance_source is not None: + source_id = variance_source + if source_id is None: + return None + return source_id, eps + + +def _trace_rms_inv_std(graph: IRGraph, value_id: str) -> tuple[str, float] | None: + return _trace_norm_inv_std(graph, value_id, 1.0e-6, _trace_mean_square_source) + + +def _trace_layer_norm_inv_std(graph: IRGraph, value_id: str) -> tuple[str, float] | None: + def _variance_source(g: IRGraph, candidate_id: str) -> str | None: + return _trace_layer_norm_variance_source(g, candidate_id) or _trace_square_reduction_source( + g, + candidate_id, + _trace_layer_norm_centered_source, + ) + + return _trace_norm_inv_std(graph, value_id, 1.0e-5, _variance_source) + + +def _valid_last_dim_weight(graph: IRGraph, weight_id: str, source_id: str) -> str | None: + weight_id = _strip_simple_wrappers(graph, weight_id) + weight_value = graph.values.get(weight_id) + source_value = graph.values.get(source_id) + if weight_value is None or source_value is None: + return None + if weight_value.shape is None or source_value.shape is None: + return None + if len(weight_value.shape) != 1 or int(weight_value.shape[0]) != int(source_value.shape[-1]): + return None + return weight_id + + +def _rewrite_norm_node(node: IRNode, op: str, inputs: list[str], eps: float, rewritten_from: str) -> None: + node.op = op + node.inputs = inputs + node.attrs = {"eps": eps} + node.kind = "semantic" + node.meta = {**node.meta, "rewritten_from": rewritten_from} + + +def _trace_norm_branch(graph: IRGraph, value_id: str, source_trace, inv_trace) -> tuple[str, float] | None: + node = _producer(graph, _strip_simple_wrappers(graph, value_id)) + if node is None or len(node.inputs) != 2: + return None + pairs = ((node.inputs[0], node.inputs[1]),) if node.op == "divide" else () + if node.op == "multiply": + pairs = ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])) + for source_candidate, inv_candidate in pairs: + source_id = source_trace(graph, source_candidate) + inv = inv_trace(graph, inv_candidate) + if source_id is None or inv is None: + continue + inv_source_id, eps = inv + if _strip_simple_wrappers(graph, source_id) == _strip_simple_wrappers(graph, inv_source_id): + return source_id, eps + return None + + +def _trace_weighted_norm(graph: IRGraph, value_id: str, branch_trace) -> tuple[str, str, float] | None: + node = _producer(graph, _strip_simple_wrappers(graph, value_id)) + if node is None or node.op != "multiply" or len(node.inputs) != 2: + return None + for norm_id, weight_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + traced = branch_trace(graph, norm_id) + if traced is None: + continue + source_id, eps = traced + weight_id = _valid_last_dim_weight(graph, weight_id, source_id) + if weight_id is None: + continue + return source_id, weight_id, eps + return None + + +def _trace_factored_weighted_norm( + graph: IRGraph, + value_id: str, + source_trace, + inv_trace, +) -> tuple[str, str, float] | None: + node = _producer(graph, _strip_simple_wrappers(graph, value_id)) + if node is None or node.op != "multiply" or len(node.inputs) != 2: + return None + for source_candidate, scaled_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + source_id = source_trace(graph, source_candidate) + scaled_node = _producer(graph, _strip_simple_wrappers(graph, scaled_id)) + if source_id is None or scaled_node is None or scaled_node.op != "multiply" or len(scaled_node.inputs) != 2: + continue + for inv_id, weight_id in ((scaled_node.inputs[0], scaled_node.inputs[1]), (scaled_node.inputs[1], scaled_node.inputs[0])): + inv = inv_trace(graph, inv_id) + if inv is None: + continue + inv_source_id, eps = inv + if _strip_simple_wrappers(graph, source_id) != _strip_simple_wrappers(graph, inv_source_id): + continue + weight_id = _valid_last_dim_weight(graph, weight_id, source_id) + if weight_id is not None: + return source_id, weight_id, eps + return None + + +def _trace_weighted_rms_norm(graph: IRGraph, value_id: str) -> tuple[str, str, float] | None: + traced = _trace_weighted_norm( + graph, + value_id, + lambda g, candidate: _trace_norm_branch(g, candidate, _trace_raw_square_source, _trace_rms_inv_std), + ) + if traced is not None: + source_id, _, _ = traced + source_producer = _producer(graph, source_id) + if source_producer is None or source_producer.op != "subtract": + return traced + return _trace_factored_weighted_norm(graph, value_id, _trace_raw_square_source, _trace_rms_inv_std) + + +def _trace_weighted_layer_norm(graph: IRGraph, value_id: str) -> tuple[str, str, float] | None: + traced = _trace_weighted_norm( + graph, + value_id, + lambda g, candidate: _trace_norm_branch(g, candidate, _trace_layer_norm_centered_source, _trace_layer_norm_inv_std), + ) + if traced is not None: + return traced + return _trace_factored_weighted_norm( + graph, + value_id, + _trace_layer_norm_centered_source, + _trace_layer_norm_inv_std, + ) + + +def _rewrite_jax_rms_norms(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op == "multiply" and len(node.inputs) == 2: + traced = _trace_weighted_rms_norm(graph, node.outputs[0] if node.outputs else "") + if traced is not None: + source_id, weight_id, eps = traced + _rewrite_norm_node(node, "rms_norm", [source_id, weight_id], eps, "jax_rms_norm_multiply") + continue + if node.op == "divide" and len(node.inputs) == 2: + denominator = _trace_rms_inv_std(graph, node.inputs[1]) + if denominator is None: + continue + source_id, eps = denominator + numerator = _producer(graph, _strip_simple_wrappers(graph, node.inputs[0])) + if numerator is None or numerator.op != "multiply" or len(numerator.inputs) != 2: + continue + for x_id, weight_id in ((numerator.inputs[0], numerator.inputs[1]), (numerator.inputs[1], numerator.inputs[0])): + if _strip_simple_wrappers(graph, x_id) != _strip_simple_wrappers(graph, source_id): + continue + weight_id = _valid_last_dim_weight(graph, weight_id, x_id) + if weight_id is not None: + _rewrite_norm_node(node, "rms_norm", [x_id, weight_id], eps, "jax_rms_norm") + break + + +def _rewrite_jax_layer_norms(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op == "add" and len(node.inputs) == 2: + for affine_id, bias_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + bias_id = _strip_simple_wrappers(graph, bias_id) + bias_value = graph.values.get(bias_id) + if bias_value is None or bias_value.shape is None or len(bias_value.shape) != 1: + continue + traced_affine = _trace_weighted_layer_norm(graph, affine_id) + if traced_affine is None: + continue + source_id, weight_id, eps = traced_affine + source_value = graph.values.get(source_id) + if source_value is None or source_value.shape is None: + continue + if int(bias_value.shape[0]) != int(source_value.shape[-1]): + continue + _rewrite_norm_node(node, "layer_norm", [source_id, weight_id, bias_id], eps, "jax_layer_norm") + break + continue + if node.op == "multiply" and len(node.inputs) == 2: + traced = _trace_weighted_layer_norm(graph, node.outputs[0] if node.outputs else "") + if traced is not None: + source_id, weight_id, eps = traced + _rewrite_norm_node(node, "layer_norm", [source_id, weight_id], eps, "jax_layer_norm_multiply") + + +def _rewrite_jax_silus(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op != "multiply" or len(node.inputs) != 2: + continue + for source_id, sigmoid_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + sigmoid_node = _producer(graph, _strip_simple_wrappers(graph, sigmoid_id)) + if sigmoid_node is None or sigmoid_node.op not in {"sigmoid", "logistic"} or len(sigmoid_node.inputs) != 1: + continue + if _strip_simple_wrappers(graph, source_id) != _strip_simple_wrappers(graph, sigmoid_node.inputs[0]): + continue + node.op = "silu" + node.inputs = [source_id] + node.attrs = {} + node.kind = "semantic" + node.meta = {**node.meta, "rewritten_from": "jax_silu"} + break + + +def _trace_batch_norm_centered(graph: IRGraph, value_id: str) -> tuple[str, str] | None: + node = _producer(graph, _strip_simple_wrappers(graph, value_id)) + if node is None or node.op != "subtract" or len(node.inputs) != 2: + return None + source_id = _trace_precision_cast_source(graph, _strip_simple_wrappers(graph, node.inputs[0])) + mean_id = _strip_simple_wrappers(graph, node.inputs[1]) + source_value = graph.values.get(source_id) + mean_value = graph.values.get(mean_id) + if source_value is None or source_value.shape is None: + return None + if mean_value is None or mean_value.shape is None or len(mean_value.shape) != 1: + return None + if int(source_value.shape[-1]) != int(mean_value.shape[0]): + return None + return source_id, mean_id + + +def _trace_batch_norm_inv_scale(graph: IRGraph, value_id: str) -> tuple[str, str, float] | None: + node = _producer(graph, _strip_simple_wrappers(graph, value_id)) + if node is None or node.op != "multiply" or len(node.inputs) != 2: + return None + for inv_id, scale_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + scale_id = _strip_simple_wrappers(graph, scale_id) + scale_value = graph.values.get(scale_id) + if scale_value is None or scale_value.shape is None or len(scale_value.shape) != 1: + continue + inv_node = _producer(graph, _strip_simple_wrappers(graph, inv_id)) + if inv_node is None or inv_node.op not in {"pow", "divide"} or not inv_node.inputs: + continue + if inv_node.op == "pow" and float(inv_node.attrs.get("exponent", 0.0)) != -0.5: + continue + add_id = inv_node.inputs[0] if inv_node.op == "pow" else inv_node.inputs[1] + add_node = _producer(graph, _strip_simple_wrappers(graph, add_id)) + if add_node is None or add_node.op != "add" or len(add_node.inputs) != 2: + continue + eps = 1.0e-5 + var_id: str | None = None + for add_input in add_node.inputs: + scalar = _constant_scalar_or_singleton(graph, add_input) + if scalar is not None: + eps = float(scalar) + continue + candidate = _strip_simple_wrappers(graph, add_input) + candidate_value = graph.values.get(candidate) + if candidate_value is not None and candidate_value.shape is not None and len(candidate_value.shape) == 1: + var_id = candidate + if var_id is None: + continue + var_value = graph.values.get(var_id) + if var_value is None or var_value.shape is None or int(var_value.shape[0]) != int(scale_value.shape[0]): + continue + return var_id, scale_id, eps + return None + + +def _trace_batch_norm_affine(graph: IRGraph, value_id: str) -> tuple[str, str, str, str, float] | None: + node = _producer(graph, _strip_simple_wrappers(graph, value_id)) + if node is None or node.op != "multiply" or len(node.inputs) != 2: + return None + for centered_id, inv_scale_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + centered = _trace_batch_norm_centered(graph, centered_id) + inv_scale = _trace_batch_norm_inv_scale(graph, inv_scale_id) + if centered is None or inv_scale is None: + continue + source_id, mean_id = centered + var_id, scale_id, eps = inv_scale + channel_count = int(graph.values[scale_id].shape[0]) # type: ignore[index] + if int(graph.values[mean_id].shape[0]) != channel_count: # type: ignore[index] + continue + if int(graph.values[var_id].shape[0]) != channel_count: # type: ignore[index] + continue + return source_id, scale_id, mean_id, var_id, eps + return None + + +def _rewrite_jax_batch_norms(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op != "add" or len(node.inputs) != 2: + continue + for affine_id, bias_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + bias_id = _strip_simple_wrappers(graph, bias_id) + bias_value = graph.values.get(bias_id) + if bias_value is None or bias_value.shape is None or len(bias_value.shape) != 1: + continue + traced = _trace_batch_norm_affine(graph, affine_id) + if traced is None: + continue + source_id, scale_id, mean_id, var_id, eps = traced + source_value = graph.values.get(source_id) + if source_value is None or source_value.shape is None or not source_value.shape: + continue + channel_count = int(bias_value.shape[0]) + if int(source_value.shape[-1]) != channel_count: + continue + node.op = "batch_norm" + node.inputs = [source_id, scale_id, bias_id, mean_id, var_id] + node.attrs = {"axis": len(source_value.shape) - 1, "eps": eps} + node.kind = "semantic" + node.meta = {**node.meta, "rewritten_from": "jax_batch_norm"} + break + + +def _branch_uses_rms_norm_of(graph: IRGraph, branch_value_id: str, residual_value_id: str) -> bool: + residual_base = _strip_simple_wrappers(graph, residual_value_id) + stack = [branch_value_id] + visited: set[str] = set() + while stack and len(visited) < 512: + current = stack.pop() + if current in visited: + continue + visited.add(current) + node = _producer(graph, _strip_simple_wrappers(graph, current)) + if node is None: + continue + if node.op == "rms_norm" and node.inputs: + if _strip_simple_wrappers(graph, node.inputs[0]) == residual_base: + return True + stack.extend(node.inputs) + return False + + +def _rewrite_prenorm_residual_adds_to_clipped(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op != "add" or len(node.inputs) != 2: + continue + lhs, rhs = node.inputs + if _branch_uses_rms_norm_of(graph, rhs, lhs) or _branch_uses_rms_norm_of(graph, lhs, rhs): + node.op = "add_clipped" + node.kind = "semantic" + node.meta = {**node.meta, "rewritten_from": "prenorm_residual_add"} + + +def _shape(graph: IRGraph, value_id: str) -> tuple[int, ...] | None: + value = graph.values.get(value_id) + if value is None or value.shape is None: + return None + return tuple(int(dim) for dim in value.shape) + + +def _trace_view_input(graph: IRGraph, value_id: str) -> str: + current = value_id + for _ in range(8): + node = _producer(graph, current) + if node is None or node.op not in {"view", "reshape", "precision_cast"} or not node.inputs: + return current + current = node.inputs[0] + return current + + +def _trace_gqa_repeat_base(graph: IRGraph, value_id: str) -> str: + current = _trace_view_input(graph, value_id) + shape = _shape(graph, current) + if shape is None or len(shape) != 5: + return current + expand = _producer(graph, current) + if expand is None or expand.op != "expand" or not expand.inputs: + return current + base = _trace_view_input(graph, expand.inputs[0]) + base_shape = _shape(graph, base) + if base_shape is None or len(base_shape) != 4: + return current + if base_shape[0] == shape[0] and base_shape[1] == shape[1] and base_shape[2] == shape[3] and base_shape[3] == shape[4]: + return base + return current + + +def _trace_jax_attention_logits(graph: IRGraph, value_id: str) -> tuple[str, str, float, str | None, bool] | None: + current = _trace_view_input(graph, value_id) + node = _producer(graph, current) + mask_id: str | None = None + additive_mask = False + if node is not None and node.op == "where" and len(node.inputs) >= 2: + mask_id = node.inputs[0] + current = _trace_view_input(graph, node.inputs[1]) + node = _producer(graph, current) + if node is not None and node.op == "add" and len(node.inputs) == 2: + for score_id, additive_mask_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + score_node = _producer(graph, _trace_view_input(graph, score_id)) + if score_node is not None and score_node.op == "divide": + mask_id = additive_mask_id + additive_mask = True + current = _trace_view_input(graph, score_id) + node = score_node + break + if node is None or node.op != "divide" or len(node.inputs) != 2: + return None + divisor = _constant_scalar_or_singleton(graph, node.inputs[1]) + if divisor is None or divisor == 0.0: + return None + matmul = _producer(graph, _trace_view_input(graph, node.inputs[0])) + if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: + return None + query_id = _trace_view_input(graph, matmul.inputs[0]) + key_t_id = _trace_view_input(graph, matmul.inputs[1]) + key_t_node = _producer(graph, key_t_id) + if key_t_node is None or key_t_node.op != "permute" or len(key_t_node.inputs) != 1: + return None + if tuple(int(dim) for dim in key_t_node.attrs.get("permutation", ())) != (0, 1, 3, 2): + return None + key_id = _trace_gqa_repeat_base(graph, key_t_node.inputs[0]) + return query_id, key_id, 1.0 / float(divisor), mask_id, additive_mask + + +def _trace_jax_attention_probs(graph: IRGraph, value_id: str) -> tuple[str, str, float, str | None, bool] | None: + current = _trace_view_input(graph, value_id) + cast = _producer(graph, current) + if cast is not None and cast.op == "precision_cast" and cast.inputs: + current = _trace_view_input(graph, cast.inputs[0]) + divide = _producer(graph, current) + if divide is None or divide.op != "divide" or len(divide.inputs) != 2: + return None + exp_node = _producer(graph, _trace_view_input(graph, divide.inputs[0])) + if exp_node is None or exp_node.op != "scalar_exp" or not exp_node.inputs: + return None + subtract = _producer(graph, _trace_view_input(graph, exp_node.inputs[0])) + if subtract is None or subtract.op != "subtract" or not subtract.inputs: + return None + return _trace_jax_attention_logits(graph, subtract.inputs[0]) + + +def _rewrite_attention_node( + node: IRNode, + query_id: str, + key_id: str, + value_id: str, + mask_id: str | None, + *, + additive_mask: bool, + scale: float, + output_layout: str, + is_causal: bool, + rewritten_from: str, +) -> None: + node.op = "attention" + node.inputs = [query_id, key_id, value_id] + ([] if mask_id is None else [mask_id]) + node.attrs = { + "scale": float(scale), + "is_causal": bool(is_causal), + "window_size": 0, + "q_layout": "bhsd", + "k_layout": "bhsd", + "v_layout": "bhsd", + "output_layout": output_layout, + **({"additive_mask": True} if additive_mask else {}), + } + node.kind = "semantic" + node.meta = {**node.meta, "rewritten_from": rewritten_from} + + +def _trace_attention_output(graph: IRGraph, node: IRNode) -> tuple[str, str, str, float, str | None, bool] | None: + if len(node.inputs) != 1 or len(node.outputs) != 1: + return None + matmul = _producer(graph, _trace_view_input(graph, node.inputs[0])) + if matmul is None or matmul.op != "matmul" or len(matmul.inputs) != 2: + return None + traced = _trace_jax_attention_probs(graph, matmul.inputs[0]) + if traced is None: + return None + query_id, key_id, scale, mask_id, additive_mask = traced + return query_id, key_id, _trace_gqa_repeat_base(graph, matmul.inputs[1]), scale, mask_id, additive_mask + + +def _valid_attention_qkv_shapes( + query_shape: tuple[int, ...] | None, + key_shape: tuple[int, ...] | None, + value_shape: tuple[int, ...] | None, +) -> bool: + if query_shape is None or key_shape is None or value_shape is None: + return False + if len(query_shape) != 4 or len(key_shape) != 4 or len(value_shape) != 4: + return False + if query_shape[0] != key_shape[0] or key_shape[0] != value_shape[0]: + return False + return key_shape[1:] == value_shape[1:] + + +def _valid_attention_mask_shape( + mask_shape: tuple[int, ...] | None, + query_shape: tuple[int, ...], + key_shape: tuple[int, ...], +) -> bool: + return mask_shape == (query_shape[0], query_shape[1], query_shape[2], key_shape[2]) + + +def _rewrite_jax_prefill_attentions(graph: IRGraph) -> None: + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op not in {"permute", "view", "reshape"}: + continue + component = str(graph.meta.get("component", "") or "").strip().lower() + is_decoder_step = component == "decoder_step" + is_causal_prefill = component in {"decoder_prefill_chunk", "decoder_media_step"} + if node.op == "permute" and tuple(int(dim) for dim in node.attrs.get("permutation", ())) != (0, 2, 1, 3): + continue + if node.op in {"view", "reshape"} and not is_decoder_step: + continue + output_shape = _shape(graph, node.outputs[0]) + if output_shape is None or len(output_shape) != 4: + continue + traced = _trace_attention_output(graph, node) + if traced is None: + continue + query_id, key_id, value_id, scale, mask_id, additive_mask = traced + query_shape = _shape(graph, query_id) + key_shape = _shape(graph, key_id) + value_shape = _shape(graph, value_id) + if not _valid_attention_qkv_shapes(query_shape, key_shape, value_shape): + continue + assert query_shape is not None and key_shape is not None + if mask_id is not None and not _valid_attention_mask_shape(_shape(graph, mask_id), query_shape, key_shape): + continue + is_self_attention = query_shape[2] == key_shape[2] + if is_decoder_step: + if query_shape[2] != 1 or (not is_self_attention and mask_id is None): + continue + _rewrite_attention_node( + node, + query_id, + key_id, + value_id, + mask_id, + additive_mask=additive_mask, + scale=scale, + output_layout="bhsd", + is_causal=is_self_attention, + rewritten_from="jax_decoder_step_self_attention" if is_self_attention else "jax_decoder_step_cross_attention", + ) + continue + if output_shape != (query_shape[0], query_shape[2], query_shape[1], query_shape[3]): + continue + _rewrite_attention_node( + node, + query_id, + key_id, + value_id, + mask_id, + additive_mask=additive_mask, + scale=scale, + output_layout="bthd", + is_causal=is_causal_prefill and mask_id is None and is_self_attention, + rewritten_from="jax_prefill_attention", + ) + + +def _trace_precision_cast_source(graph: IRGraph, value_id: str) -> str: + producer = _producer(graph, value_id) + if producer is not None and producer.op == "precision_cast" and producer.inputs: + return producer.inputs[0] + return value_id + + +def _prune_dead_nodes(graph: IRGraph) -> None: + live_values = set(graph.outputs) + live_nodes: set[str] = set() + stack = list(graph.outputs) + while stack: + value_id = stack.pop() + value = graph.values.get(value_id) + if value is None or value.producer is None: + continue + node = graph.nodes.get(value.producer) + if node is None or node.id in live_nodes: + continue + live_nodes.add(node.id) + for input_id in node.inputs: + if input_id not in live_values: + live_values.add(input_id) + stack.append(input_id) + + removed_nodes = set(graph.nodes) - live_nodes + if not removed_nodes: + return + for node_id in removed_nodes: + node = graph.nodes.pop(node_id) + for output_id in node.outputs: + graph.values.pop(output_id, None) + graph.constants.pop(output_id, None) + graph.order = [node_id for node_id in graph.order if node_id in live_nodes] + + +def _unique_graph_id(existing: set[str], base: str) -> str: + candidate = base + suffix = 0 + while candidate in existing: + suffix += 1 + candidate = f"{base}_{suffix}" + return candidate + + +def _legalize_remaining_jax_reductions(graph: IRGraph) -> None: + new_order: list[str] = [] + for node_id in list(graph.order): + node = graph.nodes[node_id] + if node.op not in {"sum", "mean"} or len(node.inputs) != 1 or len(node.outputs) != 1: + new_order.append(node_id) + continue + input_value = graph.values.get(node.inputs[0]) + output_value = graph.values.get(node.outputs[0]) + if input_value is None or output_value is None or input_value.dtype == "fp16": + new_order.append(node_id) + continue + + cast_input_id = _unique_graph_id(set(graph.values), f"{node.outputs[0]}__reduce_fp16_input") + cast_input_node_id = _unique_graph_id(set(graph.nodes), f"{node.id}_reduce_input_cast") + graph.add_node( + IRNode( + cast_input_node_id, + "precision_cast", + [node.inputs[0]], + [cast_input_id], + attrs={"dtype": "fp16"}, + meta={ + "jax_reduce_legalization": "input_fp16", + "source_dtype": input_value.dtype, + }, + ) + ) + graph.values[cast_input_id].shape = input_value.shape + graph.values[cast_input_id].dtype = "fp16" + node.inputs = [cast_input_id] + new_order.append(cast_input_node_id) + new_order.append(node_id) + + if output_value.dtype == "fp16": + continue + + original_output_id = node.outputs[0] + reduce_output_id = _unique_graph_id(set(graph.values), f"{original_output_id}__reduce_fp16") + node.outputs = [reduce_output_id] + graph.values[reduce_output_id] = IRValue( + id=reduce_output_id, + shape=output_value.shape, + dtype="fp16", + producer=node.id, + meta={ + "jax_reduce_legalization": "fp16_output", + "source_output": original_output_id, + }, + ) + cast_output_node_id = _unique_graph_id(set(graph.nodes), f"{node.id}_reduce_output_cast") + graph.nodes[cast_output_node_id] = IRNode( + cast_output_node_id, + "precision_cast", + [reduce_output_id], + [original_output_id], + attrs={"dtype": output_value.dtype}, + meta={ + "jax_reduce_legalization": "restore_output_dtype", + "source_dtype": "fp16", + }, + ) + output_value.producer = cast_output_node_id + new_order.append(cast_output_node_id) + graph.order = new_order + + +JAX_PATTERNS = ( + JaxPattern("batch_norm", _rewrite_jax_batch_norms), + JaxPattern("layer_norm", _rewrite_jax_layer_norms), + JaxPattern("rms_norm", _rewrite_jax_rms_norms), + JaxPattern("silu", _rewrite_jax_silus), + JaxPattern("attention", _rewrite_jax_prefill_attentions), + JaxPattern("prenorm_add_clipped", _rewrite_prenorm_residual_adds_to_clipped), +) + + +def _apply_pattern(graph: IRGraph, pattern: JaxPattern) -> None: + pattern.apply(graph) + + +def apply_jax_semantic_rewrites(graph: IRGraph) -> None: + for pattern in JAX_PATTERNS: + _apply_pattern(graph, pattern) + apply_import_semantics(graph) + _prune_dead_nodes(graph) + _legalize_remaining_jax_reductions(graph) + _rebuild_users(graph) diff --git a/python/cactus/transpile/jax_user_graph_bundle.py b/python/cactus/transpile/jax_user_graph_bundle.py new file mode 100644 index 000000000..b06dc3657 --- /dev/null +++ b/python/cactus/transpile/jax_user_graph_bundle.py @@ -0,0 +1,575 @@ +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from dataclasses import fields +import json +import os +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header +from cactus.transpile.capture_jax import capture_jax_graphs +from cactus.transpile.capture_jax import JaxGraphSpec +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.lower import TranspiledGraph +from cactus.transpile.runtime_compat import Graph +from cactus.transpile.runtime_compat import Tensor + + +@dataclass(frozen=True) +class JaxUserGraphBundleResult: + bundle: Any + output_dir: Path + components_manifest_path: Path + weights_dir: Path + + +@dataclass +class LoadedJaxUserGraph: + name: str + graph: Graph + runtime_inputs: list[Tensor] + outputs: list[Tensor] + logical_inputs: list[str] + logical_outputs: list[str] + + def set_inputs(self, inputs: Sequence[Any]) -> None: + if len(inputs) != len(self.runtime_inputs): + raise ValueError( + f"graph {self.name!r} expected {len(self.runtime_inputs)} inputs, got {len(inputs)}" + ) + for tensor, value in zip(self.runtime_inputs, inputs, strict=True): + if isinstance(value, Tensor): + if tuple(int(dim) for dim in value.shape) != tuple(int(dim) for dim in tensor.shape): + raise ValueError( + f"graph {self.name!r} input shape mismatch for node {tensor.id}: " + f"expected {tensor.shape}, got {value.shape}" + ) + if int(value.dtype) != int(tensor.dtype): + raise ValueError( + f"graph {self.name!r} input dtype mismatch for node {tensor.id}: " + f"expected {tensor.dtype}, got {value.dtype}" + ) + self.graph.set_input(tensor, value.numpy()) + else: + self.graph.set_input(tensor, np.asarray(value)) + + def execute(self, *inputs: Any) -> list[Tensor]: + self.set_inputs(inputs) + self.graph.execute() + return self.outputs + + def reset(self) -> None: + self.graph.hard_reset() + + +@dataclass +class LoadedJaxUserGraphBundle: + root: Path + manifest_path: Path + manifest: dict[str, object] + graphs: dict[str, LoadedJaxUserGraph] + + def execute(self, graph_name: str, *inputs: Any) -> list[Tensor]: + if graph_name not in self.graphs: + available = ", ".join(sorted(self.graphs)) or "" + raise ValueError(f"unknown JAX user graph {graph_name!r}; available graphs: {available}") + return self.graphs[graph_name].execute(*inputs) + + def reset(self, graph_name: str | None = None) -> None: + if graph_name is not None: + if graph_name not in self.graphs: + available = ", ".join(sorted(self.graphs)) or "" + raise ValueError(f"unknown JAX user graph {graph_name!r}; available graphs: {available}") + self.graphs[graph_name].reset() + return + for graph in self.graphs.values(): + graph.reset() + + +def flatten_jax_params(params: object) -> dict[str, np.ndarray]: + import jax + + return { + _tree_path_name(path): np.asarray(value) + for path, value in jax.tree_util.tree_flatten_with_path(params)[0] + } + + +def _tree_path_name(path: Sequence[Any]) -> str: + parts: list[str] = [] + for entry in path: + if hasattr(entry, "key"): + part = str(entry.key) + elif hasattr(entry, "name"): + part = str(entry.name) + elif hasattr(entry, "idx"): + part = str(entry.idx) + else: + part = str(entry) + parts.append(part) + return ".".join(parts) or "param" + + +def write_fp16_weights_manifest( + weights_dir: str | Path, + params: dict[str, object], + *, + exclude: set[str] | None = None, +) -> Path: + weights_root = Path(weights_dir) + weights_root.mkdir(parents=True, exist_ok=True) + excluded = exclude or set() + manifest: dict[str, dict[str, str]] = {} + for name, value in params.items(): + if name in excluded: + continue + filename = f"{name}.weights" + save_tensor_with_header(np.asarray(value), weights_root / filename, precision="FP16") + manifest[name] = {"filename": filename, "kind": "weight"} + manifest_path = weights_root / "weights_manifest.json" + _atomic_write_text(manifest_path, json.dumps(manifest, indent=2) + "\n") + return manifest_path + + +def load_jax_user_graph_bundle(bundle_dir_or_manifest: str | Path) -> LoadedJaxUserGraphBundle: + manifest_path = _resolve_components_manifest(bundle_dir_or_manifest) + root = manifest_path.parent.parent + manifest = json.loads(manifest_path.read_text()) + weights_dir = _resolve_manifest_path(root, manifest.get("weights_dir")) + graphs: dict[str, LoadedJaxUserGraph] = {} + for component in manifest.get("components", []): + if not isinstance(component, dict): + continue + name = str(component.get("component", "") or "") + graph_relpath = component.get("graph") + if not name or not isinstance(graph_relpath, str) or not graph_relpath: + continue + graph = Graph.load(_resolve_manifest_path(root, graph_relpath)) + for binding in component.get("bound_constant_bindings", []) or []: + if not isinstance(binding, dict): + continue + node_id = binding.get("node_id") + path = binding.get("path") + if not isinstance(node_id, int) or not isinstance(path, str): + continue + graph.bind_mmap_weights( + graph._tensor_from_node(int(node_id)), + _resolve_weight_path(root, weights_dir, path), + ) + graphs[name] = LoadedJaxUserGraph( + name=name, + graph=graph, + runtime_inputs=[ + graph._tensor_from_node(int(node_id)) + for node_id in component.get("runtime_input_node_ids", []) or [] + ], + outputs=[ + graph._tensor_from_node(int(node_id)) + for node_id in component.get("output_node_ids", []) or [] + ], + logical_inputs=[str(value) for value in component.get("logical_inputs", []) or []], + logical_outputs=[str(value) for value in component.get("logical_outputs", []) or []], + ) + return LoadedJaxUserGraphBundle( + root=root, + manifest_path=manifest_path, + manifest=manifest, + graphs=graphs, + ) + + +def _resolve_components_manifest(bundle_dir_or_manifest: str | Path) -> Path: + path = Path(bundle_dir_or_manifest).expanduser() + if path.is_dir(): + components_manifest = path / "components" / "manifest.json" + if components_manifest.exists(): + return components_manifest + if path.name == "components" and (path / "manifest.json").exists(): + return path / "manifest.json" + if path.exists(): + return path + raise FileNotFoundError(f"JAX user graph bundle manifest not found: {path}") + + +def _reject_relative_escape(value: str, *, field: str) -> None: + """Reject relative manifest entries that try to escape via `..`. Absolute + paths are a legitimate workflow (shared weights cache) and stay allowed.""" + if not value: + return + p = Path(value) + if p.is_absolute(): + return + if any(part == ".." for part in p.parts): + raise RuntimeError( + f"refusing to follow bundle manifest relative path with '..' segments: " + f"{field}={value}" + ) + + +def _resolve_manifest_path(root: Path, value: object) -> str: + if not isinstance(value, str) or not value: + return str(root) + _reject_relative_escape(value, field="manifest path") + path = Path(value).expanduser() + return str(path if path.is_absolute() else root / path) + + +def _resolve_weight_path(root: Path, weights_dir: str, value: str) -> str: + _reject_relative_escape(value, field="weight path") + path = Path(value).expanduser() + if path.is_absolute(): + return str(path) + candidate = Path(weights_dir) / path + if candidate.exists(): + return str(candidate) + return str(root / path) + + +def _embed_materialized_bound_constants(transpiled_graph: TranspiledGraph) -> None: + external_node_ids = { + int(binding["node_id"]) + for binding in transpiled_graph.bound_constant_bindings + } + for constant in transpiled_graph.bound_constants: + if int(constant.id) in external_node_ids: + continue + transpiled_graph.graph.mark_embedded_input(constant) + + +def _serialize_json_compatible(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): _serialize_json_compatible(inner) for key, inner in value.items()} + if isinstance(value, (list, tuple)): + return [_serialize_json_compatible(inner) for inner in value] + if isinstance(value, Path): + return str(value) + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, torch.Tensor): + payload: dict[str, object] = { + "type": "torch.Tensor", + "dtype": str(value.dtype), + "shape": list(value.shape), + } + if value.numel() == 1: + payload["data"] = value.detach().cpu().reshape(-1)[0].item() + return payload + if isinstance(value, np.ndarray): + payload = { + "type": "numpy.ndarray", + "dtype": str(value.dtype), + "shape": list(value.shape), + } + if value.size == 1: + payload["data"] = np.asarray(value).reshape(-1)[0].item() + return payload + if hasattr(value, "__dataclass_fields__"): + return { + field.name: _serialize_json_compatible(getattr(value, field.name)) + for field in fields(value) + } + try: + return repr(value) + except Exception: + return f"<{type(value).__module__}.{type(value).__name__}>" + + +def _ir_graph_to_dict(graph: IRGraph) -> dict[str, object]: + return { + "meta": _serialize_json_compatible(graph.meta), + "inputs": list(graph.inputs), + "outputs": list(graph.outputs), + "constants": { + value_id: _serialize_json_compatible(constant) + for value_id, constant in graph.constants.items() + }, + "values": { + value_id: _serialize_json_compatible(value) + for value_id, value in graph.values.items() + }, + "nodes": [ + _serialize_json_compatible(graph.nodes[node_id]) + for node_id in graph.order + ], + } + + +def _atomic_write_text(path: Path, text: str) -> None: + """Write text to path via tempfile + os.replace so a crash mid-write + cannot leave a half-written file the next loader would choke on.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(prefix=path.name + ".", suffix=".tmp", dir=str(path.parent)) + try: + with os.fdopen(fd, "w") as f: + f.write(text) + os.replace(tmp, path) + except Exception: + try: + os.unlink(tmp) + except OSError: + pass + raise + + +def _write_json(path: Path, payload: dict[str, object]) -> None: + _atomic_write_text(path, json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def build_jax_user_graph_bundle( + *, + params: Any, + specs: Sequence[JaxGraphSpec], + output_dir: str | Path, + model_id: str, + task: str = "generic", + family: str = "jax_user_graph", + inputs_metadata: dict[str, object] | None = None, + graph_meta: dict[str, object] | None = None, + weight_arrays: dict[str, object] | None = None, + exclude_weights: set[str] | None = None, + weights_dir: str | Path | None = None, +) -> JaxUserGraphBundleResult: + arrays = weight_arrays if weight_arrays is not None else flatten_jax_params(params) + root = Path(output_dir) + weights_root = Path(weights_dir) if weights_dir is not None else root + write_fp16_weights_manifest(weights_root, arrays, exclude=exclude_weights) + + bundle = capture_jax_graphs( + params, + specs, + weights_dir=str(weights_root), + graph_meta={ + "frontend": "jax", + "adapter_family": "generic", + "graph_family": "jax_user_graph_bundle", + **dict(graph_meta or {}), + }, + ) + return _write_jax_user_graph_bundle( + bundle=bundle, + specs=specs, + output_dir=root, + weights_dir=weights_root, + model_id=model_id, + task=task, + family=family, + inputs_metadata=inputs_metadata, + ) + + +def build_jax_generation_graph_bundle( + *, + params: Any, + output_dir: str | Path, + model_id: str, + encoder: JaxGraphSpec | None = None, + decoder_prefill: JaxGraphSpec | None = None, + decoder_step: JaxGraphSpec | None = None, + task: str = "text-generation", + family: str = "jax_user_graph", + inputs_metadata: dict[str, object] | None = None, + graph_meta: dict[str, object] | None = None, + weight_arrays: dict[str, object] | None = None, + exclude_weights: set[str] | None = None, + weights_dir: str | Path | None = None, +) -> JaxUserGraphBundleResult: + root = Path(output_dir) + weights_root = Path(weights_dir) if weights_dir is not None else root + arrays = weight_arrays if weight_arrays is not None else flatten_jax_params(params) + write_fp16_weights_manifest(weights_root, arrays, exclude=exclude_weights) + + input_specs = tuple(spec for spec in (encoder, decoder_prefill, decoder_step) if spec is not None) + bundle = capture_jax_generation_graphs( + params, + encoder=encoder, + decoder_prefill=decoder_prefill, + decoder_step=decoder_step, + weights_dir=str(weights_root), + graph_meta={ + "frontend": "jax", + "adapter_family": "generic", + "graph_family": "jax_user_graph_bundle", + **dict(graph_meta or {}), + }, + ) + return _write_jax_user_graph_bundle( + bundle=bundle, + specs=input_specs, + output_dir=root, + weights_dir=weights_root, + model_id=model_id, + task=task, + family=family, + inputs_metadata=inputs_metadata, + ) + + +def capture_jax_generation_graphs( + params: Any, + *, + encoder: JaxGraphSpec | None = None, + decoder_prefill: JaxGraphSpec | None = None, + decoder_step: JaxGraphSpec | None = None, + weights_dir: str | None = None, + weight_bindings: dict[str, dict[str, str]] | None = None, + graph_meta: dict[str, object] | None = None, +) -> Any: + specs: list[JaxGraphSpec] = [] + + def _with_meta( + spec: JaxGraphSpec | None, + *, + name: str, + role: str, + component: str, + ) -> None: + if spec is None: + return + specs.append( + JaxGraphSpec( + name=spec.name or name, + role=spec.role if spec.role != "generic" else role, + fn=spec.fn, + example_args=spec.example_args, + input_names=spec.input_names, + output_names=spec.output_names, + graph_meta={ + "component": component, + **dict(spec.graph_meta or {}), + }, + ) + ) + + _with_meta(encoder, name="encoder", role="encoder", component="encoder") + _with_meta( + decoder_prefill, + name="decoder_prefill", + role="decoder_prefill", + component="decoder_prefill_chunk", + ) + _with_meta( + decoder_step, + name="decoder_step", + role="decoder_step", + component="decoder_step", + ) + if not specs: + raise ValueError("capture_jax_generation_graphs requires at least one graph spec") + return capture_jax_graphs( + params, + specs, + weights_dir=weights_dir, + weight_bindings=weight_bindings, + graph_meta={ + "frontend": "jax", + "adapter_family": "generic", + "graph_family": "generation", + **dict(graph_meta or {}), + }, + ) + + +def _write_jax_user_graph_bundle( + *, + bundle: Any, + specs: Sequence[JaxGraphSpec], + output_dir: Path, + weights_dir: Path, + model_id: str, + task: str, + family: str, + inputs_metadata: dict[str, object] | None, +) -> JaxUserGraphBundleResult: + root = output_dir + weights_root = weights_dir + component_root = root / "components" + component_root.mkdir(parents=True, exist_ok=True) + + components = [] + component_order = [spec.name for spec in specs] + base_payload = { + "model_id": model_id, + "model_source": "jax_user_graph", + "task": task, + "family": family, + "inputs": _serialize_json_compatible(dict(inputs_metadata or {})), + } + for name, captured in bundle.graphs.items(): + component_dir = component_root / name + component_dir.mkdir(parents=True, exist_ok=True) + graph_path = component_dir / "graph.cactus" + raw_ir_path = component_dir / "raw_ir.json" + optimized_ir_path = component_dir / "optimized_ir.json" + _write_json( + raw_ir_path, + { + **base_payload, + "component": name, + "graph": _ir_graph_to_dict(captured.raw_ir_graph), + }, + ) + _write_json( + optimized_ir_path, + { + **base_payload, + "component": name, + "graph": _ir_graph_to_dict(captured.ir_graph), + }, + ) + _embed_materialized_bound_constants(captured.graph) + captured.graph.graph.save(str(graph_path)) + components.append( + { + "component": name, + "directory": str(component_dir.relative_to(root)), + "raw_ir": str(raw_ir_path.relative_to(root)), + "optimized_ir": str(optimized_ir_path.relative_to(root)), + "graph": str(graph_path.relative_to(root)), + "inputs": list(captured.ir_graph.inputs), + "outputs": list(captured.ir_graph.outputs), + "logical_inputs": list(captured.spec.input_names or ()), + "logical_outputs": list(captured.spec.output_names or ()), + "node_count": len(captured.ir_graph.order), + "weight_binding_count": len(captured.graph.bound_constant_bindings), + "runtime_input_node_ids": [int(tensor.id) for tensor in captured.graph.runtime_inputs], + "output_node_ids": [int(tensor.id) for tensor in captured.graph.outputs], + "cache_state_node_ids": [ + { + "layer_key": str(layer_key), + "key": int(key_tensor.id), + "value": int(value_tensor.id), + } + for layer_key, key_tensor, value_tensor in getattr(captured.graph, "cache_state_tensors", []) + ], + "bound_constant_bindings": captured.graph.bound_constant_bindings, + } + ) + + manifest = { + "model_id": model_id, + "model_source": "jax_user_graph", + "task": task, + "family": family, + "component_order": component_order, + "inputs": _serialize_json_compatible(dict(inputs_metadata or {})), + "components": components, + "weights_dir": str(weights_root), + "weights_manifest": str(weights_root / "weights_manifest.json"), + } + manifest_path = component_root / "manifest.json" + _atomic_write_text(manifest_path, json.dumps(manifest, indent=2) + "\n") + + return JaxUserGraphBundleResult( + bundle=bundle, + output_dir=root, + components_manifest_path=manifest_path, + weights_dir=weights_root, + ) diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py new file mode 100644 index 000000000..fab23d6ad --- /dev/null +++ b/python/cactus/transpile/lower.py @@ -0,0 +1,3533 @@ +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import field +import math +import os +from typing import Any + +import numpy as np +import torch + +from cactus.transpile.runtime_compat import Graph +from cactus.transpile.runtime_compat import Tensor +from cactus.transpile.capture_pytorch import CapturedModel +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.optimize_graph import optimize_graph +from cactus.transpile.weight_compat import ensure_binding_compatible +from cactus.transpile.weight_binding import WeightBinding + + +_DYNAMIC_KV_POSITION_OFFSET = (1 << 64) - 1 +_DYNAMIC_KV_QUERY_POSITION_OFFSET = (1 << 64) - 2 +_KV_CACHE_TRANSPARENT_OPS = { + "expand", + "precision_cast", + "reshape", + "slice", + "transpose", + "unsqueeze", + "view", +} + + +@dataclass +class TranspiledGraph: + graph: Graph + runtime_inputs: list[Tensor] + bound_constants: list[Tensor] + bound_constant_bindings: list[dict[str, object]] + bound_constant_value_ids: dict[int, str] + outputs: list[Tensor] + cache_state_tensors: list[tuple[str, Tensor, Tensor]] = field(default_factory=list) + + def set_input(self, index: int, data: Any, *, dtype: int | None = None) -> None: + if index < 0 or index >= len(self.runtime_inputs): + raise IndexError( + f"runtime input index out of range: {index} (have {len(self.runtime_inputs)})" + ) + self.graph.set_input(self.runtime_inputs[index], data, dtype=dtype) + + def set_inputs(self, inputs: list[Any] | tuple[Any, ...]) -> None: + if len(inputs) != len(self.runtime_inputs): + raise ValueError( + f"expected {len(self.runtime_inputs)} runtime inputs, got {len(inputs)}" + ) + for index, value in enumerate(inputs): + self.set_input(index, value) + + def execute(self) -> list[Tensor]: + self.graph.execute() + return self.outputs + + +@dataclass +class BroadcastAlias: + tensor: Tensor + logical_shape: tuple[int, ...] + kind: str + + +def _is_quantized_runtime_dtype(dtype: int) -> bool: + return int(dtype) in { + int(Graph.INT8), + int(getattr(Graph, "CQ1", 3)), + int(getattr(Graph, "CQ2", 4)), + int(getattr(Graph, "CQ3", 5)), + int(getattr(Graph, "CQ4", 6)), + } + + +def transpile_captured(captured: CapturedModel) -> TranspiledGraph: + return transpile_ir(captured.ir_graph) + + +def transpile_ir(ir: IRGraph) -> TranspiledGraph: + verify_ir(ir) + canonicalize_exported_graph(ir) + optimize_graph(ir) + return transpile_preoptimized_ir(ir) + + +def transpile_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: + verify_ir(ir) + g = Graph() + g._transpile_materialized_constants = [] # type: ignore[attr-defined] + trace_lower = os.environ.get("CACTUS_LOWER_TRACE_NODES", "0") != "0" + env: dict[str, Any] = {} + runtime_inputs: list[Tensor] = [] + bound_constants: list[Tensor] = [] + bound_constant_bindings: list[dict[str, object]] = [] + bound_constant_value_ids: dict[int, str] = {} + + for input_index, value_id in enumerate(ir.inputs): + value = ir.values[value_id] + tensor = _lower_input_value(g, value, ir=ir, input_index=input_index) + env[value_id] = tensor + runtime_inputs.append(tensor) + + for value_id, const in ir.constants.items(): + value = ir.values[value_id] + binding = _lookup_weight_binding(value) + if binding is not None: + binding = ensure_binding_compatible(binding, source_tensor=const) + lowered_const = _lower_constant_value(g, value, const, binding=binding) + env[value_id] = lowered_const + if isinstance(lowered_const, Tensor): + bound_constants.append(lowered_const) + bound_constant_value_ids[int(lowered_const.id)] = str(value_id) + if binding is not None: + bound_constant_bindings.append( + { + "node_id": int(lowered_const.id), + "value_id": str(value_id), + "path": binding.path, + "kind": binding.kind, + "source_name": binding.source_name, + } + ) + + for node_id in ir.order: + node = ir.nodes[node_id] + outputs = _lower_ir_node(g, node, env, ir) + if len(outputs) != len(node.outputs): + raise ValueError( + f"node {node.id} produced {len(outputs)} outputs, expected {len(node.outputs)}" + ) + for output_id, tensor in zip(node.outputs, outputs): + env[output_id] = tensor + if trace_lower and isinstance(tensor, Tensor): + print( + "lower_trace" + f" cactus_id={int(tensor.id)}" + f" ir_node={node.id}" + f" op={node.op}" + f" output={output_id}" + f" shape={tuple(int(dim) for dim in tensor.shape)}", + flush=True, + ) + + outputs = [env[value_id] for value_id in ir.outputs] + seen_bound_constant_ids = {int(tensor.id) for tensor in bound_constants} + for tensor in getattr(g, "_transpile_materialized_constants", []): + tensor_id = int(tensor.id) + if tensor_id in seen_bound_constant_ids: + continue + bound_constants.append(tensor) + seen_bound_constant_ids.add(tensor_id) + cache_state_tensors = list(env.get("__internal_kv_cache_state_entries", [])) + return TranspiledGraph( + graph=g, + runtime_inputs=runtime_inputs, + bound_constants=bound_constants, + bound_constant_bindings=bound_constant_bindings, + bound_constant_value_ids=bound_constant_value_ids, + outputs=outputs, + cache_state_tensors=cache_state_tensors, + ) + + +def _cross_kv_cache_buffer_elements(max_seq_len: int, num_kv_heads: int, head_dim: int) -> int: + kv_quant_group_size = 32 + num_groups = (int(head_dim) + kv_quant_group_size - 1) // kv_quant_group_size + kv_values = int(max_seq_len) * int(num_kv_heads) * int(head_dim) + kv_scales = int(max_seq_len) * int(num_kv_heads) * num_groups * 4 + return int(64 + kv_values + kv_scales) + + +def _external_cross_kv_cache_input_indices(ir: IRGraph) -> set[int]: + if not bool(ir.meta.get("external_cross_kv_cache_inputs", False)): + return set() + start = int(ir.meta.get("cross_kv_input_start_index", -1)) + count = int(ir.meta.get("cross_kv_input_count", 0)) + if start < 0 or count <= 0: + return set() + return set(range(start, start + count)) + + +def _is_external_cross_kv_cache_value(ir: IRGraph, value_id: str) -> bool: + return any( + input_value_id == value_id and input_index in _external_cross_kv_cache_input_indices(ir) + for input_index, input_value_id in enumerate(ir.inputs) + ) + + +def _lower_input_value(g: Graph, value: IRValue, *, ir: IRGraph | None = None, input_index: int | None = None) -> Tensor: + if value.shape is None or value.dtype is None: + raise ValueError(f"IR input missing shape or dtype: {value.id}") + if ir is not None and input_index is not None and input_index in _external_cross_kv_cache_input_indices(ir): + shape = tuple(int(dim) for dim in value.shape) + if len(shape) != 4: + raise ValueError(f"external cross-KV cache input expects original rank-4 K/V tensor, got {shape}") + layout = str(ir.meta.get("cross_kv_input_layout", "bthd") or "bthd").lower() + if layout == "bhsd": + max_seq_len, num_kv_heads, head_dim = shape[2], shape[1], shape[3] + else: + max_seq_len, num_kv_heads, head_dim = shape[1], shape[2], shape[3] + cache_elements = _cross_kv_cache_buffer_elements(max_seq_len, num_kv_heads, head_dim) + return g.input(shape=(cache_elements,), dtype=Graph.INT8) + dynamic_dims = value.meta.get("dynamic_dims") if value.meta else None + return g.input(shape=value.shape, dtype=_map_ir_dtype(value.dtype), dynamic_dims=dynamic_dims) + + +def _should_lower_attention_with_internal_kv_cache(ir: IRGraph, node: IRNode) -> bool: + if not bool(ir.meta.get("use_internal_kv_cache", False)): + return False + if node.op not in {"attention", "scaled_dot_product_attention"}: + return False + component = str(ir.meta.get("component", "") or "").strip().lower() + if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step", + "decoder_embed_chunk"}: + return False + if len(node.inputs) < 3: + return False + + def _sequence_dim(layout: str) -> int: + return 1 if layout == "bthd" else 2 + + q_layout = str(node.attrs.get("q_layout", node.attrs.get("qkv_layout", "bhsd")) or "bhsd").lower() + k_layout = str(node.attrs.get("k_layout", node.attrs.get("qkv_layout", "bhsd")) or "bhsd").lower() + query_value = ir.values.get(node.inputs[0]) + key_value = ir.values.get(node.inputs[1]) + query_shape = tuple(int(dim) for dim in (query_value.shape or ())) if query_value is not None else () + key_shape = tuple(int(dim) for dim in (key_value.shape or ())) if key_value is not None else () + if len(query_shape) >= 4 and len(key_shape) >= 4: + query_seq = int(query_shape[_sequence_dim(q_layout)]) + key_seq = int(key_shape[_sequence_dim(k_layout)]) + # Seq2seq decoders interleave self-attention (Q/K over the current + # decoder tokens) with cross-attention (Q over decoder tokens, K/V over + # encoder frames). Only the self-attention stream belongs in the + # autoregressive KV cache; cross-attention must continue to read the + # full encoder output every step. + if query_seq != key_seq: + return False + return True + + +def _kv_cache_layer_key(node: IRNode) -> str: + for meta_key in ("attention_layer_index", "layer_index"): + value = node.meta.get(meta_key) + if value is not None: + return str(value) + return str(node.id) + + +def _lower_attention_with_internal_kv_cache( + g: Graph, + ir: IRGraph, + env: dict[str, Any], + node: IRNode, + *, + query: Tensor, + key: Tensor, + value: Tensor, + output_layout: str, +) -> list[Tensor]: + key_shape = tuple(int(dim) for dim in key.shape) + value_shape = tuple(int(dim) for dim in value.shape) + if len(key_shape) != 4 or len(value_shape) != 4: + raise NotImplementedError( + f"cached attention expects key/value as [batch, seq, heads, dim], got {key_shape} and {value_shape}" + ) + if int(key_shape[0]) != 1: + raise NotImplementedError("cached attention currently expects batch size 1") + + cache_states: dict[str, tuple[Tensor, Tensor]] = env.setdefault("__internal_kv_cache_states", {}) # type: ignore[assignment] + layer_key = _kv_cache_layer_key(node) + if layer_key not in cache_states: + default_cache_len = max(512, int(key_shape[1])) + requested_cache_len = node.meta.get("max_cache_seq_len", ir.meta.get("max_cache_seq_len", default_cache_len)) + max_cache_seq_len = max(default_cache_len, int(requested_cache_len or default_cache_len)) + sink_size = int(ir.meta.get("cache_sink_size", 4) or 4) + window_size = int(node.attrs.get("window_size", 0) or 0) + cache_num_slots = int(ir.meta.get("cache_num_slots", 1) or 1) + k_cache = g.kv_cache_state( + max_cache_seq_len, + int(key_shape[2]), + int(key_shape[3]), + window_size=window_size, + sink_size=sink_size, + num_slots=cache_num_slots, + ) + v_cache = g.kv_cache_state( + max_cache_seq_len, + int(value_shape[2]), + int(value_shape[3]), + window_size=window_size, + sink_size=sink_size, + num_slots=cache_num_slots, + ) + cache_states[layer_key] = (k_cache, v_cache) + cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] + cache_entries.append((layer_key, k_cache, v_cache)) + + k_cache, v_cache = cache_states[layer_key] + window_size = int(node.attrs.get("window_size", 0) or 0) + sink_size = int(ir.meta.get("cache_sink_size", 4) or 4) + g.kv_cache_append(key, k_cache, window_size=window_size, sink_size=sink_size) + g.kv_cache_append(value, v_cache, window_size=window_size, sink_size=sink_size) + kv_input_cache_states: dict[tuple[str, str], tuple[Tensor, Tensor, int, int]] = env.setdefault( + "__internal_kv_cache_by_kv_inputs", + {}, + ) # type: ignore[assignment] + cache_entry = (k_cache, v_cache, window_size, int(value_shape[3])) + kv_input_cache_states[(node.inputs[1], node.inputs[2])] = cache_entry + kv_input_cache_states[_canonical_kv_cache_input_key(ir, node)] = cache_entry + out = g.attention_cached( + query, + key, + value, + k_cache, + v_cache, + scale=_resolve_attention_scale(node, query), + position_offset=_DYNAMIC_KV_POSITION_OFFSET, + window_size=window_size, + v_head_dim=int(value_shape[3]), + ) + if output_layout == "bthd": + return [out] + return [g.permute(out, (0, 2, 1, 3))] + + +def _lower_attention_with_shared_internal_kv_cache( + g: Graph, + ir: IRGraph, + env: dict[str, Any], + node: IRNode, + *, + query: Tensor, + key: Tensor, + value: Tensor, + output_layout: str, +) -> list[Tensor] | None: + kv_input_cache_states: dict[tuple[str, str], tuple[Tensor, Tensor, int, int]] = env.get( + "__internal_kv_cache_by_kv_inputs", + {}, + ) + cache_state = kv_input_cache_states.get((node.inputs[1], node.inputs[2])) + if cache_state is None: + cache_state = kv_input_cache_states.get(_canonical_kv_cache_input_key(ir, node)) + if cache_state is None: + return None + k_cache, v_cache, window_size, v_head_dim = cache_state + out = g.attention_cached( + query, + key, + value, + k_cache, + v_cache, + scale=_resolve_attention_scale(node, query), + position_offset=_DYNAMIC_KV_QUERY_POSITION_OFFSET, + window_size=window_size, + v_head_dim=v_head_dim, + ) + if output_layout == "bthd": + return [out] + return [g.permute(out, (0, 2, 1, 3))] + + +def _canonical_kv_cache_input_key(ir: IRGraph, node: IRNode) -> tuple[str, str]: + return ( + _canonical_kv_cache_value_id(ir, node.inputs[1]), + _canonical_kv_cache_value_id(ir, node.inputs[2]), + ) + + +def _canonical_kv_cache_value_id(ir: IRGraph, value_id: str) -> str: + current = value_id + seen: set[str] = set() + while current not in seen: + seen.add(current) + value = ir.values.get(current) + producer_id = None if value is None else value.producer + if producer_id is None: + break + producer = ir.nodes.get(producer_id) + if producer is None or producer.op not in _KV_CACHE_TRANSPARENT_OPS or len(producer.inputs) != 1: + break + current = producer.inputs[0] + return current + + +def _should_lower_conv1d_with_internal_conv_cache(ir: IRGraph, node: IRNode, x: Tensor, weight: Tensor) -> bool: + if not bool(ir.meta.get("use_internal_conv_cache", False)): + return False + if node.op != "conv1d": + return False + component = str(ir.meta.get("component", "") or "").strip().lower() + if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step", + "decoder_embed_chunk"}: + return False + stride = int(node.attrs.get("stride", 1)) + padding = int(node.attrs.get("padding", 0)) + dilation = int(node.attrs.get("dilation", 1)) + groups = int(node.attrs.get("groups", 1)) + return ( + len(x.shape) == 3 + and len(weight.shape) == 3 + and int(x.shape[0]) == 1 + and groups == int(x.shape[1]) == int(weight.shape[0]) + and int(weight.shape[1]) == 1 + and stride == 1 + and dilation == 1 + and padding == max(int(weight.shape[2]) - 1, 0) + ) + + +def _conv_cache_layer_key(ir: IRGraph, node: IRNode) -> str: + if len(node.inputs) > 1: + value = ir.values.get(node.inputs[1]) + if value is not None and isinstance(value.meta, dict): + source_name = value.meta.get("source_name") + if isinstance(source_name, str) and source_name: + return source_name + return str(node.id) + + +def _lower_conv1d_with_internal_conv_cache( + g: Graph, + ir: IRGraph, + env: dict[str, Any], + node: IRNode, + *, + x: Tensor, + weight: Tensor, + bias: Tensor | None, +) -> list[Tensor]: + channels = int(x.shape[1]) + seq_len = int(x.shape[2]) + kernel_size = int(weight.shape[2]) + x_nlc = g.permute(x, (0, 2, 1)) + + layer_key = f"conv:{_conv_cache_layer_key(ir, node)}" + conv_states: dict[str, Tensor] = env.setdefault("__internal_conv_cache_states", {}) # type: ignore[assignment] + if layer_key not in conv_states: + cache_state = g.conv_cache_state(kernel_size, channels) + conv_states[layer_key] = cache_state + cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] + # Reuse the existing manifest/cache-transfer shape. Conv state is a + # single mutable tensor, so key/value intentionally point to the same node. + cache_entries.append((layer_key, cache_state, cache_state)) + cache_state = conv_states[layer_key] + + if seq_len > 1: + # Prefill chunks still compute the chunk output with the regular causal + # convolution; the cache append is the side-effect needed by later step + # graphs. The cache only needs the last K rows for the next token. + x_rows = g.reshape(x_nlc, (seq_len, channels)) + g.conv_cache_append(x_rows, cache_state) + out_nlc = g.conv1d_causal(x_nlc, weight, kernel_size=kernel_size, dilation=1) + else: + x_rows = g.reshape(x_nlc, (1, channels)) + window = g.conv_cache_append(x_rows, cache_state) + window_nlc = g.reshape(window, (1, kernel_size, channels)) + out_nlc = g.conv1d_causal(window_nlc, weight, kernel_size=kernel_size, dilation=1) + out_nlc = g.slice(out_nlc, axis=1, start=kernel_size - 1, length=1) + + if bias is not None: + bias_reshaped = g.reshape(bias, (1, 1, int(weight.shape[0]))) + out_nlc, bias_reshaped = _legalize_elementwise_binary_inputs(g, out_nlc, bias_reshaped) + out_nlc = g.add(out_nlc, bias_reshaped) + return [g.permute(out_nlc, (0, 2, 1))] + + +def _should_lower_gated_deltanet_with_internal_cache( + ir: IRGraph, + op: str, + *, + seq_len: int, +) -> bool: + if not bool(ir.meta.get("use_internal_gated_deltanet_cache", False)): + return False + if op not in {"gated_deltanet_prefill", "gated_deltanet_decode"}: + return False + component = str(ir.meta.get("component", "") or "").strip().lower() + return component in {"decoder_prefill_chunk", "decoder_step", "decoder_media_step"} + + +def _gated_deltanet_layer_key(ir: IRGraph, node: IRNode) -> str: + if len(node.inputs) > 1: + value = ir.values.get(node.inputs[1]) + if value is not None and isinstance(value.meta, dict): + source_name = value.meta.get("source_name") + if isinstance(source_name, str) and source_name: + return source_name + return str(node.id) + + +def _lower_gated_deltanet_initial_state( + g: Graph, + ir: IRGraph, + env: dict[str, Any], + node: IRNode, + *, + batch_size: int, + key_dim: int, + num_v_heads: int, + value_dim: int, + seq_len: int, + op: str, +) -> tuple[str | None, Tensor]: + state_shape = (batch_size, key_dim, num_v_heads, value_dim) + if not _should_lower_gated_deltanet_with_internal_cache(ir, op, seq_len=seq_len): + return None, _materialize_constant_tensor(g, torch.zeros(state_shape, dtype=torch.float16)) + + layer_key = f"gdn:{_gated_deltanet_layer_key(ir, node)}" + gdn_states: dict[str, Tensor] = env.setdefault("__internal_gated_deltanet_cache_states", {}) # type: ignore[assignment] + if layer_key not in gdn_states: + gdn_states[layer_key] = g.recurrent_cache_state(state_shape, dtype=g.FP16) + return layer_key, gdn_states[layer_key] + + +def _lower_gated_deltanet_conv1d( + g: Graph, + ir: IRGraph, + env: dict[str, Any], + node: IRNode, + *, + mixed_qkv: Tensor, + conv_weight: Tensor, + batch_size: int, + seq_len: int, + mixed_qkv_dim: int, + op: str, +) -> Tensor: + kernel_size = int(conv_weight.shape[2]) + if not _should_lower_gated_deltanet_with_internal_cache(ir, op, seq_len=seq_len): + return g.conv1d_causal(mixed_qkv, conv_weight, kernel_size=kernel_size, dilation=1) + + layer_key = f"conv:gdn:{_gated_deltanet_layer_key(ir, node)}" + conv_states: dict[str, Tensor] = env.setdefault("__internal_conv_cache_states", {}) # type: ignore[assignment] + if layer_key not in conv_states: + cache_state = g.conv_cache_state(kernel_size, mixed_qkv_dim) + conv_states[layer_key] = cache_state + cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] + cache_entries.append((layer_key, cache_state, cache_state)) + + cache_state = conv_states[layer_key] + if batch_size != 1: + raise NotImplementedError( + f"gated_deltanet conv cache lowering currently supports batch_size=1 only, got {batch_size}" + ) + x_rows = g.reshape(mixed_qkv, (batch_size * seq_len, mixed_qkv_dim)) + + if seq_len == 1: + window = g.conv_cache_append(x_rows, cache_state) + window_nlc = g.reshape(window, (batch_size, kernel_size, mixed_qkv_dim)) + conv_out = g.conv1d_causal(window_nlc, conv_weight, kernel_size=kernel_size, dilation=1) + return g.slice(conv_out, axis=1, start=kernel_size - 1, length=1) + + conv_out = g.conv1d_causal(mixed_qkv, conv_weight, kernel_size=kernel_size, dilation=1) + g.conv_cache_initialize(x_rows, cache_state) + return conv_out + + +def _lower_constant_value( + g: Graph, + value: IRValue, + const: Any, + *, + binding: WeightBinding | None = None, +) -> Any: + if binding is None: + binding = _lookup_weight_binding(value) + if binding is not None: + _debug_mmap_binding(value.id, binding) + return g.mmap_weights(binding.path) + + _debug_constant_fallback(value, const) + + if isinstance(const, torch.nn.Parameter): + const = const.detach() + if isinstance(const, torch.Tensor): + tensor_value = const.detach().cpu() + else: + raise NotImplementedError( + f"unsupported IR constant type for {value.id}: {type(const).__name__}" + ) + + if tensor_value.numel() == 1 and value.shape in {None, ()}: + return tensor_value.item() + + return _materialize_constant_tensor(g, tensor_value) + + +def _lookup_weight_binding(value: IRValue) -> WeightBinding | None: + meta = getattr(value, "meta", None) + if isinstance(meta, dict): + path = meta.get("path") + kind = meta.get("kind") + source_name = meta.get("source_name") + if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): + return WeightBinding(path=path, kind=kind, source_name=source_name) + return None + + +GEMMA4_RESIDUAL_DOWNSCALE = 16.0 + + +def _activation_enum(name: object, *, default: int = 0) -> int: + value = str(name or "").strip().lower() + if value in {"gelu", "gelu_tanh", "gelu_pytorch_tanh"}: + return 1 + if value in {"gelu_erf", "gelu_none"}: + return 2 + if value == "relu": + return 3 + if value == "silu": + return 0 + return default + + +def _debug_mmap_binding(value_id: str, binding: WeightBinding) -> None: + if os.environ.get("CACTUS_TRANSPILER_DEBUG_MMAP") != "1": + return + print( + "[transpile:mmap] " + f"value={value_id} " + f"kind={binding.kind} " + f"source={binding.source_name} " + f"path={binding.path}" + ) + + +def _debug_constant_fallback(value: IRValue, const: Any) -> None: + if os.environ.get("CACTUS_TRANSPILER_DEBUG_MMAP") != "1": + return + const_type = type(const).__name__ + shape = None + dtype = None + if isinstance(const, torch.Tensor): + shape = tuple(const.shape) + dtype = str(const.dtype) + source_name = None + if isinstance(value.meta, dict): + source_name = value.meta.get("source_name") + print( + "[transpile:fallback] " + f"value={value.id} " + f"source={source_name} " + f"const_type={const_type} " + f"shape={shape} " + f"dtype={dtype}" + ) + + +def _debug_embedding_lowering(node: IRNode, embedding_tensor: Tensor, indices_tensor: Tensor) -> None: + if os.environ.get("CACTUS_TRANSPILER_DEBUG_MMAP") != "1": + return + print( + "[transpile:embedding] " + f"node={node.id} " + f"embedding_id={embedding_tensor.id} " + f"embedding_shape={tuple(embedding_tensor.shape)} " + f"embedding_dtype={embedding_tensor.dtype} " + f"indices_id={indices_tensor.id} " + f"indices_shape={tuple(indices_tensor.shape)} " + f"indices_dtype={indices_tensor.dtype}" + ) + + +def _matmul_with_quantized_rhs_legalization( + g: Graph, + lhs: Tensor, + rhs: Tensor, + *, + pretransposed_rhs: bool = False, + output_dtype: int | None = None, +) -> Tensor: + if _is_quantized_runtime_dtype(rhs.dtype) and lhs.dtype == Graph.FP32: + lhs = g.precision_cast(lhs, Graph.FP16) + if ( + _is_quantized_runtime_dtype(rhs.dtype) + and pretransposed_rhs + and len(lhs.shape) >= 2 + and len(rhs.shape) == 2 + and int(rhs.shape[-1]) > int(lhs.shape[-1]) + ): + pad = int(rhs.shape[-1]) - int(lhs.shape[-1]) + pad_shape = tuple(int(dim) for dim in lhs.shape[:-1]) + (pad,) + pad_dtype = torch.float16 if lhs.dtype == Graph.FP16 else torch.float32 + zeros = _materialize_constant_tensor(g, torch.zeros(pad_shape, dtype=pad_dtype)) + lhs = g.cat([lhs, zeros], axis=len(lhs.shape) - 1) + return g.matmul(lhs, rhs, pretransposed_rhs=pretransposed_rhs, output_dtype=output_dtype) + + +def _matmul_output_dtype( + lhs: Tensor, + rhs: Tensor, + *, + output_dtype: int | None, +) -> int | None: + if output_dtype != Graph.FP16: + return output_dtype + if _is_quantized_runtime_dtype(rhs.dtype): + return output_dtype + if lhs.dtype == Graph.FP32 or rhs.dtype == Graph.FP32: + return Graph.FP32 + return output_dtype + + +def _trim_padded_last_dim(g: Graph, tensor: Tensor, expected_shape: tuple[int, ...] | None) -> Tensor: + if expected_shape is None: + return tensor + actual_shape = tuple(int(dim) for dim in tensor.shape) + target_shape = tuple(int(dim) for dim in expected_shape) + if actual_shape == target_shape: + return tensor + if ( + len(actual_shape) == len(target_shape) + and len(actual_shape) >= 1 + and actual_shape[:-1] == target_shape[:-1] + and actual_shape[-1] > target_shape[-1] + ): + return g.slice(tensor, len(actual_shape) - 1, 0, target_shape[-1]) + return tensor + + +def _legalize_for_transpose(g: Graph, x: Tensor) -> Tensor: + return x + + +def _reshape_attention_output_for_linear( + g: Graph, + attn_out: Tensor, + target_shape: tuple[int, ...], +) -> Tensor: + actual_shape = tuple(int(dim) for dim in attn_out.shape) + target_shape = tuple(int(dim) for dim in target_shape) + if len(actual_shape) == 4 and len(target_shape) == 3 and actual_shape[0] == target_shape[0]: + batch, dim1, dim2, dim3 = actual_shape + _, target_seq, target_hidden = target_shape + if dim1 == target_seq and dim2 * dim3 == target_hidden: + return g.reshape(attn_out, target_shape) + if dim2 == target_seq and dim1 * dim3 == target_hidden: + return g.reshape(g.permute(attn_out, (0, 2, 1, 3)), target_shape) + return g.reshape(attn_out, target_shape) + + +def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> list[Any]: + op = node.op + + if op == "arange": + start = int(node.attrs.get("start", 0)) + end = int(node.attrs["end"]) + step = int(node.attrs.get("step", 1)) + dtype = _materialize_constant_torch_dtype(node.attrs.get("dtype")) + tensor_value = torch.arange(start, end, step=step, dtype=dtype) + return [_materialize_constant_tensor(g, tensor_value)] + + if op == "add": + return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "add")] + + if op == "add_clipped": + lhs, rhs = _legalize_elementwise_binary_inputs(g, _tensor(env, node.inputs[0]), _tensor(env, node.inputs[1])) + return [g.add_clipped(lhs, rhs)] + + if op == "subtract": + return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "subtract")] + + if op == "multiply": + return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "multiply")] + + if op == "multiply_inplace": + return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "multiply")] + + if op == "divide": + return [_lower_binary_op(g, env[node.inputs[0]], env[node.inputs[1]], "divide")] + + if op == "not_equal": + return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "not_equal")] + + if op == "equal": + return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "equal")] + + if op == "greater": + return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "greater")] + + if op == "greater_equal": + return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "greater_equal")] + + if op == "less": + return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "less")] + + if op == "less_equal": + return [_lower_compare_op(g, env[node.inputs[0]], env[node.inputs[1]], "less_equal")] + + if op == "logical_and": + lhs = _lower_compare_op(g, env[node.inputs[0]], 0.0, "not_equal") + rhs = _lower_compare_op(g, env[node.inputs[1]], 0.0, "not_equal") + return [_lower_binary_op(g, lhs, rhs, "multiply")] + + if op == "logical_or": + lhs = _lower_compare_op(g, env[node.inputs[0]], 0.0, "not_equal") + rhs = _lower_compare_op(g, env[node.inputs[1]], 0.0, "not_equal") + return [_lower_compare_op(g, _lower_binary_op(g, lhs, rhs, "add"), 0.0, "not_equal")] + + if op == "logical_not": + return [_lower_compare_op(g, env[node.inputs[0]], 0.0, "equal")] + + if op == "where": + return [_lower_where_op(g, node, env)] + + if op == "scalar_add": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [g.scalar_add(x, float(node.attrs["value"]))] + + if op == "scalar_subtract": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [g.scalar_subtract(x, float(node.attrs["value"]))] + + if op == "scalar_subtract_reverse": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [g.scalar_add(g.scalar_multiply(x, -1.0), float(node.attrs["value"]))] + + if op == "scalar_multiply": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [g.scalar_multiply(x, float(node.attrs["value"]))] + + if op == "scalar_divide": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [g.scalar_divide(x, float(node.attrs["value"]))] + + if op == "scalar_floor_divide": + source = _tensor(env, node.inputs[0]) + x = _ensure_scalar_math_tensor(g, source) + y = g.scalar_floor_divide(x, float(node.attrs["value"])) + output_dtype = ir.values[node.outputs[0]].dtype + if output_dtype is None: + return [y] + target_dtype = _map_ir_dtype(output_dtype) + if y.dtype == target_dtype: + return [y] + return [g.precision_cast(y, target_dtype)] + + if op == "scalar_not_equal": + value = float(node.attrs["value"]) + source = _tensor(env, node.inputs[0]) + target_dtype = Graph.FP32 if source.dtype != Graph.FP16 or abs(value) > 65504.0 else Graph.FP16 + x = _ensure_tensor_dtype(g, source, target_dtype) + return [g.scalar_not_equal(x, value)] + + if op == "scalar_equal": + value = float(node.attrs["value"]) + source = _tensor(env, node.inputs[0]) + target_dtype = Graph.FP32 if source.dtype != Graph.FP16 or abs(value) > 65504.0 else Graph.FP16 + x = _ensure_tensor_dtype(g, source, target_dtype) + not_equal = g.scalar_not_equal(x, value) + return [g.scalar_not_equal(not_equal, 1.0)] + + if op == "scalar_greater": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [_lower_compare_op(g, x, float(node.attrs["value"]), "greater")] + + if op == "scalar_greater_equal": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [_lower_compare_op(g, x, float(node.attrs["value"]), "greater_equal")] + + if op == "scalar_less": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [_lower_compare_op(g, x, float(node.attrs["value"]), "less")] + + if op == "scalar_less_equal": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + return [_lower_compare_op(g, x, float(node.attrs["value"]), "less_equal")] + + if op == "scalar_divide_reverse": + raise NotImplementedError("scalar_divide_reverse is not directly supported by Cactus graph ops") + + if op == "precision_cast": + target_dtype = node.attrs.get("dtype") + if target_dtype is None: + return [_tensor(env, node.inputs[0])] + return [g.precision_cast(_tensor(env, node.inputs[0]), _map_ir_or_torch_dtype(target_dtype))] + + if op == "type_as": + source = _tensor(env, node.inputs[0]) + target = _tensor(env, node.inputs[1]) + return [g.precision_cast(source, target.dtype)] + + if op == "abs": + return [g.abs(_tensor(env, node.inputs[0]))] + + if op == "clamp": + x = _ensure_fp16_tensor(g, _tensor(env, node.inputs[0])) + min_value = node.attrs.get("min") + max_value = node.attrs.get("max") + input_index = 1 + if min_value is None and ( + bool(node.attrs.get("has_min_tensor", False)) or len(node.inputs) > input_index + ): + min_value = _constant_scalar_from_ir(ir, node.inputs[input_index]) + input_index += 1 + if max_value is None and ( + bool(node.attrs.get("has_max_tensor", False)) or len(node.inputs) > input_index + ): + max_value = _constant_scalar_from_ir(ir, node.inputs[input_index]) + lo = float(min_value) if min_value is not None and math.isfinite(float(min_value)) else -65504.0 + hi = float(max_value) if max_value is not None and math.isfinite(float(max_value)) else 65504.0 + return [g.clamp(x, lo, hi)] + + if op == "negate": + return [g.scalar_multiply(_tensor(env, node.inputs[0]), -1.0)] + + if op == "pow": + exponent = node.attrs.get("exponent") + if exponent is None: + raise NotImplementedError("tensor-tensor pow is not supported by Cactus graph ops") + return [g.pow(_tensor(env, node.inputs[0]), float(exponent))] + + if op == "scalar_exp": + return [g.scalar_exp(_tensor(env, node.inputs[0]))] + + if op == "scalar_sqrt": + return [g.scalar_sqrt(_tensor(env, node.inputs[0]))] + + if op == "scalar_log": + return [g.scalar_log(_tensor(env, node.inputs[0]))] + + if op == "scalar_cos": + return [g.scalar_cos(_tensor(env, node.inputs[0]))] + + if op == "scalar_sin": + return [g.scalar_sin(_tensor(env, node.inputs[0]))] + + if op in {"reshape", "view"}: + source = env[node.inputs[0]] + if isinstance(source, BroadcastAlias): + input_shape = tuple(source.logical_shape) + else: + input_shape = tuple(_tensor(env, node.inputs[0]).shape) + target_shape = _resolve_reshape_shape(input_shape, tuple(node.attrs["shape"])) + if isinstance(source, BroadcastAlias): + if source.kind != "gqa_repeat_kv": + raise NotImplementedError(f"unsupported broadcast alias kind in reshape: {source.kind}") + base_shape = tuple(source.tensor.shape) + if len(base_shape) != 4: + raise NotImplementedError(f"gqa_repeat_kv alias requires 4D base tensor, got {base_shape}") + if len(target_shape) != 4: + return [g.reshape(_materialize_broadcast_alias(g, source), target_shape)] + batch, kv_heads, seq_len, head_dim = base_shape + if target_shape[0] != batch or target_shape[2] != seq_len or target_shape[3] != head_dim: + raise NotImplementedError( + f"gqa_repeat_kv reshape mismatch: base {base_shape}, target {target_shape}" + ) + if target_shape[1] % max(kv_heads, 1) != 0: + raise NotImplementedError( + f"gqa_repeat_kv target head count must be a multiple of kv heads: {base_shape} -> {target_shape}" + ) + return [BroadcastAlias(source.tensor, target_shape, source.kind)] + return [g.reshape(_tensor(env, node.inputs[0]), target_shape)] + + if op == "flatten": + return [ + g.flatten( + _tensor(env, node.inputs[0]), + start_dim=int(node.attrs["start_dim"]), + end_dim=int(node.attrs["end_dim"]), + ) + ] + + if op == "unsqueeze": + x = _tensor(env, node.inputs[0]) + target_shape = node.attrs.get("shape") + if target_shape is None: + dim = _normalize_dim(int(node.attrs["dim"]), len(x.shape) + 1) + shape_list = list(x.shape) + shape_list.insert(dim, 1) + target_shape = tuple(shape_list) + return [g.reshape(x, tuple(target_shape))] + + if op == "expand": + matched_base = _match_gqa_expand_alias(ir, node) + if matched_base is not None: + return [BroadcastAlias(_tensor(env, matched_base), tuple(node.attrs["shape"]), "gqa_repeat_kv")] + x = _tensor(env, node.inputs[0]) + target_shape = _resolve_expand_shape(tuple(x.shape), tuple(node.attrs["shape"])) + if target_shape == tuple(x.shape): + return [x] + return [g.expand(x, target_shape)] + + if op == "repeat": + x = _tensor(env, node.inputs[0]) + repeats = tuple(int(value) for value in node.attrs["repeats"]) + return [_lower_repeat(g, x, repeats)] + + if op == "one_hot": + indices = _tensor(env, node.inputs[0]) + num_classes = int(node.attrs["num_classes"]) + eye = torch.eye(num_classes, dtype=torch.float16) + embedding = _materialize_constant_tensor(g, eye) + return [g.embedding_from_tensor(embedding, indices)] + + if op == "transpose": + x = _legalize_for_transpose(g, _tensor(env, node.inputs[0])) + dim0 = _normalize_dim(int(node.attrs["dim0"]), len(x.shape)) + dim1 = _normalize_dim(int(node.attrs["dim1"]), len(x.shape)) + rank = len(x.shape) + permutation = list(range(rank)) + permutation[dim0], permutation[dim1] = permutation[dim1], permutation[dim0] + if rank == 2 and permutation == [1, 0]: + return [g.transpose(x)] + return [g.permute(x, permutation)] + + if op == "permute": + source = env[node.inputs[0]] + if isinstance(source, BroadcastAlias): + x = _materialize_broadcast_alias(g, source) + else: + x = _tensor(env, node.inputs[0]) + x = _legalize_for_transpose(g, x) + permutation = tuple(_normalize_dim(int(dim), len(x.shape)) for dim in node.attrs["permutation"]) + if len(permutation) == 2 and permutation == (1, 0): + return [g.transpose(x)] + return [g.permute(x, permutation)] + + if op == "matmul": + lhs = _tensor_or_materialized_alias(g, env, node.inputs[0]) + rhs = _tensor_or_materialized_alias(g, env, node.inputs[1]) + output_value = ir.values.get(node.outputs[0]) + output_dtype = ( + _map_ir_dtype(output_value.dtype) + if output_value is not None and output_value.dtype is not None + else None + ) + matmul_output_dtype = _matmul_output_dtype(lhs, rhs, output_dtype=output_dtype) + if len(lhs.shape) == 2 and len(rhs.shape) == 2: + out = _matmul_with_quantized_rhs_legalization(g, lhs, rhs, output_dtype=matmul_output_dtype) + if output_dtype is not None and out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + legalized = _legalize_matmul_inputs(g, lhs, rhs, node) + if legalized is None: + static_batched = _lower_static_batched_matmul(g, lhs, rhs, output_dtype=matmul_output_dtype) + if static_batched is not None: + if output_dtype is not None and static_batched.dtype != output_dtype: + static_batched = g.precision_cast(static_batched, output_dtype) + return [static_batched] + out = _matmul_with_quantized_rhs_legalization(g, lhs, rhs, output_dtype=matmul_output_dtype) + if output_dtype is not None and out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + lhs_2d, rhs_2d, output_shape = legalized + out = _matmul_with_quantized_rhs_legalization(g, lhs_2d, rhs_2d, output_dtype=matmul_output_dtype) + out = g.reshape(out, output_shape) + if output_dtype is not None and out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op == "linear": + x = _tensor(env, node.inputs[0]) + weight = _tensor(env, node.inputs[1]) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else x.dtype + reshape_back: tuple[int, ...] | None = None + if len(x.shape) > 2: + x = _flatten_to_2d_for_linear(g, x) + output_shape = output_value.shape if output_value is not None else None + if output_shape is None: + output_shape = node.meta.get("shape") + if not isinstance(output_shape, tuple): + raise NotImplementedError(f"linear missing output shape metadata for node {node.id}") + reshape_back = tuple(int(v) for v in output_shape) + + matmul_output_dtype = _matmul_output_dtype(x, weight, output_dtype=output_dtype) + if node.attrs.get("has_bias"): + bias = _tensor(env, node.inputs[2]) + if ( + x.dtype == Graph.FP16 + and weight.dtype == Graph.FP16 + and output_dtype == Graph.FP16 + ): + matmul_output_dtype = Graph.FP32 + elif bias.dtype == Graph.FP32 and not _is_quantized_runtime_dtype(weight.dtype): + matmul_output_dtype = Graph.FP32 + out = _matmul_with_quantized_rhs_legalization( + g, + x, + weight, + pretransposed_rhs=True, + output_dtype=matmul_output_dtype, + ) + out = _trim_padded_last_dim(g, out, output_value.shape if output_value is not None else None) + if node.attrs.get("has_bias"): + bias = _ensure_tensor_dtype(g, bias, matmul_output_dtype) + if output_value is not None and output_value.shape is not None: + expected_last_dim = int(output_value.shape[-1]) + if len(bias.shape) >= 1 and int(bias.shape[-1]) > expected_last_dim: + bias = g.slice(bias, len(bias.shape) - 1, 0, expected_last_dim) + out, bias = _legalize_elementwise_binary_inputs(g, out, bias) + try: + out = g.add(out, bias) + except RuntimeError as exc: + raise RuntimeError( + "linear bias add failed while lowering " + f"{node.id}: out_shape={tuple(out.shape)} bias_shape={tuple(bias.shape)} " + f"expected_shape={None if output_value is None else output_value.shape}" + ) from exc + if reshape_back is not None: + out = g.reshape(out, reshape_back) + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op == "dense_mlp_tq_fused": + hidden = _tensor(env, node.inputs[0]) + gate_weight = _tensor(env, node.inputs[1]) + up_weight = _tensor(env, node.inputs[2]) + down_weight = _tensor(env, node.inputs[3]) + product_scale = ( + node.attrs.get("product_scale") + if "product_scale" in node.attrs + else node.meta.get("product_scale_from_export", node.meta.get("product_scale_elided_by_post_norm", 1.0)) + ) + out = g.dense_mlp_tq_fused( + hidden, + gate_weight, + up_weight, + down_weight, + product_scale=float(product_scale), + ) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op in ("lfm2_moe_layer_gated", "qwen2_moe_layer_gated", "gemma4_moe_layer_gated"): + num_experts = int(node.attrs["num_experts"]) + num_experts_per_tok = int(node.attrs["num_experts_per_tok"]) + w1_start = 3 if op == "lfm2_moe_layer_gated" else 2 + expected_inputs = w1_start + 3 * num_experts + if len(node.inputs) != expected_inputs: + raise NotImplementedError( + f"{op} expects {expected_inputs} inputs for {num_experts} experts, got {len(node.inputs)}" + ) + hidden = _ensure_tensor_dtype(g, _tensor(env, node.inputs[0]), Graph.FP16) + router_logits = _ensure_tensor_dtype(g, _tensor(env, node.inputs[1]), Graph.FP16) + + if op == "lfm2_moe_layer_gated": + routing_probs = g.sigmoid(router_logits) + topk_source = routing_probs + if bool(node.attrs.get("use_expert_bias", False)): + expert_bias = _ensure_tensor_dtype(g, _tensor(env, node.inputs[2]), routing_probs.dtype) + topk_source = _lower_binary_op(g, routing_probs, expert_bias, "add") + topk_indices = _ensure_tensor_dtype(g, g.index(g.topk(topk_source, num_experts_per_tok), 0, axis=0), Graph.FP32) + normalize_default = False + activation = None + elif op == "qwen2_moe_layer_gated": + routing_probs = g.softmax(router_logits, axis=-1) + topk_indices = _ensure_tensor_dtype(g, g.index(g.topk(routing_probs, num_experts_per_tok), 0, axis=0), Graph.FP32) + normalize_default = False + activation = None + else: + hidden = g.scalar_multiply(hidden, float(GEMMA4_RESIDUAL_DOWNSCALE)) + router_logits = g.scalar_multiply(router_logits, float(GEMMA4_RESIDUAL_DOWNSCALE)) + routing_probs = g.softmax(router_logits, axis=-1) + topk_indices = _ensure_tensor_dtype(g, g.index(g.topk(router_logits, num_experts_per_tok), 0, axis=0), Graph.FP32) + normalize_default = True + activation = _activation_enum(node.attrs.get("activation"), default=0) + + w3_start = w1_start + num_experts + w2_start = w3_start + num_experts + moe_kwargs = dict( + normalize_routing=bool(node.attrs.get("normalize_routing", normalize_default)), + epsilon=float(node.attrs.get("epsilon", 1.0e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + if activation is not None: + moe_kwargs["activation"] = activation + out = g.moe_layer_gated( + hidden, + routing_probs, + topk_indices, + [_tensor(env, v) for v in node.inputs[w1_start:w3_start]], + [_tensor(env, v) for v in node.inputs[w3_start:w2_start]], + [_tensor(env, v) for v in node.inputs[w2_start:]], + num_experts, + num_experts_per_tok, + **moe_kwargs, + ) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op == "addmm": + bias = _tensor(env, node.inputs[0]) + lhs = _tensor(env, node.inputs[1]) + rhs = _tensor(env, node.inputs[2]) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else lhs.dtype + matmul_output_dtype = _matmul_output_dtype(lhs, rhs, output_dtype=output_dtype) + if lhs.dtype == Graph.FP16 and rhs.dtype == Graph.FP16 and bias.dtype == Graph.FP16 and output_dtype == Graph.FP16: + matmul_output_dtype = Graph.FP32 + elif bias.dtype == Graph.FP32 and not _is_quantized_runtime_dtype(rhs.dtype): + matmul_output_dtype = Graph.FP32 + out = _matmul_with_quantized_rhs_legalization(g, lhs, rhs, output_dtype=matmul_output_dtype) + out = _trim_padded_last_dim(g, out, output_value.shape if output_value is not None else None) + bias = _ensure_tensor_dtype(g, bias, matmul_output_dtype) + out = g.add(bias, out) + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op == "relu": + return [g.relu(_tensor(env, node.inputs[0]))] + + if op == "silu": + return [g.silu(_tensor(env, node.inputs[0]))] + + if op == "gelu": + return [g.gelu(_tensor(env, node.inputs[0]))] + + if op == "gelu_erf": + return [g.gelu_erf(_tensor(env, node.inputs[0]))] + + if op == "sigmoid": + return [g.sigmoid(_tensor(env, node.inputs[0]))] + + if op == "glu": + return [g.glu(_tensor(env, node.inputs[0]), axis=int(node.attrs.get("axis", -1)))] + + if op == "softplus": + x = _tensor(env, node.inputs[0]) + return [_lower_softplus(g, x)] + + if op == "tanh": + return [g.tanh(_tensor(env, node.inputs[0]))] + + if op == "softmax": + return [g.softmax(_tensor(env, node.inputs[0]), axis=int(node.attrs.get("axis", -1)))] + + if op == "lstm_cell": + x = _tensor(env, node.inputs[0]) + h_prev = _tensor(env, node.inputs[1]) + c_prev = _tensor(env, node.inputs[2]) + weight_ih = _tensor(env, node.inputs[3]) + weight_hh = _tensor(env, node.inputs[4]) + bias_ih = _tensor(env, node.inputs[5]) + bias_hh = _tensor(env, node.inputs[6]) + + lstm_out = g.lstm_cell( + _ensure_fp16_tensor(g, x), + _ensure_fp16_tensor(g, h_prev), + _ensure_fp16_tensor(g, c_prev), + _ensure_fp16_tensor(g, weight_ih), + _ensure_fp16_tensor(g, weight_hh), + _ensure_fp16_tensor(g, bias_ih), + _ensure_fp16_tensor(g, bias_hh), + ) + h_new = g.slice(lstm_out, axis=2, start=0, length=1) + c_new = g.slice(lstm_out, axis=2, start=1, length=1) + h_shape = ir.values[node.outputs[0]].shape + c_shape = ir.values[node.outputs[1]].shape + if h_shape is not None: + h_new = g.reshape(h_new, tuple(int(dim) for dim in h_shape)) + if c_shape is not None: + c_new = g.reshape(c_new, tuple(int(dim) for dim in c_shape)) + return [h_new, c_new] + + if op in {"scaled_dot_product_attention", "attention"}: + if _should_lower_gemma4_decoder_attention_without_kernel(ir, node): + return [_lower_gemma4_decoder_attention_without_kernel(g, ir, env, node)] + mask = node.attrs.get("mask") + mask_tensor: Tensor | None = None + additive_mask = bool(node.attrs.get("additive_mask", False)) + literal_mask_is_causal = False + if len(node.inputs) > 3: + raw_mask = env.get(node.inputs[3]) + if isinstance(raw_mask, bool): + literal_mask_is_causal = bool(raw_mask) + mask_tensor = None + else: + mask_tensor = _tensor(env, node.inputs[3]) + elif mask is not None: + raise NotImplementedError(f"{op} with literal mask is not supported yet") + if op == "scaled_dot_product_attention": + dropout_p = float(node.attrs.get("dropout_p", 0.0)) + if dropout_p != 0.0: + raise NotImplementedError("scaled_dot_product_attention with dropout is not supported yet") + + # PyTorch SDPA exports tensors as [batch, heads, seq, dim], while Cactus attention + # expects [batch, seq, heads, dim]. The optimizer may strip surrounding layout + # permutes and mark the node as already-native to avoid large round-trip copies. + qkv_layout = str(node.attrs.get("qkv_layout", "bhsd") or "bhsd").lower() + q_layout = str(node.attrs.get("q_layout", qkv_layout) or qkv_layout).lower() + k_layout = str(node.attrs.get("k_layout", qkv_layout) or qkv_layout).lower() + v_layout = str(node.attrs.get("v_layout", qkv_layout) or qkv_layout).lower() + output_layout = str(node.attrs.get("output_layout", "bhsd") or "bhsd").lower() + query = _attention_tensor(env, node.inputs[0]) + key = _attention_tensor(env, node.inputs[1]) + value = _attention_tensor(env, node.inputs[2]) + if q_layout != "bthd": + query = g.permute(query, (0, 2, 1, 3)) + if _is_external_cross_kv_cache_value(ir, node.inputs[1]) and _is_external_cross_kv_cache_value(ir, node.inputs[2]): + key_value = ir.values.get(node.inputs[1]) + value_value = ir.values.get(node.inputs[2]) + if key_value is None or value_value is None or key_value.shape is None or value_value.shape is None: + raise ValueError("external cross-KV cache attention requires original key/value shapes") + key_shape = tuple(int(dim) for dim in key_value.shape) + value_shape = tuple(int(dim) for dim in value_value.shape) + if len(key_shape) != 4 or len(value_shape) != 4: + raise ValueError(f"external cross-KV cache attention expects rank-4 key/value tensors, got {key_shape} and {value_shape}") + if k_layout == "bhsd": + num_kv_heads, head_dim = key_shape[1], key_shape[3] + else: + num_kv_heads, head_dim = key_shape[2], key_shape[3] + v_head_dim = value_shape[3] + dummy_key = _materialize_constant_tensor( + g, + torch.zeros((1, 1, int(num_kv_heads), int(head_dim)), dtype=torch.float16), + ) + dummy_value = _materialize_constant_tensor( + g, + torch.zeros((1, 1, int(num_kv_heads), int(v_head_dim)), dtype=torch.float16), + ) + out = g.attention_cached( + query, + dummy_key, + dummy_value, + key, + value, + scale=_resolve_attention_scale(node, query), + position_offset=_DYNAMIC_KV_QUERY_POSITION_OFFSET, + window_size=int(node.attrs.get("window_size", 0)), + v_head_dim=int(v_head_dim), + ) + if output_layout == "bthd": + return [out] + return [g.permute(out, (0, 2, 1, 3))] + if k_layout != "bthd": + key = g.permute(key, (0, 2, 1, 3)) + if v_layout != "bthd": + value = g.permute(value, (0, 2, 1, 3)) + if mask_tensor is not None: + mask_tensor = _ensure_fp16_tensor(g, mask_tensor) + mask_tensor = _normalize_attention_mask_for_cactus(g, mask_tensor, query) + shared_cached = _lower_attention_with_shared_internal_kv_cache( + g, + ir, + env, + node, + query=query, + key=key, + value=value, + output_layout=output_layout, + ) + if shared_cached is not None: + return shared_cached + if _should_lower_attention_with_internal_kv_cache(ir, node): + return _lower_attention_with_internal_kv_cache( + g, + ir, + env, + node, + query=query, + key=key, + value=value, + output_layout=output_layout, + ) + out = g.attention( + query, + key, + value, + scale=_resolve_attention_scale(node, query), + is_causal=bool(node.attrs.get("is_causal", False)) or literal_mask_is_causal, + window_size=int(node.attrs.get("window_size", 0)), + mask=mask_tensor, + additive_mask=additive_mask, + ) + if output_layout == "bthd": + return [out] + return [g.permute(out, (0, 2, 1, 3))] + + if op == "attention_block": + has_mask = bool(node.attrs.get("has_mask", False)) + has_gate = bool(node.attrs.get("has_gate", False)) + has_bias = bool(node.attrs.get("has_bias", False)) + qkv_layout = str(node.attrs.get("qkv_layout", "bhsd") or "bhsd").lower() + input_index = 0 + query = _attention_tensor(env, node.inputs[input_index]) + if qkv_layout != "bthd": + query = g.permute(query, (0, 2, 1, 3)) + input_index += 1 + key = _attention_tensor(env, node.inputs[input_index]) + if qkv_layout != "bthd": + key = g.permute(key, (0, 2, 1, 3)) + input_index += 1 + value = _attention_tensor(env, node.inputs[input_index]) + if qkv_layout != "bthd": + value = g.permute(value, (0, 2, 1, 3)) + input_index += 1 + mask_tensor: Tensor | None = None + literal_mask_is_causal = False + if has_mask: + raw_mask = env.get(node.inputs[input_index]) + input_index += 1 + if isinstance(raw_mask, bool): + literal_mask_is_causal = bool(raw_mask) + mask_tensor = None + else: + mask_tensor = _tensor(env, node.inputs[input_index - 1]) + mask_tensor = _ensure_fp16_tensor(g, mask_tensor) + mask_tensor = _normalize_attention_mask_for_cactus(g, mask_tensor, query) + + attn_out = g.attention( + query, + key, + value, + scale=_resolve_attention_scale(node, query), + is_causal=bool(node.attrs.get("is_causal", False)) or literal_mask_is_causal, + window_size=int(node.attrs.get("window_size", 0)), + mask=mask_tensor, + additive_mask=bool(node.attrs.get("additive_mask", False)), + ) + flat_shape_attr = tuple(int(v) for v in node.attrs.get("attention_output_shape", ())) + flat_shape = _resolve_reshape_shape(tuple(attn_out.shape), flat_shape_attr) if flat_shape_attr else () + if flat_shape: + attn_out = _reshape_attention_output_for_linear(g, attn_out, flat_shape) + + if has_gate: + gate = _tensor(env, node.inputs[input_index]) + input_index += 1 + attn_out, gate = _legalize_elementwise_binary_inputs(g, attn_out, gate) + attn_out = g.multiply(attn_out, gate) + + weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + reshape_back: tuple[int, ...] | None = None + linear_input = attn_out + if len(linear_input.shape) > 2: + linear_input = _flatten_to_2d_for_linear(g, linear_input) + output_value = ir.values.get(node.outputs[0]) + output_shape = output_value.shape if output_value is not None else None + if output_shape is None: + output_shape = node.meta.get("shape") + if not isinstance(output_shape, tuple): + raise NotImplementedError(f"attention_block missing output shape metadata for node {node.id}") + reshape_back = tuple(int(v) for v in output_shape) + + out = _matmul_with_quantized_rhs_legalization(g, linear_input, weight, pretransposed_rhs=True) + if has_bias: + out = g.add(out, _tensor(env, node.inputs[input_index])) + if reshape_back is not None: + out = g.reshape(out, reshape_back) + return [out] + + if op == "self_attention_block": + has_mask = bool(node.attrs.get("has_mask", False)) + has_gate = bool(node.attrs.get("has_gate", False)) + has_bias = bool(node.attrs.get("has_bias", False)) + has_query_projection_bias = bool(node.attrs.get("has_query_projection_bias", False)) + has_query_add = bool(node.attrs.get("has_query_add", False)) + has_rel_query_add = bool(node.attrs.get("has_rel_query_add", False)) + has_key_projection_bias = bool(node.attrs.get("has_key_projection_bias", False)) + has_value_projection_bias = bool(node.attrs.get("has_value_projection_bias", False)) + has_rel_pos_bias = bool(node.attrs.get("has_rel_pos_bias", False)) + has_relative_key_projection_bias = bool(node.attrs.get("has_relative_key_projection_bias", False)) + + input_index = 0 + hidden = _tensor(env, node.inputs[input_index]) + input_index += 1 + + query_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + query_projection_bias = None + if has_query_projection_bias: + query_projection_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + query_add = None + if has_query_add: + query_add = _tensor(env, node.inputs[input_index]) + input_index += 1 + rel_query_add = None + if has_rel_query_add: + rel_query_add = _tensor(env, node.inputs[input_index]) + input_index += 1 + + key_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + key_projection_bias = None + if has_key_projection_bias: + key_projection_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + value_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + value_projection_bias = None + if has_value_projection_bias: + value_projection_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + mask_tensor: Tensor | None = None + if has_mask: + mask_tensor = _tensor(env, node.inputs[input_index]) + input_index += 1 + + relative_key_input = None + relative_key_weight = None + relative_key_projection_bias = None + if has_rel_pos_bias: + relative_key_input = _tensor(env, node.inputs[input_index]) + input_index += 1 + relative_key_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + if has_relative_key_projection_bias: + relative_key_projection_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + gate = None + if has_gate: + gate = _tensor(env, node.inputs[input_index]) + input_index += 1 + + output_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + output_bias = None + if has_bias: + output_bias = _tensor(env, node.inputs[input_index]) + + query_shape = tuple(int(v) for v in node.attrs.get("query_shape", ())) + key_shape = tuple(int(v) for v in node.attrs.get("key_shape", ())) + value_shape = tuple(int(v) for v in node.attrs.get("value_shape", ())) + relative_key_shape = tuple(int(v) for v in node.attrs.get("relative_key_shape", ())) + + query_base = _lower_projected_attention_tensor(g, hidden, query_weight, query_projection_bias, query_shape) + query = query_base + if query_add is not None: + query_add = _normalize_attention_add_tensor(g, query_add, query_shape) + query, query_add = _legalize_elementwise_binary_inputs(g, query, query_add) + query = g.add(query, query_add) + + key = _lower_projected_attention_tensor(g, hidden, key_weight, key_projection_bias, key_shape) + value = _lower_projected_attention_tensor(g, hidden, value_weight, value_projection_bias, value_shape) + + if has_rel_pos_bias: + if relative_key_input is None or relative_key_weight is None or not relative_key_shape: + raise NotImplementedError(f"self_attention_block missing relative position bias inputs for node {node.id}") + rel_query = query_base + if rel_query_add is not None: + rel_query_add = _normalize_attention_add_tensor(g, rel_query_add, query_shape) + rel_query, rel_query_add = _legalize_elementwise_binary_inputs(g, rel_query, rel_query_add) + rel_query = g.add(rel_query, rel_query_add) + relative_key = _lower_projected_attention_tensor( + g, + relative_key_input, + relative_key_weight, + relative_key_projection_bias, + relative_key_shape, + ) + rel_mask = g.rel_pos_bias(rel_query, relative_key, float(node.attrs.get("rel_pos_scale", 1.0))) + if mask_tensor is None: + mask_tensor = rel_mask + else: + mask_tensor, rel_mask = _legalize_elementwise_binary_inputs(g, mask_tensor, rel_mask) + mask_tensor = g.add(mask_tensor, rel_mask) + mask_tensor = _ensure_fp16_tensor(g, mask_tensor) + elif mask_tensor is not None: + mask_tensor = _ensure_fp16_tensor(g, mask_tensor) + if mask_tensor is not None: + mask_tensor = _normalize_attention_mask_for_cactus(g, mask_tensor, query) + + attn_out = g.attention( + query, + key, + value, + scale=_resolve_attention_scale(node, query), + is_causal=bool(node.attrs.get("is_causal", False)), + window_size=int(node.attrs.get("window_size", 0)), + mask=mask_tensor, + additive_mask=bool(node.attrs.get("additive_mask", False)), + ) + flat_shape_attr = tuple(int(v) for v in node.attrs.get("attention_output_shape", ())) + flat_shape = _resolve_reshape_shape(tuple(attn_out.shape), flat_shape_attr) if flat_shape_attr else () + if flat_shape: + attn_out = _reshape_attention_output_for_linear(g, attn_out, flat_shape) + + if gate is not None: + attn_out, gate = _legalize_elementwise_binary_inputs(g, attn_out, gate) + attn_out = g.multiply(attn_out, gate) + + linear_input = attn_out + reshape_back: tuple[int, ...] | None = None + if len(linear_input.shape) > 2: + linear_input = _flatten_to_2d_for_linear(g, linear_input) + output_value = ir.values.get(node.outputs[0]) + output_shape = output_value.shape if output_value is not None else None + if output_shape is None: + output_shape = node.meta.get("shape") + if not isinstance(output_shape, tuple): + raise NotImplementedError(f"self_attention_block missing output shape metadata for node {node.id}") + reshape_back = tuple(int(v) for v in output_shape) + + out = _matmul_with_quantized_rhs_legalization(g, linear_input, output_weight, pretransposed_rhs=True) + if output_bias is not None: + out = g.add(out, output_bias) + if reshape_back is not None: + out = g.reshape(out, reshape_back) + return [out] + + if op in ("sum", "mean", "variance", "min", "max"): + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + axes = _normalize_reduction_axes(node.attrs.get("axis"), len(x.shape)) + keepdim = bool(node.attrs.get("keepdim", False)) + out = _lower_reduction(g, op, x, axes=axes) + if keepdim: + new_shape = list(x.shape) + for reduced_axis in axes: + new_shape[reduced_axis] = 1 + out = g.reshape(out, tuple(new_shape)) + return [out] + + if op == "cumsum": + x = _ensure_scalar_math_tensor(g, _tensor(env, node.inputs[0])) + raw_axis = node.attrs.get("axis") + if not isinstance(raw_axis, int): + args = node.attrs.get("args", ()) + if isinstance(args, (list, tuple)) and len(args) > 1 and isinstance(args[1], int): + raw_axis = int(args[1]) + else: + raw_axis = -1 + axis = _normalize_dim(int(raw_axis), len(x.shape)) + return [g.cumsum(x, axis)] + + if op == "cat": + tensors = [_tensor(env, value_id) for value_id in node.inputs] + output_dtype: int | None = None + if node.outputs: + output_value = ir.values.get(node.outputs[0]) + if output_value is not None and output_value.dtype is not None: + output_dtype = _map_ir_or_torch_dtype(output_value.dtype) + return [_cat_with_legalized_dtype(g, tensors, axis=int(node.attrs.get("axis", 0)), dtype=output_dtype)] + + if op in {"stack", "aten.stack.default"}: + tensors = [_tensor(env, value_id) for value_id in node.inputs] + if not tensors: + raise NotImplementedError("stack requires at least one input tensor") + + raw_axis = node.attrs.get("axis") + if not isinstance(raw_axis, int): + args = node.attrs.get("args", ()) + if isinstance(args, (list, tuple)) and len(args) > 1 and isinstance(args[1], int): + raw_axis = int(args[1]) + else: + raw_axis = 0 + + base_shape = tuple(int(dim) for dim in tensors[0].shape) + axis = _normalize_dim(int(raw_axis), len(base_shape) + 1) + reshaped: list[Tensor] = [] + for tensor in tensors: + if tuple(int(dim) for dim in tensor.shape) != base_shape: + raise NotImplementedError( + f"stack requires equal input shapes, got {[tuple(int(dim) for dim in t.shape) for t in tensors]}" + ) + expanded_shape = list(base_shape) + expanded_shape.insert(axis, 1) + reshaped.append(g.reshape(tensor, tuple(expanded_shape))) + return [g.cat(reshaped, axis=axis)] + + if op == "slice": + source = env[node.inputs[0]] + if isinstance(source, BroadcastAlias): + x = _materialize_broadcast_alias(g, source) + else: + x = _tensor(env, node.inputs[0]) + axis = _normalize_dim(int(node.attrs["axis"]), len(x.shape)) + start = int(node.attrs["start"]) + end = int(node.attrs["end"]) + step = int(node.attrs.get("step", 1)) + dim_size = x.shape[axis] + if step == 0: + raise NotImplementedError("slice with step == 0 is invalid") + if step != 1: + indices = list(range(*slice(start, end, step).indices(dim_size))) + if axis == 0 or x.dtype == Graph.FP16: + return [_lower_static_strided_slice_via_gather(g, x, axis=axis, indices=indices)] + if not indices: + return [g.slice(x, axis=axis, start=0, length=0)] + expanded_shape = list(int(dim) for dim in x.shape) + expanded_shape[axis] = 1 + pieces: list[Tensor] = [] + for idx in indices: + piece = g.index(x, index_value=idx, axis=axis) + piece = g.reshape(piece, tuple(expanded_shape)) + pieces.append(piece) + if len(pieces) == 1: + return [pieces[0]] + return [g.cat(pieces, axis=axis)] + start = _normalize_index(start, dim_size) + end = _normalize_slice_end(end, dim_size) + length = max(0, end - start) + return [g.slice(x, axis=axis, start=start, length=length)] + + if op == "split_with_sizes": + x = _tensor(env, node.inputs[0]) + axis = _normalize_dim(int(node.attrs.get("axis", -1)), len(x.shape)) + sizes = tuple(int(v) for v in node.attrs["sizes"]) + start = 0 + outputs: list[Tensor] = [] + for size in sizes: + outputs.append(g.slice(x, axis=axis, start=start, length=int(size))) + start += int(size) + return [outputs] + + if op == "chunk": + x = _tensor(env, node.inputs[0]) + axis = _normalize_dim(int(node.attrs.get("axis", 0)), len(x.shape)) + chunks = int(node.attrs["chunks"]) + if chunks <= 0: + raise NotImplementedError(f"chunk requires a positive chunk count, got {chunks}") + dim_size = int(x.shape[axis]) + if dim_size <= 0: + return [[g.slice(x, axis=axis, start=0, length=0)]] + chunk_size = (dim_size + chunks - 1) // chunks + outputs: list[Tensor] = [] + start = 0 + while start < dim_size: + length = min(chunk_size, dim_size - start) + outputs.append(g.slice(x, axis=axis, start=start, length=length)) + start += length + return [outputs] + + if op in {"unbind", "aten.unbind.int"}: + x = _tensor(env, node.inputs[0]) + axis = _normalize_dim(int(node.attrs.get("axis", 0)), len(x.shape)) + outputs: list[Tensor] = [] + out_shape = tuple(int(dim) for index, dim in enumerate(x.shape) if index != axis) + for index in range(int(x.shape[axis])): + piece = g.slice(x, axis=axis, start=index, length=1) + outputs.append(g.reshape(piece, out_shape)) + return [outputs] + + if op == "ones": + shape = tuple(int(v) for v in node.attrs.get("shape", ())) + if not shape: + raise NotImplementedError(f"ones requires a static shape, got {shape}") + torch_dtype = _materialize_constant_torch_dtype(node.attrs.get("dtype")) + return [_materialize_constant_tensor(g, torch.ones(shape, dtype=torch_dtype))] + + if op == "new_empty": + shape = tuple(int(v) for v in node.attrs.get("shape", ())) + if not shape: + raise NotImplementedError(f"new_empty requires a static shape, got {shape}") + torch_dtype = _materialize_constant_torch_dtype(node.attrs.get("dtype")) + return [_materialize_constant_tensor(g, torch.zeros(shape, dtype=torch_dtype))] + + if op == "tril": + x = _tensor(env, node.inputs[0]) + shape = tuple(int(dim) for dim in x.shape) + if len(shape) < 2: + raise NotImplementedError(f"tril expects rank >= 2, got shape {shape}") + diagonal = int(node.attrs.get("diagonal", 0)) + mask_dtype = _torch_dtype_for_graph_dtype(x.dtype) + mask = torch.ones(shape, dtype=mask_dtype).tril(diagonal=diagonal) + mask_tensor = _materialize_constant_tensor(g, mask) + return [_lower_binary_op(g, x, mask_tensor, "multiply")] + + if op == "unfold": + x = _tensor(env, node.inputs[0]) + rank = len(x.shape) + axis = _normalize_dim(int(node.attrs["dimension"]), rank) + size = int(node.attrs["size"]) + step = int(node.attrs["step"]) + if size <= 0: + raise NotImplementedError(f"unfold requires a positive size, got {size}") + if step <= 0: + raise NotImplementedError(f"unfold requires a positive step, got {step}") + dim_size = int(x.shape[axis]) + if size > dim_size: + raise NotImplementedError( + f"unfold size {size} exceeds dimension size {dim_size} for axis {axis}" + ) + + window_starts = list(range(0, dim_size - size + 1, step)) + if not window_starts: + raise NotImplementedError( + f"unfold produced no windows for axis={axis}, size={size}, step={step}, dim_size={dim_size}" + ) + + base_shape = [int(dim) for dim in x.shape] + trailing_shape = [base_shape[idx] for idx in range(rank) if idx != axis] + [size] + expanded_shape = tuple(trailing_shape[:axis] + [1] + trailing_shape[axis:]) + permutation = tuple(idx for idx in range(rank) if idx != axis) + (axis,) + + windows: list[Tensor] = [] + for start in window_starts: + piece = g.slice(x, axis=axis, start=start, length=size) + if axis != rank - 1: + piece = g.permute(piece, permutation=permutation) + piece = g.reshape(piece, expanded_shape) + windows.append(piece) + if len(windows) == 1: + return [windows[0]] + return [g.cat(windows, axis=axis)] + + if op == "pad": + x = _tensor(env, node.inputs[0]) + mode = str(node.attrs.get("mode", "constant")) + if mode != "constant": + raise NotImplementedError(f"pad mode is unsupported: {mode}") + pads = tuple(int(v) for v in node.attrs.get("pads", ())) + if len(pads) % 2 != 0: + raise NotImplementedError(f"pad expects an even-length pads tuple, got {pads}") + value = float(node.attrs.get("value", 0.0)) + + current = x + current_shape = list(int(dim) for dim in x.shape) + pad_dims = len(pads) // 2 + if pad_dims > len(current_shape): + raise NotImplementedError( + f"pad rank mismatch: pads={pads} for input shape {tuple(current_shape)}" + ) + + for pad_index in range(pad_dims): + before = pads[2 * pad_index] + after = pads[2 * pad_index + 1] + axis = len(current_shape) - 1 - pad_index + if before < 0 or after < 0: + raise NotImplementedError(f"negative pad is unsupported: {pads}") + pieces = [] + torch_dtype = _torch_dtype_for_graph_dtype(current.dtype) + if before > 0: + left_shape = list(current_shape) + left_shape[axis] = before + left = _materialize_constant_tensor( + g, + torch.full(tuple(left_shape), value, dtype=torch_dtype), + ) + pieces.append(left) + pieces.append(current) + if after > 0: + right_shape = list(current_shape) + right_shape[axis] = after + right = _materialize_constant_tensor( + g, + torch.full(tuple(right_shape), value, dtype=torch_dtype), + ) + pieces.append(right) + if len(pieces) > 1: + current = g.cat(pieces, axis=axis) + current_shape[axis] += before + after + return [current] + + if op == "index": + x = _tensor(env, node.inputs[0]) + axis = _normalize_dim(int(node.attrs.get("axis", 0)), len(x.shape)) + index_value = _normalize_index(int(node.attrs["index_value"]), x.shape[axis]) + return [g.index(x, index_value=index_value, axis=axis)] + + if op == "gather": + return [ + g.gather( + _tensor(env, node.inputs[0]), + _tensor(env, node.inputs[1]), + axis=int(node.attrs.get("axis", 0)), + ) + ] + + if op == "embedding": + embedding_tensor = _tensor(env, node.inputs[0]) + indices_tensor = _tensor(env, node.inputs[1]) + _debug_embedding_lowering(node, embedding_tensor, indices_tensor) + out = g.embedding_from_tensor(embedding_tensor, indices_tensor) + return [out] + + if op == "conv_module": + input_index = 0 + x_nlc = _tensor(env, node.inputs[input_index]) + input_index += 1 + + pointwise1_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + pointwise1_bias = None + if bool(node.attrs.get("has_pointwise1_bias", False)): + pointwise1_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + depthwise_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + depthwise_bias = None + if bool(node.attrs.get("has_depthwise_bias", False)): + depthwise_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + + batch_norm_weight = _tensor(env, node.inputs[input_index]) + batch_norm_bias = _tensor(env, node.inputs[input_index + 1]) + batch_norm_running_mean = _tensor(env, node.inputs[input_index + 2]) + batch_norm_running_var = _tensor(env, node.inputs[input_index + 3]) + input_index += 4 + + pointwise2_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + pointwise2_bias = None + if bool(node.attrs.get("has_pointwise2_bias", False)): + pointwise2_bias = _tensor(env, node.inputs[input_index]) + + if len(x_nlc.shape) != 3: + raise NotImplementedError(f"conv_module expects rank-3 NLC input, got {x_nlc.shape}") + kernel_size = int(node.attrs.get("depthwise_kernel_size", 0)) + padding = int(node.attrs.get("depthwise_padding", 0)) + if len(depthwise_weight.shape) != 3 or kernel_size != 9 or padding != 4: + raise NotImplementedError( + f"conv_module currently supports same depthwise kernel-9 lowering, got " + f"kernel_size={kernel_size} padding={padding} weight_shape={depthwise_weight.shape}" + ) + + current = g.conv1d_pointwise(x_nlc, pointwise1_weight, bias=pointwise1_bias) + current = g.glu(current, axis=-1) + current = g.conv1d_same_depthwise_k9(current, depthwise_weight, bias=depthwise_bias) + current = g.batch_norm( + current, + batch_norm_weight, + batch_norm_bias, + batch_norm_running_mean, + batch_norm_running_var, + axis=len(current.shape) - 1, + eps=float(node.attrs.get("eps", 1e-5)), + ) + current = g.silu(current) + current = g.conv1d_pointwise(current, pointwise2_weight, bias=pointwise2_bias) + return [current] + + if op == "conv1d": + x = _tensor(env, node.inputs[0]) + weight = _tensor(env, node.inputs[1]) + bias = _tensor(env, node.inputs[2]) if len(node.inputs) > 2 else None + stride = int(node.attrs.get("stride", 1)) + padding = int(node.attrs.get("padding", 0)) + dilation = int(node.attrs.get("dilation", 1)) + groups = int(node.attrs.get("groups", 1)) + + if _should_lower_conv1d_with_internal_conv_cache(ir, node, x, weight): + return _lower_conv1d_with_internal_conv_cache( + g, + ir, + env, + node, + x=x, + weight=weight, + bias=bias, + ) + + if ( + len(x.shape) == 3 + and len(weight.shape) == 3 + and weight.shape[2] == 1 + and stride == 1 + and padding == 0 + and dilation == 1 + and groups == 1 + ): + x_nlc = g.permute(x, (0, 2, 1)) + out_nlc = g.conv1d_pointwise(x_nlc, weight, bias=bias) + return [g.permute(out_nlc, (0, 2, 1))] + + if ( + len(x.shape) == 3 + and len(weight.shape) == 3 + and groups == x.shape[1] == weight.shape[0] + and weight.shape[1] == 1 + and weight.shape[2] == 9 + and stride == 1 + and padding == 4 + and dilation == 1 + ): + x_nlc = g.permute(x, (0, 2, 1)) + out_nlc = g.conv1d_same_depthwise_k9(x_nlc, weight, bias=bias) + return [g.permute(out_nlc, (0, 2, 1))] + + if ( + len(x.shape) == 3 + and len(weight.shape) == 3 + and groups == x.shape[1] == weight.shape[0] + and weight.shape[1] == 1 + and stride == 1 + and padding == dilation * max(weight.shape[2] - 1, 0) + ): + x_nlc = g.permute(x, (0, 2, 1)) + out_nlc = g.conv1d_causal(x_nlc, weight, kernel_size=weight.shape[2], dilation=dilation) + if bias is not None: + bias_reshaped = g.reshape(bias, (1, 1, int(weight.shape[0]))) + out_nlc, bias_reshaped = _legalize_elementwise_binary_inputs(g, out_nlc, bias_reshaped) + out_nlc = g.add(out_nlc, bias_reshaped) + return [g.permute(out_nlc, (0, 2, 1))] + + if ( + len(x.shape) == 3 + and len(weight.shape) == 3 + and groups == 1 + and dilation == 1 + and weight.shape[2] == 3 + and padding == 1 + and stride in {1, 2} + ): + out = g.conv1d_k3(x, weight, stride=stride) + if bias is not None: + bias_reshaped = g.reshape(bias, (1, int(weight.shape[0]), 1)) + out, bias_reshaped = _legalize_elementwise_binary_inputs(g, out, bias_reshaped) + out = g.add(out, bias_reshaped) + return [out] + + if dilation != 1: + raise NotImplementedError(f"conv1d with dilation != 1 is unsupported by generic lowering: {dilation}") + if padding != 0: + if len(x.shape) != 3: + raise NotImplementedError(f"generic conv1d padding expects rank-3 input, got {x.shape}") + batch_size, channels, _ = (int(dim) for dim in x.shape) + pad_dtype = _torch_dtype_for_graph_dtype(x.dtype) + pieces: list[Tensor] = [] + if padding > 0: + left = _materialize_constant_tensor( + g, + torch.zeros((batch_size, channels, padding), dtype=pad_dtype), + ) + pieces.append(left) + pieces.append(x) + if padding > 0: + right = _materialize_constant_tensor( + g, + torch.zeros((batch_size, channels, padding), dtype=pad_dtype), + ) + pieces.append(right) + x = g.cat(pieces, axis=2) + if groups != 1: + if len(x.shape) != 3 or len(weight.shape) != 3: + raise NotImplementedError( + f"grouped conv1d lowering expects rank-3 input and weight, got {x.shape} and {weight.shape}" + ) + input_channels = int(x.shape[1]) + if input_channels % groups != 0: + raise NotImplementedError( + f"grouped conv1d requires input channels divisible by groups, got C_in={input_channels}, groups={groups}" + ) + output_channels = int(weight.shape[0]) + if output_channels % groups != 0: + raise NotImplementedError( + f"grouped conv1d requires output channels divisible by groups, got C_out={output_channels}, groups={groups}" + ) + input_channels_per_group = input_channels // groups + output_channels_per_group = output_channels // groups + if int(weight.shape[1]) != input_channels_per_group: + raise NotImplementedError( + "grouped conv1d weight has incompatible per-group input channels: " + f"weight C_in/group={int(weight.shape[1])}, expected {input_channels_per_group}" + ) + outputs: list[Tensor] = [] + for group_index in range(groups): + x_group = g.slice( + x, + axis=1, + start=group_index * input_channels_per_group, + length=input_channels_per_group, + ) + weight_group = g.slice( + weight, + axis=0, + start=group_index * output_channels_per_group, + length=output_channels_per_group, + ) + bias_group = None + if bias is not None: + bias_group = g.slice( + bias, + axis=0, + start=group_index * output_channels_per_group, + length=output_channels_per_group, + ) + outputs.append(g.conv1d(x_group, weight_group, bias=bias_group, stride=stride)) + if len(outputs) == 1: + return [outputs[0]] + return [g.cat(outputs, axis=1)] + return [g.conv1d(x, weight, bias=bias, stride=stride)] + + if op == "conv2d": + x = _tensor(env, node.inputs[0]) + weight = _tensor(env, node.inputs[1]) + bias = _tensor(env, node.inputs[2]) if len(node.inputs) > 2 else None + stride = int(node.attrs.get("stride", 1)) + padding = int(node.attrs.get("padding", 0)) + dilation = int(node.attrs.get("dilation", 1)) + groups = int(node.attrs.get("groups", 1)) + + if len(x.shape) != 4 or len(weight.shape) != 4: + raise NotImplementedError(f"conv2d lowering expects rank-4 tensors, got {x.shape} and {weight.shape}") + if dilation != 1: + raise NotImplementedError(f"conv2d with dilation != 1 is unsupported: {dilation}") + + if stride == 2 and padding == 0 and groups == 1 and weight.shape[2:] == (3, 3): + batch_size, channels, height, width = (int(dim) for dim in x.shape) + pad_dtype = _torch_dtype_for_graph_dtype(x.dtype) + top = _materialize_constant_tensor( + g, + torch.zeros((batch_size, channels, 1, width), dtype=pad_dtype), + ) + bottom = _materialize_constant_tensor( + g, + torch.zeros((batch_size, channels, 1, width), dtype=pad_dtype), + ) + x = g.cat([top, x, bottom], axis=2) + left = _materialize_constant_tensor( + g, + torch.zeros((batch_size, channels, height + 2, 1), dtype=pad_dtype), + ) + right = _materialize_constant_tensor( + g, + torch.zeros((batch_size, channels, height + 2, 1), dtype=pad_dtype), + ) + x = g.cat([left, x, right], axis=3) + y = g.conv2d_k3s2p1(x, weight, bias=bias) + output_height = ((height - 3) // 2) + 1 + output_width = ((width - 3) // 2) + 1 + y = g.slice(y, axis=2, start=1, length=output_height) + y = g.slice(y, axis=3, start=1, length=output_width) + return [y] + + if stride == 2 and padding == 1 and groups == 1 and weight.shape[2:] == (3, 3): + return [g.conv2d_k3s2p1(x, weight, bias=bias)] + + if ( + stride == 2 + and padding == 1 + and groups == x.shape[1] == weight.shape[0] + and weight.shape[1] == 1 + and weight.shape[2:] == (3, 3) + ): + return [g.conv2d_depthwise_k3s2p1(x, weight, bias=bias)] + + if ( + stride == 1 + and padding == 0 + and groups == 1 + and weight.shape[2:] == (1, 1) + ): + return [g.conv2d_pointwise_1x1(x, weight, bias=bias)] + + if stride == 1 and padding == 1 and groups == 1 and weight.shape[2:] == (3, 3): + return [g.conv2d_k3s1p1(x, weight, bias=bias)] + + return [g.conv2d(x, weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)] + + if op == "layer_norm": + x = _tensor(env, node.inputs[0]) + weight = _tensor(env, node.inputs[1]) + bias = _tensor(env, node.inputs[2]) if len(node.inputs) > 2 else None + return [g.layer_norm(x, weight, bias=bias, eps=float(node.attrs["eps"]))] + + if op == "rms_norm": + x = _tensor(env, node.inputs[0]) + weight = _tensor(env, node.inputs[1]) + weight_offset = float(node.attrs.get("weight_offset", 0.0) or 0.0) + if weight_offset != 0.0: + weight = g.scalar_add(weight, weight_offset) + reshape_back: tuple[int, ...] | None = None + if len(x.shape) > 2: + reshape_back = tuple(int(dim) for dim in x.shape) + x = _flatten_to_2d_for_linear(g, x) + out = g.rms_norm(x, weight, eps=float(node.attrs["eps"])) + if reshape_back is not None: + out = g.reshape(out, reshape_back) + return [out] + + if op == "rope": + rope_input = _tensor(env, node.inputs[0]) + if _rope_input_is_bhsd(ir, node.inputs[0]): + rope_input = g.permute(rope_input, (0, 2, 1, 3)) + rope_out = g.rope( + rope_input, + float(node.attrs["theta"]), + position_offset=int(node.attrs.get("position_offset", 0)), + ) + return [g.permute(rope_out, (0, 2, 1, 3))] + return [ + g.rope( + rope_input, + float(node.attrs["theta"]), + position_offset=int(node.attrs.get("position_offset", 0)), + ) + ] + + if op == "rel_pos_bias": + query = _attention_tensor(env, node.inputs[0]) + if len(query.shape) != 4: + raise NotImplementedError(f"rel_pos_bias expects rank-4 query input, got {query.shape}") + relative_key = _tensor(env, node.inputs[1]) + if len(relative_key.shape) != 4: + raise NotImplementedError(f"rel_pos_bias expects rank-4 relative_key input, got {relative_key.shape}") + if int(query.shape[2]) == int(relative_key.shape[2]) and int(query.shape[3]) == int(relative_key.shape[3]): + pass + elif int(query.shape[1]) == int(relative_key.shape[2]) and int(query.shape[3]) == int(relative_key.shape[3]): + query = g.permute(query, (0, 2, 1, 3)) + else: + raise NotImplementedError( + "rel_pos_bias query/relative_key layout mismatch: " + f"query={query.shape}, relative_key={relative_key.shape}" + ) + return [g.rel_pos_bias(query, relative_key, float(node.attrs.get("scale", 1.0)))] + + if op in {"gated_deltanet_prefill", "gated_deltanet_decode"}: + x = _tensor(env, node.inputs[0]) + qkv_weight = _tensor(env, node.inputs[1]) + a_weight = _tensor(env, node.inputs[2]) + b_weight = _tensor(env, node.inputs[3]) + norm_weight = _tensor(env, node.inputs[4]) + + input_index = 5 + z_weight = None + if bool(node.attrs.get("has_z", False)): + z_weight = _tensor(env, node.inputs[input_index]) + input_index += 1 + dt_bias = None + if bool(node.attrs.get("has_dt_bias", False)): + dt_bias = _tensor(env, node.inputs[input_index]) + input_index += 1 + a_log = None + if bool(node.attrs.get("has_a_log", False)): + a_log = _tensor(env, node.inputs[input_index]) + input_index += 1 + conv_weight = None + if bool(node.attrs.get("has_conv", False)): + conv_weight = _tensor(env, node.inputs[input_index]) + + if len(x.shape) != 3: + raise NotImplementedError(f"{op} currently expects rank-3 normalized input, got {x.shape}") + + batch_size, seq_len, hidden_dim = (int(dim) for dim in x.shape) + if batch_size != 1: + raise NotImplementedError(f"{op} currently supports batch size 1, got {x.shape}") + + num_k_heads = int(node.attrs["num_k_heads"]) + num_v_heads = int(node.attrs["num_v_heads"]) + key_dim = int(node.attrs["key_dim"]) + value_dim = int(node.attrs["value_dim"]) + eps = float(node.attrs.get("eps", 1e-6)) + chunk_size = int(node.attrs.get("chunk_size", 64)) + + mixed_qkv = _matmul_with_quantized_rhs_legalization( + g, + _flatten_to_2d_for_linear(g, x), + qkv_weight, + pretransposed_rhs=True, + ) + mixed_qkv_dim = int(qkv_weight.shape[0]) + mixed_qkv = g.reshape(mixed_qkv, (batch_size, seq_len, mixed_qkv_dim)) + if conv_weight is not None: + mixed_qkv = _lower_gated_deltanet_conv1d( + g, + ir, + env, + node, + mixed_qkv=mixed_qkv, + conv_weight=conv_weight, + batch_size=batch_size, + seq_len=seq_len, + mixed_qkv_dim=mixed_qkv_dim, + op=op, + ) + mixed_qkv = g.silu(mixed_qkv) + + q_proj_dim = num_k_heads * key_dim + v_proj_dim = num_v_heads * value_dim + k_proj_dim = num_k_heads * key_dim + q_proj = g.slice(mixed_qkv, axis=2, start=0, length=q_proj_dim) + k_proj = g.slice(mixed_qkv, axis=2, start=q_proj_dim, length=k_proj_dim) + v_proj = g.slice(mixed_qkv, axis=2, start=q_proj_dim + k_proj_dim, length=v_proj_dim) + + q_4d = g.reshape(q_proj, (batch_size, seq_len, num_k_heads, key_dim)) + k_4d = g.reshape(k_proj, (batch_size, seq_len, num_k_heads, key_dim)) + v_4d = g.reshape(v_proj, (batch_size, seq_len, num_v_heads, value_dim)) + + q_norm = g.sum(g.multiply(q_4d, q_4d), axis=3) + q_norm = g.scalar_sqrt(g.scalar_add(q_norm, eps)) + q_norm = g.reshape(q_norm, (batch_size, seq_len, num_k_heads, 1)) + q_4d = g.divide(q_4d, q_norm) + + k_norm = g.sum(g.multiply(k_4d, k_4d), axis=3) + k_norm = g.scalar_sqrt(g.scalar_add(k_norm, eps)) + k_norm = g.reshape(k_norm, (batch_size, seq_len, num_k_heads, 1)) + k_4d = g.divide(k_4d, k_norm) + + a_logits = _matmul_with_quantized_rhs_legalization( + g, + _flatten_to_2d_for_linear(g, x), + a_weight, + pretransposed_rhs=True, + ) + a_logits = g.reshape(a_logits, (batch_size, seq_len, int(a_weight.shape[0]))) + b_logits = _matmul_with_quantized_rhs_legalization( + g, + _flatten_to_2d_for_linear(g, x), + b_weight, + pretransposed_rhs=True, + ) + b_logits = g.reshape(b_logits, (batch_size, seq_len, int(b_weight.shape[0]))) + + if dt_bias is not None: + dt_bias_2d = g.reshape(dt_bias, (1, int(dt_bias.shape[0]))) + a_logits, dt_bias_2d = _legalize_elementwise_binary_inputs(g, a_logits, dt_bias_2d) + a_logits = g.add(a_logits, dt_bias_2d) + a_softplus = _lower_softplus(g, a_logits) + + if a_log is not None: + a_log_2d = g.reshape(a_log, (1, int(a_log.shape[0]))) + neg_exp_a = g.scalar_multiply(g.scalar_exp(a_log_2d), -1.0) + neg_exp_a, a_softplus = _legalize_elementwise_binary_inputs(g, neg_exp_a, a_softplus) + gate_log = g.multiply(neg_exp_a, a_softplus) + else: + gate_log = g.scalar_multiply(a_softplus, -1.0) + beta = g.sigmoid(b_logits) + + cache_layer_key, initial_state = _lower_gated_deltanet_initial_state( + g, + ir, + env, + node, + batch_size=batch_size, + key_dim=key_dim, + num_v_heads=num_v_heads, + value_dim=value_dim, + seq_len=seq_len, + op=op, + ) + + deltanet_scale = 1.0 / math.sqrt(float(key_dim)) + if seq_len == 1: + deltanet_out = g.gated_deltanet_decode(q_4d, k_4d, v_4d, gate_log, beta, initial_state, deltanet_scale) + else: + deltanet_out = g.gated_deltanet_prefill(q_4d, k_4d, v_4d, gate_log, beta, initial_state, chunk_size, deltanet_scale) + if cache_layer_key is not None: + final_state = g.slice(deltanet_out, axis=1, start=seq_len, length=key_dim) + g.recurrent_cache_write(final_state, initial_state) + cache_entries: list[tuple[str, Tensor, Tensor]] = env.setdefault("__internal_kv_cache_state_entries", []) # type: ignore[assignment] + cache_entries.append((cache_layer_key, initial_state, initial_state)) + + y_4d = g.slice(deltanet_out, axis=1, start=0, length=seq_len) + y_2d = g.reshape(y_4d, (seq_len * num_v_heads, value_dim)) + + if z_weight is not None: + z_proj = _matmul_with_quantized_rhs_legalization( + g, + _flatten_to_2d_for_linear(g, x), + z_weight, + pretransposed_rhs=True, + ) + z_proj = g.reshape(z_proj, (seq_len * num_v_heads, value_dim)) + y_2d = g.multiply(g.rms_norm(y_2d, norm_weight, eps=eps), g.silu(z_proj)) + + return [g.reshape(y_2d, (batch_size, seq_len, num_v_heads * value_dim))] + + if op == "group_norm": + return [ + g.group_norm( + _tensor(env, node.inputs[0]), + _tensor(env, node.inputs[1]), + _tensor(env, node.inputs[2]), + num_groups=int(node.attrs["num_groups"]), + eps=float(node.attrs["eps"]), + ) + ] + + if op == "batch_norm": + axis = int(node.attrs.get("axis", 1)) + return [ + g.batch_norm( + _tensor(env, node.inputs[0]), + _tensor(env, node.inputs[1]), + _tensor(env, node.inputs[2]), + _tensor(env, node.inputs[3]), + _tensor(env, node.inputs[4]), + axis=axis, + eps=float(node.attrs["eps"]), + ) + ] + + if op == "identity": + return [env[node.inputs[0]]] + + if op == "contiguous": + return [_tensor(env, node.inputs[0])] + + if op == "advanced_index": + lowered = _try_lower_advanced_index(g, node, env, ir) + if lowered is not None: + return [lowered] + raise NotImplementedError(f"unsupported advanced_index pattern for node {node.id}") + + if op in {"masked_scatter", "aten.masked_scatter.default"}: + base = _tensor(env, node.inputs[0]) + mask = _tensor(env, node.inputs[1]) + source = _tensor(env, node.inputs[2]) + if source.dtype != base.dtype: + source = _ensure_tensor_dtype(g, source, base.dtype) + return [g.masked_scatter(base, mask, source)] + + if op == "masked_fill": + base = _tensor(env, node.inputs[0]) + mask = _lower_compare_op(g, env[node.inputs[1]], 0.0, "not_equal") + inverse_mask = g.scalar_add(g.scalar_multiply(mask, -1.0), 1.0) + fill_value = _clamp_scalar_for_dtype(float(node.attrs["value"]), base.dtype) + fill_term = g.scalar_multiply(_ensure_tensor_dtype(g, mask, base.dtype), fill_value) + base_term = _lower_binary_op(g, base, _ensure_tensor_dtype(g, inverse_mask, base.dtype), "multiply") + return [_lower_binary_op(g, fill_term, base_term, "add")] + + if op == "getitem": + source = env[node.inputs[0]] + index = int(node.attrs["index"]) + if isinstance(source, (tuple, list)): + return [source[index]] + if isinstance(source, Tensor): + normalized_index = _normalize_index(index, int(source.shape[0])) + return [g.index(source, normalized_index, axis=0)] + raise NotImplementedError(f"getitem source is not tuple/list for node {node.id}") + + raise NotImplementedError(f"unsupported IR op in lowering: {op}") + + +def _tensor(env: dict[str, Any], value_id: str) -> Tensor: + try: + value = env[value_id] + except KeyError as exc: + raise NotImplementedError(f"missing IR value during lowering: {value_id}") from exc + if not isinstance(value, Tensor): + raise TypeError(f"expected lowered tensor for {value_id}, got {type(value).__name__}") + return value + + +def _tensor_or_materialized_alias(g: Graph, env: dict[str, Any], value_id: str) -> Tensor: + value = env.get(value_id) + if isinstance(value, BroadcastAlias): + return _materialize_broadcast_alias(g, value) + return _tensor(env, value_id) + + +def _materialize_broadcast_alias(g: Graph, value: BroadcastAlias) -> Tensor: + if value.kind != "gqa_repeat_kv": + raise TypeError(f"unsupported broadcast alias: {value.kind}") + base_shape = tuple(int(dim) for dim in value.tensor.shape) + logical_shape = tuple(int(dim) for dim in value.logical_shape) + if len(base_shape) != 4 or len(logical_shape) not in {4, 5}: + raise TypeError(f"gqa_repeat_kv alias expects 4D->4D/5D shape, got {base_shape} -> {logical_shape}") + batch, kv_heads, seq_len, head_dim = base_shape + if len(logical_shape) == 5: + if logical_shape[0] != batch or logical_shape[1] != kv_heads or logical_shape[3] != seq_len or logical_shape[4] != head_dim: + raise TypeError(f"gqa_repeat_kv alias shape mismatch: {base_shape} -> {logical_shape}") + base = g.reshape(value.tensor, (batch, kv_heads, 1, seq_len, head_dim)) + return _lower_repeat(g, base, (1, 1, logical_shape[2], 1, 1)) + if logical_shape[0] != batch or logical_shape[2] != seq_len or logical_shape[3] != head_dim: + raise TypeError(f"gqa_repeat_kv alias shape mismatch: {base_shape} -> {logical_shape}") + if logical_shape[1] % max(kv_heads, 1) != 0: + raise TypeError(f"gqa_repeat_kv head count mismatch: {base_shape} -> {logical_shape}") + repeats = logical_shape[1] // max(kv_heads, 1) + expanded = g.reshape(value.tensor, (batch, kv_heads, 1, seq_len, head_dim)) + repeated = _lower_repeat(g, expanded, (1, 1, repeats, 1, 1)) + return g.reshape(repeated, logical_shape) + + +def _attention_tensor(env: dict[str, Any], value_id: str) -> Tensor: + try: + value = env[value_id] + except KeyError as exc: + raise NotImplementedError(f"missing IR value during lowering: {value_id}") from exc + if isinstance(value, BroadcastAlias): + if value.kind != "gqa_repeat_kv": + raise TypeError(f"unsupported broadcast alias for attention input {value_id}: {value.kind}") + return value.tensor + if not isinstance(value, Tensor): + raise TypeError(f"expected lowered tensor for {value_id}, got {type(value).__name__}") + return value + + +def _resolve_attention_scale(node: IRNode, query: Tensor) -> float: + scale = float(node.attrs.get("scale", 0.0) or 0.0) + if scale > 0.0: + return scale + if len(query.shape) >= 1 and int(query.shape[-1]) > 0: + return float(int(query.shape[-1]) ** -0.5) + return 1.0 + + +def _normalize_attention_mask_for_cactus(g: Graph, mask: Tensor, query: Tensor) -> Tensor: + """Convert broadcastable HF masks into the native Cactus attention contract.""" + + mask_shape = tuple(int(dim) for dim in mask.shape) + query_shape = tuple(int(dim) for dim in query.shape) + if len(query_shape) != 4: + return mask + + batch_size, seq_len, num_heads, _ = query_shape + if len(mask_shape) == 2 and batch_size == 1 and mask_shape[0] == seq_len: + return g.reshape(mask, (1, mask_shape[0], mask_shape[1])) + + if len(mask_shape) == 4: + if ( + mask_shape[0] == batch_size + and mask_shape[1] == 1 + and mask_shape[2] == seq_len + ): + return g.reshape(mask, (mask_shape[0], mask_shape[2], mask_shape[3])) + if ( + mask_shape[0] == batch_size + and mask_shape[1] == num_heads + and mask_shape[2] == seq_len + ): + return mask + + return mask + + +def _lower_projected_attention_tensor( + g: Graph, + hidden: Tensor, + weight: Tensor, + bias: Tensor | None, + target_shape: tuple[int, ...], +) -> Tensor: + if len(target_shape) != 4: + raise NotImplementedError(f"projected attention tensor expects rank-4 target shape, got {target_shape}") + + linear_input = _flatten_to_2d_for_linear(g, hidden) + out = _matmul_with_quantized_rhs_legalization(g, linear_input, weight, pretransposed_rhs=True) + if bias is not None: + out = g.add(out, bias) + return g.reshape(out, target_shape) + + +def _normalize_attention_add_tensor(g: Graph, tensor: Tensor, target_shape: tuple[int, ...]) -> Tensor: + if tuple(tensor.shape) == target_shape: + return tensor + + if len(tensor.shape) != 4 or len(target_shape) != 4: + raise NotImplementedError( + f"unsupported attention add tensor shape {tuple(tensor.shape)} for target {target_shape}" + ) + + batch, seq_len, heads, head_dim = target_shape + shape = tuple(int(v) for v in tensor.shape) + + if shape[0] in (1, batch) and shape[1] == heads and shape[2] in (1, seq_len) and shape[3] == head_dim: + return g.permute(tensor, (0, 2, 1, 3)) + if shape[0] in (1, batch) and shape[1] in (1, seq_len) and shape[2] == heads and shape[3] == head_dim: + return tensor + + raise NotImplementedError( + f"unsupported attention add tensor shape {shape} for target {target_shape}" + ) + + +def _rope_input_is_bhsd(ir: IRGraph, value_id: str) -> bool: + value = ir.values.get(value_id) + if value is None or value.producer is None: + return False + node = ir.nodes.get(value.producer) + if node is None: + return False + if node.op == "permute": + permutation = tuple(int(dim) for dim in node.attrs.get("permutation", ())) + return permutation == (0, 2, 1, 3) + if node.op == "transpose": + return int(node.attrs.get("dim0", -1)) == 1 and int(node.attrs.get("dim1", -1)) == 2 + return False + + +def _lower_softplus(g: Graph, x: Tensor) -> Tensor: + # Stable fp16 softplus: relu(x) + log(1 + exp(-abs(x))). + abs_x = g.abs(x) + neg_abs_x = g.scalar_multiply(abs_x, -1.0) + exp_term = g.scalar_exp(neg_abs_x) + log_term = g.scalar_log(g.scalar_add(exp_term, 1.0)) + return g.add(g.relu(x), log_term) + + +def _normalize_dim(dim: int, rank: int) -> int: + if dim < 0: + dim += rank + return dim + + +def _normalize_reduction_axes(axis: Any, rank: int) -> tuple[int, ...]: + if axis is None: + return tuple(range(rank)) + if isinstance(axis, int): + return (_normalize_dim(axis, rank),) + if isinstance(axis, (list, tuple)): + normalized: list[int] = [] + seen: set[int] = set() + for raw_axis in axis: + if not isinstance(raw_axis, int): + raise NotImplementedError("reduction axes must be integers") + reduced_axis = _normalize_dim(raw_axis, rank) + if reduced_axis in seen: + continue + seen.add(reduced_axis) + normalized.append(reduced_axis) + if normalized: + return tuple(sorted(normalized)) + raise NotImplementedError("reduction axes must be an int, a tuple/list of ints, or None") + + +def _lower_reduction(g: Graph, op: str, x: Tensor, *, axes: tuple[int, ...]) -> Tensor: + if not axes: + return x + + fn = getattr(g, op) + if len(axes) == 1: + return fn(x, axes[0]) + + if op == "variance": + flattened = g.flatten(x, start_dim=axes[0], end_dim=axes[-1]) + collapsed_axes = set(axes) + expected_axes = set(range(axes[0], axes[-1] + 1)) + if collapsed_axes != expected_axes: + raise NotImplementedError("variance currently requires contiguous multi-axis reductions") + return fn(flattened, axes[0]) + + reduced = x + for axis in sorted(axes, reverse=True): + reduced = fn(reduced, axis) + return reduced + + +def _normalize_index(index: int, dim_size: int) -> int: + if index < 0: + index += dim_size + return index + + +def _static_shape(shape: Any) -> tuple[int, ...] | None: + if not isinstance(shape, tuple): + return None + dims: list[int] = [] + for dim in shape: + try: + dims.append(int(dim)) + except Exception: + return None + return tuple(dims) + + +def _ir_value_dtype(ir: IRGraph, value_id: str) -> str | None: + value = ir.values.get(value_id) + if value is None or value.dtype is None: + return None + return str(value.dtype).strip().lower() + + +def _is_integer_index_dtype(dtype: str | None) -> bool: + return dtype in {"uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"} + + +def _constant_tensor_from_ir(ir: IRGraph, value_id: str) -> torch.Tensor | None: + const = ir.constants.get(value_id) + if isinstance(const, torch.nn.Parameter): + const = const.detach() + if isinstance(const, torch.Tensor): + return const.detach().cpu() + return None + + +def _constant_scalar_from_ir(ir: IRGraph, value_id: str) -> float | None: + const = _constant_tensor_from_ir(ir, value_id) + if const is None: + return None + if const.numel() != 1: + raise NotImplementedError( + f"clamp bound {value_id!r} must be a scalar constant, got shape={tuple(const.shape)}" + ) + return float(const.reshape(()).item()) + + +def _try_lower_identity_broadcast_advanced_index( + g: Graph, + node: IRNode, + env: dict[str, Any], + ir: IRGraph, +) -> Tensor | None: + if len(node.inputs) != 3: + return None + + source = _tensor(env, node.inputs[0]) + source_shape = tuple(int(dim) for dim in source.shape) + if len(source_shape) != 2 or int(source_shape[0]) != 1: + return None + + output_shape = _static_shape(ir.values[node.outputs[0]].shape if node.outputs and node.outputs[0] in ir.values else None) + if output_shape is None: + return None + + batch_index = _constant_tensor_from_ir(ir, node.inputs[1]) + token_index = _constant_tensor_from_ir(ir, node.inputs[2]) + if batch_index is None or token_index is None: + index_values = [env.get(value_id) for value_id in node.inputs[1:]] + index_shapes = [ + tuple(int(dim) for dim in value.shape) + for value in index_values + if isinstance(value, Tensor) + ] + source_elems = 1 + for dim in source_shape: + source_elems *= int(dim) + output_elems = 1 + for dim in output_shape: + output_elems *= int(dim) + if ( + int(source_shape[0]) == 1 + and output_elems == source_elems + and index_shapes + and all(shape and all(int(dim) in {1, int(source_shape[1])} for dim in shape) for shape in index_shapes) + ): + return g.reshape(source, output_shape) + return None + + try: + batch_index = torch.broadcast_to(batch_index.to(torch.int64), output_shape) + token_index = torch.broadcast_to(token_index.to(torch.int64), output_shape) + except RuntimeError: + return None + + if torch.count_nonzero(batch_index).item() != 0: + return None + + flat_index = token_index.reshape(-1) + expected = torch.arange(int(source_shape[1]), dtype=torch.int64) + if flat_index.numel() != expected.numel() or not torch.equal(flat_index, expected): + return None + + return g.reshape(source, output_shape) + + +def _try_lower_embedding_advanced_index( + g: Graph, + node: IRNode, + env: dict[str, Any], + ir: IRGraph, +) -> Tensor | None: + if len(node.inputs) != 2: + return None + if not _is_integer_index_dtype(_ir_value_dtype(ir, node.inputs[1])): + return None + + source = _tensor(env, node.inputs[0]) + indices = _tensor(env, node.inputs[1]) + source_shape = tuple(int(dim) for dim in source.shape) + if len(source_shape) != 2: + return None + + output_shape = _static_shape(ir.values[node.outputs[0]].shape if node.outputs and node.outputs[0] in ir.values else None) + gathered = g.embedding_from_tensor(source, indices) + if output_shape is not None and tuple(int(dim) for dim in gathered.shape) != tuple(int(dim) for dim in output_shape): + gathered = g.reshape(gathered, output_shape) + return gathered + + +def _try_lower_prefix_mask_advanced_index( + g: Graph, + node: IRNode, + env: dict[str, Any], + ir: IRGraph, +) -> Tensor | None: + if len(node.inputs) != 2: + return None + if _is_integer_index_dtype(_ir_value_dtype(ir, node.inputs[1])): + return None + + source = _tensor(env, node.inputs[0]) + mask = _tensor(env, node.inputs[1]) + source_shape = tuple(int(dim) for dim in source.shape) + mask_shape = tuple(int(dim) for dim in mask.shape) + if not mask_shape or len(mask_shape) > len(source_shape): + return None + if source_shape[: len(mask_shape)] != mask_shape: + return None + + output_shape = _static_shape(ir.values[node.outputs[0]].shape if node.outputs and node.outputs[0] in ir.values else None) + if output_shape: + prefix_elems = 1 + for dim in mask_shape: + prefix_elems *= int(dim) + trailing_shape = source_shape[len(mask_shape) :] + flattened_shape = (prefix_elems, *trailing_shape) + flattened = source if source_shape == flattened_shape else g.reshape(source, flattened_shape) + expected_prefix = int(output_shape[0]) + if 0 <= expected_prefix <= prefix_elems: + selected = g.slice(flattened, axis=0, start=0, length=expected_prefix) + if tuple(int(dim) for dim in selected.shape) != tuple(int(dim) for dim in output_shape): + selected = g.reshape(selected, output_shape) + return selected + + selected = g.masked_select_prefix(source, mask) + if output_shape is None or not output_shape: + return selected + if len(output_shape) != len(selected.shape): + return selected + expected_prefix = int(output_shape[0]) + if expected_prefix < int(selected.shape[0]): + selected = g.slice(selected, axis=0, start=0, length=expected_prefix) + return selected + + +def _try_lower_advanced_index( + g: Graph, + node: IRNode, + env: dict[str, Any], + ir: IRGraph, +) -> Tensor | None: + lowered = _try_lower_identity_broadcast_advanced_index(g, node, env, ir) + if lowered is not None: + return lowered + lowered = _try_lower_prefix_mask_advanced_index(g, node, env, ir) + if lowered is not None: + return lowered + return _try_lower_embedding_advanced_index(g, node, env, ir) + + +def _normalize_slice_end(end: int, dim_size: int) -> int: + if end < 0: + end += dim_size + return max(0, min(end, dim_size)) + + +def _legalize_elementwise_binary_inputs(g: Graph, lhs: Tensor, rhs: Tensor) -> tuple[Tensor, Tensor]: + if lhs.shape == rhs.shape: + return lhs, rhs + + lhs_rank = len(lhs.shape) + rhs_rank = len(rhs.shape) + + if lhs_rank > rhs_rank: + rhs = _reshape_for_trailing_broadcast(g, rhs, lhs.shape) + elif rhs_rank > lhs_rank: + lhs = _reshape_for_trailing_broadcast(g, lhs, rhs.shape) + + if lhs.shape == rhs.shape: + return lhs, rhs + + return lhs, rhs + + +def _is_scalar_like(value: Any) -> bool: + return isinstance(value, (int, float, bool)) + + +def _lower_binary_op(g: Graph, lhs_value: Any, rhs_value: Any, op: str) -> Tensor: + if isinstance(lhs_value, Tensor) and isinstance(rhs_value, Tensor): + lhs, rhs = _legalize_elementwise_binary_inputs(g, lhs_value, rhs_value) + target_dtype = Graph.FP16 if lhs.dtype == Graph.FP16 and rhs.dtype == Graph.FP16 else Graph.FP32 + lhs = _ensure_tensor_dtype(g, lhs, target_dtype) + rhs = _ensure_tensor_dtype(g, rhs, target_dtype) + if op == "add": + return g.add(lhs, rhs) + if op == "subtract": + return g.subtract(lhs, rhs) + if op == "multiply": + return g.multiply(lhs, rhs) + if op == "divide": + return g.divide(lhs, rhs) + raise NotImplementedError(f"unsupported binary op: {op}") + + if isinstance(lhs_value, Tensor) and _is_scalar_like(rhs_value): + lhs_value = _ensure_scalar_math_tensor(g, lhs_value) + scalar = float(rhs_value) + if op == "add": + return g.scalar_add(lhs_value, scalar) + if op == "subtract": + return g.scalar_subtract(lhs_value, scalar) + if op == "multiply": + return g.scalar_multiply(lhs_value, scalar) + if op == "divide": + return g.scalar_divide(lhs_value, scalar) + raise NotImplementedError(f"unsupported binary op: {op}") + + if _is_scalar_like(lhs_value) and isinstance(rhs_value, Tensor): + rhs_value = _ensure_scalar_math_tensor(g, rhs_value) + scalar = float(lhs_value) + if op == "add": + return g.scalar_add(rhs_value, scalar) + if op == "subtract": + return g.scalar_add(g.scalar_multiply(rhs_value, -1.0), scalar) + if op == "multiply": + return g.scalar_multiply(rhs_value, scalar) + if op == "divide": + raise NotImplementedError("scalar/tensor divide is not directly supported by Cactus graph ops") + raise NotImplementedError(f"unsupported binary op: {op}") + + raise TypeError( + f"unsupported lowered operand types for {op}: " + f"{type(lhs_value).__name__}, {type(rhs_value).__name__}" + ) + + +def _lower_compare_op(g: Graph, lhs_value: Any, rhs_value: Any, op: str) -> Tensor: + if op == "equal": + not_equal = _lower_compare_op(g, lhs_value, rhs_value, "not_equal") + return g.scalar_not_equal(not_equal, 1.0) + + if op == "greater": + delta = _lower_binary_op(g, lhs_value, rhs_value, "subtract") + delta = _ensure_fp16_tensor(g, delta) + return g.scalar_not_equal(g.relu(delta), 0.0) + + if op == "less": + return _lower_compare_op(g, rhs_value, lhs_value, "greater") + + if op == "greater_equal": + less = _lower_compare_op(g, lhs_value, rhs_value, "less") + return g.scalar_not_equal(less, 1.0) + + if op == "less_equal": + greater = _lower_compare_op(g, lhs_value, rhs_value, "greater") + return g.scalar_not_equal(greater, 1.0) + + if op != "not_equal": + raise NotImplementedError(f"unsupported compare op: {op}") + + if isinstance(lhs_value, Tensor) and isinstance(rhs_value, Tensor): + lhs, rhs = _legalize_elementwise_binary_inputs(g, lhs_value, rhs_value) + target_dtype = Graph.FP16 if lhs.dtype == Graph.FP16 and rhs.dtype == Graph.FP16 else Graph.FP32 + lhs = _ensure_tensor_dtype(g, lhs, target_dtype) + rhs = _ensure_tensor_dtype(g, rhs, target_dtype) + return g.not_equal(lhs, rhs) + + if isinstance(lhs_value, Tensor) and _is_scalar_like(rhs_value): + lhs_value = _ensure_tensor_dtype( + g, + lhs_value, + Graph.FP16 if lhs_value.dtype == Graph.FP16 else Graph.FP32, + ) + return g.scalar_not_equal(lhs_value, float(rhs_value)) + + if _is_scalar_like(lhs_value) and isinstance(rhs_value, Tensor): + rhs_value = _ensure_tensor_dtype( + g, + rhs_value, + Graph.FP16 if rhs_value.dtype == Graph.FP16 else Graph.FP32, + ) + return g.scalar_not_equal(rhs_value, float(lhs_value)) + + if _is_scalar_like(lhs_value) and _is_scalar_like(rhs_value): + result = 1.0 if float(lhs_value) != float(rhs_value) else 0.0 + return _materialize_constant_tensor(g, torch.tensor([result], dtype=torch.float32)) + + raise TypeError( + f"unsupported lowered operand types for {op}: " + f"{type(lhs_value).__name__}, {type(rhs_value).__name__}" + ) + + +def _reshape_for_trailing_broadcast(g: Graph, tensor: Tensor, target_shape: tuple[int, ...]) -> Tensor: + tensor_shape = tuple(tensor.shape) + target_rank = len(target_shape) + tensor_rank = len(tensor_shape) + + if tensor_rank > target_rank: + return tensor + + padded_shape = (1,) * (target_rank - tensor_rank) + tensor_shape + + # Only legalize cases that are valid trailing broadcasts, e.g.: + # (H,) -> (1, H), (1, 1, H), etc. + for src_dim, tgt_dim in zip(padded_shape, target_shape): + if src_dim != 1 and src_dim != tgt_dim: + return tensor + + if padded_shape == tensor_shape: + return tensor + + return g.reshape(tensor, padded_shape) + + +def _ensure_fp16_tensor(g: Graph, tensor: Tensor) -> Tensor: + if tensor.dtype == Graph.FP16: + return tensor + return g.precision_cast(tensor, Graph.FP16) + + +def _ensure_scalar_math_tensor(g: Graph, tensor: Tensor) -> Tensor: + return tensor + + +def _ensure_tensor_dtype(g: Graph, tensor: Tensor, dtype: int) -> Tensor: + if tensor.dtype == dtype: + return tensor + return g.precision_cast(tensor, dtype) + + +def _cat_with_legalized_dtype(g: Graph, tensors: list[Tensor], *, axis: int, dtype: int | None) -> Tensor: + if not tensors: + raise NotImplementedError("cat requires at least one tensor") + if dtype is None: + dtypes = {tensor.dtype for tensor in tensors} + dtype = tensors[0].dtype if len(dtypes) == 1 else Graph.FP16 + legalized = [_ensure_tensor_dtype(g, tensor, dtype) for tensor in tensors] + return g.cat(legalized, axis=axis) + + +def _invert_permutation(permutation: tuple[int, ...]) -> tuple[int, ...]: + inverse = [0] * len(permutation) + for index, source_axis in enumerate(permutation): + inverse[int(source_axis)] = index + return tuple(inverse) + + +def _lower_static_strided_slice_via_gather( + g: Graph, + x: Tensor, + *, + axis: int, + indices: list[int], +) -> Tensor: + if not indices: + return g.slice(x, axis=axis, start=0, length=0) + + indices_tensor = _materialize_constant_tensor(g, torch.tensor(indices, dtype=torch.float32)) + if axis == 0: + return g.gather(x, indices_tensor) + + permutation = (axis, *[dim for dim in range(len(x.shape)) if dim != axis]) + transposed = g.permute(x, permutation) + gathered = g.gather(transposed, indices_tensor) + inverse_permutation = _invert_permutation(permutation) + return g.permute(gathered, inverse_permutation) + + +def _lower_where_branch_term(g: Graph, branch_value: Any, mask: Tensor, *, dtype: int | None) -> Tensor: + if dtype is not None: + mask = _ensure_tensor_dtype(g, mask, dtype) + if isinstance(branch_value, Tensor): + branch_mask = mask if dtype is None else _ensure_tensor_dtype(g, mask, branch_value.dtype) + return _lower_binary_op(g, branch_value, branch_mask, "multiply") + return g.scalar_multiply(mask, _clamp_scalar_for_dtype(float(branch_value), mask.dtype)) + + +def _clamp_scalar_for_dtype(value: float, dtype: int | None) -> float: + if dtype != Graph.FP16: + return value + if math.isnan(value): + return value + if value == math.inf: + return 65504.0 + if value == -math.inf: + return -65504.0 + return max(-65504.0, min(65504.0, value)) + + +def _lower_where_op(g: Graph, node: IRNode, env: dict[str, Any]) -> Tensor: + condition = _lower_compare_op(g, env[node.inputs[0]], 0.0, "not_equal") + false_mask = g.scalar_add(g.scalar_multiply(condition, -1.0), 1.0) + + input_index = 1 + if bool(node.attrs.get("true_is_scalar", False)): + true_value: Any = float(node.attrs["true_value"]) + else: + true_value = env[node.inputs[input_index]] + input_index += 1 + + if bool(node.attrs.get("false_is_scalar", False)): + false_value: Any = float(node.attrs["false_value"]) + else: + false_value = env[node.inputs[input_index]] + + result_dtype: int | None = None + if isinstance(true_value, Tensor): + result_dtype = true_value.dtype + elif isinstance(false_value, Tensor): + result_dtype = false_value.dtype + + true_term = _lower_where_branch_term(g, true_value, condition, dtype=result_dtype) + false_term = _lower_where_branch_term(g, false_value, false_mask, dtype=result_dtype) + return _lower_binary_op(g, true_term, false_term, "add") + + +def _flatten_to_2d_for_linear(g: Graph, tensor: Tensor) -> Tensor: + shape = tuple(tensor.shape) + if len(shape) <= 2: + return tensor + leading = 1 + for dim in shape[:-1]: + leading *= int(dim) + return g.view(tensor, (leading, int(shape[-1]))) + + +def _broadcast_batch_shape(lhs_batch: tuple[int, ...], rhs_batch: tuple[int, ...]) -> tuple[int, ...] | None: + rank = max(len(lhs_batch), len(rhs_batch)) + lhs_padded = (1,) * (rank - len(lhs_batch)) + lhs_batch + rhs_padded = (1,) * (rank - len(rhs_batch)) + rhs_batch + result: list[int] = [] + for lhs_dim, rhs_dim in zip(lhs_padded, rhs_padded): + if lhs_dim == rhs_dim: + result.append(int(lhs_dim)) + elif lhs_dim == 1: + result.append(int(rhs_dim)) + elif rhs_dim == 1: + result.append(int(lhs_dim)) + else: + return None + return tuple(result) + + +def _project_broadcast_coord(coord: tuple[int, ...], operand_batch: tuple[int, ...]) -> tuple[int, ...]: + if not operand_batch: + return () + offset = len(coord) - len(operand_batch) + projected: list[int] = [] + for axis, dim in enumerate(operand_batch): + value = coord[offset + axis] + projected.append(0 if int(dim) == 1 else int(value)) + return tuple(projected) + + +def _index_batch_prefix(g: Graph, tensor: Tensor, coord: tuple[int, ...]) -> Tensor: + for index_value in coord: + tensor = g.index(tensor, index_value, axis=0) + return tensor + + +def _lower_static_batched_matmul( + g: Graph, + lhs: Tensor, + rhs: Tensor, + *, + output_dtype: int | None = None, +) -> Tensor | None: + lhs_shape = tuple(int(dim) for dim in lhs.shape) + rhs_shape = tuple(int(dim) for dim in rhs.shape) + if len(lhs_shape) < 3 or len(rhs_shape) < 3: + return None + if lhs_shape[-1] != rhs_shape[-2]: + return None + + lhs_batch = lhs_shape[:-2] + rhs_batch = rhs_shape[:-2] + batch_shape = _broadcast_batch_shape(lhs_batch, rhs_batch) + if batch_shape is None: + return None + + output_matrix_shape = (int(lhs_shape[-2]), int(rhs_shape[-1])) + rank = len(batch_shape) + + def _build(dim: int, coord_prefix: tuple[int, ...]) -> Tensor: + if dim == rank: + lhs_coord = _project_broadcast_coord(coord_prefix, lhs_batch) + rhs_coord = _project_broadcast_coord(coord_prefix, rhs_batch) + lhs_slice = _index_batch_prefix(g, lhs, lhs_coord) + rhs_slice = _index_batch_prefix(g, rhs, rhs_coord) + out = _matmul_with_quantized_rhs_legalization( + g, + lhs_slice, + rhs_slice, + output_dtype=output_dtype, + ) + return g.reshape(out, (1,) * rank + output_matrix_shape) + + pieces = [_build(dim + 1, coord_prefix + (index_value,)) for index_value in range(int(batch_shape[dim]))] + if len(pieces) == 1: + return pieces[0] + return g.cat(pieces, axis=dim) + + return _build(0, ()) + + +def _should_lower_gemma4_decoder_attention_without_kernel(ir: IRGraph, node: IRNode) -> bool: + if os.environ.get("CACTUS_GEMMA4_DECODER_MANUAL_ATTENTION", "0") != "1": + return False + family = str(ir.meta.get("adapter_family") or ir.meta.get("family") or "").strip().lower() + component = str(ir.meta.get("component", "") or "").strip().lower() + if family != "gemma4" or component != "decoder": + return False + return node.op in {"attention", "scaled_dot_product_attention"} + + +def _lower_gemma4_decoder_attention_without_kernel( + g: Graph, + ir: IRGraph, + env: dict[str, Any], + node: IRNode, +) -> Tensor: + query = _attention_tensor(env, node.inputs[0]) + key = _attention_tensor(env, node.inputs[1]) + value = _attention_tensor(env, node.inputs[2]) + + key_transposed = g.permute(key, (0, 1, 3, 2)) + scores = _lower_static_batched_matmul(g, query, key_transposed, output_dtype=Graph.FP16) + if scores is None: + raise NotImplementedError("Gemma4 decoder attention requires static batched matmul support") + + scale = float(node.attrs.get("scale", 1.0)) + if scale != 1.0: + scores = g.scalar_multiply(_ensure_scalar_math_tensor(g, scores), scale) + + additive_mask: Tensor | None = None + if len(node.inputs) > 3: + mask_tensor = _ensure_fp16_tensor(g, _tensor(env, node.inputs[3])) + if bool(node.attrs.get("additive_mask", False)): + additive_mask = mask_tensor + else: + additive_mask = _lower_binary_op(g, mask_tensor, 1.0, "subtract") + additive_mask = g.scalar_multiply(_ensure_scalar_math_tensor(g, additive_mask), 1.0e4) + else: + query_shape = tuple(int(dim) for dim in query.shape) + key_shape = tuple(int(dim) for dim in key.shape) + if len(query_shape) >= 4 and len(key_shape) >= 4: + query_seq = int(query_shape[-2]) + key_seq = int(key_shape[-2]) + additive_mask = _materialize_gemma4_attention_mask( + g, + query_seq=query_seq, + key_seq=key_seq, + is_causal=bool(node.attrs.get("is_causal", False)), + window_size=int(node.attrs.get("window_size", 0)), + dtype=torch.float16, + ) + + if additive_mask is not None: + scores = _lower_binary_op(g, scores, additive_mask, "add") + + probs = g.softmax(scores, axis=-1) + probs = _ensure_tensor_dtype(g, probs, value.dtype) + out = _lower_static_batched_matmul(g, probs, value, output_dtype=value.dtype) + if out is None: + raise NotImplementedError("Gemma4 decoder attention output requires static batched matmul support") + + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return out + + +def _materialize_gemma4_attention_mask( + g: Graph, + *, + query_seq: int, + key_seq: int, + is_causal: bool, + window_size: int, + dtype: torch.dtype, +) -> Tensor | None: + if query_seq <= 0 or key_seq <= 0: + return None + if not is_causal and window_size <= 0: + return None + + q_index = np.arange(query_seq, dtype=np.int32)[:, None] + k_index = np.arange(key_seq, dtype=np.int32)[None, :] + allowed = np.ones((query_seq, key_seq), dtype=np.bool_) + + if is_causal: + allowed &= k_index <= q_index + if window_size > 0 and window_size < key_seq: + allowed &= k_index >= (q_index - (window_size - 1)) + + mask = np.where(allowed, 0.0, -1.0e4).astype(np.float16) + tensor = torch.from_numpy(mask.reshape(1, 1, query_seq, key_seq)).to(dtype=dtype) + return _materialize_constant_tensor(g, tensor) + + +def _legalize_matmul_inputs( + g: Graph, + lhs: Tensor, + rhs: Tensor, + node: IRNode, +) -> tuple[Tensor, Tensor, tuple[int, ...]] | None: + lhs_shape = tuple(lhs.shape) + rhs_shape = tuple(rhs.shape) + if len(lhs_shape) <= 2 and len(rhs_shape) <= 2: + return None + + if len(lhs_shape) > 2 and len(rhs_shape) == 2 and lhs_shape[-1] == rhs_shape[0]: + leading = 1 + for dim in lhs_shape[:-1]: + leading *= int(dim) + lhs_2d = g.reshape(lhs, (leading, int(lhs_shape[-1]))) + output_shape = lhs_shape[:-1] + (int(rhs_shape[1]),) + return lhs_2d, rhs, output_shape + + # Cactus matmul is 2D-only. Legalize the narrow case where both operands have + # only singleton leading dims, e.g. rotary helper matmuls: + # (1, M, K) @ (1, K, N) -> reshape to (M, K) @ (K, N) -> reshape back. + if any(dim != 1 for dim in lhs_shape[:-2]): + return None + if any(dim != 1 for dim in rhs_shape[:-2]): + return None + + lhs_2d_shape = lhs_shape[-2:] + rhs_2d_shape = rhs_shape[-2:] + if lhs_2d_shape[-1] != rhs_2d_shape[0]: + return None + + output_shape = node.meta.get("shape") + if not isinstance(output_shape, tuple): + output_shape = lhs_shape[:-2] + (lhs_2d_shape[0], rhs_2d_shape[1]) + output_shape = tuple(int(v) for v in output_shape) + + lhs_2d = g.reshape(lhs, lhs_2d_shape) + rhs_2d = g.reshape(rhs, rhs_2d_shape) + return lhs_2d, rhs_2d, output_shape + + +def _resolve_expand_shape(input_shape: tuple[int, ...], requested_shape: tuple[int, ...]) -> tuple[int, ...]: + if len(requested_shape) < len(input_shape): + raise NotImplementedError(f"expand cannot reduce rank: {input_shape} -> {requested_shape}") + + padded_input = (1,) * (len(requested_shape) - len(input_shape)) + input_shape + resolved: list[int] = [] + for in_dim, req_dim in zip(padded_input, requested_shape): + if req_dim == -1: + resolved.append(in_dim) + continue + if req_dim < -1: + raise NotImplementedError(f"invalid expand dimension: {req_dim}") + resolved.append(int(req_dim)) + return tuple(resolved) + + +def _lower_repeat(g: Graph, x: Tensor, repeats: tuple[int, ...]) -> Tensor: + if not repeats: + return x + + input_shape = tuple(int(dim) for dim in x.shape) + if len(repeats) < len(input_shape): + raise NotImplementedError(f"repeat cannot reduce rank: {input_shape} with repeats={repeats}") + + if len(repeats) > len(input_shape): + padded_shape = (1,) * (len(repeats) - len(input_shape)) + input_shape + x = g.reshape(x, padded_shape) + + result = x + for axis, factor in enumerate(repeats): + if factor < 0: + raise NotImplementedError(f"repeat does not support negative factors: {repeats}") + if factor == 0: + raise NotImplementedError(f"repeat does not support zero factors: {repeats}") + if factor == 1: + continue + result = g.cat([result] * factor, axis=axis) + return result + + +def _resolve_reshape_shape(input_shape: tuple[int, ...], requested_shape: tuple[int, ...]) -> tuple[int, ...]: + resolved = [int(v) for v in requested_shape] + unknown_indices = [idx for idx, dim in enumerate(resolved) if dim == -1] + if len(unknown_indices) > 1: + raise NotImplementedError(f"reshape with multiple inferred dimensions is unsupported: {requested_shape}") + if not unknown_indices: + return tuple(resolved) + + input_elements = 1 + for dim in input_shape: + input_elements *= int(dim) + + known_elements = 1 + for dim in resolved: + if dim != -1: + known_elements *= int(dim) + + if known_elements == 0 or input_elements % known_elements != 0: + raise NotImplementedError(f"cannot infer reshape target {requested_shape} from input shape {input_shape}") + + resolved[unknown_indices[0]] = input_elements // known_elements + return tuple(resolved) + + +def _match_gqa_expand_alias(ir: IRGraph, expand_node: IRNode) -> str | None: + target_shape = tuple(int(v) for v in expand_node.attrs["shape"]) + if len(target_shape) != 5: + return None + + current_value_id = expand_node.inputs[0] + base_value_id: str | None = None + + while True: + producer_id = ir.values[current_value_id].producer + if producer_id is None: + return None + producer = ir.nodes[producer_id] + if producer.op == "slice": + axis = int(producer.attrs.get("axis", -1)) + input_shape = ir.values[producer.inputs[0]].shape + if input_shape is None: + return None + normalized_axis = _normalize_dim(axis, len(input_shape)) + start = int(producer.attrs.get("start", 0)) + end = int(producer.attrs.get("end", input_shape[normalized_axis])) + if start != 0 or end < input_shape[normalized_axis]: + return None + current_value_id = producer.inputs[0] + continue + if producer.op == "unsqueeze": + base_value_id = producer.inputs[0] + base_shape = ir.values[base_value_id].shape + if base_shape is None: + return None + unsqueezed_dim = _normalize_dim(int(producer.attrs.get("dim", 0)), len(base_shape) + 1) + if unsqueezed_dim != 2: + return None + break + if producer.op in {"reshape", "view"}: + base_value_id = producer.inputs[0] + base_shape = ir.values[base_value_id].shape + view_shape = ir.values[current_value_id].shape + if base_shape is None or view_shape is None: + return None + if len(base_shape) != 4 or len(view_shape) != 5: + return None + if tuple(int(v) for v in view_shape) != ( + int(base_shape[0]), + int(base_shape[1]), + 1, + int(base_shape[2]), + int(base_shape[3]), + ): + return None + break + return None + + if base_value_id is None: + return None + base_shape = ir.values[base_value_id].shape + if base_shape is None or len(base_shape) != 4: + return None + + expected_target = ( + int(base_shape[0]), + int(base_shape[1]), + int(target_shape[2]), + int(base_shape[2]), + int(base_shape[3]), + ) + if target_shape != expected_target: + return None + if int(target_shape[2]) <= 1: + return None + + return base_value_id + + +def _map_ir_dtype(dtype: str) -> int: + if dtype == "bf16": + return Graph.FP32 + if dtype == "fp16": + return Graph.FP16 + if dtype in ("fp32", "fp64"): + return Graph.FP32 + if dtype == "int8": + return Graph.INT8 + if dtype in ("int32", "int64", "bool"): + return Graph.FP32 + + raise NotImplementedError(f"unsupported IR dtype: {dtype}") + + +def _map_ir_or_torch_dtype(dtype: Any) -> int: + if dtype is None: + raise NotImplementedError("missing dtype for precision_cast") + if isinstance(dtype, str): + if dtype.startswith("torch."): + return _map_torch_dtype(getattr(torch, dtype.split(".", 1)[1])) + return _map_ir_dtype(dtype) + return _map_torch_dtype(dtype) + + +def _map_torch_dtype(dtype: Any) -> int: + if dtype == torch.bfloat16: + return Graph.FP32 + if dtype == torch.float16: + return Graph.FP16 + if dtype == torch.float32: + return Graph.FP32 + if dtype == torch.float64: + return Graph.FP32 + if dtype == torch.int8: + return Graph.INT8 + if dtype in (torch.int16, torch.int32, torch.int64, torch.bool): + return Graph.FP32 + raise NotImplementedError(f"unsupported torch dtype: {dtype}") + + +def _materialize_constant_tensor(g: Graph, tensor_value: torch.Tensor) -> Tensor: + graph_dtype = _map_torch_dtype(tensor_value.dtype) + materialized = tensor_value.detach().cpu() + if graph_dtype == Graph.FP32 and materialized.dtype not in (torch.float32, torch.float64): + materialized = materialized.to(torch.float32) + elif graph_dtype == Graph.FP16 and materialized.dtype != torch.float16: + materialized = materialized.to(torch.float16) + elif graph_dtype == Graph.INT8 and materialized.dtype != torch.int8: + materialized = materialized.to(torch.int8) + + tensor = g.input(shape=tuple(materialized.shape), dtype=graph_dtype) + g.set_input(tensor, materialized, dtype=graph_dtype) + materialized_constants = getattr(g, "_transpile_materialized_constants", None) + if isinstance(materialized_constants, list): + materialized_constants.append(tensor) + return tensor + + +def _materialize_constant_torch_dtype(dtype: Any) -> torch.dtype: + if dtype is None: + return torch.float32 + if isinstance(dtype, str): + if dtype.startswith("torch."): + return getattr(torch, dtype.split(".", 1)[1]) + if dtype in ("fp16", "bf16"): + return torch.float16 + if dtype in ("fp32", "fp64", "int16", "int32", "int64", "bool"): + return torch.float32 + if dtype == "int8": + return torch.int8 + if isinstance(dtype, torch.dtype): + return dtype + raise NotImplementedError(f"unsupported dtype for constant materialization: {dtype}") + + +def _torch_dtype_for_graph_dtype(dtype: int) -> torch.dtype: + if dtype == Graph.FP16: + return torch.float16 + if dtype == Graph.FP32: + return torch.float32 + if dtype == Graph.INT8: + return torch.int8 + raise NotImplementedError(f"unsupported graph dtype for constant materialization: {dtype}") diff --git a/python/cactus/transpile/media_limits.py b/python/cactus/transpile/media_limits.py new file mode 100644 index 000000000..8b7910cf9 --- /dev/null +++ b/python/cactus/transpile/media_limits.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import os + +_DEFAULT_IMAGE_SIZE = 256 + + +def static_image_size() -> int: + raw = os.environ.get("CACTUS_TRANSPILER_IMAGE_SIZE", str(_DEFAULT_IMAGE_SIZE)) + try: + return max(1, int(raw)) + except (TypeError, ValueError): + return _DEFAULT_IMAGE_SIZE + + +def resize_static_image(image: object) -> object: + size = static_image_size() + target = (size, size) + if getattr(image, "size", None) == target or not hasattr(image, "resize"): + return image + try: + from PIL import Image # type: ignore + from PIL import ImageOps # type: ignore + + resample = Image.Resampling.BILINEAR + except AttributeError: # pragma: no cover + resample = Image.BILINEAR # type: ignore[name-defined] + from PIL import ImageOps # type: ignore + except Exception: + return image.resize(target) + + try: + return ImageOps.pad(image, target, method=resample, color=(255, 255, 255)) + except Exception: + return image.resize(target, resample=resample) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py new file mode 100644 index 000000000..ba9e3bf04 --- /dev/null +++ b/python/cactus/transpile/model_adapters.py @@ -0,0 +1,7292 @@ +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +import inspect +import math +import struct +import os +from pathlib import Path +from typing import Callable + +import numpy as np +import torch +import torch.nn.functional as F + +from cactus.transpile.component_pipeline import ComponentModuleSpec +from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC +from cactus.convert.cactus_adapters.tensor_io import align_offset + + +_GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE = 1.0 / 64.0 +_UNSET = object() +try: + from transformers.models.gemma4.modeling_gemma4 import create_bidirectional_mask as _GEMMA4_CREATE_BIDIRECTIONAL_MASK # type: ignore +except Exception: + try: + from transformers.masking_utils import create_bidirectional_mask as _GEMMA4_CREATE_BIDIRECTIONAL_MASK # type: ignore + except Exception: + _GEMMA4_CREATE_BIDIRECTIONAL_MASK = None + + +class UnsupportedComponentsError(RuntimeError): + pass + + +@dataclass +class CanonicalizedModel: + module: torch.nn.Module + task: str + family: str + input_names: tuple[str, ...] = () + + +@dataclass(frozen=True) +class _Gemma4NativeMergeSegment: + kind: str + input_start: int + length: int + feature_start: int = 0 + + +@dataclass(frozen=True) +class _Gemma4NativeMergePlan: + segments: tuple[_Gemma4NativeMergeSegment, ...] + pli_token_ids: tuple[int, ...] + + +def _model_name_or_path(model: torch.nn.Module) -> str: + value = getattr(model, "name_or_path", None) + if isinstance(value, str) and value: + return value + config = getattr(model, "config", None) + value = getattr(config, "_name_or_path", None) + if isinstance(value, str) and value: + return value + return "" + + +def _transpile_graph_meta(model: torch.nn.Module, *, adapter_family: str, adapter_type: str, input_names: tuple[str, ...]) -> dict[str, object]: + meta: dict[str, object] = { + "adapter_family": adapter_family, + "adapter_type": adapter_type, + "model_name_or_path": _model_name_or_path(model), + "input_names": input_names, + } + # Rope-table precompute falls back to this on components that don't reserve a KV cache. + for candidate in _config_candidates_for_context_length(model): + for key in ("max_position_embeddings", "context_length", "model_max_length", "n_positions"): + value = _config_int_value(candidate, key) + if value is not None: + meta["max_position_embeddings"] = value + return meta + return meta + + +def _extract_tensor_output(output: object, *, preferred_field: str | None = None) -> torch.Tensor: + if preferred_field is not None: + value = getattr(output, preferred_field, None) + if isinstance(value, torch.Tensor): + return value + + if isinstance(output, torch.Tensor): + return output + if isinstance(output, tuple) and output and isinstance(output[0], torch.Tensor): + return output[0] + + for field_name in ("last_hidden_state", "logits"): + value = getattr(output, field_name, None) + if isinstance(value, torch.Tensor): + return value + + raise TypeError(f"could not extract tensor output from {type(output).__name__}") + + +def _gemma4_get_placeholder_masks( + get_placeholder_mask: Callable[..., object], + *, + token_type_ids: torch.Tensor | None, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + kwargs: dict[str, torch.Tensor] = { + "input_ids": input_ids, + "inputs_embeds": inputs_embeds, + } + try: + signature = inspect.signature(get_placeholder_mask) + except (TypeError, ValueError): + signature = None + if token_type_ids is not None and signature is not None: + parameters = signature.parameters + accepts_kwargs = any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in parameters.values() + ) + if accepts_kwargs or "token_type_ids" in parameters: + kwargs["token_type_ids"] = token_type_ids + + accepts_token_type_ids = "token_type_ids" in kwargs + result = get_placeholder_mask(**kwargs) + if not isinstance(result, tuple) or len(result) != 3: + raise TypeError( + "Gemma4 get_placeholder_mask must return " + "three placeholder masks" + ) + first_mask, second_mask, third_mask = result + if not ( + isinstance(first_mask, torch.Tensor) + and isinstance(second_mask, torch.Tensor) + and isinstance(third_mask, torch.Tensor) + ): + raise TypeError("Gemma4 get_placeholder_mask returned non-tensor masks") + if accepts_token_type_ids: + text_mask, image_mask, audio_mask = first_mask, second_mask, third_mask + else: + image_mask, video_mask, audio_mask = first_mask, second_mask, third_mask + text_mask = ~(image_mask | video_mask | audio_mask) + return text_mask, image_mask, audio_mask + + +def _module_or_config_attr(module: object, name: str, default: object | None = None) -> object: + value = getattr(module, name, None) + if value is not None: + return value + config = getattr(module, "config", None) + value = getattr(config, name, None) + if value is not None: + return value + return default + + +def _gemma4_get_per_layer_inputs( + backbone: torch.nn.Module, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor, +) -> torch.Tensor: + get_per_layer_inputs = getattr(backbone, "get_per_layer_inputs") + try: + signature = inspect.signature(get_per_layer_inputs) + except (TypeError, ValueError): + signature = None + if signature is not None and "inputs_embeds" not in signature.parameters: + return get_per_layer_inputs(input_ids) + return get_per_layer_inputs(input_ids, inputs_embeds) + + +def _select_last_active_token(hidden_or_logits: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor: + if attention_mask is None or hidden_or_logits.ndim < 3: + return hidden_or_logits[:, -1:, ...] + + if attention_mask.ndim != 2: + raise ValueError(f"expected 2D attention mask, got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[1] != hidden_or_logits.shape[1]: + raise ValueError( + "attention mask / hidden sequence length mismatch: " + f"{tuple(attention_mask.shape)} vs {tuple(hidden_or_logits.shape)}" + ) + + trailing_zero = attention_mask[:, :1] - attention_mask[:, :1] + shifted_mask = torch.cat((attention_mask[:, 1:], trailing_zero), dim=1) + last_active_mask = torch.logical_and(attention_mask != 0, shifted_mask == 0) + expanded_mask = last_active_mask.to(dtype=hidden_or_logits.dtype) + for _ in range(hidden_or_logits.ndim - 2): + expanded_mask = expanded_mask.unsqueeze(-1) + return (hidden_or_logits * expanded_mask).sum(dim=1, keepdim=True) + + +def _select_last_non_pad_token( + hidden_or_logits: torch.Tensor, + input_ids: torch.Tensor | None, + *, + pad_token_id: int | None, +) -> torch.Tensor: + if hidden_or_logits.ndim < 3: + return hidden_or_logits + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: + return hidden_or_logits[:, -1:, ...] + if input_ids.shape[1] != hidden_or_logits.shape[1]: + raise ValueError( + "input id / hidden sequence length mismatch: " + f"{tuple(input_ids.shape)} vs {tuple(hidden_or_logits.shape)}" + ) + if pad_token_id is None: + return hidden_or_logits[:, -1:, ...] + attention_mask = (input_ids != int(pad_token_id)).to(dtype=torch.int64) + return _select_last_active_token(hidden_or_logits, attention_mask) + + +def _tile_to_length(tensor: torch.Tensor, length: int) -> torch.Tensor: + reps = -(-length // int(tensor.shape[1])) + return torch.cat([tensor] * reps, dim=1)[:, :length].contiguous() + + +def _resolve_model_pad_token_id(model: torch.nn.Module) -> int | None: + def _coerce(value: object) -> int | None: + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + config = getattr(model, "config", None) + for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): + resolved = _coerce(getattr(config, attr_name, None)) + if resolved is not None: + return resolved + generation_config = getattr(model, "generation_config", None) + for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): + resolved = _coerce(getattr(generation_config, attr_name, None)) + if resolved is not None: + return resolved + return None + + +def _filter_supported_kwargs(module: torch.nn.Module, kwargs: dict[str, object]) -> dict[str, object]: + try: + signature = inspect.signature(module.forward) + except (TypeError, ValueError): + return kwargs + accepted = signature.parameters + supports_var_kwargs = any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in accepted.values() + ) + if supports_var_kwargs: + return kwargs + return {name: value for name, value in kwargs.items() if name in accepted} + + +def _module_floating_dtype(module: torch.nn.Module) -> torch.dtype | None: + for parameter in module.parameters(): + if parameter.is_floating_point(): + return parameter.dtype + for buffer in module.buffers(): + if buffer.is_floating_point(): + return buffer.dtype + return None + + +def _module_device(module: torch.nn.Module) -> torch.device | None: + for parameter in module.parameters(): + return parameter.device + for buffer in module.buffers(): + return buffer.device + return None + + +def _module_has_meta_tensors(module: torch.nn.Module) -> bool: + return any(parameter.is_meta for parameter in module.parameters()) or any( + buffer.is_meta for buffer in module.buffers() + ) + + +def _to_meta_example_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None: + if tensor is None or tensor.is_meta: + return tensor + return torch.empty_strided( + tuple(tensor.shape), + tuple(tensor.stride()), + dtype=tensor.dtype, + device="meta", + ) + + +def _unique_modules(*candidates: object) -> tuple[torch.nn.Module, ...]: + modules: list[torch.nn.Module] = [] + seen: set[int] = set() + for candidate in candidates: + if not isinstance(candidate, torch.nn.Module): + continue + candidate_id = id(candidate) + if candidate_id in seen: + continue + seen.add(candidate_id) + modules.append(candidate) + return tuple(modules) + + +def _gemma4_vision_modules(multimodal_backbone: torch.nn.Module) -> tuple[torch.nn.Module, ...]: + return _unique_modules( + getattr(multimodal_backbone, "vision_tower", None), + getattr(multimodal_backbone, "embed_vision", None), + ) + + +def _gemma4_audio_modules(multimodal_backbone: torch.nn.Module) -> tuple[torch.nn.Module, ...]: + return _unique_modules( + getattr(multimodal_backbone, "audio_tower", None), + getattr(multimodal_backbone, "embed_audio", None), + ) + + +def _gemma4_text_config(multimodal_backbone: torch.nn.Module) -> object: + config = getattr(multimodal_backbone, "config", None) + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + return get_text_config() + except Exception: + pass + return config + + +def _parse_cache_context_length(value: str | int | None) -> int | None: + if value is None: + return None + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"", "auto"}: + return None + parsed = int(normalized) + else: + parsed = int(value) + if parsed <= 0: + raise ValueError("--cache-context-length must be a positive integer or auto") + return parsed + + +def _config_int_value(config: object, key: str) -> int | None: + if config is None: + return None + if isinstance(config, dict): + raw_value = config.get(key) + else: + raw_value = getattr(config, key, None) + if raw_value is None: + return None + parsed = int(raw_value) + return parsed if parsed > 0 else None + + +def _config_candidates_for_context_length(model: torch.nn.Module) -> tuple[object, ...]: + candidates: list[object] = [] + seen: set[int] = set() + + def _add(candidate: object) -> None: + if candidate is None: + return + candidate_id = id(candidate) + if candidate_id in seen: + return + seen.add(candidate_id) + candidates.append(candidate) + + for owner in ( + model, + getattr(model, "model", None), + getattr(getattr(model, "model", None), "language_model", None), + ): + config = getattr(owner, "config", None) + _add(config) + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + _add(get_text_config()) + for attr in ("text_config", "llm_config", "language_config", "decoder_config"): + _add(getattr(config, attr, None)) + + return tuple(candidates) + + +def _cache_context_length( + model: torch.nn.Module, + *, + input_seq_len: int, + cache_context_length: str | int | None, + fallback_extra_tokens: int, +) -> int: + explicit_length = _parse_cache_context_length(cache_context_length) + if explicit_length is not None: + return explicit_length + + for candidate in _config_candidates_for_context_length(model): + for key in ("max_position_embeddings", "context_length", "model_max_length", "n_positions"): + value = _config_int_value(candidate, key) + if value is not None: + return value + return input_seq_len + int(fallback_extra_tokens) + + +def _max_cache_seq_len( + model: torch.nn.Module, + input_ids: torch.Tensor, + cache_context_length: str | int | None, + *, + fallback_extra_tokens: int, +) -> int: + return max(1024, _cache_context_length( + model, + input_seq_len=int(input_ids.shape[1]), + cache_context_length=cache_context_length, + fallback_extra_tokens=fallback_extra_tokens, + )) + + +def _gemma4_special_token_ids(multimodal_backbone: torch.nn.Module) -> tuple[int, int, int]: + config = getattr(multimodal_backbone, "config", None) + text_config = _gemma4_text_config(multimodal_backbone) + image_token_id = int(getattr(config, "image_token_id", 0) or 0) + audio_token_id = int(getattr(config, "audio_token_id", 0) or 0) + pad_token_id = getattr(text_config, "pad_token_id", None) + if pad_token_id is None: + pad_token_id = getattr(config, "pad_token_id", 0) + return image_token_id, audio_token_id, int(pad_token_id or 0) + + +@contextmanager +def _temporary_cpu_float32_modules(modules: tuple[torch.nn.Module, ...]): + promoted: list[tuple[torch.nn.Module, torch.dtype]] = [] + for module in modules: + device = _module_device(module) + dtype = _module_floating_dtype(module) + if device is None or device.type != "cpu" or dtype is None or dtype == torch.float32: + continue + promoted.append((module, dtype)) + module.to(dtype=torch.float32) + try: + yield + finally: + for module, dtype in reversed(promoted): + module.to(dtype=dtype) + + +def _torch_is_compiling() -> bool: + dynamo = getattr(torch, "_dynamo", None) + is_compiling = getattr(dynamo, "is_compiling", None) + if callable(is_compiling): + try: + return bool(is_compiling()) + except Exception: + return False + return False + + +_CACTUS_TORCH_LIBRARIES: list[object] = [] +_LFM2_MOE_CUSTOM_OP_REGISTERED = False +_LFM2_MOE_CUSTOM_OP = "cactus_transpile::lfm2_moe_layer_gated" +_QWEN2_MOE_CUSTOM_OP_REGISTERED = False +_QWEN2_MOE_CUSTOM_OP = "cactus_transpile::qwen2_moe_layer_gated" +_GEMMA4_MOE_CUSTOM_OP_REGISTERED = False +_GEMMA4_MOE_CUSTOM_OP = "cactus_transpile::gemma4_moe_layer_gated" + + +def _ignore_duplicate_torch_registration(exc: RuntimeError) -> bool: + text = str(exc).lower() + return "already" in text or "duplicate" in text or "previously" in text + + +def _moe_normalize_and_scale(routing_weights, normalize_routing, epsilon, routed_scaling_factor): + if normalize_routing: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + float(epsilon)) + return routing_weights * float(routed_scaling_factor) + + +def _moe_scatter_experts(hidden, selected_experts, routing_weights, w1_weights, w3_weights, w2_weights, num_experts, activation): + output = torch.zeros_like(hidden) + for expert_idx in range(int(num_experts)): + selected = selected_experts == expert_idx + if not bool(selected.any()): + continue + token_idx, topk_pos = torch.where(selected) + current = hidden[token_idx] + gate = F.linear(current, w1_weights[expert_idx]) + up = F.linear(current, w3_weights[expert_idx]) + expert_out = F.linear(activation(gate) * up, w2_weights[expert_idx]) + expert_out = expert_out * routing_weights[token_idx, topk_pos, None] + output.index_add_(0, token_idx, expert_out.to(dtype=output.dtype)) + return output + + +def _gelu_tanh(x): + return F.gelu(x, approximate="tanh") + + +def _register_moe_custom_op(schema, op_name, impl, fake): + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define(schema) + _CACTUS_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + try: + torch.library.impl(op_name, "CompositeExplicitAutograd")(impl) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + try: + torch.library.register_fake(op_name)(fake) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + + +def _ensure_lfm2_moe_custom_op_registered() -> None: + global _LFM2_MOE_CUSTOM_OP_REGISTERED + if _LFM2_MOE_CUSTOM_OP_REGISTERED: + return + + schema = ( + "lfm2_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, Tensor expert_bias, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool use_expert_bias, " + "bool normalize_routing, float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + + def _impl(hidden, router_logits, expert_bias, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, use_expert_bias, normalize_routing, + epsilon, routed_scaling_factor): + routing_probs = router_logits.sigmoid() + if use_expert_bias: + scores = routing_probs + expert_bias.to(device=routing_probs.device, dtype=routing_probs.dtype) + _, selected_experts = torch.topk(scores, k=int(num_experts_per_tok), dim=-1) + routing_weights = torch.gather(routing_probs, dim=1, index=selected_experts).type_as(router_logits) + else: + routing_weights, selected_experts = torch.topk(routing_probs, k=int(num_experts_per_tok), dim=-1) + routing_weights = _moe_normalize_and_scale(routing_weights, normalize_routing, epsilon, routed_scaling_factor) + return _moe_scatter_experts(hidden, selected_experts, routing_weights, w1_weights, w3_weights, w2_weights, num_experts, F.silu) + + def _fake(hidden, router_logits, expert_bias, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, use_expert_bias, normalize_routing, + epsilon, routed_scaling_factor): + return hidden.new_empty(hidden.shape) + + _register_moe_custom_op(schema, _LFM2_MOE_CUSTOM_OP, _impl, _fake) + _LFM2_MOE_CUSTOM_OP_REGISTERED = True + + +def _ensure_qwen2_moe_custom_op_registered() -> None: + global _QWEN2_MOE_CUSTOM_OP_REGISTERED + if _QWEN2_MOE_CUSTOM_OP_REGISTERED: + return + + schema = ( + "qwen2_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool normalize_routing, " + "float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + + def _impl(hidden, router_logits, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor): + routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32).to(dtype=router_logits.dtype) + routing_weights, selected_experts = torch.topk(routing_probs, k=int(num_experts_per_tok), dim=-1) + routing_weights = _moe_normalize_and_scale(routing_weights, normalize_routing, epsilon, routed_scaling_factor) + return _moe_scatter_experts(hidden, selected_experts, routing_weights, w1_weights, w3_weights, w2_weights, num_experts, F.silu) + + def _fake(hidden, router_logits, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor): + return hidden.new_empty(hidden.shape) + + _register_moe_custom_op(schema, _QWEN2_MOE_CUSTOM_OP, _impl, _fake) + _QWEN2_MOE_CUSTOM_OP_REGISTERED = True + + +def _ensure_gemma4_moe_custom_op_registered() -> None: + global _GEMMA4_MOE_CUSTOM_OP_REGISTERED + if _GEMMA4_MOE_CUSTOM_OP_REGISTERED: + return + + schema = ( + "gemma4_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool normalize_routing, " + "float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + + def _impl(hidden, router_logits, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor): + routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32).to(dtype=router_logits.dtype) + selected_experts = torch.topk(router_logits, k=int(num_experts_per_tok), dim=-1).indices + routing_weights = torch.gather(routing_probs, -1, selected_experts) + routing_weights = _moe_normalize_and_scale(routing_weights, normalize_routing, epsilon, routed_scaling_factor) + return _moe_scatter_experts(hidden, selected_experts, routing_weights, w1_weights, w3_weights, w2_weights, num_experts, _gelu_tanh) + + def _fake(hidden, router_logits, w1_weights, w3_weights, w2_weights, + num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor): + return hidden.new_empty(hidden.shape) + + _register_moe_custom_op(schema, _GEMMA4_MOE_CUSTOM_OP, _impl, _fake) + _GEMMA4_MOE_CUSTOM_OP_REGISTERED = True + + +class _Lfm2MoeRuntimeMoeBlock(torch.nn.Module): + + def __init__(self, block: torch.nn.Module, *, layer_index: int): + super().__init__() + _ensure_lfm2_moe_custom_op_registered() + experts = getattr(block, "experts", None) + gate_up = getattr(experts, "gate_up_proj", None) + down = getattr(experts, "down_proj", None) + gate = getattr(block, "gate", None) + if not isinstance(gate, torch.nn.Module) or not isinstance(gate_up, torch.Tensor) or not isinstance(down, torch.Tensor): + raise TypeError("LFM2-MoE runtime block requires gate, gate_up_proj, and down_proj") + + self.layer_index = int(layer_index) + self.gate = gate + self.num_experts = int(getattr(block, "num_experts", int(gate_up.shape[0])) or int(gate_up.shape[0])) + self.top_k = int(getattr(block, "top_k", 0) or 0) + self.use_expert_bias = bool(getattr(block, "use_expert_bias", False)) + self.normalize_routing = bool(getattr(block, "norm_topk_prob", False)) + self.routed_scaling_factor = float(getattr(block, "routed_scaling_factor", 1.0) or 1.0) + self.epsilon = 1.0e-6 + self.hidden_dim = int(gate_up.shape[2]) + self.intermediate_dim = int(gate_up.shape[1]) // 2 + + if self.num_experts <= 0 or self.top_k <= 0: + raise ValueError("LFM2-MoE runtime block requires positive num_experts and top_k") + if int(gate_up.shape[0]) != self.num_experts or int(down.shape[0]) != self.num_experts: + raise ValueError("LFM2-MoE expert tensor count mismatch") + if int(gate_up.shape[1]) != 2 * self.intermediate_dim: + raise ValueError("LFM2-MoE gate_up_proj must have an even expert dimension") + + self.w1_weights = torch.nn.ParameterList() + self.w3_weights = torch.nn.ParameterList() + self.w2_weights = torch.nn.ParameterList() + for expert_idx in range(self.num_experts): + self.w1_weights.append(torch.nn.Parameter( + gate_up[expert_idx, : self.intermediate_dim, :].detach().contiguous(), + requires_grad=False, + )) + self.w3_weights.append(torch.nn.Parameter( + gate_up[expert_idx, self.intermediate_dim :, :].detach().contiguous(), + requires_grad=False, + )) + self.w2_weights.append(torch.nn.Parameter(down[expert_idx].detach().contiguous(), requires_grad=False)) + + expert_bias = getattr(block, "expert_bias", None) + if self.use_expert_bias and isinstance(expert_bias, torch.Tensor): + bias_value = expert_bias.detach().to(dtype=torch.float32) + else: + bias_value = torch.zeros(self.num_experts, dtype=torch.float32, device=gate_up.device) + self.register_buffer("expert_bias", bias_value, persistent=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.reshape(-1, hidden_dim) + router_logits = self.gate(hidden_flat) + output = torch.ops.cactus_transpile.lfm2_moe_layer_gated( + hidden_flat, + router_logits, + self.expert_bias, + list(self.w1_weights), + list(self.w3_weights), + list(self.w2_weights), + self.num_experts, + self.top_k, + self.use_expert_bias, + self.normalize_routing, + self.epsilon, + self.routed_scaling_factor, + ) + return output.reshape(batch_size, sequence_length, hidden_dim) + + +def _is_lfm2_moe_sparse_block(block: object) -> bool: + experts = getattr(block, "experts", None) + return ( + isinstance(block, torch.nn.Module) + and isinstance(getattr(block, "gate", None), torch.nn.Module) + and isinstance(getattr(experts, "gate_up_proj", None), torch.Tensor) + and isinstance(getattr(experts, "down_proj", None), torch.Tensor) + ) + + +def _is_lfm2_moe_model(model: torch.nn.Module) -> bool: + config = getattr(model, "config", None) + model_type = str(getattr(config, "model_type", "") or "").lower() + if model_type == "lfm2_moe": + return True + return type(model).__module__.startswith("transformers.models.lfm2_moe.") + + +def _lfm2_backbone_from_model(model: torch.nn.Module) -> torch.nn.Module | None: + model_root = getattr(model, "model", None) + language_model = getattr(model_root, "language_model", None) + backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root + return backbone if isinstance(backbone, torch.nn.Module) else None + + +def _lfm2_position_embeddings( + backbone: torch.nn.Module, + hidden_states: torch.Tensor, + *, + position_ids: torch.Tensor, +) -> object: + rotary_emb = getattr(backbone, "rotary_emb", None) + if callable(rotary_emb): + return rotary_emb(hidden_states, position_ids=position_ids) + pos_emb = getattr(backbone, "pos_emb", None) + if callable(pos_emb): + return pos_emb(hidden_states, position_ids=position_ids) + raise AttributeError(f"{type(backbone).__name__} has neither rotary_emb nor pos_emb") + + +def _lfm2_causal_mask_for_capture( + create_causal_mask: Callable[..., object], + *, + backbone: torch.nn.Module, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + position_ids: torch.Tensor, +) -> object | None: + if inputs_embeds.is_meta: + return None + return create_causal_mask( + config=backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=None, + position_ids=position_ids, + ) + + +def _patch_lfm2_moe_feed_forwards(model: torch.nn.Module) -> None: + if not _is_lfm2_moe_model(model): + return + backbone = _lfm2_backbone_from_model(model) + layers = getattr(backbone, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + return + for layer_index, layer in enumerate(layers): + feed_forward = getattr(layer, "feed_forward", None) + if isinstance(feed_forward, _Lfm2MoeRuntimeMoeBlock): + continue + if _is_lfm2_moe_sparse_block(feed_forward): + setattr(layer, "feed_forward", _Lfm2MoeRuntimeMoeBlock(feed_forward, layer_index=layer_index)) + + +class _Qwen2MoeRuntimeMoeBlock(torch.nn.Module): + + def __init__(self, block: torch.nn.Module, *, layer_index: int): + super().__init__() + _ensure_qwen2_moe_custom_op_registered() + experts = getattr(block, "experts", None) + gate_up = getattr(experts, "gate_up_proj", None) + down = getattr(experts, "down_proj", None) + gate = getattr(block, "gate", None) + gate_weight = getattr(gate, "weight", None) + shared_expert = getattr(block, "shared_expert", None) + shared_expert_gate = getattr(block, "shared_expert_gate", None) + if ( + not isinstance(gate, torch.nn.Module) + or not isinstance(gate_weight, torch.Tensor) + or not isinstance(gate_up, torch.Tensor) + or not isinstance(down, torch.Tensor) + or not isinstance(shared_expert, torch.nn.Module) + or not isinstance(shared_expert_gate, torch.nn.Module) + ): + raise TypeError( + "Qwen2-MoE runtime block requires gate, packed experts, shared_expert, and shared_expert_gate" + ) + + self.layer_index = int(layer_index) + self.gate = gate + self.shared_expert = shared_expert + self.shared_expert_gate = shared_expert_gate + config = getattr(block, "config", None) + def _route_attr(*names, default=None): + for obj in (block, config): + if obj is None: + continue + for name in names: + val = getattr(obj, name, None) + if val is not None: + return val + return default + self.num_experts = int(_route_attr("num_experts", default=int(gate_up.shape[0])) or int(gate_up.shape[0])) + self.top_k = int(_route_attr("top_k", "num_experts_per_tok", default=0) or 0) + self.normalize_routing = bool(_route_attr("norm_topk_prob", default=False)) + self.routed_scaling_factor = float(getattr(block, "routed_scaling_factor", 1.0) or 1.0) + self.epsilon = 1.0e-6 + self.hidden_dim = int(gate_up.shape[2]) + self.intermediate_dim = int(gate_up.shape[1]) // 2 + + if self.num_experts <= 0 or self.top_k <= 0: + raise ValueError("Qwen2-MoE runtime block requires positive num_experts and top_k") + if int(gate_up.shape[0]) != self.num_experts or int(down.shape[0]) != self.num_experts: + raise ValueError("Qwen2-MoE expert tensor count mismatch") + if int(gate_up.shape[1]) != 2 * self.intermediate_dim: + raise ValueError("Qwen2-MoE gate_up_proj must have an even expert dimension") + + self.w1_weights = torch.nn.ParameterList() + self.w3_weights = torch.nn.ParameterList() + self.w2_weights = torch.nn.ParameterList() + for expert_idx in range(self.num_experts): + self.w1_weights.append(torch.nn.Parameter( + gate_up[expert_idx, : self.intermediate_dim, :].detach().contiguous(), + requires_grad=False, + )) + self.w3_weights.append(torch.nn.Parameter( + gate_up[expert_idx, self.intermediate_dim :, :].detach().contiguous(), + requires_grad=False, + )) + self.w2_weights.append(torch.nn.Parameter(down[expert_idx].detach().contiguous(), requires_grad=False)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.reshape(-1, hidden_dim) + router_logits = F.linear(hidden_flat, self.gate.weight) + routed_output = torch.ops.cactus_transpile.qwen2_moe_layer_gated( + hidden_flat, + router_logits, + list(self.w1_weights), + list(self.w3_weights), + list(self.w2_weights), + self.num_experts, + self.top_k, + self.normalize_routing, + self.epsilon, + self.routed_scaling_factor, + ) + shared_output = self.shared_expert(hidden_flat) + shared_output = torch.sigmoid(self.shared_expert_gate(hidden_flat)) * shared_output + output = routed_output + shared_output + return output.reshape(batch_size, sequence_length, hidden_dim) + + +def _is_qwen2_moe_sparse_block(block: object) -> bool: + experts = getattr(block, "experts", None) + gate = getattr(block, "gate", None) + return ( + isinstance(block, torch.nn.Module) + and isinstance(gate, torch.nn.Module) + and isinstance(getattr(gate, "weight", None), torch.Tensor) + and isinstance(getattr(experts, "gate_up_proj", None), torch.Tensor) + and isinstance(getattr(experts, "down_proj", None), torch.Tensor) + and isinstance(getattr(block, "shared_expert", None), torch.nn.Module) + and isinstance(getattr(block, "shared_expert_gate", None), torch.nn.Module) + ) + + +def _is_qwen2_moe_model(model: torch.nn.Module) -> bool: + config = getattr(model, "config", None) + model_type = str(getattr(config, "model_type", "") or "").lower() + if model_type == "qwen2_moe": + return True + return type(model).__module__.startswith("transformers.models.qwen2_moe.") + + +def _qwen2_moe_backbone_from_model(model: torch.nn.Module) -> torch.nn.Module | None: + backbone = getattr(model, "model", None) + return backbone if isinstance(backbone, torch.nn.Module) else None + + +def _patch_qwen2_moe_mlps(model: torch.nn.Module) -> None: + if not _is_qwen2_moe_model(model): + return + backbone = _qwen2_moe_backbone_from_model(model) + layers = getattr(backbone, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + return + for layer_index, layer in enumerate(layers): + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, _Qwen2MoeRuntimeMoeBlock): + continue + if _is_qwen2_moe_sparse_block(mlp): + setattr(layer, "mlp", _Qwen2MoeRuntimeMoeBlock(mlp, layer_index=layer_index)) + + +class _Gemma4RuntimeMoeBlock(torch.nn.Module): + + def __init__( + self, + block: torch.nn.Module, + *, + layer_index: int, + per_expert_scale: torch.Tensor | None = None, + ): + super().__init__() + _ensure_gemma4_moe_custom_op_registered() + gate_up = getattr(block, "gate_up_proj", None) + down = getattr(block, "down_proj", None) + if per_expert_scale is None: + per_expert_scale = getattr(block, "per_expert_scale", None) + if ( + not isinstance(gate_up, torch.Tensor) + or not isinstance(down, torch.Tensor) + or not isinstance(per_expert_scale, torch.Tensor) + ): + raise TypeError("Gemma4 MoE runtime block requires gate_up_proj, down_proj, and per_expert_scale") + + self.layer_index = int(layer_index) + self.num_experts = int(getattr(block, "num_experts", int(gate_up.shape[0])) or int(gate_up.shape[0])) + self.hidden_dim = int(gate_up.shape[2]) + self.intermediate_dim = int(gate_up.shape[1]) // 2 + if self.num_experts <= 0: + raise ValueError("Gemma4 MoE runtime block requires positive num_experts") + if int(gate_up.shape[0]) != self.num_experts or int(down.shape[0]) != self.num_experts: + raise ValueError("Gemma4 MoE expert tensor count mismatch") + if int(gate_up.shape[1]) != 2 * self.intermediate_dim: + raise ValueError("Gemma4 MoE gate_up_proj must have an even expert dimension") + + self.w1_weights = torch.nn.ParameterList() + self.w3_weights = torch.nn.ParameterList() + self.w2_weights = torch.nn.ParameterList() + for expert_idx in range(self.num_experts): + self.w1_weights.append(torch.nn.Parameter( + gate_up[expert_idx, : self.intermediate_dim, :].detach().contiguous(), + requires_grad=False, + )) + self.w3_weights.append(torch.nn.Parameter( + gate_up[expert_idx, self.intermediate_dim :, :].detach().contiguous(), + requires_grad=False, + )) + scale = per_expert_scale[expert_idx].detach().to(device=down.device, dtype=down.dtype) + self.w2_weights.append(torch.nn.Parameter( + (down[expert_idx].detach() * scale).contiguous(), + requires_grad=False, + )) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + top_k: int, + normalize_routing: bool = True, + epsilon: float = 1.0e-6, + ) -> torch.Tensor: + return torch.ops.cactus_transpile.gemma4_moe_layer_gated( + hidden_states, + router_logits, + list(self.w1_weights), + list(self.w3_weights), + list(self.w2_weights), + self.num_experts, + int(top_k), + bool(normalize_routing), + float(epsilon), + 1.0, + ) + + +class _Gemma4MoeRuntimeDecoderLayer(torch.nn.Module): + """Gemma4 decoder layer that preserves HF math while exporting MoE as one semantic op.""" + + def __init__(self, layer: torch.nn.Module, *, layer_index: int): + super().__init__() + self.config = layer.config + self.hidden_size = layer.hidden_size + self.layer_idx = layer.layer_idx + self.self_attn = layer.self_attn + config_layer_types = tuple(getattr(self.config, "layer_types", ())) + self.attention_type = getattr(layer, "attention_type", getattr(self.self_attn, "attention_type", None)) + if self.attention_type is None: + self.attention_type = ( + config_layer_types[layer_index] + if layer_index < len(config_layer_types) + else "full_attention" + ) + self.mlp = layer.mlp + self.input_layernorm = layer.input_layernorm + self.post_attention_layernorm = layer.post_attention_layernorm + self.pre_feedforward_layernorm = layer.pre_feedforward_layernorm + self.post_feedforward_layernorm = layer.post_feedforward_layernorm + self.register_buffer("layer_scalar", layer.layer_scalar.detach().clone(), persistent=True) + self.hidden_size_per_layer_input = getattr(layer, "hidden_size_per_layer_input", 0) + if self.hidden_size_per_layer_input: + self.act_fn = layer.act_fn + self.per_layer_input_gate = layer.per_layer_input_gate + self.per_layer_projection = layer.per_layer_projection + self.post_per_layer_input_norm = layer.post_per_layer_input_norm + self.enable_moe_block = True + self.router = layer.router + self.moe = _Gemma4RuntimeMoeBlock( + layer.experts, + layer_index=layer_index, + per_expert_scale=layer.router.per_expert_scale, + ) + self.post_feedforward_layernorm_1 = layer.post_feedforward_layernorm_1 + self.post_feedforward_layernorm_2 = layer.post_feedforward_layernorm_2 + self.pre_feedforward_layernorm_2 = layer.pre_feedforward_layernorm_2 + self.top_k_experts = int(getattr(self.config, "top_k_experts", 0) or 0) + self.router_epsilon = float(getattr(self.config, "rms_norm_eps", 1.0e-6) or 1.0e-6) + + def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + router_states = self.router.norm(hidden_states) + scalar_root_size = getattr(self.router, "scalar_root_size", None) + if scalar_root_size is None: + root_size = getattr(self.router, "root_size", 1.0) + scalar_root_size = root_size.to(router_states.dtype) if isinstance(root_size, torch.Tensor) else float(root_size) + router_states = router_states * self.router.scale.to(router_states.dtype) + router_states = router_states * scalar_root_size + return self.router.proj(router_states) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: torch.Tensor = None, + per_layer_input: torch.Tensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: object | None = None, + use_cache: bool | None = False, + shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None = None, + **kwargs, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + shared_kv_states=shared_kv_states, + use_cache=use_cache, + **kwargs, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states_1 = self.pre_feedforward_layernorm(hidden_states) + hidden_states_1 = self.mlp(hidden_states_1) + hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states_1) + + hidden_states_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + router_logits = self._router_logits(hidden_states_flat) + hidden_states_2 = self.pre_feedforward_layernorm_2(hidden_states_flat) + hidden_states_2 = self.moe( + hidden_states_2, + router_logits, + top_k=self.top_k_experts, + normalize_routing=True, + epsilon=self.router_epsilon, + ) + hidden_states_2 = hidden_states_2.reshape(hidden_states.shape) + hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) + + hidden_states = hidden_states_1 + hidden_states_2 + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + if self.hidden_size_per_layer_input: + residual = hidden_states + hidden_states = self.per_layer_input_gate(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = torch.multiply(hidden_states, per_layer_input) + hidden_states = self.per_layer_projection(hidden_states) + hidden_states = self.post_per_layer_input_norm(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = hidden_states * self.layer_scalar + return hidden_states + + +def _is_gemma4_moe_decoder_layer(layer: object) -> bool: + return ( + isinstance(layer, torch.nn.Module) + and bool(getattr(layer, "enable_moe_block", False)) + and hasattr(layer, "router") + and hasattr(layer, "experts") + and not isinstance(layer, _Gemma4MoeRuntimeDecoderLayer) + ) + + +def _gemma4_backbone_from_model(model: torch.nn.Module) -> torch.nn.Module | None: + model_root = getattr(model, "model", None) + backbone = getattr(model_root, "language_model", model_root) + return backbone if isinstance(backbone, torch.nn.Module) else None + + +def _patch_gemma4_moe_layers(model: torch.nn.Module) -> None: + if _family_key(model) != "gemma4": + return + backbone = _gemma4_backbone_from_model(model) + layers = getattr(backbone, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + return + for layer_index, layer in enumerate(layers): + if isinstance(layer, _Gemma4MoeRuntimeDecoderLayer): + continue + if _is_gemma4_moe_decoder_layer(layer): + layers[layer_index] = _Gemma4MoeRuntimeDecoderLayer(layer, layer_index=layer_index) + + +def _gemma4_cpu_safe_text_mlp_enabled(layer: torch.nn.Module, hidden_states: torch.Tensor) -> bool: + if hidden_states.device.type != "cpu" or hidden_states.dtype != torch.float16: + return False + return _gemma4_text_mlp_manual_available(layer) + + +def _gemma4_text_mlp_manual_available(layer: torch.nn.Module) -> bool: + if bool(getattr(layer, "enable_moe_block", False)): + return False + mlp = getattr(layer, "mlp", None) + if mlp is None: + return False + for attr_name in ("gate_proj", "up_proj", "down_proj", "act_fn"): + if not hasattr(mlp, attr_name): + return False + return True + + +def _gemma4_cpu_safe_text_mlp_forward(mlp: torch.nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + gate = F.linear( + hidden_states, + mlp.gate_proj.weight, + mlp.gate_proj.bias, + ) + up = F.linear( + hidden_states, + mlp.up_proj.weight, + mlp.up_proj.bias, + ) + # Gemma4 text MLPs can overflow in FP16 at the gated product. + # Scaling one branch here keeps the product and down-projection finite. + # The following post-feedforward RMSNorm is scale-invariant, so this + # constant factor cancels back out in the normalized output. + activated = mlp.act_fn(gate) * _GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE + return F.linear( + activated * up, + mlp.down_proj.weight, + mlp.down_proj.bias, + ) + + +def _gemma4_text_attention_forward( + attn: torch.nn.Module, + hidden_states: torch.Tensor, + *, + position_embeddings: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None, + attention_mask: torch.Tensor | None, + position_ids: torch.LongTensor | None, + past_key_values: object | None, + use_cache: bool, + shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None, +) -> torch.Tensor: + if shared_kv_states is None: + attn_out, _ = attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + shared_kv_states={}, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + ) + return attn_out + + from transformers.models.gemma4.modeling_gemma4 import ALL_ATTENTION_FUNCTIONS # type: ignore + from transformers.models.gemma4.modeling_gemma4 import apply_rotary_pos_emb # type: ignore + from transformers.models.gemma4.modeling_gemma4 import eager_attention_forward # type: ignore + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, attn.head_dim) + if position_embeddings is None: + raise ValueError("Gemma4 text attention requires position embeddings") + cos, sin = position_embeddings + + query_states = attn.q_proj(hidden_states).view(hidden_shape) + query_states = attn.q_norm(query_states) + query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2) + query_states = query_states.transpose(1, 2) + + kv_shared_key = getattr(attn, "layer_type", getattr(attn, "kv_shared_layer_index", None)) + if bool(getattr(attn, "is_kv_shared_layer", False)) and kv_shared_key in shared_kv_states: + key_states, value_states = shared_kv_states[kv_shared_key] + key_states = key_states.to(query_states.device) + value_states = value_states.to(query_states.device) + else: + key_states = attn.k_proj(hidden_states).view(hidden_shape) + value_states = attn.v_proj(hidden_states).view(hidden_shape) if attn.v_proj is not None else key_states + + key_states = attn.k_norm(key_states) + key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2) + key_states = key_states.transpose(1, 2) + + value_states = attn.v_norm(value_states) + value_states = value_states.transpose(1, 2) + + layer_idx = getattr(attn, "layer_idx", None) + if bool(getattr(attn, "store_full_length_kv", False)): + shared_kv_states[kv_shared_key] = (key_states, value_states) + elif layer_idx is not None and not bool(getattr(attn, "is_kv_shared_layer", False)): + shared_kv_states[int(layer_idx)] = (key_states, value_states) + + attention_interface = eager_attention_forward + if getattr(attn.config, "_attn_implementation", "eager") != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[attn.config._attn_implementation] + + attn_output, _ = attention_interface( + attn, + query_states, + key_states, + value_states, + attention_mask, + dropout=attn.attention_dropout if attn.training else 0.0, + scaling=attn.scaling, + sliding_window=attn.sliding_window, + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return attn.o_proj(attn_output) + + +def _gemma4_text_decoder_layer_forward( + layer: torch.nn.Module, + hidden_states: torch.Tensor, + *, + position_embeddings: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None, + per_layer_input: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: object | None = None, + use_cache: bool = False, + shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None = None, +) -> torch.Tensor: + if isinstance(layer, _Gemma4MoeRuntimeDecoderLayer): + return layer( + hidden_states, + position_embeddings=position_embeddings, + per_layer_input=per_layer_input, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + shared_kv_states=shared_kv_states, + ) + manual_attention_required = shared_kv_states is not None + if not manual_attention_required and not _gemma4_cpu_safe_text_mlp_enabled(layer, hidden_states): + return layer( + hidden_states, + position_embeddings=position_embeddings, + per_layer_input=per_layer_input, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + ) + if not _gemma4_text_mlp_manual_available(layer): + raise RuntimeError("Gemma4 KV sharing requires a supported text MLP structure") + + residual = hidden_states + + hidden_states = layer.input_layernorm(hidden_states) + hidden_states = _gemma4_text_attention_forward( + layer.self_attn, + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + shared_kv_states=shared_kv_states, + ) + hidden_states = layer.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.pre_feedforward_layernorm(hidden_states) + mlp_out = _gemma4_cpu_safe_text_mlp_forward(layer.mlp, hidden_states) + hidden_states = layer.post_feedforward_layernorm(mlp_out) + hidden_states = residual.float() + hidden_states + hidden_states = hidden_states.to(dtype=residual.dtype) + + if getattr(layer, "hidden_size_per_layer_input", 0): + residual = hidden_states + hidden_states = layer.per_layer_input_gate(hidden_states) + hidden_states = layer.act_fn(hidden_states) + if per_layer_input is None: + raise ValueError("Gemma4 layer expected per_layer_input but none was provided") + hidden_states = torch.multiply(hidden_states, per_layer_input) + hidden_states = layer.per_layer_projection(hidden_states) + hidden_states = layer.post_per_layer_input_norm(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = hidden_states * layer.layer_scalar + return hidden_states + + +def _gemma4_text_backbone_forward( + backbone: torch.nn.Module, + *, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor | None, + causal_mask_mapping: dict[str, torch.Tensor], + position_ids: torch.LongTensor, + layer_start: int = 0, + layer_end: int | None = None, + apply_norm: bool = True, + shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] | None | object = _UNSET, + capture_layer_index: int | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + hidden_states = inputs_embeds + captured_hidden: torch.Tensor | None = None + layer_types = tuple(dict.fromkeys(getattr(backbone.config, "layer_types", ()))) + position_embeddings = { + layer_type: backbone.rotary_emb(hidden_states, position_ids, layer_type) + for layer_type in layer_types + } + if shared_kv_states is _UNSET: + num_shared_layers = int(getattr(backbone.config, "num_kv_shared_layers", 0) or 0) + has_shared_kv_layers = num_shared_layers > 0 or any( + bool(getattr(getattr(layer, "self_attn", None), "is_kv_shared_layer", False)) + or bool(getattr(getattr(layer, "self_attn", None), "store_full_length_kv", False)) + for layer in getattr(backbone, "layers", ()) + ) + shared_kv_states = {} if has_shared_kv_layers else None + + config_layer_types = tuple(getattr(backbone.config, "layer_types", ())) + end = backbone.config.num_hidden_layers if layer_end is None else min(int(layer_end), int(backbone.config.num_hidden_layers)) + for layer_index in range(max(0, int(layer_start)), end): + decoder_layer = backbone.layers[layer_index] + layer_per_input = None + if per_layer_inputs is not None: + layer_per_input = per_layer_inputs[:, :, decoder_layer.layer_idx, :] + attention_type = getattr( + decoder_layer, + "attention_type", + config_layer_types[layer_index] if layer_index < len(config_layer_types) else "full_attention", + ) + layer_shared_kv_states = shared_kv_states + if shared_kv_states is not None: + layer_attn = getattr(decoder_layer, "self_attn", None) + layer_uses_shared_kv = ( + bool(getattr(layer_attn, "is_kv_shared_layer", False)) + or bool(getattr(layer_attn, "store_full_length_kv", False)) + ) + if not layer_uses_shared_kv: + layer_shared_kv_states = None + hidden_states = _gemma4_text_decoder_layer_forward( + decoder_layer, + hidden_states, + per_layer_input=layer_per_input, + attention_mask=causal_mask_mapping[attention_type], + position_embeddings=position_embeddings[attention_type], + position_ids=position_ids, + past_key_values=None, + use_cache=False, + shared_kv_states=layer_shared_kv_states, + ) + if capture_layer_index is not None and layer_index == int(capture_layer_index): + captured_hidden = hidden_states + + output_hidden = backbone.norm(hidden_states) if apply_norm else hidden_states + if capture_layer_index is None: + return output_hidden + if captured_hidden is None: + captured_hidden = hidden_states + return output_hidden, captured_hidden + + +def _gemma4_strip_audio_padding(audio_output: object) -> torch.Tensor: + audio_features = getattr(audio_output, "pooler_output", None) + if not isinstance(audio_features, torch.Tensor): + raise TypeError("Gemma4 audio output did not expose tensor pooler_output") + + audio_mask_from_encoder = getattr(audio_output, "attention_mask", None) + if isinstance(audio_mask_from_encoder, torch.Tensor): + return audio_features[audio_mask_from_encoder].unsqueeze(0) + + audio_mask_from_encoder = getattr(audio_output, "audio_mel_mask", None) + if not isinstance(audio_mask_from_encoder, torch.Tensor): + return audio_features + all_real_tokens: list[torch.Tensor] = [] + for encodings, padding_mask in zip(audio_features, audio_mask_from_encoder, strict=True): + all_real_tokens.append(encodings[~padding_mask]) + return torch.cat(all_real_tokens, dim=0).unsqueeze(0) + + +def _gemma4_rms_norm_no_scale(hidden_states: torch.Tensor, *, eps: float) -> torch.Tensor: + hidden_states_f32 = hidden_states.float() + mean_squared = hidden_states_f32.pow(2).mean(-1, keepdim=True) + eps + return hidden_states_f32 * torch.pow(mean_squared, -0.5) + + +def _gemma4_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, *, eps: float) -> torch.Tensor: + normed = _gemma4_rms_norm_no_scale(hidden_states, eps=eps) + return normed * weight.to(device=normed.device, dtype=normed.dtype) + + +def _gemma4_layer_norm_eps(multimodal_backbone: torch.nn.Module) -> float: + config = getattr(multimodal_backbone, "config", None) + text_config = None + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_config = get_text_config() + except Exception: + text_config = None + for candidate in (text_config, config): + value = getattr(candidate, "rms_norm_eps", None) + if value is not None: + return float(value) + value = getattr(candidate, "layer_norm_eps", None) + if value is not None: + return float(value) + return 1e-6 + + +def _gemma4_load_cactus_fp_tensor(path: Path) -> torch.Tensor: + with path.open("rb") as handle: + header = handle.read(84) + if len(header) < 84 or header[:4] != CACTUS_MAGIC: + raise RuntimeError(f"Gemma4 Cactus tensor is missing a valid header: {path}") + + alignment = max(1, int(struct.unpack_from(" 0) + precision = int(struct.unpack_from(" 0 else 0 + data_offset = ( + align_offset(scales_offset + scales_bytes, alignment) + if scales_bytes > 0 + else aligned_header + ) + data_count = byte_size // np.dtype(dtype).itemsize + array = np.memmap(path, mode="r", dtype=dtype, offset=data_offset, shape=(data_count,)) + tensor = torch.from_numpy(np.array(array, copy=True)) + if shape: + tensor = tensor.reshape(shape) + return tensor + + +def _gemma4_load_vision_post_proj_norm(weights_dir: str | Path | None) -> torch.Tensor | None: + if weights_dir is None: + return None + path = Path(weights_dir).expanduser() / "embed_vision_post_proj_norm.weights" + if not path.exists(): + return None + tensor = _gemma4_load_cactus_fp_tensor(path) + if tensor.ndim != 1: + raise RuntimeError(f"Gemma4 vision post-projection norm must be 1D, got shape={tuple(tensor.shape)}") + return tensor.float() + + +def _gemma4_can_use_native_like_vision_features(multimodal_backbone: torch.nn.Module) -> bool: + vision_tower = getattr(multimodal_backbone, "vision_tower", None) + embed_vision = getattr(multimodal_backbone, "embed_vision", None) + return ( + isinstance(vision_tower, torch.nn.Module) + and isinstance(embed_vision, torch.nn.Module) + and hasattr(vision_tower, "patch_embedder") + and hasattr(vision_tower, "encoder") + and hasattr(embed_vision, "embedding_projection") + ) + + +def _gemma4_can_use_native_like_audio_features(multimodal_backbone: torch.nn.Module) -> bool: + audio_tower = getattr(multimodal_backbone, "audio_tower", None) + embed_audio = getattr(multimodal_backbone, "embed_audio", None) + layers = getattr(audio_tower, "layers", None) if isinstance(audio_tower, torch.nn.Module) else None + conformer = getattr(audio_tower, "conformer", None) if isinstance(audio_tower, torch.nn.Module) else None + return ( + isinstance(audio_tower, torch.nn.Module) + and isinstance(embed_audio, torch.nn.Module) + and hasattr(audio_tower, "subsample_conv_projection") + and (isinstance(layers, torch.nn.ModuleList) or isinstance(conformer, torch.nn.ModuleList)) + and hasattr(embed_audio, "embedding_projection") + ) + + +def _gemma4_pool_vision_hidden_native_like( + vision_tower: torch.nn.Module, + hidden_states: torch.Tensor, + pixel_position_ids: torch.Tensor, + *, + image_soft_token_counts: tuple[int, ...] | None = None, + image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, +) -> torch.Tensor: + output_length = int(_module_or_config_attr(vision_tower, "default_output_length", 280) or 280) + pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3) + padding_positions = (pixel_position_ids == -1).all(dim=-1) + pooled_batches: list[torch.Tensor] = [] + for row_idx in range(int(hidden_states.shape[0])): + hidden_row = hidden_states[row_idx] + position_row = pixel_position_ids[row_idx] + padding_row = padding_positions[row_idx] + if image_pool_shapes is not None and row_idx < len(image_pool_shapes): + grid_h, grid_w, pooled_count = image_pool_shapes[row_idx] + valid_patch_count = int(grid_h) * int(grid_w) + channels = int(hidden_row.shape[-1]) + valid_hidden = hidden_row[:valid_patch_count].float() + pooled = valid_hidden.reshape( + int(grid_h) // pooling_kernel_size, + pooling_kernel_size, + int(grid_w) // pooling_kernel_size, + pooling_kernel_size, + channels, + ).mean(dim=(1, 3)) + pooled_batches.append(pooled.reshape(int(pooled_count), channels)) + continue + + # Avoid boolean advanced indexing here. Gemma4 image padding is not + # guaranteed to be prefix-shaped, and the generic lowerer optimizes + # some masks as prefix slices. Zero-weighting padded patches preserves + # native pooling semantics while staying easy to lower. + clamped_positions = position_row.clamp(min=0) + max_x = clamped_positions[:, 0].max() + 1 + kernel_positions = torch.div(clamped_positions, pooling_kernel_size, rounding_mode="floor") + kernel_indices = kernel_positions[:, 0] + torch.div( + max_x, + pooling_kernel_size, + rounding_mode="floor", + ) * kernel_positions[:, 1] + valid_patch_weights = torch.logical_not(padding_row).float().unsqueeze(-1) + weights = ( + F.one_hot(kernel_indices.long(), output_length).float() + * valid_patch_weights + / float(pooling_kernel_size**2) + ) + pooled_full = weights.transpose(0, 1) @ hidden_row.float() + if image_soft_token_counts is not None and row_idx < len(image_soft_token_counts): + pooled_batches.append(pooled_full[: int(image_soft_token_counts[row_idx])]) + else: + valid_bins = torch.logical_not((weights == 0).all(dim=0)) + pooled_batches.append(pooled_full[valid_bins]) + pooled_hidden = pooled_batches[0] if len(pooled_batches) == 1 else torch.cat(pooled_batches, dim=0) + hidden_size = int(_module_or_config_attr(vision_tower, "hidden_size", pooled_hidden.shape[-1]) or pooled_hidden.shape[-1]) + pooled_hidden = pooled_hidden * math.sqrt(float(hidden_size)) + if bool(_module_or_config_attr(vision_tower, "standardize", False)): + std_bias = getattr(vision_tower, "std_bias", None) + std_scale = getattr(vision_tower, "std_scale", None) + if std_bias is not None and std_scale is not None: + pooled_hidden = (pooled_hidden - std_bias) * std_scale + return pooled_hidden + + +def _gemma4_vision_encoder_hidden_states( + vision_encoder: torch.nn.Module, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor, + pixel_position_ids: torch.Tensor | None, +) -> torch.Tensor: + config = getattr(vision_encoder, "config", None) + hidden_states = inputs_embeds + rotary_emb = getattr(vision_encoder, "rotary_emb") + layers = getattr(vision_encoder, "layers") + num_layers = int(getattr(config, "num_hidden_layers", len(layers))) + layer_types = tuple(str(value) for value in (getattr(config, "layer_types", ()) or ())) + if layer_types and all(hasattr(rotary_emb, f"{layer_type}_inv_freq") for layer_type in layer_types): + # Newer Gemma4 builds use per-layer vision RoPE tables keyed by the + # layer attention type instead of a single vision RoPE table. + attention_mask_4d = (attention_mask.unsqueeze(-1) * attention_mask.unsqueeze(1)).unsqueeze(1) + position_embeddings_by_type = { + layer_type: rotary_emb(hidden_states, pixel_position_ids, layer_type) + for layer_type in layer_types + } + for decoder_layer in layers[:num_layers]: + layer_type = str(getattr(decoder_layer, "attention_type", layer_types[0])) + hidden_states = decoder_layer( + hidden_states, + attention_mask=attention_mask_4d, + position_embeddings=position_embeddings_by_type[layer_type], + position_ids=pixel_position_ids, + ) + return hidden_states + + if _GEMMA4_CREATE_BIDIRECTIONAL_MASK is None: + raise RuntimeError("Gemma4 bidirectional mask helper is unavailable in this transformers install") + if inputs_embeds.is_meta: + attention_mask = (attention_mask.unsqueeze(-1) * attention_mask.unsqueeze(1)).unsqueeze(1) + else: + attention_mask = _GEMMA4_CREATE_BIDIRECTIONAL_MASK( + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + ) + position_embeddings = rotary_emb(hidden_states, pixel_position_ids) + for decoder_layer in layers[:num_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + position_ids=pixel_position_ids, + ) + return hidden_states + + +def _gemma4_convert_audio_mask_to_blocked_static( + mask_4d: torch.Tensor, + *, + chunk_size: int, + left_context: int, + right_context: int, +) -> torch.Tensor: + """Build Gemma4's blocked local audio mask without tensor index expand. + + Hugging Face's helper uses an expanded gather-index tensor. The v2 graph + runtime does not have a first-class expand op, so we express the same fixed + chunk extraction as static slices and cats. + """ + _, _, seq_len, _ = mask_4d.shape + seq_len_int = int(seq_len) + num_blocks = (seq_len_int + int(chunk_size) - 1) // int(chunk_size) + padded_seq_len = num_blocks * int(chunk_size) + pad_amount = padded_seq_len - seq_len_int + if pad_amount: + mask_4d = F.pad(mask_4d, (0, pad_amount, 0, pad_amount), value=False) + mask_4d = F.pad(mask_4d, (int(left_context), int(right_context)), value=False) + + context_size = int(chunk_size) + int(left_context) + int(right_context) + blocks: list[torch.Tensor] = [] + for block_idx in range(num_blocks): + start = block_idx * int(chunk_size) + block = mask_4d[ + :, + :, + start : start + int(chunk_size), + start : start + context_size, + ] + blocks.append(block.unsqueeze(2)) + return torch.cat(blocks, dim=2) + + +def _gemma4_compute_native_like_image_features( + multimodal_backbone: torch.nn.Module, + pixel_values: torch.Tensor, + pixel_position_ids: torch.Tensor | None, + *, + post_proj_norm_weight: torch.Tensor | None = None, + image_soft_token_counts: tuple[int, ...] | None = None, + image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, +) -> torch.Tensor: + if pixel_position_ids is None: + raise TypeError("Gemma4 native-like vision feature path requires pixel_position_ids") + vision_tower = getattr(multimodal_backbone, "vision_tower", None) + embed_vision = getattr(multimodal_backbone, "embed_vision", None) + if not isinstance(vision_tower, torch.nn.Module) or not isinstance(embed_vision, torch.nn.Module): + raise TypeError("Gemma4 multimodal backbone is missing native-like vision modules") + + padding_positions = (pixel_position_ids == -1).all(dim=-1) + vision_inputs = vision_tower.patch_embedder( + pixel_values, + pixel_position_ids, + padding_positions, + ) + vision_hidden = _gemma4_vision_encoder_hidden_states( + vision_tower.encoder, + vision_inputs, + ~padding_positions, + pixel_position_ids, + ) + pooled_hidden = _gemma4_pool_vision_hidden_native_like( + vision_tower, + vision_hidden, + pixel_position_ids, + image_soft_token_counts=image_soft_token_counts, + image_pool_shapes=image_pool_shapes, + ) + projection = getattr(embed_vision, "embedding_projection", None) + if not isinstance(projection, torch.nn.Linear): + raise TypeError("Gemma4 vision embedder is missing embedding_projection") + pre_projection_norm = getattr(embed_vision, "embedding_pre_projection_norm", None) + if isinstance(pre_projection_norm, torch.nn.Module): + pooled_hidden = pre_projection_norm(pooled_hidden) + pooled_hidden = pooled_hidden.to(dtype=projection.weight.dtype) + projected = F.linear( + pooled_hidden, + projection.weight, + projection.bias, + ) + if post_proj_norm_weight is not None and post_proj_norm_weight.numel() > 0: + return _gemma4_rms_norm( + projected, + post_proj_norm_weight, + eps=_gemma4_layer_norm_eps(multimodal_backbone), + ) + return projected + + +def _gemma4_compute_native_like_audio_features( + multimodal_backbone: torch.nn.Module, + input_features: torch.Tensor, + input_features_mask: torch.Tensor, +) -> torch.Tensor: + audio_tower = getattr(multimodal_backbone, "audio_tower", None) + embed_audio = getattr(multimodal_backbone, "embed_audio", None) + if not isinstance(audio_tower, torch.nn.Module) or not isinstance(embed_audio, torch.nn.Module): + raise TypeError("Gemma4 multimodal backbone is missing native-like audio modules") + + config = getattr(audio_tower, "config", None) + if config is None: + raise TypeError("Gemma4 audio tower is missing config") + hidden_states, output_mask = audio_tower.subsample_conv_projection(input_features, input_features_mask) + + if not hasattr(audio_tower, "rel_pos_enc"): + # Newer Gemma4 audio towers keep relative-position handling inside each + # conformer block. Follow that native path with fixed-shape masks. + chunk_size = int(getattr(config, "conf_attention_chunk_size", getattr(config, "attention_chunk_size", 12))) + right_context = int(getattr(config, "conf_attention_context_right", getattr(config, "attention_context_right", 0))) + left_context = int(getattr(config, "conf_attention_context_left", getattr(config, "attention_context_left", 13))) + max_past_horizon = max(0, left_context - 1) + upper_diagonal = max_past_horizon + right_context + context_size = chunk_size + max_past_horizon + right_context + lower_causal_mask = torch.tril( + torch.ones((context_size, chunk_size), dtype=torch.bool, device=hidden_states.device), + diagonal=0, + ).T + upper_causal_mask = torch.tril( + torch.ones((chunk_size, context_size), dtype=torch.bool, device=hidden_states.device), + diagonal=upper_diagonal, + ) + causal_valid_mask = ( + torch.ones((chunk_size, context_size), dtype=torch.bool, device=hidden_states.device) + * lower_causal_mask + * upper_causal_mask + ) + + layers = getattr(audio_tower, "conformer", None) + if not isinstance(layers, torch.nn.ModuleList): + raise TypeError("Gemma4 audio tower is missing conformer layers") + num_layers = int(getattr(config, "num_hidden_layers", len(layers))) + for encoder_layer in layers[:num_layers]: + hidden_states = encoder_layer(hidden_states, output_mask, causal_valid_mask) + + reduction_factor = int(getattr(config, "conf_reduction_factor", 1) or 1) + if reduction_factor > 1: + hidden_states = hidden_states[:, ::reduction_factor] + if output_mask is not None: + output_mask = output_mask[:, ::reduction_factor] + + output_proj = getattr(audio_tower, "output_proj", None) + audio_encodings = output_proj(hidden_states) if isinstance(output_proj, torch.nn.Linear) else hidden_states + else: + position_embeddings = audio_tower.rel_pos_enc(hidden_states) + seq_len = hidden_states.shape[1] + query_positions = torch.arange(seq_len, device=hidden_states.device)[:, None] + key_positions = torch.arange(seq_len, device=hidden_states.device)[None, :] + distance = query_positions - key_positions + left_context = int(getattr(config, "attention_context_left", 13)) - 1 + right_context = int(getattr(config, "attention_context_right", 0)) + local_mask = ((distance >= 0) & (distance < left_context)) | ((distance < 0) & ((-distance) < right_context)) + attention_mask = local_mask.unsqueeze(0).unsqueeze(0) + if output_mask is not None: + attention_mask = attention_mask & output_mask[:, None, None, :].to(dtype=torch.bool) + attention_mask = _gemma4_convert_audio_mask_to_blocked_static( + attention_mask, + chunk_size=int(getattr(config, "attention_chunk_size", 12)), + left_context=left_context, + right_context=right_context, + ) + + layers = getattr(audio_tower, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + layers = getattr(audio_tower, "conformer", None) + if not isinstance(layers, torch.nn.ModuleList): + raise TypeError("Gemma4 audio tower is missing layers/conformer") + num_layers = int(getattr(config, "num_hidden_layers", len(layers))) + for encoder_layer in layers[:num_layers]: + hidden_states = encoder_layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + + output_proj = getattr(audio_tower, "output_proj", None) + if not isinstance(output_proj, torch.nn.Linear): + raise TypeError("Gemma4 audio tower is missing output_proj") + audio_encodings = output_proj(hidden_states) + + if not callable(getattr(embed_audio, "forward", None)): + raise TypeError("Gemma4 audio embedder is not callable") + projected = embed_audio(inputs_embeds=audio_encodings) + # The transpiled bundle is shape-specialized from representative media. + # Native Gemma4 audio preprocessing emits an unpadded feature tensor for + # that media, so the post-subsampling sequence is already the real token + # sequence. Returning the dense sequence avoids dynamic boolean indexing in + # torch.export and matches the prompt's static audio soft-token count. + return projected + + +def _gemma4_feature_token_count(features: torch.Tensor | None) -> int: + if features is None or features.numel() == 0: + return 0 + if features.ndim == 3: + return int(features.shape[1]) + if features.ndim == 2: + return int(features.shape[0]) + if features.ndim >= 1: + return int(features.reshape(-1, features.shape[-1]).shape[0]) + return 0 + + +def _gemma4_static_image_soft_token_counts( + pixel_position_ids: torch.Tensor | None, + *, + pooling_kernel_size: int, +) -> tuple[int, ...] | None: + if pixel_position_ids is None or pixel_position_ids.ndim != 3: + return None + counts: list[int] = [] + for positions in pixel_position_ids.detach().cpu(): + valid = positions[(positions != -1).any(dim=-1)] + if valid.numel() == 0: + counts.append(0) + continue + max_x = int(valid[:, 0].max().item()) + 1 + max_y = int(valid[:, 1].max().item()) + 1 + counts.append((max_x // int(pooling_kernel_size)) * (max_y // int(pooling_kernel_size))) + return tuple(counts) + + +def _gemma4_static_image_pool_shapes( + pixel_position_ids: torch.Tensor | None, + *, + pooling_kernel_size: int, +) -> tuple[tuple[int, int, int], ...] | None: + if pixel_position_ids is None or pixel_position_ids.ndim != 3: + return None + shapes: list[tuple[int, int, int]] = [] + for positions in pixel_position_ids.detach().cpu(): + valid = positions[(positions != -1).any(dim=-1)] + if valid.numel() == 0: + shapes.append((0, 0, 0)) + continue + grid_w = int(valid[:, 0].max().item()) + 1 + grid_h = int(valid[:, 1].max().item()) + 1 + pooled_count = (grid_w // int(pooling_kernel_size)) * (grid_h // int(pooling_kernel_size)) + shapes.append((grid_h, grid_w, pooled_count)) + return tuple(shapes) + + +def _gemma4_build_native_merge_plan( + multimodal_backbone: torch.nn.Module, + input_ids: torch.Tensor, + *, + image_feature_count: int, + audio_feature_count: int, +) -> _Gemma4NativeMergePlan: + if input_ids.ndim != 2 or int(input_ids.shape[0]) != 1: + raise ValueError( + "Gemma4 native multimodal merge currently expects a static batch-1 input_ids tensor, " + f"got shape {tuple(input_ids.shape)}" + ) + + image_token_id, audio_token_id, pad_token_id = _gemma4_special_token_ids(multimodal_backbone) + token_ids = [int(token) for token in input_ids.detach().cpu().reshape(-1).tolist()] + segments: list[_Gemma4NativeMergeSegment] = [] + pli_token_ids: list[int] = [] + image_offset = 0 + audio_offset = 0 + index = 0 + + while index < len(token_ids): + token_id = token_ids[index] + is_image = image_token_id != 0 and token_id == image_token_id + is_audio = audio_token_id != 0 and token_id == audio_token_id + + if is_image or is_audio: + region_start = index + while index < len(token_ids) and token_ids[index] == token_id: + index += 1 + placeholder_count = index - region_start + if is_image: + insert_count = min(placeholder_count, max(0, image_feature_count - image_offset)) + if insert_count > 0: + segments.append( + _Gemma4NativeMergeSegment( + kind="image", + input_start=region_start, + length=insert_count, + feature_start=image_offset, + ) + ) + pli_token_ids.extend([pad_token_id] * insert_count) + image_offset += insert_count + else: + insert_count = min(placeholder_count, max(0, audio_feature_count - audio_offset)) + if insert_count > 0: + segments.append( + _Gemma4NativeMergeSegment( + kind="audio", + input_start=region_start, + length=insert_count, + feature_start=audio_offset, + ) + ) + pli_token_ids.extend([pad_token_id] * insert_count) + audio_offset += insert_count + continue + + text_start = index + while index < len(token_ids): + next_token = token_ids[index] + if image_token_id != 0 and next_token == image_token_id: + break + if audio_token_id != 0 and next_token == audio_token_id: + break + index += 1 + text_tokens = token_ids[text_start:index] + if text_tokens: + segments.append( + _Gemma4NativeMergeSegment( + kind="text", + input_start=text_start, + length=len(text_tokens), + ) + ) + pli_token_ids.extend(text_tokens) + + if not segments: + raise RuntimeError("Gemma4 native multimodal merge built no input segments") + return _Gemma4NativeMergePlan( + segments=tuple(segments), + pli_token_ids=tuple(pli_token_ids), + ) + + +def _gemma4_feature_sequence( + features: torch.Tensor, + *, + batch_size: int, + feature_dim: int, +) -> torch.Tensor: + if features.ndim == 3: + if int(features.shape[0]) == batch_size: + return features + if int(features.shape[0]) == 1: + return features.expand(batch_size, -1, -1) + if features.ndim == 2: + if int(features.shape[-1]) != feature_dim: + raise ValueError(f"Gemma4 feature dim mismatch: expected {feature_dim}, got {tuple(features.shape)}") + return features.unsqueeze(0).expand(batch_size, -1, -1) + if features.ndim >= 1 and int(features.shape[-1]) == feature_dim: + flattened = features.reshape(-1, feature_dim) + return flattened.unsqueeze(0).expand(batch_size, -1, -1) + raise ValueError(f"unsupported Gemma4 feature tensor shape: {tuple(features.shape)}") + + +def _gemma4_text_embedding_scale(embedding: torch.nn.Module, fallback_scale: float) -> float: + """Return the extra scale needed after calling the HF embedding module. + + Native Cactus scales raw token embedding weights by sqrt(hidden_dim). The + HF Gemma4 embedding module already applies that scale in forward(), so the + transpiler must not multiply a second time when it traces the HF module. + """ + for attr_name in ("scalar_embed_scale", "embed_scale"): + value = getattr(embedding, attr_name, None) + if isinstance(value, torch.Tensor): + try: + if value.numel() > 0 and abs(float(value.reshape(-1)[0].item()) - float(fallback_scale)) < 1e-3: + return 1.0 + except Exception: + continue + elif isinstance(value, (float, int)) and abs(float(value) - float(fallback_scale)) < 1e-3: + return 1.0 + return float(fallback_scale) + + +def _gemma4_apply_native_merge_plan( + model: torch.nn.Module, + *, + input_ids: torch.Tensor, + image_features: torch.Tensor, + audio_features: torch.Tensor, + merge_plan: _Gemma4NativeMergePlan, + pli_token_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + embedding = model.get_input_embeddings() + if not isinstance(embedding, torch.nn.Module): + raise TypeError("Gemma4 model is missing input embeddings") + batch_size = int(input_ids.shape[0]) + feature_dim = int(getattr(embedding, "embedding_dim", 0) or image_features.shape[-1] or audio_features.shape[-1]) + target_dtype = _module_floating_dtype(embedding) or image_features.dtype + image_sequence = _gemma4_feature_sequence( + image_features, + batch_size=batch_size, + feature_dim=feature_dim, + ).to(device=input_ids.device, dtype=target_dtype) + audio_sequence = _gemma4_feature_sequence( + audio_features, + batch_size=batch_size, + feature_dim=feature_dim, + ).to(device=input_ids.device, dtype=target_dtype) + + embedded_segments: list[torch.Tensor] = [] + model_config = getattr(model, "config", None) + text_config = getattr(model_config, "text_config", None) + hidden_scale = float(getattr(model_config, "hidden_size", 0) or 0) + if hidden_scale <= 0.0: + hidden_scale = float(getattr(text_config, "hidden_size", 0) or 0) + if hidden_scale <= 0.0: + hidden_scale = float(feature_dim) + hidden_scale = float(hidden_scale) ** 0.5 + text_extra_scale = _gemma4_text_embedding_scale(embedding, hidden_scale) + for segment in merge_plan.segments: + if segment.kind == "text": + text_tokens = input_ids[:, segment.input_start : segment.input_start + segment.length] + text_embeds = embedding(text_tokens) + if text_extra_scale != 1.0: + text_embeds = text_embeds * text_extra_scale + embedded_segments.append(text_embeds) + elif segment.kind == "image": + embedded_segments.append( + image_sequence[:, segment.feature_start : segment.feature_start + segment.length, :] + ) + elif segment.kind == "audio": + embedded_segments.append( + audio_sequence[:, segment.feature_start : segment.feature_start + segment.length, :] + ) + else: + raise RuntimeError(f"unknown Gemma4 merge segment kind: {segment.kind!r}") + + inputs_embeds = torch.cat(embedded_segments, dim=1) + if pli_token_ids.numel() != inputs_embeds.shape[1]: + raise RuntimeError( + "Gemma4 native merge PLI token count mismatch: " + f"{int(pli_token_ids.numel())} vs {int(inputs_embeds.shape[1])}" + ) + pli_tokens = pli_token_ids.to(device=input_ids.device, dtype=input_ids.dtype).unsqueeze(0) + if batch_size != 1: + pli_tokens = pli_tokens.expand(batch_size, -1) + return inputs_embeds, pli_tokens + + +def _gemma4_remap_sequence_tensor( + tensor: torch.Tensor | None, + *, + merge_plan: _Gemma4NativeMergePlan, +) -> torch.Tensor | None: + if tensor is None: + return None + if tensor.ndim != 2: + raise ValueError(f"Gemma4 merge remap expects a rank-2 tensor, got shape {tuple(tensor.shape)}") + remapped_segments: list[torch.Tensor] = [] + for segment in merge_plan.segments: + remapped_segments.append( + tensor[:, segment.input_start : segment.input_start + segment.length] + ) + if not remapped_segments: + return tensor[:, :0] + return torch.cat(remapped_segments, dim=1) + + +def _gemma4_build_standard_causal_mask_mapping( + *, + create_causal_mask: Callable[..., torch.Tensor] | None, + create_sliding_window_causal_mask: Callable[..., torch.Tensor] | None, + config: object, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + position_ids: torch.Tensor, +) -> dict[str, torch.Tensor | None]: + if not callable(create_causal_mask) or not callable(create_sliding_window_causal_mask): + raise RuntimeError("Gemma4 standard causal mask helpers are unavailable") + if inputs_embeds.is_meta or position_ids.is_meta or (attention_mask is not None and attention_mask.is_meta): + return { + "full_attention": None, + "sliding_attention": None, + } + mask_kwargs = { + "config": config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": None, + "position_ids": position_ids, + } + return { + "full_attention": create_causal_mask(**mask_kwargs), + "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), + } + + +def _infer_input_names(module: torch.nn.Module, *, preferred: tuple[str, ...]) -> tuple[str, ...]: + try: + signature = inspect.signature(module.forward) + except (TypeError, ValueError): + return preferred[:1] + + control_names = { + "self", + "return_dict", + "use_cache", + "past_key_values", + "cache_position", + "position_ids", + "labels", + "decoder_input_ids", + "decoder_attention_mask", + "output_attentions", + "output_hidden_states", + } + available = [ + name + for name, parameter in signature.parameters.items() + if name not in control_names + and parameter.kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + ] + matched = [name for name in preferred if name in available] + if matched: + return tuple(matched) + if available: + return tuple(available[: min(2, len(available))]) + return preferred[:1] + + +class BoundInputAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], family: str, metadata_task: str): + super().__init__() + self.model = model + self.input_names = tuple(input_names) + self.family = family + self.metadata_task = metadata_task + + def _kwargs_from_bound_inputs(self, *bound_inputs: torch.Tensor | None) -> dict[str, torch.Tensor]: + provided = tuple(bound_inputs) + if len(self.input_names) > len(provided): + raise ValueError( + f"adapter expected at most {len(provided)} bound inputs, got {len(self.input_names)} names" + ) + kwargs: dict[str, torch.Tensor] = {} + for index, name in enumerate(self.input_names): + value = provided[index] + if value is None: + raise ValueError(f"missing required bound input {index} for {name}") + kwargs[name] = value + return kwargs + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family=self.family, + adapter_type=type(self).__name__, + input_names=self.input_names, + ), + "task": self.metadata_task, + } + } + + +class CausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = getattr(model, "model", None) + self.lm_head = getattr(model, "lm_head", None) + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + + def forward(self, input_ids: torch.Tensor): + backbone = self.backbone + lm_head = self.lm_head + if isinstance(backbone, torch.nn.Module) and isinstance(lm_head, torch.nn.Module): + backbone_kwargs: dict[str, object] = { + "input_ids": input_ids, + "attention_mask": (input_ids != int(self.pad_token_id)).long() + if self.pad_token_id is not None + else None, + "use_cache": False, + "return_dict": False, + } + outputs = backbone(**_filter_supported_kwargs(backbone, backbone_kwargs)) + hidden_states = _extract_tensor_output(outputs, preferred_field="last_hidden_state") + hidden_states = _select_last_non_pad_token( + hidden_states, + input_ids, + pad_token_id=self.pad_token_id, + ) + return lm_head(hidden_states) + + outputs = self.model( + input_ids=input_ids, + use_cache=False, + return_dict=False, + ) + logits = _extract_tensor_output(outputs, preferred_field="logits") + return _select_last_non_pad_token( + logits, + input_ids, + pad_token_id=self.pad_token_id, + ) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="generic", + adapter_type=type(self).__name__, + input_names=("input_ids",), + ), + } + } + + +class Lfm2CausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + _patch_lfm2_moe_feed_forwards(model) + self.model = model + model_root = getattr(model, "model", None) + language_model = getattr(model_root, "language_model", None) + self.backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root + self.lm_head = getattr(model, "lm_head", None) + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) + from transformers.models.lfm2.modeling_lfm2 import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward(self, input_ids: torch.Tensor): + backbone = self.backbone + lm_head = self.lm_head + if not isinstance(backbone, torch.nn.Module) or not isinstance(lm_head, torch.nn.Module): + raise TypeError("LFM2 causal logits adapter requires backbone and lm_head modules") + + inputs_embeds = backbone.embed_tokens(input_ids) + attention_mask = ( + (input_ids != int(self.pad_token_id)).to(dtype=torch.int64) + if self.pad_token_id is not None + else None + ) + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + causal_mask = _lfm2_causal_mask_for_capture( + self._create_causal_mask, + backbone=backbone, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = _lfm2_position_embeddings(backbone, hidden_states, position_ids=position_ids) + layer_types = tuple(getattr(backbone.config, "layer_types", ())) + linear_attention = attention_mask if inputs_embeds.shape[1] != 1 else None + + for layer_index, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): + layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" + layer_mask = causal_mask if layer_type == "full_attention" else linear_attention + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + ) + + hidden_states = backbone.embedding_norm(hidden_states) + hidden_states = _select_last_non_pad_token( + hidden_states, + input_ids, + pad_token_id=self.pad_token_id, + ) + return lm_head(hidden_states) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family=self.adapter_family, + adapter_type=type(self).__name__, + input_names=("input_ids",), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + } + } + + +class Lfm2CausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + _patch_lfm2_moe_feed_forwards(model) + self.model = model + model_root = getattr(model, "model", None) + language_model = getattr(model_root, "language_model", None) + self.backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root + self.lm_head = getattr(model, "lm_head", None) + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + backbone = self.backbone + lm_head = self.lm_head + if not isinstance(backbone, torch.nn.Module) or not isinstance(lm_head, torch.nn.Module): + raise TypeError("LFM2 causal step adapter requires backbone and lm_head modules") + + inputs_embeds = backbone.embed_tokens(input_ids) + text_position_ids = position_ids.to(dtype=torch.int64) + causal_mask = None + + hidden_states = inputs_embeds + position_embeddings = _lfm2_position_embeddings(backbone, hidden_states, position_ids=text_position_ids) + layer_types = tuple(getattr(backbone.config, "layer_types", ())) + + for layer_index, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): + layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" + layer_mask = causal_mask if layer_type == "full_attention" else None + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_embeddings=position_embeddings, + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + + hidden_states = backbone.embedding_norm(hidden_states) + return lm_head(hidden_states[:, -1:, :]) + + def get_transpile_metadata(self): + layer_types = tuple(str(value) for value in (getattr(self.backbone.config, "layer_types", ()) or ())) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family=self.adapter_family, + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": layer_types, + } + } + + +def _lfm2_vl_model_root(model: torch.nn.Module) -> torch.nn.Module: + root = getattr(model, "model", None) + if not isinstance(root, torch.nn.Module): + raise TypeError("LFM2-VL adapter requires a model.model module") + return root + + +def _lfm2_language_backbone(model: torch.nn.Module) -> torch.nn.Module: + _patch_lfm2_moe_feed_forwards(model) + root = getattr(model, "model", None) + language_model = getattr(root, "language_model", None) + backbone = language_model if isinstance(language_model, torch.nn.Module) else root + if not isinstance(backbone, torch.nn.Module): + raise TypeError("LFM2 adapter requires a language model backbone") + return backbone + + +class Lfm2VlVisionEncoderAdapter(torch.nn.Module): + def __init__( + self, + model: torch.nn.Module, + *, + weights_dir: str | None = None, + ): + super().__init__() + self.model = model + self.weights_dir = weights_dir + root = _lfm2_vl_model_root(model) + vision_tower = getattr(root, "vision_tower", None) + if not isinstance(vision_tower, torch.nn.Module): + raise TypeError("LFM2-VL model is missing a vision_tower module") + vision_model = getattr(vision_tower, "vision_model", None) + if not isinstance(vision_model, torch.nn.Module): + raise TypeError("LFM2-VL vision_tower is missing vision_model") + + self.embeddings = vision_model.embeddings + self.encoder = vision_model.encoder + self.post_layernorm = vision_model.post_layernorm + + def forward( + self, + pixel_values: torch.Tensor, + pixel_attention_mask: torch.Tensor, + positional_embeddings: torch.Tensor, + ) -> torch.Tensor: + target_dtype = self.embeddings.patch_embedding.weight.dtype + hidden_states = self.embeddings.patch_embedding(pixel_values.to(dtype=target_dtype)) + hidden_states = hidden_states + positional_embeddings.to( + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + masked_bias = -30000.0 + seq_len = hidden_states.shape[1] + key_valid = pixel_attention_mask[:, None, None, :].to(dtype=hidden_states.dtype) + encoder_attention_mask = (key_valid * (-masked_bias) + masked_bias).expand( + hidden_states.shape[0], 1, seq_len, seq_len + ) + for encoder_layer in self.encoder.layers: + hidden_states = encoder_layer(hidden_states, encoder_attention_mask) + hidden_states = self.post_layernorm(hidden_states) + return hidden_states + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="lfm2_vl", + adapter_type=type(self).__name__, + input_names=("pixel_values", "pixel_attention_mask", "positional_embeddings"), + ), + "weights_dir": self.weights_dir, + } + } + + +class Lfm2VlVisionProjectorAdapter(torch.nn.Module): + def __init__( + self, + model: torch.nn.Module, + *, + weights_dir: str | None = None, + ): + super().__init__() + self.model = model + self.weights_dir = weights_dir + root = _lfm2_vl_model_root(model) + projector = getattr(root, "multi_modal_projector", None) + if not isinstance(projector, torch.nn.Module): + raise TypeError("LFM2-VL model is missing multi_modal_projector") + self.projector = projector + + def forward(self, vision_features: torch.Tensor) -> torch.Tensor: + hidden_states = vision_features + if self.projector.use_layer_norm: + hidden_states = self.projector.layer_norm(hidden_states) + hidden_states = self.projector.linear_1(hidden_states) + hidden_states = self.projector.act(hidden_states) + hidden_states = self.projector.linear_2(hidden_states) + return hidden_states + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="lfm2_vl", + adapter_type=type(self).__name__, + input_names=("vision_features",), + ), + "weights_dir": self.weights_dir, + } + } + + +class Lfm2VlLMEncoderAdapter(torch.nn.Module): + def __init__( + self, + model: torch.nn.Module, + *, + weights_dir: str | None = None, + ): + super().__init__() + self.model = model + self.weights_dir = weights_dir + self.backbone = _lfm2_language_backbone(model) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + return inputs_embeds, attention_mask.to(dtype=torch.int64), position_ids + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="lfm2_vl", + adapter_type=type(self).__name__, + input_names=("input_ids", "attention_mask"), + ), + "weights_dir": self.weights_dir, + } + } + + +class Lfm2VlLMEncoderStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): + super().__init__() + self.model = model + self.weights_dir = weights_dir + self.backbone = _lfm2_language_backbone(model) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + attention_mask = input_ids.to(dtype=torch.int64) * 0 + 1 + return inputs_embeds, attention_mask, position_ids.to(dtype=torch.int64) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="lfm2_vl", + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + "weights_dir": self.weights_dir, + } + } + + +class Lfm2VlLMEncoderTextChunkAdapter(Lfm2VlLMEncoderStepAdapter): + pass + + +class Lfm2VlDecoderAdapter(torch.nn.Module): + def __init__( + self, + model: torch.nn.Module, + *, + weights_dir: str | None = None, + last_token_only: bool = True, + return_hidden: bool = False, + ): + super().__init__() + self.model = model + self.weights_dir = weights_dir + self.last_token_only = bool(last_token_only) + self.return_hidden = bool(return_hidden) + self.backbone = _lfm2_language_backbone(model) + self.lm_head = getattr(model, "lm_head", None) + if not isinstance(self.lm_head, torch.nn.Module): + raise TypeError("LFM2-VL decoder adapter requires an lm_head module") + from transformers.models.lfm2.modeling_lfm2 import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + causal_mask = _lfm2_causal_mask_for_capture( + self._create_causal_mask, + backbone=self.backbone, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + hidden_states = inputs_embeds + position_embeddings = _lfm2_position_embeddings(self.backbone, hidden_states, position_ids=position_ids) + layer_types = tuple(getattr(self.backbone.config, "layer_types", ())) + linear_attention = attention_mask if inputs_embeds.shape[1] != 1 else None + + for layer_index, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" + layer_mask = causal_mask if layer_type == "full_attention" else linear_attention + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + ) + + hidden_states = self.backbone.embedding_norm(hidden_states) + if self.return_hidden: + return hidden_states + if self.last_token_only: + hidden_states = hidden_states[:, -1:, :] + return self.lm_head(hidden_states) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="lfm2_vl", + adapter_type=type(self).__name__, + input_names=("inputs_embeds", "attention_mask", "position_ids"), + ), + "weights_dir": self.weights_dir, + } + } + + +class Lfm2VlMultimodalCausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...]): + super().__init__() + self.model = model + self.input_names = tuple(input_names) + + def forward(self, *args: torch.Tensor) -> torch.Tensor: + kwargs = { + name: value + for name, value in zip(self.input_names, args, strict=True) + } + outputs = self.model( + **kwargs, + use_cache=False, + return_dict=True, + logits_to_keep=1, + ) + return _extract_tensor_output(outputs, preferred_field="logits") + + def get_transpile_metadata(self): + return { + "graph": _transpile_graph_meta( + self.model, + adapter_family="lfm2_vl", + adapter_type=type(self).__name__, + input_names=self.input_names, + ), + } + + +class GemmaCausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = model.model + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + from transformers.models.gemma.modeling_gemma import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward(self, input_ids: torch.Tensor): + return self.debug_forward(input_ids)[0] + + def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + causal_mask = self._create_causal_mask( + self.backbone.config, + inputs_embeds, + None, + past_key_values=None, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.backbone.rotary_emb(hidden_states, position_ids=position_ids) + checkpoints: list[torch.Tensor] = [] + + for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + position_embeddings=position_embeddings, + ) + checkpoints.append(hidden_states) + + hidden_states = self.backbone.norm(hidden_states) + hidden_states = _select_last_non_pad_token( + hidden_states, + input_ids, + pad_token_id=self.pad_token_id, + ) + checkpoints.append(hidden_states) + return _gemma4_apply_final_logit_softcapping(self.model, self.model.lm_head(hidden_states)), checkpoints + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="gemma", + adapter_type=type(self).__name__, + input_names=("input_ids",), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + } + } + + +def _gemma3_text_backbone_forward(backbone, inputs_embeds, causal_mask_mapping, position_ids, + checkpoints=None): + hidden_states = inputs_embeds + layer_types = backbone.config.layer_types + position_embeddings = { + layer_type: backbone.rotary_emb(hidden_states, position_ids, layer_type) + for layer_type in set(layer_types) + } + for i, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask_mapping[layer_types[i]], + position_embeddings=position_embeddings[layer_types[i]], + position_ids=position_ids, + past_key_values=None, + ) + if checkpoints is not None: + checkpoints.append(hidden_states) + return backbone.norm(hidden_states) + + +def _gemma3_decoder_graph_meta(model, backbone, adapter_type, input_names): + sliding_window = getattr(backbone.config, "sliding_window", None) + return { + "graph": { + **_transpile_graph_meta( + model, + adapter_family="gemma3", + adapter_type=adapter_type, + input_names=input_names, + ), + "num_hidden_layers": int(backbone.config.num_hidden_layers), + "layer_types": tuple(getattr(backbone.config, "layer_types", [])), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + +class Gemma3CausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = model.model + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + from transformers.models.gemma3.modeling_gemma3 import create_causal_mask # type: ignore + from transformers.models.gemma3.modeling_gemma3 import create_sliding_window_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + self._create_sliding_window_causal_mask = create_sliding_window_causal_mask + + def forward(self, input_ids: torch.Tensor): + return self.debug_forward(input_ids)[0] + + def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + causal_mask_mapping = { + "full_attention": self._create_causal_mask( + self.backbone.config, + inputs_embeds, + None, + None, + past_key_values=None, + position_ids=position_ids, + ), + "sliding_attention": self._create_sliding_window_causal_mask( + self.backbone.config, + inputs_embeds, + None, + None, + past_key_values=None, + position_ids=position_ids, + ), + } + + checkpoints: list[torch.Tensor] = [] + hidden_states = _gemma3_text_backbone_forward( + self.backbone, inputs_embeds, causal_mask_mapping, position_ids, checkpoints=checkpoints, + ) + checkpoints.append(hidden_states) + return self.model.lm_head(hidden_states), checkpoints + + def get_transpile_metadata(self): + return _gemma3_decoder_graph_meta(self.model, self.backbone, type(self).__name__, ("input_ids",)) + + +class Gemma3LMEncoderStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.backbone = model.model + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + return inputs_embeds.to(torch.float16), position_ids.to(dtype=torch.int64) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.backbone, + adapter_family="gemma3", + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + } + } + + +class Gemma3EmbedsCausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + self.backbone = model.model + + def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch.Tensor]: + seq_len = int(inputs_embeds.shape[1]) + zero_mask = torch.zeros( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + return {layer_type: zero_mask for layer_type in set(self.backbone.config.layer_types)} + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + hidden_states = inputs_embeds.to(self.backbone.norm.weight.dtype) + causal_mask_mapping = self._additive_mask_mapping(hidden_states) + text_position_ids = position_ids.to(dtype=torch.int64) + hidden_states = _gemma3_text_backbone_forward( + self.backbone, hidden_states, causal_mask_mapping, text_position_ids, + ) + logits = self.model.lm_head(hidden_states[:, -1:, :]) + return _gemma4_apply_final_logit_softcapping(self.model, logits) + + def get_transpile_metadata(self): + return _gemma3_decoder_graph_meta(self.model, self.backbone, type(self).__name__, ("inputs_embeds", "position_ids")) + + +class Gemma3EmbedsCausalLMPrefillChunkAdapter(Gemma3EmbedsCausalLMStepAdapter): + def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch.Tensor]: + seq_len = int(inputs_embeds.shape[1]) + blocked_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * torch.finfo(inputs_embeds.dtype).min + allowed_values = torch.zeros( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + causal = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + ).view(1, 1, seq_len, seq_len) + mapping = {"full_attention": torch.where(causal, allowed_values, blocked_values)} + if "sliding_attention" in set(self.backbone.config.layer_types): + window = int(self.backbone.config.sliding_window) + below_window = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + diagonal=-window, + ).view(1, 1, seq_len, seq_len) + mapping["sliding_attention"] = torch.where( + below_window, blocked_values, mapping["full_attention"], + ) + return mapping + + +class Gemma4CausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + _patch_gemma4_moe_layers(model) + self.model = model + model_backbone = model.model + self.backbone = getattr(model_backbone, "language_model", model_backbone) + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + from transformers.models.gemma4.modeling_gemma4 import create_causal_mask # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_sliding_window_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + self._create_sliding_window_causal_mask = create_sliding_window_causal_mask + + def forward(self, input_ids: torch.Tensor): + return self.debug_forward(input_ids)[0] + + def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + per_layer_inputs = None + if self.backbone.hidden_size_per_layer_input: + per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) + per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) + + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + mask_kwargs = { + "config": self.backbone.config, + "inputs_embeds": inputs_embeds, + "attention_mask": None, + "past_key_values": None, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": self._create_causal_mask(**mask_kwargs), + "sliding_attention": self._create_sliding_window_causal_mask(**mask_kwargs), + } + + hidden_states = inputs_embeds + checkpoints: list[torch.Tensor] = [] + position_embeddings = { + layer_type: self.backbone.rotary_emb(hidden_states, position_ids, layer_type) + for layer_type in self.backbone.unique_layer_types + } + shared_kv_states: dict[str, torch.Tensor] = {} + + for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + per_layer_input = per_layer_inputs[:, :, i, :] if per_layer_inputs is not None else None + hidden_states = decoder_layer( + hidden_states, + per_layer_input, + shared_kv_states=shared_kv_states, + position_embeddings=position_embeddings[self.backbone.config.layer_types[i]], + attention_mask=causal_mask_mapping[self.backbone.config.layer_types[i]], + position_ids=position_ids, + past_key_values=None, + ) + checkpoints.append(hidden_states) + + hidden_states = self.backbone.norm(hidden_states) + checkpoints.append(hidden_states) + return self.model.lm_head(hidden_states), checkpoints + + def debug_first_block(self, input_ids: torch.Tensor) -> dict[str, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + per_layer_inputs = None + if self.backbone.hidden_size_per_layer_input: + per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) + per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) + + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + mask_kwargs = { + "config": self.backbone.config, + "inputs_embeds": inputs_embeds, + "attention_mask": None, + "past_key_values": None, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": self._create_causal_mask(**mask_kwargs), + "sliding_attention": self._create_sliding_window_causal_mask(**mask_kwargs), + } + + hidden_states = inputs_embeds + position_embeddings = { + layer_type: self.backbone.rotary_emb(hidden_states, position_ids, layer_type) + for layer_type in self.backbone.unique_layer_types + } + shared_kv_states: dict[str, torch.Tensor] = {} + layer = self.backbone.layers[0] + layer_type = self.backbone.config.layer_types[0] + per_layer_input = per_layer_inputs[:, :, 0, :] if per_layer_inputs is not None else None + + checkpoints: dict[str, torch.Tensor] = {} + + residual = hidden_states + normed = layer.input_layernorm(hidden_states) + checkpoints["pre_attn_norm"] = normed + + attn_out = layer.self_attn( + normed, + position_embeddings=position_embeddings[layer_type], + attention_mask=causal_mask_mapping[layer_type], + position_ids=position_ids, + past_key_values=None, + shared_kv_states=shared_kv_states, + ) + if isinstance(attn_out, tuple): + attn_out = attn_out[0] + checkpoints["attn_o_proj"] = attn_out + + post_attn_norm = layer.post_attention_layernorm(attn_out) + checkpoints["post_attn_norm"] = post_attn_norm + + after_attention = residual + post_attn_norm + checkpoints["after_attention_residual"] = after_attention + + pre_ffn_norm = layer.pre_feedforward_layernorm(after_attention) + checkpoints["pre_ffn_norm"] = pre_ffn_norm + + mlp_out = layer.mlp(pre_ffn_norm) + checkpoints["mlp_down"] = mlp_out + + post_ffn_norm = layer.post_feedforward_layernorm(mlp_out) + checkpoints["post_ffn_norm"] = post_ffn_norm + + after_ffn = after_attention + post_ffn_norm + checkpoints["after_ffn_residual"] = after_ffn + + if per_layer_input is not None: + gated = layer.per_layer_input_gate(after_ffn) + gated = layer.act_fn(gated) + projected = gated * per_layer_input + per_layer_proj = layer.per_layer_projection(projected) + checkpoints["per_layer_input_proj"] = per_layer_proj + post_per_layer_input_norm = layer.post_per_layer_input_norm(per_layer_proj) + checkpoints["post_per_layer_input_norm"] = post_per_layer_input_norm + after_ffn = after_ffn + post_per_layer_input_norm + + layer_scalar = getattr(layer, "layer_scalar", None) + if layer_scalar is not None: + after_ffn = after_ffn * layer_scalar + checkpoints["layer_scalar_out"] = after_ffn + return checkpoints + + def get_transpile_metadata(self): + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="gemma4", + adapter_type=type(self).__name__, + input_names=("input_ids",), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + +def _gemma4_apply_final_logit_softcapping(model: torch.nn.Module, logits: torch.Tensor) -> torch.Tensor: + config = getattr(model, "config", None) + text_config = None + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_config = get_text_config() + except Exception: + text_config = None + cap = getattr(text_config, "final_logit_softcapping", None) + if cap is None: + cap = getattr(config, "final_logit_softcapping", None) + if cap is None: + return logits + cap_value = float(cap) + if cap_value <= 0.0: + return logits + return torch.tanh(logits / cap_value) * cap_value + + +class Gemma4MultimodalCausalLMLogitsAdapter(BoundInputAdapter): + def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], weights_dir: str | None = None): + _patch_gemma4_moe_layers(model) + super().__init__( + model, + input_names=input_names, + family="gemma4", + metadata_task="multimodal_causal_lm_logits", + ) + model_backbone = model.model + self.multimodal_backbone = model_backbone + self.backbone = getattr(model_backbone, "language_model", model_backbone) + self.last_token_logits_only = False + self._use_cached_multimodal_features = False + self._native_merge_plan: _Gemma4NativeMergePlan | None = None + self.register_buffer("_cached_image_features", torch.empty(0), persistent=False) + self.register_buffer("_cached_audio_features", torch.empty(0), persistent=False) + self.register_buffer("_native_merge_pli_token_ids", torch.empty(0, dtype=torch.long), persistent=False) + vision_post_proj_norm = _gemma4_load_vision_post_proj_norm(weights_dir) + if vision_post_proj_norm is None: + vision_post_proj_norm = torch.empty(0) + self.register_buffer("_cactus_vision_post_proj_norm", vision_post_proj_norm, persistent=False) + self._capture_cpu_float32_text_modules: list[tuple[torch.nn.Module, torch.dtype]] = [] + self._create_causal_mask_mapping = None + self._create_masks_for_generate = None + self._create_causal_mask = None + self._create_sliding_window_causal_mask = None + try: + from transformers.models.gemma4.modeling_gemma4 import create_causal_mask # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_causal_mask_mapping # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_masks_for_generate # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_sliding_window_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + self._create_causal_mask_mapping = create_causal_mask_mapping + self._create_masks_for_generate = create_masks_for_generate + self._create_sliding_window_causal_mask = create_sliding_window_causal_mask + except Exception: + pass + + def _apply_final_logit_softcapping(self, logits: torch.Tensor) -> torch.Tensor: + return _gemma4_apply_final_logit_softcapping(self.model, logits) + + def _capture_text_modules(self) -> tuple[torch.nn.Module, ...]: + return _unique_modules( + self.backbone, + getattr(self.model, "get_input_embeddings", lambda: None)(), + getattr(self.model, "lm_head", None), + ) + + def _compute_image_features( + self, + pixel_values: torch.Tensor, + pixel_position_ids: torch.Tensor | None, + get_image_features: Callable[..., object], + ) -> torch.Tensor: + multimodal_backbone = self.multimodal_backbone + if _gemma4_can_use_native_like_vision_features(multimodal_backbone): + return _gemma4_compute_native_like_image_features( + multimodal_backbone, + pixel_values, + pixel_position_ids, + post_proj_norm_weight=self._cactus_vision_post_proj_norm, + ) + vision_tower = getattr(multimodal_backbone, "vision_tower", None) + embed_vision = getattr(multimodal_backbone, "embed_vision", None) + vision_modules = _gemma4_vision_modules(multimodal_backbone) + if ( + pixel_values.device.type != "cpu" + or pixel_values.dtype != torch.float16 + or len(vision_modules) != 2 + or _torch_is_compiling() + ): + return get_image_features( + pixel_values, + pixel_position_ids, + None, + return_dict=True, + ).pooler_output + + vision_dtype = _module_floating_dtype(vision_tower) + embed_dtype = _module_floating_dtype(embed_vision) + if vision_dtype != torch.float16 and embed_dtype != torch.float16: + return get_image_features( + pixel_values, + pixel_position_ids, + None, + return_dict=True, + ).pooler_output + + # Gemma4's CPU float16 vision path can emit non-finite soft tokens; upcast only + # the static image feature extraction path and restore the original module dtypes. + with _temporary_cpu_float32_modules(vision_modules): + return get_image_features( + pixel_values.float(), + pixel_position_ids, + None, + return_dict=True, + ).pooler_output + + def _compute_audio_features( + self, + input_features: torch.Tensor, + input_features_mask: torch.Tensor, + get_audio_features: Callable[..., object], + ) -> torch.Tensor: + multimodal_backbone = self.multimodal_backbone + if _gemma4_can_use_native_like_audio_features(multimodal_backbone): + return _gemma4_compute_native_like_audio_features( + multimodal_backbone, + input_features, + input_features_mask, + ) + audio_modules = _gemma4_audio_modules(multimodal_backbone) + if ( + input_features.device.type != "cpu" + or input_features.dtype != torch.float16 + or len(audio_modules) != 2 + or _torch_is_compiling() + ): + audio_output = get_audio_features(input_features, input_features_mask, return_dict=True) + return _gemma4_strip_audio_padding(audio_output) + + if all(_module_floating_dtype(module) != torch.float16 for module in audio_modules): + audio_output = get_audio_features(input_features, input_features_mask, return_dict=True) + return _gemma4_strip_audio_padding(audio_output) + + with _temporary_cpu_float32_modules(audio_modules): + audio_output = get_audio_features(input_features.float(), input_features_mask, return_dict=True) + return _gemma4_strip_audio_padding(audio_output) + + def prepare_cpu_float32_capture(self) -> None: + self._capture_cpu_float32_text_modules.clear() + if os.environ.get("CACTUS_GEMMA4_CAPTURE_FP32") != "1": + return + for module in self._capture_text_modules(): + device = _module_device(module) + dtype = _module_floating_dtype(module) + if device is None or device.type != "cpu" or dtype is None or dtype == torch.float32: + continue + self._capture_cpu_float32_text_modules.append((module, dtype)) + module.to(dtype=torch.float32) + + def restore_cpu_float32_capture(self) -> None: + for module, dtype in reversed(self._capture_cpu_float32_text_modules): + module.to(dtype=dtype) + self._capture_cpu_float32_text_modules.clear() + + def _resolve_image_features( + self, + *, + inputs_embeds: torch.Tensor, + pixel_values: torch.Tensor | None, + pixel_position_ids: torch.Tensor | None, + get_image_features: Callable[..., object], + ) -> torch.Tensor | None: + if pixel_values is None: + return None + if self._use_cached_multimodal_features and self._cached_image_features.numel() > 0: + image_features = self._cached_image_features + else: + image_features = self._compute_image_features( + pixel_values, + pixel_position_ids, + get_image_features, + ) + return image_features.to(inputs_embeds.device, inputs_embeds.dtype) + + def _resolve_audio_features( + self, + *, + inputs_embeds: torch.Tensor, + input_features: torch.Tensor | None, + input_features_mask: torch.Tensor | None, + get_audio_features: Callable[..., object], + ) -> torch.Tensor | None: + if input_features is None or input_features_mask is None: + return None + if self._use_cached_multimodal_features and self._cached_audio_features.numel() > 0: + audio_features = self._cached_audio_features + else: + audio_features = self._compute_audio_features( + input_features, + input_features_mask, + get_audio_features, + ) + return audio_features.to(inputs_embeds.device, inputs_embeds.dtype) + + def _prepare_text_backbone_inputs( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + token_type_ids: torch.Tensor | None, + pixel_values: torch.Tensor | None, + pixel_position_ids: torch.Tensor | None, + input_features: torch.Tensor | None, + input_features_mask: torch.Tensor | None, + get_placeholder_mask: Callable[..., object], + get_image_features: Callable[..., object], + get_audio_features: Callable[..., object], + ) -> tuple[torch.Tensor, torch.Tensor | None, dict[str, torch.Tensor], torch.LongTensor]: + inputs_embeds = self.model.get_input_embeddings()(input_ids) + + image_features = self._resolve_image_features( + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_position_ids=pixel_position_ids, + get_image_features=get_image_features, + ) + audio_features = self._resolve_audio_features( + inputs_embeds=inputs_embeds, + input_features=input_features, + input_features_mask=input_features_mask, + get_audio_features=get_audio_features, + ) + + if image_features is not None and audio_features is not None and self._native_merge_plan is not None: + inputs_embeds, per_layer_inputs_tokens = _gemma4_apply_native_merge_plan( + self.model, + input_ids=input_ids, + image_features=image_features, + audio_features=audio_features, + merge_plan=self._native_merge_plan, + pli_token_ids=self._native_merge_pli_token_ids, + ) + attention_mask = _gemma4_remap_sequence_tensor( + attention_mask, + merge_plan=self._native_merge_plan, + ) + token_type_ids = _gemma4_remap_sequence_tensor( + token_type_ids, + merge_plan=self._native_merge_plan, + ) + if token_type_ids is None: + raise RuntimeError("Gemma4 native merge requires token_type_ids for multimodal attention masking") + else: + text_mask, image_mask, audio_mask = _gemma4_get_placeholder_masks( + get_placeholder_mask, + token_type_ids=token_type_ids, + input_ids=input_ids, + inputs_embeds=inputs_embeds, + ) + if image_features is not None: + inputs_embeds = inputs_embeds.masked_scatter( + image_mask.unsqueeze(-1).expand_as(inputs_embeds), + image_features, + ) + if audio_features is not None: + inputs_embeds = inputs_embeds.masked_scatter( + audio_mask.unsqueeze(-1).expand_as(inputs_embeds), + audio_features, + ) + per_layer_inputs_tokens = input_ids * text_mask.to(dtype=input_ids.dtype) + + per_layer_inputs = None + text_config = _gemma4_text_config(self.multimodal_backbone) + if getattr(text_config, "hidden_size_per_layer_input", None): + per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, per_layer_inputs_tokens, inputs_embeds) + per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) + + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + if self._native_merge_plan is not None: + causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( + create_causal_mask=self._create_causal_mask, + create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + elif getattr(text_config, "use_bidirectional_attention", None) == "vision": + causal_mask_mapping = self._create_causal_mask_mapping( + self.multimodal_backbone.config, + inputs_embeds, + attention_mask, + None, + position_ids, + token_type_ids, + pixel_values, + is_training=self.training, + ) + else: + causal_mask_mapping = self._create_masks_for_generate( + self.multimodal_backbone.config, + inputs_embeds, + attention_mask, + None, + position_ids, + ) + return inputs_embeds, per_layer_inputs, causal_mask_mapping, position_ids + + def _forward_hidden_states( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + token_type_ids: torch.Tensor | None, + pixel_values: torch.Tensor | None, + pixel_position_ids: torch.Tensor | None, + input_features: torch.Tensor | None, + input_features_mask: torch.Tensor | None, + get_placeholder_mask: Callable[..., object], + get_image_features: Callable[..., object], + get_audio_features: Callable[..., object], + ) -> torch.Tensor: + inputs_embeds, per_layer_inputs, causal_mask_mapping, position_ids = self._prepare_text_backbone_inputs( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_position_ids=pixel_position_ids, + input_features=input_features, + input_features_mask=input_features_mask, + get_placeholder_mask=get_placeholder_mask, + get_image_features=get_image_features, + get_audio_features=get_audio_features, + ) + return _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + ) + + def prime_static_multimodal_features(self, *bound_inputs: torch.Tensor | None) -> None: + kwargs = self._kwargs_from_bound_inputs(*bound_inputs) + input_ids = kwargs["input_ids"] + token_type_ids = kwargs.get("token_type_ids") + pixel_values = kwargs.get("pixel_values") + pixel_position_ids = kwargs.get("pixel_position_ids") + input_features = kwargs.get("input_features") + input_features_mask = kwargs.get("input_features_mask") + + multimodal_backbone = self.multimodal_backbone + get_placeholder_mask = getattr(multimodal_backbone, "get_placeholder_mask", None) + get_image_features = getattr(multimodal_backbone, "get_image_features", None) + get_audio_features = getattr(multimodal_backbone, "get_audio_features", None) + if not callable(get_placeholder_mask): + return + + with torch.no_grad(): + inputs_embeds = self.model.get_input_embeddings()(input_ids) + _, image_mask, audio_mask = _gemma4_get_placeholder_masks( + get_placeholder_mask, + token_type_ids=token_type_ids, + input_ids=input_ids, + inputs_embeds=inputs_embeds, + ) + + if pixel_values is not None and callable(get_image_features) and image_mask.any(): + image_features = self._compute_image_features( + pixel_values, + pixel_position_ids, + get_image_features, + ) + self._cached_image_features = image_features.detach() + else: + self._cached_image_features = self._cached_image_features.new_empty(0) + + if ( + input_features is not None + and input_features_mask is not None + and callable(get_audio_features) + and audio_mask.any() + ): + self._cached_audio_features = self._compute_audio_features( + input_features, + input_features_mask, + get_audio_features, + ).detach() + else: + self._cached_audio_features = self._cached_audio_features.new_empty(0) + + if self._cached_image_features.numel() > 0 and self._cached_audio_features.numel() > 0: + plan = _gemma4_build_native_merge_plan( + self.multimodal_backbone, + input_ids, + image_feature_count=_gemma4_feature_token_count(self._cached_image_features), + audio_feature_count=_gemma4_feature_token_count(self._cached_audio_features), + ) + self._native_merge_plan = plan + self._native_merge_pli_token_ids = torch.tensor( + plan.pli_token_ids, + dtype=torch.long, + ) + else: + self._native_merge_plan = None + self._native_merge_pli_token_ids = self._native_merge_pli_token_ids.new_empty(0) + + self._use_cached_multimodal_features = True + + def forward(self, *bound_inputs: torch.Tensor | None) -> torch.Tensor: + kwargs = self._kwargs_from_bound_inputs(*bound_inputs) + input_ids = kwargs["input_ids"] + attention_mask = kwargs.get("attention_mask") + token_type_ids = kwargs.get("token_type_ids") + pixel_values = kwargs.get("pixel_values") + pixel_position_ids = kwargs.get("pixel_position_ids") + input_features = kwargs.get("input_features") + input_features_mask = kwargs.get("input_features_mask") + + multimodal_backbone = self.multimodal_backbone + get_placeholder_mask = getattr(multimodal_backbone, "get_placeholder_mask", None) + get_image_features = getattr(multimodal_backbone, "get_image_features", None) + get_audio_features = getattr(multimodal_backbone, "get_audio_features", None) + lm_head = getattr(self.model, "lm_head", None) + if ( + not callable(get_placeholder_mask) + or not callable(get_image_features) + or not callable(get_audio_features) + or not callable(self._create_causal_mask_mapping) + or not callable(self._create_masks_for_generate) + or not isinstance(lm_head, torch.nn.Module) + ): + outputs = self.model( + return_dict=True, + use_cache=False, + **kwargs, + ) + logits = _extract_tensor_output(outputs, preferred_field="logits") + if self.last_token_logits_only and logits.ndim >= 3: + return _select_last_active_token(logits, attention_mask) + return logits + + hidden_states = self._forward_hidden_states( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_position_ids=pixel_position_ids, + input_features=input_features, + input_features_mask=input_features_mask, + get_placeholder_mask=get_placeholder_mask, + get_image_features=get_image_features, + get_audio_features=get_audio_features, + ) + if self.last_token_logits_only: + hidden_states = _select_last_active_token(hidden_states, attention_mask) + return self._apply_final_logit_softcapping(lm_head(hidden_states)) + + def get_transpile_metadata(self): + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="gemma4", + adapter_type=type(self).__name__, + input_names=self.input_names, + ), + "task": self.metadata_task, + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + "last_token_logits_only": bool(self.last_token_logits_only), + } + } + + +_GEMMA4_DECODER_PIPELINE_IO_KEYS = ( + "inputs_embeds", + "per_layer_inputs", + "position_ids", +) + + +class _Gemma4MultimodalComponentBase(torch.nn.Module): + def __init__( + self, + model: torch.nn.Module, + *, + input_names: tuple[str, ...], + weights_dir: str | None = None, + native_merge_plan: _Gemma4NativeMergePlan | None = None, + native_image_soft_token_counts: tuple[int, ...] | None = None, + native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, + ): + super().__init__() + _patch_gemma4_moe_layers(model) + self.model = model + self.input_names = input_names + self.multimodal_backbone = model.model + model_backbone = model.model + self.backbone = getattr(model_backbone, "language_model", model_backbone) + self._native_merge_plan = native_merge_plan + self._native_image_soft_token_counts = native_image_soft_token_counts + self._native_image_pool_shapes = native_image_pool_shapes + self.register_buffer( + "_native_merge_pli_token_ids", + torch.tensor(native_merge_plan.pli_token_ids, dtype=torch.long) + if native_merge_plan is not None + else torch.empty(0, dtype=torch.long), + persistent=False, + ) + vision_post_proj_norm = _gemma4_load_vision_post_proj_norm(weights_dir) + if vision_post_proj_norm is None: + vision_post_proj_norm = torch.empty(0) + self.register_buffer("_cactus_vision_post_proj_norm", vision_post_proj_norm, persistent=False) + self._create_causal_mask_mapping = None + self._create_masks_for_generate = None + self._create_causal_mask = None + self._create_sliding_window_causal_mask = None + try: + from transformers.models.gemma4.modeling_gemma4 import create_causal_mask # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_causal_mask_mapping # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_masks_for_generate # type: ignore + from transformers.models.gemma4.modeling_gemma4 import create_sliding_window_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + self._create_causal_mask_mapping = create_causal_mask_mapping + self._create_masks_for_generate = create_masks_for_generate + self._create_sliding_window_causal_mask = create_sliding_window_causal_mask + except Exception: + pass + self._capture_modules: list[tuple[torch.nn.Module, torch.dtype]] = [] + + def _base_graph_meta(self, *, adapter_type: str, input_names: tuple[str, ...]) -> dict[str, object]: + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + **_transpile_graph_meta( + self.model, + adapter_family="gemma4", + adapter_type=adapter_type, + input_names=input_names, + ), + "task": "multimodal_causal_lm_logits", + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + + def _compute_image_features( + self, + pixel_values: torch.Tensor, + pixel_position_ids: torch.Tensor | None, + ) -> torch.Tensor: + if _gemma4_can_use_native_like_vision_features(self.multimodal_backbone): + return _gemma4_compute_native_like_image_features( + self.multimodal_backbone, + pixel_values, + pixel_position_ids, + post_proj_norm_weight=self._cactus_vision_post_proj_norm, + image_soft_token_counts=self._native_image_soft_token_counts, + image_pool_shapes=self._native_image_pool_shapes, + ) + get_image_features = getattr(self.multimodal_backbone, "get_image_features", None) + vision_tower = getattr(self.multimodal_backbone, "vision_tower", None) + embed_vision = getattr(self.multimodal_backbone, "embed_vision", None) + vision_modules = _gemma4_vision_modules(self.multimodal_backbone) + if not callable(get_image_features): + raise TypeError("Gemma4 multimodal backbone is missing get_image_features") + if ( + pixel_values.device.type != "cpu" + or pixel_values.dtype != torch.float16 + or len(vision_modules) != 2 + or _torch_is_compiling() + ): + return get_image_features( + pixel_values, + pixel_position_ids, + None, + return_dict=True, + ).pooler_output + + vision_dtype = _module_floating_dtype(vision_tower) + embed_dtype = _module_floating_dtype(embed_vision) + if vision_dtype != torch.float16 and embed_dtype != torch.float16: + return get_image_features( + pixel_values, + pixel_position_ids, + None, + return_dict=True, + ).pooler_output + + with _temporary_cpu_float32_modules(vision_modules): + return get_image_features( + pixel_values.float(), + pixel_position_ids, + None, + return_dict=True, + ).pooler_output + + def _compute_audio_features( + self, + input_features: torch.Tensor, + input_features_mask: torch.Tensor, + ) -> torch.Tensor: + if _gemma4_can_use_native_like_audio_features(self.multimodal_backbone): + return _gemma4_compute_native_like_audio_features( + self.multimodal_backbone, + input_features, + input_features_mask, + ) + get_audio_features = getattr(self.multimodal_backbone, "get_audio_features", None) + if not callable(get_audio_features): + raise TypeError("Gemma4 multimodal backbone is missing get_audio_features") + audio_modules = _gemma4_audio_modules(self.multimodal_backbone) + if ( + input_features.device.type == "cpu" + and input_features.dtype == torch.float16 + and len(audio_modules) == 2 + and not _torch_is_compiling() + and any(_module_floating_dtype(module) == torch.float16 for module in audio_modules) + ): + with _temporary_cpu_float32_modules(audio_modules): + audio_output = get_audio_features(input_features.float(), input_features_mask, return_dict=True) + return _gemma4_strip_audio_padding(audio_output) + + audio_output = get_audio_features(input_features, input_features_mask, return_dict=True) + return _gemma4_strip_audio_padding(audio_output) + + def _modules_to_prepare_for_capture(self) -> tuple[torch.nn.Module, ...]: + return () + + def prepare_for_capture(self, **_: object) -> None: + self.restore_after_capture() + modules = self._modules_to_prepare_for_capture() + if not modules: + return + promoted: list[tuple[torch.nn.Module, torch.dtype]] = [] + for module in modules: + device = _module_device(module) + dtype = _module_floating_dtype(module) + if device is None or device.type != "cpu" or dtype is None or dtype == torch.float32: + continue + promoted.append((module, dtype)) + module.to(dtype=torch.float32) + self._capture_modules = promoted + + def restore_after_capture(self, **_: object) -> None: + for module, dtype in reversed(self._capture_modules): + module.to(dtype=dtype) + self._capture_modules.clear() + + def _prepare_decoder_inputs( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + token_type_ids: torch.Tensor, + image_features: torch.Tensor, + audio_features: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + if self._native_merge_plan is not None: + audio_features_for_merge = audio_features + if audio_features_for_merge is None: + audio_features_for_merge = image_features.new_empty( + (int(input_ids.shape[0]), 0, int(image_features.shape[-1])) + ) + inputs_embeds, per_layer_inputs_tokens = _gemma4_apply_native_merge_plan( + self.model, + input_ids=input_ids, + image_features=image_features, + audio_features=audio_features_for_merge, + merge_plan=self._native_merge_plan, + pli_token_ids=self._native_merge_pli_token_ids, + ) + attention_mask = _gemma4_remap_sequence_tensor( + attention_mask, + merge_plan=self._native_merge_plan, + ) + token_type_ids = _gemma4_remap_sequence_tensor( + token_type_ids, + merge_plan=self._native_merge_plan, + ) + if token_type_ids is None: + raise RuntimeError("Gemma4 native merge requires token_type_ids for multimodal attention masking") + else: + get_placeholder_mask = getattr(self.multimodal_backbone, "get_placeholder_mask", None) + if not callable(get_placeholder_mask): + raise TypeError("Gemma4 multimodal backbone is missing get_placeholder_mask") + + inputs_embeds = self.model.get_input_embeddings()(input_ids) + text_mask, image_mask, audio_mask = _gemma4_get_placeholder_masks( + get_placeholder_mask, + token_type_ids=token_type_ids, + input_ids=input_ids, + inputs_embeds=inputs_embeds, + ) + inputs_embeds = inputs_embeds.masked_scatter( + image_mask.unsqueeze(-1).expand_as(inputs_embeds), + image_features.to(inputs_embeds.device, inputs_embeds.dtype), + ) + if audio_features is not None: + inputs_embeds = inputs_embeds.masked_scatter( + audio_mask.unsqueeze(-1).expand_as(inputs_embeds), + audio_features.to(inputs_embeds.device, inputs_embeds.dtype), + ) + per_layer_inputs_tokens = input_ids * text_mask.to(dtype=input_ids.dtype) + + per_layer_inputs: torch.Tensor + text_config = _gemma4_text_config(self.multimodal_backbone) + if getattr(text_config, "hidden_size_per_layer_input", None): + per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, per_layer_inputs_tokens, inputs_embeds) + per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) + else: + per_layer_inputs = inputs_embeds.new_empty((inputs_embeds.shape[0], inputs_embeds.shape[1], 0, 0)) + + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + return ( + inputs_embeds, + per_layer_inputs, + position_ids, + ) + + def _prepare_text_decoder_step_inputs( + self, + *, + input_ids: torch.Tensor, + position_ids: torch.LongTensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + embedding = self.model.get_input_embeddings() + if not isinstance(embedding, torch.nn.Module): + raise TypeError("Gemma4 model is missing input embeddings") + inputs_embeds = embedding(input_ids) + + model_config = getattr(self.model, "config", None) + text_config = getattr(model_config, "text_config", None) + hidden_scale = float(getattr(model_config, "hidden_size", 0) or 0) + if hidden_scale <= 0.0: + hidden_scale = float(getattr(text_config, "hidden_size", 0) or 0) + if hidden_scale <= 0.0: + hidden_scale = float(inputs_embeds.shape[-1]) + text_extra_scale = _gemma4_text_embedding_scale(embedding, hidden_scale ** 0.5) + if text_extra_scale != 1.0: + inputs_embeds = inputs_embeds * text_extra_scale + + text_config = _gemma4_text_config(self.multimodal_backbone) + if getattr(text_config, "hidden_size_per_layer_input", None): + per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) + per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) + else: + per_layer_inputs = inputs_embeds.new_empty((inputs_embeds.shape[0], inputs_embeds.shape[1], 0, 0)) + return inputs_embeds, per_layer_inputs, position_ids + + +class Gemma4VisionEncoderAdapter(_Gemma4MultimodalComponentBase): + def __init__( + self, + model: torch.nn.Module, + *, + weights_dir: str | None = None, + native_image_soft_token_counts: tuple[int, ...] | None = None, + native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, + ): + super().__init__( + model, + input_names=("pixel_values", "pixel_position_ids"), + weights_dir=weights_dir, + native_image_soft_token_counts=native_image_soft_token_counts, + native_image_pool_shapes=native_image_pool_shapes, + ) + + def _modules_to_prepare_for_capture(self) -> tuple[torch.nn.Module, ...]: + return () + + def forward(self, pixel_values: torch.Tensor, pixel_position_ids: torch.Tensor | None) -> torch.Tensor: + return self._compute_image_features(pixel_values, pixel_position_ids) + + def get_transpile_metadata(self): + return { + "graph": self._base_graph_meta( + adapter_type=type(self).__name__, + input_names=("pixel_values", "pixel_position_ids"), + ), + } + + +class Gemma4AudioEncoderAdapter(_Gemma4MultimodalComponentBase): + def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): + super().__init__(model, input_names=("input_features", "input_features_mask"), weights_dir=weights_dir) + + def _modules_to_prepare_for_capture(self) -> tuple[torch.nn.Module, ...]: + return () + + def forward(self, input_features: torch.Tensor, input_features_mask: torch.Tensor) -> torch.Tensor: + return self._compute_audio_features(input_features, input_features_mask) + + def get_transpile_metadata(self): + return { + "graph": self._base_graph_meta( + adapter_type=type(self).__name__, + input_names=("input_features", "input_features_mask"), + ), + } + + +class Gemma4LMEncoderAdapter(_Gemma4MultimodalComponentBase): + def __init__( + self, + model: torch.nn.Module, + *, + weights_dir: str | None = None, + native_merge_plan: _Gemma4NativeMergePlan | None = None, + include_audio: bool = True, + ): + self.include_audio = bool(include_audio) + input_names = ("input_ids", "attention_mask", "token_type_ids", "image_features") + if self.include_audio: + input_names = (*input_names, "audio_features") + super().__init__( + model, + input_names=input_names, + weights_dir=weights_dir, + native_merge_plan=native_merge_plan, + ) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + token_type_ids: torch.Tensor, + image_features: torch.Tensor, + audio_features: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.LongTensor]: + return self._prepare_decoder_inputs( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + image_features=image_features, + audio_features=audio_features, + ) + + def get_transpile_metadata(self): + return { + "graph": self._base_graph_meta( + adapter_type=type(self).__name__, + input_names=self.input_names, + ), + } + + +class Gemma4LMEncoderStepAdapter(_Gemma4MultimodalComponentBase): + def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): + super().__init__(model, input_names=("input_ids", "position_ids"), weights_dir=weights_dir) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.LongTensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + return self._prepare_text_decoder_step_inputs( + input_ids=input_ids, + position_ids=position_ids, + ) + + def get_transpile_metadata(self): + return { + "graph": self._base_graph_meta( + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + } + + +class Gemma4LMEncoderTextChunkAdapter(Gemma4LMEncoderStepAdapter): + pass + + +class Gemma4LMEncoderMediaStepAdapter(_Gemma4MultimodalComponentBase): + def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): + super().__init__( + model, + input_names=("inputs_embeds", "input_ids", "position_ids"), + weights_dir=weights_dir, + ) + + def forward( + self, + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + position_ids: torch.LongTensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + text_config = _gemma4_text_config(self.multimodal_backbone) + if getattr(text_config, "hidden_size_per_layer_input", None): + per_layer_inputs = _gemma4_get_per_layer_inputs(self.backbone, input_ids, inputs_embeds) + per_layer_inputs = self.backbone.project_per_layer_inputs(inputs_embeds, per_layer_inputs) + else: + per_layer_inputs = inputs_embeds.new_empty((inputs_embeds.shape[0], inputs_embeds.shape[1], 0, 0)) + return inputs_embeds, per_layer_inputs, position_ids + + def get_transpile_metadata(self): + return { + "graph": self._base_graph_meta( + adapter_type=type(self).__name__, + input_names=("inputs_embeds", "input_ids", "position_ids"), + ), + } + + +class Gemma4LMEncoderMediaChunkAdapter(Gemma4LMEncoderMediaStepAdapter): + pass + + +class Gemma4DecoderAdapter(_Gemma4MultimodalComponentBase): + def __init__(self, model: torch.nn.Module, *, weights_dir: str | None = None): + super().__init__(model, input_names=_GEMMA4_DECODER_PIPELINE_IO_KEYS, weights_dir=weights_dir) + + def forward( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor, + position_ids: torch.LongTensor, + ) -> torch.Tensor: + normalized_per_layer_inputs = per_layer_inputs + if normalized_per_layer_inputs.numel() == 0: + normalized_per_layer_inputs = None + attention_mask = torch.ones( + position_ids.shape, + dtype=torch.long, + device=position_ids.device, + ) + causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( + create_causal_mask=self._create_causal_mask, + create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + hidden_states = _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=normalized_per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + ) + logits = self.model.lm_head(hidden_states[:, -1:, :]) + return _gemma4_apply_final_logit_softcapping(self.model, logits) + + def get_transpile_metadata(self): + return { + "graph": self._base_graph_meta( + adapter_type=type(self).__name__, + input_names=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + ) + } + + +class Gemma4DecoderStepAdapter(Gemma4DecoderAdapter): + def forward( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor, + position_ids: torch.LongTensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + normalized_per_layer_inputs = per_layer_inputs + if normalized_per_layer_inputs.numel() == 0: + normalized_per_layer_inputs = None + attention_mask = torch.ones( + position_ids.shape, + dtype=torch.long, + device=position_ids.device, + ) + causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( + create_causal_mask=self._create_causal_mask, + create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + hidden_states, probe_hidden = _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=normalized_per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + capture_layer_index=28, + ) + logits = self.model.lm_head(hidden_states[:, -1:, :]) + logits = _gemma4_apply_final_logit_softcapping(self.model, logits) + return logits, probe_hidden[:, -1:, :] + + +class Gemma4DecoderPrefillChunkAdapter(Gemma4DecoderAdapter): + def forward( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor, + position_ids: torch.LongTensor, + ) -> torch.Tensor: + normalized_per_layer_inputs = per_layer_inputs + if normalized_per_layer_inputs.numel() == 0: + normalized_per_layer_inputs = None + attention_mask = torch.ones( + position_ids.shape, + dtype=torch.long, + device=position_ids.device, + ) + causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( + create_causal_mask=self._create_causal_mask, + create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + num_layers = int(getattr(self.backbone.config, "num_hidden_layers", len(self.backbone.layers))) + num_shared_layers = int(getattr(self.backbone.config, "num_kv_shared_layers", 0) or 0) + first_shared_layer = max(0, num_layers - num_shared_layers) + if first_shared_layer <= 0 or first_shared_layer >= num_layers or int(inputs_embeds.shape[1]) <= 1: + hidden_states = _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=normalized_per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + ) + else: + shared_kv_states: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + hidden_states = _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=normalized_per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + layer_end=first_shared_layer, + apply_norm=False, + shared_kv_states=shared_kv_states, + ) + tail_masks = { + key: value[:, :, -1:, :] if value is not None and value.ndim == 4 else value + for key, value in causal_mask_mapping.items() + } + tail_per_layer_inputs = ( + None + if normalized_per_layer_inputs is None + else normalized_per_layer_inputs[:, -1:, :, :] + ) + hidden_states = _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=hidden_states[:, -1:, :], + per_layer_inputs=tail_per_layer_inputs, + causal_mask_mapping=tail_masks, + position_ids=position_ids[:, -1:], + layer_start=first_shared_layer, + apply_norm=True, + shared_kv_states=shared_kv_states, + ) + logits = self.model.lm_head(hidden_states[:, -1:, :]) + return _gemma4_apply_final_logit_softcapping(self.model, logits) + + +class Gemma4DecoderEmbedChunkAdapter(Gemma4DecoderAdapter): + """Embedding-readout variant of the prefill chunk: runs the full backbone over + ALL tokens with final norm (no shared-KV-tail last-token shortcut) and returns + the full-sequence last hidden state (no lm_head).""" + + def forward( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor, + position_ids: torch.LongTensor, + ) -> torch.Tensor: + normalized_per_layer_inputs = per_layer_inputs + if normalized_per_layer_inputs.numel() == 0: + normalized_per_layer_inputs = None + attention_mask = torch.ones( + position_ids.shape, + dtype=torch.long, + device=position_ids.device, + ) + causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( + create_causal_mask=self._create_causal_mask, + create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + return _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=normalized_per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + ) + + +def _build_gemma3_causal_lm_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, + dynamic_batch: bool = True, + max_slots: int = 1, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + if input_ids is None: + return None + pad_token_id = _resolve_model_pad_token_id(model) + requested = tuple(components or ( + "decoder", + "lm_encoder_step", + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + "decoder_step", + )) + requested_set = set(requested) + common_graph_meta = { + "weights_dir": weights_dir, + "task": "causal_lm_logits", + "adapter_family": "gemma3", + } + metadata = {"family": "gemma3", "task": "causal_lm_logits"} + + specs: list[ComponentModuleSpec] = [] + if "decoder" in requested_set: + decoder = Gemma3CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=(input_ids,), + input_keys=("input_ids",), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata=metadata, + )) + + cached_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk", "decoder_step"} + if not (cached_components & requested_set): + return specs + + lm_encoder_step = Gemma3LMEncoderStepAdapter(model).eval() + lm_encoder_text_chunk = Gemma3LMEncoderStepAdapter(model).eval() + decoder_step = Gemma3EmbedsCausalLMStepAdapter(model).eval() + decoder_prefill_chunk = Gemma3EmbedsCausalLMPrefillChunkAdapter(model).eval() + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + cached_graph_meta = { + "use_internal_kv_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + } + prefill_chunk_size = max(2, int(os.environ.get("CACTUS_GEMMA3_PREFILL_CHUNK", "128") or "128")) + + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) + chunk_position_ids = torch.arange( + prefill_chunk_size, dtype=torch.int64, device=input_ids.device, + ).unsqueeze(0) + with torch.no_grad(): + step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) + chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + if "lm_encoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata=metadata, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + if "lm_encoder_text_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={ + **common_graph_meta, + "component": "lm_encoder_text_chunk", + "encoder_chunk_size": prefill_chunk_size, + }, + metadata=metadata, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill_chunk, + example_inputs=(chunk_embeds, chunk_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + **cached_graph_meta, + "component": "decoder_prefill_chunk", + "prefill_chunk_size": prefill_chunk_size, + }, + metadata=metadata, + )) + if "decoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=(step_embeds, step_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + **cached_graph_meta, + "component": "decoder_step", + "cache_num_slots": (max_slots if dynamic_batch else 1), + }, + metadata=metadata, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + return specs + + +def _build_gemma4_causal_lm_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, + dynamic_batch: bool = True, + max_slots: int = 1, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + if input_ids is None: + return None + pad_token_id = _resolve_model_pad_token_id(model) + requested_set = set(components or ()) + if not requested_set or "decoder" in requested_set: + requested_set |= { + "decoder", + "lm_encoder_step", + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + "decoder_step", + } + common_graph_meta = { + "weights_dir": weights_dir, + "task": "causal_lm_logits", + "adapter_family": "gemma4", + } + metadata = {"family": "gemma4", "task": "causal_lm_logits"} + + specs: list[ComponentModuleSpec] = [] + if "decoder" in requested_set: + decoder = Gemma4CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=(input_ids,), + input_keys=("input_ids",), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata=metadata, + )) + + cached_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk", "decoder_step"} + if not (cached_components & requested_set): + return specs + + lm_encoder_step = Gemma4LMEncoderStepAdapter(model, weights_dir=weights_dir).eval() + lm_encoder_text_chunk = Gemma4LMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_prefill_chunk = Gemma4DecoderPrefillChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_step = Gemma4DecoderStepAdapter(model, weights_dir=weights_dir).eval() + cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + prefill_chunk_size = max(2, int(os.environ.get("CACTUS_GEMMA4_PREFILL_CHUNK", "128") or "128")) + + step_input_ids = input_ids[:, -1:].contiguous() + step_position_ids = torch.full( + (int(input_ids.shape[0]), 1), + max(0, int(input_ids.shape[1]) - 1), + dtype=torch.long, + device=input_ids.device, + ) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) + chunk_position_ids = torch.arange( + prefill_chunk_size, dtype=torch.long, device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + with torch.no_grad(): + decoder_step_inputs = lm_encoder_step(step_input_ids, step_position_ids) + prefill_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + if "lm_encoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata=metadata, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + if "lm_encoder_text_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={ + **common_graph_meta, + "component": "lm_encoder_text_chunk", + "encoder_chunk_size": prefill_chunk_size, + }, + metadata=metadata, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill_chunk, + example_inputs=tuple(prefill_decoder_inputs), + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_prefill_chunk", + "use_internal_kv_cache": True, + "max_cache_seq_len": cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata=metadata, + )) + if "decoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=tuple(decoder_step_inputs), + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("logits", "probe_hidden"), + graph_meta={ + **common_graph_meta, + "component": "decoder_step", + "use_internal_kv_cache": True, + "max_cache_seq_len": cache_seq_len, + "cache_sink_size": 4, + "cache_num_slots": (max_slots if dynamic_batch else 1), + }, + metadata=metadata, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + return specs + + +def _gemma4_make_audio_mask_fp16_safe(model: torch.nn.Module) -> None: + backbone = getattr(model, "model", model) + audio_tower = getattr(backbone, "audio_tower", None) + if audio_tower is None: + return + seen: set[int] = set() + for mod in audio_tower.modules(): + cfg = getattr(mod, "config", None) + if cfg is None or id(cfg) in seen: + continue + seen.add(id(cfg)) + value = getattr(cfg, "attention_invalid_logits_value", None) + if isinstance(value, (int, float)) and value < -1e4: + cfg.attention_invalid_logits_value = -1e4 + + +def _build_gemma4_multimodal_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, + dynamic_batch: bool = True, + max_slots: int = 1, +) -> list[ComponentModuleSpec]: + pixel_values = named_tensors.get("pixel_values") + pixel_position_ids = named_tensors.get("pixel_position_ids") + input_features = named_tensors.get("input_features") + input_features_mask = named_tensors.get("input_features_mask") + input_ids = named_tensors["input_ids"] + attention_mask = named_tensors.get("attention_mask") + token_type_ids = named_tensors["token_type_ids"] + planning_input_ids = input_ids + low_memory_meta = _module_has_meta_tensors(model) + multimodal_backbone = getattr(model, "model", model) + model_has_audio = len(_gemma4_audio_modules(multimodal_backbone)) >= 2 + + has_vision_inputs = pixel_values is not None + has_audio_inputs = model_has_audio and input_features is not None and input_features_mask is not None + default_components = ["lm_encoder", "decoder"] + if has_vision_inputs: + default_components.insert(0, "vision_encoder") + if has_audio_inputs: + insert_at = 1 if has_vision_inputs else 0 + default_components.insert(insert_at, "audio_encoder") + requested_components = tuple(components or tuple(default_components)) + requested_set = set(requested_components) + if not requested_set: + return [] + include_audio = "audio_encoder" in requested_set or (components is None and has_audio_inputs) + + expanded_components: list[str] = [] + + def _require(component: str) -> None: + if component not in expanded_components: + expanded_components.append(component) + + if "decoder" in requested_set: + if has_vision_inputs or "vision_encoder" in requested_set: + _require("vision_encoder") + if include_audio: + _require("audio_encoder") + _require("lm_encoder") + _require("decoder_prefill_chunk") + _require("decoder_embed_chunk") + _require("lm_encoder_text_chunk") + _require("lm_encoder_media_chunk") + _require("lm_encoder_step") + _require("lm_encoder_media_step") + _require("decoder_step") + if "lm_encoder" in requested_set: + if has_vision_inputs or "vision_encoder" in requested_set: + _require("vision_encoder") + if include_audio: + _require("audio_encoder") + _require("lm_encoder") + if "decoder_prefill_chunk" in requested_set: + _require("decoder_prefill_chunk") + if "vision_encoder" in requested_set: + _require("vision_encoder") + if "audio_encoder" in requested_set: + _require("audio_encoder") + if "lm_encoder_step" in requested_set: + _require("lm_encoder_step") + if "lm_encoder_text_chunk" in requested_set: + _require("lm_encoder_text_chunk") + if "lm_encoder_media_step" in requested_set: + _require("lm_encoder_media_step") + if "lm_encoder_media_chunk" in requested_set: + _require("lm_encoder_media_chunk") + if "decoder_step" in requested_set: + _require("lm_encoder_step") + _require("decoder_step") + + if "vision_encoder" in expanded_components and pixel_values is None: + raise RuntimeError("Gemma4 vision component requested but pixel_values were not prepared") + if "audio_encoder" in expanded_components and (input_features is None or input_features_mask is None): + raise RuntimeError("Gemma4 audio component requested but input_features/input_features_mask were not prepared") + + vision_encoder: Gemma4VisionEncoderAdapter | None = None + audio_encoder: Gemma4AudioEncoderAdapter | None = None + native_image_soft_token_counts: tuple[int, ...] | None = None + native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None + if "vision_encoder" in expanded_components: + vision_tower = getattr(getattr(model, "model", model), "vision_tower", None) + pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3) + native_image_soft_token_counts = _gemma4_static_image_soft_token_counts( + pixel_position_ids, + pooling_kernel_size=pooling_kernel_size, + ) + native_image_pool_shapes = _gemma4_static_image_pool_shapes( + pixel_position_ids, + pooling_kernel_size=pooling_kernel_size, + ) + vision_encoder = Gemma4VisionEncoderAdapter( + model, + weights_dir=weights_dir, + native_image_soft_token_counts=native_image_soft_token_counts, + native_image_pool_shapes=native_image_pool_shapes, + ).eval() + if "audio_encoder" in expanded_components: + _gemma4_make_audio_mask_fp16_safe(model) + audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval() + + if low_memory_meta: + pixel_values = _to_meta_example_tensor(pixel_values) + pixel_position_ids = _to_meta_example_tensor(pixel_position_ids) + input_features = _to_meta_example_tensor(input_features) + input_features_mask = _to_meta_example_tensor(input_features_mask) + input_ids = _to_meta_example_tensor(input_ids) + attention_mask = _to_meta_example_tensor(attention_mask) + token_type_ids = _to_meta_example_tensor(token_type_ids) + + image_features: torch.Tensor | None = None + audio_features: torch.Tensor | None = None + decoder_inputs: tuple[torch.Tensor, ...] | None = None + prefill_decoder_inputs: tuple[torch.Tensor, ...] | None = None + decoder_step_inputs: tuple[torch.Tensor, ...] | None = None + step_input_ids: torch.Tensor | None = None + step_position_ids: torch.Tensor | None = None + media_step_embeds: torch.Tensor | None = None + native_merge_plan: _Gemma4NativeMergePlan | None = None + + with torch.no_grad(): + if "vision_encoder" in expanded_components and ( + "lm_encoder" in expanded_components + or "decoder" in expanded_components + ): + if vision_encoder is None or pixel_values is None: + raise RuntimeError("Gemma4 vision features requested but vision encoder inputs are unavailable") + image_features = vision_encoder(pixel_values, pixel_position_ids) + if "audio_encoder" in expanded_components and ( + "lm_encoder" in expanded_components + or "decoder" in expanded_components + ): + if audio_encoder is None or input_features is None or input_features_mask is None: + raise RuntimeError("Gemma4 audio features requested but audio encoder inputs are unavailable") + audio_features = audio_encoder(input_features, input_features_mask) + if "lm_encoder" in expanded_components or "decoder" in expanded_components: + if image_features is None and vision_encoder is not None and pixel_values is not None: + image_features = vision_encoder(pixel_values, pixel_position_ids) + if include_audio and audio_features is None and audio_encoder is not None and input_features is not None and input_features_mask is not None: + audio_features = audio_encoder(input_features, input_features_mask) + if image_features is not None: + native_merge_plan = _gemma4_build_native_merge_plan( + getattr(model, "model", model), + planning_input_ids, + image_feature_count=_gemma4_feature_token_count(image_features), + audio_feature_count=_gemma4_feature_token_count(audio_features), + ) + lm_encoder = Gemma4LMEncoderAdapter( + model, + weights_dir=weights_dir, + native_merge_plan=native_merge_plan, + include_audio=include_audio, + ).eval() + encoder_chunk_size = max(1, int(os.environ.get("CACTUS_GEMMA4_ENCODER_CHUNK", "128") or "128")) + encoder_chunk_size = min(encoder_chunk_size, int(input_ids.shape[1])) + lm_encoder_step = Gemma4LMEncoderStepAdapter(model, weights_dir=weights_dir).eval() + lm_encoder_text_chunk = Gemma4LMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() + lm_encoder_media_step = Gemma4LMEncoderMediaStepAdapter(model, weights_dir=weights_dir).eval() + lm_encoder_media_chunk = Gemma4LMEncoderMediaChunkAdapter(model, weights_dir=weights_dir).eval() + decoder = Gemma4DecoderAdapter(model, weights_dir=weights_dir).eval() + decoder_prefill = Gemma4DecoderPrefillChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_embed = Gemma4DecoderEmbedChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_step = Gemma4DecoderStepAdapter(model, weights_dir=weights_dir).eval() + + with torch.no_grad(): + if "decoder" in expanded_components: + if image_features is None: + raise RuntimeError("Gemma4 decoder spec requires precomputed image features") + if include_audio and audio_features is None: + raise RuntimeError("Gemma4 decoder spec requires precomputed audio features") + lm_encoder_inputs: tuple[torch.Tensor | None, ...] + if include_audio: + lm_encoder_inputs = ( + input_ids, + attention_mask, + token_type_ids, + image_features, + audio_features, + ) + else: + lm_encoder_inputs = ( + input_ids, + attention_mask, + token_type_ids, + image_features, + ) + decoder_inputs = lm_encoder(*lm_encoder_inputs) + if "lm_encoder_step" in expanded_components or "decoder_step" in expanded_components: + step_input_ids = input_ids[:, -1:].contiguous() + step_position_ids = torch.full( + (int(input_ids.shape[0]), 1), + max(0, int(input_ids.shape[1]) - 1), + dtype=torch.long, + device=input_ids.device, + ) + decoder_step_inputs = lm_encoder_step(step_input_ids, step_position_ids) + if "lm_encoder_media_step" in expanded_components: + if audio_features is not None and _gemma4_feature_token_count(audio_features) > 0: + media_source = audio_features if audio_features.ndim == 3 else audio_features.unsqueeze(0) + media_step_embeds = media_source[:, :1, :].contiguous() + elif image_features is not None and _gemma4_feature_token_count(image_features) > 0: + media_source = image_features if image_features.ndim == 3 else image_features.unsqueeze(0) + media_step_embeds = media_source[:, :1, :].contiguous() + else: + text_config = _gemma4_text_config(getattr(model, "model", model)) + hidden_size = int(getattr(text_config, "hidden_size", 0) or 0) + if hidden_size <= 0: + raise RuntimeError("Gemma4 lm_encoder_media_step spec could not infer hidden size") + dtype = _module_floating_dtype(model) or torch.float16 + media_step_embeds = torch.zeros( + (int(input_ids.shape[0]), 1, hidden_size), + dtype=dtype, + device=input_ids.device, + ) + chunk_input_ids = input_ids[:, :encoder_chunk_size].contiguous() + chunk_position_ids = torch.arange( + encoder_chunk_size, + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + if "decoder_prefill_chunk" in expanded_components and decoder_inputs is None: + prefill_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + if media_step_embeds is not None: + hidden_size = int(media_step_embeds.shape[-1]) + media_chunk_embeds = torch.zeros( + (int(input_ids.shape[0]), encoder_chunk_size, hidden_size), + dtype=media_step_embeds.dtype, + device=media_step_embeds.device, + ) + media_chunk_embeds[:, :1, :] = media_step_embeds + else: + text_config = _gemma4_text_config(getattr(model, "model", model)) + hidden_size = int(getattr(text_config, "hidden_size", 0) or 0) + if hidden_size <= 0: + hidden_size = int(chunk_input_ids.shape[-1]) + dtype = _module_floating_dtype(model) or torch.float16 + media_chunk_embeds = torch.zeros( + (int(input_ids.shape[0]), encoder_chunk_size, hidden_size), + dtype=dtype, + device=input_ids.device, + ) + + common_graph_meta = { + "weights_dir": weights_dir, + "task": "multimodal_causal_lm_logits", + "adapter_family": "gemma4", + } + specs: list[ComponentModuleSpec] = [] + if "vision_encoder" in expanded_components: + if vision_encoder is None or pixel_values is None: + raise RuntimeError("Gemma4 vision_encoder spec requires prepared pixel_values") + specs.append(ComponentModuleSpec( + component="vision_encoder", + module=vision_encoder, + example_inputs=(pixel_values, pixel_position_ids), + input_keys=("pixel_values", "pixel_position_ids"), + output_keys=("image_features",), + graph_meta={**common_graph_meta, "component": "vision_encoder"}, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "audio_encoder" in expanded_components: + if audio_encoder is None or input_features is None or input_features_mask is None: + raise RuntimeError("Gemma4 audio_encoder spec requires prepared audio features") + specs.append(ComponentModuleSpec( + component="audio_encoder", + module=audio_encoder, + example_inputs=(input_features, input_features_mask), + input_keys=("input_features", "input_features_mask"), + output_keys=("audio_features",), + graph_meta={**common_graph_meta, "component": "audio_encoder"}, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder" in expanded_components: + if image_features is None: + raise RuntimeError("Gemma4 lm_encoder spec requires precomputed image features") + if include_audio and audio_features is None: + raise RuntimeError("Gemma4 lm_encoder spec requires precomputed audio features") + lm_encoder_example_inputs: tuple[torch.Tensor | None, ...] + lm_encoder_input_keys = ("input_ids", "attention_mask", "token_type_ids", "image_features") + if include_audio: + lm_encoder_example_inputs = (input_ids, attention_mask, token_type_ids, image_features, audio_features) + lm_encoder_input_keys = (*lm_encoder_input_keys, "audio_features") + else: + lm_encoder_example_inputs = (input_ids, attention_mask, token_type_ids, image_features) + specs.append(ComponentModuleSpec( + component="lm_encoder", + module=lm_encoder, + example_inputs=lm_encoder_example_inputs, + input_keys=lm_encoder_input_keys, + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={**common_graph_meta, "component": "lm_encoder"}, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=256) + prefill_chunk_size = max(1, int(os.environ.get("CACTUS_GEMMA4_PREFILL_CHUNK", "128") or "128")) + prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) + if "decoder" in expanded_components: + if decoder_inputs is None: + raise RuntimeError("Gemma4 decoder spec requires precomputed decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=tuple(decoder_inputs), + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder_prefill_chunk" in expanded_components: + if decoder_inputs is not None: + chunk_inputs = tuple( + value[:, :prefill_chunk_size, ...].contiguous() + if value.ndim >= 2 + else value + for value in decoder_inputs + ) + elif prefill_decoder_inputs is not None: + prefill_chunk_size = int(prefill_decoder_inputs[0].shape[1]) + chunk_inputs = tuple(prefill_decoder_inputs) + else: + raise RuntimeError("Gemma4 decoder_prefill_chunk spec requires precomputed decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill, + example_inputs=chunk_inputs, + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_prefill_chunk", + "use_internal_kv_cache": True, + "max_cache_seq_len": cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder_embed_chunk" in expanded_components: + if decoder_inputs is not None: + embed_chunk_inputs = tuple( + value[:, :prefill_chunk_size, ...].contiguous() + if value.ndim >= 2 + else value + for value in decoder_inputs + ) + elif prefill_decoder_inputs is not None: + embed_chunk_inputs = tuple(prefill_decoder_inputs) + else: + raise RuntimeError("Gemma4 decoder_embed_chunk spec requires precomputed decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed, + example_inputs=embed_chunk_inputs, + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("last_hidden_state",), + graph_meta={ + **common_graph_meta, + "component": "decoder_embed_chunk", + "use_internal_kv_cache": True, + "max_cache_seq_len": cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder_step" in expanded_components: + if step_input_ids is None or step_position_ids is None: + raise RuntimeError("Gemma4 lm_encoder_step spec requires step token inputs") + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + if "lm_encoder_text_chunk" in expanded_components: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={ + **common_graph_meta, + "component": "lm_encoder_text_chunk", + "encoder_chunk_size": encoder_chunk_size, + }, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder_media_step" in expanded_components: + if media_step_embeds is None or step_input_ids is None or step_position_ids is None: + raise RuntimeError("Gemma4 lm_encoder_media_step spec requires step token inputs") + specs.append(ComponentModuleSpec( + component="lm_encoder_media_step", + module=lm_encoder_media_step, + example_inputs=(media_step_embeds, torch.zeros_like(step_input_ids), step_position_ids), + input_keys=("inputs_embeds", "input_ids", "position_ids"), + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={**common_graph_meta, "component": "lm_encoder_media_step"}, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder_media_chunk" in expanded_components: + specs.append(ComponentModuleSpec( + component="lm_encoder_media_chunk", + module=lm_encoder_media_chunk, + example_inputs=(media_chunk_embeds, torch.zeros_like(chunk_input_ids), chunk_position_ids), + input_keys=("inputs_embeds", "input_ids", "position_ids"), + output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + graph_meta={ + **common_graph_meta, + "component": "lm_encoder_media_chunk", + "encoder_chunk_size": encoder_chunk_size, + }, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder_step" in expanded_components: + if decoder_step_inputs is None: + raise RuntimeError("Gemma4 decoder_step spec requires precomputed step decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=tuple(decoder_step_inputs), + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("logits", "probe_hidden"), + graph_meta={ + **common_graph_meta, + "component": "decoder_step", + "use_internal_kv_cache": True, + "max_cache_seq_len": cache_seq_len, + "cache_sink_size": 4, + "cache_num_slots": (max_slots if dynamic_batch else 1), + }, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + return specs + + +class Qwen35CausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = _resolve_qwen35_text_backbone(model) + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward(self, input_ids: torch.Tensor): + return self.debug_forward(input_ids)[0] + + def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + attention_mask = ( + (input_ids != int(self.pad_token_id)).to(dtype=torch.int64) + if self.pad_token_id is not None + else None + ) + base_position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).view(1, 1, -1) + position_ids = torch.cat( + (base_position_ids, base_position_ids, base_position_ids, base_position_ids), + dim=0, + ) + text_position_ids = position_ids[0] + multimodal_position_ids = position_ids[1:] + + causal_mask = self._create_causal_mask( + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=None, + position_ids=text_position_ids, + ) + linear_attn_mask = attention_mask + + hidden_states = inputs_embeds + checkpoints: list[torch.Tensor] = [] + position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + + for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=_qwen35_attention_mask_for_layer( + self.backbone.config, + i, + causal_mask=causal_mask, + linear_attention_mask=linear_attn_mask, + ), + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + checkpoints.append(hidden_states) + + hidden_states = self.backbone.norm(hidden_states) + hidden_states = _select_last_non_pad_token( + hidden_states, + input_ids, + pad_token_id=self.pad_token_id, + ) + checkpoints.append(hidden_states) + return self.model.lm_head(hidden_states), checkpoints + + def get_transpile_metadata(self): + layer_types = _qwen35_layer_types(self.backbone.config) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("input_ids",), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": layer_types, + } + } + + +class Qwen35CausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = _resolve_qwen35_text_backbone(model) + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + inputs_embeds = self.backbone.embed_tokens(input_ids) + text_position_ids = position_ids.to(dtype=torch.int64) + base_position_ids = text_position_ids.view(1, text_position_ids.shape[0], -1) + multimodal_position_ids = torch.cat( + (base_position_ids, base_position_ids, base_position_ids), + dim=0, + ) + causal_mask = self._create_causal_mask( + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=None, + past_key_values=None, + position_ids=text_position_ids, + ) + linear_attn_mask = None + + hidden_states = inputs_embeds + position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=_qwen35_attention_mask_for_layer( + self.backbone.config, + i, + causal_mask=causal_mask, + linear_attention_mask=linear_attn_mask, + ), + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + + hidden_states = self.backbone.norm(hidden_states) + return self.model.lm_head(hidden_states[:, -1:, :]) + + def get_transpile_metadata(self): + layer_types = _qwen35_layer_types(self.backbone.config) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": layer_types, + } + } + + +def _resolve_qwen35_text_backbone(model: torch.nn.Module) -> torch.nn.Module: + backbone = getattr(model, "model", None) + language_model = getattr(backbone, "language_model", None) + if isinstance(language_model, torch.nn.Module): + return language_model + if isinstance(backbone, torch.nn.Module): + return backbone + raise AttributeError(f"{type(model).__name__} does not expose a Qwen3.5 text backbone") + + +def _resolve_qwen35_multimodal_backbone(model: torch.nn.Module) -> torch.nn.Module: + backbone = getattr(model, "model", None) + if isinstance(backbone, torch.nn.Module) and hasattr(backbone, "visual"): + return backbone + raise AttributeError(f"{type(model).__name__} does not expose a Qwen3.5 multimodal backbone") + + +def _qwen35_layer_types(config) -> tuple[str, ...]: + layer_types = tuple(str(value) for value in (getattr(config, "layer_types", ()) or ())) + if layer_types: + return layer_types + num_layers = int(getattr(config, "num_hidden_layers", 0) or 0) + return tuple("full_attention" for _ in range(num_layers)) + + +def _qwen35_attention_mask_for_layer( + config, + layer_index: int, + *, + causal_mask, + linear_attention_mask, +): + layer_types = _qwen35_layer_types(config) + layer_type = layer_types[layer_index] if layer_index < len(layer_types) else "full_attention" + return linear_attention_mask if layer_type == "linear_attention" else causal_mask + + +class Qwen35VisionEncoderAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, image_grid_thw: torch.Tensor): + super().__init__() + self.backbone = _resolve_qwen35_multimodal_backbone(model) + self.visual = self.backbone.visual + self.register_buffer("_image_grid_thw", image_grid_thw.detach().clone().to(dtype=torch.long), persistent=False) + with torch.no_grad(): + rotary_pos_emb = self.visual.rot_pos_emb(self._image_grid_thw) + emb = torch.cat((rotary_pos_emb.reshape(rotary_pos_emb.shape[0], -1),) * 2, dim=-1) + self.register_buffer("_vision_cos", emb.cos(), persistent=False) + self.register_buffer("_vision_sin", emb.sin(), persistent=False) + self.register_buffer( + "_vision_pos_embed", + self.visual.fast_pos_embed_interpolate(self._image_grid_thw), + persistent=False, + ) + + def _vision_attention(self, attn: torch.nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + attn.qkv(hidden_states).reshape(seq_length, 3, attn.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb_vision # type: ignore + + query_states, key_states = apply_rotary_pos_emb_vision( + query_states, + key_states, + self._vision_cos, + self._vision_sin, + ) + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * float(attn.scaling) + attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1).contiguous() + return attn.proj(attn_output) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + proj = self.visual.patch_embed.proj + patch_weight = proj.weight.reshape(int(proj.weight.shape[0]), -1) + hidden_states = torch.nn.functional.linear( + pixel_values.to(dtype=patch_weight.dtype), + patch_weight, + proj.bias, + ) + hidden_states = hidden_states + self._vision_pos_embed.to(dtype=hidden_states.dtype) + for block in self.visual.blocks: + hidden_states = hidden_states + self._vision_attention(block.attn, block.norm1(hidden_states)) + hidden_states = hidden_states + block.mlp(block.norm2(hidden_states)) + return self.visual.merger(hidden_states) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.backbone, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("pixel_values",), + ), + } + } + + +class Qwen35LMEncoderAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, input_ids: torch.Tensor): + super().__init__() + self.backbone = _resolve_qwen35_multimodal_backbone(model) + self.language_model = self.backbone.language_model + self.image_token_id = int(self.backbone.config.image_token_id) + image_positions = (input_ids[0] == self.image_token_id).nonzero(as_tuple=False).flatten() + if int(image_positions.numel()) <= 0: + raise ValueError("Qwen3.5 multimodal LM encoder requires image placeholder tokens") + self.image_token_start = int(image_positions[0].item()) + self.image_token_count = int(image_positions.numel()) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + image_features: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inputs_embeds = self.language_model.embed_tokens(input_ids) + image_embeds = image_features.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).reshape( + 1, + self.image_token_count, + int(inputs_embeds.shape[-1]), + ) + prefix = inputs_embeds[:, : self.image_token_start, :] + suffix = inputs_embeds[:, self.image_token_start + self.image_token_count :, :] + inputs_embeds = torch.cat((prefix, image_embeds, suffix), dim=1) + return inputs_embeds, attention_mask, position_ids + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.backbone, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("input_ids", "attention_mask", "position_ids", "image_features"), + ), + } + } + + +class Qwen35EmbedsCausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + self.backbone = _resolve_qwen35_text_backbone(model) + from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + if position_ids.ndim == 2: + multimodal_position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + text_position_ids = position_ids + elif position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + multimodal_position_ids = position_ids[1:] + else: + text_position_ids = None + multimodal_position_ids = position_ids + causal_mask = self._create_causal_mask( + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=None, + position_ids=text_position_ids, + ) + hidden_states = inputs_embeds + position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=_qwen35_attention_mask_for_layer( + self.backbone.config, + i, + causal_mask=causal_mask, + linear_attention_mask=attention_mask, + ), + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + hidden_states = self.backbone.norm(hidden_states) + return self.model.lm_head(hidden_states) + + def get_transpile_metadata(self): + layer_types = _qwen35_layer_types(self.backbone.config) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("inputs_embeds", "attention_mask", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": layer_types, + } + } + + +class Qwen35LMEncoderStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.backbone = _resolve_qwen35_text_backbone(model) + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + base_position_ids = position_ids.to(dtype=torch.int64).view(1, position_ids.shape[0], -1) + multimodal_position_ids = torch.cat( + (base_position_ids, base_position_ids, base_position_ids), + dim=0, + ) + return inputs_embeds, multimodal_position_ids + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.backbone, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + } + } + + +class Qwen35LMEncoderTextChunkAdapter(Qwen35LMEncoderStepAdapter): + pass + + +class Qwen35EmbedsCausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + self.backbone = _resolve_qwen35_text_backbone(model) + from transformers.models.qwen3_5.modeling_qwen3_5 import create_causal_mask # type: ignore + + self._create_causal_mask = create_causal_mask + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + if position_ids.ndim == 2: + text_position_ids = position_ids.to(dtype=torch.int64) + base_position_ids = text_position_ids.view(1, text_position_ids.shape[0], -1) + multimodal_position_ids = torch.cat( + (base_position_ids, base_position_ids, base_position_ids), + dim=0, + ) + else: + text_position_ids = None + multimodal_position_ids = position_ids.to(dtype=torch.int64) + causal_mask = self._create_causal_mask( + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=None, + past_key_values=None, + position_ids=text_position_ids, + ) + hidden_states = inputs_embeds + position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=_qwen35_attention_mask_for_layer( + self.backbone.config, + i, + causal_mask=causal_mask, + linear_attention_mask=None, + ), + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + hidden_states = self.backbone.norm(hidden_states) + return self.model.lm_head(hidden_states[:, -1:, :]) + + def get_transpile_metadata(self): + layer_types = _qwen35_layer_types(self.backbone.config) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="qwen3_5", + adapter_type=type(self).__name__, + input_names=("inputs_embeds", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": layer_types, + } + } + + +class Qwen35EmbedsCausalLMPrefillChunkAdapter(Qwen35EmbedsCausalLMStepAdapter): + pass + + +class Qwen3CausalLMLogitsAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = model.model + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) + + def forward(self, input_ids: torch.Tensor): + return self.debug_forward(input_ids)[0] + + def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + seq_len = int(inputs_embeds.shape[1]) + allowed_positions = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + ).view(1, 1, seq_len, seq_len) + if self.pad_token_id is not None: + key_mask = (input_ids != int(self.pad_token_id)).view(input_ids.shape[0], 1, 1, seq_len) + allowed_positions = torch.logical_and(allowed_positions, key_mask) + allowed_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * 0.0 + blocked_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * torch.finfo(inputs_embeds.dtype).min + causal_mask = torch.where( + allowed_positions, + allowed_values, + blocked_values, + ) + + hidden_states = inputs_embeds + checkpoints: list[torch.Tensor] = [] + position_embeddings = self.backbone.rotary_emb(hidden_states, position_ids) + + for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + ) + checkpoints.append(hidden_states) + + hidden_states = self.backbone.norm(hidden_states) + hidden_states = _select_last_non_pad_token( + hidden_states, + input_ids, + pad_token_id=self.pad_token_id, + ) + checkpoints.append(hidden_states) + return self.model.lm_head(hidden_states), checkpoints + + def get_transpile_metadata(self): + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family=self.adapter_family, + adapter_type=type(self).__name__, + input_names=("input_ids",), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + +class Qwen3CausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + super().__init__() + self.model = model + self.backbone = model.model + self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + inputs_embeds = self.backbone.embed_tokens(input_ids) + seq_len = int(inputs_embeds.shape[1]) + allowed_positions = torch.ones( + (1, 1, seq_len, seq_len), + dtype=torch.bool, + device=inputs_embeds.device, + ) + allowed_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * 0.0 + blocked_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * torch.finfo(inputs_embeds.dtype).min + causal_mask = torch.where(allowed_positions, allowed_values, blocked_values) + + hidden_states = inputs_embeds + text_position_ids = position_ids.to(dtype=torch.int64) + position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) + for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + + hidden_states = self.backbone.norm(hidden_states) + return self.model.lm_head(hidden_states[:, -1:, :]) + + def get_transpile_metadata(self): + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family=self.adapter_family, + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + +class Qwen2MoeCausalLMLogitsAdapter(Qwen3CausalLMLogitsAdapter): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + _patch_qwen2_moe_mlps(model) + super().__init__(model, pad_token_id=pad_token_id) + + +class Qwen2MoeCausalLMStepAdapter(Qwen3CausalLMStepAdapter): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + _patch_qwen2_moe_mlps(model) + super().__init__(model, pad_token_id=pad_token_id) + + +class Qwen3LMEncoderStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.backbone = model.model + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + return inputs_embeds, position_ids.to(dtype=torch.int64) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.backbone, + adapter_family="qwen3", + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + } + } + + +class Qwen3LMEncoderTextChunkAdapter(Qwen3LMEncoderStepAdapter): + pass + + +class Qwen3EmbedsCausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + self.backbone = model.model + + def _encode_hidden_states(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + seq_len = int(inputs_embeds.shape[1]) + allowed_positions = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + ).view(1, 1, seq_len, seq_len) + allowed_values = torch.zeros( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + blocked_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * torch.finfo(inputs_embeds.dtype).min + causal_mask = torch.where(allowed_positions, allowed_values, blocked_values) + + hidden_states = inputs_embeds + text_position_ids = position_ids.to(dtype=torch.int64) + position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) + for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + return self.backbone.norm(hidden_states) + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + hidden_states = self._encode_hidden_states(inputs_embeds, position_ids) + return self.model.lm_head(hidden_states[:, -1:, :]) + + def get_transpile_metadata(self): + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="qwen3", + adapter_type=type(self).__name__, + input_names=("inputs_embeds", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + +class Qwen3EmbedsCausalLMPrefillChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): + pass + + +class Qwen3EmbedsCausalLMEmbedChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): + """Like the prefill chunk, but emits the full-sequence last hidden state + (no lm_head, no last-token slice) for embedding readout.""" + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + return self._encode_hidden_states(inputs_embeds, position_ids) + + +def _build_qwen_causal_lm_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, + dynamic_batch: bool = True, + max_slots: int = 1, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + if input_ids is None: + return None + family = _family_key(model) + pad_token_id = _resolve_model_pad_token_id(model) + if family == "qwen3_5": + decoder = Qwen35CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + decoder_step = Qwen35CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() + elif family == "qwen2_moe": + decoder = Qwen2MoeCausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + decoder_step = Qwen2MoeCausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() + elif family == "qwen3": + decoder = Qwen3CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + decoder_step = Qwen3CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() + else: + return None + + requested = tuple(components or ("decoder", "decoder_step")) + requested_set = set(requested) + common_graph_meta = { + "weights_dir": weights_dir, + "task": "causal_lm_logits", + "adapter_family": family, + } + + chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_media_step", "decoder_prefill_chunk", "decoder_embed_chunk"} + wants_chunked = bool(chunk_components & requested_set) + if wants_chunked and family != "qwen3": + raise UnsupportedComponentsError( + f"text chunked-prefill components are only supported for the qwen3 family, got {family}" + ) + + specs: list[ComponentModuleSpec] = [] + if "decoder" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=(input_ids,), + input_keys=("input_ids",), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_step" in requested_set: + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_step", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "cache_num_slots": (max_slots if dynamic_batch else 1), + }, + metadata={"family": family, "task": "causal_lm_logits"}, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + + if wants_chunked: + lm_encoder_step = Qwen3LMEncoderStepAdapter(model).eval() + lm_encoder_text_chunk = Qwen3LMEncoderTextChunkAdapter(model).eval() + decoder_media_step = Qwen3EmbedsCausalLMStepAdapter(model).eval() + decoder_prefill_chunk = Qwen3EmbedsCausalLMPrefillChunkAdapter(model).eval() + decoder_embed_chunk = Qwen3EmbedsCausalLMEmbedChunkAdapter(model).eval() + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) + chunk_position_ids = torch.arange( + prefill_chunk_size, + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + with torch.no_grad(): + chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + with torch.no_grad(): + step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) + + if "lm_encoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata={"family": family, "task": "causal_lm_logits"}, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + if "lm_encoder_text_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_text_chunk"}, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill_chunk, + example_inputs=(chunk_embeds, chunk_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_prefill_chunk", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_embed_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed_chunk, + example_inputs=(chunk_embeds, chunk_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("last_hidden_state",), + graph_meta={ + **common_graph_meta, + "component": "decoder_embed_chunk", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_media_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_media_step", + module=decoder_media_step, + example_inputs=(step_embeds, step_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_media_step", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "cache_num_slots": (max_slots if dynamic_batch else 1), + }, + metadata={"family": family, "task": "causal_lm_logits"}, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + return specs + + +def _build_qwen3_5_multimodal_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + attention_mask = named_tensors.get("attention_mask") + position_ids = named_tensors.get("position_ids") + pixel_values = named_tensors.get("pixel_values") + image_grid_thw = named_tensors.get("image_grid_thw") + if not all(isinstance(value, torch.Tensor) for value in (input_ids, attention_mask, position_ids, pixel_values, image_grid_thw)): + return None + + requested = tuple(components or ( + "vision_encoder", + "lm_encoder", + "decoder", + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + "lm_encoder_step", + "decoder_media_step", + "decoder_step", + )) + requested_set = set(requested) + expanded_components: list[str] = [] + + def _require(component: str) -> None: + if component not in expanded_components: + expanded_components.append(component) + + if "decoder" in requested_set: + _require("vision_encoder") + _require("lm_encoder") + _require("decoder") + _require("lm_encoder_text_chunk") + _require("decoder_prefill_chunk") + if "lm_encoder" in requested_set: + _require("vision_encoder") + _require("lm_encoder") + if "lm_encoder_text_chunk" in requested_set: + _require("lm_encoder_text_chunk") + if "decoder_prefill_chunk" in requested_set: + _require("lm_encoder_text_chunk") + _require("decoder_prefill_chunk") + if "vision_encoder" in requested_set: + _require("vision_encoder") + if "decoder_step" in requested_set: + _require("decoder_step") + if "decoder_media_step" in requested_set: + _require("decoder_media_step") + if "lm_encoder_step" in requested_set or "decoder_media_step" in requested_set: + _require("lm_encoder_step") + + vision_encoder = Qwen35VisionEncoderAdapter(model, image_grid_thw=image_grid_thw).eval() + lm_encoder = Qwen35LMEncoderAdapter(model, input_ids=input_ids).eval() + decoder = Qwen35EmbedsCausalLMLogitsAdapter(model).eval() + lm_encoder_step = Qwen35LMEncoderStepAdapter(model).eval() + lm_encoder_text_chunk = Qwen35LMEncoderTextChunkAdapter(model).eval() + decoder_media_step = Qwen35EmbedsCausalLMStepAdapter(model).eval() + decoder_prefill_chunk = Qwen35EmbedsCausalLMPrefillChunkAdapter(model).eval() + decoder_step = Qwen35CausalLMStepAdapter(model, pad_token_id=_resolve_model_pad_token_id(model)).eval() + + image_features: torch.Tensor | None = None + decoder_inputs: tuple[torch.Tensor, ...] | None = None + prefill_decoder_inputs: tuple[torch.Tensor, ...] | None = None + with torch.no_grad(): + if "vision_encoder" in expanded_components and ("lm_encoder" in expanded_components or "decoder" in expanded_components): + image_features = vision_encoder(pixel_values) + if "lm_encoder" in expanded_components or "decoder" in expanded_components: + if image_features is None: + image_features = vision_encoder(pixel_values) + decoder_inputs = lm_encoder(input_ids, attention_mask, position_ids, image_features) + prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) + chunk_position_ids = torch.arange( + prefill_chunk_size, + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + if "decoder_prefill_chunk" in expanded_components: + prefill_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + common_graph_meta = { + "weights_dir": weights_dir, + "task": "multimodal_causal_lm_logits", + "adapter_family": "qwen3_5", + } + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + specs: list[ComponentModuleSpec] = [] + if "vision_encoder" in expanded_components: + specs.append(ComponentModuleSpec( + component="vision_encoder", + module=vision_encoder, + example_inputs=(pixel_values,), + input_keys=("pixel_values",), + output_keys=("image_features",), + graph_meta={**common_graph_meta, "component": "vision_encoder"}, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder" in expanded_components: + if image_features is None: + raise RuntimeError("Qwen3.5 lm_encoder spec requires precomputed image features") + specs.append(ComponentModuleSpec( + component="lm_encoder", + module=lm_encoder, + example_inputs=(input_ids, attention_mask, position_ids, image_features), + input_keys=("input_ids", "attention_mask", "position_ids", "image_features"), + output_keys=("inputs_embeds", "attention_mask", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder"}, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder" in expanded_components: + if decoder_inputs is None: + raise RuntimeError("Qwen3.5 decoder spec requires precomputed decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=decoder_inputs, + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder_step" in expanded_components: + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "lm_encoder_text_chunk" in expanded_components: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_text_chunk"}, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder_prefill_chunk" in expanded_components: + if prefill_decoder_inputs is None: + raise RuntimeError("Qwen3.5 decoder_prefill_chunk spec requires precomputed text chunk decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill_chunk, + example_inputs=prefill_decoder_inputs, + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_prefill_chunk", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder_media_step" in expanded_components: + if decoder_inputs is None: + raise RuntimeError("Qwen3.5 decoder_media_step spec requires precomputed decoder inputs") + step_inputs_embeds = decoder_inputs[0][:, :1, :] + step_position_ids = decoder_inputs[2][:, :, :1] + specs.append(ComponentModuleSpec( + component="decoder_media_step", + module=decoder_media_step, + example_inputs=(step_inputs_embeds, step_position_ids), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_media_step", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + }, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + if "decoder_step" in expanded_components: + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_step", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + }, + metadata={"family": "qwen3_5", "task": "multimodal_causal_lm_logits"}, + )) + return specs + + +class CTCLogitsAdapter(BoundInputAdapter): + def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], family: str): + super().__init__(model, input_names=input_names, family=family, metadata_task="ctc_logits") + + def forward(self, *bound_inputs: torch.Tensor | None) -> torch.Tensor: + outputs = self.model(return_dict=True, **self._kwargs_from_bound_inputs(*bound_inputs)) + return _extract_tensor_output(outputs, preferred_field="logits") + + +class EncoderHiddenStatesAdapter(BoundInputAdapter): + def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], family: str): + encoder = None + get_encoder = getattr(model, "get_encoder", None) + if callable(get_encoder): + encoder = get_encoder() + if encoder is None: + encoder = getattr(model, "encoder", None) + if encoder is None: + model_attr = getattr(model, "model", None) + if model_attr is not None: + encoder = getattr(model_attr, "encoder", None) + if not isinstance(encoder, torch.nn.Module): + raise NotImplementedError(f"{type(model).__name__} does not expose an encoder module") + super().__init__(model, input_names=input_names, family=family, metadata_task="encoder_hidden_states") + self.encoder = encoder + + def forward(self, *bound_inputs: torch.Tensor | None) -> torch.Tensor: + outputs = self.encoder(return_dict=True, **self._kwargs_from_bound_inputs(*bound_inputs)) + return _extract_tensor_output(outputs, preferred_field="last_hidden_state") + + +def _nomic_rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat([-x2, x1], dim=-1) + + +class NomicTextEmbeddingAdapter(torch.nn.Module): + """Export-friendly nomic-bert encoder that returns last_hidden_state. + + Reuses the HF submodules (embeddings, layernorms, attention/MLP weights) but + reimplements the forward to avoid the HF rotary cache and the data-dependent + MoE dispatch (`torch.where`/`index_add_`) which do not survive torch.export. + Pooling and L2 normalization are applied downstream in the C++ engine. + """ + + def __init__(self, model: torch.nn.Module) -> None: + super().__init__() + for name in ("embeddings", "emb_ln", "encoder"): + if not isinstance(getattr(model, name, None), torch.nn.Module): + raise NotImplementedError(f"{type(model).__name__} is not a nomic-bert model") + self.embeddings = model.embeddings + self.emb_ln = model.emb_ln + self.encoder = model.encoder + config = model.config + self.n_heads = int(config.n_head) + self.head_dim = int(config.n_embd) // int(config.n_head) + self.hidden_dim = int(config.n_embd) + self.n_experts = int(getattr(config, "num_experts", 0) or 0) + self.top_k = int(getattr(config, "moe_top_k", 0) or 0) + self.ffn_dim = int(config.n_inner) + self.rope_base = float(getattr(config, "rotary_emb_base", 10000.0) or 10000.0) + # Pre-transpose each MoE expert w2 to [hidden, n_experts*ffn] so the second + # expert matmul consumes it as a direct linear weight — transposing a packed + # quantized weight inside the graph is not lowerable. Idempotent by shape so + # repeated adapter construction on the same model is safe. + packed_rows = self.n_experts * self.ffn_dim + for layer in self.encoder.layers: + if getattr(layer, "moe", False): + w2 = layer.mlp.experts.mlp.w2 + if w2.shape[0] == packed_rows: + layer.mlp.experts.mlp.w2 = torch.nn.Parameter( + w2.detach().t().contiguous(), requires_grad=False + ) + + def _rope_cos_sin(self, seq_len: int, device, dtype): + inv_freq = 1.0 / (self.rope_base ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim)) + positions = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = positions.unsqueeze(1) * inv_freq.unsqueeze(0) + cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).to(dtype) + sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).to(dtype) + return cos, sin + + def _attention(self, layer, hidden_states, additive_mask, cos, sin): + batch, seq, _ = hidden_states.shape + qkv = layer.attn.Wqkv(hidden_states).view(batch, seq, 3, self.n_heads, self.head_dim) + query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + cos_b = cos[None, :, None, :] + sin_b = sin[None, :, None, :] + query = query * cos_b + _nomic_rotate_half(query) * sin_b + key = key * cos_b + _nomic_rotate_half(key) * sin_b + query = query.permute(0, 2, 1, 3) + key = key.permute(0, 2, 1, 3) + value = value.permute(0, 2, 1, 3) + scores = torch.matmul(query, key.transpose(-1, -2)) / (self.head_dim ** 0.5) + scores = scores + additive_mask + context = torch.matmul(F.softmax(scores, dim=-1), value) + context = context.permute(0, 2, 1, 3).reshape(batch, seq, self.hidden_dim) + return layer.attn.out_proj(context) + + def _topk_gate(self, weights): + # torch.topk does not lower; build the top-k softmax gate via iterative + # max + masking (top_k is small). gate[n, e] = softmax weight when expert e + # is among the top-k for token n, else 0. + neg = torch.finfo(weights.dtype).min + remaining = weights + gate = weights * 0.0 + for _ in range(self.top_k): + current_max = torch.amax(remaining, dim=-1, keepdim=True) + selected = (remaining >= current_max).to(weights.dtype) + gate = gate + selected * weights + remaining = remaining + selected * neg + return gate + + def _moe_mlp(self, moe, hidden_states): + # Dense MoE: run all experts via two fused matmuls against the packed + # expert weights (no per-expert indexing — slicing a quantized weight is + # not lowerable), then gate by scaling each expert's activations. The + # gated sum is identical to top-k routing because non-selected experts + # have gate 0. + batch, seq, _ = hidden_states.shape + flat = hidden_states.reshape(-1, self.hidden_dim) + weights = moe.router.layer(flat).softmax(dim=-1, dtype=torch.float32) + gate = self._topk_gate(weights).to(hidden_states.dtype) + all_hidden = F.linear(flat, moe.experts.mlp.w1) # [tokens, n_experts * ffn] + activated = F.gelu(all_hidden).reshape(-1, self.n_experts, self.ffn_dim) + activated = activated * gate.unsqueeze(-1) + activated = activated.reshape(-1, self.n_experts * self.ffn_dim) + # w2 is pre-transposed to [hidden, n_experts*ffn] in __init__. + out = F.linear(activated, moe.experts.mlp.w2) + out = out + moe.experts.bias + return out.reshape(batch, seq, self.hidden_dim) + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + hidden_states = self.emb_ln(self.embeddings.word_embeddings(input_ids)) + seq_len = input_ids.shape[1] + cos, sin = self._rope_cos_sin(seq_len, input_ids.device, hidden_states.dtype) + mask_float = attention_mask[:, None, None, :].to(hidden_states.dtype) + additive_mask = (mask_float - 1.0) * (-torch.finfo(hidden_states.dtype).min) + for layer in self.encoder.layers: + hidden_states = layer.norm1(self._attention(layer, hidden_states, additive_mask, cos, sin) + hidden_states) + mlp_out = self._moe_mlp(layer.mlp, hidden_states) if layer.moe else layer.mlp(hidden_states) + hidden_states = layer.norm2(mlp_out + hidden_states) + return hidden_states + + +class WhisperEncoderComponentAdapter(torch.nn.Module): + def __init__(self, encoder: torch.nn.Module): + super().__init__() + self.encoder = encoder + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + hidden_states = F.gelu(self.encoder.conv1(input_features)) + hidden_states = F.gelu(self.encoder.conv2(hidden_states)) + hidden_states = hidden_states.permute(0, 2, 1) + + position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + positions = self.encoder.embed_positions(position_ids) + hidden_states = hidden_states + positions.to( + device=hidden_states.device, + dtype=hidden_states.dtype, + ).unsqueeze(0) + + for layer in self.encoder.layers: + residual = hidden_states + hidden_states = layer.self_attn_layer_norm(hidden_states) + hidden_states, _ = layer.self_attn( + hidden_states=hidden_states, + attention_mask=None, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.final_layer_norm(hidden_states) + hidden_states = layer.activation_fn(layer.fc1(hidden_states)) + hidden_states = layer.fc2(hidden_states) + hidden_states = residual + hidden_states + + return self.encoder.layer_norm(hidden_states) + + +class WhisperDecoderCrossKVComponentAdapter(torch.nn.Module): + def __init__(self, decoder: torch.nn.Module): + super().__init__() + self.decoder = decoder + + def forward(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, ...]: + outputs: list[torch.Tensor] = [] + batch = int(encoder_hidden_states.shape[0]) + for layer in self.decoder.layers: + attn = layer.encoder_attn + kv_shape = (batch, -1, int(attn.num_heads), int(attn.head_dim)) + outputs.append(attn.k_proj(encoder_hidden_states).view(kv_shape).contiguous()) + outputs.append(attn.v_proj(encoder_hidden_states).view(kv_shape).contiguous()) + return tuple(outputs) + + +class WhisperDecoderStepWithCrossKVComponentAdapter(torch.nn.Module): + def __init__(self, decoder: torch.nn.Module, proj_out: torch.nn.Module, *, pad_token_id: int | None): + super().__init__() + self.decoder = decoder + self.proj_out = proj_out + self.pad_token_id = pad_token_id + + def forward( + self, + decoder_input_ids: torch.Tensor, + position_ids: torch.Tensor, + *cross_kv: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.decoder.embed_tokens(decoder_input_ids) + positions = self.decoder.embed_positions( + decoder_input_ids, + past_key_values_length=0, + position_ids=position_ids, + ) + hidden_states = hidden_states + positions.to( + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + for layer_index, layer in enumerate(self.decoder.layers): + residual = hidden_states + hidden_states = layer.self_attn_layer_norm(hidden_states) + hidden_states, _ = layer.self_attn( + hidden_states, + attention_mask=None, + past_key_values=None, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.encoder_attn_layer_norm(hidden_states) + attn = layer.encoder_attn + query_shape = (*hidden_states.shape[:-1], int(attn.num_heads), int(attn.head_dim)) + query_states = (attn.q_proj(hidden_states) * attn.scaling).view(query_shape).transpose(1, 2).contiguous() + key_states = cross_kv[layer_index * 2].transpose(1, 2).contiguous() + value_states = cross_kv[layer_index * 2 + 1].transpose(1, 2).contiguous() + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=None, + dropout_p=0.0, + scale=1.0, + ) + attn_output = attn_output.transpose(1, 2).reshape(*hidden_states.shape[:-1], -1).contiguous() + hidden_states = residual + attn.out_proj(attn_output) + + residual = hidden_states + hidden_states = layer.final_layer_norm(hidden_states) + hidden_states = layer.activation_fn(layer.fc1(hidden_states)) + hidden_states = layer.fc2(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = self.decoder.layer_norm(hidden_states) + hidden_states = _select_last_non_pad_token( + hidden_states, + decoder_input_ids, + pad_token_id=self.pad_token_id, + ) + return self.proj_out(hidden_states) + + +class NeedleSourceEncoderComponentAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + encode = getattr(self.model, "cactus_source_encode", None) + if not callable(encode): + raise TypeError(f"{type(self.model).__name__} does not expose cactus_source_encode") + return encode(input_ids, attention_mask) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="needle", + adapter_type=type(self).__name__, + input_names=("input_ids", "attention_mask"), + ), + "task": "causal_lm_logits", + } + } + + +class NeedleDecoderCrossKVComponentAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + + def forward( + self, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + cross_kv = getattr(self.model, "cactus_decoder_cross_kv", None) + if not callable(cross_kv): + raise TypeError(f"{type(self.model).__name__} does not expose cactus_decoder_cross_kv") + return cross_kv(encoder_hidden_states, encoder_attention_mask) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="needle", + adapter_type=type(self).__name__, + input_names=("encoder_hidden_states", "encoder_attention_mask"), + ), + "task": "causal_lm_logits", + } + } + + +class NeedleDecoderStepComponentAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + + def forward( + self, + decoder_input_ids: torch.Tensor, + position_ids: torch.Tensor, + encoder_attention_mask: torch.Tensor, + *cross_kv: torch.Tensor, + ) -> torch.Tensor: + step = getattr(self.model, "cactus_decoder_step", None) + if not callable(step): + raise TypeError(f"{type(self.model).__name__} does not expose cactus_decoder_step") + return step(decoder_input_ids, position_ids, encoder_attention_mask, *cross_kv) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="needle", + adapter_type=type(self).__name__, + input_names=("decoder_input_ids", "position_ids", "encoder_attention_mask"), + ), + "task": "causal_lm_logits", + } + } + + +def _family_key(model: torch.nn.Module) -> str: + explicit_family = getattr(model, "family", None) + if isinstance(explicit_family, str) and explicit_family: + return explicit_family.lower() + module_name = type(model).__module__ + if module_name.startswith("transformers.models.whisper."): + return "whisper" + if module_name.startswith("transformers.models.gemma4."): + return "gemma4" + if module_name.startswith("transformers.models.gemma3."): + return "gemma3" + if module_name.startswith("transformers.models.gemma."): + return "gemma" + if module_name.startswith("transformers.models.qwen3_5."): + return "qwen3_5" + if module_name.startswith("transformers.models.qwen3_vl."): + return "qwen3_5" + if module_name.startswith("transformers.models.qwen2_moe."): + return "qwen2_moe" + if module_name.startswith("transformers.models.qwen3."): + return "qwen3" + if module_name.startswith("transformers.models.lfm2_vl."): + return "lfm2_vl" + if module_name.startswith("transformers.models.lfm2_moe."): + return "lfm2_moe" + if module_name.startswith("transformers.models.lfm2."): + return "lfm2" + model_type = str(getattr(getattr(model, "config", None), "model_type", "") or "").lower() + if model_type == "needle": + return "needle" + if model_type == "qwen2_moe": + return "qwen2_moe" + if "nomic" in module_name.lower() or "nomic" in model_type: + return "nomic" + return "generic" + + +_ENCODER_CROSS_KV_ROUTE = "encoder_cross_kv_decoder_step" + + +def _cross_kv_output_keys(num_layers: int) -> tuple[str, ...]: + return tuple( + name + for layer_index in range(int(num_layers)) + for name in (f"cross_k_{layer_index}", f"cross_v_{layer_index}") + ) + + +def _component_spec( + *, + component: str, + module: torch.nn.Module, + example_inputs: tuple[torch.Tensor, ...], + input_keys: tuple[str, ...], + output_keys: tuple[str, ...], + common_graph_meta: dict[str, object], + family: str, + task: str, + graph_meta: dict[str, object] | None = None, + runtime_role: str | None = None, + source_kind: str | None = None, +) -> ComponentModuleSpec: + component_metadata: dict[str, object] = {"family": family, "task": task} + if runtime_role is not None: + component_metadata["runtime_route"] = _ENCODER_CROSS_KV_ROUTE + component_metadata["runtime_role"] = runtime_role + if source_kind is not None: + component_metadata["source_kind"] = source_kind + return ComponentModuleSpec( + component=component, + module=module, + example_inputs=example_inputs, + input_keys=input_keys, + output_keys=output_keys, + graph_meta={**common_graph_meta, "component": component, **(graph_meta or {})}, + metadata=component_metadata, + ) + + +def _encoder_cross_kv_step_specs( + *, + source_component: str, + source_module: torch.nn.Module, + source_inputs: tuple[torch.Tensor, ...], + source_input_keys: tuple[str, ...], + source_output_keys: tuple[str, ...], + source_kind: str, + cross_kv_module: torch.nn.Module, + cross_kv_inputs: tuple[torch.Tensor, ...], + cross_kv_input_keys: tuple[str, ...], + cross_kv_output_keys: tuple[str, ...], + decoder_step_module: torch.nn.Module, + decoder_step_inputs: tuple[torch.Tensor, ...], + decoder_step_input_keys: tuple[str, ...], + common_graph_meta: dict[str, object], + family: str, + task: str, + max_cache_seq_len: int, + decoder_step_graph_meta: dict[str, object] | None = None, +) -> list[ComponentModuleSpec]: + return [ + _component_spec( + component=source_component, + module=source_module, + example_inputs=source_inputs, + input_keys=source_input_keys, + output_keys=source_output_keys, + common_graph_meta=common_graph_meta, + family=family, + task=task, + runtime_role="source_encoder", + source_kind=source_kind, + ), + _component_spec( + component="decoder_cross_kv", + module=cross_kv_module, + example_inputs=cross_kv_inputs, + input_keys=cross_kv_input_keys, + output_keys=cross_kv_output_keys, + common_graph_meta=common_graph_meta, + family=family, + task=task, + runtime_role="decoder_cross_kv", + ), + _component_spec( + component="decoder_step", + module=decoder_step_module, + example_inputs=decoder_step_inputs, + input_keys=decoder_step_input_keys, + output_keys=("logits",), + common_graph_meta=common_graph_meta, + family=family, + task=task, + graph_meta={ + "use_internal_kv_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 0, + **(decoder_step_graph_meta or {}), + }, + runtime_role="decoder_step", + ), + ] + + +def _build_whisper_seq2seq_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + inputs_metadata: dict[str, object] | None = None, + weights_dir: str | None = None, +) -> list[ComponentModuleSpec]: + input_features = named_tensors.get("input_features") + if not isinstance(input_features, torch.Tensor): + raise RuntimeError("Whisper component transpile requires input_features") + + get_encoder = getattr(model, "get_encoder", None) + encoder = get_encoder() if callable(get_encoder) else None + if not isinstance(encoder, torch.nn.Module): + encoder = getattr(getattr(model, "model", None), "encoder", None) + decoder = getattr(getattr(model, "model", None), "decoder", None) + proj_out = getattr(model, "proj_out", None) + if not isinstance(encoder, torch.nn.Module) or not isinstance(decoder, torch.nn.Module): + raise RuntimeError(f"{type(model).__name__} does not expose Whisper encoder/decoder modules") + if not isinstance(proj_out, torch.nn.Module): + raise RuntimeError(f"{type(model).__name__} does not expose a Whisper projection head") + + metadata = dict(inputs_metadata or {}) + decoder_prompt_ids = metadata.get("decoder_input_ids") + if not isinstance(decoder_prompt_ids, list) or not decoder_prompt_ids: + decoder_start_token_id = getattr(getattr(model, "config", None), "decoder_start_token_id", None) + if not isinstance(decoder_start_token_id, int): + raise RuntimeError("Whisper component transpile requires decoder_input_ids metadata") + decoder_prompt_ids = [int(decoder_start_token_id)] + else: + decoder_prompt_ids = [int(value) for value in decoder_prompt_ids] + + target_token_count = int(metadata.get("target_token_count", len(decoder_prompt_ids)) or len(decoder_prompt_ids)) + target_token_count = max(target_token_count, len(decoder_prompt_ids)) + pad_token_id = int(metadata.get("pad_token_id", getattr(getattr(model, "config", None), "pad_token_id", 0)) or 0) + encoder_adapter = WhisperEncoderComponentAdapter(encoder).eval() + decoder_cross_kv_adapter = WhisperDecoderCrossKVComponentAdapter(decoder).eval() + decoder_step_adapter = WhisperDecoderStepWithCrossKVComponentAdapter( + decoder, + proj_out, + pad_token_id=pad_token_id, + ).eval() + + with torch.no_grad(): + encoder_hidden_states = encoder_adapter(input_features) + decoder_cross_kv = decoder_cross_kv_adapter(encoder_hidden_states) + + decoder_input_ids = torch.tensor([[decoder_prompt_ids[0]]], dtype=torch.int64, device=input_features.device) + step_position_ids = torch.zeros_like(decoder_input_ids, dtype=torch.int64) + + common_graph_meta = { + **_transpile_graph_meta( + model, + adapter_family="whisper", + adapter_type="component_pipeline", + input_names=("input_features",), + ), + "task": "seq2seq_transcription", + "adapter_family": "whisper", + } + if weights_dir: + common_graph_meta["weights_dir"] = weights_dir + cross_kv_output_keys = _cross_kv_output_keys(len(decoder.layers)) + return _encoder_cross_kv_step_specs( + source_component="audio_encoder", + source_module=encoder_adapter, + source_inputs=(input_features,), + source_input_keys=("input_features",), + source_output_keys=("encoder_hidden_states",), + source_kind="audio_features", + cross_kv_module=decoder_cross_kv_adapter, + cross_kv_inputs=(encoder_hidden_states,), + cross_kv_input_keys=("encoder_hidden_states",), + cross_kv_output_keys=cross_kv_output_keys, + decoder_step_module=decoder_step_adapter, + decoder_step_inputs=(decoder_input_ids, step_position_ids, *decoder_cross_kv), + decoder_step_input_keys=("decoder_input_ids", "position_ids", *cross_kv_output_keys), + common_graph_meta=common_graph_meta, + family="whisper", + task="seq2seq_transcription", + max_cache_seq_len=max(16, int(target_token_count)), + ) + + +def _resolve_decoder_start_token_id(model: torch.nn.Module, *, fallback: int = 0) -> int: + config = getattr(model, "config", None) + generation_config = getattr(model, "generation_config", None) + for owner in (config, generation_config): + for attr_name in ("decoder_start_token_id", "bos_token_id", "eos_token_id", "pad_token_id"): + value = getattr(owner, attr_name, None) if owner is not None else None + if value is not None: + return int(value) + return int(fallback) + +NEEDLE_CONTEXT_LENGTH = 1024 + +def _build_needle_causal_lm_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + if input_ids is None: + return None + for method_name in ("cactus_source_encode", "cactus_decoder_cross_kv", "cactus_decoder_step"): + if not callable(getattr(model, method_name, None)): + raise RuntimeError(f"{type(model).__name__} does not expose {method_name}") + + pad_token_id = _resolve_model_pad_token_id(model) + current_len = int(input_ids.shape[1]) + if current_len < NEEDLE_CONTEXT_LENGTH: + pad_value = int(pad_token_id) if pad_token_id is not None else 0 + padding = torch.full( + (int(input_ids.shape[0]), NEEDLE_CONTEXT_LENGTH - current_len), + pad_value, + dtype=input_ids.dtype, + device=input_ids.device, + ) + input_ids = torch.cat([input_ids, padding], dim=1) + if pad_token_id is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.int64) + else: + attention_mask = (input_ids != int(pad_token_id)).to(dtype=torch.int64) + source_encoder = NeedleSourceEncoderComponentAdapter(model).eval() + decoder_cross_kv = NeedleDecoderCrossKVComponentAdapter(model).eval() + decoder_step = NeedleDecoderStepComponentAdapter(model).eval() + + with torch.no_grad(): + encoder_hidden_states, encoder_attention_mask = source_encoder(input_ids, attention_mask) + cross_kv = decoder_cross_kv(encoder_hidden_states, encoder_attention_mask) + + decoder_start_token_id = _resolve_decoder_start_token_id(model, fallback=int(input_ids[0, 0].item())) + decoder_input_ids = torch.full( + (int(input_ids.shape[0]), 1), + decoder_start_token_id, + dtype=torch.int64, + device=input_ids.device, + ) + step_position_ids = torch.zeros_like(decoder_input_ids, dtype=torch.int64) + cross_kv_output_keys = _cross_kv_output_keys(len(cross_kv) // 2) + + common_graph_meta = { + **_transpile_graph_meta( + model, + adapter_family="needle", + adapter_type="component_pipeline", + input_names=("input_ids",), + ), + "task": "causal_lm_logits", + "adapter_family": "needle", + } + if weights_dir: + common_graph_meta["weights_dir"] = weights_dir + + return _encoder_cross_kv_step_specs( + source_component="source_encoder", + source_module=source_encoder, + source_inputs=(input_ids, attention_mask), + source_input_keys=("input_ids", "attention_mask"), + source_output_keys=("encoder_hidden_states", "encoder_attention_mask"), + source_kind="text_tokens", + cross_kv_module=decoder_cross_kv, + cross_kv_inputs=(encoder_hidden_states, encoder_attention_mask), + cross_kv_input_keys=("encoder_hidden_states", "encoder_attention_mask"), + cross_kv_output_keys=cross_kv_output_keys, + decoder_step_module=decoder_step, + decoder_step_inputs=(decoder_input_ids, step_position_ids, encoder_attention_mask, *cross_kv), + decoder_step_input_keys=("decoder_input_ids", "position_ids", "encoder_attention_mask", *cross_kv_output_keys), + common_graph_meta=common_graph_meta, + family="needle", + task="causal_lm_logits", + max_cache_seq_len=max(16, int(input_ids.shape[1])), + decoder_step_graph_meta={ + "external_cross_kv_cache_inputs": True, + "cross_kv_input_start_index": 3, + "cross_kv_input_count": len(cross_kv_output_keys), + "cross_kv_input_layout": "bthd", + }, + ) + + +def _build_nomic_text_embedding_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, +) -> list[ComponentModuleSpec]: + input_ids = named_tensors.get("input_ids") + if not isinstance(input_ids, torch.Tensor): + raise RuntimeError("nomic text_embedding transpile requires input_ids") + attention_mask = named_tensors.get("attention_mask") + if not isinstance(attention_mask, torch.Tensor): + attention_mask = torch.ones_like(input_ids) + + adapter = NomicTextEmbeddingAdapter(model).eval() + + common_graph_meta = { + **_transpile_graph_meta( + model, + adapter_family="nomic", + adapter_type="component_pipeline", + input_names=("input_ids", "attention_mask"), + ), + "task": "text_embedding", + "adapter_family": "nomic", + } + if weights_dir: + common_graph_meta["weights_dir"] = weights_dir + + return [ + ComponentModuleSpec( + component="text_embedding", + module=adapter, + example_inputs=(input_ids, attention_mask), + input_keys=("input_ids", "attention_mask"), + output_keys=("last_hidden_state",), + graph_meta={**common_graph_meta, "component": "text_embedding"}, + metadata={"family": "nomic", "task": "text_embedding"}, + ) + ] + + +def _build_lfm2_causal_lm_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + if input_ids is None: + return None + requested = tuple(components or ("decoder", "decoder_step")) + requested_set = set(requested) + family = _family_key(model) + common_graph_meta = { + "weights_dir": weights_dir, + "task": "causal_lm_logits", + "adapter_family": family, + } + metadata = {"family": family, "task": "causal_lm_logits"} + chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk"} + if chunk_components & requested_set: + return _lfm2_chunked_pipeline_specs( + model, + requested_set=requested_set, + input_ids=input_ids, + weights_dir=weights_dir, + cache_context_length=cache_context_length, + common_graph_meta=common_graph_meta, + metadata=metadata, + ) + + pad_token_id = _resolve_model_pad_token_id(model) + decoder = Lfm2CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + decoder_step = Lfm2CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() + specs: list[ComponentModuleSpec] = [] + if "decoder" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=(input_ids,), + input_keys=("input_ids",), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata=metadata, + )) + if "decoder_step" in requested_set: + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_step", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + }, + metadata=metadata, + )) + return specs + + +def _lfm2_chunked_pipeline_specs( + model: torch.nn.Module, + *, + requested_set: set[str], + input_ids: torch.Tensor, + weights_dir: str | None, + cache_context_length: str | int | None, + common_graph_meta: dict[str, object], + metadata: dict[str, str], + decoder_inputs_fn: Callable[[], tuple[torch.Tensor, ...]] | None = None, +) -> list[ComponentModuleSpec]: + lm_encoder_step = Lfm2VlLMEncoderStepAdapter(model, weights_dir=weights_dir).eval() + lm_encoder_text_chunk = Lfm2VlLMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_last_token = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=True).eval() + decoder_embed = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False, return_hidden=True).eval() + + prefill_chunk = max(2, int(os.environ.get("CACTUS_LFM2_PREFILL_CHUNK", "128") or "128")) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk) + chunk_position_ids = torch.arange( + prefill_chunk, dtype=torch.long, device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + + decoder_inputs: tuple[torch.Tensor, ...] | None = None + if {"decoder_prefill_chunk", "decoder_embed_chunk", "decoder_step"} & requested_set: + if decoder_inputs_fn is not None: + decoder_inputs = decoder_inputs_fn() + else: + with torch.no_grad(): + decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + cache_graph_meta = { + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "max_cache_seq_len": _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512), + "cache_sink_size": 4, + } + + specs: list[ComponentModuleSpec] = [] + if "lm_encoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "attention_mask", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata=metadata, + )) + if "lm_encoder_text_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "attention_mask", "position_ids"), + graph_meta={ + **common_graph_meta, + "component": "lm_encoder_text_chunk", + "encoder_chunk_size": prefill_chunk, + }, + metadata=metadata, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_last_token, + example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder_prefill_chunk", **cache_graph_meta}, + metadata=metadata, + )) + if "decoder_embed_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed, + example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("last_hidden_state",), + graph_meta={**common_graph_meta, "component": "decoder_embed_chunk", **cache_graph_meta}, + metadata=metadata, + )) + if "decoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_last_token, + example_inputs=tuple(tensor[:, :1, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder_step", **cache_graph_meta}, + metadata=metadata, + )) + return specs + + +def _build_lfm2_vl_multimodal_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, +) -> list[ComponentModuleSpec]: + input_ids = named_tensors["input_ids"] + attention_mask = named_tensors["attention_mask"] + pixel_values = named_tensors["pixel_values"] + spatial_shapes = named_tensors["spatial_shapes"] + pixel_attention_mask = named_tensors["pixel_attention_mask"] + text_context = max(8, int(os.environ.get("CACTUS_LFM2_TEXT_CONTEXT", "128") or "128")) + text_context = min(text_context, int(input_ids.shape[1])) + text_input_ids = torch.zeros( + (int(input_ids.shape[0]), text_context), + dtype=input_ids.dtype, + device=input_ids.device, + ) + text_attention_mask = torch.ones( + (int(attention_mask.shape[0]), text_context), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + requested_components = tuple(components or ("vision_encoder", "lm_encoder", "text_lm_encoder", "decoder")) + requested_set = set(requested_components) + if not requested_set: + return [] + + expanded_components: list[str] = [] + + def _require(component: str) -> None: + if component not in expanded_components: + expanded_components.append(component) + + if "decoder" in requested_set: + _require("vision_encoder") + _require("vision_projector") + _require("lm_encoder") + _require("text_lm_encoder") + _require("decoder") + _require("text_decoder") + _require("lm_encoder_step") + _require("lm_encoder_text_chunk") + _require("decoder_prefill_chunk") + _require("decoder_embed_chunk") + _require("decoder_step") + if "lm_encoder" in requested_set: + _require("vision_encoder") + _require("vision_projector") + _require("lm_encoder") + if "lm_encoder_step" in requested_set: + _require("lm_encoder_step") + if "lm_encoder_text_chunk" in requested_set: + _require("lm_encoder_text_chunk") + if "decoder_step" in requested_set: + _require("lm_encoder_step") + _require("decoder_step") + if "decoder_prefill_chunk" in requested_set: + _require("lm_encoder_text_chunk") + _require("decoder_prefill_chunk") + if "text_lm_encoder" in requested_set: + _require("text_lm_encoder") + if "text_decoder" in requested_set: + _require("text_lm_encoder") + _require("text_decoder") + if "vision_encoder" in requested_set: + _require("vision_encoder") + _require("vision_projector") + + vision_encoder = Lfm2VlVisionEncoderAdapter(model, weights_dir=weights_dir).eval() + vision_projector = Lfm2VlVisionProjectorAdapter(model, weights_dir=weights_dir).eval() + + _vemb = vision_encoder.embeddings + with torch.no_grad(): + _pos_grid = _vemb.position_embedding.weight.reshape( + _vemb.position_embedding_size, _vemb.position_embedding_size, -1 + ) + example_positional_embeddings = _vemb.resize_positional_embeddings( + _pos_grid, spatial_shapes.detach().cpu(), max_length=int(pixel_attention_mask.shape[1]) + ).to(pixel_values.dtype) + _root = _lfm2_vl_model_root(model) + _factor = int(_root.multi_modal_projector.factor) + _proj_in = int(_vemb.embed_dim) * _factor * _factor + _max_subimage_tokens = int(getattr(_root.config, "max_image_tokens", 256) or 256) + example_vision_features = torch.zeros((_max_subimage_tokens, _proj_in), dtype=pixel_values.dtype) + + lm_encoder = Lfm2VlLMEncoderAdapter(model, weights_dir=weights_dir).eval() + text_lm_encoder = Lfm2VlLMEncoderAdapter(model, weights_dir=weights_dir).eval() + decoder = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False).eval() + + decoder_inputs: tuple[torch.Tensor, ...] | None = None + text_decoder_inputs: tuple[torch.Tensor, ...] | None = None + with torch.no_grad(): + if "lm_encoder" in expanded_components or "decoder" in expanded_components: + decoder_inputs = lm_encoder(input_ids, attention_mask) + if "text_lm_encoder" in expanded_components: + text_decoder_inputs = text_lm_encoder(text_input_ids, text_attention_mask) + + def _ensure_decoder_inputs() -> tuple[torch.Tensor, ...]: + nonlocal decoder_inputs + if decoder_inputs is None: + decoder_inputs = lm_encoder(input_ids, attention_mask) + return decoder_inputs + + common_graph_meta = { + "weights_dir": weights_dir, + "task": "multimodal_causal_lm_logits", + "adapter_family": "lfm2_vl", + } + metadata = {"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"} + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + lm_seq = max(int(input_ids.shape[1]), int(max_cache_seq_len)) + lm_example_input_ids = torch.zeros((int(input_ids.shape[0]), lm_seq), dtype=input_ids.dtype, device=input_ids.device) + lm_example_attention_mask = torch.ones((int(attention_mask.shape[0]), lm_seq), dtype=attention_mask.dtype, device=attention_mask.device) + specs: list[ComponentModuleSpec] = [] + if "vision_encoder" in expanded_components: + specs.append(ComponentModuleSpec( + component="vision_encoder", + module=vision_encoder, + example_inputs=(pixel_values, pixel_attention_mask, example_positional_embeddings), + input_keys=("pixel_values", "pixel_attention_mask", "positional_embeddings"), + output_keys=("last_hidden_state",), + graph_meta={**common_graph_meta, "component": "vision_encoder"}, + metadata=metadata, + )) + if "vision_projector" in expanded_components: + specs.append(ComponentModuleSpec( + component="vision_projector", + module=vision_projector, + example_inputs=(example_vision_features,), + input_keys=("vision_features",), + output_keys=("image_features",), + graph_meta={**common_graph_meta, "component": "vision_projector"}, + metadata=metadata, + )) + if "lm_encoder" in expanded_components: + specs.append(ComponentModuleSpec( + component="lm_encoder", + module=lm_encoder, + example_inputs=(lm_example_input_ids, lm_example_attention_mask), + input_keys=("input_ids", "attention_mask"), + output_keys=("inputs_embeds", "attention_mask", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder"}, + metadata=metadata, + )) + if "text_lm_encoder" in expanded_components: + if text_decoder_inputs is None: + raise RuntimeError("LFM2-VL text_lm_encoder spec requires precomputed text decoder inputs") + specs.append(ComponentModuleSpec( + component="text_lm_encoder", + module=text_lm_encoder, + example_inputs=(text_input_ids, text_attention_mask), + input_keys=("input_ids", "attention_mask"), + output_keys=("inputs_embeds", "attention_mask", "position_ids"), + graph_meta={**common_graph_meta, "component": "text_lm_encoder"}, + metadata=metadata, + )) + if "decoder" in expanded_components: + if decoder_inputs is None: + decoder_inputs = text_decoder_inputs + if decoder_inputs is None: + raise RuntimeError("LFM2-VL decoder spec requires precomputed decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=decoder_inputs, + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata=metadata, + )) + if "text_decoder" in expanded_components: + if text_decoder_inputs is None: + raise RuntimeError("LFM2-VL text_decoder spec requires precomputed text decoder inputs") + specs.append(ComponentModuleSpec( + component="text_decoder", + module=decoder, + example_inputs=text_decoder_inputs, + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "text_decoder"}, + metadata=metadata, + )) + specs.extend(_lfm2_chunked_pipeline_specs( + model, + requested_set=set(expanded_components), + input_ids=input_ids, + weights_dir=weights_dir, + cache_context_length=cache_context_length, + common_graph_meta=common_graph_meta, + metadata=metadata, + decoder_inputs_fn=_ensure_decoder_inputs, + )) + return specs + + +def build_component_module_specs( + model: torch.nn.Module, + *, + task: str, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, + inputs_metadata: dict[str, object] | None = None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, + dynamic_batch: bool = True, + max_slots: int = 1, +) -> list[ComponentModuleSpec] | None: + family = _family_key(model) + if family == "qwen3_5" and task == "multimodal_causal_lm_logits": + return _build_qwen3_5_multimodal_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + ) + if family in {"qwen2_moe", "qwen3", "qwen3_5"} and task == "causal_lm_logits": + return _build_qwen_causal_lm_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + dynamic_batch=dynamic_batch, + max_slots=max_slots, + ) + if family == "gemma3" and task == "causal_lm_logits": + return _build_gemma3_causal_lm_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + dynamic_batch=dynamic_batch, + max_slots=max_slots, + ) + if family == "gemma4" and task == "causal_lm_logits": + return _build_gemma4_causal_lm_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + dynamic_batch=dynamic_batch, + max_slots=max_slots, + ) + if family == "gemma4" and task == "multimodal_causal_lm_logits": + return _build_gemma4_multimodal_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + dynamic_batch=dynamic_batch, + max_slots=max_slots, + ) + if family == "lfm2_vl" and task == "multimodal_causal_lm_logits": + return _build_lfm2_vl_multimodal_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + ) + if family in {"lfm2", "lfm2_moe"} and task == "causal_lm_logits": + return _build_lfm2_causal_lm_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + ) + if family == "needle" and task == "causal_lm_logits": + return _build_needle_causal_lm_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + ) + if family == "parakeet_tdt" and task == "tdt_transcription": + from cactus.transpile.tdt_runtime import build_parakeet_tdt_component_specs + + return build_parakeet_tdt_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + ) + if family == "whisper" and task == "seq2seq_transcription": + return _build_whisper_seq2seq_component_specs( + model, + named_tensors=named_tensors, + inputs_metadata=inputs_metadata, + weights_dir=weights_dir, + ) + if family == "nomic" and task == "text_embedding": + return _build_nomic_text_embedding_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + ) + return None + + +def canonicalize_model_interface( + model: torch.nn.Module, + task: str = "causal_lm_logits", + *, + input_names: tuple[str, ...] | None = None, + weights_dir: str | None = None, + inputs_metadata: dict[str, object] | None = None, +) -> CanonicalizedModel: + family = _family_key(model) + adapter_factory: Callable[[torch.nn.Module], torch.nn.Module] + resolved_input_names = tuple(input_names or ()) + padding_token_id_value = None + if isinstance(inputs_metadata, dict): + raw_padding_token_id = inputs_metadata.get("padding_token_id") + if isinstance(raw_padding_token_id, int): + padding_token_id_value = int(raw_padding_token_id) + + if task == "causal_lm_logits": + if family == "gemma": + adapter_factory = lambda inner_model: GemmaCausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + elif family == "gemma4": + adapter_factory = lambda inner_model: Gemma4CausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + elif family == "gemma3": + adapter_factory = lambda inner_model: Gemma3CausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + elif family == "qwen3_5": + adapter_factory = lambda inner_model: Qwen35CausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + elif family == "qwen2_moe": + adapter_factory = lambda inner_model: Qwen2MoeCausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + elif family == "qwen3": + adapter_factory = lambda inner_model: Qwen3CausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + elif family in {"lfm2", "lfm2_vl", "lfm2_moe"}: + adapter_factory = lambda inner_model: Lfm2CausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + else: + adapter_factory = lambda inner_model: CausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) + resolved_input_names = ("input_ids",) + elif task == "multimodal_causal_lm_logits": + if family not in {"gemma4", "lfm2_vl", "qwen3_5"}: + raise NotImplementedError(f"{type(model).__name__} does not support task={task}") + if family == "qwen3_5": + if not resolved_input_names: + resolved_input_names = ( + "input_ids", + "attention_mask", + "position_ids", + "pixel_values", + "image_grid_thw", + ) + adapter_factory = lambda inner_model: BoundInputAdapter( # type: ignore[assignment] + inner_model, + input_names=resolved_input_names, + family="qwen3_5", + metadata_task="multimodal_causal_lm_logits", + ) + elif family == "lfm2_vl": + if not resolved_input_names: + resolved_input_names = ( + "input_ids", + "attention_mask", + "pixel_values", + "spatial_shapes", + "pixel_attention_mask", + ) + adapter_factory = lambda inner_model: Lfm2VlMultimodalCausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + input_names=resolved_input_names, + ) + else: + if not resolved_input_names: + resolved_input_names = ( + "input_ids", + "attention_mask", + "token_type_ids", + "pixel_values", + "pixel_position_ids", + "input_features", + "input_features_mask", + ) + adapter_factory = lambda inner_model: Gemma4MultimodalCausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + input_names=resolved_input_names, + weights_dir=weights_dir, + ) + elif task == "ctc_logits": + if not resolved_input_names: + resolved_input_names = _infer_input_names( + model, + preferred=("input_values", "input_features", "attention_mask"), + ) + adapter_factory = lambda inner_model: CTCLogitsAdapter( # type: ignore[assignment] + inner_model, + input_names=resolved_input_names, + family=family, + ) + elif task == "encoder_hidden_states": + if not resolved_input_names: + resolved_input_names = _infer_input_names( + model, + preferred=("input_features", "input_values", "attention_mask"), + ) + adapter_factory = lambda inner_model: EncoderHiddenStatesAdapter( # type: ignore[assignment] + inner_model, + input_names=resolved_input_names, + family=family, + ) + elif task == "text_embedding": + if not resolved_input_names: + resolved_input_names = ("input_ids", "attention_mask") + adapter_factory = NomicTextEmbeddingAdapter + else: + raise NotImplementedError(f"unsupported task={task}") + + return CanonicalizedModel( + module=adapter_factory(model).eval(), + task=task, + family=family, + input_names=resolved_input_names, + ) diff --git a/python/cactus/transpile/model_patterns.py b/python/cactus/transpile/model_patterns.py new file mode 100644 index 000000000..fa226e9b0 --- /dev/null +++ b/python/cactus/transpile/model_patterns.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class GoldPattern: + name: str + semantic_ops: tuple[str, ...] + description: str + + +GOLD_PATTERNS: tuple[GoldPattern, ...] = ( + GoldPattern( + name="decoder_attention_gqa", + semantic_ops=("rms_norm", "rope", "attention", "attention_int8_hybrid"), + description=( + "Q/K/V projections from the same normalized hidden state; Q and K optionally get per-head RMSNorm, " + "Q/K are rotary-embedded, then Cactus attention consumes native Q-head and KV-head counts." + ), + ), + GoldPattern( + name="gated_mlp_gelu", + semantic_ops=("matmul", "gelu", "multiply"), + description="Two projection branches from the same normalized input: GELU(gate) * up, then a down projection.", + ), + GoldPattern( + name="gated_mlp_silu", + semantic_ops=("matmul", "silu", "multiply"), + description="Two projection branches from the same normalized input: SiLU(gate) * up, then a down projection.", + ), + GoldPattern( + name="decoder_block_post_attn_norm", + semantic_ops=("rms_norm", "attention", "add_clipped"), + description=( + "Pre-norm attention block with post-attention RMSNorm before the residual add, followed by a second " + "pre-FFN RMSNorm and a post-FFN RMSNorm before the final residual add." + ), + ), + GoldPattern( + name="decoder_block_simple_residual", + semantic_ops=("rms_norm", "attention", "add"), + description=( + "Pre-norm attention block with a direct residual add after attention, followed by a second RMSNorm and " + "a gated MLP branch with another direct residual add." + ), + ), + GoldPattern( + name="gemma4_partial_rope_attention", + semantic_ops=("rms_norm", "rope", "concat", "attention"), + description=( + "Gemma4 attention where only a prefix of the head dimension is rotary-embedded, optionally with shared " + "K/V heads and sliding-window attention." + ), + ), + GoldPattern( + name="sliding_window_attention_mask", + semantic_ops=("attention",), + description=( + "Exported PyTorch may build a boolean mask with diff/cumsum/comparisons, but the handwritten Cactus " + "models encode the same behavior directly as an attention `window_size`." + ), + ), +) diff --git a/python/cactus/transpile/model_profiles.py b/python/cactus/transpile/model_profiles.py new file mode 100644 index 000000000..7d132182b --- /dev/null +++ b/python/cactus/transpile/model_profiles.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +from dataclasses import dataclass +import re +from typing import Callable + +import torch + + +@dataclass(frozen=True) +class ModelProfile: + family: str + model_types: tuple[str, ...] = () + transformer_module: str | None = None + multimodal_context_tokens: int | None = None + input_combinations: tuple[tuple[str, ...], ...] = () + model_id_aliases: tuple[tuple[str, str], ...] = () + model_id_markers: tuple[str, ...] = () + family_aliases: tuple[str, ...] = () + stop_tokens: tuple[str, ...] = () + avoid_native_loader: bool = False + cached_step_components: tuple[str, ...] = () + cached_step_skip_components: tuple[str, ...] = () + dynamic_media_step_component: str | None = None + fp16_kv_cache_components: tuple[str, ...] = () + ir_replay_components: tuple[str, ...] = () + prompt_style: str = "auto" + multimodal_preprocessor: str = "auto" + text_only_component: str | None = None + default_task: str | None = None + default_components: tuple[str, ...] = () + text_component_plans: tuple[tuple[str, tuple[str, ...]], ...] = () + default_max_new_tokens: int | None = None + needs_image: bool = False + needs_audio: bool = False + force_component_pipeline: bool = False + aliases: tuple[tuple[str, str], ...] = () + regex_aliases: tuple[tuple[str, str], ...] = () + + +GEMMA4_PROFILE = ModelProfile( + family="gemma4", + model_types=("gemma4",), + transformer_module="transformers.models.gemma4.modeling_gemma4", + multimodal_context_tokens=2048, + input_combinations=((), ("image",), ("audio",), ("image", "audio")), + model_id_aliases=( + ("gemma4", "google/gemma-4-E2B-it"), + ("gemma4-e2b", "google/gemma-4-E2B-it"), + ), + model_id_markers=("gemma-4", "gemma4"), + family_aliases=("gemma",), + stop_tokens=("", ""), + avoid_native_loader=True, + cached_step_components=("lm_encoder_step", "decoder_step"), + cached_step_skip_components=("decoder",), + dynamic_media_step_component="lm_encoder_media_step", + fp16_kv_cache_components=("decoder_prefill_chunk", "decoder_step"), + ir_replay_components=( + "vision_encoder", + "audio_encoder", + "decoder_prefill_chunk", + "decoder_step", + ), + prompt_style="gemma4", + multimodal_preprocessor="gemma4", + default_task="multimodal_causal_lm_logits", + default_components=("vision_encoder", "audio_encoder", "lm_encoder", "decoder"), + needs_image=True, + needs_audio=True, + force_component_pipeline=True, +) + + +LFM2_PROFILE = ModelProfile( + family="lfm2", + model_types=("lfm2_vl", "lfm2", "lfm2_moe"), + multimodal_context_tokens=256, + input_combinations=((), ("image",)), + model_id_aliases=(("lfm", "LiquidAI/LFM2-VL-450M"),), + model_id_markers=("lfm2-vl", "lfm-vl"), + family_aliases=("lfm2_vl", "lfm"), + stop_tokens=("<|im_end|>",), + avoid_native_loader=True, + cached_step_components=("lm_encoder_step", "decoder_step"), + cached_step_skip_components=("decoder",), + fp16_kv_cache_components=("decoder_prefill_chunk", "decoder_step"), + ir_replay_components=("decoder_prefill_chunk", "decoder_step"), + prompt_style="lfm_chat", + multimodal_preprocessor="lfm2_vl", + text_only_component="text_lm_encoder", + default_task="multimodal_causal_lm_logits", + default_components=("vision_encoder", "lm_encoder", "decoder"), + text_component_plans=( + ("lfm2forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk")), + ("lfm2moeforcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk")), + ), + needs_image=True, + force_component_pipeline=True, +) + + +PARAKEET_TDT_PROFILE = ModelProfile( + family="parakeet_tdt", + model_types=("parakeet_tdt",), + model_id_aliases=( + ("parakeet", "nvidia/parakeet-tdt-0.6b-v3"), + ("parakeet-tdt", "nvidia/parakeet-tdt-0.6b-v3"), + ), + model_id_markers=("parakeet-tdt",), + avoid_native_loader=True, + default_task="tdt_transcription", + default_components=("audio_encoder", "decoder"), + needs_audio=True, + force_component_pipeline=True, + aliases=( + ("decoder.prediction.embed.weight", "decoder.embedding.weight"), + ("joint.enc.weight", "encoder_projector.weight"), + ("joint.enc.bias", "encoder_projector.bias"), + ("joint.pred.weight", "decoder.decoder_projector.weight"), + ("joint.pred.bias", "decoder.decoder_projector.bias"), + ("joint.joint_net.2.weight", "joint.head.weight"), + ("joint.joint_net.2.bias", "joint.head.bias"), + ("encoder.pre_encode.out.weight", "encoder.subsampling.linear.weight"), + ("encoder.pre_encode.out.bias", "encoder.subsampling.linear.bias"), + ), + regex_aliases=( + ( + r"encoder\.layers\.(\d+)\.self_attn\.(q|k|v)_proj\.weight", + r"encoder.layers.\1.self_attn.linear_\2.weight", + ), + ( + r"encoder\.layers\.(\d+)\.self_attn\.o_proj\.weight", + r"encoder.layers.\1.self_attn.linear_out.weight", + ), + ( + r"encoder\.layers\.(\d+)\.self_attn\.relative_k_proj\.weight", + r"encoder.layers.\1.self_attn.linear_pos.weight", + ), + ( + r"encoder\.layers\.(\d+)\.self_attn\.bias_u", + r"encoder.layers.\1.self_attn.pos_bias_u", + ), + ( + r"encoder\.layers\.(\d+)\.self_attn\.bias_v", + r"encoder.layers.\1.self_attn.pos_bias_v", + ), + ( + r"encoder\.layers\.(\d+)\.conv\.norm\.(.+)", + r"encoder.layers.\1.conv.batch_norm.\2", + ), + ), +) + + +WHISPER_PROFILE = ModelProfile( + family="whisper", + model_types=("whisper",), + model_id_aliases=(("whisper", "openai/whisper-small"),), + model_id_markers=("whisper",), + family_aliases=("whisperforconditionalgeneration",), + avoid_native_loader=True, + cached_step_components=("decoder_step",), + cached_step_skip_components=("decoder",), + fp16_kv_cache_components=("decoder_step",), + default_task="seq2seq_transcription", + default_components=("audio_encoder", "decoder"), + needs_audio=True, + force_component_pipeline=True, +) + + +QWEN_PROFILE = ModelProfile( + family="qwen", + model_types=("qwen", "qwen2", "qwen2_moe", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), + multimodal_context_tokens=512, + input_combinations=((), ("image",)), + model_id_aliases=( + ("qwen", "Qwen/Qwen3.5-0.8B"), + ("qwen3.5", "Qwen/Qwen3.5-0.8B"), + ("qwen35", "Qwen/Qwen3.5-0.8B"), + ("qwen3.5-0.8b", "Qwen/Qwen3.5-0.8B"), + ("qwen3", "Qwen/Qwen3-1.7B"), + ), + model_id_markers=("qwen",), + family_aliases=("qwen2", "qwen2_moe", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), + stop_tokens=("<|im_end|>",), + avoid_native_loader=True, + cached_step_components=("decoder_media_step",), + cached_step_skip_components=("decoder",), + fp16_kv_cache_components=("decoder_prefill_chunk", "decoder_media_step", "decoder_step"), + prompt_style="qwen_chat", + multimodal_preprocessor="qwen3_5", + default_task="multimodal_causal_lm_logits", + default_components=( + "vision_encoder", + "lm_encoder", + "decoder", + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + "lm_encoder_step", + "decoder_media_step", + "decoder_step", + ), + text_component_plans=( + ("qwen3forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk", "decoder_embed_chunk", "decoder_media_step")), + ), + needs_image=True, + force_component_pipeline=True, +) + + +NEEDLE_PROFILE = ModelProfile( + family="needle", + model_types=("needle",), + model_id_markers=("needle",), + avoid_native_loader=True, + default_task="causal_lm_logits", + default_components=("source_encoder", "decoder_cross_kv", "decoder_step"), + default_max_new_tokens=1022, + force_component_pipeline=True, +) + + +PROFILES: tuple[ModelProfile, ...] = ( + GEMMA4_PROFILE, + LFM2_PROFILE, + PARAKEET_TDT_PROFILE, + WHISPER_PROFILE, + QWEN_PROFILE, + NEEDLE_PROFILE, +) + + +def profile_for_model_type(model_type: str) -> ModelProfile | None: + normalized = str(model_type or "").strip().lower() + for profile in PROFILES: + if normalized in profile.model_types: + return profile + return None + + +def profile_for_family(family: str) -> ModelProfile | None: + normalized = str(family or "").strip().lower() + for profile in PROFILES: + if normalized == profile.family or normalized in profile.family_aliases: + return profile + return None + + +def profile_for_model_id(model_id: str) -> ModelProfile | None: + normalized = str(model_id or "").strip().lower() + if not normalized: + return None + alias_map = model_id_alias_map() + normalized = alias_map.get(normalized, normalized).lower() + for profile in PROFILES: + for _, profile_model_id in profile.model_id_aliases: + if normalized == profile_model_id.lower(): + return profile + if any(marker and marker in normalized for marker in profile.model_id_markers): + return profile + return None + + +def model_id_alias_map() -> dict[str, str]: + aliases: dict[str, str] = {} + for profile in PROFILES: + for alias, model_id in profile.model_id_aliases: + aliases[str(alias).strip().lower()] = str(model_id).strip() + return aliases + + +def multimodal_context_tokens_for_model_type(model_type: str, default: int) -> int: + profile = profile_for_model_type(model_type) + if profile is not None and profile.multimodal_context_tokens is not None: + return max(0, int(profile.multimodal_context_tokens)) + return max(0, int(default)) + + +def add_tensor_aliases( + state_dict: dict[str, torch.Tensor], + profile: ModelProfile, + *, + derived_aliases: Callable[[dict[str, torch.Tensor]], None] | None = None, +) -> dict[str, torch.Tensor]: + def alias(target: str, source: str) -> None: + if target not in state_dict and source in state_dict: + state_dict[target] = state_dict[source] + + for target, source in profile.aliases: + alias(target, source) + + for source_key in list(state_dict): + for source_pattern, target_template in profile.regex_aliases: + match = re.fullmatch(source_pattern, source_key) + if match is None: + continue + alias(match.expand(target_template), source_key) + + if derived_aliases is not None: + derived_aliases(state_dict) + + return state_dict diff --git a/python/cactus/transpile/multimodal_runtime.py b/python/cactus/transpile/multimodal_runtime.py new file mode 100644 index 000000000..3747afcf6 --- /dev/null +++ b/python/cactus/transpile/multimodal_runtime.py @@ -0,0 +1,722 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from cactus.transpile.audio_preprocess import audio_duration_limit_seconds +from cactus.transpile.audio_preprocess import prepare_native_gemma4_audio_features +from cactus.transpile.media_limits import resize_static_image +from cactus.transpile.runtime_support import patch_torch_flex_attention_compat +from cactus.transpile.runtime_support import patch_transformers_torchvision_probe +from cactus.transpile.runtime_support import PreparedInputs + + +def _resolve_audio_sample_rate(processor: object) -> int: + for attr_name in ("feature_extractor", "tokenizer"): + child = getattr(processor, attr_name, None) + sample_rate = getattr(child, "sampling_rate", None) + if isinstance(sample_rate, int) and sample_rate > 0: + return sample_rate + sample_rate = getattr(processor, "sampling_rate", None) + if isinstance(sample_rate, int) and sample_rate > 0: + return sample_rate + return 16000 + + +_GEMMA4_MULTIMODAL_INPUT_ORDER = ( + "input_ids", + "attention_mask", + "token_type_ids", + "pixel_values", + "pixel_position_ids", + "input_features", + "input_features_mask", +) + + +def _normalize_multimodal_prompt( + prompt: str, + *, + image_token: str | None, + num_images: int, + audio_token: str | None, + num_audio_segments: int, +) -> str: + normalized = prompt.strip() + prefixes: list[str] = [] + + if image_token and num_images > 0: + image_count = normalized.count(image_token) + if image_count < num_images: + prefixes.append(" ".join(image_token for _ in range(num_images - image_count))) + if audio_token and num_audio_segments > 0: + audio_count = normalized.count(audio_token) + if audio_count < num_audio_segments: + prefixes.append(" ".join(audio_token for _ in range(num_audio_segments - audio_count))) + + if prefixes: + prefix = "\n".join(part for part in prefixes if part) + if normalized: + return f"{prefix}\n{normalized}" + return prefix + return normalized + + +def _build_gemma4_chat_prompt( + *, + prompt: str, + image_token: str | None, + num_images: int, + audio_token: str | None, + num_audio_segments: int, + system_prompt: str = "", + enable_thinking_if_supported: bool = False, +) -> str: + result = "" + normalized_system = system_prompt.strip() + if enable_thinking_if_supported or normalized_system: + result += "<|turn>system\n" + if enable_thinking_if_supported: + result += "<|think|>" + result += normalized_system + result += "\n" + + result += "<|turn>user\n" + if image_token and num_images > 0: + for _ in range(num_images): + result += f"\n\n{image_token}\n\n" + result += prompt + if audio_token and num_audio_segments > 0: + result += "".join(audio_token for _ in range(num_audio_segments)) + result += "\n" + result += "<|turn>model\n" + return result + + +def _gemma4_split_cactus_newline_token_merges(batch: object) -> None: + input_ids = batch.get("input_ids") if hasattr(batch, "get") else None + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: + return + + expansions = { + 108: (107, 107), + 109: (107, 107, 107), + } + if not any(int(token) in expansions for token in input_ids.reshape(-1).tolist()): + return + + lengths: list[list[int]] = [] + max_len = 0 + for row in input_ids.detach().cpu().tolist(): + row_lengths = [len(expansions.get(int(token), (int(token),))) for token in row] + lengths.append(row_lengths) + max_len = max(max_len, sum(row_lengths)) + + for key in ("input_ids", "attention_mask", "token_type_ids"): + value = batch.get(key) if hasattr(batch, "get") else None + if not isinstance(value, torch.Tensor) or value.ndim != 2: + continue + expanded = torch.full( + (value.shape[0], max_len), + 0, + dtype=value.dtype, + device=value.device, + ) + for row_idx, row in enumerate(value.detach().cpu().tolist()): + out: list[int] = [] + for token_idx, item in enumerate(row): + if key == "input_ids": + out.extend(expansions.get(int(item), (int(item),))) + else: + out.extend([int(item)] * lengths[row_idx][token_idx]) + expanded[row_idx, : len(out)] = torch.tensor(out, dtype=value.dtype, device=value.device) + batch[key] = expanded + + +def _load_image_inputs(image_files: tuple[str, ...]) -> list[object]: + if not image_files: + return [] + try: + from PIL import Image # type: ignore + except Exception as exc: # pragma: no cover + raise RuntimeError(f"Pillow is required for --image-file: {exc}") from exc + + images: list[object] = [] + for image_file in image_files: + path = Path(image_file).resolve() + if not path.exists(): + raise RuntimeError(f"image_file does not exist: {path}") + with Image.open(path) as image: + images.append(resize_static_image(image.convert("RGB")).copy()) + return images + + +def _get_processor_image_attr(processor: object, name: str, default: object) -> object: + image_processor = getattr(processor, "image_processor", None) + if image_processor is not None and hasattr(image_processor, name): + return getattr(image_processor, name) + if isinstance(image_processor, dict) and name in image_processor: + return image_processor[name] + return default + + +def _image_channel_array(value: object, default: float) -> np.ndarray: + if isinstance(value, (list, tuple, np.ndarray)): + array = np.asarray(value, dtype=np.float32).reshape(-1) + if array.size >= 3: + return array[:3] + if array.size == 1: + return np.full((3,), float(array[0]), dtype=np.float32) + if isinstance(value, (int, float)): + return np.full((3,), float(value), dtype=np.float32) + return np.full((3,), float(default), dtype=np.float32) + + +def _prepare_gemma4_native_image_tensors( + processor: object, + image_files: tuple[str, ...], +) -> tuple[torch.Tensor, torch.Tensor] | None: + if not image_files: + return None + try: + from PIL import Image # type: ignore + except Exception as exc: # pragma: no cover + raise RuntimeError(f"Pillow is required for Gemma4 native image preprocessing: {exc}") from exc + + patch_size = int(_get_processor_image_attr(processor, "patch_size", 16)) + pooling_kernel_size = int(_get_processor_image_attr(processor, "pooling_kernel_size", 3)) + max_soft_tokens = int(_get_processor_image_attr(processor, "max_soft_tokens", 280)) + do_rescale = bool(_get_processor_image_attr(processor, "do_rescale", True)) + do_normalize = bool(_get_processor_image_attr(processor, "do_normalize", False)) + rescale_factor = float(_get_processor_image_attr(processor, "rescale_factor", 1.0 / 255.0)) if do_rescale else 1.0 + image_mean = ( + _image_channel_array(_get_processor_image_attr(processor, "image_mean", 0.0), 0.0) + if do_normalize + else np.zeros((3,), dtype=np.float32) + ) + image_std = ( + _image_channel_array(_get_processor_image_attr(processor, "image_std", 1.0), 1.0) + if do_normalize + else np.ones((3,), dtype=np.float32) + ) + max_patches = max_soft_tokens * pooling_kernel_size * pooling_kernel_size + side_multiple = pooling_kernel_size * patch_size + patch_dim = 3 * patch_size * patch_size + if patch_size <= 0 or pooling_kernel_size <= 0 or max_patches <= 0: + return None + + try: + resample_bilinear = Image.Resampling.BILINEAR + except AttributeError: # pragma: no cover + resample_bilinear = Image.BILINEAR + + pixel_batches: list[np.ndarray] = [] + position_batches: list[np.ndarray] = [] + for image_file in image_files: + path = Path(image_file).resolve() + with Image.open(path) as image: + rgb = image.convert("RGB") + width, height = rgb.size + target_pixels = float(max_patches * patch_size * patch_size) + factor = float(np.sqrt(target_pixels / max(1.0, float(width * height)))) + target_h = int(np.floor(factor * height / side_multiple)) * side_multiple + target_w = int(np.floor(factor * width / side_multiple)) * side_multiple + if target_h == 0: + target_h = side_multiple + if target_w == 0: + target_w = side_multiple + if (target_w, target_h) != rgb.size: + rgb = rgb.resize((target_w, target_h), resample=resample_bilinear) + array = np.asarray(rgb, dtype=np.float32) * rescale_factor + array = (array - image_mean.reshape(1, 1, 3)) / image_std.reshape(1, 1, 3) + + patch_h = target_h // patch_size + patch_w = target_w // patch_size + num_patches = patch_h * patch_w + if num_patches > max_patches: + raise RuntimeError( + f"Gemma4 native image preprocessing produced {num_patches} patches, " + f"but max_patches={max_patches}" + ) + chw = np.transpose(array, (2, 0, 1)) + patches = ( + chw.reshape(3, patch_h, patch_size, patch_w, patch_size) + .transpose(1, 3, 2, 4, 0) + .reshape(num_patches, patch_dim) + ) + + padded_patches = np.zeros((max_patches, patch_dim), dtype=np.float32) + padded_patches[:num_patches] = patches + positions = np.full((max_patches, 2), -1, dtype=np.int64) + valid_positions = np.zeros((num_patches, 2), dtype=np.int64) + for patch_y in range(patch_h): + row_start = patch_y * patch_w + valid_positions[row_start : row_start + patch_w, 0] = np.arange(patch_w, dtype=np.int64) + valid_positions[row_start : row_start + patch_w, 1] = patch_y + positions[:num_patches] = valid_positions + + pixel_batches.append(padded_patches) + position_batches.append(positions) + + return ( + torch.from_numpy(np.stack(pixel_batches, axis=0)), + torch.from_numpy(np.stack(position_batches, axis=0)), + ) + + +def _gemma4_image_soft_token_counts( + pixel_position_ids: torch.Tensor, + *, + pooling_kernel_size: int, +) -> list[int]: + counts: list[int] = [] + for positions in pixel_position_ids: + valid = positions[(positions != -1).any(dim=-1)] + if valid.numel() == 0: + counts.append(0) + continue + max_x = int(valid[:, 0].max().item()) + 1 + max_y = int(valid[:, 1].max().item()) + 1 + counts.append((max_x // pooling_kernel_size) * (max_y // pooling_kernel_size)) + return counts + + +def _gemma4_audio_soft_token_count(num_frames: int) -> int: + after_stage1 = (int(num_frames) + 1) // 2 + return (after_stage1 + 1) // 2 + + +def _resolve_gemma4_audio_mels(processor: object, default: int = 128) -> int: + for child_name in ("feature_extractor", "audio_processor", "audio_feature_extractor"): + child = getattr(processor, child_name, None) + if child is None: + continue + for attr_name in ("feature_size", "num_mel_bins", "n_mels", "input_feat_size"): + value = getattr(child, attr_name, None) + if isinstance(value, int) and value > 0: + return int(value) + for attr_name in ("feature_size", "num_mel_bins", "n_mels", "input_feat_size"): + value = getattr(processor, attr_name, None) + if isinstance(value, int) and value > 0: + return int(value) + return default + + +def _prepare_gemma4_processor_audio_tensors( + processor: object, + audio_waveforms: list[np.ndarray], + *, + torch_dtype: torch.dtype, + max_seconds: float | None, +) -> tuple[torch.Tensor, torch.Tensor, int] | None: + if not audio_waveforms: + return None + extractor = ( + getattr(processor, "audio_processor", None) + or getattr(processor, "feature_extractor", None) + or getattr(processor, "audio_feature_extractor", None) + ) + if not callable(extractor): + return None + batch = extractor(audio_waveforms, return_tensors="pt") + input_features = batch.get("input_features") if hasattr(batch, "get") else None + input_features_mask = batch.get("input_features_mask") if hasattr(batch, "get") else None + if not isinstance(input_features, torch.Tensor) or not isinstance(input_features_mask, torch.Tensor): + return None + if input_features.ndim != 3 or input_features_mask.ndim != 2: + return None + + active_frames = int(input_features_mask[0].sum().item()) + output_frames = int(input_features.shape[1]) + if max_seconds is not None and max_seconds > 0.0: + # Preserve the existing static Gemma4 graph capacity so a single bundle + # can accept any <=30s clip while the active mask still matches HF. + max_samples = int(round(16000.0 * float(max_seconds))) + output_frames = max(output_frames, int(np.ceil((max_samples + 640) / 160.0)) + 8) + + if output_frames > int(input_features.shape[1]): + padded_features = torch.zeros( + (input_features.shape[0], output_frames, input_features.shape[2]), + dtype=input_features.dtype, + ) + padded_mask = torch.zeros( + (input_features_mask.shape[0], output_frames), + dtype=input_features_mask.dtype, + ) + padded_features[:, : input_features.shape[1], :] = input_features + padded_mask[:, : input_features_mask.shape[1]] = input_features_mask + input_features = padded_features + input_features_mask = padded_mask + + return ( + input_features.to(dtype=torch_dtype), + input_features_mask.to(dtype=torch.bool), + active_frames, + ) + + +def _pad_gemma4_audio_batch_to_static_limit( + batch: object, + *, + max_seconds: float | None, +) -> int | None: + input_features = batch.get("input_features") if hasattr(batch, "get") else None + input_features_mask = batch.get("input_features_mask") if hasattr(batch, "get") else None + if not isinstance(input_features, torch.Tensor) or not isinstance(input_features_mask, torch.Tensor): + return None + active_frames = int(input_features_mask[0].sum().item()) + if max_seconds is None or max_seconds <= 0.0: + return active_frames + + max_samples = int(round(16000.0 * float(max_seconds))) + output_frames = max( + int(input_features.shape[1]), + int(np.ceil((max_samples + 640) / 160.0)) + 8, + ) + if output_frames > int(input_features.shape[1]): + padded_features = torch.zeros( + (input_features.shape[0], output_frames, input_features.shape[2]), + dtype=input_features.dtype, + device=input_features.device, + ) + padded_mask = torch.zeros( + (input_features_mask.shape[0], output_frames), + dtype=input_features_mask.dtype, + device=input_features_mask.device, + ) + padded_features[:, : input_features.shape[1], :] = input_features + padded_mask[:, : input_features_mask.shape[1]] = input_features_mask + batch["input_features"] = padded_features + batch["input_features_mask"] = padded_mask + return active_frames + + +def _build_gemma4_processor_chat_prompt( + processor: object, + *, + prompt: str, + image_files: tuple[str, ...], + audio_file: str | None, + system_prompt: str = "", +) -> str | None: + apply_chat_template = getattr(processor, "apply_chat_template", None) + if not callable(apply_chat_template): + return None + + content: list[dict[str, object]] = [] + for image_file in image_files: + content.append({"type": "image", "image": str(Path(image_file).expanduser().resolve())}) + if prompt.strip(): + content.append({"type": "text", "text": prompt.strip()}) + if audio_file: + content.append({"type": "audio", "audio": str(Path(audio_file).expanduser().resolve())}) + + messages: list[dict[str, object]] = [] + if system_prompt.strip(): + messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt.strip()}]}) + messages.append({"role": "user", "content": content or [{"type": "text", "text": prompt.strip()}]}) + try: + return str(apply_chat_template(messages, tokenize=False, add_generation_prompt=True)) + except Exception: + return None + + +def _build_gemma4_native_chat_prompt( + *, + prompt: str, + image_soft_token_counts: tuple[int, ...], + audio_soft_token_count: int, + system_prompt: str = "", + enable_thinking_if_supported: bool = False, +) -> str: + result = "" + normalized_system = system_prompt.strip() + if enable_thinking_if_supported or normalized_system: + result += "<|turn>system\n" + if enable_thinking_if_supported: + result += "<|think|>" + result += normalized_system + result += "\n" + + result += "<|turn>user\n" + for count in image_soft_token_counts: + if count > 0: + result += "\n\n<|image>" + result += "<|image|>" * int(count) + result += "\n\n" + result += prompt + if audio_soft_token_count > 0: + result += "<|audio>" + result += "<|audio|>" * int(audio_soft_token_count) + result += "" + result += "\n" + result += "<|turn>model\n" + return result + + +def _tokenize_gemma4_native_prompt( + processor: object, + prompt: str, + *, + image_token: str | None, + audio_token: str | None, +) -> dict[str, torch.Tensor]: + tokenizer = getattr(processor, "tokenizer", processor) + if not callable(tokenizer): + raise RuntimeError("Gemma4 native prompt path requires a tokenizer") + batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=False) + input_ids = batch["input_ids"] + if "attention_mask" not in batch: + batch["attention_mask"] = torch.ones_like(input_ids) + + token_type_ids = torch.zeros_like(input_ids) + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + if callable(convert): + if image_token: + image_id = convert(image_token) + if isinstance(image_id, int) and image_id >= 0: + token_type_ids = torch.where(input_ids == image_id, torch.ones_like(token_type_ids), token_type_ids) + if audio_token: + audio_id = convert(audio_token) + if isinstance(audio_id, int) and audio_id >= 0: + token_type_ids = torch.where(input_ids == audio_id, torch.full_like(token_type_ids, 2), token_type_ids) + batch["token_type_ids"] = token_type_ids + return {key: value for key, value in batch.items() if isinstance(value, torch.Tensor)} + + +def prepare_gemma4_multimodal_inputs( + processor: object | None, + *, + prompt: str, + image_files: tuple[str, ...], + audio_file: str | None, + torch_dtype: torch.dtype, + system_prompt: str = "", + enable_thinking_if_supported: bool = False, + use_gemma4_chat_template: bool = False, + static_image_soft_token_counts: tuple[int, ...] | None = None, + static_audio_soft_token_count: int | None = None, + force_static_multimodal_placeholders: bool = False, +) -> PreparedInputs: + if processor is None: + raise RuntimeError("multimodal Gemma4 transpile requires an AutoProcessor with image and audio support") + + images = _load_image_inputs(image_files) + audio_waveforms: list[np.ndarray] = [] + sample_rate: int | None = None + if audio_file: + sample_rate = _resolve_audio_sample_rate(processor) + from cactus.transpile.audio_preprocess import load_audio_waveform + + audio_waveforms.append( + load_audio_waveform( + audio_file, + target_sample_rate=sample_rate, + max_seconds=audio_duration_limit_seconds(), + ) + ) + + image_token = getattr(processor, "image_token", None) + audio_token = getattr(processor, "audio_token", None) + processor_prompt = _normalize_multimodal_prompt( + prompt, + image_token=image_token if isinstance(image_token, str) else None, + num_images=len(images), + audio_token=audio_token if isinstance(audio_token, str) else None, + num_audio_segments=len(audio_waveforms), + ) + normalized_prompt = processor_prompt + if use_gemma4_chat_template: + normalized_prompt = prompt.strip() + native_image_tensors = _prepare_gemma4_native_image_tensors(processor, image_files) + native_audio: torch.Tensor | None = None + native_audio_mask: torch.Tensor | None = None + native_audio_frames: int | None = None + if audio_file: + expected_mels = _resolve_gemma4_audio_mels(processor) + try: + native_audio, native_audio_mask, native_audio_frames = prepare_native_gemma4_audio_features( + audio_file, + expected_mels=expected_mels, + torch_dtype=torch_dtype, + max_seconds=audio_duration_limit_seconds(), + pad_to_max_seconds=True, + ) + except Exception as exc: + try: + processor_audio = _prepare_gemma4_processor_audio_tensors( + processor, + audio_waveforms, + torch_dtype=torch_dtype, + max_seconds=audio_duration_limit_seconds(), + ) + except Exception: + processor_audio = None + if processor_audio is None: + print(f"note=falling back to processor gemma4 audio prompt packing: {exc}") + native_audio = None + native_audio_mask = None + native_audio_frames = None + else: + native_audio, native_audio_mask, native_audio_frames = processor_audio + + has_native_media = ( + native_image_tensors is not None + or native_audio is not None + or ( + force_static_multimodal_placeholders + and ( + bool(static_image_soft_token_counts) + or (static_audio_soft_token_count is not None and int(static_audio_soft_token_count) > 0) + ) + ) + ) + if has_native_media and (not audio_file or native_audio is not None): + image_soft_counts: tuple[int, ...] = () + if native_image_tensors is not None: + pixel_values, pixel_position_ids = native_image_tensors + pooling_kernel_size = int(_get_processor_image_attr(processor, "pooling_kernel_size", 3)) + computed_image_counts = tuple( + _gemma4_image_soft_token_counts( + pixel_position_ids, + pooling_kernel_size=pooling_kernel_size, + ) + ) + if ( + static_image_soft_token_counts is not None + and len(static_image_soft_token_counts) == len(computed_image_counts) + ): + image_soft_counts = tuple(int(count) for count in static_image_soft_token_counts) + else: + image_soft_counts = computed_image_counts + elif force_static_multimodal_placeholders and static_image_soft_token_counts is not None: + image_soft_counts = tuple(int(count) for count in static_image_soft_token_counts) + audio_soft_count = 0 + if audio_file: + audio_soft_count = ( + int(static_audio_soft_token_count) + if static_audio_soft_token_count is not None and int(static_audio_soft_token_count) > 0 + else _gemma4_audio_soft_token_count(native_audio_frames or 0) + ) + elif ( + force_static_multimodal_placeholders + and static_audio_soft_token_count is not None + and int(static_audio_soft_token_count) > 0 + ): + audio_soft_count = int(static_audio_soft_token_count) + processor_prompt = _build_gemma4_native_chat_prompt( + prompt=normalized_prompt, + image_soft_token_counts=image_soft_counts, + audio_soft_token_count=audio_soft_count, + system_prompt=system_prompt, + enable_thinking_if_supported=enable_thinking_if_supported, + ) + batch = _tokenize_gemma4_native_prompt( + processor, + processor_prompt, + image_token=image_token if isinstance(image_token, str) else None, + audio_token=audio_token if isinstance(audio_token, str) else None, + ) + if native_image_tensors is not None: + batch["pixel_values"] = pixel_values + batch["pixel_position_ids"] = pixel_position_ids + if native_audio is not None and native_audio_mask is not None: + batch["input_features"] = native_audio + batch["input_features_mask"] = native_audio_mask + batch["native_audio_frames"] = int(native_audio_frames or native_audio.shape[1]) + else: + processor_chat_prompt = None + if image_files or audio_file: + processor_chat_prompt = _build_gemma4_processor_chat_prompt( + processor, + prompt=normalized_prompt, + image_files=image_files, + audio_file=audio_file, + system_prompt=system_prompt, + ) + processor_prompt = ( + processor_chat_prompt + if processor_chat_prompt is not None + else _build_gemma4_chat_prompt( + prompt=normalized_prompt, + image_token=image_token if isinstance(image_token, str) else None, + num_images=len(images), + audio_token=audio_token if isinstance(audio_token, str) else None, + num_audio_segments=len(audio_waveforms), + system_prompt=system_prompt, + enable_thinking_if_supported=enable_thinking_if_supported, + ) + ) + batch = processor( + text=processor_prompt, + images=images or None, + audio=audio_waveforms or None, + return_tensors="pt", + ) + if "mm_token_type_ids" in batch: + batch["token_type_ids"] = batch["mm_token_type_ids"] + if "image_position_ids" in batch: + batch["pixel_position_ids"] = batch["image_position_ids"] + native_audio_frames = _pad_gemma4_audio_batch_to_static_limit( + batch, + max_seconds=audio_duration_limit_seconds(), + ) + if isinstance(native_audio_frames, int): + batch["native_audio_frames"] = native_audio_frames + _gemma4_split_cactus_newline_token_merges(batch) + if native_image_tensors is not None: + batch["pixel_values"], batch["pixel_position_ids"] = native_image_tensors + else: + batch = processor( + text=processor_prompt, + images=images or None, + audio=audio_waveforms or None, + return_tensors="pt", + ) + + ordered_keys = [ + key + for key in _GEMMA4_MULTIMODAL_INPUT_ORDER + if isinstance(batch.get(key), torch.Tensor) + ] + if not ordered_keys: + raise RuntimeError("Gemma4 multimodal processor did not return any tensor inputs") + + tensors: list[torch.Tensor] = [] + for key in ordered_keys: + value = batch[key] + if not isinstance(value, torch.Tensor): + continue + if torch.is_floating_point(value): + value = value.to(dtype=torch_dtype) + tensors.append(value) + + metadata: dict[str, object] = { + "prompt": normalized_prompt, + "processor_prompt": processor_prompt, + "image_files": [str(Path(path).resolve()) for path in image_files], + "input_shapes": { + name: list(tensor.shape) + for name, tensor in zip(ordered_keys, tensors) + }, + } + if audio_file: + metadata["audio_file"] = str(Path(audio_file).resolve()) + if sample_rate is not None: + metadata["sample_rate"] = sample_rate + native_audio_frames = batch.get("native_audio_frames") + if isinstance(native_audio_frames, int): + metadata["native_audio_frames"] = native_audio_frames + + return PreparedInputs( + names=tuple(ordered_keys[: len(tensors)]), + tensors=tuple(tensors), + metadata=metadata, + ) + + +_patch_transformers_torchvision_probe = patch_transformers_torchvision_probe +_patch_torch_flex_attention_compat = patch_torch_flex_attention_compat +_prepare_gemma4_multimodal_inputs = prepare_gemma4_multimodal_inputs diff --git a/python/cactus/transpile/normalize.py b/python/cactus/transpile/normalize.py new file mode 100644 index 000000000..c7a0382a8 --- /dev/null +++ b/python/cactus/transpile/normalize.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from typing import Any + +import torch + +from cactus.transpile.aten_ops import canonical_ir_op + + +def normalize_target(target: Any) -> str: + return canonical_ir_op(target) + + +def dtype_to_ir(dtype: Any | None) -> str | None: + if dtype is None: + return None + + dtype_map = { + torch.float16: "fp16", + torch.float32: "fp32", + torch.float64: "fp64", + torch.bfloat16: "bf16", + torch.int8: "int8", + torch.int16: "int16", + torch.int32: "int32", + torch.int64: "int64", + torch.bool: "bool", + } + return dtype_map.get(dtype, str(dtype)) diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py new file mode 100644 index 000000000..5e3ee5436 --- /dev/null +++ b/python/cactus/transpile/ops.py @@ -0,0 +1,292 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class OpSchema: + name: str + num_inputs: int | None = None + attrs: tuple[str, ...] = () + backend_op: str | None = None + is_layout: bool = False + is_reduction: bool = False + is_fusible: bool = False + + +# Small canonical op registry for the transpiler. +# This is metadata for canonical IR ops, not another IR layer. + +OPS: tuple[OpSchema, ...] = ( + OpSchema("add", num_inputs=2, backend_op="add", is_fusible=True), + OpSchema("add_clipped", num_inputs=2, backend_op="add_clipped", is_fusible=True), + OpSchema("subtract", num_inputs=2, backend_op="subtract", is_fusible=True), + OpSchema("multiply", num_inputs=2, backend_op="multiply", is_fusible=True), + OpSchema("divide", num_inputs=2, backend_op="divide", is_fusible=True), + OpSchema("not_equal", num_inputs=2, backend_op="not_equal"), + OpSchema("equal", num_inputs=2), + OpSchema("greater", num_inputs=2), + OpSchema("greater_equal", num_inputs=2), + OpSchema("less", num_inputs=2), + OpSchema("less_equal", num_inputs=2), + OpSchema("logical_and", num_inputs=2), + OpSchema("logical_or", num_inputs=2), + OpSchema("logical_not", num_inputs=1), + OpSchema("abs", num_inputs=1, backend_op="abs", is_fusible=True), + OpSchema("clamp", num_inputs=1, attrs=("min", "max")), + OpSchema("pow", num_inputs=1, attrs=("exponent",), backend_op="pow"), + OpSchema("negate", num_inputs=1, is_fusible=True), + OpSchema("scalar_add", num_inputs=1, attrs=("value",), backend_op="scalar_add", is_fusible=True), + OpSchema("scalar_subtract", num_inputs=1, attrs=("value",), backend_op="scalar_subtract", is_fusible=True), + OpSchema("scalar_subtract_reverse", num_inputs=1, attrs=("value",)), + OpSchema("scalar_multiply", num_inputs=1, attrs=("value",), backend_op="scalar_multiply", is_fusible=True), + OpSchema("scalar_divide", num_inputs=1, attrs=("value",), backend_op="scalar_divide", is_fusible=True), + OpSchema("scalar_floor_divide", num_inputs=1, attrs=("value",), backend_op="scalar_floor_divide"), + OpSchema("scalar_divide_reverse", num_inputs=1, attrs=("value",)), + OpSchema("scalar_not_equal", num_inputs=1, attrs=("value",), backend_op="scalar_not_equal"), + OpSchema("scalar_equal", num_inputs=1, attrs=("value",)), + OpSchema("scalar_greater", num_inputs=1, attrs=("value",)), + OpSchema("scalar_greater_equal", num_inputs=1, attrs=("value",)), + OpSchema("scalar_less", num_inputs=1, attrs=("value",)), + OpSchema("scalar_less_equal", num_inputs=1, attrs=("value",)), + OpSchema("scalar_exp", num_inputs=1, backend_op="scalar_exp", is_fusible=True), + OpSchema("scalar_sqrt", num_inputs=1, backend_op="scalar_sqrt", is_fusible=True), + OpSchema("scalar_log", num_inputs=1, backend_op="scalar_log", is_fusible=True), + OpSchema("scalar_cos", num_inputs=1, backend_op="scalar_cos"), + OpSchema("scalar_sin", num_inputs=1, backend_op="scalar_sin"), + OpSchema("precision_cast", num_inputs=1, attrs=("dtype",), backend_op="precision_cast"), + OpSchema("view", num_inputs=1, attrs=("shape",), backend_op="reshape", is_layout=True), + OpSchema("permute", num_inputs=1, attrs=("permutation",), backend_op="permute", is_layout=True), + OpSchema("squeeze", num_inputs=1, attrs=("dim",), is_layout=True), + OpSchema("movedim", num_inputs=1, attrs=("source", "destination"), is_layout=True), + OpSchema("slice", num_inputs=1, attrs=("axis", "start", "end", "step"), backend_op="slice"), + OpSchema("split_with_sizes", num_inputs=1, attrs=("sizes", "axis")), + OpSchema("chunk", num_inputs=1, attrs=("chunks", "axis")), + OpSchema("ones", num_inputs=0, attrs=("shape", "dtype")), + OpSchema("one_hot", num_inputs=1, attrs=("num_classes",)), + OpSchema("tril", num_inputs=1, attrs=("diagonal",)), + OpSchema("unfold", num_inputs=1, attrs=("dimension", "size", "step")), + OpSchema("index", num_inputs=1, attrs=("axis", "index_value"), backend_op="index"), + OpSchema("advanced_index", num_inputs=None), + OpSchema( + "where", + num_inputs=None, + attrs=("true_is_scalar", "true_value", "false_is_scalar", "false_value"), + ), + OpSchema("masked_fill", num_inputs=2, attrs=("value",)), + OpSchema("masked_scatter", num_inputs=3), + OpSchema("getitem", num_inputs=1, attrs=("index",)), + OpSchema("cat", num_inputs=None, attrs=("axis",), backend_op="cat"), + OpSchema("matmul", num_inputs=2, backend_op="matmul"), + OpSchema("linear", num_inputs=None, attrs=("has_bias",), backend_op="matmul"), + OpSchema("addmm", num_inputs=3), + OpSchema("gather", num_inputs=2, attrs=("axis",), backend_op="gather"), + OpSchema("embedding", num_inputs=2, backend_op="embedding_from_tensor"), + OpSchema("sum", num_inputs=1, attrs=("axis", "keepdim"), backend_op="sum", is_reduction=True), + OpSchema("mean", num_inputs=1, attrs=("axis", "keepdim"), backend_op="mean", is_reduction=True), + OpSchema("variance", num_inputs=1, attrs=("axis", "keepdim"), backend_op="variance", is_reduction=True), + OpSchema("min", num_inputs=1, attrs=("axis", "keepdim"), backend_op="min", is_reduction=True), + OpSchema("max", num_inputs=1, attrs=("axis", "keepdim"), backend_op="max", is_reduction=True), + OpSchema("cumsum", num_inputs=1, attrs=("axis",), backend_op="cumsum"), + OpSchema("relu", num_inputs=1, backend_op="relu", is_fusible=True), + OpSchema("silu", num_inputs=1, backend_op="silu", is_fusible=True), + OpSchema("gelu", num_inputs=1, backend_op="gelu", is_fusible=True), + OpSchema("gelu_erf", num_inputs=1, backend_op="gelu_erf", is_fusible=True), + OpSchema("sigmoid", num_inputs=1, backend_op="sigmoid", is_fusible=True), + OpSchema("glu", num_inputs=1, attrs=("axis",), backend_op="glu"), + OpSchema("tanh", num_inputs=1, backend_op="tanh", is_fusible=True), + OpSchema("softplus", num_inputs=1), + OpSchema("softmax", num_inputs=1, attrs=("axis",), backend_op="softmax"), + OpSchema("layer_norm", num_inputs=None, attrs=("eps",), backend_op="layer_norm"), + OpSchema("group_norm", num_inputs=3, attrs=("num_groups", "eps"), backend_op="group_norm"), + OpSchema("batch_norm", num_inputs=5, attrs=("axis", "eps"), backend_op="batch_norm"), + OpSchema("conv1d", num_inputs=2, attrs=("stride", "padding", "dilation", "groups"), backend_op="conv1d"), + OpSchema("conv2d", num_inputs=2, attrs=("stride", "padding", "dilation", "groups"), backend_op="conv2d"), + OpSchema("pad", num_inputs=1, attrs=("pads", "value", "mode")), + OpSchema("lstm_cell", num_inputs=7, backend_op="lstm_cell"), + + # CUSTOM / FUSED OPS + OpSchema("dense_mlp_tq_fused", num_inputs=4, attrs=("product_scale",), backend_op="dense_mlp_tq_fused"), + OpSchema( + "lfm2_moe_layer_gated", + num_inputs=None, + attrs=( + "num_experts", + "num_experts_per_tok", + "use_expert_bias", + "normalize_routing", + "epsilon", + "routed_scaling_factor", + "activation", + ), + ), + OpSchema( + "qwen2_moe_layer_gated", + num_inputs=None, + attrs=( + "num_experts", + "num_experts_per_tok", + "normalize_routing", + "epsilon", + "routed_scaling_factor", + "activation", + ), + ), + OpSchema( + "gemma4_moe_layer_gated", + num_inputs=None, + attrs=( + "num_experts", + "num_experts_per_tok", + "normalize_routing", + "epsilon", + "routed_scaling_factor", + "activation", + ), + ), + OpSchema("rms_norm", num_inputs=2, attrs=("eps",), backend_op="rms_norm"), + OpSchema("rope", num_inputs=1, attrs=("theta", "position_offset"), backend_op="rope"), + OpSchema("rel_pos_bias", num_inputs=2, attrs=("scale",), backend_op="rel_pos_bias"), + OpSchema( + "attention", + num_inputs=3, + attrs=( + "scale", + "is_causal", + "window_size", + "mask", + "additive_mask", + "qkv_layout", + "q_layout", + "k_layout", + "v_layout", + "output_layout", + ), + backend_op="attention", + ), + OpSchema( + "attention_block", + num_inputs=None, + attrs=( + "scale", + "is_causal", + "window_size", + "has_mask", + "additive_mask", + "has_gate", + "has_bias", + "attention_output_shape", + "qkv_layout", + ), + is_fusible=True, + ), + OpSchema( + "self_attention_block", + num_inputs=None, + attrs=( + "scale", + "is_causal", + "window_size", + "has_mask", + "additive_mask", + "has_gate", + "has_bias", + "has_query_projection_bias", + "has_query_add", + "has_rel_query_add", + "has_key_projection_bias", + "has_value_projection_bias", + "has_rel_pos_bias", + "has_relative_key_projection_bias", + "query_shape", + "key_shape", + "value_shape", + "relative_key_shape", + "rel_pos_scale", + "attention_output_shape", + ), + is_fusible=True, + ), + OpSchema( + "conv_module", + num_inputs=None, + attrs=( + "eps", + "has_pointwise1_bias", + "has_depthwise_bias", + "has_pointwise2_bias", + "depthwise_kernel_size", + "depthwise_padding", + ), + is_fusible=True, + ), + OpSchema( + "gated_deltanet_prefill", + num_inputs=None, + attrs=( + "num_k_heads", + "num_v_heads", + "key_dim", + "value_dim", + "eps", + "chunk_size", + "has_z", + "has_dt_bias", + "has_a_log", + "has_conv", + ), + is_fusible=True, + ), + OpSchema( + "gated_deltanet_decode", + num_inputs=None, + attrs=( + "num_k_heads", + "num_v_heads", + "key_dim", + "value_dim", + "eps", + "has_z", + "has_dt_bias", + "has_a_log", + "has_conv", + ), + is_fusible=True, + ), +) + + +ALL_OPS: dict[str, OpSchema] = {op.name: op for op in OPS} + + +# Import/capture spellings reduced into the canonical registry above. +ALIASES: dict[str, str] = { + "reshape": "view", + "flatten": "view", + "unsqueeze": "view", + "expand": "view", + "contiguous": "view", + "identity": "view", + "transpose": "permute", + "concat": "cat", + "type_as": "precision_cast", + "scaled_dot_product_attention": "attention", + "cos": "scalar_cos", + "sin": "scalar_sin", +} + + +PRECISION_OP = "precision_cast" +SUPPORTED_PRECISIONS: tuple[str, ...] = ("int4", "int8", "fp16", "fp32") + + +def canonicalize_op(name: str) -> str: + return ALIASES.get(name, name) + + +def get_op(name: str) -> OpSchema: + return ALL_OPS[canonicalize_op(name)] + + +def has_op(name: str) -> bool: + return canonicalize_op(name) in ALL_OPS diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py new file mode 100644 index 000000000..e66d80881 --- /dev/null +++ b/python/cactus/transpile/optimize_graph.py @@ -0,0 +1,1656 @@ +from __future__ import annotations + +from dataclasses import dataclass +import re + +import torch + +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.canonicalize.utils import normalize_dtype_name +from cactus.transpile.canonicalize.utils import remove_node +from cactus.transpile.canonicalize.utils import rebuild_graph +from cactus.transpile.fusion import match_attention +from cactus.transpile.fusion import match_attention_block +from cactus.transpile.fusion import match_conv_module +from cactus.transpile.fusion import match_gated_deltanet +from cactus.transpile.fusion import match_gated_mlp +from cactus.transpile.fusion import match_lstm_cell +from cactus.transpile.fusion import match_rel_pos_bias +from cactus.transpile.fusion import match_rms_norm +from cactus.transpile.fusion import match_rope +from cactus.transpile.fusion import match_self_attention_block +from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_layout_passthrough +from cactus.transpile.fusion.common import strip_passthrough +from cactus.transpile.fusion.linear import match_linear +from cactus.transpile.fusion.rope import _extract_rope_angle_source +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.graph_ir import verify_ir +from cactus.transpile.model_patterns import GOLD_PATTERNS +from cactus.transpile.model_profiles import profile_for_family +from cactus.transpile.normalize import dtype_to_ir + + +@dataclass(frozen=True) +class DetectedPattern: + name: str + anchor_node_id: str + node_ids: tuple[str, ...] + value_ids: tuple[str, ...] + details: dict[str, object] + + +@dataclass(frozen=True) +class FusionConfig: + enable_gated_deltanet: bool = True + enable_lstm_cell: bool = True + enable_rms_norm: bool = True + enable_rope: bool = True + enable_rel_pos_bias: bool = True + enable_attention: bool = True + enable_self_attention_block: bool = True + enable_attention_block: bool = True + enable_conv_module: bool = True + enable_add_clipped: bool = True + enable_dense_mlp_tq_fused: bool = True + + +def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None, + precompute_rope: bool = True) -> IRGraph: + config = config or FusionConfig() + verify_ir(graph) + canonicalize_exported_graph(graph) + enable_attention_block_fusions = not ( + _is_gemma4_graph(graph) or _is_whisper_seq2seq_decoder_graph(graph) + ) + + for _ in range(max_passes): + changed = False + if config.enable_gated_deltanet and fuse_gated_deltanet(graph): + changed = True + if config.enable_lstm_cell and fuse_lstm_cells(graph): + changed = True + if config.enable_rms_norm: + if fuse_rms_norm(graph): + changed = True + if fuse_rms_norm_scale_multiply(graph): + changed = True + if config.enable_rope and fuse_rope(graph): + changed = True + if config.enable_rel_pos_bias and fuse_rel_pos_bias(graph): + changed = True + if config.enable_attention and fuse_attention(graph): + changed = True + if normalize_attention_layouts(graph): + changed = True + if ( + enable_attention_block_fusions + and config.enable_self_attention_block + and fuse_self_attention_blocks(graph) + ): + changed = True + if ( + enable_attention_block_fusions + and config.enable_attention_block + and fuse_attention_blocks(graph) + ): + changed = True + if normalize_gemma4_decoder_attention_semantics(graph): + changed = True + if normalize_cached_decoder_attention_hints(graph): + changed = True + if config.enable_dense_mlp_tq_fused and fuse_dense_mlp_tq(graph): + changed = True + if config.enable_conv_module and fuse_conv_modules(graph): + changed = True + if config.enable_add_clipped and fuse_add_clipped(graph): + changed = True + if not changed: + break + canonicalize_exported_graph(graph) + + if precompute_rope and precompute_rope_tables(graph): + canonicalize_exported_graph(graph) + + annotate_gold_patterns(graph) + _prune_unused_inputs(graph) + verify_ir(graph) + return graph + + +def fuse_gated_deltanet(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_gated_deltanet(graph, node) + if match is None: + continue + + inputs = [ + match.normalized_input_value_id, + match.qkv_weight_value_id, + match.a_weight_value_id, + match.b_weight_value_id, + match.norm_weight_value_id, + ] + if match.z_weight_value_id is not None: + inputs.append(match.z_weight_value_id) + if match.dt_bias_value_id is not None: + inputs.append(match.dt_bias_value_id) + if match.a_log_value_id is not None: + inputs.append(match.a_log_value_id) + if match.conv_weight_value_id is not None: + inputs.append(match.conv_weight_value_id) + + node.op = f"gated_deltanet_{match.mode}" + node.inputs = inputs + node.attrs = { + "num_k_heads": int(match.num_k_heads), + "num_v_heads": int(match.num_v_heads), + "key_dim": int(match.key_dim), + "value_dim": int(match.value_dim), + "eps": float(match.eps), + "chunk_size": int(match.chunk_size), + "has_z": bool(match.z_weight_value_id is not None), + "has_dt_bias": bool(match.dt_bias_value_id is not None), + "has_a_log": bool(match.a_log_value_id is not None), + "has_conv": bool(match.conv_weight_value_id is not None), + } + node.kind = "semantic" + node.meta["gated_deltanet_mode"] = match.mode + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_lstm_cells(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op != "multiply": + continue + + match = match_lstm_cell(graph, node) + if match is None: + continue + + bias_hh_value_id = match.bias_hh_value_id + if bias_hh_value_id is None: + bias_hh_value_id = _materialize_zero_like_constant( + graph, + match.bias_ih_value_id, + suffix="lstm_bias_hh_zero", + ) + + node.op = "lstm_cell" + node.inputs = [ + match.x_value_id, + match.h_prev_value_id, + match.c_prev_value_id, + match.weight_ih_value_id, + match.weight_hh_value_id, + match.bias_ih_value_id, + bias_hh_value_id, + ] + node.outputs = [match.h_output_value_id, match.c_output_value_id] + node.attrs = {} + node.kind = "semantic" + node.meta["lstm_cell_nodes"] = match.node_ids + node.meta["lstm_cell_bias_hh_zero"] = bool(match.bias_hh_value_id is None) + + for other_node_id in match.node_ids: + if other_node_id == node.id: + continue + graph.nodes.pop(other_node_id, None) + if other_node_id in graph.order: + graph.order.remove(other_node_id) + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def annotate_gold_patterns(graph: IRGraph) -> list[DetectedPattern]: + _clear_gold_pattern_annotations(graph) + patterns = [ + *_detect_gated_mlps(graph), + *_detect_decoder_attentions(graph), + *_detect_transformer_blocks(graph), + ] + + for pattern in patterns: + for node_id in pattern.node_ids: + node = graph.nodes.get(node_id) + if node is None: + continue + node.meta.setdefault("gold_patterns", []) + node.meta["gold_patterns"].append(pattern.name) + + anchor = graph.nodes.get(pattern.anchor_node_id) + if anchor is not None: + anchor.meta["gold_pattern_anchor"] = True + anchor.meta["gold_pattern_details"] = pattern.details + + graph.meta["gold_patterns_catalog"] = tuple(pattern.name for pattern in GOLD_PATTERNS) + graph.meta["detected_gold_patterns"] = patterns + return patterns + + +def summarize_detected_gold_patterns(graph: IRGraph) -> dict[str, int]: + summary: dict[str, int] = {} + for pattern in annotate_gold_patterns(graph): + summary[pattern.name] = summary.get(pattern.name, 0) + 1 + return summary + + +def fuse_rms_norm(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"multiply", "type_as", "precision_cast"}: + continue + + match = match_rms_norm(graph, node) + if match is None: + continue + input_value = graph.values.get(match.input_value_id) + if input_value is not None and normalize_dtype_name(input_value.dtype) == "fp32": + # Keep FP32 normalization paths unfused so they can execute with FP32 + # reductions and elementwise math instead of the FP16-only RMS kernel. + continue + + weight_value_id = match.weight_value_id + if weight_value_id is None: + if input_value is None or input_value.shape is None or not input_value.shape: + continue + hidden_dim = input_value.shape[-1] + if not isinstance(hidden_dim, int) or hidden_dim <= 0: + continue + weight_value_id = _materialize_ones_constant( + graph, + hidden_dim, + dtype=input_value.dtype, + suffix="rms_norm_ones", + ) + node.op = "rms_norm" + node.inputs = [match.input_value_id, weight_value_id] + node.attrs = { + "eps": float(match.eps), + "weight_offset": float(match.weight_offset), + } + node.kind = "semantic" + node.meta["rms_weight_offset"] = float(match.weight_offset) + node.meta["rms_input_value_id"] = match.input_value_id + node.meta["rms_weight_value_id"] = weight_value_id + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_rms_norm_scale_multiply(graph: IRGraph) -> bool: + if not _is_gemma4_graph(graph): + return False + + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op != "multiply" or len(node.inputs) != 2 or len(node.outputs) != 1: + continue + + rms_node: IRNode | None = None + rms_output_id: str | None = None + scale_value_id: str | None = None + for input_id, other_id in ((node.inputs[0], node.inputs[1]), (node.inputs[1], node.inputs[0])): + candidate = producer(graph, input_id) + if candidate is None or candidate.op != "rms_norm" or len(candidate.inputs) != 2: + continue + if not _is_materialized_ones_constant(graph, candidate.inputs[1]): + continue + rms_node = candidate + rms_output_id = input_id + scale_value_id = other_id + break + + if rms_node is None or rms_output_id is None or scale_value_id is None: + continue + if rms_node.outputs != [rms_output_id]: + continue + if _has_other_users(graph, rms_output_id, excluding=node.id) or rms_output_id in graph.outputs: + continue + + node.op = "rms_norm" + node.inputs = [rms_node.inputs[0], scale_value_id] + node.attrs = dict(rms_node.attrs) + node.kind = "semantic" + node.meta["rms_folded_scale_multiply"] = True + node.meta["rms_input_value_id"] = rms_node.inputs[0] + node.meta["rms_weight_value_id"] = scale_value_id + remove_node(graph, rms_node.id) + graph.values.pop(rms_output_id, None) + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_rope(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or not node.outputs: + continue + + match = match_rope(graph, node.outputs[0]) + if match is None or match.partial: + continue + + node.op = "rope" + node.inputs = [match.input_value_id] + node.attrs = { + "theta": float(match.theta), + "position_offset": int(match.position_offset), + } + node.kind = "semantic" + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def precompute_rope_tables(graph: IRGraph) -> bool: + """Replace the runtime cos/sin(position * inv_freq) rope angle with an fp64-precomputed fp16 table + gathered by position id, fixing fp16's inability to represent angles past ~2048 in long context.""" + + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"scalar_cos", "scalar_sin"} or len(node.inputs) != 1: + continue + # Only the text decoder's rope. Vision/audio encoders carry cos/sin with + # the same shape but an axial/2D position this scalar table cannot model. + if node.meta.get("component") != "decoder": + continue + + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + if output_value is None or output_value.shape is None: + continue + index_shape = tuple(int(dim) for dim in output_value.shape[:-1]) + + match = _match_rope_angle_table_source(graph, node.inputs[0], index_shape) + if match is None: + continue + inv_freq_value_id, position_value_id = match + + inv_freq_constant = graph.constants[inv_freq_value_id] + if isinstance(inv_freq_constant, torch.Tensor) and inv_freq_constant.is_meta: + continue + inv_freq = inv_freq_constant.detach().cpu().to(torch.float64).reshape(-1) + head_dim = int(inv_freq.numel()) * 2 + max_seq = _rope_table_max_seq(graph) + positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) + freqs = positions * inv_freq.reshape(1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + table = (torch.cos(emb) if node.op == "scalar_cos" else torch.sin(emb)).to(torch.float16) + assert table.shape[-1] == head_dim + + table_value_id = _materialize_rope_table_constant(graph, node.id, table) + output_value.dtype = "fp16" + node.op = "embedding" + node.inputs = [table_value_id, position_value_id] + node.attrs = {} + node.kind = "generic" + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def _stored_inv_freq_constant(graph: IRGraph, value_id: str) -> str | None: + stripped = strip_layout_passthrough(graph, value_id) + constant = graph.constants.get(stripped) + if isinstance(constant, torch.Tensor): + return stripped + return None + + +def _match_rope_angle_table_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> tuple[str, str] | None: + cat_node = producer(graph, strip_passthrough(graph, value_id)) + if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: + return None + if strip_layout_passthrough(graph, cat_node.inputs[0]) != strip_layout_passthrough(graph, cat_node.inputs[1]): + return None + angle_info = _extract_rope_angle_source(graph, cat_node.inputs[0]) + if angle_info is None: + return None + angle_node = angle_info["matmul_node"] + if len(angle_node.inputs) != 2: + return None + + inv_freq_value_id: str | None = None + position_value_id: str | None = None + for input_id in angle_node.inputs: + const_id = _stored_inv_freq_constant(graph, input_id) + if const_id is not None: + inv_freq_value_id = const_id + continue + position_value_id = _trace_integer_position_source(graph, input_id, index_shape) + + if inv_freq_value_id is None or position_value_id is None: + return None + return inv_freq_value_id, position_value_id + + +def _trace_integer_position_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> str | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + value = graph.values.get(current) + if ( + value is not None + and str(value.dtype) in {"int32", "int64", "i32", "i64"} + and value.shape is not None + and tuple(int(dim) for dim in value.shape) == index_shape + ): + return current + node = producer(graph, current) + if node is None or node.op not in {"precision_cast", "view", "reshape", "unsqueeze", "squeeze", "contiguous"} or len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _rope_table_max_seq(graph: IRGraph) -> int: + max_seq = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) + if max_seq is None or max_seq <= 0: + max_seq = _coerce_optional_int(graph.meta.get("max_position_embeddings")) + if max_seq is None or max_seq <= 0: + raise ValueError("rope table precompute requires max_cache_seq_len or max_position_embeddings in graph.meta") + return int(max_seq) + + +def _materialize_rope_table_constant(graph: IRGraph, node_id: str, table: torch.Tensor) -> str: + value_id = f"c_rope_table_{node_id}" + graph.constants[value_id] = table + graph.values[value_id] = IRValue( + id=value_id, + shape=tuple(int(dim) for dim in table.shape), + dtype=dtype_to_ir(table.dtype), + producer=None, + users=[], + ) + return value_id + + +def fuse_attention(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"scaled_dot_product_attention", "attention"}: + continue + + match = match_attention(graph, node) + if match is None: + continue + + window_size = int(node.attrs.get("window_size", 0)) + if len(node.inputs) > 3: + mask_info = _extract_sliding_window_mask(graph, node.inputs[3]) + if mask_info is not None: + node.meta["mask_window_size_hint"] = int(mask_info["window_size"]) + if window_size == 0: + window_size = int(mask_info["window_size"]) + node.meta["window_size_source"] = "mask_pattern" + + if window_size != 0 and "window_size_source" not in node.meta: + node.meta["window_size_source"] = "import_attr" + + semantic_attrs = { + key: value + for key, value in node.attrs.items() + if key not in {"mask", "dropout_p", "scale", "is_causal"} + } + semantic_attrs.update( + { + "scale": float(node.attrs.get("scale", 0.0)), + "is_causal": bool(node.attrs.get("is_causal", True)), + "window_size": window_size, + } + ) + + if ( + node.op == "attention" + and node.inputs[:3] == [match.query_value_id, match.key_value_id, match.value_value_id] + and node.attrs == semantic_attrs + ): + continue + + node.op = "attention" + node.inputs = [match.query_value_id, match.key_value_id, match.value_value_id] + node.attrs = semantic_attrs + node.kind = "semantic" + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_rel_pos_bias(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_rel_pos_bias(graph, node) + if match is None: + continue + + node.op = "rel_pos_bias" + node.inputs = [match.query_value_id, match.relative_key_value_id] + node.attrs = {"scale": float(match.scale)} + node.kind = "semantic" + node.meta["rel_pos_bias_nodes"] = match.node_ids + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_attention_blocks(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_attention_block(graph, node) + if match is None: + continue + + qkv_inputs = [match.query_value_id, match.key_value_id, match.value_value_id] + qkv_layout = "bhsd" + native_qkv_inputs = [ + _strip_bhsd_to_bthd_permute_input(graph, value_id) + for value_id in qkv_inputs + ] + if all(value_id is not None for value_id in native_qkv_inputs): + qkv_inputs = [str(value_id) for value_id in native_qkv_inputs] + qkv_layout = "bthd" + + inputs = qkv_inputs + if match.mask_value_id is not None: + inputs.append(match.mask_value_id) + if match.gate_value_id is not None: + inputs.append(match.gate_value_id) + inputs.append(match.output_projection_weight_value_id) + if match.output_projection_bias_value_id is not None: + inputs.append(match.output_projection_bias_value_id) + + node.op = "attention_block" + node.inputs = inputs + node.attrs = { + "scale": float(match.scale), + "is_causal": bool(match.is_causal), + "window_size": int(match.window_size), + "has_mask": bool(match.mask_value_id is not None), + "additive_mask": bool(match.additive_mask), + "has_gate": bool(match.gate_value_id is not None), + "has_bias": bool(match.output_projection_bias_value_id is not None), + "attention_output_shape": tuple(int(dim) for dim in match.attention_output_shape), + "qkv_layout": qkv_layout, + } + node.kind = "semantic" + node.meta["attention_block_source"] = match.attention_node_id + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def _strip_bhsd_to_bthd_permute_input(graph: IRGraph, value_id: str) -> str | None: + value = graph.values.get(value_id) + if value is None or value.producer is None: + return None + node = graph.nodes.get(value.producer) + if node is None or node.op != "permute" or len(node.inputs) != 1: + return None + permutation = tuple(int(dim) for dim in node.attrs.get("permutation", ())) + if permutation != (0, 2, 1, 3): + return None + source_id = strip_passthrough(graph, node.inputs[0]) + source_value = graph.values.get(source_id) + if source_value is None or source_value.shape is None or value.shape is None: + return None + source_shape = tuple(int(dim) for dim in source_value.shape) + output_shape = tuple(int(dim) for dim in value.shape) + if len(source_shape) != 4 or len(output_shape) != 4: + return None + if ( + output_shape[0] == source_shape[0] + and output_shape[1] == source_shape[2] + and output_shape[2] == source_shape[1] + and output_shape[3] == source_shape[3] + ): + return source_id + return None + + +def normalize_attention_layouts(graph: IRGraph) -> bool: + """Erase PyTorch BHSD<->BTHD layout round-trips around attention ops.""" + + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"attention", "scaled_dot_product_attention"}: + continue + if len(node.inputs) < 3 or len(node.outputs) != 1: + continue + + node_changed = False + layouts = ["bhsd", "bhsd", "bhsd"] + for index, value_id in enumerate(node.inputs[:3]): + native_value_id = _strip_bhsd_to_bthd_permute_input(graph, value_id) + if native_value_id is None: + continue + node.inputs[index] = str(native_value_id) + layouts[index] = "bthd" + node_changed = True + if node_changed: + node.attrs["q_layout"] = layouts[0] + node.attrs["k_layout"] = layouts[1] + node.attrs["v_layout"] = layouts[2] + changed = True + + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + output_users = list(output_value.users) if output_value is not None else [] + if len(output_users) != 1 or output_id in graph.outputs: + continue + output_user = graph.nodes.get(output_users[0]) + if output_user is None or output_user.op != "permute" or len(output_user.inputs) != 1 or len(output_user.outputs) != 1: + continue + permutation = tuple(int(dim) for dim in output_user.attrs.get("permutation", ())) + if permutation != (0, 2, 1, 3): + continue + + replacement_id = output_user.outputs[0] + replacement_value = graph.values.get(replacement_id) + if replacement_value is not None: + old_value = graph.values.get(output_id) + if old_value is not None: + old_value.shape = replacement_value.shape + old_value.dtype = replacement_value.dtype + old_value.meta.update(replacement_value.meta) + node.outputs = [replacement_id] + node.attrs["output_layout"] = "bthd" + remove_node(graph, output_user.id) + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def normalize_gemma4_decoder_attention_semantics(graph: IRGraph) -> bool: + if not _is_gemma4_graph(graph): + return False + + changed = False + if _assign_gemma4_decoder_attention_hints_from_graph_meta(graph): + changed = True + sliding_window = _graph_sliding_window(graph) + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: + continue + layer_type = str(node.meta.get("attention_layer_type") or "").strip().lower() + if layer_type == "full_attention": + mask_input_index: int | None = None + if node.op == "attention_block": + if bool(node.attrs.get("has_mask", False)) and len(node.inputs) > 3: + mask_input_index = 3 + elif len(node.inputs) > 3: + mask_input_index = 3 + + if mask_input_index is not None: + mask_value_id = node.inputs[mask_input_index] + if _gemma4_can_elide_attention_mask(graph, mask_value_id): + del node.inputs[mask_input_index] + node.meta["gemma4_full_mask_elided"] = True + changed = True + + if node.op == "attention_block" and bool(node.attrs.get("has_mask", False)): + node.attrs["has_mask"] = False + changed = True + if not bool(node.attrs.get("is_causal", False)): + node.attrs["is_causal"] = True + changed = True + if "additive_mask" in node.attrs and not bool(node.attrs.get("additive_mask", False)): + node.attrs.pop("additive_mask", None) + changed = True + if int(node.attrs.get("window_size", 0)) == 0 and not bool(graph.meta.get("use_internal_kv_cache", False)): + seq_len = 0 + attention_output_shape = node.attrs.get("attention_output_shape") + if isinstance(attention_output_shape, (list, tuple)) and len(attention_output_shape) >= 2: + try: + seq_len = int(attention_output_shape[1]) + except (TypeError, ValueError): + seq_len = 0 + if seq_len <= 0 and node.inputs: + query_value = graph.values.get(node.inputs[0]) + if query_value is not None and query_value.shape is not None and len(query_value.shape) >= 3: + seq_len = int(query_value.shape[2]) + if seq_len > 0: + # Gemma4 full-attention layers are semantically unwindowed, but using a + # window_size >= seq_len avoids the Apple Accelerate full-window fast path + # that has been unstable for transpiled grouped-query Gemma4 graphs. + node.attrs["window_size"] = seq_len + node.meta["gemma4_full_attention_window_compat"] = True + changed = True + continue + if layer_type != "sliding_attention": + continue + if bool(node.attrs.get("additive_mask", False)): + continue + + mask_input_index: int | None = None + if node.op == "attention_block": + if bool(node.attrs.get("has_mask", False)) and len(node.inputs) > 3: + mask_input_index = 3 + elif len(node.inputs) > 3: + mask_input_index = 3 + + if mask_input_index is not None: + mask_value_id = node.inputs[mask_input_index] + if _gemma4_can_elide_attention_mask(graph, mask_value_id): + mask_info = _extract_sliding_window_mask(graph, mask_value_id) + mask_input_name = _graph_input_name(graph, mask_value_id) + del node.inputs[mask_input_index] + node.meta["gemma4_sliding_mask_elided"] = True + if mask_info is not None: + node.meta["mask_window_size_hint"] = int(mask_info["window_size"]) + if int(node.attrs.get("window_size", 0)) == 0 and int(mask_info["window_size"]) > 0: + node.attrs["window_size"] = int(mask_info["window_size"]) + node.meta.setdefault("window_size_source", "mask_pattern") + elif mask_input_name is not None: + node.meta["gemma4_elided_mask_input_name"] = mask_input_name + changed = True + + if node.op == "attention_block" and bool(node.attrs.get("has_mask", False)): + node.attrs["has_mask"] = False + changed = True + if not bool(node.attrs.get("is_causal", False)): + node.attrs["is_causal"] = True + changed = True + if ( + int(node.attrs.get("window_size", 0)) == 0 + and sliding_window is not None + and sliding_window > 0 + ): + node.attrs["window_size"] = int(sliding_window) + node.meta.setdefault("window_size_source", "layer_type_config") + changed = True + if "additive_mask" in node.attrs and not bool(node.attrs.get("additive_mask", False)): + node.attrs.pop("additive_mask", None) + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def normalize_cached_decoder_attention_hints(graph: IRGraph) -> bool: + if _is_gemma4_graph(graph): + return False + if not bool(graph.meta.get("use_internal_kv_cache", False)): + return False + changed = _assign_gemma4_decoder_attention_hints_from_graph_meta(graph) + sliding_window = _graph_sliding_window(graph) + if sliding_window is not None and sliding_window > 0: + for node_id in graph.order: + node = graph.nodes.get(node_id) + if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: + continue + layer_type = str(node.meta.get("attention_layer_type") or "").strip().lower() + if layer_type not in {"sliding", "sliding_attention"}: + continue + if int(node.attrs.get("window_size", 0) or 0) == 0: + node.attrs["window_size"] = int(sliding_window) + node.meta.setdefault("window_size_source", "layer_type_config") + changed = True + if changed: + rebuild_graph(graph) + return changed + + +def _assign_gemma4_decoder_attention_hints_from_graph_meta(graph: IRGraph) -> bool: + component = str(graph.meta.get("component", "") or "").strip().lower() + if component not in {"decoder", "decoder_step", "decoder_prefill_chunk", "decoder_media_step"}: + return False + + layer_types = _graph_layer_types(graph) + if not layer_types: + return False + + attention_nodes: list[IRNode] = [] + for node_id in graph.order: + node = graph.nodes.get(node_id) + if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: + continue + attention_nodes.append(node) + + if len(attention_nodes) != len(layer_types): + return False + + use_internal_cache = bool(graph.meta.get("use_internal_kv_cache", False)) + sliding_window = _graph_sliding_window(graph) + graph_cache_seq_len = _graph_max_cache_seq_len(graph) + changed = False + for layer_index, (node, layer_type) in enumerate(zip(attention_nodes, layer_types, strict=True)): + if "attention_layer_type" not in node.meta: + node.meta["attention_layer_type"] = str(layer_type) + changed = True + if "attention_layer_index" not in node.meta: + node.meta["attention_layer_index"] = int(layer_index) + changed = True + if use_internal_cache: + normalized_layer_type = str(layer_type).strip().lower() + node_cache_seq_len = None + if normalized_layer_type in {"sliding", "sliding_attention"} and sliding_window is not None and sliding_window > 0: + node_cache_seq_len = int(sliding_window) + elif normalized_layer_type in {"global", "full", "full_attention"} and graph_cache_seq_len is not None: + node_cache_seq_len = int(graph_cache_seq_len) + if node_cache_seq_len is not None and node.meta.get("max_cache_seq_len") != node_cache_seq_len: + node.meta["max_cache_seq_len"] = node_cache_seq_len + changed = True + return changed + + +def _gemma4_can_elide_attention_mask( + graph: IRGraph, + value_id: str, + *, + _visited: set[str] | None = None, +) -> bool: + mask_value = graph.values.get(value_id) + if mask_value is None: + return False + if mask_value.dtype in (None, "bool"): + return True + + input_name = _graph_input_name(graph, value_id) + if isinstance(input_name, str) and "attention_mask" in input_name: + return True + + if _visited is None: + _visited = set() + if value_id in _visited: + return False + _visited.add(value_id) + + node = producer(graph, strip_passthrough(graph, value_id)) + if node is None: + return False + if node.op in { + "logical_and", + "logical_or", + "scalar_equal", + "scalar_not_equal", + "greater", + "greater_equal", + "less", + "less_equal", + "aten.__and__.Tensor", + }: + return True + if node.op in {"precision_cast", "view", "expand", "permute", "slice", "index", "where"}: + return any(_gemma4_can_elide_attention_mask(graph, input_id, _visited=_visited) for input_id in node.inputs) + return False + + +def fuse_self_attention_blocks(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_self_attention_block(graph, node) + if match is None: + continue + + inputs = [match.hidden_value_id, match.query_weight_value_id] + if match.query_projection_bias_value_id is not None: + inputs.append(match.query_projection_bias_value_id) + if match.query_add_value_id is not None: + inputs.append(match.query_add_value_id) + if match.rel_query_add_value_id is not None: + inputs.append(match.rel_query_add_value_id) + inputs.append(match.key_weight_value_id) + if match.key_projection_bias_value_id is not None: + inputs.append(match.key_projection_bias_value_id) + inputs.append(match.value_weight_value_id) + if match.value_projection_bias_value_id is not None: + inputs.append(match.value_projection_bias_value_id) + if match.mask_value_id is not None: + inputs.append(match.mask_value_id) + if match.relative_key_input_value_id is not None and match.relative_key_weight_value_id is not None: + inputs.append(match.relative_key_input_value_id) + inputs.append(match.relative_key_weight_value_id) + if match.relative_key_projection_bias_value_id is not None: + inputs.append(match.relative_key_projection_bias_value_id) + if match.gate_value_id is not None: + inputs.append(match.gate_value_id) + inputs.append(match.output_projection_weight_value_id) + if match.output_projection_bias_value_id is not None: + inputs.append(match.output_projection_bias_value_id) + + node.op = "self_attention_block" + node.inputs = inputs + node.attrs = { + "scale": float(match.scale), + "is_causal": bool(match.is_causal), + "window_size": int(match.window_size), + "has_mask": bool(match.mask_value_id is not None), + "additive_mask": bool(match.additive_mask), + "has_gate": bool(match.gate_value_id is not None), + "has_bias": bool(match.output_projection_bias_value_id is not None), + "has_query_projection_bias": bool(match.query_projection_bias_value_id is not None), + "has_query_add": bool(match.query_add_value_id is not None), + "has_rel_query_add": bool(match.rel_query_add_value_id is not None), + "has_key_projection_bias": bool(match.key_projection_bias_value_id is not None), + "has_value_projection_bias": bool(match.value_projection_bias_value_id is not None), + "has_rel_pos_bias": bool(match.relative_key_input_value_id is not None and match.relative_key_weight_value_id is not None), + "has_relative_key_projection_bias": bool(match.relative_key_projection_bias_value_id is not None), + "query_shape": tuple(int(dim) for dim in match.query_shape), + "key_shape": tuple(int(dim) for dim in match.key_shape), + "value_shape": tuple(int(dim) for dim in match.value_shape), + "relative_key_shape": tuple(int(dim) for dim in match.relative_key_shape) if match.relative_key_shape is not None else (), + "rel_pos_scale": float(match.rel_pos_scale) if match.rel_pos_scale is not None else 1.0, + "attention_output_shape": tuple(int(dim) for dim in match.attention_output_shape), + } + node.kind = "semantic" + node.meta["self_attention_block_source"] = match.attention_node_id + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_add_clipped(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op != "add" or len(node.inputs) != 2: + continue + + lhs = strip_passthrough(graph, node.inputs[0]) + rhs = strip_passthrough(graph, node.inputs[1]) + if not (_looks_like_gemma_residual_add(graph, lhs, rhs) or _looks_like_gemma_residual_add(graph, rhs, lhs)): + continue + + node.op = "add_clipped" + node.kind = "semantic" + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def fuse_dense_mlp_tq(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = _match_dense_mlp_tq(graph, node) + if match is None: + continue + + matched_node_ids = set(match["node_ids"]) + if not _matched_nodes_are_private(graph, matched_node_ids, keep_node_id=node.id): + continue + + node.op = "dense_mlp_tq_fused" + node.inputs = [ + str(match["input_value_id"]), + str(match["gate_weight_value_id"]), + str(match["up_weight_value_id"]), + str(match["down_weight_value_id"]), + ] + node.attrs = {"product_scale": float(match.get("product_scale") or 1.0)} + node.kind = "semantic" + node.meta["dense_mlp_tq_fused"] = True + node.meta["dense_mlp_tq_nodes"] = tuple(sorted(matched_node_ids)) + if match.get("product_scale") is not None: + node.meta["product_scale_from_export"] = float(match["product_scale"]) + + for other_node_id in matched_node_ids - {node.id}: + graph.nodes.pop(other_node_id, None) + if other_node_id in graph.order: + graph.order.remove(other_node_id) + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def _match_dense_mlp_tq(graph: IRGraph, node: IRNode) -> dict[str, object] | None: + down = match_linear(graph, node) + if down is None or down.bias_value_id is not None: + return None + if not _is_cq_weight_value(graph, down.weight_value_id): + return None + + mul_node = producer(graph, down.input_value_id) + if mul_node is None or mul_node.op != "multiply" or len(mul_node.inputs) != 2: + return None + + for activated_value_id, up_value_id in ((mul_node.inputs[0], mul_node.inputs[1]), (mul_node.inputs[1], mul_node.inputs[0])): + activation_node, product_scale, scale_node_id = _unwrap_gemma4_scaled_activation(graph, activated_value_id) + if activation_node is None or activation_node.op != "gelu" or len(activation_node.inputs) != 1: + continue + + gate_input_producer = producer(graph, strip_passthrough(graph, activation_node.inputs[0])) + up_producer = producer(graph, strip_passthrough(graph, up_value_id)) + if gate_input_producer is None or up_producer is None: + continue + + gate = match_linear(graph, gate_input_producer) + up = match_linear(graph, up_producer) + if gate is None or up is None: + continue + if gate.bias_value_id is not None or up.bias_value_id is not None: + continue + if strip_passthrough(graph, gate.input_value_id) != strip_passthrough(graph, up.input_value_id): + continue + if not _is_cq_weight_value(graph, gate.weight_value_id): + continue + if not _is_cq_weight_value(graph, up.weight_value_id): + continue + if not _dense_mlp_weight_shapes_match(graph, gate.weight_value_id, up.weight_value_id, down.weight_value_id): + continue + + node_ids = { + *gate.node_ids, + activation_node.id, + *up.node_ids, + mul_node.id, + *down.node_ids, + } + if scale_node_id is not None: + node_ids.add(scale_node_id) + + return { + "input_value_id": strip_passthrough(graph, gate.input_value_id), + "gate_weight_value_id": gate.weight_value_id, + "up_weight_value_id": up.weight_value_id, + "down_weight_value_id": down.weight_value_id, + "product_scale": product_scale, + "node_ids": tuple(sorted(node_ids)), + } + + return None + + +def _unwrap_gemma4_scaled_activation(graph: IRGraph, value_id: str) -> tuple[IRNode | None, float | None, str | None]: + node = producer(graph, strip_passthrough(graph, value_id)) + if node is None: + return None, None, None + if node.op == "scalar_multiply" and len(node.inputs) == 1: + scale = float(node.attrs.get("value", 1.0)) + inner = producer(graph, strip_passthrough(graph, node.inputs[0])) + return inner, scale, node.id + return node, None, None + + +def _is_cq_weight_value(graph: IRGraph, value_id: str) -> bool: + value = graph.values.get(value_id) + if value is None: + return False + path = value.meta.get("path") if isinstance(value.meta, dict) else None + if isinstance(path, str) and ".cq" in path: + return True + return False + + +def _dense_mlp_weight_shapes_match( + graph: IRGraph, + gate_weight_value_id: str, + up_weight_value_id: str, + down_weight_value_id: str, +) -> bool: + gate = graph.values.get(gate_weight_value_id) + up = graph.values.get(up_weight_value_id) + down = graph.values.get(down_weight_value_id) + if gate is None or up is None or down is None: + return False + if gate.shape is None or up.shape is None or down.shape is None: + return False + if len(gate.shape) != 2 or len(up.shape) != 2 or len(down.shape) != 2: + return False + return ( + int(gate.shape[0]) == int(up.shape[0]) + and int(gate.shape[1]) == int(up.shape[1]) + and int(down.shape[1]) == int(gate.shape[0]) + ) + + +def _matched_nodes_are_private(graph: IRGraph, matched_node_ids: set[str], *, keep_node_id: str) -> bool: + for node_id in matched_node_ids: + if node_id == keep_node_id: + continue + node = graph.nodes.get(node_id) + if node is None: + continue + for output_id in node.outputs: + value = graph.values.get(output_id) + if value is None: + continue + for user_id in value.users: + if user_id not in matched_node_ids: + return False + return True + + +def fuse_conv_modules(graph: IRGraph) -> bool: + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_conv_module(graph, node) + if match is None: + continue + + inputs = [ + match.input_value_id, + match.pointwise1_weight_value_id, + ] + if match.pointwise1_bias_value_id is not None: + inputs.append(match.pointwise1_bias_value_id) + inputs.append(match.depthwise_weight_value_id) + if match.depthwise_bias_value_id is not None: + inputs.append(match.depthwise_bias_value_id) + inputs.extend( + [ + match.batch_norm_weight_value_id, + match.batch_norm_bias_value_id, + match.batch_norm_running_mean_value_id, + match.batch_norm_running_var_value_id, + match.pointwise2_weight_value_id, + ] + ) + if match.pointwise2_bias_value_id is not None: + inputs.append(match.pointwise2_bias_value_id) + + node.op = "conv_module" + node.inputs = inputs + node.attrs = { + "eps": float(match.eps), + "has_pointwise1_bias": bool(match.pointwise1_bias_value_id is not None), + "has_depthwise_bias": bool(match.depthwise_bias_value_id is not None), + "has_pointwise2_bias": bool(match.pointwise2_bias_value_id is not None), + "depthwise_kernel_size": int(match.depthwise_kernel_size), + "depthwise_padding": int(match.depthwise_padding), + } + node.kind = "semantic" + node.meta["conv_module_nodes"] = match.node_ids + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def _materialize_shifted_constant(graph: IRGraph, value_id: str, offset: float, *, suffix: str) -> str: + base = graph.constants[value_id] + if not isinstance(base, torch.Tensor): + raise NotImplementedError(f"expected tensor constant for {value_id}, got {type(base).__name__}") + + new_value_id = f"{value_id}_{suffix}" + if new_value_id in graph.constants: + return new_value_id + + shifted = base.detach().cpu() + offset + graph.constants[new_value_id] = shifted + meta: dict[str, object] = {} + source_value = graph.values.get(value_id) + if _is_gemma4_graph(graph) and float(offset) == 1.0 and source_value is not None: + path = source_value.meta.get("path") + kind = source_value.meta.get("kind") + source_name = source_value.meta.get("source_name") + if isinstance(path, str) and isinstance(kind, str) and isinstance(source_name, str): + meta = dict(source_value.meta) + meta["materialized_from_value_id"] = value_id + meta["materialized_by_op"] = "scalar_add" + graph.values[new_value_id] = IRValue( + id=new_value_id, + shape=tuple(shifted.shape), + dtype=dtype_to_ir(shifted.dtype), + producer=None, + users=[], + meta=meta, + ) + return new_value_id + + +def _materialize_zero_like_constant(graph: IRGraph, value_id: str, *, suffix: str) -> str: + base = graph.constants[value_id] + new_value_id = f"{value_id}_{suffix}" + if new_value_id in graph.constants: + return new_value_id + + if isinstance(base, torch.Tensor): + zero = torch.zeros_like(base.detach().cpu()) + else: + value = graph.values.get(value_id) + if value is None or value.shape is None or value.dtype is None: + raise NotImplementedError( + f"expected tensor constant metadata for {value_id}, got {type(base).__name__}" + ) + dtype_map = { + "bool": torch.bool, + "u8": torch.uint8, + "i8": torch.int8, + "i16": torch.int16, + "i32": torch.int32, + "i64": torch.int64, + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp64": torch.float64, + } + torch_dtype = dtype_map.get(str(value.dtype)) + if torch_dtype is None: + raise NotImplementedError(f"unsupported IR dtype for zero-like constant: {value.dtype}") + zero = torch.zeros(tuple(int(dim) for dim in value.shape), dtype=torch_dtype) + + graph.constants[new_value_id] = zero + graph.values[new_value_id] = IRValue( + id=new_value_id, + shape=tuple(zero.shape), + dtype=dtype_to_ir(zero.dtype), + producer=None, + users=[], + ) + return new_value_id + + +def _materialize_ones_constant(graph: IRGraph, size: int, *, dtype: str | None, suffix: str) -> str: + dtype_map = { + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp64": torch.float64, + } + torch_dtype = dtype_map.get(dtype, torch.float32) + new_value_id = f"c_{suffix}_{size}_{dtype or 'fp32'}" + if new_value_id in graph.constants: + return new_value_id + + value = torch.ones((size,), dtype=torch_dtype) + graph.constants[new_value_id] = value + graph.values[new_value_id] = IRValue( + id=new_value_id, + shape=(size,), + dtype=dtype_to_ir(value.dtype), + producer=None, + users=[], + ) + return new_value_id + + +def _is_gemma4_graph(graph: IRGraph) -> bool: + family = str(graph.meta.get("adapter_family") or graph.meta.get("family") or "").lower() + profile = profile_for_family(family) + return profile is not None and profile.family == "gemma4" + + +def _is_whisper_seq2seq_decoder_graph(graph: IRGraph) -> bool: + family = str(graph.meta.get("adapter_family") or graph.meta.get("family") or "").lower() + profile = profile_for_family(family) + task = str(graph.meta.get("task") or "").lower() + component = str(graph.meta.get("component") or "").lower() + return profile is not None and profile.family == "whisper" and task == "seq2seq_transcription" and component == "decoder" + + +def _prune_unused_inputs(graph: IRGraph) -> bool: + kept_inputs: list[str] = [] + changed = False + for value_id in graph.inputs: + value = graph.values.get(value_id) + if value_id in graph.outputs: + kept_inputs.append(value_id) + continue + if value is not None and value.users: + kept_inputs.append(value_id) + continue + changed = True + + if not changed: + return False + + graph.inputs = kept_inputs + rebuild_graph(graph) + return True + + +def _graph_input_name(graph: IRGraph, value_id: str) -> str | None: + input_names = _coerce_string_tuple(graph.meta.get("input_names")) + match = re.match(r"^v_args_(\d+)$", value_id) + if match is not None: + input_index = int(match.group(1)) + if input_index < len(input_names): + return input_names[input_index] + try: + input_index = list(graph.inputs).index(value_id) + except ValueError: + return None + if input_index < len(input_names): + return input_names[input_index] + return None + + +def _is_materialized_ones_constant(graph: IRGraph, value_id: str) -> bool: + value = graph.constants.get(value_id) + if not isinstance(value, torch.Tensor): + return False + if value.numel() == 0: + return False + return bool(torch.all(value.detach().cpu() == 1).item()) + + +def _has_other_users(graph: IRGraph, value_id: str, *, excluding: str) -> bool: + for other_id, other in graph.nodes.items(): + if other_id == excluding: + continue + if value_id in other.inputs: + return True + return False + + +def _clear_gold_pattern_annotations(graph: IRGraph) -> None: + for node in graph.nodes.values(): + node.meta.pop("gold_patterns", None) + node.meta.pop("gold_pattern_anchor", None) + node.meta.pop("gold_pattern_details", None) + + +def _detect_gated_mlps(graph: IRGraph) -> list[DetectedPattern]: + patterns: list[DetectedPattern] = [] + for node_id in graph.order: + node = graph.nodes.get(node_id) + if node is None: + continue + + match = match_gated_mlp(graph, node) + if match is None: + continue + + pattern_name = "gated_mlp_gelu" if match.activation == "gelu" else "gated_mlp_silu" + patterns.append( + DetectedPattern( + name=pattern_name, + anchor_node_id=node.id, + node_ids=match.node_ids, + value_ids=(match.input_value_id,), + details={ + "activation": match.activation, + "input_value_id": match.input_value_id, + "gate_weight_value_id": match.gate_weight_value_id, + "up_weight_value_id": match.up_weight_value_id, + "down_weight_value_id": match.down_weight_value_id, + }, + ) + ) + return patterns + + +def _detect_decoder_attentions(graph: IRGraph) -> list[DetectedPattern]: + patterns: list[DetectedPattern] = [] + for node_id in graph.order: + node = graph.nodes.get(node_id) + if node is None or node.op not in {"scaled_dot_product_attention", "attention"}: + continue + + match = match_attention(graph, node) + if match is None: + continue + + q_input, k_input, v_input = match.source_input_value_ids + pattern_name = "decoder_attention_gqa" if match.has_gqa_repeat else "decoder_attention" + + node_ids = set(match.node_ids) + details = { + "input_value_ids": (q_input, k_input, v_input), + "q_linear_weight_value_id": match.weight_value_ids[0], + "k_linear_weight_value_id": match.weight_value_ids[1], + "v_linear_weight_value_id": match.weight_value_ids[2], + "has_rope": bool(match.has_rope), + "has_qk_rms_norm": bool(match.has_qk_norm), + "has_gqa_repeat": bool(match.has_gqa_repeat), + "is_causal": bool(match.is_causal), + "scale": float(match.scale), + "window_size_hint": int(match.window_size), + } + + if len(node.inputs) > 3: + details["mask_value_id"] = node.inputs[3] + mask_info = _extract_sliding_window_mask(graph, node.inputs[3]) + if mask_info is not None: + details["mask_pattern"] = "sliding_window_attention_mask" + details["window_size_hint"] = int(mask_info["window_size"]) + node_ids.update(mask_info["node_ids"]) + + patterns.append( + DetectedPattern( + name=pattern_name, + anchor_node_id=node.id, + node_ids=tuple(sorted(node_ids)), + value_ids=(q_input, k_input, v_input), + details=details, + ) + ) + return patterns + + +def _detect_transformer_blocks(graph: IRGraph) -> list[DetectedPattern]: + patterns: list[DetectedPattern] = [] + for node_id in graph.order: + node = graph.nodes.get(node_id) + if node is None or node.op not in {"add", "add_clipped"} or len(node.inputs) != 2: + continue + + rhs_pattern = _find_anchor_pattern(graph, node.inputs[1], {"gated_mlp_gelu", "gated_mlp_silu"}) + lhs_pattern = _find_anchor_pattern(graph, node.inputs[0], {"gated_mlp_gelu", "gated_mlp_silu"}) + mlp_pattern = rhs_pattern or lhs_pattern + if mlp_pattern is None: + continue + + residual_value_id = node.inputs[0] if rhs_pattern is not None else node.inputs[1] + attn_pattern = _find_anchor_pattern( + graph, + residual_value_id, + {"decoder_attention_gqa", "decoder_attention", "gemma4_partial_rope_attention"}, + ) + if attn_pattern is None: + continue + + node_ids = {node.id} + for candidate in graph.nodes.values(): + gold_patterns = candidate.meta.get("gold_patterns", []) + if attn_pattern in gold_patterns or mlp_pattern in gold_patterns: + node_ids.add(candidate.id) + + patterns.append( + DetectedPattern( + name="decoder_block_post_attn_norm" if node.op == "add_clipped" else "decoder_block_simple_residual", + anchor_node_id=node.id, + node_ids=tuple(sorted(node_ids)), + value_ids=(residual_value_id,), + details={ + "residual_value_id": residual_value_id, + "attention_pattern": attn_pattern, + "mlp_pattern": mlp_pattern, + "residual_op": node.op, + }, + ) + ) + return patterns + + +def _looks_like_gemma_residual_add(graph: IRGraph, residual_value_id: str, branch_value_id: str) -> bool: + branch_node = producer(graph, strip_passthrough(graph, branch_value_id)) + if branch_node is None or branch_node.op != "rms_norm": + return False + return strip_passthrough(graph, residual_value_id) != strip_passthrough(graph, branch_node.inputs[0]) + + +def _find_anchor_pattern(graph: IRGraph, value_id: str, names: set[str]) -> str | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + current = strip_passthrough(graph, current) + node = producer(graph, current) + if node is None: + return None + + for name in node.meta.get("gold_patterns", []): + if name in names and node.meta.get("gold_pattern_anchor"): + return name + + if len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _extract_sliding_window_mask(graph: IRGraph, value_id: str) -> dict[str, object] | None: + node_ids: set[str] = set() + top_and = producer(graph, strip_passthrough(graph, value_id)) + if top_and is None or top_and.op not in {"aten.__and__.Tensor", "logical_and"}: + return None + node_ids.add(top_and.id) + + stack = list(top_and.inputs) + saw_diff = False + saw_cumsum = False + window_candidates: list[int] = [] + + while stack: + current = strip_passthrough(graph, stack.pop()) + node = producer(graph, current) + if node is None or node.id in node_ids: + continue + node_ids.add(node.id) + + if node.op in {"aten.diff.default", "diff"}: + saw_diff = True + elif node.op in {"aten.cumsum.default", "cumsum"}: + saw_cumsum = True + elif node.op == "scalar_subtract": + maybe_window = int(node.attrs.get("value", 0)) + if maybe_window > 0: + window_candidates.append(maybe_window) + + stack.extend(node.inputs) + + if not saw_diff or not saw_cumsum: + return None + + return { + "window_size": max((candidate for candidate in window_candidates if candidate > 1), default=0), + "node_ids": tuple(sorted(node_ids)), + } + + +def _graph_layer_types(graph: IRGraph) -> tuple[str, ...]: + layer_types = _coerce_string_tuple(graph.meta.get("layer_types")) + if layer_types: + return layer_types + + providers = graph.meta.get("transpile_metadata_providers") + if isinstance(providers, dict): + for provider_meta in providers.values(): + if not isinstance(provider_meta, dict): + continue + layer_types = _coerce_string_tuple(provider_meta.get("layer_types")) + if layer_types: + return layer_types + return () + + +def _graph_sliding_window(graph: IRGraph) -> int | None: + sliding_window = _coerce_optional_int(graph.meta.get("sliding_window")) + if sliding_window is not None: + return sliding_window + + providers = graph.meta.get("transpile_metadata_providers") + if isinstance(providers, dict): + for provider_meta in providers.values(): + if not isinstance(provider_meta, dict): + continue + sliding_window = _coerce_optional_int(provider_meta.get("sliding_window")) + if sliding_window is not None: + return sliding_window + return None + + +def _graph_max_cache_seq_len(graph: IRGraph) -> int | None: + max_cache_seq_len = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) + if max_cache_seq_len is not None: + return max_cache_seq_len + + providers = graph.meta.get("transpile_metadata_providers") + if isinstance(providers, dict): + for provider_meta in providers.values(): + if not isinstance(provider_meta, dict): + continue + max_cache_seq_len = _coerce_optional_int(provider_meta.get("max_cache_seq_len")) + if max_cache_seq_len is not None: + return max_cache_seq_len + return None + + +def _coerce_string_tuple(value: object) -> tuple[str, ...]: + if not isinstance(value, (tuple, list)): + return () + result: list[str] = [] + for item in value: + if not isinstance(item, str): + return () + result.append(item) + return tuple(result) + + +def _coerce_optional_int(value: object) -> int | None: + if value is None or isinstance(value, bool): + return None + if isinstance(value, int): + return int(value) + return None diff --git a/python/cactus/transpile/runtime_compat.py b/python/cactus/transpile/runtime_compat.py new file mode 100644 index 000000000..1fd53f0b4 --- /dev/null +++ b/python/cactus/transpile/runtime_compat.py @@ -0,0 +1,375 @@ +from __future__ import annotations + +import ctypes +import importlib +import sys +from typing import Any + +import numpy as np + + +class _MissingFFIFunction: + _cactus_missing_symbol = True + + def __init__(self, name: str): + self.__name__ = name + self.argtypes = None + self.restype = ctypes.c_int + + def __call__(self, *args: Any, **kwargs: Any) -> int: + raise RuntimeError(f"Cactus runtime is missing required symbol: {self.__name__}") + + +_ORIG_CDLL_GETATTR = ctypes.CDLL.__getattr__ + + +def _patched_cdll_getattr(self: ctypes.CDLL, name: str): + try: + return _ORIG_CDLL_GETATTR(self, name) + except AttributeError: + if not name.startswith("cactus_"): + raise + missing = _MissingFFIFunction(name) + setattr(self, name, missing) + return missing + + +def _load_runtime_module(): + if "cactus.bindings.cactus" in sys.modules: + return sys.modules["cactus.bindings.cactus"] + + ctypes.CDLL.__getattr__ = _patched_cdll_getattr + try: + return importlib.import_module("cactus.bindings.cactus") + finally: + ctypes.CDLL.__getattr__ = _ORIG_CDLL_GETATTR + + +def _patch_graph_runtime(cactus_module) -> None: + Graph = cactus_module.Graph + if getattr(Graph, "_transpile_runtime_compat_patched", False): + return + + _lib = cactus_module._lib + _err = cactus_module._err + cactus_node_t = cactus_module.cactus_node_t + + def _has_symbol(name: str) -> bool: + symbol = getattr(_lib, name, None) + return symbol is not None and not getattr(symbol, "_cactus_missing_symbol", False) + + def _ensure_compare_tensor(self, tensor): + tensor = self._ensure_tensor(tensor) + if int(tensor.dtype) == int(Graph.FP16): + return tensor + return self.precision_cast(tensor, Graph.FP16) + + def _ensure_scalar_tensor(self, tensor): + tensor = self._ensure_tensor(tensor) + if int(tensor.dtype) == int(Graph.FP16): + return tensor + return self.precision_cast(tensor, Graph.FP16) + + def _ensure_fp16_activation(self, tensor): + tensor = self._ensure_tensor(tensor) + if int(tensor.dtype) == int(Graph.FP32): + return self.precision_cast(tensor, Graph.FP16) + return tensor + + def _approx_nonzero_mask(self, tensor): + tensor = _ensure_scalar_tensor(self, tensor) + # Gemma4 compare paths in v2 only need stable 0/1-style masks for + # discrete values such as token-type ids and prebuilt boolean masks. + magnitude = self.abs(tensor) + shifted = self.scalar_add(magnitude, -0.5) + sharpened = self.scalar_multiply(shifted, 16.0) + return self.sigmoid(sharpened) + + orig_conv2d = Graph.conv2d + orig_not_equal = Graph.not_equal + orig_scalar_not_equal = Graph.scalar_not_equal + orig_transpose = Graph.transpose + orig_permute = Graph.permute + orig_scalar_add = Graph.scalar_add + orig_scalar_subtract = Graph.scalar_subtract + orig_scalar_multiply = Graph.scalar_multiply + orig_scalar_divide = Graph.scalar_divide + orig_scalar_exp = Graph.scalar_exp + orig_scalar_sqrt = Graph.scalar_sqrt + orig_scalar_cos = Graph.scalar_cos + orig_scalar_sin = Graph.scalar_sin + orig_scalar_log = Graph.scalar_log + orig_clamp = getattr(Graph, "clamp", None) + orig_add = Graph.add + orig_subtract = Graph.subtract + orig_multiply = Graph.multiply + orig_divide = Graph.divide + orig_concat = Graph.concat + orig_cat = Graph.cat + orig_abs = Graph.abs + orig_pow = Graph.pow + orig_expand = Graph.expand + + def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None): + a = _ensure_fp16_activation(self, a) + b = _ensure_fp16_activation(self, b) + out = cactus_node_t() + rc = _lib.cactus_graph_matmul( + self.h, + cactus_node_t(a.id), + cactus_node_t(b.id), + ctypes.c_bool(bool(pretransposed_rhs)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_matmul failed")) + result = self._apply_backend(self._tensor_from_node(out.value), backend) + if output_dtype is not None and int(output_dtype) != int(result.dtype): + result = self.precision_cast(result, int(output_dtype)) + return result + + def gather(self, tensor, indices, axis=0, backend=None): + tensor = self._ensure_tensor(tensor) + indices = self._ensure_tensor(indices) + if int(axis) != 0: + raise NotImplementedError( + f"transpiler runtime compatibility only supports gather(axis=0), got axis={axis}" + ) + out = cactus_node_t() + rc = _lib.cactus_graph_gather( + self.h, + cactus_node_t(tensor.id), + cactus_node_t(indices.id), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError("graph_gather failed") + return self._apply_backend(self._tensor_from_node(out.value), backend) + + def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, backend=None): + def _pair(value: Any) -> tuple[int, int]: + if isinstance(value, (tuple, list)): + if len(value) != 2: + raise ValueError(f"expected pair for conv2d parameter, got {value!r}") + return int(value[0]), int(value[1]) + return int(value), int(value) + + stride_hw = _pair(stride) + padding_hw = _pair(padding) + dilation_hw = _pair(dilation) + groups_int = int(groups) + kernel_hw = tuple(int(dim) for dim in self._ensure_tensor(weight).shape[-2:]) + + if _has_symbol("cactus_graph_conv2d"): + return orig_conv2d( + self, + x, + weight, + bias=bias, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + backend=backend, + ) + + if ( + kernel_hw == (3, 3) + and stride_hw == (2, 2) + and padding_hw == (1, 1) + and dilation_hw == (1, 1) + ): + if groups_int == 1: + return self.conv2d_k3s2p1(x, weight, bias=bias, backend=backend) + if len(weight.shape) >= 1 and groups_int == int(weight.shape[0]): + return self.conv2d_depthwise_k3s2p1(x, weight, bias=bias, backend=backend) + + if ( + kernel_hw == (3, 3) + and stride_hw == (1, 1) + and padding_hw == (1, 1) + and dilation_hw == (1, 1) + and groups_int == 1 + and _has_symbol("cactus_graph_conv2d_k3s1p1") + ): + return self.conv2d_k3s1p1(x, weight, bias=bias, backend=backend) + + if ( + kernel_hw == (1, 1) + and stride_hw == (1, 1) + and padding_hw == (0, 0) + and dilation_hw == (1, 1) + and groups_int == 1 + ): + return self.conv2d_pointwise_1x1(x, weight, bias=bias, backend=backend) + + raise NotImplementedError( + "v2 runtime does not expose generic conv2d for this configuration: " + f"kernel={kernel_hw} stride={stride_hw} padding={padding_hw} " + f"dilation={dilation_hw} groups={groups_int}" + ) + + def scalar_not_equal(self, x, value, backend=None): + if _has_symbol("cactus_graph_scalar_not_equal"): + return orig_scalar_not_equal(self, x, value, backend=backend) + x = _ensure_compare_tensor(self, x) + delta = self.scalar_add(x, -float(value)) + return _approx_nonzero_mask(self, delta) + + def not_equal(self, a, b, backend=None): + if _has_symbol("cactus_graph_not_equal"): + return orig_not_equal(self, a, b, backend=backend) + a = _ensure_compare_tensor(self, a) + b = _ensure_compare_tensor(self, b) + delta = self.subtract(a, b) + return _approx_nonzero_mask(self, delta) + + def transpose(self, x, backend=None): + return orig_transpose(self, _ensure_scalar_tensor(self, x), backend=backend) + + def permute(self, x, permutation, backend=None): + return orig_permute(self, _ensure_scalar_tensor(self, x), permutation, backend=backend) + + def scalar_add(self, x, value, backend=None): + return orig_scalar_add(self, self._ensure_tensor(x), value, backend=backend) + + def scalar_subtract(self, x, value, backend=None): + return orig_scalar_subtract(self, self._ensure_tensor(x), value, backend=backend) + + def scalar_multiply(self, x, value, backend=None): + return orig_scalar_multiply(self, self._ensure_tensor(x), value, backend=backend) + + def scalar_divide(self, x, value, backend=None): + return orig_scalar_divide(self, self._ensure_tensor(x), value, backend=backend) + + def scalar_exp(self, x, backend=None): + return orig_scalar_exp(self, self._ensure_tensor(x), backend=backend) + + def scalar_sqrt(self, x, backend=None): + return orig_scalar_sqrt(self, self._ensure_tensor(x), backend=backend) + + def scalar_cos(self, x, backend=None): + return orig_scalar_cos(self, self._ensure_tensor(x), backend=backend) + + def scalar_sin(self, x, backend=None): + return orig_scalar_sin(self, self._ensure_tensor(x), backend=backend) + + def scalar_log(self, x, backend=None): + return orig_scalar_log(self, self._ensure_tensor(x), backend=backend) + + def clamp(self, x, lo, hi, backend=None): + if orig_clamp is None: + raise RuntimeError("Cactus runtime is missing required symbol: cactus_graph_clamp") + return orig_clamp(self, _ensure_scalar_tensor(self, x), lo, hi, backend=backend) + + def add(self, a, b, backend=None): + return orig_add(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) + + def subtract(self, a, b, backend=None): + return orig_subtract(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) + + def multiply(self, a, b, backend=None): + return orig_multiply(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) + + def divide(self, a, b, backend=None): + return orig_divide(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) + + def concat(self, a, b, axis=0, backend=None): + return orig_concat(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), axis=axis, backend=backend) + + def cat(self, tensors, axis=0, backend=None): + legalized = [_ensure_scalar_tensor(self, tensor) for tensor in tensors] + return orig_cat(self, legalized, axis=axis, backend=backend) + + def abs(self, x, backend=None): + return orig_abs(self, self._ensure_tensor(x), backend=backend) + + def pow(self, x, exponent, backend=None): + return orig_pow(self, self._ensure_tensor(x), exponent, backend=backend) + + def expand(self, x, shape, backend=None): + x = self._ensure_tensor(x) + shape = tuple(int(v) for v in shape) + if tuple(int(dim) for dim in x.shape) == shape: + return x + if _has_symbol("cactus_graph_expand"): + return orig_expand(self, x, shape, backend=backend) + + # v2 builds do not always expose the generic expand symbol. For + # broadcast-only expands, adding an embedded zero tensor with the target + # shape gives the same logical result while staying inside existing + # Cactus graph ops. + dtype_map = { + int(Graph.FP16): np.float16, + int(Graph.FP32): np.float32, + } + x_dtype = int(x.dtype) + if x_dtype not in dtype_map: + raise NotImplementedError( + "expand fallback only supports floating-point tensors; " + f"got dtype={x_dtype} shape={tuple(int(dim) for dim in x.shape)} -> {shape}" + ) + zero = self.input(shape, dtype=x_dtype) + self.set_input(zero, np.zeros(shape, dtype=dtype_map[x_dtype]), dtype=x_dtype) + self.mark_embedded_input(zero) + return self.add(x, zero, backend=backend) + + Graph.matmul = matmul + Graph.gather = gather + Graph.conv2d = conv2d + Graph.scalar_not_equal = scalar_not_equal + Graph.not_equal = not_equal + Graph.transpose = transpose + Graph.permute = permute + Graph.scalar_add = scalar_add + Graph.scalar_subtract = scalar_subtract + Graph.scalar_multiply = scalar_multiply + Graph.scalar_divide = scalar_divide + Graph.scalar_exp = scalar_exp + Graph.scalar_sqrt = scalar_sqrt + Graph.scalar_cos = scalar_cos + Graph.scalar_sin = scalar_sin + Graph.scalar_log = scalar_log + Graph.clamp = clamp + Graph.add = add + Graph.subtract = subtract + Graph.multiply = multiply + Graph.divide = divide + Graph.concat = concat + Graph.cat = cat + Graph.abs = abs + Graph.pow = pow + Graph.expand = expand + Graph._transpile_runtime_compat_patched = True + + +_cactus_module = _load_runtime_module() + +if hasattr(_cactus_module, "_lib"): + _lib_obj = _cactus_module._lib + if hasattr(_lib_obj, "cactus_graph_matmul"): + _lib_obj.cactus_graph_matmul.argtypes = [ + ctypes.c_void_p, + _cactus_module.cactus_node_t, + _cactus_module.cactus_node_t, + ctypes.c_bool, + ctypes.POINTER(_cactus_module.cactus_node_t), + ] + _lib_obj.cactus_graph_matmul.restype = ctypes.c_int + if hasattr(_lib_obj, "cactus_graph_gather"): + _lib_obj.cactus_graph_gather.argtypes = [ + ctypes.c_void_p, + _cactus_module.cactus_node_t, + _cactus_module.cactus_node_t, + ctypes.POINTER(_cactus_module.cactus_node_t), + ] + _lib_obj.cactus_graph_gather.restype = ctypes.c_int + +_patch_graph_runtime(_cactus_module) + +_lib = _cactus_module._lib +_err = _cactus_module._err +cactus_node_t = _cactus_module.cactus_node_t +cactus_tensor_info_t = _cactus_module.cactus_tensor_info_t +Graph = _cactus_module.Graph +Tensor = _cactus_module.Tensor diff --git a/python/cactus/transpile/runtime_support.py b/python/cactus/transpile/runtime_support.py new file mode 100644 index 000000000..8b966f77d --- /dev/null +++ b/python/cactus/transpile/runtime_support.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import builtins +import importlib.util +import sys +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path + +import torch + +from cactus.transpile.model_profiles import GEMMA4_PROFILE +from cactus.transpile.model_profiles import ModelProfile +from cactus.transpile.model_profiles import profile_for_model_type + + +_TORCHVISION_COMPAT_LIBRARIES: list[object] = [] + + +@dataclass +class PreparedInputs: + names: tuple[str, ...] + tensors: tuple[torch.Tensor, ...] + metadata: dict[str, object] + + +def _transformers_supports_model_module(module_name: str) -> bool: + try: + return importlib.util.find_spec(module_name) is not None + except Exception: + return False + + +def _candidate_external_site_packages() -> list[Path]: + candidates: list[Path] = [] + major_minor = f"python{sys.version_info.major}.{sys.version_info.minor}" + pyenv_versions = Path.home() / ".pyenv" / "versions" + if not pyenv_versions.exists(): + return candidates + for version_dir in sorted(pyenv_versions.iterdir(), reverse=True): + site_packages = version_dir / "lib" / major_minor / "site-packages" + if site_packages.exists(): + candidates.append(site_packages) + return candidates + + +def ensure_transformers_supports_profile(profile: ModelProfile | None) -> str | None: + target_module = None if profile is None else profile.transformer_module + if not target_module or _transformers_supports_model_module(target_module): + return None + + for site_packages in _candidate_external_site_packages(): + candidate = site_packages / Path(*target_module.split(".")) + if not candidate.with_suffix(".py").exists() and not candidate.is_dir(): + continue + sys.path.insert(0, str(site_packages)) + for module_name in list(sys.modules): + root_name = module_name.split(".", 1)[0] + if root_name in {"transformers", "huggingface_hub", "tokenizers"}: + del sys.modules[module_name] + if _transformers_supports_model_module(target_module): + return str(site_packages) + try: + sys.path.remove(str(site_packages)) + except ValueError: + pass + return None + + +def ensure_transformers_supports_model_type(model_type: str) -> str | None: + return ensure_transformers_supports_profile(profile_for_model_type(model_type)) + + +def ensure_transformers_supports_gemma4() -> str | None: + return ensure_transformers_supports_profile(GEMMA4_PROFILE) + + +def _patch_torchvision_missing_nms_op() -> str | None: + """Keep mismatched torch/torchvision installs importable for processors.""" + + try: + import torchvision # type: ignore # noqa: F401 + return None + except RuntimeError as exc: + if "operator torchvision::nms does not exist" not in str(exc): + return None + except Exception: + return None + + try: + import torchvision.extension as tv_extension # type: ignore + + if bool(getattr(tv_extension, "_HAS_OPS", False)): + return None + except Exception: + pass + + try: + library = torch.library.Library("torchvision", "DEF") + library.define("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor") + _TORCHVISION_COMPAT_LIBRARIES.append(library) + return "defined missing torchvision::nms operator for torchvision import compatibility" + except Exception: + return None + + +def patch_transformers_torchvision_probe() -> str | None: + has_torchvision = importlib.util.find_spec("torchvision") is not None + has_lzma = importlib.util.find_spec("_lzma") is not None + + if not has_torchvision: + return None + + if has_lzma: + nms_patch_note = _patch_torchvision_missing_nms_op() + return nms_patch_note + + base_note: str | None = None + try: + import backports.lzma as backports_lzma # type: ignore + + sys.modules.setdefault("lzma", backports_lzma) + base_note = "using backports.lzma because this Python build is missing _lzma" + nms_patch_note = _patch_torchvision_missing_nms_op() + if nms_patch_note: + return f"{base_note}; {nms_patch_note}" + return base_note + except Exception: + pass + + class _InterpolationModeStub: + NEAREST = "nearest" + BILINEAR = "bilinear" + BICUBIC = "bicubic" + LANCZOS = "lanczos" + + class _TorchvisionFunctionalStub: + InterpolationMode = _InterpolationModeStub + + def __getattr__(self, name: str): + raise RuntimeError( + "torchvision functionality is unavailable because this Python build is missing _lzma; " + f"attempted to access torchvision.transforms.functional.{name}" + ) + + builtins.F = _TorchvisionFunctionalStub() + builtins.tvF = builtins.F + + import transformers.utils as tf_utils # type: ignore + import transformers.utils.import_utils as tf_import_utils # type: ignore + + @lru_cache + def _disabled() -> bool: + return False + + tf_import_utils.is_torchvision_available = _disabled + tf_import_utils.is_torchvision_v2_available = _disabled + tf_utils.is_torchvision_available = _disabled + tf_utils.is_torchvision_v2_available = _disabled + return "disabled torchvision import checks because this Python build is missing _lzma" + + +def patch_torch_flex_attention_compat() -> str | None: + try: + import torch.nn.attention.flex_attention as flex_attention # type: ignore + except Exception: + return None + + if hasattr(flex_attention, "AuxRequest"): + return None + + class _AuxRequest: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + flex_attention.AuxRequest = _AuxRequest # type: ignore[attr-defined] + return "installed torch flex_attention AuxRequest compatibility stub" diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py new file mode 100644 index 000000000..9f0b25dfb --- /dev/null +++ b/python/cactus/transpile/tdt_runtime.py @@ -0,0 +1,893 @@ +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import replace +import json +from pathlib import Path +import re +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from cactus.transpile.component_pipeline import ComponentModuleSpec +from cactus.transpile.audio_preprocess import prepare_native_parakeet_audio_features +from cactus.transpile.model_profiles import add_tensor_aliases +from cactus.transpile.model_profiles import PARAKEET_TDT_PROFILE + + +def _cfg_get(config: dict[str, Any], key: str, default: Any = None) -> Any: + value = config.get(key, default) + return default if value is None else value + + +def _load_tensor_state_dict(model_source: str) -> dict[str, torch.Tensor]: + root = Path(model_source) + safetensors_path = root / "model.safetensors" + if safetensors_path.exists(): + from safetensors.torch import load_file + + return _with_parakeet_tdt_aliases(dict(load_file(str(safetensors_path)))) + + bin_path = root / "pytorch_model.bin" + if bin_path.exists(): + loaded = torch.load(bin_path, map_location="cpu") + if isinstance(loaded, dict): + return _with_parakeet_tdt_aliases({ + str(key): value + for key, value in loaded.items() + if isinstance(value, torch.Tensor) + }) + raise RuntimeError(f"unsupported Parakeet TDT checkpoint format in {model_source}") + + +def _add_tdt_derived_aliases(state_dict: dict[str, torch.Tensor]) -> None: + def alias(target: str, source: str) -> None: + if target not in state_dict and source in state_dict: + state_dict[target] = state_dict[source] + + for index in range(4): + alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wx", f"decoder.lstm.weight_ih_l{index}") + alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wh", f"decoder.lstm.weight_hh_l{index}") + alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wx", f"decoder.prediction.dec_rnn.lstm.weight_ih_l{index}") + alias(f"decoder.prediction.dec_rnn.lstm.{index}.Wh", f"decoder.prediction.dec_rnn.lstm.weight_hh_l{index}") + bias_key = f"decoder.prediction.dec_rnn.lstm.{index}.bias" + bias_ih = state_dict.get(f"decoder.lstm.bias_ih_l{index}") + bias_hh = state_dict.get(f"decoder.lstm.bias_hh_l{index}") + if bias_ih is None: + bias_ih = state_dict.get(f"decoder.prediction.dec_rnn.lstm.bias_ih_l{index}") + if bias_hh is None: + bias_hh = state_dict.get(f"decoder.prediction.dec_rnn.lstm.bias_hh_l{index}") + if bias_key not in state_dict and bias_ih is not None and bias_hh is not None: + state_dict[bias_key] = bias_ih + bias_hh + + for source_index, target_index in ((0, 0), (2, 2), (3, 3), (5, 5), (6, 6)): + alias( + f"encoder.pre_encode.conv.{target_index}.weight", + f"encoder.subsampling.layers.{source_index}.weight", + ) + alias( + f"encoder.pre_encode.conv.{target_index}.bias", + f"encoder.subsampling.layers.{source_index}.bias", + ) + alias("encoder.pre_encode.out.weight", "encoder.subsampling.linear.weight") + alias("encoder.pre_encode.out.bias", "encoder.subsampling.linear.bias") + + + +def _with_parakeet_tdt_aliases(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Apply profile-declared aliases needed by the local TDT fallback runtime.""" + + add_tensor_aliases( + state_dict, + PARAKEET_TDT_PROFILE, + derived_aliases=_add_tdt_derived_aliases, + ) + return state_dict + + +@dataclass +class ParakeetTDTConfig: + model_source: str + sample_rate: int + num_mel_bins: int + hidden_dim: int + num_layers: int + attention_heads: int + attention_head_dim: int + attention_scale: float + ff_intermediate_dim: int + conv_kernel_size: int + subsampling_factor: int + subsampling_conv_channels: int + predictor_hidden_dim: int + predictor_num_layers: int + joint_dim: int + num_tdt_durations: int + tdt_durations: tuple[int, ...] + blank_id: int + vocabulary: tuple[str, ...] + encoder_hidden_act: str + + +def prepare_parakeet_tdt_audio_features( + audio_file: str | Path, + *, + expected_frames: int | None, + expected_mels: int, + torch_dtype: torch.dtype, +) -> tuple[torch.Tensor, int]: + return prepare_native_parakeet_audio_features( + audio_file, + expected_frames=expected_frames, + expected_mels=expected_mels, + torch_dtype=torch_dtype, + ) + + +def greedy_decode_parakeet_tdt_token_ids( + *, + config: ParakeetTDTConfig, + encoder_hidden_states: np.ndarray, + initial_states: tuple[np.ndarray, ...], + step, +) -> list[int]: + hidden = np.ascontiguousarray(np.asarray(encoder_hidden_states)) + if hidden.ndim != 3 or hidden.shape[0] != 1: + raise ValueError( + f"expected encoder_hidden_states with shape [1, T, D], got {tuple(hidden.shape)}" + ) + predictor_layers = int(config.predictor_num_layers) + if len(initial_states) != predictor_layers * 2: + raise ValueError( + f"expected {predictor_layers * 2} initial state tensors, got {len(initial_states)}" + ) + + states = tuple(np.ascontiguousarray(np.asarray(state)) for state in initial_states) + durations = [int(value) for value in config.tdt_durations] + if not durations and int(config.num_tdt_durations) > 0: + durations = list(range(int(config.num_tdt_durations))) + duration_class_count = max(len(durations), int(config.num_tdt_durations)) + last_token = _initial_parakeet_tdt_blank_id(config) + emitted: list[int] = [] + time_index = 0 + + while time_index < int(hidden.shape[1]): + frame = np.ascontiguousarray(hidden[:, time_index, :]) + advanced = False + symbols_added = 0 + + while symbols_added < 10: + logits, next_states = step(frame, last_token, states) + logits_array = np.asarray(logits, dtype=np.float32) + total_classes = int(logits_array.shape[-1]) + active_duration_count = duration_class_count + if active_duration_count <= 0 or active_duration_count >= total_classes: + active_duration_count = max(1, min(total_classes - 1, int(config.num_tdt_durations) or len(durations) or 1)) + if len(durations) != active_duration_count: + durations = list(range(active_duration_count)) + token_class_count = total_classes - active_duration_count + blank_id = _effective_parakeet_tdt_blank_id( + config, + token_class_count=token_class_count, + ) + if last_token < 0 or last_token >= token_class_count: + last_token = blank_id + token_scores = logits_array[:, :token_class_count] + duration_scores = logits_array[:, token_class_count:] + next_token = int(np.argmax(token_scores[0])) + duration_index = int(np.argmax(duration_scores[0])) + skip = int(durations[min(duration_index, len(durations) - 1)]) if durations else 1 + + if next_token != blank_id: + emitted.append(next_token) + last_token = next_token + states = tuple(np.ascontiguousarray(np.asarray(state).copy()) for state in next_states) + + symbols_added += 1 + + if skip > 0: + time_index += skip + advanced = True + break + if next_token == blank_id: + time_index += 1 + advanced = True + break + + if not advanced: + time_index += 1 + + return emitted + + +def _initial_parakeet_tdt_blank_id(config: ParakeetTDTConfig) -> int: + blank_id = int(config.blank_id) + vocab_size = len(config.vocabulary) + if vocab_size > 0 and blank_id == vocab_size - 1: + return vocab_size + return blank_id + + +def _effective_parakeet_tdt_blank_id( + config: ParakeetTDTConfig, + *, + token_class_count: int, +) -> int: + blank_id = int(config.blank_id) + vocab_size = len(config.vocabulary) + if vocab_size > 0 and blank_id == vocab_size - 1 and token_class_count == vocab_size + 1: + return token_class_count - 1 + if 0 <= blank_id < token_class_count: + return blank_id + if vocab_size > 0 and token_class_count == vocab_size + 1: + return token_class_count - 1 + return max(0, token_class_count - 1) + + +def load_parakeet_tdt_config(model_source: str) -> ParakeetTDTConfig: + root_path = Path(model_source) + config_path = root_path / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"missing config.json for Parakeet TDT: {config_path}") + root = json.loads(config_path.read_text()) + encoder = root.get("encoder") or root.get("encoder_config") or {} + decoder = root.get("decoder") or {} + prediction = decoder.get("prediction") or decoder.get("prednet") or {} + joint = root.get("joint") or {} + jointnet = joint.get("jointnet") or {} + model_defaults = root.get("model_defaults") or {} + preprocessor = root.get("preprocessor") or {} + + hidden_dim = int(_cfg_get(root, "hidden_dim", _cfg_get(encoder, "d_model", _cfg_get(encoder, "hidden_size", 0)))) + attention_heads = int(_cfg_get(encoder, "n_heads", _cfg_get(encoder, "num_attention_heads", 0))) + attention_head_dim = hidden_dim // max(attention_heads, 1) + tdt_durations = tuple( + int(value) + for value in _cfg_get(root, "tdt_durations", _cfg_get(root, "durations", _cfg_get(model_defaults, "tdt_durations", (0, 1, 2, 3, 4)))) + ) + vocabulary = tuple(str(value) for value in _cfg_get(joint, "vocabulary", ())) + if not vocabulary: + vocabulary = _load_parakeet_vocabulary(root_path) + decoder_vocab_size = int(_cfg_get(root, "vocab_size", _cfg_get(decoder, "vocab_size", len(vocabulary)))) + decoding = root.get("decoding") or {} + blank_id_raw = _cfg_get( + root, + "tdt_blank_id", + _cfg_get(root, "blank_token_id", _cfg_get(decoding, "blank_id", _cfg_get(decoder, "blank_id", None))), + ) + if blank_id_raw is None: + blank_id = decoder_vocab_size + else: + try: + blank_id = int(blank_id_raw) + except (TypeError, ValueError): + blank_id = decoder_vocab_size + if blank_id < 0: + blank_id = decoder_vocab_size + + return ParakeetTDTConfig( + model_source=model_source, + sample_rate=int(_cfg_get(preprocessor, "sample_rate", 16000)), + num_mel_bins=int(_cfg_get(root, "num_mel_bins", _cfg_get(preprocessor, "features", _cfg_get(encoder, "feat_in", _cfg_get(encoder, "num_mel_bins", 128))))), + hidden_dim=hidden_dim, + num_layers=int(_cfg_get(root, "num_layers", _cfg_get(encoder, "n_layers", _cfg_get(encoder, "num_hidden_layers", 0)))), + attention_heads=attention_heads, + attention_head_dim=attention_head_dim, + attention_scale=float(attention_head_dim ** -0.5 if attention_head_dim > 0 else 1.0), + ff_intermediate_dim=int( + _cfg_get( + root, + "ffn_intermediate_dim", + _cfg_get( + encoder, + "ffn_hidden_size", + _cfg_get(encoder, "intermediate_size", round(hidden_dim * float(_cfg_get(encoder, "ff_expansion_factor", 4.0)))), + ), + ) + ), + conv_kernel_size=int(_cfg_get(root, "conv_kernel_size", _cfg_get(encoder, "conv_kernel_size", 9))), + subsampling_factor=int(_cfg_get(root, "subsampling_factor", _cfg_get(encoder, "subsampling_factor", 8))), + subsampling_conv_channels=int( + _cfg_get(root, "subsampling_conv_channels", _cfg_get(encoder, "subsampling_conv_channels", 256)) + ), + predictor_hidden_dim=int( + _cfg_get(root, "predictor_hidden_dim", _cfg_get(root, "decoder_hidden_size", _cfg_get(prediction, "pred_hidden", _cfg_get(model_defaults, "pred_hidden", 640)))) + ), + predictor_num_layers=int( + _cfg_get(root, "predictor_num_layers", _cfg_get(root, "num_decoder_layers", _cfg_get(prediction, "pred_rnn_layers", 1))) + ), + joint_dim=int(_cfg_get(root, "tdt_joint_dim", _cfg_get(jointnet, "joint_hidden", _cfg_get(model_defaults, "joint_hidden", 640)))), + num_tdt_durations=int(_cfg_get(root, "tdt_num_durations", _cfg_get(model_defaults, "num_tdt_durations", len(tdt_durations)))), + tdt_durations=tdt_durations, + blank_id=blank_id, + vocabulary=vocabulary, + encoder_hidden_act=str(_cfg_get(root, "encoder_hidden_act", _cfg_get(encoder, "activation", _cfg_get(encoder, "hidden_act", "silu")))).lower(), + ) + + +def _load_parakeet_vocabulary(root: Path) -> tuple[str, ...]: + vocab_txt = root / "vocab.txt" + if vocab_txt.exists(): + pieces: list[str] = [] + for line in vocab_txt.read_text(encoding="utf-8").splitlines(): + if "\t" in line: + _, token = line.split("\t", 1) + pieces.append(token) + elif line: + pieces.append(line) + if pieces: + return tuple(pieces) + + tokenizer_json = root / "tokenizer.json" + if tokenizer_json.exists(): + try: + loaded = json.loads(tokenizer_json.read_text(encoding="utf-8")) + except Exception: + return () + model = loaded.get("model") if isinstance(loaded, dict) else None + vocab = model.get("vocab") if isinstance(model, dict) else None + if isinstance(vocab, dict): + ordered = sorted( + ((int(index), str(token)) for token, index in vocab.items()), + key=lambda item: item[0], + ) + return tuple(token for _, token in ordered) + return () + + +def _copy_linear_weight(linear: nn.Linear, weight: torch.Tensor, *, bias: torch.Tensor | None = None) -> None: + linear.weight.data.copy_(weight.to(dtype=linear.weight.dtype, device=linear.weight.device)) + if linear.bias is not None: + if bias is None: + linear.bias.data.zero_() + else: + linear.bias.data.copy_(bias.to(dtype=linear.bias.dtype, device=linear.bias.device)) + + +def _copy_conv2d_weight(conv: nn.Conv2d, weight: torch.Tensor, *, bias: torch.Tensor | None = None) -> None: + tensor = weight + if tuple(tensor.shape) != tuple(conv.weight.shape): + tensor = tensor.permute(0, 3, 1, 2).contiguous() + conv.weight.data.copy_(tensor.to(dtype=conv.weight.dtype, device=conv.weight.device)) + if conv.bias is not None: + if bias is None: + conv.bias.data.zero_() + else: + conv.bias.data.copy_(bias.to(dtype=conv.bias.dtype, device=conv.bias.device)) + + +def _copy_conv1d_weight(conv: nn.Conv1d, weight: torch.Tensor, *, bias: torch.Tensor | None = None) -> None: + tensor = weight + if tensor.ndim == 3 and tensor.shape[1] != conv.in_channels // conv.groups: + tensor = tensor.permute(0, 2, 1).contiguous() + conv.weight.data.copy_(tensor.to(dtype=conv.weight.dtype, device=conv.weight.device)) + if conv.bias is not None: + if bias is None: + conv.bias.data.zero_() + else: + conv.bias.data.copy_(bias.to(dtype=conv.bias.dtype, device=conv.bias.device)) + + +def _apply_activation(x: torch.Tensor, activation: str) -> torch.Tensor: + if "gelu" in activation: + return F.gelu(x) + if activation == "relu": + return F.relu(x) + return F.silu(x) + + +def _relative_position_embeddings(*, seq_len: int, hidden_dim: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + rel_positions = torch.arange(seq_len - 1, -seq_len, -1, device=device, dtype=torch.float32) + half_dim = hidden_dim // 2 + inv_freq = 1.0 / ( + 10000.0 + ** (torch.arange(0, half_dim, device=device, dtype=torch.float32) * 2.0 / float(max(hidden_dim, 1))) + ) + angles = rel_positions.unsqueeze(1) * inv_freq.unsqueeze(0) + sin = torch.sin(angles) + cos = torch.cos(angles) + embeddings = torch.stack((sin, cos), dim=-1).reshape(2 * seq_len - 1, half_dim * 2) + if embeddings.shape[1] < hidden_dim: + pad = torch.zeros( + (embeddings.shape[0], hidden_dim - embeddings.shape[1]), + device=device, + dtype=embeddings.dtype, + ) + embeddings = torch.cat((embeddings, pad), dim=1) + elif embeddings.shape[1] > hidden_dim: + embeddings = embeddings[:, :hidden_dim] + return embeddings.to(dtype=dtype) + + +def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: + batch, heads, seq_len, _ = query.shape + scores = torch.matmul(query, relative_key.transpose(-1, -2)) + rel_index = ( + torch.arange(seq_len, device=query.device).view(seq_len, 1) + - torch.arange(seq_len, device=query.device).view(1, seq_len) + + (seq_len - 1) + ) + rel_index = rel_index.view(1, 1, seq_len, seq_len).expand(batch, heads, seq_len, seq_len) + gathered = scores.gather(-1, rel_index) + return gathered * float(scale) + + +class ParakeetTDTFeedForward(nn.Module): + def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.linear1 = nn.Linear(config.hidden_dim, config.ff_intermediate_dim, bias=f"{prefix}.linear1.bias" in state_dict) + self.linear2 = nn.Linear(config.ff_intermediate_dim, config.hidden_dim, bias=f"{prefix}.linear2.bias" in state_dict) + _copy_linear_weight( + self.linear1, + state_dict[f"{prefix}.linear1.weight"], + bias=state_dict.get(f"{prefix}.linear1.bias"), + ) + _copy_linear_weight( + self.linear2, + state_dict[f"{prefix}.linear2.weight"], + bias=state_dict.get(f"{prefix}.linear2.bias"), + ) + + def forward(self, x: torch.Tensor, *, activation: str) -> torch.Tensor: + return self.linear2(_apply_activation(self.linear1(x), activation)) + + +class ParakeetTDTSelfAttention(nn.Module): + def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.linear_q = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_q.bias" in state_dict) + self.linear_k = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_k.bias" in state_dict) + self.linear_v = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_v.bias" in state_dict) + self.linear_out = nn.Linear(config.hidden_dim, config.hidden_dim, bias=f"{prefix}.linear_out.bias" in state_dict) + self.linear_pos = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False) + _copy_linear_weight(self.linear_q, state_dict[f"{prefix}.linear_q.weight"], bias=state_dict.get(f"{prefix}.linear_q.bias")) + _copy_linear_weight(self.linear_k, state_dict[f"{prefix}.linear_k.weight"], bias=state_dict.get(f"{prefix}.linear_k.bias")) + _copy_linear_weight(self.linear_v, state_dict[f"{prefix}.linear_v.weight"], bias=state_dict.get(f"{prefix}.linear_v.bias")) + _copy_linear_weight( + self.linear_out, + state_dict[f"{prefix}.linear_out.weight"], + bias=state_dict.get(f"{prefix}.linear_out.bias"), + ) + _copy_linear_weight(self.linear_pos, state_dict[f"{prefix}.linear_pos.weight"]) + self.pos_bias_u = nn.Parameter(state_dict[f"{prefix}.pos_bias_u"].clone()) + self.pos_bias_v = nn.Parameter(state_dict[f"{prefix}.pos_bias_v"].clone()) + self.config = config + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch, seq_len, _ = x.shape + num_heads = self.config.attention_heads + head_dim = self.config.attention_head_dim + + q = self.linear_q(x).view(batch, seq_len, num_heads, head_dim) + k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim) + v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim) + + rel_pos = _relative_position_embeddings( + seq_len=seq_len, + hidden_dim=self.config.hidden_dim, + device=x.device, + dtype=x.dtype, + ) + rel_k = self.linear_pos(rel_pos).view(1, 2 * seq_len - 1, num_heads, head_dim) + + q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) + q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) + + q_u_heads = q_u.permute(0, 2, 1, 3) + q_v_heads = q_v.permute(0, 2, 1, 3) + k_heads = k.permute(0, 2, 1, 3) + v_heads = v.permute(0, 2, 1, 3) + rel_k_heads = rel_k.permute(0, 2, 1, 3) + + rel_bias = _relative_position_bias( + q_v_heads, + rel_k_heads, + scale=self.config.attention_scale, + ) + attn = F.scaled_dot_product_attention( + q_u_heads, + k_heads, + v_heads, + attn_mask=rel_bias, + dropout_p=0.0, + is_causal=False, + ) + attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim) + return self.linear_out(attn) + + +class ParakeetTDTConformerConv(nn.Module): + def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): + super().__init__() + hidden_dim = config.hidden_dim + kernel = config.conv_kernel_size + self.pointwise_conv1 = nn.Conv1d( + hidden_dim, + hidden_dim * 2, + kernel_size=1, + bias=f"{prefix}.pointwise_conv1.bias" in state_dict, + ) + self.depthwise_conv = nn.Conv1d( + hidden_dim, + hidden_dim, + kernel_size=kernel, + padding=kernel // 2, + groups=hidden_dim, + bias=f"{prefix}.depthwise_conv.bias" in state_dict, + ) + self.batch_norm = nn.BatchNorm1d(hidden_dim, eps=1e-5, affine=True, track_running_stats=True) + self.pointwise_conv2 = nn.Conv1d( + hidden_dim, + hidden_dim, + kernel_size=1, + bias=f"{prefix}.pointwise_conv2.bias" in state_dict, + ) + _copy_conv1d_weight( + self.pointwise_conv1, + state_dict[f"{prefix}.pointwise_conv1.weight"], + bias=state_dict.get(f"{prefix}.pointwise_conv1.bias"), + ) + _copy_conv1d_weight( + self.depthwise_conv, + state_dict[f"{prefix}.depthwise_conv.weight"], + bias=state_dict.get(f"{prefix}.depthwise_conv.bias"), + ) + _copy_conv1d_weight( + self.pointwise_conv2, + state_dict[f"{prefix}.pointwise_conv2.weight"], + bias=state_dict.get(f"{prefix}.pointwise_conv2.bias"), + ) + self.batch_norm.weight.data.copy_(state_dict[f"{prefix}.batch_norm.weight"].to(dtype=self.batch_norm.weight.dtype)) + self.batch_norm.bias.data.copy_(state_dict[f"{prefix}.batch_norm.bias"].to(dtype=self.batch_norm.bias.dtype)) + self.batch_norm.running_mean.data.copy_( + state_dict[f"{prefix}.batch_norm.running_mean"].to(dtype=self.batch_norm.running_mean.dtype) + ) + self.batch_norm.running_var.data.copy_( + state_dict[f"{prefix}.batch_norm.running_var"].to(dtype=self.batch_norm.running_var.dtype) + ) + self.config = config + + def forward(self, x: torch.Tensor, pad_mask: torch.Tensor | None = None) -> torch.Tensor: + x = x.transpose(1, 2) + x = self.pointwise_conv1(x) + x = F.glu(x, dim=1) + if pad_mask is not None: + x = x * pad_mask + x = self.depthwise_conv(x) + x = self.batch_norm(x) + x = _apply_activation(x, self.config.encoder_hidden_act) + x = self.pointwise_conv2(x) + return x.transpose(1, 2) + + +class ParakeetTDTEncoderLayer(nn.Module): + def __init__(self, config: ParakeetTDTConfig, layer_index: int, state_dict: dict[str, torch.Tensor]): + super().__init__() + prefix = f"encoder.layers.{layer_index}" + self.feed_forward1 = ParakeetTDTFeedForward(config, f"{prefix}.feed_forward1", state_dict) + self.feed_forward2 = ParakeetTDTFeedForward(config, f"{prefix}.feed_forward2", state_dict) + self.self_attn = ParakeetTDTSelfAttention(config, f"{prefix}.self_attn", state_dict) + self.conv = ParakeetTDTConformerConv(config, f"{prefix}.conv", state_dict) + self.norm_feed_forward1 = nn.LayerNorm(config.hidden_dim) + self.norm_self_att = nn.LayerNorm(config.hidden_dim) + self.norm_conv = nn.LayerNorm(config.hidden_dim) + self.norm_feed_forward2 = nn.LayerNorm(config.hidden_dim) + self.norm_out = nn.LayerNorm(config.hidden_dim) + for name in ("norm_feed_forward1", "norm_self_att", "norm_conv", "norm_feed_forward2", "norm_out"): + module = getattr(self, name) + module.weight.data.copy_(state_dict[f"{prefix}.{name}.weight"].to(dtype=module.weight.dtype)) + module.bias.data.copy_(state_dict[f"{prefix}.{name}.bias"].to(dtype=module.bias.dtype)) + self.config = config + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) + x = x + self.self_attn(self.norm_self_att(x)) + x = x + self.conv(self.norm_conv(x)) + x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) + return self.norm_out(x) + + +class ParakeetTDTPreEncode(nn.Module): + def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + channels = config.subsampling_conv_channels + self.conv = nn.ModuleList( + [ + nn.Conv2d(1, channels, kernel_size=3, stride=2, padding=1), + nn.Identity(), + nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1, groups=channels), + nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0), + nn.Identity(), + nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1, groups=channels), + nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0), + ] + ) + _copy_conv2d_weight(self.conv[0], state_dict["encoder.pre_encode.conv.0.weight"], bias=state_dict["encoder.pre_encode.conv.0.bias"]) + _copy_conv2d_weight(self.conv[2], state_dict["encoder.pre_encode.conv.2.weight"], bias=state_dict["encoder.pre_encode.conv.2.bias"]) + _copy_conv2d_weight(self.conv[3], state_dict["encoder.pre_encode.conv.3.weight"], bias=state_dict["encoder.pre_encode.conv.3.bias"]) + _copy_conv2d_weight(self.conv[5], state_dict["encoder.pre_encode.conv.5.weight"], bias=state_dict["encoder.pre_encode.conv.5.bias"]) + _copy_conv2d_weight(self.conv[6], state_dict["encoder.pre_encode.conv.6.weight"], bias=state_dict["encoder.pre_encode.conv.6.bias"]) + projected_width = config.num_mel_bins + for _ in range(3): + projected_width = (projected_width + 2 - 3) // 2 + 1 + self.out = nn.Linear(channels * projected_width, config.hidden_dim) + _copy_linear_weight( + self.out, + state_dict["encoder.pre_encode.out.weight"], + bias=state_dict["encoder.pre_encode.out.bias"], + ) + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + if input_features.ndim != 3: + raise ValueError(f"expected input_features [batch, frames, mels], got {tuple(input_features.shape)}") + x = input_features.unsqueeze(1) + x = F.relu(self.conv[0](x)) + x = F.relu(self.conv[3](self.conv[2](x))) + x = F.relu(self.conv[6](self.conv[5](x))) + x = x.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[2], -1) + return self.out(x) + + +class ParakeetTDTEncoder(nn.Module): + def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.pre_encode = ParakeetTDTPreEncode(config, state_dict) + self.layers = nn.ModuleList( + [ParakeetTDTEncoderLayer(config, index, state_dict) for index in range(config.num_layers)] + ) + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + x = self.pre_encode(input_features) + for layer in self.layers: + x = layer(x) + return x + + +class ParakeetTDTDecoderCell(nn.Module): + def __init__(self, hidden_dim: int, prefix: str, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.Wx = nn.Parameter(state_dict[f"{prefix}.Wx"].clone()) + self.Wh = nn.Parameter(state_dict[f"{prefix}.Wh"].clone()) + self.bias = nn.Parameter(state_dict[f"{prefix}.bias"].clone()) + self.hidden_dim = hidden_dim + + def forward(self, x: torch.Tensor, h: torch.Tensor, c: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + gates = F.linear(x, self.Wx) + F.linear(h, self.Wh) + self.bias + i, f, g, o = torch.chunk(gates, 4, dim=-1) + i = torch.sigmoid(i) + f = torch.sigmoid(f) + g = torch.tanh(g) + o = torch.sigmoid(o) + c_next = f * c + i * g + h_next = o * torch.tanh(c_next) + return h_next, c_next + + +class ParakeetTDTDecoderPrediction(nn.Module): + def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + vocab_size = max(config.blank_id + 1, state_dict["decoder.prediction.embed.weight"].shape[0]) + self.embed = nn.Embedding(vocab_size, config.predictor_hidden_dim) + self.embed.weight.data.copy_(state_dict["decoder.prediction.embed.weight"].to(dtype=self.embed.weight.dtype)) + self.dec_rnn = nn.Module() + self.dec_rnn.lstm = nn.ModuleList( + [ + ParakeetTDTDecoderCell( + config.predictor_hidden_dim, + f"decoder.prediction.dec_rnn.lstm.{index}", + state_dict, + ) + for index in range(config.predictor_num_layers) + ] + ) + + def forward( + self, + token_ids: torch.Tensor, + state_h: tuple[torch.Tensor, ...], + state_c: tuple[torch.Tensor, ...], + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + x = self.embed(token_ids) + if x.ndim == 3 and x.shape[1] == 1: + x = x[:, 0, :] + next_h: list[torch.Tensor] = [] + next_c: list[torch.Tensor] = [] + for index, cell in enumerate(self.dec_rnn.lstm): + h_new, c_new = cell(x, state_h[index], state_c[index]) + next_h.append(h_new) + next_c.append(c_new) + x = h_new + return x, tuple(next_h), tuple(next_c) + + +class ParakeetTDTJoint(nn.Module): + def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.enc = nn.Linear(config.hidden_dim, config.joint_dim) + self.pred = nn.Linear(config.predictor_hidden_dim, config.joint_dim) + output_classes = int(state_dict["joint.joint_net.2.weight"].shape[0]) + self.joint_net = nn.Sequential( + nn.Identity(), + nn.ReLU(), + nn.Linear( + config.joint_dim, + output_classes, + ), + ) + _copy_linear_weight(self.enc, state_dict["joint.enc.weight"], bias=state_dict["joint.enc.bias"]) + _copy_linear_weight(self.pred, state_dict["joint.pred.weight"], bias=state_dict["joint.pred.bias"]) + _copy_linear_weight( + self.joint_net[2], + state_dict["joint.joint_net.2.weight"], + bias=state_dict["joint.joint_net.2.bias"], + ) + + def forward(self, encoder_frame: torch.Tensor, predictor_hidden: torch.Tensor) -> torch.Tensor: + return self.joint_net(F.relu(self.enc(encoder_frame) + self.pred(predictor_hidden))) + + +class ParakeetTDTDecoderStep(nn.Module): + def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.prediction = ParakeetTDTDecoderPrediction(config, state_dict) + self.joint = ParakeetTDTJoint(config, state_dict) + self.config = config + + def forward(self, encoder_frame: torch.Tensor, token_ids: torch.Tensor, *state_tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: + if len(state_tensors) != self.config.predictor_num_layers * 2: + raise ValueError( + f"expected {self.config.predictor_num_layers * 2} predictor state tensors, got {len(state_tensors)}" + ) + state_h = tuple(state_tensors[0::2]) + state_c = tuple(state_tensors[1::2]) + predictor_hidden, next_h, next_c = self.prediction(token_ids, state_h, state_c) + logits = self.joint(encoder_frame, predictor_hidden) + outputs: list[torch.Tensor] = [logits] + for h, c in zip(next_h, next_c, strict=True): + outputs.append(h) + outputs.append(c) + return tuple(outputs) + + +class ParakeetTDTLocalModel(nn.Module): + def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor]): + super().__init__() + self.name_or_path = config.model_source + self.family = "parakeet_tdt" + self.config = config + self.encoder = ParakeetTDTEncoder(config, state_dict) + self.decoder = nn.Module() + self.decoder.prediction = ParakeetTDTDecoderPrediction(config, state_dict) + self.joint = ParakeetTDTJoint(config, state_dict) + self.decoder_step = ParakeetTDTDecoderStep(config, state_dict) + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + return self.encoder(input_features) + + def initial_decoder_state(self, *, batch_size: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, ...]: + state: list[torch.Tensor] = [] + for _ in range(self.config.predictor_num_layers): + state.append(torch.zeros((batch_size, self.config.predictor_hidden_dim), device=device, dtype=dtype)) + state.append(torch.zeros((batch_size, self.config.predictor_hidden_dim), device=device, dtype=dtype)) + return tuple(state) + + def greedy_decode_token_ids(self, input_features: torch.Tensor) -> list[int]: + with torch.no_grad(): + encoder_hidden = self.encoder(input_features) + batch = int(encoder_hidden.shape[0]) + if batch != 1: + raise ValueError("Parakeet TDT local greedy decode currently expects batch size 1") + states = self.initial_decoder_state( + batch_size=batch, + device=encoder_hidden.device, + dtype=encoder_hidden.dtype, + ) + state_arrays = tuple(state.detach().cpu().numpy() for state in states) + encoder_hidden_np = encoder_hidden.detach().cpu().numpy() + + def _step( + frame: np.ndarray, + token_id: int, + state_values: tuple[np.ndarray, ...], + ) -> tuple[np.ndarray, tuple[np.ndarray, ...]]: + frame_tensor = torch.from_numpy(frame).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) + token_tensor = torch.tensor([[token_id]], device=encoder_hidden.device, dtype=torch.long) + state_tensors = tuple( + torch.from_numpy(value).to(device=encoder_hidden.device, dtype=encoder_hidden.dtype) + for value in state_values + ) + outputs = self.decoder_step(frame_tensor, token_tensor, *state_tensors) + logits = outputs[0].detach().cpu().numpy() + next_states = tuple(output.detach().cpu().numpy() for output in outputs[1:]) + return logits, next_states + + return greedy_decode_parakeet_tdt_token_ids( + config=self.config, + encoder_hidden_states=encoder_hidden_np, + initial_states=state_arrays, + step=_step, + ) + + def decode_token_ids(self, token_ids: list[int]) -> str: + pieces: list[str] = [] + for token_id in token_ids: + if token_id < 0 or token_id >= len(self.config.vocabulary): + continue + piece = self.config.vocabulary[token_id] + if piece.startswith("<|") and piece.endswith("|>"): + continue + pieces.append(piece) + text = "".join(pieces).replace("▁", " ") + return re.sub(r"\s+", " ", text).strip() + + +def load_tdt_local_model(model_source: str, *, torch_dtype: torch.dtype) -> ParakeetTDTLocalModel: + config = load_parakeet_tdt_config(model_source) + state_dict = _load_tensor_state_dict(model_source) + predictor_vocab_size = int(state_dict["decoder.prediction.embed.weight"].shape[0]) + if config.blank_id < 0 or config.blank_id >= predictor_vocab_size: + config = replace(config, blank_id=predictor_vocab_size - 1) + elif config.blank_id == len(config.vocabulary) - 1 and predictor_vocab_size == len(config.vocabulary) + 1: + config = replace(config, blank_id=predictor_vocab_size - 1) + model = ParakeetTDTLocalModel(config, state_dict).eval() + model.to(dtype=torch_dtype) + return model + + +def build_parakeet_tdt_component_specs( + model: ParakeetTDTLocalModel, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, +) -> list[ComponentModuleSpec]: + input_features = named_tensors["input_features"] + example_hidden = model.encoder(input_features) + batch_size = int(example_hidden.shape[0]) + initial_states = model.initial_decoder_state( + batch_size=batch_size, + device=example_hidden.device, + dtype=example_hidden.dtype, + ) + example_frame = example_hidden[:, :1, :].reshape(batch_size, example_hidden.shape[-1]) + example_token_id = torch.full( + (batch_size,), + int(model.config.blank_id), + device=example_hidden.device, + dtype=torch.long, + ) + + state_input_keys: list[str] = [] + state_output_keys: list[str] = [] + for index in range(model.config.predictor_num_layers): + state_input_keys.extend((f"state_h_{index}", f"state_c_{index}")) + state_output_keys.extend((f"state_h_{index}", f"state_c_{index}")) + + common_graph_meta = { + "weights_dir": weights_dir, + "task": "tdt_transcription", + "adapter_family": "parakeet_tdt", + } + + return [ + ComponentModuleSpec( + component="audio_encoder", + module=model.encoder, + example_inputs=(input_features,), + input_keys=("input_features",), + output_keys=("encoder_hidden_states",), + graph_meta={**common_graph_meta, "component": "audio_encoder"}, + metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, + ), + ComponentModuleSpec( + component="decoder", + module=model.decoder_step, + example_inputs=(example_frame, example_token_id, *initial_states), + input_keys=("encoder_frame", "token_ids", *state_input_keys), + output_keys=("step_logits", *state_output_keys), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, + ), + ] diff --git a/python/cactus/transpile/weight_binding.py b/python/cactus/transpile/weight_binding.py new file mode 100644 index 000000000..a82f14c6a --- /dev/null +++ b/python/cactus/transpile/weight_binding.py @@ -0,0 +1,363 @@ +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +import json +import os +from pathlib import Path +import re +from typing import Any + + +@dataclass(frozen=True) +class WeightBinding: + path: str + kind: str # "weight" | "embedding" + source_name: str + + +_TOKEN_EMBEDDING_OUTPUT_NAMES = { + "token_embeddings.weights", + "token_embeddings.cq2.weights", + "token_embeddings.cq3.weights", + "token_embeddings.cq4.weights", + "decoder_token_embeddings.weights", + "decoder_token_embeddings.cq2.weights", + "decoder_token_embeddings.cq3.weights", + "decoder_token_embeddings.cq4.weights", +} + + +def _candidate_model_dir_names(model_name_or_path: str) -> list[str]: + candidates: list[str] = [] + + def _add(name: str) -> None: + name = name.strip().lower() + if name and name not in candidates: + candidates.append(name) + + raw = model_name_or_path.strip() + if not raw: + return candidates + + path = Path(raw) + _add(path.name) + if "/" in raw: + _add(raw.split("/")[-1]) + + slug = raw.replace("/", "--") + _add(slug) + _add(slug.replace("--", "-")) + + for part in path.parts: + if not part.startswith("models--"): + continue + cache_name = part[len("models--") :] + _add(cache_name) + _add(cache_name.replace("--", "-")) + _add(cache_name.split("--")[-1]) + break + + return candidates + + +def _default_weights_dir_for_model_name(model_name_or_path: str) -> str | None: + if not model_name_or_path: + return None + from ..cli.common import weights_root + root = weights_root() + for model_dir_name in _candidate_model_dir_names(model_name_or_path): + candidate = root / model_dir_name + if candidate.exists(): + return str(candidate) + return None + + +def resolve_transpile_weights_dir(graph_meta: dict[str, object]) -> str | None: + explicit = graph_meta.get("weights_dir") + if isinstance(explicit, str) and explicit: + return explicit + + family = str(graph_meta.get("adapter_family", "")).upper() + family_env = f"CACTUS_TRANSPILER_WEIGHTS_DIR_{family}" + if family and family_env in os.environ and os.environ[family_env]: + return os.environ[family_env] + + generic = os.environ.get("CACTUS_TRANSPILER_WEIGHTS_DIR") + if generic: + return generic + + model_name_or_path = graph_meta.get("model_name_or_path") + if isinstance(model_name_or_path, str) and model_name_or_path: + return _default_weights_dir_for_model_name(model_name_or_path) + return None + + +def _manifest_source_aliases(name: str) -> tuple[str, ...]: + aliases: list[str] = [] + + def _add(candidate: str) -> None: + if candidate and candidate not in aliases: + aliases.append(candidate) + + def _add_with_wrappers(candidate: str) -> None: + _add(candidate) + for prefix in ( + "module.", + "module.model.", + "adapter.", + "adapter.model.", + ): + _add(f"{prefix}{candidate}") + + raw = name.strip() + if not raw: + return () + + _add_with_wrappers(raw) + + unwrapped = raw + changed = True + while changed: + changed = False + for wrapper in ("module.model.", "adapter.model.", "module.", "adapter."): + if unwrapped.startswith(wrapper): + unwrapped = unwrapped[len(wrapper) :] + changed = True + break + + if unwrapped.endswith("lm_head.weight"): + prefix = unwrapped[: -len("lm_head.weight")] + for tied_name in ( + f"{prefix}embed_tokens.weight", + f"{prefix}language_model.embed_tokens.weight", + f"model.{prefix}embed_tokens.weight", + f"model.{prefix}language_model.embed_tokens.weight", + "embed_tokens.weight", + "language_model.embed_tokens.weight", + "model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", + ): + _add_with_wrappers(tied_name) + + def _allow_multimodal_tail_alias(tail: str) -> bool: + if tail.startswith(("audio_tower.", "vision_tower.", "embed_audio.", "embed_vision.")): + return os.environ.get("CACTUS_TRANSPILER_BIND_MULTIMODAL_TOWER_WEIGHTS") != "0" + return True + + def _add_multimodal_backbone_aliases(candidate: str) -> None: + for prefix in ( + "module.multimodal_backbone.", + "module.model.multimodal_backbone.", + "adapter.multimodal_backbone.", + "adapter.model.multimodal_backbone.", + "multimodal_backbone.", + ): + if candidate.startswith(prefix): + tail = candidate[len(prefix) :] + if not _allow_multimodal_tail_alias(tail): + continue + _add_with_wrappers(tail) + _add_with_wrappers(f"model.{tail}") + + for prefix in ( + "model.audio_tower.", + "model.vision_tower.", + "model.embed_audio.", + "model.embed_vision.", + "audio_tower.", + "vision_tower.", + "embed_audio.", + "embed_vision.", + ): + if candidate.startswith(prefix): + if candidate.startswith("model."): + tail = candidate[len("model.") :] + else: + tail = candidate + if not _allow_multimodal_tail_alias(tail): + continue + _add_with_wrappers(f"multimodal_backbone.{tail}") + _add_with_wrappers(f"module.multimodal_backbone.{tail}") + + for prefix in ("encoder.", "decoder.", "model.encoder.", "model.decoder."): + if raw.startswith(prefix): + _add_with_wrappers(raw[len(prefix) :]) + + for prefix in ("model.language_model.", "language_model."): + if raw.startswith(prefix): + tail = raw[len(prefix) :] + _add_with_wrappers(tail) + _add_with_wrappers(f"backbone.{tail}") + _add_with_wrappers(f"model.{tail}") + + if raw.startswith("model."): + tail = raw[len("model.") :] + _add_with_wrappers(tail) + + def _add_text_backbone_aliases(candidate: str) -> None: + for prefix in ( + "module.backbone.", + "module.model.backbone.", + "adapter.backbone.", + "adapter.model.backbone.", + "backbone.", + ): + if candidate.startswith(prefix): + tail = candidate[len(prefix) :] + _add_with_wrappers(tail) + _add_with_wrappers(f"model.{tail}") + _add_with_wrappers(f"model.language_model.{tail}") + + for candidate in tuple(aliases): + _add_text_backbone_aliases(candidate) + _add_multimodal_backbone_aliases(candidate) + + for candidate in tuple(aliases): + normalized = candidate + for prefix in ( + "module.model.backbone.", + "adapter.model.backbone.", + "module.backbone.", + "adapter.backbone.", + "backbone.", + "module.model.", + "adapter.model.", + "module.", + "adapter.", + ): + if normalized.startswith(prefix): + normalized = normalized[len(prefix) :] + break + match = re.match( + r"^(?:model\.)?layers\.(\d+)\.feed_forward\.(w[123])_weights\.(\d+)$", + normalized, + ) + if match is not None: + layer, weight_name, expert = match.groups() + _add_with_wrappers(f"model.layers.{layer}.feed_forward.experts.{expert}.{weight_name}.weight") + _add_with_wrappers(f"layers.{layer}.feed_forward.experts.{expert}.{weight_name}.weight") + + return tuple(aliases) + + +def _manifest_entry_kind(output_name: str, explicit_kind: Any, explicit_names: list[object]) -> str: + if isinstance(explicit_kind, str) and explicit_kind: + return explicit_kind + + output_text = output_name.lower() + explicit_text = " ".join(str(name).lower() for name in explicit_names if isinstance(name, str)) + if ( + output_name in _TOKEN_EMBEDDING_OUTPUT_NAMES + or "embedding" in output_text + or "token_embeddings" in output_text + or "embed_tokens.weight" in explicit_text + or "embed_positions.weight" in explicit_text + ): + return "embedding" + return "weight" + + +def _flatten_convert_manifest(root_manifest: dict[str, Any]) -> dict[str, object]: + flattened: dict[str, object] = {} + rows = root_manifest.get("weights") + if not isinstance(rows, list): + for name, entry in root_manifest.items(): + if not isinstance(name, str): + continue + for alias in _manifest_source_aliases(name): + flattened.setdefault(alias, entry) + return flattened + + for row in rows: + if not isinstance(row, dict): + continue + output_name = row.get("output_name") or row.get("filename") + if not isinstance(output_name, str) or not output_name: + continue + + explicit_names: list[object] = [ + row.get("source_name"), + row.get("hf_name"), + row.get("adapter_name"), + ] + source_names = row.get("source_names") + if isinstance(source_names, list): + explicit_names.extend(source_names) + + entry = { + "filename": output_name, + "kind": _manifest_entry_kind(output_name, row.get("kind"), explicit_names), + } + for source_name in explicit_names: + if not isinstance(source_name, str) or not source_name: + continue + for alias in _manifest_source_aliases(source_name): + flattened.setdefault(alias, entry) + return flattened + + +@lru_cache(maxsize=16) +def _load_weights_manifest_cached( + manifest_path_text: str, + *, + mtime_ns: int, + size: int, +) -> dict[str, object]: + del mtime_ns, size + manifest_path = Path(manifest_path_text) + if not manifest_path.exists(): + return {} + try: + loaded_manifest = json.loads(manifest_path.read_text()) + except json.JSONDecodeError as exc: + raise RuntimeError(f"corrupt weights_manifest.json at {manifest_path}: {exc}") from exc + if not isinstance(loaded_manifest, dict): + raise RuntimeError(f"weights_manifest.json at {manifest_path} is not a JSON object") + return _flatten_convert_manifest(loaded_manifest) + + +def _load_weights_manifest(root: Path) -> dict[str, object]: + manifest_path = root / "weights_manifest.json" + if not manifest_path.exists(): + return {} + stat = manifest_path.stat() + return _load_weights_manifest_cached( + str(manifest_path.resolve()), + mtime_ns=int(stat.st_mtime_ns), + size=int(stat.st_size), + ) + + +def _binding_from_manifest_entry(root: Path, source_name: str, entry: object) -> WeightBinding | None: + if not isinstance(entry, dict): + return None + filename = entry.get("filename") + kind = entry.get("kind", "weight") + if not isinstance(filename, str) or not isinstance(kind, str): + return None + candidate = root / filename + if not candidate.exists(): + return None + return WeightBinding(path=str(candidate), kind=kind, source_name=source_name) + + +def resolve_weight_binding(*, weights_dir: str | None, source_name: str) -> WeightBinding | None: + """Resolve an exported parameter to a converted Cactus tensor file. + + The resolver intentionally trusts only `weights_manifest.json`. It may add + wrapper-local aliases such as `module.` or `adapter.model.`, but it does not + guess old model-specific filenames from layer names. + """ + + if not weights_dir: + return None + root = Path(weights_dir) + if not root.exists(): + return None + manifest = _load_weights_manifest(root) + for alias in _manifest_source_aliases(source_name): + binding = _binding_from_manifest_entry(root, alias, manifest.get(alias)) + if binding is not None: + return WeightBinding(path=binding.path, kind=binding.kind, source_name=source_name) + return None diff --git a/python/cactus/transpile/weight_compat.py b/python/cactus/transpile/weight_compat.py new file mode 100644 index 000000000..61521890e --- /dev/null +++ b/python/cactus/transpile/weight_compat.py @@ -0,0 +1,568 @@ +from __future__ import annotations + +from dataclasses import replace +import json +import shutil +import struct +from pathlib import Path + +import numpy as np + +from cactus.convert.cactus_adapters.tensor_io import CACTUS_ALIGNMENT +from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC +from cactus.convert.cactus_adapters.tensor_io import FLAG_INTERLEAVED +from cactus.convert.cactus_adapters.tensor_io import align_offset +from cactus.convert.model_adapters.naming import gemma4_scale_factor +from cactus.transpile.runtime_compat import Graph +from cactus.transpile.weight_binding import WeightBinding +from cactus.convert.quantization.cq import FLAG_ORTHOGONAL_ROTATION +from cactus.convert.quantization.cq import FLAG_INTERLEAVED_4ROW +from cactus.convert.quantization.cq import GROUP_SIZE as CQ_GROUP_SIZE +from cactus.convert.quantization.cq import PRECISION_CQ +from cactus.convert.quantization.cq import make_codebook +from cactus.convert.quantization.cq import make_hadamard_components +from cactus.convert.quantization.cq import make_hadamard_matrix +from cactus.convert.quantization.cq import make_orthogonal_rotation +from cactus.convert.quantization.cq import pack_indices_interleaved_4row +from cactus.convert.quantization.cq import pack_indices_lsb +from cactus.convert.quantization.cq import quantize_hadamard +from cactus.convert.quantization.cq import quantize_orthogonal +from cactus.convert.quantization.cq import write_cq_tensor + + +_HEADER_SIZE = 84 +_FLAG_EXTENDED_SHAPE = 1 << 4 +_INT8 = int(Graph.INT8) +_INT4 = int(Graph.CQ1) + +_TOKEN_EMBEDDING_FILENAMES = { + "token_embeddings.weights", + "decoder_token_embeddings.weights", +} +_PER_LAYER_EMBEDDING_FILENAMES = { + "embed_tokens_per_layer.weights", +} + + +class _OpenedTensor: + def __init__( + self, + *, + path: Path, + precision: int, + shape: tuple[int, ...], + data: np.memmap, + scales: np.memmap | None, + group_size: int, + num_groups: int, + is_interleaved: bool, + original_n: int, + alignment: int, + ) -> None: + self.path = path + self.precision = precision + self.shape = shape + self.data = data + self.scales = scales + self.group_size = group_size + self.num_groups = num_groups + self.is_interleaved = is_interleaved + self.original_n = original_n + self.alignment = alignment + + +def ensure_binding_compatible(binding: WeightBinding, source_tensor: object | None = None) -> WeightBinding: + binding = ensure_embedding_binding_compatible(binding) + return _ensure_legacy_int4_weight_binding_compatible(binding, source_tensor=source_tensor) + + +def ensure_embedding_binding_compatible(binding: WeightBinding) -> WeightBinding: + if binding.kind != "embedding": + return binding + + source_path = Path(binding.path).expanduser().resolve() + opened = _open_cactus_tensor_file(source_path) + if opened.precision != _INT8: + return binding + if len(opened.shape) != 2 or opened.scales is None or opened.group_size <= 0: + return binding + + config = _embedding_cache_config(source_path.name, source_name=binding.source_name) + if config is None: + return binding + + compat_path = source_path.with_name(source_path.stem + f".cq{config['bits']}.weights") + _materialize_cq_embedding_cache( + opened, + compat_path, + bits=int(config["bits"]), + rotation=str(config["rotation"]), + layout=str(config.get("layout", "row_major")), + ) + _cleanup_legacy_fp16_cache(source_path) + return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) + + +def _embedding_cache_config(filename: str, *, source_name: str) -> dict[str, object] | None: + normalized_source = str(source_name or "") + if filename in _PER_LAYER_EMBEDDING_FILENAMES or normalized_source.endswith("embed_tokens_per_layer.weight"): + return {"bits": 2, "rotation": "hadamard"} + if filename in _TOKEN_EMBEDDING_FILENAMES: + return {"bits": 4, "rotation": "orthogonal", "layout": "interleaved_4row"} + return {"bits": 4, "rotation": "orthogonal"} + + +def _ensure_legacy_int4_weight_binding_compatible( + binding: WeightBinding, + *, + source_tensor: object | None, +) -> WeightBinding: + source_path = Path(binding.path).expanduser().resolve() + opened = _open_cactus_tensor_file(source_path) + if not _is_legacy_packed_int4_tensor(opened): + return binding + + compat_path = source_path.with_name(source_path.stem + ".cq4.weights") + if compat_path.exists(): + if _can_materialize_compat_weight(source_tensor): + _materialize_source_cq_weight( + source_tensor, + compat_path, + bits=4, + rotation="hadamard", + scale_factor=_compat_weight_scale_factor(source_path), + source_mtime_ns=source_path.stat().st_mtime_ns, + ) + return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) + + if not _can_materialize_compat_weight(source_tensor): + raise RuntimeError( + "legacy packed INT4 weight is not directly executable in the v2 runtime and is " + f"missing a CQ companion file: {source_path.name}. " + "Rerun `cactus convert ...` so the CQ weights can be materialized, " + "or generate them through `cq_convert` first." + ) + + _materialize_source_cq_weight( + source_tensor, + compat_path, + bits=4, + rotation="hadamard", + scale_factor=_compat_weight_scale_factor(source_path), + source_mtime_ns=source_path.stat().st_mtime_ns, + ) + return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) + + +def _is_legacy_packed_int4_tensor(opened: _OpenedTensor) -> bool: + if opened.precision != _INT4: + return False + if len(opened.shape) != 2 or opened.scales is None or opened.group_size <= 0: + return False + if not opened.is_interleaved: + return False + return True + + +def _can_materialize_compat_weight(source_tensor: object | None) -> bool: + if source_tensor is None: + return False + if getattr(source_tensor, "is_meta", False): + return False + if isinstance(source_tensor, (str, bytes, bytearray, Path)): + return False + if np.isscalar(source_tensor): + return False + return hasattr(source_tensor, "shape") + + +def _compat_weight_scale_factor(source_path: Path) -> float: + manifest_path = source_path.parent / "conversion_manifest.json" + if manifest_path.exists(): + try: + rows = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError(f"unreadable conversion_manifest.json at {manifest_path}: {exc}") from exc + if isinstance(rows, list): + for row in rows: + if isinstance(row, dict) and row.get("output_file") == source_path.name: + scale = row.get("scale_factor") + if scale is None: + return 1.0 + try: + return float(scale) + except (TypeError, ValueError) as exc: + raise ValueError( + f"non-numeric scale_factor {scale!r} for {source_path.name} in {manifest_path}" + ) from exc + parent_name = source_path.parent.name.lower() + if "gemma-4" in parent_name or "gemma4" in parent_name: + return gemma4_scale_factor(source_path.name) + return 1.0 + +def _cleanup_legacy_fp16_cache(source_path: Path) -> None: + legacy = source_path.with_name(source_path.stem + ".fp16.weights") + if not legacy.exists(): + return + try: + legacy.unlink() + except OSError: + pass + + +def _materialize_source_cq_weight( + tensor: object, + out_path: Path, + *, + bits: int, + rotation: str, + scale_factor: float = 1.0, + source_mtime_ns: int | None = None, +) -> None: + if source_mtime_ns is not None and out_path.exists() and out_path.stat().st_mtime_ns >= source_mtime_ns: + return + + if rotation == "orthogonal": + cq = quantize_orthogonal(tensor, bits=bits) + else: + cq = quantize_hadamard(tensor, bits=bits, use_gptq=False) + if float(scale_factor) != 1.0: + cq = replace(cq, norms=(cq.norms.astype(np.float32) * float(scale_factor)).astype(np.float16)) + write_cq_tensor(out_path, cq) + + +def _open_cactus_tensor_file(path: str | Path) -> _OpenedTensor: + tensor_path = Path(path).expanduser().resolve() + with tensor_path.open("rb") as handle: + header = handle.read(_HEADER_SIZE) + if len(header) < _HEADER_SIZE: + raise RuntimeError(f"tensor file is too small for a Cactus header: {tensor_path}") + if header[:4] != CACTUS_MAGIC: + raise RuntimeError(f"tensor file is missing the CACT header: {tensor_path}") + + flags = struct.unpack_from(" 0) + + dtype_map = { + int(Graph.INT8): np.int8, + int(Graph.FP16): np.float16, + int(Graph.FP32): np.float32, + int(Graph.CQ1): np.uint8, + 3: np.uint8, + 4: np.uint8, + 5: np.uint8, + 6: np.uint8, + } + dtype = dtype_map.get(precision) + if dtype is None: + raise RuntimeError(f"unsupported tensor precision {precision} in {tensor_path}") + + aligned_header = align_offset(header_size, alignment) + scales_offset = aligned_header if scales_bytes > 0 else 0 + data_offset = align_offset(scales_offset + scales_bytes, alignment) if scales_bytes > 0 else aligned_header + + data_element_count = byte_size // np.dtype(dtype).itemsize + data = np.memmap(tensor_path, mode="r", dtype=dtype, offset=data_offset, shape=(data_element_count,)) + scales = None + if scales_bytes > 0: + scales = np.memmap( + tensor_path, + mode="r", + dtype=np.float16, + offset=scales_offset, + shape=(scales_bytes // np.dtype(np.float16).itemsize,), + ) + return _OpenedTensor( + path=tensor_path, + precision=precision, + shape=shape, + data=data, + scales=scales, + group_size=group_size, + num_groups=num_groups, + is_interleaved=bool(flags & FLAG_INTERLEAVED), + original_n=original_n, + alignment=alignment, + ) + + +def _materialize_cq_embedding_cache( + opened: _OpenedTensor, + out_path: Path, + *, + bits: int, + rotation: str, + layout: str = "row_major", +) -> None: + src_mtime_ns = opened.path.stat().st_mtime_ns + if out_path.exists() and out_path.stat().st_mtime_ns >= src_mtime_ns: + return + + if len(opened.shape) != 2: + raise RuntimeError(f"expected rank-2 embedding tensor, got shape={opened.shape}") + + rows = int(opened.original_n or opened.shape[0]) + cols = int(opened.shape[1]) + interleaved = layout == "interleaved_4row" + if interleaved: + if rotation != "orthogonal" or bits != 4: + raise RuntimeError(f"INTERLEAVED_4ROW embedding cache requires orthogonal CQ4, got rotation={rotation} bits={bits}") + if rows % 4 != 0 or cols % 32 != 0: + raise RuntimeError(f"INTERLEAVED_4ROW embedding cache requires rows % 4 == 0 and cols % 32 == 0, got {(rows, cols)}") + if rotation == "hadamard" and cols % CQ_GROUP_SIZE != 0: + raise RuntimeError( + f"Hadamard CQ embedding conversion requires hidden dim divisible by {CQ_GROUP_SIZE}, " + f"got {cols} for {opened.path.name}" + ) + + input_scale = _estimate_embedding_input_scale(opened, rows=rows, cols=cols) + precision = int(PRECISION_CQ[int(bits)]) + group_size = cols if rotation == "orthogonal" else CQ_GROUP_SIZE + num_groups = cols // group_size + codebook = make_codebook(group_size, bits).astype(np.float16) + input_scale = np.asarray(input_scale, dtype=np.float16).reshape(cols) + recip = np.minimum(1.0 / np.maximum(input_scale.astype(np.float32), 1e-8), 65504.0).astype(np.float16) + + chunk_rows = _recommended_chunk_rows(cols=cols, rotation=rotation) + norms_tmp = out_path.with_suffix(out_path.suffix + ".norms.tmp") + data_tmp = out_path.with_suffix(out_path.suffix + ".data.tmp") + final_tmp = out_path.with_suffix(out_path.suffix + ".tmp") + + flags = 0 + trailer_parts: list[bytes] = [] + if rotation == "orthogonal": + flags |= FLAG_ORTHOGONAL_ROTATION + if interleaved: + flags |= FLAG_INTERLEAVED_4ROW + trailer_parts.append(make_orthogonal_rotation(group_size, seed=1234).astype(np.float16).tobytes()) + else: + left, right, perm = make_hadamard_components(group_size, seed=1234) + trailer_parts.extend([ + left.tobytes(), + right.tobytes(), + perm.astype(" current: + handle.write(b"\x00" * (data_offset - current)) + with data_tmp.open("rb") as packed_handle: + shutil.copyfileobj(packed_handle, handle, length=8 * 1024 * 1024) + final_tmp.replace(out_path) + finally: + for temp_path in (norms_tmp, data_tmp): + try: + temp_path.unlink() + except FileNotFoundError: + pass + + +def _recommended_chunk_rows(*, cols: int, rotation: str) -> int: + if rotation == "orthogonal": + if cols >= 8192: + return 8 + if cols >= 4096: + return 16 + return 64 + if cols >= 8192: + return 32 + if cols >= 4096: + return 48 + return 96 + + +def _estimate_embedding_input_scale( + opened: _OpenedTensor, + *, + rows: int, + cols: int, + max_sample_rows: int = 2048, + sample_chunk_rows: int = 64, +) -> np.ndarray: + if rows <= 0 or cols <= 0: + return np.ones((cols,), dtype=np.float16) + + stride = max(1, rows // max(1, max_sample_rows)) + sampled = 0 + accum = np.zeros((cols,), dtype=np.float64) + for start in range(0, rows, stride * sample_chunk_rows): + stop = min(start + sample_chunk_rows, rows) + chunk = _dequantize_int8_rows(opened, start, stop).astype(np.float32, copy=False) + accum += np.abs(chunk).sum(axis=0, dtype=np.float64) + sampled += int(stop - start) + if sampled >= max_sample_rows: + break + + if sampled <= 0: + return np.ones((cols,), dtype=np.float16) + + mean_abs = np.maximum(accum / float(sampled), 1e-6) + raw = np.power(mean_abs, -0.5, dtype=np.float64) + raw /= np.exp(np.mean(np.log(np.clip(raw, 1e-6, None)))) + return np.clip(raw, 1.0 / 8.0, 8.0).astype(np.float16) + + +def _quantize_hadamard_chunk( + chunk: np.ndarray, + *, + bits: int, + input_scale: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + work = chunk.astype(np.float32, copy=False) * input_scale.astype(np.float32, copy=False)[None, :] + rows, cols = work.shape + group_size = CQ_GROUP_SIZE + groups = cols // group_size + rotation = make_hadamard_matrix(group_size, seed=1234).astype(np.float32) + codebook = make_codebook(group_size, bits).astype(np.float32) + + indices = np.empty((rows, cols), dtype=np.uint8) + norms = np.empty((rows, groups), dtype=np.float16) + for group_index in range(groups): + lo = group_index * group_size + hi = lo + group_size + group = work[:, lo:hi] + row_norms = np.linalg.norm(group, axis=1).clip(min=1e-8).astype(np.float32, copy=False) + rotated = (group / row_norms[:, None]) @ rotation + idx = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) + indices[:, lo:hi] = idx + norms[:, group_index] = row_norms.astype(np.float16) + return pack_indices_lsb(indices, group_size, bits), norms + + +def _quantize_orthogonal_chunk( + chunk: np.ndarray, + *, + bits: int, + input_scale: np.ndarray, + interleaved: bool = False, +) -> tuple[np.ndarray, np.ndarray]: + work = chunk.astype(np.float32, copy=False) * input_scale.astype(np.float32, copy=False)[None, :] + rows, cols = work.shape + rotation = make_orthogonal_rotation(cols, seed=1234).astype(np.float32) + codebook = make_codebook(cols, bits).astype(np.float32) + norms = np.linalg.norm(work, axis=1).clip(min=1e-8).astype(np.float32, copy=False) + rotated = (work / norms[:, None]) @ rotation + indices = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) + if interleaved: + return pack_indices_interleaved_4row(indices, cols, bits), norms.astype(np.float16).reshape(rows, 1) + return pack_indices_lsb(indices, cols, bits), norms.astype(np.float16).reshape(rows, 1) + + +def _write_cq_header( + handle, + *, + rows: int, + cols: int, + precision: int, + packed_bytes: int, + scales_bytes: int, + group_size: int, + num_groups: int, + flags: int, +) -> None: + handle.write(CACTUS_MAGIC) + handle.write(struct.pack(" np.ndarray: + if opened.scales is None or opened.group_size <= 0: + raise RuntimeError(f"INT8 embedding file is missing grouped scales: {opened.path}") + + rows = stop - start + cols = int(opened.shape[1]) + num_groups = int(opened.num_groups) + row_group = int(opened.group_size) + padded_rows = int(opened.shape[0]) + + if opened.is_interleaved: + block = 4 + padded_cols = ((cols + block - 1) // block) * block + q_blocks = opened.data.reshape(padded_rows // block, padded_cols // block, block, block) + s_blocks = opened.scales.reshape(padded_rows // block, num_groups, block) + block_start = start // block + block_end = (stop + block - 1) // block + q = np.asarray(q_blocks[block_start:block_end], dtype=np.int8).transpose(0, 2, 1, 3).reshape(-1, padded_cols) + s = np.asarray(s_blocks[block_start:block_end], dtype=np.float16).astype(np.float32).transpose(0, 2, 1).reshape(-1, num_groups) + row_offset = start - block_start * block + q = q[row_offset : row_offset + rows, :cols] + s = s[row_offset : row_offset + rows] + else: + q = np.asarray(opened.data, dtype=np.int8).reshape(padded_rows, cols)[start:stop] + s = np.asarray(opened.scales, dtype=np.float16).astype(np.float32).reshape(padded_rows, num_groups)[start:stop] + + out = np.empty((rows, cols), dtype=np.float16) + for group_index in range(num_groups): + lo = group_index * row_group + hi = min(lo + row_group, cols) + if lo >= hi: + break + scales = s[:, group_index : group_index + 1] + out[:, lo:hi] = (q[:, lo:hi].astype(np.float32) * scales).astype(np.float16) + return out diff --git a/python/example.py b/python/example.py deleted file mode 100644 index d1b128161..000000000 --- a/python/example.py +++ /dev/null @@ -1,78 +0,0 @@ -#!/usr/bin/env python3 -""" -Cactus Python FFI Example - -Usage: - 1. cactus build - 2. cactus download LiquidAI/LFM2-VL-450M - 3. cactus download openai/whisper-small - 4. cd tools && python example.py -""" - -import sys -import json -from pathlib import Path - -SCRIPT_DIR = Path(__file__).parent -PROJECT_ROOT = SCRIPT_DIR.parent -sys.path.insert(0, str(SCRIPT_DIR / "src")) - -from cactus import ( - cactus_init, - cactus_complete, - cactus_transcribe, - cactus_embed, - cactus_image_embed, - cactus_audio_embed, - cactus_reset, - cactus_destroy, - cactus_set_pro_key -) - -cactus_set_pro_key("") # email founders@cactuscompute.com - -WEIGHTS_DIR = PROJECT_ROOT / "weights" -ASSETS_DIR = PROJECT_ROOT / "tests" / "assets" - -# Load model -print("Loading LFM2-VL-450M...") -vlm = cactus_init(str(WEIGHTS_DIR / "lfm2-vl-450m")) - -# Text completion -messages = json.dumps([{"role": "user", "content": "What is 2+2?"}]) -response = cactus_complete(vlm, messages) -print("\nCompletion:") -print(json.dumps(json.loads(response), indent=2)) - -# Text embedding -embedding = cactus_embed(vlm, "Hello world") -print(f"\nText embedding dim: {len(embedding)}") - -# Image embedding -embedding = cactus_image_embed(vlm, str(ASSETS_DIR / "test_monkey.png")) -print(f"\nImage embedding dim: {len(embedding)}") - -# VLM - describe image -messages = json.dumps([{"role": "user", "content": "Describe this image", "images": [str(ASSETS_DIR / "test_monkey.png")]}]) -response = cactus_complete(vlm, messages) -print("\nVLM Image Description:") -print(json.dumps(json.loads(response), indent=2)) - -cactus_reset(vlm) -cactus_destroy(vlm) - -# Transcription -print("\nLoading whisper-small...") -whisper = cactus_init(str(WEIGHTS_DIR / "whisper-small")) -whisper_prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" -response = cactus_transcribe(whisper, str(ASSETS_DIR / "test.wav"), prompt=whisper_prompt) -print("Transcription:") -print(json.dumps(json.loads(response), indent=2)) - -# Audio embedding -embedding = cactus_audio_embed(whisper, str(ASSETS_DIR / "test.wav")) -print(f"\nAudio embedding dim: {len(embedding)}") - -cactus_destroy(whisper) - -print("\nDone!") diff --git a/python/pyproject.toml b/python/pyproject.toml index 1306ce4f1..0198bb301 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -1,25 +1,77 @@ [project] -name = "cactus" -version = "0.1.0" -description = "Cactus Python tools" -requires-python = ">=3.8" +name = "cactus-compute" +dynamic = ["version"] +description = "On-device AI inference — LLM, vision, speech, embeddings, and RAG" +readme = "README.md" +requires-python = ">=3.10,<3.14" +license = {text = "Cactus Compute License"} +authors = [{name = "Cactus Compute"}] +keywords = [ + "ai", "inference", "on-device", "llm", "embeddings", + "transcription", "rag", "quantization", "edge-ai", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Environment :: Console", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: Other/Proprietary License", + "Operating System :: MacOS", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Multimedia :: Sound/Audio :: Speech", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Scientific/Engineering :: Image Recognition", +] dependencies = [ - "torch", - "transformers", - "numpy", + "numpy>=1.26.0,<3", + "huggingface-hub>=1.5.0,<2.0", + "fastapi>=0.100.0", + "uvicorn>=0.20.0", + "python-multipart>=0.0.5", + "Pillow>=11.0.0,<12", ] [project.optional-dependencies] -vlm = ["Pillow", "num2words", "torchvision"] -hub = ["huggingface_hub"] +dev = ["pytest>=8.0", "httpx>=0.25.0", "einops>=0.7"] +lora = ["peft>=0.15"] +convert = [ + "torch>=2.8.0,<3", + "torchvision>=0.23.0,<1", + "transformers==5.5.4", + "scipy>=1.13", + "sentencepiece>=0.2.0", +] + +[project.urls] +Homepage = "https://cactuscompute.com" +Repository = "https://github.com/cactus-compute/cactus" +Documentation = "https://docs.cactuscompute.com" +"Bug Tracker" = "https://github.com/cactus-compute/cactus/issues" [project.scripts] -cactus = "src.cli:main" +cactus = "cactus.cli:main" [build-system] -requires = ["setuptools>=61.0"] +requires = ["setuptools>=61.0", "wheel"] build-backend = "setuptools.build_meta" +[tool.setuptools.dynamic] +version = {attr = "cactus._version.__version__"} + [tool.setuptools.packages.find] where = ["."] -include = ["src*"] +include = ["cactus*"] +exclude = ["cactus.code*"] + +[tool.setuptools.package-data] +"cactus.bindings" = ["lib/*.dylib", "lib/*.so"] +"cactus" = ["bin/run", "bin/transcribe", "assets/*", "code/**/*"] +"cactus.convert" = ["assets/*/*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/python/requirements.txt b/python/requirements.txt index 64f2bc312..de9259adb 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -1,8 +1,12 @@ -torch>=2.8.0 -torchvision>=0.23.0 -transformers>=4.57.0 -pillow>=11.3.0 -num2words>=0.5.12 -peft>=0.15.0 -einops>=0.8.1 -huggingface-hub==0.36.0 \ No newline at end of file +torch==2.8.0 +torchvision==0.23.0 +transformers==5.5.4 +numpy==2.4.6 +pillow==11.3.0 +huggingface_hub==1.15.0 +scipy==1.17.1 +sentencepiece==0.2.2 +fastapi>=0.100.0 +uvicorn>=0.20.0 +python-multipart>=0.0.5 +httpx>=0.25.0 diff --git a/python/src/__init__.py b/python/src/__init__.py deleted file mode 100644 index fbf8e03ab..000000000 --- a/python/src/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -__version__ = "0.1.0" - - -def __getattr__(name): - if name == "main": - from .cli import main - return main - if name == "convert_hf_to_cactus": - from .cli import convert_hf_to_cactus - return convert_hf_to_cactus - if name == "get_model_dir_name": - from .cli import get_model_dir_name - return get_model_dir_name - if name == "save_tensor_with_header": - from .tensor_io import save_tensor_with_header - return save_tensor_with_header - if name == "convert_hf_tokenizer": - from .tokenizer import convert_hf_tokenizer - return convert_hf_tokenizer - if name == "MixedPrecisionConfig": - from .precision_config import MixedPrecisionConfig - return MixedPrecisionConfig - if name == "determine_precision": - from .precision_config import determine_precision - return determine_precision - if name == "count_model_parameters": - from .precision_config import count_model_parameters - return count_model_parameters - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/src/cactus.py b/python/src/cactus.py deleted file mode 100644 index fdac21dfe..000000000 --- a/python/src/cactus.py +++ /dev/null @@ -1,745 +0,0 @@ -""" -Cactus FFI Python Bindings - -Python bindings for Cactus Engine via FFI. Provides access to: -- Text completion with LLMs (including cloud handoff detection) -- Audio transcription with Whisper models -- Text, image, and audio embeddings -- RAG (Retrieval-Augmented Generation) queries -- Tool RAG (automatic tool selection based on query relevance) -- Streaming transcription -- Vector index for similarity search - -Response Format: -All completion responses use a unified JSON format with all fields always present: -{ - "success": bool, # True if generation succeeded - "error": str|null, # Error message if failed, null otherwise - "cloud_handoff": bool, # True if model recommends deferring to cloud - "response": str|null, # Generated text, null if cloud_handoff or error - "function_calls": [], # List of function calls if tools were used - "confidence": float, # Model confidence (1.0 - normalized_entropy) - "time_to_first_token_ms": float, - "total_time_ms": float, - "prefill_tps": float, - "decode_tps": float, - "ram_usage_mb": float, - "prefill_tokens": int, - "decode_tokens": int, - "total_tokens": int -} -""" -import ctypes -import json -import platform -from pathlib import Path - -TokenCallback = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_uint32, ctypes.c_void_p) - -_DIR = Path(__file__).parent.parent.parent -if platform.system() == "Darwin": - _LIB_PATH = _DIR / "cactus" / "build" / "libcactus.dylib" -else: - _LIB_PATH = _DIR / "cactus" / "build" / "libcactus.so" - -if not _LIB_PATH.exists(): - raise RuntimeError( - f"Cactus library not found at {_LIB_PATH}\n" - f"Please build first: cactus build --python" - ) - -_lib = ctypes.CDLL(str(_LIB_PATH)) - -_lib.cactus_init.argtypes = [ctypes.c_char_p, ctypes.c_char_p] -_lib.cactus_init.restype = ctypes.c_void_p - -_lib.cactus_complete.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_size_t, - ctypes.c_char_p, ctypes.c_char_p, TokenCallback, ctypes.c_void_p -] -_lib.cactus_complete.restype = ctypes.c_int - -_lib.cactus_transcribe.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p, - ctypes.c_size_t, ctypes.c_char_p, TokenCallback, ctypes.c_void_p, - ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t -] -_lib.cactus_transcribe.restype = ctypes.c_int - -_lib.cactus_embed.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), - ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_bool -] -_lib.cactus_embed.restype = ctypes.c_int - -_lib.cactus_image_embed.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), - ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t) -] -_lib.cactus_image_embed.restype = ctypes.c_int - -_lib.cactus_audio_embed.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), - ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t) -] -_lib.cactus_audio_embed.restype = ctypes.c_int - -_lib.cactus_reset.argtypes = [ctypes.c_void_p] -_lib.cactus_reset.restype = None - -_lib.cactus_stop.argtypes = [ctypes.c_void_p] -_lib.cactus_stop.restype = None - -_lib.cactus_destroy.argtypes = [ctypes.c_void_p] -_lib.cactus_destroy.restype = None - -_lib.cactus_get_last_error.argtypes = [] -_lib.cactus_get_last_error.restype = ctypes.c_char_p - -_lib.cactus_set_telemetry_token.argtypes = [ctypes.c_char_p] -_lib.cactus_set_telemetry_token.restype = None - -_lib.cactus_set_pro_key.argtypes = [ctypes.c_char_p] -_lib.cactus_set_pro_key.restype = None - -_lib.cactus_tokenize.argtypes = [ - ctypes.c_void_p, - ctypes.c_char_p, - ctypes.POINTER(ctypes.c_uint32), - ctypes.c_size_t, - ctypes.POINTER(ctypes.c_size_t), -] -_lib.cactus_tokenize.restype = ctypes.c_int - -_lib.cactus_score_window.argtypes = [ - ctypes.c_void_p, - ctypes.POINTER(ctypes.c_uint32), - ctypes.c_size_t, - ctypes.c_size_t, - ctypes.c_size_t, - ctypes.c_size_t, - ctypes.c_char_p, - ctypes.c_size_t, -] -_lib.cactus_score_window.restype = ctypes.c_int - -_lib.cactus_rag_query.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, - ctypes.c_size_t, ctypes.c_size_t -] -_lib.cactus_rag_query.restype = ctypes.c_int - -_lib.cactus_stream_transcribe_init.argtypes = [ctypes.c_void_p] -_lib.cactus_stream_transcribe_init.restype = ctypes.c_void_p - -_lib.cactus_stream_transcribe_insert.argtypes = [ - ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t -] -_lib.cactus_stream_transcribe_insert.restype = ctypes.c_int - -_lib.cactus_stream_transcribe_process.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t, ctypes.c_char_p -] -_lib.cactus_stream_transcribe_process.restype = ctypes.c_int - -_lib.cactus_stream_transcribe_finalize.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t -] -_lib.cactus_stream_transcribe_finalize.restype = ctypes.c_int - -_lib.cactus_stream_transcribe_destroy.argtypes = [ctypes.c_void_p] -_lib.cactus_stream_transcribe_destroy.restype = None - -_lib.cactus_index_init.argtypes = [ctypes.c_char_p, ctypes.c_size_t] -_lib.cactus_index_init.restype = ctypes.c_void_p - -_lib.cactus_index_add.argtypes = [ - ctypes.c_void_p, - ctypes.POINTER(ctypes.c_int), - ctypes.POINTER(ctypes.c_char_p), - ctypes.POINTER(ctypes.c_char_p), - ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), - ctypes.c_size_t, - ctypes.c_size_t -] -_lib.cactus_index_add.restype = ctypes.c_int - -_lib.cactus_index_delete.argtypes = [ - ctypes.c_void_p, - ctypes.POINTER(ctypes.c_int), - ctypes.c_size_t -] -_lib.cactus_index_delete.restype = ctypes.c_int - -_lib.cactus_index_query.argtypes = [ - ctypes.c_void_p, - ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), - ctypes.c_size_t, - ctypes.c_size_t, - ctypes.c_char_p, - ctypes.POINTER(ctypes.POINTER(ctypes.c_int)), - ctypes.POINTER(ctypes.c_size_t), - ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), - ctypes.POINTER(ctypes.c_size_t) -] -_lib.cactus_index_query.restype = ctypes.c_int - -_lib.cactus_index_compact.argtypes = [ctypes.c_void_p] -_lib.cactus_index_compact.restype = ctypes.c_int - -_lib.cactus_index_destroy.argtypes = [ctypes.c_void_p] -_lib.cactus_index_destroy.restype = None - - -def cactus_init(model_path, corpus_dir=None): - """ - Initialize a model and return its handle. - - Args: - model_path: Path to model weights directory - corpus_dir: Optional path to RAG corpus directory for document Q&A - - Returns: - Model handle (opaque pointer) or None if initialization failed. - Use cactus_get_last_error() to get error details. - """ - return _lib.cactus_init( - model_path.encode() if isinstance(model_path, str) else model_path, - corpus_dir.encode() if corpus_dir else None - ) - - -def cactus_complete( - model, - messages, - tools=None, - temperature=None, - top_p=None, - top_k=None, - max_tokens=None, - stop_sequences=None, - force_tools=False, - tool_rag_top_k=None, - confidence_threshold=None, - callback=None -): - """ - Run chat completion on a model. - - Args: - model: Model handle from cactus_init - messages: List of message dicts or JSON string - tools: Optional list of tool definitions for function calling - temperature: Sampling temperature - top_p: Top-p sampling - top_k: Top-k sampling - max_tokens: Maximum tokens to generate - stop_sequences: List of stop sequences - force_tools: Constrain output to tool call format - tool_rag_top_k: Select top-k relevant tools via Tool RAG (default: 2, 0 = disabled) - confidence_threshold: Minimum confidence for local generation (default: 0.7, triggers cloud_handoff when below) - callback: Streaming callback fn(token, token_id, user_data) - - Returns: - JSON string with unified response format (all fields always present): - { - "success": bool, # True if generation succeeded - "error": str|null, # Error message if failed, null otherwise - "cloud_handoff": bool, # True if model confidence too low, should defer to cloud - "response": str|null, # Generated text, null if cloud_handoff or error - "function_calls": [], # List of function calls if tools were used - "confidence": float, # Model confidence (1.0 - normalized_entropy) - "time_to_first_token_ms": float, - "total_time_ms": float, - "prefill_tps": float, - "decode_tps": float, - "ram_usage_mb": float, - "prefill_tokens": int, - "decode_tokens": int, - "total_tokens": int - } - - When cloud_handoff is True, the model confidence dropped below confidence_threshold - and recommends deferring to a cloud-based model for better results. - """ - if isinstance(messages, list): - messages_json = json.dumps(messages) - else: - messages_json = messages - - tools_json = None - if tools is not None: - if isinstance(tools, list): - tools_json = json.dumps(tools) - else: - tools_json = tools - - options = {} - if temperature is not None: - options["temperature"] = temperature - if top_p is not None: - options["top_p"] = top_p - if top_k is not None: - options["top_k"] = top_k - if max_tokens is not None: - options["max_tokens"] = max_tokens - if stop_sequences is not None: - options["stop_sequences"] = stop_sequences - if force_tools: - options["force_tools"] = True - if tool_rag_top_k is not None: - options["tool_rag_top_k"] = tool_rag_top_k - if confidence_threshold is not None: - options["confidence_threshold"] = confidence_threshold - - options_json = json.dumps(options) if options else None - - buf = ctypes.create_string_buffer(65536) - cb = TokenCallback(callback) if callback else TokenCallback() - _lib.cactus_complete( - model, - messages_json.encode() if isinstance(messages_json, str) else messages_json, - buf, len(buf), - options_json.encode() if options_json else None, - tools_json.encode() if tools_json else None, - cb, None - ) - return buf.value.decode("utf-8", errors="ignore") - - -def cactus_transcribe(model, audio_path, prompt="", callback=None): - """ - Transcribe audio using a Whisper model. - - Args: - model: Whisper model handle from cactus_init - audio_path: Path to audio file (WAV format) - prompt: Whisper prompt for language/task (e.g., "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>") - callback: Optional streaming callback fn(token, token_id, user_data) - - Returns: - JSON string with response: {"success": bool, "response": str, ...} - """ - buf = ctypes.create_string_buffer(65536) - cb = TokenCallback(callback) if callback else TokenCallback() - _lib.cactus_transcribe( - model, - audio_path.encode() if isinstance(audio_path, str) else audio_path, - prompt.encode() if isinstance(prompt, str) else prompt, - buf, len(buf), - None, cb, None, None, 0 - ) - return buf.value.decode() - - -def cactus_embed(model, text, normalize=False): - """ - Get text embeddings. - - Args: - model: Model handle from cactus_init - text: Text to embed - normalize: L2-normalize embeddings (default: False) - - Returns: - List of floats representing the embedding vector. - """ - buf = (ctypes.c_float * 4096)() - dim = ctypes.c_size_t() - _lib.cactus_embed( - model, - text.encode() if isinstance(text, str) else text, - buf, ctypes.sizeof(buf), ctypes.byref(dim), normalize - ) - return list(buf[:dim.value]) - - -def cactus_image_embed(model, image_path): - """ - Get image embeddings from a VLM. - - Args: - model: Model handle from cactus_init - image_path: Path to image file - - Returns: - List of floats representing the image embedding vector. - """ - buf = (ctypes.c_float * 4096)() - dim = ctypes.c_size_t() - _lib.cactus_image_embed( - model, - image_path.encode() if isinstance(image_path, str) else image_path, - buf, ctypes.sizeof(buf), ctypes.byref(dim) - ) - return list(buf[:dim.value]) - - -def cactus_audio_embed(model, audio_path): - """ - Get audio embeddings from a Whisper model. - - Args: - model: Whisper model handle from cactus_init - audio_path: Path to audio file (WAV format) - - Returns: - List of floats representing the audio embedding vector. - """ - buf = (ctypes.c_float * 4096)() - dim = ctypes.c_size_t() - _lib.cactus_audio_embed( - model, - audio_path.encode() if isinstance(audio_path, str) else audio_path, - buf, ctypes.sizeof(buf), ctypes.byref(dim) - ) - return list(buf[:dim.value]) - - -def cactus_reset(model): - """Reset model state (clear KV cache). Call between unrelated conversations.""" - _lib.cactus_reset(model) - - -def cactus_stop(model): - """Stop an ongoing generation (useful with streaming callbacks).""" - _lib.cactus_stop(model) - - -def cactus_destroy(model): - """Free model memory. Always call when done.""" - _lib.cactus_destroy(model) - - -def cactus_get_last_error(): - """Get the last error message, or None if no error.""" - result = _lib.cactus_get_last_error() - return result.decode() if result else None - - -def cactus_set_telemetry_token(token): - """Set telemetry token for usage tracking.""" - _lib.cactus_set_telemetry_token( - token.encode() if isinstance(token, str) else token - ) - - -def cactus_set_pro_key(pro_key): - """Set Cactus Pro key for NPU acceleration (Apple devices).""" - _lib.cactus_set_pro_key( - pro_key.encode() if isinstance(pro_key, str) else pro_key - ) - - -def cactus_tokenize(model, text: str): - """ - Tokenize text. - - Args: - model: Model handle from cactus_init - text: Text to tokenize - - Returns: - List of token IDs. - """ - needed = ctypes.c_size_t(0) - rc = _lib.cactus_tokenize( - model, - text.encode("utf-8"), - None, - 0, - ctypes.byref(needed), - ) - if rc != 0: - raise RuntimeError(f"cactus_tokenize length query failed rc={rc}") - - n = needed.value - arr = (ctypes.c_uint32 * n)() - - rc = _lib.cactus_tokenize( - model, - text.encode("utf-8"), - arr, - n, - ctypes.byref(needed), - ) - if rc != 0: - raise RuntimeError(f"cactus_tokenize fetch failed rc={rc}") - - return [arr[i] for i in range(n)] - - -def cactus_score_window(model, tokens, start, end, context): - """ - Score a window of tokens for perplexity/log probability. - - Args: - model: Model handle from cactus_init - tokens: List of token IDs - start: Start index of window to score - end: End index of window to score - context: Context size for scoring - - Returns: - Dict with "success", "logprob", and "tokens" keys. - """ - buf = ctypes.create_string_buffer(4096) - n = len(tokens) - arr = (ctypes.c_uint32 * n)(*tokens) - - _lib.cactus_score_window( - model, - arr, - n, - start, - end, - context, - buf, - len(buf), - ) - return json.loads(buf.value.decode("utf-8", errors="ignore")) - - -def cactus_rag_query(model, query, top_k=5): - """ - Query RAG corpus for relevant text chunks. - - Args: - model: Model handle (must have been initialized with corpus_dir) - query: Query text - top_k: Number of chunks to retrieve (default: 5) - - Returns: - List of dicts with "score" and "text" keys, or empty list on error. - """ - buf = ctypes.create_string_buffer(65536) - result = _lib.cactus_rag_query( - model, - query.encode() if isinstance(query, str) else query, - buf, len(buf), top_k - ) - if result != 0: - return [] - return json.loads(buf.value.decode("utf-8", errors="ignore")) - - -def cactus_stream_transcribe_init(model): - """ - Initialize streaming transcription session. - - Args: - model: Whisper model handle from cactus_init - - Returns: - Stream handle for use with other stream_transcribe functions. - """ - return _lib.cactus_stream_transcribe_init(model) - - -def cactus_stream_transcribe_insert(stream, pcm_data): - """ - Insert audio data into streaming transcription buffer. - - Args: - stream: Stream handle from cactus_stream_transcribe_init - pcm_data: PCM audio data as bytes or list of uint8 - - Returns: - 0 on success, -1 on error. - """ - if isinstance(pcm_data, bytes): - arr = (ctypes.c_uint8 * len(pcm_data)).from_buffer_copy(pcm_data) - else: - arr = (ctypes.c_uint8 * len(pcm_data))(*pcm_data) - return _lib.cactus_stream_transcribe_insert(stream, arr, len(arr)) - - -def cactus_stream_transcribe_process(stream, options=None): - """ - Process buffered audio and return transcription. - - Args: - stream: Stream handle from cactus_stream_transcribe_init - options: Optional JSON string with options (e.g., {"confirmation_threshold": 0.95}) - - Returns: - JSON string with "success", "confirmed", and "pending" keys. - """ - buf = ctypes.create_string_buffer(65536) - _lib.cactus_stream_transcribe_process( - stream, buf, len(buf), - options.encode() if options else None - ) - return buf.value.decode("utf-8", errors="ignore") - - -def cactus_stream_transcribe_finalize(stream): - """ - Finalize streaming transcription and get final result. - - Args: - stream: Stream handle from cactus_stream_transcribe_init - - Returns: - JSON string with "success" and "confirmed" keys. - """ - buf = ctypes.create_string_buffer(65536) - _lib.cactus_stream_transcribe_finalize(stream, buf, len(buf)) - return buf.value.decode("utf-8", errors="ignore") - - -def cactus_stream_transcribe_destroy(stream): - """Free streaming transcription resources.""" - _lib.cactus_stream_transcribe_destroy(stream) - - -def cactus_index_init(index_dir, embedding_dim): - """ - Initialize a vector index. - - Args: - index_dir: Path to directory for index storage - embedding_dim: Dimension of embedding vectors - - Returns: - Index handle (opaque pointer) or None if initialization failed. - """ - return _lib.cactus_index_init( - index_dir.encode() if isinstance(index_dir, str) else index_dir, - embedding_dim - ) - - -def cactus_index_add(index, ids, documents, embeddings, metadatas=None): - """ - Add documents to the index. - - Args: - index: Index handle from cactus_index_init - ids: List of integer document IDs - documents: List of document strings - embeddings: List of embedding vectors (list of floats each) - metadatas: Optional list of metadata strings - - Returns: - 0 on success, -1 on error. - """ - count = len(ids) - embedding_dim = len(embeddings[0]) if embeddings else 0 - - ids_arr = (ctypes.c_int * count)(*ids) - - docs_arr = (ctypes.c_char_p * count)() - for i, doc in enumerate(documents): - docs_arr[i] = doc.encode() if isinstance(doc, str) else doc - - meta_arr = None - if metadatas: - meta_arr = (ctypes.c_char_p * count)() - for i, meta in enumerate(metadatas): - meta_arr[i] = meta.encode() if isinstance(meta, str) else meta - - emb_ptrs = (ctypes.POINTER(ctypes.c_float) * count)() - emb_arrays = [] - for i, emb in enumerate(embeddings): - arr = (ctypes.c_float * len(emb))(*emb) - emb_arrays.append(arr) - emb_ptrs[i] = ctypes.cast(arr, ctypes.POINTER(ctypes.c_float)) - - return _lib.cactus_index_add( - index, - ids_arr, - docs_arr, - meta_arr, - emb_ptrs, - count, - embedding_dim - ) - - -def cactus_index_delete(index, ids): - """ - Delete documents from the index. - - Args: - index: Index handle from cactus_index_init - ids: List of document IDs to delete - - Returns: - 0 on success, -1 on error. - """ - count = len(ids) - ids_arr = (ctypes.c_int * count)(*ids) - return _lib.cactus_index_delete(index, ids_arr, count) - - -def cactus_index_query(index, embedding, top_k=5, options=None): - """ - Query the index for similar documents. - - Args: - index: Index handle from cactus_index_init - embedding: Query embedding vector (list of floats) - top_k: Number of results to return (default: 5) - options: Optional JSON string with query options - - Returns: - List of dicts with "id" and "score" keys. - """ - embedding_dim = len(embedding) - - emb_arr = (ctypes.c_float * embedding_dim)(*embedding) - emb_ptr = ctypes.cast(emb_arr, ctypes.POINTER(ctypes.c_float)) - emb_ptr_ptr = ctypes.pointer(emb_ptr) - - id_buffer = (ctypes.c_int * top_k)() - score_buffer = (ctypes.c_float * top_k)() - - id_ptr = ctypes.cast(id_buffer, ctypes.POINTER(ctypes.c_int)) - score_ptr = ctypes.cast(score_buffer, ctypes.POINTER(ctypes.c_float)) - - id_size = ctypes.c_size_t(top_k) - score_size = ctypes.c_size_t(top_k) - - id_ptr_ptr = ctypes.pointer(id_ptr) - score_ptr_ptr = ctypes.pointer(score_ptr) - - options_encoded = options.encode() if options else None - - result = _lib.cactus_index_query( - index, - emb_ptr_ptr, - 1, - embedding_dim, - options_encoded, - id_ptr_ptr, - ctypes.byref(id_size), - score_ptr_ptr, - ctypes.byref(score_size) - ) - - if result < 0: - return [] - - return [ - {"id": id_buffer[i], "score": score_buffer[i]} - for i in range(id_size.value) - ] - - -def cactus_index_compact(index): - """ - Compact the index to optimize storage and query performance. - - Args: - index: Index handle from cactus_index_init - - Returns: - 0 on success, -1 on error. - """ - return _lib.cactus_index_compact(index) - - -def cactus_index_destroy(index): - """Free index resources. Always call when done.""" - _lib.cactus_index_destroy(index) diff --git a/python/src/cli.py b/python/src/cli.py deleted file mode 100644 index b51fc4e4e..000000000 --- a/python/src/cli.py +++ /dev/null @@ -1,984 +0,0 @@ -#!/usr/bin/env python3 -import sys -import os -import argparse -import re -import subprocess -import shutil -import platform -from pathlib import Path - -SCRIPT_DIR = Path(__file__).parent.resolve() -PROJECT_ROOT = SCRIPT_DIR.parent.parent -DEFAULT_MODEL_ID = "LiquidAI/LFM2-1.2B" - -RED = '\033[0;31m' -GREEN = '\033[0;32m' -YELLOW = '\033[1;33m' -BLUE = '\033[0;34m' -NC = '\033[0m' - - -def print_color(color, message): - """Print a message with ANSI color codes.""" - print(f"{color}{message}{NC}") - - -def get_model_dir_name(model_id): - """Convert HuggingFace model ID to local directory name.""" - model_name = model_id.split('/')[-1] - model_name = model_name.lower() - return model_name - - -def get_weights_dir(model_id): - """Get the weights directory path for a model.""" - model_dir = get_model_dir_name(model_id) - return PROJECT_ROOT / "weights" / model_dir - - -def check_command(cmd): - """Check if a command is available in PATH.""" - return shutil.which(cmd) is not None - - -def run_command(cmd, cwd=None, check=True): - """Run a shell command and optionally exit on failure.""" - result = subprocess.run(cmd, cwd=cwd, shell=isinstance(cmd, str)) - if check and result.returncode != 0: - sys.exit(result.returncode) - return result - - -def cmd_download(args): - """Download and convert HuggingFace model weights to Cactus format.""" - model_id = args.model_id - weights_dir = get_weights_dir(model_id) - - if weights_dir.exists() and (weights_dir / "config.txt").exists(): - print_color(GREEN, f"Model weights found at {weights_dir}") - return 0 - - print() - print_color(YELLOW, f"Model weights not found. Downloading {model_id}...") - print("=" * 45) - - try: - import torch - from transformers import AutoTokenizer - except ImportError: - print_color(RED, "Error: Required Python packages not found.") - print("Please run: ./setup") - return 1 - - from .converter_llm import convert_hf_model_weights - from .converter_vlm import convert_processors - from .tokenizer import convert_hf_tokenizer - from .tensor_io import format_config_value - from .config_utils import is_lfm2_vl, pick_dtype, vision_weight_sanity_check - - weights_dir.mkdir(parents=True, exist_ok=True) - - precision = getattr(args, 'precision', 'MIXED') - cache_dir = getattr(args, 'cache_dir', None) - token = getattr(args, 'token', None) - - print(f"Converting {model_id} to {precision}...") - - from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, AutoModelForImageTextToText, AutoModel, AutoConfig - - try: - from transformers import Lfm2VlForConditionalGeneration - except ImportError: - Lfm2VlForConditionalGeneration = None - - is_vlm = 'vl' in model_id.lower() or 'vlm' in model_id.lower() - is_whisper = 'whisper' in model_id.lower() - - try: - if is_vlm: - missing_deps = [] - try: - from PIL import Image - except Exception: - missing_deps.append('Pillow') - try: - import num2words - except Exception: - missing_deps.append('num2words') - try: - import torchvision - except Exception: - missing_deps.append('torchvision') - - if missing_deps: - print_color(RED, f"Error: Missing packages for VLM: {', '.join(missing_deps)}") - print(f"Install with: pip install {' '.join(missing_deps)}") - return 1 - - processor = None - try: - processor = AutoProcessor.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - except Exception as proc_err: - if "TokenizersBackend" in str(proc_err) or "does not exist or is not currently imported" in str(proc_err): - print(f" Note: AutoProcessor failed, using fallback tokenizer loading...") - else: - raise - - cfg = AutoConfig.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - dtype = pick_dtype() - - if is_lfm2_vl(model_id, cfg) and Lfm2VlForConditionalGeneration is not None: - model = Lfm2VlForConditionalGeneration.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, dtype=dtype, token=token) - else: - model = AutoModelForImageTextToText.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, dtype=dtype, token=token) - - tokenizer = getattr(processor, "tokenizer", None) if processor else None - if tokenizer is None: - try: - tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - except Exception as tok_err: - if "TokenizersBackend" in str(tok_err) or "does not exist or is not currently imported" in str(tok_err): - from transformers import PreTrainedTokenizerFast - print(f" Note: Using PreTrainedTokenizerFast fallback for invalid tokenizer_class...") - tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id, cache_dir=cache_dir, token=token) - else: - raise - - if is_lfm2_vl(model_id, cfg) and not vision_weight_sanity_check(model): - print_color(RED, "Vision embeddings look randomly initialized.") - return 1 - - try: - convert_processors(processor, model_id, weights_dir, token=token) - except Exception as e: - print(f" Warning: convert_processors failed: {e}") - - elif is_whisper: - tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - model = AutoModel.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - - else: - tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - try: - model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - except ValueError: - model = AutoModel.from_pretrained(model_id, cache_dir=cache_dir, trust_remote_code=True, token=token) - - config = convert_hf_model_weights(model, weights_dir, precision, args) - - model_name_l = model_id.lower() - if 'extract' in model_name_l: - config['model_variant'] = 'extract' - elif 'vlm' in model_name_l: - config['model_variant'] = 'vlm' - elif 'rag' in model_name_l: - config['model_variant'] = 'rag' - else: - config.setdefault('model_variant', 'default') - - # Config precision stores the compute precision (weights are quantized, activations stay FP16) - if precision in ('INT8', 'INT4', 'MIXED'): - config['precision'] = "FP16" - else: - config['precision'] = precision - - config_path = weights_dir / "config.txt" - with open(config_path, 'w') as f: - for key, value in config.items(): - f.write(f"{key}={format_config_value(value)}\n") - - convert_hf_tokenizer(tokenizer, weights_dir, token=token) - - del model - del tokenizer - import torch - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - print_color(GREEN, f"Successfully downloaded and converted weights to {weights_dir}") - return 0 - - except Exception as e: - print_color(RED, f"Error: {e}") - return 1 - - -def cmd_build(args): - """Build the Cactus library and chat binary.""" - if getattr(args, 'apple', False): - return cmd_build_apple(args) - if getattr(args, 'android', False): - return cmd_build_android(args) - if getattr(args, 'flutter', False): - return cmd_build_flutter(args) - if getattr(args, 'python', False): - return cmd_build_python(args) - - print_color(BLUE, "Building Cactus chat...") - print("=" * 23) - - if not check_command('cmake'): - print_color(RED, "Error: CMake is not installed") - print(" macOS: brew install cmake") - print(" Ubuntu: sudo apt-get install cmake") - return 1 - - cactus_dir = PROJECT_ROOT / "cactus" - lib_path = cactus_dir / "build" / "libcactus.a" - - print_color(YELLOW, "Building Cactus library...") - build_script = cactus_dir / "build.sh" - if not build_script.exists(): - print_color(RED, f"Error: build.sh not found at {build_script}") - return 1 - result = run_command(str(build_script), cwd=cactus_dir, check=False) - if result.returncode != 0: - print_color(RED, "Failed to build cactus library") - return 1 - - tests_dir = PROJECT_ROOT / "tests" - build_dir = tests_dir / "build" - build_dir.mkdir(parents=True, exist_ok=True) - - print("Compiling chat.cpp...") - - chat_cpp = tests_dir / "chat.cpp" - if not chat_cpp.exists(): - print_color(RED, f"Error: chat.cpp not found at {chat_cpp}") - return 1 - - is_darwin = platform.system() == "Darwin" - - if is_darwin: - compiler = "clang++" - cmd = [ - compiler, "-std=c++20", "-O3", - f"-I{PROJECT_ROOT}", - str(chat_cpp), - str(lib_path), - "-o", "chat", - "-lcurl", - "-framework", "Accelerate", - "-framework", "CoreML", - "-framework", "Foundation" - ] - else: - compiler = "g++" - cmd = [ - compiler, "-std=c++20", "-O3", - f"-I{PROJECT_ROOT}", - str(chat_cpp), - str(lib_path), - "-o", "chat", - "-lcurl", - "-pthread" - ] - - if not check_command(compiler): - print_color(RED, f"Error: {compiler} is not installed") - return 1 - - result = subprocess.run(cmd, cwd=build_dir) - if result.returncode != 0: - print_color(RED, "Build failed") - return 1 - - print_color(GREEN, f"Build complete: {build_dir / 'chat'}") - return 0 - - -def cmd_build_apple(args): - """Build Cactus for Apple platforms (iOS/macOS).""" - print_color(BLUE, "Building Cactus for Apple platforms...") - print("=" * 40) - - if platform.system() != "Darwin": - print_color(RED, "Error: Apple builds require macOS") - return 1 - - build_script = PROJECT_ROOT / "apple" / "build.sh" - if not build_script.exists(): - print_color(RED, f"Error: build.sh not found at {build_script}") - return 1 - - result = run_command(str(build_script), cwd=PROJECT_ROOT / "apple", check=False) - if result.returncode != 0: - print_color(RED, "Apple build failed") - return 1 - - print_color(GREEN, "Apple build complete!") - return 0 - - -def cmd_build_android(args): - """Build Cactus for Android.""" - print_color(BLUE, "Building Cactus for Android...") - print("=" * 32) - - build_script = PROJECT_ROOT / "android" / "build.sh" - if not build_script.exists(): - print_color(RED, f"Error: build.sh not found at {build_script}") - return 1 - - result = run_command(str(build_script), cwd=PROJECT_ROOT / "android", check=False) - if result.returncode != 0: - print_color(RED, "Android build failed") - return 1 - - print_color(GREEN, "Android build complete!") - return 0 - - -def cmd_build_flutter(args): - """Build Cactus for Flutter (iOS, macOS, Android).""" - print_color(BLUE, "Building Cactus for Flutter...") - print("=" * 32) - - build_script = PROJECT_ROOT / "flutter" / "build.sh" - if not build_script.exists(): - print_color(RED, f"Error: build.sh not found at {build_script}") - return 1 - - result = run_command(str(build_script), cwd=PROJECT_ROOT / "flutter", check=False) - if result.returncode != 0: - print_color(RED, "Flutter build failed") - return 1 - - print_color(GREEN, "Flutter build complete!") - print() - print("Output:") - print(f" flutter/libcactus.so") - print(f" flutter/cactus-ios.xcframework") - print(f" flutter/cactus-macos.xcframework") - return 0 - - -def cmd_build_python(args): - """Build Cactus shared library for Python FFI.""" - print_color(BLUE, "Building Cactus for Python...") - print("=" * 30) - - if not check_command('cmake'): - print_color(RED, "Error: CMake is not installed") - print(" macOS: brew install cmake") - print(" Ubuntu: sudo apt-get install cmake") - return 1 - - cactus_dir = PROJECT_ROOT / "cactus" - build_script = cactus_dir / "build.sh" - if not build_script.exists(): - print_color(RED, f"Error: build.sh not found at {build_script}") - return 1 - - result = run_command(str(build_script), cwd=cactus_dir, check=False) - if result.returncode != 0: - print_color(RED, "Build failed") - return 1 - - if platform.system() == "Darwin": - lib_name = "libcactus.dylib" - else: - lib_name = "libcactus.so" - - lib_path = cactus_dir / "build" / lib_name - if not lib_path.exists(): - print_color(RED, f"Shared library not found at {lib_path}") - return 1 - - print_color(GREEN, "Python build complete!") - print(f"Library: {lib_path}") - return 0 - - -def cmd_run(args): - """Build, download model if needed, and start interactive chat.""" - model_id = args.model_id - - if not getattr(args, 'no_build', False): - build_result = cmd_build(args) - if build_result != 0: - return build_result - - local_path = Path(model_id) - if local_path.exists() and (local_path / "config.txt").exists(): - weights_dir = local_path - print_color(GREEN, f"Using local model: {weights_dir}") - else: - download_result = cmd_download(args) - if download_result != 0: - return download_result - weights_dir = get_weights_dir(model_id) - - chat_binary = PROJECT_ROOT / "tests" / "build" / "chat" - - if not chat_binary.exists(): - print_color(RED, f"Error: Chat binary not found at {chat_binary}") - return 1 - - os.system('clear' if platform.system() != 'Windows' else 'cls') - print_color(GREEN, f"Starting Cactus Chat with model: {model_id}") - print() - - os.execv(str(chat_binary), [str(chat_binary), str(weights_dir)]) - - -def cmd_eval(args): - model_id = getattr(args, 'model_id', DEFAULT_MODEL_ID) - - if PROJECT_ROOT.parent.name != 'evals': - print_color(RED, "Skipping internal eval checks: companion repo not found.") - return 1 - - if not getattr(args, 'no_build', False): - build_result = cmd_build(args) - if build_result != 0: - return build_result - - class DownloadArgs: - pass - - dlargs = DownloadArgs() - dlargs.model_id = model_id - dlargs.precision = getattr(args, 'precision', 'MIXED') - dlargs.cache_dir = getattr(args, 'cache_dir', None) - dlargs.token = getattr(args, 'token', None) - - download_result = cmd_download(dlargs) - if download_result != 0: - return download_result - - weights_dir = get_weights_dir(model_id) - extra = getattr(args, 'extra_args', None) or [] - - def extra_has_flag(flag: str) -> bool: - for a in extra: - if a == flag or a.startswith(flag + "="): - return True - return False - - mode_flags = [] - if getattr(args, 'tools', False): mode_flags.append('tools') - if getattr(args, 'llm', False): mode_flags.append('llm') - if getattr(args, 'stt', False): mode_flags.append('stt') - if getattr(args, 'vlm', False): mode_flags.append('vlm') - if getattr(args, 'embed', False): mode_flags.append('embed') - - if len(mode_flags) > 1: - print_color(RED, f"Error: choose only one eval mode flag, got: {' '.join(mode_flags)}") - return 1 - - mode = mode_flags[0] if mode_flags else "tools" - repo_root = PROJECT_ROOT.parent # evals/ - cwd = repo_root - - if mode == "tools": - eval_runner = repo_root / "tool-evals" / "run_eval_berk.py" - elif mode == "stt": - eval_runner = repo_root / "speech-evals" / "speech_eval.py" - elif mode == "llm": - eval_runner = repo_root / "text-evals" / "perplexity_eval.py" - elif mode == "vlm": - eval_runner = repo_root / "video-evals" / "run_benchmarks.py" - elif mode == "embed": - print_color(RED, f"Error: eval mode '{mode}' is not supported in this repo layout") - return 1 - else: - print_color(RED, f"Error: unknown eval mode '{mode}'") - return 1 - - if not eval_runner.exists(): - print_color(RED, f"Eval runner not found at {eval_runner}") - return 1 - - cmd = [sys.executable, str(eval_runner)] - - if mode == "vlm": - if not extra_has_flag("--model"): - cmd += ["--model", str(weights_dir)] - if not extra_has_flag("--all") and not extra_has_flag("--benchmarks"): - cmd += ["--all"] - else: - if not extra_has_flag("--model-path"): - cmd += ["--model-path", str(weights_dir)] - - if mode == "llm" and not extra_has_flag("--model-id"): - cmd += ["--model-id", str(model_id)] - - if mode == "stt" and not extra_has_flag("--dataset-path"): - default_dataset_path = repo_root / "speech-evals" / "dataset-retrieval" - cmd += ["--dataset-path", str(default_dataset_path)] - - if not extra_has_flag("--output-dir"): - if mode == "tools": - default_out = repo_root / "tool-evals" / "results" - elif mode == "stt": - default_out = repo_root / "speech-evals" / "results" - elif mode == "llm": - default_out = repo_root / "text-evals" / "results" - else: - default_out = None - if default_out is not None: - cmd += ["--output-dir", str(default_out)] - - cmd += extra - - print_color(BLUE, f"[cactus] launching {mode} eval runner") - print(" ".join(cmd)) - - env = os.environ.copy() - if mode == "vlm": - ffi_dir = str(repo_root / "cactus" / "tools" / "src") - existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ffi_dir if not existing else (ffi_dir + os.pathsep + existing) - - r = subprocess.run(cmd, cwd=str(cwd), env=env) - return r.returncode - - - -def cmd_test(args): - """Run the Cactus test suite.""" - print_color(BLUE, "Running test suite...") - print("=" * 20) - - precision = getattr(args, 'precision', None) - if precision: - model_id = getattr(args, 'model', 'LiquidAI/LFM2-VL-450M') - weights_dir = get_weights_dir(model_id) - - if weights_dir.exists(): - print_color(YELLOW, f"Removing existing weights at {weights_dir} to regenerate with {precision}...") - shutil.rmtree(weights_dir) - - class DownloadArgs: - pass - dl_args = DownloadArgs() - dl_args.model_id = model_id - dl_args.precision = precision - dl_args.cache_dir = None - dl_args.token = getattr(args, 'token', None) - - download_result = cmd_download(dl_args) - if download_result != 0: - return download_result - - test_script = PROJECT_ROOT / "tests" / "run.sh" - - if not test_script.exists(): - print_color(RED, f"Error: Test script not found at {test_script}") - return 1 - - cmd = [str(test_script)] - - if args.model: - cmd.extend(["--model", args.model]) - if args.transcribe_model: - cmd.extend(["--transcribe_model", args.transcribe_model]) - if getattr(args, 'no_rebuild', False): - cmd.append("--no-rebuild") - if args.android: - cmd.append("--android") - if args.ios: - cmd.append("--ios") - - result = subprocess.run(cmd, cwd=PROJECT_ROOT / "tests") - return result.returncode - - -def cmd_clean(args): - """Remove all build artifacts, caches, and downloaded weights.""" - print_color(BLUE, "Cleaning all build artifacts from Cactus project...") - print(f"Project root: {PROJECT_ROOT}") - print() - - def remove_if_exists(path): - if path.is_dir(): - print(f"Removing: {path}") - shutil.rmtree(path) - else: - print(f"Not found: {path}") - - remove_if_exists(PROJECT_ROOT / "cactus" / "build") - - remove_if_exists(PROJECT_ROOT / "android" / "build") - remove_if_exists(PROJECT_ROOT / "android" / "libs") - remove_if_exists(PROJECT_ROOT / "android" / "arm64-v8a") - - remove_if_exists(PROJECT_ROOT / "apple" / "build") - - remove_if_exists(PROJECT_ROOT / "tests" / "build") - - remove_if_exists(PROJECT_ROOT / "venv") - - remove_if_exists(PROJECT_ROOT / "weights") - - print() - print("Removing compiled libraries and frameworks...") - - so_count = 0 - for so_file in PROJECT_ROOT.rglob("*.so"): - so_file.unlink() - so_count += 1 - print(f"Removed {so_count} .so files" if so_count else "No .so files found") - - a_count = 0 - for a_file in PROJECT_ROOT.rglob("*.a"): - a_file.unlink() - a_count += 1 - print(f"Removed {a_count} .a files" if a_count else "No .a files found") - - bin_count = 0 - for bin_file in PROJECT_ROOT.rglob("*.bin"): - bin_file.unlink() - bin_count += 1 - print(f"Removed {bin_count} .bin files" if bin_count else "No .bin files found") - - xcf_count = 0 - for xcf_dir in PROJECT_ROOT.rglob("*.xcframework"): - if xcf_dir.is_dir(): - shutil.rmtree(xcf_dir) - xcf_count += 1 - print(f"Removed {xcf_count} .xcframework directories" if xcf_count else "No .xcframework directories found") - - pycache_count = 0 - for pycache_dir in PROJECT_ROOT.rglob("__pycache__"): - if pycache_dir.is_dir(): - shutil.rmtree(pycache_dir) - pycache_count += 1 - print(f"Removed {pycache_count} __pycache__ directories" if pycache_count else "No __pycache__ directories found") - - egg_count = 0 - for egg_dir in PROJECT_ROOT.rglob("*.egg-info"): - if egg_dir.is_dir(): - shutil.rmtree(egg_dir) - egg_count += 1 - print(f"Removed {egg_count} .egg-info directories" if egg_count else "No .egg-info directories found") - - print() - print_color(GREEN, "Clean complete!") - print("All build artifacts have been removed.") - print() - - # Re-run setup automatically - print_color(BLUE, "Re-running setup...") - setup_script = PROJECT_ROOT / "setup" - result = subprocess.run( - ["bash", "-c", f"source {setup_script} && pip install -e {PROJECT_ROOT / 'tools'} --quiet"], - cwd=PROJECT_ROOT - ) - if result.returncode == 0: - print_color(GREEN, "Setup complete!") - else: - print_color(YELLOW, "Setup had issues. Please run manually:") - print(" source ./setup") - return 0 - - -def merge_lora_adapter(base_model_id, lora_path, cache_dir=None, token=None): - """Merge a LoRA adapter into a base model and return the merged model.""" - try: - from peft import PeftModel - except ImportError: - print_color(RED, "Error: peft package required for LoRA merging") - print("Install with: pip install peft") - return None, None - - from transformers import AutoModelForCausalLM, AutoTokenizer - - print_color(YELLOW, f"Loading base model: {base_model_id}") - base_model = AutoModelForCausalLM.from_pretrained( - base_model_id, - cache_dir=cache_dir, - trust_remote_code=True, - token=token - ) - tokenizer = AutoTokenizer.from_pretrained( - base_model_id, - cache_dir=cache_dir, - trust_remote_code=True, - token=token - ) - - print_color(YELLOW, f"Loading LoRA adapter: {lora_path}") - model = PeftModel.from_pretrained(base_model, lora_path, token=token) - - print_color(YELLOW, "Merging LoRA weights into base model...") - merged_model = model.merge_and_unload() - - print_color(GREEN, "LoRA merge complete") - return merged_model, tokenizer - - -def cmd_convert(args): - """Convert a HuggingFace model to a custom output directory.""" - import tempfile - - model_id = args.model_name - output_dir = args.output_dir - lora_path = getattr(args, 'lora', None) - - if output_dir is None: - output_dir = get_weights_dir(model_id) - else: - output_dir = Path(output_dir) - - cache_dir = getattr(args, 'cache_dir', None) - token = getattr(args, 'token', None) - - temp_merged_dir = None - - if lora_path: - merged_model, tokenizer = merge_lora_adapter(model_id, lora_path, cache_dir, token) - if merged_model is None: - return 1 - - temp_merged_dir = tempfile.mkdtemp(prefix="cactus_lora_merged_") - print_color(YELLOW, f"Saving merged model to temp directory: {temp_merged_dir}") - merged_model.save_pretrained(temp_merged_dir) - tokenizer.save_pretrained(temp_merged_dir) - - del merged_model - import torch - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - model_id = temp_merged_dir - - class DownloadArgs: - pass - - download_args = DownloadArgs() - download_args.model_id = model_id - download_args.precision = args.precision - download_args.cache_dir = cache_dir - download_args.token = token - - original_get_weights = get_weights_dir - - def custom_weights_dir(mid): - return output_dir - - import src.cli as cli_module - cli_module.get_weights_dir = custom_weights_dir - - try: - result = cmd_download(download_args) - return result - finally: - cli_module.get_weights_dir = original_get_weights - if temp_merged_dir and Path(temp_merged_dir).exists(): - print_color(YELLOW, "Cleaning up temp directory...") - shutil.rmtree(temp_merged_dir) - - -def create_parser(): - """Create the argument parser with all subcommands.""" - parser = argparse.ArgumentParser( - formatter_class=argparse.RawDescriptionHelpFormatter, - usage=argparse.SUPPRESS, - description=""" - - ----------------------------------------------------------------- - - How to use the Cactus Repo/CLI: - - ----------------------------------------------------------------- - - cactus run opens playground for the model - auto downloads and spins up - - Optional flags: - --precision MIXED|INT4|INT8|FP16 default: MIXED - --token HF token (for gated models) - - ----------------------------------------------------------------- - - cactus download downloads model to ./weights - see supported weights on ReadMe - - Optional flags: - --precision MIXED|INT4|INT8|FP16 quantization (default: MIXED) - --token HuggingFace API token - - ----------------------------------------------------------------- - - cactus convert [output_dir] converts model to custom directory - supports LoRA adapter merging - - Optional flags: - --precision MIXED|INT4|INT8|FP16 quantization (default: MIXED) - --lora LoRA adapter path to merge - --token HuggingFace API token - - ----------------------------------------------------------------- - - cactus build builds cactus for ARM chips - output: build/libcactus.a - - Optional flags: - --apple build for Apple (iOS/macOS) - --android build for Android - --flutter build for Flutter (all platforms) - --python build shared lib for Python FFI - - ----------------------------------------------------------------- - - cactus test runs unit tests and benchmarks - all must pass for contributions - - Optional flags: - --model default: LFM2-VL-450M - --transcribe_model default: whisper-small - --precision INT4|INT8|FP16 regenerates weights with precision - --no-rebuild skip building library and tests - --ios run on connected iPhone - --android run on connected Android - - ----------------------------------------------------------------- - - cactus clean removes all build artifacts - - ----------------------------------------------------------------- - - cactus --help shows these instructions - - ----------------------------------------------------------------- - - Python bindings: - - Cactus python package is auto installed for researchers and testing - Please see python/example.py and run the following instructions. - - 1. cactus build - 2. cactus download LiquidAI/LFM2-VL-450M - 3. python python/example.py - - Note: Use any supported model - - ----------------------------------------------------------------- -""" - ) - - subparsers = parser.add_subparsers(dest='command') - subparsers.required = False - - for action in parser._actions: - if isinstance(action, argparse._SubParsersAction): - action.help = argparse.SUPPRESS - - parser._action_groups = [] - - download_parser = subparsers.add_parser('download', help='Download and convert model weights') - download_parser.add_argument('model_id', nargs='?', default=DEFAULT_MODEL_ID, - help=f'HuggingFace model ID (default: {DEFAULT_MODEL_ID})') - download_parser.add_argument('--precision', choices=['MIXED', 'INT4', 'INT8', 'FP16'], default='MIXED', - help='Quantization precision (default: MIXED)') - download_parser.add_argument('--cache-dir', help='Cache directory for HuggingFace models') - download_parser.add_argument('--token', help='HuggingFace API token') - - build_parser = subparsers.add_parser('build', help='Build the chat application') - build_parser.add_argument('--apple', action='store_true', - help='Build for Apple platforms (iOS/macOS)') - build_parser.add_argument('--android', action='store_true', - help='Build for Android') - build_parser.add_argument('--flutter', action='store_true', - help='Build for Flutter (iOS, macOS, Android)') - build_parser.add_argument('--python', action='store_true', - help='Build shared library for Python FFI') - - run_parser = subparsers.add_parser('run', help='Build, download (if needed), and run chat') - run_parser.add_argument('model_id', nargs='?', default=DEFAULT_MODEL_ID, - help=f'HuggingFace model ID (default: {DEFAULT_MODEL_ID})') - run_parser.add_argument('--precision', choices=['MIXED', 'INT4', 'INT8', 'FP16'], default='MIXED', - help='Quantization precision (default: MIXED)') - run_parser.add_argument('--cache-dir', help='Cache directory for HuggingFace models') - run_parser.add_argument('--token', help='HuggingFace API token') - run_parser.add_argument('--no-build', action='store_true', help='Skip building Cactus before running') - - eval_parser = subparsers.add_parser('eval', help='Run evaluation scripts located outside the cactus submodule') - eval_parser.add_argument('model_id', nargs='?', default=DEFAULT_MODEL_ID, - help=f'HuggingFace model ID (default: {DEFAULT_MODEL_ID})') - eval_parser.add_argument('--precision', choices=['MIXED', 'INT4', 'INT8', 'FP16'], default='MIXED', - help='Quantization precision (default: MIXED)') - eval_parser.add_argument('--cache-dir', help='Cache directory for HuggingFace models') - eval_parser.add_argument('--token', help='HuggingFace API token') - eval_parser.add_argument('--no-build', action='store_true', help='Skip building Cactus before running evals') - eval_parser.add_argument('--tools', action='store_true', help='Run tools evals (default)') - eval_parser.add_argument('--vlm', action='store_true', help='Run VLM-specific evals') - eval_parser.add_argument('--stt', action='store_true', help='Run speech-to-text evals') - eval_parser.add_argument('--llm', action='store_true', help='Run LLM evals') - eval_parser.add_argument('--embed', action='store_true', help='Run embedding evals') - - test_parser = subparsers.add_parser('test', help='Run the test suite') - test_parser.add_argument('--model', default='LiquidAI/LFM2-VL-450M', - help='Model to use for tests') - test_parser.add_argument('--transcribe_model', default='openai/whisper-small', - help='Transcribe model to use') - test_parser.add_argument('--precision', choices=['MIXED', 'INT4', 'INT8', 'FP16'], - help='Regenerate weights with this precision (deletes existing weights)') - test_parser.add_argument('--no-rebuild', action='store_true', - help='Skip building cactus library and tests') - test_parser.add_argument('--token', help='HuggingFace API token') - test_parser.add_argument('--android', action='store_true', - help='Run tests on Android') - test_parser.add_argument('--ios', action='store_true', - help='Run tests on iOS') - - clean_parser = subparsers.add_parser('clean', help='Remove all build artifacts') - - convert_parser = subparsers.add_parser('convert', help='Convert model to custom output directory') - convert_parser.add_argument('model_name', help='HuggingFace model name') - convert_parser.add_argument('output_dir', nargs='?', default=None, - help='Output directory (default: weights/)') - convert_parser.add_argument('--precision', choices=['MIXED', 'INT4', 'INT8', 'FP16'], default='MIXED', - help='Quantization precision (default: MIXED)') - convert_parser.add_argument('--cache-dir', help='Cache directory for HuggingFace models') - convert_parser.add_argument('--token', help='HuggingFace API token') - convert_parser.add_argument('--lora', help='Path to LoRA adapter (local path or HuggingFace ID) to merge before conversion') - - return parser - - -def preprocess_eval_args(parser, argv): - args, unknown = parser.parse_known_args(argv) - - if getattr(args, 'command', None) == 'eval': - setattr(args, 'extra_args', unknown) - return args - - if unknown: - parser.error(f"unrecognized arguments: {' '.join(unknown)}") - - return args - - - -def main(): - """Main entry point for the Cactus CLI.""" - parser = create_parser() - - argv = sys.argv[1:] - args = preprocess_eval_args(parser, argv) - - if args.command == 'download': - sys.exit(cmd_download(args)) - elif args.command == 'build': - sys.exit(cmd_build(args)) - elif args.command == 'run': - sys.exit(cmd_run(args)) - elif args.command == 'test': - sys.exit(cmd_test(args)) - elif args.command == 'eval': - sys.exit(cmd_eval(args)) - elif args.command == 'clean': - sys.exit(cmd_clean(args)) - elif args.command == 'convert': - sys.exit(cmd_convert(args)) - else: - parser.print_help() - sys.exit(1) - - -if __name__ == '__main__': - main() diff --git a/python/src/config_utils.py b/python/src/config_utils.py deleted file mode 100644 index 3146872d3..000000000 --- a/python/src/config_utils.py +++ /dev/null @@ -1,148 +0,0 @@ -from typing import Optional, Any, Dict - - -def cfg_get(c, key, default=None): - """Get a config value from a dict or object, with fallback default.""" - if c is None: - return default - try: - if isinstance(c, dict): - return c.get(key, default) - except Exception: - pass - try: - return getattr(c, key, default) - except Exception: - return default - - -def detect_model_type(cfg, config, output_dir=None): - """Detect the model architecture type from config.""" - model_type_str = cfg_get(cfg, 'model_type', cfg_get(config, 'model_type', '')).lower() - - if 'gemma' in model_type_str: - return 'gemma' - elif 'lfm2' in model_type_str: - return 'lfm2' - elif 'qwen' in model_type_str: - return 'qwen' - elif 'llama' in model_type_str: - if output_dir and 'smol' in str(output_dir): - return 'smol' - else: - return 'llama' - elif 'bert' in model_type_str: - return 'bert' - elif 'whisper' in model_type_str: - return 'whisper' - else: - if model_type_str: - print(f" Warning: Unknown model type '{model_type_str}', defaulting to 'qwen'") - return 'qwen' - - -def extract_base_config(cfg, config): - """Extract base model configuration parameters.""" - return { - 'vocab_size': cfg_get(cfg, 'vocab_size', cfg_get(config, 'vocab_size', 0)), - 'hidden_dim': cfg_get(cfg, 'hidden_size', cfg_get(cfg, 'hidden_dim', 0)), - 'num_layers': int(cfg_get(cfg, 'num_hidden_layers', cfg_get(cfg, 'num_layers', 0) or 0)), - 'attention_heads': cfg_get(cfg, 'num_attention_heads', 0), - 'attention_kv_heads': cfg_get(cfg, 'num_key_value_heads', cfg_get(cfg, 'num_attention_heads', 0)), - 'ffn_intermediate_dim': cfg_get(cfg, 'intermediate_size', cfg_get(cfg, 'n_inner', 0)), - 'context_length': cfg_get(cfg, 'max_position_embeddings', cfg_get(cfg, 'max_sequence_length', 0)), - 'rope_theta': cfg_get(cfg, 'rope_theta', cfg_get(config, 'rope_theta', 10000.0)), - 'attention_head_dim': int(cfg_get(cfg, 'head_dim', int(cfg_get(cfg, 'hidden_size', cfg_get(cfg, 'hidden_dim', 0)) // max(1, cfg_get(cfg, 'num_attention_heads', 1))))), - 'layer_norm_eps': cfg_get(cfg, 'layer_norm_eps', cfg_get(cfg, 'layer_norm_epsilon', cfg_get(cfg, 'rms_norm_eps', 1e-6))), - 'num_experts': cfg_get(cfg, 'num_experts', 0), - 'num_shared_experts': cfg_get(cfg, 'num_shared_experts', 0), - 'num_top_experts': cfg_get(cfg, 'moe_top_k', cfg_get(cfg, 'num_top_experts', 0)), - 'moe_every_n_layers': cfg_get(cfg, 'moe_every_n_layers', 0), - } - - -def extract_vision_config(config, vision_cfg): - """Extract vision encoder configuration for VLM models.""" - vision_hidden = int(cfg_get(vision_cfg, 'hidden_size', 0)) - vision_image_size = cfg_get(vision_cfg, 'image_size', cfg_get(vision_cfg, 'size', {}).get('longest_edge', 0) if isinstance(cfg_get(vision_cfg, 'size', {}), dict) else cfg_get(vision_cfg, 'image_size', 0)) - vision_patch = int(cfg_get(vision_cfg, 'patch_size', 0)) - vision_heads = int(cfg_get(vision_cfg, 'num_attention_heads', 0)) - vision_num_layers = int(cfg_get(vision_cfg, 'num_hidden_layers', cfg_get(vision_cfg, 'num_layers', 0) or 0)) - num_channels = int(cfg_get(vision_cfg, 'num_channels', 3)) - visual_tokens_per_img = 0 - try: - if vision_patch > 0 and vision_image_size > 0: - per_side = vision_image_size // vision_patch - visual_tokens_per_img = per_side * per_side - except Exception: - visual_tokens_per_img = 0 - - pixel_shuffle_factor = int(cfg_get(config, 'scale_factor', cfg_get(vision_cfg, 'scale_factor', 1) or 1)) - downsample_factor = int(cfg_get(config, 'downsample_factor', 2)) - - return { - 'vision_hidden_size': int(vision_hidden), - 'vision_num_layers': int(vision_num_layers), - 'vision_image_size': int(vision_image_size), - 'vision_patch_size': int(vision_patch), - 'vision_attention_heads': int(vision_heads), - 'vision_embed_dim': int(vision_hidden), - 'num_channels': int(num_channels), - 'visual_tokens_per_img': int(visual_tokens_per_img), - 'use_pixel_shuffle': bool(pixel_shuffle_factor > 1), - 'pixel_shuffle_factor': int(pixel_shuffle_factor), - 'use_image_tokens': bool(cfg_get(config, 'image_token_id', None) is not None), - 'use_layout_tags': False, - 'downsample_factor': int(downsample_factor), - } - - -def extract_lfm2_config(cfg): - """Extract LFM2-specific configuration parameters.""" - layer_types = getattr(cfg, 'layer_types', []) - conv_L_cache = getattr(cfg, 'conv_L_cache', 0) - return { - 'layer_types': layer_types, - 'conv_L_cache': conv_L_cache, - } - - -def is_vlm_model(config): - """Check if a model config indicates a vision-language model.""" - text_cfg = cfg_get(config, 'text_config', None) - vision_cfg = cfg_get(config, 'vision_config', None) - return text_cfg is not None or vision_cfg is not None - - -def is_lfm2_vl(model_name, cfg): - """Check if the model is an LFM2 vision-language model.""" - if getattr(cfg, "model_type", None) == "lfm2-vl": - return True - name = (model_name or "").lower() - return "lfm2-vl" in name - - -def pick_dtype(): - """Select the best torch dtype based on hardware capabilities.""" - import torch - if torch.cuda.is_available(): - if torch.cuda.is_bf16_supported(): - return torch.bfloat16 - return torch.float16 - return torch.float32 - - -def vision_weight_sanity_check(model): - """Verify vision tower weights are properly initialized.""" - ok = True - vt = getattr(model, "vision_tower", None) - try: - emb = vt.vision_model.embeddings - w_mean = emb.patch_embedding.weight.detach().abs().mean().item() - p_mean = emb.position_embedding.weight.detach().abs().mean().item() - print(f"[sanity] |patch W| mean={w_mean:.5f} |pos W| mean={p_mean:.5f}") - if w_mean < 1e-3 or p_mean < 1e-3: - ok = False - except Exception: - pass - return ok diff --git a/python/src/converter_llm.py b/python/src/converter_llm.py deleted file mode 100644 index 8a26ff7dc..000000000 --- a/python/src/converter_llm.py +++ /dev/null @@ -1,249 +0,0 @@ -import re -from pathlib import Path - -try: - import torch -except ImportError: - torch = None - -from .tensor_io import save_tensor_with_header, create_quantization_stats, print_quantization_summary -from .config_utils import cfg_get, detect_model_type, extract_base_config, extract_vision_config, extract_lfm2_config, is_vlm_model -from .weight_patterns import ( - EMBED_NAMES, OUTPUT_NAMES, OUTPUT_NORM_NAMES, LAYER_PREFIXES, - VISION_ITEMS, PROJECTOR_WEIGHTS, WHISPER_GLOBAL_WEIGHTS, - get_layer_weight_patterns, get_vision_layer_weights -) -from .precision_config import count_model_parameters, MixedPrecisionConfig, SMALL_MODEL_THRESHOLD, LARGE_MODEL_THRESHOLD - - -def convert_hf_model_weights(model, output_dir, precision='INT8', args=None): - """Convert HuggingFace model weights to Cactus binary format.""" - quantization_stats = create_quantization_stats() - - param_count = count_model_parameters(model) - - print(f"Model size: {param_count / 1e6:.1f}M parameters") - - if precision == 'MIXED': - if param_count > 0 and param_count <= SMALL_MODEL_THRESHOLD: - print(f"Small model (<= 300M) - using INT8 for all quantized tensors.") - mixed_config = MixedPrecisionConfig(force_int8=True) - elif param_count >= LARGE_MODEL_THRESHOLD: - print(f"Large model (>= 1B) - using INT4 for all quantized tensors.") - mixed_config = MixedPrecisionConfig(force_int4=True) - else: - print(f"Medium model (300M-1B) - using mixed precision (INT4 tolerant, INT8 sensitive)") - mixed_config = MixedPrecisionConfig() - else: - mixed_config = MixedPrecisionConfig() - - state_dict = model.state_dict() - config = model.config - saved_tensor_full_names = set() - - text_cfg = cfg_get(config, 'text_config', None) - vision_cfg = cfg_get(config, 'vision_config', None) - is_vlm = text_cfg is not None or vision_cfg is not None - - cfg = text_cfg if text_cfg is not None else config - - tie_word_embeddings = getattr(config, 'tie_word_embeddings', False) - model_type_str = cfg_get(cfg, 'model_type', cfg_get(config, 'model_type', '')).lower() - - detected_model_type = detect_model_type(cfg, config, output_dir) - - model_config = extract_base_config(cfg, config) - model_config['tie_word_embeddings'] = tie_word_embeddings - model_config['model_type'] = detected_model_type - - if is_vlm and vision_cfg is not None: - model_config.update(extract_vision_config(config, vision_cfg)) - - if detected_model_type == 'lfm2': - model_config.update(extract_lfm2_config(cfg)) - - num_layers = model_config['num_layers'] - - embedding_found = False - for name in EMBED_NAMES: - if name in state_dict: - embedding_tensor = state_dict[name] - save_tensor_with_header(embedding_tensor, output_dir / "token_embeddings.weights", precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(name) - embedding_found = True - break - - if model_type_str == 'nomic_bert': - if 'embeddings.word_embeddings.weight' in state_dict: - fused_embedding_tensor = state_dict['embeddings.word_embeddings.weight'] + state_dict.get('embeddings.token_type_embeddings.weight', torch.zeros([1])) - save_tensor_with_header(fused_embedding_tensor, output_dir / "token_embeddings.weights", precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add('embeddings.word_embeddings.weight') - if 'embeddings.token_type_embeddings.weight' in state_dict: - saved_tensor_full_names.add('embeddings.token_type_embeddings.weight') - embedding_found = True - - elif model_type_str == 'whisper': - for name, save_name in WHISPER_GLOBAL_WEIGHTS: - if name in state_dict: - save_tensor_with_header(state_dict[name], output_dir / save_name, precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(name) - embedding_found = True - - if embedding_found: - embedding_norm_names = {'emb_ln.weight': 'embedding_layernorm.weight', 'emb_ln.bias': 'embedding_layernorm.bias'} - for name, file_name in embedding_norm_names.items(): - if name in state_dict: - save_tensor_with_header(state_dict[name], output_dir / file_name, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(name) - - if not tie_word_embeddings or is_vlm: - for name in OUTPUT_NAMES: - if name in state_dict: - tensor = state_dict[name] - save_tensor_with_header(tensor, output_dir / "output_weight.weights", precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(name) - break - - for name in OUTPUT_NORM_NAMES: - if name in state_dict: - tensor = state_dict[name] - save_tensor_with_header(tensor, output_dir / "output_norm.weights", precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(name) - break - - if is_vlm: - for key, outname in VISION_ITEMS: - if key in state_dict: - save_tensor_with_header(state_dict[key], output_dir / outname, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(key) - - for key, outname in PROJECTOR_WEIGHTS: - if key in state_dict: - save_tensor_with_header(state_dict[key], output_dir / outname, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(key) - - max_v_idx = -1 - vision_prefix = None - for k in state_dict.keys(): - m = re.search(r'model\.vision_tower\.vision_model\.encoder\.layers\.(\d+)\.', k) - if m: - vision_prefix = 'model.vision_tower.vision_model.encoder.layers.' - try: - idx = int(m.group(1)) - if idx > max_v_idx: - max_v_idx = idx - except Exception: - pass - if not vision_prefix: - m = re.search(r'model\.vision_model\.encoder\.layers\.(\d+)\.', k) - if m: - vision_prefix = 'model.vision_model.encoder.layers.' - try: - idx = int(m.group(1)) - if idx > max_v_idx: - max_v_idx = idx - except Exception: - pass - - if not vision_prefix: - vision_prefix = 'model.vision_model.encoder.layers.' - - vision_layers = max_v_idx + 1 if max_v_idx >= 0 else 0 - - for i_v in range(vision_layers): - vpref = f'{vision_prefix}{i_v}.' - vision_layer_weights = get_vision_layer_weights(i_v, vpref) - for fname, out in vision_layer_weights: - if fname in state_dict: - save_tensor_with_header(state_dict[fname], output_dir / out, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(fname) - missing_tensors = [] - for i in range(num_layers): - layer_prefixes = [p.format(i=i) for p in LAYER_PREFIXES] - - existing_prefixes = set() - for prefix in layer_prefixes: - for key in state_dict.keys(): - if key.startswith(prefix): - existing_prefixes.add(prefix) - - if not existing_prefixes: - missing_tensors.append((i, "", [""])) - continue - - weight_patterns = get_layer_weight_patterns(i, precision, model_type_str) - - for layer_prefix in existing_prefixes: - for name_patterns, tensor_precision, output_name, should_transpose in weight_patterns: - found = False - for pattern in name_patterns: - full_name = layer_prefix + pattern - if full_name in state_dict: - tensor = state_dict[full_name] - if pattern.startswith('attn.Wqkv.') and model_type_str == 'nomic_bert': - if tensor.ndim == 1: - tensor = tensor.reshape(3, -1) - elif tensor.ndim == 2: - tensor = tensor.reshape(3, -1, tensor.size(-1)) - else: - raise ValueError(f"Invalid tensor shape: {tensor.shape}") - for j, ch in enumerate(['q', 'k', 'v']): - channel_output_name = output_name.replace('{channel}', ch) - save_tensor_with_header(tensor[j], output_dir / channel_output_name, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(full_name) - found = True - break - elif model_type_str == 'nomic_bert' and pattern.startswith('mlp.experts.') and 'bias' not in pattern: - num_experts = model_config['num_experts'] - if tensor.ndim != 2: - raise ValueError(f"Invalid tensor shape: {tensor.shape}") - tensor = tensor.reshape(num_experts, -1, tensor.size(-1)) - for expert_idx in range(num_experts): - expert_tensor = tensor[expert_idx] - expert_output_name = output_name.replace('{channel}', str(expert_idx)) - save_tensor_with_header(expert_tensor, output_dir / expert_output_name, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(full_name) - found = True - break - if model_type_str == 'whisper': - temp = layer_prefix[:layer_prefix.find('.')] + "." + output_name - save_tensor_with_header(tensor, output_dir / temp, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - else: - save_tensor_with_header(tensor, output_dir / output_name, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(full_name) - found = True - break - - if not found and 'c_attn.weight' in name_patterns[0]: - attn_name = layer_prefix + 'attn.c_attn.weight' - if attn_name in state_dict: - combined_weight = state_dict[attn_name] - hidden_size = combined_weight.shape[0] - q_weight = combined_weight[:, :hidden_size] - k_weight = combined_weight[:, hidden_size:2*hidden_size] - v_weight = combined_weight[:, 2*hidden_size:] - - save_tensor_with_header(q_weight, output_dir / f'layer_{i}_attn_q.weights', precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - save_tensor_with_header(k_weight, output_dir / f'layer_{i}_attn_k.weights', precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - save_tensor_with_header(v_weight, output_dir / f'layer_{i}_attn_v.weights', precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type, layer_idx=i, num_layers=num_layers, mixed_config=mixed_config) - saved_tensor_full_names.add(attn_name) - found = True - - if saved_tensor_full_names != set(state_dict.keys()): - print(f"Warning: Unsaved tensors: {set(state_dict.keys()) - saved_tensor_full_names}") - - if not found: - missing_tensors.append((i, output_name, name_patterns)) - - if missing_tensors: - missing_report = output_dir / "missing_weights.txt" - with open(missing_report, 'w') as fh: - fh.write("# Missing tensors during conversion\n") - for layer_idx, output_name, patterns in missing_tensors: - pattern_list = ', '.join(patterns) - fh.write(f"layer={layer_idx}, output={output_name}, patterns=[{pattern_list}]\n") - print(f"Warning: {len(missing_tensors)} tensors were not exported. See {missing_report.name} for details.") - - print_quantization_summary(quantization_stats, args) - - return model_config diff --git a/python/src/converter_vlm.py b/python/src/converter_vlm.py deleted file mode 100644 index 67d6b8310..000000000 --- a/python/src/converter_vlm.py +++ /dev/null @@ -1,335 +0,0 @@ -import re -import sys -from pathlib import Path - -try: - import torch -except ImportError: - torch = None - -from .tensor_io import save_tensor_with_header, create_quantization_stats, print_quantization_summary -from .config_utils import cfg_get, is_lfm2_vl, pick_dtype, vision_weight_sanity_check -from .weight_patterns import ( - EMBED_NAMES, OUTPUT_NAMES, OUTPUT_NORM_NAMES, - VISION_ITEMS, PROJECTOR_WEIGHTS, CONNECTOR_KEYS, - get_layer_weight_patterns, get_vision_layer_weights -) - - -def convert_hf_model_weights_vlm(model, output_dir, precision='INT8', args=None): - """Convert VLM (Vision-Language Model) weights to Cactus binary format.""" - quantization_stats = create_quantization_stats() - - state_dict = model.state_dict() - config = model.config - saved_tensor_full_names = set() - - tie_word_embeddings = getattr(config, 'tie_word_embeddings', False) - - text_cfg = cfg_get(config, 'text_config', None) - vision_cfg = cfg_get(config, 'vision_config', None) - - text_vocab = cfg_get(text_cfg, 'vocab_size', cfg_get(config, 'vocab_size', 0)) - text_hidden = cfg_get(text_cfg, 'hidden_size', cfg_get(text_cfg, 'hidden_dim', 0)) - text_num_layers = int(cfg_get(text_cfg, 'num_hidden_layers', cfg_get(text_cfg, 'num_layers', 0) or 0)) - text_attention_heads = int(cfg_get(text_cfg, 'num_attention_heads', 0)) - text_attention_kv_heads = int(cfg_get(text_cfg, 'num_key_value_heads', cfg_get(text_cfg, 'num_attention_heads', 0))) - text_ffn = int(cfg_get(text_cfg, 'intermediate_size', 0)) - text_context = int(cfg_get(text_cfg, 'max_position_embeddings', cfg_get(text_cfg, 'max_sequence_length', 0))) - text_rope = cfg_get(text_cfg, 'rope_theta', cfg_get(config, 'rope_theta', 10000.0)) - text_head_dim = int(cfg_get(text_cfg, 'head_dim', int(text_hidden // max(1, text_attention_heads)))) - - vision_hidden = int(cfg_get(vision_cfg, 'hidden_size', 0)) - vision_image_size = cfg_get(vision_cfg, 'image_size', cfg_get(vision_cfg, 'size', {}).get('longest_edge', 0) if isinstance(cfg_get(vision_cfg, 'size', {}), dict) else cfg_get(vision_cfg, 'image_size', 0)) - vision_patch = int(cfg_get(vision_cfg, 'patch_size', 0)) - vision_heads = int(cfg_get(vision_cfg, 'num_attention_heads', 0)) - vision_num_layers = int(cfg_get(vision_cfg, 'num_hidden_layers', cfg_get(vision_cfg, 'num_layers', 0) or 0)) - num_channels = int(cfg_get(vision_cfg, 'num_channels', cfg_get(vision_cfg, 'num_channels', 3))) - vision_embed_dim = int(vision_hidden) - visual_tokens_per_img = 0 - try: - if vision_patch > 0 and vision_image_size > 0: - per_side = vision_image_size // vision_patch - visual_tokens_per_img = per_side * per_side - except Exception: - visual_tokens_per_img = 0 - - pixel_shuffle_factor = int(cfg_get(config, 'scale_factor', cfg_get(vision_cfg, 'scale_factor', 1) or 1)) - use_pixel_shuffle = bool(pixel_shuffle_factor > 1) - use_image_tokens = bool(cfg_get(config, 'image_token_id', None) is not None) - use_layout_tags = False - - model_type_str = cfg_get(text_cfg, 'model_type', None) or cfg_get(config, 'model_type', '') - if 'smolvlm' in model_type_str: - detected_model_type = 'smolvlm' - else: - detected_model_type = 'smolvlm' - print(f" Warning: Unknown VLM model type '{model_type_str}', defaulting to 'smolvlm'") - - - model_config = { - 'vocab_size': int(text_vocab), - 'model_type': detected_model_type, - 'hidden_dim': int(text_hidden), - 'num_layers': int(text_num_layers), - 'attention_heads': int(text_attention_heads), - 'attention_kv_heads': int(text_attention_kv_heads), - 'ffn_intermediate_dim': int(text_ffn), - 'context_length': int(text_context), - 'rope_theta': float(text_rope), - 'attention_head_dim': int(text_head_dim), - 'vision_hidden_size': int(vision_hidden), - 'vision_num_layers': int(vision_num_layers), - 'vision_image_size': int(vision_image_size), - 'vision_patch_size': int(vision_patch), - 'vision_attention_heads': int(vision_heads), - 'vision_embed_dim': int(vision_embed_dim), - 'num_channels': int(num_channels), - 'visual_tokens_per_img': int(visual_tokens_per_img), - 'use_pixel_shuffle': bool(use_pixel_shuffle), - 'pixel_shuffle_factor': int(pixel_shuffle_factor), - 'use_image_tokens': bool(use_image_tokens), - 'use_layout_tags': bool(use_layout_tags), - 'tie_word_embeddings': tie_word_embeddings - } - - embed_names = ['model.embed_tokens.weight', 'embed_tokens.weight', 'embeddings.weight', 'transformer.wte.weight', 'model.text_model.embed_tokens.weight'] - for name in embed_names: - if name in state_dict: - save_tensor_with_header(state_dict[name], output_dir / "token_embeddings.weights", precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - break - - if not tie_word_embeddings: - output_names = ['lm_head.weight', 'output.weight', 'transformer.lm_head.weight', 'model.text_model.lm_head.weight'] - for name in output_names: - if name in state_dict: - tensor = state_dict[name] - save_tensor_with_header(tensor, output_dir / "output_weight.weights", precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - break - - output_norm_names = ['model.norm.weight', 'norm.weight', 'final_layernorm.weight', 'transformer.ln_f.weight', 'model.text_model.norm.weight'] - for name in output_norm_names: - if name in state_dict: - tensor = state_dict[name] - save_tensor_with_header(tensor, output_dir / "output_norm.weights", precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - break - - vision_items = [ - ('model.vision_model.embeddings.patch_embedding.weight', 'vision_patch_embedding.weights'), - ('model.vision_model.embeddings.patch_embedding.bias', 'vision_patch_embedding.bias.weights'), - ('model.vision_model.embeddings.position_embedding.weight', 'vision_position_embedding.weights'), - ('model.vision_model.post_layernorm.weight', 'vision_post_layernorm.weights'), - ('model.vision_model.post_layernorm.bias', 'vision_post_layernorm.bias.weights') - ] - for key, outname in vision_items: - if key in state_dict: - save_tensor_with_header(state_dict[key], output_dir / outname, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - - max_v_idx = -1 - vision_prefix = None - for k in state_dict.keys(): - m = re.search(r'model\.vision_tower\.vision_model\.encoder\.layers\.(\d+)\.', k) - if m: - vision_prefix = 'model.vision_tower.vision_model.encoder.layers.' - try: - idx = int(m.group(1)) - if idx > max_v_idx: - max_v_idx = idx - except Exception: - pass - if not vision_prefix: - m = re.search(r'model\.vision_model\.encoder\.layers\.(\d+)\.', k) - if m: - vision_prefix = 'model.vision_model.encoder.layers.' - try: - idx = int(m.group(1)) - if idx > max_v_idx: - max_v_idx = idx - except Exception: - pass - - if not vision_prefix: - vision_prefix = 'model.vision_model.encoder.layers.' - - vision_layers = max_v_idx + 1 if max_v_idx >= 0 else 0 - - for i_v in range(vision_layers): - vpref = f'{vision_prefix}{i_v}.' - for fname, out in [ - (vpref + 'layer_norm1.weight', f'vision_layer_{i_v}_layer_norm1.weights'), - (vpref + 'layer_norm1.bias', f'vision_layer_{i_v}_layer_norm1.bias.weights'), - (vpref + 'layer_norm2.weight', f'vision_layer_{i_v}_layer_norm2.weights'), - (vpref + 'layer_norm2.bias', f'vision_layer_{i_v}_layer_norm2.bias.weights') - ]: - if fname in state_dict: - save_tensor_with_header(state_dict[fname], output_dir / out, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(fname) - - for fname, out in [ - (vpref + 'mlp.fc1.weight', f'vision_layer_{i_v}_ffn_fc1.weights'), - (vpref + 'mlp.fc1.bias', f'vision_layer_{i_v}_ffn_fc1.bias.weights'), - (vpref + 'mlp.fc2.weight', f'vision_layer_{i_v}_ffn_fc2.weights'), - (vpref + 'mlp.fc2.bias', f'vision_layer_{i_v}_ffn_fc2.bias.weights') - ]: - if fname in state_dict: - save_tensor_with_header(state_dict[fname], output_dir / out, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(fname) - - for fname, out in [ - (vpref + 'self_attn.q_proj.weight', f'vision_layer_{i_v}_self_attn_q.weights'), - (vpref + 'self_attn.k_proj.weight', f'vision_layer_{i_v}_self_attn_k.weights'), - (vpref + 'self_attn.v_proj.weight', f'vision_layer_{i_v}_self_attn_v.weights'), - (vpref + 'self_attn.out_proj.weight', f'vision_layer_{i_v}_self_attn_out.weights'), - (vpref + 'self_attn.q_proj.bias', f'vision_layer_{i_v}_self_attn_q.bias.weights'), - (vpref + 'self_attn.k_proj.bias', f'vision_layer_{i_v}_self_attn_k.bias.weights'), - (vpref + 'self_attn.v_proj.bias', f'vision_layer_{i_v}_self_attn_v.bias.weights'), - (vpref + 'self_attn.out_proj.bias', f'vision_layer_{i_v}_self_attn_out.bias.weights') - ]: - if fname in state_dict: - save_tensor_with_header(state_dict[fname], output_dir / out, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(fname) - - for key, outname in PROJECTOR_WEIGHTS: - if key in state_dict: - save_tensor_with_header(state_dict[key], output_dir / outname, precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(key) - - for ck in CONNECTOR_KEYS: - if ck in state_dict: - save_tensor_with_header(state_dict[ck], output_dir / 'connector_proj.weights', precision, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(ck) - break - - num_layers = model_config['num_layers'] - missing_tensors = [] - for i in range(num_layers): - - layer_prefixes = [f'model.language_model.layers.{i}.', f'model.text_model.layers.{i}.', - f'model.layers.{i}.', f'layers.{i}.', f'transformer.h.{i}.', f'encoder.layers.{i}.'] - - layer_prefix = None - for prefix in layer_prefixes: - if any(key.startswith(prefix) for key in state_dict.keys()): - layer_prefix = prefix - break - - if not layer_prefix: - continue - - conv_patterns = [ - ('conv.conv.weight', f'layer_{i}_conv_depthwise.weights'), - ('conv.in_proj.weight', f'layer_{i}_conv_in_proj.weights'), - ('conv.out_proj.weight', f'layer_{i}_conv_out_proj.weights'), - ] - for suffix, outname in conv_patterns: - fname = layer_prefix + suffix - if fname in state_dict: - save_tensor_with_header(state_dict[fname], output_dir / outname, 'FP16', stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(fname) - - weight_patterns = [ - (['self_attn.q_proj.weight', 'attn.q_proj.weight', 'attn.c_attn.weight'], precision, f'layer_{i}_attn_q.weights', False), - (['self_attn.k_proj.weight', 'attn.k_proj.weight'], precision, f'layer_{i}_attn_k.weights', False), - (['self_attn.v_proj.weight', 'attn.v_proj.weight'], precision, f'layer_{i}_attn_v.weights', False), - (['self_attn.out_proj.weight', 'self_attn.o_proj.weight', 'attn.o_proj.weight', 'attn.c_proj.weight'], precision, f'layer_{i}_attn_output.weights', False), - (['operator_norm.weight', 'input_layernorm.weight', 'ln_1.weight'], precision, f'layer_{i}_input_norm.weights', False), - (['self_attn.q_norm.weight', 'self_attn.q_layernorm.weight'], precision, f'layer_{i}_attn_q_norm.weights', False), - (['self_attn.k_norm.weight', 'self_attn.k_layernorm.weight'], precision, f'layer_{i}_attn_k_norm.weights', False), - (['feed_forward.w1.weight', 'mlp.gate_proj.weight', 'mlp.c_fc.weight'], precision, f'layer_{i}_ffn_gate.weights', False), - (['feed_forward.w3.weight', 'mlp.up_proj.weight'], precision, f'layer_{i}_ffn_up.weights', False), - (['feed_forward.w2.weight', 'mlp.down_proj.weight', 'mlp.c_proj.weight'], precision, f'layer_{i}_ffn_down.weights', False), - (['ffn_norm.weight', 'post_attention_layernorm.weight', 'ln_2.weight'], precision, f'layer_{i}_post_attn_norm.weights', False), - (['pre_feedforward_layernorm.weight'], precision, f'layer_{i}_pre_ffn_norm.weights', False), - (['post_feedforward_layernorm.weight'], precision, f'layer_{i}_post_ffn_norm.weights', False), - (['attn.Wqkv.bias'], precision, f'layer_{i}_attn_{{channel}}.bias', False), - (['attn.Wqkv.weight'], precision, f'layer_{i}_attn_{{channel}}.weights', False), - (['attn.out_proj.bias'], precision, f'layer_{i}_attn_output.bias', False), - (['attn.out_proj.weight'], precision, f'layer_{i}_attn_output.weights', False), - (['mlp.fc1.bias'], precision, f'layer_{i}_mlp_fc1.bias', False), - (['mlp.fc1.weight'], precision, f'layer_{i}_mlp_fc1.weights', False), - (['mlp.fc2.bias'], precision, f'layer_{i}_mlp_fc2.bias', False), - (['mlp.fc2.weight'], precision, f'layer_{i}_mlp_fc2.weights', False), - (['norm1.bias'], precision, f'layer_{i}_norm1.bias', False), - (['norm1.weight'], precision, f'layer_{i}_norm1.weights', False), - (['norm2.bias'], precision, f'layer_{i}_norm2.bias', False), - (['norm2.weight'], precision, f'layer_{i}_norm2.weights', False), - (['mlp.experts.bias'], precision, f'layer_{i}_mlp_experts.bias', False), - (['mlp.experts.mlp.w1'], precision, f'layer_{i}_mlp_expert_{{channel}}.mlp1.weights', False), - (['mlp.experts.mlp.w2'], precision, f'layer_{i}_mlp_expert_{{channel}}.mlp2.weights', True), - (['mlp.router.layer.weight'], precision, f'layer_{i}_mlp_router.layer.weights', False), - ] - - for name_patterns, tensor_precision, output_name, should_transpose in weight_patterns: - found = False - for pattern in name_patterns: - full_name = layer_prefix + pattern - if full_name in state_dict: - tensor = state_dict[full_name] - if pattern.startswith('attn.Wqkv.') and model_type_str == 'nomic_bert': - if tensor.ndim == 1: - tensor = tensor.reshape(3, -1) - elif tensor.ndim == 2: - tensor = tensor.reshape(3, -1, tensor.size(-1)) - else: - raise ValueError(f"Invalid tensor shape: {tensor.shape}") - for j, ch in enumerate(['q', 'k', 'v']): - channel_output_name = output_name.replace('{channel}', ch) - save_tensor_with_header(tensor[j], output_dir / channel_output_name, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(full_name) - found = True - break - elif model_type_str == 'nomic_bert' and pattern.startswith('mlp.experts.') and 'bias' not in pattern: - num_experts = model_config['num_experts'] - if tensor.ndim != 2: - raise ValueError(f"Invalid tensor shape: {tensor.shape}") - tensor = tensor.reshape(num_experts, -1, tensor.size(-1)) - for expert_idx in range(num_experts): - expert_tensor = tensor[expert_idx] - expert_output_name = output_name.replace('{channel}', str(expert_idx)) - save_tensor_with_header(expert_tensor, output_dir / expert_output_name, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(full_name) - found = True - break - save_tensor_with_header(tensor, output_dir / output_name, tensor_precision, transpose=should_transpose, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(full_name) - found = True - break - - if not found and 'c_attn.weight' in name_patterns[0]: - attn_name = layer_prefix + 'attn.c_attn.weight' - if attn_name in state_dict: - combined_weight = state_dict[attn_name] - hidden_size = combined_weight.shape[0] - q_weight = combined_weight[:, :hidden_size] - k_weight = combined_weight[:, hidden_size:2*hidden_size] - v_weight = combined_weight[:, 2*hidden_size:] - - save_tensor_with_header(q_weight, output_dir / f'layer_{i}_attn_q.weights', precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - save_tensor_with_header(k_weight, output_dir / f'layer_{i}_attn_k.weights', precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - save_tensor_with_header(v_weight, output_dir / f'layer_{i}_attn_v.weights', precision, transpose=False, stats_tracker=quantization_stats, args=args, model_type=detected_model_type) - saved_tensor_full_names.add(attn_name) - found = True - - if saved_tensor_full_names != set(state_dict.keys()): - print(f"Warning: Unsaved tensors: {set(state_dict.keys()) - saved_tensor_full_names}") - - if not found: - missing_tensors.append((i, output_name, name_patterns)) - - if missing_tensors: - missing_report = output_dir / "missing_weights.txt" - with open(missing_report, 'w') as fh: - fh.write("# Missing tensors during conversion\n") - for layer_idx, output_name, patterns in missing_tensors: - pattern_list = ', '.join(patterns) - fh.write(f"layer={layer_idx}, output={output_name}, patterns=[{pattern_list}]\n") - print(f"Warning: {len(missing_tensors)} tensors were not exported. See {missing_report.name} for details.") - - print_quantization_summary(quantization_stats, args) - - return model_config - - -def convert_processors(processor, model_name, output_dir, token=None): - """Save VLM processor config files to the output directory.""" - pass diff --git a/python/src/precision_config.py b/python/src/precision_config.py deleted file mode 100644 index 22829a77c..000000000 --- a/python/src/precision_config.py +++ /dev/null @@ -1,248 +0,0 @@ -"""Mixed precision configuration for Cactus weight conversion. - -This module determines which tensors should use INT8 vs INT4 quantization -based on their sensitivity to quantization error. -""" - -import numpy as np -from typing import Optional -from dataclasses import dataclass - -try: - import torch -except ImportError: - torch = None - - -GROUP_SIZE = 128 -SMALL_MODEL_THRESHOLD = 300_000_000 -LARGE_MODEL_THRESHOLD = 1_000_000_000 - - -def count_model_parameters(model) -> int: - """Count trainable parameters in a PyTorch model. - - Args: - model: PyTorch model - - Returns: - Total number of trainable parameters - """ - if torch is None: - return 0 - - if hasattr(model, 'parameters'): - return sum(p.numel() for p in model.parameters()) - - if hasattr(model, 'state_dict'): - state_dict = model.state_dict() - elif isinstance(model, dict): - state_dict = model - else: - return 0 - - total_params = 0 - for name, tensor in state_dict.items(): - if hasattr(tensor, 'numel'): - total_params += tensor.numel() - - return total_params - - -SENSITIVE_PATTERNS = [ - 'embed', - 'lm_head', - 'output_weight', - 'attn_q', - 'attn_k', - 'layer_0_', - 'layer_1_', -] - -TOLERANT_PATTERNS = [ - 'attn_v', - 'ffn_gate', - 'ffn_up', - 'ffn_down', - 'attn_output', -] - -FP16_PATTERNS = [ - 'norm', - 'bias', - 'vision', -] - - -@dataclass -class MixedPrecisionConfig: - """Configuration for mixed precision quantization.""" - snr_threshold: float = 25.0 - sensitive_as_int8: bool = True - tolerant_as_int4: bool = True - last_n_layers_int8: int = 2 - force_int8: bool = False - force_int4: bool = False - - -def is_sensitive_tensor(output_name: str, layer_idx: Optional[int], num_layers: Optional[int], - config: MixedPrecisionConfig = None) -> bool: - """Determine if a tensor should use INT8 (more sensitive) or INT4 (more tolerant). - - Args: - output_name: The output filename for the tensor - layer_idx: The layer index (None for non-layer weights like embeddings) - num_layers: Total number of layers in the model - config: Mixed precision configuration - - Returns: - True if tensor should use INT8, False for INT4 - """ - if config is None: - config = MixedPrecisionConfig() - - name_lower = output_name.lower() - - for pattern in FP16_PATTERNS: - if pattern in name_lower: - return True - - if config.sensitive_as_int8: - for pattern in SENSITIVE_PATTERNS: - if pattern in name_lower: - return True - - if layer_idx is not None and num_layers is not None: - if layer_idx >= num_layers - config.last_n_layers_int8: - return True - - if config.tolerant_as_int4: - for pattern in TOLERANT_PATTERNS: - if pattern in name_lower: - return False - - return True - - -def determine_precision(output_name: str, base_precision: str, - layer_idx: Optional[int] = None, - num_layers: Optional[int] = None, - config: MixedPrecisionConfig = None) -> str: - """Determine the actual precision to use for a tensor. - - Args: - output_name: The output filename for the tensor - base_precision: The base precision setting ('MIXED', 'INT4', 'INT8', 'FP16') - layer_idx: The layer index (None for non-layer weights) - num_layers: Total number of layers in the model - config: Mixed precision configuration - - Returns: - The precision to use: 'INT8', 'INT4', or 'FP16' - """ - if base_precision != 'MIXED': - return base_precision - - name_lower = output_name.lower() - for pattern in FP16_PATTERNS: - if pattern in name_lower: - return 'FP16' - - if config is not None and config.force_int8: - return 'INT8' - - if config is not None and config.force_int4: - return 'INT4' - - if is_sensitive_tensor(output_name, layer_idx, num_layers, config): - return 'INT8' - else: - return 'INT4' - - -def extract_layer_index(output_name: str) -> Optional[int]: - """Extract the layer index from an output filename. - - Args: - output_name: The output filename (e.g., 'layer_5_attn_q.weights') - - Returns: - The layer index, or None if not a layer weight - """ - import re - match = re.search(r'layer_(\d+)_', output_name) - if match: - return int(match.group(1)) - return None - - -def trial_quantize_int4(data: np.ndarray) -> tuple: - """Trial INT4 quantization to compute quality metrics. - - Args: - data: The tensor data to quantize - - Returns: - Tuple of (mse, snr_db, cos_sim) - """ - if data.ndim != 2: - return 0.0, float('inf'), 1.0 - - N, K = data.shape - - if K % GROUP_SIZE != 0: - pad_k = GROUP_SIZE - (K % GROUP_SIZE) - data = np.pad(data, ((0, 0), (0, pad_k)), mode='constant', constant_values=0) - K = data.shape[1] - - num_groups = K // GROUP_SIZE - data_grouped = data.reshape(N, num_groups, GROUP_SIZE) - - group_abs_max = np.max(np.abs(data_grouped), axis=2) - scales = (group_abs_max / 7.0).astype(np.float32) - scales = np.maximum(scales, 1e-10) - - quantized = np.clip( - np.round(data_grouped / scales[:, :, np.newaxis]), - -8, 7 - ).astype(np.int8) - - dequantized = (quantized.astype(np.float32) * scales[:, :, np.newaxis]).reshape(N, K) - - original = data[:, :K] # Match shape after padding - mse = np.mean((original - dequantized) ** 2) - variance = np.var(original) - snr_db = 10 * np.log10(variance / mse) if mse > 0 else float('inf') - - original_flat = original.flatten() - dequant_flat = dequantized.flatten() - cos_sim = np.dot(original_flat, dequant_flat) / ( - np.linalg.norm(original_flat) * np.linalg.norm(dequant_flat) + 1e-10 - ) - - return mse, snr_db, cos_sim - - -def adaptive_precision_from_snr(data: np.ndarray, output_name: str, - snr_threshold: float = 25.0) -> str: - """Determine precision based on trial INT4 quantization SNR. - - If INT4 quantization yields SNR below threshold, use INT8 instead. - - Args: - data: The tensor data - output_name: The output filename - snr_threshold: Minimum acceptable SNR in dB for INT4 - - Returns: - 'INT4' if quality is acceptable, 'INT8' otherwise - """ - if data.ndim != 2: - return 'INT8' # Non-2D tensors use INT8 or FP16 - - mse, snr_db, cos_sim = trial_quantize_int4(data) - - if snr_db >= snr_threshold: - return 'INT4' - else: - return 'INT8' diff --git a/python/src/publish_to_hf.py b/python/src/publish_to_hf.py deleted file mode 100644 index d850873fb..000000000 --- a/python/src/publish_to_hf.py +++ /dev/null @@ -1,233 +0,0 @@ -import argparse -import hashlib -import json -import os -import shutil -import subprocess -from huggingface_hub import HfApi, hf_hub_download -from .cli import cmd_convert, get_weights_dir, PROJECT_ROOT - -STAGE_DIR = PROJECT_ROOT / "stage" - -MODELS = [ - "google/gemma-3-270m-it", - "google/functiongemma-270m-it", - "openai/whisper-small", - "LiquidAI/LFM2-350M", - "LiquidAI/LFM2-VL-450M", - "nomic-ai/nomic-embed-text-v2-moe", - "Qwen/Qwen3-0.6B", - "Qwen/Qwen3-Embedding-0.6B", - "LiquidAI/LFM2-700M", - "google/gemma-3-1b-it", - "LiquidAI/LFM2.5-1.2B-Instruct", - "LiquidAI/LFM2-1.2B-RAG", - "LiquidAI/LFM2-1.2B-Tool", - "openai/whisper-medium", - "LiquidAI/LFM2.5-VL-1.6B", - "Qwen/Qwen3-1.7B", -] - -PRO_MODELS = [ - "openai/whisper-small", - "LiquidAI/LFM2-VL-450M", - "openai/whisper-medium", - "LiquidAI/LFM2.5-VL-1.6B", -] - - -def sha256(file): - h = hashlib.sha256() - with open(file, "rb") as f: - for block in iter(lambda: f.read(1024 * 1024), b""): - h.update(block) - return h.hexdigest() - - -def zip_dir(source_dir, output_path): - subprocess.run( - ["find", ".", "-exec", "touch", "-t", "200310131122", "{}", "+"], - cwd=source_dir, - check=True, - ) - - subprocess.run( - ["zip", "-X", "-o", "-r", "-9", str(output_path), "."], - cwd=source_dir, - check=True, - capture_output=True, - ) - - -def get_model_name(model_id): - return model_id.split("/")[-1] - - -def export_model(model_id, token, precision): - args = argparse.Namespace( - model_name=model_id, output_dir=None, precision=precision, token=token - ) - if cmd_convert(args) != 0: - return None - return get_weights_dir(model_id) - - -def export_pro_weights(model_id, bits): - pro_repo = PROJECT_ROOT / "cactus-pro" - if not pro_repo.exists(): - return None - - build_script = pro_repo / "apple" / "build.sh" - if not build_script.exists(): - return None - - result = subprocess.run( - ["bash", str(build_script), "--model", model_id, "--bits", bits], - cwd=pro_repo, - capture_output=True, - ) - - if result.returncode != 0: - return None - - mlpackage = pro_repo / "apple" / "build" / "model.mlpackage" - return mlpackage if mlpackage.exists() else None - - -def stage_model(model_id, weights_dir, precision, bits): - model_name = get_model_name(model_id) - stage = STAGE_DIR / model_name - - if stage.exists(): - shutil.rmtree(stage) - stage.mkdir(parents=True, exist_ok=True) - - model_name_lower = model_name.lower() - weights_out = stage / "weights" / model_name_lower - shutil.move(str(weights_dir), str(weights_out)) - - model_zip = stage / "weights" / f"{model_name_lower}.zip" - zip_dir(weights_out, model_zip) - - fingerprint = hashlib.sha256() - fingerprint.update(sha256(model_zip).encode()) - - config = { - "model_type": model_name, - "precision": precision, - } - - if model_id in PRO_MODELS: - try: - mlpackage = export_pro_weights(model_id, bits) - - shutil.move(mlpackage, weights_out) - - model_pro_zip = stage / "weights" / f"{model_name_lower}-pro.zip" - zip_dir(weights_out, model_pro_zip) - - fingerprint.update(sha256(model_pro_zip).encode()) - config["bits"] = bits - except Exception: - print("Failed to export pro weights") - - shutil.rmtree(weights_out) - - config["fingerprint"] = fingerprint.hexdigest() - - with open(stage / "config.json", "w") as f: - json.dump(config, f, indent=2) - - return stage, config - - -def get_prev_config(api, repo, current): - try: - tags = api.list_repo_refs(repo_id=repo, repo_type="model").tags - versions = sorted([t.name for t in tags], reverse=True) - prev_ver = next((v for v in versions if v != current), None) - if not prev_ver: - return None - local = hf_hub_download( - repo_id=repo, - filename="config.json", - revision=prev_ver, - repo_type="model", - ) - with open(local) as f: - return json.load(f) - except Exception: - return None - - -def changed(curr, prev): - if not prev: - return True - return curr.get("fingerprint") != prev.get("fingerprint") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--version", required=True) - parser.add_argument("--org", required=True) - parser.add_argument("--precision", required=True) - parser.add_argument("--bits", required=True) - args = parser.parse_args() - - token = os.environ.get("HF_TOKEN") - if not token: - print("Error: HF_TOKEN not set") - return 1 - - api = HfApi(token=token) - - for model_id in MODELS: - name = get_model_name(model_id) - repo_id = f"{args.org}/{name}" - - stage_dir = None - try: - weights_dir = export_model(model_id, token, args.precision) - if not weights_dir: - print("Export failed") - continue - - stage_dir, config = stage_model( - model_id, weights_dir, args.precision, args.bits - ) - prev = get_prev_config(api, repo_id, args.version) - - api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True) - - if changed(config, prev): - api.upload_folder( - folder_path=str(stage_dir), - path_in_repo=".", - repo_id=repo_id, - repo_type="model", - commit_message=f"Upload {args.version}", - ) - print("Uploaded") - else: - print("Unchanged") - - info = api.repo_info(repo_id=repo_id, repo_type="model") - api.create_tag( - repo_id=repo_id, - tag=args.version, - revision=info.sha, - repo_type="model", - tag_message=f"Release {args.version}", - ) - print("Tagged release") - - except Exception: - print("Model processing failed") - finally: - if stage_dir and stage_dir.exists(): - shutil.rmtree(stage_dir) - print("Cleaned up stage directory") - - -if __name__ == "__main__": - main() diff --git a/python/src/tensor_io.py b/python/src/tensor_io.py deleted file mode 100644 index 84642aec3..000000000 --- a/python/src/tensor_io.py +++ /dev/null @@ -1,365 +0,0 @@ -import numpy as np -import struct -from pathlib import Path -from typing import Optional, Dict, Any, List - -try: - import torch -except ImportError: - torch = None - -from .precision_config import determine_precision, MixedPrecisionConfig - - -GROUP_SIZE = 128 - -CACTUS_MAGIC = b'CACT' -CACTUS_VERSION = 1 -CACTUS_ALIGNMENT = 32 - -FLAG_HAS_SCALES = 1 << 0 -FLAG_PAGE_ALIGNED = 1 << 1 -FLAG_TRANSPOSED = 1 << 2 - - -def align_offset(offset: int, alignment: int) -> int: - """Round up offset to next alignment boundary.""" - remainder = offset % alignment - if remainder == 0: - return offset - return offset + (alignment - remainder) - - -def compute_padding(current_offset: int, alignment: int) -> bytes: - """Compute padding bytes needed to reach alignment boundary.""" - aligned = align_offset(current_offset, alignment) - padding_size = aligned - current_offset - return b'\x00' * padding_size - - -def save_tensor_with_header(tensor, output_path, precision='FP16', transpose=False, stats_tracker=None, args=None, model_type=None, layer_idx=None, num_layers=None, mixed_config=None): - """Save a tensor to binary format with header metadata and group-wise INT8/INT4 quantization. - - Args: - tensor: The tensor to save (PyTorch or NumPy) - output_path: Path to save the tensor - precision: Quantization precision ('MIXED', 'INT4', 'INT8', 'FP16') - transpose: Whether to transpose 2D tensors - stats_tracker: Optional dict to track quantization statistics - args: Optional args object with additional settings - model_type: Model type string (e.g., 'gemma', 'llama') - layer_idx: Layer index for mixed precision (None for non-layer weights) - num_layers: Total number of layers for mixed precision calculation - mixed_config: MixedPrecisionConfig for controlling quantization behavior - """ - if torch is not None and isinstance(tensor, torch.Tensor): - data = tensor.detach().cpu().numpy() - else: - data = np.array(tensor) - - original_data = data.copy() - - if model_type == 'gemma' and 'norm' in str(output_path): - data = data + 1.0 - original_data = data.copy() - - # Determine actual precision for MIXED mode - output_name = str(output_path.name) if hasattr(output_path, 'name') else str(output_path) - if precision == 'MIXED': - precision = determine_precision(output_name, 'MIXED', layer_idx, num_layers, mixed_config) - - if precision in ('INT8', 'INT4'): - filename = output_path.name - if any(x in filename for x in ['norm', 'bias', 'vision']): - precision = 'FP16' - - shape = list(data.shape) - if transpose and len(shape) == 2: - data = data.T - original_data = original_data.T - shape = [shape[1], shape[0]] - - if precision == 'INT8': - if len(shape) == 2: - N, K = shape - - if K % GROUP_SIZE != 0: - pad_k = GROUP_SIZE - (K % GROUP_SIZE) - data = np.pad(data, ((0, 0), (0, pad_k)), mode='constant', constant_values=0) - original_data = np.pad(original_data, ((0, 0), (0, pad_k)), mode='constant', constant_values=0) - K = data.shape[1] - shape = [N, K] - - num_groups = K // GROUP_SIZE - - data_grouped = data.reshape(N, num_groups, GROUP_SIZE) - original_grouped = original_data.reshape(N, num_groups, GROUP_SIZE) - - group_abs_max = np.max(np.abs(data_grouped), axis=2) - scales = (group_abs_max / 127.0).astype(np.float32) - scales = np.maximum(scales, 1e-10) - - quantized = np.clip( - np.round(data_grouped / scales[:, :, np.newaxis]), - -128, 127 - ).astype(np.int8) - quantized_flat = quantized.reshape(N, K) - - dequantized = (quantized.astype(np.float32) * scales[:, :, np.newaxis]).reshape(N, K) - mse_error = np.mean((original_data - dequantized) ** 2) - snr_db = 10 * np.log10(np.var(original_data) / mse_error) if mse_error > 0 else float('inf') - - original_flat = original_data.flatten() - dequant_flat = dequantized.flatten() - cos_sim = np.dot(original_flat, dequant_flat) / (np.linalg.norm(original_flat) * np.linalg.norm(dequant_flat) + 1e-10) - - scales_fp16 = scales.astype(np.float16) - - elif len(shape) == 1: - K = shape[0] - - if K % GROUP_SIZE != 0: - pad_k = GROUP_SIZE - (K % GROUP_SIZE) - data = np.pad(data, (0, pad_k), mode='constant', constant_values=0) - original_data = np.pad(original_data, (0, pad_k), mode='constant', constant_values=0) - K = data.shape[0] - shape = [K] - - num_groups = K // GROUP_SIZE - N = 1 - - data_grouped = data.reshape(1, num_groups, GROUP_SIZE) - original_grouped = original_data.reshape(1, num_groups, GROUP_SIZE) - - group_abs_max = np.max(np.abs(data_grouped), axis=2) - scales = (group_abs_max / 127.0).astype(np.float32) - scales = np.maximum(scales, 1e-10) - - quantized = np.clip( - np.round(data_grouped / scales[:, :, np.newaxis]), - -128, 127 - ).astype(np.int8) - quantized_flat = quantized.reshape(K) - - dequantized = (quantized.astype(np.float32) * scales[:, :, np.newaxis]).reshape(K) - mse_error = np.mean((original_data - dequantized) ** 2) - snr_db = 10 * np.log10(np.var(original_data) / mse_error) if mse_error > 0 else float('inf') - cos_sim = np.dot(original_data, dequantized) / (np.linalg.norm(original_data) * np.linalg.norm(dequantized) + 1e-10) - - scales_fp16 = scales.astype(np.float16) - else: - precision = 'FP16' - - if precision == 'INT8': - if stats_tracker: - stats_tracker['int8_tensors'] += 1 - stats_tracker['quantized_parameters'] += original_data.size - stats_tracker['mse_values'].append(mse_error) - stats_tracker['snr_values'].append(snr_db) - stats_tracker['cos_sim_values'].append(cos_sim) - stats_tracker['total_tensors'] += 1 - stats_tracker['total_parameters'] += original_data.size - - with open(output_path, 'wb') as f: - ndim = len(shape) - data_bytes = quantized_flat.size - scales_bytes = scales_fp16.size * 2 - flags = FLAG_HAS_SCALES - if transpose: - flags |= FLAG_TRANSPOSED - - # Fixed 64-byte header - f.write(CACTUS_MAGIC) # 4 bytes - f.write(struct.pack(' 0 else float('inf') - original_flat = original_data.flatten() - dequant_flat = dequantized.flatten() - cos_sim = np.dot(original_flat, dequant_flat) / (np.linalg.norm(original_flat) * np.linalg.norm(dequant_flat) + 1e-10) - - scales_fp16 = scales.astype(np.float16) - - if stats_tracker: - stats_tracker['int4_tensors'] += 1 - stats_tracker['quantized_parameters'] += original_data.size - stats_tracker['mse_values'].append(mse_error) - stats_tracker['snr_values'].append(snr_db) - stats_tracker['cos_sim_values'].append(cos_sim) - stats_tracker['total_tensors'] += 1 - stats_tracker['total_parameters'] += original_data.size - - with open(output_path, 'wb') as f: - ndim = len(shape) - data_bytes = packed.size - scales_bytes = scales_fp16.size * 2 - flags = FLAG_HAS_SCALES - if transpose: - flags |= FLAG_TRANSPOSED - - # Fixed header - f.write(CACTUS_MAGIC) # 4 bytes - f.write(struct.pack(' 0: - mse_values = np.array(quantization_stats['mse_values']) - snr_values = np.array(quantization_stats['snr_values']) - cos_sim_values = np.array(quantization_stats['cos_sim_values']) - - print("\nQuantization Summary:") - print(f"MSE - Mean: {np.mean(mse_values):.2e}, Max: {np.max(mse_values):.2e}, Median: {np.median(mse_values):.2e}, Min: {np.min(mse_values):.2e}") - print(f"SNR - Mean: {np.mean(snr_values):.1f}dB, Max: {np.max(snr_values):.1f}dB, Median: {np.median(snr_values):.1f}dB, Min: {np.min(snr_values):.1f}dB") - print(f"CosSim - Mean: {np.mean(cos_sim_values):.6f}, Max: {np.max(cos_sim_values):.6f}, Median: {np.median(cos_sim_values):.6f}, Min: {np.min(cos_sim_values):.6f}") - print(f"Processed {int8_count} INT8 tensors, {int4_count} INT4 tensors, {fp16_count} FP16 tensors") diff --git a/python/src/tokenizer.py b/python/src/tokenizer.py deleted file mode 100644 index 4bc2f0815..000000000 --- a/python/src/tokenizer.py +++ /dev/null @@ -1,281 +0,0 @@ -import json -from pathlib import Path -from typing import Optional - -try: - from huggingface_hub import hf_hub_download -except ImportError: - hf_hub_download = None - - -def convert_hf_tokenizer(tokenizer, output_dir, token=None): - """Convert a HuggingFace tokenizer to Cactus format.""" - is_sentencepiece = False - tokenizer_model_path = None - - if hasattr(tokenizer, 'vocab_file'): - vocab_file = tokenizer.vocab_file - if vocab_file and vocab_file.endswith('.model'): - is_sentencepiece = True - tokenizer_model_path = vocab_file - - if not is_sentencepiece and hasattr(tokenizer, 'sp_model'): - is_sentencepiece = True - if hf_hub_download: - try: - tokenizer_model_path = hf_hub_download( - repo_id=tokenizer.name_or_path, - filename="tokenizer.model", - token=token, - ) - except Exception: - pass - - - tokenizer_json_data = {} - tokenizer_json_path = output_dir / "tokenizer.json" - try: - tokenizer.save_pretrained(output_dir) - if tokenizer_json_path.exists(): - with open(tokenizer_json_path, 'r', encoding='utf-8') as f: - tokenizer_json_data = json.load(f) - - unused_files = [ - "tokenizer_config.json", - "special_tokens_map.json", - "added_tokens.json", - "chat_template.jinja", - ] - for filename in unused_files: - filepath = output_dir / filename - if filepath.exists(): - filepath.unlink() - except Exception as e: - print(f" Warning: Could not save tokenizer JSON: {e}") - - vocab = tokenizer.get_vocab() - - id_to_token = [""] * len(vocab) - for token_str, token_id in vocab.items(): - if token_id < len(id_to_token): - id_to_token[token_id] = token_str - - vocab_output = output_dir / "vocab.txt" - - if is_sentencepiece: - with open(vocab_output, 'w', encoding='utf-8') as f: - for token_id, token_str in enumerate(id_to_token): - if token_str: - f.write(f"{token_id}\t{token_str}\n") - print(f" Saved SentencePiece vocabulary (ID\\ttoken format)") - else: - with open(vocab_output, 'w', encoding='utf-8') as f: - for token_str in id_to_token: - f.write(token_str + '\n') - print(f" Saved BPE vocabulary (line-by-line format)") - - - merges_output = output_dir / "merges.txt" - - def write_merges_file(merges_list): - with open(merges_output, 'w', encoding='utf-8', newline='') as f: - f.write("#version: 0.2\n") - for merge in merges_list: - f.write(f"{' '.join(merge)}\n") - - merges_written = False - - if not is_sentencepiece and tokenizer_json_data: - merges_from_json = tokenizer_json_data.get("model", {}).get("merges", []) or [] - write_merges_file(merges_from_json) - merges_written = True - - if not merges_written and hf_hub_download: - try: - import shutil - merges_file = hf_hub_download(repo_id=tokenizer.name_or_path, filename="merges.txt", token=token) - shutil.copy2(merges_file, merges_output) - merges_written = True - except Exception: - pass - - if not merges_written and hasattr(tokenizer, 'backend_tokenizer') and tokenizer.backend_tokenizer: - backend = tokenizer.backend_tokenizer - merges = [] - - if hasattr(backend, 'model'): - model = backend.model - if hasattr(model, 'merges'): - merges = model.merges - - write_merges_file(merges) - merges_written = True - - if not merges_written: - write_merges_file([]) - - - special_tokens = {} - special_token_ids = {} - - if hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id is not None: - special_token_ids['eos_token_id'] = tokenizer.eos_token_id - special_tokens[tokenizer.eos_token_id] = tokenizer.eos_token or "<|endoftext|>" - - if hasattr(tokenizer, 'pad_token_id') and tokenizer.pad_token_id is not None: - special_token_ids['pad_token_id'] = tokenizer.pad_token_id - special_tokens[tokenizer.pad_token_id] = tokenizer.pad_token or "<|endoftext|>" - - if hasattr(tokenizer, 'bos_token_id') and tokenizer.bos_token_id is not None: - special_token_ids['bos_token_id'] = tokenizer.bos_token_id - special_tokens[tokenizer.bos_token_id] = tokenizer.bos_token or "<|startoftext|>" - - if hasattr(tokenizer, 'unk_token_id') and tokenizer.unk_token_id is not None: - special_token_ids['unk_token_id'] = tokenizer.unk_token_id - special_tokens[tokenizer.unk_token_id] = tokenizer.unk_token or "<|unknown|>" - - additional_special_tokens = [] - if hasattr(tokenizer, 'additional_special_tokens'): - for token_str in tokenizer.additional_special_tokens or []: - token_id = tokenizer.convert_tokens_to_ids(token_str) - if token_id != tokenizer.unk_token_id: - special_tokens[token_id] = token_str - additional_special_tokens.append({"token": token_str, "id": token_id}) - - model_type = getattr(tokenizer, 'name_or_path', '').lower() - if 'gemma' in model_type: - gemma_special_tokens = { - '': None, - '': None, - '': None, - '': None, - # Gemma 3 function calling tokens - '': None, - '': None, - '': None, - '': None, - '': None, - '': None, - '': None - } - - vocab = tokenizer.get_vocab() - for token_str in gemma_special_tokens.keys(): - if token_str in vocab: - token_id = vocab[token_str] - gemma_special_tokens[token_str] = token_id - special_tokens[token_id] = token_str - print(f" Found Gemma special token: {token_str} (ID: {token_id})") - - missing_tokens = [k for k, v in gemma_special_tokens.items() if v is None] - if missing_tokens and is_sentencepiece and tokenizer_model_path: - try: - import sentencepiece as spm - sp = spm.SentencePieceProcessor(model_file=tokenizer_model_path) - for token_str in missing_tokens: - token_id = sp.piece_to_id(token_str) - if token_id != sp.unk_id(): - gemma_special_tokens[token_str] = token_id - special_tokens[token_id] = token_str - print(f" Found Gemma special token via SentencePiece: {token_str} (ID: {token_id})") - except Exception as e: - print(f" Warning: Could not check SentencePiece for Gemma tokens: {e}") - - if gemma_special_tokens[''] is None: - hardcoded_ids = { - '': 105, - '': 106 - } - for token_str, token_id in hardcoded_ids.items(): - if token_str in gemma_special_tokens and gemma_special_tokens[token_str] is None: - if token_id not in special_tokens: - gemma_special_tokens[token_str] = token_id - special_tokens[token_id] = token_str - print(f" Using hardcoded Gemma special token: {token_str} (ID: {token_id})") - - chat_template_data = {} - if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template: - chat_template_output = output_dir / "chat_template.jinja2" - with open(chat_template_output, 'w', encoding='utf-8') as f: - f.write(tokenizer.chat_template) - chat_template_data["chat_template"] = tokenizer.chat_template - - tokenizer_full_config = {} - added_tokens_decoder = {} - tool_tokens = {} - - try: - config_path = None - if hasattr(tokenizer, 'name_or_path') and hf_hub_download: - try: - config_path = hf_hub_download(repo_id=tokenizer.name_or_path, filename="tokenizer_config.json", token=token) - with open(config_path, 'r') as f: - tokenizer_full_config = json.load(f) - - if 'chat_template' in tokenizer_full_config and not chat_template_data: - chat_template_output = output_dir / "chat_template.jinja2" - with open(chat_template_output, 'w', encoding='utf-8') as f: - f.write(tokenizer_full_config['chat_template']) - chat_template_data["chat_template"] = tokenizer_full_config['chat_template'] - - if 'added_tokens_decoder' in tokenizer_full_config: - added_tokens_decoder = tokenizer_full_config['added_tokens_decoder'] - - print(" Extracting special tokens from tokenizer_config.json...") - for token_id_str, token_info in added_tokens_decoder.items(): - content = token_info.get('content', '') - token_id = int(token_id_str) - - tool_related = ['', '', - '', '', - '', '', - '', '', - # Gemma 3 function calling tokens - '', '', - '', '', - '', '', - ''] - - if any(x == content for x in tool_related): - tool_tokens[token_id] = token_info - print(f" Found tool token: {content} (ID: {token_id})") - special_tokens[token_id] = content - - except Exception as e: - print(f" Note: Could not load full tokenizer config: {e}") - pass - except Exception: - pass - - - special_tokens_output = output_dir / "special_tokens.json" - with open(special_tokens_output, 'w', encoding='utf-8') as f: - json.dump({ - **special_token_ids, - "vocab_size": len(vocab), - "model_max_length": getattr(tokenizer, 'model_max_length', 131072), - "special_tokens": special_tokens, - "additional_special_tokens": additional_special_tokens, - **chat_template_data - }, f, indent=2, ensure_ascii=False) - - - tokenizer_config_output = output_dir / "tokenizer_config.txt" - with open(tokenizer_config_output, 'w') as f: - f.write(f"vocab_size={len(vocab)}\n") - for key, value in special_token_ids.items(): - f.write(f"{key}={value}\n") - f.write(f"model_max_length={getattr(tokenizer, 'model_max_length', 131072)}\n") - - if is_sentencepiece: - f.write("tokenizer_type=sentencepiece\n") - else: - f.write("tokenizer_type=bpe\n") - - if chat_template_data: - f.write("has_chat_template=true\n") - else: - f.write("has_chat_template=false\n") - if len(tool_tokens) > 0: - f.write(f"has_tool_support=true\n") - f.write(f"tool_token_count={len(tool_tokens)}\n") diff --git a/python/src/weight_patterns.py b/python/src/weight_patterns.py deleted file mode 100644 index c248a54a4..000000000 --- a/python/src/weight_patterns.py +++ /dev/null @@ -1,169 +0,0 @@ -EMBED_NAMES = [ - 'model.language_model.embed_tokens.weight', - 'model.text_model.embed_tokens.weight', - 'model.embed_tokens.weight', - 'embed_tokens.weight', - 'embeddings.weight', - 'transformer.wte.weight' -] - -OUTPUT_NAMES = [ - 'lm_head.weight', - 'output.weight', - 'transformer.lm_head.weight', - 'model.text_model.lm_head.weight' -] - -OUTPUT_NORM_NAMES = [ - 'model.norm.weight', - 'norm.weight', - 'final_layernorm.weight', - 'transformer.ln_f.weight', - 'model.embedding_norm.weight', - 'model.language_model.embedding_norm.weight', - 'model.text_model.norm.weight' -] - -LAYER_PREFIXES = [ - 'model.language_model.layers.{i}.', - 'model.text_model.layers.{i}.', - 'model.layers.{i}.', - 'layers.{i}.', - 'transformer.h.{i}.', - 'encoder.layers.{i}.', - 'decoder.layers.{i}.', - 'model.encoder.layers.{i}.', - 'model.decoder.layers.{i}.' -] - -VISION_ITEMS = [ - ('model.vision_tower.vision_model.embeddings.patch_embedding.weight', 'vision_patch_embedding.weights'), - ('model.vision_model.embeddings.patch_embedding.weight', 'vision_patch_embedding.weights'), - ('model.vision_tower.vision_model.embeddings.patch_embedding.bias', 'vision_patch_embedding.bias.weights'), - ('model.vision_model.embeddings.patch_embedding.bias', 'vision_patch_embedding.bias.weights'), - ('model.vision_tower.vision_model.embeddings.position_embedding.weight', 'vision_position_embedding.weights'), - ('model.vision_model.embeddings.position_embedding.weight', 'vision_position_embedding.weights'), - ('model.vision_tower.vision_model.post_layernorm.weight', 'vision_post_layernorm.weights'), - ('model.vision_model.post_layernorm.weight', 'vision_post_layernorm.weights'), - ('model.vision_tower.vision_model.post_layernorm.bias', 'vision_post_layernorm.bias.weights'), - ('model.vision_model.post_layernorm.bias', 'vision_post_layernorm.bias.weights') -] - -PROJECTOR_WEIGHTS = [ - ('model.multi_modal_projector.linear_1.weight', 'projector_linear1.weights'), - ('model.multi_modal_projector.linear_1.bias', 'projector_linear1.bias.weights'), - ('model.multi_modal_projector.linear_2.weight', 'projector_linear2.weights'), - ('model.multi_modal_projector.linear_2.bias', 'projector_linear2.bias.weights'), - ('model.multi_modal_projector.layer_norm.weight', 'projector_layer_norm.weights'), - ('model.multi_modal_projector.layer_norm.bias', 'projector_layer_norm.bias.weights'), -] - -CONNECTOR_KEYS = [ - 'model.connector.modality_projection.proj.weight', - 'connector.modality_projection.proj.weight', - 'model.connector.proj.weight', - 'connector.proj.weight' -] - -WHISPER_GLOBAL_WEIGHTS = [ - ('decoder.embed_tokens.weight', 'decoder_token_embeddings.weights'), - ('decoder.embed_positions.weight', 'decoder_position_embeddings.weights'), - ('decoder.layer_norm.weight', 'decoder_norm.weights'), - ('decoder.layer_norm.bias', 'decoder_norm.bias'), - ('proj_out.weight', 'output_layer.weights'), - ('encoder.embed_positions.weight', 'encoder_position_embeddings.weights'), - ('encoder.conv1.bias', 'encoder_conv1_bias.bias'), - ('encoder.conv1.weight', 'encoder_conv1_weight.weights'), - ('encoder.conv2.bias', 'encoder_conv2_bias.bias'), - ('encoder.conv2.weight', 'encoder_conv2_weight.weights'), - ('encoder.layer_norm.bias', 'encoder_norm_bias.bias'), - ('encoder.layer_norm.weight', 'encoder_norm_weight.weights') -] - - -def get_layer_weight_patterns(i, precision, model_type=None): - is_whisper = model_type == 'whisper' - - patterns = [ - (['self_attn.q_proj.weight', 'attn.q_proj.weight', 'attn.c_attn.weight'], precision, f'layer_{i}_attn_q.weights', False) if not is_whisper else None, - (['self_attn.k_proj.weight', 'attn.k_proj.weight'], precision, f'layer_{i}_attn_k.weights', False) if not is_whisper else None, - (['self_attn.v_proj.weight', 'attn.v_proj.weight'], precision, f'layer_{i}_attn_v.weights', False) if not is_whisper else None, - (['self_attn.o_proj.weight', 'attn.o_proj.weight', 'attn.c_proj.weight', 'self_attn.out_proj.weight'], precision, f'layer_{i}_attn_output.weights', False) if not is_whisper else None, - (['input_layernorm.weight', 'ln_1.weight', 'operator_norm.weight'], precision, f'layer_{i}_input_norm.weights', False), - (['self_attn.q_norm.weight', 'self_attn.q_layernorm.weight'], precision, f'layer_{i}_attn_q_norm.weights', False), - (['self_attn.k_norm.weight', 'self_attn.k_layernorm.weight'], precision, f'layer_{i}_attn_k_norm.weights', False), - (['mlp.gate_proj.weight', 'mlp.c_fc.weight', 'feed_forward.w1.weight'], precision, f'layer_{i}_ffn_gate.weights', False), - (['mlp.up_proj.weight', 'feed_forward.w3.weight'], precision, f'layer_{i}_ffn_up.weights', False), - (['mlp.down_proj.weight', 'mlp.c_proj.weight', 'feed_forward.w2.weight'], precision, f'layer_{i}_ffn_down.weights', False), - (['post_attention_layernorm.weight', 'ln_2.weight', 'ffn_norm.weight'], precision, f'layer_{i}_post_attn_norm.weights', False), - (['pre_feedforward_layernorm.weight'], precision, f'layer_{i}_pre_ffn_norm.weights', False), - (['post_feedforward_layernorm.weight'], precision, f'layer_{i}_post_ffn_norm.weights', False), - (['conv.in_proj.weight'], precision, f'layer_{i}_conv_in_proj.weights', False), - (['conv.out_proj.weight'], precision, f'layer_{i}_conv_out_proj.weights', False), - (['conv.conv.weight'], precision, f'layer_{i}_conv_depthwise.weights', False), - (['attn.Wqkv.bias'], precision, f'layer_{i}_attn_{{channel}}.bias', False), - (['attn.Wqkv.weight'], precision, f'layer_{i}_attn_{{channel}}.weights', False), - (['attn.out_proj.bias'], precision, f'layer_{i}_attn_output.bias', False), - (['attn.out_proj.weight'], precision, f'layer_{i}_attn_output.weights', False), - (['mlp.fc1.bias'], precision, f'layer_{i}_mlp_fc1.bias', False), - (['mlp.fc1.weight'], precision, f'layer_{i}_mlp_fc1.weights', False), - (['mlp.fc2.bias'], precision, f'layer_{i}_mlp_fc2.bias', False), - (['mlp.fc2.weight'], precision, f'layer_{i}_mlp_fc2.weights', False), - (['norm1.bias'], precision, f'layer_{i}_norm1.bias', False), - (['norm1.weight'], precision, f'layer_{i}_norm1.weights', False), - (['norm2.bias'], precision, f'layer_{i}_norm2.bias', False), - (['norm2.weight'], precision, f'layer_{i}_norm2.weights', False), - (['mlp.experts.bias'], precision, f'layer_{i}_mlp_experts.bias', False), - (['mlp.experts.mlp.w1'], precision, f'layer_{i}_mlp_expert_{{channel}}.mlp1.weights', False), - (['mlp.experts.mlp.w2'], precision, f'layer_{i}_mlp_expert_{{channel}}.mlp2.weights', True), - (['mlp.router.layer.weight'], precision, f'layer_{i}_mlp_router.layer.weights', False), - (['encoder_attn.q_proj.weight'], precision, f'layer_{i}_encoder_attn_q.weights', False), - (['encoder_attn.k_proj.weight'], precision, f'layer_{i}_encoder_attn_k.weights', False), - (['encoder_attn.v_proj.weight'], precision, f'layer_{i}_encoder_attn_v.weights', False), - (['encoder_attn.out_proj.weight'], precision, f'layer_{i}_encoder_attn_output.weights', False), - (['encoder_attn.q_proj.bias'], precision, f'layer_{i}_encoder_attn_q.bias', False), - (['encoder_attn.v_proj.bias'], precision, f'layer_{i}_encoder_attn_v.bias', False), - (['encoder_attn.out_proj.bias'], precision, f'layer_{i}_encoder_attn_output.bias', False), - (['encoder_attn_layer_norm.weight'], precision, f'layer_{i}_encoder_attn_norm.weights', False), - (['encoder_attn_layer_norm.bias'], precision, f'layer_{i}_encoder_attn_norm.bias', False), - (['fc1.weight'], precision, f'layer_{i}_mlp_fc1.weights', False), - (['fc1.bias'], precision, f'layer_{i}_mlp_fc1.bias', False), - (['fc2.weight'], precision, f'layer_{i}_mlp_fc2.weights', False), - (['fc2.bias'], precision, f'layer_{i}_mlp_fc2.bias', False), - (['final_layer_norm.weight'], precision, f'layer_{i}_final_norm.weights', False), - (['final_layer_norm.bias'], precision, f'layer_{i}_final_norm.bias', False), - - # Whisper-only: separate self_attn_* outputs (non-Whisper uses attn_* above) - (['self_attn.q_proj.weight'], precision, f'layer_{i}_self_attn_q.weights', False) if is_whisper else None, - (['self_attn.k_proj.weight'], precision, f'layer_{i}_self_attn_k.weights', False) if is_whisper else None, - (['self_attn.v_proj.weight'], precision, f'layer_{i}_self_attn_v.weights', False) if is_whisper else None, - (['self_attn.q_proj.bias'], precision, f'layer_{i}_self_attn_q.bias', False) if is_whisper else None, - (['self_attn.v_proj.bias'], precision, f'layer_{i}_self_attn_v.bias', False) if is_whisper else None, - (['self_attn.out_proj.weight'], precision, f'layer_{i}_self_attn_output.weights', False) if is_whisper else None, - (['self_attn.out_proj.bias'], precision, f'layer_{i}_self_attn_output.bias', False) if is_whisper else None, - (['self_attn_layer_norm.weight'], precision, f'layer_{i}_self_attn_norm.weights', False), - (['self_attn_layer_norm.bias'], precision, f'layer_{i}_self_attn_norm.bias', False), - ] - - return [p for p in patterns if p is not None] - - -def get_vision_layer_weights(i_v, vpref): - return [ - (vpref + 'layer_norm1.weight', f'vision_layer_{i_v}_layer_norm1.weights'), - (vpref + 'layer_norm1.bias', f'vision_layer_{i_v}_layer_norm1.bias.weights'), - (vpref + 'layer_norm2.weight', f'vision_layer_{i_v}_layer_norm2.weights'), - (vpref + 'layer_norm2.bias', f'vision_layer_{i_v}_layer_norm2.bias.weights'), - (vpref + 'mlp.fc1.weight', f'vision_layer_{i_v}_ffn_fc1.weights'), - (vpref + 'mlp.fc1.bias', f'vision_layer_{i_v}_ffn_fc1.bias.weights'), - (vpref + 'mlp.fc2.weight', f'vision_layer_{i_v}_ffn_fc2.weights'), - (vpref + 'mlp.fc2.bias', f'vision_layer_{i_v}_ffn_fc2.bias.weights'), - (vpref + 'self_attn.q_proj.weight', f'vision_layer_{i_v}_self_attn_q.weights'), - (vpref + 'self_attn.k_proj.weight', f'vision_layer_{i_v}_self_attn_k.weights'), - (vpref + 'self_attn.v_proj.weight', f'vision_layer_{i_v}_self_attn_v.weights'), - (vpref + 'self_attn.out_proj.weight', f'vision_layer_{i_v}_self_attn_out.weights'), - (vpref + 'self_attn.q_proj.bias', f'vision_layer_{i_v}_self_attn_q.bias.weights'), - (vpref + 'self_attn.k_proj.bias', f'vision_layer_{i_v}_self_attn_k.bias.weights'), - (vpref + 'self_attn.v_proj.bias', f'vision_layer_{i_v}_self_attn_v.bias.weights'), - (vpref + 'self_attn.out_proj.bias', f'vision_layer_{i_v}_self_attn_out.bias.weights'), - ] diff --git a/python/test.py b/python/test.py new file mode 100644 index 000000000..5957eeb6c --- /dev/null +++ b/python/test.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python3 +"""Run all cactus python tests. + +Usage: + python python/test.py # run all tests + python python/test.py -v # verbose + python python/test.py -k graph # only graph tests +""" +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +if __name__ == "__main__": + loader = unittest.TestLoader() + suite = loader.discover(str(Path(__file__).parent / "tests"), pattern="test_*.py") + runner = unittest.TextTestRunner(verbosity=2 if "-v" in sys.argv else 1) + result = runner.run(suite) + sys.exit(0 if result.wasSuccessful() else 1) diff --git a/python/tests/__init__.py b/python/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/tests/bundles.py b/python/tests/bundles.py new file mode 100644 index 000000000..5a5aaf950 --- /dev/null +++ b/python/tests/bundles.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +WEIGHTS = PROJECT_ROOT / "weights" + + +def _read_model_type(bundle: Path) -> str: + for line in (bundle / "config.txt").read_text(encoding="utf-8").splitlines(): + if line.startswith("model_type="): + return line.split("=", 1)[1].strip() + return "" + + +def _valid_bundle(path: Path) -> bool: + return (path / "config.txt").exists() and (path / "components" / "manifest.json").exists() + + +def _iter_bundle_candidates(name: str): + """Yield the bare bundle dir for `name` plus any suffixed `name-cq*` variants + (e.g. `gemma-4-e2b-it-cq4`) so callers that know a model's bare stem still + find bundles built under the suffixed convention.""" + bare = WEIGHTS / name + if bare.exists(): + yield bare + if WEIGHTS.is_dir(): + for candidate in sorted(WEIGHTS.glob(f"{name}-cq*"), reverse=True): + if candidate.is_dir(): + yield candidate + + +def _find_bundle(preferred: list[str], types: set[str], on_missing=pytest.fail) -> Path: + for name in preferred: + for candidate in _iter_bundle_candidates(name): + if _valid_bundle(candidate) and _read_model_type(candidate) in types: + return candidate + if not WEIGHTS.is_dir(): + on_missing(f"Weights directory not found: {WEIGHTS}") + for candidate in sorted(WEIGHTS.iterdir()): + if candidate.is_dir() and _valid_bundle(candidate) and _read_model_type(candidate) in types: + return candidate + on_missing(f"No valid live-test bundle found under {WEIGHTS} for model types: {sorted(types)}") diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py new file mode 100644 index 000000000..60f832fd5 --- /dev/null +++ b/python/tests/test_bindings.py @@ -0,0 +1,494 @@ +"""Comprehensive tests for cactus Python bindings. + +Tests the helper functions, auto-serialization, edge cases, and package +structure. The actual C FFI calls require a compiled libcactus — those +are integration tests run via `cactus test`. +""" +import json +import ctypes +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent)) + + +# ── Import and structure tests ────────────────────────────────────── + + +class TestPackageStructure: + """Verify the package is importable and has the right public API.""" + + def test_version_exists(self): + from cactus import __version__ + assert __version__ and __version__[0].isdigit() + + def test_cli_module_importable(self): + from cactus.cli import main, create_parser + assert callable(main) + assert callable(create_parser) + + def test_model_module_importable(self): + from cactus.cli.model import ensure_weights, ensure_bundle, TranspileOptions + assert callable(ensure_weights) + assert callable(ensure_bundle) + + def test_public_api_reexports(self): + from cactus import ensure_model, get_weights_dir, get_model_dir_name + assert callable(ensure_model) + assert callable(get_weights_dir) + assert callable(get_model_dir_name) + + def test_py_typed_marker_exists(self): + marker = Path(__file__).parent.parent / "cactus" / "py.typed" + assert marker.exists(), "py.typed PEP 561 marker missing" + + +# ── Helper function tests ─────────────────────────────────────────── + + +class TestToJson: + """Test the _to_json auto-serialization helper.""" + + def setup_method(self): + from cactus.bindings.cactus import _to_json + self.to_json = _to_json + + def test_none_passthrough(self): + assert self.to_json(None) is None + + def test_dict_serialized(self): + result = self.to_json({"temperature": 0.7, "max_tokens": 100}) + assert isinstance(result, bytes) + parsed = json.loads(result) + assert parsed == {"temperature": 0.7, "max_tokens": 100} + + def test_list_serialized(self): + messages = [{"role": "user", "content": "Hello"}] + result = self.to_json(messages) + assert isinstance(result, bytes) + parsed = json.loads(result) + assert parsed == messages + + def test_string_encoded(self): + result = self.to_json('{"already": "json"}') + assert isinstance(result, bytes) + assert result == b'{"already": "json"}' + + def test_bytes_passthrough(self): + raw = b'{"raw": "bytes"}' + result = self.to_json(raw) + assert result is raw + + def test_empty_dict(self): + result = self.to_json({}) + assert json.loads(result) == {} + + def test_empty_list(self): + result = self.to_json([]) + assert json.loads(result) == [] + + def test_nested_structure(self): + data = {"messages": [{"role": "user", "content": "test"}], "options": {"temp": 0.5}} + result = self.to_json(data) + assert json.loads(result) == data + + def test_unicode_content(self): + data = [{"role": "user", "content": "Привет мир 你好世界"}] + result = self.to_json(data) + parsed = json.loads(result) + assert parsed[0]["content"] == "Привет мир 你好世界" + + +class TestFromJson: + """Test the _from_json buffer decoder helper.""" + + def setup_method(self): + from cactus.bindings.cactus import _from_json + self.from_json = _from_json + + def _make_buf(self, text): + encoded = text.encode() if isinstance(text, str) else text + buf = ctypes.create_string_buffer(len(encoded) + 1) + buf.value = encoded + return buf + + def test_valid_json(self): + result = self.from_json(self._make_buf('{"key": "value"}')) + assert result == {"key": "value"} + + def test_empty_buffer_returns_empty_dict(self): + buf = ctypes.create_string_buffer(64) + result = self.from_json(buf) + assert result == {} + + def test_array_json(self): + result = self.from_json(self._make_buf('[1, 2, 3]')) + assert result == [1, 2, 3] + + def test_numeric_values(self): + result = self.from_json(self._make_buf('{"score": 0.95, "count": 42}')) + assert result["score"] == 0.95 + assert result["count"] == 42 + + def test_nested_json(self): + text = '{"response": "Hi", "metrics": {"tokens": 5, "latency_ms": 123.4}}' + result = self.from_json(self._make_buf(text)) + assert result["response"] == "Hi" + assert result["metrics"]["tokens"] == 5 + + def test_unicode_response(self): + result = self.from_json(self._make_buf('{"text": "日本語テスト"}')) + assert result["text"] == "日本語テスト" + + +class TestPreparePcm: + """Test the _prepare_pcm audio marshaling helper.""" + + def setup_method(self): + from cactus.bindings.cactus import _prepare_pcm + self.prepare_pcm = _prepare_pcm + + def test_none_returns_null(self): + ptr, size = self.prepare_pcm(None) + assert ptr is None + assert size == 0 + + def test_bytes_marshaled(self): + data = b"\x00\x01\x02\x03\xff" + ptr, size = self.prepare_pcm(data) + assert ptr is not None + assert size == 5 + + def test_empty_bytes(self): + ptr, size = self.prepare_pcm(b"") + assert size == 0 + + def test_large_buffer(self): + data = bytes(range(256)) * 100 # 25600 bytes + ptr, size = self.prepare_pcm(data) + assert size == 25600 + + +class TestEnc: + """Test the _enc string encoding helper.""" + + def setup_method(self): + from cactus.bindings.cactus import _enc + self.enc = _enc + + def test_none_passthrough(self): + assert self.enc(None) is None + + def test_string_encoded(self): + result = self.enc("hello") + assert result == b"hello" + + def test_bytes_passthrough(self): + raw = b"hello" + result = self.enc(raw) + assert result is raw + + def test_unicode_encoded(self): + result = self.enc("日本語") + assert result == "日本語".encode() + + def test_empty_string(self): + result = self.enc("") + assert result == b"" + + +# ── Model ID resolution tests ────────────────────────────────────── + + +# ── TranspileOptions tests ────────────────────────────────────────── + + +class TestTranspileOptions: + """Test the TranspileOptions dataclass.""" + + def test_defaults(self): + from cactus.cli.model import TranspileOptions + opts = TranspileOptions() + assert opts.task == "auto" + assert opts.prompt is None + assert opts.image_files is None + assert opts.audio_file is None + assert opts.max_new_tokens is None + assert opts.component_pipeline == "auto" + assert opts.components is None + assert opts.system_prompt is None + assert opts.trust_remote_code is False + assert opts.local_files_only is False + assert opts.cache_context_length is None + + def test_custom_values(self): + from cactus.cli.model import TranspileOptions + opts = TranspileOptions( + task="causal_lm_logits", + prompt="Hello", + max_new_tokens=256, + trust_remote_code=True, + cache_context_length="131072", + ) + assert opts.task == "causal_lm_logits" + assert opts.prompt == "Hello" + assert opts.max_new_tokens == 256 + assert opts.trust_remote_code is True + assert opts.cache_context_length == "131072" + + def test_frozen(self): + from cactus.cli.model import TranspileOptions + opts = TranspileOptions() + with pytest.raises(AttributeError): + opts.task = "something" + + +# ── CLI parser tests ──────────────────────────────────────────────── + + +class TestCliParser: + """Test CLI argument parsing.""" + + def setup_method(self): + from cactus.cli import create_parser + self.parser = create_parser() + + def test_download_command(self): + args = self.parser.parse_args(["download", "google/gemma-4-E2B-it"]) + assert args.command == "download" + assert args.model_id == "google/gemma-4-E2B-it" + + def test_download_defaults(self): + args = self.parser.parse_args(["download"]) + assert args.token is None + assert args.bits == 4 + + def test_run_command(self): + args = self.parser.parse_args(["run", "google/gemma-4-E2B-it", "--prompt", "hi"]) + assert args.command == "run" + assert args.model_id == "google/gemma-4-E2B-it" + assert args.prompt == "hi" + + def test_run_command_chunked_bundle_flags(self): + args = self.parser.parse_args([ + "run", "Foo/Bar", + "--audio", "audio.wav", + "--image", "image.png", + "--input-ids", "1,2,3", + "--input-ids-file", "tokens.txt", + "--max-new-tokens", "4", + "--result-json", "result.json", + ]) + assert args.command == "run" + assert args.audio == "audio.wav" + assert args.image == "image.png" + assert args.input_ids == "1,2,3" + assert args.input_ids_file == "tokens.txt" + assert args.max_new_tokens == 4 + assert args.result_json == "result.json" + + def test_convert_command(self): + args = self.parser.parse_args(["convert", "Qwen/Qwen3-0.6B", "--bits", "2"]) + assert args.command == "convert" + assert args.model_id == "Qwen/Qwen3-0.6B" + assert args.bits == 2 + + def test_build_command(self): + args = self.parser.parse_args(["build", "--apple"]) + assert args.command == "build" + assert args.apple is True + + def test_test_command_suite(self): + args = self.parser.parse_args(["test", "--suite", "llm"]) + assert args.command == "test" + assert args.suite == "llm" + assert args.component == "all" + + def test_test_command_component(self): + args = self.parser.parse_args(["test", "--component", "kernels"]) + assert args.command == "test" + assert args.component == "kernels" + assert args.suite is None + + def test_test_command_defaults(self): + args = self.parser.parse_args(["test"]) + assert args.command == "test" + assert args.component == "all" + assert args.suite is None + + def test_auth_command(self): + args = self.parser.parse_args(["auth", "--status"]) + assert args.command == "auth" + assert args.status is True + + def test_clean_command(self): + args = self.parser.parse_args(["clean"]) + assert args.command == "clean" + + def test_no_command_prints_help(self): + args = self.parser.parse_args([]) + assert args.command is None + + def test_run_rejects_bare_name(self): + import pytest + with pytest.raises(SystemExit): + self.parser.parse_args(["run", "whisper-base", "--prompt", "hi"]) + + def test_download_bits_flag(self): + args = self.parser.parse_args(["download", "Foo/Bar", "--bits", "2"]) + assert args.bits == 2 + + + +# ── Index API tests ──────────────────────────────────────────────── + + +class TestIndexApi: + """Test the vector index API (no model weights needed).""" + + def setup_method(self): + import tempfile + from cactus.bindings.cactus import ( + cactus_index_init, cactus_index_add, cactus_index_query, + cactus_index_get, cactus_index_delete, cactus_index_compact, + cactus_index_destroy, + ) + self.tmpdir = tempfile.mkdtemp() + self.init = cactus_index_init + self.add = cactus_index_add + self.query = cactus_index_query + self.get = cactus_index_get + self.delete = cactus_index_delete + self.compact = cactus_index_compact + self.destroy = cactus_index_destroy + self._handles = [] + + def _create_index(self, dim): + idx = self.init(self.tmpdir, dim) + self._handles.append(idx) + return idx + + def teardown_method(self): + import shutil + for h in self._handles: + try: + self.destroy(h) + except Exception: + pass + self._handles.clear() + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def test_init_destroy(self): + idx = self._create_index(4) + assert idx is not None + + def test_add_and_query(self): + idx = self._create_index(4) + self.add(idx, [0, 1, 2], ["doc a", "doc b", "doc c"], + embeddings=[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0]]) + result = self.query(idx, [1, 0, 0, 0], {"n_results": 2}) + ids = [r["id"] for r in result["results"]] + assert 0 in ids + + def test_get_by_id(self): + idx = self._create_index(3) + self.add(idx, [42], ["hello world"], embeddings=[[1, 0, 0]]) + result = self.get(idx, [42]) + assert result["results"][0]["document"] == "hello world" + + def test_delete(self): + idx = self._create_index(3) + self.add(idx, [0, 1], ["a", "b"], embeddings=[[1, 0, 0], [0, 1, 0]]) + self.delete(idx, [0]) + result = self.query(idx, [1, 0, 0], {"n_results": 10}) + ids = [r["id"] for r in result["results"]] + assert 0 not in ids + + def test_compact(self): + idx = self._create_index(3) + self.add(idx, [0], ["x"], embeddings=[[1, 0, 0]]) + self.compact(idx) + result = self.query(idx, [1, 0, 0], {"n_results": 1}) + assert len(result["results"]) >= 1 + + def test_persistence(self): + idx = self._create_index(3) + self.add(idx, [0], ["persistent"], embeddings=[[1, 0, 0]]) + self.destroy(idx) + self._handles.remove(idx) + idx2 = self._create_index(3) + result = self.get(idx2, [0]) + assert result["results"][0]["document"] == "persistent" + + +# ── Streaming callback tests ────────────────────────────────────── + + +class TestStreamingCallbacks: + """Test the token callback wrapper (no model weights needed).""" + + def setup_method(self): + from cactus.bindings.cactus import _make_token_callback, TokenCallback + self.make_cb = _make_token_callback + self.TokenCallback = TokenCallback + + def test_none_returns_valid_callback(self): + cb = self.make_cb(None) + assert cb is not None + + def test_false_returns_valid_callback(self): + cb = self.make_cb(False) + assert cb is not None + + def test_callable_returns_valid_callback(self): + cb = self.make_cb(lambda text, tid: None) + assert cb is not None + + def test_callback_is_correct_ctypes_type(self): + cb = self.make_cb(lambda text, tid: None) + assert isinstance(cb, self.TokenCallback) + + def test_multiple_callbacks_independent(self): + results_a = [] + results_b = [] + cb_a = self.make_cb(lambda t, i: results_a.append(t)) + cb_b = self.make_cb(lambda t, i: results_b.append(t)) + assert cb_a is not cb_b + + +# ── Error path tests ────────────────────────────────────────────── + + +class TestErrorPaths: + """Test error handling (no model weights needed).""" + + def test_init_bad_path_raises(self): + from cactus.bindings.cactus import cactus_init + with pytest.raises(RuntimeError): + cactus_init("/nonexistent/path/to/model") + + def test_init_empty_path_raises(self): + from cactus.bindings.cactus import cactus_init + with pytest.raises(RuntimeError): + cactus_init("") + + def test_get_last_error_returns_string(self): + from cactus.bindings.cactus import cactus_init, cactus_get_last_error + try: + cactus_init("/nonexistent") + except RuntimeError: + pass + err = cactus_get_last_error() + assert isinstance(err, str) + assert len(err) > 0 + + def test_init_bad_path_error_message(self): + from cactus.bindings.cactus import cactus_init, cactus_get_last_error + with pytest.raises(RuntimeError, match="config"): + cactus_init("/nonexistent/model") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/python/tests/test_capture_jax.py b/python/tests/test_capture_jax.py new file mode 100644 index 000000000..4267d8560 --- /dev/null +++ b/python/tests/test_capture_jax.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +from cactus.transpile.capture_jax import capture_jax_function +from cactus.transpile.capture_jax import capture_jax_function_with_params +from cactus.transpile.lower import transpile_ir + + +def _execute_ir(ir, *inputs: object) -> list[np.ndarray]: + graph = transpile_ir(ir) + graph.set_inputs([np.asarray(value) for value in inputs]) + return [output.numpy() for output in graph.execute()] + + +def _assert_close(actual: object, expected: object, *, atol: float = 8e-2, rtol: float = 8e-2) -> None: + np.testing.assert_allclose(np.asarray(actual, dtype=np.float32), np.asarray(expected, dtype=np.float32), atol=atol, rtol=rtol) + + +def test_capture_jax_handles_scalar_and_broadcast_boundaries() -> None: + scalar = jnp.asarray(2.0, dtype=jnp.float16) + x = jnp.asarray([[1.0, -2.0, 3.0]], dtype=jnp.float16) + + def fn(runtime_scalar, values): + constant = jnp.asarray(0.25, dtype=jnp.float16) + return values * runtime_scalar + jnp.broadcast_to(constant, values.shape) + + ir = capture_jax_function(fn, (scalar, x)) + got = _execute_ir(ir, scalar, x)[0] + + _assert_close(got, fn(scalar, x)) + + +def test_capture_jax_gqa_repeat_then_matmul_matches_jax() -> None: + q = jnp.ones((1, 4, 3, 2), dtype=jnp.float16) + k = jnp.asarray(np.arange(12, dtype=np.float16).reshape(1, 2, 3, 2) / 10.0) + + def fn(query, key): + repeated = jnp.repeat(key, 2, axis=1) + return jnp.matmul(query, jnp.swapaxes(repeated, -1, -2)) + + ir = capture_jax_function(fn, (q, k)) + got = _execute_ir(ir, q, k)[0] + + _assert_close(got, fn(q, k)) + + +@pytest.mark.parametrize("variant", ["mean_rsqrt", "sum_div_rsqrt", "pow_neg_half", "inv_sqrt", "weight_first"]) +def test_capture_jax_rms_norm_common_spellings_lower_to_rms_norm(variant: str) -> None: + x = jnp.asarray(np.linspace(-650.0, 640.0, 16, dtype=np.float16).reshape(1, 2, 8)) + weight = jnp.asarray(np.linspace(0.8, 1.2, 8, dtype=np.float16)) + + def fn(values): + squared = values.astype(jnp.float32) * values.astype(jnp.float32) + mean_square = jnp.mean(squared, axis=-1, keepdims=True) + if variant == "sum_div_rsqrt": + inv = jax.lax.rsqrt(jnp.sum(squared, axis=-1, keepdims=True) / values.shape[-1] + 1.0e-5) + return values * inv * weight + if variant == "pow_neg_half": + return values * ((mean_square + 1.0e-5) ** -0.5) * weight + if variant == "inv_sqrt": + return values * (1.0 / jnp.sqrt(mean_square + 1.0e-5)) * weight + if variant == "weight_first": + return weight * (values * jax.lax.rsqrt(mean_square + 1.0e-5)) + return values * jax.lax.rsqrt(mean_square + 1.0e-5) * weight + + ir = capture_jax_function(fn, (x,), constant_names=("weight",)) + got = _execute_ir(ir, x)[0] + + assert any(ir.nodes[node_id].op == "rms_norm" for node_id in ir.order) + _assert_close(got, fn(x)) + + +def test_capture_jax_zero_centered_rms_norm_lowers_to_rms_norm() -> None: + x = jnp.asarray(np.linspace(-73.0, 73.0, 512, dtype=np.float16).reshape(1, 1, 512)) + scale = jnp.zeros((512,), dtype=jnp.float16) + + def fn(values): + rms = jnp.sqrt(jnp.mean(values.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + 1.0e-6) + return ((1.0 + scale) * values / rms).astype(jnp.float16) + + ir = capture_jax_function(fn, (x,), constant_names=("scale",)) + got = _execute_ir(ir, x)[0] + + assert any(ir.nodes[node_id].op == "rms_norm" for node_id in ir.order) + _assert_close(got, fn(x)) + + +@pytest.mark.parametrize("variant", ["mean_rsqrt", "sum_div_rsqrt", "sqrt_div", "no_bias"]) +def test_capture_jax_layer_norm_common_spellings_lower_to_layer_norm(variant: str) -> None: + x = jnp.asarray(np.linspace(-650.0, 640.0, 16, dtype=np.float16).reshape(1, 2, 8)) + weight = jnp.asarray(np.linspace(0.8, 1.2, 8, dtype=np.float16)) + bias = jnp.asarray(np.linspace(-0.2, 0.2, 8, dtype=np.float16)) + + def fn(values): + centered = values - jnp.mean(values.astype(jnp.float32), axis=-1, keepdims=True) + if variant == "sum_div_rsqrt": + var = jnp.sum(centered.astype(jnp.float32) * centered.astype(jnp.float32), axis=-1, keepdims=True) / values.shape[-1] + normed = centered * jax.lax.rsqrt(var + 1.0e-5) + elif variant == "sqrt_div": + var = jnp.mean(centered.astype(jnp.float32) * centered.astype(jnp.float32), axis=-1, keepdims=True) + normed = centered / jnp.sqrt(var + 1.0e-5) + else: + var = jnp.mean(centered.astype(jnp.float32) * centered.astype(jnp.float32), axis=-1, keepdims=True) + normed = centered * jax.lax.rsqrt(var + 1.0e-5) + if variant == "no_bias": + return normed * weight + return normed * weight + bias + + ir = capture_jax_function(fn, (x,), constant_names=("weight", "bias")) + got = _execute_ir(ir, x)[0] + + assert any(ir.nodes[node_id].op == "layer_norm" for node_id in ir.order) + _assert_close(got, fn(x)) + + +def test_capture_jax_prenorm_residual_add_lowers_to_add_clipped() -> None: + x = jnp.asarray(np.linspace(-0.5, 0.5, 16, dtype=np.float16).reshape(1, 2, 8)) + weight = jnp.asarray(np.linspace(0.8, 1.2, 8, dtype=np.float16)) + + def fn(values): + rms = jnp.sqrt(jnp.mean(values.astype(jnp.float32) ** 2, axis=-1, keepdims=True) + 1.0e-6) + branch = (values / rms) * weight + return values + branch + + ir = capture_jax_function(fn, (x,), constant_names=("weight",)) + got = _execute_ir(ir, x)[0] + + assert any(ir.nodes[node_id].op == "add_clipped" for node_id in ir.order) + _assert_close(got, fn(x)) + + +def test_capture_jax_prefill_attention_pattern_lowers_to_attention() -> None: + q = jnp.asarray(np.linspace(-0.5, 0.5, 1 * 2 * 3 * 4, dtype=np.float16).reshape(1, 2, 3, 4)) + k = jnp.asarray(np.linspace(0.25, -0.25, 1 * 2 * 3 * 4, dtype=np.float16).reshape(1, 2, 3, 4)) + v = jnp.asarray(np.linspace(-0.75, 0.75, 1 * 2 * 3 * 4, dtype=np.float16).reshape(1, 2, 3, 4)) + mask = jnp.asarray(np.where(np.tril(np.ones((1, 2, 3, 3), dtype=np.bool_)), 0.0, -1.0e4), dtype=jnp.float16) + + def fn(query, key, value, additive_mask): + scores = (query @ key.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(4.0, dtype=jnp.float32)) + probs = jax.nn.softmax(scores + additive_mask, axis=-1) + return (probs @ value).transpose(0, 2, 1, 3).reshape(1, 3, 8) + + ir = capture_jax_function(fn, (q, k, v, mask)) + got = _execute_ir(ir, q, k, v, mask)[0] + + assert any(ir.nodes[node_id].op == "attention" for node_id in ir.order) + _assert_close(got, fn(q, k, v, mask)) + + +def test_capture_jax_attention_keeps_broadcast_mask_decomposed() -> None: + q = jnp.asarray(np.linspace(-0.5, 0.5, 2 * 4 * 5 * 8, dtype=np.float16).reshape(2, 4, 5, 8)) + k = jnp.asarray(np.linspace(0.25, -0.25, 2 * 4 * 5 * 8, dtype=np.float16).reshape(2, 4, 5, 8)) + v = jnp.asarray(np.linspace(-0.75, 0.75, 2 * 4 * 5 * 8, dtype=np.float16).reshape(2, 4, 5, 8)) + mask = jnp.asarray(np.where(np.tril(np.ones((1, 1, 5, 5), dtype=np.bool_)), 0.0, -1.0e4), dtype=jnp.float16) + + def fn(query, key, value, additive_mask): + scores = (query @ key.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(8.0, dtype=jnp.float32)) + probs = jax.nn.softmax(scores + additive_mask, axis=-1) + return (probs @ value).transpose(0, 2, 1, 3).reshape(2, 5, 32) + + ir = capture_jax_function(fn, (q, k, v, mask)) + got = _execute_ir(ir, q, k, v, mask)[0] + + assert not any(ir.nodes[node_id].op == "attention" for node_id in ir.order) + _assert_close(got, fn(q, k, v, mask)) + + +def test_capture_flax_module_smoke_when_available() -> None: + pytest.importorskip("flax") + from flax import linen as nn + + class TinyFlaxModule(nn.Module): + @nn.compact + def __call__(self, token_ids, features): + embedded = nn.Embed(8, 4)(token_ids) + projected = nn.Dense(4)(features) + hidden = nn.LayerNorm()(embedded + projected) + return nn.Dense(3)(jax.nn.silu(hidden)) + + model = TinyFlaxModule() + token_ids = jnp.asarray([[1, 2]], dtype=jnp.int32) + features = jnp.asarray([[[0.2, -0.1], [0.4, 0.3]]], dtype=jnp.float16) + params = model.init(jax.random.PRNGKey(0), token_ids, features)["params"] + + def fn(model_params, ids, feats): + return model.apply({"params": model_params}, ids, feats) + + ir = capture_jax_function_with_params(fn, params, (token_ids, features)) + got = _execute_ir(ir, token_ids, features)[0] + + assert any(ir.nodes[node_id].op == "layer_norm" for node_id in ir.order) + _assert_close(got, fn(params, token_ids, features), atol=1.2e-1, rtol=1.2e-1) diff --git a/python/tests/test_capture_jax_user_graph_bundle.py b/python/tests/test_capture_jax_user_graph_bundle.py new file mode 100644 index 000000000..ac272cf74 --- /dev/null +++ b/python/tests/test_capture_jax_user_graph_bundle.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +from cactus.transpile.capture_jax import JaxGraphSpec +from cactus.transpile.jax_user_graph_bundle import build_jax_generation_graph_bundle +from cactus.transpile.jax_user_graph_bundle import build_jax_user_graph_bundle +from cactus.transpile.jax_user_graph_bundle import load_jax_user_graph_bundle + + +def _assert_close(actual: object, expected: object, *, atol: float = 8e-2, rtol: float = 8e-2) -> None: + np.testing.assert_allclose(np.asarray(actual, dtype=np.float32), np.asarray(expected, dtype=np.float32), atol=atol, rtol=rtol) + + +def _toy_export(): + params = { + "encoder_w": jnp.asarray([[0.2, -0.1, 0.4], [0.3, 0.5, -0.2]], dtype=jnp.float16), + "decoder_w": jnp.asarray([[0.1, -0.3], [0.4, 0.2], [-0.2, 0.6]], dtype=jnp.float16), + "out_b": jnp.asarray([0.01, -0.02], dtype=jnp.float16), + } + source = jnp.asarray([[[1.0, -0.5], [0.25, 0.75]]], dtype=jnp.float16) + target = jnp.asarray([[[0.4, -0.1, 0.2]]], dtype=jnp.float16) + + def encoder_fn(model_params, source_features): + encoded = source_features @ model_params["encoder_w"] + return encoded * jax.nn.sigmoid(encoded) + + def decoder_fn(model_params, target_features, encoder_out): + context = jnp.mean(encoder_out, axis=1, keepdims=True) + return (target_features + context) @ model_params["decoder_w"] + model_params["out_b"] + + encoder_out = encoder_fn(params, source) + specs = ( + JaxGraphSpec( + name="encoder", role="encoder", fn=encoder_fn, example_args=(source,), + input_names=("source_features",), output_names=("encoder_out",), + ), + JaxGraphSpec( + name="decoder", role="generic", fn=decoder_fn, example_args=(target, encoder_out), + input_names=("target_features", "encoder_out"), output_names=("logits",), + ), + ) + return params, specs, source, target, encoder_fn, decoder_fn + + +def _build_toy_bundle(tmp_path: Path, **kwargs): + params, specs, source, target, encoder_fn, decoder_fn = _toy_export() + result = build_jax_user_graph_bundle( + params=params, + specs=specs, + output_dir=tmp_path / "bundle", + model_id="toy-jax", + **kwargs, + ) + return result, params, source, target, encoder_fn, decoder_fn + + +def test_jax_user_graph_bundle_writes_manifest_graphs_and_weights(tmp_path: Path) -> None: + result, params, source, target, encoder_fn, decoder_fn = _build_toy_bundle( + tmp_path, + task="generic", + inputs_metadata={"owner": "client"}, + ) + + encoder_out = result.bundle.execute("encoder", source)[0].numpy() + logits = result.bundle.execute("decoder", target, encoder_out)[0].numpy() + manifest = json.loads(result.components_manifest_path.read_text()) + + assert manifest["model_source"] == "jax_user_graph" + assert manifest["component_order"] == ["encoder", "decoder"] + assert (tmp_path / "bundle/components/encoder/graph.cactus").exists() + assert (tmp_path / "bundle/components/decoder/graph.cactus").exists() + for component in manifest["components"]: + raw_ir_path, optimized_ir_path = (tmp_path / "bundle" / component[key] for key in ("raw_ir", "optimized_ir")) + assert raw_ir_path.exists() and optimized_ir_path.exists() + assert json.loads(raw_ir_path.read_text())["graph"]["meta"]["frontend"] == "jax" + assert json.loads(optimized_ir_path.read_text())["graph"]["outputs"] == component["outputs"] + assert (tmp_path / "bundle/weights_manifest.json").exists() + _assert_close(encoder_out, encoder_fn(params, source)) + _assert_close(logits, decoder_fn(params, target, encoder_fn(params, source))) + + +def test_jax_user_graph_bundle_loads_saved_graphs_and_mmap_weights(tmp_path: Path) -> None: + result, params, source, target, encoder_fn, decoder_fn = _build_toy_bundle(tmp_path) + loaded = load_jax_user_graph_bundle(result.output_dir) + + encoder_out = loaded.execute("encoder", source)[0] + logits = loaded.execute("decoder", target, encoder_out)[0].numpy() + + assert set(loaded.graphs) == {"encoder", "decoder"} + assert loaded.graphs["encoder"].logical_inputs == ["source_features"] + assert loaded.graphs["decoder"].logical_outputs == ["logits"] + _assert_close(encoder_out.numpy(), encoder_fn(params, source)) + _assert_close(logits, decoder_fn(params, target, encoder_fn(params, source))) + loaded.reset() + + +def test_jax_user_graph_bundle_supports_external_weights_dir(tmp_path: Path) -> None: + params = { + "w": jnp.asarray([[0.5], [-0.25]], dtype=jnp.float16), + "b": jnp.asarray([0.1], dtype=jnp.float16), + } + x = jnp.asarray([[2.0, -1.0]], dtype=jnp.float16) + + def fn(model_params, values): + return values @ model_params["w"] + model_params["b"] + + result = build_jax_user_graph_bundle( + params=params, + specs=(JaxGraphSpec(name="project", fn=fn, example_args=(x,), input_names=("x",), output_names=("y",)),), + output_dir=tmp_path / "bundle", + weights_dir=tmp_path / "weights", + model_id="external-weights", + ) + loaded = load_jax_user_graph_bundle(result.components_manifest_path) + manifest = json.loads(result.components_manifest_path.read_text()) + + assert result.weights_dir == tmp_path / "weights" + assert (tmp_path / "weights/weights_manifest.json").exists() + assert manifest["weights_dir"] == str(tmp_path / "weights") + _assert_close(loaded.execute("project", np.asarray(x, dtype=np.float32))[0].numpy(), fn(params, x)) + + +def test_jax_generation_decoder_step_fuses_attention_without_internal_cache(tmp_path: Path) -> None: + params = { + "wq": jnp.eye(16, dtype=jnp.float16), + "wk": jnp.eye(16, dtype=jnp.float16), + "wv": jnp.eye(16, dtype=jnp.float16), + } + x = jnp.asarray(np.linspace(-0.75, 0.75, 16, dtype=np.float16).reshape(1, 1, 16)) + + def decoder_step(model_params, values): + q = (values @ model_params["wq"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) + k = (values @ model_params["wk"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) + v = (values @ model_params["wv"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) + scores = (q @ k.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(8.0, dtype=jnp.float32)) + probs = jax.nn.softmax(scores, axis=-1).astype(jnp.float16) + return (probs @ v).transpose(0, 2, 1, 3).reshape(1, 1, 16) + + result = build_jax_generation_graph_bundle( + params=params, + decoder_step=JaxGraphSpec(name="decoder_step", fn=decoder_step, example_args=(x,), output_names=("hidden",)), + output_dir=tmp_path / "bundle", + model_id="decoder-step-attention", + ) + graph = result.bundle.graphs["decoder_step"].ir_graph + attention_nodes = [node for node in graph.nodes.values() if node.op == "attention"] + + assert "use_internal_kv_cache" not in graph.meta + assert len(attention_nodes) == 1 + assert result.bundle.graphs["decoder_step"].graph.cache_state_tensors == [] + _assert_close(result.bundle.execute("decoder_step", x)[0].numpy(), decoder_step(params, x)) + + +def test_jax_generation_decoder_step_fuses_cross_attention(tmp_path: Path) -> None: + params = { + "wq": jnp.eye(16, dtype=jnp.float16), + } + x = jnp.asarray(np.linspace(-0.75, 0.75, 16, dtype=np.float16).reshape(1, 1, 16)) + key = jnp.asarray(np.linspace(-0.5, 0.5, num=1 * 2 * 4 * 8).reshape(1, 2, 4, 8), dtype=jnp.float16) + value = jnp.asarray(np.linspace(0.75, -0.25, num=1 * 2 * 4 * 8).reshape(1, 2, 4, 8), dtype=jnp.float16) + mask = jnp.asarray([[[[True, True, True, False]]]], dtype=jnp.bool_) + + def decoder_step(model_params, values, cross_k, cross_v, cross_mask): + q = (values @ model_params["wq"]).reshape(1, 1, 2, 8).transpose(0, 2, 1, 3) + scores = (q @ cross_k.transpose(0, 1, 3, 2)) / jnp.sqrt(jnp.asarray(8.0, dtype=jnp.float32)) + scores = jnp.where(cross_mask, scores, jnp.finfo(scores.dtype).min) + probs = jax.nn.softmax(scores, axis=-1).astype(jnp.float16) + return (probs @ cross_v).transpose(0, 2, 1, 3).reshape(1, 1, 16) + + result = build_jax_generation_graph_bundle( + params=params, + decoder_step=JaxGraphSpec( + name="decoder_step", fn=decoder_step, example_args=(x, key, value, mask), output_names=("hidden",), + ), + output_dir=tmp_path / "bundle", + model_id="cross-attention-step", + ) + graph = result.bundle.graphs["decoder_step"].ir_graph + attention_nodes = [node for node in graph.nodes.values() if node.op == "attention"] + + assert len(attention_nodes) == 1 + assert attention_nodes[0].meta["rewritten_from"] == "jax_decoder_step_cross_attention" + assert len(attention_nodes[0].inputs) == 4 + assert result.bundle.graphs["decoder_step"].graph.cache_state_tensors == [] + _assert_close( + result.bundle.execute("decoder_step", x, key, value, mask)[0].numpy(), + decoder_step(params, x, key, value, mask), + ) diff --git a/python/tests/test_chat_template_golden.py b/python/tests/test_chat_template_golden.py new file mode 100644 index 000000000..78077dd2c --- /dev/null +++ b/python/tests/test_chat_template_golden.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +transformers = pytest.importorskip("transformers") + +from .bundles import PROJECT_ROOT, WEIGHTS, _iter_bundle_candidates, _read_model_type, _valid_bundle + +MONKEY_IMAGE = PROJECT_ROOT / "cactus-engine" / "tests" / "assets" / "test_monkey.png" + +FAMILIES = [ + ("gemma-4-e2b-it", "google/gemma-4-E2B-it"), + ("qwen3-0.6b", "Qwen/Qwen3-0.6B"), + ("lfm2-350m", "LiquidAI/LFM2-350M"), + ("functiongemma-270m-it", "google/functiongemma-270m-it"), +] +FAMILY_IDS = ["gemma4", "qwen3", "lfm2", "functiongemma"] +THINKING_TYPES = {"gemma4", "qwen"} + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name."}, + }, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the current time in a timezone.", + "parameters": { + "type": "object", + "properties": { + "timezone": {"type": "string", "description": "IANA timezone name."}, + }, + "required": ["timezone"], + }, + }, + }, +] + +SYSTEM_CHAT = [ + {"role": "system", "content": "You are a concise assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "And of Germany?"}, +] + +TOOLS_USER = [ + {"role": "user", "content": "What is the weather in Paris?"}, +] + +TOOL_CALL_RESPONSE = [ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "call_1", "type": "function", + "function": {"name": "get_weather", "arguments": {"city": "Paris"}}}, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "name": "get_weather", "content": "Sunny, 22C"}, +] + +PARALLEL_TOOL_CALLS = [ + {"role": "user", "content": "Weather in Paris and time in UTC?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "call_1", "type": "function", + "function": {"name": "get_weather", "arguments": {"city": "Paris"}}}, + {"id": "call_2", "type": "function", + "function": {"name": "get_time", "arguments": {"timezone": "UTC"}}}, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "name": "get_weather", "content": "Sunny, 22C"}, + {"role": "tool", "tool_call_id": "call_2", "name": "get_time", "content": "12:00"}, +] + +THINKING_CHAT = [ + {"role": "user", "content": "Briefly explain entropy."}, +] + +CASES = [ + ("system_chat", SYSTEM_CHAT, None, False), + ("tools_user", TOOLS_USER, TOOLS, False), + ("tool_call_response", TOOL_CALL_RESPONSE, TOOLS, False), + ("parallel_tool_calls", PARALLEL_TOOL_CALLS, TOOLS, False), + ("thinking_enabled", THINKING_CHAT, None, True), + ("thinking_disabled", THINKING_CHAT, None, False), +] + + +def _pythonic_call(call: dict) -> str: + fn = call["function"] + args = ", ".join(f"{key}={json.dumps(value)}" for key, value in fn["arguments"].items()) + return f"{fn['name']}({args})" + + +def _embed_lfm2_tool_calls(message: dict) -> dict: + calls = ", ".join(_pythonic_call(call) for call in message["tool_calls"]) + embedded = dict(message) + embedded["content"] = message["content"] + f"<|tool_call_start|>[{calls}]<|tool_call_end|>" + del embedded["tool_calls"] + return embedded + + +def _to_hf_messages(messages: list[dict], model_type: str) -> list[dict]: + out = [] + for message in messages: + if message.get("images"): + message = dict(message) + message["content"] = [{"type": "image"} for _ in message.pop("images")] + [ + {"type": "text", "text": message["content"]} + ] + if model_type == "lfm2" and message.get("tool_calls"): + message = _embed_lfm2_tool_calls(message) + out.append(message) + return out + + +_LOADED: dict[str, tuple] = {} + + +@pytest.fixture(scope="session") +def load_family(): + from cactus import cactus_destroy, cactus_init + + def _load(bundle_name: str, hf_id: str): + if bundle_name not in _LOADED: + bundle = next( + (c for c in _iter_bundle_candidates(bundle_name) if _valid_bundle(c)), + WEIGHTS / bundle_name, + ) + if not _valid_bundle(bundle): + pytest.skip(f"bundle not prepared: {bundle}") + try: + tokenizer = transformers.AutoTokenizer.from_pretrained(hf_id, local_files_only=True) + except OSError as exc: + pytest.skip(f"HF tokenizer not cached for {hf_id}: {exc}") + _LOADED[bundle_name] = (cactus_init(str(bundle)), tokenizer, _read_model_type(bundle)) + return _LOADED[bundle_name] + + yield _load + while _LOADED: + model, _, _ = _LOADED.popitem()[1] + cactus_destroy(model) + + +def _assert_render_matches(model, tokenizer, engine_render: str, hf_render: str, label: str): + from cactus import cactus_tokenize + + assert engine_render == hf_render, ( + f"{label}: engine chat-template render diverges from the HF reference template" + ) + engine_ids = cactus_tokenize(model, engine_render) + hf_ids = tokenizer(engine_render, add_special_tokens=False).input_ids + assert engine_ids == hf_ids, ( + f"{label}: engine tokenization of the render diverges from the HF tokenizer" + ) + + +@pytest.mark.parametrize(("case_name", "messages", "tools", "thinking"), CASES, + ids=[case[0] for case in CASES]) +@pytest.mark.parametrize(("bundle_name", "hf_id"), FAMILIES, ids=FAMILY_IDS) +def test_render_matches_hf_template(load_family, bundle_name, hf_id, case_name, messages, + tools, thinking): + from cactus import cactus_render_prompt + + model, tokenizer, model_type = load_family(bundle_name, hf_id) + if case_name.startswith("thinking") and not any(t in model_type for t in THINKING_TYPES): + pytest.skip(f"{model_type} has no thinking template support") + engine_render = cactus_render_prompt( + model, messages, options={"enable_thinking_if_supported": thinking}, tools=tools + ) + hf_render = tokenizer.apply_chat_template( + _to_hf_messages(messages, model_type), + tools=tools, + add_generation_prompt=True, + tokenize=False, + enable_thinking=thinking, + ) + _assert_render_matches(model, tokenizer, engine_render, hf_render, + f"{bundle_name}/{case_name}") + + +def _monkey_soft_tokens(hf_id: str) -> int: + from huggingface_hub.constants import HF_HUB_CACHE + from PIL import Image + from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor + + repo_dir = Path(HF_HUB_CACHE) / f"models--{hf_id.replace('/', '--')}" + configs = sorted(repo_dir.glob("snapshots/*/processor_config.json")) + if not configs: + pytest.skip(f"no cached processor_config.json for {hf_id}") + image_processor = Gemma4ImageProcessor( + **json.loads(configs[0].read_text(encoding="utf-8"))["image_processor"] + ) + with Image.open(MONKEY_IMAGE) as image: + return image_processor(image)["num_soft_tokens_per_image"][0] + + +def test_gemma4_image_render_matches_hf_template(load_family): + from cactus import cactus_render_prompt + + bundle_name, hf_id = FAMILIES[0] + model, tokenizer, model_type = load_family(bundle_name, hf_id) + messages = [ + {"role": "user", "content": "What is in this image?", "images": [str(MONKEY_IMAGE)]}, + ] + engine_render = cactus_render_prompt( + model, messages, options={"enable_thinking_if_supported": False} + ) + hf_render = tokenizer.apply_chat_template( + _to_hf_messages(messages, model_type), + add_generation_prompt=True, + tokenize=False, + enable_thinking=False, + ) + soft_tokens = _monkey_soft_tokens(hf_id) + expanded = hf_render.replace( + tokenizer.image_token, + tokenizer.boi_token + tokenizer.image_token * soft_tokens + tokenizer.eoi_token, + ) + _assert_render_matches(model, tokenizer, engine_render, expanded, f"{bundle_name}/image") diff --git a/python/tests/test_cli_run.py b/python/tests/test_cli_run.py new file mode 100644 index 000000000..15b976b40 --- /dev/null +++ b/python/tests/test_cli_run.py @@ -0,0 +1,69 @@ +from argparse import Namespace +from pathlib import Path +from types import SimpleNamespace + +import cactus.cli.common as common_mod +import cactus.cli.model as model_mod +import cactus.cli.run as run_mod + + +def test_cmd_run_forwards_chunked_bundle_flags(monkeypatch, tmp_path: Path) -> None: + bundle_dir = tmp_path / "bundle" + (bundle_dir / "components").mkdir(parents=True) + (bundle_dir / "components" / "manifest.json").write_text("{}", encoding="utf-8") + + fake_bin = tmp_path / "bin" + fake_run = fake_bin / "run" + fake_bin.mkdir(parents=True) + fake_run.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8") + monkeypatch.setattr(common_mod, "BIN_DIR", fake_bin) + + image_file = tmp_path / "image.png" + audio_file = tmp_path / "audio.wav" + result_json = tmp_path / "result.json" + input_ids_file = tmp_path / "tokens.txt" + image_file.write_bytes(b"image") + audio_file.write_bytes(b"audio") + input_ids_file.write_text("1 2 3", encoding="utf-8") + + calls = [] + + def fake_subprocess_run(cmd): + calls.append(cmd) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(common_mod.subprocess, "run", fake_subprocess_run) + monkeypatch.setattr(model_mod, "ensure_runnable_bundle", lambda *a, **k: bundle_dir) + + args = Namespace( + no_cloud_tele=False, + model_id="org/model", + bits=4, + token=None, + reconvert=False, + system=None, + prompt="hi", + image=str(image_file), + audio=str(audio_file), + input_ids="1,2,3", + input_ids_file=str(input_ids_file), + max_new_tokens=4, + result_json=str(result_json), + confidence_threshold=None, + cloud_timeout_ms=None, + backend="auto", + thinking=False, + no_cloud_handoff=False, + ) + + assert run_mod.cmd_run(args) == 0 + assert len(calls) == 1 + cmd = calls[0] + assert cmd[:2] == [str(fake_run), str(bundle_dir)] + assert cmd[cmd.index("--prompt") + 1] == "hi" + assert cmd[cmd.index("--image") + 1] == str(image_file) + assert cmd[cmd.index("--audio") + 1] == str(audio_file) + assert cmd[cmd.index("--input-ids") + 1] == "1,2,3" + assert cmd[cmd.index("--input-ids-file") + 1] == str(input_ids_file) + assert cmd[cmd.index("--max-new-tokens") + 1] == "4" + assert cmd[cmd.index("--result-json") + 1] == str(result_json) diff --git a/python/tests/test_cli_transpile_defaults.py b/python/tests/test_cli_transpile_defaults.py new file mode 100644 index 000000000..ee1802364 --- /dev/null +++ b/python/tests/test_cli_transpile_defaults.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest +from cactus import cli +from cactus.cli import convert as convert_cli +from cactus.transpile.model_adapters import _cache_context_length + + +class _ConfigWithText: + def get_text_config(self): + return SimpleNamespace(max_position_embeddings=128000) + + +def test_cache_context_length_uses_explicit_value() -> None: + model = SimpleNamespace(config=SimpleNamespace(max_position_embeddings=40960)) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length="32768", + fallback_extra_tokens=512, + ) == 32768 + + +def test_cache_context_length_reads_top_level_config() -> None: + model = SimpleNamespace(config=SimpleNamespace(max_position_embeddings=40960)) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length=None, + fallback_extra_tokens=512, + ) == 40960 + + +def test_cache_context_length_reads_text_config() -> None: + model = SimpleNamespace(config=_ConfigWithText()) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length="auto", + fallback_extra_tokens=512, + ) == 128000 + + +def test_cache_context_length_falls_back_to_capture_size() -> None: + model = SimpleNamespace(config=SimpleNamespace()) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length=None, + fallback_extra_tokens=512, + ) == 2560 + + +def test_cache_context_length_rejects_non_positive_explicit_value() -> None: + with pytest.raises(ValueError): + _cache_context_length( + SimpleNamespace(config=SimpleNamespace()), + input_seq_len=2048, + cache_context_length="0", + fallback_extra_tokens=512, + ) + + +def test_cmd_convert_builds_graph_after_weights(monkeypatch, tmp_path: Path) -> None: + """`cactus convert` quantizes weights and then builds the runtime graph in + one step; the graph build is pointed at the weights just written.""" + parser = cli.create_parser() + out = tmp_path / "out" + args = parser.parse_args(["convert", "google/gemma-4-E2B-it", str(out), "--reconvert"]) + + weight_calls: list[dict] = [] + + def _fake_ensure_weights(model_id, **kw): + weight_calls.append(kw) + Path(kw["output_dir"]).mkdir(parents=True, exist_ok=True) + return kw["output_dir"] + + transpile_calls: list[dict] = [] + + def _fake_run_transpile(model_id, **kw): + transpile_calls.append(kw) + return 0 + + import cactus.cli.model as model_mod + import cactus.cli.transpile as transpile_mod + monkeypatch.setattr(model_mod, "ensure_weights", _fake_ensure_weights) + monkeypatch.setattr(model_mod, "_default_multimodal_assets", lambda: ([], None)) + monkeypatch.setattr(transpile_mod, "run_transpile", _fake_run_transpile) + + rc = convert_cli.cmd_convert(args) + + assert rc == 0 + assert len(weight_calls) == 1 + assert len(transpile_calls) == 1 + extra_args = transpile_calls[0]["extra_args"] + assert "--weights-dir" in extra_args + assert str(out) in extra_args + + +def test_cmd_convert_weights_only_skips_graph(monkeypatch, tmp_path: Path) -> None: + """`--weights-only` stops after CQ quantization, before the runtime graph.""" + parser = cli.create_parser() + out = tmp_path / "out" + args = parser.parse_args(["convert", "google/gemma-4-E2B-it", str(out), "--weights-only"]) + + def _fake_ensure_weights(model_id, **kw): + Path(kw["output_dir"]).mkdir(parents=True, exist_ok=True) + return kw["output_dir"] + + def _fail_if_called(*a, **kw): + raise AssertionError("--weights-only must skip the runtime graph build") + + import cactus.cli.model as model_mod + import cactus.cli.transpile as transpile_mod + monkeypatch.setattr(model_mod, "ensure_weights", _fake_ensure_weights) + monkeypatch.setattr(transpile_mod, "run_transpile", _fail_if_called) + + rc = convert_cli.cmd_convert(args) + assert rc == 0 + + +def test_cli_convert_absorbs_graph_flags() -> None: + """The former `transpile` options now live on `cactus convert`; the separate + `transpile` subcommand no longer exists.""" + parser = cli.create_parser() + args = parser.parse_args(["convert", "google/gemma-4-E2B-it", + "--weights-dir", "/tmp/x", "--task", "causal_lm_logits", + "--cache-context-length", "4096", "--component-pipeline", "on"]) + assert args.command == "convert" + assert args.model_id == "google/gemma-4-E2B-it" + assert args.weights_dir == "/tmp/x" + assert args.task == "causal_lm_logits" + assert args.cache_context_length == "4096" + assert args.component_pipeline == "on" + + with pytest.raises(SystemExit): + parser.parse_args(["transpile", "google/gemma-4-E2B-it"]) + + +def test_cli_registers_low_memory_transpile_flag() -> None: + parser = cli.create_parser() + args = parser.parse_args([ + "convert", + "Qwen/Qwen3-0.6B", + "--weights-dir", + "/tmp/x", + "--low-memory-load", + ]) + + assert args.low_memory_load is True + + +def test_benchmark_aliases_engine_benchmark_suite(monkeypatch) -> None: + """`cactus benchmark` is `cactus test --component engine --suite benchmark`.""" + from cactus.cli import test as test_cli + + parser = cli.create_parser() + args = parser.parse_args(["benchmark", "--backend", "cpu", "--bits", "2"]) + assert args.command == "benchmark" + assert args.backend == "cpu" + assert args.bits == 2 + + calls = [] + monkeypatch.setattr(test_cli, "cmd_test", lambda a: calls.append(a) or 0) + assert test_cli.cmd_benchmark(args) == 0 + assert calls[0].component == "engine" + assert calls[0].suite == "benchmark" + + +def test_run_accepts_local_bundle_path() -> None: + """`cactus run` accepts a HF id (org/model) OR a local path. Bare names + like 'whisper-base' (no slash) are rejected.""" + parser = cli.create_parser() + args = parser.parse_args(["run", "/tmp/bundle", + "--prompt", "hi", "--system", "be brief", "--thinking"]) + assert args.command == "run" + assert args.model_id == "/tmp/bundle" + assert args.prompt == "hi" + assert args.system == "be brief" + assert args.thinking is True diff --git a/python/tests/test_cloud_handoff_live.py b/python/tests/test_cloud_handoff_live.py new file mode 100644 index 000000000..3ca32f4f5 --- /dev/null +++ b/python/tests/test_cloud_handoff_live.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import json +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +from .bundles import _find_bundle + +PREFERRED_BUNDLES = ["qwen3-0.6b", "lfm2-350m"] +OPTIONS = {"max_tokens": 8, "temperature": 0, "confidence_threshold": 1.0} +MESSAGES = [{"role": "user", "content": "What is 2+2?"}] + + +class _AuthFailHandler(BaseHTTPRequestHandler): + def do_POST(self): + with self.server.hits_lock: + self.server.hits += 1 + self.rfile.read(int(self.headers.get("Content-Length", 0))) + body = json.dumps({"error": "invalid or expired api key"}).encode() + self.send_response(401) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + +def _hits(server) -> int: + with server.hits_lock: + return server.hits + + +def test_auth_failure_disables_cloud_handoff_for_session(monkeypatch, tmp_path): + bundle = _find_bundle(PREFERRED_BUNDLES, {"qwen", "lfm2"}, on_missing=pytest.skip) + server = ThreadingHTTPServer(("127.0.0.1", 0), _AuthFailHandler) + server.hits = 0 + server.hits_lock = threading.Lock() + threading.Thread(target=server.serve_forever, daemon=True).start() + + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("CACTUS_CLOUD_API_BASE", f"http://127.0.0.1:{server.server_address[1]}/api/v1") + monkeypatch.setenv("CACTUS_CLOUD_KEY", "stale-test-key") + monkeypatch.setenv("CACTUS_NO_CLOUD_TELE", "1") + monkeypatch.delenv("CACTUS_DISABLE_CLOUD_HANDOFF", raising=False) + + from cactus.bindings.cactus import cactus_complete, cactus_destroy, cactus_init + + model = cactus_init(str(bundle)) + try: + with pytest.raises(RuntimeError): + cactus_complete(model, MESSAGES, options=OPTIONS) + assert _hits(server) == 1 + + start = time.monotonic() + try: + result = cactus_complete(model, MESSAGES, options=OPTIONS) + except RuntimeError: + result = None + elapsed = time.monotonic() - start + + assert _hits(server) == 1 + assert result is not None + assert result["success"] is True + assert elapsed < 5.0 + finally: + cactus_destroy(model) + server.shutdown() + server.server_close() diff --git a/python/tests/test_component_partition.py b/python/tests/test_component_partition.py new file mode 100644 index 000000000..a87ca03bd --- /dev/null +++ b/python/tests/test_component_partition.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from cactus.transpile.component_partition import COMPONENT_LM_ENCODER +from cactus.transpile.component_partition import COMPONENT_VISION_ENCODER +from cactus.transpile.component_partition import classify_node_component +from cactus.transpile.component_partition import classify_value_component +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue + + +def test_component_partition_ignores_torch_layer_names() -> None: + value = IRValue( + id="weight", + meta={ + "source_name": "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight", + "torch_name": "model.layers.0.mlp.gate_proj.weight", + }, + ) + + assert classify_value_component( + "weight", + value, + family="", + task="multimodal_causal_lm_logits", + ) is None + + +def test_component_partition_uses_semantic_input_names() -> None: + assert ( + classify_value_component( + "pixel_values", + IRValue(id="pixel_values"), + family="", + task="multimodal_causal_lm_logits", + ) + == COMPONENT_VISION_ENCODER + ) + + +def test_component_partition_uses_ops_for_multimodal_merge() -> None: + node = IRNode( + id="embedding_0", + op="embedding", + inputs=["input_ids", "weight"], + outputs=["inputs_embeds"], + ) + + assert ( + classify_node_component( + node, + family="", + task="multimodal_causal_lm_logits", + ) + == COMPONENT_LM_ENCODER + ) diff --git a/python/tests/test_component_plan.py b/python/tests/test_component_plan.py new file mode 100644 index 000000000..d09cd0e10 --- /dev/null +++ b/python/tests/test_component_plan.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from cactus.transpile.component_plan import infer_component_plan_from_config + + +def test_gemma4_vision_only_config_does_not_request_audio_component() -> None: + config = { + "model_type": "gemma4", + "architectures": ["Gemma4ForConditionalGeneration"], + "vision_config": {"hidden_size": 1152}, + "text_config": {"num_hidden_layers": 30}, + } + + plan = infer_component_plan_from_config(config, model_id="google/gemma-4-26B-A4B") + + assert plan is not None + assert plan.task == "multimodal_causal_lm_logits" + assert plan.components == ("vision_encoder", "lm_encoder", "decoder") + assert plan.needs_image + assert not plan.needs_audio + + +def test_gemma4_image_audio_config_keeps_audio_component() -> None: + config = { + "model_type": "gemma4", + "architectures": ["Gemma4ForConditionalGeneration"], + "vision_config": {"hidden_size": 1152}, + "audio_config": {"hidden_size": 1024}, + "text_config": {"num_hidden_layers": 30}, + } + + plan = infer_component_plan_from_config(config, model_id="google/gemma-4-E2B-it") + + assert plan is not None + assert plan.components == ("vision_encoder", "audio_encoder", "lm_encoder", "decoder") + assert plan.needs_audio + + +def test_needle_config_uses_component_pipeline_plan() -> None: + config = { + "model_type": "needle", + "architectures": ["NeedleForCausalLM"], + "num_encoder_layers": 12, + "num_decoder_layers": 8, + } + + plan = infer_component_plan_from_config(config, model_id="Cactus-Compute/needle") + + assert plan is not None + assert plan.task == "causal_lm_logits" + assert plan.components == ("source_encoder", "decoder_cross_kv", "decoder_step") + assert plan.force_component_pipeline diff --git a/python/tests/test_download.py b/python/tests/test_download.py new file mode 100644 index 000000000..23c90df64 --- /dev/null +++ b/python/tests/test_download.py @@ -0,0 +1,241 @@ +import shutil +import sys +import tarfile +import tempfile +import unittest +import zipfile +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from cactus.cli import utils as cq_utils +from cactus.cli.utils import ( + archives_from_repo_files, + parse_cq_variant, + parse_version, + promote_single_root, + resolve_archive, + resolve_weight_revision, + safe_extract_archive, + suggested_cq_repo, + validate_extracted_bundle, + verify_archive_sha256, +) + + +class TestCqDownloadResolver(unittest.TestCase): + def test_suggested_cq_repo_strips_org(self): + self.assertEqual( + suggested_cq_repo("LiquidAI/LFM2.5-350M"), + "Cactus-Compute/LFM2.5-350M", + ) + self.assertEqual( + suggested_cq_repo("google/gemma-4-E2B-it"), + "Cactus-Compute/gemma-4-E2B-it", + ) + + def test_suggested_cq_repo_strips_existing_cq_suffix(self): + self.assertEqual( + suggested_cq_repo("Cactus-Compute/model-cq4"), + "Cactus-Compute/model", + ) + self.assertEqual( + suggested_cq_repo("Cactus-Compute/model-cq"), + "Cactus-Compute/model", + ) + + def test_parse_cq_variant(self): + self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq1.zip"), 1) + self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq4.zip"), 4) + self.assertEqual(parse_cq_variant("gemma-4-E2B-it-CQ3.tar.gz"), 3) + self.assertIsNone(parse_cq_variant("README.md")) + self.assertIsNone(parse_cq_variant("L4V4A4.zip")) + + def test_archives_from_repo_files(self): + files = [ + "gemma-4-E2B-it-cq1.zip", + "gemma-4-E2B-it-cq2.zip", + "gemma-4-E2B-it-cq3.zip", + "gemma-4-E2B-it-cq4.zip", + "README.md", + "config.json", + ] + sizes = {"gemma-4-E2B-it-cq4.zip": 500_000_000} + archives = archives_from_repo_files(files, sizes=sizes) + self.assertEqual(len(archives), 4) + self.assertEqual(archives[0].bits, 1) + self.assertEqual(archives[3].bits, 4) + self.assertEqual(archives[3].size, 500_000_000) + + def test_resolve_archive_finds_match(self): + files = ["model-cq1.zip", "model-cq2.zip", "model-cq3.zip", "model-cq4.zip"] + archives = archives_from_repo_files(files) + resolution = resolve_archive("Cactus-Compute/model", "model", archives, 4) + self.assertEqual(resolution.archive.filename, "model-cq4.zip") + self.assertEqual(resolution.archive.bits, 4) + + def test_resolve_archive_missing_variant_errors(self): + files = ["model-cq1.zip", "model-cq2.zip"] + archives = archives_from_repo_files(files) + with self.assertRaisesRegex(RuntimeError, "cq4 not found"): + resolve_archive("Cactus-Compute/model", "model", archives, 4) + + def test_resolve_archive_empty_errors(self): + with self.assertRaisesRegex(RuntimeError, "No CQ archives"): + resolve_archive("Cactus-Compute/model", "model", [], 4) + + +class TestWeightRevisionResolver(unittest.TestCase): + def test_parse_version(self): + self.assertEqual(parse_version("v2.0"), (2, 0, 0)) + self.assertEqual(parse_version("2.0.0"), (2, 0, 0)) + self.assertEqual(parse_version("v1.14"), (1, 14, 0)) + self.assertEqual(parse_version("V1.2.3"), (1, 2, 3)) + self.assertIsNone(parse_version("main")) + self.assertIsNone(parse_version("v2.0-rc1")) + self.assertIsNone(parse_version(None)) + self.assertIsNone(parse_version("")) + + def _patch_tags(self, tags): + names = [(t, parse_version(t)) for t in tags if parse_version(t) is not None] + original = cq_utils.list_version_tags + cq_utils.list_version_tags = lambda repo_id, token=None: tuple(names) + self.addCleanup(setattr, cq_utils, "list_version_tags", original) + + def test_picks_latest_tag_at_or_below_runtime(self): + self._patch_tags(["v1.12", "v1.13", "v1.14", "v2.0"]) + self.assertEqual(resolve_weight_revision("Cactus-Compute/m", "2.0.0"), "v2.0") + + def test_skips_tags_above_runtime(self): + self._patch_tags(["v1.12", "v2.0", "v3.0"]) + self.assertEqual(resolve_weight_revision("Cactus-Compute/m", "2.0.0"), "v2.0") + + def test_minor_version_ordering(self): + self._patch_tags(["v1.8", "v1.9", "v1.10", "v1.14"]) + self.assertEqual(resolve_weight_revision("Cactus-Compute/m", "2.0.0"), "v1.14") + + def test_no_eligible_tag_falls_back_to_main(self): + self._patch_tags(["v3.0", "v4.1"]) + self.assertIsNone(resolve_weight_revision("Cactus-Compute/m", "2.0.0")) + + def test_no_version_tags_falls_back_to_main(self): + self._patch_tags(["main", "experimental"]) + self.assertIsNone(resolve_weight_revision("Cactus-Compute/m", "2.0.0")) + + def test_unparseable_runtime_falls_back_to_main(self): + self._patch_tags(["v1.0", "v2.0"]) + self.assertIsNone(resolve_weight_revision("Cactus-Compute/m", "garbage")) + + +class TestCqSafeExtraction(unittest.TestCase): + def setUp(self): + self.root = Path(tempfile.mkdtemp(prefix="cactus_cq_test_")) + + def tearDown(self): + shutil.rmtree(self.root, ignore_errors=True) + + def _write_minimal_package(self, base: Path): + base.mkdir(parents=True, exist_ok=True) + (base / "config.txt").write_text("model_type=test\n", encoding="utf-8") + (base / "token_embeddings.weights").write_bytes(b"x") + (base / "vocab.txt").write_text("0\t\n", encoding="utf-8") + (base / "tokenizer_config.txt").write_text("tokenizer_type=bpe\nvocab_format=id_tab_token\n", encoding="utf-8") + (base / "special_tokens.json").write_text("{}", encoding="utf-8") + (base / "tokenizer.json").write_text("{}", encoding="utf-8") + (base / "merges.txt").write_text("", encoding="utf-8") + components = base / "components" + components.mkdir(parents=True, exist_ok=True) + (components / "manifest.json").write_text("{}", encoding="utf-8") + + def test_zip_extract_and_promote_single_root(self): + package = self.root / "pkg" + self._write_minimal_package(package / "model-cq4") + archive = self.root / "model-cq4.zip" + src = package / "model-cq4" + with zipfile.ZipFile(archive, "w") as zf: + for path in src.rglob("*"): + if path.is_file(): + zf.write(path, f"model-cq4/{path.relative_to(src)}") + + out = self.root / "out" + out.mkdir() + safe_extract_archive(archive, out) + promote_single_root(out) + validate_extracted_bundle(out) + self.assertTrue((out / "config.txt").exists()) + + def test_promote_single_root_rejects_ambiguous_layout(self): + out = self.root / "out" + (out / "a").mkdir(parents=True) + (out / "b").mkdir(parents=True) + with self.assertRaisesRegex(RuntimeError, "archive root or under one top-level directory"): + promote_single_root(out) + + def test_validate_extracted_bundle_requires_tokenizer_sidecars(self): + package = self.root / "pkg" + self._write_minimal_package(package) + (package / "tokenizer.json").unlink() + with self.assertRaisesRegex(RuntimeError, "missing tokenizer sidecar"): + validate_extracted_bundle(package) + + def test_validate_extracted_bundle_accepts_sentencepiece_without_tokenizer_json(self): + package = self.root / "pkg" + self._write_minimal_package(package) + (package / "tokenizer.json").unlink() + (package / "tokenizer_config.txt").write_text( + "tokenizer_type=sentencepiece\nsp_model_type=bpe\n", encoding="utf-8" + ) + validate_extracted_bundle(package) + + def test_validate_extracted_bundle_accepts_seq2seq_embeddings(self): + package = self.root / "pkg" + self._write_minimal_package(package) + (package / "token_embeddings.weights").rename(package / "decoder_token_embeddings.weights") + validate_extracted_bundle(package) + + def test_validate_extracted_bundle_accepts_audio_bundle_without_tokenizer_config(self): + package = self.root / "pkg" + self._write_minimal_package(package) + for sidecar in ("tokenizer_config.txt", "special_tokens.json", "tokenizer.json", "merges.txt"): + (package / sidecar).unlink() + validate_extracted_bundle(package) + + def test_validate_extracted_bundle_requires_weights(self): + package = self.root / "pkg" + self._write_minimal_package(package) + (package / "token_embeddings.weights").unlink() + with self.assertRaisesRegex(RuntimeError, "no weight files"): + validate_extracted_bundle(package) + + def test_verify_archive_sha256_rejects_mismatch(self): + archive = self.root / "model-cq4.zip" + archive.write_bytes(b"archive") + with self.assertRaisesRegex(RuntimeError, "checksum mismatch"): + verify_archive_sha256(archive, "0" * 64) + + def test_zip_rejects_traversal(self): + archive = self.root / "bad.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("../escape.txt", "bad") + out = self.root / "out" + out.mkdir() + with self.assertRaisesRegex(RuntimeError, "Unsafe path"): + safe_extract_archive(archive, out) + + def test_tar_rejects_symlink(self): + archive = self.root / "bad.tar" + target = self.root / "target.txt" + target.write_text("x", encoding="utf-8") + with tarfile.open(archive, "w") as tf: + info = tarfile.TarInfo("link") + info.type = tarfile.SYMTYPE + info.linkname = str(target) + tf.addfile(info) + out = self.root / "out" + out.mkdir() + with self.assertRaisesRegex(RuntimeError, "Refusing link"): + safe_extract_archive(archive, out) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_encoder_cross_kv_route.py b/python/tests/test_encoder_cross_kv_route.py new file mode 100644 index 000000000..bba3627e9 --- /dev/null +++ b/python/tests/test_encoder_cross_kv_route.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from cactus.transpile.model_adapters import build_component_module_specs + + +class ToyNeedle(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace( + model_type="needle", + decoder_start_token_id=1, + pad_token_id=0, + ) + + def cactus_source_encode( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden = input_ids.to(dtype=torch.float32).unsqueeze(-1).expand(-1, -1, 2).contiguous() + return hidden, attention_mask[:, None, None, :].to(dtype=torch.bool) + + def cactus_decoder_cross_kv( + self, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + del encoder_attention_mask + batch, seq, _ = encoder_hidden_states.shape + return tuple(torch.zeros((batch, seq, 1, 2), dtype=torch.float32) for _ in range(4)) + + def cactus_decoder_step( + self, + decoder_input_ids: torch.Tensor, + position_ids: torch.Tensor, + encoder_attention_mask: torch.Tensor, + *cross_kv: torch.Tensor, + ) -> torch.Tensor: + del position_ids, encoder_attention_mask, cross_kv + return torch.zeros((decoder_input_ids.shape[0], decoder_input_ids.shape[1], 8), dtype=torch.float32) + + +def test_needle_encoder_cross_kv_route_specs_are_metadata_driven() -> None: + specs = build_component_module_specs( + ToyNeedle(), + task="causal_lm_logits", + named_tensors={"input_ids": torch.tensor([[1, 2, 0]], dtype=torch.long)}, + ) + + assert specs is not None + assert [spec.component for spec in specs] == ["source_encoder", "decoder_cross_kv", "decoder_step"] + + source, cross_kv, decoder_step = specs + assert source.input_keys == ("input_ids", "attention_mask") + assert source.output_keys == ("encoder_hidden_states", "encoder_attention_mask") + assert source.metadata["runtime_route"] == "encoder_cross_kv_decoder_step" + assert source.metadata["runtime_role"] == "source_encoder" + assert source.metadata["source_kind"] == "text_tokens" + + assert cross_kv.input_keys == ("encoder_hidden_states", "encoder_attention_mask") + assert cross_kv.output_keys == ("cross_k_0", "cross_v_0", "cross_k_1", "cross_v_1") + assert cross_kv.metadata["runtime_role"] == "decoder_cross_kv" + + assert decoder_step.input_keys == ( + "decoder_input_ids", + "position_ids", + "encoder_attention_mask", + "cross_k_0", + "cross_v_0", + "cross_k_1", + "cross_v_1", + ) + assert decoder_step.output_keys == ("logits",) + assert decoder_step.metadata["runtime_role"] == "decoder_step" + assert decoder_step.graph_meta["use_internal_kv_cache"] is True diff --git a/python/tests/test_graph.py b/python/tests/test_graph.py new file mode 100644 index 000000000..59d41e414 --- /dev/null +++ b/python/tests/test_graph.py @@ -0,0 +1,745 @@ +import unittest + +import numpy as np +import sys +import tempfile +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) +from cactus import Graph, Tensor + + +class TestGraphElementwise(unittest.TestCase): + + def test_pow_abs(self): + g = Graph() + a = g.input((4,)) + y = a.pow(2.0).abs() + + g.set_input(a, np.array([-2.0, -1.0, 2.0, 3.0], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([4.0, 1.0, 4.0, 9.0], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_subtract(self): + g = Graph() + a = g.input((4,)) + b = g.input((4,)) + y = a - b + + g.set_input(a, np.array([5, 6, 7, 8], dtype=np.float16)) + g.set_input(b, np.array([1, 2, 3, 4], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([4, 4, 4, 4], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_add(self): + g = Graph() + a = g.input((4,)) + b = g.input((4,)) + y = a + b + + g.set_input(a, np.array([1, 2, 3, 4], dtype=np.float16)) + g.set_input(b, np.array([10, 20, 30, 40], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([11, 22, 33, 44], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_multiply(self): + g = Graph() + a = g.input((4,)) + b = g.input((4,)) + y = a * b + + g.set_input(a, np.array([2, 3, 4, 5], dtype=np.float16)) + g.set_input(b, np.array([3, 4, 5, 6], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([6, 12, 20, 30], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_divide(self): + g = Graph() + a = g.input((4,)) + b = g.input((4,)) + y = a / b + + g.set_input(a, np.array([10, 20, 30, 40], dtype=np.float16)) + g.set_input(b, np.array([2, 4, 5, 8], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([5, 5, 6, 5], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + +class TestGraphComposed(unittest.TestCase): + + def test_composed_inference_graph(self): + g = Graph() + + a = g.input((2, 2)) + b = g.input((2, 2)) + + diff = a - b + summ = a + b + mixed = (diff * summ) / b + feat = mixed.abs().pow(2.0).view((4,)) + out_node = g.cat([feat, feat], axis=0) + + a_data = np.array([[2.0, 4.0], [6.0, 8.0]], dtype=np.float16) + b_data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float16) + g.set_input(a, a_data) + g.set_input(b, b_data) + g.execute() + + out = out_node.numpy() + base = ((a_data - b_data) * (a_data + b_data)) / b_data + expected = np.concatenate( + [np.abs(base).reshape(4) ** 2, np.abs(base).reshape(4) ** 2] + ).astype(np.float16) + + self.assertEqual(out.shape, (8,)) + np.testing.assert_allclose(out, expected, atol=1e-2) + + +class TestGraphTensorOps(unittest.TestCase): + + def test_topk_from_fp16_outputs_fp32_indices_after_roundtrip(self): + g = Graph() + a = g.input((1, 5)) + topk = g.topk(a, 3) + indices = g.index(topk, 0, axis=0) + + data = np.array([[0.1, 4.0, 2.0, 5.0, 3.0]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + self.assertEqual(g.output_info(topk)["precision"], Graph.FP32) + self.assertEqual(g.output_info(indices)["precision"], Graph.FP32) + np.testing.assert_allclose(indices.numpy(), np.array([[3.0, 1.0, 4.0]], dtype=np.float32)) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "topk_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_a = Tensor(loaded, a.id, a.shape, a.dtype) + loaded_indices = Tensor(loaded, indices.id, indices.shape, indices.dtype) + loaded.set_input(loaded_a, data) + loaded.execute() + + self.assertEqual(loaded.output_info(loaded_indices)["precision"], Graph.FP32) + np.testing.assert_allclose(loaded_indices.numpy(), np.array([[3.0, 1.0, 4.0]], dtype=np.float32)) + + def test_view(self): + g = Graph() + a = g.input((2, 3)) + y = a.view((6,)) + + g.set_input(a, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([1, 2, 3, 4, 5, 6], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_flatten(self): + g = Graph() + a = g.input((2, 3)) + y = a.flatten() + + g.set_input(a, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([1, 2, 3, 4, 5, 6], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_reshape(self): + g = Graph() + a = g.input((2, 3)) + y = a.reshape((3, 2)) + + g.set_input(a, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_transpose(self): + g = Graph() + a = g.input((2, 3)) + y = a.transpose() + + data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data.T.astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_permute(self): + g = Graph() + a = g.input((2, 3, 4)) + y = a.permute((1, 0, 2)) + + data = np.arange(24, dtype=np.float16).reshape(2, 3, 4) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = np.transpose(data, (1, 0, 2)).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_concat(self): + g = Graph() + a = g.input((2, 2)) + b = g.input((2, 2)) + y = a.concat(b, axis=0) + + g.set_input(a, np.array([[1, 2], [3, 4]], dtype=np.float16)) + g.set_input(b, np.array([[5, 6], [7, 8]], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_cat_1d(self): + g = Graph() + a = g.input((3,)) + b = g.input((2,)) + y = g.cat([a, b], axis=0) + + g.set_input(a, np.array([1, 2, 3], dtype=np.float16)) + g.set_input(b, np.array([4, 5], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([1, 2, 3, 4, 5], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_cat_2d_axis1(self): + g = Graph() + a = g.input((2, 2)) + b = g.input((2, 3)) + y = g.cat([a, b], axis=1) + + g.set_input(a, np.array([[1, 2], [3, 4]], dtype=np.float16)) + g.set_input(b, np.array([[5, 6, 7], [8, 9, 10]], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([[1, 2, 5, 6, 7], [3, 4, 8, 9, 10]], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_matmul(self): + g = Graph() + a = g.input((2, 3)) + b = g.input((3, 2)) + y = a.matmul(b) + + a_data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16) + b_data = np.array([[7, 8], [9, 10], [11, 12]], dtype=np.float16) + g.set_input(a, a_data) + g.set_input(b, b_data) + g.execute() + + out = y.numpy() + expected = (a_data.astype(np.float32) @ b_data.astype(np.float32)).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_gather(self): + g = Graph() + embeddings = g.input((4, 2)) + indices = g.input((3,), dtype=Graph.FP32) + y = g.gather(embeddings, indices) + + embedding_data = np.array( + [[10, 11], [20, 21], [30, 31], [40, 41]], + dtype=np.float16, + ) + index_data = np.array([2, 0, 3], dtype=np.float32) + g.set_input(embeddings, embedding_data) + g.set_input(indices, index_data, dtype=Graph.FP32) + g.execute() + + out = y.numpy() + expected = embedding_data[[2, 0, 3]] + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_slice(self): + g = Graph() + a = g.input((4, 3)) + y = a.slice(axis=0, start=1, length=2) + + data = np.arange(12, dtype=np.float16).reshape(4, 3) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data[1:3] + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_index(self): + g = Graph() + a = g.input((4, 3)) + y = a.index(2, axis=0) + + data = np.arange(12, dtype=np.float16).reshape(4, 3) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data[2] + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_mean(self): + g = Graph() + a = g.input((2, 3)) + y = a.mean(axis=1) + + data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data.astype(np.float32).mean(axis=1).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_variance(self): + g = Graph() + a = g.input((2, 3)) + y = a.variance(axis=1) + + data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data.astype(np.float32).var(axis=1).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_min(self): + g = Graph() + a = g.input((2, 3)) + y = a.min(axis=1) + + data = np.array([[3, 1, 2], [6, 4, 5]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data.min(axis=1).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_max(self): + g = Graph() + a = g.input((2, 3)) + y = a.max(axis=1) + + data = np.array([[3, 1, 2], [6, 4, 5]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = data.max(axis=1).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + +class TestGraphNorms(unittest.TestCase): + + def test_layernorm_matches_alias(self): + g = Graph() + x = g.input((2, 3)) + weight = g.input((3,)) + bias = g.input((3,)) + y_primary = g.layernorm(x, weight, bias, eps=1e-5) + y_alias = g.layer_norm(x, weight, bias, eps=1e-5) + + x_data = np.array([[1.0, 3.0, 5.0], [2.0, 4.0, 8.0]], dtype=np.float16) + weight_data = np.array([1.0, 1.5, 0.5], dtype=np.float16) + bias_data = np.array([0.25, -0.5, 1.0], dtype=np.float16) + g.set_input(x, x_data) + g.set_input(weight, weight_data) + g.set_input(bias, bias_data) + g.execute() + + np.testing.assert_allclose(y_primary.numpy(), y_alias.numpy(), atol=1e-2) + + def test_groupnorm_matches_alias(self): + g = Graph() + x = g.input((1, 4)) + weight = g.input((4,)) + bias = g.input((4,)) + y_primary = g.groupnorm(x, weight, bias, num_groups=2, eps=1e-5) + y_alias = g.group_norm(x, weight, bias, num_groups=2, eps=1e-5) + + x_data = np.array([[1.0, 3.0, 2.0, 4.0]], dtype=np.float16) + weight_data = np.array([1.0, 1.0, 0.5, 0.5], dtype=np.float16) + bias_data = np.array([0.0, 0.5, 1.0, -1.0], dtype=np.float16) + g.set_input(x, x_data) + g.set_input(weight, weight_data) + g.set_input(bias, bias_data) + g.execute() + + np.testing.assert_allclose(y_primary.numpy(), y_alias.numpy(), atol=1e-2) + + def test_batchnorm_matches_alias(self): + g = Graph() + x = g.input((1, 3)) + weight = g.input((3,)) + bias = g.input((3,)) + running_mean = g.input((3,)) + running_var = g.input((3,)) + y_primary = g.batchnorm(x, weight, bias, running_mean, running_var, axis=1, eps=1e-5) + y_alias = g.batch_norm(x, weight, bias, running_mean, running_var, axis=1, eps=1e-5) + + x_data = np.array([[1.0, 2.0, 3.0]], dtype=np.float16) + weight_data = np.array([1.0, 1.5, 2.0], dtype=np.float16) + bias_data = np.array([0.0, -1.0, 0.5], dtype=np.float16) + mean_data = np.array([0.5, 1.0, 1.5], dtype=np.float16) + var_data = np.array([1.0, 4.0, 9.0], dtype=np.float16) + g.set_input(x, x_data) + g.set_input(weight, weight_data) + g.set_input(bias, bias_data) + g.set_input(running_mean, mean_data) + g.set_input(running_var, var_data) + g.execute() + + np.testing.assert_allclose(y_primary.numpy(), y_alias.numpy(), atol=1e-2) + + +class TestGraphActivations(unittest.TestCase): + + def test_relu(self): + g = Graph() + a = g.input((4,)) + y = a.relu() + + g.set_input(a, np.array([-2, -1, 0, 3], dtype=np.float16)) + g.execute() + + out = y.numpy() + expected = np.array([0, 0, 0, 3], dtype=np.float16) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_sigmoid(self): + g = Graph() + a = g.input((3,)) + y = a.sigmoid() + + data = np.array([0, 2, -2], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = (1.0 / (1.0 + np.exp(-data.astype(np.float32)))).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=5e-2) + + def test_tanh(self): + g = Graph() + a = g.input((3,)) + y = a.tanh() + + data = np.array([0, 1, -1], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = np.tanh(data.astype(np.float32)).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=5e-2) + + +class TestGraphSoftmax(unittest.TestCase): + + def test_softmax_2d(self): + g = Graph() + a = g.input((2, 4)) + y = a.softmax() + + data = np.array([[1, 2, 3, 4], [2, 4, 6, 8]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + # softmax over last axis + f = data.astype(np.float32) + e = np.exp(f - f.max(axis=-1, keepdims=True)) + expected = (e / e.sum(axis=-1, keepdims=True)).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=5e-2) + + +class TestGraphCumsum(unittest.TestCase): + + def test_cumsum_last_axis(self): + g = Graph() + a = g.input((2, 4)) + y = a.cumsum(axis=-1) + + data = np.array([[1, 2, 3, 4], [0.5, 1.5, 2.5, 3.5]], dtype=np.float16) + g.set_input(a, data) + g.execute() + + out = y.numpy() + expected = np.cumsum(data.astype(np.float32), axis=-1).astype(np.float16) + np.testing.assert_allclose(out, expected, atol=5e-2) + + def test_save_load_roundtrip_cumsum(self): + g = Graph() + a = g.input((3, 2)) + y = a.cumsum(axis=0) + + data = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float16) + g.set_input(a, data) + g.execute() + expected = y.numpy() + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "cumsum_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_a = Tensor(loaded, a.id, a.shape, a.dtype) + loaded_y = Tensor(loaded, y.id, y.shape, y.dtype) + + loaded.set_input(loaded_a, data) + loaded.execute() + + out = loaded_y.numpy() + np.testing.assert_allclose(out, expected, atol=1e-2) + + +class TestGraphSaveLoad(unittest.TestCase): + + def _rebind_tensor(self, graph, tensor): + return Tensor(graph, tensor.id, tensor.shape, tensor.dtype) + + def test_save_load_roundtrip_composed_graph(self): + g = Graph() + + a = g.input((2, 2)) + b = g.input((2, 2)) + diff = a - b + summ = a + b + mixed = (diff * summ) / b + feat = mixed.abs().pow(2.0).view((4,)) + out_node = g.cat([feat, feat], axis=0) + + a_data = np.array([[2.0, 4.0], [6.0, 8.0]], dtype=np.float16) + b_data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float16) + + g.set_input(a, a_data) + g.set_input(b, b_data) + g.execute() + expected = out_node.numpy() + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "roundtrip_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_a = self._rebind_tensor(loaded, a) + loaded_b = self._rebind_tensor(loaded, b) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_a, a_data) + loaded.set_input(loaded_b, b_data) + loaded.execute() + + out = loaded_out.numpy() + self.assertEqual(out.shape, expected.shape) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_save_load_roundtrip_softmax_cat(self): + g = Graph() + + a = g.input((2, 3)) + b = g.input((2, 3)) + joined = g.cat([a, b], axis=1) + out_node = joined.softmax(axis=1) + + a_data = np.array([[1.0, 2.0, 3.0], [0.5, 1.5, 2.5]], dtype=np.float16) + b_data = np.array([[4.0, 5.0, 6.0], [3.5, 4.5, 5.5]], dtype=np.float16) + + g.set_input(a, a_data) + g.set_input(b, b_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "softmax_cat_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_a = self._rebind_tensor(loaded, a) + loaded_b = self._rebind_tensor(loaded, b) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_a, a_data) + loaded.set_input(loaded_b, b_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=5e-2) + + def test_save_load_roundtrip_matmul(self): + g = Graph() + + a = g.input((2, 3)) + b = g.input((3, 2)) + out_node = a.matmul(b) + + a_data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float16) + b_data = np.array([[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]], dtype=np.float16) + + g.set_input(a, a_data) + g.set_input(b, b_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "matmul_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_a = self._rebind_tensor(loaded, a) + loaded_b = self._rebind_tensor(loaded, b) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_a, a_data) + loaded.set_input(loaded_b, b_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_save_load_roundtrip_permute(self): + g = Graph() + + a = g.input((2, 1, 2)) + out_node = a.permute((1, 0, 2)) + + data = np.array([[[58.0, 64.0]], [[139.0, 154.0]]], dtype=np.float16) + + g.set_input(a, data) + g.execute() + expected = out_node.numpy() + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "permute_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_a = self._rebind_tensor(loaded, a) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_a, data) + loaded.execute() + + out = loaded_out.numpy() + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_save_load_roundtrip_gather(self): + g = Graph() + + embeddings = g.input((4, 2)) + indices = g.input((3,), dtype=Graph.FP32) + out_node = g.gather(embeddings, indices) + + embedding_data = np.array( + [[10, 11], [20, 21], [30, 31], [40, 41]], + dtype=np.float16, + ) + index_data = np.array([2, 0, 3], dtype=np.float32) + + g.set_input(embeddings, embedding_data) + g.set_input(indices, index_data, dtype=Graph.FP32) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "gather_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_embeddings = self._rebind_tensor(loaded, embeddings) + loaded_indices = self._rebind_tensor(loaded, indices) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_embeddings, embedding_data) + loaded.set_input(loaded_indices, index_data, dtype=Graph.FP32) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-2) + + def test_save_load_roundtrip_layernorm(self): + g = Graph() + + x = g.input((2, 3)) + weight = g.input((3,)) + bias = g.input((3,)) + out_node = g.layernorm(x, weight, bias, eps=1e-5) + + x_data = np.array([[1.0, 3.0, 5.0], [2.0, 4.0, 8.0]], dtype=np.float16) + weight_data = np.array([1.0, 1.5, 0.5], dtype=np.float16) + bias_data = np.array([0.25, -0.5, 1.0], dtype=np.float16) + + g.set_input(x, x_data) + g.set_input(weight, weight_data) + g.set_input(bias, bias_data) + g.execute() + expected = out_node.numpy() + expected_info = g.output_info(out_node) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "layernorm_graph.cg" + g.save(path) + + loaded = Graph.load(path) + loaded_x = self._rebind_tensor(loaded, x) + loaded_weight = self._rebind_tensor(loaded, weight) + loaded_bias = self._rebind_tensor(loaded, bias) + loaded_out = self._rebind_tensor(loaded, out_node) + + loaded.set_input(loaded_x, x_data) + loaded.set_input(loaded_weight, weight_data) + loaded.set_input(loaded_bias, bias_data) + loaded.execute() + + out = loaded_out.numpy() + info = loaded.output_info(loaded_out) + + self.assertEqual(info["shape"], expected_info["shape"]) + self.assertEqual(info["precision"], expected_info["precision"]) + np.testing.assert_allclose(out, expected, atol=1e-2) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_model.py b/python/tests/test_model.py new file mode 100644 index 000000000..543243ae7 --- /dev/null +++ b/python/tests/test_model.py @@ -0,0 +1,160 @@ +import unittest +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +PROJECT_ROOT = Path(__file__).parent.parent.parent + +from cactus import ( + cactus_init, + cactus_destroy, + cactus_complete, + cactus_embed, + cactus_image_embed, + cactus_audio_embed, + cactus_transcribe, +) +from cactus.cli.model import ensure_bundle + + +_ASSETS_DIR = PROJECT_ROOT / "cactus-engine" / "tests" / "assets" +_TEST_IMAGE = _ASSETS_DIR / "test_monkey.png" +_TEST_AUDIO = _ASSETS_DIR / "test.wav" + +class TestVLMModel(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.weights_dir = ensure_bundle("LiquidAI/LFM2-VL-450M") + cls.model = cactus_init(str(cls.weights_dir), None, False) + + @classmethod + def tearDownClass(cls): + cactus_destroy(cls.model) + + def test_text_completion(self): + messages = [{"role": "user", "content": "What is 2+2?"}] + result = cactus_complete(self.model, messages, None, None, None) + print(f"\n completion: {json.dumps(result, indent=2)}") + self.assertIsInstance(result, dict) + self.assertTrue(result.get("success", False)) + self.assertGreater(len(result.get("response", "")), 0) + + def test_image_embedding(self): + embedding = cactus_image_embed(self.model, str(_TEST_IMAGE)) + self.assertIsInstance(embedding, list) + self.assertGreater(len(embedding), 0) + + def test_vlm_image_completion(self): + messages = [{ + "role": "user", + "content": "Describe this image", + "images": [str(_TEST_IMAGE)], + }] + result = cactus_complete(self.model, messages, None, None, None) + print(f"\n vlm completion: {json.dumps(result, indent=2)}") + self.assertIsInstance(result, dict) + self.assertTrue(result.get("success", False)) + self.assertGreater(len(result.get("response", "")), 0) + + +class TestWhisperModel(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.weights_dir = ensure_bundle("openai/whisper-small") + cls.model = cactus_init(str(cls.weights_dir), None, False) + + @classmethod + def tearDownClass(cls): + cactus_destroy(cls.model) + + def test_transcription(self): + prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" + result = cactus_transcribe( + self.model, + str(_TEST_AUDIO), + prompt, + None, + None, + None, + ) + print(f"\n transcription: {json.dumps(result, indent=2)}") + self.assertIsInstance(result, dict) + self.assertTrue(result.get("success", False)) + self.assertGreater(len(result.get("response", "")), 0) + + def test_audio_embedding(self): + embedding = cactus_audio_embed(self.model, str(_TEST_AUDIO)) + self.assertIsInstance(embedding, list) + self.assertGreater(len(embedding), 0) + + +def _cosine(a, b): + import math + dot = sum(x * y for x, y in zip(a, b)) + na = math.sqrt(sum(x * x for x in a)) + nb = math.sqrt(sum(y * y for y in b)) + return dot / (na * nb + 1e-12) + + +class TestNomicEmbedding(unittest.TestCase): + MODEL_ID = "nomic-ai/nomic-embed-text-v2-moe" + + @classmethod + def setUpClass(cls): + cls.weights_dir = ensure_bundle(cls.MODEL_ID) + cls.model = cactus_init(str(cls.weights_dir), None, False) + + @classmethod + def tearDownClass(cls): + cactus_destroy(cls.model) + + def test_text_embedding_shape_and_determinism(self): + v1 = cactus_embed(self.model, "Paris is the capital of France.", True) + v2 = cactus_embed(self.model, "Paris is the capital of France.", True) + self.assertIsInstance(v1, list) + self.assertGreater(len(v1), 0) + self.assertEqual(len(v1), len(v2)) + self.assertLess(max(abs(a - b) for a, b in zip(v1, v2)), 1e-5) + + def test_retrieval_discrimination(self): + query = cactus_embed(self.model, "search_query: What is the capital of France?", True) + relevant = cactus_embed(self.model, "search_document: Paris is the capital of France.", True) + unrelated = cactus_embed(self.model, "search_document: The Great Barrier Reef is off Australia.", True) + self.assertGreater(_cosine(query, relevant), _cosine(query, unrelated)) + + def test_hf_parity(self): + import torch + import torch.nn.functional as F + from transformers import AutoModel, AutoTokenizer + + texts = [ + "search_query: What is the capital of France?", + "search_document: Paris is the capital and largest city of France.", + "search_document: The Great Barrier Reef is located off Australia.", + ] + cactus_vecs = [cactus_embed(self.model, t, True) for t in texts] + + tok = AutoTokenizer.from_pretrained(self.MODEL_ID) + hf = AutoModel.from_pretrained(self.MODEL_ID, trust_remote_code=True, torch_dtype=torch.float32).eval() + + def hf_embed(text): + enc = tok(text, return_tensors="pt", truncation=True, max_length=256) + with torch.no_grad(): + hs = hf(**enc).last_hidden_state + mask = enc["attention_mask"][..., None].float() + pooled = (hs * mask).sum(1) / mask.sum(1).clamp(min=1e-9) + return F.normalize(pooled, p=2, dim=1)[0].tolist() + + for text, cactus_vec in zip(texts, cactus_vecs): + cos = _cosine(cactus_vec, hf_embed(text)) + print(f"\n HF parity cos ({text[:30]}...): {cos:.4f}") + # CQ4-quantized weights vs FP reference; FP parity is ~1.0 (see plan notes). + self.assertGreater(cos, 0.85) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_nomic_text_embedding.py b/python/tests/test_nomic_text_embedding.py new file mode 100644 index 000000000..716c89286 --- /dev/null +++ b/python/tests/test_nomic_text_embedding.py @@ -0,0 +1,37 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from cactus.transpile.component_plan import infer_component_plan_from_config + + +def test_nomic_config_routes_to_text_embedding(): + config = {"model_type": "nomic_bert", "architectures": ["NomicBertModel"]} + plan = infer_component_plan_from_config(config, model_id="nomic-ai/nomic-embed-text-v2-moe") + assert plan is not None + assert plan.task == "text_embedding" + assert plan.components == ("text_embedding",) + assert plan.force_component_pipeline + + +def test_nomic_model_id_routes_to_text_embedding(): + plan = infer_component_plan_from_config({}, model_id="nomic-ai/nomic-embed-text-v2-moe") + assert plan is not None + assert plan.task == "text_embedding" + + +def test_family_key_detects_nomic(): + import torch + + from cactus.transpile.model_adapters import _family_key + + class _Cfg: + model_type = "nomic_bert" + + class _StubNomic(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = _Cfg() + + assert _family_key(_StubNomic()) == "nomic" diff --git a/python/tests/test_optimize_gemma4_attention.py b/python/tests/test_optimize_gemma4_attention.py new file mode 100644 index 000000000..36fddc19d --- /dev/null +++ b/python/tests/test_optimize_gemma4_attention.py @@ -0,0 +1,328 @@ +import torch + +from cactus.transpile.canonicalize.utils import rebuild_graph +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.optimize_graph import normalize_gemma4_decoder_attention_semantics +from cactus.transpile.optimize_graph import precompute_rope_tables + + +def test_normalize_gemma4_full_attention_uses_sequence_window_compat() -> None: + graph = IRGraph( + values={ + "query": IRValue(id="query", shape=(1, 8, 800, 512), dtype="fp16"), + "key": IRValue(id="key", shape=(1, 1, 800, 512), dtype="fp16"), + "value": IRValue(id="value", shape=(1, 1, 800, 512), dtype="fp16"), + "full_attention_mask": IRValue( + id="full_attention_mask", + shape=(1, 1, 800, 800), + dtype="bool", + ), + "proj": IRValue(id="proj", shape=(1536, 4096), dtype="fp16"), + "out": IRValue(id="out", shape=(1, 800, 1536), dtype="fp16", producer="attn"), + }, + nodes={ + "attn": IRNode( + id="attn", + op="attention_block", + inputs=["query", "key", "value", "full_attention_mask", "proj"], + outputs=["out"], + attrs={ + "has_mask": True, + "is_causal": False, + "window_size": 0, + "attention_output_shape": (1, 800, 4096), + }, + meta={"attention_layer_type": "full_attention"}, + ) + }, + order=["attn"], + inputs=["query", "key", "value", "full_attention_mask"], + outputs=["out"], + constants={}, + meta={ + "adapter_family": "gemma4", + "component": "decoder", + "input_names": ("query", "key", "value", "full_attention_mask"), + }, + ) + + changed = normalize_gemma4_decoder_attention_semantics(graph) + + assert changed is True + assert graph.nodes["attn"].inputs == ["query", "key", "value", "proj"] + assert graph.nodes["attn"].attrs["has_mask"] is False + assert graph.nodes["attn"].attrs["is_causal"] is True + assert graph.nodes["attn"].attrs["window_size"] == 800 + assert graph.nodes["attn"].meta["gemma4_full_attention_window_compat"] is True + + +def test_normalize_gemma4_sliding_attention_elides_runtime_mask() -> None: + graph = IRGraph( + values={ + "query": IRValue(id="query", shape=(1, 8, 800, 256), dtype="fp16"), + "key": IRValue(id="key", shape=(1, 8, 800, 256), dtype="fp16"), + "value": IRValue(id="value", shape=(1, 8, 800, 256), dtype="fp16"), + "sliding_attention_mask": IRValue( + id="sliding_attention_mask", + shape=(1, 1, 800, 800), + dtype="bool", + ), + "proj": IRValue(id="proj", shape=(1536, 2048), dtype="fp16"), + "out": IRValue(id="out", shape=(1, 800, 1536), dtype="fp16", producer="attn"), + }, + nodes={ + "attn": IRNode( + id="attn", + op="attention_block", + inputs=["query", "key", "value", "sliding_attention_mask", "proj"], + outputs=["out"], + attrs={ + "has_mask": True, + "is_causal": False, + "window_size": 0, + "attention_output_shape": (1, 800, 2048), + }, + meta={"attention_layer_type": "sliding_attention"}, + ) + }, + order=["attn"], + inputs=["query", "key", "value", "sliding_attention_mask"], + outputs=["out"], + constants={}, + meta={ + "adapter_family": "gemma4", + "component": "decoder", + "sliding_window": 512, + "input_names": ("query", "key", "value", "sliding_attention_mask"), + }, + ) + + changed = normalize_gemma4_decoder_attention_semantics(graph) + + assert changed is True + assert graph.nodes["attn"].inputs == ["query", "key", "value", "proj"] + assert graph.nodes["attn"].attrs["has_mask"] is False + assert graph.nodes["attn"].attrs["is_causal"] is True + assert graph.nodes["attn"].attrs["window_size"] == 512 + + +def test_normalize_gemma4_attention_recovers_layer_hints_from_graph_meta() -> None: + graph = IRGraph( + values={ + "query_0": IRValue(id="query_0", shape=(1, 8, 800, 256), dtype="fp16"), + "key_0": IRValue(id="key_0", shape=(1, 8, 800, 256), dtype="fp16"), + "value_0": IRValue(id="value_0", shape=(1, 8, 800, 256), dtype="fp16"), + "sliding_attention_mask": IRValue( + id="sliding_attention_mask", + shape=(1, 1, 800, 800), + dtype="bool", + ), + "out_0": IRValue(id="out_0", shape=(1, 8, 800, 256), dtype="fp16", producer="attn_0"), + "query_1": IRValue(id="query_1", shape=(1, 8, 800, 256), dtype="fp16"), + "key_1": IRValue(id="key_1", shape=(1, 8, 800, 256), dtype="fp16"), + "value_1": IRValue(id="value_1", shape=(1, 8, 800, 256), dtype="fp16"), + "full_attention_mask": IRValue( + id="full_attention_mask", + shape=(1, 1, 800, 800), + dtype="bool", + ), + "out_1": IRValue(id="out_1", shape=(1, 8, 800, 256), dtype="fp16", producer="attn_1"), + }, + nodes={ + "attn_0": IRNode( + id="attn_0", + op="attention", + inputs=["query_0", "key_0", "value_0", "sliding_attention_mask"], + outputs=["out_0"], + attrs={"is_causal": False, "window_size": 0}, + ), + "attn_1": IRNode( + id="attn_1", + op="attention", + inputs=["query_1", "key_1", "value_1", "full_attention_mask"], + outputs=["out_1"], + attrs={"is_causal": False, "window_size": 0}, + ), + }, + order=["attn_0", "attn_1"], + inputs=[ + "query_0", + "key_0", + "value_0", + "sliding_attention_mask", + "query_1", + "key_1", + "value_1", + "full_attention_mask", + ], + outputs=["out_1"], + constants={}, + meta={ + "adapter_family": "gemma4", + "component": "decoder", + "sliding_window": 512, + "layer_types": ("sliding_attention", "full_attention"), + }, + ) + + changed = normalize_gemma4_decoder_attention_semantics(graph) + + assert changed is True + assert graph.nodes["attn_0"].meta["attention_layer_type"] == "sliding_attention" + assert graph.nodes["attn_0"].meta["attention_layer_index"] == 0 + assert graph.nodes["attn_0"].inputs == ["query_0", "key_0", "value_0"] + assert graph.nodes["attn_0"].attrs["is_causal"] is True + assert graph.nodes["attn_0"].attrs["window_size"] == 512 + assert graph.nodes["attn_1"].meta["attention_layer_type"] == "full_attention" + assert graph.nodes["attn_1"].meta["attention_layer_index"] == 1 + assert graph.nodes["attn_1"].inputs == ["query_1", "key_1", "value_1"] + assert graph.nodes["attn_1"].attrs["is_causal"] is True + assert graph.nodes["attn_1"].attrs["window_size"] == 800 + + +def test_normalize_gemma4_cached_attention_sets_per_layer_cache_lengths() -> None: + graph = IRGraph( + values={ + "query_0": IRValue(id="query_0", shape=(1, 8, 1, 256), dtype="fp16"), + "key_0": IRValue(id="key_0", shape=(1, 8, 1, 256), dtype="fp16"), + "value_0": IRValue(id="value_0", shape=(1, 8, 1, 256), dtype="fp16"), + "out_0": IRValue(id="out_0", shape=(1, 8, 1, 256), dtype="fp16", producer="attn_0"), + "query_1": IRValue(id="query_1", shape=(1, 8, 1, 256), dtype="fp16"), + "key_1": IRValue(id="key_1", shape=(1, 8, 1, 256), dtype="fp16"), + "value_1": IRValue(id="value_1", shape=(1, 8, 1, 256), dtype="fp16"), + "out_1": IRValue(id="out_1", shape=(1, 8, 1, 256), dtype="fp16", producer="attn_1"), + }, + nodes={ + "attn_0": IRNode( + id="attn_0", + op="attention", + inputs=["query_0", "key_0", "value_0"], + outputs=["out_0"], + attrs={"is_causal": False, "window_size": 0}, + ), + "attn_1": IRNode( + id="attn_1", + op="attention", + inputs=["query_1", "key_1", "value_1"], + outputs=["out_1"], + attrs={"is_causal": False, "window_size": 0}, + ), + }, + order=["attn_0", "attn_1"], + inputs=["query_0", "key_0", "value_0", "query_1", "key_1", "value_1"], + outputs=["out_1"], + constants={}, + meta={ + "adapter_family": "gemma4", + "component": "decoder_step", + "use_internal_kv_cache": True, + "sliding_window": 512, + "max_cache_seq_len": 131072, + "layer_types": ("sliding_attention", "full_attention"), + }, + ) + + changed = normalize_gemma4_decoder_attention_semantics(graph) + + assert changed is True + assert graph.nodes["attn_0"].attrs["window_size"] == 512 + assert graph.nodes["attn_0"].meta["max_cache_seq_len"] == 512 + assert graph.nodes["attn_1"].attrs["window_size"] == 0 + assert graph.nodes["attn_1"].meta["max_cache_seq_len"] == 131072 + + +def test_normalize_gemma4_attention_elides_logical_fp16_mask() -> None: + graph = IRGraph( + values={ + "query": IRValue(id="query", shape=(1, 8, 800, 256), dtype="fp16"), + "key": IRValue(id="key", shape=(1, 8, 800, 256), dtype="fp16"), + "value": IRValue(id="value", shape=(1, 8, 800, 256), dtype="fp16"), + "mask_a": IRValue(id="mask_a", shape=(1, 1, 800, 800), dtype="bool"), + "mask_b": IRValue(id="mask_b", shape=(1, 1, 800, 800), dtype="bool"), + "mask_fp16": IRValue(id="mask_fp16", shape=(1, 1, 800, 800), dtype="fp16", producer="mask_and"), + "out": IRValue(id="out", shape=(1, 8, 800, 256), dtype="fp16", producer="attn"), + }, + nodes={ + "mask_and": IRNode( + id="mask_and", + op="logical_and", + inputs=["mask_a", "mask_b"], + outputs=["mask_fp16"], + ), + "attn": IRNode( + id="attn", + op="attention", + inputs=["query", "key", "value", "mask_fp16"], + outputs=["out"], + attrs={"is_causal": False, "window_size": 0}, + meta={"attention_layer_type": "sliding_attention"}, + ), + }, + order=["mask_and", "attn"], + inputs=["query", "key", "value", "mask_a", "mask_b"], + outputs=["out"], + constants={}, + meta={ + "adapter_family": "gemma4", + "component": "decoder", + "sliding_window": 512, + }, + ) + + changed = normalize_gemma4_decoder_attention_semantics(graph) + + assert changed is True + assert graph.nodes["attn"].inputs == ["query", "key", "value"] + assert graph.nodes["attn"].attrs["is_causal"] is True + assert graph.nodes["attn"].attrs["window_size"] == 512 + + +def test_precompute_rope_tables_replaces_runtime_angle_with_fp16_lookup() -> None: + head_dim = 8 + max_seq = 4096 + theta = 10000.0 + inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + + graph = IRGraph( + values={ + "positions": IRValue(id="positions", shape=(1, 4), dtype="int32"), + "angle": IRValue(id="angle", shape=(1, 4, head_dim // 2), dtype="fp32", producer="matmul"), + "emb": IRValue(id="emb", shape=(1, 4, head_dim), dtype="fp32", producer="cat"), + "cos": IRValue(id="cos", shape=(1, 4, head_dim), dtype="fp32", producer="cos_node"), + }, + nodes={ + "matmul": IRNode(id="matmul", op="matmul", inputs=["positions", "inv_freq"], outputs=["angle"]), + "cat": IRNode(id="cat", op="cat", inputs=["angle", "angle"], outputs=["emb"], attrs={"dim": -1}), + "cos_node": IRNode( + id="cos_node", op="scalar_cos", inputs=["emb"], outputs=["cos"], + meta={"component": "decoder"}, + ), + }, + order=["matmul", "cat", "cos_node"], + inputs=["positions"], + outputs=["cos"], + constants={"inv_freq": inv_freq}, + meta={"component": "decoder", "max_cache_seq_len": max_seq}, + ) + rebuild_graph(graph) + + assert precompute_rope_tables(graph) is True + + cos_node = graph.nodes["cos_node"] + assert cos_node.op == "embedding" + table_id, position_id = cos_node.inputs + assert table_id.startswith("c_rope_table_") + assert position_id == "positions" + + table = graph.constants[table_id] + assert table.dtype == torch.float16 + assert tuple(table.shape) == (max_seq, head_dim) + assert graph.values["cos"].dtype == "fp16" + + # Positions past the fp16-angle range (>2048) stay distinct/representable in the table. + positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) + freqs = positions * inv_freq.to(torch.float64).reshape(1, -1) + expected = torch.cos(torch.cat((freqs, freqs), dim=-1)).to(torch.float16) + assert torch.equal(table[[0, 1, 2049, max_seq - 1]], expected[[0, 1, 2049, max_seq - 1]]) diff --git a/python/tests/test_server.py b/python/tests/test_server.py new file mode 100644 index 000000000..d2137a34d --- /dev/null +++ b/python/tests/test_server.py @@ -0,0 +1,272 @@ +from pathlib import Path +from types import SimpleNamespace +import asyncio + +import pytest + +pytest.importorskip("fastapi") + +from fastapi import HTTPException +from fastapi.testclient import TestClient + +import cactus.server as server + + +def _bundle(root: Path, name: str, model_type: str, *, manifest: bool = True) -> Path: + path = root / name + path.mkdir(parents=True) + (path / "config.txt").write_text( + f"model_type={model_type}\ncontext_length=1024\n", + encoding="utf-8", + ) + if manifest: + (path / "components").mkdir() + (path / "components" / "manifest.json").write_text('{"components":[]}', encoding="utf-8") + return path + + +def _app(tmp_path: Path, monkeypatch, *, preload: bool = False): + _bundle(tmp_path, "llm", "gemma4") + monkeypatch.setattr(server, "cactus_init", lambda *args, **kwargs: SimpleNamespace(args=args, kwargs=kwargs)) + monkeypatch.setattr(server, "cactus_destroy", lambda handle: None) + return server.create_app(weights_root=tmp_path, default_model="llm", preload=preload) + + +def test_models_lists_only_v2_bundles(tmp_path: Path, monkeypatch) -> None: + _bundle(tmp_path, "valid", "gemma4") + _bundle(tmp_path, "old_weights", "lfm2", manifest=False) + monkeypatch.setattr(server, "cactus_init", lambda *args, **kwargs: object()) + monkeypatch.setattr(server, "cactus_destroy", lambda handle: None) + + app = server.create_app(weights_root=tmp_path, default_model="valid", preload=False) + with TestClient(app) as client: + data = client.get("/v1/models").json()["data"] + + assert [m["id"] for m in data] == ["valid"] + assert data[0]["context_window"] == 1024 + + +def test_chat_completion_shapes_openai_response(tmp_path: Path, monkeypatch) -> None: + app = _app(tmp_path, monkeypatch) + + def fake_complete(handle, messages, options, tools, callback): + assert messages == [{"role": "user", "content": "hello"}] + assert options["max_tokens"] == 3 + return { + "success": True, + "response": "hi", + "function_calls": [], + "prefill_tokens": 2, + "decode_tokens": 1, + } + + monkeypatch.setattr(server, "cactus_complete", fake_complete) + monkeypatch.setattr(server, "cactus_reset", lambda handle: None) + + with TestClient(app) as client: + res = client.post("/v1/chat/completions", json={ + "model": "llm", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 3, + }) + + assert res.status_code == 200 + body = res.json() + assert body["object"] == "chat.completion" + assert body["choices"][0]["message"]["content"] == "hi" + assert body["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} + + +_WRAPPED_LOOKUP = [{"type": "function", "function": {"name": "lookup", "description": "", "parameters": {}}}] + + +def _capturing_complete(captured: dict): + def fake_complete(handle, messages, options, tools, callback): + captured["tools"] = tools + return { + "success": True, + "response": "", + "function_calls": [{"name": "lookup", "arguments": {"q": "x"}}], + "prefill_tokens": 1, + "decode_tokens": 1, + } + return fake_complete + + +def _tool_capture_app(tmp_path, monkeypatch): + app = _app(tmp_path, monkeypatch) + monkeypatch.setattr(server, "cactus_reset", lambda handle: None) + captured = {} + monkeypatch.setattr(server, "cactus_complete", _capturing_complete(captured)) + return app, captured + + +def test_chat_tool_calls_are_translated(tmp_path: Path, monkeypatch) -> None: + app, captured = _tool_capture_app(tmp_path, monkeypatch) + + with TestClient(app) as client: + res = client.post("/v1/chat/completions", json={ + "model": "llm", + "messages": [{"role": "user", "content": "call tool"}], + "tools": [{"type": "function", "function": {"name": "lookup", "parameters": {}}}], + "tool_choice": "required", + }) + + choice = res.json()["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] is None + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + assert choice["message"]["tool_calls"][0]["function"]["arguments"] == '{"q": "x"}' + assert captured["tools"] == _WRAPPED_LOOKUP + + +def test_streaming_chat_sends_wrapped_tools(tmp_path: Path, monkeypatch) -> None: + app, captured = _tool_capture_app(tmp_path, monkeypatch) + + with TestClient(app) as client: + with client.stream("POST", "/v1/chat/completions", json={ + "model": "llm", + "messages": [{"role": "user", "content": "call tool"}], + "tools": [{"type": "function", "function": {"name": "lookup", "parameters": {}}}], + "stream": True, + }) as res: + body = "".join(res.iter_text()) + + assert '"finish_reason": "tool_calls"' in body + assert "data: [DONE]" in body + assert captured["tools"] == _WRAPPED_LOOKUP + + +def test_tool_choice_object_selects_wrapped_tool(tmp_path: Path, monkeypatch) -> None: + app, captured = _tool_capture_app(tmp_path, monkeypatch) + + with TestClient(app) as client: + res = client.post("/v1/chat/completions", json={ + "model": "llm", + "messages": [{"role": "user", "content": "call tool"}], + "tools": [ + {"type": "function", "function": {"name": "lookup", "parameters": {}}}, + {"type": "function", "function": {"name": "other", "parameters": {}}}, + ], + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + }) + + assert res.status_code == 200 + assert captured["tools"] == _WRAPPED_LOOKUP + + +def test_streaming_completion_error_terminates(tmp_path: Path, monkeypatch) -> None: + app = _app(tmp_path, monkeypatch) + monkeypatch.setattr(server, "cactus_reset", lambda handle: None) + + def fail_complete(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(server, "cactus_complete", fail_complete) + + with TestClient(app) as client: + with client.stream("POST", "/v1/chat/completions", json={ + "model": "llm", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + }) as res: + body = "".join(res.iter_text()) + + assert res.status_code == 200 + # OpenAI streaming spec carries errors in a `data:` JSON line, not in + # `event:` lines. We also emit a final chunk with finish_reason so + # clients don't hang before [DONE]. + assert '"error"' in body + assert "boom" in body + assert '"finish_reason": "stop"' in body + assert "data: [DONE]" in body + + +def test_transcription_rejects_non_wav(tmp_path: Path, monkeypatch) -> None: + _bundle(tmp_path, "llm", "gemma4") + _bundle(tmp_path, "stt", "parakeet_tdt") + monkeypatch.setattr(server, "cactus_init", lambda *args, **kwargs: object()) + monkeypatch.setattr(server, "cactus_destroy", lambda handle: None) + app = server.create_app(weights_root=tmp_path, default_model="llm", preload=False) + + with TestClient(app) as client: + res = client.post( + "/v1/audio/transcriptions", + files={"file": ("audio.mp3", b"not wav", "audio/mpeg")}, + data={"model": "stt"}, + ) + + assert res.status_code == 400 + assert "Only .wav" in res.json()["detail"] + + +def test_transcription_verbose_json_with_segments(tmp_path: Path, monkeypatch) -> None: + _bundle(tmp_path, "llm", "gemma4") + _bundle(tmp_path, "stt", "parakeet_tdt") + monkeypatch.setattr(server, "cactus_init", lambda *args, **kwargs: object()) + monkeypatch.setattr(server, "cactus_destroy", lambda handle: None) + monkeypatch.setattr(server, "cactus_reset", lambda handle: None) + monkeypatch.setattr(server, "cactus_transcribe", lambda *args, **kwargs: { + "success": True, + "response": "hello", + "segments": [{"start": 0.0, "end": 1.0, "text": "hello"}], + }) + app = server.create_app(weights_root=tmp_path, default_model="llm", preload=False) + + with TestClient(app) as client: + res = client.post( + "/v1/audio/transcriptions", + files={"file": ("audio.wav", b"wav", "audio/wav")}, + data={ + "model": "stt", + "response_format": "verbose_json", + "timestamp_granularities[]": "segment", + }, + ) + + assert res.status_code == 200 + assert res.json()["segments"] == [{"id": 0, "start": 0.0, "end": 1.0, "text": "hello"}] + + +def test_transcription_rejects_word_timestamps(tmp_path: Path, monkeypatch) -> None: + _bundle(tmp_path, "llm", "gemma4") + _bundle(tmp_path, "stt", "parakeet_tdt") + monkeypatch.setattr(server, "cactus_init", lambda *args, **kwargs: object()) + monkeypatch.setattr(server, "cactus_destroy", lambda handle: None) + app = server.create_app(weights_root=tmp_path, default_model="llm", preload=False) + + with TestClient(app) as client: + res = client.post( + "/v1/audio/transcriptions", + files={"file": ("audio.wav", b"wav", "audio/wav")}, + data={"model": "stt", "timestamp_granularities[]": "word"}, + ) + + assert res.status_code == 400 + assert "Word timestamp" in res.json()["detail"] + + +def test_model_manager_does_not_evict_active_slots(tmp_path: Path, monkeypatch) -> None: + for name in ("a", "b", "c"): + _bundle(tmp_path, name, "gemma4") + destroyed = [] + monkeypatch.setattr(server, "cactus_init", lambda path, **kwargs: Path(path).name) + monkeypatch.setattr(server, "cactus_destroy", lambda handle: destroyed.append(handle)) + + registry = server.ModelRegistry(tmp_path) + manager = server.ModelManager(registry, max_warm=2) + + async def run(): + async with manager.acquire("a"): + async with manager.acquire("b"): + with pytest.raises(HTTPException) as exc: + async with manager.acquire("c"): + pass + assert exc.value.status_code == 503 + assert destroyed == [] + + async with manager.acquire("c"): + pass + + asyncio.run(run()) + assert len(destroyed) == 1 diff --git a/python/tests/test_server_live.py b/python/tests/test_server_live.py new file mode 100644 index 000000000..e4091c334 --- /dev/null +++ b/python/tests/test_server_live.py @@ -0,0 +1,518 @@ +from __future__ import annotations + +import json +import os +import socket +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path + +import httpx +import pytest + +from .bundles import ( + PROJECT_ROOT, WEIGHTS, _find_bundle, _iter_bundle_candidates, + _read_model_type, _valid_bundle, +) + +ASSETS = PROJECT_ROOT / "cactus-engine" / "tests" / "assets" +LLM_TYPES = {"gemma", "gemma3n", "gemma4", "lfm2", "qwen", "qwen3p5", "needle", "youtu"} +STT_TYPES = {"whisper", "parakeet_tdt", "parakeet-tdt"} + + +def _default_llm_bundle() -> Path: + """Locate the canonical gemma-4-e2b-it LLM bundle under whichever convention + it was built with (bare `gemma-4-e2b-it` or suffixed `...-cq4`).""" + for candidate in _iter_bundle_candidates("gemma-4-e2b-it"): + if _valid_bundle(candidate) and _read_model_type(candidate) in LLM_TYPES: + return candidate + return WEIGHTS / "gemma-4-e2b-it" # fall through so _require_bundle reports it + + +def _require_bundle(relative: Path, types: set[str]) -> Path: + candidate = PROJECT_ROOT / relative + if not candidate.exists(): + pytest.fail( + f"Live-test model not found: {relative}\n" + f"Prepare it with `cactus convert google/gemma-4-E2B-it` from {PROJECT_ROOT}." + ) + if not _valid_bundle(candidate): + pytest.fail( + f"Live-test model is not a prepared v2 bundle: {relative}\n" + "Expected config.txt and components/manifest.json.\n" + f"Prepare it with `cactus convert google/gemma-4-E2B-it` from {PROJECT_ROOT}." + ) + model_type = _read_model_type(candidate) + if model_type not in types: + pytest.fail( + f"Live-test model has unsupported type {model_type!r}: {relative}\n" + f"Expected one of: {sorted(types)}." + ) + return candidate + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _start_server(bundle: Path, port: int) -> subprocess.Popen: + env = os.environ.copy() + python_path = str(PROJECT_ROOT / "python") + env["PYTHONPATH"] = python_path + os.pathsep + env.get("PYTHONPATH", "") + return subprocess.Popen( + [ + sys.executable, + "-m", + "cactus", + "serve", + str(bundle), + "--host", + "127.0.0.1", + "--port", + str(port), + ], + cwd=PROJECT_ROOT, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + +def _wait_ready(proc: subprocess.Popen, base_url: str) -> None: + deadline = time.time() + 120 + last_error = None + while time.time() < deadline: + if proc.poll() is not None: + output = proc.stdout.read() if proc.stdout else "" + pytest.fail(f"server exited before becoming ready:\n{output}") + try: + with httpx.Client(timeout=2) as client: + res = client.get(f"{base_url}/v1/models") + if res.status_code == 200: + return + except Exception as exc: + last_error = exc + time.sleep(0.5) + pytest.fail(f"server did not become ready: {last_error}") + + +@contextmanager +def _serve(bundle: Path): + port = _free_port() + proc = _start_server(bundle, port) + base_url = f"http://127.0.0.1:{port}" + try: + _wait_ready(proc, base_url) + yield base_url, bundle.name + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=5) + + +@pytest.fixture(scope="module") +def live_server(): + bundle = _default_llm_bundle() + _require_bundle(bundle, LLM_TYPES) + with _serve(bundle) as server: + yield server + + +def test_live_models(live_server) -> None: + base_url, _ = live_server + res = httpx.get(f"{base_url}/v1/models", timeout=10) + assert res.status_code == 200 + models = res.json()["data"] + assert models + for model in models: + bundle = WEIGHTS / model["id"] + assert _valid_bundle(bundle) + + +def test_live_chat_rejects_unknown_model(live_server) -> None: + base_url, _ = live_server + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={ + "model": "definitely-not-a-live-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 4, + }, + timeout=30, + ) + assert res.status_code == 404 + + +def test_live_chat_completion(live_server) -> None: + base_url, model = live_server + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": "Reply with one word."}], "max_tokens": 4}, + timeout=120, + ) + assert res.status_code == 200 + body = res.json() + assert body["choices"][0]["message"]["content"] + assert body["usage"]["completion_tokens"] > 0 + + +def test_live_chat_text_content_parts(live_server) -> None: + base_url, model = live_server + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={ + "model": model, + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "Reply with exactly "}, + {"type": "text", "text": "OK"}, + ], + }], + "max_tokens": 8, + "temperature": 0.0, + }, + timeout=180, + ) + assert res.status_code == 200 + assert res.json()["choices"][0]["message"]["content"] + + +def test_live_long_gemma_generation(live_server) -> None: + base_url, model = live_server + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={ + "model": model, + "messages": [{ + "role": "user", + "content": ( + "Write a numbered list of ten short items about why local HTTP " + "inference servers need end-to-end tests. Keep each item concise." + ), + }], + "max_tokens": 96, + "temperature": 0.0, + }, + timeout=180, + ) + assert res.status_code == 200 + body = res.json() + text = body["choices"][0]["message"]["content"] + assert isinstance(text, str) + assert len(text.split()) >= 20 + assert body["usage"]["completion_tokens"] >= 16 + assert body["choices"][0]["finish_reason"] == "stop" + + +def test_live_multi_turn_conversation(live_server) -> None: + base_url, model = live_server + messages = [ + {"role": "system", "content": "You answer directly and remember the user's chosen keyword."}, + {"role": "user", "content": "My keyword is cactus. Reply with only: remembered."}, + {"role": "assistant", "content": "remembered."}, + {"role": "user", "content": "What keyword did I give you? Reply with only the keyword."}, + ] + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={ + "model": model, + "messages": messages, + "max_tokens": 12, + "temperature": 0.0, + }, + timeout=180, + ) + assert res.status_code == 200 + body = res.json() + text = body["choices"][0]["message"]["content"].lower() + assert "cactus" in text + assert body["usage"]["prompt_tokens"] > body["usage"]["completion_tokens"] + + +def test_live_tool_call_completion(live_server) -> None: + base_url, model = live_server + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": "Use the weather tool for Paris."}], + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }], + "tool_choice": "required", + "max_tokens": 64, + "temperature": 0.0, + }, + timeout=180, + ) + assert res.status_code == 200 + body = res.json() + choice = body["choices"][0] + assert choice["finish_reason"] in {"tool_calls", "stop"} + if choice["finish_reason"] == "tool_calls": + calls = choice["message"]["tool_calls"] + assert calls + assert calls[0]["type"] == "function" + assert calls[0]["function"]["name"] + + +def test_live_concurrent_chat_requests(live_server) -> None: + base_url, model = live_server + + def request(i: int) -> int: + res = httpx.post( + f"{base_url}/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": f"Reply with the number {i}."}], + "max_tokens": 8, + "temperature": 0.0, + }, + timeout=180, + ) + assert res.status_code == 200 + assert res.json()["choices"][0]["message"]["content"] + return res.status_code + + with ThreadPoolExecutor(max_workers=2) as pool: + assert list(pool.map(request, [1, 2])) == [200, 200] + + +def test_live_chat_stream(live_server) -> None: + base_url, model = live_server + with httpx.stream( + "POST", + f"{base_url}/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": "Say hi."}], "max_tokens": 4, "stream": True}, + timeout=120, + ) as res: + assert res.status_code == 200 + text = "".join(res.iter_text()) + assert "chat.completion.chunk" in text + assert "data: [DONE]" in text + + +def test_live_long_streaming_generation(live_server) -> None: + base_url, model = live_server + with httpx.stream( + "POST", + f"{base_url}/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": "Write one short paragraph about robust server testing."}], + "max_tokens": 64, + "stream": True, + }, + timeout=180, + ) as res: + assert res.status_code == 200 + text = "".join(res.iter_text()) + assert "chat.completion.chunk" in text + assert text.count('"delta"') >= 3 + assert "data: [DONE]" in text + + +def test_live_transcription_wav(live_server) -> None: + base_url, _ = live_server + stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + audio = ASSETS / "test.wav" + assert audio.exists() + with audio.open("rb") as f: + res = httpx.post( + f"{base_url}/v1/audio/transcriptions", + data={"model": stt.name, "response_format": "json"}, + files={"file": ("test.wav", f, "audio/wav")}, + timeout=180, + ) + assert res.status_code == 200 + assert res.json()["text"] + + +def test_live_transcription_text_response(live_server) -> None: + base_url, _ = live_server + stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + audio = ASSETS / "test.wav" + with audio.open("rb") as f: + res = httpx.post( + f"{base_url}/v1/audio/transcriptions", + data={"model": stt.name, "response_format": "text"}, + files={"file": ("test.wav", f, "audio/wav")}, + timeout=180, + ) + assert res.status_code == 200 + assert res.text.strip() + assert res.headers["content-type"].startswith("text/plain") + + +def test_live_transcription_verbose_json_segments(live_server) -> None: + base_url, _ = live_server + stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + audio = ASSETS / "test.wav" + with audio.open("rb") as f: + res = httpx.post( + f"{base_url}/v1/audio/transcriptions", + data={ + "model": stt.name, + "response_format": "verbose_json", + "timestamp_granularities[]": "segment", + }, + files={"file": ("test.wav", f, "audio/wav")}, + timeout=180, + ) + assert res.status_code == 200 + body = res.json() + assert body["text"] + assert "segments" in body + + +def test_live_transcription_rejects_unsupported_format(live_server) -> None: + base_url, _ = live_server + stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + audio = ASSETS / "test.wav" + with audio.open("rb") as f: + res = httpx.post( + f"{base_url}/v1/audio/transcriptions", + data={"model": stt.name, "response_format": "srt"}, + files={"file": ("test.wav", f, "audio/wav")}, + timeout=30, + ) + assert res.status_code == 400 + assert "Unsupported transcription response_format" in json.dumps(res.json()) + + +def test_live_transcription_rejects_word_timestamps(live_server) -> None: + base_url, _ = live_server + stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + audio = ASSETS / "test.wav" + with audio.open("rb") as f: + res = httpx.post( + f"{base_url}/v1/audio/transcriptions", + data={ + "model": stt.name, + "response_format": "verbose_json", + "timestamp_granularities[]": "word", + }, + files={"file": ("test.wav", f, "audio/wav")}, + timeout=30, + ) + assert res.status_code == 400 + assert "Word timestamp" in json.dumps(res.json()) + + +def test_live_transcription_rejects_non_wav(live_server) -> None: + base_url, _ = live_server + stt = _find_bundle(["parakeet-tdt-0.6b-v3-transpiled"], STT_TYPES) + res = httpx.post( + f"{base_url}/v1/audio/transcriptions", + data={"model": stt.name}, + files={"file": ("test.mp3", b"nope", "audio/mpeg")}, + timeout=30, + ) + assert res.status_code == 400 + assert "Only .wav" in json.dumps(res.json()) + + +EMBED_TYPES = {"bert", "nomic"} + + +def _find_embed_bundle() -> Path: + for candidate in sorted(WEIGHTS.iterdir()) if WEIGHTS.exists() else []: + if candidate.is_dir() and _valid_bundle(candidate) and _read_model_type(candidate) in EMBED_TYPES: + return candidate + raise RuntimeError( + "No embedding (nomic/bert) bundle under weights/; " + "run `cactus convert nomic-ai/nomic-embed-text-v2-moe`" + ) + + +def _supports_embedding(bundle: Path) -> bool: + try: + data = json.loads((bundle / "components" / "manifest.json").read_text(encoding="utf-8")) + except Exception: + return False + names = {str(c.get("component", "")) for c in data.get("components", []) if isinstance(c, dict)} + return "text_embedding" in names or "decoder_embed_chunk" in names + + +def _find_non_embed_llm_bundle() -> Path | None: + for candidate in sorted(WEIGHTS.iterdir()) if WEIGHTS.exists() else []: + if (candidate.is_dir() and _valid_bundle(candidate) + and _read_model_type(candidate) in LLM_TYPES + and not _supports_embedding(candidate)): + return candidate + return None + + +@pytest.fixture(scope="module") +def embed_server(): + with _serve(_find_embed_bundle()) as server: + yield server + + +@pytest.fixture(scope="module") +def non_embed_server(): + bundle = _find_non_embed_llm_bundle() + if bundle is None: + pytest.skip("No non-embedding LLM bundle under weights/") + with _serve(bundle) as server: + yield server + + +def test_live_embeddings_string(embed_server) -> None: + base_url, model = embed_server + res = httpx.post( + f"{base_url}/v1/embeddings", + json={"model": model, "input": "Paris is the capital of France."}, + timeout=60, + ) + assert res.status_code == 200, res.text + body = res.json() + assert body["object"] == "list" + assert len(body["data"]) == 1 + assert body["data"][0]["object"] == "embedding" + assert body["data"][0]["index"] == 0 + assert len(body["data"][0]["embedding"]) > 0 + assert all(isinstance(x, float) for x in body["data"][0]["embedding"][:8]) + + +def test_live_embeddings_list(embed_server) -> None: + base_url, model = embed_server + res = httpx.post( + f"{base_url}/v1/embeddings", + json={"model": model, "input": ["hello world", "a second sentence"]}, + timeout=60, + ) + assert res.status_code == 200, res.text + data = res.json()["data"] + assert [d["index"] for d in data] == [0, 1] + assert len(data[0]["embedding"]) == len(data[1]["embedding"]) > 0 + + +def test_live_embeddings_rejects_llm_model(non_embed_server) -> None: + base_url, model = non_embed_server + res = httpx.post( + f"{base_url}/v1/embeddings", + json={"model": model, "input": "hello"}, + timeout=30, + ) + assert res.status_code == 400 + assert "not an embedding model" in json.dumps(res.json()) diff --git a/python/tests/test_tool_constraint_state.py b/python/tests/test_tool_constraint_state.py new file mode 100644 index 000000000..7a1096855 --- /dev/null +++ b/python/tests/test_tool_constraint_state.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +import json + +import pytest + +from cactus import cactus_complete, cactus_destroy, cactus_init + +from .bundles import WEIGHTS, _iter_bundle_candidates, _valid_bundle + +BUNDLE = next( + (c for c in _iter_bundle_candidates("gemma-4-e2b-it") if _valid_bundle(c)), + WEIGHTS / "gemma-4-e2b-it", +) + +def _tool(name, description, properties, required): + return { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": {"type": "object", "properties": properties, "required": required}, + }, + } + + +GET_WEATHER = _tool("get_weather", "Get the current weather for a location", + {"location": {"type": "string", "description": "City name"}}, ["location"]) +SET_ALARM = _tool("set_alarm", "Set an alarm for a time", + {"time": {"type": "string", "description": "HH:MM"}}, ["time"]) +SEND = _tool("send", "Send a quick ping (no body) to a user id", + {"user": {"type": "string"}}, ["user"]) +SEND_MESSAGE = _tool("send_message", "Send a text message with a body to a contact", + {"recipient": {"type": "string"}, "body": {"type": "string"}}, ["recipient", "body"]) +SEARCH = _tool("search", "Search local notes", {"query": {"type": "string"}}, ["query"]) +SEARCH_WEB = _tool("search_web", "Search the public web", {"query": {"type": "string"}}, ["query"]) +GET = _tool("get", "Get a raw value by key", {"key": {"type": "string"}}, ["key"]) +GET_USER = _tool("get_user", "Fetch a user profile by id", {"id": {"type": "string"}}, ["id"]) +PING_SERVER = _tool("ping_server", "Check server health, takes no arguments", {}, []) +WEATHER_UNIT = _tool("get_weather", "Get weather for a city in a unit", + {"location": {"type": "string", "description": "city"}, + "unit": {"type": "string", "description": "temperature unit", + "enum": ["celsius", "fahrenheit"]}}, + ["location", "unit"]) +CREATE_EVENT = _tool("create_event", "Create a calendar event", + {"title": {"type": "string"}, "date": {"type": "string", "description": "YYYY-MM-DD"}, + "location": {"type": "string"}}, + ["title", "date"]) + + +def _forced_calls(model, tools, user): + result = cactus_complete( + model, + [ + {"role": "system", "content": "You are a helpful assistant that can use tools."}, + {"role": "user", "content": user}, + ], + {"force_tools": True, "max_tokens": 128, "temperature": 0}, + tools, + ) + assert result.get("success"), result + calls = result.get("function_calls") or [] + return [c if isinstance(c, dict) else json.loads(c) for c in calls] + + +@pytest.fixture(scope="module") +def model(): + if not _valid_bundle(BUNDLE): + pytest.skip(f"Live-test model not found: {BUNDLE}") + handle = cactus_init(str(BUNDLE)) + yield handle + cactus_destroy(handle) + + +def test_turn_after_force_tools_is_not_degenerate(model): + tools = [GET_WEATHER] + messages = [ + {"role": "system", "content": "You are a helpful assistant that can use tools."}, + {"role": "user", "content": "What's the weather in Tokyo right now?"}, + ] + forced = cactus_complete( + model, messages, {"force_tools": True, "max_tokens": 200, "temperature": 0}, tools + ) + assert forced.get("success"), forced + assert forced.get("function_calls"), forced + + messages.append({ + "role": "assistant", + "content": forced.get("response", ""), + "tool_calls": [{"type": "function", "function": { + "name": "get_weather", "arguments": {"location": "Tokyo"}, + }}], + }) + messages.append({ + "role": "tool", + "name": "get_weather", + "content": json.dumps({"temperature_c": 18}), + }) + followup = cactus_complete(model, messages, {"max_tokens": 150, "temperature": 0}, tools) + assert followup.get("success"), followup + response = followup.get("response", "") + assert response.strip() + assert "<|tool_call>" not in response + + +def test_prefix_name_is_not_shadowed(model): + calls = _forced_calls(model, [SEND, SEND_MESSAGE], "Message my contact Alice: running late, start without me.") + assert calls + assert calls[0]["name"] == "send_message" + + +@pytest.mark.parametrize("tools, user, want", [ + ([SEARCH, SEARCH_WEB], "Search the public web for the best ramen in Osaka.", "search_web"), + ([GET, GET_USER], "Fetch the user profile for account 42.", "get_user"), +]) +def test_prefix_pairs_reach_longer_name(model, tools, user, want): + calls = _forced_calls(model, tools, user) + assert calls + assert calls[0]["name"] == want + + +def test_shorter_prefix_name_still_reachable(model): + calls = _forced_calls(model, [SEND, SEND_MESSAGE], "Ping user u_42.") + assert calls + assert calls[0]["name"] == "send" + + +def test_forced_name_is_always_a_declared_tool(model): + tools = [GET_WEATHER, SET_ALARM] + declared = {t["function"]["name"] for t in tools} + calls = _forced_calls(model, tools, "What's the weather in Tokyo right now?") + assert calls + assert calls[0]["name"] in declared + + +def test_no_argument_function_emits_empty_object(model): + calls = _forced_calls(model, [PING_SERVER], "Check whether the server is healthy.") + assert calls + assert calls[0]["name"] == "ping_server" + assert calls[0]["arguments"] == {} + + +def test_argument_keys_are_declared_parameters(model): + calls = _forced_calls(model, [GET_WEATHER], "What's the weather in Berlin?") + assert calls + assert calls[0]["name"] == "get_weather" + assert set(calls[0]["arguments"]).issubset({"location"}) + + +@pytest.mark.parametrize("user", [ + "What's the weather in Tokyo in kelvin?", + "Weather in Paris, in Kelvin please.", + "Weather in Berlin.", +]) +def test_enum_argument_restricted_to_allowed_values(model, user): + calls = _forced_calls(model, [WEATHER_UNIT], user) + assert calls + assert calls[0]["arguments"].get("unit") in ("celsius", "fahrenheit") + + +@pytest.mark.parametrize("user", [ + "Make a calendar event.", + "Add an event called Standup.", + "Schedule a dentist appointment.", +]) +def test_required_arguments_are_always_present(model, user): + calls = _forced_calls(model, [CREATE_EVENT], user) + assert calls + args = calls[0]["arguments"] + assert "title" in args and "date" in args + + +def test_enum_with_shared_prefixes_can_reach_longest_value(model): + tool = _tool("pick_plan", "pick a plan", + {"plan": {"type": "string", "enum": ["pro", "pro_plus", "pro_plus_max"]}}, ["plan"]) + calls = _forced_calls(model, [tool], "Pick the pro plus max plan.") + assert calls + assert calls[0]["arguments"].get("plan") in ("pro", "pro_plus", "pro_plus_max") + + +def test_enum_value_containing_spaces(model): + tool = _tool("move_card", "move a kanban card", + {"card": {"type": "string"}, + "column": {"type": "string", "enum": ["to do", "in progress", "done"]}}, + ["card", "column"]) + calls = _forced_calls(model, [tool], "Move the login card to in progress.") + assert calls + assert calls[0]["arguments"].get("column") in ("to do", "in progress", "done") + + +def test_single_value_enum_is_forced(model): + tool = _tool("finalize", "finalize", {"status": {"type": "string", "enum": ["confirmed"]}}, ["status"]) + calls = _forced_calls(model, [tool], "Finalize it however you want.") + assert calls + assert calls[0]["arguments"].get("status") == "confirmed" + + +def test_prefix_parameter_names_both_required_present(model): + tool = _tool("make_range", "range", {"to": {"type": "string"}, "total": {"type": "string"}}, ["to", "total"]) + calls = _forced_calls(model, [tool], "Make a range up to 10 with a total of 100.") + assert calls + assert {"to", "total"} <= set(calls[0]["arguments"]) + + +def test_many_required_arguments_all_present(model): + tool = _tool("register", "register a device", + {k: {"type": "string"} for k in ("name", "os", "owner", "ip", "room")}, + ["name", "os", "owner", "ip", "room"]) + calls = _forced_calls(model, [tool], "Register a device.") + assert calls + assert {"name", "os", "owner", "ip", "room"} <= set(calls[0]["arguments"]) + + +def test_prefix_tool_with_required_enum_on_longer_name(model): + send = _tool("send", "ping a user", {"user": {"type": "string"}}, ["user"]) + send_message = _tool("send_message", "send a message", + {"to": {"type": "string"}, "body": {"type": "string"}, + "urgency": {"type": "string", "enum": ["normal", "urgent"]}}, + ["to", "body", "urgency"]) + calls = _forced_calls(model, [send, send_message], "Send an urgent message to Alice saying the server is down.") + assert calls + assert calls[0]["name"] == "send_message" + assert calls[0]["arguments"].get("urgency") in ("normal", "urgent") + assert {"to", "body", "urgency"} <= set(calls[0]["arguments"]) diff --git a/python/tests/test_transpile_aten_ops.py b/python/tests/test_transpile_aten_ops.py new file mode 100644 index 000000000..d10c785bd --- /dev/null +++ b/python/tests/test_transpile_aten_ops.py @@ -0,0 +1,272 @@ +from __future__ import annotations + +from collections import Counter + +import torch + +from cactus.transpile.aten_ops import canonical_torch_op +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.normalize import normalize_target +from cactus.transpile.optimize_graph import fuse_dense_mlp_tq + + +class OddlyNamedLinearBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.pear_tree = torch.nn.Linear(4, 3, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(self.pear_tree(x)) + + +class SameOpsDifferentNames(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.not_a_transformer_layer = torch.nn.ModuleDict( + {"banana": torch.nn.Linear(4, 3, bias=False)} + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(self.not_a_transformer_layer["banana"](x)) + + +def _op_counts(module: torch.nn.Module) -> Counter[str]: + example = torch.randn(2, 4) + captured = capture_model(module, (example,)) + return Counter(captured.ir_graph.nodes[node_id].op for node_id in captured.ir_graph.order) + + +def test_aten_targets_normalize_from_real_op_overloads() -> None: + assert canonical_torch_op(torch.ops.aten.add.Tensor) == "aten.add.Tensor" + assert normalize_target(torch.ops.aten.add.Tensor) == "add" + assert normalize_target(torch.ops.aten.linear.default) == "linear" + assert normalize_target("torch.ops.aten.silu.default") == "silu" + + +def test_import_is_based_on_aten_ops_not_module_names() -> None: + first = _op_counts(OddlyNamedLinearBlock()) + second = _op_counts(SameOpsDifferentNames()) + + assert first == second + assert first["linear"] == 1 + assert first["silu"] == 1 + + +def test_import_records_canonical_aten_op_metadata() -> None: + captured = capture_model(OddlyNamedLinearBlock(), (torch.randn(2, 4),)) + node_meta = [captured.ir_graph.nodes[node_id].meta for node_id in captured.ir_graph.order] + + assert {meta["aten_op"] for meta in node_meta} >= {"aten.linear.default", "aten.silu.default"} + assert all("torch_name" in meta for meta in node_meta) + assert all("module_paths" not in meta for meta in node_meta) + + +def test_dense_mlp_fusion_uses_topology_not_layer_names() -> None: + graph = IRGraph( + values={ + "x": IRValue(id="x", shape=(1, 4), dtype="fp16"), + "pear_weight": IRValue(id="pear_weight", shape=(8, 4), dtype="fp16", meta={"path": "pear.cq4.weights"}), + "banana_weight": IRValue(id="banana_weight", shape=(8, 4), dtype="fp16", meta={"path": "banana.cq4.weights"}), + "plum_weight": IRValue(id="plum_weight", shape=(4, 8), dtype="fp16", meta={"path": "plum.cq4.weights"}), + "gate": IRValue(id="gate", shape=(1, 8), dtype="fp16", producer="gate_linear"), + "activated": IRValue(id="activated", shape=(1, 8), dtype="fp16", producer="gate_gelu"), + "up": IRValue(id="up", shape=(1, 8), dtype="fp16", producer="up_linear"), + "product": IRValue(id="product", shape=(1, 8), dtype="fp16", producer="product"), + "out": IRValue(id="out", shape=(1, 4), dtype="fp16", producer="down_linear"), + }, + nodes={ + "gate_linear": IRNode( + id="gate_linear", + op="linear", + inputs=["x", "pear_weight"], + outputs=["gate"], + attrs={"has_bias": False}, + ), + "gate_gelu": IRNode( + id="gate_gelu", + op="gelu", + inputs=["gate"], + outputs=["activated"], + ), + "up_linear": IRNode( + id="up_linear", + op="linear", + inputs=["x", "banana_weight"], + outputs=["up"], + attrs={"has_bias": False}, + ), + "product": IRNode( + id="product", + op="multiply", + inputs=["activated", "up"], + outputs=["product"], + ), + "down_linear": IRNode( + id="down_linear", + op="linear", + inputs=["product", "plum_weight"], + outputs=["out"], + attrs={"has_bias": False}, + ), + }, + order=["gate_linear", "gate_gelu", "up_linear", "product", "down_linear"], + inputs=["x"], + outputs=["out"], + constants={}, + meta={}, + ) + + assert fuse_dense_mlp_tq(graph) is True + assert graph.nodes["down_linear"].op == "dense_mlp_tq_fused" + assert graph.nodes["down_linear"].inputs == ["x", "pear_weight", "banana_weight", "plum_weight"] + + + + + + +def _ensure_test_lfm2_moe_op_registered() -> None: + from cactus.transpile.model_adapters import _ensure_lfm2_moe_custom_op_registered + _ensure_lfm2_moe_custom_op_registered() + + +def _ensure_test_qwen2_moe_op_registered() -> None: + from cactus.transpile.model_adapters import _ensure_qwen2_moe_custom_op_registered + _ensure_qwen2_moe_custom_op_registered() + + +def _ensure_test_gemma4_moe_op_registered() -> None: + from cactus.transpile.model_adapters import _ensure_gemma4_moe_custom_op_registered + _ensure_gemma4_moe_custom_op_registered() + + +class _FakeLfm2MoeBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate = torch.nn.Linear(4, 2, bias=False) + self.w1 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w3 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w2 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(4, 3)) for _ in range(2)]) + self.register_buffer("expert_bias", torch.zeros(2, dtype=torch.float32)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + return torch.ops.cactus_transpile.lfm2_moe_layer_gated( + hidden_flat, + self.gate(hidden_flat), + self.expert_bias, + list(self.w1), + list(self.w3), + list(self.w2), + 2, + 1, + True, + True, + 1.0e-6, + 1.0, + ).reshape(hidden_states.shape) + + +class _FakeQwen2MoeBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate = torch.nn.Linear(4, 2, bias=False) + self.w1 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w3 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w2 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(4, 3)) for _ in range(2)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + return torch.ops.cactus_transpile.qwen2_moe_layer_gated( + hidden_flat, + self.gate(hidden_flat), + list(self.w1), + list(self.w3), + list(self.w2), + 2, + 1, + False, + 1.0e-6, + 1.0, + ).reshape(hidden_states.shape) + + +class _FakeGemma4MoeBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.router = torch.nn.Linear(4, 2, bias=False) + self.w1 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w3 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w2 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(4, 3)) for _ in range(2)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + return torch.ops.cactus_transpile.gemma4_moe_layer_gated( + hidden_flat, + self.router(hidden_flat), + list(self.w1), + list(self.w3), + list(self.w2), + 2, + 1, + True, + 1.0e-6, + 1.0, + ).reshape(hidden_states.shape) + + +def test_lfm2_moe_runtime_block_exports_single_semantic_op() -> None: + _ensure_test_lfm2_moe_op_registered() + block = _FakeLfm2MoeBlock().eval() + captured = capture_model(block, (torch.randn(1, 2, 4),), strict=False) + + moe_nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "lfm2_moe_layer_gated" + ] + + assert len(moe_nodes) == 1 + assert moe_nodes[0].attrs["num_experts"] == 2 + assert moe_nodes[0].attrs["num_experts_per_tok"] == 1 + assert moe_nodes[0].attrs["use_expert_bias"] is True + assert len(moe_nodes[0].inputs) == 3 + 3 * 2 + + +def test_qwen2_moe_runtime_block_exports_single_semantic_op() -> None: + _ensure_test_qwen2_moe_op_registered() + block = _FakeQwen2MoeBlock().eval() + captured = capture_model(block, (torch.randn(1, 2, 4),), strict=False) + + moe_nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "qwen2_moe_layer_gated" + ] + + assert len(moe_nodes) == 1 + assert moe_nodes[0].attrs["num_experts"] == 2 + assert moe_nodes[0].attrs["num_experts_per_tok"] == 1 + assert moe_nodes[0].attrs["normalize_routing"] is False + assert len(moe_nodes[0].inputs) == 2 + 3 * 2 + + +def test_gemma4_moe_runtime_block_exports_single_semantic_op() -> None: + _ensure_test_gemma4_moe_op_registered() + block = _FakeGemma4MoeBlock().eval() + captured = capture_model(block, (torch.randn(1, 2, 4),), strict=False) + + moe_nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "gemma4_moe_layer_gated" + ] + + assert len(moe_nodes) == 1 + assert moe_nodes[0].attrs["num_experts"] == 2 + assert moe_nodes[0].attrs["num_experts_per_tok"] == 1 + assert moe_nodes[0].attrs["normalize_routing"] is True + assert len(moe_nodes[0].inputs) == 2 + 3 * 2 diff --git a/python/tests/test_transpile_importers.py b/python/tests/test_transpile_importers.py new file mode 100644 index 000000000..cc3f53de9 --- /dev/null +++ b/python/tests/test_transpile_importers.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.importers import ImportContext +from cactus.transpile.importers import import_get_attr + + +def test_import_get_attr_preserves_meta_tensor_constant() -> None: + ir = IRGraph(values={}, nodes={}, order=[], inputs=[], outputs=[]) + ctx = ImportContext() + node = SimpleNamespace(name="p_weight", op="get_attr") + value = torch.empty((3, 4), device="meta", dtype=torch.float16) + + import_get_attr( + ir, + node, + ctx, + value, + shape=(3, 4), + dtype="float16", + source_name="linear.weight", + ) + + constant = ir.constants["v_p_weight"] + assert isinstance(constant, torch.Tensor) + assert constant.is_meta + assert tuple(constant.shape) == (3, 4) + assert ir.values["v_p_weight"].meta["source_name"] == "linear.weight" diff --git a/python/tests/test_transpile_layer_keys.py b/python/tests/test_transpile_layer_keys.py new file mode 100644 index 000000000..1bd157ef0 --- /dev/null +++ b/python/tests/test_transpile_layer_keys.py @@ -0,0 +1,68 @@ +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.lower import _kv_cache_layer_key +from cactus.transpile.optimize_graph import normalize_cached_decoder_attention_hints + + +def _decoder_graph(component: str) -> IRGraph: + values: dict[str, IRValue] = {} + nodes: dict[str, IRNode] = {} + for index in range(2): + node_id = "n_scaled_dot_product_attention" + ("" if index == 0 else f"_{index}") + for name in ("query", "key", "value"): + values[f"{name}_{index}"] = IRValue(id=f"{name}_{index}", shape=(1, 8, 1, 64), dtype="fp16") + values[f"out_{index}"] = IRValue(id=f"out_{index}", shape=(1, 8, 1, 64), dtype="fp16", producer=node_id) + nodes[node_id] = IRNode( + id=node_id, + op="scaled_dot_product_attention", + inputs=[f"query_{index}", f"key_{index}", f"value_{index}"], + outputs=[f"out_{index}"], + ) + return IRGraph( + values=values, + nodes=nodes, + order=list(nodes), + inputs=[value_id for value_id in values if not value_id.startswith("out_")], + outputs=["out_1"], + constants={}, + meta={ + "component": component, + "use_internal_kv_cache": True, + "layer_types": ["full_attention"] * 2, + }, + ) + + +def test_kv_cache_layer_key_keeps_zero_layer_index() -> None: + node = IRNode( + id="n_scaled_dot_product_attention", + op="scaled_dot_product_attention", + inputs=["query_0", "key_0", "value_0"], + outputs=["out_0"], + ) + assert _kv_cache_layer_key(node) == "n_scaled_dot_product_attention" + node.meta["layer_index"] = 5 + assert _kv_cache_layer_key(node) == "5" + node.meta["attention_layer_index"] = 0 + assert _kv_cache_layer_key(node) == "0" + + +def test_media_step_receives_attention_layer_indices() -> None: + graph = _decoder_graph("decoder_media_step") + + changed = normalize_cached_decoder_attention_hints(graph) + + assert changed is True + assert graph.nodes["n_scaled_dot_product_attention"].meta["attention_layer_index"] == 0 + assert graph.nodes["n_scaled_dot_product_attention_1"].meta["attention_layer_index"] == 1 + + +def test_step_and_media_step_layer_keys_match() -> None: + keys = [] + for component in ("decoder_step", "decoder_media_step"): + graph = _decoder_graph(component) + normalize_cached_decoder_attention_hints(graph) + keys.append([_kv_cache_layer_key(graph.nodes[node_id]) for node_id in graph.order]) + + assert keys[0] == keys[1] == ["0", "1"] diff --git a/python/tests/test_transpile_weight_compat.py b/python/tests/test_transpile_weight_compat.py new file mode 100644 index 000000000..89235b988 --- /dev/null +++ b/python/tests/test_transpile_weight_compat.py @@ -0,0 +1,386 @@ +from __future__ import annotations + +from pathlib import Path +import json + +import numpy as np + +from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header +from cactus.transpile.weight_binding import WeightBinding +from cactus.transpile.weight_binding import resolve_weight_binding +from cactus.transpile.weight_compat import _open_cactus_tensor_file +from cactus.transpile.weight_compat import ensure_binding_compatible +from cactus.transpile.weight_compat import ensure_embedding_binding_compatible + + +def _write_grouped_int8_embedding(path: Path, rows: int, cols: int) -> None: + rng = np.random.default_rng(1234) + tensor = rng.standard_normal((rows, cols), dtype=np.float32) + save_tensor_with_header(tensor, path, precision="INT8", model_type="generic") + + +def test_token_embedding_binding_upgrades_to_cq4(tmp_path: Path) -> None: + source = tmp_path / "token_embeddings.weights" + _write_grouped_int8_embedding(source, rows=8, cols=128) + + legacy_fp16 = tmp_path / "token_embeddings.fp16.weights" + legacy_fp16.write_bytes(b"legacy") + + binding = WeightBinding( + path=str(source), + kind="embedding", + source_name="model.embed_tokens.weight", + ) + compat = ensure_embedding_binding_compatible(binding) + + assert compat.path.endswith(".cq4.weights") + assert compat.kind == "embedding" + assert not legacy_fp16.exists() + + opened = _open_cactus_tensor_file(compat.path) + assert opened.precision == 6 + assert opened.shape == (8, 128) + assert opened.scales is not None + + +def test_per_layer_embedding_binding_upgrades_to_cq2(tmp_path: Path) -> None: + source = tmp_path / "embed_tokens_per_layer.weights" + _write_grouped_int8_embedding(source, rows=8, cols=256) + + binding = WeightBinding( + path=str(source), + kind="embedding", + source_name="model.language_model.embed_tokens_per_layer.weight", + ) + compat = ensure_embedding_binding_compatible(binding) + + assert compat.path.endswith(".cq2.weights") + + opened = _open_cactus_tensor_file(compat.path) + assert opened.precision == 4 + assert opened.shape == (8, 256) + assert opened.scales is not None + + +def test_gemma4_per_layer_projection_binding_upgrades_legacy_int4(tmp_path: Path) -> None: + rng = np.random.default_rng(1234) + source = tmp_path / "per_layer_model_proj.weights" + tensor = rng.standard_normal((8, 128), dtype=np.float32) + save_tensor_with_header(tensor, source, precision="INT4", model_type="generic") + + binding = WeightBinding( + path=str(source), + kind="weight", + source_name="module.model.model.language_model.per_layer_model_projection.weight", + ) + compat = ensure_binding_compatible(binding, source_tensor=tensor) + + assert compat.path.endswith(".cq4.weights") + opened = _open_cactus_tensor_file(compat.path) + assert opened.precision == 6 + assert opened.shape == (8, 128) + assert opened.scales is not None + + +def test_decoder_binding_prefers_existing_cq4_companion(tmp_path: Path) -> None: + rng = np.random.default_rng(4321) + source = tmp_path / "layer_0_attn_q.weights" + tensor = rng.standard_normal((8, 128), dtype=np.float32) + save_tensor_with_header(tensor, source, precision="INT4", model_type="generic") + + companion = tmp_path / "layer_0_attn_q.cq4.weights" + save_tensor_with_header(tensor, companion, precision="FP16", model_type="generic") + + binding = WeightBinding( + path=str(source), + kind="weight", + source_name="module.backbone.layers.0.self_attn.q_proj.weight", + ) + compat = ensure_binding_compatible(binding, source_tensor=0) + + assert compat.path == str(companion) + + +def test_resolve_weight_binding_does_not_guess_legacy_filenames(tmp_path: Path) -> None: + (tmp_path / "token_embeddings.weights").write_bytes(b"") + (tmp_path / "layer_0_attn_q.weights").write_bytes(b"") + (tmp_path / "embed_tokens_per_layer.weights").write_bytes(b"") + (tmp_path / "per_layer_model_proj.weights").write_bytes(b"") + + assert ( + resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="model.embed_tokens.weight", + ) + is None + ) + assert ( + resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="model.layers.0.self_attn.q_proj.weight", + ) + is None + ) + assert ( + resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="model.language_model.embed_tokens_per_layer.weight", + ) + is None + ) + assert ( + resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="model.language_model.per_layer_model_projection.weight", + ) + is None + ) + + +def test_resolve_weight_binding_prefers_exact_manifest_without_name_guessing(tmp_path: Path) -> None: + (tmp_path / "exact.weights").write_bytes(b"") + (tmp_path / "token_embeddings.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "v_adapter_named_anything": { + "filename": "exact.weights", + "kind": "weight", + } + } + ) + ) + + exact = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="v_adapter_named_anything", + ) + assert exact is not None + assert exact.path.endswith("exact.weights") + + guessed = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="model.embed_tokens.weight", + ) + assert guessed is None + + +def test_resolve_weight_binding_reads_convert_manifest_aliases(tmp_path: Path) -> None: + (tmp_path / "decoder_attn_q.cq4.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "weights": [ + { + "source_name": "model.layers.0.self_attn.q_proj.weight", + "hf_name": "model.layers.0.self_attn.q_proj.weight", + "adapter_name": "backbone.layers.0.self_attn.q_proj.weight", + "source_names": ["adapter.backbone.layers.0.self_attn.q_proj.weight"], + "output_name": "decoder_attn_q.cq4.weights", + "component": "language", + } + ] + } + ) + ) + + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="adapter.backbone.layers.0.self_attn.q_proj.weight", + ) + + assert binding is not None + assert binding.path.endswith("decoder_attn_q.cq4.weights") + + +def test_resolve_weight_binding_expands_manifest_wrapper_prefixes(tmp_path: Path) -> None: + (tmp_path / "layer_0_q.cq4.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "model.layers.0.self_attn.q_proj.weight": { + "filename": "layer_0_q.cq4.weights", + "kind": "weight", + } + } + ) + ) + + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="adapter.model.model.layers.0.self_attn.q_proj.weight", + ) + + assert binding is not None + assert binding.path.endswith("layer_0_q.cq4.weights") + + +def test_resolve_weight_binding_expands_component_local_manifest_aliases(tmp_path: Path) -> None: + (tmp_path / "encoder_layer_0_q.cq4.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "encoder.layers.0.self_attn.q_proj.weight": { + "filename": "encoder_layer_0_q.cq4.weights", + "kind": "weight", + } + } + ) + ) + + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="module.layers.0.self_attn.q_proj.weight", + ) + + assert binding is not None + assert binding.path.endswith("encoder_layer_0_q.cq4.weights") + + +def test_resolve_weight_binding_expands_language_model_backbone_aliases(tmp_path: Path) -> None: + (tmp_path / "layer_0_scalar.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "model.language_model.layers.0.layer_scalar": { + "filename": "layer_0_scalar.weights", + "kind": "weight", + } + } + ) + ) + + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="module.backbone.layers.0.layer_scalar", + ) + + assert binding is not None + assert binding.path.endswith("layer_0_scalar.weights") + + +def test_resolve_weight_binding_expands_lfm2_moe_runtime_expert_aliases(tmp_path: Path) -> None: + (tmp_path / "layer_2_moe_expert_0_w1.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "weights": [ + { + "source_name": "model.layers.2.feed_forward.experts.0.w1.weight", + "output_name": "layer_2_moe_expert_0_w1.weights", + "component": "language", + } + ] + } + ) + ) + + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="module.backbone.layers.2.feed_forward.w1_weights.0", + ) + + assert binding is not None + assert binding.path.endswith("layer_2_moe_expert_0_w1.weights") + + +def test_resolve_weight_binding_maps_module_backbone_to_model_manifest(tmp_path: Path) -> None: + (tmp_path / "layer_0_q.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "weights": [ + { + "source_name": "model.layers.0.self_attn.q_proj.weight", + "hf_name": "model.layers.0.self_attn.q_proj.weight", + "adapter_name": "model.layers.0.self_attn.q_proj.weight", + "output_name": "layer_0_q.weights", + "component": "language", + } + ] + } + ) + ) + + for source_name in ( + "module.backbone.layers.0.self_attn.q_proj.weight", + "adapter.backbone.layers.0.self_attn.q_proj.weight", + ): + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name=source_name, + ) + + assert binding is not None, source_name + assert binding.path.endswith("layer_0_q.weights"), source_name + + +def test_resolve_weight_binding_maps_tied_lm_head_to_token_embedding_manifest(tmp_path: Path) -> None: + (tmp_path / "token_embeddings.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "weights": [ + { + "source_name": "model.language_model.embed_tokens.weight", + "hf_name": "model.language_model.embed_tokens.weight", + "adapter_name": "model.language_model.embed_tokens.weight", + "source_names": ["model.language_model.embed_tokens.weight"], + "output_name": "token_embeddings.weights", + "component": "embedding", + } + ] + } + ) + ) + + for source_name in ("lm_head.weight", "model.lm_head.weight"): + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name=source_name, + ) + + assert binding is not None + assert binding.path.endswith("token_embeddings.weights") + assert binding.kind == "embedding" + + +def test_resolve_weight_binding_expands_nested_model_tower_aliases(tmp_path: Path) -> None: + (tmp_path / "vision_q.weights").write_bytes(b"") + (tmp_path / "weights_manifest.json").write_text( + json.dumps( + { + "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight": { + "filename": "vision_q.weights", + "kind": "weight", + } + } + ) + ) + + binding = resolve_weight_binding( + weights_dir=str(tmp_path), + source_name="module.model.model.vision_tower.encoder.layers.0.self_attn.q_proj.weight", + ) + + assert binding is not None + assert binding.path.endswith("vision_q.weights") + + +def test_rank5_materialized_constant_embeds_in_cactus_graph(tmp_path: Path) -> None: + from cactus.transpile.runtime_compat import Graph + from cactus.transpile.runtime_compat import Tensor + + tensor = np.zeros((1, 2, 3, 4, 5), dtype=np.float16) + graph_path = tmp_path / "rank5.cactus" + + graph = Graph() + constant = graph.input(tensor.shape, Graph.FP16) + graph.set_input(constant, tensor) + graph.mark_embedded_input(constant) + graph.save(graph_path) + + loaded = Graph.load(graph_path) + loaded_constant = Tensor(loaded, constant.id, constant.shape, constant.dtype) + np.testing.assert_array_equal(loaded_constant.numpy(), tensor) + assert not list(tmp_path.glob("*.weights")) diff --git a/python/tests/transpile/__init__.py b/python/tests/transpile/__init__.py new file mode 100644 index 000000000..c0dd68e2d --- /dev/null +++ b/python/tests/transpile/__init__.py @@ -0,0 +1 @@ +"""Transpiler regression tests and developer tools.""" diff --git a/python/tests/transpile/tools/__init__.py b/python/tests/transpile/tools/__init__.py new file mode 100644 index 000000000..b58eb5601 --- /dev/null +++ b/python/tests/transpile/tools/__init__.py @@ -0,0 +1 @@ +"""Developer-only transpiler comparison tools.""" diff --git a/python/tests/transpile/tools/compare_gemma4_multimodal.py b/python/tests/transpile/tools/compare_gemma4_multimodal.py new file mode 100644 index 000000000..14bee44e5 --- /dev/null +++ b/python/tests/transpile/tools/compare_gemma4_multimodal.py @@ -0,0 +1,932 @@ +from __future__ import annotations + +import argparse +import copy +import gc +import json +import os +import sys +import time +from dataclasses import asdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +TOOLS_DIR = Path(__file__).resolve().parent +PYTHON_ROOT = TOOLS_DIR.parents[2] +PROJECT_ROOT = PYTHON_ROOT.parent + +sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.bindings.cactus import cactus_complete +from cactus.bindings.cactus import cactus_destroy +from cactus.bindings.cactus import cactus_init +from cactus.bindings.cactus import cactus_reset +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.model_adapters import canonicalize_model_interface +from cactus.transpile.optimize_graph import FusionConfig +from cactus.transpile.optimize_graph import optimize_graph +from cactus.transpile.hf_model import TranspileWrapper +from cactus.transpile.hf_model import _graph_to_dict +from cactus.transpile.hf_model import _load_transformers_bundle +from cactus.transpile.hf_model import _lower_preoptimized_ir +from cactus.transpile.hf_model import _parse_dtype +from cactus.transpile.hf_model import _prepare_gemma4_multimodal_inputs +from cactus.transpile.hf_model import _serialize_json_compatible +from cactus.transpile.hf_model import _validate_weights_dir +from cactus.transpile.hf_model import _write_json + + +_DEFAULT_STOP_SEQUENCES = ("", "", "", "<|im_end|>") + + +@dataclass +class NumericSummary: + mean: float + min: float + max: float + median: float + + +@dataclass +class HandwrittenRun: + wall_ms: float + total_time_ms: float + time_to_first_token_ms: float + prefill_tps: float + decode_tps: float + decode_tokens: int + total_tokens: int + response: str + response_json: dict[str, Any] + + +@dataclass +class TranspiledRun: + preprocess_ms: float + execute_ms: float + total_ms: float + time_to_first_token_ms: float + prefill_tps: float + decode_tps: float + decode_tokens: int + total_tokens: int + response: str + stop_reason: str + generated_token_ids: list[int] + input_shapes: dict[str, list[int]] + + +def _numeric_summary(values: list[float]) -> NumericSummary: + arr = np.asarray(values, dtype=np.float64) + return NumericSummary( + mean=float(np.mean(arr)), + min=float(np.min(arr)), + max=float(np.max(arr)), + median=float(np.median(arr)), + ) + + +def _roundtrip_jsonable(value: Any) -> Any: + return json.loads(json.dumps(_serialize_json_compatible(value))) + + +def _resolve_default_weights_dir(model_id: str) -> Path | None: + env_path = os.environ.get("CACTUS_TEST_GEMMA4_MODEL") + if env_path: + candidate = Path(env_path).resolve() + if (candidate / "config.txt").exists(): + return candidate + + candidates = ( + PROJECT_ROOT / "weights" / "gemma4_int4", + PROJECT_ROOT / "weights" / model_id.split("/")[-1].lower(), + PROJECT_ROOT / "weights" / model_id.split("/")[-1].lower().replace("_", "-"), + ) + for candidate in candidates: + if (candidate / "config.txt").exists(): + return candidate + return None + + +def _require_processor_tokenizer(processor_or_tokenizer: object) -> object: + tokenizer_like = getattr(processor_or_tokenizer, "tokenizer", processor_or_tokenizer) + if not hasattr(tokenizer_like, "decode") or not hasattr(tokenizer_like, "encode"): + raise RuntimeError( + "Gemma4 compare script requires a tokenizer-like object with encode/decode support." + ) + return tokenizer_like + + +def _build_messages_payload( + *, + prompt: str, + image_files: tuple[str, ...], + audio_file: str | None, + system_prompt: str = "", +) -> list[dict[str, object]]: + messages: list[dict[str, object]] = [] + normalized_system = system_prompt.strip() + if normalized_system: + messages.append({"role": "system", "content": normalized_system}) + + user_message: dict[str, object] = { + "role": "user", + "content": prompt, + } + if image_files: + user_message["images"] = [str(Path(path).resolve()) for path in image_files] + if audio_file: + user_message["audio"] = [str(Path(audio_file).resolve())] + messages.append(user_message) + return messages + + +def _encode_stop_sequences(tokenizer: object, stop_sequences: tuple[str, ...]) -> list[list[int]]: + encoded: list[list[int]] = [] + encode = getattr(tokenizer, "encode", None) + if not callable(encode): + return encoded + for stop_sequence in stop_sequences: + try: + token_ids = list(encode(stop_sequence, add_special_tokens=False)) + except TypeError: + token_ids = list(encode(stop_sequence)) + if token_ids: + encoded.append([int(token_id) for token_id in token_ids]) + return encoded + + +def _has_suffix(token_ids: list[int], suffix: list[int]) -> bool: + if not suffix or len(token_ids) < len(suffix): + return False + return token_ids[-len(suffix) :] == suffix + + +def _trim_stop_suffix(token_ids: list[int], stop_sequences: list[list[int]]) -> bool: + for stop_sequence in stop_sequences: + if _has_suffix(token_ids, stop_sequence): + del token_ids[-len(stop_sequence) :] + return True + return False + + +def _decode_generated_text(tokenizer: object, token_ids: list[int], *, skip_special_tokens: bool) -> str: + decode = getattr(tokenizer, "decode", None) + if not callable(decode): + raise RuntimeError(f"tokenizer does not expose decode(): {type(tokenizer).__name__}") + try: + return str(decode(token_ids, skip_special_tokens=skip_special_tokens)) + except TypeError: + return str(decode(token_ids)) + + +def _summarize_handwritten_runs( + runs: list[HandwrittenRun], + *, + weights_dir: Path, +) -> dict[str, Any]: + if not runs: + raise ValueError("expected at least one handwritten run") + return { + "weights_dir": str(weights_dir), + "latest_response": runs[-1].response, + "total_time_ms": asdict(_numeric_summary([run.total_time_ms for run in runs])), + "time_to_first_token_ms": asdict(_numeric_summary([run.time_to_first_token_ms for run in runs])), + "prefill_tps": asdict(_numeric_summary([run.prefill_tps for run in runs])), + "decode_tps": asdict(_numeric_summary([run.decode_tps for run in runs])), + "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), + "total_tokens": asdict(_numeric_summary([float(run.total_tokens) for run in runs])), + "wall_ms": asdict(_numeric_summary([run.wall_ms for run in runs])), + "runs": [_roundtrip_jsonable(asdict(run)) for run in runs], + } + + +def _summarize_transpiled_runs( + runs: list[TranspiledRun], + *, + compile_time_ms: float, + model_source: str, + weight_bindings: int, + prepared_input_shapes: dict[str, list[int]], +) -> dict[str, Any]: + if not runs: + raise ValueError("expected at least one transpiled run") + return { + "compile_time_ms": compile_time_ms, + "model_source": model_source, + "weight_bindings": weight_bindings, + "prepared_input_shapes": prepared_input_shapes, + "latest_response": runs[-1].response, + "preprocess_ms": asdict(_numeric_summary([run.preprocess_ms for run in runs])), + "execute_ms": asdict(_numeric_summary([run.execute_ms for run in runs])), + "total_ms": asdict(_numeric_summary([run.total_ms for run in runs])), + "time_to_first_token_ms": asdict(_numeric_summary([run.time_to_first_token_ms for run in runs])), + "prefill_tps": asdict(_numeric_summary([run.prefill_tps for run in runs])), + "decode_tps": asdict(_numeric_summary([run.decode_tps for run in runs])), + "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), + "total_tokens": asdict(_numeric_summary([float(run.total_tokens) for run in runs])), + "runs": [_roundtrip_jsonable(asdict(run)) for run in runs], + } + + +def _prepare_repeat_inputs( + processor: object, + *, + tokenizer: object, + prompt: str, + image_files: tuple[str, ...], + audio_file: str | None, + torch_dtype: torch.dtype, + system_prompt: str, + enable_thinking_if_supported: bool, + max_new_tokens: int, +) -> tuple[object, int, float]: + start = time.perf_counter() + prepared = _prepare_gemma4_multimodal_inputs( + processor, + prompt=prompt, + image_files=image_files, + audio_file=audio_file, + torch_dtype=torch_dtype, + system_prompt=system_prompt, + enable_thinking_if_supported=enable_thinking_if_supported, + use_gemma4_chat_template=True, + ) + pad_token_id = getattr(tokenizer, "pad_token_id", None) + if pad_token_id is None: + pad_token_id = 0 + prompt_tokens = 0 + for name, tensor in zip(prepared.names, prepared.tensors): + if name == "input_ids": + prompt_tokens = int(tensor.shape[1]) + break + if prompt_tokens <= 0: + raise RuntimeError("Gemma4 multimodal prepared inputs did not include a non-empty input_ids tensor") + target_tokens = prompt_tokens + int(max_new_tokens) + + input_tensors: list[torch.Tensor] = [] + for name, tensor in zip(prepared.names, prepared.tensors): + if name in {"input_ids", "attention_mask", "token_type_ids"}: + padded_shape = list(tensor.shape) + padded_shape[1] = target_tokens + if name == "input_ids": + padded = torch.full( + padded_shape, + int(pad_token_id), + dtype=tensor.dtype, + device=tensor.device, + ) + else: + padded = torch.zeros( + padded_shape, + dtype=tensor.dtype, + device=tensor.device, + ) + padded[:, : tensor.shape[1]] = tensor + input_tensors.append(padded) + else: + input_tensors.append(tensor) + metadata = dict(prepared.metadata) + metadata["prompt_token_count"] = prompt_tokens + metadata["target_token_count"] = int(target_tokens) + metadata["input_shapes"] = { + name: list(tensor.shape) + for name, tensor in zip(prepared.names, input_tensors) + } + prepared = type(prepared)( + names=prepared.names, + tensors=tuple(input_tensors), + metadata=metadata, + ) + preprocess_ms = (time.perf_counter() - start) * 1000.0 + return prepared, prompt_tokens, preprocess_ms + + +def _run_handwritten_once( + *, + model_handle: Any, + messages_json: str, + options_json: str, +) -> HandwrittenRun: + cactus_reset(model_handle) + streamed_token_ids: list[int] = [] + + def _collect(_text: str, token_id: int) -> None: + streamed_token_ids.append(int(token_id)) + + start = time.perf_counter() + response_json_text = cactus_complete( + model_handle, + messages_json, + options_json, + "[]", + _collect, + pcm_data=None, + ) + wall_ms = (time.perf_counter() - start) * 1000.0 + response_payload = json.loads(response_json_text) + return HandwrittenRun( + wall_ms=wall_ms, + total_time_ms=float(response_payload.get("total_time_ms", wall_ms)), + time_to_first_token_ms=float(response_payload.get("time_to_first_token_ms", 0.0)), + prefill_tps=float(response_payload.get("prefill_tps", 0.0)), + decode_tps=float(response_payload.get("decode_tps", 0.0)), + decode_tokens=len(streamed_token_ids), + total_tokens=int(response_payload.get("total_tokens", len(streamed_token_ids))), + response=str(response_payload.get("response", "")), + response_json=response_payload, + ) + + +def _run_transpiled_generation_once( + *, + tg, + processor: object, + tokenizer: object, + prompt: str, + image_files: tuple[str, ...], + audio_file: str | None, + torch_dtype: torch.dtype, + system_prompt: str, + enable_thinking_if_supported: bool, + max_new_tokens: int, + stop_sequences: tuple[str, ...], + expected_input_shapes: dict[str, list[int]] | None = None, + prepared_override: object | None = None, + progress_label: str = "transpiled", + progress_every_tokens: int = 8, +) -> TranspiledRun: + if prepared_override is not None: + prepared = type(prepared_override)( + names=prepared_override.names, + tensors=tuple(tensor.clone() for tensor in prepared_override.tensors), + metadata=dict(prepared_override.metadata), + ) + prompt_tokens = int(prepared.metadata.get("prompt_token_count", 0)) + preprocess_ms = 0.0 + else: + prepared, prompt_tokens, preprocess_ms = _prepare_repeat_inputs( + processor, + tokenizer=tokenizer, + prompt=prompt, + image_files=image_files, + audio_file=audio_file, + torch_dtype=torch_dtype, + system_prompt=system_prompt, + enable_thinking_if_supported=enable_thinking_if_supported, + max_new_tokens=max_new_tokens, + ) + inputs_by_name = { + name: tensor.detach().cpu().numpy() + for name, tensor in zip(prepared.names, prepared.tensors) + } + if expected_input_shapes is not None: + for name, tensor in zip(prepared.names, prepared.tensors): + expected_shape = expected_input_shapes.get(name) + if expected_shape is None: + continue + actual_shape = list(tensor.shape) + if actual_shape != list(expected_shape): + raise RuntimeError( + f"runtime prepared input shape mismatch for {name}: " + f"compiled with {list(expected_shape)}, got {actual_shape}" + ) + if "input_ids" not in inputs_by_name: + raise RuntimeError("transpiled Gemma4 generation requires input_ids") + + target_tokens = int(prepared.metadata["target_token_count"]) + + eos_token_id = getattr(tokenizer, "eos_token_id", None) + + padded_inputs: dict[str, np.ndarray] = dict(inputs_by_name) + generated_ids: list[int] = [] + encoded_stop_sequences = _encode_stop_sequences(tokenizer, stop_sequences) + current_length = prompt_tokens + first_token_ms = 0.0 + execute_start = time.perf_counter() + stop_reason = "max_new_tokens" + progress_interval = max(1, int(progress_every_tokens)) + + for step_index in range(max_new_tokens): + runtime_inputs = [padded_inputs[name] for name in prepared.names] + if step_index == 0: + print("transpiled_step0_set_inputs_begin=true", flush=True) + tg.set_inputs(runtime_inputs) + if step_index == 0: + print("transpiled_step0_set_inputs_done=true", flush=True) + print("transpiled_step0_execute_begin=true", flush=True) + logits = tg.execute()[0].numpy() + if step_index == 0: + print("transpiled_step0_execute_done=true", flush=True) + if logits.ndim == 3: + token_position = current_length - 1 + if logits.shape[1] == 1: + token_position = 0 + next_token_id = int(np.argmax(logits[0, token_position])) + elif logits.ndim == 2: + next_token_id = int(np.argmax(logits[0])) + else: + raise RuntimeError(f"unexpected transpiled logits rank: {logits.ndim}") + generated_ids.append(next_token_id) + + if step_index == 0: + first_token_ms = preprocess_ms + (time.perf_counter() - execute_start) * 1000.0 + + if eos_token_id is not None and next_token_id == int(eos_token_id): + stop_reason = "eos_token" + break + if _trim_stop_suffix(generated_ids, encoded_stop_sequences): + stop_reason = "stop_sequence" + break + + if current_length >= target_tokens: + stop_reason = "context_limit" + break + + padded_inputs["input_ids"][0, current_length] = next_token_id + padded_inputs["attention_mask"][0, current_length] = 1 + padded_inputs["token_type_ids"][0, current_length] = 0 + if step_index == 0: + print("transpiled_step0_state_update_done=true", flush=True) + current_length += 1 + generated_count = len(generated_ids) + if generated_count % progress_interval == 0 or generated_count == max_new_tokens: + elapsed_ms = (time.perf_counter() - execute_start) * 1000.0 + print( + f"{progress_label}_decode_progress={generated_count}/{max_new_tokens} " + f"elapsed_ms={elapsed_ms:.3f}", + flush=True, + ) + + execute_ms = (time.perf_counter() - execute_start) * 1000.0 + response = _decode_generated_text(tokenizer, generated_ids, skip_special_tokens=True).strip() + if not response: + response = _decode_generated_text(tokenizer, generated_ids, skip_special_tokens=False).strip() + decode_time_ms = max(0.0, execute_ms + preprocess_ms - first_token_ms) + decode_tps = ( + ((len(generated_ids) - 1) * 1000.0) / decode_time_ms + if len(generated_ids) > 1 and decode_time_ms > 0.0 + else 0.0 + ) + prefill_tps = (prompt_tokens * 1000.0) / first_token_ms if first_token_ms > 0.0 else 0.0 + print( + f"{progress_label}_decode_done=true tokens={len(generated_ids)} " + f"stop_reason={stop_reason} execute_ms={execute_ms:.3f}", + flush=True, + ) + + return TranspiledRun( + preprocess_ms=preprocess_ms, + execute_ms=execute_ms, + total_ms=preprocess_ms + execute_ms, + time_to_first_token_ms=first_token_ms, + prefill_tps=prefill_tps, + decode_tps=decode_tps, + decode_tokens=len(generated_ids), + total_tokens=prompt_tokens + len(generated_ids), + response=response, + stop_reason=stop_reason, + generated_token_ids=generated_ids, + input_shapes={ + name: list(tensor.shape) + for name, tensor in zip(prepared.names, prepared.tensors) + }, + ) + + +def _print_summary(summary: dict[str, Any]) -> None: + print(f"model_id={summary['model_id']}") + print(f"prompt={summary['prompt']!r}") + if summary.get("image_files"): + print(f"image_files={summary['image_files']}") + if summary.get("audio_file"): + print(f"audio_file={summary['audio_file']}") + + handwritten = summary.get("handwritten") + if handwritten is not None: + hw_total = handwritten["total_time_ms"]["mean"] + hw_ttft = handwritten["time_to_first_token_ms"]["mean"] + hw_tokens = handwritten["decode_tokens"]["mean"] + hw_tps = hw_tokens * 1000.0 / max(hw_total, 1e-9) + print("handwritten:") + print(f" total_ms mean={hw_total:.3f}") + print(f" time_to_first_token_ms mean={hw_ttft:.3f}") + print(f" decode_tokens mean={hw_tokens:.3f}") + print(f" output_tokens_per_sec={hw_tps:.3f}") + print(f" latest_response={handwritten['latest_response']!r}") + + transpiled = summary.get("transpiled") + if transpiled is not None: + tr_total = transpiled["total_ms"]["mean"] + tr_ttft = transpiled["time_to_first_token_ms"]["mean"] + tr_tokens = transpiled["decode_tokens"]["mean"] + tr_tps = tr_tokens * 1000.0 / max(tr_total, 1e-9) + print("transpiled:") + print(f" compile_time_ms={transpiled['compile_time_ms']:.3f}") + print(f" total_ms mean={tr_total:.3f}") + print(f" time_to_first_token_ms mean={tr_ttft:.3f}") + print(f" decode_tokens mean={tr_tokens:.3f}") + print(f" output_tokens_per_sec={tr_tps:.3f}") + print(f" weight_bindings={transpiled['weight_bindings']}") + print(f" latest_response={transpiled['latest_response']!r}") + + comparison = summary.get("comparison") + if comparison is not None: + print("comparison:") + print( + " end_to_end_speedup_vs_handwritten=" + f"{comparison['end_to_end_speedup_vs_handwritten']:.6f}" + ) + print(f" responses_match={comparison['responses_match']}") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=( + "Compare handwritten Cactus Gemma4 multimodal generation against the " + "transpiled Gemma4 multimodal graph on the same prompt, image, and audio." + ) + ) + parser.add_argument("--model-id", default="google/gemma-4-E2B-it") + parser.add_argument("--prompt", required=True) + parser.add_argument("--system-prompt", default="") + parser.add_argument("--image-file", action="append", default=[]) + parser.add_argument("--audio-file", default="") + parser.add_argument( + "--weights-dir", + default="", + help="Converted Cactus weights directory. Required for handwritten comparison and reused for transpiled mmap bindings.", + ) + parser.add_argument("--artifact-dir", default="", help="Optional directory for summary and transpile artifacts.") + parser.add_argument("--graph-filename", default="graph.cactus") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--repeats", type=int, default=1) + parser.add_argument("--max-new-tokens", type=int, default=96) + parser.add_argument( + "--transpiled-progress-every-tokens", + type=int, + default=1, + help=( + "Emit transpiled decode progress every N generated tokens. " + "Use 1 for per-token visibility on long CPU-bound Gemma4 runs." + ), + ) + parser.add_argument("--torch-dtype", default="float16") + parser.add_argument("--token", default=os.environ.get("HF_TOKEN")) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--skip-handwritten", action="store_true") + parser.add_argument("--skip-transpiled", action="store_true") + parser.add_argument("--enable-thinking", action="store_true") + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=1) + parser.add_argument("--stop-sequence", action="append", default=[]) + parser.add_argument("--no-fuse-gated-deltanet", action="store_true") + parser.add_argument("--no-fuse-rms-norm", action="store_true") + parser.add_argument("--no-fuse-rope", action="store_true") + parser.add_argument("--no-fuse-attention", action="store_true") + parser.add_argument("--no-fuse-attention-block", action="store_true") + parser.add_argument("--no-fuse-add-clipped", action="store_true") + args = parser.parse_args() + + if args.skip_handwritten and args.skip_transpiled: + raise RuntimeError("cannot skip both handwritten and transpiled runs") + + image_files = tuple(str(Path(path).resolve()) for path in args.image_file if str(path).strip()) + audio_file = str(Path(args.audio_file).resolve()) if args.audio_file.strip() else None + if audio_file and not Path(audio_file).exists(): + raise FileNotFoundError(f"audio file does not exist: {audio_file}") + for image_file in image_files: + if not Path(image_file).exists(): + raise FileNotFoundError(f"image file does not exist: {image_file}") + + torch_dtype = _parse_dtype(args.torch_dtype) + artifact_dir = Path(args.artifact_dir).resolve() if args.artifact_dir else None + if artifact_dir is not None: + artifact_dir.mkdir(parents=True, exist_ok=True) + + weights_dir_path = None + if args.weights_dir.strip(): + weights_dir_path = _validate_weights_dir(args.weights_dir.strip(), model_id=args.model_id) + else: + weights_dir_path = _resolve_default_weights_dir(args.model_id) + + if not args.skip_handwritten and weights_dir_path is None: + raise RuntimeError( + "handwritten Gemma4 comparison requires a converted weights directory.\n" + "\n" + "Create one with something like:\n" + f" cactus convert {args.model_id} /path/to/gemma4_weights\n" + "\n" + "Then rerun with --weights-dir /path/to/gemma4_weights, or rerun with --skip-handwritten." + ) + + stop_sequences = tuple(args.stop_sequence) if args.stop_sequence else _DEFAULT_STOP_SEQUENCES + fusion_config = FusionConfig( + enable_gated_deltanet=not args.no_fuse_gated_deltanet, + enable_rms_norm=not args.no_fuse_rms_norm, + enable_rope=not args.no_fuse_rope, + enable_attention=not args.no_fuse_attention, + enable_attention_block=not args.no_fuse_attention_block, + enable_add_clipped=not args.no_fuse_add_clipped, + ) + + summary: dict[str, Any] = { + "model_id": args.model_id, + "prompt": args.prompt, + "system_prompt": args.system_prompt, + "image_files": list(image_files), + "audio_file": audio_file, + "weights_dir": str(weights_dir_path) if weights_dir_path is not None else None, + "warmup": args.warmup, + "repeats": args.repeats, + "max_new_tokens": args.max_new_tokens, + "transpiled_progress_every_tokens": int(args.transpiled_progress_every_tokens), + "torch_dtype": args.torch_dtype, + "trust_remote_code": bool(args.trust_remote_code), + "local_files_only": bool(args.local_files_only), + "enable_thinking": bool(args.enable_thinking), + "temperature": float(args.temperature), + "top_p": float(args.top_p), + "top_k": int(args.top_k), + "stop_sequences": list(stop_sequences), + } + + messages_payload = _build_messages_payload( + prompt=args.prompt, + image_files=image_files, + audio_file=audio_file, + system_prompt=args.system_prompt, + ) + messages_json = json.dumps(messages_payload) + options_json = json.dumps( + { + "temperature": float(args.temperature), + "top_p": float(args.top_p), + "top_k": int(args.top_k), + "max_tokens": int(args.max_new_tokens), + "stop_sequences": list(stop_sequences), + "enable_thinking_if_supported": bool(args.enable_thinking), + "auto_handoff": False, + "telemetry_enabled": False, + } + ) + + transpiled_runs: list[TranspiledRun] = [] + if not args.skip_transpiled: + model_source, processor_or_tokenizer, model, _ = _load_transformers_bundle( + model_id=args.model_id, + task="multimodal_causal_lm_logits", + torch_dtype=torch_dtype, + token=args.token, + trust_remote_code=args.trust_remote_code, + local_files_only=args.local_files_only, + ) + processor = processor_or_tokenizer + tokenizer = _require_processor_tokenizer(processor_or_tokenizer) + initial_prepared, _, initial_prepare_ms = _prepare_repeat_inputs( + processor, + tokenizer=tokenizer, + prompt=args.prompt, + image_files=image_files, + audio_file=audio_file, + torch_dtype=torch_dtype, + system_prompt=args.system_prompt, + enable_thinking_if_supported=args.enable_thinking, + max_new_tokens=args.max_new_tokens, + ) + canonical = canonicalize_model_interface( + model, + task="multimodal_causal_lm_logits", + input_names=initial_prepared.names, + weights_dir=str(weights_dir_path) if weights_dir_path is not None else None, + ) + if hasattr(canonical.module, "last_token_logits_only"): + canonical.module.last_token_logits_only = True + prime_static_features = getattr(canonical.module, "prime_static_multimodal_features", None) + if callable(prime_static_features): + prime_static_features(*initial_prepared.tensors) + wrapper = TranspileWrapper( + canonical.module, + weights_dir=str(weights_dir_path) if weights_dir_path is not None else None, + ).eval() + + compile_start = time.perf_counter() + print("transpile_capture_begin=true", flush=True) + prepare_cpu_float32_capture = getattr(canonical.module, "prepare_cpu_float32_capture", None) + restore_cpu_float32_capture = getattr(canonical.module, "restore_cpu_float32_capture", None) + if callable(prepare_cpu_float32_capture): + prepare_cpu_float32_capture() + try: + captured = capture_model(wrapper, initial_prepared.tensors) + finally: + if callable(restore_cpu_float32_capture): + restore_cpu_float32_capture() + print("transpile_capture_done=true", flush=True) + raw_ir_graph = copy.deepcopy(captured.ir_graph) + print("transpile_canonicalize_begin=true", flush=True) + canonicalize_exported_graph(captured.ir_graph) + print("transpile_canonicalize_done=true", flush=True) + print("transpile_optimize_begin=true", flush=True) + optimize_graph(captured.ir_graph, config=fusion_config) + print("transpile_optimize_done=true", flush=True) + optimized_ir_graph = copy.deepcopy(captured.ir_graph) + print("transpile_lower_begin=true", flush=True) + tg = _lower_preoptimized_ir(captured.ir_graph) + print("transpile_lower_done=true", flush=True) + compile_time_ms = (time.perf_counter() - compile_start) * 1000.0 + + op_counts: dict[str, int] = {} + for node_id in optimized_ir_graph.order: + op = optimized_ir_graph.nodes[node_id].op + op_counts[op] = op_counts.get(op, 0) + 1 + weight_bindings = sum( + 1 + for value in optimized_ir_graph.values.values() + if isinstance(value.meta, dict) and isinstance(value.meta.get("path"), str) + ) + + summary["transpile_compile"] = { + "model_source": model_source, + "raw_ir_nodes": len(raw_ir_graph.order), + "optimized_ir_nodes": len(optimized_ir_graph.order), + "weight_bindings": weight_bindings, + "compile_time_ms": compile_time_ms, + "initial_prepare_ms": initial_prepare_ms, + "ops": op_counts, + } + prompt_token_count = int(initial_prepared.metadata.get("prompt_token_count", 0)) + target_token_count = int(initial_prepared.metadata.get("target_token_count", 0)) + total_transpiled_runs = int(args.warmup) + int(args.repeats) + print( + "transpiled_decode_plan=" + f"prompt_tokens:{prompt_token_count} " + f"target_tokens:{target_token_count} " + f"max_new_tokens:{int(args.max_new_tokens)} " + f"runs:{total_transpiled_runs} " + f"progress_every_tokens:{int(args.transpiled_progress_every_tokens)}", + flush=True, + ) + + if weights_dir_path is not None and weight_bindings == 0: + raise RuntimeError( + f"No weight bindings were resolved from {weights_dir_path}\n" + "\n" + "The converted weights folder exists, but none of the captured constants matched weights_manifest.json.\n" + "Re-convert Gemma4 with the current converter before benchmarking." + ) + + if artifact_dir is not None: + _write_json( + artifact_dir / "raw_ir.json", + { + "model_id": args.model_id, + "model_source": model_source, + "task": "multimodal_causal_lm_logits", + "inputs": _serialize_json_compatible(initial_prepared.metadata), + "graph": _graph_to_dict(raw_ir_graph), + }, + ) + _write_json( + artifact_dir / "optimized_ir.json", + { + "model_id": args.model_id, + "model_source": model_source, + "task": "multimodal_causal_lm_logits", + "inputs": _serialize_json_compatible(initial_prepared.metadata), + "graph": _graph_to_dict(optimized_ir_graph), + }, + ) + print(f"saved_raw_ir={artifact_dir / 'raw_ir.json'}") + print(f"saved_optimized_ir={artifact_dir / 'optimized_ir.json'}") + if weight_bindings > 0: + skip_reason = ( + "skipped saved_graph because graph serialization does not yet preserve " + "mmap-bound weights/embeddings required by this transpiled Gemma4 graph" + ) + summary["transpile_compile"]["saved_graph_skipped_reason"] = skip_reason + print(f"note={skip_reason}") + else: + graph_path = artifact_dir / args.graph_filename + tg.graph.save(graph_path) + print(f"saved_graph={graph_path}") + + for _ in range(args.warmup): + warmup_index = _ + 1 + print(f"transpiled_warmup_begin={warmup_index}/{args.warmup}", flush=True) + _run_transpiled_generation_once( + tg=tg, + processor=processor, + tokenizer=tokenizer, + prompt=args.prompt, + image_files=image_files, + audio_file=audio_file, + torch_dtype=torch_dtype, + system_prompt=args.system_prompt, + enable_thinking_if_supported=args.enable_thinking, + max_new_tokens=args.max_new_tokens, + stop_sequences=stop_sequences, + expected_input_shapes=dict(initial_prepared.metadata.get("input_shapes", {})), + prepared_override=initial_prepared, + progress_label=f"transpiled_warmup{warmup_index}", + progress_every_tokens=args.transpiled_progress_every_tokens, + ) + for _ in range(args.repeats): + repeat_index = _ + 1 + print(f"transpiled_repeat_begin={repeat_index}/{args.repeats}", flush=True) + transpiled_runs.append( + _run_transpiled_generation_once( + tg=tg, + processor=processor, + tokenizer=tokenizer, + prompt=args.prompt, + image_files=image_files, + audio_file=audio_file, + torch_dtype=torch_dtype, + system_prompt=args.system_prompt, + enable_thinking_if_supported=args.enable_thinking, + max_new_tokens=args.max_new_tokens, + stop_sequences=stop_sequences, + expected_input_shapes=dict(initial_prepared.metadata.get("input_shapes", {})), + prepared_override=initial_prepared, + progress_label=f"transpiled_repeat{repeat_index}", + progress_every_tokens=args.transpiled_progress_every_tokens, + ) + ) + + summary["transpiled"] = _summarize_transpiled_runs( + transpiled_runs, + compile_time_ms=compile_time_ms, + model_source=model_source, + weight_bindings=weight_bindings, + prepared_input_shapes=dict(initial_prepared.metadata.get("input_shapes", {})), + ) + + # Keep the HF/export capture stack alive until the transpiled graph has + # completed execution. Gemma4 teardown before first execute has been + # unstable in practice, even though the lowered graph is already built. + del model + del wrapper + del canonical + del captured + gc.collect() + + handwritten_runs: list[HandwrittenRun] = [] + if not args.skip_handwritten and weights_dir_path is not None: + model_handle = cactus_init(str(weights_dir_path), None, False) + try: + for _ in range(args.warmup): + warmup_index = _ + 1 + print(f"handwritten_warmup_begin={warmup_index}/{args.warmup}", flush=True) + _run_handwritten_once( + model_handle=model_handle, + messages_json=messages_json, + options_json=options_json, + ) + for _ in range(args.repeats): + repeat_index = _ + 1 + print(f"handwritten_repeat_begin={repeat_index}/{args.repeats}", flush=True) + handwritten_runs.append( + _run_handwritten_once( + model_handle=model_handle, + messages_json=messages_json, + options_json=options_json, + ) + ) + finally: + cactus_destroy(model_handle) + summary["handwritten"] = _summarize_handwritten_runs( + handwritten_runs, + weights_dir=weights_dir_path, + ) + + if handwritten_runs and transpiled_runs: + handwritten_latest = handwritten_runs[-1].response.strip() + transpiled_latest = transpiled_runs[-1].response.strip() + summary["comparison"] = { + "end_to_end_speedup_vs_handwritten": ( + summary["handwritten"]["total_time_ms"]["mean"] + / max(summary["transpiled"]["total_ms"]["mean"], 1e-9) + ), + "responses_match": handwritten_latest == transpiled_latest, + "handwritten_latest_response": handwritten_latest, + "transpiled_latest_response": transpiled_latest, + } + + if artifact_dir is not None: + summary_path = artifact_dir / "summary.json" + summary_path.write_text(json.dumps(_roundtrip_jsonable(summary), indent=2, sort_keys=True) + "\n") + print(f"saved_summary={summary_path}") + + _print_summary(summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/tests/transpile/tools/compare_parakeet_performance.py b/python/tests/transpile/tools/compare_parakeet_performance.py new file mode 100644 index 000000000..eeaf06066 --- /dev/null +++ b/python/tests/transpile/tools/compare_parakeet_performance.py @@ -0,0 +1,792 @@ +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from collections import defaultdict +from dataclasses import asdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from scipy.io import wavfile + +TOOLS_DIR = Path(__file__).resolve().parent +PYTHON_ROOT = TOOLS_DIR.parents[2] +PROJECT_ROOT = PYTHON_ROOT.parent + +sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.bindings.cactus import cactus_destroy +from cactus.bindings.cactus import cactus_init +from cactus.bindings.cactus import cactus_transcribe +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.model_adapters import canonicalize_model_interface +from cactus.transpile.optimize_graph import FusionConfig +from cactus.transpile.optimize_graph import optimize_graph +from cactus.transpile.hf_model import TranspileWrapper +from cactus.transpile.hf_model import _ctc_greedy_decode_token_ids +from cactus.transpile.hf_model import _decode_token_ids +from cactus.transpile.hf_model import _infer_task_from_config +from cactus.transpile.hf_model import _load_optional_json +from cactus.transpile.hf_model import _load_optional_tokenizer +from cactus.transpile.hf_model import _load_transformers_bundle +from cactus.transpile.hf_model import _lower_preoptimized_ir +from cactus.transpile.hf_model import _parse_dtype +from cactus.transpile.hf_model import _prepare_audio_inputs +from cactus.transpile.hf_model import _validate_weights_dir + + +@dataclass +class NumericSummary: + mean: float + min: float + max: float + median: float + + +@dataclass +class ProfileEntry: + op: str + time_ms: float + shape: str + raw_line: str + + +@dataclass +class ProfileSection: + entries: list[ProfileEntry] + total_ms: float + + +@dataclass +class HandwrittenRun: + wall_ms: float + total_time_ms: float + time_to_first_token_ms: float + prefill_tps: float + decode_tps: float + decode_tokens: int + prefill_tokens: int + total_tokens: int + confidence: float + transcript: str + response_json: dict[str, Any] + + +@dataclass +class TranspiledRun: + preprocess_ms: float + execute_ms: float + decode_ms: float + total_ms: float + graph_only_ms: float + output_frames: int + output_vocab: int + decode_tokens: int + transcript: str | None + input_shapes: dict[str, list[int]] + + +def _numeric_summary(values: list[float]) -> NumericSummary: + arr = np.asarray(values, dtype=np.float64) + return NumericSummary( + mean=float(np.mean(arr)), + min=float(np.min(arr)), + max=float(np.max(arr)), + median=float(np.median(arr)), + ) + + +def _resolve_default_weights_dir(model_id: str) -> Path | None: + candidate = PROJECT_ROOT / "weights" / model_id.split("/")[-1].lower() + if (candidate / "config.txt").exists(): + return candidate + return None + + +def _audio_duration_seconds(audio_file: str) -> float: + sample_rate, samples = wavfile.read(audio_file) + if sample_rate <= 0: + raise ValueError(f"invalid sample rate in {audio_file}: {sample_rate}") + return float(len(samples)) / float(sample_rate) + + +def _profile_env(profile_file: Path | None): + class _ProfileContext: + def __enter__(self_inner): + self_inner.old_profile = os.environ.get("CACTUS_PROFILE") + self_inner.old_profile_file = os.environ.get("CACTUS_PROFILE_FILE") + if profile_file is None: + os.environ.pop("CACTUS_PROFILE", None) + os.environ.pop("CACTUS_PROFILE_FILE", None) + return self_inner + profile_file.parent.mkdir(parents=True, exist_ok=True) + if profile_file.exists(): + profile_file.unlink() + os.environ.pop("CACTUS_PROFILE", None) + os.environ["CACTUS_PROFILE_FILE"] = str(profile_file) + return self_inner + + def __exit__(self_inner, exc_type, exc, tb): + if self_inner.old_profile is None: + os.environ.pop("CACTUS_PROFILE", None) + else: + os.environ["CACTUS_PROFILE"] = self_inner.old_profile + + if self_inner.old_profile_file is None: + os.environ.pop("CACTUS_PROFILE_FILE", None) + else: + os.environ["CACTUS_PROFILE_FILE"] = self_inner.old_profile_file + return False + + return _ProfileContext() + + +def _parse_profile_file(path: Path) -> list[ProfileSection]: + if not path.exists(): + return [] + + lines = path.read_text().splitlines() + sections: list[ProfileSection] = [] + idx = 0 + while idx < len(lines): + if lines[idx].strip() != "=== Graph Execution Profile ===": + idx += 1 + continue + + entries: list[ProfileEntry] = [] + total_ms = 0.0 + idx += 1 + while idx < len(lines): + line = lines[idx].rstrip() + stripped = line.strip() + if stripped.startswith("Total execution time:"): + try: + total_ms = float(stripped.split(":", 1)[1].strip().split()[0]) + except Exception: + total_ms = 0.0 + elif stripped == "================================": + break + elif ( + stripped + and not stripped.startswith("Operation") + and not set(stripped).issubset({"-"}) + and not stripped.startswith("Total execution time:") + ): + parts = stripped.split() + if len(parts) >= 3: + op = parts[0] + try: + time_ms = float(parts[1]) + except ValueError: + time_ms = 0.0 + shape = parts[2] if parts[2].startswith("[") else "" + entries.append(ProfileEntry(op=op, time_ms=time_ms, shape=shape, raw_line=stripped)) + idx += 1 + + sections.append(ProfileSection(entries=entries, total_ms=total_ms)) + idx += 1 + + return sections + + +def _summarize_profile_sections(sections: list[ProfileSection]) -> dict[str, Any]: + op_total_ms: dict[str, float] = defaultdict(float) + op_count: dict[str, int] = defaultdict(int) + all_entries: list[ProfileEntry] = [] + total_ms_values: list[float] = [] + + for section in sections: + total_ms_values.append(section.total_ms) + all_entries.extend(section.entries) + for entry in section.entries: + op_total_ms[entry.op] += entry.time_ms + op_count[entry.op] += 1 + + by_op = [] + for op, total_ms in sorted(op_total_ms.items(), key=lambda item: item[1], reverse=True): + count = op_count[op] + by_op.append( + { + "op": op, + "count": count, + "total_ms": round(total_ms, 6), + "mean_ms": round(total_ms / max(count, 1), 6), + } + ) + + slowest_nodes = [ + { + "op": entry.op, + "time_ms": round(entry.time_ms, 6), + "shape": entry.shape, + "raw_line": entry.raw_line, + } + for entry in sorted(all_entries, key=lambda item: item.time_ms, reverse=True)[:50] + ] + + return { + "runs_profiled": len(sections), + "graph_total_ms": asdict(_numeric_summary(total_ms_values)) if total_ms_values else None, + "entries_per_run": [len(section.entries) for section in sections], + "by_op": by_op, + "slowest_nodes": slowest_nodes, + } + + +def _run_handwritten_once( + *, + model_handle: Any, + audio_file: str, +) -> HandwrittenRun: + options_json = json.dumps( + { + "use_vad": False, + "max_tokens": 4096, + } + ) + start = time.perf_counter() + response = cactus_transcribe( + model_handle, + audio_file, + "", + options_json, + None, + None, + ) + end = time.perf_counter() + + payload = json.loads(response) + return HandwrittenRun( + wall_ms=(end - start) * 1000.0, + total_time_ms=float(payload.get("total_time_ms", 0.0)), + time_to_first_token_ms=float(payload.get("time_to_first_token_ms", 0.0)), + prefill_tps=float(payload.get("prefill_tps", 0.0)), + decode_tps=float(payload.get("decode_tps", 0.0)), + decode_tokens=int(payload.get("decode_tokens", 0)), + prefill_tokens=int(payload.get("prefill_tokens", 0)), + total_tokens=int(payload.get("total_tokens", 0)), + confidence=float(payload.get("confidence", 0.0)), + transcript=str(payload.get("response", "")), + response_json=payload, + ) + + +def _run_transpiled_once( + *, + tg: Any, + processor: object | None, + input_names: tuple[str, ...], + model_config: dict[str, object], + preprocessor_config: dict[str, object], + model: torch.nn.Module, + torch_dtype: torch.dtype, + audio_file: str, + tokenizer: object | None, + blank_token_id: int | None, +) -> TranspiledRun: + start_pre = time.perf_counter() + prepared = _prepare_audio_inputs( + processor, + input_names=input_names, + config=model_config, + preprocessor_config=preprocessor_config, + model=model, + task="ctc_logits", + audio_file=audio_file, + torch_dtype=torch_dtype, + ) + end_pre = time.perf_counter() + + input_arrays = [tensor.detach().cpu().numpy() for tensor in prepared.tensors] + + start_exec = time.perf_counter() + tg.set_inputs(input_arrays) + outputs = tg.execute() + end_exec = time.perf_counter() + + start_decode = time.perf_counter() + logits = outputs[0].numpy().astype(np.float32) + token_ids = _ctc_greedy_decode_token_ids(logits, blank_token_id=blank_token_id) + transcript = _decode_token_ids(tokenizer, token_ids) if tokenizer is not None else None + end_decode = time.perf_counter() + + return TranspiledRun( + preprocess_ms=(end_pre - start_pre) * 1000.0, + execute_ms=(end_exec - start_exec) * 1000.0, + decode_ms=(end_decode - start_decode) * 1000.0, + total_ms=(end_decode - start_pre) * 1000.0, + graph_only_ms=(end_exec - start_exec) * 1000.0, + output_frames=int(logits.shape[1]), + output_vocab=int(logits.shape[2]), + decode_tokens=len(token_ids), + transcript=transcript, + input_shapes={ + name: list(tensor.shape) + for name, tensor in zip(prepared.names, prepared.tensors) + }, + ) + + +def _run_transpiled_graph_only_once( + *, + tg: Any, + cached_inputs: list[np.ndarray], +) -> float: + start = time.perf_counter() + tg.set_inputs(cached_inputs) + tg.execute() + end = time.perf_counter() + return (end - start) * 1000.0 + + +def _warmup(fn, count: int) -> None: + for _ in range(max(count, 0)): + fn() + + +def _roundtrip_jsonable(obj: Any) -> Any: + if isinstance(obj, dict): + return {str(key): _roundtrip_jsonable(value) for key, value in obj.items()} + if isinstance(obj, list): + return [_roundtrip_jsonable(value) for value in obj] + if isinstance(obj, tuple): + return [_roundtrip_jsonable(value) for value in obj] + if isinstance(obj, Path): + return str(obj) + if isinstance(obj, NumericSummary): + return asdict(obj) + if hasattr(obj, "__dataclass_fields__"): + return asdict(obj) + return obj + + +def _build_transpiled_graph( + *, + model_id: str, + audio_file: str, + weights_dir: str | None, + torch_dtype: torch.dtype, + token: str | None, + trust_remote_code: bool, + local_files_only: bool, + fusion_config: FusionConfig, +): + task = _infer_task_from_config(model_id) + if task != "ctc_logits": + raise RuntimeError(f"expected a CTC model, got task={task}") + + model_source, processor, model, model_config = _load_transformers_bundle( + model_id=model_id, + task=task, + torch_dtype=torch_dtype, + token=token, + trust_remote_code=trust_remote_code, + local_files_only=local_files_only, + ) + + preprocessor_config = _load_optional_json(model_source, "preprocessor_config.json") + if not preprocessor_config: + preprocessor_config = _load_optional_json(model_id, "preprocessor_config.json") + + tokenizer = _load_optional_tokenizer( + model_id=model_id, + model_source=model_source, + token=token, + trust_remote_code=trust_remote_code, + local_files_only=local_files_only, + ) + + canonical = canonicalize_model_interface(model, task=task) + prepared = _prepare_audio_inputs( + processor, + input_names=canonical.input_names, + config=model_config, + preprocessor_config=preprocessor_config, + model=model, + task=task, + audio_file=audio_file, + torch_dtype=torch_dtype, + ) + canonical = canonicalize_model_interface(model, task=task, input_names=prepared.names) + + wrapper = TranspileWrapper(canonical.module, weights_dir=weights_dir).eval() + + capture_started = time.perf_counter() + captured = capture_model(wrapper, prepared.tensors) + canonicalize_exported_graph(captured.ir_graph) + optimize_graph(captured.ir_graph, config=fusion_config) + tg = _lower_preoptimized_ir(captured.ir_graph) + capture_finished = time.perf_counter() + + cached_inputs = [tensor.detach().cpu().numpy() for tensor in prepared.tensors] + blank_token_id = getattr(getattr(model, "config", None), "pad_token_id", None) + + return { + "tg": tg, + "model": model, + "processor": processor, + "model_config": model_config, + "preprocessor_config": preprocessor_config, + "tokenizer": tokenizer, + "blank_token_id": int(blank_token_id) if blank_token_id is not None else None, + "compile_time_ms": (capture_finished - capture_started) * 1000.0, + "cached_inputs": cached_inputs, + "input_names": tuple(prepared.names), + "prepared_input_shapes": { + name: list(tensor.shape) + for name, tensor in zip(prepared.names, prepared.tensors) + }, + "model_source": model_source, + "weight_bindings": sum( + 1 + for value in captured.ir_graph.values.values() + if isinstance(value.meta, dict) and isinstance(value.meta.get("path"), str) + ), + } + + +def _benchmark_handwritten( + *, + weights_dir: Path, + audio_file: str, + warmup: int, + repeats: int, + profile_repeats: int, + artifact_dir: Path | None, +) -> dict[str, Any]: + model_handle = cactus_init(str(weights_dir), None, False) + try: + _warmup(lambda: _run_handwritten_once(model_handle=model_handle, audio_file=audio_file), warmup) + + runs: list[HandwrittenRun] = [] + for _ in range(max(repeats, 1)): + runs.append(_run_handwritten_once(model_handle=model_handle, audio_file=audio_file)) + + profile_path = ( + artifact_dir / "handwritten_profile.txt" + if artifact_dir is not None and profile_repeats > 0 + else None + ) + with _profile_env(profile_path): + for _ in range(max(profile_repeats, 0)): + _run_handwritten_once(model_handle=model_handle, audio_file=audio_file) + + profile_sections = _parse_profile_file(profile_path) if profile_path is not None else [] + profile_summary = _summarize_profile_sections(profile_sections) + + return { + "weights_dir": str(weights_dir), + "runs": [_roundtrip_jsonable(run) for run in runs], + "wall_ms": asdict(_numeric_summary([run.wall_ms for run in runs])), + "internal_total_time_ms": asdict(_numeric_summary([run.total_time_ms for run in runs])), + "time_to_first_token_ms": asdict(_numeric_summary([run.time_to_first_token_ms for run in runs])), + "decode_tps": asdict(_numeric_summary([run.decode_tps for run in runs])), + "prefill_tps": asdict(_numeric_summary([run.prefill_tps for run in runs])), + "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), + "confidence": asdict(_numeric_summary([run.confidence for run in runs])), + "latest_transcript": runs[-1].transcript if runs else "", + "profile_file": str(profile_path) if profile_path is not None else None, + "profile_summary": profile_summary, + } + finally: + cactus_destroy(model_handle) + + +def _benchmark_transpiled( + *, + model_id: str, + audio_file: str, + weights_dir: str | None, + warmup: int, + repeats: int, + profile_repeats: int, + torch_dtype: torch.dtype, + token: str | None, + trust_remote_code: bool, + local_files_only: bool, + artifact_dir: Path | None, + fusion_config: FusionConfig, +) -> dict[str, Any]: + built = _build_transpiled_graph( + model_id=model_id, + audio_file=audio_file, + weights_dir=weights_dir, + torch_dtype=torch_dtype, + token=token, + trust_remote_code=trust_remote_code, + local_files_only=local_files_only, + fusion_config=fusion_config, + ) + + tg = built["tg"] + model = built["model"] + processor = built["processor"] + input_names = built["input_names"] + model_config = built["model_config"] + preprocessor_config = built["preprocessor_config"] + tokenizer = built["tokenizer"] + blank_token_id = built["blank_token_id"] + cached_inputs = built["cached_inputs"] + + _warmup( + lambda: _run_transpiled_once( + tg=tg, + processor=processor, + input_names=input_names, + model_config=model_config, + preprocessor_config=preprocessor_config, + model=model, + torch_dtype=torch_dtype, + audio_file=audio_file, + tokenizer=tokenizer, + blank_token_id=blank_token_id, + ), + warmup, + ) + _warmup(lambda: _run_transpiled_graph_only_once(tg=tg, cached_inputs=cached_inputs), warmup) + + runs: list[TranspiledRun] = [] + graph_only_ms: list[float] = [] + for _ in range(max(repeats, 1)): + runs.append( + _run_transpiled_once( + tg=tg, + processor=processor, + input_names=input_names, + model_config=model_config, + preprocessor_config=preprocessor_config, + model=model, + torch_dtype=torch_dtype, + audio_file=audio_file, + tokenizer=tokenizer, + blank_token_id=blank_token_id, + ) + ) + graph_only_ms.append(_run_transpiled_graph_only_once(tg=tg, cached_inputs=cached_inputs)) + + profile_path = ( + artifact_dir / "transpiled_profile.txt" + if artifact_dir is not None and profile_repeats > 0 + else None + ) + with _profile_env(profile_path): + for _ in range(max(profile_repeats, 0)): + _run_transpiled_graph_only_once(tg=tg, cached_inputs=cached_inputs) + + profile_sections = _parse_profile_file(profile_path) if profile_path is not None else [] + profile_summary = _summarize_profile_sections(profile_sections) + + return { + "model_source": built["model_source"], + "compile_time_ms": built["compile_time_ms"], + "weight_bindings": built["weight_bindings"], + "prepared_input_shapes": built["prepared_input_shapes"], + "runs": [_roundtrip_jsonable(run) for run in runs], + "preprocess_ms": asdict(_numeric_summary([run.preprocess_ms for run in runs])), + "execute_ms": asdict(_numeric_summary([run.execute_ms for run in runs])), + "decode_ms": asdict(_numeric_summary([run.decode_ms for run in runs])), + "total_ms": asdict(_numeric_summary([run.total_ms for run in runs])), + "graph_only_ms": asdict(_numeric_summary(graph_only_ms)), + "decode_tokens": asdict(_numeric_summary([float(run.decode_tokens) for run in runs])), + "output_frames": asdict(_numeric_summary([float(run.output_frames) for run in runs])), + "latest_transcript": runs[-1].transcript if runs else "", + "profile_file": str(profile_path) if profile_path is not None else None, + "profile_summary": profile_summary, + } + + +def _print_summary( + *, + audio_duration_sec: float, + handwritten: dict[str, Any] | None, + transpiled: dict[str, Any] | None, +) -> None: + print() + print(f"audio_duration_sec={audio_duration_sec:.3f}") + + if handwritten is not None: + hw_total = handwritten["internal_total_time_ms"]["mean"] + hw_graph = None + profile_total = handwritten.get("profile_summary", {}).get("graph_total_ms") + if isinstance(profile_total, dict): + hw_graph = profile_total.get("mean") + hw_tok = handwritten["decode_tokens"]["mean"] + hw_rtf = hw_total / max(audio_duration_sec * 1000.0, 1e-9) + hw_xrt = (audio_duration_sec * 1000.0) / max(hw_total, 1e-9) + hw_tok_s = hw_tok * 1000.0 / max(hw_total, 1e-9) + print("handwritten:") + print(f" total_ms mean={hw_total:.3f}") + print(f" graph_profile_ms mean={hw_graph:.3f}" if hw_graph is not None else " graph_profile_ms unavailable") + print(f" decode_tokens mean={hw_tok:.3f}") + print(f" output_tokens_per_sec={hw_tok_s:.3f}") + print(f" real_time_factor={hw_rtf:.3f}") + print(f" audio_x_realtime={hw_xrt:.3f}") + + if transpiled is not None: + tr_total = transpiled["total_ms"]["mean"] + tr_graph = transpiled["graph_only_ms"]["mean"] + tr_tok = transpiled["decode_tokens"]["mean"] + tr_rtf = tr_total / max(audio_duration_sec * 1000.0, 1e-9) + tr_xrt = (audio_duration_sec * 1000.0) / max(tr_total, 1e-9) + tr_tok_s = tr_tok * 1000.0 / max(tr_total, 1e-9) + tr_graph_tok_s = tr_tok * 1000.0 / max(tr_graph, 1e-9) + print("transpiled:") + print(f" total_ms mean={tr_total:.3f}") + print(f" graph_only_ms mean={tr_graph:.3f}") + print(f" decode_tokens mean={tr_tok:.3f}") + print(f" output_tokens_per_sec={tr_tok_s:.3f}") + print(f" graph_only_tokens_per_sec={tr_graph_tok_s:.3f}") + print(f" real_time_factor={tr_rtf:.3f}") + print(f" audio_x_realtime={tr_xrt:.3f}") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=( + "Compare handwritten Cactus Parakeet performance against the transpiled " + "Parakeet graph on the same WAV input, including per-op graph profiles." + ) + ) + parser.add_argument("--model-id", default="nvidia/parakeet-ctc-1.1b") + parser.add_argument("--audio-file", required=True) + parser.add_argument( + "--weights-dir", + default="", + help="Converted Cactus weights directory. Reused for handwritten runtime and transpiled mmap bindings.", + ) + parser.add_argument("--artifact-dir", default="", help="Optional directory for JSON summaries and raw profiles.") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--profile-repeats", type=int, default=1) + parser.add_argument("--torch-dtype", default="float16") + parser.add_argument("--token", default=os.environ.get("HF_TOKEN")) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--skip-handwritten", action="store_true") + parser.add_argument("--skip-transpiled", action="store_true") + parser.add_argument("--no-fuse-gated-deltanet", action="store_true") + parser.add_argument("--no-fuse-rms-norm", action="store_true") + parser.add_argument("--no-fuse-rope", action="store_true") + parser.add_argument("--no-fuse-attention", action="store_true") + parser.add_argument("--no-fuse-attention-block", action="store_true") + parser.add_argument("--no-fuse-add-clipped", action="store_true") + args = parser.parse_args() + + audio_file = str(Path(args.audio_file).resolve()) + if not Path(audio_file).exists(): + raise FileNotFoundError(f"audio file does not exist: {audio_file}") + + torch_dtype = _parse_dtype(args.torch_dtype) + artifact_dir = Path(args.artifact_dir).resolve() if args.artifact_dir else None + if artifact_dir is not None: + artifact_dir.mkdir(parents=True, exist_ok=True) + + weights_dir_path = None + if args.weights_dir.strip(): + weights_dir_path = _validate_weights_dir(args.weights_dir.strip(), model_id=args.model_id) + else: + weights_dir_path = _resolve_default_weights_dir(args.model_id) + + if not args.skip_handwritten and weights_dir_path is None: + raise RuntimeError( + "handwritten comparison requires a converted weights directory.\n" + "\n" + f"Create one with:\n" + f" cactus convert {args.model_id} {PROJECT_ROOT / 'weights' / args.model_id.split('/')[-1].lower()}\n" + "\n" + "Or rerun with --skip-handwritten." + ) + + fusion_config = FusionConfig( + enable_gated_deltanet=not args.no_fuse_gated_deltanet, + enable_rms_norm=not args.no_fuse_rms_norm, + enable_rope=not args.no_fuse_rope, + enable_attention=not args.no_fuse_attention, + enable_attention_block=not args.no_fuse_attention_block, + enable_add_clipped=not args.no_fuse_add_clipped, + ) + + summary: dict[str, Any] = { + "model_id": args.model_id, + "audio_file": audio_file, + "audio_duration_sec": _audio_duration_seconds(audio_file), + "weights_dir": str(weights_dir_path) if weights_dir_path is not None else None, + "warmup": args.warmup, + "repeats": args.repeats, + "profile_repeats": args.profile_repeats, + "torch_dtype": args.torch_dtype, + "trust_remote_code": bool(args.trust_remote_code), + "local_files_only": bool(args.local_files_only), + } + + handwritten_summary = None + if not args.skip_handwritten and weights_dir_path is not None: + handwritten_summary = _benchmark_handwritten( + weights_dir=weights_dir_path, + audio_file=audio_file, + warmup=args.warmup, + repeats=args.repeats, + profile_repeats=args.profile_repeats, + artifact_dir=artifact_dir, + ) + summary["handwritten"] = handwritten_summary + + transpiled_summary = None + if not args.skip_transpiled: + transpiled_summary = _benchmark_transpiled( + model_id=args.model_id, + audio_file=audio_file, + weights_dir=str(weights_dir_path) if weights_dir_path is not None else None, + warmup=args.warmup, + repeats=args.repeats, + profile_repeats=args.profile_repeats, + torch_dtype=torch_dtype, + token=args.token, + trust_remote_code=args.trust_remote_code, + local_files_only=args.local_files_only, + artifact_dir=artifact_dir, + fusion_config=fusion_config, + ) + summary["transpiled"] = transpiled_summary + + if handwritten_summary is not None and transpiled_summary is not None: + hw_total = handwritten_summary["internal_total_time_ms"]["mean"] + tr_total = transpiled_summary["total_ms"]["mean"] + hw_graph = handwritten_summary.get("profile_summary", {}).get("graph_total_ms") + tr_graph = transpiled_summary["graph_only_ms"]["mean"] + summary["comparison"] = { + "end_to_end_speedup_vs_handwritten": hw_total / max(tr_total, 1e-9), + "graph_only_speedup_vs_handwritten": ( + float(hw_graph["mean"]) / max(tr_graph, 1e-9) + if isinstance(hw_graph, dict) and "mean" in hw_graph + else None + ), + } + + if artifact_dir is not None: + summary_path = artifact_dir / "summary.json" + summary_path.write_text(json.dumps(_roundtrip_jsonable(summary), indent=2, sort_keys=True) + "\n") + print(f"saved_summary={summary_path}") + if handwritten_summary is not None and handwritten_summary.get("profile_file"): + print(f"saved_handwritten_profile={handwritten_summary['profile_file']}") + if transpiled_summary is not None and transpiled_summary.get("profile_file"): + print(f"saved_transpiled_profile={transpiled_summary['profile_file']}") + + _print_summary( + audio_duration_sec=float(summary["audio_duration_sec"]), + handwritten=handwritten_summary, + transpiled=transpiled_summary, + ) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/tests/transpile/tools/inspect_gemma4_cleanup.py b/python/tests/transpile/tools/inspect_gemma4_cleanup.py new file mode 100644 index 000000000..835311f2e --- /dev/null +++ b/python/tests/transpile/tools/inspect_gemma4_cleanup.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +import torch +from transformers import AutoModelForCausalLM + +PYTHON_ROOT = Path(__file__).resolve().parents[3] +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.model_adapters import canonicalize_model_interface + + +class Gemma4FullModelWrapper(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = canonicalize_model_interface(model, task="causal_lm_logits").module + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model(input_ids) + + +class Gemma4FirstBlockCheckpointWrapper(torch.nn.Module): + def __init__(self, model: torch.nn.Module, checkpoint_name: str): + super().__init__() + adapter = canonicalize_model_interface(model, task="causal_lm_logits") + if adapter.family != "gemma4": + raise ValueError(f"expected gemma4 adapter, got family={adapter.family}") + if not hasattr(adapter.module, "debug_first_block"): + raise ValueError("Gemma4 adapter does not expose debug_first_block()") + self.adapter = adapter.module + self.checkpoint_name = checkpoint_name + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + checkpoints = self.adapter.debug_first_block(input_ids) + if self.checkpoint_name not in checkpoints: + available = ", ".join(sorted(checkpoints.keys())) + raise KeyError(f"unknown checkpoint {self.checkpoint_name!r}; available: {available}") + return checkpoints[self.checkpoint_name] + + +def format_ir_graph(graph: IRGraph, *, max_nodes: int | None = None) -> str: + lines: list[str] = [] + lines.append("IR Summary") + lines.append(f" inputs={graph.inputs}") + lines.append(f" outputs={graph.outputs}") + lines.append(f" nodes={len(graph.order)}") + if graph.meta: + lines.append(f" meta={graph.meta}") + lines.append("") + lines.append("IR Nodes") + + node_ids = graph.order if max_nodes is None else graph.order[:max_nodes] + for index, node_id in enumerate(node_ids): + node = graph.nodes[node_id] + lines.append(f"[{index}] {node.id} op={node.op}") + lines.append(f" inputs={node.inputs}") + lines.append(f" outputs={node.outputs}") + if node.attrs: + lines.append(f" attrs={node.attrs}") + if node.kind != "generic": + lines.append(f" kind={node.kind}") + if node.meta: + lines.append(f" meta={node.meta}") + for output_id in node.outputs: + value = graph.values.get(output_id) + if value is None: + continue + lines.append( + f" value[{output_id}] shape={value.shape} dtype={value.dtype} users={value.users}" + ) + + if max_nodes is not None and len(graph.order) > max_nodes: + lines.append(f"... ({len(graph.order) - max_nodes} more nodes omitted)") + + if graph.constants: + lines.append("") + lines.append("IR Constants") + constant_ids = sorted(graph.constants.keys()) + preview_ids = constant_ids if max_nodes is None else constant_ids[: max(8, min(len(constant_ids), max_nodes))] + for value_id in preview_ids: + value = graph.constants[value_id] + value_meta = graph.values.get(value_id) + shape = None if value_meta is None else value_meta.shape + dtype = None if value_meta is None else value_meta.dtype + lines.append(f"- {value_id}: type={type(value).__name__} shape={shape} dtype={dtype}") + if len(preview_ids) < len(constant_ids): + lines.append(f"... ({len(constant_ids) - len(preview_ids)} more constants omitted)") + + return "\n".join(lines) + + +def format_unsupported_ops(graph: IRGraph) -> str: + counts = graph.meta.get("canonical_unsupported_op_counts", {}) + if not isinstance(counts, dict): + counts = {} + + examples: dict[str, list[str]] = {} + for node_id in graph.order: + node = graph.nodes[node_id] + print(f"node {node_id} op={node.op}") + if node.op not in counts: + continue + bucket = examples.setdefault(node.op, []) + if len(bucket) < 5: + bucket.append(node.id) + + total = sum(int(count) for count in counts.values()) + lines = ["Unsupported Ops", f" unique={len(counts)}", f" total_nodes={total}"] + if not counts: + lines.append(" none") + return "\n".join(lines) + + for op_name, count in counts.items(): + sample_nodes = ", ".join(examples.get(op_name, ())) + lines.append(f"- {op_name}: {count}") + if sample_nodes: + lines.append(f" examples={sample_nodes}") + return "\n".join(lines) + + +def resolve_local_model_path(model_id: str) -> str: + model_path = Path(model_id) + if model_path.exists(): + return str(model_path) + + cache_root = Path.home() / ".cache" / "huggingface" / "hub" / ("models--" + model_id.replace("/", "--")) + snapshots_dir = cache_root / "snapshots" + if not snapshots_dir.exists(): + raise FileNotFoundError( + f"could not find local snapshot for {model_id!r} under {snapshots_dir}" + ) + snapshots = sorted(path for path in snapshots_dir.iterdir() if path.is_dir()) + if not snapshots: + raise FileNotFoundError(f"no snapshots found for {model_id!r} under {snapshots_dir}") + return str(snapshots[-1]) + + +def load_model(model_id: str) -> torch.nn.Module: + token = os.environ.get("HF_TOKEN") + common_kwargs: dict[str, object] = { + "local_files_only": True, + } + if token: + common_kwargs["token"] = token + + local_model_path = resolve_local_model_path(model_id) + model = AutoModelForCausalLM.from_pretrained( + local_model_path, + torch_dtype=torch.float16, + device_map=None, + low_cpu_mem_usage=True, + **common_kwargs, + ).eval() + return model + + +def build_module( + model: torch.nn.Module, + *, + mode: str, + checkpoint: str, +) -> torch.nn.Module: + if mode == "full": + return Gemma4FullModelWrapper(model).eval() + if mode == "first_block": + return Gemma4FirstBlockCheckpointWrapper(model, checkpoint).eval() + raise ValueError(f"unknown mode: {mode}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Capture Gemma4, run canonicalize/cleanup.py, and report unsupported ops.", + ) + parser.add_argument( + "--model-id", + default=os.environ.get("CACTUS_GEMMA_HF_MODEL_ID", "google/gemma-4-E2B"), + help="Hugging Face model id. Uses local cache only.", + ) + parser.add_argument( + "--mode", + choices=("first_block", "full"), + default="first_block", + help="Capture a first-block checkpoint or the full Gemma4 adapter graph.", + ) + parser.add_argument( + "--checkpoint", + default="after_ffn_residual", + help="First-block checkpoint name when --mode=first_block.", + ) + parser.add_argument( + "--input-ids", + default="2,818,5279,529,7001,563", + help="Comma-separated integer token ids to capture.", + ) + parser.add_argument( + "--max-nodes", + type=int, + default=120, + help="Maximum IR nodes to print when --print-ir is enabled.", + ) + parser.add_argument( + "--print-ir", + action="store_true", + help="Also print the IR before and after canonical cleanup.", + ) + args = parser.parse_args() + + model = load_model(args.model_id) + module = build_module(model, mode=args.mode, checkpoint=args.checkpoint) + token_ids = [int(token.strip()) for token in args.input_ids.split(",") if token.strip()] + input_ids = torch.tensor([token_ids], dtype=torch.long) + + print(f"Model: {args.model_id}") + print(f"Mode: {args.mode}") + print(f"Checkpoint: {args.checkpoint if args.mode == 'first_block' else '(full model)'}") + print(f"Input IDs shape: {tuple(input_ids.shape)}") + print("") + + captured = capture_model(module, (input_ids,)) + + if args.print_ir: + print("=== IR Before Canonical Cleanup ===") + print(format_ir_graph(captured.ir_graph, max_nodes=args.max_nodes)) + print("") + + canonicalize_exported_graph(captured.ir_graph) + + print("=== Unsupported Ops After Canonical Cleanup ===") + print(format_unsupported_ops(captured.ir_graph)) + + if args.print_ir: + print("") + print("=== IR After Canonical Cleanup ===") + print(format_ir_graph(captured.ir_graph, max_nodes=args.max_nodes)) + + +if __name__ == "__main__": + main() diff --git a/python/tests/transpile/tools/inspect_transpile_cleanup.py b/python/tests/transpile/tools/inspect_transpile_cleanup.py new file mode 100644 index 000000000..b5b2f3d40 --- /dev/null +++ b/python/tests/transpile/tools/inspect_transpile_cleanup.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import torch +import torch.nn as nn + +PYTHON_ROOT = Path(__file__).resolve().parents[3] +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.capture_pytorch import dump_graph +from cactus.transpile.canonicalize.cleanup import canonicalize_exported_graph +from cactus.transpile.graph_ir import IRGraph + + +class Toy(nn.Module): + def forward(self, x, y): + z = x + y + z = z * 0.5 + z = torch.nn.functional.gelu(z) + z = torch.softmax(z, dim=-1) + return z + + +class AttentionBlockToy(nn.Module): + def __init__(self, hidden_size: int = 8): + super().__init__() + self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False, dtype=torch.float16) + + def forward(self, q, k, v, gate): + attn = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True, scale=1.0) + attn = attn.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1) + attn = attn * torch.sigmoid(gate) + attn = attn.to(self.out_proj.weight.dtype) + return self.out_proj(attn) + + +def build_case(name: str) -> tuple[nn.Module, tuple[torch.Tensor, ...]]: + if name == "toy": + model = Toy().eval() + args = ( + torch.randn(2, 4, dtype=torch.float16), + torch.randn(2, 4, dtype=torch.float16), + ) + return model, args + + if name == "attention_block": + model = AttentionBlockToy().eval() + args = ( + torch.randn(1, 2, 3, 4, dtype=torch.float16), + torch.randn(1, 2, 3, 4, dtype=torch.float16), + torch.randn(1, 2, 3, 4, dtype=torch.float16), + torch.randn(1, 3, 8, dtype=torch.float16), + ) + return model, args + + raise ValueError(f"unknown case: {name}") + + +def format_ir_graph(graph: IRGraph) -> str: + lines: list[str] = [] + lines.append("IR Summary") + lines.append(f" inputs={graph.inputs}") + lines.append(f" outputs={graph.outputs}") + lines.append(f" nodes={len(graph.order)}") + if graph.meta: + lines.append(f" meta={graph.meta}") + lines.append("") + lines.append("IR Nodes") + + for index, node_id in enumerate(graph.order): + node = graph.nodes[node_id] + lines.append(f"[{index}] {node.id} op={node.op}") + lines.append(f" inputs={node.inputs}") + lines.append(f" outputs={node.outputs}") + if node.attrs: + lines.append(f" attrs={node.attrs}") + if node.kind != "generic": + lines.append(f" kind={node.kind}") + if node.meta: + lines.append(f" meta={node.meta}") + for output_id in node.outputs: + value = graph.values.get(output_id) + if value is None: + continue + lines.append( + f" value[{output_id}] shape={value.shape} dtype={value.dtype} users={value.users}" + ) + + if graph.constants: + lines.append("") + lines.append("IR Constants") + for value_id, value in graph.constants.items(): + value_meta = graph.values.get(value_id) + shape = None if value_meta is None else value_meta.shape + dtype = None if value_meta is None else value_meta.dtype + typename = type(value).__name__ + lines.append(f"- {value_id}: type={typename} shape={shape} dtype={dtype}") + + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Capture a small PyTorch model, dump the exported graph, and print IR before/after cleanup.", + ) + parser.add_argument( + "--case", + choices=("toy", "attention_block"), + default="toy", + help="Which built-in model to inspect.", + ) + parser.add_argument( + "--no-meta", + action="store_true", + help="Hide exported graph metadata to keep the dump shorter.", + ) + args = parser.parse_args() + + model, example_args = build_case(args.case) + captured = capture_model(model, example_args) + + print(f"Case: {args.case}") + print(f"Model: {type(model).__name__}") + print("") + print("=== Exported Graph ===") + print(dump_graph(captured, include_meta=not args.no_meta)) + print("") + print("=== IR Before Cleanup ===") + print(format_ir_graph(captured.ir_graph)) + print("") + + canonicalize_exported_graph(captured.ir_graph) + + print("=== IR After Cleanup ===") + print(format_ir_graph(captured.ir_graph)) + + +if __name__ == "__main__": + main() diff --git a/python/tests/transpile/tools/torch_capture_smoke.py b/python/tests/transpile/tools/torch_capture_smoke.py new file mode 100644 index 000000000..f3be0f51e --- /dev/null +++ b/python/tests/transpile/tools/torch_capture_smoke.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import torch +import torch.nn as nn + +PYTHON_ROOT = Path(__file__).resolve().parents[3] +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.transpile.capture_pytorch import ( + capture_model_with_fallback, + dump_graph, + get_dtype, + get_shape, +) +from cactus.transpile.lower import transpile_captured + + +class TinyBlock(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(8, 16) + self.fc2 = nn.Linear(16, 4) + + def forward(self, x): + h = self.fc1(x) + h = torch.relu(h) + y = self.fc2(h) + return y + + +def main(): + model = TinyBlock().eval() + x = torch.randn(2, 8) + + captured = capture_model_with_fallback(model, args=(x,)) + + print("strict:", captured.strict) + print("graph_module type:", type(captured.graph_module).__name__) + print("state_dict keys:", list(captured.state_dict.keys())) + print() + + print("=== FX Graph ===") + print(captured.graph) + print() + + print("=== Node Walk ===") + for i, node in enumerate(captured.graph.nodes): + print(f"[{i}] op={node.op} name={node.name} target={node.target}") + print(f" args={node.args}") + print(f" kwargs={node.kwargs}") + print(f" shape={get_shape(node)} dtype={get_dtype(node)}") + print() + + print("=== Full Dump ===") + print(dump_graph(captured)) + print() + + print("=== Lower To Cactus ===") + try: + transpiled = transpile_captured(captured) + print("lowering succeeded") + print("lowered graph type:", type(transpiled.graph).__name__) + print("runtime inputs:", transpiled.runtime_inputs) + print("bound constants:", transpiled.bound_constants) + print("outputs:", transpiled.outputs) + print() + print("binding runtime input 0 from example tensor") + transpiled.set_input(0, x) + except NotImplementedError as exc: + print("lowering stopped on unimplemented case:") + print(f" {exc}") + + +if __name__ == "__main__": + main() diff --git a/python/tests/transpile/tools/transpile_hf_causal_lm.py b/python/tests/transpile/tools/transpile_hf_causal_lm.py new file mode 100644 index 000000000..a1982228e --- /dev/null +++ b/python/tests/transpile/tools/transpile_hf_causal_lm.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +PYTHON_ROOT = Path(__file__).resolve().parents[3] +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + +from cactus.transpile.hf_model import main as transpile_main + + +def main() -> int: + if "--task" not in sys.argv[1:]: + sys.argv[1:1] = ["--task", "causal_lm_logits"] + return transpile_main() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/setup b/setup index 9f106bebf..4539a168c 100755 --- a/setup +++ b/setup @@ -2,11 +2,9 @@ if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then echo "Please run: source ./setup" - exit 1 + return 1 2>/dev/null || exit 1 fi -set -e - PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" VENV_DIR="$PROJECT_ROOT/venv" REQUIREMENTS_FILE="${CACTUS_REQUIREMENTS:-$PROJECT_ROOT/python/requirements.txt}" @@ -26,8 +24,8 @@ echo -e "${BLUE}Step 1: Configuring git hooks for DCO...${NC}" git config core.hooksPath .githooks echo -e "${GREEN}✓ Git hooks configured${NC}" -name=$(git config user.name) -email=$(git config user.email) +name=$(git config user.name || true) +email=$(git config user.email || true) if [ -z "$name" ] || [ -z "$email" ]; then echo "" @@ -44,22 +42,22 @@ fi echo "" echo -e "${BLUE}Step 2: Setting up Python virtual environment...${NC}" -if ! command -v python3 &> /dev/null; then - echo -e "${RED}Error: Python3 is not installed${NC}" +if ! command -v python3.12 &> /dev/null; then + echo -e "${RED}Error: python3.12 is not installed${NC}" echo "" - echo "Please install Python3:" - echo " macOS: brew install python3" - echo " Ubuntu/Debian: sudo apt-get install python3 python3-venv" - exit 1 + echo "Please install python3.12:" + echo " macOS: brew install python@3.12" + echo " Ubuntu/Debian: sudo apt-get install python3.12 python3.12-venv" + return 1 fi if [ ! -d "$VENV_DIR" ]; then - if ! python3 -m venv "$VENV_DIR"; then + if ! python3.12 -m venv "$VENV_DIR"; then echo -e "${RED}Failed to create virtual environment${NC}" - echo "Please ensure python3-venv is installed:" - echo " Ubuntu/Debian: sudo apt-get install python3-venv" - echo " macOS: Should be included with Python3" - exit 1 + echo "Please ensure python3.12-venv is installed:" + echo " Ubuntu/Debian: sudo apt-get install python3.12-venv" + echo " macOS: Should be included with Python3.12" + return 1 fi echo -e "${GREEN}✓ Virtual environment created${NC}" else @@ -71,14 +69,14 @@ source "$VENV_DIR/bin/activate" echo "" echo -e "${BLUE}Step 3: Installing Python dependencies...${NC}" -pip install --upgrade pip -q +python3 -m pip install --upgrade pip -q if [ -f "$REQUIREMENTS_FILE" ]; then - if pip install -r "$REQUIREMENTS_FILE" -q; then + if python3 -m pip install -r "$REQUIREMENTS_FILE" -q; then echo -e "${GREEN}✓ Dependencies installed${NC}" else echo -e "${RED}Failed to install dependencies${NC}" - exit 1 + return 1 fi else echo -e "${YELLOW}Warning: requirements.txt not found${NC}" @@ -88,7 +86,7 @@ if [ -f "$PARENT_REQUIREMENTS" ] && [ "$PARENT_REQUIREMENTS" != "$REQUIREMENTS_F echo "" echo -e "${BLUE}Detected parent repo requirements at: $PARENT_REQUIREMENTS${NC}" echo -e "${BLUE}Installing parent repo requirements...${NC}" - if pip install -r "$PARENT_REQUIREMENTS" -q; then + if python3 -m pip install -r "$PARENT_REQUIREMENTS" -q; then echo -e "${GREEN}✓ Parent repo dependencies installed${NC}" else echo -e "${YELLOW}Warning: failed to install parent repo requirements${NC}" @@ -98,11 +96,11 @@ fi echo "" echo -e "${BLUE}Step 4: Installing cactus CLI tools...${NC}" -if pip install -e "$PROJECT_ROOT/python" --quiet; then +if python3 -m pip install -e "$PROJECT_ROOT/python" --quiet; then echo -e "${GREEN}✓ Cactus CLI installed${NC}" else echo -e "${RED}Failed to install cactus CLI${NC}" - exit 1 + return 1 fi echo "" @@ -111,5 +109,3 @@ echo -e "${GREEN}Setup complete!${NC}" echo -e "${GREEN}=============================================${NC}" echo "" cactus --help - -set +e diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt deleted file mode 100644 index 60ae849f2..000000000 --- a/tests/CMakeLists.txt +++ /dev/null @@ -1,88 +0,0 @@ -cmake_minimum_required(VERSION 3.10) -project(CactusTests LANGUAGES CXX) - -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED True) - -if(APPLE) - set(CMAKE_OSX_ARCHITECTURES "arm64") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -arch arm64 -pthread -Wall -Wextra -pedantic -O3") - find_package(CURL REQUIRED) -else() -endif() - -set(CACTUS_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../cactus) -set(CACTUS_BUILD_DIR ${CACTUS_ROOT}/build) - -find_library(CACTUS_LIB - NAMES cactus - PATHS ${CACTUS_BUILD_DIR} ${CACTUS_BUILD_DIR}/lib - REQUIRED -) - -set(MMAN_LIB "") -if(NOT APPLE) - find_library(MMAN_LIB mman) - if(MMAN_LIB) - message(STATUS "Found libmman: ${MMAN_LIB}") - else() - message(STATUS "libmman not found; assuming mmap symbols are provided by the C runtime") - endif() -else() - find_library(COREML_FRAMEWORK CoreML REQUIRED) - find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED) -endif() - -find_package(SDL2) -if(SDL2_FOUND) - message(STATUS "Found SDL2: ${SDL2_LIBRARIES}") -else() - message(STATUS "SDL2 not found; microphone tests will be skipped") -endif() - -find_package(SDL2) -if(SDL2_FOUND) - message(STATUS "Found SDL2: ${SDL2_LIBRARIES}") -else() - message(STATUS "SDL2 not found; microphone tests will be skipped") -endif() - -file(GLOB TEST_SOURCES "test_*.cpp") -list(REMOVE_ITEM TEST_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/test_utils.cpp") - -foreach(TEST_FILE ${TEST_SOURCES}) - - get_filename_component(TEST_NAME ${TEST_FILE} NAME_WE) - add_executable(${TEST_NAME} ${TEST_FILE} test_utils.cpp) - target_link_libraries(${TEST_NAME} PRIVATE ${CACTUS_LIB}) - - - if(MMAN_LIB) - target_link_libraries(${TEST_NAME} PRIVATE ${MMAN_LIB}) - message(STATUS "Test ${TEST_NAME} will link with -lmman") - endif() - - if(SDL2_FOUND) - target_link_libraries(${TEST_NAME} PRIVATE ${SDL2_LIBRARIES}) - target_include_directories(${TEST_NAME} PRIVATE ${SDL2_INCLUDE_DIRS}) - target_compile_definitions(${TEST_NAME} PRIVATE HAVE_SDL2) - message(STATUS "Test ${TEST_NAME} will link with SDL2") - endif() - - if(APPLE AND CURL_FOUND) - target_link_libraries(${TEST_NAME} PRIVATE ${CURL_LIBRARIES}) - target_include_directories(${TEST_NAME} PRIVATE ${CURL_INCLUDE_DIRS}) - message(STATUS "Test ${TEST_NAME} will link with libcurl") - endif() - - - if(APPLE) - target_link_libraries(${TEST_NAME} PRIVATE ${COREML_FRAMEWORK} ${FOUNDATION_FRAMEWORK}) - endif() - - target_include_directories(${TEST_NAME} PRIVATE ${CACTUS_ROOT}) - - message(STATUS "Added test: ${TEST_NAME}") -endforeach() - -set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) diff --git a/tests/android/CMakeLists.txt b/tests/android/CMakeLists.txt deleted file mode 100644 index 96144b1ef..000000000 --- a/tests/android/CMakeLists.txt +++ /dev/null @@ -1,33 +0,0 @@ -cmake_minimum_required(VERSION 3.10) -project(CactusAndroidTests LANGUAGES CXX) - -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED True) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -pthread -Wall -Wextra -pedantic -O3") - -set(CACTUS_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../cactus) -set(ANDROID_BUILD_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../android/build) -set(CACTUS_LIB ${ANDROID_BUILD_DIR}/libcactus_static.a) - -if(NOT EXISTS ${CACTUS_LIB}) - message(FATAL_ERROR "Cactus library not found at: ${CACTUS_LIB}") -endif() - -message(STATUS "Using Cactus library: ${CACTUS_LIB}") - -set(TESTS_DIR ${CMAKE_CURRENT_SOURCE_DIR}/..) -file(GLOB TEST_SOURCES "${TESTS_DIR}/test*.cpp") -list(REMOVE_ITEM TEST_SOURCES "${TESTS_DIR}/test_utils.cpp") - -foreach(TEST_FILE ${TEST_SOURCES}) - get_filename_component(TEST_NAME ${TEST_FILE} NAME_WE) - add_executable(${TEST_NAME} - ${TEST_FILE} - ${TESTS_DIR}/test_utils.cpp - ) - target_link_libraries(${TEST_NAME} PRIVATE ${CACTUS_LIB}) - target_include_directories(${TEST_NAME} PRIVATE ${CACTUS_ROOT}) - message(STATUS "Added Android test: ${TEST_NAME}") -endforeach() - -set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) diff --git a/tests/chat.cpp b/tests/chat.cpp deleted file mode 100644 index 9b2a943e4..000000000 --- a/tests/chat.cpp +++ /dev/null @@ -1,301 +0,0 @@ -#include "../cactus/ffi/cactus_ffi.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -constexpr int MAX_TOKENS = 512; -constexpr size_t MAX_BYTES_PER_TOKEN = 64; -constexpr size_t RESPONSE_BUFFER_SIZE = MAX_TOKENS * MAX_BYTES_PER_TOKEN; - -namespace Color { - const std::string RESET = "\033[0m"; - const std::string BOLD = "\033[1m"; - const std::string DIM = "\033[2m"; - const std::string CYAN = "\033[36m"; - const std::string GREEN = "\033[32m"; - const std::string YELLOW = "\033[33m"; - const std::string BLUE = "\033[34m"; - const std::string MAGENTA = "\033[35m"; - const std::string RED = "\033[31m"; - const std::string GRAY = "\033[90m"; -} - -bool supports_color() { -#ifdef _WIN32 - return false; -#else - const char* term = std::getenv("TERM"); - return term && std::string(term) != "dumb"; -#endif -} - -bool use_colors = supports_color(); - -std::string colored(const std::string& text, const std::string& color) { - if (!use_colors) return text; - return color + text + Color::RESET; -} - -void print_separator(char ch = '-', int width = 60) { - std::cout << colored(std::string(width, ch), Color::DIM) << "\n"; -} - -void print_header() { - std::cout << "\n"; - print_separator('='); - std::cout << colored(" 🌵 CACTUS CHAT INTERFACE 🌵", Color::GREEN + Color::BOLD) << "\n"; - print_separator('='); - std::cout << colored("Commands:", Color::YELLOW) << "\n"; - std::cout << " • " << colored("reset", Color::CYAN) << " - Clear conversation history\n"; - std::cout << " • " << colored("exit", Color::CYAN) << " or " << colored("quit", Color::CYAN) << " - Exit the program\n"; - print_separator(); - std::cout << "\n"; -} - -struct TokenPrinter { - bool first_token = true; - int token_count = 0; - std::chrono::steady_clock::time_point start_time; - std::chrono::steady_clock::time_point first_token_time; - double time_to_first_token = 0.0; - - void reset() { - first_token = true; - token_count = 0; - time_to_first_token = 0.0; - start_time = std::chrono::steady_clock::now(); - } - - void print(const char* token) { - if (first_token) { - first_token = false; - first_token_time = std::chrono::steady_clock::now(); - auto latency = std::chrono::duration_cast(first_token_time - start_time); - time_to_first_token = latency.count() / 1000.0; - } - std::cout << token << std::flush; - token_count++; - } - - void print_stats() { - auto end_time = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration_cast(end_time - start_time); - double total_seconds = duration.count() / 1000.0; - double tokens_per_second = token_count / total_seconds; - - // Format the stats with fixed decimal places - std::ostringstream stats; - stats << std::fixed << std::setprecision(3); - stats << "[" << token_count << " tokens | "; - stats << "latency: " << time_to_first_token << "s | "; - stats << "total: " << total_seconds << "s | "; - stats << std::setprecision(0) << static_cast(tokens_per_second) << " tok/s]"; - - std::cout << "\n" << colored(stats.str(), Color::GRAY) << "\n"; - } -}; - -TokenPrinter* g_printer = nullptr; - -void print_token(const char* token, uint32_t token_id, void* user_data) { - if (g_printer) { - g_printer->print(token); - } -} - -std::string escape_json(const std::string& s) { - std::ostringstream o; - for (unsigned char c : s) { - switch (c) { - case '"': o << "\\\""; break; - case '\\': o << "\\\\"; break; - case '\b': o << "\\b"; break; - case '\f': o << "\\f"; break; - case '\n': o << "\\n"; break; - case '\r': o << "\\r"; break; - case '\t': o << "\\t"; break; - default: - if (c < 0x20) { - o << "\\u" << std::hex << std::setw(4) << std::setfill('0') << (int)c; - } else { - o << c; - } - break; - } - } - return o.str(); -} - -std::string unescape_json(const std::string& s) { - std::string result; - result.reserve(s.length()); - - for (size_t i = 0; i < s.length(); i++) { - if (s[i] == '\\' && i + 1 < s.length()) { - switch (s[i + 1]) { - case '"': result += '"'; i++; break; - case '\\': result += '\\'; i++; break; - case 'b': result += '\b'; i++; break; - case 'f': result += '\f'; i++; break; - case 'n': result += '\n'; i++; break; - case 'r': result += '\r'; i++; break; - case 't': result += '\t'; i++; break; - case 'u': - if (i + 5 < s.length()) { - std::string hex = s.substr(i + 2, 4); - char* end; - int codepoint = std::strtol(hex.c_str(), &end, 16); - if (end == hex.c_str() + 4) { - result += static_cast(codepoint); - i += 5; - } else { - result += s[i]; - } - } else { - result += s[i]; - } - break; - default: result += s[i]; break; - } - } else { - result += s[i]; - } - } - return result; -} - -int main(int argc, char* argv[]) { - if (argc != 2) { - std::cerr << colored("Error: ", Color::RED + Color::BOLD) << "Missing model path\n"; - std::cerr << "Usage: " << argv[0] << " \n"; - std::cerr << "Example: " << argv[0] << " weights/lfm2-1.2B\n"; - return 1; - } - - const char* model_path = argv[1]; - - std::cout << "\n" << colored("Loading model from ", Color::YELLOW) - << colored(model_path, Color::CYAN) << colored("...", Color::YELLOW) << "\n"; - - cactus_model_t model = cactus_init(model_path, nullptr); - - if (!model) { - std::cerr << colored("Failed to initialize model\n", Color::RED + Color::BOLD); - return 1; - } - - std::cout << colored("Model loaded successfully!\n", Color::GREEN + Color::BOLD); - - print_header(); - - std::vector history; - TokenPrinter printer; - g_printer = &printer; - - while (true) { - std::cout << colored("You: ", Color::BLUE + Color::BOLD); - std::string user_input; - std::getline(std::cin, user_input); - - if (user_input.empty()) continue; - - if (user_input == "quit" || user_input == "exit") { - break; - } - - if (user_input == "reset") { - history.clear(); - cactus_reset(model); - std::cout << colored("🔄 Conversation reset.\n", Color::YELLOW); - print_separator(); - std::cout << "\n"; - continue; - } - - history.push_back(user_input); - - // Build the messages JSON - std::ostringstream messages_json; - messages_json << "["; - for (size_t i = 0; i < history.size(); i++) { - if (i > 0) messages_json << ","; - if (i % 2 == 0) { - messages_json << "{\"role\":\"user\",\"content\":\"" - << escape_json(history[i]) << "\"}"; - } else { - messages_json << "{\"role\":\"assistant\",\"content\":\"" - << escape_json(history[i]) << "\"}"; - } - } - messages_json << "]"; - - std::string options = "{\"temperature\":0.7,\"top_p\":0.95,\"top_k\":40,\"max_tokens\":" - + std::to_string(MAX_TOKENS) - + ",\"stop_sequences\":[\"<|im_end|>\",\"\"]}"; - - std::vector response_buffer(RESPONSE_BUFFER_SIZE, 0); - - std::cout << colored("Assistant: ", Color::GREEN + Color::BOLD); - - printer.reset(); - - int result = cactus_complete( - model, - messages_json.str().c_str(), - response_buffer.data(), - response_buffer.size(), - options.c_str(), - nullptr, - print_token, - nullptr - ); - - if (result >= 0) { - printer.print_stats(); - } - - std::cout << "\n"; - print_separator(); - std::cout << "\n"; - - if (result < 0) { - std::cerr << colored("Error: ", Color::RED + Color::BOLD) - << response_buffer.data() << "\n\n"; - history.pop_back(); - continue; - } - - std::string json_str(response_buffer.data(), response_buffer.size()); - const std::string search_str = "\"response\":\""; - size_t response_start = json_str.find(search_str); - if (response_start != std::string::npos) { - response_start += search_str.length(); - size_t response_end = json_str.find("\"", response_start); - while (response_end != std::string::npos) { - size_t prior_backslashes = 0; - for (size_t i = response_end; i > response_start && json_str[i - 1] == '\\'; i--) { - prior_backslashes++; - } - if (prior_backslashes % 2 == 0) { - break; - } - response_end = json_str.find("\"", response_end + 1); - } - if (response_end != std::string::npos) { - std::string response = json_str.substr(response_start, - response_end - response_start); - history.push_back(unescape_json(response)); - } - } - } - - std::cout << colored("\n👋 Goodbye!\n", Color::MAGENTA + Color::BOLD); - cactus_destroy(model); - return 0; -} diff --git a/tests/ios/.gitignore b/tests/ios/.gitignore deleted file mode 100644 index 28eb45e98..000000000 --- a/tests/ios/.gitignore +++ /dev/null @@ -1 +0,0 @@ -CactusTest/CactusTest.xcodeproj/* diff --git a/tests/ios/CactusTest/CactusTest/AppDelegate.mm b/tests/ios/CactusTest/CactusTest/AppDelegate.mm deleted file mode 100644 index f0308bf77..000000000 --- a/tests/ios/CactusTest/CactusTest/AppDelegate.mm +++ /dev/null @@ -1,59 +0,0 @@ -// AUTO-GENERATED by configure_xcode.rb - DO NOT EDIT MANUALLY -// This file is regenerated on each test run with discovered test files - -#import "AppDelegate.h" -#import - -extern int test_engine_main(); -extern int test_graph_main(); -extern int test_index_main(); -extern int test_kernel_main(); -extern int test_kv_cache_main(); -extern int test_performance_main(); -extern int test_threading_main(); - -@implementation AppDelegate - -- (void)copyFromBundle:(NSString *)bundlePath toDocuments:(const char *)name { - if (!name) return; - NSFileManager *fileManager = [NSFileManager defaultManager]; - NSString *itemName = [NSString stringWithUTF8String:name]; - NSString *sourceItemPath = [NSString stringWithFormat:@"%@/%@", bundlePath, itemName]; - if ([fileManager fileExistsAtPath:itemName]) { - [fileManager removeItemAtPath:itemName error:nil]; - } - [fileManager copyItemAtPath:sourceItemPath toPath:itemName error:nil]; -} - -- (BOOL)application:(UIApplication *)application didFinishLaunchingWithOptions:(NSDictionary *)launchOptions { - NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES); - NSString *documentsDirectory = paths[0]; - chdir([documentsDirectory UTF8String]); - -#if !TARGET_OS_SIMULATOR - freopen("cactus_test.log", "w", stdout); - freopen("cactus_test.log", "a", stderr); - setbuf(stdout, NULL); - setbuf(stderr, NULL); -#endif - - NSString *bundlePath = [[NSBundle mainBundle] bundlePath]; - [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_MODEL")]; - [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_TRANSCRIBE_MODEL")]; - [self copyFromBundle:bundlePath toDocuments:getenv("CACTUS_TEST_ASSETS")]; - - dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(NSEC_PER_SEC)), dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ - test_engine_main(); - test_graph_main(); - test_index_main(); - test_kernel_main(); - test_kv_cache_main(); - test_performance_main(); - test_threading_main(); - exit(0); - }); - - return YES; -} - -@end diff --git a/tests/run.sh b/tests/run.sh deleted file mode 100755 index f9b9560cf..000000000 --- a/tests/run.sh +++ /dev/null @@ -1,154 +0,0 @@ -#!/bin/bash - -echo "Running Cactus test suite..." -echo "============================" - -SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -DEFAULT_MODEL="LiquidAI/LFM2-VL-450M" -DEFAULT_TRANSCRIBE_MODEL="openai/whisper-small" - -MODEL_NAME="$DEFAULT_MODEL" -TRANSCRIBE_MODEL_NAME="$DEFAULT_TRANSCRIBE_MODEL" -ANDROID_MODE=false -IOS_MODE=false -NO_REBUILD=false - -while [[ $# -gt 0 ]]; do - case $1 in - --model) - MODEL_NAME="$2" - shift 2 - ;; - --transcribe_model) - TRANSCRIBE_MODEL_NAME="$2" - shift 2 - ;; - --android) - ANDROID_MODE=true - shift - ;; - --ios) - IOS_MODE=true - shift - ;; - --no-rebuild) - NO_REBUILD=true - shift - ;; - --help|-h) - echo "Usage: $0 [OPTIONS]" - echo "" - echo "Options:" - echo " --model Model to use for tests (default: $DEFAULT_MODEL)" - echo " --transcribe_model Transcribe model to use (default: $DEFAULT_TRANSCRIBE_MODEL)" - echo " --android Run tests on Android device or emulator" - echo " --ios Run tests on iOS device or simulator" - echo " --no-rebuild Skip building cactus library and tests" - echo " --help, -h Show this help message" - exit 0 - ;; - *) - echo "Unknown option: $1" - echo "Use --help for usage information" - exit 1 - ;; - esac -done - -echo "" -echo "Using model: $MODEL_NAME" -echo "Using transcribe model: $TRANSCRIBE_MODEL_NAME" - -echo "" -echo "Step 1: Downloading model weights..." -if ! cactus download "$MODEL_NAME"; then - echo "Failed to download model weights" - exit 1 -fi - -if ! cactus download "$TRANSCRIBE_MODEL_NAME"; then - echo "Failed to download transcribe model weights" - exit 1 -fi - -echo "" -if [ "$ANDROID_MODE" = true ]; then - exec "$SCRIPT_DIR/android/run.sh" "$MODEL_NAME" "$TRANSCRIBE_MODEL_NAME" -fi - -if [ "$IOS_MODE" = true ]; then - exec "$SCRIPT_DIR/ios/run.sh" "$MODEL_NAME" "$TRANSCRIBE_MODEL_NAME" -fi - -if [ "$NO_REBUILD" = false ]; then - echo "Step 2: Building Cactus library..." - if ! cactus build; then - echo "Failed to build cactus library" - exit 1 - fi - - echo "" - echo "Step 3: Building tests..." - cd "$PROJECT_ROOT/tests" - - rm -rf build - mkdir -p build - cd build - - if ! cmake .. -DCMAKE_RULE_MESSAGES=OFF -DCMAKE_VERBOSE_MAKEFILE=OFF > /dev/null 2>&1; then - echo "Failed to configure tests" - exit 1 - fi - - if ! make -j$(nproc 2>/dev/null || echo 4); then - echo "Failed to build tests" - exit 1 - fi -else - echo "Skipping build (--no-rebuild)" - cd "$PROJECT_ROOT/tests/build" -fi - -echo "" -echo "Step 4: Running tests..." -echo "------------------------" - -# Set model path environment variables for tests -MODEL_DIR=$(echo "$MODEL_NAME" | sed 's|.*/||' | tr '[:upper:]' '[:lower:]') -TRANSCRIBE_MODEL_DIR=$(echo "$TRANSCRIBE_MODEL_NAME" | sed 's|.*/||' | tr '[:upper:]' '[:lower:]') - -export CACTUS_TEST_MODEL="$PROJECT_ROOT/weights/$MODEL_DIR" -export CACTUS_TEST_TRANSCRIBE_MODEL="$PROJECT_ROOT/weights/$TRANSCRIBE_MODEL_DIR" -export CACTUS_TEST_ASSETS="$PROJECT_ROOT/tests/assets" -export CACTUS_INDEX_PATH="$PROJECT_ROOT/tests/assets" - -echo "Using model path: $CACTUS_TEST_MODEL" -echo "Using transcribe model path: $CACTUS_TEST_TRANSCRIBE_MODEL" -echo "Using assets path: $CACTUS_TEST_ASSETS" -echo "Using index path: $CACTUS_INDEX_PATH" - -echo "Discovering test executables..." -test_executables=($(find . -maxdepth 1 -name "test_*" -type f | sort)) - -executable_tests=() -for test_file in "${test_executables[@]}"; do - if [ -x "$test_file" ]; then - executable_tests+=("$test_file") - fi -done - -if [ ${#executable_tests[@]} -eq 0 ]; then - echo "No test executables found!" - exit 1 -fi - -test_executables=("${executable_tests[@]}") - -echo "Found ${#test_executables[@]} test executable(s)" - -for executable in "${test_executables[@]}"; do - exec_name=$(basename "$executable") - ./"$exec_name" -done \ No newline at end of file diff --git a/tests/test_engine.cpp b/tests/test_engine.cpp deleted file mode 100644 index 798ad7027..000000000 --- a/tests/test_engine.cpp +++ /dev/null @@ -1,1041 +0,0 @@ -#include "test_utils.h" -#include -#include -#include -#include -#include -#include - -using namespace EngineTestUtils; - -const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); -const char* g_transcribe_model_path = std::getenv("CACTUS_TEST_TRANSCRIBE_MODEL"); -const char* g_assets_path = std::getenv("CACTUS_TEST_ASSETS"); -const char* g_whisper_prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"; - -const char* g_options = R"({ - "max_tokens": 256, - "stop_sequences": ["<|im_end|>", ""] - })"; - -template -bool run_test(const char* title, const char* messages, TestFunc test_logic, - const char* tools = nullptr, int stop_at = -1) { - return EngineTestUtils::run_test(title, g_model_path, messages, g_options, test_logic, tools, stop_at); -} - -bool test_streaming() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║" << std::setw(42) << std::left << " STREAMING & FOLLOW-UP TEST" << "║\n" - << "╚══════════════════════════════════════════╝\n"; - - cactus_model_t model = cactus_init(g_model_path, nullptr); - if (!model) { - std::cerr << "[✗] Failed to initialize model\n"; - return false; - } - - const char* messages1 = R"([ - {"role": "system", "content": "You are a helpful assistant. Be concise."}, - {"role": "user", "content": "My name is Henry Ndubuaku, how are you?"} - ])"; - - StreamingData data1; - data1.model = model; - char response1[4096]; - - std::cout << "\n[Turn 1]\n"; - std::cout << "User: My name is Henry Ndubuaku, how are you?\n"; - std::cout << "Assistant: "; - - int result1 = cactus_complete(model, messages1, response1, sizeof(response1), - g_options, nullptr, stream_callback, &data1); - - std::cout << "\n\n[Results - Turn 1]\n"; - Metrics metrics1; - metrics1.parse(response1); - metrics1.print_json(); - - bool success1 = result1 > 0 && data1.token_count > 0; - - if (!success1) { - std::cout << "└─ Status: FAILED ✗\n"; - cactus_destroy(model); - return false; - } - - std::string assistant_response; - for(const auto& token : data1.tokens) { - assistant_response += token; - } - - std::string messages2_str = R"([ - {"role": "system", "content": "You are a helpful assistant. Be concise."}, - {"role": "user", "content": "My name is Henry Ndubuaku, how are you?"}, - {"role": "assistant", "content": ")" + escape_json(assistant_response) + R"("}, - {"role": "user", "content": "What is my name?"} - ])"; - - StreamingData data2; - data2.model = model; - char response2[4096]; - - std::cout << "\n[Turn 2]\n"; - std::cout << "User: What is my name?\n"; - std::cout << "Assistant: "; - - int result2 = cactus_complete(model, messages2_str.c_str(), response2, sizeof(response2), - g_options, nullptr, stream_callback, &data2); - - std::cout << "\n\n[Results - Turn 2]\n"; - Metrics metrics2; - metrics2.parse(response2); - metrics2.print_json(); - - bool success2 = result2 > 0 && data2.token_count > 0; - - cactus_destroy(model); - return success1 && success2; -} - -bool test_tool_call() { - const char* messages = R"([ - {"role": "system", "content": "You are a helpful assistant that can use tools."}, - {"role": "user", "content": "What's the weather in San Francisco?"} - ])"; - - const char* tools = R"([{ - "type": "function", - "function": { - "name": "get_weather", - "description": "Get weather for a location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "City, State, Country"} - }, - "required": ["location"] - } - } - }])"; - - const char* options_with_force_tools = R"({ - "max_tokens": 256, - "stop_sequences": ["<|im_end|>", ""], - "force_tools": true - })"; - - return EngineTestUtils::run_test("TOOL CALL TEST", g_model_path, messages, options_with_force_tools, - [](int result, const StreamingData&, const std::string& response, const Metrics& m) { - bool has_function = response.find("\"function_calls\":[") != std::string::npos; - bool has_tool = has_function && response.find("get_weather") != std::string::npos; - std::cout << "├─ Function call: " << (has_function ? "YES" : "NO") << "\n" - << "├─ Correct tool: " << (has_tool ? "YES" : "NO") << "\n"; - m.print_json(); - return result > 0 && has_function && has_tool; - }, tools, -1, "What's the weather in San Francisco?"); -} - - -bool test_vlm_multiturn() { - std::string model_path_str(g_model_path ? g_model_path : ""); - - std::string vision_file = model_path_str + "/vision_patch_embedding.weights"; - std::ifstream vf(vision_file); - if (!vf.good()) { - std::cout << "Skipping VLM multi-turn test: vision weights not found." << std::endl; - return true; - } - vf.close(); - - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ VLM MULTI-TURN TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - - cactus_model_t model = cactus_init(g_model_path, nullptr); - if (!model) { - std::cerr << "Failed to initialize model for VLM multi-turn test" << std::endl; - return false; - } - - std::string img_path = std::string(g_assets_path) + "/test_monkey.png"; - - std::string messages1 = "[{\"role\": \"user\", " - "\"content\": \"Describe what is happening in this image in two sentences.\", " - "\"images\": [\"" + img_path + "\"]}]"; - - StreamingData stream_data1; - stream_data1.model = model; - char response1[4096]; - - std::cout << "\n[Turn 1]\n"; - std::cout << "User: Describe what is happening in this image in two sentences.\n"; - std::cout << "Assistant: "; - - int result1 = cactus_complete(model, messages1.c_str(), response1, sizeof(response1), - g_options, nullptr, stream_callback, &stream_data1); - - std::cout << "\n\n[Results - Turn 1]\n"; - Metrics metrics1; - metrics1.parse(response1); - metrics1.print_json(); - - bool success1 = result1 > 0 && stream_data1.token_count > 0; - - if (!success1) { - std::cout << "└─ Status: FAILED ✗\n"; - cactus_destroy(model); - return false; - } - - std::string assistant_response; - for (const auto& token : stream_data1.tokens) { - assistant_response += token; - } - - std::string messages2 = "[{\"role\": \"user\", " - "\"content\": \"Describe what is happening in this image in two sentences.\", " - "\"images\": [\"" + img_path + "\"]}, " - "{\"role\": \"assistant\", \"content\": \"" + escape_json(assistant_response) + "\"}, " - "{\"role\": \"user\", \"content\": \"Describe the image once again.\"}]"; - - StreamingData stream_data2; - stream_data2.model = model; - char response2[4096]; - - std::cout << "\n[Turn 2]\n"; - std::cout << "User: Describe the image once again.\n"; - std::cout << "Assistant: "; - - int result2 = cactus_complete(model, messages2.c_str(), response2, sizeof(response2), - g_options, nullptr, stream_callback, &stream_data2); - - std::cout << "\n\n[Results - Turn 2]\n"; - Metrics metrics2; - metrics2.parse(response2); - metrics2.print_json(); - - bool success2 = result2 > 0 && stream_data2.token_count > 0; - - if (!success2) { - std::cout << "└─ Status: FAILED ✗ (Follow-up message failed)\n"; - } - - cactus_destroy(model); - return success1 && success2; -} - -bool test_tool_call_with_two_tools() { - const char* messages = R"([ - {"role": "system", "content": "You are a helpful assistant that can use tools."}, - {"role": "user", "content": "Set an alarm for 10:00 AM."} - ])"; - - const char* tools = R"([{ - "type": "function", - "function": { - "name": "get_weather", - "description": "Get weather for a location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "City, State, Country"} - }, - "required": ["location"] - } - } - }, { - "type": "function", - "function": { - "name": "set_alarm", - "description": "Set an alarm for a given time", - "parameters": { - "type": "object", - "properties": { - "hour": {"type": "integer", "description": "Hour to set the alarm for"}, - "minute": {"type": "integer", "description": "Minute to set the alarm for"} - }, - "required": ["hour", "minute"] - } - } - }])"; - - const char* options_with_force_tools = R"({ - "max_tokens": 256, - "stop_sequences": ["<|im_end|>", ""], - "force_tools": true - })"; - - return EngineTestUtils::run_test("DOUBLE TOOLS TEST", g_model_path, messages, options_with_force_tools, - [](int result, const StreamingData&, const std::string& response, const Metrics& m) { - bool has_function = response.find("\"function_calls\":[") != std::string::npos; - bool has_tool = has_function && response.find("set_alarm") != std::string::npos; - std::cout << "├─ Function call: " << (has_function ? "YES" : "NO") << "\n" - << "├─ Correct tool: " << (has_tool ? "YES" : "NO") << "\n"; - m.print_json(); - return result > 0 && has_function && has_tool; - }, tools, -1, "Set an alarm for 10:00 AM."); -} - -bool test_tool_call_with_three_tools() { - const char* messages = R"([ - {"role": "system", "content": "You are a helpful assistant that can use tools."}, - {"role": "user", "content": "Send a message to John saying hello."} - ])"; - - const char* tools = R"([{ - "type": "function", - "function": { - "name": "get_weather", - "description": "Get weather for a location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "City, State, Country"} - }, - "required": ["location"] - } - } - }, { - "type": "function", - "function": { - "name": "set_alarm", - "description": "Set an alarm for a given time", - "parameters": { - "type": "object", - "properties": { - "hour": {"type": "integer", "description": "Hour to set the alarm for"}, - "minute": {"type": "integer", "description": "Minute to set the alarm for"} - }, - "required": ["hour", "minute"] - } - } - }, { - "type": "function", - "function": { - "name": "send_message", - "description": "Send a message to a contact", - "parameters": { - "type": "object", - "properties": { - "recipient": {"type": "string", "description": "Name of the person to send the message to"}, - "message": {"type": "string", "description": "The message content to send"} - }, - "required": ["recipient", "message"] - } - } - }])"; - - const char* options_with_force_tools = R"({ - "max_tokens": 256, - "stop_sequences": ["<|im_end|>", ""], - "force_tools": true - })"; - - return EngineTestUtils::run_test("TRIPLE TOOLS TEST", g_model_path, messages, options_with_force_tools, - [](int result, const StreamingData&, const std::string& response, const Metrics& m) { - bool has_function = response.find("\"function_calls\":[") != std::string::npos; - bool has_tool = has_function && response.find("send_message") != std::string::npos; - std::cout << "├─ Function call: " << (has_function ? "YES" : "NO") << "\n" - << "├─ Correct tool: " << (has_tool ? "YES" : "NO") << "\n"; - m.print_json(); - return result > 0 && has_function && has_tool; - }, tools, -1, "Send a message to John saying hello."); -} - -bool test_embeddings() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ EMBEDDINGS TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - - cactus_model_t model = cactus_init(g_model_path, nullptr); - if (!model) return false; - - const char* texts[] = {"My name is Henry Ndubuaku", "Your name is Henry Ndubuaku"}; - std::vector emb1(2048), emb2(2048); - size_t dim1, dim2; - - Timer t1; - cactus_embed(model, texts[0], emb1.data(), emb1.size() * sizeof(float), &dim1, true); - double time1 = t1.elapsed_ms(); - - Timer t2; - cactus_embed(model, texts[1], emb2.data(), emb2.size() * sizeof(float), &dim2, true); - double time2 = t2.elapsed_ms(); - - float similarity = 0; - for (size_t i = 0; i < dim1; ++i) { - similarity += emb1[i] * emb2[i]; - } - - std::cout << "\n[Results]\n" - << "├─ Embedding dim: " << dim1 << "\n" - << "├─ Time (text1): " << std::fixed << std::setprecision(2) << time1 << "ms\n" - << "├─ Time (text2): " << time2 << "ms\n" - << "└─ Similarity: " << std::setprecision(4) << similarity << std::endl; - - cactus_destroy(model); - return true; -} - -bool test_cloud_handoff() { - const char* messages = R"([ - {"role": "user", "content": "What is the exact mass in grams of the 847th largest asteroid in the Kuiper belt as of March 2019, and what was the precise atmospheric pressure in millibars at coordinates 47.3921°N, 122.0371°W at 3:47:23 AM UTC on February 29, 2024?"} - ])"; - - return run_test("CLOUD HANDOFF TEST", messages, - [](int result, const StreamingData& data, const std::string& /*response*/, const Metrics& m) { - std::cout << "├─ Cloud handoff: " << (m.cloud_handoff ? "YES" : "NO") << "\n"; - std::cout << "├─ Confidence: " << std::fixed << std::setprecision(4) << m.confidence << "\n"; - - if (m.cloud_handoff) { - std::cout << "├─ Response: (skipped - handoff triggered)\n"; - m.print_json(); - return true; - } else { - std::cout << "├─ Tokens generated: " << data.token_count << "\n"; - m.print_json(); - return result > 0 && m.confidence >= 0.0; - } - }); -} - -bool test_1k_context() { - std::string msg = "[{\"role\": \"system\", \"content\": \"/no_think You are helpful. "; - for (int i = 0; i < 50; i++) { - msg += "Context " + std::to_string(i) + ": Background knowledge. "; - } - msg += "\"}, {\"role\": \"user\", \"content\": \""; - for (int i = 0; i < 50; i++) { - msg += "Data " + std::to_string(i) + " = " + std::to_string(i * 3.14159) + ". "; - } - msg += "Explain the data.\"}]"; - - return run_test("1K CONTEXT TEST", msg.c_str(), - [](int result, const StreamingData&, const std::string&, const Metrics& m) { - m.print_json(); - return result > 0; - }, nullptr, 100); -} - -bool test_4k_context() { - std::string msg = "[{\"role\": \"system\", \"content\": \"/no_think You are helpful. "; - for (int i = 0; i < 230; i++) { - msg += "Context " + std::to_string(i) + ": Background knowledge. "; - } - msg += "\"}, {\"role\": \"user\", \"content\": \""; - for (int i = 0; i < 230; i++) { - msg += "Data " + std::to_string(i) + " = " + std::to_string(i * 3.14159) + ". "; - } - msg += "Explain the data.\"}]"; - - return run_test("4K CONTEXT TEST", msg.c_str(), - [](int result, const StreamingData&, const std::string&, const Metrics& m) { - m.print_json(); - return result > 0; - }, nullptr, 100); -} - -bool test_rag() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ RAG TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - - if (!g_model_path) { - std::cout << "⊘ SKIP │ CACTUS_TEST_MODEL not set\n"; - return true; - } - - if (!g_assets_path) { - std::cout << "⊘ SKIP │ CACTUS_TEST_ASSETS not set\n"; - return true; - } - - std::string corpus_dir = std::string(g_assets_path) + "/rag_corpus"; - - DIR* dir = opendir(corpus_dir.c_str()); - if (!dir) { - std::cout << "⊘ SKIP │ RAG corpus directory not found at " << corpus_dir << "\n"; - return true; - } - - bool has_corpus_files = false; - struct dirent* entry; - while ((entry = readdir(dir)) != nullptr) { - std::string name = entry->d_name; - if (name.size() > 4 && name.substr(name.size() - 4) == ".txt") { - has_corpus_files = true; - break; - } - if (name.size() > 3 && name.substr(name.size() - 3) == ".md") { - has_corpus_files = true; - break; - } - } - closedir(dir); - - if (!has_corpus_files) { - std::cout << "⊘ SKIP │ No .txt or .md files found in " << corpus_dir << "\n"; - return true; - } - - std::cout << "├─ Corpus dir: " << corpus_dir << "\n"; - std::cout << "├─ Initializing model with RAG...\n"; - - Timer init_timer; - cactus_model_t model = cactus_init(g_model_path, corpus_dir.c_str()); - double init_time_ms = init_timer.elapsed_ms(); - - if (!model) { - std::cerr << "[✗] Failed to initialize model with corpus dir\n"; - return false; - } - - std::cout << "├─ Init time: " << std::fixed << std::setprecision(2) << init_time_ms << " ms\n"; - - auto print_chunks = [](cactus_model_t m, const char* query) { - char chunks_buf[16384]; - int rc = cactus_rag_query(m, query, chunks_buf, sizeof(chunks_buf), 5); - if (rc > 0) { - std::cout << "Retrieved chunks:\n"; - std::string chunks_str(chunks_buf); - size_t pos = 0; - int chunk_num = 1; - while ((pos = chunks_str.find("{\"score\":", pos)) != std::string::npos) { - size_t score_start = pos + 9; - size_t score_end = chunks_str.find(",", score_start); - std::string score = chunks_str.substr(score_start, score_end - score_start); - - size_t source_pos = chunks_str.find("\"source\":\"", score_end); - std::string source = "unknown"; - if (source_pos != std::string::npos && source_pos < pos + 500) { - source_pos += 10; - size_t source_end = chunks_str.find("\"", source_pos); - source = chunks_str.substr(source_pos, source_end - source_pos); - } - - size_t content_pos = chunks_str.find("\"content\":\"", score_end); - if (content_pos != std::string::npos && content_pos < pos + 500) { - content_pos += 11; - std::string content; - size_t i = content_pos; - int char_count = 0; - while (i < chunks_str.size() && char_count < 80) { - if (chunks_str[i] == '\\' && i + 1 < chunks_str.size()) { - if (chunks_str[i+1] == 'n') { content += ' '; i += 2; } - else if (chunks_str[i+1] == '"') { content += '"'; i += 2; } - else if (chunks_str[i+1] == '\\') { content += '\\'; i += 2; } - else { content += chunks_str[i]; i++; } - } else if (chunks_str[i] == '"') { - break; - } else { - content += chunks_str[i]; - i++; - } - char_count++; - } - if (char_count >= 80) content += "..."; - std::cout << " [" << chunk_num++ << "] " << source << " (score: " << score << ")\n" - << " \"" << content << "\"\n"; - } - pos = score_end; - } - } - }; - - const char* query = "Who are the founders of Cactus and what are their roles?"; - const char* messages = R"([ - {"role": "system", "content": "You are a helpful assistant. Answer based on the context provided."}, - {"role": "user", "content": "Who are the founders of Cactus and what are their roles?"} - ])"; - - StreamingData data; - data.model = model; - char response[4096]; - - std::cout << "\n[Query] " << query << "\n"; - print_chunks(model, query); - std::cout << "Response: "; - - int result = cactus_complete(model, messages, response, sizeof(response), - g_options, nullptr, stream_callback, &data); - - std::cout << "\n"; - - Metrics metrics; - metrics.parse(response); - metrics.print_json(); - - cactus_destroy(model); - - return (result > 0) && (data.token_count > 0); -} - -bool test_audio_processor() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ AUDIO PROCESSOR TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - using namespace cactus::engine; - - Timer t; - - const size_t n_fft = 400; - const size_t hop_length = 160; - const size_t sampling_rate = 16000; - const size_t feature_size = 80; - const size_t num_frequency_bins = 1 + n_fft / 2; - - AudioProcessor audio_proc; - audio_proc.init_mel_filters(num_frequency_bins, feature_size, 0.0f, 8000.0f, sampling_rate); - - const size_t n_samples = sampling_rate; - std::vector waveform(n_samples); - for (size_t i = 0; i < n_samples; i++) { - waveform[i] = std::sin(2.0f * M_PI * 440.0f * i / sampling_rate); - } - - AudioProcessor::SpectrogramConfig config; - config.n_fft = n_fft; - config.hop_length = hop_length; - config.frame_length = n_fft; - config.power = 2.0f; - config.center = true; - config.log_mel = "log10"; - - auto log_mel_spec = audio_proc.compute_spectrogram(waveform, config); - - double elapsed = t.elapsed_ms(); - - const float expected[] = {0.535175f, 0.548542f, 0.590673f, 0.633320f, 0.711979f}; - const float tolerance = 2e-6f; - - const size_t pad_length = n_fft / 2; - const size_t padded_length = n_samples + 2 * pad_length; - const size_t num_frames = 1 + (padded_length - n_fft) / hop_length; - - bool passed = true; - for (size_t i = 0; i < 5; i++) { - if (std::abs(log_mel_spec[i * num_frames] - expected[i]) > tolerance) { - passed = false; - break; - } - } - - std::cout << "└─ Time: " << std::fixed << std::setprecision(2) << elapsed << "ms" << std::endl; - - return passed; -} - -template -bool run_whisper_test(const char* title, const char* options_json, Predicate check) { - if (!g_transcribe_model_path) { - std::cout << "⊘ SKIP │ " << std::left << std::setw(25) << title - << " │ CACTUS_TEST_TRANSCRIBE_MODEL not set\n"; - return true; - } - - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║" << std::setw(42) << std::left << std::string(" ") + title << "║\n" - << "╚══════════════════════════════════════════╝\n"; - - cactus_model_t model = cactus_init(g_transcribe_model_path, nullptr); - if (!model) { - std::cerr << "[✗] Failed to initialize Whisper model\n"; - return false; - } - - char response[1 << 15] = {0}; - StreamingData stream; - stream.model = model; - - std::string audio_path = std::string(g_assets_path) + "/test.wav"; - std::cout << "Transcript: "; - int rc = cactus_transcribe(model, audio_path.c_str(), g_whisper_prompt, - response, sizeof(response), options_json, - stream_callback, &stream, nullptr, 0); - - std::cout << "\n\n[Results]\n"; - if (rc <= 0) { - std::cerr << "failed\n"; - cactus_destroy(model); - return false; - } - - Metrics m; - m.parse(response); - m.print_json(); - - bool ok = check(rc, m); - cactus_destroy(model); - return ok; -} - -static bool test_transcription() { - return run_whisper_test("TRANSCRIPTION", R"({"max_tokens": 100})", - [](int rc, const Metrics& m) { return rc > 0 && m.completion_tokens >= 8; }); -} - -static bool test_stream_transcription() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ STREAM TRANSCRIPTION TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - - if (!g_transcribe_model_path) { - std::cout << "⊘ SKIP │ CACTUS_TEST_TRANSCRIBE_MODEL not set\n"; - return true; - } - - cactus_model_t model = cactus_init(g_transcribe_model_path, nullptr); - if (!model) { - std::cerr << "[✗] Failed to initialize Whisper model\n"; - return false; - } - - cactus_stream_transcribe_t stream = cactus_stream_transcribe_init(model); - if (!stream) { - std::cerr << "[✗] Failed to initialize stream transcribe\n"; - cactus_destroy(model); - return false; - } - - std::string audio_path = std::string(g_assets_path) + "/test.wav"; - FILE* wav_file = fopen(audio_path.c_str(), "rb"); - if (!wav_file) { - std::cerr << "[✗] Failed to open audio file\n"; - cactus_stream_transcribe_destroy(stream); - cactus_destroy(model); - return false; - } - - fseek(wav_file, 44, SEEK_SET); - std::vector pcm_samples; - int16_t sample; - while (fread(&sample, sizeof(int16_t), 1, wav_file) == 1) { - pcm_samples.push_back(sample); - } - fclose(wav_file); - - const size_t chunk_size = 96000; - Timer timer; - std::string full_transcription; - - for (size_t offset = 0; offset < pcm_samples.size(); offset += chunk_size) { - size_t size = std::min(chunk_size, pcm_samples.size() - offset); - - cactus_stream_transcribe_insert( - stream, - reinterpret_cast(pcm_samples.data() + offset), - size * sizeof(int16_t) - ); - - char response[1 << 15] = {0}; - int result = cactus_stream_transcribe_process( - stream, - response, - sizeof(response), R"({"confirmation_threshold": 0.90})" - ); - - if (result < 0) { - std::cerr << "\n[✗] Processing failed\n"; - cactus_stream_transcribe_destroy(stream); - cactus_destroy(model); - return false; - } - - std::string response_str(response); - std::string confirmed = json_string(response_str, "confirmed"); - std::string pending = json_string(response_str, "pending"); - - full_transcription += confirmed; - if (!confirmed.empty()) std::cout << "├─ confirmed: " << confirmed << "\n"; - if (!pending.empty()) std::cout << "├─ pending: " << pending << "\n"; - } - - char final_response[1 << 15] = {0}; - int final_result = cactus_stream_transcribe_finalize( - stream, - final_response, - sizeof(final_response) - ); - - if (final_result < 0) { - std::cerr << "[✗] Finalization failed\n"; - cactus_stream_transcribe_destroy(stream); - cactus_destroy(model); - return false; - } - - std::string final_str(final_response); - std::string final_confirmed = json_string(final_str, "confirmed"); - - if (!final_confirmed.empty()) { - full_transcription += final_confirmed; - std::cout << "└─ confirmed: " << final_confirmed << "\n"; - } - - double elapsed = timer.elapsed_ms(); - - size_t word_count = 0; - bool in_word = false; - for (char c : full_transcription) { - if (std::isspace(c)) { - in_word = false; - } else if (!in_word) { - in_word = true; - word_count++; - } - } - - std::cout << "\n[Results]\n" - << " \"success\": true,\n" - << " \"total_time_ms\": " << std::fixed << std::setprecision(2) << elapsed << ",\n" - << " \"audio_chunks\": " << ((pcm_samples.size() + chunk_size - 1) / chunk_size) << ",\n" - << " \"pcm_samples\": " << pcm_samples.size() << ",\n" - << " \"duration_sec\": " << std::setprecision(2) << (pcm_samples.size() / 16000.0) << ",\n" - << " \"words_transcribed\": " << word_count << "\n" - << "├─ Full transcription: \"" << full_transcription << "\"" << std::endl; - - cactus_stream_transcribe_destroy(stream); - cactus_destroy(model); - return true; -} - -static bool test_image_embeddings() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ IMAGE EMBEDDING TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - - if (!g_model_path) { - std::cout << "⊘ SKIP │ CACTUS_TEST_MODEL not set\n"; - return true; - } - - std::string image_path = std::string(g_assets_path) + "/test_monkey.png"; - const size_t buffer_size = 1024 * 1024 * 4; - std::vector embeddings(buffer_size / sizeof(float)); - size_t embedding_dim = 0; - - cactus_model_t model = cactus_init(g_model_path, nullptr); - if (!model) { - std::cout << "⊘ SKIP │ Model doesn't support image embeddings\n"; - return true; - } - - Timer t; - int result = cactus_image_embed(model, image_path.c_str(), embeddings.data(), buffer_size, &embedding_dim); - double elapsed = t.elapsed_ms(); - - cactus_destroy(model); - - if (result == -1) { - std::cout << "⊘ SKIP │ Model doesn't support image embeddings\n"; - return true; - } - - std::cout << "├─ Embedding dim: " << embedding_dim << "\n" - << "└─ Time: " << std::fixed << std::setprecision(2) << elapsed << "ms" << std::endl; - - return result > 0 && embedding_dim > 0; -} - -static bool test_audio_embeddings() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ AUDIO EMBEDDING TEST ║\n" - << "╚══════════════════════════════════════════╝\n"; - - if (!g_transcribe_model_path) { - std::cout << "⊘ SKIP │ CACTUS_TEST_TRANSCRIBE_MODEL not set\n"; - return true; - } - - const size_t buffer_size = 1024 * 1024; - std::vector embeddings(buffer_size / sizeof(float)); - size_t embedding_dim = 0; - - cactus_model_t model = cactus_init(g_transcribe_model_path, nullptr); - if (!model) { - std::cout << "⊘ SKIP │ Failed to init Whisper model\n"; - return true; - } - - std::string audio_path = std::string(g_assets_path) + "/test.wav"; - Timer t; - int result = cactus_audio_embed(model, audio_path.c_str(), embeddings.data(), buffer_size, &embedding_dim); - double elapsed = t.elapsed_ms(); - - cactus_destroy(model); - - if (result == -1) { - std::cout << "⊘ SKIP │ Model doesn't support audio embeddings\n"; - return true; - } - - std::cout << "├─ Embedding dim: " << embedding_dim << "\n" - << "└─ Time: " << std::fixed << std::setprecision(2) << elapsed << "ms" << std::endl; - - return result > 0 && embedding_dim > 0; -} - -static bool test_pcm_transcription() { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║ PCM BUFFER TRANSCRIPTION ║\n" - << "╚══════════════════════════════════════════╝\n"; - - if (!g_transcribe_model_path) { - std::cout << "⊘ SKIP │ CACTUS_TEST_TRANSCRIBE_MODEL not set\n"; - return true; - } - - cactus_model_t model = cactus_init(g_transcribe_model_path, nullptr); - if (!model) { - std::cerr << "[✗] Failed to initialize Whisper model\n"; - return false; - } - - const size_t sample_rate = 16000; - bool use_microphone = false; - bool test_passed = false; - -#ifdef HAVE_SDL2 - { - std::cout << "Using microphone input (SDL2)...\n"; - - AudioCapture audio_capture(10000); - if (audio_capture.init(0, sample_rate)) { - std::cout << "\n🎤 Recording for 10 seconds... Speak now!\n\n"; - - audio_capture.resume(); - use_microphone = true; - - std::this_thread::sleep_for(std::chrono::seconds(10)); - - audio_capture.pause(); - - std::vector audio_float; - size_t num_samples = audio_capture.get_all(audio_float); - - if (num_samples == 0) { - std::cerr << "[!] No audio captured\n"; - use_microphone = false; - } else { - std::cout << "Captured " << (num_samples / sample_rate) - << " seconds of audio, transcribing...\n"; - - std::vector pcm_samples(num_samples); - for (size_t i = 0; i < num_samples; i++) { - float clamped = std::max(-1.0f, std::min(1.0f, audio_float[i])); - pcm_samples[i] = static_cast(clamped * 32767.0f); - } - - // Transcribe - char response[1 << 15] = {0}; - StreamingData stream; - stream.model = model; - - std::cout << "Transcript: "; - int rc = cactus_transcribe( - model, - nullptr, - g_whisper_prompt, - response, - sizeof(response), - R"({"max_tokens": 100})", - stream_callback, - &stream, - reinterpret_cast(pcm_samples.data()), - pcm_samples.size() * sizeof(int16_t) - ); - - std::cout << "\n\n[Results]\n"; - if (rc > 0) { - Metrics m; - m.parse(response); - m.print_json(); - test_passed = (rc > 0 && m.completion_tokens >= 1); - } else { - std::cerr << "Transcription failed\n"; - } - } - } else { - std::cerr << "[!] Failed to initialize audio capture, falling back to synthetic audio\n"; - } - } -#endif - if (!use_microphone) { - std::cout << "Using synthetic audio (440Hz sine wave)...\n"; - const size_t duration_seconds = 3; - const size_t num_samples = sample_rate * duration_seconds; - std::vector pcm_samples(num_samples); - - for (size_t i = 0; i < num_samples; i++) { - float t = static_cast(i) / sample_rate; - float amplitude = 0.3f; - float value = amplitude * std::sin(2.0f * M_PI * 440.0f * t); - pcm_samples[i] = static_cast(value * 32767.0f); - } - - char response[1 << 15] = {0}; - StreamingData stream; - stream.model = model; - - std::cout << "Transcript: "; - int rc = cactus_transcribe( - model, - nullptr, - g_whisper_prompt, - response, - sizeof(response), - R"({"max_tokens": 100})", - stream_callback, - &stream, - reinterpret_cast(pcm_samples.data()), - pcm_samples.size() * sizeof(int16_t) - ); - - std::cout << "\n\n[Results]\n"; - if (rc <= 0) { - std::cerr << "failed\n"; - cactus_destroy(model); - return false; - } - - Metrics m; - m.parse(response); - m.print_json(); - - std::cout << "├─ PCM samples: " << pcm_samples.size() << "\n" - << "├─ Duration: " << duration_seconds << "s\n" - << "└─ Sample rate: " << sample_rate << "Hz\n"; - - test_passed = (rc > 0 && m.completion_tokens >= 1); - } - - cactus_destroy(model); - return test_passed; -} - -int main() { -#ifdef __APPLE__ - cactus_set_telemetry_token("973e4aaa-5ee4-4947-a128-757bb66be75b"); - cactus_set_pro_key(""); // email founders@cactuscompute.com -#endif - - TestUtils::TestRunner runner("Engine Tests"); - runner.run_test("streaming", test_streaming()); - runner.run_test("tool_calls", test_tool_call()); - runner.run_test("tool_calls_with_two_tools", test_tool_call_with_two_tools()); - runner.run_test("tool_calls_with_three_tools", test_tool_call_with_three_tools()); - runner.run_test("cloud_handoff", test_cloud_handoff()); - runner.run_test("embeddings", test_embeddings()); - runner.run_test("image_embeddings", test_image_embeddings()); - runner.run_test("audio_embeddings", test_audio_embeddings()); - runner.run_test("vlm_multiturn", test_vlm_multiturn()); - runner.run_test("audio_processor", test_audio_processor()); - runner.run_test("transcription", test_transcription()); - runner.run_test("pcm_transcription", test_pcm_transcription()); - runner.run_test("stream_transcription", test_stream_transcription()); - runner.run_test("rag_preprocessing", test_rag()); - runner.run_test("4k_context", test_4k_context()); - runner.print_summary(); - return runner.all_passed() ? 0 : 1; -} \ No newline at end of file diff --git a/tests/test_graph.cpp b/tests/test_graph.cpp deleted file mode 100644 index ede2ee1aa..000000000 --- a/tests/test_graph.cpp +++ /dev/null @@ -1,907 +0,0 @@ -#include "test_utils.h" -#include -#include -#include -#include -#include -#include -#include - -bool test_basic_operations() { - TestUtils::FP16TestFixture fixture("Basic Operations"); - - size_t input_a = fixture.create_input({2, 3}); - size_t input_b = fixture.create_input({2, 3}); - size_t add_result = fixture.graph().add(input_a, input_b); - size_t mul_result = fixture.graph().multiply(add_result, input_a); - size_t scalar_result = fixture.graph().scalar_multiply(mul_result, 2.0f); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - std::vector<__fp16> data_b = {2, 3, 4, 5, 6, 7}; - - fixture.set_input_data(input_a, data_a); - fixture.set_input_data(input_b, data_b); - fixture.execute(); - - std::vector<__fp16> expected(6); - for (int i = 0; i < 6; i++) { - float result = ((static_cast(data_a[i]) + static_cast(data_b[i])) * static_cast(data_a[i])) * 2.0f; - expected[i] = static_cast<__fp16>(result); - } - - return fixture.verify_output(scalar_result, expected); -} - -bool test_basic_addition() { - return TestUtils::test_basic_operation( - "Addition", - [](CactusGraph& graph, size_t a, size_t b) { return graph.add(a, b); }, - {1, 2, 3, 4}, - {5, 6, 7, 8}, - {6, 8, 10, 12} - ); -} - -bool test_basic_subtraction() { - return TestUtils::test_basic_operation( - "Subtraction", - [](CactusGraph& graph, size_t a, size_t b) { return graph.subtract(a, b); }, - {10, 8, 6, 4}, - {2, 3, 1, 2}, - {8, 5, 5, 2} - ); -} - -bool test_basic_multiplication() { - return TestUtils::test_basic_operation( - "Multiplication", - [](CactusGraph& graph, size_t a, size_t b) { return graph.multiply(a, b); }, - {2, 3, 4, 5}, - {3, 4, 2, 2}, - {6, 12, 8, 10} - ); -} - -bool test_basic_division() { - return TestUtils::test_basic_operation( - "Division", - [](CactusGraph& graph, size_t a, size_t b) { return graph.divide(a, b); }, - {12, 15, 8, 9}, - {3, 5, 2, 3}, - {4, 3, 4, 3} - ); -} - -bool test_matrix_multiplication() { - TestUtils::FP16TestFixture fixture("Matrix Multiplication"); - - size_t input_a = fixture.create_input({2, 3}); - size_t input_b = fixture.create_input({3, 2}); - size_t matmul_result = fixture.graph().matmul(input_a, input_b, false); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - std::vector<__fp16> data_b = {1, 2, 3, 4, 5, 6}; - fixture.set_input_data(input_a, data_a); - fixture.set_input_data(input_b, data_b); - fixture.execute(); - - std::vector<__fp16> expected = {22, 28, 49, 64}; - return fixture.verify_output(matmul_result, expected); -} - -bool test_transpose() { - TestUtils::FP16TestFixture fixture("Transpose"); - - size_t input_a = fixture.create_input({2, 3}); - size_t transpose_result = fixture.graph().transpose(input_a); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected = {1, 4, 2, 5, 3, 6}; - return fixture.verify_output(transpose_result, expected); -} - -bool test_reshape() { - TestUtils::FP16TestFixture fixture("Reshape"); - - size_t input_a = fixture.create_input({2, 3}); - size_t reshape_result = fixture.graph().reshape(input_a, {3, 2}); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - return fixture.verify_output(reshape_result, data_a); -} - - -bool test_scalar_operations() { - TestUtils::FP16TestFixture fixture("Scalar Operations"); - - size_t input_a = fixture.create_input({4}); - size_t add_result = fixture.graph().scalar_add(input_a, 5.0f); - size_t mul_result = fixture.graph().scalar_multiply(add_result, 2.0f); - - std::vector<__fp16> data_a = {1, 2, 3, 4}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected = {12, 14, 16, 18}; - return fixture.verify_output(mul_result, expected); -} - -bool test_scalar_subtract_divide() { - TestUtils::FP16TestFixture fixture("Scalar Subtract/Divide"); - - size_t input_a = fixture.create_input({4}); - size_t sub_result = fixture.graph().scalar_subtract(input_a, 2.0f); - size_t div_result = fixture.graph().scalar_divide(input_a, 2.0f); - - std::vector<__fp16> data_a = {10, 8, 6, 4}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected_sub = {8, 6, 4, 2}; - std::vector<__fp16> expected_div = {5, 4, 3, 2}; - return fixture.verify_output(sub_result, expected_sub) && - fixture.verify_output(div_result, expected_div); -} - -bool test_scalar_math_functions() { - TestUtils::FP16TestFixture fixture("Scalar Math Functions"); - - size_t input_a = fixture.create_input({3}); - size_t exp_result = fixture.graph().scalar_exp(input_a); - size_t sqrt_result = fixture.graph().scalar_sqrt(input_a); - size_t cos_result = fixture.graph().scalar_cos(input_a); - size_t sin_result = fixture.graph().scalar_sin(input_a); - - std::vector<__fp16> input_data = {0.0f, 1.0f, 4.0f}; - fixture.set_input_data(input_a, input_data); - fixture.execute(); - - std::vector<__fp16> exp_expected = {1.0f, 2.71828f, 54.5982f}; - std::vector<__fp16> sqrt_expected = {0.0f, 1.0f, 2.0f}; - std::vector<__fp16> cos_expected = {1.0f, 0.54030f, -0.65364f}; - std::vector<__fp16> sin_expected = {0.0f, 0.84147f, -0.75680f}; - - return fixture.verify_output(exp_result, exp_expected, 0.01f) && - fixture.verify_output(sqrt_result, sqrt_expected, 0.01f) && - fixture.verify_output(cos_result, cos_expected, 0.01f) && - fixture.verify_output(sin_result, sin_expected, 0.01f); -} - -bool test_rms_norm() { - TestUtils::FP16TestFixture fixture("RMS Norm"); - - size_t input_a = fixture.create_input({1, 8}); - size_t weight = fixture.create_input({8}); - size_t norm_result = fixture.graph().rms_norm(input_a, weight); - - std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; - std::vector<__fp16> weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; - - fixture.set_input_data(input_a, input_data); - fixture.set_input_data(weight, weight_data); - fixture.execute(); - - float sum_squares = 0.0f; - for (auto val : input_data) { - float v = static_cast(val); - sum_squares += v * v; - } - float rms = sqrtf(sum_squares / 8.0f + 1e-5f); - float inv_rms = 1.0f / rms; - - std::vector<__fp16> expected; - for (size_t i = 0; i < input_data.size(); i++) { - expected.push_back(static_cast<__fp16>(static_cast(input_data[i]) * inv_rms * static_cast(weight_data[i]))); - } - - return fixture.verify_output(norm_result, expected, 0.01f); -} - -bool test_softmax() { - TestUtils::FP16TestFixture fixture("Softmax"); - - size_t input_a = fixture.create_input({2, 3}); - size_t softmax_result = fixture.graph().softmax(input_a, -1); - - std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; - fixture.set_input_data(input_a, input_data); - fixture.execute(); - - std::vector<__fp16> expected = {0.09003f, 0.24473f, 0.66524f, 0.09003f, 0.24473f, 0.66524f}; - return fixture.verify_output(softmax_result, expected, 0.01f); -} - -bool test_attention() { - TestUtils::FP16TestFixture fixture("Attention"); - - size_t query = fixture.create_input({1, 2, 1, 4}); - size_t key = fixture.create_input({1, 2, 1, 4}); - size_t value = fixture.create_input({1, 2, 1, 4}); - - size_t attention_result = fixture.graph().attention(query, key, value, 0.5f); - (void)attention_result; - - std::vector<__fp16> q_data = {1, 0, 0, 0, 0, 1, 0, 0}; - std::vector<__fp16> k_data = {1, 0, 0, 0, 0, 1, 0, 0}; - std::vector<__fp16> v_data = {1, 2, 3, 4, 5, 6, 7, 8}; - - fixture.set_input_data(query, q_data); - fixture.set_input_data(key, k_data); - fixture.set_input_data(value, v_data); - fixture.execute(); - - return true; -} - -bool test_reduction_operations() { - TestUtils::FP16TestFixture fixture("Reduction Operations"); - - size_t input_a = fixture.create_input({2, 3}); - size_t sum_all = fixture.graph().sum(input_a, -1); - size_t sum_axis0 = fixture.graph().sum(input_a, 0); - size_t sum_axis1 = fixture.graph().sum(input_a, 1); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected_all = {21}; - std::vector<__fp16> expected_axis0 = {5, 7, 9}; - std::vector<__fp16> expected_axis1 = {6, 15}; - - return fixture.verify_output(sum_all, expected_all) && - fixture.verify_output(sum_axis0, expected_axis0) && - fixture.verify_output(sum_axis1, expected_axis1); -} - -bool test_fp16_reduction_operations() { - CactusGraph graph; - - size_t input_a = graph.input({2, 3}, Precision::FP16); - size_t sum_all = graph.sum(input_a, -1); - - std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; - graph.set_input(input_a, input_data.data(), Precision::FP16); - graph.execute(); - - __fp16* output = static_cast<__fp16*>(graph.get_output(sum_all)); - double result = static_cast(output[0]); - double expected = 21.0; - - bool success = std::abs(result - expected) < 0.1f; // FP16 has lower precision - - graph.hard_reset(); - return success; -} - -bool test_mean_operations() { - TestUtils::FP16TestFixture fixture("Mean Operations"); - - size_t input_a = fixture.create_input({2, 4}); - size_t mean_all = fixture.graph().mean(input_a, -1); - size_t mean_axis0 = fixture.graph().mean(input_a, 0); - size_t mean_axis1 = fixture.graph().mean(input_a, 1); - - std::vector<__fp16> data_a = {2, 4, 6, 8, 10, 12, 14, 16}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected_all = {9}; - std::vector<__fp16> expected_axis0 = {6, 8, 10, 12}; - std::vector<__fp16> expected_axis1 = {5, 13}; - - return fixture.verify_output(mean_all, expected_all) && - fixture.verify_output(mean_axis0, expected_axis0) && - fixture.verify_output(mean_axis1, expected_axis1); -} - -bool test_variance_operations() { - TestUtils::FP16TestFixture fixture("Variance Operations"); - - size_t input_a = fixture.create_input({1, 4}); - size_t var_axis1 = fixture.graph().variance(input_a, 1); - - std::vector<__fp16> input_data = {1.0f, 2.0f, 3.0f, 4.0f}; - fixture.set_input_data(input_a, input_data); - fixture.execute(); - - std::vector<__fp16> expected = {1.25f}; - return fixture.verify_output(var_axis1, expected, 0.01f); -} - -bool test_min_max_operations() { - TestUtils::FP16TestFixture fixture("Min/Max Operations"); - - size_t input_a = fixture.create_input({2, 3}); - size_t min_axis0 = fixture.graph().min(input_a, 0); - size_t max_axis0 = fixture.graph().max(input_a, 0); - size_t min_axis1 = fixture.graph().min(input_a, 1); - size_t max_axis1 = fixture.graph().max(input_a, 1); - - std::vector<__fp16> data_a = {6, 2, 8, 1, 5, 3}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected_min_axis0 = {1, 2, 3}; - std::vector<__fp16> expected_max_axis0 = {6, 5, 8}; - std::vector<__fp16> expected_min_axis1 = {2, 1}; - std::vector<__fp16> expected_max_axis1 = {8, 5}; - - return fixture.verify_output(min_axis0, expected_min_axis0) && - fixture.verify_output(max_axis0, expected_max_axis0) && - fixture.verify_output(min_axis1, expected_min_axis1) && - fixture.verify_output(max_axis1, expected_max_axis1); -} - -bool test_fp16_precision() { - TestUtils::FP16TestFixture fixture("FP16 Precision"); - - size_t input_a = fixture.create_input({3}); - size_t input_b = fixture.create_input({3}); - size_t result_id = fixture.graph().add(input_a, input_b); - - std::vector<__fp16> data_a = {1.5f, 2.5f, 3.5f}; - std::vector<__fp16> data_b = {0.5f, 1.5f, 2.5f}; - fixture.set_input_data(input_a, data_a); - fixture.set_input_data(input_b, data_b); - fixture.execute(); - - std::vector<__fp16> expected = {2.0f, 4.0f, 6.0f}; - return fixture.verify_output(result_id, expected); -} - -bool test_broadcast_shape_compatibility() { - TestUtils::FP16TestFixture fixture("Broadcast Shape Compatibility"); - - size_t a_id = fixture.create_input({2, 3}); - size_t b_id = fixture.create_input({2, 1}); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - std::vector<__fp16> data_b = {10, 20}; - fixture.set_input_data(a_id, data_a); - fixture.set_input_data(b_id, data_b); - - size_t result_id = fixture.graph().add(a_id, b_id); - fixture.execute(); - - std::vector<__fp16> expected = {11, 12, 13, 24, 25, 26}; - return fixture.verify_output(result_id, expected); -} - -bool test_broadcast_scalar_tensor() { - TestUtils::FP16TestFixture fixture("Broadcast Scalar Tensor"); - - size_t a_id = fixture.create_input({2, 2}); - size_t b_id = fixture.create_input({1}); - - std::vector<__fp16> data_a = {1, 2, 3, 4}; - std::vector<__fp16> data_b = {5}; - fixture.set_input_data(a_id, data_a); - fixture.set_input_data(b_id, data_b); - - size_t result_id = fixture.graph().add(a_id, b_id); - fixture.execute(); - - std::vector<__fp16> expected = {6, 7, 8, 9}; - return fixture.verify_output(result_id, expected); -} - -bool test_broadcast_different_ranks() { - TestUtils::FP16TestFixture fixture("Broadcast Different Ranks"); - - size_t a_id = fixture.create_input({2, 2, 3}); - size_t b_id = fixture.create_input({2, 3}); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; - std::vector<__fp16> data_b = {1, 1, 1, 2, 2, 2}; - fixture.set_input_data(a_id, data_a); - fixture.set_input_data(b_id, data_b); - - size_t result_id = fixture.graph().add(a_id, b_id); - fixture.execute(); - - std::vector<__fp16> expected = {2, 3, 4, 6, 7, 8, 8, 9, 10, 12, 13, 14}; - return fixture.verify_output(result_id, expected); -} - -bool test_broadcast_fp16_precision() { - TestUtils::FP16TestFixture fixture("Broadcast FP16 Precision"); - - size_t a_id = fixture.create_input({2, 2}); - size_t b_id = fixture.create_input({1}); - - std::vector<__fp16> data_a = {1.5f, 2.5f, 3.5f, 4.5f}; - std::vector<__fp16> data_b = {0.5f}; - fixture.set_input_data(a_id, data_a); - fixture.set_input_data(b_id, data_b); - - size_t result_id = fixture.graph().add(a_id, b_id); - fixture.execute(); - - std::vector<__fp16> expected = {2.0f, 3.0f, 4.0f, 5.0f}; - return fixture.verify_output(result_id, expected); -} - -bool test_precision_traits() { - assert(PrecisionTraits::size_of(Precision::INT8) == 1); - assert(PrecisionTraits::size_of(Precision::FP32) == 4); - return true; -} - -bool test_graph_precision_construction() { - TestUtils::FP16TestFixture fixture("Graph Precision Construction"); - - size_t fp16_id = fixture.create_input({2, 3}, Precision::FP16); - size_t fp32_id = fixture.create_input({3, 4}, Precision::FP32); - - const auto& fp16_buffer = fixture.graph().get_output_buffer(fp16_id); - const auto& fp32_buffer = fixture.graph().get_output_buffer(fp32_id); - - assert(fp16_buffer.precision == Precision::FP16); - assert(fp16_buffer.shape[0] == 2); - assert(fp16_buffer.shape[1] == 3); - assert(fp16_buffer.byte_size == 12); // 6 elements * 2 bytes - - assert(fp32_buffer.precision == Precision::FP32); - assert(fp32_buffer.shape[0] == 3); - assert(fp32_buffer.shape[1] == 4); - assert(fp32_buffer.byte_size == 48); - - return true; -} - -bool test_precision_conversion() { - TestUtils::FP16TestFixture fixture("Precision Conversion"); - - size_t fp16_id = fixture.create_input({2, 2}, Precision::FP16); - std::vector<__fp16> data = {1, 2, 3, 4}; - fixture.set_input_data(fp16_id, data); - - size_t fp32_converted_id = fixture.graph().precision_cast(fp16_id, Precision::FP32); - fixture.execute(); - - float* fp32_data = static_cast(fixture.graph().get_output(fp32_converted_id)); - - for (size_t i = 0; i < 4; ++i) { - assert(std::abs(fp32_data[i] - static_cast(data[i])) < 1e-3f); - } - - return true; -} - -bool test_graph_save_load() { - try { - CactusGraph graph; - - size_t input_a = graph.input({2, 3}, Precision::FP16); - size_t input_b = graph.input({2, 3}, Precision::FP16); - size_t result_id = graph.add(input_a, input_b); - - std::vector<__fp16> data_a = {1, 2, 3, 4, 5, 6}; - std::vector<__fp16> data_b = {10, 20, 30, 40, 50, 60}; - - graph.set_input(input_a, const_cast(static_cast(data_a.data())), Precision::FP16); - graph.set_input(input_b, const_cast(static_cast(data_b.data())), Precision::FP16); - graph.execute(); - - std::string filename = "test_graph_save_load.bin"; - GraphFile::save_node(graph, result_id, filename); - - CactusGraph new_graph; - auto loaded = GraphFile::load_into_graph(new_graph, filename); - new_graph.execute(); - - __fp16* original_data = static_cast<__fp16*>(graph.get_output(result_id)); - __fp16* loaded_data = static_cast<__fp16*>(new_graph.get_output(loaded.node_id)); - - for (size_t i = 0; i < 6; ++i) { - if (std::abs(static_cast(original_data[i]) - static_cast(loaded_data[i])) > 1e-3f) { - graph.hard_reset(); - new_graph.hard_reset(); - std::remove(filename.c_str()); - return false; - } - } - - bool result = (loaded.shape == std::vector{2, 3}) && - (loaded.precision == Precision::FP16) && - (loaded.byte_size == 12); - - graph.hard_reset(); - new_graph.hard_reset(); - std::remove(filename.c_str()); - return result; - } catch (const std::exception& e) { - return false; - } -} - -bool test_complex_graph_structure() { - TestUtils::FP16TestFixture fixture("Complex Graph Structure"); - - size_t input_a = fixture.create_input({2, 2}); - size_t input_b = fixture.create_input({2, 2}); - size_t input_c = fixture.create_input({2, 2}); - - size_t add_ab = fixture.graph().add(input_a, input_b); - size_t mul_result = fixture.graph().multiply(add_ab, input_c); - size_t scalar_result = fixture.graph().scalar_add(mul_result, 1.0f); - - std::vector<__fp16> data_a = {1, 2, 3, 4}; - std::vector<__fp16> data_b = {2, 3, 4, 5}; - std::vector<__fp16> data_c = {2, 2, 2, 2}; - fixture.set_input_data(input_a, data_a); - fixture.set_input_data(input_b, data_b); - fixture.set_input_data(input_c, data_c); - - fixture.execute(); - - std::vector<__fp16> expected = {7, 11, 15, 19}; - return fixture.verify_output(scalar_result, expected); -} - -bool test_multiple_outputs() { - TestUtils::FP16TestFixture fixture("Multiple Outputs"); - - size_t input_a = fixture.create_input({3}); - size_t add_result = fixture.graph().scalar_add(input_a, 10.0f); - size_t mul_result = fixture.graph().scalar_multiply(input_a, 2.0f); - size_t combine_result = fixture.graph().add(add_result, mul_result); - - std::vector<__fp16> data_a = {1, 2, 3}; - fixture.set_input_data(input_a, data_a); - fixture.execute(); - - std::vector<__fp16> expected_add = {11, 12, 13}; - std::vector<__fp16> expected_mul = {2, 4, 6}; - std::vector<__fp16> expected_combine = {13, 16, 19}; - - return fixture.verify_output(add_result, expected_add) && - fixture.verify_output(mul_result, expected_mul) && - fixture.verify_output(combine_result, expected_combine); -} - -bool test_graph_reset() { - CactusGraph graph; - - size_t input_a = graph.input({2}, Precision::FP16); - size_t result_id = graph.scalar_add(input_a, 5.0f); - - std::vector<__fp16> data_a = {1, 2}; - graph.set_input(input_a, data_a.data(), Precision::FP16); - graph.execute(); - - __fp16* output1 = static_cast<__fp16*>(graph.get_output(result_id)); - if (std::abs(static_cast(output1[0]) - 6.0f) > 1e-2f || - std::abs(static_cast(output1[1]) - 7.0f) > 1e-2f) return false; - - graph.hard_reset(); - if (graph.get_node_count() != 0) return false; - - size_t new_input = graph.input({2}, Precision::FP16); - size_t new_result = graph.scalar_add(new_input, 5.0f); - - std::vector<__fp16> data_b = {10, 20}; - graph.set_input(new_input, data_b.data(), Precision::FP16); - graph.execute(); - - __fp16* output2 = static_cast<__fp16*>(graph.get_output(new_result)); - return (std::abs(static_cast(output2[0]) - 15.0f) < 1e-2f && - std::abs(static_cast(output2[1]) - 25.0f) < 1e-2f); -} - -bool test_gather_operation() { - // Gather uses INT8 indices but can work with FP16 data - CactusGraph graph; - - size_t embeddings = graph.input({5, 3}, Precision::FP16); - size_t indices = graph.input({2, 2}, Precision::INT8); - size_t gathered = graph.gather(embeddings, indices); - - std::vector<__fp16> emb_data = { - 1, 2, 3, - 4, 5, 6, - 7, 8, 9, - 10, 11, 12, - 13, 14, 15 - }; - std::vector idx_data = {0, 2, 4, 1}; - - graph.set_input(embeddings, emb_data.data(), Precision::FP16); - graph.set_input(indices, idx_data.data(), Precision::INT8); - graph.execute(); - - __fp16* output = static_cast<__fp16*>(graph.get_output(gathered)); - std::vector<__fp16> expected = { - 1, 2, 3, - 7, 8, 9, - 13, 14, 15, - 4, 5, 6 - }; - - for (size_t i = 0; i < expected.size(); ++i) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-3f) { - graph.hard_reset(); - return false; - } - } - - graph.hard_reset(); - return true; -} - -bool test_gather_1d_tensor() { - CactusGraph graph; - - size_t tensor = graph.input({8}, Precision::FP16); - size_t indices = graph.input({3}, Precision::INT8); - size_t gathered = graph.gather(tensor, indices); - - std::vector<__fp16> tensor_data = {10, 20, 30, 40, 50, 60, 70, 80}; - std::vector idx_data = {7, 2, 0}; - - graph.set_input(tensor, tensor_data.data(), Precision::FP16); - graph.set_input(indices, idx_data.data(), Precision::INT8); - graph.execute(); - - __fp16* output = static_cast<__fp16*>(graph.get_output(gathered)); - std::vector<__fp16> expected = {80, 30, 10}; - - for (size_t i = 0; i < expected.size(); ++i) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-3f) { - graph.hard_reset(); - return false; - } - } - - graph.hard_reset(); - return true; -} - -bool test_gather_3d_tensor() { - TestUtils::FP16TestFixture fixture("Gather 3D Tensor"); - - size_t tensor = fixture.create_input({3, 2, 4}); - size_t indices = fixture.graph().input({2}, Precision::INT8); - size_t gathered = fixture.graph().gather(tensor, indices); - - std::vector<__fp16> tensor_data = { - // First 2x4 slice - 1.0f, 2.0f, 3.0f, 4.0f, - 5.0f, 6.0f, 7.0f, 8.0f, - // Second 2x4 slice - 9.0f, 10.0f, 11.0f, 12.0f, - 13.0f, 14.0f, 15.0f, 16.0f, - // Third 2x4 slice - 17.0f, 18.0f, 19.0f, 20.0f, - 21.0f, 22.0f, 23.0f, 24.0f - }; - std::vector idx_data = {2, 0}; - - fixture.set_input_data(tensor, tensor_data); - fixture.graph().set_input(indices, idx_data.data(), Precision::INT8); - fixture.execute(); - - std::vector<__fp16> expected = { - // Third slice (index 2) - 17.0f, 18.0f, 19.0f, 20.0f, - 21.0f, 22.0f, 23.0f, 24.0f, - // First slice (index 0) - 1.0f, 2.0f, 3.0f, 4.0f, - 5.0f, 6.0f, 7.0f, 8.0f - }; - - return fixture.verify_output(gathered, expected); -} - -bool test_gather_fp16() { - TestUtils::FP16TestFixture fixture("Gather FP16"); - - size_t embeddings = fixture.create_input({4, 2}); - CactusGraph& graph = fixture.graph(); - size_t indices = graph.input({3}, Precision::INT8); - size_t gathered = graph.gather(embeddings, indices); - - std::vector<__fp16> emb_data = { - 1.0f, 2.0f, - 3.0f, 4.0f, - 5.0f, 6.0f, - 7.0f, 8.0f - }; - std::vector idx_data = {2, 0, 3}; - - fixture.set_input_data(embeddings, emb_data); - graph.set_input(indices, idx_data.data(), Precision::INT8); - fixture.execute(); - - std::vector<__fp16> expected = { - 5.0f, 6.0f, - 1.0f, 2.0f, - 7.0f, 8.0f - }; - - return fixture.verify_output(gathered, expected); -} - -bool test_mmap_gather() { - CactusGraph graph; - - std::vector<__fp16> embeddings_data = { - 1.0f, 2.0f, 3.0f, - 4.0f, 5.0f, 6.0f, - 7.0f, 8.0f, 9.0f, - 10.0f, 11.0f, 12.0f - }; - - size_t temp_embeddings = graph.input({4, 3}, Precision::FP16); - graph.set_input(temp_embeddings, embeddings_data.data(), Precision::FP16); - - const std::string temp_file = "test_embeddings.bin"; - GraphFile::save_node(graph, temp_embeddings, temp_file); - - graph.hard_reset(); - - size_t mmap_embeddings = graph.mmap_embeddings(temp_file); - size_t indices = graph.input({3}, Precision::INT8); - size_t gathered = graph.gather(mmap_embeddings, indices); - - std::vector idx_data = {2, 0, 3}; - graph.set_input(indices, idx_data.data(), Precision::INT8); - graph.execute(); - - std::vector<__fp16> expected = { - 7.0f, 8.0f, 9.0f, - 1.0f, 2.0f, 3.0f, - 10.0f, 11.0f, 12.0f - }; - - __fp16* output = static_cast<__fp16*>(graph.get_output(gathered)); - bool passed = true; - for (size_t i = 0; i < expected.size(); i++) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 0.01f) { - passed = false; - break; - } - } - - std::remove(temp_file.c_str()); - - return passed; -} - -bool test_embedding_operation() { - CactusGraph graph; - - // INT8 embeddings → FP16 output (correct pattern for quantized storage) - size_t embeddings = graph.input({4, 3}, Precision::INT8); - size_t indices = graph.input({2, 2}, Precision::INT8); - size_t embedded = graph.embedding(embeddings, indices); - - std::vector emb_data = { - 1, 5, 9, - 2, 6, 10, - 3, 7, 11, - 4, 8, 12 - }; - std::vector idx_data = {0, 2, 3, 1}; - - graph.set_input(embeddings, emb_data.data(), Precision::INT8); - graph.set_input(indices, idx_data.data(), Precision::INT8); - graph.execute(); - - // Embedding converts INT8 to FP16, so we need to check FP16 output - __fp16* output = static_cast<__fp16*>(graph.get_output(embedded)); - - // Expected values (as FP16) - std::vector<__fp16> expected = { - 1, 5, 9, - 3, 7, 11, - 4, 8, 12, - 2, 6, 10 - }; - - for (size_t i = 0; i < expected.size(); ++i) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-6f) { - return false; - } - } - - return true; -} - -bool test_embedding_from_file() { - CactusGraph graph; - - std::vector<__fp16> embeddings_data = { - 1.0f, 5.0f, 9.0f, - 2.0f, 6.0f, 10.0f, - 3.0f, 7.0f, 11.0f, - 4.0f, 8.0f, 12.0f - }; - - size_t temp_embeddings = graph.input({4, 3}, Precision::FP16); - graph.set_input(temp_embeddings, embeddings_data.data(), Precision::FP16); - - const std::string temp_file = "test_embedding.bin"; - GraphFile::save_node(graph, temp_embeddings, temp_file); - - graph.hard_reset(); - - size_t indices = graph.input({3}, Precision::INT8); - size_t embedded = graph.embedding(temp_file, indices); - - std::vector idx_data = {2, 0, 3}; - graph.set_input(indices, idx_data.data(), Precision::INT8); - graph.execute(); - - std::vector<__fp16> expected = { - 3.0f, 7.0f, 11.0f, - 1.0f, 5.0f, 9.0f, - 4.0f, 8.0f, 12.0f - }; - - __fp16* output = static_cast<__fp16*>(graph.get_output(embedded)); - bool passed = true; - for (size_t i = 0; i < expected.size(); i++) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 0.01f) { - passed = false; - break; - } - } - - std::remove(temp_file.c_str()); - - return passed; -} - -int main() { - TestUtils::TestRunner runner("Graph Operations Tests"); - - runner.run_test("Basic Operations", test_basic_operations()); - runner.run_test("Basic Addition", test_basic_addition()); - runner.run_test("Basic Subtraction", test_basic_subtraction()); - runner.run_test("Basic Multiplication", test_basic_multiplication()); - runner.run_test("Basic Division", test_basic_division()); - runner.run_test("Matrix Multiplication", test_matrix_multiplication()); - runner.run_test("Transpose", test_transpose()); - runner.run_test("Reshape", test_reshape()); - runner.run_test("Scalar Operations", test_scalar_operations()); - runner.run_test("Scalar Subtract/Divide", test_scalar_subtract_divide()); - runner.run_test("Scalar Math Functions", test_scalar_math_functions()); - runner.run_test("Reduction Operations", test_reduction_operations()); - runner.run_test("FP16 Reduction Operations", test_fp16_reduction_operations()); - runner.run_test("Mean Operations", test_mean_operations()); - runner.run_test("Variance Operations", test_variance_operations()); - runner.run_test("Min/Max Operations", test_min_max_operations()); - runner.run_test("RMS Norm", test_rms_norm()); - runner.run_test("Softmax", test_softmax()); - runner.run_test("Attention", test_attention()); - runner.run_test("FP16 Precision", test_fp16_precision()); - runner.run_test("Broadcast Shape Compatibility", test_broadcast_shape_compatibility()); - runner.run_test("Broadcast Scalar Tensor", test_broadcast_scalar_tensor()); - runner.run_test("Broadcast Different Ranks", test_broadcast_different_ranks()); - runner.run_test("Broadcast FP16 Precision", test_broadcast_fp16_precision()); - runner.run_test("Precision Traits", test_precision_traits()); - runner.run_test("Graph Precision Construction", test_graph_precision_construction()); - runner.run_test("Precision Conversion", test_precision_conversion()); - runner.run_test("Graph Save/Load", test_graph_save_load()); - runner.run_test("Complex Graph Structure", test_complex_graph_structure()); - runner.run_test("Multiple Outputs", test_multiple_outputs()); - runner.run_test("Graph Reset", test_graph_reset()); - runner.run_test("Gather Operation", test_gather_operation()); - runner.run_test("Gather 1D Tensor", test_gather_1d_tensor()); - runner.run_test("Gather 3D Tensor", test_gather_3d_tensor()); - runner.run_test("Gather FP16", test_gather_fp16()); - runner.run_test("Memory-Mapped Gather", test_mmap_gather()); - runner.run_test("Embedding Operation", test_embedding_operation()); - runner.run_test("Embedding from File", test_embedding_from_file()); - - runner.print_summary(); - return runner.all_passed() ? 0 : 1; -} diff --git a/tests/test_kernel.cpp b/tests/test_kernel.cpp deleted file mode 100644 index f6c13b87a..000000000 --- a/tests/test_kernel.cpp +++ /dev/null @@ -1,424 +0,0 @@ -#include "test_utils.h" -#include -#include -#include -#include - -bool test_neon_add_fp16_correctness() { - const size_t size = 16; - std::vector<__fp16> a(size), b(size), result(size), expected(size); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_real_distribution dis(-1.0f, 1.0f); - for (size_t i = 0; i < size; ++i) { - a[i] = static_cast<__fp16>(dis(gen)); - b[i] = static_cast<__fp16>(dis(gen)); - expected[i] = a[i] + b[i]; - } - - cactus_add_f16(a.data(), b.data(), result.data(), size); - - for (size_t i = 0; i < size; ++i) { - if (std::abs(static_cast(result[i] - expected[i])) > 1e-3f) { - return false; - } - } - return true; -} - -bool test_neon_subtract_fp16_correctness() { - const size_t size = 16; - std::vector<__fp16> a(size), b(size), result(size), expected(size); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_real_distribution dis(-1.0f, 1.0f); - for (size_t i = 0; i < size; ++i) { - a[i] = static_cast<__fp16>(dis(gen)); - b[i] = static_cast<__fp16>(dis(gen)); - expected[i] = a[i] - b[i]; - } - - cactus_subtract_f16(a.data(), b.data(), result.data(), size); - - for (size_t i = 0; i < size; ++i) { - if (std::abs(static_cast(result[i] - expected[i])) > 1e-3f) { - return false; - } - } - return true; -} - -bool test_neon_multiply_fp16_correctness() { - const size_t size = 16; - std::vector<__fp16> a(size), b(size), result(size), expected(size); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_real_distribution dis(-1.0f, 1.0f); - for (size_t i = 0; i < size; ++i) { - a[i] = static_cast<__fp16>(dis(gen)); - b[i] = static_cast<__fp16>(dis(gen)); - expected[i] = a[i] * b[i]; - } - - cactus_multiply_f16(a.data(), b.data(), result.data(), size); - - for (size_t i = 0; i < size; ++i) { - if (std::abs(static_cast(result[i] - expected[i])) > 1e-3f) { - return false; - } - } - return true; -} - -bool test_neon_scalar_operations_fp16_correctness() { - const size_t size = 8; - std::vector<__fp16> input = {1.0f, 2.0f, 3.0f, 4.0f, -1.0f, -2.0f, -3.0f, -4.0f}; - std::vector<__fp16> result(size); - const float scalar = 2.0f; - - std::vector<__fp16> expected_add(size); - for (size_t i = 0; i < size; ++i) { - expected_add[i] = static_cast<__fp16>(static_cast(input[i]) + scalar); - } - - cactus_scalar_op_f16(input.data(), result.data(), size, scalar, ScalarOpType::ADD); - - if (!TestUtils::compare_arrays(result.data(), expected_add.data(), size)) { - return false; - } - - std::vector<__fp16> expected_mul(size); - for (size_t i = 0; i < size; ++i) { - expected_mul[i] = static_cast<__fp16>(static_cast(input[i]) * scalar); - } - - cactus_scalar_op_f16(input.data(), result.data(), size, scalar, ScalarOpType::MULTIPLY); - - return TestUtils::compare_arrays(result.data(), expected_mul.data(), size); -} - -bool test_neon_reduction_correctness() { - std::vector<__fp16> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; - - double sum_result = cactus_sum_all_f16(input.data(), input.size()); - double expected_sum = 36.0; - - if (std::abs(sum_result - expected_sum) > 1e-2) { - return false; - } - - double mean_result = cactus_mean_all_f16(input.data(), input.size()); - double expected_mean = 4.5; - - if (std::abs(mean_result - expected_mean) > 1e-2) { - return false; - } - - return true; -} - -bool test_neon_transpose_fp16_correctness() { - const size_t M = 3, N = 4; - std::vector<__fp16> input = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; - std::vector<__fp16> result(M * N); - std::vector<__fp16> expected = {1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12}; - - size_t shape[] = {M, N}; - size_t perm[] = {1, 0}; - - cactus_transpose_f16(input.data(), result.data(), shape, perm, 2, 0, M); - - return TestUtils::compare_arrays(result.data(), expected.data(), M * N); -} - -bool test_neon_softmax_correctness() { - const size_t batch_size = 1, seq_len = 4, vocab_size = 3; - std::vector<__fp16> input = {1.0f, 2.0f, 3.0f, - 2.0f, 3.0f, 4.0f, - 3.0f, 4.0f, 5.0f, - 4.0f, 5.0f, 6.0f}; - std::vector<__fp16> result(input.size()); - - cactus_softmax_f16(input.data(), result.data(), batch_size, seq_len, vocab_size); - - for (size_t i = 0; i < seq_len; ++i) { - float row_sum = 0.0f; - for (size_t j = 0; j < vocab_size; ++j) { - row_sum += static_cast(result[i * vocab_size + j]); - } - if (std::abs(row_sum - 1.0f) > 1e-2f) { - return false; - } - } - - return true; -} - - -bool test_neon_rope_correctness() { - const size_t batch_size = 1, seq_len = 2, num_heads = 1, head_dim = 4; - const size_t start_pos = 0; - const float theta = 10000.0f; - const size_t total_elements = batch_size * seq_len * num_heads * head_dim; - - std::vector<__fp16> input(total_elements); - std::vector<__fp16> result(total_elements); - - TestUtils::fill_random_fp16(input); - - cactus_rope_f16(input.data(), result.data(), - batch_size, seq_len, num_heads, head_dim, start_pos, theta); - - bool different_from_input = false; - for (size_t i = 0; i < total_elements; ++i) { - if (std::abs(static_cast(result[i]) - static_cast(input[i])) > 1e-3f) { - different_from_input = true; - break; - } - } - - return different_from_input; -} - -bool test_neon_attention_fp16_correctness() { - const size_t batch_size = 1, seq_len = 2, num_heads = 1, head_dim = 8; - const float scale = 1.0f / sqrtf(static_cast(head_dim)); - const size_t total_elements = batch_size * seq_len * num_heads * head_dim; - - std::vector<__fp16> queries(total_elements); - std::vector<__fp16> keys(total_elements); - std::vector<__fp16> values(total_elements); - std::vector<__fp16> result(total_elements); - - TestUtils::fill_random_fp16(queries); - TestUtils::fill_random_fp16(keys); - TestUtils::fill_random_fp16(values); - - cactus_attention_f16(queries.data(), keys.data(), values.data(), result.data(), - batch_size, seq_len, seq_len, num_heads, num_heads, head_dim, scale, nullptr); - - bool has_non_zero = false; - for (size_t i = 0; i < total_elements; ++i) { - if (static_cast(result[i]) != 0) { - has_non_zero = true; - break; - } - } - - return has_non_zero; -} - -bool test_matmul_int8_grouped_correctness() { - // Kernel processes 64 K elements at a time, so group_size must be >= 64 - const size_t M = 2, K = 128, N = 4; - const size_t group_size = 64; - const size_t num_groups = K / group_size; - - std::vector<__fp16> A(M * K); - for (size_t i = 0; i < M * K; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 0.5f); - } - - std::vector B(N * K); - for (size_t i = 0; i < N * K; ++i) { - B[i] = static_cast((rand() % 128) - 64); - } - - std::vector<__fp16> B_scales(N * num_groups); - for (size_t n = 0; n < N; ++n) { - for (size_t g = 0; g < num_groups; ++g) { - float max_abs = 0.0f; - for (size_t k = 0; k < group_size; ++k) { - float val = std::abs(static_cast(B[n * K + g * group_size + k])); - if (val > max_abs) max_abs = val; - } - float scale = max_abs / 127.0f; - if (scale < 1e-6f) scale = 1e-6f; - B_scales[n * num_groups + g] = static_cast<__fp16>(scale); - } - } - - // Pre-quantize A to INT8 + scales - std::vector A_quant(M * K); - std::vector A_scales(M); - for (size_t m = 0; m < M; ++m) { - float max_abs = cactus_fp16_max_abs(A.data() + m * K, K); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - A_scales[m] = scale; - cactus_fp16_to_int8(A.data() + m * K, A_quant.data() + m * K, K, scale); - } - - std::vector<__fp16> C(M * N); - - cactus_matmul_int8(A_quant.data(), A_scales.data(), B.data(), B_scales.data(), C.data(), - M, K, N, group_size); - - std::vector C_ref(M * N, 0.0f); - for (size_t m = 0; m < M; ++m) { - float a_max_abs = 0.0f; - for (size_t k = 0; k < K; ++k) { - float val = std::abs(static_cast(A[m * K + k])); - if (val > a_max_abs) a_max_abs = val; - } - float a_scale = a_max_abs / 127.0f; - if (a_scale < 1e-10f) a_scale = 1e-10f; - - std::vector A_quant(K); - for (size_t k = 0; k < K; ++k) { - float val = static_cast(A[m * K + k]) / a_scale; - int32_t q = static_cast(std::round(val)); - q = std::max(-128, std::min(127, q)); - A_quant[k] = static_cast(q); - } - - for (size_t n = 0; n < N; ++n) { - float acc = 0.0f; - for (size_t g = 0; g < num_groups; ++g) { - float b_scale = static_cast(B_scales[n * num_groups + g]); - float combined_scale = a_scale * b_scale; - - int32_t group_sum = 0; - for (size_t k = 0; k < group_size; ++k) { - size_t k_idx = g * group_size + k; - group_sum += static_cast(A_quant[k_idx]) * - static_cast(B[n * K + k_idx]); - } - acc += static_cast(group_sum) * combined_scale; - } - C_ref[m * N + n] = acc; - } - } - - float max_abs_error = 0.0f; - for (size_t i = 0; i < M * N; ++i) { - float error = std::abs(static_cast(C[i]) - C_ref[i]); - if (error > max_abs_error) max_abs_error = error; - } - - return max_abs_error < 0.1f; -} - -bool test_matmul_int4_grouped_correctness() { - // Kernel processes 64 K elements at a time, so group_size must be >= 64 - const size_t M = 2, K = 128, N = 4; - const size_t group_size = 64; - const size_t num_groups = K / group_size; - const size_t K_packed = K / 2; - - std::vector<__fp16> A(M * K); - for (size_t i = 0; i < M * K; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 0.5f); - } - - std::vector B_unpacked(N * K); - for (size_t i = 0; i < N * K; ++i) { - B_unpacked[i] = static_cast((rand() % 16) - 8); - } - - std::vector B_packed(N * K_packed); - for (size_t n = 0; n < N; ++n) { - for (size_t k = 0; k < K; k += 2) { - int8_t lo = B_unpacked[n * K + k]; - int8_t hi = B_unpacked[n * K + k + 1]; - uint8_t packed = (static_cast(lo & 0x0F)) | (static_cast(hi & 0x0F) << 4); - B_packed[n * K_packed + k/2] = packed; - } - } - - std::vector<__fp16> B_scales(N * num_groups); - for (size_t n = 0; n < N; ++n) { - for (size_t g = 0; g < num_groups; ++g) { - float max_abs = 0.0f; - for (size_t k = 0; k < group_size; ++k) { - float val = std::abs(static_cast(B_unpacked[n * K + g * group_size + k])); - if (val > max_abs) max_abs = val; - } - float scale = max_abs / 7.0f; - if (scale < 1e-6f) scale = 1e-6f; - B_scales[n * num_groups + g] = static_cast<__fp16>(scale); - } - } - - // Pre-quantize A to INT8 + scales - std::vector A_quant(M * K); - std::vector A_scales(M); - for (size_t m = 0; m < M; ++m) { - float max_abs = cactus_fp16_max_abs(A.data() + m * K, K); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - A_scales[m] = scale; - cactus_fp16_to_int8(A.data() + m * K, A_quant.data() + m * K, K, scale); - } - - std::vector<__fp16> C(M * N); - - cactus_matmul_int4(A_quant.data(), A_scales.data(), B_packed.data(), B_scales.data(), C.data(), - M, K, N, group_size); - - std::vector C_ref(M * N, 0.0f); - for (size_t m = 0; m < M; ++m) { - float a_max_abs = 0.0f; - for (size_t k = 0; k < K; ++k) { - float val = std::abs(static_cast(A[m * K + k])); - if (val > a_max_abs) a_max_abs = val; - } - float a_scale = a_max_abs / 127.0f; - if (a_scale < 1e-10f) a_scale = 1e-10f; - - std::vector A_quant(K); - for (size_t k = 0; k < K; ++k) { - float val = static_cast(A[m * K + k]) / a_scale; - int32_t q = static_cast(std::round(val)); - q = std::max(-128, std::min(127, q)); - A_quant[k] = static_cast(q); - } - - for (size_t n = 0; n < N; ++n) { - float acc = 0.0f; - for (size_t g = 0; g < num_groups; ++g) { - float b_scale = static_cast(B_scales[n * num_groups + g]); - float combined_scale = a_scale * b_scale; - - int32_t group_sum = 0; - for (size_t k = 0; k < group_size; ++k) { - size_t k_idx = g * group_size + k; - group_sum += static_cast(A_quant[k_idx]) * - static_cast(B_unpacked[n * K + k_idx]); - } - acc += static_cast(group_sum) * combined_scale; - } - C_ref[m * N + n] = acc; - } - } - - float max_abs_error = 0.0f; - for (size_t i = 0; i < M * N; ++i) { - float error = std::abs(static_cast(C[i]) - C_ref[i]); - if (error > max_abs_error) max_abs_error = error; - } - - return max_abs_error < 0.1f; -} - -int main() { - TestUtils::TestRunner runner("Kernel Backend Tests"); - - runner.run_test("Kernel Add FP16 Correctness", test_neon_add_fp16_correctness()); - runner.run_test("Kernel Subtract FP16 Correctness", test_neon_subtract_fp16_correctness()); - runner.run_test("Kernel Multiply FP16 Correctness", test_neon_multiply_fp16_correctness()); - runner.run_test("Kernel Scalar Ops FP16 Correctness", test_neon_scalar_operations_fp16_correctness()); - runner.run_test("Kernel Reduction Correctness", test_neon_reduction_correctness()); - runner.run_test("Kernel Transpose FP16 Correctness", test_neon_transpose_fp16_correctness()); - runner.run_test("Kernel Softmax Correctness", test_neon_softmax_correctness()); - runner.run_test("Kernel RoPE Correctness", test_neon_rope_correctness()); - runner.run_test("Kernel Attention FP16 Correctness", test_neon_attention_fp16_correctness()); - runner.run_test("Kernel Grouped INT8 MatMul Correctness", test_matmul_int8_grouped_correctness()); - runner.run_test("Kernel Grouped INT4 MatMul Correctness", test_matmul_int4_grouped_correctness()); - - runner.print_summary(); - return runner.all_passed() ? 0 : 1; -} diff --git a/tests/test_kv_cache.cpp b/tests/test_kv_cache.cpp deleted file mode 100644 index 6de6cdb45..000000000 --- a/tests/test_kv_cache.cpp +++ /dev/null @@ -1,267 +0,0 @@ -#include "../cactus/cactus.h" -#include "test_utils.h" -#include -#include -#include -#include -#include -#include - -using namespace cactus::engine; -using namespace std; - -void fill_fp16(vector& buffer, size_t count, float value) { - buffer.resize(count * sizeof(__fp16)); - __fp16* ptr = reinterpret_cast<__fp16*>(buffer.data()); - for (size_t i = 0; i < count; i++) { - ptr[i] = static_cast<__fp16>(value); - } -} - -bool test_sliding_window_basic() { - const size_t num_layers = 2; - const size_t max_seq = 2048; - const size_t num_kv_heads = 8; - const size_t head_dim = 64; - const size_t window_size = 16; - const size_t sink_size = 4; - - KVCache cache; - cache.init(num_layers, max_seq, num_kv_heads, head_dim, Precision::INT8); - cache.set_window_size(window_size, sink_size); - - CactusGraph graph; - - { - size_t seq_len = 10; - vector k_nodes, v_nodes; - - for (size_t layer = 0; layer < num_layers; layer++) { - size_t k_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16); - size_t v_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16); - - vector k_data, v_data; - fill_fp16(k_data, seq_len * num_kv_heads * head_dim, layer + 1.0f); - fill_fp16(v_data, seq_len * num_kv_heads * head_dim, layer + 2.0f); - graph.set_input(k_node, k_data.data(), Precision::FP16); - graph.set_input(v_node, v_data.data(), Precision::FP16); - - k_nodes.push_back(k_node); - v_nodes.push_back(v_node); - } - - graph.execute(); - cache.update_from_graph(&graph, k_nodes, v_nodes, seq_len, num_layers, num_kv_heads, head_dim); - - assert(cache.get_effective_seq_len() == seq_len); - } - - { - size_t additional_seq = 6; - vector k_nodes, v_nodes; - - for (size_t layer = 0; layer < num_layers; layer++) { - size_t k_node = graph.input({additional_seq, num_kv_heads, head_dim}, Precision::FP16); - size_t v_node = graph.input({additional_seq, num_kv_heads, head_dim}, Precision::FP16); - - vector k_data, v_data; - fill_fp16(k_data, additional_seq * num_kv_heads * head_dim, layer + 10.0f); - fill_fp16(v_data, additional_seq * num_kv_heads * head_dim, layer + 20.0f); - graph.set_input(k_node, k_data.data(), Precision::FP16); - graph.set_input(v_node, v_data.data(), Precision::FP16); - - k_nodes.push_back(k_node); - v_nodes.push_back(v_node); - } - - graph.execute(); - cache.update_from_graph(&graph, k_nodes, v_nodes, additional_seq, num_layers, num_kv_heads, head_dim); - - assert(cache.get_effective_seq_len() == window_size); - } - - { - size_t additional_seq = 10; - vector k_nodes, v_nodes; - - float new_token_value = 100.0f; - - for (size_t layer = 0; layer < num_layers; layer++) { - size_t k_node = graph.input({additional_seq, num_kv_heads, head_dim}, Precision::FP16); - size_t v_node = graph.input({additional_seq, num_kv_heads, head_dim}, Precision::FP16); - - vector k_data, v_data; - fill_fp16(k_data, additional_seq * num_kv_heads * head_dim, new_token_value); - fill_fp16(v_data, additional_seq * num_kv_heads * head_dim, new_token_value); - graph.set_input(k_node, k_data.data(), Precision::FP16); - graph.set_input(v_node, v_data.data(), Precision::FP16); - - k_nodes.push_back(k_node); - v_nodes.push_back(v_node); - } - - graph.execute(); - cache.update_from_graph(&graph, k_nodes, v_nodes, additional_seq, num_layers, num_kv_heads, head_dim); - - assert(cache.get_effective_seq_len() == window_size); - - const int8_t* key_data = cache.get_keys_int8(0); - const float* key_scales = cache.get_key_scales(0); - if (key_data && key_scales) { - bool sink_preserved = true; - - float expected_sink_value = 1.0f; - float sink_scale = key_scales[0]; - int8_t expected_sink_q = static_cast(std::round(expected_sink_value / sink_scale)); - - for (size_t i = 0; i < min((size_t)8, sink_size * num_kv_heads * head_dim); i++) { - if (abs(key_data[i] - expected_sink_q) > 2) { - sink_preserved = false; - break; - } - } - assert(sink_preserved); - } - - assert(cache.get_total_seq_len() == 26); - } - - return true; -} - -bool test_incremental_updates() { - const size_t num_layers = 1; - const size_t num_kv_heads = 4; - const size_t head_dim = 32; - const size_t window_size = 8; - const size_t sink_size = 2; - - KVCache cache; - cache.init(num_layers, 2048, num_kv_heads, head_dim, Precision::FP16); - cache.set_window_size(window_size, sink_size); - - CactusGraph graph; - - for (size_t token = 0; token < 12; token++) { - vector k_nodes, v_nodes; - - size_t seq_len = 1; - size_t k_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16); - size_t v_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16); - - vector k_data, v_data; - fill_fp16(k_data, seq_len * num_kv_heads * head_dim, float(token + 1)); - fill_fp16(v_data, seq_len * num_kv_heads * head_dim, float(token + 101)); - - graph.set_input(k_node, k_data.data(), Precision::FP16); - graph.set_input(v_node, v_data.data(), Precision::FP16); - - k_nodes.push_back(k_node); - v_nodes.push_back(v_node); - - graph.execute(); - cache.update_from_graph(&graph, k_nodes, v_nodes, seq_len, num_layers, num_kv_heads, head_dim); - - size_t expected_len = min(token + 1, window_size); - assert(cache.get_effective_seq_len() == expected_len); - } - - assert(cache.get_effective_seq_len() == window_size); - assert(cache.get_total_seq_len() == 12); - - const int8_t* key_data = cache.get_keys_int8(0); - const float* key_scales = cache.get_key_scales(0); - assert(key_data != nullptr); - assert(key_scales != nullptr); - - return true; -} - -bool test_reset_functionality() { - KVCache cache; - cache.init(2, 1024, 8, 64, Precision::FP16); - cache.set_window_size(16, 4); - - CactusGraph graph; - vector k_nodes, v_nodes; - - for (size_t layer = 0; layer < 2; layer++) { - size_t k_node = graph.input({10, 8, 64}, Precision::FP16); - size_t v_node = graph.input({10, 8, 64}, Precision::FP16); - - vector data; - fill_fp16(data, 10 * 8 * 64, 1.0f); - graph.set_input(k_node, data.data(), Precision::FP16); - graph.set_input(v_node, data.data(), Precision::FP16); - - k_nodes.push_back(k_node); - v_nodes.push_back(v_node); - } - - graph.execute(); - cache.update_from_graph(&graph, k_nodes, v_nodes, 10, 2, 8, 64); - - assert(cache.get_effective_seq_len() == 10); - assert(cache.get_total_seq_len() == 10); - - cache.reset(); - - assert(cache.get_effective_seq_len() == 0); - assert(cache.get_total_seq_len() == 0); - assert(cache.get_keys_int8(0) == nullptr); - assert(cache.get_values_int8(0) == nullptr); - - return true; -} - -bool test_large_window() { - const size_t num_layers = 4; - const size_t num_kv_heads = 8; - const size_t head_dim = 64; - const size_t window_size = 512; - - KVCache cache; - cache.init(num_layers, 2048, num_kv_heads, head_dim, Precision::FP16); - cache.set_window_size(window_size, 4); - - CactusGraph graph; - - size_t seq_len = 600; - vector k_nodes, v_nodes; - - for (size_t layer = 0; layer < num_layers; layer++) { - size_t k_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16); - size_t v_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16); - - vector k_data, v_data; - fill_fp16(k_data, seq_len * num_kv_heads * head_dim, float(layer + 1)); - fill_fp16(v_data, seq_len * num_kv_heads * head_dim, float(layer + 101)); - - graph.set_input(k_node, k_data.data(), Precision::FP16); - graph.set_input(v_node, v_data.data(), Precision::FP16); - - k_nodes.push_back(k_node); - v_nodes.push_back(v_node); - } - - graph.execute(); - cache.update_from_graph(&graph, k_nodes, v_nodes, seq_len, num_layers, num_kv_heads, head_dim); - - assert(cache.get_effective_seq_len() == window_size); - assert(cache.get_total_seq_len() == seq_len); - - return true; -} - -int main() { - TestUtils::TestRunner runner("KV Cache Sliding Window Tests"); - runner.run_test("Basic Sliding Window", test_sliding_window_basic()); - runner.run_test("Incremental Updates", test_incremental_updates()); - runner.run_test("Reset Functionality", test_reset_functionality()); - runner.run_test("Large Window (512 tokens)", test_large_window()); - - cout << "────────────────────────────────────────────────────────────────────────────────────────\n"; - runner.print_summary(); - - return runner.all_passed() ? 0 : 1; -} \ No newline at end of file diff --git a/tests/test_performance.cpp b/tests/test_performance.cpp deleted file mode 100644 index a8573fa48..000000000 --- a/tests/test_performance.cpp +++ /dev/null @@ -1,1084 +0,0 @@ -#include "test_utils.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -struct BenchmarkConfig { - std::vector dimensions = {1024}; - std::vector precisions = {Precision::FP16}; - std::vector backends = {ComputeBackend::CPU}; - int iterations = 1; - - BenchmarkConfig() { - } -}; - -template -double time_operation(std::function operation, int iterations) { - auto start = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < iterations; ++i) { - operation(); - } - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - return duration.count() / 1000.0; -} - -template -void setup_random_data(std::vector& data) { - if constexpr (std::is_same_v) { - TestUtils::fill_random_int8(data); - } else if constexpr (std::is_same_v) { - TestUtils::fill_random_fp16(data); - } else { - TestUtils::fill_random_float(data); - } -} - -std::string precision_to_string(Precision prec) { - switch (prec) { - case Precision::INT8: return "INT8"; - case Precision::FP16: return "FP16"; - case Precision::FP32: return "FP32"; - default: return "UNKNOWN"; - } -} - -std::string backend_to_string(ComputeBackend backend) { - return (backend == ComputeBackend::CPU) ? "CPU" : "NPU"; -} - -double calculate_gflops(size_t ops, double time_ms) { - return ops / (time_ms * 1e6); -} - -void benchmark_streaming_stores(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - std::vector sizes = {1024*1024, 4*1024*1024}; - - for (size_t num_elements : sizes) { - std::vector<__fp16> A(num_elements), B(num_elements), C(num_elements); - for (size_t i = 0; i < num_elements; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - B[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - - double time_ms = time_operation<__fp16>([&]() { - cactus_add_f16(A.data(), B.data(), C.data(), num_elements); - }, config.iterations); - - double gb_per_sec = (num_elements * 2 * 3) / (time_ms * 1e6); - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance( - "Add (stream) " + std::to_string(num_elements / (1024*1024)) + "M elements", - details.str()); - } -} - -template -void benchmark_binary_elementwise_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - const std::vector>> ops = { - {"Add", [](CactusGraph& b, size_t a, size_t c) { return b.add(a, c); }}, - {"Subtract", [](CactusGraph& b, size_t a, size_t c) { return b.subtract(a, c); }}, - {"Multiply", [](CactusGraph& b, size_t a, size_t c) { return b.multiply(a, c); }}, - {"Divide", [](CactusGraph& b, size_t a, size_t c) { return b.divide(a, c); }} - }; - - Precision precision = TestUtils::default_precision(); - std::string prec_str = precision_to_string(precision); - - for (const auto& [op_name, op_func] : ops) { - for (size_t dim : config.dimensions) { - size_t total_elements = dim * dim; - - TestUtils::TestFixture fixture(op_name); - size_t input_a = fixture.create_input({dim, dim}, precision); - size_t input_b = fixture.create_input({dim, dim}, precision); - - std::vector data_a(total_elements), data_b(total_elements); - setup_random_data(data_a); - setup_random_data(data_b); - - fixture.set_input_data(input_a, data_a, precision); - fixture.set_input_data(input_b, data_b, precision); - - op_func(fixture.graph(), input_a, input_b); - - double time_ms = time_operation([&]() { - fixture.execute(); - }, config.iterations); - - double gflops = calculate_gflops(total_elements, time_ms); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance(op_name + " " + std::to_string(dim) + "x" + std::to_string(dim), - details.str()); - } - } -} - -void benchmark_conv1d_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - std::vector> shapes = { - {1, 3000, 80, 512}, - {1, 1500, 512, 512}, - }; - - for (const auto& [N, L, C_in, C_out] : shapes) { - size_t stride = 1; - size_t out_len = ((L - 1) / stride) + 1; - - std::vector<__fp16> input(N * C_in * L); - for (size_t i = 0; i < input.size(); ++i) { - input[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - - std::vector<__fp16> weight(C_out * C_in * 3); - for (size_t i = 0; i < weight.size(); ++i) { - weight[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 0.1f); - } - - std::vector<__fp16> output(N * C_out * out_len); - - double time_ms = time_operation<__fp16>([&]() { - cactus_conv1d_f16_k3(input.data(), weight.data(), output.data(), - N, L, C_in, C_out, stride); - }, config.iterations); - - size_t flops = 2ULL * N * out_len * C_out * C_in * 3; - double gflops = flops / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance( - "Conv1D k3 " + std::to_string(N) + "x" + std::to_string(L) + "x" + std::to_string(C_in) + "→" + std::to_string(C_out), - details.str()); - } -} - -void benchmark_broadcast_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - for (size_t dim : config.dimensions) { - size_t rows = dim; - size_t cols = dim; - size_t total_elements = rows * cols; - - std::vector<__fp16> A(total_elements); - for (size_t i = 0; i < total_elements; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - - std::vector<__fp16> B_vec(cols); - for (size_t i = 0; i < cols; ++i) { - B_vec[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 0.1f); - } - - std::vector<__fp16> C(total_elements); - - { - size_t a_strides[2] = {cols, 1}; - size_t b_strides[2] = {0, 1}; - size_t output_shape[2] = {rows, cols}; - - double time_ms = time_operation<__fp16>([&]() { - cactus_add_broadcast_f16(A.data(), B_vec.data(), C.data(), - a_strides, b_strides, output_shape, 2); - }, config.iterations); - - double gflops = calculate_gflops(total_elements, time_ms); - double gb_per_sec = (total_elements * 2 * 3) / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance( - "Broadcast Add " + std::to_string(dim) + "x" + std::to_string(dim) + "+[" + std::to_string(dim) + "]", - details.str()); - } - } -} - -template -void benchmark_scalar_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - const std::vector>> ops = { - {"Scalar Add", [](CactusGraph& b, size_t a) { return b.scalar_add(a, 2.5f); }}, - {"Scalar Sub", [](CactusGraph& b, size_t a) { return b.scalar_subtract(a, 2.5f); }}, - {"Scalar Mul", [](CactusGraph& b, size_t a) { return b.scalar_multiply(a, 2.5f); }}, - {"Scalar Div", [](CactusGraph& b, size_t a) { return b.scalar_divide(a, 2.5f); }}, - {"Scalar Exp", [](CactusGraph& b, size_t a) { return b.scalar_exp(a); }}, - {"Scalar Sqrt", [](CactusGraph& b, size_t a) { return b.scalar_sqrt(a); }}, - {"Scalar Cos", [](CactusGraph& b, size_t a) { return b.scalar_cos(a); }}, - {"Scalar Sin", [](CactusGraph& b, size_t a) { return b.scalar_sin(a); }} - }; - - Precision precision = TestUtils::default_precision(); - - for (const auto& [op_name, op_func] : ops) { - for (size_t dim : config.dimensions) { - size_t total_elements = dim * dim; - - TestUtils::TestFixture fixture(op_name); - size_t input_a = fixture.create_input({dim, dim}, precision); - - std::vector data_a(total_elements); - setup_random_data(data_a); - fixture.set_input_data(input_a, data_a, precision); - - op_func(fixture.graph(), input_a); - - double time_ms = time_operation([&]() { - fixture.execute(); - }, config.iterations); - - double gflops = calculate_gflops(total_elements, time_ms); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance(op_name + " " + std::to_string(dim) + "x" + std::to_string(dim), - details.str()); - } - } -} - -template -void benchmark_matmul_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - Precision precision = TestUtils::default_precision(); - - for (ComputeBackend backend : config.backends) { - std::string backend_str = backend_to_string(backend); - - for (size_t dim : config.dimensions) { - try { - TestUtils::TestFixture fixture("MatMul"); - size_t input_a = fixture.create_input({dim, dim}, precision); - size_t input_b = fixture.create_input({dim, dim}, precision); - - std::vector data_a(dim * dim), data_b(dim * dim); - setup_random_data(data_a); - setup_random_data(data_b); - - fixture.set_input_data(input_a, data_a, precision); - fixture.set_input_data(input_b, data_b, precision); - - fixture.graph().matmul(input_a, input_b, false, backend); - - double time_ms = time_operation([&]() { - fixture.execute(); - }, config.iterations); - - double gflops = calculate_gflops(2ULL * dim * dim * dim, time_ms); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance("MatMul " + std::to_string(dim) + "³ " + backend_str, - details.str()); - } catch (const std::exception& e) { - runner.log_performance("MatMul " + std::to_string(dim) + "³ " + backend_str, - "SKIP: " + std::string(e.what())); - } - } - } -} - -void benchmark_matmul_int8_grouped(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - const size_t group_size = 64; - - std::vector> shapes = { - {1, 1024, 1024}, - {1024, 1024, 1024}, - }; - - for (const auto& [M, K, N] : shapes) { - size_t K_aligned = ((K + group_size - 1) / group_size) * group_size; - size_t num_groups = K_aligned / group_size; - - std::vector<__fp16> A(M * K_aligned); - for (size_t i = 0; i < M * K_aligned; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - - std::vector B(N * K_aligned); - for (size_t i = 0; i < N * K_aligned; ++i) { - B[i] = static_cast((rand() % 256) - 128); - } - - std::vector<__fp16> B_scales(N * num_groups); - for (size_t i = 0; i < N * num_groups; ++i) { - B_scales[i] = static_cast<__fp16>(0.01f + (static_cast(rand()) / RAND_MAX) * 0.05f); - } - - std::vector A_quant(M * K_aligned); - std::vector A_scales(M); - for (size_t m = 0; m < M; ++m) { - float max_abs = cactus_fp16_max_abs(A.data() + m * K_aligned, K_aligned); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - A_scales[m] = scale; - cactus_fp16_to_int8(A.data() + m * K_aligned, A_quant.data() + m * K_aligned, K_aligned, scale); - } - - std::vector<__fp16> C(M * N); - - double time_ms = time_operation<__fp16>([&]() { - cactus_matmul_int8(A_quant.data(), A_scales.data(), B.data(), B_scales.data(), C.data(), - M, K_aligned, N, group_size); - }, config.iterations); - - double gflops = calculate_gflops(2ULL * M * K_aligned * N, time_ms); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance( - "MatMul INT8 " + std::to_string(M) + "x" + std::to_string(K_aligned) + "x" + std::to_string(N), - details.str()); - } -} - -void benchmark_matmul_int4_grouped(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - const size_t group_size = 64; - - std::vector> shapes = { - {1, 1024, 1024}, - {1024, 1024, 1024}, - }; - - for (const auto& [M, K, N] : shapes) { - size_t K_aligned = ((K + group_size - 1) / group_size) * group_size; - size_t num_groups = K_aligned / group_size; - size_t K_packed = K_aligned / 2; - - std::vector<__fp16> A(M * K_aligned); - for (size_t i = 0; i < M * K_aligned; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - - std::vector B_packed(N * K_packed); - for (size_t i = 0; i < N * K_packed; ++i) { - B_packed[i] = static_cast(rand() % 256); - } - - std::vector<__fp16> B_scales(N * num_groups); - for (size_t i = 0; i < N * num_groups; ++i) { - B_scales[i] = static_cast<__fp16>(0.01f + (static_cast(rand()) / RAND_MAX) * 0.05f); - } - - // Pre-quantize A to INT8 + scales - std::vector A_quant(M * K_aligned); - std::vector A_scales(M); - for (size_t m = 0; m < M; ++m) { - float max_abs = cactus_fp16_max_abs(A.data() + m * K_aligned, K_aligned); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - A_scales[m] = scale; - cactus_fp16_to_int8(A.data() + m * K_aligned, A_quant.data() + m * K_aligned, K_aligned, scale); - } - - std::vector<__fp16> C(M * N); - - double time_ms = time_operation<__fp16>([&]() { - cactus_matmul_int4(A_quant.data(), A_scales.data(), B_packed.data(), B_scales.data(), C.data(), - M, K_aligned, N, group_size); - }, config.iterations); - - double gflops = calculate_gflops(2ULL * M * K_aligned * N, time_ms); - - double bytes_loaded = M * K_aligned + N * K_packed + N * num_groups * 2 + M * N * 2; - double gb_per_sec = bytes_loaded / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance( - "MatMul INT4 " + std::to_string(M) + "x" + std::to_string(K_aligned) + "x" + std::to_string(N), - details.str()); - } -} - -template -void benchmark_unary_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - const std::vector>> ops = { - {"Transpose", [](CactusGraph& b, size_t a) { return b.transpose(a); }} - }; - - Precision precision = TestUtils::default_precision(); - - for (const auto& [op_name, op_func] : ops) { - for (size_t dim : config.dimensions) { - size_t total_elements = dim * dim; - - TestUtils::TestFixture fixture(op_name); - size_t input_a = fixture.create_input({dim, dim}, precision); - - std::vector data_a(total_elements); - setup_random_data(data_a); - fixture.set_input_data(input_a, data_a, precision); - - op_func(fixture.graph(), input_a); - - double time_ms = time_operation([&]() { - fixture.execute(); - }, config.iterations); - - double gb_per_sec = (total_elements * PrecisionTraits::size_of(precision) * 2) / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance(op_name + " " + std::to_string(dim) + "x" + std::to_string(dim), - details.str()); - } - } -} - -template -void benchmark_reduction_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - const std::vector>> ops = { - {"Sum", [](CactusGraph& b, size_t a) { return b.sum(a, -1); }}, - {"Mean", [](CactusGraph& b, size_t a) { return b.mean(a, -1); }}, - {"Variance", [](CactusGraph& b, size_t a) { return b.variance(a, -1); }}, - {"Min", [](CactusGraph& b, size_t a) { return b.min(a, -1); }}, - {"Max", [](CactusGraph& b, size_t a) { return b.max(a, -1); }} - }; - - Precision precision = TestUtils::default_precision(); - - for (const auto& [op_name, op_func] : ops) { - for (size_t dim : config.dimensions) { - size_t total_elements = dim * dim; - - TestUtils::TestFixture fixture(op_name); - size_t input_a = fixture.create_input({dim, dim}, precision); - - std::vector data_a(total_elements); - setup_random_data(data_a); - fixture.set_input_data(input_a, data_a, precision); - - op_func(fixture.graph(), input_a); - - double time_ms = time_operation([&]() { - fixture.execute(); - }, config.iterations); - - double gb_per_sec = (total_elements * PrecisionTraits::size_of(precision)) / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance(op_name + " " + std::to_string(dim) + "x" + std::to_string(dim), - details.str()); - } - } -} - -template -void benchmark_advanced_ops(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - Precision precision = TestUtils::default_precision(); - - for (size_t dim : config.dimensions) { - size_t total_elements = dim * dim; - - TestUtils::TestFixture fixture("Softmax"); - size_t input_a = fixture.create_input({dim, dim}, precision); - - std::vector data_a(total_elements); - setup_random_data(data_a); - fixture.set_input_data(input_a, data_a, precision); - - fixture.graph().softmax(input_a, -1); - - double time_ms = time_operation([&]() { - fixture.execute(); - }, config.iterations); - - double gb_per_sec = (total_elements * PrecisionTraits::size_of(precision)) / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance("Softmax " + std::to_string(dim) + "x" + std::to_string(dim), - details.str()); - } -} - -template -void benchmark_rms_norm(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - Precision precision = Precision::FP16; - - for (size_t dim : config.dimensions) { - size_t total_elements = dim * dim; - - CactusGraph graph; - size_t input_a = graph.input({dim, dim}, precision); - size_t weight = graph.input({dim}, precision); - - std::vector<__fp16> data_a(total_elements); - std::vector<__fp16> weight_data(dim, static_cast<__fp16>(1.0f)); - for (size_t i = 0; i < total_elements; ++i) { - data_a[i] = static_cast<__fp16>(static_cast(rand()) / RAND_MAX * 2.0f - 1.0f); - } - - graph.set_input(input_a, const_cast(static_cast(data_a.data())), precision); - graph.set_input(weight, const_cast(static_cast(weight_data.data())), precision); - - graph.rms_norm(input_a, weight); - - double time_ms = time_operation<__fp16>([&]() { - graph.execute(); - }, config.iterations); - - double gb_per_sec = (total_elements * 2) / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance("RMSNorm " + std::to_string(dim) + "x" + std::to_string(dim), - details.str()); - - graph.hard_reset(); - } -} - -template -void benchmark_rope(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - Precision precision = Precision::FP16; - - for (size_t dim : config.dimensions) { - size_t batch_size = 1; - size_t seq_len = dim / 4; - size_t num_heads = 4; - size_t head_dim = dim / 4; - size_t total_elements = batch_size * seq_len * num_heads * head_dim; - - TestUtils::FP16TestFixture fixture("RoPE"); - size_t input_a = fixture.create_input({batch_size, seq_len, num_heads, head_dim}, precision); - - std::vector<__fp16> data_a(total_elements); - for (size_t i = 0; i < total_elements; ++i) { - data_a[i] = static_cast<__fp16>(static_cast(rand()) / RAND_MAX * 2.0f - 1.0f); - } - fixture.set_input_data(input_a, data_a, precision); - - fixture.graph().rope(input_a, 10000.0f); - - double time_ms = time_operation<__fp16>([&]() { - fixture.execute(); - }, config.iterations); - - double gb_per_sec = (total_elements * 2 * 2) / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance("RoPE " + std::to_string(seq_len) + "x" + std::to_string(num_heads) + "x" + std::to_string(head_dim), - details.str()); - } -} - -template -void benchmark_attention(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - Precision precision = TestUtils::default_precision(); - - for (size_t dim : config.dimensions) { - size_t batch_size = 1; - size_t seq_len = 1024; - size_t num_heads = 16; - size_t head_dim = dim / 16; - size_t total_elements = batch_size * seq_len * num_heads * head_dim; - - CactusGraph graph; - size_t query = graph.input({batch_size, seq_len, num_heads, head_dim}, precision); - size_t key = graph.input({batch_size, seq_len, num_heads, head_dim}, precision); - size_t value = graph.input({batch_size, seq_len, num_heads, head_dim}, precision); - - std::vector q_data(total_elements), k_data(total_elements), v_data(total_elements); - setup_random_data(q_data); - setup_random_data(k_data); - setup_random_data(v_data); - - graph.set_input(query, const_cast(static_cast(q_data.data())), precision); - graph.set_input(key, const_cast(static_cast(k_data.data())), precision); - graph.set_input(value, const_cast(static_cast(v_data.data())), precision); - - float scale = 1.0f / sqrtf(static_cast(head_dim)); - graph.attention(query, key, value, scale); - - double time_ms = time_operation([&]() { - graph.execute(); - }, config.iterations); - - double gflops = calculate_gflops(2ULL * batch_size * num_heads * seq_len * seq_len * head_dim, time_ms); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance("Attention " + std::to_string(seq_len) + "x" + std::to_string(num_heads) + "x" + std::to_string(head_dim), - details.str()); - - graph.hard_reset(); - } -} - - -template -void benchmark_embedding_ops(TestUtils::TestRunner& runner, BenchmarkConfig& config) { - std::vector vocab_sizes = {65000}; - std::vector embedding_dims = {1024}; - std::vector sequence_lengths = {1000}; - - std::string precision_str = precision_to_string(TestUtils::default_precision()); - - for (size_t vocab_size : vocab_sizes) { - for (size_t embedding_dim : embedding_dims) { - for (size_t seq_len : sequence_lengths) { - CactusGraph graph; - - std::vector embeddings_data(vocab_size * embedding_dim); - setup_random_data(embeddings_data); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution dis(0, vocab_size - 1); - - std::vector indices_data(seq_len); - for (auto& idx : indices_data) { - int val = dis(gen); - idx = static_cast(std::min(val, 127)); - } - - Precision precision = TestUtils::default_precision(); - - size_t embeddings_id = graph.input({vocab_size, embedding_dim}, precision); - size_t indices_id = graph.input({seq_len}, Precision::INT8); - graph.embedding(embeddings_id, indices_id); - - graph.set_input(embeddings_id, embeddings_data.data(), precision); - graph.set_input(indices_id, indices_data.data(), Precision::INT8); - - double time_ms = time_operation([&]() { - graph.execute(); - }, config.iterations); - - double throughput = (seq_len * embedding_dim * sizeof(T)) / (time_ms * 1e3); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << throughput << " GB/s"; - runner.log_performance( - "Embedding " + std::to_string(vocab_size/1000) + "k vocab x" + - std::to_string(embedding_dim) + " dim " + precision_str, - details.str()); - } - } - } -} - -void benchmark_mmap_embedding(TestUtils::TestRunner& runner, BenchmarkConfig& config) { - std::vector vocab_sizes = {100}; - std::vector embedding_dims = {64}; - std::vector sequence_lengths = {32}; - - for (size_t vocab_size : vocab_sizes) { - for (size_t embedding_dim : embedding_dims) { - for (size_t seq_len : sequence_lengths) { - CactusGraph graph; - - std::vector<__fp16> embeddings_data(vocab_size * embedding_dim); - for (size_t i = 0; i < embeddings_data.size(); ++i) { - embeddings_data[i] = static_cast<__fp16>(static_cast(rand()) / RAND_MAX * 2.0f - 1.0f); - } - - size_t temp_embeddings = graph.input({vocab_size, embedding_dim}, Precision::FP16); - graph.set_input(temp_embeddings, embeddings_data.data(), Precision::FP16); - - std::filesystem::path tmpdir; - try { - tmpdir = std::filesystem::temp_directory_path(); - } catch (...) { - tmpdir = std::filesystem::current_path(); - } - std::filesystem::path tmpfile = tmpdir / (std::string("perf_embeddings_") + - std::to_string(vocab_size) + "_" + std::to_string(embedding_dim) + ".bin"); - const std::string temp_file = tmpfile.string(); - - GraphFile::save_node(graph, temp_embeddings, temp_file); - graph.hard_reset(); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution dis(0, vocab_size - 1); - - std::vector indices_data(seq_len); - for (auto& idx : indices_data) { - int val = dis(gen); - idx = static_cast(std::min(val, 127)); - } - - size_t indices_id = graph.input({seq_len}, Precision::INT8); - graph.embedding(temp_file, indices_id); - - graph.set_input(indices_id, indices_data.data(), Precision::INT8); - - double time_ms = time_operation<__fp16>([&]() { - graph.execute(); - }, config.iterations); - - double throughput = (seq_len * embedding_dim * sizeof(__fp16)) / (time_ms * 1e3); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << throughput << " GB/s"; - runner.log_performance( - "MMap Embedding " + std::to_string(vocab_size) + " vocab x" + - std::to_string(embedding_dim) + " dim", - details.str()); - - std::error_code ec; - std::filesystem::remove(tmpfile, ec); - } - } - } -} - -template -void benchmark_gather_ops(TestUtils::TestRunner& runner, BenchmarkConfig& config) { - std::string precision_str = precision_to_string(TestUtils::default_precision()); - - { - std::vector tensor_sizes = {127}; - std::vector index_counts = {132}; - - for (size_t tensor_size : tensor_sizes) { - for (size_t index_count : index_counts) { - CactusGraph graph; - - std::vector tensor_data(tensor_size); - setup_random_data(tensor_data); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution dis(0, tensor_size - 1); - - std::vector indices_data(index_count); - for (auto& idx : indices_data) { - int val = dis(gen); - idx = static_cast(std::min(val, 127)); - } - - Precision precision = TestUtils::default_precision(); - - size_t tensor_id = graph.input({tensor_size}, precision); - size_t indices_id = graph.input({index_count}, Precision::INT8); - graph.gather(tensor_id, indices_id); - - graph.set_input(tensor_id, tensor_data.data(), precision); - graph.set_input(indices_id, indices_data.data(), Precision::INT8); - - double time_ms = time_operation([&]() { - graph.execute(); - }, config.iterations); - - double throughput = (index_count * sizeof(T)) / (time_ms * 1e3); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << throughput << " GB/s"; - runner.log_performance( - "Gather 1D " + std::to_string(tensor_size) + "→" + std::to_string(index_count) + " " + precision_str, - details.str()); - } - } - } - - { - std::vector> tensor_shapes = {{64, 16, 8}}; - std::vector index_counts = {12}; - - for (const auto& shape : tensor_shapes) { - for (size_t index_count : index_counts) { - CactusGraph graph; - - size_t total_elements = shape[0] * shape[1] * shape[2]; - std::vector tensor_data(total_elements); - setup_random_data(tensor_data); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution dis(0, shape[0] - 1); - - std::vector indices_data(index_count); - for (auto& idx : indices_data) { - int val = dis(gen); - idx = static_cast(std::min(val, 127)); - } - - Precision precision = TestUtils::default_precision(); - - size_t tensor_id = graph.input(shape, precision); - size_t indices_id = graph.input({index_count}, Precision::INT8); - graph.gather(tensor_id, indices_id); - - graph.set_input(tensor_id, tensor_data.data(), precision); - graph.set_input(indices_id, indices_data.data(), Precision::INT8); - - double time_ms = time_operation([&]() { - graph.execute(); - }, config.iterations); - - double throughput = (index_count * shape[1] * shape[2] * sizeof(T)) / (time_ms * 1e3); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << throughput << " GB/s"; - runner.log_performance( - "Gather 3D " + std::to_string(shape[0]) + "x" + std::to_string(shape[1]) + "x" + std::to_string(shape[2]) + - "→" + std::to_string(index_count) + " " + precision_str, - details.str()); - } - } - } -} - -void benchmark_mel_filter_bank(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - using namespace cactus::engine; - - const size_t sampling_rate = 16000; - const size_t feature_size = 80; - const size_t num_frequency_bins = 201; - - AudioProcessor audio_proc; - - double time_ms = time_operation([&]() { - audio_proc.init_mel_filters(num_frequency_bins, feature_size, 0.0f, 8000.0f, sampling_rate); - }, config.iterations); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms"; - runner.log_performance( - "Mel Filter Bank " + std::to_string(feature_size) + "x" + std::to_string(num_frequency_bins), - details.str()); -} - -void benchmark_spectrogram(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - using namespace cactus::engine; - - const size_t n_fft = 400; - const size_t hop_length = 160; - const size_t chunk_length = 30; - const size_t sampling_rate = 16000; - const size_t feature_size = 80; - const size_t num_frequency_bins = 1 + n_fft / 2; - - const size_t n_samples = chunk_length * sampling_rate; - - AudioProcessor audio_proc; - audio_proc.init_mel_filters(num_frequency_bins, feature_size, 0.0f, 8000.0f, sampling_rate); - - std::vector waveform(n_samples); - for (size_t i = 0; i < n_samples; i++) { - waveform[i] = std::sin(2.0f * M_PI * 440.0f * i / sampling_rate); - } - - AudioProcessor::SpectrogramConfig spec_config; - spec_config.n_fft = n_fft; - spec_config.hop_length = hop_length; - spec_config.frame_length = n_fft; - spec_config.power = 2.0f; - spec_config.center = true; - spec_config.log_mel = "log10"; - - double time_ms = time_operation([&]() { - audio_proc.compute_spectrogram(waveform, spec_config); - }, config.iterations); - - const size_t pad_length = n_fft / 2; - const size_t padded_length = n_samples + 2 * pad_length; - const size_t num_frames = 1 + (padded_length - n_fft) / hop_length; - - double real_time_factor = (chunk_length * 1000.0) / time_ms; - - size_t bytes_processed = (n_samples * sizeof(float)) + (feature_size * num_frames * sizeof(float)); - double gb_per_sec = bytes_processed / (time_ms * 1e6); - - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(1) << real_time_factor << "x RT, " - << std::setprecision(2) << gb_per_sec << " GB/s"; - runner.log_performance( - "Spectrogram " + std::to_string(chunk_length) + "s (" + std::to_string(num_frames) + " frames)", - details.str()); -} - -void benchmark_gemm_f16_direct(TestUtils::TestRunner& runner, const BenchmarkConfig& config) { - std::vector> shapes = { - {1, 1024, 1024}, - {1024, 1024, 1024}, - }; - - for (const auto& [M, K, N] : shapes) { - std::vector<__fp16> A(M * K), B_T(N * K), C(M * N); - for (size_t i = 0; i < M * K; ++i) { - A[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - for (size_t i = 0; i < N * K; ++i) { - B_T[i] = static_cast<__fp16>((static_cast(rand()) / RAND_MAX - 0.5f) * 2.0f); - } - - double time_ms = time_operation<__fp16>([&]() { - cactus_matmul_f16(A.data(), B_T.data(), C.data(), M, K, N); - }, config.iterations); - - double gflops = calculate_gflops(2ULL * M * K * N, time_ms); - std::ostringstream details; - details << std::fixed << std::setprecision(3) << time_ms << "ms, " - << std::setprecision(2) << gflops << " GFLOPS"; - runner.log_performance("MatMul F16 " + std::to_string(M) + "x" + std::to_string(K) + "x" + std::to_string(N), - details.str()); - } -} - -bool test_gemm_f16_direct_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_gemm_f16_direct(runner, config); - return true; -} - -bool test_streaming_stores_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_streaming_stores(runner, config); - return true; -} - -bool test_conv1d_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_conv1d_ops(runner, config); - return true; -} - -bool test_broadcast_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_broadcast_ops(runner, config); - return true; -} - -bool test_binary_elementwise_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_binary_elementwise_ops<__fp16>(runner, config); - return true; -} - -bool test_scalar_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_scalar_ops<__fp16>(runner, config); - return true; -} - -bool test_matrix_multiplication_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_matmul_ops<__fp16>(runner, config); - return true; -} - -bool test_grouped_int8_matmul_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_matmul_int8_grouped(runner, config); - return true; -} - -bool test_grouped_int4_matmul_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_matmul_int4_grouped(runner, config); - return true; -} - -bool test_unary_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_unary_ops<__fp16>(runner, config); - return true; -} - -bool test_reduction_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_reduction_ops<__fp16>(runner, config); - return true; -} - -bool test_advanced_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_advanced_ops<__fp16>(runner, config); - return true; -} - -bool test_engine_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - benchmark_rms_norm<__fp16>(runner, config); - benchmark_rope<__fp16>(runner, config); - benchmark_attention<__fp16>(runner, config); - return true; -} - -bool test_gather_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - config.iterations = 10; - // INT8 for storage, FP16 for computation - benchmark_gather_ops(runner, config); - benchmark_gather_ops<__fp16>(runner, config); - return true; -} - -bool test_embedding_operations_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - config.iterations = 10; - // INT8 embeddings → FP16 output - benchmark_embedding_ops(runner, config); - benchmark_embedding_ops<__fp16>(runner, config); - benchmark_mmap_embedding(runner, config); - return true; -} - -bool test_signals_performance(TestUtils::TestRunner& runner) { - BenchmarkConfig config; - - benchmark_mel_filter_bank(runner, config); - benchmark_spectrogram(runner, config); - - return true; -} - - -int main() { - TestUtils::TestRunner runner("Performance Benchmarks"); - - runner.run_test("Streaming Stores", test_streaming_stores_performance(runner)); - runner.run_test("Conv1D Operations", test_conv1d_operations_performance(runner)); - runner.run_test("Broadcast Operations", test_broadcast_operations_performance(runner)); - runner.run_test("Binary Element-wise Operations", test_binary_elementwise_performance(runner)); - runner.run_test("Scalar Operations", test_scalar_operations_performance(runner)); - runner.run_test("Matrix Multiplication", test_matrix_multiplication_performance(runner)); - runner.run_test("F16 MatMul", test_gemm_f16_direct_performance(runner)); - runner.run_test("Grouped INT8 MatMul", test_grouped_int8_matmul_performance(runner)); - runner.run_test("Grouped INT4 MatMul", test_grouped_int4_matmul_performance(runner)); - runner.run_test("Unary Operations", test_unary_operations_performance(runner)); - runner.run_test("Reduction Operations", test_reduction_operations_performance(runner)); - runner.run_test("Advanced Operations", test_advanced_operations_performance(runner)); - runner.run_test("Engine Operations", test_engine_operations_performance(runner)); - runner.run_test("Gather Operations", test_gather_operations_performance(runner)); - runner.run_test("Embedding Operations", test_embedding_operations_performance(runner)); - runner.run_test("Signals Operations", test_signals_performance(runner)); - - runner.print_summary(); - return runner.all_passed() ? 0 : 1; -} \ No newline at end of file diff --git a/tests/test_threading.cpp b/tests/test_threading.cpp deleted file mode 100644 index 1074ec10e..000000000 --- a/tests/test_threading.cpp +++ /dev/null @@ -1,203 +0,0 @@ -#include "test_utils.h" -#include "../cactus/kernel/kernel_utils.h" -#include -#include -#include -#include -#include -#include -#include -#include - -struct BenchResult { - size_t M, K, N; - size_t total_tiles; - size_t num_threads; - double avg_time_ms; - double gflops; -}; - -std::vector benchmark_gemm_threading() { - const std::vector M_values = {1, 8, 64, 256, 1024}; - const std::vector K_values = {256, 1024}; - const std::vector N_values = {1024}; - const int iterations = 5; - const size_t group_size = 128; - - auto& pool = CactusThreading::get_thread_pool(); - size_t max_threads = pool.num_workers(); - std::vector thread_counts; - for (size_t t = 1; t <= max_threads; t++) { - thread_counts.push_back(t); - } - - #if defined(__APPLE__) && defined(__arm64__) - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 8; - #else - constexpr size_t TILE_M = 4; - constexpr size_t TILE_N = 4; - #endif - - std::mt19937 gen(42); - std::uniform_real_distribution float_dist(-1.0f, 1.0f); - std::uniform_int_distribution int_dist(-127, 127); - - std::vector results; - size_t total_dims = M_values.size() * K_values.size() * N_values.size(); - size_t dim_idx = 0; - - std::cout << "\n"; - std::cout << "┌─────────────────────────────────────────────────────────────────────────────────┐\n"; - std::cout << "│ GEMM INT8 M-Based Threading Benchmark │\n"; - std::cout << "├─────────────────────────────────────────────────────────────────────────────────┤\n"; - std::cout << "│ Testing thread counts: 1 to " << max_threads << std::setw(51) << "│\n"; - std::cout << "│ Dimensions: " << total_dims << " combos, " << iterations << " iterations each" - << std::setw(40) << "│\n"; - std::cout << "│ Tile size: " << TILE_M << "x" << TILE_N - << std::setw(60) << "│\n"; - std::cout << "└─────────────────────────────────────────────────────────────────────────────────┘\n"; - std::cout << "\n"; - - for (size_t M : M_values) { - for (size_t K : K_values) { - for (size_t N : N_values) { - dim_idx++; - size_t K_aligned = ((K + group_size - 1) / group_size) * group_size; - size_t num_groups = K_aligned / group_size; - - size_t num_row_tiles = (M + TILE_M - 1) / TILE_M; - size_t num_col_tiles = (N + TILE_N - 1) / TILE_N; - size_t total_tiles = num_row_tiles * num_col_tiles; - - std::vector<__fp16> A(M * K_aligned); - std::vector B(N * K_aligned); - - std::vector<__fp16> B_scales(N * num_groups); - std::vector<__fp16> C(M * N); - - for (size_t i = 0; i < M * K_aligned; ++i) { - A[i] = static_cast<__fp16>(float_dist(gen)); - } - for (size_t i = 0; i < N * K_aligned; ++i) { - B[i] = static_cast(int_dist(gen)); - } - for (size_t i = 0; i < N * num_groups; ++i) { - B_scales[i] = static_cast<__fp16>(0.01f + std::abs(float_dist(gen)) * 0.05f); - } - - // Pre-quantize A to INT8 + scales - std::vector A_quant(M * K_aligned); - std::vector A_scales(M); - for (size_t m = 0; m < M; ++m) { - float max_abs = cactus_fp16_max_abs(A.data() + m * K_aligned, K_aligned); - float scale = max_abs / 127.0f; - if (scale < 1e-10f) scale = 1e-10f; - A_scales[m] = scale; - cactus_fp16_to_int8(A.data() + m * K_aligned, A_quant.data() + m * K_aligned, K_aligned, scale); - } - - BenchResult best = {M, K, N, total_tiles, 0, std::numeric_limits::max(), 0}; - - std::cout << "[" << dim_idx << "/" << total_dims << "] " - << "M=" << std::setw(4) << M - << " K=" << std::setw(4) << K - << " N=" << std::setw(4) << N - << " (tiles=" << std::setw(5) << total_tiles << ") ... " << std::flush; - - for (size_t num_threads : thread_counts) { - CactusThreading::set_gemm_threads(num_threads); - - cactus_matmul_int8(A_quant.data(), A_scales.data(), B.data(), B_scales.data(), C.data(), - M, K_aligned, N, group_size); - - auto start = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < iterations; ++i) { - cactus_matmul_int8(A_quant.data(), A_scales.data(), B.data(), B_scales.data(), C.data(), - M, K_aligned, N, group_size); - } - auto end = std::chrono::high_resolution_clock::now(); - double avg_ms = std::chrono::duration(end - start).count() / iterations; - - if (avg_ms < best.avg_time_ms) { - best.num_threads = num_threads; - best.avg_time_ms = avg_ms; - } - } - - best.gflops = (2.0 * M * K_aligned * N) / (best.avg_time_ms * 1e6); - - std::cout << "threads=" << std::setw(2) << best.num_threads - << " -> " << std::fixed << std::setprecision(3) << best.avg_time_ms << "ms" - << " (" << std::setprecision(1) << best.gflops << " GFLOPS)\n"; - - results.push_back(best); - } - } - } - - CactusThreading::reset_gemm_threads(); - return results; -} - -void print_results_table(const std::vector& results) { - std::cout << "\n"; - std::cout << "┌────────────────────────────────────────────────────────────────────────────────┐\n"; - std::cout << "│ OPTIMAL THREAD COUNTS │\n"; - std::cout << "├──────┬──────┬──────┬────────┬─────────┬───────────┬──────────────────────────┤\n"; - std::cout << "│ M │ K │ N │ Tiles │ Threads │ Time(ms) │ GFLOPS │\n"; - std::cout << "├──────┼──────┼──────┼────────┼─────────┼───────────┼──────────────────────────┤\n"; - - for (const auto& r : results) { - std::cout << "│" << std::setw(5) << r.M << " " - << "│" << std::setw(5) << r.K << " " - << "│" << std::setw(5) << r.N << " " - << "│" << std::setw(7) << r.total_tiles << " " - << "│" << std::setw(8) << r.num_threads << " " - << "│" << std::setw(10) << std::fixed << std::setprecision(3) << r.avg_time_ms << " " - << "│" << std::setw(10) << std::setprecision(1) << r.gflops << " │\n"; - } - - std::cout << "└──────┴──────┴──────┴────────┴─────────┴───────────┴──────────────────────────┘\n"; - - std::cout << "\n"; - std::cout << "┌─────────────────────────────────────────────────────────────────────────────────┐\n"; - std::cout << "│ RECOMMENDED THREAD COUNTS BY BATCH SIZE │\n"; - std::cout << "└─────────────────────────────────────────────────────────────────────────────────┘\n"; - - std::map> by_batch; - for (const auto& r : results) { - by_batch[r.M].push_back(r); - } - - std::cout << "\n GemmThreading::get_num_threads(M, pool_size) should return:\n\n"; - - for (const auto& [M, batch_results] : by_batch) { - std::map thread_freq; - double total_gflops = 0; - for (const auto& r : batch_results) { - thread_freq[r.num_threads]++; - total_gflops += r.gflops; - } - - size_t best_threads = 0; - int max_freq = 0; - for (const auto& [t, f] : thread_freq) { - if (f > max_freq) { max_freq = f; best_threads = t; } - } - - std::string workload = (M <= 1) ? "decode" : (M <= 8) ? "small" : (M <= 64) ? "medium" : "prefill"; - std::cout << " M=" << std::setw(4) << M << " (" << std::setw(6) << workload << "): " - << std::setw(2) << best_threads << " threads" - << " (avg " << std::fixed << std::setprecision(1) - << (total_gflops / batch_results.size()) << " GFLOPS)\n"; - } - - std::cout << "\n"; -} - -int main() { - auto results = benchmark_gemm_threading(); - print_results_table(results); - return 0; -} diff --git a/tests/test_utils.cpp b/tests/test_utils.cpp deleted file mode 100644 index 29dde8b03..000000000 --- a/tests/test_utils.cpp +++ /dev/null @@ -1,457 +0,0 @@ -#include "test_utils.h" -#include -#include - -namespace TestUtils { - -static std::mt19937 gen(42); - -size_t random_graph_input(CactusGraph& graph, const std::vector& shape, Precision precision) { - size_t node_id = graph.input(shape, precision); - size_t total_elements = 1; - for (size_t dim : shape) total_elements *= dim; - - if (precision == Precision::INT8) { - std::uniform_int_distribution dist(-50, 50); - std::vector data(total_elements); - for (size_t i = 0; i < total_elements; ++i) data[i] = static_cast(dist(gen)); - graph.set_input(node_id, data.data(), precision); - } else { - std::uniform_real_distribution dist(-2.0f, 2.0f); - std::vector data(total_elements); - for (size_t i = 0; i < total_elements; ++i) data[i] = dist(gen); - graph.set_input(node_id, data.data(), precision); - } - return node_id; -} - -bool verify_graph_outputs(CactusGraph& graph, size_t node_a, size_t node_b, float tolerance) { - graph.execute(); - const auto& buffer_a = graph.get_output_buffer(node_a); - const auto& buffer_b = graph.get_output_buffer(node_b); - - if (buffer_a.shape != buffer_b.shape || buffer_a.precision != buffer_b.precision) return false; - - void* data_a = graph.get_output(node_a); - void* data_b = graph.get_output(node_b); - size_t total_elements = 1; - for (size_t dim : buffer_a.shape) total_elements *= dim; - - if (buffer_a.precision == Precision::INT8) { - const int8_t* ptr_a = static_cast(data_a); - const int8_t* ptr_b = static_cast(data_b); - for (size_t i = 0; i < total_elements; ++i) - if (std::abs(ptr_a[i] - ptr_b[i]) > tolerance) return false; - } else { - const float* ptr_a = static_cast(data_a); - const float* ptr_b = static_cast(data_b); - for (size_t i = 0; i < total_elements; ++i) - if (std::abs(ptr_a[i] - ptr_b[i]) > tolerance) return false; - } - - graph.hard_reset(); - return true; -} - -bool verify_graph_against_data(CactusGraph& graph, size_t node_id, const void* expected_data, size_t byte_size, float tolerance) { - graph.execute(); - void* actual_data = graph.get_output(node_id); - const auto& buffer = graph.get_output_buffer(node_id); - - if (buffer.precision == Precision::INT8) { - const int8_t* actual = static_cast(actual_data); - const int8_t* expected = static_cast(expected_data); - for (size_t i = 0; i < byte_size; ++i) - if (std::abs(actual[i] - expected[i]) > tolerance) return false; - } else { - const float* actual = static_cast(actual_data); - const float* expected = static_cast(expected_data); - size_t count = byte_size / sizeof(float); - for (size_t i = 0; i < count; ++i) - if (std::abs(actual[i] - expected[i]) > tolerance) return false; - } - - graph.hard_reset(); - return true; -} - -void fill_random_int8(std::vector& data) { - std::uniform_int_distribution dist(-50, 50); - for (auto& val : data) val = static_cast(dist(gen)); -} - -void fill_random_float(std::vector& data) { - std::uniform_real_distribution dist(-2.0f, 2.0f); - for (auto& val : data) val = dist(gen); -} - -void fill_random_fp16(std::vector<__fp16>& data) { - std::uniform_real_distribution dist(-2.0f, 2.0f); - for (auto& val : data) val = static_cast<__fp16>(dist(gen)); -} - -TestRunner::TestRunner(const std::string& suite_name) - : suite_name_(suite_name), passed_count_(0), total_count_(0) { - std::cout << "\n╔══════════════════════════════════════════════════════════════════════════════════════╗\n" - << "║ Running " << std::left << std::setw(76) << suite_name_ << " ║\n" - << "╚══════════════════════════════════════════════════════════════════════════════════════╝\n"; -} - -void TestRunner::run_test(const std::string& test_name, bool result) { - total_count_++; - if (result) { - passed_count_++; - std::cout << "✓ PASS │ " << std::left << std::setw(25) << test_name << "\n"; - } else { - std::cout << "✗ FAIL │ " << std::left << std::setw(25) << test_name << "\n"; - } -} - -void TestRunner::log_performance(const std::string& test_name, const std::string& details) { - std::cout << "⚡PERF │ " << std::left << std::setw(38) << test_name << " │ " << details << "\n"; -} - -void TestRunner::log_skip(const std::string& test_name, const std::string& reason) { - std::cout << "⊘ SKIP │ " << std::left << std::setw(25) << test_name << " │ " << reason << "\n"; -} - -void TestRunner::print_summary() { - std::cout << "────────────────────────────────────────────────────────────────────────────────────────\n"; - if (all_passed()) - std::cout << "✓ All " << total_count_ << " tests passed!\n"; - else - std::cout << "✗ " << (total_count_ - passed_count_) << " of " << total_count_ << " tests failed!\n"; - std::cout << "\n"; -} - -bool TestRunner::all_passed() const { - return passed_count_ == total_count_; -} - -bool test_basic_operation(const std::string& op_name, - std::function op_func, - const std::vector<__fp16>& data_a, - const std::vector<__fp16>& data_b, - const std::vector<__fp16>& expected, - const std::vector& shape) { - (void)op_name; - CactusGraph graph; - size_t input_a = graph.input(shape, Precision::FP16); - size_t input_b = graph.input(shape, Precision::FP16); - size_t result_id = op_func(graph, input_a, input_b); - - graph.set_input(input_a, const_cast(static_cast(data_a.data())), Precision::FP16); - graph.set_input(input_b, const_cast(static_cast(data_b.data())), Precision::FP16); - graph.execute(); - - __fp16* output = static_cast<__fp16*>(graph.get_output(result_id)); - for (size_t i = 0; i < expected.size(); ++i) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-2f) { - graph.hard_reset(); - return false; - } - } - graph.hard_reset(); - return true; -} - -bool test_scalar_operation(const std::string& op_name, - std::function op_func, - const std::vector<__fp16>& data, - float scalar, - const std::vector<__fp16>& expected, - const std::vector& shape) { - (void)op_name; - CactusGraph graph; - size_t input_a = graph.input(shape, Precision::FP16); - size_t result_id = op_func(graph, input_a, scalar); - - graph.set_input(input_a, const_cast(static_cast(data.data())), Precision::FP16); - graph.execute(); - - __fp16* output = static_cast<__fp16*>(graph.get_output(result_id)); - for (size_t i = 0; i < expected.size(); ++i) { - if (std::abs(static_cast(output[i]) - static_cast(expected[i])) > 1e-2f) { - graph.hard_reset(); - return false; - } - } - graph.hard_reset(); - return true; -} - -} - -namespace EngineTestUtils { - -Timer::Timer() : start(std::chrono::high_resolution_clock::now()) {} - -double Timer::elapsed_ms() const { - auto end = std::chrono::high_resolution_clock::now(); - return std::chrono::duration_cast(end - start).count() / 1000.0; -} - -double json_number(const std::string& json, const std::string& key, double def) { - std::string pattern = "\"" + key + "\":"; - size_t pos = json.find(pattern); - if (pos == std::string::npos) return def; - size_t start = pos + pattern.size(); - while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; - size_t end = start; - while (end < json.size() && std::string(",}] \t\n\r").find(json[end]) == std::string::npos) ++end; - try { return std::stod(json.substr(start, end - start)); } - catch (...) { return def; } -} - -std::string json_string(const std::string& json, const std::string& key) { - std::string pattern = "\"" + key + "\":"; - size_t pos = json.find(pattern); - if (pos == std::string::npos) return {}; - size_t start = pos + pattern.size(); - - while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; - if (start >= json.size() || json[start] != '"') return {}; - size_t q1 = start; - size_t q2 = json.find('"', q1 + 1); - if (q2 == std::string::npos) return {}; - return json.substr(q1 + 1, q2 - q1 - 1); -} - -std::string escape_json(const std::string& s) { - std::ostringstream o; - for (auto c : s) { - switch (c) { - case '"': o << "\\\""; break; - case '\\': o << "\\\\"; break; - case '\n': o << "\\n"; break; - case '\r': o << "\\r"; break; - default: o << c; break; - } - } - return o.str(); -} - -void stream_callback(const char* token, uint32_t token_id, void* user_data) { - auto* data = static_cast(user_data); - data->tokens.push_back(token ? token : ""); - data->token_ids.push_back(token_id); - data->token_count++; - - std::string out = token ? token : ""; - for (char& c : out) if (c == '\n') c = ' '; - std::cout << out << std::flush; - - if (data->stop_at > 0 && data->token_count >= data->stop_at) { - std::cout << " [-> stopped]" << std::flush; - cactus_stop(data->model); - } -} - -bool json_bool(const std::string& json, const std::string& key, bool def = false) { - std::string pattern = "\"" + key + "\":"; - size_t pos = json.find(pattern); - if (pos == std::string::npos) return def; - size_t start = pos + pattern.size(); - while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; - if (start + 4 <= json.size() && json.substr(start, 4) == "true") return true; - if (start + 5 <= json.size() && json.substr(start, 5) == "false") return false; - return def; -} - -std::string json_array(const std::string& json, const std::string& key) { - std::string pattern = "\"" + key + "\":"; - size_t pos = json.find(pattern); - if (pos == std::string::npos) return "[]"; - size_t start = pos + pattern.size(); - while (start < json.size() && (json[start] == ' ' || json[start] == '\t')) ++start; - if (start >= json.size() || json[start] != '[') return "[]"; - int depth = 1; - size_t end = start + 1; - while (end < json.size() && depth > 0) { - if (json[end] == '[') depth++; - else if (json[end] == ']') depth--; - end++; - } - return json.substr(start, end - start); -} - -void Metrics::parse(const std::string& json) { - success = json_bool(json, "success", false); - error = json_string(json, "error"); - cloud_handoff = json_bool(json, "cloud_handoff", false); - response = json_string(json, "response"); - function_calls = json_array(json, "function_calls"); - confidence = json_number(json, "confidence", -1.0); - ttft = json_number(json, "time_to_first_token_ms"); - total_ms = json_number(json, "total_time_ms"); - prefill_tps = json_number(json, "prefill_tps"); - decode_tps = json_number(json, "decode_tps"); - ram_mb = json_number(json, "ram_usage_mb"); - prefill_tokens = json_number(json, "prefill_tokens"); - completion_tokens = json_number(json, "decode_tokens"); - total_tokens = json_number(json, "total_tokens"); -} - -void Metrics::print_json() const { - std::cout << " \"success\": " << (success ? "true" : "false") << ",\n" - << " \"error\": " << (error.empty() ? "null" : "\"" + error + "\"") << ",\n" - << " \"cloud_handoff\": " << (cloud_handoff ? "true" : "false") << ",\n" - << " \"response\": \"" << response << "\",\n" - << " \"function_calls\": " << function_calls << ",\n" - << " \"confidence\": " << std::fixed << std::setprecision(4) << confidence << ",\n" - << " \"time_to_first_token_ms\": " << std::setprecision(2) << ttft << ",\n" - << " \"total_time_ms\": " << total_ms << ",\n" - << " \"prefill_tps\": " << prefill_tps << ",\n" - << " \"decode_tps\": " << decode_tps << ",\n" - << " \"ram_usage_mb\": " << ram_mb << ",\n" - << " \"prefill_tokens\": " << std::setprecision(0) << prefill_tokens << ",\n" - << " \"decode_tokens\": " << completion_tokens << ",\n" - << " \"total_tokens\": " << total_tokens << std::endl; -} - -} - -#ifdef HAVE_SDL2 - -AudioCapture::AudioCapture(int len_ms) - : m_len_ms(len_ms) - , m_running(false) - , m_dev_id_in(0) - , m_sdl_initialized(false) - , m_audio_pos(0) - , m_audio_len(0) - , m_total_samples_received(0) {} - -AudioCapture::~AudioCapture() { - if (m_dev_id_in) SDL_CloseAudioDevice(m_dev_id_in); - if (m_sdl_initialized) SDL_Quit(); -} - -bool AudioCapture::init(int capture_id, int sample_rate) { - static bool sdl_globally_initialized = false; - - if (!sdl_globally_initialized) { - if (SDL_Init(SDL_INIT_AUDIO) < 0) { - std::cerr << "SDL_Init failed: " << SDL_GetError() << std::endl; - return false; - } - sdl_globally_initialized = true; - m_sdl_initialized = true; - } - - SDL_SetHintWithPriority(SDL_HINT_AUDIO_RESAMPLING_MODE, "medium", SDL_HINT_OVERRIDE); - m_audio.resize((m_len_ms * sample_rate) / 1000); - - int num_devices = SDL_GetNumAudioDevices(SDL_TRUE); - std::cout << "\nAvailable audio capture devices:\n"; - for (int i = 0; i < num_devices; i++) - std::cout << " [" << i << "] " << SDL_GetAudioDeviceName(i, SDL_TRUE) << "\n"; - - if (capture_id >= num_devices) { - std::cerr << "Invalid capture device ID: " << capture_id << std::endl; - return false; - } - - std::cout << "Selected device: [" << capture_id << "] " - << SDL_GetAudioDeviceName(capture_id, SDL_TRUE) << "\n\n"; - - SDL_AudioSpec capture_spec_requested; - SDL_zero(capture_spec_requested); - capture_spec_requested.freq = sample_rate; - capture_spec_requested.format = AUDIO_F32; - capture_spec_requested.channels = 1; - capture_spec_requested.samples = 1024; - capture_spec_requested.callback = [](void* userdata, uint8_t* stream, int len) { - static_cast(userdata)->callback(stream, len); - }; - capture_spec_requested.userdata = this; - - SDL_AudioSpec capture_spec_obtained; - m_dev_id_in = SDL_OpenAudioDevice( - SDL_GetAudioDeviceName(capture_id, SDL_TRUE), - SDL_TRUE, &capture_spec_requested, &capture_spec_obtained, 0); - - if (!m_dev_id_in) { - std::cerr << "SDL_OpenAudioDevice failed: " << SDL_GetError() << std::endl; - return false; - } - - std::cout << "Audio capture initialized:\n" - << " Sample rate: " << capture_spec_obtained.freq << " Hz\n" - << " Channels: " << (int)capture_spec_obtained.channels << "\n" - << " Samples: " << capture_spec_obtained.samples << "\n" - << " Buffer length: " << m_len_ms << " ms\n"; - - return true; -} - -void AudioCapture::resume() { - std::lock_guard lock(m_mutex); - if (!m_running && m_dev_id_in) { - SDL_PauseAudioDevice(m_dev_id_in, 0); - m_running = true; - } -} - -void AudioCapture::pause() { - std::lock_guard lock(m_mutex); - if (m_running && m_dev_id_in) { - SDL_PauseAudioDevice(m_dev_id_in, 1); - m_running = false; - } -} - -void AudioCapture::clear() { - std::lock_guard lock(m_mutex); - m_audio_pos = 0; - m_audio_len = 0; -} - -size_t AudioCapture::get(int duration_ms, std::vector& result) { - std::lock_guard lock(m_mutex); - const size_t n_samples = (duration_ms * m_audio.size()) / m_len_ms; - if (n_samples > m_audio_len) return 0; - - result.resize(n_samples); - size_t start_pos = (m_audio_pos + m_audio.size() - m_audio_len) % m_audio.size(); - for (size_t i = 0; i < n_samples; i++) - result[i] = m_audio[(start_pos + i) % m_audio.size()]; - - m_audio_len = (m_audio_len > n_samples) ? (m_audio_len - n_samples) : 0; - return n_samples; -} - -size_t AudioCapture::get_all(std::vector& result) { - std::lock_guard lock(m_mutex); - if (m_audio_len == 0) return 0; - - result.resize(m_audio_len); - size_t start_pos = (m_audio_pos + m_audio.size() - m_audio_len) % m_audio.size(); - for (size_t i = 0; i < m_audio_len; i++) - result[i] = m_audio[(start_pos + i) % m_audio.size()]; - - size_t n_samples = m_audio_len; - m_audio_len = 0; - return n_samples; -} - -size_t AudioCapture::get_buffer_length() const { - std::lock_guard lock(m_mutex); - return m_audio_len; -} - -void AudioCapture::callback(uint8_t* stream, int len) { - const size_t n_samples = len / sizeof(float); - const float* samples = reinterpret_cast(stream); - if (!m_running) return; - - std::lock_guard lock(m_mutex); - for (size_t i = 0; i < n_samples; i++) { - m_audio[m_audio_pos] = samples[i]; - m_audio_pos = (m_audio_pos + 1) % m_audio.size(); - if (m_audio_len < m_audio.size()) m_audio_len++; - } - m_total_samples_received += n_samples; -} - -#endif diff --git a/tests/test_utils.h b/tests/test_utils.h deleted file mode 100644 index 08479dcc8..000000000 --- a/tests/test_utils.h +++ /dev/null @@ -1,247 +0,0 @@ -#ifndef TEST_UTILS_H -#define TEST_UTILS_H - -#include "../cactus/cactus.h" -#include "../cactus/ffi/cactus_ffi.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - - -namespace TestUtils { - -size_t random_graph_input(CactusGraph& graph, const std::vector& shape, Precision precision = Precision::INT8); -bool verify_graph_outputs(CactusGraph& graph, size_t node_a, size_t node_b, float tolerance = 1e-6f); -bool verify_graph_against_data(CactusGraph& graph, size_t node_id, const void* expected_data, size_t byte_size, float tolerance = 1e-6f); -void fill_random_int8(std::vector& data); -void fill_random_float(std::vector& data); - -template -double time_function(Func&& func, int iterations = 1) { - auto start = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < iterations; ++i) func(); - auto end = std::chrono::high_resolution_clock::now(); - return std::chrono::duration(end - start).count(); -} - -class TestRunner { -public: - TestRunner(const std::string& suite_name); - void run_test(const std::string& test_name, bool result); - void log_performance(const std::string& test_name, const std::string& details); - void log_skip(const std::string& test_name, const std::string& reason); - void print_summary(); - bool all_passed() const; - -private: - std::string suite_name_; - int passed_count_; - int total_count_; -}; - -template -constexpr Precision default_precision() { - if constexpr (std::is_same_v) return Precision::INT8; - else if constexpr (std::is_same_v) return Precision::FP16; - else return Precision::FP32; -} - -template -constexpr float default_tolerance() { - if constexpr (std::is_same_v) return 1e-2f; - else return 1e-6f; -} - -template -bool compare_arrays(const T* actual, const T* expected, size_t count, float tolerance = default_tolerance()) { - for (size_t i = 0; i < count; ++i) { - if constexpr (std::is_same_v) { - if (std::abs(static_cast(actual[i]) - static_cast(expected[i])) > tolerance) return false; - } else if constexpr (std::is_floating_point_v) { - if (std::abs(actual[i] - expected[i]) > tolerance) return false; - } else { - if (actual[i] != expected[i]) return false; - } - } - return true; -} - -template -class TestFixture { -public: - TestFixture(const std::string& = "") {} - ~TestFixture() { graph_.hard_reset(); } - - CactusGraph& graph() { return graph_; } - - size_t create_input(const std::vector& shape, Precision precision = default_precision()) { - return graph_.input(shape, precision); - } - - void set_input_data(size_t input_id, const std::vector& data, Precision precision = default_precision()) { - graph_.set_input(input_id, const_cast(static_cast(data.data())), precision); - } - - void execute() { graph_.execute(); } - - T* get_output(size_t node_id) { - return static_cast(graph_.get_output(node_id)); - } - - bool verify_output(size_t node_id, const std::vector& expected, float tolerance = default_tolerance()) { - return compare_arrays(get_output(node_id), expected.data(), expected.size(), tolerance); - } - -private: - CactusGraph graph_; -}; - -using Int8TestFixture = TestFixture; -using FP16TestFixture = TestFixture<__fp16>; - -void fill_random_fp16(std::vector<__fp16>& data); - -bool test_basic_operation(const std::string& op_name, - std::function op_func, - const std::vector<__fp16>& data_a, - const std::vector<__fp16>& data_b, - const std::vector<__fp16>& expected, - const std::vector& shape = {4}); - -bool test_scalar_operation(const std::string& op_name, - std::function op_func, - const std::vector<__fp16>& data, - float scalar, - const std::vector<__fp16>& expected, - const std::vector& shape = {4}); - -} - -namespace EngineTestUtils { - -struct Timer { - std::chrono::high_resolution_clock::time_point start; - Timer(); - double elapsed_ms() const; -}; - -double json_number(const std::string& json, const std::string& key, double def = 0.0); -std::string json_string(const std::string& json, const std::string& key); -std::string escape_json(const std::string& s); - -struct StreamingData { - std::vector tokens; - std::vector token_ids; - int token_count = 0; - cactus_model_t model = nullptr; - int stop_at = -1; -}; - -void stream_callback(const char* token, uint32_t token_id, void* user_data); - -struct Metrics { - bool success = false; - std::string error; - bool cloud_handoff = false; - std::string response; - std::string function_calls; - double confidence = -1.0; - double ttft = 0.0; - double total_ms = 0.0; - double prefill_tps = 0.0; - double decode_tps = 0.0; - double ram_mb = 0.0; - double prefill_tokens = 0.0; - double completion_tokens = 0.0; - double total_tokens = 0.0; - - void parse(const std::string& json); - void print_json() const; -}; - -template -bool run_test(const char* title, const char* model_path, const char* messages, - const char* options, TestFunc test_logic, - const char* tools = nullptr, int stop_at = -1, - const char* user_prompt = nullptr) { - std::cout << "\n╔══════════════════════════════════════════╗\n" - << "║" << std::setw(42) << std::left << std::string(" ") + title << "║\n" - << "╚══════════════════════════════════════════╝\n"; - - if (user_prompt) { - std::cout << "├─ User prompt: " << user_prompt << "\n"; - } - - cactus_model_t model = cactus_init(model_path, nullptr); - if (!model) { - std::cerr << "[✗] Failed to initialize model\n"; - return false; - } - - StreamingData data; - data.model = model; - data.stop_at = stop_at; - - char response[4096]; - std::cout << "Response: "; - - int result = cactus_complete(model, messages, response, sizeof(response), - options, tools, stream_callback, &data); - - std::cout << "\n\n[Results]\n"; - - Metrics metrics; - metrics.parse(response); - - bool success = test_logic(result, data, response, metrics); - std::cout << "└─ Status: " << (success ? "PASSED ✓" : "FAILED ✗") << std::endl; - - cactus_destroy(model); - return success; -} - -} - -#ifdef HAVE_SDL2 - -#include -#include - -class AudioCapture { -public: - AudioCapture(int len_ms = 10000); - ~AudioCapture(); - - bool init(int capture_id, int sample_rate); - void resume(); - void pause(); - void clear(); - size_t get(int duration_ms, std::vector& result); - size_t get_all(std::vector& result); - bool is_running() const { return m_running; } - size_t get_total_samples_received() const { return m_total_samples_received; } - size_t get_buffer_length() const; - -private: - void callback(uint8_t* stream, int len); - - int m_len_ms; - std::atomic m_running; - SDL_AudioDeviceID m_dev_id_in; - bool m_sdl_initialized; - std::vector m_audio; - size_t m_audio_pos; - size_t m_audio_len; - std::atomic m_total_samples_received; - mutable std::mutex m_mutex; -}; - -#endif - -#endif